Fix ESRGAN model structure to match checkpoint exact state dict

app.py

@@ -185,53 +185,53 @@ FG_ESRGAN_PATH = CACHE_DIR / "esrgan"
 _fg_esrgan_model = None
 _fg_esrgan_loading = False
 
-def _math_ceil(x): return int(x + 0.999) if x > 0 else 0
-
-def _conv_block(in_nc: int, out_nc: int) -> torch.nn.Sequential:
+def _conv_block(in_nc, out_nc):
     return torch.nn.Sequential(
         torch.nn.Conv2d(in_nc, out_nc, kernel_size=3, padding=1),
         torch.nn.LeakyReLU(negative_slope=0.2, inplace=True),
     )
 
+class _ResidualDenseBlock5C(torch.nn.Module):
+    def __init__(self, nf=64, gc=32):
+        super().__init__()
+        self.conv1 = _conv_block(nf, gc)
+        self.conv2 = _conv_block(nf + gc, gc)
+        self.conv3 = _conv_block(nf + 2 * gc, gc)
+        self.conv4 = _conv_block(nf + 3 * gc, gc)
+        self.conv5 = torch.nn.Sequential(torch.nn.Conv2d(nf + 4 * gc, nf, kernel_size=3, padding=1))
+
+    def forward(self, x):
+        x1 = self.conv1(x)
+        x2 = self.conv2(torch.cat((x, x1), 1))
+        x3 = self.conv3(torch.cat((x, x1, x2), 1))
+        x4 = 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
+
 class _RRDB(torch.nn.Module):
-    def __init__(self, nf: int):
+    def __init__(self, nf):
         super().__init__()
-        gc = 32
-        self.RDB1 = _conv_block(nf, gc)
-        self.conv1_2 = _conv_block(nf + gc, gc)
-        self.conv1_3 = _conv_block(nf + 2 * gc, gc)
-        self.conv1_4 = _conv_block(nf + 3 * gc, gc)
-        self.conv1_5 = torch.nn.Sequential(torch.nn.Conv2d(nf + 4 * gc, nf, kernel_size=3, padding=1))
-        self.RDB2 = _conv_block(nf, gc)
-        self.conv2_2 = _conv_block(nf + gc, gc)
-        self.conv2_3 = _conv_block(nf + 2 * gc, gc)
-        self.conv2_4 = _conv_block(nf + 3 * gc, gc)
-        self.conv2_5 = torch.nn.Sequential(torch.nn.Conv2d(nf + 4 * gc, nf, kernel_size=3, padding=1))
-        self.RDB3 = _conv_block(nf, gc)
-        self.conv3_2 = _conv_block(nf + gc, gc)
-        self.conv3_3 = _conv_block(nf + 2 * gc, gc)
-        self.conv3_4 = _conv_block(nf + 3 * gc, gc)
-        self.conv3_5 = torch.nn.Sequential(torch.nn.Conv2d(nf + 4 * gc, nf, kernel_size=3, padding=1))
-
-    def _dense_block(self, x, conv2, conv3, conv4, conv5):
-        x1 = self.RDB1(x)
-        x2 = conv2(torch.cat((x, x1), 1))
-        x3 = conv3(torch.cat((x, x1, x2), 1))
-        x4 = conv4(torch.cat((x, x1, x2, x3), 1))
-        x5 = conv5(torch.cat((x, x1, x2, x3, x4), 1))
-        return x5 * 0.2
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        out = self._dense_block(x, self.conv1_2, self.conv1_3, self.conv1_4, self.conv1_5)
-        out = self._dense_block(out + x, self.conv2_2, self.conv2_3, self.conv2_4, self.conv2_5)
-        out = self._dense_block(out + (out + x), self.conv3_2, self.conv3_3, self.conv3_4, self.conv3_5)
+        self.RDB1 = _ResidualDenseBlock5C(nf)
+        self.RDB2 = _ResidualDenseBlock5C(nf)
+        self.RDB3 = _ResidualDenseBlock5C(nf)
+
+    def forward(self, x):
+        out = self.RDB1(x)
+        out = self.RDB2(out)
+        out = self.RDB3(out)
         return out * 0.2 + x
 
+class _SkipBlock(torch.nn.Module):
+    def __init__(self, sub):
+        super().__init__()
+        self.sub = sub
+    def forward(self, x): return x + self.sub(x)
+
 class _RRDBNet(torch.nn.Module):
-    def __init__(self, nf: int = 64, nb: int = 23):
+    def __init__(self, in_nc=3, out_nc=3, nf=64, nb=23):
         super().__init__()
         self.model = torch.nn.Sequential(
-            torch.nn.Conv2d(3, nf, kernel_size=3, padding=1),
+            torch.nn.Conv2d(in_nc, nf, kernel_size=3, padding=1),
             _SkipBlock(torch.nn.Sequential(
                 *(_RRDB(nf) for _ in range(nb)),
                 torch.nn.Conv2d(nf, nf, kernel_size=3, padding=1),
@@ -244,18 +244,12 @@ class _RRDBNet(torch.nn.Module):
             torch.nn.LeakyReLU(negative_slope=0.2, inplace=True),
             torch.nn.Conv2d(nf, nf, kernel_size=3, padding=1),
             torch.nn.LeakyReLU(negative_slope=0.2, inplace=True),
-            torch.nn.Conv2d(nf, 3, kernel_size=3, padding=1),
+            torch.nn.Conv2d(nf, out_nc, kernel_size=3, padding=1),
         )
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x):
         return self.model(x)
 
-class _SkipBlock(torch.nn.Module):
-    def __init__(self, sub: torch.nn.Module):
-        super().__init__()
-        self.sub = sub
-    def forward(self, x): return x + self.sub(x)
-
 def _load_fg_esrgan():
     global _fg_esrgan_model, _fg_esrgan_loading
     if _fg_esrgan_model is not None: return _fg_esrgan_model
@@ -267,9 +261,8 @@ def _load_fg_esrgan():
         ckpt = hf_hub_download("philz1337x/upscaler", "4x-UltraSharp.pth", cache_dir=str(FG_ESRGAN_PATH))
         logger.info("Loading ESRGAN...")
         state = torch.load(ckpt, map_location="cpu", weights_only=True)
-        in_nc, out_nc, nf, nb = 3, 3, 64, 23  # defaults for 4x-UltraSharp
-        model = _RRDBNet(nf=nf, nb=nb)
-        model.load_state_dict(state)
+        model = _RRDBNet(in_nc=3, out_nc=3, nf=64, nb=23)
+        model.load_state_dict(state, strict=False)
         model.eval()
         _fg_esrgan_model = model
         logger.info("ESRGAN 4X ready (CPU)")
@@ -304,8 +297,8 @@ def upscale_finegrain(image_bytes: bytes, use_sd_refinement: bool = False) -> tu
     else:
         # Tiled upscale
         stride = tile_size - overlap
-        cols = _math_ceil((w - overlap) / stride) if w > tile_size else 1
-        rows = _math_ceil((h - overlap) / stride) if h > tile_size else 1
+        cols = -(-max(0, w - overlap) // stride) if w > tile_size else 1
+        rows = -(-max(0, h - overlap) // stride) if h > tile_size else 1
         out_arr = np.zeros((h * 4, w * 4, 3), dtype=np.float32)
         weight = np.zeros((h * 4, w * 4, 1), dtype=np.float32)
         for row in range(rows):