Create a list of timesteps to use from an original diffusion process, given the number of timesteps we want to take from equally-sized portio

Args: x1: shape (bs, *dim), represents the target minibatch (data) x0: shape (bs, *dim), represents the source miniba

x_start = x_0

Update the reweighting using losses from a model. Sub-classes should override this method to update the reweighting using los

Get a numpy array of weights, one per diffusion step. The weights needn't be normalized, but must be positive.

Assumes target and pred in [-1, 1] range

Assumes target and pred in [-1, 1] range

Sampler class for the FlowModel

Logit-Normal sampler from the paper 'Scaling Rectified Flow Transformers for High-Resolution Image Synthesis' - Esser et al.

Flow Matching, Stochastic Interpolants, or Rectified Flow model. :) Args: net: Neural network that takes in x an

Initialize pretrained TAESD on the given device from the given checkpoints.

Customized iteration-wise exponential scheduler. Re-calculate for every step, to reduce error accumulation Args:

Args: fm_cfg: Flow matching model config. noising_step: Forward diffusion noising step with linear schedule

Compute the entire variational lower-bound, measured in bits-per-dim, as well as other related quantities. Returns:

Compute the entire variational lower-bound, measured in bits-per-dim, as well as other related quantities. :param model: the

Colorize a depth map using a matplotlib colormap. Args: depth: Depth tensor of shape (b, 1, h, w) or (b, h, w) with

Special purposed function for computing numerical stabled d_alpha_t / alpha_t