From c998ca0f3ecaeb1b0257d4ba48356ead884e08b6 Mon Sep 17 00:00:00 2001 From: guishiw Date: Tue, 27 Oct 2026 16:21:53 +0800 Subject: [PATCH 1/2] fix(wan): load serialized UMT5 tokenizer pipeline --- .../model_manager/load/model_loaders/wan.py | 11 +- invokeai/backend/wan/tokenizer.py | 112 ++++++++++++++ .../model_manager/load/test_wan_tokenizer.py | 137 ++++++++++++++++++ 3 files changed, 257 insertions(+), 3 deletions(-) create mode 100644 invokeai/backend/wan/tokenizer.py create mode 100644 tests/backend/model_manager/load/test_wan_tokenizer.py diff --git a/invokeai/backend/model_manager/load/model_loaders/wan.py b/invokeai/backend/model_manager/load/model_loaders/wan.py index faf21e528c5..38031001138 100644 --- a/invokeai/backend/model_manager/load/model_loaders/wan.py +++ b/invokeai/backend/model_manager/load/model_loaders/wan.py @@ -72,6 +72,11 @@ def _load_model( patch_wan_causal_conv3d_for_rocm() model_path = Path(config.path) + if submodel_type is SubModelType.Tokenizer: + from invokeai.backend.wan.tokenizer import load_wan_tokenizer + + return load_wan_tokenizer(model_path / submodel_type.value) + load_class = self.get_hf_load_class(model_path, submodel_type) repo_variant = config.repo_variant if isinstance(config, Diffusers_Config_Base) else None variant = repo_variant.value if repo_variant else None @@ -581,7 +586,7 @@ class WanT5EncoderLoader(ModelLoader): what ``Wan-AI/Wan2.2-T2V-A14B::text_encoder+tokenizer`` produces. 2. A flat ``text_encoder/`` folder with ``config.json`` declaring ``model_type: umt5`` directly at the root. In this case the tokenizer - is loaded from the same folder via ``AutoTokenizer.from_pretrained``. + is loaded from the same folder. """ def _load_model( @@ -606,7 +611,7 @@ def _load_model( local_files_only=True, ) if submodel_type == SubModelType.Tokenizer: - from transformers import AutoTokenizer + from invokeai.backend.wan.tokenizer import load_wan_tokenizer # Prefer a sibling tokenizer/ directory; fall back to the encoder dir # itself, which is normal for "flat" downloads. @@ -615,7 +620,7 @@ def _load_model( if nested_tokenizer.exists() else (nested_text_encoder if nested_text_encoder.exists() else root) ) - return AutoTokenizer.from_pretrained(str(target), local_files_only=True) + return load_wan_tokenizer(target) raise ValueError( f"Unsupported submodel type for WanT5Encoder: {submodel_type.value if submodel_type else 'None'}" diff --git a/invokeai/backend/wan/tokenizer.py b/invokeai/backend/wan/tokenizer.py new file mode 100644 index 00000000000..129fd573d99 --- /dev/null +++ b/invokeai/backend/wan/tokenizer.py @@ -0,0 +1,112 @@ +import json +from pathlib import Path +from typing import Any + +_UMT5_SPECIAL_TOKENS = ( + ("pad_token", ""), + ("eos_token", ""), + ("bos_token", ""), + ("unk_token", ""), +) + + +def _apply_umt5_special_token_defaults(tokenizer: Any) -> Any: + """Fill UMT5's fixed special tokens without overriding folder metadata.""" + for token_name, token_value in _UMT5_SPECIAL_TOKENS: + if getattr(tokenizer, token_name) is None: + setattr(tokenizer, token_name, token_value) + return tokenizer + + +def _install_umt5_eos_post_processor(tokenizer: Any) -> None: + """Install the T5 EOS template when a serialized backend does not provide one.""" + from tokenizers.processors import TemplateProcessing + + tokenizer.backend_tokenizer.post_processor = TemplateProcessing( + single=["$A", tokenizer.eos_token], + pair=["$A", tokenizer.eos_token, "$B", tokenizer.eos_token], + special_tokens=[(tokenizer.eos_token, tokenizer.eos_token_id)], + ) + + +def _load_sentencepiece_tokenizer(tokenizer_path: Path) -> Any: + """Convert a standalone UMT5 SentencePiece model without rebuilding it as T5.""" + from sentencepiece import sentencepiece_model_pb2 + from tokenizers import AddedToken, Regex, Tokenizer, decoders, normalizers, pre_tokenizers, processors + from tokenizers.models import Unigram + from transformers import PreTrainedTokenizerFast + + model_proto = sentencepiece_model_pb2.ModelProto() + model_proto.ParseFromString((tokenizer_path / "spiece.model").read_bytes()) + trainer_spec = model_proto.trainer_spec + if trainer_spec.model_type != sentencepiece_model_pb2.TrainerSpec.UNIGRAM: + raise ValueError(f"Wan UMT5 requires a Unigram SentencePiece model: {tokenizer_path / 'spiece.model'}") + + backend = Tokenizer( + Unigram( + [(piece.piece, piece.score) for piece in model_proto.pieces], + unk_id=trainer_spec.unk_id, + # The published UMT5 tokenizer.json deliberately disables the + # SentencePiece model's byte fallback. Wan was trained against that + # behavior, where unsupported text maps to rather than bytes. + byte_fallback=False, + ) + ) + # Only UMT5's named special tokens are AddedTokens in the published + # tokenizer.json. SentencePiece USER_DEFINED pieces such as ``[eod]`` must + # remain ordinary Unigram pieces so metaspace boundaries are preserved. + named_special_tokens = {token_value for _, token_value in _UMT5_SPECIAL_TOKENS} + added_tokens = [ + AddedToken(piece.piece, normalized=False, special=True) + for piece in model_proto.pieces + if piece.piece in named_special_tokens + ] + if added_tokens: + backend.add_tokens(added_tokens) + + tokenizer_normalizers = [] + if model_proto.normalizer_spec.precompiled_charsmap: + tokenizer_normalizers.append(normalizers.Precompiled(model_proto.normalizer_spec.precompiled_charsmap)) + tokenizer_normalizers.append(normalizers.Replace(Regex(" {2,}"), " ")) + backend.normalizer = normalizers.Sequence(tokenizer_normalizers) + backend.pre_tokenizer = pre_tokenizers.Metaspace(replacement="▁", prepend_scheme="always") + backend.decoder = decoders.Metaspace(replacement="▁", prepend_scheme="always") + + if trainer_spec.eos_id >= 0: + eos_token = trainer_spec.eos_piece or "" + backend.post_processor = processors.TemplateProcessing( + single=["$A", eos_token], + pair=["$A", eos_token, "$B", eos_token], + special_tokens=[(eos_token, trainer_spec.eos_id)], + ) + + tokenizer = PreTrainedTokenizerFast.from_pretrained( + str(tokenizer_path), tokenizer_object=backend, local_files_only=True + ) + return _apply_umt5_special_token_defaults(tokenizer) + + +def load_wan_tokenizer(tokenizer_path: Path) -> Any: + """Load a tokenizer for a standalone Wan UMT5 encoder.""" + if (tokenizer_path / "tokenizer.json").is_file(): + from transformers import PreTrainedTokenizerFast + + with (tokenizer_path / "tokenizer.json").open(encoding="utf-8") as tokenizer_file: + has_serialized_post_processor = json.load(tokenizer_file).get("post_processor") is not None + + # Transformers 5's T5Tokenizer rebuilds a Unigram pipeline with unk_id=2. + # UMT5 uses id 2 for and serializes the correct unk_id=3 in tokenizer.json, + # so load that pipeline directly while still applying the folder's token config. + tokenizer = PreTrainedTokenizerFast.from_pretrained(str(tokenizer_path), local_files_only=True) + tokenizer = _apply_umt5_special_token_defaults(tokenizer) + if not has_serialized_post_processor: + _install_umt5_eos_post_processor(tokenizer) + return tokenizer + + if (tokenizer_path / "spiece.model").is_file(): + return _load_sentencepiece_tokenizer(tokenizer_path) + + from transformers import AutoTokenizer + + # Keep support for non-standard tokenizer layouts handled by Transformers. + return AutoTokenizer.from_pretrained(str(tokenizer_path), local_files_only=True) diff --git a/tests/backend/model_manager/load/test_wan_tokenizer.py b/tests/backend/model_manager/load/test_wan_tokenizer.py new file mode 100644 index 00000000000..1dc98821330 --- /dev/null +++ b/tests/backend/model_manager/load/test_wan_tokenizer.py @@ -0,0 +1,137 @@ +import json +from types import SimpleNamespace + +import sentencepiece as spm +from tokenizers import Tokenizer +from tokenizers.models import Unigram + +from invokeai.backend.wan.tokenizer import load_wan_tokenizer + + +def test_load_wan_tokenizer_preserves_serialized_unknown_token_id(tmp_path) -> None: + """Unknown text must use UMT5's id, not its id.""" + backend = Tokenizer( + Unigram( + [("", 0.0), ("", 0.0), ("", 0.0), ("", 0.0), ("music", -1.0)], + unk_id=3, + ) + ) + backend.save(str(tmp_path / "tokenizer.json")) + (tmp_path / "tokenizer_config.json").write_text( + json.dumps( + { + "tokenizer_class": "T5Tokenizer", + "extra_ids": 0, + "pad_token": "", + "eos_token": "", + "bos_token": "", + "unk_token": "", + } + ), + encoding="utf-8", + ) + + tokenizer = load_wan_tokenizer(tmp_path) + + assert tokenizer.unk_token_id == 3 + assert tokenizer.encode("𝄞", add_special_tokens=False) == [3] + + +def test_load_wan_tokenizer_supplies_missing_umt5_special_token_config(tmp_path) -> None: + backend = Tokenizer( + Unigram( + [("", 0.0), ("", 0.0), ("", 0.0), ("", 0.0), ("music", -1.0)], + unk_id=3, + ) + ) + backend.save(str(tmp_path / "tokenizer.json")) + + tokenizer = load_wan_tokenizer(tmp_path) + + assert tokenizer.pad_token_id == 0 + assert tokenizer.eos_token_id == 1 + assert tokenizer.bos_token_id == 2 + assert tokenizer.unk_token_id == 3 + encoded = tokenizer("music", padding="max_length", max_length=3) + assert encoded["input_ids"] == [4, 1, 0] + assert encoded["attention_mask"] == [1, 1, 0] + + +def test_load_wan_tokenizer_reconstructs_umt5_pipeline_without_tokenizer_json(tmp_path) -> None: + corpus = tmp_path / "corpus.txt" + corpus.write_text("music melody video prompt\nhello world generation\na b\n", encoding="utf-8") + tokenizer_dir = tmp_path / "tokenizer" + tokenizer_dir.mkdir() + spm.SentencePieceTrainer.train( + input=str(corpus), + model_prefix=str(tokenizer_dir / "spiece"), + model_type="unigram", + vocab_size=300, + hard_vocab_limit=False, + byte_fallback=True, + pad_id=0, + pad_piece="", + eos_id=1, + eos_piece="", + bos_id=2, + bos_piece="", + unk_id=3, + unk_piece="", + user_defined_symbols=["[eod]"], + minloglevel=2, + ) + sentencepiece = spm.SentencePieceProcessor(model_file=str(tokenizer_dir / "spiece.model")) + + tokenizer = load_wan_tokenizer(tokenizer_dir) + + for text in ("music", " music", "[eod]", "a[eod]b"): + expected = sentencepiece.encode(text, out_type=int) + assert tokenizer.encode(text, add_special_tokens=False) == expected + assert tokenizer.encode(text) == [*expected, sentencepiece.eos_id()] + unknown_ids = tokenizer.encode("𝄞", add_special_tokens=False) + assert unknown_ids[-1] == sentencepiece.unk_id() == 3 + assert 2 not in unknown_ids + assert tokenizer.encode("𝄞") == [*unknown_ids, sentencepiece.eos_id()] + assert tokenizer.unk_token_id == sentencepiece.unk_id() == 3 + + +def test_wan_diffusers_loader_uses_serialized_tokenizer_pipeline(tmp_path) -> None: + from invokeai.backend.model_manager.load.model_loaders.wan import WanDiffusersModel + from invokeai.backend.model_manager.taxonomy import SubModelType + + tokenizer_dir = tmp_path / "tokenizer" + tokenizer_dir.mkdir() + backend = Tokenizer( + Unigram( + [("", 0.0), ("", 0.0), ("", 0.0), ("", 0.0), ("music", -1.0)], + unk_id=3, + ) + ) + backend.save(str(tokenizer_dir / "tokenizer.json")) + (tokenizer_dir / "tokenizer_config.json").write_text( + json.dumps( + { + "tokenizer_class": "T5Tokenizer", + "extra_ids": 0, + "pad_token": "", + "eos_token": "", + "bos_token": "", + "unk_token": "", + } + ), + encoding="utf-8", + ) + (tmp_path / "model_index.json").write_text( + json.dumps({"_class_name": "WanPipeline", "tokenizer": ["transformers", "T5TokenizerFast"]}), + encoding="utf-8", + ) + loader = object.__new__(WanDiffusersModel) + + def fail_if_model_index_class_is_used(*_args, **_kwargs): + raise AssertionError("the Wan tokenizer must not be rebuilt through the model-index T5 class") + + loader.get_hf_load_class = fail_if_model_index_class_is_used + tokenizer = loader._load_model(SimpleNamespace(path=str(tmp_path)), SubModelType.Tokenizer) + + assert tokenizer.unk_token_id == 3 + assert tokenizer.encode("𝄞", add_special_tokens=False) == [3] From 5c95308f82205680d9dffe51871d930da15763fc Mon Sep 17 00:00:00 2001 From: guishiw Date: Sun, 1 Nov 2026 12:41:51 +0800 Subject: [PATCH 2/2] test(wan): cover standalone tokenizer layouts --- .../model_manager/load/test_wan_tokenizer.py | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/tests/backend/model_manager/load/test_wan_tokenizer.py b/tests/backend/model_manager/load/test_wan_tokenizer.py index 1dc98821330..32a4c24924f 100644 --- a/tests/backend/model_manager/load/test_wan_tokenizer.py +++ b/tests/backend/model_manager/load/test_wan_tokenizer.py @@ -1,6 +1,8 @@ import json +from pathlib import Path from types import SimpleNamespace +import pytest import sentencepiece as spm from tokenizers import Tokenizer from tokenizers.models import Unigram @@ -8,6 +10,29 @@ from invokeai.backend.wan.tokenizer import load_wan_tokenizer +def _write_serialized_umt5_tokenizer(path: Path) -> None: + backend = Tokenizer( + Unigram( + [("", 0.0), ("", 0.0), ("", 0.0), ("", 0.0), ("music", -1.0)], + unk_id=3, + ) + ) + backend.save(str(path / "tokenizer.json")) + (path / "tokenizer_config.json").write_text( + json.dumps( + { + "tokenizer_class": "T5Tokenizer", + "extra_ids": 0, + "pad_token": "", + "eos_token": "", + "bos_token": "", + "unk_token": "", + } + ), + encoding="utf-8", + ) + + def test_load_wan_tokenizer_preserves_serialized_unknown_token_id(tmp_path) -> None: """Unknown text must use UMT5's id, not its id.""" backend = Tokenizer( @@ -95,6 +120,31 @@ def test_load_wan_tokenizer_reconstructs_umt5_pipeline_without_tokenizer_json(tm assert tokenizer.unk_token_id == sentencepiece.unk_id() == 3 +@pytest.mark.parametrize("nested", [True, False], ids=["nested", "flat"]) +def test_wan_t5_encoder_loader_uses_serialized_tokenizer_pipeline(tmp_path, nested: bool) -> None: + from invokeai.backend.model_manager.load.model_loaders.wan import WanT5EncoderLoader + from invokeai.backend.model_manager.taxonomy import SubModelType + + model_path = tmp_path / "umt5" + model_path.mkdir() + if nested: + text_encoder_dir = model_path / "text_encoder" + text_encoder_dir.mkdir() + (text_encoder_dir / "config.json").write_text(json.dumps({"model_type": "umt5"}), encoding="utf-8") + tokenizer_dir = model_path / "tokenizer" + tokenizer_dir.mkdir() + else: + (model_path / "config.json").write_text(json.dumps({"model_type": "umt5"}), encoding="utf-8") + tokenizer_dir = model_path + _write_serialized_umt5_tokenizer(tokenizer_dir) + + loader = object.__new__(WanT5EncoderLoader) + tokenizer = loader._load_model(SimpleNamespace(path=str(model_path)), SubModelType.Tokenizer) + + assert tokenizer.unk_token_id == 3 + assert tokenizer.encode("𝄞", add_special_tokens=False) == [3] + + def test_wan_diffusers_loader_uses_serialized_tokenizer_pipeline(tmp_path) -> None: from invokeai.backend.model_manager.load.model_loaders.wan import WanDiffusersModel from invokeai.backend.model_manager.taxonomy import SubModelType