[tab] sample_resolution_batches.py (python)
from dataclasses import dataclass
import time
import os

import dm_pix as pix
import einops
import jax
import jax.numpy as jnp
import numpy as np
from orbax.checkpoint import PyTreeCheckpointer
from PIL import Image, ImageDraw
import tyro

from genie import Genie
from utils.dataloader import get_dataloader


@dataclass
class Args:
    # Experiment
    seed: int = 0
    seq_len: int = 16
    image_channels: int = 3
    image_height: int = 90
    image_width: int = 160
    data_dir: str = "data/coinrun_episodes"
    checkpoint: str = ""
    # Sampling
    batch_size: int = 1
    maskgit_steps: int = 25
    temperature: float = 1.0
    sample_argmax: bool = True
    start_frame: int = 0
    output_dir: str = "sample_results"
    # Tokenizer checkpoint
    tokenizer_dim: int = 512
    latent_patch_dim: int = 32
    num_patch_latents: int = 1024
    patch_size: int = 4
    tokenizer_num_blocks: int = 8
    tokenizer_num_heads: int = 8
    # LAM checkpoint
    lam_dim: int = 512
    latent_action_dim: int = 32
    num_latent_actions: int = 6
    lam_patch_size: int = 16
    lam_num_blocks: int = 8
    lam_num_heads: int = 8
    # Dynamics checkpoint
    dyna_dim: int = 512
    dyna_num_blocks: int = 12
    dyna_num_heads: int = 8


args = tyro.cli(Args)
rng = jax.random.PRNGKey(args.seed)

# --- Load Genie checkpoint ---
genie = Genie(
    # Tokenizer
    in_dim=args.image_channels,
    tokenizer_dim=args.tokenizer_dim,
    latent_patch_dim=args.latent_patch_dim,
    num_patch_latents=args.num_patch_latents,
    patch_size=args.patch_size,
    tokenizer_num_blocks=args.tokenizer_num_blocks,
    tokenizer_num_heads=args.tokenizer_num_heads,
    # LAM
    lam_dim=args.lam_dim,
    latent_action_dim=args.latent_action_dim,
    num_latent_actions=args.num_latent_actions,
    lam_patch_size=args.lam_patch_size,
    lam_num_blocks=args.lam_num_blocks,
    lam_num_heads=args.lam_num_heads,
    # Dynamics
    dyna_dim=args.dyna_dim,
    dyna_num_blocks=args.dyna_num_blocks,
    dyna_num_heads=args.dyna_num_heads,
)
rng, _rng = jax.random.split(rng)
image_shape = (args.image_height, args.image_width, args.image_channels)
dummy_inputs = dict(
    videos=jnp.zeros((args.batch_size, args.seq_len, *image_shape), dtype=jnp.float32),
    mask_rng=_rng,
)
rng, _rng = jax.random.split(rng)
params = genie.init(_rng, dummy_inputs)
ckpt = PyTreeCheckpointer().restore(args.checkpoint)["model"]["params"]["params"]
params["params"].update(ckpt)


# --- Define autoregressive sampling loop ---
def _autoreg_sample(rng, video_batch, action_batch):
    vid = video_batch[:, : args.start_frame + 1]
    for frame_idx in range(args.start_frame + 1, args.seq_len):
        # --- Sample next frame ---
        print("Frame", frame_idx)
        rng, _rng = jax.random.split(rng)
        batch = dict(videos=vid, latent_actions=action_batch[:, :frame_idx], rng=_rng)
        new_frame = genie.apply(
            params,
            batch,
            args.maskgit_steps,
            args.temperature,
            args.sample_argmax,
            method=Genie.sample,
        )
        vid = jnp.concatenate([vid, new_frame], axis=1)
    return vid


# --- Get video + latent actions ---

tfrecord_files = [
    os.path.join(args.data_dir, x)
    for x in os.listdir(args.data_dir)
    if x.endswith(".tfrecord")
]
image_shape = (args.image_height, args.image_width, args.image_channels)

dataloader = get_dataloader(tfrecord_files, args.seq_len, args.batch_size, *image_shape)
video_batch = next(iter(dataloader))
# Get latent actions from first video only
first_video = video_batch[:1]
batch = dict(videos=first_video)
action_batch = genie.apply(params, batch, False, method=Genie.vq_encode)
action_batch = action_batch.reshape(1, args.seq_len - 1, 1)
# Use actions from first video for all videos
action_batch = jnp.repeat(action_batch, video_batch.shape[0], axis=0)

# --- Sample + evaluate video ---
vid = _autoreg_sample(rng, video_batch, action_batch)
gt = video_batch[:, : vid.shape[1]].clip(0, 1).reshape(-1, *video_batch.shape[2:])
recon = vid.clip(0, 1).reshape(-1, *vid.shape[2:])
ssim = pix.ssim(gt[:, args.start_frame + 1 :], recon[:, args.start_frame + 1 :]).mean()
print(f"SSIM: {ssim}")

# --- Construct and save videos ---
# Save first video comparison (true vs predicted)
first_true = (video_batch[0:1] * 255).astype(np.uint8)
first_pred = (vid[0:1] * 255).astype(np.uint8)
first_video_comparison = np.zeros((2, *vid.shape[1:5]), dtype=np.uint8)
first_video_comparison[0] = first_true[:, : vid.shape[1]]
first_video_comparison[1] = first_pred

# Create comparison frames for first video
comparison_frames = einops.rearrange(first_video_comparison, "n t h w c -> t h (n w) c")
imgs = [Image.fromarray(img) for img in comparison_frames]
# Write actions on each frame
for img, action in zip(imgs[1:], action_batch[0, :, 0]):
    d = ImageDraw.Draw(img)
    d.text((2, 2), f"{action}", fill=255)
imgs[0].save(
    f"generation_video_0_{time.time()}.gif",
    save_all=True,
    append_images=imgs[1:],
    duration=250,
    loop=0,
)

# Save each other video in its own GIF

output_dir = f"{args.output_dir}/{time.time()}"
os.makedirs(output_dir, exist_ok=True)
for i in range(1, video_batch.shape[0]):
    # Create comparison: true video vs predicted video
    true_video = (video_batch[i:i+1] * 255).astype(np.uint8)
    pred_video = (vid[i:i+1] * 255).astype(np.uint8)
    video_comparison = np.zeros((2, *vid.shape[1:5]), dtype=np.uint8)
    video_comparison[0] = true_video[:, : vid.shape[1]]
    video_comparison[1] = pred_video
    
    # Create comparison frames
    comparison_frames = einops.rearrange(video_comparison, "n t h w c -> t h (n w) c")
    imgs = [Image.fromarray(img) for img in comparison_frames]
    # Write actions on each frame
    for img, action in zip(imgs[1:], action_batch[i, :, 0]):
        d = ImageDraw.Draw(img)
        d.text((2, 2), f"{action}", fill=255)

    imgs[0].save(
        f"{output_dir}/generation_video_{i}.gif",
        save_all=True,
        append_images=imgs[1:],
        duration=250,
        loop=0,
    )