From c97611e7d7da214f45b311c9999a2a3152c81c23 Mon Sep 17 00:00:00 2001 From: DustyShoe Date: Sat, 28 Feb 2026 19:26:57 +0200 Subject: [PATCH] fix(model-install): persist remote access_token for resume after restart --- .../model_install/model_install_default.py | 12 +++- .../model_install/test_model_install.py | 63 +++++++++++++++++++ 2 files changed, 73 insertions(+), 2 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index c47267eab5f..714ae9329a4 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -132,6 +132,9 @@ def _write_install_marker(self, job: ModelInstallJob, status: Optional[InstallSt marker = { "version": INSTALL_MARKER_VERSION, "source": str(job.source), + "access_token": ( + job.source.access_token if isinstance(job.source, (HFModelSource, URLModelSource)) else None + ), "config_in": job.config_in.model_dump(), "status": (status or job.status).value, "updated_at": get_iso_timestamp(), @@ -200,7 +203,13 @@ def _restore_incomplete_installs(self) -> None: continue try: - source_str = marker["source"] + source_str = marker.get("source") + if not isinstance(source_str, str): + raise ValueError("Missing source in install marker") + source = self._guess_source(source_str) + access_token = marker.get("access_token") + if isinstance(source, (HFModelSource, URLModelSource)) and isinstance(access_token, str): + source.access_token = access_token if source_str in active_sources: # This tmpdir belongs to an install already in progress; leave it alone. self._logger.debug(f"Skipping restore for {source_str} - already being tracked") @@ -210,7 +219,6 @@ def _restore_incomplete_installs(self) -> None: self._safe_rmtree(tmpdir, self._logger) continue seen_sources.add(source_str) - source = self._guess_source(source_str) except Exception as e: self._logger.warning(f"Skipping install marker in {tmpdir}: {e}") continue diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index d19eb95a8c2..9a7cc6aaaaf 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -4,6 +4,7 @@ import gc import platform +import shutil import uuid from pathlib import Path from typing import Any, Dict @@ -23,6 +24,7 @@ ) from invokeai.app.services.model_install import ( HFModelSource, + ModelInstallService, ModelInstallServiceBase, ) from invokeai.app.services.model_install.model_install_common import ( @@ -343,6 +345,67 @@ def test_huggingface_repo_id(mm2_installer: ModelInstallServiceBase, mm2_app_con assert job.total_bytes == sum(x["total_bytes"] for x in downloading_events[-1].parts) +def test_restore_paused_hf_install_preserves_access_token( + mm2_installer: ModelInstallServiceBase, + mm2_app_config: InvokeAIAppConfig, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + assert isinstance(mm2_installer, ModelInstallService) + + access_token = "hf_test_access_token" + tmpdir = mm2_app_config.models_path / f"tmpinstall_resume_token_{uuid.uuid4().hex}" + tmpdir.mkdir(parents=True, exist_ok=True) + + try: + paused_job = ModelInstallJob( + id=99999, + source=HFModelSource( + repo_id="stabilityai/sdxl-turbo", + variant=ModelRepoVariant.Default, + access_token=access_token, + ), + config_in=ModelRecordChanges(), + local_path=tmpdir, + ) + paused_job._install_tmpdir = tmpdir + paused_job.status = InstallStatus.PAUSED + + mm2_installer._write_install_marker(paused_job, status=InstallStatus.PAUSED) + + marker = mm2_installer._read_install_marker(tmpdir) + assert marker is not None + assert marker["access_token"] == access_token + + restored_installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_installer.record_store, + download_queue=mm2_download_queue, + session=mm2_session, + ) + restored_installer._restore_incomplete_installs() + restored_jobs = restored_installer.list_jobs() + assert len(restored_jobs) == 1 + + restored_job = restored_jobs[0] + assert restored_job.paused + assert isinstance(restored_job.source, HFModelSource) + assert restored_job.source.access_token == access_token + + captured: dict[str, str | None] = {} + + def _capture_resume(job: ModelInstallJob) -> None: + assert isinstance(job.source, HFModelSource) + captured["access_token"] = job.source.access_token + + monkeypatch.setattr(restored_installer, "_resume_remote_download", _capture_resume) + restored_installer.resume_job(restored_job) + assert captured["access_token"] == access_token + finally: + shutil.rmtree(tmpdir, ignore_errors=True) + + def test_404_download(mm2_installer: ModelInstallServiceBase, mm2_app_config: InvokeAIAppConfig) -> None: source = URLModelSource(url=Url("https://test.com/missing_model.safetensors")) job = mm2_installer.import_model(source)