Update app.py

app.py

@@ -336,15 +336,14 @@ def generate_images_with_progress(prompt, negative_prompt, batch_count, use_cont
                 preprocessed_images.append(img)
 
         images = []
-        for _ in range(batch_count):
-            for i in range(0, len(preprocessed_images), chunk_size):
-                chunk = preprocessed_images[i:i+chunk_size]
-                if controlnet_type == "Reference":
-                    images_chunk = process_image_batch(chunk, reference_pipe, prompt, negative_prompt, num_inference_steps, progress)
-                else:
-                    images_chunk = process_image_batch(chunk, controlnet_pipe, prompt, negative_prompt, num_inference_steps, progress)
-                images.extend(images_chunk)
-                clear_memory()
+        for i in range(0, len(preprocessed_images), chunk_size):
+            chunk = preprocessed_images[i:i+chunk_size]
+            if controlnet_type == "Reference":
+                images_chunk = process_image_batch(chunk, reference_pipe, prompt, negative_prompt, num_inference_steps, progress)
+            else:
+                images_chunk = process_image_batch(chunk, controlnet_pipe, prompt, negative_prompt, num_inference_steps, progress)
+            images.extend(images_chunk)
+            clear_memory()
 
     else:
         if 'controlnet_pipe' in globals() and controlnet_pipe is not None:
@@ -493,18 +492,18 @@ with gr.Blocks() as demo:
     clear_selection_button.click(fn=clear_selected_folder_images, inputs=[], outputs=selected_folder_images)
 
     def generate_images_with_folder_images(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, use_control_folder, selected_folder_images, batch_images_input, num_inference_steps, control_image, progress=gr.Progress(track_tqdm=True)):
-        if mode == "Single Image" and control_image is not None:
-            selected_images = [control_image]
-        elif mode == "Multiselect":
-            selected_images = [img[1] for img in selected_folder_images]
-        elif mode == "Batch" and use_control_folder:
+        if mode == "Batch" and use_control_folder:
             selected_images = [img[1] for img in loaded_images]
         elif mode == "Batch":
             if not batch_images_input:
                 raise ValueError("No images uploaded for batch processing.")
             selected_images = [resize_image(Image.open(img).convert("RGB")) for img in batch_images_input]
+        elif mode == "Single Image" and control_image is not None:
+            selected_images = [control_image]
         else:
-            raise ValueError("Invalid mode or inputs for the selected mode.")
+            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)
 
     generate_button = gr.Button("Generate Images")
@@ -550,3 +549,4 @@ if __name__ == "__main__":
     # Your Gradio interface setup here
     demo.launch(auth=("roland", "roland"), debug=True)
     clear_memory()
+