Update

.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']

README.md

@@ -4,9 +4,12 @@ emoji: ⚡
 colorFrom: indigo
 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/2112.05142

app.py

@@ -2,24 +2,16 @@
 
 from __future__ import annotations
 
-import argparse
 import pathlib
 
 import gradio as gr
 
 from model import Model
 
+DESCRIPTION = '''# [HairCLIP](https://github.com/wty-ustc/HairCLIP)
 
-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()
+<center><img id="teaser" src="https://raw.githubusercontent.com/wty-ustc/HairCLIP/main/assets/teaser.png" alt="teaser"></center>
+'''
 
 
 def load_hairstyle_list() -> list[str]:
@@ -40,106 +32,73 @@ def update_step2_components(choice: str) -> tuple[dict, dict]:
     )
 
 
-def main():
-    args = parse_args()
-    model = Model(device=args.device)
-
-    css = '''
-h1#title {
-  text-align: center;
-}
-img#teaser {
-  max-width: 1000px;
-  max-height: 600px;
-}
-'''
-
-    with gr.Blocks(theme=args.theme, css=css) as demo:
-        gr.Markdown('''<h1 id="title">HairCLIP</h1>
-
-This is an unofficial demo for <a href="https://github.com/wty-ustc/HairCLIP">https://github.com/wty-ustc/HairCLIP</a>.
-
-<center><img id="teaser" src="https://raw.githubusercontent.com/wty-ustc/HairCLIP/main/assets/teaser.png" alt="teaser"></center>
-''')
-        with gr.Box():
-            gr.Markdown('## Step 1')
-            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():
-                    aligned_face = gr.Image(label='Aligned Face',
-                                            type='pil',
-                                            interactive=False)
-                with gr.Column():
-                    reconstructed_face = gr.Image(label='Reconstructed Face',
-                                                  type='numpy')
-                    latent = 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')
-            with gr.Row():
-                with gr.Column():
-                    with gr.Row():
-                        editing_type = gr.Radio(['hairstyle', 'color', 'both'],
-                                                value='both',
-                                                type='value',
-                                                label='Editing Type')
-                    with gr.Row():
-                        hairstyles = load_hairstyle_list()
-                        hairstyle_index = gr.Dropdown(hairstyles,
-                                                      value='afro',
-                                                      type='index',
-                                                      label='Hairstyle')
-                    with gr.Row():
-                        color_description = gr.Textbox(value='red',
-                                                       label='Color')
-                    with gr.Row():
-                        run_button = gr.Button('Run')
-
-                with gr.Column():
-                    result = gr.Image(label='Result')
-
-        gr.Markdown(
-            '<center><img src="https://visitor-badge.glitch.me/badge?page_id=gradio-blocks.hairclip" alt="visitor badge"/></center>'
-        )
-
-        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],
-                            outputs=[reconstructed_face, latent])
-        editing_type.change(fn=update_step2_components,
-                            inputs=[editing_type],
-                            outputs=[hairstyle_index, color_description])
-        run_button.click(fn=model.generate,
-                         inputs=[
-                             editing_type,
-                             hairstyle_index,
-                             color_description,
-                             latent,
-                         ],
-                         outputs=[result])
-        example_images.click(fn=set_example_image,
-                             inputs=example_images,
-                             outputs=example_images.components)
-
-    demo.launch(
-        enable_queue=args.enable_queue,
-        server_port=args.port,
-        share=args.share,
-    )
-
-
-if __name__ == '__main__':
-    main()
+model = Model()
+
+with gr.Blocks(css='style.css') as demo:
+    gr.Markdown(DESCRIPTION)
+    with gr.Box():
+        gr.Markdown('## Step 1')
+        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():
+                aligned_face = gr.Image(label='Aligned Face',
+                                        type='pil',
+                                        interactive=False)
+            with gr.Column():
+                reconstructed_face = gr.Image(label='Reconstructed Face',
+                                              type='numpy')
+                latent = 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')
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    editing_type = gr.Radio(
+                        label='Editing Type',
+                        choices=['hairstyle', 'color', 'both'],
+                        value='both',
+                        type='value')
+                with gr.Row():
+                    hairstyles = load_hairstyle_list()
+                    hairstyle_index = gr.Dropdown(label='Hairstyle',
+                                                  choices=hairstyles,
+                                                  value='afro',
+                                                  type='index')
+                with gr.Row():
+                    color_description = gr.Textbox(label='Color', value='red')
+                with gr.Row():
+                    run_button = gr.Button('Run')
+
+            with gr.Column():
+                result = gr.Image(label='Result')
+
+    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,
+                        outputs=[reconstructed_face, latent])
+    editing_type.change(fn=update_step2_components,
+                        inputs=editing_type,
+                        outputs=[hairstyle_index, color_description])
+    run_button.click(fn=model.generate,
+                     inputs=[
+                         editing_type,
+                         hairstyle_index,
+                         color_description,
+                         latent,
+                     ],
+                     outputs=result)
+
+demo.queue(max_size=10).launch()

images/README.md

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

model.py

@@ -15,7 +15,7 @@ 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.e4e') as f:
         subprocess.run('patch -p1'.split(), cwd='encoder4editing', stdin=f)
     with open('patch.hairclip') as f:
@@ -37,12 +37,11 @@ sys.path.insert(0, mapper_dir.as_posix())
 from mapper.datasets.latents_dataset_inference import LatentsDatasetInference
 from mapper.hairclip_mapper import HairCLIPMapper
 
-HF_TOKEN = os.environ['HF_TOKEN']
-
 
 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.e4e = self._load_e4e()
         self.hairclip = self._load_hairclip()
@@ -51,15 +50,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_e4e(self) -> nn.Module:
-        ckpt_path = huggingface_hub.hf_hub_download('hysts/e4e',
-                                                    'e4e_ffhq_encode.pt',
-                                                    use_auth_token=HF_TOKEN)
+        ckpt_path = huggingface_hub.hf_hub_download('public-data/e4e',
+                                                    'e4e_ffhq_encode.pt')
         ckpt = torch.load(ckpt_path, map_location='cpu')
         opts = ckpt['opts']
         opts['device'] = self.device.type
@@ -71,9 +68,8 @@ class Model:
         return model
 
     def _load_hairclip(self) -> nn.Module:
-        ckpt_path = huggingface_hub.hf_hub_download('hysts/HairCLIP',
-                                                    'hairclip.pt',
-                                                    use_auth_token=HF_TOKEN)
+        ckpt_path = huggingface_hub.hf_hub_download('public-data/HairCLIP',
+                                                    'hairclip.pt')
         ckpt = torch.load(ckpt_path, map_location='cpu')
         opts = ckpt['opts']
         opts['device'] = self.device.type
@@ -98,8 +94,8 @@ class Model:
         ])
         return transform
 
-    def detect_and_align_face(self, image) -> PIL.Image.Image:
-        image = align_face(filepath=image.name, predictor=self.landmark_model)
+    def detect_and_align_face(self, image: str) -> PIL.Image.Image:
+        image = align_face(filepath=image, predictor=self.landmark_model)
         return image
 
     @staticmethod

requirements.txt

@@ -1,8 +1,8 @@
 dlib==19.23.0
+git+https://github.com/openai/CLIP.git
 numpy==1.22.3
 opencv-python-headless==4.5.5.64
 Pillow==9.1.0
 scipy==1.8.0
 torch==1.11.0
 torchvision==0.12.0
-git+https://github.com/openai/CLIP.git

style.css

@@ -0,0 +1,8 @@
+h1 {
+  text-align: center;
+}
+
+img#teaser {
+  max-width: 1000px;
+  max-height: 600px;
+}