fix rgbs outputted by unet upsampler to discriminator, and validate c…

…orrectly on GigaGAN instantiation, thanks to @XavierXiao

gigagan_pytorch/gigagan_pytorch.py

@@ -1614,6 +1614,7 @@ def forward(
         x = images
 
         image_size = (self.image_size, self.image_size)
+
         assert x.shape[-2:] == image_size
 
         batch = x.shape[0]
@@ -1834,6 +1835,11 @@ def __init__(
         self.D = discriminator
         self.VD = vision_aided_discriminator
 
+        # validate multiscale input resolutions
+
+        if train_upsampler:
+            assert is_empty(set(discriminator.multiscale_input_resolutions) - set(generator.allowable_rgb_resolutions)), f'only multiscale input resolutions of {generator.allowable_rgb_resolutions} is allowed based on the unet input and output image size'
+
         # ema
 
         self.has_ema_generator = False

gigagan_pytorch/unet_upsampler.py

@@ -449,6 +449,13 @@ def __init__(
         self.style_to_conv_modulations = nn.Linear(style_network.dim, sum(style_embed_split_dims))
         self.style_embed_split_dims = style_embed_split_dims
 
+    @property
+    def allowable_rgb_resolutions(self):
+        input_res_base = int(log2(self.input_image_size))
+        output_res_base = int(log2(self.image_size))
+        allowed_rgb_res_base = list(range(input_res_base + 1, output_res_base))
+        return [*map(lambda p: 2 ** p, allowed_rgb_res_base)]
+
     @property
     def device(self):
         return next(self.parameters()).device
@@ -584,7 +591,7 @@ def forward(
 
         # only keep those rgbs whose feature map is greater than the input image to be upsampled
 
-        rgbs = list(filter(lambda t: t.shape[-1] <= shape[-1], rgbs))
+        rgbs = list(filter(lambda t: t.shape[-1] > shape[-1], rgbs))
 
         if not replace_rgb_with_input_lowres_image:
             return rgb, rgbs

gigagan_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.2.5'
+__version__ = '0.2.6'