Update app.py

app.py

@@ -1,239 +1,349 @@
-import gradio as gr
+#!/usr/bin/env python
+
+import os
+import re
+import tempfile
+from collections.abc import Iterator
 from threading import Thread
-import time
-from PIL import Image
-import torch
-import spaces
+
 import cv2
-import numpy as np
-from transformers import (
-     AutoProcessor, 
-     AutoModelForVision2Seq, 
-     TextIteratorStreamer
+import gradio as gr
+import spaces
+import torch
+from loguru import logger
+from PIL import Image
+from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextIteratorStreamer
+
+model_id = os.getenv("MODEL_ID", "google/gemma-3-12b-it")
+processor = AutoProcessor.from_pretrained(model_id, padding_side="left")
+model = Gemma3ForConditionalGeneration.from_pretrained(
+    model_id, device_map="auto", torch_dtype=torch.bfloat16, attn_implementation="eager"
 )
-# Helper function to return a progress bar HTML snippet.
-def progress_bar_html(label: str) -> str:
-    return f'''
-<div style="display: flex; align-items: center;">
-    <span style="margin-right: 10px; font-size: 14px;">{label}</span>
-    <div style="width: 110px; height: 5px; background-color: #FFB6C1; border-radius: 2px; overflow: hidden;">
-        <div style="width: 100%; height: 100%; background-color: #FF69B4; animation: loading 1.5s linear infinite;"></div>
-    </div>
-</div>
-<style>
-@keyframes loading {{
-    0% {{ transform: translateX(-100%); }}
-    100% {{ transform: translateX(100%); }}
-}}
-</style>
-    '''
-
-#adding examples
-examples=[
-        [{"text": "@video-infer Explain the content of the Advertisement", "files": ["examples/videoplayback.mp4"]}],
-        [{"text": "Explain the Image", "files": ["examples/3.jpg"]}],
-        [{"text": "@video-infer Explain the content of the video in detail", "files": ["examples/breakfast.mp4"]}],
-        [{"text": "@video-infer Describe the video", "files": ["examples/Missing.mp4"]}],
-        [{"text": "@video-infer Explain what is happening in this video ?", "files": ["examples/oreo.mp4"]}],
-        [{"text": "@video-infer Summarize the events in this video", "files": ["examples/sky.mp4"]}],
-        [{"text": "@video-infer What is in the video ?", "files": ["examples/redlight.mp4"]}],
-        [{"text": "Transcription of the letter", "files": ["examples/222.png"]}],
 
-]
+MAX_NUM_IMAGES = int(os.getenv("MAX_NUM_IMAGES", "5"))
+
+
+def count_files_in_new_message(paths: list[str]) -> tuple[int, int]:
+    image_count = 0
+    video_count = 0
+    for path in paths:
+        if path.endswith(".mp4"):
+            video_count += 1
+        else:
+            image_count += 1
+    return image_count, video_count
+
+
+def count_files_in_history(history: list[dict]) -> tuple[int, int]:
+    image_count = 0
+    video_count = 0
+    for item in history:
+        if item["role"] != "user" or isinstance(item["content"], str):
+            continue
+        if item["content"][0].endswith(".mp4"):
+            video_count += 1
+        else:
+            image_count += 1
+    return image_count, video_count
 
-# Helper: Downsample video to extract a fixed number of frames.
-def downsample_video(video_path, num_frames=22):
-    cap = cv2.VideoCapture(video_path)
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    # Calculate evenly spaced frame indices.
-    frame_indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
+
+def validate_media_constraints(message: dict, history: list[dict]) -> bool:
+    new_image_count, new_video_count = count_files_in_new_message(message["files"])
+    history_image_count, history_video_count = count_files_in_history(history)
+    image_count = history_image_count + new_image_count
+    video_count = history_video_count + new_video_count
+    if video_count > 1:
+        gr.Warning("Only one video is supported.")
+        return False
+    if video_count == 1:
+        if image_count > 0:
+            gr.Warning("Mixing images and videos is not allowed.")
+            return False
+        if "<image>" in message["text"]:
+            gr.Warning("Using <image> tags with video files is not supported.")
+            return False
+        # TODO: Add frame count validation for videos similar to image count limits  # noqa: FIX002, TD002, TD003
+    if video_count == 0 and image_count > MAX_NUM_IMAGES:
+        gr.Warning(f"You can upload up to {MAX_NUM_IMAGES} images.")
+        return False
+    if "<image>" in message["text"] and message["text"].count("<image>") != new_image_count:
+        gr.Warning("The number of <image> tags in the text does not match the number of images.")
+        return False
+    return True
+
+
+def downsample_video(video_path: str) -> list[tuple[Image.Image, float]]:
+    vidcap = cv2.VideoCapture(video_path)
+    fps = vidcap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    frame_interval = int(fps / 3)
     frames = []
-    for idx in frame_indices:
-        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
-        ret, frame = cap.read()
-        if ret:
-            # Convert BGR to RGB and then to a PIL image.
-            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            frame = Image.fromarray(frame)
-            frames.append(frame)
-    cap.release()
+
+    for i in range(0, total_frames, frame_interval):
+        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
+        success, image = vidcap.read()
+        if success:
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            pil_image = Image.fromarray(image)
+            timestamp = round(i / fps, 2)
+            frames.append((pil_image, timestamp))
+
+    vidcap.release()
     return frames
 
-# Load processor and model.
-processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-Instruct")
-model = AutoModelForVision2Seq.from_pretrained(
-    "HuggingFaceTB/SmolVLM-Instruct", 
-    torch_dtype=torch.bfloat16,
-).to("cuda")
-
-@spaces.GPU
-def model_inference(
-    input_dict, history, decoding_strategy, temperature, max_new_tokens,
-    repetition_penalty, top_p
-): 
-    text = input_dict["text"]
-    
-    # --- Video Inference Branch ---
-    if text.lower().startswith("@video-infer"):
-        # Remove the command prefix to get the prompt.
-        prompt_text = text[len("@video-infer"):].strip()
-        if not input_dict["files"]:
-            yield "Error: Please provide a video file for @video-infer."
-            return
-        # Assume the first file is a video.
-        video_file = input_dict["files"][0]
-        frames = downsample_video(video_file)
-        if not frames:
-            yield "Error: Could not extract frames from the video."
-            return
-        # Build a chat content: include the user prompt and then each frame labeled.
-        content = [{"type": "text", "text": prompt_text}]
-        for idx, frame in enumerate(frames):
-            content.append({"type": "text", "text": f"Frame {idx+1}:"})
-            content.append({"type": "image", "image": frame})
-        resulting_messages = [{
-            "role": "user",
-            "content": content
-        }]
-        prompt = processor.apply_chat_template(resulting_messages, add_generation_prompt=True)
-        # Process the extracted frames as images.
-        inputs = processor(text=prompt, images=[frames], return_tensors="pt")
-        inputs = {k: v.to("cuda") for k, v in inputs.items()}
-        
-        # Setup generation parameters.
-        generation_args = {
-            "max_new_tokens": max_new_tokens,
-            "repetition_penalty": repetition_penalty,
-        }
-        assert decoding_strategy in ["Greedy", "Top P Sampling"]
-        if decoding_strategy == "Greedy":
-            generation_args["do_sample"] = False
-        elif decoding_strategy == "Top P Sampling":
-            generation_args["temperature"] = temperature
-            generation_args["do_sample"] = True
-            generation_args["top_p"] = top_p
-        
-        generation_args.update(inputs)
-        
-        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-        generation_args = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
-        buffer = ""
-        thread = Thread(target=model.generate, kwargs=generation_args)
-        thread.start()
-        yield progress_bar_html("Processing Video with SmolVLM")
-        for new_text in streamer:
-            buffer += new_text
-            time.sleep(0.01)
-            yield buffer
+
+def process_video(video_path: str) -> list[dict]:
+    content = []
+    frames = downsample_video(video_path)
+    for frame in frames:
+        pil_image, timestamp = frame
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
+            pil_image.save(temp_file.name)
+            content.append({"type": "text", "text": f"Frame {timestamp}:"})
+            content.append({"type": "image", "url": temp_file.name})
+    logger.debug(f"{content=}")
+    return content
+
+
+def process_interleaved_images(message: dict) -> list[dict]:
+    logger.debug(f"{message['files']=}")
+    parts = re.split(r"(<image>)", message["text"])
+    logger.debug(f"{parts=}")
+
+    content = []
+    image_index = 0
+    for part in parts:
+        logger.debug(f"{part=}")
+        if part == "<image>":
+            content.append({"type": "image", "url": message["files"][image_index]})
+            logger.debug(f"file: {message['files'][image_index]}")
+            image_index += 1
+        elif part.strip():
+            content.append({"type": "text", "text": part.strip()})
+        elif isinstance(part, str) and part != "<image>":
+            content.append({"type": "text", "text": part})
+    logger.debug(f"{content=}")
+    return content
+
+
+def process_new_user_message(message: dict) -> list[dict]:
+    if not message["files"]:
+        return [{"type": "text", "text": message["text"]}]
+
+    if message["files"][0].endswith(".mp4"):
+        return [{"type": "text", "text": message["text"]}, *process_video(message["files"][0])]
+
+    if "<image>" in message["text"]:
+        return process_interleaved_images(message)
+
+    return [
+        {"type": "text", "text": message["text"]},
+        *[{"type": "image", "url": path} for path in message["files"]],
+    ]
+
+
+def process_history(history: list[dict]) -> list[dict]:
+    messages = []
+    current_user_content: list[dict] = []
+    for item in history:
+        if item["role"] == "assistant":
+            if current_user_content:
+                messages.append({"role": "user", "content": current_user_content})
+                current_user_content = []
+            messages.append({"role": "assistant", "content": [{"type": "text", "text": item["content"]}]})
+        else:
+            content = item["content"]
+            if isinstance(content, str):
+                current_user_content.append({"type": "text", "text": content})
+            else:
+                current_user_content.append({"type": "image", "url": content[0]})
+    return messages
+
+
+@spaces.GPU(duration=120)
+def run(message: dict, history: list[dict], system_prompt: str = "", max_new_tokens: int = 512) -> Iterator[str]:
+    if not validate_media_constraints(message, history):
+        yield ""
         return
 
-    # --- Default Image Inference Branch ---
-    # Process input images if provided.
-    if len(input_dict["files"]) > 1:
-        images = [Image.open(image).convert("RGB") for image in input_dict["files"]]
-    elif len(input_dict["files"]) == 1:
-        images = [Image.open(input_dict["files"][0]).convert("RGB")]
-    else:
-        images = []
-    
-    # Validate input.
-    if text == "" and not images:
-        gr.Error("Please input a query and optionally image(s).")
-    if text == "" and images:
-        gr.Error("Please input a text query along with the image(s).")
-    
-    resulting_messages = [{
-        "role": "user",
-        "content": [{"type": "image"} for _ in range(len(images))] + [
-            {"type": "text", "text": text}
-        ]
-    }]
-    prompt = processor.apply_chat_template(resulting_messages, add_generation_prompt=True)
-    inputs = processor(text=prompt, images=[images], return_tensors="pt")
-    inputs = {k: v.to("cuda") for k, v in inputs.items()}
-    
-    generation_args = {
-        "max_new_tokens": max_new_tokens,
-        "repetition_penalty": repetition_penalty,
-    }
-    assert decoding_strategy in ["Greedy", "Top P Sampling"]
-    if decoding_strategy == "Greedy":
-        generation_args["do_sample"] = False
-    elif decoding_strategy == "Top P Sampling":
-        generation_args["temperature"] = temperature
-        generation_args["do_sample"] = True
-        generation_args["top_p"] = top_p
-
-    generation_args.update(inputs)
-    
-    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-    generation_args = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
-    buffer = ""
-    thread = Thread(target=model.generate, kwargs=generation_args)
-    thread.start()
-    yield progress_bar_html("Processing Video with SmolVLM")
-    for new_text in streamer:
-        buffer += new_text
-        time.sleep(0.01)
-        yield buffer
-
-# Gradio ChatInterface: Allow both image and video file types.
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt}]})
+    messages.extend(process_history(history))
+    messages.append({"role": "user", "content": process_new_user_message(message)})
+
+    inputs = processor.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+    ).to(device=model.device, dtype=torch.bfloat16)
+
+    streamer = TextIteratorStreamer(processor, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        inputs,
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+    )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    output = ""
+    for delta in streamer:
+        output += delta
+        yield output
+
+
+examples = [
+    [
+        {
+            "text": "I need to be in Japan for 10 days, going to Tokyo, Kyoto and Osaka. Think about number of attractions in each of them and allocate number of days to each city. Make public transport recommendations.",
+            "files": [],
+        }
+    ],
+    [
+        {
+            "text": "Write the matplotlib code to generate the same bar chart.",
+            "files": ["assets/additional-examples/barchart.png"],
+        }
+    ],
+    [
+        {
+            "text": "What is odd about this video?",
+            "files": ["assets/additional-examples/tmp.mp4"],
+        }
+    ],
+    [
+        {
+            "text": "I already have this supplement <image> and I want to buy this one <image>. Any warnings I should know about?",
+            "files": ["assets/additional-examples/pill1.png", "assets/additional-examples/pill2.png"],
+        }
+    ],
+    [
+        {
+            "text": "Write a poem inspired by the visual elements of the images.",
+            "files": ["assets/sample-images/06-1.png", "assets/sample-images/06-2.png"],
+        }
+    ],
+    [
+        {
+            "text": "Compose a short musical piece inspired by the visual elements of the images.",
+            "files": [
+                "assets/sample-images/07-1.png",
+                "assets/sample-images/07-2.png",
+                "assets/sample-images/07-3.png",
+                "assets/sample-images/07-4.png",
+            ],
+        }
+    ],
+    [
+        {
+            "text": "Write a short story about what might have happened in this house.",
+            "files": ["assets/sample-images/08.png"],
+        }
+    ],
+    [
+        {
+            "text": "Create a short story based on the sequence of images.",
+            "files": [
+                "assets/sample-images/09-1.png",
+                "assets/sample-images/09-2.png",
+                "assets/sample-images/09-3.png",
+                "assets/sample-images/09-4.png",
+                "assets/sample-images/09-5.png",
+            ],
+        }
+    ],
+    [
+        {
+            "text": "Describe the creatures that would live in this world.",
+            "files": ["assets/sample-images/10.png"],
+        }
+    ],
+    [
+        {
+            "text": "Read text in the image.",
+            "files": ["assets/additional-examples/1.png"],
+        }
+    ],
+    [
+        {
+            "text": "When is this ticket dated and how much did it cost?",
+            "files": ["assets/additional-examples/2.png"],
+        }
+    ],
+    [
+        {
+            "text": "Read the text in the image into markdown.",
+            "files": ["assets/additional-examples/3.png"],
+        }
+    ],
+    [
+        {
+            "text": "Evaluate this integral.",
+            "files": ["assets/additional-examples/4.png"],
+        }
+    ],
+    [
+        {
+            "text": "caption this image",
+            "files": ["assets/sample-images/01.png"],
+        }
+    ],
+    [
+        {
+            "text": "What's the sign says?",
+            "files": ["assets/sample-images/02.png"],
+        }
+    ],
+    [
+        {
+            "text": "Compare and contrast the two images.",
+            "files": ["assets/sample-images/03.png"],
+        }
+    ],
+    [
+        {
+            "text": "List all the objects in the image and their colors.",
+            "files": ["assets/sample-images/04.png"],
+        }
+    ],
+    [
+        {
+            "text": "Describe the atmosphere of the scene.",
+            "files": ["assets/sample-images/05.png"],
+        }
+    ],
+]
+
+DESCRIPTION = """\
+<img src='https://huggingface.co/spaces/huggingface-projects/gemma-3-12b-it/resolve/main/assets/logo.png' id='logo' />
+
+This is a demo of Gemma 3 12B it, a vision language model with outstanding performance on a wide range of tasks.
+You can upload images, interleaved images and videos. Note that video input only supports single-turn conversation and mp4 input.
+"""
+
 demo = gr.ChatInterface(
-    fn=model_inference,
-    description="# **SmolVLM Video Infer `@video-infer for video understanding`**",
-    examples=examples,
-    textbox=gr.MultimodalTextbox(
-        label="Query Input", 
-        file_types=["image", "video"], 
-        file_count="multiple"
-    ),
-    stop_btn="Stop Generation",
+    fn=run,
+    type="messages",
+    chatbot=gr.Chatbot(type="messages", scale=1, allow_tags=["image"]),
+    textbox=gr.MultimodalTextbox(file_types=["image", ".mp4"], file_count="multiple", autofocus=True),
     multimodal=True,
     additional_inputs=[
-        gr.Radio(
-            ["Top P Sampling", "Greedy"],
-            value="Greedy",
-            label="Decoding strategy",
-            info="Higher values is equivalent to sampling more low-probability tokens.",
-        ),
-        gr.Slider(
-            minimum=0.0,
-            maximum=5.0,
-            value=0.4,
-            step=0.1,
-            interactive=True,
-            label="Sampling temperature",
-            info="Higher values will produce more diverse outputs.",
-        ),
-        gr.Slider(
-            minimum=8,
-            maximum=1024,
-            value=512,
-            step=1,
-            interactive=True,
-            label="Maximum number of new tokens to generate",
-        ),
-        gr.Slider(
-            minimum=0.01,
-            maximum=5.0,
-            value=1.2,
-            step=0.01,
-            interactive=True,
-            label="Repetition penalty",
-            info="1.0 is equivalent to no penalty",
-        ),
-        gr.Slider(
-            minimum=0.01,
-            maximum=0.99,
-            value=0.8,
-            step=0.01,
-            interactive=True,
-            label="Top P",
-            info="Higher values is equivalent to sampling more low-probability tokens.",
-        )
-    ],
-    cache_examples=False
+        gr.Textbox(label="System Prompt", value="You are a helpful assistant."),
+        gr.Slider(label="Max New Tokens", minimum=100, maximum=2000, step=10, value=700),
+    ],
+    stop_btn=False,
+    title="Gemma 3 12B IT",
+    description=DESCRIPTION,
+    examples=examples,
+    run_examples_on_click=False,
+    cache_examples=False,
+    css_paths="style.css",
+    delete_cache=(1800, 1800),
 )
 
-demo.launch(debug=True)
+if __name__ == "__main__":
+    demo.launch()