mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
67 lines
2.7 KiB
Python
67 lines
2.7 KiB
Python
# Copyright (c) 2023 The InvokeAI Development Team
|
|
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
from omegaconf import OmegaConf
|
|
|
|
from invokeai.app.services.config import InvokeAIAppConfig
|
|
from invokeai.backend.util.logging import InvokeAILogger
|
|
|
|
from ..config import BaseModelType, MainCheckpointConfig, MainConfig, ModelType
|
|
from .base import CONFIG_FILE_VERSION
|
|
|
|
|
|
def migrate_models_store(config: InvokeAIAppConfig):
|
|
"""Migrate models from v1 models.yaml to v3.2 models.yaml."""
|
|
# avoid circular import
|
|
from invokeai.backend.model_manager import DuplicateModelException, ModelInstall
|
|
from invokeai.backend.model_manager.storage import get_config_store
|
|
|
|
app_config = InvokeAIAppConfig.get_config()
|
|
logger = InvokeAILogger.get_logger()
|
|
old_file: Path = app_config.model_conf_path
|
|
new_file: Path = old_file.with_name("models3_2.yaml")
|
|
|
|
old_conf = OmegaConf.load(old_file)
|
|
store = get_config_store(new_file)
|
|
installer = ModelInstall(store=store)
|
|
logger.info(f"Migrating old models file at {old_file} to new {CONFIG_FILE_VERSION} format")
|
|
|
|
for model_key, stanza in old_conf.items():
|
|
if model_key == "__metadata__":
|
|
assert (
|
|
stanza["version"] == "3.0.0"
|
|
), f"This script works on version 3.0.0 yaml files, but your configuration points to a {stanza['version']} version"
|
|
continue
|
|
|
|
base_type, model_type, model_name = str(model_key).split("/")
|
|
new_key = "<NOKEY>"
|
|
|
|
try:
|
|
path = app_config.models_path / stanza["path"]
|
|
new_key = installer.register_path(path)
|
|
except DuplicateModelException:
|
|
# if model already installed, then we just update its info
|
|
models = store.search_by_name(
|
|
model_name=model_name, base_model=BaseModelType(base_type), model_type=ModelType(model_type)
|
|
)
|
|
if len(models) != 1:
|
|
continue
|
|
new_key = models[0].key
|
|
except Exception as excp:
|
|
print(str(excp))
|
|
|
|
model_info = store.get_model(new_key)
|
|
if vae := stanza.get("vae") and isinstance(model_info, MainConfig):
|
|
model_info.vae = (app_config.models_path / vae).as_posix()
|
|
if model_config := stanza.get("config") and isinstance(model_info, MainCheckpointConfig):
|
|
model_info.config = (app_config.root_path / model_config).as_posix()
|
|
model_info.description = stanza.get("description")
|
|
store.update_model(new_key, model_info)
|
|
store.update_model(new_key, model_info)
|
|
|
|
logger.info(f"Original version of models config file saved as {str(old_file) + '.orig'}")
|
|
shutil.move(old_file, str(old_file) + ".orig")
|
|
shutil.move(new_file, old_file)
|