predict_diffusion.py

#!/usr/bin/env python3
"""
Conditional Latent Diffusion Model Inference for Himawari Satellite Images

This script uses a trained conditional latent diffusion model to generate
satellite imagery predictions based on input metadata conditioning.

Usage:
    # Generate from metadata only
    python predict_diffusion.py --metadata input.json --output prediction.png

    # Use reference image for guidance
    python predict_diffusion.py --metadata input.json --reference-image ref.png --output prediction.png

    # Batch prediction on multiple metadata files
    python predict_diffusion.py --metadata-dir ./metadata --output-dir ./predictions
"""

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

import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from torchvision import transforms

# Diffusers
from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler, DDIMScheduler


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,
        unet_out_channels: int = 4,
        conditioning_dim: int = 12,
        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)

        # Conditioning encoder
        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",
        )

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

        Args:
            conditioning: Metadata conditioning tensor (batch_size, conditioning_dim)
            num_inference_steps: Number of denoising steps
            guidance_scale: Classifier-free guidance scale (not fully implemented)
            generator: Random generator for reproducibility
            reference_latent: Optional reference latent for guidance
            reference_strength: Strength of reference guidance (0-1)
        """
        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 or initialize from reference
        latent_shape = (batch_size, 4, 64, 64)  # For 512x512 images

        if reference_latent is not None and reference_strength > 0:
            # Start from partially noised reference latent
            noise = torch.randn(latent_shape, device=device, generator=generator)
            # Mix reference with noise based on strength
            start_timestep = int(num_inference_steps * reference_strength)
            timestep = scheduler.timesteps[start_timestep]
            latents = scheduler.add_noise(reference_latent, noise, timestep)
            # Skip early steps
            timesteps = scheduler.timesteps[start_timestep:]
        else:
            # Pure random initialization
            latents = torch.randn(latent_shape, device=device, generator=generator)
            timesteps = scheduler.timesteps

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

        # Denoising loop
        for t in 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 encode_temporal_features(obs_time_str: str) -> torch.Tensor:
    """Encode temporal features from observation time."""
    try:
        dt = datetime.fromisoformat(obs_time_str.replace('Z', '+00:00'))

        hour = dt.hour
        day = dt.timetuple().tm_yday
        month = dt.month

        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(metadata: Dict) -> torch.Tensor:
    """Encode spatial features (geographic bounds)."""
    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)

    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(metadata: Dict) -> torch.Tensor:
    """Encode all conditioning features from metadata."""
    temporal = encode_temporal_features(metadata.get('observation_time_utc', ''))
    spatial = encode_spatial_features(metadata)

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

    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)

    conditioning = torch.cat([temporal, spatial, enhanced, satellite_feat])

    return conditioning


def load_and_preprocess_image(image_path: str, image_size: int = 512, device: str = "cuda") -> torch.Tensor:
    """Load and preprocess reference image."""
    image = Image.open(image_path)

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

    # Transform
    transform = transforms.Compose([
        transforms.Resize((image_size, image_size)),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])

    image_tensor = transform(image).unsqueeze(0).to(device)

    return image_tensor


def encode_image_to_latent(model: ConditionalLatentDiffusion, image_tensor: torch.Tensor) -> torch.Tensor:
    """Encode image to latent space using VAE."""
    with torch.no_grad():
        latent = model.vae.encode(image_tensor).latent_dist.sample()
        latent = latent * model.vae.config.scaling_factor
    return latent


def predict_single(
    model: ConditionalLatentDiffusion,
    metadata: Dict,
    output_path: str,
    reference_image_path: Optional[str] = None,
    reference_strength: float = 0.3,
    num_inference_steps: int = 50,
    image_size: int = 512,
    device: str = "cuda",
    seed: Optional[int] = None
):
    """
    Generate prediction for a single metadata file.

    Args:
        model: Trained diffusion model
        metadata: Metadata dictionary with conditioning information
        output_path: Path to save predicted image
        reference_image_path: Optional reference image for guidance
        reference_strength: How much to use reference (0=none, 1=full)
        num_inference_steps: Number of denoising steps
        image_size: Output image size
        device: Device to use
        seed: Random seed for reproducibility
    """
    model.eval()

    # Set random seed if provided
    generator = None
    if seed is not None:
        generator = torch.Generator(device=device).manual_seed(seed)
        print(f"Using seed: {seed}")

    # Encode metadata to conditioning
    conditioning = encode_metadata(metadata).unsqueeze(0).to(device)

    # Load and encode reference image if provided
    reference_latent = None
    if reference_image_path and os.path.exists(reference_image_path):
        print(f"Using reference image: {reference_image_path}")
        reference_tensor = load_and_preprocess_image(reference_image_path, image_size, device)
        reference_latent = encode_image_to_latent(model, reference_tensor)

    # Generate
    print(f"Generating with {num_inference_steps} steps...")
    generated = model.generate(
        conditioning,
        num_inference_steps=num_inference_steps,
        generator=generator,
        reference_latent=reference_latent,
        reference_strength=reference_strength
    )

    # Denormalize from [-1, 1] to [0, 1]
    generated = (generated + 1) / 2
    generated = torch.clamp(generated, 0, 1)

    # Convert to PIL and save
    generated_np = generated[0].permute(1, 2, 0).cpu().numpy()
    generated_np = (generated_np * 255).astype(np.uint8)

    output_image = Image.fromarray(generated_np)
    output_image.save(output_path)

    print(f"Saved prediction to: {output_path}")


def predict_batch(
    model: ConditionalLatentDiffusion,
    metadata_dir: str,
    output_dir: str,
    reference_image_dir: Optional[str] = None,
    reference_strength: float = 0.3,
    num_inference_steps: int = 50,
    image_size: int = 512,
    device: str = "cuda",
):
    """Batch prediction on multiple metadata files."""
    metadata_path = Path(metadata_dir)
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    # Find all metadata files
    metadata_files = list(metadata_path.glob("*.json"))
    metadata_files = [f for f in metadata_files if not f.name.startswith("_")]

    print(f"Found {len(metadata_files)} metadata files")

    for metadata_file in metadata_files:
        print(f"\nProcessing: {metadata_file.name}")

        # Load metadata
        with open(metadata_file) as f:
            metadata = json.load(f)

        # Output filename
        output_filename = metadata_file.stem + "_predicted.png"
        output_file = output_path / output_filename

        # Check for reference image
        reference_image = None
        if reference_image_dir:
            ref_dir = Path(reference_image_dir)
            # Try to find matching image
            ref_candidates = [
                ref_dir / (metadata_file.stem + ext)
                for ext in [".png", ".jpg", ".jpeg"]
            ]
            for ref_path in ref_candidates:
                if ref_path.exists():
                    reference_image = str(ref_path)
                    break

        # Generate prediction
        try:
            predict_single(
                model,
                metadata,
                str(output_file),
                reference_image_path=reference_image,
                reference_strength=reference_strength,
                num_inference_steps=num_inference_steps,
                image_size=image_size,
                device=device
            )
        except Exception as e:
            print(f"Error processing {metadata_file.name}: {e}")
            continue


def main():
    parser = argparse.ArgumentParser(description="Predict cloud patterns using trained diffusion model")

    # Input/output
    parser.add_argument("--metadata", type=str, help="Path to metadata JSON file")
    parser.add_argument("--metadata-dir", type=str, help="Directory with metadata files (batch mode)")
    parser.add_argument("--image-dir", type=str, help="Directory with input images (batch mode)")
    parser.add_argument("--output", type=str, help="Output image path (single mode)")
    parser.add_argument("--output-dir", type=str, default="./predictions", help="Output directory (batch mode)")

    # Model
    parser.add_argument("--model-path", type=str, required=True, help="Path to trained model checkpoint")
    parser.add_argument("--image-size", type=int, default=256, help="Image size (should match training)")

    # Reference image (optional)
    parser.add_argument("--reference-image", type=str, help="Reference image for guidance (single mode)")
    parser.add_argument("--reference-strength", type=float, default=0.3, help="Reference image strength (0-1)")

    # Generation parameters
    parser.add_argument("--num-steps", type=int, default=50, help="Number of inference steps")
    parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility")

    # Device
    parser.add_argument("--gpu-id", type=int, default=0, help="GPU ID to use")

    args = parser.parse_args()

    # Validate arguments
    if not args.metadata and not args.metadata_dir:
        parser.error("Either --metadata or --metadata-dir must be provided")

    if args.metadata and not args.output:
        parser.error("--output must be provided when using --metadata")

    # Setup device
    if torch.cuda.is_available():
        if args.gpu_id >= torch.cuda.device_count():
            print(f"WARNING: GPU {args.gpu_id} not available. Using GPU 0.")
            args.gpu_id = 0
        torch.cuda.set_device(args.gpu_id)
        device = f"cuda:{args.gpu_id}"
        print(f"Using GPU {args.gpu_id}: {torch.cuda.get_device_name(args.gpu_id)}")
    else:
        device = "cpu"
        print("WARNING: CUDA not available. Using CPU.")

    # Load model
    print(f"Loading model from: {args.model_path}")
    model = ConditionalLatentDiffusion(conditioning_dim=12)

    # Load checkpoint
    checkpoint = torch.load(args.model_path, map_location=device)
    if 'model_state_dict' in checkpoint:
        model.load_state_dict(checkpoint['model_state_dict'])
    else:
        model.load_state_dict(checkpoint)

    model = model.to(device)
    model.eval()
    print("Model loaded successfully")

    # Single or batch mode
    if args.metadata:
        # Single prediction
        print("\nSingle prediction mode")
        with open(args.metadata) as f:
            metadata = json.load(f)

        predict_single(
            model,
            metadata,
            args.output,
            reference_image_path=args.reference_image,
            reference_strength=args.reference_strength,
            num_inference_steps=args.num_steps,
            image_size=args.image_size,
            device=device,
            seed=args.seed
        )
    else:
        # Batch prediction
        print("\nBatch prediction mode")
        predict_batch(
            model,
            args.metadata_dir,
            args.output_dir,
            reference_image_dir=args.image_dir,
            reference_strength=args.reference_strength,
            num_inference_steps=args.num_steps,
            image_size=args.image_size,
            device=device
        )

    print("\nPrediction complete!")


if __name__ == "__main__":
    main()