Update app.py

app.py

@@ -14,16 +14,19 @@ device = infer_device()
 def process_image(image, prompt):
    # Process the image and prompt using the processor
    inputs = processor(image, text=prompt, return_tensors="pt").to(device, torch.float16)
-  
+
    try:
        # Generate output from the model
        output = model.generate(**inputs, max_new_tokens=10)
-      
+
        # Decode and return the output
        decoded_output = processor.batch_decode(output, skip_special_tokens=True)[0].strip()
-      
+
        # Return the answer (exclude the prompt part from output)
-       return decoded_output[len(prompt):]
+       # The model's output often includes the prompt, so we remove it.
+       if decoded_output.startswith(prompt):
+           return decoded_output[len(prompt):].strip()
+       return decoded_output
    except IndexError as e:
        print(f"IndexError: {e}")
        return "An error occurred during processing."