Upload app.py with huggingface_hub

app.py

@@ -28,22 +28,86 @@ except ImportError:
     spaces = _SpacesShim()
 
 
-# === CPU MODE OVERRIDE ===
-# Monkey-patch torch to prevent any accidental .cuda() calls on CPU-only build
-_original_cuda = torch.Tensor.cuda
-def _safe_cuda(self, *args, **kwargs):
-    if not torch.cuda.is_available():
-        return self  # silently stay on CPU
-    return _original_cuda(self, *args, **kwargs)
-torch.Tensor.cuda = _safe_cuda
-
-# Also override .half() to return float32 on CPU (half is not well-supported on CPU)
-_original_half = torch.Tensor.half
-def _safe_half(self, *args, **kwargs):
-    if self.device.type == "cpu":
+# === CPU MODE OVERRIDE (comprehensive) ===
+import functools
+
+if not torch.cuda.is_available():
+    # 1. Tensor.cuda() -> noop
+    _orig_tensor_cuda = torch.Tensor.cuda
+    def _safe_tensor_cuda(self, *a, **kw):
+        return self
+    torch.Tensor.cuda = _safe_tensor_cuda
+
+    # 2. Tensor.half() -> float() on CPU
+    _orig_half = torch.Tensor.half
+    def _safe_half(self, *a, **kw):
         return self.float()
-    return _original_half(self, *args, **kwargs)
-torch.Tensor.half = _safe_half
+    torch.Tensor.half = _safe_half
+
+    # 3. Module.cuda() -> noop
+    _orig_module_cuda = torch.nn.Module.cuda
+    def _safe_module_cuda(self, *a, **kw):
+        return self
+    torch.nn.Module.cuda = _safe_module_cuda
+
+    # 4. Module.to() -> force cpu
+    _orig_module_to = torch.nn.Module.to
+    def _safe_module_to(self, *args, **kwargs):
+        # Replace any "cuda" device with "cpu"
+        new_args = []
+        for a in args:
+            if isinstance(a, (str,)) and "cuda" in a:
+                new_args.append("cpu")
+            elif isinstance(a, torch.device) and a.type == "cuda":
+                new_args.append(torch.device("cpu"))
+            elif a == torch.float16:
+                new_args.append(torch.float32)
+            else:
+                new_args.append(a)
+        if "device" in kwargs:
+            d = kwargs["device"]
+            if isinstance(d, str) and "cuda" in d:
+                kwargs["device"] = "cpu"
+            elif isinstance(d, torch.device) and d.type == "cuda":
+                kwargs["device"] = torch.device("cpu")
+        if "dtype" in kwargs and kwargs["dtype"] == torch.float16:
+            kwargs["dtype"] = torch.float32
+        return _orig_module_to(self, *new_args, **kwargs)
+    torch.nn.Module.to = _safe_module_to
+
+    # 5. Tensor.to() -> force cpu
+    _orig_tensor_to = torch.Tensor.to
+    def _safe_tensor_to(self, *args, **kwargs):
+        new_args = []
+        for a in args:
+            if isinstance(a, (str,)) and "cuda" in a:
+                new_args.append("cpu")
+            elif isinstance(a, torch.device) and a.type == "cuda":
+                new_args.append(torch.device("cpu"))
+            elif a == torch.float16:
+                new_args.append(torch.float32)
+            else:
+                new_args.append(a)
+        if "device" in kwargs:
+            d = kwargs["device"]
+            if isinstance(d, str) and "cuda" in d:
+                kwargs["device"] = "cpu"
+            elif isinstance(d, torch.device) and d.type == "cuda":
+                kwargs["device"] = torch.device("cpu")
+        if "dtype" in kwargs and kwargs["dtype"] == torch.float16:
+            kwargs["dtype"] = torch.float32
+        return _orig_tensor_to(self, *new_args, **kwargs)
+    torch.Tensor.to = _safe_tensor_to
+
+    # 6. torch.load -> force map_location=cpu
+    _orig_load = torch.load
+    @functools.wraps(_orig_load)
+    def _safe_load(*args, **kwargs):
+        kwargs["map_location"] = "cpu"
+        return _orig_load(*args, **kwargs)
+    torch.load = _safe_load
+
+    print("[CPU OVERRIDE] All CUDA calls redirected to CPU", flush=True)
 # === END CPU MODE OVERRIDE ===
 
 from diffusers import AutoencoderKL, DDIMScheduler