app.py

import argparse
import gradio as gr
import os
from PIL import Image  # This was missing
import spaces
import copy
import numpy as np  # Required if you're doing image array work

from kimi_vl.serve.frontend import reload_javascript
from kimi_vl.serve.utils import (
    configure_logger,
    pil_to_base64,
    parse_ref_bbox,
    strip_stop_words,
    is_variable_assigned,
)
from kimi_vl.serve.gradio_utils import (
    cancel_outputing,
    delete_last_conversation,
    reset_state,
    reset_textbox,
    transfer_input,
    wrap_gen_fn,
)
from kimi_vl.serve.chat_utils import (
    generate_prompt_with_history,
    convert_conversation_to_prompts,
    to_gradio_chatbot,
    to_gradio_history,
)
from kimi_vl.serve.inference import kimi_vl_generate, load_model
from kimi_vl.serve.examples import get_examples

TITLE = """<h1 align="left" style="min-width:200px; margin-top:0;">Chat with Kimi-VL-A3B-Thinking🤔 </h1>"""
DESCRIPTION_TOP = """<a href="https://github.com/MoonshotAI/Kimi-VL" target="_blank">Kimi-VL-A3B-Thinking</a> is a multi-modal LLM that can understand text, video images, and generate text with thinking processes. \n this specific space was hacked to also accept videos, and the system prompt has been changed to favor video analysis."""
DESCRIPTION = """"""
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
DEPLOY_MODELS = dict()
logger = configure_logger()

def resize_image(image: Image.Image, max_size: int = 640, min_size: int = 28):
    width, height = image.size
    if width < min_size or height < min_size:
        scale = min_size / min(width, height)
        new_width = int(width * scale)
        new_height = int(height * scale)
        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    elif max_size > 0 and (width > max_size or height > max_size):
        scale = max_size / max(width, height)
        new_width = int(width * scale)
        new_height = int(height * scale)
        image = image.resize((new_width, new_height))
    return image

def load_frames(video_file, max_num_frames=64, long_edge=448):
    from decord import VideoReader
    vr = VideoReader(video_file)
    duration = len(vr)
    fps = vr.get_avg_fps()
    length = int(duration / fps)
    num_frames = min(max_num_frames, length)

    frame_indices = [int(duration / num_frames * (i + 0.5)) for i in range(num_frames)]
    frames_data = vr.get_batch(frame_indices).asnumpy()

    imgs = []
    for idx in range(num_frames):
        img = resize_image(Image.fromarray(frames_data[idx]).convert("RGB"), long_edge)
        imgs.append(img)

    return imgs

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, default="Kimi-VL-A3B-Thinking")
    parser.add_argument("--local-path", type=str, default="", help="huggingface ckpt, optional")
    parser.add_argument("--ip", type=str, default="0.0.0.0")
    parser.add_argument("--port", type=int, default=7860)
    return parser.parse_args()

def fetch_model(model_name: str):
    global args, DEPLOY_MODELS
    if args.local_path:
        model_path = args.local_path
    else:
        model_path = f"moonshotai/{args.model}"

    if model_name in DEPLOY_MODELS:
        model_info = DEPLOY_MODELS[model_name]
        print(f"{model_name} has been loaded.")
    else:
        print(f"{model_name} is loading...")
        DEPLOY_MODELS[model_name] = load_model(model_path)
        print(f"Load {model_name} successfully...")
        model_info = DEPLOY_MODELS[model_name]
    return model_info

def highlight_thinking(msg: str) -> str:
    msg = copy.deepcopy(msg)
    if "◁think▷" in msg:
        msg = msg.replace("◁think▷", "<b style='color:blue;'>🤔Thinking...</b>\n")
    if "◁/think▷" in msg:
        msg = msg.replace("◁/think▷", "\n<b style='color:purple;'>💡Summary</b>\n")
    return msg

@wrap_gen_fn
@spaces.GPU(duration=180)
def predict(
    text,
    images,
    chatbot,
    history,
    top_p,
    temperature,
    max_length_tokens,
    max_context_length_tokens,
    video_num_frames,
    video_long_edge,
    chunk_size: int = 512,
):
    print("running the prediction function")
    try:
        model, processor = fetch_model(args.model)
        if text == "":
            yield chatbot, history, "Empty context."
            return
    except KeyError:
        yield [[text, "No Model Found"]], [], "No Model Found"
        return

    if images is None:
        images = []

    pil_images = []
    for img_or_file in images:
        try:
            if isinstance(img_or_file, Image.Image):
                pil_images.append(img_or_file)
            else:
                image = Image.open(img_or_file.name).convert("RGB")
                pil_images.append(image)
        except:
            try:
                pil_images = load_frames(img_or_file, video_num_frames, video_long_edge)
                break
            except Exception as e:
                print(f"Error loading image or video: {e}")

    conversation = generate_prompt_with_history(
        text=text,
        images=pil_images,
        history=history,
        processor=processor,
        max_length=max_context_length_tokens,
    )

    all_conv, last_image = convert_conversation_to_prompts(conversation)
    stop_words = conversation.stop_str
    gradio_chatbot_output = to_gradio_chatbot(conversation)

    full_response = ""
    for x in kimi_vl_generate(
        conversations=all_conv,
        model=model,
        processor=processor,
        stop_words=stop_words,
        max_length=max_length_tokens,
        temperature=temperature,
        top_p=top_p,
    ):
        full_response += x
        response = strip_stop_words(full_response, stop_words)
        conversation.update_last_message(response)
        gradio_chatbot_output[-1][1] = highlight_thinking(response)
        yield gradio_chatbot_output, to_gradio_history(conversation), "Generating..."

    if last_image is not None:
        vg_image = parse_ref_bbox(response, last_image)
        if vg_image is not None:
            vg_base64 = pil_to_base64(vg_image, "vg", max_size=800, min_size=400)
            gradio_chatbot_output[-1][1] += vg_base64
            yield gradio_chatbot_output, to_gradio_history(conversation), "Generating..."

    logger.info("flushed result to gradio")

    if is_variable_assigned("x"):
        print(f"temperature: {temperature}, top_p: {top_p}, max_length_tokens: {max_length_tokens}")

    yield gradio_chatbot_output, to_gradio_history(conversation), "Generate: Success"

# ------------------------------
# Interface + Launch
# ------------------------------

if __name__ == "__main__":
    args = parse_args()
    reload_javascript()

    with gr.Blocks(title="Kimi-VL") as demo:
        gr.Markdown(TITLE)
        gr.Markdown(DESCRIPTION_TOP)

        with gr.Row():
            text_input = gr.Textbox(label="Enter your message", scale=4)
            image_input = gr.File(label="Upload image or video", file_types=["image", "video"], file_count="multiple")

        chatbot_output = gr.Chatbot(label="Kimi-VL Output")
        history_state = gr.State([])
        top_p = gr.Slider(0, 1, value=0.9, label="Top-p")
        temperature = gr.Slider(0.1, 1.5, value=0.6, label="Temperature")
        max_length_tokens = gr.Slider(16, 4096, value=2048, step=64, label="Max Length")
        max_context_length_tokens = gr.Slider(128, 4096, value=2048, step=64, label="Max Context")
        video_num_frames = gr.Slider(4, 64, value=24, step=4, label="Frames (for video)")
        video_long_edge = gr.Slider(128, 1024, value=1024, step=32, label="Long edge resize (video)")

        submit_btn = gr.Button("Submit")

        submit_btn.click(
            predict,
            inputs=[
                text_input,
                image_input,
                chatbot_output,
                history_state,
                top_p,
                temperature,
                max_length_tokens,
                max_context_length_tokens,
                video_num_frames,
                video_long_edge,
            ],
            outputs=[chatbot_output, history_state, gr.Textbox(visible=False)],
        )

    demo.queue().launch(server_name=args.ip, server_port=args.port)