Method step

cake-core/src/models/vibevoice/ddpm.rs:166–212

Run a complete diffusion step, managing solver order automatically. `step_idx`: 0-based index into timesteps `x0_buffer`: mutable buffer of previous x0 predictions (caller maintains) `ts_buffer`: mutable buffer of corresponding timesteps (caller maintains) Returns the updated sample.

(
        &self,
        v_pred: &Tensor,
        step_idx: usize,
        sample: &Tensor,
        x0_buffer: &mut Vec<Tensor>,
        ts_buffer: &mut Vec<usize>,
    )

    ///
    /// Returns the updated sample.
    pub fn step(
        &self,
        v_pred: &Tensor,
        step_idx: usize,
        sample: &Tensor,
        x0_buffer: &mut Vec<Tensor>,
        ts_buffer: &mut Vec<usize>,
    ) -> Result<Tensor> {
        let t = self.timesteps[step_idx];
        let num_steps = self.timesteps.len();

        // Convert v-prediction to x0
        let x0 = self.convert_v_to_x0(v_pred, t, sample)?;

        // Target: next timestep, or final_sigma=0 (perfectly clean) for last step
        let is_final = step_idx + 1 >= num_steps;
        let next_t = if !is_final {
            Some(self.timesteps[step_idx + 1])
        } else {
            None // final_sigmas_type="zero": target sigma=0, alpha=1
        };

        // Use first order for: first step, last step (lower_order_final with final_sigma=0)
        let use_first_order = step_idx == 0
            || is_final
            || x0_buffer.is_empty();

        let result = if is_final {
            // final_sigmas_type="zero": last step returns x0 directly
            // (sigma_t=0 means x_t = 0*sample - 1*(0-1)*x0 = x0)
            x0.clone()
        } else if use_first_order {
            self.first_order_update(&x0, t, next_t.unwrap(), sample)?
        } else {
            let m1 = x0_buffer.last().unwrap();
            let s1 = *ts_buffer.last().unwrap();
            self.second_order_update(&x0, m1, t, s1, next_t.unwrap(), sample)?
        };

        // Update buffers (keep only last entry for second-order)
        x0_buffer.clear();
        x0_buffer.push(x0);
        ts_buffer.clear();
        ts_buffer.push(t);

        Ok(result)
    }
}

// Keep the old name as an alias for backward compatibility