WildAi
8404100f31
- streaming per-tensor assign (tensor.clone) replaces merged-dict load_state_dict_sharded — mmap views were keeping all shard files resident after partial offload (+15GB ghost RSS, Pin-error flood); private RAM copies pin reliably and release shard mappings early - console progress bar during TTS/ASR generation (progress_utils) - GGUF read-only mmap no longer warns (copy before from_numpy) - demote internal load-narration logs to DEBUG (success-only INFO) - drop dead deps (einops, tokenizers, s3tokenizer, conformer), add accelerate to pyproject (4-bit path needs it), requirements trimmed to non-ComfyUI-bundled packages; parity guard tests
72 lines
2.7 KiB
Python
72 lines
2.7 KiB
Python
"""Progress reporting helpers for VibeVoice nodes.
|
|
|
|
ComfyUI surfaces inference progress in two places:
|
|
|
|
* the **frontend** bar, driven by ``comfy.utils.ProgressBar`` (it posts
|
|
updates to a server hook which relays them to the UI), and
|
|
* the **console** bar, the standard ``tqdm`` bar that every sampler prints to
|
|
the terminal (the k-diffusion loops run under ``comfy.utils.model_trange``).
|
|
|
|
``comfy.utils.ProgressBar`` only drives the frontend bar — it never writes to
|
|
the terminal. VibeVoice's long-running ``generate()`` loops therefore showed a
|
|
UI bar but no console bar, unlike the samplers. This module provides a
|
|
drop-in wrapper that drives both.
|
|
"""
|
|
|
|
import comfy.utils
|
|
from tqdm.auto import tqdm
|
|
|
|
|
|
class ProgressBarWithConsole:
|
|
"""``comfy.utils.ProgressBar`` plus the standard tqdm console bar.
|
|
|
|
Mirrors the ``ProgressBar`` interface (``update_absolute`` / ``update`` /
|
|
``total`` / ``current``) so it can replace ``ProgressBar`` at the call
|
|
sites unchanged. Every update is forwarded to the real frontend bar and
|
|
also mirrored to a lazily-created ``tqdm`` console bar, matching the
|
|
sampler bar's look (no description, absolute position, auto total
|
|
re-sync).
|
|
|
|
The console bar is created lazily on the first update so loops that only
|
|
learn their real total at runtime (the streaming path starts at
|
|
``total=1`` and corrects it on the first callback) do not flash a wrong
|
|
total. It is gated on ``comfy.utils.PROGRESS_BAR_ENABLED`` — the same flag
|
|
``nodes.py`` reads to build the samplers' ``disable_pbar`` — so disabling
|
|
the ComfyUI progress bar silences the console bar here too.
|
|
"""
|
|
|
|
def __init__(self, total, node_id=None):
|
|
self._ui = comfy.utils.ProgressBar(total, node_id=node_id)
|
|
self._console = None
|
|
|
|
@property
|
|
def total(self):
|
|
return self._ui.total
|
|
|
|
@property
|
|
def current(self):
|
|
return self._ui.current
|
|
|
|
def update_absolute(self, value, total=None, preview=None):
|
|
self._ui.update_absolute(value, total=total, preview=preview)
|
|
self._sync_console()
|
|
|
|
def update(self, value):
|
|
self.update_absolute(self._ui.current + value)
|
|
|
|
def _sync_console(self):
|
|
if self._console is None:
|
|
if not comfy.utils.PROGRESS_BAR_ENABLED:
|
|
return
|
|
self._console = tqdm(total=self._ui.total)
|
|
elif self._console.total != self._ui.total:
|
|
self._console.total = self._ui.total
|
|
self._console.n = self._ui.current
|
|
self._console.refresh()
|
|
|
|
def close(self):
|
|
"""Finalize the console bar. Safe to call when no bar was created."""
|
|
if self._console is not None:
|
|
self._console.close()
|
|
self._console = None
|