FLUX.2-klein-base-4B on Ascend NPU

1. 简介

本文档记录 FLUX.2-klein-base-4B 在华为昇腾 NPU（Ascend 910B）环境的快速部署与验证结果。

FLUX.2-klein-base-4B 是 Black Forest Labs 推出的 FLUX.2 系列轻量级文生图模型，采用 4B 参数 transformer 骨干网络，兼顾生成质量与推理效率。本验证基于 diffusers 框架在昇腾 NPU 上完成端到端推理适配，涵盖模型加载、文本编码、VAE 解码及图像生成全流程。

模型特点：

• 基于 Flow Matching 的扩散训练范式
• 采用 Qwen3 作为文本编码器（Qwen3ForCausalLM）
• 轻量级 transformer 架构（5 层 joint attention + 20 层 single attention）
• 支持 512x512 至 1024x1024 分辨率生成

相关获取地址：

• 权重下载地址（ModelScope）：https://modelscope.cn/models/black-forest-labs/FLUX.2-klein-base-4B
• 权重下载地址（HuggingFace）：https://huggingface.co/black-forest-labs/FLUX.2-klein-base-4B

参考文档：

• https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux
• https://www.hiascend.com/document/detail/zh/canncommercial/800/quickstart/quickstart_0001.html

2. 验证环境

• NPU：1 逻辑卡（Ascend 910B4，32GB HBM）
• 模型路径：/opt/atomgit/models
• 运行环境：Python 3.11

3. 模型加载与推理

3.1 环境准备

pip install diffusers transformers accelerate safetensors
export ASCEND_RT_VISIBLE_DEVICES=0

3.2 推理脚本

已验证通过的加载与推理方式：

import torch
import torch_npu
from diffusers import DiffusionPipeline

# 指定 NPU 设备
device = torch.device("npu:0")

# 加载 pipeline（Diffusers 格式）
pipe = DiffusionPipeline.from_pretrained(
    "/opt/atomgit/models",
    torch_dtype=torch.bfloat16,
    trust_remote_code=False,
)

# 迁移至 NPU
pipe = pipe.to(device)

# 执行文生图推理
image = pipe(
    prompt="A cute cat sitting on a sofa, high quality, photorealistic",
    num_inference_steps=20,
    guidance_scale=3.5,
    height=512,
    width=512,
    max_sequence_length=512,
).images[0]

image.save("flux_output_npu.png")

3.3 关键参数说明

4. Smoke 验证

4.1 模型加载验证

from diffusers import DiffusionPipeline
import torch

pipe = DiffusionPipeline.from_pretrained(
    "/opt/atomgit/models",
    torch_dtype=torch.bfloat16,
)
print(type(pipe).__name__)  # 预期输出：Flux2KleinPipeline

验证结果：

• Pipeline 类正确加载为 Flux2KleinPipeline
• 5 个组件（scheduler、text_encoder、tokenizer、transformer、vae）全部加载成功
• 权重加载耗时约 1.5s

4.2 图像生成验证

使用上述推理脚本运行，验证结果：

• 图像成功生成并保存
• 输出格式：PIL.Image.Image
• 输出尺寸：512 x 512

5. 性能参考

测试条件：单卡 Ascend 910B4，bfloat16 精度，512x512 分辨率。

注：当前为未启用任何图优化或算子融合的基线性能。首次运行因需将权重从 CPU 迁往 NPU，耗时较长；后续同进程连续推理约 4.35s/张。后续可通过 torch.compile（NPU 支持）或 torch_npu 亲和算子替换进一步提升吞吐。

6. 自验证

本仓库提供完整的自验证脚本，可一键复现所有适配验证项。

6.1 运行验证

# 进入仓库目录
cd flux-2-klein-base-4b

# 直接运行（输出到终端）
python3 flux_npu_validate.py

# 重定向保存完整日志
python3 flux_npu_validate.py > flux_validation_log.txt 2>&1

验证完成后会生成：

• flux_validation_log.txt — 完整终端输出日志
• flux_validation_report.json — 结构化 JSON 报告

6.2 验证结果摘要

7. 精度与功能验证

7.1 功能覆盖

7.2 算子兼容性

经扫描 Flux2Transformer2DModel 及 Flux2KleinPipeline 源码，模型核心算子均为 PyTorch 原生算子（nn.Linear、nn.LayerNorm、F.scaled_dot_product_attention 等），无 CUDA/Triton 硬编码内核。在昇腾 NPU 上通过 torch_npu 自动映射，无需额外算子适配。

8. 仓库文件结构

flux-2-klein-base-4b/
├── README.md                        # 本适配文档
├── flux_npu_inference.py            # NPU 推理示例脚本
├── flux_npu_validate.py             # 自验证脚本（adapt-agent 10 步验证流程）
├── flux_validation_log.txt          # 完整验证日志（shell 重定向生成）
└── flux_validation_report.json      # 结构化验证报告

9. 注意事项

1. Pipeline 参数顺序：Flux2KleinPipeline.__call__ 的第一个位置参数为 image，文本 prompt 需显式使用关键字参数 prompt= 传入，否则会被解析为 image 导致 ValueError。

2. 内存优化：若需降低显存占用，可考虑：

   • 使用 pipe.enable_model_cpu_offload()（需确认 NPU 兼容性）
   • 降低分辨率至 256x256 进行快速验证
   • 减少 num_inference_steps 至 10-15 步

3. 文本编码器依赖：本模型使用 Qwen3ForCausalLM 作为文本编码器，需确保 transformers >= 4.57.0，否则可能无法识别该模型类。

4. 自定义 Pipeline 类：Flux2KleinPipeline 与 Flux2Transformer2DModel 属于 diffusers 0.38.0 引入的 FLUX 2 系列新类，旧版本 diffusers 无法直接加载，请务必升级。

5. 性能优化方向：

   • 启用 torch.compile（NPU 后端）加速 transformer 前向
   • 使用 torch_npu 提供的融合算子替换 nn.MultiheadAttention
   • 开启 ACL Graph 捕获（若后续接入 vLLM-Ascend 图像生成服务化框架）