File size: 5,531 Bytes
33cef83
 
 
 
26286d1
ce50ac7
4049f95
33cef83
ce50ac7
33cef83
 
 
 
 
 
 
 
 
811cb03
 
 
33cef83
26286d1
811cb03
26286d1
33cef83
 
 
 
 
 
 
ce50ac7
811cb03
33cef83
 
 
 
 
811cb03
 
 
 
 
 
 
33cef83
811cb03
 
33cef83
 
 
 
811cb03
 
33cef83
 
 
811cb03
 
 
 
 
33cef83
 
 
 
 
 
 
 
811cb03
 
 
 
 
 
 
 
33cef83
 
 
 
811cb03
 
 
33cef83
 
 
 
 
 
811cb03
 
33cef83
811cb03
 
33cef83
 
 
ce50ac7
 
33cef83
 
 
 
 
 
 
 
 
 
 
811cb03
33cef83
 
811cb03
 
 
 
 
 
 
 
33cef83
 
 
 
 
811cb03
 
 
 
 
 
 
 
 
33cef83
 
 
 
 
 
 
 
 
 
811cb03
 
 
 
 
 
 
 
 
 
 
 
33cef83
 
 
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
from __future__ import annotations

import argparse
import os
import pathlib
import shlex
import subprocess
import sys
from typing import Callable

import dlib
import huggingface_hub
import numpy as np
import PIL.Image
import torch
import torch.nn as nn
import torchvision.transforms as T

if os.getenv("SYSTEM") == "spaces" and not torch.cuda.is_available():
    with open("patch") as f:
        subprocess.run(shlex.split("patch -p1"), cwd="DualStyleGAN", stdin=f)

app_dir = pathlib.Path(__file__).parent
submodule_dir = app_dir / "DualStyleGAN"
sys.path.insert(0, submodule_dir.as_posix())

from model.dualstylegan import DualStyleGAN
from model.encoder.align_all_parallel import align_face
from model.encoder.psp import pSp


class Model:
    def __init__(self):
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.landmark_model = self._create_dlib_landmark_model()
        self.encoder = self._load_encoder()
        self.transform = self._create_transform()

        self.style_types = [
            "cartoon",
            "caricature",
            "anime",
            "arcane",
            "comic",
            "pixar",
            "slamdunk",
        ]
        self.generator_dict = {style_type: self._load_generator(style_type) for style_type in self.style_types}
        self.exstyle_dict = {style_type: self._load_exstylecode(style_type) for style_type in self.style_types}

    @staticmethod
    def _create_dlib_landmark_model():
        path = huggingface_hub.hf_hub_download(
            "public-data/dlib_face_landmark_model", "shape_predictor_68_face_landmarks.dat"
        )
        return dlib.shape_predictor(path)

    def _load_encoder(self) -> nn.Module:
        ckpt_path = huggingface_hub.hf_hub_download("public-data/DualStyleGAN", "models/encoder.pt")
        ckpt = torch.load(ckpt_path, map_location="cpu")
        opts = ckpt["opts"]
        opts["device"] = self.device.type
        opts["checkpoint_path"] = ckpt_path
        opts = argparse.Namespace(**opts)
        model = pSp(opts)
        model.to(self.device)
        model.eval()
        return model

    @staticmethod
    def _create_transform() -> Callable:
        transform = T.Compose(
            [
                T.Resize(256),
                T.CenterCrop(256),
                T.ToTensor(),
                T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
            ]
        )
        return transform

    def _load_generator(self, style_type: str) -> nn.Module:
        model = DualStyleGAN(1024, 512, 8, 2, res_index=6)
        ckpt_path = huggingface_hub.hf_hub_download("public-data/DualStyleGAN", f"models/{style_type}/generator.pt")
        ckpt = torch.load(ckpt_path, map_location="cpu")
        model.load_state_dict(ckpt["g_ema"])
        model.to(self.device)
        model.eval()
        return model

    @staticmethod
    def _load_exstylecode(style_type: str) -> dict[str, np.ndarray]:
        if style_type in ["cartoon", "caricature", "anime"]:
            filename = "refined_exstyle_code.npy"
        else:
            filename = "exstyle_code.npy"
        path = huggingface_hub.hf_hub_download("public-data/DualStyleGAN", f"models/{style_type}/{filename}")
        exstyles = np.load(path, allow_pickle=True).item()
        return exstyles

    def detect_and_align_face(self, image: str) -> np.ndarray:
        image = align_face(filepath=image, predictor=self.landmark_model)
        return image

    @staticmethod
    def denormalize(tensor: torch.Tensor) -> torch.Tensor:
        return torch.clamp((tensor + 1) / 2 * 255, 0, 255).to(torch.uint8)

    def postprocess(self, tensor: torch.Tensor) -> np.ndarray:
        tensor = self.denormalize(tensor)
        return tensor.cpu().numpy().transpose(1, 2, 0)

    @torch.inference_mode()
    def reconstruct_face(self, image: np.ndarray) -> tuple[np.ndarray, torch.Tensor]:
        image = PIL.Image.fromarray(image)
        input_data = self.transform(image).unsqueeze(0).to(self.device)
        img_rec, instyle = self.encoder(
            input_data,
            randomize_noise=False,
            return_latents=True,
            z_plus_latent=True,
            return_z_plus_latent=True,
            resize=False,
        )
        img_rec = torch.clamp(img_rec.detach(), -1, 1)
        img_rec = self.postprocess(img_rec[0])
        return img_rec, instyle

    @torch.inference_mode()
    def generate(
        self,
        style_type: str,
        style_id: int,
        structure_weight: float,
        color_weight: float,
        structure_only: bool,
        instyle: torch.Tensor,
    ) -> np.ndarray:
        generator = self.generator_dict[style_type]
        exstyles = self.exstyle_dict[style_type]

        style_id = int(style_id)
        stylename = list(exstyles.keys())[style_id]

        latent = torch.tensor(exstyles[stylename]).to(self.device)
        if structure_only:
            latent[0, 7:18] = instyle[0, 7:18]
        exstyle = generator.generator.style(
            latent.reshape(latent.shape[0] * latent.shape[1], latent.shape[2])
        ).reshape(latent.shape)

        img_gen, _ = generator(
            [instyle],
            exstyle,
            z_plus_latent=True,
            truncation=0.7,
            truncation_latent=0,
            use_res=True,
            interp_weights=[structure_weight] * 7 + [color_weight] * 11,
        )
        img_gen = torch.clamp(img_gen.detach(), -1, 1)
        img_gen = self.postprocess(img_gen[0])
        return img_gen