Implement TCD, avoid repeated allocation in DDIM, implement eta parameter for DDIM and TCD, minor comment clarification

Author: yslai

denoiser.hpp

@@ -474,7 +474,8 @@ static void sample_k_diffusion(sample_method_t method,
                                ggml_context* work_ctx,
                                ggml_tensor* x,
                                std::vector<float> sigmas,
-                               std::shared_ptr<RNG> rng) {
+                               std::shared_ptr<RNG> rng,
+                               float eta) {
     size_t steps = sigmas.size() - 1;
     // sample_euler_ancestral
     switch (method) {
@@ -1014,6 +1015,8 @@ static void sample_k_diffusion(sample_method_t method,
             // structure hides from the denoiser), and the sigmas are
             // also needed to invert the behavior of CompVisDenoiser
             // (k-diffusion's LMSDiscreteScheduler)
+            float beta_start = 0.00085f;
+            float beta_end = 0.0120f;
             std::vector<double> alphas_cumprod;
             std::vector<double> compvis_sigmas;
 
@@ -1023,21 +1026,41 @@ static void sample_k_diffusion(sample_method_t method,
                 alphas_cumprod[i] =
                     (i == 0 ? 1.0f : alphas_cumprod[i - 1]) *
                     (1.0f -
-                     std::pow(sqrtf(0.00085f) +
-                              (sqrtf(0.0120f) - sqrtf(0.00085f)) *
+                     std::pow(sqrtf(beta_start) +
+                              (sqrtf(beta_end) - sqrtf(beta_start)) *
                               ((float)i / (TIMESTEPS - 1)), 2));
                 compvis_sigmas[i] =
                     std::sqrt((1 - alphas_cumprod[i]) /
                               alphas_cumprod[i]);
             }
+
+            struct ggml_tensor* pred_original_sample =
+                ggml_dup_tensor(work_ctx, x);
+            struct ggml_tensor* variance_noise =
+                ggml_dup_tensor(work_ctx, x);
+
             for (int i = 0; i < steps; i++) {
                 // The "trailing" DDIM timestep, see S. Lin et al.,
                 // "Common Diffusion Noise Schedules and Sample Steps
                 // are Flawed", arXiv:2305.08891 [cs], p. 4, Table
-                // 2. Most variables below follow Diffusers naming.
+                // 2. Most variables below follow Diffusers naming
+                //
+                // Diffuser naming vs. J. Song et al., "Denoising
+                // Diffusion Implicit Models", arXiv:2010.02502, p. 5,
+                // (12) and p. 16, (16) (<variable name> -> <name in
+                // paper>):
+                //
+                // - pred_noise_t -> epsilon_theta^(t)(x_t)
+                // - pred_original_sample -> f_theta^(t)(x_t) or x_0
+                // - std_dev_t -> sigma_t (not the LMS sigma)
+                // - eta -> eta (set to 0 at the moment)
+                // - pred_sample_direction -> "direction pointing to
+                //   x_t"
+                // - pred_prev_sample -> "x_t-1"
                 int timestep =
                     roundf(TIMESTEPS -
                            i * ((float)TIMESTEPS / steps)) - 1;
+                // 1. get previous step value (=t-1)
                 int prev_timestep = timestep - TIMESTEPS / steps;
                 // The sigma here is chosen to cause the
                 // CompVisDenoiser to produce t = timestep
@@ -1066,51 +1089,53 @@ static void sample_k_diffusion(sample_method_t method,
                 }
                 else {
                     // For the subsequent steps after the first one,
-                    // at this point x = latents (pipeline) or x =
-                    // sample (scheduler), and needs to be prescaled
-                    // with x <- latents / c_in to compensate for
-                    // model() applying the scale c_in before the
-                    // U-net F_theta
+                    // at this point x = latents or x = sample, and
+                    // needs to be prescaled with x <- sample / c_in
+                    // to compensate for model() applying the scale
+                    // c_in before the U-net F_theta
                     float* vec_x = (float*)x->data;
                     for (int j = 0; j < ggml_nelements(x); j++) {
                         vec_x[j] *= std::sqrt(sigma * sigma + 1);
                     }
                 }
-                // Note model() is the D(x, sigma) as defined in
-                // T. Karras et al., arXiv:2206.00364, p. 3, Table 1
-                // and p. 8 (7)
-                struct ggml_tensor* noise_pred =
+                // Note (also noise_pred in Diffuser's pipeline)
+                // model_output = model() is the D(x, sigma) as
+                // defined in T. Karras et al., arXiv:2206.00364,
+                // p. 3, Table 1 and p. 8 (7), compare also p. 38
+                // (226) therein.
+                struct ggml_tensor* model_output =
                     model(x, sigma, i + 1);
-                // Here noise_pred is still the k-diffusion denoiser
+                // Here model_output is still the k-diffusion denoiser
                 // output, not the U-net output F_theta(c_in(sigma) x;
                 // ...) in Karras et al. (2022), whereas Diffusers'
-                // noise_pred is F_theta(...). Recover the actual
-                // noise_pred, which is also referred to as the
+                // model_output is F_theta(...). Recover the actual
+                // model_output, which is also referred to as the
                 // "Karras ODE derivative" d or d_cur in several
                 // samplers above.
                 {
                     float* vec_x = (float*)x->data;
-                    float* vec_noise_pred = (float*)noise_pred->data;
+                    float* vec_model_output =
+                        (float*)model_output->data;
                     for (int j = 0; j < ggml_nelements(x); j++) {
-                        vec_noise_pred[j] =
-                            (vec_x[j] - vec_noise_pred[j]) *
+                        vec_model_output[j] =
+                            (vec_x[j] - vec_model_output[j]) *
                             (1 / sigma);
                     }
                 }
                 // 2. compute alphas, betas
                 float alpha_prod_t = alphas_cumprod[timestep];
-                // Note final_alpha_cumprod = alphas_cumprod[0]
+                // Note final_alpha_cumprod = alphas_cumprod[0] due to
+                // trailing timestep spacing
                 float alpha_prod_t_prev = prev_timestep >= 0 ?
                     alphas_cumprod[prev_timestep] : alphas_cumprod[0];
                 float beta_prod_t = 1 - alpha_prod_t;
                 // 3. compute predicted original sample from predicted
                 // noise also called "predicted x_0" of formula (12)
                 // from https://arxiv.org/pdf/2010.02502.pdf
-                struct ggml_tensor* pred_original_sample =
-                    ggml_dup_tensor(work_ctx, x);
                 {
                     float* vec_x = (float*)x->data;
-                    float* vec_noise_pred = (float*)noise_pred->data;
+                    float* vec_model_output =
+                        (float*)model_output->data;
                     float* vec_pred_original_sample =
                         (float*)pred_original_sample->data;
                     // Note the substitution of latents or sample = x
@@ -1119,12 +1144,12 @@ static void sample_k_diffusion(sample_method_t method,
                         vec_pred_original_sample[j] =
                             (vec_x[j] / std::sqrt(sigma * sigma + 1) -
                              std::sqrt(beta_prod_t) *
-                             vec_noise_pred[j]) *
+                             vec_model_output[j]) *
                             (1 / std::sqrt(alpha_prod_t));
                     }
                 }
                 // Assuming the "epsilon" prediction type, where below
-                // pred_epsilon = noise_pred is inserted, and is not
+                // pred_epsilon = model_output is inserted, and is not
                 // defined/copied explicitly.
                 //
                 // 5. compute variance: "sigma_t(eta)" -> see formula
@@ -1135,34 +1160,35 @@ static void sample_k_diffusion(sample_method_t method,
                 float beta_prod_t_prev = 1 - alpha_prod_t_prev;
                 float variance = (beta_prod_t_prev / beta_prod_t) *
                     (1 - alpha_prod_t / alpha_prod_t_prev);
-                float std_dev_t = 0 * std::sqrt(variance);
+                float std_dev_t = eta * std::sqrt(variance);
                 // 6. compute "direction pointing to x_t" of formula
                 // (12) from https://arxiv.org/pdf/2010.02502.pdf
-                struct ggml_tensor* pred_sample_direction =
-                    ggml_dup_tensor(work_ctx, noise_pred);
-                {
-                    float* vec_noise_pred = (float*)noise_pred->data;
-                    float* vec_pred_sample_direction =
-                        (float*)pred_sample_direction->data;
-                    for (int j = 0; j < ggml_nelements(x); j++) {
-                        vec_pred_sample_direction[j] =
-                            std::sqrt(1 - alpha_prod_t_prev -
-                                      std::pow(std_dev_t, 2)) *
-                            vec_noise_pred[j];
-                    }
-                }
                 // 7. compute x_t without "random noise" of formula
                 // (12) from https://arxiv.org/pdf/2010.02502.pdf
                 {
+                    float* vec_model_output = (float*)model_output->data;
                     float* vec_pred_original_sample =
                         (float*)pred_original_sample->data;
-                    float* vec_pred_sample_direction =
-                        (float*)pred_sample_direction->data;
                     float* vec_x = (float*)x->data;
                     for (int j = 0; j < ggml_nelements(x); j++) {
+                        // Two step inner loop without an explicit
+                        // tensor
+                        float pred_sample_direction =
+                            std::sqrt(1 - alpha_prod_t_prev -
+                                      std::pow(std_dev_t, 2)) *
+                            vec_model_output[j];
                         vec_x[j] = std::sqrt(alpha_prod_t_prev) *
                             vec_pred_original_sample[j] +
-                            vec_pred_sample_direction[j];
+                            pred_sample_direction;
+                    }
+                }
+                if (eta > 0) {
+                    ggml_tensor_set_f32_randn(variance_noise, rng);
+                    float* vec_variance_noise =
+                        (float*)variance_noise->data;
+                    float* vec_x = (float*)x->data;
+                    for (int j = 0; j < ggml_nelements(x); j++) {
+                        vec_x[j] += std_dev_t * vec_variance_noise[j];
                     }
                 }
                 // See the note above: x = latents or sample here, and
@@ -1173,6 +1199,174 @@ static void sample_k_diffusion(sample_method_t method,
                 // factor c_in.
             }
         } break;
+        case TCD:  // Strategic Stochastic Sampling (Algorithm 4) in
+                   // Trajectory Consistency Distillation
+        {
+            float beta_start = 0.00085f;
+            float beta_end = 0.0120f;
+            std::vector<double> alphas_cumprod;
+            std::vector<double> compvis_sigmas;
+
+            alphas_cumprod.reserve(TIMESTEPS);
+            compvis_sigmas.reserve(TIMESTEPS);
+            for (int i = 0; i < TIMESTEPS; i++) {
+                alphas_cumprod[i] =
+                    (i == 0 ? 1.0f : alphas_cumprod[i - 1]) *
+                    (1.0f -
+                     std::pow(sqrtf(beta_start) +
+                              (sqrtf(beta_end) - sqrtf(beta_start)) *
+                              ((float)i / (TIMESTEPS - 1)), 2));
+                compvis_sigmas[i] =
+                    std::sqrt((1 - alphas_cumprod[i]) /
+                              alphas_cumprod[i]);
+            }
+            int original_steps = 50;
+
+            struct ggml_tensor* pred_original_sample =
+                ggml_dup_tensor(work_ctx, x);
+            struct ggml_tensor* noise =
+                ggml_dup_tensor(work_ctx, x);
+
+            for (int i = 0; i < steps; i++) {
+                // Analytic form for TCD timesteps
+                int timestep = TIMESTEPS - 1 -
+                    (TIMESTEPS / original_steps) *
+                    (int)floor(i * ((float)original_steps / steps));
+                // 1. get previous step value
+                int prev_timestep = i >= steps - 1 ? 0 :
+                    TIMESTEPS - 1 - (TIMESTEPS / original_steps) *
+                    (int)floor((i + 1) *
+                               ((float)original_steps / steps));
+                // Here timestep_s is tau_n' in Algorithm 4. The _s
+                // notation appears to be that from DPM-Solver, C. Lu,
+                // arXiv:2206.00927 [cs.LG], but this notation is not
+                // continued in Algorithm 4, where _n' is used.
+                int timestep_s =
+                    (int)floor((1 - eta) * prev_timestep);
+                // Begin k-diffusion specific workaround for
+                // evaluating F_theta(x; ...) from D(x, sigma), same
+                // as in DDIM (and see there for detailed comments)
+                float sigma = compvis_sigmas[timestep];
+                if (i == 0) {
+                    float* vec_x = (float*)x->data;
+                    for (int j = 0; j < ggml_nelements(x); j++) {
+                        vec_x[j] *= std::sqrt(sigma * sigma + 1) /
+                            sigma;
+                    }
+                }
+                else {
+                    float* vec_x = (float*)x->data;
+                    for (int j = 0; j < ggml_nelements(x); j++) {
+                        vec_x[j] *= std::sqrt(sigma * sigma + 1);
+                    }
+                }
+                struct ggml_tensor* model_output =
+                    model(x, sigma, i + 1);
+                {
+                    float* vec_x = (float*)x->data;
+                    float* vec_model_output =
+                        (float*)model_output->data;
+                    for (int j = 0; j < ggml_nelements(x); j++) {
+                        vec_model_output[j] =
+                            (vec_x[j] - vec_model_output[j]) *
+                            (1 / sigma);
+                    }
+                }
+                // 2. compute alphas, betas
+                //
+                // When comparing TCD with DDPM/DDIM note that Zheng
+                // et al. (2024) follows the DPM-Solver notation for
+                // alpha. One can find the following comment in the
+                // original DPM-Solver code
+                // (https://github.com/LuChengTHU/dpm-solver/):
+                // "**Important**: Please pay special attention for
+                // the args for `alphas_cumprod`: The `alphas_cumprod`
+                // is the \hat{alpha_n} arrays in the notations of
+                // DDPM. [...] Therefore, the notation \hat{alpha_n}
+                // is different from the notation alpha_t in
+                // DPM-Solver. In fact, we have alpha_{t_n} =
+                // \sqrt{\hat{alpha_n}}, [...]"
+                float alpha_prod_t = alphas_cumprod[timestep];
+                float beta_prod_t = 1 - alpha_prod_t;
+                // Note final_alpha_cumprod = alphas_cumprod[0] since
+                // TCD is always "trailing"
+                float alpha_prod_t_prev = prev_timestep >= 0 ?
+                    alphas_cumprod[prev_timestep] : alphas_cumprod[0];
+                // The subscript _s are the only portion in this
+                // section (2) unique to TCD
+                float alpha_prod_s = alphas_cumprod[timestep_s];
+                float beta_prod_s = 1 - alpha_prod_s;
+                // 3. Compute the predicted noised sample x_s based on
+                // the model parameterization
+                //
+                // This section is also exactly the same as DDIM
+                {
+                    float* vec_x = (float*)x->data;
+                    float* vec_model_output =
+                        (float*)model_output->data;
+                    float* vec_pred_original_sample =
+                        (float*)pred_original_sample->data;
+                    for (int j = 0; j < ggml_nelements(x); j++) {
+                        vec_pred_original_sample[j] =
+                            (vec_x[j] / std::sqrt(sigma * sigma + 1) -
+                             std::sqrt(beta_prod_t) *
+                             vec_model_output[j]) *
+                            (1 / std::sqrt(alpha_prod_t));
+                    }
+                }
+                // This consistency function step can be difficult to
+                // decipher from Algorithm 4, as it involves a
+                // difficult notation ("|->"). In Diffusers it is
+                // borrowed verbatim (with the same comments below for
+                // step (4)) from LCMScheduler's noise injection step,
+                // compare in S. Luo et al., arXiv:2310.04378 p. 14,
+                // Algorithm 3.
+                {
+                    float* vec_pred_original_sample =
+                        (float*)pred_original_sample->data;
+                    float* vec_model_output =
+                        (float*)model_output->data;
+                    float* vec_x = (float*)x->data;
+                    for (int j = 0; j < ggml_nelements(x); j++) {
+                        // Substituting x = pred_noised_sample and
+                        // pred_epsilon = model_output
+                        vec_x[j] =
+                            std::sqrt(alpha_prod_s) *
+                            vec_pred_original_sample[j] +
+                            std::sqrt(beta_prod_s) *
+                            vec_model_output[j];
+                    }
+                }
+                // 4. Sample and inject noise z ~ N(0, I) for
+                // MultiStep Inference Noise is not used on the final
+                // timestep of the timestep schedule. This also means
+                // that noise is not used for one-step sampling. Eta
+                // (referred to as "gamma" in the paper) was
+                // introduced to control the stochasticity in every
+                // step. When eta = 0, it represents deterministic
+                // sampling, whereas eta = 1 indicates full stochastic
+                // sampling.
+                if (eta > 0 && i != steps - 1) {
+                    // In this case, x is still pred_noised_sample,
+                    // continue in-place
+                    ggml_tensor_set_f32_randn(noise, rng);
+                    float* vec_x = (float*)x->data;
+                    float* vec_noise = (float*)noise->data;
+                    for (int j = 0; j < ggml_nelements(x); j++) {
+                        // Corresponding to (35) in Zheng et
+                        // al. (2024), substituting x =
+                        // pred_noised_sample
+                        vec_x[j] =
+                            std::sqrt(alpha_prod_t_prev /
+                                      alpha_prod_s) *
+                            vec_x[j] +
+                            std::sqrt(1 - alpha_prod_t_prev /
+                                      alpha_prod_s) *
+                            vec_noise[j];
+                    }
+                }
+            }
+        } break;
 
         default:
             LOG_ERROR("Attempting to sample with nonexisting sample method %i", method);