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

refactor big _install_next_item() loop

Browse Source
This commit is contained in:
Lincoln Stein 2024-03-19 23:45:02 -04:00 committed by psychedelicious
parent 30283a4767
commit 9bd7dabed3
1 changed files with 33 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -438,7 +438,24 @@ class ModelInstallService(ModelInstallServiceBase):
elif (
job.waiting or job.downloads_done
): # local jobs will be in waiting state, remote jobs will be downloading state
):
self._register_or_install(job)
except InvalidModelConfigException as excp:
self._set_error(job, excp)
except (OSError, DuplicateModelException) as excp:
self._set_error(job, excp)
finally:
# if this is an install of a remote file, then clean up the temporary directory
if job._install_tmpdir is not None:
rmtree(job._install_tmpdir)
self._install_completed_event.set()
self._install_queue.task_done()
def _register_or_install(self, job: ModelInstallJob) -> None:
# local jobs will be in waiting state, remote jobs will be downloading state
job.total_bytes = self._stat_size(job.local_path)
job.bytes = job.total_bytes
self._signal_job_running(job)
@ -455,7 +472,7 @@ class ModelInstallService(ModelInstallServiceBase):
job.config_out = self.record_store.get_model(key)
self._signal_job_completed(job)
except InvalidModelConfigException as excp:
def _set_error(self, job: ModelInstallJob, excp: Exception) -> None:
if any(x.content_type is not None and "text/html" in x.content_type for x in job.download_parts):
job.set_error(
InvalidModelConfigException(
@ -466,17 +483,6 @@ class ModelInstallService(ModelInstallServiceBase):
job.set_error(excp)
self._signal_job_errored(job)
except (OSError, DuplicateModelException) as excp:
job.set_error(excp)
self._signal_job_errored(job)
finally:
# if this is an install of a remote file, then clean up the temporary directory
if job._install_tmpdir is not None:
rmtree(job._install_tmpdir)
self._install_completed_event.set()
self._install_queue.task_done()
# --------------------------------------------------------------------------------------------
# Internal functions that manage the models directory
# --------------------------------------------------------------------------------------------