Update

.pre-commit-config.yaml

@@ -1,4 +1,4 @@
-exclude: ^(DualStyleGAN|patch)
+exclude: ^patch
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
   rev: v4.2.0
@@ -21,26 +21,17 @@ repos:
   - id: docformatter
     args: ['--in-place']
 - repo: https://github.com/pycqa/isort
-  rev: 5.10.1
+  rev: 5.12.0
   hooks:
     - id: isort
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v0.812
+  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']
-- repo: https://github.com/kynan/nbstripout
-  rev: 0.5.0
-  hooks:
-    - id: nbstripout
-      args: ['--extra-keys', 'metadata.interpreter metadata.kernelspec cell.metadata.pycharm']
-- repo: https://github.com/nbQA-dev/nbQA
-  rev: 1.3.1
-  hooks:
-    - id: nbqa-isort
-    - id: nbqa-yapf

README.md

@@ -4,9 +4,12 @@ emoji: 😻
 colorFrom: purple
 colorTo: red
 sdk: gradio
-sdk_version: 3.0.17
+sdk_version: 3.35.2
 app_file: app.py
 pinned: false
+suggested_hardware: t4-small
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference
+
+https://arxiv.org/abs/2203.13248

app.py

@@ -13,19 +13,6 @@ DESCRIPTION = '''# Portrait Style Transfer with <a href="https://github.com/will
 
 <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:
@@ -83,154 +70,127 @@ def set_example_weights(example: list) -> list[dict]:
     ]
 
 
-def main():
-    args = parse_args()
-    model = Model(device=args.device)
+model = Model()
 
-    with gr.Blocks(theme=args.theme, css='style.css') as demo:
-        gr.Markdown(DESCRIPTION)
+with gr.Blocks(css='style.css') as demo:
+    gr.Markdown(DESCRIPTION)
 
-        with gr.Box():
-            gr.Markdown('''## Step 1 (Preprocess Input Image)
+    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.
+- 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.
+- 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():
-                        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'))
-                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)
+        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(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)
+        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.
+- 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,
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    structure_weight = gr.Slider(label='Structure Weight',
+                                                 minimum=0,
+                                                 maximum=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)
-
-        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)
-        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()
+                                                 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

@@ -3,9 +3,10 @@ from __future__ import annotations
 import argparse
 import os
 import pathlib
+import shlex
 import subprocess
 import sys
-from typing import Callable, Union
+from typing import Callable
 
 import dlib
 import huggingface_hub
@@ -15,9 +16,9 @@ import torch
 import torch.nn as nn
 import torchvision.transforms as T
 
-if os.getenv('SYSTEM') == 'spaces':
+if os.getenv('SYSTEM') == 'spaces' and not torch.cuda.is_available():
     with open('patch') as f:
-        subprocess.run('patch -p1'.split(), cwd='DualStyleGAN', stdin=f)
+        subprocess.run(shlex.split('patch -p1'), cwd='DualStyleGAN', stdin=f)
 
 app_dir = pathlib.Path(__file__).parent
 submodule_dir = app_dir / 'DualStyleGAN'
@@ -27,13 +28,11 @@ from model.dualstylegan import DualStyleGAN
 from model.encoder.align_all_parallel import align_face
 from model.encoder.psp import pSp
 
-HF_TOKEN = os.environ['HF_TOKEN']
-MODEL_REPO = 'hysts/DualStyleGAN'
-
 
 class Model:
-    def __init__(self, device: Union[torch.device, str]):
-        self.device = torch.device(device)
+    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()
@@ -59,15 +58,13 @@ class Model:
     @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=HF_TOKEN)
+            '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(MODEL_REPO,
-                                                    'models/encoder.pt',
-                                                    use_auth_token=HF_TOKEN)
+        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
@@ -91,9 +88,7 @@ class Model:
     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=HF_TOKEN)
+            '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)
@@ -107,14 +102,12 @@ class Model:
         else:
             filename = 'exstyle_code.npy'
         path = huggingface_hub.hf_hub_download(
-            MODEL_REPO,
-            f'models/{style_type}/{filename}',
-            use_auth_token=HF_TOKEN)
+            '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) -> np.ndarray:
-        image = align_face(filepath=image.name, predictor=self.landmark_model)
+    def detect_and_align_face(self, image: str) -> np.ndarray:
+        image = align_face(filepath=image, predictor=self.landmark_model)
         return image
 
     @staticmethod

requirements.txt

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