Create app.py

app.py

@@ -0,0 +1,35 @@
+import io
+from PIL import Image
+import torch
+from transformers import AutoProcessor, AutoModelForVisionEncoderDecoder
+
+# Load the model and processor
+model_name = "colt12/maxcushion"
+processor = AutoProcessor.from_pretrained(model_name)
+model = AutoModelForVisionEncoderDecoder.from_pretrained(model_name)
+
+def predict(image_bytes):
+    # Open the image using PIL
+    image = Image.open(io.BytesIO(image_bytes))
+    
+    # Preprocess the image
+    pixel_values = processor(images=image, return_tensors="pt").pixel_values
+
+    # Generate the caption
+    generated_ids = model.generate(pixel_values, max_length=50)
+    generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    
+    return generated_caption
+
+def run(raw_image):
+    # Input validation
+    if not raw_image:
+        raise ValueError("No image provided")
+    
+    try:
+        # Process the image and generate the caption
+        result = predict(raw_image)
+        return {"caption": result}
+    except Exception as e:
+        # Error handling
+        return {"error": str(e)}