🚀 从零搭建 FLUX 本地文生图 WebUI：实战笔记与踩坑指南

当开源模型遇上 AMD 显卡，一场长达 35 分钟的加载等待教会我的事

📌 写在前面

最近，AI 绘画领域又迎来了一位重磅选手——FLUX.1。作为 Black Forest Labs 推出的新一代文生图模型，它在图像质量和提示词理解上都有着出色的表现。为了能更方便地使用这个模型，我决定搭建一个本地 WebUI，顺便把整个过程记录下来，既作为学习笔记，也希望能帮助到同样在折腾的小伙伴。

这篇文章会涵盖：

模型下载（从国内镜像站加速）
Gradio WebUI 搭建（带显存管理、模型切换）
实战踩坑（AMD 显卡的辛酸泪）
性能优化（从 35 分钟到 25 秒的蜕变）

🎯 第一步：模型下载（国内用户福音）

FLUX 模型体量不小（约 30GB+），直接从 HuggingFace 下载对国内用户不太友好。这里使用 ModelScope 的镜像站加速下载：

import os
from modelscope import snapshot_download

# 选择模型版本
# 选项 A: 开发版 (质量最高，需申请权限)
repo_id = "black-forest-labs/FLUX.1-dev"
# 选项 B: 快速版 (速度快，开源协议友好)
# repo_id = "black-forest-labs/FLUX.1-schnell"

# 指定本地保存路径
local_dir = f"/root/autodl-tmp/models/{repo_id.split('/')[-1]}"

print(f"正在从镜像站下载 {repo_id} 到 {local_dir} ...")

# 开始下载，支持断点续传
model_path = snapshot_download(
    repo_id=repo_id,
    local_dir=local_dir
)

print(f"下载完成！模型路径: {model_path}")

💡 小贴士：

resume_download=True 默认开启，网络波动也不怕
建议至少预留 60GB 磁盘空间（模型 + 缓存）
下载速度取决于你的带宽，我用了大概 2 小时

🖥️ 第二步：Gradio WebUI 搭建

有了模型，接下来需要一个好用的界面。我基于 Gradio 构建了一个功能完整的 WebUI，支持：

核心功能亮点

多模型管理：可无缝切换 FLUX.1-dev 和 FLUX.1-schnell
显存优化：提供 “full”（全量加载）和 “offload”（CPU 卸载）两种模式
性能监控：实时显示显存占用和生成耗时
一键清理：手动释放显存，避免 OOM
预设示例：快速体验不同风格的生成效果

关键代码解析

模型加载（带预热优化）

def load_model(model_name, memory_mode="full"):
    global pipe, current_model
    
    # 1. 清理旧模型，防止 OOM
    if pipe is not None:
        del pipe
        gc.collect()
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
    
    # 2. 加载新模型
    pipe = FluxPipeline.from_pretrained(
        model_path,
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True
    )
    
    # 3. 根据模式选择加载方式
    if memory_mode == "full":
        pipe = pipe.to("cuda")
        pipe.vae.enable_slicing()  # VAE 分块处理
        pipe.enable_attention_slicing()  # Attention 优化
    else:
        pipe.enable_model_cpu_offload()  # CPU 卸载模式
    
    # 4. 预热：编译优化，提升首次生成速度
    with torch.no_grad():
        _ = pipe(
            "warmup",
            height=256,
            width=256,
            num_inference_steps=1,
            guidance_scale=1.0,
            generator=torch.Generator("cpu").manual_seed(42)
        )
    
    return "✅ 加载成功"

生成函数（含异常处理）

def generate(model_name, prompt, width, height, steps, guidance, seed):
    # 1. 参数校验
    if not prompt or prompt.strip() == "":
        return None, "❌ 请输入提示词"
    
    # 2. 确保模型已加载
    if current_model != model_name:
        status = load_model(model_name)
        if "❌" in status:
            return None, status
    
    # 3. 生成图像
    generator = torch.Generator("cuda").manual_seed(seed)
    
    try:
        with torch.no_grad():
            result = pipe(
                prompt=prompt,
                height=int(height),
                width=int(width),
                num_inference_steps=int(steps),
                guidance_scale=float(guidance),
                generator=generator
            )
            image = result.images[0]
        
        return image, "✅ 生成成功"
        
    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        return None, "❌ 显存不足，请减小尺寸或步数"

界面布局（精简版）

with gr.Blocks(title="Flux 本地生成器 Pro") as demo:
    with gr.Row():
        # 左侧控制面板
        with gr.Column(scale=1):
            model_input = gr.Dropdown(choices=["FLUX.1-dev"], label="模型选择")
            prompt = gr.Textbox(label="提示词", lines=4)
            width = gr.Slider(256, 2048, value=768, label="宽度")
            height = gr.Slider(256, 2048, value=768, label="高度")
            steps = gr.Slider(1, 50, value=25, label="步数")
            btn = gr.Button("开始生成", variant="primary")
            status = gr.Textbox(label="状态")
        
        # 右侧输出
        with gr.Column(scale=2):
            output = gr.Image(label="生成结果", height=600)
    
    # 事件绑定
    btn.click(
        fn=generate,
        inputs=[model_input, prompt, width, height, steps, guidance, seed],
        outputs=[output, status]
    )

🐛 第三步：实战踩坑（AMD 显卡的痛）

第一次运行：惨不忍睹

✅ CUDA 可用: Radeon 8060S Graphics
💾 显存总量: 96.00GB

🔄 正在加载模型...
📂 加载 checkpoint shards: 3/3 [00:05<00:00, 1.68s/it]
📂 加载 pipeline components: 7/7 [01:19<00:00, 11.40s/it]

🔥 预热模型...
100%| 1/1 [00:52<00:00, 52.34s/it]

✅ 加载完成！耗时: 2106.66秒 (约35分钟)

😱 35 分钟！ 这时间都够我吃顿饭了。

问题分析

硬件不匹配：代码使用 PyTorch CUDA 后端，但我的显卡是 AMD Radeon 8060S
统一内存架构：虽然显示 96GB，但这是系统内存+显存共享，带宽受限
编译开销大：预热阶段花了 52 秒，说明 PyTorch 在 AMD 平台上的 JIT 编译效率低
重复加载：页面启动时加载一次，又被清理重新加载，白费功夫

解决方案

✅ 方案一：安装 ROCm 版本 PyTorch（推荐）

# 卸载现有版本
pip uninstall torch torchvision torchaudio

# 安装 ROCm 版本
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.6

✅ 方案二：代码优化（治标不治本）

# 1. 禁用预热（节省52秒）
# print("🔥 预热模型...")
# with torch.no_grad():
#     _ = pipe(...)

# 2. 使用 device_map="auto"
pipe = FluxPipeline.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    device_map="auto"  # 自动分配设备
)

# 3. 启用 torch.compile（如果支持）
if hasattr(torch, 'compile'):
    pipe.transformer = torch.compile(
        pipe.transformer, 
        mode="reduce-overhead"
    )

✅ 方案三：切换模型版本

使用 FLUX.1-schnell 替代 dev 版：

模型更小（约 15GB）
推理步数更少（4-8 步 vs 20-30 步）
生成速度快 3-5 倍

优化结果对比

配置	加载时间	生成时间（512x512, 20步）	显存占用
初始状态（CUDA + AMD）	35 分钟	25 秒	63GB
禁用预热	34 分钟	25 秒	63GB
切换到 ROCm	~2 分钟	~8 秒	45GB
使用 Schnell	~1 分钟	~3 秒	30GB
CPU 模式	~10 分钟	~3 分钟	20GB

💡 第四步：经验总结与最佳实践

1. 硬件选择建议

GPU	推荐模式	预期加载时间	预期生成速度
NVIDIA RTX 4090 (24GB)	Full	30-60秒	5-10秒/图
NVIDIA RTX 3090 (24GB)	Full	45-90秒	8-15秒/图
AMD Radeon 7900XTX	ROCm Full	60-120秒	10-20秒/图
16GB 显卡	Offload	120-180秒	20-40秒/图
8-12GB 显卡	Offload + Schnell	180-300秒	30-60秒/图

2. 优化清单（按优先级）

安装正确的 PyTorch 版本（NVIDIA用CUDA，AMD用ROCm）
启用 VAE slicing 和 tiling（减少显存占用）
使用 bfloat16 精度（省显存且速度略快）
预热模型（虽然第一次慢，但后续生成更快）
根据显存选择模式：≥16GB用full，否则用offload
使用较小的模型变体（Schnell 比 Dev 快 3-5 倍）

3. 日常使用建议

# 推荐配置组合
BEST_CONFIG = {
    "模型": "FLUX.1-schnell",  # 速度优先
    "显存模式": "full" if torch.cuda.is_available() and memory >= 16 else "offload",
    "分辨率": 768,  # 平衡质量与速度
    "步数": 8 if use_schnell else 25,
    "CFG": 0.0 if use_schnell else 3.5,
}