Add files using upload-large-folder tool

README.md

@@ -0,0 +1,109 @@
+---
+language:
+- en
+tags:
+- compact
+- image-tokenization
+- world-model
+- robotics
+- pytorch
+- image-to-image
+---
+
+# CompACT — 16-token checkpoints
+
+Final 16-token checkpoints for **Planning in 8 Tokens: A Compact Discrete Tokenizer for Latent World Model** (CVPR 2026).
+
+[Code and setup instructions](https://github.com/kdwonn/CompACT) · [Paper](https://arxiv.org/abs/2603.05438) · [Project](https://kdwonn.github.io/CompACT)
+
+## Included models
+
+| Directory | Model | Resolution | Checkpoint |
+|---|---|---|---|
+| `tokenizer-16-224` | CompACT, 16 tokens | 224 × 224 | `checkpoints/epoch=24-step=500000.ckpt` |
+| `tokenizer-16-256` | CompACT, 16 tokens | 256 × 256 | `checkpoints/epoch=24-step=500000.ckpt` |
+| `cdit-b-16` | CDiT-B world model | 224 × 224 | `checkpoints/latest.pth.tar` |
+| `cdit-l-16` | CDiT-L world model | 224 × 224 | `checkpoints/latest.pth.tar` |
+
+Both world models use **`tokenizer-16-224`**. The 256-resolution tokenizer is provided separately. These are the original full training checkpoint files, including training state; world-model inference uses the `ema` weights. Only the final checkpoint for each variant is included. Exact training steps, original experiment names, file sizes, and SHA-256 checksums are in [manifest.json](manifest.json).
+
+## Download
+
+Install the environment following the [code repository](https://github.com/kdwonn/CompACT). Run from its root:
+
+```python
+from huggingface_hub import snapshot_download
+
+snapshot_download(repo_id="kdwon/CompACT", local_dir="checkpoints/CompACT")
+```
+
+To download just one variant, include its config:
+
+```python
+snapshot_download(
+    repo_id="kdwon/CompACT",
+    local_dir="checkpoints/CompACT",
+    allow_patterns=["tokenizer-16-224/**", "manifest.json", "SHA256SUMS"],
+)
+```
+
+## Configuration
+
+Each variant includes `.hydra/config.yaml`, in the format expected by the code repository. Set these environment variables to your local paths:
+
+```bash
+export COMPACT_CKPT_ROOT="$(pwd)/checkpoints/CompACT"
+export BASE_TOKENIZER_CKPT=/absolute/path/to/base-checkpoints
+export DATASET_PREFIX=/absolute/path/to/datasets
+```
+
+The current constructors also require these initialization files:
+
+- `$BASE_TOKENIZER_CKPT/mage_vqgan/vqgan_jax_strongaug.ckpt`
+- `$BASE_TOKENIZER_CKPT/dinov3/dinov3_vitb16_pretrain_lvd1689m-73cec8be.pth`
+
+See the code repository for obtaining the MAGE and DINOv3 files. `DATASET_PREFIX` must be defined even for standalone tokenizer loading because the world-model loader resolves the complete saved tokenizer configuration.
+
+World-model configs use `${oc.env:COMPACT_CKPT_ROOT}/tokenizer-16-224` and `${oc.env:DATASET_PREFIX}/nwm/{recon,sacson,scand}` instead of the original machine's paths. Adjust dataset overrides to match your local layout.
+
+## Load a tokenizer
+
+```bash
+uv run load_tokenizer_checkpoint.py "$COMPACT_CKPT_ROOT/tokenizer-16-224" --no-test
+```
+
+Use `tokenizer-16-256` for the 256-resolution variant. For image encoding, use the saved DINO normalization (`dinov2_mean` and `dinov2_std`); the tokenizer's output normalization buffers describe decoded images.
+
+## Planning evaluation
+
+With the required navigation datasets installed:
+
+```bash
+uv run bash scripts/plan.sh --nproc=1 -- \
+  ++exp_dir="$COMPACT_CKPT_ROOT/cdit-b-16" \
+  ++tokenizer_path="$COMPACT_CKPT_ROOT/tokenizer-16-224" \
+  ++ckp=latest
+```
+
+Replace `cdit-b-16` with `cdit-l-16` to use the larger world model. Dataset setup and evaluation options are documented in the code repository.
+
+## Validation
+
+These files passed strict weight-loading checks in the public CompACT codebase. Both tokenizers encoded and reconstructed a real image. Both the regular and EMA weights of the world models passed forward checks, and EMA models completed image-to-predicted-image inference. Validation used PyTorch 2.6.0+cu124 and bfloat16 inference on an RTX 6000 Ada GPU. These smoke checks do not constitute a rerun of the paper's evaluation benchmarks.
+
+Verify downloaded checkpoint bytes from the download directory:
+
+```bash
+sha256sum -c SHA256SUMS
+```
+
+## Citation
+
+```bibtex
+@inproceedings{kim2026planning,
+  title={Planning in 8 Tokens: A Compact Discrete Tokenizer for Latent World Model},
+  author={Kim, Dongwon and Seo, Gawon and Lee, Jinsung and Cho, Minsu and Kwak, Suha},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+  year={2026}
+}
+```

SHA256SUMS

@@ -0,0 +1,4 @@
+946156a657430b3aedfa005e2a83729642c679a48f739667970e9dbedba09354  tokenizer-16-224/checkpoints/epoch=24-step=500000.ckpt
+f73ffacaceed9d8389410d398fe9e458cbb9b55aa744736fc5a468821bb3ddf0  tokenizer-16-256/checkpoints/epoch=24-step=500000.ckpt
+d52268676201c96e6d1586d6a49ad60b5798ea59997301cb2f27afde2aa2212f  cdit-b-16/checkpoints/latest.pth.tar
+168f3d0cd5784464f056288c9757fc12391c98804a34d4956d36b52ce287f501  cdit-l-16/checkpoints/latest.pth.tar

cdit-b-16/.hydra/config.yaml

@@ -0,0 +1,82 @@
+seed: 0
+epochs: 100
+log_every: 100
+ckpt_every: 10000
+eval_every: 5000
+bfloat16: 1
+torch_compile: 0
+resume_from_dir: null
+dataset:
+  image_size: 224
+  context_size: 4
+  normalize: true
+  distance:
+    min_dist_cat: -64
+    max_dist_cat: 64
+  len_traj_pred: 64
+  action_stats:
+    min:
+    - -2.5
+    - -4
+    max:
+    - 5
+    - 4
+  datasets:
+    recon:
+      data_folder: ${oc.env:DATASET_PREFIX}/nwm/recon
+      train: data_splits/recon/train/
+      test: data_splits/recon/test/
+      goals_per_obs: 4
+      metric_waypoint_spacing: 0.25
+    sacson:
+      data_folder: ${oc.env:DATASET_PREFIX}/nwm/sacson
+      train: data_splits/sacson/train
+      test: data_splits/sacson/test
+      goals_per_obs: 4
+      metric_waypoint_spacing: 0.255
+    scand:
+      data_folder: ${oc.env:DATASET_PREFIX}/nwm/scand
+      train: data_splits/scand/train
+      test: data_splits/scand/test
+      goals_per_obs: 4
+      metric_waypoint_spacing: 0.38
+model:
+  generator:
+    _target_: disc_models.DiscreteCDiT
+    input_size: ${model.tokenizer.num_tokens}
+    codebook_size: ${model.tokenizer.codebook_size}
+    context_size: 2
+    hidden_size: 768
+    depth: 12
+    num_heads: 12
+    use_action_cfg: ${model.diffusion.cfg_prob}
+  diffusion:
+    _target_: diffusion.disc_diffusion.DiscreteDiffusionWithHistoryMasking
+    mask_token_id: ${model.tokenizer.codebook_size}
+    num_timesteps: 8
+    schedule_type: linear
+    temperature: 1.0
+    history_key: x_cond
+    history_sample_t: 0.8
+    cfg_prob: 0
+    cfg_scale: 2.0
+tokenizer_path: ${oc.env:COMPACT_CKPT_ROOT}/tokenizer-16-224
+training:
+  notes: newcfg_acfcfg-0_historymasking-linear_qformer-freeze
+  run_name: ${now:%Y%m%d_%H%M%S}_${training.notes}
+  results_dir: logs/nwm_training
+  batch_size: 32
+  num_workers: 8
+  grad_clip_val: 2.0
+  from_checkpoint: null
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 0.0001
+    weight_decay: 0.01
+  wandb_enabled: true
+  wandb_project: nwm-training
+  wandb_notes: ${training.run_name}
+  wandb_tags: []
+scheduler:
+  scheduler:
+    _target_: null

cdit-b-16/checkpoints/latest.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d52268676201c96e6d1586d6a49ad60b5798ea59997301cb2f27afde2aa2212f
+size 3887081502

cdit-l-16/.hydra/config.yaml

@@ -0,0 +1,83 @@
+seed: 0
+epochs: 100
+log_every: 100
+ckpt_every: 10000
+eval_every: 5000
+bfloat16: 1
+torch_compile: 0
+resume_from_dir: null
+dataset:
+  image_size: 224
+  context_size: 4
+  normalize: true
+  distance:
+    min_dist_cat: -64
+    max_dist_cat: 64
+  len_traj_pred: 64
+  action_stats:
+    min:
+    - -2.5
+    - -4
+    max:
+    - 5
+    - 4
+  datasets:
+    recon:
+      data_folder: ${oc.env:DATASET_PREFIX}/nwm/recon
+      train: data_splits/recon/train/
+      test: data_splits/recon/test/
+      goals_per_obs: 4
+      metric_waypoint_spacing: 0.25
+    sacson:
+      data_folder: ${oc.env:DATASET_PREFIX}/nwm/sacson
+      train: data_splits/sacson/train
+      test: data_splits/sacson/test
+      goals_per_obs: 4
+      metric_waypoint_spacing: 0.255
+    scand:
+      data_folder: ${oc.env:DATASET_PREFIX}/nwm/scand
+      train: data_splits/scand/train
+      test: data_splits/scand/test
+      goals_per_obs: 4
+      metric_waypoint_spacing: 0.38
+model:
+  generator:
+    _target_: disc_models.DiscreteCDiT
+    input_size: ${model.tokenizer.num_tokens}
+    codebook_size: ${model.tokenizer.codebook_size}
+    context_size: 2
+    hidden_size: 1024
+    depth: 24
+    num_heads: 16
+    use_action_cfg: ${model.diffusion.cfg_prob}
+  diffusion:
+    _target_: diffusion.disc_diffusion.DiscreteDiffusionWithHistoryMasking
+    mask_token_id: ${model.tokenizer.codebook_size}
+    num_timesteps: 8
+    schedule_type: linear
+    temperature: 1.0
+    history_key: x_cond
+    history_sample_t: 0.8
+    cfg_prob: 0.1
+    cfg_scale: 2.0
+    without_history_masking: false
+tokenizer_path: ${oc.env:COMPACT_CKPT_ROOT}/tokenizer-16-224
+training:
+  notes: cdit-l_newcfg_acfcfg-0_historymasking-linear_qformer-freeze
+  run_name: ${now:%Y%m%d_%H%M%S}_${training.notes}
+  results_dir: logs/nwm_training
+  batch_size: 32
+  num_workers: 16
+  grad_clip_val: 2.0
+  from_checkpoint: null
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 0.0001
+    weight_decay: 0.01
+  wandb_enabled: true
+  wandb_project: nwm-training
+  wandb_notes: ${training.run_name}
+  wandb_tags: []
+scheduler:
+  scheduler:
+    _target_: null

cdit-l-16/checkpoints/latest.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:168f3d0cd5784464f056288c9757fc12391c98804a34d4956d36b52ce287f501
+size 11996868852

manifest.json

@@ -0,0 +1,34 @@
+[
+  {
+    "variant": "tokenizer-16-224",
+    "checkpoint": "tokenizer-16-224/checkpoints/epoch=24-step=500000.ckpt",
+    "bytes": 11572743012,
+    "sha256": "946156a657430b3aedfa005e2a83729642c679a48f739667970e9dbedba09354",
+    "step": 500000,
+    "original_experiment": "20250930_072404_dinov3-b-qformer-freeze_mmdit-l_code2exp16_lr1e-4_repa-dec-8_nomasking"
+  },
+  {
+    "variant": "tokenizer-16-256",
+    "checkpoint": "tokenizer-16-256/checkpoints/epoch=24-step=500000.ckpt",
+    "bytes": 11573725860,
+    "sha256": "f73ffacaceed9d8389410d398fe9e458cbb9b55aa744736fc5a468821bb3ddf0",
+    "step": 500000,
+    "original_experiment": "20251003_100611_dinov3-b-qformer-freeze_256"
+  },
+  {
+    "variant": "cdit-b-16",
+    "checkpoint": "cdit-b-16/checkpoints/latest.pth.tar",
+    "bytes": 3887081502,
+    "sha256": "d52268676201c96e6d1586d6a49ad60b5798ea59997301cb2f27afde2aa2212f",
+    "step": 230000,
+    "original_experiment": "20251005_020501_newcfg_acfcfg-0_historymasking-linear_qformer-freeze"
+  },
+  {
+    "variant": "cdit-l-16",
+    "checkpoint": "cdit-l-16/checkpoints/latest.pth.tar",
+    "bytes": 11996868852,
+    "sha256": "168f3d0cd5784464f056288c9757fc12391c98804a34d4956d36b52ce287f501",
+    "step": 200000,
+    "original_experiment": "20251102_074208_cdit-l_newcfg_acfcfg-0_historymasking-linear_qformer-freeze"
+  }
+]

tokenizer-16-224/.hydra/config.yaml

@@ -0,0 +1,207 @@
+seed: 42
+max_steps: 500000
+log_every: 500
+eval_every: 100000
+bfloat16: true
+devices: auto
+experiment:
+  n_compact_tokens: 16
+  n_2d_tokens: 196
+  compact_tokenizer_dim: 512
+  compact_tokenizer_code_dim: 6
+  compact_tokenizer_num_blocks: 8
+  process_attn_depth: 3
+  n_registers: 0
+  enable_encoder_masking: false
+  cutoff_enable: false
+  cutoff_prob: 0.5
+  cutoff_downsample_factor: 2
+model:
+  tokenizer:
+    _target_: token_distill.token_distill_model.TokenDistillModel
+    codebook_size: ${model.tokenizer.compact_tokenizer.quantizer.codebook_size}
+    num_tokens: ${experiment.n_compact_tokens}
+    normalize_mean: ${dataset.mean}
+    normalize_std: ${dataset.std}
+    base_tokenizer:
+      _target_: token_distill.mage_vqgan.vqgan.VQModel
+      encoder:
+        _target_: token_distill.mage_vqgan.model.Encoder
+        ch: 128
+        out_ch: 3
+        ch_mult:
+        - 1
+        - 1
+        - 2
+        - 2
+        - 4
+        num_res_blocks: 2
+        attn_resolutions:
+        - 16
+        dropout: 0.0
+        in_channels: 3
+        resolution: 256
+        z_channels: 256
+        double_z: false
+      decoder:
+        _target_: token_distill.mage_vqgan.model.Decoder
+        ch: 128
+        out_ch: 3
+        ch_mult:
+        - 1
+        - 1
+        - 2
+        - 2
+        - 4
+        num_res_blocks: 2
+        attn_resolutions:
+        - 16
+        dropout: 0.0
+        in_channels: 3
+        resolution: 256
+        z_channels: 256
+        double_z: false
+      n_embed: 1024
+      embed_dim: 256
+      ckpt_path: ${oc.env:BASE_TOKENIZER_CKPT}/mage_vqgan/vqgan_jax_strongaug.ckpt
+    compact_tokenizer:
+      _target_: token_distill.compact_tokenizer.CompactTokenizer
+      encoder:
+        _target_: token_distill.vit_components.DINOv3EncoderQFormer
+        in_dim: 768
+        latent_dim: 768
+        num_2d_tokens: ${experiment.n_2d_tokens}
+        dinov3_model_name: dinov3_vitb16
+        dinov3_weights_path: ${oc.env:BASE_TOKENIZER_CKPT}/dinov3/dinov3_vitb16_pretrain_lvd1689m-73cec8be.pth
+        num_1d_tokens: ${experiment.n_compact_tokens}
+        num_registers: ${experiment.n_registers}
+        final_norm: true
+        out_proj_dim: ${experiment.compact_tokenizer_code_dim}
+        final_norm_context: ${model.tokenizer.compact_tokenizer.repa_on_enc}
+        enable_masking: ${experiment.enable_encoder_masking}
+        masking_prob: 0.9
+        min_masking_ratio: 0.0
+        max_masking_ratio: 0.6
+        freeze_dinov3: true
+        num_qformer_layers: 5
+        num_attention_heads: 8
+      decoder:
+        _target_: token_distill.vit_components.DiTDecoder
+        in_dim: 1024
+        latent_dim: 1024
+        target_codebook_size: ${model.tokenizer.base_tokenizer.n_embed}
+        process_attn_depth: ${experiment.process_attn_depth}
+        num_blocks: 16
+        num_2d_tokens: ${experiment.n_2d_tokens}
+        final_norm: false
+        in_proj_dim: ${experiment.compact_tokenizer_code_dim}
+        final_norm_context: true
+        repa_target_block_num: 8
+        num_1d_tokens: ${experiment.n_compact_tokens}
+      quantizer:
+        _target_: token_distill.fsq.FSQWrapper
+        levels:
+        - 8
+        - 8
+        - 8
+        - 5
+        - 5
+        - 5
+        codebook_size: 64000
+      diffusion:
+        _target_: diffusion.disc_diffusion.DiscreteDiffusionWithCFG
+        mask_token_id: ${model.tokenizer.base_tokenizer.n_embed}
+        num_timesteps: 16
+        schedule_type: cosine
+        temperature: 1.0
+        output_key: decoder_2d_logits
+        cfg_prob: 0.1
+      num_1d_tokens: ${experiment.n_compact_tokens}
+      num_2d_tokens: ${experiment.n_2d_tokens}
+      masking_type: cosine
+      target_codebook_size: ${model.tokenizer.base_tokenizer.n_embed}
+      repa_on_enc: false
+  rep_guidance:
+    _target_: token_distill.rep_guidance.RepGuidanceLoss
+    in_dim: ${model.tokenizer.compact_tokenizer.decoder.latent_dim}
+    repa_target_enc: dinov3_vitb16
+    dinov3_weights_path: ${oc.env:BASE_TOKENIZER_CKPT}/dinov3/dinov3_vitb16_pretrain_lvd1689m-73cec8be.pth
+    use_mask: false
+dataset:
+  context_size: 1
+  batch_size: 64
+  image_size: 224
+  dinov2_image_size: 224
+  num_workers: 8
+  data_root: ${oc.env:DATASET_PREFIX}/imagenet_wds
+  train_shards: ${dataset.data_root}/train/imagenet-train-{000000..001281}.tar
+  val_shards: ${dataset.data_root}/val/imagenet-val-{000000..000049}.tar
+  test_shards: ${dataset.data_root}/val/imagenet-val-{000000..000049}.tar
+  shuffle_buffer_size: 5000
+  shard_shuffle_initial: 1000
+  estimated_train_size: 1281167
+  estimated_val_size: 50000
+  estimated_test_size: 50000
+  use_augmentation: true
+  mean:
+  - 0
+  - 0
+  - 0
+  std:
+  - 1
+  - 1
+  - 1
+  dinov2_mean:
+  - 0.485
+  - 0.456
+  - 0.406
+  dinov2_std:
+  - 0.229
+  - 0.224
+  - 0.225
+  use_dual_normalization: true
+  use_dino_image_for_compact: true
+  base_tokenizer_name: mage
+training:
+  notes: dinov3-b-qformer-freeze_mmdit-l_code2exp16_lr1e-4_repa-dec-8_nomasking
+  run_name: ${now:%Y%m%d_%H%M%S}_${training.notes}
+  results_dir: logs/tokenizer_training
+  grad_clip_val: 1.0
+  accumulate_grad_batches: 1
+  from_checkpoint: null
+  recon_loss_weight: 1.0
+  commit_loss_weight: 0.25
+  rep_guidance_loss_weight: 1.0
+  wandb_project: token-distill
+  log_images_locally: true
+  local_image_log_dir: null
+  model_summary_depth: 3
+  log_metrics_every_n_steps: 50
+  profiler:
+    enabled: false
+    type: pytorch
+    record_module_names: true
+    file_name: debug
+    export_to_chrome: true
+    sort_by_key: cpu_time
+    row_limit: 100
+    schedule:
+      wait: 10
+      warmup: 20
+      active: 10
+      repeat: 2
+    profile_memory: true
+    with_stack: true
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 0.0001
+    weight_decay: 0.001
+    betas:
+    - 0.9
+    - 0.999
+scheduler:
+  scheduler:
+    _target_: schedulers.CosineAnnealingWithPlateau
+    T_max: ${max_steps}
+    eta_min: 1.0e-06
+    warmup_steps: 10000

tokenizer-16-224/checkpoints/epoch=24-step=500000.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:946156a657430b3aedfa005e2a83729642c679a48f739667970e9dbedba09354
+size 11572743012

tokenizer-16-256/.hydra/config.yaml

@@ -0,0 +1,207 @@
+seed: 42
+max_steps: 500000
+log_every: 500
+eval_every: 100000
+bfloat16: true
+devices: auto
+experiment:
+  n_compact_tokens: 16
+  n_2d_tokens: 256
+  compact_tokenizer_dim: 512
+  compact_tokenizer_code_dim: 6
+  compact_tokenizer_num_blocks: 8
+  process_attn_depth: 3
+  n_registers: 0
+  enable_encoder_masking: false
+  cutoff_enable: false
+  cutoff_prob: 0.5
+  cutoff_downsample_factor: 2
+model:
+  tokenizer:
+    _target_: token_distill.token_distill_model.TokenDistillModel
+    codebook_size: ${model.tokenizer.compact_tokenizer.quantizer.codebook_size}
+    num_tokens: ${experiment.n_compact_tokens}
+    normalize_mean: ${dataset.mean}
+    normalize_std: ${dataset.std}
+    base_tokenizer:
+      _target_: token_distill.mage_vqgan.vqgan.VQModel
+      encoder:
+        _target_: token_distill.mage_vqgan.model.Encoder
+        ch: 128
+        out_ch: 3
+        ch_mult:
+        - 1
+        - 1
+        - 2
+        - 2
+        - 4
+        num_res_blocks: 2
+        attn_resolutions:
+        - 16
+        dropout: 0.0
+        in_channels: 3
+        resolution: 256
+        z_channels: 256
+        double_z: false
+      decoder:
+        _target_: token_distill.mage_vqgan.model.Decoder
+        ch: 128
+        out_ch: 3
+        ch_mult:
+        - 1
+        - 1
+        - 2
+        - 2
+        - 4
+        num_res_blocks: 2
+        attn_resolutions:
+        - 16
+        dropout: 0.0
+        in_channels: 3
+        resolution: 256
+        z_channels: 256
+        double_z: false
+      n_embed: 1024
+      embed_dim: 256
+      ckpt_path: ${oc.env:BASE_TOKENIZER_CKPT}/mage_vqgan/vqgan_jax_strongaug.ckpt
+    compact_tokenizer:
+      _target_: token_distill.compact_tokenizer.CompactTokenizer
+      encoder:
+        _target_: token_distill.vit_components.DINOv3EncoderQFormer
+        in_dim: 768
+        latent_dim: 768
+        num_2d_tokens: ${experiment.n_2d_tokens}
+        dinov3_model_name: dinov3_vitb16
+        dinov3_weights_path: ${oc.env:BASE_TOKENIZER_CKPT}/dinov3/dinov3_vitb16_pretrain_lvd1689m-73cec8be.pth
+        num_1d_tokens: ${experiment.n_compact_tokens}
+        num_registers: ${experiment.n_registers}
+        final_norm: true
+        out_proj_dim: ${experiment.compact_tokenizer_code_dim}
+        final_norm_context: ${model.tokenizer.compact_tokenizer.repa_on_enc}
+        enable_masking: ${experiment.enable_encoder_masking}
+        masking_prob: 0.9
+        min_masking_ratio: 0.0
+        max_masking_ratio: 0.6
+        freeze_dinov3: true
+        num_qformer_layers: 5
+        num_attention_heads: 8
+      decoder:
+        _target_: token_distill.vit_components.DiTDecoder
+        in_dim: 1024
+        latent_dim: 1024
+        target_codebook_size: ${model.tokenizer.base_tokenizer.n_embed}
+        process_attn_depth: ${experiment.process_attn_depth}
+        num_blocks: 16
+        num_2d_tokens: ${experiment.n_2d_tokens}
+        final_norm: false
+        in_proj_dim: ${experiment.compact_tokenizer_code_dim}
+        final_norm_context: true
+        repa_target_block_num: 8
+        num_1d_tokens: ${experiment.n_compact_tokens}
+      quantizer:
+        _target_: token_distill.fsq.FSQWrapper
+        levels:
+        - 8
+        - 8
+        - 8
+        - 5
+        - 5
+        - 5
+        codebook_size: 64000
+      diffusion:
+        _target_: diffusion.disc_diffusion.DiscreteDiffusionWithCFG
+        mask_token_id: ${model.tokenizer.base_tokenizer.n_embed}
+        num_timesteps: 16
+        schedule_type: cosine
+        temperature: 1.0
+        output_key: decoder_2d_logits
+        cfg_prob: 0.1
+      num_1d_tokens: ${experiment.n_compact_tokens}
+      num_2d_tokens: ${experiment.n_2d_tokens}
+      masking_type: cosine
+      target_codebook_size: ${model.tokenizer.base_tokenizer.n_embed}
+      repa_on_enc: false
+  rep_guidance:
+    _target_: token_distill.rep_guidance.RepGuidanceLoss
+    in_dim: ${model.tokenizer.compact_tokenizer.decoder.latent_dim}
+    repa_target_enc: dinov3_vitb16
+    dinov3_weights_path: ${oc.env:BASE_TOKENIZER_CKPT}/dinov3/dinov3_vitb16_pretrain_lvd1689m-73cec8be.pth
+    use_mask: false
+dataset:
+  context_size: 1
+  batch_size: 64
+  image_size: 256
+  dinov2_image_size: 256
+  num_workers: 8
+  data_root: ${oc.env:DATASET_PREFIX}/imagenet_wds
+  train_shards: ${dataset.data_root}/train/imagenet-train-{000000..001281}.tar
+  val_shards: ${dataset.data_root}/val/imagenet-val-{000000..000049}.tar
+  test_shards: ${dataset.data_root}/val/imagenet-val-{000000..000049}.tar
+  shuffle_buffer_size: 5000
+  shard_shuffle_initial: 1000
+  estimated_train_size: 1281167
+  estimated_val_size: 50000
+  estimated_test_size: 50000
+  use_augmentation: true
+  mean:
+  - 0
+  - 0
+  - 0
+  std:
+  - 1
+  - 1
+  - 1
+  dinov2_mean:
+  - 0.485
+  - 0.456
+  - 0.406
+  dinov2_std:
+  - 0.229
+  - 0.224
+  - 0.225
+  use_dual_normalization: true
+  use_dino_image_for_compact: true
+  base_tokenizer_name: mage
+training:
+  notes: dinov3-b-qformer-freeze_256
+  run_name: ${now:%Y%m%d_%H%M%S}_${training.notes}
+  results_dir: logs/tokenizer_training
+  grad_clip_val: 1.0
+  accumulate_grad_batches: 1
+  from_checkpoint: null
+  recon_loss_weight: 1.0
+  commit_loss_weight: 0.25
+  rep_guidance_loss_weight: 1.0
+  wandb_project: token-distill
+  log_images_locally: true
+  local_image_log_dir: null
+  model_summary_depth: 3
+  log_metrics_every_n_steps: 50
+  profiler:
+    enabled: false
+    type: pytorch
+    record_module_names: true
+    file_name: debug
+    export_to_chrome: true
+    sort_by_key: cpu_time
+    row_limit: 100
+    schedule:
+      wait: 10
+      warmup: 20
+      active: 10
+      repeat: 2
+    profile_memory: true
+    with_stack: true
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 0.0001
+    weight_decay: 0.001
+    betas:
+    - 0.9
+    - 0.999
+scheduler:
+  scheduler:
+    _target_: schedulers.CosineAnnealingWithPlateau
+    T_max: ${max_steps}
+    eta_min: 1.0e-06
+    warmup_steps: 10000

tokenizer-16-256/checkpoints/epoch=24-step=500000.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f73ffacaceed9d8389410d398fe9e458cbb9b55aa744736fc5a468821bb3ddf0
+size 11573725860