address #29

README.md

@@ -20,7 +20,7 @@ Please join <a href="https://discord.gg/xBPBXfcFHd"><img alt="Join us on Discord
 
 - All the maintainers at <a href="https://github.com/mlfoundations/open_clip">OpenClip</a>, for their SOTA open sourced contrastive learning text-image models
 
-- <a href="https://github.com/XavierXiao">Xavier</a> for reviewing the discriminator code and pointing out that the scale invariance was not correctly built!
+- <a href="https://github.com/XavierXiao">Xavier</a> for the very helpful code review, and for discussions on how the scale invariance in the discriminator should be built!
 
 - <a href="https://github.com/CerebralSeed">@CerebralSeed</a> for pull requesting the initial sampling code for both the generator and upsampler!
 

gigagan_pytorch/unet_upsampler.py

@@ -434,7 +434,6 @@ def __init__(
         self.final_res_block = block_klass(dim, dim)
 
         self.final_to_rgb = nn.Conv2d(dim, channels, 1)
-        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
 
         # resize mode
 
@@ -518,13 +517,13 @@ def forward(
             if exists(cross_attn):
                 x = cross_attn(x, context = fine_text_tokens, mask = text_mask)
 
-            x = attn(x) + x
+            x = attn(x)
             h.append(x)
 
             x = downsample(x)
 
         x = self.mid_block1(x, conv_mods_iter = conv_mods)
-        x = self.mid_attn(x) + x
+        x = self.mid_attn(x)
         x = self.mid_block2(x, conv_mods_iter = conv_mods)
 
         # rgbs
@@ -563,7 +562,7 @@ def forward(
             if exists(cross_attn):
                 x = cross_attn(x, context = fine_text_tokens, mask = text_mask)
 
-            x = attn(x) + x
+            x = attn(x)
 
             rgb = rgb + to_rgb(x)
             rgbs.append(rgb)