Fixed upscale

chaiNNer-org · Nov 23, 2023 · 7fbc3ad · 7fbc3ad
1 parent 8152764
commit 7fbc3ad
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/spandrel/architectures/SwiftSRGAN/__init__.py b/src/spandrel/architectures/SwiftSRGAN/__init__.py
@@ -12,7 +12,7 @@ def load(state: StateDict) -> SRModelDescriptor[SwiftSRGAN]:
     in_channels = state["initial.cnn.depthwise.weight"].shape[0]
     num_channels = state["initial.cnn.pointwise.weight"].shape[0]
     num_blocks = get_seq_len(state, "residual")
-    upscale_factor = get_seq_len(state, "upsampler") * 2
+    upscale_factor = 2 ** get_seq_len(state, "upsampler")
 
     model = SwiftSRGAN(
         in_channels=in_channels,

diff --git a/src/spandrel/architectures/SwiftSRGAN/arch/SwiftSRGAN.py b/src/spandrel/architectures/SwiftSRGAN/arch/SwiftSRGAN.py
@@ -1,6 +1,8 @@
 # type: ignore
 # From https://github.com/Koushik0901/Swift-SRGAN/blob/master/swift-srgan/models.py
 
+import math
+
 import torch
 from torch import nn
 
@@ -124,7 +126,7 @@ def __init__(
         self.upsampler = nn.Sequential(
             *[
                 UpsampleBlock(num_channels, scale_factor=2)
-                for _ in range(upscale_factor // 2)
+                for _ in range(int(math.log2(upscale_factor)))
             ]
         )
         self.final_conv = SeperableConv2d(