CPU fallback: avoid CUDA JIT by disabling fused ops when build unavailable; force CPU

models/psp/stylegan2/op/fused_act.py

@@ -3,15 +3,25 @@ import os
 import torch
 from torch import nn
 from torch.autograd import Function
-from torch.utils.cpp_extension import load
-module_path = os.path.dirname(__file__)
-fused = load(
-    "fused",
-    sources=[
-        os.path.join(module_path, "fused_bias_act.cpp"),
-        os.path.join(module_path, "fused_bias_act_kernel.cu"),
-    ],
-)
+
+_USE_FUSED = False
+fused = None
+try:
+    # Prefer fused CUDA extension when available
+    from torch.utils.cpp_extension import load
+    module_path = os.path.dirname(__file__)
+    fused = load(
+        "fused",
+        sources=[
+            os.path.join(module_path, "fused_bias_act.cpp"),
+            os.path.join(module_path, "fused_bias_act_kernel.cu"),
+        ],
+        verbose=False,
+    )
+    _USE_FUSED = True
+except Exception:
+    # CPU-only or build-restricted environments (e.g., HF Spaces) fallback
+    _USE_FUSED = False
 
 class FusedLeakyReLUFunctionBackward(Function):
     @staticmethod
@@ -22,9 +32,13 @@ class FusedLeakyReLUFunctionBackward(Function):
 
         empty = grad_output.new_empty(0)
 
-        grad_input = fused.fused_bias_act(
-            grad_output, empty, out, 3, 1, negative_slope, scale
-        )
+        if _USE_FUSED:
+            grad_input = fused.fused_bias_act(
+                grad_output, empty, out, 3, 1, negative_slope, scale
+            )
+        else:
+            # Fallback: autograd-friendly composite ops
+            grad_input = torch.nn.functional.leaky_relu(grad_output, negative_slope) * scale
 
         dim = [0]
 
@@ -38,9 +52,12 @@ class FusedLeakyReLUFunctionBackward(Function):
     @staticmethod
     def backward(ctx, gradgrad_input, gradgrad_bias):
         (out,) = ctx.saved_tensors
-        gradgrad_out = fused.fused_bias_act(
-            gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale
-        )
+        if _USE_FUSED:
+            gradgrad_out = fused.fused_bias_act(
+                gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale
+            )
+        else:
+            gradgrad_out = torch.nn.functional.leaky_relu(gradgrad_input, ctx.negative_slope) * ctx.scale
 
         return gradgrad_out, None, None, None
 
@@ -48,8 +65,13 @@ class FusedLeakyReLUFunctionBackward(Function):
 class FusedLeakyReLUFunction(Function):
     @staticmethod
     def forward(ctx, input, bias, negative_slope, scale):
-        empty = input.new_empty(0)
-        out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale)
+        if _USE_FUSED:
+            empty = input.new_empty(0)
+            out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale)
+        else:
+            if bias is not None:
+                input = input + bias.view(1, -1, *([1] * (input.dim() - 2)))
+            out = torch.nn.functional.leaky_relu(input, negative_slope) * scale
         ctx.save_for_backward(out)
         ctx.negative_slope = negative_slope
         ctx.scale = scale

models/psp/stylegan2/op/upfirdn2d.py

@@ -1,17 +1,25 @@
 import os
 
 import torch
+import torch.nn.functional as F
 from torch.autograd import Function
-from torch.utils.cpp_extension import load
 
-module_path = os.path.dirname(__file__)
-upfirdn2d_op = load(
-    "upfirdn2d",
-    sources=[
-        os.path.join(module_path, "upfirdn2d.cpp"),
-        os.path.join(module_path, "upfirdn2d_kernel.cu"),
-    ],
-)
+_USE_EXT = False
+upfirdn2d_op = None
+try:
+    from torch.utils.cpp_extension import load
+    module_path = os.path.dirname(__file__)
+    upfirdn2d_op = load(
+        "upfirdn2d",
+        sources=[
+            os.path.join(module_path, "upfirdn2d.cpp"),
+            os.path.join(module_path, "upfirdn2d_kernel.cu"),
+        ],
+        verbose=False,
+    )
+    _USE_EXT = True
+except Exception:
+    _USE_EXT = False
 
 
 class UpFirDn2dBackward(Function):
@@ -25,18 +33,24 @@ class UpFirDn2dBackward(Function):
 
         grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1)
 
-        grad_input = upfirdn2d_op.upfirdn2d(
-            grad_output,
-            grad_kernel,
-            down_x,
-            down_y,
-            up_x,
-            up_y,
-            g_pad_x0,
-            g_pad_x1,
-            g_pad_y0,
-            g_pad_y1,
-        )
+        if _USE_EXT:
+            grad_input = upfirdn2d_op.upfirdn2d(
+                grad_output,
+                grad_kernel,
+                down_x,
+                down_y,
+                up_x,
+                up_y,
+                g_pad_x0,
+                g_pad_x1,
+                g_pad_y0,
+                g_pad_y1,
+            )
+        else:
+            # Fallback approximation using native ops
+            grad_input = upfirdn2d_native(
+                grad_output, grad_kernel, up_x, up_y, down_x, down_y, g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1
+            )
         grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3])
 
         ctx.save_for_backward(kernel)
@@ -62,18 +76,23 @@ class UpFirDn2dBackward(Function):
 
         gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1)
 
-        gradgrad_out = upfirdn2d_op.upfirdn2d(
-            gradgrad_input,
-            kernel,
-            ctx.up_x,
-            ctx.up_y,
-            ctx.down_x,
-            ctx.down_y,
-            ctx.pad_x0,
-            ctx.pad_x1,
-            ctx.pad_y0,
-            ctx.pad_y1,
-        )
+        if _USE_EXT:
+            gradgrad_out = upfirdn2d_op.upfirdn2d(
+                gradgrad_input,
+                kernel,
+                ctx.up_x,
+                ctx.up_y,
+                ctx.down_x,
+                ctx.down_y,
+                ctx.pad_x0,
+                ctx.pad_x1,
+                ctx.pad_y0,
+                ctx.pad_y1,
+            )
+        else:
+            gradgrad_out = upfirdn2d_native(
+                gradgrad_input, kernel, ctx.up_x, ctx.up_y, ctx.down_x, ctx.down_y, ctx.pad_x0, ctx.pad_x1, ctx.pad_y0, ctx.pad_y1
+            )
         # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3])
         gradgrad_out = gradgrad_out.view(
             ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1]
@@ -112,9 +131,12 @@ class UpFirDn2d(Function):
 
         ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1)
 
-        out = upfirdn2d_op.upfirdn2d(
-            input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
-        )
+        if _USE_EXT:
+            out = upfirdn2d_op.upfirdn2d(
+                input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
+            )
+        else:
+            out = upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1)
         # out = out.view(major, out_h, out_w, minor)
         out = out.view(-1, channel, out_h, out_w)