Upload 4 files

data_utils.py

@@ -0,0 +1,99 @@
+import os
+
+import random
+import torch
+import imageio.v3 as imageio
+import numpy as np
+import skimage.morphology as morph
+import torchvision.transforms.v2.functional as T_F
+
+from skimage.filters import sato
+from pathlib import Path
+from scipy.ndimage import zoom
+from torchvision.datasets.folder import has_file_allowed_extension
+
+
+def make_dataset_t(image_dir, extensions=(".tif", ".tiff")):
+    image_dir = Path(image_dir)
+    images = [
+        (path, image_dir / f'Ridge_{path.name}', image_dir / f'Basins_{path.name}')
+        for path in sorted(image_dir.iterdir())
+        if (has_file_allowed_extension(path.name, extensions)
+            and (not path.name.startswith('Ridge_')) and (not path.name.startswith('Basins_')))
+    ]
+    return images
+
+def make_dataset_t_v(image_dir, extensions=(".tif", ".tiff")):
+    image_dir = Path(image_dir)
+    # Use list comprehension for faster filtering
+    images = [
+        (path, image_dir / f'Ridge_{path.name}', image_dir / f'Basins_{path.name}')
+        for path in sorted(image_dir.iterdir())
+        if (has_file_allowed_extension(path.name, extensions)
+            and (not path.name.startswith('Ridge_')) and (not path.name.startswith('Basins_')))
+    ]
+
+    # Shuffle in place
+    random.shuffle(images)
+
+    # Calculate split index once
+    split_idx = int(0.95 * len(images))
+    return images[:split_idx], images[split_idx:]
+
+def augmentations(image, label1, label2):
+    if random.random() < 0.5:
+        image, label1, label2 = T_F.vflip(image), T_F.vflip(label1), T_F.vflip(label2)
+    if random.random() < 0.5:
+        image, label1, label2 = T_F.hflip(image), T_F.hflip(label1), T_F.vflip(label2)
+    angles = [90, 180, 270]
+    angle = random.choice(angles)
+    if random.random() < 0.75:
+        image, label1, label2 = T_F.rotate(image, angle), T_F.rotate(label1, angle), T_F.rotate(label2, angle)
+    return image, label1, label2
+
+mean, std = (149.95293407563648, 330.8314960521203)
+target_water_level_range = [-100, 300]
+
+class TrainDataset(torch.utils.data.Dataset):
+    def __init__(self, train_split):
+        self.train_split = train_split
+
+    def __len__(self):
+        return len(self.train_split)
+
+    def __getitem__(self, index):
+        pair = self.train_split[index]
+        img = torch.from_numpy(imageio.imread(str(pair[0])))[None, :]
+        img = (img - mean) / std
+        ridge = torch.from_numpy(imageio.imread(str(pair[1])))[None, :].to(torch.float16)
+        basins = torch.from_numpy(imageio.imread(str(pair[2])))[None, :]
+        water_level = random.randint(*target_water_level_range)
+        basins = (basins >= water_level).to(torch.float16)
+        img, ridge, basins = augmentations(img, ridge, basins)
+        return img, ridge, basins, torch.tensor(water_level, dtype=torch.float16)
+
+class ValDataset(torch.utils.data.Dataset):
+    def __init__(self, val_split):
+        self.val_split = val_split
+
+    def __len__(self):
+        return len(self.val_split)
+
+    def __getitem__(self, index):
+        pair = self.val_split[index]
+        img = torch.from_numpy(imageio.imread(str(pair[0])))[None, :]
+        img = (img - mean) / std
+        ridge = torch.from_numpy(imageio.imread(str(pair[1])))[None, :].to(torch.float16)
+        basins = torch.from_numpy(imageio.imread(str(pair[2])))[None, :]
+        target_level = random.randint(*target_water_level_range)
+        basins = (basins >= target_level).to(torch.float16)
+        return img, ridge, basins, torch.tensor(target_level, dtype=torch.float16)
+
+if __name__ == '__main__':
+    train_split, val_split = make_dataset_t_v('dataset')
+
+    train_dataset = TrainDataset(train_split)
+    val_dataset = ValDataset(val_split)
+
+    print(train_dataset.__getitem__(0))
+    print(val_dataset.__getitem__(0))

mass_generate_examples.py

@@ -0,0 +1,122 @@
+import torch
+import math
+
+import torch.utils.data
+
+import imageio.v3 as imageio
+import lightning.pytorch as pl
+import matplotlib.pyplot as plt
+
+from network_diffusion_unet import ConditionalUNetDiT
+from safetensors.torch import load_file
+
+class PLModule(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.model = ConditionalUNetDiT(8, 16)
+
+    @torch.no_grad()
+    def inference_step(self, ridge_map, basin_map, water_level, num_steps=50):
+        device = self.device
+        b = ridge_map.shape[0]
+        x = torch.randn_like(ridge_map, device=device, dtype=torch.float16)
+        water_level = torch.tensor((water_level,), device=device, dtype=torch.float16).expand(b, )
+        time = torch.linspace(0, 1, num_steps + 1, device=device, dtype=torch.float16)
+
+        for i in range(num_steps):
+            t = torch.full((b,), time[i], device=device, dtype=torch.float16)
+            dt = torch.full((b, 1, 1, 1), time[i + 1] - time[i], device=device, dtype=torch.float16)
+
+            v = self.model(x, ridge_map, basin_map, water_level, t)
+
+            x = x + dt * v
+
+        return x
+
+if __name__ == "__main__":
+    #model = PLModule.load_from_checkpoint('FlashScape.ckpt').to(device='cuda', dtype=torch.float16)
+    model = PLModule()
+    model.model.load_state_dict(load_file('FlashScape.safetensors'))
+    model.to(device='cuda', dtype=torch.float16)
+    model.eval()
+
+    test_ridge = torch.from_numpy(imageio.imread('dataset_large/Ridge_11417648.tiff'))[None, None, :].to(dtype=torch.float16, device='cuda')
+    test_basin = torch.from_numpy(imageio.imread('dataset_large/Basins_11417648.tiff'))[None, None, :].to(dtype=torch.float16, device='cuda')
+    gt = torch.from_numpy(imageio.imread('dataset_large/11417648.tiff'))[None, None, :].to(dtype=torch.float16, device='cuda')
+    water_level = 300.0
+    num_steps = 10
+    num_images = 4
+
+    test_basin = (test_basin >= water_level).to(torch.float16)
+    test_ridge = test_ridge.expand(num_images, -1, -1, -1)
+    test_basin = test_basin.expand(num_images, -1, -1, -1)
+    generated = model.inference_step(test_ridge, test_basin, water_level, num_steps)
+    # Back to original range
+    generated = generated * 330.8314960521203 + 149.95293407563648
+
+    # Prepare images for visualization
+    ridge_display = test_ridge[0, 0].cpu().float()
+    basin_display = test_basin[0, 0].cpu().float()
+    gt_display = gt[0, 0].cpu().float()
+    generated_display = generated[:, 0].cpu()  # Remove channel dim
+
+    # Calculate optimal grid layout
+    total_images = num_images + 3  # condition1+ condition2 + gt + generated images
+    image_size = ridge_display.shape[0]  # assuming square images
+
+    # Determine optimal number of columns (aim for roughly 4:3 aspect ratio)
+    max_cols = min(6, total_images)  # Maximum 6 columns for readability
+    cols = min(max_cols, total_images)
+    rows = math.ceil(total_images / cols)
+
+    # Calculate figure size based on image dimensions and grid layout
+    base_height_per_image = 5  # inches per image height
+    base_width_per_image = 5  # inches per image width
+
+    fig_width = cols * base_width_per_image + 0.1  # +1 for colorbar space
+    fig_height = rows * base_height_per_image
+
+    # Create figure with subplots
+    fig, axes = plt.subplots(rows, cols, figsize=(fig_width, fig_height))
+
+    # Flatten axes array for easier indexing
+    if rows > 1 and cols > 1:
+        axes = axes.flatten()
+    elif rows == 1 and cols > 1:
+        axes = axes
+    elif rows > 1 and cols == 1:
+        axes = axes[:, 0]
+    else:
+        axes = [axes]
+
+    # Hide unused subplots
+    for i in range(total_images, len(axes)):
+        axes[i].set_visible(False)
+
+    # Plot condition image
+    im0 = axes[0].imshow(ridge_display, cmap='gray')
+    axes[0].set_title('Ridge Condition', fontsize=12, pad=2)
+    axes[0].set_axis_off()
+
+    # Plot condition image
+    im1 = axes[1].imshow(basin_display, cmap='gray')
+    axes[1].set_title('Basin Condition', fontsize=12, pad=2)
+    axes[1].set_axis_off()
+
+    # Plot ground truth image
+    im2 = axes[2].imshow(gt_display, cmap='gray')
+    axes[2].set_title('Ground Truth', fontsize=12, pad=2)
+    axes[2].set_axis_off()
+
+    # Plot generated images
+    for i in range(num_images):
+        im = axes[i + 3].imshow(generated_display[i], cmap='gray')
+        axes[i + 3].set_title(f'Generated {i + 1}', fontsize=10, pad=2)
+        axes[i + 3].set_axis_off()
+
+    # Add colorbar
+    cbar = fig.colorbar(im, ax=axes.ravel().tolist(), shrink=0.8, location='right')
+    cbar.set_label('Elevation', fontsize=14)
+
+    plt.savefig('result_grid.png', bbox_inches='tight', dpi=300)
+    plt.show()

network_diffusion_unet.py

@@ -0,0 +1,366 @@
+import math
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+
+
+class SinusoidalEmbedding(nn.Module):
+    def __init__(self, embedding_dim=128, scaling=1000):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        half_dim = embedding_dim // 2
+        freqs = torch.exp(-math.log(10000) * torch.arange(0, half_dim) / half_dim)
+        self.scaling = nn.parameter.Buffer(torch.tensor(scaling))
+        self.freqs = nn.parameter.Buffer(freqs)
+
+    def forward(self, scaler):
+        scaler = scaler * self.scaling
+        args = scaler[:, None] * self.freqs[None]
+        embedding = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
+        return embedding
+
+
+class SinusoidalPositionalEmbedding2D(nn.Module):
+
+    def __init__(self, embedding_dim):
+        super().__init__()
+        assert embedding_dim % 2 == 0, "embedding_dim must be even"
+        self.embedding_dim = embedding_dim
+        half_dim = self.embedding_dim // 2
+        div_term = torch.exp(torch.arange(0, half_dim, 2, dtype=torch.float32) * (-math.log(10000.0) / half_dim))
+        self.div_term = nn.parameter.Buffer(div_term)
+
+    def forward(self, height, width):
+        """Generate embeddings for a grid of size (height, width)."""
+
+        # Generate grid coordinates
+        y_pos = torch.arange(height, dtype=torch.float32, device=self.div_term.device)
+        x_pos = torch.arange(width, dtype=torch.float32, device=self.div_term.device)
+
+        # Compute sinusoidal components for height and width
+        y_sin = torch.sin(y_pos[:, None] * self.div_term[None, :])
+        y_cos = torch.cos(y_pos[:, None] * self.div_term[None, :])
+        x_sin = torch.sin(x_pos[:, None] * self.div_term[None, :])
+        x_cos = torch.cos(x_pos[:, None] * self.div_term[None, :])
+
+        # Interleave sin and cos components
+        y_embed = torch.stack([y_sin, y_cos], dim=-1).view(height, -1)
+        x_embed = torch.stack([x_sin, x_cos], dim=-1).view(width, -1)
+
+        # Combine height and width embeddings
+        pos_embed = torch.cat([y_embed[:, None, :].expand(-1, width, -1),
+                               x_embed[None, :, :].expand(height, -1, -1)], dim=-1)
+        return pos_embed.view(height * width, self.embedding_dim)
+
+
+class ImageLinearAttention(nn.Module):
+    def __init__(self, chan, kernel_size=3, heads=4, norm_queries=True, embd_dim=None):
+        super().__init__()
+        self.chan = chan
+        self.heads = heads
+        self.key_dim = key_dim = chan // heads
+        self.value_dim = value_dim = chan // heads
+        self.norm_queries = norm_queries
+
+        # Convolutional projections for Q, K, V
+        self.to_q = nn.Conv2d(chan, key_dim * heads, kernel_size, padding='same', padding_mode='replicate')
+        self.to_k = nn.Conv2d(chan, key_dim * heads, kernel_size, padding='same', padding_mode='replicate')
+        self.to_v = nn.Conv2d(chan, value_dim * heads, kernel_size, padding='same', padding_mode='replicate')
+        self.to_out = nn.Conv2d(value_dim * heads, chan, kernel_size, padding='same', padding_mode='replicate')
+
+        # Adaptive normalization: Project embedding to scale/shift for group norm
+        if embd_dim is not None:
+            self.norm = nn.GroupNorm(1, key_dim * heads, affine=False)  # Normalize without inherent affine params
+            self.emb_proj = nn.Linear(embd_dim, 2 * key_dim * heads)  # Project emb to scale/shift
+        else:
+            self.norm = nn.GroupNorm(1, key_dim * heads, affine=True)
+            self.emb_proj = None
+
+    def forward(self, x, emb=None):
+        b, c, h, w = x.shape
+        heads = self.heads
+        key_dim = self.key_dim
+
+        # Project input to queries, keys, and values
+        q = self.to_q(x)
+        k = self.to_k(x)
+        v = self.to_v(x)
+
+        # Apply adaptive normalization if embedding is provided
+        if emb is not None and self.emb_proj is not None:
+            emb_params = self.emb_proj(emb).view(b, 2, -1)  # (b, 2, key_dim * heads)
+            scale, shift = emb_params[:, 0], emb_params[:, 1]  # Split into scale and shift
+            # Normalize and modulate Q, K, V
+            q = self.norm(q)
+            k = self.norm(k)
+            v = self.norm(v)
+            # Apply scale and shift across spatial dimensions
+            q = q * (1 + scale[:, :, None, None]) + shift[:, :, None, None]
+            k = k * (1 + scale[:, :, None, None]) + shift[:, :, None, None]
+            v = v * (1 + scale[:, :, None, None]) + shift[:, :, None, None]
+
+        # Reshape Q, K, V for multi-head attention
+        q = q.view(b, heads, key_dim, h * w)
+        k = k.view(b, heads, key_dim, h * w)
+        v = v.view(b, heads, self.value_dim, h * w)
+
+        # Scale queries and keys
+        q = q * (key_dim ** -0.25)
+        k = k * (key_dim ** -0.25)
+
+        # Softmax on keys along the sequence dimension
+        k = k.softmax(dim=-1)
+        if self.norm_queries:
+            q = q.softmax(dim=-2)
+
+        # Compute context and output
+        context = torch.einsum('bhdn,bhen->bhde', k, v)
+        out = torch.einsum('bhdn,bhde->bhen', q, context)
+        out = out.reshape(b, -1, h, w)
+        out = self.to_out(out)
+        return x + out
+
+
+class ResConvBlock(nn.Module):
+    def __init__(self, channels, time_dim):
+        super().__init__()
+        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1, bias=False, padding_mode='replicate')
+        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1, padding_mode='replicate')
+        self.gn1 = nn.GroupNorm(8, channels, affine=True)
+        self.gn2 = nn.GroupNorm(8, channels, affine=False)
+        self.time_affine = nn.Linear(time_dim, channels * 2)
+        self.act = nn.LeakyReLU(inplace=True)
+
+    def forward(self, x, t_emb):
+        # Get affine parameters from time embedding
+        affine_params = self.time_affine(t_emb)
+        scale, shift = affine_params.chunk(2, dim=1)
+
+        # First convolution path
+        h = self.conv1(self.act(self.gn1(x)))
+
+        # Second convolution path with adaptive normalization
+        h = self.gn2(h)
+        h = h * (1 + scale[:, :, None, None]) + shift[:, :, None, None]
+        h = self.conv2(self.act(h))
+
+        return x + h
+
+
+class DiTLayer(nn.Module):
+    def __init__(self, d_model, nhead, dim_feedforward=1024):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(d_model)
+        self.attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, dim_feedforward),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(dim_feedforward, d_model),
+        )
+
+    def forward(self, src):
+        # Self-attention block
+        attn_output, _ = self.attn(self.norm1(src), self.norm1(src), self.norm1(src))
+        src = src + attn_output
+
+        # Feedforward block
+        ffn_output = self.ffn(self.norm2(src))
+        src = src + ffn_output
+        return src
+
+
+class DiTBlock(nn.Module):
+    def __init__(self, channels, patch_size, hidden_size, nhead, num_layers=2):
+        super().__init__()
+        self.patch_size = patch_size
+        self.patchify = nn.Unfold(kernel_size=patch_size, stride=patch_size)
+        self.patch_embedding_in = nn.Linear(channels * patch_size**2, hidden_size)
+        self.pos_embd = SinusoidalPositionalEmbedding2D(hidden_size)
+        self.waterlevel_embd = SinusoidalEmbedding(hidden_size, 10)
+        self.patch_embedding_out = nn.Linear(hidden_size, channels * patch_size**2)
+        self.dit_layers = nn.ModuleList([
+            DiTLayer(hidden_size, nhead, 2*hidden_size)
+            for _ in range(num_layers)
+        ])
+        self.norm = nn.GroupNorm(8, channels)
+
+    def forward(self, src, water_level):
+        B, C, H, W = src.shape
+        H_p, W_p = H // self.patch_size, W // self.patch_size
+        x = self.norm(src)
+        x = self.patchify(x).permute(0, 2, 1)
+        x = self.patch_embedding_in(x)
+        pos_embd = self.pos_embd(H_p, W_p).to(dtype=x.dtype)
+        x = x + pos_embd.unsqueeze(0)
+        water_level_cls = self.waterlevel_embd(water_level).unsqueeze(1)
+        x = torch.cat((x, water_level_cls), dim=1)
+        for dit_layer in self.dit_layers:
+            x = dit_layer(x)
+        x = self.patch_embedding_out(x).permute(0, 2, 1)
+        x = x[:, :, :-1]
+        x = nn.functional.fold(x, (H, W), (self.patch_size, self.patch_size), stride=(self.patch_size, self.patch_size))
+        return src + x
+
+
+class UpBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, time_dim, cat):
+        super().__init__()
+        self.res = ResConvBlock(in_ch, time_dim)
+        self.up = nn.ConvTranspose2d(in_ch, out_ch, 4, stride=2, padding=1)
+        self.cat = cat
+
+    def forward(self, x, t_emb, skip=None):
+        x = self.res(x, t_emb)
+        x = self.up(x)
+        if self.cat:
+            x = torch.cat([x, skip], dim=1)
+        else:
+            x = x + skip
+        return x
+
+class UpBlockWithDit(nn.Module):
+    def __init__(self, in_ch, out_ch, patch_size, hidden_size, nhead, time_dim, cat):
+        super().__init__()
+        self.res = ResConvBlock(in_ch, time_dim)
+        self.dit = DiTBlock(in_ch, patch_size, hidden_size, nhead, 4)
+        self.up = nn.ConvTranspose2d(in_ch, out_ch, 4, stride=2, padding=1)
+        self.cat = cat
+
+    def forward(self, x, t_emb, water_level, skip=None):
+        x = self.res(x, t_emb)
+        x = self.dit(x, water_level)
+        x = self.up(x)
+        if self.cat:
+            x = torch.cat([x, skip], dim=1)
+        else:
+            x = x + skip
+        return x
+
+
+def run_block(module, *args):
+    return module(*args)
+
+
+class ConditionalUNet(nn.Module):
+    def __init__(self, base_ch=16, embd_dim=64, depth=5):
+        super().__init__()
+        self.depth = depth
+        self.time_embd = SinusoidalEmbedding(embd_dim)
+        self.waterlevel_embd = SinusoidalEmbedding(embd_dim, 10)
+        embd_dim *= 2
+
+        # Input channels = noisy height (1) + ridge map (1) + lake map (1)
+        self.expand = nn.Conv2d(4, base_ch, 3, padding=1, padding_mode='replicate')
+
+        # Encoder layers
+        self.enc_blocks = nn.ModuleList()
+        self.enc_dit_blocks = nn.ModuleList()
+        self.down_convs = nn.ModuleList()
+        current_ch = base_ch
+
+        for i in range(depth):
+            self.enc_blocks.append(ResConvBlock(current_ch, embd_dim))
+            if i < depth - 1:
+                self.down_convs.append(
+                    nn.Conv2d(current_ch, current_ch * 2, 4, stride=2, padding=1, padding_mode='replicate')
+                )
+                current_ch *= 2
+
+        # Bottleneck
+        self.bottleneck = nn.Conv2d(current_ch, current_ch * 2, 4, stride=2, padding=1, padding_mode='replicate')
+        current_ch *= 2
+
+        # Decoder layers
+        self.up_blocks = nn.ModuleList()
+        for i in range(depth):
+            cat = (i == depth - 1)  # Only concatenate in the final up block
+            self.up_blocks.append(UpBlock(current_ch, current_ch // 2, embd_dim, cat))
+            current_ch = current_ch // 2 * (2 if cat else 1)
+
+        self.out = ResConvBlock(current_ch, embd_dim)
+        self.final = nn.Conv2d(current_ch, 1, 1)
+
+
+    def forward(self, x, map_average, ridge_map, basin_map, water_level, t):
+        t_embed = self.time_embd(t).to(x.dtype)
+        waterlevel_embd = self.waterlevel_embd(water_level).to(x.dtype)
+        embeds = torch.cat([t_embed, waterlevel_embd], dim=1)
+
+        h = torch.cat([x, ridge_map, basin_map, map_average], dim=1)
+        h = checkpoint(run_block, self.expand, h, use_reentrant=False) if self.training else self.expand(h)
+
+        # Encoder
+        skips = []
+        for i in range(self.depth):
+            h = checkpoint(run_block, self.enc_blocks[i], h, embeds, use_reentrant=False) if self.training else self.enc_blocks[i](h, embeds)
+            skips.append(h)
+            if i < self.depth - 1:
+                h = checkpoint(run_block, self.down_convs[i], h, use_reentrant=False) if self.training else self.down_convs[i](h)
+
+        # Bottleneck
+        h = checkpoint(run_block, self.bottleneck, h, use_reentrant=False) if self.training else self.bottleneck(h)
+
+        # Decoder
+        for i in range(self.depth):
+            h = checkpoint(run_block, self.up_blocks[i], h, embeds, skips[-(i + 1)], use_reentrant=False) if self.training else self.up_blocks[i](h, embeds, skips[-(i + 1)])
+
+        h = checkpoint(run_block, self.out, h, embeds, use_reentrant=False) if self.training else self.out(h, embeds)
+        h = checkpoint(run_block, self.final, h, use_reentrant=False) if self.training else self.final(h)
+        return h
+
+
+class ConditionalUNetDiT(nn.Module):
+    def __init__(self, base_ch=8, embd_dim=16):
+        super().__init__()
+        self.time_embd = SinusoidalEmbedding(embd_dim, 1000)
+
+        # Input channels = noisy height (1) + ridge map (1) + lake map (1)
+        self.expand = nn.Conv2d(3, base_ch, 3, padding=1, padding_mode='replicate')
+        self.enc_0 = ResConvBlock(base_ch, embd_dim)
+
+        self.down0 = nn.Conv2d(base_ch, base_ch * 2, 4, stride=2, padding=1, padding_mode='replicate') # 1024->512
+        self.enc_1 = ResConvBlock(base_ch * 2, embd_dim)
+        self.enc_1_dit = DiTBlock(base_ch * 2, 16, 1024, 8, 4)
+
+        self.down1 = nn.Conv2d(base_ch * 2, base_ch * 4, 4, stride=2, padding=1, padding_mode='replicate') # 512->256
+
+        self.up1 = UpBlockWithDit(base_ch * 4, base_ch * 2, 8, 1024, 8, embd_dim, False) # 256->512
+        self.up0 = UpBlockWithDit(base_ch * 2, base_ch, 16, 1024, 8, embd_dim, True) # 512->1024
+        self.out = ResConvBlock(base_ch * 2, embd_dim)
+        self.final = nn.Conv2d(base_ch * 2, 1, 1)
+
+    def forward(self, x, ridge_map, basin_map, water_level, t):
+        t_embed = self.time_embd(t).to(x.dtype)
+        # x: noisy height map, ridge_map: binary edges, basin_map: binary basins, water_level: the estimate sea level
+        h0 = torch.cat([x, ridge_map, basin_map], dim=1)  # concat condition
+        # encode
+        h0 = checkpoint(run_block, self.expand, h0, use_reentrant=False) if self.training else self.expand(h0)
+        h0 = checkpoint(run_block, self.enc_0, h0, t_embed, use_reentrant=False) if self.training else self.enc_0(h0, t_embed)
+        h1 = checkpoint(run_block, self.down0, h0, use_reentrant=False) if self.training else self.down0(h0)
+        h1 = checkpoint(run_block, self.enc_1, h1, t_embed, use_reentrant=False) if self.training else self.enc_1(h1, t_embed)
+        h1 = checkpoint(run_block, self.enc_1_dit, h1, water_level, use_reentrant=False) if self.training else self.enc_1_dit(h1, water_level) # 512x512
+        h2 = checkpoint(run_block, self.down1, h1, use_reentrant=False) if self.training else self.down1(h1)  # 256x256
+        # decode with skip connections
+        out = checkpoint(run_block, self.up1, h2, t_embed, water_level, h1, use_reentrant=False) if self.training else self.up1(h2, t_embed, water_level, h1)  # 512x512
+        out = checkpoint(run_block, self.up0, out, t_embed, water_level, h0, use_reentrant=False) if self.training else self.up0(out, t_embed, water_level, h0)  # 1024x1024
+        out = checkpoint(run_block, self.out, out, t_embed, use_reentrant=False) if self.training else self.out(out, t_embed)
+        out = self.final(out)
+        return out  # predicted noise for diffusion loss
+
+
+
+if __name__ == "__main__":
+    #a = ConditionalUNet()
+    #t = SinusoidalEmbedding(256)
+    #t_embd = t(torch.randint(0, 100, (1,)))
+    #x = torch.randn(1, 1, 256, 256)
+    #r = torch.randn(1, 1, 256, 256)
+    #c = a(x, r, t_embd)
+    #print(c)
+    #print(c.shape)
+    network = ConditionalUNetDiT()
+    for name, m in network.named_modules():
+        if isinstance(m, nn.Linear) and 'time_affine':
+            m.weight.data.zero_()
+            m.bias.data.zero_()

pl_module_rectifiedflow.py

@@ -0,0 +1,181 @@
+import torch
+import math
+
+import data_utils
+import torch.utils.data
+
+import imageio.v3 as imageio
+import lightning.pytorch as pl
+import torch.nn as nn
+import torch.distributions as dist
+import numpy as np
+import safetensors.torch as st
+
+from network_diffusion_unet import ConditionalUNet, ConditionalUNetDiT
+from loss_fn import L1andGDL
+from adam_atan2_pytorch import AdamAtan2
+from lightning.pytorch.loggers.tensorboard import TensorBoardLogger
+from lightning.pytorch.utilities import grad_norm
+from lightning.pytorch.callbacks import LearningRateMonitor, StochasticWeightAveraging, LearningRateFinder
+from torchvision.utils import make_grid
+
+def convert_uniform_to_custom(u):
+    #return 0.5 - torch.cos((1/3) * torch.acos(1 - 2 * u) + math.pi / 3)
+    return 0.5 + 2 * torch.cos((2 * math.pi - torch.arccos((11/16)*(1-2*u)))/3)
+
+class PLModule(pl.LightningModule):
+    def __init__(self, mid_visual_ridge, mid_visual_basins, mid_visual_gt):
+        super().__init__()
+        self.save_hyperparameters()
+        self.lr = 6e-4
+        self.wd = 5e-5
+        self.model = ConditionalUNetDiT(base_ch=8, embd_dim=16)
+        #self.map_average = torch.from_numpy(imageio.imread(map_average)).unsqueeze(0)
+        #self.map_average = (self.map_average - self.map_average.mean()) / self.map_average.std()
+        self.loss_fn = L1andGDL()
+        self.val_metrics = []
+        self.mid_visual_ridge, self.mid_visual_basins = mid_visual_ridge, mid_visual_basins
+        self.mid_visual_gt = mid_visual_gt
+        self.initialize_model()
+
+    def initialize_model(self):
+        for name, m in self.model.named_modules():
+            if isinstance(m, nn.Linear) and ('time_affine' in name or 'water_level_affine' in name):
+                m.weight.data.zero_()
+                m.bias.data.zero_()
+
+    def configure_optimizers(self):
+        opt = AdamAtan2(self.parameters(), lr=self.lr, decoupled_wd=True, weight_decay=self.wd)
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, 100, eta_min=1e-7)
+        return {
+            "optimizer": opt,
+            "lr_scheduler": {"scheduler": scheduler, "interval": "epoch", "frequency": 1},
+        }
+
+    def _step(self, batch, batch_idx):
+        x0, ridge_map, basin_map, water_level = batch
+        b = water_level.shape[0]
+        #map_average = self.map_average.expand((b, -1, -1, -1)).to(self.device)
+
+        noise = torch.randn_like(x0, device=self.device, dtype=x0.dtype)
+        t = torch.rand((b,), device=self.device)
+        t = convert_uniform_to_custom(t).to(x0.dtype)
+
+        xt = t.view(-1, 1, 1, 1) * x0 + (1 - t.view(-1, 1, 1, 1)) * noise
+        v = x0 - noise
+
+        predicted_v = self.model(xt, ridge_map, basin_map, water_level, t)  # Predict velocity v
+        loss = self.loss_fn(predicted_v, v)  # Loss between predicted and target v
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        loss = self._step(batch, batch_idx)
+        self.logger.experiment.add_scalar(f"Train/Loss", loss.detach(), self.global_step)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        loss = self._step(batch, batch_idx)
+        self.val_metrics.append(loss.detach())
+        return loss
+
+    @torch.no_grad()
+    def inference_step(self, ridge_map, basin_map, water_level, num_steps=50):
+        device = self.device
+        b = ridge_map.shape[0]
+        x = torch.randn_like(ridge_map, device=device)
+        water_level = torch.tensor((water_level,), device=device).expand(b,)
+        time = torch.linspace(0, 1, num_steps + 1, device=device)
+
+        for i in range(num_steps):
+            t = torch.full((b,), time[i], device=device)
+            dt = torch.full((b, 1, 1, 1), time[i+1] - time[i], device=device)
+
+            v = self.model(x, ridge_map, basin_map, water_level, t)
+
+            x = x + dt * v
+
+        return x
+
+    def on_train_epoch_end(self):
+        sea_level = 0.0
+        ridge_map = torch.from_numpy(imageio.imread(self.mid_visual_ridge))[None,None,:].to(device=self.device, dtype=torch.float32)
+
+        basin_map = torch.from_numpy(imageio.imread(self.mid_visual_basins))[None,None,:].to(device=self.device)
+        basin_map = (basin_map>=sea_level).to(torch.float32)
+        output = self.inference_step(ridge_map, basin_map, sea_level)
+        mid_visual_result = output.squeeze([1])
+        self.logger.experiment.add_scalar("Visualize/Min", mid_visual_result.min(), self.current_epoch)
+        self.logger.experiment.add_scalar("Visualize/Max", mid_visual_result.max(), self.current_epoch)
+        self.logger.experiment.add_scalar("Visualize/Mean", mid_visual_result.mean(), self.current_epoch)
+        mid_visual_result = (mid_visual_result - mid_visual_result.min()) / (mid_visual_result.max() - mid_visual_result.min())
+        self.logger.experiment.add_image(f'Visualize/Model Output', mid_visual_result, self.current_epoch)
+
+        vram_data = torch.cuda.mem_get_info()
+        vram_usage = (vram_data[1] - vram_data[0]) / (1024 ** 2)
+        self.logger.experiment.add_scalar(f"Other/VRAM Usage", vram_usage, self.current_epoch)
+        torch.cuda.reset_peak_memory_stats()
+        if self.current_epoch == 0:
+            mid_visual_gt = torch.from_numpy(imageio.imread(self.mid_visual_gt))[None,:]
+            mid_visual_gt = (mid_visual_gt - mid_visual_gt.min()) / (mid_visual_gt.max() - mid_visual_gt.min())
+            self.logger.experiment.add_image(f'Visualize/Ridge', ridge_map.squeeze([1]), self.current_epoch)
+            self.logger.experiment.add_image(f'Visualize/Basin', basin_map.squeeze([1]), self.current_epoch)
+            self.logger.experiment.add_image(f'Visualize/GT', mid_visual_gt, self.current_epoch)
+
+    def on_validation_epoch_end(self):
+        epoch_averages = torch.stack(self.val_metrics).nanmean(dim=0)
+        self.logger.experiment.add_scalar("Val/Loss", epoch_averages, self.current_epoch)
+        self.val_metrics.clear()
+    #def on_before_optimizer_step(self, optimizer):
+    #    norms = grad_norm(self.model, norm_type=2)
+    #    self.log_dict(norms, logger=True)
+
+
+
+# Example usage
+if __name__ == "__main__":
+    torch.set_float32_matmul_precision('medium')
+    if torch.cuda.is_available() and torch.version.cuda:
+        print('Optimising computing and memory use via cuDNN! (NVIDIA GPU only).')
+        torch.backends.cudnn.enabled = True
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cudnn.allow_tf32 = True
+    elif torch.cuda.is_available() and torch.version.hip:
+        print('Optimising computing using TunableOp! (AMD GPU only).')
+        torch.cuda.tunable.enable()
+        torch.cuda.tunable.set_filename('TunableOp_results')
+
+    train_split, val_split = data_utils.make_dataset_t_v('dataset_large')
+
+    callbacks = []
+    callbacks.append(LearningRateMonitor(logging_interval='epoch'))
+    model_checkpoint = pl.callbacks.ModelCheckpoint(dirpath="", filename="FlashScape",
+                                                    save_weights_only=False,
+                                                    enable_version_counter=False, save_last=False)
+    callbacks.append(model_checkpoint)
+    swa_callback = StochasticWeightAveraging(1e-5, 0.8, int(0.2 * 100 - 1))
+    callbacks.append(swa_callback)
+    #lr_finder = LearningRateFinder(1e-5, 0.1)
+    #callbacks.append(lr_finder)
+    #model = PLModule.load_from_checkpoint('FlashScape V2.ckpt')
+    trainer = pl.Trainer(max_epochs=100, log_every_n_steps=1, logger=TensorBoardLogger(f'lightning_logs', name='FlashScape Dit No MapAvg Zero Init'),
+                         accelerator="gpu", enable_checkpointing=True,
+                         precision='16-mixed', enable_progress_bar=True, num_sanity_val_steps=0, callbacks=callbacks)
+    with trainer.init_module():
+        model = PLModule('dataset_large/Ridge_11417648.tiff',
+                         'dataset_large/Basins_11417648.tiff',
+                         'dataset_large/11417648.tiff')
+    model = torch.compile(model)
+
+
+    train_dataset = data_utils.TrainDataset(train_split)
+    val_dataset = data_utils.ValDataset(val_split)
+
+    train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=8,
+                                               num_workers=8, pin_memory=False, persistent_workers=True, shuffle=True)
+    val_loader = torch.utils.data.DataLoader(dataset=val_dataset, batch_size=8,
+                                             num_workers=8, pin_memory=False, persistent_workers=True)
+
+    trainer.fit(model,
+                val_dataloaders=val_loader,
+                train_dataloaders=train_loader)
+                #ckpt_path='FlashScape.ckpt')