Update app.py

app.py

@@ -93,15 +93,19 @@ def get_demo():
         gr.Markdown("## HairFastGan")
         gr.Markdown(
             '<div style="display: flex; align-items: center; gap: 10px;">'
-            '<span>TDNM DEMO:</span>'
+            '<span>Official HairFastGAN Gradio demo:</span>'
+            '<a href="https://arxiv.org/abs/2404.01094"><img src="https://img.shields.io/badge/arXiv-2404.01094-b31b1b.svg" height=22.5></a>'
+            '<a href="https://github.com/AIRI-Institute/HairFastGAN"><img src="https://img.shields.io/badge/github-%23121011.svg?style=for-the-badge&logo=github&logoColor=white" height=22.5></a>'
+            '<a href="https://huggingface.co/AIRI-Institute/HairFastGAN"><img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/model-on-hf-md.svg" height=22.5></a>'
+            '<a href="https://colab.research.google.com/#fileId=https://huggingface.co/AIRI-Institute/HairFastGAN/blob/main/notebooks/HairFast_inference.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" height=22.5></a>'
             '</div>'
         )
         with gr.Row():
             with gr.Column():
-                source = gr.Image(label="Nguồn ảnh để thử kiểu tóc", type="pil")
+                source = gr.Image(label="Source photo to try on the hairstyle", type="pil")
                 with gr.Row():
-                    shape = gr.Image(label="Ảnh với kiểu tóc mong muốn (optional)", type="pil")
-                    color = gr.Image(label="Màu tóc mong muốn (optional)", type="pil")
+                    shape = gr.Image(label="Shape photo with desired hairstyle (optional)", type="pil")
+                    color = gr.Image(label="Color photo with desired hair color (optional)", type="pil")
                 with gr.Accordion("Advanced Options", open=False):
                     blending = gr.Radio(["Article", "Alternative_v1", "Alternative_v2"], value='Article',
                                         label="Color Encoder version", info="Selects a model for hair color transfer.")
@@ -128,7 +132,15 @@ def get_demo():
         btn.click(fn=swap_hair, inputs=[source, shape, color, blending, poisson_iters, poisson_erosion],
                   outputs=[output, error_message])
 
-        gr.Markdown('''Thank You
+        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}
+        }
+    ```
         ''')
     return demo
 
@@ -136,4 +148,4 @@ def get_demo():
 if __name__ == '__main__':
     align_cache = LRUCache(maxsize=10)
     demo = get_demo()
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)