| |
| |
| |
| |
| |
| import os, sys, time, math, json |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from torch.utils.data import Dataset, DataLoader |
| from PIL import Image |
| sys.path.insert(0, "/home/z/my-project/scripts/os6") |
| from os6_models import FrameGenerator |
| from os6_draw import rasterize_hint, build_state |
|
|
| DATA = "/home/z/my-project/data/einkos6" |
| CKPT = "/home/z/my-project/checkpoints6/gen_ckpt.pt" |
| PHASE0_STEPS = 4200 |
| PHASE1_STEPS = 700 |
|
|
| torch.set_num_threads(2) |
| try: |
| torch.multiprocessing.set_sharing_strategy("file_system") |
| except Exception: |
| pass |
|
|
|
|
| class GlassDS(Dataset): |
| def __init__(self, phase, recs): |
| self.phase = phase |
| self.recs = recs |
|
|
| def __len__(self): |
| return len(self.recs) |
|
|
| def __getitem__(self, i): |
| rec = self.recs[i] |
| fr = np.asarray(Image.open(os.path.join( |
| DATA, "frames", "%06d.jpg" % rec["i"])).convert("RGB")) |
| if self.phase == 0: |
| fr = fr.reshape(544, 2, 256, 2, 3).mean(axis=(1, 3)) |
| hw, hh = 128, 272 |
| else: |
| hw, hh = 256, 544 |
| hint = rasterize_hint(rec["shapes"], rec["ripples"], hw, hh) |
| st = build_state(rec["theme"], rec["shimmer"], rec["day"], rec["pressed"]) |
| t = torch.from_numpy(np.ascontiguousarray(fr)).permute(2, 0, 1).float() / 255.0 |
| h = torch.from_numpy(hint) |
| s = torch.from_numpy(st) |
| return h, s, t |
|
|
|
|
| def sobel_mag(x): |
| kx = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], |
| dtype=torch.float32).view(1, 1, 3, 3) |
| ky = kx.transpose(2, 3).contiguous() |
| g = x.mean(1, keepdim=True) |
| gx = F.conv2d(g, kx, padding=1) |
| gy = F.conv2d(g, ky, padding=1) |
| return torch.sqrt(gx * gx + gy * gy + 1e-6) |
|
|
|
|
| def loss_fn(pred, tgt): |
| l1 = (pred - tgt).abs().mean() |
| e = (sobel_mag(pred) - sobel_mag(tgt)).abs().mean() |
| return l1 + 0.5 * e, l1.item(), e.item() |
|
|
|
|
| def main(): |
| recs = [json.loads(l) for l in open(os.path.join(DATA, "layouts.jsonl"))] |
| dev = "cpu" |
| model = FrameGenerator().to(dev) |
| step = 0 |
| phase = 0 |
| best = 1e9 |
| ck = None |
| if os.path.exists(CKPT): |
| ck = torch.load(CKPT, map_location=dev, weights_only=False) |
| model.load_state_dict(ck["m"]) |
| phase = ck["phase"]; step = ck["step"]; best = ck.get("best", 1e9) |
| print("resume step %d phase %d best %.4f" % (step, phase, best), flush=True) |
| lr = 2.5e-4 if phase == 0 else 8e-5 |
| lr_end = 1.2e-4 if phase == 0 else 2.5e-5 |
| extra = int(sys.argv[1]) if len(sys.argv) > 1 else 0 |
| p0 = step if phase == 0 else 0 |
| p1 = step - PHASE0_STEPS if phase == 1 else 0 |
| bs = 6 if phase == 0 else 2 |
| total = PHASE0_STEPS if phase == 0 else PHASE1_STEPS + extra |
| done = p0 if phase == 0 else p1 |
|
|
| idxs = list(range(len(recs))) |
| import random |
| random.Random(7).shuffle(idxs) |
| ds = GlassDS(phase, [recs[i] for i in idxs]) |
| dl = DataLoader(ds, batch_size=bs, shuffle=True, num_workers=0, |
| drop_last=True) |
| opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5) |
| if ck is not None and phase == ck["phase"]: |
| try: |
| opt.load_state_dict(ck["o"]) |
| except Exception: |
| pass |
|
|
| def val_psnr(idxs_val): |
| ps = [] |
| with torch.no_grad(): |
| for vi in idxs_val: |
| h, s, t = ds[vi] |
| p = model(h.unsqueeze(0), s.unsqueeze(0)) |
| mse = ((p - t.unsqueeze(0)) ** 2).mean().item() |
| ps.append(10 * math.log10(1.0 / max(1e-8, mse))) |
| return float(np.mean(ps)) |
|
|
| t0 = time.time() |
| ep = 0 |
| while True: |
| for hb, sb, tb in dl: |
| prog = min(1.0, done / float(total)) |
| cur = lr_end + 0.5 * (lr - lr_end) * (1 + math.cos(math.pi * prog)) |
| for g in opt.param_groups: |
| g["lr"] = cur |
| pred = model(hb, sb) |
| loss, l1v, ev = loss_fn(pred, tb) |
| opt.zero_grad() |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| opt.step() |
| done += 1; step += 1 |
| if done % 30 == 0: |
| tmp = CKPT + ".tmp" |
| torch.save({"m": model.state_dict(), "o": opt.state_dict(), |
| "step": step, "phase": phase, "best": best}, tmp) |
| os.replace(tmp, CKPT) |
| if done % 150 == 0: |
| el = time.time() - t0 |
| print("p%d step %d/%d loss %.4f (l1 %.4f e %.4f) lr %.1e %.2f s/it" % |
| (phase, done, total, loss.item(), l1v, ev, cur, el / 150.0), |
| flush=True) |
| t0 = time.time() |
| if done % 600 == 0 or done == total: |
| model.eval() |
| vidx = [i for i in range(len(ds)) if i % 16 == 5][:24] |
| v = val_psnr(vidx) |
| print(" [val] phase %d step %d PSNR %.2f dB" % (phase, done, v), |
| flush=True) |
| if v > best: |
| best = v |
| torch.save({"m": model.state_dict(), "step": step, |
| "phase": phase, "best": best, "opt": None}, |
| CKPT.replace(".pt", "_best.pt")) |
| model.train() |
| if done >= total: |
| if phase == 0: |
| ck2 = torch.load(CKPT, map_location=dev, weights_only=False) |
| ck2["phase"] = 1; ck2["step"] = PHASE0_STEPS |
| tmp = CKPT + ".tmp" |
| torch.save(ck2, tmp); os.replace(tmp, CKPT) |
| print("PHASE 0 COMPLETE -> phase 1 (app-res fine-tune)", |
| flush=True) |
| return |
| else: |
| torch.save({"m": model.state_dict(), "step": step, |
| "phase": 2, "best": best}, |
| CKPT.replace(".pt", "_final.pt")) |
| print("TRAINING DONE step %d best-val %.2f dB" % (step, best), |
| flush=True) |
| return |
| ep += 1 |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
