[tab] genie.py (python)
from typing import Dict, Any

import optax
import jax
import jax.numpy as jnp
import flax.linen as nn
from flax.training.train_state import TrainState
import orbax.checkpoint as ocp

from flax.core import unfreeze, freeze
from models.dynamics import DynamicsMaskGIT, DynamicsAutoregressive
from models.lam import LatentActionModel
from models.tokenizer import TokenizerVQVAE

import grain


class Genie(nn.Module):
    """Genie model"""

    # --- Tokenizer ---
    in_dim: int
    tokenizer_dim: int
    latent_patch_dim: int
    num_patch_latents: int
    patch_size: int
    tokenizer_num_blocks: int
    tokenizer_num_heads: int
    # --- LAM ---
    lam_dim: int
    latent_action_dim: int
    num_latent_actions: int
    lam_patch_size: int
    lam_num_blocks: int
    lam_num_heads: int
    lam_co_train: bool
    # --- Dynamics ---
    dyna_dim: int
    dyna_num_blocks: int
    dyna_num_heads: int
    use_maskgit: bool
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
    is_training: bool
    dropout: float = 0.0
    mask_limit: float = 0.0

    def setup(self):
        self.tokenizer = TokenizerVQVAE(
            in_dim=self.in_dim,
            model_dim=self.tokenizer_dim,
            latent_dim=self.latent_patch_dim,
            num_latents=self.num_patch_latents,
            patch_size=self.patch_size,
            num_blocks=self.tokenizer_num_blocks,
            num_heads=self.tokenizer_num_heads,
            dropout=0.0,
            codebook_dropout=0.0,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            use_flash_attention=self.use_flash_attention,
        )
        self.lam = LatentActionModel(
            in_dim=self.in_dim,
            model_dim=self.lam_dim,
            latent_dim=self.latent_patch_dim,
            num_latents=self.num_latent_actions,
            patch_size=self.lam_patch_size,
            num_blocks=self.lam_num_blocks,
            num_heads=self.lam_num_heads,
            dropout=0.0,
            codebook_dropout=0.0,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            use_flash_attention=self.use_flash_attention,
        )

        if self.use_maskgit:
            self.dynamics = DynamicsMaskGIT(
                model_dim=self.dyna_dim,
                num_latents=self.num_patch_latents,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                dropout=self.dropout,
                mask_limit=self.mask_limit,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
            ) 
        else:
            self.dynamics = DynamicsAutoregressive(
                model_dim=self.dyna_dim,
                num_latents=self.num_patch_latents,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                is_training=self.is_training,
            )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        tokenizer_outputs = self.tokenizer.vq_encode(batch["videos"], training=False)
        lam_outputs = self.lam.vq_encode(batch["videos"], training=False)
        latent_actions = jax.lax.cond(
            self.lam_co_train,
            lambda: lam_outputs["z_q"],
            lambda: jax.lax.stop_gradient(lam_outputs["z_q"])
        )
        outputs = dict(
            video_tokens=jax.lax.stop_gradient(tokenizer_outputs["indices"]),
            latent_actions=latent_actions,
        )
        outputs["mask_rng"] = batch["mask_rng"]
        dyna_outputs = self.dynamics(outputs, training)
        outputs.update(dyna_outputs)
        mle_indices = jnp.argmax(outputs["token_logits"], axis=-1)
        outputs["recon"] = self.tokenizer.decode(
            mle_indices, batch["videos"].shape[2:4]
        )
        outputs["lam_indices"] = lam_outputs["indices"]
        return outputs


    def sample_causal(
        self,
        batch: Dict[str, Any],
        seq_len: int,
        temperature: float = 1.0,
        sample_argmax: bool = False,
    ):
        """
        Autoregressively samples with a KV cache, resetting the spatial cache each frame.
        """
        assert self.decode, "This sampling method requires the model to be in decode mode."

        # --- 1. Initial Setup ---
        # Get the initial token context
        tokenizer_out = self.tokenizer.vq_encode(batch["videos"], training=False)
        context_tokens = tokenizer_out["indices"]  # (B, T, N)
        B, T_context, N = context_tokens.shape

        # Get latent actions for the full generation sequence
        action_tokens = self.lam.vq.get_codes(batch["latent_actions"]) # (B, S-1, D)
        
        # The full token sequence we will fill up
        generated_tokens = jnp.zeros((B, seq_len, N), dtype=context_tokens.dtype)
        generated_tokens = generated_tokens.at[:, :T_context].set(context_tokens)

        # --- 2. Define the Body of the Per-Frame Generation Loop ---
        
        # This function generates one full frame `t`
        def generate_frame_body(carry, t):
            # Unpack the carry
            current_tokens, variables, rng = carry
            
            # --- a. Reset the spatial cache ---
            # Create a fresh, zero-initialized cache for just the spatial layers
            # We do this by filtering the original initialized cache from `model.init()`
            
            # This is a bit of a hack: we rely on having access to the initial variables
            # We'll pass `initial_variables` into the scan via closure.
            fresh_spatial_cache = jax.tree_util.tree_map_with_path(
                lambda path, x: x if 'spatial_attn' in str(path) else jnp.array([]), 
                initial_variables['cache']
            )
            # Filter out the empty arrays
            fresh_spatial_cache = {k:v for k,v in jax.tree_util.tree_leaves_with_path(fresh_spatial_cache) if v.size > 0}
            fresh_spatial_cache = jax.tree_util.tree_unflatten(jax.tree_util.tree_structure(initial_variables['cache']), fresh_spatial_cache.values())
            
            # Merge the fresh spatial cache into our current cache
            # This overwrites the old spatial cache from the previous frame
            # while keeping the temporal cache intact.
            mutable_vars = unfreeze(variables)
            mutable_vars['cache'] = jax.tree_util.tree_map(
                lambda old, new: new if new.size > 0 else old,
                mutable_vars['cache'], fresh_spatial_cache
            )
            variables = freeze(mutable_vars)
            
            # Get the input for the dynamics model: the *last* generated frame
            prev_frame_tokens = current_tokens[:, t - 1 : t, :] # (B, 1, N)
            current_action = action_tokens[:, t - 1 : t, :]

            # --- b. Inner loop for patches ---
            new_frame_tokens = jnp.zeros((B, N), dtype=current_tokens.dtype)
            
            for n in range(N):
                # Prepare input for this patch step: previous frame + patches so far
                # NOTE: This part is inefficient, a patch-level cache would be better,
                # but let's stick to frame-level for now to fix the main issue.
                # A better approach would be to have a single loop over all tokens. Let's pivot to that.
                pass # This approach is too complex, let's simplify.
            
            # The hierarchical reset is tricky. Let's try a simpler, more direct approach.
            # New Strategy: One loop for all tokens, but we pass the cache in and out.
            # We will manually reset the cache index for spatial layers.
            
            return carry, generated_tokens

        # Let's restart with a better, single-loop strategy.
        # This is the most efficient and correct way.
        
        # --- 2. Initialize Cache with Context ---
        # The first call populates the cache for both spatial and temporal layers
        # for all context frames.
        dyna_inputs = {
            "video_tokens": context_tokens,
            "latent_actions": action_tokens[:, :T_context-1]
        }
        # The variables dict must be passed in from the outside.
        # Let's assume `variables` is available in this scope.
        logits, updated_vars = self.dynamics.apply(
            variables, dyna_inputs, training=False, mutable=['cache']
        )
        # We need the logits for the very last token of the context
        next_token_logits = logits["token_logits"][:, T_context-1, N-1, :]
        
        
        # --- 3. Single Autoregressive Generation Loop (over all tokens) ---
        
        # This function generates one token and updates the state
        def generate_token_body(carry, step_idx):
            # Unpack carry
            tokens, variables, rng = carry
            
            # Figure out which frame (t) and patch (n) we are generating
            t = (T_context * N + step_idx) // N
            n = (T_context * N + step_idx) % N
            
            # --- Reset spatial cache index at the start of a new frame ---
            def reset_spatial_cache_index(leaf):
                # This is a bit brittle, relies on the variable path containing 'cache_index'
                return jnp.array(0, dtype=jnp.int32)
            
            def no_op(leaf):
                return leaf

            is_first_patch = (n == 0)
            
            # Conditionally reset the spatial cache's `cache_index` to 0
            mutable_vars = unfreeze(variables)
            mutable_vars['cache'] = jax.tree_util.tree_map_with_path(
                lambda path, leaf: jax.lax.cond(is_first_patch, reset_spatial_cache_index, no_op, leaf)
                                   if 'spatial_attn' in str(path) and 'cache_index' in str(path)
                                   else leaf,
                mutable_vars['cache']
            )
            variables = freeze(mutable_vars)
            
            # Sample the token
            if 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 / temperature)
            
            # Update the token sequence
            tokens = tokens.at[:, t, n].set(next_token)
            
            # Prepare input for the next step: just the token we generated
            # The model will use its cache for all previous tokens.
            next_input_token = next_token[:, None, None] # Reshape to (B, 1, 1)
            
            # Get the action for the current frame
            current_action = action_tokens[:, t - 1 : t, :]

            next_dyna_inputs = {"video_tokens": next_input_token, "latent_actions": current_action}

            # This call uses and updates the cache, and is very fast!
            next_logits, updated_vars = self.dynamics.apply(
                variables, next_dyna_inputs, training=False, mutable=['cache']
            )
            next_token_logits = next_logits["token_logits"][:, 0, 0, :]
            
            new_carry = (tokens, updated_vars, rng)
            return new_carry, next_token_logits # We don't need to return anything here
        
        # This is where the initial variables must be passed in.
        # The `apply` method requires the full `variables` dict.
        # This is the main change needed in the calling script.
        variables = self.scope.variables # This won't work directly, needs to be passed in.
        initial_carry = (generated_tokens, variables, rng)
        
        # Run the generation loop
        num_steps_to_generate = (seq_len - T_context) * N
        final_carry, _ = jax.lax.scan(
            generate_token_body, initial_carry, jnp.arange(num_steps_to_generate)
        )
        
        final_tokens = final_carry[0]
        
        # --- Decode all tokens at once at the end ---
        final_frames = self.tokenizer.decode(
            final_tokens, video_hw=batch["videos"].shape[2:4]
        )
        return final_frames

    @nn.compact
    def sample_maskgit(
        self,
        batch: Dict[str, Any],
        seq_len: int,
        steps: int = 25,
        temperature: float = 1,
        sample_argmax: bool = False,
    ) -> Any:
        """
        Autoregressively samples up to `seq_len` future frames, following Figure 8 of the paper.

        - Input frames are tokenized once.
        - Future frames are generated autoregressively in token space.
        - All frames are detokenized in a single pass.

        Note:
        - For interactive or step-wise sampling, detokenization should occur after each action.
        - To maintain consistent tensor shapes across timesteps, all current and future frames are decoded at every step.
        - Temporal causal structure is preserved by 
            a) reapplying the mask before each decoding step.
            b) a temporal causal mask is applied within each ST-transformer block.

        Dimension keys:
            B: batch size  
            T: number of input (conditioning) frames  
            N: patches per frame  
            S: sequence length  
            A: action space  
            D: model latent dimension
        """
        # --- Encode videos and actions ---
        tokenizer_out = self.tokenizer.vq_encode(batch["videos"], training=False)
        token_idxs = tokenizer_out["indices"] # (B, T, N)
        B, T, N = token_idxs.shape
        pad_shape = (B, seq_len - T, N)
        pad = jnp.zeros(pad_shape, dtype=token_idxs.dtype)
        token_idxs = jnp.concatenate([token_idxs, pad], axis=1) # (B, S, N)
        action_tokens = self.lam.vq.get_codes(batch["latent_actions"]) 

        MaskGITLoop = nn.scan(
            MaskGITStep,
            variable_broadcast="params",
            split_rngs={"params": False},
            in_axes=0,
            out_axes=0,
            length=steps,
        )
    
        loop_fn = MaskGITLoop(
            dynamics=self.dynamics,
            tokenizer=self.tokenizer,
            temperature=temperature,
            sample_argmax=sample_argmax,
            steps=steps,
        )

        def generation_step_fn(carry, step_t):
            rng, current_token_idxs = carry
            rng, step_rng = jax.random.split(rng)

            # Mask current and future frames (i.e., t >= step_t)
            mask = jnp.arange(seq_len) >= step_t # (S,)
            mask = jnp.broadcast_to(mask[None, :, None], (B, seq_len, N)) # (B, S, N)
            mask = mask.astype(bool)
            masked_token_idxs = current_token_idxs * ~mask

            # --- Initialize and run MaskGIT loop ---
            init_carry_maskgit = (
                step_rng,
                masked_token_idxs,
                mask,
                action_tokens,
            )
            final_carry_maskgit, _ = loop_fn(init_carry_maskgit, jnp.arange(steps))
            updated_token_idxs = final_carry_maskgit[1]
            new_carry = (rng, updated_token_idxs)
            return new_carry, None

        # --- Run the autoregressive generation using scan ---
        initial_carry = (batch["rng"], token_idxs)
        timesteps_to_scan = jnp.arange(T, seq_len)
        final_carry, _ = jax.lax.scan(
            generation_step_fn,
            initial_carry,
            timesteps_to_scan
        )
        final_token_idxs = final_carry[1]

        # --- Decode all tokens at once at the end ---
        final_frames = self.tokenizer.decode(
            final_token_idxs,
            video_hw=batch["videos"].shape[2:4],
        )
        return final_frames

    def vq_encode(self, batch, training) -> Dict[str, Any]:
        # --- Preprocess videos ---
        lam_output = self.lam.vq_encode(batch["videos"], training=training)
        return lam_output["indices"]


class MaskGITStep(nn.Module):
    dynamics: nn.Module
    tokenizer: nn.Module
    temperature: float
    sample_argmax: bool
    steps: int

    @nn.compact
    def __call__(self, carry, x):
        rng, token_idxs, mask, action_tokens = carry
        step = x
        N = token_idxs.shape[2]

        # --- Construct + encode video ---
        vid_embed = self.dynamics.patch_embed(token_idxs) # (B, S, N, D)
        mask_token = self.dynamics.mask_token  # (1, 1, 1, D,)
        mask_expanded = mask[..., None] # (B, S, N, 1) 
        vid_embed = jnp.where(mask_expanded, mask_token, vid_embed)

        # --- Predict transition ---
        act_embed = self.dynamics.action_up(action_tokens)
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        unmasked_ratio = jnp.cos(jnp.pi * (step + 1) / (self.steps * 2))
        step_temp = self.temperature * (1.0 - unmasked_ratio)
        final_logits = self.dynamics.dynamics(vid_embed) / step_temp

        # --- Sample new tokens for final frame ---
        if self.sample_argmax:
            sampled_token_idxs = jnp.argmax(final_logits, axis=-1)
        else:
            rng, _rng = jax.random.split(rng)
            sampled_token_idxs = jax.random.categorical(_rng, final_logits)
        gather_fn = jax.vmap(jax.vmap(jax.vmap(lambda x, y: x[y])))
        final_token_probs = gather_fn(jax.nn.softmax(final_logits), sampled_token_idxs)
        final_token_probs += ~mask
        # Update masked tokens only
        token_idxs = jnp.where(mask, sampled_token_idxs, token_idxs)

        # --- Update mask ---
        num_unmasked_tokens = jnp.round(N * (1.0 - unmasked_ratio)).astype(int)
        idx_mask = jnp.arange(final_token_probs.shape[-1]) > num_unmasked_tokens
        sorted_idxs = jnp.argsort(final_token_probs, axis=-1, descending=True)
        mask_update_fn = jax.vmap(lambda msk, ids: msk.at[ids].set(idx_mask))
        new_mask = mask_update_fn(mask, sorted_idxs)

        new_carry = (rng, token_idxs, new_mask, action_tokens)
        return new_carry, None

def restore_genie_components(
    train_state: TrainState,
    sharding: jax.sharding.NamedSharding,
    grain_iterator: grain.DataLoaderIterator,
    inputs: Dict[str, jax.Array],
    rng: jax.Array,
    args,
):
    """Restore pre-trained Genie components"""
    rng, _rng = jax.random.split(rng)

    # dummy values since we only use tx to initialize the dummy train states
    dummy_tx = optax.adamw(
        learning_rate=optax.constant_schedule(args.max_lr),
        b1=0.9,
        b2=0.9,
        weight_decay=1e-4,
    )
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler)
    

    checkpoint_options = ocp.CheckpointManagerOptions(
        step_format_fixed_length=6,
    )
    tokenizer_checkpoint_manager = ocp.CheckpointManager(
        directory=args.tokenizer_checkpoint,
        options=checkpoint_options,
        handler_registry=handler_registry,
    )
    dummy_tokenizer = TokenizerVQVAE(
        in_dim=args.image_channels,
        model_dim=args.tokenizer_dim,
        latent_dim=args.latent_patch_dim,
        num_latents=args.num_patch_latents,
        patch_size=args.patch_size,
        num_blocks=args.tokenizer_num_blocks,
        num_heads=args.tokenizer_num_heads,
        dropout=args.dropout,
        codebook_dropout=args.dropout,
        param_dtype=args.param_dtype,
        dtype=args.dtype,
        use_flash_attention=args.use_flash_attention,
    )
    tokenizer_init_params = dummy_tokenizer.init(_rng, inputs)
    dummy_tokenizer_train_state = TrainState.create(
        apply_fn=dummy_tokenizer.apply, params=tokenizer_init_params, tx=dummy_tx
    )
    abstract_sharded_tokenizer_state = _create_abstract_sharded_pytree(
        dummy_tokenizer_train_state, sharding
    )
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
        step=tokenizer_checkpoint_manager.latest_step(),
        args=ocp.args.Composite(
            model_state=ocp.args.StandardRestore(abstract_sharded_tokenizer_state),
        ),
    )["model_state"]
    restored_tokenizer_params = restored_tokenizer.params["params"]
    train_state.params["params"]["tokenizer"].update(restored_tokenizer_params)
    tokenizer_checkpoint_manager.close()

    if args.lam_checkpoint:
        lam_checkpoint_manager = ocp.CheckpointManager(
            directory=args.lam_checkpoint,
            options=checkpoint_options,
            handler_registry=handler_registry,
        )
        dummy_lam = LatentActionModel(
            in_dim=args.image_channels,
            model_dim=args.lam_dim,
            latent_dim=args.latent_patch_dim,
            num_latents=args.num_latent_actions,
            patch_size=args.lam_patch_size,
            num_blocks=args.lam_num_blocks,
            num_heads=args.lam_num_heads,
            dropout=args.dropout,
            codebook_dropout=args.dropout,
            param_dtype=args.param_dtype,
            dtype=args.dtype,
            use_flash_attention=args.use_flash_attention,
        )
        lam_init_params = dummy_lam.init(_rng, inputs)
        dummy_lam_train_state = TrainState.create(
            apply_fn=dummy_lam.apply, params=lam_init_params, tx=dummy_tx
        )
        abstract_sharded_lam_state = _create_abstract_sharded_pytree(
            dummy_lam_train_state, sharding
        )
        restored_lam = lam_checkpoint_manager.restore(
            step=lam_checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_sharded_lam_state),
            ),
        )["model_state"]
        restored_lam_params = restored_lam.params["params"]
        # Genie does not initialize all LAM modules, thus we omit those extra modules during restoration
        # (f.srambical) FIXME: Currently, this is a small HBM memory crunch since the LAM's decoder is loaded into HBM and immediately dicarded.
        # A workaround would be to restore to host memory first, and only move the weights to HBM after pruning the decoder
        restored_lam_params = {
            k: v
            for k, v in restored_lam_params.items()
            if k in train_state.params["params"]["lam"]
        }
        train_state.params["params"]["lam"].update(restored_lam_params)
        lam_checkpoint_manager.close()

    return train_state

def _create_abstract_sharded_pytree(pytree_template, sharding_spec):
    """Replaces arrays in a pytree with ShapeDtypeStructs having the given sharding."""

    def map_fn(leaf_template):
        if hasattr(leaf_template, "shape") and hasattr(leaf_template, "dtype"):
            return jax.ShapeDtypeStruct(
                leaf_template.shape, leaf_template.dtype, sharding=sharding_spec
            )
        return leaf_template

    return jax.tree_util.tree_map(map_fn, pytree_template)