diff --git a/data/reference-base.json b/data/reference-base.json index 5ad8fad56..39f430c3b 100644 --- a/data/reference-base.json +++ b/data/reference-base.json @@ -296,6 +296,14 @@ "date": "2026 July", "size": 4.99 }, + "Anima 2.9B Preview v1": { + "path": "CalamitousFelicitousness/Anima-2.9B-Preview-v1-Diffusers", + "preview": "CalamitousFelicitousness--Anima-2.9B-Preview-v1-Diffusers.jpg", + "desc": "Anima 2.9B preview v1 by Gazingstars, a depth-expanded fine-tune of Anima 1.0 Base: the transformer is grown from 28 to 40 layers and trained on 1.7M additional anime and illustration samples with mixed tag and natural-language captions. Same Qwen3-0.6B text encoder, LLM adapter and VAE as Anima 1.0.", + "extras": "sampler: Default, cfg_scale: 4.0, steps: 30", + "date": "2026 August", + "size": 7.32 + }, "Meituan LongCat Image": { "path": "meituan-longcat/LongCat-Image", "preview": "meituan-longcat--LongCat-Image.jpg", diff --git a/models/Reference/CalamitousFelicitousness--Anima-2.9B-Preview-v1-Diffusers.jpg b/models/Reference/CalamitousFelicitousness--Anima-2.9B-Preview-v1-Diffusers.jpg new file mode 100644 index 000000000..a403db079 Binary files /dev/null and b/models/Reference/CalamitousFelicitousness--Anima-2.9B-Preview-v1-Diffusers.jpg differ diff --git a/pipelines/anima/__init__.py b/pipelines/anima/__init__.py index 6d7953090..4dc2b339e 100644 --- a/pipelines/anima/__init__.py +++ b/pipelines/anima/__init__.py @@ -12,16 +12,31 @@ knobs that differ from the native-loader defaults: - 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. +- ``num_layers`` comes from the block indices in the file, so depth-expanded + finetunes (Anima-2.9B: 40 blocks) load against the 28-layer base config. - All other knobs (prefixes, ``acceptable_missing`` buffers) use the defaults from :mod:`pipelines.native_transformer`. """ +import re + import diffusers from diffusers.loaders.single_file_utils import convert_cosmos_transformer_checkpoint_to_diffusers from pipelines.native_transformer import TransformerSpec, SiblingSpec +BLOCK_INDEX = re.compile(r'^(?:transformer_)?blocks\.(\d+)\.') + + +def infer_config(state_dict: dict) -> dict: + """Derive ``num_layers`` from the highest transformer block index in the file.""" + indices = [int(m.group(1)) for key in state_dict if (m := BLOCK_INDEX.match(key))] + if not indices: + return {} + return {'num_layers': max(indices) + 1} + + ANIMA_SPEC = TransformerSpec( cls=diffusers.CosmosTransformer3DModel, converter=convert_cosmos_transformer_checkpoint_to_diffusers, @@ -37,4 +52,5 @@ ANIMA_SPEC = TransformerSpec( 'unsupported Cosmos 1.0 structure', ), ), + infer_config=infer_config, ) diff --git a/pipelines/model_anima.py b/pipelines/model_anima.py index cbf10885c..c3e20dc66 100644 --- a/pipelines/model_anima.py +++ b/pipelines/model_anima.py @@ -1,9 +1,11 @@ +import os import importlib.util import transformers import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae, errors from modules.logger import log from pipelines import generic +from pipelines.generic_map import transformers_map def _import_from_file(module_name, file_path): @@ -13,17 +15,18 @@ def _import_from_file(module_name, file_path): return mod -def init_transformer_component(repo_id, diffusers_load_config, adapter_cls): +def init_transformer_component(repo_id, diffusers_load_config, adapter_cls, local_file=None): """Load (transformer, llm_adapter_or_none). - 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. + A UNET dropdown selection, else ``local_file`` (a single-file checkpoint), + goes through :mod:`pipelines.native_transformer`, which also extracts a + bundled ``llm_adapter``. Without either, the transformer comes from the base + repo and the adapter is ``None`` for the caller to load. """ + from modules import sd_unet from pipelines import native_transformer - local_file = native_transformer.resolve_path() + override = native_transformer.resolve_path() + local_file = override or local_file if local_file is not None: from pipelines.anima import ANIMA_SPEC try: @@ -31,11 +34,13 @@ def init_transformer_component(repo_id, diffusers_load_config, adapter_cls): 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') return None, None + if override is not None: + sd_unet.loaded_unet = shared.opts.sd_unet + return transformer, siblings.get('llm_adapter') transformer = generic.load_transformer( repo_id, cls_name=diffusers.CosmosTransformer3DModel, @@ -51,9 +56,15 @@ def load_anima(checkpoint_info, diffusers_load_config=None): repo_id = sd_models.path_to_repo(checkpoint_info) sd_models.hf_auth_check(checkpoint_info) + # single-file checkpoint: transformer (and bundled llm_adapter) from the file, everything else from the base repo + local_file = None + if repo_id is not None and os.path.isfile(repo_id) and repo_id.lower().endswith('.safetensors'): + local_file = repo_id + repo_id = transformers_map['AnimaTextToImagePipeline'] + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) load_args.pop('cache_dir', None) - log.debug(f'Load model: type=Anima repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + log.debug(f'Load model: type=Anima repo="{repo_id}" file="{local_file}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') if repo_id is None or repo_id.lower() == 'none': return None @@ -69,9 +80,8 @@ def load_anima(checkpoint_info, diffusers_load_config=None): diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["anima"] = AnimaInpaintPipeline generic.set_pipeline('Anima', AnimaTextToImagePipeline) - # UNET dropdown (shared.opts.sd_unet) may redirect the transformer to a - # community file that bundles both the transformer and the llm_adapter. - transformer, llm_adapter = init_transformer_component(repo_id, diffusers_load_config, modeling_llm_adapter.AnimaLLMAdapter) + # UNET dropdown or single-file checkpoint may bundle transformer and llm_adapter + transformer, llm_adapter = init_transformer_component(repo_id, diffusers_load_config, modeling_llm_adapter.AnimaLLMAdapter, local_file=local_file) if transformer is None: return None text_encoder = generic.load_text_encoder( diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index fc844bf5e..3c77d01cd 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -28,9 +28,10 @@ Algorithm: in a Cosmos 2.0 loader). 5. Partition off sibling component keys (e.g. Anima's bundled ``llm_adapter.*``). 6. Run the spec's converter if present (else pass through unchanged). -7. 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. Fetch ``/config.json`` from the base repo, apply the spec's + ``infer_config`` overrides, instantiate via ``cls.from_config``, + ``load_state_dict(strict=False)``, validate, dtype-cast, quantize, and + offload-place. 8. Repeat the build for each populated sibling (no converter, no quant by default; sibling weights are read raw from the bundled file). @@ -134,6 +135,11 @@ class TransformerSpec: ``converter_handles_quant`` runs the converter before comfy_quant detection; such converters must translate marker/scale sidecar keys along with the weights. Float-oriented converters keep the default. + + ``infer_config`` receives the transformer state dict after prefix strip and + sibling partition, before the converter, and returns config overrides + merged over the base repo config (e.g. ``num_layers`` of a depth-expanded + finetune). """ cls: type @@ -146,6 +152,7 @@ class TransformerSpec: acceptable_missing: tuple[str, ...] = DEFAULT_ACCEPTABLE_MISSING zero_init_missing: tuple[str, ...] = () forbidden_markers: tuple[tuple[str, str], ...] = () + infer_config: Callable[[dict], dict] | None = None def make_default_spec(cls: type) -> TransformerSpec: @@ -296,6 +303,12 @@ def load( effective_dtype = dtype if dtype is not None else devices.dtype transformer_cfg = fetch_component_config(repo_id, spec.subfolder) + if spec.infer_config is not None: + inferred = spec.infer_config(transformer_sd) + overrides = {k: v for k, v in inferred.items() if transformer_cfg.get(k) != v} + if overrides: + log.info(f'Load model: type={spec.cls.__name__} native_transformer config={overrides}') + transformer_cfg = {**transformer_cfg, **overrides} transformer = build_component( component_name="transformer", state_dict=transformer_sd, diff --git a/test/test-anima-native-adapters.py b/test/test-anima-native-adapters.py index e2923ad32..08eae0963 100644 --- a/test/test-anima-native-adapters.py +++ b/test/test-anima-native-adapters.py @@ -821,6 +821,10 @@ def test_cosmos_rename_full_coverage(): ('blocks_0_adaln_modulation_cross_attn_2', 'transformer_blocks_0_norm2_linear_2'), ('blocks_0_adaln_modulation_mlp_1', 'transformer_blocks_0_norm3_linear_1'), ('blocks_0_adaln_modulation_mlp_2', 'transformer_blocks_0_norm3_linear_2'), + # Depth-expanded checkpoints (Anima-2.9B carries 40 blocks) + ('blocks_39_self_attn_q_proj', 'transformer_blocks_39_attn1_to_q'), + ('blocks_39_cross_attn_output_proj', 'transformer_blocks_39_attn2_to_out_0'), + ('blocks_39_mlp_layer2', 'transformer_blocks_39_ff_net_2'), ] for src, expected in cases: got = A.cosmos_rename_flat(src) @@ -837,6 +841,9 @@ def test_resolve_targets_per_prefix(): assert A.resolve_targets('diffusion_model.', 'blocks.0.self_attn.q_proj') == [ ('transformer_blocks_0_attn1_to_q', None), ] + assert A.resolve_targets('diffusion_model.', 'blocks.39.cross_attn.k_proj') == [ + ('transformer_blocks_39_attn2_to_k', None), + ] assert A.resolve_targets('diffusion_model.llm_adapter.', 'input_proj') == [ ('input_proj', None), ] diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index 84775b51e..d45cac649 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -843,6 +843,7 @@ def test_transformer_spec_defaults(): assert spec.acceptable_missing == ('rope.', 'pos_embedder.', 'learnable_pos_embed.') assert spec.ignored_prefixes == ('cond_stage_model.', 'conditioner.', 'first_stage_model.', 'text_encoders.', 'vae.') assert spec.forbidden_markers == () + assert spec.infer_config is None def test_sibling_spec_defaults(): @@ -903,6 +904,23 @@ class MockKwargsTransformer(MockMiniTransformer): return cls(dim=config['dim']) +class MockLayeredTransformer(torch.nn.Module): + """Depth comes from config, mirroring ``num_layers`` on real DiTs.""" + + @classmethod + def from_config(cls, config: dict) -> 'MockLayeredTransformer': + return cls(dim=config['dim'], num_layers=config['num_layers']) + + def __init__(self, dim: int, num_layers: int): + super().__init__() + self.blocks = torch.nn.ModuleList([torch.nn.Linear(dim, dim) for _ in range(num_layers)]) + + +def infer_num_layers(state_dict: dict) -> dict: + indices = [int(k.split('.')[1]) for k in state_dict if k.startswith('blocks.')] + return {'num_layers': max(indices) + 1} if indices else {} + + def write_fixture(state_dict_keys: dict, fd: int, path: str, metadata: dict | None = None) -> str: os.close(fd) safetensors.torch.save_file(state_dict_keys, path, metadata=metadata) @@ -1012,6 +1030,51 @@ def test_load_forwards_kwargs_to_from_config(): os.unlink(path) +def test_load_infer_config_sizes_module_to_file(): + """A file deeper than the base config loads only with the hook; without it + the extra block is an arch mismatch.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + file_layers = 3 + raw = {} + for i in range(file_layers): + raw[f'model.diffusion_model.blocks.{i}.weight'] = torch.randn(dim, dim) + raw[f'model.diffusion_model.blocks.{i}.bias'] = torch.zeros(dim) + write_fixture(raw, fd, path) + + orig_fetch = nt.fetch_component_config + nt.fetch_component_config = lambda repo, sub: {'dim': dim, 'num_layers': 2} + 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: + try: + nt.load(local_file=path, repo_id='fake/repo', spec=nt.TransformerSpec(cls=MockLayeredTransformer), diffusers_cfg={}) + raise AssertionError('expected OverrideArchMismatch without infer_config') + except nt.OverrideArchMismatch: + pass + spec = nt.TransformerSpec(cls=MockLayeredTransformer, infer_config=infer_num_layers) + transformer, _ = 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 len(transformer.blocks) == file_layers + loaded_w = transformer.blocks[2].weight.detach().cpu() + assert torch.allclose(loaded_w, raw['model.diffusion_model.blocks.2.weight'].to(loaded_w.dtype)) + 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 @@ -2204,6 +2267,61 @@ def test_build_component_comfy_preempts_sdnq_fresh_quant(): assert component.in_proj.weight.dtype == torch.int8 +# ============================================================ +# anima config inference (depth-expanded checkpoints) +# ============================================================ + +def anima_block_keys(count: int, prefix: str = 'blocks') -> dict: + sd = {} + for i in range(count): + sd[f'{prefix}.{i}.self_attn.q_proj.weight'] = torch.zeros(2, 2) + sd[f'{prefix}.{i}.mlp.layer2.weight'] = torch.zeros(2, 2) + return sd + + +def test_anima_infer_config_counts_cosmos_blocks(): + """40-block file reports num_layers=40; adapter ``blocks.N`` keys and + top-level tensors do not count.""" + from pipelines import anima + sd = anima_block_keys(40) + sd['llm_adapter.blocks.5.self_attn.q_proj.weight'] = torch.zeros(2, 2) + sd['x_embedder.proj.1.weight'] = torch.zeros(2, 2) + sd['final_layer.linear.weight'] = torch.zeros(2, 2) + assert anima.infer_config(sd) == {'num_layers': 40} + + +def test_anima_infer_config_matches_base_depth(): + from pipelines import anima + assert anima.infer_config(anima_block_keys(28)) == {'num_layers': 28} + + +def test_anima_infer_config_diffusers_layout(): + """Files already in diffusers key layout count ``transformer_blocks.N``.""" + from pipelines import anima + assert anima.infer_config(anima_block_keys(40, prefix='transformer_blocks')) == {'num_layers': 40} + + +def test_anima_infer_config_comfy_sidecars_do_not_inflate(): + """Per-block quant sidecars share the block index and do not change the count.""" + from pipelines import anima + sd = anima_block_keys(40) + for i in range(40): + sd[f'blocks.{i}.self_attn.q_proj.weight_scale'] = torch.zeros(2, 1) + sd[f'blocks.{i}.self_attn.q_proj.comfy_quant'] = torch.zeros(8, dtype=torch.uint8) + assert anima.infer_config(sd) == {'num_layers': 40} + + +def test_anima_infer_config_no_blocks_returns_empty(): + from pipelines import anima + assert anima.infer_config({'x_embedder.proj.1.weight': torch.zeros(2, 2)}) == {} + assert anima.infer_config({}) == {} + + +def test_anima_spec_registers_infer_config(): + from pipelines import anima + assert anima.ANIMA_SPEC.infer_config is anima.infer_config + + # ============================================================ # Run # ============================================================ @@ -2353,6 +2471,7 @@ def run_all(): for fn in [ test_load_end_to_end_with_bfl_prefix_no_converter, test_load_forwards_kwargs_to_from_config, + test_load_infer_config_sizes_module_to_file, test_load_end_to_end_with_sibling_partition, test_load_raises_on_missing_sibling_class, test_load_rejects_non_safetensors, @@ -2401,6 +2520,18 @@ def run_all(): ]: run_test(cat, fn) + log.warning('=== anima config inference ===') + cat = category('anima') + for fn in [ + test_anima_infer_config_counts_cosmos_blocks, + test_anima_infer_config_matches_base_depth, + test_anima_infer_config_diffusers_layout, + test_anima_infer_config_comfy_sidecars_do_not_inflate, + test_anima_infer_config_no_blocks_returns_empty, + test_anima_spec_registers_infer_config, + ]: + run_test(cat, fn) + log.warning('=== Results ===') total_passed = 0 total_failed = 0