| |
|
|
| |
|
|
| from __future__ import annotations |
|
|
| from abc import ABC, abstractmethod |
| from typing import Any |
|
|
| import ttnn |
|
|
| from .schedule import Schedule |
|
|
|
|
| class Solver(ABC): |
| """On-device solver for one denoising trajectory. |
| |
| A solver owns the scheduler defining its discretization (where it has one) and the |
| `Schedule` derived from it, so the schedule has a single owner across a run. |
| """ |
|
|
| def __init__(self, *, scheduler: Any = None) -> None: |
| self._scheduler = scheduler |
| self._schedule = None |
|
|
| @property |
| def scheduler(self) -> Any: |
| """The scheduler this solver was built from, or `None` if it has none.""" |
| return self._scheduler |
|
|
| @property |
| def schedule(self) -> Schedule: |
| if self._schedule is None: |
| msg = "schedule must be set before stepping" |
| raise ValueError(msg) |
| return self._schedule |
|
|
| @property |
| def sigmas(self) -> tuple[float, ...]: |
| return self.schedule.sigmas |
|
|
| @property |
| def alphas(self) -> tuple[float, ...]: |
| return self.schedule.alphas |
|
|
| @property |
| def timesteps(self) -> tuple[float, ...]: |
| return self.schedule.timesteps |
|
|
| @abstractmethod |
| 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 construction-time value is used when |
| omitted. Each solver maps this onto its own scheduler's spelling. |
| kwargs: Solver-specific schedule arguments. |
| """ |
|
|
| @abstractmethod |
| def step(self, *, step: int, latent: ttnn.Tensor, velocity_pred: ttnn.Tensor) -> ttnn.Tensor: |
| """Advance the latent one step toward the clean data. |
| |
| Args: |
| step: Current step index into the sigmas/alphas schedule. |
| latent: Noisy latent at the current step. |
| velocity_pred: Predicted velocity at the current step. |
| |
| Returns: |
| The predicted latent at the next step. |
| """ |
|
|
