Fix ZeroGPU: add @spaces.GPU, import spaces first, clean up requirements

app.py

@@ -1,5 +1,6 @@
+import spaces  # MUST be the very first import for ZeroGPU to work
+
 import io
-import os
 import numpy as np
 from PIL import Image, UnidentifiedImageError
 import torch
@@ -9,85 +10,59 @@ from fastapi.responses import Response
 from transformers import DepthProImageProcessor, DepthProForDepthEstimation
 
 
-# 1. Initialize FastAPI & Load Apple Depth Pro
-fastapi_app = FastAPI(title="True Metric ToF Server")
+# =====================================================================
+# 1. Load Apple Depth Pro at module level.
+#    ZeroGPU safely emulates CUDA during startup so .to("cuda") works here.
+# =====================================================================
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Updated to avoid the transformers deprecation crash
+dtype  = torch.float16 if torch.cuda.is_available() else torch.float32
 
 print("Loading Apple Depth Pro components...")
 image_processor = DepthProImageProcessor.from_pretrained("apple/DepthPro-hf")
-model = DepthProForDepthEstimation.from_pretrained("apple/DepthPro-hf").to(device)
-
+model = DepthProForDepthEstimation.from_pretrained(
+    "apple/DepthPro-hf",
+    torch_dtype=dtype,
+).to(device)
 model.eval()
 print(f"Model successfully loaded on target device: {device}")
 
 
-@fastapi_app.get("/")
-def home():
-    return {"status": "running"}
-
-
-@fastapi_app.post("/get_metric_depth")
-def get_metric_depth(file: UploadFile = File(...)):
-    if not (file.content_type or "").startswith("image/"):
-        raise HTTPException(status_code=400, detail="Uploaded file must be an image.")
-
-    try:
-        image_bytes = file.file.read()
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    except (UnidentifiedImageError, OSError):
-        raise HTTPException(status_code=400, detail="Could not decode image.")
-
-    try:
-        inputs = image_processor(images=image, return_tensors="pt").to(device)
-        with torch.no_grad():
-            outputs = model(**inputs)
-
-        post_processed = image_processor.post_process_depth_estimation(
-            outputs, target_sizes=[(image.height, image.width)],
-        )
-
-        metric_depth_map = (
-            post_processed["predicted_depth"]
-            .detach().cpu().numpy().astype(np.float32)
-        )
-
-        buffer = io.BytesIO()
-        np.save(buffer, metric_depth_map)
-        return Response(content=buffer.getvalue(), media_type="application/octet-stream")
-
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")
+# =====================================================================
+# 2. Core inference helper (shared by Gradio UI and the REST endpoint)
+# =====================================================================
+def _run_depth_inference(image: Image.Image):
+    """Returns post-processed depth output list. Caller handles indexing."""
+    inputs = image_processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    return image_processor.post_process_depth_estimation(
+        outputs,
+        target_sizes=[(image.height, image.width)],
+    )
 
 
 # =====================================================================
-# 2. GRADIO INTERACTIVE INTERFACE (Does the 3D Projection in the Cloud)
+# 3. Gradio function - decorated with @spaces.GPU so ZeroGPU allocates
+#    a real GPU for this call.
 # =====================================================================
-
+@spaces.GPU(duration=120)
 def create_3d_point_cloud(input_pil_image):
     """Processes the image, calculates 3D vertices, and writes a .ply file."""
     if input_pil_image is None:
         return None
 
-    # Step A: Run inference locally inside the space
+    # Step A: Depth inference
     image = input_pil_image.convert("RGB")
-    inputs = image_processor(images=image, return_tensors="pt").to(device)
-    
-    with torch.no_grad():
-        outputs = model(**inputs)
-    
-    post_processed = image_processor.post_process_depth_estimation(
-    outputs, target_sizes=[(image.height, image.width)],
-)
-# Your list index fix: grabbing the first dictionary element out of the post-processed array
-    depth_matrix = post_processed[0]["predicted_depth"].detach().cpu().numpy().astype(np.float32)
+    post_processed = _run_depth_inference(image)
+    depth_matrix = (
+        post_processed[0]["predicted_depth"]
+        .detach().cpu().numpy().astype(np.float32)
+    )
 
-
-    # Step B: Pinhole camera math projection (u, v, depth) -> (X, Y, Z)
+    # Step B: Pinhole camera projection (u, v, depth) -> (X, Y, Z)
     height, width = depth_matrix.shape
     color_np = np.array(image) / 255.0
-    
+
     focal_length = max(width, height)
     cx, cy = width / 2.0, height / 2.0
     u, v = np.meshgrid(np.arange(width), np.arange(height))
@@ -95,18 +70,18 @@ def create_3d_point_cloud(input_pil_image):
     z = depth_matrix.flatten()
     x = ((u.flatten() - cx) * z) / focal_length
     y = ((v.flatten() - cy) * z) / focal_length
-    
-    # Clean up infinite backgrounds (keep objects between 10cm and 30 meters)
+
+    # Keep objects between 10 cm and 30 m; discard sky / infinite background
     valid_mask = (z > 0.1) & (z < 30.0)
     x, y, z = x[valid_mask], y[valid_mask], z[valid_mask]
-    
+
     rgb = color_np.reshape(-1, 3)[valid_mask]
     rgb_ints = (rgb * 255).astype(np.uint8)
 
-    # Step C: Write raw points directly to a standard .ply file format
+    # Step C: Write ASCII PLY file
     ply_filename = "output_scene.ply"
     num_points = len(z)
-    
+
     with open(ply_filename, "w") as f:
         f.write("ply\n")
         f.write("format ascii 1.0\n")
@@ -118,29 +93,78 @@ def create_3d_point_cloud(input_pil_image):
         f.write("property uchar green\n")
         f.write("property uchar blue\n")
         f.write("end_header\n")
-        
         for i in range(num_points):
-            # Invert Y axis slightly so the 3D model loads right-side up
-            f.write(f"{x[i]:.4f} {-y[i]:.4f} {z[i]:.4f} {rgb_ints[i][0]} {rgb_ints[i][1]} {rgb_ints[i][2]}\n")
+            # Invert Y axis so the model loads right-side up
+            f.write(
+                f"{x[i]:.4f} {-y[i]:.4f} {z[i]:.4f} "
+                f"{rgb_ints[i][0]} {rgb_ints[i][1]} {rgb_ints[i][2]}\n"
+            )
 
     return ply_filename
 
 
-# Define web dashboard interface layout
+# =====================================================================
+# 4. FastAPI backend for the REST /get_metric_depth endpoint
+# =====================================================================
+fastapi_app = FastAPI(title="True Metric ToF Server")
+
+
+@fastapi_app.get("/status")
+def status():
+    return {"status": "running", "device": str(device)}
+
+
+@fastapi_app.post("/get_metric_depth")
+async def get_metric_depth(file: UploadFile = File(...)):
+    if not (file.content_type or "").startswith("image/"):
+        raise HTTPException(status_code=400, detail="Uploaded file must be an image.")
+
+    try:
+        image_bytes = await file.read()
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    except (UnidentifiedImageError, OSError):
+        raise HTTPException(status_code=400, detail="Could not decode image.")
+
+    try:
+        post_processed = _run_depth_inference(image)
+        metric_depth_map = (
+            post_processed[0]["predicted_depth"]
+            .detach().cpu().numpy().astype(np.float32)
+        )
+        buffer = io.BytesIO()
+        np.save(buffer, metric_depth_map)
+        return Response(content=buffer.getvalue(), media_type="application/octet-stream")
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")
+
+
+# =====================================================================
+# 5. Gradio UI
+# =====================================================================
 with gr.Blocks(title="Software-ToF 3D Cloud Visualizer") as demo:
     gr.Markdown("# 📷 Software-ToF Camera Simulation")
-    gr.Markdown("Upload any image to calculate absolute metric depths and generate an interactive 3D Point Cloud visualization.")
-    
+    gr.Markdown(
+        "Upload any image to calculate absolute metric depths and generate "
+        "an interactive 3D Point Cloud visualization."
+    )
+
     with gr.Row():
         with gr.Column():
-            input_img = gr.Image(type="pil", label="Input Scene Photo")
+            input_img  = gr.Image(type="pil", label="Input Scene Photo")
             submit_btn = gr.Button("Generate ToF 3D Structure", variant="primary")
-        
         with gr.Column():
-            # Model3D renders an interactive canvas in the web browser
-            output_3d = gr.Model3D(label="Interactive 3D Point Cloud Canvas", display_mode="point_cloud")
+            output_3d = gr.Model3D(
+                label="Interactive 3D Point Cloud Canvas",
+                display_mode="point_cloud",
+            )
 
     submit_btn.click(fn=create_3d_point_cloud, inputs=[input_img], outputs=[output_3d])
 
-# Mount your unified architecture onto Hugging Face port 7860
+
+# =====================================================================
+# 6. Mount Gradio onto FastAPI.
+#    Hugging Face (sdk: gradio) auto-discovers this `app` variable and
+#    serves it via uvicorn on port 7860 - no manual uvicorn.run() needed.
+# =====================================================================
 app = gr.mount_gradio_app(fastapi_app, demo, path="/")

requirements.txt

@@ -1,10 +1,10 @@
-fastapi
-uvicorn[standard]
-python-multipart
-torch --index-url https://pytorch.org
-torchvision --index-url https://pytorch.org
+torch
+torchvision
 transformers
 accelerate
 numpy
 pillow
 gradio
+fastapi
+uvicorn[standard]
+python-multipart