170 lines
8.8 KiB
Python
170 lines
8.8 KiB
Python
"""Read-only ComfyUI telemetry, embedded in the existing private normalizers.
|
|
|
|
Only the requesting Huobao job's client connection is observed. No queue mutation
|
|
or image previews are used. Sampling ETA is never presented as whole-video ETA.
|
|
"""
|
|
import asyncio as _hp_asyncio
|
|
import json as _hp_json
|
|
import time as _hp_time
|
|
import math as _hp_math
|
|
import statistics as _hp_stats
|
|
from urllib.parse import quote as _hp_quote
|
|
import httpx as _hp_httpx
|
|
import websockets as _hp_ws
|
|
|
|
_hp_watchers = {}
|
|
_hp_histories = {}
|
|
|
|
|
|
def _hp_phase(kind):
|
|
kind = str(kind).lower()
|
|
if 'sampler' in kind: return 'sampling'
|
|
if 'decode' in kind: return 'decoding'
|
|
if 'savevideo' in kind or 'combine' in kind or 'createvideo' in kind: return 'encoding'
|
|
if 'loader' in kind or kind == 'loadimage': return 'loading'
|
|
if 'caption' in kind: return 'postprocessing'
|
|
if 'encode' in kind or 'tovideo' in kind: return 'conditioning'
|
|
return 'running'
|
|
|
|
|
|
def _hp_signature(graph):
|
|
# Keep graph topology, model weights, image count, dimensions, FPS, duration,
|
|
# steps and sampler. Exclude content, random seeds and output filenames.
|
|
omitted = {'text', 'prompt', 'positive', 'negative', 'seed', 'noise_seed',
|
|
'filename_prefix', 'image', 'audio', 'video', 'file', 'path'}
|
|
result = {}
|
|
for node, value in graph.items():
|
|
inputs = {}
|
|
for key, item in value.get('inputs', {}).items():
|
|
if key in omitted and not isinstance(item, list): continue
|
|
if any(word in key.lower() for word in ['prompt', 'text', 'caption']) and not isinstance(item, list): continue
|
|
inputs[key] = item
|
|
result[node] = {'class_type': value.get('class_type'), 'inputs': inputs}
|
|
return _hp_json.dumps(result, sort_keys=True, separators=(',', ':'))
|
|
|
|
|
|
def _hp_history_estimate(graph, history):
|
|
signature = _hp_signature(graph)
|
|
durations = []
|
|
for entry in history.values():
|
|
status = entry.get('status', {})
|
|
prompt = entry.get('prompt', [])
|
|
if status.get('status_str') != 'success' or not status.get('completed') or len(prompt) < 3: continue
|
|
previous = prompt[2]
|
|
if not isinstance(previous, dict) or _hp_signature(previous) != signature: continue
|
|
times = {}; cached = []
|
|
for event in status.get('messages', []):
|
|
if len(event) < 2 or not isinstance(event[1], dict): continue
|
|
times[event[0]] = event[1].get('timestamp')
|
|
if event[0] == 'execution_cached': cached = event[1].get('nodes', [])
|
|
if any('sampler' in str(previous.get(str(n), {}).get('class_type', '')).lower() for n in cached): continue
|
|
start, end = times.get('execution_start'), times.get('execution_success')
|
|
if not isinstance(start, (int, float)) or not isinstance(end, (int, float)): continue
|
|
duration = (end - start) / 1000
|
|
if 2 <= duration <= 86400: durations.append(duration)
|
|
result = {'historySamples': len(durations)}
|
|
if len(durations) >= 2:
|
|
# A range, not a guarantee; model loading and shared GPU contention vary.
|
|
result.update(historyMinSeconds=round(min(durations) * .8), historyMaxSeconds=round(max(durations) * 1.2))
|
|
return result
|
|
|
|
|
|
def _hp_event(state, message, timestamp):
|
|
data = message.get('data', {})
|
|
if not isinstance(data, dict): return
|
|
event = message.get('type')
|
|
if data.get('prompt_id') not in (None, state['prompt']): return
|
|
if event == 'executing':
|
|
node = data.get('node')
|
|
if node is None: return
|
|
node = str(node)
|
|
if node != state.get('node'):
|
|
state.update(node=node, samples=[], value=None, total=None)
|
|
state['phase'] = _hp_phase(state['classes'].get(node))
|
|
state['updated'] = timestamp
|
|
elif event == 'progress' and data.get('prompt_id') == state['prompt']:
|
|
value, total = data.get('value'), data.get('max')
|
|
if not all(type(x) in (int, float) and _hp_math.isfinite(x) for x in (value, total)): return
|
|
if not 0 <= value <= total <= 1000000 or total <= 0: return
|
|
node = str(data.get('node', state.get('node', '')))
|
|
if node != state.get('node') or total != state.get('total') or value < (state.get('value') or 0):
|
|
state.update(node=node, samples=[])
|
|
if state.get('value') != value or not state['samples']:
|
|
state['samples'] = (state['samples'] + [(timestamp, value)])[-8:]
|
|
state.update(value=value, total=total, updated=timestamp, phase=_hp_phase(state['classes'].get(node)))
|
|
|
|
|
|
def _hp_snapshot(state, timestamp):
|
|
result = {'phase': state.get('phase', 'running')}
|
|
if not state.get('connected') or timestamp - state.get('updated', 0) > 600:
|
|
return {'phase': 'running'}
|
|
if state.get('value') is None: return result
|
|
result.update(step=state['value'], totalSteps=state['total'], stepUpdatedAt=round(state['updated'] * 1000))
|
|
samples = state['samples']
|
|
if state.get('phase') == 'sampling' and len(samples) >= 3:
|
|
seconds = samples[-1][0] - samples[0][0]; steps = samples[-1][1] - samples[0][1]
|
|
if seconds >= 5 and steps > 0 and state['value'] < state['total']:
|
|
speed = seconds / steps
|
|
if timestamp - state['updated'] <= max(30, speed * 2.5):
|
|
remaining = max(1, speed * (state['total'] - state['value']) - (timestamp - state['updated']))
|
|
result.update(samplingMinSeconds=round(remaining * .8), samplingMaxSeconds=round(remaining * 1.3))
|
|
return result
|
|
|
|
|
|
async def _hp_watch(state, base, headers, client_id):
|
|
uri = base.replace('https://', 'wss://').replace('http://', 'ws://') + '/ws?clientId=' + _hp_quote(client_id, safe='')
|
|
while _hp_time.time() - state['requested'] < 120:
|
|
try:
|
|
async with _hp_ws.connect(uri, additional_headers=headers, max_size=32 * 1024 * 1024, open_timeout=5) as ws:
|
|
state['connected'] = True
|
|
# This observer needs text progress only, never preview images.
|
|
await ws.send(_hp_json.dumps({'type':'feature_flags','data':{'supports_preview_metadata':True}}))
|
|
while _hp_time.time() - state['requested'] < 120:
|
|
try: message = await _hp_asyncio.wait_for(ws.recv(), 15)
|
|
except _hp_asyncio.TimeoutError: continue
|
|
if isinstance(message, str):
|
|
try: _hp_event(state, _hp_json.loads(message), _hp_time.time())
|
|
except (ValueError, TypeError, AttributeError): pass
|
|
except _hp_asyncio.CancelledError: raise
|
|
except Exception:
|
|
await _hp_asyncio.sleep(5)
|
|
finally: state['connected'] = False
|
|
|
|
|
|
async def huobao_progress(provider, prompt_id, queue, headers):
|
|
"""Failure of progress reporting must never fail or resubmit a generation."""
|
|
timestamp = _hp_time.time()
|
|
running = queue.get('queue_running', []); pending = queue.get('queue_pending', [])
|
|
row = next((x for x in running + pending if len(x) > 3 and x[1] == prompt_id), None)
|
|
if row is None: return None
|
|
result = {'observedAt': round(timestamp * 1000), 'phase': 'queued' if row in pending else 'running'}
|
|
if row in pending:
|
|
result.update(queuePosition=pending.index(row) + 1, queueAhead=len(running) + pending.index(row))
|
|
try:
|
|
previous = _hp_watchers.get(provider.id)
|
|
if previous and (previous['prompt'] != prompt_id and row in running or previous['task'].done()):
|
|
previous['task'].cancel(); _hp_watchers.pop(provider.id, None); previous = None
|
|
if row in running:
|
|
client_id = row[3].get('client_id')
|
|
# Huobao's normalizer generates unique canvas-* client IDs. Never
|
|
# replace a human's browser websocket or guess another job's state.
|
|
if isinstance(client_id, str) and client_id.startswith('canvas-'):
|
|
if previous is None:
|
|
previous = {'prompt':prompt_id, 'classes':{str(k):v.get('class_type') for k,v in row[2].items()},
|
|
'requested':timestamp, 'samples':[], 'connected':False}
|
|
previous['task'] = _hp_asyncio.create_task(_hp_watch(previous, provider.base_url, headers, client_id))
|
|
_hp_watchers[provider.id] = previous
|
|
previous['requested'] = timestamp
|
|
if previous['prompt'] == prompt_id: result.update(_hp_snapshot(previous, timestamp))
|
|
cached = _hp_histories.get(provider.id)
|
|
if cached is None or timestamp - cached[0] > 90:
|
|
async with _hp_httpx.AsyncClient(timeout=5) as client:
|
|
response = await client.get(provider.base_url + '/history', params={'max_items':80}, headers=headers)
|
|
response.raise_for_status(); history = response.json()
|
|
_hp_histories[provider.id] = (timestamp, history)
|
|
else: history = cached[1]
|
|
result.update(_hp_history_estimate(row[2], history))
|
|
except Exception:
|
|
pass
|
|
return result
|