【diffusers极速入门（十）】Flux-pipe 推理，完美利用任何显存大小，GPU显存节省终极方案（附代码）

多恩Stone

已于 2025-02-20 14:34:28 修改

阅读量715

点赞数 6

分类专栏： AIGC Diffusion 编程学习文章标签： python AIGC stable diffusion 人工智能人工智能作画

于 2025-02-20 14:30:02 首次发布

本文链接：https://blog.csdn.net/weixin_44212848/article/details/145752050

版权

编程学习同时被 3 个专栏收录

108 篇文章

订阅专栏

AIGC

92 篇文章

订阅专栏

Diffusion

55 篇文章

订阅专栏

系列文章目录

文章目录

1-完美利用方案

该方案可以根据所使用的 GPU 有多少显存空闲，即可设置对应的max_memory，最大化利用 GPU 的同时，保证不出现 OOM 问题。
可以减少对 GPU 的显存需求，推理时间也不会过于慢。

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"  # 减少显存碎片化

torch.inference_mode()  # 启用推理模式，减少显存占用

local_path = '/path/to/models/FLUX.1-Canny-dev'
# online_path: "black-forest-labs/FLUX.1-Canny-dev"

max_memory = {
    0: "39GiB",  # 限制 GPU 0 使用 39GB 显存 
}

torch_dtype = torch.bfloat16
pipe = FluxControlPipeline.from_pretrained(local_path, 
                                           torch_dtype=torch_dtype,
                                           device_map="balanced",
                                           max_memory=max_memory)  # 自动决定模型组件的设备分配    
processor = CannyDetector()

# pipe.to("cuda")
pipe.enable_vae_slicing()
pipe.enable_vae_tiling()  # 新增分块策略
pipe.enable_attention_slicing("max")  # 新增分块策略

基本概念

在 diffusers 库中，device_map 是一个非常有用的参数，它主要用于控制模型在不同计算设备（如 CPU、GPU 等）之间的分布，这对于处理大型模型时的内存管理和性能优化至关重要。
device_map 本质上是一个字典或字符串，用于指定模型的各个部分（如不同的层）应该被放置到哪个计算设备上。通过合理配置 device_map，可以将模型分布到多个 GPU 上，或者在 GPU 和 CPU 之间进行混合部署，以充分利用不同设备的计算资源并避免内存溢出。

使用场景

多 GPU 系统：当你有多个 GPU 可用时，可以使用 device_map 将模型均匀地分布到这些 GPU 上，从而加速推理过程。
内存受限：如果单个 GPU 的内存不足以容纳整个模型，你可以将部分模型放置到 CPU 上，或者在多个 GPU 之间拆分模型，以避免内存溢出错误。

其他场景示例

device_map 可以接受以下字符串或字典两种类型的值。

使用字符串作为 `device_map`

"auto"：自动将模型分布到可用的 GPU 上，如果没有足够的 GPU 内存，会将部分模型放到 CPU 上。
"balanced"：尝试在所有可用的 GPU 之间均匀地分布模型，以平衡内存使用。
"balanced_low_0"：与 "balanced" 类似，但会尽量减少第一个 GPU 的内存使用。
"sequential"：按顺序将模型层依次分配到可用的 GPU 上。
"cpu"：将整个模型加载到 CPU 上。
"cuda"：将整个模型加载到第一个可用的 GPU 上。

from diffusers import StableDiffusionPipeline

# 使用 "auto" 自动分配模型到可用设备
pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", device_map="auto")

# 生成图像
image = pipeline("A beautiful landscape").images[0]
image.save("landscape.png")

使用字典作为 `device_map`

字典的键是模型的模块名称，值是设备名称（如 "cpu"、"cuda:0"、"cuda:1" 等）。通过这种方式，你可以精确地指定每个模块应该放置的设备。

from diffusers import StableDiffusionPipeline

# 自定义设备映射
device_map = {
    "text_encoder": "cuda:0",
    "unet": "cuda:1",
    "vae": "cuda:0"
}

pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", device_map=device_map)

# 生成图像
image = pipeline("A cute cat").images[0]
image.save("cat.png")

注意事项

设备可用性：确保你指定的设备（如 cuda:1）是实际可用的，否则会抛出错误。
内存管理：合理配置 device_map 可以避免内存溢出，但不同的分配策略可能会影响推理性能，需要根据实际情况进行调整。
兼容性：某些模型可能不支持所有的 device_map 选项，使用时需要参考模型的文档。

2-超节省方案

几乎不需要 GPU 的显存占有 (好像是纯 CPU 跑?)，但推理时间会变慢很多。

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"  # 减少显存碎片化

torch.inference_mode()  # 启用推理模式，减少显存占用

local_path = '/path/to/models/FLUX.1-Canny-dev'
# online_path: "black-forest-labs/FLUX.1-Canny-dev"

torch_dtype = torch.bfloat16
pipe = FluxControlPipeline.from_pretrained(local_path, 
                                           torch_dtype=torch_dtype,
                                           )  
pipe.enable_sequential_cpu_offload() # 推理几乎 0 显存占用