p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] experiments/dynamics_grain_tok_restore.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/maskgit_dynamics_openai_grain_tok_restore"
tokenizer_ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing"

export XLA_FLAGS=--xla_gpu_autotune_level=0
export PYTHONUNBUFFERED=1
srun python train_dynamics.py \
    --dyna_type 'causal' \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --batch_size 12 \
    --tokenizer_checkpoint $tokenizer_ckpt_dir \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] experiments/sample.sh (shellscript)
source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/maskgit_dynamics_openai_grain_tok_restore"

export XLA_FLAGS='--xla_gpu_autotune_level=0'
    #--dynamics_type "causal" \
srun python sample.py \
    --batch_size 1 \
    --seq_len 2 \
    --start_frame 1 \
    --checkpoint $ckpt_dir \
    --data_dir $data_dir
[content] utils/nn.py (python)
        z_BTNF = self.ffn_dense1(z_BTNM)
        z_BTNF = jax.nn.gelu(z_BTNF)
        z_BTNM = self.ffn_dense2(z_BTNF)
        x_BTNM = x_BTNM + z_BTNM

        return x_BTNM


class STTransformer(nnx.Module):
    """
    Dimension keys:
        B: batch size
        T: number of frames
        N: number of patches per frame
        I: number of input features
        M: model dimension
        F: FFN dimension
        O: number of output features
    """
    def __init__(
        self,
        input_dim: int,
        model_dim: int,
        ffn_dim: int,
        out_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.input_dim = input_dim
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.out_dim = out_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.input_norm1 = nnx.LayerNorm(
            num_features=self.input_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )
        self.input_dense = nnx.Linear(
            in_features=self.input_dim,
            out_features=self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )
        self.input_norm2 = nnx.LayerNorm(
            num_features=self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

        self.blocks = []
        for _ in range(self.num_blocks):
            self.blocks.append(
                STBlock(
                    dim=self.model_dim,
                    ffn_dim=self.ffn_dim,
                    num_heads=self.num_heads,
                    dropout=self.dropout,
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                    use_flash_attention=self.use_flash_attention,
                    rngs=rngs,
                )
            )

        self.output_dense = nnx.Linear(
            in_features=self.model_dim,
            out_features=self.out_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    def __call__(self, x_BTNI: jax.Array) -> jax.Array:
        x_BTNI = self.input_norm1(x_BTNI)
        x_BTNM = self.input_dense(x_BTNI)
        x_BTNM = self.input_norm2(x_BTNM)

        for block in self.blocks:
            x_BTNM = block(x_BTNM)

        x_BTNO = self.output_dense(x_BTNM)
        return x_BTNO
[content] /fast/home/franz.srambical/jafar/sample.py (python)
        generated_vid_BSHWC = _sampling_fn(genie, batch)
        return generated_vid_BSHWC

    # --- 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,
        # We don't use workers in order to avoid grain shutdown issues (https://github.com/google/grain/issues/398)
        num_workers=0,
        prefetch_buffer_size=1,
        seed=args.seed,
    )
    dataloader = iter(dataloader)
    video_batch_BSHWC = next(dataloader)
    gt_video = jnp.asarray(video_batch_BSHWC, dtype=jnp.float32) / 255.0
    video_batch_BSHWC = gt_video.astype(args.dtype)
    # Get latent actions for all videos in the batch
    batch = dict(videos=video_batch_BSHWC)
    action_batch_E = genie.vq_encode(batch, training=False)

    # --- Sample + evaluate video ---
    recon_video_BSHWC = _autoreg_sample(rng, video_batch_BSHWC, action_batch_E)
    recon_video_BSHWC = recon_video_BSHWC.astype(jnp.float32)
    gt = gt_video[:, : recon_video_BSHWC.shape[1]].clip(0, 1).reshape(-1, *gt_video.shape[2:])
    recon = recon_video_BSHWC.clip(0, 1).reshape(-1, *recon_video_BSHWC.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 = (gt_video * 255).astype(np.uint8)
    pred_videos = (recon_video_BSHWC * 255).astype(np.uint8)
    video_comparison = np.zeros((2, *recon_video_BSHWC.shape), dtype=np.uint8)