switch model

app.py

@@ -1,49 +1,59 @@
 import torch
 import spaces
-from diffusers import ControlNetUnionModel, AutoencoderKL
+from diffusers import ControlNetUnionModel, AutoencoderKL, UNet2DConditionModel
 import gradio as gr
 
 from pipeline.mod_controlnet_tile_sr_sdxl import StableDiffusionXLControlNetTileSRPipeline, calculate_overlap
 from pipeline.util import (
     SAMPLERS,    
     create_hdr_effect,
+    optionally_disable_offloading,
     progressive_upscale,
+    quantize_8bit,
     select_scheduler,
-    torch_gc,
+    torch_gc,    
 )
 
 device = "cuda"
 pipe = None
 last_loaded_model = None
-
-# Initialize the models and pipeline
-controlnet = ControlNetUnionModel.from_pretrained(
-    "brad-twinkl/controlnet-union-sdxl-1.0-promax", torch_dtype=torch.float16
-).to(device=device)
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to(device=device)
+MODELS = {"RealVisXL 5 Lightning": "SG161222/RealVisXL_V5.0_Lightning", 
+          "RealVisXL 5": "SG161222/RealVisXL_V5.0"
+         }
 
 def load_model(model_id):
     global pipe, last_loaded_model
-    
+
     if model_id != last_loaded_model:
-        pipe = None
-              
+        
+        # Initialize the models and pipeline
+        controlnet = ControlNetUnionModel.from_pretrained(
+            "brad-twinkl/controlnet-union-sdxl-1.0-promax", torch_dtype=torch.float16
+        )
+        vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+        if pipe is not None:
+            optionally_disable_offloading(pipe)
+            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
+            MODELS[model_id], controlnet=controlnet, vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
+        )
+        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
+
+        unet = UNet2DConditionModel.from_pretrained(MODELS[model_id], subfolder="unet", variant="fp16", use_safetensors=True)
+        quantize_8bit(unet)  # << Enable this if you have limited VRAM
+        pipe.unet = unet
+
         last_loaded_model = model_id
 
-load_model("SG161222/RealVisXL_V5.0_Lightning")
+load_model("RealVisXL 5 Lightning")
 
 # region functions
 @spaces.GPU(duration=120)
 def predict(
-    model_id,
     image,
+    model_id,
     prompt,
     negative_prompt,
     resolution,
@@ -124,38 +134,6 @@ 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 = """
@@ -174,7 +152,7 @@ body {
     text-shadow: 1px 1px 2px rgba(0, 0, 0, 0.2);
 }
 .fillable {
-    width: 95% !important;
+    width: 100% !important;
     max-width: unset !important;
 }
 #examples_container {
@@ -279,7 +257,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=["SG161222/RealVisXL_V5.0_Lightning", "SG161222/RealVisXL_V5.0"], value="SG161222/RealVisXL_V5.0_Lightning"
+                label="Model", choices=MODELS.keys(), value=list(MODELS.keys())[0]
             )
             tile_weighting_method = gr.Dropdown(
                 label="Tile Weighting Method", choices=["Cosine", "Gaussian"], value="Cosine"
@@ -303,9 +281,10 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
     with gr.Accordion(label="Example Images", open=True):
         with gr.Row(elem_id="examples_row"):
             with gr.Column(scale=12, elem_id="examples_container"):
-                gr.Examples(
+                eg = gr.Examples(
                     examples=[
                         [   "./examples/1.jpg",
+                            "RealVisXL 5 Lightning",
                             prompt.value,
                             negative_prompt.value,
                             4096,
@@ -320,6 +299,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             "Cosine"
                         ],
                         [   "./examples/1.jpg",
+                            "RealVisXL 5",
                             prompt.value,
                             negative_prompt.value,
                             4096,
@@ -334,6 +314,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             "Cosine"
                         ],
                         [   "./examples/2.jpg",
+                            "RealVisXL 5 Lightning",
                             prompt.value,
                             negative_prompt.value,
                             4096,
@@ -348,6 +329,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             "Cosine"
                         ],
                         [   "./examples/2.jpg",
+                            "RealVisXL 5",                       
                             prompt.value,
                             negative_prompt.value,
                             4096,
@@ -362,6 +344,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             "Cosine"
                         ],
                         [   "./examples/3.jpg",
+                            "RealVisXL 5 Lightning",
                             prompt.value,
                             negative_prompt.value,
                             5120,
@@ -376,6 +359,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             "Gaussian"
                         ],
                         [   "./examples/3.jpg",
+                            "RealVisXL 5",
                             prompt.value,
                             negative_prompt.value,
                             5120,
@@ -390,6 +374,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             "Gaussian"
                         ],
                         [   "./examples/4.jpg",
+                            "RealVisXL 5 Lightning",
                             prompt.value,
                             negative_prompt.value,
                             8192,
@@ -404,6 +389,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             "Gaussian"
                         ],
                         [   "./examples/4.jpg",
+                            "RealVisXL 5",
                             prompt.value,
                             negative_prompt.value,
                             8192,
@@ -418,6 +404,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             "Gaussian"
                         ],
                         [   "./examples/5.jpg",
+                            "RealVisXL 5 Lightning",
                             prompt.value,
                             negative_prompt.value,
                             8192,
@@ -432,6 +419,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                             "Cosine"
                         ],
                         [   "./examples/5.jpg",
+                            "RealVisXL 5",
                             prompt.value,
                             negative_prompt.value,
                             8192,
@@ -448,6 +436,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                     ],
                     inputs=[
                         input_image,
+                        model,
                         prompt,
                         negative_prompt,
                         resolution,
@@ -461,13 +450,13 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
                         max_tile_size,
                         tile_weighting_method,
                     ],                    
-                    fn=run_for_examples,
+                    fn=predict,
                     outputs=result,
                     cache_examples=False,
                 )
 
     max_tile_size.select(fn=set_maximum_resolution, inputs=[max_tile_size, resolution], outputs=resolution)
-    tile_weighting_method.select(fn=select_tile_weighting_method, inputs=tile_weighting_method, outputs=tile_gaussian_sigma)
+    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,
@@ -475,8 +464,8 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(), title="MoD ControlNet Tile Upsc
     ).then(
         fn=predict,
         inputs=[
-            model,
             input_image,
+            model,
             prompt,
             negative_prompt,
             resolution,

pipeline/util.py

@@ -16,6 +16,8 @@
 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
 
@@ -96,6 +98,32 @@ 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):
@@ -185,7 +213,7 @@ def torch_gc():
     if torch.cuda.is_available():
         with torch.cuda.device("cuda"):
             torch.cuda.empty_cache()
-            torch.cuda.ipc_collect()
+            #torch.cuda.ipc_collect()
 
     gc.collect()