Update app.py

app.py

@@ -6,10 +6,20 @@ import time
 import torch
 import spaces
 import subprocess
+import uuid
+import cv2
+import numpy as np
+from PIL import Image
+from io import BytesIO
 
-# Install flash-attn with no CUDA build
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+# Install flash-attn
+subprocess.run(
+    'pip install flash-attn --no-build-isolation',
+    env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
+    shell=True
+)
 
+# Load processor and model.
 processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct")
 model = AutoModelForImageTextToText.from_pretrained(
     "HuggingFaceTB/SmolVLM2-2.2B-Instruct", 
@@ -17,66 +27,112 @@ model = AutoModelForImageTextToText.from_pretrained(
     torch_dtype=torch.bfloat16
 ).to("cuda:0")
 
+def downsample_video(video_path):
+    """
+    Extracts 10 evenly spaced frames from the video at video_path.
+    Each frame is converted from BGR to RGB and returned as a PIL Image.
+    """
+    vidcap = cv2.VideoCapture(video_path)
+    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = vidcap.get(cv2.CAP_PROP_FPS)
+    frames = []
+    if total_frames <= 0 or fps <= 0:
+        vidcap.release()
+        return frames
+    frame_indices = np.linspace(0, total_frames - 1, 10, dtype=int)
+    for i in frame_indices:
+        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
+        success, frame = vidcap.read()
+        if success:
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            pil_image = Image.fromarray(frame)
+            frames.append((pil_image, round(i / fps, 2)))
+    vidcap.release()
+    return frames
+
 @spaces.GPU
-def model_inference(input_dict, history, max_tokens): 
-    text = input_dict.get("text", "").strip()
-    media_queue = []
+def model_inference(input_dict, history, max_tokens):
+    text = input_dict["text"]
     user_content = []
-
-    # Process uploaded media files
+    media_queue = []
+    
+    # Process input files.
     for file in input_dict.get("files", []):
         if file.endswith((".png", ".jpg", ".jpeg", ".gif", ".bmp")):
             media_queue.append({"type": "image", "path": file})
         elif file.endswith((".mp4", ".mov", ".avi", ".mkv", ".flv")):
-            media_queue.append({"type": "video", "path": file})
-
-    # Construct user content with placeholders
-    if "<image>" in text or "<video>" in text:
-        parts = re.split(r'(<image>|<video>)', text)  
-        for part in parts:
-            if part == "<image>" and media_queue:
-                user_content.append(media_queue.pop(0)) 
-            elif part == "<video>" and media_queue:
-                user_content.append(media_queue.pop(0))  
-            elif part.strip():
-                user_content.append({"type": "text", "text": part.strip()})
+            # Extract frames from video using OpenCV.
+            frames = downsample_video(file)
+            for frame, timestamp in frames:
+                temp_file = f"video_frame_{uuid.uuid4().hex}.png"
+                frame.save(temp_file)
+                media_queue.append({"type": "image", "path": temp_file})
+    
+    # Build the conversation messages.
+    if not history:
+        text = text.strip()
+        # Use only the "<image>" token for inserting images.
+        if "<image>" in text:
+            parts = re.split(r'(<image>)', text)
+            for part in parts:
+                if part == "<image>" and media_queue:
+                    user_content.append(media_queue.pop(0))
+                elif part.strip():
+                    user_content.append({"type": "text", "text": part.strip()})
+        else:
+            user_content.append({"type": "text", "text": text})
+            for media in media_queue:
+                user_content.append(media)
+        resulting_messages = [{"role": "user", "content": user_content}]
     else:
-        user_content.append({"type": "text", "text": text})
-        user_content.extend(media_queue)  
-
-    resulting_messages = [{"role": "user", "content": user_content}]
-
-    # Process history
-    if history:
+        resulting_messages = []
+        user_content = []
+        media_queue = []
+        # Process history: now only image files are expected.
         for hist in history:
-            if hist["role"] == "user":
-                if isinstance(hist["content"], tuple) and len(hist["content"]) > 0:
-                    file_name = hist["content"][0]
-                    if file_name.endswith((".png", ".jpg", ".jpeg")):
-                        media_queue.append({"type": "image", "path": file_name})
-                    elif file_name.endswith(".mp4"):
-                        media_queue.append({"type": "video", "path": file_name})
-
+            if hist["role"] == "user" and isinstance(hist["content"], tuple): 
+                file_name = hist["content"][0]
+                if file_name.endswith((".png", ".jpg", ".jpeg", ".gif", ".bmp")):
+                    media_queue.append({"type": "image", "path": file_name})
+        for hist in history:
+            if hist["role"] == "user" and isinstance(hist["content"], str):
+                text = hist["content"]
+                parts = re.split(r'(<image>)', text)
+                for part in parts:
+                    if part == "<image>" and media_queue:
+                        user_content.append(media_queue.pop(0))
+                    elif part.strip():
+                        user_content.append({"type": "text", "text": part.strip()})
             elif hist["role"] == "assistant":
-                resulting_messages.append({"role": "assistant", "content": [{"type": "text", "text": hist["content"]}]})
-
-    if not text and not media_queue:
-        gr.Warning("Please provide text or an image/video.")
-
-    # Process inputs
+                resulting_messages.append({
+                    "role": "user",
+                    "content": user_content
+                })
+                resulting_messages.append({
+                    "role": "assistant",
+                    "content": [{"type": "text", "text": hist["content"]}]
+                })
+                user_content = []
+    
+    if text == "":
+        gr.Error("Please input a query and optionally image(s).")
+    
+    print("resulting_messages", resulting_messages)
     inputs = processor.apply_chat_template(
         resulting_messages,
         add_generation_prompt=True,
         tokenize=True,
         return_dict=True,
-        return_tensors="pt"
-    ).to(model.device, dtype=torch.bfloat16)  # Ensure dtype consistency
-
-    # Generate output
-    streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
-    thread = Thread(target=model.generate, kwargs={"input_ids": inputs["input_ids"], "streamer": streamer, "max_new_tokens": max_tokens})
+        return_tensors="pt",
+    )
+    inputs = inputs.to(model.device)
+    
+    # Generate response with streaming.
+    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+    generation_args = dict(inputs, streamer=streamer, max_new_tokens=max_tokens)
+    thread = Thread(target=model.generate, kwargs=generation_args)
     thread.start()
-
+    
     yield "..."
     buffer = ""
     for new_text in streamer:
@@ -84,16 +140,30 @@ def model_inference(input_dict, history, max_tokens):
         time.sleep(0.01)
         yield buffer
 
+examples = [
+    [{"text": "Where do the severe droughts happen according to this diagram?", "files": ["example_images/examples_weather_events.png"]}],
+    [{"text": "What art era does this artpiece <image> belong to?", "files": ["example_images/rococo.jpg"]}],
+    [{"text": "Describe this image.", "files": ["example_images/mosque.jpg"]}],
+    [{"text": "When was this purchase made and how much did it cost?", "files": ["example_images/fiche.jpg"]}],
+    [{"text": "What is the date in this document?", "files": ["example_images/document.jpg"]}],
+    [{"text": "What is happening in the video?", "files": ["example_images/short.mp4"]}],
+]
+
 demo = gr.ChatInterface(
     fn=model_inference,
-    title="SmolVLM2: The Smollest Video Model Ever 📺", 
-    description="Play with [SmolVLM2-2.2B-Instruct](https://huggingface.co/HuggingFaceTB/SmolVLM2-2.2B-Instruct) in this demo. To get started, upload an image and text.",
-    textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image", ".mp4"], file_count="multiple"), 
-    stop_btn="Stop Generation", 
+    title="SmolVLM2: The Smollest Video Model Ever 📺",
+    description=(
+        "Play with [SmolVLM2-2.2B-Instruct](https://huggingface.co/HuggingFaceTB/SmolVLM2-2.2B-Instruct) in this demo. "
+        "To get started, upload an image and text or try one of the examples. "
+        "This demo doesn't use history for the chat, so every chat you start is a new conversation."
+    ),
+    examples=examples,
+    textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image", "video"], file_count="multiple"),
+    stop_btn="Stop Generation",
     multimodal=True,
     cache_examples=False,
     additional_inputs=[gr.Slider(minimum=100, maximum=500, step=50, value=200, label="Max Tokens")],
     type="messages"
 )
 
-demo.launch(debug=True, share=True)
+demo.launch(debug=True)