Skip to content

fix(backend): ignore ComfyUI model_sampling keys when loading Anima checkpoints - #9404

Merged
lstein merged 3 commits into
invoke-ai:mainfrom
kappacommit:fix/anima-comfyui-model-sampling-key
Aug 2, 2026
Merged

fix(backend): ignore ComfyUI model_sampling keys when loading Anima checkpoints#9404
lstein merged 3 commits into
invoke-ai:mainfrom
kappacommit:fix/anima-comfyui-model-sampling-key

Conversation

@kappacommit

@kappacommit kappacommit commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Summary

Some Anima finetunes exported from ComfyUI (e.g. Arthemy Comics Anima v2) include a model_sampling.sigmas tensor in the checkpoint. This is exporter metadata, not a model weight, so the Anima single-file loader rejected the whole checkpoint with:

RuntimeError: Checkpoint contains 1 unexpected keys. This may indicate a corrupted or incompatible checkpoint. First 5 unexpected keys: ['model_sampling.sigmas']

The fix filters model_sampling.* keys out of the state dict before loading.

Small refactor included: the existing non-weight key filtering was consolidated into a _filter_non_model_keys() helper driven by suffix/prefix lists, so future cases like this are a one-line addition. Unit tests added.

Related Issues / Discussions

Closes #9402

QA Instructions

  1. Install the Arthemy Comics Anima v2 checkpoint.
  2. Generate an image with it — previously this failed at model load; it now loads and generates normally. Verified locally on Windows/CUDA.

Merge Plan

No special instructions.

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Tests added / updated (if applicable)
  • Documentation added / updated (if applicable)

🤖 Generated with Claude Code

…heckpoints

Some Anima finetunes exported from ComfyUI carry a model_sampling.sigmas
tensor, which is ComfyUI sampling metadata rather than a model weight. The
Anima single-file loader treated it as an unexpected key and refused to load
the checkpoint. Filter model_sampling.* keys out alongside the other
runtime-only tensors.

Closes invoke-ai#9402

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@github-actions github-actions Bot added python PRs that change python files backend PRs that change backend files labels Jul 30, 2026
Replace the inline filter with a _filter_non_model_keys helper driven by
module-level suffix/prefix tuples, so future non-weight checkpoint keys can be
handled by extending the tuples. Filtering now happens right after prefix
stripping, before the RAM estimate and dtype conversion. Add unit tests for
the helper.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@github-actions github-actions Bot added the python-tests PRs that change python tests label Jul 30, 2026
@lstein lstein self-assigned this Jul 31, 2026
@lstein lstein added the 6.14.0 label Jul 31, 2026
@lstein lstein moved this to 6.14.x Theme: USER EXPERIENCE in Invoke - Community Roadmap Jul 31, 2026

@lstein lstein left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approving. I ran an adversarial review of this diff — the working assumption was that it's broken and the job was to prove it. I couldn't. Findings below, with the verification that backs them.

The fix is necessary and targets the right layout

I instantiated the real AnimaTransformer on meta using the exact constructor kwargs from _load_from_singlefile, then replayed every checkpoint layout the Anima probe will actually admit (configs/main.py::_has_anima_keys only accepts bare, net., and model.diffusion_model.) through _strip_anima_bundle_prefix_filter_non_model_keys:

layout main this PR
bare DiT keys + model_sampling.sigmas unexpected=['model_sampling.sigmas'] — issue #9402 verbatim clean
net.* + net.model_sampling.sigmas unexpected=['model_sampling.sigmas'] clean
ComfyUI bundle + top-level model_sampling.sigmas already clean clean
net.* + derived pos_embedder/inv_freq buffers + model_sampling.sigmas fails clean

Worth noting for whoever touches this next: for a genuinely bundled ComfyUI checkpoint (model.diffusion_model.* + first_stage_model.* + cond_stage_model.*), _strip_anima_bundle_prefix already discards model_sampling.*, so the new prefix rule is inert there. The layouts it actually rescues are the bare/unprefixed one and net.model_sampling.* — which is consistent with the reporter seeing exactly one unexpected key in post-strip form.

No legitimate key can be swallowed

model.state_dict() has 685 keys, exactly matching the official-checkpoint fixture. Zero match any entry in _NON_MODEL_KEY_SUFFIXES or _NON_MODEL_KEY_PREFIXES; zero named parameters match. grep -rn model_sampling across the repo hits only this loader and its test, so the new prefix can't collide with a real module now or shadow one by accident.

No silent data loss

AnimaDenoise._get_sigmas derives the schedule analytically from a hardcoded ANIMA_SHIFT = 3.0; there is no code path that has ever read a sigma tensor out of a checkpoint. Dropping model_sampling.sigmas discards nothing InvokeAI could have honored, so there's no risk of trading a hard error for a silently-wrong image.

The relocation is behaviorally inert

Moving the filter above make_room/dtype-conversion hands load_state_dict a bit-identical key set and the identical tensor objects — the filter is unconditional, and dtype conversion is per-tensor and order-independent. Both the unexpected_keys raise and the missing_keys warning are unchanged. The only delta is new_sd_size, now short by ~2 KB against a ~4 GB model; _make_room_internal frees max(0, bytes_needed - ram_available), so that's not an OOM vector, and it's arguably more accurate since rebinding sd releases the pre-filter dict before make_room runs.

Attacks that held

  • Non-str keys reach _filter_non_model_keys without the isinstance(key, str) guard its sibling has — but the only producer is safetensors.load_file, and the format cannot express non-string keys. Unreachable.
  • str.endswith/startswith with tuples can't over-match any of the 685 real keys or the 145 fixture keys.
  • LoRA (lora_model_from_anima_state_dict) and ControlNet-LLLite key spaces never flow through this helper; AnimaControlNetLLLite.from_state_dict pulls only the keys it wants, so it has no equivalent rejection path to fix.

Non-blocking suggestions

  1. No regression test for the actual bug. Nothing composes strip → filter on the issue-#9402 layout (bare DiT keys + model_sampling.sigmas). A one-line composed test would also pin the ordering contract this PR just changed — _load_from_singlefile has no coverage of its own.
  2. Uneven suffix coverage. The new tests exercise .inv_freq and pos_embedder.seq, but not pos_embedder.dim_spatial_range / dim_temporal_range. A typo in either would silently reintroduce the original RuntimeError for animaCatTower-style checkpoints with CI staying green.
  3. test_model_weights_are_untouched asserts less than it appears to. CPython's dict_equal short-circuits on object identity via PyObject_RichCompareBool, so assert out == sd never calls torch.Tensor.__eq__ — it verifies identity-preservation, not values. It can still fail on a dropped key, so it isn't vacuous, just weaker than it reads.

Two things for the maintainers, not for this PR

  • This is the second round of the same whack-a-mole after #9201, and Anima is the only loader in the repo that raises on unexpected keys — qwen_image.py warns, z_image.py ignores them. Every new exporter quirk therefore becomes a user-facing crash that needs a code change and a release. Worth deciding whether that raise should be downgraded to a warning; the tuple-driven refactor here is the right shape either way.
  • Pre-existing: the missing-keys warning says "(expected for inv_freq buffers)", but inv_freq is registered persistent=False and so can never appear in missing_keys. The message is misleading whenever it fires.

CI is fully green, ruff check/format are clean, the six new tests pass, and neither file has drifted against main.

@lstein
lstein enabled auto-merge (squash) August 2, 2026 00:58
@lstein
lstein merged commit 01ed515 into invoke-ai:main Aug 2, 2026
17 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

6.14.0 backend PRs that change backend files python PRs that change python files python-tests PRs that change python tests

Projects

Status: 6.14.x Theme: USER EXPERIENCE

Development

Successfully merging this pull request may close these issues.

[bug]: Server error when generating images using Arthemy Comics Anima v2 checkpoint

2 participants