ethanNeuralImage/inversion_testing

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 31from models.pti.manipulator import Manipulator32from models.pti.wrapper import 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 89G = Generator_wrapper('./pretrained_models/pti/ffhq.pkl', device)90manipulator = Manipulator(G, device)91 92 93def map_latent(mapper, inputs, stylespace=False, weight_deltas=None, strength=0.1):94    w = inputs.to(device)95    with torch.no_grad():96        if stylespace:97            delta = mapper.mapper(w)98            w_hat = [c + strength * delta_c for (c, delta_c) in zip(w, delta)]99            x_hat, _, w_hat = mapper.decoder([w_hat], input_is_latent=True, return_latents=True,100                                                randomize_noise=False, truncation=1, input_is_stylespace=True, weights_deltas=weight_deltas)101        else:102            delta = mapper.mapper(w)103            w_hat = w + strength * delta104            x_hat, w_hat, _ = mapper.decoder([w_hat], input_is_latent=True, return_latents=True,105                                                randomize_noise=False, truncation=1, weights_deltas=weight_deltas)106        result_batch = (x_hat, w_hat)107    return result_batch108def run_metrics(base_img, edited_img):109    #print(base_img.shape, edited_img.shape)110    #base_img = base_img.unsqueeze(0)111    #edited_img = edited_img.unqueeze(0)112    lpips_score = lpips_metric(base_img, edited_img)[0]113    ssim_score = ssim_metric(base_img, edited_img)[0]114    id_score = id_metric(base_img, edited_img)[0]115 116    return lpips_score, ssim_score, id_score117def clip_text_metric(tensor, text):118    clip_embed = torch.cat([clip.tokenize(text)]).cuda()119    clip_score = 1-clip_text(tensor.unsqueeze(0), clip_embed).item()120    return clip_score121 122def submit(123    src, align_img, inverter_bools, n_iterations, invert_bool,124    mapper_bool, mapper_choice, mapper_alpha, 125    gd_bool, neutral_text, target_text, alpha, beta,126    ris_bool, ref_img,127    ):128    if device == 'cuda': torch.cuda.empty_cache()129    opts.checkpoint_path = mapper_dict[mapper_choice]130    ckpt = torch.load(mapper_dict[mapper_choice], map_location='cpu')131    mapper_args = ckpt['opts']132    mapper_args.update(vars(opts))133    mapper_args = Namespace(**mapper_args)134    mapper = StyleCLIPMapper(mapper_args)135    mapper.eval()136    mapper.to(device)137    resize_to = (256, 256) if hyperstyle_args.resize_outputs else (hyperstyle_args.output_size, hyperstyle_args.output_size)138    with torch.no_grad():139        output_imgs = []140        if align_img:141            input_img = align_face(src, predictor)142        else:143            input_img = Image.open(src).convert('RGB')144        input_img = im2tensor_transforms(input_img).to(device)145 146        if gd_bool:147            opts.neutral_text = neutral_text148            opts.target_text = target_text149            opts.alpha = alpha150            opts.beta = beta151        152        if ris_bool:153            if align_img:154                ref_input = align_face(ref_img, predictor)155            else:156                ref_input = Image.open(src).convert('RGB')157            ref_input = im2tensor_transforms(ref_input).to(device)158        hyperstyle_metrics_text = ''159        if 'Hyperstyle' in inverter_bools:160            hyperstyle_batch, hyperstyle_latents, hyperstyle_deltas, _ = run_inversion(input_img.unsqueeze(0), hyperstyle, hyperstyle_args, return_intermediate_results=False)161            invert_hyperstyle = tensor2im(hyperstyle_batch[0])162            if mapper_bool:163                mapped_hyperstyle, _ = map_latent(mapper, hyperstyle_latents, stylespace=False, weight_deltas=hyperstyle_deltas, strength=mapper_alpha)164                clip_score = clip_text_metric(mapped_hyperstyle[0], mapper_args.description)165                mapped_hyperstyle = tensor2im(mapped_hyperstyle[0])166                lpips_score, ssim_score, id_score = run_metrics(invert_hyperstyle.resize(resize_to), mapped_hyperstyle.resize(resize_to))167                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}'168            else:169                mapped_hyperstyle = None170            171            if gd_bool:172                gd_hyperstyle = edit_image(_, hyperstyle_latents[0], hyperstyle.decoder, direction_calculator, opts, hyperstyle_deltas)173                clip_score = clip_text_metric(gd_hyperstyle[0], opts.target_text)174                gd_hyperstyle = tensor2im(gd_hyperstyle[0])175                lpips_score, ssim_score, id_score = run_metrics(invert_hyperstyle.resize(resize_to), gd_hyperstyle.resize(resize_to))176                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}'177            else:178                gd_hyperstyle = None179            180            if ris_bool:181 182                ref_hyperstyle_batch, ref_hyperstyle_latents, ref_hyperstyle_deltas, _ = run_inversion(ref_input.unsqueeze(0), hyperstyle, hyperstyle_args, return_intermediate_results=False)183                blend_hyperstyle, blend_hyperstyle_latents = blend_latents(hyperstyle_latents, ref_hyperstyle_latents,184                                                src_deltas=hyperstyle_deltas, ref_deltas=ref_hyperstyle_deltas,185                                                generator=ris_gen, device=device)186                ris_hyperstyle = tensor2im(blend_hyperstyle[0])187 188                lpips_score, ssim_score, id_score = run_metrics(invert_hyperstyle.resize(resize_to), ris_hyperstyle.resize(resize_to))189                clip_score = clip_hair(invert_hyperstyle.resize(resize_to), ris_hyperstyle.resize(resize_to))[1]190                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}'191            else:192                ris_hyperstyle=None193 194            hyperstyle_output = [invert_hyperstyle, mapped_hyperstyle,gd_hyperstyle, ris_hyperstyle, hyperstyle_metrics_text]195        else:196            hyperstyle_output = [None, None, None, None, hyperstyle_metrics_text]197        output_imgs.extend(hyperstyle_output)198        e4e_metrics_text = ''199        if 'E4E' in inverter_bools:200            e4e_batch, e4e_latents = hyperstyle.w_invert(input_img.unsqueeze(0))201            e4e_deltas = None202            invert_e4e = tensor2im(e4e_batch[0])203            if mapper_bool:204                mapped_e4e, _ = map_latent(mapper, e4e_latents, stylespace=False, weight_deltas=e4e_deltas, strength=mapper_alpha)205                clip_score = clip_text_metric(mapped_e4e[0], mapper_args.description)206                mapped_e4e = tensor2im(mapped_e4e[0])207                lpips_score, ssim_score, id_score = run_metrics(invert_e4e.resize(resize_to), mapped_e4e.resize(resize_to))208                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}'209            210            else:211                mapped_e4e = None212            213            if gd_bool:214                gd_e4e = edit_image(_, e4e_latents[0], hyperstyle.decoder, direction_calculator, opts, e4e_deltas)215                clip_score = clip_text_metric(gd_e4e[0], opts.target_text)216                gd_e4e = tensor2im(gd_e4e[0])217                lpips_score, ssim_score, id_score = run_metrics(invert_e4e.resize(resize_to), gd_e4e.resize(resize_to))218                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}'219            220            else:221                gd_e4e = None222            223            if ris_bool:224                ref_e4e_batch, ref_e4e_latents, = hyperstyle.w_invert(ref_input.unsqueeze(0))225                ref_e4e_deltas= None226                blend_e4e, blend_e4e_latents = blend_latents(e4e_latents, ref_e4e_latents,227                                                src_deltas=None, ref_deltas=None,228                                                generator=ris_gen, device=device)229                ris_e4e = tensor2im(blend_e4e[0])230 231                lpips_score, ssim_score, id_score = run_metrics(invert_e4e.resize(resize_to), ris_e4e.resize(resize_to))232                clip_score = clip_hair(invert_e4e.resize(resize_to), ris_e4e.resize(resize_to))[1]233                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}'234            else:235                ris_e4e=None236            237            e4e_output = [invert_e4e, mapped_e4e, gd_e4e, ris_e4e, e4e_metrics_text]238        else:239            e4e_output = [None, None, None, None, e4e_metrics_text]240        output_imgs.extend(e4e_output)241        if 'PTI' in inverter_bools:242            pti_output = None, None, None, None243            manipulator.set_real_img_projection(src, inv_mode='w+', pti_mode='s')244        else:245            pti_output = None, None, None, None246        output_imgs.extend(pti_output)247    return output_imgs