Kelix-AEDIT

Paper | GitHub | 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 and the DC-AE VAE from SANA1.5. 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

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.

PYTHONPATH=. python checkpoints/Kelix-AEDIT/infer.py \
  --ckpt checkpoints/Kelix-AEDIT \
  --image input.jpg \
  --output reconstruction.png

To use existing local checkpoints:

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

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

@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.