Stable Diffusion — generating images by denoising

Stable Diffusion Deep Dive - 01: Architecture Overview

This series dissects Stable Diffusion v1.5 (a Latent Diffusion Model, Rombach et al., CVPR 2022) using the Hugging Face diffusers PyTorch implementation and real downloaded weights at /mnt/mydata/huggingface/.

The big idea. Pixel-space diffusion is expensive: denoising a 512x512x3 image directly means every U-Net step operates on ~800k values. Latent Diffusion instead runs diffusion in the compressed latent space of a pretrained autoencoder (VAE). A 512x512 image becomes a 64x64x4 latent (a 48x spatial-element reduction), so the same U-Net does far less work while the VAE decoder restores high-frequency detail.

Four components

Series map

1. The generation pipeline, end to end

For text-to-image the data flow at inference is:

`` prompt --(CLIP tokenizer+text encoder)--> c in R^{77 x 768} z_T ~ N(0, I) # 4x64x64 latent noise for t = T ... 1: # scheduler timesteps eps_hat = UNet(z_t, t, c) # predict the noise z_{t-1} = scheduler.step(eps_hat, t, z_t) x = VAE.decode(z_0 / scaling) # 3x512x512 image ``

Classifier-free guidance (CFG) runs the U-Net twice per step - once with the prompt embedding c and once with an empty/"" embedding c_uncond - and extrapolates: $$\hat\epsilon = \epsilon_\text{uncond} + s\,(\epsilon_\text{cond} - \epsilon_\text{uncond})$$ with guidance scale s (typically 7.5). We visualize its effect in NB05.

2. U-Net structure (why cross-attention?)

The U-Net is an encoder-bottleneck-decoder with residual (ResNet) blocks for spatial processing and Transformer blocks for attention. Each Transformer block has:

• self-attention over spatial latent positions (global spatial coherence),
• cross-attention where queries come from the image latent and keys/values come from the 77x768 text embedding - this is how the prompt steers generation,
• a feed-forward (GEGLU) MLP.

Timestep t enters through a sinusoidal embedding -> MLP added to every ResNet block (adaptive conditioning on the noise level). The channel widths follow block_out_channels = (320, 640, 1280, 1280) with attention at the 8x/16x/32x down-sampled stages.

3. Confirm the architecture from the real model_index.json

If the weights are downloaded we read the pipeline manifest directly; otherwise we show the canonical layout. This proves which classes diffusers will instantiate.

4. Takeaways

• SD is latent diffusion: a VAE compresses 8x so the U-Net is cheap.
• The U-Net predicts noise (epsilon) and is conditioned on text via cross-attention and on the noise level via a timestep embedding.
• The scheduler is a swappable numerical integrator of the reverse SDE/ODE.
• Everything is modular in diffusers - NB03 loads each real sub-module and counts its parameters.

Stable Diffusion Deep Dive - 02: Design Choices & Training

How SD is trained and why each design choice was made. We keep the math concrete and runnable - every claim below is checked with a tiny numeric demo.

1. The DDPM forward process

A fixed variance schedule beta_1..beta_T gradually noises a clean latent z_0. With alpha_t = 1 - beta_t and alpha_bar_t = prod alpha_s, the closed form for any timestep is $$z_t = \sqrt{\bar\alpha_t}\,z_0 + \sqrt{1-\bar\alpha_t}\,\epsilon,\quad \epsilon\sim N(0,I).$$ SD1.5 uses T=1000 and a scaled-linear beta schedule from 0.00085 to 0.012.

2. The training objective (epsilon-prediction)

The network eps_theta(z_t, t, c) is trained to predict the noise that was added. The (simplified) loss from Ho et al. 2020, used by SD, is just an MSE: $$L = \mathbb{E}_{z_0,\epsilon,t}\big\|\epsilon - \epsilon_\theta(z_t,t,c)\big\|^2.$$ Why predict noise rather than the image? It (a) keeps the target unit-variance across all t (stable optimization), and (b) is equivalent to score matching - eps_theta is proportional to the score grad log p(z_t).

We verify the equivalence x0 <-> eps numerically:

3. Classifier-free guidance (CFG) - trained by prompt dropout

During training the text condition is randomly replaced by the empty prompt ~10% of the time. This yields a single network that models both the conditional and unconditional score, enabling guidance at sample time: $$\hat\epsilon = \epsilon_\text{uncond} + s\,(\epsilon_\text{cond}-\epsilon_\text{uncond}).$$ Larger s -> stronger prompt adherence but less diversity and more saturation.

4. Samplers: DDPM vs DDIM vs PNDM

• DDPM (stochastic): reverse the Markov chain step by step - needs many steps.
• DDIM (deterministic): a non-Markovian ODE view; same model, far fewer steps (20-50) and reproducible outputs for a fixed seed.
• PNDM / DPM-Solver: higher-order ODE solvers - SD1.5's default is PNDM; DPM-Solver++ gives good quality in ~20 steps.

The scheduler is decoupled from the network, which is why you can swap samplers without retraining (shown in NB05).

5. How SD 1.5 was actually trained (facts)

• Data: subsets of LAION-5B (LAION-2B-en, aesthetics-filtered).
• Init: resumed from SD 1.2 and trained ~595k steps at 512x512 on laion-aesthetics v2 5+, with 10% text-conditioning dropout for CFG.
• VAE: trained separately with an L1/L2 + LPIPS perceptual loss and a patch adversarial term, plus a tiny KL penalty on the latent (hence AutoencoderKL).
• Text encoder: the CLIP ViT-L/14 text tower, frozen during LDM training.
• Optimizer: AdamW; the U-Net is the only large module updated for the diffusion objective.

This modular, frozen-encoder recipe is exactly what makes fine-tuning cheap - the subject of NB06/07.

Stable Diffusion Deep Dive - 03: The Real PyTorch Model

We load the downloaded SD 1.5 weights with diffusers, then inspect each sub-module as a plain torch.nn.Module: parameter counts, dtypes, the U-Net block tree, and a single real forward pass through the VAE and U-Net to confirm the shapes derived symbolically in NB01.

If the weights / diffusers are missing, the notebook builds a tiny synthetic UNet2DConditionModel + AutoencoderKL with the same interface so every cell still runs (in 'schematic mode').

1. Load the pipeline (or a synthetic fallback)

We load in float16 on GPU when available (halves memory), else float32 on CPU.

2. Parameter counts per component

This is the concrete answer to 'how big is each part of SD?'.

3. The U-Net block tree

We print the top-level down/mid/up blocks and the class of each - this is the encoder-bottleneck-decoder with cross-attention discussed in NB01.

4. A real forward pass (shape check)

Encode a random image to a latent with the VAE, add noise, and run one U-Net denoising step. This exercises the actual PyTorch graph and confirms the (4, H/8, W/8) latent shape.

5. Takeaways

• diffusers exposes each component as an ordinary nn.Module you can index, hook, and profile.
• The U-Net dominates the parameter budget (~86% of SD1.5).
• Latent shape is exactly (4, H/8, W/8); the U-Net is a shape-preserving noise predictor. NB04 hooks these modules to see what they compute.

Stable Diffusion Deep Dive - 04: Visualizing the Layers

We open the black box with forward hooks and direct decoding:

1. VAE latent channels - what the 4-channel compressed code looks like.
2. U-Net feature maps - activations at several down/up stages.
3. Cross-attention maps - where each prompt word attends in the image grid (the mechanism behind prompt-to-region control, a la Prompt-to-Prompt).

All figures are saved to stable_diffusion/notebooks/figures/. In schematic mode the same visualizations run on the tiny synthetic model.

figures -> /home/nadurthi/repos/deeplearning/stable_diffusion/notebooks/figures

`torch_dtype` is deprecated! Use `dtype` instead!

You have disabled the safety checker for <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> by passing `safety_checker=None`. Ensure that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling it only for use-cases that involve analyzing network behavior or auditing its results. For more information, please have a look at https://github.com/huggingface/diffusers/pull/254 .

real model loaded

1. VAE latent channels

Encode a synthetic RGB test pattern and display the four latent channels. Each is a 64x64 (real) heatmap; the VAE has learned an edge/colour/texture decomposition.

saved /home/nadurthi/repos/deeplearning/stable_diffusion/notebooks/figures/04_vae_latents.png | latent shape (4, 64, 64)

2. U-Net feature maps via forward hooks

We register hooks on a few down/up blocks and capture the activation from a single denoising step, plotting the mean over channels at each stage. Deeper stages have more channels at coarser spatial resolution (the U-Net funnel).

captured stages: [('down0', (1, 320, 32, 32)), ('down1', (1, 640, 16, 16)), ('down2', (1, 1280, 8, 8)), ('down3', (1, 1280, 8, 8)), ('mid', (1, 1280, 8, 8)), ('up0', (1, 1280, 16, 16)), ('up1', (1, 1280, 32, 32)), ('up2', (1, 640, 64, 64)), ('up3', (1, 320, 64, 64))]

saved /home/nadurthi/repos/deeplearning/stable_diffusion/notebooks/figures/04_unet_features.png

3. Cross-attention maps (prompt -> image regions)

Cross-attention is where the text conditions the image. For each word token we extract the attention weight over the spatial grid: bright regions are where that word 'paints'. This only produces meaningful maps with the real model + a real prompt; in schematic mode we show the mechanism on random tensors.

We use a lightweight attention processor that records the softmax attention probabilities during a single U-Net call.

captured cross-attn layers: ['down_blocks.2.attentions.0.transformer_blocks.0.attn2.processor', 'down_blocks.2.attentions.1.transformer_blocks.0.attn2.processor', 'up_blocks.1.attentions.0.transformer_blocks.0.attn2.processor'] ...

saved /home/nadurthi/repos/deeplearning/stable_diffusion/notebooks/figures/04_cross_attention.png

4. Takeaways

• The VAE latent is a compact, semantically-structured 4-channel code.
• U-Net stages trade spatial resolution for channel depth (classic U-Net funnel).
• Cross-attention maps localize each prompt token to image regions - the basis for editing techniques (Prompt-to-Prompt, attend-and-excite) and for the reference-image editing we fine-tune in NB06.

Stable Diffusion Deep Dive - 05: Inference Examples

Concrete generations with the real pipeline:

1. Text-to-image on several prompts (seeded, reproducible).
2. Guidance-scale sweep - visualize the CFG effect from NB02.
3. Scheduler swap - PNDM vs DDIM vs DPM-Solver++ at low step counts.
4. Image-to-image - partial noising + denoising for edits.

Outputs are saved to figures/. Without weights, the notebook prints the exact calls it would run and generates placeholder tensors so it still executes.

1. Text-to-image

2. Guidance-scale sweep (CFG)

Low s = diverse but weak prompt adherence; high s = strong adherence, more saturation.

3. Scheduler swap at low steps

Same U-Net, different ODE/SDE solver. DPM-Solver++ is strong at ~20 steps.

4. Image-to-image (edit an existing image)

strength controls how much of the input is kept (low = subtle edit, high = mostly regenerate).

5. Takeaways

• One pipeline call hides the full loop from NB01; every knob (steps, guidance, scheduler, seed, strength) is exposed.
• CFG and scheduler choice trade quality vs speed vs diversity.
• img2img is 'partial diffusion': start from a noised real image instead of pure noise - the conceptual bridge to the editing task fine-tuned in NB06.

Stable Diffusion Deep Dive - 06: Fine-tuning for Image Editing

We fine-tune SD 1.5 for instruction-guided / reference image editing using the Hugging Face diffusers training API and the molbal/multi_reference_image_editing dataset.

Approach. The InstructPix2Pix recipe adapts SD so the U-Net is additionally conditioned on an input image: the U-Net input channels are expanded from 4 (noisy latent) to 8 (noisy latent + the VAE-encoded reference latent), and the text carries the edit instruction. We train with the same epsilon-MSE objective. This is the exact structure of diffusers' train_instruct_pix2pix.py.

1. Inspect the dataset

We stream a few rows to see the schema (reference image(s), edit instruction, target image). Field names vary by dataset revision, so we detect them dynamically.

2. Column mapping

Editing datasets typically expose an input image, an edit prompt, and an edited image. We map whatever this dataset calls them onto those three roles so the training code is dataset-agnostic.

3. Build the InstructPix2Pix-style U-Net (4->8 input channels)

We load SD 1.5 and surgically widen conv_in to accept the concatenated reference latent, copying the pretrained weights into the first 4 channels and zero-initializing the new 4 (so the model starts as the identity-conditioned base).

4. Preprocessing -> latents

For each example: VAE-encode both the reference and target images to latents, tokenize the edit prompt. We build a tiny in-memory batch (a handful of examples) to keep the demo fast; swap in the streamed dataset for real training.

5. The training loop (HF-style, epsilon-MSE + InstructPix2Pix conditioning)

This is the same objective as diffusers/examples/instruct_pix2pix/train_instruct_pix2pix.py, reduced to its essence. For each step we noise the target latent, concatenate the reference latent, and train the U-Net to predict the added noise. We log the loss curve. We use accelerate conventions (AdamW + grad clipping).

6. Save the fine-tuned U-Net

We persist to stable_diffusion/checkpoints/ so it can be reloaded into a pipeline.

6b. Fine-tuning hyperparameters reference

The DreamBooth / diffusers fine-tuning algorithm exposes the hyperparameters below. The teaching loop above uses a subset (LR, batch size, AdamW settings, grad-accum, grad clipping, seed); the rest map directly onto the official train_dreambooth.py / train_instruct_pix2pix.py flags for a full run.

Notes. prior_loss_weight / num_class_images / class_prompt apply only when with_prior_preservation=True (DreamBooth-style regularization). Max_steps overriding Epochs is why the teaching loop uses a step budget (SD_FT_STEPS). The AdamW betas/epsilon, weight_decay, and max_grad_norm=1.0 used in the loop above are exactly these defaults.

7. What a real run looks like (the official script)

For a full run, use the maintained example rather than this teaching loop: ``bash accelerate launch diffusers/examples/instruct_pix2pix/train_instruct_pix2pix.py \ --pretrained_model_name_or_path=$SD_MODEL_DIR \ --dataset_name=molbal/multi_reference_image_editing \ --resolution=256 --random_flip \ --train_batch_size=4 --gradient_accumulation_steps=4 --gradient_checkpointing \ --max_train_steps=15000 --checkpointing_steps=2500 \ --learning_rate=5e-05 --lr_warmup_steps=0 \ --mixed_precision=fp16 --use_ema \ --output_dir=stable_diffusion/checkpoints/ip2p_full `` The next notebook (07) documents everything that tends to go wrong here.

Stable Diffusion Deep Dive - 07: Fine-tuning Issues & Diagnostics

Fine-tuning SD is notoriously finicky. This notebook is a field guide to the failure modes seen in NB06 - what they look like, why they happen, and concrete fixes - with small runnable demos where a demo helps.

1. Convergence: the loss barely moves

Symptom. The epsilon-MSE loss looks flat/noisy and never clearly trends down.

Why. The training target is ||eps - eps_hat||^2 with eps ~ N(0,I) and a uniformly random timestep each step. The irreducible noise variance dominates, so per-step loss is inherently high-variance and a good model still sits around a plateau. Visual quality improves long before the scalar loss looks 'nice'.

Fixes.

• Judge progress by generated samples / validation edits, not the raw loss.
• Use an EMA copy of the weights for evaluation (--use_ema).
• Try min-SNR-gamma loss weighting to reduce timestep variance.
• Average the loss over many steps (a moving average) before reading a trend.

2. Overfitting on a small edit set

Symptom. Training loss drops nicely but the model memorizes the few training pairs: it reproduces training targets and ignores new instructions; outputs collapse to a narrow style; prompt diversity vanishes.

Why. SD's U-Net has ~860M parameters; a few hundred image pairs is trivially memorizable, especially at higher LR / many epochs.

Fixes.

• LoRA / partial fine-tuning instead of full-U-Net updates (far fewer trainable params → strong regularization). See the demo below.
• Lower LR (1e-6..5e-5), fewer steps, early stopping on a held-out edit set.
• Data augmentation (flip, crop) and more/varied data.
• EMA and weight decay; keep the text encoder frozen.

3. Out-of-memory (OOM)

Symptom. CUDA out of memory during the backward pass, often only after a few steps (fragmentation) or at higher resolution/batch.

Why. Activations for an 860M U-Net at 512x512 + Adam's two moment buffers (2x the parameter memory) add up fast.

Fixes (roughly in order of impact).

• --gradient_checkpointing (trade compute for memory - biggest win).
• --mixed_precision=fp16/bf16.
• Smaller --train_batch_size + larger --gradient_accumulation_steps.
• 8-bit Adam (bitsandbytes) - halves optimizer state.
• LoRA (few trainable params → tiny optimizer state).
• Lower --resolution (256 instead of 512).
• xformers / SDPA memory-efficient attention.

4. Instability / NaNs

Symptom. Loss becomes NaN/Inf; images turn to noise or solid colour.

Why. fp16 underflow/overflow, LR too high, or an un-scaled loss.

Fixes. Use bf16 if the GPU supports it (no loss scaler needed), keep LR small, gradient clipping (max_norm=1.0), and ensure the VAE runs in fp32 when encoding (the SD VAE is known to overflow in fp16).

5. Data / label issues specific to editing datasets

• Misaligned pairs: input/edited images not registered → the model learns noise. Verify pairs visually (NB06 §1).
• Weak/ambiguous instructions: 'make it nicer' gives no gradient signal. Prefer concrete edits.
• Resolution/aspect mismatch: resizing non-square images distorts geometry; center-crop instead.
• Channel-conditioning bug: forgetting to zero-init the extra 4 conv_in channels (NB06 §3) makes early training diverge.

6. Evaluation is hard

There is no single scalar that tracks 'good edits'. Practical signals:

• CLIP similarity between the output and the target instruction (directional CLIP).
• LPIPS between output and target for reconstruction-style edits.
• Fixed-seed qualitative panels across training checkpoints (eyeball drift).

Log all three; do not trust the training loss alone.

7. A pragmatic recipe that usually works

1. Start with LoRA (rank 8-16), LR 1e-4, fp16, 256px, batch 4 + grad-accum 4.
2. Turn on gradient_checkpointing and (optionally) 8-bit Adam.
3. Freeze the text encoder and VAE; train only U-Net (or LoRA adapters).
4. Evaluate every N steps on a held-out set with fixed seeds; keep EMA weights.
5. Only move to full fine-tuning if LoRA plateaus and you have enough data.

Bottom line. For our tiny NB06 run, expect a noisy/flat loss and weak edits - that is the expected outcome at this scale and is exactly why these mitigations exist. The loop itself is correct and scales to the full recipe in NB06 §7.

08 - Fine-tuning Stable Diffusion with DreamBooth

DreamBooth (Ruiz et al., 2023) teaches a text-to-image model a new subject from just 3-5 images by binding it to a rare token identifier (e.g. sks) inside a class prompt (a photo of sks dog). The whole U-Net (and optionally the text encoder) is fine-tuned.

The key trick that separates DreamBooth from naive fine-tuning is the prior-preservation loss: we also train on generic class images (a photo of a dog) generated by the frozen base model, so the model keeps its broad notion of the class and does not collapse onto the few instance images (a failure mode called language drift / overfitting).

Dataset: diffusers/dog-example - the canonical 5-image DreamBooth dog set used in the official docs.

1. DreamBooth in one picture

``` instance images (5) --VAE.encode--> z0 'a photo of sks dog' --CLIP--> c_instance

class/prior images --VAE.encode--> z0' 'a photo of a dog' --CLIP--> c_class

loss = MSE(eps_hat(z_t , c_instance), eps) # instance term + lambda * MSE(eps_hat(z_t', c_class), eps') # prior-preservation term ```

• sks is a rare token - it carries almost no prior meaning, so the optimizer is free to attach the new subject to it.
• lambda = prior_loss_weight (default 1.0) balances the two terms.
• Only 3-5 instance images are needed; num_class_images (~200) prior images are generated once by the frozen model.

loaded 5 instance images from diffusers/dog-example
sizes: [(1815, 1967), (2469, 2558), (2796, 2656), (2476, 2612), (2732, 2736)]

saved figures/08_dreambooth_instances.png

2. Load the components we will fine-tune

DreamBooth trains the U-Net (VAE + text encoder frozen in the simplest recipe). We load the real SD 1.5 pieces when available, otherwise a tiny synthetic stand-in so the loop still runs.

loaded real SD 1.5 components; U-Net params = 859.520964 M

3. Preprocess: images + prompts -> latents + text embeddings

Standard SD training pipeline: resize/normalize to [-1, 1], encode to the 4x64x64 latent space with the VAE (scaled by 0.18215), tokenize the prompt and run the CLIP text encoder once (it is frozen).

latents (5, 4, 64, 64) | text embeddings (5, 77, 768)

4. A tiny real DreamBooth training loop (with prior preservation)

We run only DB_STEPS steps (env-overridable, default 20) just to show the loss dropping and confirm gradients flow. For simplicity here we build the prior batch from a couple of dataset images tied to the class prompt a photo of a dog; a real run instead generates ~200 class images once from the frozen base model.

The combined objective is

$$ \mathcal{L} = \mathbb{E}\,\lVert \epsilon - \hat\epsilon_\theta(z_t, c_{inst})\rVert^2 + \lambda\, \mathbb{E}\,\lVert \epsilon' - \hat\epsilon_\theta(z'_t, c_{cls})\rVert^2. $$

step 000  loss=0.3819  (inst=0.3609 prior=0.0210)

step 004  loss=0.6364  (inst=0.1834 prior=0.4530)

step 008  loss=0.0216  (inst=0.0143 prior=0.0073)

step 012  loss=0.2127  (inst=0.0465 prior=0.1662)

step 016  loss=0.4221  (inst=0.3235 prior=0.0987)

saved figures/08_dreambooth_loss.png ; final loss = 0.0286

5. DreamBooth hyperparameters (what to actually turn)

Failure modes to watch: subject overfitting (every prompt returns the exact training pose), language drift (the class word stops meaning the generic class - fixed by prior preservation), and editability loss (cannot recontextualize the subject - lower LR / fewer steps).

6. Full run: the official train_dreambooth.py

For a real DreamBooth run use the maintained script from the diffusers repo (examples/dreambooth/train_dreambooth.py). The cell below prints a ready-to-paste command that points at the local SD 1.5 weights.

Cheaper variant: train_dreambooth_lora.py

For a single consumer GPU, DreamBooth-LoRA trains tiny rank-r adapters on the attention layers instead of the full U-Net - same script family, add --rank=4 and raise the LR to ~1e-4. Output is a small pytorch_lora_weights.safetensors you load with pipe.load_lora_weights(...).

See also: notebook 09 for standard caption-driven text->image fine-tuning, and 07 for a catalogue of fine-tuning failure modes.

09 - Fine-tuning Stable Diffusion with the HF diffusers text-to-image API

This is the standard caption-driven fine-tuning recipe implemented by diffusers/examples/text_to_image/train_text_to_image.py. Unlike DreamBooth (which specializes on one subject from ~5 images), this trains on a large (image, caption) dataset to shift the model toward a whole style/domain.

Dataset: lambdalabs/naruto-blip-captions - a popular, small captioned dataset (BLIP captions of Naruto-style art) that is the drop-in successor to the classic pokemon-blip-captions demo.

1. The text-to-image fine-tuning objective

Exactly the base LDM training objective, now on your (image, caption) pairs:

`` image --VAE.encode--> z0 --add_noise(t)--> z_t caption --CLIP--> c loss = MSE( eps_hat_theta(z_t, t, c), eps ) # train the U-Net ``

The VAE and CLIP text encoder stay frozen; only the U-Net is updated (full fine-tune, or LoRA adapters for the cheap variant).

Repo card metadata block was not found. Setting CardData to empty.

streamed 16 rows from lambdalabs/naruto-blip-captions
  caption: a man with dark hair and brown eyes
  caption: a man in a hoodie with a fire in the background
  caption: a man with a red hair and a black shirt

saved figures/09_t2i_samples.png

2. Load the components (U-Net trainable, VAE + text encoder frozen)

loaded real SD 1.5; U-Net params = 859.5 M (trainable)

3. Preprocess: (image, caption) -> latents + text embeddings

train_text_to_image.py resizes to resolution, applies a random crop/flip, normalizes to [-1, 1], VAE-encodes to latents (x0.18215), and tokenizes captions. We do the same here for our small batch.

latents (16, 4, 64, 64) | text (16, 77, 768)

4. A short real text-to-image training loop

T2I_STEPS steps (env-overridable, default 25) of the epsilon-MSE objective over mini-batches drawn from the streamed rows.

step 000  loss=0.0418

step 005  loss=0.0359

step 010  loss=0.0856

step 015  loss=0.2824

step 020  loss=0.1894

saved figures/09_t2i_loss.png ; final loss = 0.4136

5. Hyperparameters for train_text_to_image.py

6. Full run: the official train_text_to_image.py

Cheaper LoRA variant: train_text_to_image_lora.py

`` accelerate launch train_text_to_image_lora.py \ --pretrained_model_name_or_path=<local SD1.5> \ --dataset_name=lambdalabs/naruto-blip-captions \ --resolution=512 --random_flip --train_batch_size=1 \ --gradient_accumulation_steps=4 --max_train_steps=15000 \ --learning_rate=1e-4 --lr_scheduler=cosine --lr_warmup_steps=0 \ --rank=4 --mixed_precision=fp16 \ --output_dir=./checkpoints/sd-naruto-lora ``

LoRA trains only ~1-3M adapter params (vs 860M for the full U-Net), so it fits comfortably on a single GPU and produces a small pytorch_lora_weights.safetensors loadable via pipe.load_lora_weights.

See also: notebook 08 for subject-driven DreamBooth, and 07 for fine-tuning failure modes (overfitting, NaN, OOM).

10 - Stable Diffusion Inpainting (text + image + mask)

The inpainting model (runwayml/stable-diffusion-inpainting) fills a masked region of an input image, guided by a text prompt, while keeping the unmasked pixels intact. Its inputs are a triple:

1. prompt (text) - what to paint into the hole;
2. image (RGB) - the picture to edit;
3. mask (grayscale) - white = repaint, black = keep.

This notebook explains how the U-Net was modified to consume the mask and image, walks through every tensor shape, summarizes the training process, and runs real inference examples.

1. The architecture change: a 9-channel U-Net conv_in

Standard SD 1.5 text2img: the U-Net's first conv (conv_in) takes the 4-channel noisy latent z_t and predicts the noise. Inpainting concatenates two extra conditioning tensors along the channel axis before conv_in, so its input has 9 channels:

`` in_channels = 4 + 1 + 4 = 9 +-----------------------------+---------+----------------------------+ | noisy latent z_t (4 ch) | mask m | masked-image latent (4 ch)| | = what we denoise | (1 ch) | = VAE(image * (1-mask)) | +-----------------------------+---------+----------------------------+ | conv_in (9 -> 320) | ... rest of the U-Net (unchanged) ... | conv_out (320 -> 4) # still predicts 4-ch epsilon ``

Key points

• Only conv_in changes (in_channels 4->9). Every downsample/upsample block, the cross-attention to the text embedding, and conv_out (still 4 channels) are identical to the base model.
• The mask is a single channel, bilinearly downsampled to the latent resolution (H/8 x W/8) - it is not VAE-encoded.
• The masked-image latent is VAE.encode(image * (1 - mask)) - the known pixels are shown to the network as clean context; the hole is zeroed before encoding.
• The 5 new input channels of conv_in.weight are zero-initialized when converting a text2img checkpoint, so training starts as a no-op perturbation of the base model.

An error occurred while trying to fetch /mnt/mydata/huggingface/runwayml_stable-diffusion-inpainting/unet: Error no file named diffusion_pytorch_model.safetensors found in directory /mnt/mydata/huggingface/runwayml_stable-diffusion-inpainting/unet.

Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead.

`torch_dtype` is deprecated! Use `dtype` instead!

real inpaint load failed -> schematic mode: ValueError('Due to a serious vulnerability issue in `torch.load`, even with `weights_only=True`, we now require users to upgrade torch to at least v2.6 in order to use the function. This version restr
schematic 9-channel UNet: conv_in (9->320), conv_out (320->4)

2. Tensor shapes at every stage

Work an example at the native SD resolution 512x512 (latent 64x64, VAE downsample factor 8). Batch size B=1, text tokens L=77, CLIP hidden D=768.

With classifier-free guidance the batch dim is doubled (uncond + cond), so the U-Net actually sees (2B, 9, 64, 64).

U-Net input (schematic): (1, 9, 64, 64) <- 9 channels

U-Net input : (1, 9, 64, 64)
U-Net output: (1, 4, 64, 64)   <- predicted 4-channel epsilon

Visualize the 9 input channels

The figure shows the three tensors that are concatenated: the noisy latent (4 ch), the downsampled mask (1 ch), and the masked-image latent (4 ch).

saved figures/10_unet9_channels.png

3. How the inpainting model was trained

The inpainting checkpoint is initialized from the text2img SD1.5 weights (with the 5 extra conv_in channels zeroed) and then fine-tuned with the same latent-diffusion objective, but on (image, mask, prompt) triples:

1. Sample a mask. During training a random mask is drawn per image - a mixture of random rectangles, brush strokes, and occasionally the whole image (so the model still works as a plain text2img generator when the mask covers everything). This diverse masking is what makes the model robust to arbitrary hole shapes.
2. Build the 9-channel input. x = concat(z_t, downsample(mask), VAE.encode(image*(1-mask))).
3. Add noise only to z_t (the target latent z_0 = VAE.encode(image)), at a random timestep t, exactly as in base DDPM training. The mask and masked-image-latent conditioning channels are noise-free.
4. Predict epsilon and take the MSE loss on the full latent (the loss is not restricted to the hole; the known region is easy and the model quickly learns to copy it through):

$$\mathcal{L} = \mathbb{E}_{z_0,\,\epsilon,\,t,\,m}\; \big\lVert \epsilon - \hat\epsilon_\theta\big(\,[\,z_t,\,m,\,\mathrm{VAE}(x\odot(1-m))\,],\; t,\; c_{\text{text}}\big)\big\rVert_2^2.$$

1. Classifier-free guidance is trained by randomly dropping the text prompt (~10% of steps), identical to base SD.

At inference the scheduler additionally re-injects the known region at each denoising step in some pipelines, but the runwayml model was trained end-to-end so the network itself learns to preserve context.

schematic loss curve 8 steps
saved figures/10_inpaint_train_loss.png

4. Inference: text + image + mask -> inpainted image

Pull a small demo (image, mask) pair. We try a popular HF inpainting demo dataset and fall back to the synthetic gradient+square if offline.

demo dataset unavailable -> synthetic image+mask: ConnectionError("Couldn't reach 'diffusers/inpaint-example' on the Hub (OfflineModeIsEnabled)")
saved figures/10_inpaint_inputs.png

schematic mode: no real generation (need the inpainting weights).
run:  INCLUDE_INPAINT=1 bash download_weights.sh

5. Prompt sweep on the same hole

Same image + mask, different prompts - shows the model paints different content into the identical region while preserving the surroundings.

schematic mode: prompt sweep needs real weights.

6. Summary

• Architecture: identical to SD1.5 except conv_in is widened 4 -> 9 channels to also ingest the downsampled mask (1) and the masked-image latent (4); conv_out still emits 4-channel epsilon.
• Shapes: at 512x512 the U-Net input is (B, 9, 64, 64) (or (2B, 9, 64, 64) with CFG) and the output is (B, 4, 64, 64).
• Training: initialized from text2img (new channels zeroed), then fine-tuned with random masks + the standard epsilon-MSE objective on the full latent; text dropped ~10% for CFG.
• Inference: StableDiffusionInpaintPipeline(prompt, image, mask_image); white mask = repaint, black = keep.

See also: NB01/02 for the base architecture & objective, NB05 for txt2img/img2img inference, and NB06/08/09 for fine-tuning recipes (the same conv_in widening trick powers InstructPix2Pix in NB06).