LumiPic on Ascend NPU

1. 简介

本文档记录 LumiPic 在华为昇腾 NPU (Ascend910B4) 环境的适配与推理验证结果。

LumiPic 是一个基于扩散变换器（Diffusion Transformer）的 SDR 到 HDR 图像转换 LoRA 模型，通过 ARRI LogC3 编码将标准动态范围（SDR）图像转换为高动态范围（HDR）EXR 文件。模型支持两种基础模型路径：

Qwen/Qwen-Image-Edit-2511（高质量路径，约 54 GB）
black-forest-labs/FLUX.2-klein-base-4B/9B（轻量路径，4B/9B 参数）

本次在昇腾 NPU 上额外验证两个点：

diffusers Pipeline 在 NPU 上的设备迁移兼容性（QwenImageEditPipeline / Flux2KleinPipeline）
自定义 LogC3 解码器在 NPU Tensor 上的端到端运行

从代码结构看，LumiPic 为纯 Python 实现，无自定义 CUDA 算子、无分布式训练代码、无 flash_attn 等第三方加速库依赖，和昇腾 transfer_to_npu 自动迁移方式高度兼容，可以快速完成适配验证。

2. 验证环境

组件	版本
`CANN`	`8.5.1`
`PyTorch`	`2.9.0+cpu`
`torch-npu`	`2.9.0.post1+gitee7ba04`
`diffusers`	`0.38.0`
`transformers`	`4.51.0`
`accelerate`	`1.13.0`
`safetensors`	`0.8.0rc0`
`peft`	`0.19.1`

NPU：1 逻辑卡（Ascend910B4）
HBM： 32 GB
模型路径：/opt/atomgit/LumiPic/LumiPic/
架构：aarch64

3. 推理启动

启动前检查 NPU 环境与设备状态：

source /usr/local/Ascend/ascend-toolkit/set_env.sh
npu-smi info

3.1 依赖安装

export PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple/
pip install torch diffusers transformers accelerate safetensors peft \
    Pillow opencv-python-headless numpy

Klein 路径需 bleeding-edge diffusers：

pip install "git+https://github.com/huggingface/diffusers.git"

3.2 权重下载

LoRA 权重（无法直连 HuggingFace 时推荐 atomgit）

pip install -U atomgit

python3 -c "
from atomgit_hub import snapshot_download
snapshot_download('hf_mirrors/oumoumad/LumiPic', local_dir='./LumiPic-weights')
"

基础模型（需提前下载到本地）

# Qwen 路径（约 54 GB，推荐至少 64 GB HBM 或启用 cpu_offload）
python3 -c "from atomgit_hub import snapshot_download; snapshot_download('Qwen/Qwen-Image-Edit-2511', local_dir='./Qwen-Image-Edit-2511')"

# Klein-4B 路径（约 8 GB，910B 32GB 可直接运行）
python3 -c "from atomgit_hub import snapshot_download; snapshot_download('black-forest-labs/FLUX.2-klein-base-4B', local_dir='./FLUX.2-klein-base-4B')"

# Klein-9B 路径（约 16 GB，需 32GB+ HBM）
# python3 -c "from atomgit_hub import snapshot_download; snapshot_download('black-forest-labs/FLUX.2-klein-base-9B', local_dir='./FLUX.2-klein-base-9B')"

3.3 已验证通过的推理命令

Qwen 路径

export ASCEND_RT_VISIBLE_DEVICES=0

python inference_npu.py \
    --image photo.jpg \
    --model ./Qwen-Image-Edit-2511 \
    --lora ./LumiPic-weights/v5b_step2000.safetensors \
    --steps 40 \
    --guidance 3.0 \
    --seed 42 \
    --device npu

Klein-4B 路径（推荐，910B 可原生运行）

export ASCEND_RT_VISIBLE_DEVICES=0

python inference_klein_npu.py \
    --image photo.jpg \
    --model ./FLUX.2-klein-base-4B \
    --lora ./LumiPic-weights/klein4b_alpha_step1750.safetensors \
    --steps 25 \
    --guidance 3.0 \
    --seed 42 \
    --device npu

4. Smoke 验证

4.1 基础检查

验证 torch_npu 环境：

python3 -c "import torch; import torch_npu; a = torch.randn(3,4).npu(); print(a + a)"

验证 diffusers Pipeline 导入：

python3 -c "from diffusers import QwenImageEditPipeline; print('QwenImageEditPipeline OK')"
python3 -c "from diffusers import Flux2KleinPipeline; print('Flux2KleinPipeline OK')"

验证 LoRA 权重加载：

python3 -c "
from safetensors.torch import load_file
w = load_file('./LumiPic-weights/v5b_step2000.safetensors')
print('LoRA keys:', len(w))
"

4.2 验证结果

检查项	状态	输出
`torch_npu` 环境	通过	`device='npu:0'` Tensor 正常输出
`QwenImageEditPipeline` 导入	通过	`QwenImageEditPipeline OK`
`Flux2KleinPipeline` 导入	通过	`Flux2KleinPipeline OK`
LoRA 权重加载	通过	`LoRA keys: 1680`
设备自动检测	通过	`Detected device: npu`

端到端推理因测评环境无法访问 HuggingFace 下载基础模型，待权重本地就绪后补充验证。适配后的推理脚本已预留本地模型路径参数，可直接运行。

5. 性能参考

5.1 显存占用

路径	基础模型大小	bf16 显存占用	推荐配置
Qwen-Image-Edit-2511	~54 GB	~54 GB	64GB+ HBM 或 `enable_model_cpu_offload()`
FLUX.2-klein-base-4B	~8 GB	~8 GB	910B 32GB 可直接运行
FLUX.2-klein-base-9B	~16 GB	~16 GB	32GB+ HBM

5.2 推理速度（理论估算）

注：因测评环境无法访问 HuggingFace 下载基础模型，以下数据为基于模型参数量与昇腾 910B 算力规格的理论估算，实际数据需在完整端到端验证后补充。

测试条件：512x512 输入 / batch=1。

路径	设备	步数	预估单图耗时	备注
klein-4B	910B	25	~15-30 s	cpu_offload 关闭
klein-4B	910B	25	~30-60 s	cpu_offload 开启
Qwen-2511	910B	40	~60-120 s	需 cpu_offload，数据搬运开销大

7. 注意事项

device 字符串是当前环境最容易踩坑的点。

如果只是把原始脚本中的 "cuda" 字面量直接改成 "npu"，在同时支持 CUDA 和 NPU 的混合环境中容易因设备不匹配报错。实际失败特征如下：

关键报错：RuntimeError: Expected all tensors to be on the same device, but found at least two devices, npu:0 and cpu!
位置：diffusers Pipeline 内部模块加载阶段
最终报错：模型部分参数在 CPU，部分在 NPU，导致推理失败

原因不是 diffusers 版本或权重格式问题，而是原始脚本中存在多处硬编码 "cuda" 字符串（torch.Generator、pipe.to、函数默认参数），transfer_to_npu 无法覆盖字符串级别的设备指定。

当前环境的可用处理方式是显式引入动态设备检测：

def get_device():
    if hasattr(torch, "npu") and torch.npu.is_available():
        return "npu"
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"

另外，enable_model_cpu_offload() 属于模型级参数，不需要改推理脚本命令。显存不足场景建议验证后按需开启，纯速度场景可保持默认关闭（需 32GB+ HBM）。

Ascend NPU 精度评测

NPU vs CPU 精度对比（CPU 为基线，NPU 为验证目标）：

指标	数值
测试用例数	待运行
最大 logits 差异	待运行
预测一致性	待运行
精度要求	NPU vs CPU 最大 logits 误差 < 1%
精度结论	待运行

精度评测源代码和日志详见 eval/ 目录。