Update app.py

app.py

@@ -1619,7 +1619,7 @@ if __name__ == "__main__":
 
     # GRADIO MODE
 
-    def infer(prompt, progress=gr.Progress()):
+    def infer(prompt, progress=gr.Progress(track_tqdm=True)):
         set_seed(args.seed)
         width,height = list(map(int, args.resolution.split(',')))
         pipe = AccDiffusionSDXLPipeline.from_pretrained(args.model_ckpt, torch_dtype=torch.float16).to("cuda")
@@ -1663,7 +1663,7 @@ if __name__ == "__main__":
                     )
         print(images)
 
-        return images
+        return images[0]
     
     css = """
     #col-container{
@@ -1674,9 +1674,30 @@ if __name__ == "__main__":
     with gr.Blocks(css=css) as demo:
         with gr.Column(elem_id="col-container"):
             gr.Markdown("# AccDiffusion: An Accurate Method for Higher-Resolution Image Generation")
+            gr.HTML("""
+            <div style="display:flex;column-gap:4px;">
+                <a href='https://lzhxmu.github.io/accdiffusion/accdiffusion.html'>
+                    <img src='https://img.shields.io/badge/Project-Page-blue'>
+                </a>
+                <a href='https://github.com/lzhxmu/AccDiffusion'>
+                    <img src='https://img.shields.io/badge/Code-blue'>
+                </a>
+                <a href='https://arxiv.org/abs/2407.10738v1'>
+                    <img src='https://img.shields.io/badge/Paper-Arxiv-red'>
+                </a>
+            </div>
+            """)
             prompt = gr.Textbox(label="Prompt")
             submit_btn = gr.Button("Submit")
-            output_images = gr.Gallery(format="png")
+            output_images = gr.Image(format="png")
+            gr.Examples(
+                examples = [
+                    ["Astronaut on Mars during sunset."],
+                    ["A fow peeking out from behind a bush."],
+                    ["A cute corgi on the lawn."],
+                ],
+                inputs = [prompt]
+            )
         submit_btn.click(
             fn = infer,
             inputs = [prompt],