WAM-Diff2-B32-2B

WAM-Diff2-B32-2B is the 2B-parameter Block-32 checkpoint for WAM-Diff2: Hierarchical AR-to-Diffusion Distillation for Highly Efficient Autonomous Driving VLA.

• Code
• Project page
• Paper

WAM-Diff2 converts an autoregressive vision-language-action model into a block-diffusion model. This checkpoint supports single-image and multi-image inputs across autonomous-driving planning, visual question answering (VQA), and object detection (DET) tasks.

Model details

The default inference path uses dynamic low-confidence remasking with a confidence threshold of 0.9.

Supported tasks

• Autonomous-driving planning: predict future driving waypoints from images, navigation instructions, and ego-vehicle state.
• Visual question answering (VQA): answer questions grounded in one or more input images.
• Object detection (DET): identify and localize requested objects from visual inputs.

All tasks use the same multimodal conversational interface. Prompt templates and output representations are task-dependent; use the template and decoding convention associated with the corresponding training or evaluation dataset.

Installation

This checkpoint uses the custom wam_diff2 architecture. It is not loadable with an unmodified Transformers installation alone. Install the official WAM-Diff2 repository before using the checkpoint.

Python 3.11–3.12 is supported. The reference environment uses PyTorch 2.8 and CUDA 12.x.

git clone https://github.com/fudan-generative-vision/WAM-Diff2.git
cd WAM-Diff2

conda create -n wam-diff2 python=3.12 -y
conda activate wam-diff2
pip install -r environment/requirements_cuda.txt
pip install -e .

PyTorch SDPA is used by default. FlashAttention is optional and must match the installed PyTorch and CUDA versions.

Inference

Prepare an evaluation JSON file and run the official inference entry point:

scripts/infer_gpu.sh \
  --model_id fudan-generative-ai/WAM-Diff2 \
  --input_file /path/to/eval.json \
  --output_file outputs/predictions.json

For multiple GPUs, set NUM_GPUS:

NUM_GPUS=8 scripts/infer_gpu.sh \
  --model_id fudan-generative-ai/WAM-Diff2 \
  --input_file /path/to/eval.json \
  --output_file outputs/predictions.json

Ascend NPU inference is also available with a compatible torch-npu installation:

scripts/infer_npu.sh \
  --model_id fudan-generative-ai/WAM-Diff2 \
  --input_file /path/to/eval.json \
  --output_file outputs/predictions.json

The main decoding options are:

• --block_size (default: 32)
• --denoising_steps (default: 32)
• --remasking_strategy (default: low_confidence_dynamic)
• --confidence_threshold (default: 0.9)
• --max_new_tokens (default: 128)

Input format

The inference script accepts a JSON list for driving, VQA, and DET requests. The following example uses a single driving image:

[
  {
    "datasource": "Navsim",
    "id": "sample-id",
    "image": ["/path/to/front-camera-image.jpg"],
    "conversations": [
      {
        "from": "human",
        "value": "Here is a front-view image from a driving vehicle: <image>\nThe navigation information is: right\nThe current position is (0.00,0.00)\nThe current velocity is: (4.66,-0.09) and current acceleration is: (0.17,-1.53)\nInstruction: Predict the optimal driving action for the next 4 seconds with 8 new waypoints."
      }
    ]
  }
]

Relative image paths are resolved by the WAM-Diff2 data-loading code; use absolute paths if the working directory is uncertain.

License

WAM-Diff2 is released under the Apache License 2.0. Users must also comply with the licenses and terms of the upstream model, datasets, and other dependencies used in their application.

Citation

@article{zhu2026wam,
  title={WAM-Diff2: Hierarchical AR-to-Diffusion Distillation for Highly Efficient Autonomous Driving VLA},
  author={Zhu, Zhihao and Shang, Hanlin and Xu, Mingwang and Cai, Feipeng and He, Zhuolin and Li, Yaoyi and Han, Jianhua and Xu, Hang and Zhu, Siyu},
  journal={arXiv preprint arXiv:2608.01035},
  year={2026}
}