File size: 5,942 Bytes
fab088d a6684d5 fab088d a6684d5 fab088d a6684d5 fab088d a6684d5 fab088d a6684d5 fab088d a6684d5 fab088d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | # GenFirst model release manifest
All paths below are relative to the root of this model release. These are the
checkpoints used for the paper results, not the later validation runs.
## VAE bundles
- `vae/e2e-f8d4/`: ImageNet f8d4 VAE initialized from SD-VAE.
- `vae/e2e-f16d32/`: decoder-finetuned f16d32 VAE used by the ImageNet-256
class-conditional results.
- `vae/e2e-f16d32-nodecft/`: pre-decoder-finetune f16d32 VAE used by the
ImageNet-512 results.
- `vae/e2e-flux-vae-sit-imagenet-s2/`: T2I variant A VAE (FLUX.1 init).
- `vae/e2e-flux-vae-mmdit-l2-t2i-s2/`: T2I variant B VAE (FLUX.1 init).
- `vae/qwenvit-siglip-b256-s2/`: Qwen3-VL ViT + ViT decoder VAE shaped with
E2E + SigLIP supervision (Table 9).
- `vae/qwenvit-nlltrunk-s2/`: Qwen3-VL ViT + ViT decoder VAE shaped with
E2E + VLM-NLL supervision (Table 9).
- `vae/e2e-joint-gan-l2-330k/`: Jointly E2E-shaped FLUX.1 image VAE + Cola text
VAE (Table 10).
- `vae/e2e-vae-flux2-sitxl2-auxkl2p5e8-s3-20k/`: 32-channel f8d32 E2E FLUX.2 VAE
(three-stage aspect-ratio-bucketed E2E + decoder finetune).
Every VAE directory contains:
- `<name>.pt`: directly loadable VAE state dict.
- `<name>-latents-stats.pt`: required latent normalization statistics.
- `<name>-discriminator-ckpt.pt`: only needed to resume end-to-end training.
In addition:
- `vae/e2e-joint-gan-l2-330k/` includes `e2e-joint-gan-l2-330k-text-vae.pt`
(standalone E2E-reshaped Cola text VAE state dict).
- `vae/e2e-vae-flux2-sitxl2-auxkl2p5e8-s3-20k/` includes:
- `e2e-vae-flux2-sitxl2-auxkl2p5e8-s3-20k.safetensors` and `diffusion_pytorch_model.safetensors`:
diffusers-compatible safetensors weights.
- `config.json` and `vae_params.json`: diffusers configuration and explicit per-channel
`scaling_factor` and `shift_factor`:
- `--vae_pth=vae/e2e-vae-flux2-sitxl2-auxkl2p5e8-s3-20k/e2e-vae-flux2-sitxl2-auxkl2p5e8-s3-20k.safetensors`
- `--model.vae.scaling_factor=[0.887221,0.886926,0.887096,0.881914,0.886901,0.887307,0.886685,0.886833,0.886540,0.886646,0.887031,0.887073,0.887381,0.886768,0.887299,0.885674,0.886466,0.886795,0.886979,0.886519,0.886949,0.886949,0.887036,0.886828,0.887344,0.887004,0.886816,0.886799,0.887354,0.887523,0.887012,0.886643]`
- `--model.vae.shift_factor=[-0.000623,0.000143,0.000282,0.001599,-0.001093,-0.000477,0.000855,-0.000830,-0.000883,0.001076,-0.000417,0.000200,0.000626,-0.000649,-0.000137,0.001453,-0.000816,-0.000307,-0.000290,0.000662,0.000484,0.000993,0.000574,-0.000139,-0.000914,0.000578,-0.000691,-0.000323,0.000619,-0.000606,-0.000372,0.000437]`
- `e2e-vae-flux2-sitxl2-auxkl2p5e8-s3-20k-stage2-bn-latents-stats.pt`: Stage-2 70k SiT EMA BatchNorm statistics.
Keep each VAE and its latent statistics together.
## Class-conditional priors
ImageNet-256, paired with `vae/e2e-f16d32/`:
- `class-conditional/imagenet256/eit-imagenet256-weak-20ep.pt`: weak model
required for the paper's autoguidance evaluation.
- `class-conditional/imagenet256/eit-imagenet256-80ep.pt`
- `class-conditional/imagenet256/eit-imagenet256-480ep.pt`
- `class-conditional/imagenet256/eit-imagenet256-800ep.pt`
ImageNet-512, paired with `vae/e2e-f16d32-nodecft/`:
- `class-conditional/imagenet512/eit-imagenet512-80ep.pt`
- `class-conditional/imagenet512/eit-imagenet512-400ep.pt`
These prior checkpoints are directly accepted by
`Prior_only/generate.py --ckpt` and `Prior_only/scripts/eval_fid.sh` through its
`CKPT` override. Do not convert or strip them.
## Text-to-image priors
- `text-to-image/variant-a/eit-t2i-variant-a-512.pt`, paired with
`vae/e2e-flux-vae-sit-imagenet-s2/`.
- `text-to-image/variant-b/eit-t2i-variant-b-512.pt`, paired with
`vae/e2e-flux-vae-mmdit-l2-t2i-s2/`.
These checkpoints are directly accepted by `Prior_only/infer_t2i.py` and the
GenEval/DPG generation scripts. Pass the matching VAE, latent statistics and
`--text-encoder-name openai/clip-vit-large-patch14`; no checkpoint conversion
is required.
## Shared visual latents (generation + understanding, Table 9)
- `shared-latents/siglip/eit-qwenvit-siglip-t2i-512.pt`, paired with
`vae/qwenvit-siglip-b256-s2/` (`0.90` GenEval, `83.76` DPG-Bench, `79.25%` IN1K linear probe).
- `shared-latents/vlm-nll/eit-qwenvit-vlm-nll-t2i-512.pt`, paired with
`vae/qwenvit-nlltrunk-s2/` (`0.89` GenEval, `83.49` DPG-Bench, `81.69%` IN1K linear probe).
Directly accepted by `Prior_only/infer_t2i.py` and `Prior_only/gen_for_bench.py`
with `--vae-arch qwenvit --vae-vit-path Qwen/Qwen3-VL-2B-Instruct`.
## Unified text-image modeling (`JointMMDiT-XL/2`, Table 10)
- `unified-text-image/e2e-joint/eit-joint-e2e-text-e2e-img-512-sft.pt`: Phase-3
512x512 SFT checkpoint for text-to-image (`86.57` GenEval, `72.68` DPG-Bench).
- `unified-text-image/e2e-joint/eit-joint-e2e-text-e2e-img-512.pt`: Phase-2
512x512 checkpoint before SFT for image-to-text COCO captioning (`33.28` CIDEr,
`57.97` CLIPScore).
- Both pair with `vae/e2e-joint-gan-l2-330k/` (`--image-vae-ckpt vae/e2e-joint-gan-l2-330k/e2e-joint-gan-l2-330k.pt`;
the reshaped text VAE weights are already embedded inside `ckpt["text_vae"]`).
## Checkpoint formats produced by the repository
- `train_prior_only.py`, `train_t2i_frozen.py`, and `train_joint_frozen.py` save
final checkpoints in direct evaluation format.
- Class-conditional `train_e2e.py` checkpoints can be used directly by the
top-level `generate.py`; all `train_e2e.py`, `train_e2e_vae.py`, and
`train_decoder_ft.py` checkpoints can be evaluated directly for reconstruction.
Before using either checkpoint's VAE with a separate prior, run
`Prior_only/scripts/run_prior_only_pipeline.sh extract` (or `save_joint_vae_weights.py`
for `train_joint_e2e.py`), which produces a standalone VAE bundle.
- A downloaded VAE bundle is already extracted and must not be extracted again.
Weights are released under CC BY-NC 4.0. See the repository's
`MODEL_CARD.md`, `NOTICE`, and documentation for complete terms and commands.
|