diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md
index 3c58c1181..945f83caa 100644
--- a/docs/user_guide/diffusion_features.md
+++ b/docs/user_guide/diffusion_features.md
@@ -145,6 +145,7 @@ The following tables show which models support each feature:
 | Model                        | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
 | **Wan2.2**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ✅         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
+| **Wan2.2-S2V**               |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |         ❌         |   ❌    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
diff --git a/examples/offline_inference/speech_to_video/speech_to_video.py b/examples/offline_inference/speech_to_video/speech_to_video.py
index 97ce63700..d783ee2f7 100644
--- a/examples/offline_inference/speech_to_video/speech_to_video.py
+++ b/examples/offline_inference/speech_to_video/speech_to_video.py
@@ -16,9 +16,10 @@ Usage:
 """
 
 import argparse
-import os
+import json
 import time
 from pathlib import Path
+from typing import Any
 
 import torch
 
@@ -29,6 +30,16 @@ from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
 
+def parse_profiler_config(value: str) -> dict[str, Any]:
+    try:
+        config = json.loads(value)
+    except json.JSONDecodeError as e:
+        raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
+    if not isinstance(config, dict):
+        raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
+    return config
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Generate a talking-head video from a reference image and audio (Wan2.2 S2V)."
@@ -87,6 +98,12 @@ def parse_args() -> argparse.Namespace:
         default=3.0,
         help="Scheduler flow shift (default: 3.0).",
     )
+    parser.add_argument(
+        "--boundary-ratio",
+        type=float,
+        default=None,
+        help="Boundary split ratio for low/high DiT stages (S2V). Default varies by model.",
+    )
     parser.add_argument(
         "--output",
         type=str,
@@ -178,30 +195,10 @@ def parse_args() -> argparse.Namespace:
         help="Enable cache-dit summary logging after diffusion forward passes.",
     )
     parser.add_argument(
-        "--profile-dir",
-        type=str,
+        "--profiler-config",
+        type=parse_profiler_config,
         default=None,
-        help="Enable torch profiler and save traces to this directory.",
-    )
-    parser.add_argument(
-        "--profile-record-shapes",
-        action="store_true",
-        help="Record tensor shapes in profiler (increases trace size).",
-    )
-    parser.add_argument(
-        "--profile-with-stack",
-        action="store_true",
-        help="Record stack traces in profiler (increases overhead).",
-    )
-    parser.add_argument(
-        "--profile-with-memory",
-        action="store_true",
-        help="Profile memory usage.",
-    )
-    parser.add_argument(
-        "--profile-with-flops",
-        action="store_true",
-        help="Estimate FLOPs for operations.",
+        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
     )
     return parser.parse_args()
 
@@ -239,23 +236,9 @@ def main():
         ulysses_degree=args.ulysses_degree,
     )
 
-    # Check if profiling is requested via CLI args or environment variable
-    profile_dir = args.profile_dir or os.getenv("VLLM_TORCH_PROFILER_DIR")
-    profiler_enabled = bool(profile_dir)
-    profiler_config = None
-    if profiler_enabled:
-        from vllm.config import ProfilerConfig
+    profiler_enabled = args.profiler_config is not None
 
-        profiler_config = ProfilerConfig(
-            profiler="torch",
-            torch_profiler_dir=profile_dir,
-            torch_profiler_record_shapes=args.profile_record_shapes,
-            torch_profiler_with_stack=args.profile_with_stack,
-            torch_profiler_with_memory=args.profile_with_memory,
-            torch_profiler_with_flops=args.profile_with_flops,
-        )
-
-    omni = Omni(
+    omni_kwargs = dict(
         model=args.model,
         model_class_name="WanS2VPipeline",
         flow_shift=args.flow_shift,
@@ -269,8 +252,12 @@ def main():
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         enable_cache_dit_summary=args.enable_cache_dit_summary,
-        profiler_config=profiler_config,
+        profiler_config=args.profiler_config,
     )
+    if args.boundary_ratio is not None:
+        omni_kwargs["boundary_ratio"] = args.boundary_ratio
+
+    omni = Omni(**omni_kwargs)
 
     # Print generation configuration
     print(f"\n{'=' * 60}")
@@ -325,13 +312,16 @@ def main():
     if profiler_enabled:
         print("\n[Profiler] Stopping profiler and collecting results...")
         profile_results = omni.stop_profile()
-        if profile_results and isinstance(profile_results, list):
+        if profile_results and isinstance(profile_results, dict):
+            traces = profile_results.get("traces", [])
             print("\n" + "=" * 60)
             print("PROFILING RESULTS:")
-            for i, result in enumerate(profile_results):
-                print(f"\nStage {i}:")
-                if result:
-                    print(f"  {result}")
+            for rank, trace in enumerate(traces):
+                print(f"\nRank {rank}:")
+                if trace:
+                    print(f"  • Trace: {trace}")
+            if not traces:
+                print("  No traces collected.")
             print("=" * 60)
         else:
             print("[Profiler] No valid profiling data returned.")
diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py
index 6a2600e9f..ff1eea6ae 100644
--- a/vllm_omni/diffusion/cache/cache_dit_backend.py
+++ b/vllm_omni/diffusion/cache/cache_dit_backend.py
@@ -242,6 +242,80 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int],
     return refresh_cache_context
 
 
+def enable_cache_for_wan22_s2v(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for Wan2.2 S2V.
+
+    S2V uses a single transformer, but unlike the other Wan2.2 variants its
+    block loop calls each block as ``block(hidden_states, **kwargs)`` and keeps
+    the timestep modulation state in ``e`` rather than a second positional
+    tensor. CacheDiT Pattern_3 matches that contract: cache hidden states only
+    and pass the remaining conditioning through kwargs unchanged.
+
+    The S2V transformer has an ``after_transformer_block`` method that injects
+    audio embeddings after specific layers. The cached blocks wrapper
+    (Wan22S2VCachedBlocks._run_block) calls the original internally, so we
+    permanently replace it with a no-op on the transformer to prevent double
+    injection from the main forward loop.
+    """
+    db_cache_config = _build_db_cache_config(cache_config)
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    # Save the original after_transformer_block before cache-dit wrapping
+    transformer = pipeline.transformer
+    if hasattr(transformer, "after_transformer_block"):
+        transformer._cache_dit_original_after_transformer_block = transformer.after_transformer_block
+
+    Wan22S2VCachedAdapter.apply(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=[transformer.blocks],
+            forward_pattern=[ForwardPattern.Pattern_3],
+            params_modifiers=[
+                ParamsModifier(cache_config=db_cache_config, calibrator_config=calibrator_config),
+            ],
+            has_separate_cfg=True,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    # Permanently replace after_transformer_block with a no-op.
+    # The cached blocks wrapper (Wan22S2VCachedBlocks._run_block) already calls
+    # the original via _cache_dit_original_after_transformer_block.
+    def _noop_after_transformer_block(block_idx: int, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states
+
+    transformer.after_transformer_block = _noop_after_transformer_block
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache context for the S2V transformer."""
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                num_inference_steps=num_inference_steps,
+                verbose=verbose,
+            )
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
 def enable_cache_for_longcat_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
     """Enable cache-dit for LongCatImage pipeline.
 
@@ -986,6 +1060,151 @@ class BagelCachedBlocks(CachedBlocks_Pattern_0_1_2):
         return hidden_states, encoder_hidden_states
 
 
+class Wan22S2VCachedBlocks(CachedBlocks_Pattern_3_4_5):
+    """CacheDiT blocks wrapper that preserves S2V per-layer audio injection."""
+
+    def _run_block(self, block_id: int, block: torch.nn.Module, hidden_states: torch.Tensor, *args, **kwargs):
+        hidden_states = block(hidden_states, *args, **kwargs)
+        hidden_states, new_encoder_hidden_states = self._process_block_outputs(hidden_states)
+        original_after_transformer_block = getattr(
+            self.transformer,
+            "_cache_dit_original_after_transformer_block",
+            getattr(self.transformer, "after_transformer_block", None),
+        )
+        if original_after_transformer_block is not None:
+            hidden_states = original_after_transformer_block(block_id, hidden_states)
+        return hidden_states, new_encoder_hidden_states
+
+    def call_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        new_encoder_hidden_states = None
+        for block_id, block in enumerate(self.transformer_blocks):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+        return hidden_states, new_encoder_hidden_states
+
+    def call_Fn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        new_encoder_hidden_states = None
+        for block_id, block in enumerate(self._Fn_blocks()):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+        return hidden_states, new_encoder_hidden_states
+
+    def call_Mn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        original_hidden_states = hidden_states
+        new_encoder_hidden_states = None
+        start_idx = self.context_manager.Fn_compute_blocks()
+        for block_id, block in enumerate(self._Mn_blocks(), start=start_idx):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+
+        hidden_states = hidden_states.contiguous()
+        hidden_states_residual = hidden_states - original_hidden_states.to(hidden_states.device)
+
+        return (
+            hidden_states,
+            new_encoder_hidden_states,
+            hidden_states_residual,
+        )
+
+    def call_Bn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        new_encoder_hidden_states = None
+        if self.context_manager.Bn_compute_blocks() == 0:
+            return hidden_states, new_encoder_hidden_states
+
+        start_idx = len(self.transformer_blocks) - self.context_manager.Bn_compute_blocks()
+        for block_id, block in enumerate(self._Bn_blocks(), start=start_idx):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+
+        return hidden_states, new_encoder_hidden_states
+
+
+class Wan22S2VCachedAdapter(CachedAdapter):
+    """CacheDiT adapter that uses Wan22S2VCachedBlocks for S2V audio injection.
+
+    Only overrides collect_unified_blocks to use Wan22S2VCachedBlocks (which
+    calls after_transformer_block per-layer internally). The base class
+    mock_transformer handles the forward wrapping — after_transformer_block is
+    permanently replaced with a no-op in enable_cache_for_wan22_s2v() to prevent
+    double injection.
+    """
+
+    @classmethod
+    def collect_unified_blocks(
+        cls,
+        block_adapter: BlockAdapter,
+        contexts_kwargs: list[dict],
+    ) -> list[dict[str, torch.nn.ModuleList]]:
+        BlockAdapter.assert_normalized(block_adapter)
+
+        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
+        assert hasattr(block_adapter.pipe, "_context_manager")
+
+        for i in range(len(block_adapter.transformer)):
+            unified_blocks_bind_context = {}
+            for j in range(len(block_adapter.blocks[i])):
+                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
+                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
+                    [
+                        Wan22S2VCachedBlocks(
+                            block_adapter.blocks[i][j],
+                            transformer=block_adapter.transformer[i],
+                            forward_pattern=block_adapter.forward_pattern[i][j],
+                            check_forward_pattern=block_adapter.check_forward_pattern,
+                            check_num_outputs=block_adapter.check_num_outputs,
+                            cache_prefix=block_adapter.blocks_name[i][j],
+                            cache_context=block_adapter.unique_blocks_name[i][j],
+                            context_manager=block_adapter.pipe._context_manager,
+                            cache_type=cache_config.cache_type,
+                        )
+                    ]
+                )
+
+            total_cached_blocks.append(unified_blocks_bind_context)
+
+        return total_cached_blocks
+
+
 class BagelCachedAdapter(CachedAdapter):
     """
     Custom CachedAdapter for Bagel that uses BagelCachedContextManager and BagelCachedBlocks.
@@ -1808,6 +2027,7 @@ CUSTOM_DIT_ENABLERS.update(
     {
         "Wan22Pipeline": enable_cache_for_wan22,
         "Wan22I2VPipeline": enable_cache_for_wan22,
+        "Wan22S2VPipeline": enable_cache_for_wan22_s2v,
         "HunyuanImage3Pipeline": enable_cache_for_hunyuan_image3,
         "FluxPipeline": enable_cache_for_flux,
         "Flux2KleinPipeline": enable_cache_for_flux2_klein,
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index 54e93c718..4e3853a2f 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -96,6 +96,8 @@ def _move_tensor_tree_to_cpu(value: object) -> object:
 
 
 def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool) -> int:
+    """Get num_frames for the dummy warmup run. Returns 0 to skip warmup."""
+
     model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
     if model_cls is not None and hasattr(model_cls, "dummy_run_num_frames"):
         return int(getattr(model_cls, "dummy_run_num_frames"))
@@ -722,6 +724,9 @@ class DiffusionEngine:
             prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio
 
         num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input)
+        if num_frames <= 0:
+            logger.info("Skipping dummy warmup run (num_frames=0)")
+            return
         req = OmniDiffusionRequest(
             prompts=[prompt],
             request_id=DUMMY_DIFFUSION_REQUEST_ID,
diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py
index b4fa1209f..4cb714465 100644
--- a/vllm_omni/diffusion/layers/rope.py
+++ b/vllm_omni/diffusion/layers/rope.py
@@ -5,6 +5,7 @@ from einops import rearrange, repeat
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.layers.custom_op import CustomOp
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -300,3 +301,234 @@ def apply_rope_to_qk(
         query = rope(query, cos, sin)
         key = rope(key, cos, sin)
     return query, key
+
+
+class WanS2VRotaryPosEmbed(torch.nn.Module):
+    """Precompute complex-valued RoPE embeddings for S2V multi-grid positions.
+
+    Owns the base frequency buffer and provides forward() to compute position
+    embeddings given hidden_states (for shape) and grid_sizes.
+    """
+
+    def __init__(self, num_heads: int, head_dim: int, max_seq_len: int = 1024):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        d = head_dim
+        freqs = torch.cat(
+            [
+                self._rope_params(max_seq_len, d - 4 * (d // 6)),
+                self._rope_params(max_seq_len, 2 * (d // 6)),
+                self._rope_params(max_seq_len, 2 * (d // 6)),
+            ],
+            dim=1,
+        )
+        self.register_buffer("freqs", freqs.to(torch.complex64), persistent=False)
+
+    @staticmethod
+    @torch.amp.autocast(current_omni_platform.device_type, enabled=False)
+    def _rope_params(max_seq_len, dim, theta=10000):
+        if dim % 2 != 0:
+            raise ValueError(f"dim ({dim}) must be even")
+        freqs = torch.outer(
+            torch.arange(max_seq_len), 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim))
+        )
+        return torch.polar(torch.ones_like(freqs), freqs)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        grid_sizes: list,
+        trainable_freqs: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Precompute RoPE embeddings for the given grid layout.
+
+        Args:
+            hidden_states: Tensor [B, S, ...] (used for batch/seq shape and device)
+            grid_sizes: Grid specification (list of [offsets, sizes, totals])
+            trainable_freqs: Optional trainable frequency overrides for t_f < 0
+
+        Returns:
+            Complex tensor [B, S, 1, head_dim//2] of precomputed position embeddings
+        """
+        b, s = hidden_states.shape[0], hidden_states.shape[1]
+        c = self.head_dim // 2
+        device = hidden_states.device
+
+        freqs = self.freqs.to(device)
+        if trainable_freqs is not None:
+            freqs_input = [freqs, trainable_freqs]
+        else:
+            freqs_input = freqs
+
+        if isinstance(freqs_input, list):
+            trainable_f = freqs_input[1]
+            freqs_split = freqs_input[0]
+        else:
+            trainable_f = None
+            freqs_split = freqs_input
+        freqs_split = freqs_split.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+        output = torch.empty((b, s, 1, c), device=device, dtype=torch.complex64)
+        seq_bucket = [0]
+        if not isinstance(grid_sizes, list):
+            grid_sizes = [grid_sizes]
+        for g in grid_sizes:
+            if not isinstance(g, list):
+                g = [torch.zeros_like(g), g]
+            batch_size = g[0].shape[0]
+            for i in range(batch_size):
+                f_o, h_o, w_o = g[0][i]
+                f, h, w = g[1][i]
+                t_f, t_h, t_w = g[2][i]
+                seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o
+                seq_len = int(seq_f * seq_h * seq_w)
+                if seq_len > 0:
+                    if t_f > 0:
+                        assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0
+                        seq_f_int = int(seq_f)
+                        seq_h_int = int(seq_h)
+                        seq_w_int = int(seq_w)
+
+                        if f_o >= 0:
+                            f_sam = torch.linspace(int(f_o), int(t_f + f_o) - 1, seq_f_int, device=device).long()
+                        else:
+                            f_sam = torch.linspace(int(-f_o), int(-t_f - f_o) + 1, seq_f_int, device=device).long()
+                        h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long()
+                        w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long()
+
+                        freqs_0 = freqs_split[0][f_sam] if f_o >= 0 else freqs_split[0][f_sam].conj()
+                        freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1)
+
+                        freqs_i = torch.cat(
+                            [
+                                freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                freqs_split[1][h_sam]
+                                .view(1, seq_h_int, 1, -1)
+                                .expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                freqs_split[2][w_sam]
+                                .view(1, 1, seq_w_int, -1)
+                                .expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                            ],
+                            dim=-1,
+                        ).reshape(seq_len, 1, -1)
+                    elif t_f < 0:
+                        freqs_i = trainable_f.unsqueeze(1)
+                    output[i, seq_bucket[-1] : seq_bucket[-1] + seq_len] = freqs_i
+            seq_bucket.append(seq_bucket[-1] + seq_len)
+        return output
+
+
+class RotaryEmbeddingWanS2V(RotaryEmbeddingWan):
+    """Apply RoPE using precomputed complex freqs for Wan S2V main transformer.
+
+    Converts complex freqs (from WanS2VRotaryPosEmbed) to cos/sin and delegates
+    to RotaryEmbeddingWan for platform-optimized application (float32 kernel).
+    Under TP, freqs has 1 head — broadcasts automatically via cos/sin.
+    """
+
+    def __init__(self) -> None:
+        super().__init__(is_neox_style=False, half_head_dim=True)
+
+    def forward(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+        freqs_sliced = freqs[:, : x.size(1)]
+        cos = freqs_sliced.real.to(x.dtype)
+        sin = freqs_sliced.imag.to(x.dtype)
+        return super().forward(x, cos, sin)
+
+
+class RotaryEmbeddingS2VGrid(torch.nn.Module):
+    """Grid-based RoPE for S2V motioner/init attention.
+
+    Applies complex-valued rotary embeddings using 3D grid sampling
+    (frame, height, width). Used by SimpleSelfAttention, SwinSelfAttention,
+    CausalSelfAttention in motioner blocks.
+    """
+
+    @staticmethod
+    @torch.amp.autocast(current_omni_platform.device_type, enabled=False)
+    def precompute(
+        seq_len: int, num_heads: int, head_dim: int, grid_sizes, freqs: torch.Tensor, device: torch.device, start=None
+    ) -> torch.Tensor:
+        """Precompute position frequency tensor from grid specification.
+
+        Returns a complex tensor that can be reused across layers via apply_precomputed().
+        """
+        c = head_dim // 2
+
+        trainable_freqs = None
+        if isinstance(freqs, list):
+            trainable_freqs = freqs[1]
+            freqs = freqs[0]
+        freqs = freqs.to(device)
+        freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+        if not isinstance(grid_sizes, list):
+            grid_sizes = [grid_sizes]
+
+        batch_size = grid_sizes[0][0].shape[0] if isinstance(grid_sizes[0], list) else grid_sizes[0].shape[0]
+        precomputed = torch.empty((batch_size, seq_len, 1, c), device=device, dtype=torch.complex64)
+        seq_bucket = [0]
+
+        for g in grid_sizes:
+            if not isinstance(g, list):
+                g = [torch.zeros_like(g), g]
+            g_batch_size = g[0].shape[0]
+            for i in range(g_batch_size):
+                if start is None:
+                    f_o, h_o, w_o = g[0][i]
+                else:
+                    f_o, h_o, w_o = start[i]
+
+                f, h, w = g[1][i]
+                t_f, t_h, t_w = g[2][i]
+                seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o
+                seg_len = int(seq_f * seq_h * seq_w)
+                if seg_len > 0:
+                    if t_f > 0:
+                        seq_f_int = int(seq_f)
+                        seq_h_int = int(seq_h)
+                        seq_w_int = int(seq_w)
+
+                        if f_o >= 0:
+                            f_sam = torch.linspace(int(f_o), int(t_f + f_o) - 1, seq_f_int, device=device).long()
+                        else:
+                            f_sam = torch.linspace(int(-f_o), int(-t_f - f_o) + 1, seq_f_int, device=device).long()
+                        h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long()
+                        w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long()
+
+                        freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj()
+                        freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1)
+
+                        freqs_i = torch.cat(
+                            [
+                                freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                freqs[1][h_sam].view(1, seq_h_int, 1, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                freqs[2][w_sam].view(1, 1, seq_w_int, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                            ],
+                            dim=-1,
+                        ).reshape(seg_len, 1, -1)
+                    elif t_f < 0:
+                        freqs_i = trainable_freqs.unsqueeze(1)
+                    precomputed[i, seq_bucket[-1] : seq_bucket[-1] + seg_len] = freqs_i
+            seq_bucket.append(seq_bucket[-1] + seg_len)
+        return precomputed
+
+    @staticmethod
+    @torch.amp.autocast(current_omni_platform.device_type, enabled=False)
+    def apply_precomputed(x: torch.Tensor, precomputed_freqs: torch.Tensor) -> torch.Tensor:
+        """Apply precomputed position frequencies to input tensor."""
+        n = x.size(2)
+        input_dtype = x.dtype
+        seq_len = x.size(1)
+        precomputed_freqs = precomputed_freqs[:, :seq_len]
+        x_c = torch.view_as_complex(x.to(torch.float64).reshape(x.size(0), seq_len, n, -1, 2))
+        x_c = torch.view_as_real(x_c * precomputed_freqs).flatten(3)
+        return x_c.reshape(x.size(0), seq_len, n, -1).to(input_dtype)
+
+    @staticmethod
+    @torch.amp.autocast(current_omni_platform.device_type, enabled=False)
+    def forward(x: torch.Tensor, grid_sizes, freqs: torch.Tensor, start=None) -> torch.Tensor:
+        n, d = x.size(2), x.size(3)
+        precomputed = RotaryEmbeddingS2VGrid.precompute(x.size(1), n, d, grid_sizes, freqs, x.device, start=start)
+        return RotaryEmbeddingS2VGrid.apply_precomputed(x, precomputed)
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py
index 0817d6f48..1bd32bb42 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py
@@ -13,7 +13,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any
+from typing import Any, ClassVar
 
 import numpy as np
 import PIL.Image
@@ -482,6 +482,7 @@ class Wan22S2VPipeline(
     _DEFAULT_MOTION_FRAMES = 73
     _DEFAULT_INFER_FRAMES = 80
     _DEFAULT_FPS = 16
+    dummy_run_num_frames: ClassVar[int] = 0
 
     def __init__(
         self,
@@ -929,21 +930,19 @@ class Wan22S2VPipeline(
     def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
         """Forward pass through the S2V transformer to predict noise.
 
-        The S2V model returns a list of tensors (one per batch element).
-        We take the first element since S2V processes one sample at a time.
+        With return_dict=False, the model returns (output,) where output is [B, C, T, H, W].
+        We squeeze batch dim since S2V processes one sample at a time.
         """
         if current_model is None:
             current_model = self.transformer
-        # WanModel_S2V's norm layers compute in float32; autocast ensures
-        # the float32→bfloat16 casts happen automatically (matching the
-        # original Wan2.2 inference path).
         param_dtype = next(current_model.parameters()).dtype
         with torch.amp.autocast(self.device.type, dtype=param_dtype):
             result = current_model(**kwargs)
-        # WanModel_S2V.forward returns a list of tensors
         if isinstance(result, list):
             return result[0]
-        return result[0] if isinstance(result, tuple) else result
+        if isinstance(result, tuple):
+            return result[0].squeeze(0)
+        return result.sample.squeeze(0)
 
     def diffuse(
         self,
@@ -996,32 +995,29 @@ class Wan22S2VPipeline(
                 latent_model_input = [latents.to(device)]
                 timestep = t.unsqueeze(0).to(device) if t.dim() == 0 else t.to(device)
 
-                # -- Positive (conditional) prediction kwargs --
                 positive_kwargs = {
-                    "x": latent_model_input,
-                    "t": timestep,
-                    "context": prompt_embeds[0:1],
-                    "seq_len": max_seq_len,
+                    "hidden_states": latent_model_input,
+                    "timestep": timestep,
+                    "encoder_hidden_states": prompt_embeds[0:1],
                     "cond_states": cond_latents,
                     "motion_latents": input_motion_latents,
                     "ref_latents": ref_latents,
-                    "motion_frames": motion_frames,
                     "drop_motion_frames": drop_first_motion,
-                    "audio_emb": positive_audio_emb,
+                    "encoder_hidden_states_audio": positive_audio_emb,
+                    "return_dict": False,
                 }
 
                 if do_true_cfg:
                     negative_kwargs = {
-                        "x": latent_model_input,
-                        "t": timestep,
-                        "context": negative_prompt_embeds[0:1],
-                        "seq_len": max_seq_len,
+                        "hidden_states": latent_model_input,
+                        "timestep": timestep,
+                        "encoder_hidden_states": negative_prompt_embeds[0:1],
                         "cond_states": cond_latents,
                         "motion_latents": input_motion_latents,
                         "ref_latents": ref_latents,
-                        "motion_frames": motion_frames,
                         "drop_motion_frames": drop_first_motion,
-                        "audio_emb": negative_audio_emb,
+                        "encoder_hidden_states_audio": negative_audio_emb,
+                        "return_dict": False,
                     }
                 else:
                     negative_kwargs = None
@@ -1174,7 +1170,7 @@ class Wan22S2VPipeline(
         ref_latents = self.encode_ref_image(image, height, width, device=device).to(dtype=dtype)
 
         # ---- 4. Initial motion latents (zeros) ----
-        motion_frames = self.motion_frames
+        motion_frames = min(self.motion_frames, infer_frames - 1)
         motion_pixels = torch.zeros(
             [1, 3, motion_frames, height, width],
             dtype=dtype,
@@ -1288,6 +1284,10 @@ class Wan22S2VPipeline(
             )
 
             # ---- Decode this clip ----
+            if self.od_config.enable_cpu_offload:
+                self.transformer.to("cpu")
+                current_omni_platform.empty_cache()
+
             latents_for_decode = latents.unsqueeze(0)  # [1, C, T, H, W]
             if not (drop_first_motion and r == 0):
                 decode_latents = torch.cat([motion_latents, latents_for_decode], dim=2)
diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py
index cc0bb78c1..8fef41f74 100644
--- a/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py
+++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py
@@ -10,13 +10,13 @@ import math
 from collections.abc import Iterable
 from copy import deepcopy
 
-import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from diffusers.models.attention import AdaLayerNorm
+from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps
 from diffusers.models.normalization import FP32LayerNorm
-from einops import rearrange, repeat
+from einops import rearrange
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
@@ -31,6 +31,7 @@ from vllm.model_executor.layers.linear import (
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.layers.rope import RotaryEmbeddingS2VGrid, RotaryEmbeddingWanS2V, WanS2VRotaryPosEmbed
 from vllm_omni.platforms import current_omni_platform
 
 from .wan2_2_transformer import DistributedRMSNorm, WanFeedForward
@@ -53,170 +54,6 @@ def sinusoidal_embedding_1d(dim, position):
     return x
 
 
-@torch.amp.autocast(current_omni_platform.device_type, enabled=False)
-def rope_params(max_seq_len, dim, theta=10000):
-    if dim % 2 != 0:
-        raise ValueError(f"dim ({dim}) must be even")
-    freqs = torch.outer(
-        torch.arange(max_seq_len), 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim))
-    )
-    freqs = torch.polar(torch.ones_like(freqs), freqs)
-    return freqs
-
-
-def rope_precompute(x, grid_sizes, freqs, start=None):
-    """Precompute complex-valued RoPE embeddings for S2V multi-grid positions."""
-    b, s, n, c = x.size(0), x.size(1), x.size(2), x.size(3) // 2
-
-    if isinstance(freqs, list):
-        trainable_freqs = freqs[1]
-        freqs = freqs[0]
-    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
-
-    output = torch.view_as_complex(x.detach().reshape(b, s, n, -1, 2).to(torch.float64))
-    seq_bucket = [0]
-    if not isinstance(grid_sizes, list):
-        grid_sizes = [grid_sizes]
-    for g in grid_sizes:
-        if not isinstance(g, list):
-            g = [torch.zeros_like(g), g]
-        batch_size = g[0].shape[0]
-        for i in range(batch_size):
-            if start is None:
-                f_o, h_o, w_o = g[0][i]
-            else:
-                f_o, h_o, w_o = start[i]
-
-            f, h, w = g[1][i]
-            t_f, t_h, t_w = g[2][i]
-            seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o
-            seq_len = int(seq_f * seq_h * seq_w)
-            if seq_len > 0:
-                if t_f > 0:
-                    if f_o >= 0:
-                        f_sam = np.linspace(f_o.item(), (t_f + f_o).item() - 1, seq_f).astype(int).tolist()
-                    else:
-                        f_sam = np.linspace(-f_o.item(), (-t_f - f_o).item() + 1, seq_f).astype(int).tolist()
-                    h_sam = np.linspace(h_o.item(), (t_h + h_o).item() - 1, seq_h).astype(int).tolist()
-                    w_sam = np.linspace(w_o.item(), (t_w + w_o).item() - 1, seq_w).astype(int).tolist()
-
-                    assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0
-                    freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj()
-                    freqs_0 = freqs_0.view(seq_f, 1, 1, -1)
-
-                    freqs_i = torch.cat(
-                        [
-                            freqs_0.expand(seq_f, seq_h, seq_w, -1),
-                            freqs[1][h_sam].view(1, seq_h, 1, -1).expand(seq_f, seq_h, seq_w, -1),
-                            freqs[2][w_sam].view(1, 1, seq_w, -1).expand(seq_f, seq_h, seq_w, -1),
-                        ],
-                        dim=-1,
-                    ).reshape(seq_len, 1, -1)
-                elif t_f < 0:
-                    freqs_i = trainable_freqs.unsqueeze(1)
-                output[i, seq_bucket[-1] : seq_bucket[-1] + seq_len] = freqs_i
-        seq_bucket.append(seq_bucket[-1] + seq_len)
-    return output
-
-
-@torch.amp.autocast(current_omni_platform.device_type, enabled=False)
-def rope_apply(x, grid_sizes, freqs, start=None):
-    """Apply RoPE using grid-based frequency computation (for motioner/init)."""
-    n, c = x.size(2), x.size(3) // 2
-    input_dtype = x.dtype
-
-    if isinstance(freqs, list):
-        trainable_freqs = freqs[1]
-        freqs = freqs[0]
-    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
-
-    output = x.clone()
-    seq_bucket = [0]
-    if not isinstance(grid_sizes, list):
-        grid_sizes = [grid_sizes]
-    for g in grid_sizes:
-        if not isinstance(g, list):
-            g = [torch.zeros_like(g), g]
-        batch_size = g[0].shape[0]
-        for i in range(batch_size):
-            if start is None:
-                f_o, h_o, w_o = g[0][i]
-            else:
-                f_o, h_o, w_o = start[i]
-
-            f, h, w = g[1][i]
-            t_f, t_h, t_w = g[2][i]
-            seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o
-            seq_len = int(seq_f * seq_h * seq_w)
-            if seq_len > 0:
-                if t_f > 0:
-                    if f_o >= 0:
-                        f_sam = np.linspace(f_o.item(), (t_f + f_o).item() - 1, seq_f).astype(int).tolist()
-                    else:
-                        f_sam = np.linspace(-f_o.item(), (-t_f - f_o).item() + 1, seq_f).astype(int).tolist()
-                    h_sam = np.linspace(h_o.item(), (t_h + h_o).item() - 1, seq_h).astype(int).tolist()
-                    w_sam = np.linspace(w_o.item(), (t_w + w_o).item() - 1, seq_w).astype(int).tolist()
-
-                    assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0
-                    freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj()
-                    freqs_0 = freqs_0.view(seq_f, 1, 1, -1)
-
-                    freqs_i = torch.cat(
-                        [
-                            freqs_0.expand(seq_f, seq_h, seq_w, -1),
-                            freqs[1][h_sam].view(1, seq_h, 1, -1).expand(seq_f, seq_h, seq_w, -1),
-                            freqs[2][w_sam].view(1, 1, seq_w, -1).expand(seq_f, seq_h, seq_w, -1),
-                        ],
-                        dim=-1,
-                    ).reshape(seq_len, 1, -1)
-                elif t_f < 0:
-                    freqs_i = trainable_freqs.unsqueeze(1)
-                x_i = torch.view_as_complex(
-                    x[i, seq_bucket[-1] : seq_bucket[-1] + seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)
-                )
-                x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
-                output[i, seq_bucket[-1] : seq_bucket[-1] + seq_len] = x_i
-        seq_bucket.append(seq_bucket[-1] + seq_len)
-    return output.to(input_dtype)
-
-
-@torch.compiler.disable
-@torch.amp.autocast(current_omni_platform.device_type, enabled=False)
-def rope_apply_s2v(x, grid_sizes, freqs, start=None):
-    """Apply RoPE using precomputed complex freqs (for S2V main transformer).
-
-    Under TP, x has fewer heads than freqs (which is precomputed with all heads).
-    Since RoPE frequencies are identical across heads, we slice freqs to match.
-    """
-    n = x.size(2)
-    input_dtype = x.dtype
-    output = []
-    for i, _ in enumerate(x):
-        s = x.size(1)
-        x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(s, n, -1, 2))
-        freqs_i = freqs[i, :s, :n]
-        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
-        x_i = torch.cat([x_i, x[i, s:]])
-        output.append(x_i)
-    return torch.stack(output).to(input_dtype)
-
-
-@torch.amp.autocast(current_omni_platform.device_type, enabled=False)
-def rope_apply_usp(x, grid_sizes, freqs):
-    """Apply RoPE for Ulysses sequence parallel (context parallel mode)."""
-    n = x.size(2)
-    input_dtype = x.dtype
-    output = []
-    for i, _ in enumerate(x):
-        s = x.size(1)
-        x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(s, n, -1, 2))
-        freqs_i = freqs[i, :, :n]
-        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
-        x_i = torch.cat([x_i, x[i, s:]])
-        output.append(x_i)
-    return torch.stack(output).to(input_dtype)
-
-
 def zero_module(module):
     for p in module.parameters():
         p.detach().zero_()
@@ -292,6 +129,8 @@ class WanS2VSelfAttention(nn.Module):
             causal=False,
         )
 
+        self.rotary_embedding = RotaryEmbeddingWanS2V()
+
     def forward(self, x, seq_lens, grid_sizes, freqs):
         b, s = x.shape[:2]
 
@@ -311,15 +150,12 @@ class WanS2VSelfAttention(nn.Module):
         value = value.view(b, s, self.num_kv_heads, self.head_dim)
 
         # Apply S2V complex-valued RoPE
-        query = rope_apply_s2v(query, grid_sizes, freqs)
-        key = rope_apply_s2v(key, grid_sizes, freqs)
+        query = self.rotary_embedding(query, freqs)
+        key = self.rotary_embedding(key, freqs)
 
-        # Reshape back after rope (rope_apply_s2v preserves [B, S, N, D])
+        # Reshape back after RoPE (preserves [B, S, N, D])
         query = query.view(b, s, self.num_heads, self.head_dim)
         key = key.view(b, s, self.num_kv_heads, self.head_dim)
-
-        # Attention — rope_apply_s2v preserves input dtype (bf16), so
-        # query/key/value are all bf16 here, matching FlashAttention requirements.
         hidden_states = self.attn(query, key, value)
         hidden_states = hidden_states.flatten(2, 3)
         hidden_states = hidden_states.type_as(x)
@@ -411,7 +247,24 @@ class WanS2VTransformerBlock(nn.Module):
         # 6-way scale-shift modulation
         self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
 
-    def forward(self, x, e, seq_lens, grid_sizes, freqs, context, context_lens):
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        seq_lens=None,
+        grid_sizes=None,
+        freqs=None,
+        context=None,
+        context_lens=None,
+        e=None,
+    ):
+        x = hidden_states
+        if encoder_hidden_states is None:
+            encoder_hidden_states = e
+        if encoder_hidden_states is None:
+            raise ValueError("S2V transformer block requires modulation states.")
+
+        e = encoder_hidden_states
         seg_boundary = e[1]
         seg_idx = [0, min(max(0, seg_boundary), x.size(1)), x.size(1)]
         e = e[0]
@@ -450,25 +303,68 @@ class WanS2VTransformerBlock(nn.Module):
         return x
 
 
-class WanS2VHead(nn.Module):
-    """S2V output head with modulation."""
+class WanS2VConditionEmbedder(nn.Module):
+    """Combined time and text condition embeddings for S2V."""
 
-    def __init__(self, dim, out_dim, patch_size, eps=1e-6):
+    def __init__(self, dim: int, freq_dim: int, text_dim: int, text_len: int):
+        super().__init__()
+        self.text_len = text_len
+        self.timesteps_proj = Timesteps(num_channels=freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.time_embedder = TimestepEmbedding(in_channels=freq_dim, time_embed_dim=dim)
+        self.act_fn = nn.SiLU()
+        self.time_proj = nn.Linear(dim, dim * 6)
+        self.text_embedder = PixArtAlphaTextProjection(text_dim, dim, act_fn="gelu_tanh")
+
+    def forward(
+        self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, hidden_states_dtype: torch.dtype
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        timestep = self.timesteps_proj(timestep)
+        temb = self.time_embedder(timestep.to(hidden_states_dtype))
+        timestep_proj = self.time_proj(self.act_fn(temb))
+        encoder_hidden_states = self.text_embedder(
+            torch.stack(
+                [torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in encoder_hidden_states]
+            )
+        )
+        return temb, timestep_proj, encoder_hidden_states
+
+
+class WanS2VTimestepProjPrepare(nn.Module):
+    """Prepares timestep_proj with segment modulation for S2V."""
+
+    def __init__(self, dim: int):
         super().__init__()
         self.dim = dim
-        self.out_dim = out_dim
-        self.patch_size = patch_size
 
-        out_dim_full = math.prod(patch_size) * out_dim
-        self.norm = FP32LayerNorm(dim, eps, elementwise_affine=False)
-        self.head = nn.Linear(dim, out_dim_full)
-        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+    def forward(self, timestep_proj: torch.Tensor, temb: torch.Tensor, zero_timestep: bool, seg_boundary: int):
+        timestep_proj = timestep_proj.unflatten(1, (6, self.dim))
+        if zero_timestep:
+            temb = temb[:-1]
+            zero_timestep_proj = timestep_proj[-1:]
+            timestep_proj = timestep_proj[:-1]
+            timestep_proj = torch.cat(
+                [
+                    timestep_proj.unsqueeze(2),
+                    zero_timestep_proj.unsqueeze(2).repeat(timestep_proj.size(0), 1, 1, 1),
+                ],
+                dim=2,
+            )
+            return temb, [timestep_proj, seg_boundary]
+        else:
+            timestep_proj = timestep_proj.unsqueeze(2).repeat(1, 1, 2, 1)
+            return temb, [timestep_proj, 0]
 
-    def forward(self, x, e):
-        e = (self.modulation.to(e.dtype) + e.unsqueeze(1)).chunk(2, dim=1)
-        x = self.norm(x).type_as(x) * (1 + e[1]) + e[0]
-        x = self.head(x)
-        return x
+
+class WanS2VOutputScaleShiftPrepare(nn.Module):
+    """Prepares output scale/shift from temb for S2V output modulation."""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        shift, scale = (self.scale_shift_table.to(temb.dtype) + temb.unsqueeze(1)).chunk(2, dim=1)
+        return shift, scale
 
 
 # ---------------------------------------------------------------------------
@@ -573,43 +469,10 @@ class CausalAudioEncoder(nn.Module):
         return res
 
 
-class AudioCrossAttention(nn.Module):
-    """Audio cross-attention using nn.Linear (no TP needed)."""
+class AudioCrossAttention(WanS2VCrossAttention):
+    """Audio cross-attention — same as WanS2VCrossAttention (TP-aware)."""
 
-    def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6):
-        if dim % num_heads != 0:
-            raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})")
-        super().__init__()
-        self.dim = dim
-        self.num_heads = num_heads
-        self.head_dim = dim // num_heads
-
-        self.q = nn.Linear(dim, dim)
-        self.k = nn.Linear(dim, dim)
-        self.v = nn.Linear(dim, dim)
-        self.o = nn.Linear(dim, dim)
-        self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-
-        self.attn = Attention(
-            num_heads=num_heads,
-            head_size=self.head_dim,
-            num_kv_heads=num_heads,
-            softmax_scale=1.0 / (self.head_dim**0.5),
-            causal=False,
-        )
-
-    def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: list[int] | None = None) -> torch.Tensor:
-        b, n, d = x.size(0), self.num_heads, self.head_dim
-
-        q = self.norm_q(self.q(x)).view(b, -1, n, d)
-        k = self.norm_k(self.k(context)).view(b, -1, n, d)
-        v = self.v(context).view(b, -1, n, d)
-
-        x = self.attn(q, k, v)
-        x = x.flatten(2, 3)
-        x = self.o(x)
-        return x
+    pass
 
 
 class AudioInjector_WAN(nn.Module):
@@ -692,8 +555,15 @@ class SimpleSelfAttention(nn.Module):
             causal=False,
         )
 
+        self.rotary_embedding = RotaryEmbeddingS2VGrid()
+
     def forward(
-        self, x: torch.Tensor, seq_lens: list[int], grid_sizes: list[tuple[int, int, int]], freqs: torch.Tensor
+        self,
+        x: torch.Tensor,
+        seq_lens: list[int],
+        grid_sizes: list[tuple[int, int, int]],
+        freqs: torch.Tensor,
+        precomputed_freqs: torch.Tensor | None = None,
     ) -> torch.Tensor:
         b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
 
@@ -701,10 +571,13 @@ class SimpleSelfAttention(nn.Module):
         k = self.norm_k(self.k(x)).view(b, s, n, d)
         v = self.v(x).view(b, s, n, d)
 
-        q = rope_apply(q, grid_sizes, freqs)
-        k = rope_apply(k, grid_sizes, freqs)
+        if precomputed_freqs is not None:
+            q = RotaryEmbeddingS2VGrid.apply_precomputed(q, precomputed_freqs)
+            k = RotaryEmbeddingS2VGrid.apply_precomputed(k, precomputed_freqs)
+        else:
+            q = self.rotary_embedding(q, grid_sizes, freqs)
+            k = self.rotary_embedding(k, grid_sizes, freqs)
 
-        # rope_apply preserves [B, S, N, D] shape and input dtype
         q = q.view(b, s, n, d)
         k = k.view(b, s, n, d)
 
@@ -716,45 +589,54 @@ class SimpleSelfAttention(nn.Module):
 
 class SwinSelfAttention(SimpleSelfAttention):
     def forward(
-        self, x: torch.Tensor, seq_lens: list[int], grid_sizes: list[tuple[int, int, int]], freqs: torch.Tensor
+        self,
+        x: torch.Tensor,
+        seq_lens: list[int],
+        grid_sizes: list[tuple[int, int, int]],
+        freqs: torch.Tensor,
+        precomputed_freqs: torch.Tensor | None = None,
     ) -> torch.Tensor:
         b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
-        if b != 1:
-            raise ValueError(
-                f"SwinSelfAttention only supports batch_size=1, got batch_size={b}. "
-                "Batched inference requires refactoring the frame reference and rearrange logic."
-            )
 
         q = self.norm_q(self.q(x)).view(b, s, n, d)
         k = self.norm_k(self.k(x)).view(b, s, n, d)
         v = self.v(x).view(b, s, n, d)
 
-        q = rope_apply(q, grid_sizes, freqs)
-        k = rope_apply(k, grid_sizes, freqs)
+        if precomputed_freqs is not None:
+            q = RotaryEmbeddingS2VGrid.apply_precomputed(q, precomputed_freqs)
+            k = RotaryEmbeddingS2VGrid.apply_precomputed(k, precomputed_freqs)
+        else:
+            q = self.rotary_embedding(q, grid_sizes, freqs)
+            k = self.rotary_embedding(k, grid_sizes, freqs)
         q = q.view(b, s, n, d)
         k = k.view(b, s, n, d)
         T, H, W = grid_sizes[0].tolist()
 
-        q = rearrange(q, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W)
-        k = rearrange(k, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W)
-        v = rearrange(v, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W)
+        q = rearrange(q, "b (t h w) n d -> b t (h w) n d", t=T, h=H, w=W)
+        k = rearrange(k, "b (t h w) n d -> b t (h w) n d", t=T, h=H, w=W)
+        v = rearrange(v, "b (t h w) n d -> b t (h w) n d", t=T, h=H, w=W)
 
-        # Skip ref_q (last frame query) - unused in current implementation
-        q = q[:-1]
+        q_frames = q[:, :-1]
+        ref_k = k[:, -1:].expand(-1, T - 1, -1, -1, -1)
+        ref_v = v[:, -1:].expand(-1, T - 1, -1, -1, -1)
 
-        ref_k = repeat(k[-1:], "1 s n d -> t s n d", t=k.shape[0] - 1)
-        k = k[:-1]
-        k = torch.cat([k[:1], k, k[-1:]])
-        k = torch.cat([k[1:-1], k[2:], k[:-2], ref_k], dim=1)
+        k_frames = k[:, :-1]
+        k_pad = torch.cat([k_frames[:, :1], k_frames, k_frames[:, -1:]], dim=1)
+        k_frames = torch.cat([k_pad[:, 1:-1], k_pad[:, 2:], k_pad[:, :-2], ref_k], dim=2)
 
-        ref_v = repeat(v[-1:], "1 s n d -> t s n d", t=v.shape[0] - 1)
-        v = v[:-1]
-        v = torch.cat([v[:1], v, v[-1:]])
-        v = torch.cat([v[1:-1], v[2:], v[:-2], ref_v], dim=1)
+        v_frames = v[:, :-1]
+        v_pad = torch.cat([v_frames[:, :1], v_frames, v_frames[:, -1:]], dim=1)
+        v_frames = torch.cat([v_pad[:, 1:-1], v_pad[:, 2:], v_pad[:, :-2], ref_v], dim=2)
 
-        out = self.attn(q, k, v)
-        out = torch.cat([out, ref_v[:1]], axis=0)
-        out = rearrange(out, "(b t) (h w) n d -> b (t h w) n d", t=T, h=H, w=W)
+        q_flat = rearrange(q_frames, "b t s n d -> (b t) s n d")
+        k_flat = rearrange(k_frames, "b t s n d -> (b t) s n d")
+        v_flat = rearrange(v_frames, "b t s n d -> (b t) s n d")
+
+        out = self.attn(q_flat, k_flat, v_flat)
+        out = rearrange(out, "(b t) s n d -> b t s n d", b=b)
+        ref_out = v[:, -1:]
+        out = torch.cat([out, ref_out], dim=1)
+        out = rearrange(out, "b t (h w) n d -> b (t h w) n d", h=H, w=W)
         x = out.flatten(2, 3)
         x = self.o(x)
         return x
@@ -766,11 +648,6 @@ class CausalSelfAttention(SimpleSelfAttention):
     ) -> torch.Tensor:
         shifting = 3
         b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
-        if b != 1:
-            raise ValueError(
-                f"CausalSelfAttention only supports batch_size=1, got batch_size={b}. "
-                "Batched inference requires refactoring the causal masking and frame reference logic."
-            )
 
         q = self.norm_q(self.q(x)).view(b, s, n, d)
         k = self.norm_k(self.k(x)).view(b, s, n, d)
@@ -778,55 +655,61 @@ class CausalSelfAttention(SimpleSelfAttention):
 
         T, H, W = grid_sizes[0].tolist()
 
-        q = rearrange(q, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W)
-        k = rearrange(k, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W)
-        v = rearrange(v, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W)
+        q = rearrange(q, "b (t h w) n d -> b t (h w) n d", t=T, h=H, w=W)
+        k = rearrange(k, "b (t h w) n d -> b t (h w) n d", t=T, h=H, w=W)
+        v = rearrange(v, "b (t h w) n d -> b t (h w) n d", t=T, h=H, w=W)
 
-        # Skip ref_q (last frame query) - unused in current causal implementation
-        q = q[:-1]
+        num_frames = T - 1
 
-        grid_sizes_q = torch.tensor([[1, H, W]] * q.shape[0], dtype=torch.long)
-        start = [[shifting, 0, 0]] * q.shape[0]
-        q = rope_apply(q, grid_sizes_q, freqs, start=start)
-        q = q.view(q.shape[0], -1, n, d)
+        # Query: skip ref frame, apply RoPE per frame
+        q_frames = rearrange(q[:, :-1], "b t s n d -> (b t) s n d")
+        grid_sizes_q = torch.tensor([[1, H, W]] * (b * num_frames), dtype=torch.long)
+        start = [[shifting, 0, 0]] * (b * num_frames)
+        q_frames = self.rotary_embedding(q_frames, grid_sizes_q, freqs, start=start)
+        q_frames = q_frames.view(b * num_frames, -1, n, d)
 
-        ref_k = k[-1:]
-        grid_sizes_ref = torch.tensor([[1, H, W]], dtype=torch.long)
-        start_ref = [[shifting + 10, 0, 0]]
-        ref_k = rope_apply(ref_k, grid_sizes_ref, freqs, start_ref)
-        ref_k = ref_k.view(1, -1, n, d)
-        ref_k = repeat(ref_k, "1 s n d -> t s n d", t=k.shape[0] - 1)
+        # Ref key: one per batch, RoPE with shifted start, expand to all frames
+        ref_k = rearrange(k[:, -1:], "b 1 s n d -> b s n d")
+        grid_sizes_ref = torch.tensor([[1, H, W]] * b, dtype=torch.long)
+        start_ref = [[shifting + 10, 0, 0]] * b
+        ref_k = self.rotary_embedding(ref_k, grid_sizes_ref, freqs, start_ref)
+        ref_k = ref_k.view(b, 1, -1, n, d).expand(-1, num_frames, -1, -1, -1)
+        ref_k = rearrange(ref_k, "b t s n d -> (b t) s n d")
 
-        k = k[:-1]
-        k = torch.cat([*([k[:1]] * shifting), k])
+        # Causal key: pad with shifting duplicates, gather shifting+1 neighboring frames
+        k_frames = k[:, :-1]
+        k_pad = torch.cat([k_frames[:, :1].expand(-1, shifting, -1, -1, -1), k_frames], dim=1)
         cat_k = []
         for i in range(shifting):
-            cat_k.append(k[i : i - shifting])
-        cat_k.append(k[shifting:])
-        k = torch.cat(cat_k, dim=1)
+            cat_k.append(k_pad[:, i : i + num_frames])
+        cat_k.append(k_pad[:, shifting:])
+        k_causal = torch.cat(cat_k, dim=2)
+        k_causal = rearrange(k_causal, "b t s n d -> (b t) s n d")
 
-        grid_sizes_k = torch.tensor([[shifting + 1, H, W]] * (q.shape[0]), dtype=torch.long)
-        k = rope_apply(k, grid_sizes_k, freqs)
-        k = k.view(k.shape[0], -1, n, d)
-        k = torch.cat([k, ref_k], dim=1)
+        grid_sizes_k = torch.tensor([[shifting + 1, H, W]] * (b * num_frames), dtype=torch.long)
+        k_causal = self.rotary_embedding(k_causal, grid_sizes_k, freqs)
+        k_causal = k_causal.view(b * num_frames, -1, n, d)
+        k_causal = torch.cat([k_causal, ref_k], dim=1)
 
-        ref_v = repeat(v[-1:], "1 s n d -> t s n d", t=q.shape[0])
-        v = v[:-1]
-        v = torch.cat([*([v[:1]] * shifting), v])
+        # Causal value: same padding structure as key, plus ref
+        ref_v = v[:, -1:].expand(-1, num_frames, -1, -1, -1)
+        ref_v_flat = rearrange(ref_v, "b t s n d -> (b t) s n d")
+        v_frames = v[:, :-1]
+        v_pad = torch.cat([v_frames[:, :1].expand(-1, shifting, -1, -1, -1), v_frames], dim=1)
         cat_v = []
         for i in range(shifting):
-            cat_v.append(v[i : i - shifting])
-        cat_v.append(v[shifting:])
-        v = torch.cat(cat_v, dim=1)
-        v = torch.cat([v, ref_v], dim=1)
+            cat_v.append(v_pad[:, i : i + num_frames])
+        cat_v.append(v_pad[:, shifting:])
+        v_causal = torch.cat(cat_v, dim=2)
+        v_causal = rearrange(v_causal, "b t s n d -> (b t) s n d")
+        v_causal = torch.cat([v_causal, ref_v_flat], dim=1)
 
-        outs = []
-        for i in range(q.shape[0]):
-            out = self.attn(q[i : i + 1], k[i : i + 1], v[i : i + 1])
-            outs.append(out)
-        out = torch.cat(outs, dim=0)
-        out = torch.cat([out, ref_v[:1]], axis=0)
-        out = rearrange(out, "(b t) (h w) n d -> b (t h w) n d", t=T, h=H, w=W)
+        # Batched attention (all frames at once)
+        out = self.attn(q_frames, k_causal, v_causal)
+        out = rearrange(out, "(b t) s n d -> b t s n d", b=b)
+        ref_out = v[:, -1:]
+        out = torch.cat([out, ref_out], dim=1)
+        out = rearrange(out, "b t (h w) n d -> b (t h w) n d", h=H, w=W)
         x = out.flatten(2, 3)
         x = self.o(x)
         return x
@@ -861,9 +744,14 @@ class MotionerAttentionBlock(nn.Module):
         self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim))
 
     def forward(
-        self, x: torch.Tensor, seq_lens: list[int], grid_sizes: list[tuple[int, int, int]], freqs: torch.Tensor
+        self,
+        x: torch.Tensor,
+        seq_lens: list[int],
+        grid_sizes: list[tuple[int, int, int]],
+        freqs: torch.Tensor,
+        precomputed_freqs: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        y = self.self_attn(self.norm1(x).type_as(x), seq_lens, grid_sizes, freqs)
+        y = self.self_attn(self.norm1(x).type_as(x), seq_lens, grid_sizes, freqs, precomputed_freqs=precomputed_freqs)
         x = x + y
         y = self.ffn(self.norm2(x).type_as(x))
         x = x + y
@@ -934,11 +822,8 @@ class MotionerTransformers(nn.Module):
             raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})")
         if (dim // num_heads) % 2 != 0:
             raise ValueError(f"dim // num_heads ({dim // num_heads}) must be even")
-        d = dim // num_heads
-        self.freqs = torch.cat(
-            [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))],
-            dim=1,
-        )
+        self.rope = WanS2VRotaryPosEmbed(num_heads, dim // num_heads)
+        self.freqs = self.rope.freqs
 
         self.gradient_checkpointing = False
         self.motion_side_len = int(math.sqrt(motion_token_num))
@@ -948,6 +833,7 @@ class MotionerTransformers(nn.Module):
 
         self.trainable_token_pos_emb = trainable_token_pos_emb
         if trainable_token_pos_emb:
+            d = dim // num_heads
             x = torch.zeros([1, motion_token_num, num_heads, d])
             x[..., ::2] = 1
             gride_sizes = [
@@ -957,13 +843,13 @@ class MotionerTransformers(nn.Module):
                     torch.tensor([1, self.motion_side_len, self.motion_side_len]).unsqueeze(0).repeat(1, 1),
                 ]
             ]
-            token_freqs = rope_apply(x, gride_sizes, self.freqs)
+            token_freqs = RotaryEmbeddingS2VGrid.forward(x, gride_sizes, self.rope.freqs)
             token_freqs = token_freqs[0, :, 0].reshape(motion_token_num, -1, 2)
             token_freqs = token_freqs * 0.01
             self.token_freqs = torch.nn.Parameter(token_freqs)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # x[0].shape[1] contains motion frames count, but not used directly here
+        # x: [B, C, T, H, W] tensor
         device = self.patch_embedding.weight.device
         freqs = self.freqs
         if freqs.device != device:
@@ -977,18 +863,20 @@ class MotionerTransformers(nn.Module):
 
         if self.enable_tsm:
             sample_idx = [
-                _sample_indices(u.shape[1], stride=self.motion_stride, expand_ratio=self.expand_ratio, c=self.sample_c)
-                for u in x
+                _sample_indices(
+                    x[i].shape[1], stride=self.motion_stride, expand_ratio=self.expand_ratio, c=self.sample_c
+                )
+                for i in range(x.shape[0])
             ]
-            x = [torch.flip(torch.flip(u, [1])[:, idx], [1]) for idx, u in zip(sample_idx, x)]
+            x = torch.stack([torch.flip(torch.flip(x[i], [1])[:, idx], [1]) for i, idx in enumerate(sample_idx)])
 
-        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+        x = self.patch_embedding(x)
 
-        seq_f, seq_h, seq_w = x[0].shape[-3:]
-        batch_size = len(x)
+        batch_size = x.shape[0]
+        seq_f, seq_h, seq_w = x.shape[-3:]
         if not self.enable_tsm:
-            grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
-            grid_sizes = [[torch.zeros_like(grid_sizes), grid_sizes, grid_sizes]]
+            grid_sizes_t = torch.tensor([list(x.shape[2:])], dtype=torch.long).repeat(batch_size, 1)
+            grid_sizes = [[torch.zeros_like(grid_sizes_t), grid_sizes_t, grid_sizes_t]]
             seq_f = 0
         else:
             grid_sizes = []
@@ -1003,10 +891,9 @@ class MotionerTransformers(nn.Module):
                 grid_sizes += tsm_frame_grid_sizes
             seq_f = sample_idx[0][-1] + 1
 
-        x = [u.flatten(2).transpose(1, 2) for u in x]
-        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
-        x = torch.cat([u for u in x])
-        batch_size = len(x)
+        x = x.flatten(2).transpose(1, 2)
+        seq_lens = torch.full((batch_size,), x.shape[1], dtype=torch.long)
+        batch_size = x.shape[0]
 
         token_grid_sizes = [
             [
@@ -1022,11 +909,15 @@ class MotionerTransformers(nn.Module):
 
         grid_sizes = grid_sizes + token_grid_sizes
         token_len = self.token.shape[1]
-        token = self.token.clone().repeat(x.shape[0], 1, 1).contiguous()
-        seq_lens = seq_lens + torch.tensor([t.size(0) for t in token], dtype=torch.long)
+        token = self.token.clone().repeat(batch_size, 1, 1).contiguous()
+        seq_lens = seq_lens + token_len
         x = torch.cat([x, token], dim=1)
 
-        kwargs = dict(seq_lens=seq_lens, grid_sizes=grid_sizes, freqs=freqs)
+        precomputed_freqs = RotaryEmbeddingS2VGrid.precompute(
+            x.shape[1], self.num_heads, self.dim // self.num_heads, grid_sizes, freqs, x.device
+        )
+
+        kwargs = dict(seq_lens=seq_lens, grid_sizes=grid_sizes, freqs=freqs, precomputed_freqs=precomputed_freqs)
         for block in self.blocks:
             x = block(x, **kwargs)
 
@@ -1062,11 +953,8 @@ class FramePackMotioner(nn.Module):
             raise ValueError(f"inner_dim ({inner_dim}) must be divisible by num_heads ({num_heads})")
         if (inner_dim // num_heads) % 2 != 0:
             raise ValueError(f"inner_dim // num_heads ({inner_dim // num_heads}) must be even")
-        d = inner_dim // num_heads
-        self.freqs = torch.cat(
-            [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))],
-            dim=1,
-        )
+        self.rope = WanS2VRotaryPosEmbed(num_heads, inner_dim // num_heads)
+        self.freqs = self.rope.freqs
         self.drop_mode = drop_mode
 
     def forward(self, motion_latents: torch.Tensor, add_last_motion: int = 2) -> torch.Tensor:
@@ -1147,12 +1035,7 @@ class FramePackMotioner(nn.Module):
 
             grid_sizes = grid_sizes + grid_sizes_2x + grid_sizes_4x
 
-            motion_rope_emb = rope_precompute(
-                motion_lat.detach().view(1, motion_lat.shape[1], self.num_heads, self.inner_dim // self.num_heads),
-                grid_sizes,
-                self.freqs,
-                start=None,
-            )
+            motion_rope_emb = self.rope(motion_lat.detach(), grid_sizes)
 
             mot.append(motion_lat)
             mot_remb.append(motion_rope_emb)
@@ -1264,9 +1147,7 @@ class WanS2VTransformer3DModel(nn.Module):
 
         # Embeddings
         self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size)
-        self.text_embedding = nn.Sequential(nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim))
-        self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
-        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
+        self.condition_embedder = WanS2VConditionEmbedder(dim, freq_dim, text_dim, text_len)
 
         # Transformer blocks (TP-enabled)
         self.blocks = nn.ModuleList(
@@ -1277,18 +1158,18 @@ class WanS2VTransformer3DModel(nn.Module):
         )
 
         # Output head
-        self.head = WanS2VHead(dim, out_dim, patch_size, eps)
+        self.norm_out = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.proj_out = nn.Linear(dim, math.prod(patch_size) * out_dim)
+        self.output_scale_shift_prepare = WanS2VOutputScaleShiftPrepare(dim)
+        self.timestep_proj_prepare = WanS2VTimestepProjPrepare(dim)
 
         # RoPE base frequencies (complex-valued)
         if dim % num_heads != 0:
             raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})")
         if (dim // num_heads) % 2 != 0:
             raise ValueError(f"dim // num_heads ({dim // num_heads}) must be even")
-        d = dim // num_heads
-        self.freqs = torch.cat(
-            [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))],
-            dim=1,
-        )
+        self.rope = WanS2VRotaryPosEmbed(num_heads, dim // num_heads)
+        self.freqs = self.rope.freqs
 
         self.use_context_parallel = False
 
@@ -1323,6 +1204,8 @@ class WanS2VTransformer3DModel(nn.Module):
 
         self.zero_timestep = zero_timestep
         self.add_last_motion = add_last_motion
+        self._cached_rope_key = None
+        self._cached_rope_emb = None
 
         # Motion modules
         if enable_motioner and enable_framepack:
@@ -1369,7 +1252,7 @@ class WanS2VTransformer3DModel(nn.Module):
                         .repeat(1, 1),
                     ]
                 ]
-                token_freqs = rope_apply(x, gride_sizes, self.freqs)
+                token_freqs = RotaryEmbeddingS2VGrid.forward(x, gride_sizes, self.rope.freqs)
                 token_freqs = token_freqs[0, :, 0].reshape(motion_token_num, -1, 2)
                 token_freqs = token_freqs * 0.01
                 self.token_freqs = nn.Parameter(token_freqs)
@@ -1386,12 +1269,26 @@ class WanS2VTransformer3DModel(nn.Module):
             if hasattr(self, "cond_encoder"):
                 self.cond_encoder = zero_module(self.cond_encoder)
             for i in range(len(self.audio_injector.injector)):
-                self.audio_injector.injector[i].o = zero_module(self.audio_injector.injector[i].o)
+                self.audio_injector.injector[i].to_out = zero_module(self.audio_injector.injector[i].to_out)
                 if self.enable_adain:
                     self.audio_injector.injector_adain_layers[i].linear = zero_module(
                         self.audio_injector.injector_adain_layers[i].linear
                     )
 
+    # ------------------------------------------------------------------
+    # RoPE preparation
+    # ------------------------------------------------------------------
+
+    def prepare_rope(self, x, grid_sizes):
+        """Precompute RoPE for the concatenated hidden states (cached across steps)."""
+        cache_key = (x.shape[0], x.shape[1], x.device)
+        if self._cached_rope_key == cache_key and self._cached_rope_emb is not None:
+            return self._cached_rope_emb
+        result = self.rope(x.detach(), grid_sizes)
+        self._cached_rope_key = cache_key
+        self._cached_rope_emb = result
+        return result
+
     # ------------------------------------------------------------------
     # Motion processing
     # ------------------------------------------------------------------
@@ -1414,12 +1311,7 @@ class WanS2VTransformer3DModel(nn.Module):
                     torch.tensor([self.lat_motion_frames, height, width]).unsqueeze(0).repeat(1, 1),
                 ]
             ]
-            motion_rope_emb = rope_precompute(
-                flat_mot.detach().view(1, flat_mot.shape[1], self.num_heads, self.dim // self.num_heads),
-                motion_grid_sizes,
-                self.freqs,
-                start=None,
-            )
+            motion_rope_emb = self.rope(flat_mot.detach(), motion_grid_sizes)
             mot_remb.append(motion_rope_emb)
             flatten_mot.append(flat_mot)
         return flatten_mot, mot_remb
@@ -1432,25 +1324,21 @@ class WanS2VTransformer3DModel(nn.Module):
             return flatten_mot, mot_remb
 
     def process_motion_transformer_motioner(self, motion_latents, drop_motion_frames=False, add_last_motion=True):
-        batch_size = len(motion_latents)
-        height = motion_latents[0].shape[2] // self.patch_size[1]
-        width = motion_latents[0].shape[3] // self.patch_size[2]
+        batch_size = motion_latents.shape[0]
+        height = motion_latents.shape[3] // self.patch_size[1]
+        width = motion_latents.shape[4] // self.patch_size[2]
 
-        freqs = self.freqs
-        device = self.patch_embedding.weight.device
-        if freqs.device != device:
-            freqs = freqs.to(device)
+        trainable_f = None
         if self.trainable_token_pos_emb:
             with torch.amp.autocast(current_omni_platform.device_type, dtype=torch.float64):
                 token_freqs = self.token_freqs.to(torch.float64)
                 token_freqs = token_freqs / token_freqs.norm(dim=-1, keepdim=True)
-                freqs = [freqs, torch.view_as_complex(token_freqs)]
+                trainable_f = torch.view_as_complex(token_freqs)
 
         if not drop_motion_frames and add_last_motion:
-            last_motion_latent = [u[:, -1:] for u in motion_latents]
-            last_mot = [self.patch_embedding(m.unsqueeze(0)) for m in last_motion_latent]
-            last_mot = [m.flatten(2).transpose(1, 2) for m in last_mot]
-            last_mot = torch.cat(last_mot)
+            last_motion_latent = motion_latents[:, :, -1:]
+            last_mot = self.patch_embedding(last_motion_latent)
+            last_mot = last_mot.flatten(2).transpose(1, 2)
             gride_sizes = [
                 [
                     torch.tensor([-1, 0, 0]).unsqueeze(0).repeat(batch_size, 1),
@@ -1459,9 +1347,7 @@ class WanS2VTransformer3DModel(nn.Module):
                 ]
             ]
         else:
-            last_mot = torch.zeros(
-                [batch_size, 0, self.dim], device=motion_latents[0].device, dtype=motion_latents[0].dtype
-            )
+            last_mot = torch.zeros([batch_size, 0, self.dim], device=motion_latents.device, dtype=motion_latents.dtype)
             gride_sizes = []
 
         zip_motion = self.motioner(motion_latents)
@@ -1483,12 +1369,7 @@ class WanS2VTransformer3DModel(nn.Module):
         mot = torch.cat([last_mot, zip_motion], dim=1)
         gride_sizes = gride_sizes + zip_motion_grid_sizes
 
-        motion_rope_emb = rope_precompute(
-            mot.detach().view(batch_size, mot.shape[1], self.num_heads, self.dim // self.num_heads),
-            gride_sizes,
-            freqs,
-            start=None,
-        )
+        motion_rope_emb = self.rope(mot.detach(), gride_sizes, trainable_freqs=trainable_f)
         return [m.unsqueeze(0) for m in mot], [r.unsqueeze(0) for r in motion_rope_emb]
 
     def inject_motion(
@@ -1521,12 +1402,11 @@ class WanS2VTransformer3DModel(nn.Module):
             audio_emb = self.merged_audio_emb
             num_frames = audio_emb.shape[1]
 
-            input_hidden_states = hidden_states[:, : self.original_seq_len].clone()
+            input_hidden_states = hidden_states[:, : self.original_seq_len]
             input_hidden_states = rearrange(input_hidden_states, "b (t n) c -> (b t) n c", t=num_frames)
 
             if self.enable_adain and self.adain_mode == "attn_norm":
-                audio_emb_global = self.audio_emb_global
-                audio_emb_global = rearrange(audio_emb_global, "b t n c -> (b t) n c")
+                audio_emb_global = rearrange(self.audio_emb_global, "b t n c -> (b t) n c")
                 adain_hidden_states = self.audio_injector.injector_adain_layers[audio_attn_id](
                     input_hidden_states, temb=audio_emb_global[:, 0]
                 )
@@ -1534,8 +1414,7 @@ class WanS2VTransformer3DModel(nn.Module):
             else:
                 attn_hidden_states = self.audio_injector.injector_pre_norm_feat[audio_attn_id](input_hidden_states)
 
-            audio_emb = rearrange(audio_emb, "b t n c -> (b t) n c", t=num_frames)
-            attn_audio_emb = audio_emb
+            attn_audio_emb = rearrange(audio_emb, "b t n c -> (b t) n c")
             residual_out = self.audio_injector.injector[audio_attn_id](
                 x=attn_hidden_states,
                 context=attn_audio_emb,
@@ -1543,7 +1422,7 @@ class WanS2VTransformer3DModel(nn.Module):
                 * attn_audio_emb.shape[1],
             )
             residual_out = rearrange(residual_out, "(b t) n c -> b (t n) c", t=num_frames)
-            hidden_states[:, : self.original_seq_len] = hidden_states[:, : self.original_seq_len] + residual_out
+            hidden_states[:, : self.original_seq_len].add_(residual_out)
 
         return hidden_states
 
@@ -1570,7 +1449,7 @@ class WanS2VTransformer3DModel(nn.Module):
         result = {}
         if self.enable_adain:
             audio_emb_global, audio_emb = audio_emb_res
-            result["audio_emb_global"] = audio_emb_global[:, motion_frames[1] :].clone()
+            result["audio_emb_global"] = audio_emb_global[:, motion_frames[1] :]
         else:
             audio_emb = audio_emb_res
         result["audio_emb"] = audio_emb[:, motion_frames[1] :, :]
@@ -1578,60 +1457,51 @@ class WanS2VTransformer3DModel(nn.Module):
 
     def forward(
         self,
-        x: torch.Tensor,
-        t: torch.Tensor,
-        context: torch.Tensor,
-        seq_len: int,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
         ref_latents: torch.Tensor,
         motion_latents: torch.Tensor,
         cond_states: torch.Tensor,
-        audio_input: torch.Tensor | None = None,
-        motion_frames: list[int] = [17, 5],
+        encoder_hidden_states_audio: dict[str, torch.Tensor] | None = None,
         add_last_motion: int = 2,
         drop_motion_frames: bool = False,
-        audio_emb: dict[str, torch.Tensor] | None = None,
-        *extra_args: object,
+        return_dict: bool = True,
         **extra_kwargs: object,
     ) -> torch.Tensor:
         """Forward pass for S2V transformer.
 
         Args:
-            x: Noisy latents [B, C, T, H, W]
-            t: Timesteps [B]
-            context: Text embeddings [B, seq_len, dim]
-            seq_len: Sequence length
+            hidden_states: Noisy latents [B, C, T, H, W]
+            timestep: Timesteps [B]
+            encoder_hidden_states: Text embeddings [B, seq_len, dim]
             ref_latents: Reference image latents [B, C, 1, H, W]
             motion_latents: Motion context latents [B, C, T_motion, H, W]
-            cond_states: Pose condition latents [B, C, T, H, W] — typically zero tensor
-                        for no pose control. Iterates over batch dimension (dim 0).
-            audio_input: Raw audio features (optional if audio_emb provided)
-            motion_frames: [num_frames, num_latent_frames] for audio alignment
+            cond_states: Pose condition latents [B, C, T, H, W]
+            encoder_hidden_states_audio: Precomputed audio embeddings from encode_audio()
             add_last_motion: Number of last motion frames to repeat
             drop_motion_frames: Whether to drop first motion frame
-            audio_emb: Precomputed audio embeddings from encode_audio()
+            return_dict: Whether to return Transformer2DModelOutput
 
         Returns:
             Predicted noise [B, C, T, H, W]
         """
+        # Clear cached state from previous forward (prevents stale data from
+        # dummy warmup or prior clips polluting this inference)
+        self._cached_rope_key = None
+        self._cached_rope_emb = None
+
+        x = hidden_states
+        t = timestep
+        context = encoder_hidden_states
+        audio_emb = encoder_hidden_states_audio
         add_last_motion = self.add_last_motion * add_last_motion
 
-        # Audio encoding — use precomputed embeddings if available
-        if audio_emb is not None:
-            self.merged_audio_emb = audio_emb["audio_emb"]
-            if self.enable_adain:
-                self.audio_emb_global = audio_emb["audio_emb_global"]
-        else:
-            audio_input = torch.cat(
-                [audio_input[..., 0:1].repeat(1, 1, 1, motion_frames[0]), audio_input],
-                dim=-1,
-            )
-            audio_emb_res = self.casual_audio_encoder(audio_input)
-            if self.enable_adain:
-                audio_emb_global, audio_emb_local = audio_emb_res
-                self.audio_emb_global = audio_emb_global[:, motion_frames[1] :].clone()
-            else:
-                audio_emb_local = audio_emb_res
-            self.merged_audio_emb = audio_emb_local[:, motion_frames[1] :, :]
+        # Audio embeddings — must be precomputed by pipeline via encode_audio()
+        assert audio_emb is not None, "S2V requires precomputed audio_emb (encoder_hidden_states_audio)"
+        self.merged_audio_emb = audio_emb["audio_emb"]
+        if self.enable_adain:
+            self.audio_emb_global = audio_emb["audio_emb_global"]
 
         # Patch embedding + condition
         # NOTE: x and cond_states are 5D tensors [B, C, T, H, W]. Iterating over them
@@ -1674,8 +1544,7 @@ class WanS2VTransformer3DModel(nn.Module):
 
         # Precompute RoPE
         x = torch.cat(x)
-        b, s, n, d = x.size(0), x.size(1), self.num_heads, self.dim // self.num_heads
-        self.pre_compute_freqs = rope_precompute(x.detach().view(b, s, n, d), grid_sizes, self.freqs, start=None)
+        self.pre_compute_freqs = self.prepare_rope(x, grid_sizes)
 
         x = [u.unsqueeze(0) for u in x]
         self.pre_compute_freqs = [u.unsqueeze(0) for u in self.pre_compute_freqs]
@@ -1697,29 +1566,13 @@ class WanS2VTransformer3DModel(nn.Module):
 
         x = x + self.trainable_cond_mask(mask_input).to(x.dtype)
 
-        # Time embeddings — sinusoidal_embedding_1d uses float64 internally
-        # for precision, then we cast to the model dtype (bf16) since the
-        # downstream scale/shift modulation doesn't need float32.
+        # Condition embeddings (time + text)
         if self.zero_timestep:
             t = torch.cat([t, torch.zeros([1], dtype=t.dtype, device=t.device)])
-        e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t).to(x.dtype))
-        e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+        e, timestep_proj, context = self.condition_embedder(t, context, x.dtype)
+        e, e0 = self.timestep_proj_prepare(timestep_proj, e, self.zero_timestep, self.original_seq_len)
 
-        if self.zero_timestep:
-            e = e[:-1]
-            zero_e0 = e0[-1:]
-            e0 = e0[:-1]
-            e0 = torch.cat([e0.unsqueeze(2), zero_e0.unsqueeze(2).repeat(e0.size(0), 1, 1, 1)], dim=2)
-            e0 = [e0, self.original_seq_len]
-        else:
-            e0 = e0.unsqueeze(2).repeat(1, 1, 2, 1)
-            e0 = [e0, 0]
-
-        # Context embedding
         context_lens = None
-        context = self.text_embedding(
-            torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context])
-        )
 
         # Transformer blocks
         kwargs = dict(
@@ -1735,11 +1588,19 @@ class WanS2VTransformer3DModel(nn.Module):
             x = block(x, **kwargs)
             x = self.after_transformer_block(idx, x)
 
-        # Output
+        # Output norm, projection & unpatchify
         x = x[:, : self.original_seq_len]
-        x = self.head(x, e)
+        shift, scale = self.output_scale_shift_prepare(e)
+        x = self.norm_out(x).type_as(x) * (1 + scale) + shift
+        x = self.proj_out(x)
         x = self.unpatchify(x, original_grid_sizes)
-        return [u.float() for u in x]
+        output = torch.stack([u.float() for u in x])
+
+        if not return_dict:
+            return (output,)
+        from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+        return Transformer2DModelOutput(sample=output)
 
     def unpatchify(self, x, grid_sizes):
         c = self.out_dim
@@ -1802,26 +1663,44 @@ class WanS2VTransformer3DModel(nn.Module):
                 name = name.replace(".ffn.0.", ".ffn.net_0.proj.")
                 name = name.replace(".ffn.2.", ".ffn.net_2.")
 
+            # Audio injector weight remapping (q/k/v/o → to_q/to_k/to_v/to_out)
+            if "audio_injector.injector." in name:
+                name = name.replace(".q.", ".to_q.")
+                name = name.replace(".k.", ".to_k.")
+                name = name.replace(".v.", ".to_v.")
+                name = name.replace(".o.", ".to_out.")
+
+            # Output head remapping (head.* → norm_out/proj_out/output_scale_shift_prepare)
+            if name.startswith("head.head."):
+                name = "proj_out." + name[len("head.head.") :]
+            elif name == "head.modulation":
+                name = "output_scale_shift_prepare.scale_shift_table"
+
+            # Condition embedder remapping (only top-level)
+            # time_embedding.0.* → condition_embedder.time_embedder.linear_1.*
+            # time_embedding.2.* → condition_embedder.time_embedder.linear_2.*
+            if name.startswith("time_embedding.0."):
+                name = "condition_embedder.time_embedder.linear_1." + name[len("time_embedding.0.") :]
+            elif name.startswith("time_embedding.2."):
+                name = "condition_embedder.time_embedder.linear_2." + name[len("time_embedding.2.") :]
+            # time_projection.1.* → condition_embedder.time_proj.*
+            elif name.startswith("time_projection.1."):
+                name = "condition_embedder.time_proj." + name[len("time_projection.1.") :]
+            # text_embedding.0.* → condition_embedder.text_embedder.linear_1.*
+            # text_embedding.2.* → condition_embedder.text_embedder.linear_2.*
+            elif name.startswith("text_embedding.0."):
+                name = "condition_embedder.text_embedder.linear_1." + name[len("text_embedding.0.") :]
+            elif name.startswith("text_embedding.2."):
+                name = "condition_embedder.text_embedder.linear_2." + name[len("text_embedding.2.") :]
+
             if name not in params_dict:
                 logger.warning(f"Skipping weight {original_name} -> {name}")
                 continue
 
             param = params_dict[name]
 
-            # TP-shard norm weights for main blocks
-            if (
-                tp_size > 1
-                and is_main_block
-                and any(
-                    norm_name in name
-                    for norm_name in [
-                        ".self_attn.norm_q.",
-                        ".self_attn.norm_k.",
-                        ".cross_attn.norm_q.",
-                        ".cross_attn.norm_k.",
-                    ]
-                )
-            ):
+            # TP-shard norm weights for TP-aware attention modules
+            if tp_size > 1 and (".norm_q." in name or ".norm_k." in name) and loaded_weight.shape[0] != param.shape[0]:
                 shard_size = loaded_weight.shape[0] // tp_size
                 loaded_weight = loaded_weight[tp_rank * shard_size : (tp_rank + 1) * shard_size]