Update app.py

app.py

@@ -11,6 +11,7 @@ from PIL import Image
 import os
 import random
 import gc
+import tempfile
 
 def clear_memory():
     gc.collect()
@@ -311,6 +312,16 @@ def process_image_batch(images, pipe, prompt, negative_prompt, num_inference_ste
         clear_memory()  # Clear memory after each batch
     return all_processed_images
 
+# Function to save images as PNG and return their paths
+def save_images_as_png(images):
+    temp_dir = tempfile.mkdtemp()
+    png_paths = []
+    for i, img in enumerate(images):
+        png_path = os.path.join(temp_dir, f"image_{i}.png")
+        img.save(png_path, "PNG")
+        png_paths.append(png_path)
+    return png_paths
+
 # Define the function to generate images
 def generate_images_with_progress(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, control_images, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
     global controlnet_pipe, pipe, reference_pipe
@@ -378,17 +389,10 @@ def generate_images_with_progress(prompt, negative_prompt, batch_count, use_cont
 
     clear_memory()
 
-    # 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]
-
-    # 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)
+    # Save images as PNG and return their paths
+    png_paths = save_images_as_png(images)
 
-    return pil_images, output_files
+    return png_paths
 
 # Function to extract PNG metadata
 def extract_png_info(image_path):
@@ -469,8 +473,7 @@ 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=580)
-            download_links = gr.Files(label="Download Generated Images")
+            gallery = gr.Gallery(label="Generated Images", show_label=False, elem_id="gallery", height=820)
 
     selected_style = gr.State(value="Anime Studio Dance")
 
@@ -524,13 +527,12 @@ 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
-        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
+        return generate_images_with_progress(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, selected_images, num_inference_steps, progress)
 
     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, download_links]
+        outputs=gallery
     )
 
     metadata_button = gr.Button("Extract Metadata")