Duplicate from Gradio-Blocks/HairCLIP

Co-authored-by: hysts <[email protected]>

.gitattributes

@@ -0,0 +1,28 @@
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitmodules

@@ -0,0 +1,6 @@
+[submodule "HairCLIP"]
+	path = HairCLIP
+	url = https://github.com/wty-ustc/HairCLIP
+[submodule "encoder4editing"]
+	path = encoder4editing
+	url = https://github.com/omertov/encoder4editing

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

.style.yapf

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

HairCLIP

@@ -0,0 +1 @@
+Subproject commit 29290cf5bdca0f21ff27e0ec2e93bdd1ebbe3605

README.md

@@ -0,0 +1,16 @@
+---
+title: HairCLIP
+emoji: ⚡
+colorFrom: indigo
+colorTo: red
+sdk: gradio
+sdk_version: 3.36.1
+app_file: app.py
+pinned: false
+suggested_hardware: t4-small
+duplicated_from: Gradio-Blocks/HairCLIP
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference
+
+https://arxiv.org/abs/2112.05142

app.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import pathlib
+
+import gradio as gr
+
+from model import Model
+
+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:
+        lines = [line.strip() for line in f.readlines()]
+        lines = [line[:-10] for line in lines]
+    return lines
+
+
+def set_example_image(example: list) -> dict:
+    return gr.Image.update(value=example[0])
+
+
+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']),
+    )
+
+
+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()

encoder4editing

@@ -0,0 +1 @@
+Subproject commit 99ea50578695d2e8a1cf7259d8ee89b23eea942b

images/README.md

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

model.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import argparse
+import os
+import pathlib
+import subprocess
+import sys
+from typing import Callable, Union
+
+import dlib
+import huggingface_hub
+import numpy as np
+import PIL.Image
+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)
+
+app_dir = pathlib.Path(__file__).parent
+
+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'
+sys.path.insert(0, hairclip_dir.as_posix())
+sys.path.insert(0, mapper_dir.as_posix())
+
+from mapper.datasets.latents_dataset_inference import LatentsDatasetInference
+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.landmark_model = self._create_dlib_landmark_model()
+        self.e4e = self._load_e4e()
+        self.hairclip = self._load_hairclip()
+        self.transform = self._create_transform()
+
+    @staticmethod
+    def _create_dlib_landmark_model():
+        path = huggingface_hub.hf_hub_download(
+            '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
+        opts = argparse.Namespace(**opts)
+        model = pSp(opts)
+        model.to(self.device)
+        model.eval()
+        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'
+        opts = argparse.Namespace(**opts)
+        model = HairCLIPMapper(opts)
+        model.to(self.device)
+        model.eval()
+        return 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]),
+        ])
+        return transform
+
+    def detect_and_align_face(self, image: str) -> PIL.Image.Image:
+        image = align_face(filepath=image, predictor=self.landmark_model)
+        return image
+
+    @staticmethod
+    def denormalize(tensor: torch.Tensor) -> torch.Tensor:
+        return torch.clamp((tensor + 1) / 2 * 255, 0, 255).to(torch.uint8)
+
+    def postprocess(self, tensor: torch.Tensor) -> np.ndarray:
+        tensor = self.denormalize(tensor)
+        return tensor.cpu().numpy().transpose(1, 2, 0)
+
+    @torch.inference_mode()
+    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 = 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:
+        opts = self.hairclip.opts
+        opts.editing_type = editing_type
+        opts.color_description = color_description
+
+        if editing_type == 'color':
+            hairstyle_index = 0
+
+        device = torch.device(opts.device)
+
+        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)
+        color_text_inputs = color_text_inputs_list[0].unsqueeze(0).to(device)
+
+        hairstyle_tensor_hairmasked = torch.Tensor([0]).unsqueeze(0).to(device)
+        color_tensor_hairmasked = torch.Tensor([0]).unsqueeze(0).to(device)
+
+        w_hat = w + 0.1 * self.hairclip.mapper(
+            w,
+            hairstyle_text_inputs,
+            color_text_inputs,
+            hairstyle_tensor_hairmasked,
+            color_tensor_hairmasked,
+        )
+        x_hat, _ = self.hairclip.decoder(
+            [w_hat],
+            input_is_latent=True,
+            return_latents=True,
+            randomize_noise=False,
+            truncation=1,
+        )
+        res = torch.clamp(x_hat[0].detach(), -1, 1)
+        res = self.postprocess(res)
+        return res

packages.txt

@@ -0,0 +1,2 @@
+cmake
+ninja-build

patch.e4e

@@ -0,0 +1,131 @@
+diff --git a/models/stylegan2/op/fused_act.py b/models/stylegan2/op/fused_act.py
+index 973a84f..6854b97 100644
+--- a/models/stylegan2/op/fused_act.py
++++ b/models/stylegan2/op/fused_act.py
+@@ -2,17 +2,18 @@ import os
+ 
+ import torch
+ from torch import nn
++from torch.nn import functional as F
+ from torch.autograd import Function
+ from torch.utils.cpp_extension import load
+ 
+-module_path = os.path.dirname(__file__)
+-fused = load(
+-    'fused',
+-    sources=[
+-        os.path.join(module_path, 'fused_bias_act.cpp'),
+-        os.path.join(module_path, 'fused_bias_act_kernel.cu'),
+-    ],
+-)
++#module_path = os.path.dirname(__file__)
++#fused = load(
++#    'fused',
++#    sources=[
++#        os.path.join(module_path, 'fused_bias_act.cpp'),
++#        os.path.join(module_path, 'fused_bias_act_kernel.cu'),
++#    ],
++#)
+ 
+ 
+ class FusedLeakyReLUFunctionBackward(Function):
+@@ -82,4 +83,18 @@ class FusedLeakyReLU(nn.Module):
+ 
+ 
+ def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
+-    return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale)
++    if input.device.type == "cpu":
++        if bias is not None:
++            rest_dim = [1] * (input.ndim - bias.ndim - 1)
++            return (
++                F.leaky_relu(
++                    input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2
++                )
++                * scale
++            )
++
++        else:
++            return F.leaky_relu(input, negative_slope=0.2) * scale
++
++    else:
++        return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale)
+diff --git a/models/stylegan2/op/upfirdn2d.py b/models/stylegan2/op/upfirdn2d.py
+index 7bc5a1e..5465d1a 100644
+--- a/models/stylegan2/op/upfirdn2d.py
++++ b/models/stylegan2/op/upfirdn2d.py
+@@ -1,17 +1,18 @@
+ import os
+ 
+ import torch
++from torch.nn import functional as F
+ from torch.autograd import Function
+ from torch.utils.cpp_extension import load
+ 
+-module_path = os.path.dirname(__file__)
+-upfirdn2d_op = load(
+-    'upfirdn2d',
+-    sources=[
+-        os.path.join(module_path, 'upfirdn2d.cpp'),
+-        os.path.join(module_path, 'upfirdn2d_kernel.cu'),
+-    ],
+-)
++#module_path = os.path.dirname(__file__)
++#upfirdn2d_op = load(
++#    'upfirdn2d',
++#    sources=[
++#        os.path.join(module_path, 'upfirdn2d.cpp'),
++#        os.path.join(module_path, 'upfirdn2d_kernel.cu'),
++#    ],
++#)
+ 
+ 
+ class UpFirDn2dBackward(Function):
+@@ -97,8 +98,8 @@ class UpFirDn2d(Function):
+ 
+         ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1]))
+ 
+-        out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
+-        out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1
++        out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h + down_y) // down_y
++        out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w + down_x) // down_x
+         ctx.out_size = (out_h, out_w)
+ 
+         ctx.up = (up_x, up_y)
+@@ -140,9 +141,13 @@ class UpFirDn2d(Function):
+ 
+ 
+ def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
+-    out = UpFirDn2d.apply(
+-        input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1])
+-    )
++    if input.device.type == "cpu":
++        out = upfirdn2d_native(input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1])
++
++    else:
++        out = UpFirDn2d.apply(
++            input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1])
++        )
+ 
+     return out
+ 
+@@ -150,6 +155,9 @@ def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
+ def upfirdn2d_native(
+         input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
+ ):
++    _, channel, in_h, in_w = input.shape
++    input = input.reshape(-1, in_h, in_w, 1)
++
+     _, in_h, in_w, minor = input.shape
+     kernel_h, kernel_w = kernel.shape
+ 
+@@ -180,5 +188,9 @@ def upfirdn2d_native(
+         in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
+     )
+     out = out.permute(0, 2, 3, 1)
++    out = out[:, ::down_y, ::down_x, :]
++
++    out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h + down_y) // down_y
++    out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w + down_x) // down_x
+ 
+-    return out[:, ::down_y, ::down_x, :]
++    return out.view(-1, channel, out_h, out_w)

patch.hairclip

@@ -0,0 +1,61 @@
+diff --git a/mapper/latent_mappers.py b/mapper/latent_mappers.py
+index 56b9c55..f0dd005 100644
+--- a/mapper/latent_mappers.py
++++ b/mapper/latent_mappers.py
+@@ -19,7 +19,7 @@ class ModulationModule(Module):
+ 
+     def forward(self, x, embedding, cut_flag):
+         x = self.fc(x)
+-        x = self.norm(x) 	
++        x = self.norm(x)
+         if cut_flag == 1:
+             return x
+         gamma = self.gamma_function(embedding.float())
+@@ -39,20 +39,20 @@ class SubHairMapper(Module):
+     def forward(self, x, embedding, cut_flag=0):
+         x = self.pixelnorm(x)
+         for modulation_module in self.modulation_module_list:
+-        	x = modulation_module(x, embedding, cut_flag)        
++        	x = modulation_module(x, embedding, cut_flag)
+         return x
+ 
+-class HairMapper(Module): 
++class HairMapper(Module):
+     def __init__(self, opts):
+         super(HairMapper, self).__init__()
+         self.opts = opts
+-        self.clip_model, self.preprocess = clip.load("ViT-B/32", device="cuda")
++        self.clip_model, self.preprocess = clip.load("ViT-B/32", device=opts.device)
+         self.transform = transforms.Compose([transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))])
+         self.face_pool = torch.nn.AdaptiveAvgPool2d((224, 224))
+         self.hairstyle_cut_flag = 0
+         self.color_cut_flag = 0
+ 
+-        if not opts.no_coarse_mapper: 
++        if not opts.no_coarse_mapper:
+             self.course_mapping = SubHairMapper(opts, 4)
+         if not opts.no_medium_mapper:
+             self.medium_mapping = SubHairMapper(opts, 4)
+@@ -70,13 +70,13 @@ class HairMapper(Module):
+         elif hairstyle_tensor.shape[1] != 1:
+             hairstyle_embedding = self.gen_image_embedding(hairstyle_tensor, self.clip_model, self.preprocess).unsqueeze(1).repeat(1, 18, 1).detach()
+         else:
+-            hairstyle_embedding = torch.ones(x.shape[0], 18, 512).cuda()
++            hairstyle_embedding = torch.ones(x.shape[0], 18, 512).to(self.opts.device)
+         if color_text_inputs.shape[1] != 1:
+             color_embedding = self.clip_model.encode_text(color_text_inputs).unsqueeze(1).repeat(1, 18, 1).detach()
+         elif color_tensor.shape[1] != 1:
+             color_embedding = self.gen_image_embedding(color_tensor, self.clip_model, self.preprocess).unsqueeze(1).repeat(1, 18, 1).detach()
+         else:
+-            color_embedding = torch.ones(x.shape[0], 18, 512).cuda()
++            color_embedding = torch.ones(x.shape[0], 18, 512).to(self.opts.device)
+ 
+ 
+         if (hairstyle_text_inputs.shape[1] == 1) and (hairstyle_tensor.shape[1] == 1):
+@@ -106,4 +106,4 @@ class HairMapper(Module):
+             x_fine = torch.zeros_like(x_fine)
+ 
+         out = torch.cat([x_coarse, x_medium, x_fine], dim=1)
+-        return out
+
++        return out

requirements.txt

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

style.css

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