# Implementing LoRA Fine-Tuning with DreamBooth for Sana

> Learn to implement LoRA fine-tuning with DreamBooth for NVlabs Sana. Create custom adapters efficiently on consumer GPUs without altering base model parameters.

- Repository: [NVIDIA Research Projects/Sana](https://github.com/NVlabs/Sana)
- Tags: how-to-guide
- Published: 2026-05-19

---

**This guide explains how to create personalized LoRA adapters for the Sana diffusion model using DreamBooth techniques, enabling efficient fine-tuning on consumer GPUs while keeping the base model parameters frozen.**

The NVlabs/Sana repository provides a complete implementation for customizing the Sana text-to-image diffusion model through Low-Rank Adaptation (LoRA). By combining DreamBooth data augmentation with parameter-efficient fine-tuning, you can teach the model new visual concepts using minimal compute resources and storage.

## Dataset Preparation and Validation

The training script [`train_scripts/train_dreambooth_lora_sana.py`](https://github.com/NVlabs/Sana/blob/main/train_scripts/train_dreambooth_lora_sana.py) accepts data from either a Hugging Face dataset (`--dataset_name`) or a local image folder (`--instance_data_dir`). Validation logic at lines 79-99 ensures exactly one data source is specified.

The `DreamBoothDataset` class handles two concurrent image streams: **instance images** for your target concept and **class images** for prior preservation. It applies standard augmentation including resizing, center-cropping, random horizontal flips, and normalization to the `[-1, 1]` range (lines 24-84).

### Prior Preservation Configuration

When using `--with_prior_preservation`, the dataset automatically manages class images to prevent overfitting to your specific instances. Specify the class prompt via `--class_prompt` and the storage directory via `--class_data_dir`. The dataset repeats instance images according to the `--repeats` parameter to balance the sampling ratio.

## Injecting LoRA Adapters into the Sana Transformer

Only the transformer backbone receives LoRA weights; the VAE and text encoder remain frozen for computational efficiency. The adapter configuration uses PEFT's `LoraConfig` with default targeting of attention layers:

```python
target_modules = (
    [layer.strip() for layer in args.lora_layers.split(",")]
    if args.lora_layers is not None
    else ["to_k", "to_q", "to_v"]
)
transformer_lora_config = LoraConfig(
    r=args.rank,
    lora_alpha=args.rank,
    init_lora_weights="gaussian",
    target_modules=target_modules,
)
transformer.add_adapter(transformer_lora_config)

```

This injection occurs at lines 86-99 in [`train_scripts/train_dreambooth_lora_sana.py`](https://github.com/NVlabs/Sana/blob/main/train_scripts/train_dreambooth_lora_sana.py). The default rank is 4, but you can adjust this via the `--rank` argument to trade off between adapter capacity and file size.

## Training Loop and Optimization

The training loop runs for `args.max_train_steps` and supports text-embedding caching (`compute_text_embeddings`) and latent caching (`--cache_latents`) to accelerate training by avoiding redundant encodings (starting around line 311).

### Custom Save and Load Hooks

The script registers accelerator hooks to handle LoRA state dicts separately from the base model weights. The `save_model_hook` extracts trainable parameters using `get_peft_model_state_dict` and persists them via `SanaPipeline.save_lora_weights` (lines 106-124). During resumption, the `load_model_hook` restores the adapter using `SanaPipeline.lora_state_dict` and `set_peft_model_state_dict` (lines 124-143), ensuring seamless checkpoint recovery.

### Optimizer Selection

You can choose between **AdamW** (with optional 8-bit quantization via bitsandbytes) or **Prodigy** for adaptive learning rate optimization. Configure this through the script's argument parser to match your hardware constraints and convergence requirements.

## Running DreamBooth Training

Execute the full training pipeline with this command structure:

```bash
python -m train_scripts.train_dreambooth_lora_sana \
  --pretrained_model_name_or_path path/to/sana_base \
  --instance_data_dir ./my_photos \
  --instance_prompt "photo of a my-cat" \
  --output_dir ./sana-cat-lora \
  --rank 8 \
  --train_batch_size 4 \
  --num_train_epochs 3 \
  --learning_rate 5e-5 \
  --with_prior_preservation \
  --class_prompt "photo of a cat" \
  --class_data_dir ./class_images \
  --push_to_hub \
  --hub_model_id myusername/sana-cat-lora

```

Key parameters include `--rank` for adapter dimensionality and `--instance_prompt` which serves as the trigger word for subsequent generation. The `--with_prior_preservation` flag is essential for maintaining the model's generalization ability for the base class.

## Inference with LoRA Weights

Load the fine-tuned adapter using `SanaPipeline` from [`app/sana_pipeline.py`](https://github.com/NVlabs/Sana/blob/main/app/sana_pipeline.py):

```python
from app.sana_pipeline import SanaPipeline
import torch

# Load base Sana model

pipe = SanaPipeline().from_pretrained("output/Sana_D20/SANA.pth")

# Attach the LoRA weights produced by the training script

pipe.load_lora_weights("./sana-cat-lora")

# Generate an image with the custom token

result = pipe(
    prompt="photo of a my-cat playing with a ball",
    guidance_scale=7.5,
    num_inference_steps=50,
    generator=torch.Generator().manual_seed(123),
)

result[0].save("cat_generated.png")

```

The pipeline automatically switches between **classifier-free guidance** and **Perturbed Attention Guidance (PAG)** based on the `pag_scale` and `attn_type` parameters (lines 66-89 in [`app/sana_pipeline.py`](https://github.com/NVlabs/Sana/blob/main/app/sana_pipeline.py)).

## Exporting and Sharing Models

### Automatic Model Card Generation

The `save_model_card` function (lines 66-94) generates a [`README.md`](https://github.com/NVlabs/Sana/blob/main/README.md) inside the output directory containing validation images, the trigger word, and ready-to-use code snippets. When using `--push_to_hub`, this card accompanies the LoRA weights to the Hugging Face Hub, enabling immediate community sharing and version control.

## Summary

- **Dataset Flexibility**: Use either Hugging Face datasets or local folders with automatic prior-preservation image handling to maintain model fidelity.
- **Efficient Architecture**: LoRA adapters attach only to transformer attention layers (`to_q`, `to_k`, `to_v` by default), leaving VAE and text encoder weights frozen and reducing memory overhead.
- **Robust Checkpointing**: Custom hooks in [`train_dreambooth_lora_sana.py`](https://github.com/NVlabs/Sana/blob/main/train_dreambooth_lora_sana.py) handle conversion between PEFT state dicts and the pipeline's expected format via `SanaPipeline` utilities.
- **Seamless Inference**: The `SanaPipeline` class loads LoRA weights via `load_lora_weights()` without requiring base model reinitialization, supporting both classifier-free and PAG guidance modes.
- **Hub Integration**: Built-in model card generation and `--push_to_hub` support streamline distribution of custom adapters through the Hugging Face ecosystem.

## Frequently Asked Questions

### What is the default LoRA rank for Sana DreamBooth training?

The default rank is **4**, but you can specify any dimensionality using the `--rank` argument. Higher ranks increase the adapter's capacity to learn complex features but require more GPU memory and produce larger checkpoint files.

### Which layers of the Sana transformer receive LoRA adapters?

By default, adapters attach to the query, key, and value projection layers (`to_q`, `to_k`, `to_v`). You can customize the target modules via the `--lora_layers` parameter to experiment with different fine-tuning strategies or reduce the parameter count further.

### How do I resume training from a checkpoint?

The script handles resumption automatically through custom accelerator hooks. The `load_model_hook` restores adapter weights using `SanaPipeline.lora_state_dict` and `set_peft_model_state_dict`, ensuring training continues seamlessly from the exact optimization state where it stopped.

### Can I use 8-bit Adam for memory-efficient training?

Yes, the script supports 8-bit AdamW optimization via the bitsandbytes library. Specify this optimizer choice through the command-line arguments to significantly reduce VRAM requirements during training, enabling fine-tuning on consumer-grade GPUs with limited memory.