Skip to content

test(model-loaders): add regression tests for checkpoint key-remapping#9336

Open
Pfannkuchensack wants to merge 2 commits into
invoke-ai:mainfrom
Pfannkuchensack:tests/model-loaders
Open

test(model-loaders): add regression tests for checkpoint key-remapping#9336
Pfannkuchensack wants to merge 2 commits into
invoke-ai:mainfrom
Pfannkuchensack:tests/model-loaders

Conversation

@Pfannkuchensack

Copy link
Copy Markdown
Collaborator

Summary

This PR fixes a FLUX/Qwen single-file loading regression and adds the regression-test coverage that would have caught it.

The fix. Qwen2_5_VLForConditionalGeneration._checkpoint_conversion_mapping is empty in transformers 5.x, so the if key_mapping: guard skipped remapping and ComfyUI single-file Qwen2.5-VL encoder checkpoints (visual.*, model.*) failed to load: their keys were reported as unexpected while model.visual.* params stayed as meta tensors. We fall back to the legacy mapping when transformers no longer provides one; a negative lookahead keeps new-layout checkpoints untouched.

The tests. Extract the state-dict key surgery in the qwen_image, flux, and anima loaders into pure module-level helpers and cover them with unit tests, so a regression like this one fails in CI instead of at load time.

  • qwen_image: extract _remap_qwen_vl_checkpoint_keys
  • flux: move the FLUX.2 BFL->diffusers transformer/VAE converters into a new flux2_state_dict_utils module, deduplicating the two identical copies that lived in the checkpoint and GGUF loader classes
  • anima: extract _strip_anima_bundle_prefix

Tests use fixtures captured from real single-file checkpoints and, where the diffusers architecture builds cheaply, assert on the meta device that every converted key maps to a real parameter (the FLUX.2 VAE is checked for full bijective coverage). Behavior of all loaders is unchanged.

Related Issues / Discussions

N/A

QA Instructions

Automated (no network, no GPU, no real weights):
uv run --extra cuda --extra test python -m pytest tests/backend/model_manager/load/ --no-cov

  • 28 new tests pass; the full load/ suite stays green (246 passed, 122 skipped).
  • Regression counter-check: temporarily remove the fallback mapping in _remap_qwen_vl_checkpoint_keystest_qwen_image_state_dict_utils.py goes red (incl. the bug-report-keys test), proving the tests catch the original bug.

Manual (optional): load a ComfyUI single-file Qwen2.5-VL encoder (e.g. qwen_2.5_vl_7b_fp8_scaled.safetensors) and run a qwen_image_text_encoder invocation — before the fix it raised Failed to load all parameters from checkpoint. Meta tensors remain: [...] with a 728 unexpected keys warning; after the fix it loads cleanly.

Merge Plan

Standard merge — no DB/schema changes, no migrations. Loader behavior is unchanged (converters were only extracted/deduplicated), validated end-to-end against real FLUX.2 transformer/VAE checkpoints (0 unmatched keys).

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Tests added / updated (if applicable)
  • ❗Changes to a redux slice have a corresponding migration
  • Documentation added / updated (if applicable)
  • Updated What's New copy (if doing a release after this PR)

Extract the state-dict key surgery in the qwen_image, flux, and anima loaders
into pure module-level helpers and cover them with unit tests, so a regression
like the transformers-5.x Qwen-VL remap bug fails in CI instead of at load time.

- qwen_image: extract _remap_qwen_vl_checkpoint_keys
- flux: move the FLUX.2 BFL->diffusers transformer/VAE converters into a new
  flux2_state_dict_utils module, deduplicating the two identical copies that
  lived in the checkpoint and GGUF loader classes
- anima: extract _strip_anima_bundle_prefix

Tests use fixtures captured from real single-file checkpoints and, where the
diffusers architecture builds cheaply, assert on the meta device that every
converted key maps to a real parameter (the FLUX.2 VAE is checked for full
bijective coverage). Behavior of all loaders is unchanged.
@github-actions github-actions Bot added python PRs that change python files backend PRs that change backend files python-tests PRs that change python tests labels Jul 5, 2026
@JPPhoto JPPhoto self-assigned this Jul 6, 2026

@JPPhoto JPPhoto 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.

LGTM!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants