CVPR/DualStyleGAN

1from __future__ import annotations2 3import argparse4import os5import pathlib6import subprocess7import sys8from typing import Callable9 10import dlib11import huggingface_hub12import numpy as np13import PIL.Image14import torch15import torch.nn as nn16import torchvision.transforms as T17 18if os.getenv('SYSTEM') == 'spaces':19    os.system("sed -i '10,17d' DualStyleGAN/model/stylegan/op/fused_act.py")20    os.system("sed -i '10,17d' DualStyleGAN/model/stylegan/op/upfirdn2d.py")21 22app_dir = pathlib.Path(__file__).parent23submodule_dir = app_dir / 'DualStyleGAN'24sys.path.insert(0, submodule_dir.as_posix())25 26from model.dualstylegan import DualStyleGAN27from model.encoder.align_all_parallel import align_face28from model.encoder.psp import pSp29 30MODEL_REPO = 'CVPR/DualStyleGAN'31 32 33class Model:34    def __init__(self):35        self.device = torch.device(36            'cuda:0' if torch.cuda.is_available() else 'cpu')37        self.landmark_model = self._create_dlib_landmark_model()38        self.encoder_dict = self._load_encoder()39        self.transform = self._create_transform()40        self.encoder_type = 'z+'41 42        self.style_types = [43            'cartoon',44            'caricature',45            'anime',46            'arcane',47            'comic',48            'pixar',49            'slamdunk',50        ]51        self.generator_dict = {52            style_type: self._load_generator(style_type)53            for style_type in self.style_types54        }55        self.exstyle_dict = {56            style_type: self._load_exstylecode(style_type)57            for style_type in self.style_types58        }59 60    @staticmethod61    def _create_dlib_landmark_model():62        url = 'http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2'63        path = pathlib.Path('shape_predictor_68_face_landmarks.dat')64        if not path.exists():65            bz2_path = 'shape_predictor_68_face_landmarks.dat.bz2'66            torch.hub.download_url_to_file(url, bz2_path)67            subprocess.run(f'bunzip2 -d {bz2_path}'.split())68        return dlib.shape_predictor(path.as_posix())69 70    def _load_encoder(self) -> nn.Module:71        ckpt_path = huggingface_hub.hf_hub_download(MODEL_REPO,72                                                    'models/encoder.pt')73        ckpt = torch.load(ckpt_path, map_location='cpu')74        opts = ckpt['opts']75        opts['device'] = self.device.type76        opts['checkpoint_path'] = ckpt_path77        opts = argparse.Namespace(**opts)78        model = pSp(opts)79        model.to(self.device)80        model.eval()81 82        ckpt_path = huggingface_hub.hf_hub_download(MODEL_REPO,83                                                    'models/encoder_wplus.pt')84        ckpt = torch.load(ckpt_path, map_location='cpu')85        opts = ckpt['opts']86        opts['device'] = self.device.type87        opts['checkpoint_path'] = ckpt_path88        opts['output_size'] = 102489        opts = argparse.Namespace(**opts)90        model2 = pSp(opts)91        model2.to(self.device)92        model2.eval()93 94        return {'z+': model, 'w+': model2}95 96    @staticmethod97    def _create_transform() -> Callable:98        transform = T.Compose([99            T.Resize(256),100            T.CenterCrop(256),101            T.ToTensor(),102            T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),103        ])104        return transform105 106    def _load_generator(self, style_type: str) -> nn.Module:107        model = DualStyleGAN(1024, 512, 8, 2, res_index=6)108        ckpt_path = huggingface_hub.hf_hub_download(109            MODEL_REPO, f'models/{style_type}/generator.pt')110        ckpt = torch.load(ckpt_path, map_location='cpu')111        model.load_state_dict(ckpt['g_ema'])112        model.to(self.device)113        model.eval()114        return model115 116    @staticmethod117    def _load_exstylecode(style_type: str) -> dict[str, np.ndarray]:118        if style_type in ['cartoon', 'caricature', 'anime']:119            filename = 'refined_exstyle_code.npy'120        else:121            filename = 'exstyle_code.npy'122        path = huggingface_hub.hf_hub_download(123            MODEL_REPO, f'models/{style_type}/{filename}')124        exstyles = np.load(path, allow_pickle=True).item()125        return exstyles126 127    def detect_and_align_face(self, image_path) -> np.ndarray:128        image = align_face(filepath=image_path, predictor=self.landmark_model)129        x, y = np.random.randint(255), np.random.randint(255)130        r, g, b = image.getpixel((x, y))131        image.putpixel(132            (x, y), (r, g + 1, b)133        )  # trick to make sure run reconstruct_face() once any input setting changes134        return image135 136    @staticmethod137    def denormalize(tensor: torch.Tensor) -> torch.Tensor:138        return torch.clamp((tensor + 1) / 2 * 255, 0, 255).to(torch.uint8)139 140    def postprocess(self, tensor: torch.Tensor) -> np.ndarray:141        tensor = self.denormalize(tensor)142        return tensor.cpu().numpy().transpose(1, 2, 0)143 144    @torch.inference_mode()145    def reconstruct_face(self, image: np.ndarray,146                         encoder_type: str) -> tuple[np.ndarray, torch.Tensor]:147        if encoder_type == 'Z+ encoder (better stylization)':148            self.encoder_type = 'z+'149            z_plus_latent = True150            return_z_plus_latent = True151        else:152            self.encoder_type = 'w+'153            z_plus_latent = False154            return_z_plus_latent = False155        image = PIL.Image.fromarray(image)156        input_data = self.transform(image).unsqueeze(0).to(self.device)157        img_rec, instyle = self.encoder_dict[self.encoder_type](158            input_data,159            randomize_noise=False,160            return_latents=True,161            z_plus_latent=z_plus_latent,162            return_z_plus_latent=return_z_plus_latent,163            resize=False)164        img_rec = torch.clamp(img_rec.detach(), -1, 1)165        img_rec = self.postprocess(img_rec[0])166        return img_rec, instyle167 168    @torch.inference_mode()169    def generate(self, style_type: str, style_id: int, structure_weight: float,170                 color_weight: float, structure_only: bool,171                 instyle: torch.Tensor) -> np.ndarray:172 173        if self.encoder_type == 'z+':174            z_plus_latent = True175            input_is_latent = False176        else:177            z_plus_latent = False178            input_is_latent = True179 180        generator = self.generator_dict[style_type]181        exstyles = self.exstyle_dict[style_type]182 183        style_id = int(style_id)184        stylename = list(exstyles.keys())[style_id]185 186        latent = torch.tensor(exstyles[stylename]).to(self.device)187        if structure_only and self.encoder_type == 'z+':188            latent[0, 7:18] = instyle[0, 7:18]189        exstyle = generator.generator.style(190            latent.reshape(latent.shape[0] * latent.shape[1],191                           latent.shape[2])).reshape(latent.shape)192        if structure_only and self.encoder_type == 'w+':193            exstyle[:, 7:18] = instyle[:, 7:18]194 195        img_gen, _ = generator([instyle],196                               exstyle,197                               input_is_latent=input_is_latent,198                               z_plus_latent=z_plus_latent,199                               truncation=0.7,200                               truncation_latent=0,201                               use_res=True,202                               interp_weights=[structure_weight] * 7 +203                               [color_weight] * 11)204        img_gen = torch.clamp(img_gen.detach(), -1, 1)205        img_gen = self.postprocess(img_gen[0])206        return img_gen207