Skip to content

Commit 712cca3

Browse files
authored
feat: throttle ProgressBar updates to reduce WebSocket flooding (#11504)
1 parent ac4d8ea commit 712cca3

File tree

1 file changed

+30
-1
lines changed

1 file changed

+30
-1
lines changed

comfy/utils.py

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
from einops import rearrange
3131
from comfy.cli_args import args
3232
import json
33+
import time
3334

3435
MMAP_TORCH_FILES = args.mmap_torch_files
3536
DISABLE_MMAP = args.disable_mmap
@@ -1097,13 +1098,19 @@ def set_progress_bar_global_hook(function):
10971098
global PROGRESS_BAR_HOOK
10981099
PROGRESS_BAR_HOOK = function
10991100

1101+
# Throttle settings for progress bar updates to reduce WebSocket flooding
1102+
PROGRESS_THROTTLE_MIN_INTERVAL = 0.1 # 100ms minimum between updates
1103+
PROGRESS_THROTTLE_MIN_PERCENT = 0.5 # 0.5% minimum progress change
1104+
11001105
class ProgressBar:
11011106
def __init__(self, total, node_id=None):
11021107
global PROGRESS_BAR_HOOK
11031108
self.total = total
11041109
self.current = 0
11051110
self.hook = PROGRESS_BAR_HOOK
11061111
self.node_id = node_id
1112+
self._last_update_time = 0.0
1113+
self._last_sent_value = -1
11071114

11081115
def update_absolute(self, value, total=None, preview=None):
11091116
if total is not None:
@@ -1112,7 +1119,29 @@ def update_absolute(self, value, total=None, preview=None):
11121119
value = self.total
11131120
self.current = value
11141121
if self.hook is not None:
1115-
self.hook(self.current, self.total, preview, node_id=self.node_id)
1122+
current_time = time.perf_counter()
1123+
is_first = (self._last_sent_value < 0)
1124+
is_final = (value >= self.total)
1125+
has_preview = (preview is not None)
1126+
1127+
# Always send immediately for previews, first update, or final update
1128+
if has_preview or is_first or is_final:
1129+
self.hook(self.current, self.total, preview, node_id=self.node_id)
1130+
self._last_update_time = current_time
1131+
self._last_sent_value = value
1132+
return
1133+
1134+
# Apply throttling for regular progress updates
1135+
if self.total > 0:
1136+
percent_changed = ((value - max(0, self._last_sent_value)) / self.total) * 100
1137+
else:
1138+
percent_changed = 100
1139+
time_elapsed = current_time - self._last_update_time
1140+
1141+
if time_elapsed >= PROGRESS_THROTTLE_MIN_INTERVAL and percent_changed >= PROGRESS_THROTTLE_MIN_PERCENT:
1142+
self.hook(self.current, self.total, preview, node_id=self.node_id)
1143+
self._last_update_time = current_time
1144+
self._last_sent_value = value
11161145

11171146
def update(self, value):
11181147
self.update_absolute(self.current + value)

0 commit comments

Comments
 (0)