89 lines
2.6 KiB
Python
89 lines
2.6 KiB
Python
#original code from https://github.com/genmoai/models under apache 2.0 license
|
|
|
|
# import functools
|
|
import math
|
|
|
|
import torch
|
|
|
|
|
|
def centers(start: float, stop, num, dtype=None, device=None):
|
|
"""linspace through bin centers.
|
|
|
|
Args:
|
|
start (float): Start of the range.
|
|
stop (float): End of the range.
|
|
num (int): Number of points.
|
|
dtype (torch.dtype): Data type of the points.
|
|
device (torch.device): Device of the points.
|
|
|
|
Returns:
|
|
centers (Tensor): Centers of the bins. Shape: (num,).
|
|
"""
|
|
edges = torch.linspace(start, stop, num + 1, dtype=dtype, device=device)
|
|
return (edges[:-1] + edges[1:]) / 2
|
|
|
|
|
|
# @functools.lru_cache(maxsize=1)
|
|
def create_position_matrix(
|
|
T: int,
|
|
pH: int,
|
|
pW: int,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
*,
|
|
target_area: float = 36864,
|
|
):
|
|
"""
|
|
Args:
|
|
T: int - Temporal dimension
|
|
pH: int - Height dimension after patchify
|
|
pW: int - Width dimension after patchify
|
|
|
|
Returns:
|
|
pos: [T * pH * pW, 3] - position matrix
|
|
"""
|
|
# Create 1D tensors for each dimension
|
|
t = torch.arange(T, dtype=dtype)
|
|
|
|
# Positionally interpolate to area 36864.
|
|
# (3072x3072 frame with 16x16 patches = 192x192 latents).
|
|
# This automatically scales rope positions when the resolution changes.
|
|
# We use a large target area so the model is more sensitive
|
|
# to changes in the learned pos_frequencies matrix.
|
|
scale = math.sqrt(target_area / (pW * pH))
|
|
w = centers(-pW * scale / 2, pW * scale / 2, pW)
|
|
h = centers(-pH * scale / 2, pH * scale / 2, pH)
|
|
|
|
# Use meshgrid to create 3D grids
|
|
grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij")
|
|
|
|
# Stack and reshape the grids.
|
|
pos = torch.stack([grid_t, grid_h, grid_w], dim=-1) # [T, pH, pW, 3]
|
|
pos = pos.view(-1, 3) # [T * pH * pW, 3]
|
|
pos = pos.to(dtype=dtype, device=device)
|
|
|
|
return pos
|
|
|
|
|
|
def compute_mixed_rotation(
|
|
freqs: torch.Tensor,
|
|
pos: torch.Tensor,
|
|
):
|
|
"""
|
|
Project each 3-dim position into per-head, per-head-dim 1D frequencies.
|
|
|
|
Args:
|
|
freqs: [3, num_heads, num_freqs] - learned rotation frequency (for t, row, col) for each head position
|
|
pos: [N, 3] - position of each token
|
|
num_heads: int
|
|
|
|
Returns:
|
|
freqs_cos: [N, num_heads, num_freqs] - cosine components
|
|
freqs_sin: [N, num_heads, num_freqs] - sine components
|
|
"""
|
|
assert freqs.ndim == 3
|
|
freqs_sum = torch.einsum("Nd,dhf->Nhf", pos.to(freqs), freqs)
|
|
freqs_cos = torch.cos(freqs_sum)
|
|
freqs_sin = torch.sin(freqs_sum)
|
|
return freqs_cos, freqs_sin
|