Update

app.py

@@ -18,11 +18,57 @@ MAX_IMAGE_SIZE = (1024, 1024)
 class RRDBNet(torch.nn.Module):
     def __init__(self, in_nc=3, out_nc=3, nf=64, nb=23, gc=32):
         super(RRDBNet, self).__init__()
-        self.model = self._make_network(in_nc, out_nc, nf, nb, gc)
+        self.conv_first = torch.nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
+        self.RRDB_trunk = torch.nn.ModuleList([RRDB(nf, gc=gc) for _ in range(nb)])
+        self.trunk_conv = torch.nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.upconv1 = torch.nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.upconv2 = torch.nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.HRconv = torch.nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.conv_last = torch.nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
+        self.lrelu = torch.nn.LeakyReLU(negative_slope=0.2, inplace=True)
 
-    def _make_network(self, in_nc, out_nc, nf, nb, gc):
-        # [Original architecture implementation here...]
-        # Full implementation: https://github.com/xinntao/ESRGAN/blob/master/RRDBNet_arch.py
+    def forward(self, x):
+        fea = self.conv_first(x)
+        trunk = fea.clone()
+        for block in self.RRDB_trunk:
+            trunk = block(trunk)
+        trunk = self.trunk_conv(trunk)
+        fea = fea + trunk
+        fea = self.lrelu(self.upconv1(torch.nn.functional.interpolate(fea, scale_factor=2, mode='nearest')))
+        fea = self.lrelu(self.upconv2(torch.nn.functional.interpolate(fea, scale_factor=2, mode='nearest')))
+        out = self.conv_last(self.lrelu(self.HRconv(fea)))
+        return out
+
+class RRDB(torch.nn.Module):
+    def __init__(self, nf, gc=32):
+        super(RRDB, self).__init__()
+        self.RDB1 = ResidualDenseBlock(nf, gc)
+        self.RDB2 = ResidualDenseBlock(nf, gc)
+        self.RDB3 = ResidualDenseBlock(nf, gc)
+
+    def forward(self, x):
+        out = self.RDB1(x)
+        out = self.RDB2(out)
+        out = self.RDB3(out)
+        return out * 0.2 + x
+
+class ResidualDenseBlock(torch.nn.Module):
+    def __init__(self, nf=64, gc=32, bias=True):
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = torch.nn.Conv2d(nf, gc, 3, 1, 1, bias=bias)
+        self.conv2 = torch.nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias)
+        self.conv3 = torch.nn.Conv2d(nf + 2*gc, gc, 3, 1, 1, bias=bias)
+        self.conv4 = torch.nn.Conv2d(nf + 3*gc, gc, 3, 1, 1, bias=bias)
+        self.conv5 = torch.nn.Conv2d(nf + 4*gc, nf, 3, 1, 1, bias=bias)
+        self.lrelu = torch.nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x):
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        return x5 * 0.2 + x
 
 def load_model() -> torch.nn.Module:
     """Download and load ESRGAN model"""