bridge.diffusion.models.wan.rope_utils

Module Contents

Classes

Wan3DRopeEmbeddings

Wan 3D RoPE embeddings implementation. Implements Wan’s 3D RoPE embeddings for Mcore Attention based on Wan’s implementation at https://github.com/Wan-Video/Wan2.1/blob/main/wan/modules/model.py.

API

class bridge.diffusion.models.wan.rope_utils.Wan3DRopeEmbeddings(dim_head, max_position_len)

Bases: torch.nn.Module

Wan 3D RoPE embeddings implementation. Implements Wan’s 3D RoPE embeddings for Mcore Attention based on Wan’s implementation at https://github.com/Wan-Video/Wan2.1/blob/main/wan/modules/model.py.

Initialization

rope_params(max_position_len, dim_head, theta=10000)

forward(n_head, dim_head, cu_seqlens_q_padded, grid_sizes, device)