dRAE

KO | EN

gitlite — search

#react #security #python #react-native #generative-ai #go #css #chatbot #csharp #ai

(★ 9)

No description available.

File Explorer

ZIP Download

- settings.json
- pixabay_cat.png
- train_drae_384.sh
- train_drae_512.sh
- - __init__.py
  - convnext_perceptual.py
  - diffaug.py
  - dinodisc.py
  - discriminator.py
  - gan_loss.py
  - lpips.py
  - lpips_utils.py
  - utils.py
- - __init__.py
  - fid.py
  - ref_iqa.py
  - utils.py
- - __init__.py
  - drae.py
  - hsq.py
  - rae.py
- - dist_utils.py
  - ds_imagefolder.py
  - ds_parquet.py
  - hf_dataset.py
  - model_utils.py
  - optim_utils.py
  - resume_utils.py
  - save_val.py
  - train_utils.py
  - wandb_utils.py
- calculate_stat.py
- train.py
- train_drae_e2e.py
.gitignore
init.sh
LICENSE
pack_images.py
README.md
requirements.txt
tokenize_demo.py

# Use via CDN

jsDelivr

jsDelivr serves any public GitHub repository as a CDN with zero setup. Pick a version and a file to get a ready-to-paste link and snippet.

Version

File

Link

Example

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

Training code for **dRAE**, a visual tokenizer that applies Hyper-Spherical Quantization on top of pre-trained vision encoders (SigLIP2, DINOv2). Built upon the [RAE](https://github.com/bytetriper/RAE) codebase.

> [**dRAE: Representation Autoencoder with Hyper-Spherical Codes**](https://arxiv.org/abs/2607.22148)<br>
> Tianren Ma, Lin Long, Chuyan Chen, Mu Zhang, Junbo Zhao, Tong Zhang and Qixiang Ye.

## Environment

```bash
conda create -n drae python=3.10 -y
conda activate drae

# PyTorch 2.6+ with CUDA 12.1 (adjust for your CUDA version)
conda install pytorch==2.6.0 torchvision==0.21.0 pytorch-cuda=12.1 -c pytorch -c nvidia

pip install -r requirements.txt
```

## Preparation

Before launching training, three things need to be ready:

**1. Encoder weights** — downloaded automatically from HuggingFace on first run, or pre-download to avoid DDP race conditions:

```bash
python src/stage1/encoders/download_encoders.py
```

**2. Discriminator checkpoint** — the GAN discriminator uses a frozen DINO ViT-S backbone. Download from the [RAE repo](https://huggingface.co/nyu-visionx/RAE-collections/resolve/main/discs/dino_vit_small_patch8_224.pth) and place under `pretrained/`:

**3. Training data** — we train on the BLIP3o image-caption dataset. Download the following from HuggingFace:

```bash
pip install huggingface_hub
huggingface-cli download BLIP3o/BLIP3o-Pretrain-Long-Caption --repo-type dataset --local-dir data/Blip3o-LongCap
huggingface-cli download BLIP3o/BLIP3o-Pretrain-JourneyDB --repo-type dataset --local-dir data/Blip3o-DB
huggingface-cli download BLIP3o/BLIP3o-Pretrain-Short-Caption --repo-type dataset --local-dir data/Blip3o-ShortCap
```

These datasets are stored as `.tar` shards (WebDataset format). Point `--data-path` to the downloaded directories:

```bash
--data-path data/Blip3o-LongCap data/Blip3o-DB data/Blip3o-ShortCap
```

The dataloader (`src/utils/ds_parquet.py`) auto-detects the format and supports:

| Format | How to use |
|--------|-----------|
| **WebDataset** (default) | Directories of `.tar` shards containing image files (png/jpg/webp). Pass via `--data-path dir1 dir2 ...` |
| **Parquet** | Directories of `.parquet` files with an `image_bytes` column. Pass via `--data-path dir1 dir2 ...` |
| **HuggingFace** | Any HF image dataset (e.g. `ILSVRC/imagenet-1k`). Pass via `--hf-dataset <dataset_id>` |

## Training

Edit the `--data-path` in the launch script, then run:

```bash
# 384px, SigLIP2-Large encoder, ViT-B decoder
bash scripts/train_drae_384.sh

# 512px, SigLIP2-SO400M encoder, ViT-XL decoder
bash scripts/train_drae_512.sh
```

Key variables at the top of each script:

| Variable | Description |
|----------|-------------|
| `CODEBOOK_SIZE` | VQ codebook size (appears in experiment name) |
| `ACCUM_STEPS` | Gradient accumulation steps — increase if OOM |
| `DISTILL_WEIGHT` | Encoder distillation loss weight (e2e mode) |
| `ENCODER_LR_SCALE` | LR multiplier for the encoder (e.g. 0.1 = 10x smaller) |

Checkpoints and logs are saved to `--results-dir/$EXPERIMENT_NAME/`.

**Resuming** — if the checkpoint directory already exists, training resumes from the latest checkpoint automatically.

**W&B logging** — set `WANDB_KEY`, `ENTITY`, `PROJECT` as environment variables, and keep `--wandb` in the script. Set `WANDB_MODE=offline` to log locally only.

## Optional: Online Evaluation

To compute PSNR/SSIM/rFID during training, prepare ImageNet-1k validation data and reference statistics:

```bash
python src/utils/save_val.py --data-dir /path/to/imagenet_val_parquet --output-dir data/imagenet/ref
```

Then set `eval.data_path` and `eval.reference_npz_path` in the YAML config accordingly.

## Acknowledgement

- [RAE](https://github.com/bytetriper/RAE) — base codebase for representation autoencoders.
- [JiT](https://github.com/LTH14/JiT) — modern ViT decoder design.
- [TiTok](https://github.com/bytedance/1d-tokenizer) — perceptual loss design.
- Awesome Coding Agents

## BibTex
If you find this work helpful, please cite:
```
@article{ma2026dRAE,
      title={dRAE: Representation Autoencoder with Hyper-Spherical Codes},
      author={Ma, Tianren and Long, Lin and Chen, Chuyan and Zhang, Mu and Zhao, Junbo and Zhang, Tong and Ye, Qixiang.},
      year={2026},
      note={arXiv:2607.22148}
}
```

## License

MIT. See [LICENSE](LICENSE).

dRAE

Owner Repos

File Explorer

# Use via CDN

Was this content helpful?

Recent Feedback

dRAE

Owner Repos

File Explorer

# Use via CDN

Was this content helpful?

Recent Feedback

Download README