Duplicate from CVPR/DualStyleGAN

Co-authored-by: hysts <[email protected]>

.gitattributes

@@ -0,0 +1,28 @@
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+shape_predictor_68_face_landmarks.dat*

.gitmodules

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

.pre-commit-config.yaml

@@ -0,0 +1,35 @@
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.2.0
+  hooks:
+  - id: check-executables-have-shebangs
+  - id: check-json
+  - id: check-merge-conflict
+  - id: check-shebang-scripts-are-executable
+  - id: check-toml
+  - id: check-yaml
+  - id: double-quote-string-fixer
+  - id: end-of-file-fixer
+  - id: mixed-line-ending
+    args: ['--fix=lf']
+  - id: requirements-txt-fixer
+  - id: trailing-whitespace
+- repo: https://github.com/myint/docformatter
+  rev: v1.4
+  hooks:
+  - id: docformatter
+    args: ['--in-place']
+- repo: https://github.com/pycqa/isort
+  rev: 5.10.1
+  hooks:
+    - id: isort
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v0.812
+  hooks:
+    - id: mypy
+      args: ['--ignore-missing-imports']
+- repo: https://github.com/google/yapf
+  rev: v0.32.0
+  hooks:
+  - id: yapf
+    args: ['--parallel', '--in-place']

.style.yapf

@@ -0,0 +1,5 @@
+[style]
+based_on_style = pep8
+blank_line_before_nested_class_or_def = false
+spaces_before_comment = 2
+split_before_logical_operator = true

DualStyleGAN

@@ -0,0 +1 @@
+Subproject commit d9c52c2313913352cd2e35707f72fd450bf16630

README.md

@@ -0,0 +1,13 @@
+---
+title: Portrait Style Transfer with DualStyleGAN
+emoji: 😻
+colorFrom: purple
+colorTo: red
+sdk: gradio
+sdk_version: 3.0.15
+app_file: app.py
+pinned: false
+duplicated_from: CVPR/DualStyleGAN
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference

app.py

@@ -0,0 +1,242 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import argparse
+import pathlib
+
+import gradio as gr
+
+from dualstylegan import Model
+
+DESCRIPTION = '''# Portrait Style Transfer with <a href="https://github.com/williamyang1991/DualStyleGAN">DualStyleGAN</a>
+
+<img id="overview" alt="overview" src="https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/overview.jpg" />
+'''
+FOOTER = '<img id="visitor-badge" alt="visitor badge" src="https://visitor-badge.glitch.me/badge?page_id=gradio-blocks.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('--share', action='store_true')
+    parser.add_argument('--port', type=int)
+    parser.add_argument('--disable-queue',
+                        dest='enable_queue',
+                        action='store_false')
+    return parser.parse_args()
+
+
+def get_style_image_url(style_name: str) -> str:
+    base_url = 'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images'
+    filenames = {
+        'cartoon': 'cartoon_overview.jpg',
+        'caricature': 'caricature_overview.jpg',
+        'anime': 'anime_overview.jpg',
+        'arcane': 'Reconstruction_arcane_overview.jpg',
+        'comic': 'Reconstruction_comic_overview.jpg',
+        'pixar': 'Reconstruction_pixar_overview.jpg',
+        'slamdunk': 'Reconstruction_slamdunk_overview.jpg',
+    }
+    return f'{base_url}/{filenames[style_name]}'
+
+
+def get_style_image_markdown_text(style_name: str) -> str:
+    url = get_style_image_url(style_name)
+    return f'<img id="style-image" src="{url}" alt="style image">'
+
+
+def update_slider(choice: str) -> dict:
+    max_vals = {
+        'cartoon': 316,
+        'caricature': 198,
+        'anime': 173,
+        'arcane': 99,
+        'comic': 100,
+        'pixar': 121,
+        'slamdunk': 119,
+    }
+    return gr.Slider.update(maximum=max_vals[choice])
+
+
+def update_style_image(style_name: str) -> dict:
+    text = get_style_image_markdown_text(style_name)
+    return gr.Markdown.update(value=text)
+
+
+def set_example_image(example: list) -> dict:
+    return gr.Image.update(value=example[0])
+
+
+def set_example_styles(example: list) -> list[dict]:
+    return [
+        gr.Radio.update(value=example[0]),
+        gr.Slider.update(value=example[1]),
+    ]
+
+
+def set_example_weights(example: list) -> list[dict]:
+    return [
+        gr.Slider.update(value=example[0]),
+        gr.Slider.update(value=example[1]),
+    ]
+
+
+def main():
+    args = parse_args()
+    model = Model(device=args.device)
+
+    with gr.Blocks(theme=args.theme, css='style.css') as demo:
+        gr.Markdown(DESCRIPTION)
+
+        with gr.Box():
+            gr.Markdown('''## Step 1 (Preprocess Input Image)
+
+- 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 **Preprocess** button.
+    - Choose the encoder version. Default is Z+ encoder which has better stylization performance. W+ encoder better reconstructs the input image to preserve more details.
+    - 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():
+                encoder_type = gr.Radio(choices=['Z+ encoder (better stylization)', 'W+ encoder (better reconstruction)'],
+                                         value='Z+ encoder (better stylization)',
+                                          label='Encoder Type')                
+            with gr.Row():
+                with gr.Column():
+                    with gr.Row():
+                        input_image = gr.Image(label='Input Image',
+                                               type='file')
+                    with gr.Row():
+                        preprocess_button = gr.Button('Preprocess')
+                with gr.Column():
+                    with gr.Row():
+                        aligned_face = gr.Image(label='Aligned Face',
+                                                type='numpy',
+                                                interactive=False)
+                with gr.Column():
+                    reconstructed_face = gr.Image(label='Reconstructed Face',
+                                                  type='numpy')
+                    instyle = gr.Variable()
+
+            with gr.Row():
+                paths = sorted(pathlib.Path('images').glob('*.jpg'))
+                example_images = gr.Dataset(components=[input_image],
+                                            samples=[[path.as_posix()]
+                                                     for path in paths])
+
+        with gr.Box():
+            gr.Markdown('''## Step 2 (Select Style Image)
+
+- Select **Style Type**.
+- Select **Style Image Index** from the image table below.
+''')
+            with gr.Row():
+                with gr.Column():
+                    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,
+                                            316,
+                                            value=26,
+                                            step=1,
+                                            label='Style Image Index')
+
+            with gr.Row():
+                example_styles = gr.Dataset(
+                    components=[style_type, style_index],
+                    samples=[
+                        ['cartoon', 26],
+                        ['caricature', 65],
+                        ['arcane', 63],
+                        ['pixar', 80],
+                    ])
+
+        with gr.Box():
+            gr.Markdown('''## Step 3 (Generate Style Transferred Image)
+
+- 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.
+    - Tips: For W+ encoder, better way of (Structure Only) is to uncheck (Structure Only) and set Color weight to 0. 
+- 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():
+                    result = gr.Image(label='Result')
+
+            with gr.Row():
+                example_weights = gr.Dataset(
+                    components=[structure_weight, color_weight],
+                    samples=[
+                        [0.6, 1.0],
+                        [0.3, 1.0],
+                        [0.0, 1.0],
+                        [1.0, 0.0],
+                    ])
+
+        gr.Markdown(FOOTER)
+
+        preprocess_button.click(fn=model.detect_and_align_face,
+                                inputs=[input_image],
+                                outputs=aligned_face)
+        aligned_face.change(fn=model.reconstruct_face,
+                            inputs=[aligned_face, encoder_type],
+                            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=model.generate,
+                              inputs=[
+                                  style_type,
+                                  style_index,
+                                  structure_weight,
+                                  color_weight,
+                                  structure_only,
+                                  instyle,
+                              ],
+                              outputs=result)
+        example_images.click(fn=set_example_image,
+                             inputs=example_images,
+                             outputs=example_images.components)
+        example_styles.click(fn=set_example_styles,
+                             inputs=example_styles,
+                             outputs=example_styles.components)
+        example_weights.click(fn=set_example_weights,
+                              inputs=example_weights,
+                              outputs=example_weights.components)
+
+    demo.launch(
+        enable_queue=args.enable_queue,
+        server_port=args.port,
+        share=args.share,
+    )
+
+
+if __name__ == '__main__':
+    main()

dualstylegan.py

@@ -0,0 +1,203 @@
+from __future__ import annotations
+
+import argparse
+import os
+import pathlib
+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':
+    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")
+
+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
+
+MODEL_REPO = 'CVPR/DualStyleGAN'
+
+
+class Model:
+    def __init__(self, device: torch.device | str):
+        self.device = torch.device(device)
+        self.landmark_model = self._create_dlib_landmark_model()
+        self.encoder_dict = self._load_encoder()
+        self.transform = self._create_transform()
+        self.encoder_type = 'z+'
+
+        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():
+        url = 'http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2'
+        path = pathlib.Path('shape_predictor_68_face_landmarks.dat')
+        if not path.exists():
+            bz2_path = 'shape_predictor_68_face_landmarks.dat.bz2'
+            torch.hub.download_url_to_file(url, bz2_path)
+            subprocess.run(f'bunzip2 -d {bz2_path}'.split())
+        return dlib.shape_predictor(path.as_posix())
+
+    def _load_encoder(self) -> nn.Module:
+        ckpt_path = huggingface_hub.hf_hub_download(MODEL_REPO,
+                                                    '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()
+        
+        ckpt_path = huggingface_hub.hf_hub_download(MODEL_REPO,
+                                                    'models/encoder_wplus.pt')
+        ckpt = torch.load(ckpt_path, map_location='cpu')
+        opts = ckpt['opts']
+        opts['device'] = self.device.type
+        opts['checkpoint_path'] = ckpt_path
+        opts['output_size'] = 1024   
+        opts = argparse.Namespace(**opts)
+        model2 = pSp(opts)
+        model2.to(self.device)
+        model2.eval()
+
+        return {'z+': model, 'w+': model2}
+
+    @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')
+        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}')
+        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)
+        x, y = np.random.randint(255), np.random.randint(255)
+        r, g, b = image.getpixel((x, y))
+        image.putpixel((x, y), (r, g+1, b)) # trick to make sure run reconstruct_face() once any input setting changes
+        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, encoder_type: str) -> tuple[np.ndarray, torch.Tensor]:
+        if encoder_type == 'Z+ encoder (better stylization)':
+            self.encoder_type = 'z+'
+            z_plus_latent = True
+            return_z_plus_latent = True 
+        else:
+            self.encoder_type = 'w+'
+            z_plus_latent = False
+            return_z_plus_latent = False             
+        image = PIL.Image.fromarray(image)
+        input_data = self.transform(image).unsqueeze(0).to(self.device)
+        img_rec, instyle = self.encoder_dict[self.encoder_type](input_data,
+                                        randomize_noise=False,
+                                        return_latents=True,
+                                        z_plus_latent=z_plus_latent,
+                                        return_z_plus_latent=return_z_plus_latent,
+                                        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:
+
+
+        if self.encoder_type == 'z+':
+            z_plus_latent = True
+            input_is_latent = False
+        else:
+            z_plus_latent = False
+            input_is_latent = True 
+
+        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 and self.encoder_type == 'z+':
+            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)
+        if structure_only and self.encoder_type == 'w+':
+            exstyle[:,7:18] = instyle[:,7:18]
+
+        img_gen, _ = generator([instyle],
+                               exstyle,
+                               input_is_latent=input_is_latent,
+                               z_plus_latent=z_plus_latent,
+                               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

images/README.md

@@ -0,0 +1,6 @@
+These images are freely-usable ones from [Unsplash](https://unsplash.com/).
+
+- https://unsplash.com/photos/rDEOVtE7vOs
+- https://unsplash.com/photos/et_78QkMMQs
+- https://unsplash.com/photos/ILip77SbmOE
+- https://unsplash.com/photos/95UF6LXe-Lo

packages.txt

@@ -0,0 +1,3 @@
+bzip2
+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

style.css

@@ -0,0 +1,19 @@
+h1 {
+  text-align: center;
+}
+img#overview {
+  max-width: 1000px;
+  max-height: 600px;
+  display: block;
+  margin: auto;
+}
+img#style-image {
+  max-width: 1000px;
+  max-height: 600px;
+  display: block;
+  margin: auto;
+}
+img#visitor-badge {
+  display: block;
+  margin: auto;
+}