Update app.py

app.py

@@ -1,7 +1,93 @@
 import gradio as gr
+import torch
+from PIL import Image
+import requests
+from transformers import AutoProcessor, Idefics3ForConditionalGeneration, TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList
 
-def greet(name):
-    return "Hello " + name + "!!"
+base_model_id = "Andres77872/SmolVLM-500M-anime-caption-v0.1"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+processor = AutoProcessor.from_pretrained(base_model_id)
+model = Idefics3ForConditionalGeneration.from_pretrained(
+    base_model_id,
+    device_map="auto",
+    torch_dtype=torch.bfloat16
+)
+
+class StopOnTokens(StoppingCriteria):
+    def __init__(self, tokenizer, stop_sequence):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.stop_sequence = stop_sequence
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        new_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
+        max_keep = len(self.stop_sequence) + 10
+        if len(new_text) > max_keep:
+            new_text = new_text[-max_keep:]
+        return self.stop_sequence in new_text
+
+def prepare_inputs(image: Image.Image):
+    question = "describe the image"
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image"},
+                {"type": "text", "text": question}
+            ]
+        }
+    ]
+    max_image_size = processor.image_processor.max_image_size["longest_edge"]
+    size = processor.image_processor.size.copy()
+    if "longest_edge" in size and size["longest_edge"] > max_image_size:
+        size["longest_edge"] = max_image_size
+    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = processor(text=[prompt], images=[[image]], return_tensors='pt', padding=True, size=size)
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    return inputs
+
+def caption_anime_image_stream(image):
+    if image is None:
+        yield "Please upload an image."
+        return
+    inputs = prepare_inputs(image)
+    stop_sequence = "</QUERY>"
+    streamer = TextIteratorStreamer(
+        processor.tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True,
+    )
+    custom_stopping_criteria = StoppingCriteriaList([
+        StopOnTokens(processor.tokenizer, stop_sequence)
+    ])
+    with torch.no_grad():
+        generation_kwargs = dict(
+            **inputs,
+            streamer=streamer,
+            do_sample=False,
+            max_new_tokens=512,
+            pad_token_id=processor.tokenizer.pad_token_id,
+            stopping_criteria=custom_stopping_criteria,
+        )
+        import threading
+        generation_thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
+        generation_thread.start()
+        caption = ""
+        for new_text in streamer:
+            caption += new_text
+            yield caption.strip()
+        generation_thread.join()
+
+demo = gr.Interface(
+    caption_anime_image_stream,
+    inputs=gr.Image(type="pil", label="Anime Image"),
+    outputs=gr.Textbox(lines=8, label="Caption"),
+    title="SmolVLM-500M-Anime-Caption Demo",
+    description="Upload an anime-style image to generate a caption.",
+    # Enable live streaming:
+    allow_flagging="auto",
+    examples=None,
+)
+
+demo.queue()
+demo.launch()