lingbot_map/layers/rope.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.


# Implementation of 2D Rotary Position Embeddings (RoPE).

# This module provides a clean implementation of 2D Rotary Position Embeddings,
# which extends the original RoPE concept to handle 2D spatial positions.

# Inspired by:
#         https://github.com/meta-llama/codellama/blob/main/llama/model.py
#         https://github.com/naver-ai/rope-vit


import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple

from typing import List, Optional, Tuple, Union


class PositionGetter:
    """Generates and caches 2D spatial positions for patches in a grid.

    This class efficiently manages the generation of spatial coordinates for patches
    in a 2D grid, caching results to avoid redundant computations.

    Attributes:
        position_cache: Dictionary storing precomputed position tensors for different
            grid dimensions.
    """

    def __init__(self):
        """Initializes the position generator with an empty cache."""
        self.position_cache: Dict[Tuple[int, int], torch.Tensor] = {}

    def __call__(self, batch_size: int, height: int, width: int, device: torch.device) -> torch.Tensor:
        """Generates spatial positions for a batch of patches.

        Args:
            batch_size: Number of samples in the batch.
            height: Height of the grid in patches.
            width: Width of the grid in patches.
            device: Target device for the position tensor.

        Returns:
            Tensor of shape (batch_size, height*width, 2) containing y,x coordinates
            for each position in the grid, repeated for each batch item.
        """
        if (height, width) not in self.position_cache:
            y_coords = torch.arange(height, device=device)
            x_coords = torch.arange(width, device=device)
            positions = torch.cartesian_prod(y_coords, x_coords)
            self.position_cache[height, width] = positions

        cached_positions = self.position_cache[height, width]
        return cached_positions.view(1, height * width, 2).expand(batch_size, -1, -1).clone()


class RotaryPositionEmbedding2D(nn.Module):
    """2D Rotary Position Embedding implementation.

    This module applies rotary position embeddings to input tokens based on their
    2D spatial positions. It handles the position-dependent rotation of features
    separately for vertical and horizontal dimensions.

    Args:
        frequency: Base frequency for the position embeddings. Default: 100.0
        scaling_factor: Scaling factor for frequency computation. Default: 1.0

    Attributes:
        base_frequency: Base frequency for computing position embeddings.
        scaling_factor: Factor to scale the computed frequencies.
        frequency_cache: Cache for storing precomputed frequency components.
    """

    def __init__(self, frequency: float = 100.0, scaling_factor: float = 1.0):
        """Initializes the 2D RoPE module."""
        super().__init__()
        self.base_frequency = frequency
        self.scaling_factor = scaling_factor
        self.frequency_cache: Dict[Tuple, Tuple[torch.Tensor, torch.Tensor]] = {}

    def _compute_frequency_components(
        self, dim: int, seq_len: int, device: torch.device, dtype: torch.dtype
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Computes frequency components for rotary embeddings.

        Args:
            dim: Feature dimension (must be even).
            seq_len: Maximum sequence length.
            device: Target device for computations.
            dtype: Data type for the computed tensors.

        Returns:
            Tuple of (cosine, sine) tensors for frequency components.
        """
        cache_key = (dim, seq_len, device, dtype)
        if cache_key not in self.frequency_cache:
            # Compute frequency bands
            exponents = torch.arange(0, dim, 2, device=device).float() / dim
            inv_freq = 1.0 / (self.base_frequency**exponents)

            # Generate position-dependent frequencies
            positions = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
            angles = torch.einsum("i,j->ij", positions, inv_freq)

            # Compute and cache frequency components
            angles = angles.to(dtype)
            angles = torch.cat((angles, angles), dim=-1)
            cos_components = angles.cos().to(dtype)
            sin_components = angles.sin().to(dtype)
            self.frequency_cache[cache_key] = (cos_components, sin_components)

        return self.frequency_cache[cache_key]

    @staticmethod
    def _rotate_features(x: torch.Tensor) -> torch.Tensor:
        """Performs feature rotation by splitting and recombining feature dimensions.

        Args:
            x: Input tensor to rotate.

        Returns:
            Rotated feature tensor.
        """
        feature_dim = x.shape[-1]
        x1, x2 = x[..., : feature_dim // 2], x[..., feature_dim // 2 :]
        return torch.cat((-x2, x1), dim=-1)

    def _apply_1d_rope(
        self, tokens: torch.Tensor, positions: torch.Tensor, cos_comp: torch.Tensor, sin_comp: torch.Tensor
    ) -> torch.Tensor:
        """Applies 1D rotary position embeddings along one dimension.

        Args:
            tokens: Input token features.
            positions: Position indices.
            cos_comp: Cosine components for rotation.
            sin_comp: Sine components for rotation.

        Returns:
            Tokens with applied rotary position embeddings.
        """
        # Embed positions with frequency components
        cos = F.embedding(positions, cos_comp)[:, None, :, :]
        sin = F.embedding(positions, sin_comp)[:, None, :, :]

        # Apply rotation
        return (tokens * cos) + (self._rotate_features(tokens) * sin)

    def forward(self, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        """Applies 2D rotary position embeddings to input tokens.

        Args:
            tokens: Input tensor of shape (batch_size, n_heads, n_tokens, dim).
                   The feature dimension (dim) must be divisible by 4.
            positions: Position tensor of shape (batch_size, n_tokens, 2) containing
                      the y and x coordinates for each token.

        Returns:
            Tensor of same shape as input with applied 2D rotary position embeddings.

        Raises:
            AssertionError: If input dimensions are invalid or positions are malformed.
        """
        # Validate inputs
        assert tokens.size(-1) % 2 == 0, "Feature dimension must be even"
        assert positions.ndim == 3 and positions.shape[-1] == 2, "Positions must have shape (batch_size, n_tokens, 2)"

        # Compute feature dimension for each spatial direction
        feature_dim = tokens.size(-1) // 2

        # Get frequency components
        max_position = int(positions.max()) + 1
        cos_comp, sin_comp = self._compute_frequency_components(feature_dim, max_position, tokens.device, tokens.dtype)

        # Split features for vertical and horizontal processing
        vertical_features, horizontal_features = tokens.chunk(2, dim=-1)

        # Apply RoPE separately for each dimension
        vertical_features = self._apply_1d_rope(vertical_features, positions[..., 0], cos_comp, sin_comp)
        horizontal_features = self._apply_1d_rope(horizontal_features, positions[..., 1], cos_comp, sin_comp)

        # Combine processed features
        return torch.cat((vertical_features, horizontal_features), dim=-1)
    


def get_1d_rotary_pos_embed(
    dim: int,
    pos: Union[np.ndarray, int],
    theta: float = 10000.0,
    use_real=False,
    linear_factor=1.0,
    ntk_factor=1.0,
    repeat_interleave_real=True,
    freqs_dtype=torch.float32,  #  torch.float32, torch.float64 (flux)
):
    """
    计算1D旋转位置编码（RoPE）的频率张量。
    
    RoPE的核心思想：使用旋转矩阵来编码位置信息，使得相对位置关系保持不变。
    公式：对于位置m和维度i，频率为 θ_i = θ^(-2i/d)，其中θ是基础频率（默认10000）
    
    Args:
        dim: 特征维度，必须是偶数（因为要成对处理）
        pos: 位置索引，可以是整数（自动生成0到pos-1的序列）或位置数组 [S]
        theta: 基础频率，控制位置编码的周期性（默认10000）
        use_real: 是否返回实数形式（cos和sin分开）还是复数形式
        linear_factor: 线性缩放因子，用于上下文扩展
        ntk_factor: NTK-Aware缩放因子，用于处理更长的序列
        repeat_interleave_real: 当use_real=True时，是否交错重复（用于某些模型架构）
        freqs_dtype: 频率张量的数据类型
        
    Returns:
        复数形式：[S, D/2] 的复数张量，表示 e^(i*m*θ_j)
        实数形式：两个 [S, D] 的张量（cos和sin）
    """
    # 确保维度是偶数（RoPE需要成对处理维度）
    assert dim % 2 == 0

    # 将位置转换为torch张量
    if isinstance(pos, int):
        pos = torch.arange(pos)  # 生成 [0, 1, 2, ..., pos-1]
    if isinstance(pos, np.ndarray):
        pos = torch.from_numpy(pos)  # [S]

    # 应用NTK缩放（Neural Tangent Kernel，用于处理训练时未见过的长序列）
    theta = theta * ntk_factor
    
    # 步骤1：计算频率 θ_i = 1 / (θ^(2i/d))
    # 其中 i ∈ {0, 2, 4, ..., dim-2}（只取偶数索引，因为成对处理）
    # 公式：freq_i = 1 / (theta^(2i/d) * linear_factor)
    freqs = (
        1.0
        / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=pos.device)[: (dim // 2)] / dim))
        / linear_factor
    )  # [D/2]，每个频率对应一个维度对
    
    # 步骤2：计算位置-频率矩阵
    # 使用外积：pos[m] * freqs[i] = m * θ_i
    # 结果：每个位置m和每个频率i的组合
    freqs = torch.outer(pos, freqs)  # [S, D/2]
    
    # 步骤3：根据返回格式转换
    if use_real and repeat_interleave_real:
        # 方式1：交错重复（用于flux, hunyuan-dit, cogvideox等模型）
        # 将每个频率的cos和sin交错排列：[cos_0, cos_0, cos_1, cos_1, ...]
        freqs_cos = freqs.cos().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float()  # [S, D]
        freqs_sin = freqs.sin().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float()  # [S, D]
        return freqs_cos, freqs_sin
    elif use_real:
        # 方式2：拼接重复（用于stable audio, allegro等模型）
        # 将所有cos拼接，然后是所有sin：[cos_0, cos_1, ..., cos_n, cos_0, cos_1, ..., cos_n]
        freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float()  # [S, D]
        freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float()  # [S, D]
        return freqs_cos, freqs_sin
    else:
        # 方式3：复数形式（用于lumina等模型）
        # 使用欧拉公式：e^(iθ) = cos(θ) + i*sin(θ)
        # torch.polar(r, θ) 返回 r * e^(iθ)，这里r=1，所以就是 e^(i*freqs)
        freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64: [S, D/2]
        return freqs_cis


class WanRotaryPosEmbed(nn.Module):
    """
    3D旋转位置编码（3D RoPE）模块
    
    核心思想：将RoPE扩展到3D空间（时间、高度、宽度），为视频或3D数据提供位置编码。
    每个维度（t, h, w）独立使用RoPE，然后拼接起来。
    
    公式：
    对于3D位置 (f, h, w)（帧、高度、宽度）：
    - 帧维度使用 dim_f 个特征维度
    - 高度维度使用 dim_h 个特征维度  
    - 宽度维度使用 dim_w 个特征维度
    其中 dim_f + dim_h + dim_w = attention_head_dim
    """
    def __init__(
        self,
        attention_head_dim: int,
        patch_size: Tuple[int, int, int],
        max_seq_len: int = 1024,
        theta: float = 10000.0,
        fhw_dim: Optional[Tuple[int, int, int]] = [20, 22, 22],
    ):
        super().__init__()

        self.attention_head_dim = attention_head_dim  # 注意力头的总维度
        self.patch_size = patch_size  # patch大小 (patch_f, patch_h, patch_w)
        self.max_seq_len = max_seq_len  # 最大序列长度（用于预计算频率）

        # 步骤1：分配维度给三个空间维度
        if fhw_dim is not None:
            # 如果指定了维度分配，使用指定的
            assert attention_head_dim == sum(
                fhw_dim
            ), f"attention_head_dim {attention_head_dim} must match sum(fhw_dim) {sum(fhw_dim)}"
            t_dim, h_dim, w_dim = fhw_dim
        else:
            # 否则自动分配：h和w各占1/3，t占剩余
            # 例如：如果attention_head_dim=64，则 h_dim=w_dim=21，t_dim=22
            h_dim = w_dim = 2 * (attention_head_dim // 6)
            t_dim = attention_head_dim - h_dim - w_dim
        
        # 保存维度分配以便在forward中使用
        self.fhw_dim = (t_dim, h_dim, w_dim)

        # 步骤2：为每个维度预计算频率
        # 分别计算时间、高度、宽度三个维度的RoPE频率
        freqs = []
        for dim in [t_dim, h_dim, w_dim]:
            # 每个维度独立调用1D RoPE
            # 返回复数形式的频率: [max_seq_len, dim//2]
            freq = get_1d_rotary_pos_embed(
                dim, max_seq_len, theta, use_real=False, repeat_interleave_real=False, freqs_dtype=torch.float64
            )
            freqs.append(freq)
        # 将三个维度的频率在最后一维拼接: [max_seq_len, (t_dim + h_dim + w_dim)//2]
        self.freqs = torch.cat(freqs, dim=1)

    def forward(self, ppf, pph, ppw, patch_start_idx, device: torch.device, f_start: int = 0, f_end: Optional[int] = None) -> torch.Tensor:
        """
        前向传播：为3D输入（视频帧+patch）生成旋转位置编码
        
        参数：
        - ppf (int): 帧数（patches per frame），当f_end为None时使用
        - pph (int): 每帧的patch高度数量
        - ppw (int): 每帧的patch宽度数量  
        - patch_start_idx (int): 每帧的特殊token数量（在patches之前）
        - device: 计算设备（CPU/GPU）
        - f_start (int): 起始帧索引（用于causal模式），默认为0
        - f_end (Optional[int]): 结束帧索引（用于causal模式），如果为None则使用ppf作为帧数
        
        返回：
        - freqs: [1, 1, ppf * (patch_start_idx + pph * ppw), head_dim//2] 复数频率tensor
        
        Token排列顺序：
        [frame0_special_token_0, ..., frame0_special_token_N,
         frame0_patch_0, ..., frame0_patch_M,
         frame1_special_token_0, ..., frame1_special_token_N,
         frame1_patch_0, ..., frame1_patch_M,
         ...]
        
        模式：
        - 非causal模式：f_end=None，使用ppf作为帧数，从位置0开始
        - Causal模式：f_end不为None，使用[f_start, f_end)范围的帧，ppf会被重新计算
        """

        # 步骤1：将预计算的频率移到目标设备，并分割成三个维度
        self.freqs = self.freqs.to(device)
        # 获取实际的维度分配
        if hasattr(self, 'fhw_dim') and self.fhw_dim is not None:
            t_dim, h_dim, w_dim = self.fhw_dim
        else:
            # 自动分配的情况
            h_dim = w_dim = 2 * (self.attention_head_dim // 6)
            t_dim = self.attention_head_dim - h_dim - w_dim
        
        # 使用正确的split sizes（每个维度的一半）
        freqs = self.freqs.split_with_sizes(
            [
                t_dim // 2,  # 时间维度
                h_dim // 2,  # 高度维度
                w_dim // 2,  # 宽度维度
            ],
            dim=1,
        )
        
        # 处理causal模式：如果指定了f_end，重新计算ppf和帧范围
        if f_end is not None:
            ppf = f_end - f_start
            frame_slice = slice(f_start, f_end)
        else:
            # 非causal模式：使用从0开始的ppf个帧
            frame_slice = slice(0, ppf)
        
        # 步骤2：处理特殊token（如果存在）
        ## For other tokens
        if patch_start_idx > 0:
            # 2.1 为特殊token生成位置编码
            # 特殊token位于对角线位置 (f, i, i)，每个特殊token有唯一位置
            # camera: (f, 0, 0), register_0: (f, 1, 1), ..., scale: (f, 5, 5)
            # Shape: (ppf, patch_start_idx, dim)
            freqs_special_f = freqs[0][frame_slice].reshape(ppf, 1, -1).expand(ppf, patch_start_idx, -1)  # (ppf, patch_start_idx, dim_f) 帧维度变化
            freqs_special_h = freqs[1][:patch_start_idx].reshape(1, patch_start_idx, -1).expand(ppf, patch_start_idx, -1)  # (ppf, patch_start_idx, dim_h) 高度=0,1,2,...
            freqs_special_w = freqs[2][:patch_start_idx].reshape(1, patch_start_idx, -1).expand(ppf, patch_start_idx, -1)  # (ppf, patch_start_idx, dim_w) 宽度=0,1,2,...
            freqs_special = torch.cat([freqs_special_f, freqs_special_h, freqs_special_w], dim=-1)  # (ppf, patch_start_idx, dim) 拼接三维
            freqs_special = freqs_special.reshape(ppf, patch_start_idx, -1)  # (ppf, patch_start_idx, dim)

            # 2.2 为图像patch生成位置编码
            # Patch位于 (f, patch_start_idx+h, patch_start_idx+w)，h,w 整体偏移 patch_start_idx
            # 这样 patches 与 special tokens 位置不冲突，且 h,w 对称处理
            # Shape: (ppf, pph, ppw, dim)
            freqs_f = freqs[0][frame_slice].reshape(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1)  # (ppf, pph, ppw, dim_f) 帧维度
            freqs_h = freqs[1][patch_start_idx : patch_start_idx + pph].reshape(1, pph, 1, -1).expand(ppf, pph, ppw, -1)  # (ppf, pph, ppw, dim_h) 高度从patch_start_idx开始
            freqs_w = freqs[2][patch_start_idx : patch_start_idx + ppw].reshape(1, 1, ppw, -1).expand(ppf, pph, ppw, -1)  # (ppf, pph, ppw, dim_w) 宽度从patch_start_idx开始
            freqs_patches = torch.cat([freqs_f, freqs_h, freqs_w], dim=-1)  # (ppf, pph, ppw, dim) 拼接三维
            freqs_patches = freqs_patches.reshape(ppf, pph * ppw, -1)  # (ppf, pph * ppw, dim) 展平空间维度
            
            # 步骤3：按照正确的顺序组合特殊token和patches
            # 每帧内部顺序：[特殊tokens, patches]
            # Concatenate special tokens and patches for each frame along the second dimension
            # Shape: (ppf, patch_start_idx + pph * ppw, dim)
            freqs = torch.cat([freqs_special, freqs_patches], dim=1)  # (ppf, patch_start_idx + pph * ppw, dim)
            
            # 步骤4：展平为最终形状并添加batch和head维度
            # Flatten to get final shape: (ppf * (patch_start_idx + pph * ppw), dim)
            freqs = freqs.reshape(ppf * (patch_start_idx + pph * ppw), -1)
            freqs = freqs.unsqueeze(0).unsqueeze(0)  # (1, 1, ppf * (patch_start_idx + pph * ppw), dim) 添加batch和head维度
            return freqs
        
        # 如果没有特殊token（patch_start_idx == 0），只处理图像patches
        # 所有patches位于 (f, 0:pph, 0:ppw)
        freqs_f = freqs[0][frame_slice].reshape(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1)  # (ppf, pph, ppw, dim_f) 帧维度
        freqs_h = freqs[1][:pph].reshape(1, pph, 1, -1).expand(ppf, pph, ppw, -1)  # (ppf, pph, ppw, dim_h) 高度从0开始
        freqs_w = freqs[2][:ppw].reshape(1, 1, ppw, -1).expand(ppf, pph, ppw, -1)  # (ppf, pph, ppw, dim_w) 宽度从0开始
        freqs = torch.cat([freqs_f, freqs_h, freqs_w], dim=-1).reshape(1, 1, ppf * pph * ppw, -1)  # (1, 1, ppf * pph * ppw, dim)
        return freqs
    
def apply_rotary_emb(x, freqs):
    """
    应用旋转位置编码到输入特征
    
    核心思想：使用复数乘法实现特征旋转，保持相对位置信息
    
    数学原理：
    对于2D向量 [x1, x2]，旋转θ角度可以表示为复数乘法：
    (x1 + ix2) * e^(iθ) = (x1 + ix2) * (cos(θ) + i*sin(θ))
                        = (x1*cos(θ) - x2*sin(θ)) + i*(x1*sin(θ) + x2*cos(θ))
    
    这等价于旋转矩阵：
    [cos(θ)  -sin(θ)] [x1]
    [sin(θ)   cos(θ)] [x2]
    
    参数：
    - x: 输入特征 [batch, heads, seq_len, head_dim]
    - freqs: 旋转频率（复数） [1, 1, seq_len, head_dim//2]
    
    返回：
    - x_out: 旋转后的特征 [batch, heads, seq_len, head_dim]
    
    实现步骤：
    1. 将x的每两个连续特征看作一个复数 (real, imag)
    2. 与预计算的复数频率 e^(iθ) 相乘
    3. 转回实数表示
    """
    # 步骤1：reshape成 [..., head_dim//2, 2] 形式，最后一维表示(real, imag)
    # 例如：[b, h, seq, 64] -> [b, h, seq, 32, 2]
    x_reshaped = x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2)
    
    # 步骤2：转换为复数表示 [b, h, seq, 32]
    # 每个元素是 real + imag*i
    x_complex = torch.view_as_complex(x_reshaped)
    
    # 步骤3：复数乘法实现旋转
    # x_complex * freqs 相当于将每对特征旋转θ角度
    # freqs已经是 e^(iθ) = cos(θ) + i*sin(θ) 的形式
    x_rotated = x_complex * freqs
    
    # 步骤4：转回实数表示 [b, h, seq, 32, 2]
    x_real = torch.view_as_real(x_rotated)
    
    # 步骤5：展平最后两维 [b, h, seq, 64]
    x_out = x_real.flatten(3)
    
    # 步骤6：转回原始数据类型
    return x_out.to(x.dtype)