Update app.py

app.py

@@ -131,7 +131,10 @@ with gr.Blocks(css=css) as demo:
             gr.Markdown("""# Rainbow Media X""")
 
             # Display result image at the top
-            result = gr.Image(label="Result", show_label=False, elem_id="result")
+            result = gr.Image(label="Result", show_label=False, elem_id="result", type="pil")
+
+            # Add a download button
+            download_button = gr.Button("Download as PNG", elem_id="download-button")
 
             # Add a textbox to display the prompts used for generation
             prompt_info = gr.Textbox(label="Prompts Used", lines=3, interactive=False, elem_id="prompt-info")
@@ -228,6 +231,7 @@ with gr.Blocks(css=css) as demo:
                     selected_atmosphere_tags = gr.CheckboxGroup(choices=list(atmosphere_tags.keys()), label="Atmosphere Tags")
                     tag_tab.select(lambda: "Tag Selection", inputs=None, outputs=active_tab)
 
+        # Link the "Run" button to the infer function
         run_button.click(
             infer,
             inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps,
@@ -238,4 +242,11 @@ with gr.Blocks(css=css) as demo:
             outputs=[result, seed, prompt_info]
         )
 
+        # Enable download functionality
+        download_button.click(
+            download_image,
+            inputs=[result],
+            outputs=gr.File()
+        )
+
 demo.queue().launch()