Split file

app.py

@@ -5,27 +5,10 @@ from __future__ import annotations
 import argparse
 import os
 import pathlib
-import sys
-from typing import Callable
 
-import dlib
 import gradio as gr
-import huggingface_hub
-import numpy as np
-import PIL.Image
-import torch
-import torch.nn as nn
-import torchvision.transforms as T
 
-if os.environ.get('SYSTEM') == 'spaces':
-    os.system("sed -i '10,17d' DualStyleGAN/model/stylegan/op/fused_act.py")
-    os.system("sed -i '10,17d' DualStyleGAN/model/stylegan/op/upfirdn2d.py")
-
-sys.path.insert(0, 'DualStyleGAN')
-
-from model.dualstylegan import DualStyleGAN
-from model.encoder.align_all_parallel import align_face
-from model.encoder.psp import pSp
+from dualstylegan import Model
 
 TOKEN = os.environ['TOKEN']
 MODEL_REPO = 'hysts/DualStyleGAN'
@@ -43,146 +26,6 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-class App:
-
-    def __init__(self, device: torch.device):
-        self.device = device
-        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(
-            'hysts/dlib_face_landmark_model',
-            'shape_predictor_68_face_landmarks.dat',
-            use_auth_token=TOKEN)
-        return dlib.shape_predictor(path)
-
-    def _load_encoder(self) -> nn.Module:
-        ckpt_path = huggingface_hub.hf_hub_download(MODEL_REPO,
-                                                    'models/encoder.pt',
-                                                    use_auth_token=TOKEN)
-        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(
-            MODEL_REPO,
-            f'models/{style_type}/generator.pt',
-            use_auth_token=TOKEN)
-        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(
-            MODEL_REPO,
-            f'models/{style_type}/{filename}',
-            use_auth_token=TOKEN)
-        exstyles = np.load(path, allow_pickle=True).item()
-        return exstyles
-
-    def detect_and_align_face(self, image) -> np.ndarray:
-        image = align_face(filepath=image.name, 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
-
-
 def get_style_image_url(style_name: str) -> str:
     base_url = 'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images'
     filenames = {
@@ -240,7 +83,7 @@ def set_example_weights(example: list) -> list[dict]:
 
 def main():
     args = parse_args()
-    app = App(device=torch.device(args.device))
+    model = Model(device=args.device)
 
     css = '''
 h1#title {
@@ -304,7 +147,8 @@ img#style-image {
 ''')
             with gr.Row():
                 with gr.Column():
-                    style_type = gr.Radio(app.style_types, label='Style Type')
+                    style_type = gr.Radio(model.style_types,
+                                          label='Style Type')
                     text = get_style_image_markdown_text('cartoon')
                     style_image = gr.Markdown(value=text)
                     style_index = gr.Slider(0,
@@ -366,10 +210,10 @@ img#style-image {
             '<center><img src="https://visitor-badge.glitch.me/badge?page_id=gradio-blocks.dualstylegan" alt="visitor badge"/></center>'
         )
 
-        detect_button.click(fn=app.detect_and_align_face,
+        detect_button.click(fn=model.detect_and_align_face,
                             inputs=input_image,
                             outputs=aligned_face)
-        reconstruct_button.click(fn=app.reconstruct_face,
+        reconstruct_button.click(fn=model.reconstruct_face,
                                  inputs=aligned_face,
                                  outputs=[reconstructed_face, instyle])
         style_type.change(fn=update_slider,
@@ -378,7 +222,7 @@ img#style-image {
         style_type.change(fn=update_style_image,
                           inputs=style_type,
                           outputs=style_image)
-        generate_button.click(fn=app.generate,
+        generate_button.click(fn=model.generate,
                               inputs=[
                                   style_type,
                                   style_index,

dualstylegan.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from typing import Callable, Union
+
+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.environ.get('SYSTEM') == 'spaces':
+    os.system("sed -i '10,17d' DualStyleGAN/model/stylegan/op/fused_act.py")
+    os.system("sed -i '10,17d' DualStyleGAN/model/stylegan/op/upfirdn2d.py")
+
+sys.path.insert(0, 'DualStyleGAN')
+
+from model.dualstylegan import DualStyleGAN
+from model.encoder.align_all_parallel import align_face
+from model.encoder.psp import pSp
+
+TOKEN = os.environ['TOKEN']
+MODEL_REPO = 'hysts/DualStyleGAN'
+
+
+class Model:
+
+    def __init__(self, device: Union[torch.device, str]):
+        self.device = torch.device(device)
+        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(
+            'hysts/dlib_face_landmark_model',
+            'shape_predictor_68_face_landmarks.dat',
+            use_auth_token=TOKEN)
+        return dlib.shape_predictor(path)
+
+    def _load_encoder(self) -> nn.Module:
+        ckpt_path = huggingface_hub.hf_hub_download(MODEL_REPO,
+                                                    'models/encoder.pt',
+                                                    use_auth_token=TOKEN)
+        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(
+            MODEL_REPO,
+            f'models/{style_type}/generator.pt',
+            use_auth_token=TOKEN)
+        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(
+            MODEL_REPO,
+            f'models/{style_type}/{filename}',
+            use_auth_token=TOKEN)
+        exstyles = np.load(path, allow_pickle=True).item()
+        return exstyles
+
+    def detect_and_align_face(self, image) -> np.ndarray:
+        image = align_face(filepath=image.name, 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