| 273 | |
| 274 | |
| 275 | def ddim_update( |
| 276 | model_output: torch.Tensor, |
| 277 | sigmas, |
| 278 | step_index, |
| 279 | sample: torch.Tensor = None, |
| 280 | noise: Optional[torch.Tensor] = None, |
| 281 | eta: float = 1.0, |
| 282 | ) -> torch.Tensor: |
| 283 | |
| 284 | t, s = sigmas[step_index + 1], sigmas[step_index] |
| 285 | |
| 286 | std_dev_t = eta * t |
| 287 | dt_sqrt = torch.sqrt(1.0 - t**2 * (1 - s) ** 2 / (s**2 * (1 - t) ** 2)) |
| 288 | rho_t = std_dev_t * dt_sqrt |
| 289 | noise_pred = (sample - (1 - s) * model_output) / s |
| 290 | if noise is None: |
| 291 | noise = torch.randn_like(model_output) |
| 292 | prev_mean = (1 - t) * model_output + torch.sqrt(t**2 - rho_t**2) * noise_pred |
| 293 | x_t = prev_mean + rho_t * noise |
| 294 | |
| 295 | return x_t, prev_mean, std_dev_t, dt_sqrt |
| 296 | |
| 297 | |
| 298 | def dpm_solver_first_order_update( |