Spaces:
Running
on
A10G
Running
on
A10G
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() |