Args: x: source minibatch (bs, *dim) sample_kwargs: dict, additional sampling arguments for the solver num_steps: int, number of steps to take cfg_scale: float, scale for the classifier-free guidance uc_cond: torch.Tens
(self, x: Tensor, sample_kwargs=None, reverse=False, return_intermediates=False, **kwargs)
| 452 | return self(x=x, t=t, **kwargs) |
| 453 | |
| 454 | def generate(self, x: Tensor, sample_kwargs=None, reverse=False, return_intermediates=False, **kwargs): |
| 455 | """ |
| 456 | Args: |
| 457 | x: source minibatch (bs, *dim) |
| 458 | sample_kwargs: dict, additional sampling arguments for the solver |
| 459 | num_steps: int, number of steps to take |
| 460 | cfg_scale: float, scale for the classifier-free guidance |
| 461 | uc_cond: torch.Tensor, unconditional conditioning information (1, *dim) or (bs, *dim) |
| 462 | cond_key: str, key for the conditional information |
| 463 | intermediate_freq: int, frequency of intermediate outputs |
| 464 | use_sde: if true, use SDE sampling instead of ODE |
| 465 | __ ODE Sampler __: |
| 466 | method: str, method for the ODE solver (see torchdiffeq) |
| 467 | atol/rtol: float, absolute and relative tolerance for the ODE solver |
| 468 | __ SDE Sampler __: |
| 469 | method: str, method for the SDE solver (euler, heun) |
| 470 | diffusion_form: str, form of the diffusion coefficient (sigma, SBDM, ...) |
| 471 | diffusion_norm: float, magnitude of the diffusion coefficient (default 1.0) |
| 472 | last_step: str, type of the last step (Mean, Tweedie, Euler) |
| 473 | last_step_size: float, size of the last step (default 0.04) |
| 474 | progress: bool, whether to show a progress bar |
| 475 | reverse: bool, whether to reverse the direction of the flow. If True, |
| 476 | we map from x1 -> x0, otherwise we map from x0 -> x1. |
| 477 | n_intermediates: int, number of intermediate points to return. |
| 478 | kwargs: additional arguments for the network (e.g. conditioning information). |
| 479 | """ |
| 480 | sample_kwargs = sample_kwargs or {} |
| 481 | |
| 482 | # timesteps |
| 483 | num_steps = sample_kwargs.get("num_steps", 50) |
| 484 | t = torch.linspace(0, 1, num_steps, dtype=x.dtype).to(x.device) |
| 485 | t = 1 - t if reverse else t |
| 486 | |
| 487 | # include classifier-free guidance |
| 488 | cfg_kwargs = dict( |
| 489 | cfg_scale=sample_kwargs.get("cfg_scale", 1.0), |
| 490 | uc_cond=sample_kwargs.get("uc_cond", None), |
| 491 | cond_key=sample_kwargs.get("cond_key", "y"), |
| 492 | ) |
| 493 | |
| 494 | # SDE sampling |
| 495 | if sample_kwargs.get("use_sde", False): |
| 496 | results = self.sde_sampler.sample( |
| 497 | init=x, |
| 498 | model=self.net, # sde_sampler already includes CFG |
| 499 | sampling_method=sample_kwargs.get("method", "euler"), |
| 500 | diffusion_form=sample_kwargs.get("diffusion_form", "sigma"), |
| 501 | diffusion_norm=sample_kwargs.get("diffusion_norm", 1.0), |
| 502 | last_step=sample_kwargs.get("last_step", "Mean"), |
| 503 | last_step_size=sample_kwargs.get("last_step_size", 0.04), |
| 504 | num_steps=num_steps, |
| 505 | progress=sample_kwargs.get("progress", False), |
| 506 | return_intermediates=True, |
| 507 | **cfg_kwargs, |
| 508 | **kwargs |
| 509 | ) |
| 510 | |
| 511 | # ODE sampling |