Update app.py

app.py

@@ -1,144 +1,101 @@
 import gradio as gr
-import spaces
-from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor, TextIteratorStreamer
-from qwen_vl_utils import process_vision_info
-import torch
-from PIL import Image
-import subprocess
-import numpy as np
-import os
+from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration, TextIteratorStreamer
+from transformers.image_utils import load_image
 from threading import Thread
-import uuid
-import io
+import time
+import torch
+import spaces
 
-# Model and Processor Loading (Done once at startup)
+# Load the Qwen2.5-VL-3B-Instruct model and processor
 MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct"
+processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
 model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     MODEL_ID,
     trust_remote_code=True,
     torch_dtype=torch.float16
 ).to("cuda").eval()
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
-
-DESCRIPTION = "# **Qwen2.5-VL-3B-Instruct**"
-
-image_extensions = Image.registered_extensions()
-video_extensions = ("avi", "mp4", "mov", "mkv", "flv", "wmv", "mjpeg", "wav", "gif", "webm", "m4v", "3gp")
-
-
-def identify_and_save_blob(blob_path):
-    """Identifies if the blob is an image or video and saves it accordingly."""
-    try:
-        with open(blob_path, 'rb') as file:
-            blob_content = file.read()
-            
-            # Try to identify if it's an image
-            try:
-                Image.open(io.BytesIO(blob_content)).verify()  # Check if it's a valid image
-                extension = ".png"  # Default to PNG for saving
-                media_type = "image"
-            except (IOError, SyntaxError):
-                # If it's not a valid image, assume it's a video
-                extension = ".mp4"  # Default to MP4 for saving
-                media_type = "video"
-            
-            # Create a unique filename
-            filename = f"temp_{uuid.uuid4()}_media{extension}"
-            with open(filename, "wb") as f:
-                f.write(blob_content)
-                
-            return filename, media_type
-            
-    except FileNotFoundError:
-        raise ValueError(f"The file {blob_path} was not found.")
-    except Exception as e:
-        raise ValueError(f"An error occurred while processing the file: {e}")
-
 
 @spaces.GPU
-def qwen_inference(media_input, text_input=None):
-    if isinstance(media_input, str):  # If it's a filepath
-        media_path = media_input
-        if media_path.endswith(tuple([i for i, f in image_extensions.items()])):
-            media_type = "image"
-        elif media_path.endswith(video_extensions): 
-            media_type = "video"
-        else:
-            try:
-                media_path, media_type = identify_and_save_blob(media_input)
-                print(media_path, media_type)
-            except Exception as e:
-                print(e)
-                raise ValueError(
-                    "Unsupported media type. Please upload an image or video."
-                )
-        
-
-    print(media_path)
-
+def model_inference(input_dict, history):
+    text = input_dict["text"]
+    files = input_dict["files"]
+
+    # Load images if provided
+    if len(files) > 1:
+        images = [load_image(image) for image in files]
+    elif len(files) == 1:
+        images = [load_image(files[0])]
+    else:
+        images = []
+
+    # Validate input
+    if text == "" and not images:
+        gr.Error("Please input a query and optionally image(s).")
+        return
+    if text == "" and images:
+        gr.Error("Please input a text query along with the image(s).")
+        return
+
+    # Prepare messages for the model
     messages = [
         {
             "role": "user",
             "content": [
-                {
-                    "type": media_type,
-                    media_type: media_path,
-                    **({"fps": 8.0} if media_type == "video" else {}),
-                },
-                {"type": "text", "text": text_input},
+                *[{"type": "image", "image": image} for image in images],
+                {"type": "text", "text": text},
             ],
         }
     ]
 
-    text = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-    image_inputs, video_inputs = process_vision_info(messages)
+    # Apply chat template and process inputs
+    prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = processor(
-        text=[text],
-        images=image_inputs,
-        videos=video_inputs,
-        padding=True,
+        text=[prompt],
+        images=images if images else None,
         return_tensors="pt",
+        padding=True,
     ).to("cuda")
 
-    streamer = TextIteratorStreamer(
-        processor, skip_prompt=True, **{"skip_special_tokens": True}
-    )
+    # Set up streamer for real-time output
+    streamer = TextIteratorStreamer(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=model.generate, kwargs=generation_kwargs)
     thread.start()
 
+    # Stream the output
     buffer = ""
+    yield "..."
     for new_text in streamer:
         buffer += new_text
+        time.sleep(0.01)
         yield buffer
 
-css = """
-  #output {
-    height: 500px; 
-    overflow: auto; 
-    border: 1px solid #ccc; 
-  }
-"""
-
-with gr.Blocks(css=css) as demo:
-    gr.Markdown(DESCRIPTION)
-
-    with gr.Tab(label="Image/Video Input"):
-        with gr.Row():
-            with gr.Column():
-                input_media = gr.File(
-                    label="Upload Image or Video", type="filepath" 
-                )
-                text_input = gr.Textbox(label="Question")
-                submit_btn = gr.Button(value="Submit")
-            with gr.Column():
-                output_text = gr.Textbox(label="Output Text")
-
-        submit_btn.click(
-            qwen_inference, [input_media, text_input], [output_text]
-        )
 
+# Example inputs
+examples = [
+    [{"text": "Can you describe this image?", "files": ["example_images/newyork.jpg"]}],
+    [{"text": "Can you describe this image?", "files": ["example_images/dogs.jpg"]}],
+    [{"text": "Where do the severe droughts happen according to this diagram?", "files": ["example_images/examples_weather_events.png"]}],
+    [{"text": "What art era do these artpieces belong to?", "files": ["example_images/rococo.jpg", "example_images/rococo_1.jpg"]}],
+    [{"text": "Describe this image.", "files": ["example_images/campeones.jpg"]}],
+    [{"text": "What does this say?", "files": ["example_images/math.jpg"]}],
+    [{"text": "What is the date in this document?", "files": ["example_images/document.jpg"]}],
+    [{"text": "What is this UI about?", "files": ["example_images/s2w_example.png"]}],
+]
+
+# Gradio interface
+demo = gr.ChatInterface(
+    fn=model_inference,
+    title="# **Qwen2.5-VL-3B-Instruc**",
+    description="Interact with [Qwen2.5-VL-3B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct) in this demo. Upload an image and text, or try one of the examples. Each chat starts a new conversation.",
+    examples=examples,
+    textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image"], file_count="multiple"),
+    stop_btn="Stop Generation",
+    multimodal=True,
+    cache_examples=False,
+)
+
+# Launch the demo
 demo.launch(debug=True)