| 1 | import functools |
| 2 | import math |
| 3 | from enum import Enum |
| 4 | from typing import Callable, Tuple |
| 5 | |
| 6 | import numpy as np |
| 7 | import torch |
| 8 | from einops import rearrange |
| 9 | |
| 10 | |
| 11 | class LTXRopeType(Enum): |
| 12 | INTERLEAVED = "interleaved" |
| 13 | SPLIT = "split" |
| 14 | |
| 15 | |
| 16 | def apply_rotary_emb( |
| 17 | input_tensor: torch.Tensor, |
| 18 | freqs_cis: Tuple[torch.Tensor, torch.Tensor], |
| 19 | rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, |
| 20 | ) -> torch.Tensor: |
| 21 | if rope_type == LTXRopeType.INTERLEAVED: |
| 22 | return apply_interleaved_rotary_emb(input_tensor, *freqs_cis) |
| 23 | elif rope_type == LTXRopeType.SPLIT: |
| 24 | return apply_split_rotary_emb(input_tensor, *freqs_cis) |
| 25 | else: |
| 26 | raise ValueError(f"Invalid rope type: {rope_type}") |
| 27 | |
| 28 | |
| 29 | def apply_interleaved_rotary_emb( |
| 30 | input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor |
| 31 | ) -> torch.Tensor: |
| 32 | t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2) |
| 33 | t1, t2 = t_dup.unbind(dim=-1) |
| 34 | t_dup = torch.stack((-t2, t1), dim=-1) |
| 35 | input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)") |
| 36 | |
| 37 | out = input_tensor * cos_freqs + input_tensor_rot * sin_freqs |
| 38 | |
| 39 | return out |
| 40 | |
| 41 | |
| 42 | def apply_split_rotary_emb( |
| 43 | input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor |
| 44 | ) -> torch.Tensor: |
| 45 | needs_reshape = False |
| 46 | if input_tensor.ndim != 4 and cos_freqs.ndim == 4: |
| 47 | b, h, t, _ = cos_freqs.shape |
| 48 | input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2) |
| 49 | needs_reshape = True |
| 50 | |
| 51 | split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2) |
| 52 | first_half_input = split_input[..., :1, :] |
| 53 | second_half_input = split_input[..., 1:, :] |
| 54 | |
| 55 | output = split_input * cos_freqs.unsqueeze(-2) |
| 56 | first_half_output = output[..., :1, :] |
| 57 | second_half_output = output[..., 1:, :] |
| 58 | |
| 59 | first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input) |
| 60 | second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input) |
| 61 | |
| 62 | output = rearrange(output, "... d r -> ... (d r)") |
| 63 | if needs_reshape: |
| 64 | output = output.swapaxes(1, 2).reshape(b, t, -1) |
| 65 | |
| 66 | return output |
| 67 | |
| 68 | |
| 69 | @functools.lru_cache(maxsize=5) |
| 70 | def generate_freq_grid_np( |
| 71 | positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int |
| 72 | ) -> torch.Tensor: |
| 73 | theta = positional_embedding_theta |
| 74 | start = 1 |
| 75 | end = theta |
| 76 | |
| 77 | n_elem = 2 * positional_embedding_max_pos_count |
| 78 | pow_indices = np.power( |
| 79 | theta, |
| 80 | np.linspace( |
| 81 | np.log(start) / np.log(theta), |
| 82 | np.log(end) / np.log(theta), |
| 83 | inner_dim // n_elem, |
| 84 | dtype=np.float64, |
| 85 | ), |
| 86 | ) |
| 87 | return torch.tensor(pow_indices * math.pi / 2, dtype=torch.float32) |
| 88 | |
| 89 | |
| 90 | @functools.lru_cache(maxsize=5) |
| 91 | def generate_freq_grid_pytorch( |
| 92 | positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int |
| 93 | ) -> torch.Tensor: |
| 94 | theta = positional_embedding_theta |
| 95 | start = 1 |
| 96 | end = theta |
| 97 | n_elem = 2 * positional_embedding_max_pos_count |
| 98 | |
| 99 | indices = theta ** ( |
| 100 | torch.linspace( |
| 101 | math.log(start, theta), |
| 102 | math.log(end, theta), |
| 103 | inner_dim // n_elem, |
| 104 | dtype=torch.float32, |
| 105 | ) |
| 106 | ) |
| 107 | indices = indices.to(dtype=torch.float32) |
| 108 | |
| 109 | indices = indices * math.pi / 2 |
| 110 | |
| 111 | return indices |
| 112 | |
| 113 | |
| 114 | def get_fractional_positions(indices_grid: torch.Tensor, max_pos: list[int]) -> torch.Tensor: |
| 115 | n_pos_dims = indices_grid.shape[1] |
| 116 | assert n_pos_dims == len(max_pos), ( |
| 117 | f"Number of position dimensions ({n_pos_dims}) must match max_pos length ({len(max_pos)})" |
| 118 | ) |
| 119 | fractional_positions = torch.stack( |
| 120 | [indices_grid[:, i] / max_pos[i] for i in range(n_pos_dims)], |
| 121 | dim=-1, |
| 122 | ) |
| 123 | return fractional_positions |
| 124 | |
| 125 | |
| 126 | def generate_freqs( |
| 127 | indices: torch.Tensor, indices_grid: torch.Tensor, max_pos: list[int], use_middle_indices_grid: bool |
| 128 | ) -> torch.Tensor: |
| 129 | if use_middle_indices_grid: |
| 130 | assert len(indices_grid.shape) == 4 |
| 131 | assert indices_grid.shape[-1] == 2 |
| 132 | indices_grid_start, indices_grid_end = indices_grid[..., 0], indices_grid[..., 1] |
| 133 | indices_grid = (indices_grid_start + indices_grid_end) / 2.0 |
| 134 | elif len(indices_grid.shape) == 4: |
| 135 | indices_grid = indices_grid[..., 0] |
| 136 | |
| 137 | fractional_positions = get_fractional_positions(indices_grid, max_pos) |
| 138 | indices = indices.to(device=fractional_positions.device) |
| 139 | |
| 140 | freqs = (indices * (fractional_positions.unsqueeze(-1) * 2 - 1)).transpose(-1, -2).flatten(2) |
| 141 | return freqs |
| 142 | |
| 143 | |
| 144 | def split_freqs_cis(freqs: torch.Tensor, pad_size: int, num_attention_heads: int) -> tuple[torch.Tensor, torch.Tensor]: |
| 145 | cos_freq = freqs.cos() |
| 146 | sin_freq = freqs.sin() |
| 147 | |
| 148 | if pad_size != 0: |
| 149 | cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) |
| 150 | sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) |
| 151 | |
| 152 | cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1) |
| 153 | sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1) |
| 154 | |
| 155 | # Reshape freqs to be compatible with multi-head attention |
| 156 | b = cos_freq.shape[0] |
| 157 | t = cos_freq.shape[1] |
| 158 | |
| 159 | cos_freq = cos_freq.reshape(b, t, num_attention_heads, -1) |
| 160 | sin_freq = sin_freq.reshape(b, t, num_attention_heads, -1) |
| 161 | |
| 162 | cos_freq = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2) |
| 163 | sin_freq = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2) |
| 164 | return cos_freq, sin_freq |
| 165 | |
| 166 | |
| 167 | def interleaved_freqs_cis(freqs: torch.Tensor, pad_size: int) -> tuple[torch.Tensor, torch.Tensor]: |
| 168 | cos_freq = freqs.cos().repeat_interleave(2, dim=-1) |
| 169 | sin_freq = freqs.sin().repeat_interleave(2, dim=-1) |
| 170 | if pad_size != 0: |
| 171 | cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) |
| 172 | sin_padding = torch.zeros_like(cos_freq[:, :, :pad_size]) |
| 173 | cos_freq = torch.cat([cos_padding, cos_freq], dim=-1) |
| 174 | sin_freq = torch.cat([sin_padding, sin_freq], dim=-1) |
| 175 | return cos_freq, sin_freq |
| 176 | |
| 177 | |
| 178 | def precompute_freqs_cis( |
| 179 | indices_grid: torch.Tensor, |
| 180 | dim: int, |
| 181 | out_dtype: torch.dtype, |
| 182 | theta: float = 10000.0, |
| 183 | max_pos: list[int] | None = None, |
| 184 | use_middle_indices_grid: bool = False, |
| 185 | num_attention_heads: int = 32, |
| 186 | rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, |
| 187 | freq_grid_generator: Callable[[float, int, int, torch.device], torch.Tensor] = generate_freq_grid_pytorch, |
| 188 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 189 | if max_pos is None: |
| 190 | max_pos = [20, 2048, 2048] |
| 191 | |
| 192 | indices = freq_grid_generator(theta, indices_grid.shape[1], dim) |
| 193 | freqs = generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid) |
| 194 | |
| 195 | if rope_type == LTXRopeType.SPLIT: |
| 196 | expected_freqs = dim // 2 |
| 197 | current_freqs = freqs.shape[-1] |
| 198 | pad_size = expected_freqs - current_freqs |
| 199 | cos_freq, sin_freq = split_freqs_cis(freqs, pad_size, num_attention_heads) |
| 200 | else: |
| 201 | # 2 because of cos and sin by 3 for (t, x, y), 1 for temporal only |
| 202 | n_elem = 2 * indices_grid.shape[1] |
| 203 | cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem) |
| 204 | return cos_freq.to(out_dtype), sin_freq.to(out_dtype) |
| 205 |