mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
120 lines
4.4 KiB
Python
120 lines
4.4 KiB
Python
# Copyright 2024 Lincoln D. Stein and the InvokeAI Development Team
|
|
|
|
"""
|
|
Utility class for migrating among versions of the InvokeAI app config schema.
|
|
"""
|
|
|
|
from abc import ABC, abstractmethod
|
|
from typing import TYPE_CHECKING, Any, Callable, List, TypeVar
|
|
|
|
from packaging.version import Version
|
|
from pydantic import BaseModel, ConfigDict, field_validator
|
|
|
|
if TYPE_CHECKING:
|
|
pass
|
|
|
|
AppConfigDict = TypeVar("AppConfigDict", bound=dict[str, Any])
|
|
|
|
|
|
class ConfigMigratorBase(ABC):
|
|
"""This class allows migrators to register their input and output versions."""
|
|
|
|
@classmethod
|
|
@abstractmethod
|
|
def register(
|
|
cls, from_version: Version, to_version: Version
|
|
) -> Callable[[Callable[[AppConfigDict], AppConfigDict]], Callable[[AppConfigDict], AppConfigDict]]:
|
|
"""Define a decorator which registers the migration between two versions."""
|
|
|
|
@classmethod
|
|
@abstractmethod
|
|
def migrate(cls, config: AppConfigDict) -> AppConfigDict:
|
|
"""
|
|
Use the registered migration steps to bring config up to latest version.
|
|
|
|
:param config: The original configuration.
|
|
:return: The new configuration, lifted up to the latest version.
|
|
|
|
As a side effect, the new configuration will be written to disk.
|
|
"""
|
|
|
|
|
|
class MigrationEntry(BaseModel):
|
|
"""Defines an individual migration."""
|
|
|
|
model_config = ConfigDict(arbitrary_types_allowed=True)
|
|
|
|
from_version: Version
|
|
to_version: Version
|
|
function: Callable[[AppConfigDict], AppConfigDict]
|
|
|
|
@field_validator("from_version", "to_version", mode="before")
|
|
@classmethod
|
|
def _string_to_version(cls, v: str | Version) -> Version: # noqa D102
|
|
if isinstance(v, str):
|
|
return Version(v)
|
|
else:
|
|
return v
|
|
|
|
|
|
class ConfigMigrator(ConfigMigratorBase):
|
|
"""This class allows migrators to register their input and output versions."""
|
|
|
|
_migrations: List[MigrationEntry] = []
|
|
|
|
@classmethod
|
|
def register(
|
|
cls,
|
|
from_version: Version | str,
|
|
to_version: Version | str,
|
|
) -> Callable[[Callable[[AppConfigDict], AppConfigDict]], Callable[[AppConfigDict], AppConfigDict]]:
|
|
"""Define a decorator which registers the migration between two versions."""
|
|
|
|
def decorator(function: Callable[[AppConfigDict], AppConfigDict]) -> Callable[[AppConfigDict], AppConfigDict]:
|
|
if from_version in cls._migrations:
|
|
raise ValueError(
|
|
f"function {function.__name__} is trying to register a migration for version {str(from_version)}, but this migration has already been registered."
|
|
)
|
|
cls._migrations.append(MigrationEntry(from_version=from_version, to_version=to_version, function=function))
|
|
return function
|
|
|
|
return decorator
|
|
|
|
@staticmethod
|
|
def _check_for_overlaps(migrations: List[MigrationEntry]) -> None:
|
|
current_version = Version("0.0.0")
|
|
for m in migrations:
|
|
if current_version > m.from_version:
|
|
raise ValueError(f"Version range overlap detected while processing function {m.function.__name__}")
|
|
|
|
@classmethod
|
|
def migrate(cls, config_dict: AppConfigDict) -> AppConfigDict:
|
|
"""
|
|
Use the registered migration steps to bring config up to latest version.
|
|
|
|
:param config: The original configuration.
|
|
:return: The new configuration, lifted up to the latest version.
|
|
|
|
As a side effect, the new configuration will be written to disk.
|
|
If an inconsistency in the registered migration steps' `from_version`
|
|
and `to_version` parameters are identified, this will raise a
|
|
ValueError exception.
|
|
"""
|
|
# Sort migrations by version number and raise a ValueError if
|
|
# any version range overlaps are detected. Discontinuities are ok
|
|
sorted_migrations = sorted(cls._migrations, key=lambda x: x.from_version)
|
|
cls._check_for_overlaps(sorted_migrations)
|
|
|
|
if "InvokeAI" in config_dict:
|
|
version = Version("3.0.0")
|
|
else:
|
|
version = Version(config_dict["schema_version"])
|
|
|
|
for migration in sorted_migrations:
|
|
if version >= migration.from_version and version < migration.to_version:
|
|
config_dict = migration.function(config_dict)
|
|
version = migration.to_version
|
|
|
|
config_dict["schema_version"] = str(version)
|
|
return config_dict
|