nodes/Azornes-Comfyui-Model-Resolver
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Improve HuggingFace Xet downloads with native progress snapshots, transfer metrics, cancellation, and finalizing UI state

Browse Source
This commit is contained in:
Dariusz L
2026-07-14 11:02:57 +02:00
parent 1f0edaaa7e
commit 280ca2e6de
4 changed files with 346 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/downloader.py
+245 -34
View File
@@ -53,6 +53,8 @@ aria2_action_locks: Dict[str, threading.Lock] = {}
aria2_desired_states: Dict[str, Dict[str, Any]] = {}
aria2_idle_timer: Optional[threading.Timer] = None
aria2_process_started_by_resolver = False
xet_transfers: Dict[str, Dict[str, Any]] = {}
xet_transfers_lock = threading.Lock()
# Speed calculation settings
SPEED_HISTORY_SIZE = 5 # Number of samples for smoothing
@@ -1270,6 +1272,25 @@ def _delete_partial_download_files(dest_path: str) -> None:
log.warning(f"Could not delete incomplete download file {path}: {exc}")
def _delete_xet_partial_file(partial_path: str, attempts: int = 5) -> bool:
"""Delete a stopped Xet partial file, retrying while Windows releases it."""
attempts = max(1, int(attempts or 1))
last_error: Optional[Exception] = None
for attempt in range(attempts):
try:
if not os.path.exists(partial_path):
return True
os.remove(partial_path)
return True
except Exception as exc:
last_error = exc
if attempt + 1 < attempts:
time.sleep(0.25)
log.warning(f"Could not delete incomplete Xet file {partial_path}: {last_error}")
return False
class _HuggingFaceXetDownloadCancelled(Exception):
"""Raised by the Xet progress adapter when a download is cancelled."""
@@ -1281,7 +1302,65 @@ class _HuggingFaceXetProgressAdapter:
self.download_id = download_id
self.total_size = max(0, int(total_size or 0))
self.downloaded = 0
self.transfer_downloaded = 0
self.transfer_total_size = 0
self.last_update = start_time
self.speed = 0
def _publish(
self,
downloaded: int,
speed: int,
transfer_downloaded: int = 0,
transfer_total_size: int = 0,
) -> None:
downloaded = max(0, int(downloaded or 0))
if self.total_size > 0:
downloaded = min(downloaded, self.total_size)
self.downloaded = max(self.downloaded, downloaded)
progress = (
min(100.0, round((self.downloaded / self.total_size) * 100, 1))
if self.total_size > 0
else 0
)
self.transfer_downloaded = max(
self.transfer_downloaded,
max(0, int(transfer_downloaded or 0)),
)
self.transfer_total_size = max(
self.transfer_total_size,
max(0, int(transfer_total_size or 0)),
)
transfer_progress = (
min(
100.0,
round(
(self.transfer_downloaded / self.transfer_total_size) * 100,
1,
),
)
if self.transfer_total_size > 0
else 0
)
self.speed = max(0, int(speed or 0))
with download_lock:
state = download_progress.get(self.download_id)
if state is not None:
state.update(
{
"status": "downloading",
"progress": progress,
"downloaded": self.downloaded,
"total_size": self.total_size,
"speed": self.speed,
"transfer_downloaded": self.transfer_downloaded,
"transfer_total_size": self.transfer_total_size,
"transfer_progress": transfer_progress,
"download_backend": "huggingface_xet",
}
)
def update(self, byte_delta: Any) -> None:
if self.download_id in cancelled_downloads:
@@ -1292,30 +1371,157 @@ class _HuggingFaceXetProgressAdapter:
except (TypeError, ValueError):
delta = 0
now = time.time()
elapsed = max(now - self.last_update, 0.001)
self.last_update = now
self.downloaded += delta
speed = int(delta / elapsed) if delta else 0
progress = (
min(100, int((self.downloaded / self.total_size) * 100))
if self.total_size > 0
else 0
downloaded = self.downloaded + delta
# The legacy one-argument callback exposes logical byte increments only,
# not network-transfer speed. Reporting a derived rate here would confuse
# fast cache/file reconstruction with actual download throughput.
self._publish(downloaded, 0)
def __call__(self, total_update: Any, item_updates: Any) -> None:
"""Receive hf_xet's detailed 200 ms progress snapshots when available."""
if self.download_id in cancelled_downloads:
raise _HuggingFaceXetDownloadCancelled("Download cancelled")
downloaded = int(getattr(total_update, "total_bytes_completed", 0) or 0)
if downloaded <= 0:
items = item_updates.values() if isinstance(item_updates, dict) else item_updates
for item in items or []:
downloaded = max(
downloaded,
int(getattr(item, "bytes_completed", 0) or 0),
)
transfer_downloaded = int(
getattr(total_update, "total_transfer_bytes_completed", 0) or 0
)
transfer_total_size = int(
getattr(total_update, "total_transfer_bytes", 0) or 0
)
transfer_rate = getattr(
total_update,
"total_transfer_bytes_completion_rate",
None,
)
try:
speed = max(0, int(float(transfer_rate or 0)))
except (TypeError, ValueError):
speed = 0
self.last_update = time.time()
self._publish(
downloaded,
speed,
transfer_downloaded,
transfer_total_size,
)
with download_lock:
state = download_progress.get(self.download_id)
if state is not None:
state.update(
{
"status": "downloading",
"progress": progress,
"downloaded": self.downloaded,
"total_size": self.total_size,
"speed": speed,
"download_backend": "huggingface_xet",
}
def _run_huggingface_xet_transfer(
incomplete_path: Path,
xet_file_data: Any,
request_headers: Dict[str, str],
expected_size: int,
filename: str,
progress_adapter: _HuggingFaceXetProgressAdapter,
) -> None:
"""Use detailed native Xet progress when supported, with legacy fallback."""
import hf_xet
from huggingface_hub.file_download import xet_get
from huggingface_hub.utils import refresh_xet_connection_info
supports_session_progress = all(
hasattr(hf_xet, name)
for name in ("XetFileInfo", "XetSession")
)
if supports_session_progress:
try:
from huggingface_hub.utils._xet import (
get_xet_session,
xet_headers_without_auth,
)
except ImportError:
supports_session_progress = False
if supports_session_progress:
session = get_xet_session()
xet_headers = xet_headers_without_auth(request_headers)
group = session.new_file_download_group(
token_refresh_url=xet_file_data.refresh_route,
token_refresh_headers=request_headers,
custom_headers=xet_headers,
progress_callback=progress_adapter,
progress_interval_ms=200,
)
try:
with group:
handle = group.start_download_file(
hf_xet.XetFileInfo(xet_file_data.file_hash, expected_size or None),
str(incomplete_path.absolute()),
)
with xet_transfers_lock:
xet_transfers[progress_adapter.download_id] = {
"handle": handle,
"partial_path": str(incomplete_path),
}
if progress_adapter.download_id in cancelled_downloads:
cancel = getattr(handle, "cancel", None)
if callable(cancel):
cancel()
finally:
with xet_transfers_lock:
xet_transfers.pop(progress_adapter.download_id, None)
return
supports_detailed_progress = all(
hasattr(hf_xet, name)
for name in (
"PyItemProgressUpdate",
"PyTotalProgressUpdate",
"PyXetDownloadInfo",
"download_files",
)
)
if not supports_detailed_progress:
xet_get(
incomplete_path=incomplete_path,
xet_file_data=xet_file_data,
headers=request_headers,
expected_size=expected_size or None,
displayed_filename=filename,
_tqdm_bar=progress_adapter,
)
return
connection_info = refresh_xet_connection_info(
file_data=xet_file_data,
headers=request_headers,
)
if connection_info is None:
raise ValueError("Failed to refresh Hugging Face Xet connection info")
def token_refresher() -> tuple[str, int]:
refreshed = refresh_xet_connection_info(
file_data=xet_file_data,
headers=request_headers,
)
if refreshed is None:
raise ValueError("Failed to refresh Hugging Face Xet access token")
return refreshed.access_token, refreshed.expiration_unix_epoch
download_info = hf_xet.PyXetDownloadInfo(
destination_path=str(incomplete_path.absolute()),
hash=xet_file_data.file_hash,
file_size=expected_size or None,
)
hf_xet.download_files(
[download_info],
endpoint=connection_info.endpoint,
token_info=(
connection_info.access_token,
connection_info.expiration_unix_epoch,
),
token_refresher=token_refresher,
progress_updater=[progress_adapter],
)
def _download_huggingface_xet(
@@ -1337,7 +1543,7 @@ def _download_huggingface_xet(
try:
__import__("hf_xet")
from huggingface_hub.file_download import get_hf_file_metadata, xet_get
from huggingface_hub.file_download import get_hf_file_metadata
except ImportError:
return None
@@ -1403,13 +1609,13 @@ def _download_huggingface_xet(
os.remove(partial_path)
log.info(f"Starting Hugging Face Xet download: {filename}")
xet_get(
incomplete_path=Path(partial_path),
xet_file_data=xet_file_data,
headers=request_headers,
expected_size=expected_size or None,
displayed_filename=filename,
_tqdm_bar=progress_adapter,
_run_huggingface_xet_transfer(
Path(partial_path),
xet_file_data,
request_headers,
expected_size,
filename,
progress_adapter,
)
if download_id in cancelled_downloads:
@@ -1465,11 +1671,7 @@ def _download_huggingface_xet(
isinstance(exc, _HuggingFaceXetDownloadCancelled)
or download_id in cancelled_downloads
)
try:
if os.path.exists(partial_path):
os.remove(partial_path)
except Exception as cleanup_exc:
log.warning(f"Could not delete incomplete Xet file: {cleanup_exc}")
_delete_xet_partial_file(partial_path)
error_msg = (
"Download cancelled" if was_cancelled else _sanitize_download_error(exc)
@@ -2354,6 +2556,15 @@ def cancel_download(download_id: str) -> bool:
args=(download_id, transfer["gid"]),
daemon=True,
).start()
with xet_transfers_lock:
xet_transfer = dict(xet_transfers.get(download_id) or {})
xet_handle = xet_transfer.get("handle")
xet_cancel = getattr(xet_handle, "cancel", None)
if callable(xet_cancel):
try:
xet_cancel()
except Exception as exc:
log.warning(f"Could not cancel Hugging Face Xet transfer {download_id}: {exc}")
return True
web/resolver/actions/queue_methods.js
+11 -5
View File
@@ -1253,7 +1253,8 @@ export const queueMethods = {
const progress = typeof this.applyPendingDownloadStatus === 'function'
? this.applyPendingDownloadStatus(info, info.lastProgress || {})
: (info.lastProgress || {});
const percent = Math.max(0, Math.min(100, Number(progress.progress) || 0));
const displayProgress = this.getDownloadDisplayProgress(progress);
const percent = displayProgress.percent;
const filename = progress.filename
|| info.filename
|| info.missing?.download_source?.filename
@@ -1261,17 +1262,22 @@ export const queueMethods = {
|| 'model';
const category = this.getCategoryDisplayName?.(info.category || info.missing?.category || '') || info.category || '';
const nodeLabel = info.missing?.subgraph_name || info.missing?.node_type || (info.missing?.subgraph_id ? 'Subgraph' : 'Node');
const downloaded = this.formatBytes(progress.downloaded || 0);
const total = progress.total_size ? this.formatBytes(progress.total_size) : '';
const downloaded = this.formatBytes(displayProgress.downloaded);
const total = displayProgress.totalSize ? this.formatBytes(displayProgress.totalSize) : '';
const progressMeta = this.formatDownloadProgressMeta?.(progress) || '';
const status = progress.status || info.lastStatus || 'starting';
const backend = String(progress.download_backend || info.downloadBackend || '').toLowerCase();
const isAria2 = backend === 'aria2';
const isPaused = status === 'paused';
const percentLabel = this.formatDownloadPercent?.(percent) ?? String(Math.round(percent));
const statusLabel = status === 'downloading'
? `${Math.round(percent)}%`
? (displayProgress.isFinalizing ? 'Finalizing' : `${percentLabel}%`)
: (status === 'starting' ? 'Starting' : (isPaused ? 'Paused' : status.replace(/_/g, ' ')));
const sizeText = total ? `${downloaded} / ${total}` : downloaded;
const logicalDownloaded = this.formatBytes(progress.downloaded || 0);
const logicalTotal = progress.total_size ? this.formatBytes(progress.total_size) : '';
const sizeText = displayProgress.isFinalizing && logicalTotal
? `Finalizing file: ${logicalDownloaded} / ${logicalTotal}`
: (total ? `${downloaded} / ${total}` : downloaded);
const targetPath = progress.directory || info.downloadDirectory || info.downloadPath || '';
const targetLabel = targetPath ? targetPath.split(/[\/\\]/).filter(Boolean).pop() || targetPath : '';
const workflowLabel = this.getDownloadWorkflowLabel?.(info) || 'Unknown workflow';
web/resolver/actions/resolve_download_methods.js
+47 -19
View File
@@ -463,7 +463,7 @@ export const resolveDownloadMethods = {
name: progress?.filename || info?.filename || 'Download',
path: targetPath,
resolved_path: targetPath,
open_path: filePath || directory,
open_path: directory || filePath,
folder_path: directory || targetPath,
download_directory: directory || '',
download_path: filePath || '',
@@ -1282,8 +1282,12 @@ export const resolveDownloadMethods = {
});
},
getDownloadProgressStatusLabel(status = '', percent = 0) {
if (status === 'downloading') return `${Math.round(percent)}%`;
getDownloadProgressStatusLabel(status = '', percent = 0, isFinalizing = false) {
if (status === 'downloading') {
if (isFinalizing) return 'Finalizing';
const percentLabel = this.formatDownloadPercent?.(percent) ?? String(Math.round(percent));
return `${percentLabel}%`;
}
if (status === 'starting') return 'Starting';
if (status === 'paused') return 'Paused';
if (status === 'completed_checking') return 'Checking';
@@ -1309,13 +1313,19 @@ export const resolveDownloadMethods = {
const backend = String(progress.download_backend || snapshot.downloadBackend || snapshot.download_backend || '').toLowerCase();
const isAria2 = backend === 'aria2';
const isPaused = status === 'paused';
const percent = Math.max(0, Math.min(100, Number(progress.progress) || 0));
const downloaded = this.formatBytes(progress.downloaded || 0);
const total = this.formatBytes(progress.total_size || 0);
const displayProgress = this.getDownloadDisplayProgress(progress);
const percent = displayProgress.percent;
const downloaded = this.formatBytes(displayProgress.downloaded);
const total = this.formatBytes(displayProgress.totalSize);
const progressMeta = this.formatDownloadProgressMeta(progress);
const leftText = progress.total_size
? `${downloaded} / ${total} (${percent}%)${isPaused ? ' - Paused' : ''}`
: '<span class="mr-info-accent-text">Connecting...</span>';
const percentLabel = this.formatDownloadPercent?.(percent) ?? String(Math.round(percent));
const logicalDownloaded = this.formatBytes(progress.downloaded || 0);
const logicalTotal = progress.total_size ? this.formatBytes(progress.total_size) : 'Unknown';
const leftText = displayProgress.isFinalizing
? `Finalizing file: ${logicalDownloaded} / ${logicalTotal}`
: (displayProgress.totalSize
? `${downloaded} / ${total} (${percentLabel}%)${isPaused ? ' - Paused' : ''}`
: '<span class="mr-info-accent-text">Connecting...</span>');
let actionClass = canCancel ? 'cancel-download-btn mr-btn mr-btn-danger mr-btn-sm' : '';
let actionText = canCancel ? 'Cancel' : '';
let actionsHtml = '';
@@ -1390,9 +1400,10 @@ export const resolveDownloadMethods = {
renderDownloadProgressGroupItem(downloadId, snapshot) {
const progress = snapshot.progress || {};
const status = snapshot.status || progress.status || '';
const percent = Math.max(0, Math.min(100, Number(progress.progress) || 0));
const displayProgress = this.getDownloadDisplayProgress(progress);
const percent = displayProgress.percent;
const filename = progress.filename || snapshot.filename || 'Download';
const statusLabel = this.getDownloadProgressStatusLabel(status, percent);
const statusLabel = this.getDownloadProgressStatusLabel(status, percent, displayProgress.isFinalizing);
const itemIdAttr = downloadId ? ` data-download-id="${this.escapeHtml(String(downloadId))}"` : '';
return `
<div class="mr-download-progress-item"${itemIdAttr}>
@@ -1432,20 +1443,28 @@ export const resolveDownloadMethods = {
if (!downloadBtn) return;
if (shouldRenderProgress) {
const percent = Number(progress.progress);
const displayProgress = this.getDownloadDisplayProgress(progress);
const percent = displayProgress.percent;
downloadBtn.disabled = true;
downloadBtn.classList.remove('mr-is-success-action', 'mr-btn-primary');
const label = status === 'paused'
? `Paused ${Number.isFinite(percent) && percent > 0 ? `${Math.round(percent)}%` : ''}`.trim()
: (Number.isFinite(percent) && percent > 0 ? `Downloading ${Math.round(percent)}%` : 'Starting download...');
const percentLabel = Number.isFinite(percent) && percent > 0
? `${this.formatDownloadPercent?.(percent) ?? Math.round(percent)}%`
: '';
const label = displayProgress.isFinalizing
? 'Finalizing download...'
: (status === 'paused'
? `Paused ${percentLabel}`.trim()
: (percentLabel ? `Downloading ${percentLabel}` : 'Starting download...'));
if (downloadBtn.classList.contains('search-download-btn')) {
downloadBtn.innerHTML = getSvgIcon('download');
downloadBtn.setAttribute('data-tooltip', label);
downloadBtn.setAttribute('aria-label', label);
} else {
downloadBtn.textContent = status === 'paused'
? 'Paused'
: (Number.isFinite(percent) && percent > 0 ? `${Math.round(percent)}%` : 'Starting...');
downloadBtn.textContent = displayProgress.isFinalizing
? 'Finalizing...'
: (status === 'paused'
? 'Paused'
: (Number.isFinite(percent) && percent > 0 ? `${Math.round(percent)}%` : 'Starting...'));
}
} else if (status === 'cancelling') {
downloadBtn.disabled = true;
@@ -1992,6 +2011,14 @@ export const resolveDownloadMethods = {
/**
* Poll download progress
*/
getDownloadProgressPollDelay(progress = {}, info = {}) {
if (progress.status === 'paused') return 1500;
const backend = String(
progress.download_backend || info.downloadBackend || info.download_backend || ''
).toLowerCase();
return backend === 'huggingface_xet' ? 200 : 1000;
},
async pollDownloadProgress(downloadId) {
const info = this.activeDownloads[downloadId];
if (!info) return;
@@ -2017,7 +2044,8 @@ export const resolveDownloadMethods = {
}
// Continue polling
setTimeout(() => this.pollDownloadProgress(downloadId), progress.status === 'paused' ? 1500 : 1000);
const pollDelay = this.getDownloadProgressPollDelay(progress, info);
setTimeout(() => this.pollDownloadProgress(downloadId), pollDelay);
} else if (progress.status === 'cancelling') {
this.renderDownloadSnapshot(downloadId, snapshot, { progressDiv, downloadBtn });
web/resolver/utils/render_format_methods.js
+43 -1
View File
@@ -350,9 +350,51 @@ export const renderFormatMethods = {
return remainingSeconds > 0 ? `${minutes}m ${remainingSeconds}s` : `${minutes}m`;
},
formatDownloadPercent(value) {
const percent = Math.max(0, Math.min(100, Number(value) || 0));
if (percent > 0 && percent < 0.1) return '<0.1';
if (percent > 0 && percent < 10) {
return percent.toFixed(1).replace(/\.0$/, '');
}
return String(Math.round(percent));
},
getDownloadDisplayProgress(progress = {}) {
const logicalDownloaded = Math.max(0, Number(progress.downloaded) || 0);
const logicalTotalSize = Math.max(0, Number(progress.total_size) || 0);
const logicalPercent = Math.max(0, Math.min(100, Number(progress.progress) || 0));
const backend = String(progress.download_backend || '').toLowerCase();
const transferDownloaded = Math.max(0, Number(progress.transfer_downloaded) || 0);
const transferTotalSize = Math.max(0, Number(progress.transfer_total_size) || 0);
if (backend === 'huggingface_xet' && logicalTotalSize > 0) {
return {
percent: Math.max(0, Math.min(100, (transferDownloaded / logicalTotalSize) * 100)),
downloaded: transferDownloaded,
totalSize: logicalTotalSize,
isTransfer: true,
isFinalizing: transferTotalSize > 0
&& transferDownloaded >= transferTotalSize
&& logicalTotalSize > 0
&& logicalDownloaded < logicalTotalSize
};
}
return {
percent: logicalPercent,
downloaded: logicalDownloaded,
totalSize: logicalTotalSize,
isTransfer: false,
isFinalizing: false
};
},
getDownloadEtaText(progress = {}) {
const backend = String(progress.download_backend || '').toLowerCase();
const totalSize = Number(progress.total_size) || 0;
const downloaded = Number(progress.downloaded) || 0;
const downloaded = backend === 'huggingface_xet'
? (Number(progress.transfer_downloaded) || 0)
: (Number(progress.downloaded) || 0);
const speed = Number(progress.speed) || 0;
if (totalSize <= 0 || downloaded <= 0 || speed <= 0 || downloaded >= totalSize) return '';