modeling/lance/modeling_utils.py

# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
#
# 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.
# coding: utf-8

import math

import numpy as np
import torch
from torch import nn
from transformers.activations import ACT2FN

# --------------------------------------------------------
# 2D sine-cosine position embedding
# References:
# DiT: https://github.com/facebookresearch/DiT/blob/main/models.py
# --------------------------------------------------------
def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
    grid_h = np.arange(grid_size, dtype=np.float32)
    grid_w = np.arange(grid_size, dtype=np.float32)
    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
    grid = np.stack(grid, axis=0)

    grid = grid.reshape([2, 1, grid_size, grid_size])
    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
    if cls_token and extra_tokens > 0:
        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
    return pos_embed


def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
    assert embed_dim % 2 == 0

    # use half of dimensions to encode grid_h
    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)

    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
    return emb


def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
    """
    embed_dim: output dimension for each position
    pos: a list of positions to be encoded: size (M,)
    out: (M, D)
    """
    assert embed_dim % 2 == 0
    omega = np.arange(embed_dim // 2, dtype=np.float64)
    omega /= embed_dim / 2.
    omega = 1. / 10000**omega  # (D/2,)

    pos = pos.reshape(-1)  # (M,)
    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product

    emb_sin = np.sin(out) # (M, D/2)
    emb_cos = np.cos(out) # (M, D/2)

    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
    return emb


def get_3d_sincos_pos_embed_from_grid(embed_dim, grid):
    """
    Get 3D sine-cosine positional embeddings from a grid.
    """
    assert embed_dim % 2 == 0, "Embedding dimension must be even for 3D embeddings"

    # 维度分配策略保持不变（确保每轴维度为偶数）
    d = embed_dim // 3
    d = d if d % 2 == 0 else d - 1
    dim_t, dim_h = d, d
    dim_w = embed_dim - 2 * d
    assert dim_w % 2 == 0

    emb_t = get_1d_sincos_pos_embed_from_grid(dim_t, grid[0])  # (T*H*W, Dt)
    emb_h = get_1d_sincos_pos_embed_from_grid(dim_h, grid[1])  # (T*H*W, Dh)
    emb_w = get_1d_sincos_pos_embed_from_grid(dim_w, grid[2])  # (T*H*W, Dw)
    return np.concatenate([emb_t, emb_h, emb_w], axis=1)


def get_3d_sincos_pos_embed(embed_dim, t, h, w):
    """
    Get 3D sine-cosine positional embeddings (v2 version, using thw indexing).
    """
    grid_t = np.arange(t, dtype=np.float32)
    grid_h = np.arange(h, dtype=np.float32)
    grid_w = np.arange(w, dtype=np.float32)
    tt, hh, ww = np.meshgrid(grid_t, grid_h, grid_w, indexing="ij")  # (t,h,w)

    grid = np.stack([tt, hh, ww], axis=0)  # [3, t, h, w]
    return get_3d_sincos_pos_embed_from_grid(embed_dim, grid)


# --------------------------------------------------------
# TimestepEmbedder
# Reference:
# DiT: https://github.com/facebookresearch/DiT/blob/main/models.py
# --------------------------------------------------------
class TimestepEmbedder(nn.Module):
    """
    Embeds scalar timesteps into vector representations.
    """
    def __init__(self, hidden_size, frequency_embedding_size=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size, bias=True),
        )
        self.frequency_embedding_size = frequency_embedding_size

    @staticmethod
    def timestep_embedding(t, dim, max_period=10000):
        """
        Create sinusoidal timestep embeddings.
        :param t: a 1-D Tensor of N indices, one per batch element.
                          These may be fractional.
        :param dim: the dimension of the output.
        :param max_period: controls the minimum frequency of the embeddings.
        :return: an (N, D) Tensor of positional embeddings.
        """
        half = dim // 2
        freqs = torch.exp(
            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
        ).to(device=t.device)
        args = t[:, None].float() * freqs[None]
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if dim % 2:
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
        return embedding

    def forward(self, t):
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
        t_emb = self.mlp(t_freq) # 跟llm的hidden size对齐
        return t_emb


class MLPconnector(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, hidden_act: str):
        super().__init__()
        self.activation_fn = ACT2FN[hidden_act]
        self.fc1 = nn.Linear(in_dim, out_dim)
        self.fc2 = nn.Linear(out_dim, out_dim)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.fc1(hidden_states)
        hidden_states = self.activation_fn(hidden_states)
        hidden_states = self.fc2(hidden_states)
        return hidden_states


class PositionEmbedding(nn.Module):
    def __init__(self, max_num_patch_per_side, hidden_size):
        super().__init__()
        self.max_num_patch_per_side = max_num_patch_per_side
        self.hidden_size = hidden_size
        self.pos_embed = nn.Parameter(
            torch.zeros(max_num_patch_per_side ** 2, hidden_size),
            requires_grad=False
        )
        self._init_weights()

    def _init_weights(self):
        # Initialize (and freeze) pos_embed by sin-cos embedding:
        pos_embed = get_2d_sincos_pos_embed(self.hidden_size, self.max_num_patch_per_side)
        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float())

    def forward(self, position_ids):
        return self.pos_embed[position_ids]


class PositionEmbedding3D(nn.Module):
    def __init__(self, max_latent_num_frames, max_latent_size, hidden_size):
        super().__init__()
        self.max_num_latent_frames = max_latent_num_frames  # t
        self.max_latent_size = max_latent_size  # h, w
        self.hidden_size = hidden_size
        self.temporal_dim, self.height_dim, self.width_dim = self._split_hidden_dims(hidden_size)

    @staticmethod
    def _split_hidden_dims(embed_dim: int) -> tuple[int, int, int]:
        assert embed_dim % 2 == 0, "Embedding dimension must be even for 3D embeddings"
        d = embed_dim // 3
        d = d if d % 2 == 0 else d - 1
        dim_t = d
        dim_h = d
        dim_w = embed_dim - 2 * d
        assert dim_w % 2 == 0
        return dim_t, dim_h, dim_w

    @staticmethod
    def _build_1d_sincos(coords: torch.Tensor, embed_dim: int) -> torch.Tensor:
        assert embed_dim % 2 == 0, "Embedding dimension must be even for 1D embeddings"
        half = embed_dim // 2
        omega = torch.arange(half, device=coords.device, dtype=torch.float32)
        omega = omega / (embed_dim / 2.0)
        omega = 1.0 / (10000.0 ** omega)
        args = coords.to(dtype=torch.float32)[:, None] * omega[None, :]
        return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)

    def forward(self, position_ids):
        position_ids = position_ids.reshape(-1).to(dtype=torch.long)
        plane_size = self.max_latent_size * self.max_latent_size
        t = position_ids // plane_size
        rem = position_ids % plane_size
        h = rem // self.max_latent_size
        w = rem % self.max_latent_size

        emb_t = self._build_1d_sincos(t, self.temporal_dim)
        emb_h = self._build_1d_sincos(h, self.height_dim)
        emb_w = self._build_1d_sincos(w, self.width_dim)
        return torch.cat([emb_t, emb_h, emb_w], dim=-1)