From d72f272f1682763170af3627f7089cbe5a4aee64 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Wed, 24 Apr 2024 23:53:30 -0400 Subject: [PATCH] Address change requests in first round of PR reviews. Pending: - Move model install calls into model manager and create passthrus in invocation_context. - Consider splitting load_model_from_url() into a call to get the path and a call to load the path. --- invokeai/app/invocations/upscale.py | 28 ++++----- .../model_install/model_install_default.py | 6 +- .../app/services/shared/invocation_context.py | 61 ------------------ .../migrations/migration_10.py | 62 ++++++++----------- .../model_manager/load/load_default.py | 3 +- invokeai/backend/util/util.py | 21 +++++++ .../app/services/model_load/test_load_api.py | 6 -- 7 files changed, 64 insertions(+), 123 deletions(-) diff --git a/invokeai/app/invocations/upscale.py b/invokeai/app/invocations/upscale.py index b8acfcb7bf..29cf7819de 100644 --- a/invokeai/app/invocations/upscale.py +++ b/invokeai/app/invocations/upscale.py @@ -11,7 +11,6 @@ from invokeai.app.invocations.primitives import ImageOutput from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.image_util.basicsr.rrdbnet_arch import RRDBNet from invokeai.backend.image_util.realesrgan.realesrgan import RealESRGAN -from invokeai.backend.util.devices import TorchDevice from .baseinvocation import BaseInvocation, invocation from .fields import InputField, WithBoard, WithMetadata @@ -96,22 +95,21 @@ class ESRGANInvocation(BaseInvocation, WithMetadata, WithBoard): source=ESRGAN_MODEL_URLS[self.model_name], ) - upscaler = RealESRGAN( - scale=netscale, - loadnet=loadnet.model, - model=rrdbnet_model, - half=False, - tile=self.tile_size, - ) + with loadnet as loadnet_model: + upscaler = RealESRGAN( + scale=netscale, + loadnet=loadnet_model, + model=rrdbnet_model, + half=False, + tile=self.tile_size, + ) - # prepare image - Real-ESRGAN uses cv2 internally, and cv2 uses BGR vs RGB for PIL - # TODO: This strips the alpha... is that okay? - cv2_image = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR) - upscaled_image = upscaler.upscale(cv2_image) + # prepare image - Real-ESRGAN uses cv2 internally, and cv2 uses BGR vs RGB for PIL + # TODO: This strips the alpha... is that okay? + cv2_image = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR) + upscaled_image = upscaler.upscale(cv2_image) - pil_image = Image.fromarray(cv2.cvtColor(upscaled_image, cv2.COLOR_BGR2RGB)).convert("RGBA") - - TorchDevice.empty_cache() + pil_image = Image.fromarray(cv2.cvtColor(upscaled_image, cv2.COLOR_BGR2RGB)).convert("RGBA") image_dto = context.images.save(image=pil_image) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index e259cebace..c4127acf7a 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -6,7 +6,6 @@ import re import signal import threading import time -from hashlib import sha256 from pathlib import Path from queue import Empty, Queue from shutil import copyfile, copytree, move, rmtree @@ -44,6 +43,7 @@ from invokeai.backend.model_manager.probe import ModelProbe from invokeai.backend.model_manager.search import ModelSearch from invokeai.backend.util import InvokeAILogger from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.util import slugify from .model_install_base import ( MODEL_SOURCE_TO_TYPE_MAP, @@ -396,8 +396,8 @@ class ModelInstallService(ModelInstallServiceBase): @classmethod def _download_cache_path(cls, source: Union[str, AnyHttpUrl], app_config: InvokeAIAppConfig) -> Path: - model_hash = sha256(str(source).encode("utf-8")).hexdigest()[0:32] - return app_config.download_cache_path / model_hash + escaped_source = slugify(str(source)) + return app_config.download_cache_path / escaped_source def download_and_cache( self, diff --git a/invokeai/app/services/shared/invocation_context.py b/invokeai/app/services/shared/invocation_context.py index 0d27b2520b..50551efa31 100644 --- a/invokeai/app/services/shared/invocation_context.py +++ b/invokeai/app/services/shared/invocation_context.py @@ -430,67 +430,6 @@ class ModelsInterface(InvocationContextInterface): model_format=format, ) - def install_model( - self, - source: str, - config: Optional[Dict[str, Any]] = None, - access_token: Optional[str] = None, - inplace: Optional[bool] = False, - timeout: Optional[int] = 0, - ) -> str: - """Install and register a model in the database. - - Args: - source: String source; see below - config: Optional dict. Any fields in this dict - will override corresponding autoassigned probe fields in the - model's config record. - access_token: Optional access token for remote sources. - inplace: If true, installs a local model in place rather than copying - it into the models directory - timeout: How long to wait on install (in seconds). A value of 0 (default) - blocks indefinitely - - The source can be: - 1. A local file path in posix() format (`/foo/bar` or `C:\foo\bar`) - 2. An http or https URL (`https://foo.bar/foo`) - 3. A HuggingFace repo_id (`foo/bar`, `foo/bar:fp16`, `foo/bar:fp16:vae`) - - We extend the HuggingFace repo_id syntax to include the variant and the - subfolder or path. The following are acceptable alternatives: - stabilityai/stable-diffusion-v4 - stabilityai/stable-diffusion-v4:fp16 - stabilityai/stable-diffusion-v4:fp16:vae - stabilityai/stable-diffusion-v4::/checkpoints/sd4.safetensors - stabilityai/stable-diffusion-v4:onnx:vae - - Because a local file path can look like a huggingface repo_id, the logic - first checks whether the path exists on disk, and if not, it is treated as - a parseable huggingface repo. - - Returns: - Key to the newly installed model. - - May Raise: - ValueError -- bad source - UnknownModelException -- remote model not found - InvalidModelException -- what was retrieved from remote is not a model - TimeoutError -- model could not be installed within timeout - Exception -- another error condition - """ - installer = self._services.model_manager.install - job = installer.heuristic_import( - source=source, - config=config, - access_token=access_token, - inplace=inplace, - ) - installer.wait_for_job(job, timeout) - if job.errored: - raise Exception(job.error) - key: str = job.config_out.key - return key - def download_and_cache_ckpt( self, source: Union[str, AnyHttpUrl], diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_10.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_10.py index df341a6f2e..4c4f742d4c 100644 --- a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_10.py +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_10.py @@ -1,28 +1,26 @@ -import pathlib import shutil import sqlite3 from logging import Logger from invokeai.app.services.config import InvokeAIAppConfig -from invokeai.app.services.model_install.model_install_default import ModelInstallService from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration -LEGACY_CORE_MODELS = { +LEGACY_CORE_MODELS = [ # OpenPose - "https://huggingface.co/yzd-v/DWPose/resolve/main/yolox_l.onnx?download=true": "any/annotators/dwpose/yolox_l.onnx", - "https://huggingface.co/yzd-v/DWPose/resolve/main/dw-ll_ucoco_384.onnx?download=true": "any/annotators/dwpose/dw-ll_ucoco_384.onnx", + "any/annotators/dwpose/yolox_l.onnx", + "any/annotators/dwpose/dw-ll_ucoco_384.onnx", # DepthAnything - "https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vitl14.pth?download=true": "any/annotators/depth_anything/depth_anything_vitl14.pth", - "https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vitb14.pth?download=true": "any/annotators/depth_anything/depth_anything_vitb14.pth", - "https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vits14.pth?download=true": "any/annotators/depth_anything/depth_anything_vits14.pth", + "any/annotators/depth_anything/depth_anything_vitl14.pth", + "any/annotators/depth_anything/depth_anything_vitb14.pth", + "any/annotators/depth_anything/depth_anything_vits14.pth", # Lama inpaint - "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt": "core/misc/lama/lama.pt", + "core/misc/lama/lama.pt", # RealESRGAN upscale - "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth": "core/upscaling/realesrgan/RealESRGAN_x4plus.pth", - "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth": "core/upscaling/realesrgan/RealESRGAN_x4plus_anime_6B.pth", - "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth": "core/upscaling/realesrgan/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth", - "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth": "core/upscaling/realesrgan/RealESRGAN_x2plus.pth", -} + "core/upscaling/realesrgan/RealESRGAN_x4plus.pth", + "core/upscaling/realesrgan/RealESRGAN_x4plus_anime_6B.pth", + "core/upscaling/realesrgan/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth", + "core/upscaling/realesrgan/RealESRGAN_x2plus.pth", +] class Migration10Callback: @@ -31,34 +29,24 @@ class Migration10Callback: self._logger = logger def __call__(self, cursor: sqlite3.Cursor) -> None: - self._rename_convert_cache() - self._migrate_downloaded_models_cache() + self._remove_convert_cache() + self._remove_downloaded_models() self._remove_unused_core_models() - def _rename_convert_cache(self) -> None: + def _remove_convert_cache(self) -> None: """Rename models/.cache to models/.convert_cache.""" + self._logger.info("Removing .cache directory. Converted models will now be cached in .convert_cache.") legacy_convert_path = self._app_config.root_path / "models" / ".cache" - configured_convert_dir = self._app_config.convert_cache_dir - configured_convert_path = self._app_config.convert_cache_path - # old convert dir was in use, and current convert dir has not been changed - if legacy_convert_path.exists() and configured_convert_dir == pathlib.Path("models/.convert_cache"): - self._logger.info( - f"Migrating legacy convert cache directory from {str(legacy_convert_path)} to {str(configured_convert_path)}" - ) - shutil.rmtree(configured_convert_path, ignore_errors=True) # shouldn't be needed, but just in case... - shutil.move(legacy_convert_path, configured_convert_path) + shutil.rmtree(legacy_convert_path, ignore_errors=True) - def _migrate_downloaded_models_cache(self) -> None: - """Move used core models to modsl/.download_cache.""" - self._logger.info(f"Migrating legacy core models to {str(self._app_config.download_cache_path)}") - for url, legacy_dest in LEGACY_CORE_MODELS.items(): - legacy_dest_path = self._app_config.models_path / legacy_dest - if not legacy_dest_path.exists(): - continue - # this returns a unique directory path - new_path = ModelInstallService._download_cache_path(url, self._app_config) - new_path.mkdir(parents=True, exist_ok=True) - shutil.move(legacy_dest_path, new_path / legacy_dest_path.name) + def _remove_downloaded_models(self) -> None: + """Remove models from their old locations; they will re-download when needed.""" + self._logger.info( + "Removing legacy just-in-time models. Downloaded models will now be cached in .download_cache." + ) + for model_path in LEGACY_CORE_MODELS: + legacy_dest_path = self._app_config.models_path / model_path + legacy_dest_path.unlink(missing_ok=True) def _remove_unused_core_models(self) -> None: """Remove unused core models and their directories.""" diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index 459808b455..16b9e3646e 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/invokeai/backend/model_manager/load/load_default.py @@ -19,6 +19,7 @@ from invokeai.backend.model_manager.load.model_cache.model_cache_base import Mod from invokeai.backend.model_manager.load.model_util import calc_model_size_by_fs from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.util import slugify # TO DO: The loader is not thread safe! @@ -84,7 +85,7 @@ class ModelLoader(ModelLoaderBase): except IndexError: pass - cache_path: Path = self._convert_cache.cache_path(config.key) + cache_path: Path = self._convert_cache.cache_path(slugify(model_path)) if self._needs_conversion(config, model_path, cache_path): loaded_model = self._do_convert(config, model_path, cache_path, submodel_type) else: diff --git a/invokeai/backend/util/util.py b/invokeai/backend/util/util.py index 7d0d9d03f7..2b2bee3458 100644 --- a/invokeai/backend/util/util.py +++ b/invokeai/backend/util/util.py @@ -1,6 +1,8 @@ import base64 import io import os +import re +import unicodedata import warnings from pathlib import Path @@ -12,6 +14,25 @@ from transformers import logging as transformers_logging GIG = 1073741824 +def slugify(value: str, allow_unicode: bool = False) -> str: + """ + Convert to ASCII if 'allow_unicode' is False. Convert spaces or repeated + dashes to single dashes. Remove characters that aren't alphanumerics, + underscores, or hyphens. Convert to lowercase. Also strip leading and + trailing whitespace, dashes, and underscores. + + Adapted from Django: https://github.com/django/django/blob/main/django/utils/text.py + """ + value = str(value) + if allow_unicode: + value = unicodedata.normalize("NFKC", value) + else: + value = unicodedata.normalize("NFKD", value).encode("ascii", "ignore").decode("ascii") + value = re.sub(r"[/]", "_", value.lower()) + value = re.sub(r"[^\w\s-]", "", value.lower()) + return re.sub(r"[-\s]+", "-", value).strip("-_") + + def directory_size(directory: Path) -> int: """ Return the aggregate size of all files in a directory (bytes). diff --git a/tests/app/services/model_load/test_load_api.py b/tests/app/services/model_load/test_load_api.py index 59b207f2ee..463be86c68 100644 --- a/tests/app/services/model_load/test_load_api.py +++ b/tests/app/services/model_load/test_load_api.py @@ -49,9 +49,3 @@ def test_download_and_load(mock_context: InvocationContext): assert isinstance(model_1, dict) -def test_install_model(mock_context: InvocationContext): - key = mock_context.models.install_model("https://www.test.foo/download/test_embedding.safetensors") - assert key is not None - model = mock_context.models.load(key) - assert model is not None - assert model.config.key == key