Merge pull request #1 from tcapelle/refactor

Refactor and simplify

Author: tcapelle

README.md

@@ -4,8 +4,7 @@
 
 # Cloud Diffusion Experiment
 
-This codebase contains an implementation of a deep diffusion model applied to cloud images. It was developed as part of a research project exploring the potential of diffusion models
-for image generation and forecasting.
+This codebase contains an implementation of a deep diffusion model applied to cloud images. It was developed as part of a research project exploring the potential of diffusion models for image generation and forecasting.
 
 ## Setup
 
@@ -18,6 +17,7 @@ for image generation and forecasting.
 To train the model, run `python train.py`. You can play with the parameters on top of the file to change the model architecture, training parameters, etc.
 
 You can also override the configuration parameters by passing them as command-line arguments, e.g.
+
 ```bash
 > python train.py --epochs=10 --batch_size=32
 ```
@@ -27,7 +27,6 @@ You can also override the configuration parameters by passing them as command-li
 This training is based on a Transformer based Unet (UViT), you can train the default model by running:
 
 ```bash
-
 > python train_uvit.py
 ```
 

cloud_diffusion/dataset.py

@@ -11,6 +11,20 @@
 PROJECT_NAME = "ddpm_clouds"
 DATASET_ARTIFACT = 'capecape/gtc/np_dataset:v0'
 
+class DummyNextFrameDataset:
+    "Dataset that returns random images"
+    def __init__(self, num_frames=4, img_size=64, N=1000):
+        self.img_size = img_size
+        self.num_frames = num_frames
+        self.N = N
+
+    def __getitem__(self, idx):
+        return torch.randn(self.num_frames, self.img_size, self.img_size)
+    
+    def __len__(self):
+        return self.N
+
+
 class CloudDataset:
     """Dataset for cloud images
     It loads numpy files from wandb artifact and stacks them into a single array

cloud_diffusion/ddpm.py

@@ -1,17 +1,10 @@
-from pathlib import Path
 from functools import partial
 
-import torch, wandb
-from torch.nn import init
-from torch.utils.data.dataloader import default_collate
-
-import fastcore.all as fc
+import torch
 from fastprogress import progress_bar
 
 from diffusers.schedulers import DDIMScheduler
 
-from diffusers import UNet2DModel
-
 
 ## DDPM params
 ## From fastai V2 Course DDPM notebooks
@@ -21,52 +14,15 @@
 alphabar = alpha.cumprod(dim=0)
 sigma = beta.sqrt()
 
-def noisify(x0, ᾱ):
-    "Noise only the last frame"
-    past_frames = x0[:,:-1]
-    x0 = x0[:,-1:]
+def noisify_ddpm(x0):
+    "Noise by ddpm"
     device = x0.device
     n = len(x0)
     t = torch.randint(0, n_steps, (n,), dtype=torch.long)
     ε = torch.randn(x0.shape, device=device)
-    ᾱ_t = ᾱ[t].reshape(-1, 1, 1, 1).to(device)
+    ᾱ_t = alphabar[t].reshape(-1, 1, 1, 1).to(device)
     xt = ᾱ_t.sqrt()*x0 + (1-ᾱ_t).sqrt()*ε
-    return torch.cat([past_frames, xt], dim=1), t.to(device), ε
-
-def collate_ddpm(b): 
-    "Collate function that noisifies the last frame"
-    return noisify(default_collate(b), alphabar)
-
-def get_unet_params(model_name="unet_small", num_frames=4):
-    "Return the parameters for the diffusers UNet2d model"
-    if model_name == "unet_small":
-        return dict(
-            block_out_channels=(16, 32, 64, 128), # number of channels for each block
-            norm_num_groups=8, # number of groups for the normalization layer
-            in_channels=num_frames, # number of input channels
-            out_channels=1, # number of output channels
-            )
-    elif model_name == "unet_big":
-        return dict(
-            block_out_channels=(32, 64, 128, 256), # number of channels for each block
-            norm_num_groups=8, # number of groups for the normalization layer
-            in_channels=num_frames, # number of input channels
-            out_channels=1, # number of output channels
-            )
-    else:
-        raise(f"Model name not found: {model_name}, choose between 'unet_small' or 'unet_big'")
-
-def init_ddpm(model):
-    "From Jeremy's bag of tricks on fastai V2 2023"
-    for o in model.down_blocks:
-        for p in o.resnets:
-            p.conv2.weight.data.zero_()
-            for p in fc.L(o.downsamplers): init.orthogonal_(p.conv.weight)
-
-    for o in model.up_blocks:
-        for p in o.resnets: p.conv2.weight.data.zero_()
-
-    model.conv_out.weight.data.zero_()
+    return xt, t.to(device), ε
 
 @torch.no_grad()
 def diffusers_sampler(model, past_frames, sched, **kwargs):
@@ -88,23 +44,3 @@ def ddim_sampler(steps=350, eta=1.):
     ddim_sched = DDIMScheduler()
     ddim_sched.set_timesteps(steps)
     return partial(diffusers_sampler, sched=ddim_sched, eta=eta)
-
-class UNet2D(UNet2DModel):
-    def forward(self, *x, **kwargs):
-        return super().forward(*x, **kwargs).sample ## Diffusers's UNet2DOutput class
-    
-    @classmethod
-    def from_checkpoint(cls, model_params, checkpoint_file):
-        "Load a UNet2D model from a checkpoint file"
-        model = cls(**model_params)
-        model.load_state_dict(torch.load(checkpoint_file, map_location="cpu"))
-        return model
-
-
-    @classmethod
-    def from_artifact(cls, model_params, artifact_name):
-        "Load a UNet2D model from a wandb.Artifact, need to be run in a wandb run"
-        artifact = wandb.use_artifact(artifact_name, type='model')
-        artifact_dir = Path(artifact.download())
-        chpt_file = list(artifact_dir.glob("*.pth"))[0]
-        return cls.from_checkpoint(model_params, chpt_file)

cloud_diffusion/models.py

@@ -0,0 +1,99 @@
+from pathlib import Path
+
+import wandb
+import fastcore.all as fc
+
+import torch
+from torch import nn
+from diffusers import UNet2DModel
+
+try:
+    from denoising_diffusion_pytorch.simple_diffusion import UViT
+except:
+    raise ImportError("Please install denoising_diffusion_pytorch with `pip install denoising_diffusion_pytorch`")
+
+
+def init_unet(model):
+    "From Jeremy's bag of tricks on fastai V2 2023"
+    for o in model.down_blocks:
+        for p in o.resnets:
+            p.conv2.weight.data.zero_()
+            for p in fc.L(o.downsamplers): nn.init.orthogonal_(p.conv.weight)
+
+    for o in model.up_blocks:
+        for p in o.resnets: p.conv2.weight.data.zero_()
+
+    model.conv_out.weight.data.zero_()
+
+class WandbModel:
+    "A model that can be saved to wandb"
+    @classmethod
+    def from_checkpoint(cls, model_params, checkpoint_file):
+        "Load a UNet2D model from a checkpoint file"
+        model = cls(**model_params)
+        print(f"Loading model from: {checkpoint_file}")
+        model.load_state_dict(torch.load(checkpoint_file))
+        return model
+
+    @classmethod
+    def from_artifact(cls, model_params, artifact_name):
+        "Load a UNet2D model from a wandb.Artifact, need to be run in a wandb run"
+        artifact = wandb.use_artifact(artifact_name, type='model')
+        artifact_dir = Path(artifact.download())
+        chpt_file = list(artifact_dir.glob("*.pth"))[0]
+        return cls.from_checkpoint(model_params, chpt_file)
+
+def get_unet_params(model_name="unet_small", num_frames=4):
+    "Return the parameters for the diffusers UNet2d model"
+    if model_name == "unet_small":
+        return dict(
+            block_out_channels=(16, 32, 64, 128), # number of channels for each block
+            norm_num_groups=8, # number of groups for the normalization layer
+            in_channels=num_frames, # number of input channels
+            out_channels=1, # number of output channels
+            )
+    elif model_name == "unet_big":
+        return dict(
+            block_out_channels=(32, 64, 128, 256), # number of channels for each block
+            norm_num_groups=8, # number of groups for the normalization layer
+            in_channels=num_frames, # number of input channels
+            out_channels=1, # number of output channels
+            )
+    else:
+        raise(f"Model name not found: {model_name}, choose between 'unet_small' or 'unet_big'")
+
+class UNet2D(UNet2DModel, WandbModel):
+    def __init__(self, *x, **kwargs):
+        super().__init__(*x, **kwargs)
+        init_unet(self)
+
+    def forward(self, *x, **kwargs):
+        return super().forward(*x, **kwargs).sample ## Diffusers's UNet2DOutput class
+
+
+## Simple Diffusion paper
+
+def get_uvit_params(model_name="uvit_small", num_frames=4):
+    "Return the parameters for the diffusers UViT model"
+    if model_name == "uvit_small":
+        return dict(
+            dim=512,
+            ff_mult=2,
+            vit_depth=4,
+            channels=4, 
+            patch_size=4,
+            final_img_itransform=nn.Conv2d(num_frames,1,1)
+            )
+    elif model_name == "uvit_big":
+        return dict(
+            dim=1024,
+            ff_mult=4,
+            vit_depth=8,
+            channels=4, 
+            patch_size=4,
+            final_img_itransform=nn.Conv2d(num_frames,1,1)
+            )
+    else:
+        raise(f"Model name not found: {model_name}, choose between 'uvit_small' or 'uvit_big'")
+
+class UViTModel(UViT, WandbModel): pass

cloud_diffusion/simple_diffusion.py

@@ -1,19 +1,15 @@
 from functools import partial
 
-import torch, math
-from torch import nn, sqrt
+import torch
+from torch import sqrt
 from torch.special import expm1
-from torch.utils.data import DataLoader
-from torch.utils.data.dataloader import default_collate
 
 from fastprogress import progress_bar
 
 from einops import repeat
 
 try:
-    from denoising_diffusion_pytorch.simple_diffusion import (
-        UViT, right_pad_dims_to, logsnr_schedule_cosine
-    )
+    from denoising_diffusion_pytorch.simple_diffusion import right_pad_dims_to, logsnr_schedule_cosine
 except:
     raise ImportError("Please install denoising_diffusion_pytorch with `pip install denoising_diffusion_pytorch`")
 
@@ -26,51 +22,24 @@ def q_sample(x_start, times, noise):
 
     return x_noised, log_snr
 
-def noisify(frames, pred_objective="v"):
-    past_frames = frames[:,:-1]
-    last_frame  = frames[:,-1:]
-    device = frames.device
+def noisify_uvit(x0, pred_objective="v"):
+    device = x0.device
 
-    noise =  torch.randn_like(last_frame)
-    times = torch.zeros((last_frame.shape[0],), device = device).float().uniform_(0, 1)
-    x, log_snr = q_sample(last_frame, times, noise)
+    noise =  torch.randn_like(x0)
+    times = torch.zeros((x0.shape[0],), device = device).float().uniform_(0, 1)
+    x, log_snr = q_sample(x0, times, noise)
 
     if pred_objective == 'v':
         padded_log_snr = right_pad_dims_to(x, log_snr)
         alpha, sigma = padded_log_snr.sigmoid().sqrt(), (-padded_log_snr).sigmoid().sqrt()
-        target = alpha * noise - sigma * last_frame
+        target = alpha * noise - sigma * x0
 
     elif pred_objective == 'eps':
         target = noise
 
-    return torch.cat([past_frames, x], dim=1), log_snr, target
-
-def collate_simple_diffusion(b): 
-    "Collate function that noisifies the last frame"
-    return noisify(default_collate(b))
-
-def get_uvit_params(model_name="uvit_small", num_frames=4):
-    "Return the parameters for the diffusers UViT model"
-    if model_name == "uvit_small":
-        return dict(
-            dim=512,
-            ff_mult=2,
-            vit_depth=4,
-            channels=4, 
-            patch_size=4,
-            final_img_itransform=nn.Conv2d(num_frames,1,1)
-            )
-    elif model_name == "uvit_big":
-        return dict(
-            dim=1024,
-            ff_mult=4,
-            vit_depth=8,
-            channels=4, 
-            patch_size=4,
-            final_img_itransform=nn.Conv2d(num_frames,1,1)
-            )
-    else:
-        raise(f"Model name not found: {model_name}, choose between 'uvit_small' or 'uvit_big'")
+    return x, log_snr, target
+
+
 
 # Sampling functions
 
@@ -138,4 +107,4 @@ def p_sample_loop(model, past_frames, steps=500):
 def simple_diffusion_sampler(steps=500):
     """Returns a function that samples from the diffusion model using
     the simple diffusion sampling scheme"""
-    return partial(p_sample_loop, steps=500)
+    return partial(p_sample_loop, steps=steps)