| """Fine-tune LiquidInbetween from RIFE's weights. |
| |
| Each sample is two frames 2–12 steps apart at 15 fps (0.13–0.8 s: photos |
| taken moments apart), and a real frame somewhere between them as the target |
| at its own t. The held-out set is a different video entirely. |
| |
| python train.py <frames dir> <eval frames dir> <teacher.safetensors> <out dir> [iterations] |
| """ |
| import json, math, random, sys, time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| sys.path.insert(0, str(Path(__file__).parent)) |
| from liquid_model import load_teacher, student_from_teacher |
|
|
| frames_dir, eval_dir, teacher_path, out_dir = map(Path, sys.argv[1:5]) |
| iterations = int(sys.argv[5]) if len(sys.argv) > 5 else 6000 |
| out_dir.mkdir(parents=True, exist_ok=True) |
| device = torch.device("mps") |
| random.seed(7) |
| np.random.seed(7) |
| torch.manual_seed(7) |
|
|
|
|
| def load_set(d): |
| index = json.loads((d / "index.json").read_text()) |
| return [(np.load(d / f"{v['name']}.npy", mmap_mode="r"), set(v["cuts"])) for v in index] |
|
|
|
|
| train_set, eval_set = load_set(frames_dir), load_set(eval_dir) |
|
|
|
|
| def triplet(videos, gap_range, size=None): |
| """Two frames `gap` apart with no cut between them, and one in between.""" |
| while True: |
| frames, cuts = random.choices(videos, weights=[len(v[0]) for v in videos])[0] |
| gap = random.randint(*gap_range) |
| if len(frames) <= gap: |
| continue |
| i = random.randrange(0, len(frames) - gap) |
| if any(i < c <= i + gap for c in cuts): |
| continue |
| j = random.randint(1, gap - 1) |
| a, m, b = frames[i], frames[i + j], frames[i + gap] |
| t = j / gap |
| if size: |
| h, w = a.shape[:2] |
| y, x = random.randrange(0, h - size + 1), random.randrange(0, w - size + 1) |
| a, m, b = (f[y:y + size, x:x + size] for f in (a, m, b)) |
| return a, m, b, t |
|
|
|
|
| def to_tensor(arrays): |
| return torch.from_numpy(np.stack(arrays)).permute(0, 3, 1, 2).float().div(255) |
|
|
|
|
| def batch(n, size=224): |
| items = [] |
| for _ in range(n): |
| a, m, b, t = triplet(train_set, (2, 12), size) |
| if random.random() < 0.5: |
| a, m, b = a[:, ::-1], m[:, ::-1], b[:, ::-1] |
| if random.random() < 0.5: |
| a, b, t = b, a, 1 - t |
| items.append((np.ascontiguousarray(a), np.ascontiguousarray(m), np.ascontiguousarray(b), t)) |
| a, m, b = (to_tensor([x[k] for x in items]).to(device) for k in range(3)) |
| t = torch.tensor([x[3] for x in items], dtype=torch.float32, device=device).view(-1, 1, 1, 1) |
| return a, m, b, t |
|
|
|
|
| def charbonnier(x, y): |
| return torch.sqrt((x - y) ** 2 + 1e-6).mean() |
|
|
|
|
| def laplacian_loss(x, y, levels=4): |
| loss = 0 |
| for _ in range(levels): |
| xd, yd = F.avg_pool2d(x, 2), F.avg_pool2d(y, 2) |
| loss = loss + (x - F.interpolate(xd, size=x.shape[-2:], mode="bilinear", align_corners=False) |
| - (y - F.interpolate(yd, size=y.shape[-2:], mode="bilinear", align_corners=False))).abs().mean() |
| x, y = xd, yd |
| return loss |
|
|
|
|
| teacher = load_teacher(teacher_path).to(device).eval() |
| student = student_from_teacher(load_teacher(teacher_path)).to(device) |
| for p in teacher.parameters(): |
| p.requires_grad_(False) |
|
|
| |
| random.seed(123) |
| held = [] |
| while len(held) < 120: |
| a, m, b, t = triplet(eval_set, (6, 10)) |
| if t in (0.25, 0.5, 0.75) or len(held) % 3 == 0: |
| held.append((a, m, b, t)) |
| random.seed(7) |
|
|
|
|
| def psnr(x, y): |
| return (-10 * torch.log10(((x - y) ** 2).mean(dim=(1, 2, 3)) + 1e-10)).mean().item() |
|
|
|
|
| @torch.no_grad() |
| def evaluate(net, timestep): |
| scores = {"blend": [], "model": []} |
| for a, m, b, t in held: |
| a_, m_, b_ = (to_tensor([np.ascontiguousarray(f)]).to(device) for f in (a, m, b)) |
| h, w = a_.shape[-2:] |
| H, W = (h + 31) // 32 * 32, (w + 31) // 32 * 32 |
| pad = lambda x: F.pad(x, (0, W - w, 0, H - h), mode="replicate") |
| if timestep: |
| out = net(pad(a_), pad(b_), torch.tensor([[[[t]]]], device=device)) |
| else: |
| |
| def at(lo, hi, x0, x1, depth=0): |
| mid = net(x0, x1) |
| c = (lo + hi) / 2 |
| if abs(c - t) < 1e-6 or depth > 4: |
| return mid |
| return at(lo, c, x0, mid, depth + 1) if t < c else at(c, hi, mid, x1, depth + 1) |
| out = at(0.0, 1.0, pad(a_), pad(b_)) |
| out = out[..., :h, :w].clamp(0, 1) |
| scores["model"].append(psnr(out, m_)) |
| scores["blend"].append(psnr(a_ * (1 - t) + b_ * t, m_)) |
| return {k: float(np.mean(v)) for k, v in scores.items()} |
|
|
|
|
| log = [] |
| baseline = evaluate(teacher, timestep=False) |
| start = evaluate(student, timestep=True) |
| print(f"held-out PSNR: blend {baseline['blend']:.2f}, RIFE {baseline['model']:.2f}, ours at start {start['model']:.2f}", flush=True) |
| log.append({"iteration": 0, "blend": baseline["blend"], "rife": baseline["model"], "ours": start["model"]}) |
|
|
| optimizer = torch.optim.AdamW(student.parameters(), lr=1e-4, weight_decay=1e-4) |
| best = start["model"] |
| began = time.time() |
| for it in range(1, iterations + 1): |
| lr = 1e-5 + 0.5 * (1e-4 - 1e-5) * (1 + math.cos(math.pi * it / iterations)) |
| for g in optimizer.param_groups: |
| g["lr"] = lr |
| a, m, b, t = batch(8) |
| outs = student(a, b, t, all_scales=True) |
| loss = charbonnier(outs[-1], m) + 0.5 * laplacian_loss(outs[-1], m) + 0.25 * charbonnier(outs[-2], m) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0) |
| optimizer.step() |
| if it % 100 == 0: |
| print(f"it {it} loss {loss.item():.4f} lr {lr:.2e} {(time.time() - began) / it:.2f}s/it", flush=True) |
| if it % 500 == 0 or it == iterations: |
| student.eval() |
| score = evaluate(student, timestep=True)["model"] |
| student.train() |
| log.append({"iteration": it, "ours": score}) |
| print(f"== it {it}: held-out PSNR ours {score:.2f} (RIFE {baseline['model']:.2f}, blend {baseline['blend']:.2f})", flush=True) |
| if score > best: |
| best = score |
| torch.save(student.state_dict(), out_dir / "liquid_inbetween.pt") |
| (out_dir / "log.json").write_text(json.dumps(log, indent=1)) |
| torch.save(student.state_dict(), out_dir / "liquid_inbetween_last.pt") |
| print("done; best held-out PSNR", best, flush=True) |
|
|
