| """ |
| 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}") |
|
|
