Update server.py

server.py

@@ -55,7 +55,7 @@ print(f"✓ Image transform defined (size: {CONFIG['image_size']}x{CONFIG['image
 def preprocess_image(image_path: str) -> torch.Tensor:
     """Load and preprocess image."""
     image = Image.open(image_path).convert('RGB')
-    return transform(image).unsqueeze(0)  # Add batch dimension
+    return transform(image)
 
 # ─────────────────────────────────────────────────────────────────────────────
 # ARCHITECTURE 1 — CoAtNet Encoder  (shared by all three models)
@@ -434,7 +434,7 @@ async def sft_inference(file: UploadFile = File(...)):
 
         try:
             # Use file path preprocessing (exact Colab match)
-            tensor = preprocess_image(temp_path)
+            tensor = preprocess_image(temp_path).unsqueeze(0).to(device)
             report = sft_model.generate_reports(tensor)[0]
             print(f"[SFT] Generated: {report}")
             return {"report": report[:81]}
@@ -458,7 +458,7 @@ async def reward_inference(file: UploadFile = File(...)):
 
         try:
             # Use file path preprocessing (exact Colab match)
-            tensor = preprocess_image(temp_path)
+            tensor = preprocess_image(temp_path).unsqueeze(0).to(device)
             # First get the SFT report to score
             sft_report = sft_model.generate_reports(tensor)[0]
             print(f"[REWARD] Scoring SFT report: {sft_report}")
@@ -528,7 +528,7 @@ async def ppo_inference(file: UploadFile = File(...)):
 
         try:
             # Use file path preprocessing (exact Colab match)
-            tensor = preprocess_image(temp_path)
+            tensor = preprocess_image(temp_path).unsqueeze(0).to(device)
             report = ppo_model.generate_reports(tensor)[0]
             print(f"[PPO] Generated: {report}")
             return {"report": report}