Update app.py

app.py

@@ -381,7 +381,14 @@ def generate_images_with_progress(prompt, negative_prompt, batch_count, use_cont
     # Ensure all generated images are in PIL format
     pil_images = [img if isinstance(img, Image.Image) else Image.fromarray(np.array(img)) for img in images]
 
-    return pil_images
+    # Save images as PNG files for download
+    output_files = []
+    for i, img in enumerate(pil_images):
+        output_path = f"output_{i}.png"
+        img.save(output_path, format="PNG")
+        output_files.append(output_path)
+
+    return pil_images, output_files
 
 # Function to extract PNG metadata
 def extract_png_info(image_path):
@@ -462,7 +469,8 @@ with gr.Blocks() as demo:
                 allow_preview=False
             )
             generate_button = gr.Button("Generate Images")
-            gallery = gr.Gallery(label="Generated Images", show_label=False, elem_id="gallery", height=820)
+            gallery = gr.Gallery(label="Generated Images", show_label=False, elem_id="gallery", height=520)
+            download_links = gr.Files(label="Download Generated Images")
 
     selected_style = gr.State(value="Anime Studio Dance")
 
@@ -516,12 +524,13 @@ with gr.Blocks() as demo:
             selected_images = [img[1] for img in selected_folder_images]
         # Adjust the batch_count here to generate the desired number of images
         selected_images = selected_images * batch_count
-        return generate_images_with_progress(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, selected_images, num_inference_steps, progress)
+        pil_images, output_files = generate_images_with_progress(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, selected_images, num_inference_steps, progress)
+        return pil_images, output_files
 
     generate_button.click(
         generate_images_with_folder_images,
         inputs=[prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, use_control_folder, selected_folder_images, batch_images_input, num_inference_steps, control_image],
-        outputs=gallery
+        outputs=[gallery, download_links]
     )
 
     metadata_button = gr.Button("Extract Metadata")