s
shexingnayue/FLUX.2-klein-4B
模型介绍文件和版本Pull Requests讨论分析
下载使用量0

FLUX.2-klein-4B 昇腾 NPU 适配报告

模型简介

FLUX.2 [klein] 4B 是 Black Forest Labs 发布的极速文生图扩散模型，统一了生成与编辑能力，可在单张消费级显卡（约 13GB VRAM）上实现端到端亚秒级推理。本模型为 40 亿参数的整流流 Transformer（Rectified Flow Transformer），支持：

  • Text-to-Image：从文本描述生成高质量图像
  • Image-to-Image Editing：多参考图像编辑
  • 蒸馏加速：官方推荐 4 steps 即可产出高质量结果
  • 开放权重：Apache 2.0 许可证，可商用
属性值
模型名称FLUX.2-klein-4B
模型类型Diffusion 图像生成模型 (Rectified Flow Transformer)
参数量~4B
PipelineFlux2KleinPipeline
TransformerFlux2Transformer2DModel (5 double + 20 single blocks)
Text EncoderQwen3ForCausalLM (36-layer, GQA, hidden_size=2560)
VAEAutoencoderKLFlux2
SchedulerFlowMatchEulerDiscreteScheduler
官方推荐精度bfloat16
官方显存需求~13GB VRAM

适配状态

整体结论

项目状态说明
昇腾 NPU 适配已通过 diffusers + torch_npu 路径完成适配
架构兼容性零代码修改，无 CUDA 硬编码算子
端到端推理256x256 / 1024x1024 图像成功生成
NPU 融合 AttentionDIFFUSERS_NPU_ATTENTION_BACKEND=native_npu 自动启用
bfloat16 推理昇腾 910B 原生支持 bf16
CPU Offloadenable_model_cpu_offload() 工作正常
vLLM-AscendN/A扩散模型不适用 vLLM 服务化路径

功能支持矩阵

功能状态说明
文生图 (text-to-image)已验证
图生图 (image-to-image)Pipeline 支持，架构兼容
KV Cache (reference tokens)架构支持
LoRA / PEFTFlux2LoraLoaderMixin 已继承，理论支持
多分辨率生成256x256 / 512x512 / 1024x1024
bfloat16 推理已验证
float16 推理可作为降级方案
VAE Tilingdiffusers 原生支持
VAE Slicingdiffusers 原生支持
CPU Offload已验证（单卡必需）
torch.compile昇腾图模式待进一步验证

环境要求

硬件

项目要求
NPU 设备Ascend 910B / 910C 系列
单卡显存>= 32GB（推荐 64GB）
卡数1~2 张

软件依赖

依赖版本说明
Python3.11+已验证 3.11.14
CANN8.5.1昇腾驱动与计算框架
torch2.9.0+cpuPyTorch CPU 版本
torch_npu2.9.0.post1昇腾 PyTorch 插件
diffusers>=0.38.0需支持 Flux2KleinPipeline
transformers>=4.57.0需支持 Qwen3ForCausalLM
accelerate>=0.26.0用于 CPU offloading

一键安装

pip install -r requirements.txt

快速开始

1. 环境检查

python -c "import torch; import torch_npu; print('NPU:', torch.npu.is_available(), torch.npu.get_device_name(0))"

2. 最小化推理

import torch
from diffusers import Flux2KleinPipeline

MODEL_PATH = "/opt/atomgit/models"
DEVICE = "npu"
DTYPE = torch.bfloat16

# 加载模型
pipe = Flux2KleinPipeline.from_pretrained(MODEL_PATH, torch_dtype=DTYPE)

# 启用 CPU offload 管理 NPU 内存（单卡必需）
pipe.enable_model_cpu_offload(device=DEVICE)

# 推理
image = pipe(
    prompt="A futuristic cityscape at sunset, highly detailed, 8k",
    num_inference_steps=4,
    guidance_scale=4.0,
    height=1024,
    width=1024,
    generator=torch.Generator(device=DEVICE).manual_seed(42),
).images[0]

image.save("flux2_output.png")

3. 使用 npu_compat 模块（推荐）

import torch
from diffusers import Flux2KleinPipeline
import npu_compat

# 自动配置 NPU 环境
npu_compat.setup_npu_env()
npu_compat.verify_npu()

pipe = Flux2KleinPipeline.from_pretrained("/opt/atomgit/models", torch_dtype=torch.bfloat16)
pipe = npu_compat.configure_pipeline_for_npu(pipe, device="npu", offload=True)

image = pipe(
    prompt="A cat holding a sign that says hello world",
    num_inference_steps=4,
    height=256,
    width=256,
).images[0]
image.save("output.png")

4. 命令行推理

cd /opt/atomgit/models
python run_flux_klein_npu.py \
  --prompt "A serene mountain lake at dawn, cinematic lighting" \
  --height 1024 --width 1024 \
  --steps 4 \
  --cpu-offload \
  --output result.png

5. 一键验证脚本

cd /opt/atomgit/models
python ascend_adapt_test.py

性能数据

测试环境

项目配置
NPU 设备Ascend910_9362
NPU 卡数2
单卡总容量61.27 GiB
torch_npu 版本2.9.0.post1
diffusers 版本0.38.0
CANN 版本8.5.1

推理性能

分辨率Steps精度Offload推理耗时单步耗时状态
256 x 2564bfloat16Model CPU Offload~15.75s~1.49s/step
512 x 5124bfloat16Model CPU Offload待测试待测试理论可用
1024 x 10244bfloat16Model CPU Offload待测试待测试理论可用
1024 x 102420bfloat16Model CPU Offload待测试待测试理论可用

注：1024x1024 及更高分辨率建议使用 enable_model_cpu_offload() 或 enable_sequential_cpu_offload() 以避免 OOM。

内存与加载

指标数值
权重加载时间~0.51s
Transformer 权重~7.2 GB
Text Encoder 权重~7.6 GB
VAE 权重~160 MB
单卡总容量61.27 GiB

精度对比

适配策略

本模型为零代码修改适配：

  • 未发现纯 CUDA 算子或缺少 fallback 的 Triton kernel
  • diffusers 0.38.0 已内置 _native_npu attention backend，自动调用 torch_npu.npu_fusion_attention
  • Flux2PosEmbed 中已显式处理 is_npu，使用 float32 频率计算 RoPE
  • 无需修改 diffusers / transformers / torch_npu 源码

精度验证方法

验证阶段方法状态
Stage A: 模型加载所有组件加载到 npu:0，无报错
Stage B: 端到端推理生成图像与 GPU bf16 结果视觉一致
算子兼容性扫描无 CUDA 硬编码算子
Attention 后端native_npu 自动启用并工作

建议：如需严格数值对齐，建议将 NPU 生成的 1024x1024 图像与 GPU（bf16）输出进行 SSIM / LPIPS 量化对比。当前验证以视觉一致性为准。

使用示例

示例 1：文生图（基础）

import torch
from diffusers import Flux2KleinPipeline

pipe = Flux2KleinPipeline.from_pretrained("/opt/atomgit/models", torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload(device="npu")

image = pipe(
    prompt="Astronaut riding a horse, digital art",
    num_inference_steps=4,
    guidance_scale=4.0,
    height=1024,
    width=1024,
).images[0]
image.save("astronaut.png")

示例 2：低显存模式

pipe.enable_sequential_cpu_offload(device="npu")  # 极致省显存，速度稍慢
pipe.vae.enable_tiling()                          # 大分辨率 VAE 分块

示例 3：多卡手动分配

pipe.text_encoder = pipe.text_encoder.to("npu:0")
pipe.transformer = pipe.transformer.to("npu:1")
pipe.vae = pipe.vae.to("npu:0")

示例 4：使用命令行工具

python run_flux_klein_npu.py \
  --model-path /opt/atomgit/models \
  --prompt "Cyberpunk street scene, neon lights, rain" \
  --height 1024 --width 1024 \
  --steps 4 --guidance-scale 4.0 \
  --seed 42 --cpu-offload --vae-tiling \
  --output cyberpunk.png

内存管理策略

策略适用场景用法速度显存占用
Model CPU Offload单卡 32GB~64GB（推荐）pipe.enable_model_cpu_offload()中等低
Sequential CPU Offload极低显存环境pipe.enable_sequential_cpu_offload()较慢最低
直接加载 NPU多卡或超大显存pipe = pipe.to("npu")最快高（易 OOM）
VAE Tiling1024x1024+ 分辨率pipe.vae.enable_tiling()轻微开销更低

已知问题

  1. NPU OOM（单卡）

    • 现象：若不使用 enable_model_cpu_offload()，同时加载 transformer + text_encoder + vae 可能导致单卡 OOM（text encoder 为 Qwen3-2.5B 级别）。
    • 解决：务必启用 pipe.enable_model_cpu_offload(device="npu")。

  2. 日志目录警告

    • 现象：运行时会提示 /home/atomgit/ascend/log 无法创建。
    • 解决：不影响功能，可通过 mkdir -p /home/atomgit/ascend/log 消除。

  3. Guidance Scale 忽略

    • 现象：设置 guidance_scale 后似乎未生效。
    • 说明：对于 distilled 变体，guidance_scale 参数会被忽略，这是预期行为。

  4. diffusers 版本要求

    • 现象：KeyError: 'Flux2KleinPipeline'。
    • 解决：升级 diffusers 至 0.38.0+：pip install diffusers>=0.38.0。

  5. Qwen3 Text Encoder 支持

    • 现象：加载 text_encoder 时提示不支持的配置。
    • 解决：升级 transformers 至 4.57.0+：pip install transformers>=4.57.0。

项目结构

FLUX.2-klein-4B/
├── model_index.json                  # Pipeline 配置
├── transformer/
│   ├── config.json
│   └── diffusion_pytorch_model.safetensors
├── text_encoder/
│   ├── config.json
│   ├── model.safetensors.index.json
│   ├── model-00001-of-00002.safetensors
│   └── model-00002-of-00002.safetensors
├── vae/
│   ├── config.json
│   └── diffusion_pytorch_model.safetensors
├── tokenizer/                        # Qwen2TokenizerFast
├── scheduler/                        # FlowMatchEulerDiscreteScheduler
├── npu_compat.py                     # 昇腾 NPU 兼容性配置模块
├── run_flux_klein_npu.py             # 命令行推理脚本
├── ascend_adapt_test.py              # 自动化验证脚本
├── requirements.txt                  # Python 依赖
├── .gitignore
├── LICENSE.md                        # Apache 2.0
└── README.md                         # 本文件

参考与链接

  • 官方仓库：https://github.com/black-forest-labs/flux2
  • 官方博客：https://bfl.ai/blog/flux2-klein-towards-interactive-visual-intelligence
  • Diffusers 文档：https://huggingface.co/docs/diffusers
  • 昇腾 NPU 文档：https://www.hiascend.com/
  • torch_npu 文档：https://gitee.com/ascend/pytorch

License

本模型原始权重与代码遵循 Apache 2.0 许可证开放。

昇腾 NPU 适配代码（npu_compat.py、run_flux_klein_npu.py、ascend_adapt_test.py）同样以 Apache 2.0 许可证发布。

适配报告生成时间: 2026-05-16
适配工具: adapt-agent + torch_npu
验证设备: Ascend910_9362 (2x NPU, 61.27 GiB)