Spaces:
Running
on
Zero
Running
on
Zero
| import gradio as gr | |
| from transformers import AutoProcessor, AutoModelForImageTextToText, 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 | |
| processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct") | |
| model = AutoModelForImageTextToText.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct", | |
| _attn_implementation="flash_attention_2", | |
| torch_dtype=torch.bfloat16).to("cuda:0") | |
| def model_inference( | |
| input_dict, history, max_tokens | |
| ): | |
| 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).") | |
| print("resulting_messages", resulting_messages) | |
| inputs = processor.apply_chat_template( | |
| resulting_messages, | |
| add_generation_prompt=True, | |
| tokenize=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| ) | |
| 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 "..." | |
| 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": "Where do the severe droughts happen according to this diagram?", "files": ["example_images/examples_weather_events.png"]}], | |
| [{"text": "What art era this artpiece <image> and this artpiece <image> belong to?", "files": ["example_images/rococo.jpg", "example_images/rococo_1.jpg"]}], | |
| [{"text": "Describe this image.", "files": ["example_images/mosque.jpg"]}], | |
| [{"text": "When was this purchase made and how much did it cost?", "files": ["example_images/fiche.jpg"]}], | |
| [{"text": "What is the date in this document?", "files": ["example_images/document.jpg"]}], | |
| [{"text": "What is happening in the video?", "files": ["example_images/short.mp4"]}], | |
| ] | |
| 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, | |
| additional_inputs=[gr.Slider(minimum=100, maximum=500, step=50, value=200, label="Max Tokens")], | |
| type="messages" | |
| ) | |
| demo.launch(debug=True) | |