mirror of
https://github.com/vladmandic/automatic
synced 2026-08-27 15:41:00 +02:00
CalamitousFelicitousness
8c884c1e02
One UNET override cannot serve dual-transformer arches: ideogram4 conditional/unconditional and wan combined-stage experts need separate files, and previously a single override landed on both experts. - sd_unet_secondary option with per-slot tracking, consumed-state sync, arch-change reset, and incompatible-override fallback - dropdown renders beside the primary, follows it into quicksettings, and is visible only for dual-transformer model types - ideogram4 native single-file spec with a quant-aware fused-qkv converter; such converters run before comfy_quant detection via TransformerSpec.converter_handles_quant - quicksettings render in configured order (sort keyed on the option object and always fell back to alphabetical) - post-load dtype warning skips quantized transformers
99 lines
4.7 KiB
Python
99 lines
4.7 KiB
Python
import dataclasses
|
|
|
|
import diffusers
|
|
import transformers
|
|
from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
|
|
from modules.logger import log
|
|
from pipelines import generic
|
|
|
|
|
|
def pin_transformers(transformer, unconditional_transformer) -> bool:
|
|
"""Keep both transformers resident under balanced offload when they fit the budget.
|
|
Every denoise step runs both transformers, so balanced offload ping-pongs them across
|
|
PCIe each step. ``offload_never`` makes ``offload_allowed`` skip the per-step pre-sweep so
|
|
they stay resident, but only when both fit ``gpu_memory * max watermark`` (which leaves the
|
|
watermark headroom for activations); otherwise the normal offload path is kept.
|
|
"""
|
|
if shared.opts.diffusers_offload_mode != 'balanced' or shared.gpu_memory <= 0 or not shared.opts.model_ideogram4_pin:
|
|
return False
|
|
if transformer is None or unconditional_transformer is None: # if cg is disabled we dont need to pin
|
|
return False
|
|
if id(transformer) == id(unconditional_transformer): # if we're using the same transformer for both, no need to pin
|
|
return False
|
|
size_gb = sum(p.numel() * p.element_size() for m in (transformer, unconditional_transformer) for p in m.parameters()) / (1024 ** 3)
|
|
budget_gb = shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory
|
|
fits = size_gb <= budget_gb
|
|
if fits:
|
|
transformer.offload_never = True
|
|
unconditional_transformer.offload_never = True
|
|
log.debug(f'Load model: type=Ideogram4 offload=balanced transformers={size_gb:.2f} budget={budget_gb:.2f} action={"pin" if fits else "default"}')
|
|
return fits
|
|
|
|
|
|
def load_ideogram4(checkpoint_info, diffusers_load_config=None):
|
|
if diffusers_load_config is None:
|
|
diffusers_load_config = {}
|
|
repo_id = sd_models.path_to_repo(checkpoint_info)
|
|
sd_models.hf_auth_check(checkpoint_info)
|
|
load_args, _ = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
|
|
log.debug(f'Load model: type=Ideogram4 repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args} conditional={shared.opts.model_ideogram4_enable_cg} enhance={shared.opts.model_ideogram4_enable_pe} pin={shared.opts.model_ideogram4_pin}')
|
|
|
|
from pipelines.ideogram.ideogram4 import Ideogram4Pipeline
|
|
generic.set_pipeline('Ideogram4', Ideogram4Pipeline)
|
|
if repo_id is None or repo_id.lower() == 'none':
|
|
return None
|
|
|
|
from pipelines.ideogram import IDEOGRAM4_SPEC
|
|
transformer_cls = diffusers.Ideogram4Transformer2DModel
|
|
|
|
transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, subfolder="transformer", load_config=diffusers_load_config, native_spec=IDEOGRAM4_SPEC)
|
|
if shared.opts.model_ideogram4_enable_cg:
|
|
# the spec subfolder swap makes an unconditional override fetch its config from the matching subfolder
|
|
unconditional_transformer = generic.load_transformer(
|
|
repo_id, cls_name=transformer_cls, subfolder="unconditional_transformer", load_config=diffusers_load_config,
|
|
native_spec=dataclasses.replace(IDEOGRAM4_SPEC, subfolder="unconditional_transformer"), override_slot='secondary',
|
|
)
|
|
else:
|
|
unconditional_transformer = None
|
|
text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3VLModel, load_config=diffusers_load_config)
|
|
|
|
components = {
|
|
'transformer': transformer,
|
|
'unconditional_transformer': unconditional_transformer,
|
|
'text_encoder': text_encoder,
|
|
}
|
|
|
|
prompt_enhancer_head = None
|
|
if shared.opts.model_ideogram4_enable_pe:
|
|
enhancer_repo_id = "diffusers/qwen3-vl-8b-instruct-lm-head"
|
|
pe_load_args, pe_quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
|
|
enhancer_cls = diffusers.Ideogram4PromptEnhancerHead
|
|
log.debug(f'Load model: enhancer="{enhancer_repo_id}" cls={enhancer_cls.__name__}')
|
|
prompt_enhancer_head = enhancer_cls.from_pretrained(
|
|
enhancer_repo_id,
|
|
cache_dir=shared.opts.hfcache_dir,
|
|
**pe_load_args,
|
|
**pe_quant_args,
|
|
)
|
|
components['prompt_enhancer_head'] = prompt_enhancer_head
|
|
|
|
pin_transformers(transformer, unconditional_transformer)
|
|
pipe = Ideogram4Pipeline.from_pretrained(
|
|
repo_id,
|
|
cache_dir=shared.opts.diffusers_dir,
|
|
**components,
|
|
**load_args,
|
|
)
|
|
|
|
pipe.task_args = {
|
|
'output_type': 'pil',
|
|
'prompt_upsampling': shared.opts.model_ideogram4_enable_pe,
|
|
}
|
|
|
|
del transformer, unconditional_transformer, text_encoder, prompt_enhancer_head
|
|
sd_hijack_te.init_hijack(pipe)
|
|
sd_hijack_vae.init_hijack(pipe)
|
|
devices.torch_gc(force=True, reason='load')
|
|
|
|
return pipe
|