SD 3.5 Medium na M1: optimalizace do 11.6 GB

Optimalní využití paměti pro Stable Diffusion 3.5 Medium na Apple Silicon M1

Model Stable Diffusion 3.5 Medium s 2,5 miliardami parametrů vyžaduje až 21 GB unified memory na Apple Silicon kvůli T5-XXL enkodéru (10,75 GB ve formátu fp16). Optimalizace snižuje vrcholovou spotřebu na 11,6 GB, což umožňuje spuštění i na M1 s 16 GB paměti. Přehled problému a postupné načítání komponent.

Příčiny vrcholové spotřeby:

• Duplikace váh: načtení do CPU a následné kopírování do MPS.
• enable_model_cpu_offload() není efektivní – CPU a MPS sdílí unified memory.
• Hooky z knihovny accelerate s device_map udržují silné reference, blokují garbage collector.

Standardní kód:

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-medium",
    torch_dtype=torch.float16, variant="fp16")
pipe.to("mps")

Vrchol: 21,4 GB bez ohledu na počet kroků generace.

Řešení: postupné načítání

Rozdělení na fáze: kódování promptu (T5), vyčištění, generace (DiT + VAE).

1. Načtení s device_map="balanced"

Přímé načtení do MPS bez kopírování z CPU:

pipe = StableDiffusion3Pipeline.from_pretrained(
    path, torch_dtype=torch.float16, device_map="balanced")

Paměťový vrchol T5: ~11 GB.

2. Správné odstranění hooků

from accelerate.hooks import remove_hook_from_submodules

for attr in list(vars(pipe)):
    comp = getattr(pipe, attr, None)
    if isinstance(comp, torch.nn.Module):
        remove_hook_from_submodules(comp)
        setattr(pipe, attr, None)

del pipe
gc.collect()
torch.mps.synchronize()
torch.mps.empty_cache()

Bez tohoto kroku zůstane T5 v paměti (10,75 GB).

3. Kompletní postup

Fáze 1: Enkodér

enc_pipe = StableDiffusion3Pipeline.from_pretrained(path,
    transformer=None, vae=None,
    text_encoder=None, text_encoder_2=None,
    tokenizer=None, tokenizer_2=None,
    variant="fp16", torch_dtype=torch.float16,
    device_map="balanced")

prompt_embeds, neg_embeds, pooled, neg_pooled = enc_pipe.encode_prompt(
    prompt="a warrior with a sword",
    prompt_2="a warrior with a sword",
    prompt_3="a warrior with a sword",
    device="mps", num_images_per_prompt=1)

Vyčištění podle kroku 2.

Fáze 2: Generátor

gen_pipe = StableDiffusion3Pipeline.from_pretrained(path,
    text_encoder=None, text_encoder_2=None, text_encoder_3=None,
    tokenizer=None, tokenizer_2=None, tokenizer_3=None,
    variant="fp16", torch_dtype=torch.float16)
gen_pipe.to("mps")

image = gen_pipe(
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=neg_embeds,
    pooled_prompt_embeds=pooled,
    negative_pooled_prompt_embeds=neg_pooled,
    num_inference_steps=40).images[0]

.to("mps") reutilizuje cache T5, snižuje footprint DiT z 11,5 na 6 GB.

Srovnání spotřeby paměti

| Etapa | Před optimalizací | Po optimalizaci |

|-------------------|-------------------|------------------|

| Načtení T5 | 15,3 GB | 11,5 GB |

| T5 po vyčištění | 10,75 GB | 0,01 GB |

| DiT + difúze | 8,9 GB | 8,9 GB |

| Vrchol | 21,4 GB | 11,6 GB |

Výkon na hardwaru

• M1 16 GB: 4 min 54 s (T5: 1 min 6 s, DiT: 3 min 34 s), 11–13 GB.
• M3 Pro 36 GB: 1 min 47 s (T5: 14 s, DiT: 1 min 27 s), 9–11 GB.

Rozlišení 512x512, 40 kroků.

Důležité:

• T5-XXL je hlavním spotřebitelem paměti (71 %).
• device_map="balanced" + odstranění hooků je klíčové pro úspěšné vyčištění.
• Postupné načítání snižuje vrcholovou spotřebu o 45 %.
• Funguje na M1 16 GB bez swapování.
• .to("mps") pro DiT řeší reutilizaci cache.

— Editorial Team