Update app.py

app.py

@@ -1,9 +1,7 @@
 import gradio as gr
 from diffusers import DiffusionPipeline, ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL, UniPCMultistepScheduler
 from stable_diffusion_xl_reference import StableDiffusionXLReferencePipeline
-from controlnet_aux import OpenposeDetector
-#from transformers import DPTFeatureExtractor, DPTForDepthEstimation
-from controlnet_aux import MidasDetector, ZoeDetector
+from controlnet_aux import OpenposeDetector, MidasDetector, ZoeDetector
 from tqdm import tqdm
 
 import torch
@@ -25,6 +23,7 @@ def clear_memory():
 controlnet_pipe = None
 reference_pipe = None
 pipe = None
+current_controlnet_type = None
 
 # Load the base model
 model = "aicollective1/aicollective"
@@ -40,18 +39,15 @@ controlnet_models = {
 }
 
 # Load necessary models and feature extractors for depth estimation and OpenPose
-#feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
-#depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
-
 processor_zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
 processor_midas = MidasDetector.from_pretrained("lllyasviel/Annotators")
-
-
 openpose_processor = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
 
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, use_safetensors=True)
-controlnet_pipe = None  # Initial placeholder, will be loaded dynamically
-reference_pipe = None  # Initial placeholder for reference pipeline
+
+controlnet_model_shared = ControlNetModel.from_pretrained(
+    "xinsir/controlnet-union-sdxl-1.0", torch_dtype=torch.float16, use_safetensors=True
+)
 
 # Define the prompts and negative prompts for each style
 styles = {
@@ -130,20 +126,21 @@ style_images = {
 
 # Function to load ControlNet models dynamically
 def load_controlnet_model(controlnet_type):
-    global controlnet_pipe, pipe, reference_pipe, controlnet_models, vae, model
+    global controlnet_pipe, pipe, reference_pipe, controlnet_models, vae, model, current_controlnet_type, controlnet_model_shared
 
     clear_memory()
 
     if controlnet_models[controlnet_type] is None:
         if controlnet_type in ["Canny", "Depth", "OpenPose"]:
-            controlnet_models[controlnet_type] = ControlNetModel.from_pretrained(
-                "xinsir/controlnet-union-sdxl-1.0", torch_dtype=torch.float16, use_safetensors=True
-            )
+            controlnet_models[controlnet_type] = controlnet_model_shared
         elif controlnet_type == "Reference":
             controlnet_models[controlnet_type] = StableDiffusionXLReferencePipeline.from_pretrained(
                 model, torch_dtype=torch.float16, use_safetensors=True
             )
 
+    if current_controlnet_type == controlnet_type:
+        return f"{controlnet_type} model already loaded."
+
     if 'controlnet_pipe' in globals() and controlnet_pipe is not None:
         controlnet_pipe.to("cpu")
         del controlnet_pipe
@@ -172,10 +169,10 @@ def load_controlnet_model(controlnet_type):
         controlnet_pipe.to("cuda")
         globals()['controlnet_pipe'] = controlnet_pipe
 
+    current_controlnet_type = controlnet_type
     clear_memory()
     return f"Loaded {controlnet_type} model."
 
-
 # Preprocessing functions for each ControlNet type
 def preprocess_canny(image):
     if isinstance(image, Image.Image):
@@ -384,10 +381,8 @@ with gr.Blocks() as demo:
     with gr.Row():
         generate_button
 
-
-
 # At the end of your script:
 if __name__ == "__main__":
     # Your Gradio interface setup here
     demo.launch(debug=True)
-    clear_memory()
+    clear_memory()