Update

.pre-commit-config.yaml

@@ -20,11 +20,11 @@ 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']

README.md

@@ -4,9 +4,10 @@ emoji: 😻
 colorFrom: purple
 colorTo: red
 sdk: gradio
-sdk_version: 3.0.15
+sdk_version: 3.36.1
 app_file: app.py
 pinned: false
+suggested_hardware: t4-small
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference

app.py

@@ -2,30 +2,16 @@
 
 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>
+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" />
 '''
-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:
@@ -57,186 +43,162 @@ def update_slider(choice: str) -> dict:
         'pixar': 121,
         'slamdunk': 119,
     }
-    return gr.Slider.update(maximum=max_vals[choice])
+    return gr.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]),
-    ]
-
+    return gr.update(value=text)
 
-def set_example_weights(example: list) -> list[dict]:
-    return [
-        gr.Slider.update(value=example[0]),
-        gr.Slider.update(value=example[1]),
-    ]
 
+model = Model()
 
-def main():
-    args = parse_args()
-    model = Model(device=args.device)
+with gr.Blocks(css='style.css') as demo:
+    gr.Markdown(DESCRIPTION)
 
-    with gr.Blocks(theme=args.theme, 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 **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.
+- 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)
+        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.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,
+                                      value=model.style_types[0])
+                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.
-    - Tips: For W+ encoder, better way of (Structure Only) is to uncheck (Structure Only) and set Color weight to 0. 
+- 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,
+        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)
-
-        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,
+                                                 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,
     )
-
-
-if __name__ == '__main__':
-    main()
+    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,
+    )
+demo.queue(max_size=20).launch()

dualstylegan.py

@@ -31,8 +31,9 @@ MODEL_REPO = 'CVPR/DualStyleGAN'
 
 
 class Model:
-    def __init__(self, device: 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_dict = self._load_encoder()
         self.transform = self._create_transform()
@@ -77,14 +78,14 @@ class Model:
         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['output_size'] = 1024
         opts = argparse.Namespace(**opts)
         model2 = pSp(opts)
         model2.to(self.device)
@@ -123,11 +124,13 @@ class Model:
         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_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
+        image.putpixel(
+            (x, y), (r, g + 1, b)
+        )  # trick to make sure run reconstruct_face() once any input setting changes
         return image
 
     @staticmethod
@@ -139,24 +142,25 @@ class Model:
         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]:
+    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 
+            return_z_plus_latent = True
         else:
             self.encoder_type = 'w+'
             z_plus_latent = False
-            return_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, 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
@@ -166,13 +170,12 @@ class Model:
                  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 
+            input_is_latent = True
 
         generator = self.generator_dict[style_type]
         exstyles = self.exstyle_dict[style_type]
@@ -187,7 +190,7 @@ class Model:
             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]
+            exstyle[:, 7:18] = instyle[:, 7:18]
 
         img_gen, _ = generator([instyle],
                                exstyle,
@@ -200,4 +203,4 @@ class Model:
                                [color_weight] * 11)
         img_gen = torch.clamp(img_gen.detach(), -1, 1)
         img_gen = self.postprocess(img_gen[0])
-        return img_gen
+        return img_gen

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.11.1
+torch==2.0.1
+torchvision==0.15.2

style.css

@@ -1,19 +1,17 @@
 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;
-}