refactor: centralize background task handling and deduplicate title matching
- Add run_in_background_thread() in __init__.py to unify success/cancel/error handling for hash calculation and local metadata build progress - Extract match_model_by_title_generic() in core/matcher.py and use it from both CivitAI and CivArchive providers to remove duplicated version/file selection logic - Enable safetensors header SHA256 extraction by default and report hash source through progress updates - Update options_methods.js to match backend progress behavior - Add tests/test_title_matching_hashing.py covering title matching and hashing - Minor metadata updates
This commit is contained in:
Dariusz L
+209
-154
@@ -416,6 +416,57 @@ class ModelResolverExtension:
|
||||
"progress_id": job_id,
|
||||
})
|
||||
|
||||
def run_in_background_thread(
|
||||
tracker,
|
||||
progress_id,
|
||||
task_func,
|
||||
on_success,
|
||||
on_cancel=None,
|
||||
on_error=None,
|
||||
error_log_msg="Background task failed",
|
||||
):
|
||||
def wrapper():
|
||||
try:
|
||||
result = task_func()
|
||||
if tracker.is_cancelled(progress_id):
|
||||
if on_cancel:
|
||||
on_cancel(result)
|
||||
else:
|
||||
tracker.update(
|
||||
progress_id,
|
||||
status="cancelled",
|
||||
stage="cancelled",
|
||||
message="Task cancelled",
|
||||
)
|
||||
return
|
||||
on_success(result)
|
||||
except (HashCalculationCancelled, asyncio.CancelledError):
|
||||
if on_cancel:
|
||||
on_cancel()
|
||||
else:
|
||||
tracker.update(
|
||||
progress_id,
|
||||
status="cancelled",
|
||||
stage="cancelled",
|
||||
message="Task cancelled",
|
||||
)
|
||||
except Exception as exc:
|
||||
self.logger.exception(f"{error_log_msg}: {exc}")
|
||||
if on_error:
|
||||
on_error(exc)
|
||||
else:
|
||||
tracker.update(
|
||||
progress_id,
|
||||
status="error",
|
||||
stage="error",
|
||||
message=str(exc) or "Task failed",
|
||||
percent=100,
|
||||
error=str(exc) or "Task failed",
|
||||
)
|
||||
|
||||
import threading
|
||||
threading.Thread(target=wrapper, daemon=True).start()
|
||||
|
||||
async def get_override_settings_from_request(request):
|
||||
settings = await asyncio.to_thread(load_resolver_settings)
|
||||
if request.method == "POST":
|
||||
@@ -1035,57 +1086,29 @@ class ModelResolverExtension:
|
||||
if self.hash_tracker.is_cancelled(progress_id):
|
||||
raise HashCalculationCancelled()
|
||||
|
||||
is_safetensors = normalized_path.lower().endswith(".safetensors")
|
||||
if is_safetensors:
|
||||
self.logger.info(
|
||||
f"Checking safetensors header for SHA256: {normalized_path}"
|
||||
)
|
||||
else:
|
||||
self.logger.info(
|
||||
"Skipping safetensors header lookup for non-safetensors file; "
|
||||
f"calculating full SHA256: {normalized_path}"
|
||||
)
|
||||
|
||||
header_sha256 = extract_safetensors_header_sha256(normalized_path)
|
||||
if header_sha256:
|
||||
self.logger.info(
|
||||
f"Found SHA256 in safetensors header: {normalized_path}"
|
||||
)
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="header",
|
||||
message="Found SHA256 in safetensors header",
|
||||
percent=98,
|
||||
bytes_read=0,
|
||||
total_bytes=total_bytes,
|
||||
sha256=header_sha256,
|
||||
hash=header_sha256,
|
||||
sha256_source="safetensors_header",
|
||||
)
|
||||
return header_sha256, "safetensors_header"
|
||||
|
||||
if is_safetensors:
|
||||
self.logger.info(
|
||||
"No SHA256 found in safetensors header; calculating full "
|
||||
f"file SHA256: {normalized_path}"
|
||||
)
|
||||
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="hashing",
|
||||
message="Calculating SHA256...",
|
||||
percent=0,
|
||||
bytes_read=0,
|
||||
total_bytes=total_bytes,
|
||||
)
|
||||
|
||||
hash_source = ["file"]
|
||||
stage_transitioned = [False]
|
||||
_bytes_read = [0]
|
||||
_last_update = [0.0]
|
||||
|
||||
def on_progress(bytes_read, total_bytes):
|
||||
_bytes_read[0] = bytes_read
|
||||
if not stage_transitioned[0]:
|
||||
stage_transitioned[0] = True
|
||||
self.logger.info(
|
||||
"No SHA256 found in safetensors header; calculating full "
|
||||
f"file SHA256: {normalized_path}"
|
||||
)
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="hashing",
|
||||
message="Calculating SHA256...",
|
||||
percent=0,
|
||||
bytes_read=0,
|
||||
total_bytes=total_bytes,
|
||||
)
|
||||
|
||||
now = time.time()
|
||||
if now - _last_update[0] >= 0.15 or bytes_read >= total_bytes:
|
||||
percent = 98 if total_bytes <= 0 else min(
|
||||
@@ -1121,14 +1144,36 @@ class ModelResolverExtension:
|
||||
return True
|
||||
return False
|
||||
|
||||
def on_hash_source(source_name):
|
||||
hash_source[0] = source_name
|
||||
if source_name == "safetensors_header":
|
||||
self.logger.info(
|
||||
f"Found SHA256 in safetensors header: {normalized_path}"
|
||||
)
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="header",
|
||||
message="Found SHA256 in safetensors header",
|
||||
percent=98,
|
||||
bytes_read=0,
|
||||
total_bytes=total_bytes,
|
||||
sha256_source="safetensors_header",
|
||||
)
|
||||
|
||||
sha256 = _calculate_file_sha256_core(
|
||||
normalized_path,
|
||||
chunk_size=1024 * 1024 * 4,
|
||||
on_progress=on_progress,
|
||||
is_cancelled=is_cancelled,
|
||||
use_safetensors_header=False,
|
||||
use_safetensors_header=True,
|
||||
on_hash_source=on_hash_source,
|
||||
)
|
||||
return sha256, "file"
|
||||
|
||||
if is_cancelled():
|
||||
raise HashCalculationCancelled()
|
||||
|
||||
return sha256, hash_source[0]
|
||||
|
||||
@routes.post("/model_resolver/calculate-file-hash")
|
||||
@json_api_endpoint("calculate-file-hash")
|
||||
@@ -1200,57 +1245,58 @@ class ModelResolverExtension:
|
||||
)
|
||||
|
||||
def run_hash_task():
|
||||
try:
|
||||
sha256, sha256_source = calculate_sha256_with_progress(
|
||||
normalized_path,
|
||||
progress_id=progress_id,
|
||||
)
|
||||
if self.hash_tracker.is_cancelled(progress_id):
|
||||
raise HashCalculationCancelled()
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="metadata",
|
||||
message="Saving metadata...",
|
||||
percent=99,
|
||||
)
|
||||
resolved_metadata_path, metadata_updated = write_calculated_hash_metadata(
|
||||
normalized_path,
|
||||
sha256,
|
||||
sha256_source,
|
||||
)
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="done",
|
||||
stage="done",
|
||||
message="Hash calculated",
|
||||
percent=100,
|
||||
sha256=sha256,
|
||||
hash=sha256,
|
||||
sha256_source=sha256_source,
|
||||
file_path=normalized_path,
|
||||
metadata_path=resolved_metadata_path,
|
||||
metadata_updated=metadata_updated,
|
||||
)
|
||||
except HashCalculationCancelled:
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="cancelled",
|
||||
stage="cancelled",
|
||||
message="Hash calculation cancelled",
|
||||
)
|
||||
except Exception as exc:
|
||||
self.logger.exception(f"Hash calculation failed for {normalized_path}: {exc}")
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="error",
|
||||
stage="error",
|
||||
message=str(exc) or "Hash calculation failed",
|
||||
percent=100,
|
||||
error=str(exc) or "Hash calculation failed",
|
||||
)
|
||||
sha256, sha256_source = calculate_sha256_with_progress(
|
||||
normalized_path,
|
||||
progress_id=progress_id,
|
||||
)
|
||||
if self.hash_tracker.is_cancelled(progress_id):
|
||||
raise HashCalculationCancelled()
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="metadata",
|
||||
message="Saving metadata...",
|
||||
percent=99,
|
||||
)
|
||||
resolved_metadata_path, metadata_updated = write_calculated_hash_metadata(
|
||||
normalized_path,
|
||||
sha256,
|
||||
sha256_source,
|
||||
)
|
||||
return sha256, sha256_source, resolved_metadata_path, metadata_updated
|
||||
|
||||
threading.Thread(target=run_hash_task, daemon=True).start()
|
||||
def on_success(res):
|
||||
sha256, sha256_source, resolved_metadata_path, metadata_updated = res
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="done",
|
||||
stage="done",
|
||||
message="Hash calculated",
|
||||
percent=100,
|
||||
sha256=sha256,
|
||||
hash=sha256,
|
||||
sha256_source=sha256_source,
|
||||
file_path=normalized_path,
|
||||
metadata_path=resolved_metadata_path,
|
||||
metadata_updated=metadata_updated,
|
||||
)
|
||||
|
||||
def on_cancel(res=None):
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="cancelled",
|
||||
stage="cancelled",
|
||||
message="Hash calculation cancelled",
|
||||
)
|
||||
|
||||
run_in_background_thread(
|
||||
self.hash_tracker,
|
||||
progress_id,
|
||||
run_hash_task,
|
||||
on_success,
|
||||
on_cancel,
|
||||
error_log_msg=f"Hash calculation failed for {normalized_path}",
|
||||
)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
@@ -1372,64 +1418,73 @@ class ModelResolverExtension:
|
||||
return self.metadata_builder_progress.is_cancelled(progress_id)
|
||||
|
||||
def run_metadata_build_task():
|
||||
try:
|
||||
result = build_missing_local_metadata(
|
||||
force_rescan=force_rescan,
|
||||
worker_count=worker_count,
|
||||
progress_callback=update_metadata_build_progress,
|
||||
is_cancelled=is_metadata_build_cancelled,
|
||||
)
|
||||
if result.get("cancelled") or is_metadata_build_cancelled():
|
||||
self.metadata_builder_progress.update(
|
||||
progress_id,
|
||||
status="cancelled",
|
||||
stage="cancelled",
|
||||
message="Metadata build cancelled",
|
||||
percent=100,
|
||||
active_models=[],
|
||||
active_worker_count=0,
|
||||
current_model="",
|
||||
current_path="",
|
||||
bytes_read=0,
|
||||
total_bytes=0,
|
||||
result=result,
|
||||
**result,
|
||||
)
|
||||
return
|
||||
result = build_missing_local_metadata(
|
||||
force_rescan=force_rescan,
|
||||
worker_count=worker_count,
|
||||
progress_callback=update_metadata_build_progress,
|
||||
is_cancelled=is_metadata_build_cancelled,
|
||||
)
|
||||
return result
|
||||
|
||||
self.metadata_builder_progress.update(
|
||||
progress_id,
|
||||
status="done",
|
||||
stage="done",
|
||||
message="Local metadata build completed.",
|
||||
percent=100,
|
||||
active_models=[],
|
||||
active_worker_count=0,
|
||||
current_model="",
|
||||
current_path="",
|
||||
bytes_read=0,
|
||||
total_bytes=0,
|
||||
result=result,
|
||||
**result,
|
||||
)
|
||||
except Exception as exc:
|
||||
self.logger.exception(f"Metadata build failed: {exc}")
|
||||
self.metadata_builder_progress.update(
|
||||
progress_id,
|
||||
status="error",
|
||||
stage="error",
|
||||
message=str(exc) or "Metadata build failed",
|
||||
percent=100,
|
||||
active_models=[],
|
||||
active_worker_count=0,
|
||||
current_model="",
|
||||
current_path="",
|
||||
bytes_read=0,
|
||||
total_bytes=0,
|
||||
error=str(exc) or "Metadata build failed",
|
||||
)
|
||||
def on_success(result):
|
||||
self.metadata_builder_progress.update(
|
||||
progress_id,
|
||||
status="done",
|
||||
stage="done",
|
||||
message="Local metadata build completed.",
|
||||
percent=100,
|
||||
active_models=[],
|
||||
active_worker_count=0,
|
||||
current_model="",
|
||||
current_path="",
|
||||
bytes_read=0,
|
||||
total_bytes=0,
|
||||
result=result,
|
||||
**result,
|
||||
)
|
||||
|
||||
threading.Thread(target=run_metadata_build_task, daemon=True).start()
|
||||
def on_cancel(result=None):
|
||||
self.metadata_builder_progress.update(
|
||||
progress_id,
|
||||
status="cancelled",
|
||||
stage="cancelled",
|
||||
message="Metadata build cancelled",
|
||||
percent=100,
|
||||
active_models=[],
|
||||
active_worker_count=0,
|
||||
current_model="",
|
||||
current_path="",
|
||||
bytes_read=0,
|
||||
total_bytes=0,
|
||||
result=result or {},
|
||||
**(result or {}),
|
||||
)
|
||||
|
||||
def on_error(exc):
|
||||
self.metadata_builder_progress.update(
|
||||
progress_id,
|
||||
status="error",
|
||||
stage="error",
|
||||
message=str(exc) or "Metadata build failed",
|
||||
percent=100,
|
||||
active_models=[],
|
||||
active_worker_count=0,
|
||||
current_model="",
|
||||
current_path="",
|
||||
bytes_read=0,
|
||||
total_bytes=0,
|
||||
error=str(exc) or "Metadata build failed",
|
||||
)
|
||||
|
||||
run_in_background_thread(
|
||||
self.metadata_builder_progress,
|
||||
progress_id,
|
||||
run_metadata_build_task,
|
||||
on_success,
|
||||
on_cancel,
|
||||
on_error,
|
||||
error_log_msg="Metadata build failed",
|
||||
)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
|
||||
+67
-2
@@ -7,7 +7,7 @@ Implements fuzzy string matching to find similar model names.
|
||||
import os
|
||||
import re
|
||||
import heapq
|
||||
from typing import List, Dict, Tuple, Optional
|
||||
from typing import List, Dict, Tuple, Optional, Any
|
||||
from difflib import SequenceMatcher
|
||||
|
||||
from .log_system import create_module_logger
|
||||
@@ -287,7 +287,7 @@ def base_model_score(candidate: str, preferred: Optional[str]) -> float:
|
||||
|
||||
# ==================== CENTRALIZED MATCHING & CONFIDENCE HELPERS ====================
|
||||
|
||||
from .type_utils import MODEL_EXTENSIONS, PRECISION_FORMAT_SUFFIXES
|
||||
from .type_utils import MODEL_EXTENSIONS, PRECISION_FORMAT_SUFFIXES, get_version_sort_key
|
||||
MODEL_FILE_EXTENSIONS = MODEL_EXTENSIONS
|
||||
|
||||
|
||||
@@ -437,3 +437,68 @@ def build_filename_search_queries(filename: str) -> List[str]:
|
||||
queries.append(query_val)
|
||||
|
||||
return queries
|
||||
|
||||
|
||||
def match_model_by_title_generic(
|
||||
model_id: int,
|
||||
title_query: str,
|
||||
model_name: str,
|
||||
versions: list,
|
||||
base_model_context: Optional[str],
|
||||
get_base_model_fn,
|
||||
select_file_fn,
|
||||
build_result_fn,
|
||||
hydrate_version_fn=None,
|
||||
log_prefix: str = "ModelResolver",
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Generic helper to resolve extensionless workflow values by model title across CivitAI and CivArchive."""
|
||||
from typing import Callable, Any
|
||||
|
||||
title_confidence = calculate_model_title_confidence(title_query, model_name)
|
||||
if title_confidence < MODEL_TITLE_MATCH_THRESHOLD:
|
||||
log.debug(
|
||||
f"{log_prefix} title candidate rejected: model_id={model_id}, query={title_query}, model_name={model_name}, confidence={title_confidence}"
|
||||
)
|
||||
return None
|
||||
|
||||
if not versions:
|
||||
return None
|
||||
|
||||
# Sort versions
|
||||
versions = sorted(versions, key=get_version_sort_key, reverse=True)
|
||||
rejected_by_base_model = False
|
||||
|
||||
for version in versions:
|
||||
hydrated_version = hydrate_version_fn(version) if hydrate_version_fn else version
|
||||
|
||||
# Check base model context
|
||||
if base_model_context:
|
||||
base_model = get_base_model_fn(hydrated_version)
|
||||
if not base_model_matches(base_model, base_model_context):
|
||||
rejected_by_base_model = True
|
||||
continue
|
||||
|
||||
# Select file
|
||||
file_info = select_file_fn(hydrated_version)
|
||||
if not file_info:
|
||||
continue
|
||||
|
||||
# Build result
|
||||
result = build_result_fn(hydrated_version, file_info, title_confidence)
|
||||
if not result:
|
||||
continue
|
||||
|
||||
result["confidence"] = title_confidence
|
||||
result["title_confidence"] = title_confidence
|
||||
|
||||
log.info(
|
||||
f"{log_prefix} model-title match: query={title_query}, model_id={model_id}, model_name={model_name}, version_id={result.get('version_id')}, filename={result.get('filename')}, confidence={title_confidence}, base={result.get('base_model')}"
|
||||
)
|
||||
return result
|
||||
|
||||
if rejected_by_base_model:
|
||||
log.debug(
|
||||
f"{log_prefix} title candidate rejected by base model: model_id={model_id}, base={base_model_context}"
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
+23
-39
@@ -1215,14 +1215,9 @@ def _find_model_title_match_in_model_details(
|
||||
base_model_context: Optional[str] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""For extensionless workflow values, resolve by CivArchive model page title."""
|
||||
model_name = model_details.get("name", "")
|
||||
title_confidence = calculate_model_title_confidence(title_query, model_name)
|
||||
if title_confidence < MODEL_TITLE_MATCH_THRESHOLD:
|
||||
log.debug(
|
||||
f"CivArchive title candidate rejected: model_id={model_id}, query={title_query}, model_name={model_name}, confidence={title_confidence}"
|
||||
)
|
||||
return None
|
||||
from ..matcher import match_model_by_title_generic
|
||||
|
||||
model_name = model_details.get("name", "")
|
||||
versions = [
|
||||
version
|
||||
for version in (model_details.get("versions") or [])
|
||||
@@ -1234,46 +1229,35 @@ def _find_model_title_match_in_model_details(
|
||||
if not any(_coerce_int(version.get("id")) == selected_id for version in versions):
|
||||
versions.append(selected_version)
|
||||
|
||||
if not versions:
|
||||
return None
|
||||
def hydrate_version(v):
|
||||
return _hydrate_civarchive_version_with_files(model_id, v)
|
||||
|
||||
versions = sorted(versions, key=get_version_sort_key, reverse=True)
|
||||
rejected_by_base_model = False
|
||||
def get_base_model(v):
|
||||
return v.get("base_model")
|
||||
|
||||
for version in versions:
|
||||
hydrated_version = _hydrate_civarchive_version_with_files(model_id, version)
|
||||
if base_model_context and not _base_model_matches(
|
||||
hydrated_version.get("base_model"), base_model_context
|
||||
):
|
||||
rejected_by_base_model = True
|
||||
continue
|
||||
def select_file(v):
|
||||
return _select_primary_model_file(v.get("files") or [])
|
||||
|
||||
file_info = _select_primary_model_file(hydrated_version.get("files") or [])
|
||||
if not file_info:
|
||||
continue
|
||||
|
||||
result = _build_result_from_normalized_version(
|
||||
def build_result(v, file_info, confidence):
|
||||
return _build_result_from_normalized_version(
|
||||
model_details=model_details,
|
||||
version=hydrated_version,
|
||||
version=v,
|
||||
file_info=file_info,
|
||||
match_type="model_title",
|
||||
)
|
||||
if not result:
|
||||
continue
|
||||
|
||||
result["confidence"] = title_confidence
|
||||
result["title_confidence"] = title_confidence
|
||||
log.info(
|
||||
f"CivArchive model-title match: query={title_query}, model_id={model_id}, model_name={model_name}, version_id={result.get('version_id')}, filename={result.get('filename')}, confidence={title_confidence}, base={result.get('base_model')}"
|
||||
)
|
||||
return result
|
||||
|
||||
if rejected_by_base_model:
|
||||
log.debug(
|
||||
f"CivArchive title candidate rejected by base model: model_id={model_id}, base={base_model_context}"
|
||||
)
|
||||
|
||||
return None
|
||||
return match_model_by_title_generic(
|
||||
model_id=model_id,
|
||||
title_query=title_query,
|
||||
model_name=model_name,
|
||||
versions=versions,
|
||||
base_model_context=base_model_context,
|
||||
get_base_model_fn=get_base_model,
|
||||
select_file_fn=select_file,
|
||||
build_result_fn=build_result,
|
||||
hydrate_version_fn=hydrate_version,
|
||||
log_prefix="CivArchive",
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
+23
-36
@@ -190,59 +190,46 @@ def _find_model_title_match_in_model(
|
||||
base_model_context: Optional[str] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""For extensionless workflow values, resolve by CivitAI model page title."""
|
||||
model_name = model_data.get("name", "")
|
||||
title_confidence = calculate_model_title_confidence(title_query, model_name)
|
||||
if title_confidence < MODEL_TITLE_MATCH_THRESHOLD:
|
||||
log.debug(
|
||||
f"CivitAI title candidate rejected: model_id={model_id}, query={title_query}, model_name={model_name}, confidence={title_confidence}"
|
||||
)
|
||||
return None
|
||||
from ..matcher import match_model_by_title_generic
|
||||
|
||||
model_name = model_data.get("name", "")
|
||||
versions = [
|
||||
version
|
||||
for version in (model_data.get("modelVersions") or [])
|
||||
if isinstance(version, dict)
|
||||
]
|
||||
if not versions:
|
||||
return None
|
||||
|
||||
if base_model_context:
|
||||
versions = [
|
||||
version
|
||||
for version in versions
|
||||
if _base_model_matches(version.get("baseModel"), base_model_context)
|
||||
]
|
||||
if not versions:
|
||||
log.debug(
|
||||
f"CivitAI title candidate rejected by base model: model_id={model_id}, base={base_model_context}"
|
||||
)
|
||||
return None
|
||||
def get_base_model(v):
|
||||
return v.get("baseModel")
|
||||
|
||||
versions = sorted(versions, key=get_version_sort_key, reverse=True)
|
||||
def select_file(v):
|
||||
return select_primary_model_file(v.get("files") or [])
|
||||
|
||||
for version in versions:
|
||||
file_info = select_primary_model_file(version.get("files") or [])
|
||||
if not file_info:
|
||||
continue
|
||||
|
||||
result = _build_civitai_result_from_version(
|
||||
def build_result(v, file_info, confidence):
|
||||
res = _build_civitai_result_from_version(
|
||||
model_id=model_id,
|
||||
model_name=model_name,
|
||||
model_type=model_data.get("type", ""),
|
||||
version=version,
|
||||
version=v,
|
||||
file_info=file_info,
|
||||
tags=model_data.get("tags", []),
|
||||
match_type="model_title",
|
||||
)
|
||||
result["confidence"] = title_confidence
|
||||
result["title_confidence"] = title_confidence
|
||||
result["version_name"] = version.get("name", "")
|
||||
log.info(
|
||||
f"CivitAI model-title match: query={title_query}, model_id={model_id}, model_name={model_name}, version_id={result.get('version_id')}, filename={result.get('filename')}, confidence={title_confidence}, base={result.get('base_model')}"
|
||||
)
|
||||
return result
|
||||
if res:
|
||||
res["version_name"] = v.get("name", "")
|
||||
return res
|
||||
|
||||
return None
|
||||
return match_model_by_title_generic(
|
||||
model_id=model_id,
|
||||
title_query=title_query,
|
||||
model_name=model_name,
|
||||
versions=versions,
|
||||
base_model_context=base_model_context,
|
||||
get_base_model_fn=get_base_model,
|
||||
select_file_fn=select_file,
|
||||
build_result_fn=build_result,
|
||||
log_prefix="CivitAI",
|
||||
)
|
||||
|
||||
|
||||
def _filename_base_partial_match(target_base: str, candidate_base: str) -> bool:
|
||||
|
||||
@@ -1,10 +1,7 @@
|
||||
{
|
||||
"updated_at": "2026-07-12T19:07:42+00:00",
|
||||
"last_checked_at": "2026-07-12T19:07:42+00:00",
|
||||
"updated_at": "2026-07-15T11:23:31+00:00",
|
||||
"last_checked_at": "2026-07-15T11:23:31+00:00",
|
||||
"local_count": 96,
|
||||
"new_models_added": 2,
|
||||
"new_models_added_list": [
|
||||
"Tripo",
|
||||
"Hunyuan3D"
|
||||
]
|
||||
"new_models_added": 0,
|
||||
"new_models_added_list": []
|
||||
}
|
||||
|
||||
@@ -5,5 +5,5 @@
|
||||
"sha": "5d9308f50a08cf3cfb47e6a25cc525425714d3ad",
|
||||
"size": 274526,
|
||||
"model_count": 538,
|
||||
"updated_at": "2026-07-12T19:07:40+00:00"
|
||||
"updated_at": "2026-07-15T11:23:29+00:00"
|
||||
}
|
||||
|
||||
@@ -0,0 +1,123 @@
|
||||
import sys
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch, AsyncMock
|
||||
|
||||
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
if parent_dir not in sys.path:
|
||||
sys.path.insert(0, parent_dir)
|
||||
|
||||
from core.sources.civitai import _find_model_title_match_in_model
|
||||
from core.sources.civarchive import _find_model_title_match_in_model_details
|
||||
from core.path_utils import calculate_file_sha256
|
||||
|
||||
class TestTitleMatchingAndHashing(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
def test_civitai_find_model_title_match_success(self):
|
||||
model_data = {
|
||||
"name": "SDXL Base Model",
|
||||
"type": "checkpoint",
|
||||
"tags": ["anime", "stylized"],
|
||||
"modelVersions": [
|
||||
{
|
||||
"id": 123,
|
||||
"name": "v1.0",
|
||||
"baseModel": "SDXL 1.0",
|
||||
"files": [
|
||||
{
|
||||
"id": 456,
|
||||
"name": "sdxl_base.safetensors",
|
||||
"sizeKB": 2048,
|
||||
"primary": True,
|
||||
"hashes": {"sha256": "abcdef123456"},
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
# SDXL Base Model vs sdxl base model should be a high confidence match
|
||||
result = _find_model_title_match_in_model(
|
||||
model_id=1,
|
||||
model_data=model_data,
|
||||
title_query="sdxl base model",
|
||||
base_model_context="SDXL 1.0"
|
||||
)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result["model_id"], 1)
|
||||
self.assertEqual(result["version_id"], 123)
|
||||
self.assertEqual(result["filename"], "sdxl_base.safetensors")
|
||||
self.assertEqual(result["base_model"], "SDXL 1.0")
|
||||
self.assertIn("anime", result["tags"])
|
||||
|
||||
def test_civitai_find_model_title_match_rejected_confidence(self):
|
||||
model_data = {
|
||||
"name": "Totally Different Model Name",
|
||||
"modelVersions": [{"id": 123, "name": "v1.0"}]
|
||||
}
|
||||
result = _find_model_title_match_in_model(
|
||||
model_id=1,
|
||||
model_data=model_data,
|
||||
title_query="sdxl base model"
|
||||
)
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_civarchive_find_model_title_match_success(self):
|
||||
model_details = {
|
||||
"model_id": 10,
|
||||
"name": "SD 1.5 Pruned",
|
||||
"type": "checkpoint",
|
||||
"tags": ["base"],
|
||||
"versions": [
|
||||
{
|
||||
"id": 999,
|
||||
"name": "v1.5 pruned",
|
||||
"base_model": "SD 1.5",
|
||||
"files": [
|
||||
{
|
||||
"name": "v1-5-pruned.safetensors",
|
||||
"size": 10240,
|
||||
"download_url": "https://civarchive.com/api/download/models/999",
|
||||
"primary": True,
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
result = _find_model_title_match_in_model_details(
|
||||
model_id=10,
|
||||
model_details=model_details,
|
||||
title_query="SD 1.5 Pruned",
|
||||
base_model_context="SD 1.5"
|
||||
)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result["model_id"], 10)
|
||||
self.assertEqual(result["version_id"], 999)
|
||||
self.assertEqual(result["filename"], "v1-5-pruned.safetensors")
|
||||
self.assertEqual(result["base_model"], "SD 1.5")
|
||||
|
||||
async def test_calculate_file_sha256_with_mock_progress(self):
|
||||
import tempfile
|
||||
on_progress = MagicMock()
|
||||
on_source = MagicMock()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
file_path = os.path.join(tmpdir, "dummy.safetensors")
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(b"a" * 1024 * 512) # 512 KB
|
||||
|
||||
# Non-safetensors or empty header safetensors
|
||||
sha256 = calculate_file_sha256(
|
||||
file_path,
|
||||
chunk_size=65536,
|
||||
on_progress=on_progress,
|
||||
use_safetensors_header=True,
|
||||
on_hash_source=on_source
|
||||
)
|
||||
self.assertIsNotNone(sha256)
|
||||
on_progress.assert_called()
|
||||
on_source.assert_called_with("file")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1159,23 +1159,24 @@ export const optionsMethods = {
|
||||
|
||||
const shortSha = (value) => value ? String(value).slice(0, 8) : '';
|
||||
|
||||
const setModelListBusy = (busy) => {
|
||||
if (modelListCheckBtn) modelListCheckBtn.disabled = busy;
|
||||
if (modelListUpdateBtn) modelListUpdateBtn.disabled = busy;
|
||||
const setControlsBusyState = (buttons, busy, cancelBtn = null) => {
|
||||
const btnList = Array.isArray(buttons) ? buttons : [buttons];
|
||||
btnList.forEach(btn => {
|
||||
if (btn) btn.disabled = busy;
|
||||
});
|
||||
if (cancelBtn) {
|
||||
cancelBtn.hidden = !busy;
|
||||
cancelBtn.disabled = false;
|
||||
}
|
||||
};
|
||||
|
||||
const setBaseModelsBusy = (busy) => {
|
||||
if (baseModelsCheckBtn) baseModelsCheckBtn.disabled = busy;
|
||||
if (baseModelsUpdateBtn) baseModelsUpdateBtn.disabled = busy;
|
||||
};
|
||||
const setModelListBusy = (busy) => setControlsBusyState([modelListCheckBtn, modelListUpdateBtn], busy);
|
||||
|
||||
const setHfIndexBusy = (busy) => {
|
||||
if (hfIndexRefreshBtn) hfIndexRefreshBtn.disabled = busy;
|
||||
};
|
||||
const setBaseModelsBusy = (busy) => setControlsBusyState([baseModelsCheckBtn, baseModelsUpdateBtn], busy);
|
||||
|
||||
const setMetadataSizeAuditBusy = (busy) => {
|
||||
if (metadataSizeAuditBtn) metadataSizeAuditBtn.disabled = busy;
|
||||
};
|
||||
const setHfIndexBusy = (busy) => setControlsBusyState(hfIndexRefreshBtn, busy);
|
||||
|
||||
const setMetadataSizeAuditBusy = (busy) => setControlsBusyState(metadataSizeAuditBtn, busy);
|
||||
|
||||
const setStatusMode = (element, mode = '') => {
|
||||
if (!element) return;
|
||||
@@ -1188,12 +1189,14 @@ export const optionsMethods = {
|
||||
return Number.isFinite(number) ? number.toLocaleString() : '0';
|
||||
};
|
||||
|
||||
const setMetadataSizeAuditStatus = (text, mode = '') => {
|
||||
if (!metadataSizeAuditStatus) return;
|
||||
metadataSizeAuditStatus.textContent = text;
|
||||
setStatusMode(metadataSizeAuditStatus, mode);
|
||||
const setElementStatusText = (element, text, mode = '') => {
|
||||
if (!element) return;
|
||||
element.textContent = text;
|
||||
setStatusMode(element, mode);
|
||||
};
|
||||
|
||||
const setMetadataSizeAuditStatus = (text, mode = '') => setElementStatusText(metadataSizeAuditStatus, text, mode);
|
||||
|
||||
const setMetadataSizeSummaryValue = (element, value, mode = '') => {
|
||||
if (!element) return;
|
||||
element.textContent = formatAuditNumber(value);
|
||||
@@ -1394,19 +1397,9 @@ export const optionsMethods = {
|
||||
}
|
||||
};
|
||||
|
||||
const setMetadataBuildBusy = (busy) => {
|
||||
if (metadataBuildStartBtn) metadataBuildStartBtn.disabled = busy;
|
||||
if (metadataBuildCancelBtn) {
|
||||
metadataBuildCancelBtn.hidden = !busy;
|
||||
metadataBuildCancelBtn.disabled = false;
|
||||
}
|
||||
};
|
||||
const setMetadataBuildBusy = (busy) => setControlsBusyState(metadataBuildStartBtn, busy, metadataBuildCancelBtn);
|
||||
|
||||
const setMetadataBuildStatus = (text, mode = '') => {
|
||||
if (!metadataBuildStatus) return;
|
||||
metadataBuildStatus.textContent = text;
|
||||
setStatusMode(metadataBuildStatus, mode);
|
||||
};
|
||||
const setMetadataBuildStatus = (text, mode = '') => setElementStatusText(metadataBuildStatus, text, mode);
|
||||
|
||||
const setMetadataBuildSummaryValue = (element, value, mode = '') => {
|
||||
if (!element) return;
|
||||
|
||||
Reference in New Issue
Block a user