Args: input_noise: random noise, of the same shape as the desired latent representation. autoencoder_model: first stage model. diffusion_model: model to sample from. scheduler: diffusion scheduler. If none provided will use the class attribute
( # type: ignore[override]
self,
input_noise: torch.Tensor,
autoencoder_model: AutoencoderKL | VQVAE,
diffusion_model: DiffusionModelUNet,
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,
)
| 1251 | |
| 1252 | @torch.no_grad() |
| 1253 | def sample( # type: ignore[override] |
| 1254 | self, |
| 1255 | input_noise: torch.Tensor, |
| 1256 | autoencoder_model: AutoencoderKL | VQVAE, |
| 1257 | diffusion_model: DiffusionModelUNet, |
| 1258 | scheduler: Scheduler | None = None, |
| 1259 | save_intermediates: bool | None = False, |
| 1260 | intermediate_steps: int | None = 100, |
| 1261 | conditioning: torch.Tensor | None = None, |
| 1262 | mode: str = "crossattn", |
| 1263 | verbose: bool = True, |
| 1264 | seg: torch.Tensor | None = None, |
| 1265 | cfg: float | None = None, |
| 1266 | cfg_fill_value: float = -1.0, |
| 1267 | ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: |
| 1268 | """ |
| 1269 | Args: |
| 1270 | input_noise: random noise, of the same shape as the desired latent representation. |
| 1271 | autoencoder_model: first stage model. |
| 1272 | diffusion_model: model to sample from. |
| 1273 | scheduler: diffusion scheduler. If none provided will use the class attribute scheduler. |
| 1274 | save_intermediates: whether to return intermediates along the sampling change |
| 1275 | intermediate_steps: if save_intermediates is True, saves every n steps |
| 1276 | conditioning: Conditioning for network input. |
| 1277 | mode: Conditioning mode for the network. |
| 1278 | verbose: if true, prints the progression bar of the sampling process. |
| 1279 | seg: if diffusion model is instance of SPADEDiffusionModel, or autoencoder_model |
| 1280 | is instance of SPADEAutoencoderKL, segmentation must be provided. |
| 1281 | cfg: classifier-free-guidance scale, which indicates the level of strengthening on the conditioning. |
| 1282 | cfg_fill_value: the fill value to use for the unconditioned input when using classifier-free guidance. |
| 1283 | """ |
| 1284 | |
| 1285 | if ( |
| 1286 | isinstance(autoencoder_model, SPADEAutoencoderKL) |
| 1287 | and isinstance(diffusion_model, SPADEDiffusionModelUNet) |
| 1288 | and autoencoder_model.decoder.label_nc != diffusion_model.label_nc |
| 1289 | ): |
| 1290 | raise ValueError( |
| 1291 | f"If both autoencoder_model and diffusion_model implement SPADE, the number of semantic" |
| 1292 | f"labels for each must be compatible, but got {autoencoder_model.decoder.label_nc} and" |
| 1293 | f"{diffusion_model.label_nc}" |
| 1294 | ) |
| 1295 | |
| 1296 | outputs = super().sample( |
| 1297 | input_noise=input_noise, |
| 1298 | diffusion_model=diffusion_model, |
| 1299 | scheduler=scheduler, |
| 1300 | save_intermediates=save_intermediates, |
| 1301 | intermediate_steps=intermediate_steps, |
| 1302 | conditioning=conditioning, |
| 1303 | mode=mode, |
| 1304 | verbose=verbose, |
| 1305 | seg=seg, |
| 1306 | cfg=cfg, |
| 1307 | cfg_fill_value=cfg_fill_value, |
| 1308 | ) |
| 1309 | |
| 1310 | if save_intermediates: |