train_diffusion_truecolor.py

#!/usr/bin/env python3
"""
Conditional Latent Diffusion Model Training for Himawari True Color RGB Images

This script trains a conditional latent diffusion model on true color satellite imagery.
The model is conditioned on metadata (observation time, satellite position, etc.) to
generate realistic cloud patterns and weather conditions.

Architecture:
- Uses Stable Diffusion architecture with custom conditioning
- Encoder: Compresses RGB images to latent space
- U-Net: Denoising network with cross-attention for conditioning
- Decoder: Reconstructs images from latent space

Conditioning includes:
- Temporal features (hour, day, season)
- Spatial features (lat, lon bounds)
- Satellite metadata (band info, resolution)
"""

import os
import json
import argparse
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Tuple, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
from PIL import Image
from tqdm import tqdm

# Diffusers
from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler, DDIMScheduler
from diffusers.optimization import get_cosine_schedule_with_warmup


class TrueColorSatelliteDataset(Dataset):
    """
    Dataset for True Color RGB satellite imagery with metadata conditioning.
    """

    def __init__(
        self,
        data_dir: str,
        image_size: int = 512,
        normalize: bool = True,
        augment: bool = True
    ):
        self.data_dir = Path(data_dir)
        self.image_dir = self.data_dir / "images"
        self.metadata_dir = self.data_dir / "metadata"
        self.image_size = image_size
        self.normalize = normalize

        # Load all samples
        self.samples = self._load_samples()

        # Image transforms
        transform_list = [
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
        ]

        if normalize:
            # Normalize to [-1, 1] for diffusion models
            transform_list.append(transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]))

        if augment:
            self.augment_transforms = transforms.Compose([
                transforms.RandomHorizontalFlip(p=0.5),
                transforms.RandomVerticalFlip(p=0.3),
                transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1),
            ])
        else:
            self.augment_transforms = None

        self.image_transforms = transforms.Compose(transform_list)

    def _load_samples(self) -> List[Dict]:
        """Load metadata for all samples."""
        samples = []
        for json_file in self.metadata_dir.glob("*.json"):
            if json_file.name.startswith("_"):
                continue
            with open(json_file) as f:
                metadata = json.load(f)

            # Check if corresponding image exists
            if isinstance(metadata.get('filename'), str):
                image_path = self.image_dir / metadata['filename']
            else:
                image_path = self.image_dir / (json_file.stem + ".png")

            if image_path.exists():
                metadata['_image_path'] = str(image_path)
                samples.append(metadata)

        return samples

    def _encode_temporal_features(self, obs_time_str: str) -> torch.Tensor:
        """
        Encode temporal features from observation time.
        Returns: [hour_sin, hour_cos, day_sin, day_cos, month_sin, month_cos]
        """
        try:
            # Parse timestamp
            dt = datetime.fromisoformat(obs_time_str.replace('Z', '+00:00'))

            # Cyclical encoding
            hour = dt.hour
            day = dt.timetuple().tm_yday  # Day of year
            month = dt.month

            # Sine/cosine encoding for cyclical features
            hour_sin = np.sin(2 * np.pi * hour / 24)
            hour_cos = np.cos(2 * np.pi * hour / 24)
            day_sin = np.sin(2 * np.pi * day / 365)
            day_cos = np.cos(2 * np.pi * day / 365)
            month_sin = np.sin(2 * np.pi * month / 12)
            month_cos = np.cos(2 * np.pi * month / 12)

            return torch.tensor([hour_sin, hour_cos, day_sin, day_cos, month_sin, month_cos], dtype=torch.float32)
        except Exception:
            return torch.zeros(6, dtype=torch.float32)

    def _encode_spatial_features(self, metadata: Dict) -> torch.Tensor:
        """
        Encode spatial features (geographic bounds).
        Returns: [min_lat, max_lat, min_lon, max_lon] normalized to [-1, 1]
        """
        min_lat = metadata.get('min_lat', 0.0)
        max_lat = metadata.get('max_lat', 0.0)
        min_lon = metadata.get('min_lon', 0.0)
        max_lon = metadata.get('max_lon', 0.0)

        # Normalize to [-1, 1]
        min_lat_norm = min_lat / 90.0
        max_lat_norm = max_lat / 90.0
        min_lon_norm = min_lon / 180.0
        max_lon_norm = max_lon / 180.0

        return torch.tensor([min_lat_norm, max_lat_norm, min_lon_norm, max_lon_norm], dtype=torch.float32)

    def _encode_metadata(self, metadata: Dict) -> torch.Tensor:
        """
        Encode all conditioning features.
        Returns: Concatenated feature vector
        """
        temporal = self._encode_temporal_features(metadata.get('observation_time_utc', ''))
        spatial = self._encode_spatial_features(metadata)

        # Additional features
        enhanced = torch.tensor([1.0 if metadata.get('enhanced', False) else 0.0], dtype=torch.float32)

        # Satellite ID (one-hot or continuous)
        satellite = metadata.get('satellite', 'Himawari-8')
        satellite_id = 8.0 if '8' in satellite else 9.0
        satellite_feat = torch.tensor([satellite_id / 10.0], dtype=torch.float32)  # Normalize

        # Concatenate all features
        conditioning = torch.cat([temporal, spatial, enhanced, satellite_feat])

        return conditioning

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, Dict]:
        metadata = self.samples[idx]

        # Load image - check if it's grayscale or RGB
        image = Image.open(metadata['_image_path'])

        # Convert grayscale to RGB by duplicating channels
        if image.mode == 'L':
            image = image.convert('RGB')
        elif image.mode != 'RGB':
            image = image.convert('RGB')

        # Apply augmentation before transforms
        if self.augment_transforms is not None:
            image = self.augment_transforms(image)

        # Apply transforms
        image_tensor = self.image_transforms(image)

        # Encode conditioning
        conditioning = self._encode_metadata(metadata)

        return image_tensor, conditioning, metadata


class ConditionalLatentDiffusion(nn.Module):
    """
    Conditional Latent Diffusion Model for satellite imagery.
    """

    def __init__(
        self,
        vae_model_name: str = "stabilityai/sd-vae-ft-mse",
        unet_in_channels: int = 4,  # Latent channels
        unet_out_channels: int = 4,
        conditioning_dim: int = 12,  # Metadata features
        cross_attention_dim: int = 768,
        num_train_timesteps: int = 1000,
    ):
        super().__init__()

        # VAE for encoding images to latent space
        self.vae = AutoencoderKL.from_pretrained(vae_model_name)
        self.vae.requires_grad_(False)  # Freeze VAE

        # Conditioning encoder: metadata -> embedding
        self.conditioning_encoder = nn.Sequential(
            nn.Linear(conditioning_dim, 256),
            nn.GELU(),
            nn.Linear(256, 512),
            nn.GELU(),
            nn.Linear(512, cross_attention_dim),
        )

        # U-Net denoiser
        self.unet = UNet2DConditionModel(
            in_channels=unet_in_channels,
            out_channels=unet_out_channels,
            cross_attention_dim=cross_attention_dim,
            block_out_channels=(128, 256, 512, 512),
            layers_per_block=2,
            attention_head_dim=8,
        )

        # Noise scheduler
        self.noise_scheduler = DDPMScheduler(
            num_train_timesteps=num_train_timesteps,
            beta_schedule="scaled_linear",
            prediction_type="epsilon",
        )

    def forward(
        self,
        images: torch.Tensor,
        conditioning: torch.Tensor,
        noise: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Forward pass for training.

        Returns:
            noise_pred: Predicted noise
            noise: Ground truth noise
        """
        batch_size = images.size(0)
        device = images.device

        # Encode images to latent space
        with torch.no_grad():
            latents = self.vae.encode(images).latent_dist.sample()
            latents = latents * self.vae.config.scaling_factor

        # Sample noise
        if noise is None:
            noise = torch.randn_like(latents)

        # Sample random timesteps
        timesteps = torch.randint(
            0, self.noise_scheduler.config.num_train_timesteps,
            (batch_size,), device=device
        ).long()

        # Add noise to latents
        noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)

        # Encode conditioning
        encoder_hidden_states = self.conditioning_encoder(conditioning)
        encoder_hidden_states = encoder_hidden_states.unsqueeze(1)  # Add sequence dimension

        # Predict noise
        noise_pred = self.unet(
            noisy_latents,
            timesteps,
            encoder_hidden_states=encoder_hidden_states
        ).sample

        return noise_pred, noise

    @torch.no_grad()
    def generate(
        self,
        conditioning: torch.Tensor,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        generator: Optional[torch.Generator] = None
    ) -> torch.Tensor:
        """
        Generate images from conditioning.
        """
        device = conditioning.device
        batch_size = conditioning.size(0)

        # Initialize DDIM scheduler for faster inference
        scheduler = DDIMScheduler.from_config(self.noise_scheduler.config)
        scheduler.set_timesteps(num_inference_steps)

        # Random latents
        latent_shape = (batch_size, 4, 64, 64)  # 512x512 -> 64x64 latent
        latents = torch.randn(latent_shape, device=device, generator=generator)

        # Encode conditioning
        encoder_hidden_states = self.conditioning_encoder(conditioning)
        encoder_hidden_states = encoder_hidden_states.unsqueeze(1)

        # Denoising loop
        for t in scheduler.timesteps:
            # Predict noise
            noise_pred = self.unet(
                latents,
                t,
                encoder_hidden_states=encoder_hidden_states
            ).sample

            # Compute previous noisy sample
            latents = scheduler.step(noise_pred, t, latents).prev_sample

        # Decode latents to images
        latents = latents / self.vae.config.scaling_factor
        images = self.vae.decode(latents).sample

        return images


def train(
    model: ConditionalLatentDiffusion,
    train_loader: DataLoader,
    val_loader: Optional[DataLoader],
    num_epochs: int,
    learning_rate: float,
    output_dir: str,
    device: str = "cuda",
    log_interval: int = 100,
    save_interval: int = 1000,
    mixed_precision: bool = False,
    resume_from: Optional[str] = None,
):
    """Train the conditional latent diffusion model."""

    model = model.to(device)

    # Mixed precision training for faster GPU training
    scaler = torch.cuda.amp.GradScaler() if mixed_precision and device == "cuda" else None
    if mixed_precision:
        print("Using mixed precision training (FP16)")

    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=learning_rate,
        weight_decay=0.01
    )

    num_training_steps = len(train_loader) * num_epochs
    lr_scheduler = get_cosine_schedule_with_warmup(
        optimizer,
        num_warmup_steps=500,
        num_training_steps=num_training_steps
    )

    # Resume from checkpoint if specified
    start_epoch = 0
    global_step = 0
    if resume_from is not None and os.path.exists(resume_from):
        print(f"Resuming from checkpoint: {resume_from}")
        checkpoint = torch.load(resume_from, map_location=device)
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        start_epoch = checkpoint.get('epoch', 0) + 1
        global_step = checkpoint.get('global_step', 0)
        print(f"Resuming from epoch {start_epoch}, global step {global_step}")

    # TensorBoard
    writer = SummaryWriter(os.path.join(output_dir, "logs"))

    for epoch in range(start_epoch, start_epoch + num_epochs):
        model.train()
        epoch_loss = 0.0

        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}")
        for batch_idx, (images, conditioning, metadata) in enumerate(pbar):
            images = images.to(device, non_blocking=True)
            conditioning = conditioning.to(device, non_blocking=True)

            # Forward pass with mixed precision
            if scaler is not None:
                with torch.cuda.amp.autocast():
                    noise_pred, noise = model(images, conditioning)
                    loss = F.mse_loss(noise_pred, noise)

                # Backward pass with gradient scaling
                optimizer.zero_grad()
                scaler.scale(loss).backward()
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                scaler.step(optimizer)
                scaler.update()
            else:
                # Standard forward pass
                noise_pred, noise = model(images, conditioning)
                loss = F.mse_loss(noise_pred, noise)

                # Backward pass
                optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()

            lr_scheduler.step()

            epoch_loss += loss.item()
            global_step += 1

            # Logging
            if global_step % log_interval == 0:
                writer.add_scalar("train/loss", loss.item(), global_step)
                writer.add_scalar("train/lr", lr_scheduler.get_last_lr()[0], global_step)
                pbar.set_postfix({"loss": f"{loss.item():.4f}"})

            # Save checkpoint
            if global_step % save_interval == 0:
                checkpoint_path = os.path.join(output_dir, f"checkpoint-{global_step}.pt")
                torch.save({
                    'epoch': epoch,
                    'global_step': global_step,
                    'model_state_dict': model.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'loss': loss.item(),
                }, checkpoint_path)
                print(f"\nSaved checkpoint: {checkpoint_path}")

        avg_epoch_loss = epoch_loss / len(train_loader)
        print(f"Epoch {epoch+1} - Avg Loss: {avg_epoch_loss:.4f}")
        writer.add_scalar("train/epoch_loss", avg_epoch_loss, epoch)

        # Validation
        if val_loader is not None:
            val_loss = validate(model, val_loader, device)
            print(f"Validation Loss: {val_loss:.4f}")
            writer.add_scalar("val/loss", val_loss, epoch)

            # Generate samples
            model.eval()
            with torch.no_grad():
                sample_conditioning = conditioning[:4].to(device)
                generated = model.generate(sample_conditioning, num_inference_steps=20)
                generated = (generated + 1) / 2  # Denormalize
                writer.add_images("val/generated", generated, epoch)

    # Save final model
    final_path = os.path.join(output_dir, "final_model.pt")
    torch.save(model.state_dict(), final_path)
    print(f"Training complete. Model saved to {final_path}")

    writer.close()


def validate(model: ConditionalLatentDiffusion, val_loader: DataLoader, device: str) -> float:
    """Validate the model."""
    model.eval()
    total_loss = 0.0

    with torch.no_grad():
        for images, conditioning, _ in val_loader:
            images = images.to(device)
            conditioning = conditioning.to(device)

            noise_pred, noise = model(images, conditioning)
            loss = F.mse_loss(noise_pred, noise)
            total_loss += loss.item()

    return total_loss / len(val_loader)


def main():
    parser = argparse.ArgumentParser(description="Train Conditional Latent Diffusion on Satellite Images")
    parser.add_argument("--data-dir", default="./true_color_output", help="Data directory (works with grayscale IR from create_true_color.py)")
    parser.add_argument("--output-dir", default="./models/diffusion", help="Output directory")
    parser.add_argument("--image-size", type=int, default=256, help="Image size (use 256 for faster training)")
    parser.add_argument("--batch-size", type=int, default=8, help="Batch size")
    parser.add_argument("--num-epochs", type=int, default=25, help="Number of epochs")
    parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate")
    parser.add_argument("--gpu-id", type=int, default=0, help="GPU ID to use (default: 0)")
    parser.add_argument("--val-split", type=float, default=0.1, help="Validation split ratio")
    parser.add_argument("--mixed-precision", action="store_true", help="Use mixed precision training (faster on GPU)")
    parser.add_argument("--resume-from", type=str, default=None, help="Path to checkpoint to resume training from")

    args = parser.parse_args()

    # Check GPU availability and set specific GPU
    if not torch.cuda.is_available():
        print("WARNING: CUDA not available. Falling back to CPU.")
        args.device = "cpu"
    else:
        # Check if requested GPU ID exists
        if args.gpu_id >= torch.cuda.device_count():
            print(f"WARNING: GPU {args.gpu_id} not available. Using GPU 0 instead.")
            args.gpu_id = 0

        # Set specific GPU
        torch.cuda.set_device(args.gpu_id)
        args.device = f"cuda:{args.gpu_id}"

        print(f"Using GPU {args.gpu_id}: {torch.cuda.get_device_name(args.gpu_id)}")
        print(f"GPU Memory: {torch.cuda.get_device_properties(args.gpu_id).total_memory / 1e9:.2f} GB")

    os.makedirs(args.output_dir, exist_ok=True)

    print("Loading dataset...")
    full_dataset = TrueColorSatelliteDataset(
        args.data_dir,
        image_size=args.image_size,
        augment=True
    )

    # Train/val split
    val_size = int(len(full_dataset) * args.val_split)
    train_size = len(full_dataset) - val_size
    train_dataset, val_dataset = torch.utils.data.random_split(
        full_dataset, [train_size, val_size]
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=4,
        pin_memory=True
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=4,
        pin_memory=True
    )

    print(f"Train samples: {len(train_dataset)}, Val samples: {len(val_dataset)}")

    print("Initializing model...")
    model = ConditionalLatentDiffusion(conditioning_dim=12)

    print("Starting training...")
    train(
        model,
        train_loader,
        val_loader,
        num_epochs=args.num_epochs,
        learning_rate=args.lr,
        output_dir=args.output_dir,
        device=args.device,
        mixed_precision=args.mixed_precision,
        resume_from=args.resume_from
    )


if __name__ == "__main__":
    main()