Lingmin-Ran/SADAE

STARFlow: Scalable Transformer Auto-Regressive Flow

This is the official open source release of STARFlow and STARFlow-V, state-of-the-art transformer autoregressive flow models for high-quality image and video generation.

📖 Overview

STARFlow introduces a novel transformer autoregressive flow architecture that combines the expressiveness of autoregressive models with the efficiency of normalizing flows. The model achieves state-of-the-art results in both text-to-image and text-to-video generation tasks.

• STARFlow: Scaling Latent Normalizing Flows for High-resolution Image Synthesis (NeurIPS 2025 Spotlight)
• STARFlow-V: End-to-End Video Generative Modeling with Normalizing Flows (CVPR 2026 Highlight)
• STARFlow2: STARFlow2: Bridging Language Models and Normalizing Flows for Unified Multimodal Generation (Code TBD)
• NTM: Normalizing Trajectory Models (Code TBD)

🎬 View Video Results Gallery - See examples of generated videos and comparisons

🚀 Quick Start

Environment Setup

# Clone the repository
git clone https://github.com/apple/ml-starflow
cd ml-starflow

# Set up conda environment (recommended)
bash scripts/setup_conda.sh

# Or install dependencies manually
pip install -r requirements.txt

Model Checkpoints

Important: You'll need to download the pretrained model checkpoints and place them in the ckpts/ directory. For example:

• ckpts/starflow_3B_t2i_256x256.pth - For text-to-image generation
• ckpts/starflow-v_7B_t2v_caus_480p_v3.pth - For text-to-video generation

The checkpoint files are not included in this repository due to size constraints. Please download via Hugging Face.

For text-to-image inference, pre-download the checkpoint, FLAN-T5-XL text encoder, and SD-VAE with the workspace proxy settings:

bash scripts/download_t2i_models.sh
# Check the cache later without accessing the network:
bash scripts/download_t2i_models.sh --check-only

The script uses the pvd Conda environment and caches the Hugging Face models under ckpts/hf_cache by default. Run inference from this directory with the same cache in offline mode to fail promptly if any model is missing:

HF_HUB_CACHE="$PWD/ckpts/hf_cache" HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 \
    bash scripts/test_sample_image.sh "a film still of a cat playing piano"

Text-to-Image Generation

Generate high-quality images from text prompts:

# Basic image generation (256x256)
bash scripts/test_sample_image.sh "a film still of a cat playing piano"

# Custom prompt and settings
torchrun --standalone --nproc_per_node 1 sample.py \
    --model_config_path "configs/starflow_3B_t2i_256x256.yaml" \
    --checkpoint_path "ckpts/starflow_3B_t2i_256x256.pth" \
    --caption "your custom prompt here" \
    --sample_batch_size 8 \
    --cfg 3.6 \
    --aspect_ratio "1:1" \
    --seed 999

SADAE editing and PIE-Bench evaluation

SADAE 支持原版 PIE-Bench（mapping_file.json + annotation_images/）和 PIE-Bench++ Parquet。 通过 --dataset 指定原版数据集，编辑结束后自动计算指标：

PYTHON=/mnt/bn/jiangzhongtao/miniconda3/envs/pvd/bin/python \
CUDA_VISIBLE_DEVICES=0 \
bash scripts/nf_source_anchored_decoupled_affine_editing.sh \
    --dataset /mnt/bn/jiangzhongtao/lingmin/dataset/PIE_Bench \
    --max-samples 4 \
    --output-dir result/sadae_pie_bench

--max-samples 0 处理全部 700 条样本；--category 1_change_object_80 或 --image-id 140 可筛选样本， ID 输出保留原版的前导零。未指定 --dataset 时仍使用 PIE-Bench++。 当前 Shell 脚本运行三组参数：alpha=1.5、gamma=1.0、beta=0/0.2/0.4，每组单独保存图像和指标。 --output-dir 必须是尚不存在的新目录。

指标沿用 PIE-Bench 官方评测实现 的定义：

指标	含义	方向
structure_distance	DINO 特征自相似矩阵的结构距离	越小越好
psnr_unedit_part	背景 PSNR（dB）	越大越好
lpips_unedit_part	背景 LPIPS	越小越好
mse_unedit_part	背景 MSE	越小越好
ssim_unedit_part	背景 SSIM	越大越好
clip_similarity_target_image	编辑图像与目标提示词的 CLIP 相似度	越大越好
clip_similarity_target_image_edit_part	编辑区域与目标提示词的 CLIP 相似度	越大越好

每组输出的 metrics.csv 保存逐图指标，metadata/<category>/<id>.json 保存对应图像的指标， summary.json 中的 pie_bench_metrics.statistics 保存各指标的均值、总体标准差和有效样本数。 指标使用原始数值，CLIP 分数范围为 0–100。评测使用保存的 PNG 和与输入图像同步缩放、裁剪的标注掩码； 掩码只用于评测。当前默认模型在 256×256 分辨率评测，与官方常用的 512×512 结果比较时需要区分分辨率。 背景指标按官方方式将编辑区域置零后对整图归约。区域为空时对应指标为 null，完全一致背景的 PSNR 为字符串 "inf"。

默认在 GPU 评测，也可用 --metrics-device cpu，或用 --skip-metrics 仅生成图像。 评测需要 CLIP ViT-L/14、DINO ViT-B/8 和 LPIPS SqueezeNet 权重；当前离线环境分别从 ckpts/hf_cache 和 ckpts/torch_hub 读取。单张自定义图像缺少区域标注，自动跳过 PIE-Bench 评测。

PIE-Bench++ NF reconstruction and direct editing

Run source-caption inversion and compare the reverse flow with the original image. The same source image is inverted again, then reversed with the dataset's target caption:

# Reproduce image 48 with FP32 NF, preserving historical BF16 results
bash scripts/run_pie_bench_nf_baseline.sh --image-id 48 --save-z --output-dir result/baseline_fp32

# All images with the validated batch size 1
bash scripts/run_pie_bench_nf_baseline.sh --output-dir result/baseline_fp32

# If GPU memory allows, increase batch size
bash scripts/run_pie_bench_nf_baseline.sh --batch-size 4 --output-dir result/baseline_fp32

# Optional: match the VAE sampling and latent noise used during training
bash scripts/run_pie_bench_nf_baseline.sh --vae-latent sample --input-noise-std 0.3 --overwrite

By default, the script reads /mnt/bn/jiangzhongtao/lingmin/dataset/PIE_Bench_pp and writes result/baseline/reconstruct/ and result/baseline/editing/, organized by category and image ID. source/ contains the 256×256 center crop used as input; vae_only/ shows the VAE reconstruction without NF; metadata/ records prompts, NF latent round-trip error, and pixel metrics. The default batch size is 1; editing with CFG doubles the internal batch. NF defaults to true FP32, using the same shared checkpoint loader as the precision experiment: strict checkpoint loading, no FSDP casting, no NF autocast, and TF32 disabled. VAE and text encoding retain BF16. Use --precision bf16 to reproduce the historical precision policy. Completed outputs with a different precision (including historical metadata without a precision field) cause an error instead of being silently skipped; use a new output directory or --overwrite to regenerate them. The default uses the VAE posterior mean with no added noise so reconstruction can be compared directly. STARFlow was trained with sampled VAE latents and noise standard deviation 0.3, which may give different editing behavior. Direct caption replacement is a baseline experiment and does not apply the dataset masks.

For high-CFG reconstruction, opt into the closed-loop encoder:

bash scripts/run_pie_bench_nf_baseline.sh --precision fp32 \
    --reconstruct-encoder closed-loop --reconstruct-cfg 3.6 \
    --image-id 39 --save-z --diagnose-cfg \
    --output-dir result/baseline_fp32_cfg36_closed_loop

This mode computes the deepest block token by token, feeding its simulated reconstruction into subsequent steps. Network weights, prefixes, RoPE and KV caches remain FP32; scale conversion, CFG, affine updates and reconstruction z use FP64. Decoding starts with fresh caches and receives no source prefixes. Shallow blocks retain the existing FP32 path. --save-z additionally writes the FP64 reconstruction tensor under z_reconstruct/; z/ retains the original unguided editing z. Closed-loop decoding requires the FP64 tensor without a .float() conversion and uses nf_roundtrip.closed_loop_reverse; the original model(..., reverse=True) path is unchanged. --diagnose-cfg records closed_loop_diagnostics, comparing the independently decoded deep state with the source state and encoder feedback. Replay uses the same batch grouping (reconstruction_batch_ids in metadata), model precision and device/backend to keep token predictions numerically aligned. The editing branch retains its existing behavior. Use separate output directories for parallel and closed-loop runs; resume checks reject mixed encoder modes.

NF source/target likelihood ratio

Score each saved reconstruction and edit under both its source and target prompts:

bash scripts/compare_nf_target_nll.sh --score-mode ratio --batch-size 4

# Optional smoke run; use a separate output directory for the full run.
bash scripts/compare_nf_target_nll.sh --score-mode ratio --max-samples 2 \
    --output-dir result/nll_ratio_smoke

The metric is R(x) = NLL(x | source_prompt) - NLL(x | target_prompt) and delta_R = R(editing) - R(reconstruct). Positive R means the target condition assigns higher density than the source condition; positive delta_R means editing increased that relative preference. It does not measure content preservation or perceptual quality.

The default output is result/baseline_fp32/nll_ratio/0_random_140/. Select another experiment/category with --baseline-dir and --category. Output directories must not already exist. The original target-only mode remains the default and writes to nll_target/<category>/.

• scores.jsonl: four records per image ID (reconstruct/editing × source/target), including condition, conditioning_prompt, NLL, prior NLL and log determinant.
• comparison.csv: target NLLs in the existing reconstruct_nll/editing_nll columns, additional reconstruct_source_nll/editing_source_nll, reconstruct_r, editing_r, delta_r, and their per-dimension values.
• summary.json: run settings, original target-NLL statistics, mean source NLLs, mean R values, mean/median delta_R, and r_improved_count/r_improved_fraction.

Both conditions use the exact same VAE posterior mean of each saved PNG, zero added noise and FP32 NF inference without CFG or TF32; VAE/text processing retains BF16. Source and target are scored sequentially without doubling the NF batch. Scores are in nats in the scaled VAE latent space (not pixel-space likelihood), with normalized scores in nats per latent dimension. Ratio mode recomputes both conditions; it does not mix new scores with historical target-NLL results.

NF reconstruction precision experiments

scripts/run_pie_bench_nf_precision.sh compares the original BF16/FSDP inference with FP32 NF inference. It fixes the VAE posterior to its mean, adds no noise, uses the source caption in both directions, and reconstructs with CFG disabled and exact autoregressive decoding (jacobi=False). VAE/text processing keeps the baseline precision. The default runs one image with batch size 1.

# Compare the same dataset ID and batch size; ID 48 is not a row offset.
bash scripts/run_pie_bench_nf_precision.sh --image-id 48 --precision bf16
bash scripts/run_pie_bench_nf_precision.sh --image-id 48 --precision fp32 --save-z

# If FP32 still has substantial NF error: align encoder/decoder token kernels.
# Sequential encoding is slower, and math attention can use more memory.
bash scripts/run_pie_bench_nf_precision.sh --image-id 48 \
    --precision fp32 --encoder sequential --attention math --save-z

# Parallel FP32 with the same math backend isolates the encoder change.
bash scripts/run_pie_bench_nf_precision.sh --image-id 48 \
    --precision fp32 --encoder parallel --attention math

# Compare all 49 images with existing baseline outputs.
# Run once with bf16 and once with fp32, using separate output directories.
bash scripts/run_pie_bench_nf_precision.sh --category 0_random_140 \
    --max-samples 49 --batch-size 1 --precision fp32 \
    --output-dir result/nf_fp32_first49

# Optional editing: reuse the encoded z with the dataset's target caption.
bash scripts/run_pie_bench_nf_precision.sh --image-id 48 --precision fp32 \
    --edit --edit-cfg 3.6 --output-dir result/nf_fp32_edit48

# CPU checks of the numerical implementation on small random models.
/mnt/bn/jiangzhongtao/miniconda3/envs/pvd/bin/python scripts/test_nf_roundtrip.py

The default output is result/nf_<precision>_<encoder>_<attention>/, with source/, vae_only/, reconstruct/, per-image metadata/, summary.json, and run_config.json. Use a new directory or --overwrite for another run. --max-samples 0 processes all selected images; --image-id can be repeated (also set an appropriate --max-samples). --edit adds editing/ and --save-z saves the latent without reducing its storage precision.

Inspect latent_error, reconstruct_vs_source, and reconstruct_vs_vae. The last metric measures additional image error introduced by NF; a null PSNR with zero MSE means identical tensors. blocks reports cumulative inverse error against the corresponding forward state, in reverse block order 5 to 0, plus forward affine scale ranges. These cumulative errors are not independent tests of individual blocks. historical_baseline, when available, contains the old metrics; batch size/backend changes can affect numerics, so use the new BF16 run for a controlled comparison and check that VAE-only results agree.

Why precision is a concern: MetaBlock.forward casts the normalized latent to the transformer's output dtype, and MetaBlock.reverse casts each decoded block back to its incoming dtype. utils/training.py also forces BF16 parameters, inputs and outputs through FSDP. Converting an already rounded z to FP32 cannot recover those bits. The FP32 experiment loads the original checkpoint strictly, bypasses FSDP, disables NF autocast and TF32, and retains FP32 intermediate latents. It requires more memory and runs slower. The checkpoint contains about 3.86 billion parameters, roughly 15.4 GB of FP32 parameter storage before text encoder, VAE, KV caches and working memory; start with batch size 1.

The sequential encoder uses the same reverse_step and text-prefix KV-cache path as the decoder to compute affine parameters from known prefixes. Decoding uses fresh caches and only z plus the caption. This reduces discrepancies between full-sequence and single-token computation; it does not guarantee exact inversion for an ill-conditioned checkpoint.

On the 49 existing results inspected on 2026-09-30, mean VAE-only PSNR was 27.75 dB versus 20.24 dB after NF. Image 000000000048 had VAE-only PSNR 34.36 dB, NF reconstruction PSNR 9.67 dB and latent MAE 0.8030. Repeated forward z was identical for all 49 images. These establish NF round-trip error; they do not by themselves prove which numerical mechanism dominates. The CPU checks validate non-identity toy flows, sequential encoding, block alignment, and cache isolation. CUDA was unavailable in the analysis session, so the released checkpoint's improvement must be measured with the commands above.

Related reading: STARFlow describes autoregressive affine flows, latent noise and training stability; Behrmann et al., AISTATS 2021 analyzes numerical non-invertibility and unstable inverses; PyTorch numerical accuracy explains why mathematically equivalent batched and sliced computations need not produce identical floating point results. These motivate the experiments; they do not establish a measured fix for this particular image.

Text-to-Video Generation

Generate videos from text descriptions:

# Basic video generation (480p, ~5 seconds)
bash scripts/test_sample_video.sh "a corgi dog looks at the camera"

# With custom input image for TI2V video generation
bash scripts/test_sample_video.sh "a cat playing piano" "/path/to/input/image.jpg"

# Longer video generation (specify target length in frames)
bash scripts/test_sample_video.sh "a corgi dog looks at the camera" "none" 241  # ~15 seconds at 16fps
bash scripts/test_sample_video.sh "a corgi dog looks at the camera" "none" 481  # ~30 seconds at 16fps

# Advanced video generation
torchrun --standalone --nproc_per_node 8 sample.py \
    --model_config_path "configs/starflow-v_7B_t2v_caus_480p.yaml" \
    --checkpoint_path "ckpts/starflow-v_7B_t2v_caus_480p_v3.pth" \
    --caption "your video prompt here" \
    --sample_batch_size 1 \
    --cfg 3.5 \
    --aspect_ratio "16:9" \
    --out_fps 16 \
    --jacobi 1 --jacobi_th 0.001 \
    --target_length 161  # Customize video length

🛠️ Training

Image Training

Train your own STARFlow model for text-to-image generation:

# Quick training test
bash scripts/test_train_image.sh 10 16

# Full training with custom parameters
torchrun --standalone --nproc_per_node 8 train.py \
    --model_config_path "configs/starflow_3B_t2i_256x256.yaml" \
    --epochs 100 \
    --batch_size 1024 \
    --wandb_name "my_starflow_training"

Video Training

Train STARFlow-V for text-to-video generation:

# Quick training test
bash scripts/test_train_video.sh 10 8

# Resume training from checkpoint
torchrun --standalone --nproc_per_node 8 train.py \
    --model_config_path "configs/starflow-v_7B_t2v_caus_480p.yaml" \
    --resume_path "ckpts/starflow-v_7B_t2v_caus_480p_v3.pth" \
    --epochs 100 \
    --batch_size 192

🔧 Utilities

Video Processing

Extract individual frames from multi-video grids:

# Extract frames from a video containing multiple video grids
python scripts/extract_image_from_video.py --input_video path/to/video.mp4 --output_dir output/

# Extract images with custom settings
python scripts/extract_images.py input_file.mp4

📁 Model Architecture

STARFlow (3B Parameters - Text-to-Image)

• Resolution: 256×256
• Architecture: 6-block deep-shallow architecture
• Text Encoder: T5-XL
• VAE: SD-VAE
• Features: RoPE positional encoding, mixed precision training

STARFlow-V (7B Parameters - Text-to-Video)

• Resolution: Up to 640×480 (480p)
• Temporal: 81 frames (16 FPS = ~5 seconds)
• Architecture: 6-block deep-shallow architecture (full sequence)
• Text Encoder: T5-XL
• VAE: WAN2.2-VAE
• Features: Causal attention, autoregressive generation, variable length support

🔧 Key Features

• Autoregressive Flow Architecture: Novel combination of autoregressive models and normalizing flows
• High-Quality Generation: Competetive FID scores and visual quality to State-of-the-art Diffusion Models
• Flexible Resolution: Support for various aspect ratios and resolutions
• Efficient Training: FSDP support for large-scale distributed training
• Fast Sampling: Block-wise Jacobi iteration for accelerated inference
• Text Conditioning: Advanced text-to-image/video capabilities
• Video Generation: Temporal consistency and smooth motion

📊 Configuration

Key Parameters

Image Generation (starflow_3B_t2i_256x256.yaml)

• img_size: 256 - Output image resolution
• txt_size: 128 - Text sequence length
• channels: 3072 - Model hidden dimension
• cfg: 3.6 - Classifier-free guidance scale
• noise_std: 0.3 - Flow noise standard deviation

Video Generation (starflow-v_7B_t2v_caus_480p.yaml)

• img_size: 640 - Video frame resolution
• vid_size: '81:16' - Temporal dimensions (frames:downsampling)
• fps_cond: 1 - FPS conditioning enabled
• temporal_causal: 1 - Causal temporal attention

Sampling Options

• --cfg - Classifier-free guidance scale (higher = more prompt adherence)
• --jacobi - Enable Jacobi iteration for faster sampling
• --jacobi_th - Jacobi convergence threshold
• --jacobi_block_size - Block size for Jacobi iteration
• --aspect_ratio - Output aspect ratio ("1:1", "16:9", "4:3", etc.)
• --seed - Random seed for reproducible generation

📚 Project Structure

├── train.py               # Main training script
├── sample.py              # Sampling and inference
├── transformer_flow.py    # Core model implementation
├── dataset.py             # Dataset loading and preprocessing
├── finetune_decoder.py    # Decoder fine-tuning script
├── utils/                 # Utility modules
│   ├── common.py         # Core utility functions
│   ├── model_setup.py    # Model configuration and setup
│   ├── training.py       # Training utilities and metrics
│   └── inference.py      # Evaluation and metrics
├── configs/              # Model configuration files
│   ├── starflow_3B_t2i_256x256.yaml
│   └── starflow-v_7B_t2v_caus_480p.yaml
├── scripts/                 # Example training and sampling scripts
│   ├── test_sample_image.sh
│   ├── test_sample_video.sh
│   ├── test_train_image.sh
│   ├── test_train_video.sh
│   ├── setup_conda.sh
│   ├── extract_images.py
│   └── extract_image_from_video.py
└── misc/                  # Additional utilities
    ├── pe.py             # Positional encodings
    ├── lpips.py          # LPIPS loss
    └── wan_vae2.py       # Video VAE implementation

💡 Tips

Image Generation

1. Use guidance scales between 2.0-5.0 for balanced quality and diversity
2. Experiment with different aspect ratios for your use case
3. Enable Jacobi iteration (--jacobi 1) for faster sampling
4. Use higher resolution models for detailed outputs
5. The default script uses optimized settings: --jacobi_th 0.001 and --jacobi_block_size 16

Video Generation

1. Start with shorter sequences (81 frames) and gradually increase length (161, 241, 481+ frames)
2. Use input images (--input_image) for more controlled generation
3. Adjust FPS settings based on content type (8-24 FPS)
4. Consider temporal consistency when crafting prompts
5. The default script uses --jacobi_block_size 64.
6. Longer videos: Use --target_length to generate videos beyond the training length (requires --jacobi 1)
7. Frame reference: 81 frames ≈ 5s, 161 frames ≈ 10s, 241 frames ≈ 15s, 481 frames ≈ 30s (at 16fps)

Training

1. Use FSDP for efficient large model training
2. Start with smaller batch sizes and scale up
3. Monitor loss curves and adjust learning rates accordingly
4. Use gradient checkpointing to reduce memory usage
5. The test scripts include --dry_run 1 for validation

🔗 Citation

If you use STARFlow in your research, please cite:

@article{gu2025starflow,
  title={STARFlow: Scaling Latent Normalizing Flows for High-resolution Image Synthesis},
  author={Gu, Jiatao and Chen, Tianrong and Berthelot, David and Zheng, Huangjie and Wang, Yuyang and Zhang, Ruixiang and Dinh, Laurent and Bautista, Miguel Angel and Susskind, Josh and Zhai, Shuangfei},
  journal={NeurIPS},
  year={2025}
}

📄 License

LICENSE: Please check out the repository LICENSE before using the provided code and LICENSE_MODEL for the released models.