diff --git a/invokeai/backend/model_manager/load/model_loaders/z_image.py b/invokeai/backend/model_manager/load/model_loaders/z_image.py index ae80f0c5160..406d243c371 100644 --- a/invokeai/backend/model_manager/load/model_loaders/z_image.py +++ b/invokeai/backend/model_manager/load/model_loaders/z_image.py @@ -28,6 +28,7 @@ SubModelType, ) from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader +from invokeai.backend.qwen3.qwen3_tokenizer import load_bundled_qwen3_tokenizer from invokeai.backend.util.devices import TorchDevice @@ -533,10 +534,6 @@ def _load_control_adapter( class Qwen3EncoderCheckpointLoader(ModelLoader): """Class to load single-file Qwen3 Encoder models for Z-Image (safetensors format).""" - # Default HuggingFace model to load tokenizer from when using single-file Qwen3 encoder - # Must be Qwen3 (not Qwen2.5) to match Z-Image's text encoder architecture and special tokens - DEFAULT_TOKENIZER_SOURCE = "Qwen/Qwen3-4B" - def _load_model( self, config: AnyModelConfig, @@ -549,27 +546,22 @@ def _load_model( case SubModelType.TextEncoder: return self._load_from_singlefile(config) case SubModelType.Tokenizer: - # For single-file Qwen3, load tokenizer from HuggingFace - # Try local cache first to support offline usage after initial download - return self._load_tokenizer_with_offline_fallback() + # Single-file checkpoints ship no tokenizer files; use the vendored copy. + return self._load_bundled_tokenizer() raise ValueError( f"Only TextEncoder and Tokenizer submodels are supported. Received: {submodel_type.value if submodel_type else 'None'}" ) - def _load_tokenizer_with_offline_fallback(self) -> AnyModel: - """Load tokenizer with local_files_only fallback for offline support. + def _load_bundled_tokenizer(self) -> AnyModel: + """Load the Qwen3 tokenizer from the vendored, bundled copy. - First tries to load from local cache (offline), falling back to network download - if the tokenizer hasn't been cached yet. This ensures offline operation after - the initial download. + Single-file / GGUF checkpoints do not ship tokenizer files. The Qwen3 BPE + tokenizer is identical across the 0.6B / 4B / 8B variants, so we load the + self-contained copy vendored in the package — fully offline, no HuggingFace + download required. """ - try: - # Try loading from local cache first (supports offline usage) - return AutoTokenizer.from_pretrained(self.DEFAULT_TOKENIZER_SOURCE, local_files_only=True) - except OSError: - # Not in cache yet, download from HuggingFace - return AutoTokenizer.from_pretrained(self.DEFAULT_TOKENIZER_SOURCE) + return load_bundled_qwen3_tokenizer() def _load_from_singlefile( self, @@ -790,10 +782,6 @@ def _load_from_singlefile( class Qwen3EncoderGGUFLoader(ModelLoader): """Class to load GGUF-quantized Qwen3 Encoder models for Z-Image.""" - # Default HuggingFace model to load tokenizer from when using GGUF Qwen3 encoder - # Must be Qwen3 (not Qwen2.5) to match Z-Image's text encoder architecture and special tokens - DEFAULT_TOKENIZER_SOURCE = "Qwen/Qwen3-4B" - def _load_model( self, config: AnyModelConfig, @@ -806,27 +794,22 @@ def _load_model( case SubModelType.TextEncoder: return self._load_from_gguf(config) case SubModelType.Tokenizer: - # For GGUF Qwen3, load tokenizer from HuggingFace - # Try local cache first to support offline usage after initial download - return self._load_tokenizer_with_offline_fallback() + # GGUF checkpoints ship no tokenizer files; use the vendored copy. + return self._load_bundled_tokenizer() raise ValueError( f"Only TextEncoder and Tokenizer submodels are supported. Received: {submodel_type.value if submodel_type else 'None'}" ) - def _load_tokenizer_with_offline_fallback(self) -> AnyModel: - """Load tokenizer with local_files_only fallback for offline support. + def _load_bundled_tokenizer(self) -> AnyModel: + """Load the Qwen3 tokenizer from the vendored, bundled copy. - First tries to load from local cache (offline), falling back to network download - if the tokenizer hasn't been cached yet. This ensures offline operation after - the initial download. + Single-file / GGUF checkpoints do not ship tokenizer files. The Qwen3 BPE + tokenizer is identical across the 0.6B / 4B / 8B variants, so we load the + self-contained copy vendored in the package — fully offline, no HuggingFace + download required. """ - try: - # Try loading from local cache first (supports offline usage) - return AutoTokenizer.from_pretrained(self.DEFAULT_TOKENIZER_SOURCE, local_files_only=True) - except OSError: - # Not in cache yet, download from HuggingFace - return AutoTokenizer.from_pretrained(self.DEFAULT_TOKENIZER_SOURCE) + return load_bundled_qwen3_tokenizer() def _load_from_gguf( self, diff --git a/invokeai/backend/qwen3/__init__.py b/invokeai/backend/qwen3/__init__.py new file mode 100644 index 00000000000..f5cfc572ab0 --- /dev/null +++ b/invokeai/backend/qwen3/__init__.py @@ -0,0 +1,6 @@ +"""Qwen3 encoder backend module. + +Shared assets for the standalone Qwen3 text encoders used by Z-Image (4B/8B) and +Anima (0.6B). The Qwen3 BPE tokenizer is identical across all variants, so a single +vendored copy is bundled here and reused by every Qwen3 encoder loader. +""" diff --git a/invokeai/backend/qwen3/qwen3_tokenizer.py b/invokeai/backend/qwen3/qwen3_tokenizer.py new file mode 100644 index 00000000000..7ed651057a3 --- /dev/null +++ b/invokeai/backend/qwen3/qwen3_tokenizer.py @@ -0,0 +1,37 @@ +"""Bundled Qwen3 tokenizer for single-file / GGUF Qwen3 encoders. + +Single-file (safetensors) and GGUF Qwen3 encoder checkpoints ship weights only — +no tokenizer files. Previously the tokenizer was pulled from HuggingFace +(``Qwen/Qwen3-4B``) on first use, which fails in offline / airgapped setups and +whenever the HF cache is not persisted. The Qwen3 BPE tokenizer is identical +across the 0.6B / 4B / 8B variants, so a single self-contained copy (Apache-2.0, +vendored from ``Qwen/Qwen3-4B``) serves every Qwen3 encoder, fully offline. +""" + +import gzip +import shutil +import tempfile +from functools import lru_cache +from pathlib import Path + +from transformers import AutoTokenizer, PreTrainedTokenizerBase + +_TOKENIZER_DIR = Path(__file__).parent / "tokenizer" + +# tokenizer.json is ~11MB uncompressed, over the repo's 10MB LFS threshold, so it is +# vendored gzip-compressed and expanded into a temp dir at load time. +_TOKENIZER_JSON_GZ = _TOKENIZER_DIR / "tokenizer.json.gz" + + +@lru_cache(maxsize=1) +def load_bundled_qwen3_tokenizer() -> PreTrainedTokenizerBase: + """Load the vendored Qwen3 fast tokenizer. Result is cached for the process.""" + # The fast tokenizer reads the vocab fully into memory at construction, so the + # decompressed files only need to exist for the duration of from_pretrained(). + with tempfile.TemporaryDirectory() as tmp: + tmp_dir = Path(tmp) + with gzip.open(_TOKENIZER_JSON_GZ, "rb") as src, open(tmp_dir / "tokenizer.json", "wb") as dst: + shutil.copyfileobj(src, dst) + for name in ("tokenizer_config.json", "special_tokens_map.json"): + shutil.copyfile(_TOKENIZER_DIR / name, tmp_dir / name) + return AutoTokenizer.from_pretrained(tmp_dir, local_files_only=True) diff --git a/invokeai/backend/qwen3/tokenizer/special_tokens_map.json b/invokeai/backend/qwen3/tokenizer/special_tokens_map.json new file mode 100644 index 00000000000..ac23c0aaa24 --- /dev/null +++ b/invokeai/backend/qwen3/tokenizer/special_tokens_map.json @@ -0,0 +1,31 @@ +{ + "additional_special_tokens": [ + "<|im_start|>", + "<|im_end|>", + "<|object_ref_start|>", + "<|object_ref_end|>", + "<|box_start|>", + "<|box_end|>", + "<|quad_start|>", + "<|quad_end|>", + "<|vision_start|>", + "<|vision_end|>", + "<|vision_pad|>", + "<|image_pad|>", + "<|video_pad|>" + ], + "eos_token": { + "content": "<|im_end|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false + }, + "pad_token": { + "content": "<|endoftext|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false + } +} diff --git a/invokeai/backend/qwen3/tokenizer/tokenizer.json.gz b/invokeai/backend/qwen3/tokenizer/tokenizer.json.gz new file mode 100644 index 00000000000..c83504c9bf3 Binary files /dev/null and b/invokeai/backend/qwen3/tokenizer/tokenizer.json.gz differ diff --git a/invokeai/backend/qwen3/tokenizer/tokenizer_config.json b/invokeai/backend/qwen3/tokenizer/tokenizer_config.json new file mode 100644 index 00000000000..f17e6b5d5a5 --- /dev/null +++ b/invokeai/backend/qwen3/tokenizer/tokenizer_config.json @@ -0,0 +1,240 @@ +{ + "add_bos_token": false, + "add_prefix_space": false, + "added_tokens_decoder": { + "151643": { + "content": "<|endoftext|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151644": { + "content": "<|im_start|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151645": { + "content": "<|im_end|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151646": { + "content": "<|object_ref_start|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151647": { + "content": "<|object_ref_end|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151648": { + "content": "<|box_start|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151649": { + "content": "<|box_end|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151650": { + "content": "<|quad_start|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151651": { + "content": "<|quad_end|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151652": { + "content": "<|vision_start|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151653": { + "content": "<|vision_end|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151654": { + "content": "<|vision_pad|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151655": { + "content": "<|image_pad|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151656": { + "content": "<|video_pad|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "151657": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "151658": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "151659": { + "content": "<|fim_prefix|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "151660": { + "content": "<|fim_middle|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "151661": { + "content": "<|fim_suffix|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "151662": { + "content": "<|fim_pad|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "151663": { + "content": "<|repo_name|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "151664": { + "content": "<|file_sep|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "151665": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "151666": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "151667": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "151668": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + } + }, + "additional_special_tokens": [ + "<|im_start|>", + "<|im_end|>", + "<|object_ref_start|>", + "<|object_ref_end|>", + "<|box_start|>", + "<|box_end|>", + "<|quad_start|>", + "<|quad_end|>", + "<|vision_start|>", + "<|vision_end|>", + "<|vision_pad|>", + "<|image_pad|>", + "<|video_pad|>" + ], + "bos_token": null, + "clean_up_tokenization_spaces": false, + "eos_token": "<|im_end|>", + "errors": "replace", + "extra_special_tokens": {}, + "model_max_length": 131072, + "pad_token": "<|endoftext|>", + "split_special_tokens": false, + "tokenizer_class": "Qwen2Tokenizer", + "unk_token": null, + "chat_template": "{%- if tools %}\n {{- '<|im_start|>system\\n' }}\n {%- if messages[0].role == 'system' %}\n {{- messages[0].content + '\\n\\n' }}\n {%- endif %}\n {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within XML tags:\\n\" }}\n {%- for tool in tools %}\n {{- \"\\n\" }}\n {{- tool | tojson }}\n {%- endfor %}\n {{- \"\\n\\n\\nFor each function call, return a json object with function name and arguments within XML tags:\\n\\n{\\\"name\\\": , \\\"arguments\\\": }\\n<|im_end|>\\n\" }}\n{%- else %}\n {%- if messages[0].role == 'system' %}\n {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n {%- set index = (messages|length - 1) - loop.index0 %}\n {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('') and message.content.endswith('')) %}\n {%- set ns.multi_step_tool = false %}\n {%- set ns.last_query_index = index %}\n {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n {%- if message.content is string %}\n {%- set content = message.content %}\n {%- else %}\n {%- set content = '' %}\n {%- endif %}\n {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n {%- elif message.role == \"assistant\" %}\n {%- set reasoning_content = '' %}\n {%- if message.reasoning_content is string %}\n {%- set reasoning_content = message.reasoning_content %}\n {%- else %}\n {%- if '' in content %}\n {%- set reasoning_content = content.split('')[0].rstrip('\\n').split('')[-1].lstrip('\\n') %}\n {%- set content = content.split('')[-1].lstrip('\\n') %}\n {%- endif %}\n {%- endif %}\n {%- if loop.index0 > ns.last_query_index %}\n {%- if loop.last or (not loop.last and reasoning_content) %}\n {{- '<|im_start|>' + message.role + '\\n\\n' + reasoning_content.strip('\\n') + '\\n\\n\\n' + content.lstrip('\\n') }}\n {%- else %}\n {{- '<|im_start|>' + message.role + '\\n' + content }}\n {%- endif %}\n {%- else %}\n {{- '<|im_start|>' + message.role + '\\n' + content }}\n {%- endif %}\n {%- if message.tool_calls %}\n {%- for tool_call in message.tool_calls %}\n {%- if (loop.first and content) or (not loop.first) %}\n {{- '\\n' }}\n {%- endif %}\n {%- if tool_call.function %}\n {%- set tool_call = tool_call.function %}\n {%- endif %}\n {{- '\\n{\"name\": \"' }}\n {{- tool_call.name }}\n {{- '\", \"arguments\": ' }}\n {%- if tool_call.arguments is string %}\n {{- tool_call.arguments }}\n {%- else %}\n {{- tool_call.arguments | tojson }}\n {%- endif %}\n {{- '}\\n' }}\n {%- endfor %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"tool\" %}\n {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n {{- '<|im_start|>user' }}\n {%- endif %}\n {{- '\\n\\n' }}\n {{- content }}\n {{- '\\n' }}\n {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n {{- '<|im_start|>assistant\\n' }}\n {%- if enable_thinking is defined and enable_thinking is false %}\n {{- '\\n\\n\\n\\n' }}\n {%- endif %}\n{%- endif %}" +} diff --git a/pyproject.toml b/pyproject.toml index bab9b8631cf..ca4a8b5f65f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -233,6 +233,8 @@ version = { attr = "invokeai.version.__version__" } [tool.setuptools.package-data] "invokeai.app.assets" = ["**/*.png"] +"invokeai.backend.anima" = ["tokenizer/*.json"] +"invokeai.backend.qwen3" = ["tokenizer/*.json", "tokenizer/*.json.gz"] "invokeai.backend.t5" = ["tokenizer/*.json"] "invokeai.app.services.workflow_records.default_workflows" = ["*.json"] "invokeai.app.services.style_preset_records" = ["*.json"] diff --git a/tests/backend/qwen3/test_qwen3_tokenizer.py b/tests/backend/qwen3/test_qwen3_tokenizer.py new file mode 100644 index 00000000000..e24c7282cbb --- /dev/null +++ b/tests/backend/qwen3/test_qwen3_tokenizer.py @@ -0,0 +1,55 @@ +"""Tests for the bundled Qwen3 tokenizer used by single-file / GGUF Qwen3 encoders. + +Single-file and GGUF Qwen3 encoder checkpoints (Z-Image, Anima) ship weights only. +The tokenizer is vendored in the package so the encoder works fully offline instead +of pulling ``Qwen/Qwen3-4B`` from HuggingFace on first use. +""" + +from invokeai.backend.qwen3.qwen3_tokenizer import load_bundled_qwen3_tokenizer + +# The Qwen3 BPE tokenizer is shared across the 0.6B / 4B / 8B variants. +QWEN3_VOCAB_SIZE = 151643 + + +def test_bundled_tokenizer_is_fast() -> None: + tokenizer = load_bundled_qwen3_tokenizer() + assert tokenizer.is_fast + + +def test_bundled_tokenizer_known_ids() -> None: + tokenizer = load_bundled_qwen3_tokenizer() + ids = tokenizer("A cinematic photo of a cat").input_ids + assert ids == [32, 64665, 6548, 315, 264, 8251] + + +def test_bundled_tokenizer_roundtrip() -> None: + tokenizer = load_bundled_qwen3_tokenizer() + prompt = "A cinematic photo of a cat" + ids = tokenizer(prompt).input_ids + assert tokenizer.decode(ids) == prompt + + +def test_bundled_tokenizer_ids_within_vocab() -> None: + tokenizer = load_bundled_qwen3_tokenizer() + ids = tokenizer( + "a very long and unusual prompt with rare tokens: zxqwv 12345", + ).input_ids + assert all(0 <= i < QWEN3_VOCAB_SIZE for i in ids) + + +def test_bundled_tokenizer_is_cached() -> None: + assert load_bundled_qwen3_tokenizer() is load_bundled_qwen3_tokenizer() + + +def test_bundled_tokenizer_has_chat_template() -> None: + # The Z-Image encoder formats prompts via apply_chat_template(), which raises if the + # tokenizer_config.json ships without a chat_template. Guard against dropping it again. + tokenizer = load_bundled_qwen3_tokenizer() + assert tokenizer.chat_template + formatted = tokenizer.apply_chat_template( + [{"role": "user", "content": "A cinematic photo of a cat"}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=True, + ) + assert "<|im_start|>" in formatted