Update model card: latent-space table + LDM HF links, SDXL/Qwen-Image VAE, flux2 2kto4k _2606 ckpt, updated demo cmd

README.md

@@ -39,7 +39,7 @@ super-resolved image in one pass. This repository hosts the released decoder
 checkpoints, plus the encoder/decoder ("VAE") weights they depend on.
 
 All `PiD_*` checkpoints in this repo are **4-step distilled**. The non-`PiD_*`
-entries (`ae.safetensors`, `flux2_ae.safetensors`, `sd3_vae/`, `rae/`,
+entries (`ae.safetensors`, `flux2_ae.safetensors`, `sdxl_vae.safetensors`, `QwenImage_VAE_2d.pth`, `sd3_vae/`, `rae/`,
 `scale_rae/`) are **the corresponding encoder/decoder VAE weights** that PiD
 plugs into — they're not PiD checkpoints themselves.
 
@@ -56,41 +56,56 @@ Two variants are released for each diffusers-style backbone:
 
 - **`2k`** — trained at 2048px, used as a 4× decoder (512 LDM → 2048 px), or as
   an 8× decoder for the Scale-RAE backbone (256 → 2048).
-- **`2kto4k`** — trained with multi-resolution data bucketing 2048→3840 and an
+- **`2kto4k`** — trained with multi-resolution data bucketing 2048→4096 and an
   SD3-style dynamic shift; designed for 1024 LDM → 4K (4096 px) decoding.
 
 Both checkpoint variants support multiple aspect ratios.
 
-| Path                                                          | Backbone (encoder side)                    | SR factor | Variant   |
-|---------------------------------------------------------------|--------------------------------------------|-----------|-----------|
-| `checkpoints/PiD_res2k_sr4x_official_flux_distill_4step`      | Flux1-dev (16-ch VAE)                      | 4×        | 2k        |
-| `checkpoints/PiD_res2k_sr4x_official_flux2_distill_4step`     | Flux2-dev (128-ch BN VAE)                  | 4×        | 2k        |
-| `checkpoints/PiD_res2k_sr4x_official_sd3_distill_4step`       | SD3 medium (16-ch VAE)                     | 4×        | 2k        |
-| `checkpoints/PiD_res2k_sr4x_official_dinov2_distill_4step`    | DINOv2-B + RAE ViT-XL (768-ch)             | 4×        | 2k        |
-| `checkpoints/PiD_res2k_sr8x_official_siglip_distill_4step`    | SigLIP-2 So400M + Scale-RAE ViT-XL (1152)  | 8×        | 2k        |
-| `checkpoints/PiD_res2kto4k_sr4x_official_flux_distill_4step`  | Flux1-dev (16-ch VAE)                      | 4×        | 2kto4k    |
-| `checkpoints/PiD_res2kto4k_sr4x_official_flux2_distill_4step` | Flux2-dev (128-ch BN VAE)                  | 4×        | 2kto4k    |
-| `checkpoints/PiD_res2kto4k_sr4x_official_sd3_distill_4step`   | SD3 medium (16-ch VAE)                     | 4×        | 2kto4k    |
-
-Z-Image shares Flux1's VAE, so its inference path reuses the `flux` checkpoints
-(both `2k` and `2kto4k`) — no separate `zimage` checkpoint is shipped.
+| Path                                                            | Latent space | SR factor | Variant |
+|-----------------------------------------------------------------|--------------|-----------|---------|
+| `checkpoints/PiD_res2k_sr4x_official_flux_distill_4step`        | Flux1-dev    | 4×        | 2k      |
+| `checkpoints/PiD_res2k_sr4x_official_flux2_distill_4step`       | Flux2-dev    | 4×        | 2k      |
+| `checkpoints/PiD_res2k_sr4x_official_sd3_distill_4step`         | SD3 medium   | 4×        | 2k      |
+| `checkpoints/PiD_res2k_sr4x_official_dinov2_distill_4step`      | DINOv2-B     | 4×        | 2k      |
+| `checkpoints/PiD_res2k_sr8x_official_siglip_distill_4step`      | SigLIP-2     | 8×        | 2k      |
+| `checkpoints/PiD_res2kto4k_sr4x_official_flux_distill_4step`    | Flux1-dev    | 4×        | 2kto4k  |
+| `checkpoints/PiD_res2kto4k_sr4x_official_flux2_distill_4step_2606` | Flux2-dev  | 4×        | 2kto4k  |
+| `checkpoints/PiD_res2kto4k_sr4x_official_sd3_distill_4step`     | SD3 medium   | 4×        | 2kto4k  |
+| `checkpoints/PiD_res2kto4k_sr4x_official_sdxl_distill_4step`    | SDXL         | 4×        | 2kto4k  |
+| `checkpoints/PiD_res2kto4k_sr4x_official_qwenimage_distill_4step` | Qwen-Image  | 4×        | 2kto4k  |
 
 Each directory contains a single file, `model_ema_bf16.pth`, which is the EMA
 weights cast to bfloat16 — the format the inference scripts load by default.
 
-## VAE / encoder weights
-
-These are the per-backbone encoder (and, where applicable, original decoder)
-weights that PiD pairs with. They're hosted here so a single download brings
-everything needed end-to-end.
-
-| Path                            | Description                                                                          |
-|---------------------------------|--------------------------------------------------------------------------------------|
-| `checkpoints/ae.safetensors`    | Flux1-dev / Z-Image 16-ch VAE (encoder + original Flux decoder).                     |
-| `checkpoints/flux2_ae.safetensors` | Flux2-dev 128-ch BN VAE.                                                          |
-| `checkpoints/sd3_vae/`          | SD3 medium 16-ch VAE in diffusers format.                                            |
-| `checkpoints/rae/`              | DINOv2-B image encoder + RAE ViT-XL decoder + ImageNet-512 normalization statistics. |
-| `checkpoints/scale_rae/`        | SigLIP-2 So400M encoder + Scale-RAE ViT-XL decoder + decoder config.                 |
+> **⚠️ Flux2-dev `2kto4k` — use the new `_2606` checkpoint.** The previous
+> `PiD_res2kto4k_sr4x_official_flux2_distill_4step` (without the `_2606` suffix)
+> suffered from a color-drifting issue. The new
+> `PiD_res2kto4k_sr4x_official_flux2_distill_4step_2606` fixes it — please use it
+> and do **not** use the old one. See the
+> [comparison](https://github.com/nv-tlabs/pid/blob/main/docs/FLUX2_2kto4k_new_ckpt_compare.md)
+> for details.
+
+### Latent space → compatible LDMs
+
+A PiD decoder is tied to a *latent space*, not to a single generative model. Any
+LDM that produces latents in that space can reuse the same checkpoint. The
+`--backbone` aliases below pick the right LDM pipeline; they all decode through
+the latent space's checkpoint above.
+
+| Latent space | VAE / vision encoder weights       | compatible `--backbone`                   | Corresponding LDM Links |
+|--------------|------------------------------------|-------------------------------------------|-----------------|
+| Flux1-dev    | `checkpoints/ae.safetensors`       | `flux`, `zimage`, `zimage-turbo`          | [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev), [Z-Image](https://huggingface.co/Tongyi-MAI/Z-Image), [Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) |
+| Flux2-dev    | `checkpoints/flux2_ae.safetensors` | `flux2`, `flux2-klein-4b`, `flux2-klein-9b` | [FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev), [FLUX.2-klein-4B](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B), [FLUX.2-klein-9B](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B) |
+| SD3 medium   | `checkpoints/sd3_vae/`             | `sd3`                                     | [SD3-medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers) |
+| SDXL         | `checkpoints/sdxl_vae.safetensors` | `sdxl`                                    | [SDXL-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) |
+| Qwen-Image   | `checkpoints/QwenImage_VAE_2d.pth` | `qwenimage`, `qwenimage-2512`             | [Qwen-Image](https://huggingface.co/Qwen/Qwen-Image), [Qwen-Image-2512](https://huggingface.co/Qwen/Qwen-Image-2512) |
+| DINOv2-B     | `checkpoints/rae/`                 | `dinov2`                                  | [RAE](https://github.com/bytetriper/RAE) (class-conditional; DINOv2-B) |
+| SigLIP-2     | `checkpoints/scale_rae/`           | `siglip`                                  | [Scale-RAE](https://github.com/ZitengWangNYU/Scale-RAE) (text-conditional; nyu-visionx/Scale-RAE-Qwen1.5B_DiT2.4B) |
+
+For example, Z-Image and Z-Image-Turbo share Flux1-dev's VAE, so they reuse the
+`flux` checkpoints (both `2k` and `2kto4k`) — no separate `zimage` checkpoint is
+shipped. Likewise `qwenimage-2512` reuses the `qwenimage` decoder (same VAE,
+different transformer). 
 
 ## Usage
 
@@ -107,11 +122,11 @@ automatically:
 hf download nvidia/PiD --local-dir . --include "checkpoints/*"
 
 # Then run any of the demos, e.g.:
-PYTHONPATH=. python -m pid._src.inference.from_ldm_flux \
-    --prompt "A photorealistic cat" \
-    --ldm_inference_steps 28 --save_xt_steps 22 24 26 \
-    --output_dir ./results/demo \
-    --cfg_scale 1 --pid_inference_steps 4 --scale 4
+PYTHONPATH=. python -m pid._src.inference.from_ldm --backbone flux \
+    --prompt "A photorealistic half-body portrait of a brown tabby cat with bold stripes sitting attentively on a rustic wooden kitchen table, soft morning light streaming sideways through a large window, fine fur detail and stripe patterns sharply visible, intense amber-green eyes in razor-sharp focus, warm farmhouse kitchen softly out of focus, cinematic shallow depth of field, ultra-detailed fur texture, photorealistic" \
+    --ldm_inference_steps 28 --save_xt_steps 24 \
+    --output_dir ./results/official_demo/flux \
+    --pid_inference_steps 4
 ```
 
 Pick the `2kto4k` variant via `--pid_ckpt_type 2kto4k` when decoding at 4K.