Spaces:
Running
Running
| """HuggingFace download progress — one monkey-patch, every `hf_hub_download` | |
| reports bytes downloaded through a central callback. | |
| `huggingface_hub` uses tqdm for progress bars; we subclass it, intercept | |
| `update()` calls, and forward (filename, downloaded_bytes, total_bytes) to | |
| whatever callback is registered. No changes to calling sites across | |
| transformers / mlx_whisper / diffusers / accelerate — they all route through | |
| `hf_hub_download`, which uses the patched tqdm. | |
| Usage: | |
| from utils.hf_progress import install, register_listener, unregister_listener | |
| install() # once at app startup | |
| listener_id = register_listener(lambda ev: print(ev)) | |
| # …models download, listener fires… | |
| unregister_listener(listener_id) | |
| """ | |
| from __future__ import annotations | |
| import contextvars | |
| import itertools | |
| import logging | |
| import threading | |
| from typing import Callable, Optional | |
| logger = logging.getLogger("omnivoice.hf_progress") | |
| # Context-scoped active repo_id. Set in the install/delete handler so every | |
| # tqdm event fired while a snapshot_download runs can be stamped with the | |
| # originating repo, letting the frontend route per-file events to the right | |
| # row instead of heuristically matching filename substrings. | |
| current_repo_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar( | |
| "omnivoice_hf_progress_repo_id", default=None, | |
| ) | |
| # Event shape forwarded to listeners. Typed loosely on purpose — SSE encodes | |
| # it as JSON so consumers read the dict directly. | |
| # { | |
| # "filename": str, # desc on the tqdm bar, usually the HF file path | |
| # "downloaded": int, # bytes pulled so far | |
| # "total": int | None, # total bytes or None if unknown | |
| # "pct": float, # 0.0-1.0 (or 0.0 if total unknown) | |
| # "phase": "start"|"progress"|"done", | |
| # } | |
| ProgressEvent = dict | |
| Listener = Callable[[ProgressEvent], None] | |
| _listeners: dict[int, Listener] = {} | |
| _listener_lock = threading.Lock() | |
| _listener_counter = itertools.count(1) | |
| _installed = False | |
| _install_lock = threading.Lock() | |
| def register_listener(cb: Listener) -> int: | |
| """Register a callback that receives progress events. Returns an id that | |
| can be passed to `unregister_listener` when the listener is done.""" | |
| with _listener_lock: | |
| lid = next(_listener_counter) | |
| _listeners[lid] = cb | |
| return lid | |
| def unregister_listener(lid: int) -> None: | |
| with _listener_lock: | |
| _listeners.pop(lid, None) | |
| def _emit(event: ProgressEvent) -> None: | |
| """Fan out to all registered listeners. Never raise — a bad listener | |
| shouldn't break a download.""" | |
| # Stamp the active repo_id so frontends can route events to the right | |
| # row. Only set when this emit is happening inside an install handler. | |
| rid = current_repo_id.get() | |
| if rid is not None and "repo_id" not in event: | |
| event = {**event, "repo_id": rid} | |
| with _listener_lock: | |
| listeners = list(_listeners.values()) | |
| for cb in listeners: | |
| try: | |
| cb(event) | |
| except Exception as e: # noqa: BLE001 | |
| logger.debug("hf_progress listener raised: %s", e) | |
| def emit(event: ProgressEvent) -> None: | |
| """Public emit — lets non-tqdm operations (delete, verify, etc.) push | |
| lifecycle events onto the same SSE stream.""" | |
| _emit(event) | |
| class SafeFileWrapper: | |
| def __init__(self, fp): | |
| self.fp = fp | |
| self._is_safe_wrapper = True | |
| def write(self, s): | |
| try: | |
| self.fp.write(s) | |
| except OSError: | |
| pass | |
| def flush(self): | |
| try: | |
| getattr(self.fp, 'flush', lambda: None)() | |
| except OSError: | |
| pass | |
| def __getattr__(self, name): | |
| return getattr(self.fp, name) | |
| def install() -> None: | |
| """Monkey-patch `huggingface_hub`'s tqdm so every download reports to our | |
| listeners. Safe to call multiple times — second call is a no-op.""" | |
| global _installed | |
| with _install_lock: | |
| if _installed: | |
| return | |
| # `huggingface_hub.utils.__init__` does `from .tqdm import tqdm`, | |
| # which shadows the `tqdm` SUBMODULE with the CLASS of the same name | |
| # when accessed via attribute lookup. Pull the real module out of | |
| # sys.modules after an explicit import so we patch the right thing. | |
| try: | |
| import sys | |
| import huggingface_hub.utils.tqdm # noqa: F401 | |
| hf_tqdm_module = sys.modules.get("huggingface_hub.utils.tqdm") | |
| if hf_tqdm_module is None: | |
| raise ImportError("huggingface_hub.utils.tqdm not in sys.modules after import") | |
| except Exception as e: # noqa: BLE001 | |
| logger.warning( | |
| "hf_progress.install: huggingface_hub.utils.tqdm missing (%s); " | |
| "progress tracking disabled.", e, | |
| ) | |
| return | |
| original = getattr(hf_tqdm_module, "tqdm", None) | |
| if original is None or not isinstance(original, type): | |
| logger.warning("hf_progress.install: no `tqdm` class on the module; aborting") | |
| return | |
| class TrackedTqdm(original): # type: ignore[misc,valid-type] | |
| """tqdm subclass that emits a progress event on every update.""" | |
| _last_emit_time: float = 0.0 | |
| def status_printer(file): | |
| if file is not None and not getattr(file, "_is_safe_wrapper", False): | |
| file = SafeFileWrapper(file) | |
| try: | |
| return original.status_printer(file) | |
| except Exception: | |
| return lambda s: None | |
| def __init__(self, *args, **kwargs): | |
| if 'file' in kwargs and kwargs['file'] is not None and not getattr(kwargs['file'], "_is_safe_wrapper", False): | |
| kwargs['file'] = SafeFileWrapper(kwargs['file']) | |
| try: | |
| super().__init__(*args, **kwargs) | |
| except OSError: | |
| pass | |
| if hasattr(self, 'fp') and getattr(self, 'fp', None) is not None and not getattr(self.fp, "_is_safe_wrapper", False): | |
| self.fp = SafeFileWrapper(self.fp) | |
| import time as _t | |
| self._last_emit_time = _t.monotonic() | |
| try: | |
| desc = getattr(self, "desc", None) | |
| total = int(getattr(self, "total", 0) or 0) | |
| _emit({ | |
| "filename": str(desc or "download"), | |
| "downloaded": 0, | |
| "total": total, | |
| "pct": 0.0, | |
| "phase": "start", | |
| }) | |
| except Exception: | |
| pass | |
| def _emit_progress(self): | |
| """Emit current state as a progress event.""" | |
| try: | |
| desc = getattr(self, "desc", None) | |
| total = int(getattr(self, "total", 0) or 0) | |
| done = int(getattr(self, "n", 0) or 0) | |
| pct = (done / total) if total > 0 else 0.0 | |
| # Pull rate from tqdm's own calculations if available | |
| rate = None | |
| try: | |
| rate = self.format_dict.get("rate") | |
| except Exception: | |
| pass | |
| event = { | |
| "filename": str(desc or "download"), | |
| "downloaded": done, | |
| "total": total, | |
| "pct": pct, | |
| "phase": "done" if (total > 0 and done >= total) else "progress", | |
| } | |
| if rate and rate > 0: | |
| event["rate"] = rate # bytes/sec from tqdm | |
| _emit(event) | |
| except Exception: | |
| pass | |
| def update(self, n=1): | |
| try: | |
| super().update(n) | |
| except OSError: | |
| pass | |
| import time as _t | |
| now = _t.monotonic() | |
| # Throttle: emit at most every 0.3s to avoid flooding SSE | |
| if (now - self._last_emit_time) >= 0.3: | |
| self._last_emit_time = now | |
| self._emit_progress() | |
| def display(self, msg=None, pos=None): | |
| """tqdm calls display() on its refresh cycle; piggyback for | |
| periodic emits even when update() intervals are large.""" | |
| import time as _t | |
| now = _t.monotonic() | |
| if (now - self._last_emit_time) >= 0.5: | |
| self._last_emit_time = now | |
| self._emit_progress() | |
| try: | |
| return super().display(msg, pos) | |
| except OSError: | |
| pass | |
| def close(self): | |
| try: | |
| super().close() | |
| except OSError: | |
| pass | |
| # Stash the original for inspection / uninstall, then swap. | |
| hf_tqdm_module._omnivoice_original_tqdm = original # type: ignore[attr-defined] | |
| hf_tqdm_module.tqdm = TrackedTqdm # type: ignore[assignment] | |
| _installed = True | |
| logger.info("hf_progress: installed tqdm patch on huggingface_hub.utils.tqdm") | |