Огромное потребление памяти с SD3.5-medium ⇐ Python

[Anonymous]

У меня есть экземпляр графического процессора AWS g4dn.xlarge, он имеет 16 ГБ памяти + 48 ГБ подкачки, а также экземпляр графического процессора Tesla T4 с 16 ГБ видеопамяти.
Согласно блогу о стабильности, так и должно быть будет достаточно для запуска модели SD3.5 Medium.


Итак, я загрузил модель с обнимающим лицом и запустил тестовую программу.Сначала я вижу, что память и подкачка увеличиваются, всего занято ~30 ГБ памяти. Затем системная память начала уменьшаться, а использование памяти графического процессора Nvidia медленно росло. Затем позже произошел сбой из-за проблемы с распределением памяти. Не удалось выделить больше памяти после выделения ~15 ГБ памяти графического процессора.
Мои вопросы:

• Это нормально? потреблять столько памяти? И на уровне системы, и на уровне графического процессора.
• Что не так с моей программой?

Прикрепил исходный код,
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)



Подробнее здесь: https://stackoverflow.com/questions/792 ... 3-5-medium