Create vlm.py

vlm.py

@@ -0,0 +1,161 @@
+"""
+File: vlm.py
+Description: Vision language model utility functions.
+
+Heavily inspired (i.e. copied) from
+    https://huggingface.co/spaces/HuggingFaceTB/SmolVLM2/blob/main/app.py
+
+Author: Didier Guillevic
+Date: 2025-04-02
+"""
+
+from transformers import AutoProcessor, AutoModelForImageTextToText
+from transformers import TextIteratorStreamer
+from threading import Thread
+import re
+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
+
+#
+# Load the model: HuggingFaceTB/SmolVLM2-2.2B-Instruct
+#
+
+model_id = "HuggingFaceTB/SmolVLM2-2.2B-Instruct"
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+processor = AutoProcessor.from_pretrained(model_id)
+model = AutoModelForImageTextToText.from_pretrained(
+    model_id, 
+    _attn_implementation="flash_attention_2",
+    torch_dtype=torch.bfloat16
+).to(device)
+
+#
+# Build messages
+#
+def build_messages(input_dict: dict, history: list[tuple]):
+    """Build messages given message & history from a **multimodal** chat interface.
+    Args:
+        input_dict: dictionary with keys: 'text', 'files'
+        history: list of tuples with (message, response)
+    
+    Returns:
+        list of messages (to be sent to the model)
+    """
+    text = input_dict["text"]
+    images = []
+    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).")
+    
+    return resulting_messages
+
+#
+# Streaming response
+#
+@spaces.GPU
+@torch.inference_mode()
+def stream_response(messages: list[dict]):
+    """Stream the model's response to the chat interface.
+    
+    Args:
+        messages: list of messages to send to the model
+    """
+    # Generate model's response
+    inputs = processor.apply_chat_template(
+        resulting_messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+    ).to(model.device, dtype=torch.bfloat16)
+    
+    # Generate
+    streamer = TextIteratorStreamer(
+        processor, skip_prompt=True, skip_special_tokens=True)
+    generation_args = dict(
+        inputs,
+        streamer=streamer,
+        max_new_tokens=2_048,
+        do_sample=True
+    )
+
+    thread = Thread(target=model.generate, kwargs=generation_args)
+    thread.start()
+
+    partial_message = ""
+    for new_text in streamer:
+        partial_message += new_text
+        yield partial_message