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 {