Principle:Huggingface Diffusers DreamBooth Export

Metadata

Knowledge Sources	Huggingface Diffusers LoRA Weight Loading
Domains	Generative_AI Model_Serialization Parameter_Efficient_Fine_Tuning
Last Updated	2026-02-13 00:00 GMT

Overview

A design principle for exporting trained DreamBooth LoRA weights from both the UNet and text encoder into a single serialized file. DreamBooth export handles multi-component LoRA state dict extraction, format conversion from PEFT to Diffusers conventions, and serialization using safetensors for safe and efficient model sharing.

Description

After DreamBooth LoRA training completes, the trained adapter weights must be extracted, converted, and saved in a format compatible with the Diffusers LoRA loading pipeline. The export process involves:

1. Model unwrapping -- When using distributed training with Accelerate, the models are wrapped in distributed containers. The unwrap_model() utility strips these wrappers to access the underlying model.
2. Precision casting -- The model is cast back to torch.float32 to ensure full-precision weight saving regardless of the training mixed-precision setting.
3. PEFT state dict extraction -- get_peft_model_state_dict(model) extracts only the LoRA adapter weights from the full model, filtering out frozen base model parameters.
4. Format conversion -- convert_state_dict_to_diffusers() converts PEFT-format state dict keys to the Diffusers naming convention, ensuring compatibility with load_lora_weights().
5. Multi-component saving -- StableDiffusionLoraLoaderMixin.save_lora_weights() combines UNet and text encoder LoRA state dicts into a single file, prefixing keys with component names.

The export produces a single pytorch_lora_weights.safetensors file that can be loaded into any compatible Stable Diffusion pipeline.

Usage

Export is typically the final step in the DreamBooth training script, performed only on the main process:

1. Ensure training is complete and all processes are synchronized.
2. Extract LoRA state dicts from both UNet and text encoder.
3. Save using StableDiffusionLoraLoaderMixin.save_lora_weights().
4. Optionally load the weights back into a fresh pipeline for validation inference.

Theoretical Basis

DreamBooth LoRA export implements a multi-component adapter serialization pattern:

EXPORT PIPELINE:
    1. UNWRAP:  model = unwrap_model(distributed_model)
    2. CAST:    model = model.to(torch.float32)
    3. EXTRACT: peft_state = get_peft_model_state_dict(model)
       -- Filters: only keys containing "lora_A", "lora_B", etc.
       -- Example key: "mid_block.attentions.0.transformer_blocks.0.attn1.to_q.lora_A.weight"
    4. CONVERT: diffusers_state = convert_state_dict_to_diffusers(peft_state)
       -- Renames PEFT keys to Diffusers convention
    5. SAVE:    save_lora_weights(
                   save_directory,
                   unet_lora_layers=unet_state,
                   text_encoder_lora_layers=te_state,  # None if text encoder not trained
               )

SERIALIZED FILE STRUCTURE:
    pytorch_lora_weights.safetensors:
        "unet.down_blocks.0.attentions.0....lora_A.weight": Tensor
        "unet.down_blocks.0.attentions.0....lora_B.weight": Tensor
        ...
        "text_encoder.text_model.encoder....lora_A.weight": Tensor  # if trained
        "text_encoder.text_model.encoder....lora_B.weight": Tensor  # if trained
        ...

Key theoretical properties:

• Adapter isolation -- get_peft_model_state_dict() extracts only the adapter weights, producing a state dict orders of magnitude smaller than the full model. A rank-4 LoRA for Stable Diffusion v1.5 produces a file of approximately 3--4 MB, compared to the ~3.5 GB full model.
• Key namespace separation -- UNet and text encoder weights are prefixed with their respective component names ("unet." and "text_encoder.") to enable unambiguous loading.
• Safe serialization -- The default safe_serialization=True uses the safetensors format, which prevents arbitrary code execution during loading (unlike pickle-based .bin files).

Related Pages

• Implementation:Huggingface_Diffusers_Save_Lora_Weights_DreamBooth