| 1 | """Utilities for building 2D self-attention masks for conditioning items.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from typing import TYPE_CHECKING |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | if TYPE_CHECKING: |
| 10 | from ltx_core.types import LatentState |
| 11 | |
| 12 | |
| 13 | def resolve_cross_mask( |
| 14 | attention_mask: float | int | torch.Tensor, |
| 15 | num_new_tokens: int, |
| 16 | batch_size: int, |
| 17 | device: torch.device, |
| 18 | dtype: torch.dtype, |
| 19 | ) -> torch.Tensor: |
| 20 | """Convert an attention_mask (scalar or tensor) to a (B, M) cross_mask tensor. |
| 21 | Args: |
| 22 | attention_mask: Scalar value applied uniformly, 1D tensor of shape (M,) |
| 23 | broadcast across batch, or 2D tensor of shape (B, M). |
| 24 | num_new_tokens: Number of new conditioning tokens M. |
| 25 | batch_size: Batch size B. |
| 26 | device: Device for the output tensor. |
| 27 | dtype: Data type for the output tensor. |
| 28 | Returns: |
| 29 | Cross-mask tensor of shape (B, M). |
| 30 | """ |
| 31 | if isinstance(attention_mask, (int, float)): |
| 32 | return torch.full( |
| 33 | (batch_size, num_new_tokens), |
| 34 | fill_value=float(attention_mask), |
| 35 | device=device, |
| 36 | dtype=dtype, |
| 37 | ) |
| 38 | mask = attention_mask.to(device=device, dtype=dtype) |
| 39 | |
| 40 | # Handle scalar (0-D) tensor like a Python scalar. |
| 41 | if mask.dim() == 0: |
| 42 | return torch.full( |
| 43 | (batch_size, num_new_tokens), |
| 44 | fill_value=float(mask.item()), |
| 45 | device=device, |
| 46 | dtype=dtype, |
| 47 | ) |
| 48 | |
| 49 | if mask.dim() == 1: |
| 50 | if mask.shape[0] != num_new_tokens: |
| 51 | raise ValueError( |
| 52 | f"1-D attention_mask length must equal num_new_tokens ({num_new_tokens}), got shape {tuple(mask.shape)}" |
| 53 | ) |
| 54 | mask = mask.unsqueeze(0).expand(batch_size, -1) |
| 55 | elif mask.dim() == 2: |
| 56 | b, m = mask.shape |
| 57 | if m != num_new_tokens: |
| 58 | raise ValueError( |
| 59 | f"2-D attention_mask second dimension must equal num_new_tokens ({num_new_tokens}), " |
| 60 | f"got shape {tuple(mask.shape)}" |
| 61 | ) |
| 62 | if b not in (batch_size, 1): |
| 63 | raise ValueError( |
| 64 | f"2-D attention_mask batch dimension must equal batch_size ({batch_size}) or 1, " |
| 65 | f"got shape {tuple(mask.shape)}" |
| 66 | ) |
| 67 | if b == 1 and batch_size > 1: |
| 68 | mask = mask.expand(batch_size, -1) |
| 69 | else: |
| 70 | raise ValueError( |
| 71 | f"attention_mask tensor must be 0-D, 1-D, or 2-D, got {mask.dim()}-D with shape {tuple(mask.shape)}" |
| 72 | ) |
| 73 | return mask |
| 74 | |
| 75 | |
| 76 | def update_attention_mask( |
| 77 | latent_state: LatentState, |
| 78 | attention_mask: float | torch.Tensor | None, |
| 79 | num_noisy_tokens: int, |
| 80 | num_new_tokens: int, |
| 81 | batch_size: int, |
| 82 | device: torch.device, |
| 83 | dtype: torch.dtype, |
| 84 | ) -> torch.Tensor | None: |
| 85 | """Build or update the self-attention mask for newly appended conditioning tokens. |
| 86 | If *attention_mask* is ``None`` and no existing mask is present, returns |
| 87 | ``None``. If *attention_mask* is ``None`` but an existing mask is present, |
| 88 | the mask is expanded with full attention (1s) for the new tokens so that |
| 89 | its dimensions stay consistent with the growing latent sequence. Otherwise, |
| 90 | resolves *attention_mask* to a per-token cross-mask and expands the 2-D |
| 91 | attention mask via :func:`build_attention_mask`. |
| 92 | Args: |
| 93 | latent_state: Current latent state (provides the existing mask and total |
| 94 | existing-token count). |
| 95 | attention_mask: Per-token attention weight. Scalar, 1-D ``(M,)``, 2-D |
| 96 | ``(B, M)`` tensor, or ``None`` (no-op). |
| 97 | num_noisy_tokens: Number of original noisy tokens (from |
| 98 | ``latent_tools.target_shape.token_count()``). |
| 99 | num_new_tokens: Number of new conditioning tokens being appended. |
| 100 | batch_size: Batch size. |
| 101 | device: Device for the output tensor. |
| 102 | dtype: Data type for the output tensor. |
| 103 | Returns: |
| 104 | Updated attention mask of shape ``(B, N+M, N+M)``, or ``None`` if no |
| 105 | masking is needed. |
| 106 | """ |
| 107 | if attention_mask is None: |
| 108 | if latent_state.attention_mask is None: |
| 109 | return None |
| 110 | # Existing mask present but no new mask requested: pad with 1s (full |
| 111 | # attention) so the mask dimensions stay consistent with the growing |
| 112 | # latent sequence. |
| 113 | cross_mask = torch.ones(batch_size, num_new_tokens, device=device, dtype=dtype) |
| 114 | return build_attention_mask( |
| 115 | existing_mask=latent_state.attention_mask, |
| 116 | num_noisy_tokens=num_noisy_tokens, |
| 117 | num_new_tokens=num_new_tokens, |
| 118 | num_existing_tokens=latent_state.latent.shape[1], |
| 119 | cross_mask=cross_mask, |
| 120 | device=device, |
| 121 | dtype=dtype, |
| 122 | ) |
| 123 | |
| 124 | cross_mask = resolve_cross_mask(attention_mask, num_new_tokens, batch_size, device, dtype) |
| 125 | return build_attention_mask( |
| 126 | existing_mask=latent_state.attention_mask, |
| 127 | num_noisy_tokens=num_noisy_tokens, |
| 128 | num_new_tokens=num_new_tokens, |
| 129 | num_existing_tokens=latent_state.latent.shape[1], |
| 130 | cross_mask=cross_mask, |
| 131 | device=device, |
| 132 | dtype=dtype, |
| 133 | ) |
| 134 | |
| 135 | |
| 136 | def build_attention_mask( |
| 137 | existing_mask: torch.Tensor | None, |
| 138 | num_noisy_tokens: int, |
| 139 | num_new_tokens: int, |
| 140 | num_existing_tokens: int, |
| 141 | cross_mask: torch.Tensor, |
| 142 | device: torch.device, |
| 143 | dtype: torch.dtype, |
| 144 | ) -> torch.Tensor: |
| 145 | """ |
| 146 | Expand the attention mask to include newly appended conditioning tokens. |
| 147 | Each conditioning item appends M new reference tokens to the sequence. This function |
| 148 | builds a (B, N+M, N+M) attention mask with the following block structure: |
| 149 | noisy prev_ref new_ref |
| 150 | (N_noisy) (N-N_noisy) (M) |
| 151 | ┌───────────┬───────────┬───────────┐ |
| 152 | noisy │ │ │ │ |
| 153 | (N_noisy) │ existing │ existing │ cross │ |
| 154 | │ │ │ │ |
| 155 | ├───────────┼───────────┼───────────┤ |
| 156 | prev_ref │ │ │ │ |
| 157 | (N-N_noisy)│ existing │ existing │ 0 │ |
| 158 | │ │ │ │ |
| 159 | ├───────────┼───────────┼───────────┤ |
| 160 | new_ref │ │ │ │ |
| 161 | (M) │ cross │ 0 │ 1 │ |
| 162 | │ │ │ │ |
| 163 | └───────────┴───────────┴───────────┘ |
| 164 | Where: |
| 165 | - **existing**: preserved from the previous mask (or 1.0 if first conditioning) |
| 166 | - **cross**: values from *cross_mask* (shape B, M), in [0, 1] |
| 167 | - **0**: no attention between different reference groups |
| 168 | Args: |
| 169 | existing_mask: Current attention mask of shape (B, N, N), or None if no mask exists yet. |
| 170 | When None, the top-left NxN block is filled with 1s (full attention between all |
| 171 | existing tokens including any prior reference tokens that had no mask). |
| 172 | num_noisy_tokens: Number of original noisy tokens (always at positions [0:num_noisy_tokens]). |
| 173 | num_new_tokens: Number of new conditioning tokens M being appended. |
| 174 | num_existing_tokens: Total number of current tokens N (noisy + any prior conditioning tokens). |
| 175 | cross_mask: Per-token attention weight of shape (B, M) controlling attention between |
| 176 | new reference tokens and noisy tokens. Values in [0, 1]. |
| 177 | device: Device for the output tensor. |
| 178 | dtype: Data type for the output tensor. |
| 179 | Returns: |
| 180 | Attention mask of shape (B, N+M, N+M) with values in [0, 1]. |
| 181 | """ |
| 182 | batch_size = cross_mask.shape[0] |
| 183 | total = num_existing_tokens + num_new_tokens |
| 184 | |
| 185 | # Start with zeros |
| 186 | mask = torch.zeros((batch_size, total, total), device=device, dtype=dtype) |
| 187 | |
| 188 | # Top-left: preserve existing mask or fill with 1s for noisy tokens |
| 189 | if existing_mask is not None: |
| 190 | mask[:, :num_existing_tokens, :num_existing_tokens] = existing_mask |
| 191 | else: |
| 192 | mask[:, :num_existing_tokens, :num_existing_tokens] = 1.0 |
| 193 | |
| 194 | # Bottom-right: new reference tokens fully attend to themselves |
| 195 | mask[:, num_existing_tokens:, num_existing_tokens:] = 1.0 |
| 196 | |
| 197 | # Cross-attention between noisy tokens and new reference tokens |
| 198 | # cross_mask shape: (B, M) -> broadcast to (B, N_noisy, M) and (B, M, N_noisy) |
| 199 | |
| 200 | # Noisy tokens attending to new reference tokens: [0:N_noisy, N:N+M] |
| 201 | # Each column j in this block gets cross_mask[:, j] |
| 202 | mask[:, :num_noisy_tokens, num_existing_tokens:] = cross_mask.unsqueeze(1) |
| 203 | |
| 204 | # New reference tokens attending to noisy tokens: [N:N+M, 0:N_noisy] |
| 205 | # Each row i in this block gets cross_mask[:, i] |
| 206 | mask[:, num_existing_tokens:, :num_noisy_tokens] = cross_mask.unsqueeze(2) |
| 207 | |
| 208 | # [N_noisy:N, N:N+M] and [N:N+M, N_noisy:N] remain 0 (no cross-ref attention) |
| 209 | |
| 210 | return mask |
| 211 |