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

@spaces.GPU
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()