ok

.gitignore

@@ -161,8 +161,8 @@ cython_debug/
 *.jpeg
 *.gz
 cifar-10-batches-py
-checkpoints
 MNIST
 *.ipynb
 data
 wandb
+/checkpoints/lightning_logs

checkpoints/{cifar.ckpt → model/celebahq.ckpt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d9efb3494902fa10ab74f65d670a114a5470cfb879fa60bdd0292956895de587
-size 278317592
+oid sha256:bfc9fa8cb71bc57bc4d1f54da56e71060609828bc6903cec3ae46418c18bf3a1
+size 99080226

checkpoints/model/cifar10.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:509e43eb3be202b3d71ef37ca5b66de501fccda8806e4a193a466bcbfcb71b83
+size 99090784

checkpoints/{mnist.ckpt → model/mnist.ckpt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:584b8c9097fc8022291ad2760512bf5cef8bee1b1f4fc52b211e112dca44c643
-size 278302296
+oid sha256:d2cde8ee89e68b413c32685145b4ad1fea10b7a6617c0f14a12a9af8afac9712
+size 99081632

diffusion/dataset/celeba.py

@@ -11,12 +11,13 @@ class CelebADataset(Dataset):
     def __init__(
             self,
             data_dir: str,
+            img_dim: int = 64
     ):
         self.list_path = os.listdir(data_dir)
         self.data_dir = data_dir
         self.transform = transforms.Compose(
             [
-                transforms.Resize((64, 64)),
+                transforms.Resize((img_dim, img_dim)),
                 transforms.ToTensor(),
                 transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
             ]
@@ -37,13 +38,15 @@ class CelebADataModule(pl.LightningDataModule):
         batch_size: int = 32,
         num_workers: int = 0,
         seed: int = 42,
-        train_ratio: float = 0.99
+        train_ratio: float = 0.99,
+        img_dim: int = 64
     ):
         super().__init__()
         self.data_dir = data_dir
         self.batch_size = batch_size
         self.num_workers = num_workers
         self.train_ratio = min(train_ratio, 0.99)
+        self.img_dim = img_dim
         self.seed = seed
 
         self.loader = partial(
@@ -56,7 +59,7 @@ class CelebADataModule(pl.LightningDataModule):
 
     def setup(self, stage: str):
         if stage == "fit":
-            dataset = CelebADataset(self.data_dir)
+            dataset = CelebADataset(self.data_dir, self.img_dim)
             self.CelebA_train, self.CelebA_val, _ = random_split(
                 dataset=dataset,
                 lengths=[self.train_ratio, 0.01, 1 - 0.01 - self.train_ratio],

diffusion/dataset/cifar10.py

@@ -13,7 +13,8 @@ class CIFAR10DataModule(pl.LightningDataModule):
         batch_size: int = 32,
         num_workers: int = 0,
         seed: int = 42,
-        train_ratio: float = 0.99
+        train_ratio: float = 0.99,
+        img_dim: int = 32
     ):
         super().__init__()
         self.data_dir = data_dir
@@ -23,7 +24,7 @@ class CIFAR10DataModule(pl.LightningDataModule):
         self.train_ratio = min(train_ratio, 0.99)
         self.transform = transforms.Compose(
             [
-                transforms.Resize((32, 32)),
+                transforms.Resize((img_dim, img_dim)),
                 transforms.ToTensor(),
                 transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
             ]

diffusion/dataset/mnist.py

@@ -13,7 +13,8 @@ class MNISTDataModule(pl.LightningDataModule):
         batch_size: int = 32,
         num_workers: int = 0,
         seed: int = 42,
-        train_ratio: float = 0.99
+        train_ratio: float = 0.99,
+        img_dim: int = 32
     ):
         super().__init__()
         self.data_dir = data_dir
@@ -23,7 +24,7 @@ class MNISTDataModule(pl.LightningDataModule):
         self.seed = seed
         self.transform = transforms.Compose(
             [
-                transforms.Resize((32, 32)),
+                transforms.Resize((img_dim, img_dim)),
                 transforms.ToTensor(),
                 transforms.Normalize(mean=(0.5), std=(0.5))
             ]

diffusion/model/diffusion/__init__.py

@@ -1,4 +1,3 @@
 from .unet import *
 from .model import *
-from .sampling import *
 from .scheduler import *

diffusion/model/diffusion/model.py

@@ -4,8 +4,10 @@ import numpy as np
 import pytorch_lightning as pl
 import diffusion
 import wandb
+import matplotlib.pyplot as plt
 from torchvision.utils import make_grid
 from torch.optim.lr_scheduler import OneCycleLR
+from IPython.display import clear_output
 
 
 class DiffusionModel(pl.LightningModule):
@@ -16,6 +18,7 @@ class DiffusionModel(pl.LightningModule):
         beta_1: float = 0.0001,
         beta_2: float = 0.02,
         in_channels: int = 3,
+        mode: str = "ddpm",
         dim: int = 32,
         num_classes: int | None = 10,
         sample_per_epochs: int = 50,
@@ -33,11 +36,17 @@ class DiffusionModel(pl.LightningModule):
         self.max_timesteps = max_timesteps
         self.in_channels = in_channels
         self.dim = dim
+        self.mode = mode
         self.num_classes = num_classes
 
-        self.scheduler = diffusion.LinearScheduler(
-            max_timesteps, beta_1, beta_2
-        )
+        if mode == "ddpm":
+            self.scheduler = diffusion.DDPMScheduler(
+                max_timesteps, beta_1, beta_2
+            )
+        elif mode == "ddim":
+            self.scheduler = diffusion.DDIMScheduler(
+                max_timesteps, beta_1, beta_2
+            )
 
         self.criterion = nn.MSELoss()
 
@@ -49,8 +58,6 @@ class DiffusionModel(pl.LightningModule):
 
         self.sampling_kwargs = {
             'model': self.model,
-            'scheduler': self.scheduler,
-            'max_timesteps': self.max_timesteps,
             'in_channels': self.in_channels,
             'dim': self.dim,
         }
@@ -75,29 +82,37 @@ class DiffusionModel(pl.LightningModule):
         x_0: torch.Tensor,
         t: torch.Tensor
     ):
-        noise = torch.randn_like(x_0, device=x_0.device)
-        new_x = self.scheduler.get('sqrt_alpha_hat', t) * x_0
-        new_noise = self.scheduler.get('sqrt_one_minus_alpha_hat', t) * noise
-        return new_x + new_noise, noise
+        return self.scheduler.noising(x_0, t)
 
-    def sampling(self, labels=None, n_samples: int = 16):
-        return diffusion.ddpm_sampling(
-            n_samples=n_samples,
-            labels=labels,
-            **self.sampling_kwargs
-        )
-
-    def sampling_demo(self, labels=None, n_samples: int = 16):
-        return diffusion.ddpm_sampling_demo(
-            n_samples=n_samples,
-            labels=labels,
-            **self.sampling_kwargs
-        )
+    def sampling(
+        self,
+        labels=None,
+        mode: int = "ddpm",
+        demo: bool = True,
+        n_samples: int = 16,
+        timesteps: int = 1000,
+    ):
+        if mode == "ddpm":
+            self.test_scheduler = diffusion.DDPMScheduler(self.max_timesteps)
+        elif mode == "ddim":
+            self.test_scheduler = diffusion.DDIMScheduler(self.max_timesteps)
+
+        kwargs = {
+            "n_samples": n_samples,
+            "labels": labels,
+            "timesteps": timesteps,
+        } | self.sampling_kwargs
+        if demo:
+            return self.test_scheduler.sampling_demo(**kwargs)
+        else:
+            return self.test_scheduler.sampling(**kwargs)
 
     def forward(self, x_0, labels):
+        n = x_0.shape[0]
         t = torch.randint(
-            low=0, high=self.max_timesteps, size=(x_0.shape[0],), device=x_0.device
+            low=0, high=self.max_timesteps, size=(n//2+1,), device=x_0.device
         )
+        t = torch.cat([t, self.max_timesteps - t - 1], dim=0)[:n]
         x_noise, noise = self.noising(x_0, t)
         noise_pred = self.model(x_noise, t, labels)
         return noise, noise_pred
@@ -108,8 +123,8 @@ class DiffusionModel(pl.LightningModule):
             labels = None
         else:
             x_0, labels = batch
-        if np.random.random() < 0.1:
-            labels = None
+            if np.random.random() < 0.1:
+                labels = None
         noise, noise_pred = self(x_0, labels)
         loss = self.criterion(noise, noise_pred)
         self.train_loss.append(loss)
@@ -135,19 +150,20 @@ class DiffusionModel(pl.LightningModule):
         )
         self.train_loss.clear()
 
-        if self.epoch_count % self.spe == 0:
-            wandblog = self.logger.experiment
-            x_t = self.sampling(n_samples=self.n_samples)
-            img_array = [x_t[i] for i in range(x_t.shape[0])]
-
-            wandblog.log(
-                {
-                    "sampling": wandb.Image(
-                        make_grid(img_array, nrow=4).permute(1, 2, 0).cpu().numpy(),
-                        caption="Sampled Image!"
-                    )
-                }
-            )
+        if self.spe > 0:
+            if self.epoch_count % self.spe == 0:
+                wandblog = self.logger.experiment
+                x_t = self.sampling(n_samples=self.n_samples, timesteps=100, demo=False)
+                img_array = [x_t[i] for i in range(x_t.shape[0])]
+
+                wandblog.log(
+                    {
+                        "sampling": wandb.Image(
+                            make_grid(img_array, nrow=4).permute(1, 2, 0).cpu().numpy(),
+                            caption="Sampled Image!"
+                        )
+                    }
+                )
 
         self.epoch_count += 1
 
@@ -173,10 +189,34 @@ class DiffusionModel(pl.LightningModule):
             total_steps=self.trainer.estimated_stepping_batches,
 
         )
-        return {
-            'optimizer': optimizer,
-            'lr_scheduler': scheduler
-        }
+        return [optimizer], [scheduler]
+
+    def draw(
+        self,
+        labels=None,
+        mode: int = "ddpm",
+        n_samples: int = 1,
+        timesteps: int = 1000,
+    ):
+        demo = self.sampling(
+            labels=labels,
+            mode=mode,
+            n_samples=n_samples,
+            timesteps=timesteps,
+            demo=True
+        )
+        idx = 0
+        length = labels.shape[0] if labels is not None else n_samples
+        for img in demo:
+            for i in range(length):
+                plt.subplot(1, length, i+1)
+                plt.imshow(img[i].permute(1, 2, 0))
+                plt.axis('off')
+            plt.title(f"{idx+1}/{timesteps}")
+            idx += 1
+            plt.show()
+            if idx < timesteps:
+                clear_output(wait=True)
 
 
 if __name__ == "__main__":

diffusion/model/diffusion/sampling.py

@@ -1,82 +0,0 @@
-import torch
-
-
-def ddpm_sampling_timestep(
-    x_t,
-    model,
-    scheduler,
-    labels,
-    t,
-    n_samples: int = 16,
-    cfg_scale: int = 3,
-):
-    time = torch.full((n_samples,), fill_value=t, device=model.device)
-    pred_noise = model(x_t, time, labels)
-    if cfg_scale > 0:
-        uncond_pred_noise = model(x_t, time, None)
-        pred_noise = torch.lerp(uncond_pred_noise, pred_noise, cfg_scale)
-    alpha = scheduler.get('alpha', time)
-    sqrt_alpha = scheduler.get('sqrt_alpha', time)
-    somah = scheduler.get('sqrt_one_minus_alpha_hat', time)
-    sqrt_beta = scheduler.get('sqrt_beta', time)
-    if t > 0:
-        noise = torch.randn_like(x_t, device=model.device)
-    else:
-        noise = torch.zeros_like(x_t, device=model.device)
-
-    x_t_new = 1 / sqrt_alpha * (x_t - (1-alpha) / somah * pred_noise) + sqrt_beta * noise
-    return x_t_new.clamp(-1, 1)
-
-
-@torch.no_grad()
-def ddpm_sampling(
-    model,
-    scheduler,
-    n_samples: int = 16,
-    max_timesteps: int = 1000,
-    in_channels: int = 3,
-    dim: int = 32,
-    cfg_scale: int = 3,
-
-    labels=None
-):
-    if labels is not None:
-        n_samples = labels.shape[0]
-
-    x_t = torch.randn(
-        n_samples, in_channels, dim, dim, device=model.device
-    )
-    model.eval()
-    for t in range(max_timesteps-1, -1, -1):
-        x_t = ddpm_sampling_timestep(x_t=x_t, model=model, scheduler=scheduler,
-                                     labels=labels, t=t, n_samples=n_samples,
-                                     cfg_scale=cfg_scale)
-
-    model.train()
-    x_t = (x_t + 1) / 2 * 255.  # range [0,255]
-    return x_t.type(torch.uint8)
-
-
-@torch.no_grad()
-def ddpm_sampling_demo(
-    model,
-    scheduler,
-    n_samples: int = 16,
-    max_timesteps: int = 1000,
-    in_channels: int = 3,
-    dim: int = 32,
-    cfg_scale: int = 3,
-    labels=None
-):
-    if labels is not None:
-        n_samples = labels.shape[0]
-
-    x_t = torch.randn(
-        n_samples, in_channels, dim, dim, device=model.device
-    )
-    model.eval()
-    for t in range(max_timesteps-1, -1, -1):
-        x_t = ddpm_sampling_timestep(x_t=x_t, model=model, scheduler=scheduler,
-                                     labels=labels, t=t, n_samples=n_samples,
-                                     cfg_scale=cfg_scale)
-        yield ((x_t + 1) / 2 * 255).type(torch.uint8)

diffusion/model/diffusion/scheduler.py

@@ -1,20 +1,184 @@
 import torch
 
 
-class LinearScheduler:
+class DDPMScheduler:
     def __init__(
         self,
         max_timesteps: int = 1000,
         beta_1: int = 0.0001,
         beta_2: int = 0.02
     ) -> None:
-        self.beta = torch.linspace(beta_1, beta_2, max_timesteps)
-        self.sqrt_beta = torch.sqrt(self.beta)[:, None, None, None]
-        self.alpha = (1 - self.beta)[:, None, None, None]
+        self.beta_1 = beta_1
+        self.beta_2 = beta_2
+        self.max_timesteps = max_timesteps
+        self._init_params()
+
+    def _init_params(self, timesteps: int | None = None):
+        self.beta = torch.linspace(self.beta_1, self.beta_2, timesteps or self.max_timesteps)
+        self.sqrt_beta = torch.sqrt(self.beta)
+        self.alpha = (1 - self.beta)
+        self.sqrt_alpha = torch.sqrt(self.alpha)
+        self.alpha_hat = torch.cumprod(1 - self.beta, dim=0)
+        self.sqrt_alpha_hat = torch.sqrt(self.alpha_hat)
+        self.sqrt_one_minus_alpha = torch.sqrt(1 - self.alpha)
+        self.sqrt_one_minus_alpha_hat = torch.sqrt(1 - self.alpha_hat)
+
+    def noising(
+        self,
+        x_0: torch.Tensor,
+        t: torch.Tensor
+    ):
+        if t.device != x_0.device:
+            t = t.to(x_0.device)
+        noise = torch.randn_like(x_0, device=x_0.device)
+        new_x = self.sqrt_alpha_hat.to(x_0.device)[t][:, None, None, None] * x_0
+        new_noise = self.sqrt_one_minus_alpha_hat.to(x_0.device)[t][:, None, None, None] * noise
+        return new_x + new_noise, noise
+
+    @torch.no_grad()
+    def sampling_t(
+        self,
+        x_t: torch.Tensor,
+        model,
+        labels: torch.Tensor,
+        timesteps: int,
+        t: int,
+        n_samples: int = 16,
+        cfg_scale: int = 3,
+    ):
+        time = torch.full((n_samples,), fill_value=t, device=model.device)
+        pred_noise = model(x_t, time, labels)
+        if cfg_scale > 0 and labels is not None:
+            uncond_pred_noise = model(x_t, time, None)
+            pred_noise = torch.lerp(uncond_pred_noise, pred_noise, cfg_scale)
+        alpha = self.alpha.to(model.device)[time][:, None, None, None]
+        sqrt_alpha = self.sqrt_alpha.to(model.device)[time][:, None, None, None]
+        somah = self.sqrt_one_minus_alpha_hat.to(model.device)[time][:, None, None, None]
+        sqrt_beta = self.sqrt_beta.to(model.device)[time][:, None, None, None]
+        if t > 1:
+            noise = torch.randn_like(x_t, device=model.device)
+        else:
+            noise = torch.zeros_like(x_t, device=model.device)
+
+        x_t_new = 1 / sqrt_alpha * (x_t - (1-alpha) / somah * pred_noise) + sqrt_beta * noise
+        return x_t_new.clamp(-1, 1)
+
+    @torch.no_grad()
+    def sampling(
+        self,
+        model,
+        n_samples: int = 16,
+        in_channels: int = 3,
+        dim: int = 32,
+        timesteps: int = 1000,
+        cfg_scale: int = 3,
+        labels=None,
+        *args, **kwargs
+    ):
+        if labels is not None:
+            n_samples = labels.shape[0]
+        model.eval()
+        x_t = torch.randn(
+            n_samples, in_channels, dim, dim, device=model.device
+        )
+        step_ratios = self.max_timesteps // timesteps
+        all_timesteps = torch.flip(torch.arange(0, timesteps) * step_ratios, dims=(0,))
+        for t in all_timesteps:
+            x_t = self.sampling_t(x_t=x_t, model=model, labels=labels, t=t, timesteps=timesteps,
+                                  n_samples=n_samples, cfg_scale=cfg_scale)
+        model.train()
+        x_t = (x_t.clamp(-1, 1) + 1) / 2 * 255.  # range [0,255]
+        return x_t.type(torch.uint8)
+
+    @torch.no_grad()
+    def sampling_demo(
+        self,
+        model,
+        n_samples: int = 16,
+        in_channels: int = 3,
+        dim: int = 32,
+        timesteps: int = 1000,
+        cfg_scale: int = 3,
+        labels=None,
+        *args, **kwargs
+    ):
+        if labels is not None:
+            n_samples = labels.shape[0]
+
+        x_t = torch.randn(
+            n_samples, in_channels, dim, dim, device=model.device
+        )
+        model.eval()
+        step_ratios = self.max_timesteps // timesteps
+        all_timesteps = torch.flip(torch.arange(0, timesteps) * step_ratios, dims=(0,))
+        for t in all_timesteps:
+            x_t = self.sampling_t(x_t=x_t, model=model, labels=labels, t=t, timesteps=timesteps,
+                                  n_samples=n_samples, cfg_scale=cfg_scale)
+            yield ((x_t.clamp(-1, 1) + 1) / 2 * 255).type(torch.uint8)
+
+
+class DDIMScheduler(DDPMScheduler):
+    def __init__(
+        self,
+        max_timesteps: int = 1000,
+        beta_1: int = 0.0001,
+        beta_2: int = 0.02
+    ) -> None:
+        super().__init__(beta_1=beta_1, beta_2=beta_2, max_timesteps=max_timesteps)
+        self._init_params()
+
+    def _init_params(self, timesteps: int | None = None):
+        self.beta = torch.linspace(self.beta_1, self.beta_2, timesteps or self.max_timesteps)
+        self.sqrt_beta = torch.sqrt(self.beta)
+        self.alpha = (1 - self.beta)
         self.sqrt_alpha = torch.sqrt(self.alpha)
-        self.alpha_hat = torch.cumprod(1 - self.beta, dim=0)[:, None, None, None]
+        self.alpha_hat = torch.cumprod(1 - self.beta, dim=0)
         self.sqrt_alpha_hat = torch.sqrt(self.alpha_hat)
+        self.sqrt_one_minus_alpha = torch.sqrt(1 - self.alpha)
         self.sqrt_one_minus_alpha_hat = torch.sqrt(1 - self.alpha_hat)
+        self.alpha_hat_prev = torch.cat([torch.tensor([1.]), self.alpha_hat], dim=0)[:-1]
+        self.variance = (1 - self.alpha_hat_prev) / (1 - self.alpha_hat) * \
+            (1 - self.alpha_hat / self.alpha_hat_prev)
+
+    @torch.no_grad()
+    def sampling_t(
+        self,
+        x_t: torch.Tensor, model, t: int,
+        timesteps: int,
+        labels: torch.Tensor | None = None,
+        n_samples: int = 16,
+        eta: float = 0.0,
+        *args, **kwargs
+    ):
+        time = torch.full((n_samples,), fill_value=t, device=model.device)
+        time_prev = time - self.max_timesteps // timesteps
+        pred_noise = model(x_t, time, labels)
+
+        sqrt_one_minus_alpha_hat = self.sqrt_one_minus_alpha_hat.to(model.device)[time][:, None, None, None]
+        sqrt_alpha_hat = self.sqrt_alpha_hat.to(model.device)[time][:, None, None, None]
+        alpha_hat_prev = self.alpha_hat[time_prev] if time_prev[0] >= 0 else torch.ones_like(time_prev)
+        alpha_hat_prev = alpha_hat_prev.to(model.device)[:, None, None, None]
+        sqrt_alpha_hat_prev = torch.sqrt(alpha_hat_prev)
+        posterior_std = torch.sqrt(self.variance)[time][:, None, None, None] * eta
+
+        if t > 0:
+            noise = torch.randn_like(x_t, device=model.device)
+        else:
+            noise = torch.zeros_like(x_t, device=model.device)
+
+        x_0_pred = (x_t - sqrt_one_minus_alpha_hat * pred_noise) / sqrt_alpha_hat
+        x_0_pred = x_0_pred.clamp(-1, 1)
+        x_t_direction = torch.sqrt(1. - alpha_hat_prev - posterior_std**2) * pred_noise
+        random_noise = posterior_std * noise
+        x_t_1 = sqrt_alpha_hat_prev * x_0_pred + x_t_direction + random_noise
+
+        return x_t_1
+
 
-    def get(self, key: str, t: torch.Tensor):
-        return self.__dict__[key].to(t.device)[t]
+if __name__ == "__main__":
+    dct = DDIMScheduler().__dict__
+    for k in dct.keys():
+        if isinstance(dct[k], torch.Tensor):
+            print(k, dct[k].shape)
+        else:
+            print(k, dct[k])

diffusion/model/diffusion/unet.py

@@ -45,10 +45,10 @@ class DoubleConv(nn.Module):
             mid_channels = out_channels
         self.double_conv = nn.Sequential(
             nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
-            nn.GroupNorm(8, mid_channels),
+            nn.GroupNorm(32, mid_channels, eps=1e-6, affine=True),
             nn.GELU(),
             nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
-            nn.GroupNorm(8, out_channels),
+            nn.GroupNorm(32, out_channels, eps=1e-6, affine=True),
         )
 
     def forward(self, x):
@@ -137,15 +137,20 @@ class UNet(pl.LightningModule):
         self.sa3 = SelfAttention(channels=256)
 
         self.mid1 = DoubleConv(in_channels=256, out_channels=512)
+        self.sa4 = SelfAttention(channels=512)
         self.mid2 = DoubleConv(in_channels=512, out_channels=512)
 
         self.up1 = UpSample(in_channels=512, out_channels=256)
-        self.sa4 = SelfAttention(channels=256)
+        self.sa5 = SelfAttention(channels=256)
         self.up2 = UpSample(in_channels=256, out_channels=128)
-        self.sa5 = SelfAttention(channels=128)
+        self.sa6 = SelfAttention(channels=128)
         self.up3 = UpSample(in_channels=128, out_channels=64)
-        self.sa6 = SelfAttention(channels=64)
-        self.outc = nn.Conv2d(64, c_out, kernel_size=1)
+        self.sa7 = SelfAttention(channels=64)
+        self.outc = nn.Sequential(
+            nn.GroupNorm(32, 64, eps=1e-6, affine=True),
+            nn.SiLU(),
+            nn.Conv2d(64, c_out, kernel_size=3, padding=1)
+        )
 
     def pos_encoding(self, t, channels):
         inv_freq = 1.0 / (
@@ -168,14 +173,15 @@ class UNet(pl.LightningModule):
         x4 = self.sa3(x4)
 
         x4 = self.mid1(x4)
+        x4 = self.sa4(x4)
         x4 = self.mid2(x4)
 
         x = self.up1(x4, x3, t)
-        x = self.sa4(x)
-        x = self.up2(x, x2, t)
         x = self.sa5(x)
-        x = self.up3(x, x1, t)
+        x = self.up2(x, x2, t)
         x = self.sa6(x)
+        x = self.up3(x, x1, t)
+        x = self.sa7(x)
         output = self.outc(x)
         return output
 

diffusion/model/ldm/model.py

@@ -2,4 +2,5 @@ import torch
 import pytorch_lightning as pl
 
 class LatentDiffusionModel(pl.LightningModule):
+    # TODO
     pass

diffusion/train/__main__.py

@@ -20,6 +20,10 @@ def main():
         '--data_dir', '-dd', type=str, default='./data/',
         help='model name'
     )
+    parser.add_argument(
+        '--mode', type=str, default='ddim',
+        help='sampling mode'
+    )
     parser.add_argument(
         '--max_epochs', '-me', type=int, default=200,
         help='max epoch'
@@ -117,7 +121,8 @@ def main():
         batch_size=args.batch_size,
         num_workers=args.num_workers,
         seed=args.seed,
-        train_ratio=args.train_ratio
+        train_ratio=args.train_ratio,
+        img_dim=img_dim
     )
 
     # MODEL
@@ -129,7 +134,8 @@ def main():
         max_timesteps=args.timesteps,
         dim=img_dim,
         num_classes=num_classes,
-        n_samples=args.n_samples
+        n_samples=args.n_samples,
+        mode=args.mode
     )
 
     # CALLBACK