model switching

app.py

@@ -33,28 +33,48 @@ def callback_cfg_cutoff(pipeline, step_index, timestep, callback_kwargs):
 
 
 MODELS = {
+    "DreamShaper XL Turbo": "Lykon/dreamshaper-xl-v2-turbo",
     "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
+    "Playground v2.5": "playgroundai/playground-v2.5-1024px-aesthetic",
+    "Juggernaut XL Lightning": "RunDiffusion/Juggernaut-XL-Lightning",
+    "LEOSAM HelloWorld XL": "Leosam/HelloWorld-SDXL-Lightning",
 }
 
+DEFAULT_MODEL = "DreamShaper XL Turbo"
+
 controlnet_model = ControlNetUnionModel.from_pretrained(
     "OzzyGT/controlnet-union-promax-sdxl-1.0", variant="fp16", torch_dtype=torch.float16
 )
 controlnet_model.to(device="cuda", dtype=torch.float16)
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to("cuda")
 
-pipe = DiffusionPipeline.from_pretrained(
-    "SG161222/RealVisXL_V5.0_Lightning",
-    torch_dtype=torch.float16,
-    vae=vae,
-    controlnet=controlnet_model,
-    custom_pipeline="OzzyGT/custom_sdxl_cnet_union",
-).to("cuda")
 
-pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
+def load_pipeline(model_name):
+    """Load a pipeline for the given model name."""
+    model_id = MODELS[model_name]
+    pipeline = DiffusionPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16,
+        vae=vae,
+        controlnet=controlnet_model,
+        custom_pipeline="OzzyGT/custom_sdxl_cnet_union",
+    ).to("cuda")
+    pipeline.scheduler = TCDScheduler.from_config(pipeline.scheduler.config)
+    return pipeline
+
+
+current_model = DEFAULT_MODEL
+pipe = load_pipeline(current_model)
 
 
 @spaces.GPU(duration=24)
 def fill_image(prompt, negative_prompt, image, model_selection, paste_back):
+    global pipe, current_model
+
+    if model_selection != current_model:
+        pipe = load_pipeline(model_selection)
+        current_model = model_selection
+
     (
         prompt_embeds,
         negative_prompt_embeds,
@@ -146,7 +166,7 @@ with gr.Blocks() as demo:
 
     use_as_input_button = gr.Button("Use as Input Image", visible=False)
 
-    model_selection = gr.Dropdown(choices=list(MODELS.keys()), value="RealVisXL V5.0 Lightning", label="Model")
+    model_selection = gr.Dropdown(choices=list(MODELS.keys()), value=DEFAULT_MODEL, label="Model")
 
     def use_output_as_input(output_image):
         return gr.update(value=output_image[1])