# rollout.py — ton modèle rêve à partir d'une seule image
import torch, cv2
from collections import deque
from model import NyvexUNet
from diffusion import Diffusion

K, SIZE, FPS, N = 3, (128, 72), 20, 400        # 400 frames = 20 s de vidéo
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):                              # re-JPEG léger = casse-dérive
    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")        # joueur immobile → rêve libre

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é")