Next: Allow in place local installs of models

This commit is contained in:
Brandon Rising 2024-02-29 15:13:38 -05:00
parent e5d9f33f7b
commit 6cb8f37a37
2 changed files with 4 additions and 1 deletions

View File

@ -451,6 +451,7 @@ async def add_model_record(
) )
async def install_model( async def install_model(
source: str = Query(description="Model source to install, can be a local path, repo_id, or remote URL"), source: str = Query(description="Model source to install, can be a local path, repo_id, or remote URL"),
inplace: Optional[bool] = Query(description="Whether or not to install a local model in place", default=False),
# TODO(MM2): Can we type this? # TODO(MM2): Can we type this?
config: Optional[Dict[str, Any]] = Body( config: Optional[Dict[str, Any]] = Body(
description="Dict of fields that override auto-probed values in the model config record, such as name, description and prediction_type ", description="Dict of fields that override auto-probed values in the model config record, such as name, description and prediction_type ",
@ -493,6 +494,7 @@ async def install_model(
source=source, source=source,
config=config, config=config,
access_token=access_token, access_token=access_token,
inplace=bool(inplace),
) )
logger.info(f"Started installation of {source}") logger.info(f"Started installation of {source}")
except UnknownModelException as e: except UnknownModelException as e:

View File

@ -181,13 +181,14 @@ class ModelInstallService(ModelInstallServiceBase):
source: str, source: str,
config: Optional[Dict[str, Any]] = None, config: Optional[Dict[str, Any]] = None,
access_token: Optional[str] = None, access_token: Optional[str] = None,
inplace: bool = False,
) -> ModelInstallJob: ) -> ModelInstallJob:
variants = "|".join(ModelRepoVariant.__members__.values()) variants = "|".join(ModelRepoVariant.__members__.values())
hf_repoid_re = f"^([^/:]+/[^/:]+)(?::({variants})?(?::/?([^:]+))?)?$" hf_repoid_re = f"^([^/:]+/[^/:]+)(?::({variants})?(?::/?([^:]+))?)?$"
source_obj: Optional[StringLikeSource] = None source_obj: Optional[StringLikeSource] = None
if Path(source).exists(): # A local file or directory if Path(source).exists(): # A local file or directory
source_obj = LocalModelSource(path=Path(source)) source_obj = LocalModelSource(path=Path(source), inplace=inplace)
elif match := re.match(hf_repoid_re, source): elif match := re.match(hf_repoid_re, source):
source_obj = HFModelSource( source_obj = HFModelSource(
repo_id=match.group(1), repo_id=match.group(1),