Update app.py

app.py

@@ -1,29 +1,29 @@
 import os
 import mimetypes
-from typing import List, Tuple
+from typing import List, Optional
 
 import cv2
 import gradio as gr
 import numpy as np
 from PIL import Image
 import torch
-from transformers import AutoProcessor, AutoModelForVision2Seq, TextStreamer
+from transformers import AutoProcessor, AutoModelForVision2Seq
 
 # =========================
 # Config
 # =========================
 TITLE = "Qwen2-VL Visual Assistant"
-DESC = "Upload one image or a short video (≤15s). Ask anything in the chat like ChatGPT."
+DESC = "Upload one image or a short video (≤15s). Ask anything like ChatGPT."
 
 MODEL_ID = "Qwen/Qwen2-VL-2B-Instruct"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
 
-MAX_VIDEO_SECONDS = 15
-MAX_FILE_MB = 60
-MAX_SIDE = 768
-MAX_FRAMES = 4
-MAX_NEW_TOKENS = 256
+MAX_VIDEO_SECONDS = 15      # hard limit for video duration we sample from
+MAX_FILE_MB = 60            # guard rail for very large files
+MAX_SIDE = 768              # resize longer side for images/frames
+MAX_FRAMES = 4              # how many frames to sample from video
+MAX_NEW_TOKENS = 256        # generation length
 
 # =========================
 # Load model & processor
@@ -36,17 +36,36 @@ model = AutoModelForVision2Seq.from_pretrained(
     trust_remote_code=True,
 ).to(DEVICE)
 
+
 # =========================
 # Helpers
 # =========================
+def get_path(file_obj) -> Optional[str]:
+    """
+    Gradio <File> may arrive as:
+    - str (tmp path)         -> return as-is
+    - dict {"name": "..."}   -> return ["name"] or ["path"]
+    - SimpleNamespace with .name -> return .name
+    """
+    if file_obj is None:
+        return None
+    if isinstance(file_obj, str):
+        return file_obj
+    if isinstance(file_obj, dict):
+        return file_obj.get("name") or file_obj.get("path")
+    return getattr(file_obj, "name", None)
+
+
 def is_image_file(path: str) -> bool:
     kind, _ = mimetypes.guess_type(path)
     return (kind or "").startswith("image/")
 
+
 def is_video_file(path: str) -> bool:
     kind, _ = mimetypes.guess_type(path)
     return (kind or "").startswith("video/")
 
+
 def resize_max_side(img: Image.Image, max_side: int = MAX_SIDE) -> Image.Image:
     w, h = img.size
     m = max(w, h)
@@ -55,20 +74,27 @@ def resize_max_side(img: Image.Image, max_side: int = MAX_SIDE) -> Image.Image:
     s = max_side / float(m)
     return img.resize((int(w * s), int(h * s)), Image.LANCZOS)
 
+
 def sample_video_frames(video_path: str, max_seconds=MAX_VIDEO_SECONDS, k=MAX_FRAMES) -> List[Image.Image]:
+    """Take K evenly spaced frames within the first `max_seconds`."""
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
         return []
+
     fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
     total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
     duration = total / max(fps, 1)
+
+    # Only consider frames within MAX_VIDEO_SECONDS (or all if shorter)
     limit_frames = int(min(duration, max_seconds) * fps) if total else 0
     nframes = min(limit_frames or total, total) if total > 0 else 0
     if nframes == 0:
         nframes = int(min(duration, max_seconds) * fps) or 1
+
     k = max(1, min(k, nframes))
     idxs = [int(i * nframes / k + nframes / (2 * k)) for i in range(k)]
     idxs = sorted(set(max(0, min(nframes - 1, x)) for x in idxs))
+
     frames: List[Image.Image] = []
     for target in idxs:
         cap.set(cv2.CAP_PROP_POS_FRAMES, target)
@@ -81,7 +107,9 @@ def sample_video_frames(video_path: str, max_seconds=MAX_VIDEO_SECONDS, k=MAX_FR
     cap.release()
     return frames
 
+
 def build_messages(prompt: str) -> list:
+    """Qwen2-VL chat template expects <image> tag in user turn."""
     return [{
         "role": "user",
         "content": [
@@ -90,72 +118,88 @@ def build_messages(prompt: str) -> list:
         ]
     }]
 
+
 # =========================
 # Chat function
 # =========================
-def chat_fn(history, file, prompt):
-    # Append user msg first
-    history = history + [[prompt, None]]
+def chat_fn(history, file_obj, prompt):
+    history = history or []
+    history.append([prompt, ""])           # add user bubble first
+    yield history                          # show immediately
 
-    if file is None:
+    path = get_path(file_obj)
+    if not path or not os.path.exists(path):
         history[-1][1] = "⚠️ Please upload an image or a short video first."
         yield history
         return
 
-    path = file.name
-    images = []
-
-    if is_image_file(path):
-        images = [resize_max_side(Image.open(path).convert("RGB"))]
-    elif is_video_file(path):
-        images = sample_video_frames(path, MAX_VIDEO_SECONDS, MAX_FRAMES)
-    else:
-        history[-1][1] = "⚠️ Unsupported file type."
-        yield history
-        return
-
+    # Size guard
     try:
+        if os.path.getsize(path) > MAX_FILE_MB * 1024 * 1024:
+            history[-1][1] = f"⚠️ File is larger than {MAX_FILE_MB} MB."
+            yield history
+            return
+    except Exception:
+        pass
+
+    # Collect images for the model
+    images: List[Image.Image] = []
+    try:
+        if is_image_file(path):
+            images = [resize_max_side(Image.open(path).convert("RGB"))]
+        elif is_video_file(path):
+            images = sample_video_frames(path, MAX_VIDEO_SECONDS, MAX_FRAMES)
+            if not images:
+                history[-1][1] = "⚠️ Could not read frames from the video."
+                yield history
+                return
+        else:
+            history[-1][1] = "⚠️ Unsupported file type."
+            yield history
+            return
+
+        # Build chat input
         messages = build_messages(prompt)
-        chat_str = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+        chat_str = processor.apply_chat_template(
+            messages, add_generation_prompt=True, tokenize=False
+        )
 
+        # Processor expects a LIST of images (even when it's 1)
         enc = processor(
             text=chat_str,
-            images=images if len(images) > 1 else images[0],
+            images=images,          # ALWAYS a list
             return_tensors="pt"
         ).to(model.device)
 
-        acc = ""
-        streamer = TextStreamer(
-            processor.tokenizer,
-            skip_prompt=True,
-            skip_special_tokens=True,
-            callback=lambda s: None
-        )
-
-        # Run generation with streaming
+        # Inference (one-shot; more robust than streaming for Spaces CPU)
         with torch.no_grad():
-            for out in model.generate(
+            out_ids = model.generate(
                 **enc,
                 max_new_tokens=MAX_NEW_TOKENS,
                 do_sample=False,
-                streamer=streamer,
-                return_dict_in_generate=True,
-                output_scores=False,
-            ):
-                text = processor.batch_decode(
-                    out[:, enc.input_ids.shape[-1]:], skip_special_tokens=True
-                )[0]
-                acc = text.strip()
-                history[-1][1] = acc
-                yield history
+            )
+
+        # Decode only the generated continuation
+        cut = enc["input_ids"].shape[-1]
+        text = processor.batch_decode(
+            out_ids[:, cut:], skip_special_tokens=True
+        )[0].strip()
+
+        history[-1][1] = text if text else "(no output)"
+        yield history
+
     except Exception as e:
+        # Show clean error in the bubble
         history[-1][1] = f"⚠️ Error: {str(e)}"
         yield history
 
+
 # =========================
 # Gradio UI
 # =========================
-CSS = "#chatbot {height: 520px; overflow:auto;}"
+CSS = """
+#chat {height: 560px}
+"""
 
 with gr.Blocks(title=TITLE, css=CSS, theme=gr.themes.Soft()) as demo:
     gr.Markdown(f"# {TITLE}")
@@ -166,30 +210,33 @@ with gr.Blocks(title=TITLE, css=CSS, theme=gr.themes.Soft()) as demo:
             uploader = gr.File(
                 label="Upload image or video",
                 file_count="single",
-                file_types=["image", "video"]
+                file_types=["image", "video"],
             )
             preview = gr.Image(label="Preview", interactive=False)
 
         with gr.Column(scale=7):
-            chatbot = gr.Chatbot(label="Visual Assistant", elem_id="chatbot")
-            prompt = gr.Textbox(
+            chat = gr.Chatbot(label="Visual Assistant", elem_id="chat")
+            prompt_in = gr.Textbox(
                 label="Prompt",
                 placeholder="Ask about the image/video…",
-                lines=2
+                lines=2,
             )
-            submit = gr.Button("Submit", variant="primary")
+            send = gr.Button("Submit", variant="primary")
 
     # Preview on file change
-    def on_file(file):
-        if file is None: return None
-        if is_image_file(file.name):
-            return resize_max_side(Image.open(file.name).convert("RGB"))
-        if is_video_file(file.name):
-            return sample_video_frames(file.name, k=1)[0]
+    def on_file(file_obj):
+        path = get_path(file_obj)
+        if not path or not os.path.exists(path):
+            return None
+        if is_image_file(path):
+            return resize_max_side(Image.open(path).convert("RGB"))
+        if is_video_file(path):
+            frames = sample_video_frames(path, k=1)
+            return frames[0] if frames else None
         return None
 
     uploader.change(on_file, inputs=[uploader], outputs=[preview])
-    submit.click(chat_fn, inputs=[chatbot, uploader, prompt], outputs=[chatbot])
+    send.click(chat_fn, inputs=[chat, uploader, prompt_in], outputs=[chat])
 
 if __name__ == "__main__":
     demo.queue().launch()