[api] Add models router and list model API.

This commit is contained in:
Kyle Schouviller 2023-03-14 22:15:53 -07:00 committed by Kent Keirsey
parent 1654269125
commit c025d0521d

View File

@ -2,11 +2,7 @@
from typing import Annotated, Any, List, Literal, Optional, Union from typing import Annotated, Any, List, Literal, Optional, Union
<<<<<<< HEAD
from fastapi.routing import APIRouter, HTTPException from fastapi.routing import APIRouter, HTTPException
=======
from fastapi.routing import APIRouter
>>>>>>> 5f92f290 ([api] Add models router and list model API.)
from pydantic import BaseModel, Field, parse_obj_as from pydantic import BaseModel, Field, parse_obj_as
from ..dependencies import ApiDependencies from ..dependencies import ApiDependencies
@ -19,17 +15,9 @@ class VaeRepo(BaseModel):
path: Optional[str] = Field(description="The path to the VAE") path: Optional[str] = Field(description="The path to the VAE")
subfolder: Optional[str] = Field(description="The subfolder to use for this VAE") subfolder: Optional[str] = Field(description="The subfolder to use for this VAE")
<<<<<<< HEAD
class ModelInfo(BaseModel): class ModelInfo(BaseModel):
description: Optional[str] = Field(description="A description of the model") description: Optional[str] = Field(description="A description of the model")
=======
class ModelInfo(BaseModel):
description: Optional[str] = Field(description="A description of the model")
>>>>>>> 5f92f290 ([api] Add models router and list model API.)
class CkptModelInfo(ModelInfo): class CkptModelInfo(ModelInfo):
format: Literal['ckpt'] = 'ckpt' format: Literal['ckpt'] = 'ckpt'
@ -39,10 +27,6 @@ class CkptModelInfo(ModelInfo):
width: Optional[int] = Field(description="The width of the model") width: Optional[int] = Field(description="The width of the model")
height: Optional[int] = Field(description="The height of the model") height: Optional[int] = Field(description="The height of the model")
<<<<<<< HEAD
=======
>>>>>>> 5f92f290 ([api] Add models router and list model API.)
class DiffusersModelInfo(ModelInfo): class DiffusersModelInfo(ModelInfo):
format: Literal['diffusers'] = 'diffusers' format: Literal['diffusers'] = 'diffusers'
@ -50,7 +34,6 @@ class DiffusersModelInfo(ModelInfo):
repo_id: Optional[str] = Field(description="The repo ID to use for this model") repo_id: Optional[str] = Field(description="The repo ID to use for this model")
path: Optional[str] = Field(description="The path to the model") path: Optional[str] = Field(description="The path to the model")
<<<<<<< HEAD
class CreateModelRequest (BaseModel): class CreateModelRequest (BaseModel):
name: str = Field(description="The name of the model") name: str = Field(description="The name of the model")
info: Union[CkptModelInfo, DiffusersModelInfo] = Field(..., discriminator="format", description="The model details and configuration") info: Union[CkptModelInfo, DiffusersModelInfo] = Field(..., discriminator="format", description="The model details and configuration")
@ -59,17 +42,10 @@ class CreateModelResponse (BaseModel):
name: str = Field(description="The name of the new model") name: str = Field(description="The name of the new model")
info: Union[CkptModelInfo, DiffusersModelInfo] = Field(..., discriminator="format", description="The model details and configuration") info: Union[CkptModelInfo, DiffusersModelInfo] = Field(..., discriminator="format", description="The model details and configuration")
status: str = Field(description="The status of the API response") status: str = Field(description="The status of the API response")
=======
>>>>>>> 5f92f290 ([api] Add models router and list model API.)
class ModelsList(BaseModel): class ModelsList(BaseModel):
models: dict[str, Annotated[Union[(CkptModelInfo,DiffusersModelInfo)], Field(discriminator="format")]] models: dict[str, Annotated[Union[(CkptModelInfo,DiffusersModelInfo)], Field(discriminator="format")]]
<<<<<<< HEAD
=======
>>>>>>> 5f92f290 ([api] Add models router and list model API.)
@models_router.get( @models_router.get(
"/", "/",
operation_id="list_models", operation_id="list_models",
@ -81,7 +57,6 @@ async def list_models() -> ModelsList:
models = parse_obj_as(ModelsList, { "models": models_raw }) models = parse_obj_as(ModelsList, { "models": models_raw })
return models return models
<<<<<<< HEAD
#Update Model #Update Model
@models_router.post( @models_router.post(
"/", "/",
@ -137,8 +112,6 @@ async def delete_model(model_name: str) -> None:
raise HTTPException(status_code=500, detail=str(e)) raise HTTPException(status_code=500, detail=str(e))
=======
>>>>>>> 5f92f290 ([api] Add models router and list model API.)
# @socketio.on("requestSystemConfig") # @socketio.on("requestSystemConfig")
# def handle_request_capabilities(): # def handle_request_capabilities():
# print(">> System config requested") # print(">> System config requested")