utils.py

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

def load_unet_model(base, repo, ckpt, device="cpu"):
    """
    Load the UNet model from Hugging Face Hub.

    Args:
        base (str): Base model name.
        repo (str): Repository name.
        ckpt (str): Checkpoint filename.
        device (str): Device to load the model on.

    Returns:
        UNet2DConditionModel: Loaded UNet model.
    """
    from diffusers import UNet2DConditionModel
    unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(device, torch.float16)
    unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device))
    return unet