ethanNeuralImage/inversion_testing
0
1import sys2import os3import torch4 5from metrics.metrics import ClipHair6 7sys.path.append(".")8 9from gradio_wrapper.gradio_options import GradioTestOptions10from models.hyperstyle.utils.model_utils import load_model11from models.hyperstyle.utils.common import tensor2im12from models.hyperstyle.utils.inference_utils import run_inversion13 14from hyperstyle_global_directions.edit import load_direction_calculator, edit_image15 16from torchvision import transforms17 18import gradio as gr19 20from utils.alignment import align_face21import dlib22 23from argparse import Namespace24 25from mapper.styleclip_mapper import StyleCLIPMapper26 27import ris.spherical_kmeans as spherical_kmeans28from ris.blend import blend_latents29from ris.model import Generator as RIS_Generator30 31#from models.pti.manipulator import Manipulator32#from models.pti.wrapper import Generator as Generator_wrapper33#from models.pti.e4e_projection import projection34 35from metrics import FaceMetric36from metrics.criteria.clip_loss import CLIPLoss37import clip38 39from PIL import Image40 41opts_args = ['--no_fine_mapper']42opts = GradioTestOptions().parse(opts_args)43device = 'cuda' if torch.cuda.is_available() else 'cpu'44opts.device= device45 46mapper_dict = {47 'afro':'./pretrained_models/styleCLIP_mappers/afro_hairstyle.pt',48 'bob':'./pretrained_models/styleCLIP_mappers/bob_hairstyle.pt',49 'bowl':'./pretrained_models/styleCLIP_mappers/bowl_hairstyle.pt',50 'buzz':'./pretrained_models/styleCLIP_mappers/buzz_hairstyle.pt',51 'caesar':'./pretrained_models/styleCLIP_mappers/caesar_hairstyle.pt',52 'crew':'./pretrained_models/styleCLIP_mappers/crew_hairstyle.pt',53 'pixie':'./pretrained_models/styleCLIP_mappers/pixie_hairstyle.pt',54 'straight':'./pretrained_models/styleCLIP_mappers/straight_hairstyle.pt',55 'undercut':'./pretrained_models/styleCLIP_mappers/undercut_hairstyle.pt',56 'wavy':'./pretrained_models/styleCLIP_mappers/wavy_hairstyle.pt'57}58 59mapper_descs = {60 'afro':'A face with an afro',61 'bob':'A face with a bob-cut hairstyle',62 'bowl':'A face with a bowl cut hairstyle',63 'buzz':'A face with a buzz cut hairstyle',64 'caesar':'A face with a caesar cut hairstyle',65 'crew':'A face with a crew cut hairstyle',66 'pixie':'A face with a pixie cut hairstyle',67 'straight':'A face with a straight hair hairstyle',68 'undercut':'A face with a undercut hairstyle',69 'wavy':'A face with a wavy hair hairstyle',70}71 72 73predictor = dlib.shape_predictor("./pretrained_models/hyperstyle/shape_predictor_68_face_landmarks.lfs.dat")74hyperstyle, hyperstyle_args = load_model(opts.hyperstyle_checkpoint_path, device=device, update_opts=opts)75resize_amount = (256, 256) if hyperstyle_args.resize_outputs else (hyperstyle_args.output_size, hyperstyle_args.output_size)76im2tensor_transforms = transforms.Compose([transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])77direction_calculator = load_direction_calculator(opts)78 79ris_gen = RIS_Generator(1024, 512, 8, channel_multiplier=2).to(device).eval()80ris_ckpt = torch.load('./pretrained_models/ris/stylegan2-ffhq-config-f.pt', map_location=lambda storage, loc: storage)81ris_gen.load_state_dict(ris_ckpt['g_ema'], strict=False)82 83lpips_metric = FaceMetric(metric_type='lpips', device=device)84ssim_metric = FaceMetric(metric_type='ms-ssim', device=device)85id_metric = FaceMetric(metric_type='id', device=device)86clip_hair = FaceMetric(metric_type='cliphair', device=device)87clip_text = CLIPLoss(hyperstyle_args)88 89#G = Generator_wrapper('./pretrained_models/pti/ffhq.pkl', device)90#manipulator = Manipulator(G, device)91 92 93with gr.Blocks() as demo: 94 with gr.Row() as row:95 with gr.Column() as inputs:96 source = gr.Image(label="Image to Map", type='filepath')97 align = gr.Checkbox(True, label='Align Image')98 inverter_bools = gr.CheckboxGroup(["Hyperstyle", "E4E"], value=['Hyperstyle'], label='Inverter Choices')99 n_hyperstyle_iterations = gr.Number(5, label='Number of Iterations For Hyperstyle', precision=0)100 with gr.Box():101 invert_bool = gr.Checkbox(False, label='Output Inverter Result')102 with gr.Box():103 mapper_bool = gr.Checkbox(True, label='Output Mapper Result')104 with gr.Box() as mapper_opts:105 mapper_choice = gr.Dropdown(list(mapper_dict.keys()), value='afro', label='What Hairstyle Mapper to Use?')106 mapper_alpha = gr.Slider(minimum=-0.5, maximum=0.5, value=0.1, step=0.01, label='Strength of Mapper Alpha',)107 with gr.Box():108 gd_bool = gr.Checkbox(False, label='Output Global Direction Result')109 with gr.Box(visible=False) as gd_opts:110 neutral_text = gr.Text(value='A face with hair', label='Neutral Text')111 target_text = gr.Text(value=mapper_descs['afro'], label='Target Text')112 alpha = gr.Slider(minimum=-10.0, maximum=10.0, value=4.1, step=0.1, label="Alpha for Global Direction")113 beta = gr.Slider(minimum=0.0, maximum=0.30, value=0.15, step=0.01, label="Beta for Global Direction")114 with gr.Box():115 ris_bool = gr.Checkbox(False, label='Output RIS Result')116 with gr.Box(visible=False) as ris_opts:117 ref_img = gr.Image(label='Refrence Image for Hair', type='filepath')118 submit_button = gr.Button("Edit Image")119 with gr.Column() as outputs:120 with gr.Row() as hyperstyle_images:121 output_hyperstyle_invert = gr.Image(type='pil', label="Hyperstyle Inverted", visible=False)122 output_hyperstyle_mapper = gr.Image(type='pil', label="Hyperstyle Mapper")123 output_hyperstyle_gd = gr.Image(type='pil', label="Hyperstyle Global Directions", visible=False)124 output_hyperstyle_ris = gr.Image(type='pil', label='Hyperstyle RIS', visible=False)125 with gr.Row() as hyperstyle_metrics:126 output_hypersyle_metrics = gr.Text(label='Hyperstyle Metrics')127 with gr.Row(visible=False) as e4e_images:128 output_e4e_invert = gr.Image(type='pil', label="E4E Inverted", visible=False)129 output_e4e_mapper = gr.Image(type='pil', label="E4E Mapper")130 output_e4e_gd = gr.Image(type='pil', label="E4E Global Directions", visible=False)131 output_e4e_ris = gr.Image(type='pil', label='E4E RIS', visible=False)132 with gr.Row(visible=False) as e4e_metrics:133 output_e4e_metrics = gr.Text(label='E4E Metrics')134 with gr.Row(visible=False) as pti_images:135 output_pti_invert = gr.Image(type='pil', label="PTI Inverted", visible=False)136 output_pti_mapper = gr.Image(type='pil', label="PTI Mapper")137 output_pti_gd = gr.Image(type='pil', label="PTI Global Directions", visible=False)138 output_pti_ris = gr.Image(type='pil', label='PTI RIS', visible=False)139 with gr.Row(visible=False) as pti_metrics:140 output_pti_metrics = gr.Text(label='PTI Metrics')141 def n_iter_change(number):142 if number < 0:143 return 0144 else:145 return number146 def mapper_change(new_mapper):147 return mapper_descs[new_mapper]148 def inverter_toggles(bools):149 e4e_bool = 'E4E' in bools150 hyperstyle_bool = 'Hyperstyle' in bools151 return {152 hyperstyle_images: gr.update(visible=hyperstyle_bool),153 hyperstyle_metrics: gr.update(visible=hyperstyle_bool),154 e4e_images: gr.update(visible=e4e_bool),155 e4e_metrics: gr.update(visible=e4e_bool),156 n_hyperstyle_iterations: gr.update(visible=hyperstyle_bool)157 }158 def outp_toggles(bool):159 return {160 output_hyperstyle_invert: gr.update(visible=bool),161 output_e4e_invert: gr.update(visible=bool)162 }163 def mapper_toggles(bool):164 return {165 mapper_opts: gr.update(visible=bool),166 output_hyperstyle_mapper: gr.update(visible=bool),167 output_e4e_mapper: gr.update(visible=bool)168 }169 def gd_toggles(bool):170 return {171 gd_opts: gr.update(visible=bool),172 output_hyperstyle_gd: gr.update(visible=bool),173 output_e4e_gd: gr.update(visible=bool)174 }175 def ris_toggles(bool):176 return {177 ris_opts: gr.update(visible=bool),178 output_hyperstyle_ris: gr.update(visible=bool),179 output_e4e_ris: gr.update(visible=bool)180 }181 182 n_hyperstyle_iterations.change(n_iter_change, n_hyperstyle_iterations, n_hyperstyle_iterations)183 mapper_choice.change(mapper_change, mapper_choice, [target_text])184 inverter_bools.change(inverter_toggles, inverter_bools, [hyperstyle_images, hyperstyle_metrics, e4e_images, e4e_metrics, n_hyperstyle_iterations])185 invert_bool.change(outp_toggles, invert_bool, [output_hyperstyle_invert, output_e4e_invert])186 mapper_bool.change(mapper_toggles, mapper_bool, [mapper_opts, output_hyperstyle_mapper, output_e4e_mapper])187 gd_bool.change(gd_toggles, gd_bool, [gd_opts, output_hyperstyle_gd, output_e4e_gd])188 ris_bool.change(ris_toggles, ris_bool, [ris_opts, output_hyperstyle_ris, output_e4e_ris])189 def map_latent(mapper, inputs, stylespace=False, weight_deltas=None, strength=0.1):190 w = inputs.to(device)191 with torch.no_grad():192 if stylespace:193 delta = mapper.mapper(w)194 w_hat = [c + strength * delta_c for (c, delta_c) in zip(w, delta)]195 x_hat, _, w_hat = mapper.decoder([w_hat], input_is_latent=True, return_latents=True,196 randomize_noise=False, truncation=1, input_is_stylespace=True, weights_deltas=weight_deltas)197 else:198 delta = mapper.mapper(w)199 w_hat = w + strength * delta200 x_hat, w_hat, _ = mapper.decoder([w_hat], input_is_latent=True, return_latents=True,201 randomize_noise=False, truncation=1, weights_deltas=weight_deltas)202 result_batch = (x_hat, w_hat)203 return result_batch204 def run_metrics(base_img, edited_img):205 #print(base_img.shape, edited_img.shape)206 #base_img = base_img.unsqueeze(0)207 #edited_img = edited_img.unqueeze(0)208 lpips_score = lpips_metric(base_img, edited_img)[0]209 ssim_score = ssim_metric(base_img, edited_img)[0]210 id_score = id_metric(base_img, edited_img)[0]211 212 return lpips_score, ssim_score, id_score213 def clip_text_metric(tensor, text):214 clip_embed = torch.cat([clip.tokenize(text)]).cpu()215 clip_score = 1-clip_text(tensor.unsqueeze(0), clip_embed).item()216 return clip_score217 218 def submit(219 src, align_img, inverter_bools, n_iterations, invert_bool,220 mapper_bool, mapper_choice, mapper_alpha, 221 gd_bool, neutral_text, target_text, alpha, beta,222 ris_bool, ref_img,223 ):224 if device == 'cuda': torch.cuda.empty_cache()225 opts.checkpoint_path = mapper_dict[mapper_choice]226 ckpt = torch.load(mapper_dict[mapper_choice], map_location='cpu')227 mapper_args = ckpt['opts']228 mapper_args.update(vars(opts))229 mapper_args = Namespace(**mapper_args)230 mapper = StyleCLIPMapper(mapper_args)231 mapper.eval()232 mapper.to(device)233 resize_to = (256, 256) if hyperstyle_args.resize_outputs else (hyperstyle_args.output_size, hyperstyle_args.output_size)234 with torch.no_grad():235 output_imgs = []236 if align_img:237 input_img = align_face(src, predictor)238 else:239 input_img = Image.open(src).convert('RGB')240 input_img = im2tensor_transforms(input_img).to(device)241 242 if gd_bool:243 opts.neutral_text = neutral_text244 opts.target_text = target_text245 opts.alpha = alpha246 opts.beta = beta247 248 if ris_bool:249 if align_img:250 ref_input = align_face(ref_img, predictor)251 else:252 ref_input = Image.open(src).convert('RGB')253 ref_input = im2tensor_transforms(ref_input).to(device)254 hyperstyle_metrics_text = ''255 if 'Hyperstyle' in inverter_bools:256 hyperstyle_batch, hyperstyle_latents, hyperstyle_deltas, _ = run_inversion(input_img.unsqueeze(0), hyperstyle, hyperstyle_args, return_intermediate_results=False)257 invert_hyperstyle = tensor2im(hyperstyle_batch[0])258 if mapper_bool:259 mapped_hyperstyle, _ = map_latent(mapper, hyperstyle_latents, stylespace=False, weight_deltas=hyperstyle_deltas, strength=mapper_alpha)260 clip_score = clip_text_metric(mapped_hyperstyle[0], mapper_args.description)261 mapped_hyperstyle = tensor2im(mapped_hyperstyle[0])262 lpips_score, ssim_score, id_score = run_metrics(invert_hyperstyle.resize(resize_to), mapped_hyperstyle.resize(resize_to))263 hyperstyle_metrics_text += f'\nMapper Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Text Score: \t{clip_score}'264 else:265 mapped_hyperstyle = None266 267 if gd_bool:268 gd_hyperstyle = edit_image(_, hyperstyle_latents[0], hyperstyle.decoder, direction_calculator, opts, hyperstyle_deltas)269 clip_score = clip_text_metric(gd_hyperstyle[0], opts.target_text)270 gd_hyperstyle = tensor2im(gd_hyperstyle[0])271 lpips_score, ssim_score, id_score = run_metrics(invert_hyperstyle.resize(resize_to), gd_hyperstyle.resize(resize_to))272 hyperstyle_metrics_text += f'\nGlobal Direction Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Text Score: \t{clip_score}'273 else:274 gd_hyperstyle = None275 276 if ris_bool:277 278 ref_hyperstyle_batch, ref_hyperstyle_latents, ref_hyperstyle_deltas, _ = run_inversion(ref_input.unsqueeze(0), hyperstyle, hyperstyle_args, return_intermediate_results=False)279 blend_hyperstyle, blend_hyperstyle_latents = blend_latents(hyperstyle_latents, ref_hyperstyle_latents,280 src_deltas=hyperstyle_deltas, ref_deltas=ref_hyperstyle_deltas,281 generator=ris_gen, device=device)282 ris_hyperstyle = tensor2im(blend_hyperstyle[0])283 284 lpips_score, ssim_score, id_score = run_metrics(invert_hyperstyle.resize(resize_to), ris_hyperstyle.resize(resize_to))285 clip_score = clip_hair(invert_hyperstyle.resize(resize_to), ris_hyperstyle.resize(resize_to))[1]286 hyperstyle_metrics_text += f'\nRIS Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Hair Score: \t{clip_score}'287 else:288 ris_hyperstyle=None289 290 hyperstyle_output = [invert_hyperstyle, mapped_hyperstyle,gd_hyperstyle, ris_hyperstyle, hyperstyle_metrics_text]291 else:292 hyperstyle_output = [None, None, None, None, hyperstyle_metrics_text]293 output_imgs.extend(hyperstyle_output)294 e4e_metrics_text = ''295 if 'E4E' in inverter_bools:296 e4e_batch, e4e_latents = hyperstyle.w_invert(input_img.unsqueeze(0))297 e4e_deltas = None298 invert_e4e = tensor2im(e4e_batch[0])299 if mapper_bool:300 mapped_e4e, _ = map_latent(mapper, e4e_latents, stylespace=False, weight_deltas=e4e_deltas, strength=mapper_alpha)301 clip_score = clip_text_metric(mapped_e4e[0], mapper_args.description)302 mapped_e4e = tensor2im(mapped_e4e[0])303 lpips_score, ssim_score, id_score = run_metrics(invert_e4e.resize(resize_to), mapped_e4e.resize(resize_to))304 e4e_metrics_text += f'\nMapper Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Text Score: \t{clip_score}'305 306 else:307 mapped_e4e = None308 309 if gd_bool:310 gd_e4e = edit_image(_, e4e_latents[0], hyperstyle.decoder, direction_calculator, opts, e4e_deltas)311 clip_score = clip_text_metric(gd_e4e[0], opts.target_text)312 gd_e4e = tensor2im(gd_e4e[0])313 lpips_score, ssim_score, id_score = run_metrics(invert_e4e.resize(resize_to), gd_e4e.resize(resize_to))314 e4e_metrics_text += f'\nGlobal Direction Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Text Score: \t{clip_score}'315 316 else:317 gd_e4e = None318 319 if ris_bool:320 ref_e4e_batch, ref_e4e_latents, = hyperstyle.w_invert(ref_input.unsqueeze(0))321 ref_e4e_deltas= None322 blend_e4e, blend_e4e_latents = blend_latents(e4e_latents, ref_e4e_latents,323 src_deltas=None, ref_deltas=None,324 generator=ris_gen, device=device)325 ris_e4e = tensor2im(blend_e4e[0])326 327 lpips_score, ssim_score, id_score = run_metrics(invert_e4e.resize(resize_to), ris_e4e.resize(resize_to))328 clip_score = clip_hair(invert_e4e.resize(resize_to), ris_e4e.resize(resize_to))[1]329 e4e_metrics_text += f'\nRIS Metrics:\n\tLPIPS: \t{lpips_score} \n\tSSIM: \t{ssim_score}\n\tID Score: \t{id_score}\n\tCLIP Hair Score: \t{clip_score}'330 else:331 ris_e4e=None332 333 e4e_output = [invert_e4e, mapped_e4e, gd_e4e, ris_e4e, e4e_metrics_text]334 else:335 e4e_output = [None, None, None, None, e4e_metrics_text]336 output_imgs.extend(e4e_output)337 return output_imgs338 submit_button.click(339 submit, 340 [341 source, align, inverter_bools, n_hyperstyle_iterations, invert_bool,342 mapper_bool, mapper_choice, mapper_alpha,343 gd_bool, neutral_text, target_text, alpha, beta,344 ris_bool, ref_img345 ],346 [347 output_hyperstyle_invert, output_hyperstyle_mapper, output_hyperstyle_gd, output_hyperstyle_ris, output_hypersyle_metrics,348 output_e4e_invert, output_e4e_mapper, output_e4e_gd, output_e4e_ris, output_e4e_metrics,349 ]350 )351 352demo.launch()