Add application

app.py

@@ -0,0 +1,836 @@
+import os
+from enum import IntEnum
+from pathlib import Path
+from tempfile import mktemp
+from typing import IO, Dict, Type
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from gradio import Interface, inputs, outputs
+
+DEVICE = "cpu"
+
+WEIGHTS_PATH = Path(__file__).parent / "weights"
+
+AVALIABLE_WEIGHTS = {
+    basename: path
+    for basename, ext in (
+        os.path.splitext(filename) for filename in os.listdir(WEIGHTS_PATH)
+    )
+    if (path := WEIGHTS_PATH / (basename + ext)).is_file() and ext.endswith("pth")
+}
+
+
+class ScaleMode(IntEnum):
+    up2x = 2
+    up3x = 3
+    up4x = 4
+
+
+class TileMode(IntEnum):
+    full = 0
+    half = 1
+    quarter = 2
+    ninth = 3
+    sixteenth = 4
+
+
+class SEBlock(nn.Module):
+    def __init__(self, in_channels, reduction=8, bias=False):
+        super(SEBlock, self).__init__()
+        self.conv1 = nn.Conv2d(
+            in_channels, in_channels // reduction, 1, 1, 0, bias=bias
+        )
+        self.conv2 = nn.Conv2d(
+            in_channels // reduction, in_channels, 1, 1, 0, bias=bias
+        )
+
+    def forward(self, x):
+        if "Half" in x.type():  # torch.HalfTensor/torch.cuda.HalfTensor
+            x0 = torch.mean(x.float(), dim=(2, 3), keepdim=True).half()
+        else:
+            x0 = torch.mean(x, dim=(2, 3), keepdim=True)
+        x0 = self.conv1(x0)
+        x0 = F.relu(x0, inplace=True)
+        x0 = self.conv2(x0)
+        x0 = torch.sigmoid(x0)
+        x = torch.mul(x, x0)
+        return x
+
+    def forward_mean(self, x, x0):
+        x0 = self.conv1(x0)
+        x0 = F.relu(x0, inplace=True)
+        x0 = self.conv2(x0)
+        x0 = torch.sigmoid(x0)
+        x = torch.mul(x, x0)
+        return x
+
+
+class UNetConv(nn.Module):
+    def __init__(self, in_channels, mid_channels, out_channels, se):
+        super(UNetConv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, 3, 1, 0),
+            nn.LeakyReLU(0.1, inplace=True),
+            nn.Conv2d(mid_channels, out_channels, 3, 1, 0),
+            nn.LeakyReLU(0.1, inplace=True),
+        )
+        if se:
+            self.seblock = SEBlock(out_channels, reduction=8, bias=True)
+        else:
+            self.seblock = None
+
+    def forward(self, x):
+        z = self.conv(x)
+        if self.seblock is not None:
+            z = self.seblock(z)
+        return z
+
+
+class UNet1(nn.Module):
+    def __init__(self, in_channels, out_channels, deconv):
+        super(UNet1, self).__init__()
+        self.conv1 = UNetConv(in_channels, 32, 64, se=False)
+        self.conv1_down = nn.Conv2d(64, 64, 2, 2, 0)
+        self.conv2 = UNetConv(64, 128, 64, se=True)
+        self.conv2_up = nn.ConvTranspose2d(64, 64, 2, 2, 0)
+        self.conv3 = nn.Conv2d(64, 64, 3, 1, 0)
+
+        if deconv:
+            self.conv_bottom = nn.ConvTranspose2d(64, out_channels, 4, 2, 3)
+        else:
+            self.conv_bottom = nn.Conv2d(64, out_channels, 3, 1, 0)
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x1 = self.conv1(x)
+        x2 = self.conv1_down(x1)
+        x2 = F.leaky_relu(x2, 0.1, inplace=True)
+        x2 = self.conv2(x2)
+        x2 = self.conv2_up(x2)
+        x2 = F.leaky_relu(x2, 0.1, inplace=True)
+
+        x1 = F.pad(x1, (-4, -4, -4, -4))
+        x3 = self.conv3(x1 + x2)
+        x3 = F.leaky_relu(x3, 0.1, inplace=True)
+        z = self.conv_bottom(x3)
+        return z
+
+    def forward_a(self, x):
+        x1 = self.conv1(x)
+        x2 = self.conv1_down(x1)
+        x2 = F.leaky_relu(x2, 0.1, inplace=True)
+        x2 = self.conv2.conv(x2)
+        return x1, x2
+
+    def forward_b(self, x1, x2):
+        x2 = self.conv2_up(x2)
+        x2 = F.leaky_relu(x2, 0.1, inplace=True)
+
+        x1 = F.pad(x1, (-4, -4, -4, -4))
+        x3 = self.conv3(x1 + x2)
+        x3 = F.leaky_relu(x3, 0.1, inplace=True)
+        z = self.conv_bottom(x3)
+        return z
+
+
+class UNet1x3(nn.Module):
+    def __init__(self, in_channels, out_channels, deconv):
+        super(UNet1x3, self).__init__()
+        self.conv1 = UNetConv(in_channels, 32, 64, se=False)
+        self.conv1_down = nn.Conv2d(64, 64, 2, 2, 0)
+        self.conv2 = UNetConv(64, 128, 64, se=True)
+        self.conv2_up = nn.ConvTranspose2d(64, 64, 2, 2, 0)
+        self.conv3 = nn.Conv2d(64, 64, 3, 1, 0)
+
+        if deconv:
+            self.conv_bottom = nn.ConvTranspose2d(64, out_channels, 5, 3, 2)
+        else:
+            self.conv_bottom = nn.Conv2d(64, out_channels, 3, 1, 0)
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x1 = self.conv1(x)
+        x2 = self.conv1_down(x1)
+        x2 = F.leaky_relu(x2, 0.1, inplace=True)
+        x2 = self.conv2(x2)
+        x2 = self.conv2_up(x2)
+        x2 = F.leaky_relu(x2, 0.1, inplace=True)
+
+        x1 = F.pad(x1, (-4, -4, -4, -4))
+        x3 = self.conv3(x1 + x2)
+        x3 = F.leaky_relu(x3, 0.1, inplace=True)
+        z = self.conv_bottom(x3)
+        return z
+
+    def forward_a(self, x):
+        x1 = self.conv1(x)
+        x2 = self.conv1_down(x1)
+        x2 = F.leaky_relu(x2, 0.1, inplace=True)
+        x2 = self.conv2.conv(x2)
+        return x1, x2
+
+    def forward_b(self, x1, x2):
+        x2 = self.conv2_up(x2)
+        x2 = F.leaky_relu(x2, 0.1, inplace=True)
+
+        x1 = F.pad(x1, (-4, -4, -4, -4))
+        x3 = self.conv3(x1 + x2)
+        x3 = F.leaky_relu(x3, 0.1, inplace=True)
+        z = self.conv_bottom(x3)
+        return z
+
+
+class UNet2(nn.Module):
+    def __init__(self, in_channels, out_channels, deconv):
+        super(UNet2, self).__init__()
+
+        self.conv1 = UNetConv(in_channels, 32, 64, se=False)
+        self.conv1_down = nn.Conv2d(64, 64, 2, 2, 0)
+        self.conv2 = UNetConv(64, 64, 128, se=True)
+        self.conv2_down = nn.Conv2d(128, 128, 2, 2, 0)
+        self.conv3 = UNetConv(128, 256, 128, se=True)
+        self.conv3_up = nn.ConvTranspose2d(128, 128, 2, 2, 0)
+        self.conv4 = UNetConv(128, 64, 64, se=True)
+        self.conv4_up = nn.ConvTranspose2d(64, 64, 2, 2, 0)
+        self.conv5 = nn.Conv2d(64, 64, 3, 1, 0)
+
+        if deconv:
+            self.conv_bottom = nn.ConvTranspose2d(64, out_channels, 4, 2, 3)
+        else:
+            self.conv_bottom = nn.Conv2d(64, out_channels, 3, 1, 0)
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x1 = self.conv1(x)
+        x2 = self.conv1_down(x1)
+        x2 = F.leaky_relu(x2, 0.1, inplace=True)
+        x2 = self.conv2(x2)
+
+        x3 = self.conv2_down(x2)
+        x3 = F.leaky_relu(x3, 0.1, inplace=True)
+        x3 = self.conv3(x3)
+        x3 = self.conv3_up(x3)
+        x3 = F.leaky_relu(x3, 0.1, inplace=True)
+
+        x2 = F.pad(x2, (-4, -4, -4, -4))
+        x4 = self.conv4(x2 + x3)
+        x4 = self.conv4_up(x4)
+        x4 = F.leaky_relu(x4, 0.1, inplace=True)
+
+        x1 = F.pad(x1, (-16, -16, -16, -16))
+        x5 = self.conv5(x1 + x4)
+        x5 = F.leaky_relu(x5, 0.1, inplace=True)
+
+        z = self.conv_bottom(x5)
+        return z
+
+    def forward_a(self, x):  # conv234结尾有se
+        x1 = self.conv1(x)
+        x2 = self.conv1_down(x1)
+        x2 = F.leaky_relu(x2, 0.1, inplace=True)
+        x2 = self.conv2.conv(x2)
+        return x1, x2
+
+    def forward_b(self, x2):  # conv234结尾有se
+        x3 = self.conv2_down(x2)
+        x3 = F.leaky_relu(x3, 0.1, inplace=True)
+        x3 = self.conv3.conv(x3)
+        return x3
+
+    def forward_c(self, x2, x3):  # conv234结尾有se
+        x3 = self.conv3_up(x3)
+        x3 = F.leaky_relu(x3, 0.1, inplace=True)
+
+        x2 = F.pad(x2, (-4, -4, -4, -4))
+        x4 = self.conv4.conv(x2 + x3)
+        return x4
+
+    def forward_d(self, x1, x4):  # conv234结尾有se
+        x4 = self.conv4_up(x4)
+        x4 = F.leaky_relu(x4, 0.1, inplace=True)
+
+        x1 = F.pad(x1, (-16, -16, -16, -16))
+        x5 = self.conv5(x1 + x4)
+        x5 = F.leaky_relu(x5, 0.1, inplace=True)
+
+        z = self.conv_bottom(x5)
+        return z
+
+
+class UpCunet2x(nn.Module):  # 完美tile，全程无损
+    def __init__(self, in_channels=3, out_channels=3):
+        super(UpCunet2x, self).__init__()
+        self.unet1 = UNet1(in_channels, out_channels, deconv=True)
+        self.unet2 = UNet2(in_channels, out_channels, deconv=False)
+
+    def forward(self, x, tile_mode):  # 1.7G
+        n, c, h0, w0 = x.shape
+        if tile_mode == 0:  # 不tile
+            ph = ((h0 - 1) // 2 + 1) * 2
+            pw = ((w0 - 1) // 2 + 1) * 2
+            x = F.pad(x, (18, 18 + pw - w0, 18, 18 + ph - h0), "reflect")  # 需要保证被2整除
+            x = self.unet1.forward(x)
+            x0 = self.unet2.forward(x)
+            x1 = F.pad(x, (-20, -20, -20, -20))
+            x = torch.add(x0, x1)
+            if w0 != pw or h0 != ph:
+                x = x[:, :, : h0 * 2, : w0 * 2]
+            return x
+        elif tile_mode == 1:  # 对长边减半
+            if w0 >= h0:
+                crop_size_w = ((w0 - 1) // 4 * 4 + 4) // 2  # 减半后能被2整除，所以要先被4整除
+                crop_size_h = (h0 - 1) // 2 * 2 + 2  # 能被2整除
+            else:
+                crop_size_h = ((h0 - 1) // 4 * 4 + 4) // 2  # 减半后能被2整除，所以要先被4整除
+                crop_size_w = (w0 - 1) // 2 * 2 + 2  # 能被2整除
+            crop_size = (crop_size_h, crop_size_w)  # 6.6G
+        elif tile_mode == 2:  # hw都减半
+            crop_size = (
+                ((h0 - 1) // 4 * 4 + 4) // 2,
+                ((w0 - 1) // 4 * 4 + 4) // 2,
+            )  # 5.6G
+        elif tile_mode == 3:  # hw都三分之一
+            crop_size = (
+                ((h0 - 1) // 6 * 6 + 6) // 3,
+                ((w0 - 1) // 6 * 6 + 6) // 3,
+            )  # 4.2G
+        elif tile_mode == 4:  # hw都四分之一
+            crop_size = (
+                ((h0 - 1) // 8 * 8 + 8) // 4,
+                ((w0 - 1) // 8 * 8 + 8) // 4,
+            )  # 3.7G
+        ph = ((h0 - 1) // crop_size[0] + 1) * crop_size[0]
+        pw = ((w0 - 1) // crop_size[1] + 1) * crop_size[1]
+        x = F.pad(x, (18, 18 + pw - w0, 18, 18 + ph - h0), "reflect")
+        n, c, h, w = x.shape
+        se_mean0 = torch.zeros((n, 64, 1, 1)).to(x.device)
+        if "Half" in x.type():
+            se_mean0 = se_mean0.half()
+        n_patch = 0
+        tmp_dict = {}
+        opt_res_dict = {}
+        for i in range(0, h - 36, crop_size[0]):
+            tmp_dict[i] = {}
+            for j in range(0, w - 36, crop_size[1]):
+                x_crop = x[:, :, i : i + crop_size[0] + 36, j : j + crop_size[1] + 36]
+                n, c1, h1, w1 = x_crop.shape
+                tmp0, x_crop = self.unet1.forward_a(x_crop)
+                if "Half" in x.type():  # torch.HalfTensor/torch.cuda.HalfTensor
+                    tmp_se_mean = torch.mean(
+                        x_crop.float(), dim=(2, 3), keepdim=True
+                    ).half()
+                else:
+                    tmp_se_mean = torch.mean(x_crop, dim=(2, 3), keepdim=True)
+                se_mean0 += tmp_se_mean
+                n_patch += 1
+                tmp_dict[i][j] = (tmp0, x_crop)
+        se_mean0 /= n_patch
+        se_mean1 = torch.zeros((n, 128, 1, 1)).to(x.device)  # 64#128#128#64
+        if "Half" in x.type():
+            se_mean1 = se_mean1.half()
+        for i in range(0, h - 36, crop_size[0]):
+            for j in range(0, w - 36, crop_size[1]):
+                tmp0, x_crop = tmp_dict[i][j]
+                x_crop = self.unet1.conv2.seblock.forward_mean(x_crop, se_mean0)
+                opt_unet1 = self.unet1.forward_b(tmp0, x_crop)
+                tmp_x1, tmp_x2 = self.unet2.forward_a(opt_unet1)
+                if "Half" in x.type():  # torch.HalfTensor/torch.cuda.HalfTensor
+                    tmp_se_mean = torch.mean(
+                        tmp_x2.float(), dim=(2, 3), keepdim=True
+                    ).half()
+                else:
+                    tmp_se_mean = torch.mean(tmp_x2, dim=(2, 3), keepdim=True)
+                se_mean1 += tmp_se_mean
+                tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x2)
+        se_mean1 /= n_patch
+        se_mean0 = torch.zeros((n, 128, 1, 1)).to(x.device)  # 64#128#128#64
+        if "Half" in x.type():
+            se_mean0 = se_mean0.half()
+        for i in range(0, h - 36, crop_size[0]):
+            for j in range(0, w - 36, crop_size[1]):
+                opt_unet1, tmp_x1, tmp_x2 = tmp_dict[i][j]
+                tmp_x2 = self.unet2.conv2.seblock.forward_mean(tmp_x2, se_mean1)
+                tmp_x3 = self.unet2.forward_b(tmp_x2)
+                if "Half" in x.type():  # torch.HalfTensor/torch.cuda.HalfTensor
+                    tmp_se_mean = torch.mean(
+                        tmp_x3.float(), dim=(2, 3), keepdim=True
+                    ).half()
+                else:
+                    tmp_se_mean = torch.mean(tmp_x3, dim=(2, 3), keepdim=True)
+                se_mean0 += tmp_se_mean
+                tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x2, tmp_x3)
+        se_mean0 /= n_patch
+        se_mean1 = torch.zeros((n, 64, 1, 1)).to(x.device)  # 64#128#128#64
+        if "Half" in x.type():
+            se_mean1 = se_mean1.half()
+        for i in range(0, h - 36, crop_size[0]):
+            for j in range(0, w - 36, crop_size[1]):
+                opt_unet1, tmp_x1, tmp_x2, tmp_x3 = tmp_dict[i][j]
+                tmp_x3 = self.unet2.conv3.seblock.forward_mean(tmp_x3, se_mean0)
+                tmp_x4 = self.unet2.forward_c(tmp_x2, tmp_x3)
+                if "Half" in x.type():  # torch.HalfTensor/torch.cuda.HalfTensor
+                    tmp_se_mean = torch.mean(
+                        tmp_x4.float(), dim=(2, 3), keepdim=True
+                    ).half()
+                else:
+                    tmp_se_mean = torch.mean(tmp_x4, dim=(2, 3), keepdim=True)
+                se_mean1 += tmp_se_mean
+                tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x4)
+        se_mean1 /= n_patch
+        for i in range(0, h - 36, crop_size[0]):
+            opt_res_dict[i] = {}
+            for j in range(0, w - 36, crop_size[1]):
+                opt_unet1, tmp_x1, tmp_x4 = tmp_dict[i][j]
+                tmp_x4 = self.unet2.conv4.seblock.forward_mean(tmp_x4, se_mean1)
+                x0 = self.unet2.forward_d(tmp_x1, tmp_x4)
+                x1 = F.pad(opt_unet1, (-20, -20, -20, -20))
+                x_crop = torch.add(x0, x1)  # x0是unet2的最终输出
+                opt_res_dict[i][j] = x_crop
+        del tmp_dict
+        torch.cuda.empty_cache()
+        res = torch.zeros((n, c, h * 2 - 72, w * 2 - 72)).to(x.device)
+        if "Half" in x.type():
+            res = res.half()
+        for i in range(0, h - 36, crop_size[0]):
+            for j in range(0, w - 36, crop_size[1]):
+                res[
+                    :, :, i * 2 : i * 2 + h1 * 2 - 72, j * 2 : j * 2 + w1 * 2 - 72
+                ] = opt_res_dict[i][j]
+        del opt_res_dict
+        torch.cuda.empty_cache()
+        if w0 != pw or h0 != ph:
+            res = res[:, :, : h0 * 2, : w0 * 2]
+        return res  #
+
+
+class UpCunet3x(nn.Module):  # 完美tile，全程无损
+    def __init__(self, in_channels=3, out_channels=3):
+        super(UpCunet3x, self).__init__()
+        self.unet1 = UNet1x3(in_channels, out_channels, deconv=True)
+        self.unet2 = UNet2(in_channels, out_channels, deconv=False)
+
+    def forward(self, x, tile_mode):  # 1.7G
+        n, c, h0, w0 = x.shape
+        if tile_mode == 0:  # 不tile
+            ph = ((h0 - 1) // 4 + 1) * 4
+            pw = ((w0 - 1) // 4 + 1) * 4
+            x = F.pad(x, (14, 14 + pw - w0, 14, 14 + ph - h0), "reflect")  # 需要保证被2整除
+            x = self.unet1.forward(x)
+            x0 = self.unet2.forward(x)
+            x1 = F.pad(x, (-20, -20, -20, -20))
+            x = torch.add(x0, x1)
+            if w0 != pw or h0 != ph:
+                x = x[:, :, : h0 * 3, : w0 * 3]
+            return x
+        elif tile_mode == 1:  # 对长边减半
+            if w0 >= h0:
+                crop_size_w = ((w0 - 1) // 8 * 8 + 8) // 2  # 减半后能被4整除，所以要先被8整除
+                crop_size_h = (h0 - 1) // 4 * 4 + 4  # 能被4整除
+            else:
+                crop_size_h = ((h0 - 1) // 8 * 8 + 8) // 2  # 减半后能被4整除，所以要先被8整除
+                crop_size_w = (w0 - 1) // 4 * 4 + 4  # 能被4整除
+            crop_size = (crop_size_h, crop_size_w)  # 6.6G
+        elif tile_mode == 2:  # hw都减半
+            crop_size = (
+                ((h0 - 1) // 8 * 8 + 8) // 2,
+                ((w0 - 1) // 8 * 8 + 8) // 2,
+            )  # 5.6G
+        elif tile_mode == 3:  # hw都三分之一
+            crop_size = (
+                ((h0 - 1) // 12 * 12 + 12) // 3,
+                ((w0 - 1) // 12 * 12 + 12) // 3,
+            )  # 4.2G
+        elif tile_mode == 4:  # hw都四分之一
+            crop_size = (
+                ((h0 - 1) // 16 * 16 + 16) // 4,
+                ((w0 - 1) // 16 * 16 + 16) // 4,
+            )  # 3.7G
+        ph = ((h0 - 1) // crop_size[0] + 1) * crop_size[0]
+        pw = ((w0 - 1) // crop_size[1] + 1) * crop_size[1]
+        x = F.pad(x, (14, 14 + pw - w0, 14, 14 + ph - h0), "reflect")
+        n, c, h, w = x.shape
+        se_mean0 = torch.zeros((n, 64, 1, 1)).to(x.device)
+        if "Half" in x.type():
+            se_mean0 = se_mean0.half()
+        n_patch = 0
+        tmp_dict = {}
+        opt_res_dict = {}
+        for i in range(0, h - 28, crop_size[0]):
+            tmp_dict[i] = {}
+            for j in range(0, w - 28, crop_size[1]):
+                x_crop = x[:, :, i : i + crop_size[0] + 28, j : j + crop_size[1] + 28]
+                n, c1, h1, w1 = x_crop.shape
+                tmp0, x_crop = self.unet1.forward_a(x_crop)
+                if "Half" in x.type():  # torch.HalfTensor/torch.cuda.HalfTensor
+                    tmp_se_mean = torch.mean(
+                        x_crop.float(), dim=(2, 3), keepdim=True
+                    ).half()
+                else:
+                    tmp_se_mean = torch.mean(x_crop, dim=(2, 3), keepdim=True)
+                se_mean0 += tmp_se_mean
+                n_patch += 1
+                tmp_dict[i][j] = (tmp0, x_crop)
+        se_mean0 /= n_patch
+        se_mean1 = torch.zeros((n, 128, 1, 1)).to(x.device)  # 64#128#128#64
+        if "Half" in x.type():
+            se_mean1 = se_mean1.half()
+        for i in range(0, h - 28, crop_size[0]):
+            for j in range(0, w - 28, crop_size[1]):
+                tmp0, x_crop = tmp_dict[i][j]
+                x_crop = self.unet1.conv2.seblock.forward_mean(x_crop, se_mean0)
+                opt_unet1 = self.unet1.forward_b(tmp0, x_crop)
+                tmp_x1, tmp_x2 = self.unet2.forward_a(opt_unet1)
+                if "Half" in x.type():  # torch.HalfTensor/torch.cuda.HalfTensor
+                    tmp_se_mean = torch.mean(
+                        tmp_x2.float(), dim=(2, 3), keepdim=True
+                    ).half()
+                else:
+                    tmp_se_mean = torch.mean(tmp_x2, dim=(2, 3), keepdim=True)
+                se_mean1 += tmp_se_mean
+                tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x2)
+        se_mean1 /= n_patch
+        se_mean0 = torch.zeros((n, 128, 1, 1)).to(x.device)  # 64#128#128#64
+        if "Half" in x.type():
+            se_mean0 = se_mean0.half()
+        for i in range(0, h - 28, crop_size[0]):
+            for j in range(0, w - 28, crop_size[1]):
+                opt_unet1, tmp_x1, tmp_x2 = tmp_dict[i][j]
+                tmp_x2 = self.unet2.conv2.seblock.forward_mean(tmp_x2, se_mean1)
+                tmp_x3 = self.unet2.forward_b(tmp_x2)
+                if "Half" in x.type():  # torch.HalfTensor/torch.cuda.HalfTensor
+                    tmp_se_mean = torch.mean(
+                        tmp_x3.float(), dim=(2, 3), keepdim=True
+                    ).half()
+                else:
+                    tmp_se_mean = torch.mean(tmp_x3, dim=(2, 3), keepdim=True)
+                se_mean0 += tmp_se_mean
+                tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x2, tmp_x3)
+        se_mean0 /= n_patch
+        se_mean1 = torch.zeros((n, 64, 1, 1)).to(x.device)  # 64#128#128#64
+        if "Half" in x.type():
+            se_mean1 = se_mean1.half()
+        for i in range(0, h - 28, crop_size[0]):
+            for j in range(0, w - 28, crop_size[1]):
+                opt_unet1, tmp_x1, tmp_x2, tmp_x3 = tmp_dict[i][j]
+                tmp_x3 = self.unet2.conv3.seblock.forward_mean(tmp_x3, se_mean0)
+                tmp_x4 = self.unet2.forward_c(tmp_x2, tmp_x3)
+                if "Half" in x.type():  # torch.HalfTensor/torch.cuda.HalfTensor
+                    tmp_se_mean = torch.mean(
+                        tmp_x4.float(), dim=(2, 3), keepdim=True
+                    ).half()
+                else:
+                    tmp_se_mean = torch.mean(tmp_x4, dim=(2, 3), keepdim=True)
+                se_mean1 += tmp_se_mean
+                tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x4)
+        se_mean1 /= n_patch
+        for i in range(0, h - 28, crop_size[0]):
+            opt_res_dict[i] = {}
+            for j in range(0, w - 28, crop_size[1]):
+                opt_unet1, tmp_x1, tmp_x4 = tmp_dict[i][j]
+                tmp_x4 = self.unet2.conv4.seblock.forward_mean(tmp_x4, se_mean1)
+                x0 = self.unet2.forward_d(tmp_x1, tmp_x4)
+                x1 = F.pad(opt_unet1, (-20, -20, -20, -20))
+                x_crop = torch.add(x0, x1)  # x0是unet2的最终输出
+                opt_res_dict[i][j] = x_crop  #
+        del tmp_dict
+        torch.cuda.empty_cache()
+        res = torch.zeros((n, c, h * 3 - 84, w * 3 - 84)).to(x.device)
+        if "Half" in x.type():
+            res = res.half()
+        for i in range(0, h - 28, crop_size[0]):
+            for j in range(0, w - 28, crop_size[1]):
+                res[
+                    :, :, i * 3 : i * 3 + h1 * 3 - 84, j * 3 : j * 3 + w1 * 3 - 84
+                ] = opt_res_dict[i][j]
+        del opt_res_dict
+        torch.cuda.empty_cache()
+        if w0 != pw or h0 != ph:
+            res = res[:, :, : h0 * 3, : w0 * 3]
+        return res
+
+
+class UpCunet4x(nn.Module):  # 完美tile，全程无损
+    def __init__(self, in_channels=3, out_channels=3):
+        super(UpCunet4x, self).__init__()
+        self.unet1 = UNet1(in_channels, 64, deconv=True)
+        self.unet2 = UNet2(64, 64, deconv=False)
+        self.ps = nn.PixelShuffle(2)
+        self.conv_final = nn.Conv2d(64, 12, 3, 1, padding=0, bias=True)
+
+    def forward(self, x, tile_mode):
+        n, c, h0, w0 = x.shape
+        x00 = x
+        if tile_mode == 0:  # 不tile
+            ph = ((h0 - 1) // 2 + 1) * 2
+            pw = ((w0 - 1) // 2 + 1) * 2
+            x = F.pad(x, (19, 19 + pw - w0, 19, 19 + ph - h0), "reflect")  # 需要保证被2整除
+            x = self.unet1.forward(x)
+            x0 = self.unet2.forward(x)
+            x1 = F.pad(x, (-20, -20, -20, -20))
+            x = torch.add(x0, x1)
+            x = self.conv_final(x)
+            x = F.pad(x, (-1, -1, -1, -1))
+            x = self.ps(x)
+            if w0 != pw or h0 != ph:
+                x = x[:, :, : h0 * 4, : w0 * 4]
+            x += F.interpolate(x00, scale_factor=4, mode="nearest")
+            return x
+        elif tile_mode == 1:  # 对长边减半
+            if w0 >= h0:
+                crop_size_w = ((w0 - 1) // 4 * 4 + 4) // 2  # 减半后能被2整除，所以要先被4整除
+                crop_size_h = (h0 - 1) // 2 * 2 + 2  # 能被2整除
+            else:
+                crop_size_h = ((h0 - 1) // 4 * 4 + 4) // 2  # 减半后能被2整除，所以要先被4整除
+                crop_size_w = (w0 - 1) // 2 * 2 + 2  # 能被2整除
+            crop_size = (crop_size_h, crop_size_w)  # 6.6G
+        elif tile_mode == 2:  # hw都减半
+            crop_size = (
+                ((h0 - 1) // 4 * 4 + 4) // 2,
+                ((w0 - 1) // 4 * 4 + 4) // 2,
+            )  # 5.6G
+        elif tile_mode == 3:  # hw都三分之一
+            crop_size = (
+                ((h0 - 1) // 6 * 6 + 6) // 3,
+                ((w0 - 1) // 6 * 6 + 6) // 3,
+            )  # 4.1G
+        elif tile_mode == 4:  # hw都四分之一
+            crop_size = (
+                ((h0 - 1) // 8 * 8 + 8) // 4,
+                ((w0 - 1) // 8 * 8 + 8) // 4,
+            )  # 3.7G
+        ph = ((h0 - 1) // crop_size[0] + 1) * crop_size[0]
+        pw = ((w0 - 1) // crop_size[1] + 1) * crop_size[1]
+        x = F.pad(x, (19, 19 + pw - w0, 19, 19 + ph - h0), "reflect")
+        n, c, h, w = x.shape
+        se_mean0 = torch.zeros((n, 64, 1, 1)).to(x.device)
+        if "Half" in x.type():
+            se_mean0 = se_mean0.half()
+        n_patch = 0
+        tmp_dict = {}
+        opt_res_dict = {}
+        for i in range(0, h - 38, crop_size[0]):
+            tmp_dict[i] = {}
+            for j in range(0, w - 38, crop_size[1]):
+                x_crop = x[:, :, i : i + crop_size[0] + 38, j : j + crop_size[1] + 38]
+                n, c1, h1, w1 = x_crop.shape
+                tmp0, x_crop = self.unet1.forward_a(x_crop)
+                if "Half" in x.type():  # torch.HalfTensor/torch.cuda.HalfTensor
+                    tmp_se_mean = torch.mean(
+                        x_crop.float(), dim=(2, 3), keepdim=True
+                    ).half()
+                else:
+                    tmp_se_mean = torch.mean(x_crop, dim=(2, 3), keepdim=True)
+                se_mean0 += tmp_se_mean
+                n_patch += 1
+                tmp_dict[i][j] = (tmp0, x_crop)
+        se_mean0 /= n_patch
+        se_mean1 = torch.zeros((n, 128, 1, 1)).to(x.device)  # 64#128#128#64
+        if "Half" in x.type():
+            se_mean1 = se_mean1.half()
+        for i in range(0, h - 38, crop_size[0]):
+            for j in range(0, w - 38, crop_size[1]):
+                tmp0, x_crop = tmp_dict[i][j]
+                x_crop = self.unet1.conv2.seblock.forward_mean(x_crop, se_mean0)
+                opt_unet1 = self.unet1.forward_b(tmp0, x_crop)
+                tmp_x1, tmp_x2 = self.unet2.forward_a(opt_unet1)
+                if "Half" in x.type():  # torch.HalfTensor/torch.cuda.HalfTensor
+                    tmp_se_mean = torch.mean(
+                        tmp_x2.float(), dim=(2, 3), keepdim=True
+                    ).half()
+                else:
+                    tmp_se_mean = torch.mean(tmp_x2, dim=(2, 3), keepdim=True)
+                se_mean1 += tmp_se_mean
+                tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x2)
+        se_mean1 /= n_patch
+        se_mean0 = torch.zeros((n, 128, 1, 1)).to(x.device)  # 64#128#128#64
+        if "Half" in x.type():
+            se_mean0 = se_mean0.half()
+        for i in range(0, h - 38, crop_size[0]):
+            for j in range(0, w - 38, crop_size[1]):
+                opt_unet1, tmp_x1, tmp_x2 = tmp_dict[i][j]
+                tmp_x2 = self.unet2.conv2.seblock.forward_mean(tmp_x2, se_mean1)
+                tmp_x3 = self.unet2.forward_b(tmp_x2)
+                if "Half" in x.type():  # torch.HalfTensor/torch.cuda.HalfTensor
+                    tmp_se_mean = torch.mean(
+                        tmp_x3.float(), dim=(2, 3), keepdim=True
+                    ).half()
+                else:
+                    tmp_se_mean = torch.mean(tmp_x3, dim=(2, 3), keepdim=True)
+                se_mean0 += tmp_se_mean
+                tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x2, tmp_x3)
+        se_mean0 /= n_patch
+        se_mean1 = torch.zeros((n, 64, 1, 1)).to(x.device)  # 64#128#128#64
+        if "Half" in x.type():
+            se_mean1 = se_mean1.half()
+        for i in range(0, h - 38, crop_size[0]):
+            for j in range(0, w - 38, crop_size[1]):
+                opt_unet1, tmp_x1, tmp_x2, tmp_x3 = tmp_dict[i][j]
+                tmp_x3 = self.unet2.conv3.seblock.forward_mean(tmp_x3, se_mean0)
+                tmp_x4 = self.unet2.forward_c(tmp_x2, tmp_x3)
+                if "Half" in x.type():  # torch.HalfTensor/torch.cuda.HalfTensor
+                    tmp_se_mean = torch.mean(
+                        tmp_x4.float(), dim=(2, 3), keepdim=True
+                    ).half()
+                else:
+                    tmp_se_mean = torch.mean(tmp_x4, dim=(2, 3), keepdim=True)
+                se_mean1 += tmp_se_mean
+                tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x4)
+        se_mean1 /= n_patch
+        for i in range(0, h - 38, crop_size[0]):
+            opt_res_dict[i] = {}
+            for j in range(0, w - 38, crop_size[1]):
+                opt_unet1, tmp_x1, tmp_x4 = tmp_dict[i][j]
+                tmp_x4 = self.unet2.conv4.seblock.forward_mean(tmp_x4, se_mean1)
+                x0 = self.unet2.forward_d(tmp_x1, tmp_x4)
+                x1 = F.pad(opt_unet1, (-20, -20, -20, -20))
+                x_crop = torch.add(x0, x1)  # x0是unet2的最终输出
+                x_crop = self.conv_final(x_crop)
+                x_crop = F.pad(x_crop, (-1, -1, -1, -1))
+                x_crop = self.ps(x_crop)
+                opt_res_dict[i][j] = x_crop
+        del tmp_dict
+        torch.cuda.empty_cache()
+        res = torch.zeros((n, c, h * 4 - 152, w * 4 - 152)).to(x.device)
+        if "Half" in x.type():
+            res = res.half()
+        for i in range(0, h - 38, crop_size[0]):
+            for j in range(0, w - 38, crop_size[1]):
+                # print(opt_res_dict[i][j].shape,res[:, :, i * 4:i * 4 + h1 * 4 - 144, j * 4:j * 4 + w1 * 4 - 144].shape)
+                res[
+                    :, :, i * 4 : i * 4 + h1 * 4 - 152, j * 4 : j * 4 + w1 * 4 - 152
+                ] = opt_res_dict[i][j]
+        del opt_res_dict
+        torch.cuda.empty_cache()
+        if w0 != pw or h0 != ph:
+            res = res[:, :, : h0 * 4, : w0 * 4]
+        res += F.interpolate(x00, scale_factor=4, mode="nearest")
+        return res  #
+
+
+models: Dict[str, Type[nn.Module]] = {
+    obj.__name__: obj
+    for obj in globals().values()
+    if isinstance(obj, type) and issubclass(obj, nn.Module)
+}
+
+
+class RealWaifuUpScaler:
+    def __init__(self, scale: int, weight_path: str, half: bool, device: str):
+        weight = torch.load(weight_path, map_location=device)
+        self.model = models[f"UpCunet{scale}x"]()
+
+        if half == True:
+            self.model = self.model.half().to(device)
+        else:
+            self.model = self.model.to(device)
+
+        self.model.load_state_dict(weight, strict=True)
+        self.model.eval()
+
+        self.half = half
+        self.device = device
+
+    def np2tensor(self, np_frame):
+        if self.half == False:
+            return (
+                torch.from_numpy(np.transpose(np_frame, (2, 0, 1)))
+                .unsqueeze(0)
+                .to(self.device)
+                .float()
+                / 255
+            )
+        else:
+            return (
+                torch.from_numpy(np.transpose(np_frame, (2, 0, 1)))
+                .unsqueeze(0)
+                .to(self.device)
+                .half()
+                / 255
+            )
+
+    def tensor2np(self, tensor):
+        if self.half == False:
+            return np.transpose(
+                (tensor.data.squeeze() * 255.0)
+                .round()
+                .clamp_(0, 255)
+                .byte()
+                .cpu()
+                .numpy(),
+                (1, 2, 0),
+            )
+        else:
+            return np.transpose(
+                (tensor.data.squeeze().float() * 255.0)
+                .round()
+                .clamp_(0, 255)
+                .byte()
+                .cpu()
+                .numpy(),
+                (1, 2, 0),
+            )
+
+    def __call__(self, frame, tile_mode):
+        with torch.no_grad():
+            tensor = self.np2tensor(frame)
+            result = self.tensor2np(self.model(tensor, tile_mode))
+        return result
+
+
+input_image = inputs.File(label="Input image")
+half_precision = inputs.Checkbox(
+    label="Half precision (NOT work for CPU)", default=False
+)
+model_weight = inputs.Dropdown(sorted(AVALIABLE_WEIGHTS), label="Choice model weight")
+tile_mode = inputs.Radio([mode.name for mode in TileMode], label="Output tile mode")
+
+output_image = outputs.Image(label="Output image preview")
+output_file = outputs.File(label="Output image file")
+
+
+def main(file: IO[bytes], half: bool, weight: str, tile: str):
+    scale = next(mode.value for mode in ScaleMode if weight.startswith(mode.name))
+    upscaler = RealWaifuUpScaler(
+        scale, weight_path=str(AVALIABLE_WEIGHTS[weight]), half=half, device=DEVICE
+    )
+
+    frame = cv2.imread(file.name)
+    result = upscaler(frame[:, :, [2, 1, 0]], TileMode[tile])
+
+    _, ext = os.path.splitext(file.name)
+    tempfile = mktemp(suffix=ext)
+    cv2.imwrite(tempfile, result)
+    return result, tempfile
+
+
+interface = Interface(
+    main,
+    inputs=[input_image, half_precision, model_weight, tile_mode],
+    outputs=[output_image, output_file],
+)
+interface.launch()