Added new dual style gann= files

.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,60 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.5.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: 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.7.5
+    hooks:
+      - id: docformatter
+        args: ["--in-place"]
+  - repo: https://github.com/pycqa/isort
+    rev: 5.13.2
+    hooks:
+      - id: isort
+        args: ["--profile", "black"]
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.8.0
+    hooks:
+      - id: mypy
+        args: ["--ignore-missing-imports"]
+        additional_dependencies:
+          [
+            "types-python-slugify",
+            "types-requests",
+            "types-PyYAML",
+            "types-pytz",
+          ]
+  - repo: https://github.com/psf/black
+    rev: 24.2.0
+    hooks:
+      - id: black
+        language_version: python3.10
+        args: ["--line-length", "119"]
+  - repo: https://github.com/kynan/nbstripout
+    rev: 0.7.1
+    hooks:
+      - id: nbstripout
+        args:
+          [
+            "--extra-keys",
+            "metadata.interpreter metadata.kernelspec cell.metadata.pycharm",
+          ]
+  - repo: https://github.com/nbQA-dev/nbQA
+    rev: 1.7.1
+    hooks:
+      - id: nbqa-black
+      - id: nbqa-pyupgrade
+        args: ["--py37-plus"]
+      - id: nbqa-isort
+        args: ["--float-to-top"]

.vscode/settings.json

@@ -0,0 +1,30 @@
+{
+    "editor.formatOnSave": true,
+    "files.insertFinalNewline": false,
+    "[python]": {
+        "editor.defaultFormatter": "ms-python.black-formatter",
+        "editor.formatOnType": true,
+        "editor.codeActionsOnSave": {
+            "source.organizeImports": "explicit"
+        }
+    },
+    "[jupyter]": {
+        "files.insertFinalNewline": false
+    },
+    "black-formatter.args": [
+        "--line-length=119"
+    ],
+    "isort.args": ["--profile", "black"],
+    "flake8.args": [
+        "--max-line-length=119"
+    ],
+    "ruff.lint.args": [
+        "--line-length=119"
+    ],
+    "notebook.output.scrolling": true,
+    "notebook.formatOnCellExecution": true,
+    "notebook.formatOnSave.enabled": true,
+        "notebook.codeActionsOnSave": {
+            "source.organizeImports": "explicit"
+        }
+}

app.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import pathlib
+
+import gradio as gr
+
+from dualstylegan import Model
+
+DESCRIPTION = """# Portrait Style Transfer with [DualStyleGAN](https://github.com/williamyang1991/DualStyleGAN)
+
+<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'<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(maximum=max_vals[choice])
+
+
+def update_style_image(style_name: str) -> dict:
+    text = get_style_image_markdown_text(style_name)
+    return gr.Markdown(value=text)
+
+
+model = Model()
+
+with gr.Blocks(css="style.css") as demo:
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Group():
+        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(
+                label="Encoder Type",
+                choices=["Z+ encoder (better stylization)", "W+ encoder (better reconstruction)"],
+                value="Z+ encoder (better stylization)",
+            )
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    input_image = gr.Image(label="Input Image", type="filepath")
+                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.State()
+
+        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.Group():
+        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, value=model.style_types[0])
+                text = get_style_image_markdown_text("cartoon")
+                style_image = gr.Markdown(value=text, latex_delimiters=[])
+                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.Group():
+        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(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", value=False)
+                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],
+            )
+
+    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,
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch()

dualstylegan.py

@@ -0,0 +1,206 @@
+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):
+        self.device = torch.device(
+            'cuda:0' if torch.cuda.is_available() else 'cpu')
+        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_path) -> np.ndarray:
+        image = align_face(filepath=image_path, 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,8 @@
+dlib==19.24.4
+gradio==4.36.1
+numpy==1.23.5
+opencv-python-headless==4.9.0.80
+Pillow==10.3.0
+scipy==1.13.1
+torch==2.0.1
+torchvision==0.15.2

style.css

@@ -0,0 +1,25 @@
+h1 {
+  text-align: center;
+  display: block;
+}
+
+#duplicate-button {
+  margin: auto;
+  color: #fff;
+  background: #1565c0;
+  border-radius: 100vh;
+}
+
+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;
+}