| |
|
|
| |
|
|
| """Ideogram 4.0 Euler flow-matching sampler. |
| |
| Faithful port of the reference ``scheduler.py`` + ``sampler_configs.py``. The |
| logit-normal schedule and step intervals are host-side math (identical to the |
| reference); only the per-step Euler update ``z <- z + v * (s - t)`` runs on |
| device. Asymmetric CFG (``v = gw*v_cond + (1-gw)*v_uncond``) is applied on device |
| by the pipeline; this module owns the schedule, the guidance-weight sequence, and |
| the device step. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
|
|
| import torch |
|
|
| import ttnn |
|
|
|
|
| @dataclass(frozen=True) |
| class LogitNormalSchedule: |
| """Reference LogitNormalSchedule: maps uniform t in [0,1] to a flow-matching time.""" |
|
|
| mean: float |
| std: float = 1.0 |
| logsnr_min: float = -15.0 |
| logsnr_max: float = 18.0 |
|
|
| def __call__(self, t: torch.Tensor) -> torch.Tensor: |
| t = t.to(torch.float64) |
| z = torch.special.ndtri(t) |
| y = self.mean + self.std * z |
| t_ = torch.special.expit(y) |
| t_ = 1 - t_ |
| t_min = 1.0 / (1 + math.exp(0.5 * self.logsnr_max)) |
| t_max = 1.0 / (1 + math.exp(0.5 * self.logsnr_min)) |
| return t_.clamp(t_min, t_max).to(torch.float32) |
|
|
|
|
| def get_schedule_for_resolution( |
| image_resolution: tuple[int, int], |
| known_resolution: tuple[int, int] = (512, 512), |
| known_mean: float = 1.0, |
| std: float = 1.0, |
| ) -> LogitNormalSchedule: |
| """Resolution-aware schedule: mean shifts by 0.5*log(num_pixels / known_pixels).""" |
| num_pixels = image_resolution[0] * image_resolution[1] |
| known_pixels = known_resolution[0] * known_resolution[1] |
| mean = known_mean + 0.5 * math.log(num_pixels / known_pixels) |
| return LogitNormalSchedule(mean=mean, std=std) |
|
|
|
|
| def make_step_intervals(num_steps: int) -> torch.Tensor: |
| return torch.linspace(0.0, 1.0, num_steps + 1, dtype=torch.float32) |
|
|
|
|
| @dataclass(frozen=True, kw_only=True) |
| class SamplerParameters: |
| num_steps: int |
| guidance_schedule: tuple[float, ...] |
| mu: float |
| std: float = 1.0 |
|
|
| def __post_init__(self) -> None: |
| if len(self.guidance_schedule) != self.num_steps: |
| raise ValueError( |
| f"guidance_schedule has length {len(self.guidance_schedule)}, expected num_steps={self.num_steps}" |
| ) |
|
|
|
|
| |
| PRESETS: dict[str, SamplerParameters] = { |
| "V4_QUALITY_48": SamplerParameters(num_steps=48, guidance_schedule=(3.0,) * 3 + (7.0,) * 45, mu=0.0, std=1.5), |
| "V4_DEFAULT_20": SamplerParameters(num_steps=20, guidance_schedule=(3.0,) * 2 + (7.0,) * 18, mu=0.0, std=1.75), |
| "V4_TURBO_12": SamplerParameters(num_steps=12, guidance_schedule=(3.0,) * 1 + (7.0,) * 11, mu=0.5, std=1.75), |
| } |
|
|
|
|
| class Ideogram4Sampler: |
| """Owns the logit-normal schedule, guidance-weight sequence, and the device Euler step. |
| |
| The denoising loop runs ``for i in reversed(range(num_steps))``; at step ``i`` the |
| flow-matching times are ``t = schedule(intervals[i+1])`` and ``s = schedule(intervals[i])`` |
| and the Euler update is ``z <- z + v * (s - t)``. |
| """ |
|
|
| def __init__(self, params: SamplerParameters, *, height: int, width: int) -> None: |
| self.params = params |
| self.schedule = get_schedule_for_resolution((height, width), known_mean=params.mu, std=params.std) |
| self.step_intervals = make_step_intervals(params.num_steps) |
|
|
| @classmethod |
| def from_preset(cls, name: str, *, height: int, width: int) -> "Ideogram4Sampler": |
| return cls(PRESETS[name], height=height, width=width) |
|
|
| @property |
| def num_steps(self) -> int: |
| return self.params.num_steps |
|
|
| def times_for_step(self, i: int) -> tuple[float, float]: |
| """Return (t, s) for loop index i — the current and next flow-matching times.""" |
| t_val = float(self.schedule(self.step_intervals[i + 1].unsqueeze(0)).item()) |
| s_val = float(self.schedule(self.step_intervals[i].unsqueeze(0)).item()) |
| return t_val, s_val |
|
|
| def guidance_weight(self, i: int) -> float: |
| return float(self.params.guidance_schedule[i]) |
|
|
| def step(self, z: ttnn.Tensor, velocity: ttnn.Tensor, i: int | None = None, *, scale=None) -> ttnn.Tensor: |
| """Device Euler update: z <- z + velocity * (s - t). |
| |
| ``scale`` is the step size ``s - t``. When omitted it is computed from ``times_for_step(i)`` |
| (Python float — used by the standalone sampler test). The traced pipeline passes ``scale`` |
| as a 1-element DEVICE tensor instead (so the same scalar can be re-patched per step inside |
| one reused trace, where a Python float would be baked in) — routing production through here. |
| """ |
| if scale is None: |
| t_val, s_val = self.times_for_step(i) |
| scale = s_val - t_val |
| return z + velocity * scale |
|
|
