| 1359 | timesteps = noise_scheduler.timesteps[indices].to(device=latents.device) |
| 1360 | |
| 1361 | def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): |
| 1362 | sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype) |
| 1363 | schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device) |
| 1364 | timesteps = timesteps.to(accelerator.device) |
| 1365 | step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] |
| 1366 | |
| 1367 | sigma = sigmas[step_indices].flatten() |
| 1368 | while len(sigma.shape) < n_dim: |
| 1369 | sigma = sigma.unsqueeze(-1) |
| 1370 | return sigma |
| 1371 | |
| 1372 | # Add noise according to flow matching. |
| 1373 | # zt = (1 - texp) * x + texp * z1 |