import os import gradio as gr os.system("gdown https://drive.google.com/uc?id=1--h-4E5LSxe6VTp9rjJhtoILxxNH0X5n") # python 3.7 """Demo.""" import io import cv2 import warnings import numpy as np import torch from PIL import Image from models import build_model warnings.filterwarnings(action='ignore', category=UserWarning) def postprocess_image(image, min_val=-1.0, max_val=1.0): """Post-processes image to pixel range [0, 255] with dtype `uint8`. This function is particularly used to handle the results produced by deep models. NOTE: The input image is assumed to be with format `NCHW`, and the returned image will always be with format `NHWC`. Args: image: The input image for post-processing. min_val: Expected minimum value of the input image. max_val: Expected maximum value of the input image. Returns: The post-processed image. """ assert isinstance(image, np.ndarray) image = image.astype(np.float64) image = (image - min_val) / (max_val - min_val) * 255 image = np.clip(image + 0.5, 0, 255).astype(np.uint8) assert image.ndim == 4 and image.shape[1] in [1, 3, 4] return image.transpose(0, 2, 3, 1) def to_numpy(data): """Converts the input data to `numpy.ndarray`.""" if isinstance(data, (int, float)): return np.array(data) if isinstance(data, np.ndarray): return data if isinstance(data, torch.Tensor): return data.detach().cpu().numpy() raise TypeError(f'Not supported data type `{type(data)}` for ' f'converting to `numpy.ndarray`!') def linear_interpolate(latent_code, boundary, layer_index=None, start_distance=-10.0, end_distance=10.0, steps=7): """Interpolate between the latent code and boundary.""" assert (len(latent_code.shape) == 3 and len(boundary.shape) == 3 and latent_code.shape[0] == 1 and boundary.shape[0] == 1 and latent_code.shape[1] == boundary.shape[1]) linspace = np.linspace(start_distance, end_distance, steps) linspace = linspace.reshape([-1, 1, 1]).astype(np.float32) inter_code = linspace * boundary is_manipulatable = np.zeros(inter_code.shape, dtype=bool) is_manipulatable[:, layer_index, :] = True mani_code = np.where(is_manipulatable, latent_code+inter_code, latent_code) return mani_code def imshow(images, col, viz_size=256): """Shows images in one figure.""" num, height, width, channels = images.shape assert num % col == 0 row = num // col fused_image = np.zeros((viz_size*row, viz_size*col, channels), dtype=np.uint8) for idx, image in enumerate(images): i, j = divmod(idx, col) y = i * viz_size x = j * viz_size if height != viz_size or width != viz_size: image = cv2.resize(image, (viz_size, viz_size)) fused_image[y:y + viz_size, x:x + viz_size] = image fused_image = np.asarray(fused_image, dtype=np.uint8) data = io.BytesIO() if channels == 4: Image.fromarray(fused_image).save(data, 'png') elif channels == 3: Image.fromarray(fused_image).save(data, 'jpeg') else: raise ValueError('Image channel error') im_data = data.getvalue() image = Image.open(io.BytesIO(im_data)) return image print('Building generator') checkpoint_path='stylegan2-ffhq-config-f-1024x1024.pth' config = dict(model_type='StyleGAN2Generator', resolution=1024, w_dim=512, fmaps_base=int(1 * (32 << 10)), fmaps_max=512,) generator = build_model(**config) print(f'Loading checkpoint from `{checkpoint_path}` ...') checkpoint = torch.load(checkpoint_path, map_location='cpu')['models'] if 'generator_smooth' in checkpoint: generator.load_state_dict(checkpoint['generator_smooth']) else: generator.load_state_dict(checkpoint['generator']) generator = generator.eval().cpu() print('Finish loading checkpoint.') print('Loading boundaries') ATTRS = ['eyebrows', 'eyesize', 'gaze_direction', 'nose_length', 'mouth', 'lipstick'] boundaries = {} for attr in ATTRS: boundary_path = os.path.join(f'directions/ffhq/stylegan2/{attr}.npy') boundary = np.load(boundary_path) boundaries[attr] = boundary print('Generator and boundaries are ready.') def inference(num_of_image,seed,trunc_psi,eyebrows,eyesize,gaze_direction,nose_length,mouth,lipstick): print('Sampling latent codes with given seed.') num_of_image = num_of_image #@param {type:"slider", min:1, max:8, step:1} seed = seed #@param {type:"slider", min:0, max:10000, step:1} trunc_psi = trunc_psi #@param {type:"slider", min:0, max:1, step:0.1} trunc_layers = 8 np.random.seed(seed) latent_z = np.random.randn(num_of_image, generator.z_dim) latent_z = torch.from_numpy(latent_z.astype(np.float32)) latent_z = latent_z.cpu() wp = generator.mapping(latent_z, None)['wp'] if trunc_psi < 1.0: w_avg = generator.w_avg w_avg = w_avg.reshape(1, -1, generator.w_dim)[:, :trunc_layers] wp[:, :trunc_layers] = w_avg.lerp(wp[:, :trunc_layers], trunc_psi) with torch.no_grad(): images_ori = generator.synthesis(wp)['image'] images_ori = postprocess_image(to_numpy(images_ori)) print('Original images are shown as belows.') imshow(images_ori, col=images_ori.shape[0]) latent_wp = to_numpy(wp) eyebrows = eyebrows #@param {type:"slider", min:-12.0, max:12.0, step:2} eyesize = eyesize #@param {type:"slider", min:-12.0, max:12.0, step:2} gaze_direction = gaze_direction #@param {type:"slider", min:-12.0, max:12.0, step:2} nose_length = nose_length #@param {type:"slider", min:-12.0, max:12.0, step:2} mouth = mouth #@param {type:"slider", min:-12.0, max:12.0, step:2} lipstick = lipstick #@param {type:"slider", min:-12.0, max:12.0, step:2} new_codes = latent_wp.copy() for attr_name in ATTRS: if attr_name in ['eyebrows', 'lipstick']: layers_idx = [8,9,10,11] else: layers_idx = [4,5,6,7] step = eval(attr_name) direction = boundaries[attr_name] direction = np.tile(direction, [1, generator.num_layers, 1]) new_codes[:, layers_idx, :] += direction[:, layers_idx, :] * step new_codes = torch.from_numpy(new_codes.astype(np.float32)).cpu() with torch.no_grad(): images_mani = generator.synthesis(new_codes)['image'] images_mani = postprocess_image(to_numpy(images_mani)) return imshow(images_mani, col=images_mani.shape[0]) title = "resefa" description = "## Gradio Demo for [Region-Based Semantic Factorization in GANs](https://github.com/zhujiapeng/resefa)" gr.Interface(inference,[gr.Slider(1, 3, value=1,label="num_of_image",step=1), gr.Slider(0, 10000, value=210,label="seed",step=1), gr.Slider(0, 1, value=0.7,step=0.1,label="truncation psi"), gr.Slider(-12, 12, value=0,label="eyebrows",step=1), gr.Slider(-12, 12, value=0,label="eyesize",step=1), gr.Slider(-12, 12, value=0,label="gaze direction",step=1), gr.Slider(-12, 12, value=0,label="nose_length",step=1), gr.Slider(-12, 12, value=0,label="mouth",step=1), gr.Slider(-12, 12, value=0,label="lipstick",step=1), ],gr.Image(type="pil"),title=title,description=description).launch()