offloading hooks

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -10,16 +10,17 @@
     if you have a compatible nVidia GPU, Nunchaku is the fastest quantization engine, currently available for Flux.1, SANA and Qwen-Image models  
     *note*: release version of `nunchaku==0.3.2` does NOT include support, so you need to build [nunchaku](https://nunchaku.tech/docs/nunchaku/installation/installation.html) from source  
   - updated [SD.Next Model Samples Gallery](https://vladmandic.github.io/sd-samples/compare.html)  
-- **Core**
-  - enable offload during pre-forward by default  
-  - improve offloading of very large models  
-  - update `requirements`  
 - **UI**  
   - improved image scaling in img2img and control interfaces  
   - add base model type to networks display, thanks @Artheriax  
   - additional hints to ui, thanks @Artheriax  
   - add video support to gallery, thanks @CalamitousFelicitousness  
   - additional artwork for reference models in networks, thanks @liutyi  
+- **Offloading**
+  - enable offload during pre-forward by default  
+  - improve offloading of models with multiple dits  
+  - improve offloading of models with impliciy vae processing  
+  - improve offloading of models with controlnet  
 - **Fixes**
   - normalize path hanlding when deleting images  
   - fix hidden model tags in networks display  

modules/control/units/controlnet.py

@@ -313,6 +313,7 @@ class ControlNet():
                             errors.display(e, 'Control')
                 if self.model is None:
                     return
+                self.model.offload_never = True
                 if self.dtype is not None:
                     self.model.to(self.dtype)
                 if "Control" in opts.sdnq_quantize_weights:
@@ -441,7 +442,7 @@ class ControlNetPipeline():
                 tokenizer=pipeline.tokenizer,
                 transformer=pipeline.transformer,
                 scheduler=pipeline.scheduler,
-                controlnet=controlnets, # can be a list
+                controlnet=controlnets[0] if isinstance(controlnets, list) else controlnets, # can be a list
             )
         elif len(loras) > 0:
             self.pipeline = pipeline
@@ -463,11 +464,13 @@ class ControlNetPipeline():
         if dtype is not None:
             self.pipeline = self.pipeline.to(dtype)
 
+        controlnet = None # free up memory
+        controlnets = None
         sd_models.copy_diffuser_options(self.pipeline, pipeline)
         if opts.diffusers_offload_mode == 'none':
             sd_models.move_model(self.pipeline, devices.device)
-        from modules.sd_models import set_diffuser_offload
-        set_diffuser_offload(self.pipeline, 'model')
+        sd_models.clear_caches()
+        sd_models.set_diffuser_offload(self.pipeline, 'model')
 
         t1 = time.time()
         debug_log(f'Control {what} pipeline: class={self.pipeline.__class__.__name__} time={t1-t0:.2f}')

modules/sd_hijack_te.py

@@ -19,15 +19,14 @@ def hijack_encode_prompt(*args, **kwargs):
         res = None
     t1 = time.time()
     timer.process.add('te', t1-t0)
-    if hasattr(shared.sd_model, "maybe_free_model_hooks"):
-        shared.sd_model.maybe_free_model_hooks()
+    # if hasattr(shared.sd_model, "maybe_free_model_hooks"):
+    #     shared.sd_model.maybe_free_model_hooks()
     shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
     shared.state.end()
     return res
 
 
 def init_hijack(pipe):
-    if shared.opts.te_hijack and pipe is not None and not hasattr(pipe, 'orig_encode_prompt') and hasattr(pipe, 'encode_prompt'):
-        # shared.log.debug(f'Model: cls={pipe.__class__.__name__} hijack encode')
+    if pipe is not None and not hasattr(pipe, 'orig_encode_prompt') and hasattr(pipe, 'encode_prompt'):
         pipe.orig_encode_prompt = pipe.encode_prompt
         pipe.encode_prompt = hijack_encode_prompt

modules/sd_hijack_vae.py

@@ -0,0 +1,64 @@
+import os
+import time
+import torch
+from modules import shared, sd_models, devices, timer, errors
+
+
+debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None
+
+
+def hijack_vae_decode(*args, **kwargs):
+    shared.state.begin('VAE')
+    t0 = time.time()
+    res = None
+    shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae'])
+    try:
+        sd_models.move_model(shared.sd_model.vae, devices.device)
+        if torch.is_tensor(args[0]):
+            latents = args[0].to(device=devices.device, dtype=shared.sd_model.vae.dtype) # upcast to vae dtype
+            res = shared.sd_model.vae.orig_decode(latents, *args[1:], **kwargs)
+            t1 = time.time()
+            shared.log.debug(f'Decode: vae={shared.sd_model.vae.__class__.__name__} slicing={getattr(shared.sd_model.vae, "use_slicing", None)} tiling={getattr(shared.sd_model.vae, "use_tiling", None)} latents={list(latents.shape)}:{latents.device}:{latents.dtype} time={t1-t0:.3f}')
+        else:
+            res = shared.sd_model.vae.orig_decode(*args, **kwargs)
+    except Exception as e:
+        shared.log.error(f'Decode: vae={shared.sd_model.vae.__class__.__name__} {e}')
+        errors.display(e, 'vae')
+        res = None
+    t1 = time.time()
+    timer.process.add('vae', t1-t0)
+    shared.state.end()
+    return res
+
+
+def hijack_vae_encode(*args, **kwargs):
+    shared.state.begin('VAE')
+    t0 = time.time()
+    res = None
+    shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae'])
+    try:
+        sd_models.move_model(shared.sd_model.vae, devices.device)
+        if torch.is_tensor(args[0]):
+            latents = args[0].to(device=devices.device, dtype=shared.sd_model.vae.dtype) # upcast to vae dtype
+            res = shared.sd_model.vae.orig_encode(latents, *args[1:], **kwargs)
+            t1 = time.time()
+            shared.log.debug(f'Encode: vae={shared.sd_model.vae.__class__.__name__} slicing={getattr(shared.sd_model.vae, "use_slicing", None)} tiling={getattr(shared.sd_model.vae, "use_tiling", None)} latents={list(latents.shape)}:{latents.device}:{latents.dtype} time={t1-t0:.3f}')
+        else:
+            res = shared.sd_model.vae.orig_encode(*args, **kwargs)
+    except Exception as e:
+        shared.log.error(f'Encode: vae={shared.sd_model.vae.__class__.__name__} {e}')
+        errors.display(e, 'vae')
+        res = None
+    t1 = time.time()
+    timer.process.add('vae', t1-t0)
+    shared.state.end()
+    return res
+
+
+def init_hijack(pipe):
+    if pipe is not None and hasattr(pipe, 'vae') and hasattr(pipe.vae, 'decode') and not hasattr(pipe.vae, 'orig_decode'):
+        pipe.vae.orig_decode = pipe.vae.decode
+        pipe.vae.decode = hijack_vae_decode
+    if pipe is not None and hasattr(pipe, 'vae') and hasattr(pipe.vae, 'encode') and not hasattr(pipe.vae, 'orig_encode'):
+        pipe.vae.orig_encode = pipe.vae.encode
+        pipe.vae.encode = hijack_vae_encode

modules/sd_models.py

@@ -57,22 +57,6 @@ i2i_pipes = [
 ]
 
 
-def copy_diffuser_options(new_pipe, orig_pipe):
-    new_pipe.sd_checkpoint_info = getattr(orig_pipe, 'sd_checkpoint_info', None)
-    new_pipe.sd_model_checkpoint = getattr(orig_pipe, 'sd_model_checkpoint', None)
-    new_pipe.embedding_db = getattr(orig_pipe, 'embedding_db', None)
-    new_pipe.sd_model_hash = getattr(orig_pipe, 'sd_model_hash', None)
-    new_pipe.has_accelerate = getattr(orig_pipe, 'has_accelerate', False)
-    new_pipe.current_attn_name = getattr(orig_pipe, 'current_attn_name', None)
-    new_pipe.default_scheduler = getattr(orig_pipe, 'default_scheduler', None)
-    new_pipe.is_sdxl = getattr(orig_pipe, 'is_sdxl', False) # a1111 compatibility item
-    new_pipe.is_sd2 = getattr(orig_pipe, 'is_sd2', False)
-    new_pipe.is_sd1 = getattr(orig_pipe, 'is_sd1', True)
-    add_noise_pred_to_diffusers_callback(new_pipe)
-    if new_pipe.has_accelerate:
-        set_accelerate(new_pipe)
-
-
 def set_huggingface_options():
     if shared.opts.diffusers_to_gpu: # and model_type.startswith('Stable Diffusion'):
         sd_hijack_accelerate.hijack_accelerate()
@@ -872,6 +856,28 @@ def clean_diffuser_pipe(pipe):
         pipe.register_to_config(**internal_dict)
 
 
+def copy_diffuser_options(new_pipe, orig_pipe):
+    new_pipe.sd_checkpoint_info = getattr(orig_pipe, 'sd_checkpoint_info', None)
+    new_pipe.sd_model_checkpoint = getattr(orig_pipe, 'sd_model_checkpoint', None)
+    new_pipe.embedding_db = getattr(orig_pipe, 'embedding_db', None)
+    new_pipe.loaded_loras = getattr(orig_pipe, 'loaded_loras', {})
+    new_pipe.sd_model_hash = getattr(orig_pipe, 'sd_model_hash', None)
+    new_pipe.has_accelerate = getattr(orig_pipe, 'has_accelerate', False)
+    new_pipe.current_attn_name = getattr(orig_pipe, 'current_attn_name', None)
+    new_pipe.default_scheduler = getattr(orig_pipe, 'default_scheduler', None)
+    new_pipe.image_encoder = getattr(orig_pipe, 'image_encoder', None)
+    new_pipe.feature_extractor = getattr(orig_pipe, 'feature_extractor', None)
+    new_pipe.mask_processor = getattr(orig_pipe, 'mask_processor', None)
+    new_pipe.restore_pipeline = getattr(orig_pipe, 'restore_pipeline', None)
+    new_pipe.task_args = getattr(orig_pipe, 'task_args', None)
+    new_pipe.is_sdxl = getattr(orig_pipe, 'is_sdxl', False) # a1111 compatibility item
+    new_pipe.is_sd2 = getattr(orig_pipe, 'is_sd2', False)
+    new_pipe.is_sd1 = getattr(orig_pipe, 'is_sd1', True)
+    add_noise_pred_to_diffusers_callback(new_pipe)
+    if new_pipe.has_accelerate:
+        set_accelerate(new_pipe)
+
+
 def backup_pipe_components(pipe):
     if pipe is None:
         return {}

modules/sd_offload.py

@@ -203,18 +203,19 @@ class OffloadHook(accelerate.hooks.ModelHook):
         return module
 
     def pre_forward(self, module, *args, **kwargs):
-        if self.last_pre != id(module): # offload every other module first time when new module starts pre-forward
-            self.last_pre = id(module)
+        _id = id(module)
+        if self.last_pre != _id and not hasattr(module, "offload_never"): # offload every other module first time when new module starts pre-forward
+            self.last_pre = _id
             if shared.opts.diffusers_offload_pre:
                 debug_move(f'Offload: type=balanced op=pre module={module.__class__.__name__}')
                 for pipe in get_pipe_variants():
                     for module_name in get_module_names(pipe):
                         module_instance = getattr(pipe, module_name, None)
                         module_cls = module_instance.__class__.__name__
-                        if (id(module) != id(module_instance)) and (module_cls not in self.offload_never) and (not devices.same_device(module_instance.device, devices.cpu)):
+                        if (_id != id(module_instance)) and (module_cls not in self.offload_never) and (not devices.same_device(module_instance.device, devices.cpu)):
                             apply_balanced_offload_to_module(module_instance, op='pre')
 
-        if not devices.same_device(module.device, devices.device):
+        if not devices.same_device(module.device, devices.device): # move-to-device
             device_index = torch.device(devices.device).index
             if device_index is None:
                 device_index = 0
@@ -233,6 +234,13 @@ class OffloadHook(accelerate.hooks.ModelHook):
             module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access
             module.balanced_offload_device_map = device_map
             module.balanced_offload_max_memory = max_memory
+
+        if debug:
+            for pipe in get_pipe_variants():
+                for module_name in get_module_names(pipe):
+                    module_instance = getattr(pipe, module_name, None)
+                    shared.log.trace(f'Offload: type=balanced op=pre check module={module_instance.__class__.__name__} device={module_instance.device} dtype={module_instance.dtype}')
+
         return args, kwargs
 
     def post_forward(self, module, output):
@@ -292,7 +300,7 @@ def get_module_sizes(pipe=None, exclude=[]):
     return modules
 
 
-def move_module_to_cpu(module, op='unk'):
+def move_module_to_cpu(module, op='unk', force:bool=False):
     try:
         module_name = getattr(module, "module_name", module.__class__.__name__)
         module_size = offload_hook_instance.offload_map.get(module_name, offload_hook_instance.model_size())
@@ -301,7 +309,11 @@ def move_module_to_cpu(module, op='unk'):
         prev_gpu = used_gpu
         module_cls = module.__class__.__name__
         op = f'{op}:skip'
-        if module_cls in offload_hook_instance.offload_never:
+        if force:
+            op = f'{op}:force'
+            module = module.to(devices.cpu)
+            used_gpu -= module_size
+        elif module_cls in offload_hook_instance.offload_never:
             op = f'{op}:never'
         elif module_cls in offload_hook_instance.offload_always:
             op = f'{op}:always'
@@ -313,7 +325,7 @@ def move_module_to_cpu(module, op='unk'):
             used_gpu -= module_size
         if debug:
             quant = getattr(module, "quantization_method", None)
-            debug_move(f'Offload: type=balanced op={op} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={module_cls} size={module_size:.3f}')
+            debug_move(f'Offload: type=balanced op={op} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f}:{shared.opts.diffusers_offload_min_gpu_memory} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={module_cls} size={module_size:.3f}')
     except Exception as e:
         if 'out of memory' in str(e):
             devices.torch_gc(fast=True, force=True, reason='oom')
@@ -325,7 +337,7 @@ def move_module_to_cpu(module, op='unk'):
             errors.display(e, f'Offload: type=balanced op=apply module={getattr(module, "__name__", None)}')
 
 
-def apply_balanced_offload_to_module(module, op="apply"):
+def apply_balanced_offload_to_module(module, op="apply", force:bool=False):
     module_name = getattr(module, "module_name", module.__class__.__name__)
     network_layer_name = getattr(module, "network_layer_name", None)
     device_map = getattr(module, "balanced_offload_device_map", None)
@@ -334,7 +346,7 @@ def apply_balanced_offload_to_module(module, op="apply"):
         module = accelerate.hooks.remove_hook_from_module(module, recurse=True)
     except Exception as e:
         shared.log.warning(f'Offload remove hook: module={module_name} {e}')
-    move_module_to_cpu(module, op=op)
+    move_module_to_cpu(module, op=op, force=force)
     try:
         module = accelerate.hooks.add_hook_to_module(module, offload_hook_instance, append=True)
     except Exception as e:
@@ -345,7 +357,7 @@ def apply_balanced_offload_to_module(module, op="apply"):
     if device_map and max_memory:
         module.balanced_offload_device_map = device_map
         module.balanced_offload_max_memory = max_memory
-    module.offload_post = shared.sd_model_type in offload_post and shared.opts.te_hijack and module_name.startswith("text_encoder")
+    module.offload_post = shared.sd_model_type in offload_post and module_name.startswith("text_encoder")
     if shared.opts.layerwise_quantization or getattr(module, 'quantization_method', None) == 'LayerWise':
         model_quant.apply_layerwise(module, quiet=True) # need to reapply since hooks were removed/readded
     devices.torch_gc(fast=True, force=True, reason='offload')

modules/shared.py

@@ -251,7 +251,6 @@ options_templates.update(options_section(('text_encoder', "Text Encoder"), {
     "sd_textencoder_cache_size": OptionInfo(4, "Text encoder cache size", gr.Slider, {"minimum": 0, "maximum": 16, "step": 1}),
     "sd_textencder_linebreak": OptionInfo(True, "Use line break as prompt segment marker", gr.Checkbox),
     "diffusers_zeros_prompt_pad": OptionInfo(False, "Use zeros for prompt padding", gr.Checkbox),
-    "te_hijack": OptionInfo(True, "Offload after prompt encode", gr.Checkbox),
     "te_optional_sep": OptionInfo("<h2>Optional</h2>", "", gr.HTML),
     "te_shared_t5": OptionInfo(True, "T5: Use shared instance of text encoder"),
     "te_pooled_embeds": OptionInfo(False, "SDXL: Use weighted pooled embeds"),

pipelines/model_bria.py

@@ -2,7 +2,7 @@ import os
 import sys
 import transformers
 import diffusers
-from modules import shared, devices, sd_models, model_quant, sd_hijack_te
+from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
 from pipelines import generic
 
 
@@ -33,11 +33,8 @@ def load_bria(checkpoint_info, diffusers_load_config={}):
 
     del text_encoder
     del transformer
-
     sd_hijack_te.init_hijack(pipe)
-    from modules.video_models import video_vae
-    pipe.vae.orig_decode = pipe.vae.decode
-    pipe.vae.decode = video_vae.hijack_vae_decode
+    sd_hijack_vae.init_hijack(pipe)
 
     devices.torch_gc()
     return pipe

pipelines/model_chroma.py

@@ -27,5 +27,6 @@ def load_chroma(checkpoint_info, diffusers_load_config={}):
     del text_encoder
     del transformer
     sd_hijack_te.init_hijack(pipe)
+
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_cogview.py

@@ -48,5 +48,6 @@ def load_cogview4(checkpoint_info, diffusers_load_config={}):
     sd_hijack_te.init_hijack(pipe)
     del transformer
     del text_encoder
+
     devices.torch_gc()
     return pipe

pipelines/model_cosmos.py

@@ -1,6 +1,6 @@
 import transformers
 import diffusers
-from modules import shared, devices, sd_models, model_quant, sd_hijack_te
+from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
 from pipelines import generic
 
 
@@ -29,9 +29,7 @@ def load_cosmos_t2i(checkpoint_info, diffusers_load_config={}):
     del transformer
 
     sd_hijack_te.init_hijack(pipe)
-    from modules.video_models import video_vae
-    pipe.vae.orig_decode = pipe.vae.decode
-    pipe.vae.decode = video_vae.hijack_vae_decode
+    sd_hijack_vae.init_hijack(pipe)
 
     devices.torch_gc()
     return pipe

pipelines/model_flex.py

@@ -29,5 +29,6 @@ def load_flex(checkpoint_info, diffusers_load_config={}):
     del text_encoder_2
     del transformer
     sd_hijack_te.init_hijack(pipe)
+
     devices.torch_gc()
     return pipe

pipelines/model_flite.py

@@ -32,5 +32,6 @@ def load_flite(checkpoint_info, diffusers_load_config={}):
     del text_encoder
     del dit_model
     sd_hijack_te.init_hijack(pipe)
+
     devices.torch_gc()
     return pipe

pipelines/model_flux.py

@@ -84,5 +84,6 @@ def load_flux(checkpoint_info, diffusers_load_config={}):
         apply_cache_on_pipe(pipe, residual_diff_threshold=0.12)
 
     sd_hijack_te.init_hijack(pipe)
+
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_hidream.py

@@ -71,5 +71,6 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
     del tokenizer_4
     del transformer
     sd_hijack_te.init_hijack(pipe)
+
     devices.torch_gc()
     return pipe

pipelines/model_hunyuandit.py

@@ -26,5 +26,6 @@ def load_hunyuandit(checkpoint_info, diffusers_load_config={}):
     del text_encoder_2
     del transformer
     sd_hijack_te.init_hijack(pipe)
+
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_kandinsky.py

@@ -1,6 +1,6 @@
 import transformers
 import diffusers
-from modules import shared, sd_models, devices, model_quant, sd_hijack_te
+from modules import shared, sd_models, devices, model_quant, sd_hijack_te, sd_hijack_vae
 from pipelines import generic
 
 
@@ -61,5 +61,7 @@ def load_kandinsky3(checkpoint_info, diffusers_load_config={}):
     del text_encoder
     del unet
     sd_hijack_te.init_hijack(pipe)
+    sd_hijack_vae.init_hijack(pipe)
+
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_kolors.py

@@ -16,6 +16,8 @@ def load_kolors(_checkpoint_info, diffusers_load_config={}):
         **diffusers_load_config,
     )
     pipe.vae.config.force_upcast = True
+
     sd_hijack_te.init_hijack(pipe)
+
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_lumina.py

@@ -40,5 +40,6 @@ def load_lumina2(checkpoint_info, diffusers_load_config={}):
     del transformer
     del text_encoder
     sd_hijack_te.init_hijack(pipe)
+
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_meissonic.py

@@ -53,5 +53,6 @@ def load_meissonic(checkpoint_info, diffusers_load_config={}):
     diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["meissonic"] = MeissonicImg2ImgPipeline
     diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["meissonic"] = MeissonicInpaintPipeline
     sd_hijack_te.init_hijack(pipe)
+
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_omnigen.py

@@ -26,5 +26,6 @@ def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=u
     )
 
     sd_hijack_te.init_hijack(pipe)
+
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_omnigen2.py

@@ -42,5 +42,6 @@ def load_omnigen2(checkpoint_info, diffusers_load_config={}): # pylint: disable=
     pipe.transformer = transformer # for omnigen2 transformer must be loaded after pipeline
 
     sd_hijack_te.init_hijack(pipe)
+
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_pixart.py

@@ -34,5 +34,6 @@ def load_pixart(checkpoint_info, diffusers_load_config={}):
     del text_encoder
     del transformer
     sd_hijack_te.init_hijack(pipe)
+
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_qwen.py

@@ -1,6 +1,6 @@
 import transformers
 import diffusers
-from modules import shared, devices, sd_models, model_quant, sd_hijack_te
+from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
 from pipelines import generic
 
 
@@ -49,11 +49,8 @@ def load_qwen(checkpoint_info, diffusers_load_config={}):
 
     del text_encoder
     del transformer
-
     sd_hijack_te.init_hijack(pipe)
-    from modules.video_models import video_vae
-    pipe.vae.orig_decode = pipe.vae.decode
-    pipe.vae.decode = video_vae.hijack_vae_decode
+    sd_hijack_vae.init_hijack(pipe)
 
     devices.torch_gc()
     return pipe

pipelines/model_sana.py

@@ -79,7 +79,6 @@ def load_sana(checkpoint_info, kwargs={}):
         shared.log.error(f'Load model: type=Sana {e}')
 
     sd_hijack_te.init_hijack(pipe)
-    t1 = time.time()
-    shared.log.debug(f'Load model: type=Sana target={devices.dtype} te={pipe.text_encoder.dtype} transformer={pipe.transformer.dtype} vae={pipe.vae.dtype} time={t1-t0:.2f}')
+
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_sd3.py

@@ -32,5 +32,6 @@ def load_sd3(checkpoint_info, diffusers_load_config={}):
     del text_encoder_3
     del transformer
     sd_hijack_te.init_hijack(pipe)
+
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_wanai.py

@@ -1,7 +1,7 @@
 import os
 import transformers
 import diffusers
-from modules import shared, devices, sd_models, model_quant, sd_hijack_te
+from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
 
 
 def load_transformer(repo_id, diffusers_load_config={}, subfolder='transformer'):
@@ -102,9 +102,7 @@ def load_wan(checkpoint_info, diffusers_load_config={}):
     del transformer_2
 
     sd_hijack_te.init_hijack(pipe)
-    from modules.video_models import video_vae
-    pipe.vae.orig_decode = pipe.vae.decode
-    pipe.vae.decode = video_vae.hijack_vae_decode
+    sd_hijack_vae.init_hijack(pipe)
 
     devices.torch_gc()
     return pipe