konfai 1.1.8__py3-none-any.whl → 1.2.0__py3-none-any.whl

konfai/models/generation/ddpm.py

@@ -1,16 +1,14 @@
+from collections.abc import Callable
 
-
-
-from typing import Union, Callable
+import numpy as np
 import torch
 import tqdm
-import numpy as np
 
-from konfai.network import network, blocks
-from konfai.utils.config import config
 from konfai.data.patching import ModelPatch
-from konfai.utils.utils import gpuInfo
 from konfai.metric.measure import Criterion
+from konfai.network import blocks, network
+from konfai.utils.utils import gpu_info
+
 
 def cosine_beta_schedule(timesteps, s=0.008):
     steps = timesteps + 1
@@ -20,198 +18,427 @@ def cosine_beta_schedule(timesteps, s=0.008):
     betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
     return torch.clip(betas, 0.0001, 0.9999)
 
+
 class DDPM(network.Network):
-    
-    class DDPM_TE(torch.nn.Module):
+
+    class DDPMTE(torch.nn.Module):
 
         def __init__(self, in_channels: int, out_channels: int) -> None:
             super().__init__()
             self.linear_0 = torch.nn.Linear(in_channels, out_channels)
             self.siLU = torch.nn.SiLU()
             self.linear_1 = torch.nn.Linear(out_channels, out_channels)
-        
-        def forward(self, input: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-            return input + self.linear_1(self.siLU(self.linear_0(t))).reshape(input.shape[0], -1, *[1 for _ in range(len(input.shape)-2)])
-            
-    class DDPM_UNetBlock(network.ModuleArgsDict):
-        
-        def __init__(self, channels: list[int], nb_conv_per_stage: int, blockConfig: blocks.BlockConfig, downSampleMode: blocks.DownSampleMode, upSampleMode: blocks.UpSampleMode, attention : bool, time_embedding_dim: int, dim: int, i : int = 0) -> None:
+
+        def forward(self, tensor: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+            return tensor + self.linear_1(self.siLU(self.linear_0(t))).reshape(
+                tensor.shape[0], -1, *[1 for _ in range(len(tensor.shape) - 2)]
+            )
+
+    class DDPMUNetBlock(network.ModuleArgsDict):
+
+        def __init__(
+            self,
+            channels: list[int],
+            nb_conv_per_stage: int,
+            block_config: blocks.BlockConfig,
+            downsample_mode: blocks.DownsampleMode,
+            upsample_mode: blocks.UpsampleMode,
+            attention: bool,
+            time_embedding_dim: int,
+            dim: int,
+            i: int = 0,
+        ) -> None:
             super().__init__()
             if i > 0:
-                self.add_module(downSampleMode.name, blocks.downSample(in_channels=channels[0], out_channels=channels[1], downSampleMode=downSampleMode, dim=dim))
-            self.add_module("Te_down", DDPM.DDPM_TE(time_embedding_dim, channels[1 if downSampleMode == blocks.DownSampleMode.CONV_STRIDE and i > 0 else 0]), in_branch=[0, 1])
-            self.add_module("DownConvBlock", blocks.ResBlock(in_channels=channels[1 if downSampleMode == blocks.DownSampleMode.CONV_STRIDE and i > 0 else 0], out_channels=channels[1], nb_conv=nb_conv_per_stage, blockConfig=blockConfig, dim=dim))
+                self.add_module(
+                    downsample_mode.name,
+                    blocks.downsample(
+                        in_channels=channels[0],
+                        out_channels=channels[1],
+                        downsample_mode=downsample_mode,
+                        dim=dim,
+                    ),
+                )
+            self.add_module(
+                "Te_down",
+                DDPM.DDPMTE(
+                    time_embedding_dim,
+                    channels[(1 if downsample_mode == blocks.DownsampleMode.CONV_STRIDE and i > 0 else 0)],
+                ),
+                in_branch=[0, 1],
+            )
+            self.add_module(
+                "DownConvBlock",
+                blocks.ResBlock(
+                    in_channels=channels[(1 if downsample_mode == blocks.DownsampleMode.CONV_STRIDE and i > 0 else 0)],
+                    out_channels=channels[1],
+                    block_configs=[block_config] * nb_conv_per_stage,
+                    dim=dim,
+                ),
+            )
             if len(channels) > 2:
-                self.add_module("UNetBlock_{}".format(i+1), DDPM.DDPM_UNetBlock(channels[1:], nb_conv_per_stage, blockConfig, downSampleMode, upSampleMode, attention, time_embedding_dim, dim, i+1), in_branch=[0,1])
-                self.add_module("Te_up", DDPM.DDPM_TE(time_embedding_dim, (channels[1]+channels[2]) if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else channels[1]*2), in_branch=[0, 1])
-                self.add_module("UpConvBlock", blocks.ResBlock(in_channels=(channels[1]+channels[2]) if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else channels[1]*2, out_channels=channels[1], nb_conv=nb_conv_per_stage, blockConfig=blockConfig, dim=dim))
+                self.add_module(
+                    f"UNetBlock_{i + 1}",
+                    DDPM.DDPMUNetBlock(
+                        channels[1:],
+                        nb_conv_per_stage,
+                        block_config,
+                        downsample_mode,
+                        upsample_mode,
+                        attention,
+                        time_embedding_dim,
+                        dim,
+                        i + 1,
+                    ),
+                    in_branch=[0, 1],
+                )
+                self.add_module(
+                    "Te_up",
+                    DDPM.DDPMTE(
+                        time_embedding_dim,
+                        (
+                            (channels[1] + channels[2])
+                            if upsample_mode != blocks.UpsampleMode.CONV_TRANSPOSE
+                            else channels[1] * 2
+                        ),
+                    ),
+                    in_branch=[0, 1],
+                )
+                self.add_module(
+                    "UpConvBlock",
+                    blocks.ResBlock(
+                        in_channels=(
+                            (channels[1] + channels[2])
+                            if upsample_mode != blocks.UpsampleMode.CONV_TRANSPOSE
+                            else channels[1] * 2
+                        ),
+                        out_channels=channels[1],
+                        block_configs=[block_config] * nb_conv_per_stage,
+                        dim=dim,
+                    ),
+                )
             if i > 0:
                 if attention:
-                    self.add_module("Attention", blocks.Attention(F_g=channels[1], F_l=channels[0], F_int=channels[0], dim=dim), in_branch=[2, 0], out_branch=[2])
-                self.add_module(upSampleMode.name, blocks.upSample(in_channels=channels[1], out_channels=channels[0], upSampleMode=upSampleMode, dim=dim))
+                    self.add_module(
+                        "Attention",
+                        blocks.Attention(f_g=channels[1], f_l=channels[0], f_int=channels[0], dim=dim),
+                        in_branch=[2, 0],
+                        out_branch=[2],
+                    )
+                self.add_module(
+                    upsample_mode.name,
+                    blocks.upsample(
+                        in_channels=channels[1],
+                        out_channels=channels[0],
+                        upsample_mode=upsample_mode,
+                        dim=dim,
+                    ),
+                )
                 self.add_module("SkipConnection", blocks.Concat(), in_branch=[0, 2])
-    
-    class DDPM_UNetHead(network.ModuleArgsDict):
+
+    class DDPMUNetHead(network.ModuleArgsDict):
 
         def __init__(self, in_channels: int, out_channels: int, dim: int) -> None:
             super().__init__()
-            self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels = in_channels, out_channels = out_channels, kernel_size = 3, stride = 1, padding = 1))
+            self.add_module(
+                "Conv",
+                blocks.get_torch_module("Conv", dim)(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    kernel_size=3,
+                    stride=1,
+                    padding=1,
+                ),
+            )
 
-    class DDPM_ForwardProcess(torch.nn.Module):
+    class DDPMForwardProcess(torch.nn.Module):
 
-        def __init__(self, noise_step: int=1000, beta_start: float = 1e-4, beta_end: float = 0.02) -> None:
+        def __init__(
+            self,
+            noise_step: int = 1000,
+            beta_start: float = 1e-4,
+            beta_end: float = 0.02,
+        ) -> None:
             super().__init__()
             self.betas = torch.linspace(beta_start, beta_end, noise_step)
-            self.betas = DDPM.DDPM_ForwardProcess.enforce_zero_terminal_snr(self.betas)
+            self.betas = DDPM.DDPMForwardProcess.enforce_zero_terminal_snr(self.betas)
             self.alphas = 1 - self.betas
             self.alpha_hat = torch.concat((torch.ones(1), torch.cumprod(self.alphas, dim=0)))
 
+        @staticmethod
         def enforce_zero_terminal_snr(betas: torch.Tensor):
             alphas = 1 - betas
             alphas_bar = alphas.cumprod(0)
             alphas_bar_sqrt = alphas_bar.sqrt()
             alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
-            alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
-            alphas_bar_sqrt -= alphas_bar_sqrt_T
-            alphas_bar_sqrt *= alphas_bar_sqrt_0 / (
-            alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
-            alphas_bar = alphas_bar_sqrt ** 2
+            alphas_bar_sqrt_t = alphas_bar_sqrt[-1].clone()
+            alphas_bar_sqrt -= alphas_bar_sqrt_t
+            alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_t)
+            alphas_bar = alphas_bar_sqrt**2
             alphas = alphas_bar[1:] / alphas_bar[:-1]
             alphas = torch.cat([alphas_bar[0:1], alphas])
             betas = 1 - alphas
             return betas
 
-        def forward(self, input: torch.Tensor, t: torch.Tensor, eta: torch.Tensor) -> torch.Tensor:
-            alpha_hat_t = self.alpha_hat[t.cpu()].to(input.device).reshape(input.shape[0], *[1 for _ in range(len(input.shape)-1)])
-            result = alpha_hat_t.sqrt() * input + (1 - alpha_hat_t).sqrt() * eta
+        def forward(self, tensor: torch.Tensor, t: torch.Tensor, eta: torch.Tensor) -> torch.Tensor:
+            alpha_hat_t = (
+                self.alpha_hat[t.cpu()]
+                .to(tensor.device)
+                .reshape(tensor.shape[0], *[1 for _ in range(len(tensor.shape) - 1)])
+            )
+            result = alpha_hat_t.sqrt() * tensor + (1 - alpha_hat_t).sqrt() * eta
             return result
-    
-    class DDPM_SampleT(torch.nn.Module):
+
+    class DDPMSampleT(torch.nn.Module):
 
         def __init__(self, noise_step: int) -> None:
             super().__init__()
             self.noise_step = noise_step
-    
-        def forward(self, input: torch.Tensor) -> torch.Tensor:
-            return torch.randint(0, self.noise_step, (input.shape[0],)).to(input.device)
-    
-    class DDPM_TimeEmbedding(torch.nn.Module):
-        
+
+        def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+            return torch.randint(0, self.noise_step, (tensor.shape[0],)).to(tensor.device)
+
+    class DDPMTimeEmbedding(torch.nn.Module):
+
+        @staticmethod
         def sinusoidal_embedding(noise_step: int, time_embedding_dim: int):
             noise_step += 1
             embedding = torch.zeros(noise_step, time_embedding_dim)
             wk = torch.tensor([1 / 10_000 ** (2 * j / time_embedding_dim) for j in range(time_embedding_dim)])
             wk = wk.reshape((1, time_embedding_dim))
             t = torch.arange(noise_step).reshape((noise_step, 1))
-            embedding[:,::2] = torch.sin(t * wk[:,::2])
-            embedding[:,1::2] = torch.cos(t * wk[:,::2])
+            embedding[:, ::2] = torch.sin(t * wk[:, ::2])
+            embedding[:, 1::2] = torch.cos(t * wk[:, ::2])
             return embedding
 
-        def __init__(self, noise_step: int=1000, time_embedding_dim: int=100) -> None:
+        def __init__(self, noise_step: int = 1000, time_embedding_dim: int = 100) -> None:
             super().__init__()
             self.time_embed = torch.nn.Embedding(noise_step, time_embedding_dim)
-            self.time_embed.weight.data = DDPM.DDPM_TimeEmbedding.sinusoidal_embedding(noise_step, time_embedding_dim)
+            self.time_embed.weight.data = DDPM.DDPMTimeEmbedding.sinusoidal_embedding(noise_step, time_embedding_dim)
             self.time_embed.requires_grad_(False)
             self.noise_step = noise_step
-        
-        def forward(self, input: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
-            return self.time_embed((p*self.noise_step).long().repeat((input.shape[0])))
-        
-    class DDPM_UNet(network.ModuleArgsDict):
 
-        def __init__(self, noise_step: int, channels: list[int], blockConfig: blocks.BlockConfig, nb_conv_per_stage: int, downSampleMode: str, upSampleMode: str, attention : bool, time_embedding_dim: int, dim : int) -> None:
+        def forward(self, tensor: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
+            return self.time_embed((p * self.noise_step).long().repeat(tensor.shape[0]))
+
+    class DDPMUNet(network.ModuleArgsDict):
+
+        def __init__(
+            self,
+            noise_step: int,
+            channels: list[int],
+            block_config: blocks.BlockConfig,
+            nb_conv_per_stage: int,
+            downsample_mode: str,
+            upsample_mode: str,
+            attention: bool,
+            time_embedding_dim: int,
+            dim: int,
+        ) -> None:
             super().__init__()
-            self.add_module("t", DDPM.DDPM_TimeEmbedding(noise_step, time_embedding_dim), in_branch=[1], out_branch=["te"])
-            self.add_module("UNetBlock_0", DDPM.DDPM_UNetBlock(channels, nb_conv_per_stage, blockConfig, downSampleMode=blocks.DownSampleMode._member_map_[downSampleMode], upSampleMode=blocks.UpSampleMode._member_map_[upSampleMode], attention=attention, time_embedding_dim=time_embedding_dim, dim=dim), in_branch=[0,"te"])
-            self.add_module("Head", DDPM.DDPM_UNetHead(in_channels=channels[1], out_channels=1, dim=dim))
+            self.add_module(
+                "t",
+                DDPM.DDPMTimeEmbedding(noise_step, time_embedding_dim),
+                in_branch=[1],
+                out_branch=["te"],
+            )
+            self.add_module(
+                "UNetBlock_0",
+                DDPM.DDPMUNetBlock(
+                    channels,
+                    nb_conv_per_stage,
+                    block_config,
+                    downsample_mode=blocks.DownsampleMode[downsample_mode],
+                    upsample_mode=blocks.UpsampleMode[upsample_mode],
+                    attention=attention,
+                    time_embedding_dim=time_embedding_dim,
+                    dim=dim,
+                ),
+                in_branch=[0, "te"],
+            )
+            self.add_module(
+                "Head",
+                DDPM.DDPMUNetHead(in_channels=channels[1], out_channels=1, dim=dim),
+            )
 
-    class DDPM_Inference(torch.nn.Module):
+    class DDPMInference(torch.nn.Module):
 
-        def __init__(self, model: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], train_noise_step: int, inference_noise_step: int, beta_start: float, beta_end: float) -> None:
+        def __init__(
+            self,
+            model: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+            train_noise_step: int,
+            inference_noise_step: int,
+            beta_start: float,
+            beta_end: float,
+        ) -> None:
             super().__init__()
             self.model = model
             self.train_noise_step = train_noise_step
             self.inference_noise_step = inference_noise_step
-            self.forwardProcess = DDPM.DDPM_ForwardProcess(train_noise_step, beta_start, beta_end)
-        
-        def forward(self, input: torch.Tensor) -> torch.Tensor:
-            x = torch.randn_like(input).to(input.device)
+            self.forwardProcess = DDPM.DDPMForwardProcess(train_noise_step, beta_start, beta_end)
+
+        def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+            x = torch.randn_like(tensor).to(tensor.device)
             result = []
             result.append(x.unsqueeze(1))
-            description = lambda : "Inference : "+gpuInfo(input.device)
             offset = self.train_noise_step // self.inference_noise_step
-            t_list = np.round(np.arange(self.train_noise_step-1, 0, -offset)).astype(int).tolist()
-            with tqdm.tqdm(iterable = enumerate(t_list), desc = description(), total=len(t_list), leave=False, disable=True) as batch_iter:
-                for i, t in batch_iter:
+            t_list = np.round(np.arange(self.train_noise_step - 1, 0, -offset)).astype(int).tolist()
+            with tqdm.tqdm(
+                iterable=enumerate(t_list),
+                desc="Inference : " + gpu_info(),
+                total=len(t_list),
+                leave=False,
+                disable=True,
+            ) as batch_iter:
+                for _, t in batch_iter:
                     alpha_t_hat = self.forwardProcess.alpha_hat[t]
-                    alpha_t_hat_prev = self.forwardProcess.alpha_hat[t - offset] if t - offset >= 0 else self.forwardProcess.alpha_hat[0]
-                    eta_theta = self.model(torch.concat((x, input), dim=1), (torch.ones(input.shape[0], 1) * t).to(input.device).long())
+                    alpha_t_hat_prev = (
+                        self.forwardProcess.alpha_hat[t - offset]
+                        if t - offset >= 0
+                        else self.forwardProcess.alpha_hat[0]
+                    )
+                    eta_theta = self.model(
+                        torch.concat((x, tensor), dim=1),
+                        (torch.ones(tensor.shape[0], 1) * t).to(tensor.device).long(),
+                    )
 
-                    predicted = (x - (1 - alpha_t_hat).sqrt() * eta_theta)/alpha_t_hat.sqrt()*alpha_t_hat_prev.sqrt()
+                    predicted = (
+                        (x - (1 - alpha_t_hat).sqrt() * eta_theta) / alpha_t_hat.sqrt() * alpha_t_hat_prev.sqrt()
+                    )
 
                     variance = ((1 - alpha_t_hat_prev) / (1 - alpha_t_hat)) * (1 - alpha_t_hat / alpha_t_hat_prev)
-                    
-                    direction = (1-alpha_t_hat_prev-variance).sqrt()*eta_theta
-                    
-                    x = predicted+direction
-                    x += variance.sqrt()*torch.randn_like(input).to(input.device)
-                    
+
+                    direction = (1 - alpha_t_hat_prev - variance).sqrt() * eta_theta
+
+                    x = predicted + direction
+                    x += variance.sqrt() * torch.randn_like(tensor).to(tensor.device)
+
                     result.append(x.unsqueeze(1))
-                    
-                    batch_iter.set_description(description())
+
+                    batch_iter.set_description("Inference : " + gpu_info())
             result[-1] = torch.clip(result[-1], -1, 1)
             return torch.concat(result, dim=1)
 
-    class DDPM_VLoss(torch.nn.Module):
+    class DDPMVLoss(torch.nn.Module):
 
         def __init__(self, alpha_hat: torch.Tensor) -> None:
             super().__init__()
             self.alpha_hat = alpha_hat
 
-        def v(self, input: torch.Tensor, noise: torch.Tensor, t: torch.Tensor):
-            alpha_hat_t = self.alpha_hat[t.cpu()].to(input.device).reshape(input.shape[0], *[1 for _ in range(len(input.shape)-1)])
-            return alpha_hat_t.sqrt()*input-(1-alpha_hat_t).sqrt()*noise
-            
-        def forward(self, input: torch.Tensor, eta: torch.Tensor, eta_hat: torch.Tensor, t: int) -> torch.Tensor:
-            return torch.concat((self.v(input, eta, t), self.v(input, eta_hat, t)), dim=1)
-
-    def __init__(   self,
-                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
-                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
-                    patch : Union[ModelPatch, None] = None,
-                    train_noise_step: int = 1000,
-                    inference_noise_step: int = 100,
-                    beta_start: float = 1e-4, 
-                    beta_end: float = 0.02,
-                    time_embedding_dim: int = 100,
-                    channels: list[int]=[1, 64, 128, 256, 512, 1024],
-                    blockConfig: blocks.BlockConfig = blocks.BlockConfig(),
-                    nb_conv_per_stage: int = 2,
-                    downSampleMode: str = "MAXPOOL",
-                    upSampleMode: str = "CONV_TRANSPOSE",
-                    attention : bool = False,
-                    dim : int = 3) -> None:
-        super().__init__(in_channels=1, optimizer=optimizer, schedulers=schedulers, outputsCriterions=outputsCriterions, patch=patch, dim=dim)
+        def v(self, tensor: torch.Tensor, noise: torch.Tensor, t: torch.Tensor):
+            alpha_hat_t = (
+                self.alpha_hat[t.cpu()]
+                .to(tensor.device)
+                .reshape(tensor.shape[0], *[1 for _ in range(len(tensor.shape) - 1)])
+            )
+            return alpha_hat_t.sqrt() * tensor - (1 - alpha_hat_t).sqrt() * noise
+
+        def forward(self, tensor: torch.Tensor, eta: torch.Tensor, eta_hat: torch.Tensor, t: int) -> torch.Tensor:
+            return torch.concat((self.v(tensor, eta, t), self.v(tensor, eta_hat, t)), dim=1)
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        patch: ModelPatch | None = None,
+        train_noise_step: int = 1000,
+        inference_noise_step: int = 100,
+        beta_start: float = 1e-4,
+        beta_end: float = 0.02,
+        time_embedding_dim: int = 100,
+        channels: list[int] = [1, 64, 128, 256, 512, 1024],
+        block_config: blocks.BlockConfig = blocks.BlockConfig(),
+        nb_conv_per_stage: int = 2,
+        downsample_mode: str = "MAXPOOL",
+        upsample_mode: str = "CONV_TRANSPOSE",
+        attention: bool = False,
+        dim: int = 3,
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            patch=patch,
+            dim=dim,
+        )
         self.add_module("Noise", blocks.NormalNoise(), out_branch=["eta"], training=True)
-        self.add_module("Sample", DDPM.DDPM_SampleT(train_noise_step), out_branch=["t"], training=True)
-        self.add_module("Forward", DDPM.DDPM_ForwardProcess(train_noise_step, beta_start, beta_end), in_branch=[0, "t", "eta"], out_branch=["x_t"], training=True)
-        self.add_module("Concat", blocks.Concat(), in_branch=["x_t",1], out_branch=["xy_t"], training=True)
-        self.add_module("UNet", DDPM.DDPM_UNet(train_noise_step, channels, blockConfig, nb_conv_per_stage, downSampleMode, upSampleMode, attention, time_embedding_dim, dim), in_branch=["xy_t", "t"], out_branch=["eta_hat"], training=True)
-
-        self.add_module("Noise_optim", DDPM.DDPM_VLoss(self["Forward"].alpha_hat), in_branch=[0, "eta", "eta_hat", "t"], out_branch=["noise"], training=True)
-        
-        self.add_module("Inference", DDPM.DDPM_Inference(self.inference, train_noise_step, inference_noise_step, beta_start, beta_end), in_branch=[1], training=False)
-        self.add_module("LastImage", blocks.Select([slice(None, None), slice(-1, None)]), training=False)
-    
-    def inference(self, input: torch.Tensor, t: torch.Tensor):
-        return self["UNet"](input, t)
+        self.add_module(
+            "Sample",
+            DDPM.DDPMSampleT(train_noise_step),
+            out_branch=["t"],
+            training=True,
+        )
+        self.add_module(
+            "Forward",
+            DDPM.DDPMForwardProcess(train_noise_step, beta_start, beta_end),
+            in_branch=[0, "t", "eta"],
+            out_branch=["x_t"],
+            training=True,
+        )
+        self.add_module(
+            "Concat",
+            blocks.Concat(),
+            in_branch=["x_t", 1],
+            out_branch=["xy_t"],
+            training=True,
+        )
+        self.add_module(
+            "UNet",
+            DDPM.DDPMUNet(
+                train_noise_step,
+                channels,
+                block_config,
+                nb_conv_per_stage,
+                downsample_mode,
+                upsample_mode,
+                attention,
+                time_embedding_dim,
+                dim,
+            ),
+            in_branch=["xy_t", "t"],
+            out_branch=["eta_hat"],
+            training=True,
+        )
+
+        self.add_module(
+            "Noise_optim",
+            DDPM.DDPMVLoss(self["Forward"].alpha_hat),
+            in_branch=[0, "eta", "eta_hat", "t"],
+            out_branch=["noise"],
+            training=True,
+        )
+
+        self.add_module(
+            "Inference",
+            DDPM.DDPMInference(
+                self.inference,
+                train_noise_step,
+                inference_noise_step,
+                beta_start,
+                beta_end,
+            ),
+            in_branch=[1],
+            training=False,
+        )
+        self.add_module(
+            "LastImage",
+            blocks.Select([slice(None, None), slice(-1, None)]),
+            training=False,
+        )
+
+    def inference(self, tensor: torch.Tensor, t: torch.Tensor):
+        return self["UNet"](tensor, t)
+
 
 class MSE(Criterion):
 
     def __init__(self):
         super().__init__()
         self.loss = torch.nn.MSELoss()
-    
-    def forward(self, input: torch.Tensor) -> torch.Tensor:
-        return self.loss(input[:, 0, ...], input[:, 1, ...])
+
+    def forward(self, tensor: torch.Tensor, *targets: list[torch.Tensor]) -> torch.Tensor:
+        return self.loss(tensor[:, 0, ...], tensor[:, 1, ...])