Spaces:
Runtime error
Runtime error
import os | |
os.environ["USE_NATIVE"] = "1" | |
import math | |
import torch | |
import torchvision | |
import gradio as gr | |
from PIL import Image | |
import torchvision | |
from test_ddgan import load_model, sample | |
from model_configs import get_model_config | |
from subprocess import call | |
def download(filename): | |
return "models/" + filename | |
device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
cache = {} | |
def load(name): | |
if name in cache: | |
return cache[name] | |
else: | |
model_config, model_path = models[name] | |
print(model_config, model_path) | |
model = load_model(model_config, model_path, device=device) | |
cache[name] = model | |
return model | |
models = { | |
"diffusion_db_128ch_1timesteps_openclip_vith14": (get_model_config('ddgan_ddb_v2'), download('diffusion_db_128ch_1timesteps_openclip_vith14.th')), | |
"diffusion_db_192ch_2timesteps_openclip_vith14": (get_model_config('ddgan_ddb_v3'), download('diffusion_db_192ch_2timesteps_openclip_vith14.th')), | |
} | |
default = "diffusion_db_128ch_1timesteps_openclip_vith14" | |
def gen(md, model_name, md2, text, seed, nb_samples, width, height): | |
print("load ", model_name) | |
model = load(model_name) | |
print(model) | |
torch.manual_seed(int(seed)) | |
nb_samples = int(nb_samples) | |
height = int(height) | |
width = int(width) | |
with torch.no_grad(): | |
cond = model.text_encoder([text]*nb_samples) | |
if text == "": | |
cond[0].normal_() | |
cond[1].normal_() | |
cond[0][1:] = cond[0][0:1] | |
cond[1][1:] = cond[1][0:1] | |
x_init = torch.randn(nb_samples, 3, height, width).to(device) | |
print(x_init.shape) | |
fake_sample = sample(model, x_init=x_init, cond=cond) | |
fake_sample = (fake_sample + 1) / 2 | |
grid = torchvision.utils.make_grid(fake_sample, nrow=4) | |
grid = grid.permute(1, 2, 0).cpu().numpy() | |
grid = (grid*255).astype("uint8") | |
return Image.fromarray(grid) | |
text = """ | |
DDGAN | |
""" | |
iface = gr.Interface( | |
fn=gen, | |
inputs=[ | |
gr.Markdown(text), | |
# text caption | |
gr.Dropdown(list(models.keys()), value=default), | |
gr.Markdown("If text caption is empty, random CLIP embeddings will be used as input"), | |
gr.Textbox( | |
lines=1, | |
placeholder="Enter text caption here, or leave empty", | |
value="Painting of a hamster king with a crown and a cape in a magical forest." | |
), | |
gr.Number(value=0), # seed | |
gr.Number(value=4), # nb_samples | |
gr.Number(value=256), # width | |
gr.Number(value=256),# height | |
], | |
outputs="image" | |
) | |
iface.launch(debug=True) | |