diff --git a/invokeai/backend/install/model_install_backend.py b/invokeai/backend/install/model_install_backend.py index f8bae1455b..82b5da6ebb 100644 --- a/invokeai/backend/install/model_install_backend.py +++ b/invokeai/backend/install/model_install_backend.py @@ -95,8 +95,6 @@ class ModelInstall(object): prediction_type_helper: Callable[[Path],SchedulerPredictionType]=None, access_token:str = None): self.config = config - with open('log.txt','w') as file: - print(config.model_conf_path,file=file) self.mgr = ModelManager(config.model_conf_path) self.datasets = OmegaConf.load(Dataset_path) self.prediction_helper = prediction_type_helper @@ -271,27 +269,36 @@ class ModelInstall(object): # 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] + location = None with TemporaryDirectory(dir=self.config.models_path) as staging: staging = Path(staging) if 'model_index.json' in files: location = self._download_hf_pipeline(repo_id, staging) # pipeline - - elif 'pytorch_lora_weights.bin' in files: - location = self._download_hf_model(repo_id, ['pytorch_lora_weights.bin'], staging) # LoRA - - elif self.config.precision=='float16' and 'diffusion_pytorch_model.fp16.safetensors' in files: # vae, controlnet or some other standalone - files = ['config.json', 'diffusion_pytorch_model.fp16.safetensors'] - location = self._download_hf_model(repo_id, files, staging) - - elif 'diffusion_pytorch_model.safetensors' in files: - files = ['config.json', 'diffusion_pytorch_model.safetensors'] - location = self._download_hf_model(repo_id, files, staging) - - elif 'learned_embeds.bin' in files: - location = self._download_hf_model(repo_id, ['learned_embeds.bin'], staging) - + else: + for suffix in ['safetensors','bin']: + if f'pytorch_lora_weights.{suffix}' in files: + location = self._download_hf_model(repo_id, ['pytorch_lora_weights.bin'], staging) # LoRA + break + elif self.config.precision=='float16' and f'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) + break + elif f'diffusion_pytorch_model.{suffix}' in files: + files = ['config.json', f'diffusion_pytorch_model.{suffix}'] + location = self._download_hf_model(repo_id, files, staging) + break + elif f'learned_embeds.{suffix}' in files: + location = self._download_hf_model(repo_id, [f'learned_embeds.suffix'], staging) + 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) diff --git a/invokeai/backend/model_management/model_install.py b/invokeai/backend/model_management/model_install.py deleted file mode 100644 index 64c52185f3..0000000000 --- a/invokeai/backend/model_management/model_install.py +++ /dev/null @@ -1,118 +0,0 @@ -""" -Routines for downloading and installing models. -""" -import json -import safetensors -import safetensors.torch -import shutil -import tempfile -import torch -import traceback -from dataclasses import dataclass -from diffusers import ModelMixin -from enum import Enum -from typing import Callable -from pathlib import Path - -import invokeai.backend.util.logging as logger -from invokeai.app.services.config import InvokeAIAppConfig -from . import ModelManager -from .models import BaseModelType, ModelType, VariantType -from .model_probe import ModelProbe, ModelVariantInfo -from .model_cache import SilenceWarnings - -class ModelInstall(object): - ''' - This class is able to download and install several different kinds of - InvokeAI models. The helper function, if provided, is called on to distinguish - between v2-base and v2-768 stable diffusion pipelines. This usually involves - asking the user to select the proper type, as there is no way of distinguishing - the two type of v2 file programmatically (as far as I know). - ''' - def __init__(self, - config: InvokeAIAppConfig, - model_base_helper: Callable[[Path],BaseModelType]=None, - clobber:bool = False - ): - ''' - :param config: InvokeAI configuration object - :param model_base_helper: A function call that accepts the Path to a checkpoint model and returns a ModelType enum - :param clobber: If true, models with colliding names will be overwritten - ''' - self.config = config - self.clogger = clobber - self.helper = model_base_helper - self.prober = ModelProbe() - - def install_checkpoint_file(self, checkpoint: Path)->dict: - ''' - Install the checkpoint file at path and return a - configuration entry that can be added to `models.yaml`. - Model checkpoints and VAEs will be converted into - diffusers before installation. Note that the model manager - does not hold entries for anything but diffusers pipelines, - and the configuration file stanzas returned from such models - can be safely ignored. - ''' - model_info = self.prober.probe(checkpoint, self.helper) - if not model_info: - raise ValueError(f"Unable to determine type of checkpoint file {checkpoint}") - - key = ModelManager.create_key( - model_name = checkpoint.stem, - base_model = model_info.base_type, - model_type = model_info.model_type, - ) - destination_path = self._dest_path(model_info) / checkpoint - destination_path.parent.mkdir(parents=True, exist_ok=True) - self._check_for_collision(destination_path) - stanza = { - key: dict( - name = checkpoint.stem, - description = f'{model_info.model_type} model {checkpoint.stem}', - base = model_info.base_model.value, - type = model_info.model_type.value, - variant = model_info.variant_type.value, - path = str(destination_path), - ) - } - - # non-pipeline; no conversion needed, just copy into right place - if model_info.model_type != ModelType.Pipeline: - shutil.copyfile(checkpoint, destination_path) - stanza[key].update({'format': 'checkpoint'}) - - # pipeline - conversion needed here - else: - destination_path = self._dest_path(model_info) / checkpoint.stem - config_file = self._pipeline_type_to_config_file(model_info.model_type) - - from .convert_ckpt_to_diffusers import convert_ckpt_to_diffusers - with SilenceWarnings: - convert_ckpt_to_diffusers( - checkpoint, - destination_path, - extract_ema=True, - original_config_file=config_file, - scan_needed=False, - ) - stanza[key].update({'format': 'folder', - 'path': destination_path, # no suffix on this - }) - - return stanza - - - def _check_for_collision(self, path: Path): - if not path.exists(): - return - if self.clobber: - shutil.rmtree(path) - else: - raise ValueError(f"Destination {path} already exists. Won't overwrite unless clobber=True.") - - def _staging_directory(self)->tempfile.TemporaryDirectory: - return tempfile.TemporaryDirectory(dir=self.config.root_path) - - - diff --git a/invokeai/backend/model_management/model_manager.py b/invokeai/backend/model_management/model_manager.py index bb46c80b77..e56ec2a0d2 100644 --- a/invokeai/backend/model_management/model_manager.py +++ b/invokeai/backend/model_management/model_manager.py @@ -703,7 +703,7 @@ class ModelManager(object): model_path = os.path.join(models_dir, entry_name) if model_path not in loaded_files: # TODO: check model_path = Path(model_path) - model_name = model_path.name if model_path.is_dir else model_path.stem + model_name = model_path.name if model_path.is_dir() else model_path.stem model_key = self.create_key(model_name, base_model, model_type) if model_key in self.models: diff --git a/invokeai/backend/model_management/model_probe.py b/invokeai/backend/model_management/model_probe.py index 59e0c8e970..ec3146a00e 100644 --- a/invokeai/backend/model_management/model_probe.py +++ b/invokeai/backend/model_management/model_probe.py @@ -401,8 +401,16 @@ class ControlNetFolderProbe(FolderProbeBase): else BaseModelType.StableDiffusion2 class LoRAFolderProbe(FolderProbeBase): - # I've never seen one of these in the wild, so this is a noop - pass + def get_base_type(self)->BaseModelType: + model_file = None + for suffix in ['safetensors','bin']: + base_file = self.folder_path / f'pytorch_lora_weights.{suffix}' + if base_file.exists(): + model_file = base_file + break + if not model_file: + raise Exception('Unknown LoRA format encountered') + return LoRACheckpointProbe(model_file,None).get_base_type() ############## register probe classes ###### ModelProbe.register_probe('diffusers', ModelType.Pipeline, PipelineFolderProbe) diff --git a/invokeai/configs/INITIAL_MODELS.yaml b/invokeai/configs/INITIAL_MODELS.yaml index cb16f3ed4b..e2288c43fd 100644 --- a/invokeai/configs/INITIAL_MODELS.yaml +++ b/invokeai/configs/INITIAL_MODELS.yaml @@ -87,7 +87,5 @@ sd-1/embedding/ahx-beta-453407d: repo_id: sd-concepts-library/ahx-beta-453407d sd-1/lora/LowRA: path: https://civitai.com/api/download/models/63006 -sd-1/lora/Ink Scenery: +sd-1/lora/Ink scenery: path: https://civitai.com/api/download/models/83390 -sd-1/lora/sd-model-finetuned-lora-t4: - repo_id: sayakpaul/sd-model-finetuned-lora-t4