app.py

import spaces  # MUST come before torch / any CUDA-touching import
import torch
from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
from threading import Thread
from transformers import TextIteratorStreamer
import gradio as gr
from PIL import Image
import os

# Set decord as the video reader backend for qwen_vl_utils
os.environ.setdefault("QWEN_VL_VIDEO_READER_BACKEND", "decord")

MODEL_ID = "jdopensource/JoyAI-VL-Interaction-Preview"

print(f"Loading model {MODEL_ID} ...")
model = Qwen3VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda").eval()
processor = AutoProcessor.from_pretrained(MODEL_ID)
print("Model loaded.")


def _file_to_content_item(file_path: str):
    """Convert a file path into a qwen_vl_utils content item, i.e.
    {"type": "image"/"video", "image"/"video": ...}. Returns None if the
    file can't be opened as an image (and isn't a recognized video ext).
    """
    ext = os.path.splitext(file_path)[1].lower()
    if ext in (".mp4", ".avi", ".mov", ".mkv", ".webm"):
        return {
            "type": "video",
            "video": file_path,
            "max_pixels": 360 * 420,
            "fps": 1.0,
        }
    try:
        img = Image.open(file_path)
        return {"type": "image", "image": img}
    except Exception as e:
        print(f"Could not open {file_path} as image: {e}")
        return None


@spaces.GPU(duration=120)
def respond(
    message: dict,
    history: list,
    max_new_tokens: int = 1024,
    temperature: float = 0.7,
    top_p: float = 0.95,
):
    """Chat with JoyAI-VL-Interaction-Preview, a vision-driven VLM based on Qwen3-VL.

    Upload an image or short video, then ask a question about it.
    The model understands both static images and video clips.

    Args:
        message: dict with 'text' and 'files' keys from the multimodal chat input.
        history: list of {'role', 'content'} dicts in Gradio's messages format.
        max_new_tokens: maximum number of tokens to generate.
        temperature: sampling temperature.
        top_p: nucleus sampling probability.
    """
    text = message.get("text", "")
    files = message.get("files", [])

    # Build messages list from history.
    #
    # NOTE: Gradio's ChatInterface stores multimodal user turns with
    # `content` as a *list* whose items are EITHER a bare str (the typed
    # text) OR a dict shaped like {"path": "/tmp/xxx.jpg"} (one per
    # uploaded file) -- see gradio.chat_interface.ChatInterface.
    # _message_as_message_dict(). That is NOT the qwen_vl_utils shape
    # (which expects every content item to be a dict with a "type" key),
    # so it must be normalized here, item by item, rather than passed
    # through as-is. Assistant turns are always a plain str.
    messages = []
    for hist_item in history:
        role = hist_item["role"]
        content = hist_item["content"]
        new_content = []
        if isinstance(content, str):
            if content:
                new_content.append({"type": "text", "text": content})
        elif isinstance(content, list):
            for item in content:
                if isinstance(item, str):
                    if item:
                        new_content.append({"type": "text", "text": item})
                elif isinstance(item, dict):
                    if "path" in item:
                        file_item = _file_to_content_item(item["path"])
                        if file_item:
                            new_content.append(file_item)
                    elif item.get("type") in ("text", "image", "video"):
                        # Already in qwen_vl_utils shape (defensive).
                        new_content.append(item)
        if new_content:
            messages.append({"role": role, "content": new_content})

    # Build current user message content
    current_content = []
    for file_path in files:
        file_item = _file_to_content_item(file_path)
        if file_item:
            current_content.append(file_item)

    if text:
        current_content.append({"type": "text", "text": text})

    if not current_content:
        yield "Please provide a message with text and/or an image/video."
        return

    messages.append({"role": "user", "content": current_content})

    # Apply chat template
    chat_text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    # Process inputs - handle videos and images
    from qwen_vl_utils import process_vision_info

    image_inputs, video_inputs = process_vision_info(messages)

    inputs = processor(
        text=[chat_text],
        images=image_inputs if image_inputs else None,
        videos=video_inputs if video_inputs else None,
        padding=True,
        return_tensors="pt",
    ).to(model.device)

    # Stream output
    streamer = TextIteratorStreamer(
        processor.tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
        timeout=180,
    )

    thread = Thread(
        target=model.generate,
        kwargs=dict(
            **inputs,
            streamer=streamer,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
        ),
    )
    thread.start()

    full_text = ""
    for tok in streamer:
        full_text += tok
        yield full_text
    thread.join()


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # JoyAI-VL-Interaction-Preview

            An 8B vision-language model from JD, built on Qwen3-VL, that understands images and video clips.
            Upload an image or short video and ask any question about it.

            [Model Card](https://huggingface.co/jdopensource/JoyAI-VL-Interaction-Preview) · [Project Page](https://joyai-vl-video-future-academy-jd.github.io/JoyAI-VL-Interaction/)
            """
        )

    chat_interface = gr.ChatInterface(
        fn=respond,
        multimodal=True,
        chatbot=gr.Chatbot(height=600),
        title="JoyAI-VL-Interaction Chat",
        # `render=False`: keep these components un-rendered at construction
        # time so ChatInterface renders their accordion immediately after
        # the chat area (not before it -- see additional_inputs_accordion
        # below), and so gr.Examples (added after) can be placed below that
        # accordion, matching the requested page order: chat -> advanced
        # options accordion -> examples.
        additional_inputs=[
            gr.Slider(
                minimum=256, maximum=4096, value=1024, step=128,
                label="Max New Tokens", render=False,
            ),
            gr.Slider(
                minimum=0.1, maximum=2.0, value=0.7, step=0.1,
                label="Temperature", render=False,
            ),
            gr.Slider(
                minimum=0.1, maximum=1.0, value=0.95, step=0.05,
                label="Top-p", render=False,
            ),
        ],
        # `render=False` here too: this instance is only used to carry the
        # label/open kwargs (via ChatInterface's internal
        # `recover_kwargs`/`get_config`); it is never itself inserted into
        # the layout, so it doesn't create a stray empty accordion.
        additional_inputs_accordion=gr.Accordion(
            "Advanced options", open=False, render=False,
        ),
    )

    # Examples are added explicitly (instead of via ChatInterface's
    # `examples=` argument) so they render below the "Advanced options"
    # accordion instead of above it (ChatInterface's default footer order
    # is examples-then-accordion). Reuses ChatInterface's own example-cache
    # function so caching/streaming behavior is unchanged.
    gr.Examples(
        examples=[
            [{"text": "What do you see in this image?", "files": ["examples/sample_image.jpg"]}, 1024, 0.7, 0.95],
            [{"text": "Describe what happens in this video.", "files": ["examples/sample_video.mp4"]}, 1024, 0.7, 0.95],
        ],
        inputs=[chat_interface.textbox] + chat_interface.additional_inputs,
        outputs=chat_interface.chatbot,
        fn=chat_interface._examples_stream_fn if chat_interface.is_generator else chat_interface._examples_fn,
        cache_examples=True,
        cache_mode="lazy",
    )


if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)