explicit te hijack in pipeline loaders

Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2025-08-15 16:55:01 -04:00
parent 65847bc022
commit 8af64b1e10
14 changed files with 27 additions and 22 deletions
+1 -1
View File
@@ -1131,7 +1131,7 @@ def clear_caches(full:bool=False):
lora_common.previously_loaded_networks.clear()
lora_load.lora_cache.clear()
if full:
shared.log.debug(f'Cache clear')
shared.log.debug('Cache clear')
sd_offload.offload_hook_instance = None
+2 -2
View File
@@ -1,6 +1,6 @@
import diffusers
import transformers
from modules import shared, devices, sd_models, model_quant
from modules import shared, devices, sd_models, model_quant, sd_hijack_te
from pipelines import generic
@@ -26,6 +26,6 @@ def load_chroma(checkpoint_info, diffusers_load_config={}):
diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["chroma"] = diffusers.ChromaImg2ImgPipeline
del text_encoder
del transformer
sd_hijack_te.init_hijack(pipe)
devices.torch_gc(force=True, reason='load')
return pipe
+1 -2
View File
@@ -26,9 +26,8 @@ def load_flex(checkpoint_info, diffusers_load_config={}):
diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["flex2"] = Flex2Pipeline
diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["flex2"] = Flex2Pipeline
sd_hijack_te.init_hijack(pipe)
del text_encoder_2
del transformer
sd_hijack_te.init_hijack(pipe)
devices.torch_gc()
return pipe
+1 -2
View File
@@ -29,9 +29,8 @@ def load_flite(checkpoint_info, diffusers_load_config={}):
**load_args,
)
sd_hijack_te.init_hijack(pipe)
del text_encoder
del dit_model
sd_hijack_te.init_hijack(pipe)
devices.torch_gc()
return pipe
+2 -1
View File
@@ -1,7 +1,7 @@
import os
import diffusers
import transformers
from modules import shared, devices, sd_models, model_quant
from modules import shared, devices, sd_models, model_quant, sd_hijack_te
from pipelines import generic
@@ -79,5 +79,6 @@ def load_flux(checkpoint_info, diffusers_load_config={}):
from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe
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
+1 -2
View File
@@ -66,11 +66,10 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
**load_args,
)
sd_hijack_te.init_hijack(pipe)
del text_encoder_3
del text_encoder_4
del tokenizer_4
del transformer
sd_hijack_te.init_hijack(pipe)
devices.torch_gc()
return pipe
+2 -1
View File
@@ -1,6 +1,6 @@
import transformers
import diffusers
from modules import shared, sd_models, devices, model_quant
from modules import shared, sd_models, devices, model_quant, sd_hijack_te
from pipelines import generic
@@ -25,5 +25,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
+2 -1
View File
@@ -1,9 +1,9 @@
import torch
import diffusers
from modules import shared, devices, sd_hijack_te
def load_kolors(_checkpoint_info, diffusers_load_config={}):
from modules import shared, devices
diffusers_load_config['variant'] = "fp16"
if 'torch_dtype' not in diffusers_load_config:
diffusers_load_config['torch_dtype'] = torch.float16
@@ -16,5 +16,6 @@ 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
+1
View File
@@ -12,6 +12,7 @@ def load_lumina(_checkpoint_info, diffusers_load_config={}):
cache_dir = shared.opts.diffusers_dir,
**load_config,
)
sd_hijack_te.init_hijack(pipe)
devices.torch_gc(force=True, reason='load')
return pipe
+2 -1
View File
@@ -1,9 +1,9 @@
import transformers
import diffusers
from modules import shared, devices, modelloader, sd_models, shared_items, sd_hijack_te
def load_meissonic(checkpoint_info, diffusers_load_config={}):
from modules import shared, devices, modelloader, sd_models, shared_items
from pipelines.meissonic.transformer import Transformer2DModel as TransformerMeissonic
from pipelines.meissonic.scheduler import Scheduler as MeissonicScheduler
from pipelines.meissonic.pipeline import MeissonicPipeline
@@ -52,5 +52,6 @@ def load_meissonic(checkpoint_info, diffusers_load_config={}):
diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["meissonic"] = MeissonicPipeline
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
+2 -1
View File
@@ -1,5 +1,5 @@
import diffusers
from modules import shared, devices, sd_models, model_quant
from modules import shared, devices, sd_models, model_quant, sd_hijack_te
def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument
@@ -25,5 +25,6 @@ def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=u
**load_config,
)
sd_hijack_te.init_hijack(pipe)
devices.torch_gc(force=True, reason='load')
return pipe
+6 -6
View File
@@ -1,16 +1,15 @@
from modules import shared, devices, sd_models, model_quant
import diffusers
from modules import shared, devices, sd_models, model_quant, sd_hijack_te
def load_omnigen2(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument
repo_id = sd_models.path_to_repo(checkpoint_info)
from pipelines.omnigen2 import OmniGen2Pipeline, OmniGen2Transformer2DModel, Qwen2_5_VLForConditionalGeneration
import diffusers
from diffusers import pipelines
diffusers.OmniGen2Pipeline = OmniGen2Pipeline # monkey-pathch
pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline
pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline
pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline
diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline
diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline
diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline
load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Model')
transformer = OmniGen2Transformer2DModel.from_pretrained(
@@ -42,5 +41,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
+2 -1
View File
@@ -1,7 +1,7 @@
import transformers
import diffusers
from huggingface_hub import file_exists
from modules import shared, devices, sd_models, model_quant
from modules import shared, devices, sd_models, model_quant, sd_hijack_te
from pipelines import generic
@@ -33,5 +33,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
+2 -1
View File
@@ -1,6 +1,6 @@
import diffusers
import transformers
from modules import shared, devices, sd_models, model_quant
from modules import shared, devices, sd_models, model_quant, sd_hijack_te
from pipelines import generic
@@ -31,5 +31,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