Update app.py

app.py

@@ -6,8 +6,10 @@ import time
 import torch
 import spaces
 import subprocess
+
+# Ensure pyav is installed
+subprocess.run('pip install pyav', shell=True, check=True)
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-subprocess.run("pip install av", shell=True)  # Install pyav for video processing
 
 from io import BytesIO
 
@@ -17,50 +19,100 @@ model = AutoModelForImageTextToText.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B
                                                torch_dtype=torch.bfloat16).to("cuda:0")
 
 @spaces.GPU
-def model_inference(input_dict, history, max_tokens): 
+def model_inference(
+    input_dict, history, max_tokens
+): 
     text = input_dict["text"]
-    media_queue = []
+    images = []
     user_content = []
-    
-    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})
+    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})
-        user_content.extend(media_queue)
+        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}]
-    
-    if not text and not media_queue:
-        return "Please provide text and/or media input."
-    
+        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).")
+    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)
+    )
+
+    inputs = inputs.to(model.device)
     
+    # Generate
     streamer = TextIteratorStreamer(processor, 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 "Generating response..."
+
+    yield "..."
     buffer = ""
     
     for new_text in streamer:
@@ -68,17 +120,12 @@ def model_inference(input_dict, history, max_tokens):
         time.sleep(0.01)
         yield buffer
 
+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. This demo doesn't use history for the chat, so every chat you start is a new conversation.",
+                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 = gr.ChatInterface(
-    fn=model_inference,
-    title="SmolVLM2: The Smallest Video Model Ever 📺", 
-    description="Play with SmolVLM2-2.2B-Instruct. Upload an image or video and ask a question.",
-    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(share=True, debug=True)
+demo.launch(debug=True, share=True)