[tab] /fast/home/franz.srambical/jafar/sample.py (python)
from dataclasses import dataclass
from typing import Optional, Mapping, Any
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
import flax.core
from flax.core import unfreeze, freeze

from genie import Genie
from models.dynamics import DynamicsAutoregressive
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)

cache_struct, _ = flax.core.pop(inference_variables, 'cache')
variables_for_sampling = freeze({
    'params': params,
    'cache': cache_struct
})

def reset_spatial_cache_in_vars(variables: Mapping[str, Any]) -> Mapping[str, Any]:
    """
    Creates a new variables dictionary with the spatial attention cache reset.
    """
    if 'cache' not in variables:
        return variables
    cache_unfrozen = unfreeze(variables['cache'])
    if 'dynamics' in cache_unfrozen and 'dynamics' in cache_unfrozen['dynamics']:
        st_transformer_cache = cache_unfrozen['dynamics']['dynamics']
        for block_name in st_transformer_cache:
            if 'STBlock' in block_name and 'MultiHeadAttention_0' in st_transformer_cache[block_name]:
                spatial_attn_cache = st_transformer_cache[block_name]['MultiHeadAttention_0']
                if 'cache_index' in spatial_attn_cache:
                    spatial_attn_cache['cache_index'] = jnp.array(0, dtype=jnp.int32)
                if 'cached_key' in spatial_attn_cache:
                    spatial_attn_cache['cached_key'] = jnp.zeros_like(spatial_attn_cache['cached_key'])
                if 'cached_value' in spatial_attn_cache:
                    spatial_attn_cache['cached_value'] = jnp.zeros_like(spatial_attn_cache['cached_value'])
    return freeze({'params': variables['params'], 'cache': cache_unfrozen})

# --- Define autoregressive sampling loop ---
def _autoreg_sample(rng, video_batch, action_batch):
    # --- 1. Define the JIT-compiled step function ---
    # @jax.jit
    def predict_token_fn(variables: Mapping[str, Any], inputs: Mapping[str, Any]) -> tuple[Mapping[str, Any], Mapping[str, Any]]:
        """
        This function calls the dynamics submodule's apply method.
        """
        return inference_genie.apply(
            variables,
            inputs,
            method = lambda module, inputs : module.dynamics(inputs, training=False),
            mutable=['cache']
        )

    # --- 2. Initialization ---
    variables = variables_for_sampling
    context_frames = video_batch[:, : args.start_frame + 1]

    # Fix parameter access - parameters are nested under params["params"]
    tokenizer_params = {'params': variables['params']['params']['tokenizer']}
    tokenizer_out = inference_genie.apply(
        tokenizer_params,
        context_frames,
        method=lambda module, videos: module.tokenizer.vq_encode(videos, training=False)
    )
    context_tokens = tokenizer_out['indices']
    B, T_context, N = context_tokens.shape
    
    token_sequence = jnp.zeros((B, args.seq_len, N), dtype=context_tokens.dtype)
    token_sequence = token_sequence.at[:, :T_context, :].set(context_tokens)

    lam_params = {'params': variables['params']['params']['lam']}
    action_codes = inference_genie.apply(
        lam_params,
        action_batch,
        method=lambda module, actions: module.lam.vq.get_codes(actions)
    )

    # --- 3. The Sampling Loop (State Management) ---
    for t in range(T_context, args.seq_len):
        print(f"Generating frame {t}... resetting spatial cache.")
        variables = reset_spatial_cache_in_vars(variables)

        for n in range(N):
            step_inputs = {
                "video_tokens": token_sequence,
                "latent_actions": action_codes,
            }
            
            dynamics_variables = {'params': variables['params']['params']['dynamics'], 'cache': variables['cache']['dynamics']}
            
            dyna_outputs, updated_dynamics_cache = predict_token_fn(dynamics_variables, step_inputs)
            
            variables_unfrozen = unfreeze(variables)
            variables_unfrozen['cache']['dynamics'] = updated_dynamics_cache['cache']['dynamics']
            variables = freeze(variables_unfrozen)
            
            next_token_logits = dyna_outputs["token_logits"][:, t, n, :]
            rng, step_rng = jax.random.split(rng)
            if args.sample_argmax:
                next_token = jnp.argmax(next_token_logits, axis=-1)
            else:
                next_token = jax.random.categorical(step_rng, next_token_logits / args.temperature)
            
            token_sequence = token_sequence.at[:, t, n].set(next_token)
    
    # --- 4. Finalization (Detokenization) ---
    generated_vid = inference_genie.tokenizer.apply(
        {'params': variables['params']['params']['tokenizer']},
        token_sequence,
        video_batch.shape[2:4],
        method=inference_genie.tokenizer.decode
    )
    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,
)