train.py

"""
LiquidGen Training Pipeline v2

Optimized for Colab free tier:
- Fast VAE encoding: batch=64 for 256px, batch=32 for 512px (~5x faster)
- Auto-limits large datasets (WikiArt capped at 10K by default)
- Latent pre-caching: train on pure tensors, no VAE during training
- Gradient checkpointing + auto batch size = no OOM
- ETA shown on every log line
- All datasets pure parquet, open SDXL VAE (no login)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from torch.amp import autocast, GradScaler
import math
import os
import json
import time
from dataclasses import dataclass, asdict


DATASET_PRESETS = {
    "cartoon": {
        "name": "Norod78/cartoon-blip-captions",
        "config": "",
        "image_column": "image",
        "label_column": "",
        "num_classes": 0,
        "max_default": 0,
        "description": "~2.5K cartoon/anime, unconditional, 181MB — fast",
    },
    "flowers": {
        "name": "huggan/flowers-102-categories",
        "config": "",
        "image_column": "image",
        "label_column": "",
        "num_classes": 0,
        "max_default": 0,
        "description": "~8K flower photos, unconditional, 331MB",
    },
    "wikiart": {
        "name": "Artificio/WikiArt",
        "config": "",
        "image_column": "image",
        "label_column": "style",
        "num_classes": 0,
        "max_default": 10000,
        "description": "~105K paintings with styles (auto-capped to 10K for speed)",
    },
    "art_painting": {
        "name": "huggan/few-shot-art-painting",
        "config": "",
        "image_column": "image",
        "label_column": "",
        "num_classes": 0,
        "max_default": 0,
        "description": "~6K art paintings, unconditional, 511MB",
    },
}


def auto_batch_size(model_size, image_size, gpu_mem_gb):
    param_mem = {"small": 0.66, "base": 1.68, "large": 3.35}
    base = param_mem.get(model_size, 1.0)
    act_per_sample = {"small": {256: 0.02, 512: 0.07},
                      "base":  {256: 0.03, 512: 0.13},
                      "large": {256: 0.05, 512: 0.21}}
    per_sample = act_per_sample.get(model_size, {}).get(image_size, 0.1)
    available = gpu_mem_gb - base - 1.5
    bs = max(1, int(available / per_sample))
    if bs >= 32: return 32
    if bs >= 16: return 16
    if bs >= 8: return 8
    if bs >= 4: return 4
    return max(1, bs)


def _fmt_time(seconds):
    """Format seconds into human readable string."""
    if seconds < 60: return f"{seconds:.0f}s"
    if seconds < 3600: return f"{seconds/60:.1f}m"
    return f"{seconds/3600:.1f}h"


@dataclass
class TrainConfig:
    model_size: str = "small"
    num_classes: int = 0
    class_drop_prob: float = 0.1
    dataset_preset: str = "cartoon"
    image_size: int = 256
    max_images: int = 0
    vae_id: str = "madebyollin/sdxl-vae-fp16-fix"
    vae_scaling_factor: float = 0.13025
    latent_channels: int = 4
    batch_size: int = 0
    gradient_accumulation_steps: int = 1
    learning_rate: float = 1e-4
    weight_decay: float = 0.01
    max_grad_norm: float = 2.0
    num_epochs: int = 100
    warmup_steps: int = 500
    ema_decay: float = 0.9999
    mixed_precision: bool = True
    gradient_checkpointing: bool = True
    min_timestep: float = 0.001
    max_timestep: float = 0.999
    output_dir: str = "./outputs"
    save_every_n_steps: int = 2000
    sample_every_n_steps: int = 500
    log_every_n_steps: int = 25
    num_sample_steps: int = 50
    cfg_scale: float = 2.0
    num_samples: int = 4
    seed: int = 42
    num_workers: int = 2
    compile_model: bool = False
    push_to_hub: bool = False
    hub_model_id: str = ""


def get_model_config(size, num_classes=0, class_drop_prob=0.1):
    configs = {
        "small": dict(embed_dim=512, depth=12, spatial_kernel=7, scan_kernel=31,
                      expand_ratio=2.0, mlp_ratio=3.0),
        "base": dict(embed_dim=640, depth=18, spatial_kernel=7, scan_kernel=31,
                     expand_ratio=2.0, mlp_ratio=4.0),
        "large": dict(embed_dim=768, depth=24, spatial_kernel=7, scan_kernel=31,
                      expand_ratio=2.5, mlp_ratio=4.0),
    }
    cfg = configs[size]
    cfg["num_classes"] = num_classes
    cfg["class_drop_prob"] = class_drop_prob
    cfg["use_zigzag"] = True
    return cfg


class CachedLatentDataset(Dataset):
    def __init__(self, cache_path):
        data = torch.load(cache_path, map_location="cpu", weights_only=True)
        self.latents = data["latents"]
        self.labels = data.get("labels", None)
        print(f"Loaded {len(self.latents)} cached latents: {self.latents.shape}")
        if self.labels is not None and (self.labels >= 0).any():
            print(f"  {self.labels[self.labels >= 0].unique().shape[0]} classes")
    def __len__(self): return len(self.latents)
    def __getitem__(self, idx):
        return self.latents[idx], (self.labels[idx] if self.labels is not None else -1)


def precache_latents(config, cache_path=None):
    if cache_path is None:
        cache_path = os.path.join(config.output_dir, "cached_latents.pt")
    if os.path.exists(cache_path):
        print(f"Cache exists: {cache_path}")
        d = torch.load(cache_path, map_location="cpu", weights_only=True)
        print(f"  {d['latents'].shape[0]} latents {d['latents'].shape[1:]}")
        return cache_path

    os.makedirs(os.path.dirname(cache_path) if os.path.dirname(cache_path) else ".", exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    print(f"Loading VAE: {config.vae_id}...")
    from diffusers import AutoencoderKL
    vae = AutoencoderKL.from_pretrained(config.vae_id, torch_dtype=torch.float16).to(device).eval()
    for p in vae.parameters(): p.requires_grad_(False)

    preset = DATASET_PRESETS[config.dataset_preset]
    print(f"Dataset: {preset['name']}")
    from datasets import load_dataset
    from torchvision import transforms

    ds_kwargs = {"split": "train"}
    if preset["config"]: ds_kwargs["name"] = preset["config"]
    dataset = load_dataset(preset["name"], **ds_kwargs)

    transform = transforms.Compose([
        transforms.Resize(config.image_size, interpolation=transforms.InterpolationMode.LANCZOS),
        transforms.CenterCrop(config.image_size), transforms.ToTensor(),
    ])

    if config.max_images > 0:
        max_imgs = config.max_images
    elif preset.get("max_default", 0) > 0:
        max_imgs = preset["max_default"]
        print(f"  Auto-capping to {max_imgs} images (set max_images to override)")
    else:
        max_imgs = len(dataset)
    max_imgs = min(max_imgs, len(dataset))

    encode_bs = 64 if config.image_size <= 256 else 32
    print(f"  Encoding {max_imgs} images (batch={encode_bs})...")

    img_col, lbl_col = preset["image_column"], preset["label_column"]
    style_to_id = {}
    all_latents, all_labels = [], []
    batch_px, batch_lb = [], []
    count = 0
    t0 = time.time()

    for item in dataset:
        if count >= max_imgs: break
        img = item[img_col]
        if img.mode != "RGB": img = img.convert("RGB")
        batch_px.append(transform(img))
        if lbl_col and lbl_col in item:
            raw = item[lbl_col]
            if isinstance(raw, str):
                if raw not in style_to_id: style_to_id[raw] = len(style_to_id)
                batch_lb.append(style_to_id[raw])
            elif isinstance(raw, int): batch_lb.append(raw)
            else: batch_lb.append(-1)
        else: batch_lb.append(-1)
        count += 1
        if len(batch_px) >= encode_bs:
            with torch.no_grad():
                px = torch.stack(batch_px).to(device, dtype=torch.float16) * 2 - 1
                lat = vae.encode(px).latent_dist.sample() * config.vae_scaling_factor
                all_latents.append(lat.cpu().float())
            all_labels.extend(batch_lb); batch_px, batch_lb = [], []
            elapsed = time.time() - t0
            speed = count / elapsed
            eta = (max_imgs - count) / speed if speed > 0 else 0
            if count % (encode_bs * 4) == 0:
                print(f"  {count}/{max_imgs} | {speed:.0f} img/s | ETA {_fmt_time(eta)}")

    if batch_px:
        with torch.no_grad():
            px = torch.stack(batch_px).to(device, dtype=torch.float16) * 2 - 1
            lat = vae.encode(px).latent_dist.sample() * config.vae_scaling_factor
            all_latents.append(lat.cpu().float())
        all_labels.extend(batch_lb)

    all_latents = torch.cat(all_latents, dim=0)
    all_labels = torch.tensor(all_labels, dtype=torch.long)
    save_data = {"latents": all_latents, "labels": all_labels}
    if style_to_id:
        save_data["style_to_id"] = style_to_id
        print(f"  {len(style_to_id)} style classes")
    torch.save(save_data, cache_path)
    mb = os.path.getsize(cache_path) / 1024**2
    print(f"Cached {count} latents -> {cache_path} ({mb:.0f}MB, {_fmt_time(time.time()-t0)})")
    del vae
    if torch.cuda.is_available(): torch.cuda.empty_cache()
    return cache_path


class EMAModel:
    def __init__(self, model, decay=0.9999):
        self.decay = decay
        self.shadow = {n: p.clone().detach() for n, p in model.named_parameters() if p.requires_grad}
    @torch.no_grad()
    def update(self, model):
        for n, p in model.named_parameters():
            if p.requires_grad and n in self.shadow:
                self.shadow[n].mul_(self.decay).add_(p.data, alpha=1 - self.decay)
    def apply(self, model):
        self.backup = {n: p.data.clone() for n, p in model.named_parameters() if p.requires_grad}
        for n, p in model.named_parameters():
            if p.requires_grad and n in self.shadow: p.data.copy_(self.shadow[n])
    def restore(self, model):
        for n, p in model.named_parameters():
            if p.requires_grad and n in self.backup: p.data.copy_(self.backup[n])
        self.backup = {}


class FlowMatchingScheduler:
    def __init__(self, min_t=0.001, max_t=0.999):
        self.min_t, self.max_t = min_t, max_t
    def sample_timesteps(self, bs, dev):
        return torch.rand(bs, device=dev) * (self.max_t - self.min_t) + self.min_t
    def add_noise(self, x0, noise, t):
        t = t.view(-1, 1, 1, 1); return (1 - t) * x0 + t * noise
    def get_velocity_target(self, x0, noise):
        return noise - x0
    @torch.no_grad()
    def sample(self, model, shape, dev, num_steps=50, labels=None, cfg=1.0):
        model.eval(); x = torch.randn(shape, device=dev); dt = 1.0 / num_steps
        for tv in torch.linspace(1.0, dt, num_steps, device=dev):
            t = torch.full((shape[0],), tv.item(), device=dev)
            with torch.amp.autocast("cuda"):
                if cfg > 1.0 and labels is not None:
                    vc = model(x, t, labels); vu = model(x, t, torch.zeros_like(labels))
                    v = vu + cfg * (vc - vu)
                else: v = model(x, t, labels)
            x = x - dt * v.float()
        return x


def cosine_schedule(opt, warmup, total):
    def lr(s):
        if s < warmup: return s / max(1, warmup)
        return max(0, 0.5 * (1 + math.cos(math.pi * (s - warmup) / max(1, total - warmup))))
    return torch.optim.lr_scheduler.LambdaLR(opt, lr)


def train(config):
    from model import LiquidGen
    torch.manual_seed(config.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    gpu_mem = 0
    if torch.cuda.is_available():
        gpu_mem = torch.cuda.get_device_properties(0).total_mem / 1024**3
        print(f"GPU: {torch.cuda.get_device_name(0)} ({gpu_mem:.1f} GB)")

    if config.batch_size <= 0:
        config.batch_size = auto_batch_size(config.model_size, config.image_size, gpu_mem) if gpu_mem > 0 else 4
        print(f"Auto batch: {config.batch_size}")

    os.makedirs(config.output_dir, exist_ok=True)
    os.makedirs(f"{config.output_dir}/samples", exist_ok=True)
    os.makedirs(f"{config.output_dir}/checkpoints", exist_ok=True)

    cache_path = precache_latents(config)
    train_ds = CachedLatentDataset(cache_path)
    train_dl = DataLoader(train_ds, batch_size=config.batch_size, shuffle=True,
                          num_workers=config.num_workers, pin_memory=True, drop_last=True)

    mcfg = get_model_config(config.model_size, config.num_classes, config.class_drop_prob)
    mcfg["in_channels"] = config.latent_channels
    model = LiquidGen(**mcfg).to(device)
    if config.gradient_checkpointing:
        model.enable_gradient_checkpointing()
    print(f"LiquidGen-{config.model_size}: {model.count_params()/1e6:.1f}M (ckpt={'ON' if config.gradient_checkpointing else 'OFF'})")

    opt = torch.optim.AdamW(model.parameters(), lr=config.learning_rate,
                            weight_decay=config.weight_decay, betas=(0.9, 0.999))
    total_steps = len(train_dl) * config.num_epochs // config.gradient_accumulation_steps
    sched = cosine_schedule(opt, config.warmup_steps, total_steps)
    ema = EMAModel(model, config.ema_decay)
    scaler = GradScaler("cuda", enabled=config.mixed_precision and torch.cuda.is_available())
    fm = FlowMatchingScheduler(config.min_timestep, config.max_timestep)
    lat_size = config.image_size // 8
    print(f"Steps: {total_steps} | Batch: {config.batch_size} | Epochs: {config.num_epochs}")

    gs = 0; la = 0.0; vae = None; vae_loaded = False
    print(f"\nTraining!\n")
    t_start = time.time()

    for epoch in range(config.num_epochs):
        model.train(); et = time.time()
        for bi, (lats, lbls) in enumerate(train_dl):
            lats = lats.to(device)
            lbls = lbls.to(device) if config.num_classes > 0 else None
            t = fm.sample_timesteps(lats.shape[0], device)
            noise = torch.randn_like(lats)
            xt = fm.add_noise(lats, noise, t)
            vtgt = fm.get_velocity_target(lats, noise)
            with autocast("cuda", enabled=config.mixed_precision and torch.cuda.is_available()):
                loss = F.mse_loss(model(xt, t, lbls), vtgt) / config.gradient_accumulation_steps
            scaler.scale(loss).backward(); la += loss.item()
            if (bi + 1) % config.gradient_accumulation_steps == 0:
                scaler.unscale_(opt)
                gn = torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)
                scaler.step(opt); scaler.update(); opt.zero_grad(); sched.step()
                ema.update(model); gs += 1
                if gs % config.log_every_n_steps == 0:
                    al = la / config.log_every_n_steps
                    elapsed = time.time() - t_start
                    sps = gs / max(elapsed, 1)
                    remaining = (total_steps - gs) / sps if sps > 0 else 0
                    vram = torch.cuda.memory_allocated()/1024**3 if torch.cuda.is_available() else 0
                    pct = gs / total_steps * 100
                    print(f"step={gs:>6d}/{total_steps} ({pct:.0f}%) | ep={epoch} | "
                          f"loss={al:.4f} | gn={gn:.2f} | lr={opt.param_groups[0]['lr']:.2e} | "
                          f"vram={vram:.1f}G | {sps:.1f} st/s | ETA {_fmt_time(remaining)}")
                    la = 0.0
                    if math.isnan(al) or al > 50: print("Diverged!"); return
                if gs % config.sample_every_n_steps == 0:
                    if not vae_loaded:
                        from diffusers import AutoencoderKL
                        vae = AutoencoderKL.from_pretrained(config.vae_id, torch_dtype=torch.float16).to(device).eval()
                        for p in vae.parameters(): p.requires_grad_(False)
                        vae_loaded = True
                    ema.apply(model); model.eval()
                    sl = torch.randint(0, max(1, config.num_classes), (config.num_samples,), device=device) if config.num_classes > 0 else None
                    samp = fm.sample(model, (config.num_samples, config.latent_channels, lat_size, lat_size),
                                     device, config.num_sample_steps, sl, config.cfg_scale)
                    with torch.no_grad():
                        imgs = ((vae.decode(samp.half() / config.vae_scaling_factor).sample + 1) / 2).clamp(0, 1).float()
                    from torchvision.utils import save_image
                    save_image(imgs, f"{config.output_dir}/samples/step_{gs:07d}.png", nrow=2)
                    print(f"  Saved samples"); ema.restore(model); model.train()
                if gs % config.save_every_n_steps == 0:
                    torch.save({"model": model.state_dict(), "ema": ema.shadow,
                                "optimizer": opt.state_dict(), "step": gs, "model_config": mcfg},
                               f"{config.output_dir}/checkpoints/step_{gs:07d}.pt")
        ep_time = time.time() - et
        ep_eta = ep_time * (config.num_epochs - epoch - 1)
        print(f"Epoch {epoch}/{config.num_epochs} done | {_fmt_time(ep_time)} | ETA {_fmt_time(ep_eta)}\n")

    final = f"{config.output_dir}/checkpoints/final.pt"
    torch.save({"model": model.state_dict(), "ema": ema.shadow, "model_config": mcfg, "step": gs}, final)
    total_time = time.time() - t_start
    print(f"\nDone! {gs} steps in {_fmt_time(total_time)} -> {final}")