Spaces:
Sleeping
Sleeping
File size: 2,294 Bytes
f379395 aa0a746 8abe14f b934874 f379395 aa0a746 f379395 aa0a746 f379395 aa0a746 f379395 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import json
import numpy as np
import torch
from huggan.pytorch.lightweight_gan.lightweight_gan import LightweightGAN
from huggingface_hub import hf_hub_download
CONFIG_NAME = "config.json"
revision = None
cache_dir = None
force_download = False
proxies = None
resume_download = False
local_files_only = False
token = None
def load_model(model_name="ceyda/butterfly_cropped_uniq1K_512"):
"""
Loads a pre-trained LightweightGAN model from Hugging Face Model Hub.
Args:
model_name (str): The name of the pre-trained model to load. Defaults to "ceyda/butterfly_cropped_uniq1K_512".
model_version (str): The version of the pre-trained model to load. Defaults to None.
Returns:
LightweightGAN: The loaded pre-trained model.
"""
# Load the config
config_file = hf_hub_download(
repo_id=str(model_name),
filename=CONFIG_NAME,
revision=revision,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
resume_download=resume_download,
token=token,
local_files_only=local_files_only,
)
with open(config_file, "r", encoding="utf-8") as f:
config = json.load(f)
# Call the _from_pretrained with all the needed arguments
gan = LightweightGAN(latent_dim=256, image_size=512)
gan = gan._from_pretrained(
model_id=str(model_name),
revision=revision,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
resume_download=resume_download,
local_files_only=local_files_only,
token=token,
use_auth_token=False,
config=config, # usually in **model_kwargs
)
gan.eval()
return gan
def generate(gan, batch_size=1):
"""
Generates images using the given GAN model.
Args:
gan (nn.Module): The GAN model to use for generating images.
batch_size (int, optional): The number of images to generate in each batch. Defaults to 1.
Returns:
numpy.ndarray: A numpy array of generated images.
"""
with torch.no_grad():
ims = gan.G(torch.randn(batch_size, gan.latent_dim)).clamp_(0.0, 1.0) * 255
ims = ims.permute(0, 2, 3, 1).detach().cpu().numpy().astype(np.uint8)
return ims
|