|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import csv
|
| import json
|
| import math
|
| import os
|
| import time
|
| from dataclasses import dataclass, asdict, field
|
| from pathlib import Path
|
|
|
| import torch
|
| from safetensors import safe_open
|
| from safetensors.torch import save_file, load_file
|
| from torch.utils.data import DataLoader
|
|
|
| from bora_dataset import BoraLatents
|
| from bora_dit import DiTConfig, BoraDiT, PRESETS, build, count_params, config_dict, flow_loss, shift_t, sample_video
|
|
|
|
|
| @dataclass
|
| class TrainConfig:
|
| data_root: str = "/content/bora_data"
|
| out_dir: str = "/content/bora_runs/v1"
|
| text_embeds: str = "text_t5v11base_32.safetensors"
|
| preset: str = "v1"
|
| model: dict = field(default_factory=dict)
|
|
|
| batch_size: int = 32
|
| num_workers: int = 4
|
| block_len: int = 4
|
| overlap: int = 1
|
| val_mod: int = 100
|
| primary_p: float = 0.5
|
|
|
| p_cond: float = 0.5
|
| p_uncond: float = 0.1
|
| time_shift: float = 3.0
|
|
|
| lr: float = 1e-4
|
| warmup: int = 1000
|
| max_steps: int = 50000
|
| min_lr_frac: float = 0.1
|
| weight_decay: float = 0.0
|
| betas: tuple = (0.9, 0.95)
|
| grad_clip: float = 1.0
|
| ema: float = 0.9999
|
| bf16: bool = True
|
| grad_checkpoint: bool = False
|
| compile: bool = False
|
| seed: int = 0
|
|
|
| log_every: int = 50
|
| val_every: int = 1000
|
| val_batches: int = 8
|
| val_t: tuple = (0.25, 0.5, 0.75, 0.9)
|
| sample_every: int = 2000
|
| sample_steps: int = 8
|
| sample_cfg: float = 4.0
|
| sample_captions: tuple = ()
|
| sample_blocks: int = 3
|
| save_every: int = 1000
|
| keep_ema_every: int = 10000
|
| overfit_batch: int = 0
|
| resume: bool = True
|
|
|
|
|
|
|
|
|
| def set_seed(seed):
|
| torch.manual_seed(seed)
|
| torch.cuda.manual_seed_all(seed)
|
|
|
|
|
| def lr_at(step, cfg):
|
| if step < cfg.warmup:
|
| return cfg.lr * (step + 1) / cfg.warmup
|
| p = min(1.0, (step - cfg.warmup) / max(1, cfg.max_steps - cfg.warmup))
|
| return cfg.lr * (cfg.min_lr_frac + (1 - cfg.min_lr_frac) * 0.5 * (1 + math.cos(math.pi * p)))
|
|
|
|
|
| def load_null_text(path):
|
| with safe_open(str(path), framework="pt") as f:
|
| return f.get_slice("emb")[0:1][0].float(), f.get_slice("mask")[0:1][0].bool()
|
|
|
|
|
| class EMA:
|
| def __init__(self, model, decay):
|
| self.decay = decay
|
| self.shadow = {k: v.detach().clone().float() for k, v in model.state_dict().items()}
|
|
|
| @torch.no_grad()
|
| def update(self, model, step=None):
|
|
|
| d = self.decay if step is None else min(self.decay, (1 + step) / (10 + step))
|
| for k, v in model.state_dict().items():
|
| s = self.shadow[k]
|
| if s.dtype.is_floating_point:
|
| s.mul_(d).add_(v.detach().float(), alpha=1 - d)
|
| else:
|
| s.copy_(v)
|
|
|
| def state_dict(self):
|
| return self.shadow
|
|
|
| def load_state_dict(self, sd):
|
| self.shadow = {k: v.clone() for k, v in sd.items()}
|
|
|
| def to_model(self, model_cfg, device, dtype=torch.bfloat16):
|
| m = BoraDiT(model_cfg)
|
| m.load_state_dict({k: v for k, v in self.shadow.items()})
|
| return m.to(device, dtype).eval()
|
|
|
|
|
| def make_model(cfg, device):
|
| model, mcfg = build(cfg.preset, time_shift=cfg.time_shift, **cfg.model)
|
| model.grad_checkpoint = cfg.grad_checkpoint
|
| model.to(device)
|
| return model, mcfg
|
|
|
|
|
| def make_data(cfg):
|
| root = Path(cfg.data_root)
|
| emb = root / cfg.text_embeds
|
| kw = dict(mode="block", block_len=cfg.block_len, overlap=cfg.overlap, val_mod=cfg.val_mod, text_embeds=emb, dtype=torch.float32)
|
| train_ds = BoraLatents(root, "train", random_block=True, primary_p=cfg.primary_p, **kw)
|
| val_ds = BoraLatents(root, "val", random_block=False, **kw)
|
| return train_ds, val_ds
|
|
|
|
|
| def infinite(loader):
|
| while True:
|
| for b in loader:
|
| yield b
|
|
|
|
|
| def to_device(batch, device):
|
| return {k: (v.to(device, non_blocking=True) if torch.is_tensor(v) else v) for k, v in batch.items()}
|
|
|
|
|
| def apply_uncond(text, mask, p, null_text, null_mask):
|
| """Replace a random fraction p of the text rows by the empty caption (classifier-free guidance training)."""
|
| if p <= 0:
|
| return text, mask
|
| drop = torch.rand(text.shape[0], device=text.device) < p
|
| text = torch.where(drop[:, None, None], null_text[None].to(text.dtype), text)
|
| mask = torch.where(drop[:, None], null_mask[None], mask)
|
| return text, mask
|
|
|
|
|
| @torch.no_grad()
|
| def evaluate(model, val_dl, cfg, device, null_text, null_mask, n_batches):
|
| """Loss of `model` on the val set at fixed timesteps with fixed noise. Returns {"val_loss": .., "t=0.25": .., ...}."""
|
| model.eval()
|
| g = torch.Generator(device=device).manual_seed(1234)
|
| sums = {t: 0.0 for t in cfg.val_t}
|
| n = 0
|
| for i, batch in enumerate(val_dl):
|
| if i >= n_batches:
|
| break
|
| batch = to_device(batch, device)
|
| x = batch["latents"]
|
| text, mask = batch["text_emb"].float(), batch["text_mask"]
|
| noise = torch.randn(x.shape, device=device, generator=g)
|
| for t in cfg.val_t:
|
| tt = torch.full((x.shape[0],), t, device=device)
|
| with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=cfg.bf16 and device.type == "cuda"):
|
| loss, _ = flow_loss(model, x, text, mask, p_cond=cfg.p_cond, shift=cfg.time_shift, t=tt, noise=noise)
|
| sums[t] += loss.item() * x.shape[0]
|
| n += x.shape[0]
|
| model.train()
|
| out = {f"t={t}": sums[t] / max(1, n) for t in cfg.val_t}
|
| out["val_loss"] = sum(out.values()) / len(cfg.val_t)
|
| return out
|
|
|
|
|
| def pick_sample_captions(cfg, val_ds, train_ds):
|
| if cfg.sample_captions:
|
| caps = train_ds.captions
|
| ids = []
|
| for c in cfg.sample_captions:
|
| assert c in caps, f"caption not in captions.json: {c!r} (sample any string by encoding it with T5 in the notebook)"
|
| ids.append(caps.index(c))
|
| return ids
|
| step = max(1, len(val_ds) // (3 * 4))
|
| return [val_ds[i * 3 * step]["caption_id"] for i in range(4)]
|
|
|
|
|
| @torch.no_grad()
|
| def sample_captions(ema_model, caption_ids, emb_path, cfg, device, null_text, null_mask, seed=0):
|
| with safe_open(str(emb_path), framework="pt") as f:
|
| text = torch.stack([f.get_slice("emb")[i:i + 1][0] for i in caption_ids]).float().to(device)
|
| mask = torch.stack([f.get_slice("mask")[i:i + 1][0] for i in caption_ids]).bool().to(device)
|
| B = len(caption_ids)
|
| g = torch.Generator(device=device).manual_seed(seed)
|
| return sample_video(ema_model, text, mask, null_text[None].expand(B, -1, -1).to(device), null_mask[None].expand(B, -1).to(device),
|
| n_blocks=cfg.sample_blocks, overlap=cfg.overlap, steps=cfg.sample_steps, cfg_scale=cfg.sample_cfg,
|
| shift=cfg.time_shift, generator=g).float().cpu()
|
|
|
|
|
| def save_checkpoint(out, model, ema, opt, step, cfg, mcfg):
|
| tmp = out / "ckpt_last.pt.tmp"
|
| torch.save({"model": model.state_dict(), "ema": ema.state_dict(), "opt": opt.state_dict(), "step": step,
|
| "cfg": asdict(cfg), "model_cfg": config_dict(mcfg)}, tmp)
|
| os.replace(tmp, out / "ckpt_last.pt")
|
| ema_sd = {k: v.to(torch.bfloat16) for k, v in ema.state_dict().items() if v.dtype.is_floating_point}
|
| ema_sd.update({k: v for k, v in ema.state_dict().items() if not v.dtype.is_floating_point})
|
| save_file(ema_sd, str(out / f"ema_step{step}.safetensors"), metadata={"step": str(step), "model_cfg": json.dumps(config_dict(mcfg))})
|
| (out / "config.json").write_text(json.dumps({"train": asdict(cfg), "model": config_dict(mcfg), "step": step}, indent=1))
|
| for p in sorted(out.glob("ema_step*.safetensors")):
|
| s = int(p.stem[len("ema_step"):])
|
| if s != step and s % cfg.keep_ema_every != 0:
|
| p.unlink()
|
|
|
|
|
| def load_ema_model(path, device, dtype=torch.bfloat16):
|
| """ema_stepN.safetensors -> BoraDiT in eval mode."""
|
| with safe_open(str(path), framework="pt") as f:
|
| mcfg = json.loads(f.metadata()["model_cfg"])
|
| mcfg["patch"] = tuple(mcfg["patch"])
|
| m = BoraDiT(DiTConfig(**mcfg))
|
| m.load_state_dict(load_file(str(path)))
|
| return m.to(device, dtype).eval()
|
|
|
|
|
|
|
|
|
| def train(cfg: TrainConfig, decode_fn=None, log=print):
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| set_seed(cfg.seed)
|
| out = Path(cfg.out_dir)
|
| out.mkdir(parents=True, exist_ok=True)
|
| (out / "samples").mkdir(exist_ok=True)
|
| torch.backends.cuda.matmul.allow_tf32 = True
|
| torch.backends.cudnn.allow_tf32 = True
|
|
|
|
|
| train_ds, val_ds = make_data(cfg)
|
| emb_path = Path(cfg.data_root) / cfg.text_embeds
|
| null_text, null_mask = load_null_text(emb_path)
|
| null_text, null_mask = null_text.to(device), null_mask.to(device)
|
| pin = device.type == "cuda"
|
| if cfg.overfit_batch:
|
| fixed = next(iter(DataLoader(train_ds, batch_size=cfg.overfit_batch, shuffle=True, num_workers=0)))
|
| fixed["text_emb"] = fixed["text_emb"].float()
|
| train_it = iter(lambda: fixed, None)
|
| save_file({"latents": fixed["latents"].contiguous(), "global_idx": fixed["global_idx"], "start": fixed["start"]},
|
| str(out / "overfit_batch.safetensors"), metadata={"prompts": json.dumps(fixed["text"])})
|
| log(f"overfit mode: one fixed batch of {cfg.overfit_batch} | captions: {[t[:40] for t in fixed['text']]}")
|
| else:
|
| train_dl = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.num_workers, drop_last=True,
|
| pin_memory=pin, persistent_workers=cfg.num_workers > 0, prefetch_factor=4 if cfg.num_workers > 0 else None)
|
| train_it = infinite(train_dl)
|
| val_dl = DataLoader(val_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=min(2, cfg.num_workers), pin_memory=pin)
|
| log(f"data: train {len(train_ds)} clips / {train_ds.n_videos} videos | val {len(val_ds)} blocks / {val_ds.n_videos} videos | captions {len(train_ds.captions) - 1}")
|
|
|
|
|
| model, mcfg = make_model(cfg, device)
|
| log(f"model {cfg.preset}: {count_params(model) / 1e6:.1f}M params | dim {mcfg.dim} depth {mcfg.depth} heads {mcfg.heads} attn {mcfg.attn} | tokens {mcfg.tokens}")
|
| opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr, betas=cfg.betas, weight_decay=cfg.weight_decay, fused=device.type == "cuda")
|
| ema = EMA(model, cfg.ema)
|
| step = 0
|
| ck = out / "ckpt_last.pt"
|
| if cfg.resume and ck.exists():
|
| sd = torch.load(ck, map_location=device, weights_only=False)
|
| model.load_state_dict(sd["model"]); ema.load_state_dict(sd["ema"]); opt.load_state_dict(sd["opt"]); step = sd["step"]
|
| log(f"resumed from {ck} at step {step}")
|
| fwd = torch.compile(model) if cfg.compile else model
|
|
|
| sample_ids = pick_sample_captions(cfg, val_ds, train_ds) if not cfg.overfit_batch else [int(i) for i in fixed["caption_id"][:4]]
|
| sample_prompts = [train_ds.captions[i] for i in sample_ids]
|
| log(f"sample prompts: {[p[:50] for p in sample_prompts]}")
|
|
|
| def write_row(name, row):
|
| p = out / name
|
| new = not p.exists()
|
| with open(p, "a", newline="") as f:
|
| w = csv.DictWriter(f, fieldnames=list(row.keys()))
|
| if new:
|
| w.writeheader()
|
| w.writerow(row)
|
|
|
| def do_sample(step):
|
| ema_model = ema.to_model(mcfg, device)
|
| lat = sample_captions(ema_model, sample_ids, emb_path, cfg, device, null_text, null_mask)
|
| save_file({"latents": lat}, str(out / "samples" / f"step{step}.safetensors"), metadata={"prompts": json.dumps(sample_prompts)})
|
| del ema_model
|
| if decode_fn is not None:
|
| decode_fn(lat, sample_prompts, step)
|
| torch.cuda.empty_cache() if device.type == "cuda" else None
|
|
|
|
|
| model.train()
|
| t0, n_since, loss_acc, loss_n = time.time(), 0, 0.0, 0
|
| log(f"training from step {step} to {cfg.max_steps} | batch {cfg.overfit_batch or cfg.batch_size} | bf16 {cfg.bf16} | grad_checkpoint {cfg.grad_checkpoint} | device {device}")
|
| while step < cfg.max_steps:
|
| batch = to_device(next(train_it), device)
|
| x = batch["latents"]
|
| text, mask = apply_uncond(batch["text_emb"].float(), batch["text_mask"], cfg.p_uncond, null_text, null_mask)
|
| for g in opt.param_groups:
|
| g["lr"] = lr_at(step, cfg)
|
| with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=cfg.bf16 and device.type == "cuda"):
|
| loss, parts = flow_loss(fwd, x, text, mask, p_cond=cfg.p_cond, shift=cfg.time_shift)
|
| opt.zero_grad(set_to_none=True)
|
| loss.backward()
|
| gn = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip) if cfg.grad_clip else torch.tensor(0.0)
|
| if not torch.isfinite(loss):
|
| log(f"step {step}: non-finite loss, batch skipped")
|
| continue
|
| opt.step()
|
| ema.update(model, step)
|
| step += 1
|
| loss_acc += loss.item(); loss_n += 1; n_since += x.shape[0]
|
|
|
| if step % cfg.log_every == 0:
|
| dt = time.time() - t0
|
| row = {"step": step, "loss": loss_acc / loss_n, "lr": lr_at(step - 1, cfg), "grad_norm": float(gn), "samples_per_s": n_since / dt, "t_mean": parts["t_mean"]}
|
| write_row("log.csv", row)
|
| log(f"step {step} | loss {row['loss']:.4f} | lr {row['lr']:.2e} | gn {row['grad_norm']:.2f} | {row['samples_per_s']:.1f} samples/s")
|
| t0, n_since, loss_acc, loss_n = time.time(), 0, 0.0, 0
|
| if cfg.val_every and step % cfg.val_every == 0 and not cfg.overfit_batch:
|
| ema_model = ema.to_model(mcfg, device, dtype=torch.float32)
|
| v = evaluate(ema_model, val_dl, cfg, device, null_text, null_mask, cfg.val_batches)
|
| del ema_model
|
| write_row("val.csv", {"step": step, **v})
|
| log(f"step {step} | val {v['val_loss']:.4f} | " + " ".join(f"{k} {val:.3f}" for k, val in v.items() if k != "val_loss"))
|
| if cfg.sample_every and step % cfg.sample_every == 0:
|
| do_sample(step)
|
| if cfg.save_every and step % cfg.save_every == 0:
|
| save_checkpoint(out, model, ema, opt, step, cfg, mcfg)
|
| log(f"step {step} | saved ckpt_last.pt and ema_step{step}.safetensors")
|
| if not (cfg.save_every and step % cfg.save_every == 0):
|
| save_checkpoint(out, model, ema, opt, step, cfg, mcfg)
|
| if cfg.sample_every and step % cfg.sample_every != 0:
|
| do_sample(step)
|
| log(f"done at step {step}")
|
| return out
|
|
|
