def get_sampling_fn()

in jcm/sampling.py [0:0]

• 58 lines of code
• 9 McCabe index (conditional complexity)

def get_sampling_fn(config, sde, model, shape, eps=1e-3):
    """Create a sampling function.

    Args:
      config: A `ml_collections.ConfigDict` object that contains all configuration information.
      sde: A `sde_lib.SDE` object that represents the forward SDE.
      model: A `flax.linen.Module` object that represents the architecture of a time-dependent score-based model.
      shape: A sequence of integers representing the expected shape of a single sample.
      eps: A `float` number. The reverse-time SDE is only integrated to `eps` for numerical stability.

    Returns:
      A function that takes random states and a replicated training state and outputs samples with the
        trailing dimensions matching `shape`.
    """

    sampler_name = config.sampling.method
    # Probability flow ODE sampling with black-box ODE solvers
    if sampler_name.lower() == "ode":
        sampling_fn = get_ode_sampler(
            sde=sde,
            model=model,
            shape=shape,
            denoise=config.sampling.noise_removal,
            eps=eps,
        )
    # Predictor-Corrector sampling. Predictor-only and Corrector-only samplers are special cases.
    elif sampler_name.lower() == "pc":
        predictor = get_predictor(config.sampling.predictor.lower())
        corrector = get_corrector(config.sampling.corrector.lower())
        sampling_fn = get_pc_sampler(
            sde=sde,
            model=model,
            shape=shape,
            predictor=predictor,
            corrector=corrector,
            snr=config.sampling.snr,
            n_steps=config.sampling.n_steps_each,
            probability_flow=config.sampling.probability_flow,
            denoise=config.sampling.noise_removal,
            eps=eps,
        )
    elif sampler_name.lower() == "heun":
        sampling_fn = get_heun_sampler(
            sde=sde, model=model, shape=shape, denoise=config.sampling.denoise
        )
    elif sampler_name.lower() == "euler":
        sampling_fn = get_euler_sampler(
            sde=sde, model=model, shape=shape, denoise=config.sampling.denoise
        )
    elif sampler_name.lower() == "onestep":
        sampling_fn = get_onestep_sampler(
            config=config,
            sde=sde,
            model=model,
            shape=shape,
        )
    elif sampler_name.lower() == "seeded_sampler":
        sampling_fn = get_seeded_sampler(
            config=config,
            sde=sde,
            model=model,
            shape=shape,
        )
    elif sampler_name.lower() == "progressive_distillation":
        sampling_fn = get_progressive_distillation_sampler(
            config=config,
            sde=sde,
            model=model,
            shape=shape,
            denoise=config.sampling.denoise,
        )
    else:
        raise ValueError(f"Sampler name {sampler_name} unknown.")

    return sampling_fn