mirror of
https://github.com/vladmandic/automatic
synced 2026-09-20 01:31:13 +02:00
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.
This commit is contained in:
CalamitousFelicitousness
@@ -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',
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
+13
-25
@@ -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')
|
||||
|
||||
@@ -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/<arch>/__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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user