FLUX.2-klein-4B Spritesheet LoRA - NPU

概述

本仓库提供 FLUX.2-klein-4B Spritesheet LoRA 模型在华为昇腾 Ascend910 NPU 上的适配与推理方案。该 LoRA 用于将单个对象图像转换为 2x2 精灵图（Sprite Sheet），包含 4 个不同视角的视图。

原始模型: fal/flux-2-klein-4b-spritesheet-lora
基础模型: black-forest-labs/FLUX.2-klein-4B
任务类型: 图像编辑 / 精灵图生成（Sprite Sheet）
模型框架: PyTorch + Diffusers
输入格式: RGB 图像（单个对象）
输出格式: RGB 图像（2x2 网格精灵图）
LoRA 权重: ~76MB（safetensors 格式）
基础模型: ~23GB
许可证: Apache 2.0

环境要求

组件	版本
Python	3.11+
PyTorch	>= 2.0.0
torch_npu	>= 2.9.0
Diffusers	>= 0.38.0
Ascend NPU	Ascend910

NPU 适配说明

使用 Flux2KleinPipeline + load_lora_weights() 加载。输入为单个对象的图像，输出为等距、正视图、侧视图、俯视图等 4 种视角组成的精灵图。

快速开始

python3 inference.py --device npu:0 --steps 20 --seed 42 --height 512 --width 512
python3 compare_cpu_npu.py --steps 4 --seed 42 --height 128 --width 128

推理结果

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

指标	值
推理耗时	3.97 秒
设备	Ascend910

CPU vs NPU 性能对比

指标	CPU	NPU	加速比
推理耗时	572.00 秒	0.86 秒	667x

CPU/NPU 精度测试

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

指标	值
MAE（0-255）	2.3232
最大绝对误差	168.0000
相对误差	0.9111%
RMSE	7.4911
PSNR	30.64 dB
SSIM	0.995596
像素一致率（<1%）	76.48%

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

核心推理代码

import torch, 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-spritesheet-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="2x2 sprite sheet",
        height=512, width=512,
        num_inference_steps=20,
        guidance_scale=3.5,
        generator=generator,
    )
output.images[0].save("output.png")

概述

原始模型: fal/flux-2-klein-4b-spritesheet-lora

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

任务类型: 图像编辑 / 精灵图生成（Sprite Sheet）

模型框架: PyTorch + Diffusers

输入格式: RGB 图像（单个对象）

输出格式: RGB 图像（2x2 网格精灵图）

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

基础模型: ~23GB

许可证: Apache 2.0

组件

版本

Python

3.11+

PyTorch

>= 2.0.0

torch_npu

>= 2.9.0

Diffusers

>= 0.38.0

Ascend NPU

Ascend910

指标

值

推理耗时

3.97 秒

设备

Ascend910

指标

CPU

NPU

加速比

推理耗时

572.00 秒

0.86 秒

667x

指标

值

MAE（0-255）

2.3232

最大绝对误差

168.0000

相对误差

0.9111%

RMSE

7.4911

PSNR

30.64 dB

SSIM

0.995596

像素一致率（<1%）

76.48%

核心推理代码

import torch, 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-spritesheet-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="2x2 sprite sheet",
        height=512, width=512,
        num_inference_steps=20,
        guidance_scale=3.5,
        generator=generator,
    )
output.images[0].save("output.png")