trainer.py

"""
    SD15 Flow-Matching trainer
    Author: AbstractPhil

    Loads the current format pt and ensures through multiple validations that the process is correct for training.

    Trains flow matching for sd15.

    License: MIT
    If you use my work, a cite wouldnt hurt.

"""

import os
import json
import datetime
from dataclasses import dataclass, asdict
from tqdm.auto import tqdm
import matplotlib.pyplot as plt

import torch
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader

import datasets
from diffusers import UNet2DConditionModel
from huggingface_hub import HfApi, create_repo, hf_hub_download


@dataclass
class TrainConfig:
    output_dir: str = "./outputs"
    model_repo: str = "AbstractPhil/sd15-flow-matching-try2"
    checkpoint_filename: str = "sd15_flowmatch_david_weighted_2_e34.pt"
    dataset_name: str = "AbstractPhil/sd15-latent-distillation-500k"
    
    # HuggingFace upload settings
    hf_repo_id: str = "AbstractPhil/sd15-flow-lune"
    upload_to_hub: bool = True
    
    seed: int = 42
    batch_size: int = 16
    base_lr: float = 2e-6
    shift: float = 2.0
    dropout: float = 0.1
    
    max_train_steps: int = 50_000
    checkpointing_steps: int = 1000
    num_workers: int = 0
    
    # VAE scaling factor - multiply raw latents
    vae_scale: float = 0.18215


def load_student_unet(repo_id: str, filename: str, device="cuda") -> UNet2DConditionModel:
    """Load UNet from .pt checkpoint containing student state_dict"""
    # Download checkpoint from HuggingFace
    print(f"Downloading checkpoint from {repo_id}/{filename}...")
    checkpoint_path = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        repo_type="model"
    )
    print(f"✓ Downloaded to: {checkpoint_path}")
    
    checkpoint = torch.load(checkpoint_path, map_location="cpu")
    
    # Initialize UNet with SD1.5 config in fp32
    print("Loading SD1.5 UNet architecture...")
    unet = UNet2DConditionModel.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        subfolder="unet",
        torch_dtype=torch.float32
    )
    
    # Get original state for comparison
    original_state_dict = {k: v.clone() for k, v in unet.state_dict().items()}
    
    # Load student weights and strip "unet." prefix
    student_state_dict = checkpoint["student"]
    
    # Strip prefix if present
    cleaned_student_dict = {}
    for key, value in student_state_dict.items():
        if key.startswith("unet."):
            cleaned_key = key[5:]  # Remove "unet." prefix
            cleaned_student_dict[cleaned_key] = value
        else:
            cleaned_student_dict[key] = value
    
    print(f"\n{'='*70}")
    print("WEIGHT VERIFICATION")
    print(f"{'='*70}")
    
    # 1. Compare keys
    original_keys = set(original_state_dict.keys())
    student_keys = set(cleaned_student_dict.keys())
    
    matching_keys = original_keys & student_keys
    
    print(f"Original UNet keys: {len(original_keys)}")
    print(f"Student checkpoint keys: {len(student_keys)}")
    print(f"Matching keys: {len(matching_keys)}")
    
    # 2. Compare student weights vs original BEFORE loading
    total_params = 0
    different_params = 0
    mean_diff_sum = 0.0
    max_diff = 0.0
    
    for key in matching_keys:
        if key not in original_state_dict or key not in cleaned_student_dict:
            continue
            
        orig = original_state_dict[key]
        student = cleaned_student_dict[key].float()  # Convert to fp32 for comparison
        
        if orig.shape != student.shape:
            print(f"⚠ Shape mismatch for {key}: {orig.shape} vs {student.shape}")
            continue
        
        total_params += orig.numel()
        
        # Check if weights are different
        diff = (orig - student).abs()
        if diff.max() > 1e-6:
            different_params += orig.numel()
            mean_diff_sum += diff.sum().item()
            max_diff = max(max_diff, diff.max().item())
    
    pct_different = (different_params / total_params * 100) if total_params > 0 else 0
    avg_diff = mean_diff_sum / different_params if different_params > 0 else 0
    
    print(f"\nStudent vs Original (BEFORE loading):")
    print(f"  Total parameters: {total_params:,}")
    print(f"  Parameters different: {different_params:,} ({pct_different:.1f}%)")
    print(f"  Average difference: {avg_diff:.6f}")
    print(f"  Max difference: {max_diff:.6f}")
    
    # 3. Load weights
    load_result = unet.load_state_dict(cleaned_student_dict, strict=False)
    
    if load_result.missing_keys:
        print(f"\n⚠ Missing keys during load: {len(load_result.missing_keys)}")
        for key in load_result.missing_keys[:3]:
            print(f"  - {key}")
    
    if load_result.unexpected_keys:
        print(f"⚠ Unexpected keys during load: {len(load_result.unexpected_keys)}")
        for key in load_result.unexpected_keys[:3]:
            print(f"  - {key}")
    
    # 4. Verify weights actually changed after loading
    loaded_state_dict = unet.state_dict()
    
    total_params_after = 0
    changed_params = 0
    mean_diff_after = 0.0
    max_diff_after = 0.0
    
    for key in matching_keys:
        if key not in original_state_dict or key not in loaded_state_dict:
            continue
            
        orig = original_state_dict[key]
        loaded = loaded_state_dict[key]
        
        total_params_after += orig.numel()
        
        diff = (orig - loaded).abs()
        if diff.max() > 1e-6:
            changed_params += orig.numel()
            mean_diff_after += diff.sum().item()
            max_diff_after = max(max_diff_after, diff.max().item())
    
    pct_changed = (changed_params / total_params_after * 100) if total_params_after > 0 else 0
    avg_diff_after = mean_diff_after / changed_params if changed_params > 0 else 0
    
    print(f"\nOriginal vs Loaded (AFTER loading):")
    print(f"  Parameters changed: {changed_params:,} ({pct_changed:.1f}%)")
    print(f"  Average difference: {avg_diff_after:.6f}")
    print(f"  Max difference: {max_diff_after:.6f}")
    
    print(f"\n{'='*70}")
    # Verification checks
    if pct_different < 50:
        print(f"⚠️ WARNING: Student weights only {pct_different:.1f}% different from base!")
        print("   This checkpoint may not be trained.")
    elif pct_changed < 90:
        print(f"⚠️ WARNING: Only {pct_changed:.1f}% of weights changed after loading!")
        print("   The load may have failed.")
    else:
        print(f"✅ Weights loaded successfully!")
        print(f"   Checkpoint step: {checkpoint.get('gstep', 'unknown')}")
        print(f"   {pct_different:.1f}% of weights differ from base SD1.5")
    
    print(f"{'='*70}\n")
    
    return unet.to(device)


def train(config: TrainConfig):
    device = "cuda"
    torch.backends.cuda.matmul.allow_tf32 = True
    
    torch.manual_seed(config.seed)
    torch.cuda.manual_seed(config.seed)
    
    # Setup output directory
    date_time = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    real_output_dir = os.path.join(config.output_dir, date_time)
    os.makedirs(real_output_dir, exist_ok=True)
    t_writer = SummaryWriter(log_dir=real_output_dir, flush_secs=60)
    
    # Initialize HuggingFace API
    hf_api = None
    if config.upload_to_hub:
        try:
            hf_api = HfApi()
            create_repo(
                repo_id=config.hf_repo_id,
                repo_type="model",
                exist_ok=True,
                private=False
            )
            print(f"✓ HuggingFace repo ready: {config.hf_repo_id}")
        except Exception as e:
            print(f"⚠ Hub upload disabled: {e}")
            config.upload_to_hub = False
    
    # Save config locally and to hub
    config_path = os.path.join(real_output_dir, "config.json")
    with open(config_path, "w") as f:
        json.dump(asdict(config), f, indent=2)
    
    if config.upload_to_hub:
        hf_api.upload_file(
            path_or_fileobj=config_path,
            path_in_repo="config.json",
            repo_id=config.hf_repo_id,
            repo_type="model"
        )
    
    # Load dataset in streaming mode
    print(f"\nLoading dataset (streaming): {config.dataset_name}")
    train_dataset = datasets.load_dataset(
        config.dataset_name,
        split="train",
        streaming=True,
        trust_remote_code=True
    )
    train_dataset = train_dataset.shuffle(seed=config.seed, buffer_size=1000)
    print(f"✓ Dataset loaded in streaming mode")
    
    def collate_fn(examples):
        # Latents are RAW from VAE - need to scale them
        latents = torch.stack([torch.tensor(ex["latent"]) for ex in examples])
        latents = latents * config.vae_scale  # Scale: ~[-6, 6] -> ~[-1, 1]
        
        clip_embeddings = torch.stack([torch.tensor(ex["clip_embedding"]) for ex in examples])
        ids = [ex["id"] for ex in examples]
        prompts = [ex["prompt"] for ex in examples]
        
        return latents, clip_embeddings, ids, prompts
    
    train_dataloader = DataLoader(
        dataset=train_dataset,
        batch_size=config.batch_size,
        collate_fn=collate_fn,
        num_workers=config.num_workers,
    )
    
    # Verify first batch latent range (on GPU for speed)
    print("\nVerifying latent scaling on first batch...")
    first_batch = next(iter(train_dataloader))
    latents_check, _, _, _ = first_batch
    print(f"Raw latent range: [{latents_check.min():.3f}, {latents_check.max():.3f}]")
    latents_check = latents_check.to(device)
    print(f"After GPU transfer: [{latents_check.min():.3f}, {latents_check.max():.3f}]")
    print(f"Expected: ~[-1, 1] for properly scaled latents")
    del latents_check
    
    # Load pretrained student UNet
    print(f"\nLoading model from HuggingFace...")
    unet = load_student_unet(config.model_repo, config.checkpoint_filename, device=device)
    unet.requires_grad_(True)
    unet.enable_gradient_checkpointing()
    unet.train()
    
    optimizer = torch.optim.Adam(
        unet.parameters(),
        lr=config.base_lr * (config.batch_size ** 0.5),
    )
    
    global_step = 0
    train_logs = {
        "train_step": [],
        "train_loss": [],
        "train_timestep": [],
        "trained_images": []
    }
    
    def get_prediction(batch, log_to=None):
        latents, encoder_hidden_states, ids, prompts = batch
        
        # Everything in fp32
        latents = latents.to(dtype=torch.float32, device=device)
        encoder_hidden_states = encoder_hidden_states.to(dtype=torch.float32, device=device)
        
        batch_size = latents.shape[0]
        
        # Apply dropout to conditioning for CFG support
        dropout_mask = torch.rand(batch_size, device=device) < config.dropout
        encoder_hidden_states = encoder_hidden_states.clone()
        encoder_hidden_states[dropout_mask] = 0
        
        # Sample timesteps with shift
        sigmas = torch.rand(batch_size, device=device)
        sigmas = (config.shift * sigmas) / (1 + (config.shift - 1) * sigmas)
        timesteps = sigmas * 1000
        sigmas = sigmas[:, None, None, None]
        
        # Flow matching forward process
        noise = torch.randn_like(latents)
        noisy_latents = noise * sigmas + latents * (1 - sigmas)
        target = noise - latents
        
        # Predict velocity
        pred = unet(noisy_latents, timesteps, encoder_hidden_states, return_dict=False)[0]
        
        loss = F.mse_loss(pred, target, reduction="none")
        loss = loss.mean(dim=list(range(1, len(loss.shape))))
        
        if log_to is not None:
            for i in range(batch_size):
                log_to["train_step"].append(global_step)
                log_to["train_loss"].append(loss[i].item())
                log_to["train_timestep"].append(timesteps[i].item())
                log_to["trained_images"].append({
                    "step": global_step,
                    "id": ids[i],
                    "prompt": prompts[i]
                })
        
        return loss.mean()
    
    def plot_logs(log_dict):
        plt.figure(figsize=(10, 6))
        plt.scatter(
            log_dict["train_timestep"],
            log_dict["train_loss"],
            s=3,
            c=log_dict["train_step"],
            marker=".",
            cmap='cool'
        )
        plt.xlabel("timestep")
        plt.ylabel("loss")
        plt.yscale("log")
        plt.colorbar(label="step")
    
    def save_checkpoint(step):
        checkpoint_path = os.path.join(real_output_dir, f"checkpoint-{step:08}")
        os.makedirs(checkpoint_path, exist_ok=True)
        
        # Save UNet weights as diffusers format
        unet.save_pretrained(
            os.path.join(checkpoint_path, "unet"),
            safe_serialization=True
        )
        
        # Save complete checkpoint in .pt format
        pt_filename = f"sd15_flow_lune_e{step//1000}_s{step}.pt"
        pt_path = os.path.join(checkpoint_path, pt_filename)
        
        torch.save({
            "cfg": asdict(config),
            "student": unet.state_dict(),
            "opt": optimizer.state_dict(),
            "gstep": step
        }, pt_path)
        
        # Save training metadata
        metadata = {
            "step": step,
            "trained_images": train_logs["trained_images"]
        }
        metadata_path = os.path.join(checkpoint_path, "trained_images.json")
        with open(metadata_path, "w") as f:
            json.dump(metadata, f, indent=2)
        
        print(f"✓ Checkpoint saved at step {step}")
        
        # Upload to HuggingFace Hub
        if config.upload_to_hub and hf_api is not None:
            try:
                hf_api.upload_file(
                    path_or_fileobj=pt_path,
                    path_in_repo=pt_filename,
                    repo_id=config.hf_repo_id,
                    repo_type="model"
                )
                
                hf_api.upload_folder(
                    folder_path=os.path.join(checkpoint_path, "unet"),
                    path_in_repo=f"checkpoint-{step:08}/unet",
                    repo_id=config.hf_repo_id,
                    repo_type="model"
                )
                
                hf_api.upload_file(
                    path_or_fileobj=metadata_path,
                    path_in_repo=f"checkpoint-{step:08}/trained_images.json",
                    repo_id=config.hf_repo_id,
                    repo_type="model"
                )
                
                print(f"✓ Uploaded to hub: {config.hf_repo_id}")
            except Exception as e:
                print(f"⚠ Upload failed: {e}")
    
    print("\nStarting training...")
    progress_bar = tqdm(range(0, config.max_train_steps))
    
    for batch in train_dataloader:
        loss = get_prediction(batch, log_to=train_logs)
        t_writer.add_scalar("train/loss", loss.detach().item(), global_step)
        
        loss.backward()
        
        grad_norm = torch.nn.utils.clip_grad_norm_(unet.parameters(), 2.0)
        t_writer.add_scalar("train/grad_norm", grad_norm.detach().item(), global_step)
        
        optimizer.step()
        optimizer.zero_grad()
        
        progress_bar.update(1)
        progress_bar.set_postfix({"loss": f"{loss.item():.4f}"})
        global_step += 1
        
        if global_step % 100 == 0:
            plot_logs(train_logs)
            t_writer.add_figure("train_loss", plt.gcf(), global_step)
            plt.close()
        
        if global_step % config.checkpointing_steps == 0:
            save_checkpoint(global_step)
        
        if global_step >= config.max_train_steps:
            save_checkpoint(global_step)
            print("\n✅ Training complete!")
            return


if __name__ == "__main__":
    config = TrainConfig()
    train(config)