import hashlib import os from io import BytesIO import gradio as gr import grpc from PIL import Image from cachetools import LRUCache 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): if not face and not shape and not color: return gr.update(visible=False), gr.update(value="Need to upload a face and at least a shape or color ❗", visible=True) elif not face: return gr.update(visible=False), gr.update(value="Need to upload a face ❗", visible=True) elif not shape and not color: return gr.update(visible=False), gr.update(value="Need to upload at least a shape or color ❗", visible=True) 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' if os.environ.get('https_proxy'): del os.environ['https_proxy'] if os.environ.get('http_proxy'): del os.environ['http_proxy'] os.environ['SERVER'] = '172.16.4.26:7860' # with grpc.insecure_channel(os.environ['SERVER'], options=(('grpc.enable_http_proxy', 0),)) 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 gr.update(value=output, visible=True), gr.update(visible=False) def get_demo(): with gr.Blocks() as demo: gr.Markdown("## HairFastGan") gr.Markdown( '
' 'Official HairFastGAN Gradio demo:' '' '' '' '' '
' ) with gr.Row(): with gr.Column(): source = gr.Image(label="Source photo to try on the hairstyle", type="pil") with gr.Row(): shape = gr.Image(label="Shape photo with desired hairstyle (optional)", type="pil") color = gr.Image(label="Color photo with desired hair color (optional)", type="pil") with gr.Accordion("Advanced Options", open=False): blending = gr.Radio(["Article", "Alternative_v1", "Alternative_v2"], value='Article', label="Color Encoder 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") error_message = gr.Textbox(label="⚠️ Error ⚠️", visible=False, elem_classes="error-message") 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, error_message]) 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, share=True)