mirror of
https://github.com/calesthio/OpenMontage.git
synced 2026-08-25 17:42:50 +08:00
feat: wait on ComfyUI websocket feed instead of polling for completion
Resolves the "async generation" open question from the adapter plan. generate() now watches ComfyUI's websocket events (executing/progress/ execution_error) and reacts immediately instead of sleeping between REST polls, with an optional on_progress callback that comfyui_video uses to print step progress on long renders. websocket-client is an optional import; _wait() falls back to the original poll() loop (with the remaining time budget, not a fresh one) when it's unavailable or the connection drops, so resume_prompt_id recovery is unaffected either way. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
@@ -217,6 +217,23 @@ class ComfyUIVideo(BaseTool):
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._client = ComfyUIClient()
|
||||
self._last_progress_log = 0.0
|
||||
|
||||
def _log_progress(self, data: dict) -> None:
|
||||
"""Print a throttled progress line for long video renders.
|
||||
|
||||
Video jobs can run for tens of minutes; without this the process
|
||||
looks hung. Throttled to once per 10s since ComfyUI pushes a
|
||||
``progress`` event per sampling step, which would otherwise flood
|
||||
stdout on fast GPUs.
|
||||
"""
|
||||
now = time.monotonic()
|
||||
if now - self._last_progress_log < 10:
|
||||
return
|
||||
self._last_progress_log = now
|
||||
value, max_value = data.get("value"), data.get("max")
|
||||
if value is not None and max_value:
|
||||
print(f"[comfyui_video] step {value}/{max_value}")
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if not self._client.is_available():
|
||||
@@ -341,6 +358,7 @@ class ComfyUIVideo(BaseTool):
|
||||
timeout=inputs.get("timeout_seconds", 3600),
|
||||
interval=10,
|
||||
resume_prompt_id=inputs.get("resume_prompt_id"),
|
||||
on_progress=self._log_progress,
|
||||
)
|
||||
|
||||
except ComfyUIError as exc:
|
||||
|
||||
Reference in New Issue
Block a user