added load model pipeline

app.py

@@ -8,34 +8,51 @@ from pipeline.util import (
     SAMPLERS,    
     create_hdr_effect,    
     progressive_upscale,    
-    select_scheduler    
+    select_scheduler,
+    torch_gc,    
 )
 
 device = "cuda"
-pipe = None
-last_loaded_model = None
 MODELS = {"RealVisXL 5 Lightning": "SG161222/RealVisXL_V5.0_Lightning", 
           "RealVisXL 5": "SG161222/RealVisXL_V5.0"
          }
+class Pipeline:
+    def __init__(self):
+        self.pipe = None
+        self.controlnet = None
+        self.vae = None
+        self.last_loaded_model = None
 
-def load_model(model_id):
-    global pipe, last_loaded_model
+    def load_model(self, model_id):
+        if model_id != self.last_loaded_model:
+            print(f"\n--- Loading model: {model_id} ---")
+            if self.pipe is not None:                
+                self.pipe.to("cpu")                
+                del self.pipe
+                self.pipe = None
+                del self.controlnet
+                self.controlnet = None
+                del self.vae
+                self.vae = None                
+                torch_gc()            
+            
+            self.controlnet = ControlNetUnionModel.from_pretrained(
+                    "brad-twinkl/controlnet-union-sdxl-1.0-promax", torch_dtype=torch.float16
+                ).to(device=device)
+            self.vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to(device=device)
 
-    if model_id != last_loaded_model:        
-        # Initialize the models and pipeline
-        controlnet = ControlNetUnionModel.from_pretrained(
-            "brad-twinkl/controlnet-union-sdxl-1.0-promax", torch_dtype=torch.float16
-        ).to(device)
-        vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to(device)
-        pipe = StableDiffusionXLControlNetTileSRPipeline.from_pretrained(
-            MODELS[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
+            self.pipe = StableDiffusionXLControlNetTileSRPipeline.from_pretrained(
+                MODELS[model_id], controlnet=self.controlnet, vae=self.vae, torch_dtype=torch.float16, variant="fp16"
+            ).to(device=device)
 
-load_model("RealVisXL 5 Lightning")
+            self.pipe.enable_model_cpu_offload()
+            self.pipe.enable_vae_tiling()
+            self.pipe.enable_vae_slicing()
+            self.last_loaded_model = model_id
+            print(f"Model {model_id} loaded.")
+
+    def __call__(self, *args, **kwargs):
+        return self.pipe(*args, **kwargs)
 
 # region functions
 @spaces.GPU(duration=120)
@@ -56,14 +73,12 @@ def predict(
     tile_weighting_method,
     progress=gr.Progress(track_tqdm=True),
 ):
-    global pipe
-
     # Load model if changed
     load_model(model_id)
-
+    
     # Set selected scheduler
     print(f"Using scheduler: {scheduler}...")
-    pipe.scheduler = select_scheduler(pipe, scheduler)
+    pipeline.pipe.scheduler = select_scheduler(pipeline.pipe, scheduler)
 
     # Get current image size
     original_height = image.height
@@ -86,7 +101,7 @@ def predict(
 
     # Image generation
     print("Diffusion kicking in... almost done, coffee's on you!")
-    image = pipe(
+    image = pipeline(
         image=image,
         control_image=control_image,
         control_mode=[6],
@@ -112,6 +127,14 @@ def predict(
 def clear_result():
     return gr.update(value=None)
 
+def load_model(model_name, on_load=False):
+    global pipeline  # Declare pipeline as global
+    if on_load and 'pipeline' not in globals(): # Prevent reload page
+        pipeline = Pipeline()  # Create pipeline inside the function
+        pipeline.load_model(model_name) # Load the initial model
+    elif pipeline is not None and not on_load:
+        pipeline.load_model(model_name) # Switch model       
+    
 def set_maximum_resolution(max_tile_size, current_value):
     max_scale = 8  # <- you can try increase it to 12x, 16x if you wish!
     maximum_value = max_tile_size * max_scale
@@ -213,7 +236,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
         with gr.Column(scale=3):                        
             with gr.Row():
                 with gr.Column():
-                    input_image = gr.Image(type="pil", label="Input Image",sources=["upload"], height=500)
+                    input_image = gr.Image(type="pil", label="Input Image", sources=["upload"], height=500)
                 with gr.Column():
                     result = gr.Image(
                         label="Generated Image", show_label=True, format="png", interactive=False, scale=1, height=500, min_width=670
@@ -245,7 +268,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
         with gr.Row(elem_id="parameters_row"):
             gr.Markdown("### General parameters")
             model = gr.Dropdown(
-                label="Model", choices=MODELS.keys(), value=list(MODELS.keys())[0]
+                label="Model", choices=list(MODELS.keys()), value=list(MODELS.keys())[0]
             )
             tile_weighting_method = gr.Dropdown(
                 label="Tile Weighting Method", choices=["Cosine", "Gaussian"], value="Cosine"
@@ -446,6 +469,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
 
     max_tile_size.select(fn=set_maximum_resolution, inputs=[max_tile_size, resolution], outputs=resolution)
     tile_weighting_method.change(fn=select_tile_weighting_method, inputs=tile_weighting_method, outputs=tile_gaussian_sigma)    
+    
     generate_button.click(
         fn=clear_result,
         inputs=None,
@@ -468,8 +492,8 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
             max_tile_size,
             tile_weighting_method,
         ],
-        outputs=result,
-        show_progress="full"
+        outputs=result,     
     )
     gr.Markdown(about)
+    app.load(fn=load_model, inputs=[model, gr.State(value=True)], outputs=None, concurrency_limit=1) # Load initial model on app load
 app.launch(share=False)

pipeline/util.py

@@ -16,8 +16,6 @@
 import gc
 import cv2
 import numpy as np
-from torch import nn
-from accelerate.hooks import AlignDevicesHook, CpuOffload, remove_hook_from_module
 import torch
 from PIL import Image
 
@@ -98,32 +96,6 @@ def select_scheduler(pipe, selected_sampler):
 
     return scheduler.from_config(config, **add_kwargs)
 
-def optionally_disable_offloading(_pipeline):
-    """
-    Optionally removes offloading in case the pipeline has been already sequentially offloaded to CPU.
-
-    Args:
-        _pipeline (`DiffusionPipeline`):
-            The pipeline to disable offloading for.
-
-    Returns:
-        tuple:
-            A tuple indicating if `is_model_cpu_offload` or `is_sequential_cpu_offload` is True.
-    """
-    is_model_cpu_offload = False
-    is_sequential_cpu_offload = False   
-    if _pipeline is not None:
-        for _, component in _pipeline.components.items():
-            if isinstance(component, nn.Module) and hasattr(component, "_hf_hook"):
-                if not is_model_cpu_offload:
-                    is_model_cpu_offload = isinstance(component._hf_hook, CpuOffload)
-                if not is_sequential_cpu_offload:
-                    is_sequential_cpu_offload = isinstance(component._hf_hook, AlignDevicesHook)
-
-               
-                remove_hook_from_module(component, recurse=True)
-
-    return (is_model_cpu_offload, is_sequential_cpu_offload)
 
 # This function was copied and adapted from https://huggingface.co/spaces/gokaygokay/TileUpscalerV2, licensed under Apache 2.0.
 def progressive_upscale(input_image, target_resolution, steps=3):
@@ -210,14 +182,12 @@ def create_hdr_effect(original_image, hdr):
     
 
 def torch_gc():
+    gc.collect()
     if torch.cuda.is_available():
         with torch.cuda.device("cuda"):
             torch.cuda.empty_cache()
             torch.cuda.ipc_collect()
 
-    gc.collect()
-
-
 def quantize_8bit(unet):
     if unet is None:
         return