Upload 8 files

utils/common.py

@@ -0,0 +1,87 @@
+import cv2
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+
+
+# Log images
+def log_input_image(x, opts):
+	if opts.label_nc == 0:
+		return tensor2im(x)
+	elif opts.label_nc == 1:
+		return tensor2sketch(x)
+	else:
+		return tensor2map(x)
+
+
+def tensor2im(var):
+	var = var.cpu().detach().transpose(0, 2).transpose(0, 1).numpy()
+	var = ((var + 1) / 2)
+	var[var < 0] = 0
+	var[var > 1] = 1
+	var = var * 255
+	return Image.fromarray(var.astype('uint8'))
+
+
+def tensor2map(var):
+	mask = np.argmax(var.data.cpu().numpy(), axis=0)
+	colors = get_colors()
+	mask_image = np.ones(shape=(mask.shape[0], mask.shape[1], 3))
+	for class_idx in np.unique(mask):
+		mask_image[mask == class_idx] = colors[class_idx]
+	mask_image = mask_image.astype('uint8')
+	return Image.fromarray(mask_image)
+
+
+def tensor2sketch(var):
+	im = var[0].cpu().detach().numpy()
+	im = cv2.cvtColor(im, cv2.COLOR_GRAY2BGR)
+	im = (im * 255).astype(np.uint8)
+	return Image.fromarray(im)
+
+
+# Visualization utils
+def get_colors():
+	# currently support up to 19 classes (for the celebs-hq-mask dataset)
+	colors = [[0, 0, 0], [204, 0, 0], [76, 153, 0], [204, 204, 0], [51, 51, 255], [204, 0, 204], [0, 255, 255],
+			  [255, 204, 204], [102, 51, 0], [255, 0, 0], [102, 204, 0], [255, 255, 0], [0, 0, 153], [0, 0, 204],
+			  [255, 51, 153], [0, 204, 204], [0, 51, 0], [255, 153, 51], [0, 204, 0]]
+	return colors
+
+
+def vis_faces(log_hooks):
+	display_count = len(log_hooks)
+	fig = plt.figure(figsize=(8, 4 * display_count))
+	gs = fig.add_gridspec(display_count, 3)
+	for i in range(display_count):
+		hooks_dict = log_hooks[i]
+		fig.add_subplot(gs[i, 0])
+		if 'diff_input' in hooks_dict:
+			vis_faces_with_id(hooks_dict, fig, gs, i)
+		else:
+			vis_faces_no_id(hooks_dict, fig, gs, i)
+	plt.tight_layout()
+	return fig
+
+
+def vis_faces_with_id(hooks_dict, fig, gs, i):
+	plt.imshow(hooks_dict['input_face'])
+	plt.title('Input\nOut Sim={:.2f}'.format(float(hooks_dict['diff_input'])))
+	fig.add_subplot(gs[i, 1])
+	plt.imshow(hooks_dict['target_face'])
+	plt.title('Target\nIn={:.2f}, Out={:.2f}'.format(float(hooks_dict['diff_views']),
+	                                                 float(hooks_dict['diff_target'])))
+	fig.add_subplot(gs[i, 2])
+	plt.imshow(hooks_dict['output_face'])
+	plt.title('Output\n Target Sim={:.2f}'.format(float(hooks_dict['diff_target'])))
+
+
+def vis_faces_no_id(hooks_dict, fig, gs, i):
+	plt.imshow(hooks_dict['input_face'], cmap="gray")
+	plt.title('Input')
+	fig.add_subplot(gs[i, 1])
+	plt.imshow(hooks_dict['target_face'])
+	plt.title('Target')
+	fig.add_subplot(gs[i, 2])
+	plt.imshow(hooks_dict['output_face'])
+	plt.title('Output')

utils/data_utils.py

@@ -0,0 +1,25 @@
+"""
+Code adopted from pix2pixHD:
+https://github.com/NVIDIA/pix2pixHD/blob/master/data/image_folder.py
+"""
+import os
+
+IMG_EXTENSIONS = [
+    '.jpg', '.JPG', '.jpeg', '.JPEG',
+    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tiff'
+]
+
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def make_dataset(dir):
+    images = []
+    assert os.path.isdir(dir), '%s is not a valid directory' % dir
+    for root, _, fnames in sorted(os.walk(dir)):
+        for fname in fnames:
+            if is_image_file(fname):
+                path = os.path.join(root, fname)
+                images.append(path)
+    return images

utils/inference_utils.py

@@ -0,0 +1,182 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from PIL import Image
+import cv2
+import random
+import math
+import argparse
+import torch
+from torch.utils import data
+from torch.nn import functional as F
+from torch import autograd
+from torch.nn import init
+import torchvision.transforms as transforms
+from scripts.align_all_parallel import get_landmark
+    
+def visualize(img_arr, dpi):
+    plt.figure(figsize=(10,10),dpi=dpi)
+    plt.imshow(((img_arr.detach().cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8))
+    plt.axis('off')
+    plt.show()
+
+def save_image(img, filename):
+    tmp = ((img.detach().cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8)
+    cv2.imwrite(filename, cv2.cvtColor(tmp, cv2.COLOR_RGB2BGR))
+    
+def load_image(filename):
+    transform = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5, 0.5, 0.5],std=[0.5,0.5,0.5]),
+    ])
+    
+    img = Image.open(filename)
+    img = transform(img)
+    return img.unsqueeze(dim=0)   
+
+def get_video_crop_parameter(filepath, predictor, padding=[256,256,256,256]):
+    if type(filepath) == str:
+        img = dlib.load_rgb_image(filepath)
+    else:
+        img = filepath
+    lm = get_landmark(img, predictor)
+    if lm is None:
+        return None
+    lm_chin          = lm[0  : 17]  # left-right
+    lm_eyebrow_left  = lm[17 : 22]  # left-right
+    lm_eyebrow_right = lm[22 : 27]  # left-right
+    lm_nose          = lm[27 : 31]  # top-down
+    lm_nostrils      = lm[31 : 36]  # top-down
+    lm_eye_left      = lm[36 : 42]  # left-clockwise
+    lm_eye_right     = lm[42 : 48]  # left-clockwise
+    lm_mouth_outer   = lm[48 : 60]  # left-clockwise
+    lm_mouth_inner   = lm[60 : 68]  # left-clockwise
+    
+    scale = 64. / (np.mean(lm_eye_right[:,0])-np.mean(lm_eye_left[:,0]))
+    center = ((np.mean(lm_eye_right, axis=0)+np.mean(lm_eye_left, axis=0)) / 2) * scale
+    h, w = round(img.shape[0] * scale), round(img.shape[1] * scale)
+    left = max(round(center[0] - padding[0]), 0) // 8 * 8
+    right = min(round(center[0] + padding[1]), w) // 8 * 8
+    top = max(round(center[1] - padding[2]), 0) // 8 * 8
+    bottom = min(round(center[1] + padding[3]), h) // 8 * 8
+    return h,w,top,bottom,left,right,scale
+
+def tensor2cv2(img):
+    tmp = ((img.cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8)
+    return cv2.cvtColor(tmp, cv2.COLOR_RGB2BGR)
+
+def noise_regularize(noises):
+    loss = 0
+
+    for noise in noises:
+        size = noise.shape[2]
+
+        while True:
+            loss = (
+                loss
+                + (noise * torch.roll(noise, shifts=1, dims=3)).mean().pow(2)
+                + (noise * torch.roll(noise, shifts=1, dims=2)).mean().pow(2)
+            )
+
+            if size <= 8:
+                break
+
+            #noise = noise.reshape([-1, 1, size // 2, 2, size // 2, 2])
+            #noise = noise.mean([3, 5])
+            noise = F.interpolate(noise, scale_factor=0.5, mode='bilinear')
+            size //= 2
+
+    return loss
+
+
+def noise_normalize_(noises):
+    for noise in noises:
+        mean = noise.mean()
+        std = noise.std()
+
+        noise.data.add_(-mean).div_(std)
+
+
+def get_lr(t, initial_lr, rampdown=0.25, rampup=0.05):
+    lr_ramp = min(1, (1 - t) / rampdown)
+    lr_ramp = 0.5 - 0.5 * math.cos(lr_ramp * math.pi)
+    lr_ramp = lr_ramp * min(1, t / rampup)
+
+    return initial_lr * lr_ramp
+
+
+def latent_noise(latent, strength):
+    noise = torch.randn_like(latent) * strength
+
+    return latent + noise
+
+
+def make_image(tensor):
+    return (
+        tensor.detach()
+        .clamp_(min=-1, max=1)
+        .add(1)
+        .div_(2)
+        .mul(255)
+        .type(torch.uint8)
+        .permute(0, 2, 3, 1)
+        .to("cpu")
+        .numpy()
+    )
+
+
+# from pix2pixeHD
+# Converts a one-hot tensor into a colorful label map
+def tensor2label(label_tensor, n_label, imtype=np.uint8):
+    if n_label == 0:
+        return tensor2im(label_tensor, imtype)
+    label_tensor = label_tensor.cpu().float()    
+    if label_tensor.size()[0] > 1:
+        label_tensor = label_tensor.max(0, keepdim=True)[1]
+    label_tensor = Colorize(n_label)(label_tensor)
+    label_numpy = np.transpose(label_tensor.numpy(), (1, 2, 0))
+    return label_numpy.astype(imtype)
+
+def uint82bin(n, count=8):
+    """returns the binary of integer n, count refers to amount of bits"""
+    return ''.join([str((n >> y) & 1) for y in range(count-1, -1, -1)])
+
+def labelcolormap(N):
+    if N == 35: # cityscape
+        cmap = np.array([(  0,  0,  0), (  0,  0,  0), (  0,  0,  0), (  0,  0,  0), (  0,  0,  0), (111, 74,  0), ( 81,  0, 81),
+                     (128, 64,128), (244, 35,232), (250,170,160), (230,150,140), ( 70, 70, 70), (102,102,156), (190,153,153),
+                     (180,165,180), (150,100,100), (150,120, 90), (153,153,153), (153,153,153), (250,170, 30), (220,220,  0),
+                     (107,142, 35), (152,251,152), ( 70,130,180), (220, 20, 60), (255,  0,  0), (  0,  0,142), (  0,  0, 70),
+                     (  0, 60,100), (  0,  0, 90), (  0,  0,110), (  0, 80,100), (  0,  0,230), (119, 11, 32), (  0,  0,142)], 
+                     dtype=np.uint8)
+    else:
+        cmap = np.zeros((N, 3), dtype=np.uint8)
+        for i in range(N):
+            r, g, b = 0, 0, 0
+            id = i
+            for j in range(7):
+                str_id = uint82bin(id)
+                r = r ^ (np.uint8(str_id[-1]) << (7-j))
+                g = g ^ (np.uint8(str_id[-2]) << (7-j))
+                b = b ^ (np.uint8(str_id[-3]) << (7-j))
+                id = id >> 3
+            cmap[i, 0] = r
+            cmap[i, 1] = g
+            cmap[i, 2] = b
+    return cmap
+
+class Colorize(object):
+    def __init__(self, n=35):
+        self.cmap = labelcolormap(n)
+        self.cmap = torch.from_numpy(self.cmap[:n])
+
+    def __call__(self, gray_image):
+        size = gray_image.size()
+        color_image = torch.ByteTensor(3, size[1], size[2]).fill_(0)
+
+        for label in range(0, len(self.cmap)):
+            mask = (label == gray_image[0]).cpu()
+            color_image[0][mask] = self.cmap[label][0]
+            color_image[1][mask] = self.cmap[label][1]
+            color_image[2][mask] = self.cmap[label][2]
+
+        return color_image

utils/train_utils.py

@@ -0,0 +1,13 @@
+
+def aggregate_loss_dict(agg_loss_dict):
+	mean_vals = {}
+	for output in agg_loss_dict:
+		for key in output:
+			mean_vals[key] = mean_vals.setdefault(key, []) + [output[key]]
+	for key in mean_vals:
+		if len(mean_vals[key]) > 0:
+			mean_vals[key] = sum(mean_vals[key]) / len(mean_vals[key])
+		else:
+			print('{} has no value'.format(key))
+			mean_vals[key] = 0
+	return mean_vals

utils/wandb_utils.py

@@ -0,0 +1,47 @@
+import datetime
+import os
+import numpy as np
+import wandb
+
+from utils import common
+
+
+class WBLogger:
+
+    def __init__(self, opts):
+        wandb_run_name = os.path.basename(opts.exp_dir)
+        wandb.init(project="pixel2style2pixel", config=vars(opts), name=wandb_run_name)
+
+    @staticmethod
+    def log_best_model():
+        wandb.run.summary["best-model-save-time"] = datetime.datetime.now()
+
+    @staticmethod
+    def log(prefix, metrics_dict, global_step):
+        log_dict = {f'{prefix}_{key}': value for key, value in metrics_dict.items()}
+        log_dict["global_step"] = global_step
+        wandb.log(log_dict)
+
+    @staticmethod
+    def log_dataset_wandb(dataset, dataset_name, n_images=16):
+        idxs = np.random.choice(a=range(len(dataset)), size=n_images, replace=False)
+        data = [wandb.Image(dataset.source_paths[idx]) for idx in idxs]
+        wandb.log({f"{dataset_name} Data Samples": data})
+
+    @staticmethod
+    def log_images_to_wandb(x, y, y_hat, id_logs, prefix, step, opts):
+        im_data = []
+        column_names = ["Source", "Target", "Output"]
+        if id_logs is not None:
+            column_names.append("ID Diff Output to Target")
+        for i in range(len(x)):
+            cur_im_data = [
+                wandb.Image(common.log_input_image(x[i], opts)),
+                wandb.Image(common.tensor2im(y[i])),
+                wandb.Image(common.tensor2im(y_hat[i])),
+            ]
+            if id_logs is not None:
+                cur_im_data.append(id_logs[i]["diff_target"])
+            im_data.append(cur_im_data)
+        outputs_table = wandb.Table(data=im_data, columns=column_names)
+        wandb.log({f"{prefix.title()} Step {step} Output Samples": outputs_table})

webUI/app_task.py

@@ -0,0 +1,305 @@
+from __future__ import annotations
+from huggingface_hub import hf_hub_download
+import numpy as np
+import gradio as gr
+
+
+def create_demo_sr(process):
+    with gr.Blocks() as demo:
+        with gr.Row():
+            gr.Markdown('## Face Super Resolution')
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(source='upload', type='filepath')
+                model_type = gr.Radio(label='Model Type', choices=['SR for 32x','SR for 4x-48x'], value='SR for 32x')
+                resize_scale = gr.Slider(label='Resize Scale',
+                                            minimum=4,
+                                            maximum=48,
+                                            value=32,
+                                            step=4)                                
+                run_button = gr.Button(label='Run')
+                gr.Examples(
+                    examples =[['pexels-daniel-xavier-1239291.jpg', 'SR for 32x', 32],
+                               ['ILip77SbmOE.png', 'SR for 32x', 32],
+                               ['ILip77SbmOE.png', 'SR for 4x-48x', 48],
+                              ],
+                    inputs = [input_image, model_type, resize_scale],
+                )                
+            with gr.Column():
+                #lrinput = gr.Image(label='Low-resolution input',type='numpy', interactive=False)
+                #result = gr.Image(label='Output',type='numpy', interactive=False)
+                result = gr.Gallery(label='LR input and Output',
+                                    elem_id='gallery').style(grid=2,
+                                                             height='auto')    
+
+        inputs = [
+            input_image,
+            resize_scale,
+            model_type,
+        ]
+        run_button.click(fn=process,
+                         inputs=inputs,
+                         outputs=[result],
+                         api_name='sr')
+    return demo
+    
+def create_demo_s2f(process):
+    with gr.Blocks() as demo:
+        with gr.Row():
+            gr.Markdown('## Sketch-to-Face Translation')
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(source='upload', type='filepath')
+                gr.Markdown("""Note: Input will be cropped if larger than 512x512.""")
+                seed = gr.Slider(label='Seed for appearance',
+                                    minimum=0,
+                                    maximum=2147483647,
+                                    step=1,
+                                    randomize=True)
+                #input_info = gr.Textbox(label='Process Information', interactive=False, value='n.a.')
+                run_button = gr.Button(label='Run')
+                gr.Examples(
+                    examples =[['234_sketch.jpg', 1024]],
+                    inputs = [input_image, seed],
+                )                
+            with gr.Column():
+                result = gr.Image(label='Output',type='numpy', interactive=False) 
+
+        inputs = [
+            input_image, seed
+        ]
+        run_button.click(fn=process,
+                         inputs=inputs,
+                         outputs=[result],
+                         api_name='s2f')
+    return demo
+
+
+def create_demo_m2f(process):
+    with gr.Blocks() as demo:
+        with gr.Row():
+            gr.Markdown('## Mask-to-Face Translation')
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(source='upload', type='filepath') 
+                input_type = gr.Radio(label='Input Type', choices=['color image','parsing mask'], value='color image')
+                seed = gr.Slider(label='Seed for appearance',
+                                    minimum=0,
+                                    maximum=2147483647,
+                                    step=1,
+                                    randomize=True)
+                #input_info = gr.Textbox(label='Process Information', interactive=False, value='n.a.')
+                run_button = gr.Button(label='Run')
+                gr.Examples(
+                    examples =[['ILip77SbmOE.png', 'color image', 4], ['ILip77SbmOE_mask.png', 'parsing mask', 4]],
+                    inputs = [input_image, input_type, seed],
+                )                
+            with gr.Column():
+                #vizmask = gr.Image(label='Visualized mask',type='numpy', interactive=False)
+                #result = gr.Image(label='Output',type='numpy', interactive=False) 
+                result = gr.Gallery(label='Visualized mask and Output',
+                                    elem_id='gallery').style(grid=2,
+                                                             height='auto')                
+
+        inputs = [
+            input_image, input_type, seed
+        ]
+        run_button.click(fn=process,
+                         inputs=inputs,
+                         outputs=[result],
+                         api_name='m2f')
+    return demo
+
+def create_demo_editing(process):
+    with gr.Blocks() as demo:
+        with gr.Row():
+            gr.Markdown('## Video Face Editing (for image input)')
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(source='upload', type='filepath') 
+                model_type = gr.Radio(label='Editing Type', choices=['reduce age','light hair color'], value='color image')
+                scale_factor = gr.Slider(label='editing degree (-2~2)',
+                                    minimum=-2,
+                                    maximum=2,
+                                    value=1,
+                                    step=0.1)
+                #input_info = gr.Textbox(label='Process Information', interactive=False, value='n.a.')
+                run_button = gr.Button(label='Run')
+                gr.Examples(
+                    examples =[['ILip77SbmOE.png', 'reduce age', -2], 
+                               ['ILip77SbmOE.png', 'light hair color', 1]],
+                    inputs = [input_image, model_type, scale_factor],
+                )                
+            with gr.Column():
+                result = gr.Image(label='Output',type='numpy', interactive=False) 
+
+        inputs = [
+            input_image, scale_factor, model_type
+        ]
+        run_button.click(fn=process,
+                         inputs=inputs,
+                         outputs=[result],
+                         api_name='editing')
+    return demo
+
+def create_demo_toonify(process):
+    with gr.Blocks() as demo:
+        with gr.Row():
+            gr.Markdown('## Video Face Toonification (for image input)')
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(source='upload', type='filepath') 
+                style_type = gr.Radio(label='Style Type', choices=['Pixar','Cartoon','Arcane'], value='Pixar')                
+                #input_info = gr.Textbox(label='Process Information', interactive=False, value='n.a.')
+                run_button = gr.Button(label='Run')
+                gr.Examples(
+                    examples =[['ILip77SbmOE.png', 'Pixar'], ['ILip77SbmOE.png', 'Cartoon'], ['ILip77SbmOE.png', 'Arcane']],
+                    inputs = [input_image, style_type],
+                )                
+            with gr.Column():
+                result = gr.Image(label='Output',type='numpy', interactive=False) 
+
+        inputs = [
+            input_image, style_type
+        ]
+        run_button.click(fn=process,
+                         inputs=inputs,
+                         outputs=[result],
+                         api_name='toonify')
+    return demo
+
+
+def create_demo_vediting(process, max_frame_num = 4):
+    with gr.Blocks() as demo:
+        with gr.Row():
+            gr.Markdown('## Video Face Editing (for video input)')
+        with gr.Row():
+            with gr.Column():
+                input_video = gr.Video(source='upload', mirror_webcam=False, type='filepath') 
+                model_type = gr.Radio(label='Editing Type', choices=['reduce age','light hair color'], value='color image')
+                scale_factor = gr.Slider(label='editing degree (-2~2)',
+                                    minimum=-2,
+                                    maximum=2,
+                                    value=1,
+                                    step=0.1)
+                frame_num = gr.Slider(label='Number of frames to edit (full video editing is not allowed so as not to slow down the demo, \
+                            but you can duplicate the Space to modify the number limit from 4 to a large value)',
+                                    minimum=1,
+                                    maximum=max_frame_num,
+                                    value=2,
+                                    step=1)   
+                #input_info = gr.Textbox(label='Process Information', interactive=False, value='n.a.')
+                run_button = gr.Button(label='Run')
+                gr.Examples(
+                    examples =[['684.mp4', 'reduce age', 1.5, 2], 
+                               ['684.mp4', 'light hair color', 0.7, 2]],
+                    inputs = [input_video, model_type, scale_factor],
+                )                
+            with gr.Column():
+                viz_result = gr.Gallery(label='Several edited frames', elem_id='gallery').style(grid=2, height='auto')  
+                result = gr.Video(label='Output', type='mp4', interactive=False) 
+
+        inputs = [
+            input_video, scale_factor, model_type, frame_num
+        ]
+        run_button.click(fn=process,
+                         inputs=inputs,
+                         outputs=[viz_result, result],
+                         api_name='vediting')
+    return demo
+
+def create_demo_vtoonify(process, max_frame_num = 4):
+    with gr.Blocks() as demo:
+        with gr.Row():
+            gr.Markdown('## Video Face Toonification (for video input)')
+        with gr.Row():
+            with gr.Column():
+                input_video = gr.Video(source='upload', mirror_webcam=False, type='filepath') 
+                style_type = gr.Radio(label='Style Type', choices=['Pixar','Cartoon','Arcane'], value='Pixar')
+                frame_num = gr.Slider(label='Number of frames to toonify (full video toonification is not allowed so as not to slow down the demo, \
+                            but you can duplicate the Space to modify the number limit from 4 to a large value)',
+                                    minimum=1,
+                                    maximum=max_frame_num,
+                                    value=2,
+                                    step=1)                
+                #input_info = gr.Textbox(label='Process Information', interactive=False, value='n.a.')
+                run_button = gr.Button(label='Run')
+                gr.Examples(
+                    examples =[['529_2.mp4', 'Arcane'],
+                               ['pexels-anthony-shkraba-production-8136210.mp4', 'Pixar'], 
+                               ['684.mp4', 'Cartoon']],
+                    inputs = [input_video, style_type],
+                )                
+            with gr.Column():
+                viz_result = gr.Gallery(label='Several toonified frames', elem_id='gallery').style(grid=2, height='auto')  
+                result = gr.Video(label='Output', type='mp4', interactive=False) 
+
+        inputs = [
+            input_video, style_type, frame_num
+        ]
+        run_button.click(fn=process,
+                         inputs=inputs,
+                         outputs=[viz_result, result],
+                         api_name='vtoonify')
+    return demo
+
+def create_demo_inversion(process, allow_optimization=False):
+    with gr.Blocks() as demo:
+        with gr.Row():
+            gr.Markdown('## StyleGANEX Inversion for Editing')
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(source='upload', type='filepath') 
+                optimize = gr.Radio(label='Whether optimize latent (latent optimization is not allowed so as not to slow down the demo, \
+                but you can duplicate the Space to modify the option or directly upload an optimized latent file. \
+                The file can be computed by inversion.py from the github page or colab)', choices=['No optimization','Latent optimization'], 
+                                    value='No optimization', interactive=allow_optimization)
+                input_latent = gr.File(label='Optimized latent code (optional)', file_types=[".pt"])
+                editing_options = gr.Dropdown(['None', 'Style Mixing', 
+                                               'Attribute Editing: smile', 
+                                               'Attribute Editing: open_eye', 
+                                               'Attribute Editing: open_mouth', 
+                                               'Attribute Editing: pose', 
+                                               'Attribute Editing: reduce_age', 
+                                               'Attribute Editing: glasses', 
+                                               'Attribute Editing: light_hair_color', 
+                                               'Attribute Editing: slender', 
+                                               'Domain Transfer: disney_princess',
+                                               'Domain Transfer: vintage_comics',
+                                               'Domain Transfer: pixar',
+                                               'Domain Transfer: edvard_munch',
+                                               'Domain Transfer: modigliani',
+                                              ], 
+                                              label="editing options",
+                                              value='None')
+                scale_factor = gr.Slider(label='editing degree (-2~2) for Attribute Editing',
+                                    minimum=-2,
+                                    maximum=2,
+                                    value=2,
+                                    step=0.1)
+                seed = gr.Slider(label='Appearance Seed for Style Mixing',
+                                    minimum=0,
+                                    maximum=2147483647,
+                                    step=1,
+                                    randomize=True)
+                #input_info = gr.Textbox(label='Process Information', interactive=False, value='n.a.')
+                run_button = gr.Button(label='Run')
+                gr.Examples(
+                    examples =[['ILip77SbmOE.png', 'ILip77SbmOE_inversion.pt', 'Domain Transfer: vintage_comics'],
+                               ['ILip77SbmOE.png', 'ILip77SbmOE_inversion.pt', 'Attribute Editing: smile'],
+                               ['ILip77SbmOE.png', 'ILip77SbmOE_inversion.pt', 'Style Mixing'],
+                              ], 
+                    inputs = [input_image, input_latent, editing_options],
+                )                
+            with gr.Column():
+                result = gr.Image(label='Inversion output',type='numpy', interactive=False) 
+                editing_result = gr.Image(label='Editing output',type='numpy', interactive=False) 
+
+        inputs = [
+            input_image, optimize, input_latent, editing_options, scale_factor, seed
+        ]
+        run_button.click(fn=process,
+                         inputs=inputs,
+                         outputs=[result, editing_result],
+                         api_name='inversion')
+    return demo

webUI/styleganex_model.py

@@ -0,0 +1,492 @@
+from __future__ import annotations
+import numpy as np
+import gradio as gr
+
+import os
+import pathlib
+import gc
+import torch
+import dlib
+import cv2
+import PIL
+from tqdm import tqdm
+import numpy as np
+import torch.nn.functional as F
+import torchvision
+from torchvision import transforms, utils
+from argparse import Namespace
+from datasets import augmentations
+from huggingface_hub import hf_hub_download
+from scripts.align_all_parallel import align_face
+from latent_optimization import latent_optimization
+from utils.inference_utils import save_image, load_image, visualize, get_video_crop_parameter, tensor2cv2, tensor2label, labelcolormap
+from models.psp import pSp
+from models.bisenet.model import BiSeNet
+from models.stylegan2.model import Generator
+
+class Model():
+    def __init__(self, device):
+        super().__init__()
+        
+        self.device = device
+        self.task_name = None
+        self.editing_w = None
+        self.pspex = None
+        self.landmarkpredictor = dlib.shape_predictor(hf_hub_download('PKUWilliamYang/VToonify', 'models/shape_predictor_68_face_landmarks.dat'))
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5],std=[0.5,0.5,0.5]),
+            ])
+        self.to_tensor = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+        ])
+        self.maskpredictor = BiSeNet(n_classes=19)
+        self.maskpredictor.load_state_dict(torch.load(hf_hub_download('PKUWilliamYang/VToonify', 'models/faceparsing.pth'), map_location='cpu'))
+        self.maskpredictor.to(self.device).eval()
+        self.parameters = {}
+        self.parameters['inversion'] = {'path':'pretrained_models/styleganex_inversion.pt', 'image_path':'./data/ILip77SbmOE.png'}
+        self.parameters['sr-32'] = {'path':'pretrained_models/styleganex_sr32.pt', 'image_path':'./data/pexels-daniel-xavier-1239291.jpg'}
+        self.parameters['sr'] = {'path':'pretrained_models/styleganex_sr.pt', 'image_path':'./data/pexels-daniel-xavier-1239291.jpg'}
+        self.parameters['sketch2face'] = {'path':'pretrained_models/styleganex_sketch2face.pt', 'image_path':'./data/234_sketch.jpg'}
+        self.parameters['mask2face'] = {'path':'pretrained_models/styleganex_mask2face.pt', 'image_path':'./data/540.jpg'}
+        self.parameters['edit_age'] = {'path':'pretrained_models/styleganex_edit_age.pt', 'image_path':'./data/390.mp4'}
+        self.parameters['edit_hair'] = {'path':'pretrained_models/styleganex_edit_hair.pt', 'image_path':'./data/390.mp4'}
+        self.parameters['toonify_pixar'] = {'path':'pretrained_models/styleganex_toonify_pixar.pt', 'image_path':'./data/pexels-anthony-shkraba-production-8136210.mp4'}
+        self.parameters['toonify_cartoon'] = {'path':'pretrained_models/styleganex_toonify_cartoon.pt', 'image_path':'./data/pexels-anthony-shkraba-production-8136210.mp4'}
+        self.parameters['toonify_arcane'] = {'path':'pretrained_models/styleganex_toonify_arcane.pt', 'image_path':'./data/pexels-anthony-shkraba-production-8136210.mp4'}        
+        self.print_log = True
+        self.editing_dicts = torch.load(hf_hub_download('PKUWilliamYang/StyleGANEX', 'direction_dics.pt'))
+        self.generator = Generator(1024, 512, 8)
+        self.model_type = None
+        self.error_info = 'Error: no face detected! \
+                               StyleGANEX uses dlib.get_frontal_face_detector but sometimes it fails to detect a face. \
+                               You can try several times or use other images until a face is detected, \
+                               then switch back to the original image.'
+        
+    def load_model(self, task_name: str) -> None:
+        if task_name == self.task_name:
+            return
+        if self.pspex is not None:
+            del self.pspex
+        torch.cuda.empty_cache()
+        gc.collect()
+        path = self.parameters[task_name]['path']
+        local_path = hf_hub_download('PKUWilliamYang/StyleGANEX', path)
+        ckpt = torch.load(local_path, map_location='cpu')
+        opts = ckpt['opts']
+        opts['checkpoint_path'] = local_path
+        opts['device'] = self.device
+        opts = Namespace(**opts)
+        self.pspex = pSp(opts, ckpt).to(self.device).eval()
+        self.pspex.latent_avg = self.pspex.latent_avg.to(self.device)
+        if 'editing_w' in ckpt.keys():
+            self.editing_w = ckpt['editing_w'].clone().to(self.device)   
+        self.task_name = task_name  
+        torch.cuda.empty_cache()
+        gc.collect()
+
+    def load_G_model(self, model_type: str) -> None:
+        if model_type == self.model_type:
+            return
+        torch.cuda.empty_cache()
+        gc.collect()
+        local_path = hf_hub_download('rinong/stylegan-nada-models', model_type+'.pt')
+        self.generator.load_state_dict(torch.load(local_path, map_location='cpu')['g_ema'], strict=False)
+        self.generator.to(self.device).eval()
+        self.model_type = model_type  
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+    def tensor2np(self, img):
+        tmp = ((img.cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8)
+        return tmp
+
+    def process_sr(self, input_image: str, resize_scale: int, model: str) -> list[np.ndarray]:
+        #false_image = np.zeros((256,256,3), np.uint8)
+        #info = 'Error: no face detected! Please retry or change the photo.'
+            
+        if input_image is None:
+            #return [false_image, false_image], 'Error: fail to load empty file.'
+            raise gr.Error("Error: fail to load empty file.")
+        frame = cv2.imread(input_image)
+        if frame is None:
+            #return [false_image, false_image], 'Error: fail to load the image.'
+            raise gr.Error("Error: fail to load the image.")
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)            
+            
+        if model is None or model == 'SR for 32x':
+            task_name = 'sr-32'
+            resize_scale = 32
+        else:
+            task_name = 'sr'
+            
+        with torch.no_grad():
+            paras = get_video_crop_parameter(frame, self.landmarkpredictor)
+            if paras is None:
+                #return [false_image, false_image], info
+                raise gr.Error(self.error_info)
+            h,w,top,bottom,left,right,scale = paras
+            H, W = int(bottom-top), int(right-left)
+            frame = cv2.resize(frame, (w, h))[top:bottom, left:right]
+            x1 = PIL.Image.fromarray(np.uint8(frame))
+            x1 = augmentations.BilinearResize(factors=[resize_scale//4])(x1)
+            x1_up = x1.resize((W, H))
+            x2_up = align_face(np.array(x1_up), self.landmarkpredictor)
+            if x2_up is None:
+                #return [false_image, false_image], 'Error: no face detected! Please retry or change the photo.'
+                raise gr.Error(self.error_info)
+            x1_up = transforms.ToTensor()(x1_up).unsqueeze(dim=0).to(self.device) * 2 - 1
+            x2_up = self.transform(x2_up).unsqueeze(dim=0).to(self.device)
+            if self.print_log: print('image loaded')
+            self.load_model(task_name)
+            if self.print_log: print('model %s loaded'%(task_name))
+            y_hat = torch.clamp(self.pspex(x1=x1_up, x2=x2_up, use_skip=self.pspex.opts.use_skip, resize=False), -1, 1)
+
+        return [self.tensor2np(x1_up[0]), self.tensor2np(y_hat[0])]
+    
+    
+    def process_s2f(self, input_image: str, seed: int) -> np.ndarray:
+        task_name = 'sketch2face'
+        with torch.no_grad():
+            x1 = transforms.ToTensor()(PIL.Image.open(input_image)).unsqueeze(0).to(self.device)
+            if x1.shape[2] > 513:
+                x1 = x1[:,:,(x1.shape[2]//2-256)//8*8:(x1.shape[2]//2+256)//8*8]
+            if x1.shape[3] > 513:
+                x1 = x1[:,:,:,(x1.shape[3]//2-256)//8*8:(x1.shape[3]//2+256)//8*8]
+            x1 = x1[:,0:1] # uploaded files will be transformed to 3-channel RGB image!
+            if self.print_log: print('image loaded')
+            self.load_model(task_name)
+            if self.print_log: print('model %s loaded'%(task_name))
+            self.pspex.train()
+            torch.manual_seed(seed)
+            y_hat = self.pspex(x1=x1, resize=False, latent_mask=[8,9,10,11,12,13,14,15,16,17], use_skip=self.pspex.opts.use_skip,
+                                  inject_latent= self.pspex.decoder.style(torch.randn(1, 512).to(self.device)).unsqueeze(1).repeat(1,18,1) * 0.7)  
+            y_hat = torch.clamp(y_hat, -1, 1)
+            self.pspex.eval()
+        return self.tensor2np(y_hat[0])
+    
+    def process_m2f(self, input_image: str, input_type: str, seed: int) -> list[np.ndarray]:
+        #false_image = np.zeros((256,256,3), np.uint8)
+        if input_image is None:
+            raise gr.Error('Error: fail to load empty file.' )
+            #return [false_image, false_image], 'Error: fail to load empty file.'        
+        task_name = 'mask2face'
+        with torch.no_grad():
+            if input_type == 'parsing mask':
+                x1 = PIL.Image.open(input_image).getchannel(0) # uploaded files will be transformed to 3-channel RGB image!
+                x1 = augmentations.ToOneHot(19)(x1)
+                x1 = transforms.ToTensor()(x1).unsqueeze(dim=0).float().to(self.device)
+                #print(x1.shape)
+            else:
+                frame = cv2.imread(input_image)
+                if frame is None:
+                    #return [false_image, false_image], 'Error: fail to load the image.' 
+                    raise gr.Error('Error: fail to load the image.' )
+                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                paras = get_video_crop_parameter(frame, self.landmarkpredictor)        
+                if paras is None:
+                    #return [false_image, false_image], 'Error: no face detected! Please retry or change the photo.'
+                    raise gr.Error(self.error_info)
+                h,w,top,bottom,left,right,scale = paras
+                H, W = int(bottom-top), int(right-left)
+                frame = cv2.resize(frame, (w, h))[top:bottom, left:right]        
+                # convert face image to segmentation mask
+                x1 = self.to_tensor(frame).unsqueeze(0).to(self.device)
+                # upsample image for precise segmentation
+                x1 = F.interpolate(x1, scale_factor=2, mode='bilinear') 
+                x1 = self.maskpredictor(x1)[0]
+                x1 = F.interpolate(x1, scale_factor=0.5).argmax(dim=1)
+                x1 = F.one_hot(x1, num_classes=19).permute(0, 3, 1, 2).float().to(self.device)   
+                
+            if x1.shape[2] > 513:
+                x1 = x1[:,:,(x1.shape[2]//2-256)//8*8:(x1.shape[2]//2+256)//8*8]
+            if x1.shape[3] > 513:
+                x1 = x1[:,:,:,(x1.shape[3]//2-256)//8*8:(x1.shape[3]//2+256)//8*8]
+
+            x1_viz = (tensor2label(x1[0], 19) / 192 * 256).astype(np.uint8)
+
+            if self.print_log: print('image loaded')
+            self.load_model(task_name)
+            if self.print_log: print('model %s loaded'%(task_name))
+            self.pspex.train()
+            torch.manual_seed(seed)
+            y_hat = self.pspex(x1=x1, resize=False, latent_mask=[8,9,10,11,12,13,14,15,16,17], use_skip=self.pspex.opts.use_skip,
+                                  inject_latent= self.pspex.decoder.style(torch.randn(1, 512).to(self.device)).unsqueeze(1).repeat(1,18,1) * 0.7)  
+            y_hat = torch.clamp(y_hat, -1, 1)
+            self.pspex.eval()
+        return [x1_viz, self.tensor2np(y_hat[0])]
+    
+    
+    def process_editing(self, input_image: str, scale_factor: float, model_type: str) -> np.ndarray:
+        #false_image = np.zeros((256,256,3), np.uint8)
+        #info = 'Error: no face detected! Please retry or change the photo.'
+            
+        if input_image is None:
+            #return false_image, false_image, 'Error: fail to load empty file.'
+            raise gr.Error('Error: fail to load empty file.')
+        frame = cv2.imread(input_image)
+        if frame is None:
+            #return false_image, false_image, 'Error: fail to load the image.'
+            raise gr.Error('Error: fail to load the image.')
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)            
+            
+        if model_type is None or model_type == 'reduce age':
+            task_name = 'edit_age'
+        else:
+            task_name = 'edit_hair'
+            
+        with torch.no_grad():
+            paras = get_video_crop_parameter(frame, self.landmarkpredictor)
+            if paras is None:
+                #return false_image, false_image, info
+                raise gr.Error(self.error_info)
+            h,w,top,bottom,left,right,scale = paras
+            H, W = int(bottom-top), int(right-left)
+            frame = cv2.resize(frame, (w, h))[top:bottom, left:right]
+            x1 = self.transform(frame).unsqueeze(0).to(self.device)
+            x2 = align_face(frame, self.landmarkpredictor)
+            if x2 is None:
+                #return false_image, 'Error: no face detected! Please retry or change the photo.'
+                raise gr.Error(self.error_info)
+            x2 = self.transform(x2).unsqueeze(dim=0).to(self.device)
+            if self.print_log: print('image loaded')
+            self.load_model(task_name)
+            if self.print_log: print('model %s loaded'%(task_name))
+            y_hat = self.pspex(x1=x1, x2=x2, use_skip=self.pspex.opts.use_skip, zero_noise=True, 
+                        resize=False, editing_w= - scale_factor* self.editing_w[0:1])
+            y_hat = torch.clamp(y_hat, -1, 1)
+
+        return self.tensor2np(y_hat[0])
+
+    def process_vediting(self, input_video: str, scale_factor: float, model_type: str, frame_num: int) -> tuple[list[np.ndarray], str]:
+        #false_image = np.zeros((256,256,3), np.uint8)
+        #info = 'Error: no face detected! Please retry or change the video.'
+            
+        if input_video is None:
+            #return [false_image], 'default.mp4', 'Error: fail to load empty file.'
+            raise gr.Error('Error: fail to load empty file.')
+        video_cap = cv2.VideoCapture(input_video)
+        success, frame = video_cap.read()
+        if success is False:
+            #return [false_image], 'default.mp4', 'Error: fail to load the video.' 
+            raise gr.Error('Error: fail to load the video.')
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)         
+        
+        if model_type is None or model_type == 'reduce age':
+            task_name = 'edit_age'
+        else:
+            task_name = 'edit_hair'
+            
+        with torch.no_grad():
+            paras = get_video_crop_parameter(frame, self.landmarkpredictor)
+            if paras is None:
+                #return [false_image], 'default.mp4', info
+                raise gr.Error(self.error_info)
+            h,w,top,bottom,left,right,scale = paras
+            H, W = int(bottom-top), int(right-left)
+            frame = cv2.resize(frame, (w, h))[top:bottom, left:right]
+            x1 = self.transform(frame).unsqueeze(0).to(self.device)
+            x2 = align_face(frame, self.landmarkpredictor)
+            if x2 is None:
+                #return [false_image], 'default.mp4', info
+                raise gr.Error(self.error_info)
+            x2 = self.transform(x2).unsqueeze(dim=0).to(self.device)
+            if self.print_log: print('first frame loaded')
+            self.load_model(task_name)
+            if self.print_log: print('model %s loaded'%(task_name))
+            
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            videoWriter = cv2.VideoWriter('output.mp4', fourcc, video_cap.get(5), (4*W, 4*H))
+            
+            viz_frames = []
+            for i in range(frame_num):
+                if i > 0:
+                    success, frame = video_cap.read()
+                    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                    frame = cv2.resize(frame, (w, h))[top:bottom, left:right]
+                    x1 = self.transform(frame).unsqueeze(0).to(self.device)
+                y_hat = self.pspex(x1=x1, x2=x2, use_skip=self.pspex.opts.use_skip, zero_noise=True, 
+                        resize=False, editing_w= - scale_factor * self.editing_w[0:1])
+                y_hat = torch.clamp(y_hat, -1, 1)
+                videoWriter.write(tensor2cv2(y_hat[0].cpu()))
+                if i < min(frame_num, 4):
+                    viz_frames += [self.tensor2np(y_hat[0])]
+                
+            videoWriter.release()    
+            
+        return viz_frames, 'output.mp4'
+    
+    
+    def process_toonify(self, input_image: str, style_type: str) -> np.ndarray:
+        #false_image = np.zeros((256,256,3), np.uint8)
+        #info = 'Error: no face detected! Please retry or change the photo.'
+            
+        if input_image is None:
+            raise gr.Error('Error: fail to load empty file.')
+            #return false_image, false_image, 'Error: fail to load empty file.'
+        frame = cv2.imread(input_image)
+        if frame is None:
+            raise gr.Error('Error: fail to load the image.')
+            #return false_image, false_image, 'Error: fail to load the image.'
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)            
+            
+        if style_type is None or style_type == 'Pixar':
+            task_name = 'toonify_pixar'
+        elif style_type == 'Cartoon':
+            task_name = 'toonify_cartoon'
+        else:
+            task_name = 'toonify_arcane'
+            
+        with torch.no_grad():
+            paras = get_video_crop_parameter(frame, self.landmarkpredictor)
+            if paras is None:
+                raise gr.Error(self.error_info)
+                #return false_image, false_image, info
+            h,w,top,bottom,left,right,scale = paras
+            H, W = int(bottom-top), int(right-left)
+            frame = cv2.resize(frame, (w, h))[top:bottom, left:right]
+            x1 = self.transform(frame).unsqueeze(0).to(self.device)
+            x2 = align_face(frame, self.landmarkpredictor)
+            if x2 is None:
+                raise gr.Error(self.error_info)
+                #return false_image, 'Error: no face detected! Please retry or change the photo.'
+            x2 = self.transform(x2).unsqueeze(dim=0).to(self.device)
+            if self.print_log: print('image loaded')
+            self.load_model(task_name)
+            if self.print_log: print('model %s loaded'%(task_name))
+            y_hat = self.pspex(x1=x1, x2=x2, use_skip=self.pspex.opts.use_skip, zero_noise=True, resize=False)
+            y_hat = torch.clamp(y_hat, -1, 1)
+
+        return self.tensor2np(y_hat[0]), 'Done!'    
+
+
+    def process_vtoonify(self, input_video: str, style_type: str, frame_num: int) -> tuple[list[np.ndarray], str]:
+        #false_image = np.zeros((256,256,3), np.uint8)
+        #info = 'Error: no face detected! Please retry or change the video.'
+            
+        if input_video is None:
+            raise gr.Error('Error: fail to load empty file.')
+            #return [false_image], 'default.mp4', 'Error: fail to load empty file.'
+        video_cap = cv2.VideoCapture(input_video)
+        success, frame = video_cap.read()
+        if success is False:
+            raise gr.Error('Error: fail to load the video.')
+            #return [false_image], 'default.mp4', 'Error: fail to load the video.'        
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)         
+        
+        if style_type is None or style_type == 'Pixar':
+            task_name = 'toonify_pixar'
+        elif style_type == 'Cartoon':
+            task_name = 'toonify_cartoon'
+        else:
+            task_name = 'toonify_arcane'
+            
+        with torch.no_grad():
+            paras = get_video_crop_parameter(frame, self.landmarkpredictor)
+            if paras is None:
+                raise gr.Error(self.error_info)
+                #return [false_image], 'default.mp4', info
+            h,w,top,bottom,left,right,scale = paras
+            H, W = int(bottom-top), int(right-left)
+            frame = cv2.resize(frame, (w, h))[top:bottom, left:right]
+            x1 = self.transform(frame).unsqueeze(0).to(self.device)
+            x2 = align_face(frame, self.landmarkpredictor)
+            if x2 is None:
+                raise gr.Error(self.error_info)
+                #return [false_image], 'default.mp4', info
+            x2 = self.transform(x2).unsqueeze(dim=0).to(self.device)
+            if self.print_log: print('first frame loaded')
+            self.load_model(task_name)
+            if self.print_log: print('model %s loaded'%(task_name))
+            
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            videoWriter = cv2.VideoWriter('output.mp4', fourcc, video_cap.get(5), (4*W, 4*H))
+            
+            viz_frames = []
+            for i in range(frame_num):
+                if i > 0:
+                    success, frame = video_cap.read()
+                    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                    frame = cv2.resize(frame, (w, h))[top:bottom, left:right]
+                    x1 = self.transform(frame).unsqueeze(0).to(self.device)
+                y_hat = self.pspex(x1=x1, x2=x2, use_skip=self.pspex.opts.use_skip, zero_noise=True, resize=False)
+                y_hat = torch.clamp(y_hat, -1, 1)
+                videoWriter.write(tensor2cv2(y_hat[0].cpu()))
+                if i < min(frame_num, 4):
+                    viz_frames += [self.tensor2np(y_hat[0])]
+                
+            videoWriter.release()    
+            
+        return viz_frames, 'output.mp4'   
+    
+    
+    def process_inversion(self, input_image: str, optimize: str, input_latent: file-object, editing_options: str, 
+                          scale_factor: float, seed: int) -> tuple[np.ndarray, np.ndarray]:
+        #false_image = np.zeros((256,256,3), np.uint8)
+        #info = 'Error: no face detected! Please retry or change the photo.'
+            
+        if input_image is None:
+            raise gr.Error('Error: fail to load empty file.')
+            #return false_image, false_image, 'Error: fail to load empty file.'
+        frame = cv2.imread(input_image)
+        if frame is None:
+            raise gr.Error('Error: fail to load the image.')
+            #return false_image, false_image, 'Error: fail to load the image.'
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        
+        task_name = 'inversion'
+        self.load_model(task_name)
+        if self.print_log: print('model %s loaded'%(task_name))
+        if input_latent is not None:
+            if '.pt' not in input_latent.name:
+                raise gr.Error('Error: the latent format is wrong')
+                #return false_image, false_image, 'Error: the latent format is wrong'
+            latents = torch.load(input_latent.name)
+            if 'wplus' not in latents.keys() or 'f' not in latents.keys():
+                raise gr.Error('Error: the latent format is wrong')
+                #return false_image, false_image, 'Error: the latent format is wrong'
+            wplus = latents['wplus'].to(self.device) # w+
+            f = [latents['f'][0].to(self.device)]    # f
+        elif optimize == 'Latent optimization':
+            wplus, f, _, _, _ = latent_optimization(frame, self.pspex, self.landmarkpredictor, 
+                                                                     step=500, device=self.device)
+        else:
+            with torch.no_grad():
+                paras = get_video_crop_parameter(frame, self.landmarkpredictor)
+                if paras is None:
+                    raise gr.Error(self.error_info)
+                    #return false_image, false_image, info
+                h,w,top,bottom,left,right,scale = paras
+                H, W = int(bottom-top), int(right-left)
+                frame = cv2.resize(frame, (w, h))[top:bottom, left:right]
+                x1 = self.transform(frame).unsqueeze(0).to(self.device)
+                x2 = align_face(frame, self.landmarkpredictor)
+                if x2 is None:
+                    raise gr.Error(self.error_info)
+                    #return false_image, false_image, 'Error: no face detected! Please retry or change the photo.'
+                x2 = self.transform(x2).unsqueeze(dim=0).to(self.device)
+                if self.print_log: print('image loaded')
+                wplus = self.pspex.encoder(x2) + self.pspex.latent_avg.unsqueeze(0)
+                _, f =  self.pspex.encoder(x1, return_feat=True)
+            
+        with torch.no_grad():
+            y_hat, _ = self.pspex.decoder([wplus], input_is_latent=True, first_layer_feature=f)
+            y_hat = torch.clamp(y_hat, -1, 1)
+
+            if 'Style Mixing' in editing_options:
+                torch.manual_seed(seed)
+                wplus[:, 8:] = self.pspex.decoder.style(torch.randn(1, 512).to(self.device)).unsqueeze(1).repeat(1,10,1) * 0.7
+                y_hat_edit, _ = self.pspex.decoder([wplus], input_is_latent=True, first_layer_feature=f)
+            elif 'Attribute Editing' in editing_options: 
+                editing_w = self.editing_dicts[editing_options[19:]].to(self.device)
+                y_hat_edit, _ = self.pspex.decoder([wplus+scale_factor*editing_w], input_is_latent=True, first_layer_feature=f)
+            elif 'Domain Transfer' in editing_options: 
+                self.load_G_model(editing_options[17:])
+                if self.print_log: print('model %s loaded'%(editing_options[17:]))
+                y_hat_edit, _ = self.generator([wplus], input_is_latent=True, first_layer_feature=f)
+            else:
+                y_hat_edit = y_hat
+            y_hat_edit = torch.clamp(y_hat_edit, -1, 1)
+            
+        return self.tensor2np(y_hat[0]), self.tensor2np(y_hat_edit[0])