Args: input_noise: random noise, of the same shape as the desired sample. diffusion_model: model to sample from. controlnet: controlnet sub-network. cn_cond: conditioning image for the ControlNet. scheduler: diffusion scheduler. If
( # type: ignore[override]
self,
input_noise: torch.Tensor,
diffusion_model: DiffusionModelUNet,
controlnet: ControlNet,
cn_cond: torch.Tensor,
scheduler: Scheduler | None = None,
save_intermediates: bool | None = False,
intermediate_steps: int | None = 100,
conditioning: torch.Tensor | None = None,
mode: str = "crossattn",
verbose: bool = True,
seg: torch.Tensor | None = None,
cfg: float | None = None,
cfg_fill_value: float = -1.0,
)
| 1471 | |
| 1472 | @torch.no_grad() |
| 1473 | def sample( # type: ignore[override] |
| 1474 | self, |
| 1475 | input_noise: torch.Tensor, |
| 1476 | diffusion_model: DiffusionModelUNet, |
| 1477 | controlnet: ControlNet, |
| 1478 | cn_cond: torch.Tensor, |
| 1479 | scheduler: Scheduler | None = None, |
| 1480 | save_intermediates: bool | None = False, |
| 1481 | intermediate_steps: int | None = 100, |
| 1482 | conditioning: torch.Tensor | None = None, |
| 1483 | mode: str = "crossattn", |
| 1484 | verbose: bool = True, |
| 1485 | seg: torch.Tensor | None = None, |
| 1486 | cfg: float | None = None, |
| 1487 | cfg_fill_value: float = -1.0, |
| 1488 | ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: |
| 1489 | """ |
| 1490 | Args: |
| 1491 | input_noise: random noise, of the same shape as the desired sample. |
| 1492 | diffusion_model: model to sample from. |
| 1493 | controlnet: controlnet sub-network. |
| 1494 | cn_cond: conditioning image for the ControlNet. |
| 1495 | scheduler: diffusion scheduler. If none provided will use the class attribute scheduler |
| 1496 | save_intermediates: whether to return intermediates along the sampling change |
| 1497 | intermediate_steps: if save_intermediates is True, saves every n steps |
| 1498 | conditioning: Conditioning for network input. |
| 1499 | mode: Conditioning mode for the network. |
| 1500 | verbose: if true, prints the progression bar of the sampling process. |
| 1501 | seg: if diffusion model is instance of SPADEDiffusionModel, segmentation must be provided. |
| 1502 | cfg: classifier-free-guidance scale, which indicates the level of strengthening on the conditioning. |
| 1503 | cfg_fill_value: the fill value to use for the unconditioned input when using classifier-free guidance. |
| 1504 | """ |
| 1505 | if mode not in ["crossattn", "concat"]: |
| 1506 | raise NotImplementedError(f"{mode} condition is not supported") |
| 1507 | |
| 1508 | if not scheduler: |
| 1509 | scheduler = self.scheduler |
| 1510 | image = input_noise |
| 1511 | |
| 1512 | all_next_timesteps = torch.cat((scheduler.timesteps[1:], torch.tensor([0], dtype=scheduler.timesteps.dtype))) |
| 1513 | if verbose and has_tqdm: |
| 1514 | progress_bar = tqdm( |
| 1515 | zip(scheduler.timesteps, all_next_timesteps), |
| 1516 | total=min(len(scheduler.timesteps), len(all_next_timesteps)), |
| 1517 | ) |
| 1518 | else: |
| 1519 | progress_bar = iter(zip(scheduler.timesteps, all_next_timesteps)) |
| 1520 | intermediates = [] |
| 1521 | |
| 1522 | if cfg is not None: |
| 1523 | cn_cond = torch.cat([cn_cond] * 2, dim=0) |
| 1524 | |
| 1525 | for t, next_t in progress_bar: |
| 1526 | # Controlnet prediction |
| 1527 | if cfg is not None: |
| 1528 | model_input = torch.cat([image] * 2, dim=0) |
| 1529 | if conditioning is not None: |
| 1530 | uncondition = torch.ones_like(conditioning) |