Update app.py

app.py

@@ -12,7 +12,6 @@ import os
 import random
 import gc
 
-
 def clear_memory():
     gc.collect()
     if torch.cuda.is_available():
@@ -208,6 +207,8 @@ def load_controlnet_model(controlnet_type):
 
 # Preprocessing functions for each ControlNet type
 def preprocess_canny(image):
+    if isinstance(image, str):
+        image = Image.open(image).convert("RGB")
     if isinstance(image, Image.Image):
         image = np.array(image)
     if image.dtype != np.uint8:
@@ -217,8 +218,9 @@ def preprocess_canny(image):
     image = np.concatenate([image, image, image], axis=2)
     return Image.fromarray(image)
 
-
 def preprocess_depth(image, target_size=(1024, 1024)):
+    if isinstance(image, str):
+        image = Image.open(image).convert("RGB")
     if isinstance(image, Image.Image):
         img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
     else:
@@ -234,6 +236,8 @@ def preprocess_depth(image, target_size=(1024, 1024)):
     return Image.fromarray(depth_img_resized)
 
 def preprocess_openpose(image):
+    if isinstance(image, str):
+        image = Image.open(image).convert("RGB")
     if isinstance(image, Image.Image):
         image = np.array(image)
     image = openpose_processor(image, hand_and_face=False, output_type='cv2')
@@ -273,21 +277,24 @@ def process_image_batch(images, pipe, prompt, negative_prompt, progress, batch_s
         clear_memory()  # Clear memory after each batch
     return all_processed_images
 
-
 # Define the function to generate images
-def generate_images_with_progress(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, control_image, batch_images_input, progress=gr.Progress(track_tqdm=True)):
+def generate_images_with_progress(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, control_images, progress=gr.Progress(track_tqdm=True)):
     global controlnet_pipe, pipe, reference_pipe
 
     clear_memory()
 
+    chunk_size = 2  # Adjust this number based on your memory capacity
+
     if use_controlnet:
         if controlnet_type not in controlnet_models or controlnet_models[controlnet_type] is None:
             raise ValueError(f"{controlnet_type} model not loaded. Please load the model first.")
 
         if mode == "Single Image":
-            control_images = [control_image]
+            control_images = [control_images] if isinstance(control_images, Image.Image) else control_images
         else:
-            control_images = [Image.open(img).convert("RGB") for img in batch_images_input]
+            if not control_images:
+                raise ValueError("No images provided for batch processing.")
+            control_images = [Image.open(img).convert("RGB") if isinstance(img, str) else img for img in control_images]
 
         preprocessed_images = []
         for img in tqdm(control_images, desc="Preprocessing images"):
@@ -300,10 +307,16 @@ def generate_images_with_progress(prompt, negative_prompt, batch_count, use_cont
             else:  # Reference
                 preprocessed_images.append(img)
 
-        if controlnet_type == "Reference":
-            images = process_image_batch(preprocessed_images, reference_pipe, prompt, negative_prompt, progress)
-        else:
-            images = process_image_batch(preprocessed_images, controlnet_pipe, prompt, negative_prompt, progress)
+        images = []
+        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, progress)
+            else:
+                images_chunk = process_image_batch(chunk, controlnet_pipe, prompt, negative_prompt, progress)
+            images.extend(images_chunk)
+            clear_memory()
+
     else:
         if 'controlnet_pipe' in globals() and controlnet_pipe is not None:
             controlnet_pipe.to("cpu")
@@ -331,13 +344,27 @@ def generate_images_with_progress(prompt, negative_prompt, batch_count, use_cont
     clear_memory()
     return images
 
-
-
 # Function to extract PNG metadata
 def extract_png_info(image_path):
     metadata = image_path.info  # This is a dictionary containing key-value pairs of metadata
     return metadata
 
+# Load images from the specified folder
+def load_images_from_folder(folder_path):
+    images = []
+    for filename in os.listdir(folder_path):
+        if filename.endswith(('.png', '.jpg', '.jpeg')):
+            img_path = os.path.join(folder_path, filename)
+            img = Image.open(img_path).convert("RGB")
+            images.append((filename, img))
+    return images
+
+# Folder path where images are stored
+image_folder_path = "control"  # Update this path to your folder
+
+# Load images from folder
+loaded_images = load_images_from_folder(image_folder_path)
+
 # Define the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Image Generation with Custom Prompts and Styles")
@@ -356,6 +383,7 @@ with gr.Blocks() as demo:
             controlnet_type = gr.Dropdown(choices=["Canny", "Depth", "OpenPose", "Reference"], label="ControlNet Type")
             controlnet_status = gr.Textbox(label="ControlNet Status", value="", interactive=False)
             mode = gr.Radio(choices=["Single Image", "Batch"], label="Mode", value="Single Image")
+            use_control_folder = gr.Checkbox(label="Use Control Folder for Batch Processing", value=False)
 
             with gr.Tabs() as tabs:
                 with gr.TabItem("Single Image"):
@@ -368,6 +396,18 @@ with gr.Blocks() as demo:
                     png_image = gr.Image(label="Upload PNG Image", type='pil')
                     metadata_output = gr.JSON(label="PNG Metadata")
 
+                with gr.TabItem("Select from Folder"):
+                    folder_images_gallery = gr.Gallery(
+                        label="Images from Folder",
+                        value=[img[1] for img in loaded_images],
+                        interactive=True,
+                        elem_id="folder-gallery",
+                        columns=5,
+                        object_fit="contain",
+                        height=235,
+                        allow_preview=False
+                    )
+
         with gr.Column(scale=2):
             style_images_gallery = gr.Gallery(
                 label="Choose a Style",
@@ -398,10 +438,37 @@ with gr.Blocks() as demo:
 
     controlnet_type.change(fn=update_controlnet, inputs=controlnet_type, outputs=controlnet_status)
 
+    selected_folder_images = gr.State(value=[])
+
+    def select_folder_image(evt: gr.SelectData):
+        folder_image_names = [img[0] for img in loaded_images]
+        if evt.index < 0 or evt.index >= len(folder_image_names):
+            raise ValueError(f"Invalid index: {evt.index}")
+        selected_image_name = folder_image_names[evt.index]
+        selected_image = next(img for img in loaded_images if img[0] == selected_image_name)
+        current_images = selected_folder_images.value
+        if selected_image not in current_images:
+            current_images.append(selected_image)
+        return current_images
+
+    folder_images_gallery.select(fn=select_folder_image, 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, progress=gr.Progress(track_tqdm=True)):
+        if mode == "Batch":
+            if use_control_folder:
+                selected_images = [img[1] for img in loaded_images]
+            else:
+                if not batch_images_input:
+                    raise ValueError("No images uploaded for batch processing.")
+                selected_images = [Image.open(img).convert("RGB") for img in batch_images_input]
+        else:
+            selected_images = [img[1] for img in selected_folder_images]
+        return generate_images_with_progress(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, selected_images, progress)
+
     generate_button = gr.Button("Generate Images")
     generate_button.click(
-        generate_images_with_progress,
-        inputs=[prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, control_image, batch_images_input],
+        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],
         outputs=gallery
     )
 
@@ -419,4 +486,4 @@ with gr.Blocks() as demo:
 if __name__ == "__main__":
     # Your Gradio interface setup here
     demo.launch(auth=("roland", "roland"), debug=True)
-    clear_memory()
+    clear_memory()