Update

.pre-commit-config.yaml

@@ -1,37 +1,61 @@
-exclude: ^patch.*
+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']
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.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.10.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.4.2
+    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.8.5
+    hooks:
+      - id: nbqa-black
+      - id: nbqa-pyupgrade
+        args: ["--py37-plus"]
+      - id: nbqa-isort
+        args: ["--float-to-top"]

.style.yapf

@@ -1,5 +0,0 @@
-[style]
-based_on_style = pep8
-blank_line_before_nested_class_or_def = false
-spaces_before_comment = 2
-split_before_logical_operator = true

.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

@@ -8,14 +8,14 @@ import gradio as gr
 
 from model import Model
 
-DESCRIPTION = '''# [HairCLIP](https://github.com/wty-ustc/HairCLIP)
+DESCRIPTION = """# [HairCLIP](https://github.com/wty-ustc/HairCLIP)
 
 <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]:
-    with open('HairCLIP/mapper/hairstyle_list.txt') as f:
+    with open("HairCLIP/mapper/hairstyle_list.txt") as f:
         lines = [line.strip() for line in f.readlines()]
         lines = [line[:-10] for line in lines]
     return lines
@@ -27,78 +27,64 @@ def set_example_image(example: list) -> dict:
 
 def update_step2_components(choice: str) -> tuple[dict, dict]:
     return (
-        gr.Dropdown.update(visible=choice in ['hairstyle', 'both']),
-        gr.Textbox.update(visible=choice in ['color', 'both']),
+        gr.Dropdown.update(visible=choice in ["hairstyle", "both"]),
+        gr.Textbox.update(visible=choice in ["color", "both"]),
     )
 
 
 model = Model()
 
-with gr.Blocks(css='style.css') as demo:
+with gr.Blocks(css="style.css") as demo:
     gr.Markdown(DESCRIPTION)
     with gr.Box():
-        gr.Markdown('## Step 1')
+        gr.Markdown("## Step 1")
         with gr.Row():
             with gr.Column():
                 with gr.Row():
-                    input_image = gr.Image(label='Input Image',
-                                           type='filepath')
+                    input_image = gr.Image(label="Input Image", type="filepath")
                 with gr.Row():
-                    preprocess_button = gr.Button('Preprocess')
+                    preprocess_button = gr.Button("Preprocess")
             with gr.Column():
-                aligned_face = gr.Image(label='Aligned Face',
-                                        type='pil',
-                                        interactive=False)
+                aligned_face = gr.Image(label="Aligned Face", type="pil", interactive=False)
             with gr.Column():
-                reconstructed_face = gr.Image(label='Reconstructed Face',
-                                              type='numpy')
+                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)
+            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')
+        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')
+                        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')
+                    hairstyle_index = gr.Dropdown(label="Hairstyle", choices=hairstyles, value="afro", type="index")
                 with gr.Row():
-                    color_description = gr.Textbox(label='Color', value='red')
+                    color_description = gr.Textbox(label="Color", value="red")
                 with gr.Row():
-                    run_button = gr.Button('Run')
+                    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)
+                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()

model.py

@@ -15,22 +15,22 @@ 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.e4e') as f:
-        subprocess.run('patch -p1'.split(), cwd='encoder4editing', stdin=f)
-    with open('patch.hairclip') as f:
-        subprocess.run('patch -p1'.split(), cwd='HairCLIP', stdin=f)
+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:
+        subprocess.run("patch -p1".split(), cwd="HairCLIP", stdin=f)
 
 app_dir = pathlib.Path(__file__).parent
 
-e4e_dir = app_dir / 'encoder4editing'
+e4e_dir = app_dir / "encoder4editing"
 sys.path.insert(0, e4e_dir.as_posix())
 
 from models.psp import pSp
 from utils.alignment import align_face
 
-hairclip_dir = app_dir / 'HairCLIP'
-mapper_dir = hairclip_dir / 'mapper'
+hairclip_dir = app_dir / "HairCLIP"
+mapper_dir = hairclip_dir / "mapper"
 sys.path.insert(0, hairclip_dir.as_posix())
 sys.path.insert(0, mapper_dir.as_posix())
 
@@ -40,8 +40,7 @@ from mapper.hairclip_mapper import HairCLIPMapper
 
 class Model:
     def __init__(self):
-        self.device = torch.device(
-            'cuda:0' if torch.cuda.is_available() else 'cpu')
+        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()
@@ -50,17 +49,16 @@ class Model:
     @staticmethod
     def _create_dlib_landmark_model():
         path = huggingface_hub.hf_hub_download(
-            'public-data/dlib_face_landmark_model',
-            'shape_predictor_68_face_landmarks.dat')
+            "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('public-data/e4e',
-                                                    'e4e_ffhq_encode.pt')
-        ckpt = torch.load(ckpt_path, map_location='cpu')
-        opts = ckpt['opts']
-        opts['device'] = self.device.type
-        opts['checkpoint_path'] = ckpt_path
+        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
+        opts["checkpoint_path"] = ckpt_path
         opts = argparse.Namespace(**opts)
         model = pSp(opts)
         model.to(self.device)
@@ -68,16 +66,15 @@ class Model:
         return model
 
     def _load_hairclip(self) -> nn.Module:
-        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
-        opts['checkpoint_path'] = ckpt_path
-        opts['editing_type'] = 'both'
-        opts['input_type'] = 'text'
-        opts['hairstyle_description'] = 'HairCLIP/mapper/hairstyle_list.txt'
-        opts['color_description'] = 'red'
+        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
+        opts["checkpoint_path"] = ckpt_path
+        opts["editing_type"] = "both"
+        opts["input_type"] = "text"
+        opts["hairstyle_description"] = "HairCLIP/mapper/hairstyle_list.txt"
+        opts["color_description"] = "red"
         opts = argparse.Namespace(**opts)
         model = HairCLIPMapper(opts)
         model.to(self.device)
@@ -86,12 +83,14 @@ class 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]),
-        ])
+        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 detect_and_align_face(self, image: str) -> PIL.Image.Image:
@@ -107,35 +106,31 @@ class Model:
         return tensor.cpu().numpy().transpose(1, 2, 0)
 
     @torch.inference_mode()
-    def reconstruct_face(
-            self, image: PIL.Image.Image) -> tuple[np.ndarray, torch.Tensor]:
+    def reconstruct_face(self, image: PIL.Image.Image) -> tuple[np.ndarray, torch.Tensor]:
         input_data = self.transform(image).unsqueeze(0).to(self.device)
-        reconstructed_images, latents = self.e4e(input_data,
-                                                 randomize_noise=False,
-                                                 return_latents=True)
+        reconstructed_images, latents = self.e4e(input_data, randomize_noise=False, return_latents=True)
         reconstructed = torch.clamp(reconstructed_images[0].detach(), -1, 1)
         reconstructed = self.postprocess(reconstructed)
         return reconstructed, latents[0]
 
     @torch.inference_mode()
-    def generate(self, editing_type: str, hairstyle_index: int,
-                 color_description: str, latent: torch.Tensor) -> np.ndarray:
+    def generate(
+        self, editing_type: str, hairstyle_index: int, color_description: str, latent: torch.Tensor
+    ) -> np.ndarray:
         opts = self.hairclip.opts
         opts.editing_type = editing_type
         opts.color_description = color_description
 
-        if editing_type == 'color':
+        if editing_type == "color":
             hairstyle_index = 0
 
         device = torch.device(opts.device)
 
-        dataset = LatentsDatasetInference(latents=latent.unsqueeze(0).cpu(),
-                                          opts=opts)
+        dataset = LatentsDatasetInference(latents=latent.unsqueeze(0).cpu(), opts=opts)
         w, hairstyle_text_inputs_list, color_text_inputs_list = dataset[0][:3]
 
         w = w.unsqueeze(0).to(device)
-        hairstyle_text_inputs = hairstyle_text_inputs_list[
-            hairstyle_index].unsqueeze(0).to(device)
+        hairstyle_text_inputs = hairstyle_text_inputs_list[hairstyle_index].unsqueeze(0).to(device)
         color_text_inputs = color_text_inputs_list[0].unsqueeze(0).to(device)
 
         hairstyle_tensor_hairmasked = torch.Tensor([0]).unsqueeze(0).to(device)

style.css

@@ -1,5 +1,6 @@
 h1 {
   text-align: center;
+  display: block;
 }
 
 img#teaser {