Update app.py

app.py

@@ -1,42 +1,38 @@
 from flask import Flask, request, jsonify
-from transformers import BlipImageProcessor, AutoTokenizer, BlipForConditionalGeneration
+from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
 import torch
+import io
 
 app = Flask(__name__)
 
-# Load components properly
-image_processor = BlipImageProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-tokenizer = AutoTokenizer.from_pretrained("Salesforce/blip-image-captioning-base")
+# Load model and processor
+processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
 model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 model.eval()
 
-# Set device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
 @app.route("/caption", methods=["POST"])
 def caption_image():
-    try:
-        if 'image' not in request.files:
-            return jsonify({"error": "No image file provided"}), 400
+    if 'image' not in request.files:
+        return jsonify({"error": "No image file provided"}), 400
 
+    try:
         image_file = request.files['image']
-        image = Image.open(image_file.stream).convert("RGB")
+        image = Image.open(io.BytesIO(image_file.read())).convert("RGB")
 
-        # Process image and caption
-        pixel_values = image_processor(images=image, return_tensors="pt").pixel_values.to(device)
-        generated_ids = model.generate(pixel_values)
-        caption = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+        processed_image = processor(images=image, return_tensors="pt").pixel_values.to(device)
 
-        return jsonify({"caption": caption})
+        with torch.no_grad():
+            output = model.generate(processed_image)
 
+        caption = processor.decode(output[0], skip_special_tokens=True)
+        return jsonify({"caption": caption})
+    
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
-@app.route("/")
-def home():
-    return "BLIP Captioning API is running!"
-
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)