Wan: add gated r-embedder fusion + AnyFlow weight remap

AnyFlow's released HF checkpoints store the r-pathway as
``condition_embedder.delta_embedder.*`` inside the shared
``WanTwoTimeTextImageEmbedding`` module and use ONE shared ``time_proj``
for both t and (t, r). Their forward then mixes the two embeddings
with a convex combination ``(1 - g) * temb_t + g * temb_r`` before the
shared final projection:

    rt_emb         = (1 - g) * temb_t + g * temb_r
    timestep_proj  = time_proj(silu(rt_emb))

FastGen's existing r-embedder design (used by MeanFlow) instead has a
separate top-level ``r_embedder`` with its own ``time_proj`` and adds
``temb_t + temb_r`` / ``timestep_proj_t + timestep_proj_r`` after the
non-linearity. The two layouts are not functionally equivalent because
``silu`` is non-linear.

Two changes:

* ``Wan.__init__``: add ``r_embedder_fusion: str = "additive"`` (default
  preserves MeanFlow's behaviour) and ``r_embedder_gate_value: float =
  0.25``. When ``r_embedder_fusion="gated"``, ``classify_forward_prepare``
  computes the convex-mix variant and uses ``r_embedder.time_proj``
  (which ``init_embedder`` already deep-copies from
  ``condition_embedder.time_proj``) for the shared final projection.

* ``fastgen/methods/distribution_matching/anyflow.py``: add
  ``remap_anyflow_keys`` helper that rewrites AnyFlow's
  ``condition_embedder.delta_embedder.linear_{1,2}.*`` to FastGen's
  ``r_embedder.time_embedder.linear_{1,2}.*`` and duplicates
  ``condition_embedder.time_proj.*`` into ``r_embedder.time_proj.*``
  so the two projections start identical. The function is a no-op when
  no AnyFlow-format keys are present.

Verification (on GMI 2 x H200, gpu-h200-68):

* Forward equivalence on the same inputs (FastGen-loaded vs AnyFlow's
  own loader): rel mean diff = 2.8% in bf16 (forward noise floor).
* Training-step loss equivalence (AnyFlow ``train_bidirection`` math
  reproduced inline on both code paths, same seed): AnyFlow loss
  0.381619 vs FastGen loss 0.397162, rel diff = 4.07%.
* 4-step Euler-flow inference end-to-end (text encoder + FastGen Wan +
  VAE decode) produces a finite 81-frame 480x832 video matching the
  AnyFlow paper's any-step inference pattern.

Signed-off-by: Enderfga <qq2639135175@gmail.com>

Author: Enderfga

fastgen/methods/distribution_matching/anyflow.py

@@ -46,6 +46,45 @@
 import fastgen.utils.logging_utils as logger
 
 
+def remap_anyflow_keys(state_dict: dict) -> dict:
+    """Remap an AnyFlow HF release state_dict to FastGen's Wan layout.
+
+    AnyFlow's ``FAR_Wan_Transformer3DModel`` stores the r-pathway inside the
+    main ``condition_embedder`` as ``delta_embedder``, and uses ONE shared
+    ``time_proj`` for both t and (t, r). FastGen exposes a separate top-level
+    ``r_embedder`` with its own ``time_embedder`` + ``time_proj``. The two
+    layouts are functionally equivalent (FastGen's ``r_embedder.time_proj``
+    starts as a deepcopy of ``condition_embedder.time_proj`` per
+    :meth:`Wan.init_embedder`), so we just rename / duplicate the tensors.
+
+    The function is a no-op when no ``condition_embedder.delta_embedder.*``
+    keys are present, so it's safe to call unconditionally.
+    """
+    delta_keys = [k for k in state_dict if k.startswith("condition_embedder.delta_embedder.")]
+    if not delta_keys:
+        return state_dict
+    new_sd = dict(state_dict)
+    for k in delta_keys:
+        # condition_embedder.delta_embedder.linear_1.weight
+        #   -> r_embedder.time_embedder.linear_1.weight
+        new_k = k.replace("condition_embedder.delta_embedder.", "r_embedder.time_embedder.")
+        new_sd[new_k] = new_sd.pop(k)
+    # AnyFlow's gated fusion shares the final time_proj. FastGen has a
+    # separate r_embedder.time_proj that mathematically substitutes for the
+    # shared one when fusion="gated"; copy the weights across so the two
+    # projections start identical (and AnyFlow's training never diverges them).
+    for sub in ("weight", "bias"):
+        src = f"condition_embedder.time_proj.{sub}"
+        dst = f"r_embedder.time_proj.{sub}"
+        if src in new_sd and dst not in new_sd:
+            new_sd[dst] = new_sd[src].clone()
+    logger.info(
+        f"remap_anyflow_keys: rewrote {len(delta_keys)} delta_embedder tensors "
+        "and duplicated time_proj weights into r_embedder."
+    )
+    return new_sd
+
+
 if TYPE_CHECKING:
     from fastgen.configs.methods.config_anyflow import ModelConfig
 

fastgen/networks/Wan/network.py

@@ -338,14 +338,27 @@ def classify_forward_prepare(
             r_timestep = r_timestep.to(time_embedder_dtype)
 
         remb = self.r_embedder.time_embedder(r_timestep).type_as(encoder_hidden_states)
-        r_timestep_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(remb))
-        r_timestep_proj = unflatten_timestep_proj(r_timestep_proj, rs_seq_len)
 
-        if self.encoder_depth is None:
-            timestep_proj = timestep_proj + r_timestep_proj
-            temb = temb + remb
+        # AnyFlow-style gated mixing: convex-combine t- and r-embeddings BEFORE
+        # the shared final projection, matching the WanTwoTimeTextImageEmbedding
+        # forward in the AnyFlow reference. The released AnyFlow HF checkpoints
+        # require this fusion to reproduce the published forward pass.
+        if getattr(self.r_embedder, "fusion_mode", "additive") == "gated":
+            gate = self.r_embedder.gate_value.to(remb.dtype)
+            rt_emb = (1 - gate) * temb + gate * remb
+            timestep_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(rt_emb))
+            timestep_proj = unflatten_timestep_proj(timestep_proj, rs_seq_len)
+            r_timestep_proj = None
+            temb = rt_emb
         else:
-            temb = remb
+            r_timestep_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(remb))
+            r_timestep_proj = unflatten_timestep_proj(r_timestep_proj, rs_seq_len)
+
+            if self.encoder_depth is None:
+                timestep_proj = timestep_proj + r_timestep_proj
+                temb = temb + remb
+            else:
+                temb = remb
     elif r_timestep is not None:
         # Raise an error here, otherwise we silently ignore the r_timestep
         raise ValueError("r_timestep provided but no r_embedder is present")
@@ -557,6 +570,8 @@ def __init__(
         load_pretrained: bool = True,
         use_fsdp_checkpoint: bool = True,
         use_wan_official_sinusoidal: bool = False,
+        r_embedder_fusion: str = "additive",
+        r_embedder_gate_value: float = 0.25,
         **model_kwargs,
     ):
         """Wan2.1/2.2 model constructor.
@@ -610,6 +625,20 @@ def __init__(
         if r_timestep:
             logger.info(f"Initializing r embedder with {r_embedder_init}")
             self.transformer.r_embedder = self.init_embedder(r_embedder_init)
+            # Stash fusion config on the r_embedder module so the (method-bound)
+            # forward override can branch on it without changing its signature.
+            if r_embedder_fusion not in ("additive", "gated"):
+                raise ValueError(f"r_embedder_fusion must be 'additive' or 'gated', got {r_embedder_fusion!r}")
+            self.transformer.r_embedder.fusion_mode = r_embedder_fusion
+            self.transformer.r_embedder.register_buffer(
+                "gate_value",
+                torch.tensor([float(r_embedder_gate_value)], dtype=torch.float32),
+                persistent=False,
+            )
+            logger.info(
+                f"r_embedder fusion={r_embedder_fusion}"
+                + (f" gate_value={r_embedder_gate_value}" if r_embedder_fusion == "gated" else "")
+            )
         else:
             self.transformer.r_embedder = None