MONAI/brain_image_synthesis_latent_diffusion_model

1from __future__ import annotations2 3import torch4import torch.nn as nn5from monai.utils import optional_import6from torch.cuda.amp import autocast7 8tqdm, has_tqdm = optional_import("tqdm", name="tqdm")9 10 11class Sampler:12    def __init__(self) -> None:13        super().__init__()14 15    @torch.no_grad()16    def sampling_fn(17        self,18        input_noise: torch.Tensor,19        autoencoder_model: nn.Module,20        diffusion_model: nn.Module,21        scheduler: nn.Module,22        conditioning: torch.Tensor,23    ) -> torch.Tensor:24        if has_tqdm:25            progress_bar = tqdm(scheduler.timesteps)26        else:27            progress_bar = iter(scheduler.timesteps)28 29        image = input_noise30        cond_concat = conditioning.squeeze(1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)31        cond_concat = cond_concat.expand(list(cond_concat.shape[0:2]) + list(input_noise.shape[2:]))32        for t in progress_bar:33            with torch.no_grad():34                model_output = diffusion_model(35                    torch.cat((image, cond_concat), dim=1),36                    timesteps=torch.Tensor((t,)).to(input_noise.device).long(),37                    context=conditioning,38                )39                image, _ = scheduler.step(model_output, t, image)40 41        with torch.no_grad():42            with autocast():43                sample = autoencoder_model.decode_stage_2_outputs(image)44 45        return sample46