def __init__(self, resolution, latent_size=512, dlatent_size=512,
truncation_psi=0.7, truncation_cutoff=8, dlatent_avg_beta=0.995,
style_mixing_prob=0.9, **kwargs):
"""
# Style-based generator used in the StyleGAN paper.
# Composed of two sub-networks (G_mapping and G_synthesis).
:param resolution:
:param latent_size:
:param dlatent_size:
:param truncation_psi: Style strength multiplier for the truncation trick. None = disable.
:param truncation_cutoff: Number of layers for which to apply the truncation trick. None = disable.
:param dlatent_avg_beta: Decay for tracking the moving average of W during training. None = disable.
:param style_mixing_prob: Probability of mixing styles during training. None = disable.
:param kwargs: Arguments for sub-networks (G_mapping and G_synthesis).
"""
super(Generator, self).__init__()
self.style_mixing_prob = style_mixing_prob
self.style_dim = dlatent_size
# Setup components.
self.num_layers = (int(np.log2(resolution)) - 1) * 2
self.g_mapping = GMapping(latent_size, dlatent_size, dlatent_broadcast=self.num_layers, **kwargs)
self.g_synthesis = GSynthesis(resolution=resolution, **kwargs)
if truncation_psi > 0:
self.truncation = Truncation(avg_latent=torch.zeros(dlatent_size),
max_layer=truncation_cutoff,
threshold=truncation_psi,
beta=dlatent_avg_beta)
else:
self.truncation = None
class Generator(nn.Module):
def __init__(self, resolution, latent_size=512, dlatent_size=512,
truncation_psi=0.7, truncation_cutoff=8, dlatent_avg_beta=0.995,
style_mixing_prob=0.9, **kwargs):
"""
# Style-based generator used in the StyleGAN paper.
# Composed of two sub-networks (G_mapping and G_synthesis).
:param resolution:
:param latent_size:
:param dlatent_size:
:param truncation_psi: Style strength multiplier for the truncation trick. None = disable.
:param truncation_cutoff: Number of layers for which to apply the truncation trick. None = disable.
:param dlatent_avg_beta: Decay for tracking the moving average of W during training. None = disable.
:param style_mixing_prob: Probability of mixing styles during training. None = disable.
:param kwargs: Arguments for sub-networks (G_mapping and G_synthesis).
"""
super(Generator, self).__init__()
self.style_mixing_prob = style_mixing_prob
self.style_dim = dlatent_size
# Setup components.
self.num_layers = (int(np.log2(resolution)) - 1) * 2
self.g_mapping = GMapping(latent_size, dlatent_size, dlatent_broadcast=self.num_layers, **kwargs)
self.g_synthesis = GSynthesis(resolution=resolution, **kwargs)
if truncation_psi > 0:
self.truncation = Truncation(avg_latent=torch.zeros(dlatent_size),
max_layer=truncation_cutoff,
threshold=truncation_psi,
beta=dlatent_avg_beta)
else:
self.truncation = None
def forward(self, latents_in, depth, alpha, labels_in=None):
"""
:param latents_in: First input: Latent vectors (Z) [mini_batch, latent_size].
:param depth: current depth from where output is required
:param alpha: value of alpha for fade-in effect
:param labels_in: Second input: Conditioning labels [mini_batch, label_size].
:return:
"""
dlatents_in = self.g_mapping(latents_in)
if self.training:
# Update moving average of W(dlatent).
# TODO
if self.truncation is not None:
self.truncation.update(dlatents_in[0, 0].detach())
# Perform style mixing regularization.
if self.style_mixing_prob is not None and self.style_mixing_prob > 0:
latents2 = torch.randn(latents_in.shape).to(latents_in.device)
dlatents2 = self.g_mapping(latents2)
layer_idx = torch.from_numpy(np.arange(self.num_layers)[np.newaxis, :, np.newaxis]).to(
latents_in.device)
cur_layers = 2 * (depth + 1)
mixing_cutoff = random.randint(1,
cur_layers) if random.random() < self.style_mixing_prob else cur_layers
dlatents_in = torch.where(layer_idx < mixing_cutoff, dlatents_in, dlatents2)
# Apply truncation trick.
if self.truncation is not None:
dlatents_in = self.truncation(dlatents_in)
fake_images = self.g_synthesis(dlatents_in, depth, alpha)
return fake_images
def mean_latent(self, n_latent):
latent_in = torch.randn(
n_latent, self.style_dim
)
latent = self.g_mapping(latent_in.cuda()).mean(0, keepdim=True)
return latent