Update app.py

app.py

@@ -377,6 +377,31 @@ 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")
@@ -433,7 +458,7 @@ with gr.Blocks() as demo:
                 height=235,
                 allow_preview=False
             )
-            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=850)
 
     selected_style = gr.State(value="Anime Studio Dance")
 
@@ -497,8 +522,16 @@ with gr.Blocks() as demo:
         outputs=metadata_output
     )
 
+    restart_button = gr.Button("Restart")
+    restart_button.click(
+        fn=restart,
+        inputs=[],
+        outputs=controlnet_status
+    )
+
     with gr.Row():
         generate_button
+        restart_button
 
 # At the end of your script:
 if __name__ == "__main__":