Image Generation with GANs

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import numpy as np
5
6class GANFundamentals:
7    """
8    Understanding the core GAN training paradigm.
9
10    Minimax objective:
11    min_G max_D V(D, G) = E[log(D(x))] + E[log(1 - D(G(z)))]
12
13    The Generator tries to fool the Discriminator.
14    The Discriminator tries to correctly classify real vs fake.
15    """
16
17    @staticmethod
18    def gan_loss_discriminator(real_output: torch.Tensor,
19                               fake_output: torch.Tensor) -> torch.Tensor:
20        """
21        Discriminator loss: maximize log(D(real)) + log(1 - D(fake))
22
23        In practice, we minimize: -log(D(real)) - log(1 - D(fake))
24        Which is equivalent to BCE loss.
25        """
26        # Real images should be classified as 1
27        real_loss = F.binary_cross_entropy_with_logits(
28            real_output,
29            torch.ones_like(real_output)
30        )
31
32        # Fake images should be classified as 0
33        fake_loss = F.binary_cross_entropy_with_logits(
34            fake_output,
35            torch.zeros_like(fake_output)
36        )
37
38        return real_loss + fake_loss
39
40    @staticmethod
41    def gan_loss_generator(fake_output: torch.Tensor) -> torch.Tensor:
42        """
43        Generator loss: minimize log(1 - D(fake))
44
45        In practice, we maximize log(D(fake)) instead
46        (non-saturating loss for better gradients).
47        """
48        # Generator wants fake images classified as 1 (real)
49        return F.binary_cross_entropy_with_logits(
50            fake_output,
51            torch.ones_like(fake_output)
52        )
53
54
55class SimpleGenerator(nn.Module):
56    """
57    Basic Generator network.
58    Maps latent noise vector to image.
59    """
60
61    def __init__(self, latent_dim: int = 100, img_channels: int = 3,
62                 img_size: int = 64):
63        super().__init__()
64        self.latent_dim = latent_dim
65        self.img_size = img_size
66
67        # Calculate initial feature map size
68        self.init_size = img_size // 16  # 4 for 64x64
69
70        # Project and reshape
71        self.fc = nn.Linear(latent_dim, 512 * self.init_size ** 2)
72
73        # Upsampling blocks
74        self.conv_blocks = nn.Sequential(
75            nn.BatchNorm2d(512),
76            nn.Upsample(scale_factor=2),
77            nn.Conv2d(512, 256, 3, stride=1, padding=1),
78            nn.BatchNorm2d(256),
79            nn.LeakyReLU(0.2, inplace=True),
80
81            nn.Upsample(scale_factor=2),
82            nn.Conv2d(256, 128, 3, stride=1, padding=1),
83            nn.BatchNorm2d(128),
84            nn.LeakyReLU(0.2, inplace=True),
85
86            nn.Upsample(scale_factor=2),
87            nn.Conv2d(128, 64, 3, stride=1, padding=1),
88            nn.BatchNorm2d(64),
89            nn.LeakyReLU(0.2, inplace=True),
90
91            nn.Upsample(scale_factor=2),
92            nn.Conv2d(64, img_channels, 3, stride=1, padding=1),
93            nn.Tanh()  # Output in [-1, 1]
94        )
95
96    def forward(self, z: torch.Tensor) -> torch.Tensor:
97        """Generate image from noise vector z."""
98        x = self.fc(z)
99        x = x.view(-1, 512, self.init_size, self.init_size)
100        img = self.conv_blocks(x)
101        return img
102
103
104class SimpleDiscriminator(nn.Module):
105    """
106    Basic Discriminator network.
107    Classifies images as real or fake.
108    """
109
110    def __init__(self, img_channels: int = 3, img_size: int = 64):
111        super().__init__()
112
113        def discriminator_block(in_channels, out_channels,
114                               batch_norm: bool = True):
115            layers = [
116                nn.Conv2d(in_channels, out_channels, 4, stride=2, padding=1),
117                nn.LeakyReLU(0.2, inplace=True)
118            ]
119            if batch_norm:
120                layers.insert(1, nn.BatchNorm2d(out_channels))
121            return layers
122
123        self.model = nn.Sequential(
124            *discriminator_block(img_channels, 64, batch_norm=False),
125            *discriminator_block(64, 128),
126            *discriminator_block(128, 256),
127            *discriminator_block(256, 512),
128        )
129
130        # Calculate output size after convolutions
131        ds_size = img_size // 16
132        self.fc = nn.Sequential(
133            nn.Flatten(),
134            nn.Linear(512 * ds_size ** 2, 1)
135            # No sigmoid - using BCEWithLogitsLoss
136        )
137
138    def forward(self, img: torch.Tensor) -> torch.Tensor:
139        """Classify image as real (1) or fake (0)."""
140        features = self.model(img)
141        validity = self.fc(features)
142        return validity
143
144
145# Example usage
146latent_dim = 100
147img_size = 64
148
149generator = SimpleGenerator(latent_dim=latent_dim, img_size=img_size)
150discriminator = SimpleDiscriminator(img_size=img_size)
151
152# Generate fake images
153z = torch.randn(16, latent_dim)
154fake_images = generator(z)
155print(f"Generated images shape: {fake_images.shape}")
156
157# Discriminator output
158real_images = torch.randn(16, 3, img_size, img_size)
159real_output = discriminator(real_images)
160fake_output = discriminator(fake_images)
161print(f"Discriminator outputs - Real: {real_output.mean():.4f}, Fake: {fake_output.mean():.4f}")

1import torch
2import torch.nn as nn
3
4class DCGANGenerator(nn.Module):
5    """
6    DCGAN Generator with transposed convolutions.
7
8    Architecture guidelines:
9    - Use transposed convolutions for upsampling
10    - Batch normalization in all layers except output
11    - ReLU activation in generator
12    - Tanh output activation
13    """
14
15    def __init__(self, latent_dim: int = 100, feature_maps: int = 64,
16                 img_channels: int = 3):
17        super().__init__()
18        self.latent_dim = latent_dim
19        nf = feature_maps
20
21        self.main = nn.Sequential(
22            # Input: latent_dim x 1 x 1
23            nn.ConvTranspose2d(latent_dim, nf * 8, 4, 1, 0, bias=False),
24            nn.BatchNorm2d(nf * 8),
25            nn.ReLU(True),
26            # State: (nf*8) x 4 x 4
27
28            nn.ConvTranspose2d(nf * 8, nf * 4, 4, 2, 1, bias=False),
29            nn.BatchNorm2d(nf * 4),
30            nn.ReLU(True),
31            # State: (nf*4) x 8 x 8
32
33            nn.ConvTranspose2d(nf * 4, nf * 2, 4, 2, 1, bias=False),
34            nn.BatchNorm2d(nf * 2),
35            nn.ReLU(True),
36            # State: (nf*2) x 16 x 16
37
38            nn.ConvTranspose2d(nf * 2, nf, 4, 2, 1, bias=False),
39            nn.BatchNorm2d(nf),
40            nn.ReLU(True),
41            # State: nf x 32 x 32
42
43            nn.ConvTranspose2d(nf, img_channels, 4, 2, 1, bias=False),
44            nn.Tanh()
45            # Output: img_channels x 64 x 64
46        )
47
48        # Weight initialization
49        self.apply(self._init_weights)
50
51    def _init_weights(self, m):
52        """DCGAN weight initialization."""
53        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
54            nn.init.normal_(m.weight, 0.0, 0.02)
55        elif isinstance(m, nn.BatchNorm2d):
56            nn.init.normal_(m.weight, 1.0, 0.02)
57            nn.init.constant_(m.bias, 0)
58
59    def forward(self, z: torch.Tensor) -> torch.Tensor:
60        # Reshape z for convolution: [B, latent_dim] -> [B, latent_dim, 1, 1]
61        z = z.view(-1, self.latent_dim, 1, 1)
62        return self.main(z)
63
64
65class DCGANDiscriminator(nn.Module):
66    """
67    DCGAN Discriminator with strided convolutions.
68
69    Architecture guidelines:
70    - Use strided convolutions (no pooling)
71    - Batch normalization except first layer
72    - LeakyReLU activation
73    """
74
75    def __init__(self, img_channels: int = 3, feature_maps: int = 64):
76        super().__init__()
77        nf = feature_maps
78
79        self.main = nn.Sequential(
80            # Input: img_channels x 64 x 64
81            nn.Conv2d(img_channels, nf, 4, 2, 1, bias=False),
82            nn.LeakyReLU(0.2, inplace=True),
83            # State: nf x 32 x 32
84
85            nn.Conv2d(nf, nf * 2, 4, 2, 1, bias=False),
86            nn.BatchNorm2d(nf * 2),
87            nn.LeakyReLU(0.2, inplace=True),
88            # State: (nf*2) x 16 x 16
89
90            nn.Conv2d(nf * 2, nf * 4, 4, 2, 1, bias=False),
91            nn.BatchNorm2d(nf * 4),
92            nn.LeakyReLU(0.2, inplace=True),
93            # State: (nf*4) x 8 x 8
94
95            nn.Conv2d(nf * 4, nf * 8, 4, 2, 1, bias=False),
96            nn.BatchNorm2d(nf * 8),
97            nn.LeakyReLU(0.2, inplace=True),
98            # State: (nf*8) x 4 x 4
99
100            nn.Conv2d(nf * 8, 1, 4, 1, 0, bias=False),
101            # Output: 1 x 1 x 1
102        )
103
104        self.apply(self._init_weights)
105
106    def _init_weights(self, m):
107        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
108            nn.init.normal_(m.weight, 0.0, 0.02)
109        elif isinstance(m, nn.BatchNorm2d):
110            nn.init.normal_(m.weight, 1.0, 0.02)
111            nn.init.constant_(m.bias, 0)
112
113    def forward(self, img: torch.Tensor) -> torch.Tensor:
114        return self.main(img).view(-1, 1)
115
116
117class DCGANTrainer:
118    """Complete DCGAN training pipeline."""
119
120    def __init__(self, latent_dim: int = 100, lr: float = 0.0002,
121                 beta1: float = 0.5, device: str = 'cuda'):
122        self.latent_dim = latent_dim
123        self.device = device if torch.cuda.is_available() else 'cpu'
124
125        # Initialize networks
126        self.generator = DCGANGenerator(latent_dim).to(self.device)
127        self.discriminator = DCGANDiscriminator().to(self.device)
128
129        # Optimizers (Adam with specific betas for GAN training)
130        self.opt_g = torch.optim.Adam(
131            self.generator.parameters(),
132            lr=lr,
133            betas=(beta1, 0.999)
134        )
135        self.opt_d = torch.optim.Adam(
136            self.discriminator.parameters(),
137            lr=lr,
138            betas=(beta1, 0.999)
139        )
140
141        # Loss function
142        self.criterion = nn.BCEWithLogitsLoss()
143
144        # Fixed noise for visualization
145        self.fixed_noise = torch.randn(64, latent_dim, device=self.device)
146
147    def train_step(self, real_images: torch.Tensor):
148        """Single training step for both networks."""
149        batch_size = real_images.size(0)
150        real_images = real_images.to(self.device)
151
152        # Labels
153        real_label = torch.ones(batch_size, 1, device=self.device)
154        fake_label = torch.zeros(batch_size, 1, device=self.device)
155
156        # ---------------------
157        # Train Discriminator
158        # ---------------------
159        self.opt_d.zero_grad()
160
161        # Real images
162        output_real = self.discriminator(real_images)
163        loss_d_real = self.criterion(output_real, real_label)
164
165        # Fake images
166        noise = torch.randn(batch_size, self.latent_dim, device=self.device)
167        fake_images = self.generator(noise)
168        output_fake = self.discriminator(fake_images.detach())
169        loss_d_fake = self.criterion(output_fake, fake_label)
170
171        # Combined loss
172        loss_d = loss_d_real + loss_d_fake
173        loss_d.backward()
174        self.opt_d.step()
175
176        # ---------------------
177        # Train Generator
178        # ---------------------
179        self.opt_g.zero_grad()
180
181        # Generator wants discriminator to think fake is real
182        output_fake = self.discriminator(fake_images)
183        loss_g = self.criterion(output_fake, real_label)
184
185        loss_g.backward()
186        self.opt_g.step()
187
188        return {
189            'loss_d': loss_d.item(),
190            'loss_g': loss_g.item(),
191            'd_real': torch.sigmoid(output_real).mean().item(),
192            'd_fake': torch.sigmoid(output_fake).mean().item()
193        }
194
195    @torch.no_grad()
196    def generate_samples(self, num_samples: int = 64):
197        """Generate sample images."""
198        self.generator.eval()
199        noise = torch.randn(num_samples, self.latent_dim, device=self.device)
200        samples = self.generator(noise)
201        self.generator.train()
202        return samples
203
204# Example
205trainer = DCGANTrainer(latent_dim=100)
206print("DCGAN Trainer initialized")
207print(f"Generator parameters: {sum(p.numel() for p in trainer.generator.parameters()):,}")
208print(f"Discriminator parameters: {sum(p.numel() for p in trainer.discriminator.parameters()):,}")

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5class MappingNetwork(nn.Module):
6    """
7    StyleGAN Mapping Network.
8    Maps latent code z to intermediate latent w.
9    This disentangles the latent space.
10    """
11
12    def __init__(self, latent_dim: int = 512, mapping_layers: int = 8):
13        super().__init__()
14
15        layers = []
16        for _ in range(mapping_layers):
17            layers.extend([
18                nn.Linear(latent_dim, latent_dim),
19                nn.LeakyReLU(0.2)
20            ])
21
22        self.mapping = nn.Sequential(*layers)
23
24    def forward(self, z: torch.Tensor) -> torch.Tensor:
25        """Map z to w."""
26        # Normalize z
27        z = F.normalize(z, dim=1)
28        return self.mapping(z)
29
30
31class AdaIN(nn.Module):
32    """
33    Adaptive Instance Normalization.
34    Injects style information into feature maps.
35    """
36
37    def __init__(self, style_dim: int, num_features: int):
38        super().__init__()
39        self.norm = nn.InstanceNorm2d(num_features)
40        self.style_scale = nn.Linear(style_dim, num_features)
41        self.style_bias = nn.Linear(style_dim, num_features)
42
43    def forward(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
44        """
45        Apply AdaIN.
46
47        Args:
48            x: Feature maps [B, C, H, W]
49            w: Style vector [B, style_dim]
50
51        Returns:
52            Styled feature maps
53        """
54        # Instance normalize
55        x = self.norm(x)
56
57        # Get style parameters
58        gamma = self.style_scale(w).unsqueeze(-1).unsqueeze(-1)
59        beta = self.style_bias(w).unsqueeze(-1).unsqueeze(-1)
60
61        # Apply affine transformation
62        return gamma * x + beta
63
64
65class NoiseInjection(nn.Module):
66    """Inject scaled noise into feature maps."""
67
68    def __init__(self, channels: int):
69        super().__init__()
70        self.weight = nn.Parameter(torch.zeros(1, channels, 1, 1))
71
72    def forward(self, x: torch.Tensor,
73                noise: torch.Tensor = None) -> torch.Tensor:
74        if noise is None:
75            batch, _, height, width = x.shape
76            noise = torch.randn(batch, 1, height, width, device=x.device)
77        return x + self.weight * noise
78
79
80class StyleConvBlock(nn.Module):
81    """
82    StyleGAN synthesis block with style injection.
83    """
84
85    def __init__(self, in_channels: int, out_channels: int,
86                 style_dim: int = 512, upsample: bool = True):
87        super().__init__()
88
89        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear',
90                                    align_corners=False) if upsample else None
91
92        self.conv = nn.Conv2d(in_channels, out_channels, 3, padding=1)
93        self.noise = NoiseInjection(out_channels)
94        self.adain = AdaIN(style_dim, out_channels)
95        self.activation = nn.LeakyReLU(0.2)
96
97    def forward(self, x: torch.Tensor, w: torch.Tensor,
98                noise: torch.Tensor = None) -> torch.Tensor:
99        if self.upsample:
100            x = self.upsample(x)
101
102        x = self.conv(x)
103        x = self.noise(x, noise)
104        x = self.adain(x, w)
105        x = self.activation(x)
106
107        return x
108
109
110class StyleGANGenerator(nn.Module):
111    """
112    Simplified StyleGAN-style Generator.
113    """
114
115    def __init__(self, latent_dim: int = 512, style_dim: int = 512,
116                 img_channels: int = 3):
117        super().__init__()
118
119        self.mapping = MappingNetwork(latent_dim, mapping_layers=8)
120
121        # Learned constant input (4x4)
122        self.constant = nn.Parameter(torch.randn(1, 512, 4, 4))
123
124        # Synthesis blocks
125        self.blocks = nn.ModuleList([
126            StyleConvBlock(512, 512, style_dim, upsample=False),  # 4x4
127            StyleConvBlock(512, 512, style_dim),  # 8x8
128            StyleConvBlock(512, 256, style_dim),  # 16x16
129            StyleConvBlock(256, 128, style_dim),  # 32x32
130            StyleConvBlock(128, 64, style_dim),   # 64x64
131        ])
132
133        # ToRGB layers
134        self.to_rgb = nn.Conv2d(64, img_channels, 1)
135
136    def forward(self, z: torch.Tensor,
137                truncation: float = 1.0) -> torch.Tensor:
138        """
139        Generate image from latent code.
140
141        Args:
142            z: Latent code [B, latent_dim]
143            truncation: Truncation trick factor (1.0 = no truncation)
144
145        Returns:
146            Generated image [B, 3, 64, 64]
147        """
148        batch_size = z.size(0)
149
150        # Map z to w
151        w = self.mapping(z)
152
153        # Apply truncation trick
154        if truncation < 1.0:
155            # In practice, compute w_avg during training
156            w_avg = torch.zeros_like(w)
157            w = w_avg + truncation * (w - w_avg)
158
159        # Start from constant
160        x = self.constant.repeat(batch_size, 1, 1, 1)
161
162        # Apply synthesis blocks
163        for block in self.blocks:
164            x = block(x, w)
165
166        # Final RGB output
167        img = self.to_rgb(x)
168        img = torch.tanh(img)
169
170        return img
171
172    def style_mixing(self, z1: torch.Tensor, z2: torch.Tensor,
173                    crossover_layer: int = 3) -> torch.Tensor:
174        """
175        Generate image with style mixing.
176        Use w1 for early layers, w2 for later layers.
177        """
178        batch_size = z1.size(0)
179
180        w1 = self.mapping(z1)
181        w2 = self.mapping(z2)
182
183        x = self.constant.repeat(batch_size, 1, 1, 1)
184
185        for i, block in enumerate(self.blocks):
186            w = w1 if i < crossover_layer else w2
187            x = block(x, w)
188
189        img = torch.tanh(self.to_rgb(x))
190        return img
191
192
193# Example usage
194stylegan = StyleGANGenerator(latent_dim=512)
195
196z = torch.randn(4, 512)
197images = stylegan(z, truncation=0.7)
198print(f"StyleGAN output shape: {images.shape}")
199
200# Style mixing example
201z1, z2 = torch.randn(4, 512), torch.randn(4, 512)
202mixed = stylegan.style_mixing(z1, z2, crossover_layer=3)
203print(f"Style mixed output shape: {mixed.shape}")

1import torch
2import torch.nn as nn
3
4class UNetDown(nn.Module):
5    """Encoder block for U-Net generator."""
6
7    def __init__(self, in_channels: int, out_channels: int,
8                 normalize: bool = True):
9        super().__init__()
10
11        layers = [nn.Conv2d(in_channels, out_channels, 4, 2, 1, bias=False)]
12        if normalize:
13            layers.append(nn.InstanceNorm2d(out_channels))
14        layers.append(nn.LeakyReLU(0.2))
15
16        self.model = nn.Sequential(*layers)
17
18    def forward(self, x: torch.Tensor) -> torch.Tensor:
19        return self.model(x)
20
21
22class UNetUp(nn.Module):
23    """Decoder block for U-Net generator."""
24
25    def __init__(self, in_channels: int, out_channels: int,
26                 dropout: float = 0.0):
27        super().__init__()
28
29        layers = [
30            nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1, bias=False),
31            nn.InstanceNorm2d(out_channels),
32            nn.ReLU(inplace=True)
33        ]
34        if dropout > 0:
35            layers.append(nn.Dropout(dropout))
36
37        self.model = nn.Sequential(*layers)
38
39    def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
40        x = self.model(x)
41        x = torch.cat([x, skip], dim=1)
42        return x
43
44
45class Pix2PixGenerator(nn.Module):
46    """
47    U-Net based Generator for Pix2Pix.
48    Takes conditional input and generates target image.
49    """
50
51    def __init__(self, in_channels: int = 3, out_channels: int = 3):
52        super().__init__()
53
54        # Encoder (downsampling)
55        self.down1 = UNetDown(in_channels, 64, normalize=False)  # 128
56        self.down2 = UNetDown(64, 128)   # 64
57        self.down3 = UNetDown(128, 256)  # 32
58        self.down4 = UNetDown(256, 512)  # 16
59        self.down5 = UNetDown(512, 512)  # 8
60        self.down6 = UNetDown(512, 512)  # 4
61        self.down7 = UNetDown(512, 512)  # 2
62        self.down8 = UNetDown(512, 512, normalize=False)  # 1
63
64        # Decoder (upsampling with skip connections)
65        self.up1 = UNetUp(512, 512, dropout=0.5)    # 2
66        self.up2 = UNetUp(1024, 512, dropout=0.5)   # 4
67        self.up3 = UNetUp(1024, 512, dropout=0.5)   # 8
68        self.up4 = UNetUp(1024, 512)   # 16
69        self.up5 = UNetUp(1024, 256)   # 32
70        self.up6 = UNetUp(512, 128)    # 64
71        self.up7 = UNetUp(256, 64)     # 128
72
73        # Final layer
74        self.final = nn.Sequential(
75            nn.ConvTranspose2d(128, out_channels, 4, 2, 1),
76            nn.Tanh()
77        )
78
79    def forward(self, x: torch.Tensor) -> torch.Tensor:
80        # Encoder with skip connections
81        d1 = self.down1(x)
82        d2 = self.down2(d1)
83        d3 = self.down3(d2)
84        d4 = self.down4(d3)
85        d5 = self.down5(d4)
86        d6 = self.down6(d5)
87        d7 = self.down7(d6)
88        d8 = self.down8(d7)
89
90        # Decoder with skip connections
91        u1 = self.up1(d8, d7)
92        u2 = self.up2(u1, d6)
93        u3 = self.up3(u2, d5)
94        u4 = self.up4(u3, d4)
95        u5 = self.up5(u4, d3)
96        u6 = self.up6(u5, d2)
97        u7 = self.up7(u6, d1)
98
99        return self.final(u7)
100
101
102class PatchGANDiscriminator(nn.Module):
103    """
104    PatchGAN Discriminator (70x70 receptive field).
105    Classifies each NxN patch as real or fake.
106    """
107
108    def __init__(self, in_channels: int = 6):  # Input + condition
109        super().__init__()
110
111        def discriminator_block(in_ch, out_ch, normalize=True):
112            layers = [nn.Conv2d(in_ch, out_ch, 4, stride=2, padding=1)]
113            if normalize:
114                layers.append(nn.InstanceNorm2d(out_ch))
115            layers.append(nn.LeakyReLU(0.2, inplace=True))
116            return layers
117
118        self.model = nn.Sequential(
119            *discriminator_block(in_channels, 64, normalize=False),
120            *discriminator_block(64, 128),
121            *discriminator_block(128, 256),
122            *discriminator_block(256, 512),
123            nn.Conv2d(512, 1, 4, padding=1)  # PatchGAN output
124        )
125
126    def forward(self, img_A: torch.Tensor,
127                img_B: torch.Tensor) -> torch.Tensor:
128        """
129        Args:
130            img_A: Condition image (e.g., semantic map)
131            img_B: Target or generated image
132
133        Returns:
134            Patch-wise predictions
135        """
136        # Concatenate condition and target
137        combined = torch.cat([img_A, img_B], dim=1)
138        return self.model(combined)
139
140
141class Pix2PixTrainer:
142    """Complete Pix2Pix training pipeline."""
143
144    def __init__(self, lambda_l1: float = 100.0, device: str = 'cuda'):
145        self.device = device if torch.cuda.is_available() else 'cpu'
146        self.lambda_l1 = lambda_l1
147
148        self.generator = Pix2PixGenerator().to(self.device)
149        self.discriminator = PatchGANDiscriminator().to(self.device)
150
151        self.opt_g = torch.optim.Adam(
152            self.generator.parameters(), lr=0.0002, betas=(0.5, 0.999)
153        )
154        self.opt_d = torch.optim.Adam(
155            self.discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999)
156        )
157
158        self.criterion_gan = nn.BCEWithLogitsLoss()
159        self.criterion_l1 = nn.L1Loss()
160
161    def train_step(self, condition: torch.Tensor,
162                   target: torch.Tensor) -> dict:
163        """
164        Single training step.
165
166        Args:
167            condition: Input condition (e.g., edge map)
168            target: Ground truth target image
169        """
170        condition = condition.to(self.device)
171        target = target.to(self.device)
172
173        # Generate fake image
174        fake = self.generator(condition)
175
176        # ---------------------
177        # Train Discriminator
178        # ---------------------
179        self.opt_d.zero_grad()
180
181        # Real pair
182        pred_real = self.discriminator(condition, target)
183        loss_d_real = self.criterion_gan(
184            pred_real, torch.ones_like(pred_real)
185        )
186
187        # Fake pair
188        pred_fake = self.discriminator(condition, fake.detach())
189        loss_d_fake = self.criterion_gan(
190            pred_fake, torch.zeros_like(pred_fake)
191        )
192
193        loss_d = (loss_d_real + loss_d_fake) * 0.5
194        loss_d.backward()
195        self.opt_d.step()
196
197        # ---------------------
198        # Train Generator
199        # ---------------------
200        self.opt_g.zero_grad()
201
202        pred_fake = self.discriminator(condition, fake)
203
204        # GAN loss
205        loss_g_gan = self.criterion_gan(
206            pred_fake, torch.ones_like(pred_fake)
207        )
208
209        # L1 loss (reconstruction)
210        loss_g_l1 = self.criterion_l1(fake, target) * self.lambda_l1
211
212        loss_g = loss_g_gan + loss_g_l1
213        loss_g.backward()
214        self.opt_g.step()
215
216        return {
217            'loss_d': loss_d.item(),
218            'loss_g': loss_g.item(),
219            'loss_g_gan': loss_g_gan.item(),
220            'loss_g_l1': loss_g_l1.item()
221        }
222
223    @torch.no_grad()
224    def translate(self, condition: torch.Tensor) -> torch.Tensor:
225        """Generate target from condition."""
226        self.generator.eval()
227        condition = condition.to(self.device)
228        output = self.generator(condition)
229        self.generator.train()
230        return output
231
232# Example
233pix2pix = Pix2PixTrainer()
234print("Pix2Pix initialized")
235print(f"Generator params: {sum(p.numel() for p in pix2pix.generator.parameters()):,}")
236print(f"Discriminator params: {sum(p.numel() for p in pix2pix.discriminator.parameters()):,}")