Train a LoRA on a Mac

Why you train on your own Mac

The training set is the most sensitive part of this job. If you put the model on a phone, keep those images on the machine that already has permission to see them. MPS or CPU is slower than a rented GPU. It is also a machine you control. There is no CUDA step in this course. If a script assumes CUDA, rewrite it.

Training produces a LoRA checkpoint. That file is not the on-device model. You still have to merge and export. Do not skip ahead and try to load the adapter in Core AI.

The steps, one box at a time

Load the dataset you already wrote down in your contract.

The frozen base and a tiny trainable patch share one forward pass (one trip through the model).

The optimizer updates only the adapter.

Save the adapter, a preview sheet, and a manifest. Then compare holdout prompts to the untouched base. Loop if the concept is missing or the background is memorized.

Environment setup

Use Python 3.11 or newer on macOS. Create the project before the first training command.

curl -LsSf https://astral.sh/uv/install.sh | sh
mkdir -p diffusion-lora-coreai && cd diffusion-lora-coreai
uv init --python 3.11
uv venv
uv add torch torchvision diffusers peft transformers accelerate datasets safetensors coreai-torch
uv run python -c "import torch; print(torch.__version__); print('MPS:', torch.backends.mps.is_available())"

Use uv run for every Python command. Pin the base revision from your contract. Never add a CUDA extra.

Worked path

Save this as train_lora.py. It is a PEFT-on-UNet sketch for sd-1.5. Adjust rank, steps, and resolution to fit memory. On 16 GB machines start at 512, batch 1, rank 8. If MPS throws an unsupported op, set the device to CPU or move that op.

from pathlib import Path
import json
import torch
from diffusers import StableDiffusionPipeline
from peft import LoraConfig, get_peft_model
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

DEVICE = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
print("training on", DEVICE)

BASE = "runwayml/stable-diffusion-v1-5"
REVISION = "pin-this"
TRAIN_DIR = Path("data/train")
OUT = Path("artifacts/style-lora")
OUT.mkdir(parents=True, exist_ok=True)

class CaptionFolder(Dataset):
    def __init__(self, root):
        self.paths = sorted(p for p in Path(root).iterdir() if p.suffix.lower() in {".jpg", ".jpeg", ".png"})
        self.tf = transforms.Compose([
            transforms.Resize(512, interpolation=transforms.InterpolationMode.BILINEAR),
            transforms.CenterCrop(512),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ])

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, i):
        path = self.paths[i]
        image = self.tf(Image.open(path).convert("RGB"))
        caption = path.with_suffix(".txt").read_text(encoding="utf-8").strip()
        return {"pixel_values": image, "caption": caption}

pipe = StableDiffusionPipeline.from_pretrained(BASE, revision=REVISION, torch_dtype=torch.float32)
pipe.to(DEVICE)
pipe.vae.requires_grad_(False)
pipe.text_encoder.requires_grad_(False)
pipe.unet.requires_grad_(False)

lora = LoraConfig(r=8, lora_alpha=16, target_modules=["to_q", "to_k", "to_v", "to_out.0"])
unet = get_peft_model(pipe.unet, lora)
unet.print_trainable_parameters()

loader = DataLoader(CaptionFolder(TRAIN_DIR), batch_size=1, shuffle=True)
optimizer = torch.optim.AdamW((p for p in unet.parameters() if p.requires_grad), lr=1e-4)

unet.train()
for step, batch in enumerate(loader, start=1):
    pixels = batch["pixel_values"].to(DEVICE)
    tokens = pipe.tokenizer(list(batch["caption"]), padding=True, truncation=True, return_tensors="pt")
    tokens = {k: v.to(DEVICE) for k, v in tokens.items()}
    with torch.no_grad():
        latents = pipe.vae.encode(pixels).latent_dist.sample() * pipe.vae.config.scaling_factor
        noise = torch.randn_like(latents)
        timesteps = torch.randint(0, pipe.scheduler.config.num_train_timesteps, (latents.size(0),), device=DEVICE)
        noisy = pipe.scheduler.add_noise(latents, noise, timesteps)
        hidden = pipe.text_encoder(**tokens).last_hidden_state
    pred = unet(noisy, timesteps, encoder_hidden_states=hidden).sample
    loss = torch.nn.functional.mse_loss(pred, noise)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    print(f"step {step} loss={loss.item():.4f}")
    if step >= 200:
        break

unet.save_pretrained(OUT)
manifest = {
    "base": BASE,
    "revision": REVISION,
    "device": str(DEVICE),
    "rank": 8,
    "steps": step,
    "runtime_lora_swap": False,
}
(OUT / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
print("wrote", OUT)

Run it with the project environment, not a global interpreter.

uv run python train_lora.py
ls artifacts/style-lora
# Expect adapter weights plus manifest.json

Generate the same holdout prompts with the frozen base and with the adapter attached. Keep the seed fixed. If the subject is absent, train longer or clean the captions. If the background is a copy of one training photo, you overfit. Fix data before you merge.

from diffusers import StableDiffusionPipeline
from peft import PeftModel
import torch

device = "mps" if torch.backends.mps.is_available() else "cpu"
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", revision="pin-this"
).to(device)
base_image = pipe("sks mug on a hotel nightstand, warm lamp, 35mm", generator=torch.manual_seed(0)).images[0]
base_image.save("artifacts/holdout-base.png")
pipe.unet = PeftModel.from_pretrained(pipe.unet, "artifacts/style-lora")
lora_image = pipe("sks mug on a hotel nightstand, warm lamp, 35mm", generator=torch.manual_seed(0)).images[0]
lora_image.save("artifacts/holdout-lora.png")

Failure modes

• CUDA in the script. Remove it. This course is MPS or CPU.
• MPS op missing. Move the failing module to CPU, or train that step on CPU. Do not invent a CUDA fallback.
• Saving the whole pipeline as if it were a LoRA. Save the adapter directory. Keep the base revision in the manifest.
• No holdout comparison. You cannot tell personalization from memorization.
• Trying to load the adapter in Xcode. There is no documented Core AI LoRA slot. Merge first.

Done when

• artifacts/style-lora contains adapter weights and manifest.json with base, revision, device, rank, and steps.
• The device string is mps or cpu.
• Holdout previews exist for the same prompt and seed, base versus adapter.
• You have not claimed the adapter is the Core AI model file.

Keep learning

PyTorch to Core AI in Xcode · From the metal to the model · Core AI models, typed · Model architectures in plain English · LLM LoRA for your writing style