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.