Update models/gan_generator.py

models/gan_generator.py

@@ -1,6 +1,6 @@
 """
-GAN Generator model for creating new Kolam designs.
-Uses a deep convolutional architecture to generate high-quality Kolam patterns.
+Enhanced GAN Generator for Kolam designs.
+Adds style-conditioning and more diverse outputs.
 """
 
 import torch
@@ -9,10 +9,7 @@ import torch.nn.functional as F
 
 
 class KolamGenerator(nn.Module):
-    """
-    Generator network for creating Kolam designs.
-    Takes random noise and optional style features as input.
-    """
+    """Base generator network for Kolam designs."""
     
     def __init__(self, noise_dim=100, feature_dim=128, output_channels=1, image_size=64):
         super(KolamGenerator, self).__init__()
@@ -20,174 +17,104 @@ class KolamGenerator(nn.Module):
         self.noise_dim = noise_dim
         self.feature_dim = feature_dim
         self.image_size = image_size
-        
-        # Calculate the starting size after upsampling
-        # Assuming we start from 4x4 and upsample to 64x64
         self.start_size = 4
-        self.num_upsamples = int(torch.log2(torch.tensor(image_size / self.start_size)).item())
+        self.num_upsamples = int(torch.log2(torch.tensor(image_size // self.start_size)).item())
         
-        # Input projection layer
+        # Input projection
         self.input_projection = nn.Linear(noise_dim + feature_dim, 256 * self.start_size * self.start_size)
         
-        # Upsampling layers
-        self.upsample_layers = nn.ModuleList()
-        self.conv_layers = nn.ModuleList()
-        self.bn_layers = nn.ModuleList()
-        
-        # Build upsampling blocks
+        # Upsampling blocks
+        self.upsample_layers, self.conv_layers, self.bn_layers = nn.ModuleList(), nn.ModuleList(), nn.ModuleList()
         in_channels = 256
         for i in range(self.num_upsamples):
             out_channels = in_channels // 2 if i < self.num_upsamples - 1 else 64
-            
-            self.upsample_layers.append(nn.ConvTranspose2d(in_channels, out_channels, 
-                                                          kernel_size=4, stride=2, padding=1))
-            self.conv_layers.append(nn.Conv2d(out_channels, out_channels, 
-                                            kernel_size=3, padding=1))
+            self.upsample_layers.append(nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1))
+            self.conv_layers.append(nn.Conv2d(out_channels, out_channels, 3, padding=1))
             self.bn_layers.append(nn.BatchNorm2d(out_channels))
-            
             in_channels = out_channels
         
-        # Final output layer
-        self.final_conv = nn.Conv2d(64, output_channels, kernel_size=3, padding=1)
-        
+        # Final output
+        self.final_conv = nn.Conv2d(64, output_channels, 3, padding=1)
+    
     def forward(self, noise, features=None):
-        """
-        Generate Kolam images from noise and optional features.
-        
-        Args:
-            noise: Random noise tensor of shape (batch_size, noise_dim)
-            features: Optional feature tensor of shape (batch_size, feature_dim)
-        
-        Returns:
-            Generated images of shape (batch_size, 1, image_size, image_size)
-        """
         batch_size = noise.size(0)
-        
-        # Combine noise and features
         if features is not None:
             x = torch.cat([noise, features], dim=1)
         else:
-            # If no features provided, use zero features
             zero_features = torch.zeros(batch_size, self.feature_dim, device=noise.device)
             x = torch.cat([noise, zero_features], dim=1)
         
-        # Project to initial feature map
         x = self.input_projection(x)
         x = x.view(batch_size, 256, self.start_size, self.start_size)
         
-        # Upsample and refine
         for i in range(self.num_upsamples):
             x = self.upsample_layers[i](x)
             x = self.bn_layers[i](x)
             x = F.relu(x)
-            
             x = self.conv_layers[i](x)
             x = self.bn_layers[i](x)
             x = F.relu(x)
         
-        # Final output
-        x = self.final_conv(x)
-        x = torch.tanh(x)  # Output in range [-1, 1]
-        
-        return x
-    
-    def generate(self, num_samples=1, features=None, device='cpu'):
-        """
-        Generate samples without gradients (for inference).
-        
-        Args:
-            num_samples: Number of samples to generate
-            features: Optional feature tensor
-            device: Device to generate on
-        
-        Returns:
-            Generated images
-        """
-        self.eval()
-        with torch.no_grad():
-            noise = torch.randn(num_samples, self.noise_dim, device=device)
-            return self.forward(noise, features)
+        return torch.tanh(self.final_conv(x))  # [-1, 1]
 
 
 class StyleConditionedGenerator(KolamGenerator):
-    """
-    Style-conditioned generator that can generate Kolam designs
-    in specific styles based on input features.
-    """
+    """Generator with style-conditioning for more variety."""
     
-    def __init__(self, noise_dim=100, feature_dim=128, style_dim=32, 
-                 output_channels=1, image_size=64):
+    def __init__(self, noise_dim=100, feature_dim=128, style_dim=32, output_channels=1, image_size=64):
         super().__init__(noise_dim, feature_dim, output_channels, image_size)
-        
-        # Style embedding layer
         self.style_embedding = nn.Sequential(
             nn.Linear(style_dim, 64),
             nn.ReLU(),
             nn.Linear(64, 128)
         )
-        
-        # Update input projection to include style
-        self.input_projection = nn.Linear(noise_dim + feature_dim + 128, 
-                                        256 * self.start_size * self.start_size)
+        self.input_projection = nn.Linear(noise_dim + feature_dim + 128, 256 * self.start_size * self.start_size)
     
     def forward(self, noise, features=None, style=None):
-        """
-        Generate with style conditioning.
-        
-        Args:
-            noise: Random noise
-            features: Design features
-            style: Style vector
-        """
         batch_size = noise.size(0)
-        
-        # Process style
         if style is not None:
             style_embed = self.style_embedding(style)
         else:
             style_embed = torch.zeros(batch_size, 128, device=noise.device)
         
-        # Combine all inputs
         if features is not None:
             x = torch.cat([noise, features, style_embed], dim=1)
         else:
             zero_features = torch.zeros(batch_size, self.feature_dim, device=noise.device)
             x = torch.cat([noise, zero_features, style_embed], dim=1)
         
-        # Continue with parent forward pass
         x = self.input_projection(x)
         x = x.view(batch_size, 256, self.start_size, self.start_size)
         
-        # Upsample and refine
         for i in range(self.num_upsamples):
             x = self.upsample_layers[i](x)
             x = self.bn_layers[i](x)
             x = F.relu(x)
-            
             x = self.conv_layers[i](x)
             x = self.bn_layers[i](x)
             x = F.relu(x)
         
-        # Final output
-        x = self.final_conv(x)
-        x = torch.tanh(x)
-        
-        return x
+        return torch.tanh(self.final_conv(x))
+
+
+# -------------------------------
+# Utility: easy generation method
+# -------------------------------
+def generate_kolam_samples(generator, num_samples=4, device="cpu"):
+    """Generate sample Kolams with random noise + styles."""
+    generator.eval()
+    with torch.no_grad():
+        noise = torch.randn(num_samples, generator.noise_dim, device=device)
+        features = torch.randn(num_samples, generator.feature_dim, device=device)
+        
+        if isinstance(generator, StyleConditionedGenerator):
+            style = torch.randn(num_samples, 32, device=device)
+            return generator(noise, features, style)
+        else:
+            return generator(noise, features)
 
 
 if __name__ == "__main__":
-    # Test the generator
-    generator = KolamGenerator()
-    noise = torch.randn(4, 100)  # Batch of 4, 100-dim noise
-    features = torch.randn(4, 128)  # Batch of 4, 128-dim features
-    
-    generated = generator(noise, features)
-    print(f"Noise shape: {noise.shape}")
-    print(f"Features shape: {features.shape}")
-    print(f"Generated shape: {generated.shape}")
-    
-    # Test style-conditioned generator
-    style_gen = StyleConditionedGenerator()
-    style = torch.randn(4, 32)
-    style_generated = style_gen(noise, features, style)
-    print(f"Style-generated shape: {style_generated.shape}")
+    gen = StyleConditionedGenerator()
+    samples = generate_kolam_samples(gen, num_samples=2)
+    print("Generated:", samples.shape)