File size: 1,302 Bytes
f7494c3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | import torch
import torch.nn as nn
import torch.nn.functional as F
try:
from huggingface_hub import PyTorchModelHubMixin
except ImportError: # the model still works without the hub helpers
class PyTorchModelHubMixin:
pass
class ResBlock(nn.Module):
def __init__(self, nf):
super().__init__()
self.c1 = nn.Conv2d(nf, nf, 3, 1, 1)
self.c2 = nn.Conv2d(nf, nf, 3, 1, 1)
def forward(self, x):
return x + self.c2(F.relu(self.c1(x), inplace=True))
class SRNet(nn.Module, PyTorchModelHubMixin):
"""x{scale} super-resolution. Input: float RGB (B,3,H,W) in [0,1]. Output: (B,3,H*scale,W*scale), clamp to [0,1]."""
def __init__(self, scale: int = 4, nf: int = 64, nb: int = 16):
super().__init__()
self.scale, self.nf, self.nb = scale, nf, nb
self.head = nn.Conv2d(3, nf, 3, 1, 1)
self.body = nn.Sequential(*[ResBlock(nf) for _ in range(nb)])
self.tail = nn.Conv2d(nf, nf, 3, 1, 1)
self.up = nn.Conv2d(nf, 3 * scale * scale, 3, 1, 1)
self.shuffle = nn.PixelShuffle(scale)
def forward(self, x):
f = self.head(x)
f = f + self.tail(self.body(f))
return self.shuffle(self.up(f)) + F.interpolate(x, scale_factor=self.scale, mode="bicubic", align_corners=False)
|