port concept of musubi-tuner wan_force_2_1_time_embedding

Author: bghira

documentation/quickstart/WAN.md

@@ -42,6 +42,10 @@ You'll need:
 - **a realistic minimum** is 16GB or, a single 3090 or V100 GPU
 - **ideally** multiple 4090, A6000, L40S, or better
 
+If you encounter shape mismatches in the time embedding layers when running Wan 2.2 checkpoints, enable the new
+`wan_force_2_1_time_embedding` flag. This forces the transformer to fall back to Wan 2.1 style time embeddings and
+resolves the compatibility issue.
+
 Apple silicon systems do not work super well with Wan 2.1 so far, something like 10 minutes for a single training step can be expected..
 
 ### Prerequisites

simpletuner/helpers/models/wan/model.py

@@ -92,6 +92,8 @@ class Wan(VideoModelFoundation):
     def __init__(self, config, accelerator):
         super().__init__(config, accelerator)
         self._wan_cached_stage_modules: Dict[str, WanTransformer3DModel] = {}
+        if not hasattr(self.config, "wan_force_2_1_time_embedding"):
+            self.config.wan_force_2_1_time_embedding = False
 
     def requires_conditioning_image_embeds(self) -> bool:
         return self._wan_stage_info() is not None
@@ -122,6 +124,14 @@ def _wan_stage_info(self) -> Optional[Dict[str, object]]:
         flavour = getattr(self.config, "model_flavour", None)
         return self.WAN_STAGE_OVERRIDES.get(flavour)
 
+    def _apply_time_embedding_override(self, transformer: Optional[WanTransformer3DModel]) -> None:
+        if transformer is None:
+            return
+        target = self.unwrap_model(transformer)
+        setter = getattr(target, "set_time_embedding_v2_1", None)
+        if callable(setter):
+            setter(bool(getattr(self.config, "wan_force_2_1_time_embedding", False)))
+
     def _should_load_other_stage(self) -> bool:
         stage_info = self._wan_stage_info()
         if stage_info is None:
@@ -142,13 +152,19 @@ def _get_or_load_wan_stage_module(self, subfolder: str) -> WanTransformer3DModel
         stage.requires_grad_(False)
         stage.to(self.accelerator.device, dtype=self.config.weight_dtype)
         stage.eval()
+        self._apply_time_embedding_override(stage)
         self._wan_cached_stage_modules[subfolder] = stage
         return stage
 
     def unload_model(self):
         super().unload_model()
         self._wan_cached_stage_modules.clear()
 
+    def set_prepared_model(self, model, base_model: bool = False):
+        super().set_prepared_model(model, base_model)
+        if not base_model:
+            self._apply_time_embedding_override(self.model)
+
     def get_group_offload_components(self, pipeline):
         components = dict(super().get_group_offload_components(pipeline))
         transformer_2 = getattr(pipeline, "transformer_2", None)
@@ -183,6 +199,10 @@ def get_pipeline(self, pipeline_type: str = PipelineTypes.TEXT2IMG, load_base_mo
         else:
             pipeline.config.boundary_ratio = None
 
+        self._apply_time_embedding_override(getattr(pipeline, "transformer", None))
+        if getattr(pipeline, "transformer_2", None) is not None:
+            self._apply_time_embedding_override(pipeline.transformer_2)
+
         return pipeline
 
     def tread_init(self):

simpletuner/helpers/models/wan/transformer.py

@@ -454,6 +454,18 @@ def __init__(
         self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5)
 
         self.gradient_checkpointing = False
+        self.force_v2_1_time_embedding: bool = False
+
+    def set_time_embedding_v2_1(self, force_2_1_time_embedding: bool) -> None:
+        """
+        Force the Wan transformer to use 2.1-style time embeddings even when running Wan 2.2 checkpoints.
+
+        Args:
+            force_2_1_time_embedding: Whether to override the default time embedding behaviour.
+        """
+        self.force_v2_1_time_embedding = bool(force_2_1_time_embedding)
+        if self.force_v2_1_time_embedding:
+            logger.info("WanTransformer3DModel: Forcing Wan 2.1 style time embedding.")
 
     def set_router(self, router: TREADRouter, routes: List[Dict[str, Any]]):
         """Set the TREAD router and routing configuration."""
@@ -519,6 +531,11 @@ def forward(
         hidden_states = self.patch_embedding(hidden_states)
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
 
+        if self.force_v2_1_time_embedding and timestep.dim() > 1:
+            # Wan 2.1 uses a single timestep per batch entry. When forcing 2.1 behaviour with Wan 2.2
+            # checkpoints we fall back to the first timestep value which matches the reference implementation.
+            timestep = timestep[..., 0].contiguous()
+
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
             timestep, encoder_hidden_states, encoder_hidden_states_image
         )

simpletuner/simpletuner_sdk/server/services/field_registry/sections/model.py

@@ -506,6 +506,24 @@ def _quant_label(value: str) -> str:
         )
     )
 
+    registry._add_field(
+        ConfigField(
+            name="wan_force_2_1_time_embedding",
+            arg_name="--wan_force_2_1_time_embedding",
+            ui_label="Force Wan 2.1 Time Embedding",
+            field_type=FieldType.CHECKBOX,
+            tab="model",
+            section="model_config",
+            subsection="wan_specific",
+            default_value=False,
+            dependencies=[FieldDependency(field="model_family", operator="equals", value="wan", action="show")],
+            help_text="Use Wan 2.1 style time embeddings even when running Wan 2.2 checkpoints.",
+            tooltip="Enable this if Wan 2.2 checkpoints report shape mismatches in the time embedding layers.",
+            importance=ImportanceLevel.ADVANCED,
+            order=30,
+        )
+    )
+
     # Fused QKV Projections
     registry._add_field(
         ConfigField(