Skip to content

Commit b9f9015

Browse files
DustyShoeCopilotlstein
authored
Feat(Model Manager): Add improved download manager with pause/resume partial download. (invoke-ai#8864)
* Refine messaging and pause behavior * Improved resume download behavior * Syntax fix * Formatting * Improved partial download recovering * fix(downloads): resume integrity, serialized parts, and UI feedback * Fix download test expectations and multifile totals * Ruff appease * schema updates * schema fix * Added toast msg if partial file was deleted. * Formatting * Fixed "missing temp file" message pop up * Update invokeai/app/services/download/download_default.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Fix: Add bulk action buttons and force resync on backend reconnect. --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com>
1 parent ddaa12b commit b9f9015

17 files changed

Lines changed: 2367 additions & 365 deletions

File tree

invokeai/app/api/routers/model_manager.py

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -765,6 +765,7 @@ async def list_model_installs() -> List[ModelInstallJob]:
765765
* "waiting" -- Job is waiting in the queue to run
766766
* "downloading" -- Model file(s) are downloading
767767
* "running" -- Model has downloaded and the model probing and registration process is running
768+
* "paused" -- Job is paused and can be resumed
768769
* "completed" -- Installation completed successfully
769770
* "error" -- An error occurred. Details will be in the "error_type" and "error" fields.
770771
* "cancelled" -- Job was cancelled before completion.
@@ -818,6 +819,89 @@ async def cancel_model_install_job(id: int = Path(description="Model install job
818819
installer.cancel_job(job)
819820

820821

822+
@model_manager_router.post(
823+
"/install/{id}/pause",
824+
operation_id="pause_model_install_job",
825+
responses={
826+
201: {"description": "The job was paused successfully"},
827+
415: {"description": "No such job"},
828+
},
829+
status_code=201,
830+
)
831+
async def pause_model_install_job(id: int = Path(description="Model install job ID")) -> ModelInstallJob:
832+
"""Pause the model install job corresponding to the given job ID."""
833+
installer = ApiDependencies.invoker.services.model_manager.install
834+
try:
835+
job = installer.get_job_by_id(id)
836+
except ValueError as e:
837+
raise HTTPException(status_code=415, detail=str(e))
838+
installer.pause_job(job)
839+
return job
840+
841+
842+
@model_manager_router.post(
843+
"/install/{id}/resume",
844+
operation_id="resume_model_install_job",
845+
responses={
846+
201: {"description": "The job was resumed successfully"},
847+
415: {"description": "No such job"},
848+
},
849+
status_code=201,
850+
)
851+
async def resume_model_install_job(id: int = Path(description="Model install job ID")) -> ModelInstallJob:
852+
"""Resume a paused model install job corresponding to the given job ID."""
853+
installer = ApiDependencies.invoker.services.model_manager.install
854+
try:
855+
job = installer.get_job_by_id(id)
856+
except ValueError as e:
857+
raise HTTPException(status_code=415, detail=str(e))
858+
installer.resume_job(job)
859+
return job
860+
861+
862+
@model_manager_router.post(
863+
"/install/{id}/restart_failed",
864+
operation_id="restart_failed_model_install_job",
865+
responses={
866+
201: {"description": "Failed files restarted successfully"},
867+
415: {"description": "No such job"},
868+
},
869+
status_code=201,
870+
)
871+
async def restart_failed_model_install_job(id: int = Path(description="Model install job ID")) -> ModelInstallJob:
872+
"""Restart failed or non-resumable file downloads for the given job."""
873+
installer = ApiDependencies.invoker.services.model_manager.install
874+
try:
875+
job = installer.get_job_by_id(id)
876+
except ValueError as e:
877+
raise HTTPException(status_code=415, detail=str(e))
878+
installer.restart_failed(job)
879+
return job
880+
881+
882+
@model_manager_router.post(
883+
"/install/{id}/restart_file",
884+
operation_id="restart_model_install_file",
885+
responses={
886+
201: {"description": "File restarted successfully"},
887+
415: {"description": "No such job"},
888+
},
889+
status_code=201,
890+
)
891+
async def restart_model_install_file(
892+
id: int = Path(description="Model install job ID"),
893+
file_source: AnyHttpUrl = Body(description="File download URL to restart"),
894+
) -> ModelInstallJob:
895+
"""Restart a specific file download for the given job."""
896+
installer = ApiDependencies.invoker.services.model_manager.install
897+
try:
898+
job = installer.get_job_by_id(id)
899+
except ValueError as e:
900+
raise HTTPException(status_code=415, detail=str(e))
901+
installer.restart_file(job, str(file_source))
902+
return job
903+
904+
821905
@model_manager_router.delete(
822906
"/install",
823907
operation_id="prune_model_install_jobs",

invokeai/app/services/download/download_base.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ class DownloadJobStatus(str, Enum):
1818

1919
WAITING = "waiting" # not enqueued, will not run
2020
RUNNING = "running" # actively downloading
21+
PAUSED = "paused" # paused, can be resumed
2122
COMPLETED = "completed" # finished running
2223
CANCELLED = "cancelled" # user cancelled
2324
ERROR = "error" # terminated with an error message
@@ -61,6 +62,7 @@ class DownloadJobBase(BaseModel):
6162

6263
# internal flag
6364
_cancelled: bool = PrivateAttr(default=False)
65+
_paused: bool = PrivateAttr(default=False)
6466

6567
# optional event handlers passed in on creation
6668
_on_start: Optional[DownloadEventHandler] = PrivateAttr(default=None)
@@ -72,6 +74,12 @@ class DownloadJobBase(BaseModel):
7274
def cancel(self) -> None:
7375
"""Call to cancel the job."""
7476
self._cancelled = True
77+
self._paused = False
78+
79+
def pause(self) -> None:
80+
"""Pause the job, preserving partial downloads."""
81+
self._paused = True
82+
self._cancelled = True
7583

7684
# cancelled and the callbacks are private attributes in order to prevent
7785
# them from being serialized and/or used in the Json Schema
@@ -80,6 +88,11 @@ def cancelled(self) -> bool:
8088
"""Call to cancel the job."""
8189
return self._cancelled
8290

91+
@property
92+
def paused(self) -> bool:
93+
"""Return true if job is paused."""
94+
return self._paused
95+
8396
@property
8497
def complete(self) -> bool:
8598
"""Return true if job completed without errors."""
@@ -161,6 +174,17 @@ class DownloadJob(DownloadJobBase):
161174
default=None, description="Timestamp for when the download job ende1d (completed or errored)"
162175
)
163176
content_type: Optional[str] = Field(default=None, description="Content type of downloaded file")
177+
canonical_url: Optional[str] = Field(default=None, description="Canonical URL to request on resume")
178+
etag: Optional[str] = Field(default=None, description="ETag from the remote server, if available")
179+
last_modified: Optional[str] = Field(default=None, description="Last-Modified from the remote server, if available")
180+
final_url: Optional[str] = Field(default=None, description="Final resolved URL after redirects, if available")
181+
expected_total_bytes: Optional[int] = Field(default=None, description="Expected total size of the download")
182+
resume_required: bool = Field(default=False, description="True if server refused resume; restart required")
183+
resume_message: Optional[str] = Field(default=None, description="Message explaining why resume is required")
184+
resume_from_scratch: bool = Field(
185+
default=False,
186+
description="True if resume metadata existed but the partial file was missing and the download restarted from the beginning",
187+
)
164188

165189
def __hash__(self) -> int:
166190
"""Return hash of the string representation of this object, for indexing."""
@@ -321,6 +345,10 @@ def cancel_job(self, job: DownloadJobBase) -> None:
321345
"""Cancel the job, clearing partial downloads and putting it into ERROR state."""
322346
pass
323347

348+
def pause_job(self, job: DownloadJobBase) -> None: # noqa D401
349+
"""Pause the job, preserving partial downloads."""
350+
raise NotImplementedError
351+
324352
@abstractmethod
325353
def join(self) -> None:
326354
"""Wait until all jobs are off the queue."""

0 commit comments

Comments
 (0)