File size: 3,063 Bytes
e43bbf3 cd015de e43bbf3 c58bbe8 cd015de c58bbe8 |
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 80 81 |
import numpy as np
import pickle as pickle
import os
import sys
import wget
#os.system("git clone https://github.com/NVlabs/stylegan3")
#sys.path.append('./stylegan3')
model_names = {
'AFHQv2-512-R': 'stylegan3-r-afhqv2-512x512.pkl',
'FFHQ-1024-R': 'stylegan3-r-ffhq-1024x1024.pkl',
'FFHQ-U-256-R': 'stylegan3-r-ffhqu-256x256.pkl',
'FFHQ-U-1024-R': 'stylegan3-r-ffhqu-1024x1024.pkl',
'MetFaces-1024-R': 'stylegan3-r-metfaces-1024x1024.pkl',
'MetFaces-U-1024-R': 'stylegan3-r-metfacesu-1024x1024.pkl',
'AFHQv2-512-T': 'stylegan3-t-afhqv2-512x512.pkl',
'FFHQ-1024-T': 'stylegan3-t-ffhq-1024x1024.pkl',
'FFHQ-U-256-T': 'stylegan3-t-ffhqu-256x256.pkl',
'FFHQ-U-1024-T': 'stylegan3-t-ffhqu-1024x1024.pkl',
'MetFaces-1024-T': 'stylegan3-t-metfaces-1024x1024.pkl',
'MetFaces-U-1024-T': 'stylegan3-t-metfacesu-1024x1024.pkl',
}
model_dict = {
name: file_name
for name, file_name in model_names.items()
}
def fetch_model(url_or_path):
basename = os.path.basename(url_or_path)
if os.path.exists(basename):
return basename
else:
wget.download(url_or_path)
return basename
def load_model(file_name: str, device: torch.device):
#path = torch.hub.download_url_to_file('https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/'+f'{file_name}',
# f'{file_name}')
base_url = "https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/"
network_url = base_url + f'{file_name}'
local_path = '/content/'f'{file_name}'
print(local_path)
with open(fetch_model(network_url), 'rb') as f:
model = pickle.load(f)['G_ema']
model.eval()
model.to(device)
with torch.inference_mode():
z = torch.zeros((1, model.z_dim)).to(device)
label = torch.zeros([1, model.c_dim], device=device)
model(z, label)
return model
def generate_image(model_name: str, seed: int, truncation_psi: float):
device = 'cuda'
model = model_dict[model_name]
model = load_model(model, device)
seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max))
z = torch.from_numpy(np.random.RandomState(seed).randn(1, model.z_dim)).to(device)
label = torch.zeros([1, model.c_dim], device=device)
out = model(z, label, truncation_psi=truncation_psi)
out = (out.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
return out[0].cpu().numpy()
import gradio as gr
gr.Interface(
generate_image,
[
gr.inputs.Radio(list(model_names.keys()),
type='value',
default='FFHQ-1024-R',
label='Model'),
gr.inputs.Number(default=0, label='Seed'),
gr.inputs.Slider(
0, 2, step=0.05, default=0.7, label='Truncation psi')
],
gr.outputs.Image(type='numpy', label='Output')
).launch(debug=True
) |