changing around sg2/op

app.py

@@ -27,7 +27,7 @@ from PIL import Image
 opts_args = ['--no_fine_mapper']
 opts = GradioTestOptions().parse(opts_args)
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-opts.device=device
+opts.device= device if opts.device is None else opts.device
 
 mapper_dict = {
     'afro':'./pretrained_models/styleCLIP_mappers/afro_hairstyle.pt',

models/stylegan2/model.py

@@ -5,12 +5,8 @@ import torch
 from torch import nn
 from torch.nn import functional as F
 
-if torch.cuda.is_available():
-    from models.stylegan2.op.fused_act import FusedLeakyReLU, fused_leaky_relu
-    from op.upfirdn2d import upfirdn2d
-else:
-    from models.stylegan2.op.fused_act_cpu import FusedLeakyReLU, fused_leaky_relu
-    from models.stylegan2.op.upfirdn2d_cpu import upfirdn2d
+from models.stylegan2.op.fused_act import FusedLeakyReLU, fused_leaky_relu
+from models.stylegan2.op.upfirdn2d import upfirdn2d
 
 class PixelNorm(nn.Module):
     def __init__(self):

models/stylegan2/op/conv2d_gradfix.py

@@ -1,227 +0,0 @@
-import contextlib
-import warnings
-
-import torch
-from torch import autograd
-from torch.nn import functional as F
-
-enabled = True
-weight_gradients_disabled = False
-
-
-@contextlib.contextmanager
-def no_weight_gradients():
-    global weight_gradients_disabled
-
-    old = weight_gradients_disabled
-    weight_gradients_disabled = True
-    yield
-    weight_gradients_disabled = old
-
-
-def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
-    if could_use_op(input):
-        return conv2d_gradfix(
-            transpose=False,
-            weight_shape=weight.shape,
-            stride=stride,
-            padding=padding,
-            output_padding=0,
-            dilation=dilation,
-            groups=groups,
-        ).apply(input, weight, bias)
-
-    return F.conv2d(
-        input=input,
-        weight=weight,
-        bias=bias,
-        stride=stride,
-        padding=padding,
-        dilation=dilation,
-        groups=groups,
-    )
-
-
-def conv_transpose2d(
-    input,
-    weight,
-    bias=None,
-    stride=1,
-    padding=0,
-    output_padding=0,
-    groups=1,
-    dilation=1,
-):
-    if could_use_op(input):
-        return conv2d_gradfix(
-            transpose=True,
-            weight_shape=weight.shape,
-            stride=stride,
-            padding=padding,
-            output_padding=output_padding,
-            groups=groups,
-            dilation=dilation,
-        ).apply(input, weight, bias)
-
-    return F.conv_transpose2d(
-        input=input,
-        weight=weight,
-        bias=bias,
-        stride=stride,
-        padding=padding,
-        output_padding=output_padding,
-        dilation=dilation,
-        groups=groups,
-    )
-
-
-def could_use_op(input):
-    if (not enabled) or (not torch.backends.cudnn.enabled):
-        return False
-
-    if input.device.type != "cuda":
-        return False
-
-    if any(torch.__version__.startswith(x) for x in ["1.7.", "1.8."]):
-        return True
-
-    warnings.warn(
-        f"conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d()."
-    )
-
-    return False
-
-
-def ensure_tuple(xs, ndim):
-    xs = tuple(xs) if isinstance(xs, (tuple, list)) else (xs,) * ndim
-
-    return xs
-
-
-conv2d_gradfix_cache = dict()
-
-
-def conv2d_gradfix(
-    transpose, weight_shape, stride, padding, output_padding, dilation, groups
-):
-    ndim = 2
-    weight_shape = tuple(weight_shape)
-    stride = ensure_tuple(stride, ndim)
-    padding = ensure_tuple(padding, ndim)
-    output_padding = ensure_tuple(output_padding, ndim)
-    dilation = ensure_tuple(dilation, ndim)
-
-    key = (transpose, weight_shape, stride, padding, output_padding, dilation, groups)
-    if key in conv2d_gradfix_cache:
-        return conv2d_gradfix_cache[key]
-
-    common_kwargs = dict(
-        stride=stride, padding=padding, dilation=dilation, groups=groups
-    )
-
-    def calc_output_padding(input_shape, output_shape):
-        if transpose:
-            return [0, 0]
-
-        return [
-            input_shape[i + 2]
-            - (output_shape[i + 2] - 1) * stride[i]
-            - (1 - 2 * padding[i])
-            - dilation[i] * (weight_shape[i + 2] - 1)
-            for i in range(ndim)
-        ]
-
-    class Conv2d(autograd.Function):
-        @staticmethod
-        def forward(ctx, input, weight, bias):
-            if not transpose:
-                out = F.conv2d(input=input, weight=weight, bias=bias, **common_kwargs)
-
-            else:
-                out = F.conv_transpose2d(
-                    input=input,
-                    weight=weight,
-                    bias=bias,
-                    output_padding=output_padding,
-                    **common_kwargs,
-                )
-
-            ctx.save_for_backward(input, weight)
-
-            return out
-
-        @staticmethod
-        def backward(ctx, grad_output):
-            input, weight = ctx.saved_tensors
-            grad_input, grad_weight, grad_bias = None, None, None
-
-            if ctx.needs_input_grad[0]:
-                p = calc_output_padding(
-                    input_shape=input.shape, output_shape=grad_output.shape
-                )
-                grad_input = conv2d_gradfix(
-                    transpose=(not transpose),
-                    weight_shape=weight_shape,
-                    output_padding=p,
-                    **common_kwargs,
-                ).apply(grad_output, weight, None)
-
-            if ctx.needs_input_grad[1] and not weight_gradients_disabled:
-                grad_weight = Conv2dGradWeight.apply(grad_output, input)
-
-            if ctx.needs_input_grad[2]:
-                grad_bias = grad_output.sum((0, 2, 3))
-
-            return grad_input, grad_weight, grad_bias
-
-    class Conv2dGradWeight(autograd.Function):
-        @staticmethod
-        def forward(ctx, grad_output, input):
-            op = torch._C._jit_get_operation(
-                "aten::cudnn_convolution_backward_weight"
-                if not transpose
-                else "aten::cudnn_convolution_transpose_backward_weight"
-            )
-            flags = [
-                torch.backends.cudnn.benchmark,
-                torch.backends.cudnn.deterministic,
-                torch.backends.cudnn.allow_tf32,
-            ]
-            grad_weight = op(
-                weight_shape,
-                grad_output,
-                input,
-                padding,
-                stride,
-                dilation,
-                groups,
-                *flags,
-            )
-            ctx.save_for_backward(grad_output, input)
-
-            return grad_weight
-
-        @staticmethod
-        def backward(ctx, grad_grad_weight):
-            grad_output, input = ctx.saved_tensors
-            grad_grad_output, grad_grad_input = None, None
-
-            if ctx.needs_input_grad[0]:
-                grad_grad_output = Conv2d.apply(input, grad_grad_weight, None)
-
-            if ctx.needs_input_grad[1]:
-                p = calc_output_padding(
-                    input_shape=input.shape, output_shape=grad_output.shape
-                )
-                grad_grad_input = conv2d_gradfix(
-                    transpose=(not transpose),
-                    weight_shape=weight_shape,
-                    output_padding=p,
-                    **common_kwargs,
-                ).apply(grad_output, grad_grad_weight, None)
-
-            return grad_grad_output, grad_grad_input
-
-    conv2d_gradfix_cache[key] = Conv2d
-
-    return Conv2d

models/stylegan2/op/fused_act.py

@@ -2,72 +2,10 @@ import os
 
 import torch
 from torch import nn
-from torch.autograd import Function
-from torch.utils.cpp_extension import load
-
+from torch.nn import functional as F
 
 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'),
-    ],
-)
-
-
-class FusedLeakyReLUFunctionBackward(Function):
-    @staticmethod
-    def forward(ctx, grad_output, out, negative_slope, scale):
-        ctx.save_for_backward(out)
-        ctx.negative_slope = negative_slope
-        ctx.scale = scale
-
-        empty = grad_output.new_empty(0)
-
-        grad_input = fused.fused_bias_act(
-            grad_output, empty, out, 3, 1, negative_slope, scale
-        )
-
-        dim = [0]
-
-        if grad_input.ndim > 2:
-            dim += list(range(2, grad_input.ndim))
-
-        grad_bias = grad_input.sum(dim).detach()
-
-        return grad_input, grad_bias
-
-    @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
-        )
-
-        return gradgrad_out, None, None, None
-
 
-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)
-        ctx.save_for_backward(out)
-        ctx.negative_slope = negative_slope
-        ctx.scale = scale
-
-        return out
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        out, = ctx.saved_tensors
-
-        grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply(
-            grad_output, out, ctx.negative_slope, ctx.scale
-        )
-
-        return grad_input, grad_bias, None, None
 
 
 class FusedLeakyReLU(nn.Module):
@@ -83,4 +21,20 @@ class FusedLeakyReLU(nn.Module):
 
 
 def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
-    return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale)
+    rest_dim = [1] * (input.ndim - bias.ndim - 1)
+    input = input.cuda()
+    if input.ndim == 3:
+        return (
+            F.leaky_relu(
+                input + bias.view(1, *rest_dim, bias.shape[0]), negative_slope=negative_slope
+            )
+            * scale
+        )
+    else:
+        return (
+            F.leaky_relu(
+                input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=negative_slope
+            )
+            * scale
+        )
+

models/stylegan2/op/upfirdn2d.py

@@ -1,149 +1,16 @@
 import os
 
 import torch
-from torch.autograd import Function
-from torch.utils.cpp_extension import load
+from torch.nn import functional as F
 
 
 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'),
-    ],
-)
 
 
-class UpFirDn2dBackward(Function):
-    @staticmethod
-    def forward(
-        ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size
-    ):
-
-        up_x, up_y = up
-        down_x, down_y = down
-        g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad
-
-        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,
-        )
-        grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3])
-
-        ctx.save_for_backward(kernel)
-
-        pad_x0, pad_x1, pad_y0, pad_y1 = pad
-
-        ctx.up_x = up_x
-        ctx.up_y = up_y
-        ctx.down_x = down_x
-        ctx.down_y = down_y
-        ctx.pad_x0 = pad_x0
-        ctx.pad_x1 = pad_x1
-        ctx.pad_y0 = pad_y0
-        ctx.pad_y1 = pad_y1
-        ctx.in_size = in_size
-        ctx.out_size = out_size
-
-        return grad_input
-
-    @staticmethod
-    def backward(ctx, gradgrad_input):
-        kernel, = ctx.saved_tensors
-
-        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,
-        )
-        # 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]
-        )
-
-        return gradgrad_out, None, None, None, None, None, None, None, None
-
-
-class UpFirDn2d(Function):
-    @staticmethod
-    def forward(ctx, input, kernel, up, down, pad):
-        up_x, up_y = up
-        down_x, down_y = down
-        pad_x0, pad_x1, pad_y0, pad_y1 = pad
-
-        kernel_h, kernel_w = kernel.shape
-        batch, channel, in_h, in_w = input.shape
-        ctx.in_size = input.shape
-
-        input = input.reshape(-1, in_h, in_w, 1)
-
-        ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1]))
-
-        out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
-        out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1
-        ctx.out_size = (out_h, out_w)
-
-        ctx.up = (up_x, up_y)
-        ctx.down = (down_x, down_y)
-        ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1)
-
-        g_pad_x0 = kernel_w - pad_x0 - 1
-        g_pad_y0 = kernel_h - pad_y0 - 1
-        g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1
-        g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1
-
-        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
-        )
-        # out = out.view(major, out_h, out_w, minor)
-        out = out.view(-1, channel, out_h, out_w)
-
-        return out
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        kernel, grad_kernel = ctx.saved_tensors
-
-        grad_input = UpFirDn2dBackward.apply(
-            grad_output,
-            kernel,
-            grad_kernel,
-            ctx.up,
-            ctx.down,
-            ctx.pad,
-            ctx.g_pad,
-            ctx.in_size,
-            ctx.out_size,
-        )
-
-        return grad_input, None, None, None, None
-
 
 def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
-    out = UpFirDn2d.apply(
-        input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1])
+    out = upfirdn2d_native(
+        input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1]
     )
 
     return out
@@ -152,6 +19,9 @@ def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
 def upfirdn2d_native(
     input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
 ):
+    _, channel, in_h, in_w = input.shape
+    input = input.reshape(-1, in_h, in_w, 1)
+
     _, in_h, in_w, minor = input.shape
     kernel_h, kernel_w = kernel.shape
 
@@ -182,6 +52,9 @@ def upfirdn2d_native(
         in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
     )
     out = out.permute(0, 2, 3, 1)
+    out = out[:, ::down_y, ::down_x, :]
 
-    return out[:, ::down_y, ::down_x, :]
+    out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
+    out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1
 
+    return out.view(-1, channel, out_h, out_w)