cds006/Progan
0
1"""
2Progressive GAN Generator and Discriminator
3"""
4
5import torch
6import torch.nn as nn
7import torch.nn.functional as F
8from layers import (
9 EqualizedConv2d, GeneratorBlock, DiscriminatorBlock,
10 PixelNorm, MinibatchStddev
11)
12
13
14class Generator(nn.Module):
15 """Progressive GAN Generator"""
16 def __init__(self, latent_dim=512, max_resolution=1024, feature_maps=None):
17 super().__init__()
18
19 self.latent_dim = latent_dim
20 self.max_resolution = max_resolution
21
22 if feature_maps is None:
23 feature_maps = {
24 4: 512, 8: 512, 16: 512, 32: 512,
25 64: 256, 128: 128, 256: 64, 512: 32, 1024: 16
26 }
27 self.feature_maps = feature_maps
28
29 # Build progressive blocks
30 self.blocks = nn.ModuleDict()
31 self.to_rgb = nn.ModuleDict()
32
33 # Initial block (4x4)
34 self.blocks['4x4'] = GeneratorBlock(latent_dim, feature_maps[4], initial_block=True)
35 self.to_rgb['4x4'] = EqualizedConv2d(feature_maps[4], 3, kernel_size=1)
36
37 # Progressive blocks
38 resolutions = [8, 16, 32, 64, 128, 256, 512, 1024]
39 for res in resolutions:
40 if res > max_resolution:
41 break
42
43 prev_res = res // 2
44 in_ch = feature_maps[prev_res]
45 out_ch = feature_maps[res]
46
47 self.blocks[f'{res}x{res}'] = GeneratorBlock(in_ch, out_ch)
48 self.to_rgb[f'{res}x{res}'] = EqualizedConv2d(out_ch, 3, kernel_size=1)
49
50 self.current_resolution = 4
51 self.alpha = 1.0 # For fade-in
52
53 def forward(self, z, resolution=None, alpha=None):
54 """
55 Args:
56 z: Latent vector [B, latent_dim, 1, 1]
57 resolution: Target resolution (if None, uses current_resolution)
58 alpha: Fade-in parameter (if None, uses self.alpha)
59 """
60 if resolution is None:
61 resolution = self.current_resolution
62 if alpha is None:
63 alpha = self.alpha
64
65 # Initial block
66 x = self.blocks['4x4'](z)
67
68 # Progressive blocks
69 res = 4
70 x_prev = None
71 while res < resolution:
72 # Save previous feature map before processing next block
73 if res == resolution//2:
74 x_prev = x
75
76 next_res = res * 2
77 x = self.blocks[f'{next_res}x{next_res}'](x)
78 res = next_res
79
80 # Handle fade-in during transition
81 if alpha < 1.0 and resolution > 4 and x_prev is not None:
82 # Get output from previous resolution
83 prev_res = resolution // 2
84
85 # RGB conversion for current resolution
86 rgb_new = self.to_rgb[f'{resolution}x{resolution}'](x)
87
88 # RGB conversion for previous resolution, then upsample
89 prev_key = f'{prev_res}x{prev_res}'
90 if prev_key in self.to_rgb:
91
92 rgb_prev = self.to_rgb[prev_key](x_prev)
93 rgb_prev = F.interpolate(rgb_prev, scale_factor=2, mode='nearest')
94
95 # Blend based on alpha
96 rgb = alpha * rgb_new + (1 - alpha) * rgb_prev
97 else:
98 rgb = rgb_new
99 else:
100 rgb = self.to_rgb[f'{resolution}x{resolution}'](x)
101
102 return torch.tanh(rgb)
103
104 def grow(self, new_resolution):
105 """Grow to new resolution"""
106 self.current_resolution = new_resolution
107 self.alpha = 0.0
108
109
110class Discriminator(nn.Module):
111 """Progressive GAN Discriminator"""
112 def __init__(self, max_resolution=1024, feature_maps=None):
113 super().__init__()
114
115 self.max_resolution = max_resolution
116
117 if feature_maps is None:
118 feature_maps = {
119 4: 512, 8: 512, 16: 512, 32: 512,
120 64: 256, 128: 128, 256: 64, 512: 32, 1024: 16
121 }
122 self.feature_maps = feature_maps
123
124 # Build progressive blocks (in reverse order)
125 self.blocks = nn.ModuleDict()
126 self.from_rgb = nn.ModuleDict()
127
128 # Build all resolution blocks
129 resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
130 for i, res in enumerate(resolutions):
131 if res > max_resolution:
132 break
133
134 # From RGB layer
135 self.from_rgb[f'{res}x{res}'] = EqualizedConv2d(3, feature_maps[res], kernel_size=1)
136
137 # Discriminator block
138 if res == 4:
139 self.blocks[f'{res}x{res}'] = DiscriminatorBlock(
140 feature_maps[res], feature_maps[res], final_block=True
141 )
142 else:
143 next_res = res // 2
144 self.blocks[f'{res}x{res}'] = DiscriminatorBlock(
145 feature_maps[res], feature_maps[next_res]
146 )
147
148 self.current_resolution = 4
149 self.alpha = 1.0
150
151 def forward(self, x, resolution=None, alpha=None):
152 """
153 Args:
154 x: Input image [B, 3, H, W]
155 resolution: Current resolution (if None, uses current_resolution)
156 alpha: Fade-in parameter (if None, uses self.alpha)
157 """
158 if resolution is None:
159 resolution = self.current_resolution
160 if alpha is None:
161 alpha = self.alpha
162
163 # Handle fade-in during transition
164 if alpha < 1.0 and resolution > 4:
165 # New path: from_rgb at current resolution
166 x_new = self.from_rgb[f'{resolution}x{resolution}'](x)
167
168 # Old path: downsample then from_rgb
169 prev_res = resolution // 2
170 x_old = F.avg_pool2d(x, kernel_size=2, stride=2)
171 x_old = self.from_rgb[f'{prev_res}x{prev_res}'](x_old)
172
173 # Process new path through one block
174 x_new = self.blocks[f'{resolution}x{resolution}'](x_new)
175
176 # Blend
177 x = alpha * x_new + (1 - alpha) * x_old
178
179 # Continue from previous resolution
180 res = prev_res
181 else:
182 # Standard path
183 x = self.from_rgb[f'{resolution}x{resolution}'](x)
184 res = resolution
185
186 # Process through remaining blocks
187 while res >= 4:
188 x = self.blocks[f'{res}x{res}'](x)
189 if res == 4:
190 break
191 res = res // 2
192
193 return x
194
195 def grow(self, new_resolution):
196 """Grow to new resolution"""
197 self.current_resolution = new_resolution
198 self.alpha = 0.0
199
200
201class ProgressiveGAN(nn.Module):
202 """Combined Progressive GAN model"""
203 def __init__(self, latent_dim=512, max_resolution=1024, feature_maps=None):
204 super().__init__()
205
206 self.generator = Generator(latent_dim, max_resolution, feature_maps)
207 self.discriminator = Discriminator(max_resolution, feature_maps)
208
209 self.latent_dim = latent_dim
210 self.current_resolution = 4
211
212 def generate(self, batch_size, device='cuda'):
213 """Generate random images"""
214 z = torch.randn(batch_size, self.latent_dim, 1, 1, device=device)
215 return self.generator(z)
216
217 def grow(self, new_resolution):
218 """Grow both networks to new resolution"""
219 self.generator.grow(new_resolution)
220 self.discriminator.grow(new_resolution)
221 self.current_resolution = new_resolution
222
223 def set_alpha(self, alpha):
224 """Set fade-in alpha for both networks"""
225 self.generator.alpha = alpha
226 self.discriminator.alpha = alpha