fix

app.py

@@ -15,9 +15,6 @@ from tqdm import tqdm
 import lpips
 import time
 
-
-#from e4e_projection import projection as e4e_projection
-
 from copy import deepcopy
 import imageio
 

e4e_projection.py

@@ -1,38 +0,0 @@
-import os
-import sys
-import numpy as np
-from PIL import Image
-import torch
-import torchvision.transforms as transforms
-from argparse import Namespace
-from e4e.models.psp import pSp
-from util import *
-
-
-
-@ torch.no_grad()
-def projection(img, name, device='cuda'):
-
- 
-    model_path = 'e4e_ffhq_encode.pt'
-    ckpt = torch.load(model_path, map_location='cpu')
-    opts = ckpt['opts']
-    opts['checkpoint_path'] = model_path
-    opts= Namespace(**opts)
-    net = pSp(opts, device).eval().to(device)
-
-    transform = transforms.Compose(
-        [
-            transforms.Resize(256),
-            transforms.CenterCrop(256),
-            transforms.ToTensor(),
-            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
-        ]
-    )
-
-    img = transform(img).unsqueeze(0).to(device)
-    images, w_plus = net(img, randomize_noise=False, return_latents=True)
-    result_file = {}
-    result_file['latent'] = w_plus[0]
-    torch.save(result_file, name)
-    return w_plus[0]

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

op/fused_act.py

@@ -1,86 +0,0 @@
-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'),
-    ],
-)
-
-
-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):
-    def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
-        super().__init__()
-
-        self.bias = nn.Parameter(torch.zeros(channel))
-        self.negative_slope = negative_slope
-        self.scale = scale
-
-    def forward(self, input):
-        return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
-
-
-def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
-    return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale)

op/fused_act_cpu.py

@@ -1,41 +0,0 @@
-import os
-
-import torch
-from torch import nn
-from torch.autograd import Function
-from torch.nn import functional as F
-
-
-module_path = os.path.dirname(__file__)
-
-
-class FusedLeakyReLU(nn.Module):
-    def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
-        super().__init__()
-
-        self.bias = nn.Parameter(torch.zeros(channel))
-        self.negative_slope = negative_slope
-        self.scale = scale
-
-    def forward(self, input):
-        return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
-
-def fused_leaky_relu(input, bias=None, negative_slope=0.2, scale=2 ** 0.5):
-    if input.device.type == "cpu":
-        if bias is not None:
-            rest_dim = [1] * (input.ndim - bias.ndim - 1)
-            return (
-                F.leaky_relu(
-                    input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2
-                )
-                * scale
-            )
-
-        else:
-            return F.leaky_relu(input, negative_slope=0.2) * scale
-
-    else:
-        return FusedLeakyReLUFunction.apply(
-            input.contiguous(), bias, negative_slope, scale
-        )
-

op/fused_bias_act.cpp

@@ -1,21 +0,0 @@
-#include <torch/extension.h>
-
-
-torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
-    int act, int grad, float alpha, float scale);
-
-#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
-#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
-#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
-
-torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
-    int act, int grad, float alpha, float scale) {
-    CHECK_CUDA(input);
-    CHECK_CUDA(bias);
-
-    return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale);
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)");
-}

op/fused_bias_act_kernel.cu

@@ -1,99 +0,0 @@
-// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
-//
-// This work is made available under the Nvidia Source Code License-NC.
-// To view a copy of this license, visit
-// https://nvlabs.github.io/stylegan2/license.html
-
-#include <torch/types.h>
-
-#include <ATen/ATen.h>
-#include <ATen/AccumulateType.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <ATen/cuda/CUDAApplyUtils.cuh>
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-
-template <typename scalar_t>
-static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref,
-    int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) {
-    int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x;
-
-    scalar_t zero = 0.0;
-
-    for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) {
-        scalar_t x = p_x[xi];
-
-        if (use_bias) {
-            x += p_b[(xi / step_b) % size_b];
-        }
-
-        scalar_t ref = use_ref ? p_ref[xi] : zero;
-
-        scalar_t y;
-
-        switch (act * 10 + grad) {
-            default:
-            case 10: y = x; break;
-            case 11: y = x; break;
-            case 12: y = 0.0; break;
-
-            case 30: y = (x > 0.0) ? x : x * alpha; break;
-            case 31: y = (ref > 0.0) ? x : x * alpha; break;
-            case 32: y = 0.0; break;
-        }
-
-        out[xi] = y * scale;
-    }
-}
-
-
-torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
-    int act, int grad, float alpha, float scale) {
-    int curDevice = -1;
-    cudaGetDevice(&curDevice);
-    cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);
-
-    auto x = input.contiguous();
-    auto b = bias.contiguous();
-    auto ref = refer.contiguous();
-
-    int use_bias = b.numel() ? 1 : 0;
-    int use_ref = ref.numel() ? 1 : 0;
-
-    int size_x = x.numel();
-    int size_b = b.numel();
-    int step_b = 1;
-
-    for (int i = 1 + 1; i < x.dim(); i++) {
-        step_b *= x.size(i);
-    }
-
-    int loop_x = 4;
-    int block_size = 4 * 32;
-    int grid_size = (size_x - 1) / (loop_x * block_size) + 1;
-
-    auto y = torch::empty_like(x);
-
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] {
-        fused_bias_act_kernel<scalar_t><<<grid_size, block_size, 0, stream>>>(
-            y.data_ptr<scalar_t>(),
-            x.data_ptr<scalar_t>(),
-            b.data_ptr<scalar_t>(),
-            ref.data_ptr<scalar_t>(),
-            act,
-            grad,
-            alpha,
-            scale,
-            loop_x,
-            size_x,
-            step_b,
-            size_b,
-            use_bias,
-            use_ref
-        );
-    });
-
-    return y;
-}

op/upfirdn2d.cpp

@@ -1,23 +0,0 @@
-#include <torch/extension.h>
-
-
-torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel,
-                            int up_x, int up_y, int down_x, int down_y,
-                            int pad_x0, int pad_x1, int pad_y0, int pad_y1);
-
-#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
-#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
-#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
-
-torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel,
-                        int up_x, int up_y, int down_x, int down_y,
-                        int pad_x0, int pad_x1, int pad_y0, int pad_y1) {
-    CHECK_CUDA(input);
-    CHECK_CUDA(kernel);
-
-    return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1);
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)");
-}

op/upfirdn2d.py

@@ -1,187 +0,0 @@
-import os
-
-import torch
-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'),
-    ],
-)
-
-
-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])
-    )
-
-    return out
-
-
-def upfirdn2d_native(
-    input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
-):
-    _, in_h, in_w, minor = input.shape
-    kernel_h, kernel_w = kernel.shape
-
-    out = input.view(-1, in_h, 1, in_w, 1, minor)
-    out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
-    out = out.view(-1, in_h * up_y, in_w * up_x, minor)
-
-    out = F.pad(
-        out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]
-    )
-    out = out[
-        :,
-        max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),
-        max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),
-        :,
-    ]
-
-    out = out.permute(0, 3, 1, 2)
-    out = out.reshape(
-        [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]
-    )
-    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
-    out = F.conv2d(out, w)
-    out = out.reshape(
-        -1,
-        minor,
-        in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
-        in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
-    )
-    out = out.permute(0, 2, 3, 1)
-
-    return out[:, ::down_y, ::down_x, :]
-

op/upfirdn2d_cpu.py

@@ -1,60 +0,0 @@
-import os
-
-import torch
-from torch.autograd import Function
-from torch.nn import functional as F
-
-
-
-module_path = os.path.dirname(__file__)
-
-def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
-    out = upfirdn2d_native(
-        input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1]
-    )
-
-    return out
-
-
-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
-
-    out = input.view(-1, in_h, 1, in_w, 1, minor)
-    out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
-    out = out.view(-1, in_h * up_y, in_w * up_x, minor)
-
-    out = F.pad(
-        out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]
-    )
-    out = out[
-        :,
-        max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),
-        max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),
-        :,
-    ]
-
-    out = out.permute(0, 3, 1, 2)
-    out = out.reshape(
-        [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]
-    )
-    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
-    out = F.conv2d(out, w)
-    out = out.reshape(
-        -1,
-        minor,
-        in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
-        in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
-    )
-    out = out.permute(0, 2, 3, 1)
-    out = out[:, ::down_y, ::down_x, :]
-
-    out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h + down_y) // down_y
-    out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w + down_x) // down_x
-
-    return out.view(-1, channel, out_h, out_w)

op/upfirdn2d_kernel.cu

@@ -1,272 +0,0 @@
-// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
-//
-// This work is made available under the Nvidia Source Code License-NC.
-// To view a copy of this license, visit
-// https://nvlabs.github.io/stylegan2/license.html
-
-#include <torch/types.h>
-
-#include <ATen/ATen.h>
-#include <ATen/AccumulateType.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <ATen/cuda/CUDAApplyUtils.cuh>
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-
-static __host__ __device__ __forceinline__ int floor_div(int a, int b) {
-    int c = a / b;
-
-    if (c * b > a) {
-        c--;
-    }
-
-    return c;
-}
-
-
-struct UpFirDn2DKernelParams {
-    int up_x;
-    int up_y;
-    int down_x;
-    int down_y;
-    int pad_x0;
-    int pad_x1;
-    int pad_y0;
-    int pad_y1;
-
-    int major_dim;
-    int in_h;
-    int in_w;
-    int minor_dim;
-    int kernel_h;
-    int kernel_w;
-    int out_h;
-    int out_w;
-    int loop_major;
-    int loop_x;
-};
-
-
-template <typename scalar_t, int up_x, int up_y, int down_x, int down_y, int kernel_h, int kernel_w, int tile_out_h, int tile_out_w>
-__global__ void upfirdn2d_kernel(scalar_t* out, const scalar_t* input, const scalar_t* kernel, const UpFirDn2DKernelParams p) {
-    const int tile_in_h = ((tile_out_h - 1) * down_y + kernel_h - 1) / up_y + 1;
-    const int tile_in_w = ((tile_out_w - 1) * down_x + kernel_w - 1) / up_x + 1;
-
-    __shared__ volatile float sk[kernel_h][kernel_w];
-    __shared__ volatile float sx[tile_in_h][tile_in_w];
-
-    int minor_idx = blockIdx.x;
-    int tile_out_y = minor_idx / p.minor_dim;
-    minor_idx -= tile_out_y * p.minor_dim;
-    tile_out_y *= tile_out_h;
-    int tile_out_x_base = blockIdx.y * p.loop_x * tile_out_w;
-    int major_idx_base = blockIdx.z * p.loop_major;
-
-    if (tile_out_x_base >= p.out_w | tile_out_y >= p.out_h | major_idx_base >= p.major_dim) {
-        return;
-    }
-
-    for (int tap_idx = threadIdx.x; tap_idx < kernel_h * kernel_w; tap_idx += blockDim.x) {
-        int ky = tap_idx / kernel_w;
-        int kx = tap_idx - ky * kernel_w;
-        scalar_t v = 0.0;
-
-        if (kx < p.kernel_w & ky < p.kernel_h) {
-            v = kernel[(p.kernel_h - 1 - ky) * p.kernel_w + (p.kernel_w - 1 - kx)];
-        }
-
-        sk[ky][kx] = v;
-    }
-
-    for (int loop_major = 0, major_idx = major_idx_base; loop_major < p.loop_major & major_idx < p.major_dim; loop_major++, major_idx++) {
-        for (int loop_x = 0, tile_out_x = tile_out_x_base; loop_x < p.loop_x & tile_out_x < p.out_w; loop_x++, tile_out_x += tile_out_w) {
-            int tile_mid_x = tile_out_x * down_x + up_x - 1 - p.pad_x0;
-            int tile_mid_y = tile_out_y * down_y + up_y - 1 - p.pad_y0;
-            int tile_in_x = floor_div(tile_mid_x, up_x);
-            int tile_in_y = floor_div(tile_mid_y, up_y);
-
-            __syncthreads();
-
-            for (int in_idx = threadIdx.x; in_idx < tile_in_h * tile_in_w; in_idx += blockDim.x) {
-                int rel_in_y = in_idx / tile_in_w;
-                int rel_in_x = in_idx - rel_in_y * tile_in_w;
-                int in_x = rel_in_x + tile_in_x;
-                int in_y = rel_in_y + tile_in_y;
-
-                scalar_t v = 0.0;
-
-                if (in_x >= 0 & in_y >= 0 & in_x < p.in_w & in_y < p.in_h) {
-                    v = input[((major_idx * p.in_h + in_y) * p.in_w + in_x) * p.minor_dim + minor_idx];
-                }
-
-                sx[rel_in_y][rel_in_x] = v;
-            }
-
-            __syncthreads();
-            for (int out_idx = threadIdx.x; out_idx < tile_out_h * tile_out_w; out_idx += blockDim.x) {
-                int rel_out_y = out_idx / tile_out_w;
-                int rel_out_x = out_idx - rel_out_y * tile_out_w;
-                int out_x = rel_out_x + tile_out_x;
-                int out_y = rel_out_y + tile_out_y;
-
-                int mid_x = tile_mid_x + rel_out_x * down_x;
-                int mid_y = tile_mid_y + rel_out_y * down_y;
-                int in_x = floor_div(mid_x, up_x);
-                int in_y = floor_div(mid_y, up_y);
-                int rel_in_x = in_x - tile_in_x;
-                int rel_in_y = in_y - tile_in_y;
-                int kernel_x = (in_x + 1) * up_x - mid_x - 1;
-                int kernel_y = (in_y + 1) * up_y - mid_y - 1;
-
-                scalar_t v = 0.0;
-
-                #pragma unroll
-                for (int y = 0; y < kernel_h / up_y; y++)
-                    #pragma unroll
-                    for (int x = 0; x < kernel_w / up_x; x++)
-                        v += sx[rel_in_y + y][rel_in_x + x] * sk[kernel_y + y * up_y][kernel_x + x * up_x];
-
-                if (out_x < p.out_w & out_y < p.out_h) {
-                    out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim + minor_idx] = v;
-                }
-            }
-        }
-    }
-}
-
-
-torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel,
-    int up_x, int up_y, int down_x, int down_y,
-    int pad_x0, int pad_x1, int pad_y0, int pad_y1) {
-    int curDevice = -1;
-    cudaGetDevice(&curDevice);
-    cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);
-
-    UpFirDn2DKernelParams p;
-
-    auto x = input.contiguous();
-    auto k = kernel.contiguous();
-
-    p.major_dim = x.size(0);
-    p.in_h = x.size(1);
-    p.in_w = x.size(2);
-    p.minor_dim = x.size(3);
-    p.kernel_h = k.size(0);
-    p.kernel_w = k.size(1);
-    p.up_x = up_x;
-    p.up_y = up_y;
-    p.down_x = down_x;
-    p.down_y = down_y;
-    p.pad_x0 = pad_x0;
-    p.pad_x1 = pad_x1;
-    p.pad_y0 = pad_y0;
-    p.pad_y1 = pad_y1;
-
-    p.out_h = (p.in_h * p.up_y + p.pad_y0 + p.pad_y1 - p.kernel_h + p.down_y) / p.down_y;
-    p.out_w = (p.in_w * p.up_x + p.pad_x0 + p.pad_x1 - p.kernel_w + p.down_x) / p.down_x;
-
-    auto out = at::empty({p.major_dim, p.out_h, p.out_w, p.minor_dim}, x.options());
-
-    int mode = -1;
-
-    int tile_out_h;
-    int tile_out_w;
-
-    if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 4 && p.kernel_w <= 4) {
-        mode = 1;
-        tile_out_h = 16;
-        tile_out_w = 64;
-    }
-
-    if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 3 && p.kernel_w <= 3) {
-        mode = 2;
-        tile_out_h = 16;
-        tile_out_w = 64;
-    }
-
-    if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 4 && p.kernel_w <= 4) {
-        mode = 3;
-        tile_out_h = 16;
-        tile_out_w = 64;
-    }
-
-    if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 2 && p.kernel_w <= 2) {
-        mode = 4;
-        tile_out_h = 16;
-        tile_out_w = 64;
-    }
-
-    if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 && p.kernel_h <= 4 && p.kernel_w <= 4) {
-        mode = 5;
-        tile_out_h = 8;
-        tile_out_w = 32;
-    }
-
-    if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 && p.kernel_h <= 2 && p.kernel_w <= 2) {
-        mode = 6;
-        tile_out_h = 8;
-        tile_out_w = 32;
-    }
-
-    dim3 block_size;
-    dim3 grid_size;
-
-    if (tile_out_h > 0 && tile_out_w) {
-        p.loop_major = (p.major_dim - 1) / 16384 + 1;
-        p.loop_x = 1;
-        block_size = dim3(32 * 8, 1, 1);
-        grid_size = dim3(((p.out_h - 1) / tile_out_h + 1) * p.minor_dim,
-                         (p.out_w - 1) / (p.loop_x * tile_out_w) + 1,
-                         (p.major_dim - 1) / p.loop_major + 1);
-    }
-
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] {
-        switch (mode) {
-        case 1:
-            upfirdn2d_kernel<scalar_t, 1, 1, 1, 1, 4, 4, 16, 64><<<grid_size, block_size, 0, stream>>>(
-                out.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(), k.data_ptr<scalar_t>(), p
-            );
-
-            break;
-
-        case 2:
-            upfirdn2d_kernel<scalar_t, 1, 1, 1, 1, 3, 3, 16, 64><<<grid_size, block_size, 0, stream>>>(
-                out.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(), k.data_ptr<scalar_t>(), p
-            );
-
-            break;
-
-        case 3:
-            upfirdn2d_kernel<scalar_t, 2, 2, 1, 1, 4, 4, 16, 64><<<grid_size, block_size, 0, stream>>>(
-                out.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(), k.data_ptr<scalar_t>(), p
-            );
-
-            break;
-
-        case 4:
-            upfirdn2d_kernel<scalar_t, 2, 2, 1, 1, 2, 2, 16, 64><<<grid_size, block_size, 0, stream>>>(
-                out.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(), k.data_ptr<scalar_t>(), p
-            );
-
-            break;
-
-        case 5:
-            upfirdn2d_kernel<scalar_t, 1, 1, 2, 2, 4, 4, 8, 32><<<grid_size, block_size, 0, stream>>>(
-                out.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(), k.data_ptr<scalar_t>(), p
-            );
-
-            break;
-
-        case 6:
-            upfirdn2d_kernel<scalar_t, 1, 1, 2, 2, 4, 4, 8, 32><<<grid_size, block_size, 0, stream>>>(
-                out.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(), k.data_ptr<scalar_t>(), p
-            );
-
-            break;
-        }
-    });
-
-    return out;
-}