fix avantgan

app.py

@@ -39,7 +39,6 @@ with gr.Blocks(css="style.css") as demo:
                     run_button = gr.Button()
                 with gr.Column():
                     result = gr.Image(label="Result", elem_id="result", width=300, height=300)
-                    print("RESULT", result, type(result), result.__dict__)
 
         with gr.TabItem("Sample Images"):
             with gr.Row():

model.py

@@ -27,66 +27,66 @@ dcgan_img_size = 64
 nc = 3
 
 
-class Generator(nn.Module):
-    def __init__(self, ngpu, nz):
-        super(Generator, self).__init__()
-        self.ngpu = ngpu
-        self.main = nn.Sequential(
-            # input is Z, going into a convolution
-            nn.ConvTranspose2d(     nz, ngf * 8, 4, 1, 0, bias=False),
-            nn.BatchNorm2d(ngf * 8),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size. (ngf*8) x 4 x 4
-            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(ngf * 4),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size. (ngf*4) x 8 x 8
-            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(ngf * 2),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size. (ngf*2) x 16 x 16
-            nn.ConvTranspose2d(ngf * 2,     ngf, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(ngf),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size. (ngf) x 32 x 32
-            nn.ConvTranspose2d(    ngf,      nc, 4, 2, 1, bias=False),
-            nn.Tanh()
-            # state size. (nc) x 64 x 64
-        )
-
-    def forward(self, input):
-        return self.main(input)
-
 # class Generator(nn.Module):
-#     def __init__(self, n_gen_feats, n_gpu, z_dim, n_channels):
+#     def __init__(self, ngpu, nz):
 #         super(Generator, self).__init__()
-#         self.n_gpu = n_gpu
+#         self.ngpu = ngpu
 #         self.main = nn.Sequential(
 #             # input is Z, going into a convolution
-#             nn.ConvTranspose2d(z_dim, n_gen_feats * 8, 4, 1, 0, bias=False),
-#             nn.BatchNorm2d(n_gen_feats * 8),
+#             nn.ConvTranspose2d(     nz, ngf * 8, 4, 1, 0, bias=False),
+#             nn.BatchNorm2d(ngf * 8),
 #             nn.LeakyReLU(0.2, inplace=True),
-#             # state size. (n_gen_feats*8) x 4 x 4
-#             nn.ConvTranspose2d(n_gen_feats * 8, n_gen_feats * 4, 4, 2, 1, bias=False),
-#             nn.BatchNorm2d(n_gen_feats * 4),
+#             # state size. (ngf*8) x 4 x 4
+#             nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
+#             nn.BatchNorm2d(ngf * 4),
 #             nn.LeakyReLU(0.2, inplace=True),
-#             # state size. (n_gen_feats*4) x 8 x 8
-#             nn.ConvTranspose2d(n_gen_feats * 4, n_gen_feats * 2, 4, 2, 1, bias=False),
-#             nn.BatchNorm2d(n_gen_feats * 2),
+#             # state size. (ngf*4) x 8 x 8
+#             nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
+#             nn.BatchNorm2d(ngf * 2),
 #             nn.LeakyReLU(0.2, inplace=True),
-#             # state size. (n_gen_feats*2) x 16 x 16
-#             nn.ConvTranspose2d(n_gen_feats * 2,     n_gen_feats, 4, 2, 1, bias=False),
-#             nn.BatchNorm2d(n_gen_feats),
+#             # state size. (ngf*2) x 16 x 16
+#             nn.ConvTranspose2d(ngf * 2,     ngf, 4, 2, 1, bias=False),
+#             nn.BatchNorm2d(ngf),
 #             nn.LeakyReLU(0.2, inplace=True),
-#             # state size. (n_gen_feats) x 32 x 32
-#             nn.ConvTranspose2d(n_gen_feats, n_channels, 4, 2, 1, bias=False),
+#             # state size. (ngf) x 32 x 32
+#             nn.ConvTranspose2d(    ngf,      nc, 4, 2, 1, bias=False),
 #             nn.Tanh()
-#             # state size. (n_channels) x 64 x 64
+#             # state size. (nc) x 64 x 64
 #         )
 
 #     def forward(self, input):
 #         return self.main(input)
 
+class Generator(nn.Module):
+    def __init__(self, n_gen_feats, n_gpu, z_dim, n_channels):
+        super(Generator, self).__init__()
+        self.n_gpu = n_gpu
+        self.main = nn.Sequential(
+            # input is Z, going into a convolution
+            nn.ConvTranspose2d(z_dim, n_gen_feats * 8, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(n_gen_feats * 8),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (n_gen_feats*8) x 4 x 4
+            nn.ConvTranspose2d(n_gen_feats * 8, n_gen_feats * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(n_gen_feats * 4),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (n_gen_feats*4) x 8 x 8
+            nn.ConvTranspose2d(n_gen_feats * 4, n_gen_feats * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(n_gen_feats * 2),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (n_gen_feats*2) x 16 x 16
+            nn.ConvTranspose2d(n_gen_feats * 2,     n_gen_feats, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(n_gen_feats),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (n_gen_feats) x 32 x 32
+            nn.ConvTranspose2d(n_gen_feats, n_channels, 4, 2, 1, bias=False),
+            nn.Tanh()
+            # state size. (n_channels) x 64 x 64
+        )
+
+    def forward(self, input):
+        return self.main(input)
+
 class Model:
     MODEL_DICT = {
         "stylegan3-abstract": {"name": "abstract-560eps.pkl", "repo": "avantStyleGAN3"},
@@ -109,9 +109,9 @@ class Model:
                 model = pickle.load(f)["G_ema"]
         else:
             # todo (elle): don't hardcode the config
-            # model = Generator(dcgan_gen_feats, 1, dcgan_z_dim, 3)
-            print("WAS HERE")
-            model = Generator(0, 100)
+            model = Generator(dcgan_gen_feats, 1, dcgan_z_dim, 3)
+            # model = Generator(0, 100)
+            model.load_state_dict(torch.load(path, map_location=self.device))
 
         model.eval()
         model.to(self.device)
@@ -149,10 +149,6 @@ class Model:
         tensor = (tensor.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
         return tensor.cpu().numpy()
 
-    def dcgan_postprocess(self, tensor: torch.Tensor) -> np.ndarray:
-        tensor = (tensor.permute(0, 2, 3, 1)).clamp(0, 255).to(torch.uint8)
-        return tensor.cpu().numpy()
-
     def set_transform(self, tx: float = 0, ty: float = 0, angle: float = 0) -> None:
         mat = self.make_transform((tx, ty), angle)
         mat = np.linalg.inv(mat)
@@ -173,20 +169,18 @@ class Model:
         return out[0]
 
     def dcgan_generate_image(self, seed: int) -> np.ndarray:
-        dcgan_img_size = 64
-        dcgan_z_dim = 100
+        torch.manual_seed(seed)
+        if self.device == 'cuda':
+            torch.cuda.manual_seed(seed)
 
         with torch.no_grad():
             n_images = 1
             z = torch.randn(n_images, dcgan_z_dim, 1, 1, device=self.device)
             fake_images = self.model(z.to(self.device)).cpu()
             fake_images = fake_images.view(fake_images.size(0), 3, dcgan_img_size, dcgan_img_size)
-        
-            print('fake', fake_images)
-            print(fake_images.min(), fake_images.max())
             # Create a grid of images
             grid = vutils.make_grid(fake_images, normalize=True)
-            print('grid', grid)
+
             # Plot the grid and save it to a buffer
             fig, ax = plt.subplots()
             ax.imshow(grid.permute(1, 2, 0))  # Convert from CHW to HWC for imshow