model_delete method now working

This commit is contained in:
Lincoln Stein 2023-07-04 10:40:32 -04:00
parent c1c49d9a76
commit 752b4d50cf
2 changed files with 20 additions and 20 deletions

View File

@ -2,7 +2,7 @@
from typing import Literal, Optional, Union from typing import Literal, Optional, Union
from fastapi import Query, Body from fastapi import Query, Body, Path
from fastapi.routing import APIRouter, HTTPException from fastapi.routing import APIRouter, HTTPException
from pydantic import BaseModel, Field, parse_obj_as from pydantic import BaseModel, Field, parse_obj_as
from ..dependencies import ApiDependencies from ..dependencies import ApiDependencies
@ -153,7 +153,7 @@ async def import_model(
@models_router.delete( @models_router.delete(
"/{model_name}", "/{base_model}/{model_type}/{model_name}",
operation_id="del_model", operation_id="del_model",
responses={ responses={
204: { 204: {
@ -164,24 +164,27 @@ async def import_model(
} }
}, },
) )
async def delete_model(model_name: str) -> None: async def delete_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"),
) -> None:
"""Delete Model""" """Delete Model"""
model_names = ApiDependencies.invoker.services.model_manager.model_names()
logger = ApiDependencies.invoker.services.logger logger = ApiDependencies.invoker.services.logger
model_exists = model_name in model_names
# check if model exists try:
logger.info(f"Checking for model {model_name}...") ApiDependencies.invoker.services.model_manager.del_model(model_name,
base_model = base_model,
if model_exists: model_type = model_type
logger.info(f"Deleting Model: {model_name}") )
ApiDependencies.invoker.services.model_manager.del_model(model_name, delete_files=True) logger.info(f"Deleted model: {model_name}")
logger.info(f"Model Deleted: {model_name}")
raise HTTPException(status_code=204, detail=f"Model '{model_name}' deleted successfully") raise HTTPException(status_code=204, detail=f"Model '{model_name}' deleted successfully")
except KeyError:
else: logger.error(f"Model not found: {model_name}")
logger.error("Model not found")
raise HTTPException(status_code=404, detail=f"Model '{model_name}' not found") raise HTTPException(status_code=404, detail=f"Model '{model_name}' not found")
else:
logger.info(f"Model deleted: {model_name}")
raise HTTPException(status_code=204, detail=f"Model '{model_name}' deleted successfully")
# @socketio.on("convertToDiffusers") # @socketio.on("convertToDiffusers")

View File

@ -546,10 +546,7 @@ class ModelManager(object):
model_cfg = self.models.pop(model_key, None) model_cfg = self.models.pop(model_key, None)
if model_cfg is None: if model_cfg is None:
self.logger.error( raise KeyError(f"Unknown model {model_key}")
f"Unknown model {model_key}"
)
return
# note: it not garantie to release memory(model can has other references) # note: it not garantie to release memory(model can has other references)
cache_ids = self.cache_keys.pop(model_key, []) cache_ids = self.cache_keys.pop(model_key, [])