p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] jasmine/test_dit.py (python)
from models.dynamics import DynamicsDiffusion
import jax.numpy as jnp
import jax
from flax import nnx

if __name__ == "__main__":
    model = DynamicsDiffusion(
        model_dim=128,
        ffn_dim=128,
        num_latents=128,
        latent_action_dim=128,
        num_blocks=128,
        num_heads=128,
        dropout=0.0,
        param_dtype=jnp.float32,
        dtype=jnp.float32,
        use_flash_attention=False,
        rngs=nnx.Rngs(params=0, dropout=jax.random.key(1))
    )
    batch = {
        "images": jnp.zeros((4, 10, 64, 64, 3)),
        "rng": jax.random.key(0),
    }
    x_t, mask = model(batch)
    print(x_t.shape)
    print(mask.shape)
[content] jasmine/utils/nn.py (python)
        self.temporal_attention = nnx.MultiHeadAttention(
            num_heads=self.num_heads,
            in_features=self.model_dim,
            qkv_features=self.model_dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(
                self.use_flash_attention, is_causal=True
            ),
            rngs=rngs,
            decode=self.decode,
        )
        self.spatial_attention = nnx.MultiHeadAttention(
            num_heads=self.num_heads,
            in_features=self.model_dim,
            qkv_features=self.model_dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(
                self.use_flash_attention, is_causal=True
            ),
            rngs=rngs,
            decode=self.decode,
        )
[content] jasmine/utils/nn.py (python)

        # --- Spatial attention ---
        B, T, N, M = x_BTNM.shape
        z_FNM = einops.rearrange(x_BTNM, "b t n m -> (b t) n m")
        z_FNM = self.spatial_norm(z_FNM)
        z_FNM = self.spatial_attention(z_FNM, sow_weights=self.sow_weights)
        z_BTNM = einops.rearrange(z_FNM, "(b t) n m -> b t n m", t=T)
        x_BTNM = x_BTNM + z_BTNM
        # --- Temporal attention ---
        z_PTM = einops.rearrange(x_BTNM, "b t n m -> (b n) t m")
        z_PTM = self.temporal_norm(z_PTM)
        z_PTM = self.temporal_attention(z_PTM, sow_weights=self.sow_weights)
        z_BTNM = einops.rearrange(z_PTM, "(b n) t m -> b t n m", n=N)
        x_BTNM = x_BTNM + z_BTNM
        # --- Feedforward ---
        z_BTNM = self.ffn_norm(x_BTNM)
        z_BTND = self.ffn_dense1(z_BTNM)
        z_BTND = jax.nn.gelu(z_BTND)
        z_BTNM = self.ffn_dense2(z_BTND)
        x_BTNM = x_BTNM + z_BTNM
        if self.sow_activations:
            self.sow(nnx.Intermediate, "activations", x_BTNM)

        return x_BTNM
[content] jasmine/utils/nn.py (python)
        # channels_per_head = self.model_dim // self.num_heads
        # k_BTNM = self.K_proj(x_BTNM)
        # q_BTNM = self.Q_proj(x_BTNM)
        # v_BTNM = self.V_proj(x_BTNM)
        # k_BTNM = jnp.reshape(k_BTNM, (B, T, self.num_heads, channels_per_head))
        # q_BTNM = jnp.reshape(q_BTNM, (B, T, self.num_heads, channels_per_head))
        # v_BTNM = jnp.reshape(v_BTNM, (B, T, self.num_heads, channels_per_head))
        # q_BTNM = q_BTNM / q_BTNM.shape[3] # (1/d) scaling.
        # w_BTNM = jnp.einsum('bqhc,bkhc->bhqk', q_BTNM, k_BTNM) # [B, HW, HW, num_heads]
        # w_BTNM = w_BTNM.astype(jnp.float32)
        # w_BTNM = nnx.softmax(w_BTNM, axis=-1)
        # y_BTNM = jnp.einsum('bhqk,bkhc->bqhc', w_BTNM, v_BTNM) # [B, HW, num_heads, channels_per_head]
        # y_BTNM = jnp.reshape(y_BTNM, x_BTNM.shape) # [B, H, W, C] (C = heads * channels_per_head)
        # attn_x_BTNM = self.O_proj(y_BTNM)
        # x_BTNM = x_BTNM + (gate_msa[:, None] * attn_x_BTNM)

        # MLP Residual.
        x_norm2 = self.layer_norm2(x_BTNM)
        x_modulated2 = modulate(x_norm2, shift_mlp, scale_mlp)

        mlp_x = self.dense1(x_modulated2)
        mlp_x = nnx.gelu(mlp_x)
        mlp_x = self.dense2(mlp_x)
        x = x + (gate_mlp[:, None] * mlp_x)
        return x
[content] jasmine/utils/nn.py (python)

        self.temporal_norm = nnx.LayerNorm(
            num_features=self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.param_dtype,  # layer norm in full precision
            rngs=rngs,
        )
        self.spatial_norm = nnx.LayerNorm(
            num_features=self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.param_dtype,  # layer norm in full precision
            rngs=rngs,
        )
        self.ffn_norm = nnx.LayerNorm(
            num_features=self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.param_dtype,  # layer norm in full precision
            rngs=rngs,
        )
[content] jasmine/utils/nn.py (python)

        self.ffn_dense1 = nnx.Linear(
            in_features=self.model_dim,
            out_features=self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )
        self.ffn_dense2 = nnx.Linear(
            in_features=self.ffn_dim,
            out_features=self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )
[content] jasmine/utils/nn.py (python)

        scale_spatial = jnp.tile(scale_spatial, (N, 1))
        gate_spatial = jnp.tile(gate_spatial, (N, 1))
        shift_temporal = jnp.tile(shift_temporal, (N, 1))
        scale_temporal = jnp.tile(scale_temporal, (N, 1))
        gate_temporal = jnp.tile(gate_temporal, (N, 1))
        shift_mlp = jnp.tile(shift_mlp, (N, 1))
        scale_mlp = jnp.tile(scale_mlp, (N, 1))
        gate_mlp = jnp.tile(gate_mlp, (N, 1))
[content] jasmine/genie.py (python)

            self.dynamics = DynamicsDiffusion(
                model_dim=self.dyna_dim,
                ffn_dim=self.dyna_ffn_dim,
                num_latents=self.num_patch_latents,
                latent_action_dim=self.latent_action_dim,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
            )
[content] jasmine/genie.py (python)

                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                decode=decode,
                rngs=rngs,
[content] jasmine/genie.py (python)

        outputs = dict(
            video_tokens=jax.lax.stop_gradient(token_indices_BTN),
            latent_actions=(
                action_embeddings_BTm11L
                if self.use_gt_actions
                else latent_actions_BTm11L
            ),
        )
[content] jasmine/train_dynamics.py (python)
    dyna_type: str = "maskgit"  # supported options: maskgit, causal
    dyna_dim: int = 512
    dyna_ffn_dim: int = 2048
    dyna_num_blocks: int = 6
    dyna_num_heads: int = 8
    dropout: float = 0.0
    mask_limit: float = 0.5
    z_loss_weight: float = 0.0
    param_dtype = jnp.float32
    dtype = jnp.bfloat16
    use_flash_attention: bool = True
    use_gt_actions: bool = False
    # Logging
    log: bool = True
    entity: str = ""
    project: str = ""
    name: str = "train_dynamics"
    tags: list[str] = field(default_factory=lambda: ["dynamics"])
    log_interval: int = 50
    log_image_interval: int = 1000
    ckpt_dir: str = ""
    log_checkpoint_interval: int = 5000
    log_checkpoint_keep_period: int = 20_000
    log_gradients: bool = False
    val_data_dir: str = ""
    val_interval: int = 20_000
    val_steps: int = 50
    eval_full_frame: bool = True
    val_maskgit_steps: int = 25
    val_temperature: float = 1
    val_sample_argmax: bool = False
    wandb_id: str = ""


def build_model(args: Args, rng: jax.Array) -> tuple[Genie, jax.Array]:
    rng, _rng = jax.random.split(rng)
    rngs = nnx.Rngs(_rng)
    genie = Genie(
        # Tokenizer
        in_dim=args.image_channels,
        tokenizer_dim=args.tokenizer_dim,
        tokenizer_ffn_dim=args.tokenizer_ffn_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,
        lam_ffn_dim=args.lam_ffn_dim,
        latent_action_dim=args.latent_action_dim,
        num_actions=args.num_actions,
        lam_patch_size=args.lam_patch_size,
        lam_num_blocks=args.lam_num_blocks,
        lam_num_heads=args.lam_num_heads,
        lam_co_train=not args.lam_checkpoint,
        use_gt_actions=args.use_gt_actions,
        # Dynamics
        dyna_type=args.dyna_type,
        dyna_dim=args.dyna_dim,
        dyna_ffn_dim=args.dyna_ffn_dim,
        dyna_num_blocks=args.dyna_num_blocks,
        dyna_num_heads=args.dyna_num_heads,
        dropout=args.dropout,
        mask_limit=args.mask_limit,
[content] jasmine/train_dynamics.py (python)

    _, index_counts_tokenizer = jnp.unique_counts(
        jnp.ravel(outputs["video_tokens"]),
        size=num_patch_latents,
        fill_value=0,
    )
    codebook_usage_tokenizer = (index_counts_tokenizer != 0).mean()