135 lines
5.7 KiB
Python
135 lines
5.7 KiB
Python
"""Per-step latent reporting for sample generation (inference engine).
|
|
|
|
A holder exposes `sample_step_hook(step_index, num_steps, latents)`; while it
|
|
is set, the sampling scheduler's `step` is wrapped on the instance so every
|
|
step reaches the hook without per-pipeline code. The latent handed to the
|
|
hook is the step's predicted CLEAN sample (x0) — what a preview should show
|
|
— falling back to the scheduler's next sample when x0 cannot be derived. Shared by
|
|
BaseModel and the legacy StableDiffusion monolith.
|
|
"""
|
|
|
|
from typing import Callable, List
|
|
|
|
|
|
def install_sample_step_hooks(holder, pipeline) -> Callable[[], None]:
|
|
"""Wrap scheduler.step on the sampling scheduler(s). Returns an unwrap
|
|
callable. No-op without a hook. Instance-level wrap: the class method is
|
|
untouched and a stale wrapper is replaced rather than stacked."""
|
|
if getattr(holder, "sample_step_hook", None) is None:
|
|
return lambda: None
|
|
schedulers: List = []
|
|
for cand in (
|
|
getattr(pipeline, "scheduler", None),
|
|
getattr(getattr(holder, "pipeline", None), "scheduler", None),
|
|
getattr(holder, "noise_scheduler", None),
|
|
):
|
|
if cand is not None and hasattr(cand, "step") and all(cand is not s for s in schedulers):
|
|
schedulers.append(cand)
|
|
unwraps = []
|
|
for scheduler in schedulers:
|
|
current = scheduler.step
|
|
orig = getattr(current, "_aitk_orig_step", None) or current
|
|
|
|
def make(orig, scheduler):
|
|
def step(*args, **kwargs):
|
|
# capture the pre-step state: x0 is derived from (model_output,
|
|
# timestep, sample) with the sigma/alpha of THIS step, which
|
|
# the scheduler advances past inside orig()
|
|
model_output = args[0] if len(args) > 0 else kwargs.get("model_output")
|
|
timestep = args[1] if len(args) > 1 else kwargs.get("timestep")
|
|
sample_in = args[2] if len(args) > 2 else kwargs.get("sample")
|
|
coeffs = _x0_coefficients(scheduler, timestep)
|
|
out = orig(*args, **kwargs)
|
|
if isinstance(out, tuple):
|
|
prev_sample = out[0]
|
|
pred_x0 = None
|
|
else:
|
|
prev_sample = getattr(out, "prev_sample", None)
|
|
pred_x0 = getattr(out, "pred_original_sample", None)
|
|
hook = getattr(holder, "sample_step_hook", None)
|
|
if hook is not None and prev_sample is not None:
|
|
if pred_x0 is None:
|
|
pred_x0 = _predict_x0(coeffs, model_output, sample_in)
|
|
timesteps = getattr(scheduler, "timesteps", None)
|
|
num_steps = len(timesteps) if timesteps is not None else None
|
|
idx = getattr(holder, "_sample_step_index", 0)
|
|
holder._sample_step_index = idx + 1
|
|
hook(idx, num_steps, pred_x0 if pred_x0 is not None else prev_sample)
|
|
return out
|
|
|
|
step._aitk_orig_step = orig
|
|
return step
|
|
|
|
scheduler.step = make(orig, scheduler)
|
|
|
|
def unwrap(scheduler=scheduler):
|
|
if "step" in scheduler.__dict__:
|
|
del scheduler.__dict__["step"]
|
|
|
|
unwraps.append(unwrap)
|
|
return lambda: [u() for u in unwraps]
|
|
|
|
|
|
def _x0_coefficients(scheduler, timestep):
|
|
"""Per-step numbers needed to turn the model output into a clean-sample
|
|
estimate, read BEFORE scheduler.step advances its index.
|
|
Returns ("flow", sigma) | ("ddpm", prediction_type, alpha_bar) | None."""
|
|
try:
|
|
import torch
|
|
|
|
sigmas = getattr(scheduler, "sigmas", None)
|
|
if sigmas is not None and hasattr(scheduler, "index_for_timestep"):
|
|
idx = getattr(scheduler, "step_index", None)
|
|
if idx is None:
|
|
idx = scheduler.index_for_timestep(timestep)
|
|
return ("flow", float(sigmas[idx]))
|
|
alphas_cumprod = getattr(scheduler, "alphas_cumprod", None)
|
|
if alphas_cumprod is not None:
|
|
t = timestep
|
|
if isinstance(t, torch.Tensor):
|
|
t = int(t.flatten()[0].item())
|
|
prediction_type = getattr(getattr(scheduler, "config", None), "prediction_type", "epsilon")
|
|
return ("ddpm", prediction_type, float(alphas_cumprod[int(t)]))
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
|
|
def _predict_x0(coeffs, model_output, sample):
|
|
"""Clean-sample estimate for the preview: flow matching x0 = x_t - sigma*v;
|
|
DDPM-style eps / v / sample predictions via alpha_bar. None when unknown."""
|
|
if coeffs is None or model_output is None or sample is None:
|
|
return None
|
|
try:
|
|
x = sample.tensor if hasattr(sample, "tensor") else sample
|
|
v = model_output.tensor if hasattr(model_output, "tensor") else model_output
|
|
if v.shape != x.shape:
|
|
return None
|
|
v = v.to(x.dtype)
|
|
if coeffs[0] == "flow":
|
|
sigma = coeffs[1]
|
|
return x - sigma * v
|
|
_, prediction_type, alpha_bar = coeffs
|
|
a = alpha_bar**0.5
|
|
b = (1.0 - alpha_bar) ** 0.5
|
|
if prediction_type == "epsilon":
|
|
return (x - b * v) / a
|
|
if prediction_type == "v_prediction":
|
|
return a * x - b * v
|
|
if prediction_type == "sample":
|
|
return v
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
|
|
def emit_sample_step(holder, latents, step_index=None, num_steps=None):
|
|
"""For sampling loops that bypass scheduler.step: report one denoised
|
|
latent to the holder's hook (no-op when unset)."""
|
|
hook = getattr(holder, "sample_step_hook", None)
|
|
if hook is None:
|
|
return
|
|
idx = getattr(holder, "_sample_step_index", 0) if step_index is None else step_index
|
|
holder._sample_step_index = idx + 1
|
|
hook(idx, num_steps, latents)
|