Implementation:Huggingface Diffusers ModelMixin From Pretrained Frozen

Metadata

Knowledge Sources	Huggingface Diffusers Diffusers ModelMixin API
Domains	Generative_AI Transfer_Learning Model_Loading
Last Updated	2026-02-13 00:00 GMT

Overview

The concrete API for loading pretrained diffusion model components and freezing their parameters. This combines ModelMixin.from_pretrained() for weight loading with model.requires_grad_(False) for gradient disabling, preparing the models for LoRA adapter injection.

Description

The DreamBooth LoRA training script loads four pretrained components in sequence:

1. DDPMScheduler -- Loaded from the "scheduler" subfolder. No trainable parameters.
2. Text encoder -- Loaded via the dynamically resolved text_encoder_cls.from_pretrained() from the "text_encoder" subfolder.
3. AutoencoderKL -- Loaded from the "vae" subfolder. Wrapped in a try/except to handle models without a VAE (e.g., IF models).
4. UNet2DConditionModel -- Loaded from the "unet" subfolder.

After loading, all model parameters are frozen:

if vae is not None:
    vae.requires_grad_(False)
text_encoder.requires_grad_(False)
unet.requires_grad_(False)

The models are then cast to the appropriate inference dtype and moved to the accelerator device.

Usage

This pattern is used in the initialization phase of DreamBooth training, before LoRA adapters are added:

from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel

# Load models
noise_scheduler = DDPMScheduler.from_pretrained(model_path, subfolder="scheduler")
text_encoder = text_encoder_cls.from_pretrained(model_path, subfolder="text_encoder")
vae = AutoencoderKL.from_pretrained(model_path, subfolder="vae")
unet = UNet2DConditionModel.from_pretrained(model_path, subfolder="unet")

# Freeze all parameters
vae.requires_grad_(False)
text_encoder.requires_grad_(False)
unet.requires_grad_(False)

# Cast to inference dtype and move to device
unet.to(device, dtype=weight_dtype)
vae.to(device, dtype=weight_dtype)
text_encoder.to(device, dtype=weight_dtype)

Code Reference

Source Location

• Repository: huggingface/diffusers
• File: examples/dreambooth/train_dreambooth_lora.py (lines 880--915)
• File: src/diffusers/models/modeling_utils.py (lines 836--1370, ModelMixin.from_pretrained definition)

Signature

# ModelMixin.from_pretrained (from src/diffusers/models/modeling_utils.py)
@classmethod
def from_pretrained(
    cls,
    pretrained_model_name_or_path: str | os.PathLike | None,
    **kwargs,
) -> Self:
    """
    Instantiate a pretrained PyTorch model from a pretrained model configuration.

    Parameters:
        pretrained_model_name_or_path: Model id (e.g. 'google/ddpm-celebahq-256')
            or path to a directory containing model weights.
        cache_dir: Path to cache directory.
        torch_dtype: Override the default torch.dtype.
        force_download: Force re-download of model weights.
        subfolder: Subfolder location within a larger model repository.
        revision: Specific model version (branch, tag, or commit id).
        variant: Variant of model files (e.g., 'fp16').
        device_map: Device placement map for model submodules.
    """
    ...

# Freezing pattern (from train_dreambooth_lora.py)
unet.requires_grad_(False)       # torch.nn.Module method
text_encoder.requires_grad_(False)
vae.requires_grad_(False)

Import

from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
from transformers import AutoTokenizer

I/O Contract

Inputs

Input Contract

Name	Type	Description
pretrained_model_name_or_path	os.PathLike	Hub model id or local path to the pretrained model directory.
subfolder	str	Subfolder within the model repo: "unet", "vae", "text_encoder", or "scheduler".
revision	str	Model version identifier (branch, tag, commit).
variant	str	Model file variant, e.g., "fp16" for half-precision checkpoints.
torch_dtype	torch.dtype	Override dtype for loading (e.g., torch.float16).

Outputs

Output Contract

Name	Type	Description
return	Model instance (Self)	The instantiated model with pretrained weights loaded, set to evaluation mode.

Usage Examples

Example 1: Loading and Freezing the Full Pipeline

from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
import torch

model_id = "runwayml/stable-diffusion-v1-5"
device = "cuda"
weight_dtype = torch.float16

# Load pretrained components
noise_scheduler = DDPMScheduler.from_pretrained(model_id, subfolder="scheduler")
vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", variant="fp16")
unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet", variant="fp16")

# Freeze all parameters (no gradients stored)
vae.requires_grad_(False)
unet.requires_grad_(False)

# Move to device with inference precision
vae.to(device, dtype=weight_dtype)
unet.to(device, dtype=weight_dtype)

# Verify: all parameters frozen
assert all(not p.requires_grad for p in unet.parameters())
assert all(not p.requires_grad for p in vae.parameters())

Example 2: Handling Models Without a VAE

The IF (DeepFloyd) model does not include a VAE, so the script gracefully handles the OSError.

try:
    vae = AutoencoderKL.from_pretrained(
        model_id, subfolder="vae", revision=revision, variant=variant
    )
except OSError:
    # IF does not have a VAE so let's just set it to None
    vae = None

if vae is not None:
    vae.requires_grad_(False)
    vae.to(device, dtype=weight_dtype)

Related Pages

• Principle:Huggingface_Diffusers_Model_Freezing