diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py index a454eff91..d87654c95 100644 --- a/modules/lora/lora_load.py +++ b/modules/lora/lora_load.py @@ -19,7 +19,7 @@ exclude_errors = [ # shared.sd_model_type -> dotted module path of a pipeline native loader # exposing ``try_load(name, network_on_disk, lora_scale)``. New archs add an -# entry here and ship a per-arch ``try_load`` (either binding native_loader's +# entry here and ship a per-arch ``try_load`` (either binding native_adapter's # generic helpers via try_load_chain, or rolling their own). _NATIVE_DISPATCH = { 'zimage': 'pipelines.z_image.zimage_lora', diff --git a/modules/lora/native_loader.py b/modules/lora/native_adapter.py similarity index 98% rename from modules/lora/native_loader.py rename to modules/lora/native_adapter.py index d0a1679ef..f2bed9ec3 100644 --- a/modules/lora/native_loader.py +++ b/modules/lora/native_adapter.py @@ -1,7 +1,6 @@ """Shared scaffolding for native adapter loaders. -The four native adapter loaders (z-image, chroma, ernie, flux2) all implement -the same algorithm: +Each per-arch native adapter loader implements the same algorithm: 1. Read the safetensors state dict 2. Test for family-specific markers; bail out if absent @@ -25,7 +24,7 @@ diffusers paths plus optional chunk descriptors). Per-arch loader modules import this module and pass their own ``prefixes``, ``bare_prefixes``, ``bare_diffusers_prefixes``, and ``resolve_targets`` to the -generic helpers. Loader business logic itself lands in subsequent commits. +generic helpers. """ import os @@ -45,7 +44,8 @@ from modules.lora import lora_common as l # Universal prefix list shared by every native arch loader. Per-arch loaders -# extend this with arch-specific entries (e.g. flux2 adds ``"lycoris_"``). +# extend this with arch-specific entries when their files use additional +# vendor-specific naming conventions. KNOWN_PREFIXES_DEFAULT = ("diffusion_model.", "transformer.", "lora_unet_") @@ -56,9 +56,10 @@ KNOWN_PREFIXES_DEFAULT = ("diffusion_model.", "transformer.", "lora_unet_") BARE_DIFFUSERS_PREFIX_USED = "bare_diffusers" -# Default network-key prefix. Single-component arches (flux2, zimage, chroma, -# ernie) keep this default; multi-component arches (anima: transformer plus -# llm_adapter plus text_encoder) pass a callable that picks per ``prefix_used``. +# Default network-key prefix. Single-component arches keep this default; +# multi-component arches (those with separate text-encoder or adapter +# components alongside the transformer) pass a callable that picks per +# ``prefix_used``. NETWORK_PREFIX_DEFAULT = "lora_transformer_" diff --git a/modules/sd_models.py b/modules/sd_models.py index cc36d57ea..206bdea09 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1412,6 +1412,19 @@ def reload_model_weights(sd_model=None, info=None, op='model', force=False, revi jobid = shared.state.begin('Load model') if sd_model is None: sd_model = model_data.sd_model if op == 'model' or op == 'dict' else model_data.sd_refiner + loaded_ckpt = getattr(sd_model, 'sd_checkpoint_info', None) if sd_model is not None else None + changed_checkpoint = loaded_ckpt is None or checkpoint_info is None or loaded_ckpt.filename != checkpoint_info.filename + if op == 'model' and sd_model is not None and changed_checkpoint and shared.opts.sd_unet not in (None, 'Default', 'None'): + old_class = type(sd_model).__name__ + try: + new_pipeline, _ = sd_detect.detect_pipeline(checkpoint_info.path, op) + except Exception: + new_pipeline = None + new_class = getattr(new_pipeline, '__name__', None) + if new_class is not None and new_class != old_class: + log.info(f'Load model: pipeline cls={old_class} changed={new_class} unet="{shared.opts.sd_unet}" set to default') + shared.opts.data["sd_unet"] = 'Default' + sd_unet.loaded_unet = None if sd_model is None: # previous model load failed current_checkpoint_info = None else: diff --git a/modules/sd_models_utils.py b/modules/sd_models_utils.py index de241c8ca..773d955b0 100644 --- a/modules/sd_models_utils.py +++ b/modules/sd_models_utils.py @@ -100,26 +100,29 @@ def read_state_dict(checkpoint_file, map_location=None, what:str='model'): # pyl if not os.path.isfile(checkpoint_file): log.error(f'Load dict: path="{checkpoint_file}" not a file') return None + _, extension = os.path.splitext(checkpoint_file) + if extension.lower() == ".ckpt" and shared.opts.sd_disable_ckpt: + log.warning(f"Checkpoint loading disabled: {checkpoint_file}") + return None try: pl_sd = None - with progress.open(checkpoint_file, 'rb', description=f'[cyan]Load {what}: [yellow]{checkpoint_file}', auto_refresh=True, console=console) as f: - _, extension = os.path.splitext(checkpoint_file) - if extension.lower() == ".ckpt" and shared.opts.sd_disable_ckpt: - log.warning(f"Checkpoint loading disabled: {checkpoint_file}") - return None - if shared.opts.stream_load: - if extension.lower() == ".safetensors": - buffer = f.read() - pl_sd = safetensors.torch.load(buffer) - else: - buffer = io.BytesIO(f.read()) - pl_sd = torch.load(buffer, map_location='cpu') - else: - if extension.lower() == ".safetensors": - pl_sd = safetensors.torch.load_file(checkpoint_file, device='cpu') + # safetensors.torch.load_file opens its own handle by path, so wrapping + # with progress.open leaves the bar stuck at 0/total. Skip the wrapper + # on that path; other paths actually read through f and update. + if extension.lower() == ".safetensors" and not shared.opts.stream_load: + pl_sd = safetensors.torch.load_file(checkpoint_file, device='cpu') + else: + with progress.open(checkpoint_file, 'rb', description=f'[cyan]Load {what}: [yellow]{checkpoint_file}', auto_refresh=True, console=console) as f: + if shared.opts.stream_load: + if extension.lower() == ".safetensors": + buffer = f.read() + pl_sd = safetensors.torch.load(buffer) + else: + buffer = io.BytesIO(f.read()) + pl_sd = torch.load(buffer, map_location='cpu') else: pl_sd = torch.load(f, map_location='cpu') - sd = get_state_dict_from_checkpoint(pl_sd) + sd = get_state_dict_from_checkpoint(pl_sd) del pl_sd except Exception as e: errors.display(e, f'Load model: {checkpoint_file}') diff --git a/modules/sd_unet.py b/modules/sd_unet.py index db0ffbe22..a371fe8ea 100644 --- a/modules/sd_unet.py +++ b/modules/sd_unet.py @@ -46,6 +46,12 @@ def load_unet(model, repo_id: str | None = None): return if shared.opts.sd_unet == 'Default' or shared.opts.sd_unet == 'None': + # Switching back to Default reverts a previously-loaded custom transformer. + if loaded_unet in (None, 'Default', 'None'): + return + log.info(f'Load module: type=UNet name="Default" (was="{loaded_unet}") reverting to base transformer') + loaded_unet = shared.opts.sd_unet + sd_models.reload_model_weights(force=True) return if shared.opts.sd_unet not in list(unet_dict): @@ -74,7 +80,7 @@ def load_unet(model, repo_id: str | None = None): model.prior_pipe.text_encoder = prior_text_encoder.to(devices.device, dtype=devices.dtype) elif any([m in model.__class__.__name__ for m in dit_models]) or hasattr(model, 'transformer'): # noqa: C419 # pylint: disable=use-a-generator loaded_unet = shared.opts.sd_unet - sd_models.load_diffuser() # TODO model load: force-reloading entire model as loading transformers only leads to massive memory usage + sd_models.reload_model_weights(force=True) # full reload: in-place transformer swap leaks memory else: if not hasattr(model, 'unet') or model.unet is None: log.error('Load module: type=UNET not found in current model') diff --git a/modules/ui_settings.py b/modules/ui_settings.py index 6441c2e90..aa9e6a4a9 100644 --- a/modules/ui_settings.py +++ b/modules/ui_settings.py @@ -391,6 +391,10 @@ def create_quicksettings(interfaces): if shared.opts.notification_audio_enable and os.path.exists(os.path.join(paths.script_path, shared.opts.notification_audio_path)): gr.Audio(interactive=False, value=os.path.join(paths.script_path, shared.opts.notification_audio_path), elem_id="audio_notification", visible=False) + def sync_checkpoint_unet(value, progress=False, force=False): + checkpoint_update, settings_text = run_settings_single(value, key='sd_model_checkpoint', progress=progress, force=force) + return checkpoint_update, get_value_for_setting('sd_unet'), settings_text + for k, _item in quicksettings_list: component = shared.settings_components[k] info = shared.opts.data_labels[k] @@ -405,20 +409,32 @@ def create_quicksettings(interfaces): change_handlers = [component.blur] else: change_handlers = [component.release if hasattr(component, 'release') else component.change] + progress_flag = info.refresh is not None + if k == 'sd_model_checkpoint': + def fn(value, progress=progress_flag): + return sync_checkpoint_unet(value, progress=progress) + outputs = [component, shared.settings_components['sd_unet'], text_settings] + else: + def fn(value, k=k, progress=progress_flag): + return run_settings_single(value, key=k, progress=progress) + outputs = [component, text_settings] for change_handler in change_handlers: change_handler( - fn=lambda value, k=k, progress=info.refresh is not None: run_settings_single(value, key=k, progress=progress), + fn=fn, inputs=[component], - outputs=[component, text_settings], + outputs=outputs, show_progress='full' if info.refresh is not None else 'hidden', ) + def sync_checkpoint_unet_forced(value, _dummy): + return sync_checkpoint_unet(value, force=True) + button_set_checkpoint = gr.Button('Change model', elem_id='change_checkpoint', visible=False) button_set_checkpoint.click( - fn=lambda value, _: run_settings_single(value, key='sd_model_checkpoint', force=True), + fn=sync_checkpoint_unet_forced, _js="consumeDesiredCheckpointName", inputs=[shared.settings_components['sd_model_checkpoint'], dummy_component], - outputs=[shared.settings_components['sd_model_checkpoint'], text_settings], + outputs=[shared.settings_components['sd_model_checkpoint'], shared.settings_components['sd_unet'], text_settings], ) button_set_refiner = gr.Button('Change refiner', elem_id='change_refiner', visible=False) button_set_refiner.click( diff --git a/pipelines/anima/__init__.py b/pipelines/anima/__init__.py index e69de29bb..6d7953090 100644 --- a/pipelines/anima/__init__.py +++ b/pipelines/anima/__init__.py @@ -0,0 +1,40 @@ +"""Anima pipeline package. + +Exports :data:`ANIMA_SPEC` for use by :mod:`pipelines.model_anima` together +with :mod:`pipelines.native_transformer`. The spec captures the Anima-specific +knobs that differ from the native-loader defaults: + +- The bundled ``llm_adapter`` sibling: Anima community files frequently inline + the custom AnimaLLMAdapter weights in the same safetensors as the + transformer. The resolved adapter class is supplied at load time via + ``sibling_classes`` because AnimaLLMAdapter is loaded dynamically through + ``trust_remote_code`` and is not available at import time. +- Cosmos 1.0 structural marker: any community file whose state dict contains + a Cosmos 1.0 nested key (``net.blocks.block1.*``) is rejected with a clear + error since Anima is Cosmos 2.0 only. +- All other knobs (prefixes, ``acceptable_missing`` buffers) use the defaults + from :mod:`pipelines.native_transformer`. +""" + +import diffusers +from diffusers.loaders.single_file_utils import convert_cosmos_transformer_checkpoint_to_diffusers + +from pipelines.native_transformer import TransformerSpec, SiblingSpec + + +ANIMA_SPEC = TransformerSpec( + cls=diffusers.CosmosTransformer3DModel, + converter=convert_cosmos_transformer_checkpoint_to_diffusers, + siblings={ + 'llm_adapter': SiblingSpec( + subfolder='llm_adapter', + inline_prefix='llm_adapter.', + ), + }, + forbidden_markers=( + ( + 'net.blocks.block1.blocks.0.block.attn.to_q.0.weight', + 'unsupported Cosmos 1.0 structure', + ), + ), +) diff --git a/pipelines/anima/anima_lora.py b/pipelines/anima/anima_lora.py index 39530c243..cc9c01dbf 100644 --- a/pipelines/anima/anima_lora.py +++ b/pipelines/anima/anima_lora.py @@ -25,23 +25,23 @@ already stamped on the network_layer_mapping. Adapter and TE paths bypass the rename and are flattened verbatim. Network-key construction (transformer vs llm_adapter vs te) is parameterized -in :mod:`modules.lora.native_loader` via the ``network_prefix`` kwarg; this +in :mod:`modules.lora.native_adapter` via the ``network_prefix`` kwarg; this module supplies :func:`network_prefix_for` to pick per ``prefix_used``. Family-specific dispatch (LoRA, LoHA, LoKR, OFT, IA3, GLoRA, Norm, Full) is -inherited from native_loader's generics; alpha / scale / DoRA flow through +inherited from native_adapter's generics; alpha / scale / DoRA flow through the standard ``NetworkWeights.w`` slots rather than being baked into the factor weights at load time. """ from collections import OrderedDict -from modules.lora import native_loader +from modules.lora import native_adapter # === Arch-specific prefix configuration === # # Order matters: longer / more-specific prefixes must precede shorter ones, -# because :func:`native_loader.parse_key` returns the first match. Both +# because :func:`native_adapter.parse_key` returns the first match. Both # ``diffusion_model.llm_adapter.`` and ``text_encoders.qwen3_06b.transformer.model.`` # start with ``diffusion_model.`` / ``text_encoders.`` so they must be listed first. @@ -58,37 +58,37 @@ ANIMA_PREFIXES = ( # Tests address these through the anima_lora module surface; sibling pipelines # do the same (see flux2_lora / zimage_lora / ernie_lora). -LORA_SUFFIXES = native_loader.LORA_SUFFIXES -LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES -LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES -OFT_SUFFIXES = native_loader.OFT_SUFFIXES -IA3_SUFFIXES = native_loader.IA3_SUFFIXES -GLORA_SUFFIXES = native_loader.GLORA_SUFFIXES -NORM_SUFFIXES = native_loader.NORM_SUFFIXES -FULL_SUFFIXES = native_loader.FULL_SUFFIXES +LORA_SUFFIXES = native_adapter.LORA_SUFFIXES +LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES +LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES +OFT_SUFFIXES = native_adapter.OFT_SUFFIXES +IA3_SUFFIXES = native_adapter.IA3_SUFFIXES +GLORA_SUFFIXES = native_adapter.GLORA_SUFFIXES +NORM_SUFFIXES = native_adapter.NORM_SUFFIXES +FULL_SUFFIXES = native_adapter.FULL_SUFFIXES -LORA_MARKERS = native_loader.LORA_MARKERS -LOKR_MARKERS = native_loader.LOKR_MARKERS -LOHA_MARKERS = native_loader.LOHA_MARKERS -OFT_MARKERS = native_loader.OFT_MARKERS -IA3_MARKERS = native_loader.IA3_MARKERS -GLORA_MARKERS = native_loader.GLORA_MARKERS -NORM_MARKERS = native_loader.NORM_MARKERS -FULL_MARKERS = native_loader.FULL_MARKERS +LORA_MARKERS = native_adapter.LORA_MARKERS +LOKR_MARKERS = native_adapter.LOKR_MARKERS +LOHA_MARKERS = native_adapter.LOHA_MARKERS +OFT_MARKERS = native_adapter.OFT_MARKERS +IA3_MARKERS = native_adapter.IA3_MARKERS +GLORA_MARKERS = native_adapter.GLORA_MARKERS +NORM_MARKERS = native_adapter.NORM_MARKERS +FULL_MARKERS = native_adapter.FULL_MARKERS -SUFFIX_NORMALIZE = native_loader.SUFFIX_NORMALIZE -BARE_DIFFUSERS_PREFIX_USED = native_loader.BARE_DIFFUSERS_PREFIX_USED -has_marker = native_loader.has_marker +SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE +BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED +has_marker = native_adapter.has_marker def parse_key(key, suffixes): - """Anima-bound :func:`native_loader.parse_key`.""" - return native_loader.parse_key(key, suffixes, prefixes=ANIMA_PREFIXES) + """Anima-bound :func:`native_adapter.parse_key`.""" + return native_adapter.parse_key(key, suffixes, prefixes=ANIMA_PREFIXES) def group_by_suffixes(state_dict, suffixes): - """Anima-bound :func:`native_loader.group_by_suffixes`.""" - return native_loader.group_by_suffixes(state_dict, suffixes, prefixes=ANIMA_PREFIXES) + """Anima-bound :func:`native_adapter.group_by_suffixes`.""" + return native_adapter.group_by_suffixes(state_dict, suffixes, prefixes=ANIMA_PREFIXES) # === Cosmos 2.0 path rename (transformer only) === @@ -174,7 +174,7 @@ def network_prefix_for(prefix_used): return "lora_transformer_" -# === Native loaders (thin wrappers over native_loader generics) === +# === Native loaders (thin wrappers over native_adapter generics) === _BIND_KWARGS = dict( resolve_targets=resolve_targets, @@ -185,40 +185,40 @@ _BIND_KWARGS = dict( def try_load_lora(name, network_on_disk, lora_scale): - return native_loader.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_lokr(name, network_on_disk, lora_scale): - return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_loha(name, network_on_disk, lora_scale): - return native_loader.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_oft(name, network_on_disk, lora_scale): - return native_loader.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_ia3(name, network_on_disk, lora_scale): - return native_loader.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_glora(name, network_on_disk, lora_scale): - return native_loader.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_norm(name, network_on_disk, lora_scale): - return native_loader.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_full(name, network_on_disk, lora_scale): - return native_loader.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load(name, network_on_disk, lora_scale): """Run every Anima family loader, merge any that match.""" - return native_loader.try_load_chain( + return native_adapter.try_load_chain( name, network_on_disk, lora_scale, family_loaders=( try_load_lora, try_load_lokr, try_load_loha, try_load_oft, diff --git a/pipelines/anima/anima_transformer.py b/pipelines/anima/anima_transformer.py deleted file mode 100644 index 06692dc4e..000000000 --- a/pipelines/anima/anima_transformer.py +++ /dev/null @@ -1,217 +0,0 @@ -"""Anima custom-transformer loader. - -Called from :func:`pipelines.model_anima.load_anima` when the user has selected -a transformer file via the UNET dropdown (``shared.opts.sd_unet``). Reads the -safetensors directly, strips the BFL-style prefix, splits off the bundled -``llm_adapter.*`` keys, and routes the two halves into the diffusers -``CosmosTransformer3DModel`` and the remote ``AnimaLLMAdapter`` respectively. - -The transformer half is run through diffusers' -``convert_cosmos_transformer_checkpoint_to_diffusers`` (Cosmos 2.0 branch), -whose rename table covers Anima's native key fragments exactly, so the -converted state dict drops cleanly into ``CosmosTransformer3DModel`` with no -ad-hoc renames needed here. The adapter half matches the base repo's -``llm_adapter/diffusion_pytorch_model.safetensors`` exactly, so it loads -as-is. - -Supported input formats (safetensors only; GGUF and .pth are rejected early): - -- Bare BFL keys: ``blocks.0.self_attn.q_proj.weight`` (e.g. ``rdbtAnima_v027``) -- ``model.diffusion_model.`` prefix (e.g. ``animaika_v35``) -- ``diffusion_model.`` prefix (ComfyUI-style export) -- ``net.`` prefix (NVIDIA/Cosmos native export, e.g. ``animayume_v04``) - -Quantization: SDNQ (pre/post/auto) and ``layerwise_quantization`` are honored. -SDNQ pre-mode is applied post-load here because this path bypasses -``from_pretrained``, where ``quantization_config`` normally takes effect. -TensorRT (``NVIDIAModelOptConfig``) is not supported and is skipped with a -warning. GGUF would require a separate converter and is not supported. -""" - -import os -import time -import diffusers -import huggingface_hub as hf -from modules import shared, devices, sd_models, model_quant, errors -from modules.logger import log - - -KNOWN_PREFIXES = ("model.diffusion_model.", "diffusion_model.", "net.") -ADAPTER_PREFIX = "llm_adapter." -COSMOS_1_MARKER = "net.blocks.block1.blocks.0.block.attn.to_q.0.weight" - -# Buffer keys that CosmosTransformer3DModel creates at __init__ time and do -# not appear in trainer state dicts. Acceptable in the "missing" set. -ACCEPTABLE_MISSING = ("rope.", "pos_embedder.", "learnable_pos_embed.") - - -def load_custom_transformer(repo_id, local_file, diffusers_load_config, adapter_cls): - """Load a custom Anima transformer (and optional bundled adapter) from a safetensors file. - - Returns ``(transformer, llm_adapter_or_none)``. If the file does not bundle - an adapter, the second element is ``None`` and the caller should fall back - to the base repo's adapter via ``AnimaLLMAdapter.from_pretrained``. - Raises on any hard failure (prefix mix, shape mismatch, missing configs). - """ - t0 = time.time() - - if not local_file.lower().endswith('.safetensors'): - raise ValueError(f'Load model: type=Anima custom transformer requires .safetensors, got "{local_file}"') - - # from_config + load_state_dict does not consume load_args (device_map, - # torch_dtype, etc.); dtype is applied via explicit .to() below. Only - # quant_type is read from this call. - _, quant_args = model_quant.get_dit_args( - diffusers_load_config, module='Model', device_map=True, allow_quant=True, - ) - quant_type = model_quant.get_quant_type(quant_args) - - transformer_cfg = fetch_component_config(repo_id, 'transformer/config.json') - adapter_cfg = fetch_component_config(repo_id, 'llm_adapter/config.json') - - state_dict = sd_models.read_state_dict(local_file, what='transformer') - state_dict = strip_prefix(state_dict) - transformer_sd, adapter_sd = partition_adapter(state_dict) - del state_dict - - if COSMOS_1_MARKER in transformer_sd: - raise ValueError(f'Load model: type=Anima custom transformer has unsupported Cosmos 1.0 structure (file="{local_file}")') - - log.info(f'Load model: type=Anima custom="{os.path.basename(local_file)}" transformer_keys={len(transformer_sd)} adapter_keys={len(adapter_sd)}') - - transformer = build_transformer(transformer_sd, transformer_cfg, quant_args, quant_type) - del transformer_sd - devices.torch_gc() - - if adapter_sd: - llm_adapter = build_adapter(adapter_sd, adapter_cfg, adapter_cls) - else: - log.info('Load model: type=Anima custom transformer has no bundled adapter, caller will load from base repo') - llm_adapter = None - - sd_models.allow_post_quant = False # transformer already quantized above - devices.torch_gc() - log.debug(f'Load model: type=Anima custom transformer time={time.time()-t0:.2f}') - return transformer, llm_adapter - - -def fetch_component_config(repo_id, relative_path): - """Download and parse a component config.json from the base repo.""" - try: - local = hf.hf_hub_download(repo_id, filename=relative_path, cache_dir=shared.opts.diffusers_dir) - except Exception as e: - raise RuntimeError(f'Load model: type=Anima failed to download {relative_path} from repo="{repo_id}": {e}') from e - return shared.readfile(local, as_type='dict') - - -def strip_prefix(state_dict): - """Detect and uniformly strip the BFL-style prefix from all keys. - - Supported prefixes (longest first, so ``model.diffusion_model.`` beats ``diffusion_model.``): - ``model.diffusion_model.``, ``diffusion_model.``, or no prefix. Raises - ValueError if some keys match the dominant prefix and others do not, - since mixed prefixes indicate a malformed file. - """ - counts = {p: sum(1 for k in state_dict if k.startswith(p)) for p in KNOWN_PREFIXES} - total = len(state_dict) - dominant = max(counts, key=counts.get) - if counts[dominant] == 0: - log.debug('Load model: type=Anima custom transformer prefix=bare') - return state_dict - if counts[dominant] != total: - raise ValueError( - f'Load model: type=Anima custom transformer has mixed prefixes ' - f'(total={total} {dominant}={counts[dominant]})' - ) - log.debug(f'Load model: type=Anima custom transformer prefix="{dominant}"') - offset = len(dominant) - return {k[offset:]: v for k, v in state_dict.items()} - - -def partition_adapter(state_dict): - """Split into (transformer_sd, adapter_sd) by the ``llm_adapter.`` prefix.""" - transformer_sd = {} - adapter_sd = {} - for key, value in state_dict.items(): - if key.startswith(ADAPTER_PREFIX): - adapter_sd[key[len(ADAPTER_PREFIX):]] = value - else: - transformer_sd[key] = value - return transformer_sd, adapter_sd - - -def build_transformer(transformer_sd, transformer_cfg, quant_args, quant_type): - """Convert, instantiate, load, dtype-cast, quantize, and (if offloading) move to CPU.""" - from diffusers.loaders.single_file_utils import convert_cosmos_transformer_checkpoint_to_diffusers - try: - converted = convert_cosmos_transformer_checkpoint_to_diffusers(transformer_sd) - transformer = diffusers.CosmosTransformer3DModel.from_config(transformer_cfg) - missing, unexpected = transformer.load_state_dict(converted, strict=False) - validate_state_dict_load('transformer', missing, unexpected) - del converted - devices.torch_gc() - transformer = transformer.to(dtype=devices.dtype) - except Exception as e: - log.error(f'Load model: type=Anima transformer load failed: {e}') - errors.display(e, 'Load') - raise - - apply_quant(transformer, quant_type) - - if shared.opts.diffusers_offload_mode != 'none': - sd_models.move_model(transformer, devices.cpu) - - if not hasattr(transformer, 'quantization_config'): - if hasattr(transformer, 'config') and hasattr(transformer.config, 'quantization_config'): - transformer.quantization_config = transformer.config.quantization_config - elif (quant_type is not None) and (quant_args.get('quantization_config', None) is not None): - transformer.quantization_config = quant_args.get('quantization_config', None) - return transformer - - -def build_adapter(adapter_sd, adapter_cfg, adapter_cls): - """Instantiate AnimaLLMAdapter from the base repo config and load bundled weights.""" - try: - adapter = adapter_cls.from_config(adapter_cfg) - missing, unexpected = adapter.load_state_dict(adapter_sd, strict=False) - validate_state_dict_load('adapter', missing, unexpected) - adapter = adapter.to(dtype=devices.dtype) - except Exception as e: - log.error(f'Load model: type=Anima adapter load failed: {e}') - errors.display(e, 'Load') - raise - if shared.opts.diffusers_offload_mode != 'none': - sd_models.move_model(adapter, devices.cpu) - return adapter - - -def validate_state_dict_load(component, missing, unexpected): - """Raise ValueError if load_state_dict produced unexpected keys or non-buffer missing keys.""" - if unexpected: - sample = ', '.join(unexpected[:5]) - raise ValueError(f'Load model: type=Anima {component} has {len(unexpected)} unexpected keys (sample: {sample})') - hard_missing = [k for k in missing if not any(k.startswith(p) for p in ACCEPTABLE_MISSING)] - if hard_missing: - sample = ', '.join(hard_missing[:5]) - raise ValueError(f'Load model: type=Anima {component} missing {len(hard_missing)} required keys (sample: {sample})') - if missing: - log.debug(f'Load model: type=Anima {component} ignored {len(missing)} buffer-only missing keys') - - -def apply_quant(transformer, quant_type): - """Apply SDNQ / layerwise quantization to the bare transformer. - - SDNQ 'pre' and 'auto' would normally route through ``quantization_config`` - at ``from_pretrained`` time; since we bypass that boundary, we call the - per-module quant path directly. SDNQ 'post' and ``layerwise_quantization`` - go through ``do_post_load_quant`` as usual. - """ - if quant_type == 'NVIDIAModelOptConfig': - log.warning('Load model: type=Anima quant=TRT not supported on custom transformer path, skipping') - elif quant_type == 'SDNQConfig': - if shared.opts.sdnq_quantize_mode == 'pre': - log.info('Load model: type=Anima quant=SDNQ pre-mode applied post-load on custom transformer path') - model_quant.sdnq_quantize_model(transformer, op='transformer') - # allow=False avoids double-applying SDNQ in auto mode (applied directly - # above); post mode fires regardless of allow, and layerwise always fires. - model_quant.do_post_load_quant(transformer, allow=False) diff --git a/pipelines/bria/__init__.py b/pipelines/bria/__init__.py index e69de29bb..e6eb9aaee 100644 --- a/pipelines/bria/__init__.py +++ b/pipelines/bria/__init__.py @@ -0,0 +1,27 @@ +"""Bria pipeline package. + +Exports :data:`BRIA_SPEC` and :data:`BRIA_FIBO_SPEC` for use by +:mod:`pipelines.model_bria` together with :mod:`pipelines.native_transformer`. + +Bria ships two transformer variants: + +- The original Bria family uses a custom :class:`BriaTransformer2DModel` + imported from :mod:`pipelines.bria.transformer_bria`. The custom class + has no entry in diffusers' ``SINGLE_FILE_LOADABLE_CLASSES`` table. +- Bria FIBO (and FIBO Edit) use the upstream + :class:`diffusers.BriaFiboTransformer2DModel`, which also lacks a + ``SINGLE_FILE_LOADABLE_CLASSES`` entry. + +Both specs use the default 3-prefix detection +(``model.diffusion_model.``, ``diffusion_model.``, ``net.``), no +converter, and no siblings. +""" + +import diffusers + +from pipelines.bria.transformer_bria import BriaTransformer2DModel +from pipelines.native_transformer import TransformerSpec + + +BRIA_SPEC = TransformerSpec(cls=BriaTransformer2DModel) +BRIA_FIBO_SPEC = TransformerSpec(cls=diffusers.BriaFiboTransformer2DModel) diff --git a/pipelines/chroma/__init__.py b/pipelines/chroma/__init__.py index e69de29bb..b8eaa68c6 100644 --- a/pipelines/chroma/__init__.py +++ b/pipelines/chroma/__init__.py @@ -0,0 +1,18 @@ +"""Chroma pipeline package. + +Exports :data:`CHROMA_SPEC`. Chroma community files use BFL-style +``model.diffusion_model.``-prefixed keys that need renaming into the +diffusers naming convention, so the spec plugs in +:func:`convert_chroma_transformer_checkpoint_to_diffusers` explicitly. +""" + +import diffusers +from diffusers.loaders.single_file_utils import convert_chroma_transformer_checkpoint_to_diffusers + +from pipelines.native_transformer import TransformerSpec + + +CHROMA_SPEC = TransformerSpec( + cls=diffusers.ChromaTransformer2DModel, + converter=convert_chroma_transformer_checkpoint_to_diffusers, +) diff --git a/pipelines/chroma/chroma_lora.py b/pipelines/chroma/chroma_lora.py index 8354a98b5..965d5bf52 100644 --- a/pipelines/chroma/chroma_lora.py +++ b/pipelines/chroma/chroma_lora.py @@ -30,13 +30,13 @@ any ``_mod_lin`` / ``_modulation_lin`` keys land in ``unmapped``. LoRAs targeting the approximator pass through unchanged. """ -from modules.lora import native_loader -from modules.lora.native_loader import ChunkSpec +from modules.lora import native_adapter +from modules.lora.native_adapter import ChunkSpec # === Arch-specific prefix configuration === -KNOWN_PREFIXES = native_loader.KNOWN_PREFIXES_DEFAULT +KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT BARE_FLUX_PREFIXES = ("double_blocks.", "single_blocks.") @@ -57,24 +57,24 @@ LINEAR1_DIMS = [3072, 3072, 3072, 12288] # === Re-exports for test/back-compat === -LORA_SUFFIXES = native_loader.LORA_SUFFIXES -LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES -LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES -OFT_SUFFIXES = native_loader.OFT_SUFFIXES +LORA_SUFFIXES = native_adapter.LORA_SUFFIXES +LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES +LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES +OFT_SUFFIXES = native_adapter.OFT_SUFFIXES -LORA_MARKERS = native_loader.LORA_MARKERS -LOKR_MARKERS = native_loader.LOKR_MARKERS -LOHA_MARKERS = native_loader.LOHA_MARKERS -OFT_MARKERS = native_loader.OFT_MARKERS +LORA_MARKERS = native_adapter.LORA_MARKERS +LOKR_MARKERS = native_adapter.LOKR_MARKERS +LOHA_MARKERS = native_adapter.LOHA_MARKERS +OFT_MARKERS = native_adapter.OFT_MARKERS -SUFFIX_NORMALIZE = native_loader.SUFFIX_NORMALIZE -BARE_DIFFUSERS_PREFIX_USED = native_loader.BARE_DIFFUSERS_PREFIX_USED -has_marker = native_loader.has_marker +SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE +BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED +has_marker = native_adapter.has_marker def parse_key(key, suffixes): - """Chroma-bound :func:`native_loader.parse_key`.""" - return native_loader.parse_key( + """Chroma-bound :func:`native_adapter.parse_key`.""" + return native_adapter.parse_key( key, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, @@ -83,8 +83,8 @@ def parse_key(key, suffixes): def group_by_suffixes(state_dict, suffixes): - """Chroma-bound :func:`native_loader.group_by_suffixes`.""" - return native_loader.group_by_suffixes( + """Chroma-bound :func:`native_adapter.group_by_suffixes`.""" + return native_adapter.group_by_suffixes( state_dict, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, @@ -202,7 +202,7 @@ def _split_single_linear1(block_idx): return targets -# === Native loaders (thin wrappers over native_loader generics) === +# === Native loaders (thin wrappers over native_adapter generics) === _BIND_KWARGS = dict( @@ -215,24 +215,24 @@ _BIND_KWARGS = dict( def try_load_lora(name, network_on_disk, lora_scale): - return native_loader.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_lokr(name, network_on_disk, lora_scale): - return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_loha(name, network_on_disk, lora_scale): - return native_loader.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_oft(name, network_on_disk, lora_scale): - return native_loader.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load(name, network_on_disk, lora_scale): """Run every Chroma family loader, merge any that match.""" - return native_loader.try_load_chain( + return native_adapter.try_load_chain( name, network_on_disk, lora_scale, family_loaders=(try_load_lora, try_load_lokr, try_load_loha, try_load_oft), ) diff --git a/pipelines/chrono/__init__.py b/pipelines/chrono/__init__.py new file mode 100644 index 000000000..b9fe64887 --- /dev/null +++ b/pipelines/chrono/__init__.py @@ -0,0 +1,15 @@ +"""ChronoEdit pipeline package. + +Exports :data:`CHRONOEDIT_SPEC`. The minimum +``TransformerSpec(cls=ChronoEditTransformer3DModel)`` works because +ChronoEdit community files use BFL-style ``model.diffusion_model.`` +prefixed keys whose names match the diffusers state_dict verbatim after +prefix strip. No siblings, no converter, no forbidden markers. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +CHRONOEDIT_SPEC = TransformerSpec(cls=diffusers.ChronoEditTransformer3DModel) diff --git a/pipelines/cogview/__init__.py b/pipelines/cogview/__init__.py new file mode 100644 index 000000000..7163f543f --- /dev/null +++ b/pipelines/cogview/__init__.py @@ -0,0 +1,15 @@ +"""CogView pipeline package. + +Exports :data:`COGVIEW3_SPEC` and :data:`COGVIEW4_SPEC`. Both default +specs work because CogView community files use BFL-style +``model.diffusion_model.``-prefixed keys whose names match the diffusers +state_dict verbatim after prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +COGVIEW3_SPEC = TransformerSpec(cls=diffusers.CogView3PlusTransformer2DModel) +COGVIEW4_SPEC = TransformerSpec(cls=diffusers.CogView4Transformer2DModel) diff --git a/pipelines/ernie/__init__.py b/pipelines/ernie/__init__.py index e69de29bb..f01ec13ff 100644 --- a/pipelines/ernie/__init__.py +++ b/pipelines/ernie/__init__.py @@ -0,0 +1,16 @@ +"""ERNIE-Image pipeline package. + +Exports :data:`ERNIE_SPEC`. The minimum +``TransformerSpec(cls=ErnieImageTransformer2DModel)`` works because Ernie +trainer dumps use BFL-style ``model.diffusion_model.``-prefixed keys +whose names match the diffusers state_dict verbatim after prefix strip +(probed against a community finetune: full key overlap with zero +missing or unexpected). No siblings, no converter, no forbidden markers. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +ERNIE_SPEC = TransformerSpec(cls=diffusers.ErnieImageTransformer2DModel) diff --git a/pipelines/ernie/ernie_lora.py b/pipelines/ernie/ernie_lora.py index c7ad8d497..e76d716aa 100644 --- a/pipelines/ernie/ernie_lora.py +++ b/pipelines/ernie/ernie_lora.py @@ -16,12 +16,12 @@ modules (no fused QKV) and ``ErnieImageFeedForward`` exposes ``gate_proj``, straight passthrough; no chunking, no renames, no dispatch table. """ -from modules.lora import native_loader +from modules.lora import native_adapter # === Arch-specific prefix configuration === -KNOWN_PREFIXES = native_loader.KNOWN_PREFIXES_DEFAULT +KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT BARE_DIFFUSERS_PREFIXES = ( "layers.", "adaLN_modulation.", "final_norm.", "final_linear.", @@ -30,24 +30,24 @@ BARE_DIFFUSERS_PREFIXES = ( # === Re-exports for test/back-compat === -LORA_SUFFIXES = native_loader.LORA_SUFFIXES -LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES -LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES -OFT_SUFFIXES = native_loader.OFT_SUFFIXES +LORA_SUFFIXES = native_adapter.LORA_SUFFIXES +LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES +LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES +OFT_SUFFIXES = native_adapter.OFT_SUFFIXES -LORA_MARKERS = native_loader.LORA_MARKERS -LOKR_MARKERS = native_loader.LOKR_MARKERS -LOHA_MARKERS = native_loader.LOHA_MARKERS -OFT_MARKERS = native_loader.OFT_MARKERS +LORA_MARKERS = native_adapter.LORA_MARKERS +LOKR_MARKERS = native_adapter.LOKR_MARKERS +LOHA_MARKERS = native_adapter.LOHA_MARKERS +OFT_MARKERS = native_adapter.OFT_MARKERS -SUFFIX_NORMALIZE = native_loader.SUFFIX_NORMALIZE -BARE_DIFFUSERS_PREFIX_USED = native_loader.BARE_DIFFUSERS_PREFIX_USED -has_marker = native_loader.has_marker +SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE +BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED +has_marker = native_adapter.has_marker def parse_key(key, suffixes): - """ERNIE-bound :func:`native_loader.parse_key`.""" - return native_loader.parse_key( + """ERNIE-bound :func:`native_adapter.parse_key`.""" + return native_adapter.parse_key( key, suffixes, prefixes=KNOWN_PREFIXES, bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, @@ -55,8 +55,8 @@ def parse_key(key, suffixes): def group_by_suffixes(state_dict, suffixes): - """ERNIE-bound :func:`native_loader.group_by_suffixes`.""" - return native_loader.group_by_suffixes( + """ERNIE-bound :func:`native_adapter.group_by_suffixes`.""" + return native_adapter.group_by_suffixes( state_dict, suffixes, prefixes=KNOWN_PREFIXES, bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, @@ -75,7 +75,7 @@ def resolve_targets(prefix_used, base): return [] -# === Native loaders (thin wrappers over native_loader generics) === +# === Native loaders (thin wrappers over native_adapter generics) === _BIND_KWARGS = dict( @@ -87,24 +87,24 @@ _BIND_KWARGS = dict( def try_load_lora(name, network_on_disk, lora_scale): - return native_loader.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_lokr(name, network_on_disk, lora_scale): - return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_loha(name, network_on_disk, lora_scale): - return native_loader.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_oft(name, network_on_disk, lora_scale): - return native_loader.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load(name, network_on_disk, lora_scale): """Run every ERNIE family loader, merge any that match.""" - return native_loader.try_load_chain( + return native_adapter.try_load_chain( name, network_on_disk, lora_scale, family_loaders=(try_load_lora, try_load_lokr, try_load_loha, try_load_oft), ) diff --git a/pipelines/f_lite/__init__.py b/pipelines/f_lite/__init__.py index e4ddb20f5..5c4a23513 100644 --- a/pipelines/f_lite/__init__.py +++ b/pipelines/f_lite/__init__.py @@ -1,5 +1,9 @@ +from pipelines.native_transformer import TransformerSpec from .pipeline import FLitePipeline, FLitePipelineOutput, APGConfig from .model import DiT -__all__ = ["APGConfig", "DiT", "FLitePipeline", "FLitePipelineOutput"] +FLITE_SPEC = TransformerSpec(cls=DiT, subfolder='dit_model') + + +__all__ = ["APGConfig", "DiT", "FLITE_SPEC", "FLitePipeline", "FLitePipelineOutput"] diff --git a/pipelines/flux/flux2_lora.py b/pipelines/flux/flux2_lora.py index 81994cf82..1c69a66b4 100644 --- a/pipelines/flux/flux2_lora.py +++ b/pipelines/flux/flux2_lora.py @@ -16,12 +16,12 @@ produced by ``Flux2Transformer2DModel.save_lora_adapter()``). Diffusers-PEFT BFL/kohya keys are mapped to diffusers paths via ``F2_SINGLE_MAP`` / ``F2_DOUBLE_MAP`` / ``F2_QKV_MAP``. Fused QKV in double_blocks emits three -Q/K/V targets each carrying a :class:`modules.lora.native_loader.ChunkSpec` +Q/K/V targets each carrying a :class:`modules.lora.native_adapter.ChunkSpec` that the generic loaders use to chunk the up-weight or instantiate the appropriate ``NetworkModule*Chunk`` variant. Per-family fused-QKV handling is inherited from -:mod:`modules.lora.native_loader`; see the loader-by-loader notes there. +:mod:`modules.lora.native_adapter`; see the loader-by-loader notes there. LyCORIS algorithm coverage relative to upstream ``KohakuBlueleaf/LyCORIS/lycoris/modules/``: @@ -49,13 +49,13 @@ to inject the ``diffusion_model.`` prefix for bare-BFL keys and bake kohya import os from modules.logger import log -from modules.lora import native_loader -from modules.lora.native_loader import ChunkSpec +from modules.lora import native_adapter +from modules.lora.native_adapter import ChunkSpec # === Arch-specific prefix configuration === -KNOWN_PREFIXES = native_loader.KNOWN_PREFIXES_DEFAULT + ("lycoris_",) +KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT + ("lycoris_",) BARE_FLUX_PREFIXES = ( "single_blocks.", "double_blocks.", "img_in.", "txt_in.", @@ -108,34 +108,34 @@ KOHYA_SUFFIX_MAP = { # === Re-exports for backward compatibility === # The offline test suite addresses these via the flux2_lora module surface. -# Re-export rather than asking tests to import native_loader directly. +# Re-export rather than asking tests to import native_adapter directly. -LORA_SUFFIXES = native_loader.LORA_SUFFIXES -LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES -LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES -OFT_SUFFIXES = native_loader.OFT_SUFFIXES -IA3_SUFFIXES = native_loader.IA3_SUFFIXES -GLORA_SUFFIXES = native_loader.GLORA_SUFFIXES -NORM_SUFFIXES = native_loader.NORM_SUFFIXES -FULL_SUFFIXES = native_loader.FULL_SUFFIXES +LORA_SUFFIXES = native_adapter.LORA_SUFFIXES +LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES +LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES +OFT_SUFFIXES = native_adapter.OFT_SUFFIXES +IA3_SUFFIXES = native_adapter.IA3_SUFFIXES +GLORA_SUFFIXES = native_adapter.GLORA_SUFFIXES +NORM_SUFFIXES = native_adapter.NORM_SUFFIXES +FULL_SUFFIXES = native_adapter.FULL_SUFFIXES -LORA_MARKERS = native_loader.LORA_MARKERS -LOKR_MARKERS = native_loader.LOKR_MARKERS -LOHA_MARKERS = native_loader.LOHA_MARKERS -OFT_MARKERS = native_loader.OFT_MARKERS -IA3_MARKERS = native_loader.IA3_MARKERS -GLORA_MARKERS = native_loader.GLORA_MARKERS -NORM_MARKERS = native_loader.NORM_MARKERS -FULL_MARKERS = native_loader.FULL_MARKERS +LORA_MARKERS = native_adapter.LORA_MARKERS +LOKR_MARKERS = native_adapter.LOKR_MARKERS +LOHA_MARKERS = native_adapter.LOHA_MARKERS +OFT_MARKERS = native_adapter.OFT_MARKERS +IA3_MARKERS = native_adapter.IA3_MARKERS +GLORA_MARKERS = native_adapter.GLORA_MARKERS +NORM_MARKERS = native_adapter.NORM_MARKERS +FULL_MARKERS = native_adapter.FULL_MARKERS -SUFFIX_NORMALIZE = native_loader.SUFFIX_NORMALIZE -BARE_DIFFUSERS_PREFIX_USED = native_loader.BARE_DIFFUSERS_PREFIX_USED -has_marker = native_loader.has_marker +SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE +BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED +has_marker = native_adapter.has_marker def parse_key(key, suffixes): - """Flux2-bound :func:`native_loader.parse_key`. Returns ``(prefix_used, base, suffix)`` or ``None``.""" - return native_loader.parse_key( + """Flux2-bound :func:`native_adapter.parse_key`. Returns ``(prefix_used, base, suffix)`` or ``None``.""" + return native_adapter.parse_key( key, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, @@ -144,8 +144,8 @@ def parse_key(key, suffixes): def group_by_suffixes(state_dict, suffixes): - """Flux2-bound :func:`native_loader.group_by_suffixes`.""" - return native_loader.group_by_suffixes( + """Flux2-bound :func:`native_adapter.group_by_suffixes`.""" + return native_adapter.group_by_suffixes( state_dict, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, @@ -231,7 +231,7 @@ def _bfl_to_diffusers_targets(base): return targets -# === Native loaders (thin wrappers over native_loader generics) === +# === Native loaders (thin wrappers over native_adapter generics) === _BIND_KWARGS = dict( @@ -244,40 +244,40 @@ _BIND_KWARGS = dict( def try_load_lora(name, network_on_disk, lora_scale): - return native_loader.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_lokr(name, network_on_disk, lora_scale): - return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_loha(name, network_on_disk, lora_scale): - return native_loader.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_oft(name, network_on_disk, lora_scale): - return native_loader.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_ia3(name, network_on_disk, lora_scale): - return native_loader.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_glora(name, network_on_disk, lora_scale): - return native_loader.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_norm(name, network_on_disk, lora_scale): - return native_loader.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_full(name, network_on_disk, lora_scale): - return native_loader.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load(name, network_on_disk, lora_scale): """Single dispatcher entry point: run every family loader, merge any that match.""" - return native_loader.try_load_chain( + return native_adapter.try_load_chain( name, network_on_disk, lora_scale, family_loaders=( try_load_lora, try_load_lokr, try_load_loha, try_load_oft, diff --git a/pipelines/flux2_klein/__init__.py b/pipelines/flux2_klein/__init__.py new file mode 100644 index 000000000..72d898a92 --- /dev/null +++ b/pipelines/flux2_klein/__init__.py @@ -0,0 +1,26 @@ +"""Flux 2 Klein pipeline package. + +Exports :data:`FLUX2_KLEIN_SPEC`. Klein shares +:class:`Flux2Transformer2DModel` with full Flux 2 but uses a smaller +config (hidden_size and friends). diffusers' ``from_single_file`` picks +the class default (= Flux 2 full), so loading a Klein-shaped community +file crashes at ``load_model_dict_into_meta`` with a shape mismatch like +``expected (36864, 6144), got (24576, 4096)``. + +Routing through :mod:`pipelines.native_transformer` pulls the Klein +``transformer/config.json`` from the base repo first and instantiates +``Flux2Transformer2DModel`` at the right size, then runs the diffusers +Flux 2 converter to split fused QKV blocks and rename BFL keys into the +diffusers-expected names. +""" + +import diffusers +from diffusers.loaders.single_file_utils import convert_flux2_transformer_checkpoint_to_diffusers + +from pipelines.native_transformer import TransformerSpec + + +FLUX2_KLEIN_SPEC = TransformerSpec( + cls=diffusers.Flux2Transformer2DModel, + converter=convert_flux2_transformer_checkpoint_to_diffusers, +) diff --git a/pipelines/generic.py b/pipelines/generic.py index b6aa3e92b..8ba75214b 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -19,7 +19,14 @@ def _loader(component): return 'runai' if shared.opts.runai_streamer_transformers else 'default' -def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer", allow_quant=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None, **kwargs): +def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer", allow_quant=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None, native_spec=None, **kwargs): + """Load a DiT transformer from the base repo, or from a user-selected + single file when the UNET dropdown (``shared.opts.sd_unet``) is set. + + With ``native_spec`` set and a .safetensors override selected, dispatches + to :func:`pipelines.native_transformer.load`. Without a spec, a single-file + override falls back to ``from_single_file``. + """ if shared.state.interrupted: return None transformer = None @@ -55,6 +62,19 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer **load_args, ) transformer = model_quant.do_post_load_quant(transformer, allow=quant_type is not None) + elif local_file is not None and local_file.lower().endswith('.safetensors') and native_spec is not None: + from pipelines import native_transformer + log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader=native args={load_args}') + transformer, _ = native_transformer.load( + local_file, repo_id, native_spec, load_config, + allow_quant=allow_quant, + dtype=dtype, + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, + quant_args=quant_args, + quant_type=quant_type, + **kwargs, + ) elif local_file is not None and local_file.lower().endswith('.safetensors'): log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("diffusers")} args={load_args}') if dtype is not None: diff --git a/pipelines/glm/__init__.py b/pipelines/glm/__init__.py new file mode 100644 index 000000000..9de6a1cc6 --- /dev/null +++ b/pipelines/glm/__init__.py @@ -0,0 +1,15 @@ +"""GLM-Image pipeline package. + +Exports :data:`GLM_IMAGE_SPEC`. The minimum +``TransformerSpec(cls=GlmImageTransformer2DModel)`` works because +GLM-Image community files use BFL-style ``model.diffusion_model.`` +prefixed keys whose names match the diffusers state_dict verbatim after +prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +GLM_IMAGE_SPEC = TransformerSpec(cls=diffusers.GlmImageTransformer2DModel) diff --git a/pipelines/hunyuandit/__init__.py b/pipelines/hunyuandit/__init__.py new file mode 100644 index 000000000..da6fbe23d --- /dev/null +++ b/pipelines/hunyuandit/__init__.py @@ -0,0 +1,14 @@ +"""HunyuanDiT pipeline package. + +Exports :data:`HUNYUANDIT_SPEC`. The minimum +``TransformerSpec(cls=HunyuanDiT2DModel)`` works because HunyuanDiT +community files use BFL-style ``model.diffusion_model.``-prefixed keys +whose names match the diffusers state_dict verbatim after prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +HUNYUANDIT_SPEC = TransformerSpec(cls=diffusers.HunyuanDiT2DModel) diff --git a/pipelines/hyimage/__init__.py b/pipelines/hyimage/__init__.py new file mode 100644 index 000000000..42cc55c6e --- /dev/null +++ b/pipelines/hyimage/__init__.py @@ -0,0 +1,19 @@ +"""HunyuanImage pipeline package. + +Exports :data:`HUNYUANIMAGE_SPEC`. The minimum +``TransformerSpec(cls=HunyuanImageTransformer2DModel)`` works because +HunyuanImage 2.1 community files use BFL-style +``model.diffusion_model.``-prefixed keys whose names match the diffusers +state_dict verbatim after prefix strip. + +The HunyuanImage 3 path is a transformers ``AutoModelForCausalLM`` and +does not go through the native transformer loader, so it does not need a +spec here. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +HUNYUANIMAGE_SPEC = TransformerSpec(cls=diffusers.HunyuanImageTransformer2DModel) diff --git a/pipelines/joy/__init__.py b/pipelines/joy/__init__.py new file mode 100644 index 000000000..1ee025cbe --- /dev/null +++ b/pipelines/joy/__init__.py @@ -0,0 +1,15 @@ +"""Joy-Image-Edit pipeline package. + +Exports :data:`JOY_SPEC`. The minimum +``TransformerSpec(cls=JoyImageEditTransformer3DModel)`` works because +Joy community files use BFL-style ``model.diffusion_model.``-prefixed +keys whose names match the diffusers state_dict verbatim after prefix +strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +JOY_SPEC = TransformerSpec(cls=diffusers.JoyImageEditTransformer3DModel) diff --git a/pipelines/kandinsky3/__init__.py b/pipelines/kandinsky3/__init__.py new file mode 100644 index 000000000..da46fa862 --- /dev/null +++ b/pipelines/kandinsky3/__init__.py @@ -0,0 +1,13 @@ +"""Kandinsky 3.0 pipeline package. + +Exports :data:`KANDINSKY3_UNET_SPEC`. Kandinsky 3 uses +:class:`Kandinsky3UNet` in the ``unet`` subfolder of the repo, hence +``subfolder='unet'`` instead of the default ``'transformer'``. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +KANDINSKY3_UNET_SPEC = TransformerSpec(cls=diffusers.Kandinsky3UNet, subfolder='unet') diff --git a/pipelines/kandinsky5/__init__.py b/pipelines/kandinsky5/__init__.py new file mode 100644 index 000000000..a2927f1c9 --- /dev/null +++ b/pipelines/kandinsky5/__init__.py @@ -0,0 +1,12 @@ +"""Kandinsky 5.0 pipeline package. + +Exports :data:`KANDINSKY5_SPEC`. Kandinsky 5 uses +:class:`Kandinsky5Transformer3DModel` and follows the default layout. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +KANDINSKY5_SPEC = TransformerSpec(cls=diffusers.Kandinsky5Transformer3DModel) diff --git a/pipelines/lens/__init__.py b/pipelines/lens/__init__.py index bcfca204f..cae35a683 100644 --- a/pipelines/lens/__init__.py +++ b/pipelines/lens/__init__.py @@ -2,6 +2,7 @@ import diffusers as _diffusers import transformers as _transformers +from pipelines.native_transformer import TransformerSpec from .pipeline import LensPipeline, LensPipelineOutput from .pipeline_image import LensImg2ImgPipeline, LensInpaintPipeline from .reasoner import PromptReasoner @@ -9,6 +10,9 @@ from .resolution import RESOLUTION_BUCKETS, resolve_resolution from .text_encoder import LensGptOssEncoder from .transformer import LensTransformer2DModel + +LENS_SPEC = TransformerSpec(cls=LensTransformer2DModel) + # --------------------------------------------------------------------------- # Make our custom subclasses discoverable by ``diffusers.DiffusionPipeline``. # @@ -38,6 +42,7 @@ if not hasattr(_diffusers, "LensInpaintPipeline"): del _diffusers, _transformers __all__ = [ + "LENS_SPEC", "LensPipeline", "LensPipelineOutput", "LensImg2ImgPipeline", diff --git a/pipelines/longcat/__init__.py b/pipelines/longcat/__init__.py new file mode 100644 index 000000000..f68f49c9c --- /dev/null +++ b/pipelines/longcat/__init__.py @@ -0,0 +1,15 @@ +"""LongCat pipeline package. + +Exports :data:`LONGCAT_SPEC`. The minimum +``TransformerSpec(cls=LongCatImageTransformer2DModel)`` works because +LongCat community files use BFL-style ``model.diffusion_model.`` +prefixed keys whose names match the diffusers state_dict verbatim after +prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +LONGCAT_SPEC = TransformerSpec(cls=diffusers.LongCatImageTransformer2DModel) diff --git a/pipelines/model_anima.py b/pipelines/model_anima.py index 0e9908eed..4ed2228fb 100644 --- a/pipelines/model_anima.py +++ b/pipelines/model_anima.py @@ -16,37 +16,25 @@ def _import_from_file(module_name, file_path): return mod -def resolve_custom_transformer_path(): - """Return an absolute path if the user selected a transformer in the UNET - dropdown and the file is resolvable, else ``None``. - """ - sel = shared.opts.sd_unet - if sel is None or sel in ('Default', 'None'): - return None - from modules import sd_unet - if sel not in list(sd_unet.unet_dict): - log.error(f'Load module: type=transformer file="{sel}" not found') - return None - path = sd_unet.unet_dict[sel] - if not os.path.exists(path): - log.error(f'Load module: type=transformer path="{path}" does not exist') - return None - return path - - def load_transformer_components(repo_id, diffusers_load_config, adapter_cls): """Load (transformer, llm_adapter_or_none). - If the UNET dropdown points at a valid safetensors, route through the - custom-transformer helper, which also extracts the bundled adapter - weights. Otherwise fall back to ``generic.load_transformer`` and return - ``None`` for the adapter so the caller loads it from the base repo. + If the UNET dropdown points at a valid safetensors, route through + :mod:`pipelines.native_transformer` with :data:`pipelines.anima.ANIMA_SPEC`, + which extracts any bundled ``llm_adapter`` weights inline with the + transformer. Otherwise fall back to :func:`generic.load_transformer` and + return ``None`` for the adapter so the caller loads it from the base repo. """ - local_file = resolve_custom_transformer_path() + from pipelines import native_transformer + local_file = native_transformer.resolve_path() if local_file is not None: - from pipelines.anima import anima_transformer + from pipelines.anima import ANIMA_SPEC try: - return anima_transformer.load_custom_transformer(repo_id, local_file, diffusers_load_config, adapter_cls) + transformer, siblings = native_transformer.load( + local_file, repo_id, ANIMA_SPEC, diffusers_load_config, + sibling_classes={'llm_adapter': adapter_cls}, + ) + return transformer, siblings.get('llm_adapter') except Exception as e: log.error(f'Load model: type=Anima custom transformer="{local_file}": {e}') errors.display(e, 'Load') diff --git a/pipelines/model_bria.py b/pipelines/model_bria.py index ae9c9fa7a..961f9c26f 100644 --- a/pipelines/model_bria.py +++ b/pipelines/model_bria.py @@ -20,11 +20,13 @@ def load_bria(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=BriaFibo repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.bria import BRIA_FIBO_SPEC transformer = generic.load_transformer( repo_id, cls_name=diffusers.BriaFiboTransformer2DModel, load_config=diffusers_load_config, allow_quant=False, + native_spec=BRIA_FIBO_SPEC, ) text_encoder = generic.load_text_encoder( repo_id, @@ -66,7 +68,8 @@ def load_bria(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=Bria repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=BriaTransformer2DModel, load_config=diffusers_load_config) + from pipelines.bria import BRIA_SPEC + transformer = generic.load_transformer(repo_id, cls_name=BriaTransformer2DModel, load_config=diffusers_load_config, native_spec=BRIA_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config) pipe = BriaPipeline.from_pretrained( diff --git a/pipelines/model_chroma.py b/pipelines/model_chroma.py index 1e7f9ea8a..ee33d2109 100644 --- a/pipelines/model_chroma.py +++ b/pipelines/model_chroma.py @@ -14,7 +14,8 @@ def load_chroma(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=Chroma repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.ChromaTransformer2DModel, load_config=diffusers_load_config, modules_to_not_convert=["distilled_guidance_layer"]) + from pipelines.chroma import CHROMA_SPEC + transformer = generic.load_transformer(repo_id, cls_name=diffusers.ChromaTransformer2DModel, load_config=diffusers_load_config, modules_to_not_convert=["distilled_guidance_layer"], native_spec=CHROMA_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config) pipe = diffusers.ChromaPipeline.from_pretrained( diff --git a/pipelines/model_chrono.py b/pipelines/model_chrono.py index 502e7d89d..be66af407 100644 --- a/pipelines/model_chrono.py +++ b/pipelines/model_chrono.py @@ -21,7 +21,8 @@ def load_chrono(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=ChronoEdit repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.ChronoEditTransformer3DModel, load_config=diffusers_load_config, subfolder="transformer") + from pipelines.chrono import CHRONOEDIT_SPEC + transformer = generic.load_transformer(repo_id, cls_name=diffusers.ChronoEditTransformer3DModel, load_config=diffusers_load_config, subfolder="transformer", native_spec=CHRONOEDIT_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.UMT5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder") try: diff --git a/pipelines/model_cogview.py b/pipelines/model_cogview.py index 673c41565..a48489e1f 100644 --- a/pipelines/model_cogview.py +++ b/pipelines/model_cogview.py @@ -14,7 +14,8 @@ def load_cogview3(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=CogView3 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.CogView3PlusTransformer2DModel, load_config=diffusers_load_config, subfolder="transformer") + from pipelines.cogview import COGVIEW3_SPEC + transformer = generic.load_transformer(repo_id, cls_name=diffusers.CogView3PlusTransformer2DModel, load_config=diffusers_load_config, subfolder="transformer", native_spec=COGVIEW3_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder") pipe = diffusers.CogView3PlusPipeline.from_pretrained( @@ -40,7 +41,8 @@ def load_cogview4(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=CogView4 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.CogView4Transformer2DModel, load_config=diffusers_load_config, subfolder="transformer") + from pipelines.cogview import COGVIEW4_SPEC + transformer = generic.load_transformer(repo_id, cls_name=diffusers.CogView4Transformer2DModel, load_config=diffusers_load_config, subfolder="transformer", native_spec=COGVIEW4_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.GlmModel, load_config=diffusers_load_config, subfolder="text_encoder", allow_quant=True) pipe = diffusers.CogView4Pipeline.from_pretrained( diff --git a/pipelines/model_ernie.py b/pipelines/model_ernie.py index 382b1fe80..471c081be 100644 --- a/pipelines/model_ernie.py +++ b/pipelines/model_ernie.py @@ -14,10 +14,12 @@ def load_ernie_image(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=ERNIE-Image repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args} pe={shared.opts.model_ernie_enable_pe}') + from pipelines.ernie import ERNIE_SPEC transformer = generic.load_transformer( repo_id, cls_name=diffusers.ErnieImageTransformer2DModel, load_config=diffusers_load_config, + native_spec=ERNIE_SPEC, ) text_encoder = generic.load_text_encoder( repo_id, diff --git a/pipelines/model_flite.py b/pipelines/model_flite.py index 52e7f2f81..aa63969c3 100644 --- a/pipelines/model_flite.py +++ b/pipelines/model_flite.py @@ -19,7 +19,7 @@ def load_flite(checkpoint_info, diffusers_load_config=None): diffusers.FLitePipeline = f_lite.FLitePipeline sys.modules['f_lite'] = f_lite - dit_model = generic.load_transformer(repo_id, cls_name=f_lite.DiT, load_config=diffusers_load_config, subfolder="dit_model") + dit_model = generic.load_transformer(repo_id, cls_name=f_lite.DiT, load_config=diffusers_load_config, subfolder="dit_model", native_spec=f_lite.FLITE_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder") pipe = f_lite.FLitePipeline.from_pretrained( diff --git a/pipelines/model_flux2_klein.py b/pipelines/model_flux2_klein.py index 3a238805b..00c666df6 100644 --- a/pipelines/model_flux2_klein.py +++ b/pipelines/model_flux2_klein.py @@ -15,7 +15,8 @@ def load_flux2_klein(checkpoint_info, diffusers_load_config=None): log.debug(f'Load model: type=Flux2Klein repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') # Load transformer - Klein uses Flux2Transformer2DModel (same class as Flux2, different size) - transformer = generic.load_transformer(repo_id, cls_name=diffusers.Flux2Transformer2DModel, load_config=diffusers_load_config) + from pipelines.flux2_klein import FLUX2_KLEIN_SPEC + transformer = generic.load_transformer(repo_id, cls_name=diffusers.Flux2Transformer2DModel, load_config=diffusers_load_config, native_spec=FLUX2_KLEIN_SPEC) # Load text encoder - Klein uses Qwen3 (4B for Klein-4B, 8B for Klein-9B) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3ForCausalLM, load_config=diffusers_load_config) diff --git a/pipelines/model_glm.py b/pipelines/model_glm.py index c58358f88..c46b0ae63 100644 --- a/pipelines/model_glm.py +++ b/pipelines/model_glm.py @@ -95,10 +95,12 @@ def load_glm_image(checkpoint_info, diffusers_load_config=None): log.debug(f'Load model: type=GLM-Image repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') # Load transformer (DiT decoder - 7B) with quantization support + from pipelines.glm import GLM_IMAGE_SPEC transformer = generic.load_transformer( repo_id, cls_name=diffusers.GlmImageTransformer2DModel, - load_config=diffusers_load_config + load_config=diffusers_load_config, + native_spec=GLM_IMAGE_SPEC, ) # Load text encoder (ByT5 for glyph) - cannot use shared T5 as GLM-Image requires specific ByT5 encoder (1472 hidden size) diff --git a/pipelines/model_hunyuandit.py b/pipelines/model_hunyuandit.py index b2933a94f..1caa06f16 100644 --- a/pipelines/model_hunyuandit.py +++ b/pipelines/model_hunyuandit.py @@ -19,7 +19,8 @@ def load_hunyuandit(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=HunyuanDiT repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.HunyuanDiT2DModel, load_config=diffusers_load_config) + from pipelines.hunyuandit import HUNYUANDIT_SPEC + transformer = generic.load_transformer(repo_id, cls_name=diffusers.HunyuanDiT2DModel, load_config=diffusers_load_config, native_spec=HUNYUANDIT_SPEC) repo_te = 'Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers' if 'HunyuanDiT-v1' in repo_id else repo_id text_encoder_2 = generic.load_text_encoder(repo_te, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder_2", allow_shared=False) # this is not normal t5 diff --git a/pipelines/model_hyimage.py b/pipelines/model_hyimage.py index 4cf9776ee..83b1f9c1c 100644 --- a/pipelines/model_hyimage.py +++ b/pipelines/model_hyimage.py @@ -16,7 +16,8 @@ def load_hyimage(checkpoint_info, diffusers_load_config=None): # pylint: disable load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=HunyuanImage21 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.HunyuanImageTransformer2DModel, load_config=diffusers_load_config, subfolder="transformer") + from pipelines.hyimage import HUNYUANIMAGE_SPEC + transformer = generic.load_transformer(repo_id, cls_name=diffusers.HunyuanImageTransformer2DModel, load_config=diffusers_load_config, subfolder="transformer", native_spec=HUNYUANIMAGE_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config, subfolder="text_encoder") text_encoder_2 = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder_2", allow_shared=False) diff --git a/pipelines/model_joy.py b/pipelines/model_joy.py index 5ed363315..e0c25f6c7 100644 --- a/pipelines/model_joy.py +++ b/pipelines/model_joy.py @@ -14,10 +14,12 @@ def load_joy(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=JoyImageEdit repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.joy import JOY_SPEC transformer = generic.load_transformer( repo_id, cls_name=diffusers.JoyImageEditTransformer3DModel, load_config=diffusers_load_config, + native_spec=JOY_SPEC, ) text_encoder = generic.load_text_encoder( repo_id, diff --git a/pipelines/model_kandinsky.py b/pipelines/model_kandinsky.py index d8e950177..3fec7f510 100644 --- a/pipelines/model_kandinsky.py +++ b/pipelines/model_kandinsky.py @@ -50,7 +50,8 @@ def load_kandinsky3(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=Kandinsky30 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - unet = generic.load_transformer(repo_id, cls_name=diffusers.Kandinsky3UNet, load_config=diffusers_load_config, subfolder="unet", variant="fp16") + from pipelines.kandinsky3 import KANDINSKY3_UNET_SPEC + unet = generic.load_transformer(repo_id, cls_name=diffusers.Kandinsky3UNet, load_config=diffusers_load_config, subfolder="unet", variant="fp16", native_spec=KANDINSKY3_UNET_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder", variant="fp16", allow_shared=False) pipe = diffusers.Kandinsky3Pipeline.from_pretrained( @@ -83,7 +84,8 @@ def load_kandinsky5(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=Kandinsky50 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.Kandinsky5Transformer3DModel, load_config=diffusers_load_config) + from pipelines.kandinsky5 import KANDINSKY5_SPEC + transformer = generic.load_transformer(repo_id, cls_name=diffusers.Kandinsky5Transformer3DModel, load_config=diffusers_load_config, native_spec=KANDINSKY5_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config) if 'I2I' in repo_id: diff --git a/pipelines/model_lens.py b/pipelines/model_lens.py index a4d07a4d9..f6e64abfb 100644 --- a/pipelines/model_lens.py +++ b/pipelines/model_lens.py @@ -15,7 +15,7 @@ def load_lens(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=Lens repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} reasoner={shared.opts.model_lens_enable_pe} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=lens.LensTransformer2DModel, load_config=diffusers_load_config) + transformer = generic.load_transformer(repo_id, cls_name=lens.LensTransformer2DModel, load_config=diffusers_load_config, native_spec=lens.LENS_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=lens.LensGptOssEncoder, load_config=diffusers_load_config, allow_quant=False) pipe = lens.LensPipeline.from_pretrained( diff --git a/pipelines/model_longcat.py b/pipelines/model_longcat.py index ac8a27322..397d9c0de 100644 --- a/pipelines/model_longcat.py +++ b/pipelines/model_longcat.py @@ -14,7 +14,8 @@ def load_longcat(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=LongCat repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.LongCatImageTransformer2DModel, load_config=diffusers_load_config) + from pipelines.longcat import LONGCAT_SPEC + transformer = generic.load_transformer(repo_id, cls_name=diffusers.LongCatImageTransformer2DModel, load_config=diffusers_load_config, native_spec=LONGCAT_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config) text_processor = transformers.Qwen2VLProcessor.from_pretrained(repo_id, subfolder='tokenizer', cache_dir=shared.opts.hfcache_dir) diff --git a/pipelines/model_nucleus.py b/pipelines/model_nucleus.py index b7087fa59..e465f2216 100644 --- a/pipelines/model_nucleus.py +++ b/pipelines/model_nucleus.py @@ -14,10 +14,12 @@ def load_nucleus(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=NucleusMoEImage repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.nucleus import NUCLEUS_SPEC transformer = generic.load_transformer( repo_id, cls_name=diffusers.NucleusMoEImageTransformer2DModel, load_config=diffusers_load_config, + native_spec=NUCLEUS_SPEC, ) text_encoder = generic.load_text_encoder( repo_id, diff --git a/pipelines/model_ovis.py b/pipelines/model_ovis.py index e4e65270f..753237610 100644 --- a/pipelines/model_ovis.py +++ b/pipelines/model_ovis.py @@ -14,7 +14,8 @@ def load_ovis(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=OvisImage repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.OvisImageTransformer2DModel, load_config=diffusers_load_config) + from pipelines.ovis import OVIS_SPEC + transformer = generic.load_transformer(repo_id, cls_name=diffusers.OvisImageTransformer2DModel, load_config=diffusers_load_config, native_spec=OVIS_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3Model, load_config=diffusers_load_config) pipe = diffusers.OvisImagePipeline.from_pretrained( diff --git a/pipelines/model_pixart.py b/pipelines/model_pixart.py index 1ced8b659..0d2ff1127 100644 --- a/pipelines/model_pixart.py +++ b/pipelines/model_pixart.py @@ -24,7 +24,8 @@ def load_pixart(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=PixArtSigma repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.PixArtTransformer2DModel, load_config=diffusers_load_config) + from pipelines.pixart import PIXART_SPEC + transformer = generic.load_transformer(repo_id, cls_name=diffusers.PixArtTransformer2DModel, load_config=diffusers_load_config, native_spec=PIXART_SPEC) text_encoder = generic.load_text_encoder(repo_id_tenc, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config) pipe = diffusers.PixArtSigmaPipeline.from_pretrained( diff --git a/pipelines/model_prx.py b/pipelines/model_prx.py index 52088444c..4c886f231 100644 --- a/pipelines/model_prx.py +++ b/pipelines/model_prx.py @@ -14,7 +14,8 @@ def load_prx(checkpoint_info, diffusers_load_config=None): log.debug(f'Load model: type=PRX repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder - transformer = generic.load_transformer(repo_id, cls_name=diffusers.PRXTransformer2DModel, load_config=diffusers_load_config) + from pipelines.prx import PRX_SPEC + transformer = generic.load_transformer(repo_id, cls_name=diffusers.PRXTransformer2DModel, load_config=diffusers_load_config, native_spec=PRX_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=T5GemmaEncoder, load_config=diffusers_load_config) pipe = diffusers.PRXPipeline.from_pretrained( diff --git a/pipelines/model_qwen.py b/pipelines/model_qwen.py index 52c9fd251..094d23ed3 100644 --- a/pipelines/model_qwen.py +++ b/pipelines/model_qwen.py @@ -57,12 +57,14 @@ def load_qwen(checkpoint_info, diffusers_load_config=None): transformer_subfolder = "transformer" if transformer is None: + from pipelines.qwen import QWEN_SPEC transformer = generic.load_transformer( repo_transformer, subfolder=transformer_subfolder, cls_name=diffusers.QwenImageTransformer2DModel, load_config=diffusers_load_config, modules_to_not_convert=["transformer_blocks.0.img_mod.1.weight"], + native_spec=QWEN_SPEC, ) repo_te = 'Qwen/Qwen-Image' diff --git a/pipelines/model_step1x_edit.py b/pipelines/model_step1x_edit.py index aba2bd5ab..47d0c8adb 100644 --- a/pipelines/model_step1x_edit.py +++ b/pipelines/model_step1x_edit.py @@ -22,7 +22,8 @@ def load_step1x_edit(checkpoint_info, diffusers_load_config=None): text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config) processor = transformers.Qwen2_5_VLProcessor.from_pretrained(repo_id, cache_dir=shared.opts.hfcache_dir, subfolder='processor') - transformer = generic.load_transformer(repo_id, cls_name=Step1XEditTransformer2DModel, load_config=diffusers_load_config) + from pipelines.step1x import STEP1X_SPEC + transformer = generic.load_transformer(repo_id, cls_name=Step1XEditTransformer2DModel, load_config=diffusers_load_config, native_spec=STEP1X_SPEC) pipe = Step1XEditPipeline.from_pretrained( repo_id, diff --git a/pipelines/model_vibe.py b/pipelines/model_vibe.py index ecac08867..999ab9dd7 100644 --- a/pipelines/model_vibe.py +++ b/pipelines/model_vibe.py @@ -21,11 +21,13 @@ def load_vibe(checkpoint_info, diffusers_load_config=None): sys.modules['vibe.transformer.vibe_sana_editing'] = diffusers # monkey patch since hf model_index.json points to custom class path + from pipelines.vibe import VIBE_SPEC transformer = generic.load_transformer( repo_id, cls_name=VIBESanaEditingModel, load_config=diffusers_load_config, allow_quant=False, + native_spec=VIBE_SPEC, ) text_encoder = generic.load_text_encoder( repo_id, diff --git a/pipelines/model_wanai.py b/pipelines/model_wanai.py index f7a4413a8..e62fb165d 100644 --- a/pipelines/model_wanai.py +++ b/pipelines/model_wanai.py @@ -1,51 +1,8 @@ -import os import transformers import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log - - -def load_transformer(repo_id, diffusers_load_config=None, subfolder='transformer'): - if diffusers_load_config is None: - diffusers_load_config = {} - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True) - fn = None - - if 'VACE' in repo_id: - transformer_cls = diffusers.WanVACETransformer3DModel - else: - transformer_cls = diffusers.WanTransformer3DModel - - if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default': - from modules import sd_unet - if shared.opts.sd_unet not in list(sd_unet.unet_dict): - log.error(f'Load module: type=Transformer not found: {shared.opts.sd_unet}') - return None - fn = sd_unet.unet_dict[shared.opts.sd_unet] if os.path.exists(sd_unet.unet_dict[shared.opts.sd_unet]) else None - - if fn is not None and 'gguf' in fn.lower(): - log.error('Load model: type=WanAI format="gguf" unsupported') - transformer = None - elif fn is not None and 'safetensors' in fn.lower(): - log.debug(f'Load model: type=WanAI {subfolder}="{fn}" quant="{model_quant.get_quant(repo_id)}" args={load_args}') - transformer = transformer_cls.from_single_file( - fn, - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - else: - log.debug(f'Load model: type=WanAI {subfolder}="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - transformer = transformer_cls.from_pretrained( - repo_id, - subfolder=subfolder, - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - if shared.opts.diffusers_offload_mode != 'none' and transformer is not None: - sd_models.move_model(transformer, devices.cpu) - return transformer +from pipelines import generic def load_text_encoder(repo_id, diffusers_load_config=None): @@ -71,26 +28,27 @@ def load_wan(checkpoint_info, diffusers_load_config=None): diffusers_load_config = {} repo_id = sd_models.path_to_repo(checkpoint_info) sd_models.hf_auth_check(checkpoint_info) + transformer_cls = diffusers.WanVACETransformer3DModel if 'VACE' in repo_id else diffusers.WanTransformer3DModel boundary_ratio = None if 'a14b' in repo_id.lower() or 'fun-14b' in repo_id.lower(): if shared.opts.model_wan_stage == 'high noise' or shared.opts.model_wan_stage == 'first': - transformer = load_transformer(repo_id, diffusers_load_config, 'transformer') + transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer') transformer_2 = None boundary_ratio = 0.0 elif shared.opts.model_wan_stage == 'low noise' or shared.opts.model_wan_stage == 'second': transformer = None - transformer_2 = load_transformer(repo_id, diffusers_load_config, 'transformer_2') + transformer_2 = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer_2') boundary_ratio = 1000.0 elif shared.opts.model_wan_stage == 'combined' or shared.opts.model_wan_stage == 'both': - transformer = load_transformer(repo_id, diffusers_load_config, 'transformer') - transformer_2 = load_transformer(repo_id, diffusers_load_config, 'transformer_2') + transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer') + transformer_2 = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer_2') boundary_ratio = shared.opts.model_wan_boundary else: log.error(f'Load model: type=WanAI stage="{shared.opts.model_wan_stage}" unsupported') return None else: - transformer = load_transformer(repo_id, diffusers_load_config, 'transformer') + transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer') transformer_2 = None text_encoder = load_text_encoder(repo_id, diffusers_load_config) diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py new file mode 100644 index 000000000..b52326e5a --- /dev/null +++ b/pipelines/native_transformer.py @@ -0,0 +1,624 @@ +"""Generic native loader for DiT transformers and bundled sibling components. + +Loads a single-safetensors file into a diffusers (or custom) transformer class +when the user selects an override via the UNET dropdown (``shared.opts.sd_unet``). +Bypasses :func:`diffusers.loaders.FromOriginalModelMixin.from_single_file` so +sdnext owns prefix detection, optional sibling partitioning, dtype/quant/offload +handling, and explicit validation of missing/unexpected keys. + +The per-arch knobs are captured in :class:`TransformerSpec`. Each pipeline +defines its spec in ``pipelines//__init__.py`` and passes it explicitly +to :func:`load` (or to :func:`pipelines.generic.load_transformer` via the +``native_spec`` kwarg). No class-keyed registry: two pipelines may share a +transformer class but need different specs (e.g. Anima vs raw Cosmos both +use ``CosmosTransformer3DModel`` but Anima has a bundled ``llm_adapter`` +sibling). Pipelines without a custom spec fall back to +:func:`make_default_spec`, which opportunistically picks up a real converter +from diffusers' ``SINGLE_FILE_LOADABLE_CLASSES`` table. + +Algorithm: + +1. Read the safetensors state dict (.gguf and .pth are rejected up front). +2. Detect and strip one of the spec's known prefixes (raises on mixed prefixes). +3. Check forbidden markers (catches structural mismatches like Cosmos 1.0 keys + in a Cosmos 2.0 loader). +4. Partition off sibling component keys (e.g. Anima's bundled ``llm_adapter.*``). +5. Run the spec's converter if present (else pass through unchanged). +6. Fetch ``/config.json`` from the base repo, instantiate via + ``cls.from_config``, ``load_state_dict(strict=False)``, validate, dtype-cast, + quantize, and offload-place. +7. Repeat the build for each populated sibling (no converter, no quant by + default; sibling weights are read raw from the bundled file). + +Returns ``(transformer, sibling_components_dict)``. The dict is empty for +arches with no siblings; for Anima it carries the ``llm_adapter`` if the +community file bundled one, else the empty dict and the caller falls back to +loading the adapter from the base repo. +""" + +import os +import time +from dataclasses import dataclass, field +from typing import Callable + +import huggingface_hub as hf +import torch + +from modules import shared, devices, sd_models, model_quant, errors +from modules.logger import log, console + + +DEFAULT_PREFIXES: tuple[str, ...] = ( + "model.diffusion_model.", + "diffusion_model.", + "net.", +) +DEFAULT_ACCEPTABLE_MISSING: tuple[str, ...] = ( + "rope.", + "pos_embedder.", + "learnable_pos_embed.", +) + + +@dataclass(frozen=True) +class SiblingSpec: + """Describes a non-transformer component that may ship inline in the same + safetensors as the transformer (e.g. Anima's ``llm_adapter``). + + ``subfolder`` names the base repo subfolder holding the canonical config and + weights when the sibling is NOT bundled inline; ``inline_prefix`` is the + key prefix that identifies the sibling's weights within the bundled file + (after the transformer's prefix has already been stripped). + """ + + subfolder: str + inline_prefix: str + acceptable_missing: tuple[str, ...] = () + + +@dataclass(frozen=True) +class TransformerSpec: + """Per-arch configuration for the native loader. + + Most arches only need to override ``cls`` (and rely on the default + prefixes, no converter, no siblings). Arches with bundled-sibling + components (Anima) or unusual key conventions (custom converters, + Cosmos-style structural markers) override the relevant fields. + """ + + cls: type + subfolder: str = "transformer" + prefixes: tuple[str, ...] = DEFAULT_PREFIXES + converter: Callable[[dict], dict] | None = None + siblings: dict[str, SiblingSpec] = field(default_factory=dict) + acceptable_missing: tuple[str, ...] = DEFAULT_ACCEPTABLE_MISSING + forbidden_markers: tuple[tuple[str, str], ...] = () + + +def make_default_spec(cls: type) -> TransformerSpec: + """Synthesize a default spec for ``cls``: default prefixes, no siblings, + no forbidden markers, and a converter picked up automatically from + diffusers' ``SINGLE_FILE_LOADABLE_CLASSES`` table if one exists (and is + not the no-op identity lambda that ``QwenImageTransformer2DModel`` and + a few other classes register). + + Used by callers (notably :func:`pipelines.generic.load_transformer`) when + a pipeline does not supply a custom ``TransformerSpec`` of its own. + """ + return TransformerSpec(cls=cls, converter=auto_pickup_converter(cls)) + + +def auto_pickup_converter(cls: type) -> Callable[[dict], dict] | None: + """Pull a checkpoint converter from diffusers for ``cls`` when one exists. + + Skipped for the no-op identity lambda some classes register + (notably ``QwenImageTransformer2DModel``), because using it would silently + accept whatever key naming the file happens to have. + """ + try: + from diffusers.loaders.single_file_model import SINGLE_FILE_LOADABLE_CLASSES + except ImportError: + return None + entry = SINGLE_FILE_LOADABLE_CLASSES.get(cls.__name__) + if entry is None: + return None + fn = entry.get("checkpoint_mapping_fn") + if fn is None or is_noop_converter(fn): + return None + return fn + + +def is_noop_converter(fn: Callable) -> bool: + """Detect ``lambda checkpoint, **kwargs: checkpoint`` and equivalents. + + Strips inline ``#`` comments from the source line before inspecting the + body, so that diagnostic markers like ``# noqa`` on the lambda's source + line do not defeat the detection. + """ + try: + import inspect + src = inspect.getsource(fn).strip() + except (OSError, TypeError): + return False + if "lambda" not in src: + return False + if "#" in src: + src = src.split("#", 1)[0].rstrip() + body = src.split(":", 1)[-1].strip().rstrip(",").rstrip(")") + return body.endswith("checkpoint") + + +def resolve_path() -> str | None: + """Return the absolute path of the UNET dropdown selection, or None if + no selection is active or the file is unresolvable. + """ + sel = shared.opts.sd_unet + if sel is None or sel in ("Default", "None"): + return None + from modules import sd_unet + if sel not in list(sd_unet.unet_dict): + log.error(f'Load module: type=transformer file="{sel}" not found') + return None + path = sd_unet.unet_dict[sel] + if not os.path.exists(path): + log.error(f'Load module: type=transformer path="{path}" does not exist') + return None + return path + + +def load( + local_file: str, + repo_id: str, + spec: TransformerSpec, + diffusers_cfg: dict | None = None, + sibling_classes: dict[str, type] | None = None, + *, + allow_quant: bool = True, + dtype=None, + modules_to_not_convert: list | None = None, + modules_dtype_dict: dict | None = None, + quant_args: dict | None = None, + quant_type: str | None = None, + **kwargs, +) -> tuple[object, dict[str, object]]: + """Load the transformer (and any bundled siblings) from ``local_file``. + + ``sibling_classes`` supplies the runtime class for each sibling named in + ``spec.siblings``. Required when a sibling has a dynamic class (e.g. + Anima's ``AnimaLLMAdapter`` is loaded from remote_code at runtime). + Missing sibling classes raise ``ValueError`` if the corresponding sibling + keys are present in the bundled file. + + Keyword-only arguments ``allow_quant``, ``dtype``, ``modules_to_not_convert``, + and ``modules_dtype_dict`` mirror the corresponding kwargs of + :func:`pipelines.generic.load_transformer` and are forwarded unchanged. + ``quant_args`` and ``quant_type`` are precomputed by the caller; when + ``None`` they are derived here via ``model_quant.get_dit_args``. Extra + ``**kwargs`` reach the transformer's ``cls.from_config``; siblings do not + receive them. + + Returns ``(transformer, siblings_dict)``. ``siblings_dict`` is keyed by + sibling name and is empty for non-sibling specs, or for sibling specs + whose keys are absent from the bundled file. + """ + if diffusers_cfg is None: + diffusers_cfg = {} + if sibling_classes is None: + sibling_classes = {} + + t0 = time.time() + if not local_file.lower().endswith(".safetensors"): + raise ValueError( + f"Load model: type={spec.cls.__name__} custom transformer requires .safetensors, " + f'got "{local_file}"' + ) + + if quant_args is None: + _, quant_args = model_quant.get_dit_args( + diffusers_cfg, module="Model", device_map=True, + allow_quant=allow_quant, + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, + ) + quant_type = model_quant.get_quant_type(quant_args) + + log.debug(f'Load model: native_transformer reading state_dict cls={spec.cls.__name__} file="{os.path.basename(local_file)}"') + state_dict = sd_models.read_state_dict(local_file, what="transformer") + state_dict = strip_prefix(state_dict, spec.prefixes, spec.cls.__name__) + check_forbidden_markers(state_dict, spec.forbidden_markers, spec.cls.__name__, local_file) + transformer_sd, sibling_sds = partition_siblings(state_dict, spec.siblings) + del state_dict + + sibling_counts = {name: len(sd) for name, sd in sibling_sds.items() if sd} + log.info( + f'Load model: type={spec.cls.__name__} custom="{os.path.basename(local_file)}" ' + f"transformer_keys={len(transformer_sd)} siblings={sibling_counts or '{}'}" + ) + + effective_dtype = dtype if dtype is not None else devices.dtype + transformer_cfg = fetch_component_config(repo_id, spec.subfolder) + transformer = build_component( + component_name="transformer", + state_dict=transformer_sd, + config=transformer_cfg, + cls=spec.cls, + converter=spec.converter, + acceptable_missing=spec.acceptable_missing, + quant_args=quant_args, + quant_type=quant_type, + dtype=effective_dtype, + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, + **kwargs, + ) + del transformer_sd + devices.torch_gc() + + loaded_siblings: dict[str, object] = {} + for name, sibling_sd in sibling_sds.items(): + if not sibling_sd: + continue + sibling_spec = spec.siblings[name] + sibling_cls = sibling_classes.get(name) + if sibling_cls is None: + raise ValueError( + f"Load model: type={spec.cls.__name__} bundled sibling '{name}' present in " + f"file but no class was supplied via sibling_classes" + ) + sibling_cfg = fetch_component_config(repo_id, sibling_spec.subfolder) + loaded_siblings[name] = build_component( + component_name=name, + state_dict=sibling_sd, + config=sibling_cfg, + cls=sibling_cls, + converter=None, + acceptable_missing=sibling_spec.acceptable_missing, + quant_args={}, + quant_type=None, + dtype=effective_dtype, + ) + + sd_models.allow_post_quant = False + devices.torch_gc() + log.debug(f"Load model: type={spec.cls.__name__} native_transformer time={time.time() - t0:.2f}") + return transformer, loaded_siblings + + +def strip_prefix(state_dict: dict, prefixes: tuple[str, ...], type_name: str) -> dict: + """Detect and uniformly strip the most common known prefix from every key. + + Order matters: longer prefixes win over shorter ones with the same suffix + (e.g. ``model.diffusion_model.`` beats ``diffusion_model.``). If some keys + match the dominant prefix and others do not, raises ValueError because + mixed prefixes indicate a malformed file rather than a recoverable export + quirk. + """ + sorted_prefixes = sorted(prefixes, key=len, reverse=True) + counts: dict[str, int] = {} + seen = 0 + for key in state_dict: + for prefix in sorted_prefixes: + if key.startswith(prefix): + counts[prefix] = counts.get(prefix, 0) + 1 + seen += 1 + break + total = len(state_dict) + if seen == 0: + log.debug(f"Load model: type={type_name} native_transformer prefix=bare") + return state_dict + dominant = max(counts, key=counts.get) + if counts[dominant] != total: + raise ValueError( + f"Load model: type={type_name} native_transformer has mixed prefixes " + f"(total={total} {dominant}={counts[dominant]})" + ) + log.debug(f'Load model: type={type_name} native_transformer prefix="{dominant}"') + offset = len(dominant) + return {key[offset:]: value for key, value in state_dict.items()} + + +def check_forbidden_markers( + state_dict: dict, + forbidden_markers: tuple[tuple[str, str], ...], + type_name: str, + local_file: str, +) -> None: + """Raise if any forbidden marker key is present in the state_dict. + + Catches structural mismatches that pass prefix detection but indicate the + file is from an incompatible architecture variant (e.g. Cosmos 1.0 keys + showing up in a Cosmos 2.0 loader path). + """ + for marker, description in forbidden_markers: + if marker in state_dict: + raise ValueError( + f"Load model: type={type_name} native_transformer rejects " + f'"{os.path.basename(local_file)}" ({description}; marker key {marker!r})' + ) + + +def partition_siblings( + state_dict: dict, + siblings: dict[str, SiblingSpec], +) -> tuple[dict, dict[str, dict]]: + """Split state_dict into (transformer_sd, {sibling_name: sibling_sd}). + + Keys matching a sibling's ``inline_prefix`` go into that sibling's dict + with the prefix stripped; everything else stays in the transformer dict. + Sibling names with no matching keys still appear in the output dict but + map to an empty dict, so the caller can iterate uniformly. + """ + sibling_sds: dict[str, dict] = {name: {} for name in siblings} + transformer_sd: dict = {} + if not siblings: + return state_dict, sibling_sds + sibling_lookups = [(name, siblings[name].inline_prefix) for name in siblings] + for key, value in state_dict.items(): + matched = False + for name, prefix in sibling_lookups: + if key.startswith(prefix): + sibling_sds[name][key[len(prefix):]] = value + matched = True + break + if not matched: + transformer_sd[key] = value + return transformer_sd, sibling_sds + + +def fetch_component_config(repo_id: str, subfolder: str) -> dict: + """Download and parse ``/config.json`` from the base repo.""" + relative_path = f"{subfolder}/config.json" + try: + local = hf.hf_hub_download( + repo_id, filename=relative_path, cache_dir=shared.opts.diffusers_dir, + ) + except Exception as e: + raise RuntimeError( + f'Load model: native_transformer failed to download {relative_path} ' + f'from repo="{repo_id}": {e}' + ) from e + return shared.readfile(local, as_type="dict") + + +def build_component_quantized( + *, + component_name: str, + state_dict: dict, + config: dict, + cls: type, + quant_args: dict, + dtype, + acceptable_missing: tuple[str, ...], + **kwargs, +) -> object: + """Build a component with per-tensor SDNQ quantization during load. + + Mirrors the per-tensor loop in + :func:`diffusers.models.model_loading_utils.load_model_dict_into_meta`: + iterates the (already-converted) state_dict and dispatches each tensor + through ``SDNQQuantizer.check_if_quantized_param`` / + ``create_quantized_param`` so Linear/Conv/Embed weights are packed to + uint4 as they land, while biases, LayerNorms, and other non-quantizable + parameters go through ``accelerate.set_module_tensor_to_device``. + + The component is constructed inside ``init_empty_weights(include_buffers= + False)`` so parameter slots are meta tensors (no full bf16 instantiation + upfront) while computed buffers like ``rope.freqs`` materialize normally + during ``cls.from_config(config)``. After the loop, the quantizer's + ``_process_model_after_weight_loading`` hook attaches + ``quantization_config`` and handles the CPU offload move. + + Caller is responsible for running the converter and prefix stripping + before passing ``state_dict``. + """ + import rich.progress as rp + from accelerate import init_empty_weights + from accelerate.utils import set_module_tensor_to_device + from modules.sdnq.quantizer import SDNQQuantizer + + quantization_config = quant_args.get("quantization_config") + if quantization_config is None: + raise ValueError( + f"Load model: native_transformer {component_name} " + f"per-tensor quantization requires quant_args['quantization_config']" + ) + + target_dtype = dtype if dtype is not None else devices.dtype + quantizer = SDNQQuantizer(quantization_config, pre_quantized=False) + quantizer.torch_dtype = target_dtype + + with init_empty_weights(include_buffers=False): + component = cls.from_config(config, **kwargs) + + quantizer._process_model_before_weight_loading(component, device_map=None) # pylint: disable=protected-access + + target_device = ( + devices.cpu if shared.opts.diffusers_offload_mode != "none" + else devices.device + ) + + expected_keys = set(component.state_dict().keys()) + loaded_keys: set[str] = set() + unexpected: list[str] = [] + total = len(state_dict) + + pbar = rp.Progress( + rp.TextColumn(f'[cyan]Load {component_name}:'), + rp.BarColumn(), + rp.MofNCompleteColumn(), + rp.TaskProgressColumn(), + rp.TimeRemainingColumn(), + rp.TimeElapsedColumn(), + rp.TextColumn('[cyan]{task.description}'), + console=console, + ) + with pbar: + task = pbar.add_task(total=total, description=cls.__name__) + for name, value in state_dict.items(): + if name in expected_keys: + if torch.is_floating_point(value): + value = value.to(target_dtype) + if quantizer.check_if_quantized_param(component, value, name): + quantizer.create_quantized_param(component, value, name, target_device, dtype=target_dtype) + else: + set_module_tensor_to_device(component, name, target_device, value=value, dtype=target_dtype) + loaded_keys.add(name) + else: + unexpected.append(name) + pbar.update(task, advance=1) + + missing = sorted(expected_keys - loaded_keys) + validate_state_dict_load(component_name, missing, unexpected, acceptable_missing) + + component = quantizer._process_model_after_weight_loading(component) # pylint: disable=protected-access + return component + + +def build_component( + *, + component_name: str, + state_dict: dict, + config: dict, + cls: type, + converter: Callable[[dict], dict] | None, + acceptable_missing: tuple[str, ...], + quant_args: dict, + quant_type: str | None, + dtype=None, + modules_to_not_convert: list | None = None, + modules_dtype_dict: dict | None = None, + **kwargs, +) -> object: + """Convert (if needed), instantiate, load weights, dtype-cast, quantize, + and offload-place a single component. Raises on any hard failure. + + For the transformer component under SDNQ, the per-tensor pre-mode path + in :func:`build_component_quantized` is used so quantization is applied + in flight (one layer's worth of bf16 in memory at a time). All other + cases (siblings, non-quantized loads, NVIDIAModelOptConfig, layerwise + quant) go through the standard load_state_dict + post-quantize path. + + ``dtype`` overrides ``devices.dtype`` when supplied; otherwise the global + default is used. ``modules_to_not_convert`` and ``modules_dtype_dict`` + are forwarded to :func:`apply_quant` for the post-mode path; pre-mode + receives them via the SDNQConfig in ``quant_args``. Extra ``**kwargs`` + reach ``cls.from_config`` for both construction paths. + """ + try: + if converter is not None: + log.debug(f'Load model: native_transformer {component_name} converter={converter.__name__} keys={len(state_dict)}') + sd = converter(state_dict) + else: + sd = state_dict + + if component_name == "transformer" and quant_type == "SDNQConfig": + component = build_component_quantized( + component_name=component_name, + state_dict=sd, + config=config, + cls=cls, + quant_args=quant_args, + dtype=dtype, + acceptable_missing=acceptable_missing, + **kwargs, + ) + del sd + devices.torch_gc() + return component + + log.debug(f'Load model: native_transformer {component_name} loading keys={len(sd)} cls={cls.__name__}') + component = cls.from_config(config, **kwargs) + missing, unexpected = component.load_state_dict(sd, strict=False) + validate_state_dict_load(component_name, missing, unexpected, acceptable_missing) + del sd + devices.torch_gc() + target_dtype = dtype if dtype is not None else devices.dtype + log.debug(f'Load model: native_transformer {component_name} cast dtype={target_dtype}') + component = component.to(dtype=target_dtype) + except Exception as e: + log.error(f"Load model: native_transformer {component_name} load failed: {e}") + errors.display(e, "Load") + raise + + if component_name == "transformer": + apply_quant( + component, + quant_type, + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, + ) + + if shared.opts.diffusers_offload_mode != "none": + sd_models.move_model(component, devices.cpu) + + if not hasattr(component, "quantization_config"): + if hasattr(component, "config") and hasattr(component.config, "quantization_config"): + component.quantization_config = component.config.quantization_config + elif quant_type is not None and quant_args.get("quantization_config") is not None: + component.quantization_config = quant_args.get("quantization_config") + return component + + +def validate_state_dict_load( + component_name: str, + missing: list[str], + unexpected: list[str], + acceptable_missing: tuple[str, ...], +) -> None: + """Raise ValueError if load_state_dict produced unexpected keys or + non-acceptable missing keys. Buffer-only missing keys matching the + ``acceptable_missing`` prefix list are logged at debug level and ignored. + """ + if unexpected: + sample = ", ".join(unexpected[:5]) + raise ValueError( + f"Load model: native_transformer {component_name} has {len(unexpected)} " + f"unexpected keys (sample: {sample})" + ) + hard_missing = [ + k for k in missing if not any(k.startswith(p) for p in acceptable_missing) + ] + if hard_missing: + sample = ", ".join(hard_missing[:5]) + raise ValueError( + f"Load model: native_transformer {component_name} missing " + f"{len(hard_missing)} required keys (sample: {sample})" + ) + if missing: + log.debug( + f"Load model: native_transformer {component_name} ignored " + f"{len(missing)} buffer-only missing keys" + ) + + +def apply_quant( + transformer: object, + quant_type: str | None, + modules_to_not_convert: list | None = None, + modules_dtype_dict: dict | None = None, +) -> None: + """Apply post-load quantization to a fully-loaded transformer. + + Used as a fallback for cases that don't go through the per-tensor + pre-mode path in :func:`build_component_quantized`: ``layerwise_quantization`` + (via ``do_post_load_quant``), and the no-quant case (no-op). SDNQ post + mode also reaches this path because pre-mode dispatch only triggers + when an SDNQConfig is present (which itself only happens under modes + ``pre`` or ``auto``). + + ``modules_to_not_convert`` and ``modules_dtype_dict`` are forwarded + to :func:`sdnq_quantize_model` so per-call skip lists set by the + caller are honored. + """ + if quant_type == "NVIDIAModelOptConfig": + log.warning( + "Load model: native_transformer quant=TRT not supported on native path, skipping" + ) + elif quant_type == "SDNQConfig": + model_quant.sdnq_quantize_model( + transformer, + op="transformer", + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, + ) + model_quant.do_post_load_quant(transformer, allow=False) diff --git a/pipelines/nucleus/__init__.py b/pipelines/nucleus/__init__.py new file mode 100644 index 000000000..798c82edb --- /dev/null +++ b/pipelines/nucleus/__init__.py @@ -0,0 +1,15 @@ +"""Nucleus MoE-Image pipeline package. + +Exports :data:`NUCLEUS_SPEC`. The minimum +``TransformerSpec(cls=NucleusMoEImageTransformer2DModel)`` works because +Nucleus community files use BFL-style ``model.diffusion_model.`` +prefixed keys whose names match the diffusers state_dict verbatim after +prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +NUCLEUS_SPEC = TransformerSpec(cls=diffusers.NucleusMoEImageTransformer2DModel) diff --git a/pipelines/ovis/__init__.py b/pipelines/ovis/__init__.py new file mode 100644 index 000000000..7a7b86c8e --- /dev/null +++ b/pipelines/ovis/__init__.py @@ -0,0 +1,15 @@ +"""Ovis-Image pipeline package. + +Exports :data:`OVIS_SPEC`. The minimum +``TransformerSpec(cls=OvisImageTransformer2DModel)`` works because +Ovis community files use BFL-style ``model.diffusion_model.``-prefixed +keys whose names match the diffusers state_dict verbatim after prefix +strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +OVIS_SPEC = TransformerSpec(cls=diffusers.OvisImageTransformer2DModel) diff --git a/pipelines/pixart/__init__.py b/pipelines/pixart/__init__.py new file mode 100644 index 000000000..fb0dc043a --- /dev/null +++ b/pipelines/pixart/__init__.py @@ -0,0 +1,14 @@ +"""PixArt pipeline package. + +Exports :data:`PIXART_SPEC`. The minimum +``TransformerSpec(cls=PixArtTransformer2DModel)`` works because PixArt +community files use BFL-style ``model.diffusion_model.``-prefixed keys +whose names match the diffusers state_dict verbatim after prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +PIXART_SPEC = TransformerSpec(cls=diffusers.PixArtTransformer2DModel) diff --git a/pipelines/prx/__init__.py b/pipelines/prx/__init__.py new file mode 100644 index 000000000..85d28ea23 --- /dev/null +++ b/pipelines/prx/__init__.py @@ -0,0 +1,14 @@ +"""PRX pipeline package. + +Exports :data:`PRX_SPEC`. The minimum +``TransformerSpec(cls=PRXTransformer2DModel)`` works because PRX +community files use BFL-style ``model.diffusion_model.``-prefixed keys +whose names match the diffusers state_dict verbatim after prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +PRX_SPEC = TransformerSpec(cls=diffusers.PRXTransformer2DModel) diff --git a/pipelines/qwen/__init__.py b/pipelines/qwen/__init__.py index e1455f988..aac930654 100644 --- a/pipelines/qwen/__init__.py +++ b/pipelines/qwen/__init__.py @@ -1,2 +1,19 @@ +"""Qwen-Image pipeline package. + +Exports :data:`QWEN_SPEC`. diffusers registers a no-op identity lambda +for ``QwenImageTransformer2DModel`` in ``SINGLE_FILE_LOADABLE_CLASSES``, +so ``from_single_file`` silently accepts whatever key naming the file +uses and loads with mismatched weights. The spec sets ``converter=None`` +explicitly to skip that no-op; validation then surfaces mismatches as +clear errors. A real converter can be plugged in here if a trainer +format that needs one is encountered. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec from pipelines.qwen.qwen_nunchaku import load_qwen_nunchaku from pipelines.qwen.qwen_pruning import check_qwen_pruning + + +QWEN_SPEC = TransformerSpec(cls=diffusers.QwenImageTransformer2DModel, converter=None) diff --git a/pipelines/step1x/__init__.py b/pipelines/step1x/__init__.py index f7983b96f..2de77fcb0 100644 --- a/pipelines/step1x/__init__.py +++ b/pipelines/step1x/__init__.py @@ -1,8 +1,14 @@ +from pipelines.native_transformer import TransformerSpec from pipelines.step1x.pipeline_output import Step1XEditPipelineOutput from pipelines.step1x.pipeline_step1x_edit import Step1XEditPipeline from pipelines.step1x.transformer_step1x_edit import Step1XEditTransformer2DModel + +STEP1X_SPEC = TransformerSpec(cls=Step1XEditTransformer2DModel) + + __all__ = [ + "STEP1X_SPEC", "Step1XEditPipeline", "Step1XEditPipelineOutput", "Step1XEditTransformer2DModel", diff --git a/pipelines/vibe/__init__.py b/pipelines/vibe/__init__.py index 49b0390cb..5aa63c70c 100644 --- a/pipelines/vibe/__init__.py +++ b/pipelines/vibe/__init__.py @@ -1,6 +1,11 @@ """VIBE pipeline components for SD.Next.""" +from pipelines.native_transformer import TransformerSpec from .vibe_sana_editing import VIBESanaEditingModel from .vibe_sana_pipeline import VIBESanaEditingPipeline, VIBESanaImagePipeline -__all__ = ["VIBESanaEditingModel", "VIBESanaEditingPipeline", "VIBESanaImagePipeline"] + +VIBE_SPEC = TransformerSpec(cls=VIBESanaEditingModel) + + +__all__ = ["VIBE_SPEC", "VIBESanaEditingModel", "VIBESanaEditingPipeline", "VIBESanaImagePipeline"] diff --git a/pipelines/z_image/zimage_lora.py b/pipelines/z_image/zimage_lora.py index 49c35f9d8..f3d39e536 100644 --- a/pipelines/z_image/zimage_lora.py +++ b/pipelines/z_image/zimage_lora.py @@ -22,37 +22,37 @@ OFT block structure is tied to the target module's ``out_features`` so a Q/K/V split is not a drop-in. """ -from modules.lora import native_loader -from modules.lora.native_loader import ChunkSpec +from modules.lora import native_adapter +from modules.lora.native_adapter import ChunkSpec # === Arch-specific prefix configuration === -KNOWN_PREFIXES = native_loader.KNOWN_PREFIXES_DEFAULT +KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT BARE_DIFFUSERS_PREFIXES = ("layers.", "noise_refiner.", "context_refiner.") # === Re-exports for test/back-compat === -LORA_SUFFIXES = native_loader.LORA_SUFFIXES -LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES -LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES -OFT_SUFFIXES = native_loader.OFT_SUFFIXES +LORA_SUFFIXES = native_adapter.LORA_SUFFIXES +LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES +LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES +OFT_SUFFIXES = native_adapter.OFT_SUFFIXES -LORA_MARKERS = native_loader.LORA_MARKERS -LOKR_MARKERS = native_loader.LOKR_MARKERS -LOHA_MARKERS = native_loader.LOHA_MARKERS -OFT_MARKERS = native_loader.OFT_MARKERS +LORA_MARKERS = native_adapter.LORA_MARKERS +LOKR_MARKERS = native_adapter.LOKR_MARKERS +LOHA_MARKERS = native_adapter.LOHA_MARKERS +OFT_MARKERS = native_adapter.OFT_MARKERS -SUFFIX_NORMALIZE = native_loader.SUFFIX_NORMALIZE -BARE_DIFFUSERS_PREFIX_USED = native_loader.BARE_DIFFUSERS_PREFIX_USED -has_marker = native_loader.has_marker +SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE +BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED +has_marker = native_adapter.has_marker def parse_key(key, suffixes): - """Z-Image-bound :func:`native_loader.parse_key`.""" - return native_loader.parse_key( + """Z-Image-bound :func:`native_adapter.parse_key`.""" + return native_adapter.parse_key( key, suffixes, prefixes=KNOWN_PREFIXES, bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, @@ -60,8 +60,8 @@ def parse_key(key, suffixes): def group_by_suffixes(state_dict, suffixes): - """Z-Image-bound :func:`native_loader.group_by_suffixes`.""" - return native_loader.group_by_suffixes( + """Z-Image-bound :func:`native_adapter.group_by_suffixes`.""" + return native_adapter.group_by_suffixes( state_dict, suffixes, prefixes=KNOWN_PREFIXES, bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, @@ -127,7 +127,7 @@ def _underscore_to_diffusers_targets(base): return [(base, None)] -# === Native loaders (thin wrappers over native_loader generics) === +# === Native loaders (thin wrappers over native_adapter generics) === _BIND_KWARGS = dict( @@ -139,24 +139,24 @@ _BIND_KWARGS = dict( def try_load_lora(name, network_on_disk, lora_scale): - return native_loader.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_lokr(name, network_on_disk, lora_scale): - return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_loha(name, network_on_disk, lora_scale): - return native_loader.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_oft(name, network_on_disk, lora_scale): - return native_loader.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load(name, network_on_disk, lora_scale): """Run every Z-Image family loader, merge any that match.""" - return native_loader.try_load_chain( + return native_adapter.try_load_chain( name, network_on_disk, lora_scale, family_loaders=(try_load_lora, try_load_lokr, try_load_loha, try_load_oft), ) diff --git a/test/test-anima-native-adapters.py b/test/test-anima-native-adapters.py index ad77366a1..abbe25213 100644 --- a/test/test-anima-native-adapters.py +++ b/test/test-anima-native-adapters.py @@ -5,10 +5,10 @@ Offline unit tests for Anima native adapter loaders. Anima is the only native arch with a multi-component network namespace: keys route into ``lora_transformer_*`` (Cosmos 2.0 DiT), ``lora_llm_adapter_*`` (a custom Qwen3-projection MLP), or ``lora_te_*`` (Qwen3 text encoder). Routing -is parameterized in ``modules.lora.native_loader`` via the ``network_prefix`` +is parameterized in ``modules.lora.native_adapter`` via the ``network_prefix`` callable that ``pipelines.anima.anima_lora`` supplies. -Covers the eight families exposed through native_loader's generics (LoRA, +Covers the families exposed through native_adapter's generics (LoRA, LoKR, LoHA, OFT, IA3, GLoRA, Norm, Full), focused on: - LoRA across all five recognized prefixes (BFL transformer / BFL llm_adapter / diff --git a/test/test-flux2-native-adapters.py b/test/test-flux2-native-adapters.py index 9400033a8..196a95763 100644 --- a/test/test-flux2-native-adapters.py +++ b/test/test-flux2-native-adapters.py @@ -2,8 +2,8 @@ """ Offline unit tests for Flux2/Klein native adapter loaders. -Covers the nine native families (LoRA, LoKR, LoHA, OFT, BOFT, IA3, -GLoRA, Norm, Full) plus DoRA threading via the universal +Covers the native families (LoRA, LoKR, LoHA, OFT, BOFT, IA3, GLoRA, +Norm, Full) plus DoRA threading via the universal NetworkModule.finalize_updown hook, and ex_bias accumulation across stacked Norm adapters. @@ -580,7 +580,7 @@ def test_parse_key_all_prefixes(): def test_resolve_targets_qkv_chunking(): - from modules.lora.native_loader import ChunkSpec + from modules.lora.native_adapter import ChunkSpec # Kohya double_blocks fused QKV → three chunks targeting Q/K/V. targets = F.resolve_targets('lora_unet_', 'double_blocks_0_img_attn_qkv') assert targets == [ diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py new file mode 100644 index 000000000..538ce84a5 --- /dev/null +++ b/test/test-native-transformer.py @@ -0,0 +1,799 @@ +#!/usr/bin/env python +""" +Offline unit tests for pipelines.native_transformer. + +Covers the pure helpers that own per-arch knob handling: + +- ``strip_prefix`` for single/multi prefix detection and mixed-prefix rejection +- ``partition_siblings`` for inline-sibling key partitioning +- ``check_forbidden_markers`` for structural-mismatch rejection +- ``is_noop_converter`` for diffusers no-op lambda detection +- ``validate_state_dict_load`` for unexpected / missing key handling +- ``make_default_spec`` default-spec synthesis with diffusers converter pickup +- ``auto_pickup_converter`` for diffusers ``SINGLE_FILE_LOADABLE_CLASSES`` integration +- ``TransformerSpec`` / ``SiblingSpec`` defaults + +Plus one end-to-end ``load`` test against a tiny mock module that exercises +the read -> strip -> convert -> from_config -> load_state_dict -> validate +pipeline without needing a real diffusers transformer or hf_hub_download. + +No running server required. + +Usage: + python test/test-native-transformer.py +""" + +import os +import sys +import tempfile + +import torch +import safetensors.torch + +script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, script_dir) +os.chdir(script_dir) + +os.environ['SD_INSTALL_QUIET'] = '1' + +# Bootstrap cmd_args before any module that pulls in shared.py. +import modules.cmd_args # pylint: disable=wrong-import-position +import installer # pylint: disable=wrong-import-position +_orig_argv = sys.argv +sys.argv = [sys.argv[0]] +try: + modules.cmd_args.parse_args() +finally: + sys.argv = _orig_argv +installer.add_args(modules.cmd_args.parser) +modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([]) + +from modules.errors import log # pylint: disable=wrong-import-position +from pipelines import native_transformer as nt # pylint: disable=wrong-import-position + + +# ============================================================ +# Test infrastructure +# ============================================================ + +results: dict[str, dict] = {} + + +def category(name: str): + if name not in results: + results[name] = {'passed': 0, 'failed': 0, 'tests': []} + return name + + +def record(cat: str, passed: bool, name: str, detail: str = ''): + status = 'PASS' if passed else 'FAIL' + results[cat]['passed' if passed else 'failed'] += 1 + results[cat]['tests'].append((status, name)) + msg = f' {status}: {name}' + if detail: + msg += f' ({detail})' + if passed: + log.info(msg) + else: + log.error(msg) + + +def run_test(cat: str, fn): + name = fn.__name__ + try: + ok = fn() + if ok is False: + record(cat, False, name) + else: + record(cat, True, name) + except AssertionError as e: + record(cat, False, name, str(e)) + except Exception as e: # pylint: disable=broad-except + record(cat, False, name, f'exception: {e}') + import traceback + traceback.print_exc() + + +# ============================================================ +# strip_prefix +# ============================================================ + +def test_strip_prefix_bare_keys_pass_through(): + sd = {'layers.0.weight': 1, 'layers.0.bias': 2} + out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test') + assert out == sd, 'bare keys must pass through unchanged' + + +def test_strip_prefix_dominant_single_variant(): + sd = {f'model.diffusion_model.layers.{i}.weight': i for i in range(10)} + out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test') + assert all(k.startswith('layers.') for k in out) + assert len(out) == 10 + + +def test_strip_prefix_picks_longest_match_first(): + """``model.diffusion_model.`` must beat ``diffusion_model.`` when both match.""" + sd = { + 'model.diffusion_model.layers.0.weight': 1, + 'model.diffusion_model.layers.1.weight': 2, + } + out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test') + # If shorter prefix matched, keys would start with 'model.' + assert 'layers.0.weight' in out + assert 'layers.1.weight' in out + assert not any(k.startswith('model.') for k in out) + + +def test_strip_prefix_mixed_prefixes_raises(): + sd = { + 'model.diffusion_model.layers.0.weight': 1, + 'net.layers.0.weight': 2, + } + try: + nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test') + raise AssertionError('expected ValueError') + except ValueError as e: + assert 'mixed prefixes' in str(e) + + +def test_strip_prefix_net_variant(): + sd = {'net.layers.0.weight': 1, 'net.layers.1.bias': 2} + out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test') + assert set(out.keys()) == {'layers.0.weight', 'layers.1.bias'} + + +def test_strip_prefix_diffusion_model_variant(): + sd = {'diffusion_model.layers.0.weight': 1, 'diffusion_model.layers.1.bias': 2} + out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test') + assert set(out.keys()) == {'layers.0.weight', 'layers.1.bias'} + + +def test_strip_prefix_custom_prefix_set(): + sd = {'lora_unet_blocks_0.weight': 1, 'lora_unet_blocks_1.weight': 2} + out = nt.strip_prefix(sd, ('lora_unet_',), 'Test') + assert set(out.keys()) == {'blocks_0.weight', 'blocks_1.weight'} + + +# ============================================================ +# partition_siblings +# ============================================================ + +def test_partition_siblings_empty_spec_returns_state_dict_unchanged(): + sd = {'a': 1, 'b': 2} + transformer_sd, siblings = nt.partition_siblings(sd, {}) + assert transformer_sd == sd + assert siblings == {} + + +def test_partition_siblings_no_matches_keeps_all_in_transformer(): + sd = {'layers.0.weight': 1, 'layers.1.weight': 2} + siblings_spec = {'llm_adapter': nt.SiblingSpec(subfolder='llm_adapter', inline_prefix='llm_adapter.')} + transformer_sd, siblings = nt.partition_siblings(sd, siblings_spec) + assert transformer_sd == sd + assert siblings == {'llm_adapter': {}} + + +def test_partition_siblings_single_sibling_split(): + sd = { + 'layers.0.weight': 'tx0', + 'layers.1.weight': 'tx1', + 'llm_adapter.input_proj.weight': 'ad0', + 'llm_adapter.output_proj.weight': 'ad1', + } + siblings_spec = {'llm_adapter': nt.SiblingSpec(subfolder='llm_adapter', inline_prefix='llm_adapter.')} + transformer_sd, siblings = nt.partition_siblings(sd, siblings_spec) + assert set(transformer_sd.keys()) == {'layers.0.weight', 'layers.1.weight'} + assert set(siblings['llm_adapter'].keys()) == {'input_proj.weight', 'output_proj.weight'} + assert siblings['llm_adapter']['input_proj.weight'] == 'ad0' + + +def test_partition_siblings_multiple_siblings(): + sd = { + 'layers.0.weight': 'tx', + 'sibling_a.x.weight': 'a0', + 'sibling_b.y.weight': 'b0', + 'sibling_b.z.weight': 'b1', + } + siblings_spec = { + 'sibling_a': nt.SiblingSpec(subfolder='a', inline_prefix='sibling_a.'), + 'sibling_b': nt.SiblingSpec(subfolder='b', inline_prefix='sibling_b.'), + } + transformer_sd, siblings = nt.partition_siblings(sd, siblings_spec) + assert list(transformer_sd.keys()) == ['layers.0.weight'] + assert set(siblings['sibling_a'].keys()) == {'x.weight'} + assert set(siblings['sibling_b'].keys()) == {'y.weight', 'z.weight'} + + +# ============================================================ +# check_forbidden_markers +# ============================================================ + +def test_forbidden_markers_passes_when_absent(): + sd = {'layers.0.weight': 1} + markers = (('legacy.marker.weight', 'old format'),) + nt.check_forbidden_markers(sd, markers, 'Test', '/tmp/x.safetensors') + # no exception = pass + + +def test_forbidden_markers_raises_when_present(): + sd = {'layers.0.weight': 1, 'legacy.marker.weight': 2} + markers = (('legacy.marker.weight', 'old Cosmos 1.0 structure'),) + try: + nt.check_forbidden_markers(sd, markers, 'Test', '/tmp/x.safetensors') + raise AssertionError('expected ValueError') + except ValueError as e: + msg = str(e) + assert 'old Cosmos 1.0 structure' in msg + assert 'legacy.marker.weight' in msg + + +def test_forbidden_markers_empty_tuple_no_op(): + sd = {'layers.0.weight': 1} + nt.check_forbidden_markers(sd, (), 'Test', '/tmp/x.safetensors') + + +# ============================================================ +# is_noop_converter +# ============================================================ + +def test_noop_converter_identity_lambda(): + fn = lambda checkpoint, **kwargs: checkpoint # pylint: disable=unnecessary-lambda-assignment + assert nt.is_noop_converter(fn) is True + + +def test_noop_converter_real_function(): + def real(checkpoint, **kwargs): # pylint: disable=unused-argument + return {k.replace('a.', 'b.'): v for k, v in checkpoint.items()} + assert nt.is_noop_converter(real) is False + + +def test_noop_converter_lambda_with_modification(): + fn = lambda checkpoint, **kwargs: {k: v.float() for k, v in checkpoint.items()} # pylint: disable=unnecessary-lambda-assignment + assert nt.is_noop_converter(fn) is False + + +# ============================================================ +# validate_state_dict_load +# ============================================================ + +def test_validate_accepts_buffer_only_missing(): + nt.validate_state_dict_load( + component_name='transformer', + missing=['rope.freqs', 'pos_embedder.pos'], + unexpected=[], + acceptable_missing=('rope.', 'pos_embedder.'), + ) + + +def test_validate_rejects_unexpected(): + try: + nt.validate_state_dict_load( + component_name='transformer', + missing=[], + unexpected=['some.junk.weight'], + acceptable_missing=(), + ) + raise AssertionError('expected ValueError') + except ValueError as e: + assert 'unexpected' in str(e) + assert 'some.junk.weight' in str(e) + + +def test_validate_rejects_hard_missing(): + try: + nt.validate_state_dict_load( + component_name='transformer', + missing=['layers.0.weight', 'rope.freqs'], + unexpected=[], + acceptable_missing=('rope.',), + ) + raise AssertionError('expected ValueError') + except ValueError as e: + msg = str(e) + assert 'missing' in msg + assert 'layers.0.weight' in msg + # Buffer-only missing must not show up in the hard-missing list + assert msg.count('rope.freqs') == 0 + + +def test_validate_empty_passes(): + nt.validate_state_dict_load( + component_name='transformer', + missing=[], + unexpected=[], + acceptable_missing=(), + ) + + +# ============================================================ +# make_default_spec +# ============================================================ + +class FakeTransformer: + """Minimal stand-in for a diffusers transformer class.""" + + +def test_make_default_spec_for_unknown_class(): + spec = nt.make_default_spec(FakeTransformer) + assert spec.cls is FakeTransformer + assert spec.subfolder == 'transformer' + assert spec.prefixes == nt.DEFAULT_PREFIXES + assert spec.converter is None # no diffusers entry for FakeTransformer + assert spec.siblings == {} + assert spec.forbidden_markers == () + + +def test_make_default_spec_picks_up_real_diffusers_converter(): + import diffusers + spec = nt.make_default_spec(diffusers.FluxTransformer2DModel) + assert spec.converter is not None + assert spec.converter.__name__ == 'convert_flux_transformer_checkpoint_to_diffusers' + + +def test_make_default_spec_skips_qwen_image_noop(): + """QwenImageTransformer2DModel's diffusers entry is a no-op lambda; the + default spec must NOT pick it up, leaving converter=None so the caller + sees only their own (potentially absent) override.""" + import diffusers + spec = nt.make_default_spec(diffusers.QwenImageTransformer2DModel) + assert spec.converter is None + + +# ============================================================ +# auto_pickup_converter +# ============================================================ + +def test_auto_pickup_returns_none_for_unknown_class(): + assert nt.auto_pickup_converter(FakeTransformer) is None + + +def test_auto_pickup_returns_real_diffusers_converter(): + import diffusers + fn = nt.auto_pickup_converter(diffusers.FluxTransformer2DModel) + assert fn is not None + assert callable(fn) + assert fn.__name__ == 'convert_flux_transformer_checkpoint_to_diffusers' + + +def test_auto_pickup_skips_noop_converter_qwen(): + """QwenImageTransformer2DModel registers a no-op lambda in diffusers; + auto_pickup_converter must return None so the spec falls back to no + converter (the user-registered spec can override with a real converter).""" + import diffusers + assert nt.auto_pickup_converter(diffusers.QwenImageTransformer2DModel) is None + + +# ============================================================ +# TransformerSpec / SiblingSpec defaults +# ============================================================ + +def test_transformer_spec_defaults(): + spec = nt.TransformerSpec(cls=FakeTransformer) + assert spec.subfolder == 'transformer' + assert spec.prefixes == ('model.diffusion_model.', 'diffusion_model.', 'net.') + assert spec.converter is None + assert spec.siblings == {} + assert spec.acceptable_missing == ('rope.', 'pos_embedder.', 'learnable_pos_embed.') + assert spec.forbidden_markers == () + + +def test_sibling_spec_defaults(): + spec = nt.SiblingSpec(subfolder='llm_adapter', inline_prefix='llm_adapter.') + assert spec.subfolder == 'llm_adapter' + assert spec.inline_prefix == 'llm_adapter.' + assert spec.acceptable_missing == () + + +def test_transformer_spec_is_frozen(): + spec = nt.TransformerSpec(cls=FakeTransformer) + try: + spec.subfolder = 'changed' # type: ignore[misc] + except Exception as e: # pylint: disable=broad-except + assert 'FrozenInstanceError' in type(e).__name__ or 'frozen' in str(e).lower() + return + raise AssertionError('expected FrozenInstanceError') + + +# ============================================================ +# Integration: end-to-end load() with a tiny mock module +# ============================================================ +# We sidestep diffusers + hf_hub_download by patching: +# - ``fetch_component_config`` to return a hand-rolled config dict +# - ``model_quant.get_dit_args`` / ``model_quant.get_quant_type`` / quant +# application to no-ops (we only want to test the load path itself). +# The mock cls is a torch.nn.Module subclass whose ``from_config`` constructs +# a fresh module of the expected shape; ``load_state_dict`` is the standard +# PyTorch method. + +class MockMiniTransformer(torch.nn.Module): + """Tiny stand-in: linear in -> linear out, plus a nested rope sub-module + holding a buffer the trainer state dict won't carry. Nested mirrors how + real DiTs structure rope / pos_embedder buffers.""" + + @classmethod + def from_config(cls, config: dict) -> 'MockMiniTransformer': + return cls(dim=config['dim']) + + def __init__(self, dim: int): + super().__init__() + self.in_proj = torch.nn.Linear(dim, dim) + self.out_proj = torch.nn.Linear(dim, dim) + self.rope = torch.nn.Module() + self.rope.register_buffer('freqs', torch.zeros(dim)) + + +class MockKwargsTransformer(MockMiniTransformer): + """Records the kwargs from_config received, so a test can assert the native + path forwards caller kwargs to construction. Mirrors diffusers from_config, + which accepts **kwargs.""" + + last_kwargs: dict = {} + + @classmethod + def from_config(cls, config: dict, **kwargs) -> 'MockKwargsTransformer': + cls.last_kwargs = dict(kwargs) + return cls(dim=config['dim']) + + +def write_fixture(state_dict_keys: dict, fd: int, path: str) -> str: + os.close(fd) + safetensors.torch.save_file(state_dict_keys, path) + return path + + +def test_load_end_to_end_with_bfl_prefix_no_converter(): + """Exercise the full load pipeline: read .safetensors, strip prefix, + no converter, instantiate via from_config, load weights, validate. + """ + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + # Save with model.diffusion_model. prefix; in_proj.* and out_proj.* + # are the real weights the mock cls expects after the strip. + raw = { + 'model.diffusion_model.in_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.in_proj.bias': torch.zeros(dim), + 'model.diffusion_model.out_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.out_proj.bias': torch.zeros(dim), + } + write_fixture(raw, fd, path) + + # Patch fetch_component_config to return our hand-rolled config. + orig_fetch = nt.fetch_component_config + nt.fetch_component_config = lambda repo, sub: {'dim': dim} + + # Patch quant helpers (we only care about the load path). + from modules import model_quant + orig_get_dit = model_quant.get_dit_args + orig_get_qtype = model_quant.get_quant_type + orig_do_post = model_quant.do_post_load_quant + model_quant.get_dit_args = lambda *a, **k: ({}, {}) + model_quant.get_quant_type = lambda *a, **k: None + model_quant.do_post_load_quant = lambda *a, **k: None + + try: + spec = nt.TransformerSpec(cls=MockMiniTransformer) + transformer, siblings = nt.load( + local_file=path, + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + ) + finally: + nt.fetch_component_config = orig_fetch + model_quant.get_dit_args = orig_get_dit + model_quant.get_quant_type = orig_get_qtype + model_quant.do_post_load_quant = orig_do_post + + assert isinstance(transformer, MockMiniTransformer) + assert transformer.in_proj.weight.shape == (dim, dim) + # Weights from the fixture should match what was loaded. + loaded_in_w = transformer.in_proj.weight.detach().cpu() + fixture_in_w = raw['model.diffusion_model.in_proj.weight'].to(loaded_in_w.dtype) + assert torch.allclose(loaded_in_w, fixture_in_w) + assert siblings == {} + finally: + if os.path.exists(path): + os.unlink(path) + + +def test_load_forwards_kwargs_to_from_config(): + """Caller **kwargs reach cls.from_config through the native load path + rather than being dropped.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + raw = { + 'model.diffusion_model.in_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.in_proj.bias': torch.zeros(dim), + 'model.diffusion_model.out_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.out_proj.bias': torch.zeros(dim), + } + write_fixture(raw, fd, path) + + orig_fetch = nt.fetch_component_config + nt.fetch_component_config = lambda repo, sub: {'dim': dim} + from modules import model_quant + orig_get_dit = model_quant.get_dit_args + orig_get_qtype = model_quant.get_quant_type + orig_do_post = model_quant.do_post_load_quant + model_quant.get_dit_args = lambda *a, **k: ({}, {}) + model_quant.get_quant_type = lambda *a, **k: None + model_quant.do_post_load_quant = lambda *a, **k: None + + MockKwargsTransformer.last_kwargs = {} + try: + spec = nt.TransformerSpec(cls=MockKwargsTransformer) + transformer, _ = nt.load( + local_file=path, + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + low_cpu_mem_usage=True, + ) + finally: + nt.fetch_component_config = orig_fetch + model_quant.get_dit_args = orig_get_dit + model_quant.get_quant_type = orig_get_qtype + model_quant.do_post_load_quant = orig_do_post + + assert MockKwargsTransformer.last_kwargs == {'low_cpu_mem_usage': True} + assert isinstance(transformer, MockKwargsTransformer) + finally: + if os.path.exists(path): + os.unlink(path) + + +def test_load_end_to_end_with_sibling_partition(): + """Bundled-sibling case: file carries both transformer and sibling weights, + sibling_classes supplies the runtime sibling class, partition routes each + half into its target.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + sibling_dim = 4 + raw = { + # Transformer half (after strip). + 'model.diffusion_model.in_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.in_proj.bias': torch.zeros(dim), + 'model.diffusion_model.out_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.out_proj.bias': torch.zeros(dim), + # Sibling half (after strip + sibling partition). + 'model.diffusion_model.sibling.in_proj.weight': torch.randn(sibling_dim, sibling_dim), + 'model.diffusion_model.sibling.in_proj.bias': torch.zeros(sibling_dim), + 'model.diffusion_model.sibling.out_proj.weight': torch.randn(sibling_dim, sibling_dim), + 'model.diffusion_model.sibling.out_proj.bias': torch.zeros(sibling_dim), + } + write_fixture(raw, fd, path) + + orig_fetch = nt.fetch_component_config + + def patched_fetch(_repo, sub): + return {'dim': dim if sub == 'transformer' else sibling_dim} + + nt.fetch_component_config = patched_fetch + + from modules import model_quant + orig_get_dit = model_quant.get_dit_args + orig_get_qtype = model_quant.get_quant_type + orig_do_post = model_quant.do_post_load_quant + model_quant.get_dit_args = lambda *a, **k: ({}, {}) + model_quant.get_quant_type = lambda *a, **k: None + model_quant.do_post_load_quant = lambda *a, **k: None + + try: + spec = nt.TransformerSpec( + cls=MockMiniTransformer, + siblings={ + 'sibling': nt.SiblingSpec( + subfolder='sibling', + inline_prefix='sibling.', + acceptable_missing=('rope.',), + ), + }, + ) + transformer, siblings = nt.load( + local_file=path, + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + sibling_classes={'sibling': MockMiniTransformer}, + ) + finally: + nt.fetch_component_config = orig_fetch + model_quant.get_dit_args = orig_get_dit + model_quant.get_quant_type = orig_get_qtype + model_quant.do_post_load_quant = orig_do_post + + assert isinstance(transformer, MockMiniTransformer) + assert transformer.in_proj.weight.shape == (dim, dim) + assert 'sibling' in siblings + assert isinstance(siblings['sibling'], MockMiniTransformer) + assert siblings['sibling'].in_proj.weight.shape == (sibling_dim, sibling_dim) + finally: + if os.path.exists(path): + os.unlink(path) + + +def test_load_raises_on_missing_sibling_class(): + """Sibling keys present in file but caller forgot to supply the class.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + sibling_dim = 4 + raw = { + 'model.diffusion_model.in_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.in_proj.bias': torch.zeros(dim), + 'model.diffusion_model.out_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.out_proj.bias': torch.zeros(dim), + 'model.diffusion_model.sibling.in_proj.weight': torch.randn(sibling_dim, sibling_dim), + } + write_fixture(raw, fd, path) + + orig_fetch = nt.fetch_component_config + nt.fetch_component_config = lambda repo, sub: {'dim': dim} + + from modules import model_quant + orig_get_dit = model_quant.get_dit_args + orig_get_qtype = model_quant.get_quant_type + orig_do_post = model_quant.do_post_load_quant + model_quant.get_dit_args = lambda *a, **k: ({}, {}) + model_quant.get_quant_type = lambda *a, **k: None + model_quant.do_post_load_quant = lambda *a, **k: None + + try: + spec = nt.TransformerSpec( + cls=MockMiniTransformer, + siblings={'sibling': nt.SiblingSpec(subfolder='s', inline_prefix='sibling.')}, + ) + raised = False + try: + nt.load( + local_file=path, + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + sibling_classes={}, # missing! + ) + except ValueError as e: + raised = True + assert "'sibling'" in str(e) + assert 'sibling_classes' in str(e) + assert raised, 'expected ValueError' + finally: + nt.fetch_component_config = orig_fetch + model_quant.get_dit_args = orig_get_dit + model_quant.get_quant_type = orig_get_qtype + model_quant.do_post_load_quant = orig_do_post + finally: + if os.path.exists(path): + os.unlink(path) + + +def test_load_rejects_non_safetensors(): + spec = nt.TransformerSpec(cls=MockMiniTransformer) + try: + nt.load( + local_file='/tmp/some.gguf', + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + ) + raise AssertionError('expected ValueError') + except ValueError as e: + assert '.safetensors' in str(e) + + +# ============================================================ +# Run +# ============================================================ + +def run_all(): + log.warning('=== strip_prefix ===') + cat = category('strip') + for fn in [ + test_strip_prefix_bare_keys_pass_through, + test_strip_prefix_dominant_single_variant, + test_strip_prefix_picks_longest_match_first, + test_strip_prefix_mixed_prefixes_raises, + test_strip_prefix_net_variant, + test_strip_prefix_diffusion_model_variant, + test_strip_prefix_custom_prefix_set, + ]: + run_test(cat, fn) + + log.warning('=== partition_siblings ===') + cat = category('partition') + for fn in [ + test_partition_siblings_empty_spec_returns_state_dict_unchanged, + test_partition_siblings_no_matches_keeps_all_in_transformer, + test_partition_siblings_single_sibling_split, + test_partition_siblings_multiple_siblings, + ]: + run_test(cat, fn) + + log.warning('=== forbidden_markers ===') + cat = category('forbidden') + for fn in [ + test_forbidden_markers_passes_when_absent, + test_forbidden_markers_raises_when_present, + test_forbidden_markers_empty_tuple_no_op, + ]: + run_test(cat, fn) + + log.warning('=== noop_converter detection ===') + cat = category('noop') + for fn in [ + test_noop_converter_identity_lambda, + test_noop_converter_real_function, + test_noop_converter_lambda_with_modification, + ]: + run_test(cat, fn) + + log.warning('=== validate_state_dict_load ===') + cat = category('validate') + for fn in [ + test_validate_accepts_buffer_only_missing, + test_validate_rejects_unexpected, + test_validate_rejects_hard_missing, + test_validate_empty_passes, + ]: + run_test(cat, fn) + + log.warning('=== make_default_spec ===') + cat = category('default_spec') + for fn in [ + test_make_default_spec_for_unknown_class, + test_make_default_spec_picks_up_real_diffusers_converter, + test_make_default_spec_skips_qwen_image_noop, + ]: + run_test(cat, fn) + + log.warning('=== auto_pickup_converter ===') + cat = category('autopickup') + for fn in [ + test_auto_pickup_returns_none_for_unknown_class, + test_auto_pickup_returns_real_diffusers_converter, + test_auto_pickup_skips_noop_converter_qwen, + ]: + run_test(cat, fn) + + log.warning('=== TransformerSpec / SiblingSpec ===') + cat = category('specs') + for fn in [ + test_transformer_spec_defaults, + test_sibling_spec_defaults, + test_transformer_spec_is_frozen, + ]: + run_test(cat, fn) + + log.warning('=== end-to-end load ===') + cat = category('load') + for fn in [ + test_load_end_to_end_with_bfl_prefix_no_converter, + test_load_forwards_kwargs_to_from_config, + test_load_end_to_end_with_sibling_partition, + test_load_raises_on_missing_sibling_class, + test_load_rejects_non_safetensors, + ]: + run_test(cat, fn) + + log.warning('=== Results ===') + total_passed = 0 + total_failed = 0 + for cat_name, info in results.items(): + ok = info['failed'] == 0 + status = 'PASS' if ok else 'FAIL' + log.info(f" {cat_name}: {info['passed']} passed, {info['failed']} failed [{status}]") + total_passed += info['passed'] + total_failed += info['failed'] + log.warning(f'Total: {total_passed} passed, {total_failed} failed') + return total_failed == 0 + + +if __name__ == '__main__': + import time + t0 = time.time() + ok = run_all() + log.warning(f'Total time: {time.time() - t0:.2f}s') + sys.exit(0 if ok else 1)