def get_progressive_distillation_sampler()

in jcm/sampling.py [0:0]

• 32 lines of code
• 2 McCabe index (conditional complexity)

def get_progressive_distillation_sampler(config, sde, model, shape, denoise=True):
    ema_scales_fn = get_ema_scales_fn(config)

    def progressive_distillation_sampler(rng, state):
        denoiser_fn = mutils.get_denoiser_fn(
            sde, model, state.params_ema, state.model_state, train=False
        )
        _, num_scales = ema_scales_fn(state.step)

        rng = hk.PRNGSequence(rng)
        x = sde.prior_sampling(next(rng), shape)

        t_start = sde.t_max ** (1 / sde.rho)
        t_end = sde.t_min ** (1 / sde.rho)

        def loop_body(i, val):
            x = val
            t = (t_start + i / num_scales * (t_end - t_start)) ** sde.rho
            vec_t = jnp.ones((shape[0],)) * t
            denoiser = denoiser_fn(x, vec_t)
            d = 1 / t * x - 1 / t * denoiser
            next_t = (t_start + (i + 1) / num_scales * (t_end - t_start)) ** sde.rho
            samples = x + (next_t - t) * d
            return samples

        x = jax.lax.fori_loop(0, num_scales, loop_body, x)
        if denoise:
            t = sde.t_min
            vec_t = jnp.ones((shape[0],)) * t
            denoiser = denoiser_fn(x, vec_t)
            d = 1 / t * x - 1 / t * denoiser
            next_t = 0.0
            samples = x + (next_t - t) * d
        else:
            samples = x
        return samples, num_scales

    return jax.pmap(progressive_distillation_sampler, axis_name="batch")