Update app.py

app.py

@@ -6,123 +6,118 @@ import time
 import torch
 import spaces
 import subprocess
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-
 from io import BytesIO
 
-processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct")
-model = AutoModelForImageTextToText.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct", 
-                                               _attn_implementation="flash_attention_2",
-                                               torch_dtype=torch.bfloat16).to("cuda:0")
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM2-256M-Video-Instruct")
+model = AutoModelForImageTextToText.from_pretrained(
+    "HuggingFaceTB/SmolVLM2-256M-Video-Instruct",
+    _attn_implementation="flash_attention_2",
+    torch_dtype=torch.bfloat16
+).to("cuda:0")
 
 
 @spaces.GPU
 def model_inference(
     input_dict, history, max_tokens
-): 
-    text = input_dict["text"]
-    images = []
+):
+
+    text = input_dict["text"].strip()
     user_content = []
     media_queue = []
-    if history == []:
-        text = input_dict["text"].strip() 
-        
-        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})
-
-        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()})
-        else:
-            user_content.append({"type": "text", "text": text})
-            
-            for media in media_queue:
-                user_content.append(media)
-
-        resulting_messages = [{"role": "user", "content": user_content}]
-
-    elif len(history) > 0:
-        resulting_messages = []
-        user_content = []
-        media_queue = []
-        for hist in history:
-            if hist["role"] == "user" and isinstance(hist["content"], tuple): 
-                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})
-
-
-        for hist in history:
-            if hist["role"] == "user" and isinstance(hist["content"], str): 
-                text = hist["content"]
-                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()})
-            
-            elif hist["role"] == "assistant":
-                resulting_messages.append({
-                    "role": "user",
-                    "content": user_content
-                })
-                resulting_messages.append({
-                    "role": "assistant",
-                    "content": [{"type": "text", "text": hist["content"]}]
-                })
-                user_content = [] 
-
-
-    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 the images(s).")
+
+    for file_path in input_dict.get("files", []):
+        if file_path.lower().endswith((".png", ".jpg", ".jpeg", ".gif", ".bmp")):
+            media_queue.append({"type": "image", "path": file_path})
+        elif file_path.lower().endswith((".mp4", ".mov", ".avi", ".mkv", ".flv")):
+            media_queue.append({"type": "video", "path": file_path})
+
+
+    if not text and not media_queue:
+        gr.Warning("Please input a query and optionally image(s)/video(s).")
+        return
+
+    if not text and media_queue:
+        gr.Warning("Please input a text query along with the image(s)/video(s).")
+        return
+
+
+    if "<image>" in text or "<video>" in text:
+        parts = re.split(r'(<image>|<video>)', text)
+        temp_media_queue = list(media_queue)
+        for part in parts:
+            if part == "<image>" and temp_media_queue:
+                media_item = temp_media_queue.pop(0)
+                if media_item["type"] == "image":
+                    user_content.append(media_item)
+                else:
+                    gr.Warning(f"Placeholder <image> found, but next media is a video: {media_item['path']}. Skipping placeholder.")
+                    user_content.append({"type": "text", "text": part})
+                    temp_media_queue.insert(0, media_item)
+            elif part == "<video>" and temp_media_queue:
+                media_item = temp_media_queue.pop(0)
+                if media_item["type"] == "video":
+                     user_content.append(media_item)
+                else:
+                    gr.Warning(f"Placeholder <video> found, but next media is an image: {media_item['path']}. Skipping placeholder.")
+                    user_content.append({"type": "text", "text": part})
+                    temp_media_queue.insert(0, media_item)
+            elif part.strip():
+                user_content.append({"type": "text", "text": part.strip()})
+            elif part in ["<image>", "<video>"] and not temp_media_queue:
+                 gr.Warning(f"Placeholder {part} found, but no more media items available.")
+                 user_content.append({"type": "text", "text": part})
+
+        user_content.extend(temp_media_queue)
+
+    else:
+        if text:
+             user_content.append({"type": "text", "text": text})
+        user_content.extend(media_queue)
+
+
+    resulting_messages = [{"role": "user", "content": user_content}]
+
+
     print("resulting_messages", resulting_messages)
-    inputs = processor.apply_chat_template(
-    resulting_messages,
-    add_generation_prompt=True,
-    tokenize=True,
-    return_dict=True,
-    return_tensors="pt",
-    )
+
+    try:
+        inputs = processor.apply_chat_template(
+            resulting_messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+        )
+    except Exception as e:
+        gr.Error(f"Error during input processing: {e}")
+        print(f"Processor Error: {e}")
+        print("Problematic message structure:", resulting_messages)
+        return
+
 
     inputs = inputs.to(model.device)
-    
 
-    # Generate
-    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+
+    if "pixel_values" in inputs:
+        inputs["pixel_values"] = inputs["pixel_values"].to(model.dtype)
+
+
+    streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
     generation_args = dict(inputs, streamer=streamer, max_new_tokens=max_tokens)
     generated_text = ""
 
     thread = Thread(target=model.generate, kwargs=generation_args)
     thread.start()
 
-    yield "..."
     buffer = ""
-    
-      
+
     for new_text in streamer:
-    
-      buffer += new_text
-      generated_text_without_prompt = buffer#[len(ext_buffer):]
-      time.sleep(0.01)
-      yield buffer
+        buffer += new_text
+        yield buffer
+
+    thread.join()
 
 
 examples=[
@@ -133,16 +128,15 @@ examples=[
               [{"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 or try one of the examples. This demo doesn't use history for the chat, so every chat you start is a new conversation.",
+demo = gr.ChatInterface(fn=model_inference, title="SmolVLM2: The Smollest Video Model Ever 📺",
+                description="Play with [SmolVLM2-256M-Video-Instruct](https://huggingface.co/HuggingFaceTB/SmolVLM2-256M-Video-Instruct) in this demo. To get started, upload an image/video 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", ".mp4"], 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)
-        
+
+
+demo.launch(debug=True)