Spaces:
Runtime error
Runtime error
| import subprocess | |
| import shutil | |
| import os | |
| import gradio as gr | |
| import torchvision.transforms as T | |
| import sys | |
| import spaces | |
| from PIL import Image | |
| subprocess.run(["git", "clone", "https://github.com/AIRI-Institute/HairFastGAN"], check=True) | |
| os.chdir("HairFastGAN") | |
| subprocess.run(["git", "clone", "https://huggingface.co/AIRI-Institute/HairFastGAN"], check=True) | |
| os.chdir("HairFastGAN") | |
| subprocess.run(["git", "lfs", "pull"], check=True) | |
| os.chdir("..") | |
| shutil.move("HairFastGAN/pretrained_models", "pretrained_models") | |
| shutil.move("HairFastGAN/input", "input") | |
| shutil.rmtree("HairFastGAN") | |
| items = os.listdir() | |
| for item in items: | |
| print(item) | |
| shutil.move(item, os.path.join('..', item)) | |
| os.chdir("..") | |
| shutil.rmtree("HairFastGAN") | |
| from hair_swap import HairFast, get_parser | |
| hair_fast = HairFast(get_parser().parse_args([])) | |
| def resize(image_path): | |
| img = Image.open(image_path) | |
| square_size = 1024 | |
| left = (img.width - square_size) / 2 | |
| top = (img.height - square_size) / 2 | |
| right = (img.width + square_size) / 2 | |
| bottom = (img.height + square_size) / 2 | |
| img_cropped = img.crop((left, top, right, bottom)) | |
| return img_cropped | |
| def swap_hair(source, target_1, target_2, progress=gr.Progress(track_tqdm=True)): | |
| target_2 = target_2 if target_2 else target_1 | |
| final_image = hair_fast.swap(source, target_1, target_2) | |
| return T.functional.to_pil_image(final_image) | |
| with gr.Blocks() as demo: | |
| gr.Markdown("## HairFastGan") | |
| gr.Markdown("Gradio demo for [AIRI Institute](https://github.com/AIRI-Institute)'s HairFastGan: [Paper](https://huggingface.co/papers/2404.01094) | [GitHub](https://github.com/AIRI-Institute/HairFastGAN) | [Weights 🤗](https://huggingface.co/AIRI-Institute/HairFastGAN) | [Colab](https://colab.research.google.com/#fileId=https%3A//huggingface.co/AIRI-Institute/HairFastGAN/blob/main/notebooks/HairFast_inference.ipynb)") | |
| with gr.Row(): | |
| with gr.Column(): | |
| with gr.Row(): | |
| source = gr.Image(label="Photo that you want to replace the hair", type="filepath") | |
| target_1 = gr.Image(label="Reference hair you want to get", type="filepath") | |
| with gr.Accordion("Reference hair color", open=False): | |
| target_2 = gr.Image(label="Reference color hair you want to get (optional)", type="filepath") | |
| btn = gr.Button("Get the haircut") | |
| with gr.Column(): | |
| output = gr.Image(label="Your result") | |
| gr.Examples(examples=[["michael_cera-min.png", "leo_square-min.png", "pink_hair_celeb-min.png"]], inputs=[source, target_1, target_2], outputs=output) | |
| source.upload(fn=resize, inputs=source, outputs=source) | |
| target_1.upload(fn=resize, inputs=target_1, outputs=target_1) | |
| target_2.upload(fn=resize, inputs=target_2, outputs=target_2) | |
| btn.click(fn=swap_hair, inputs=[source, target_1, target_2], 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} | |
| } | |
| ``` | |
| ''') | |
| demo.launch() |