from typing import * import torch import numpy as np from tqdm import tqdm from easydict import EasyDict as edict from .base import Sampler from .classifier_free_guidance_mixin import ClassifierFreeGuidanceSamplerMixin from .guidance_interval_mixin import GuidanceIntervalSamplerMixin class FlowEulerSampler(Sampler): """ Generate samples from a flow-matching model using Euler sampling. Args: sigma_min: The minimum scale of noise in flow. """ def __init__( self, sigma_min: float, ): self.sigma_min = sigma_min def _eps_to_xstart(self, x_t, t, eps): assert x_t.shape == eps.shape return (x_t - (self.sigma_min + (1 - self.sigma_min) * t) * eps) / (1 - t) def _xstart_to_eps(self, x_t, t, x_0): assert x_t.shape == x_0.shape return (x_t - (1 - t) * x_0) / (self.sigma_min + (1 - self.sigma_min) * t) def _v_to_xstart_eps(self, x_t, t, v): assert x_t.shape == v.shape eps = (1 - t) * v + x_t x_0 = (1 - self.sigma_min) * x_t - (self.sigma_min + (1 - self.sigma_min) * t) * v return x_0, eps def _inference_model(self, model, x_t, t, cond=None, **kwargs): t = torch.tensor([1000 * t] * x_t.shape[0], device=x_t.device, dtype=torch.float32) return model(x_t, t, cond, **kwargs) def _get_model_prediction(self, model, x_t, t, cond=None, **kwargs): pred_v = self._inference_model(model, x_t, t, cond, **kwargs) pred_x_0, pred_eps = self._v_to_xstart_eps(x_t=x_t, t=t, v=pred_v) return pred_x_0, pred_eps, pred_v @torch.no_grad() def sample_once( self, model, x_t, t: float, t_prev: float, cond: Optional[Any] = None, **kwargs ): """ Sample x_{t-1} from the model using Euler method. Args: model: The model to sample from. x_t: The [N x C x ...] tensor of noisy inputs at time t. t: The current timestep. t_prev: The previous timestep. cond: conditional information. **kwargs: Additional arguments for model inference. Returns: a dict containing the following - 'pred_x_prev': x_{t-1}. - 'pred_x_0': a prediction of x_0. """ pred_x_0, pred_eps, pred_v = self._get_model_prediction(model, x_t, t, cond, **kwargs) pred_x_prev = x_t - (t - t_prev) * pred_v return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0}) @torch.no_grad() def sample( self, model, noise, cond: Optional[Any] = None, steps: int = 50, rescale_t: float = 1.0, verbose: bool = True, **kwargs ): """ Generate samples from the model using Euler method. Args: model: The model to sample from. noise: The initial noise tensor. cond: conditional information. steps: The number of steps to sample. rescale_t: The rescale factor for t. verbose: If True, show a progress bar. **kwargs: Additional arguments for model_inference. Returns: a dict containing the following - 'samples': the model samples. - 'pred_x_t': a list of prediction of x_t. - 'pred_x_0': a list of prediction of x_0. """ sample = noise t_seq = np.linspace(1, 0, steps + 1) t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq) t_pairs = list((t_seq[i], t_seq[i + 1]) for i in range(steps)) ret = edict({"samples": None, "pred_x_t": [], "pred_x_0": []}) for t, t_prev in tqdm(t_pairs, desc="Sampling", disable=not verbose): out = self.sample_once(model, sample, t, t_prev, cond, **kwargs) sample = out.pred_x_prev ret.pred_x_t.append(out.pred_x_prev) ret.pred_x_0.append(out.pred_x_0) ret.samples = sample return ret class FlowEulerCfgSampler(ClassifierFreeGuidanceSamplerMixin, FlowEulerSampler): """ Generate samples from a flow-matching model using Euler sampling with classifier-free guidance. """ @torch.no_grad() def sample( self, model, noise, cond, neg_cond, steps: int = 50, rescale_t: float = 1.0, cfg_strength: float = 3.0, verbose: bool = True, **kwargs ): """ Generate samples from the model using Euler method. Args: model: The model to sample from. noise: The initial noise tensor. cond: conditional information. neg_cond: negative conditional information. steps: The number of steps to sample. rescale_t: The rescale factor for t. cfg_strength: The strength of classifier-free guidance. verbose: If True, show a progress bar. **kwargs: Additional arguments for model_inference. Returns: a dict containing the following - 'samples': the model samples. - 'pred_x_t': a list of prediction of x_t. - 'pred_x_0': a list of prediction of x_0. """ return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, **kwargs) class FlowEulerGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, FlowEulerSampler): """ Generate samples from a flow-matching model using Euler sampling with classifier-free guidance and interval. """ @torch.no_grad() def sample( self, model, noise, cond, neg_cond, steps: int = 50, rescale_t: float = 1.0, cfg_strength: float = 3.0, cfg_interval: Tuple[float, float] = (0.0, 1.0), verbose: bool = True, **kwargs ): """ Generate samples from the model using Euler method. Args: model: The model to sample from. noise: The initial noise tensor. cond: conditional information. neg_cond: negative conditional information. steps: The number of steps to sample. rescale_t: The rescale factor for t. cfg_strength: The strength of classifier-free guidance. cfg_interval: The interval for classifier-free guidance. verbose: If True, show a progress bar. **kwargs: Additional arguments for model_inference. Returns: a dict containing the following - 'samples': the model samples. - 'pred_x_t': a list of prediction of x_t. - 'pred_x_0': a list of prediction of x_0. """ return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, cfg_interval=cfg_interval, **kwargs)