diff --git a/configs/INITIAL_MODELS.yaml b/configs/INITIAL_MODELS.yaml index 94502855da..ea43280340 100644 --- a/configs/INITIAL_MODELS.yaml +++ b/configs/INITIAL_MODELS.yaml @@ -93,6 +93,7 @@ voxel_art-1.0: format: ckpt vae: repo_id: stabilityai/sd-vae-ft-mse + file: vae-ft-mse-840000-ema-pruned.ckpt recommended: False width: 512 height: 512 @@ -102,7 +103,7 @@ ft-mse-improved-autoencoder-840000: format: ckpt config: VAE/default file: vae-ft-mse-840000-ema-pruned.ckpt - recommended: False + recommended: True width: 512 height: 512 trinart_vae: diff --git a/ldm/generate.py b/ldm/generate.py index c94eafc465..1afb642721 100644 --- a/ldm/generate.py +++ b/ldm/generate.py @@ -852,6 +852,7 @@ class Generate: model_data = cache.get_model(model_name) except Exception as e: print(f'** model {model_name} could not be loaded: {str(e)}') + print(traceback.format_exc(), file=sys.stderr) if previous_model_name is None: raise e print(f'** trying to reload previous model') diff --git a/ldm/invoke/CLI.py b/ldm/invoke/CLI.py index 58df4bdc92..4d637e6ea2 100644 --- a/ldm/invoke/CLI.py +++ b/ldm/invoke/CLI.py @@ -578,7 +578,7 @@ def import_model(model_path:str, gen, opt, completer): elif re.match('^[\w.+-]+/[\w.+-]+$',model_path): model_name = import_diffuser_model(model_path, gen, opt, completer) elif os.path.isdir(model_path): - model_name = import_diffuser_model(model_path, gen, opt, completer) + model_name = import_diffuser_model(Path(model_path), gen, opt, completer) else: print(f'** {model_path} is neither the path to a .ckpt file nor a diffusers repository id. Can\'t import.') @@ -590,7 +590,7 @@ def import_model(model_path:str, gen, opt, completer): gen.model_manager.del_model(model_name) return - if input('Make this the default model? [n] ') in ('y','Y'): + if input('Make this the default model? [n] ').strip() in ('y','Y'): gen.model_manager.set_default_model(model_name) gen.model_manager.commit(opt.conf) @@ -607,10 +607,14 @@ def import_diffuser_model(path_or_repo:str, gen, opt, completer)->str: model_name=default_name, model_description=default_description ) + vae = None + if input('Replace this model\'s VAE with "stabilityai/sd-vae-ft-mse"? [n] ').strip() in ('y','Y'): + vae = dict(repo_id='stabilityai/sd-vae-ft-mse') if not manager.import_diffuser_model( path_or_repo, model_name = model_name, + vae = vae, description = model_description): print('** model failed to import') return None @@ -628,17 +632,28 @@ def import_ckpt_model(path_or_url:str, gen, opt, completer)->str: ) config_file = None default = Path(Globals.root,'configs/stable-diffusion/v1-inference.yaml') + completer.complete_extensions(('.yaml','.yml')) completer.set_line(str(default)) done = False while not done: config_file = input('Configuration file for this model: ').strip() done = os.path.exists(config_file) + + completer.complete_extensions(('.ckpt','.safetensors')) + vae = None + default = Path(Globals.root,'models/ldm/stable-diffusion-v1/vae-ft-mse-840000-ema-pruned.ckpt') + completer.set_line(str(default)) + done = False + while not done: + vae = input('VAE file for this model (leave blank for none): ').strip() or None + done = (not vae) or os.path.exists(vae) completer.complete_extensions(None) if not manager.import_ckpt_model( path_or_url, config = config_file, + vae = vae, model_name = model_name, model_description = model_description, commit_to_conf = opt.conf, @@ -710,7 +725,7 @@ def optimize_model(model_name_or_path:str, gen, opt, completer): return completer.update_models(gen.model_manager.list_models()) - if input(f'Load optimized model {model_name}? [y] ') not in ('n','N'): + if input(f'Load optimized model {model_name}? [y] ').strip() not in ('n','N'): gen.set_model(model_name) response = input(f'Delete the original .ckpt file at ({ckpt_path} ? [n] ') @@ -726,7 +741,12 @@ def del_config(model_name:str, gen, opt, completer): if model_name not in gen.model_manager.config: print(f"** Unknown model {model_name}") return - gen.model_manager.del_model(model_name) + + if input(f'Remove {model_name} from the list of models known to InvokeAI? [y] ').strip().startswith(('n','N')): + return + + delete_completely = input('Completely remove the model file or directory from disk? [n] ').startswith(('y','Y')) + gen.model_manager.del_model(model_name,delete_files=delete_completely) gen.model_manager.commit(opt.conf) print(f'** {model_name} deleted') completer.update_models(gen.model_manager.list_models()) diff --git a/ldm/invoke/model_manager.py b/ldm/invoke/model_manager.py index 88d2b6ccac..585e37355f 100644 --- a/ldm/invoke/model_manager.py +++ b/ldm/invoke/model_manager.py @@ -18,7 +18,9 @@ import traceback import warnings import safetensors.torch from pathlib import Path +from shutil import move, rmtree from typing import Union, Any +from huggingface_hub import scan_cache_dir from ldm.util import download_with_progress_bar import torch @@ -35,6 +37,9 @@ from ldm.invoke.globals import Globals, global_models_dir, global_autoscan_dir, from ldm.util import instantiate_from_config, ask_user DEFAULT_MAX_MODELS=2 +VAE_TO_REPO_ID = { # hack, see note in convert_and_import() + 'vae-ft-mse-840000-ema-pruned': 'stabilityai/sd-vae-ft-mse', + } class ModelManager(object): def __init__(self, @@ -230,7 +235,7 @@ class ModelManager(object): line = f'\033[1m{line}\033[0m' print(line) - def del_model(self, model_name:str) -> None: + def del_model(self, model_name:str, delete_files:bool=False) -> None: ''' Delete the named model. ''' @@ -238,9 +243,25 @@ class ModelManager(object): if model_name not in omega: print(f'** Unknown model {model_name}') return + # save these for use in deletion later + conf = omega[model_name] + repo_id = conf.get('repo_id',None) + path = self._abs_path(conf.get('path',None)) + weights = self._abs_path(conf.get('weights',None)) + del omega[model_name] if model_name in self.stack: self.stack.remove(model_name) + if delete_files: + if weights: + print(f'** deleting file {weights}') + Path(weights).unlink(missing_ok=True) + elif path: + print(f'** deleting directory {path}') + rmtree(path,ignore_errors=True) + elif repo_id: + print(f'** deleting the cached model directory for {repo_id}') + self._delete_model_from_cache(repo_id) def add_model(self, model_name:str, model_attributes:dict, clobber:bool=False) -> None: ''' @@ -417,7 +438,7 @@ class ModelManager(object): safety_checker=None, local_files_only=not Globals.internet_available ) - if 'vae' in mconfig: + if 'vae' in mconfig and mconfig['vae'] is not None: vae = self._load_vae(mconfig['vae']) pipeline_args.update(vae=vae) if not isinstance(name_or_path,Path): @@ -523,11 +544,12 @@ class ModelManager(object): print('>> Model scanned ok!') def import_diffuser_model(self, - repo_or_path:Union[str,Path], - model_name:str=None, - description:str=None, - commit_to_conf:Path=None, - )->bool: + repo_or_path:Union[str,Path], + model_name:str=None, + description:str=None, + vae:dict=None, + commit_to_conf:Path=None, + )->bool: ''' Attempts to install the indicated diffuser model and returns True if successful. @@ -543,6 +565,7 @@ class ModelManager(object): description = description or f'imported diffusers model {model_name}' new_config = dict( description=description, + vae=vae, format='diffusers', ) if isinstance(repo_or_path,Path) and repo_or_path.exists(): @@ -556,18 +579,22 @@ class ModelManager(object): return True def import_ckpt_model(self, - weights:Union[str,Path], - config:Union[str,Path]='configs/stable-diffusion/v1-inference.yaml', - model_name:str=None, - model_description:str=None, - commit_to_conf:Path=None, - )->bool: + weights:Union[str,Path], + config:Union[str,Path]='configs/stable-diffusion/v1-inference.yaml', + vae:Union[str,Path]=None, + model_name:str=None, + model_description:str=None, + commit_to_conf:Path=None, + )->bool: ''' Attempts to install the indicated ckpt file and returns True if successful. "weights" can be either a path-like object corresponding to a local .ckpt file or a http/https URL pointing to a remote model. + "vae" is a Path or str object pointing to a ckpt or safetensors file to be used + as the VAE for this model. + "config" is the model config file to use with this ckpt file. It defaults to v1-inference.yaml. If a URL is provided, the config will be downloaded. @@ -594,6 +621,8 @@ class ModelManager(object): width=512, height=512 ) + if vae: + new_config['vae'] = vae self.add_model(model_name, new_config, True) if commit_to_conf: self.commit(commit_to_conf) @@ -633,7 +662,7 @@ class ModelManager(object): def convert_and_import(self, ckpt_path:Path, - diffuser_path:Path, + diffusers_path:Path, model_name=None, model_description=None, commit_to_conf:Path=None, @@ -645,46 +674,56 @@ class ModelManager(object): new_config = None from ldm.invoke.ckpt_to_diffuser import convert_ckpt_to_diffuser import transformers - if diffuser_path.exists(): - print(f'ERROR: The path {str(diffuser_path)} already exists. Please move or remove it and try again.') + if diffusers_path.exists(): + print(f'ERROR: The path {str(diffusers_path)} already exists. Please move or remove it and try again.') return - model_name = model_name or diffuser_path.name + model_name = model_name or diffusers_path.name model_description = model_description or 'Optimized version of {model_name}' - print(f'>> {model_name}: optimizing (30-60s).') + print(f'>> Optimizing {model_name} (30-60s)') try: verbosity =transformers.logging.get_verbosity() transformers.logging.set_verbosity_error() - convert_ckpt_to_diffuser(ckpt_path, diffuser_path,extract_ema=True) + convert_ckpt_to_diffuser(ckpt_path, diffusers_path,extract_ema=True) transformers.logging.set_verbosity(verbosity) - print(f'>> Success. Optimized model is now located at {str(diffuser_path)}') - print(f'>> Writing new config file entry for {model_name}...',end='') + print(f'>> Success. Optimized model is now located at {str(diffusers_path)}') + print(f'>> Writing new config file entry for {model_name}') new_config = dict( - path=str(diffuser_path), + path=str(diffusers_path), description=model_description, format='diffusers', ) + + # HACK (LS): in the event that the original entry is using a custom ckpt VAE, we try to + # map that VAE onto a diffuser VAE using a hard-coded dictionary. + # I would prefer to do this differently: We load the ckpt model into memory, swap the + # VAE in memory, and then pass that to convert_ckpt_to_diffuser() so that the swapped + # VAE is built into the model. However, when I tried this I got obscure key errors. + if model_name in self.config and (vae_ckpt_path := self.model_info(model_name)['vae']): + vae_basename = Path(vae_ckpt_path).stem + diffusers_vae = None + if (diffusers_vae := VAE_TO_REPO_ID.get(vae_basename,None)): + print(f'>> {vae_basename} VAE corresponds to known {diffusers_vae} diffusers version') + new_config.update( + vae = {'repo_id': diffusers_vae} + ) + else: + print(f'** Custom VAE "{vae_basename}" found, but corresponding diffusers model unknown') + print(f'** Using "stabilityai/sd-vae-ft-mse"; If this isn\'t right, please edit the model config') + new_config.update( + vae = {'repo_id': 'stabilityai/sd-vae-ft-mse'} + ) + self.del_model(model_name) self.add_model(model_name, new_config, True) if commit_to_conf: self.commit(commit_to_conf) + print('>> Conversion succeeded') except Exception as e: print(f'** Conversion failed: {str(e)}') - traceback.print_exc() - print('done.') return new_config - def del_config(self, model_name:str, gen, opt, completer): - current_model = gen.model_name - if model_name == current_model: - print("** Can't delete active model. !switch to another model first. **") - return - gen.model_manager.del_model(model_name) - gen.model_manager.commit(opt.conf) - print(f'** {model_name} deleted') - completer.del_model(model_name) - def search_models(self, search_folder): print(f'>> Finding Models In: {search_folder}') models_folder_ckpt = Path(search_folder).glob('**/*.ckpt') @@ -766,7 +805,6 @@ class ModelManager(object): print('** Legacy version <= 2.2.5 model directory layout detected. Reorganizing.') print('** This is a quick one-time operation.') - from shutil import move, rmtree # transformer files get moved into the hub directory if cls._is_huggingface_hub_directory_present(): @@ -982,6 +1020,27 @@ class ModelManager(object): return vae + @staticmethod + def _delete_model_from_cache(repo_id): + cache_info = scan_cache_dir(global_cache_dir('diffusers')) + + # I'm sure there is a way to do this with comprehensions + # but the code quickly became incomprehensible! + hashes_to_delete = set() + for repo in cache_info.repos: + if repo.repo_id==repo_id: + for revision in repo.revisions: + hashes_to_delete.add(revision.commit_hash) + strategy = cache_info.delete_revisions(*hashes_to_delete) + print(f'** deletion of this model is expected to free {strategy.expected_freed_size_str}') + strategy.execute() + + @staticmethod + def _abs_path(path:Union(str,Path))->Path: + if path is None or Path(path).is_absolute(): + return path + return Path(Globals.root,path).resolve() + @staticmethod def _is_huggingface_hub_directory_present() -> bool: return os.getenv('HF_HOME') is not None or os.getenv('XDG_CACHE_HOME') is not None