feat: add SHA256 hash calculation with cancellation, progress tracking, and remote size probing
This commit is contained in:
Dariusz L
+409
-1
@@ -54,8 +54,11 @@ class ModelResolverExtension:
|
||||
self.logger = create_module_logger(__name__)
|
||||
self.analysis_progress = {}
|
||||
self.search_progress = {}
|
||||
self.hash_progress = {}
|
||||
self.cancelled_search_progress = set()
|
||||
self.cancelled_hash_progress = set()
|
||||
self.search_progress_lock = threading.Lock()
|
||||
self.hash_progress_lock = threading.Lock()
|
||||
self.search_result_timestamps = {}
|
||||
|
||||
def initialize(self):
|
||||
@@ -103,7 +106,11 @@ class ModelResolverExtension:
|
||||
)
|
||||
from .core.path_templates import infer_download_path_templates
|
||||
from .core.scanner import get_model_files, invalidate_model_files_cache
|
||||
from .core.path_utils import get_filename_from_path
|
||||
from .core.path_utils import (
|
||||
get_filename_from_path,
|
||||
read_json_safe,
|
||||
write_json_atomic,
|
||||
)
|
||||
from .core.type_utils import first_non_empty, to_int, to_bool
|
||||
from .core.settings import (
|
||||
bool_setting as resolver_bool_setting,
|
||||
@@ -124,7 +131,9 @@ class ModelResolverExtension:
|
||||
get_all_progress,
|
||||
cancel_download,
|
||||
get_download_directory,
|
||||
get_metadata_sidecar_path,
|
||||
normalize_download_category,
|
||||
write_lora_manager_metadata,
|
||||
)
|
||||
from .core.sources.popular import (
|
||||
get_popular_model_url,
|
||||
@@ -544,6 +553,349 @@ class ModelResolverExtension:
|
||||
|
||||
return web.json_response({"success": True})
|
||||
|
||||
def cleanup_hash_progress(max_age_seconds=300):
|
||||
cutoff = time.time() - max_age_seconds
|
||||
with self.hash_progress_lock:
|
||||
expired = [
|
||||
progress_id
|
||||
for progress_id, progress in self.hash_progress.items()
|
||||
if progress.get("updated_at", progress.get("created_at", 0)) < cutoff
|
||||
]
|
||||
for progress_id in expired:
|
||||
self.hash_progress.pop(progress_id, None)
|
||||
self.cancelled_hash_progress.difference_update(expired)
|
||||
|
||||
def update_hash_progress(progress_id, **payload):
|
||||
if not progress_id:
|
||||
return
|
||||
with self.hash_progress_lock:
|
||||
current = self.hash_progress.get(progress_id, {})
|
||||
self.hash_progress[progress_id] = {
|
||||
**current,
|
||||
**payload,
|
||||
"progress_id": progress_id,
|
||||
"updated_at": time.time(),
|
||||
}
|
||||
|
||||
def is_hash_progress_cancelled(progress_id):
|
||||
if not progress_id:
|
||||
return False
|
||||
with self.hash_progress_lock:
|
||||
return progress_id in self.cancelled_hash_progress
|
||||
|
||||
def mark_hash_progress_cancelled(progress_id):
|
||||
if not progress_id:
|
||||
return False
|
||||
cleanup_hash_progress()
|
||||
with self.hash_progress_lock:
|
||||
current = self.hash_progress.get(progress_id)
|
||||
if not current:
|
||||
return False
|
||||
self.cancelled_hash_progress.add(progress_id)
|
||||
self.hash_progress[progress_id] = {
|
||||
**current,
|
||||
"status": "cancelling",
|
||||
"stage": "cancelling",
|
||||
"message": "Stopping hash calculation...",
|
||||
"updated_at": time.time(),
|
||||
}
|
||||
return True
|
||||
|
||||
class HashCalculationCancelled(Exception):
|
||||
pass
|
||||
|
||||
def resolve_hash_file_request(data):
|
||||
import os as _os
|
||||
|
||||
file_path = (
|
||||
data.get("file_path")
|
||||
or data.get("resolved_path")
|
||||
or data.get("path")
|
||||
or ""
|
||||
)
|
||||
if not file_path:
|
||||
return "", "file_path is required"
|
||||
|
||||
normalized_path = _os.path.abspath(_os.path.normpath(file_path))
|
||||
if not _os.path.exists(normalized_path) or not _os.path.isfile(normalized_path):
|
||||
return "", "file does not exist"
|
||||
return normalized_path, ""
|
||||
|
||||
def write_calculated_hash_metadata(normalized_path, metadata_path, sha256):
|
||||
import os as _os
|
||||
|
||||
model_dir = _os.path.abspath(_os.path.dirname(normalized_path))
|
||||
resolved_metadata_path = ""
|
||||
if metadata_path:
|
||||
candidate_path = _os.path.abspath(_os.path.normpath(metadata_path))
|
||||
try:
|
||||
same_dir = _os.path.commonpath([model_dir, candidate_path]) == model_dir
|
||||
except ValueError:
|
||||
same_dir = False
|
||||
if same_dir:
|
||||
resolved_metadata_path = candidate_path
|
||||
|
||||
if not resolved_metadata_path:
|
||||
resolved_metadata_path = get_metadata_sidecar_path(normalized_path)
|
||||
|
||||
metadata_updated = False
|
||||
try:
|
||||
metadata = read_json_safe(resolved_metadata_path, {})
|
||||
if not isinstance(metadata, dict):
|
||||
metadata = {}
|
||||
|
||||
filename = _os.path.basename(normalized_path)
|
||||
stem, _ext = _os.path.splitext(filename)
|
||||
hashes = metadata.get("hashes")
|
||||
if not isinstance(hashes, dict):
|
||||
hashes = {}
|
||||
hashes["SHA256"] = sha256
|
||||
|
||||
metadata["sha256"] = sha256
|
||||
metadata["hashes"] = hashes
|
||||
metadata["hash_status"] = "completed"
|
||||
metadata["last_checked_at"] = time.time()
|
||||
metadata.setdefault("file_name", stem)
|
||||
metadata.setdefault("model_name", stem)
|
||||
metadata.setdefault("file_path", normalized_path.replace("\\", "/"))
|
||||
try:
|
||||
metadata.setdefault("size", _os.path.getsize(normalized_path))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
write_json_atomic(resolved_metadata_path, metadata, indent=2)
|
||||
metadata_updated = True
|
||||
log_info(
|
||||
f"Calculated SHA256 and updated metadata: {resolved_metadata_path}"
|
||||
)
|
||||
except Exception as metadata_error:
|
||||
log_warn(
|
||||
f"Could not update metadata with calculated SHA256 for {normalized_path}: {metadata_error}"
|
||||
)
|
||||
|
||||
return resolved_metadata_path, metadata_updated
|
||||
|
||||
def calculate_sha256_with_progress(normalized_path, progress_id=""):
|
||||
import hashlib
|
||||
import os as _os
|
||||
|
||||
total_bytes = max(0, _os.path.getsize(normalized_path))
|
||||
bytes_read = 0
|
||||
sha256_hash = hashlib.sha256()
|
||||
last_update = 0.0
|
||||
chunk_size = 1024 * 1024 * 4
|
||||
|
||||
update_hash_progress(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="hashing",
|
||||
message="Calculating SHA256...",
|
||||
percent=0,
|
||||
bytes_read=0,
|
||||
total_bytes=total_bytes,
|
||||
)
|
||||
|
||||
with open(normalized_path, "rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(chunk_size), b""):
|
||||
if not chunk:
|
||||
continue
|
||||
sha256_hash.update(chunk)
|
||||
bytes_read += len(chunk)
|
||||
if is_hash_progress_cancelled(progress_id):
|
||||
percent = 0 if total_bytes <= 0 else min(
|
||||
98,
|
||||
(bytes_read / total_bytes) * 98,
|
||||
)
|
||||
update_hash_progress(
|
||||
progress_id,
|
||||
status="cancelled",
|
||||
stage="cancelled",
|
||||
message="Hash calculation cancelled",
|
||||
percent=percent,
|
||||
bytes_read=bytes_read,
|
||||
total_bytes=total_bytes,
|
||||
)
|
||||
raise HashCalculationCancelled()
|
||||
now = time.time()
|
||||
if now - last_update >= 0.15 or bytes_read >= total_bytes:
|
||||
percent = 98 if total_bytes <= 0 else min(
|
||||
98,
|
||||
(bytes_read / total_bytes) * 98,
|
||||
)
|
||||
update_hash_progress(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="hashing",
|
||||
message="Calculating SHA256...",
|
||||
percent=percent,
|
||||
bytes_read=bytes_read,
|
||||
total_bytes=total_bytes,
|
||||
)
|
||||
last_update = now
|
||||
|
||||
return sha256_hash.hexdigest()
|
||||
|
||||
@routes.post("/model_resolver/calculate-file-hash")
|
||||
@json_api_endpoint("calculate-file-hash")
|
||||
async def calculate_file_hash_route(request):
|
||||
"""Calculate SHA256 for a local model and persist it to sidecar metadata."""
|
||||
data = await request.json()
|
||||
metadata_path = data.get("metadata_path") or ""
|
||||
|
||||
normalized_path, error = resolve_hash_file_request(data)
|
||||
if error == "file_path is required":
|
||||
return web.json_response(
|
||||
{"error": "file_path is required"}, status=400
|
||||
)
|
||||
if error:
|
||||
return web.json_response(
|
||||
{"error": error}, status=404
|
||||
)
|
||||
|
||||
sha256 = calculate_sha256_with_progress(normalized_path)
|
||||
if not sha256:
|
||||
return web.json_response(
|
||||
{"error": "could not calculate hash"}, status=500
|
||||
)
|
||||
resolved_metadata_path, metadata_updated = write_calculated_hash_metadata(
|
||||
normalized_path,
|
||||
metadata_path,
|
||||
sha256,
|
||||
)
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"sha256": sha256,
|
||||
"hash": sha256,
|
||||
"file_path": normalized_path,
|
||||
"metadata_path": resolved_metadata_path,
|
||||
"metadata_updated": metadata_updated,
|
||||
}
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/calculate-file-hash/start")
|
||||
@json_api_endpoint("calculate-file-hash-start")
|
||||
async def calculate_file_hash_start_route(request):
|
||||
"""Start SHA256 calculation in a background thread."""
|
||||
import uuid
|
||||
|
||||
data = await request.json()
|
||||
metadata_path = data.get("metadata_path") or ""
|
||||
normalized_path, error = resolve_hash_file_request(data)
|
||||
if error == "file_path is required":
|
||||
return web.json_response(
|
||||
{"error": "file_path is required"}, status=400
|
||||
)
|
||||
if error:
|
||||
return web.json_response(
|
||||
{"error": error}, status=404
|
||||
)
|
||||
|
||||
cleanup_hash_progress()
|
||||
progress_id = f"hash_{uuid.uuid4().hex}"
|
||||
with self.hash_progress_lock:
|
||||
self.hash_progress[progress_id] = {
|
||||
"progress_id": progress_id,
|
||||
"status": "queued",
|
||||
"stage": "queued",
|
||||
"message": "Preparing hash calculation...",
|
||||
"percent": 0,
|
||||
"file_path": normalized_path,
|
||||
"created_at": time.time(),
|
||||
"updated_at": time.time(),
|
||||
}
|
||||
|
||||
def run_hash_task():
|
||||
try:
|
||||
sha256 = calculate_sha256_with_progress(
|
||||
normalized_path,
|
||||
progress_id=progress_id,
|
||||
)
|
||||
if is_hash_progress_cancelled(progress_id):
|
||||
raise HashCalculationCancelled()
|
||||
update_hash_progress(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="metadata",
|
||||
message="Saving metadata...",
|
||||
percent=99,
|
||||
)
|
||||
resolved_metadata_path, metadata_updated = write_calculated_hash_metadata(
|
||||
normalized_path,
|
||||
metadata_path,
|
||||
sha256,
|
||||
)
|
||||
update_hash_progress(
|
||||
progress_id,
|
||||
status="done",
|
||||
stage="done",
|
||||
message="Hash calculated",
|
||||
percent=100,
|
||||
sha256=sha256,
|
||||
hash=sha256,
|
||||
file_path=normalized_path,
|
||||
metadata_path=resolved_metadata_path,
|
||||
metadata_updated=metadata_updated,
|
||||
)
|
||||
except HashCalculationCancelled:
|
||||
update_hash_progress(
|
||||
progress_id,
|
||||
status="cancelled",
|
||||
stage="cancelled",
|
||||
message="Hash calculation cancelled",
|
||||
)
|
||||
except Exception as exc:
|
||||
log_exception(f"Hash calculation failed for {normalized_path}: {exc}")
|
||||
update_hash_progress(
|
||||
progress_id,
|
||||
status="error",
|
||||
stage="error",
|
||||
message=str(exc) or "Hash calculation failed",
|
||||
percent=100,
|
||||
error=str(exc) or "Hash calculation failed",
|
||||
)
|
||||
|
||||
threading.Thread(target=run_hash_task, daemon=True).start()
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"progress_id": progress_id,
|
||||
}
|
||||
)
|
||||
|
||||
@routes.get("/model_resolver/calculate-file-hash/progress/{progress_id}")
|
||||
@json_api_endpoint("calculate-file-hash-progress")
|
||||
async def calculate_file_hash_progress_route(request):
|
||||
"""Return progress for a background SHA256 calculation."""
|
||||
progress_id = request.match_info.get("progress_id", "").strip()
|
||||
cleanup_hash_progress()
|
||||
with self.hash_progress_lock:
|
||||
progress = dict(self.hash_progress.get(progress_id) or {})
|
||||
if not progress:
|
||||
return web.json_response(
|
||||
{"error": "progress not found"}, status=404
|
||||
)
|
||||
return web.json_response(progress)
|
||||
|
||||
@routes.post("/model_resolver/calculate-file-hash/cancel/{progress_id}")
|
||||
@json_api_endpoint("calculate-file-hash-cancel")
|
||||
async def calculate_file_hash_cancel_route(request):
|
||||
"""Cancel a background SHA256 calculation."""
|
||||
progress_id = request.match_info.get("progress_id", "").strip()
|
||||
if not progress_id:
|
||||
return web.json_response(
|
||||
{"error": "progress_id is required"}, status=400
|
||||
)
|
||||
cancelled = mark_hash_progress_cancelled(progress_id)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"cancelled": cancelled,
|
||||
"progress_id": progress_id,
|
||||
}
|
||||
)
|
||||
|
||||
@routes.get("/model_resolver/models")
|
||||
@json_api_endpoint("get_models")
|
||||
async def get_models(request):
|
||||
@@ -820,11 +1172,67 @@ class ModelResolverExtension:
|
||||
or result.get("from_metadata")
|
||||
or result.get("trained_words")
|
||||
):
|
||||
metadata_path = result.get("metadata_path") or ""
|
||||
metadata_saved = False
|
||||
if not result.get("from_metadata"):
|
||||
try:
|
||||
sidecar_path = get_metadata_sidecar_path(file_path)
|
||||
if sidecar_path and not _os.path.exists(sidecar_path):
|
||||
metadata_payload = {
|
||||
"source": "civitai",
|
||||
"details_source": "civitai",
|
||||
"filename": filename,
|
||||
"category": category,
|
||||
"model_name": result.get("model_name", clean_name),
|
||||
"name": result.get("model_name", clean_name),
|
||||
"model_type": result.get("model_type", ""),
|
||||
"type": result.get("model_type", ""),
|
||||
"model_id": result.get("model_id"),
|
||||
"version_id": result.get("version_id"),
|
||||
"version_name": result.get("version_name", ""),
|
||||
"sha256": result.get("sha256"),
|
||||
"size": result.get("size"),
|
||||
"base_model": result.get("base_model"),
|
||||
"tags": result.get("tags", []),
|
||||
"trained_words": result.get("trained_words", []),
|
||||
"images": result.get("images", []),
|
||||
"clip_skip": result.get("clip_skip"),
|
||||
"description": result.get("description", ""),
|
||||
"model_description": result.get("model_description", ""),
|
||||
"download_url": result.get("download_url"),
|
||||
"url": result.get("version_url") or result.get("url"),
|
||||
"path_metadata": {
|
||||
"filename": filename,
|
||||
"category": category,
|
||||
"source": "civitai",
|
||||
"model_id": result.get("model_id"),
|
||||
"version_id": result.get("version_id"),
|
||||
},
|
||||
}
|
||||
metadata_path = write_lora_manager_metadata(
|
||||
file_path,
|
||||
metadata_payload,
|
||||
category,
|
||||
result.get("version_url")
|
||||
or result.get("url")
|
||||
or result.get("download_url")
|
||||
or "",
|
||||
) or ""
|
||||
metadata_saved = bool(metadata_path)
|
||||
elif sidecar_path:
|
||||
metadata_path = sidecar_path
|
||||
except Exception as metadata_error:
|
||||
self.logger.warning(
|
||||
f"CivitAI metadata sidecar save failed: {metadata_error}"
|
||||
)
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"filename": filename,
|
||||
"file_path": file_path,
|
||||
"resolved_path": file_path,
|
||||
"metadata_path": metadata_path,
|
||||
"metadata_saved": metadata_saved,
|
||||
"location": result.get("location")
|
||||
or file_location,
|
||||
"url": result.get("url"),
|
||||
|
||||
@@ -1699,6 +1699,7 @@ def _metadata_to_model_info(metadata: Dict[str, Any]) -> Dict[str, Any]:
|
||||
model_info = _as_metadata_dict(civitai_data.get("model"))
|
||||
file_infos = _as_metadata_list(civitai_data.get("files"))
|
||||
file_info = _as_metadata_dict(file_infos[0]) if file_infos else {}
|
||||
metadata_hashes = _as_metadata_dict(metadata.get("hashes"))
|
||||
|
||||
# Build model_id from CivitAI data
|
||||
model_id = civitai_data.get("modelId") or civitai_data.get("id")
|
||||
@@ -1722,7 +1723,13 @@ def _metadata_to_model_info(metadata: Dict[str, Any]) -> Dict[str, Any]:
|
||||
civitai_data.get("name"),
|
||||
)
|
||||
or "",
|
||||
"sha256": metadata.get("sha256", ""),
|
||||
"sha256": _first_metadata_value(
|
||||
metadata.get("sha256"),
|
||||
metadata.get("hash"),
|
||||
metadata_hashes.get("SHA256"),
|
||||
metadata_hashes.get("sha256"),
|
||||
)
|
||||
or "",
|
||||
"size": _metadata_size_to_bytes(
|
||||
_first_metadata_value(
|
||||
metadata.get("size"),
|
||||
@@ -1801,6 +1808,7 @@ def get_model_info_for_file(
|
||||
log_info(f"Using metadata file for {file_path}")
|
||||
result = _metadata_to_model_info(metadata)
|
||||
result["file_path"] = file_path
|
||||
result["metadata_path"] = metadata_path
|
||||
result["location"] = _format_model_location(file_path)
|
||||
if not result.get("size"):
|
||||
try:
|
||||
|
||||
+189
-14
@@ -37,6 +37,7 @@ HF_AUTHOR_INDEX_CACHE_PATH = os.path.join(
|
||||
# Cache for search results
|
||||
_search_cache: Dict[str, Any] = {}
|
||||
_author_index_cache: Dict[str, Dict[str, Any]] = {}
|
||||
_download_size_cache: Dict[str, Optional[int]] = {}
|
||||
_author_index_lock = threading.RLock()
|
||||
|
||||
|
||||
@@ -44,6 +45,7 @@ def clear_search_cache():
|
||||
"""Clear cached HuggingFace search results and in-memory indexes."""
|
||||
global _search_cache, _author_index_cache
|
||||
_search_cache.clear()
|
||||
_download_size_cache.clear()
|
||||
with _author_index_lock:
|
||||
_author_index_cache.clear()
|
||||
|
||||
@@ -350,6 +352,128 @@ def get_huggingface_download_url(repo: str, filename: str, branch: str = "main")
|
||||
return f"https://huggingface.co/{repo}/resolve/{branch}/{quote(filename)}"
|
||||
|
||||
|
||||
def _parse_size_header(value: Any) -> Optional[int]:
|
||||
if value is None or value == "":
|
||||
return None
|
||||
try:
|
||||
size = int(str(value).strip())
|
||||
return size if size > 0 else None
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _parse_content_range_size(value: Any) -> Optional[int]:
|
||||
if not value:
|
||||
return None
|
||||
match = re.search(r"/(\d+)\s*$", str(value))
|
||||
if not match:
|
||||
return None
|
||||
return _parse_size_header(match.group(1))
|
||||
|
||||
|
||||
def _extract_file_info_size(file_info: Dict[str, Any]) -> Optional[int]:
|
||||
if not isinstance(file_info, dict):
|
||||
return None
|
||||
|
||||
for key in ("size", "sizeBytes", "size_bytes", "fileSize", "file_size", "bytes"):
|
||||
size = _parse_size_header(file_info.get(key))
|
||||
if size:
|
||||
return size
|
||||
|
||||
lfs_info = file_info.get("lfs") if isinstance(file_info.get("lfs"), dict) else {}
|
||||
for key in ("size", "sizeBytes", "size_bytes"):
|
||||
size = _parse_size_header(lfs_info.get(key))
|
||||
if size:
|
||||
return size
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_huggingface_size_probe_url(url: str) -> Optional[str]:
|
||||
if not isinstance(url, str) or not url.strip():
|
||||
return None
|
||||
value = url.strip()
|
||||
parsed = urlparse(value)
|
||||
if "huggingface.co" not in parsed.netloc:
|
||||
return None
|
||||
if "/blob/" in parsed.path:
|
||||
value = value.replace("/blob/", "/resolve/", 1)
|
||||
return value
|
||||
|
||||
|
||||
def _response_file_size(response: requests.Response) -> Optional[int]:
|
||||
size = _parse_content_range_size(response.headers.get("Content-Range"))
|
||||
if size:
|
||||
return size
|
||||
|
||||
for key in (
|
||||
"x-linked-size",
|
||||
"x-file-size",
|
||||
"x-goog-stored-content-length",
|
||||
"content-length",
|
||||
):
|
||||
size = _parse_size_header(response.headers.get(key))
|
||||
if size:
|
||||
content_type = (response.headers.get("content-type") or "").lower()
|
||||
if key == "content-length" and "text/html" in content_type:
|
||||
continue
|
||||
return size
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _fetch_remote_file_size_bytes(
|
||||
url: str,
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
timeout: int = 15,
|
||||
) -> Optional[int]:
|
||||
probe_url = _normalize_huggingface_size_probe_url(url)
|
||||
if not probe_url:
|
||||
return None
|
||||
|
||||
request_headers = {**(headers or {}), "Accept-Encoding": "identity"}
|
||||
cache_scope = "token" if request_headers.get("Authorization") else "public"
|
||||
cache_key = f"{cache_scope}::{probe_url}"
|
||||
if cache_key in _download_size_cache:
|
||||
return _download_size_cache[cache_key]
|
||||
|
||||
size = None
|
||||
try:
|
||||
response = requests.head(
|
||||
probe_url,
|
||||
headers=request_headers,
|
||||
allow_redirects=True,
|
||||
timeout=timeout,
|
||||
)
|
||||
try:
|
||||
if response.status_code < 400:
|
||||
size = _response_file_size(response)
|
||||
finally:
|
||||
response.close()
|
||||
except Exception as e:
|
||||
log_debug(f"HuggingFace size HEAD failed: url={probe_url}, error={e}")
|
||||
|
||||
if not size:
|
||||
try:
|
||||
response = requests.get(
|
||||
probe_url,
|
||||
headers={**request_headers, "Range": "bytes=0-0"},
|
||||
allow_redirects=True,
|
||||
stream=True,
|
||||
timeout=timeout,
|
||||
)
|
||||
try:
|
||||
if response.status_code < 400:
|
||||
size = _response_file_size(response)
|
||||
finally:
|
||||
response.close()
|
||||
except Exception as e:
|
||||
log_debug(f"HuggingFace size range request failed: url={probe_url}, error={e}")
|
||||
|
||||
_download_size_cache[cache_key] = size
|
||||
return size
|
||||
|
||||
|
||||
def clean_filename_for_search(filename: str) -> str:
|
||||
"""
|
||||
Clean up filename for better search results.
|
||||
@@ -367,7 +491,11 @@ def clean_filename_for_search(filename: str) -> str:
|
||||
|
||||
|
||||
def _build_huggingface_result(
|
||||
repo_id: str, file_path: str, file_info: Dict[str, Any], match_type: str
|
||||
repo_id: str,
|
||||
file_path: str,
|
||||
file_info: Dict[str, Any],
|
||||
match_type: str,
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
lfs_info = file_info.get("lfs") if isinstance(file_info.get("lfs"), dict) else {}
|
||||
sha256 = (
|
||||
@@ -376,13 +504,17 @@ def _build_huggingface_result(
|
||||
or lfs_info.get("sha256")
|
||||
or lfs_info.get("oid")
|
||||
)
|
||||
download_url = get_huggingface_download_url(repo_id, file_path)
|
||||
size = _extract_file_info_size(file_info)
|
||||
if not size:
|
||||
size = _fetch_remote_file_size_bytes(download_url, headers=headers)
|
||||
return {
|
||||
"source": "huggingface",
|
||||
"repo_id": repo_id,
|
||||
"filename": os.path.basename(file_path),
|
||||
"path": file_path,
|
||||
"url": get_huggingface_download_url(repo_id, file_path),
|
||||
"size": file_info.get("size"),
|
||||
"url": download_url,
|
||||
"size": size,
|
||||
"match_type": match_type,
|
||||
"sha256": sha256,
|
||||
}
|
||||
@@ -416,6 +548,7 @@ def _find_matching_file_in_repo(
|
||||
files: List[Dict[str, Any]],
|
||||
filename: str,
|
||||
exact_only: bool = False,
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
filename_lower = filename.lower()
|
||||
filename_base = os.path.splitext(filename_lower)[0]
|
||||
@@ -431,7 +564,13 @@ def _find_matching_file_in_repo(
|
||||
file_base = os.path.splitext(file_name_lower)[0]
|
||||
|
||||
if file_name_lower == filename_lower or file_path.lower().endswith(filename_lower):
|
||||
return _build_huggingface_result(repo_id, file_path, file_info, "exact")
|
||||
return _build_huggingface_result(
|
||||
repo_id,
|
||||
file_path,
|
||||
file_info,
|
||||
"exact",
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
if exact_only:
|
||||
return None
|
||||
@@ -449,7 +588,11 @@ def _find_matching_file_in_repo(
|
||||
if filename_base in file_base or file_base in filename_base:
|
||||
if file_path.endswith(".safetensors") or file_path.endswith(".ckpt"):
|
||||
partial_match = _build_huggingface_result(
|
||||
repo_id, file_path, file_info, "partial"
|
||||
repo_id,
|
||||
file_path,
|
||||
file_info,
|
||||
"partial",
|
||||
headers=headers,
|
||||
)
|
||||
break
|
||||
|
||||
@@ -457,7 +600,10 @@ def _find_matching_file_in_repo(
|
||||
|
||||
|
||||
def _find_matching_file_in_author_index(
|
||||
index: Dict[str, Any], filename: str, exact_only: bool = False
|
||||
index: Dict[str, Any],
|
||||
filename: str,
|
||||
exact_only: bool = False,
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
files = index.get("files") or []
|
||||
filename_lower = filename.lower()
|
||||
@@ -475,7 +621,13 @@ def _find_matching_file_in_author_index(
|
||||
if file_name_lower == filename_lower or file_path.lower().endswith(
|
||||
filename_lower
|
||||
):
|
||||
return _build_huggingface_result(repo_id, file_path, file_info, "exact")
|
||||
return _build_huggingface_result(
|
||||
repo_id,
|
||||
file_path,
|
||||
file_info,
|
||||
"exact",
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
if exact_only:
|
||||
return None
|
||||
@@ -496,7 +648,11 @@ def _find_matching_file_in_author_index(
|
||||
".ckpt"
|
||||
):
|
||||
partial_match = _build_huggingface_result(
|
||||
repo_id, file_path, file_info, "partial"
|
||||
repo_id,
|
||||
file_path,
|
||||
file_info,
|
||||
"partial",
|
||||
headers=headers,
|
||||
)
|
||||
break
|
||||
|
||||
@@ -736,7 +892,11 @@ def search_huggingface_for_file(
|
||||
|
||||
log_debug(f"HuggingFace tree repo={repo_id} files={len(files)}")
|
||||
result = _find_matching_file_in_repo(
|
||||
repo_id, files, filename, exact_only=exact_only
|
||||
repo_id,
|
||||
files,
|
||||
filename,
|
||||
exact_only=exact_only,
|
||||
headers=headers,
|
||||
)
|
||||
if result:
|
||||
_search_cache[cache_key] = result
|
||||
@@ -774,7 +934,10 @@ def search_huggingface_for_file(
|
||||
)
|
||||
if index:
|
||||
result = _find_matching_file_in_author_index(
|
||||
index, filename, exact_only=exact_only
|
||||
index,
|
||||
filename,
|
||||
exact_only=exact_only,
|
||||
headers={},
|
||||
)
|
||||
author_fallback_repo_count += int(index.get("repo_count") or 0)
|
||||
if result:
|
||||
@@ -831,7 +994,11 @@ def search_huggingface_for_file(
|
||||
continue
|
||||
|
||||
result = _find_matching_file_in_repo(
|
||||
repo_id, files, filename, exact_only=exact_only
|
||||
repo_id,
|
||||
files,
|
||||
filename,
|
||||
exact_only=exact_only,
|
||||
headers=headers,
|
||||
)
|
||||
if result:
|
||||
_search_cache[cache_key] = result
|
||||
@@ -888,7 +1055,11 @@ def search_huggingface_for_file(
|
||||
continue
|
||||
|
||||
result = _find_matching_file_in_repo(
|
||||
repo_id, files, expected_file, exact_only=True
|
||||
repo_id,
|
||||
files,
|
||||
expected_file,
|
||||
exact_only=True,
|
||||
headers=headers,
|
||||
)
|
||||
if result and result.get("filename", "").lower() == filename.lower():
|
||||
_search_cache[cache_key] = result
|
||||
@@ -995,12 +1166,16 @@ def get_repo_files(repo: str, token: Optional[str] = None) -> List[Dict[str, Any
|
||||
if path.endswith(
|
||||
(".safetensors", ".ckpt", ".pt", ".bin", ".pth", ".onnx", ".gguf")
|
||||
):
|
||||
url = get_huggingface_download_url(repo, path)
|
||||
size = _extract_file_info_size(item)
|
||||
if not size:
|
||||
size = _fetch_remote_file_size_bytes(url, headers=headers)
|
||||
files.append(
|
||||
{
|
||||
"filename": os.path.basename(path),
|
||||
"path": path,
|
||||
"url": get_huggingface_download_url(repo, path),
|
||||
"size": item.get("size"),
|
||||
"url": url,
|
||||
"size": size,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -4046,6 +4046,84 @@
|
||||
word-break: break-word;
|
||||
}
|
||||
|
||||
.mr-info-hash-control {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 8px;
|
||||
min-width: 0;
|
||||
width: 100%;
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
|
||||
.mr-info-hash {
|
||||
min-width: 0;
|
||||
overflow-wrap: anywhere;
|
||||
}
|
||||
|
||||
.mr-info-hash.is-empty {
|
||||
display: none;
|
||||
}
|
||||
|
||||
.mr-info-hash-calculate {
|
||||
flex: 0 0 auto;
|
||||
min-height: 26px;
|
||||
padding: 4px 9px;
|
||||
border: 1px solid rgba(255, 255, 255, 0.12);
|
||||
border-radius: 6px;
|
||||
background: rgba(255, 255, 255, 0.055);
|
||||
color: rgba(255, 255, 255, 0.86);
|
||||
font-size: 11px;
|
||||
font-weight: 800;
|
||||
cursor: pointer;
|
||||
}
|
||||
|
||||
.mr-info-hash-calculate:hover:not(:disabled) {
|
||||
background: rgba(255, 255, 255, 0.09);
|
||||
border-color: rgba(255, 255, 255, 0.18);
|
||||
}
|
||||
|
||||
.mr-info-hash-calculate:disabled {
|
||||
cursor: wait;
|
||||
opacity: 0.72;
|
||||
}
|
||||
|
||||
.mr-info-hash-progress {
|
||||
display: grid;
|
||||
grid-template-columns: minmax(84px, 1fr) auto;
|
||||
align-items: center;
|
||||
gap: 7px;
|
||||
min-width: min(240px, 100%);
|
||||
flex: 1 1 220px;
|
||||
}
|
||||
|
||||
.mr-info-hash-progress.mr-hidden-initial,
|
||||
.mr-info-hash-progress[hidden] {
|
||||
display: none !important;
|
||||
}
|
||||
|
||||
.mr-info-hash-progress-bar {
|
||||
height: 7px;
|
||||
overflow: hidden;
|
||||
border-radius: 999px;
|
||||
background: rgba(255, 255, 255, 0.08);
|
||||
}
|
||||
|
||||
.mr-info-hash-progress-bar span {
|
||||
display: block;
|
||||
width: 0%;
|
||||
height: 100%;
|
||||
border-radius: inherit;
|
||||
background: linear-gradient(90deg, rgba(114,211,138,0.74), rgba(100,181,246,0.82));
|
||||
transition: width 150ms ease;
|
||||
}
|
||||
|
||||
.mr-info-hash-progress small {
|
||||
color: rgba(255, 255, 255, 0.62);
|
||||
font-size: 10px;
|
||||
font-weight: 800;
|
||||
white-space: nowrap;
|
||||
}
|
||||
|
||||
.mr-info-field-label {
|
||||
display: grid;
|
||||
grid-template-columns: 14px minmax(0, 1fr) auto;
|
||||
@@ -7838,4 +7916,4 @@ a.mr-image-preview-resource-name:hover {
|
||||
|
||||
#model-resolver-modal .mr-missing-detail-pane [id^="local-matches-body-"] .mr-local-alternatives-label {
|
||||
color: #bdbdbd;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -319,7 +319,20 @@ export const modelInfoMethods = {
|
||||
</tr>
|
||||
<tr class="mr-info-hash-row">
|
||||
<td>${this.renderInfoFieldLabel('hash', 'Hash (sha256)', 'Unique fingerprint of the local file. Model Resolver uses it to confirm the exact CivitAI version.')}</td>
|
||||
<td><span class="mr-info-hash"></span></td>
|
||||
<td>
|
||||
<div class="mr-info-hash-control">
|
||||
<span class="mr-info-hash"></span>
|
||||
<button type="button" class="mr-info-hash-calculate mr-hidden-initial">
|
||||
Calculate hash
|
||||
</button>
|
||||
<div class="mr-info-hash-progress mr-hidden-initial">
|
||||
<div class="mr-info-hash-progress-bar">
|
||||
<span></span>
|
||||
</div>
|
||||
<small>0%</small>
|
||||
</div>
|
||||
</div>
|
||||
</td>
|
||||
</tr>
|
||||
<tr class="mr-info-size-row mr-hidden-initial">
|
||||
<td>${this.renderInfoFieldLabel('hardDrive', 'Size', 'The local model file size, read from metadata or from the file on disk.')}</td>
|
||||
@@ -687,11 +700,225 @@ export const modelInfoMethods = {
|
||||
}
|
||||
},
|
||||
|
||||
getInfoDialogHash(data = {}) {
|
||||
return String(data.sha256 || data.hash || data.hashes?.SHA256 || data.hashes?.sha256 || '').trim();
|
||||
},
|
||||
|
||||
getInfoDialogModelFilePath(data = {}) {
|
||||
return data.file_path
|
||||
|| data.filePath
|
||||
|| data.resolved_path
|
||||
|| data.resolvedPath
|
||||
|| data.full_path
|
||||
|| data.fullPath
|
||||
|| data.path
|
||||
|| '';
|
||||
},
|
||||
|
||||
updateInfoDialogHashDisplay(dialog, data = {}) {
|
||||
const hashEl = dialog.querySelector('.mr-info-hash');
|
||||
const button = dialog.querySelector('.mr-info-hash-calculate');
|
||||
const progressEl = dialog.querySelector('.mr-info-hash-progress');
|
||||
const hash = this.getInfoDialogHash(data);
|
||||
const filePath = this.getInfoDialogModelFilePath(data);
|
||||
|
||||
if (hashEl) {
|
||||
hashEl.textContent = hash;
|
||||
hashEl.classList.toggle('is-empty', !hash);
|
||||
}
|
||||
if (button) {
|
||||
const canCalculate = Boolean(!hash && filePath);
|
||||
button.classList.toggle('mr-hidden-initial', !canCalculate);
|
||||
button.hidden = !canCalculate;
|
||||
button.disabled = false;
|
||||
button.textContent = 'Calculate hash';
|
||||
delete button.dataset.hashCalculating;
|
||||
delete button.dataset.progressId;
|
||||
}
|
||||
if (progressEl && hash) {
|
||||
this.hideInfoDialogHashProgress(dialog);
|
||||
}
|
||||
},
|
||||
|
||||
hideInfoDialogHashProgress(dialog) {
|
||||
const progressEl = dialog.querySelector('.mr-info-hash-progress');
|
||||
if (!progressEl) return;
|
||||
const bar = progressEl.querySelector('.mr-info-hash-progress-bar span');
|
||||
const label = progressEl.querySelector('small');
|
||||
progressEl.classList.add('mr-hidden-initial');
|
||||
progressEl.hidden = true;
|
||||
if (bar) bar.style.width = '0%';
|
||||
if (label) label.textContent = '0%';
|
||||
},
|
||||
|
||||
updateInfoDialogHashProgress(dialog, percent = 0, message = '') {
|
||||
const progressEl = dialog.querySelector('.mr-info-hash-progress');
|
||||
if (!progressEl) return;
|
||||
|
||||
const safePercent = Math.max(0, Math.min(100, Number(percent) || 0));
|
||||
const bar = progressEl.querySelector('.mr-info-hash-progress-bar span');
|
||||
const label = progressEl.querySelector('small');
|
||||
progressEl.classList.remove('mr-hidden-initial');
|
||||
progressEl.hidden = false;
|
||||
if (bar) bar.style.width = `${safePercent}%`;
|
||||
if (label) {
|
||||
label.textContent = message
|
||||
? `${Math.round(safePercent)}% · ${message}`
|
||||
: `${Math.round(safePercent)}%`;
|
||||
}
|
||||
},
|
||||
|
||||
async cancelInfoDialogHashCalculation(dialog, button) {
|
||||
const state = dialog?._hashCalculation || {};
|
||||
state.cancelRequested = true;
|
||||
dialog._hashCalculation = state;
|
||||
if (button) {
|
||||
button.disabled = true;
|
||||
button.textContent = 'Stopping...';
|
||||
}
|
||||
|
||||
if (state.progressId) {
|
||||
try {
|
||||
await this.fetchJson(
|
||||
`/model_resolver/calculate-file-hash/cancel/${encodeURIComponent(state.progressId)}`,
|
||||
{ method: 'POST', silent: true },
|
||||
'Cancel model hash calculation'
|
||||
);
|
||||
} catch (error) {
|
||||
// Keep the UI in stopping state; polling will surface the final state if possible.
|
||||
}
|
||||
}
|
||||
},
|
||||
|
||||
async calculateInfoDialogHash(dialog, button) {
|
||||
if (button?.dataset.hashCalculating === 'true') {
|
||||
await this.cancelInfoDialogHashCalculation(dialog, button);
|
||||
return;
|
||||
}
|
||||
|
||||
const data = dialog?._infoDialogData || {};
|
||||
const filePath = this.getInfoDialogModelFilePath(data);
|
||||
if (!filePath) {
|
||||
this.showNotification?.('No local file path available', 'error');
|
||||
return;
|
||||
}
|
||||
|
||||
dialog._hashCalculation = { progressId: '', cancelRequested: false };
|
||||
if (button) {
|
||||
button.disabled = false;
|
||||
button.textContent = 'Stop calculate';
|
||||
button.dataset.hashCalculating = 'true';
|
||||
delete button.dataset.progressId;
|
||||
}
|
||||
this.updateInfoDialogHashProgress(dialog, 0, 'Preparing');
|
||||
|
||||
try {
|
||||
const start = await this.fetchJson('/model_resolver/calculate-file-hash/start', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({
|
||||
file_path: filePath,
|
||||
metadata_path: data.metadata_path || data.metadataPath || ''
|
||||
})
|
||||
}, 'Start model hash calculation');
|
||||
const progressId = start?.progress_id;
|
||||
if (!progressId) {
|
||||
throw new Error('No progress id returned');
|
||||
}
|
||||
if (!dialog._hashCalculation) {
|
||||
dialog._hashCalculation = { progressId, cancelRequested: false };
|
||||
}
|
||||
dialog._hashCalculation.progressId = progressId;
|
||||
if (button) {
|
||||
button.dataset.progressId = progressId;
|
||||
button.textContent = 'Stop calculate';
|
||||
button.disabled = false;
|
||||
}
|
||||
if (dialog._hashCalculation.cancelRequested) {
|
||||
await this.cancelInfoDialogHashCalculation(dialog, button);
|
||||
}
|
||||
|
||||
let result = null;
|
||||
for (;;) {
|
||||
await new Promise(resolve => setTimeout(resolve, 250));
|
||||
if (!dialog.isConnected) {
|
||||
await this.cancelInfoDialogHashCalculation(dialog, button);
|
||||
return;
|
||||
}
|
||||
const progress = await this.fetchJson(
|
||||
`/model_resolver/calculate-file-hash/progress/${encodeURIComponent(progressId)}`,
|
||||
{ silent: true },
|
||||
'Get model hash progress'
|
||||
);
|
||||
this.updateInfoDialogHashProgress(
|
||||
dialog,
|
||||
progress?.percent || 0,
|
||||
progress?.message || ''
|
||||
);
|
||||
|
||||
if (progress?.status === 'done') {
|
||||
result = progress;
|
||||
break;
|
||||
}
|
||||
if (progress?.status === 'cancelled') {
|
||||
this.hideInfoDialogHashProgress(dialog);
|
||||
if (button) {
|
||||
button.disabled = false;
|
||||
button.textContent = 'Calculate hash';
|
||||
delete button.dataset.hashCalculating;
|
||||
delete button.dataset.progressId;
|
||||
}
|
||||
dialog._hashCalculation = null;
|
||||
this.showNotification?.('Hash calculation stopped.', 'info');
|
||||
return;
|
||||
}
|
||||
if (progress?.status === 'error') {
|
||||
throw new Error(progress.error || progress.message || 'Hash calculation failed');
|
||||
}
|
||||
}
|
||||
|
||||
const sha256 = result?.sha256 || result?.hash || '';
|
||||
if (!sha256) {
|
||||
throw new Error('No hash returned');
|
||||
}
|
||||
|
||||
const updatedData = {
|
||||
...data,
|
||||
sha256,
|
||||
hash: sha256,
|
||||
metadata_path: result.metadata_path || data.metadata_path || data.metadataPath || '',
|
||||
metadata_saved: Boolean(result.metadata_updated || data.metadata_saved)
|
||||
};
|
||||
dialog._infoDialogData = updatedData;
|
||||
dialog._hashCalculation = null;
|
||||
this.updateInfoDialogHashDisplay(dialog, updatedData);
|
||||
this.showNotification?.(
|
||||
result.metadata_updated ? 'Hash calculated and saved to metadata.' : 'Hash calculated.',
|
||||
'success'
|
||||
);
|
||||
} catch (error) {
|
||||
dialog._hashCalculation = null;
|
||||
if (button) {
|
||||
button.disabled = false;
|
||||
button.textContent = 'Calculate hash';
|
||||
delete button.dataset.hashCalculating;
|
||||
delete button.dataset.progressId;
|
||||
}
|
||||
this.hideInfoDialogHashProgress(dialog);
|
||||
this.showNotification?.(`Hash calculation failed: ${error?.message || error}`, 'error');
|
||||
}
|
||||
},
|
||||
|
||||
bindInfoDialogInteractions(dialog) {
|
||||
if (!dialog || dialog.dataset.mlInfoBound === 'true') return;
|
||||
dialog.dataset.mlInfoBound = 'true';
|
||||
|
||||
dialog.addEventListener('click', async (event) => {
|
||||
const hashButton = event.target.closest('.mr-info-hash-calculate');
|
||||
if (hashButton && dialog.contains(hashButton)) {
|
||||
await this.calculateInfoDialogHash(dialog, hashButton);
|
||||
return;
|
||||
}
|
||||
|
||||
const locationBtn = event.target.closest('.mr-info-location-button');
|
||||
if (locationBtn && dialog.contains(locationBtn)) {
|
||||
const path = locationBtn.dataset.path || '';
|
||||
@@ -860,6 +1087,7 @@ export const modelInfoMethods = {
|
||||
this.updateInfoDialogError(dialog, 'No data received');
|
||||
return;
|
||||
}
|
||||
dialog._infoDialogData = { ...(dialog._infoDialogData || {}), ...data };
|
||||
|
||||
// Update title
|
||||
const titleEl = dialog.querySelector('.mr-info-dialog-title');
|
||||
@@ -937,10 +1165,7 @@ export const modelInfoMethods = {
|
||||
}
|
||||
|
||||
// Update hash
|
||||
const hashEl = dialog.querySelector('.mr-info-hash');
|
||||
if (hashEl) {
|
||||
hashEl.textContent = data.sha256 || data.hash || '';
|
||||
}
|
||||
this.updateInfoDialogHashDisplay(dialog, data);
|
||||
|
||||
// Update CivitAI link
|
||||
const civitaiLinkEl = dialog.querySelector('.mr-info-civitai-link');
|
||||
|
||||
Reference in New Issue
Block a user