Initial commit

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,13 @@
 ---
 title: HairFastGAN
-emoji: ⚡
+emoji: 💈
 colorFrom: pink
 colorTo: blue
 sdk: gradio
 sdk_version: 4.31.5
 app_file: app.py
 pinned: false
+license: mit
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,144 @@
+import os
+from io import BytesIO
+
+import gradio as gr
+import grpc
+from PIL import Image
+from cachetools import LRUCache
+import hashlib
+
+from inference_pb2 import HairSwapRequest, HairSwapResponse
+from inference_pb2_grpc import HairSwapServiceStub
+from utils.shape_predictor import align_face
+
+
+def get_bytes(img):
+    if img is None:
+        return img
+
+    buffered = BytesIO()
+    img.save(buffered, format="JPEG")
+    return buffered.getvalue()
+
+
+def bytes_to_image(image: bytes) -> Image.Image:
+    image = Image.open(BytesIO(image))
+    return image
+
+
+def center_crop(img):
+    width, height = img.size
+    side = min(width, height)
+
+    left = (width - side) / 2
+    top = (height - side) / 2
+    right = (width + side) / 2
+    bottom = (height + side) / 2
+
+    img = img.crop((left, top, right, bottom))
+    return img
+
+
+def resize(name):
+    def resize_inner(img, align):
+        global align_cache
+
+        if name in align:
+            img_hash = hashlib.md5(get_bytes(img)).hexdigest()
+
+            if img_hash not in align_cache:
+                img = align_face(img, return_tensors=False)[0]
+                align_cache[img_hash] = img
+            else:
+                img = align_cache[img_hash]
+
+        elif img.size != (1024, 1024):
+            img = center_crop(img)
+            img = img.resize((1024, 1024), Image.Resampling.LANCZOS)
+
+        return img
+
+    return resize_inner
+
+
+def swap_hair(face, shape, color, blending, poisson_iters, poisson_erosion, progress=gr.Progress(track_tqdm=True)):
+    if not face or not shape and not color:
+        raise ValueError("Need to upload a face and at least a shape or color")
+
+    face_bytes, shape_bytes, color_bytes = map(lambda item: get_bytes(item), (face, shape, color))
+
+    if shape_bytes is None:
+        shape_bytes = b'face'
+    if color_bytes is None:
+        color_bytes = b'shape'
+
+    with grpc.insecure_channel(os.environ['SERVER']) as channel:
+        stub = HairSwapServiceStub(channel)
+
+        output: HairSwapResponse = stub.swap(
+            HairSwapRequest(face=face_bytes, shape=shape_bytes, color=color_bytes, blending=blending,
+                            poisson_iters=poisson_iters, poisson_erosion=poisson_erosion, use_cache=True)
+        )
+
+    output = bytes_to_image(output.image)
+    return output
+
+
+def get_demo():
+    with gr.Blocks() as demo:
+        gr.Markdown("## HairFastGan")
+        gr.Markdown(
+            '<div style="display: flex; align-items: center; gap: 10px;">'
+            '<span>Official HairFastGAN Gradio demo:</span>'
+            '<a href="https://arxiv.org/abs/2404.01094"><img src="https://img.shields.io/badge/arXiv-2404.01094-b31b1b.svg" height=22.5></a>'
+            '<a href="https://github.com/AIRI-Institute/HairFastGAN"><img src="https://img.shields.io/badge/github-%23121011.svg?style=for-the-badge&logo=github&logoColor=white" height=22.5></a>'
+            '<a href="https://huggingface.co/AIRI-Institute/HairFastGAN"><img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/model-on-hf-md.svg" height=22.5></a>'
+            '<a href="https://colab.research.google.com/#fileId=https://huggingface.co/AIRI-Institute/HairFastGAN/blob/main/notebooks/HairFast_inference.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" height=22.5></a>'
+            '</div>'
+        )
+        with gr.Row():
+            with gr.Column():
+                source = gr.Image(label="Photo that you want to replace the hair", type="pil")
+                with gr.Row():
+                    shape = gr.Image(label="Reference hair you want to get (optional)", type="pil")
+                    color = gr.Image(label="Reference color hair you want to get (optional)", type="pil")
+                with gr.Accordion("Advanced Options", open=False):
+                    blending = gr.Radio(["Article", "Alternative_v1", "Alternative_v2"], value='Article',
+                                        label="Blending version", info="Selects a model for hair color transfer.")
+                    poisson_iters = gr.Slider(0, 2500, value=0, step=1, label="Poisson iters",
+                                              info="The power of blending with the original image, helps to recover more details. Not included in the article, disabled by default.")
+                    poisson_erosion = gr.Slider(1, 100, value=15, step=1, label="Poisson erosion",
+                                                info="Smooths out the blending area.")
+                    align = gr.CheckboxGroup(["Face", "Shape", "Color"], value=["Face", "Shape", "Color"],
+                                             label="Image cropping [recommended]", info="Selects which images to crop by face")
+                btn = gr.Button("Get the haircut")
+            with gr.Column():
+                output = gr.Image(label="Your result")
+
+        gr.Examples(examples=[["input/0.png", "input/1.png", "input/2.png"], ["input/6.png", "input/7.png", None],
+                              ["input/10.jpg", None, "input/11.jpg"]],
+                    inputs=[source, shape, color], outputs=output)
+
+        source.upload(fn=resize('Face'), inputs=[source, align], outputs=source)
+        shape.upload(fn=resize('Shape'), inputs=[shape, align], outputs=shape)
+        color.upload(fn=resize('Color'), inputs=[color, align], outputs=color)
+
+        btn.click(fn=swap_hair, inputs=[source, shape, color, blending, poisson_iters, poisson_erosion], outputs=output)
+
+        gr.Markdown('''To cite the paper by the authors
+    ```
+        @article{nikolaev2024hairfastgan,
+          title={HairFastGAN: Realistic and Robust Hair Transfer with a Fast Encoder-Based Approach},
+          author={Nikolaev, Maxim and Kuznetsov, Mikhail and Vetrov, Dmitry and Alanov, Aibek},
+          journal={arXiv preprint arXiv:2404.01094},
+          year={2024}
+        }
+    ```
+        ''')
+    return demo
+
+
+if __name__ == '__main__':
+    align_cache = LRUCache(maxsize=10)
+    demo = get_demo()
+    demo.launch(server_name="0.0.0.0", server_port=7860)

inference_pb2.py

@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: inference.proto
+# Protobuf Python Version: 5.26.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0finference.proto\x12\tinference\"\x92\x01\n\x0fHairSwapRequest\x12\x0c\n\x04\x66\x61\x63\x65\x18\x01 \x01(\x0c\x12\r\n\x05shape\x18\x02 \x01(\x0c\x12\r\n\x05\x63olor\x18\x03 \x01(\x0c\x12\x10\n\x08\x62lending\x18\x04 \x01(\t\x12\x15\n\rpoisson_iters\x18\x05 \x01(\x05\x12\x17\n\x0fpoisson_erosion\x18\x06 \x01(\x05\x12\x11\n\tuse_cache\x18\x07 \x01(\x08\"!\n\x10HairSwapResponse\x12\r\n\x05image\x18\x01 \x01(\x0c\x32R\n\x0fHairSwapService\x12?\n\x04swap\x12\x1a.inference.HairSwapRequest\x1a\x1b.inference.HairSwapResponseb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'inference_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_HAIRSWAPREQUEST']._serialized_start=31
+  _globals['_HAIRSWAPREQUEST']._serialized_end=177
+  _globals['_HAIRSWAPRESPONSE']._serialized_start=179
+  _globals['_HAIRSWAPRESPONSE']._serialized_end=212
+  _globals['_HAIRSWAPSERVICE']._serialized_start=214
+  _globals['_HAIRSWAPSERVICE']._serialized_end=296
+# @@protoc_insertion_point(module_scope)

inference_pb2.pyi

@@ -0,0 +1,29 @@
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Optional as _Optional
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class HairSwapRequest(_message.Message):
+    __slots__ = ("face", "shape", "color", "blending", "poisson_iters", "poisson_erosion", "use_cache")
+    FACE_FIELD_NUMBER: _ClassVar[int]
+    SHAPE_FIELD_NUMBER: _ClassVar[int]
+    COLOR_FIELD_NUMBER: _ClassVar[int]
+    BLENDING_FIELD_NUMBER: _ClassVar[int]
+    POISSON_ITERS_FIELD_NUMBER: _ClassVar[int]
+    POISSON_EROSION_FIELD_NUMBER: _ClassVar[int]
+    USE_CACHE_FIELD_NUMBER: _ClassVar[int]
+    face: bytes
+    shape: bytes
+    color: bytes
+    blending: str
+    poisson_iters: int
+    poisson_erosion: int
+    use_cache: bool
+    def __init__(self, face: _Optional[bytes] = ..., shape: _Optional[bytes] = ..., color: _Optional[bytes] = ..., blending: _Optional[str] = ..., poisson_iters: _Optional[int] = ..., poisson_erosion: _Optional[int] = ..., use_cache: bool = ...) -> None: ...
+
+class HairSwapResponse(_message.Message):
+    __slots__ = ("image",)
+    IMAGE_FIELD_NUMBER: _ClassVar[int]
+    image: bytes
+    def __init__(self, image: _Optional[bytes] = ...) -> None: ...

inference_pb2_grpc.py

@@ -0,0 +1,101 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+import inference_pb2 as inference__pb2
+
+GRPC_GENERATED_VERSION = '1.63.0'
+GRPC_VERSION = grpc.__version__
+EXPECTED_ERROR_RELEASE = '1.65.0'
+SCHEDULED_RELEASE_DATE = 'June 25, 2024'
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    warnings.warn(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + f' but the generated code in inference_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+        + f' This warning will become an error in {EXPECTED_ERROR_RELEASE},'
+        + f' scheduled for release on {SCHEDULED_RELEASE_DATE}.',
+        RuntimeWarning
+    )
+
+
+class HairSwapServiceStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.swap = channel.unary_unary(
+                '/inference.HairSwapService/swap',
+                request_serializer=inference__pb2.HairSwapRequest.SerializeToString,
+                response_deserializer=inference__pb2.HairSwapResponse.FromString,
+                _registered_method=True)
+
+
+class HairSwapServiceServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def swap(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_HairSwapServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'swap': grpc.unary_unary_rpc_method_handler(
+                    servicer.swap,
+                    request_deserializer=inference__pb2.HairSwapRequest.FromString,
+                    response_serializer=inference__pb2.HairSwapResponse.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'inference.HairSwapService', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+
+
+ # This class is part of an EXPERIMENTAL API.
+class HairSwapService(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def swap(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/inference.HairSwapService/swap',
+            inference__pb2.HairSwapRequest.SerializeToString,
+            inference__pb2.HairSwapResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)

requirements.txt

@@ -0,0 +1,6 @@
+pillow==10.0.0
+face_alignment==1.3.4
+addict==2.4.0
+git+https://github.com/openai/CLIP.git
+gdown==3.12.2
+dlib==19.24.1