Second-order DPM-Solver++ midpoint update. m0 = most recent x0 prediction (at timestep s0) m1 = previous x0 prediction (at timestep s1) Goes from s0 to target t.
(
&self,
m0: &Tensor,
m1: &Tensor,
s0: usize,
s1: usize,
t: usize,
sample: &Tensor,
)
| 123 | /// m1 = previous x0 prediction (at timestep s1) |
| 124 | /// Goes from s0 to target t. |
| 125 | pub fn second_order_update( |
| 126 | &self, |
| 127 | m0: &Tensor, |
| 128 | m1: &Tensor, |
| 129 | s0: usize, |
| 130 | s1: usize, |
| 131 | t: usize, |
| 132 | sample: &Tensor, |
| 133 | ) -> Result<Tensor> { |
| 134 | let lambda_s0 = self.lambda_t[s0]; |
| 135 | let lambda_s1 = self.lambda_t[s1]; |
| 136 | let lambda_t = self.lambda_t[t]; |
| 137 | let alpha_t = self.alpha_t[t]; |
| 138 | let sigma_t = self.sigma_t[t]; |
| 139 | let sigma_s0 = self.sigma_t[s0]; |
| 140 | |
| 141 | let h = lambda_t - lambda_s0; |
| 142 | let h_0 = lambda_s0 - lambda_s1; |
| 143 | let r0 = h_0 / h; |
| 144 | |
| 145 | // D0 = m0, D1 = (1/(2r)) * (m0 - m1) [midpoint solver type] |
| 146 | let d1 = ((m0 - m1)? * (1.0 / (2.0 * r0)))?; |
| 147 | |
| 148 | let ratio = sigma_t / sigma_s0; |
| 149 | let expm1_neg_h = (-h).exp_m1(); // exp(-h) - 1 |
| 150 | |
| 151 | // x_t = ratio * x - alpha_t * expm1(-h) * D0 - 0.5 * alpha_t * expm1(-h) * D1 |
| 152 | let coeff = alpha_t * expm1_neg_h; |
| 153 | let term1 = (sample * ratio)?; |
| 154 | let term2 = (m0 * coeff)?; |
| 155 | let term3 = (&d1 * (0.5 * coeff))?; |
| 156 | (term1 - term2)? - term3 |
| 157 | } |
| 158 | |
| 159 | /// Run a complete diffusion step, managing solver order automatically. |
| 160 | /// |
no outgoing calls