| 1 | import torch |
| 2 | |
| 3 | from ltx_core.components.protocols import DiffusionStepProtocol |
| 4 | from ltx_core.utils import to_velocity |
| 5 | |
| 6 | |
| 7 | class EulerDiffusionStep(DiffusionStepProtocol): |
| 8 | """ |
| 9 | First-order Euler method for diffusion sampling. |
| 10 | Takes a single step from the current noise level (sigma) to the next by |
| 11 | computing velocity from the denoised prediction and applying: sample + velocity * dt. |
| 12 | """ |
| 13 | |
| 14 | def step( |
| 15 | self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int, **_kwargs |
| 16 | ) -> torch.Tensor: |
| 17 | sigma = sigmas[step_index] |
| 18 | sigma_next = sigmas[step_index + 1] |
| 19 | dt = sigma_next - sigma |
| 20 | velocity = to_velocity(sample, sigma, denoised_sample) |
| 21 | |
| 22 | return (sample.to(torch.float32) + velocity.to(torch.float32) * dt).to(sample.dtype) |
| 23 | |
| 24 | |
| 25 | class Res2sDiffusionStep(DiffusionStepProtocol): |
| 26 | """ |
| 27 | Second-order diffusion step for res_2s sampling with SDE noise injection. |
| 28 | Used by the res_2s denoising loop. Advances the sample from the current |
| 29 | sigma to the next by mixing a deterministic update (from the denoised |
| 30 | prediction) with injected noise via ``get_sde_coeff``, producing |
| 31 | variance-preserving transitions. |
| 32 | """ |
| 33 | |
| 34 | @staticmethod |
| 35 | def get_sde_coeff( |
| 36 | sigma_next: torch.Tensor, |
| 37 | sigma_up: torch.Tensor | None = None, |
| 38 | sigma_down: torch.Tensor | None = None, |
| 39 | sigma_max: torch.Tensor | None = None, |
| 40 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 41 | """ |
| 42 | Compute SDE coefficients (alpha_ratio, sigma_down, sigma_up) for the step. |
| 43 | Given either ``sigma_down`` or ``sigma_up``, returns the mixing |
| 44 | coefficients used for variance-preserving noise injection. If |
| 45 | ``sigma_up`` is provided, ``sigma_down`` and ``alpha_ratio`` are |
| 46 | derived; if ``sigma_down`` is provided, ``sigma_up`` and |
| 47 | ``alpha_ratio`` are derived. |
| 48 | """ |
| 49 | if sigma_down is not None: |
| 50 | alpha_ratio = (1 - sigma_next) / (1 - sigma_down) |
| 51 | sigma_up = (sigma_next**2 - sigma_down**2 * alpha_ratio**2).clamp(min=0) ** 0.5 |
| 52 | elif sigma_up is not None: |
| 53 | # Fallback to avoid sqrt(neg_num) |
| 54 | sigma_up.clamp_(max=sigma_next * 0.9999) |
| 55 | sigmax = sigma_max if sigma_max is not None else torch.ones_like(sigma_next) |
| 56 | sigma_signal = sigmax - sigma_next |
| 57 | sigma_residual = (sigma_next**2 - sigma_up**2).clamp(min=0) ** 0.5 |
| 58 | alpha_ratio = sigma_signal + sigma_residual |
| 59 | sigma_down = sigma_residual / alpha_ratio |
| 60 | else: |
| 61 | alpha_ratio = torch.ones_like(sigma_next) |
| 62 | sigma_down = sigma_next |
| 63 | sigma_up = torch.zeros_like(sigma_next) |
| 64 | |
| 65 | sigma_up = torch.nan_to_num(sigma_up if sigma_up is not None else torch.zeros_like(sigma_next), 0.0) |
| 66 | # Replace NaNs in sigma_down with corresponding sigma_next elements (float32) |
| 67 | nan_mask = torch.isnan(sigma_down) |
| 68 | sigma_down[nan_mask] = sigma_next[nan_mask].to(sigma_down.dtype) |
| 69 | alpha_ratio = torch.nan_to_num(alpha_ratio, 1.0) |
| 70 | |
| 71 | return alpha_ratio, sigma_down, sigma_up |
| 72 | |
| 73 | def step( |
| 74 | self, |
| 75 | sample: torch.Tensor, |
| 76 | denoised_sample: torch.Tensor, |
| 77 | sigmas: torch.Tensor, |
| 78 | step_index: int, |
| 79 | noise: torch.Tensor, |
| 80 | ) -> torch.Tensor: |
| 81 | """Advance one step with SDE noise injection via get_sde_coeff.""" |
| 82 | sigma = sigmas[step_index] |
| 83 | sigma_next = sigmas[step_index + 1] |
| 84 | alpha_ratio, sigma_down, sigma_up = self.get_sde_coeff(sigma_next, sigma_up=sigma_next * 0.5) |
| 85 | output_dtype = denoised_sample.dtype |
| 86 | if torch.any(sigma_up == 0) or torch.any(sigma_next == 0): |
| 87 | return denoised_sample |
| 88 | |
| 89 | # Extract epsilon prediction |
| 90 | eps_next = (sample - denoised_sample) / (sigma - sigma_next) |
| 91 | denoised_next = sample - sigma * eps_next |
| 92 | |
| 93 | # Mix deterministic and stochastic components |
| 94 | x_noised = alpha_ratio * (denoised_next + sigma_down * eps_next) + sigma_up * noise |
| 95 | return x_noised.to(output_dtype) |
| 96 |