From cec6d0dce523b64c03fa886e497744ac6e5880ec Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 05:47:56 +0100 Subject: [PATCH] refactor(anima): port anima_transformer to native_transformer 217-line bespoke loader collapses to a 40-line ANIMA_SPEC in pipelines/anima/__init__.py (Cosmos converter + llm_adapter sibling + Cosmos 1.0 forbidden marker). Drop the class-keyed REGISTRY: Anima and raw Cosmos share CosmosTransformer3DModel but need different specs. Specs pass via explicit native_spec= kwarg; make_default_spec(cls) covers the auto-converter case. --- pipelines/anima/__init__.py | 40 +++++ pipelines/anima/anima_transformer.py | 217 --------------------------- pipelines/model_anima.py | 38 ++--- pipelines/native_transformer.py | 47 +++--- test/test-native-transformer.py | 69 +++------ 5 files changed, 94 insertions(+), 317 deletions(-) delete mode 100644 pipelines/anima/anima_transformer.py 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_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/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/native_transformer.py b/pipelines/native_transformer.py index 757fbbe4c..5f46de9fa 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -6,11 +6,15 @@ 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`. Arches register a -spec at import time via :func:`register`; arches without a registration get a -default spec that handles BFL-style ``model.diffusion_model.`` prefix stripping -and opportunistically picks up a diffusers converter from -``SINGLE_FILE_LOADABLE_CLASSES`` if the class has one. +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: @@ -90,32 +94,17 @@ class TransformerSpec: forbidden_markers: tuple[tuple[str, str], ...] = () -REGISTRY: dict[type, TransformerSpec] = {} +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). - -def register(cls: type, spec: TransformerSpec | None = None) -> None: - """Register a transformer class with an explicit spec, or with the default - spec if ``spec`` is None. Idempotent: re-registering the same class - replaces the previous entry. + Used by callers (notably :func:`pipelines.generic.load_transformer`) when + a pipeline does not supply a custom ``TransformerSpec`` of its own. """ - if spec is None: - spec = TransformerSpec(cls=cls) - if spec.cls is not cls: - raise ValueError(f"register: spec.cls ({spec.cls.__name__}) does not match cls ({cls.__name__})") - REGISTRY[cls] = spec - - -def lookup(cls: type) -> TransformerSpec: - """Return the registered spec for ``cls``, or synthesize a default one. - - The synthesized default opportunistically pulls a converter from diffusers' - ``SINGLE_FILE_LOADABLE_CLASSES`` table if one exists for the class name and - is not a pass-through no-op lambda. - """ - if cls in REGISTRY: - return REGISTRY[cls] - converter = auto_pickup_converter(cls) - return TransformerSpec(cls=cls, converter=converter) + return TransformerSpec(cls=cls, converter=auto_pickup_converter(cls)) def auto_pickup_converter(cls: type) -> Callable[[dict], dict] | None: diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index bac70ad6f..b0a524b29 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -9,7 +9,7 @@ Covers the pure helpers that own per-arch knob handling: - ``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 -- ``register`` / ``lookup`` registry behavior and default spec synthesis +- ``make_default_spec`` default-spec synthesis with diffusers converter pickup - ``auto_pickup_converter`` for diffusers ``SINGLE_FILE_LOADABLE_CLASSES`` integration - ``TransformerSpec`` / ``SiblingSpec`` defaults @@ -306,58 +306,37 @@ def test_validate_empty_passes(): # ============================================================ -# register / lookup +# make_default_spec # ============================================================ class FakeTransformer: """Minimal stand-in for a diffusers transformer class.""" -class FakeTransformer2: - """Second stand-in for register/lookup tests.""" - - -def test_register_with_explicit_spec(): - nt.REGISTRY.clear() - spec = nt.TransformerSpec(cls=FakeTransformer, subfolder='custom_sub') - nt.register(FakeTransformer, spec) - assert nt.lookup(FakeTransformer) is spec - - -def test_register_with_default_spec(): - nt.REGISTRY.clear() - nt.register(FakeTransformer) - spec = nt.lookup(FakeTransformer) +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 + assert spec.converter is None # no diffusers entry for FakeTransformer assert spec.siblings == {} + assert spec.forbidden_markers == () -def test_register_idempotent_replaces(): - nt.REGISTRY.clear() - spec1 = nt.TransformerSpec(cls=FakeTransformer, subfolder='sub_one') - spec2 = nt.TransformerSpec(cls=FakeTransformer, subfolder='sub_two') - nt.register(FakeTransformer, spec1) - nt.register(FakeTransformer, spec2) - assert nt.lookup(FakeTransformer) is spec2 +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_register_rejects_mismatched_cls(): - try: - nt.register(FakeTransformer, nt.TransformerSpec(cls=FakeTransformer2)) - raise AssertionError('expected ValueError') - except ValueError as e: - assert 'does not match' in str(e) - - -def test_lookup_synthesizes_default_for_unregistered(): - nt.REGISTRY.clear() - spec = nt.lookup(FakeTransformer) - assert spec.cls is FakeTransformer - assert spec.subfolder == 'transformer' # default - assert spec.converter is None # FakeTransformer has no diffusers entry +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 # ============================================================ @@ -701,14 +680,12 @@ def run_all(): ]: run_test(cat, fn) - log.warning('=== register / lookup ===') - cat = category('registry') + log.warning('=== make_default_spec ===') + cat = category('default_spec') for fn in [ - test_register_with_explicit_spec, - test_register_with_default_spec, - test_register_idempotent_replaces, - test_register_rejects_mismatched_cls, - test_lookup_synthesizes_default_for_unregistered, + 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)