| --- |
| language: |
| - en |
| tags: |
| - kelix |
| - image-reconstruction |
| - image-tokenizer |
| - diffusion |
| - flow-matching |
| - safetensors |
| --- |
| |
| # Kelix-AEDIT |
|
|
| [Paper](https://arxiv.org/abs/2602.09843) | [GitHub](https://github.com/Kuaishou-OneRec/Kelix) | [Kelix-DiT](https://huggingface.co/OpenOneRec/Kelix-DiT) |
|
|
| ## Introduction |
|
|
| Kelix-AEDIT reconstructs images from discrete visual tokens using a customized SANA diffusion transformer. This repository releases the DiT weights, an inference script, and its dependencies. |
|
|
| The inference pipeline reuses the image tokenizer from [Kelix-SFT](https://huggingface.co/OpenOneRec/Kelix-SFT) and the [DC-AE VAE from SANA1.5](https://huggingface.co/Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers/tree/main/vae). These components are loaded separately and are not duplicated in this repository. |
|
|
| Unlike Kelix-DiT, which takes Kelix-LLM hidden states as its condition, Kelix-AEDIT operates on quantized image features and supports image reconstruction without the LLM. |
|
|
| ## Model Overview |
|
|
| | Item | Value | |
| |---|---| |
| | Released weights | Customized SANA-DiT | |
| | Image tokenizer | Reused from Kelix-SFT | |
| | Latent VAE | DC-AE from SANA1.5 | |
| | Condition | Quantized image features with mean fusion | |
| | Codebooks | 8 × 8,192 entries | |
| | Default condition grid | 18 × 18 | |
| | Default output resolution | 512 × 512 | |
| | Sampling | 20-step Euler | |
| | Guidance scale | 4.5 | |
| | Flow shift | 3.0 | |
| | Inference precision | BF16 | |
|
|
| ## Usage |
|
|
| Requires an NVIDIA CUDA GPU and a FlashAttention build compatible with your PyTorch and CUDA environment. The inference script uses the model implementations in the Kelix code repository. |
|
|
| ### Installation |
|
|
| ```bash |
| git clone https://github.com/Kuaishou-OneRec/Kelix.git |
| cd Kelix |
| pip install huggingface-hub==0.36.0 |
| hf download OpenOneRec/Kelix-AEDIT --local-dir checkpoints/Kelix-AEDIT |
| pip install -r checkpoints/Kelix-AEDIT/requirements.txt |
| pip install flash-attn==2.7.4.post1 --no-build-isolation |
| ``` |
|
|
| ### Image Reconstruction |
|
|
| Run from the Kelix repository root. The tokenizer and VAE are downloaded automatically on first use. Only the Kelix-SFT shards containing the visual tokenizer are needed; the LLM is not loaded. |
|
|
| ```bash |
| PYTHONPATH=. python checkpoints/Kelix-AEDIT/infer.py \ |
| --ckpt checkpoints/Kelix-AEDIT \ |
| --image input.jpg \ |
| --output reconstruction.png |
| ``` |
|
|
| To use existing local checkpoints: |
|
|
| ```bash |
| PYTHONPATH=. python checkpoints/Kelix-AEDIT/infer.py \ |
| --ckpt checkpoints/Kelix-AEDIT \ |
| --tokenizer /path/to/Kelix-SFT \ |
| --vae /path/to/vae \ |
| --image input.jpg \ |
| --output reconstruction.png |
| ``` |
|
|
| ### Decode Image Tokens |
|
|
| ```bash |
| PYTHONPATH=. python checkpoints/Kelix-AEDIT/infer.py \ |
| --ckpt checkpoints/Kelix-AEDIT \ |
| --tokens tokens.pt \ |
| --grid 18 \ |
| --output generated.png |
| ``` |
|
|
| `tokens.pt` must contain an integer tensor saved with `torch.save`: |
|
|
| - Shape: `[324, 8]` for an 18 × 18 grid. |
| - Values: local indices in `[0, 8192)` for each codebook. |
| - Spatial order: row-major. |
| - Dtype: `torch.int32` or `torch.int64`. |
|
|
| Tokens must use the matching image tokenizer and codebooks. Raw autoregressive vocabulary IDs must be converted to local codebook indices before decoding. |
|
|
| This model does not accept text prompts directly. Text-to-image generation requires a compatible autoregressive model to generate image tokens first. |
|
|
| ## Citation |
|
|
| ```bibtex |
| @article{kelix2026, |
| title = {Kelix Technique Report: Closing the Understanding Gap of Discrete Tokens in Unified Multimodal Models}, |
| author = {Kuaishou Technology}, |
| journal = {arXiv preprint arXiv:2602.09843}, |
| year = {2026}, |
| url = {https://arxiv.org/abs/2602.09843} |
| } |
| ``` |
|
|
| ## License |
|
|
| Please contact the OneRec Team for the license of the Kelix series. The base SANA-DiT, DC-AE, and Keye vision components are subject to their respective original licenses. |
|
|
