""" LiquidFlow Trainer — Complete training pipeline. Usage: python train.py --dataset cifar10 --image_size 128 --variant small --batch_size 32 --epochs 100 Features: - Automatic VAE loading (TAESD by default) - Physics-informed regularization - Mixed precision training (AMP) - Checkpoint saving - Sample generation during training - Colab/Kaggle compatible (T4 GPU, 15GB VRAM) Requirements: pip install torch torchvision diffusers tqdm pillow numpy """ import os import sys import math import argparse import json from datetime import datetime from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms from torchvision.utils import save_image import numpy as np from tqdm import tqdm # Add parent to path sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from liquid_flow.generator import LiquidFlowGenerator, create_liquidflow from liquid_flow.vae_wrapper import TAESDWrapper def get_dataloader(dataset_name, image_size, batch_size, data_dir='./data'): """Get training dataloader for common datasets.""" transform = transforms.Compose([ transforms.Resize((image_size, image_size)), transforms.ToTensor(), transforms.Normalize([0.5], [0.5]), # [-1, 1] ]) if dataset_name == 'cifar10': dataset = datasets.CIFAR10( root=data_dir, train=True, download=True, transform=transform ) elif dataset_name == 'cifar100': dataset = datasets.CIFAR100( root=data_dir, train=True, download=True, transform=transform ) elif dataset_name == 'stl10': dataset = datasets.STL10( root=data_dir, split='train', download=True, transform=transform ) elif dataset_name == 'celeba': dataset = datasets.CelebA( root=data_dir, split='train', download=True, transform=transform ) elif dataset_name == 'lsun': dataset = datasets.LSUN( root=data_dir, classes='bedroom_train', transform=transform ) elif dataset_name == 'imagenet': transform = transforms.Compose([ transforms.Resize((image_size, image_size)), transforms.RandomCrop(image_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.5], [0.5]), ]) dataset = datasets.ImageFolder( root=f'{data_dir}/imagenet/train', transform=transform ) else: raise ValueError(f"Unknown dataset: {dataset_name}") dataloader = DataLoader( dataset, batch_size=batch_size, shuffle=True, num_workers=min(4, os.cpu_count() or 1), pin_memory=True, drop_last=True, ) return dataloader def train(args): """Main training loop.""" device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Using device: {device}") # Create output directory os.makedirs(args.output_dir, exist_ok=True) os.makedirs(os.path.join(args.output_dir, 'samples'), exist_ok=True) os.makedirs(os.path.join(args.output_dir, 'checkpoints'), exist_ok=True) # Load VAE print("Loading VAE...") vae = TAESDWrapper.load(device) print(f"VAE loaded. Latent size: {args.image_size // 8}x{args.image_size // 8}") # Create model print(f"Creating LiquidFlow model (variant={args.variant})...") model = create_liquidflow( variant=args.variant, image_size=args.image_size, ) model = model.to(device) n_params = model.count_parameters() print(f"Model parameters: {n_params:,} (~{n_params/1e6:.1f}M)") # Calculate memory estimate latent_h = latent_w = args.image_size // 8 mem_per_sample = latent_h * latent_w * 4 * 4 / (1024**2) # in MB print(f"Estimated memory per sample: {mem_per_sample:.1f} MB") print(f"Estimated batch memory: {mem_per_sample * args.batch_size:.1f} MB") # Dataset print(f"Loading dataset: {args.dataset}") dataloader = get_dataloader(args.dataset, args.image_size, args.batch_size, args.data_dir) print(f"Dataset size: {len(dataloader.dataset)} images, {len(dataloader)} batches") # Optimizer (AdamW, following DiT/DiMSUM convention) optimizer = torch.optim.AdamW( model.parameters(), lr=args.lr, betas=(0.9, 0.999), weight_decay=args.weight_decay, ) # Learning rate scheduler if args.lr_schedule == 'cosine': scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=args.epochs * len(dataloader) ) elif args.lr_schedule == 'cosine_restart': scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts( optimizer, T_0=args.epochs * len(dataloader) // 3, ) else: scheduler = None # AMP use_amp = args.amp and device.type == 'cuda' scaler = torch.cuda.amp.GradScaler() if use_amp else None # Fixed noise for sample generation (track progress) sample_noise = torch.randn(16, 4, args.image_size // 8, args.image_size // 8, device=device) # Training state global_step = 0 best_loss = float('inf') print(f"\n{'='*60}") print(f"Starting training: {args.epochs} epochs, {args.batch_size} batch size") print(f"LR: {args.lr}, Weight Decay: {args.weight_decay}") print(f"AMP: {use_amp}, LR Schedule: {args.lr_schedule}") print(f"{'='*60}\n") for epoch in range(args.epochs): model.train() epoch_losses = {'total': 0, 'diffusion': 0, 'physics': 0} pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{args.epochs}") for batch_idx, (images, _) in enumerate(pbar): images = images.to(device) # Encode to latent space with torch.no_grad(): latents = TAESDWrapper.encode(vae, images) # Training step loss_dict = model.training_step(latents, optimizer, scaler, use_amp) # Update scheduler if scheduler is not None: scheduler.step() # Track losses for k in epoch_losses: epoch_losses[k] += loss_dict.get(k, 0) global_step += 1 # Update progress bar pbar.set_postfix({ 'loss': f"{loss_dict.get('total', 0):.4f}", 'diff': f"{loss_dict.get('diffusion', 0):.4f}", 'phys': f"{loss_dict.get('physics', 0):.4f}", 'lr': f"{optimizer.param_groups[0]['lr']:.2e}", }) # Epoch summary n_batches = len(dataloader) avg_losses = {k: v / n_batches for k, v in epoch_losses.items()} print(f"\nEpoch {epoch+1} Summary:") print(f" Total Loss: {avg_losses['total']:.4f}") print(f" Diffusion Loss: {avg_losses['diffusion']:.4f}") print(f" Physics Loss: {avg_losses['physics']:.4f}") # Generate samples if (epoch + 1) % args.sample_every == 0 or epoch == args.epochs - 1: print(f"Generating samples...") model.eval() with torch.no_grad(): # DDIM sampling latents_gen = model.sample( batch_size=16, steps=args.sample_steps, ddim=True, progress=False, ) images_gen = TAESDWrapper.decode(vae, latents_gen) # Also generate from fixed noise for tracking t_fixed = torch.full((16,), 0, device=device, dtype=torch.long) # Quick DDIM from fixed noise x_fixed = sample_noise.clone() skip = 1000 // args.sample_steps for i in reversed(range(0, 1000, skip)): t = torch.full((16,), i, device=device, dtype=torch.long) noise_pred = model(x_fixed, t) alpha_bar = model.alphas_cumprod[i] alpha_bar_prev = model.alphas_cumprod[i - skip] if i >= skip else torch.tensor(1.0, device=device) x0_pred = (x_fixed - torch.sqrt(1 - alpha_bar) * noise_pred) / torch.sqrt(alpha_bar) x0_pred = torch.clamp(x0_pred, -1, 1) x_fixed = torch.sqrt(alpha_bar_prev) * x0_pred + torch.sqrt(1 - alpha_bar_prev) * torch.randn_like(x_fixed) images_fixed = TAESDWrapper.decode(vae, x_fixed) # Save samples sample_path = os.path.join(args.output_dir, 'samples', f'epoch_{epoch+1:03d}.png') save_image(images_gen, sample_path, nrow=4, normalize=True, value_range=(-1, 1)) fixed_path = os.path.join(args.output_dir, 'samples', f'fixed_{epoch+1:03d}.png') save_image(images_fixed, fixed_path, nrow=4, normalize=True, value_range=(-1, 1)) print(f" Samples saved to {sample_path}") # Save checkpoint if (epoch + 1) % args.save_every == 0 or epoch == args.epochs - 1: checkpoint_path = os.path.join(args.output_dir, 'checkpoints', f'epoch_{epoch+1:03d}.pt') torch.save({ 'epoch': epoch + 1, 'global_step': global_step, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': avg_losses['total'], 'args': vars(args), }, checkpoint_path) print(f" Checkpoint saved to {checkpoint_path}") # Save best model if avg_losses['total'] < best_loss: best_loss = avg_losses['total'] best_path = os.path.join(args.output_dir, 'checkpoints', 'best_model.pt') torch.save(model.state_dict(), best_path) print(f" Best model saved (loss={best_loss:.4f})") print() print(f"\n{'='*60}") print(f"Training complete!") print(f"Best loss: {best_loss:.4f}") print(f"Model saved to: {args.output_dir}/checkpoints/") print(f"{'='*60}") return model def main(): parser = argparse.ArgumentParser(description='LiquidFlow Generator Training') # Dataset parser.add_argument('--dataset', type=str, default='cifar10', choices=['cifar10', 'cifar100', 'stl10', 'celeba', 'lsun', 'imagenet'], help='Training dataset') parser.add_argument('--data_dir', type=str, default='./data', help='Data directory') parser.add_argument('--image_size', type=int, default=128, choices=[64, 128, 256, 512], help='Image size (will be VAE-encoded)') # Model parser.add_argument('--variant', type=str, default='small', choices=['tiny', 'small', 'base'], help='Model size variant') # Training parser.add_argument('--batch_size', type=int, default=32, help='Batch size') parser.add_argument('--epochs', type=int, default=100, help='Number of epochs') parser.add_argument('--lr', type=float, default=2e-4, help='Learning rate') parser.add_argument('--weight_decay', type=float, default=1e-4, help='Weight decay') parser.add_argument('--lr_schedule', type=str, default='cosine', choices=['cosine', 'cosine_restart', 'none'], help='LR schedule') parser.add_argument('--amp', action='store_true', default=True, help='Use automatic mixed precision') # Generation parser.add_argument('--sample_every', type=int, default=5, help='Generate samples every N epochs') parser.add_argument('--sample_steps', type=int, default=50, help='DDIM sampling steps') # IO parser.add_argument('--output_dir', type=str, default='./outputs', help='Output directory') parser.add_argument('--save_every', type=int, default=10, help='Save checkpoint every N epochs') args = parser.parse_args() train(args) if __name__ == '__main__': main()