Update app.py

app.py

@@ -7,52 +7,46 @@ import sys
 import spaces
 from PIL import Image
 
+# Clone và di chuyển các thư mục cần thiết
 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")
 
+# Import và khởi tạo HairFast
 from hair_swap import HairFast, get_parser
-
 hair_fast = HairFast(get_parser().parse_args([]))
 
+# Hàm resize ảnh
 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
-    
+    bottom = (img.height - square_size) / 2
     img_cropped = img.crop((left, top, right, bottom))
     return img_cropped
 
+# Hàm swap hair sử dụng GPU
 @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)
-    
+
+# Xây dựng giao diện Gradio
 with gr.Blocks() as demo:
     gr.Markdown("## TDNM Demo")
     gr.Markdown("Thanks to [AIRI Institute]")
@@ -62,23 +56,15 @@ with gr.Blocks() as demo:
                 source = gr.Image(label="Ảnh nhân vật", type="filepath")
                 target_1 = gr.Image(label="Mẫu tóc", type="filepath")
             with gr.Accordion("Màu tóc (optional)", type="filepath"):
+                pass  # Đặt các thành phần tùy chọn ở đây nếu cần
             btn = gr.Button("Thực hiện")
         with gr.Column():
             output = gr.Image(label="Kết quả")
-    gr.Examples(examples=[["michael_cera-min.png", "leo_square-min.png", "pink_hair_celeb-min.png"]], inputs=[source, target_1, target_2], outputs=output)
+    gr.Examples(examples=[["michael_cera-min.png", "leo_square-min.png", "pink_hair_celeb-min.png"]], inputs=[source, target_1], 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}
-    }
-```
-    ''')
-
+    btn.click(fn=swap_hair, inputs=[source, target_1], outputs=[output])
+    gr.Markdown('''
+To cite the paper by the authors
+''')
 demo.launch()