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

Allow passing in key on register

Browse Source
This commit is contained in:
Brandon Rising 2024-02-23 12:57:54 -05:00 committed by Brandon
parent 80ad14d89f
commit 09295ae43b
2 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
invokeai/app/services/model_install/model_install_default.py
View File

@ -542,8 +542,10 @@ class ModelInstallService(ModelInstallServiceBase):
def _register( def _register(
self, model_path: Path, config: Optional[Dict[str, Any]] = None, info: Optional[AnyModelConfig] = None self, model_path: Path, config: Optional[Dict[str, Any]] = None, info: Optional[AnyModelConfig] = None
) -> str: ) -> str:
info = info or ModelProbe.probe(model_path, config)
key = self._create_key() key = self._create_key()
if config and not config.get('key', None):
config['key'] = key
info = info or ModelProbe.probe(model_path, config)
model_path = model_path.absolute() model_path = model_path.absolute()
if model_path.is_relative_to(self.app_config.models_path): if model_path.is_relative_to(self.app_config.models_path):
@ -556,8 +558,8 @@ class ModelInstallService(ModelInstallServiceBase):
# make config relative to our root # make config relative to our root
legacy_conf = (self.app_config.root_dir / self.app_config.legacy_conf_dir / info.config).resolve() legacy_conf = (self.app_config.root_dir / self.app_config.legacy_conf_dir / info.config).resolve()
info.config = legacy_conf.relative_to(self.app_config.root_dir).as_posix() info.config = legacy_conf.relative_to(self.app_config.root_dir).as_posix()
self.record_store.add_model(key, info) self.record_store.add_model(info.key, info)
return key return info.key
def _next_id(self) -> int: def _next_id(self) -> int:
with self._lock: with self._lock:

2
invokeai/backend/model_manager/probe.py
View File

@ -188,7 +188,7 @@ class ModelProbe(object):
and fields["prediction_type"] == SchedulerPredictionType.VPrediction and fields["prediction_type"] == SchedulerPredictionType.VPrediction
) )
model_info = ModelConfigFactory.make_config(fields) model_info = ModelConfigFactory.make_config(fields, key=fields.get("key", None))
return model_info return model_info
@classmethod @classmethod