Didier/Vision_Language_SmolVLM2

1"""2File: vlm.py3Description: Vision language model utility functions.4 5Heavily inspired (i.e. copied) from6    https://huggingface.co/spaces/HuggingFaceTB/SmolVLM2/blob/main/app.py7 8Author: Didier Guillevic9Date: 2025-04-0210"""11 12from transformers import AutoProcessor, AutoModelForImageTextToText13from transformers import TextIteratorStreamer14from threading import Thread15import re16import time17import torch18import spaces19import subprocess20subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)21 22from io import BytesIO23 24#25# Load the model: HuggingFaceTB/SmolVLM2-2.2B-Instruct26#27 28model_id = "HuggingFaceTB/SmolVLM2-2.2B-Instruct"29device = 'cuda' if torch.cuda.is_available() else 'cpu'30processor = AutoProcessor.from_pretrained(model_id)31model = AutoModelForImageTextToText.from_pretrained(32    model_id, 33    _attn_implementation="flash_attention_2",34    torch_dtype=torch.bfloat1635).to(device)36 37#38# Build messages39#40def build_messages(input_dict: dict, history: list[tuple]):41    """Build messages given message & history from a **multimodal** chat interface.42    Args:43        input_dict: dictionary with keys: 'text', 'files'44        history: list of tuples with (message, response)45    46    Returns:47        list of messages (to be sent to the model)48    """49    text = input_dict["text"]50    images = []51    user_content = []52    media_queue = []53    if history == []:54        text = input_dict["text"].strip() 55        56        for file in input_dict.get("files", []):57            if file.endswith((".png", ".jpg", ".jpeg", ".gif", ".bmp")):58                media_queue.append({"type": "image", "path": file})59            elif file.endswith((".mp4", ".mov", ".avi", ".mkv", ".flv")):60                media_queue.append({"type": "video", "path": file})61 62        if "<image>" in text or "<video>" in text:63            parts = re.split(r'(<image>|<video>)', text)  64            for part in parts:65                if part == "<image>" and media_queue:66                    user_content.append(media_queue.pop(0)) 67                elif part == "<video>" and media_queue:68                    user_content.append(media_queue.pop(0))  69                elif part.strip():  70                    user_content.append({"type": "text", "text": part.strip()})71        else:72            user_content.append({"type": "text", "text": text})73            74            for media in media_queue:75                user_content.append(media)76 77        resulting_messages = [{"role": "user", "content": user_content}]78 79    elif len(history) > 0:80        resulting_messages = []81        user_content = []82        media_queue = []83        for hist in history:84            if hist["role"] == "user" and isinstance(hist["content"], tuple): 85                file_name = hist["content"][0]86            if file_name.endswith((".png", ".jpg", ".jpeg")):87                media_queue.append({"type": "image", "path": file_name})88            elif file_name.endswith(".mp4"):89                media_queue.append({"type": "video", "path": file_name})90 91 92        for hist in history:93            if hist["role"] == "user" and isinstance(hist["content"], str): 94                text = hist["content"]95                parts = re.split(r'(<image>|<video>)', text)  96                97                for part in parts:98                    if part == "<image>" and media_queue:99                        user_content.append(media_queue.pop(0)) 100                    elif part == "<video>" and media_queue:101                        user_content.append(media_queue.pop(0))  102                    elif part.strip(): 103                        user_content.append({"type": "text", "text": part.strip()})104            105            elif hist["role"] == "assistant":106                resulting_messages.append({107                    "role": "user",108                    "content": user_content109                })110                resulting_messages.append({111                    "role": "assistant",112                    "content": [{"type": "text", "text": hist["content"]}]113                })114                user_content = [] 115 116 117    if text == "" and not images:118        gr.Error("Please input a query and optionally image(s).")119 120    if text == "" and images:121        gr.Error("Please input a text query along the images(s).")122    123    return resulting_messages124 125#126# Streaming response127#128@spaces.GPU129@torch.inference_mode()130def stream_response(131        messages: list[dict],132        max_new_tokens: int=1_024,133        temperature: float=0.15134    ):135    """Stream the model's response to the chat interface.136    137    Args:138        messages: list of messages to send to the model139    """140    # Generate model's response141    inputs = processor.apply_chat_template(142        messages,143        add_generation_prompt=True,144        tokenize=True,145        return_dict=True,146        return_tensors="pt",147    ).to(model.device, dtype=torch.bfloat16)148    149    # Generate150    streamer = TextIteratorStreamer(151        processor, skip_prompt=True, skip_special_tokens=True)152    generation_args = dict(153        inputs,154        streamer=streamer,155        max_new_tokens=max_new_tokens,156        temperature=temperature,157        top_p=0.9,158        do_sample=True159    )160 161    thread = Thread(target=model.generate, kwargs=generation_args)162    thread.start()163 164    partial_message = ""165    for new_text in streamer:166        partial_message += new_text167        yield partial_message168