README.md

---
license: apache-2.0
library_name: transformers
pipeline_tag: image-text-to-text
tags:
  - autonomous-driving
  - vision-language-action
  - multimodal
  - multi-image
  - visual-question-answering
  - object-detection
  - block-diffusion
  - qwen3-vl
  - pytorch
---

# 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](https://github.com/fudan-generative-vision/WAM-Diff2)
- [Project page](https://fudan-generative-vision.github.io/WAM-Diff2/)
- [Paper](https://arxiv.org/abs/2608.01035)

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

| Property | Value |
| --- | --- |
| Architecture | WAM-Diff2, based on the Qwen3-VL architecture |
| Checkpoint | Block-32, 2B |
| Parameters | 2,127,532,032 |
| Weight dtype | BF16 |
| Weight format | Safetensors |
| Primary modality | One or more images and a text prompt; telemetry is optional and task-dependent |
| Output | Task-dependent token sequences, including answers, detections, and driving waypoints |
| Default decoding | Dynamic decoding with Block-32 and up to 32 denoising steps |

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.

```bash
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:

```bash
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`:

```bash
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:

```bash
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:

```json
[
  {
    "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

```bibtex
@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}
}
```