File size: 6,465 Bytes
ebec118 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """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)
# Fixed held-out triplets: a different video, larger gaps, three moments.
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:
# RIFE only makes the middle; reach other moments by halving recursively.
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)
|