utils.py 🚧

Source: LION/models/diffusion/NCSNpp_helpers/utils.py

Warning

This file has not yet received a complete narrative and docstring audit. Its public source-level API is listed automatically below.

All functions and modules related to model definition.

register_model(cls=None, *, name=None)

A decorator for registering model classes.

get_model(name)

No docstring is available.

get_sigmas(params)

Get sigmas — the set of noise levels for SMLD from config files. Args: config: A ConfigDict object parsed from the config file Returns: sigmas: a jax numpy arrary of noise levels

get_ddpm_params(config)

Get betas and alphas — parameters used in the original DDPM paper.

create_model(config)

Create the score model.

get_model_fn(model, train=False)

Create a function to give the output of the score-based model.

Args: model: The score model. train: True for training and False for evaluation.

Returns: A model function.

get_score_fn(sde, model, train=False, continuous=False)

Wraps score_fn so that the model output corresponds to a real time-dependent score function.

Args: sde: An sde_lib.SDE object that represents the forward SDE. model: A score model. train: True for training and False for evaluation. continuous: If True, the score-based model is expected to directly take continuous time steps.

Returns: A score function.

to_flattened_numpy(x)

Flatten a torch tensor x and convert it to numpy.

from_flattened_numpy(x, shape)

Form a torch tensor with the given shape from a flattened numpy array x.