Spaces:
Build error
Build error
| import gradio as gr | |
| import torch | |
| from threading import Thread | |
| from transformers import ( | |
| SmolVLMProcessor, | |
| AutoModelForImageTextToText, | |
| TextIteratorStreamer, | |
| ) | |
| # ====================== | |
| # INIT MODÈLE | |
| # ====================== | |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| MODEL_ID = "HuggingFaceTB/SmolVLM2-2.2B-Instruct" | |
| processor = SmolVLMProcessor.from_pretrained(MODEL_ID) | |
| model = AutoModelForImageTextToText.from_pretrained( | |
| MODEL_ID, | |
| torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, | |
| ).to(DEVICE).eval() | |
| # ====================== | |
| # STREAMING INFERENCE | |
| # ====================== | |
| def analyze_stream(text, image, max_tokens): | |
| if image is None and not text.strip(): | |
| return "❌ Veuillez fournir un texte ou une image." | |
| content = [] | |
| if image: | |
| content.append({"type": "image", "path": image}) | |
| if text.strip(): | |
| content.append({"type": "text", "text": text}) | |
| messages = [{"role": "user", "content": content}] | |
| inputs = processor.apply_chat_template( | |
| messages, | |
| add_generation_prompt=True, | |
| tokenize=True, | |
| return_tensors="pt", | |
| ).to(DEVICE) | |
| streamer = TextIteratorStreamer( | |
| processor, | |
| skip_prompt=True, | |
| skip_special_tokens=True, | |
| ) | |
| Thread( | |
| target=model.generate, | |
| kwargs=dict( | |
| **inputs, | |
| streamer=streamer, | |
| max_new_tokens=max_tokens, | |
| do_sample=False, | |
| temperature=0.0, | |
| ), | |
| ).start() | |
| output = "" | |
| for token in streamer: | |
| output += token | |
| yield output | |
| # ====================== | |
| # UI GRADIO | |
| # ====================== | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("## ⚡ SmolVLM2 – Analyse Temps Réel") | |
| with gr.Row(): | |
| with gr.Column(): | |
| txt = gr.Textbox( | |
| label="Question / Description", | |
| lines=3, | |
| ) | |
| img = gr.Image(type="filepath", label="Image") | |
| max_tokens = gr.Slider( | |
| 50, 400, value=200, step=50, label="Max Tokens" | |
| ) | |
| btn = gr.Button("🚀 Analyser", variant="primary") | |
| with gr.Column(): | |
| out = gr.Textbox( | |
| label="Réponse en Temps Réel", | |
| lines=14, | |
| ) | |
| btn.click( | |
| fn=analyze_stream, | |
| inputs=[txt, img, max_tokens], | |
| outputs=out, | |
| ) | |
| demo.launch() | |