p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission_case_study_vanilla.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=01:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --job-name=coinrun_sample_maskgit_mila_submission

# Activate virtual environment
source .venv/bin/activate

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_500m_seed_w_increment/val"
CHECKPOINT_PATH="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/dynamics/dynamics_coinrun_mila_submission/29973"


echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --checkpoint $CHECKPOINT_PATH \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=4 \
    --patch_size=16 \
    --start_frame=4 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit
[content] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission_case_study_action_prepend.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=01:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --job-name=coinrun_sample_maskgit_mila_submission

# Activate virtual environment
source .venv/bin/activate

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_500m_seed_w_increment/val"
CHECKPOINT_PATH="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/dynamics/dynamics_coinrun_mila_submission/29973"


echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --checkpoint $CHECKPOINT_PATH \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=4 \
    --patch_size=16 \
    --start_frame=4 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit
[content] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission_case_study_action_prepend.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=01:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --job-name=coinrun_sample_maskgit_mila_submission_case_study_vanilla

# Activate virtual environment
source .venv/bin/activate

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10M_npy_arr_rec/array_record/test"
CHECKPOINT_PATH="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/dynamics/dynamics_case_study_dataset_10M_30031"


echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --checkpoint $CHECKPOINT_PATH \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=4 \
    --patch_size=16 \
    --start_frame=4 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit
[tab] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission_case_study_action_prepend.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --job-name=coinrun_sample_maskgit_mila_submission_case_study_action_prepend

# Activate virtual environment
source .venv/bin/activate

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10M_npy_arr_rec/array_record/test"
CHECKPOINT_PATH="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/dynamics/dynamics_case_study_dataset_10M_action_prepend_30032"


echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --tokenizer_ffn_dim=512 \
    --tokenizer_num_blocks=8 \
    --dyna_ffn_dim=512 \
    --dyna_num_blocks=12 \
    --output_dir=gifs/dynamics_case_study_dataset_10M_action_prepend \
    --checkpoint $CHECKPOINT_PATH \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=4 \
    --patch_size=4 \
    --start_frame=4 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit
[content] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission_case_study_action_prepend.sh (shellscript)
current_branch=$(git rev-parse --abbrev-ref HEAD)
if [ "$current_branch" != "prepend-action-maskgit" ]; then
    echo "This script must be run from the prepend-action-maskgit branch. Current branch is $current_branch. Exiting."
    exit 1
fi
[tab] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission_case_study_vanilla.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --job-name=coinrun_sample_maskgit_mila_submission_case_study_vanilla

# Activate virtual environment
source .venv/bin/activate

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10M_npy_arr_rec/array_record/test"
CHECKPOINT_PATH="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/dynamics/dynamics_case_study_dataset_10M_30031"


echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --tokenizer_ffn_dim=512 \
    --tokenizer_num_blocks=8 \
    --dyna_ffn_dim=512 \
    --dyna_num_blocks=12 \
    --output_dir=gifs/dynamics_case_study_dataset_10M_vanilla \
    --checkpoint $CHECKPOINT_PATH \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=4 \
    --patch_size=4 \
    --start_frame=4 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit
[content] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission_case_study_vanilla.sh (shellscript)
current_branch=$(git rev-parse --abbrev-ref HEAD)
if [ "$current_branch" != "prepend-action-maskgit" ]; then
    echo "This script must be run from the prepend-action-maskgit branch. Current branch is $current_branch. Exiting."
    exit 1
fi
[content] jasmine/sample.py (python)
    for t, img in enumerate(imgs[1:]):
        d = ImageDraw.Draw(img)
        for row in range(B):
            if args.print_action_indices:
                action = action_batch_BSm11[row, t, 0]
                y_offset = row * batch["videos"].shape[2] + 2
                d.text((2, y_offset), f"{action}", fill=255)
[content] jasmine/sample.py (python)
    video_comparison = np.zeros((2, *recon_video_BSHWC.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")
[content] jasmine/sample.py (python)
    if args.print_action_indices:
        draw = ImageDraw.Draw(canvas)
        # actions are defined for transitions between frames: indices [0..T-2] correspond to frames [1..T-1]
        for t in range(1, T):
            action = int(action_batch_BSm11[seq_idx, t - 1, 0])
            draw.text((t * W + 2, 2), f"{action}", fill=255)

    # Save PNG
    os.makedirs(args.output_dir, exist_ok=True)
    canvas.save(os.path.join(args.output_dir, f"sequence_{time.time()}.png"))
[content] jasmine/sample.py (python)
    # Move from JAX device arrays to NumPy for PIL compatibility
    pred_videos = jax.device_get(pred_videos)

    # Select sequence index in the batch (user-specified)
    seq_idx = 31
    pred_seq_Thwc = pred_videos[seq_idx]

    T, H, W, C = pred_seq_Thwc.shape

    # Create canvas with one row (Reconstruction only) and T frames concatenated horizontally
    canvas = Image.new("RGB", (T * W, H))
    for t in range(T):
        pred_img = Image.fromarray(pred_seq_Thwc[t])
        canvas.paste(pred_img, (t * W, 0))
[content] jasmine/sample.py (python)
        gt_img = Image.fromarray(gt_seq_Thwc[t])
        canvas.paste(gt_img, (t * W, 0))

    # Do not draw action indices on the ground-truth-only PNG

    # Save PNG
    os.makedirs(args.output_dir, exist_ok=True)
    canvas.save(os.path.join(args.output_dir, f"sequence_{time.time()}.png"))