[tab] diff.diff (diff)
diff --git a/genie.py b/genie.py
index 0e66676..384c84c 100644
--- a/genie.py
+++ b/genie.py
@@ -7,7 +7,7 @@ import flax.linen as nn
 from flax.training.train_state import TrainState
 import orbax.checkpoint as ocp
 
-from models.dynamics import DynamicsMaskGIT
+from models.dynamics import DynamicsMaskGIT, DynamicsAutoregressive
 from models.lam import LatentActionModel
 from models.tokenizer import TokenizerVQVAE
 
@@ -38,6 +38,7 @@ class Genie(nn.Module):
     dyna_dim: int
     dyna_num_blocks: int
     dyna_num_heads: int
+    use_maskgit: bool
     param_dtype: jnp.dtype
     dtype: jnp.dtype
     dropout: float = 0.0
@@ -70,16 +71,28 @@ class Genie(nn.Module):
             param_dtype=self.param_dtype,
             dtype=self.dtype,
         )
-        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,
-        )
+
+        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,
+            ) 
+        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,
+            )
 
     def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
         tokenizer_outputs = self.tokenizer.vq_encode(batch["videos"], training=False)
@@ -103,8 +116,81 @@ class Genie(nn.Module):
         outputs["lam_indices"] = lam_outputs["indices"]
         return outputs
 
+
+    def sample_causal(
+        self,
+        batch: Dict[str, Any],
+        seq_len: int,
+        temperature: float = 1,
+        sample_argmax: bool = False,
+    ):
+        """
+        Autoregressively samples up to `seq_len` future frames using the causal transformer backend.
+
+        - Input frames are tokenized once.
+        - Future frames are generated one at a time, each conditioned on all previous frames.
+        - All frames are detokenized in a single pass at the end.
+
+        Args:
+            batch: Dict with at least "videos" (B, T, H, W, C)
+            seq_len: total number of frames to generate (including context)
+            temperature: sampling temperature
+            sample_argmax: if True, use argmax instead of sampling
+
+        Returns:
+            Generated video frames (B, seq_len, H, W, C)
+        """
+        # --- Encode context frames ---
+        tokenizer_out = self.tokenizer.vq_encode(batch["videos"], training=False)
+        token_idxs = tokenizer_out["indices"]  # (B, T, N)
+        B, T, N = token_idxs.shape
+
+        # jax.debug.print("token_idxs shape: {}", token_idxs.shape)
+        # --- Prepare initial token sequence ---
+        # Pad with zeros for future frames
+        pad_shape = (B, 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 = self.lam.vq.get_codes(batch["latent_actions"]) # (B, S-1, )
+        # --- Autoregressive generation loop ---
+        rng = batch["rng"]
+        for t in range(T, seq_len):
+            for n in range(300):
+                dyna_inputs = {
+                    "video_tokens": token_idxs_full,
+                    "latent_actions": action_tokens
+                    }
+                # jax.debug.print("token_idxs_full 0: {}", token_idxs_full[0,:,0])
+                dyna_outputs = self.dynamics(dyna_inputs, training=False)
+                # # dyna_outputs["token_logits"]: (B, t, N, vocab_size)
+                # # We want the logits for the last time step (frame t-1 predicting t)
+                # jax.debug.breakpoint()
+                next_token_logits = dyna_outputs["token_logits"][:, t, n, :].astype(jnp.float32)  # (B, 1, vocab_size)
+
+                # Sample or argmax for each patch
+                if sample_argmax:
+                    next_token = jnp.argmax(next_token_logits, axis=-1)  # (B, 1)
+                else:
+                    rng, step_rng = jax.random.split(rng)
+                    next_token = jax.random.categorical(
+                        step_rng, next_token_logits / temperature, axis=-1
+                    )  # (B, 1)
+
+                # Insert the generated tokens into the sequence
+                token_idxs_full = token_idxs_full.at[:, t, n].set(next_token)
+
+        # --- Decode all tokens at once at the end ---
+        final_frames = self.tokenizer.decode(
+            token_idxs_full, video_hw=batch["videos"].shape[2:4]
+        )
+        return final_frames
+
+
     @nn.compact
-    def sample(
+    def sample_maskgit(
         self,
         batch: Dict[str, Any],
         seq_len: int,
@@ -141,7 +227,7 @@ class Genie(nn.Module):
         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"])
+        action_tokens = self.lam.vq.get_codes(batch["latent_actions"]) 
 
         MaskGITLoop = nn.scan(
             MaskGITStep,
diff --git a/models/dynamics.py b/models/dynamics.py
index 8b183dc..7d184d8 100644
--- a/models/dynamics.py
+++ b/models/dynamics.py
@@ -3,6 +3,7 @@ from typing import Dict, Any
 import jax
 import jax.numpy as jnp
 import flax.linen as nn
+import einops
 
 from utils.nn import STTransformer
 
@@ -28,6 +29,7 @@ class DynamicsMaskGIT(nn.Module):
             self.dropout,
             self.param_dtype,
             self.dtype,
+            spacial_bert=True,
         )
         self.patch_embed = nn.Embed(self.num_latents, self.model_dim)
         self.mask_token = self.param(
@@ -58,3 +60,47 @@ class DynamicsMaskGIT(nn.Module):
         vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
         logits = self.dynamics(vid_embed)
         return dict(token_logits=logits, mask=mask)
+
+
+class DynamicsAutoregressive(nn.Module):
+    """Autoregressive (causal) dynamics model"""
+
+    model_dim: int
+    num_latents: int
+    num_blocks: int
+    num_heads: int
+    dropout: float
+    param_dtype: jnp.dtype
+    dtype: jnp.dtype
+
+    def setup(self):
+        self.dynamics = STTransformer(
+            self.model_dim,
+            self.num_latents,
+            self.num_blocks,
+            self.num_heads,
+            self.dropout,
+            self.param_dtype,
+            self.dtype,
+            spacial_bert=False,
+        )
+        self.patch_embed = nn.Embed(self.num_latents, self.model_dim)
+        self.action_up = nn.Dense(
+            self.model_dim,
+            param_dtype=self.param_dtype,
+            dtype=self.dtype,
+        )
+
+    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
+        vid_embed = self.patch_embed(batch["video_tokens"])
+        act_embed = self.action_up(batch["latent_actions"])
+        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
+        # vid_embed_padded = jnp.pad(vid_embed, ((0, 0), (1, 0), (1, 0), (0, 0)))
+        # logits = self.dynamics(vid_embed_padded)[:, :-1, :-1]
+        vid_embed_padded = jnp.pad(vid_embed, ((0, 0), (0, 0), (1, 0), (0, 0)))
+        logits = self.dynamics(vid_embed_padded)[:, :, :-1]
+        mask = jnp.ones(vid_embed.shape[:-1])
+        # next_tokens = jnp.argmax(logits, axis=-1)
+        # print(next_tokens.shape)
+        # jax.debug.breakpoint()
+        return dict(token_logits=logits, mask=mask)
\ No newline at end of file
diff --git a/models/lam.py b/models/lam.py
index cb3cc00..86a1b53 100644
--- a/models/lam.py
+++ b/models/lam.py
@@ -85,6 +85,7 @@ class LatentActionModel(nn.Module):
         B, T = videos.shape[:2]
         patches = patchify(videos, self.patch_size)
         action_pad = jnp.broadcast_to(self.action_in, (B, T, 1, self.patch_token_dim))
+        # FIXME mihir do this the other way around
         padded_patches = jnp.concatenate((action_pad, patches), axis=2)
 
         # --- Encode ---
diff --git a/sample.py b/sample.py
index 4ef592d..3048b3a 100644
--- a/sample.py
+++ b/sample.py
@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from typing import Optional
 import time
 import os
 
@@ -8,7 +9,10 @@ import jax
 import jax.numpy as jnp
 import flax.linen as nn
 import numpy as np
-from orbax.checkpoint import PyTreeCheckpointer
+from flax.training.train_state import TrainState
+import grain
+import orbax.checkpoint as ocp
+import optax
 from PIL import Image, ImageDraw
 import tyro
 
@@ -26,6 +30,7 @@ class Args:
     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
@@ -46,6 +51,7 @@ class Args:
     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
@@ -74,11 +80,12 @@ genie = Genie(
     lam_patch_size=args.lam_patch_size,
     lam_num_blocks=args.lam_num_blocks,
     lam_num_heads=args.lam_num_heads,
-    lam_co_train=False,
+    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,
 )
@@ -90,12 +97,40 @@ dummy_inputs = dict(
 )
 rng, _rng = jax.random.split(rng)
 params = genie.init(_rng, dummy_inputs)
-ckpt = PyTreeCheckpointer().restore(args.checkpoint)["model"]["params"]["params"]
-params["params"].update(ckpt)
+
+dummy_train_state = TrainState.create(
+    apply_fn=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
 
 
 def _sampling_wrapper(module, batch):
-    return module.sample(batch, args.seq_len, args.maskgit_steps, args.temperature, args.sample_argmax)
+    # 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 ---
 def _autoreg_sample(rng, video_batch, action_batch):
@@ -109,31 +144,42 @@ def _autoreg_sample(rng, video_batch, action_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 ---
-array_record_files = [
-    os.path.join(args.data_dir, x)
-    for x in os.listdir(args.data_dir)
-    if x.endswith(".array_record")
-]
-dataloader = get_dataloader(
-    array_record_files,
-    args.seq_len,
-    args.batch_size,
-    args.image_height,
-    args.image_width,
-    args.image_channels,
-    num_workers=8,
-    prefetch_buffer_size=1,
-    seed=args.seed,
-)
-video_batch = next(iter(dataloader))
+# 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)
+batch = dict(videos=video_batch[:,:args.seq_len])
 action_batch = genie.apply(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()
diff --git a/train_dynamics.py b/train_dynamics.py
index 3865cb3..4dd49a3 100644
--- a/train_dynamics.py
+++ b/train_dynamics.py
@@ -62,6 +62,7 @@ class Args:
     dyna_num_heads: int = 8
     dropout: float = 0.0
     mask_limit: float = 0.5
+    use_maskgit: bool = False
     param_dtype: jnp.dtype = jnp.float32
     dtype: jnp.dtype = jnp.bfloat16
     # Logging
@@ -93,13 +94,21 @@ def dynamics_loss_fn(params, state, inputs):
     )
     mask = outputs["mask"]
     outputs["token_logits"] = outputs["token_logits"].astype(jnp.float32)
+    logits = outputs["token_logits"]
+    targets = outputs["video_tokens"]
+
+    # if not args.use_maskgit:
+        # logits = outputs["token_logits"][:, :, :-1]
+        # targets = outputs["video_tokens"][:, :, 1:]
+        # mask = outputs["mask"][:, :, 1:] 
+
     ce_loss = optax.softmax_cross_entropy_with_integer_labels(
-        outputs["token_logits"], outputs["video_tokens"]
+        logits, targets
     )
     ce_loss = (mask * ce_loss).sum() / mask.sum()
-    acc = outputs["token_logits"].argmax(-1) == outputs["video_tokens"]
+    acc = logits.argmax(-1) == targets
     acc = (mask * acc).sum() / mask.sum()
-    select_probs = jax.nn.softmax(outputs["token_logits"])
+    select_probs = jax.nn.softmax(logits)
     gt = inputs["videos"].clip(0, 1).reshape(-1, *inputs["videos"].shape[2:])
     recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
     psnr = pix.psnr(gt, recon).mean() # type: ignore
@@ -115,7 +124,7 @@ def dynamics_loss_fn(params, state, inputs):
     metrics = dict(
         cross_entropy_loss=ce_loss,
         masked_token_accuracy=acc,
-        select_logit=outputs["token_logits"].max(-1).mean(),
+        select_logit=logits.max(-1).mean(),
         select_p=select_probs.max(-1).mean(),
         entropy=jax.scipy.special.entr(select_probs).sum(-1).mean(),
         psnr=psnr,
@@ -180,6 +189,7 @@ if __name__ == "__main__":
         dyna_num_heads=args.dyna_num_heads,
         dropout=args.dropout,
         mask_limit=args.mask_limit,
+        use_maskgit=args.use_maskgit,
         param_dtype=args.param_dtype,
         dtype=args.dtype,
     )
@@ -315,6 +325,9 @@ if __name__ == "__main__":
     dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_iterator)  # type: ignore
     while step < args.num_steps:
         for videos in dataloader:
+        # videos = np.load("overfit_dir/corner_8repl.npy")
+        # videos = jax.make_array_from_process_local_data(videos_sharding, videos)
+        # while True:
             # --- Train step ---
             rng, _rng, _rng_dropout, _rng_mask = jax.random.split(rng, 4)
 
diff --git a/utils/nn.py b/utils/nn.py
index b7bec9f..cf53092 100644
--- a/utils/nn.py
+++ b/utils/nn.py
@@ -26,6 +26,112 @@ class PositionalEncoding(nn.Module):
         x = x + self.pe[: x.shape[2]]
         return x
 
+# class STBlock2(nn.Module):
+    # dim: int
+    # num_heads: int
+    # dropout: float
+    # param_dtype: jnp.dtype
+    # dtype: jnp.dtype
+
+    # @nn.remat
+    # @nn.compact
+    # def __call__(self, x: jax.Array) -> jax.Array:
+        # # --- Spatial attention ---
+        # z = PositionalEncoding(self.dim)(x)
+        # z = nn.LayerNorm(
+            # param_dtype=self.param_dtype,
+            # dtype=self.dtype,
+        # )(z)
+        # causal_mask = jnp.tri(z.shape[-2])
+        # z = nn.MultiHeadAttention(
+            # num_heads=self.num_heads,
+            # qkv_features=self.dim,
+            # dropout_rate=self.dropout,
+            # param_dtype=self.param_dtype,
+            # dtype=self.dtype,
+        # )(z, mask=causal_mask)
+        # x = x + z
+
+        # # --- Temporal attention ---
+        # x = x.swapaxes(1, 2)
+        # z = PositionalEncoding(self.dim)(x)
+        # z = nn.LayerNorm(
+            # param_dtype=self.param_dtype,
+            # dtype=self.dtype,
+        # )(z)
+        # causal_mask = jnp.tri(z.shape[-2])
+        # z = nn.MultiHeadAttention(
+            # num_heads=self.num_heads,
+            # qkv_features=self.dim,
+            # dropout_rate=self.dropout,
+            # param_dtype=self.param_dtype,
+            # dtype=self.dtype,
+        # )(z, mask=causal_mask)
+        # x = x + z
+        # x = x.swapaxes(1, 2)
+
+        # # --- Feedforward ---
+        # z = nn.LayerNorm(
+            # param_dtype=self.param_dtype,
+            # dtype=self.dtype,
+        # )(x)
+        # # FIXME (f.srambical): Here, the attention hidden dimension is the same as the FFN's. Usually, FFN hidden dimension is 4x model_dim
+        # z = nn.Dense(
+            # self.dim,
+            # param_dtype=self.param_dtype,
+            # dtype=self.dtype,
+        # )(z)
+        # z = nn.gelu(z)
+        # x = x + z
+
+        # return x
+
+# class CausalTransformer(nn.Module):
+    # model_dim: int
+    # out_dim: int
+    # num_blocks: int
+    # num_heads: int
+    # dropout: float
+    # param_dtype: jnp.dtype
+    # dtype: jnp.dtype
+
+    # @nn.compact
+    # def __call__(self, x: jax.Array) -> jax.Array:
+        # # Input projection and normalization
+        # x = nn.Sequential(
+            # [
+                # nn.LayerNorm(
+                    # param_dtype=self.param_dtype,
+                    # dtype=self.dtype,
+                # ),
+                # nn.Dense(self.model_dim,
+                # param_dtype=self.param_dtype,
+                # dtype=self.dtype,
+                # ),
+                # nn.LayerNorm(
+                    # param_dtype=self.param_dtype,
+                    # dtype=self.dtype,
+                # ),
+            # ]
+        # )(x)
+        # # Causal transformer blocks
+        # for _ in range(self.num_blocks):
+            # x = STBlock2(
+                # dim=self.model_dim,
+                # num_heads=self.num_heads,
+                # dropout=self.dropout,
+                # param_dtype=self.param_dtype,
+                # dtype=self.dtype,
+            # )(x)
+
+        # # Output projection
+        # x = nn.Dense(
+            # self.out_dim,
+            # param_dtype=self.param_dtype,
+            # dtype=self.dtype,
+        # )(x)
+        # return x  # (B, T, E)
+
 
 class STBlock(nn.Module):
     dim: int
@@ -33,6 +139,7 @@ class STBlock(nn.Module):
     dropout: float
     param_dtype: jnp.dtype
     dtype: jnp.dtype
+    spatial_bert: bool = True
 
     @nn.remat
     @nn.compact
@@ -43,13 +150,14 @@ class STBlock(nn.Module):
             param_dtype=self.param_dtype,
             dtype=self.dtype,
         )(z)
+        spatial_mask = None if self.spatial_bert else jnp.tri(z.shape[-2])
         z = nn.MultiHeadAttention(
             num_heads=self.num_heads,
             qkv_features=self.dim,
             dropout_rate=self.dropout,
             param_dtype=self.param_dtype,
             dtype=self.dtype,
-        )(z)
+        )(z, mask=spatial_mask)
         x = x + z
 
         # --- Temporal attention ---
@@ -95,6 +203,7 @@ class STTransformer(nn.Module):
     dropout: float
     param_dtype: jnp.dtype
     dtype: jnp.dtype
+    spacial_bert: bool = True
 
     @nn.compact
     def __call__(self, x: jax.Array) -> jax.Array:
@@ -121,6 +230,7 @@ class STTransformer(nn.Module):
                 dropout=self.dropout,
                 param_dtype=self.param_dtype,
                 dtype=self.dtype,
+                spatial_bert=self.spacial_bert,
             )(x)
         x = nn.Dense(
             self.out_dim,