Implementation:Huggingface Diffusers Save Lora Weights DreamBooth

Metadata

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

Overview

The concrete API for saving DreamBooth-trained LoRA adapter weights using StableDiffusionLoraLoaderMixin.save_lora_weights(). This implementation extracts PEFT adapter state dicts from both the UNet and text encoder, converts them to Diffusers format, and serializes them into a single safetensors file.

Description

The export flow in the DreamBooth LoRA training script performs three stages:

Stage 1: State dict extraction

unet = unwrap_model(unet)
unet = unet.to(torch.float32)
unet_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))

if args.train_text_encoder:
    text_encoder = unwrap_model(text_encoder)
    text_encoder_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(text_encoder))
else:
    text_encoder_state_dict = None

Stage 2: Multi-component saving

The save_lora_weights() classmethod merges both state dicts under their respective component namespaces and delegates to _save_lora_weights() for the actual file I/O.

Stage 3: Validation loading

After saving, the script loads the weights back into a fresh pipeline to run validation inference:

pipeline = DiffusionPipeline.from_pretrained(model_path, ...)
pipeline.load_lora_weights(output_dir, weight_name="pytorch_lora_weights.safetensors")

Usage

from diffusers.loaders import StableDiffusionLoraLoaderMixin
from diffusers.utils import convert_state_dict_to_diffusers
from peft.utils import get_peft_model_state_dict

# Extract and convert state dicts
unet_state = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
te_state = convert_state_dict_to_diffusers(get_peft_model_state_dict(text_encoder))

# Save both components
StableDiffusionLoraLoaderMixin.save_lora_weights(
    save_directory="./output/dreambooth-lora",
    unet_lora_layers=unet_state,
    text_encoder_lora_layers=te_state,
)

Code Reference

Source Location

• Repository: huggingface/diffusers
• File: examples/dreambooth/train_dreambooth_lora.py (lines 1401--1419)
• File: src/diffusers/loaders/lora_pipeline.py (lines 473--534, save_lora_weights definition)

Signature

# StableDiffusionLoraLoaderMixin.save_lora_weights
# (from src/diffusers/loaders/lora_pipeline.py)

@classmethod
def save_lora_weights(
    cls,
    save_directory: str | os.PathLike,
    unet_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
    text_encoder_lora_layers: dict[str, torch.nn.Module] = None,
    is_main_process: bool = True,
    weight_name: str = None,
    save_function: Callable = None,
    safe_serialization: bool = True,
    unet_lora_adapter_metadata=None,
    text_encoder_lora_adapter_metadata=None,
):
    """
    Save the LoRA parameters corresponding to the UNet and text encoder.

    Arguments:
        save_directory: Directory to save LoRA parameters to.
        unet_lora_layers: State dict of the LoRA layers corresponding to the `unet`.
        text_encoder_lora_layers: State dict of the LoRA layers corresponding
            to the `text_encoder`.
        is_main_process: Whether the process calling this is the main process.
        weight_name: Override filename for the saved weights.
        save_function: Custom save function (e.g., for distributed settings).
        safe_serialization: Whether to use safetensors format (default True).
    """
    lora_layers = {}
    if unet_lora_layers:
        lora_layers[cls.unet_name] = unet_lora_layers
    if text_encoder_lora_layers:
        lora_layers[cls.text_encoder_name] = text_encoder_lora_layers

    if not lora_layers:
        raise ValueError("You must pass at least one of `unet_lora_layers` "
                         "or `text_encoder_lora_layers`.")

    cls._save_lora_weights(
        save_directory=save_directory,
        lora_layers=lora_layers,
        is_main_process=is_main_process,
        weight_name=weight_name,
        save_function=save_function,
        safe_serialization=safe_serialization,
    )

Import

from diffusers.loaders import StableDiffusionLoraLoaderMixin
from diffusers.utils import convert_state_dict_to_diffusers
from peft.utils import get_peft_model_state_dict

I/O Contract

Inputs

Input Contract

Name	Type	Default	Description
save_directory	os.PathLike	(required)	Directory to save the LoRA weight file. Created if it does not exist.
unet_lora_layers	dict[str, Tensor]	None	State dict containing UNet LoRA adapter weights in Diffusers format.
text_encoder_lora_layers	dict[str, Tensor]	None	State dict containing text encoder LoRA adapter weights in Diffusers format. None when text encoder was not trained.
is_main_process	bool	True	Set to True only on the main process to avoid race conditions in distributed training.
weight_name	str	None	Custom filename for the weight file. Defaults to "pytorch_lora_weights.safetensors".
safe_serialization	bool	True	Whether to use safetensors format (recommended for security).

Outputs

Output Contract

Name	Type	Description
Weight file	File on disk	A single pytorch_lora_weights.safetensors file in save_directory containing both UNet and text encoder LoRA weights with namespaced keys.

Usage Examples

Example 1: Saving UNet-Only LoRA Weights

from diffusers.loaders import StableDiffusionLoraLoaderMixin
from diffusers.utils import convert_state_dict_to_diffusers
from peft.utils import get_peft_model_state_dict

unet_state = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))

StableDiffusionLoraLoaderMixin.save_lora_weights(
    save_directory="./output/dreambooth-lora",
    unet_lora_layers=unet_state,
    text_encoder_lora_layers=None,
)
# Creates: ./output/dreambooth-lora/pytorch_lora_weights.safetensors

Example 2: Saving Dual LoRA Weights and Loading for Inference

from diffusers import DiffusionPipeline
from diffusers.loaders import StableDiffusionLoraLoaderMixin
from diffusers.utils import convert_state_dict_to_diffusers
from peft.utils import get_peft_model_state_dict

# Extract and save
unet_state = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
te_state = convert_state_dict_to_diffusers(get_peft_model_state_dict(text_encoder))

StableDiffusionLoraLoaderMixin.save_lora_weights(
    save_directory="./output/dreambooth-lora",
    unet_lora_layers=unet_state,
    text_encoder_lora_layers=te_state,
)

# Load back for inference
pipeline = DiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
)
pipeline.load_lora_weights(
    "./output/dreambooth-lora",
    weight_name="pytorch_lora_weights.safetensors",
)
pipeline.to("cuda")

image = pipeline("a photo of sks dog in a park", num_inference_steps=25).images[0]
image.save("sks_dog_park.png")

Related Pages

• Principle:Huggingface_Diffusers_DreamBooth_Export