Spaces:
Sleeping
Sleeping
Fix #8 app.py
Browse files
app.py
CHANGED
@@ -167,8 +167,9 @@ class Generator(nn.Module):
|
|
167 |
|
168 |
# FUNCIÓN PARA CARGAR EL MODELO
|
169 |
def load_pretrained_model(ckpt_path, img_size=256, style_dim=64, num_domains=3, device='cpu'):
|
|
|
170 |
G = Generator(img_size, style_dim).to(device)
|
171 |
-
M = MappingNetwork(
|
172 |
S = StyleEncoder(img_size, style_dim, num_domains).to(device)
|
173 |
checkpoint = torch.load(ckpt_path, map_location=device)
|
174 |
G.load_state_dict(checkpoint['generator'])
|
|
|
167 |
|
168 |
# FUNCIÓN PARA CARGAR EL MODELO
|
169 |
def load_pretrained_model(ckpt_path, img_size=256, style_dim=64, num_domains=3, device='cpu'):
|
170 |
+
latent_dim_for_mapping = 13
|
171 |
G = Generator(img_size, style_dim).to(device)
|
172 |
+
M = MappingNetwork(latent_dim_for_mapping, style_dim, num_domains).to(device)
|
173 |
S = StyleEncoder(img_size, style_dim, num_domains).to(device)
|
174 |
checkpoint = torch.load(ckpt_path, map_location=device)
|
175 |
G.load_state_dict(checkpoint['generator'])
|