Files
diodiogod-TTS-Audio-Suite/utils/asr/progress_callback.py
diodiogod 273229976b Version 4.24.7
Fix CosyVoice3 narrator reference and interrupt handling

Implementation details:
- Restore connected opt_narrator resolution in unified text and SRT nodes
- Propagate connected reference text and character names through CosyVoice3 processing
- Add mid-generation interrupt polling for CosyVoice3
- Patch bundled Qwen3-ASR generate loop for reliable ComfyUI interrupts
- Mark vendored Qwen3-ASR interrupt logic as a TTS Audio Suite patch
2026-03-27 22:05:56 -03:00

52 lines
1.8 KiB
Python

"""
ASR Progress Callback for ComfyUI Integration
Token-level progress tracking with it/s and ETA.
"""
import time
import comfy.model_management as model_management
class ASRProgressStreamer:
def __init__(self, max_new_tokens: int, progress_bar=None, label: str = "ASR"):
self.max_new_tokens = max_new_tokens
self.progress_bar = progress_bar
self.label = label
self.generated_tokens = 0
self.start_time = time.time()
self.last_update = 0.0
def put(self, value):
if model_management.interrupt_processing:
raise InterruptedError("Qwen3-ASR generation interrupted by user")
# value is a batch of token ids; count tokens
try:
new_tokens = len(value)
except Exception:
new_tokens = 1
self.generated_tokens += new_tokens
if self.progress_bar is not None:
try:
self.progress_bar.update(new_tokens)
except Exception:
pass
now = time.time()
if now - self.last_update >= 0.5:
elapsed = now - self.start_time
its = self.generated_tokens / elapsed if elapsed > 0 else 0.0
bar_width = 12
total = self.max_new_tokens if self.max_new_tokens > 0 else 1
filled = int(min(1.0, self.generated_tokens / total) * bar_width)
bar = f"[{'█' * filled}{'░' * (bar_width - filled)}]"
print(f"\r {self.label} Progress: {bar} {self.generated_tokens}/{total} | {its:.1f} it/s", end="", flush=True)
self.last_update = now
def end(self):
elapsed = time.time() - self.start_time
avg = self.generated_tokens / elapsed if elapsed > 0 else 0.0
print(f"\r {self.label} Complete: {self.generated_tokens} tokens in {elapsed:.1f}s (avg {avg:.1f} it/s)" + " " * 20)