| 1 | import math |
| 2 | from collections.abc import Mapping, Sequence |
| 3 | from dataclasses import dataclass, field |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from ltx_core.components.protocols import GuiderProtocol |
| 8 | |
| 9 | |
| 10 | @dataclass(frozen=True) |
| 11 | class CFGGuider(GuiderProtocol): |
| 12 | """ |
| 13 | Classifier-free guidance (CFG) guider. |
| 14 | Computes the guidance delta as (scale - 1) * (cond - uncond), steering the |
| 15 | denoising process toward the conditioned prediction. |
| 16 | Attributes: |
| 17 | scale: Guidance strength. 1.0 means no guidance, higher values increase |
| 18 | adherence to the conditioning. |
| 19 | """ |
| 20 | |
| 21 | scale: float |
| 22 | |
| 23 | def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor: |
| 24 | return (self.scale - 1) * (cond - uncond) |
| 25 | |
| 26 | def enabled(self) -> bool: |
| 27 | return self.scale != 1.0 |
| 28 | |
| 29 | |
| 30 | @dataclass(frozen=True) |
| 31 | class CFGStarRescalingGuider(GuiderProtocol): |
| 32 | """ |
| 33 | Calculates the CFG delta between conditioned and unconditioned samples. |
| 34 | To minimize offset in the denoising direction and move mostly along the |
| 35 | conditioning axis within the distribution, the unconditioned sample is |
| 36 | rescaled in accordance with the norm of the conditioned sample. |
| 37 | Attributes: |
| 38 | scale (float): |
| 39 | Global guidance strength. A value of 1.0 corresponds to no extra |
| 40 | guidance beyond the base model prediction. Values > 1.0 increase |
| 41 | the influence of the conditioned sample relative to the |
| 42 | unconditioned one. |
| 43 | """ |
| 44 | |
| 45 | scale: float |
| 46 | |
| 47 | def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor: |
| 48 | rescaled_neg = projection_coef(cond, uncond) * uncond |
| 49 | return (self.scale - 1) * (cond - rescaled_neg) |
| 50 | |
| 51 | def enabled(self) -> bool: |
| 52 | return self.scale != 1.0 |
| 53 | |
| 54 | |
| 55 | @dataclass(frozen=True) |
| 56 | class STGGuider(GuiderProtocol): |
| 57 | """ |
| 58 | Calculates the STG delta between conditioned and perturbed denoised samples. |
| 59 | Perturbed samples are the result of the denoising process with perturbations, |
| 60 | e.g. attentions acting as passthrough for certain layers and modalities. |
| 61 | Attributes: |
| 62 | scale (float): |
| 63 | Global strength of the STG guidance. A value of 0.0 disables the |
| 64 | guidance. Larger values increase the correction applied in the |
| 65 | direction of (pos_denoised - perturbed_denoised). |
| 66 | """ |
| 67 | |
| 68 | scale: float |
| 69 | |
| 70 | def delta(self, pos_denoised: torch.Tensor, perturbed_denoised: torch.Tensor) -> torch.Tensor: |
| 71 | return self.scale * (pos_denoised - perturbed_denoised) |
| 72 | |
| 73 | def enabled(self) -> bool: |
| 74 | return self.scale != 0.0 |
| 75 | |
| 76 | |
| 77 | @dataclass(frozen=True) |
| 78 | class LtxAPGGuider(GuiderProtocol): |
| 79 | """ |
| 80 | Calculates the APG (adaptive projected guidance) delta between conditioned |
| 81 | and unconditioned samples. |
| 82 | To minimize offset in the denoising direction and move mostly along the |
| 83 | conditioning axis within the distribution, the (cond - uncond) delta is |
| 84 | decomposed into components parallel and orthogonal to the conditioned |
| 85 | sample. The `eta` parameter weights the parallel component, while `scale` |
| 86 | is applied to the orthogonal component. Optionally, a norm threshold can |
| 87 | be used to suppress guidance when the magnitude of the correction is small. |
| 88 | Attributes: |
| 89 | scale (float): |
| 90 | Strength applied to the component of the guidance that is orthogonal |
| 91 | to the conditioned sample. Controls how aggressively we move in |
| 92 | directions that change semantics but stay consistent with the |
| 93 | conditioning manifold. |
| 94 | eta (float): |
| 95 | Weight of the component of the guidance that is parallel to the |
| 96 | conditioned sample. A value of 1.0 keeps the full parallel |
| 97 | component; values in [0, 1] attenuate it, and values > 1.0 amplify |
| 98 | motion along the conditioning direction. |
| 99 | norm_threshold (float): |
| 100 | Minimum L2 norm of the guidance delta below which the guidance |
| 101 | can be reduced or ignored (depending on implementation). |
| 102 | This is useful for avoiding noisy or unstable updates when the |
| 103 | guidance signal is very small. |
| 104 | """ |
| 105 | |
| 106 | scale: float |
| 107 | eta: float = 1.0 |
| 108 | norm_threshold: float = 0.0 |
| 109 | |
| 110 | def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor: |
| 111 | guidance = cond - uncond |
| 112 | if self.norm_threshold > 0: |
| 113 | ones = torch.ones_like(guidance) |
| 114 | guidance_norm = guidance.norm(p=2, dim=[-1, -2, -3], keepdim=True) |
| 115 | scale_factor = torch.minimum(ones, self.norm_threshold / guidance_norm) |
| 116 | guidance = guidance * scale_factor |
| 117 | proj_coeff = projection_coef(guidance, cond) |
| 118 | g_parallel = proj_coeff * cond |
| 119 | g_orth = guidance - g_parallel |
| 120 | g_apg = g_parallel * self.eta + g_orth |
| 121 | |
| 122 | return g_apg * (self.scale - 1) |
| 123 | |
| 124 | def enabled(self) -> bool: |
| 125 | return self.scale != 1.0 |
| 126 | |
| 127 | |
| 128 | @dataclass(frozen=False) |
| 129 | class LegacyStatefulAPGGuider(GuiderProtocol): |
| 130 | """ |
| 131 | Calculates the APG (adaptive projected guidance) delta between conditioned |
| 132 | and unconditioned samples. |
| 133 | To minimize offset in the denoising direction and move mostly along the |
| 134 | conditioning axis within the distribution, the (cond - uncond) delta is |
| 135 | decomposed into components parallel and orthogonal to the conditioned |
| 136 | sample. The `eta` parameter weights the parallel component, while `scale` |
| 137 | is applied to the orthogonal component. Optionally, a norm threshold can |
| 138 | be used to suppress guidance when the magnitude of the correction is small. |
| 139 | Attributes: |
| 140 | scale (float): |
| 141 | Strength applied to the component of the guidance that is orthogonal |
| 142 | to the conditioned sample. Controls how aggressively we move in |
| 143 | directions that change semantics but stay consistent with the |
| 144 | conditioning manifold. |
| 145 | eta (float): |
| 146 | Weight of the component of the guidance that is parallel to the |
| 147 | conditioned sample. A value of 1.0 keeps the full parallel |
| 148 | component; values in [0, 1] attenuate it, and values > 1.0 amplify |
| 149 | motion along the conditioning direction. |
| 150 | norm_threshold (float): |
| 151 | Minimum L2 norm of the guidance delta below which the guidance |
| 152 | can be reduced or ignored (depending on implementation). |
| 153 | This is useful for avoiding noisy or unstable updates when the |
| 154 | guidance signal is very small. |
| 155 | momentum (float): |
| 156 | Exponential moving-average coefficient for accumulating guidance |
| 157 | over time. running_avg = momentum * running_avg + guidance |
| 158 | """ |
| 159 | |
| 160 | scale: float |
| 161 | eta: float |
| 162 | norm_threshold: float = 5.0 |
| 163 | momentum: float = 0.0 |
| 164 | # it is user's responsibility not to use same APGGuider for several denoisings or different modalities |
| 165 | # in order not to share accumulated average across different denoisings or modalities |
| 166 | running_avg: torch.Tensor | None = None |
| 167 | |
| 168 | def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor: |
| 169 | guidance = cond - uncond |
| 170 | if self.momentum != 0: |
| 171 | if self.running_avg is None: |
| 172 | self.running_avg = guidance.clone() |
| 173 | else: |
| 174 | self.running_avg = self.momentum * self.running_avg + guidance |
| 175 | guidance = self.running_avg |
| 176 | |
| 177 | if self.norm_threshold > 0: |
| 178 | ones = torch.ones_like(guidance) |
| 179 | guidance_norm = guidance.norm(p=2, dim=[-1, -2, -3], keepdim=True) |
| 180 | scale_factor = torch.minimum(ones, self.norm_threshold / guidance_norm) |
| 181 | guidance = guidance * scale_factor |
| 182 | |
| 183 | proj_coeff = projection_coef(guidance, cond) |
| 184 | g_parallel = proj_coeff * cond |
| 185 | g_orth = guidance - g_parallel |
| 186 | g_apg = g_parallel * self.eta + g_orth |
| 187 | |
| 188 | return g_apg * self.scale |
| 189 | |
| 190 | def enabled(self) -> bool: |
| 191 | return self.scale != 0.0 |
| 192 | |
| 193 | |
| 194 | @dataclass(frozen=True) |
| 195 | class MultiModalGuiderParams: |
| 196 | """ |
| 197 | Parameters for the multi-modal guider. |
| 198 | """ |
| 199 | |
| 200 | cfg_scale: float = 1.0 |
| 201 | "CFG (Classifier-free guidance) scale controlling how strongly the model adheres to the prompt." |
| 202 | stg_scale: float = 0.0 |
| 203 | "STG (Spatio-Temporal Guidance) scale controls how strongly the model reacts to the perturbation of the modality." |
| 204 | stg_blocks: list[int] | None = field(default_factory=list) |
| 205 | "Which transformer blocks to perturb for STG." |
| 206 | rescale_scale: float = 0.0 |
| 207 | "Rescale scale controlling how strongly the model rescales the modality after applying other guidance." |
| 208 | modality_scale: float = 1.0 |
| 209 | "Modality scale controlling how strongly the model reacts to the perturbation of the modality." |
| 210 | skip_step: int = 0 |
| 211 | "Skip step controlling how often the model skips the step." |
| 212 | |
| 213 | |
| 214 | def _params_for_sigma_from_sorted_dict( |
| 215 | sigma: float, params_by_sigma: Sequence[tuple[float, MultiModalGuiderParams]] |
| 216 | ) -> MultiModalGuiderParams: |
| 217 | """ |
| 218 | Return params for the given sigma from a sorted (sigma_upper_bound -> params) structure. |
| 219 | Keys are sorted descending (bin upper bounds). Bin i is (key_{i+1}, key_i]. |
| 220 | Get all keys >= sigma; use last in list (smallest such key = upper bound of bin containing sigma), |
| 221 | or last entry in the sequence if list is empty (sigma above max key). |
| 222 | """ |
| 223 | if not params_by_sigma: |
| 224 | raise ValueError("params_by_sigma must be non-empty") |
| 225 | sigma = float(sigma) |
| 226 | keys_desc = [k for k, _ in params_by_sigma] |
| 227 | keys_ge_sigma = [k for k in keys_desc if k >= sigma] |
| 228 | # sigma above all keys: use first bin (max key) |
| 229 | key = keys_ge_sigma[-1] if keys_ge_sigma else keys_desc[0] |
| 230 | return next(p for k, p in params_by_sigma if k == key) |
| 231 | |
| 232 | |
| 233 | @dataclass(frozen=True) |
| 234 | class MultiModalGuider: |
| 235 | """ |
| 236 | Multi-modal guider with constant params per instance. |
| 237 | For sigma-dependent params, use MultiModalGuiderFactory.build_from_sigma(sigma) to |
| 238 | obtain a guider for each step. |
| 239 | """ |
| 240 | |
| 241 | params: MultiModalGuiderParams |
| 242 | negative_context: torch.Tensor | None = None |
| 243 | |
| 244 | def calculate( |
| 245 | self, |
| 246 | cond: torch.Tensor, |
| 247 | uncond_text: torch.Tensor | float, |
| 248 | uncond_perturbed: torch.Tensor | float, |
| 249 | uncond_modality: torch.Tensor | float, |
| 250 | ) -> torch.Tensor: |
| 251 | """ |
| 252 | The guider calculates the guidance delta as (scale - 1) * (cond - uncond) for cfg and modality cfg, |
| 253 | and as scale * (cond - uncond) for stg, steering the denoising process away from the unconditioned |
| 254 | prediction. |
| 255 | """ |
| 256 | pred = ( |
| 257 | cond |
| 258 | + (self.params.cfg_scale - 1) * (cond - uncond_text) |
| 259 | + self.params.stg_scale * (cond - uncond_perturbed) |
| 260 | + (self.params.modality_scale - 1) * (cond - uncond_modality) |
| 261 | ) |
| 262 | |
| 263 | if self.params.rescale_scale != 0: |
| 264 | factor = cond.std() / pred.std() |
| 265 | factor = self.params.rescale_scale * factor + (1 - self.params.rescale_scale) |
| 266 | pred = pred * factor |
| 267 | |
| 268 | return pred |
| 269 | |
| 270 | def do_unconditional_generation(self) -> bool: |
| 271 | """Returns True if the guider is doing unconditional generation.""" |
| 272 | return not math.isclose(self.params.cfg_scale, 1.0) |
| 273 | |
| 274 | def do_perturbed_generation(self) -> bool: |
| 275 | """Returns True if the guider is doing perturbed generation.""" |
| 276 | return not math.isclose(self.params.stg_scale, 0.0) |
| 277 | |
| 278 | def do_isolated_modality_generation(self) -> bool: |
| 279 | """Returns True if the guider is doing isolated modality generation.""" |
| 280 | return not math.isclose(self.params.modality_scale, 1.0) |
| 281 | |
| 282 | def should_skip_step(self, step: int) -> bool: |
| 283 | """Returns True if the guider should skip the step.""" |
| 284 | if self.params.skip_step == 0: |
| 285 | return False |
| 286 | return step % (self.params.skip_step + 1) != 0 |
| 287 | |
| 288 | |
| 289 | @dataclass(frozen=True) |
| 290 | class MultiModalGuiderFactory: |
| 291 | """ |
| 292 | Factory that creates a MultiModalGuider for a given sigma. |
| 293 | Single source of truth: _params_by_sigma (schedule). Use constant() for |
| 294 | one params for all sigma, from_dict() for sigma-binned params. |
| 295 | """ |
| 296 | |
| 297 | negative_context: torch.Tensor | None = None |
| 298 | _params_by_sigma: tuple[tuple[float, MultiModalGuiderParams], ...] = () |
| 299 | |
| 300 | @classmethod |
| 301 | def constant( |
| 302 | cls, |
| 303 | params: MultiModalGuiderParams, |
| 304 | negative_context: torch.Tensor | None = None, |
| 305 | ) -> "MultiModalGuiderFactory": |
| 306 | """Build a factory with constant params (same guider for all sigma).""" |
| 307 | return cls( |
| 308 | negative_context=negative_context, |
| 309 | _params_by_sigma=((float("inf"), params),), |
| 310 | ) |
| 311 | |
| 312 | @classmethod |
| 313 | def from_dict( |
| 314 | cls, |
| 315 | sigma_to_params: Mapping[float, MultiModalGuiderParams], |
| 316 | negative_context: torch.Tensor | None = None, |
| 317 | ) -> "MultiModalGuiderFactory": |
| 318 | """ |
| 319 | Build a factory from a dict of sigma_value -> MultiModalGuiderParams. |
| 320 | Keys are sorted descending and used for bin lookup in params(sigma). |
| 321 | """ |
| 322 | if not sigma_to_params: |
| 323 | raise ValueError("sigma_to_params must be non-empty") |
| 324 | sorted_items = tuple(sorted(sigma_to_params.items(), key=lambda x: x[0], reverse=True)) |
| 325 | return cls(negative_context=negative_context, _params_by_sigma=sorted_items) |
| 326 | |
| 327 | def params(self, sigma: float | torch.Tensor) -> MultiModalGuiderParams: |
| 328 | """Return params effective for the given sigma (getter; single source of truth).""" |
| 329 | sigma_val = float(sigma.item() if isinstance(sigma, torch.Tensor) else sigma) |
| 330 | return _params_for_sigma_from_sorted_dict(sigma_val, self._params_by_sigma) |
| 331 | |
| 332 | def build_from_sigma(self, sigma: float | torch.Tensor) -> MultiModalGuider: |
| 333 | """Return a MultiModalGuider with params effective for the given sigma.""" |
| 334 | return MultiModalGuider( |
| 335 | params=self.params(sigma), |
| 336 | negative_context=self.negative_context, |
| 337 | ) |
| 338 | |
| 339 | |
| 340 | def create_multimodal_guider_factory( |
| 341 | params: MultiModalGuiderParams | MultiModalGuiderFactory, |
| 342 | negative_context: torch.Tensor | None = None, |
| 343 | ) -> MultiModalGuiderFactory: |
| 344 | """ |
| 345 | Create or return a MultiModalGuiderFactory. Pass constant params for a |
| 346 | single-params factory (uses MultiModalGuiderFactory.constant), or an existing |
| 347 | MultiModalGuiderFactory. When given a factory, returns it as-is unless |
| 348 | negative_context is provided. For sigma-dependent params use |
| 349 | MultiModalGuiderFactory.from_dict(...) and pass that as params. |
| 350 | """ |
| 351 | if isinstance(params, MultiModalGuiderFactory): |
| 352 | if negative_context is not None and params.negative_context is not negative_context: |
| 353 | return MultiModalGuiderFactory.from_dict(dict(params._params_by_sigma), negative_context=negative_context) |
| 354 | return params |
| 355 | return MultiModalGuiderFactory.constant(params, negative_context=negative_context) |
| 356 | |
| 357 | |
| 358 | def projection_coef(to_project: torch.Tensor, project_onto: torch.Tensor) -> torch.Tensor: |
| 359 | batch_size = to_project.shape[0] |
| 360 | positive_flat = to_project.reshape(batch_size, -1) |
| 361 | negative_flat = project_onto.reshape(batch_size, -1) |
| 362 | dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) |
| 363 | squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8 |
| 364 | return dot_product / squared_norm |
| 365 |