Configuration and model installer for new model layout

invokeai/app/api/routers/models.py

@@ -1,13 +1,13 @@
 # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) and 2023 Kent Keirsey (https://github.com/hipsterusername)
 
-from typing import Annotated, Literal, Optional, Union, Dict
+from typing import Literal, Optional, Union
 
 from fastapi import Query
 from fastapi.routing import APIRouter, HTTPException
 from pydantic import BaseModel, Field, parse_obj_as
 from ..dependencies import ApiDependencies
 from invokeai.backend import BaseModelType, ModelType
-from invokeai.backend.model_management.models import OPENAPI_MODEL_CONFIGS
+from invokeai.backend.model_management.models import OPENAPI_MODEL_CONFIGS, SchedulerPredictionType
 MODEL_CONFIGS = Union[tuple(OPENAPI_MODEL_CONFIGS)]
 
 models_router = APIRouter(prefix="/v1/models", tags=["models"])
@@ -51,11 +51,14 @@ class CreateModelResponse(BaseModel):
     info: Union[CkptModelInfo, DiffusersModelInfo] = Field(discriminator="format", description="The model info")
     status: str = Field(description="The status of the API response")
 
+class ImportModelRequest(BaseModel):
+    name: str = Field(description="A model path, repo_id or URL to import")
+    prediction_type: Optional[Literal['epsilon','v_prediction','sample']] = Field(description='Prediction type for SDv2 checkpoint files')
+
 class ConversionRequest(BaseModel):
     name: str = Field(description="The name of the new model")
     info: CkptModelInfo = Field(description="The converted model info")
     save_location: str = Field(description="The path to save the converted model weights")
-
 
 class ConvertedModelResponse(BaseModel):
     name: str = Field(description="The name of the new model")
@@ -105,6 +108,28 @@ async def update_model(
 
     return model_response
 
+@models_router.post(
+    "/",
+    operation_id="import_model",
+    responses={200: {"status": "success"}},
+)
+async def import_model(
+    model_request: ImportModelRequest
+) -> None:
+    """ Add Model """
+    items_to_import = set([model_request.name])
+    prediction_types = { x.value: x for x in SchedulerPredictionType }
+    logger = ApiDependencies.invoker.services.logger
+
+    installed_models = ApiDependencies.invoker.services.model_manager.heuristic_import(
+        items_to_import = items_to_import,
+        prediction_type_helper = lambda x: prediction_types.get(model_request.prediction_type)
+    )
+    if len(installed_models) > 0:
+        logger.info(f'Successfully imported {model_request.name}')
+    else:
+        logger.error(f'Model {model_request.name} not imported')
+        raise HTTPException(status_code=500, detail=f'Model {model_request.name} not imported')
 
 @models_router.delete(
     "/{model_name}",

invokeai/app/invocations/model.py

@@ -73,7 +73,7 @@ def invoke(self, context: InvocationContext) -> ModelLoaderOutput:
 
         base_model = self.model.base_model
         model_name = self.model.model_name
-        model_type = ModelType.Pipeline
+        model_type = ModelType.Main
 
         # TODO: not found exceptions
         if not context.services.model_manager.model_exists(

invokeai/app/services/config.py

@@ -15,7 +15,7 @@
     conf_path: configs/models.yaml
     legacy_conf_dir: configs/stable-diffusion
     outdir: outputs
-    autoconvert_dir: null
+    autoimport_dir: null
   Models:
     model: stable-diffusion-1.5
     embeddings: true
@@ -367,16 +367,17 @@ class InvokeAIAppConfig(InvokeAISettings):
 
     always_use_cpu      : bool = Field(default=False, description="If true, use the CPU for rendering even if a GPU is available.", category='Memory/Performance')
     free_gpu_mem        : bool = Field(default=False, description="If true, purge model from GPU after each generation.", category='Memory/Performance')
-    max_loaded_models   : int = Field(default=2, gt=0, description="Maximum number of models to keep in memory for rapid switching", category='Memory/Performance')
+    max_loaded_models   : int = Field(default=3, gt=0, description="Maximum number of models to keep in memory for rapid switching", category='Memory/Performance')
     precision           : Literal[tuple(['auto','float16','float32','autocast'])] = Field(default='float16',description='Floating point precision', category='Memory/Performance')
     sequential_guidance : bool = Field(default=False, description="Whether to calculate guidance in serial instead of in parallel, lowering memory requirements", category='Memory/Performance')
     xformers_enabled    : bool = Field(default=True, description="Enable/disable memory-efficient attention", category='Memory/Performance')
     tiled_decode        : bool = Field(default=False, description="Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty)", category='Memory/Performance')
 
     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')
+    autoimport_dir      : Path = Field(default='autoimport', description='Path to a directory of models files to be imported on startup.', category='Paths')
+    autoconvert_dir     : Path = Field(default=None, description='Deprecated configuration option.', category='Paths')
     conf_path           : Path = Field(default='configs/models.yaml', description='Path to models definition file', category='Paths')
-    models_dir          : Path = Field(default='./models', description='Path to the models directory', 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')
     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')