SD 3.5 Medium en M1: optimización a 11.6 GB

Optimización de memoria para Stable Diffusion 3.5 Medium en Apple Silicon M1

El modelo Stable Diffusion 3.5 Medium con 2.5B de parámetros puede consumir hasta 21 GB de memoria unificada en Apple Silicon debido al codificador T5-XXL (10.75 GB en fp16). La carga optimizada reduce el pico de uso a solo 11.6 GB, permitiendo una operación fluida en Macs M1 con 16 GB de RAM. Esta guía desglosa el problema y presenta una estrategia paso a paso para la carga.

Causas principales de los picos de memoria:

• Duplicación de pesos: el modelo se carga en CPU y luego se copia a MPS.
• enable_model_cpu_offload() es ineficaz—CPU y MPS compiten por la memoria unificada.
• Los ganchos de Accelerate con device_map mantienen referencias fuertes, impidiendo la recolección de basura.

Código estándar:

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

Máximo de memoria: 21.4 GB sin importar el número de pasos de inferencia.

Solución: Carga secuencial

Divide el proceso en fases: codificación del prompt (T5), descarga, y generación (DiT + VAE).

1. Carga con device_map="balanced"

Carga directamente en MPS sin etapa previa en CPU:

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

Pico de T5: ~11 GB.

2. Eliminación correcta de ganchos

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()

Sin esto, T5 permanece en memoria (10.75 GB).

3. Flujo completo secuencial

Fase 1: Codificador

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="un guerrero con una espada",
    prompt_2="un guerrero con una espada",
    prompt_3="un guerrero con una espada",
    device="mps", num_images_per_prompt=1)

Descarga tras esta fase.

Fase 2: Generador

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") reutiliza la caché de T5, reduciendo el consumo de DiT de 11.5 GB a solo 6 GB.

Comparación de uso de memoria

| Etapa | Antes de optimizar | Después de optimizar |

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

| Carga de T5 | 15.3 GB | 11.5 GB |

| T5 tras descarga | 10.75 GB | 0.01 GB |

| DiT + Difusión | 8.9 GB | 8.9 GB |

| Pico total | 21.4 GB | 11.6 GB |

Rendimiento en hardware

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

Resolución: 512x512, 40 pasos de inferencia.

Conclusiones clave:

• T5-XXL es el principal consumidor de memoria (~71% del total).
• device_map="balanced" combinado con eliminación de ganchos es esencial para liberar memoria.
• La carga secuencial reduce el pico de memoria en un 45%.
• Funciona de forma confiable en Macs M1 de 16 GB sin intercambio.
• .to("mps") en DiT permite una reutilización eficiente de la caché.

— Editorial Team