FLUX.2-klein-4B Background Remove LoRA - NPU

概述

本仓库提供 FLUX.2-klein-4B Background Remove LoRA 模型在华为昇腾 Ascend910 NPU 上的适配与推理方案。该 LoRA 用于从图像中移除背景并分离主体，基于 black-forest-labs/FLUX.2-klein-4B 扩散模型。

原始模型: fal/flux-2-klein-4B-background-remove-lora
基础模型: black-forest-labs/FLUX.2-klein-4B
任务类型: 图像编辑 / 背景移除
模型框架: PyTorch + Diffusers
输入格式: RGB 图像（PIL Image）
输出格式: RGB 图像（背景被移除）
LoRA 权重: ~76MB（safetensors 格式）
基础模型: ~23GB
许可证: Apache 2.0

环境要求

组件	版本
Python	3.11+
PyTorch	>= 2.0.0
torch_npu	>= 2.9.0
Diffusers	>= 0.38.0
Transformers	>= 4.57.0
Accelerate	>= 1.13.0
PEFT	>= 0.19.0
Ascend NPU	Ascend910
磁盘	>= 50GB

NPU 适配说明

FLUX.2-klein-4B 是 FLUX.2 系列的蒸馏版本（4B 参数），使用 Qwen3ForCausalLM 作为文本编码器、Flux2Transformer2DModel 作为扩散主干，通过 Flux2KleinPipeline 在 Diffusers 框架中加载。

适配要点：

使用 torch.bfloat16 精度适配 Ascend910 NPU
通过 pipe.load_lora_weights() 加载 LoRA 权重
NPU 推理无需修改模型代码
使用 CPU 生成器确保 CPU 与 NPU 初始噪声一致

快速开始

# 安装依赖
pip install torch torch_npu diffusers transformers accelerate peft safetensors numpy pillow modelscope

# NPU 推理
python3 inference.py --device npu:0 --steps 20 --seed 42 --height 512 --width 512

# CPU/NPU 精度对比
python3 compare_cpu_npu.py --steps 4 --seed 42 --height 128 --width 128

推理结果

NPU 推理性能（512x512，20步）

指标	值
推理耗时	3.90 秒
设备	Ascend910
数据类型	bfloat16

CPU vs NPU 性能对比

指标	CPU	NPU	加速比
推理耗时（4步，128x128）	481.78 秒	0.78 秒	616x

CPU/NPU 精度测试

测试条件：128x128 输入，4步推理，seed=42，相同初始潜变量

指标	值
MAE（0-255）	1.0425
最大绝对误差	25.0000
相对误差	0.4088%
RMSE	1.8729
PSNR	42.68 dB
SSIM	0.999349
像素一致率（<1%）	92.07%

结论：NPU 与 CPU 推理结果误差为 0.4088%，小于 1%。

核心推理代码

import torch
import torch_npu
from diffusers import Flux2KleinPipeline
from PIL import Image

pipe = Flux2KleinPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-4B",
    torch_dtype=torch.bfloat16,
)
pipe.load_lora_weights("flux-background-remove-lora.safetensors")
pipe = pipe.to("npu:0")

input_img = Image.open("input.jpg").convert("RGB").resize((512, 512))
generator = torch.Generator(device="cpu").manual_seed(42)

with torch.no_grad():
    output = pipe(
        image=input_img,
        prompt="Remove the background from the image",
        height=512, width=512,
        num_inference_steps=20,
        guidance_scale=3.5,
        generator=generator,
        output_type="pil",
    )
output.images[0].save("output.png")

文件结构

inference.py          # NPU/CPU 推理脚本
compare_cpu_npu.py    # CPU/NPU 精度对比脚本
requirements.txt      # Python 依赖
README.md             # 本文件
compare_results.json  # 精度测试结果

概述

原始模型: fal/flux-2-klein-4B-background-remove-lora

基础模型: black-forest-labs/FLUX.2-klein-4B

任务类型: 图像编辑 / 背景移除

模型框架: PyTorch + Diffusers

输入格式: RGB 图像（PIL Image）

输出格式: RGB 图像（背景被移除）

LoRA 权重: ~76MB（safetensors 格式）

基础模型: ~23GB

许可证: Apache 2.0

组件

版本

Python

3.11+

PyTorch

>= 2.0.0

torch_npu

>= 2.9.0

Diffusers

>= 0.38.0

Transformers

>= 4.57.0

Accelerate

>= 1.13.0

PEFT

>= 0.19.0

Ascend NPU

Ascend910

磁盘

>= 50GB

NPU 适配说明

适配要点：

使用 torch.bfloat16 精度适配 Ascend910 NPU

通过 pipe.load_lora_weights() 加载 LoRA 权重

NPU 推理无需修改模型代码

使用 CPU 生成器确保 CPU 与 NPU 初始噪声一致

快速开始

# 安装依赖
pip install torch torch_npu diffusers transformers accelerate peft safetensors numpy pillow modelscope

# NPU 推理
python3 inference.py --device npu:0 --steps 20 --seed 42 --height 512 --width 512

# CPU/NPU 精度对比
python3 compare_cpu_npu.py --steps 4 --seed 42 --height 128 --width 128

指标

值

推理耗时

3.90 秒

设备

Ascend910

数据类型

bfloat16

指标

CPU

NPU

加速比

推理耗时（4步，128x128）

481.78 秒

0.78 秒

616x

指标

值

MAE（0-255）

1.0425

最大绝对误差

25.0000

相对误差

0.4088%

RMSE

1.8729

PSNR

42.68 dB

SSIM

0.999349

像素一致率（<1%）

92.07%

核心推理代码

import torch
import torch_npu
from diffusers import Flux2KleinPipeline
from PIL import Image

pipe = Flux2KleinPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-4B",
    torch_dtype=torch.bfloat16,
)
pipe.load_lora_weights("flux-background-remove-lora.safetensors")
pipe = pipe.to("npu:0")

input_img = Image.open("input.jpg").convert("RGB").resize((512, 512))
generator = torch.Generator(device="cpu").manual_seed(42)

with torch.no_grad():
    output = pipe(
        image=input_img,
        prompt="Remove the background from the image",
        height=512, width=512,
        num_inference_steps=20,
        guidance_scale=3.5,
        generator=generator,
        output_type="pil",
    )
output.images[0].save("output.png")

inference.py # NPU/CPU 推理脚本 compare_cpu_npu.py # CPU/NPU 精度对比脚本 requirements.txt # Python 依赖 README.md # 本文件 compare_results.json # 精度测试结果