294 lines
8.2 KiB
Python
294 lines
8.2 KiB
Python
import torch
|
|
|
|
from torch import Tensor
|
|
|
|
from torch.nn import (
|
|
Module,
|
|
Sequential,
|
|
Conv2d,
|
|
Upsample,
|
|
PixelShuffle,
|
|
SiLU,
|
|
BatchNorm2d,
|
|
MaxPool2d,
|
|
AdaptiveAvgPool2d,
|
|
Identity,
|
|
Flatten,
|
|
Linear,
|
|
)
|
|
|
|
from torch.nn.utils.parametrizations import weight_norm
|
|
from torch.nn.utils.parametrize import remove_parametrizations
|
|
|
|
from huggingface_hub import PyTorchModelHubMixin
|
|
|
|
|
|
class SuperCool(Module, PyTorchModelHubMixin):
|
|
"""
|
|
A fast single-image super-resolution model with a deep low-resolution encoder network
|
|
and high-resolution sub-pixel convolutional decoder head with global residual pathway.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
base_upscaler: str,
|
|
upscale_ratio: int,
|
|
num_channels: int,
|
|
hidden_ratio: int,
|
|
num_layers: int,
|
|
):
|
|
super().__init__()
|
|
|
|
if base_upscaler not in {"bilinear", "bicubic"}:
|
|
raise ValueError(
|
|
f"Base upscaler must be bilinear or bicubic, {base_upscaler} given."
|
|
)
|
|
|
|
if upscale_ratio not in {2, 4, 8}:
|
|
raise ValueError(
|
|
f"Upscale ratio must be either 2, 4, or 8, {upscale_ratio} given."
|
|
)
|
|
|
|
if num_channels < 1:
|
|
raise ValueError(
|
|
f"Num channels must be greater than 0, {num_channels} given."
|
|
)
|
|
|
|
if hidden_ratio not in {1, 2, 4}:
|
|
raise ValueError(
|
|
f"Hidden ratio must be either 1, 2, or 4, {hidden_ratio} given."
|
|
)
|
|
|
|
if num_layers < 1:
|
|
raise ValueError(f"Num layers must be greater than 0, {num_layers} given.")
|
|
|
|
self.input = weight_norm(Conv2d(3, num_channels, kernel_size=5, padding=2))
|
|
|
|
self.skip = Upsample(scale_factor=upscale_ratio, mode=base_upscaler)
|
|
|
|
self.encoder = Sequential(
|
|
*[EncoderBlock(num_channels, hidden_ratio) for _ in range(num_layers)]
|
|
)
|
|
|
|
self.decoder = SubpixelConv2d(
|
|
num_channels, upscale_ratio, kernel_size=3, padding=1
|
|
)
|
|
|
|
self.shuffle = PixelShuffle(upscale_ratio)
|
|
|
|
@property
|
|
def num_trainable_params(self) -> int:
|
|
return sum(param.numel() for param in self.parameters() if param.requires_grad)
|
|
|
|
def remove_weight_norms(self) -> None:
|
|
for module in self.modules():
|
|
if hasattr(module, "parametrizations"):
|
|
params = [name for name in module.parametrizations.keys()]
|
|
|
|
for name in params:
|
|
remove_parametrizations(module, name, leave_parametrized=True)
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
z = self.input(x)
|
|
|
|
z = self.encoder(z)
|
|
|
|
z = self.decoder(z)
|
|
z = self.shuffle(z)
|
|
|
|
s = self.skip(x)
|
|
|
|
z += s # Global residual connection
|
|
|
|
z = torch.clamp(z, 0, 1)
|
|
|
|
return z
|
|
|
|
|
|
class EncoderBlock(Module):
|
|
"""A low-resolution encoder block with {num_channels} feature maps and wide activations."""
|
|
|
|
def __init__(self, num_channels: int, hidden_ratio: int):
|
|
super().__init__()
|
|
|
|
if hidden_ratio not in {1, 2, 4}:
|
|
raise ValueError(
|
|
f"Hidden ratio must be either 1, 2, or 4, {hidden_ratio} given."
|
|
)
|
|
|
|
hidden_channels = num_channels * hidden_ratio
|
|
|
|
conv1 = Conv2d(num_channels, hidden_channels, kernel_size=3, padding=1)
|
|
conv2 = Conv2d(hidden_channels, num_channels, kernel_size=3, padding=1)
|
|
|
|
self.conv1 = weight_norm(conv1)
|
|
self.conv2 = weight_norm(conv2)
|
|
|
|
self.silu = SiLU()
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
z = self.conv1(x)
|
|
z = self.silu(z)
|
|
z = self.conv2(z)
|
|
|
|
z += x # Local residual connection
|
|
|
|
return z
|
|
|
|
|
|
class SubpixelConv2d(Module):
|
|
"""A sub-pixel (1 / upscale_ratio) convolution layer with weight normalization."""
|
|
|
|
def __init__(
|
|
self, num_channels: int, upscale_ratio: int, kernel_size: int, padding: int
|
|
):
|
|
super().__init__()
|
|
|
|
channels_out = 3 * upscale_ratio**2
|
|
|
|
conv = Conv2d(
|
|
num_channels, channels_out, kernel_size=kernel_size, padding=padding
|
|
)
|
|
|
|
self.conv = weight_norm(conv)
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
return self.conv(x)
|
|
|
|
|
|
class Bouncer(Module):
|
|
"""A residual-style discriminator network for adversarial training."""
|
|
|
|
def __init__(self, model_size: str):
|
|
super().__init__()
|
|
|
|
if model_size not in {"small", "medium", "large"}:
|
|
raise ValueError(
|
|
f"Model size must be small, medium, or large, {model_size} given."
|
|
)
|
|
|
|
num_primary_layers = 3
|
|
num_quaternary_layers = 3
|
|
|
|
match model_size:
|
|
case "small":
|
|
num_secondary_layers = 4
|
|
num_tertiary_layers = 6
|
|
case "medium":
|
|
num_secondary_layers = 4
|
|
num_tertiary_layers = 23
|
|
case "large":
|
|
num_secondary_layers = 8
|
|
num_tertiary_layers = 36
|
|
|
|
self.input = Sequential(
|
|
Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False),
|
|
BatchNorm2d(64),
|
|
SiLU(),
|
|
MaxPool2d(3, stride=2, padding=1),
|
|
)
|
|
|
|
self.detector = Sequential(DetectorBlock(64, 64, 256))
|
|
|
|
self.detector.extend(
|
|
[DetectorBlock(256, 64, 256) for _ in range(num_primary_layers - 1)]
|
|
)
|
|
|
|
self.detector.append(DetectorBlock(256, 128, 512, stride=2))
|
|
|
|
self.detector.extend(
|
|
[DetectorBlock(512, 128, 512) for _ in range(num_secondary_layers - 1)]
|
|
)
|
|
|
|
self.detector.append(DetectorBlock(512, 256, 1024, stride=2))
|
|
|
|
self.detector.extend(
|
|
[DetectorBlock(1024, 256, 1024) for _ in range(num_tertiary_layers - 1)]
|
|
)
|
|
|
|
self.detector.append(DetectorBlock(1024, 512, 2048, stride=2))
|
|
|
|
self.detector.extend(
|
|
[DetectorBlock(2048, 512, 2048) for _ in range(num_quaternary_layers - 1)]
|
|
)
|
|
|
|
self.pool = AdaptiveAvgPool2d(1)
|
|
|
|
self.flatten = Flatten(start_dim=1)
|
|
|
|
self.linear = Linear(2048, 1)
|
|
|
|
@property
|
|
def num_trainable_params(self) -> int:
|
|
return sum(param.numel() for param in self.parameters() if param.requires_grad)
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
x = self.input(x)
|
|
|
|
x = self.detector(x)
|
|
|
|
x = self.pool(x)
|
|
x = self.flatten(x)
|
|
x = self.linear(x)
|
|
|
|
return x
|
|
|
|
|
|
class DetectorBlock(Module):
|
|
"""A residual bottleneck block with 3x3 convolutions and batch normalization."""
|
|
|
|
def __init__(
|
|
self, channels_in: int, hidden_channels: int, channels_out: int, stride: int = 1
|
|
):
|
|
super().__init__()
|
|
|
|
self.residual = Sequential(
|
|
Conv2d(
|
|
in_channels=channels_in,
|
|
out_channels=hidden_channels,
|
|
kernel_size=1,
|
|
bias=False,
|
|
),
|
|
BatchNorm2d(hidden_channels),
|
|
SiLU(),
|
|
Conv2d(
|
|
in_channels=hidden_channels,
|
|
out_channels=hidden_channels,
|
|
kernel_size=3,
|
|
stride=stride,
|
|
padding=1,
|
|
),
|
|
SiLU(),
|
|
Conv2d(
|
|
in_channels=hidden_channels,
|
|
out_channels=channels_out,
|
|
kernel_size=1,
|
|
),
|
|
)
|
|
|
|
if channels_in == channels_out:
|
|
skip = Sequential(Identity())
|
|
else:
|
|
skip = Sequential(
|
|
Conv2d(
|
|
in_channels=channels_in,
|
|
out_channels=channels_out,
|
|
kernel_size=1,
|
|
stride=stride,
|
|
bias=False,
|
|
)
|
|
)
|
|
|
|
self.skip = skip
|
|
self.silu = SiLU()
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
z = self.residual(x)
|
|
s = self.skip(x)
|
|
|
|
z += s
|
|
|
|
z = self.silu(z)
|
|
|
|
return z
|