Add MaskGIT sampling, ST-attention, LAM codebook resets (#5)

* Sampling WIP

* Add notebook

* Add sample script

* autoregressive sampling and full sequence prompting

* Maintain sampled tokens in maskgit generation

* autoregressive generation

* Clean up sampling code, add arguments

* Fix new frame shapes in sampling

* Add image logging and rng split

* MaskGitScanning

* adding psnr and interweaved video logging

* Add sweep

* Add ST positional embedding

* Add temp docker runner

* Refactor MaskGIT, fix cosine schedule + token sampling

* Reset inactive latent actions in LAM training

* Sampling code updates, log dynamics model statistics

* Refactor maskgit step

* Remove dev file paths

* Remove sample notebook

* Log gifs and refactor sampling

* Black formatting

* Remove sweep config

---------

Co-authored-by: Aidandos <[email protected]>

Author: EmptyJackson

.gitignore

@@ -0,0 +1,9 @@
+*.pyc
+*.npy
+*.png
+*.gif
+
+wandb_key
+checkpoints/
+wandb/
+__pycache__/

genie.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any, Optional
+from typing import Dict, Any
 
 from orbax.checkpoint import PyTreeCheckpointer
 import jax
@@ -32,8 +32,8 @@ class Genie(nn.Module):
     dyna_dim: int
     dyna_num_blocks: int
     dyna_num_heads: int
-    dropout: float
-    mask_limit: float
+    dropout: float = 0.0
+    mask_limit: float = 0.0
 
     def setup(self):
         self.tokenizer = TokenizerVQVAE(
@@ -83,19 +83,125 @@ def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, An
         )
         return outputs
 
+    @nn.compact
+    def sample(
+        self,
+        batch: Dict[str, Any],
+        steps: int = 25,
+        temperature: int = 1,
+        sample_argmax: bool = False,
+    ) -> Any:
+        # --- 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]
+        action_tokens = self.lam.vq.get_codes(batch["latent_actions"])
 
-def restore_genie_checkpoint(
-    params: Dict[str, Any], tokenizer: str, lam: str, dyna: Optional[str] = None
-):
+        # --- 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",
+            split_rngs={"params": False},
+            in_axes=0,
+            out_axes=0,
+            length=steps,
+        )
+
+        # --- Run MaskGIT loop ---
+        loop_fn = MaskGITLoop(
+            dynamics=self.dynamics,
+            tokenizer=self.tokenizer,
+            temperature=temperature,
+            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),
+            video_hw=batch["videos"].shape[2:4],
+        )
+        return new_frame_pixels
+
+    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, final_token_idxs, mask, token_idxs, action_tokens = carry
+        step = x
+        B, T, N = token_idxs.shape[:3]
+
+        # --- 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)
+
+        # --- 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
+
+        # --- 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 = jnp.where(
+                step == self.steps - 1,
+                jnp.argmax(final_logits, axis=-1),
+                jax.random.categorical(_rng, final_logits),
+            )
+        gather_fn = 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)
+
+        # --- 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, new_token_idxs, new_mask, token_idxs, action_tokens)
+        return new_carry, None
+
+
+def restore_genie_components(params: Dict[str, Any], tokenizer: str, lam: str):
     """Restore pre-trained Genie components"""
     params["params"]["tokenizer"].update(
         PyTreeCheckpointer().restore(tokenizer)["model"]["params"]["params"]
     )
     params["params"]["lam"].update(
         PyTreeCheckpointer().restore(lam)["model"]["params"]["params"]
     )
-    if dyna:
-        params["params"]["dyna"].update(
-            PyTreeCheckpointer().restore(dyna)["model"]["params"]["params"]
-        )
     return params

models/lam.py

@@ -74,7 +74,7 @@ def vq_encode(self, videos: Any, training: bool = True) -> Dict[str, Any]:
         # --- Encode ---
         z = self.encoder(padded_patches)  # (B, T, N, E)
         # Get latent action for all future frames
-        z = z[:, 1:, 0]  # (B, T-1, 1, E)
+        z = z[:, 1:, 0]  # (B, T-1, E)
 
         # --- Vector quantize ---
         z = z.reshape(B * (T - 1), self.latent_dim)

run_docker.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+gpu=$1
+script_and_args="${@:2}"
+WANDB_API_KEY=$(cat ./docker/wandb_key)
+git pull
+
+echo "Launching container jafar_$gpu on GPU $gpu"
+docker run \
+    --env CUDA_VISIBLE_DEVICES=$gpu \
+    --gpus all \
+    -e WANDB_API_KEY=$WANDB_API_KEY \
+    -v $(pwd):/home/duser/jafar \
+    --name jafar\_$gpu \
+    --user $(id -u) \
+    --rm \
+    -d \
+    jafar \
+    /bin/bash -c "$script_and_args"

sample.py

@@ -0,0 +1,141 @@
+from dataclasses import dataclass
+import time
+
+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 data.dataloader import get_dataloader
+from genie import Genie
+
+
+@dataclass
+class Args:
+    # Experiment
+    seed: int = 0
+    seq_len: int = 16
+    image_channels: int = 3
+    image_resolution: int = 64
+    file_path: str = "data/coinrun.npy"
+    checkpoint: str = ""
+    # Sampling
+    batch_size: int = 1
+    maskgit_steps: int = 25
+    temperature: float = 1.0
+    sample_argmax: bool = False
+    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 = 8
+    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_resolution, args.image_resolution, 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)
+
+# --- Get video + latent actions ---
+dataloader = get_dataloader(args.file_path, args.seq_len, args.batch_size)
+for vids in dataloader:
+    video_batch = jnp.array(vids, dtype=jnp.float32) / 255.0
+    break
+batch = dict(videos=video_batch)
+lam_output = genie.apply(params, batch, False, method=Genie.vq_encode)
+lam_output = lam_output.reshape(args.batch_size, args.seq_len - 1, 1)
+
+
+# --- Define autoregressive sampling loop ---
+def _autoreg_sample(rng, video_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=lam_output[:, :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
+
+
+# --- Sample + evaluate video ---
+vid = _autoreg_sample(rng, video_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}")
+
+# --- Save generated video ---
+original_frames = (video_batch * 255).astype(np.uint8)
+interweaved_frames = np.zeros((vid.shape[0] * 2, *vid.shape[1:5]), dtype=np.uint8)
+interweaved_frames[0::2] = original_frames[:, : vid.shape[1]]
+interweaved_frames[1::2] = (vid * 255).astype(np.uint8)
+flat_vid = einops.rearrange(interweaved_frames, "n t h w c -> t h (n w) c")
+imgs = [Image.fromarray(img) for img in flat_vid]
+for img, action in zip(imgs[1:], lam_output[0, :, 0]):
+    d = ImageDraw.Draw(img)
+    d.text((2, 2), f"{action}", fill=255)
+imgs[0].save(
+    f"generation_{time.time()}.gif",
+    save_all=True,
+    append_images=imgs[1:],
+    duration=250,
+    loop=0,
+)

train_dynamics.py

@@ -10,11 +10,11 @@
 import numpy as np
 import jax
 import jax.numpy as jnp
-import wandb
 import tyro
+import wandb
 
-from genie import Genie, restore_genie_checkpoint
 from data.dataloader import get_dataloader
+from genie import Genie, restore_genie_components
 
 ts = int(time.time())
 
@@ -103,7 +103,7 @@ class Args:
 )
 rng, _rng = jax.random.split(rng)
 init_params = genie.init(_rng, dummy_inputs)
-init_params = restore_genie_checkpoint(
+init_params = restore_genie_components(
     init_params, args.tokenizer_checkpoint, args.lam_checkpoint
 )
 lr_schedule = optax.warmup_cosine_decay_schedule(
@@ -113,8 +113,9 @@ class Args:
 train_state = TrainState.create(apply_fn=genie.apply, params=init_params, tx=tx)
 
 
+# --- Define dynamics loss + train step ---
 def dynamics_loss_fn(params, state, inputs):
-    # --- Compute masked loss ---
+    """Compute masked dynamics loss"""
     outputs = state.apply_fn(
         params, inputs, training=True, rngs={"dropout": inputs["dropout_rng"]}
     )
@@ -125,13 +126,20 @@ def dynamics_loss_fn(params, state, inputs):
     ce_loss = (mask * ce_loss).sum() / mask.sum()
     acc = outputs["token_logits"].argmax(-1) == outputs["video_tokens"]
     acc = (mask * acc).sum() / mask.sum()
-    metrics = dict(cross_entropy_loss=ce_loss, masked_token_accuracy=acc)
+    select_probs = jax.nn.softmax(outputs["token_logits"])
+    metrics = dict(
+        cross_entropy_loss=ce_loss,
+        masked_token_accuracy=acc,
+        select_logit=outputs["token_logits"].max(-1).mean(),
+        select_p=select_probs.max(-1).mean(),
+        entropy=jax.scipy.special.entr(select_probs).sum(-1).mean(),
+    )
     return ce_loss, (outputs["recon"], metrics)
 
 
-# --- Define train step ---
 @jax.jit
 def train_step(state, inputs):
+    """Update state and compute metrics"""
     grad_fn = jax.value_and_grad(dynamics_loss_fn, has_aux=True, allow_int=True)
     (loss, (recon, metrics)), grads = grad_fn(state.params, state, inputs)
     state = state.apply_gradients(grads=grads)