Duplicate from Gradio-Blocks/DualStyleGAN

Co-authored-by: Mamta Narang <[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
+*.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
+*.wasm 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

.gitmodules

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

.pre-commit-config.yaml

@@ -0,0 +1,37 @@
+exclude: ^patch
+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.12.0
+  hooks:
+    - id: isort
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v0.991
+  hooks:
+    - id: mypy
+      args: ['--ignore-missing-imports']
+      additional_dependencies: ['types-python-slugify']
+- 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 64285b179d0929e301a97c2f2c438546ff49e20d

README.md

@@ -0,0 +1,16 @@
+---
+title: Portrait Style Transfer
+emoji: 😻
+colorFrom: purple
+colorTo: red
+sdk: gradio
+sdk_version: 3.35.2
+app_file: app.py
+pinned: false
+suggested_hardware: t4-small
+duplicated_from: Gradio-Blocks/DualStyleGAN
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference
+
+https://arxiv.org/abs/2203.13248

app.py

@@ -0,0 +1,196 @@
+#!/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" />
+'''
+
+
+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'<center><img id="style-image" src="{url}" alt="style image"></center>'
+
+
+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]),
+    ]
+
+
+model = Model()
+
+with gr.Blocks(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 **Detect & Align Face** 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='filepath')
+                with gr.Row():
+                    detect_button = gr.Button('Detect & Align Face')
+            with gr.Column():
+                with gr.Row():
+                    aligned_face = gr.Image(label='Aligned Face',
+                                            type='numpy',
+                                            interactive=False)
+                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.Row():
+            paths = sorted(pathlib.Path('images').glob('*.jpg'))
+            gr.Examples(examples=[[path.as_posix()] for path in paths],
+                        inputs=input_image)
+
+    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(label='Style Type',
+                                      choices=model.style_types)
+                text = get_style_image_markdown_text('cartoon')
+                style_image = gr.Markdown(value=text)
+                style_index = gr.Slider(label='Style Image Index',
+                                        minimum=0,
+                                        maximum=316,
+                                        step=1,
+                                        value=26)
+
+        with gr.Row():
+            gr.Examples(examples=[
+                ['cartoon', 26],
+                ['caricature', 65],
+                ['arcane', 63],
+                ['pixar', 80],
+            ],
+                        inputs=[style_type, style_index])
+
+    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.
+- Hit the **Generate** button.
+''')
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    structure_weight = gr.Slider(label='Structure Weight',
+                                                 minimum=0,
+                                                 maximum=1,
+                                                 step=0.1,
+                                                 value=0.6)
+                with gr.Row():
+                    color_weight = gr.Slider(label='Color Weight',
+                                             minimum=0,
+                                             maximum=1,
+                                             step=0.1,
+                                             value=1)
+                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():
+            gr.Examples(examples=[
+                [0.6, 1.0],
+                [0.3, 1.0],
+                [0.0, 1.0],
+                [1.0, 0.0],
+            ],
+                        inputs=[structure_weight, color_weight])
+
+    detect_button.click(fn=model.detect_and_align_face,
+                        inputs=input_image,
+                        outputs=aligned_face)
+    reconstruct_button.click(fn=model.reconstruct_face,
+                             inputs=aligned_face,
+                             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)
+demo.queue(max_size=10).launch()

dualstylegan.py

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

images/README.md

@@ -0,0 +1,7 @@
+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,2 @@
+cmake
+ninja-build

patch

@@ -0,0 +1,38 @@
+diff --git a/model/stylegan/op/fused_act.py b/model/stylegan/op/fused_act.py
+index c3b735c..67e7972 100644
+--- a/model/stylegan/op/fused_act.py
++++ b/model/stylegan/op/fused_act.py
+@@ -7,14 +7,6 @@ from torch.autograd import Function
+ from torch.utils.cpp_extension import load
+ 
+ 
+-module_path = os.path.dirname(__file__)
+-fused = load(
+-    "fused",
+-    sources=[
+-        os.path.join(module_path, "fused_bias_act.cpp"),
+-        os.path.join(module_path, "fused_bias_act_kernel.cu"),
+-    ],
+-)
+ 
+ 
+ class FusedLeakyReLUFunctionBackward(Function):
+diff --git a/model/stylegan/op/upfirdn2d.py b/model/stylegan/op/upfirdn2d.py
+index 67e0375..97565d7 100644
+--- a/model/stylegan/op/upfirdn2d.py
++++ b/model/stylegan/op/upfirdn2d.py
+@@ -7,14 +7,6 @@ from torch.autograd import Function
+ from torch.utils.cpp_extension import load
+ 
+ 
+-module_path = os.path.dirname(__file__)
+-upfirdn2d_op = load(
+-    "upfirdn2d",
+-    sources=[
+-        os.path.join(module_path, "upfirdn2d.cpp"),
+-        os.path.join(module_path, "upfirdn2d_kernel.cu"),
+-    ],
+-)
+ 
+ 
+ class UpFirDn2dBackward(Function):

requirements.txt

@@ -0,0 +1,7 @@
+dlib==19.24.2
+numpy==1.23.5
+opencv-python-headless==4.8.0.74
+Pillow==9.5.0
+scipy==1.10.1
+torch==2.0.1
+torchvision==0.15.2

style.css

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