fix: load models on CPU at startup for Zero GPU compatibility

app.py

@@ -1,4 +1,3 @@
-import io
 import spaces
 import numpy as np
 import torch
@@ -11,58 +10,37 @@ from transformers import (
     AutoModelForDepthEstimation,
 )
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
-
 ADE20K_FLOOR_IDS = {3, 28}  # 3=floor, 28=rug/carpet
 
-seg_processor = None
-seg_model = None
-depth_processor = None
-depth_model = None
-
-
-def load_models():
-    global seg_processor, seg_model, depth_processor, depth_model
-
-    if seg_model is None:
-        seg_processor = OneFormerProcessor.from_pretrained(
-            "shi-labs/oneformer_ade20k_swin_large"
-        )
-        seg_model = OneFormerForUniversalSegmentation.from_pretrained(
-            "shi-labs/oneformer_ade20k_swin_large",
-            torch_dtype=DTYPE,
-        ).to(DEVICE)
+# Load models on CPU at startup. @spaces.GPU moves them to CUDA automatically.
+seg_processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_ade20k_swin_large")
+seg_model = OneFormerForUniversalSegmentation.from_pretrained(
+    "shi-labs/oneformer_ade20k_swin_large", torch_dtype=torch.float16
+)
 
-    if depth_model is None:
-        depth_processor = AutoImageProcessor.from_pretrained(
-            "depth-anything/Depth-Anything-V2-Large-hf"
-        )
-        depth_model = AutoModelForDepthEstimation.from_pretrained(
-            "depth-anything/Depth-Anything-V2-Large-hf",
-            torch_dtype=DTYPE,
-        ).to(DEVICE)
+depth_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Large-hf")
+depth_model = AutoModelForDepthEstimation.from_pretrained(
+    "depth-anything/Depth-Anything-V2-Large-hf", torch_dtype=torch.float16
+)
 
 
+@spaces.GPU
 @torch.inference_mode()
-def process_image(image: Image.Image):
-    """Takes a room photo, returns floor mask + depth map as images."""
+def predict(image):
     if image is None:
         raise gr.Error("No image provided")
 
-    load_models()
-
     orig_w, orig_h = image.size
     max_size = 1024
     scale = min(1.0, max_size / max(orig_w, orig_h))
     proc_w, proc_h = int(orig_w * scale), int(orig_h * scale)
     image_resized = image.resize((proc_w, proc_h), Image.LANCZOS)
 
-    # --- Segmentation ---
-    seg_inputs = seg_processor(
-        images=image_resized, task_inputs=["semantic"], return_tensors="pt"
-    )
-    seg_inputs = {k: v.to(DEVICE, dtype=DTYPE) if v.dtype == torch.float32 else v.to(DEVICE) for k, v in seg_inputs.items()}
+    device = seg_model.device
+
+    # Segmentation
+    seg_inputs = seg_processor(images=image_resized, task_inputs=["semantic"], return_tensors="pt")
+    seg_inputs = {k: v.to(device) for k, v in seg_inputs.items()}
 
     seg_outputs = seg_model(**seg_inputs)
     seg_result = seg_processor.post_process_semantic_segmentation(
@@ -74,17 +52,15 @@ def process_image(image: Image.Image):
     for class_id in ADE20K_FLOOR_IDS:
         floor_mask[seg_map == class_id] = 255
 
-    # Resize mask to original dimensions
     mask_img = Image.fromarray(floor_mask).resize((orig_w, orig_h), Image.NEAREST)
 
-    # --- Depth estimation ---
+    # Depth estimation
     depth_inputs = depth_processor(images=image_resized, return_tensors="pt")
-    depth_inputs = {k: v.to(DEVICE, dtype=DTYPE) if v.dtype == torch.float32 else v.to(DEVICE) for k, v in depth_inputs.items()}
+    depth_inputs = {k: v.to(device) for k, v in depth_inputs.items()}
 
     depth_outputs = depth_model(**depth_inputs)
     depth_map = depth_outputs.predicted_depth.squeeze().cpu().numpy()
 
-    # Normalize to 0-255
     depth_min, depth_max = depth_map.min(), depth_map.max()
     if depth_max - depth_min > 0:
         depth_norm = ((depth_map - depth_min) / (depth_max - depth_min) * 255).astype(np.uint8)
@@ -96,15 +72,8 @@ def process_image(image: Image.Image):
     return mask_img, depth_img
 
 
-@spaces.GPU
-def predict(image):
-    mask, depth = process_image(image)
-    return mask, depth
-
-
 with gr.Blocks() as demo:
     gr.Markdown("# Tile Visualizer - Segmentation API")
-    gr.Markdown("Upload a room photo to get floor mask + depth map.")
 
     with gr.Row():
         input_image = gr.Image(type="pil", label="Room photo")