Support single-file Wan 2.2 checkpoints

invokeai/app/invocations/wan_lora_loader.py

@@ -25,6 +25,13 @@
 # - ``both``: append to both lists regardless of the config.
 # - ``high``: append only to the primary list (high-noise expert).
 # - ``low``: append only to the low-noise list (low-noise expert).
+#
+# One exception applies afterwards, to whichever of the four produced a low-only
+# routing — ``auto`` on a low-tagged LoRA, or an explicit ``low``. Against a
+# single-transformer TI2V-5B main, ``_correct_inert_low_routing`` re-points it at
+# the primary list, because that model has no low-noise expert and the alternative
+# is to accept the LoRA and silently do nothing with it. ``both`` and ``high``
+# always reach the primary list, so they are never affected.
 WanLoRATarget = Literal["auto", "both", "high", "low"]
 
 
@@ -66,22 +73,34 @@ def _assert_lora_variant_matches_main(lora_config: object, main_config: object,
         )
 
 
-def _warn_if_low_routing_is_inert(
+def _correct_inert_low_routing(
     context: InvocationContext, main_config: object, lora_key: str, to_primary: bool, to_low_noise: bool
-) -> None:
-    """Warn when a LoRA is routed only to the low-noise list of a TI2V-5B main.
-
-    The single-transformer TI2V-5B denoise path consumes only the primary list, so
-    such a LoRA silently has no effect — the node would otherwise report success
-    while doing nothing.
+) -> tuple[bool, bool]:
+    """Re-point a low-only routing at the primary list when the main is TI2V-5B.
+
+    TI2V-5B is single-transformer: the denoise path only ever reads the primary LoRA
+    list, so a LoRA routed low-only has no effect at all and the node still reports
+    success. There is no ambiguity about what to do instead — the model has exactly one
+    transformer — so correct the routing rather than merely warning about it.
+
+    This is the backstop for the probe-side pin in ``LoRA_LyCORIS_Wan_Config``, which
+    can only suppress the expert tag when it managed to detect the variant.
+    ``detect_wan_lora_variant`` reads the inner dim off an ``attn1.to_q`` LoRA pair, so
+    it returns None for a LoKr/LoHa adapter or one that patches only ``to_k``/``to_v``,
+    and the tag survives. Records written before that pin existed are in the same
+    position. Here the main model's own variant is known for certain, which is the one
+    signal that cannot be wrong.
     """
     if to_primary or not to_low_noise:
-        return
-    if getattr(main_config, "variant", None) == WanVariantType.TI2V_5B:
-        context.logger.warning(
-            f"LoRA '{lora_key}' is routed only to the low-noise expert, which the single-transformer "
-            "TI2V-5B variant never uses — the LoRA will have no effect."
-        )
+        return to_primary, to_low_noise
+    if getattr(main_config, "variant", None) != WanVariantType.TI2V_5B:
+        return to_primary, to_low_noise
+    context.logger.warning(
+        f"LoRA '{lora_key}' is tagged as the low-noise expert, but the single-transformer "
+        "TI2V-5B variant has no such expert. Applying it to the transformer instead — "
+        "the alternative is to silently do nothing."
+    )
+    return True, False
 
 
 def _resolve_target(target: WanLoRATarget, lora_expert: str | None) -> tuple[bool, bool]:
@@ -127,8 +146,9 @@ class WanLoRALoaderInvocation(BaseInvocation):
     field to override.
 
     For TI2V-5B (single transformer) only the primary list is used at denoise
-    time; a LoRA routed only to the low-noise list would be inert, so that
-    routing logs a warning.
+    time, so a LoRA that would land only in the low-noise list is applied to
+    the transformer instead, with a warning. The alternative is to accept the
+    LoRA and silently have no effect.
     """
 
     lora: ModelIdentifierField = InputField(
@@ -141,7 +161,9 @@ class WanLoRALoaderInvocation(BaseInvocation):
     target: WanLoRATarget = InputField(
         default="auto",
         description="Which expert(s) to apply this LoRA to. 'auto' uses the LoRA's "
-        "recorded expert tag (or both if untagged); 'both'/'high'/'low' override it.",
+        "recorded expert tag (or both if untagged); 'both'/'high'/'low' override it. "
+        "On the single-transformer TI2V-5B, which has no low-noise expert, 'low' is "
+        "applied to the transformer instead of being discarded.",
     )
     transformer: WanTransformerField | None = InputField(
         default=None,
@@ -168,7 +190,7 @@ def invoke(self, context: InvocationContext) -> WanLoRALoaderOutput:
 
         lora_expert = getattr(lora_config, "expert", None)
         to_primary, to_low_noise = _resolve_target(self.target, lora_expert)
-        _warn_if_low_routing_is_inert(context, main_config, lora_key, to_primary, to_low_noise)
+        to_primary, to_low_noise = _correct_inert_low_routing(context, main_config, lora_key, to_primary, to_low_noise)
 
         # Reject duplicates on whichever list(s) we're about to append to.
         if to_primary and any(item.lora.key == lora_key for item in self.transformer.loras):
@@ -200,6 +222,10 @@ class WanLoRACollectionLoader(BaseInvocation):
     Each LoRA is routed to the primary and/or low-noise list based on its
     recorded ``expert`` tag (set by the probe from the filename). Untagged
     LoRAs go to both lists.
+
+    Against a TI2V-5B main, which is a single transformer with no low-noise
+    expert, a LoRA that would land only in the low-noise list is applied to
+    the transformer instead, with a warning.
     """
 
     loras: Optional[LoRAField | list[LoRAField]] = InputField(
@@ -245,7 +271,9 @@ def invoke(self, context: InvocationContext) -> WanLoRALoaderOutput:
 
             lora_expert = getattr(lora_config, "expert", None)
             to_primary, to_low_noise = _resolve_target("auto", lora_expert)
-            _warn_if_low_routing_is_inert(context, main_config, lora_key, to_primary, to_low_noise)
+            to_primary, to_low_noise = _correct_inert_low_routing(
+                context, main_config, lora_key, to_primary, to_low_noise
+            )
 
             # Reject LoRAs already applied upstream (same invariant the single loader
             # enforces) — re-appending would silently double the effective weight.

invokeai/app/invocations/wan_model_loader.py

@@ -17,6 +17,10 @@
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, SubModelType, WanVariantType
 
+# Transformer-only Wan formats: one file holds exactly one expert, so the A14B MoE
+# pair has to be wired up by hand and the VAE / T5 encoder come from elsewhere.
+_SINGLE_FILE_FORMATS = frozenset({ModelFormat.GGUFQuantized, ModelFormat.Checkpoint})
+
 
 @invocation_output("wan_model_loader_output")
 class WanModelLoaderOutput(BaseInvocationOutput):
@@ -38,6 +42,9 @@ class WanModelLoaderOutput(BaseInvocationOutput):
     title="Main Model - Wan 2.2",
     tags=["model", "wan"],
     category="model",
+    # Not bumped for the single-file-checkpoint support: no stored node data changed,
+    # only the live template's model-picker filter widened. Bumping would flag every
+    # saved Wan workflow as needing an update for no benefit.
     version="1.0.1",
     classification=Classification.Prototype,
 )
@@ -49,15 +56,16 @@ class WanModelLoaderInvocation(BaseInvocation):
     - Transformer(s):
         * Diffusers main: emits ``transformer/`` and (for A14B) ``transformer_2/``
           from the same model record.
-        * GGUF main: emits the single GGUF as the primary transformer; for A14B
-          the second-expert GGUF must be wired to ``Transformer (Low Noise)``.
+        * Single-file main (GGUF or safetensors checkpoint): emits the file as the
+          primary transformer; for A14B the second-expert file must be wired to
+          ``Transformer (Low Noise)``.
     - VAE: standalone Wan VAE > main (if Diffusers) > Component Source (Diffusers).
     - UMT5-XXL encoder: standalone Wan T5 encoder > main (if Diffusers) >
       Component Source (Diffusers).
 
     The Component Source slot lets users supply a Diffusers Wan main model purely
     for VAE / encoder extraction when the actual transformer is in a single-file
-    format. Together, the standalone VAE + standalone encoder let a GGUF
+    format. Together, the standalone VAE + standalone encoder let a single-file
     transformer run without a full ~30 GB Diffusers install.
     """
 
@@ -71,14 +79,14 @@ class WanModelLoaderInvocation(BaseInvocation):
 
     transformer_low_noise_model: Optional[ModelIdentifierField] = InputField(
         default=None,
-        description="Optional second GGUF transformer for the A14B low-noise expert. "
-        "Only relevant when the main model is a single-file GGUF and the variant is A14B; "
-        "ignored when the main is a Diffusers A14B (both experts are pulled from "
-        "transformer/ and transformer_2/ already) or when the variant is TI2V-5B.",
+        description="Optional second single-file transformer for the A14B low-noise expert. "
+        "Only relevant when the main model is a single-file GGUF or safetensors checkpoint and "
+        "the variant is A14B; ignored when the main is a Diffusers A14B (both experts are pulled "
+        "from transformer/ and transformer_2/ already) or when the variant is TI2V-5B.",
         input=Input.Direct,
         ui_model_base=BaseModelType.Wan,
         ui_model_type=ModelType.Main,
-        ui_model_format=ModelFormat.GGUFQuantized,
+        ui_model_format=[ModelFormat.GGUFQuantized, ModelFormat.Checkpoint],
         title="Transformer (Low Noise)",
     )
 
@@ -118,9 +126,9 @@ def invoke(self, context: InvocationContext) -> WanModelLoaderOutput:
         self._validate_main_config(main_config, "Wan main")
         main_format = main_config.format
         main_is_diffusers = main_format == ModelFormat.Diffusers
-        main_is_gguf = main_format == ModelFormat.GGUFQuantized
+        main_is_single_file = main_format in _SINGLE_FILE_FORMATS
         main_variant = getattr(main_config, "variant", None)
-        if main_is_gguf and self.component_source is not None:
+        if main_is_single_file and self.component_source is not None:
             self._validate_component_source_format(context, self.component_source)
 
         # Resolve transformer + dual-expert wiring + boundary_ratio.
@@ -129,11 +137,11 @@ def invoke(self, context: InvocationContext) -> WanModelLoaderOutput:
         # low-noise expert (A14B only). boundary_ratio comes from the probed
         # model_index.json.
         #
-        # GGUF main: the file itself is one expert (high or low). For A14B,
-        # the user wires the other expert to transformer_low_noise_model.
-        # We swap so the *high*-noise expert is always the primary if needed.
-        # boundary_ratio falls back to 0.875 unless a Diffusers component_source
-        # provides a recorded value.
+        # Single-file main (GGUF or safetensors checkpoint): the file itself is one
+        # expert (high or low). For A14B, the user wires the other expert to
+        # transformer_low_noise_model. We swap so the *high*-noise expert is always
+        # the primary if needed. boundary_ratio falls back to 0.875 unless a
+        # Diffusers component_source provides a recorded value.
         boundary_ratio = 0.9 if main_variant == WanVariantType.I2V_A14B else 0.875
         transformer_low_noise: Optional[ModelIdentifierField] = None
 
@@ -144,7 +152,7 @@ def invoke(self, context: InvocationContext) -> WanModelLoaderOutput:
                 recorded = getattr(main_config, "boundary_ratio", None)
                 if recorded is not None:
                     boundary_ratio = float(recorded)
-        elif main_is_gguf:
+        elif main_is_single_file:
             primary_expert = getattr(main_config, "expert", "none")
             primary_id = self.model.model_copy(update={"submodel_type": SubModelType.Transformer})
 
@@ -157,36 +165,43 @@ def invoke(self, context: InvocationContext) -> WanModelLoaderOutput:
                 if self.transformer_low_noise_model.key == self.model.key:
                     raise ValueError(
                         "The same model is wired to both 'Transformer' and 'Transformer (Low Noise)'. "
-                        "A Wan A14B expert pair needs two different GGUF models."
+                        "A Wan A14B expert pair needs two different single-file models."
                     )
                 low_config = context.models.get_config(self.transformer_low_noise_model)
                 self._validate_main_config(low_config, "Transformer (Low Noise)")
-                if low_config.format != ModelFormat.GGUFQuantized:
+                # The two experts don't have to share a format — both single-file
+                # loaders produce a plain WanTransformer3DModel, so a GGUF high-noise
+                # expert pairs fine with a safetensors low-noise one.
+                if low_config.format not in _SINGLE_FILE_FORMATS:
                     raise ValueError(
-                        f"'Transformer (Low Noise)' must be a GGUF-format Wan model. "
+                        f"'Transformer (Low Noise)' must be a single-file Wan model (GGUF or checkpoint). "
                         f"'{low_config.name}' is in {low_config.format.value} format."
                     )
                 low_id = self.transformer_low_noise_model.model_copy(update={"submodel_type": SubModelType.Transformer})
                 low_expert = getattr(low_config, "expert", "none")
 
                 if getattr(low_config, "variant", None) != main_variant:
-                    raise ValueError("The high-noise and low-noise GGUF models must use the same Wan variant.")
+                    low_variant = getattr(low_config, "variant", None)
+                    raise ValueError(
+                        "The high-noise and low-noise models must use the same Wan variant, but "
+                        f"'{main_config.name}' is {main_variant.value} and '{low_config.name}' is "
+                        f"{getattr(low_variant, 'value', low_variant)}."
+                    )
 
-                # The expert tag is a filename heuristic, so 'none' (untagged)
-                # is common on community finetunes. The wiring itself is
-                # explicit user intent — main slot = high, low-noise slot =
-                # low — so an untagged file is taken at its wired position (or
-                # inferred as the complement of its tagged partner). Only a
-                # genuine conflict, both files claiming the *same* expert, is
-                # an error.
+                # The expert tag is a filename heuristic, so 'none' (untagged) is common on
+                # community finetunes. The wiring itself is explicit user intent — main slot
+                # = high, low-noise slot = low — so an untagged file is taken at its wired
+                # position (or inferred as the complement of its tagged partner). Only a
+                # genuine conflict, both files claiming the *same* expert, is an error.
                 if primary_expert == low_expert != "none":
                     raise ValueError(
-                        f"Both selected GGUF models are tagged as the {primary_expert}-noise expert. "
-                        "A Wan A14B expert pair must contain one high and one low expert."
+                        f"Both selected models are tagged as the {primary_expert}-noise expert "
+                        f"('{main_config.name}' and '{low_config.name}'). A Wan A14B expert pair "
+                        "must contain one high and one low expert."
                     )
                 if primary_expert == "none" and low_expert == "none":
                     context.logger.warning(
-                        "Neither Wan A14B GGUF filename identifies its expert, so 'Transformer' is assumed to "
+                        "Neither Wan A14B filename identifies its expert, so 'Transformer' is assumed to "
                         "be the high-noise expert and 'Transformer (Low Noise)' the low-noise expert. If the "
                         "output looks wrong, swap the two models."
                     )
@@ -201,7 +216,7 @@ def invoke(self, context: InvocationContext) -> WanModelLoaderOutput:
                     # filename tag, so say so: a mistagged file is otherwise an
                     # invisible expert inversion.
                     context.logger.warning(
-                        f"The wired Wan A14B GGUF experts look reversed, so they were swapped: "
+                        f"The wired Wan A14B experts look reversed, so they were swapped: "
                         f"'{low_config.name}' (tagged '{low_expert}') runs as the high-noise expert and "
                         f"'{main_config.name}' (tagged '{primary_expert}') as the low-noise expert. "
                         "The tags come from the filenames — if the output looks wrong, a filename is lying."
@@ -211,16 +226,15 @@ def invoke(self, context: InvocationContext) -> WanModelLoaderOutput:
                     transformer_low_noise = low_id
             else:
                 transformer = primary_id
-                # A14B without a paired low-noise GGUF will produce degraded
-                # quality (only one expert runs). Warn but don't abort — a
-                # single wired transformer is explicit intent just like a pair
-                # is, and the tag is only a filename guess, so an untagged file
-                # must not be fatal here when the paired path accepts it.
-                # TI2V-5B GGUFs are single-expert and totally fine.
+                # A14B without a paired low-noise expert will produce degraded quality
+                # (only one expert runs). Warn but don't abort — a single wired transformer
+                # is explicit intent just like a pair is, and the tag is only a filename
+                # guess, so an untagged file must not be fatal here when the paired path
+                # accepts it. TI2V-5B is single-expert and totally fine.
                 if main_variant in (WanVariantType.T2V_A14B, WanVariantType.I2V_A14B):
                     message = (
-                        "An A14B GGUF is wired to 'Transformer' without a paired 'Transformer (Low Noise)'. "
-                        "Only this one expert will run; image quality will be reduced."
+                        "An A14B single-file main is wired to 'Transformer' without a paired "
+                        "'Transformer (Low Noise)'. Only this one expert will run; quality will be reduced."
                     )
                     if primary_expert == "low":
                         message += (
@@ -243,7 +257,7 @@ def invoke(self, context: InvocationContext) -> WanModelLoaderOutput:
         else:
             raise ValueError(
                 f"Unsupported main model format for Wan: {main_format.value}. "
-                "Use a Diffusers folder or a GGUF single-file checkpoint."
+                "Use a Diffusers folder, a GGUF file, or a single-file safetensors checkpoint."
             )
 
         # VAE: standalone override > main (if Diffusers) > component source.

invokeai/backend/model_manager/configs/factory.py

@@ -78,6 +78,7 @@
     Main_Checkpoint_SD2_Config,
     Main_Checkpoint_SDXL_Config,
     Main_Checkpoint_SDXLRefiner_Config,
+    Main_Checkpoint_Wan_Config,
     Main_Checkpoint_ZImage_Config,
     Main_Diffusers_CogView4_Config,
     Main_Diffusers_ErnieImage_Config,
@@ -286,6 +287,7 @@ def has_model_export(module: Any, name: Any, expected_bases: tuple[type, ...]) -
         Annotated[Main_Checkpoint_Flux2_Config, Main_Checkpoint_Flux2_Config.get_tag()],
         Annotated[Main_Checkpoint_FLUX_Config, Main_Checkpoint_FLUX_Config.get_tag()],
         Annotated[Main_Checkpoint_QwenImage_Config, Main_Checkpoint_QwenImage_Config.get_tag()],
+        Annotated[Main_Checkpoint_Wan_Config, Main_Checkpoint_Wan_Config.get_tag()],
         Annotated[Main_Checkpoint_ZImage_Config, Main_Checkpoint_ZImage_Config.get_tag()],
         Annotated[Main_Checkpoint_Krea2_Config, Main_Checkpoint_Krea2_Config.get_tag()],
         Annotated[Main_Checkpoint_Anima_Config, Main_Checkpoint_Anima_Config.get_tag()],