Create generator.py

generator.py

@@ -0,0 +1,388 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import random
+
+# --- Helper Modules ---
+
+class LeakyReLU(nn.Module):
+    """
+    Custom LeakyReLU implementation to allow for a fixed negative slope
+    and in-place operation.
+    """
+    def __init__(self, negative_slope=0.2, inplace=False):
+        super().__init__()
+        self.negative_slope = negative_slope
+        self.inplace = inplace
+
+    def forward(self, x):
+        return F.leaky_relu(x, self.negative_slope, self.inplace)
+
+class PixelNorm(nn.Module):
+    """
+    Pixel-wise feature vector normalization.
+    """
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        # Epsilon added for numerical stability
+        return x * torch.rsqrt(torch.mean(x**2, dim=1, keepdim=True) + 1e-8)
+
+class ModulatedConv2d(nn.Module):
+    """
+    This is the core building block of the StyleGAN2 synthesis network.
+    It applies style modulation and demodulation.
+    """
+    def __init__(self, in_channels, out_channels, kernel_size, style_dim, demodulate=True, upsample=False):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.style_dim = style_dim
+        self.demodulate = demodulate
+        self.upsample = upsample
+
+        # Standard convolution weights
+        self.weight = nn.Parameter(
+            torch.randn(1, out_channels, in_channels, kernel_size, kernel_size)
+        )
+
+        # Affine transform (A) from style vector (w)
+        self.modulation = nn.Linear(style_dim, in_channels, bias=True)
+
+        # Initialize modulation bias to 1 (identity transform)
+        nn.init.constant_(self.modulation.bias, 1.0)
+
+        # Padding for the convolution
+        self.padding = (kernel_size - 1) // 2
+
+        # Upsampling filter (if needed)
+        if self.upsample:
+            # Using a simple bilinear filter
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
+
+    def forward(self, x, style):
+        # Store initial batch_size and in_channels
+        batch_size, in_channels_original, _, _ = x.shape
+
+        # 1. Modulate (Style-based feature scaling)
+        # style shape: [batch_size, style_dim]
+        # s shape: [batch_size, 1, in_channels, 1, 1]
+        s = self.modulation(style).view(batch_size, 1, in_channels_original, 1, 1)
+
+        # Get conv weights and combine with modulation
+        # w shape: [batch_size, out_channels, in_channels, k, k]
+        w = self.weight * s
+
+        # 2. Demodulate (Normalize weights to prevent scale explosion)
+        if self.demodulate:
+            # Calculate per-weight normalization factor
+            d = torch.rsqrt(torch.sum(w**2, dim=[2, 3, 4], keepdim=True) + 1e-8)
+            w = w * d
+
+        # 3. Upsample (if applicable)
+        if self.upsample:
+            x = self.up(x)
+
+        # Get current height and width *after* potential upsampling
+        current_height = x.shape[2]
+        current_width = x.shape[3]
+
+        # 4. Convolution
+        # Because weights are now per-batch, we need to group convolutions
+        # We reshape x and w to use a single grouped convolution operation
+
+        x = x.view(1, batch_size * in_channels_original, current_height, current_width)
+        w = w.view(batch_size * self.out_channels, in_channels_original, self.kernel_size, self.kernel_size)
+
+        # padding='same' is not supported for strided/grouped conv, so we use manual padding
+        x = F.conv2d(x, w, padding=self.padding, groups=batch_size)
+
+        # Reshape back to [batch_size, out_channels, h, w]
+        _, _, new_height, new_width = x.shape
+        x = x.view(batch_size, self.out_channels, new_height, new_width)
+
+        return x
+
+class NoiseInjection(nn.Module):
+    """
+    Adds scaled noise to the feature maps.
+    """
+    def __init__(self, channels):
+        super().__init__()
+        # Learned scaling factor for the noise
+        self.weight = nn.Parameter(torch.zeros(1, channels, 1, 1))
+
+    def forward(self, x, noise=None):
+        if noise is None:
+            batch, _, height, width = x.shape
+            noise = torch.randn(batch, 1, height, width, device=x.device, dtype=x.dtype)
+
+        return x + self.weight * noise
+
+class ConstantInput(nn.Module):
+    """
+    A learned constant 4x4 feature map to start the synthesis process.
+    """
+    def __init__(self, channels, size=4):
+        super().__init__()
+        self.input = nn.Parameter(torch.randn(1, channels, size, size))
+
+    def forward(self, batch_size):
+        return self.input.repeat(batch_size, 1, 1, 1)
+
+class ToRGB(nn.Module):
+    """
+    Projects feature maps to an RGB image.
+    Uses a 1x1 modulated convolution.
+    """
+    def __init__(self, in_channels, out_channels, style_dim):
+        super().__init__()
+        # 1x1 convolution
+        self.conv = ModulatedConv2d(in_channels, out_channels, 1, style_dim, demodulate=False, upsample=False)
+        self.bias = nn.Parameter(torch.zeros(1, out_channels, 1, 1))
+
+    def forward(self, x, style, skip=None):
+        x = self.conv(x, style)
+        x = x + self.bias
+
+        if skip is not None:
+            # Upsample the previous RGB output and add
+            skip = F.interpolate(skip, scale_factor=2, mode='bilinear', align_corners=False)
+            x = x + skip
+
+        return x
+
+# --- Main Generator Components ---
+
+class MappingNetwork(nn.Module):
+    """
+    Maps the initial latent vector Z to the intermediate style vector W.
+    """
+    def __init__(self, z_dim, w_dim, num_layers=8):
+        super().__init__()
+        self.z_dim = z_dim
+        self.w_dim = w_dim
+
+        layers = [PixelNorm()]
+        for i in range(num_layers):
+            layers.extend([
+                nn.Linear(z_dim if i == 0 else w_dim, w_dim),
+                LeakyReLU(0.2, inplace=True)
+            ])
+
+        self.mapping = nn.Sequential(*layers)
+
+    def forward(self, z):
+        # z shape: [batch_size, z_dim]
+        w = self.mapping(z)
+        # w shape: [batch_size, w_dim]
+        return w
+
+class SynthesisBlock(nn.Module):
+    """
+    A single block in the Synthesis Network (e.g., 8x8 -> 16x16).
+    Contains upsampling, modulated convolutions, noise, and activation.
+    """
+    def __init__(self, in_channels, out_channels, style_dim):
+        super().__init__()
+        # First modulated conv with upsampling
+        self.conv1 = ModulatedConv2d(in_channels, out_channels, 3, style_dim, upsample=True)
+        self.noise1 = NoiseInjection(out_channels)
+        self.activate1 = LeakyReLU(0.2, inplace=True)
+
+        # Second modulated conv
+        self.conv2 = ModulatedConv2d(out_channels, out_channels, 3, style_dim, upsample=False)
+        self.noise2 = NoiseInjection(out_channels)
+        self.activate2 = LeakyReLU(0.2, inplace=True)
+
+    def forward(self, x, w, noise1, noise2):
+        x = self.conv1(x, w)
+        x = self.noise1(x, noise1)
+        x = self.activate1(x)
+
+        x = self.conv2(x, w)
+        x = self.noise2(x, noise2)
+        x = self.activate2(x)
+
+        return x
+
+class SynthesisNetwork(nn.Module):
+    """
+    Builds the image from the style vector W.
+    """
+    def __init__(self, w_dim, img_channels, img_resolution=256, start_res=4, num_blocks=None):
+        super().__init__()
+        self.w_dim = w_dim
+        self.img_channels = img_channels
+        self.start_res = start_res
+
+        if num_blocks is None:
+            self.num_blocks = int(math.log2(img_resolution) - math.log2(start_res))
+            self.img_resolution = img_resolution
+        else:
+            self.num_blocks = num_blocks
+            self.img_resolution = start_res * (2**self.num_blocks)
+            print(f"Synthesis network created with {self.num_blocks} blocks, output resolution: {self.img_resolution}x{self.img_resolution}")
+
+        channels = {
+            4: 512,
+            8: 512,
+            16: 512,
+            32: 512,
+            64: 256,
+            128: 128,
+            256: 64,
+            512: 32,
+            1024: 16,
+        }
+
+        self.input = ConstantInput(channels[start_res])
+
+        self.conv1 = ModulatedConv2d(channels[start_res], channels[start_res], 3, w_dim, upsample=False)
+        self.noise1 = NoiseInjection(channels[start_res])
+        self.activate1 = LeakyReLU(0.2, inplace=True)
+
+        self.to_rgb1 = ToRGB(channels[start_res], img_channels, w_dim)
+
+        self.blocks = nn.ModuleList()
+        self.to_rgbs = nn.ModuleList()
+
+        in_c = channels[start_res]
+
+        for i in range(self.num_blocks):
+            current_res = start_res * (2**(i+1))
+            out_c = channels.get(current_res, 16)
+            if current_res > 1024:
+                print(f"Warning: Resolution {current_res}x{current_res} not in channel map. Using {out_c} channels.")
+
+            self.blocks.append(SynthesisBlock(in_c, out_c, w_dim))
+            self.to_rgbs.append(ToRGB(out_c, img_channels, w_dim))
+
+            in_c = out_c
+
+        # Number of style vectors needed: 1 for initial conv1, 1 for initial to_rgb, and 3 per block (conv1, conv2, to_rgb)
+        self.num_styles = self.num_blocks * 3 + 2 # Corrected num_styles
+
+    def forward(self, w, noise=None):
+        # w shape: [batch_size, num_styles, w_dim]
+        if w.ndim == 2:
+            w = w.unsqueeze(1).repeat(1, self.num_styles, 1)
+
+        batch_size = w.shape[0]
+
+        # --- Handle Noise (generate if None) ---
+        if noise is None:
+            noise_list = []
+            # Noise for the initial 4x4 conv (self.conv1)
+            noise_list.append(torch.randn(batch_size, 1, self.start_res, self.start_res, device=w.device))
+
+            current_res = self.start_res
+            # Iterate through the synthesis blocks to generate noise for each
+            for i in range(self.num_blocks):
+                current_res *= 2 # This is the resolution *after* the current block's upsampling
+                # Noise for the first conv of the current block (after upsampling)
+                noise_list.append(torch.randn(batch_size, 1, current_res, current_res, device=w.device))
+                # Noise for the second conv of the current block (same resolution)
+                noise_list.append(torch.randn(batch_size, 1, current_res, current_res, device=w.device))
+            noise = noise_list
+
+        # --- 4x4 Block ---
+        x = self.input(batch_size)
+        x = self.conv1(x, w[:, 0]) # Style for initial conv1
+        x = self.noise1(x, noise[0]) # Noise for initial conv1
+        x = self.activate1(x)
+
+        skip = self.to_rgb1(x, w[:, 1]) # Style for initial ToRGB
+
+        # --- Main blocks (8x8 to img_resolution) ---
+        current_noise_idx_in_list = 1 # index for noise_list: noise[0] was used above
+        current_style_idx_in_w = 2   # index for w: w[:,0] and w[:,1] were used above
+
+        for i, (block, to_rgb) in enumerate(zip(self.blocks, self.to_rgbs)):
+            # Styles for this block
+            w_block_conv1 = w[:, current_style_idx_in_w]
+            w_block_conv2 = w[:, current_style_idx_in_w + 1]
+            w_block_to_rgb = w[:, current_style_idx_in_w + 2]
+
+            # Noises for this block
+            n_block_conv1 = noise[current_noise_idx_in_list]
+            n_block_conv2 = noise[current_noise_idx_in_list + 1]
+
+            x = block(x, w_block_conv1, n_block_conv1, n_block_conv2)
+
+            skip = to_rgb(x, w_block_to_rgb, skip)
+
+            # Increment indices for next block
+            current_style_idx_in_w += 3
+            current_noise_idx_in_list += 2
+
+        return skip # Final RGB image
+
+class Generator(nn.Module):
+    """
+    The complete StyleGAN2 Generator.
+    Combines the Mapping and Synthesis networks.
+    """
+    def __init__(self, z_dim, w_dim, img_resolution, img_channels,
+                 mapping_layers=8, num_synthesis_blocks=None):
+        super().__init__()
+        self.z_dim = z_dim
+        self.w_dim = w_dim
+
+        self.mapping = MappingNetwork(z_dim, w_dim, mapping_layers)
+
+        self.synthesis = SynthesisNetwork(
+            w_dim, img_channels, img_resolution, num_blocks=num_synthesis_blocks
+        )
+
+        self.num_styles = self.synthesis.num_styles
+        self.img_resolution = self.synthesis.img_resolution # Get final resolution
+
+        # For truncation trick
+        self.register_buffer('w_avg', torch.zeros(w_dim))
+
+    def update_w_avg(self, new_w, momentum=0.995):
+        """Helper to update the moving average of W"""
+        self.w_avg = torch.lerp(new_w.mean(0), self.w_avg, momentum)
+
+    def forward(self, z, truncation_psi=0.7, use_truncation=True,
+                style_mix_prob=0.0, noise=None):
+
+        # --- 1. Get W vector(s) ---
+
+        # Check if we're doing style mixing
+        do_style_mix = False
+        if isinstance(z, list) and len(z) == 2:
+            do_style_mix = True
+            z1, z2 = z
+            w1 = self.mapping(z1) # [batch, w_dim]
+            w2 = self.mapping(z2) # [batch, w_dim]
+        else:
+            w = self.mapping(z) # [batch, w_dim]
+            w1 = w
+            w2 = w
+
+        # --- 2. Truncation Trick ---
+        if use_truncation:
+            w1 = torch.lerp(self.w_avg, w1, truncation_psi)
+            w2 = torch.lerp(self.w_avg, w2, truncation_psi)
+
+        # --- 3. Style Mixing ---
+        # w_final shape: [batch, num_styles, w_dim]
+        w_final = torch.empty(w.shape[0], self.num_styles, self.w_dim, device=w.device)
+
+        if do_style_mix and random.random() < style_mix_prob:
+            # Select a random crossover point
+            mix_cutoff = random.randint(1, self.num_styles - 1)
+            w_final[:, :mix_cutoff] = w1.unsqueeze(1) # [batch, cutoff, w_dim]
+            w_final[:, mix_cutoff:] = w2.unsqueeze(1) # [batch, num_styles-cutoff, w_dim]
+        else:
+            # No mixing, just use w1
+            w_final = w1.unsqueeze(1).repeat(1, self.num_styles, 1)
+
+        # --- 4. Synthesis ---
+        img = self.synthesis(w_final, noise)
+        return img