| |
|
|
| |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
| from enum import Enum |
| from typing import TYPE_CHECKING, Any |
|
|
| import ttnn |
|
|
| from .base import Solver |
| from .schedule import Schedule |
|
|
| if TYPE_CHECKING: |
| from collections.abc import Sequence |
|
|
| from diffusers.schedulers import UniPCMultistepScheduler |
|
|
|
|
| @dataclass(frozen=True) |
| class _State: |
| clean_preds: tuple[ttnn.Tensor, ...] |
| corrected: ttnn.Tensor |
| oldest_idx: int = 0 |
|
|
|
|
| class UniPCVariant(Enum): |
| BH1 = "bh1" |
| BH2 = "bh2" |
|
|
| def b(self, h: float) -> float: |
| if self is UniPCVariant.BH1: |
| return h |
| return -math.expm1(-h) |
|
|
|
|
| class UniPCSolver(Solver): |
| def __init__(self, *, order: int, variant: UniPCVariant, scheduler: UniPCMultistepScheduler) -> None: |
| if order not in (1, 2): |
| msg = f"only order 1 and 2 are supported, got {order}" |
| raise ValueError(msg) |
|
|
| super().__init__(scheduler=scheduler) |
|
|
| self.order = order |
| self.variant = variant |
| self._state = None |
| self._default_shift = scheduler.config.flow_shift |
|
|
| def set_schedule( |
| self, num_inference_steps: int | None = None, *, shift: float | None = None, **kwargs: Any |
| ) -> None: |
| """Derive and adopt the schedule for one denoising run. |
| |
| Args: |
| num_inference_steps: Number of denoising steps. |
| shift: Flow shift for this run only; the scheduler's construction-time value |
| is used when omitted. |
| kwargs: Forwarded to the scheduler's `set_timesteps`. |
| """ |
| |
| self._scheduler.register_to_config(flow_shift=self._default_shift if shift is None else shift) |
| self._schedule = Schedule.from_scheduler(self._scheduler, num_inference_steps, **kwargs) |
|
|
| if self._state is not None: |
| self._state = _State(self._state.clean_preds, self._state.corrected, 0) |
|
|
| def step(self, *, step: int, latent: ttnn.Tensor, velocity_pred: ttnn.Tensor) -> ttnn.Tensor: |
| clean_pred = latent - self.sigmas[step] * velocity_pred |
|
|
| state = self._state or _State( |
| tuple(ttnn.empty_like(latent) for _ in range(self.order)), |
| ttnn.empty_like(latent), |
| ) |
| clean_preds = _ordered_clean_preds(state.clean_preds, state.oldest_idx) |
|
|
| if step != 0: |
| corrected = self._correct( |
| order=_taper(self.order, step - 1, len(self.sigmas) - 1), |
| latent=state.corrected, |
| step=step - 1, |
| clean_preds=(*clean_preds, clean_pred), |
| ) |
| else: |
| corrected = latent |
|
|
| ttnn.copy(corrected, state.corrected) |
| del corrected |
|
|
| ttnn.copy(clean_pred, state.clean_preds[state.oldest_idx]) |
| oldest_idx = (state.oldest_idx + 1) % self.order |
| clean_preds = _ordered_clean_preds(state.clean_preds, oldest_idx) |
| del clean_pred |
|
|
| predicted = self._predict( |
| order=_taper(self.order, step, len(self.sigmas) - 1), |
| latent=state.corrected, |
| step=step, |
| clean_preds=clean_preds, |
| ) |
|
|
| self._state = _State(state.clean_preds, state.corrected, oldest_idx) |
| return predicted |
|
|
| def _predict( |
| self, *, order: int, latent: ttnn.Tensor, step: int, clean_preds: Sequence[ttnn.Tensor] |
| ) -> ttnn.Tensor: |
| sigma_curr, sigma_next = self.sigmas[step : step + 2] |
| alpha_curr, alpha_next = self.alphas[step : step + 2] |
|
|
| lam_curr = _log(alpha_curr) - _log(sigma_curr) |
| lam_next = _log(alpha_next) - _log(sigma_next) |
| h = lam_next - lam_curr |
|
|
| coeff_latent = sigma_next / sigma_curr |
| coeff_curr = -alpha_next * math.expm1(-h) |
|
|
| latent = coeff_latent * latent + coeff_curr * clean_preds[-1] |
|
|
| if order == 1: |
| return latent |
|
|
| lam_prev = _log(self.alphas[step - 1]) - _log(self.sigmas[step - 1]) |
| r = (lam_curr - lam_prev) / h |
| w = alpha_next * self.variant.b(h) * 0.5 / r |
|
|
| return latent + w * (clean_preds[-1] - clean_preds[-2]) |
|
|
| def _correct( |
| self, *, order: int, latent: ttnn.Tensor, step: int, clean_preds: Sequence[ttnn.Tensor] |
| ) -> ttnn.Tensor: |
| sigma_curr, sigma_next = self.sigmas[step : step + 2] |
| alpha_curr, alpha_next = self.alphas[step : step + 2] |
|
|
| lam_curr = _log(alpha_curr) - _log(sigma_curr) |
| lam_next = _log(alpha_next) - _log(sigma_next) |
| h = lam_next - lam_curr |
|
|
| coeff_latent = sigma_next / sigma_curr |
| coeff_clean = -alpha_next * math.expm1(-h) |
|
|
| latent = coeff_latent * latent + coeff_clean * clean_preds[-2] |
|
|
| if order == 1: |
| |
| w = alpha_next * self.variant.b(h) * 0.5 |
| return latent + w * (clean_preds[-1] - clean_preds[-2]) |
|
|
| |
| exp_neg_h = math.expm1(-h) |
|
|
| lam_prev = _log(self.alphas[step - 1]) - _log(self.sigmas[step - 1]) |
| r_1 = (lam_curr - lam_prev) / h |
|
|
| g1 = (h + exp_neg_h) / h**2 |
| g2 = (h**2 - 2 * h - 2 * exp_neg_h) / h**2 |
|
|
| if math.isinf(r_1): |
| w_prev = 0.0 |
| w_pred = alpha_next * h * g1 |
| else: |
| det = h * (1 + r_1) |
| c_1 = (h * g1 - g2) / det |
| c_2 = (g2 + r_1 * h * g1) / det |
|
|
| w_prev = alpha_next * h * c_1 / r_1 |
| w_pred = alpha_next * h * c_2 |
|
|
| return latent + w_prev * (clean_preds[-2] - clean_preds[-3]) + w_pred * (clean_preds[-1] - clean_preds[-2]) |
|
|
|
|
| def _taper(order: int, step: int, num_steps: int) -> int: |
| return min(order, step + 1, num_steps - step) |
|
|
|
|
| def _ordered_clean_preds(clean_preds: tuple[ttnn.Tensor, ...], oldest_idx: int) -> tuple[ttnn.Tensor, ...]: |
| return clean_preds[oldest_idx:] + clean_preds[:oldest_idx] |
|
|
|
|
| def _log(x: float, /) -> float: |
| return math.log(x) if x != 0 else -math.inf |
|
|
