Update app.py

app.py

@@ -3,33 +3,47 @@ from __future__ import absolute_import, division, print_function
 import os, sys
 import cv2
 import yaml
-import torch
 import numpy as np
-import torch.nn as nn
 import gradio as gr
 from huggingface_hub import hf_hub_download
 
+# ✅ 必须最早 import spaces（在 torch / 任何 CUDA 初始化之前）
+try:
+    import spaces  # noqa: F401
+except Exception:
+    spaces = None  # 不影响本地跑
+
 # ========== 让 Space 能 import 你的工程 ==========
-PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))  # app.py 在仓库根目录
+PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(PROJECT_ROOT)
 
-from networks.models import make
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
+from networks.models import make  # noqa: E402
 
-# ====== HF 权重仓库配置（你已经上传成功）======
+# ====== HF 权重仓库配置 ======
 WEIGHTS_REPO = "Insta360-Research/DAP-weights"
 WEIGHTS_FILE = "model.pth"
+CONFIG_PATH = "config/infer.yaml"
+
+# 先定义全局占位
+model = None
+device = "cpu"
+
 
-# ========== 可视化 ==========
 def colorize_depth(depth, colormap=cv2.COLORMAP_JET):
     depth = depth.astype(np.float32)
     depth_norm = (depth - depth.min()) / (depth.max() - depth.min() + 1e-6)
     depth_u8 = (depth_norm * 255).astype(np.uint8)
     return cv2.applyColorMap(depth_u8, colormap)  # BGR
 
-# ========== 加载模型（只加载一次） ==========
+
 def load_model(config_path: str):
+    # ✅ torch 放到这里 import，避免在 spaces import 之前触发 CUDA
+    import torch
+    import torch.nn as nn
+
+    global device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
     with open(config_path, "r") as f:
         config = yaml.load(f, Loader=yaml.FullLoader)
 
@@ -39,46 +53,42 @@ def load_model(config_path: str):
 
     state = torch.load(model_path, map_location=device)
 
-    model = make(config["model"])
+    m = make(config["model"])
     if any(k.startswith("module") for k in state.keys()):
-        model = nn.DataParallel(model)
-
-    model = model.to(device)
+        m = nn.DataParallel(m)
 
-    model_state = model.state_dict()
-    model.load_state_dict({k: v for k, v in state.items() if k in model_state}, strict=False)
-    model.eval()
+    m = m.to(device)
+    m_state = m.state_dict()
+    m.load_state_dict({k: v for k, v in state.items() if k in m_state}, strict=False)
+    m.eval()
     print("✅ Model loaded.")
-    return model
+    return m
 
-# 这里改成你仓库里的 config 路径
-CONFIG_PATH = "config/infer.yaml"
+
+# ✅ 启动时加载一次模型
 model = load_model(CONFIG_PATH)
 
-# ========== 单张图推理 ==========
-@torch.no_grad()
+
 def predict(img_rgb: np.ndarray):
-    """
-    img_rgb: H x W x 3 (RGB), uint8
-    return: depth_color_rgb, depth_gray
-    """
     if img_rgb is None:
         return None, None
 
+    import torch  # 这里用到 torch，再 import 一次没关系
+
     img = img_rgb.astype(np.float32) / 255.0
     tensor = torch.from_numpy(img.transpose(2, 0, 1)).unsqueeze(0).to(device)
 
-    outputs = model(tensor)
+    with torch.no_grad():
+        outputs = model(tensor)
 
-    if isinstance(outputs, dict) and "pred_depth" in outputs:
-        # 你原来的 mask 逻辑
-        if "pred_mask" in outputs:
-            outputs["pred_mask"] = 1 - outputs["pred_mask"]
-            outputs["pred_mask"] = (outputs["pred_mask"] > 0.5)
-            outputs["pred_depth"][~outputs["pred_mask"]] = 1
-        pred = outputs["pred_depth"][0].detach().cpu().squeeze().numpy()
-    else:
-        pred = outputs[0].detach().cpu().squeeze().numpy()
+        if isinstance(outputs, dict) and "pred_depth" in outputs:
+            if "pred_mask" in outputs:
+                outputs["pred_mask"] = 1 - outputs["pred_mask"]
+                outputs["pred_mask"] = (outputs["pred_mask"] > 0.5)
+                outputs["pred_depth"][~outputs["pred_mask"]] = 1
+            pred = outputs["pred_depth"][0].detach().cpu().squeeze().numpy()
+        else:
+            pred = outputs[0].detach().cpu().squeeze().numpy()
 
     pred_clip = np.clip(pred, 0.001, 1.0)
     depth_gray = (pred_clip * 255).astype(np.uint8)
@@ -88,6 +98,7 @@ def predict(img_rgb: np.ndarray):
 
     return depth_color_rgb, depth_gray
 
+
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="numpy", label="Input Image"),
@@ -96,7 +107,7 @@ demo = gr.Interface(
         gr.Image(type="numpy", label="Depth (Gray)"),
     ],
     title="DAP Depth Prediction Demo",
-    description="Upload an image and get depth prediction."
+    description="Upload an image and get depth prediction.",
 )
 
 demo.launch(