fix(wan): load serialized UMT5 tokenizer pipeline

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'}"

invokeai/backend/wan/tokenizer.py

@@ -0,0 +1,112 @@
+import json
+from pathlib import Path
+from typing import Any
+
+_UMT5_SPECIAL_TOKENS = (
+    ("pad_token", "<pad>"),
+    ("eos_token", "</s>"),
+    ("bos_token", "<s>"),
+    ("unk_token", "<unk>"),
+)
+
+
+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 <unk> 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 "</s>"
+        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 <s> 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)

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 <unk> id, not its <s> id."""
+    backend = Tokenizer(
+        Unigram(
+            [("<pad>", 0.0), ("</s>", 0.0), ("<s>", 0.0), ("<unk>", 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": "<pad>",
+                "eos_token": "</s>",
+                "bos_token": "<s>",
+                "unk_token": "<unk>",
+            }
+        ),
+        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(
+            [("<pad>", 0.0), ("</s>", 0.0), ("<s>", 0.0), ("<unk>", 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="<pad>",
+        eos_id=1,
+        eos_piece="</s>",
+        bos_id=2,
+        bos_piece="<s>",
+        unk_id=3,
+        unk_piece="<unk>",
+        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(
+            [("<pad>", 0.0), ("</s>", 0.0), ("<s>", 0.0), ("<unk>", 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": "<pad>",
+                "eos_token": "</s>",
+                "bos_token": "<s>",
+                "unk_token": "<unk>",
+            }
+        ),
+        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]