diff --git a/installer/lib/main.py b/installer/lib/main.py index be9dc18f96..3dd3f2739e 100644 --- a/installer/lib/main.py +++ b/installer/lib/main.py @@ -8,9 +8,19 @@ from pathlib import Path from installer import Installer + +def find_wheel() -> Path: + dist = Path("./dist") + wheel = next(dist.glob("*.whl")) + assert wheel is not None + return wheel + + if __name__ == "__main__": parser = argparse.ArgumentParser() + wheel = find_wheel() + parser.add_argument( "-r", "--root", @@ -42,7 +52,7 @@ if __name__ == "__main__": dest="wheel", help="Specifies a wheel for the InvokeAI package. Used for troubleshooting or testing prereleases.", type=Path, - default=None, + default=wheel, ) args = parser.parse_args() diff --git a/invokeai/app/api/routers/model_manager.py b/invokeai/app/api/routers/model_manager.py index 99f00423c6..1730d383f0 100644 --- a/invokeai/app/api/routers/model_manager.py +++ b/invokeai/app/api/routers/model_manager.py @@ -6,6 +6,7 @@ import pathlib import shutil import traceback from copy import deepcopy +from enum import Enum from typing import Any, Dict, List, Optional, Type from fastapi import Body, Path, Query, Response, UploadFile @@ -16,6 +17,7 @@ from pydantic import AnyHttpUrl, BaseModel, ConfigDict, Field from starlette.exceptions import HTTPException from typing_extensions import Annotated +from invokeai.app.services.config import get_config from invokeai.app.services.model_images.model_images_common import ModelImageFileNotFoundException from invokeai.app.services.model_install.model_install_common import ModelInstallJob from invokeai.app.services.model_records import ( @@ -32,6 +34,7 @@ from invokeai.backend.model_manager.config import ( ModelType, SubModelType, ) +from invokeai.backend.model_manager.load.model_cache.model_cache_base import CacheStats from invokeai.backend.model_manager.metadata.fetch.huggingface import HuggingFaceMetadataFetch from invokeai.backend.model_manager.metadata.metadata_base import ModelMetadataWithFiles, UnknownMetadataException from invokeai.backend.model_manager.search import ModelSearch @@ -53,6 +56,13 @@ class ModelsList(BaseModel): model_config = ConfigDict(use_enum_values=True) +class CacheType(str, Enum): + """Cache type - one of vram or ram.""" + + RAM = "RAM" + VRAM = "VRAM" + + def add_cover_image_to_model_config(config: AnyModelConfig, dependencies: Type[ApiDependencies]) -> AnyModelConfig: """Add a cover image URL to a model configuration.""" cover_image = dependencies.invoker.services.model_images.get_url(config.key) @@ -174,18 +184,6 @@ async def get_model_record( raise HTTPException(status_code=404, detail=str(e)) -# @model_manager_router.get("/summary", operation_id="list_model_summary") -# async def list_model_summary( -# page: int = Query(default=0, description="The page to get"), -# per_page: int = Query(default=10, description="The number of models per page"), -# order_by: ModelRecordOrderBy = Query(default=ModelRecordOrderBy.Default, description="The attribute to order by"), -# ) -> PaginatedResults[ModelSummary]: -# """Gets a page of model summary data.""" -# record_store = ApiDependencies.invoker.services.model_manager.store -# results: PaginatedResults[ModelSummary] = record_store.list_models(page=page, per_page=per_page, order_by=order_by) -# return results - - class FoundModel(BaseModel): path: str = Field(description="Path to the model") is_installed: bool = Field(description="Whether or not the model is already installed") @@ -816,3 +814,61 @@ async def get_starter_models() -> list[StarterModel]: model.dependencies = missing_deps return starter_models + + +@model_manager_router.get( + "/model_cache", + operation_id="get_cache_size", + response_model=float, + summary="Get maximum size of model manager RAM or VRAM cache.", +) +async def get_cache_size(cache_type: CacheType = Query(description="The cache type", default=CacheType.RAM)) -> float: + """Return the current RAM or VRAM cache size setting (in GB).""" + cache = ApiDependencies.invoker.services.model_manager.load.ram_cache + return cache.max_cache_size if cache_type == CacheType.RAM else cache.max_vram_cache_size + + +@model_manager_router.put( + "/model_cache", + operation_id="set_cache_size", + response_model=float, + summary="Set maximum size of model manager RAM or VRAM cache, optionally writing new value out to invokeai.yaml config file.", +) +async def set_cache_size( + value: float = Query(description="The new value for the maximum cache size"), + cache_type: CacheType = Query(description="The cache type", default=CacheType.RAM), + persist: bool = Query(description="Write new value out to invokeai.yaml", default=False), +) -> float: + """Set the current RAM or VRAM cache size setting (in GB). .""" + cache = ApiDependencies.invoker.services.model_manager.load.ram_cache + app_config = get_config() + if cache_type == CacheType.RAM: + cache.max_cache_size = value + app_config.ram = value + elif cache_type == CacheType.VRAM: + cache.max_vram_cache_size = value + app_config.vram = value + + if persist: + config_path = app_config.config_file_path + print(f"DEBUG: config_path = {config_path}") + try: + shutil.copy(config_path, config_path.with_suffix(".yaml.bak")) + app_config.write_file(config_path) + except Exception as e: + shutil.move(config_path.with_suffix(".yaml.bak"), config_path) + raise RuntimeError(f"Failed to write modified configuration to {config_path}: {e}") from e + + return cache.max_vram_cache_size if cache_type == CacheType.VRAM else cache.max_cache_size + + +@model_manager_router.get( + "/stats", + operation_id="get_stats", + response_model=Optional[CacheStats], + summary="Get model manager RAM cache performance statistics.", +) +async def get_stats() -> Optional[CacheStats]: + """Return performance statistics on the model manager's RAM cache. Will return null if no models have been loaded.""" + + return ApiDependencies.invoker.services.model_manager.load.ram_cache.stats diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index a63cc66a86..13030395d2 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/invokeai/backend/model_manager/load/load_default.py @@ -79,8 +79,10 @@ class ModelLoader(ModelLoaderBase): def _convert_and_load( self, config: AnyModelConfig, model_path: Path, submodel_type: Optional[SubModelType] = None ) -> ModelLockerBase: + stats_name = ":".join([config.base, config.type, config.name, (submodel_type or "")]) + try: - return self._ram_cache.get(config.key, submodel_type) + return self._ram_cache.get(config.key, submodel_type, stats_name=stats_name) except IndexError: pass @@ -100,7 +102,7 @@ class ModelLoader(ModelLoaderBase): return self._ram_cache.get( key=config.key, submodel_type=submodel_type, - stats_name=":".join([config.base, config.type, config.name, (submodel_type or "")]), + stats_name=stats_name, ) def get_size_fs( diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache_base.py b/invokeai/backend/model_manager/load/model_cache/model_cache_base.py index 6e518bc9e3..ae7b86753b 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache_base.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache_base.py @@ -124,7 +124,24 @@ class ModelCacheBase(ABC, Generic[T]): @property @abstractmethod def max_cache_size(self) -> float: - """Return true if the cache is configured to lazily offload models in VRAM.""" + """Return the maximum size the RAM cache can grow to.""" + pass + + @max_cache_size.setter + @abstractmethod + def max_cache_size(self, value: float) -> None: + """Set the cap on vram cache size.""" + + @property + @abstractmethod + def max_vram_cache_size(self) -> float: + """Return the maximum size the VRAM cache can grow to.""" + pass + + @max_vram_cache_size.setter + @abstractmethod + def max_vram_cache_size(self, value: float) -> float: + """Set the maximum size the VRAM cache can grow to.""" pass @property diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache_default.py b/invokeai/backend/model_manager/load/model_cache/model_cache_default.py index 3a2d36e87a..68d7afead0 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache_default.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache_default.py @@ -72,7 +72,6 @@ class ModelCache(ModelCacheBase[AnyModel]): :param max_cache_size: Maximum size of the RAM cache [6.0 GB] :param storage_device: Torch device to save inactive model in [torch.device('cpu')] :param precision: Precision for loaded models [torch.float16] - :param sequential_offload: Conserve VRAM by loading and unloading each stage of the pipeline sequentially :param log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to disable this feature unless you are actively inspecting the model cache's @@ -182,6 +181,16 @@ class ModelCache(ModelCacheBase[AnyModel]): """Set the cap on cache size.""" self._max_cache_size = value + @property + def max_vram_cache_size(self) -> float: + """Return the cap on vram cache size.""" + return self._max_vram_cache_size + + @max_vram_cache_size.setter + def max_vram_cache_size(self, value: float) -> None: + """Set the cap on vram cache size.""" + self._max_vram_cache_size = value + @property def stats(self) -> Optional[CacheStats]: """Return collected CacheStats object.""" @@ -298,8 +307,11 @@ class ModelCache(ModelCacheBase[AnyModel]): """ self.logger.info(f"Called to move {cache_entry.key} ({type(cache_entry.model)=}) to {target_device}") - # Some models don't have a state dictionary, in which case the - # stored model will still reside in CPU + # Some models don't have a state dictionary, in which case we brute-force + # a to(), and if that doesn't work, just keep the thing in CPU. + # This is primarily to handle the IP_Adapter model, which has + # two state dicts and no load_state_dict(). This could be fixed + # more elegantly. if cache_entry.state_dict is None: if hasattr(cache_entry.model, "to"): model_in_gpu = copy.deepcopy(cache_entry.model) diff --git a/invokeai/version/invokeai_version.py b/invokeai/version/invokeai_version.py index 9b575128e6..27fdca497c 100644 --- a/invokeai/version/invokeai_version.py +++ b/invokeai/version/invokeai_version.py @@ -1 +1 @@ -__version__ = "4.2.4" +__version__ = "0.0.3" diff --git a/pyproject.toml b/pyproject.toml index fcc0aff60c..7e04c5a9ab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["setuptools~=65.5", "pip~=22.3", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "InvokeAI" +name = "InvokeAI-MGPU" description = "An implementation of Stable Diffusion which provides various new features and options to aid the image generation process" requires-python = ">=3.10, <3.12" readme = { content-type = "text/markdown", file = "README.md" }