Add training pipeline

train.py

@@ -0,0 +1,378 @@
+"""
+LiquidGen Training Pipeline
+
+Flow Matching training objective (velocity prediction):
+- Forward: x_t = (1 - t) * x_0 + t * ε   (linear interpolation)
+- Target: v = ε - x_0                      (velocity)
+- Loss: MSE(model(x_t, t), v)
+
+At inference: solve ODE from t=1 (noise) to t=0 (clean) using Euler steps.
+"""
+
+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 pathlib import Path
+from typing import Optional, Dict, Any
+from dataclasses import dataclass, field, asdict
+
+
+@dataclass
+class TrainConfig:
+    """Training configuration with sensible defaults for Colab free tier."""
+    # Model
+    model_size: str = "small"
+    num_classes: int = 0
+    class_drop_prob: float = 0.1
+    
+    # Data
+    image_size: int = 256
+    dataset_name: str = "huggan/wikiart"
+    dataset_config: str = ""
+    image_column: str = "image"
+    label_column: str = ""
+    
+    # VAE
+    vae_id: str = "black-forest-labs/FLUX.1-schnell"
+    vae_subfolder: str = "vae"
+    vae_dtype: str = "float16"
+    vae_scaling_factor: float = 0.3611
+    vae_shift_factor: float = 0.1159
+    
+    # Training
+    batch_size: int = 8
+    gradient_accumulation_steps: int = 4
+    learning_rate: float = 1e-4
+    weight_decay: float = 0.01
+    max_grad_norm: float = 2.0
+    num_epochs: int = 100
+    warmup_steps: int = 1000
+    ema_decay: float = 0.9999
+    mixed_precision: bool = True
+    
+    # Flow matching
+    min_timestep: float = 0.001
+    max_timestep: float = 0.999
+    
+    # Saving
+    output_dir: str = "./outputs"
+    save_every_n_steps: int = 5000
+    sample_every_n_steps: int = 1000
+    log_every_n_steps: int = 50
+    
+    # Sampling
+    num_sample_steps: int = 50
+    cfg_scale: float = 1.5
+    num_samples: int = 4
+    
+    # System
+    seed: int = 42
+    num_workers: int = 2
+    pin_memory: bool = True
+    compile_model: bool = False
+    
+    # Hub
+    push_to_hub: bool = False
+    hub_model_id: str = ""
+
+
+def get_model_config(size: str, num_classes: int = 0, class_drop_prob: float = 0.1) -> dict:
+    """Get model kwargs for a given size preset."""
+    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 EMAModel:
+    """Exponential Moving Average of model parameters."""
+    
+    def __init__(self, model: nn.Module, decay: float = 0.9999):
+        self.decay = decay
+        self.shadow = {name: p.clone().detach() for name, p in model.named_parameters() if p.requires_grad}
+    
+    @torch.no_grad()
+    def update(self, model: nn.Module):
+        for name, p in model.named_parameters():
+            if p.requires_grad and name in self.shadow:
+                self.shadow[name].mul_(self.decay).add_(p.data, alpha=1 - self.decay)
+    
+    def apply(self, model: nn.Module):
+        self.backup = {name: p.data.clone() for name, p in model.named_parameters() if p.requires_grad}
+        for name, p in model.named_parameters():
+            if p.requires_grad and name in self.shadow:
+                p.data.copy_(self.shadow[name])
+    
+    def restore(self, model: nn.Module):
+        for name, p in model.named_parameters():
+            if p.requires_grad and name in self.backup:
+                p.data.copy_(self.backup[name])
+        self.backup = {}
+    
+    def state_dict(self):
+        return self.shadow
+    
+    def load_state_dict(self, state_dict):
+        self.shadow = state_dict
+
+
+class FlowMatchingScheduler:
+    """
+    Flow Matching scheduler for training and sampling.
+    
+    Training: x_t = (1-t)*x_0 + t*ε, v_target = ε - x_0
+    Sampling: Euler ODE from t=1 (noise) to t=0 (clean)
+    """
+    
+    def __init__(self, min_t: float = 0.001, max_t: float = 0.999):
+        self.min_t = min_t
+        self.max_t = max_t
+    
+    def sample_timesteps(self, batch_size: int, device: torch.device) -> torch.Tensor:
+        return torch.rand(batch_size, device=device) * (self.max_t - self.min_t) + self.min_t
+    
+    def add_noise(self, x0: torch.Tensor, noise: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        t_expand = t.view(-1, 1, 1, 1)
+        return (1 - t_expand) * x0 + t_expand * noise
+    
+    def get_velocity_target(self, x0: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
+        return noise - x0
+    
+    @torch.no_grad()
+    def sample(
+        self, model: nn.Module, shape: tuple, device: torch.device,
+        num_steps: int = 50, class_labels: Optional[torch.Tensor] = None,
+        cfg_scale: float = 1.0, dtype: torch.dtype = torch.float32,
+    ) -> torch.Tensor:
+        model.eval()
+        x = torch.randn(shape, device=device, dtype=dtype)
+        dt = 1.0 / num_steps
+        times = torch.linspace(1.0, dt, num_steps, device=device)
+        
+        for t_val in times:
+            t = torch.full((shape[0],), t_val.item(), device=device, dtype=dtype)
+            
+            if cfg_scale > 1.0 and class_labels is not None:
+                with torch.amp.autocast('cuda', enabled=(dtype != torch.float32)):
+                    v_cond = model(x, t, class_labels)
+                    v_uncond = model(x, t, torch.zeros_like(class_labels))
+                v = v_uncond + cfg_scale * (v_cond - v_uncond)
+            else:
+                with torch.amp.autocast('cuda', enabled=(dtype != torch.float32)):
+                    v = model(x, t, class_labels)
+            
+            x = x - dt * v
+        
+        return x
+
+
+def get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps):
+    """Cosine LR schedule with linear warmup."""
+    def lr_lambda(current_step):
+        if current_step < warmup_steps:
+            return float(current_step) / float(max(1, warmup_steps))
+        progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
+        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+@torch.no_grad()
+def encode_images_with_vae(images, vae, scaling_factor, shift_factor):
+    """Encode pixel images to VAE latents."""
+    images = images * 2.0 - 1.0
+    latents = vae.encode(images).latent_dist.sample()
+    latents = (latents - shift_factor) * scaling_factor
+    return latents
+
+
+@torch.no_grad()
+def decode_latents_with_vae(latents, vae, scaling_factor, shift_factor):
+    """Decode VAE latents to pixel images."""
+    latents = latents / scaling_factor + shift_factor
+    images = vae.decode(latents).sample
+    images = (images + 1.0) / 2.0
+    return images.clamp(0, 1)
+
+
+def train(config: TrainConfig):
+    """Main training loop."""
+    from model import LiquidGen
+    
+    torch.manual_seed(config.seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+    
+    os.makedirs(config.output_dir, exist_ok=True)
+    os.makedirs(os.path.join(config.output_dir, "samples"), exist_ok=True)
+    os.makedirs(os.path.join(config.output_dir, "checkpoints"), exist_ok=True)
+    
+    with open(os.path.join(config.output_dir, "config.json"), "w") as f:
+        json.dump(asdict(config), f, indent=2)
+    
+    # Load VAE
+    print("Loading VAE...")
+    from diffusers import AutoencoderKL
+    vae_dtype = torch.float16 if config.vae_dtype == "float16" else torch.bfloat16
+    vae = AutoencoderKL.from_pretrained(
+        config.vae_id, subfolder=config.vae_subfolder, torch_dtype=vae_dtype
+    ).to(device).eval()
+    for p in vae.parameters():
+        p.requires_grad_(False)
+    
+    # Load Dataset
+    print(f"Loading dataset: {config.dataset_name}")
+    from datasets import load_dataset
+    from torchvision import transforms
+    
+    ds_kwargs = {}
+    if config.dataset_config:
+        ds_kwargs["name"] = config.dataset_config
+    dataset = load_dataset(config.dataset_name, split="train", **ds_kwargs)
+    
+    transform = transforms.Compose([
+        transforms.Resize(config.image_size, interpolation=transforms.InterpolationMode.LANCZOS),
+        transforms.CenterCrop(config.image_size),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+    ])
+    
+    class ImageDataset(Dataset):
+        def __init__(self, hf_dataset, transform, image_col, label_col=""):
+            self.dataset = hf_dataset
+            self.transform = transform
+            self.image_col = image_col
+            self.label_col = label_col
+        
+        def __len__(self):
+            return len(self.dataset)
+        
+        def __getitem__(self, idx):
+            item = self.dataset[idx]
+            img = item[self.image_col]
+            if img.mode != "RGB":
+                img = img.convert("RGB")
+            img = self.transform(img)
+            label = item[self.label_col] if self.label_col and self.label_col in item else -1
+            return img, label
+    
+    train_dataset = ImageDataset(dataset, transform, config.image_column, config.label_column)
+    train_loader = DataLoader(
+        train_dataset, batch_size=config.batch_size, shuffle=True,
+        num_workers=config.num_workers, pin_memory=config.pin_memory, drop_last=True,
+    )
+    
+    # Create Model
+    model_kwargs = get_model_config(config.model_size, config.num_classes, config.class_drop_prob)
+    model = LiquidGen(**model_kwargs).to(device)
+    print(f"LiquidGen-{config.model_size}: {model.count_params() / 1e6:.1f}M params")
+    
+    if config.compile_model and hasattr(torch, "compile"):
+        model = torch.compile(model)
+    
+    optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate,
+                                   weight_decay=config.weight_decay, betas=(0.9, 0.999))
+    total_steps = len(train_loader) * config.num_epochs // config.gradient_accumulation_steps
+    scheduler = get_cosine_schedule_with_warmup(optimizer, config.warmup_steps, total_steps)
+    ema = EMAModel(model, decay=config.ema_decay)
+    scaler = GradScaler('cuda', enabled=config.mixed_precision)
+    fm = FlowMatchingScheduler(min_t=config.min_timestep, max_t=config.max_timestep)
+    
+    print(f"\nTraining: {total_steps} steps, effective batch {config.batch_size * config.gradient_accumulation_steps}")
+    
+    global_step = 0
+    loss_accum = 0.0
+    
+    for epoch in range(config.num_epochs):
+        model.train()
+        t_start = time.time()
+        
+        for batch_idx, (images, labels) in enumerate(train_loader):
+            images = images.to(device)
+            labels = labels.to(device) if config.num_classes > 0 else None
+            
+            with torch.no_grad():
+                latents = encode_images_with_vae(
+                    images.to(vae_dtype), vae, config.vae_scaling_factor, config.vae_shift_factor
+                ).float()
+            
+            t = fm.sample_timesteps(latents.shape[0], device)
+            noise = torch.randn_like(latents)
+            x_t = fm.add_noise(latents, noise, t)
+            v_target = fm.get_velocity_target(latents, noise)
+            
+            with autocast('cuda', enabled=config.mixed_precision):
+                v_pred = model(x_t, t, labels)
+                loss = F.mse_loss(v_pred, v_target) / config.gradient_accumulation_steps
+            
+            scaler.scale(loss).backward()
+            loss_accum += loss.item()
+            
+            if (batch_idx + 1) % config.gradient_accumulation_steps == 0:
+                scaler.unscale_(optimizer)
+                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad()
+                scheduler.step()
+                ema.update(model)
+                global_step += 1
+                
+                if global_step % config.log_every_n_steps == 0:
+                    avg_loss = loss_accum / config.log_every_n_steps
+                    lr = optimizer.param_groups[0]["lr"]
+                    print(f"step={global_step} | epoch={epoch} | loss={avg_loss:.4f} | "
+                          f"grad_norm={grad_norm:.2f} | lr={lr:.2e}")
+                    loss_accum = 0.0
+                    
+                    if math.isnan(avg_loss) or avg_loss > 100:
+                        print("⚠️ Training diverged!")
+                        return
+                
+                if global_step % config.sample_every_n_steps == 0:
+                    ema.apply(model)
+                    model.eval()
+                    latent_size = config.image_size // 8
+                    sample_labels = None
+                    if config.num_classes > 0:
+                        sample_labels = torch.randint(0, config.num_classes, (config.num_samples,), device=device)
+                    sampled = fm.sample(model, (config.num_samples, 16, latent_size, latent_size),
+                                       device, config.num_sample_steps, sample_labels, config.cfg_scale)
+                    sample_imgs = decode_latents_with_vae(sampled.to(vae_dtype), vae,
+                                                          config.vae_scaling_factor, config.vae_shift_factor).float()
+                    from torchvision.utils import save_image
+                    save_image(sample_imgs, os.path.join(config.output_dir, "samples", f"step_{global_step:07d}.png"), nrow=2)
+                    ema.restore(model)
+                    model.train()
+                
+                if global_step % config.save_every_n_steps == 0:
+                    torch.save({
+                        "model": model.state_dict(), "ema": ema.state_dict(),
+                        "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(),
+                        "global_step": global_step, "epoch": epoch, "config": asdict(config),
+                    }, os.path.join(config.output_dir, "checkpoints", f"step_{global_step:07d}.pt"))
+        
+        print(f"Epoch {epoch} complete | time={time.time()-t_start:.0f}s")
+    
+    torch.save({"model": model.state_dict(), "ema": ema.state_dict(), "config": asdict(config),
+                "global_step": global_step}, os.path.join(config.output_dir, "checkpoints", "final.pt"))
+    print(f"Training complete! Final model saved.")
+
+
+if __name__ == "__main__":
+    config = TrainConfig(model_size="small", image_size=256, batch_size=4, num_epochs=2)
+    train(config)