EveryDream2trainer/utils/huggingface_downloader.py

import logging
import os
from typing import Optional, Tuple

import huggingface_hub
from diffusers import StableDiffusionPipeline

from utils.unet_utils import get_attn_yaml


def try_download_model_from_hf(repo_id: str) -> Tuple[StableDiffusionPipeline, str, bool, str] | None:
    """
    Attempts to download files from the following subfolders under the given repo id:
    "text_encoder", "vae", "unet", "scheduler", "tokenizer".
    :param repo_id The repository id of the model on huggingface, such as 'stabilityai/stable-diffusion-2-1' which
                    corresponds to `https://huggingface.co/stabilityai/stable-diffusion-2-1`.
    :param access_token Access token to use when fetching. If None, uses environment-saved token.
    :return: Root folder on disk to the downloaded files, or None if download failed.
    """

    access_token = os.environ.get('HF_API_TOKEN', None)
    if access_token is not None:
        huggingface_hub.login(access_token)

    # check if the model exists
    model_info = huggingface_hub.model_info(repo_id)
    if model_info is None:
        return None

    # load it to download it
    #pipe, cache_folder = StableDiffusionPipeline.from_pretrained(repo_id, return_cached_folder=True)
    cache_folder = StableDiffusionPipeline.download(repo_id)
    pipe = StableDiffusionPipeline.from_pretrained(repo_id)

    is_sd1_attn, yaml_path = get_attn_yaml(cache_folder)
    print(f"* HuggingFace Downloaded model from {repo_id} to {cache_folder}.")
    print(f"** Using attention yaml file: {yaml_path}, is_sd1_attn: {is_sd1_attn}.")

    return pipe, cache_folder, is_sd1_attn, yaml_path