Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

add directory scanning for loras, controlnets and textual_inversions

Browse Source
This commit is contained in:
Lincoln Stein 2023-06-08 23:11:53 -04:00
parent 6652f3405b
commit 887576d217
4 changed files with 173 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

61
invokeai/app/services/config.py
View File

@ -15,10 +15,7 @@ InvokeAI:
conf_path: configs/models.yaml
legacy_conf_dir: configs/stable-diffusion
outdir: outputs
embedding_dir: embeddings
lora_dir: loras
autoconvert_dir: null
gfpgan_model_dir: models/gfpgan/GFPGANv1.4.pth
Models:
model: stable-diffusion-1.5
embeddings: true
@ -171,7 +168,7 @@ from argparse import ArgumentParser
from omegaconf import OmegaConf, DictConfig
from pathlib import Path
from pydantic import BaseSettings, Field, parse_obj_as
from typing import ClassVar, Dict, List, Literal, Type, Union, get_origin, get_type_hints, get_args
from typing import ClassVar, Dict, List, Literal, Union, get_origin, get_type_hints, get_args
INIT_FILE = Path('invokeai.yaml')
DB_FILE = Path('invokeai.db')
@ -379,19 +376,15 @@ setting environment variables INVOKEAI_<setting>.
root : Path = Field(default=_find_root(), description='InvokeAI runtime root directory', category='Paths')
autoconvert_dir : Path = Field(default=None, description='Path to a directory of ckpt files to be converted into diffusers and imported on startup.', category='Paths')
conf_path : Path = Field(default='configs/models.yaml', description='Path to models definition file', category='Paths')
embedding_dir : Path = Field(default='embeddings', description='Path to InvokeAI textual inversion aembeddings directory', category='Paths')
gfpgan_model_dir : Path = Field(default="./models/gfpgan/GFPGANv1.4.pth", description='Path to GFPGAN models directory.', category='Paths')
controlnet_dir : Path = Field(default="controlnets", description='Path to directory of ControlNet models.', category='Paths')
models_dir : Path = Field(default='./models', description='Path to the models directory', category='Paths')
legacy_conf_dir : Path = Field(default='configs/stable-diffusion', description='Path to directory of legacy checkpoint config files', category='Paths')
lora_dir : Path = Field(default='loras', description='Path to InvokeAI LoRA model directory', category='Paths')
db_dir : Path = Field(default='databases', description='Path to InvokeAI databases directory', category='Paths')
outdir : Path = Field(default='outputs', description='Default folder for output images', category='Paths')
from_file : Path = Field(default=None, description='Take command input from the indicated file (command-line client only)', category='Paths')
use_memory_db : bool = Field(default=False, description='Use in-memory database for storing image metadata', category='Paths')
model : str = Field(default='stable-diffusion-1.5', description='Initial model name', category='Models')
embeddings : bool = Field(default=True, description='Load contents of embeddings directory', category='Models')
log_handlers : List[str] = Field(default=["console"], description='Log handler. Valid options are "console", "file=<path>", "syslog=path|address:host:port", "http=<url>"', category="Logging")
# note - would be better to read the log_format values from logging.py, but this creates circular dependencies issues
log_format : Literal[tuple(['plain','color','syslog','legacy'])] = Field(default="color", description='Log format. Use "plain" for text-only, "color" for colorized output, "legacy" for 2.3-style logging and "syslog" for syslog-style', category="Logging")
@ -492,46 +485,11 @@ setting environment variables INVOKEAI_<setting>.
return self._resolve(self.legacy_conf_dir)
@property
def cache_dir(self)->Path:
'''
Path to the global cache directory for HuggingFace hub-managed models
'''
return self.models_dir / "hub"
@property
def models_dir(self)->Path:
def models_path(self)->Path:
'''
Path to the models directory
'''
return self._resolve("models")
@property
def converted_ckpts_dir(self)->Path:
'''
Path to the converted models
'''
return self._resolve("models/converted_ckpts")
@property
def embedding_path(self)->Path:
'''
Path to the textual inversion embeddings directory.
'''
return self._resolve(self.embedding_dir) if self.embedding_dir else None
@property
def lora_path(self)->Path:
'''
Path to the LoRA models directory.
'''
return self._resolve(self.lora_dir) if self.lora_dir else None
@property
def controlnet_path(self)->Path:
'''
Path to the controlnet models directory.
'''
return self._resolve(self.controlnet_dir) if self.controlnet_dir else None
return self._resolve(self.models_dir)
@property
def autoconvert_path(self)->Path:
@ -540,13 +498,6 @@ setting environment variables INVOKEAI_<setting>.
'''
return self._resolve(self.autoconvert_dir) if self.autoconvert_dir else None
@property
def gfpgan_model_path(self)->Path:
'''
Path to the GFPGAN model.
'''
return self._resolve(self.gfpgan_model_dir) if self.gfpgan_model_dir else None
# the following methods support legacy calls leftover from the Globals era
@property
def full_precision(self)->bool:

5
invokeai/backend/model_management/model_cache.py
View File

@ -162,7 +162,12 @@ class SDModelType(str, Enum):
Scheduler = "scheduler"
Lora = "lora"
TextualInversion = "textual_inversion"
ControlNet = "control_net"
class BaseModel(str, Enum):
StableDiffusion1_5 = "SD-1"
StableDiffusion2Base = "SD-2-base" # 512 pixels; this will have epsilon parameterization
StableDiffusion2 = "SD-2" # 768 pixels; this will have v-prediction parameterization
class ModelInfoBase:
#model_path: str

127
invokeai/backend/model_management/model_manager.py
View File

@ -54,15 +54,17 @@ MODELS.YAML
The general format of a models.yaml section is:
type-of-model/name-of-model:
format: folder|ckpt|safetensors
repo_id: owner/repo
path: /path/to/local/file/or/directory
description: a description
format: folder|ckpt|safetensors|pt
base: SD-1|SD-2
subfolder: subfolder-name
The type of model is given in the stanza key, and is one of
{diffusers, ckpt, vae, text_encoder, tokenizer, unet, scheduler,
safety_checker, feature_extractor, lora, textual_inversion}, and
correspond to items in the SDModelType enum defined in model_cache.py
safety_checker, feature_extractor, lora, textual_inversion,
controlnet}, and correspond to items in the SDModelType enum defined
in model_cache.py
The format indicates whether the model is organized as a folder with
model subdirectories, or is contained in a single checkpoint or
@ -80,12 +82,12 @@ This example summarizes the two ways of getting a non-diffuser model:
text_encoder/clip-test-1:
format: folder
repo_id: openai/clip-vit-large-patch14
path: /path/to/folder
description: Returns standalone CLIPTextModel
text_encoder/clip-test-2:
format: folder
repo_id: stabilityai/stable-diffusion-2
repo_id: /path/to/folder
subfolder: text_encoder
description: Returns the text_encoder in the subfolder of the diffusers model (just the encoder in RAM)
@ -99,6 +101,14 @@ model. Use the `submodel` parameter to select which part:
print(type(my_vae))
# "AutoencoderKL"
DIRECTORY_SCANNING:
Loras, textual_inversion and controlnet models are usually not listed
explicitly in models.yaml, but are added to the in-memory data
structure at initialization time by scanning the models directory. The
in-memory data structure can be resynchronized by calling
`manager.scan_models_directory`.
DISAMBIGUATION:
You may wish to use the same name for a related family of models. To
@ -107,12 +117,12 @@ separated by "/". Example:
tokenizer/clip-large:
format: tokenizer
repo_id: openai/clip-vit-large-patch14
path: /path/to/folder
description: Returns standalone tokenizer
text_encoder/clip-large:
format: text_encoder
repo_id: openai/clip-vit-large-patch14
path: /path/to/folder
description: Returns standalone text encoder
You can now use the `model_type` argument to indicate which model you
@ -126,6 +136,14 @@ OTHER FUNCTIONS:
Other methods provided by ModelManager support importing, editing,
converting and deleting models.
IMPORTANT CHANGES AND LIMITATIONS SINCE 2.3:
1. Only local paths are supported. Repo_ids are no longer accepted. This
simplifies the logic.
2. VAEs can't be swapped in and out at load time. They must be baked
into the model when downloaded or converted.
"""
from __future__ import annotations
@ -133,17 +151,13 @@ import os
import re
import textwrap
import shutil
import sys
import time
import traceback
import warnings
from dataclasses import dataclass
from enum import Enum, auto
from packaging import version
from pathlib import Path
from typing import Callable, Optional, List, Tuple, Union, types
from shutil import move, rmtree
from typing import Any, Optional, Union, Callable, Dict, List, types
from typing import Callable, Dict, Optional, List, Tuple, Union, types
from shutil import rmtree
import safetensors
import safetensors.torch
@ -156,7 +170,7 @@ from omegaconf.dictconfig import DictConfig
import invokeai.backend.util.logging as logger
from invokeai.app.services.config import InvokeAIAppConfig
from invokeai.backend.util import CUDA_DEVICE, ask_user, download_with_resume
from invokeai.backend.util import CUDA_DEVICE, download_with_resume
from ..install.model_install_backend import Dataset_path, hf_download_with_resume
from .model_cache import (ModelCache, ModelLocker, SDModelType,
SilenceWarnings)
@ -197,6 +211,25 @@ class SDLegacyType(Enum):
MAX_CACHE_SIZE = 6.0 # GB
# layout of the models directory:
# models
# ├── SD-1
# │   ├── controlnet
# │   ├── lora
# │   ├── diffusers
# │   └── textual_inversion
# ├── SD-2
# │   ├── controlnet
# │   ├── lora
# │   ├── diffusers
# │   └── textual_inversion
# └── support
# ├── codeformer
# ├── gfpgan
# └── realesrgan
class ModelManager(object):
"""
High-level interface to model management.
@ -241,6 +274,9 @@ class ModelManager(object):
)
self.cache_keys = dict()
# add controlnet, lora and textual_inversion models from disk
self.scan_models_directory(include_diffusers=False)
def model_exists(
self,
model_name: str,
@ -1258,6 +1294,67 @@ class ModelManager(object):
if self.config_path:
self.commit()
def _delete_defunct_models(self):
'''
Remove models no longer on disk.
'''
config = self.config
to_delete = set()
for key in config:
if 'path' not in config[key]:
continue
path = self.globals.root_dir / config[key].path
if path.exists():
continue
to_delete.add(key)
for key in to_delete:
self.logger.warn(f'Removing model {key} from in-memory config because its path is no longer on disk')
config.pop(key)
def scan_models_directory(self, include_diffusers:bool=False):
'''
Scan the models directory for loras, textual_inversions and controlnets
and create appropriate entries in the in-memory omegaconf. Diffusers
will not be added unless include_diffusers is true.
'''
self._delete_defunct_models()
model_directory = self.globals.models_path
config = self.config
for root, dirs, files in os.walk(model_directory):
parents = root.split('/')
subpaths = parents[parents.index('models')+1:]
if len(subpaths) < 2:
continue
base, model_type, *_ = subpaths
if model_type == "diffusers" and not include_diffusers:
continue
for d in dirs:
config[f'{model_type}/{d}'] = dict(
path = os.path.join(root,d),
description = f'{model_type} model {d}',
format = 'folder',
base = base,
)
for f in files:
basename = Path(f).stem
format = Path(f).suffix[1:]
config[f'{model_type}/{basename}'] = dict(
path = os.path.join(root,f),
description = f'{model_type} model {basename}',
format = format,
base = base,
)
##### NONE OF THE METHODS BELOW WORK NOW BECAUSE OF MODEL DIRECTORY REORGANIZATION
##### AND NEED TO BE REWRITTEN
def list_lora_models(self)->Dict[str,bool]:
'''Return a dict of installed lora models; key is either the shortname
defined in INITIAL_MODELS, or the basename of the file in the LoRA

50
scripts/scan_models_directory.py Normal file
View File

@ -0,0 +1,50 @@
'''
Scan the models directory and print out a new models.yaml
'''
import os
import sys
import argparse
from pathlib import Path
from omegaconf import OmegaConf
def main():
parser = argparse.ArgumentParser(description="Model directory scanner")
parser.add_argument('models_directory')
args = parser.parse_args()
directory = args.models_directory
conf = OmegaConf.create()
conf['_version'] = '3.0.0'
for root, dirs, files in os.walk(directory):
for d in dirs:
parents = root.split('/')
subpaths = parents[parents.index('models')+1:]
if len(subpaths) < 2:
continue
base, model_type, *_ = subpaths
conf[f'{model_type}/{d}'] = dict(
path = os.path.join(root,d),
description = f'{model_type} model {d}',
format = 'folder',
base = base,
)
for f in files:
basename = Path(f).stem
format = Path(f).suffix[1:]
conf[f'{model_type}/{basename}'] = dict(
path = os.path.join(root,f),
description = f'{model_type} model {basename}',
format = format,
base = base,
)
OmegaConf.save(config=dict(sorted(conf.items())), f=sys.stdout)
if __name__ == '__main__':
main()