[content] sample.py (python)
from dataclasses import dataclass
from typing import Optional
import time
import os

import dm_pix as pix
import einops
import jax
import jax.numpy as jnp
import flax.linen as nn
import numpy as np
from flax.training.train_state import TrainState
import grain
import orbax.checkpoint as ocp
import optax
from PIL import Image, ImageDraw
import tyro
from flax.core import unfreeze, freeze

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 = ""
    checkpoint_step: Optional[int] = None
    # Sampling
    batch_size: int = 1
    maskgit_steps: int = 25
    temperature: float = 1.0
    sample_argmax: bool = True
    start_frame: int = 0
    # 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
    lam_co_train: bool = True
    # Dynamics checkpoint
    dyna_dim: int = 512
    dyna_num_blocks: int = 12
    dyna_num_heads: int = 8
    param_dtype: jnp.dtype = jnp.float32
    dtype: jnp.dtype = jnp.bfloat16
    use_flash_attention: bool = True


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

# --- Load Genie checkpoint ---
train_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,
    lam_co_train=args.lam_co_train,
    # Dynamics
    dyna_dim=args.dyna_dim,
    dyna_num_blocks=args.dyna_num_blocks,
    dyna_num_heads=args.dyna_num_heads,
    use_maskgit=False,
    param_dtype=args.param_dtype,
    dtype=args.dtype,
    use_flash_attention=args.use_flash_attention,
    is_training=True,
)
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 = train_genie.init(_rng, dummy_inputs)

dummy_train_state = TrainState.create(
    apply_fn=train_genie.apply,
    params=params,
    tx=optax.adamw(
        optax.warmup_cosine_decay_schedule(
            0, 0, 1, 2 # dummy values
        )
    ), 
)
handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler)
checkpoint_manager = ocp.CheckpointManager(
    args.checkpoint,
    options=ocp.CheckpointManagerOptions(step_format_fixed_length=6),
    handler_registry=handler_registry
)
abstract_train_state = jax.tree_util.tree_map(
    ocp.utils.to_shape_dtype_struct, dummy_train_state
)

restored = checkpoint_manager.restore(
    args.checkpoint_step or checkpoint_manager.latest_step(),
    args=ocp.args.Composite(
        model_state=ocp.args.StandardRestore(abstract_train_state),
    ),
)
restored_train_state = restored["model_state"]
params = restored_train_state.params

inference_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,
    lam_co_train=args.lam_co_train,
    # Dynamics
    dyna_dim=args.dyna_dim,
    dyna_num_blocks=args.dyna_num_blocks,
    dyna_num_heads=args.dyna_num_heads,
    use_maskgit=False,
    param_dtype=args.param_dtype,
    dtype=args.dtype,
    use_flash_attention=args.use_flash_attention,
    is_training=False,
)

rng, _rng = jax.random.split(rng)
inference_variables = inference_genie.init(_rng, dummy_inputs)

variables_for_sampling = unfreeze(inference_variables)
variables_for_sampling["params"] = params
variables_for_sampling = freeze(variables_for_sampling)


def _reset_spatial_kv_cache(variables):
    """
    Reset spatial KV cache for all STBlocks in the dynamics transformer.
    This function manipulates the cache variables directly without changing module structure.
    """
    # Get the dynamics cache
    dynamics_cache = variables["cache"]["dynamics"]
    
    # Reset cache for each STBlock
    for i in range(args.dyna_num_blocks):
        block_name = f"st_block_{i}"
        if block_name in dynamics_cache:
            spatial_attn_cache = dynamics_cache[block_name]["spatial_attn"]["cache"]
            
            # Reset cached_key, cached_value, and cache_index
            if "cached_key" in spatial_attn_cache:
                cached_key = spatial_attn_cache["cached_key"]
                spatial_attn_cache["cached_key"] = jnp.zeros_like(cached_key)
            
            if "cached_value" in spatial_attn_cache:
                cached_value = spatial_attn_cache["cached_value"]
                spatial_attn_cache["cached_value"] = jnp.zeros_like(cached_value)
            
            if "cache_index" in spatial_attn_cache:
                spatial_attn_cache["cache_index"] = jnp.array(0, dtype=jnp.int32)
    
    return variables

def _sampling_wrapper(module: Genie, batch):
    # return module.sample_maskgit(batch, args.seq_len, args.maskgit_steps, args.temperature, args.sample_argmax)
    return module.sample_causal(batch, args.seq_len, args.temperature, args.sample_argmax)

# --- Define autoregressive sampling loop with external KV cache resetting ---
def _autoreg_sample_with_kv_reset(rng, video_batch, action_batch):
    vid = video_batch[:, : args.start_frame + 1]
    sampling_fn = nn.apply(_sampling_wrapper, inference_genie)
    
    # Initialize variables for sampling
    current_variables = variables_for_sampling.copy()
    
    # Get tokenizer output for initial frames
    tokenizer_out = inference_genie.apply(
        current_variables["params"], 
        dict(videos=vid), 
        False, 
        method=inference_genie.tokenizer.vq_encode
    )
    token_idxs = tokenizer_out["indices"]  # (B, T, N)
    B, T, N = token_idxs.shape
    
    # Prepare initial token sequence
    pad_shape = (B, args.seq_len - T, N)
    token_idxs_full = jnp.concatenate(
        [token_idxs, jnp.zeros(pad_shape, dtype=token_idxs.dtype)], axis=1
    )  # (B, seq_len, N)
    
    # Prepare latent actions
    action_tokens = inference_genie.apply(
        current_variables["params"],
        dict(latent_actions=action_batch),
        False,
        method=lambda params, batch: inference_genie.lam.vq.get_codes(batch["latent_actions"])
    )
    
    # Autoregressive generation loop
    rng, _rng = jax.random.split(rng)
    for t in range(T, args.seq_len):
        for n in range(N):
            # Create batch for this step
            batch = dict(
                videos=vid, 
                latent_actions=action_batch, 
                rng=_rng,
                # Pass the current token sequence
                _token_idxs_full=token_idxs_full,
                _current_t=t,
                _current_n=n
            )
            
            # Run dynamics to get logits
            dyna_inputs = {
                "video_tokens": token_idxs_full,
                "latent_actions": action_tokens
            }
            dyna_outputs = inference_genie.apply(
                current_variables["params"],
                dyna_inputs,
                False,
                method=inference_genie.dynamics
            )
            
            # Get logits for current position
            next_token_logits = dyna_outputs["token_logits"][:, t, n, :].astype(jnp.float32)
            
            # Sample or argmax
            if args.sample_argmax:
                next_token = jnp.argmax(next_token_logits, axis=-1)
            else:
                rng, step_rng = jax.random.split(rng)
                next_token = jax.random.categorical(
                    step_rng, next_token_logits / args.temperature, axis=-1
                )
            
            # Update token sequence
            token_idxs_full = token_idxs_full.at[:, t, n].set(next_token)
        
        # Reset spatial KV cache after completing all patches for a frame
        current_variables = _reset_spatial_kv_cache(current_variables)
    
    # Decode final tokens
    final_frames = inference_genie.apply(
        current_variables["params"],
        dict(indices=token_idxs_full, video_hw=vid.shape[2:4]),
        False,
        method=inference_genie.tokenizer.decode
    )
    
    return final_frames

# --- Define autoregressive sampling loop ---
def _autoreg_sample(rng, video_batch, action_batch):
    vid = video_batch[:, : args.start_frame + 1]
    # sampling_fn = jax.jit(nn.apply(_sampling_wrapper, genie)) 
    sampling_fn = nn.apply(_sampling_wrapper, inference_genie)
    rng, _rng = jax.random.split(rng)
    batch = dict(videos=vid, latent_actions=action_batch, rng=_rng)
    generated_vid = sampling_fn(
        {'params': variables_for_sampling["params"]["params"], 'cache': variables_for_sampling["cache"]},
        batch,
    )
    return generated_vid

def _get_dataloader_iterator():
    array_record_files = [
        os.path.join(args.data_dir, x)
        for x in os.listdir(args.data_dir)
        if x.endswith(".array_record")
    ]
    grain_dataloader = get_dataloader(
        array_record_files,
        args.seq_len,
        # NOTE: We deliberately pass the global batch size
        # The dataloader shards the dataset across all processes
        args.batch_size,
        *image_shape,
        num_workers=0,
        prefetch_buffer_size=1,
        seed=args.seed,
    )
    initial_state = grain_dataloader._create_initial_state()
    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
    return grain_iterator

# --- Get video + latent actions ---
grain_iterator = _get_dataloader_iterator()
video_batch = next(grain_iterator)
# video_batch = np.load("overfit_dir/single_sample_corner.npy")

video_batch = video_batch.astype(args.dtype) / 255.0
# Get latent actions for all videos in the batch
batch = dict(videos=video_batch[:,:args.seq_len])
action_batch = inference_genie.apply(variables_for_sampling["params"], batch, False, method=Genie.vq_encode)
action_batch = action_batch.reshape(video_batch.shape[0], args.seq_len - 1, 1)

# --- Sample + evaluate video ---
print("autoreg sampling...")
vid = _autoreg_sample(rng, video_batch, action_batch)
print("autoreg sampling done. calculating ssim and saving video")
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 video ---
true_videos = (video_batch * 255).astype(np.uint8)
pred_videos = (vid * 255).astype(np.uint8)
video_comparison = np.zeros((2, *vid.shape), dtype=np.uint8)
video_comparison[0] = true_videos[:, :args.seq_len]
video_comparison[1] = pred_videos
frames = einops.rearrange(video_comparison, "n b t h w c -> t (b h) (n w) c")

# --- Save video --- 
imgs = [Image.fromarray(img) for img in frames]
# Write actions on each frame, on each row (i.e., for each video in the batch, on the GT row)
for t, img in enumerate(imgs[1:]):
    d = ImageDraw.Draw(img)
    for row in range(action_batch.shape[0]):
        action = action_batch[row, t, 0]
        y_offset = row * video_batch.shape[2] + 2
        d.text((2, y_offset), f"{action}", fill=255)
imgs[0].save(
    f"generation_{time.time()}.gif",
    save_all=True,
    append_images=imgs[1:],
    duration=250,
    loop=0,
)