test_on_gpu

app.py

@@ -67,6 +67,9 @@ try:
     model_file = hf_hub_download("KYAGABA/combined-multimodal-model", "pytorch_model.bin", token=HF_TOKEN)
     state_dict = torch.load(model_file, map_location=device)
     combined_model.load_state_dict(state_dict)
+    
+    # Move model to device
+    combined_model = combined_model.to(device)
     combined_model.eval()
 except Exception as e:
     raise SystemExit(f"Error loading models: {e}")
@@ -81,48 +84,58 @@ class_names = ["acute", "normal", "chronic", "lacunar"]
 
 @app.post("/predict/")
 async def predict(files: list[UploadFile]):
-    print(f"Received {len(files)} files")
-    n_frames = 16
-    images = []
-
-    for file in files:
-        ext = file.filename.split('.')[-1].lower()
-        try:
-            if ext in ['dcm', 'ima']:
-                dicom_img = dicom_to_png(await file.read())
-                images.append(dicom_img.convert("RGB"))
-            elif ext in ['png', 'jpeg', 'jpg']:
-                img = Image.open(io.BytesIO(await file.read())).convert("RGB")
-                images.append(img)
-            else:
-                raise HTTPException(status_code=400, detail="Unsupported file type.")
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error processing file {file.filename}: {e}")
-
-    if not images:
-        return JSONResponse(content={"error": "No valid images provided."}, status_code=400)
-
-    if len(images) >= n_frames:
-        images_sampled = [images[i] for i in np.linspace(0, len(images) - 1, n_frames, dtype=int)]
-    else:
-        images_sampled = images + [images[-1]] * (n_frames - len(images))
-
-    image_tensors = [image_transform(img) for img in images_sampled]
-    images_tensor = torch.stack(image_tensors).permute(1, 0, 2, 3).unsqueeze(0).to(device)
-
-    with torch.no_grad():
-        class_outputs, generated_report, _ = combined_model(images_tensor)
-        predicted_class = torch.argmax(class_outputs, dim=1).item()
-        predicted_class_name = class_names[predicted_class]
-
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
+    try:
+        print(f"Received {len(files)} files")
+        n_frames = 16
+        images = []
+
+        for file in files:
+            ext = file.filename.split('.')[-1].lower()
+            try:
+                if ext in ['dcm', 'ima']:
+                    dicom_img = dicom_to_png(await file.read())
+                    images.append(dicom_img.convert("RGB"))
+                elif ext in ['png', 'jpeg', 'jpg']:
+                    img = Image.open(io.BytesIO(await file.read())).convert("RGB")
+                    images.append(img)
+                else:
+                    raise HTTPException(status_code=400, detail="Unsupported file type.")
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=f"Error processing file {file.filename}: {e}")
+
+        if not images:
+            return JSONResponse(content={"error": "No valid images provided."}, status_code=400)
+
+        if len(images) >= n_frames:
+            images_sampled = [images[i] for i in np.linspace(0, len(images) - 1, n_frames, dtype=int)]
+        else:
+            images_sampled = images + [images[-1]] * (n_frames - len(images))
+
+        image_tensors = [image_transform(img) for img in images_sampled]
+        images_tensor = torch.stack(image_tensors).permute(1, 0, 2, 3).unsqueeze(0)
         
-    return {
-        "predicted_class": predicted_class_name,
-        "generated_report": generated_report[0] if generated_report else "No report generated."
-    }
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))
+        # Ensure tensor is on the same device as model
+        images_tensor = images_tensor.to(device)
+
+        with torch.no_grad():
+            class_outputs, generated_report, _ = combined_model(images_tensor)
+            predicted_class = torch.argmax(class_outputs, dim=1).item()
+            predicted_class_name = class_names[predicted_class]
+
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            
+        return {
+            "predicted_class": predicted_class_name,
+            "generated_report": generated_report[0] if generated_report else "No report generated."
+        }
+        
+    except Exception as e:
+        import traceback
+        error_details = traceback.format_exc()
+        print(f"Error during prediction: {str(e)}\n{error_details}")
+        return JSONResponse(
+            status_code=500,
+            content={"error": f"Error processing request: {str(e)}"}
+        )