InvokeAI/invokeai/backend/install/model_install_backend.py

"""
Utility (backend) functions used by model_install.py
"""
import os
import re
import shutil
import warnings
from dataclasses import dataclass, field
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Callable, Dict, List, Optional, Set, Union

import requests
import torch
from diffusers import DiffusionPipeline
from diffusers import logging as dlogging
from huggingface_hub import HfApi, HfFolder, hf_hub_url
from omegaconf import OmegaConf
from tqdm import tqdm

import invokeai.configs as configs
from invokeai.app.services.config import InvokeAIAppConfig
from invokeai.backend.model_management import AddModelResult, BaseModelType, ModelManager, ModelType, ModelVariantType
from invokeai.backend.model_management.model_probe import ModelProbe, ModelProbeInfo, SchedulerPredictionType
from invokeai.backend.util import download_with_resume
from invokeai.backend.util.devices import choose_torch_device, torch_dtype

from ..util.logging import InvokeAILogger

warnings.filterwarnings("ignore")

# --------------------------globals-----------------------
config = InvokeAIAppConfig.get_config()
logger = InvokeAILogger.get_logger(name="InvokeAI")

# the initial "configs" dir is now bundled in the `invokeai.configs` package
Dataset_path = Path(configs.__path__[0]) / "INITIAL_MODELS.yaml"

Config_preamble = """
# This file describes the alternative machine learning models
# available to InvokeAI script.
#
# To add a new model, follow the examples below. Each
# model requires a model config file, a weights file,
# and the width and height of the images it
# was trained on.
"""

LEGACY_CONFIGS = {
    BaseModelType.StableDiffusion1: {
        ModelVariantType.Normal: {
            SchedulerPredictionType.Epsilon: "v1-inference.yaml",
            SchedulerPredictionType.VPrediction: "v1-inference-v.yaml",
        },
        ModelVariantType.Inpaint: {
            SchedulerPredictionType.Epsilon: "v1-inpainting-inference.yaml",
            SchedulerPredictionType.VPrediction: "v1-inpainting-inference-v.yaml",
        },
    },
    BaseModelType.StableDiffusion2: {
        ModelVariantType.Normal: {
            SchedulerPredictionType.Epsilon: "v2-inference.yaml",
            SchedulerPredictionType.VPrediction: "v2-inference-v.yaml",
        },
        ModelVariantType.Inpaint: {
            SchedulerPredictionType.Epsilon: "v2-inpainting-inference.yaml",
            SchedulerPredictionType.VPrediction: "v2-inpainting-inference-v.yaml",
        },
    },
    BaseModelType.StableDiffusionXL: {
        ModelVariantType.Normal: "sd_xl_base.yaml",
    },
    BaseModelType.StableDiffusionXLRefiner: {
        ModelVariantType.Normal: "sd_xl_refiner.yaml",
    },
}


@dataclass
class InstallSelections:
    install_models: List[str] = field(default_factory=list)
    remove_models: List[str] = field(default_factory=list)


@dataclass
class ModelLoadInfo:
    name: str
    model_type: ModelType
    base_type: BaseModelType
    path: Optional[Path] = None
    repo_id: Optional[str] = None
    subfolder: Optional[str] = None
    description: str = ""
    installed: bool = False
    recommended: bool = False
    default: bool = False
    requires: Optional[List[str]] = field(default_factory=list)


class ModelInstall(object):
    def __init__(
        self,
        config: InvokeAIAppConfig,
        prediction_type_helper: Optional[Callable[[Path], SchedulerPredictionType]] = None,
        model_manager: Optional[ModelManager] = None,
        access_token: Optional[str] = None,
    ):
        self.config = config
        self.mgr = model_manager or ModelManager(config.model_conf_path)
        self.datasets = OmegaConf.load(Dataset_path)
        self.prediction_helper = prediction_type_helper
        self.access_token = access_token or HfFolder.get_token()
        self.reverse_paths = self._reverse_paths(self.datasets)

    def all_models(self) -> Dict[str, ModelLoadInfo]:
        """
        Return dict of model_key=>ModelLoadInfo objects.
        This method consolidates and simplifies the entries in both
        models.yaml and INITIAL_MODELS.yaml so that they can
        be treated uniformly. It also sorts the models alphabetically
        by their name, to improve the display somewhat.
        """
        model_dict = dict()

        # first populate with the entries in INITIAL_MODELS.yaml
        for key, value in self.datasets.items():
            name, base, model_type = ModelManager.parse_key(key)
            value["name"] = name
            value["base_type"] = base
            value["model_type"] = model_type
            model_info = ModelLoadInfo(**value)
            if model_info.subfolder and model_info.repo_id:
                model_info.repo_id += f":{model_info.subfolder}"
            model_dict[key] = model_info

        # supplement with entries in models.yaml
        installed_models = [x for x in self.mgr.list_models()]

        for md in installed_models:
            base = md["base_model"]
            model_type = md["model_type"]
            name = md["model_name"]
            key = ModelManager.create_key(name, base, model_type)
            if key in model_dict:
                model_dict[key].installed = True
            else:
                model_dict[key] = ModelLoadInfo(
                    name=name,
                    base_type=base,
                    model_type=model_type,
                    path=value.get("path"),
                    installed=True,
                )
        return {x: model_dict[x] for x in sorted(model_dict.keys(), key=lambda y: model_dict[y].name.lower())}

    def _is_autoloaded(self, model_info: dict) -> bool:
        path = model_info.get("path")
        if not path:
            return False
        for autodir in ["autoimport_dir", "lora_dir", "embedding_dir", "controlnet_dir"]:
            if autodir_path := getattr(self.config, autodir):
                autodir_path = self.config.root_path / autodir_path
                if Path(path).is_relative_to(autodir_path):
                    return True
        return False

    def list_models(self, model_type):
        installed = self.mgr.list_models(model_type=model_type)
        print()
        print(f"Installed models of type `{model_type}`:")
        print(f"{'Model Key':50} Model Path")
        for i in installed:
            print(f"{'/'.join([i['base_model'],i['model_type'],i['model_name']]):50} {i['path']}")
        print()

    # logic here a little reversed to maintain backward compatibility
    def starter_models(self, all_models: bool = False) -> Set[str]:
        models = set()
        for key, value in self.datasets.items():
            name, base, model_type = ModelManager.parse_key(key)
            if all_models or model_type in [ModelType.Main, ModelType.Vae]:
                models.add(key)
        return models

    def recommended_models(self) -> Set[str]:
        starters = self.starter_models(all_models=True)
        return set([x for x in starters if self.datasets[x].get("recommended", False)])

    def default_model(self) -> str:
        starters = self.starter_models()
        defaults = [x for x in starters if self.datasets[x].get("default", False)]
        return defaults[0]

    def install(self, selections: InstallSelections):
        verbosity = dlogging.get_verbosity()  # quench NSFW nags
        dlogging.set_verbosity_error()

        job = 1
        jobs = len(selections.remove_models) + len(selections.install_models)

        # remove requested models
        for key in selections.remove_models:
            name, base, mtype = self.mgr.parse_key(key)
            logger.info(f"Deleting {mtype} model {name} [{job}/{jobs}]")
            try:
                self.mgr.del_model(name, base, mtype)
            except FileNotFoundError as e:
                logger.warning(e)
            job += 1

        # add requested models
        self._remove_installed(selections.install_models)
        self._add_required_models(selections.install_models)
        for path in selections.install_models:
            logger.info(f"Installing {path} [{job}/{jobs}]")
            try:
                self.heuristic_import(path)
            except (ValueError, KeyError) as e:
                logger.error(str(e))
            job += 1

        dlogging.set_verbosity(verbosity)
        self.mgr.commit()

    def heuristic_import(
        self,
        model_path_id_or_url: Union[str, Path],
        models_installed: Set[Path] = None,
    ) -> Dict[str, AddModelResult]:
        """
        :param model_path_id_or_url: A Path to a local model to import, or a string representing its repo_id or URL
        :param models_installed: Set of installed models, used for recursive invocation
        Returns a set of dict objects corresponding to newly-created stanzas in models.yaml.
        """

        if not models_installed:
            models_installed = dict()

        model_path_id_or_url = str(model_path_id_or_url).strip("\"' ")

        # A little hack to allow nested routines to retrieve info on the requested ID
        self.current_id = model_path_id_or_url
        path = Path(model_path_id_or_url)

        # fix relative paths
        if path.exists() and not path.is_absolute():
            path = path.absolute()  # make relative to current WD

        # checkpoint file, or similar
        if path.is_file():
            models_installed.update({str(path): self._install_path(path)})

        # folders style or similar
        elif path.is_dir() and any(
            [
                (path / x).exists()
                for x in {
                    "config.json",
                    "model_index.json",
                    "learned_embeds.bin",
                    "pytorch_lora_weights.bin",
                    "pytorch_lora_weights.safetensors",
                }
            ]
        ):
            models_installed.update({str(model_path_id_or_url): self._install_path(path)})

        # recursive scan
        elif path.is_dir():
            for child in path.iterdir():
                self.heuristic_import(child, models_installed=models_installed)

        # huggingface repo
        elif len(str(model_path_id_or_url).split("/")) == 2:
            models_installed.update({str(model_path_id_or_url): self._install_repo(str(model_path_id_or_url))})

        # a URL
        elif str(model_path_id_or_url).startswith(("http:", "https:", "ftp:")):
            models_installed.update({str(model_path_id_or_url): self._install_url(model_path_id_or_url)})

        else:
            raise KeyError(f"{str(model_path_id_or_url)} is not recognized as a local path, repo ID or URL. Skipping")

        return models_installed

    def _remove_installed(self, model_list: List[str]):
        all_models = self.all_models()
        for path in model_list:
            key = self.reverse_paths.get(path)
            if key and all_models[key].installed:
                logger.warning(f"{path} already installed. Skipping.")
                model_list.remove(path)

    def _add_required_models(self, model_list: List[str]):
        additional_models = []
        all_models = self.all_models()
        for path in model_list:
            if not (key := self.reverse_paths.get(path)):
                continue
            for requirement in all_models[key].requires:
                requirement_key = self.reverse_paths.get(requirement)
                if not all_models[requirement_key].installed:
                    additional_models.append(requirement)
        model_list.extend(additional_models)

    # install a model from a local path. The optional info parameter is there to prevent
    # the model from being probed twice in the event that it has already been probed.
    def _install_path(self, path: Path, info: ModelProbeInfo = None) -> AddModelResult:
        info = info or ModelProbe().heuristic_probe(path, self.prediction_helper)
        if not info:
            logger.warning(f"Unable to parse format of {path}")
            return None
        model_name = path.stem if path.is_file() else path.name
        if self.mgr.model_exists(model_name, info.base_type, info.model_type):
            raise ValueError(f'A model named "{model_name}" is already installed.')
        attributes = self._make_attributes(path, info)
        return self.mgr.add_model(
            model_name=model_name,
            base_model=info.base_type,
            model_type=info.model_type,
            model_attributes=attributes,
        )

    def _install_url(self, url: str) -> AddModelResult:
        with TemporaryDirectory(dir=self.config.models_path) as staging:
            location = download_with_resume(url, Path(staging))
            if not location:
                logger.error(f"Unable to download {url}. Skipping.")
            info = ModelProbe().heuristic_probe(location, self.prediction_helper)
            dest = self.config.models_path / info.base_type.value / info.model_type.value / location.name
            dest.parent.mkdir(parents=True, exist_ok=True)
            models_path = shutil.move(location, dest)

        # staged version will be garbage-collected at this time
        return self._install_path(Path(models_path), info)

    def _install_repo(self, repo_id: str) -> AddModelResult:
        # hack to recover models stored in subfolders --
        # Required to get the "v2" model of monster-labs/control_v1p_sd15_qrcode_monster
        subfolder = None
        if match := re.match(r"^([^/]+/[^/]+):(\w+)$", repo_id):
            repo_id = match.group(1)
            subfolder = match.group(2)

        hinfo = HfApi().model_info(repo_id)

        # we try to figure out how to download this most economically
        # list all the files in the repo
        files = [x.rfilename for x in hinfo.siblings]
        if subfolder:
            files = [x for x in files if x.startswith(f"{subfolder}/")]
        prefix = f"{subfolder}/" if subfolder else ""

        location = None

        with TemporaryDirectory(dir=self.config.models_path) as staging:
            staging = Path(staging)
            if f"{prefix}model_index.json" in files:
                location = self._download_hf_pipeline(repo_id, staging, subfolder=subfolder)  # pipeline
            elif f"{prefix}unet/model.onnx" in files:
                location = self._download_hf_model(repo_id, files, staging)
            else:
                for suffix in ["safetensors", "bin"]:
                    if f"{prefix}pytorch_lora_weights.{suffix}" in files:
                        location = self._download_hf_model(
                            repo_id, [f"pytorch_lora_weights.{suffix}"], staging, subfolder=subfolder
                        )  # LoRA
                        break
                    elif (
                        self.config.precision == "float16" and f"{prefix}diffusion_pytorch_model.fp16.{suffix}" in files
                    ):  # vae, controlnet or some other standalone
                        files = ["config.json", f"diffusion_pytorch_model.fp16.{suffix}"]
                        location = self._download_hf_model(repo_id, files, staging, subfolder=subfolder)
                        break
                    elif f"{prefix}diffusion_pytorch_model.{suffix}" in files:
                        files = ["config.json", f"diffusion_pytorch_model.{suffix}"]
                        location = self._download_hf_model(repo_id, files, staging, subfolder=subfolder)
                        break
                    elif f"{prefix}learned_embeds.{suffix}" in files:
                        location = self._download_hf_model(
                            repo_id, [f"learned_embeds.{suffix}"], staging, subfolder=subfolder
                        )
                        break
                    elif (
                        f"{prefix}image_encoder.txt" in files and f"{prefix}ip_adapter.{suffix}" in files
                    ):  # IP-Adapter
                        files = ["image_encoder.txt", f"ip_adapter.{suffix}"]
                        location = self._download_hf_model(repo_id, files, staging, subfolder=subfolder)
                        break
                    elif f"{prefix}model.{suffix}" in files and f"{prefix}config.json" in files:
                        # This elif-condition is pretty fragile, but it is intended to handle CLIP Vision models hosted
                        # by InvokeAI for use with IP-Adapters.
                        files = ["config.json", f"model.{suffix}"]
                        location = self._download_hf_model(repo_id, files, staging, subfolder=subfolder)
                        break
            if not location:
                logger.warning(f"Could not determine type of repo {repo_id}. Skipping install.")
                return {}

            info = ModelProbe().heuristic_probe(location, self.prediction_helper)
            if not info:
                logger.warning(f"Could not probe {location}. Skipping install.")
                return {}
            dest = (
                self.config.models_path
                / info.base_type.value
                / info.model_type.value
                / self._get_model_name(repo_id, location)
            )
            if dest.exists():
                shutil.rmtree(dest)
            shutil.copytree(location, dest)
            return self._install_path(dest, info)

    def _get_model_name(self, path_name: str, location: Path) -> str:
        """
        Calculate a name for the model - primitive implementation.
        """
        if key := self.reverse_paths.get(path_name):
            (name, base, mtype) = ModelManager.parse_key(key)
            return name
        elif location.is_dir():
            return location.name
        else:
            return location.stem

    def _make_attributes(self, path: Path, info: ModelProbeInfo) -> dict:
        model_name = path.name if path.is_dir() else path.stem
        description = f"{info.base_type.value} {info.model_type.value} model {model_name}"
        if key := self.reverse_paths.get(self.current_id):
            if key in self.datasets:
                description = self.datasets[key].get("description") or description

        rel_path = self.relative_to_root(path, self.config.models_path)

        attributes = dict(
            path=str(rel_path),
            description=str(description),
            model_format=info.format,
        )
        legacy_conf = None
        if info.model_type == ModelType.Main or info.model_type == ModelType.ONNX:
            attributes.update(
                dict(
                    variant=info.variant_type,
                )
            )
            if info.format == "checkpoint":
                try:
                    possible_conf = path.with_suffix(".yaml")
                    if possible_conf.exists():
                        legacy_conf = str(self.relative_to_root(possible_conf))
                    elif info.base_type in [BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2]:
                        legacy_conf = Path(
                            self.config.legacy_conf_dir,
                            LEGACY_CONFIGS[info.base_type][info.variant_type][info.prediction_type],
                        )
                    else:
                        legacy_conf = Path(
                            self.config.legacy_conf_dir, LEGACY_CONFIGS[info.base_type][info.variant_type]
                        )
                except KeyError:
                    legacy_conf = Path(self.config.legacy_conf_dir, "v1-inference.yaml")  # best guess

        if info.model_type == ModelType.ControlNet and info.format == "checkpoint":
            possible_conf = path.with_suffix(".yaml")
            if possible_conf.exists():
                legacy_conf = str(self.relative_to_root(possible_conf))
            else:
                legacy_conf = Path(
                    self.config.root_path,
                    "configs/controlnet",
                    ("cldm_v15.yaml" if info.base_type == BaseModelType("sd-1") else "cldm_v21.yaml"),
                )

        if legacy_conf:
            attributes.update(dict(config=str(legacy_conf)))
        return attributes

    def relative_to_root(self, path: Path, root: Optional[Path] = None) -> Path:
        root = root or self.config.root_path
        if path.is_relative_to(root):
            return path.relative_to(root)
        else:
            return path

    def _download_hf_pipeline(self, repo_id: str, staging: Path, subfolder: str = None) -> Path:
        """
        Retrieve a StableDiffusion model from cache or remote and then
        does a save_pretrained() to the indicated staging area.
        """
        _, name = repo_id.split("/")
        precision = torch_dtype(choose_torch_device())
        variants = ["fp16", None] if precision == torch.float16 else [None, "fp16"]

        model = None
        for variant in variants:
            try:
                model = DiffusionPipeline.from_pretrained(
                    repo_id,
                    variant=variant,
                    torch_dtype=precision,
                    safety_checker=None,
                    subfolder=subfolder,
                )
            except Exception as e:  # most errors are due to fp16 not being present. Fix this to catch other errors
                if "fp16" not in str(e):
                    print(e)

            if model:
                break

        if not model:
            logger.error(f"Diffusers model {repo_id} could not be downloaded. Skipping.")
            return None
        model.save_pretrained(staging / name, safe_serialization=True)
        return staging / name

    def _download_hf_model(self, repo_id: str, files: List[str], staging: Path, subfolder: None) -> Path:
        _, name = repo_id.split("/")
        location = staging / name
        paths = list()
        for filename in files:
            filePath = Path(filename)
            p = hf_download_with_resume(
                repo_id,
                model_dir=location / filePath.parent,
                model_name=filePath.name,
                access_token=self.access_token,
                subfolder=filePath.parent / subfolder if subfolder else filePath.parent,
            )
            if p:
                paths.append(p)
            else:
                logger.warning(f"Could not download {filename} from {repo_id}.")

        return location if len(paths) > 0 else None

    @classmethod
    def _reverse_paths(cls, datasets) -> dict:
        """
        Reverse mapping from repo_id/path to destination name.
        """
        return {v.get("path") or v.get("repo_id"): k for k, v in datasets.items()}


# -------------------------------------
def yes_or_no(prompt: str, default_yes=True):
    default = "y" if default_yes else "n"
    response = input(f"{prompt} [{default}] ") or default
    if default_yes:
        return response[0] not in ("n", "N")
    else:
        return response[0] in ("y", "Y")


# ---------------------------------------------
def hf_download_from_pretrained(model_class: object, model_name: str, destination: Path, **kwargs):
    logger = InvokeAILogger.get_logger("InvokeAI")
    logger.addFilter(lambda x: "fp16 is not a valid" not in x.getMessage())

    model = model_class.from_pretrained(
        model_name,
        resume_download=True,
        **kwargs,
    )
    model.save_pretrained(destination, safe_serialization=True)
    return destination


# ---------------------------------------------
def hf_download_with_resume(
    repo_id: str,
    model_dir: str,
    model_name: str,
    model_dest: Path = None,
    access_token: str = None,
    subfolder: str = None,
) -> Path:
    model_dest = model_dest or Path(os.path.join(model_dir, model_name))
    os.makedirs(model_dir, exist_ok=True)

    url = hf_hub_url(repo_id, model_name, subfolder=subfolder)

    header = {"Authorization": f"Bearer {access_token}"} if access_token else {}
    open_mode = "wb"
    exist_size = 0

    if os.path.exists(model_dest):
        exist_size = os.path.getsize(model_dest)
        header["Range"] = f"bytes={exist_size}-"
        open_mode = "ab"

    resp = requests.get(url, headers=header, stream=True)
    total = int(resp.headers.get("content-length", 0))

    if resp.status_code == 416:  # "range not satisfiable", which means nothing to return
        logger.info(f"{model_name}: complete file found. Skipping.")
        return model_dest
    elif resp.status_code == 404:
        logger.warning("File not found")
        return None
    elif resp.status_code != 200:
        logger.warning(f"{model_name}: {resp.reason}")
    elif exist_size > 0:
        logger.info(f"{model_name}: partial file found. Resuming...")
    else:
        logger.info(f"{model_name}: Downloading...")

    try:
        with (
            open(model_dest, open_mode) as file,
            tqdm(
                desc=model_name,
                initial=exist_size,
                total=total + exist_size,
                unit="iB",
                unit_scale=True,
                unit_divisor=1000,
            ) as bar,
        ):
            for data in resp.iter_content(chunk_size=1024):
                size = file.write(data)
                bar.update(size)
    except Exception as e:
        logger.error(f"An error occurred while downloading {model_name}: {str(e)}")
        return None
    return model_dest