From ec0d8b22023a9a24c979971219f9aa9967aeda47 Mon Sep 17 00:00:00 2001 From: Mihir Mahajan Date: Mon, 7 Jul 2025 20:08:00 +0200 Subject: [PATCH 1/2] feat: 1) genie-like sampling in token-space, 2) consistent shapes to avoid recompile compile, 3) jitted sampling, 4) scan instead of loop --- genie.py | 26 ++++++++++---------------- sample.py | 28 +++++++++++++--------------- 2 files changed, 23 insertions(+), 31 deletions(-) diff --git a/genie.py b/genie.py index 364fba14..50729ab2 100644 --- a/genie.py +++ b/genie.py @@ -146,28 +146,22 @@ class MaskGITStep(nn.Module): @nn.compact def __call__(self, carry, x): - rng, final_token_idxs, mask, token_idxs, action_tokens = carry + rng, token_idxs, mask, action_tokens = carry step = x - B, T, N = token_idxs.shape[:3] + N = token_idxs.shape[2] # --- Construct + encode video --- - vid_token_idxs = jnp.concatenate( - (token_idxs, jnp.expand_dims(final_token_idxs, 1)), axis=1 - ) - vid_embed = self.dynamics.patch_embed(vid_token_idxs) - curr_masked_frame = jnp.where( - jnp.expand_dims(mask, -1), - self.dynamics.mask_token[0], - vid_embed[:, -1], - ) - vid_embed = vid_embed.at[:, -1].set(curr_masked_frame) + 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)[:, -1] / step_temp + final_logits = self.dynamics.dynamics(vid_embed) / step_temp # --- Sample new tokens for final frame --- if self.sample_argmax: @@ -179,11 +173,11 @@ def __call__(self, carry, x): jnp.argmax(final_logits, axis=-1), jax.random.categorical(_rng, final_logits), ) - gather_fn = jax.vmap(jax.vmap(lambda x, y: x[y])) + 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 - new_token_idxs = jnp.where(mask, sampled_token_idxs, final_token_idxs) + token_idxs = jnp.where(mask, sampled_token_idxs, token_idxs) # --- Update mask --- num_unmasked_tokens = jnp.round(N * (1.0 - unmasked_ratio)).astype(int) @@ -192,7 +186,7 @@ def __call__(self, carry, x): 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, new_token_idxs, new_mask, token_idxs, action_tokens) + new_carry = (rng, token_idxs, new_mask, action_tokens) return new_carry, None diff --git a/sample.py b/sample.py index 48866a71..42eeaa20 100644 --- a/sample.py +++ b/sample.py @@ -5,6 +5,7 @@ import einops import jax import jax.numpy as jnp +import flax.linen as nn import numpy as np from orbax.checkpoint import PyTreeCheckpointer from PIL import Image, ImageDraw @@ -85,24 +86,21 @@ class Args: ckpt = PyTreeCheckpointer().restore(args.checkpoint)["model"]["params"]["params"] params["params"].update(ckpt) + +def _sampling_wrapper(module, batch): + return module.sample(batch, args.seq_len, args.maskgit_steps, args.temperature, args.sample_argmax) + # --- 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 + sampling_fn = jax.jit(nn.apply(_sampling_wrapper, genie)) + rng, _rng = jax.random.split(rng) + batch = dict(videos=vid, latent_actions=action_batch, rng=_rng) + generated_vid = sampling_fn( + params, + batch + ) + return generated_vid # --- Get video + latent actions --- From 7fb86d7e27904bcab7784451f436f73f30edaef9 Mon Sep 17 00:00:00 2001 From: Mihir Mahajan Date: Sun, 3 Aug 2025 16:47:07 +0200 Subject: [PATCH 2/2] Refactor Genie autoregressive sampling and MaskGIT mask update logic - Overhauled Genie.sample() to support autoregressive generation of up to `seq_len` future frames, following the approach in Figure 8 of the paper. - Input frames are tokenized once, and future frames are generated autoregressively in token space. - All frames are detokenized in a single pass at the end. - Added detailed docstring explaining the sampling process and tensor dimensions. - Replaced the old MaskGIT loop with a step-wise scan over timesteps, masking and updating tokens for each future frame. - Updated MaskGITStep mask update logic: - Flattened and sorted token probabilities for mask updates using einops. - Changed mask update to operate on the flattened token dimension, then reshaped back. - Fixed mask indexing to ensure correct number of tokens are unmasked at each step. - Added einops import for tensor reshaping. - Minor type and variable name improvements for clarity and correctness. --- genie.py | 93 ++++++++++++++++++++++++++++++++++++++++++++------------ 1 file changed, 73 insertions(+), 20 deletions(-) diff --git a/genie.py b/genie.py index 50729ab2..dfa352c8 100644 --- a/genie.py +++ b/genie.py @@ -4,6 +4,7 @@ import jax import jax.numpy as jnp import flax.linen as nn +import einops from models.dynamics import DynamicsMaskGIT from models.lam import LatentActionModel @@ -87,25 +88,42 @@ def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, An def sample( self, batch: Dict[str, Any], + seq_len: int, steps: int = 25, - temperature: int = 1, + 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"] - new_frame_idxs = jnp.zeros_like(token_idxs)[:, 0] + 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"]) - # --- Initialize MaskGIT --- - init_mask = jnp.ones_like(token_idxs, dtype=bool)[:, 0] - init_carry = ( - batch["rng"], - new_frame_idxs, - init_mask, - token_idxs, - action_tokens, - ) MaskGITLoop = nn.scan( MaskGITStep, variable_broadcast="params", @@ -123,13 +141,45 @@ def sample( sample_argmax=sample_argmax, steps=steps, ) - final_carry, _ = loop_fn(init_carry, jnp.arange(steps)) - new_frame_idxs = final_carry[1] - new_frame_pixels = self.tokenizer.decode( - jnp.expand_dims(new_frame_idxs, 1), + + def generation_step_fn(carry, step_t): + rng, current_token_idxs = carry + rng, step_rng = jax.random.split(rng) + + # Mask current frame (future frames are masked by default using causal mask in ST-transformer) + 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 new_frame_pixels + return final_frames def vq_encode(self, batch, training) -> Dict[str, Any]: # --- Preprocess videos --- @@ -181,10 +231,13 @@ def __call__(self, carry, x): # --- 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) + idx_mask = jnp.arange(final_token_probs.shape[-1]) <= N - num_unmasked_tokens + final_token_probs_flat = einops.rearrange(final_token_probs, "b s n -> b (s n)") + sorted_idxs = jnp.argsort(final_token_probs_flat, axis=-1) mask_update_fn = jax.vmap(lambda msk, ids: msk.at[ids].set(idx_mask)) - new_mask = mask_update_fn(mask, sorted_idxs) + mask_flat = einops.rearrange(mask, "b s n -> b (s n)") + new_mask_flat = mask_update_fn(mask_flat, sorted_idxs) + new_mask = einops.rearrange(new_mask_flat, "b (s n) -> b s n", n=N) new_carry = (rng, token_idxs, new_mask, action_tokens) return new_carry, None