| 466 | return sigma |
| 467 | |
| 468 | def get_timestep(self, bsz, device): |
| 469 | if self.hparams.timestep_densities_type == "logit_normal": |
| 470 | # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$). |
| 471 | # In practice, we sample the random variable u from a normal distribution u ∼ N (u; m, s) |
| 472 | # and map it through the standard logistic function |
| 473 | u = torch.normal( |
| 474 | mean=self.hparams.logit_mean, |
| 475 | std=self.hparams.logit_std, |
| 476 | size=(bsz,), |
| 477 | device="cpu", |
| 478 | ) |
| 479 | u = torch.nn.functional.sigmoid(u) |
| 480 | indices = (u * self.scheduler.config.num_train_timesteps).long() |
| 481 | indices = torch.clamp( |
| 482 | indices, 0, self.scheduler.config.num_train_timesteps - 1 |
| 483 | ) |
| 484 | timesteps = self.scheduler.timesteps[indices].to(device) |
| 485 | |
| 486 | return timesteps |
| 487 | |
| 488 | def run_step(self, batch, batch_idx): |
| 489 | self.plot_step(batch, batch_idx) |