Update app.py

app.py

@@ -9,9 +9,8 @@ base_model_id = "Andres77872/SmolVLM-500M-anime-caption-v0.1"
 processor = AutoProcessor.from_pretrained(base_model_id)
 model = Idefics3ForConditionalGeneration.from_pretrained(
     base_model_id,
-    device_map="auto",
     torch_dtype=torch.bfloat16
-)
+).to("cuda:0")
 
 class StopOnTokens(StoppingCriteria):
     def __init__(self, tokenizer, stop_sequence):
@@ -26,7 +25,12 @@ class StopOnTokens(StoppingCriteria):
             new_text = new_text[-max_keep:]
         return self.stop_sequence in new_text
 
-def prepare_inputs(image: Image.Image):
+@spaces.GPU
+def caption_anime_image_stream(image):
+    if image is None:
+        yield "Please upload an image."
+        return
+    
     question = "describe the image"
     messages = [
         {
@@ -44,13 +48,7 @@ def prepare_inputs(image: Image.Image):
     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,