p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission copy.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

# 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
[tab] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission_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

# 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
[tab] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission_prepend copy.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

# 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
[tab] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission_no_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

# 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
[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"

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

echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --seed=1 \
    --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=32 \
    --patch_size=4 \
    --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_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"

current_branch=$(git rev-parse --abbrev-ref HEAD)
if [ "$current_branch" != "main" ]; then
    echo "This script must be run from the main branch. Current branch is $current_branch. Exiting."
    exit 1
fi

echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --seed=1 \
    --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=32 \
    --patch_size=4 \
    --start_frame=4 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit
[content] jasmine/sample.py (python)
    per_frame_ssim = ssim.mean(0)
    per_frame_psnr = psnr.mean(0)
    avg_ssim = ssim.mean()
    avg_psnr = psnr.mean()

    print("Per-frame SSIM:
", per_frame_ssim)
    print("Per-frame PSNR:
", per_frame_psnr)

    print(f"SSIM: {avg_ssim}")
    print(f"PSNR: {avg_psnr}")

    # --- 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)
    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)
    B = batch["videos"].shape[0]
    if action_batch_E is not None:
        action_batch_BSm11 = jnp.reshape(action_batch_E, (B, args.seq_len - 1, 1))
    else:
        action_batch_BSm11 = jnp.reshape(
            batch["actions"][:, :-1], (B, args.seq_len - 1, 1)
        )
    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)

    os.makedirs(args.output_dir, exist_ok=True)
    imgs[0].save(
        os.path.join(args.output_dir, f"generation_{time.time()}.gif"),
        save_all=True,
        append_images=imgs[1:],
        duration=250,
        loop=0,
    )