Generate samples from the model.
(
self,
model,
noise,
clip_denoised=False,
progress=False,
return_intermediates=False,
intermediate_key="sample",
intermediate_freq=50,
model_kwargs=None,
pbar_desc="Sampling",
)
| 348 | return {'sample': sample, 'pred_xstart': out['pred_xstart']} |
| 349 | |
| 350 | def p_sample_loop( |
| 351 | self, |
| 352 | model, |
| 353 | noise, |
| 354 | clip_denoised=False, |
| 355 | progress=False, |
| 356 | return_intermediates=False, |
| 357 | intermediate_key="sample", |
| 358 | intermediate_freq=50, |
| 359 | model_kwargs=None, |
| 360 | pbar_desc="Sampling", |
| 361 | ): |
| 362 | """ Generate samples from the model. """ |
| 363 | model_kwargs = model_kwargs or {} |
| 364 | |
| 365 | shape = noise.shape |
| 366 | dev = noise.device |
| 367 | |
| 368 | indices = list(range(self.num_timesteps))[::-1] |
| 369 | |
| 370 | if progress: |
| 371 | # Lazy import so that we don't depend on tqdm. |
| 372 | from tqdm.auto import tqdm |
| 373 | indices = tqdm(indices, desc=f"Sampling") |
| 374 | |
| 375 | img = noise |
| 376 | intermediates = [img.cpu()] |
| 377 | for i in indices: |
| 378 | t = torch.tensor([i] * shape[0], device=dev).long() |
| 379 | with torch.no_grad(): |
| 380 | out = self.p_sample( |
| 381 | model, img, t, |
| 382 | clip_denoised=clip_denoised, |
| 383 | model_kwargs=model_kwargs |
| 384 | ) |
| 385 | if return_intermediates: |
| 386 | if i % intermediate_freq == 0 or i == self.num_timesteps - 1: |
| 387 | intermediates.append(out[intermediate_key].cpu()) |
| 388 | |
| 389 | img = out["sample"] |
| 390 | |
| 391 | if return_intermediates: |
| 392 | return img, intermediates |
| 393 | return img |
| 394 | |
| 395 | """ VLB """ |
| 396 |