mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-25 07:05:43 -03:00
feat: propagate version info to model update record creation. Fixes #756
- Pass `version_info` parameter through download manager to model update service - Enhance `_create_record` to use version info when creating records for missing versions - Add `_extract_single_version` helper method for consistent version extraction - Improve handling of version metadata during library synchronization
This commit is contained in:
@@ -487,6 +487,7 @@ class DownloadManager:
|
|||||||
model_type,
|
model_type,
|
||||||
resolved_model_id,
|
resolved_model_id,
|
||||||
sorted(version_ids),
|
sorted(version_ids),
|
||||||
|
version_info=version_info,
|
||||||
)
|
)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
logger.debug(
|
logger.debug(
|
||||||
|
|||||||
@@ -589,6 +589,8 @@ class ModelUpdateService:
|
|||||||
model_type: str,
|
model_type: str,
|
||||||
model_id: int,
|
model_id: int,
|
||||||
version_ids: Sequence[int],
|
version_ids: Sequence[int],
|
||||||
|
*,
|
||||||
|
version_info: Optional[Mapping] = None,
|
||||||
) -> ModelUpdateRecord:
|
) -> ModelUpdateRecord:
|
||||||
"""Persist a new set of in-library version identifiers."""
|
"""Persist a new set of in-library version identifiers."""
|
||||||
|
|
||||||
@@ -600,6 +602,7 @@ class ModelUpdateService:
|
|||||||
normalized_versions,
|
normalized_versions,
|
||||||
model_type=model_type,
|
model_type=model_type,
|
||||||
model_id=model_id,
|
model_id=model_id,
|
||||||
|
version_info=version_info,
|
||||||
)
|
)
|
||||||
self._upsert_record(record)
|
self._upsert_record(record)
|
||||||
return record
|
return record
|
||||||
@@ -944,6 +947,7 @@ class ModelUpdateService:
|
|||||||
model_type: Optional[str] = None,
|
model_type: Optional[str] = None,
|
||||||
model_id: Optional[int] = None,
|
model_id: Optional[int] = None,
|
||||||
last_checked_at: Optional[float] = None,
|
last_checked_at: Optional[float] = None,
|
||||||
|
version_info: Optional[Mapping] = None,
|
||||||
) -> ModelUpdateRecord:
|
) -> ModelUpdateRecord:
|
||||||
local_set = set(normalized_local)
|
local_set = set(normalized_local)
|
||||||
versions: List[ModelVersionRecord] = []
|
versions: List[ModelVersionRecord] = []
|
||||||
@@ -965,19 +969,26 @@ class ModelUpdateService:
|
|||||||
|
|
||||||
seen_ids = {version.version_id for version in versions}
|
seen_ids = {version.version_id for version in versions}
|
||||||
for missing_id in sorted(local_set - seen_ids):
|
for missing_id in sorted(local_set - seen_ids):
|
||||||
versions.append(
|
new_version: Optional[ModelVersionRecord] = None
|
||||||
ModelVersionRecord(
|
if version_info and _normalize_int(version_info.get("id")) == missing_id:
|
||||||
version_id=missing_id,
|
new_version = self._extract_single_version(version_info, index=len(versions))
|
||||||
name=None,
|
|
||||||
base_model=None,
|
if new_version:
|
||||||
released_at=None,
|
versions.append(replace(new_version, is_in_library=True))
|
||||||
size_bytes=None,
|
else:
|
||||||
preview_url=None,
|
versions.append(
|
||||||
is_in_library=True,
|
ModelVersionRecord(
|
||||||
should_ignore=ignore_map.get(missing_id, False),
|
version_id=missing_id,
|
||||||
sort_index=len(versions),
|
name=None,
|
||||||
|
base_model=None,
|
||||||
|
released_at=None,
|
||||||
|
size_bytes=None,
|
||||||
|
preview_url=None,
|
||||||
|
is_in_library=True,
|
||||||
|
should_ignore=ignore_map.get(missing_id, False),
|
||||||
|
sort_index=len(versions),
|
||||||
|
)
|
||||||
)
|
)
|
||||||
)
|
|
||||||
|
|
||||||
return ModelUpdateRecord(
|
return ModelUpdateRecord(
|
||||||
model_type=model_type,
|
model_type=model_type,
|
||||||
@@ -1083,33 +1094,45 @@ class ModelUpdateService:
|
|||||||
return []
|
return []
|
||||||
if not isinstance(versions, Iterable):
|
if not isinstance(versions, Iterable):
|
||||||
return None
|
return None
|
||||||
|
|
||||||
extracted: List[ModelVersionRecord] = []
|
extracted: List[ModelVersionRecord] = []
|
||||||
for index, entry in enumerate(versions):
|
for index, entry in enumerate(versions):
|
||||||
if not isinstance(entry, Mapping):
|
version_record = self._extract_single_version(entry, index)
|
||||||
continue
|
if version_record:
|
||||||
version_id = _normalize_int(entry.get("id"))
|
extracted.append(version_record)
|
||||||
if version_id is None:
|
|
||||||
continue
|
|
||||||
name = _normalize_string(entry.get("name"))
|
|
||||||
base_model = _normalize_string(entry.get("baseModel"))
|
|
||||||
released_at = _normalize_string(entry.get("publishedAt") or entry.get("createdAt"))
|
|
||||||
size_bytes = self._extract_size_bytes(entry.get("files"))
|
|
||||||
preview_url = self._extract_preview_url(entry.get("images"))
|
|
||||||
extracted.append(
|
|
||||||
ModelVersionRecord(
|
|
||||||
version_id=version_id,
|
|
||||||
name=name,
|
|
||||||
base_model=base_model,
|
|
||||||
released_at=released_at,
|
|
||||||
size_bytes=size_bytes,
|
|
||||||
preview_url=preview_url,
|
|
||||||
is_in_library=False,
|
|
||||||
should_ignore=False,
|
|
||||||
sort_index=index,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
return extracted
|
return extracted
|
||||||
|
|
||||||
|
def _extract_single_version(
|
||||||
|
self, entry: Any, index: int = 0
|
||||||
|
) -> Optional[ModelVersionRecord]:
|
||||||
|
"""Convert a raw metadata entry into a structured record."""
|
||||||
|
|
||||||
|
if not isinstance(entry, Mapping):
|
||||||
|
return None
|
||||||
|
|
||||||
|
version_id = _normalize_int(entry.get("id"))
|
||||||
|
if version_id is None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
name = _normalize_string(entry.get("name"))
|
||||||
|
base_model = _normalize_string(entry.get("baseModel"))
|
||||||
|
released_at = _normalize_string(entry.get("publishedAt") or entry.get("createdAt"))
|
||||||
|
size_bytes = self._extract_size_bytes(entry.get("files"))
|
||||||
|
preview_url = self._extract_preview_url(entry.get("images"))
|
||||||
|
|
||||||
|
return ModelVersionRecord(
|
||||||
|
version_id=version_id,
|
||||||
|
name=name,
|
||||||
|
base_model=base_model,
|
||||||
|
released_at=released_at,
|
||||||
|
size_bytes=size_bytes,
|
||||||
|
preview_url=preview_url,
|
||||||
|
is_in_library=False,
|
||||||
|
should_ignore=False,
|
||||||
|
sort_index=index,
|
||||||
|
)
|
||||||
|
|
||||||
def _extract_size_bytes(self, files) -> Optional[int]:
|
def _extract_size_bytes(self, files) -> Optional[int]:
|
||||||
if not isinstance(files, Iterable):
|
if not isinstance(files, Iterable):
|
||||||
return None
|
return None
|
||||||
|
|||||||
@@ -515,7 +515,31 @@ async def test_refresh_rewrites_remote_preview_urls(tmp_path):
|
|||||||
assert record is not None
|
assert record is not None
|
||||||
assert record.versions
|
assert record.versions
|
||||||
preview_url = record.versions[0].preview_url
|
preview_url = record.versions[0].preview_url
|
||||||
assert (
|
|
||||||
preview_url
|
@pytest.mark.asyncio
|
||||||
== "https://image.civitai.com/safe/width=450,optimized=true/preview.png"
|
async def test_update_in_library_versions_populates_metadata(tmp_path):
|
||||||
)
|
db_path = tmp_path / "updates.sqlite"
|
||||||
|
service = ModelUpdateService(str(db_path))
|
||||||
|
|
||||||
|
version_info = {
|
||||||
|
"id": 123,
|
||||||
|
"name": "v1.0",
|
||||||
|
"baseModel": "SD 1.5",
|
||||||
|
"publishedAt": "2024-03-01T00:00:00Z",
|
||||||
|
"files": [{"sizeKB": 1024, "type": "Model", "primary": True}],
|
||||||
|
"images": [{"url": "https://example.com/preview.png"}],
|
||||||
|
}
|
||||||
|
|
||||||
|
await service.update_in_library_versions("lora", 1, [123], version_info=version_info)
|
||||||
|
record = await service.get_record("lora", 1)
|
||||||
|
|
||||||
|
assert record is not None
|
||||||
|
assert len(record.versions) == 1
|
||||||
|
version = record.versions[0]
|
||||||
|
assert version.version_id == 123
|
||||||
|
assert version.name == "v1.0"
|
||||||
|
assert version.base_model == "SD 1.5"
|
||||||
|
assert version.size_bytes == 1024 * 1024
|
||||||
|
assert version.preview_url == "https://example.com/preview.png"
|
||||||
|
assert version.is_in_library is True
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user