Update app.py

app.py

@@ -21,18 +21,13 @@ def generate_caption(image):
     Generate a poem from an image.
     """
     # Process the image and the prompt
-    inputs = processor(texts=["Write a poem about this picture"],
-                       images=[image],
+    inputs = processor("[INST] <image>\nWrite a poem about this picture [/INST]",
+                       image,
                        return_tensors="pt").to('cuda')
-    # Generate the output
-    with torch.inference_mode():
-        output = model.generate(
-            **inputs,
-            max_new_tokens=128,
-        )
-    prompt_len = inputs["input_ids"].shape[1]
-    decoded_text = processor.batch_decode(output[:, prompt_len:])[0]
-    return decoded_text
+    # autoregressively complete prompt
+    output = model.generate(**inputs, max_new_tokens=200)
+
+    return processor.decode(output[0], skip_special_tokens=True)
 
 
 # Define the Gradio interface