File size: 2,004 Bytes
bd18f96 bb2f22e bd18f96 bb2f22e 733e8c1 bd18f96 d7fed5a bb2f22e bd18f96 |
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 |
import streamlit as st
import numpy as np
import torch
from huggingface_hub import hf_hub_download
import json
CONFIG_NAME = "config.json"
revision = None
cache_dir = None
force_download = False
proxies = None
resume_download = False
local_files_only = False
token = None
from huggan.pytorch.lightweight_gan.lightweight_gan import LightweightGAN
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 generation(gan, batch_size=1):
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 |