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 | |
| from utils.shape_predictor import align_face | |
| hair_fast = HairFast(get_parser().parse_args([])) | |
| def resize(image_path): | |
| img = Image.open(image_path) | |
| if img.size != (1024, 1024): | |
| img = align_face(img, return_tensors=False)[0] | |
| return img | |
| 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() |