Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

all methods now return OPENAPI_MODEL_CONFIGS; convert uses PUT

Browse Source
This commit is contained in:
Lincoln Stein 2023-07-05 23:13:01 -04:00
parent 3691b55565
commit f7daa6e71d
4 changed files with 71 additions and 42 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

58
invokeai/app/api/routers/models.py
View File

@ -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):

30
invokeai/app/services/model_manager_service.py
View File

@ -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
@ -105,6 +99,19 @@ 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,

21
invokeai/backend/model_management/model_manager.py
View File

@ -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)

4
invokeai/backend/model_management/model_merge.py
View File

@ -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,
):