Update caption.py

caption.py

@@ -6,21 +6,28 @@ from PIL import Image
 model = AutoModelForCausalLM.from_pretrained("Chesscorner/git-chess-v3")
 processor = AutoProcessor.from_pretrained("Chesscorner/git-chess-v3")
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# Set up device and move model to it
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-max_length = 60
-num_beams = 2
-gen_kwargs = {'max_length': max_length, 'num_beams': num_beams}
+# Enable mixed precision if on GPU
+use_fp16 = device.type == "cuda"
+if use_fp16:
+    model.half()
 
+# Set generation parameters
+gen_kwargs = {'max_length': 10, 'num_beams': 2}  # Adjust num_beams if needed
+
+
+# Prediction function
 def predict_step(image):
-    pixel_values = processor(
-        images=[image], return_tensors='pt').pixel_values
-    pixel_values = pixel_values.to(device)
+    # Preprocess the image
+    pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device)
 
-    output_ids = model.generate(pixel_values, **gen_kwargs)
+    # Generate predictions with no_grad for efficiency
+    with torch.no_grad():
+        output_ids = model.generate(pixel_values=pixel_values, **gen_kwargs)
 
+    # Decode predictions
     preds = processor.batch_decode(output_ids, skip_special_tokens=True)
-    preds = [pred.strip() for pred in preds]
-    return preds[0]
-    
+    return preds[0].strip()