Update gradio_app.py

gradio_app.py

@@ -208,10 +208,41 @@ with open("imgs/background.png", "rb") as f:
 
 @spaces.GPU
 def inference(id_image, hair_image):
-    # Require GPU (HairMapper currently uses CUDA explicitly)
-    if not torch.cuda.is_available():
-        raise RuntimeError("This demo requires a GPU Space. Please enable a GPU in this Space.")
+    # ZeroGPU: 强制使用 'cuda' 设备（ZeroGPU 下 torch.cuda.is_available 可能为 False）。
+    device = torch.device("cuda")
 
+    # 先禁用 StyleGAN2 自定义 CUDA 算子（导入 HairMapper 前），避免触发 JIT 编译。
+    # 1) 禁用模块级开关
+    try:
+        from torch_utils.ops import bias_act as _bias_act2
+        _bias_act2.USING_CUDA_TO_SPEED_UP = False
+    except Exception:
+        pass
+    for _mod_name in ("upfirdn2d", "filtered_lrelu"):
+        try:
+            _m = __import__(f"torch_utils.ops.{_mod_name}", fromlist=["*"])
+            if hasattr(_m, 'USING_CUDA_TO_SPEED_UP'):
+                setattr(_m, 'USING_CUDA_TO_SPEED_UP', False)
+        except Exception:
+            pass
+    try:
+        from HairMapper.styleGAN2_ada_model.stylegan2_ada.torch_utils.ops import bias_act as _bias_act
+        _bias_act.USING_CUDA_TO_SPEED_UP = False
+    except Exception:
+        pass
+
+    # 2) 强制 bias_act 走 ref 实现（即便上层传 impl='cuda' 也改为 'ref'）。
+    try:
+        import types
+        from torch_utils.ops import bias_act as _ba_mod
+        _orig_bias_act = _ba_mod.bias_act
+        def _bias_act_ref(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None, impl='cuda'):
+            return _orig_bias_act(x, b, dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp, impl='ref')
+        _ba_mod.bias_act = types.FunctionType(_bias_act_ref.__code__, globals(), name='bias_act')
+    except Exception:
+        pass
+
+    # 再导入依赖（此时已关闭自定义算子与强制 ref 实现）
     (
         log_validation,
         UNet3DConditionModel,
@@ -225,26 +256,6 @@ def inference(id_image, hair_image):
         bald_head,
     ) = _import_inference_bits()
 
-    # Disable StyleGAN2 custom CUDA ops to avoid JIT compiling (needs ninja/NVCC).
-    # ZeroGPU 下建议走纯 PyTorch 引用实现，避免扩展编译失败。
-    try:
-        from HairMapper.styleGAN2_ada_model.stylegan2_ada.torch_utils.ops import bias_act as _bias_act
-        _bias_act.USING_CUDA_TO_SPEED_UP = False
-        try:
-            from HairMapper.styleGAN2_ada_model.stylegan2_ada.torch_utils.ops import upfirdn2d as _upfirdn2d
-            if hasattr(_upfirdn2d, 'USING_CUDA_TO_SPEED_UP'):
-                _upfirdn2d.USING_CUDA_TO_SPEED_UP = False
-        except Exception:
-            pass
-        try:
-            from HairMapper.styleGAN2_ada_model.stylegan2_ada.torch_utils.ops import filtered_lrelu as _fl
-            if hasattr(_fl, 'USING_CUDA_TO_SPEED_UP'):
-                _fl.USING_CUDA_TO_SPEED_UP = False
-        except Exception:
-            pass
-    except Exception:
-        pass
-
     os.makedirs("gradio_inputs", exist_ok=True)
     os.makedirs("gradio_outputs", exist_ok=True)
 
@@ -291,7 +302,7 @@ def inference(id_image, hair_image):
 
     args = Args()
 
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = torch.device("cuda")
 
     logging.basicConfig(
         format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",