Update app.py

app.py

@@ -1,10 +1,4 @@
 import gradio as gr
-from transformers import (
-    Qwen2VLForConditionalGeneration,
-    AutoProcessor,
-    TextIteratorStreamer,
-    AutoModelForImageTextToText,
-)
 from transformers.image_utils import load_image
 from threading import Thread
 import time
@@ -13,6 +7,14 @@ import spaces
 from PIL import Image
 import requests
 from io import BytesIO
+import cv2
+import numpy as np
+from transformers import (
+    Qwen2VLForConditionalGeneration,
+    AutoProcessor,
+    TextIteratorStreamer,
+    AutoModelForImageTextToText,
+)
 
 # Helper function to return a progress bar HTML snippet.
 def progress_bar_html(label: str) -> str:
@@ -20,7 +22,7 @@ def progress_bar_html(label: str) -> str:
 <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 style="width: 100%; height: 100%; background-color: #FF69B4; animation: loading 1.5s linear infinite;"></div>
     </div>
 </div>
 <style>
@@ -31,7 +33,29 @@ def progress_bar_html(label: str) -> str:
 </style>
     '''
 
-QV_MODEL_ID = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct" # or use  #prithivMLmods/Qwen2-VL-OCR2-2B-Instruct
+# Helper function to downsample a video into 10 evenly spaced frames.
+def downsample_video(video_path):
+    vidcap = cv2.VideoCapture(video_path)
+    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = vidcap.get(cv2.CAP_PROP_FPS)
+    frames = []
+    # Sample 10 evenly spaced 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, 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
+
+# Model and processor setups
+
+# Setup for Qwen2VL OCR branch (default).
+QV_MODEL_ID = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct"  # or use "prithivMLmods/Qwen2-VL-OCR2-2B-Instruct"
 qwen_processor = AutoProcessor.from_pretrained(QV_MODEL_ID, trust_remote_code=True)
 qwen_model = Qwen2VLForConditionalGeneration.from_pretrained(
     QV_MODEL_ID,
@@ -39,25 +63,77 @@ qwen_model = Qwen2VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to("cuda").eval()
 
+# Setup for Aya-Vision branch.
 AYA_MODEL_ID = "CohereForAI/aya-vision-8b"
 aya_processor = AutoProcessor.from_pretrained(AYA_MODEL_ID)
 aya_model = AutoModelForImageTextToText.from_pretrained(
     AYA_MODEL_ID, device_map="auto", torch_dtype=torch.float16
 )
 
+# ---------------------------
+# Main Inference Function
+# ---------------------------
 @spaces.GPU
 def model_inference(input_dict, history):
     text = input_dict["text"].strip()
     files = input_dict.get("files", [])
     
+    # Branch for video inference with Aya-Vision using @video-infer.
+    if text.lower().startswith("@video-infer"):
+        prompt = text[len("@video-infer"):].strip()
+        if not files:
+            yield "Error: Please provide a video for the @video-infer feature."
+            return
+        video_path = files[0]
+        frames = downsample_video(video_path)
+        if not frames:
+            yield "Error: Could not extract frames from the video."
+            return
+        # Build messages: start with the prompt then add each frame with its timestamp.
+        content_list = []
+        content_list.append({"type": "text", "text": prompt})
+        for frame, timestamp in frames:
+            content_list.append({"type": "text", "text": f"Frame {timestamp}:"})
+            content_list.append({"type": "image", "image": frame})
+        messages = [{
+            "role": "user",
+            "content": content_list,
+        }]
+        inputs = aya_processor.apply_chat_template(
+            messages,
+            padding=True,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt"
+        ).to(aya_model.device)
+        streamer = TextIteratorStreamer(aya_processor, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = dict(
+            inputs, 
+            streamer=streamer, 
+            max_new_tokens=1024, 
+            do_sample=True, 
+            temperature=0.3
+        )
+        thread = Thread(target=aya_model.generate, kwargs=generation_kwargs)
+        thread.start()
+        buffer = ""
+        yield progress_bar_html("Processing video with Aya-Vision-8b")
+        for new_text in streamer:
+            buffer += new_text
+            buffer = buffer.replace("<|im_end|>", "")
+            time.sleep(0.01)
+            yield buffer
+        return
+
+    # Branch for single image inference with Aya-Vision using @aya-vision.
     if text.lower().startswith("@aya-vision"):
-        # Remove the command prefix and trim the prompt.
         text_prompt = text[len("@aya-vision"):].strip()
         if not files:
             yield "Error: Please provide an image for the @aya-vision feature."
             return
         else:
-            # For simplicity, use the first provided image.
+            # Use the first provided image.
             image = load_image(files[0])
             yield progress_bar_html("Processing with Aya-Vision-8b")
             messages = [{
@@ -75,7 +151,6 @@ def model_inference(input_dict, history):
                 return_dict=True,
                 return_tensors="pt"
             ).to(aya_model.device)
-            # Set up a streamer for Aya-Vision output
             streamer = TextIteratorStreamer(aya_processor, skip_prompt=True, skip_special_tokens=True)
             generation_kwargs = dict(
                 inputs, 
@@ -94,7 +169,7 @@ def model_inference(input_dict, history):
                 yield buffer
             return
 
-    # Load images if provided.
+    # Default branch: Use Qwen2VL OCR for text (with optional images).
     if len(files) > 1:
         images = [load_image(image) for image in files]
     elif len(files) == 1:
@@ -102,7 +177,6 @@ def model_inference(input_dict, history):
     else:
         images = []
     
-    # Validate input: require both text and (optionally) image(s).
     if text == "" and not images:
         yield "Error: Please input a query and optionally image(s)."
         return
@@ -110,7 +184,6 @@ def model_inference(input_dict, history):
         yield "Error: Please input a text query along with the image(s)."
         return
 
-    # Prepare messages for the Qwen2-VL model.
     messages = [{
         "role": "user",
         "content": [
@@ -129,11 +202,9 @@ def model_inference(input_dict, history):
         padding=True,
     ).to("cuda")
     
-    # Set up a streamer for real-time output.
     streamer = TextIteratorStreamer(qwen_processor, skip_prompt=True, skip_special_tokens=True)
     generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
     
-    # Start generation in a separate thread.
     thread = Thread(target=qwen_model.generate, kwargs=generation_kwargs)
     thread.start()
     
@@ -145,7 +216,11 @@ def model_inference(input_dict, history):
         time.sleep(0.01)
         yield buffer
 
+
+# Gradio Interface Setup
+
 examples = [
+    [{"text": "@video-infer Summarize the video content", "files": ["examples/videoplayback.mp4"]}],
     [{"text": "@aya-vision Summarize the letter", "files": ["examples/1.png"]}],
     [{"text": "@aya-vision Extract JSON from the image", "files": ["example_images/document.jpg"]}],
     [{"text": "Extract as JSON table from the table", "files": ["examples/4.jpg"]}],
@@ -160,13 +235,13 @@ examples = [
 
 demo = gr.ChatInterface(
     fn=model_inference,
-    description="# **Multimodal OCR `@aya-vision 'prompt..'`**",
+    description="# **Multimodal OCR and Video Inference with Aya-Vision (@aya-vision for image, @video-infer for video) and Qwen2VL OCR (default)**",
     examples=examples,
     textbox=gr.MultimodalTextbox(
         label="Query Input", 
-        file_types=["image"], 
+        file_types=["image", "video"], 
         file_count="multiple", 
-        placeholder="By default, it runs Qwen2VL OCR, Tag @aya-vision for Aya Vision 8B"
+        placeholder="Tag @aya-vision for Aya-Vision image infer, @video-infer for Aya-Vision video infer, default runs Qwen2VL OCR"
     ),
     stop_btn="Stop Generation",
     multimodal=True,