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 psychedelicious
parent 65b91356d0
commit c778ab8db4
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(
self, model_path: Path, config: Optional[Dict[str, Any]] = None, info: Optional[AnyModelConfig] = None
) -> str:
info = info or ModelProbe.probe(model_path, config)
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()
if model_path.is_relative_to(self.app_config.models_path):
@ -556,8 +558,8 @@ class ModelInstallService(ModelInstallServiceBase):
# make config relative to our root
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()
self.record_store.add_model(key, info)
return key
self.record_store.add_model(info.key, info)
return info.key
def _next_id(self) -> int:
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
)
model_info = ModelConfigFactory.make_config(fields)
model_info = ModelConfigFactory.make_config(fields, key=fields.get("key", None))
return model_info
@classmethod