|
|
| import torch, cv2
|
| from collections import deque
|
| from model import NyvexUNet
|
| from diffusion import Diffusion
|
|
|
| K, SIZE, FPS, N = 3, (128, 72), 20, 400
|
| CKPT = "ckpts/nyvex-0.1.pt"
|
| PROMPT = "prompt.png"
|
|
|
| model = NyvexUNet(k=K).cuda().eval()
|
| model.load_state_dict(torch.load(CKPT, map_location="cuda")["ema"])
|
| diff = Diffusion(1000, "cuda")
|
|
|
| def to_t(bgr):
|
| rgb = cv2.cvtColor(cv2.resize(bgr, SIZE), cv2.COLOR_BGR2RGB)
|
| return torch.from_numpy(rgb).float().div(127.5).sub(1).permute(2, 0, 1)
|
|
|
| def stabilize(t):
|
| a = t.add(1).mul(127.5).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()[..., ::-1]
|
| a = cv2.imdecode(cv2.imencode(".jpg", a, [cv2.IMWRITE_JPEG_QUALITY, 90])[1], 1)[..., ::-1]
|
| return torch.from_numpy(a.copy()).float().div(127.5).sub(1).permute(2, 0, 1)
|
|
|
| first = to_t(cv2.imread(PROMPT))
|
| hist = deque([first] * K, maxlen=64)
|
| act = torch.zeros(1, 14, device="cuda")
|
|
|
| vw = cv2.VideoWriter("Nyvex-Minecraft-Bedwars-Dream-v0.mp4",
|
| cv2.VideoWriter_fourcc(*"mp4v"), FPS, SIZE)
|
|
|
| for i in range(N):
|
| cond = torch.stack(list(hist)[-K:])[None].reshape(1, -1, SIZE[1], SIZE[0]).cuda()
|
| with torch.autocast("cuda", dtype=torch.bfloat16), torch.inference_mode():
|
| nxt = stabilize(diff.sample(model, cond, act, steps=8)[0].float().cpu())
|
| hist.append(nxt)
|
| bgr = nxt.add(1).mul(127.5).byte().permute(1, 2, 0).cpu().numpy()[..., ::-1]
|
| vw.write(bgr)
|
| print(f"\rframe {i+1}/{N}", end="")
|
|
|
| vw.release()
|
| print("\n✅ Nyvex-Minecraft-Bedwars-Dream-v0.mp4 généré") |
