Flux.1 Schnell生成器单次生成后GPU显存耗尽及重复下载问题求助

一、解决二次生成显存不足问题 #

第一次生成正常、第二次爆显存，大概率是显存泄漏或重复加载导致的，按以下步骤处理：

• 每次生成后强制清缓存：生成完成后立刻调用torch.cuda.empty_cache()，再搭配gc.collect()触发Python垃圾回收，彻底清理无用张量。
• 模型只加载一次：把模型和LoRA的加载代码放在全局作用域，第一次加载完成后直接复用已初始化的pipe对象，不要每次生成都重新加载。
• 开启梯度检查点：加载FLUX模型时添加gradient_checkpointing=True参数，能大幅降低显存占用，代价是生成速度略有下降。
• 降低分辨率或单张循环：如果批量生成压力大，先将分辨率从1024降至768，或者采用单张生成循环3次的方式，比一次性批量生成更省显存。

二、避免每次启动重新下载模型 #

你已经保存了model.safetensor和训练文件，直接本地加载即可：

• 指定本地模型路径：使用FluxPipeline.from_pretrained()时，将路径改为存放模型文件的本地文件夹，同时添加local_files_only=True参数，示例：pipe = FluxPipeline.from_pretrained("./my_flux_model", local_files_only=True)，这样就不会触发远程下载。
• 本地存储LoRA权重：训练完成的LoRA用lora_model.save_pretrained("./my_lora")保存到本地，后续加载直接调用pipe.load_lora_weights("./my_lora")，无需从训练输出临时加载。
• 使用持久化存储：在Lightning AI Studio中，将模型文件放在/workspace目录下，不要放在临时目录，确保每次启动都能直接读取。

三、实现批量生成3张图 #

结合显存优化，按以下方式操作：

• 复用已加载的模型：确保模型只初始化一次，之后多次调用生成接口。
• 单张循环生成：如果直接设置num_images_per_prompt=3仍爆显存，就循环执行3次单张生成，每次生成后清理缓存。
• 开启CPU offload：调用pipe.enable_model_cpu_offload()，将模型非核心层放在CPU，需要时再移至GPU，能显著降低显存占用，仅生成速度会稍有减慢。
• 示例代码：

import torch
import gc
from diffusers import FluxPipeline

# 仅加载一次模型和LoRA
pipe = FluxPipeline.from_pretrained("./my_flux_model", local_files_only=True, torch_dtype=torch.bfloat16)
pipe.load_lora_weights("./my_lora")
pipe.enable_model_cpu_offload()

# 批量生成3张图片
def generate_batch(prompt, count=3):
    imgs = []
    for i in range(count):
        img = pipe(prompt, guidance_scale=3.5, height=768, width=768).images[0]
        imgs.append(img)
        torch.cuda.empty_cache()
        gc.collect()
    return imgs

# 调用生成并保存结果
generated_imgs = generate_batch("你的提示词内容", 3)
for idx, img in enumerate(generated_imgs):
    img.save(f"output_{idx}.png")

内容的提问来源于stack exchange，提问作者ACHINTYA GUPTA