train.py

import os, gc, random, warnings
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from pathlib import Path
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms.functional as TF
from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler, DDIMScheduler
from diffusers.training_utils import EMAModel
from diffusers.optimization import get_cosine_schedule_with_warmup
from peft import LoraConfig, get_peft_model
from accelerate import Accelerator
from accelerate.utils import set_seed
from tqdm.auto import tqdm
import wandb
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
warnings.filterwarnings('ignore')

DATA_ROOT        = '/kaggle/input/datasets/shambac/augmented-sentinel-1-2'
OUTPUT_DIR       = '/kaggle/working/checkpoints'
SD_MODEL_ID      = 'runwayml/stable-diffusion-v1-5'
IMG_SIZE         = 256
SEASONS          = ('spring', 'summer', 'fall', 'winter')
MAX_PAIRS        = 90000
VAL_FRACTION     = 0.10
LORA_RANK        = 32
LORA_ALPHA       = 32
USE_DORA         = True
NUM_TRAIN_STEPS  = 12_000
RESUME_STEP      = 0
RESUME_FROM      = None
BATCH_SIZE       = 8
GRAD_ACCUM       = 2
LR               = 5e-5
LR_WARMUP        = 1_000
MIXED_PRECISION  = 'fp16'
GRAD_CKPT        = True
EMA_DECAY        = 0.9999
NUM_TIMESTEPS    = 1000
COLOR_LOSS_W     = 0.5
PERCEPTUAL_W     = 0.1
COLOR_LOSS_FREQ  = 35
PERCEPTUAL_FREQ  = 50
LOG_EVERY        = 100
SAVE_EVERY       = 2_000
VIS_EVERY        = 3_000
SEED             = 332

def collect_pairs(root, seasons, max_pairs=None):
    root_str = Path(root).as_posix()
    season_buckets = {s: [] for s in seasons}
    for season in seasons:
        csv_files = list((Path(root) / season).glob('*.csv'))
        if not csv_files:
            continue
        df = pd.concat([pd.read_csv(f) for f in csv_files], ignore_index=True)
        df['season'] = season
        df['region'] = df['region'].str.strip().str.lower()
        df['s1_fileName'] = (df['s1_fileName']
                             .str.replace('\\', '/', regex=False)
                             .str.replace(r'(\w+?)1s2_', r'\1_s1_', regex=True))
        df['s2_fileName'] = df['s2_fileName'].str.replace('\\', '/', regex=False)
        df['s1'] = root_str + '/' + df['s1_fileName']
        df['s2'] = root_str + '/' + df['s2_fileName']
        season_buckets[season] = df[['s1','s2','season','region']].to_dict('records')
    active = [s for s in seasons if season_buckets[s]]
    if max_pairs is None:
        pairs = [p for s in active for p in season_buckets[s]]
    else:
        per_season = max_pairs // len(active)
        pairs = []
        for s in active:
            bucket = season_buckets[s].copy()
            random.shuffle(bucket)
            pairs.extend(bucket[:per_season])
    random.shuffle(pairs)
    return pairs

class SAROpticalDataset(Dataset):
    def __init__(self, pairs, img_size=256, augment=True):
        self.pairs    = pairs
        self.img_size = img_size
        self.augment  = augment
    def __len__(self):
        return len(self.pairs)
    def _load_sar(self, path):
        img = Image.open(path).convert('L')
        if img.size != (self.img_size, self.img_size):
            img = img.resize((self.img_size, self.img_size), Image.BILINEAR)
        arr = np.array(img, dtype=np.float32) / 255.0
        arr = np.stack([arr, arr, arr], axis=2)
        return torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0
    def _load_optical(self, path):
        img = Image.open(path).convert('RGB')
        if img.size != (self.img_size, self.img_size):
            img = img.resize((self.img_size, self.img_size), Image.BILINEAR)
        arr = np.array(img, dtype=np.float32) / 255.0
        return torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0
    def __getitem__(self, idx):
        pair = self.pairs[idx]
        sar  = self._load_sar(pair['s1'])
        opt  = self._load_optical(pair['s2'])
        if self.augment:
            if random.random() > 0.5:
                sar = TF.hflip(sar); opt = TF.hflip(opt)
            if random.random() > 0.5:
                sar = TF.vflip(sar); opt = TF.vflip(opt)
            k = random.randint(0, 3)
            if k > 0:
                sar = torch.rot90(sar, k, [1, 2])
                opt = torch.rot90(opt, k, [1, 2])
        return {'sar': sar, 'optical': opt, 'season': pair['season'], 'region': pair['region']}

class VGGPerceptualLoss(nn.Module):
    def __init__(self):
        super().__init__()
        import torchvision.models as models
        vgg = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1)
        self.features = nn.Sequential(*list(vgg.features)[:9]).eval()
        for p in self.parameters():
            p.requires_grad_(False)
        self.register_buffer('mean', torch.tensor([0.485, 0.456, 0.406]).view(1,3,1,1))
        self.register_buffer('std',  torch.tensor([0.229, 0.224, 0.225]).view(1,3,1,1))
    def forward(self, pred, target):
        pred   = (pred   * 0.5 + 0.5 - self.mean) / self.std
        target = (target * 0.5 + 0.5 - self.mean) / self.std
        return F.l1_loss(self.features(pred), self.features(target))

def color_supervision_loss(noise_pred, noise, noisy_latents, timesteps, scheduler, vae, opt_latents):
    low_noise_mask = timesteps < 500
    if low_noise_mask.sum() == 0:
        return torch.tensor(0.0, device=noise_pred.device)
    vae_mod        = vae.module if hasattr(vae, 'module') else vae
    alphas_cumprod = scheduler.alphas_cumprod.to(noise_pred.device)
    sqrt_alpha     = alphas_cumprod[timesteps[low_noise_mask]].sqrt().view(-1,1,1,1)
    sqrt_one_minus = (1 - alphas_cumprod[timesteps[low_noise_mask]]).sqrt().view(-1,1,1,1)
    noisy_sub = noisy_latents[low_noise_mask, :4]
    noise_sub = noise_pred[low_noise_mask]
    x0_pred   = (noisy_sub - sqrt_one_minus * noise_sub) / (sqrt_alpha + 1e-8)
    x0_pred   = x0_pred.clamp(-4, 4)
    with torch.no_grad():
        pred_img = vae_mod.decode((x0_pred / vae_mod.config.scaling_factor).to(vae_mod.dtype)).sample
        gt_img   = vae_mod.decode((opt_latents[low_noise_mask] / vae_mod.config.scaling_factor).to(vae_mod.dtype)).sample
    return F.l1_loss(pred_img.float(), gt_img.float())

@torch.no_grad()
def translate_single(sar_tensor, unet_m, vae_m, null_embed, device, num_steps=50):
    ddim = DDIMScheduler(num_train_timesteps=NUM_TIMESTEPS, beta_schedule='scaled_linear', prediction_type='epsilon')
    sar      = sar_tensor.unsqueeze(0).to(device, dtype=torch.float16)
    vae_mod  = vae_m.module if hasattr(vae_m, 'module') else vae_m
    sar_latent = vae_mod.encode(sar).latent_dist.mean * vae_mod.config.scaling_factor
    latents    = torch.randn_like(sar_latent) * ddim.init_noise_sigma
    embed      = null_embed.to(device, dtype=torch.float16).expand(1, -1, -1)
    ddim.set_timesteps(num_steps)
    for t in ddim.timesteps:
        model_in   = torch.cat([latents, sar_latent], dim=1)
        noise_pred = unet_m(model_in.float(), t.unsqueeze(0).to(device), encoder_hidden_states=embed.float()).sample
        latents    = ddim.step(noise_pred.to(torch.float16), t, latents).prev_sample
    image = vae_mod.decode(latents / vae_mod.config.scaling_factor).sample
    return image.squeeze(0).clamp(-1, 1).float().cpu()

def denorm(t):
    return ((t.clamp(-1, 1) + 1) / 2).permute(1, 2, 0).numpy()

def save_validation_grid(val_pairs, unet_m, vae_m, null_embed, device, step, out_dir):
    selected = []
    for s in SEASONS:
        pool = [p for p in val_pairs if p['season'] == s]
        if pool: selected.append(random.choice(pool))
    selected = selected[:8]
    fig, axes = plt.subplots(len(selected), 3, figsize=(12, 3.5 * len(selected)))
    for i, pair in enumerate(selected):
        ds   = SAROpticalDataset([pair], img_size=IMG_SIZE, augment=False)
        item = ds[0]
        pred = translate_single(item['sar'], unet_m, vae_m, null_embed, device)
        axes[i,0].imshow(denorm(item['sar'])[:,:,0], cmap='gray')
        axes[i,0].set_title(f"SAR | {pair['season']}", fontsize=8)
        axes[i,0].axis('off')
        axes[i,1].imshow(denorm(pred))
        axes[i,1].set_title(f'Predicted (step {step})', fontsize=8)
        axes[i,1].axis('off')
        axes[i,2].imshow(denorm(item['optical']))
        axes[i,2].set_title('Ground Truth', fontsize=8)
        axes[i,2].axis('off')
    plt.tight_layout()
    path = os.path.join(out_dir, f'val_step_{step:06d}.png')
    plt.savefig(path, dpi=100, bbox_inches='tight')
    plt.close()
    return path

def main():
    accelerator = Accelerator(
        mixed_precision=MIXED_PRECISION,
        gradient_accumulation_steps=GRAD_ACCUM,
        log_with='wandb',
        project_dir=OUTPUT_DIR,
    )
    set_seed(SEED + accelerator.process_index)
    is_main = accelerator.is_main_process
    if is_main:
        os.makedirs(OUTPUT_DIR, exist_ok=True)
        os.makedirs(os.path.join(OUTPUT_DIR, 'val_grids'), exist_ok=True)
        wandb.login(key=os.environ['WANDB_API_KEY'], relogin=True)
        accelerator.init_trackers(
            project_name='sar-optical-diffusion',
            config={'lr': LR, 'batch_size': BATCH_SIZE, 'lora_rank': LORA_RANK,
                    'steps': NUM_TRAIN_STEPS, 'resume_step': RESUME_STEP,
                    'color_loss_w': COLOR_LOSS_W, 'perceptual_w': PERCEPTUAL_W,
                    'effective_batch': BATCH_SIZE * GRAD_ACCUM * accelerator.num_processes},
            init_kwargs={'wandb': {
                'name': f'dora-r{LORA_RANK}-step{RESUME_STEP}-to-{RESUME_STEP+NUM_TRAIN_STEPS}',
                'tags': ['sentinel', 'SAR', 'diffusion', 'DoRA', 'color-supervision'],
                'resume': 'allow',
            }}
        )
    all_pairs = collect_pairs(DATA_ROOT, SEASONS, MAX_PAIRS)
    train_pairs, val_pairs = train_test_split(
        all_pairs, test_size=VAL_FRACTION,
        stratify=[f"{p['season']}_{p['region']}" for p in all_pairs],
        random_state=SEED,
    )
    train_ds = SAROpticalDataset(train_pairs, img_size=IMG_SIZE, augment=True)
    train_dl = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True,
                          num_workers=2, pin_memory=True, drop_last=True, persistent_workers=False)
    vae = AutoencoderKL.from_pretrained(SD_MODEL_ID, subfolder='vae', torch_dtype=torch.float16)
    vae.requires_grad_(False)
    unet = UNet2DConditionModel.from_pretrained(SD_MODEL_ID, subfolder='unet', torch_dtype=torch.float32)
    old_conv = unet.conv_in
    new_conv = nn.Conv2d(8, old_conv.out_channels, old_conv.kernel_size,
                         old_conv.stride, old_conv.padding, bias=(old_conv.bias is not None))
    with torch.no_grad():
        new_conv.weight[:, :4].copy_(old_conv.weight)
        new_conv.weight[:, 4:].zero_()
        if old_conv.bias is not None: new_conv.bias.copy_(old_conv.bias)
    unet.conv_in = new_conv
    unet.config.in_channels = 8
    text_embed_dim  = unet.config.cross_attention_dim
    null_text_embed = nn.Parameter(torch.randn(1, 77, text_embed_dim) * 0.01)
    lora_cfg = LoraConfig(
        r=LORA_RANK, lora_alpha=LORA_ALPHA, use_dora=USE_DORA,
        init_lora_weights='gaussian',
        target_modules=['to_q','to_k','to_v','to_out.0','add_q_proj','add_k_proj','add_v_proj','ff.net.0.proj','ff.net.2'],
        lora_dropout=0.0, bias='none',
    )
    unet = get_peft_model(unet, lora_cfg)
    unet.conv_in.weight.requires_grad_(True)
    if unet.conv_in.bias is not None: unet.conv_in.bias.requires_grad_(True)
    if GRAD_CKPT: unet.enable_gradient_checkpointing()
    perceptual_loss_fn = VGGPerceptualLoss()
    noise_scheduler = DDPMScheduler(num_train_timesteps=NUM_TIMESTEPS,
                                    beta_schedule='scaled_linear', prediction_type='epsilon')
    ema_unet = EMAModel(unet.parameters(), decay=EMA_DECAY)
    trainable_params = list(filter(lambda p: p.requires_grad, unet.parameters())) + [null_text_embed]
    total_steps_across_sessions = RESUME_STEP + NUM_TRAIN_STEPS
    optimizer = torch.optim.AdamW(trainable_params, lr=LR, betas=(0.9,0.999), weight_decay=1e-2, eps=1e-8)
    lr_scheduler = get_cosine_schedule_with_warmup(
        optimizer,
        num_warmup_steps=LR_WARMUP * GRAD_ACCUM,
        num_training_steps=total_steps_across_sessions * GRAD_ACCUM,
    )
    unet, vae, optimizer, train_dl, lr_scheduler = accelerator.prepare(
        unet, vae, optimizer, train_dl, lr_scheduler
    )
    perceptual_loss_fn = perceptual_loss_fn.to(accelerator.device)
    null_text_embed    = null_text_embed.to(accelerator.device)
    ema_unet.to(accelerator.device)
    if RESUME_FROM is not None:
        adapter_path = os.path.join(RESUME_FROM, 'unet_adapter')
        embed_path   = os.path.join(RESUME_FROM, 'null_embed.pt')
        if os.path.exists(adapter_path):
            accelerator.unwrap_model(unet).load_adapter(adapter_path, adapter_name='default')
            tqdm.write(f'Loaded adapter from {adapter_path}')
        if os.path.exists(embed_path):
            null_text_embed.data = torch.load(embed_path, map_location=accelerator.device).data
            tqdm.write(f'Loaded null_embed from {embed_path}')
        for _ in range(RESUME_STEP * GRAD_ACCUM):
            lr_scheduler.step()
        tqdm.write(f'LR scheduler fast-forwarded to step {RESUME_STEP}')

    @torch.no_grad()
    def encode_to_latent(images):
        images  = images.to(dtype=vae.dtype)
        vae_mod = vae.module if hasattr(vae, 'module') else vae
        return vae_mod.encode(images).latent_dist.sample() * vae_mod.config.scaling_factor

    steps_per_epoch = len(train_dl) // GRAD_ACCUM
    total_epochs    = (NUM_TRAIN_STEPS + steps_per_epoch - 1) // steps_per_epoch
    global_step     = RESUME_STEP
    avg_loss        = 0.0
    overall_bar = tqdm(total=NUM_TRAIN_STEPS, desc=f'Training (step {RESUME_STEP} -> {RESUME_STEP+NUM_TRAIN_STEPS})',
                       disable=not is_main, position=0, dynamic_ncols=True, leave=True)
    unet.train()
    for epoch in range(total_epochs):
        accum_loss   = 0.0
        batches_seen = 0
        for batch in train_dl:
            with accelerator.accumulate(unet):
                sar_imgs    = batch['sar'].to(accelerator.device)
                opt_imgs    = batch['optical'].to(accelerator.device)
                opt_latents = encode_to_latent(opt_imgs)
                sar_latents = encode_to_latent(sar_imgs)
                noise       = torch.randn_like(opt_latents)
                B           = opt_latents.shape[0]
                timesteps   = torch.randint(0, NUM_TIMESTEPS, (B,), device=accelerator.device, dtype=torch.long)
                noisy_latents = noise_scheduler.add_noise(opt_latents, noise, timesteps)
                model_input   = torch.cat([noisy_latents, sar_latents], dim=1)
                enc_hidden    = null_text_embed.to(opt_latents.dtype).expand(B, -1, -1)
                noise_pred    = unet(model_input, timesteps, encoder_hidden_states=enc_hidden).sample
                loss_diffusion  = F.mse_loss(noise_pred.float(), noise.float())
                loss_color      = torch.tensor(0.0, device=accelerator.device)
                loss_perceptual = torch.tensor(0.0, device=accelerator.device)
                if global_step % COLOR_LOSS_FREQ == 0:
                    loss_color = color_supervision_loss(
                        noise_pred, noise, model_input, timesteps, noise_scheduler, vae, opt_latents)
                if global_step % PERCEPTUAL_FREQ == 0:
                    low_mask = timesteps < 300
                    if low_mask.sum() > 0:
                        vae_mod   = vae.module if hasattr(vae, 'module') else vae
                        alphas_cp = noise_scheduler.alphas_cumprod.to(accelerator.device)
                        sqrt_a    = alphas_cp[timesteps[low_mask]].sqrt().view(-1,1,1,1)
                        sqrt_1ma  = (1 - alphas_cp[timesteps[low_mask]]).sqrt().view(-1,1,1,1)
                        x0_pred   = (noisy_latents[low_mask] - sqrt_1ma * noise_pred[low_mask]) / (sqrt_a + 1e-8)
                        x0_pred   = x0_pred.clamp(-4, 4)
                        with torch.no_grad():
                            pred_img = vae_mod.decode((x0_pred / vae_mod.config.scaling_factor).to(vae_mod.dtype)).sample
                            gt_img   = vae_mod.decode((opt_latents[low_mask] / vae_mod.config.scaling_factor).to(vae_mod.dtype)).sample
                        loss_perceptual = perceptual_loss_fn(pred_img.float().clamp(-1,1), gt_img.float().clamp(-1,1))
                loss = loss_diffusion + COLOR_LOSS_W * loss_color + PERCEPTUAL_W * loss_perceptual
                accum_loss   += loss.item()
                batches_seen += 1
                accelerator.backward(loss)
                if accelerator.sync_gradients:
                    accelerator.clip_grad_norm_(accelerator.unwrap_model(unet).parameters(), 1.0)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad(set_to_none=True)
            if accelerator.sync_gradients:
                ema_unet.step(accelerator.unwrap_model(unet).parameters())
                global_step += 1
                avg_loss     = accum_loss / batches_seen
                lr_now       = lr_scheduler.get_last_lr()[0]
                accum_loss   = 0.0
                batches_seen = 0
                if is_main:
                    overall_bar.update(1)
                    overall_bar.set_postfix({'loss': f'{avg_loss:.4f}', 'lr': f'{lr_now:.1e}', 'step': global_step})
                if global_step % LOG_EVERY == 0 and is_main:
                    accelerator.log({'train/loss': avg_loss, 'train/loss_color': loss_color.item(), 'train/lr': lr_now}, step=global_step)
                if global_step % VIS_EVERY == 0 and is_main:
                    unet.eval()
                    ema_unwrapped = accelerator.unwrap_model(unet)
                    ema_unet.copy_to(ema_unwrapped.parameters())
                    grid_path = save_validation_grid(
                        val_pairs, ema_unwrapped,
                        vae.module if hasattr(vae, 'module') else vae,
                        null_text_embed, accelerator.device, global_step,
                        os.path.join(OUTPUT_DIR, 'val_grids'),
                    )
                    tqdm.write(f'  Val grid -> {grid_path}')
                    unet.train()
                if global_step % SAVE_EVERY == 0 and is_main:
                    ckpt = os.path.join(OUTPUT_DIR, f'step_{global_step}')
                    os.makedirs(ckpt, exist_ok=True)
                    unwrapped = accelerator.unwrap_model(unet)
                    unwrapped.save_pretrained(os.path.join(ckpt, 'unet_adapter'))
                    torch.save(null_text_embed, os.path.join(ckpt, 'null_embed.pt'))
                    tqdm.write(f'step {global_step} | adapter saved -> {ckpt}')
            if global_step >= RESUME_STEP + NUM_TRAIN_STEPS: break
        gc.collect()
        torch.cuda.empty_cache()
        if is_main: tqdm.write(f'Epoch {epoch+1}/{total_epochs} | step {global_step} | loss {avg_loss:.4f}')
        if global_step >= RESUME_STEP + NUM_TRAIN_STEPS: break
    overall_bar.close()
    if is_main:
        final_dir = os.path.join(OUTPUT_DIR, f'session_final_step{global_step}')
        os.makedirs(final_dir, exist_ok=True)
        unwrapped = accelerator.unwrap_model(unet)
        ema_unet.copy_to(unwrapped.parameters())
        unwrapped.save_pretrained(os.path.join(final_dir, 'unet_adapter'))
        merged = unwrapped.merge_and_unload()
        merged.save_pretrained(os.path.join(final_dir, 'unet_full'))
        torch.save(null_text_embed, os.path.join(final_dir, 'null_embed.pt'))
        vae_mod = vae.module if hasattr(vae, 'module') else vae
        vae_mod.save_pretrained(os.path.join(final_dir, 'vae'))
        tqdm.write(f'Session complete — step {global_step}. Saved -> {final_dir}')
        tqdm.write(f'To resume next session set:')
        tqdm.write(f'  RESUME_STEP = {global_step}')
        tqdm.write(f'  RESUME_FROM = "{final_dir}"')
    accelerator.end_training()

if __name__ == '__main__':
    main()