[Python] Huge memory consumption with SD3.5-medium

I have a g4dn.xlarge AWS GPU instance, it has 16GB memory + 48GB swap, and a Tesla T4 GPU Instance with 16GB vRAM.

According to the stability blog, it should be sufficient to run SD3.5 Medium model.



So I downloaded the model from hugging face, and launched my test program.

At first I see the memory and swap goes up, consumed in total ~30GB memory. Then the system memory started to go down, and Nvidia GPU memory usage grows slowly. Then later it failed with memory allocation issue. It failed to allocate more memory after ~15GB GPU memory is allocated.

My questions are,

1. Is it normal to consume that much of memory? Both system and GPU level.
2. What's wrong with my program?

Attached my source code,

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)

Continue reading...