Skip to content

feat(wan): load the UMT5 text encoder from GGUF builds - #9967

Open
Pfannkuchensack wants to merge 10 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/wan-umt5-gguf
Open

Pfannkuchensack wants to merge 10 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/wan-umt5-gguf

Conversation

@Pfannkuchensack

@Pfannkuchensack Pfannkuchensack commented Oct 21, 2026 •

Copy link
Copy Markdown
Member

Summary

city96's GGUF builds of UMT5-XXL, Wan 2.2's text encoder (city96/umt5-xxl-encoder-gguf), now install as Wan T5 encoders and work in Wan 2.2 generation.

Until now, these files were claimed by the FLUX / SD 3 T5 encoder config. The keys follow the same llama.cpp t5encoder layout, so the file installed, but it produced unusable embeddings and could not be selected for Wan.

Install. A new WanT5Encoder_GGUF_Config:

  • Recognizes UMT5 by the relative attention bias that UMT5 carries in every block (T5 v1.1 has it in block 0 only).
  • Refuses with a reason:
    • incomplete files;
    • encoder-decoder t5 conversions;
    • files without the tokenizer fields.
  • T5Encoder_GGUF_Config no longer claims UMT5 or mT5. Extended-vocabulary T5 v1.1 builds still match.

Load. A new WanT5EncoderGGUFLoader builds transformers' UMT5EncoderModel:

  • Packed: quantized Linear weights stay packed, wo included.
  • Unpack at load: the shared unpack_ggml_at_load (as the LTX-2 text encoder uses) unpacks the rest one tensor at a time. The embedding is decoded in row slices and the RAM reservation comes from the helper.
  • Unpacked: relative biases, norms and the embedding (kept tied to shared).
  • F16 / F32 / BF16 files: their Linear weights are unpacked to the compute dtype. An F32 file therefore no longer stays resident at full precision.
  • Compute dtype: bf16, or float32 on a device without bf16. UMT5 overflows in fp16.

Tokenizer. The tokenizer is built from the GGUF's own fields (vocabulary, scores, flags), so no tokenizer file has to be downloaded or vendored. It matches the Wan 2.2 Diffusers tokenizer.json exactly.

Shared T5 code. The llama.cpp T5 key mapping, the config inference and the feed-forward patch move from the FLUX loader into util/llamacpp_t5.py. Both loaders now use it.

  • The feed-forward patch now covers UMT5's classes.
  • It also stops casting activations to an F16 wo's storage dtype, which made city96's F16 T5 GGUFs fail on a dtype mismatch.

Old records. A UMT5 file that an earlier version installed as a FLUX T5 encoder is refused at load, with a hint to use "Re-identify model".

Nodes.

  • The Wan model loader accepts the GGUF encoder.
  • The Wan prompt node reserves the GGUF dequantization transient, as the other text-encoder nodes do.

Starter. A new starter, "Wan T5 Encoder (UMT5-XXL, GGUF Q8_0)" (~6 GB). Bundles are unchanged.

Docs. The Wan 2.2 page notes that lower quantizations change the conditioning and that Q8_0 stays closest to the Diffusers encoder.

QA Instructions

Automated checks

  • ruff check / ruff format --check on the changed files: clean.
  • pytest -n 4 on tests/backend/model_manager, tests/backend/quantization, tests/model_identification, tests/backend/ltx2, plus the Wan model loader, packed-layer encoder, SD 3 encoder and Wan denoise tests: 3099 passed, 157 skipped, 1 xfailed.
  • openapi.json / schema.ts regenerated (new config class); the api package lint and test pass. webv2 pnpm lint (format, oxlint, tsc, architecture) passes.

New tests

  • Tiny UMT5 GGUFs written with gguf-py (Q8_0, F32, F16, BF16), with different relative biases per layer:
    • a float32 forward matches a dense UMT5EncoderModel built from the same weights;
    • quantized Linear weights stay packed;
    • unquantized weights are unpacked;
    • the embedding is tied;
    • compute dtype falls back to float32 instead of float16.
  • Tokenizer from GGUF fields: space prefix, whitespace collapsing, </s>, padding, truncation.
  • Identification:
    • UMT5 → Wan, T5 v1.1 → FLUX T5, both with exactly one match;
    • mT5 is no longer FLUX T5;
    • incomplete, encoder-decoder and non-t5-tokenizer files are refused;
    • stripped real headers of city96's UMT5 Q8_0 and T5 v1.1 Q3_K_S.
  • Mutation probes, each caught:
    • block 0's bias used for every layer;
    • the feed-forward patch removed;
    • the FLUX config claiming UMT5;
    • prepend_scheme="first";
    • the float16 fallback;
    • the unquantized unpack removed.

Tokenizer parity

Checked against the Wan 2.2 TI2V-5B Diffusers tokenizer.json, using the Wan node's tokenizer arguments: 0 mismatches.

  • 15 hand-picked prompts: EN, ZH, DE, full-width, emoji with ZWJ, runs of spaces, tabs/newlines, literal special tokens, over 512 tokens, empty.
  • A fuzz set of 4021 strings.

On an RTX 4090, through the loader with the app's eval mode

Per-token cosine of the last hidden state against the Wan 2.2 TI2V-5B Diffusers UMT5 in bf16, 8 prompts (EN / ZH / DE / full-width / emoji), token IDs identical:

city96 file cosine, mean (min) rel. L2 VRAM resident load
Q8_0 (6.0 GB) 0.9993 (0.984) 3.5 % 6.9 GiB 22–26 s
Q4_K_M (3.7 GB) 0.971 (0.40) 21 % 4.8 GiB 23 s
  • The Q4_K_M deviation is quantization noise:
    • its Q4_K weights are 7.2 % off the originals;
    • InvokeAI's dequantization matches gguf-py to rel. L2 8e-4.
  • The bf16 Diffusers encoder needs 10.8 GiB of VRAM.
  • The GGUF win is in residency, not in the cold-load RAM peak. That peak is about the same as for the folder build.

Browser (webv2)

  • The installed GGUF appears in the Wan T5 Encoder slot of the Components section.
  • A Wan 2.2 TI2V-5B (GGUF Q4_K_M) generation with it rendered a clean image matching the prompt.
  • The encoder loaded fully onto the GPU (6.7 GB).

Review

Material findings resolved:

  • The Wan GGUF loader uses unpack_ggml_at_load: the embedding is decoded in row slices, and the RAM reservation covers what the load holds, including the float32 widening on a device without bf16.
  • The Wan 2.2 docs say that lower quantizations change the conditioning.
  • F16 / F32 / BF16 files no longer keep full-precision Linear weights resident.
  • The tokenizer tests pin the space-prefix rule and truncation.
  • A float16 fallback no longer runs UMT5 in fp16.
  • Encoder-decoder files get an accurate refusal.
  • T5Encoder_GGUF_Config stops claiming mT5.

Remaining limitations:

  • Cold-load RAM peak: the load reserves what it actually holds through unpack_ggml_at_load (about 7.8 GiB estimated for Q8_0, no separate embedding reservation). The embedding is decoded in row slices, so its decode peak is lower than a whole-tensor decode. The estimate is not re-measured on a GPU run.
  • Header parsing: gguf-py parses this header (256k-token vocabulary) in ~10 s per open. A cold first prompt spends about 35 s loading tokenizer and encoder.
  • Folder tokenizer mismatch: the existing Diffusers-folder tokenizer path (AutoTokenizer under transformers 5.10) does not reproduce tokenizer.json. It maps unknown characters to <s> instead of <unk>. The GGUF tokenizer matches tokenizer.json. Not changed here.

Partial load, FLUX T5 GGUF and multi-GPU (RTX 4090, through the real ModelCache and loaders)

  • Wan, partial load: the encoder was encoded once fully on the GPU and once with max_vram_cache_size_gb=3, which left about 3.9 GiB of the 8.5 GiB (Q8_0) or 6.5 GiB (Q4_K_M) encoder on the GPU. For both Q8_0 and Q4_K_M, the two embeddings are bit-identical (max abs diff 0).
  • FLUX T5 GGUF, main vs this branch: city96's T5 v1.1 XXL Q6_K went through the FLUX loader and HFEncoder (as the FLUX text-encoder node does) for 3 prompts at 512 tokens, on main (e28937b) and on this branch. The embeddings are bit-identical, so the moved key mapping, config inference and feed-forward patch leave FLUX and SD 3 conditioning unchanged. On the branch, a partial load (2 GiB cap) gave the same bit-identical result.
  • Multi-GPU: untested, because there is no second CUDA device. The tied embedding and patched forwards follow the FLUX T5 GGUF loader's pattern.

Compatibility / Rollout

  • New config class: WanT5Encoder_GGUF_Config joins AnyModelConfig, and openapi.json / schema.ts are regenerated. No migration is needed.
  • Existing UMT5 GGUF records: these were installed as FLUX T5 encoders and keep their record. They fail at load with a message until re-identified. "Re-identify model" moves them to the Wan encoder type.

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Meaningful regression coverage added / updated where needed; obsolete tests/code removed
  • Persisted-state and API changes include required migrations / compatibility validation
  • Relevant performance/efficiency opportunities considered; material claims have evidence
  • Material review findings resolved and relevant checks rerun
  • Documentation added / updated (if applicable)
  • Updated What's New copy (if doing a release after this PR)

View with [code]smith Autofix with [code]smith
Need help on this PR? Tag @codesmith-bot with what you need.

city96's UMT5-XXL GGUFs now install as Wan T5 encoders and load with packed weights and a tokenizer built from the GGUF fields; before, they were taken for FLUX T5 encoders.
The llama.cpp T5 helpers move to util/llamacpp_t5.py, and ModelOnDisk.metadata() returns typed GGUF scalars.
@github-actions github-actions Bot added python PRs that change python files invocations PRs that change invocations backend PRs that change backend files python-tests PRs that change python tests docs PRs that change docs labels Oct 21, 2026
# Conflicts:
#	docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx
#	invokeai/frontend/api/schema.ts
# Conflicts:
#	invokeai/frontend/api/schema.ts
@lstein lstein added the 7.0.1 label Oct 29, 2026
@lstein lstein moved this to 7.0.1: bugfixes and small features in Invoke - Community Roadmap Oct 29, 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.

A few things to address before this can merge. I reviewed by reading the code and did not run the branch locally; CI is green at this head.

1. Merge conflict. invokeai/frontend/api/schema.ts conflicts with main (from #10094). Please merge main and regenerate schema.ts and openapi.json, even though openapi.json itself merges cleanly.

2. Use unpack_ggml_at_load in the Wan GGUF loader. Details are inline. This removes the embedding peak the PR lists as a limitation for a later PR, without changing the shared helper.

3. Low-VRAM, multi-GPU and FLUX paths haven't been run. The real run in the description loaded the encoder fully on an RTX 4090, and test_wan_t5_encoder_gguf_loader.py doesn't go through the model cache. Reading the code, the tied embedding and patched forwards follow the same pattern as the FLUX T5 GGUF loader, so I expect them to work, but nothing shows it yet. Low-VRAM users are the main audience for this feature. Please run and report:

  • Partial load: a Wan prompt with the encoder only partly in VRAM, for example with a VRAM cap, checking that the output matches a full load.
  • Multi-GPU: a second device adopting the shared CPU weights. If you don't have the hardware, please list it as untested in the description.
  • FLUX or SD 3: a real generation with a T5 GGUF, since this PR moves that loader's key mapping, config inference and feed-forward patch into shared code and changes the cast in that patch.

4. Docs caveat (docs/src/content/docs/Users Guide/05.Video Generation/wan-2.2.mdx:58). The page offers "any other quantization" with no warning, but your own Q4_K_M measurement shows 21 % relative L2 error and a worst-token cosine of 0.40. Please add a sentence saying that lower quantizations change the conditioning, and that Q8_0 stays closest to the Diffusers encoder.

5. Out-of-date PR description.

  • The "ModelOnDisk.metadata() for GGUF" section describes a change that isn't in this diff: model_on_disk.py isn't among the changed files, because that behavior is already on main.
  • The "Starter" section names "Wan 2.2 T5 Encoder (GGUF Q8_0)", but the code uses "Wan T5 Encoder (UMT5-XXL, GGUF Q8_0)".
  • The embedding reservation limitation will also need updating once item 2 is done.

with accelerate.init_empty_weights():
model = UMT5EncoderModel(infer_t5_encoder_config(sd, UMT5Config))

# Not `unpack_ggml_at_load` with the model's declared skip patterns, as other GGUF loaders do:

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.

The skip patterns are an argument to this helper, so they aren't a reason to avoid it. ltx2.py:746 already calls it with () for a text encoder:

unpack_ggml_at_load(sd, model, (), self._ram_cache.make_room)

Called after the model is built and before load_state_dict_ignoring_extras, it keeps quantized Linear weights packed and unpacks everything else plus unquantized Linears, one tensor at a time. That replaces the loop below and materialize_quantized_embedding; sd.clear() is no longer needed to release the originals. The embed_tokens re-tie stays.

Why it matters:

  • Decode peak. The embedding is the one tensor here decoded whole. unpack_ggml_tensor decodes in row slices, and the comment above _UNPACK_CHUNK_ELEMENTS describes this exact case. This is the limitation the PR defers, and it doesn't need a change to materialize_quantized_embedding.
  • Reservation. Every cold load currently reserves 11.7 GiB, which can evict the RAM cache. The helper reserves what the load actually holds: about 7.8 GiB for Q8_0 by my calculation.
  • Unreserved widening. On a device without bf16, an F16 file widens to float32 in this loop (about 10 GiB more) with nothing reserved. The helper accounts for it.

Numbers for Q8_0:

Now With the helper Basis
Decode rise ~3.9 GiB ~2.2 GiB Measured on CPU at 1/8 scale with get_dequantized_tensor vs the sliced decode, then scaled up (the slice overhead is fixed, so 2.2 is an upper estimate)
Total peak 9.6 GiB ~7.8 GiB Your measurement, and that minus the decode difference

For Q4_K_M, your 11.5 GiB peak implies about 6 GiB saved, but I didn't measure that. I also haven't run this change against your tests, so please confirm the packed/unpacked and tied-embedding tests still pass.

…ml_at_load

Replaces the manual unpack loop and embedding reservation with the shared helper; docs note that lower quantizations change the conditioning.
@Pfannkuchensack

Copy link
Copy Markdown
Member Author

Thanks, all addressed (82fb271, b290123):

  1. Merge: main merged; openapi.json and schema.ts regenerated.
  2. unpack_ggml_at_load: the Wan GGUF loader now calls unpack_ggml_at_load(sd, model, (), self._ram_cache.make_room) before load_state_dict_ignoring_extras. The manual loop, materialize_quantized_embedding and sd.clear() are gone; the embed_tokens re-tie stays. The embedding is decoded in row slices, and the reservation now covers what the load holds, including the float32 widening on a device without bf16. The packed/unpacked and tied-embedding tests pass unchanged.
  3. Runs on an RTX 4090, through the real ModelCache and loaders:
    • Partial load: Wan prompt encoded fully on the GPU and again with max_vram_cache_size_gb=3 (about 3.9 GiB of the encoder resident). The embeddings are bit-identical for Q8_0 and Q4_K_M.
    • FLUX T5 GGUF: city96's T5 v1.1 XXL Q6_K through the FLUX loader and HFEncoder, 3 prompts at 512 tokens, on main and on this branch. The embeddings are bit-identical, also with a partial load on the branch, so FLUX and SD 3 generations with a T5 GGUF are unchanged.
    • Multi-GPU: untested, since I have no second CUDA device. Noted in the description.
  4. Docs: the Wan 2.2 page now says lower quantizations change the conditioning and that, among the quantized builds, Q8_0 stays closest to the Diffusers encoder.
  5. Description: updated (the ModelOnDisk.metadata() section removed, the starter name corrected, the reservation limitation replaced, run results added).

This branch has not been deployed

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

Labels

7.0.1 backend PRs that change backend files docs PRs that change docs invocations PRs that change invocations python PRs that change python files python-tests PRs that change python tests

Projects

Status: 7.0.1: bugfixes and small features

Development

Successfully merging this pull request may close these issues.

2 participants