Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,71 @@
+# Matrix Game 2.0 Modular Pipeline
+
+## Set Up
+
+```shell
+uv venv -p 3.10
+uv pip install -r requirements.txt
+uv pip install git+https://github.com/huggingface/diffusers.git
+```
+
+## How to Use
+
+```python
+import torch
+from diffusers import ModularPipelineBlocks
+from diffusers.utils import export_to_video, load_image
+from diffusers.modular_pipelines import WanModularPipeline
+
+class MatrixGameWanModularPipeline(WanModularPipeline):
+    """
+    A ModularPipeline for MatrixGameWan.
+
+    <Tip warning={true}>
+
+        This is an experimental feature and is likely to change in the future.
+
+    </Tip>
+    """
+
+    @property
+    def default_sample_height(self):
+        return 44
+
+    @property
+    def default_sample_width(self):
+        return 80
+
+# Download custom blocks for the pipeline
+blocks = ModularPipelineBlocks.from_pretrained(
+    "diffusers-internal-dev/matrix-game-2-modular",
+    trust_remote_code=True,
+)
+
+# Initialize the pipeline runtime using the config in the repo
+pipe = MatrixGameWanModularPipeline(blocks, "diffusers-internal-dev/matrix-game-2-modular")
+
+# Load the model components of the pipeline
+pipe.load_components(
+    trust_remote_code=True,
+    device_map="cuda",
+    torch_dtype={"default": torch.bfloat16, "vae": torch.float32}
+)
+
+image = load_image("https://github.com/SkyworkAI/Matrix-Game/blob/main/Matrix-Game-2/demo_images/universal/0016.png?raw=true")
+output = pipe(image=image, num_frames=141)
+export_to_video(output.values['videos'][0], "matrix-game.mp4")
+```
+
+## Providing Actions as Inputs
+
+Each action is represented as a string. The available actions are:
+
+Motion Actions: ["forward", "left", "right"]
+Camera Actions: ["camera_l", "camera_r", "camera_u", "camera_d"]
+Compound Actions: Combinations of motion and camera actions with an `_` separating actions, e.g. "forward_left", "forward_left_camera_l"
+
+```py
+image = load_image("https://github.com/SkyworkAI/Matrix-Game/blob/main/Matrix-Game-2/demo_images/universal/0016.png?raw=true")
+output = pipe(image=image, actions=["forward", "camera_l"], num_frames=141)
+export_to_video(output.values['videos'][0], "matrix-game.mp4")
+```

before_denoise.py

@@ -0,0 +1,604 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import List, Optional, Union, Dict
+
+import torch
+
+from diffusers import AutoencoderKLWan
+from diffusers.configuration_utils import FrozenDict
+from diffusers.schedulers import UniPCMultistepScheduler
+from diffusers.utils import logging
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+from diffusers.modular_pipelines import ModularPipeline, ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+# Constants
+FRAME_MULTIPLE = 4
+DEFAULT_SAMPLES_PER_ACTION = 4
+DEFAULT_FRAMES_PER_ACTION = 12
+
+DEFAULT_MOUSE_DIM = 2
+DEFAULT_KEYBOARD_DIM = 4
+
+# Camera movement configuration
+CAMERA_MOVEMENT_VALUE = 0.1
+CAMERA_VALUE_MAP = {
+    "camera_up": [CAMERA_MOVEMENT_VALUE, 0],
+    "camera_down": [-CAMERA_MOVEMENT_VALUE, 0],
+    "camera_l": [0, -CAMERA_MOVEMENT_VALUE],
+    "camera_r": [0, CAMERA_MOVEMENT_VALUE],
+    "camera_ur": [CAMERA_MOVEMENT_VALUE, CAMERA_MOVEMENT_VALUE],
+    "camera_ul": [CAMERA_MOVEMENT_VALUE, -CAMERA_MOVEMENT_VALUE],
+    "camera_dr": [-CAMERA_MOVEMENT_VALUE, CAMERA_MOVEMENT_VALUE],
+    "camera_dl": [-CAMERA_MOVEMENT_VALUE, -CAMERA_MOVEMENT_VALUE],
+}
+
+# Define available actions
+MOVEMENT_ACTIONS = ["forward", "left", "right"]
+COMPOUND_MOVEMENTS = ["forward_left", "forward_right"]
+CAMERA_ACTIONS = list(CAMERA_VALUE_MAP.keys())
+
+# Keyboard action indices
+KEYBOARD_ACTION_INDICES = {"forward": 0, "back": 1, "left": 2, "right": 3}
+
+
+def sync_actions_to_frames(
+    actions: List[str],
+    num_frames: int,
+    min_duration: int = 12
+) -> List[Dict[str, Union[str, int]]]:
+    """
+    Synchronize a list of actions to fit exactly within the given number of frames
+    using equal distribution strategy.
+
+    Args:
+        actions: List of action names to perform
+        num_frames: Total frames to fill
+        min_duration: Minimum frames per action (should be multiple of frame_multiple)
+        frame_multiple: Actions must be multiples of this value
+
+    Returns:
+        List of action dictionaries with 'type', 'start_frame', and 'duration'
+    """
+
+    if not actions:
+        raise ValueError("No actions provided")
+
+    max_possible_actions = num_frames // DEFAULT_FRAMES_PER_ACTION
+    if len(actions) > max_possible_actions:
+        actions = actions[:max_possible_actions]
+
+    num_actions = len(actions)
+
+    frames_per_action = num_frames // num_actions
+    frames_per_action = (frames_per_action // FRAME_MULTIPLE) * FRAME_MULTIPLE
+    frames_per_action = max(DEFAULT_FRAMES_PER_ACTION, frames_per_action)
+
+    remaining_frames = num_frames - (frames_per_action * num_actions)
+    output = []
+    current_frame = 0
+
+    for i, action in enumerate(actions):
+        duration = frames_per_action if i != num_actions - 1 else num_frames - current_frame
+
+        output.append({
+            "action_type": action,
+            "start_frame": current_frame,
+            "duration": duration
+        })
+
+        current_frame += duration
+
+    return output
+
+
+def actions_to_condition_tensors(actions, num_frames):
+    keyboard_conditions = torch.zeros((num_frames, DEFAULT_KEYBOARD_DIM))
+    mouse_conditions = torch.zeros((num_frames, DEFAULT_MOUSE_DIM))
+
+    for action in actions:
+        action_type = action['action_type']
+        start_frame = action['start_frame']
+        end_frame = start_frame + action['duration']
+
+        action_components = action_type.split("_")
+        for component in action_components:
+            if component in KEYBOARD_ACTION_INDICES:
+                action_idx = KEYBOARD_ACTION_INDICES[component]
+                keyboard_conditions[start_frame:end_frame, action_idx] = 1.0
+
+        if not "camera" in action_type:
+            continue
+
+        mouse_x = mouse_y = 0.0
+        for idx, component in enumerate(action_components):
+            if not action_components[idx] == "camera":
+                continue
+
+            camera_action = f"camera_{action_components[idx+1]}"
+            if camera_action not in CAMERA_VALUE_MAP:
+                continue
+
+            camera_values = CAMERA_VALUE_MAP[camera_action]
+            mouse_x += camera_values[0]
+            mouse_y += camera_values[1]
+
+        mouse_conditions[start_frame:end_frame, 0] = mouse_x
+        mouse_conditions[start_frame:end_frame, 1] = mouse_y
+
+    return keyboard_conditions, mouse_conditions
+
+
+def _build_test_actions(
+    movement_actions: List[str],
+    compound_movements: List[str],
+    camera_actions: List[str],
+) -> List[str]:
+    """Build comprehensive list of test action combinations.
+
+    Args:
+        movement_actions: List of basic movement actions
+        compound_movements: List of compound movement actions
+        camera_actions: List of camera control actions
+
+    Returns:
+        List of all action combinations to test
+    """
+    # Create base test actions with repetition for variety
+    test_actions = compound_movements * 5 + camera_actions * 5 + movement_actions * 5
+
+    # Add combined movement + camera actions
+    for movement in movement_actions + compound_movements:
+        for camera in camera_actions:
+            combined_action = f"{movement}_{camera}"
+            test_actions.append(combined_action)
+
+    return test_actions
+
+
+def generate_random_condition_tensors(num_frames: int) -> Dict[str, torch.Tensor]:
+    """Generate benchmark action sequences for testing.
+
+    Args:
+        num_frames: Total number of frames to generate
+        num_samples_per_action: Number of samples per action type
+
+    Returns:
+        Dictionary containing keyboard and mouse conditions for benchmark actions
+    """
+    # Build test action combinations
+    actions = _build_test_actions(
+        MOVEMENT_ACTIONS, COMPOUND_MOVEMENTS, CAMERA_ACTIONS
+    )
+    actions = sync_actions_to_frames(actions, num_frames)
+    return actions_to_condition_tensors(actions, num_frames)
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+def retrieve_latents(
+    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+class MatrixGameWanActionInputStep(ModularPipelineBlocks):
+    model_name = "MatrixGameWan"
+
+    @property
+    def description(self) -> str:
+        return "Action Input step"
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return []
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [InputParam("num_frames", type_hint=int, required=True), InputParam("actions", type_hint=List[str])]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam(
+                "keyboard_conditions",
+                type_hint=torch.Tensor,
+                description="image embeddings used to guide the image generation",
+            ),
+            OutputParam(
+                "mouse_conditions",
+                type_hint=torch.Tensor,
+                description="image embeddings used to guide the image generation",
+            )
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
+        # Get inputs and intermediates
+        block_state = self.get_block_state(state)
+        block_state.device = components._execution_device
+        actions = block_state.actions
+
+        if actions is not None:
+            actions = sync_actions_to_frames(actions, block_state.num_frames)
+            keyboard_conditions, mouse_conditions = actions_to_condition_tensors(actions, block_state.num_frames)
+        else:
+            keyboard_conditions, mouse_conditions = generate_random_condition_tensors(block_state.num_frames)
+
+        block_state.keyboard_conditions = keyboard_conditions.to(block_state.device)
+        block_state.mouse_conditions = mouse_conditions.to(block_state.device)
+
+        # Add outputs
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MatrixGameWanSetTimestepsStep(ModularPipelineBlocks):
+    model_name = "MatrixGameWan"
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [
+            ComponentSpec("scheduler", UniPCMultistepScheduler),
+        ]
+
+    @property
+    def description(self) -> str:
+        return "Step that sets the scheduler's timesteps for inference"
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam("num_inference_steps", default=4),
+            InputParam("timesteps"),
+            InputParam("sigmas"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"),
+            OutputParam(
+                "num_inference_steps",
+                type_hint=int,
+                description="The number of denoising steps to perform at inference time",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        block_state.device = components._execution_device
+
+        block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+            components.scheduler,
+            block_state.num_inference_steps,
+            block_state.device,
+            block_state.timesteps,
+            block_state.sigmas,
+        )
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MatrixGameWanPrepareLatentsStep(ModularPipelineBlocks):
+    model_name = "MatrixGameWan"
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [ComponentSpec("vae", AutoencoderKLWan),]
+
+    @property
+    def description(self) -> str:
+        return "Prepare latents step that prepares the latents for the text-to-video generation process"
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam("height", type_hint=int),
+            InputParam("width", type_hint=int),
+            InputParam("num_frames", type_hint=int),
+            InputParam("latents", type_hint=Optional[torch.Tensor]),
+            InputParam("num_videos_per_prompt", type_hint=int, default=1),
+            InputParam("generator"),
+            InputParam("dtype", type_hint=torch.dtype, description="The dtype of the model inputs"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam(
+                "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
+            )
+        ]
+
+    @staticmethod
+    def check_inputs(components, block_state):
+        if (block_state.height is not None and block_state.height % components.vae_scale_factor_spatial != 0) or (
+            block_state.width is not None and block_state.width % components.vae_scale_factor_spatial != 0
+        ):
+            raise ValueError(
+                f"`height` and `width` have to be divisible by {components.vae_scale_factor_spatial} but are {block_state.height} and {block_state.width}."
+            )
+        if block_state.num_frames is not None and (
+            block_state.num_frames < 1 or (block_state.num_frames - 1) % components.vae_scale_factor_temporal != 0
+        ):
+            raise ValueError(
+                f"`num_frames` has to be greater than 0, and (num_frames - 1) must be divisible by {components.vae_scale_factor_temporal}, but got {block_state.num_frames}."
+            )
+
+    @staticmethod
+    def prepare_latents(
+        components,
+        batch_size: int,
+        num_channels_latents: int = 16,
+        height: int = 352,
+        width: int = 640,
+        num_frames: int = 81,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        num_latent_frames = (num_frames - 1) // components.vae_scale_factor_temporal + 1
+        shape = (
+            batch_size,
+            num_channels_latents,
+            num_latent_frames,
+            int(height) // components.vae_scale_factor_spatial,
+            int(width) // components.vae_scale_factor_spatial,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        return latents
+
+    @torch.no_grad()
+    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        block_state.height = block_state.height or components.default_height
+        block_state.width = block_state.width or components.default_width
+        block_state.num_frames = block_state.num_frames or components.default_num_frames
+        block_state.device = components._execution_device
+        block_state.dtype = torch.float32  # Wan latents should be torch.float32 for best quality
+        block_state.num_channels_latents = components.num_channels_latents
+
+        self.check_inputs(components, block_state)
+
+        block_state.latents = self.prepare_latents(
+            components,
+            1,
+            block_state.num_channels_latents,
+            block_state.height,
+            block_state.width,
+            block_state.num_frames,
+            block_state.dtype,
+            block_state.device,
+            block_state.generator,
+            block_state.latents,
+        )
+
+        self.set_block_state(state, block_state)
+
+        return components, state
+
+
+class MatrixGameWanPrepareImageMaskLatentsStep(ModularPipelineBlocks):
+    model_name = "MatrixGameWan"
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [
+            ComponentSpec("vae", AutoencoderKLWan),
+            ComponentSpec("video_processor", VideoProcessor, config=FrozenDict({"vae_scale_factor": 8}))
+        ]
+
+    @property
+    def description(self) -> str:
+        return "Prepare latents step that prepares the latents for the text-to-video generation process"
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam("image"),
+            InputParam("height", type_hint=int),
+            InputParam("width", type_hint=int),
+            InputParam("num_frames", type_hint=int),
+            InputParam("image_mask_latents", type_hint=Optional[torch.Tensor]),
+            InputParam("num_videos_per_prompt", type_hint=int, default=1),
+            InputParam("generator"),
+            InputParam("dtype", type_hint=torch.dtype, description="The dtype of the model inputs"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam(
+                "image_mask_latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
+            )
+        ]
+
+    @staticmethod
+    def check_inputs(components, block_state):
+        if (block_state.height is not None and block_state.height % components.vae_scale_factor_spatial != 0) or (
+            block_state.width is not None and block_state.width % components.vae_scale_factor_spatial != 0
+        ):
+            raise ValueError(
+                f"`height` and `width` have to be divisible by {components.vae_scale_factor_spatial} but are {block_state.height} and {block_state.width}."
+            )
+
+    @staticmethod
+    @torch.no_grad()
+    def prepare_latents(
+        components,
+        image,
+        batch_size: int,
+        num_channels_latents: int = 16,
+        height: int = 352,
+        width: int = 640,
+        num_frames: int = 81,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        image = components.video_processor.preprocess(image, height, width).to(device, torch.float32)
+        image = image.unsqueeze(2)  # [batch_size, channels, 1, height, width]
+
+        video_condition = torch.cat(
+            [image, image.new_zeros(image.shape[0], image.shape[1], num_frames - 1, height, width)], dim=2
+        )
+        video_condition = video_condition.to(device=device, dtype=components.vae.dtype)
+        latent_condition = retrieve_latents(components.vae.encode(video_condition), sample_mode="argmax")
+        latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1)
+
+        latents_mean = (
+            torch.tensor(components.vae.config.latents_mean)
+            .view(1, components.vae.config.z_dim, 1, 1, 1)
+            .to(device, dtype)
+        )
+        latents_std = 1.0 / torch.tensor(components.vae.config.latents_std).view(1, components.vae.config.z_dim, 1, 1, 1).to(
+            device, dtype
+        )
+        latent_condition = latent_condition.to(dtype)
+        latent_condition = (latent_condition - latents_mean) * latents_std
+
+        latent_height = height // components.vae_scale_factor_spatial
+        latent_width = width // components.vae_scale_factor_spatial
+
+        mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, latent_width)
+        mask_lat_size[:, :, list(range(1, num_frames))] = 0
+
+        first_frame_mask = mask_lat_size[:, :, 0:1]
+        first_frame_mask = torch.repeat_interleave(first_frame_mask, dim=2, repeats=components.vae_scale_factor_temporal)
+
+        mask_lat_size = torch.concat([first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
+        mask_lat_size = mask_lat_size.view(batch_size, -1, components.vae_scale_factor_temporal, latent_height, latent_width)
+        mask_lat_size = mask_lat_size.transpose(1, 2).to(latent_condition.device)
+
+        image_mask_latents = torch.concat([mask_lat_size, latent_condition], dim=1)
+        return image_mask_latents
+
+    @torch.no_grad()
+    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        block_state.height = block_state.height or components.default_height
+        block_state.width = block_state.width or components.default_width
+        block_state.num_frames = block_state.num_frames or components.default_num_frames
+        block_state.device = components._execution_device
+        block_state.dtype = torch.float32  # Wan latents should be torch.float32 for best quality
+        block_state.num_channels_latents = components.num_channels_latents
+
+        self.check_inputs(components, block_state)
+        block_state.image_mask_latents = self.prepare_latents(
+            components,
+            block_state.image,
+            1,
+            block_state.num_channels_latents,
+            block_state.height,
+            block_state.width,
+            block_state.num_frames,
+            block_state.dtype,
+            block_state.device,
+            block_state.generator,
+            block_state.image_mask_latents,
+        )
+
+        self.set_block_state(state, block_state)
+
+        return components, state
+

block.py

@@ -0,0 +1,6 @@
+from diffusers.modular_pipelines import SequentialPipelineBlocks
+from .modular_blocks import ACTION2VIDEO_BLOCKS
+
+class MatrixGameWanBlocks(SequentialPipelineBlocks):
+    block_classes = list(ACTION2VIDEO_BLOCKS.copy().values())
+    block_names = list(ACTION2VIDEO_BLOCKS.copy().keys())

decoders.py

@@ -0,0 +1,100 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, List, Tuple, Union
+
+import numpy as np
+import PIL
+import torch
+
+from diffusers.configuration_utils import FrozenDict
+from diffusers.models import AutoencoderKLWan
+from diffusers.utils import logging
+from diffusers.video_processor import VideoProcessor
+from diffusers.modular_pipelines import ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class MatrixGameWanDecodeStep(ModularPipelineBlocks):
+    model_name = "MatrixGameWan"
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [
+            ComponentSpec("vae", AutoencoderKLWan, repo="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", subfolder="vae"),
+            ComponentSpec(
+                "video_processor",
+                VideoProcessor,
+                config=FrozenDict({"vae_scale_factor": 8}),
+                default_creation_method="from_config",
+            ),
+        ]
+
+    @property
+    def description(self) -> str:
+        return "Step that decodes the denoised latents into images"
+
+    @property
+    def inputs(self) -> List[Tuple[str, Any]]:
+        return [
+            InputParam("output_type", default="pil"),
+            InputParam(
+                "latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The denoised latents from the denoising step",
+            )
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[str]:
+        return [
+            OutputParam(
+                "videos",
+                type_hint=Union[List[List[PIL.Image.Image]], List[torch.Tensor], List[np.ndarray]],
+                description="The generated videos, can be a PIL.Image.Image, torch.Tensor or a numpy array",
+            )
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        vae_dtype = components.vae.dtype
+
+        if not block_state.output_type == "latent":
+            latents = block_state.latents
+            latents_mean = (
+                torch.tensor(components.vae.config.latents_mean)
+                .view(1, components.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = 1.0 / torch.tensor(components.vae.config.latents_std).view(
+                1, components.vae.config.z_dim, 1, 1, 1
+            ).to(latents.device, latents.dtype)
+            latents = latents / latents_std + latents_mean
+            latents = latents.to(vae_dtype)
+            block_state.videos = components.vae.decode(latents, return_dict=False)[0]
+        else:
+            block_state.videos = block_state.latents
+
+        block_state.videos = components.video_processor.postprocess_video(
+            block_state.videos, output_type=block_state.output_type
+        )
+
+        self.set_block_state(state, block_state)
+
+        return components, state

denoise.py

@@ -0,0 +1,420 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, List, Tuple
+
+import torch
+
+from diffusers.configuration_utils import FrozenDict
+from diffusers.guiders import ClassifierFreeGuidance
+from diffusers.models import AutoModel, WanTransformer3DModel
+from diffusers.schedulers import UniPCMultistepScheduler
+from diffusers.utils import logging
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.modular_pipelines import (
+    BlockState,
+    LoopSequentialPipelineBlocks,
+    ModularPipelineBlocks,
+    PipelineState,
+    ModularPipeline
+)
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class MatrixGameWanLoopDenoiser(ModularPipelineBlocks):
+    model_name = "MatrixGameWan"
+    frame_seq_length = 880
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [
+            ComponentSpec(
+                "guider",
+                ClassifierFreeGuidance,
+                config=FrozenDict({"guidance_scale": 5.0}),
+                default_creation_method="from_config",
+            ),
+            ComponentSpec("transformer", AutoModel),
+        ]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Step within the denoising loop that denoise the latents with guidance. "
+            "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+            "object (e.g. `MatrixGameWanDenoiseLoopWrapper`)"
+        )
+
+    @property
+    def inputs(self) -> List[Tuple[str, Any]]:
+        return [
+            InputParam("attention_kwargs"),
+            InputParam(
+                "latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+            ),
+            InputParam(
+                "image_mask_latents",
+                required=True,
+                type_hint=torch.Tensor,
+            ),
+            InputParam(
+                "image_embeds",
+                required=True,
+                type_hint=torch.Tensor,
+            ),
+            InputParam(
+                "keyboard_conditions",
+                required=True,
+                type_hint=torch.Tensor,
+            ),
+            InputParam(
+                "mouse_conditions",
+                required=True,
+                type_hint=torch.Tensor,
+            ),
+            InputParam(
+                "num_inference_steps",
+                required=True,
+                type_hint=int,
+                default=4,
+                description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",
+            ),
+            InputParam(
+                kwargs_type="guider_input_fields",
+                description=(
+                    "All conditional model inputs that need to be prepared with guider. "
+                    "It should contain prompt_embeds/negative_prompt_embeds. "
+                    "Please add `kwargs_type=guider_input_fields` to their parameter spec (`OutputParam`) when they are created and added to the pipeline state"
+                ),
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(
+        self, components: ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+    ) -> PipelineState:
+        cond_concat = block_state.image_mask_latents
+        keyboard_conditions = block_state.keyboard_conditions
+        mouse_conditions = block_state.mouse_conditions
+        visual_context = block_state.image_embeds
+
+        transformer_dtype = components.transformer.dtype
+        components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+
+        # Prepare mini‐batches according to guidance method and `guider_input_fields`
+        # Each guider_state_batch will have .prompt_embeds, .time_ids, text_embeds, image_embeds.
+        # e.g. for CFG, we prepare two batches: one for uncond, one for cond
+        # for first batch, guider_state_batch.prompt_embeds correspond to block_state.prompt_embeds
+        # for second batch, guider_state_batch.prompt_embeds correspond to block_state.negative_prompt_embeds
+        guider_state = components.guider.prepare_inputs(block_state, {})
+
+        # run the denoiser for each guidance batch
+        for guider_state_batch in guider_state:
+            components.guider.prepare_models(components.transformer)
+            cond_kwargs = guider_state_batch.as_dict()
+
+            # Predict the noise residual
+            # store the noise_pred in guider_state_batch so that we can apply guidance across all batches
+            guider_state_batch.noise_pred = components.transformer(
+                x=block_state.latents.to(transformer_dtype),
+                t=t.expand(block_state.latents.shape[0], block_state.num_frames_per_block),
+                visual_context=visual_context.to(transformer_dtype),
+                cond_concat=cond_concat.to(transformer_dtype),
+                keyboard_cond=keyboard_conditions,
+                mouse_cond=mouse_conditions,
+                kv_cache=block_state.kv_cache,
+                kv_cache_mouse=block_state.kv_cache_mouse,
+                kv_cache_keyboard=block_state.kv_cache_keyboard,
+                crossattn_cache=block_state.kv_cache_cross_attn,
+                current_start=block_state.current_frame_idx * self.frame_seq_length,
+                num_frames_per_block=block_state.num_frames_per_block,
+            )[0]
+            components.guider.cleanup_models(components.transformer)
+
+        # Perform guidance
+        block_state.noise_pred = components.guider(guider_state)[0]
+
+        return components, block_state
+
+
+class MatrixGameWanLoopAfterDenoiser(ModularPipelineBlocks):
+    model_name = "MatrixGameWan"
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [
+            ComponentSpec("scheduler", UniPCMultistepScheduler),
+        ]
+
+    @property
+    def description(self) -> str:
+        return (
+            "step within the denoising loop that update the latents. "
+            "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+            "object (e.g. `MatrixGameWanDenoiseLoopWrapper`)"
+        )
+
+    @property
+    def inputs(self) -> List[Tuple[str, Any]]:
+        return []
+
+    @property
+    def intermediate_inputs(self) -> List[str]:
+        return [
+            InputParam("generator"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")]
+
+    @torch.no_grad()
+    def __call__(self, components: ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+        # Perform scheduler step using the predicted output
+        latents_dtype = block_state.latents.dtype
+
+        step_index = components.scheduler.index_for_timestep(t)
+        sigma_t = components.scheduler.sigmas[step_index]
+
+        latents = block_state.latents.double() - sigma_t.double() * block_state.noise_pred.double()
+        block_state.latents = latents
+
+        if block_state.latents.dtype != latents_dtype:
+            block_state.latents = block_state.latents.to(latents_dtype)
+
+        return components, block_state
+
+
+class MatrixGameWanDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+    model_name = "MatrixGameWan"
+    frame_seq_length = 880
+    local_attn_size = 6
+    num_transformer_blocks = 30
+
+    def _initialize_kv_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+        """
+        cache = []
+        if self.local_attn_size != -1:
+            # Use the local attention size to compute the KV cache size
+            kv_cache_size = self.local_attn_size * self.frame_seq_length
+        else:
+            # Use the default KV cache size
+            kv_cache_size = 15 * 1 * self.frame_seq_length # 32760
+
+        for _ in range(self.num_transformer_blocks):
+            cache.append({
+                "k": torch.zeros((batch_size, kv_cache_size, 12, 128), dtype=dtype, device=device),
+                "v": torch.zeros((batch_size, kv_cache_size, 12, 128), dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device)
+            })
+
+        return cache  # always store the clean cache
+
+    def _initialize_kv_cache_mouse_and_keyboard(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+        """
+        kv_cache_mouse = []
+        kv_cache_keyboard = []
+        if self.local_attn_size != -1:
+            kv_cache_size = self.local_attn_size
+        else:
+            kv_cache_size = 15 * 1
+        for _ in range(self.num_transformer_blocks):
+            kv_cache_keyboard.append({
+                "k": torch.zeros([batch_size, kv_cache_size, 16, 64], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, kv_cache_size, 16, 64], dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device)
+            })
+            kv_cache_mouse.append({
+                "k": torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, 16, 64], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, 16, 64], dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device)
+            })
+        return kv_cache_mouse, kv_cache_keyboard  # always store the clean cache
+
+    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU cross-attention cache for the Wan model.
+        """
+        crossattn_cache = []
+
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache.append({
+                "k": torch.zeros([batch_size, 257, 12, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, 257, 12, 128], dtype=dtype, device=device),
+                "is_init": False
+            })
+
+        return crossattn_cache
+
+    @property
+    def description(self) -> str:
+        return (
+            "Pipeline block that iteratively denoise the latents over `timesteps`. "
+            "The specific steps with each iteration can be customized with `sub_blocks` attributes"
+        )
+
+    @property
+    def loop_expected_components(self) -> List[ComponentSpec]:
+        return [
+            ComponentSpec(
+                "guider",
+                ClassifierFreeGuidance,
+                config=FrozenDict({"guidance_scale": 5.0}),
+                default_creation_method="from_config",
+            ),
+            ComponentSpec("scheduler", UniPCMultistepScheduler),
+            ComponentSpec("transformer", AutoModel),
+        ]
+
+    @property
+    def loop_inputs(self) -> List[InputParam]:
+        return [
+            InputParam(
+                "timesteps",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
+            ),
+            InputParam(
+                "num_inference_steps",
+                required=True,
+                type_hint=int,
+                description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",
+            ),
+            InputParam(
+                "num_frames_per_block",
+                required=True,
+                type_hint=int,
+                default=3,
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(
+        self, components: ModularPipeline, state: PipelineState
+    ) -> PipelineState:
+        block_state = self.get_block_state(state)
+        transformer_dtype = components.transformer.dtype
+
+        num_frames_per_block = block_state.num_frames_per_block
+        latents = block_state.latents.to(transformer_dtype)
+        image_mask_latents = block_state.image_mask_latents.to(transformer_dtype)
+        mouse_conditions = block_state.mouse_conditions.unsqueeze(0).to(transformer_dtype)
+        keyboard_conditions = block_state.keyboard_conditions.unsqueeze(0).to(transformer_dtype)
+        visual_context = block_state.image_embeds
+
+        batch_size, num_channels, num_frames, height, width = latents.shape
+        output = torch.zeros(
+            (batch_size, num_channels, num_frames, height, width),
+            device=latents.device,
+            dtype=latents.dtype,
+        )
+
+        current_frame_idx = 0
+        num_blocks = num_frames // num_frames_per_block
+
+        kv_cache = self._initialize_kv_cache(batch_size, latents.dtype, latents.device)
+        kv_cache_mouse, kv_cache_keyboard = self._initialize_kv_cache_mouse_and_keyboard(batch_size, latents.dtype, latents.device)
+        kv_cache_cross_attn = self._initialize_crossattn_cache(batch_size, latents.dtype, latents.device)
+
+        block_state.kv_cache = kv_cache
+        block_state.kv_cache_mouse = kv_cache_mouse
+        block_state.kv_cache_keyboard = kv_cache_keyboard
+        block_state.kv_cache_cross_attn = kv_cache_cross_attn
+
+        for _ in range(num_blocks):
+            block_state.current_frame_idx = current_frame_idx
+            block_state.image_mask_latents = image_mask_latents[
+                :, :, current_frame_idx : current_frame_idx + num_frames_per_block
+            ]
+            cond_idx = 1 + 4 * (current_frame_idx + num_frames_per_block - 1)
+            block_state.mouse_conditions = mouse_conditions[:, :cond_idx]
+            block_state.keyboard_conditions = keyboard_conditions[:, :cond_idx]
+
+            block_state.latents = latents[
+                :, :, current_frame_idx : current_frame_idx + num_frames_per_block
+            ]
+            for i, t in enumerate(block_state.timesteps):
+                components, block_state = self.loop_step(
+                    components, block_state, i=i, t=t
+                )
+
+                if i < (block_state.num_inference_steps - 1):
+                    t1 = components.scheduler.timesteps[i+1]
+                    block_state.latents = components.scheduler.add_noise(
+                        block_state.latents,
+                        randn_tensor(
+                            block_state.latents.shape,
+                            device=block_state.latents.device,
+                            dtype=block_state.latents.dtype
+                        ),
+                        t1.expand(block_state.latents.shape[0])
+                    )
+
+            output[
+                :, :, current_frame_idx : current_frame_idx + num_frames_per_block
+            ] = block_state.latents
+
+            components.transformer(
+                x=block_state.latents,
+                t=t.expand(block_state.latents.shape[0], block_state.num_frames_per_block) * 0.0,
+                visual_context=visual_context,
+                cond_concat=block_state.image_mask_latents,
+                keyboard_cond=block_state.keyboard_conditions,
+                mouse_cond=block_state.mouse_conditions,
+                kv_cache=block_state.kv_cache,
+                kv_cache_mouse=block_state.kv_cache_mouse,
+                kv_cache_keyboard=block_state.kv_cache_keyboard,
+                crossattn_cache=block_state.kv_cache_cross_attn,
+                current_start=block_state.current_frame_idx * self.frame_seq_length,
+                num_frames_per_block=block_state.num_frames_per_block,
+            )[0]
+            current_frame_idx += num_frames_per_block
+
+        block_state.latents = output
+        self.set_block_state(state, block_state)
+
+        return components, state
+
+
+class MatrixGameWanDenoiseStep(MatrixGameWanDenoiseLoopWrapper):
+    block_classes = [
+        MatrixGameWanLoopDenoiser,
+        MatrixGameWanLoopAfterDenoiser,
+    ]
+    block_names = ["denoiser", "after_denoiser"]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Denoise step that iteratively denoise the latents. \n"
+            "Its loop logic is defined in `MatrixGameWanDenoiseLoopWrapper.__call__` method \n"
+            "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n"
+            " - `MatrixGameWanLoopDenoiser`\n"
+            " - `MatrixGameWanLoopAfterDenoiser`\n"
+            "This block supports both text2vid tasks."
+        )

encoders.py

@@ -0,0 +1,89 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Dict, List
+
+import random
+import torch
+from torchvision.transforms import v2
+
+from diffusers.utils import logging
+from diffusers import ModularPipeline, ModularPipelineBlocks
+from diffusers.modular_pipelines import PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
+from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class MatrixGameWanImageEncoderStep(ModularPipelineBlocks):
+    model_name = "MatrixGameWan"
+
+    @property
+    def description(self) -> str:
+        return "Image Encoder step that generate image_embeddings to guide the video generation"
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [
+            ComponentSpec(
+                "image_encoder",
+                CLIPVisionModelWithProjection,
+                repo="laion/CLIP-ViT-H-14-laion2B-s32B-b79K",
+            ),
+            ComponentSpec(
+                "image_processor",
+                CLIPImageProcessor,
+                repo="Wan-AI/Wan2.1-I2V-14B-720P-Diffusers",
+                subfolder="image_processor"
+            ),
+        ]
+
+    @property
+    def expected_configs(self) -> List[ConfigSpec]:
+        return []
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam("image"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam(
+                "image_embeds",
+                type_hint=torch.Tensor,
+                description="image embeddings used to guide the image generation",
+            )
+        ]
+
+    def encode_image(self, components, image):
+        device = components._execution_device
+        image = components.image_processor(images=image, return_tensors="pt").to(device)
+        image_embeds = components.image_encoder(**image, output_hidden_states=True)
+        return image_embeds.hidden_states[-2]
+
+    @torch.no_grad()
+    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
+        # Get inputs and intermediates
+        block_state = self.get_block_state(state)
+        block_state.device = components._execution_device
+        #image_tensor = preprocess(block_state.image)
+        #image_tensor = image_tensor.to(block_state.device)
+        block_state.image_embeds = self.encode_image(components, block_state.image)
+
+        # Add outputs
+        self.set_block_state(state, block_state)
+        return components, state

modular_blocks.py

@@ -0,0 +1,66 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from diffusers.utils import logging
+from diffusers.modular_pipelines import SequentialPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline_utils import InsertableDict
+
+from .before_denoise import (
+    MatrixGameWanActionInputStep,
+    MatrixGameWanPrepareImageMaskLatentsStep,
+    MatrixGameWanPrepareLatentsStep,
+    MatrixGameWanSetTimestepsStep,
+)
+from .decoders import MatrixGameWanDecodeStep
+from .encoders import MatrixGameWanImageEncoderStep
+from .denoise import MatrixGameWanDenoiseStep
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class MatrixGameWanBeforeDenoiseStep(SequentialPipelineBlocks):
+    block_classes = [
+        MatrixGameWanActionInputStep,
+        MatrixGameWanSetTimestepsStep,
+        MatrixGameWanPrepareLatentsStep,
+        MatrixGameWanPrepareImageMaskLatentsStep,
+    ]
+    block_names = ["action_input", "set_timesteps", "prepare_latents", "prepare_mask_latents"]
+
+    @property
+    def description(self):
+        return (
+            "Before denoise step that prepare the inputs for the denoise step.\n"
+            + "This is a sequential pipeline blocks:\n"
+            + " - `MatrixGameWanInputStep` is used to adjust the batch size of the model inputs\n"
+            + " - `MatrixGameWanSetTimestepsStep` is used to set the timesteps\n"
+            + " - `MatrixGameWanPrepareLatentsStep` is used to prepare the latents\n"
+        )
+
+ACTION2VIDEO_BLOCKS = InsertableDict(
+    [
+        ("action_input", MatrixGameWanActionInputStep),
+        ("image_encoder", MatrixGameWanImageEncoderStep),
+        ("set_timesteps", MatrixGameWanSetTimestepsStep),
+        ("prepare_latents", MatrixGameWanPrepareLatentsStep),
+        ("prepare_masked_latents", MatrixGameWanPrepareImageMaskLatentsStep),
+        ("denoise", MatrixGameWanDenoiseStep),
+        ("decode", MatrixGameWanDecodeStep),
+    ]
+)
+
+ALL_BLOCKS = {
+    "action2video": ACTION2VIDEO_BLOCKS,
+}

modular_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "MatrixGameWanBlocks",
+  "_diffusers_version": "0.36.0.dev0",
+  "auto_map": {
+    "ModularPipelineBlocks": "block.MatrixGameWanBlocks"
+  }
+}

modular_model_index.json

@@ -0,0 +1,60 @@
+{
+  "_blocks_class_name": "SequentialPipelineBlocks",
+  "_class_name": "MatrixGameWanModularPipeline",
+  "_diffusers_version": "0.36.0.dev0",
+  "image_encoder": [
+    null,
+    null,
+    {
+      "repo": "laion/CLIP-ViT-H-14-laion2B-s32B-b79K",
+      "revision": null,
+      "type_hint": [
+        "transformers",
+        "CLIPVisionModelWithProjection"
+      ],
+      "variant": null
+    }
+  ],
+  "scheduler": [
+    null,
+    null,
+    {
+      "repo": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+      "revision": null,
+      "subfolder": "scheduler",
+      "type_hint": [
+        "diffusers",
+        "UniPCMultistepScheduler"
+      ],
+      "variant": null
+    }
+  ],
+  "transformer": [
+    null,
+    null,
+    {
+      "repo": "diffusers-internal-dev/matrix-game-2-modular",
+      "revision": null,
+      "subfolder": "transformer",
+      "type_hint": [
+        "diffusers",
+        "AutoModel"
+      ],
+      "variant": null
+    }
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLWan",
+    {
+      "repo": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+      "revision": null,
+      "subfolder": "vae",
+      "type_hint": [
+        "diffusers",
+        "AutoencoderKLWan"
+      ],
+      "variant": null
+    }
+  ]
+}

modular_pipeline.py

@@ -0,0 +1,85 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from diffusers.loaders import WanLoraLoaderMixin
+from diffusers.utils import logging
+from diffusers.modular_pipelines import ModularPipeline
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class MatrixGameWanModularPipeline(ModularPipeline, WanLoraLoaderMixin):
+    """
+    A ModularPipeline for MatrixGameWan.
+
+    <Tip warning={true}>
+
+        This is an experimental feature and is likely to change in the future.
+
+    </Tip>
+    """
+
+    @property
+    def default_height(self):
+        return self.default_sample_height * self.vae_scale_factor_spatial
+
+    @property
+    def default_width(self):
+        return self.default_sample_width * self.vae_scale_factor_spatial
+
+    @property
+    def default_num_frames(self):
+        return (self.default_sample_num_frames - 1) * self.vae_scale_factor_temporal + 1
+
+    @property
+    def default_sample_height(self):
+        return 44
+
+    @property
+    def default_sample_width(self):
+        return 80
+
+    @property
+    def default_sample_num_frames(self):
+        return 21
+
+    @property
+    def vae_scale_factor_spatial(self):
+        vae_scale_factor = 8
+        if hasattr(self, "vae") and self.vae is not None:
+            vae_scale_factor = 2 ** len(self.vae.temperal_downsample)
+        return vae_scale_factor
+
+    @property
+    def vae_scale_factor_temporal(self):
+        vae_scale_factor = 4
+        if hasattr(self, "vae") and self.vae is not None:
+            vae_scale_factor = 2 ** sum(self.vae.temperal_downsample)
+        return vae_scale_factor
+
+    @property
+    def num_channels_transformer(self):
+        num_channels_transformer = 16
+        if hasattr(self, "transformer") and self.transformer is not None:
+            num_channels_transformer = self.transformer.config.in_channels
+        return num_channels_transformer
+
+    @property
+    def num_channels_latents(self):
+        num_channels_latents = 16
+        if hasattr(self, "vae") and self.vae is not None:
+            num_channels_latents = self.vae.config.z_dim
+        return num_channels_latents

requirements.txt

@@ -0,0 +1,14 @@
+accelerate==1.10.1
+einops==0.8.1
+flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
+hf-transfer==0.1.9
+hf-xet==1.1.8
+huggingface-hub==0.34.4
+imageio==2.37.0
+imageio-ffmpeg==0.6.0
+safetensors==0.6.2
+sentencepiece==0.2.1
+torch==2.7.0
+torchao==0.12.0
+torchvision==0.22.0
+transformers==4.55.4

test_pipeline.py

@@ -0,0 +1,24 @@
+import sys
+import os
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+import torch
+from modular_pipeline import MatrixGameWanModularPipeline
+from modular_blocks import ACTION2VIDEO_BLOCKS
+from diffusers.modular_pipelines import SequentialPipelineBlocks
+from diffusers import AutoModel
+from diffusers.utils import load_image, export_to_video
+
+blocks = SequentialPipelineBlocks.from_blocks_dict(ACTION2VIDEO_BLOCKS.copy())
+pipe = MatrixGameWanModularPipeline(blocks)
+pipe.load_components(torch_dtype=torch.bfloat16)
+pipe.load_components(["vae"], repo="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
+pipe.load_components(["scheduler"], repo="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", subfolder="scheduler")
+
+transformer = AutoModel.from_pretrained("./transformer", trust_remote_code=True, torch_dtype=torch.bfloat16)
+pipe.transformer = transformer
+pipe.to("cuda")
+
+image = load_image("/home/dhruv/matrix-game-workspace/Matrix-Game/Matrix-Game-2/demo_images/universal/0000.png")
+output = pipe(image=image, num_frames=141)
+export_to_video(output.values['videos'][0], "modular-matrix-game.mp4")

transformer/action_module.py

@@ -0,0 +1,1148 @@
+from typing import Any, List, Tuple, Optional, Union, Dict
+from einops import rearrange
+from flash_attn import flash_attn_func
+import torch
+import torch.nn as nn
+import math
+from torch.nn.attention.flex_attention import flex_attention
+
+try:
+    import flash_attn
+
+except:
+    from flash_attn import flash_attn_func
+
+FLASH_ATTN_3_AVAILABLE = False
+
+
+DISABLE_COMPILE = False  # get os env
+flex_attention = torch.compile(
+    flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs"
+)
+
+import torch
+from typing import Union, Tuple, List
+
+
+def _to_tuple(x, dim=2):
+    if isinstance(x, int):
+        return (x,) * dim
+    elif len(x) == dim:
+        return x
+    else:
+        raise ValueError(f"Expected length {dim} or int, but got {x}")
+
+
+def get_meshgrid_nd(start, *args, dim=2):
+    """
+    Get n-D meshgrid with start, stop and num.
+
+    Args:
+        start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop,
+            step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num
+            should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in
+            n-tuples.
+        *args: See above.
+        dim (int): Dimension of the meshgrid. Defaults to 2.
+
+    Returns:
+        grid (np.ndarray): [dim, ...]
+    """
+    if len(args) == 0:
+        # start is grid_size
+        num = _to_tuple(start, dim=dim)
+        start = (0,) * dim
+        stop = num
+    elif len(args) == 1:
+        # start is start, args[0] is stop, step is 1
+        start = _to_tuple(start, dim=dim)
+        stop = _to_tuple(args[0], dim=dim)
+        num = [stop[i] - start[i] for i in range(dim)]
+    elif len(args) == 2:
+        # start is start, args[0] is stop, args[1] is num
+        start = _to_tuple(start, dim=dim)  # Left-Top       eg: 12,0
+        stop = _to_tuple(args[0], dim=dim)  # Right-Bottom   eg: 20,32
+        num = _to_tuple(args[1], dim=dim)  # Target Size    eg: 32,124
+    else:
+        raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}")
+
+    # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False)
+    axis_grid = []
+    for i in range(dim):
+        a, b, n = start[i], stop[i], num[i]
+        g = torch.linspace(a, b, n + 1, dtype=torch.float32, device=torch.cuda.current_device())[:n]
+        axis_grid.append(g)
+    grid = torch.meshgrid(*axis_grid, indexing="ij")  # dim x [W, H, D]
+    grid = torch.stack(grid, dim=0)  # [dim, W, H, D]
+
+    return grid
+
+
+#################################################################################
+#                   Rotary Positional Embedding Functions                       #
+#################################################################################
+# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L80
+
+
+def reshape_for_broadcast(
+    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+    x: torch.Tensor,
+    head_first=False,
+):
+    """
+    Reshape frequency tensor for broadcasting it with another tensor.
+
+    This function reshapes the frequency tensor to have the same shape as the target tensor 'x'
+    for the purpose of broadcasting the frequency tensor during element-wise operations.
+
+    Notes:
+        When using FlashMHAModified, head_first should be False.
+        When using Attention, head_first should be True.
+
+    Args:
+        freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped.
+        x (torch.Tensor): Target tensor for broadcasting compatibility.
+        head_first (bool): head dimension first (except batch dim) or not.
+
+    Returns:
+        torch.Tensor: Reshaped frequency tensor.
+
+    Raises:
+        AssertionError: If the frequency tensor doesn't match the expected shape.
+        AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions.
+    """
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+
+    if isinstance(freqs_cis, tuple):
+        # freqs_cis: (cos, sin) in real space
+        if head_first:
+            assert freqs_cis[0].shape == (
+                x.shape[-2],
+                x.shape[-1],
+            ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}"
+            shape = [
+                d if i == ndim - 2 or i == ndim - 1 else 1
+                for i, d in enumerate(x.shape)
+            ]
+        else:
+            # assert freqs_cis[0].shape == (
+            #     x.shape[1],
+            #     x.shape[-1],
+            # ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}"
+            # shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+            shape = [1, freqs_cis[0].shape[0], 1, freqs_cis[0].shape[1]]
+        return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape)
+    else:
+        # freqs_cis: values in complex space
+        if head_first:
+            assert freqs_cis.shape == (
+                x.shape[-2],
+                x.shape[-1],
+            ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}"
+            shape = [
+                d if i == ndim - 2 or i == ndim - 1 else 1
+                for i, d in enumerate(x.shape)
+            ]
+        else:
+            assert freqs_cis.shape == (
+                x.shape[1],
+                x.shape[-1],
+            ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}"
+            shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        return freqs_cis.view(*shape)
+
+
+def rotate_half(x):
+    x_real, x_imag = (
+        x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1)
+    )  # [B, S, H, D//2]
+    return torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
+    head_first: bool = False,
+    start_offset: int = 0,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor.
+
+    This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided
+    frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor
+    is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are
+    returned as real tensors.
+
+    Args:
+        xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D]
+        xk (torch.Tensor): Key tensor to apply rotary embeddings.   [B, S, H, D]
+        freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential.
+        head_first (bool): head dimension first (except batch dim) or not.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+
+    """
+    # print(freqs_cis[0].shape, xq.shape, xk.shape)
+    xk_out = None
+    assert isinstance(freqs_cis, tuple)
+    if isinstance(freqs_cis, tuple):
+        cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first)  # [S, D]
+        cos, sin = cos.to(xq.device), sin.to(xq.device)
+        # real * cos - imag * sin
+        # imag * cos + real * sin
+        xq_out = (xq.float() * cos[:, start_offset:start_offset + xq.shape[1], :, :] + rotate_half(xq.float()) * sin[:, start_offset:start_offset + xq.shape[1], :, :]).type_as(xq)
+        xk_out = (xk.float() * cos[:, start_offset:start_offset + xk.shape[1], :, :] + rotate_half(xk.float()) * sin[:, start_offset:start_offset + xk.shape[1], :, :]).type_as(xk)
+    else:
+        # view_as_complex will pack [..., D/2, 2](real) to [..., D/2](complex)
+        xq_ = torch.view_as_complex(
+            xq.float().reshape(*xq.shape[:-1], -1, 2)
+        )  # [B, S, H, D//2]
+        freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(
+            xq.device
+        )  # [S, D//2] --> [1, S, 1, D//2]
+        # (real, imag) * (cos, sin) = (real * cos - imag * sin, imag * cos + real * sin)
+        # view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real)
+        xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq)
+        xk_ = torch.view_as_complex(
+            xk.float().reshape(*xk.shape[:-1], -1, 2)
+        )  # [B, S, H, D//2]
+        xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk)
+
+    return xq_out, xk_out
+
+
+def get_nd_rotary_pos_embed(
+    rope_dim_list,
+    start,
+    *args,
+    theta=10000.0,
+    use_real=False,
+    theta_rescale_factor: Union[float, List[float]] = 1.0,
+    interpolation_factor: Union[float, List[float]] = 1.0,
+):
+    """
+    This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure.
+
+    Args:
+        rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n.
+            sum(rope_dim_list) should equal to head_dim of attention layer.
+        start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start,
+            args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num.
+        *args: See above.
+        theta (float): Scaling factor for frequency computation. Defaults to 10000.0.
+        use_real (bool): If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+            Some libraries such as TensorRT does not support complex64 data type. So it is useful to provide a real
+            part and an imaginary part separately.
+        theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0.
+
+    Returns:
+        pos_embed (torch.Tensor): [HW, D/2]
+    """
+
+    grid = get_meshgrid_nd(
+        start, *args, dim=len(rope_dim_list)
+    )  # [3, W, H, D] / [2, W, H]
+
+    if isinstance(theta_rescale_factor, int) or isinstance(theta_rescale_factor, float):
+        theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list)
+    elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1:
+        theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list)
+    assert len(theta_rescale_factor) == len(
+        rope_dim_list
+    ), "len(theta_rescale_factor) should equal to len(rope_dim_list)"
+
+    if isinstance(interpolation_factor, int) or isinstance(interpolation_factor, float):
+        interpolation_factor = [interpolation_factor] * len(rope_dim_list)
+    elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1:
+        interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list)
+    assert len(interpolation_factor) == len(
+        rope_dim_list
+    ), "len(interpolation_factor) should equal to len(rope_dim_list)"
+
+    # use 1/ndim of dimensions to encode grid_axis
+    embs = []
+    for i in range(len(rope_dim_list)):
+        emb = get_1d_rotary_pos_embed(
+            rope_dim_list[i],
+            grid[i].reshape(-1),
+            theta,
+            use_real=use_real,
+            theta_rescale_factor=theta_rescale_factor[i],
+            interpolation_factor=interpolation_factor[i],
+        )  # 2 x [WHD, rope_dim_list[i]]
+        embs.append(emb)
+
+    if use_real:
+        cos = torch.cat([emb[0] for emb in embs], dim=1)  # (WHD, D/2)
+        sin = torch.cat([emb[1] for emb in embs], dim=1)  # (WHD, D/2)
+        return cos, sin
+    else:
+        emb = torch.cat(embs, dim=1)  # (WHD, D/2)
+        return emb
+
+
+def get_1d_rotary_pos_embed(
+    dim: int,
+    pos: Union[torch.FloatTensor, int],
+    theta: float = 10000.0,
+    use_real: bool = False,
+    theta_rescale_factor: float = 1.0,
+    interpolation_factor: float = 1.0,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Precompute the frequency tensor for complex exponential (cis) with given dimensions.
+    (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.)
+
+    This function calculates a frequency tensor with complex exponential using the given dimension 'dim'
+    and the end index 'end'. The 'theta' parameter scales the frequencies.
+    The returned tensor contains complex values in complex64 data type.
+
+    Args:
+        dim (int): Dimension of the frequency tensor.
+        pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar
+        theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.
+        use_real (bool, optional): If True, return real part and imaginary part separately.
+                                   Otherwise, return complex numbers.
+        theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0.
+
+    Returns:
+        freqs_cis: Precomputed frequency tensor with complex exponential. [S, D/2]
+        freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D]
+    """
+    if isinstance(pos, int):
+        pos = torch.arange(pos, device=torch.cuda.current_device()).float()
+
+    # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+    # has some connection to NTK literature
+    if theta_rescale_factor != 1.0:
+        theta *= theta_rescale_factor ** (dim / (dim - 2))
+
+    freqs = 1.0 / (
+        theta ** (torch.arange(0, dim, 2, device=torch.cuda.current_device())[: (dim // 2)].float() / dim)
+    )  # [D/2]
+    # assert interpolation_factor == 1.0, f"interpolation_factor: {interpolation_factor}"
+    freqs = torch.outer(pos * interpolation_factor, freqs)  # [S, D/2]
+    if use_real:
+        freqs_cos = freqs.cos().repeat_interleave(2, dim=1)  # [S, D]
+        freqs_sin = freqs.sin().repeat_interleave(2, dim=1)  # [S, D]
+        return freqs_cos, freqs_sin
+    else:
+        freqs_cis = torch.polar(
+            torch.ones_like(freqs), freqs
+        )  # complex64     # [S, D/2]
+        return freqs_cis
+
+
+class MatrixGameWanRMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+
+class ActionModule(nn.Module):
+    """
+    action module from https://arxiv.org/pdf/2501.08325
+    鼠标控制信号的输入是一个 L*D 的向量
+    键盘同样
+    """
+
+    def __init__(
+        self,
+        mouse_dim_in: int = 2,
+        keyboard_dim_in: int = 6,
+        hidden_size: int = 128,
+        img_hidden_size: int = 1536,
+        keyboard_hidden_dim: int = 1024,
+        mouse_hidden_dim: int = 1024,
+        vae_time_compression_ratio: int = 4,
+        windows_size: int = 3,
+        heads_num: int = 16,
+        patch_size: list = [1, 2, 2],
+        qk_norm: bool = True,
+        qkv_bias: bool = False,
+        rope_dim_list: list = [8, 28, 28],
+        rope_theta=256,
+        mouse_qk_dim_list=[8, 28, 28],
+        enable_mouse=True,
+        enable_keyboard=True,
+        local_attn_size=6,
+        blocks=[],
+    ):
+        device = None
+
+        super().__init__()
+        self.local_attn_size = local_attn_size
+        self.enable_mouse = enable_mouse
+        self.enable_keyboard = enable_keyboard
+
+        self.rope_dim_list = rope_dim_list
+        self.rope_theta = rope_theta
+        if self.enable_keyboard:
+            self.keyboard_embed = nn.Sequential(
+                nn.Linear(keyboard_dim_in, hidden_size, bias=True),
+                nn.SiLU(),
+                nn.Linear(hidden_size, hidden_size, bias=True),
+            )
+
+        self.mouse_qk_dim_list = mouse_qk_dim_list
+        self.heads_num = heads_num
+        if self.enable_mouse:
+            c = mouse_hidden_dim
+            self.mouse_mlp = torch.nn.Sequential(
+                torch.nn.Linear(
+                    mouse_dim_in * vae_time_compression_ratio * windows_size
+                    + img_hidden_size,
+                    c,
+                    bias=True,
+                ),
+                torch.nn.GELU(approximate="tanh"),
+                torch.nn.Linear(c, c),
+                torch.nn.LayerNorm(c),
+            )
+
+            head_dim = c // heads_num
+            self.t_qkv = nn.Linear(c, c * 3, bias=qkv_bias)
+            self.img_attn_q_norm = (
+                MatrixGameWanRMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity()
+            )
+            self.img_attn_k_norm = (
+                MatrixGameWanRMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity()
+            )
+            self.proj_mouse = nn.Linear(c, img_hidden_size, bias=qkv_bias)
+
+        if self.enable_keyboard:
+            head_dim_key = keyboard_hidden_dim // heads_num
+            self.key_attn_q_norm = (
+                MatrixGameWanRMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity()
+            )
+            self.key_attn_k_norm = (
+                MatrixGameWanRMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity()
+            )
+
+            self.mouse_attn_q = nn.Linear(
+                img_hidden_size, keyboard_hidden_dim, bias=qkv_bias
+            )
+            self.keyboard_attn_kv = nn.Linear(
+                hidden_size * windows_size * vae_time_compression_ratio,
+                keyboard_hidden_dim * 2,
+                bias=qkv_bias,
+            )
+            self.proj_keyboard = nn.Linear(
+                keyboard_hidden_dim, img_hidden_size, bias=qkv_bias
+            )
+
+        self.vae_time_compression_ratio = vae_time_compression_ratio
+        self.windows_size = windows_size
+        self.patch_size = patch_size
+        self.freqs_cos, self.freqs_sin = self.get_rotary_pos_embed(
+            7500,
+            self.patch_size[1],
+            self.patch_size[2],
+            64,
+            self.mouse_qk_dim_list,
+            start_offset=0,
+        )
+
+    def patchify(self, x, patch_size):
+        """
+        x : (N C T H W)
+        """
+        pt, ph, pw = self.patch_size
+        t, h, w = x.shape[2] // pt, x.shape[3] // ph, x.shape[4] // pw
+        c = x.shape[1]
+        x = x.reshape(shape=(x.shape[0], c, t, pt, h, ph, w, pw))
+        x = torch.einsum("nctohpwq->nthwcopq", x)
+        x = x.reshape(shape=(x.shape[0], t * h * w, c * pt * ph * pw))
+        return x
+
+    def unpatchify(self, x, t, h, w, patch_size):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        c = x.shape[2] // patch_size  # self.unpatchify_channels
+        pt, ph, pw = self.patch_size
+        assert t * h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw))
+        x = torch.einsum("nthwcopq->nctohpwq", x)
+        imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw))
+
+        return imgs
+
+    def get_rotary_pos_embed(
+        self, video_length, height, width, head_dim, rope_dim_list=None, start_offset=0
+    ):
+        target_ndim = 3
+        ndim = 5 - 2
+
+        latents_size = [video_length + start_offset, height, width]
+
+        if isinstance(self.patch_size, int):
+            assert all(s % self.patch_size == 0 for s in latents_size), (
+                f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), "
+                f"but got {latents_size}."
+            )
+            rope_sizes = [s // self.patch_size for s in latents_size]
+        elif isinstance(self.patch_size, list):
+            assert all(
+                s % self.patch_size[idx] == 0 for idx, s in enumerate(latents_size)
+            ), (
+                f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), "
+                f"but got {latents_size}."
+            )
+            rope_sizes = [
+                s // self.patch_size[idx] for idx, s in enumerate(latents_size)
+            ]
+
+        if len(rope_sizes) != target_ndim:
+            rope_sizes = [1] * (target_ndim - len(rope_sizes)) + rope_sizes  # time axis
+
+        if rope_dim_list is None:
+            rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)]
+        assert (
+            sum(rope_dim_list) == head_dim
+        ), "sum(rope_dim_list) should equal to head_dim of attention layer"
+        freqs_cos, freqs_sin = get_nd_rotary_pos_embed(
+            rope_dim_list,
+            rope_sizes,
+            theta=self.rope_theta,
+            use_real=True,
+            theta_rescale_factor=1,
+        )
+        return freqs_cos[
+            -video_length * rope_sizes[1] * rope_sizes[2] // self.patch_size[0] :
+        ], freqs_sin[
+            -video_length * rope_sizes[1] * rope_sizes[2] // self.patch_size[0] :
+        ]
+
+    def forward(
+        self,
+        x,
+        tt,
+        th,
+        tw,
+        mouse_condition=None,
+        keyboard_condition=None,
+        block_mask_mouse=None,
+        block_mask_keyboard=None,
+        is_causal=False,
+        kv_cache_mouse=None,
+        kv_cache_keyboard=None,
+        start_frame=0,
+        use_rope_keyboard=True,
+        num_frame_per_block=3,
+    ):
+        """
+        hidden_states: B, tt*th*tw, C
+        mouse_condition: B, N_frames, C1
+        keyboard_condition: B, N_frames, C2
+        """
+        assert use_rope_keyboard == True
+
+        B, N_frames, C = keyboard_condition.shape
+
+        assert tt * th * tw == x.shape[1]
+        assert (
+            (N_frames - 1) + self.vae_time_compression_ratio
+        ) % self.vae_time_compression_ratio == 0
+        N_feats = int((N_frames - 1) / self.vae_time_compression_ratio) + 1
+
+        # Defined freqs_cis early so it's available for both mouse and keyboard
+        freqs_cis = (self.freqs_cos, self.freqs_sin)
+
+        assert (
+            N_feats == tt and ((is_causal and kv_cache_mouse == None) or not is_causal)
+        ) or (
+            (N_frames - 1) // self.vae_time_compression_ratio + 1 == start_frame + num_frame_per_block and is_causal
+        )
+
+        if self.enable_mouse and mouse_condition is not None:
+            hidden_states = rearrange(
+                x, "B (T S) C -> (B S) T C", T=tt, S=th * tw
+            )  # 65*272*480 -> 17*(272//16)*(480//16) -> 8670
+            B, N_frames, C = mouse_condition.shape
+        else:
+            hidden_states = x
+        # padding
+
+        pad_t = self.vae_time_compression_ratio * self.windows_size
+        if self.enable_mouse and mouse_condition is not None:
+            pad = mouse_condition[:, 0:1, :].expand(-1, pad_t, -1)
+            mouse_condition = torch.cat([pad, mouse_condition], dim=1)
+            if is_causal and kv_cache_mouse is not None:
+                mouse_condition = mouse_condition[
+                    :,
+                    self.vae_time_compression_ratio
+                    * (N_feats - num_frame_per_block - self.windows_size)
+                    + pad_t :,
+                    :,
+                ]
+                group_mouse = [
+                    mouse_condition[
+                        :,
+                        self.vae_time_compression_ratio * (i - self.windows_size)
+                        + pad_t : i * self.vae_time_compression_ratio + pad_t,
+                        :,
+                    ]
+                    for i in range(num_frame_per_block)
+                ]
+            else:
+                group_mouse = [
+                    mouse_condition[
+                        :,
+                        self.vae_time_compression_ratio * (i - self.windows_size)
+                        + pad_t : i * self.vae_time_compression_ratio + pad_t,
+                        :,
+                    ]
+                    for i in range(N_feats)
+                ]
+
+            group_mouse = torch.stack(group_mouse, dim=1)
+
+            S = th * tw
+            group_mouse = group_mouse.unsqueeze(-1).expand(
+                B, num_frame_per_block, pad_t, C, S
+            )
+            group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape(
+                B * S, num_frame_per_block, pad_t * C
+            )
+
+            group_mouse = torch.cat([hidden_states, group_mouse], dim=-1)
+            group_mouse = self.mouse_mlp(group_mouse)
+
+            # qkv
+            mouse_qkv = self.t_qkv(group_mouse)
+            q, k, v = rearrange(
+                mouse_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num
+            )  # BHW F H C
+            q = self.img_attn_q_norm(q).to(v)
+            k = self.img_attn_k_norm(k).to(v)
+            # rope embd
+
+            # freqs_cis = (self.freqs_cos, self.freqs_sin)
+
+            q, k = apply_rotary_emb(
+                q, k, freqs_cis, start_offset=start_frame, head_first=False
+            )
+            ## TODO: adding cache here
+            if is_causal:
+                if kv_cache_mouse is None:
+                    assert (
+                        q.shape[0] == k.shape[0] and q.shape[0] % 880 == 0
+                    )  # == 880, f"{q.shape[0]},{k.shape[0]}"
+                    padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1]
+                    padded_q = torch.cat(
+                        [
+                            q,
+                            torch.zeros(
+                                [q.shape[0], padded_length, q.shape[2], q.shape[3]],
+                                device=q.device,
+                                dtype=v.dtype,
+                            ),
+                        ],
+                        dim=1,
+                    )
+                    padded_k = torch.cat(
+                        [
+                            k,
+                            torch.zeros(
+                                [k.shape[0], padded_length, k.shape[2], k.shape[3]],
+                                device=k.device,
+                                dtype=v.dtype,
+                            ),
+                        ],
+                        dim=1,
+                    )
+                    padded_v = torch.cat(
+                        [
+                            v,
+                            torch.zeros(
+                                [v.shape[0], padded_length, v.shape[2], v.shape[3]],
+                                device=v.device,
+                                dtype=v.dtype,
+                            ),
+                        ],
+                        dim=1,
+                    )
+                    attn = flex_attention(
+                        query=padded_q.transpose(2, 1),  # after: B, HW, F, C
+                        key=padded_k.transpose(2, 1),
+                        value=padded_v.transpose(2, 1),
+                        block_mask=block_mask_mouse,
+                    )[:, :, :-padded_length].transpose(2, 1)
+                else:
+                    current_start = start_frame
+                    current_end = current_start + q.shape[1]
+
+                    assert q.shape[1] == num_frame_per_block
+                    sink_size = 0
+                    max_attention_size = self.local_attn_size
+                    sink_tokens = sink_size * 1
+                    kv_cache_size = kv_cache_mouse["k"].shape[1]
+                    num_new_tokens = q.shape[1]
+
+                    if (current_end > kv_cache_mouse["global_end_index"].item()) and (
+                        num_new_tokens + kv_cache_mouse["local_end_index"].item()
+                        > kv_cache_size
+                    ):
+                        num_evicted_tokens = (
+                            num_new_tokens
+                            + kv_cache_mouse["local_end_index"].item()
+                            - kv_cache_size
+                        )
+                        num_rolled_tokens = (
+                            kv_cache_mouse["local_end_index"].item()
+                            - num_evicted_tokens
+                            - sink_tokens
+                        )
+                        kv_cache_mouse["k"][
+                            :, sink_tokens : sink_tokens + num_rolled_tokens
+                        ] = kv_cache_mouse["k"][
+                            :,
+                            sink_tokens + num_evicted_tokens : sink_tokens
+                            + num_evicted_tokens
+                            + num_rolled_tokens,
+                        ].clone()
+                        kv_cache_mouse["v"][
+                            :, sink_tokens : sink_tokens + num_rolled_tokens
+                        ] = kv_cache_mouse["v"][
+                            :,
+                            sink_tokens + num_evicted_tokens : sink_tokens
+                            + num_evicted_tokens
+                            + num_rolled_tokens,
+                        ].clone()
+                        # Insert the new keys/values at the end
+                        local_end_index = (
+                            kv_cache_mouse["local_end_index"].item()
+                            + current_end
+                            - kv_cache_mouse["global_end_index"].item()
+                            - num_evicted_tokens
+                        )
+                        local_start_index = local_end_index - num_new_tokens
+                    else:
+                        local_end_index = (
+                            kv_cache_mouse["local_end_index"].item()
+                            + current_end
+                            - kv_cache_mouse["global_end_index"].item()
+                        )
+                        local_start_index = local_end_index - num_new_tokens
+
+                    kv_cache_mouse["k"][:, local_start_index:local_end_index] = k
+                    kv_cache_mouse["v"][:, local_start_index:local_end_index] = v
+
+                    if FLASH_ATTN_3_AVAILABLE:
+                        attn, attn_prob = flash_attn.flash_attn_func(
+                            q,
+                            kv_cache_mouse["k"][
+                                :,
+                                max(
+                                    0, local_end_index - max_attention_size
+                                ) : local_end_index,
+                            ],
+                            kv_cache_mouse["v"][
+                                :,
+                                max(
+                                    0, local_end_index - max_attention_size
+                                ) : local_end_index,
+                            ],
+                        )
+                    else:
+                        attn = flash_attn_func(
+                            q,
+                            kv_cache_mouse["k"][
+                                :,
+                                max(
+                                    0, local_end_index - max_attention_size
+                                ) : local_end_index,
+                            ],
+                            kv_cache_mouse["v"][
+                                :,
+                                max(
+                                    0, local_end_index - max_attention_size
+                                ) : local_end_index,
+                            ],
+                        )
+                    kv_cache_mouse["global_end_index"].fill_(current_end)
+                    kv_cache_mouse["local_end_index"].fill_(local_end_index)
+            else:
+                attn = flash_attn_func(
+                    q,  # 880, f, 16, 64
+                    k,  # 880, f, 16, 64
+                    v,  # 880, f, 16, 64
+                )
+            # Compute cu_squlens and max_seqlen for flash attention
+            # qk norm
+            attn = rearrange(attn, "(b S) T h d -> b (T S) (h d)", b=B)
+
+            hidden_states = rearrange(x, "(B S) T C -> B (T S) C", B=B)
+            attn = self.proj_mouse(attn)
+
+            hidden_states = hidden_states + attn
+
+        if self.enable_keyboard and keyboard_condition is not None:
+            pad = keyboard_condition[:, 0:1, :].expand(-1, pad_t, -1)
+            keyboard_condition = torch.cat([pad, keyboard_condition], dim=1)
+            if is_causal and kv_cache_keyboard is not None:
+                keyboard_condition = keyboard_condition[
+                    :,
+                    self.vae_time_compression_ratio
+                    * (N_feats - num_frame_per_block - self.windows_size)
+                    + pad_t :,
+                    :,
+                ]  # keyboard_condition[:, self.vae_time_compression_ratio*(start_frame - self.windows_size) + pad_t:start_frame * self.vae_time_compression_ratio + pad_t,:]
+                keyboard_condition = self.keyboard_embed(keyboard_condition)
+                group_keyboard = [
+                    keyboard_condition[
+                        :,
+                        self.vae_time_compression_ratio * (i - self.windows_size)
+                        + pad_t : i * self.vae_time_compression_ratio + pad_t,
+                        :,
+                    ]
+                    for i in range(num_frame_per_block)
+                ]
+            else:
+                keyboard_condition = self.keyboard_embed(keyboard_condition)
+                group_keyboard = [
+                    keyboard_condition[
+                        :,
+                        self.vae_time_compression_ratio * (i - self.windows_size)
+                        + pad_t : i * self.vae_time_compression_ratio + pad_t,
+                        :,
+                    ]
+                    for i in range(N_feats)
+                ]
+            group_keyboard = torch.stack(group_keyboard, dim=1)  # B F RW C
+            group_keyboard = group_keyboard.reshape(
+                shape=(group_keyboard.shape[0], group_keyboard.shape[1], -1)
+            )
+            # apply cross attn
+            mouse_q = self.mouse_attn_q(hidden_states)
+            keyboard_kv = self.keyboard_attn_kv(group_keyboard)
+
+            B, L, HD = mouse_q.shape
+            D = HD // self.heads_num
+            q = mouse_q.view(B, L, self.heads_num, D)
+
+            B, L, KHD = keyboard_kv.shape
+            k, v = keyboard_kv.view(B, L, 2, self.heads_num, D).permute(2, 0, 1, 3, 4)
+
+            # Compute cu_squlens and max_seqlen for flash attention
+            # qk norm
+
+            q = self.key_attn_q_norm(q).to(v)
+            k = self.key_attn_k_norm(k).to(v)
+            S = th * tw
+            assert S == 880
+            # position embed
+            if use_rope_keyboard:
+                B, TS, H, D = q.shape
+                T_ = TS // S
+                q = q.view(B, T_, S, H, D).transpose(1, 2).reshape(B * S, T_, H, D)
+                q, k = apply_rotary_emb(
+                    q, k, freqs_cis, start_offset=start_frame, head_first=False
+                )
+
+                k1, k2, k3, k4 = k.shape
+                k = k.expand(S, k2, k3, k4)
+                v = v.expand(S, k2, k3, k4)
+
+                if is_causal:
+                    if kv_cache_keyboard is None:
+                        assert q.shape[0] == k.shape[0] and q.shape[0] % 880 == 0
+
+                        padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1]
+                        padded_q = torch.cat(
+                            [
+                                q,
+                                torch.zeros(
+                                    [q.shape[0], padded_length, q.shape[2], q.shape[3]],
+                                    device=q.device,
+                                    dtype=v.dtype,
+                                ),
+                            ],
+                            dim=1,
+                        )
+                        padded_k = torch.cat(
+                            [
+                                k,
+                                torch.zeros(
+                                    [k.shape[0], padded_length, k.shape[2], k.shape[3]],
+                                    device=k.device,
+                                    dtype=v.dtype,
+                                ),
+                            ],
+                            dim=1,
+                        )
+                        padded_v = torch.cat(
+                            [
+                                v,
+                                torch.zeros(
+                                    [v.shape[0], padded_length, v.shape[2], v.shape[3]],
+                                    device=v.device,
+                                    dtype=v.dtype,
+                                ),
+                            ],
+                            dim=1,
+                        )
+                        attn = flex_attention(
+                            query=padded_q.transpose(2, 1),  # after: B, HW, F, C
+                            key=padded_k.transpose(2, 1),
+                            value=padded_v.transpose(2, 1),
+                            block_mask=block_mask_keyboard,
+                        )[:, :, :-padded_length].transpose(2, 1)
+                    else:
+                        current_start = start_frame
+                        current_end = current_start + k.shape[1]
+                        assert k.shape[1] == num_frame_per_block
+                        sink_size = 0
+                        max_attention_size = self.local_attn_size
+                        sink_tokens = sink_size * 1
+                        kv_cache_size = kv_cache_keyboard["k"].shape[1]
+                        num_new_tokens = k.shape[1]
+
+                        if (
+                            current_end > kv_cache_keyboard["global_end_index"].item()
+                        ) and (
+                            num_new_tokens + kv_cache_keyboard["local_end_index"].item()
+                            > kv_cache_size
+                        ):
+                            num_evicted_tokens = (
+                                num_new_tokens
+                                + kv_cache_keyboard["local_end_index"].item()
+                                - kv_cache_size
+                            )
+                            num_rolled_tokens = (
+                                kv_cache_keyboard["local_end_index"].item()
+                                - num_evicted_tokens
+                                - sink_tokens
+                            )
+                            kv_cache_keyboard["k"][
+                                :, sink_tokens : sink_tokens + num_rolled_tokens
+                            ] = kv_cache_keyboard["k"][
+                                :,
+                                sink_tokens + num_evicted_tokens : sink_tokens
+                                + num_evicted_tokens
+                                + num_rolled_tokens,
+                            ].clone()
+                            kv_cache_keyboard["v"][
+                                :, sink_tokens : sink_tokens + num_rolled_tokens
+                            ] = kv_cache_keyboard["v"][
+                                :,
+                                sink_tokens + num_evicted_tokens : sink_tokens
+                                + num_evicted_tokens
+                                + num_rolled_tokens,
+                            ].clone()
+                            # Insert the new keys/values at the end
+                            local_end_index = (
+                                kv_cache_keyboard["local_end_index"].item()
+                                + current_end
+                                - kv_cache_keyboard["global_end_index"].item()
+                                - num_evicted_tokens
+                            )
+                            local_start_index = local_end_index - num_new_tokens
+                        else:
+                            local_end_index = (
+                                kv_cache_keyboard["local_end_index"].item()
+                                + current_end
+                                - kv_cache_keyboard["global_end_index"].item()
+                            )
+                            local_start_index = local_end_index - num_new_tokens
+                        assert (
+                            k.shape[0] == 880
+                        )  # BS == 1 or the cache should not be saved/ load method should be modified
+                        kv_cache_keyboard["k"][:, local_start_index:local_end_index] = (
+                            k[:1]
+                        )
+                        kv_cache_keyboard["v"][:, local_start_index:local_end_index] = (
+                            v[:1]
+                        )
+
+                        if FLASH_ATTN_3_AVAILABLE:
+                            attn, attn_prob = flash_attn.flash_attn_func(
+                                q,
+                                kv_cache_keyboard["k"][
+                                    :,
+                                    max(
+                                        0, local_end_index - max_attention_size
+                                    ) : local_end_index,
+                                ].repeat(S, 1, 1, 1),
+                                kv_cache_keyboard["v"][
+                                    :,
+                                    max(
+                                        0, local_end_index - max_attention_size
+                                    ) : local_end_index,
+                                ].repeat(S, 1, 1, 1),
+                            )
+                        else:
+                            attn = flash_attn_func(
+                                q,
+                                kv_cache_keyboard["k"][
+                                    :,
+                                    max(
+                                        0, local_end_index - max_attention_size
+                                    ) : local_end_index,
+                                ].repeat(S, 1, 1, 1),
+                                kv_cache_keyboard["v"][
+                                    :,
+                                    max(
+                                        0, local_end_index - max_attention_size
+                                    ) : local_end_index,
+                                ].repeat(S, 1, 1, 1),
+                            )
+
+                        kv_cache_keyboard["global_end_index"].fill_(current_end)
+                        kv_cache_keyboard["local_end_index"].fill_(local_end_index)
+                else:
+                    attn = flash_attn_func(
+                        q,  # 1, f*880, 16, 64
+                        k,  # 1, f, 16, 64
+                        v,  # 1, f, 16, 64
+                        causal=False,
+                    )
+                attn = rearrange(attn, "(B S) T H D -> B (T S) (H D)", S=S)
+            else:
+                if is_causal:
+                    if kv_cache_keyboard is None:
+                        padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1]
+                        padded_q = torch.cat(
+                            [
+                                q,
+                                torch.zeros(
+                                    [q.shape[0], padded_length, q.shape[2], q.shape[3]],
+                                    device=q.device,
+                                    dtype=v.dtype,
+                                ),
+                            ],
+                            dim=1,
+                        )
+                        padded_k = torch.cat(
+                            [
+                                k,
+                                torch.zeros(
+                                    [k.shape[0], padded_length, k.shape[2], k.shape[3]],
+                                    device=k.device,
+                                    dtype=v.dtype,
+                                ),
+                            ],
+                            dim=1,
+                        )
+                        padded_v = torch.cat(
+                            [
+                                v,
+                                torch.zeros(
+                                    [v.shape[0], padded_length, v.shape[2], v.shape[3]],
+                                    device=v.device,
+                                    dtype=v.dtype,
+                                ),
+                            ],
+                            dim=1,
+                        )
+                        attn = flex_attention(
+                            query=padded_q.transpose(2, 1),  # after: B, HW, F, C
+                            key=padded_k.transpose(2, 1),
+                            value=padded_v.transpose(2, 1),
+                            block_mask=block_mask_keyboard,
+                        )[:, :, :-padded_length].transpose(2, 1)
+                    else:
+                        current_start = start_frame
+                        current_end = current_start + k.shape[1]
+                        assert k.shape[1] == num_frame_per_block
+                        sink_size = 0
+                        local_attn_size = self.local_attn_size
+                        max_attention_size = self.local_attn_size
+                        sink_tokens = sink_size * 1
+                        kv_cache_size = kv_cache_keyboard["k"].shape[1]
+                        num_new_tokens = k.shape[1]
+
+                        if (
+                            current_end > kv_cache_keyboard["global_end_index"].item()
+                        ) and (
+                            num_new_tokens + kv_cache_keyboard["local_end_index"].item()
+                            > kv_cache_size
+                        ):
+                            num_evicted_tokens = (
+                                num_new_tokens
+                                + kv_cache_keyboard["local_end_index"].item()
+                                - kv_cache_size
+                            )
+                            num_rolled_tokens = (
+                                kv_cache_keyboard["local_end_index"].item()
+                                - num_evicted_tokens
+                                - sink_tokens
+                            )
+                            kv_cache_keyboard["k"][
+                                :, sink_tokens : sink_tokens + num_rolled_tokens
+                            ] = kv_cache_keyboard["k"][
+                                :,
+                                sink_tokens + num_evicted_tokens : sink_tokens
+                                + num_evicted_tokens
+                                + num_rolled_tokens,
+                            ].clone()
+                            kv_cache_keyboard["v"][
+                                :, sink_tokens : sink_tokens + num_rolled_tokens
+                            ] = kv_cache_keyboard["v"][
+                                :,
+                                sink_tokens + num_evicted_tokens : sink_tokens
+                                + num_evicted_tokens
+                                + num_rolled_tokens,
+                            ].clone()
+                            # Insert the new keys/values at the end
+                            local_end_index = (
+                                kv_cache_keyboard["local_end_index"].item()
+                                + current_end
+                                - kv_cache_keyboard["global_end_index"].item()
+                                - num_evicted_tokens
+                            )
+                            local_start_index = local_end_index - num_new_tokens
+
+                        else:
+                            local_end_index = (
+                                kv_cache_keyboard["local_end_index"].item()
+                                + current_end
+                                - kv_cache_keyboard["global_end_index"].item()
+                            )
+                            local_start_index = local_end_index - num_new_tokens
+                        kv_cache_keyboard["k"][:, local_start_index:local_end_index] = k
+                        kv_cache_keyboard["v"][:, local_start_index:local_end_index] = v
+                        attn = flash_attn_func(
+                            q,
+                            kv_cache_keyboard["k"][
+                                :,
+                                max(
+                                    0, local_end_index - max_attention_size
+                                ) : local_end_index,
+                            ],
+                            kv_cache_keyboard["v"][
+                                :,
+                                max(
+                                    0, local_end_index - max_attention_size
+                                ) : local_end_index,
+                            ],
+                            # causal=is_causal
+                        )
+                        kv_cache_keyboard["global_end_index"].fill_(current_end)
+                        kv_cache_keyboard["local_end_index"].fill_(local_end_index)
+                else:
+                    attn = flash_attn_func(
+                        q,  # 1, f*880, 16, 64
+                        k,  # 1, f, 16, 64
+                        v,  # 1, f, 16, 64
+                        # causal=is_causal,
+                    )
+                attn = rearrange(attn, "B L H D -> B L (H D)")
+            attn = self.proj_keyboard(attn)
+            hidden_states = hidden_states + attn
+        return hidden_states

transformer/attention.py

@@ -0,0 +1,198 @@
+# Copyright 2024-2025 The Alibaba MatrixGameWan Team Authors. All rights reserved.
+import torch
+
+try:
+    import flash_attn
+
+    def is_hopper_gpu():
+        if not torch.cuda.is_available():
+            return False
+        device_name = torch.cuda.get_device_name(0).lower()
+        return (
+            "h100" in device_name
+            or "hopper" in device_name
+            or "l20y" in device_name
+            or "h800" in device_name
+        )
+
+    FLASH_ATTN_3_AVAILABLE = is_hopper_gpu()
+except ModuleNotFoundError:
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+
+    FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_2_AVAILABLE = False
+
+
+import warnings
+
+__all__ = [
+    "flash_attention",
+    "attention",
+]
+
+
+def flash_attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    version=None,
+):
+    """
+    q:              [B, Lq, Nq, C1].
+    k:              [B, Lk, Nk, C1].
+    v:              [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+    q_lens:         [B].
+    k_lens:         [B].
+    dropout_p:      float. Dropout probability.
+    softmax_scale:  float. The scaling of QK^T before applying softmax.
+    causal:         bool. Whether to apply causal attention mask.
+    window_size:    (left right). If not (-1, -1), apply sliding window local attention.
+    deterministic:  bool. If True, slightly slower and uses more memory.
+    dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+    """
+    half_dtypes = (torch.float16, torch.bfloat16)
+    assert dtype in half_dtypes
+    assert q.device.type == "cuda" and q.size(-1) <= 256
+
+    # params
+    b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    # preprocess query
+    if q_lens is None:
+        q = half(q.flatten(0, 1))
+        q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(
+            device=q.device, non_blocking=True
+        )
+    else:
+        q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+    # preprocess key, value
+    if k_lens is None:
+        k = half(k.flatten(0, 1))
+        v = half(v.flatten(0, 1))
+        k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(
+            device=k.device, non_blocking=True
+        )
+    else:
+        k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+        v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+    q = q.to(v.dtype)
+    k = k.to(v.dtype)
+
+    if q_scale is not None:
+        q = q * q_scale
+
+    if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+        warnings.warn(
+            "Flash attention 3 is not available, use flash attention 2 instead."
+        )
+
+    # apply attention
+    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+        # Note: dropout_p, window_size are not supported in FA3 now.
+        x = flash_attn.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            deterministic=deterministic,
+        )[0].unflatten(0, (b, lq))
+    else:
+        assert FLASH_ATTN_2_AVAILABLE
+        x = flash_attn.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+        ).unflatten(0, (b, lq))
+
+    # output
+    return x.type(out_dtype)
+
+
+def attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    fa_version=None,
+):
+    if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+        return flash_attention(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+            dtype=dtype,
+            version=fa_version,
+        )
+    else:
+        if q_lens is not None or k_lens is not None:
+            warnings.warn(
+                "Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance."
+            )
+        attn_mask = None
+
+        q = q.transpose(1, 2).to(dtype)
+        k = k.transpose(1, 2).to(dtype)
+        v = v.transpose(1, 2).to(dtype)
+
+        out = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p
+        )
+
+        out = out.transpose(1, 2).contiguous()
+        return out

transformer/causal_model.py

@@ -0,0 +1,949 @@
+from .attention import attention
+from .model import (
+    MatrixGameWanRMSNorm,
+    rope_apply,
+    MatrixGameWanLayerNorm,
+    MatrixGameWan_CROSSATTENTION_CLASSES,
+    rope_params,
+    MLPProj,
+    sinusoidal_embedding_1d,
+)
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from torch.nn.attention.flex_attention import BlockMask
+from diffusers.models.modeling_utils import ModelMixin
+import torch.nn as nn
+import torch
+import math
+import torch.distributed as dist
+from .action_module import ActionModule
+
+
+def causal_rope_apply(x, grid_sizes, freqs, start_frame=0):
+    n, c = x.size(2), x.size(3) // 2
+
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+    # loop over samples
+    output = []
+    f, h, w = grid_sizes.tolist()
+
+    for i in range(len(x)):
+        seq_len = f * h * w
+
+        # precompute multipliers
+        x_i = torch.view_as_complex(
+            x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)
+        )
+        freqs_i = torch.cat(
+            [
+                freqs[0][start_frame : start_frame + f]
+                .view(f, 1, 1, -1)
+                .expand(f, h, w, -1),
+                freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+                freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1),
+            ],
+            dim=-1,
+        ).reshape(seq_len, 1, -1)
+
+        # apply rotary embedding
+        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+        x_i = torch.cat([x_i, x[i, seq_len:]])
+
+        # append to collection
+        output.append(x_i)
+    return torch.stack(output).type_as(x)
+
+
+class MatrixGameWanCausalSelfAttention(nn.Module):
+    def __init__(
+        self, dim, num_heads, local_attn_size=-1, sink_size=0, qk_norm=True, eps=1e-6
+    ):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+        self.max_attention_size = (
+            15 * 1 * 880 if local_attn_size == -1 else local_attn_size * 880
+        )
+        # layers
+        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 = MatrixGameWanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = MatrixGameWanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(
+        self,
+        x,
+        seq_lens,
+        grid_sizes,
+        freqs,
+        block_mask,
+        kv_cache=None,
+        current_start=0,
+        cache_start=None,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C] # num_heads, C / num_heads]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+            block_mask (BlockMask)
+        """
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+        if cache_start is None:
+            cache_start = current_start
+
+        # query, key, value function
+        def qkv_fn(x):
+            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)
+            return q, k, v
+
+        q, k, v = qkv_fn(x)  # B, F, HW, C
+
+        if kv_cache is None:
+            roped_query = rope_apply(q, grid_sizes, freqs).type_as(v)
+            roped_key = rope_apply(k, grid_sizes, freqs).type_as(v)
+
+            padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
+            padded_roped_query = torch.cat(
+                [
+                    roped_query,
+                    torch.zeros(
+                        [q.shape[0], padded_length, q.shape[2], q.shape[3]],
+                        device=q.device,
+                        dtype=v.dtype,
+                    ),
+                ],
+                dim=1,
+            )
+
+            padded_roped_key = torch.cat(
+                [
+                    roped_key,
+                    torch.zeros(
+                        [k.shape[0], padded_length, k.shape[2], k.shape[3]],
+                        device=k.device,
+                        dtype=v.dtype,
+                    ),
+                ],
+                dim=1,
+            )
+
+            padded_v = torch.cat(
+                [
+                    v,
+                    torch.zeros(
+                        [v.shape[0], padded_length, v.shape[2], v.shape[3]],
+                        device=v.device,
+                        dtype=v.dtype,
+                    ),
+                ],
+                dim=1,
+            )
+
+            x = flex_attention(
+                query=padded_roped_query.transpose(2, 1),  # after: B, HW, F, C
+                key=padded_roped_key.transpose(2, 1),
+                value=padded_v.transpose(2, 1),
+                block_mask=block_mask,
+            )[:, :, :-padded_length].transpose(2, 1)
+        else:
+            assert grid_sizes.ndim == 1
+            frame_seqlen = math.prod(grid_sizes[1:]).item()
+            current_start_frame = current_start // frame_seqlen
+            roped_query = causal_rope_apply(
+                q, grid_sizes, freqs, start_frame=current_start_frame
+            ).type_as(v)
+            roped_key = causal_rope_apply(
+                k, grid_sizes, freqs, start_frame=current_start_frame
+            ).type_as(v)
+
+            current_end = current_start + roped_query.shape[1]
+            sink_tokens = self.sink_size * frame_seqlen
+
+            kv_cache_size = kv_cache["k"].shape[1]
+            num_new_tokens = roped_query.shape[1]
+
+            if (current_end > kv_cache["global_end_index"].item()) and (
+                num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size
+            ):
+                num_evicted_tokens = (
+                    num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
+                )
+                num_rolled_tokens = (
+                    kv_cache["local_end_index"].item()
+                    - num_evicted_tokens
+                    - sink_tokens
+                )
+                kv_cache["k"][:, sink_tokens : sink_tokens + num_rolled_tokens] = (
+                    kv_cache["k"][
+                        :,
+                        sink_tokens + num_evicted_tokens : sink_tokens
+                        + num_evicted_tokens
+                        + num_rolled_tokens,
+                    ].clone()
+                )
+                kv_cache["v"][:, sink_tokens : sink_tokens + num_rolled_tokens] = (
+                    kv_cache["v"][
+                        :,
+                        sink_tokens + num_evicted_tokens : sink_tokens
+                        + num_evicted_tokens
+                        + num_rolled_tokens,
+                    ].clone()
+                )
+                # Insert the new keys/values at the end
+                local_end_index = (
+                    kv_cache["local_end_index"].item()
+                    + current_end
+                    - kv_cache["global_end_index"].item()
+                    - num_evicted_tokens
+                )
+                local_start_index = local_end_index - num_new_tokens
+                kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+                kv_cache["v"][:, local_start_index:local_end_index] = v
+            else:
+                # Assign new keys/values directly up to current_end
+                local_end_index = (
+                    kv_cache["local_end_index"].item()
+                    + current_end
+                    - kv_cache["global_end_index"].item()
+                )
+                local_start_index = local_end_index - num_new_tokens
+
+                kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+                kv_cache["v"][:, local_start_index:local_end_index] = v
+            x = attention(
+                roped_query,
+                kv_cache["k"][
+                    :,
+                    max(0, local_end_index - self.max_attention_size) : local_end_index,
+                ],
+                kv_cache["v"][
+                    :,
+                    max(0, local_end_index - self.max_attention_size) : local_end_index,
+                ],
+            )
+            kv_cache["global_end_index"].fill_(current_end)
+            kv_cache["local_end_index"].fill_(local_end_index)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class MatrixGameWanCausalAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        cross_attn_type,
+        dim,
+        ffn_dim,
+        num_heads,
+        local_attn_size=-1,
+        sink_size=0,
+        qk_norm=True,
+        cross_attn_norm=False,
+        action_config={},
+        block_idx=0,
+        eps=1e-6,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.local_attn_size = local_attn_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+        if len(action_config) != 0 and block_idx in action_config["blocks"]:
+            self.action_model = ActionModule(
+                **action_config, local_attn_size=self.local_attn_size
+            )
+        else:
+            self.action_model = None
+        # layers
+        self.norm1 = MatrixGameWanLayerNorm(dim, eps)
+        self.self_attn = MatrixGameWanCausalSelfAttention(
+            dim, num_heads, local_attn_size, sink_size, qk_norm, eps
+        )
+        self.norm3 = (
+            MatrixGameWanLayerNorm(dim, eps, elementwise_affine=True)
+            if cross_attn_norm
+            else nn.Identity()
+        )
+        self.cross_attn = MatrixGameWan_CROSSATTENTION_CLASSES[cross_attn_type](
+            dim, num_heads, (-1, -1), qk_norm, eps
+        )
+        self.norm2 = MatrixGameWanLayerNorm(dim, eps)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, ffn_dim),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(ffn_dim, dim),
+        )
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        x,
+        e,
+        seq_lens,
+        grid_sizes,
+        freqs,
+        context,
+        block_mask,
+        block_mask_mouse,
+        block_mask_keyboard,
+        num_frame_per_block=3,
+        use_rope_keyboard=False,
+        mouse_cond=None,
+        keyboard_cond=None,
+        kv_cache=None,
+        kv_cache_mouse=None,
+        kv_cache_keyboard=None,
+        crossattn_cache=None,
+        current_start=0,
+        cache_start=None,
+        context_lens=None,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+            e(Tensor): Shape [B, F, 6, C]
+            seq_lens(Tensor): Shape [B], length of each sequence in batch
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        assert e.ndim == 4
+        num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
+
+        e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2)
+
+        y = self.self_attn(
+            (
+                self.norm1(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen))
+                * (1 + e[1])
+                + e[0]
+            ).flatten(1, 2),
+            seq_lens,
+            grid_sizes,
+            freqs,
+            block_mask,
+            kv_cache,
+            current_start,
+            cache_start,
+        )
+
+        x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[2]).flatten(
+            1, 2
+        )
+
+        # cross-attention & ffn function
+        def cross_attn_ffn(
+            x,
+            context,
+            e,
+            mouse_cond,
+            keyboard_cond,
+            block_mask_mouse,
+            block_mask_keyboard,
+            kv_cache_mouse=None,
+            kv_cache_keyboard=None,
+            crossattn_cache=None,
+            start_frame=0,
+            use_rope_keyboard=False,
+            num_frame_per_block=3,
+        ):
+            x = x + self.cross_attn(
+                self.norm3(x.to(context.dtype)),
+                context,
+                crossattn_cache=crossattn_cache,
+            )
+            if self.action_model is not None:
+                assert mouse_cond is not None or keyboard_cond is not None
+                x = self.action_model(
+                    x.to(context.dtype),
+                    grid_sizes[0],
+                    grid_sizes[1],
+                    grid_sizes[2],
+                    mouse_cond,
+                    keyboard_cond,
+                    block_mask_mouse,
+                    block_mask_keyboard,
+                    is_causal=True,
+                    kv_cache_mouse=kv_cache_mouse,
+                    kv_cache_keyboard=kv_cache_keyboard,
+                    start_frame=start_frame,
+                    use_rope_keyboard=use_rope_keyboard,
+                    num_frame_per_block=num_frame_per_block,
+                )
+
+            y = self.ffn(
+                (
+                    self.norm2(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen))
+                    * (1 + e[4])
+                    + e[3]
+                ).flatten(1, 2)
+            )
+
+            x = x + (
+                y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[5]
+            ).flatten(1, 2)
+            return x
+
+        assert grid_sizes.ndim == 1
+        x = cross_attn_ffn(
+            x,
+            context,
+            e,
+            mouse_cond,
+            keyboard_cond,
+            block_mask_mouse,
+            block_mask_keyboard,
+            kv_cache_mouse,
+            kv_cache_keyboard,
+            crossattn_cache,
+            start_frame=current_start // math.prod(grid_sizes[1:]).item(),
+            use_rope_keyboard=use_rope_keyboard,
+            num_frame_per_block=num_frame_per_block,
+        )
+        return x
+
+
+class CausalHead(nn.Module):
+    def __init__(self, dim, out_dim, patch_size, eps=1e-6):
+        super().__init__()
+        self.dim = dim
+        self.out_dim = out_dim
+        self.patch_size = patch_size
+        self.eps = eps
+
+        # layers
+        out_dim = math.prod(patch_size) * out_dim
+        self.norm = MatrixGameWanLayerNorm(dim, eps)
+        self.head = nn.Linear(dim, out_dim)
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, e):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            e(Tensor): Shape [B, F, 1, C]
+        """
+
+        num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
+        e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2)
+        x = self.head(
+            self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1])
+            + e[0]
+        )
+        return x
+
+
+class MatrixGameWanCausalModel(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin):
+    r"""
+    MatrixGameWan diffusion backbone supporting both text-to-video and image-to-video.
+    """
+
+    ignore_for_config = ["patch_size", "cross_attn_norm", "qk_norm", "text_dim"]
+    _no_split_modules = ["MatrixGameWanAttentionBlock"]
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        model_type="t2v",
+        patch_size=(1, 2, 2),
+        text_len=512,
+        in_dim=36,
+        dim=1536,
+        ffn_dim=8960,
+        freq_dim=256,
+        text_dim=4096,
+        out_dim=16,
+        num_heads=12,
+        num_layers=30,
+        local_attn_size=-1,
+        sink_size=0,
+        qk_norm=True,
+        cross_attn_norm=True,
+        action_config={},
+        eps=1e-6,
+    ):
+        r"""
+        Initialize the diffusion model backbone.
+
+        Args:
+            model_type (`str`, *optional*, defaults to 't2v'):
+                Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+            patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+                3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+            text_len (`int`, *optional*, defaults to 512):
+                Fixed length for text embeddings
+            in_dim (`int`, *optional*, defaults to 16):
+                Input video channels (C_in)
+            dim (`int`, *optional*, defaults to 2048):
+                Hidden dimension of the transformer
+            ffn_dim (`int`, *optional*, defaults to 8192):
+                Intermediate dimension in feed-forward network
+            freq_dim (`int`, *optional*, defaults to 256):
+                Dimension for sinusoidal time embeddings
+            text_dim (`int`, *optional*, defaults to 4096):
+                Input dimension for text embeddings
+            out_dim (`int`, *optional*, defaults to 16):
+                Output video channels (C_out)
+            num_heads (`int`, *optional*, defaults to 16):
+                Number of attention heads
+            num_layers (`int`, *optional*, defaults to 32):
+                Number of transformer blocks
+            local_attn_size (`int`, *optional*, defaults to -1):
+                Window size for temporal local attention (-1 indicates global attention)
+            sink_size (`int`, *optional*, defaults to 0):
+                Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache
+            qk_norm (`bool`, *optional*, defaults to True):
+                Enable query/key normalization
+            cross_attn_norm (`bool`, *optional*, defaults to False):
+                Enable cross-attention normalization
+            eps (`float`, *optional*, defaults to 1e-6):
+                Epsilon value for normalization layers
+        """
+
+        super().__init__()
+
+        assert model_type in ["i2v"]
+        self.model_type = model_type
+        self.use_action_module = len(action_config) > 0
+        self.patch_size = patch_size
+        self.text_len = text_len
+        self.in_dim = in_dim
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.freq_dim = freq_dim
+        self.text_dim = text_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.local_attn_size = local_attn_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # embeddings
+        self.patch_embedding = nn.Conv3d(
+            in_dim, dim, kernel_size=patch_size, stride=patch_size
+        )
+
+        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))
+
+        # blocks
+        cross_attn_type = "i2v_cross_attn"
+        self.blocks = nn.ModuleList(
+            [
+                MatrixGameWanCausalAttentionBlock(
+                    cross_attn_type,
+                    dim,
+                    ffn_dim,
+                    num_heads,
+                    local_attn_size,
+                    sink_size,
+                    qk_norm,
+                    cross_attn_norm,
+                    action_config=action_config,
+                    eps=eps,
+                    block_idx=idx,
+                )
+                for idx in range(num_layers)
+            ]
+        )
+
+        # head
+        self.head = CausalHead(dim, out_dim, patch_size, eps)
+
+        # buffers (don't use register_buffer otherwise dtype will be changed in to())
+        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+        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,
+        )
+
+        if model_type == "i2v":
+            self.img_emb = MLPProj(1280, dim)
+
+        self.gradient_checkpointing = False
+
+        self.block_mask = None
+        self.block_mask_keyboard = None
+        self.block_mask_mouse = None
+        self.use_rope_keyboard = True
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        self.gradient_checkpointing = value
+
+    @staticmethod
+    def _prepare_blockwise_causal_attn_mask(
+        device: torch.device | str,
+        num_frames: int = 9,
+        frame_seqlen: int = 880,
+        num_frame_per_block=1,
+        local_attn_size=-1,
+    ) -> BlockMask:
+        """
+        we will divide the token sequence into the following format
+        [1 latent frame] [1 latent frame] ... [1 latent frame]
+        We use flexattention to construct the attention mask
+        """
+        total_length = num_frames * frame_seqlen
+
+        # we do right padding to get to a multiple of 128
+        padded_length = math.ceil(total_length / 128) * 128 - total_length
+
+        ends = torch.zeros(
+            total_length + padded_length, device=device, dtype=torch.long
+        )
+
+        # Block-wise causal mask will attend to all elements that are before the end of the current chunk
+        frame_indices = torch.arange(
+            start=0,
+            end=total_length,
+            step=frame_seqlen * num_frame_per_block,
+            device=device,
+        )
+
+        for tmp in frame_indices:
+            ends[tmp : tmp + frame_seqlen * num_frame_per_block] = (
+                tmp + frame_seqlen * num_frame_per_block
+            )
+
+        def attention_mask(b, h, q_idx, kv_idx):
+            if local_attn_size == -1:
+                return (kv_idx < ends[q_idx]) | (q_idx == kv_idx)
+            else:
+                return (
+                    (kv_idx < ends[q_idx])
+                    & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))
+                ) | (q_idx == kv_idx)
+            # return ((kv_idx < total_length) & (q_idx < total_length))  | (q_idx == kv_idx) # bidirectional mask
+
+        block_mask = create_block_mask(
+            attention_mask,
+            B=None,
+            H=None,
+            Q_LEN=total_length + padded_length,
+            KV_LEN=total_length + padded_length,
+            _compile=False,
+            device=device,
+        )
+
+        import torch.distributed as dist
+
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(
+                f" cache a block wise causal mask with block size of {num_frame_per_block} frames"
+            )
+
+        return block_mask
+
+    @staticmethod
+    def _prepare_blockwise_causal_attn_mask_keyboard(
+        device: torch.device | str,
+        num_frames: int = 9,
+        frame_seqlen: int = 880,
+        num_frame_per_block=1,
+        local_attn_size=-1,
+    ) -> BlockMask:
+        """
+        we will divide the token sequence into the following format
+        [1 latent frame] [1 latent frame] ... [1 latent frame]
+        We use flexattention to construct the attention mask
+        """
+        total_length2 = num_frames * frame_seqlen
+
+        # we do right padding to get to a multiple of 128
+        padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2
+        padded_length_kv2 = math.ceil(num_frames / 32) * 32 - num_frames
+        ends2 = torch.zeros(
+            total_length2 + padded_length2, device=device, dtype=torch.long
+        )
+
+        # Block-wise causal mask will attend to all elements that are before the end of the current chunk
+        frame_indices2 = torch.arange(
+            start=0,
+            end=total_length2,
+            step=frame_seqlen * num_frame_per_block,
+            device=device,
+        )
+        cnt = num_frame_per_block
+        for tmp in frame_indices2:
+            ends2[tmp : tmp + frame_seqlen * num_frame_per_block] = cnt
+            cnt += num_frame_per_block
+
+        def attention_mask2(b, h, q_idx, kv_idx):
+            if local_attn_size == -1:
+                return (kv_idx < ends2[q_idx]) | (q_idx == kv_idx)
+            else:
+                return (
+                    (kv_idx < ends2[q_idx])
+                    & (kv_idx >= (ends2[q_idx] - local_attn_size))
+                ) | (q_idx == kv_idx)
+            # return ((kv_idx < total_length) & (q_idx < total_length))  | (q_idx == kv_idx) # bidirectional mask
+
+        block_mask2 = create_block_mask(
+            attention_mask2,
+            B=None,
+            H=None,
+            Q_LEN=total_length2 + padded_length2,
+            KV_LEN=num_frames + padded_length_kv2,
+            _compile=False,
+            device=device,
+        )
+
+        import torch.distributed as dist
+
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(
+                f" cache a block wise causal mask with block size of {num_frame_per_block} frames"
+            )
+
+        return block_mask2
+
+    @staticmethod
+    def _prepare_blockwise_causal_attn_mask_action(
+        device: torch.device | str,
+        num_frames: int = 9,
+        frame_seqlen: int = 1,
+        num_frame_per_block=1,
+        local_attn_size=-1,
+    ) -> BlockMask:
+        """
+        we will divide the token sequence into the following format
+        [1 latent frame] [1 latent frame] ... [1 latent frame]
+        We use flexattention to construct the attention mask
+        """
+        total_length2 = num_frames * frame_seqlen
+
+        # we do right padding to get to a multiple of 128
+        padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2
+        padded_length_kv2 = math.ceil(num_frames / 32) * 32 - num_frames
+        ends2 = torch.zeros(
+            total_length2 + padded_length2, device=device, dtype=torch.long
+        )
+
+        # Block-wise causal mask will attend to all elements that are before the end of the current chunk
+        frame_indices2 = torch.arange(
+            start=0,
+            end=total_length2,
+            step=frame_seqlen * num_frame_per_block,
+            device=device,
+        )
+        cnt = num_frame_per_block
+        for tmp in frame_indices2:
+            ends2[tmp : tmp + frame_seqlen * num_frame_per_block] = cnt
+            cnt += num_frame_per_block
+
+        def attention_mask2(b, h, q_idx, kv_idx):
+            if local_attn_size == -1:
+                return (kv_idx < ends2[q_idx]) | (q_idx == kv_idx)
+            else:
+                return (
+                    (kv_idx < ends2[q_idx])
+                    & (kv_idx >= (ends2[q_idx] - local_attn_size))
+                ) | (q_idx == kv_idx)
+            # return ((kv_idx < total_length) & (q_idx < total_length))  | (q_idx == kv_idx) # bidirectional mask
+
+        block_mask2 = create_block_mask(
+            attention_mask2,
+            B=None,
+            H=None,
+            Q_LEN=total_length2 + padded_length2,
+            KV_LEN=num_frames + padded_length_kv2,
+            _compile=False,
+            device=device,
+        )
+
+        import torch.distributed as dist
+
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(
+                f" cache a block wise causal mask with block size of {num_frame_per_block} frames"
+            )
+
+        return block_mask2
+
+    def _forward_inference(
+        self,
+        x,
+        t,
+        visual_context,
+        cond_concat,
+        mouse_cond=None,
+        keyboard_cond=None,
+        kv_cache: dict = None,
+        kv_cache_mouse=None,
+        kv_cache_keyboard=None,
+        crossattn_cache: dict = None,
+        current_start: int = 0,
+        cache_start: int = 0,
+        num_frames_per_block=3,
+    ):
+        r"""
+        Run the diffusion model with kv caching.
+        See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details.
+        This function will be run for num_frame times.
+        Process the latent frames one by one (1560 tokens each)
+
+        Args:
+            x (List[Tensor]):
+                List of input video tensors, each with shape [C_in, F, H, W]
+            t (Tensor):
+                Diffusion timesteps tensor of shape [B]
+            context (List[Tensor]):
+                List of text embeddings each with shape [L, C]
+            seq_len (`int`):
+                Maximum sequence length for positional encoding
+            clip_fea (Tensor, *optional*):
+                CLIP image features for image-to-video mode
+            y (List[Tensor], *optional*):
+                Conditional video inputs for image-to-video mode, same shape as x
+
+        Returns:
+            List[Tensor]:
+                List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+        """
+
+        if mouse_cond is not None or keyboard_cond is not None:
+            assert self.use_action_module == True
+        # params
+        device = self.patch_embedding.weight.device
+        if self.freqs.device != device:
+            self.freqs = self.freqs.to(device)
+
+        x = torch.cat([x, cond_concat], dim=1)  # B C' F H W
+
+        # embeddings
+        x = self.patch_embedding(x)
+        grid_sizes = torch.tensor(x.shape[2:], dtype=torch.long)
+
+        x = x.flatten(2).transpose(1, 2)  # B FHW C'
+        seq_lens = torch.tensor([u.size(0) for u in x], dtype=torch.long)
+        assert seq_lens[0] <= 15 * 1 * 880
+
+        e = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x)
+        )
+        e0 = (
+            self.time_projection(e)
+            .unflatten(1, (6, self.dim))
+            .unflatten(dim=0, sizes=t.shape)
+        )
+        # context
+        context_lens = None
+        context = self.img_emb(visual_context)
+        # arguments
+        kwargs = dict(
+            e=e0,
+            seq_lens=seq_lens,
+            grid_sizes=grid_sizes,
+            freqs=self.freqs,
+            context=context,
+            mouse_cond=mouse_cond,
+            context_lens=context_lens,
+            keyboard_cond=keyboard_cond,
+            block_mask=self.block_mask,
+            block_mask_mouse=self.block_mask_mouse,
+            block_mask_keyboard=self.block_mask_keyboard,
+            use_rope_keyboard=self.use_rope_keyboard,
+            num_frame_per_block=num_frames_per_block,
+        )
+
+        def create_custom_forward(module):
+            def custom_forward(*inputs, **kwargs):
+                return module(*inputs, **kwargs)
+
+            return custom_forward
+
+        for block_index, block in enumerate(self.blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                kwargs.update(
+                    {
+                        "kv_cache": kv_cache[block_index],
+                        "kv_cache_mouse": kv_cache_mouse[block_index],
+                        "kv_cache_keyboard": kv_cache_keyboard[block_index],
+                        "current_start": current_start,
+                        "cache_start": cache_start,
+                    }
+                )
+                x = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    x,
+                    **kwargs,
+                    use_reentrant=False,
+                )
+            else:
+                kwargs.update(
+                    {
+                        "kv_cache": kv_cache[block_index],
+                        "kv_cache_mouse": kv_cache_mouse[block_index],
+                        "kv_cache_keyboard": kv_cache_keyboard[block_index],
+                        "crossattn_cache": crossattn_cache[block_index],
+                        "current_start": current_start,
+                        "cache_start": cache_start,
+                    }
+                )
+                x = block(x, **kwargs)
+
+        # head
+        x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2))
+        # unpatchify
+        x = self.unpatchify(x, grid_sizes)
+        return x
+
+    def forward(self, *args, **kwargs):
+        return self._forward_inference(*args, **kwargs)
+
+    def unpatchify(self, x, grid_sizes):
+        r"""
+        Reconstruct video tensors from patch embeddings.
+
+        Args:
+            x (List[Tensor]):
+                List of patchified features, each with shape [L, C_out * prod(patch_size)]
+            grid_sizes (Tensor):
+                Original spatial-temporal grid dimensions before patching,
+                    shape [3] (3 dimensions correspond to F_patches, H_patches, W_patches)
+
+        Returns:
+            List[Tensor]:
+                Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
+        """
+
+        c = self.out_dim
+        bs = x.shape[0]
+        x = x.view(bs, *grid_sizes, *self.patch_size, c)
+        x = torch.einsum("bfhwpqrc->bcfphqwr", x)
+        x = x.reshape(bs, c, *[i * j for i, j in zip(grid_sizes, self.patch_size)])
+        return x
+

transformer/config.json

@@ -0,0 +1,68 @@
+{
+  "_class_name": "MatrixGameWanCausalModel",
+  "_diffusers_version": "0.35.1",
+  "auto_map": {
+    "AutoModel": "causal_model.MatrixGameWanCausalModel"
+  },
+  "action_config": {
+    "blocks": [
+      0,
+      1,
+      2,
+      3,
+      4,
+      5,
+      6,
+      7,
+      8,
+      9,
+      10,
+      11,
+      12,
+      13,
+      14
+    ],
+    "enable_keyboard": true,
+    "enable_mouse": true,
+    "heads_num": 16,
+    "hidden_size": 128,
+    "img_hidden_size": 1536,
+    "keyboard_dim_in": 4,
+    "keyboard_hidden_dim": 1024,
+    "mouse_dim_in": 2,
+    "mouse_hidden_dim": 1024,
+    "mouse_qk_dim_list": [
+      8,
+      28,
+      28
+    ],
+    "patch_size": [
+      1,
+      2,
+      2
+    ],
+    "qk_norm": true,
+    "qkv_bias": false,
+    "rope_dim_list": [
+      8,
+      28,
+      28
+    ],
+    "rope_theta": 256,
+    "vae_time_compression_ratio": 4,
+    "windows_size": 3
+  },
+  "dim": 1536,
+  "eps": 1e-06,
+  "ffn_dim": 8960,
+  "freq_dim": 256,
+  "in_dim": 36,
+  "inject_sample_info": false,
+  "local_attn_size": 6,
+  "model_type": "i2v",
+  "num_heads": 12,
+  "num_layers": 30,
+  "out_dim": 16,
+  "sink_size": 0,
+  "text_len": 512
+}

transformer/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4d7930e670e1475abdeed25fa0f1c34c47c8be51d9e7b8f637e2a4b720548a4
+size 3238601944

transformer/model.py

@@ -0,0 +1,781 @@
+# Copyright 2024-2025 The Alibaba MatrixGameWan Team Authors. All rights reserved.
+import math
+import numpy as np
+import torch
+import torch.amp as amp
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.models.modeling_utils import ModelMixin
+from einops import repeat, rearrange
+from .action_module import ActionModule
+from .attention import flash_attention
+
+DISABLE_COMPILE = False  # get os env
+__all__ = ["MatrixGameWanModel"]
+
+
+def sinusoidal_embedding_1d(dim, position):
+    # preprocess
+    assert dim % 2 == 0
+    half = dim // 2
+    position = position.type(torch.float64)
+
+    # calculation
+    sinusoid = torch.outer(
+        position, torch.pow(10000, -torch.arange(half).to(position).div(half))
+    )
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x
+
+
+# @amp.autocast(enabled=False)
+def rope_params(max_seq_len, dim, theta=10000):
+    assert dim % 2 == 0
+    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
+
+
+# @amp.autocast(enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+    n, c = x.size(2), x.size(3) // 2
+
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+    # loop over samples
+    output = []
+    # print(grid_sizes.shape, len(grid_sizes.tolist()), grid_sizes.tolist()[0])
+    f, h, w = grid_sizes.tolist()
+    for i in range(len(x)):
+        seq_len = f * h * w
+
+        # precompute multipliers
+        x_i = torch.view_as_complex(
+            x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)
+        )
+        freqs_i = torch.cat(
+            [
+                freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+                freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+                freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1),
+            ],
+            dim=-1,
+        ).reshape(seq_len, 1, -1)
+
+        # apply rotary embedding
+        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+        x_i = torch.cat([x_i, x[i, seq_len:]])
+
+        # append to collection
+        output.append(x_i)
+    return torch.stack(output).type_as(x)
+
+
+class MatrixGameWanRMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return self._norm(x.float()).type_as(x) * self.weight
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+
+class MatrixGameWanLayerNorm(nn.LayerNorm):
+    def __init__(self, dim, eps=1e-6, elementwise_affine=False):
+        super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return super().forward(x).type_as(x)
+
+
+class MatrixGameWanSelfAttention(nn.Module):
+    def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # layers
+        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 = MatrixGameWanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = MatrixGameWanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(self, x, seq_lens, grid_sizes, freqs):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, num_heads, C / num_heads]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+
+        # query, key, value function
+        def qkv_fn(x):
+            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)
+            return q, k, v
+
+        q, k, v = qkv_fn(x)
+        # print(k.shape, seq_lens)
+        x = flash_attention(
+            q=rope_apply(q, grid_sizes, freqs),
+            k=rope_apply(k, grid_sizes, freqs),
+            v=v,
+            k_lens=seq_lens,
+            window_size=self.window_size,
+        )
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+# class MatrixGameWanT2VCrossAttention(MatrixGameWanSelfAttention):
+
+#     def forward(self, x, context, context_lens, crossattn_cache=None):
+#         r"""
+#         Args:
+#             x(Tensor): Shape [B, L1, C]
+#             context(Tensor): Shape [B, L2, C]
+#             context_lens(Tensor): Shape [B]
+#             crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding.
+#         """
+#         b, n, d = x.size(0), self.num_heads, self.head_dim
+
+#         # compute query, key, value
+#         q = self.norm_q(self.q(x)).view(b, -1, n, d)
+
+#         if crossattn_cache is not None:
+#             if not crossattn_cache["is_init"]:
+#                 crossattn_cache["is_init"] = True
+#                 k = self.norm_k(self.k(context)).view(b, -1, n, d)
+#                 v = self.v(context).view(b, -1, n, d)
+#                 crossattn_cache["k"] = k
+#                 crossattn_cache["v"] = v
+#             else:
+#                 k = crossattn_cache["k"]
+#                 v = crossattn_cache["v"]
+#         else:
+#             k = self.norm_k(self.k(context)).view(b, -1, n, d)
+#             v = self.v(context).view(b, -1, n, d)
+
+#         # compute attention
+#         x = flash_attention(q, k, v, k_lens=context_lens)
+
+#         # output
+#         x = x.flatten(2)
+#         x = self.o(x)
+#         return x
+
+
+# class MatrixGameWanGanCrossAttention(MatrixGameWanSelfAttention):
+
+#     def forward(self, x, context, crossattn_cache=None):
+#         r"""
+#         Args:
+#             x(Tensor): Shape [B, L1, C]
+#             context(Tensor): Shape [B, L2, C]
+#             context_lens(Tensor): Shape [B]
+#             crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding.
+#         """
+#         b, n, d = x.size(0), self.num_heads, self.head_dim
+
+#         # compute query, key, value
+#         qq = self.norm_q(self.q(context)).view(b, 1, -1, d)
+
+#         kk = self.norm_k(self.k(x)).view(b, -1, n, d)
+#         vv = self.v(x).view(b, -1, n, d)
+
+#         # compute attention
+#         x = flash_attention(qq, kk, vv)
+
+#         # output
+#         x = x.flatten(2)
+#         x = self.o(x)
+#         return x
+
+
+class MatrixGameWanI2VCrossAttention(MatrixGameWanSelfAttention):
+    def forward(self, x, context, crossattn_cache=None):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+        if crossattn_cache is not None:
+            if not crossattn_cache["is_init"]:
+                crossattn_cache["is_init"] = True
+                k = self.norm_k(self.k(context)).view(b, -1, n, d)
+                v = self.v(context).view(b, -1, n, d)
+                crossattn_cache["k"] = k
+                crossattn_cache["v"] = v
+            else:
+                k = crossattn_cache["k"]
+                v = crossattn_cache["v"]
+        else:
+            k = self.norm_k(self.k(context)).view(b, -1, n, d)
+            v = self.v(context).view(b, -1, n, d)
+        # compute attention
+        x = flash_attention(q, k, v, k_lens=None)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+MatrixGameWan_CROSSATTENTION_CLASSES = {
+    "i2v_cross_attn": MatrixGameWanI2VCrossAttention,
+}
+
+
+def mul_add(x, y, z):
+    return x.float() + y.float() * z.float()
+
+
+def mul_add_add(x, y, z):
+    return x.float() * (1 + y) + z
+
+
+class MatrixGameWanAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        cross_attn_type,
+        dim,
+        ffn_dim,
+        num_heads,
+        window_size=(-1, -1),
+        qk_norm=True,
+        cross_attn_norm=False,
+        action_config={},
+        eps=1e-6,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+        if len(action_config) != 0:
+            self.action_model = ActionModule(**action_config)
+        else:
+            self.action_model = None
+        # layers
+        self.norm1 = MatrixGameWanLayerNorm(dim, eps)
+        self.self_attn = MatrixGameWanSelfAttention(dim, num_heads, window_size, qk_norm, eps)
+        self.norm3 = (
+            MatrixGameWanLayerNorm(dim, eps, elementwise_affine=True)
+            if cross_attn_norm
+            else nn.Identity()
+        )
+        self.cross_attn = MatrixGameWan_CROSSATTENTION_CLASSES[cross_attn_type](
+            dim, num_heads, (-1, -1), qk_norm, eps
+        )
+        self.norm2 = MatrixGameWanLayerNorm(dim, eps)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, ffn_dim),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(ffn_dim, dim),
+        )
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        x,
+        e,
+        seq_lens,
+        grid_sizes,
+        freqs,
+        context,
+        mouse_cond=None,
+        keyboard_cond=None,
+        # context_lens,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+            e(Tensor): Shape [B, 6, C]
+            seq_lens(Tensor): Shape [B], length of each sequence in batch
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        # assert e.dtype == torch.float32
+        if e.dim() == 3:
+            modulation = self.modulation
+            # with amp.autocast(dtype=torch.float32):
+            e = (self.modulation + e).chunk(6, dim=1)
+        elif e.dim() == 4:
+            modulation = self.modulation.unsqueeze(2)  # 1, 6, 1, dim
+            # with amp.autocast("cuda", dtype=torch.float32):
+            e = (modulation + e).chunk(6, dim=1)
+            e = [ei.squeeze(1) for ei in e]
+        # assert e[0].dtype == torch.float32
+
+        # self-attention
+        y = self.self_attn(
+            self.norm1(x) * (1 + e[1]) + e[0], seq_lens, grid_sizes, freqs
+        )
+        # with amp.autocast(dtype=torch.float32):
+        x = x + y * e[2]
+
+        # cross-attention & ffn function
+        def cross_attn_ffn(x, context, e, mouse_cond, keyboard_cond):
+            dtype = context.dtype
+            x = x + self.cross_attn(self.norm3(x.to(dtype)), context)
+            if self.action_model is not None:
+                assert mouse_cond is not None or keyboard_cond is not None
+                x = self.action_model(
+                    x.to(dtype),
+                    grid_sizes[0],
+                    grid_sizes[1],
+                    grid_sizes[2],
+                    mouse_cond,
+                    keyboard_cond,
+                )
+            y = self.ffn(self.norm2(x) * (1 + e[4]) + e[3])
+            # with amp.autocast(dtype=torch.float32):
+            x = x + y * e[5]
+            return x
+
+        x = cross_attn_ffn(x, context, e, mouse_cond, keyboard_cond)
+        return x
+
+
+class Head(nn.Module):
+    def __init__(self, dim, out_dim, patch_size, eps=1e-6):
+        super().__init__()
+        self.dim = dim
+        self.out_dim = out_dim
+        self.patch_size = patch_size
+        self.eps = eps
+
+        # layers
+        out_dim = math.prod(patch_size) * out_dim
+        self.norm = MatrixGameWanLayerNorm(dim, eps)
+        self.head = nn.Linear(dim, out_dim)
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, e):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            e(Tensor): Shape [B, C]
+        """
+        # assert e.dtype == torch.float32
+        # with amp.autocast(dtype=torch.float32):
+        if e.dim() == 2:
+            modulation = self.modulation  # 1, 2, dim
+            e = (modulation + e.unsqueeze(1)).chunk(2, dim=1)
+        elif e.dim() == 3:
+            modulation = self.modulation.unsqueeze(2)  # 1, 2, seq, dim
+            e = (modulation + e.unsqueeze(1)).chunk(2, dim=1)
+            e = [ei.squeeze(1) for ei in e]
+        x = self.head(self.norm(x) * (1 + e[1]) + e[0])
+        return x
+
+
+class MLPProj(torch.nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+
+        self.proj = torch.nn.Sequential(
+            torch.nn.LayerNorm(in_dim),
+            torch.nn.Linear(in_dim, in_dim),
+            torch.nn.GELU(),
+            torch.nn.Linear(in_dim, out_dim),
+            torch.nn.LayerNorm(out_dim),
+        )
+
+    def forward(self, image_embeds):
+        clip_extra_context_tokens = self.proj(image_embeds)
+        return clip_extra_context_tokens
+
+
+# class RegisterTokens(nn.Module):
+#     def __init__(self, num_registers: int, dim: int):
+#         super().__init__()
+#         self.register_tokens = nn.Parameter(torch.randn(num_registers, dim) * 0.02)
+#         self.rms_norm = MatrixGameWanRMSNorm(dim, eps=1e-6)
+
+#     def forward(self):
+#         return self.rms_norm(self.register_tokens)
+
+#     def reset_parameters(self):
+#         nn.init.normal_(self.register_tokens, std=0.02)
+
+
+class MatrixGameWanModel(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin):
+    r"""
+    MatrixGameWan diffusion backbone supporting both text-to-video and image-to-video.
+    """
+
+    ignore_for_config = [
+        "patch_size",
+        "cross_attn_norm",
+        "qk_norm",
+        "text_dim",
+        "window_size",
+    ]
+    _no_split_modules = ["MatrixGameWanAttentionBlock"]
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        model_type="i2v",
+        patch_size=(1, 2, 2),
+        text_len=512,
+        in_dim=36,
+        dim=1536,
+        ffn_dim=8960,
+        freq_dim=256,
+        text_dim=4096,
+        out_dim=16,
+        num_heads=12,
+        num_layers=30,
+        window_size=(-1, -1),
+        qk_norm=True,
+        cross_attn_norm=True,
+        inject_sample_info=False,
+        action_config={},
+        eps=1e-6,
+    ):
+        r"""
+        Initialize the diffusion model backbone.
+
+        Args:
+            model_type (`str`, *optional*, defaults to 't2v'):
+                Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+            patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+                3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+            text_len (`int`, *optional*, defaults to 512):
+                Fixed length for text embeddings
+            in_dim (`int`, *optional*, defaults to 16):
+                Input video channels (C_in)
+            dim (`int`, *optional*, defaults to 2048):
+                Hidden dimension of the transformer
+            ffn_dim (`int`, *optional*, defaults to 8192):
+                Intermediate dimension in feed-forward network
+            freq_dim (`int`, *optional*, defaults to 256):
+                Dimension for sinusoidal time embeddings
+            text_dim (`int`, *optional*, defaults to 4096):
+                Input dimension for text embeddings
+            out_dim (`int`, *optional*, defaults to 16):
+                Output video channels (C_out)
+            num_heads (`int`, *optional*, defaults to 16):
+                Number of attention heads
+            num_layers (`int`, *optional*, defaults to 32):
+                Number of transformer blocks
+            window_size (`tuple`, *optional*, defaults to (-1, -1)):
+                Window size for local attention (-1 indicates global attention)
+            qk_norm (`bool`, *optional*, defaults to True):
+                Enable query/key normalization
+            cross_attn_norm (`bool`, *optional*, defaults to False):
+                Enable cross-attention normalization
+            eps (`float`, *optional*, defaults to 1e-6):
+                Epsilon value for normalization layers
+        """
+
+        super().__init__()
+
+        assert model_type in ["i2v"]
+        self.model_type = model_type
+        self.use_action_module = len(action_config) > 0
+        assert self.use_action_module == True
+        self.patch_size = patch_size
+        self.text_len = text_len
+        self.in_dim = in_dim
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.freq_dim = freq_dim
+        self.text_dim = text_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+        self.local_attn_size = -1
+
+        # 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))
+
+        # blocks
+        cross_attn_type = "i2v_cross_attn"
+        self.blocks = nn.ModuleList(
+            [
+                MatrixGameWanAttentionBlock(
+                    cross_attn_type,
+                    dim,
+                    ffn_dim,
+                    num_heads,
+                    window_size,
+                    qk_norm,
+                    cross_attn_norm,
+                    eps=eps,
+                    action_config=action_config,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        # head
+        self.head = Head(dim, out_dim, patch_size, eps)
+
+        # buffers (don't use register_buffer otherwise dtype will be changed in to())
+        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+        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,
+        )
+
+        if model_type == "i2v":
+            self.img_emb = MLPProj(1280, dim)
+
+        # initialize weights
+        self.init_weights()
+
+        self.gradient_checkpointing = False
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        self.gradient_checkpointing = value
+
+    def forward(self, *args, **kwargs):
+        # if kwargs.get('classify_mode', False) is True:
+        # kwargs.pop('classify_mode')
+        # return self._forward_classify(*args, **kwargs)
+        # else:
+        return self._forward(*args, **kwargs)
+
+    def _forward(
+        self,
+        x,
+        t,
+        visual_context,
+        cond_concat,
+        mouse_cond=None,
+        keyboard_cond=None,
+        fps=None,
+        # seq_len,
+        # classify_mode=False,
+        # concat_time_embeddings=False,
+        # register_tokens=None,
+        # cls_pred_branch=None,
+        # gan_ca_blocks=None,
+        # clip_fea=None,
+        # y=None,
+    ):
+        r"""
+        Forward pass through the diffusion model
+
+        Args:
+            x (List[Tensor]):
+                List of input video tensors, each with shape [C_in, F, H, W]
+            t (Tensor):
+                Diffusion timesteps tensor of shape [B]
+            context (List[Tensor]):
+                List of text embeddings each with shape [L, C]
+            seq_len (`int`):
+                Maximum sequence length for positional encoding
+            clip_fea (Tensor, *optional*):
+                CLIP image features for image-to-video mode
+            y (List[Tensor], *optional*):
+                Conditional video inputs for image-to-video mode, same shape as x
+
+        Returns:
+            List[Tensor]:
+                List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+        """
+        # params
+        if mouse_cond is not None or keyboard_cond is not None:
+            assert self.use_action_module == True
+        device = self.patch_embedding.weight.device
+        if self.freqs.device != device:
+            self.freqs = self.freqs.to(device)
+
+        x = torch.cat([x, cond_concat], dim=1)
+        # embeddings
+        x = self.patch_embedding(x)
+        grid_sizes = torch.tensor(x.shape[2:], dtype=torch.long)
+        x = x.flatten(2).transpose(1, 2)
+        seq_lens = torch.tensor([u.size(0) for u in x], dtype=torch.long)
+        # seq_len = seq_lens.max()
+        # # assert seq_lens.max() <= seq_len
+        # x = torch.cat([
+        #     torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
+        #               dim=1) for u in x
+        # ])
+
+        # time embeddings
+        # with amp.autocast(dtype=torch.float32):
+        # assert t.ndim == 1
+        e = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, t).type_as(x)
+        )  # TODO: check if t ndim == 1
+
+        e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+        # assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+        # context
+        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
+        #     ]))
+
+        # if clip_fea is not None:
+        #     context_clip = self.img_emb(clip_fea)  # bs x 257 x dim
+        context = self.img_emb(visual_context)
+
+        # arguments
+        # kwargs = dict(
+        #     e=e0,
+        #     seq_lens=seq_lens,
+        #     grid_sizes=grid_sizes,
+        #     freqs=self.freqs,
+        #     context=context,
+        #     context_lens=context_lens)
+        kwargs = dict(
+            e=e0,
+            grid_sizes=grid_sizes,
+            seq_lens=seq_lens,
+            freqs=self.freqs,
+            context=context,
+            mouse_cond=mouse_cond,
+            # context_lens=context_lens,
+            keyboard_cond=keyboard_cond,
+        )
+
+        def create_custom_forward(module):
+            def custom_forward(*inputs, **kwargs):
+                return module(*inputs, **kwargs)
+
+            return custom_forward
+
+        for ii, block in enumerate(self.blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                x = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    x,
+                    **kwargs,
+                    use_reentrant=False,
+                )
+            else:
+                x = block(x, **kwargs)
+
+        # head
+        x = self.head(x, e)
+
+        # unpatchify
+        x = self.unpatchify(x, grid_sizes)
+
+        return x.float()
+
+    def unpatchify(self, x, grid_sizes):  # TODO check grid sizes
+        r"""
+        Reconstruct video tensors from patch embeddings.
+
+        Args:
+            x (List[Tensor]):
+                List of patchified features, each with shape [L, C_out * prod(patch_size)]
+            grid_sizes (Tensor):
+                Original spatial-temporal grid dimensions before patching,
+                    shape [3] (3 dimensions correspond to F_patches, H_patches, W_patches)
+
+        Returns:
+            List[Tensor]:
+                Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
+        """
+
+        c = self.out_dim
+        bs = x.shape[0]
+        x = x.view(bs, *grid_sizes, *self.patch_size, c)
+        x = torch.einsum("bfhwpqrc->bcfphqwr", x)
+        x = x.reshape(bs, c, *[i * j for i, j in zip(grid_sizes, self.patch_size)])
+        return x
+
+    def init_weights(self):
+        r"""
+        Initialize model parameters using Xavier initialization.
+        """
+
+        # basic init
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+        # init embeddings
+        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+        for m in self.time_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=0.02)
+
+        # init output layer
+        nn.init.zeros_(self.head.head.weight)
+        if self.use_action_module == True:
+            for m in self.blocks:
+                nn.init.zeros_(m.action_model.proj_mouse.weight)
+                if m.action_model.proj_mouse.bias is not None:
+                    nn.init.zeros_(m.action_model.proj_mouse.bias)
+                nn.init.zeros_(m.action_model.proj_keyboard.weight)
+                if m.action_model.proj_keyboard.bias is not None:
+                    nn.init.zeros_(m.action_model.proj_keyboard.bias)

transformer/posemb_layers.py

@@ -0,0 +1,314 @@
+import torch
+from typing import Union, Tuple, List
+
+
+def _to_tuple(x, dim=2):
+    if isinstance(x, int):
+        return (x,) * dim
+    elif len(x) == dim:
+        return x
+    else:
+        raise ValueError(f"Expected length {dim} or int, but got {x}")
+
+
+def get_meshgrid_nd(start, *args, dim=2):
+    """
+    Get n-D meshgrid with start, stop and num.
+
+    Args:
+        start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop,
+            step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num
+            should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in
+            n-tuples.
+        *args: See above.
+        dim (int): Dimension of the meshgrid. Defaults to 2.
+
+    Returns:
+        grid (np.ndarray): [dim, ...]
+    """
+    if len(args) == 0:
+        # start is grid_size
+        num = _to_tuple(start, dim=dim)
+        start = (0,) * dim
+        stop = num
+    elif len(args) == 1:
+        # start is start, args[0] is stop, step is 1
+        start = _to_tuple(start, dim=dim)
+        stop = _to_tuple(args[0], dim=dim)
+        num = [stop[i] - start[i] for i in range(dim)]
+    elif len(args) == 2:
+        # start is start, args[0] is stop, args[1] is num
+        start = _to_tuple(start, dim=dim)  # Left-Top       eg: 12,0
+        stop = _to_tuple(args[0], dim=dim)  # Right-Bottom   eg: 20,32
+        num = _to_tuple(args[1], dim=dim)  # Target Size    eg: 32,124
+    else:
+        raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}")
+
+    # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False)
+    axis_grid = []
+    for i in range(dim):
+        a, b, n = start[i], stop[i], num[i]
+        g = torch.linspace(a, b, n + 1, dtype=torch.float32, device=torch.cuda.current_device())[:n]
+        axis_grid.append(g)
+    grid = torch.meshgrid(*axis_grid, indexing="ij")  # dim x [W, H, D]
+    grid = torch.stack(grid, dim=0)  # [dim, W, H, D]
+
+    return grid
+
+
+#################################################################################
+#                   Rotary Positional Embedding Functions                       #
+#################################################################################
+# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L80
+
+
+def reshape_for_broadcast(
+    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+    x: torch.Tensor,
+    head_first=False,
+):
+    """
+    Reshape frequency tensor for broadcasting it with another tensor.
+
+    This function reshapes the frequency tensor to have the same shape as the target tensor 'x'
+    for the purpose of broadcasting the frequency tensor during element-wise operations.
+
+    Notes:
+        When using FlashMHAModified, head_first should be False.
+        When using Attention, head_first should be True.
+
+    Args:
+        freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped.
+        x (torch.Tensor): Target tensor for broadcasting compatibility.
+        head_first (bool): head dimension first (except batch dim) or not.
+
+    Returns:
+        torch.Tensor: Reshaped frequency tensor.
+
+    Raises:
+        AssertionError: If the frequency tensor doesn't match the expected shape.
+        AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions.
+    """
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+
+    if isinstance(freqs_cis, tuple):
+        # freqs_cis: (cos, sin) in real space
+        if head_first:
+            assert freqs_cis[0].shape == (
+                x.shape[-2],
+                x.shape[-1],
+            ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}"
+            shape = [
+                d if i == ndim - 2 or i == ndim - 1 else 1
+                for i, d in enumerate(x.shape)
+            ]
+        else:
+            # assert freqs_cis[0].shape == (
+            #     x.shape[1],
+            #     x.shape[-1],
+            # ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}"
+            # shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+            shape = [1, freqs_cis[0].shape[0], 1, freqs_cis[0].shape[1]]
+        return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape)
+    else:
+        # freqs_cis: values in complex space
+        if head_first:
+            assert freqs_cis.shape == (
+                x.shape[-2],
+                x.shape[-1],
+            ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}"
+            shape = [
+                d if i == ndim - 2 or i == ndim - 1 else 1
+                for i, d in enumerate(x.shape)
+            ]
+        else:
+            assert freqs_cis.shape == (
+                x.shape[1],
+                x.shape[-1],
+            ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}"
+            shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+        return freqs_cis.view(*shape)
+
+
+def rotate_half(x):
+    x_real, x_imag = (
+        x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1)
+    )  # [B, S, H, D//2]
+    return torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
+    head_first: bool = False,
+    start_offset: int = 0,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor.
+
+    This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided
+    frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor
+    is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are
+    returned as real tensors.
+
+    Args:
+        xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D]
+        xk (torch.Tensor): Key tensor to apply rotary embeddings.   [B, S, H, D]
+        freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential.
+        head_first (bool): head dimension first (except batch dim) or not.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+
+    """
+    # print(freqs_cis[0].shape, xq.shape, xk.shape)
+    xk_out = None
+    assert isinstance(freqs_cis, tuple)
+    if isinstance(freqs_cis, tuple):
+        cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first)  # [S, D]
+        cos, sin = cos.to(xq.device), sin.to(xq.device)
+        # real * cos - imag * sin
+        # imag * cos + real * sin
+        xq_out = (xq.float() * cos[:, start_offset:start_offset + xq.shape[1], :, :] + rotate_half(xq.float()) * sin[:, start_offset:start_offset + xq.shape[1], :, :]).type_as(xq)
+        xk_out = (xk.float() * cos[:, start_offset:start_offset + xk.shape[1], :, :] + rotate_half(xk.float()) * sin[:, start_offset:start_offset + xk.shape[1], :, :]).type_as(xk)
+    else:
+        # view_as_complex will pack [..., D/2, 2](real) to [..., D/2](complex)
+        xq_ = torch.view_as_complex(
+            xq.float().reshape(*xq.shape[:-1], -1, 2)
+        )  # [B, S, H, D//2]
+        freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(
+            xq.device
+        )  # [S, D//2] --> [1, S, 1, D//2]
+        # (real, imag) * (cos, sin) = (real * cos - imag * sin, imag * cos + real * sin)
+        # view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real)
+        xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq)
+        xk_ = torch.view_as_complex(
+            xk.float().reshape(*xk.shape[:-1], -1, 2)
+        )  # [B, S, H, D//2]
+        xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk)
+
+    return xq_out, xk_out
+
+
+def get_nd_rotary_pos_embed(
+    rope_dim_list,
+    start,
+    *args,
+    theta=10000.0,
+    use_real=False,
+    theta_rescale_factor: Union[float, List[float]] = 1.0,
+    interpolation_factor: Union[float, List[float]] = 1.0,
+):
+    """
+    This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure.
+
+    Args:
+        rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n.
+            sum(rope_dim_list) should equal to head_dim of attention layer.
+        start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start,
+            args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num.
+        *args: See above.
+        theta (float): Scaling factor for frequency computation. Defaults to 10000.0.
+        use_real (bool): If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+            Some libraries such as TensorRT does not support complex64 data type. So it is useful to provide a real
+            part and an imaginary part separately.
+        theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0.
+
+    Returns:
+        pos_embed (torch.Tensor): [HW, D/2]
+    """
+
+    grid = get_meshgrid_nd(
+        start, *args, dim=len(rope_dim_list)
+    )  # [3, W, H, D] / [2, W, H]
+
+    if isinstance(theta_rescale_factor, int) or isinstance(theta_rescale_factor, float):
+        theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list)
+    elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1:
+        theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list)
+    assert len(theta_rescale_factor) == len(
+        rope_dim_list
+    ), "len(theta_rescale_factor) should equal to len(rope_dim_list)"
+
+    if isinstance(interpolation_factor, int) or isinstance(interpolation_factor, float):
+        interpolation_factor = [interpolation_factor] * len(rope_dim_list)
+    elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1:
+        interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list)
+    assert len(interpolation_factor) == len(
+        rope_dim_list
+    ), "len(interpolation_factor) should equal to len(rope_dim_list)"
+
+    # use 1/ndim of dimensions to encode grid_axis
+    embs = []
+    for i in range(len(rope_dim_list)):
+        emb = get_1d_rotary_pos_embed(
+            rope_dim_list[i],
+            grid[i].reshape(-1),
+            theta,
+            use_real=use_real,
+            theta_rescale_factor=theta_rescale_factor[i],
+            interpolation_factor=interpolation_factor[i],
+        )  # 2 x [WHD, rope_dim_list[i]]
+        embs.append(emb)
+
+    if use_real:
+        cos = torch.cat([emb[0] for emb in embs], dim=1)  # (WHD, D/2)
+        sin = torch.cat([emb[1] for emb in embs], dim=1)  # (WHD, D/2)
+        return cos, sin
+    else:
+        emb = torch.cat(embs, dim=1)  # (WHD, D/2)
+        return emb
+
+
+def get_1d_rotary_pos_embed(
+    dim: int,
+    pos: Union[torch.FloatTensor, int],
+    theta: float = 10000.0,
+    use_real: bool = False,
+    theta_rescale_factor: float = 1.0,
+    interpolation_factor: float = 1.0,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Precompute the frequency tensor for complex exponential (cis) with given dimensions.
+    (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.)
+
+    This function calculates a frequency tensor with complex exponential using the given dimension 'dim'
+    and the end index 'end'. The 'theta' parameter scales the frequencies.
+    The returned tensor contains complex values in complex64 data type.
+
+    Args:
+        dim (int): Dimension of the frequency tensor.
+        pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar
+        theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.
+        use_real (bool, optional): If True, return real part and imaginary part separately.
+                                   Otherwise, return complex numbers.
+        theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0.
+
+    Returns:
+        freqs_cis: Precomputed frequency tensor with complex exponential. [S, D/2]
+        freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D]
+    """
+    if isinstance(pos, int):
+        pos = torch.arange(pos, device=torch.cuda.current_device()).float()
+
+    # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+    # has some connection to NTK literature
+    if theta_rescale_factor != 1.0:
+        theta *= theta_rescale_factor ** (dim / (dim - 2))
+
+    freqs = 1.0 / (
+        theta ** (torch.arange(0, dim, 2, device=torch.cuda.current_device())[: (dim // 2)].float() / dim)
+    )  # [D/2]
+    # assert interpolation_factor == 1.0, f"interpolation_factor: {interpolation_factor}"
+    freqs = torch.outer(pos * interpolation_factor, freqs)  # [S, D/2]
+    if use_real:
+        freqs_cos = freqs.cos().repeat_interleave(2, dim=1)  # [S, D]
+        freqs_sin = freqs.sin().repeat_interleave(2, dim=1)  # [S, D]
+        return freqs_cos, freqs_sin
+    else:
+        freqs_cis = torch.polar(
+            torch.ones_like(freqs), freqs
+        )  # complex64     # [S, D/2]
+        return freqs_cis