Update app.py

app.py

@@ -74,6 +74,11 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
     # Return image, seed, and the used prompts
     return image, seed, f"Prompt used: {final_prompt}\nNegative prompt used: {full_negative_prompt}"
 
+def download_image(image):
+    """
+    Prepares the image for download as a PNG.
+    """
+    return image
 
 css = """
 #col-container {
@@ -231,7 +236,6 @@ 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,
@@ -242,11 +246,10 @@ 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()
+            inputs=result,
+            outputs=gr.File(label="Download as PNG")
         )
 
-demo.queue().launch()
+demo.queue().launch()