Update app.py

app.py

@@ -12,26 +12,35 @@ from torchvision.utils import save_image
 
 
 class Generator(nn.Module):
-    def __init__(self, nc=4, nz=100, ngf=64):
+    def __init__(self, num_channels=4, latent_dim=100, hidden_size=64):
         super(Generator, self).__init__()
         self.model = nn.Sequential(
-            nn.ConvTranspose2d(nz, ngf * 4, 3, 1, 0, bias=False),
-            nn.BatchNorm2d(ngf * 4),
+            # input is Z, going into a convolution
+            nn.ConvTranspose2d(latent_dim, hidden_size * 8, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(hidden_size * 8),
             nn.ReLU(True),
-            nn.ConvTranspose2d(ngf * 4, ngf * 2, 3, 2, 1, bias=False),
-            nn.BatchNorm2d(ngf * 2),
+            # state size. (hidden_size*8) x 4 x 4
+            nn.ConvTranspose2d(hidden_size * 8, hidden_size * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(hidden_size * 4),
             nn.ReLU(True),
-            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 0, bias=False),
-            nn.BatchNorm2d(ngf),
+            # state size. (hidden_size*4) x 8 x 8
+            nn.ConvTranspose2d(hidden_size * 4, hidden_size * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(hidden_size * 2),
             nn.ReLU(True),
-            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
-            nn.Tanh(),
+            # state size. (hidden_size*2) x 16 x 16
+            nn.ConvTranspose2d(hidden_size * 2, hidden_size, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(hidden_size),
+            nn.ReLU(True),
+            # state size. (hidden_size) x 32 x 32
+            nn.ConvTranspose2d(hidden_size, num_channels, 4, 2, 1, bias=False),
+            nn.Tanh()
+            # state size. (num_channels) x 64 x 64
         )
 
-    def forward(self, input):
-        output = self.model(input)
-        return output
+    def forward(self, noise):
+        pixel_values = self.model(noise)
 
+        return pixel_values
 
 model = Generator()
 weights_path = hf_hub_download('huggingnft/dooggies', 'pytorch_model.bin')