import gradio as gr from transformers import AutoProcessor, AutoModelForVision2Seq, TextIteratorStreamer from transformers.models.smolvlm.video_processing_smolvlm import load_smolvlm_video from transformers.image_utils import load_image 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 from transformers.image_utils import load_image processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct") model = AutoModelForVision2Seq.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct", _attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16, device_map="auto") #@spaces.GPU def model_inference( input_dict, history ): text = input_dict["text"] # first turn input_dict {'text': 'What', 'files': ['/tmp/gradio/0350274350a64a5737e1a5732f014aee2f28bb7344bbad5105c0d0b7e7334375/cats_2.mp4', '/tmp/gradio/2dd39f382fcf5444a1a2ac57ed6f9acafa775dd855248cf273034e8ce18aeff4/IMG_2201.JPG']} # first turn history [] print("input_dict", input_dict) print("history", history) print("model.device", model.device) images = [] # first conv turn if history == []: text = input_dict["text"] resulting_messages = [{"role": "user", "content": [{"type": "text"}, {"type": "text", "text": text}]}] for file in input_dict["files"]: if file.endswith(".mp4"): resulting_messages[0]["content"].append({"type": "video"}) frames, timestamps, duration_sec = load_smolvlm_video( file, sampling_fps=1, max_frames=64 ) print("frames", frames) images.append(frames) elif file.endswith(".jpg") or file.endswith(".jpeg") or file.endswith(".png"): resulting_messages[0]["content"].append({"type": "image"}) images.append(load_image(file)) print("images", images) # second turn input_dict {'text': 'what', 'files': ['/tmp/gradio/7bafdcc4722c4b9902a4936439b3bb694927abd72106a946d773a15cc1c630d7/IMG_2198.JPG']} # second turn history [[('/tmp/gradio/7bafdcc4722c4b9902a4936439b3bb694927abd72106a946d773a15cc1c630d7/IMG_2198.JPG',), None], # [('/tmp/gradio/5b105e97e4876912b4e763902144540bd3ab00d9fd4016491337ee4f4c36f320/football.mp4',), None], ['what', None]] # later conv turn elif len(history) > 0: for hist in history: if isinstance(hist[0], tuple): if hist[0][0].endswith(".mp4"): resulting_messages.append({"role": "user", "content": [{"type": "video"}, {"type": "text", "text": hist[0][0]}]}) frames, timestamps, duration_sec = load_smolvlm_video( file, sampling_fps=1, max_frames=64 ) images.append(frames) else: resulting_messages.append({"role": "user", "content": [{"type": "image"}, {"type": "text", "text": hist[0][0]}]}) images.append(load_image(hist[0][0])) elif isinstance(hist[0], str): resulting_messages.append({"role": "user", "content": [{"type": "text"}, {"type": "text", "text": hist[0]}]}) if isinstance(hist[1], str): resulting_messages.append({"role": "user", "content": [{"type": "text"}, {"type": "text", "text": hist[0]}]}) 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 image(s).") print("resulting_messages", resulting_messages) prompt = processor.apply_chat_template(resulting_messages, add_generation_prompt=True) inputs = processor(text=prompt, images=[images], padding=True, return_tensors="pt") inputs = inputs.to(model.device) generation_args = { "input_ids": inputs.input_ids, "pixel_values": inputs.pixel_values, "attention_mask": inputs.attention_mask, "num_return_sequences": 1, "no_repeat_ngram_size": 2, "max_new_tokens": 500, "min_new_tokens": 10, } # Generate streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True) generation_args = dict(inputs, streamer=streamer, max_new_tokens=500) 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 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"]}], ] 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.", examples=examples, textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image", ".mp4"], file_count="multiple"), stop_btn="Stop Generation", multimodal=True, cache_examples=False, type="messages" ) demo.launch(debug=True)