Update app.py

app.py

@@ -19,17 +19,60 @@ class Generator(nn.Module):
     def __init__(self, nz=100, ngf=64, nt=768, nc=3):
         super().__init__()
         self.layer1 = nn.Sequential(
-            nn.ConvTranspose2d(nz + nt, ngf * 8, 4, 1, 0, bias=False),
-            nn.BatchNorm2d(ngf * 8)
+            nn.ConvTranspose2d(nz+nt, ngf*8, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(ngf*8)
         )
         self.layer2 = nn.Sequential(
-            nn.Conv2d(ngf * 8, ngf * 2, 1, 1),
+            nn.Conv2d(ngf*8, ngf*2, 1, 1),
             nn.Dropout2d(inplace=True),
-            nn.BatchNorm2d(ngf * 2),
+            nn.BatchNorm2d(ngf*2),
+            nn.ReLU(True)
+        )
+        self.layer3 = nn.Sequential(
+            nn.Conv2d(ngf*2, ngf*2, 3, 1, 1),
+            nn.Dropout2d(inplace=True),
+            nn.BatchNorm2d(ngf*2),
+            nn.ReLU(True)
+        )
+        self.layer4 = nn.Sequential(
+            nn.Conv2d(ngf*2, ngf*8, 3, 1, 1),
+            nn.Dropout2d(inplace=True),
+            nn.BatchNorm2d(ngf*8),
+            nn.ReLU(True)
+        )
+        self.layer5 = nn.Sequential(
+            nn.ConvTranspose2d(ngf*8, ngf*4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf*4),
+            nn.ReLU(True)
+        )
+        self.layer6 = nn.Sequential(
+            nn.Conv2d(ngf*4, ngf, 1, 1),
+            nn.Dropout2d(inplace=True),
+            nn.BatchNorm2d(ngf),
+            nn.ReLU(True)
+        )
+        self.layer7 = nn.Sequential(
+            nn.Conv2d(ngf, ngf, 3, 1, 1),
+            nn.Dropout2d(inplace=True),
+            nn.BatchNorm2d(ngf),
+            nn.ReLU(True)
+        )
+        self.layer8 = nn.Sequential(
+            nn.Conv2d(ngf, ngf*4, 3, 1, 1),
+            nn.Dropout2d(inplace=True),
+            nn.BatchNorm2d(ngf*4),
+            nn.ReLU(True)
+        )
+        self.layer9 = nn.Sequential(
+            nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf*2),
+            nn.ReLU(True)
+        )
+        self.layer10 = nn.Sequential(
+            nn.ConvTranspose2d(ngf*2, ngf, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf),
             nn.ReLU(True)
         )
-        # Add other layers as needed...
-
         self.layer11 = nn.Sequential(
             nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
             nn.Tanh()
@@ -39,10 +82,19 @@ class Generator(nn.Module):
         x = torch.cat([noise, encoded_text], dim=1)
         x = self.layer1(x)
         x = self.layer2(x)
-        # Pass through other layers...
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.layer5(x)
+        x = self.layer6(x)
+        x = self.layer7(x)
+        x = self.layer8(x)
+        x = self.layer9(x)
+        x = self.layer10(x)
         x = self.layer11(x)
         return x
 
+
+
 # Load the model and tokenizer
 model_path = "checkpoint.pth"  # Adjust as necessary
 tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')