Add files

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "DualStyleGAN"]
+	path = DualStyleGAN
+	url = https://github.com/williamyang1991/DualStyleGAN

DualStyleGAN

@@ -0,0 +1 @@
+Subproject commit 64285b179d0929e301a97c2f2c438546ff49e20d

app.py

@@ -0,0 +1,351 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import argparse
+import os
+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
+
+STYLE_IMAGE_PATHS = {
+    'cartoon':
+    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/cartoon_overview.jpg',
+    'caricature':
+    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/caricature_overview.jpg',
+    'anime':
+    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/anime_overview.jpg',
+    'arcane':
+    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/Reconstruction_arcane_overview.jpg',
+    'comic':
+    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/Reconstruction_comic_overview.jpg',
+    'pixar':
+    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/Reconstruction_pixar_overview.jpg',
+    'slamdunk':
+    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/Reconstruction_slamdunk_overview.jpg',
+}
+
+TOKEN = os.environ['TOKEN']
+MODEL_REPO = 'hysts/DualStyleGAN'
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', type=str, default='cpu')
+    parser.add_argument('--theme', type=str)
+    parser.add_argument('--live', action='store_true')
+    parser.add_argument('--share', action='store_true')
+    parser.add_argument('--port', type=int)
+    parser.add_argument('--disable-queue',
+                        dest='enable_queue',
+                        action='store_false')
+    parser.add_argument('--allow-flagging', type=str, default='never')
+    return parser.parse_args()
+
+
+class App:
+
+    def __init__(self, device: torch.device):
+        self.device = device
+        self.face_detector = 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.face_detector)
+        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 update_slider(choice: str):
+    max_vals = {
+        'cartoon': 316,
+        'caricature': 198,
+        'anime': 173,
+        'arcane': 99,
+        'comic': 100,
+        'pixar': 121,
+        'slamdunk': 119,
+    }
+    return gr.Slider.update(maximum=max_vals[choice] + 1, value=26)
+
+
+def update_style_image(choice: str):
+    style_image_path = STYLE_IMAGE_PATHS[choice]
+    text = f'<center><img src="{style_image_path}" alt="style image" width="800" height="400"></center>'
+    return gr.Markdown.update(value=text)
+
+
+def main():
+    args = parse_args()
+    app = App(device=torch.device(args.device))
+
+    with gr.Blocks(theme=args.theme) as demo:
+        gr.Markdown(
+            '''<center><h1>Portrait Style Transfer with DualStyleGAN</h1></center>
+
+This is an unofficial demo app for https://github.com/williamyang1991/DualStyleGAN.
+
+<center><img src="https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/overview.jpg" alt="overview" width="800" height="400"></center>
+
+Related App: https://huggingface.co/spaces/hysts/DualStyleGAN
+''')
+
+        with gr.Box():
+            gr.Markdown('''## Step 1
+
+- Drop an image containing a near-frontal face to the **Input Image**.
+    - If there are multiple faces in the image, hit the Edit button in the upper right corner and crop the input image beforehand.
+- Hit the **Detect & Align** button.
+- Hit the **Reconstruct Face** button.
+    - The final result will be based on this **Reconstructed Face**. So, if the reconstructed image is not satisfactory, you may want to change the input image.
+''')
+            with gr.Row():
+                with gr.Column():
+                    with gr.Row():
+                        input_image = gr.Image(label='Input Image',
+                                               type='file')
+                    with gr.Row():
+                        detect_button = gr.Button('Detect & Align Face')
+                with gr.Column():
+                    with gr.Row():
+                        face_image = gr.Image(label='Aligned Face',
+                                              type='numpy')
+                    with gr.Row():
+                        reconstruct_button = gr.Button('Reconstruct Face')
+                with gr.Column():
+                    reconstructed_face = gr.Image(label='Reconstructed Face',
+                                                  type='numpy')
+                    instyle = gr.Variable()
+
+        with gr.Box():
+            gr.Markdown('''## Step 2
+
+- Select **Style Type**.
+- Select **Style Image Index** from the image table below.
+''')
+            with gr.Row():
+                with gr.Column():
+                    with gr.Column():
+                        style_type = gr.Radio(app.style_types,
+                                              label='Style Type')
+                    with gr.Column():
+                        style_index = gr.Slider(0,
+                                                317,
+                                                value=26,
+                                                step=1,
+                                                label='Style Image Index',
+                                                interactive=True)
+                    style_image_path = STYLE_IMAGE_PATHS['cartoon']
+                    text = f'<center><img src="{style_image_path}" alt="style image" width="800" height="400"></center>'
+                    style_image = gr.Markdown(value=text)
+
+        with gr.Box():
+            gr.Markdown('''## Step 3
+
+- Adjust **Structure Weight** and **Color Weight**.
+    - These are weights for the style image, so the larger the value, the closer the resulting image will be to the style image.
+- Hit the **Generate** button.
+''')
+            with gr.Row():
+                with gr.Column():
+                    with gr.Row():
+                        structure_weight = gr.Slider(0,
+                                                     1,
+                                                     value=0.6,
+                                                     step=0.1,
+                                                     label='Structure Weight')
+                    with gr.Row():
+                        color_weight = gr.Slider(0,
+                                                 1,
+                                                 value=1,
+                                                 step=0.1,
+                                                 label='Color Weight')
+                    with gr.Row():
+                        structure_only = gr.Checkbox(label='Structure Only')
+                    with gr.Row():
+                        generate_button = gr.Button('Generate')
+
+                with gr.Column():
+                    output_image = gr.Image(label='Output Image')
+
+        gr.Markdown(
+            '<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,
+                            inputs=input_image,
+                            outputs=face_image)
+        reconstruct_button.click(fn=app.reconstruct_face,
+                                 inputs=face_image,
+                                 outputs=[reconstructed_face, instyle])
+        style_type.change(fn=update_slider,
+                          inputs=style_type,
+                          outputs=style_index)
+        style_type.change(fn=update_style_image,
+                          inputs=style_type,
+                          outputs=style_image)
+        generate_button.click(fn=app.generate,
+                              inputs=[
+                                  style_type,
+                                  style_index,
+                                  structure_weight,
+                                  color_weight,
+                                  structure_only,
+                                  instyle,
+                              ],
+                              outputs=output_image)
+
+    demo.launch(
+        enable_queue=args.enable_queue,
+        server_port=args.port,
+        share=args.share,
+    )
+
+
+if __name__ == '__main__':
+    main()

packages.txt

@@ -0,0 +1,2 @@
+cmake
+ninja-build

requirements.txt

@@ -0,0 +1,7 @@
+dlib==19.23.0
+numpy==1.22.3
+opencv-python-headless==4.5.5.62
+Pillow==9.0.1
+scipy==1.8.0
+torch==1.11.0
+torchvision==0.12.0