Spaces:
Build error
Build error
| # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved | |
| import os, sys | |
| os.system('pip install -r requirements.txt') | |
| import gradio as gr | |
| import numpy as np | |
| import dnnlib | |
| import time | |
| import legacy | |
| import torch | |
| import glob | |
| import cv2 | |
| from torch_utils import misc | |
| from renderer import Renderer | |
| from training.networks import Generator | |
| from huggingface_hub import hf_hub_download | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| port = int(sys.argv[1]) if len(sys.argv) > 1 else 21111 | |
| model_lists = { | |
| 'ffhq-512x512-basic': dict(repo_id='facebook/stylenerf-ffhq-config-basic', filename='ffhq_512.pkl'), | |
| 'ffhq-512x512-cc': dict(repo_id='facebook/stylenerf-ffhq-config-basic', filename='ffhq_512_cc.pkl'), | |
| 'ffhq-256x256-basic': dict(repo_id='facebook/stylenerf-ffhq-config-basic', filename='ffhq_256.pkl'), | |
| 'ffhq-1024x1024-basic': dict(repo_id='facebook/stylenerf-ffhq-config-basic', filename='ffhq_1024.pkl'), | |
| } | |
| model_names = [name for name in model_lists] | |
| def set_random_seed(seed): | |
| torch.manual_seed(seed) | |
| np.random.seed(seed) | |
| def get_camera_traj(model, pitch, yaw, fov=12, batch_size=1, model_name=None): | |
| gen = model.synthesis | |
| range_u, range_v = gen.C.range_u, gen.C.range_v | |
| if not (('car' in model_name) or ('Car' in model_name)): # TODO: hack, better option? | |
| yaw, pitch = 0.5 * yaw, 0.3 * pitch | |
| pitch = pitch + np.pi/2 | |
| u = (yaw - range_u[0]) / (range_u[1] - range_u[0]) | |
| v = (pitch - range_v[0]) / (range_v[1] - range_v[0]) | |
| else: | |
| u = (yaw + 1) / 2 | |
| v = (pitch + 1) / 2 | |
| cam = gen.get_camera(batch_size=batch_size, mode=[u, v, 0.5], device=device, fov=fov) | |
| return cam | |
| def check_name(model_name): | |
| """Gets model by name.""" | |
| if model_name in model_lists: | |
| network_pkl = hf_hub_download(**model_lists[model_name]) | |
| else: | |
| if os.path.isdir(model_name): | |
| network_pkl = sorted(glob.glob(model_name + '/*.pkl'))[-1] | |
| else: | |
| network_pkl = model_name | |
| return network_pkl | |
| def get_model(network_pkl, render_option=None): | |
| print('Loading networks from "%s"...' % network_pkl) | |
| with dnnlib.util.open_url(network_pkl) as f: | |
| network = legacy.load_network_pkl(f) | |
| G = network['G_ema'].to(device) # type: ignore | |
| with torch.no_grad(): | |
| G2 = Generator(*G.init_args, **G.init_kwargs).to(device) | |
| misc.copy_params_and_buffers(G, G2, require_all=False) | |
| print('compile and go through the initial image') | |
| G2 = G2.eval() | |
| init_z = torch.from_numpy(np.random.RandomState(0).rand(1, G2.z_dim)).to(device) | |
| init_cam = get_camera_traj(G2, 0, 0, model_name=network_pkl) | |
| dummy = G2(z=init_z, c=None, camera_matrices=init_cam, render_option=render_option, theta=0) | |
| res = dummy['img'].shape[-1] | |
| imgs = np.zeros((res, res//2, 3)) | |
| return G2, res, imgs | |
| global_states = list(get_model(check_name(model_names[0]))) | |
| wss = [None, None] | |
| def proc_seed(history, seed): | |
| if isinstance(seed, str): | |
| seed = 0 | |
| else: | |
| seed = int(seed) | |
| def f_synthesis(model_name, model_find, render_option, early, trunc, seed1, seed2, mix1, mix2, yaw, pitch, roll, fov, history): | |
| history = history or {} | |
| seeds = [] | |
| trunc = trunc / 100 | |
| if model_find != "": | |
| model_name = model_find | |
| model_name = check_name(model_name) | |
| if model_name != history.get("model_name", None): | |
| model, res, imgs = get_model(model_name, render_option) | |
| global_states[0] = model | |
| global_states[1] = res | |
| global_states[2] = imgs | |
| model, res, imgs = global_states | |
| for idx, seed in enumerate([seed1, seed2]): | |
| if isinstance(seed, str): | |
| seed = 0 | |
| else: | |
| seed = int(seed) | |
| if (seed != history.get(f'seed{idx}', -1)) or \ | |
| (model_name != history.get("model_name", None)) or \ | |
| (trunc != history.get("trunc", 0.7)) or \ | |
| (wss[idx] is None): | |
| print(f'use seed {seed}') | |
| set_random_seed(seed) | |
| z = torch.from_numpy(np.random.RandomState(int(seed)).randn(1, model.z_dim).astype('float32')).to(device) | |
| ws = model.mapping(z=z, c=None, truncation_psi=trunc) | |
| img = model.get_final_output(styles=ws, camera_matrices=get_camera_traj(model, 0, 0, model_name=model_name), render_option=render_option) | |
| ws = ws.detach().cpu().numpy() | |
| img = img[0].permute(1,2,0).detach().cpu().numpy() | |
| imgs[idx * res // 2: (1 + idx) * res // 2] = cv2.resize( | |
| np.asarray(img).clip(-1, 1) * 0.5 + 0.5, | |
| (res//2, res//2), cv2.INTER_AREA) | |
| wss[idx] = ws | |
| else: | |
| seed = history[f'seed{idx}'] | |
| seeds += [seed] | |
| history[f'seed{idx}'] = seed | |
| history['trunc'] = trunc | |
| history['model_name'] = model_name | |
| set_random_seed(sum(seeds)) | |
| # style mixing (?) | |
| ws1, ws2 = [torch.from_numpy(ws).to(device) for ws in wss] | |
| ws = ws1.clone() | |
| ws[:, :8] = ws1[:, :8] * mix1 + ws2[:, :8] * (1 - mix1) | |
| ws[:, 8:] = ws1[:, 8:] * mix2 + ws2[:, 8:] * (1 - mix2) | |
| # set visualization for other types of inputs. | |
| if early == 'Normal Map': | |
| render_option += ',normal,early' | |
| elif early == 'Gradient Map': | |
| render_option += ',gradient,early' | |
| start_t = time.time() | |
| with torch.no_grad(): | |
| cam = get_camera_traj(model, pitch, yaw, fov, model_name=model_name) | |
| image = model.get_final_output( | |
| styles=ws, camera_matrices=cam, | |
| theta=roll * np.pi, | |
| render_option=render_option) | |
| end_t = time.time() | |
| image = image[0].permute(1,2,0).detach().cpu().numpy().clip(-1, 1) * 0.5 + 0.5 | |
| if imgs.shape[0] == image.shape[0]: | |
| image = np.concatenate([imgs, image], 1) | |
| else: | |
| a = image.shape[0] | |
| b = int(imgs.shape[1] / imgs.shape[0] * a) | |
| print(f'resize {a} {b} {image.shape} {imgs.shape}') | |
| image = np.concatenate([cv2.resize(imgs, (b, a), cv2.INTER_AREA), image], 1) | |
| print(f'rendering time = {end_t-start_t:.4f}s') | |
| image = (image * 255).astype('uint8') | |
| return image, history | |
| model_name = gr.inputs.Dropdown(model_names) | |
| model_find = gr.inputs.Textbox(label="Checkpoint path (folder or .pkl file)", default="") | |
| render_option = gr.inputs.Textbox(label="Additional rendering options", default='freeze_bg,steps:50') | |
| trunc = gr.inputs.Slider(default=70, maximum=100, minimum=0, label='Truncation trick (%)') | |
| seed1 = gr.inputs.Number(default=1, label="Random seed1") | |
| seed2 = gr.inputs.Number(default=9, label="Random seed2") | |
| mix1 = gr.inputs.Slider(minimum=0, maximum=1, default=0, label="Linear mixing ratio (geometry)") | |
| mix2 = gr.inputs.Slider(minimum=0, maximum=1, default=0, label="Linear mixing ratio (apparence)") | |
| early = gr.inputs.Radio(['None', 'Normal Map', 'Gradient Map'], default='None', label='Intermedia output') | |
| yaw = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="Yaw") | |
| pitch = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="Pitch") | |
| roll = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="Roll (optional, not suggested for basic config)") | |
| fov = gr.inputs.Slider(minimum=10, maximum=14, default=12, label="Fov") | |
| css = ".output-image, .input-image, .image-preview {height: 600px !important} " | |
| gr.Interface(fn=f_synthesis, | |
| inputs=[model_name, model_find, render_option, early, trunc, seed1, seed2, mix1, mix2, yaw, pitch, roll, fov, "state"], | |
| title="Interactive Web Demo for StyleNeRF (ICLR 2022)", | |
| description="StyleNeRF: A Style-based 3D-Aware Generator for High-resolution Image Synthesis. Currently the demo runs on CPU only.", | |
| outputs=["image", "state"], | |
| layout='unaligned', | |
| css=css, theme='dark-seafoam', | |
| live=True).launch(enable_queue=True) | |