PrITTI（CVPR 2026）是基于图元（primitive）的可控可编辑 3D 语义城市场景生成框架，在 KITTI-360 上训练，用紧凑、语义明确的 3D 图元替代体素表示，实现低显存、快速推理的场景生成。
原始模型：raniatze/pritti-checkpoints（CC BY-NC 4.0） 官方代码：autonomousvision/pritti 论文：Primitive-based Generation of Controllable and Editable 3D Semantic Urban Scenes
Checkpoint 组成与选型说明：发布包含三部分——lvae.ckpt（Layout VAE，300 epoch）、ldm_b/（DiT-B 潜空间扩散模型，diffusers pipeline 格式：transformer + decoder + scheduler）、以及 LVAE 解码器包装。其中 LVAE 与 decoder 包装类（pritti.autoencoder.modeling...LVAEDecoderWrapper）依赖官方 GitHub 仓库的自定义包（本容器不可达）；ldm_b/transformer 是标准 diffusers DiTTransformer2DModel（12 层、768 维、12 heads、in/out 64 通道、patch 2、32×32 潜空间、ada_norm_zero 条件、3 类嵌入），自包含可离线加载，故选其作为昇腾适配的最小可行路径。
Demo 路径：固定种子合成潜变量 (1, 64, 32, 32) + t=500 + 类别条件 → DiT epsilon 预测 → 一步 DDPM 重建 x0（用 ldm_b/scheduler 的 DDPMScheduler alpha 累积值与 clip 范围）。CPU 与 NPU 输入完全一致。
# 1. 克隆仓库
git clone https://gitcode.com/weasonlee/pritti-checkpoints.git
cd pritti-checkpoints
# 2. 安装依赖
pip install -r requirements.txt
# 3. 下载权重到 models/（已 .gitignore 排除；lvae.ckpt 1.1GB + ldm_b 773MB）
mkdir -p models
HF_ENDPOINT=https://hf-mirror.com python -c "
from huggingface_hub import snapshot_download
snapshot_download('raniatze/pritti-checkpoints', local_dir='models')"
# 4. 推理（NPU）
python inference.py --device npu
# 或 CPU
python inference.py --device cpu
# 换时间步/类别条件
python inference.py --device npu --timestep 200 --class-label 0$ python inference.py --device npu
torch.npu.is_available() = True
torch.npu.device_count() = 2
[info] device = npu:0, torch = 2.9.0+cpu
[load] DiTTransformer2DModel: 138.6M params, timesteps=1000, prediction=epsilon
[info] latent = (1, 64, 32, 32), t = 500, class_label = 1
[time] run 0: 20 ms
[time] avg: 19 ms (latent 64x32x32, 12-layer DiT-B)
[out ] eps_pred: mean=0.0259 std=1.0785 min=-5.334 max=5.030
[out ] x0(1-step): mean=-0.0509 std=0.7113
===== CPU vs NPU comparison (t=500, cls=1) =====
eps_pred cos_sim=0.99999999 max_abs_err=1.077e-03 mean_abs_err=1.172e-04
x0_1step cos_sim=0.99999984 max_abs_err=3.508e-03 mean_abs_err=2.571e-04| 输出 | cos_sim | max_abs_err | mean_abs_err |
|---|---|---|---|
| eps_pred（DiT 噪声预测） | 0.99999999 | 1.077e-03 | 1.172e-04 |
| x0_1step（一步 DDPM 重建潜变量） | 0.99999984 | 3.508e-03 | 2.571e-04 |
| 单次前向耗时（含同步，3 次均值） | NPU ~19 ms vs CPU ~876 ms（~46× 加速） | — | — |
结论：NPU（Ascend 910, torch-npu 2.9.0.post1）与 CPU 的 DiT 噪声预测一致性达 1e-3 量级（fp32 正常差异水平），一步重建潜变量 cos_sim > 0.9999998，适配成功。
├── inference.py # 推理脚本（DiT-B epsilon 预测 + 一步 DDPM + NPU/CPU 对比）
├── readme.md
├── requirements.txt
├── models/ # 权重目录（需自行下载，已 gitignore）
│ ├── lvae.ckpt
│ └── ldm_b/（transformer + decoder + scheduler + model_index）
├── assets/ # 运行截图
│ ├── agent_workflow.png
│ ├── npu_device_call.png
│ └── model_result.png
└── results/ # 每次运行的输出/耗时产物（gitignore）权重 CC BY-NC 4.0（仅限学术与非商业研究使用）。
Rumpler et al. "PrITTI: Primitive-based Generation of Controllable and Editable 3D Semantic Urban Scenes." CVPR 2026. arXiv:2506.19117.