Create app.py

app.py

@@ -0,0 +1,43 @@
+import subprocess
+import shutil
+import os
+from gradio import gr
+import torchvision.transforms as T
+import spaces
+
+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")
+
+from hair_swap import HairFast, get_parser
+
+hair_fast = HairFast(get_parser().parse_args([]))
+
+@spaces.GPU
+def swap_hair(source, target_1, target_2):
+    result = hair_fast(face_img, shape_img, color_img)
+    final_image = hair_fast.swap(face_path, shape_path, color_path)
+    return T.functional.to_pil_image(final_image)
+    
+with gr.Blocks() as demo:
+    gr.Markdown("Start typing below and then click **Run** to see the output.")
+    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")
+        target_2 = gr.Image(label="Reference color hair you want to get (optional)", type="filepath")
+    btn = gr.Button("Get the haircut")
+    output = gr.Image(label="Your result")
+    
+    btn.click(fn=update, inputs=inp, outputs=out)
+
+demo.launch()