Args: input_noise: random noise, of the same shape as the desired sample. diffusion_model: model to sample from. scheduler: diffusion scheduler. If none provided will use the class attribute scheduler save_intermediates: whether to return inte
(
self,
input_noise: torch.Tensor,
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,
)
| 905 | |
| 906 | @torch.no_grad() |
| 907 | def sample( |
| 908 | self, |
| 909 | input_noise: torch.Tensor, |
| 910 | diffusion_model: DiffusionModelUNet, |
| 911 | scheduler: Scheduler | None = None, |
| 912 | save_intermediates: bool | None = False, |
| 913 | intermediate_steps: int | None = 100, |
| 914 | conditioning: torch.Tensor | None = None, |
| 915 | mode: str = "crossattn", |
| 916 | verbose: bool = True, |
| 917 | seg: torch.Tensor | None = None, |
| 918 | cfg: float | None = None, |
| 919 | cfg_fill_value: float = -1.0, |
| 920 | ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: |
| 921 | """ |
| 922 | Args: |
| 923 | input_noise: random noise, of the same shape as the desired sample. |
| 924 | diffusion_model: model to sample from. |
| 925 | scheduler: diffusion scheduler. If none provided will use the class attribute scheduler |
| 926 | save_intermediates: whether to return intermediates along the sampling change |
| 927 | intermediate_steps: if save_intermediates is True, saves every n steps |
| 928 | conditioning: Conditioning for network input. |
| 929 | mode: Conditioning mode for the network. |
| 930 | verbose: if true, prints the progression bar of the sampling process. |
| 931 | seg: if diffusion model is instance of SPADEDiffusionModel, segmentation must be provided. |
| 932 | cfg: classifier-free-guidance scale, which indicates the level of strengthening on the conditioning. |
| 933 | cfg_fill_value: the fill value to use for the unconditioned input when using classifier-free guidance. |
| 934 | """ |
| 935 | if mode not in ["crossattn", "concat"]: |
| 936 | raise NotImplementedError(f"{mode} condition is not supported") |
| 937 | if mode == "concat" and conditioning is None: |
| 938 | raise ValueError("Conditioning must be supplied for if condition mode is concat.") |
| 939 | if not scheduler: |
| 940 | scheduler = self.scheduler |
| 941 | image = input_noise |
| 942 | |
| 943 | all_next_timesteps = torch.cat((scheduler.timesteps[1:], torch.tensor([0], dtype=scheduler.timesteps.dtype))) |
| 944 | if verbose and has_tqdm: |
| 945 | progress_bar = tqdm( |
| 946 | zip(scheduler.timesteps, all_next_timesteps), |
| 947 | total=min(len(scheduler.timesteps), len(all_next_timesteps)), |
| 948 | ) |
| 949 | else: |
| 950 | progress_bar = iter(zip(scheduler.timesteps, all_next_timesteps)) |
| 951 | intermediates = [] |
| 952 | |
| 953 | for t, next_t in progress_bar: |
| 954 | # 1. predict noise model_output |
| 955 | diffusion_model = ( |
| 956 | partial(diffusion_model, seg=seg) |
| 957 | if isinstance(diffusion_model, SPADEDiffusionModelUNet) |
| 958 | else diffusion_model |
| 959 | ) |
| 960 | if ( |
| 961 | cfg is not None |
| 962 | ): # if classifier-free guidance is used, a conditioned and unconditioned bit is generated. |
| 963 | model_input = torch.cat([image] * 2, dim=0) |
| 964 | if conditioning is not None: |