Update app.py

app.py

@@ -55,16 +55,31 @@ def swap_hair(source, target_1, target_2, progress=gr.Progress(track_tqdm=True))
     
 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():
-        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")
+        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()