Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| class Encoder(nn.Module): | |
| def __init__(self, latent_dim): | |
| super(Encoder, self).__init__() | |
| # A simple convolutional encoder for demonstration | |
| self.conv = nn.Sequential( | |
| nn.Conv2d(3, 32, 4, 2, 1), # 64x64 -> 32x32 | |
| nn.ReLU(), | |
| nn.Conv2d(32, 64, 4, 2, 1), # 32x32 -> 16x16 | |
| nn.ReLU(), | |
| nn.Conv2d(64, 128, 4, 2, 1), # 16x16 -> 8x8 | |
| nn.ReLU(), | |
| nn.Flatten() | |
| ) | |
| # Assuming input image of 64x64 | |
| self.fc_mu = nn.Linear(128 * 8 * 8, latent_dim) | |
| self.fc_logvar = nn.Linear(128 * 8 * 8, latent_dim) | |
| def forward(self, x): | |
| features = self.conv(x) | |
| mu = self.fc_mu(features) | |
| logvar = self.fc_logvar(features) | |
| return mu, logvar | |
| class Decoder(nn.Module): | |
| def __init__(self, latent_dim, condition_dim=1): | |
| super(Decoder, self).__init__() | |
| # The decoder takes the latent vector PLUS the age condition | |
| self.fc = nn.Linear(latent_dim + condition_dim, 128 * 8 * 8) | |
| self.deconv = nn.Sequential( | |
| nn.ConvTranspose2d(128, 64, 4, 2, 1), # 8x8 -> 16x16 | |
| nn.ReLU(), | |
| nn.ConvTranspose2d(64, 32, 4, 2, 1), # 16x16 -> 32x32 | |
| nn.ReLU(), | |
| nn.ConvTranspose2d(32, 3, 4, 2, 1), # 32x32 -> 64x64 | |
| nn.Sigmoid() # Output pixels between 0 and 1 | |
| ) | |
| def forward(self, z, age_condition): | |
| # Concatenate latent identity with age condition | |
| z_cond = torch.cat((z, age_condition), dim=1) | |
| hidden = self.fc(z_cond) | |
| hidden = hidden.view(-1, 128, 8, 8) | |
| out_img = self.deconv(hidden) | |
| return out_img | |
| class GAP_CVAE(nn.Module): | |
| def __init__(self, latent_dim=128): | |
| super(GAP_CVAE, self).__init__() | |
| self.encoder = Encoder(latent_dim) | |
| self.decoder = Decoder(latent_dim, condition_dim=1) | |
| def reparameterize(self, mu, logvar): | |
| std = torch.exp(0.5 * logvar) | |
| eps = torch.randn_like(std) | |
| return mu + eps * std | |
| def forward(self, x, age): | |
| mu, logvar = self.encoder(x) | |
| z = self.reparameterize(mu, logvar) | |
| reconstructed = self.decoder(z, age) | |
| return reconstructed, mu, logvar | |
| def simulate_age(self, x, target_age): | |
| """Used for inference when we want to change the age of an image""" | |
| device = next(self.parameters()).device | |
| x = x.to(device) | |
| target_age = target_age.to(device) | |
| # 1. Extract Identity Latent (mu) | |
| mu, _ = self.encoder(x) | |
| # 2. Decode with new target age | |
| projected_image = self.decoder(mu, target_age) | |
| return projected_image | |