File size: 8,222 Bytes
3a19a1a
 
 
b879b4e
 
3a19a1a
 
 
 
 
 
 
 
 
 
 
 
 
 
db3750b
3a19a1a
 
 
 
 
 
 
 
 
 
 
 
 
5479d05
 
1408f30
 
 
 
3a19a1a
 
db3750b
 
5479d05
3a19a1a
5479d05
 
 
3a19a1a
5479d05
3a19a1a
5479d05
3a19a1a
 
 
 
 
 
 
 
 
 
 
 
5479d05
 
 
3a19a1a
 
 
 
a85cdf5
3a19a1a
 
 
 
 
5479d05
3a19a1a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ea00796
3a19a1a
 
 
5479d05
3a19a1a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5479d05
3a19a1a
 
 
5479d05
3a19a1a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
720df82
3a19a1a
 
 
 
 
 
 
 
 
 
558a786
3a19a1a
 
 
 
 
 
 
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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
import os

import torch
import gradio as gr

import os
import sys
import numpy as np

from e4e.models.psp import pSp
from util import *
from huggingface_hub import hf_hub_download

import os
import sys
import tempfile
import shutil
from argparse import Namespace
from pathlib import Path
import shutil

import dlib
import numpy as np
import torchvision.transforms as transforms
from torchvision import utils
from PIL import Image

from model.sg2_model import Generator
from generate_videos import generate_frames, video_from_interpolations, vid_to_gif

model_dir = "models"
os.makedirs(model_dir, exist_ok=True)

model_repos = {"e4e": ("akhaliq/JoJoGAN_e4e_ffhq_encode", "e4e_ffhq_encode.pt"),
               "dlib": ("akhaliq/jojogan_dlib", "shape_predictor_68_face_landmarks.dat"),
               "base": ("akhaliq/jojogan-stylegan2-ffhq-config-f", "stylegan2-ffhq-config-f.pt"),
               "anime": ("rinong/stylegan-nada-models", "anime.pt"),
               "joker": ("rinong/stylegan-nada-models", "joker.pt")
               }

def get_models():
    os.makedirs(model_dir, exist_ok=True)

    model_paths = {}

    for model_name, repo_details in model_repos.items():
        download_path = hf_hub_download(repo_id=repo_details[0], filename=repo_details[1])
        model_paths[model_name] = download_path

    return model_paths

model_paths = get_models()

class ImageEditor(object):
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        latent_size = 512
        n_mlp = 8
        channel_mult = 2
        model_size = 1024

        self.generators = {}

        self.model_list = [name for name in model_paths.keys() if name not in ["e4e", "dlib"]]

        for model in self.model_list:
            g_ema = Generator(
                model_size, latent_size, n_mlp, channel_multiplier=channel_mult
            ).to(self.device)

            checkpoint = torch.load(model_paths[model], map_location=self.device)

            g_ema.load_state_dict(checkpoint['g_ema'])

            self.generators[model] = g_ema

        self.experiment_args = {"model_path": model_paths["e4e"]}
        self.experiment_args["transform"] = transforms.Compose(
            [
                transforms.Resize((256, 256)),
                transforms.ToTensor(),
                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
            ]
        )
        self.resize_dims = (256, 256)

        model_path = self.experiment_args["model_path"]

        ckpt = torch.load(model_path, map_location="cpu")
        opts = ckpt["opts"]

        opts["checkpoint_path"] = model_path
        opts = Namespace(**opts)

        self.e4e_net = pSp(opts, self.device)
        self.e4e_net.eval()

        self.shape_predictor = dlib.shape_predictor(
            model_paths["dlib"]
        )

        print("setup complete")

    def get_style_list(self):
        style_list = ['all', 'list - enter below']

        for key in self.generators:
            style_list.append(key)

        return style_list

    def predict(
        self,
        input,              # Input image path
        output_style,       # Which output style do you want to use?
        style_list,         # Comma seperated list of models to use. Only accepts models from the output_style list
        generate_video,     # Generate a video instead of an output image
        with_editing,       # Apply latent space editing to the generated video
        video_format        # Choose gif to display in browser, mp4 for higher-quality downloadable video
    ):  

        if output_style == 'all':
            styles = self.model_list
        elif output_style == 'list - enter below':
            styles = style_list.split(",")
            for style in styles:
                if style not in self.model_list:
                    raise ValueError(f"Encountered style '{style}' in the style_list which is not an available option.")
        else:
            styles = [output_style]

        # @title Align image
        input_image = self.run_alignment(str(input))
        
        input_image = input_image.resize(self.resize_dims)

        img_transforms = self.experiment_args["transform"]
        transformed_image = img_transforms(input_image)

        with torch.no_grad():
            images, latents = self.run_on_batch(transformed_image.unsqueeze(0))
            result_image, latent = images[0], latents[0]

        inverted_latent = latent.unsqueeze(0).unsqueeze(1)
        out_dir = Path(tempfile.mkdtemp())
        out_path = out_dir / "out.jpg"
        
        generators = [self.generators[style] for style in styles]

        if not generate_video:
            with torch.no_grad():
                img_list = []
                for g_ema in generators:
                    img, _ = g_ema(inverted_latent, input_is_latent=True, truncation=1, randomize_noise=False)
                    img_list.append(img)
                
                out_img = torch.cat(img_list, axis=0)
                utils.save_image(out_img, out_path, nrow=int(np.sqrt(out_img.size(0))), normalize=True, scale_each=True, range=(-1, 1))

            return Path(out_path)
        
        return self.generate_vid(generators, inverted_latent, out_dir, video_format, with_editing)

    def generate_vid(self, generators, latent, out_dir, video_format, with_editing):      
        np_latent = latent.squeeze(0).cpu().detach().numpy()
        args = {
                'fps': 24,
                'target_latents': None,
                'edit_directions': None,
                'unedited_frames': 0 if with_editing else 40 * (len(generators) - 1)
                }

        args = Namespace(**args)
        with tempfile.TemporaryDirectory() as dirpath:

            generate_frames(args, np_latent, generators, dirpath)
            video_from_interpolations(args.fps, dirpath)
            
            gen_path = Path(dirpath) / "out.mp4"
            out_path = out_dir / f"out.{video_format}"

            if video_format == 'gif':
                vid_to_gif(gen_path, out_dir, scale=256, fps=args.fps)
            else:
                shutil.copy2(gen_path, out_path) 

        return out_path

    def run_alignment(self, image_path):
        aligned_image = align_face(filepath=image_path, predictor=self.shape_predictor)
        print("Aligned image has shape: {}".format(aligned_image.size))
        return aligned_image

    def run_on_batch(self, inputs):
        images, latents = self.e4e_net(
            inputs.to(self.device).float(), randomize_noise=False, return_latents=True
        )
        return images, latents

editor = ImageEditor()

title = "StyleGAN-NADA"
description = "Gradio Demo for StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators (SIGGRAPH 2022). To use it, upload your image and select a target style. More information about the paper and training new models can be found below."

article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.00946' target='_blank'>StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators</a> | <a href='https://stylegan-nada.github.io/' target='_blank'>Project Page</a> | <a href='https://github.com/rinongal/StyleGAN-nada' target='_blank'>Code</a></p> <center><img src='https://visitor-badge.glitch.me/badge?page_id=rinong_sgnada' alt='visitor badge'></center>"

gr.Interface(editor.predict, [gr.inputs.Image(type="filepath"), 
                              gr.inputs.Dropdown(choices=editor.get_style_list(), type="value", default='base', label="Model"), 
                              gr.inputs.Textbox(lines=1, placeholder=None, default="joker,anime,modigliani", label="Style List", optional=True),
                              gr.inputs.Checkbox(default=False, label="Generate Video?", optional=False),
                              gr.inputs.Checkbox(default=False, label="With Editing?", optional=False),
                              gr.inputs.Radio(choices=["gif", "mp4"], type="value", default='mp4', label="Video Format")],
                              gr.outputs.Image(type="file"), title=title, description=description, article=article, allow_flagging=False, allow_screenshot=False).launch()