update utils.py
Browse files
utils.py
CHANGED
@@ -2,14 +2,19 @@ import numpy as np
|
|
2 |
import torch
|
3 |
from huggan.pytorch.lightweight_gan.lightweight_gan import LightweightGAN
|
4 |
|
|
|
|
|
5 |
def carga_modelo(model_name="ceyda/butterfly_cropped_uniq1K_512", model_version=None):
|
6 |
gan = LightweightGAN.from_pretrained(model_name, version=model_version)
|
7 |
gan.eval()
|
8 |
return gan
|
9 |
|
10 |
-
|
|
|
|
|
11 |
with torch.no_grad():
|
12 |
ims = gan.G(torch.randn(batch_size, gan.latent_dim)).clamp_(0.0, 1.0) * 255
|
13 |
-
ims = ims.permute(0,2,3,1).detach().cpu().numpy().astype(np.uint8)
|
14 |
return ims
|
15 |
|
|
|
|
2 |
import torch
|
3 |
from huggan.pytorch.lightweight_gan.lightweight_gan import LightweightGAN
|
4 |
|
5 |
+
|
6 |
+
## Cargamos el modelo desde el Hub de Hugging Face
|
7 |
def carga_modelo(model_name="ceyda/butterfly_cropped_uniq1K_512", model_version=None):
|
8 |
gan = LightweightGAN.from_pretrained(model_name, version=model_version)
|
9 |
gan.eval()
|
10 |
return gan
|
11 |
|
12 |
+
|
13 |
+
## Usamos el modelo GAN para generar imágenes
|
14 |
+
def genera(gan, batch_size=1):
|
15 |
with torch.no_grad():
|
16 |
ims = gan.G(torch.randn(batch_size, gan.latent_dim)).clamp_(0.0, 1.0) * 255
|
17 |
+
ims = ims.permute(0, 2, 3, 1).detach().cpu().numpy().astype(np.uint8)
|
18 |
return ims
|
19 |
|
20 |
+
|