fix(download): allow downloading additional files of an in-library model version (#1058)

This commit is contained in:
Will Miao
2026-08-19 16:29:59 +08:00
parent 0a28500848
commit cef4129fc9
9 changed files with 1092 additions and 83 deletions
+4 -2
View File
@@ -1659,7 +1659,8 @@ class ModelDownloadHandler:
import json
try:
data["file_params"] = json.loads(file_params_json)
# Normalize falsy payloads (e.g. {}) to None (#1058)
data["file_params"] = json.loads(file_params_json) or None
except json.JSONDecodeError:
self._logger.warning(
"Invalid file_params JSON: %s", file_params_json
@@ -1811,7 +1812,8 @@ class ModelDownloadHandler:
model_id = int(model_id_str) if model_id_str else None
model_version_id = int(model_version_id_str) if model_version_id_str else None
file_params = json.loads(file_params_json) if file_params_json else None
# Normalize falsy payloads (e.g. {}) to None (#1058)
file_params = (json.loads(file_params_json) if file_params_json else None) or None
service = await DownloadQueueService.get_instance()
item = await service.add_to_queue(
+3 -1
View File
@@ -87,7 +87,9 @@ class DownloadCoordinator:
progress_callback=progress_callback,
download_id=download_id,
source=payload.get("source"),
file_params=payload.get("file_params"),
# Normalize falsy file_params (e.g. {}) to None so download gates
# treat it as "no explicit file selection" (#1058).
file_params=payload.get("file_params") or None,
)
result["download_id"] = download_id
+223 -69
View File
@@ -213,6 +213,162 @@ class DownloadManager:
)
return False
async def _get_scanner_for_model_type(self, model_type: str):
"""Return the scanner responsible for the given model type."""
if model_type == "checkpoint":
return await self._get_checkpoint_scanner()
if model_type == "embedding":
return await ServiceRegistry.get_embedding_scanner()
return await self._get_lora_scanner()
@staticmethod
def _resolve_target_file(
files: Any, file_params: Dict[str, Any] | None
) -> Optional[Dict[str, Any]]:
"""Resolve the target file within a version's file list from file_params.
Shared by the existence gate and the actual file selection so both
always agree on which file a download refers to (#1058). Returns None
when file_params is None or no file matches.
"""
if not file_params or not isinstance(files, list):
return None
target_file_id = file_params.get("id")
target_type = file_params.get("type", "Model")
target_format = file_params.get("format")
target_size = file_params.get("size")
target_fp = file_params.get("fp")
is_primary = file_params.get("isPrimary", False)
logger.debug(
"[download] file_params received: id=%s, type=%s, format=%s, size=%s, fp=%s, "
"isPrimary=%s, total_files=%d",
target_file_id, target_type, target_format, target_size, target_fp,
is_primary, len(files),
)
file_info: Optional[Dict[str, Any]] = None
if target_file_id:
target_id_str = str(target_file_id)
for f in files:
if not isinstance(f, dict):
continue
f_id = f.get("id")
if str(f_id) == target_id_str:
file_info = f
logger.debug(
"[download] MATCH by ID: id=%s name='%s'",
f_id, f.get("name"),
)
break
if not file_info:
logger.debug("[download] No file found with id=%s", target_file_id)
elif is_primary:
file_info = next(
(
f
for f in files
if isinstance(f, dict)
and f.get("primary")
and f.get("type") in MODEL_WEIGHT_FILE_TYPES
),
None,
)
else:
# Lenient metadata match: only compare fields present on both sides
for f in files:
if not isinstance(f, dict):
continue
f_type = f.get("type", "")
if f_type != target_type:
continue
f_meta = f.get("metadata", {})
f_format = f_meta.get("format") or f.get("format")
f_size = f_meta.get("size") or f.get("size")
f_fp = f_meta.get("fp") or f.get("fp")
if target_format and f_format != target_format:
continue
if target_size and f_size and f_size != target_size:
continue
if target_fp and f_fp and f_fp != target_fp:
continue
file_info = f
break
return file_info
async def _find_local_file_entry(
self,
model_type: str,
model_version_id: int,
target_file: Dict[str, Any],
) -> Optional[Dict[str, Any]]:
"""Find a local library entry for a specific file of a model version.
Matches per design rule D2 (#1058): SHA256 is only compared when both
sides carry a non-empty hash; otherwise fall back to (extension-less)
file name equality. Never let two empty hashes compare equal.
"""
try:
normalized_version_id = int(model_version_id)
except (TypeError, ValueError):
return None
try:
scanner = await self._get_scanner_for_model_type(model_type)
cache = await scanner.get_cached_data()
except Exception as exc:
logger.debug(
"Failed to scan local entries for version %s file check: %s",
model_version_id,
exc,
)
return None
raw_data = getattr(cache, "raw_data", None) if cache else None
if not raw_data:
return None
target_hash = str(
(target_file.get("hashes") or {}).get("SHA256") or ""
).strip().lower()
target_name = str(target_file.get("name") or "").strip()
target_base = os.path.splitext(target_name)[0] if target_name else ""
for item in raw_data:
if not isinstance(item, dict):
continue
civitai_data = item.get("civitai")
if not isinstance(civitai_data, dict):
continue
try:
item_version_id = int(civitai_data.get("id"))
except (TypeError, ValueError):
continue
if item_version_id != normalized_version_id:
continue
local_hash = str(item.get("sha256") or "").strip().lower()
if target_hash and local_hash:
if local_hash == target_hash:
return item
# Both sides carry hashes that differ: this is a different
# file of the same version — do not fall back to name match.
continue
if target_base:
local_name = str(item.get("file_name") or "").strip()
if local_name == target_base:
return item
return None
async def download_from_civitai(
self,
model_id: int | None = None,
@@ -242,6 +398,10 @@ class DownloadManager:
Returns:
Dict with download result
"""
# Normalize falsy file_params (e.g. an empty dict from API JSON
# parsing) to None so gate conditions behave consistently (#1058).
file_params = file_params or None
logger.debug(
"[download] download_from_civitai called: model_id=%s, model_version_id=%s, "
"source=%s, file_params=%s",
@@ -1152,9 +1312,13 @@ class DownloadManager:
use_save_dir_as_root: bool = False,
) -> Dict[str, Any]:
"""Wrapper for original download_from_civitai implementation"""
file_params = file_params or None
try:
# Check if model version already exists in library
if model_version_id is not None:
# Check if model version already exists in library.
# With an explicit file selection (file_params) the version-level
# check is deferred until after the metadata fetch, when the target
# file can be resolved and checked individually (#1058).
if model_version_id is not None and file_params is None:
# Check both scanners
lora_scanner = await self._get_lora_scanner()
checkpoint_scanner = await self._get_checkpoint_scanner()
@@ -1235,8 +1399,26 @@ class DownloadManager:
except (TypeError, ValueError):
resolved_version_id = None
# Resolve the explicitly selected file (if any) up front so the
# existence gates and the actual file selection below always agree
# on the target file (#1058).
target_file: Optional[Dict[str, Any]] = None
if file_params is not None:
target_file = self._resolve_target_file(
version_info.get("files") or [], file_params
)
if target_file is None:
logger.warning(
"[download] file_params provided but no file matched; "
"falling back to version-level checks and primary file "
"selection (model_version_id=%s)",
resolved_version_id,
)
explicit_file = target_file is not None
if (
get_settings_manager().get_skip_previously_downloaded_model_versions()
not explicit_file
and get_settings_manager().get_skip_previously_downloaded_model_versions()
and resolved_version_id is not None
and await self._has_been_downloaded(model_type, resolved_version_id)
):
@@ -1346,9 +1528,38 @@ class DownloadManager:
f"baseModel '{base_model_value}' is a known diffusion model, routing to unet folder"
)
# Case 2: model_version_id was None, check after getting version_info
if model_version_id is None:
version_id = version_info.get("id")
# Existence check after the metadata fetch (#1058):
# - An explicit file selection only blocks when THIS file is
# already in the library; other files of the same version
# remain downloadable.
# - Without file_params (or when file_params failed to resolve),
# keep version-level protection. The case "model_version_id
# given + no file_params" was already covered by the early
# gate above.
if explicit_file and resolved_version_id is not None:
existing_entry = await self._find_local_file_entry(
model_type, resolved_version_id, target_file
)
if existing_entry is not None:
error_message = (
f"File '{target_file.get('name')}' from model version "
f"{resolved_version_id} already exists in {model_type} library"
)
logger.info("[download] %s", error_message)
return {"success": False, "error": error_message}
logger.info(
"[download] File '%s' of model version %s not in %s library — "
"download allowed (other files of this version may exist locally)",
target_file.get("name"), resolved_version_id, model_type,
)
elif file_params is not None or model_version_id is None:
# Case 2: model_version_id was None, or file_params did not
# resolve to a concrete file — check at version level.
version_id = (
resolved_version_id
if resolved_version_id is not None
else version_info.get("id")
)
if model_type == "lora":
# Check lora scanner
@@ -1495,73 +1706,16 @@ class DownloadManager:
files = version_info.get("files", [])
file_info = None
# If file_params is provided, try to find matching file
if file_params and model_version_id:
target_file_id = file_params.get("id")
target_type = file_params.get("type", "Model")
target_format = file_params.get("format")
target_size = file_params.get("size")
target_fp = file_params.get("fp")
is_primary = file_params.get("isPrimary", False)
logger.debug(
"[download] file_params received: id=%s, type=%s, format=%s, size=%s, fp=%s, isPrimary=%s, "
"model_version_id=%s, total_files=%d",
target_file_id, target_type, target_format, target_size, target_fp, is_primary,
model_version_id, len(files),
)
if target_file_id:
target_id_str = str(target_file_id)
for f in files:
f_id = f.get("id")
if str(f_id) == target_id_str:
file_info = f
logger.debug(
"[download] MATCH by ID: id=%s name='%s'",
f_id, f.get("name"),
)
break
if not file_info:
logger.debug("[download] No file found with id=%s", target_file_id)
elif is_primary:
file_info = next(
(
f
for f in files
if f.get("primary")
and f.get("type") in MODEL_WEIGHT_FILE_TYPES
),
None,
)
else:
# Lenient metadata match: only compare fields present on both sides
for f in files:
f_type = f.get("type", "")
if f_type != target_type:
continue
f_meta = f.get("metadata", {})
f_format = f_meta.get("format") or f.get("format")
f_size = f_meta.get("size") or f.get("size")
f_fp = f_meta.get("fp") or f.get("fp")
if target_format and f_format != target_format:
continue
if target_size and f_size and f_size != target_size:
continue
if target_fp and f_fp and f_fp != target_fp:
continue
file_info = f
break
# If file_params is provided, reuse the file resolved right after
# the metadata fetch so the existence gate and this selection
# always agree on the target file (#1058).
if file_params is not None:
file_info = target_file
if not file_info:
logger.debug(
"[download] No match found via file_params — falling back to primary file lookup",
)
elif not file_params:
else:
logger.debug(
"[download] No file_params provided (null/None) — will use primary file lookup. "
"model_version_id=%s, total_files=%d",
+34 -8
View File
@@ -64,6 +64,7 @@ class DownloadQueueService:
model_name TEXT NOT NULL DEFAULT '',
version_name TEXT DEFAULT '',
thumbnail_url TEXT DEFAULT '',
file_params TEXT,
status TEXT NOT NULL,
error TEXT,
file_path TEXT,
@@ -120,6 +121,18 @@ class DownloadQueueService:
with self._connect() as conn:
conn.executescript(self._SCHEMA_TABLES)
# Databases created by older versions lack
# download_history.file_params; add it so retry-from-history can
# restore the originally selected file (#1058).
history_columns = {
row["name"]
for row in conn.execute("PRAGMA table_info(download_history)")
}
if "file_params" not in history_columns:
conn.execute(
"ALTER TABLE download_history ADD COLUMN file_params TEXT"
)
# Creating the unique index on download_history.download_id can
# fail if pre-existing rows have duplicate values (e.g. from a
# previous version that lacked the index). Deduplicate first so
@@ -418,6 +431,12 @@ class DownloadQueueService:
return None
now = completed_at if completed_at is not None else time.time()
# Guard against legacy databases whose download_queue table
# predates the file_params column.
queue_columns = set(row.keys())
file_params_json = (
row["file_params"] if "file_params" in queue_columns else None
)
conn.execute(
"DELETE FROM download_queue WHERE download_id = ?",
(download_id,),
@@ -426,9 +445,9 @@ class DownloadQueueService:
"""
INSERT OR IGNORE INTO download_history (
download_id, model_id, model_version_id, model_name,
version_name, thumbnail_url, status, error, file_path,
bytes_downloaded, total_bytes, completed_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
version_name, thumbnail_url, file_params, status, error,
file_path, bytes_downloaded, total_bytes, completed_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
row["download_id"],
@@ -437,6 +456,7 @@ class DownloadQueueService:
row["model_name"],
row["version_name"],
row["thumbnail_url"],
file_params_json,
status,
error,
file_path,
@@ -503,6 +523,7 @@ class DownloadQueueService:
bytes_downloaded: int = 0,
total_bytes: Optional[int] = None,
is_already_exists: int = 0,
file_params: Optional[dict[str, Any]] = None,
) -> int:
"""Insert a record into the download history.
@@ -510,6 +531,7 @@ class DownloadQueueService:
inserted row.
"""
now = time.time()
file_params_json = json.dumps(file_params) if file_params is not None else None
async with self._lock:
conn = self._get_conn()
@@ -517,9 +539,10 @@ class DownloadQueueService:
"""
INSERT INTO download_history (
download_id, model_id, model_version_id, model_name,
version_name, thumbnail_url, status, error, file_path,
bytes_downloaded, total_bytes, completed_at, is_already_exists
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
version_name, thumbnail_url, file_params, status, error,
file_path, bytes_downloaded, total_bytes, completed_at,
is_already_exists
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
download_id,
@@ -528,6 +551,7 @@ class DownloadQueueService:
model_name,
version_name,
thumbnail_url,
file_params_json,
status,
error,
file_path,
@@ -702,7 +726,7 @@ class DownloadQueueService:
download_id, model_id, model_version_id, model_name,
version_name, thumbnail_url, source, file_params,
status, priority, added_at
) VALUES (?, ?, ?, ?, ?, ?, ?, NULL, 'queued', 0, ?)
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'queued', 0, ?)
""",
(
new_id,
@@ -712,6 +736,7 @@ class DownloadQueueService:
row["version_name"],
row["thumbnail_url"],
"retry",
row["file_params"],
now,
),
)
@@ -755,7 +780,7 @@ class DownloadQueueService:
download_id, model_id, model_version_id, model_name,
version_name, thumbnail_url, source, file_params,
status, priority, added_at
) VALUES (?, ?, ?, ?, ?, ?, ?, NULL, 'queued', 0, ?)
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'queued', 0, ?)
""",
(
new_id,
@@ -765,6 +790,7 @@ class DownloadQueueService:
row["version_name"],
row["thumbnail_url"],
"retry",
row["file_params"],
now,
),
)