SD3.5-medium 内存消耗巨大

Question

我有一个 g4dn.xlarge AWS GPU 实例，它具有 16GB 内存 + 48GB 交换空间，以及一个具有 16GB vRAM 的 Tesla T4 GPU 实例。

根据稳定性博客，运行SD3.5 Medium模型应该足够了。

所以我从huggingface下载了模型，并启动了我的测试程序。

首先我看到内存和交换空间增加，总共消耗了约 30GB 内存。然后系统内存开始下降，Nvidia GPU 内存使用量缓慢增长。后来由于内存分配问题而失败。分配~15GB GPU 内存后，无法分配更多内存。

我的问题是，

1. 消耗这么多内存正常吗？系统级别和 GPU 级别。
2. 我的程序出了什么问题？

附上我的源代码，

import os
import json
import torch

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("./stable-diffusion-3.5-medium/")
if torch.cuda.is_available():
    print('use cuda')
    pipe = pipe.to("cuda")
elif torch.mps.is_available():
    print('use mps')
    pipe = pipe.to('mps')
else:
    print('use cpu')

data = []
with open('data.json', 'r') as f:
    data = json.load(f)

os.makedirs('output', exist_ok=True)
for row in data:
    prompt   = row['prompt']
    filename = 'output/%s.png' % (row['uuid'])
    height   = 1280
    width    = 1280
    
    if row['aspect_ratio'] == '16:9':
        width = 720
    elif row['aspect_ratio'] == '9:16':
        width = 720
        height = 1280
    
    print('saving', filename)
    image = pipe(prompt, height=height, width=width).images[0]
    image.save(filename)

Answer 1

我无法发表评论，所以我提交我的想法作为答案。

首先，我不认为你的代码有什么问题。

图像生成模型使用大量内存。根据我个人的经验，当使用SD2.0时，我记得内存增加了大约23GB。内存消耗随着模型变大而增加，但图像大小也会增加。 我建议您尝试减小图像尺寸。尝试从 256/512/768 开始，如果没有出现任何内存分配错误，请逐渐增加它。

希望这有帮助。