support multimodel

app.py

@@ -9,9 +9,12 @@ from pipeline.util import (
     create_hdr_effect,
     progressive_upscale,
     select_scheduler,
+    torch_gc,
 )
 
 device = "cuda"
+pipe = None
+last_loaded_model = None
 
 # Initialize the models and pipeline
 controlnet = ControlNetUnionModel.from_pretrained(
@@ -19,19 +22,28 @@ controlnet = ControlNetUnionModel.from_pretrained(
 ).to(device=device)
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to(device=device)
 
-#model_id = "SG161222/RealVisXL_V5.0"
-model_id = "SG161222/RealVisXL_V5.0_Lightning"
-pipe = StableDiffusionXLControlNetTileSRPipeline.from_pretrained(
-    model_id, controlnet=controlnet, vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
-).to(device)
+def load_model(model_id):
+    global pipe, last_loaded_model
+    
+    if model_id != last_loaded_model:
+        pipe = None
+        torch_gc()        
+       
+        pipe = StableDiffusionXLControlNetTileSRPipeline.from_pretrained(
+            model_id, controlnet=controlnet, vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
+        ).to(device)
+
+        #pipe.enable_model_cpu_offload()  # << Enable this if you have limited VRAM
+        pipe.enable_vae_tiling() # << Enable this if you have limited VRAM
+        pipe.enable_vae_slicing() # << Enable this if you have limited VRAM
+        last_loaded_model = model_id
 
-#pipe.enable_model_cpu_offload()  # << Enable this if you have limited VRAM
-pipe.enable_vae_tiling() # << Enable this if you have limited VRAM
-pipe.enable_vae_slicing() # << Enable this if you have limited VRAM
+load_model("SG161222/RealVisXL_V5.0_Lightning")
 
 # region functions
 @spaces.GPU(duration=120)
 def predict(
+    model_id,
     image,
     prompt,
     negative_prompt,
@@ -49,6 +61,9 @@ def predict(
 ):
     global pipe
 
+    # Load model if changed
+    load_model(model_id)
+
     # Set selected scheduler
     print(f"Using scheduler: {scheduler}...")
     pipe.scheduler = select_scheduler(pipe, scheduler)
@@ -97,7 +112,6 @@ def predict(
     
     return image
 
-
 def clear_result():
     return gr.update(value=None)
 
@@ -111,6 +125,38 @@ def set_maximum_resolution(max_tile_size, current_value):
 def select_tile_weighting_method(tile_weighting_method):
     return gr.update(visible=True if tile_weighting_method=="Gaussian" else False)
 
+@spaces.GPU(duration=120)
+def run_for_examples(image,
+    prompt,
+    negative_prompt,
+    resolution,
+    hdr,
+    num_inference_steps,
+    denoising_strenght,
+    controlnet_strength,
+    tile_gaussian_sigma,
+    scheduler,
+    guidance_scale,
+    max_tile_size,
+    tile_weighting_method):
+
+    predict(
+        model.value,
+        image,
+        prompt,
+        negative_prompt,
+        resolution,
+        hdr,
+        num_inference_steps,
+        denoising_strenght,
+        controlnet_strength,
+        tile_gaussian_sigma,
+        scheduler,
+        guidance_scale,
+        max_tile_size,
+        tile_weighting_method)
+
+
 # endregion
 
 css = """
@@ -207,6 +253,8 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                     result = gr.Image(
                         label="Generated Image", show_label=True, format="png", interactive=False, scale=1, height=500, min_width=670
                     )
+            with gr.Row():
+                gr.HTML("<div style='color: red;'>If you are not a Pro account, run the LCM sampler examples on the RealVisXL_V5.0_Lightning model. For best results use the UniPC sampler and RealVisXL_V5.0 model examples.</div>")
             with gr.Row():
                 with gr.Accordion("Input Prompt", open=False):
                     with gr.Column():
@@ -231,6 +279,9 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
     with gr.Sidebar(label="Parameters", open=True):
         with gr.Row(elem_id="parameters_row"):
             gr.Markdown("### General parameters")
+            model = gr.Dropdown(
+                label="Model", choices=["SG161222/RealVisXL_V5.0_Lightning", "SG161222/RealVisXL_V5.0"], value="SG161222/RealVisXL_V5.0_Lightning"
+            )
             tile_weighting_method = gr.Dropdown(
                 label="Tile Weighting Method", choices=["Cosine", "Gaussian"], value="Cosine"
             )
@@ -260,7 +311,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             negative_prompt.value,
                             4096,
                             0.0,
-                            18,
+                            25,
                             0.35,
                             1.0,
                             0.3,
@@ -282,7 +333,21 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             4,
                             1024,
                             "Cosine"
-                        ],                   
+                        ],
+                        [   "./examples/2.jpg",
+                            prompt.value,
+                            negative_prompt.value,
+                            4096,
+                            0.5,
+                            25,
+                            0.35,
+                            1.0,
+                            0.3,
+                            "LCM",
+                            4,
+                            1024,
+                            "Cosine"
+                        ],
                         [   "./examples/2.jpg",
                             prompt.value,
                             negative_prompt.value,
@@ -297,6 +362,20 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             1024,
                             "Cosine"
                         ],
+                        [   "./examples/3.jpg",
+                            prompt.value,
+                            negative_prompt.value,
+                            5120,
+                            0.5,
+                            25,
+                            0.35,
+                            1.0,
+                            0.3,
+                            "LCM",
+                            4,
+                            1280,
+                            "Gaussian"
+                        ],
                         [   "./examples/3.jpg",
                             prompt.value,
                             negative_prompt.value,
@@ -311,6 +390,20 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             1280,
                             "Gaussian"
                         ],
+                        [   "./examples/4.jpg",
+                            prompt.value,
+                            negative_prompt.value,
+                            8192,
+                            0.1,
+                            25,
+                            0.35,
+                            1.0,
+                            0.3,
+                            "LCM",
+                            4,
+                            1024,
+                            "Gaussian"
+                        ],
                         [   "./examples/4.jpg",
                             prompt.value,
                             negative_prompt.value,
@@ -325,6 +418,20 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             1024,
                             "Gaussian"
                         ],
+                        [   "./examples/5.jpg",
+                            prompt.value,
+                            negative_prompt.value,
+                            8192,
+                            0.3,
+                            25,
+                            0.35,
+                            1.0,
+                            0.3,
+                            "LCM",
+                            4,
+                            1024,
+                            "Cosine"
+                        ],
                         [   "./examples/5.jpg",
                             prompt.value,
                             negative_prompt.value,
@@ -338,7 +445,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             4,
                             1024,
                             "Cosine"
-                        ],     
+                        ]                        
                     ],
                     inputs=[
                         input_image,
@@ -355,7 +462,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                         max_tile_size,
                         tile_weighting_method,
                     ],                    
-                    fn=predict,
+                    fn=run_for_examples,
                     outputs=result,
                     cache_examples=False,
                 )
@@ -369,6 +476,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
     ).then(
         fn=predict,
         inputs=[
+            model,
             input_image,
             prompt,
             negative_prompt,