From 3ddbf546b72e7e5fb30cb138ba71f5e59e617fa1 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Fri, 6 Nov 2026 18:38:09 +0100 Subject: [PATCH 1/4] feat(wan): load ComfyUI int8 convrot checkpoints Keep marked layers in int8 for A14B experts and TI2V-5B and refuse undecodable files at install. Choose LoRA sidecar patching per loaded expert, so int8 experts are never patched directly. --- .../03.Models/02.model-families.mdx | 2 +- .../03.Models/03.model-format-support.mdx | 2 +- .../05.Video Generation/wan-2.2.mdx | 5 + invokeai/app/invocations/wan/wan_denoise.py | 29 +- .../app/invocations/wan/wan_video_denoise.py | 8 +- .../backend/model_manager/configs/main.py | 97 +++++- .../model_manager/load/model_loaders/wan.py | 163 +++++----- invokeai/backend/quantization/int8_convrot.py | 5 +- invokeai/frontend/api/openapi.json | 2 +- invokeai/frontend/api/schema.ts | 4 +- .../invocations/test_wan_expert_swapper.py | 51 +++- .../configs/test_wan_checkpoint_config.py | 121 +++++++- .../load/state_dicts/wan_int8_convrot_keys.py | 111 +++++++ .../load/test_moss_image_loader.py | 4 +- .../load/test_wan_checkpoint_loader.py | 288 ++++++++++++++++-- .../load/test_wan_int8_convrot_keys.py | 94 ++++++ .../load/test_z_image_loader_boundaries.py | 2 +- tests/backend/pid/test_pid_int8.py | 2 +- tests/backend/qwen3_5/test_qwen3_5_loading.py | 2 +- 19 files changed, 827 insertions(+), 165 deletions(-) create mode 100644 tests/backend/model_manager/load/state_dicts/wan_int8_convrot_keys.py create mode 100644 tests/backend/model_manager/load/test_wan_int8_convrot_keys.py diff --git a/docs/src/content/docs/Users Guide/03.Models/02.model-families.mdx b/docs/src/content/docs/Users Guide/03.Models/02.model-families.mdx index 6737fa60d74..61628917426 100644 --- a/docs/src/content/docs/Users Guide/03.Models/02.model-families.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/02.model-families.mdx @@ -47,7 +47,7 @@ InvokeAI recognizes the format from the file's contents, so you install them lik - **GGUF** (`.gguf`, llama.cpp quantization such as Q4_K or Q8_0): weights stay packed in memory and are unpacked layer by layer during generation. Supported for the FLUX.1, FLUX.2, FLUX.3, Z-Image, Krea-2, Qwen-Image, Qwen-Image-2.1, Ideogram 4, ERNIE-Image and Wan transformers, and for the T5, Qwen3, Qwen3-VL, Mistral and Gemma 2 text encoders. ComfyUI-GGUF's `Q8_CR` and `Q4_CR` files, which store int8 or int4 convrot layers inside a `.gguf`, load for the Krea-2 transformer; for other models installation is refused with a message saying so. - **NF4** (bitsandbytes): 4-bit, CUDA only. - **SDNQ**: see [SDNQ Quantization](/configuration/optimization/sdnq-quantization/). -- **int8** (`int8_convrot`, from ComfyUI repackages): the weights stay int8 in memory on every device, including Apple Silicon, and need no setting to turn on. Supported for FLUX.1, FLUX.2 Klein, Z-Image, Anima, Krea-2, Qwen-Image-2.1, Ideogram 4, MiniMax H3, LTX-2 and the PiD decoders, and for the Qwen3 4B and 8B, Qwen3-VL and LTX-2 Gemma encoders. +- **int8** (`int8_convrot`, from ComfyUI repackages): the weights stay int8 in memory on every device, including Apple Silicon, and need no setting to turn on. Supported for FLUX.1, FLUX.2 Klein, Z-Image, Anima, Krea-2, Qwen-Image-2.1, Ideogram 4, MiniMax H3, LTX-2, Wan 2.2 and the PiD decoders, and for the Qwen3 4B and 8B, Qwen3-VL and LTX-2 Gemma encoders. - **int4** (`convrot_w4a4` or `int4_tensorwise`, ComfyUI repackages that usually keep some layers int8): the weights stay packed in memory at half a byte each and are unpacked per layer, on every device. This saves memory, not time: InvokeAI unpacks the weights for every step rather than running 4-bit matrix kernels, and 4-bit weights keep less detail than int8. Supported for the Krea-2 transformer. - **nvfp4** (4-bit, in ComfyUI's checkpoint format, such as Comfy-Org's builds; other nvfp4 exports such as AWQ are refused): the weights stay packed in memory and are unpacked per layer, so no special GPU is needed. Supported for the Z-Image, Krea-2, Qwen-Image and LTX-2 transformers and for the Qwen3, Qwen3-VL 4B/8B, Qwen2.5-VL and FLUX.2 Mistral encoders; the Ideogram 4 and Qwen-Image-2.1 nvfp4 transformer files are refused at install. - **fp8** (plain or `fp8_scaled`): InvokeAI keeps the weights in fp8 when [FP8 Storage](/configuration/optimization/fp8-storage/) is on, which installation does for you, or when `fp8_compute` is enabled. diff --git a/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx b/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx index e15ca2ff17e..12eb949c58b 100644 --- a/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx @@ -42,7 +42,7 @@ Others are recognized by their layer names and fail, or misbehave, only when the | [Ideogram 4](/users-guide/models/local-models/ideogram4/) | ✓ | ✓ | ✓ | ✓ | R | R | ✓ | ✓ | ✗ | inside the Diffusers folder | | [ERNIE-Image](/users-guide/models/local-models/ernie-image/) | ✓ | ✓ | L | L | L | L | ✗ | ✓ | ✗ | ✗ | | [MOSS-Image](/users-guide/models/local-models/moss-image/) | ✓ | ✓ | R | ✓ | R | R | R | ✓ | ✗ | ✗ | -| [Wan 2.2](/users-guide/video-generation/wan-22/) | ✓ | ✓ | ✓ | L | L | L | L | ✓ | ✗ | ✗ | +| [Wan 2.2](/users-guide/video-generation/wan-22/) | ✓ | ✓ | ✓ | ✓ | R | L | L | ✓ | ✗ | ✗ | | [LTX-2](/users-guide/video-generation/ltx-2/) (2.5 only) | official layout ✓, Diffusers layout ✗ | ✓ | L | ✓ | L | ✓ | ✗ | ✓ | ✗ | ✗ | | [MiniMax H3](/users-guide/video-generation/minimax-h3/) | ✓ | ✓ | L | ✓ | L | L | ✗ | ✗ | ✗ | ✗ | diff --git a/docs/src/content/docs/Users Guide/05.Video Generation/wan-2.2.mdx b/docs/src/content/docs/Users Guide/05.Video Generation/wan-2.2.mdx index f569648511d..0189b07f2b0 100644 --- a/docs/src/content/docs/Users Guide/05.Video Generation/wan-2.2.mdx +++ b/docs/src/content/docs/Users Guide/05.Video Generation/wan-2.2.mdx @@ -59,3 +59,8 @@ Higher-quality Q8_0 quantizations and full Diffusers builds are available a-la-c **On a 12 GB VRAM card**: use TI2V-5B — it covers both text-to-video and image-to-video and fits comfortably. The bundle includes it; to skip the A14B downloads, install the **Wan 2.2 TI2V-5B (Q4_K_M)** starter model on its own, which brings its VAE and the UMT5-XXL encoder. The A14B variants will technically run via aggressive offloading but are slow and prone to OOM. ::: +#### ComfyUI int8 builds + +Single-file A14B and TI2V-5B transformers in ComfyUI's `int8_tensorwise` (convrot) format load like any other single file, with the same two-expert setup for A14B. The attention and feed-forward weights stay int8 in memory, which saves memory. By default they are dequantized per layer on each forward pass; [int8 Compute](/configuration/optimization/int8-compute/) can instead run them on NVIDIA int8 tensor cores. LoRAs, including Lightning, are applied alongside the int8 weights rather than merged into them, except that DoRA layers do not run on an int8 transformer, so a DoRA LoRA fails there when it is applied. Files marked only `{"convrot": true}`, without a `format` entry, and the int4 builds are refused when you install them. + +On one Windows machine with an RTX 4090 (T2V A14B, 832x480, 17 frames), an int8 expert file is about 14.5 GB against about 28.6 GB for a bf16 expert. With both experts and the UMT5 encoder kept in the RAM cache, the process peaked at about 62 GB with int8 and about 89 GB with bf16. With `max_cache_ram_gb: 16` the int8 build still runs but reloads the experts for each clip, and the process peaked at about 44 GB, because memory use while loading goes above the cache limit. int8 mainly saves memory; how much speed it changes depends on whether the model otherwise fits in VRAM. Other systems will differ. diff --git a/invokeai/app/invocations/wan/wan_denoise.py b/invokeai/app/invocations/wan/wan_denoise.py index d6aa668637b..d6f57a0211d 100644 --- a/invokeai/app/invocations/wan/wan_denoise.py +++ b/invokeai/app/invocations/wan/wan_denoise.py @@ -44,11 +44,11 @@ from invokeai.app.invocations.model import LoRAField, WanTransformerField from invokeai.app.invocations.primitives import LatentsOutput from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, WanVariantType +from invokeai.backend.model_manager.taxonomy import BaseModelType, WanVariantType from invokeai.backend.patches.layer_patcher import LayerPatcher, PatchSpec from invokeai.backend.patches.lora_conversions.wan_lora_constants import WAN_LORA_TRANSFORMER_PREFIX from invokeai.backend.patches.model_patch_raw import ModelPatchRaw -from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes +from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes, requires_sidecar_patching from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import WanConditioningInfo @@ -189,8 +189,6 @@ def __init__( inference_dtype: torch.dtype, high_lora_factory: LoRAIteratorFactory | None = None, low_lora_factory: LoRAIteratorFactory | None = None, - high_is_quantized: bool = False, - low_is_quantized: bool = False, working_mem_bytes: int | None = None, max_resident_model_bytes: int | None = None, ) -> None: @@ -200,8 +198,6 @@ def __init__( self._inference_dtype = inference_dtype self._high_lora_factory = high_lora_factory self._low_lora_factory = low_lora_factory - self._high_is_quantized = high_is_quantized - self._low_is_quantized = low_is_quantized self._working_mem_bytes = working_mem_bytes self._max_resident_model_bytes = max_resident_model_bytes self._active_label: str | None = None @@ -257,7 +253,7 @@ def get(self, label: str) -> Any: info = self._context.models.load(model_id) supports_partial_loading = getattr(info, "supports_partial_loading", None) if self._working_mem_bytes is None or supports_partial_loading is False: - # No estimate to add it to, so a GGUF expert's per-forward dequantization transient is passed + # No estimate to add it to, so a quantized expert's per-forward dequantization transient is passed # alone; the cache floors it at `device_working_mem_gb`, so zero asks for the default. device_ctx = info.model_on_device( working_mem_bytes=peak_dequant_transient_bytes(info.model, self._inference_dtype) @@ -270,7 +266,7 @@ def get(self, label: str) -> Any: # likely the LoRA patcher), the surrounding ExitStack will eventually call # ``swapper.close() -> _release()`` to clean up, and ``_release`` needs to # see the device context to exit it. Without this early stash the expert's - # weights stay GPU-resident (8-9 GB for GGUF experts) until the cache's LRU + # weights stay GPU-resident (8-9 GB for quantized experts) until the cache's LRU # finally evicts them, masquerading as a memory leak. self._active_label = label self._active_info = info @@ -291,10 +287,10 @@ def get(self, label: str) -> Any: info.unload_from_vram(vram_bytes_to_free, keep_required_weights_in_vram=True) TorchDevice.empty_cache() - # Apply LoRA patches for this expert. GGUF transformers need sidecar - # patching since direct patching of GGMLTensors isn't supported. + # Apply LoRA patches for this expert. Quantized transformers (GGUF, or an int8 checkpoint whose Linears + # hold buffers instead of parameters) need sidecar patching since direct patching cannot write into + # their weights. Decided on the loaded module tree, per expert: a pair may mix a quantized and a dense one. lora_factory = self._high_lora_factory if label == self.HIGH else self._low_lora_factory - is_quantized = self._high_is_quantized if label == self.HIGH else self._low_is_quantized lora_ctx: Any | None = None if lora_factory is not None: lora_ctx = LayerPatcher.apply_smart_model_patches( @@ -303,7 +299,7 @@ def get(self, label: str) -> Any: prefix=WAN_LORA_TRANSFORMER_PREFIX, dtype=self._inference_dtype, cached_weights=cached_weights, - force_sidecar_patching=is_quantized, + force_sidecar_patching=requires_sidecar_patching(model, info.config.format), ) lora_ctx.__enter__() @@ -608,7 +604,6 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: # ``get_config`` doesn't touch the weights cache. high_model = self.transformer.transformer low_model = self.transformer.transformer_low_noise - low_config = context.models.get_config(low_model) if low_model is not None else None # FlowMatchEulerDiscreteScheduler stores num_train_timesteps in its config # (default 1000). Diffusers' WanPipeline computes: # boundary_timestep = boundary_ratio * num_train_timesteps @@ -622,13 +617,9 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: # "no LoRAs on the low expert" — reusing the primary list here would # silently apply high-expert-only LoRAs (e.g. Lightning high-noise # distills) to the wrong expert. - # Quantized (GGUF) experts force sidecar patching so GGMLTensor weights - # aren't touched directly. + # Quantized experts (GGUF, int8) get sidecar patching; the swapper decides per expert once loaded. high_loras = self.transformer.loras low_loras = self.transformer.loras_low_noise - high_config = context.models.get_config(high_model) - high_is_quantized = high_config.format == ModelFormat.GGUFQuantized - low_is_quantized = low_config.format == ModelFormat.GGUFQuantized if low_config is not None else False def high_lora_factory() -> Iterable[PatchSpec]: return self._lora_iterator(context, high_loras) @@ -651,8 +642,6 @@ def low_lora_factory() -> Iterable[PatchSpec]: inference_dtype=inference_dtype, high_lora_factory=high_lora_factory if high_loras else None, low_lora_factory=low_lora_factory if low_loras else None, - high_is_quantized=high_is_quantized, - low_is_quantized=low_is_quantized, working_mem_bytes=working_mem_bytes, max_resident_model_bytes=( WAN_MAX_RESIDENT_TRANSFORMER_BYTES if working_mem_bytes is not None else None diff --git a/invokeai/app/invocations/wan/wan_video_denoise.py b/invokeai/app/invocations/wan/wan_video_denoise.py index 1e1bf100205..d34575b1276 100644 --- a/invokeai/app/invocations/wan/wan_video_denoise.py +++ b/invokeai/app/invocations/wan/wan_video_denoise.py @@ -37,7 +37,7 @@ _validate_spatial_dimensions, ) from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, WanVariantType +from invokeai.backend.model_manager.taxonomy import BaseModelType, WanVariantType from invokeai.backend.patches.layer_patcher import PatchSpec from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import WanConditioningInfo @@ -274,7 +274,6 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: high_model = self.transformer.transformer low_model = self.transformer.transformer_low_noise - low_config = context.models.get_config(low_model) if low_model is not None else None num_train_timesteps = int(scheduler.config.num_train_timesteps) boundary_timestep = self.transformer.boundary_ratio * num_train_timesteps if low_model is not None else None @@ -282,9 +281,6 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: # wan_denoise: an empty low list means "no LoRAs on the low expert". high_loras = self.transformer.loras low_loras = self.transformer.loras_low_noise - high_config = context.models.get_config(high_model) - high_is_quantized = high_config.format == ModelFormat.GGUFQuantized - low_is_quantized = low_config.format == ModelFormat.GGUFQuantized if low_config is not None else False def high_lora_factory() -> Iterable[PatchSpec]: return proxy._lora_iterator(context, high_loras) @@ -304,8 +300,6 @@ def low_lora_factory() -> Iterable[PatchSpec]: inference_dtype=inference_dtype, high_lora_factory=high_lora_factory if high_loras else None, low_lora_factory=low_lora_factory if low_loras else None, - high_is_quantized=high_is_quantized, - low_is_quantized=low_is_quantized, working_mem_bytes=working_mem_bytes, max_resident_model_bytes=( WAN_MAX_RESIDENT_TRANSFORMER_BYTES if working_mem_bytes is not None else None diff --git a/invokeai/backend/model_manager/configs/main.py b/invokeai/backend/model_manager/configs/main.py index b48d2af33a7..fb29f719a82 100644 --- a/invokeai/backend/model_manager/configs/main.py +++ b/invokeai/backend/model_manager/configs/main.py @@ -16,7 +16,7 @@ is_state_dict_xlabs_controlnet, ) from invokeai.backend.ltx2 import checkpoint_layout as ltx2_layout -from invokeai.backend.model_manager.checkpoint_prefix import COMFYUI_KEY_PREFIXES +from invokeai.backend.model_manager.checkpoint_prefix import COMFYUI_KEY_PREFIXES, CheckpointPrefix from invokeai.backend.model_manager.configs.base import ( Checkpoint_Config_Base, Config_Base, @@ -2749,6 +2749,84 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - return cls(**override_fields, variant=variant) +# Top-level modules that legitimately ride along in a single-file Wan checkpoint +# without being part of the transformer. Dropping these is correct: the pipeline +# sources its VAE and text encoder from separately-wired models, and the EMA copy +# is not the weight set we generate with. +# +# ``vae`` / ``text_encoders`` / ``clip`` / ``cond_stage_model`` / ``first_stage_model`` +# are the "all-in-one" packaging convention — one file holding transformer + VAE + +# CLIP so ComfyUI's ``Load Checkpoint`` node can supply all three. The +# Phr00t/WAN2.2-14B-Rapid-AllInOne family and its ~110 GGUF conversions +# (befox/WAN2.2-14B-Rapid-AllInOne-GGUF) ship this way and loaded fine before the +# unexpected-key check existed, so refusing them would be a regression. +_BENIGN_EXTRA_MODULES = frozenset( + { + "vae", + "first_stage_model", + "text_encoders", + "cond_stage_model", + "clip", + "model_ema", + } +) + +# Trailing segments marking a merged-in LoRA's leftover adapter tensors. The main-model +# probe deliberately admits checkpoints that retain these — see +# ``configs.main._has_wan_transformer_block_weights``, which uses a *positive* structural +# test precisely so merged-LoRA mains aren't turned away — so the loader has to admit +# them too, or the two halves disagree about the same file. +# +# Kept in step with the suffix set ``LoRA_LyCORIS_Wan_Config`` matches on +# (``configs/lora.py``): kohya, PEFT, DoRA and LoKr. Matched against the *last* path +# segment rather than as a substring anywhere in the key, so a future conditioning branch +# that merely contains "lora_a" in a module name still trips the backstop instead of +# being silently discarded by it. +_MERGED_LORA_SEGMENTS = frozenset( + { + "alpha", + "dora_scale", + "lora_magnitude_vector", + "lokr_w1", + "lokr_w2", + "lokr_w1_a", + "lokr_w1_b", + "lokr_w2_a", + "lokr_w2_b", + "hada_w1_a", + "hada_w1_b", + "hada_w2_a", + "hada_w2_b", + "oft_blocks", + } +) +_MERGED_LORA_PENULTIMATE = frozenset({"lora_a", "lora_b", "lora_down", "lora_up", "lora_mid", "lora_magnitude"}) + + +def _is_benign_extra_key(key: str) -> bool: + """True if an unexpected key is packaging rather than an unsupported branch.""" + parts = key.lower().split(".") + if parts[0] in _BENIGN_EXTRA_MODULES: + return True + if parts[-1] in _MERGED_LORA_SEGMENTS: + return True + # `...to_q.lora_down.weight` — the marker is the segment before the tensor name. + return len(parts) >= 2 and parts[-2] in _MERGED_LORA_PENULTIMATE + + +def _benign_extra_wan_layers(state_dict: dict[str | int, Any]) -> tuple[str, ...]: + """Layer paths of the bundled/merged weights the Wan loader drops before it reads any quantization marker. + + Keys are compared without the ComfyUI prefix, as the loader compares them after stripping it. + """ + stripped = CheckpointPrefix.detect(state_dict).strip(state_dict) + return tuple( + key[: -len(".weight")] + for key in stripped + if isinstance(key, str) and key.endswith(".weight") and _is_benign_extra_key(key) + ) + + def _has_wan_keys(state_dict: dict[str | int, Any]) -> bool: """Check if state dict contains Wan 2.2 transformer keys. @@ -3214,7 +3292,9 @@ class Main_Checkpoint_Wan_Config(Checkpoint_Config_Base, Main_Config_Base, Confi Hugging Face repos: one ``.safetensors`` per transformer, in either the native upstream key layout or the diffusers one, optionally under a ``model.diffusion_model.`` prefix, and optionally ComfyUI ``fp8_scaled`` - quantized. The loader normalises all of those. + quantized, or ComfyUI ``int8_tensorwise`` (convrot) quantized, which stays int8 in memory. The loader + normalises all of those. An int8 weight without a decodable ``int8_tensorwise`` marker (a bare + ``{"convrot": true}``, int4) is refused at install. As with GGUF, A14B's MoE arrives as two files (one per expert); ``expert`` records which one this is so the Wan model loader invocation can pair them. @@ -3275,6 +3355,19 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - if "wan21" in normalized_identity: raise NotAMatchError("Wan 2.1 models are not supported by the Wan 2.2 loader") + # After the variant and Wan 2.1 checks, so Animate and Wan 2.1 stay unknown models rather than a refusal + # about their quantization. The loader reads both the per-tensor markers and the safetensors header, and + # drops bundled VAE/text-encoder weights before it reads either. + _raise_for_undecodable_int8_weights( + mod, + sd, + accepted_convrot_formats(accept_int4=False), + read_header=True, + discarded=_benign_extra_wan_layers(sd), + what="Wan transformer", + supported="Only ComfyUI's int8_tensorwise (convrot) build is supported.", + ) + explicit_variant = override_fields.pop("variant", None) variant = explicit_variant or _detect_wan_variant_from_state_dict(sd) if variant is None: diff --git a/invokeai/backend/model_manager/load/model_loaders/wan.py b/invokeai/backend/model_manager/load/model_loaders/wan.py index faf21e528c5..272fea23882 100644 --- a/invokeai/backend/model_manager/load/model_loaders/wan.py +++ b/invokeai/backend/model_manager/load/model_loaders/wan.py @@ -22,9 +22,10 @@ from invokeai.backend.model_manager.configs.main import ( Main_Checkpoint_Wan_Config, Main_GGUF_Wan_Config, + _is_benign_extra_key, _is_native_wan_layout, ) -from invokeai.backend.model_manager.load.load_default import ModelLoader +from invokeai.backend.model_manager.load.load_default import ModelLoader, _model_declared_skip_patterns from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.model_manager.load.model_loaders.comfyui_state_dict_utils import ( _dequantize_comfyui_fp8, @@ -38,9 +39,21 @@ ModelType, SubModelType, ) +from invokeai.backend.quantization.fp8_scaled import ( + parse_quantization_metadata, + read_safetensors_metadata, + strip_layer_path_prefix, +) from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader from invokeai.backend.quantization.gguf.utils import TORCH_COMPATIBLE_QTYPES +from invokeai.backend.quantization.int8_convrot import ( + drop_unconsumed_quantization_sidecars, + extract_int8_convrot_markers, + install_int8_convrot_layers, + reject_unmarked_int8_weights, + resolve_quantized_module_paths, +) from invokeai.backend.util.devices import TorchDevice @@ -147,6 +160,13 @@ def _load_with_variant_fallback(dtype_kwarg: dict[str, torch.dtype]) -> AnyModel ) +def _rename_wan_native_key(key: str) -> str: + """The diffusers name of one native upstream Wan key.""" + for needle, replacement in _WAN_NATIVE_TO_DIFFUSERS_RENAMES: + key = key.replace(needle, replacement) + return key + + def _convert_wan_native_to_diffusers(state_dict: dict) -> dict: """Rename native upstream Wan keys (ComfyUI / QuantStack) to diffusers names. @@ -158,10 +178,7 @@ def _convert_wan_native_to_diffusers(state_dict: dict) -> dict: if not isinstance(key, str): converted[key] = value continue - new_key = key - for needle, replacement in _WAN_NATIVE_TO_DIFFUSERS_RENAMES: - new_key = new_key.replace(needle, replacement) - converted[new_key] = value + converted[_rename_wan_native_key(key)] = value return converted @@ -190,71 +207,6 @@ def _unwrap_unquantized_to_compute_dtype(state_dict: dict) -> dict: return unwrapped -# Top-level modules that legitimately ride along in a single-file Wan checkpoint -# without being part of the transformer. Dropping these is correct: the pipeline -# sources its VAE and text encoder from separately-wired models, and the EMA copy -# is not the weight set we generate with. -# -# ``vae`` / ``text_encoders`` / ``clip`` / ``cond_stage_model`` / ``first_stage_model`` -# are the "all-in-one" packaging convention — one file holding transformer + VAE + -# CLIP so ComfyUI's ``Load Checkpoint`` node can supply all three. The -# Phr00t/WAN2.2-14B-Rapid-AllInOne family and its ~110 GGUF conversions -# (befox/WAN2.2-14B-Rapid-AllInOne-GGUF) ship this way and loaded fine before the -# unexpected-key check existed, so refusing them would be a regression. -_BENIGN_EXTRA_MODULES = frozenset( - { - "vae", - "first_stage_model", - "text_encoders", - "cond_stage_model", - "clip", - "model_ema", - } -) - -# Trailing segments marking a merged-in LoRA's leftover adapter tensors. The main-model -# probe deliberately admits checkpoints that retain these — see -# ``configs.main._has_wan_transformer_block_weights``, which uses a *positive* structural -# test precisely so merged-LoRA mains aren't turned away — so the loader has to admit -# them too, or the two halves disagree about the same file. -# -# Kept in step with the suffix set ``LoRA_LyCORIS_Wan_Config`` matches on -# (``configs/lora.py``): kohya, PEFT, DoRA and LoKr. Matched against the *last* path -# segment rather than as a substring anywhere in the key, so a future conditioning branch -# that merely contains "lora_a" in a module name still trips the backstop instead of -# being silently discarded by it. -_MERGED_LORA_SEGMENTS = frozenset( - { - "alpha", - "dora_scale", - "lora_magnitude_vector", - "lokr_w1", - "lokr_w2", - "lokr_w1_a", - "lokr_w1_b", - "lokr_w2_a", - "lokr_w2_b", - "hada_w1_a", - "hada_w1_b", - "hada_w2_a", - "hada_w2_b", - "oft_blocks", - } -) -_MERGED_LORA_PENULTIMATE = frozenset({"lora_a", "lora_b", "lora_down", "lora_up", "lora_mid", "lora_magnitude"}) - - -def _is_benign_extra_key(key: str) -> bool: - """True if an unexpected key is packaging rather than an unsupported branch.""" - parts = key.lower().split(".") - if parts[0] in _BENIGN_EXTRA_MODULES: - return True - if parts[-1] in _MERGED_LORA_SEGMENTS: - return True - # `...to_q.lora_down.weight` — the marker is the segment before the tensor name. - return len(parts) >= 2 and parts[-2] in _MERGED_LORA_PENULTIMATE - - def _drop_benign_extra_keys(sd: dict, source: str, logger: Any) -> None: """Remove packaging weights the transformer has no use for, in place. @@ -469,7 +421,9 @@ class WanCheckpointModel(ModelLoader): ``model.diffusion_model.`` key prefix, the native upstream key layout as well as the diffusers one, ComfyUI ``fp8_scaled`` weights (dequantized to the compute dtype at load time), and plain ``float8_e4m3fn`` weights with no - scales (cast the same way as any other non-bf16 dtype). + scales (cast the same way as any other non-bf16 dtype). ComfyUI ``int8_tensorwise`` (convrot) builds keep + their marked layers in int8 as ``Int8ConvrotLinear`` and dequantize them per forward; layers the model + reads outside a matmul (``condition_embedder``) are widened at load instead. Like the GGUF loader, one file is one expert; A14B pairing happens at the WanModelLoaderInvocation layer. @@ -508,20 +462,35 @@ def _load_from_singlefile(self, config: Main_Checkpoint_Wan_Config) -> AnyModel: sd = CheckpointPrefix.detect(sd).strip(sd) _drop_benign_extra_keys(sd, "Wan checkpoint", logger) - dequantized = _dequantize_comfyui_fp8(sd, model_dtype, f"Wan checkpoint {model_path.name}") - if dequantized > 0: - logger.info(f"Dequantized {dequantized} ComfyUI-quantized weights") - # Drop the scale tensors themselves — they've been folded into the weights - # above and are not parameters of WanTransformer3DModel. load_state_dict - # runs with strict=False and would ignore them anyway, but dropping them - # here keeps the dtype cast and the RAM-cache reservation below honest. - _strip_quantization_metadata(sd) + # Markers and header entries name their layers in the file's native key space, so they are read before + # any rename. The refusal sits outside the "found markers" branch on purpose: an int8 weight no marker + # claims must be refused, not cast to the compute dtype as raw codes. + header_layers = strip_layer_path_prefix( + parse_quantization_metadata(read_safetensors_metadata(model_path, logger)) + ) + int8_markers = extract_int8_convrot_markers(sd, header_layers=header_layers) + reject_unmarked_int8_weights(sd, int8_markers, f"Wan 2.2 ({model_path.name})", header_layers=header_layers) + + dequantized = 0 + if int8_markers: + sd = drop_unconsumed_quantization_sidecars(sd) + else: + dequantized = _dequantize_comfyui_fp8(sd, model_dtype, f"Wan checkpoint {model_path.name}") + if dequantized > 0: + logger.info(f"Dequantized {dequantized} ComfyUI-quantized weights") + # Drop the scale tensors themselves — they've been folded into the weights + # above and are not parameters of WanTransformer3DModel. load_state_dict + # runs with strict=False and would ignore them anyway, but dropping them + # here keeps the dtype cast and the RAM-cache reservation below honest. + _strip_quantization_metadata(sd) # Community releases ship the native upstream Wan key layout # (text_embedding.0, self_attn/cross_attn, ffn.0/2, head.head, ...); # diffusers' WanTransformer3DModel expects condition_embedder.*, # attn1/attn2, ffn.net.*, proj_out. Convert if needed. + key_map: dict[str, str] = {} if _is_native_wan_layout(sd): + key_map = {key: _rename_wan_native_key(key) for key in sd if isinstance(key, str)} sd = _convert_wan_native_to_diffusers(sd) model_config = _build_wan_transformer_config(sd, source="checkpoint state dict") @@ -529,15 +498,33 @@ def _load_from_singlefile(self, config: Main_Checkpoint_Wan_Config) -> AnyModel: with accelerate.init_empty_weights(): model = WanTransformer3DModel(**model_config) - # Cast every float tensor to the compute dtype. Dequantized fp8_scaled - # weights are already there; this catches plain fp16/fp32/fp8 checkpoints - # and makes the cache reservation below reflect the post-cast sizes. - for key in list(sd.keys()): - if sd[key].is_floating_point(): - sd[key] = sd[key].to(model_dtype) - - new_sd_size = sum(t.nelement() * t.element_size() for t in sd.values()) - self._ram_cache.make_room(new_sd_size) + if int8_markers: + # Casts only the dense tensors: the blanket cast below would also cast the fp32 scales. The skip + # patterns match diffusers names, hence after the rename; `condition_embedder` reads its dtype from + # a parameter, which an int8 layer does not have, so it is widened here. + quantized = install_int8_convrot_layers( + model, + sd, + resolve_quantized_module_paths(int8_markers, key_map), + model_dtype, + architecture="Wan 2.2", + reserve=self._ram_cache.make_room, + skip_patterns=_model_declared_skip_patterns(model), + ) + logger.info( + f"Wan 2.2: kept {len(quantized)} of {len(int8_markers)} layer(s) in int8 " + "(int8_tensorwise checkpoint, dequantized per forward)" + ) + else: + # Cast every float tensor to the compute dtype. Dequantized fp8_scaled + # weights are already there; this catches plain fp16/fp32/fp8 checkpoints + # and makes the cache reservation below reflect the post-cast sizes. + for key in list(sd.keys()): + if sd[key].is_floating_point(): + sd[key] = sd[key].to(model_dtype) + + new_sd_size = sum(t.nelement() * t.element_size() for t in sd.values()) + self._ram_cache.make_room(new_sd_size) incompatible_keys = model.load_state_dict(sd, strict=False, assign=True) _raise_for_incompatible_keys(incompatible_keys, source="Wan checkpoint") diff --git a/invokeai/backend/quantization/int8_convrot.py b/invokeai/backend/quantization/int8_convrot.py index 3ceb401bb94..1aa5f0a01d1 100644 --- a/invokeai/backend/quantization/int8_convrot.py +++ b/invokeai/backend/quantization/int8_convrot.py @@ -805,8 +805,9 @@ def reject_unmarked_int8_weights( else f" A build that names its quantization only in the safetensors header is not supported for {architecture}." ) raise ValueError( - f"{architecture} checkpoint has {len(orphans)} int8 weight(s) with no `comfy_quant` marker, " - f"e.g. {orphans[:3]}. Loading them would produce a model that runs and generates noise.{unread_header}" + f"{architecture} checkpoint has {len(orphans)} int8 weight(s) with no readable int8_tensorwise marker, " + f"e.g. {orphans[:3]}. Loading them would produce a model that runs and generates noise.{unread_header} " + "Use another build of this model." ) diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json index c5abef0f9f7..fa5a2e7f15f 100644 --- a/invokeai/frontend/api/openapi.json +++ b/invokeai/frontend/api/openapi.json @@ -78973,7 +78973,7 @@ "expert" ], "title": "Main_Checkpoint_Wan_Config", - "description": "Model config for single-file Wan 2.2 transformer checkpoints (safetensors).\n\nThis is the format the community ships on CivitAI and in ComfyUI-oriented\nHugging Face repos: one ``.safetensors`` per transformer, in either the native\nupstream key layout or the diffusers one, optionally under a\n``model.diffusion_model.`` prefix, and optionally ComfyUI ``fp8_scaled``\nquantized. The loader normalises all of those.\n\nAs with GGUF, A14B's MoE arrives as two files (one per expert); ``expert``\nrecords which one this is so the Wan model loader invocation can pair them.\nTI2V-5B is single-transformer and stores ``expert='none'``." + "description": "Model config for single-file Wan 2.2 transformer checkpoints (safetensors).\n\nThis is the format the community ships on CivitAI and in ComfyUI-oriented\nHugging Face repos: one ``.safetensors`` per transformer, in either the native\nupstream key layout or the diffusers one, optionally under a\n``model.diffusion_model.`` prefix, and optionally ComfyUI ``fp8_scaled``\nquantized, or ComfyUI ``int8_tensorwise`` (convrot) quantized, which stays int8 in memory. The loader\nnormalises all of those. An int8 weight without a decodable ``int8_tensorwise`` marker (a bare\n``{\"convrot\": true}``, int4) is refused at install.\n\nAs with GGUF, A14B's MoE arrives as two files (one per expert); ``expert``\nrecords which one this is so the Wan model loader invocation can pair them.\nTI2V-5B is single-transformer and stores ``expert='none'``." }, "Main_Checkpoint_ZImage_Config": { "properties": { diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts index 4ef1a6ab029..25608ea26d5 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -32803,7 +32803,9 @@ export type components = { * Hugging Face repos: one ``.safetensors`` per transformer, in either the native * upstream key layout or the diffusers one, optionally under a * ``model.diffusion_model.`` prefix, and optionally ComfyUI ``fp8_scaled`` - * quantized. The loader normalises all of those. + * quantized, or ComfyUI ``int8_tensorwise`` (convrot) quantized, which stays int8 in memory. The loader + * normalises all of those. An int8 weight without a decodable ``int8_tensorwise`` marker (a bare + * ``{"convrot": true}``, int4) is refused at install. * * As with GGUF, A14B's MoE arrives as two files (one per expert); ``expert`` * records which one this is so the Wan model loader invocation can pair them. diff --git a/tests/app/invocations/test_wan_expert_swapper.py b/tests/app/invocations/test_wan_expert_swapper.py index 2ad62ffa944..360d2307212 100644 --- a/tests/app/invocations/test_wan_expert_swapper.py +++ b/tests/app/invocations/test_wan_expert_swapper.py @@ -14,6 +14,7 @@ LoRA factory at each step. """ +from types import SimpleNamespace from typing import Iterable, Tuple from unittest.mock import patch @@ -22,7 +23,9 @@ import torch.nn as nn from invokeai.app.invocations.wan.wan_denoise import _ExpertSwapper +from invokeai.backend.model_manager.taxonomy import ModelFormat from invokeai.backend.patches.model_patch_raw import ModelPatchRaw +from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear class _FakeModelOnDevice: @@ -72,10 +75,13 @@ class _FakeInfo: the device on swap. ``weight_bytes`` and ``resident_weight_bytes`` are the public properties the swapper sizes its unload request from.""" - def __init__(self, label: str, model: nn.Module, log: list[str]) -> None: + def __init__( + self, label: str, model: nn.Module, log: list[str], model_format: ModelFormat = ModelFormat.Checkpoint + ) -> None: self._label = label self._model = model self._log = log + self.config = SimpleNamespace(format=model_format) self._cache_record = _FakeCacheRecord(_FakeCachedModel(label, log)) self.weight_bytes = 0 self.resident_weight_bytes = 0 @@ -263,11 +269,18 @@ def test_lifecycle_dual_expert_swap(): assert calls[1]["model"] is low_nn -def test_quantized_flag_forwards_to_sidecar(): - """GGUF (quantized) experts must request sidecar patching.""" +def _int8_expert() -> nn.Module: + """A checkpoint-format expert whose Linear holds int8 buffers instead of parameters.""" + return nn.Sequential( + Int8ConvrotLinear(torch.zeros(4, 4, dtype=torch.int8), torch.ones(4, 1), convrot=False), + ) + + +def _sidecar_decisions(high: tuple[nn.Module, ModelFormat], low: tuple[nn.Module, ModelFormat]) -> list: + """What the swapper tells the LoRA patcher for each expert of a pair, in swap order.""" log: list[str] = [] - high_nn = nn.Linear(1, 1) - ctx = _FakeContext({"high": _FakeInfo("HIGH", high_nn, log)}, log) + infos = {"high": _FakeInfo("HIGH", high[0], log, high[1]), "low": _FakeInfo("LOW", low[0], log, low[1])} + ctx = _FakeContext(infos, log) stub, calls = _stub_lora_context_manager(log) with patch( @@ -277,15 +290,37 @@ def test_quantized_flag_forwards_to_sidecar(): swapper = _ExpertSwapper( context=ctx, high_model="high", - low_model=None, + low_model="low", inference_dtype=torch.bfloat16, high_lora_factory=_make_factory(log, "HIGH"), - high_is_quantized=True, + low_lora_factory=_make_factory(log, "LOW"), ) swapper.get(_ExpertSwapper.HIGH) + swapper.get(_ExpertSwapper.LOW) swapper.close() + return [call["force_sidecar_patching"] for call in calls] + + +GGUF = (nn.Linear(1, 1), ModelFormat.GGUFQuantized) +DENSE = (nn.Linear(1, 1), ModelFormat.Checkpoint) + + +def test_gguf_experts_get_sidecar_patching(): + assert _sidecar_decisions(GGUF, (nn.Linear(1, 1), ModelFormat.GGUFQuantized)) == [True, True] + + +def test_a_mixed_gguf_and_dense_pair_is_decided_per_expert(): + assert _sidecar_decisions(GGUF, DENSE) == [True, False] + assert _sidecar_decisions(DENSE, GGUF) == [False, True] + + +def test_an_int8_checkpoint_expert_gets_sidecar_patching_and_a_dense_one_does_not(): + """Direct patching cannot write into int8 buffers, and a `Checkpoint` config does not say which experts are + int8: the decision is made on the loaded module tree, per expert, so a mixed pair works in either order.""" + int8 = (_int8_expert(), ModelFormat.Checkpoint) - assert calls[0]["force_sidecar_patching"] is True + assert _sidecar_decisions(int8, DENSE) == [True, False] + assert _sidecar_decisions(DENSE, (_int8_expert(), ModelFormat.Checkpoint)) == [False, True] def test_no_lora_factory_skips_lora_context(): diff --git a/tests/backend/model_manager/configs/test_wan_checkpoint_config.py b/tests/backend/model_manager/configs/test_wan_checkpoint_config.py index ed18f655cbe..72fb55bc6a2 100644 --- a/tests/backend/model_manager/configs/test_wan_checkpoint_config.py +++ b/tests/backend/model_manager/configs/test_wan_checkpoint_config.py @@ -6,14 +6,16 @@ a matching config class. """ +import json from pathlib import Path from unittest.mock import MagicMock import gguf import pytest import torch +from safetensors.torch import save_file -from invokeai.backend.model_manager.configs.identification_utils import NotAMatchError +from invokeai.backend.model_manager.configs.identification_utils import InvalidMatchError, NotAMatchError from invokeai.backend.model_manager.configs.main import ( Main_Checkpoint_Wan_Config, _detect_wan_expert, @@ -22,6 +24,7 @@ from invokeai.backend.model_manager.model_on_disk import ModelOnDisk from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, WanVariantType from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor +from tests.fixtures.quantized_payloads import comfy_quant_marker A14B_DIM = 5120 TI2V_DIM = 3072 @@ -283,6 +286,122 @@ def test_channel_count_must_agree_with_transformer_width( _probe(tmp_path, f"wan22-{label.replace(' ', '-')}.safetensors", _native_sd(in_channels, dim=dim)) +INT8_MARKER = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256, "per_row": True} +INT8_LAYERS = ["blocks.0.self_attn.q", "blocks.0.cross_attn.q", "blocks.0.ffn.0"] +SUPPORTED = r"Only ComfyUI's int8_tensorwise \(convrot\) build is supported\." + + +def _probe_int8( + tmp_path: Path, + sd: dict, + *, + marker: dict | None = INT8_MARKER, + prefix: str = "", + header_prefix: str | None = None, +) -> Main_Checkpoint_Wan_Config: + """Probe ``sd`` with its ``INT8_LAYERS`` turned into int8 weights. + + Identification sees dtypes but not tensor bytes, so the markers are read from the file itself: a small real + safetensors file holds them (and only them), while the expanded ``sd`` stands in for the weights. With + ``header_prefix`` the layers are declared in the file's ``_quantization_metadata`` instead, under that prefix. + """ + sd = dict(sd) + for layer in INT8_LAYERS: + key = f"{prefix}{layer}.weight" + sd[key] = _t(*sd[key].shape, dtype=torch.int8) + path = tmp_path / "wan2.2_t2v_high_noise_14B_int8_convrot.safetensors" + on_disk: dict[str, torch.Tensor] = {"unrelated": torch.zeros(1)} + metadata = None + if header_prefix is not None: + entry = {key: value for key, value in INT8_MARKER.items() if key != "per_row"} + layers = {f"{header_prefix}{layer}": entry for layer in INT8_LAYERS} + metadata = {"_quantization_metadata": json.dumps({"format_version": "1.0", "layers": layers})} + elif marker is not None: + on_disk = {f"{prefix}{layer}.comfy_quant": comfy_quant_marker(marker) for layer in INT8_LAYERS} + save_file(on_disk, path, metadata=metadata) + if header_prefix is None: + mod = _make_mod(path, sd) + else: + # The real header, read by the real `ModelOnDisk`; only the (huge) weights are stood in for. + mod = ModelOnDisk(path) + mod.load_state_dict = lambda _path=None: sd + return Main_Checkpoint_Wan_Config.from_model_on_disk(mod, _build_overrides(path, path.stem)) + + +class TestInt8Convrot: + @pytest.mark.parametrize("prefix", ["", "model.diffusion_model."]) + def test_accepts_the_comfyui_int8_build(self, tmp_path: Path, prefix: str) -> None: + config = _probe_int8(tmp_path, _native_sd(16, prefix=prefix), prefix=prefix) + + assert config.variant == WanVariantType.T2V_A14B + assert config.expert == "high" + + @pytest.mark.parametrize( + ("prefix", "header_prefix"), + [("", ""), ("model.diffusion_model.", ""), ("model.diffusion_model.", "model.diffusion_model.")], + ids=["bare", "prefixed-tensors", "prefixed-tensors-and-header"], + ) + def test_accepts_a_build_that_declares_its_layers_only_in_the_header( + self, tmp_path: Path, prefix: str, header_prefix: str + ) -> None: + config = _probe_int8(tmp_path, _native_sd(16, prefix=prefix), prefix=prefix, header_prefix=header_prefix) + + assert config.variant == WanVariantType.T2V_A14B + + @pytest.mark.parametrize( + ("marker", "refusal"), + [ + pytest.param( + None, + r"3 int8 weight\(s\) in this Wan transformer carry no readable marker of a format it loads " + rf"\('int8_tensorwise'; e\.g\. 'blocks\.0\.[a-z_.0-9]+\.weight'\).*{SUPPORTED}", + id="no marker", + ), + pytest.param( + {"convrot": True}, + r"in this Wan transformer carry no readable marker of a format it loads.*" + SUPPORTED, + id="convrot flag without format", + ), + pytest.param( + {"format": "int4_tensorwise", "convrot": True}, + r"3 int8 weight\(s\) in this Wan transformer are marked 'int4_tensorwise'.*" + SUPPORTED, + id="int4", + ), + ], + ) + def test_refuses_int8_weights_the_loader_cannot_decode( + self, tmp_path: Path, marker: dict | None, refusal: str + ) -> None: + """The loader refuses these, so the installer has to say so instead of registering a model that fails at + its first render. `InvalidMatchError` keeps the reason visible rather than falling back to unknown.""" + with pytest.raises(InvalidMatchError, match=refusal): + _probe_int8(tmp_path, _native_sd(16), marker=marker) + + @pytest.mark.parametrize("prefix", ["", "model.diffusion_model."]) + def test_unmarked_int8_weights_in_a_bundled_text_encoder_do_not_block_the_install( + self, tmp_path: Path, prefix: str + ) -> None: + """An all-in-one file bundles a text encoder the loader drops before it reads any marker, so an int8 + weight there is not the transformer's business.""" + sd = _native_sd(16, prefix=prefix) + sd[f"{prefix}text_encoders.umt5xxl.encoder.block.0.layer.0.SelfAttention.q.weight"] = _t( + 64, 64, dtype=torch.int8 + ) + + config = _probe(tmp_path, "Wan2.2-T2V-A14B-high_noise-aio.safetensors", sd) + + assert config.variant == WanVariantType.T2V_A14B + + def test_animate_stays_an_unknown_model_rather_than_a_quantization_refusal(self, tmp_path: Path) -> None: + sd = _native_sd(36) + sd["face_adapter.fuser_blocks.0.k_norm.weight"] = _t(A14B_DIM) + sd["motion_encoder.dec.direction.weight"] = _t(512, 512) + sd["img_emb.proj.0.bias"] = _t(1280) + + with pytest.raises(NotAMatchError, match="Animate"): + _probe_int8(tmp_path, sd, marker=None) + + class TestWan21Marker: def test_clean_wan_2_2_state_dicts_have_no_marker(self) -> None: assert _find_wan_2_1_marker(_native_sd(16)) is None diff --git a/tests/backend/model_manager/load/state_dicts/wan_int8_convrot_keys.py b/tests/backend/model_manager/load/state_dicts/wan_int8_convrot_keys.py new file mode 100644 index 00000000000..a974f8bc163 --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/wan_int8_convrot_keys.py @@ -0,0 +1,111 @@ +"""Representative key layout of the community `int8_tensorwise` Wan 2.2 T2V-14B high-noise build. + +Captured on 2026-11-06 by range-reading the safetensors header (no tensor data) of +`FartingBackwards/wan2.2-t2v-14B-int8-convrot/wan2.2_t2v_high_noise_14B_int8_convrot.safetensors`. + +Subsetting rule: `blocks` index 0, plus every key outside a stack. Values are `(shape, dtype)`. +The full file has 40 blocks and 1895 tensors. + +These are the native Wan names (`blocks.0.self_attn.q`, `ffn.0`, `text_embedding.0`), so the layout reaches the +module tree only through `_convert_wan_native_to_diffusers`. Only the self/cross attention projections and the two +FFN linears are int8; the `condition_embedder` linears (`text_embedding`, `time_embedding`, `time_projection`), +`patch_embedding` and `head` are dense F16. + +The build declares its scheme twice: a `.comfy_quant` marker beside every quantized weight and an identical +`_quantization_metadata` header entry per layer. `convrot` and the group size are in both. +""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "blocks.0.cross_attn.k.bias": ([5120], "F16"), + "blocks.0.cross_attn.k.comfy_quant": ([89], "U8"), + "blocks.0.cross_attn.k.weight": ([5120, 5120], "I8"), + "blocks.0.cross_attn.k.weight_scale": ([5120, 1], "F32"), + "blocks.0.cross_attn.norm_k.weight": ([5120], "F16"), + "blocks.0.cross_attn.norm_q.weight": ([5120], "F16"), + "blocks.0.cross_attn.o.bias": ([5120], "F16"), + "blocks.0.cross_attn.o.comfy_quant": ([89], "U8"), + "blocks.0.cross_attn.o.weight": ([5120, 5120], "I8"), + "blocks.0.cross_attn.o.weight_scale": ([5120, 1], "F32"), + "blocks.0.cross_attn.q.bias": ([5120], "F16"), + "blocks.0.cross_attn.q.comfy_quant": ([89], "U8"), + "blocks.0.cross_attn.q.weight": ([5120, 5120], "I8"), + "blocks.0.cross_attn.q.weight_scale": ([5120, 1], "F32"), + "blocks.0.cross_attn.v.bias": ([5120], "F16"), + "blocks.0.cross_attn.v.comfy_quant": ([89], "U8"), + "blocks.0.cross_attn.v.weight": ([5120, 5120], "I8"), + "blocks.0.cross_attn.v.weight_scale": ([5120, 1], "F32"), + "blocks.0.ffn.0.bias": ([13824], "F16"), + "blocks.0.ffn.0.comfy_quant": ([89], "U8"), + "blocks.0.ffn.0.weight": ([13824, 5120], "I8"), + "blocks.0.ffn.0.weight_scale": ([13824, 1], "F32"), + "blocks.0.ffn.2.bias": ([5120], "F16"), + "blocks.0.ffn.2.comfy_quant": ([89], "U8"), + "blocks.0.ffn.2.weight": ([5120, 13824], "I8"), + "blocks.0.ffn.2.weight_scale": ([5120, 1], "F32"), + "blocks.0.modulation": ([1, 6, 5120], "F16"), + "blocks.0.norm3.bias": ([5120], "F16"), + "blocks.0.norm3.weight": ([5120], "F16"), + "blocks.0.self_attn.k.bias": ([5120], "F16"), + "blocks.0.self_attn.k.comfy_quant": ([89], "U8"), + "blocks.0.self_attn.k.weight": ([5120, 5120], "I8"), + "blocks.0.self_attn.k.weight_scale": ([5120, 1], "F32"), + "blocks.0.self_attn.norm_k.weight": ([5120], "F16"), + "blocks.0.self_attn.norm_q.weight": ([5120], "F16"), + "blocks.0.self_attn.o.bias": ([5120], "F16"), + "blocks.0.self_attn.o.comfy_quant": ([89], "U8"), + "blocks.0.self_attn.o.weight": ([5120, 5120], "I8"), + "blocks.0.self_attn.o.weight_scale": ([5120, 1], "F32"), + "blocks.0.self_attn.q.bias": ([5120], "F16"), + "blocks.0.self_attn.q.comfy_quant": ([89], "U8"), + "blocks.0.self_attn.q.weight": ([5120, 5120], "I8"), + "blocks.0.self_attn.q.weight_scale": ([5120, 1], "F32"), + "blocks.0.self_attn.v.bias": ([5120], "F16"), + "blocks.0.self_attn.v.comfy_quant": ([89], "U8"), + "blocks.0.self_attn.v.weight": ([5120, 5120], "I8"), + "blocks.0.self_attn.v.weight_scale": ([5120, 1], "F32"), + "head.head.bias": ([64], "F16"), + "head.head.weight": ([64, 5120], "F16"), + "head.modulation": ([1, 2, 5120], "F16"), + "patch_embedding.bias": ([5120], "F16"), + "patch_embedding.weight": ([5120, 16, 1, 2, 2], "F16"), + "text_embedding.0.bias": ([5120], "F16"), + "text_embedding.0.weight": ([5120, 4096], "F16"), + "text_embedding.2.bias": ([5120], "F16"), + "text_embedding.2.weight": ([5120, 5120], "F16"), + "time_embedding.0.bias": ([5120], "F16"), + "time_embedding.0.weight": ([5120, 256], "F16"), + "time_embedding.2.bias": ([5120], "F16"), + "time_embedding.2.weight": ([5120, 5120], "F16"), + "time_projection.1.bias": ([30720], "F16"), + "time_projection.1.weight": ([30720, 5120], "F16"), +} + +# The per-layer `.comfy_quant` markers, decoded from a blob read from the file (all ten blocks.0 layers share it). +MARKER: dict[str, object] = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256, "per_row": True} + +markers: dict[str, dict[str, object]] = { + "blocks.0.cross_attn.k": MARKER, + "blocks.0.cross_attn.o": MARKER, + "blocks.0.cross_attn.q": MARKER, + "blocks.0.cross_attn.v": MARKER, + "blocks.0.ffn.0": MARKER, + "blocks.0.ffn.2": MARKER, + "blocks.0.self_attn.k": MARKER, + "blocks.0.self_attn.o": MARKER, + "blocks.0.self_attn.q": MARKER, + "blocks.0.self_attn.v": MARKER, +} + +# The `_quantization_metadata` header entries for the same block. +header_layers: dict[str, dict[str, object]] = { + "blocks.0.cross_attn.k": {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}, + "blocks.0.cross_attn.o": {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}, + "blocks.0.cross_attn.q": {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}, + "blocks.0.cross_attn.v": {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}, + "blocks.0.ffn.0": {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}, + "blocks.0.ffn.2": {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}, + "blocks.0.self_attn.k": {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}, + "blocks.0.self_attn.o": {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}, + "blocks.0.self_attn.q": {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}, + "blocks.0.self_attn.v": {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}, +} diff --git a/tests/backend/model_manager/load/test_moss_image_loader.py b/tests/backend/model_manager/load/test_moss_image_loader.py index f7ebfced47a..5c92986f1fd 100644 --- a/tests/backend/model_manager/load/test_moss_image_loader.py +++ b/tests/backend/model_manager/load/test_moss_image_loader.py @@ -516,7 +516,9 @@ def test_an_int8_weight_without_its_marker_is_refused_not_cast_as_codes( path = tmp_path / "moss-int8-unmarked.safetensors" save_file(stored, str(path)) - with pytest.raises(ValueError, match=r"no `comfy_quant` marker, e\.g\. \['layers\.0\.proj\.o\.weight'\]"): + with pytest.raises( + ValueError, match=r"no readable int8_tensorwise marker, e\.g\. \['layers\.0\.proj\.o\.weight'\]" + ): self._load(monkeypatch, path) def test_fp8_storage_leaves_the_int8_layers_alone_and_still_runs( diff --git a/tests/backend/model_manager/load/test_wan_checkpoint_loader.py b/tests/backend/model_manager/load/test_wan_checkpoint_loader.py index c688a614d05..0befa57438a 100644 --- a/tests/backend/model_manager/load/test_wan_checkpoint_loader.py +++ b/tests/backend/model_manager/load/test_wan_checkpoint_loader.py @@ -6,6 +6,7 @@ expectation. """ +import json from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -20,7 +21,13 @@ _build_wan_transformer_config, ) from invokeai.backend.model_manager.taxonomy import SubModelType, WanVariantType -from tests.fixtures.quantized_payloads import comfy_quant_marker, nvfp4_signed_tensors, quantize_convrot +from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear +from tests.fixtures.quantized_payloads import ( + comfy_quant_marker, + nvfp4_signed_tensors, + quantize_convrot, + quantize_scaled_fp8, +) # A structurally faithful but tiny Wan transformer. attention_head_dim must stay # at 128 — the loader derives num_attention_heads as inner_dim // 128, matching @@ -90,7 +97,12 @@ def _make_loader() -> WanCheckpointModel: return loader -def _load(path: Path, variant: WanVariantType = WanVariantType.T2V_A14B): +def _load( + path: Path, + variant: WanVariantType = WanVariantType.T2V_A14B, + loader: WanCheckpointModel | None = None, + dtype: torch.dtype = torch.bfloat16, +): config = MagicMock() config.path = str(path) config.variant = variant @@ -99,10 +111,10 @@ def _load(path: Path, variant: WanVariantType = WanVariantType.T2V_A14B): patch("invokeai.backend.model_manager.load.model_loaders.wan.TorchDevice.choose_torch_device"), patch( "invokeai.backend.model_manager.load.model_loaders.wan.TorchDevice.choose_bfloat16_safe_dtype", - return_value=torch.bfloat16, + return_value=dtype, ), ): - return _make_loader()._load_from_singlefile(config) + return (loader or _make_loader())._load_from_singlefile(config) class TestArchitectureInference: @@ -198,31 +210,6 @@ def test_fp8_scaled_is_dequantized_and_scales_are_dropped(self, tmp_path: Path) assert loaded[target].dtype == torch.bfloat16 assert torch.allclose(loaded[target].float(), torch.full_like(loaded[target].float(), 2.0)) - def test_an_int8_convrot_checkpoint_is_refused_rather_than_folded_unrotated(self, tmp_path: Path) -> None: - """Wan has no int8 branch, and `int8_tensorwise` shares the fp8 key layout, so the fold used - to apply the scale and skip the inverse rotation. That produces a weight of the right shape, - dtype and magnitude bearing no relation to the stored one -- a model that loads and - generates noise. `test_int8_through_the_fp8_fold.py` measures how little relation; how - little depends on the rotation width, so the number is kept where the geometry is fixed. - Pinned at the loader as well, because Wan is one of the two that reach the shared fold. - """ - reference = _tiny_model() - sd = {k: v.clone() for k, v in reference.state_dict().items()} - - target = "blocks.0.attn1.to_q.weight" - payload = quantize_convrot(sd[target].float(), group_size=64) - sd[target] = payload.codes - sd["blocks.0.attn1.to_q.weight_scale"] = payload.scale - sd["blocks.0.attn1.to_q.comfy_quant"] = comfy_quant_marker( - {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 64} - ) - path = tmp_path / "Wan2.2-A14B-HighNoise-int8_convrot.safetensors" - save_file(sd, path) - - message = r"Wan checkpoint Wan2\.2-A14B-HighNoise-int8_convrot\.safetensors.*quantized with convrot" - with pytest.raises(ValueError, match=message): - _load(path) - def test_an_nvfp4_checkpoint_is_refused_rather_than_folded_over_its_packed_codes(self, tmp_path: Path) -> None: """Wan reads scaled fp8 but never calls `pop_nvfp4_layers`, and the shared fold has no dtype gate, so an nvfp4 layer was multiplied by its block-scale grid and logged as dequantized. @@ -391,3 +378,246 @@ def test_only_the_transformer_submodel_is_served(self, submodel) -> None: def test_wrong_config_class_is_rejected(self) -> None: with pytest.raises(TypeError, match="Main_Checkpoint_Wan_Config"): _make_loader()._load_model(MagicMock(), SubModelType.Transformer) + + +# --- ComfyUI int8_tensorwise (convrot) ------------------------------------------------------------------------- + +GROUP = 16 # a power of 4 that divides every tiny input width, 32 (text_dim) up to 256 (freq_dim) +MARKER = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": GROUP, "per_row": True} + +ATTENTION = [ + f"blocks.{block}.{attn}.{proj}" + for block in range(2) + for attn in ("attn1", "attn2") + for proj in ("to_q", "to_k", "to_v", "to_out.0") +] +FFN = [f"blocks.{block}.ffn.net.{index}{suffix}" for block in range(2) for index, suffix in ((0, ".proj"), (2, ""))] +BLOCK_LAYERS = ATTENTION + FFN +# `WanTimeTextImageEmbedding.forward` reads `next(iter(time_embedder.parameters())).dtype`, so an +# `Int8ConvrotLinear` there (buffers only) would raise StopIteration at the first forward. +EMBEDDERS = [ + "condition_embedder.text_embedder.linear_1", + "condition_embedder.text_embedder.linear_2", + "condition_embedder.time_embedder.linear_1", + "condition_embedder.time_embedder.linear_2", + "condition_embedder.time_proj", +] + + +def _int8_checkpoint( + layers: list[str], + *, + marker: dict | None = MARKER, + dense_dtype: torch.dtype = torch.float32, +) -> tuple[torch.nn.Module, dict[str, torch.Tensor]]: + """A diffusers-layout state dict with ``layers`` stored as int8 + convrot, and the dense model it decodes to. + + The reference carries the *round trip* of each payload (what a correct loader must reconstruct), taken from the + fixture's own quantizer rather than from the code under test. + """ + torch.manual_seed(0) + reference = _tiny_model().float() + sd = {key: value.detach().clone().to(dense_dtype) for key, value in reference.state_dict().items()} + for path in layers: + weight = reference.get_submodule(path).weight + payload = quantize_convrot(weight.detach().float(), group_size=GROUP) + with torch.no_grad(): + weight.copy_(payload.dequantized) + sd[f"{path}.weight"] = payload.codes + sd[f"{path}.weight_scale"] = payload.scale + if marker is not None: + sd[f"{path}.comfy_quant"] = comfy_quant_marker(marker) + return reference, sd + + +def _write( + tmp_path: Path, + sd: dict[str, torch.Tensor], + *, + native: bool = True, + prefix: str = "", + header_layers: list[str] | None = None, + header_prefix: str = "", +) -> Path: + """Save ``sd`` as the community ships it: native or diffusers names, optional ComfyUI prefix, optional header.""" + if native: + sd = _to_native_layout(sd) + metadata = None + if header_layers is not None: + # The release names header entries in native names; `header_prefix` is the ComfyUI wrapper some tools keep. + names = _to_native_layout({f"{path}.weight": torch.empty(0) for path in header_layers}) + entry = {key: value for key, value in MARKER.items() if key != "per_row"} + layers = {f"{header_prefix}{name[: -len('.weight')]}": entry for name in names} + metadata = {"_quantization_metadata": json.dumps({"format_version": "1.0", "layers": layers})} + path = tmp_path / "wan2.2_t2v_high_noise_14B_int8_convrot.safetensors" + save_file({f"{prefix}{key}": value.contiguous() for key, value in sd.items()}, path, metadata=metadata) + return path + + +def _forward(model: torch.nn.Module) -> torch.Tensor: + torch.manual_seed(1) + dtype = next(model.buffers()).dtype if not list(model.parameters()) else next(model.parameters()).dtype + with torch.no_grad(): + return model( + hidden_states=torch.randn(1, 16, 1, 4, 4, dtype=dtype), + timestep=torch.tensor([500]), + encoder_hidden_states=torch.randn(1, 6, 32, dtype=dtype), + return_dict=False, + )[0] + + +def _assert_decodes_like_the_reference(model: torch.nn.Module, reference: torch.nn.Module, layers: list[str]) -> None: + """The kept layers reconstruct the reference weight, and the whole model computes what the reference computes.""" + for path in layers: + layer = model.get_submodule(path) + assert isinstance(layer, Int8ConvrotLinear), path + assert layer.weight.dtype is torch.int8, path + assert layer.convrot is True and layer.group_size == GROUP, path + reconstructed = layer._dequantized_weight(torch.device("cpu"), torch.float32) + torch.testing.assert_close(reconstructed, reference.get_submodule(path).weight, rtol=1e-4, atol=1e-5) + expected = _forward(reference) + assert expected.abs().max() > 0, "reference output is degenerate, the comparison would pass vacuously" + torch.testing.assert_close(_forward(model), expected, rtol=1e-3, atol=1e-4) + + +class TestInt8Convrot: + @pytest.mark.parametrize("native", [True, False], ids=["native-keys", "diffusers-keys"]) + @pytest.mark.parametrize("prefix", ["", "model.diffusion_model."], ids=["bare", "comfyui-prefix"]) + def test_layout_variants_keep_the_layers_int8_and_decode_to_the_reference( + self, tmp_path: Path, native: bool, prefix: str + ) -> None: + reference, sd = _int8_checkpoint(BLOCK_LAYERS) + path = _write(tmp_path, sd, native=native, prefix=prefix) + + model = _load(path, dtype=torch.float32) + + _assert_decodes_like_the_reference(model, reference, BLOCK_LAYERS) + + def test_markers_declared_only_in_the_header_are_honoured(self, tmp_path: Path) -> None: + reference, sd = _int8_checkpoint(BLOCK_LAYERS, marker=None) + path = _write(tmp_path, sd, header_layers=BLOCK_LAYERS) + + model = _load(path, dtype=torch.float32) + + _assert_decodes_like_the_reference(model, reference, BLOCK_LAYERS) + + @pytest.mark.parametrize("header_prefix", ["", "model.diffusion_model."], ids=["bare-header", "prefixed-header"]) + def test_a_prefixed_file_with_header_only_markers_loads(self, tmp_path: Path, header_prefix: str) -> None: + """The loader compares header names with the keys it has already stripped, so a header that kept the ComfyUI + wrapper still has to find its layers.""" + reference, sd = _int8_checkpoint(BLOCK_LAYERS, marker=None) + path = _write( + tmp_path, sd, prefix="model.diffusion_model.", header_layers=BLOCK_LAYERS, header_prefix=header_prefix + ) + + model = _load(path, dtype=torch.float32) + + _assert_decodes_like_the_reference(model, reference, BLOCK_LAYERS) + + def test_activation_scales_are_dropped_not_loaded(self, tmp_path: Path) -> None: + """`input_scale` is W8A8 metadata this path does not use. Left in the dict it would be an unexpected key, + which the Wan loader refuses.""" + reference, sd = _int8_checkpoint(BLOCK_LAYERS) + for layer in BLOCK_LAYERS: + sd[f"{layer}.input_scale"] = torch.ones(()) + path = _write(tmp_path, sd) + + model = _load(path, dtype=torch.float32) + + _assert_decodes_like_the_reference(model, reference, BLOCK_LAYERS) + + def test_a_quantized_condition_embedder_is_widened_and_the_forward_runs(self, tmp_path: Path) -> None: + """The time embedder's dtype probe reads a parameter, so those layers cannot stay int8. They are decoded + densely at load, and the result must still be the reference weight.""" + reference, sd = _int8_checkpoint(BLOCK_LAYERS + EMBEDDERS) + path = _write(tmp_path, sd) + + model = _load(path, dtype=torch.float32) + + for embedder in EMBEDDERS: + layer = model.get_submodule(embedder) + assert not isinstance(layer, Int8ConvrotLinear), embedder + torch.testing.assert_close(layer.weight, reference.get_submodule(embedder).weight, rtol=1e-4, atol=1e-5) + _assert_decodes_like_the_reference(model, reference, BLOCK_LAYERS) + + @pytest.mark.parametrize("embedders", [[], EMBEDDERS], ids=["standard-build", "embedders-quantized"]) + def test_the_reservation_charges_int8_payloads_one_byte_and_widened_layers_the_compute_width( + self, tmp_path: Path, embedders: list[str] + ) -> None: + """A kept payload is one byte per code plus its fp32 scale column; everything else, the widened embedders and + the F16 remainder included, arrives at the compute width. The blanket cast and its size would charge the + codes two bytes and the scales two instead of four.""" + _, sd = _int8_checkpoint(BLOCK_LAYERS + embedders, dense_dtype=torch.float16) + path = _write(tmp_path, sd) + loader = _make_loader() + + _load(path, loader=loader) + + kept = set(BLOCK_LAYERS) + expected = 0 + for key, tensor in _tiny_model().state_dict().items(): + layer = key[: -len(".weight")] if key.endswith(".weight") else None + if layer in kept: + expected += tensor.numel() * 1 + tensor.shape[0] * torch.float32.itemsize + else: + expected += tensor.numel() * torch.bfloat16.itemsize + loader._ram_cache.make_room.assert_called_once_with(expected) + + def test_fp8_storage_leaves_the_int8_buffers_alone(self, tmp_path: Path) -> None: + """The loader ends in the FP8 Storage pass for every path. It walks parameters and an ``Int8ConvrotLinear`` + holds buffers, so the codes and their fp32 scales must come out untouched while the dense remainder is + cast.""" + dense_layer = BLOCK_LAYERS[-1] + quantized = BLOCK_LAYERS[:-1] + _, sd = _int8_checkpoint(quantized) + path = _write(tmp_path, sd) + loader = _make_loader() + loader._should_use_fp8 = lambda *_args: True + loader._torch_dtype = torch.bfloat16 + loader._logger = MagicMock() + + model = _load(path, loader=loader) + + for name in quantized: + layer = model.get_submodule(name) + assert isinstance(layer, Int8ConvrotLinear), name + assert layer.weight.dtype is torch.int8, name + assert layer.weight_scale.dtype is torch.float32, name + # The pass really ran: a block Linear the checkpoint left dense is now stored as fp8. + assert model.get_submodule(dense_layer).weight.dtype is torch.float8_e4m3fn + + def test_an_int8_weight_without_a_marker_is_refused(self, tmp_path: Path) -> None: + """Cast as raw codes it would be a model of small integers that loads without complaint.""" + _, sd = _int8_checkpoint(BLOCK_LAYERS, marker=None) + path = _write(tmp_path, sd) + + # The file name is in the message so a user with two experts can tell which one failed. + message = r"Wan 2\.2 \(wan2\.2_t2v_high_noise_14B_int8_convrot\.safetensors\) checkpoint has \d+ int8 weight" + with pytest.raises(ValueError, match=message + r".*no readable int8_tensorwise marker"): + _load(path) + + def test_an_int4_marked_file_is_refused(self, tmp_path: Path) -> None: + _, sd = _int8_checkpoint(BLOCK_LAYERS, marker={"format": "int4_tensorwise", "convrot": True}) + path = _write(tmp_path, sd) + + with pytest.raises(ValueError, match="marked 'int4_tensorwise'"): + _load(path) + + def test_a_marker_without_a_format_is_refused(self, tmp_path: Path) -> None: + """Files marked only `{"convrot": true}` do not say what they are, so nothing is guessed.""" + _, sd = _int8_checkpoint(BLOCK_LAYERS, marker={"convrot": True}) + path = _write(tmp_path, sd) + + with pytest.raises(ValueError, match="no readable int8_tensorwise marker"): + _load(path) + + def test_a_file_mixing_int8_with_scaled_fp8_is_refused(self, tmp_path: Path) -> None: + """Inside the int8 branch the fp8 fold is skipped, so an fp8 layer would load without its scale.""" + reference, sd = _int8_checkpoint(BLOCK_LAYERS) + fp8 = quantize_scaled_fp8(reference.get_submodule("proj_out").weight.detach().float()) + sd["proj_out.weight"] = fp8.codes + sd["proj_out.weight_scale"] = fp8.scale + path = _write(tmp_path, sd) + + with pytest.raises(ValueError, match="mixing convrot layers with scaled fp8"): + _load(path) diff --git a/tests/backend/model_manager/load/test_wan_int8_convrot_keys.py b/tests/backend/model_manager/load/test_wan_int8_convrot_keys.py new file mode 100644 index 00000000000..968cc314691 --- /dev/null +++ b/tests/backend/model_manager/load/test_wan_int8_convrot_keys.py @@ -0,0 +1,94 @@ +"""What the community `int8_tensorwise` Wan 2.2 T2V-14B build actually contains. + +`test_wan_checkpoint_loader.py` runs on synthetic tiny geometry, which can say whether the loader handles the +shape it was given and cannot say whether the released file still has that shape. These cells read the header +captured in `state_dicts/wan_int8_convrot_keys.py` and pin what the decode depends on: the marker declares the +rotation, the scales and widths fit it, and every native key lands on a parameter of the released architecture +with the quantized layers on Linears the model does not declare precision-sensitive. +""" + +import accelerate +import pytest +import torch +from diffusers import WanTransformer3DModel + +from invokeai.backend.model_manager.load.load_default import _model_declared_skip_patterns +from invokeai.backend.model_manager.load.model_loaders.wan import ( + _convert_wan_native_to_diffusers, + _rename_wan_native_key, +) +from invokeai.backend.quantization.int8_convrot import ( + CONVROT_GROUP_SIZE, + INT8_TENSORWISE_FORMAT, + marker_from_header_entry, + resolve_quantized_module_paths, +) +from tests.backend.model_manager.load.state_dicts import wan_int8_convrot_keys as fixture + +KEYS = fixture.state_dict_keys +QUANTIZED = sorted(key[: -len(".weight")] for key, (_shape, dtype) in KEYS.items() if dtype == "I8") + + +@pytest.fixture(scope="module") +def model() -> WanTransformer3DModel: + """The released T2V-A14B geometry, one block deep, on the meta device.""" + with accelerate.init_empty_weights(): + return WanTransformer3DModel( + patch_size=(1, 2, 2), + in_channels=16, + out_channels=16, + num_layers=1, + attention_head_dim=128, + num_attention_heads=40, + ffn_dim=13824, + text_dim=4096, + ) + + +def test_the_marker_is_rotated_and_the_header_entry_decodes_to_the_same_layer() -> None: + """`convrot` and its group size are what the decode cannot guess: a build that dropped them would still + declare `int8_tensorwise`, still load, and derotate nothing.""" + for path, marker in fixture.markers.items(): + assert marker["format"] == INT8_TENSORWISE_FORMAT + assert marker["convrot"] is True + assert marker["convrot_groupsize"] == CONVROT_GROUP_SIZE + entry = marker_from_header_entry(path, fixture.header_layers[path], {INT8_TENSORWISE_FORMAT}) + assert entry is not None + assert entry["convrot"] is True + assert entry["convrot_groupsize"] == CONVROT_GROUP_SIZE + + +def test_the_scales_are_per_output_channel_and_every_width_is_whole_rotation_groups() -> None: + for path in QUANTIZED: + weight_shape, _ = KEYS[f"{path}.weight"] + scale_shape, scale_dtype = KEYS[f"{path}.weight_scale"] + assert scale_shape == [weight_shape[0], 1], path + assert scale_dtype == "F32", path + assert weight_shape[1] % CONVROT_GROUP_SIZE == 0, path + + +def test_every_key_lands_on_a_parameter_of_the_released_architecture(model: WanTransformer3DModel) -> None: + expected = {key: tuple(value.shape) for key, value in model.state_dict().items()} + converted = _convert_wan_native_to_diffusers( + {key: torch.empty(shape, device="meta") for key, (shape, _) in KEYS.items()} + ) + loadable = { + key: tuple(value.shape) + for key, value in converted.items() + if not key.endswith((".comfy_quant", ".weight_scale")) + } + + assert loadable == expected + + +def test_the_quantized_layers_resolve_to_linears_the_model_does_not_declare_sensitive( + model: WanTransformer3DModel, +) -> None: + key_map = {key: _rename_wan_native_key(key) for key in KEYS} + resolved = resolve_quantized_module_paths(fixture.markers, key_map) + skip = _model_declared_skip_patterns(model) + + assert len(resolved) == len(QUANTIZED) + for path in resolved: + assert isinstance(model.get_submodule(path), torch.nn.Linear), path + assert not any(pattern in path for pattern in skip), path diff --git a/tests/backend/model_manager/load/test_z_image_loader_boundaries.py b/tests/backend/model_manager/load/test_z_image_loader_boundaries.py index 2de7454d840..89285dd3036 100644 --- a/tests/backend/model_manager/load/test_z_image_loader_boundaries.py +++ b/tests/backend/model_manager/load/test_z_image_loader_boundaries.py @@ -170,7 +170,7 @@ def test_an_int8_weight_without_a_per_tensor_marker_is_refused(monkeypatch, tmp_ }, ) - with pytest.raises(ValueError, match=r"int8 weight\(s\) with no `comfy_quant` marker"): + with pytest.raises(ValueError, match=r"int8 weight\(s\) with no readable int8_tensorwise marker"): run.load(config) diff --git a/tests/backend/pid/test_pid_int8.py b/tests/backend/pid/test_pid_int8.py index badfbb60b6c..45d5098689d 100644 --- a/tests/backend/pid/test_pid_int8.py +++ b/tests/backend/pid/test_pid_int8.py @@ -173,7 +173,7 @@ def test_an_int8_weight_no_marker_claims_is_refused(monkeypatch: pytest.MonkeyPa state_dict = {k: v for k, v in state_dict.items() if not k.endswith((".comfy_quant", ".weight_scale"))} monkeypatch.setattr(pid_decode_module, "build_pid_net", lambda backbone, version: _tiny_net()) - with pytest.raises(ValueError, match="no `comfy_quant` marker"): + with pytest.raises(ValueError, match="no readable int8_tensorwise marker"): load_pid_decoder(state_dict, BaseModelType.Flux) diff --git a/tests/backend/qwen3_5/test_qwen3_5_loading.py b/tests/backend/qwen3_5/test_qwen3_5_loading.py index 2ae0901f01b..86994c8a192 100644 --- a/tests/backend/qwen3_5/test_qwen3_5_loading.py +++ b/tests/backend/qwen3_5/test_qwen3_5_loading.py @@ -320,7 +320,7 @@ def test_an_int8_layer_without_an_int8_tensorwise_marker_is_refused(tmp_path: Pa path = tmp_path / "qwen3.5_9b_int8_convrot.safetensors" save_file(file, path) - with pytest.raises(ValueError, match="no `comfy_quant` marker"): + with pytest.raises(ValueError, match="no readable int8_tensorwise marker"): load_qwen3_5_text_encoder(loader, [path], _config(), final_norm=True, keep_fp8=False, what="test") loader._ram_cache.make_room.assert_not_called() From 564523074dbd2a38a13b9ae11c6c025da5fbdf97 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Fri, 6 Nov 2026 19:10:20 +0100 Subject: [PATCH 2/4] docs(wan): list Wan among install refusals of low-bit layouts --- .../docs/Users Guide/03.Models/03.model-format-support.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx b/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx index 12eb949c58b..eae5cc765c1 100644 --- a/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx @@ -53,7 +53,7 @@ Others are recognized by their layer names and fail, or misbehave, only when the exports such as unsloth's FP8 build are refused at install. MOSS-Image loads Comfy-Org's int8 build and GGUFs; its fp8, nvfp4 and w4a8 transformers, and unsloth's torchao builds, are refused at install, unscaled fp8 included. ComfyUI's other low-bit layouts (`w4a8`, `w6a8`, SVDQuant) are not supported; - Krea-2 refuses them at install. + Krea-2 and Wan 2.2 refuse them at install. - **GGUF** covers the llama.cpp types (Q4_0 through Q8_0, the K-quants, BF16). ComfyUI-GGUF's `Q8_CR` (int8) and `Q4_CR` (int4) files load for Krea-2 only; for every other model with GGUF support the install fails, as it does for `Q4_CR` files of ComfyUI-GGUF's retired W4A16 kind. From f147243e9e8b05d9196040b4b4ee6afd60ee257a Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 7 Nov 2026 04:49:07 +0100 Subject: [PATCH 3/4] docs(wan): report int8 RAM use as working set Commit figures on Windows include VRAM, so the page overstated RAM use. --- .../content/docs/Users Guide/05.Video Generation/wan-2.2.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/content/docs/Users Guide/05.Video Generation/wan-2.2.mdx b/docs/src/content/docs/Users Guide/05.Video Generation/wan-2.2.mdx index 0189b07f2b0..f9702639452 100644 --- a/docs/src/content/docs/Users Guide/05.Video Generation/wan-2.2.mdx +++ b/docs/src/content/docs/Users Guide/05.Video Generation/wan-2.2.mdx @@ -63,4 +63,4 @@ Higher-quality Q8_0 quantizations and full Diffusers builds are available a-la-c Single-file A14B and TI2V-5B transformers in ComfyUI's `int8_tensorwise` (convrot) format load like any other single file, with the same two-expert setup for A14B. The attention and feed-forward weights stay int8 in memory, which saves memory. By default they are dequantized per layer on each forward pass; [int8 Compute](/configuration/optimization/int8-compute/) can instead run them on NVIDIA int8 tensor cores. LoRAs, including Lightning, are applied alongside the int8 weights rather than merged into them, except that DoRA layers do not run on an int8 transformer, so a DoRA LoRA fails there when it is applied. Files marked only `{"convrot": true}`, without a `format` entry, and the int4 builds are refused when you install them. -On one Windows machine with an RTX 4090 (T2V A14B, 832x480, 17 frames), an int8 expert file is about 14.5 GB against about 28.6 GB for a bf16 expert. With both experts and the UMT5 encoder kept in the RAM cache, the process peaked at about 62 GB with int8 and about 89 GB with bf16. With `max_cache_ram_gb: 16` the int8 build still runs but reloads the experts for each clip, and the process peaked at about 44 GB, because memory use while loading goes above the cache limit. int8 mainly saves memory; how much speed it changes depends on whether the model otherwise fits in VRAM. Other systems will differ. +On one Windows machine with an RTX 4090 (T2V A14B, 832x480, 17 frames), an int8 expert file is about 14.5 GB against about 28.6 GB for a bf16 expert. With both experts and the UMT5 encoder kept in the RAM cache, the int8 build peaked at about 41 GB of RAM. With `max_cache_ram_gb: 16` it still runs but reloads the experts for each clip, and peaked at about 30 GB of RAM, because both experts are held in memory during the swap. Windows also counts VRAM toward a process's committed memory, so that figure reads higher than the RAM in use. int8 mainly saves memory; how much speed it changes depends on whether the model otherwise fits in VRAM. Other systems will differ. From b4f44bc61f6dab1a73f3527b8a2c6662be336ff4 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sun, 8 Nov 2026 04:27:08 +0100 Subject: [PATCH 4/4] fix(wan): drop stray fp8 metadata on the int8 load path A top-level scaled_fp8 flag or a weight scale beside a dense weight no longer fails the load as an unsupported Wan variant; the int8 branch now drops them like the fp8 branch does. --- .../backend/model_manager/load/model_loaders/wan.py | 6 ++++++ .../model_manager/load/test_wan_checkpoint_loader.py | 12 ++++++++++++ 2 files changed, 18 insertions(+) diff --git a/invokeai/backend/model_manager/load/model_loaders/wan.py b/invokeai/backend/model_manager/load/model_loaders/wan.py index 272fea23882..9099209b4cd 100644 --- a/invokeai/backend/model_manager/load/model_loaders/wan.py +++ b/invokeai/backend/model_manager/load/model_loaders/wan.py @@ -40,6 +40,7 @@ SubModelType, ) from invokeai.backend.quantization.fp8_scaled import ( + is_scale_metadata_key, parse_quantization_metadata, read_safetensors_metadata, strip_layer_path_prefix, @@ -511,6 +512,11 @@ def _load_from_singlefile(self, config: Main_Checkpoint_Wan_Config) -> AnyModel: reserve=self._ram_cache.make_room, skip_patterns=_model_declared_skip_patterns(model), ) + # Metadata the fp8 branch strips: a stray `scaled_fp8` flag, or a weight scale beside a dense weight. + # Left in, `_raise_for_incompatible_keys` would blame it on an Animate/S2V variant. Dropped after the + # install, whose foreign-scale check has to see (and refuse) a scale that sits on an fp8 weight. + for key in [k for k in sd if is_scale_metadata_key(k) and k.rsplit(".", 1)[0] not in quantized]: + del sd[key] logger.info( f"Wan 2.2: kept {len(quantized)} of {len(int8_markers)} layer(s) in int8 " "(int8_tensorwise checkpoint, dequantized per forward)" diff --git a/tests/backend/model_manager/load/test_wan_checkpoint_loader.py b/tests/backend/model_manager/load/test_wan_checkpoint_loader.py index 0befa57438a..dc7a4c0093a 100644 --- a/tests/backend/model_manager/load/test_wan_checkpoint_loader.py +++ b/tests/backend/model_manager/load/test_wan_checkpoint_loader.py @@ -526,6 +526,18 @@ def test_activation_scales_are_dropped_not_loaded(self, tmp_path: Path) -> None: _assert_decodes_like_the_reference(model, reference, BLOCK_LAYERS) + def test_stray_fp8_metadata_is_dropped_as_on_the_fp8_path(self, tmp_path: Path) -> None: + """The fp8 branch drops a top-level `scaled_fp8` flag and a weight scale beside a dense weight silently. Left + in the dict here they would be unexpected keys, which the Wan loader blames on an Animate/S2V variant.""" + reference, sd = _int8_checkpoint(BLOCK_LAYERS) + sd["scaled_fp8"] = torch.zeros(1, dtype=torch.float8_e4m3fn) + sd["proj_out.weight_scale"] = torch.ones(()) + path = _write(tmp_path, sd) + + model = _load(path, dtype=torch.float32) + + _assert_decodes_like_the_reference(model, reference, BLOCK_LAYERS) + def test_a_quantized_condition_embedder_is_widened_and_the_forward_runs(self, tmp_path: Path) -> None: """The time embedder's dtype probe reads a parameter, so those layers cannot stay int8. They are decoded densely at load, and the result must still be the reference weight."""