p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] sample.py (python)
    action_batch = genie.vq_encode(batch, training=False)
    action_batch = jnp.asarray(action_batch).reshape(
        video_batch.shape[0], args.seq_len - 1, 1
    )

    # --- Sample + evaluate video ---
    vid = _autoreg_sample(rng, video_batch, action_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 = jnp.asarray(
        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]
[content] sample.py (python)
    action_batch_DL = genie.vq_encode(batch, training=False)
    # FIXME (f.srambical): where is the last L dimension?
    action_batch_BSm11 = jnp.asarray(action_batch_DL).reshape(
        video_batch.shape[0], args.seq_len - 1, 1
    )

    # --- Sample + evaluate video ---
    vid = _autoreg_sample(rng, video_batch, action_batch_BSm11)
    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 = jnp.asarray(
        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_BSm11.shape[0]):
            action = action_batch_BSm11[row, t, 0]
[content] sample.py (python)

        Dimension keys:
            B: batch size
            T: number of input (conditioning) frames
            N: number of patches per frame
            S: sequence length
            H: height
            W: width
            D: B * T * N
            E: B * (T - 1)
[content] genie.py (python)
            # Create a boolean mask where True means the token should remain masked
            # We want to unmask the top num_unmasked_tokens tokens
            mask_update_fn = jax.vmap(jax.vmap(lambda msk, sorted_idx: 
                msk.at[sorted_idx[:num_unmasked_tokens]].set(False)))
[content] genie.py (python)
            # Only use mask token for MaskGIT dynamics
            if self.dyna_type == "maskgit":
                mask_token_111M = self.dynamics.mask_token.value
                mask_expanded_BSN1 = mask_BSN[..., None]
                vid_embed_BSNM = jnp.where(mask_expanded_BSN1, mask_token_111M, vid_embed_BSNM)
[content] genie.py (python)
            mask_token_111M = self.dynamics.mask_token.value
            mask_expanded_BSN1 = mask_BSN[..., None]
            vid_embed_BSNM = jnp.where(mask_expanded_BSN1, mask_token_111M, vid_embed_BSNM)
[content] sample.py (python)
    batch = dict(videos=video_batch_BSHWC)
    action_batch_E = genie.vq_encode(batch, training=False)
    # FIXME (f.srambical): I think we should omit this reshape!
    action_batch_BSm11 = jnp.asarray(action_batch_E).reshape(
        video_batch_BSHWC.shape[0], args.seq_len - 1, 1
    )
[content] models/dynamics.py (python)

class DynamicsCausal(nnx.Module):
    """Causal dynamics model"""

    def __init__(
        self,
        model_dim: int,
        ffn_dim: int,
        num_latents: int,
        latent_action_dim: int,
        num_blocks: int,
        num_heads: int,
        dropout: float,
        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: bool,
        rngs: nnx.Rngs,
    ):
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.num_latents = num_latents
        self.latent_action_dim = latent_action_dim
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention

        self.transformer = Transformer(
            self.model_dim,
            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        self.action_up = nnx.Linear(
            self.latent_action_dim,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    def __call__(
        self, batch: Dict[str, jax.Array], training: bool = True
    ) -> tuple[jax.Array, jax.Array | None]:
        vid_embed = self.patch_embed(batch["video_tokens"])  # (B, T, N, E)
        act_embed = self.action_up(batch["latent_actions"])  # (B, T-1, E)
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))

        vid_embed = einops.rearrange(vid_embed, "b t n e -> b (t n) e")
        logits = self.transformer(vid_embed)
        logits = einops.rearrange(logits, "b (t n) v -> b t n v")

        return logits, jnp.ones_like(batch["video_tokens"])

    def sample(self, batch: Dict[str, jax.Array]) -> jax.Array:
         # FIXME (f.srambical): implement this
        raise NotImplementedError("Sampling not implemented for Transformer dynamics model")