Spaces:
Build error
Build error
| import gradio as gr | |
| from transformers import AutoProcessor, AutoModelForVision2Seq, 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-500M-Instruct") | |
| model = AutoModelForVision2Seq.from_pretrained("HuggingFaceTB/SmolVLM2-500M-Instruct", | |
| _attn_implementation="flash_attention_2", | |
| torch_dtype=torch.bfloat16).to("cuda:0") | |
| #@spaces.GPU | |
| def model_inference( | |
| input_dict, history, max_tokens | |
| ): | |
| text = input_dict["text"] | |
| images = [] | |
| # first conv turn | |
| if history == []: | |
| text = input_dict["text"] | |
| resulting_messages = [{"role": "user", "content": [{"type": "text", "text": text}]}] | |
| for file in input_dict["files"]: | |
| if file.endswith(".mp4"): | |
| resulting_messages[0]["content"].append({"type": "video", "path": file}) | |
| elif file.endswith(".jpg") or file.endswith(".jpeg") or file.endswith(".png"): | |
| resulting_messages[0]["content"].append({"type": "image", "path": file}) | |
| elif len(history) > 0: | |
| resulting_messages = [] | |
| for entry in history: | |
| if entry["role"] == "user": | |
| user_content = [] | |
| if isinstance(entry["content"], tuple): | |
| file_name = entry["content"][0] | |
| if file_name.endswith((".png", ".jpg", ".jpeg", ".gif", ".bmp")): | |
| user_content.append({"type": "image", "path": file_name}) | |
| elif file_name.endswith((".mp4", ".mov", ".avi", ".mkv", ".flv")): | |
| user_content.append({"type": "video", "path": file_name}) | |
| elif isinstance(entry["content"], str): | |
| user_content.insert(0, {"type": "text", "text": entry["content"]}) | |
| elif entry["role"] == "assistant": | |
| resulting_messages.append({ | |
| "role": "user", | |
| "content": user_content | |
| }) | |
| resulting_messages.append({ | |
| "role": "assistant", | |
| "content": [{"type": "text", "text": entry["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).") | |
| 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": "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, | |
| additional_inputs=[gr.Slider(minimum=100, maximum=500, step=50, value=200, label="Max Tokens")], | |
| type="messages" | |
| ) | |
| demo.launch(debug=True) | |