ml-diffusion-finetuning

Fine-tuning diffusion models for custom styles or subjects

Diffusion Fine-Tuning

Scope: Fine-tuning diffusion models with DreamBooth, LoRA, textual inversion, dataset preparation Lines: ~360 Last Updated: 2025-10-18

When to Use This Skill

Activate this skill when:

• Fine-tuning diffusion models for custom styles or subjects
• Training DreamBooth models for specific people, objects, or concepts
• Implementing LoRA (Low-Rank Adaptation) for efficient fine-tuning
• Using textual inversion to add new concepts without model modification
• Preparing datasets for diffusion model training
• Debugging overfitting or underfitting in fine-tuned models
• Merging multiple LoRAs or checkpoints
• Optimizing training hyperparameters for quality and speed

Core Concepts

Fine-Tuning Approaches

Full Fine-Tuning:

• Train entire model (billions of parameters)
• Requires massive compute (A100/H100 for days)
• Best quality but extremely expensive
• Risk of catastrophic forgetting

DreamBooth:

• Fine-tune subset of model on 3-20 images
• Preserves general knowledge with regularization
• Moderate compute (A100 for 30-60 minutes)
• Best for specific subjects (people, pets, objects)

LoRA (Low-Rank Adaptation):

• Train small adapter matrices (rank 4-128)
• Only 0.5-5% of model parameters
• Fast training (10-30 minutes on A100)
• Easy to share and merge (10-100MB files)

Textual Inversion:

• Learn new token embedding (1-8 vectors)
• Model weights frozen, only embedding trained
• Very fast (5-15 minutes on A100)
• Limited flexibility but lightweight

DreamBooth Theory

Core Idea: Fine-tune model to associate trigger word with specific subject

Training Process:

1. Collect 3-20 images of subject
2. Choose unique trigger word (e.g., "sks person", "xyz dog")
3. Fine-tune model with caption: "[trigger] [class]"
4. Use regularization images to prevent overfitting

Regularization: Generate ~100-200 images of class (e.g., "person", "dog") to maintain general knowledge

Key Parameters:

• Learning rate: 1e-6 to 5e-6 (lower = safer)
• Training steps: 500-2000 (depends on dataset size)
• Regularization weight: 0.5-1.0

LoRA for Diffusion

Architecture: Add low-rank matrices to attention layers

W_modified = W_pretrained + α * (A × B)
where A is rxd, B is dxr (r << d)

Rank Selection:

• Rank 4-8: Simple styles, fast training, small files
• Rank 16-32: Complex styles, balanced
• Rank 64-128: Maximum flexibility, slower training

Benefits over Full Fine-Tuning:

• 100x faster training
• 1000x smaller files (10-100MB vs 2-7GB)
• No catastrophic forgetting
• Easy to merge multiple LoRAs
• Can be applied/removed at inference

Dataset Preparation

Image Requirements:

• Quantity: 3-5 (minimum), 10-20 (ideal), 50+ (style training)
• Quality: High resolution, clear subject, good lighting
• Diversity: Different angles, poses, lighting, backgrounds
• Consistency: Same subject, avoid multiple people/objects

Captioning:

• DreamBooth: "[trigger] [class]" (e.g., "sks person")
• LoRA: Descriptive captions (helps generalization)
• Textual Inversion: Fixed template with trigger word

Data Augmentation: Not recommended (model handles variation)

---

Patterns

DreamBooth Training (diffusers)

from diffusers import StableDiffusionPipeline, DreamBoothTrainer
from diffusers.utils import make_image_grid
import torch

# 1. Prepare dataset
# - instance_images/: Your 10-20 images
# - class_images/: 100-200 regularization images (generate with SD)

# 2. Generate regularization images (one-time)
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda")

class_prompt = "a photo of a person"  # Match your subject class
for i in range(200):
    image = pipe(class_prompt, num_inference_steps=50).images[0]
    image.save(f"class_images/{i:04d}.png")

# 3. Train DreamBooth
from accelerate import Accelerator
from diffusers import DDPMScheduler

accelerator = Accelerator(mixed_precision="fp16")

# Training config
config = {
    "pretrained_model": "runwayml/stable-diffusion-v1-5",
    "instance_data_dir": "./instance_images",
    "class_data_dir": "./class_images",
    "instance_prompt": "a photo of sks person",  # Your trigger
    "class_prompt": "a photo of a person",       # Class for regularization
    "learning_rate": 5e-6,
    "max_train_steps": 800,
    "train_batch_size": 1,
    "gradient_accumulation_steps": 1,
    "prior_preservation": True,
    "prior_preservation_weight": 1.0,
    "output_dir": "./dreambooth_output",
}

# Note: Full training script requires ~200 lines
# Use diffusers training script: examples/dreambooth/train_dreambooth.py

Command-line training (recommended):

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
  --instance_data_dir="./instance_images" \
  --class_data_dir="./class_images" \
  --output_dir="./dreambooth_output" \
  --instance_prompt="a photo of sks person" \
  --class_prompt="a photo of a person" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=800 \
  --prior_preservation \
  --prior_loss_weight=1.0

LoRA Training (diffusers + PEFT)

import modal
from pathlib import Path

# Modal deployment for LoRA training
app = modal.App("lora-training")

image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "diffusers==0.30.0",
        "transformers==4.44.0",
        "accelerate==0.33.0",
        "peft==0.11.0",
        "bitsandbytes==0.43.0",
        "torch==2.1.0",
        "torchvision==0.16.0",
        "xformers==0.0.27",
    )
)

@app.function(
    gpu="A100",  # 40GB VRAM recommended
    image=image,
    timeout=3600,  # 1 hour
    volumes={"/data": modal.Volume.from_name("lora-data")},
)
def train_lora(
    instance_images_dir: str,
    instance_prompt: str,
    output_name: str,
    steps: int = 1000,
    rank: int = 16,
    learning_rate: float = 1e-4,
):
    """Train LoRA for Stable Diffusion"""
    from diffusers import StableDiffusionPipeline
    from peft import LoraConfig, get_peft_model
    import torch
    from torch.utils.data import Dataset, DataLoader
    from PIL import Image
    import os

    # Load base model
    pipe = StableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        torch_dtype=torch.float16,
    ).to("cuda")

    # Configure LoRA
    lora_config = LoraConfig(
        r=rank,                    # Rank (4-128)
        lora_alpha=rank,           # Scaling factor
        target_modules=[           # Which layers to adapt
            "to_q", "to_k", "to_v", "to_out.0"
        ],
        lora_dropout=0.0,
        bias="none",
    )

    # Apply LoRA to UNet
    pipe.unet = get_peft_model(pipe.unet, lora_config)
    pipe.unet.print_trainable_parameters()

    # Dataset
    class DreamBoothDataset(Dataset):
        def __init__(self, images_dir, prompt):
            self.images = [
                os.path.join(images_dir, f)
                for f in os.listdir(images_dir)
                if f.endswith(('.png', '.jpg', '.jpeg'))
            ]
            self.prompt = prompt

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

        def __getitem__(self, idx):
            image = Image.open(self.images[idx]).convert("RGB")
            image = image.resize((512, 512))
            return {"image": image, "prompt": self.prompt}

    dataset = DreamBoothDataset(instance_images_dir, instance_prompt)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

    # Optimizer
    optimizer = torch.optim.AdamW(
        pipe.unet.parameters(),
        lr=learning_rate,
        betas=(0.9, 0.999),
        weight_decay=1e-2,
    )

    # Training loop
    pipe.unet.train()
    for step in range(steps):
        batch = next(iter(dataloader))

        # Forward pass
        latents = pipe.vae.encode(batch["image"]).latent_dist.sample()
        latents = latents * pipe.vae.config.scaling_factor

        # Add noise
        noise = torch.randn_like(latents)
        timesteps = torch.randint(0, pipe.scheduler.config.num_train_timesteps, (1,))

        noisy_latents = pipe.scheduler.add_noise(latents, noise, timesteps)

        # Predict noise
        encoder_hidden_states = pipe.encode_prompt(batch["prompt"])
        noise_pred = pipe.unet(noisy_latents, timesteps, encoder_hidden_states).sample

        # Loss
        loss = torch.nn.functional.mse_loss(noise_pred, noise)

        # Backward
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if step % 100 == 0:
            print(f"Step {step}/{steps}, Loss: {loss.item():.4f}")

    # Save LoRA weights
    output_path = f"/data/{output_name}"
    pipe.unet.save_pretrained(output_path)
    print(f"LoRA saved to {output_path}")

    return output_path

Simpler approach using diffusers training script:

accelerate launch train_dreambooth_lora.py \
  --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
  --instance_data_dir="./instance_images" \
  --output_dir="./lora_output" \
  --instance_prompt="a photo of sks person" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=1e-4 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=1000 \
  --use_lora \
  --lora_rank=16

Using Trained LoRA

from diffusers import StableDiffusionPipeline
import torch

# Load base model
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda")

# Load LoRA weights
pipe.unet.load_attn_procs("./lora_output")

# Generate with LoRA
prompt = "a photo of sks person as an astronaut"
image = pipe(
    prompt=prompt,
    num_inference_steps=30,
    guidance_scale=7.5,
).images[0]

image.save("lora_output.png")

Merging Multiple LoRAs

from diffusers import StableDiffusionPipeline
from safetensors.torch import load_file
import torch

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda")

# Load multiple LoRAs with different weights
lora_weights = [
    ("./lora_style_1", 0.7),   # 70% influence
    ("./lora_style_2", 0.5),   # 50% influence
    ("./lora_subject", 1.0),   # 100% influence
]

# Apply LoRAs
for lora_path, weight in lora_weights:
    pipe.load_lora_weights(lora_path, adapter_name=lora_path)
    pipe.set_adapters([lora_path], adapter_weights=[weight])

# Generate with merged LoRAs
prompt = "sks person in xyz style"
image = pipe(prompt).images[0]

Textual Inversion

# Textual inversion training (simplified)
from diffusers import StableDiffusionPipeline
import torch

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda")

# Initialize new token embedding
placeholder_token = "<my-concept>"
num_vectors = 1  # Can use 2-8 for complex concepts

# Add token to tokenizer
pipe.tokenizer.add_tokens([placeholder_token])
pipe.text_encoder.resize_token_embeddings(len(pipe.tokenizer))

# Initialize embedding (copy from similar word)
token_id = pipe.tokenizer.convert_tokens_to_ids(placeholder_token)
similar_token_id = pipe.tokenizer.convert_tokens_to_ids("person")

with torch.no_grad():
    pipe.text_encoder.get_input_embeddings().weight[token_id] = \
        pipe.text_encoder.get_input_embeddings().weight[similar_token_id].clone()

# Training: Only optimize new embedding, freeze rest
# (Full training code ~150 lines, use diffusers script)

Command-line training:

accelerate launch textual_inversion.py \
  --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
  --train_data_dir="./instance_images" \
  --learnable_property="object" \
  --placeholder_token="<my-concept>" \
  --initializer_token="person" \
  --resolution=512 \
  --train_batch_size=1 \
  --learning_rate=5e-4 \
  --max_train_steps=3000 \
  --output_dir="./textual_inversion_output"

Dataset Preparation Script

from PIL import Image
import os
from pathlib import Path

def prepare_dreambooth_dataset(
    input_dir: str,
    output_dir: str,
    target_size: int = 512,
    min_images: int = 5,
):
    """Prepare and validate DreamBooth dataset"""

    input_path = Path(input_dir)
    output_path = Path(output_dir)
    output_path.mkdir(exist_ok=True)

    # Supported formats
    image_extensions = {'.jpg', '.jpeg', '.png', '.webp'}

    # Collect images
    images = [
        f for f in input_path.iterdir()
        if f.suffix.lower() in image_extensions
    ]

    if len(images) < min_images:
        raise ValueError(f"Need at least {min_images} images, found {len(images)}")

    print(f"Processing {len(images)} images...")

    # Process each image
    for i, img_path in enumerate(images):
        img = Image.open(img_path).convert("RGB")

        # Resize to square, maintaining aspect ratio
        width, height = img.size
        min_dim = min(width, height)

        # Center crop to square
        left = (width - min_dim) // 2
        top = (height - min_dim) // 2
        img = img.crop((left, top, left + min_dim, top + min_dim))

        # Resize to target size
        img = img.resize((target_size, target_size), Image.LANCZOS)

        # Save
        output_file = output_path / f"{i:04d}.png"
        img.save(output_file, "PNG")
        print(f"Processed: {output_file}")

    print(f"Dataset ready: {len(images)} images in {output_path}")

    # Create caption file
    caption_file = output_path / "captions.txt"
    with open(caption_file, "w") as f:
        f.write("# Edit this file to add captions\n")
        f.write("# Format: image_name.png: caption\n\n")
        for i in range(len(images)):
            f.write(f"{i:04d}.png: a photo of sks person\n")

    print(f"Caption template created: {caption_file}")

# Usage
prepare_dreambooth_dataset(
    input_dir="./raw_images",
    output_dir="./instance_images",
    target_size=512,
)

---

Quick Reference

Method Comparison

Method            | Train Time | File Size | Quality | Flexibility | Best For
------------------|------------|-----------|---------|-------------|------------------
Full Fine-Tune    | Days       | 2-7GB     | Best    | Highest     | Max quality, budget unlimited
DreamBooth        | 30-60 min  | 2-7GB     | High    | High        | Specific subjects, traditional
LoRA              | 10-30 min  | 10-100MB  | High    | High        | Modern default, shareable
Textual Inversion | 5-15 min   | <1MB      | Medium  | Low         | Simple concepts, lightweight

Hyperparameter Guidelines

Parameter                  | DreamBooth | LoRA     | Textual Inv
---------------------------|------------|----------|-------------
Learning Rate              | 5e-6       | 1e-4     | 5e-4
Training Steps             | 800-2000   | 1000     | 3000
Batch Size                 | 1          | 1-2      | 1
Gradient Accumulation      | 1-4        | 1-2      | 1
LoRA Rank                  | N/A        | 16-32    | N/A
Number of Images           | 10-20      | 10-50    | 5-15
Regularization Weight      | 1.0        | Optional | N/A

Training Checklist

✅ Dataset: 5-20 images, high quality, diverse angles
✅ Images: Resized to 512x512, square crop, PNG format
✅ Captions: Consistent format with trigger word
✅ Regularization: 100-200 class images (DreamBooth)
✅ GPU: A100 recommended (T4 works but slower)
✅ Learning Rate: Start conservative (lower is safer)
✅ Checkpoints: Save every 100-200 steps
✅ Validation: Test generation every 200 steps
✅ Monitoring: Track loss, watch for overfitting

Overfitting Signs

❌ Training loss near zero but validation loss high
❌ Generated images identical to training images
❌ Can only generate exact training poses/angles
❌ Poor generalization to new prompts
❌ Artifacts or distortions in novel scenes

✅ Solution: Reduce steps, lower learning rate, add regularization, increase dataset diversity

---

Anti-Patterns

❌ Too few training images: Using 1-3 images ✅ Use 5-10 minimum, 10-20 ideal for DreamBooth/LoRA

❌ Overfitting with too many steps: Training until loss = 0 ✅ Monitor validation, stop at 800-1500 steps (DreamBooth), 1000 (LoRA)

❌ No regularization (DreamBooth): Model forgets general knowledge ✅ Always use prior preservation with 100-200 class images

❌ Generic trigger words: Using "person", "dog" (conflicts with model) ✅ Use unique identifiers: "sks person", "xyz dog"

❌ Wrong LoRA rank: Using rank 128 for simple styles ✅ Start with rank 16, increase only if needed

❌ High learning rate: Using 1e-3 or higher ✅ Use 5e-6 (DreamBooth), 1e-4 (LoRA), 5e-4 (textual inversion)

❌ Low quality dataset: Blurry, inconsistent, multiple subjects ✅ High resolution, clear subject, consistent lighting, single subject

❌ No validation during training: Training blind ✅ Generate test images every 200 steps to monitor progress

❌ Training on CPU or small GPU: T4 with SDXL ✅ Use A100 for SDXL, L40S/A100 for SD 1.5

---

Related Skills

• diffusion-model-basics.md - Core diffusion concepts and architecture
• stable-diffusion-deployment.md - Deploying fine-tuned models to production
• lora-peft-techniques.md - General LoRA/PEFT concepts
• modal-gpu-workloads.md - GPU training infrastructure
• modal-functions-basics.md - Modal deployment patterns

---

Last Updated: 2025-10-18 Format Version: 1.0 (Atomic)