diff --git a/invokeai/app/api/routers/models.py b/invokeai/app/api/routers/models.py index a3c0d1db50..25b227e87a 100644 --- a/invokeai/app/api/routers/models.py +++ b/invokeai/app/api/routers/models.py @@ -1,16 +1,15 @@ -# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) and 2023 Kent Keirsey (https://github.com/hipsterusername) +# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654), 2023 Kent Keirsey (https://github.com/hipsterusername), 2024 Lincoln Stein + from typing import Literal, Optional, Union from fastapi import Body, Path, Query, Response from fastapi.routing import APIRouter -from pydantic import BaseModel, Field, parse_obj_as +from pydantic import BaseModel, parse_obj_as from starlette.exceptions import HTTPException from invokeai.backend import BaseModelType, ModelType -from invokeai.backend.model_management import AddModelResult from invokeai.backend.model_management.models import ( - MODEL_CONFIGS, OPENAPI_MODEL_CONFIGS, SchedulerPredictionType ) @@ -19,13 +18,9 @@ from ..dependencies import ApiDependencies models_router = APIRouter(prefix="/v1/models", tags=["models"]) -class UpdateModelResponse(BaseModel): - model_name: str = Field(description="The name of the new model") - info: Union[tuple(MODEL_CONFIGS)] = Field(description="The model info") - -class ImportModelResponse(BaseModel): - location: str = Field(description="The path, repo_id or URL of the imported model") - info: AddModelResult = Field(description="The model info") +UpdateModelResponse = Union[tuple(OPENAPI_MODEL_CONFIGS)] +ImportModelResponse = Union[tuple(OPENAPI_MODEL_CONFIGS)] +ConvertModelResponse = Union[tuple(OPENAPI_MODEL_CONFIGS)] class ModelsList(BaseModel): models: list[Union[tuple(OPENAPI_MODEL_CONFIGS)]] @@ -62,7 +57,7 @@ async def update_model( base_model: BaseModelType = Path(default='sd-1', description="Base model"), model_type: ModelType = Path(default='main', description="The type of model"), model_name: str = Path(default=None, description="model name"), - info: Union[tuple(MODEL_CONFIGS)] = Body(description="Model configuration"), + info: Union[tuple(OPENAPI_MODEL_CONFIGS)] = Body(description="Model configuration"), ) -> UpdateModelResponse: """ Add Model """ try: @@ -72,14 +67,12 @@ async def update_model( model_type=model_type, model_attributes=info.dict() ) - model_response = UpdateModelResponse( - model_name = model_name, - info = ApiDependencies.invoker.services.model_manager.model_info( - model_name=model_name, - base_model=base_model, - model_type=model_type, - ) + model_raw = ApiDependencies.invoker.services.model_manager.list_model( + model_name=model_name, + base_model=base_model, + model_type=model_type, ) + model_response = parse_obj_as(UpdateModelResponse, model_raw) except KeyError as e: raise HTTPException(status_code=404, detail=str(e)) except ValueError as e: @@ -122,10 +115,13 @@ async def import_model( raise HTTPException(status_code=424) logger.info(f'Successfully imported {location}, got {info}') - return ImportModelResponse( - location = location, - info = info, + model_raw = ApiDependencies.invoker.services.model_manager.list_model( + model_name=info.name, + base_model=info.base_model, + model_type=info.model_type ) + return parse_obj_as(ImportModelResponse, model_raw) + except KeyError as e: logger.error(str(e)) raise HTTPException(status_code=404, detail=str(e)) @@ -165,7 +161,7 @@ async def delete_model( logger.error(f"Model not found: {model_name}") raise HTTPException(status_code=404, detail=f"Model '{model_name}' not found") -@models_router.patch( +@models_router.put( "/convert/{base_model}/{model_type}/{model_name}", operation_id="convert_model", responses={ @@ -174,26 +170,30 @@ async def delete_model( 404: { "description": "Model not found" }, }, status_code = 200, - response_model = Union[tuple(MODEL_CONFIGS)], + response_model = Union[tuple(OPENAPI_MODEL_CONFIGS)], ) async def convert_model( base_model: BaseModelType = Path(description="Base model"), model_type: ModelType = Path(description="The type of model"), model_name: str = Path(description="model name"), -) -> Union[tuple(MODEL_CONFIGS)]: +) -> ConvertModelResponse: """Convert a checkpoint model into a diffusers model""" logger = ApiDependencies.invoker.services.logger try: logger.info(f"Converting model: {model_name}") - result = ApiDependencies.invoker.services.model_manager.convert_model(model_name, + ApiDependencies.invoker.services.model_manager.convert_model(model_name, + base_model = base_model, + model_type = model_type + ) + model_raw = ApiDependencies.invoker.services.model_manager.list_model(model_name, base_model = base_model, - model_type = model_type - ) + model_type = model_type) + response = parse_obj_as(ConvertModelResponse, model_raw) except KeyError: raise HTTPException(status_code=404, detail=f"Model '{model_name}' not found") except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) - return result.config + return response # @socketio.on("mergeDiffusersModels") # def merge_diffusers_models(model_merge_info: dict): diff --git a/invokeai/app/services/model_manager_service.py b/invokeai/app/services/model_manager_service.py index 0834b11559..3e307d7e76 100644 --- a/invokeai/app/services/model_manager_service.py +++ b/invokeai/app/services/model_manager_service.py @@ -76,13 +76,7 @@ class ModelManagerServiceBase(ABC): def model_info(self, model_name: str, base_model: BaseModelType, model_type: ModelType) -> dict: """ Given a model name returns a dict-like (OmegaConf) object describing it. - """ - pass - - @abstractmethod - def model_names(self) -> List[Tuple[str, BaseModelType, ModelType]]: - """ - Returns a list of all the model names known. + Uses the exact format as the omegaconf stanza. """ pass @@ -104,7 +98,20 @@ class ModelManagerServiceBase(ABC): } """ pass + + @abstractmethod + def list_model(self, model_name: str, base_model: BaseModelType, model_type: ModelType) -> dict: + """ + Return information about the model using the same format as list_models() + """ + pass + @abstractmethod + def model_names(self) -> List[Tuple[str, BaseModelType, ModelType]]: + """ + Returns a list of all the model names known. + """ + pass @abstractmethod def add_model( @@ -339,12 +346,19 @@ class ModelManagerService(ModelManagerServiceBase): base_model: Optional[BaseModelType] = None, model_type: Optional[ModelType] = None ) -> list[dict]: - # ) -> dict: """ Return a list of models. """ return self.mgr.list_models(base_model, model_type) + def list_model(self, model_name: str, base_model: BaseModelType, model_type: ModelType) -> dict: + """ + Return information about the model using the same format as list_models() + """ + return self.mgr.list_model(model_name=model_name, + base_model=base_model, + model_type=model_type) + def add_model( self, model_name: str, diff --git a/invokeai/backend/model_management/model_manager.py b/invokeai/backend/model_management/model_manager.py index e8a4a0541c..505956ed69 100644 --- a/invokeai/backend/model_management/model_manager.py +++ b/invokeai/backend/model_management/model_manager.py @@ -480,7 +480,7 @@ class ModelManager(object): """ model_key = self.create_key(model_name, base_model, model_type) if model_key in self.models: - return self.models[model_key].dict(exclude_defaults=True, exclude={"error"}) + return self.models[model_key].dict(exclude_defaults=True) else: return None # TODO: None or empty dict on not found @@ -491,17 +491,32 @@ class ModelManager(object): """ return [(self.parse_key(x)) for x in self.models.keys()] + def list_model( + self, + model_name: str, + base_model: BaseModelType, + model_type: ModelType, + )->dict: + """ + Returns a dict describing one installed model, using + the combined format of the list_models() method. + """ + models = self.list_models(base_model,model_type,model_name) + return models[0] if models else None + def list_models( self, base_model: Optional[BaseModelType] = None, model_type: Optional[ModelType] = None, + model_name: Optional[str] = None, ) -> list[dict]: """ Return a list of models. """ + model_keys = [self.create_key(model_name, base_model, model_type)] if model_name else sorted(self.models, key=str.casefold) models = [] - for model_key in sorted(self.models, key=str.casefold): + for model_key in model_keys: model_config = self.models[model_key] cur_model_name, cur_base_model, cur_model_type = self.parse_key(model_key) @@ -653,7 +668,7 @@ class ModelManager(object): old_diffusers_path = self.app_config.models_path / model.location new_diffusers_path = self.app_config.models_path / base_model.value / model_type.value / model_name if new_diffusers_path.exists(): - raise ValueError(f"A diffusers model already exists at {new_path}") + raise ValueError(f"A diffusers model already exists at {new_diffusers_path}") try: move(old_diffusers_path,new_diffusers_path) diff --git a/invokeai/backend/model_management/model_merge.py b/invokeai/backend/model_management/model_merge.py index 1519c76f6f..1a331ff040 100644 --- a/invokeai/backend/model_management/model_merge.py +++ b/invokeai/backend/model_management/model_merge.py @@ -26,7 +26,7 @@ class MergeInterpolationMethod(str, Enum): def merge_diffusion_models( model_paths: List[Path], alpha: float = 0.5, - interp: InterpolationMethod = None, + interp: MergeInterpolationMethod = None, force: bool = False, **kwargs, ) -> DiffusionPipeline: @@ -67,7 +67,7 @@ def merge_diffusion_models_and_save ( merged_model_name: str, config: InvokeAIAppConfig, alpha: float = 0.5, - interp: InterpolationMethod = None, + interp: MergeInterpolationMethod = None, force: bool = False, **kwargs, ):