Update app.py

app.py

@@ -18,6 +18,23 @@ def clear_memory():
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
 
+def reset_ui():
+    clear_memory()
+    return (
+        "",  # Reset prompt
+        "",  # Reset negative prompt
+        1,  # Reset batch count
+        30,  # Reset number of inference steps
+        False,  # Reset use controlnet
+        None,  # Reset controlnet type
+        "Restart/Refresh completed",  # Reset controlnet status with message
+        "Single Image",  # Reset mode
+        False,  # Reset use control folder
+        None,  # Reset control image
+        [],  # Reset selected folder images
+        None,  # Reset batch images input
+    )
+
 # Function to resize images while preserving the aspect ratio
 def resize_image(image, max_size=1024):
     width, height = image.size
@@ -377,31 +394,6 @@ image_folder_path = "control"  # Update this path to your folder
 # Load images from folder
 loaded_images = load_images_from_folder(image_folder_path)
 
-# Restart function to clear memory and reinitialize models
-def restart():
-    global controlnet_pipe, reference_pipe, pipe, current_controlnet_type, controlnet_models
-    clear_memory()
-    controlnet_pipe = None
-    reference_pipe = None
-    pipe = None
-    current_controlnet_type = None
-    controlnet_models = {
-        "Canny": None,
-        "Depth": None,
-        "OpenPose": None,
-        "Reference": None
-    }
-    load_base_model()
-    return "Restarted successfully!"
-
-def load_base_model():
-    global pipe
-    model = "aicollective1/aicollective"
-    pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
-    pipe.to("cuda")
-
-load_base_model()
-
 # Define the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Image Generation with Custom Prompts and Styles")
@@ -419,8 +411,8 @@ with gr.Blocks() as demo:
             num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=30)
             use_controlnet = gr.Checkbox(label="Use ControlNet", value=False)
             controlnet_type = gr.Dropdown(choices=["Canny", "Depth", "OpenPose", "Reference"], label="ControlNet Type")
-            controlnet_status = gr.Textbox(label="Status", value="", interactive=False)
-            mode = gr.Radio(choices=["Single Image", "Batch"], label="Mode", value="Single Image")
+            controlnet_status = gr.Textbox(label="ControlNet Status", value="", interactive=False)
+            mode = gr.Radio(choices=["Single Image", "Batch", "Multiselect"], label="Mode", value="Single Image")
             use_control_folder = gr.Checkbox(label="Use Control Folder for Batch Processing", value=False)
 
             with gr.Tabs() as tabs:
@@ -458,7 +450,7 @@ with gr.Blocks() as demo:
                 height=235,
                 allow_preview=False
             )
-            gallery = gr.Gallery(label="Generated Images", show_label=False, elem_id="gallery", height=870)
+            gallery = gr.Gallery(label="Generated Images", show_label=False, elem_id="gallery", height=820)
 
     selected_style = gr.State(value="Anime Studio Dance")
 
@@ -479,31 +471,35 @@ with gr.Blocks() as demo:
 
     selected_folder_images = gr.State(value=[])
 
-    def select_folder_image(evt: gr.SelectData, selected_folder_images):
+    def select_folder_image(evt: gr.SelectData, selected_folder_images, mode):
         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 or []
-        if selected_image not in current_images:
-            current_images.append(selected_image)
+        if mode == "Single Image":
+            current_images = [selected_image]
+        else:
+            if selected_image not in current_images:
+                current_images.append(selected_image)
         return current_images
 
     def clear_selected_folder_images():
         return []
 
-    folder_images_gallery.select(fn=select_folder_image, inputs=[selected_folder_images], outputs=selected_folder_images)
+    folder_images_gallery.select(fn=select_folder_image, inputs=[selected_folder_images, mode], outputs=selected_folder_images)
     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, 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 = [resize_image(Image.open(img).convert("RGB")) for img in batch_images_input]
+    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 == "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:
             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, num_inference_steps, progress)
@@ -511,7 +507,7 @@ with gr.Blocks() as demo:
     generate_button = gr.Button("Generate Images")
     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],
+        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
     )
 
@@ -522,19 +518,32 @@ with gr.Blocks() as demo:
         outputs=metadata_output
     )
 
-    restart_button = gr.Button("Restart")
-    restart_button.click(
-        fn=restart,
+    refresh_button = gr.Button("Restart/Refresh")
+    refresh_button.click(
+        fn=reset_ui,
         inputs=[],
-        outputs=controlnet_status
+        outputs=[
+            prompt,
+            negative_prompt,
+            batch_count,
+            num_inference_steps,
+            use_controlnet,
+            controlnet_type,
+            controlnet_status,
+            mode,
+            use_control_folder,
+            control_image,
+            selected_folder_images,
+            batch_images_input
+        ]
     )
 
     with gr.Row():
         generate_button
-        restart_button
+        refresh_button
 
 # At the end of your script:
 if __name__ == "__main__":
     # Your Gradio interface setup here
     demo.launch(auth=("roland", "roland"), debug=True)
-    clear_memory()
+    clear_memory()