def denoiser_distiller_fn()

in jcm/models/utils.py [0:0]

• 26 lines of code
• 4 McCabe index (conditional complexity)

    def denoiser_distiller_fn(x, t, rng=None):
        in_x = batch_mul(x, 1 / jnp.sqrt(t**2 + sde.data_std**2))
        cond_t = 0.25 * jnp.log(t)
        if isinstance(model, NCSNpp):
            model_output, state = model_fn(in_x, cond_t, rng)
            denoiser = model_output[..., :3]
            distiller = model_output[..., 3:]
        elif isinstance(model, JointNCSNpp):
            (denoiser, distiller), state = model_fn(in_x, cond_t, rng)

        denoiser = batch_mul(
            denoiser, t * sde.data_std / jnp.sqrt(t**2 + sde.data_std**2)
        )
        skip_x = batch_mul(x, sde.data_std**2 / (t**2 + sde.data_std**2))
        denoiser = skip_x + denoiser

        distiller = batch_mul(
            distiller,
            (t - pred_t) * sde.data_std / jnp.sqrt(t**2 + sde.data_std**2),
        )
        skip_x = batch_mul(
            x, sde.data_std**2 / ((t - pred_t) ** 2 + sde.data_std**2)
        )
        distiller = skip_x + distiller

        if return_state:
            return (denoiser, distiller), state
        else:
            return denoiser, distiller