Update discriminatorModel.py

discriminatorModel.py

@@ -2,6 +2,21 @@ import torch
 import torch.nn as nn
 
 
+class Embedding(nn.Module):
+    def __init__(self, size_in, size_out):
+        super(Embedding, self).__init__()
+        self.text_embedding = nn.Sequential(
+            nn.Linear(size_in, size_out),
+            nn.BatchNorm1d(1),
+            nn.LeakyReLU(0.2, inplace=True)
+        )
+
+    def forward(self, x, text):
+        embed_out = self.text_embedding(text)
+        embed_out_resize = embed_out.repeat(4, 1, 4, 1).permute(1, 3, 0, 2)  # Resize to match the discriminator input size
+        out = torch.cat([x, embed_out_resize], 1)  # Concatenate text embedding with the input feature map
+        return out
+
 # The Discriminator model
 class Discriminator(nn.Module):
     def __init__(self, channels, embed_dim=1024, embed_out_dim=128):
@@ -23,21 +38,6 @@ class Discriminator(nn.Module):
             nn.Conv2d(512 + self.embed_out_dim, 1, 4, 1, 0, bias=False), nn.Sigmoid()
         )
         
-    class Embedding(nn.Module):
-        def __init__(self, size_in, size_out):
-            super(Embedding, self).__init__()
-            self.text_embedding = nn.Sequential(
-                nn.Linear(size_in, size_out),
-                nn.BatchNorm1d(1),
-                nn.LeakyReLU(0.2, inplace=True)
-            )
-    
-        def forward(self, x, text):
-            embed_out = self.text_embedding(text)
-            embed_out_resize = embed_out.repeat(4, 1, 4, 1).permute(1, 3, 0, 2)  # Resize to match the discriminator input size
-            out = torch.cat([x, embed_out_resize], 1)  # Concatenate text embedding with the input feature map
-            return out
-        
     def _create_layer(self, size_in, size_out, kernel_size=4, stride=2, padding=1, normalize=True):
         layers = [nn.Conv2d(size_in, size_out, kernel_size=kernel_size, stride=stride, padding=padding)]
         if normalize: