p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/dev/mihir/horeka/yolo-runs/sampling_dev.sh (shellscript)

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
# source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=$ws_dir/checkpoints/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4
dynamics_ckpt_dir=$1
echo $dynamics_ckpt_dir

env | grep SLURM

srun python sample.py \
    --checkpoint $dynamics_ckpt_dir \
    --dyna_dim=1024 \
    --dyna_num_blocks=16 \
    --dyna_num_heads=16 \
    --dyna_ffn_dim=4096 \
    --seq_len=32 \
    --seed=69 \
    --batch_size=4 \
    --maskgit_steps=2 \
    --start_frame=16 \
    --data_dir $array_records_dir

# srun python sample.py \
    # --checkpoint $dynamics_ckpt_dir \
    # --start_frame=0 \
    # --batch_size=12 \
    # --seq_len=2 \
    # --data_dir $array_records_dir
[content] sample.py (python)
rng, _rng = jax.random.split(rng)
image_shape = (args.image_height, args.image_width, 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)

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(
    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
[content] sample.py (python)
from dataclasses import dataclass
from typing import Optional
import time
import os

import dm_pix as pix
import einops
import jax
import jax.numpy as jnp
import flax.linen as nn
import numpy as np
from flax.training.train_state import TrainState
import grain
import orbax.checkpoint as ocp
import optax
from PIL import Image, ImageDraw
import tyro

from genie import Genie
from utils.dataloader import get_dataloader
[content] sample.py (python)
video_batch = jnp.array(video_batch)
print(video_batch.dtype)
video_batch = video_batch.astype(args.dtype) # / 255.0
print(video_batch.dtype)
video_batch = video_batch / 255.0
print(video_batch.dtype)
[tab] slurm/dev/mihir/horeka/yolo-runs/sampling_dev.sh (shellscript)

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
# source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=$ws_dir/checkpoints/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4
dynamics_ckpt_dir=$1
echo $dynamics_ckpt_dir

env | grep SLURM

srun python sample.py \
    --checkpoint $dynamics_ckpt_dir \
    --seq_len=16 \
    --seed=42 \
    --batch_size=4 \
    --maskgit_steps=1 \
    --start_frame=0 \
    --data_dir $array_records_dir \

    # --dyna_dim=1024 \
    # --dyna_num_blocks=16 \
    # --dyna_num_heads=16 \
    # --dyna_ffn_dim=4096 \

# srun python sample.py \
    # --checkpoint $dynamics_ckpt_dir \
    # --start_frame=0 \
    # --batch_size=12 \
    # --seq_len=2 \
    # --data_dir $array_records_dir