mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 17:24:32 +02:00
explicit te hijack in pipeline loaders
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user