Update app.py

app.py

@@ -61,7 +61,16 @@ def get_caption(image):
     with torch.no_grad():
         output = model.generate(**inputs, max_new_tokens=200)
 
-    return processor.decode(output[0], skip_special_tokens=True)
+    decoded = processor.decode(output[0], skip_special_tokens=True)
+
+    # ================= CLEANING STEP =================
+    if "ASSISTANT:" in decoded:
+        decoded = decoded.split("ASSISTANT:")[-1]
+
+    if "USER:" in decoded:
+        decoded = decoded.split("USER:")[-1]
+
+    return decoded.strip()
 
 # =====================
 # MAIN PIPELINE