Fix image prompt swipes and RunPod cold starts
| @@ -58,3 +58,5 @@ public/scripts/extensions/third-party | |||
| 58 | yarn.lock | 58 | yarn.lock |
| 59 | *.code-workspace | 59 | *.code-workspace |
| 60 | test-results/ | 60 | test-results/ |
| 61 | __pycache__/ | ||
| 62 | *.py[cod] | ||
| @@ -7,7 +7,7 @@ changing one, update both (scp to the NAS path + commit here). | |||
| 7 | | File | Deployed at | Purpose | | 7 | | File | Deployed at | Purpose | |
| 8 | |---|---|---| | 8 | |---|---|---| |
| 9 | | `whisper-lazy-proxy.py` | `/mnt/user/appdata/whisper-lazy/proxy.py` (container `whisper-lazy`, :10300) | Wyoming lazy proxy: answers HA availability checks from cache, starts `faster-whisper` on demand, stops it after 15 min idle (frees GPU VRAM). | | 9 | | `whisper-lazy-proxy.py` | `/mnt/user/appdata/whisper-lazy/proxy.py` (container `whisper-lazy`, :10300) | Wyoming lazy proxy: answers HA availability checks from cache, starts `faster-whisper` on demand, stops it after 15 min idle (frees GPU VRAM). | |
| 10 | | `runpod-lazy-proxy.py` | `/mnt/user/appdata/runpod-lazy/proxy.py` (container `runpod-lazy`, :8189) | ComfyUI-compatible lazy proxy for a volumeless RunPod GPU pod: wake-on-generate with model-group detection, `/lazy/status|warmup|shutdown|ping|catalog` control API for the SillyTavern warmup UI, NAS-side 15-min idle terminate. Model-set changes on a running pod download in place via the worker image's model-manager (port 8189, LRU disk eviction); pods are only recreated when none exists or the image predates the manager. | | 10 | | `runpod-lazy-proxy.py` | `/mnt/user/appdata/runpod-lazy/proxy.py` (container `runpod-lazy`, :8189) | ComfyUI-compatible on-demand proxy for a volumeless RunPod GPU pod: only `/lazy/warmup` starts or provisions, while cold ComfyUI requests return 503 for provider failover. `/lazy/status|warmup|shutdown|ping|catalog` provide SillyTavern controls and a NAS-side 15-min idle terminate. Model-set changes requested by Warm up download in place via the worker image's model-manager (port 8189, LRU disk eviction); pods are only recreated when none exists or the image predates the manager. | |
| 11 | | `runpod-pod.sh` | `/boot/config/runpod-pod.sh` | CLI for the runpod-lazy control API (`status|warmup|shutdown|ping|watch`). | | 11 | | `runpod-pod.sh` | `/boot/config/runpod-pod.sh` | CLI for the runpod-lazy control API (`status|warmup|shutdown|ping|watch`). | |
| 12 | 12 | ||
| 13 | Related: the RunPod worker image (ComfyUI + custom nodes + boot-time model | 13 | Related: the RunPod worker image (ComfyUI + custom nodes + boot-time model |
| @@ -14,14 +14,15 @@ Control API (CORS-enabled; ST warmup UI + runpod-pod.sh CLI): | |||
| 14 | GET /lazy/status -> {"state": red|orange|green, "model": ..., ...} | 14 | GET /lazy/status -> {"state": red|orange|green, "model": ..., ...} |
| 15 | POST /lazy/warmup -> start pod for the active catalog entry | 15 | POST /lazy/warmup -> start pod for the active catalog entry |
| 16 | POST /lazy/shutdown -> terminate pod now | 16 | POST /lazy/shutdown -> terminate pod now |
| 17 | POST /lazy/ping -> keepalive (extends idle timer; never starts a pod) | 17 | POST /lazy/ping -> keepalive for a ready pod (never starts a pod) |
| 18 | POST /lazy/catalog -> {"models": [...], "active": "<value>"} store catalog; | 18 | POST /lazy/catalog -> {"models": [...], "active": "<value>"} store catalog; |
| 19 | re-provision the pod if the active model changed | 19 | never starts or provisions a pod |
| 20 | 20 | ||
| 21 | Wake semantics: a ComfyUI /prompt arriving with no ready pod matches the | 21 | Start semantics: only POST /lazy/warmup may provision or change models. ComfyUI |
| 22 | workflow against catalog filenames and provisions for that entry (falls back | 22 | requests return 503 unless a pod is already ready, allowing SillyTavern's image |
| 23 | to the active entry). The idle timer lives here (NAS side): IDLE_SECONDS after | 23 | target fallback chain to continue without waking this target. The idle timer |
| 24 | the last activity signal the pod is terminated, frontend or no frontend. | 24 | lives here (NAS side): IDLE_SECONDS after the last activity signal the pod is |
| 25 | terminated, frontend or no frontend. | ||
| 25 | """ | 26 | """ |
| 26 | import json | 27 | import json |
| 27 | import os | 28 | import os |
| @@ -45,7 +46,6 @@ POD_NAME = os.environ.get('POD_NAME', 'comfyui-lazy') | |||
| 45 | START_TIMEOUT = int(os.environ.get('START_TIMEOUT', '1500')) | 46 | START_TIMEOUT = int(os.environ.get('START_TIMEOUT', '1500')) |
| 46 | HF_TOKEN = os.environ.get('HF_TOKEN', '') | 47 | HF_TOKEN = os.environ.get('HF_TOKEN', '') |
| 47 | COMFY_ARGS = os.environ.get('COMFY_ARGS', '--listen 0.0.0.0 --port 8188 --use-pytorch-cross-attention') | 48 | COMFY_ARGS = os.environ.get('COMFY_ARGS', '--listen 0.0.0.0 --port 8188 --use-pytorch-cross-attention') |
| 48 | CACHE_FILE = os.environ.get('CACHE_FILE', '/app/system_stats.cache') | ||
| 49 | CATALOG_FILE = os.environ.get('CATALOG_FILE', '/app/catalog.json') | 49 | CATALOG_FILE = os.environ.get('CATALOG_FILE', '/app/catalog.json') |
| 50 | 50 | ||
| 51 | state = { | 51 | state = { |
| @@ -56,6 +56,7 @@ state = { | |||
| 56 | 'since': time.time(), | 56 | 'since': time.time(), |
| 57 | 'last': time.monotonic(), | 57 | 'last': time.monotonic(), |
| 58 | 'ensuring': 0, # active ensure_pod waiters (status skips probing then) | 58 | 'ensuring': 0, # active ensure_pod waiters (status skips probing then) |
| 59 | 'control_epoch': 0, # incremented by shutdown to cancel in-flight warmups | ||
| 59 | 'lock': threading.Lock(), | 60 | 'lock': threading.Lock(), |
| 60 | } | 61 | } |
| 61 | 62 | ||
| @@ -209,9 +210,7 @@ def upstream_ready(pod_id): | |||
| 209 | req = urllib.request.Request(pod_url(pod_id) + '/system_stats', headers={'User-Agent': UA}) | 210 | req = urllib.request.Request(pod_url(pod_id) + '/system_stats', headers={'User-Agent': UA}) |
| 210 | with urllib.request.urlopen(req, timeout=5) as resp: | 211 | with urllib.request.urlopen(req, timeout=5) as resp: |
| 211 | if resp.status == 200: | 212 | if resp.status == 200: |
| 212 | data = resp.read() | 213 | resp.read() |
| 213 | with open(CACHE_FILE, 'wb') as f: | ||
| 214 | f.write(data) | ||
| 215 | return True | 214 | return True |
| 216 | except Exception: | 215 | except Exception: |
| 217 | pass | 216 | pass |
| @@ -245,49 +244,67 @@ def _create_pod_with_retries(values): | |||
| 245 | raise RuntimeError(f'could not create pod: {last_err}') | 244 | raise RuntimeError(f'could not create pod: {last_err}') |
| 246 | 245 | ||
| 247 | 246 | ||
| 248 | def ensure_pod(values=None, wait=True): | 247 | def ensure_pod(values=None, wait=True, control_epoch=None): |
| 249 | """Ensures a pod holding the catalog values (None = active selection). | 248 | """Ensures a pod holding the catalog values (None = active selection). |
| 250 | 249 | ||
| 251 | An existing pod gets missing models downloaded IN PLACE through the in-pod | 250 | An existing pod gets missing models downloaded IN PLACE through the in-pod |
| 252 | model manager; recreation only happens for old-image pods without one.""" | 251 | model manager; recreation only happens for old-image pods without one.""" |
| 252 | def check_cancelled(): | ||
| 253 | if control_epoch is not None and control_epoch != state['control_epoch']: | ||
| 254 | raise RuntimeError('pod warmup cancelled') | ||
| 255 | |||
| 253 | values = values or active_values() | 256 | values = values or active_values() |
| 254 | key = values_key(values) | 257 | key = values_key(values) |
| 255 | files = needed_files(values) | 258 | files = needed_files(values) |
| 256 | dests = [f['dest'].lstrip('/') for f in files] | 259 | dests = [f['dest'].lstrip('/') for f in files] |
| 257 | created = False | 260 | created = False |
| 258 | with state['lock']: | 261 | with state['lock']: |
| 262 | check_cancelled() | ||
| 259 | pod_id, have, gpu = (state['pod_id'], state['model'], state['gpu']) | 263 | pod_id, have, gpu = (state['pod_id'], state['model'], state['gpu']) |
| 260 | if not pod_id: | 264 | if not pod_id: |
| 261 | pod_id, have, gpu = find_pod() | 265 | pod_id, have, gpu = find_pod() |
| 262 | if not pod_id: | 266 | if not pod_id: |
| 263 | pod_id, gpu = _create_pod_with_retries(values) | 267 | pod_id, gpu = _create_pod_with_retries(values) |
| 264 | have, created = key, True | 268 | have, created = key, True |
| 265 | state['pod_id'], state['model'], state['gpu'] = pod_id, have, gpu | 269 | check_cancelled() |
| 266 | state['last'] = time.monotonic() | 270 | state.update({ |
| 271 | 'pod_id': pod_id, | ||
| 272 | 'model': have, | ||
| 273 | 'gpu': gpu, | ||
| 274 | 'last': time.monotonic(), | ||
| 275 | 'phase': 'green' if state['phase'] == 'green' else 'orange', | ||
| 276 | }) | ||
| 267 | 277 | ||
| 268 | if state['phase'] != 'green': | ||
| 269 | state['phase'] = 'orange' | ||
| 270 | if not wait: | 278 | if not wait: |
| 271 | return pod_id | 279 | return pod_id |
| 272 | # Freshly created pods boot with the right manifest; existing pods need an | 280 | # Freshly created pods boot with the right manifest; existing pods need an |
| 273 | # /ensure pushed once the manager answers. | 281 | # /ensure pushed once the manager answers. |
| 274 | ensured = created or not files | 282 | ensured = created or not files |
| 275 | no_manager_strikes = 0 | 283 | no_manager_strikes = 0 |
| 284 | with state['lock']: | ||
| 285 | check_cancelled() | ||
| 276 | state['ensuring'] += 1 | 286 | state['ensuring'] += 1 |
| 277 | try: | 287 | try: |
| 278 | deadline = time.monotonic() + START_TIMEOUT | 288 | deadline = time.monotonic() + START_TIMEOUT |
| 279 | while time.monotonic() < deadline: | 289 | while time.monotonic() < deadline: |
| 280 | # Booting/downloading counts as activity, else the reaper would | 290 | # Booting/downloading counts as activity, else the reaper would |
| 281 | # kill a warming pod mid-download. | 291 | # kill a warming pod mid-download. |
| 292 | with state['lock']: | ||
| 293 | check_cancelled() | ||
| 294 | if state['pod_id'] != pod_id: | ||
| 295 | raise RuntimeError('pod replaced while waiting') | ||
| 282 | state['last'] = time.monotonic() | 296 | state['last'] = time.monotonic() |
| 283 | if not ensured: | 297 | if not ensured: |
| 284 | try: | 298 | try: |
| 285 | mm_request(pod_id, 'POST', '/ensure', {'files': files, 'priority': True}) | 299 | mm_request(pod_id, 'POST', '/ensure', {'files': files, 'priority': True}) |
| 300 | except Exception: | ||
| 301 | pass # manager still booting, or old image without one | ||
| 302 | else: | ||
| 286 | ensured = True | 303 | ensured = True |
| 304 | with state['lock']: | ||
| 305 | check_cancelled() | ||
| 287 | state['model'] = values_key(key_values(state['model']) | set(values)) | 306 | state['model'] = values_key(key_values(state['model']) | set(values)) |
| 288 | log(f'in-place ensure requested: {key}') | 307 | log(f'in-place ensure requested: {key}') |
| 289 | except Exception: | ||
| 290 | pass # manager still booting, or old image without one | ||
| 291 | ready = upstream_ready(pod_id) | 308 | ready = upstream_ready(pod_id) |
| 292 | models_ok = True | 309 | models_ok = True |
| 293 | if not created and files and ensured: | 310 | if not created and files and ensured: |
| @@ -307,12 +324,18 @@ def ensure_pod(values=None, wait=True): | |||
| 307 | # Comfy is up but there is no manager (old image): last resort. | 324 | # Comfy is up but there is no manager (old image): last resort. |
| 308 | log(f'pod has models={state["model"]}, need {key} - recreating (no model manager)') | 325 | log(f'pod has models={state["model"]}, need {key} - recreating (no model manager)') |
| 309 | with state['lock']: | 326 | with state['lock']: |
| 327 | check_cancelled() | ||
| 310 | terminate(pod_id) | 328 | terminate(pod_id) |
| 311 | pod_id, gpu = _create_pod_with_retries(values) | 329 | pod_id, gpu = _create_pod_with_retries(values) |
| 330 | check_cancelled() | ||
| 312 | state.update({'pod_id': pod_id, 'model': key, 'gpu': gpu, 'last': time.monotonic()}) | 331 | state.update({'pod_id': pod_id, 'model': key, 'gpu': gpu, 'last': time.monotonic()}) |
| 313 | created, ensured = True, True | 332 | created, ensured = True, True |
| 314 | continue | 333 | continue |
| 315 | if ready and models_ok: | 334 | if ready and models_ok: |
| 335 | with state['lock']: | ||
| 336 | check_cancelled() | ||
| 337 | if state['pod_id'] != pod_id: | ||
| 338 | raise RuntimeError('pod replaced while waiting') | ||
| 316 | if state['phase'] != 'green': | 339 | if state['phase'] != 'green': |
| 317 | state['phase'] = 'green' | 340 | state['phase'] = 'green' |
| 318 | state['since'] = time.time() | 341 | state['since'] = time.time() |
| @@ -323,6 +346,8 @@ def ensure_pod(values=None, wait=True): | |||
| 323 | state['prefetch_pod'] = pod_id | 346 | state['prefetch_pod'] = pod_id |
| 324 | threading.Thread(target=prefetch_rest, args=(pod_id,), daemon=True).start() | 347 | threading.Thread(target=prefetch_rest, args=(pod_id,), daemon=True).start() |
| 325 | return pod_id | 348 | return pod_id |
| 349 | with state['lock']: | ||
| 350 | check_cancelled() | ||
| 326 | if state['phase'] == 'green': | 351 | if state['phase'] == 'green': |
| 327 | state['phase'] = 'orange' # in-place download in progress | 352 | state['phase'] = 'orange' # in-place download in progress |
| 328 | if state['pod_id'] != pod_id: | 353 | if state['pod_id'] != pod_id: |
| @@ -330,23 +355,10 @@ def ensure_pod(values=None, wait=True): | |||
| 330 | time.sleep(5) | 355 | time.sleep(5) |
| 331 | raise RuntimeError('pod not ready within START_TIMEOUT') | 356 | raise RuntimeError('pod not ready within START_TIMEOUT') |
| 332 | finally: | 357 | finally: |
| 358 | with state['lock']: | ||
| 333 | state['ensuring'] -= 1 | 359 | state['ensuring'] -= 1 |
| 334 | 360 | ||
| 335 | 361 | ||
| 336 | def match_workflow_models(workflow_bytes): | ||
| 337 | """Returns all catalog values whose filenames appear in the workflow.""" | ||
| 338 | try: | ||
| 339 | text = workflow_bytes.decode('utf-8', 'replace') | ||
| 340 | except Exception: | ||
| 341 | return [] | ||
| 342 | found = [] | ||
| 343 | for entry in load_catalog().get('models', []): | ||
| 344 | value = entry.get('value') | ||
| 345 | if value and value in text: | ||
| 346 | found.append(value) | ||
| 347 | return found | ||
| 348 | |||
| 349 | |||
| 350 | def idle_reaper(): | 362 | def idle_reaper(): |
| 351 | while True: | 363 | while True: |
| 352 | time.sleep(30) | 364 | time.sleep(30) |
| @@ -362,13 +374,21 @@ def idle_reaper(): | |||
| 362 | 374 | ||
| 363 | def status_body(): | 375 | def status_body(): |
| 364 | pod_id = state['pod_id'] | 376 | pod_id = state['pod_id'] |
| 377 | previous_phase = state['phase'] | ||
| 365 | if state['ensuring'] > 0: | 378 | if state['ensuring'] > 0: |
| 366 | # An ensure_pod waiter owns the phase; probing here could report green | 379 | # An ensure_pod waiter owns the phase; probing here could report green |
| 367 | # mid-download or block the status reply on a dead pod's proxy URL. | 380 | # mid-download or block the status reply on a dead pod's proxy URL. |
| 368 | pass | 381 | pass |
| 369 | elif pod_id: | 382 | elif pod_id: |
| 370 | if state['phase'] != 'green' and upstream_ready(pod_id): | 383 | if upstream_ready(pod_id): |
| 371 | state['phase'] = 'green' | 384 | state['phase'] = 'green' |
| 385 | else: | ||
| 386 | found, have, gpu = find_pod() | ||
| 387 | if found: | ||
| 388 | state.update({'pod_id': found, 'model': have, 'gpu': gpu, 'phase': 'orange'}) | ||
| 389 | else: | ||
| 390 | state.update({'pod_id': None, 'model': None, 'gpu': None, 'phase': 'red'}) | ||
| 391 | if state['phase'] == 'green' and previous_phase != 'green': | ||
| 372 | state['since'] = time.time() | 392 | state['since'] = time.time() |
| 373 | else: | 393 | else: |
| 374 | found, have, gpu = find_pod() | 394 | found, have, gpu = find_pod() |
| @@ -415,17 +435,27 @@ class Proxy(BaseHTTPRequestHandler): | |||
| 415 | 435 | ||
| 416 | # ---- control API ---- | 436 | # ---- control API ---- |
| 417 | if path == '/lazy/status': | 437 | if path == '/lazy/status': |
| 438 | if self.command != 'GET': | ||
| 439 | return self._reply(405, b'{"error": "method not allowed"}') | ||
| 418 | return self._reply(200, status_body()) | 440 | return self._reply(200, status_body()) |
| 441 | if path in ('/lazy/ping', '/lazy/warmup', '/lazy/shutdown', '/lazy/catalog') and self.command != 'POST': | ||
| 442 | return self._reply(405, b'{"error": "method not allowed"}') | ||
| 419 | if path == '/lazy/ping': | 443 | if path == '/lazy/ping': |
| 420 | if state['pod_id']: | 444 | if state['pod_id'] and state['phase'] == 'green': |
| 421 | state['last'] = time.monotonic() | 445 | state['last'] = time.monotonic() |
| 422 | return self._reply(200, b'{"ok": true}') | 446 | return self._reply(200, b'{"ok": true}') |
| 423 | if path == '/lazy/warmup': | 447 | if path == '/lazy/warmup': |
| 424 | threading.Thread(target=self._safe_ensure, args=(None,), daemon=True).start() | 448 | with state['lock']: |
| 449 | control_epoch = state['control_epoch'] | ||
| 425 | if state['phase'] != 'green': | 450 | if state['phase'] != 'green': |
| 426 | state['phase'] = 'orange' | 451 | state['phase'] = 'orange' |
| 452 | threading.Thread(target=self._safe_ensure, args=(None, control_epoch), daemon=True).start() | ||
| 427 | return self._reply(200, status_body()) | 453 | return self._reply(200, status_body()) |
| 428 | if path == '/lazy/shutdown': | 454 | if path == '/lazy/shutdown': |
| 455 | # Serialize with pod creation/recreation. Incrementing the epoch | ||
| 456 | # also cancels a warmup thread that has not entered its lock yet. | ||
| 457 | with state['lock']: | ||
| 458 | state['control_epoch'] += 1 | ||
| 429 | pod_id = state['pod_id'] or find_pod()[0] | 459 | pod_id = state['pod_id'] or find_pod()[0] |
| 430 | if pod_id: | 460 | if pod_id: |
| 431 | terminate(pod_id) | 461 | terminate(pod_id) |
| @@ -440,32 +470,24 @@ class Proxy(BaseHTTPRequestHandler): | |||
| 440 | catalog = {'models': payload.get('models', []), 'active': [a for a in (active or []) if a]} | 470 | catalog = {'models': payload.get('models', []), 'active': [a for a in (active or []) if a]} |
| 441 | save_catalog(catalog) | 471 | save_catalog(catalog) |
| 442 | log(f"catalog updated: {len(catalog['models'])} models, active={catalog['active']}") | 472 | log(f"catalog updated: {len(catalog['models'])} models, active={catalog['active']}") |
| 443 | # Fetch missing models proactively when the active set changed | ||
| 444 | # under a live pod (in place; recreate only for old images). | ||
| 445 | if state['pod_id'] and catalog['active'] and state['model'] and not (set(catalog['active']) <= key_values(state['model'])): | ||
| 446 | log('active model set changed - ensuring models on pod') | ||
| 447 | threading.Thread(target=self._safe_ensure, args=(catalog['active'],), daemon=True).start() | ||
| 448 | return self._reply(200, status_body()) | 473 | return self._reply(200, status_body()) |
| 449 | except Exception as err: | 474 | except Exception as err: |
| 450 | return self._reply(400, json.dumps({'error': str(err)}).encode()) | 475 | return self._reply(400, json.dumps({'error': str(err)}).encode()) |
| 451 | 476 | ||
| 452 | # ---- comfy proxying ---- | 477 | # ---- comfy proxying ---- |
| 453 | state['last'] = time.monotonic() | ||
| 454 | pod_id = state['pod_id'] | 478 | pod_id = state['pod_id'] |
| 455 | ready = pod_id and upstream_ready(pod_id) | 479 | if not pod_id: |
| 456 | if not ready and self.command == 'GET' and path == '/system_stats': | 480 | pod_id, have, gpu = find_pod() |
| 457 | try: | 481 | if pod_id: |
| 458 | with open(CACHE_FILE, 'rb') as f: | 482 | state.update({'pod_id': pod_id, 'model': have, 'gpu': gpu, 'phase': 'orange'}) |
| 459 | return self._reply(200, f.read()) | 483 | ready = pod_id and state['ensuring'] == 0 and upstream_ready(pod_id) |
| 460 | except FileNotFoundError: | 484 | if not ready: |
| 461 | return self._reply(200, json.dumps({'system': {'comfyui_version': 'cold'}, 'devices': []}).encode()) | 485 | return self._reply(503, json.dumps({ |
| 462 | if not ready or (self.command == 'POST' and path == '/prompt'): | 486 | 'error': 'RunPod pod is not ready; start it with SillyTavern\'s Warm up control.', |
| 463 | values = match_workflow_models(body) if (self.command == 'POST' and path == '/prompt' and body) else None | 487 | }).encode()) |
| 464 | try: | 488 | if state['phase'] != 'green': |
| 465 | pod_id = ensure_pod(values or None) | 489 | state.update({'phase': 'green', 'since': time.time()}) |
| 466 | except Exception as err: | 490 | state['last'] = time.monotonic() |
| 467 | log('ensure_pod failed:', err) | ||
| 468 | return self._reply(502, json.dumps({'error': str(err)}).encode()) | ||
| 469 | try: | 491 | try: |
| 470 | req = urllib.request.Request(pod_url(pod_id) + self.path, body, method=self.command) | 492 | req = urllib.request.Request(pod_url(pod_id) + self.path, body, method=self.command) |
| 471 | req.add_header('User-Agent', UA) | 493 | req.add_header('User-Agent', UA) |
| @@ -482,9 +504,9 @@ class Proxy(BaseHTTPRequestHandler): | |||
| 482 | finally: | 504 | finally: |
| 483 | state['last'] = time.monotonic() | 505 | state['last'] = time.monotonic() |
| 484 | 506 | ||
| 485 | def _safe_ensure(self, value): | 507 | def _safe_ensure(self, value, control_epoch): |
| 486 | try: | 508 | try: |
| 487 | ensure_pod(value) | 509 | ensure_pod(value, control_epoch=control_epoch) |
| 488 | except Exception as err: | 510 | except Exception as err: |
| 489 | log('provision failed:', err) | 511 | log('provision failed:', err) |
| 490 | 512 | ||
| @@ -494,21 +516,14 @@ class Proxy(BaseHTTPRequestHandler): | |||
| 494 | 516 | ||
| 495 | def main(): | 517 | def main(): |
| 496 | threading.Thread(target=idle_reaper, daemon=True).start() | 518 | threading.Thread(target=idle_reaper, daemon=True).start() |
| 497 | # Adopt a pre-existing pod (e.g. after a proxy restart mid-provisioning) so | 519 | # Adopt a pre-existing pod after a proxy restart without provisioning or |
| 498 | # its readiness wait - which also feeds the idle timer - keeps running. | 520 | # changing it. Read-only status polls will update its readiness phase. |
| 499 | pod_id, have, gpu = find_pod() | 521 | pod_id, have, gpu = find_pod() |
| 500 | if pod_id: | 522 | if pod_id: |
| 501 | state.update({'pod_id': pod_id, 'model': have, 'gpu': gpu, 'phase': 'orange', 'last': time.monotonic()}) | 523 | state.update({'pod_id': pod_id, 'model': have, 'gpu': gpu, 'phase': 'orange', 'last': time.monotonic()}) |
| 502 | log(f'adopting existing pod {pod_id} (model={have})') | 524 | log(f'adopting existing pod {pod_id} (model={have})') |
| 503 | |||
| 504 | def _adopt(): | ||
| 505 | try: | ||
| 506 | ensure_pod(have.split('+') if have else None) | ||
| 507 | except Exception as err: | ||
| 508 | log('adoption wait failed:', err) | ||
| 509 | threading.Thread(target=_adopt, daemon=True).start() | ||
| 510 | server = ThreadingHTTPServer(('0.0.0.0', LISTEN_PORT), Proxy) | 525 | server = ThreadingHTTPServer(('0.0.0.0', LISTEN_PORT), Proxy) |
| 511 | log(f'runpod-lazy v5 (in-place model switching) on :{LISTEN_PORT} (pod "{POD_NAME}", DCs {DATACENTERS or "any"}, idle {IDLE_SECONDS}s)') | 526 | log(f'runpod-lazy v6 (manual warmup only) on :{LISTEN_PORT} (pod "{POD_NAME}", DCs {DATACENTERS or "any"}, idle {IDLE_SECONDS}s)') |
| 512 | server.serve_forever() | 527 | server.serve_forever() |
| 513 | 528 | ||
| 514 | 529 | ||
| @@ -0,0 +1,161 @@ | |||
| 1 | import http.client | ||
| 2 | import importlib.util | ||
| 3 | import json | ||
| 4 | import os | ||
| 5 | import pathlib | ||
| 6 | import tempfile | ||
| 7 | import threading | ||
| 8 | import unittest | ||
| 9 | from unittest import mock | ||
| 10 | |||
| 11 | |||
| 12 | os.environ.setdefault('RUNPOD_KEY', 'test-key') | ||
| 13 | PROXY_PATH = pathlib.Path(__file__).parents[1] / 'runpod-lazy-proxy.py' | ||
| 14 | SPEC = importlib.util.spec_from_file_location('runpod_lazy_proxy', PROXY_PATH) | ||
| 15 | proxy = importlib.util.module_from_spec(SPEC) | ||
| 16 | SPEC.loader.exec_module(proxy) | ||
| 17 | |||
| 18 | |||
| 19 | class RunpodLazyProxyTest(unittest.TestCase): | ||
| 20 | def setUp(self): | ||
| 21 | self.temp_dir = tempfile.TemporaryDirectory() | ||
| 22 | proxy.CATALOG_FILE = str(pathlib.Path(self.temp_dir.name) / 'catalog.json') | ||
| 23 | proxy.state.update({ | ||
| 24 | 'pod_id': None, | ||
| 25 | 'model': None, | ||
| 26 | 'phase': 'red', | ||
| 27 | 'gpu': None, | ||
| 28 | 'since': 0, | ||
| 29 | 'last': 100, | ||
| 30 | 'ensuring': 0, | ||
| 31 | 'control_epoch': 0, | ||
| 32 | }) | ||
| 33 | self.find_pod = mock.patch.object(proxy, 'find_pod', return_value=(None, None, None)) | ||
| 34 | self.find_pod_mock = self.find_pod.start() | ||
| 35 | self.server = proxy.ThreadingHTTPServer(('127.0.0.1', 0), proxy.Proxy) | ||
| 36 | self.server_thread = threading.Thread(target=self.server.serve_forever, daemon=True) | ||
| 37 | self.server_thread.start() | ||
| 38 | |||
| 39 | def tearDown(self): | ||
| 40 | self.server.shutdown() | ||
| 41 | self.server.server_close() | ||
| 42 | self.server_thread.join() | ||
| 43 | self.find_pod.stop() | ||
| 44 | self.temp_dir.cleanup() | ||
| 45 | |||
| 46 | def request(self, method, path, payload=None): | ||
| 47 | body = json.dumps(payload) if payload is not None else None | ||
| 48 | headers = {'Content-Type': 'application/json'} if body is not None else {} | ||
| 49 | connection = http.client.HTTPConnection('127.0.0.1', self.server.server_port, timeout=2) | ||
| 50 | connection.request(method, path, body=body, headers=headers) | ||
| 51 | response = connection.getresponse() | ||
| 52 | data = response.read() | ||
| 53 | connection.close() | ||
| 54 | return response.status, json.loads(data or b'{}') | ||
| 55 | |||
| 56 | def test_cold_comfy_requests_fail_without_provisioning(self): | ||
| 57 | with mock.patch.object(proxy, 'ensure_pod') as ensure_pod: | ||
| 58 | for method, path, payload in [ | ||
| 59 | ('GET', '/system_stats', None), | ||
| 60 | ('GET', '/object_info', None), | ||
| 61 | ('POST', '/prompt', {'prompt': {}}), | ||
| 62 | ]: | ||
| 63 | with self.subTest(path=path): | ||
| 64 | status, body = self.request(method, path, payload) | ||
| 65 | self.assertEqual(status, 503) | ||
| 66 | self.assertIn('Warm up', body['error']) | ||
| 67 | |||
| 68 | ensure_pod.assert_not_called() | ||
| 69 | |||
| 70 | def test_status_ping_and_catalog_do_not_provision_cold_pod(self): | ||
| 71 | with mock.patch.object(proxy, 'ensure_pod') as ensure_pod: | ||
| 72 | status, body = self.request('GET', '/lazy/status') | ||
| 73 | self.assertEqual(status, 200) | ||
| 74 | self.assertEqual(body['state'], 'red') | ||
| 75 | |||
| 76 | status, _ = self.request('POST', '/lazy/ping') | ||
| 77 | self.assertEqual(status, 200) | ||
| 78 | self.assertEqual(proxy.state['last'], 100) | ||
| 79 | |||
| 80 | catalog = { | ||
| 81 | 'models': [{'name': 'Test', 'value': 'test.safetensors', 'files': []}], | ||
| 82 | 'active': ['test.safetensors'], | ||
| 83 | } | ||
| 84 | status, _ = self.request('POST', '/lazy/catalog', catalog) | ||
| 85 | self.assertEqual(status, 200) | ||
| 86 | self.assertEqual(proxy.load_catalog(), catalog) | ||
| 87 | ensure_pod.assert_not_called() | ||
| 88 | |||
| 89 | def test_status_clears_a_stale_ready_pod_without_restarting_it(self): | ||
| 90 | proxy.state.update({'pod_id': 'stopped-pod', 'phase': 'green', 'last': 100}) | ||
| 91 | |||
| 92 | with mock.patch.object(proxy, 'upstream_ready', return_value=False): | ||
| 93 | status, body = self.request('GET', '/lazy/status') | ||
| 94 | |||
| 95 | self.assertEqual(status, 200) | ||
| 96 | self.assertEqual(body['state'], 'red') | ||
| 97 | self.assertIsNone(body['pod_id']) | ||
| 98 | |||
| 99 | status, _ = self.request('POST', '/lazy/ping') | ||
| 100 | self.assertEqual(status, 200) | ||
| 101 | self.assertEqual(proxy.state['last'], 100) | ||
| 102 | |||
| 103 | def test_get_warmup_is_rejected_without_provisioning(self): | ||
| 104 | with mock.patch.object(proxy, 'ensure_pod') as ensure_pod: | ||
| 105 | status, body = self.request('GET', '/lazy/warmup') | ||
| 106 | |||
| 107 | self.assertEqual(status, 405) | ||
| 108 | self.assertEqual(body['error'], 'method not allowed') | ||
| 109 | ensure_pod.assert_not_called() | ||
| 110 | |||
| 111 | def test_warmup_is_the_explicit_provisioning_path(self): | ||
| 112 | called = threading.Event() | ||
| 113 | with mock.patch.object(proxy, 'ensure_pod') as ensure_pod: | ||
| 114 | ensure_pod.side_effect = lambda _value, control_epoch: called.set() | ||
| 115 | |||
| 116 | status, body = self.request('POST', '/lazy/warmup') | ||
| 117 | |||
| 118 | self.assertEqual(status, 200) | ||
| 119 | self.assertEqual(body['state'], 'orange') | ||
| 120 | self.assertTrue(called.wait(1)) | ||
| 121 | ensure_pod.assert_called_once_with(None, control_epoch=0) | ||
| 122 | |||
| 123 | def test_shutdown_cancels_pod_creation_in_progress(self): | ||
| 124 | create_started = threading.Event() | ||
| 125 | finish_create = threading.Event() | ||
| 126 | shutdown_finished = threading.Event() | ||
| 127 | shutdown_result = {} | ||
| 128 | |||
| 129 | def create_pod(_values): | ||
| 130 | create_started.set() | ||
| 131 | self.assertTrue(finish_create.wait(2)) | ||
| 132 | return 'created-after-warmup', 'test-gpu' | ||
| 133 | |||
| 134 | def shutdown_request(): | ||
| 135 | shutdown_result['response'] = self.request('POST', '/lazy/shutdown') | ||
| 136 | shutdown_finished.set() | ||
| 137 | |||
| 138 | with mock.patch.object(proxy, '_create_pod_with_retries', side_effect=create_pod), \ | ||
| 139 | mock.patch.object(proxy, 'terminate') as terminate: | ||
| 140 | status, body = self.request('POST', '/lazy/warmup') | ||
| 141 | self.assertEqual(status, 200) | ||
| 142 | self.assertEqual(body['state'], 'orange') | ||
| 143 | self.assertTrue(create_started.wait(1)) | ||
| 144 | |||
| 145 | shutdown_thread = threading.Thread(target=shutdown_request) | ||
| 146 | shutdown_thread.start() | ||
| 147 | self.assertFalse(shutdown_finished.wait(0.1)) | ||
| 148 | |||
| 149 | finish_create.set() | ||
| 150 | shutdown_thread.join(2) | ||
| 151 | self.assertFalse(shutdown_thread.is_alive()) | ||
| 152 | |||
| 153 | status, body = shutdown_result['response'] | ||
| 154 | self.assertEqual(status, 200) | ||
| 155 | self.assertEqual(body['state'], 'red') | ||
| 156 | self.assertIsNone(proxy.state['pod_id']) | ||
| 157 | terminate.assert_called_once_with('created-after-warmup') | ||
| 158 | |||
| 159 | |||
| 160 | if __name__ == '__main__': | ||
| 161 | unittest.main() | ||
| @@ -511,9 +511,13 @@ async function isCurrentSourceReachable() { | |||
| 511 | 511 | ||
| 512 | switch (extension_settings.sd.source) { | 512 | switch (extension_settings.sd.source) { |
| 513 | case sources.comfy: | 513 | case sources.comfy: |
| 514 | return extension_settings.sd.comfy_type === comfyTypes.standard | 514 | if (extension_settings.sd.comfy_type !== comfyTypes.standard) { |
| 515 | ? probe('/api/sd/comfy/ping', { url: extension_settings.sd.comfy_url }) | 515 | return true; |
| 516 | : true; | 516 | } |
| 517 | if (isRunpodProxyUrl(extension_settings.sd.comfy_url)) { | ||
| 518 | return isRunpodReady(SOURCE_PROBE_TIMEOUT_MS); | ||
| 519 | } | ||
| 520 | return probe('/api/sd/comfy/ping', { url: extension_settings.sd.comfy_url }); | ||
| 517 | case sources.auto: | 521 | case sources.auto: |
| 518 | case sources.vlad: | 522 | case sources.vlad: |
| 519 | case sources.drawthings: | 523 | case sources.drawthings: |
| @@ -831,7 +835,6 @@ async function loadSettings() { | |||
| 831 | renderRefImages(); | 835 | renderRefImages(); |
| 832 | renderRunpodModels(); | 836 | renderRunpodModels(); |
| 833 | setupRunpodLoops(); | 837 | setupRunpodLoops(); |
| 834 | pushRunpodCatalog(); | ||
| 835 | 838 | ||
| 836 | for (const style of extension_settings.sd.styles) { | 839 | for (const style of extension_settings.sd.styles) { |
| 837 | const option = document.createElement('option'); | 840 | const option = document.createElement('option'); |
| @@ -1538,7 +1541,7 @@ async function fetchReferenceImageBase64(refImage) { | |||
| 1538 | /** Poll cadence for the pod status indicator (faster while it is starting). */ | 1541 | /** Poll cadence for the pod status indicator (faster while it is starting). */ |
| 1539 | const RUNPOD_POLL_IDLE_MS = 30000; | 1542 | const RUNPOD_POLL_IDLE_MS = 30000; |
| 1540 | const RUNPOD_POLL_BUSY_MS = 5000; | 1543 | const RUNPOD_POLL_BUSY_MS = 5000; |
| 1541 | /** Keepalive cadence: signals "a SillyTavern tab is open" to the proxy. */ | 1544 | /** Keepalive cadence for a pod that the status API has confirmed is ready. */ |
| 1542 | const RUNPOD_PING_MS = 60000; | 1545 | const RUNPOD_PING_MS = 60000; |
| 1543 | 1546 | ||
| 1544 | let runpodStatusTimer = null; | 1547 | let runpodStatusTimer = null; |
| @@ -1550,6 +1553,36 @@ function getRunpodLazyUrl() { | |||
| 1550 | } | 1553 | } |
| 1551 | 1554 | ||
| 1552 | /** | 1555 | /** |
| 1556 | * Reads the lazy proxy status without starting or warming a pod. | ||
| 1557 | * @param {number} timeout Request timeout in milliseconds. | ||
| 1558 | * @returns {Promise<object>} Parsed /lazy/status response. | ||
| 1559 | */ | ||
| 1560 | async function getRunpodStatus(timeout = 8000) { | ||
| 1561 | const url = getRunpodLazyUrl(); | ||
| 1562 | if (!url) { | ||
| 1563 | throw new Error('RunPod lazy proxy URL is not configured.'); | ||
| 1564 | } | ||
| 1565 | const result = await fetch(`${url}/lazy/status`, { signal: AbortSignal.timeout(timeout) }); | ||
| 1566 | if (!result.ok) { | ||
| 1567 | throw new Error(`RunPod status returned ${result.status}.`); | ||
| 1568 | } | ||
| 1569 | return result.json(); | ||
| 1570 | } | ||
| 1571 | |||
| 1572 | /** | ||
| 1573 | * Checks whether the lazy proxy already has a ready pod. | ||
| 1574 | * @param {number} timeout Request timeout in milliseconds. | ||
| 1575 | * @returns {Promise<boolean>} True only when /lazy/status reports green. | ||
| 1576 | */ | ||
| 1577 | async function isRunpodReady(timeout = 8000) { | ||
| 1578 | try { | ||
| 1579 | return (await getRunpodStatus(timeout))?.state === 'green'; | ||
| 1580 | } catch { | ||
| 1581 | return false; | ||
| 1582 | } | ||
| 1583 | } | ||
| 1584 | |||
| 1585 | /** | ||
| 1553 | * Updates the status dot + text from a /lazy/status response (or an error). | 1586 | * Updates the status dot + text from a /lazy/status response (or an error). |
| 1554 | * @param {object|null} status Parsed status JSON, or null when unreachable. | 1587 | * @param {object|null} status Parsed status JSON, or null when unreachable. |
| 1555 | */ | 1588 | */ |
| @@ -1619,12 +1652,8 @@ async function pollRunpodStatus() { | |||
| 1619 | } | 1652 | } |
| 1620 | let phase = runpodLastPhase; | 1653 | let phase = runpodLastPhase; |
| 1621 | try { | 1654 | try { |
| 1622 | const result = await fetch(`${url}/lazy/status`, { signal: AbortSignal.timeout(8000) }); | ||
| 1623 | if (!result.ok) { | ||
| 1624 | throw new Error(`status ${result.status}`); | ||
| 1625 | } | ||
| 1626 | runpodPollFailures = 0; | 1655 | runpodPollFailures = 0; |
| 1627 | phase = renderRunpodStatus(await result.json()); | 1656 | phase = renderRunpodStatus(await getRunpodStatus()); |
| 1628 | } catch { | 1657 | } catch { |
| 1629 | // A single slow/failed poll (e.g. while the proxy is provisioning a pod) | 1658 | // A single slow/failed poll (e.g. while the proxy is provisioning a pod) |
| 1630 | // must not flip the dot to red; only sustained unreachability does. | 1659 | // must not flip the dot to red; only sustained unreachability does. |
| @@ -1643,8 +1672,14 @@ async function runpodControl(action) { | |||
| 1643 | return; | 1672 | return; |
| 1644 | } | 1673 | } |
| 1645 | try { | 1674 | try { |
| 1675 | if (action === 'warmup' && !(await pushRunpodCatalog())) { | ||
| 1676 | throw new Error('Could not sync the RunPod model catalog. The pod was not started.'); | ||
| 1677 | } | ||
| 1646 | const result = await fetch(`${url}/lazy/${action}`, { method: 'POST', signal: AbortSignal.timeout(10000) }); | 1678 | const result = await fetch(`${url}/lazy/${action}`, { method: 'POST', signal: AbortSignal.timeout(10000) }); |
| 1647 | renderRunpodStatus(result.ok ? await result.json() : null); | 1679 | if (!result.ok) { |
| 1680 | throw new Error(`RunPod ${action} returned ${result.status}.`); | ||
| 1681 | } | ||
| 1682 | renderRunpodStatus(await result.json()); | ||
| 1648 | if (action === 'warmup') { | 1683 | if (action === 'warmup') { |
| 1649 | toastr.info(t`Pod warmup requested. Models will pre-download; the dot turns green when ready.`, t`Image Generation`); | 1684 | toastr.info(t`Pod warmup requested. Models will pre-download; the dot turns green when ready.`, t`Image Generation`); |
| 1650 | } | 1685 | } |
| @@ -1657,10 +1692,11 @@ async function runpodControl(action) { | |||
| 1657 | 1692 | ||
| 1658 | function runpodKeepalive() { | 1693 | function runpodKeepalive() { |
| 1659 | const url = getRunpodLazyUrl(); | 1694 | const url = getRunpodLazyUrl(); |
| 1660 | if (!url) { | 1695 | if (!url || runpodLastPhase !== 'green') { |
| 1661 | return; | 1696 | return; |
| 1662 | } | 1697 | } |
| 1663 | // Only extends the pod's idle timer; the proxy never starts a pod for a ping. | 1698 | // The proxy also rejects wake-on-ping; this phase guard avoids unnecessary |
| 1699 | // requests whenever the last read-only status poll says the pod is down. | ||
| 1664 | fetch(`${url}/lazy/ping`, { method: 'POST', signal: AbortSignal.timeout(5000) }).catch(() => { }); | 1700 | fetch(`${url}/lazy/ping`, { method: 'POST', signal: AbortSignal.timeout(5000) }).catch(() => { }); |
| 1665 | } | 1701 | } |
| 1666 | 1702 | ||
| @@ -1747,11 +1783,15 @@ function getRunpodActiveModels() { | |||
| 1747 | return [config.model, config.lora].filter(v => v && known.has(v)); | 1783 | return [config.model, config.lora].filter(v => v && known.has(v)); |
| 1748 | } | 1784 | } |
| 1749 | 1785 | ||
| 1750 | /** Pushes the model catalog + active selection to the proxy (fire-and-forget). */ | 1786 | /** |
| 1787 | * Pushes the model catalog + active selection to the proxy. | ||
| 1788 | * Catalog updates only store configuration; only /lazy/warmup may provision. | ||
| 1789 | * @returns {Promise<boolean>} Whether the proxy accepted the catalog. | ||
| 1790 | */ | ||
| 1751 | async function pushRunpodCatalog() { | 1791 | async function pushRunpodCatalog() { |
| 1752 | const url = getRunpodLazyUrl(); | 1792 | const url = getRunpodLazyUrl(); |
| 1753 | if (!url) { | 1793 | if (!url) { |
| 1754 | return; | 1794 | return false; |
| 1755 | } | 1795 | } |
| 1756 | const models = getRunpodCatalog().map(m => ({ | 1796 | const models = getRunpodCatalog().map(m => ({ |
| 1757 | name: m.name || m.value, | 1797 | name: m.name || m.value, |
| @@ -1760,14 +1800,19 @@ async function pushRunpodCatalog() { | |||
| 1760 | files: parseRunpodFiles(m.downloads), | 1800 | files: parseRunpodFiles(m.downloads), |
| 1761 | })); | 1801 | })); |
| 1762 | try { | 1802 | try { |
| 1763 | await fetch(`${url}/lazy/catalog`, { | 1803 | const result = await fetch(`${url}/lazy/catalog`, { |
| 1764 | method: 'POST', | 1804 | method: 'POST', |
| 1765 | headers: { 'Content-Type': 'application/json' }, | 1805 | headers: { 'Content-Type': 'application/json' }, |
| 1766 | signal: AbortSignal.timeout(8000), | 1806 | signal: AbortSignal.timeout(8000), |
| 1767 | body: JSON.stringify({ models, active: getRunpodActiveModels() }), | 1807 | body: JSON.stringify({ models, active: getRunpodActiveModels() }), |
| 1768 | }); | 1808 | }); |
| 1809 | if (!result.ok) { | ||
| 1810 | throw new Error(`catalog returned ${result.status}`); | ||
| 1811 | } | ||
| 1812 | return true; | ||
| 1769 | } catch (error) { | 1813 | } catch (error) { |
| 1770 | console.warn('SD: runpod catalog push failed', error); | 1814 | console.warn('SD: runpod catalog push failed', error); |
| 1815 | return false; | ||
| 1771 | } | 1816 | } |
| 1772 | } | 1817 | } |
| 1773 | 1818 | ||
| @@ -2724,8 +2769,8 @@ async function onModelChange() { | |||
| 2724 | rememberWorkflowPref('model', extension_settings.sd.model); | 2769 | rememberWorkflowPref('model', extension_settings.sd.model); |
| 2725 | saveSettingsDebounced(); | 2770 | saveSettingsDebounced(); |
| 2726 | 2771 | ||
| 2727 | // Selecting a catalog model while pointed at the RunPod proxy kicks off the | 2772 | // Keep the proxy's stored selection current, but catalog updates never |
| 2728 | // download / re-provision for it right away. | 2773 | // download, provision, or start a pod. Warm up applies the selection. |
| 2729 | if (isRunpodProxyUrl(extension_settings.sd.comfy_url)) { | 2774 | if (isRunpodProxyUrl(extension_settings.sd.comfy_url)) { |
| 2730 | pushRunpodCatalog(); | 2775 | pushRunpodCatalog(); |
| 2731 | } | 2776 | } |
| @@ -3142,6 +3187,10 @@ async function loadComfySamplers() { | |||
| 3142 | if (extension_settings.sd.comfy_type === comfyTypes.runpod_serverless) { | 3187 | if (extension_settings.sd.comfy_type === comfyTypes.runpod_serverless) { |
| 3143 | return ['N/A']; | 3188 | return ['N/A']; |
| 3144 | } | 3189 | } |
| 3190 | if (isRunpodProxyUrl(extension_settings.sd.comfy_url)) { | ||
| 3191 | // Do not query /object_info through a cold lazy proxy on page load. | ||
| 3192 | return extension_settings.sd.sampler ? [extension_settings.sd.sampler] : []; | ||
| 3193 | } | ||
| 3145 | if (!extension_settings.sd.comfy_url) { | 3194 | if (!extension_settings.sd.comfy_url) { |
| 3146 | return []; | 3195 | return []; |
| 3147 | } | 3196 | } |
| @@ -3909,6 +3958,10 @@ async function loadComfySchedulers() { | |||
| 3909 | if (extension_settings.sd.comfy_type === comfyTypes.runpod_serverless) { | 3958 | if (extension_settings.sd.comfy_type === comfyTypes.runpod_serverless) { |
| 3910 | return ['N/A']; | 3959 | return ['N/A']; |
| 3911 | } | 3960 | } |
| 3961 | if (isRunpodProxyUrl(extension_settings.sd.comfy_url)) { | ||
| 3962 | // Do not query /object_info through a cold lazy proxy on page load. | ||
| 3963 | return extension_settings.sd.scheduler ? [extension_settings.sd.scheduler] : []; | ||
| 3964 | } | ||
| 3912 | if (!extension_settings.sd.comfy_url) { | 3965 | if (!extension_settings.sd.comfy_url) { |
| 3913 | return []; | 3966 | return []; |
| 3914 | } | 3967 | } |
| @@ -5742,6 +5795,9 @@ async function generateComfyImageCommon(prompt, negativePrompt, signal, basePath | |||
| 5742 | * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete. | 5795 | * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete. |
| 5743 | */ | 5796 | */ |
| 5744 | async function generateComfyImage(prompt, negativePrompt, signal) { | 5797 | async function generateComfyImage(prompt, negativePrompt, signal) { |
| 5798 | if (isRunpodProxyUrl(extension_settings.sd.comfy_url) && !(await isRunpodReady(SOURCE_PROBE_TIMEOUT_MS))) { | ||
| 5799 | throw new Error('RunPod on-demand pod is not ready. Start it with the Warm up control.'); | ||
| 5800 | } | ||
| 5745 | const placeholders = [ | 5801 | const placeholders = [ |
| 5746 | 'model', | 5802 | 'model', |
| 5747 | 'vae', | 5803 | 'vae', |
| @@ -6433,6 +6489,8 @@ async function sendMessage(prompt, image, generationType, additionalNegativePref | |||
| 6433 | media_display: MEDIA_DISPLAY.GALLERY, | 6489 | media_display: MEDIA_DISPLAY.GALLERY, |
| 6434 | media_index: 0, | 6490 | media_index: 0, |
| 6435 | inline_image: false, | 6491 | inline_image: false, |
| 6492 | sd_prompt: prompt, | ||
| 6493 | sd_prompt_message: messageText, | ||
| 6436 | }, | 6494 | }, |
| 6437 | }; | 6495 | }; |
| 6438 | context.chat.push(message); | 6496 | context.chat.push(message); |
| @@ -6480,6 +6538,7 @@ async function addSDGenButtons() { | |||
| 6480 | }); | 6538 | }); |
| 6481 | 6539 | ||
| 6482 | $(document).on('click', '.sd_message_gen', (e) => sdMessageButton($(e.currentTarget), { animate: false })); | 6540 | $(document).on('click', '.sd_message_gen', (e) => sdMessageButton($(e.currentTarget), { animate: false })); |
| 6541 | $(document).on('click', '.mes_edit', rememberImagePromptBeforeEdit); | ||
| 6483 | 6542 | ||
| 6484 | $(document).on('click touchend', function (e) { | 6543 | $(document).on('click touchend', function (e) { |
| 6485 | const target = $(e.target); | 6544 | const target = $(e.target); |
| @@ -6614,6 +6673,108 @@ function isValidState() { | |||
| 6614 | /** @type {WeakMap<HTMLElement, AbortController>} */ | 6673 | /** @type {WeakMap<HTMLElement, AbortController>} */ |
| 6615 | const buttonAbortControllers = new WeakMap(); | 6674 | const buttonAbortControllers = new WeakMap(); |
| 6616 | 6675 | ||
| 6676 | /** @type {WeakMap<ChatMessage, { message: string, prompt: string }>} */ | ||
| 6677 | const imagePromptEditSnapshots = new WeakMap(); | ||
| 6678 | |||
| 6679 | /** | ||
| 6680 | * Gets the most recent generated media attachment whose prompt backs image swipes. | ||
| 6681 | * @param {ChatMessage} message Message that owns the media gallery. | ||
| 6682 | * @returns {MediaAttachment|null} Generated media attachment, if one exists. | ||
| 6683 | */ | ||
| 6684 | function getImagePromptMedia(message) { | ||
| 6685 | const media = message?.extra?.media; | ||
| 6686 | if (!Array.isArray(media)) { | ||
| 6687 | return null; | ||
| 6688 | } | ||
| 6689 | |||
| 6690 | return media.findLast(item => item?.source === MEDIA_SOURCE.GENERATED || item?.generation_type !== undefined) ?? null; | ||
| 6691 | } | ||
| 6692 | |||
| 6693 | /** | ||
| 6694 | * Extracts an edited prompt while preserving the surrounding generated-message template. | ||
| 6695 | * Falls back to the entire edited message when the old prompt cannot be located, which | ||
| 6696 | * also supports users who intentionally replace the template along with the prompt. | ||
| 6697 | * @param {string} previousMessage Message text before editing. | ||
| 6698 | * @param {string} editedMessage Message text after editing. | ||
| 6699 | * @param {string} previousPrompt Prompt represented by the previous message. | ||
| 6700 | * @returns {string} Prompt to use for subsequent image swipes. | ||
| 6701 | */ | ||
| 6702 | function extractEditedImagePrompt(previousMessage, editedMessage, previousPrompt) { | ||
| 6703 | if (previousMessage === editedMessage) { | ||
| 6704 | return previousPrompt; | ||
| 6705 | } | ||
| 6706 | |||
| 6707 | if (previousPrompt) { | ||
| 6708 | let promptIndex = previousMessage.lastIndexOf(previousPrompt); | ||
| 6709 | while (promptIndex >= 0) { | ||
| 6710 | const prefix = previousMessage.slice(0, promptIndex); | ||
| 6711 | const suffix = previousMessage.slice(promptIndex + previousPrompt.length); | ||
| 6712 | if (editedMessage.startsWith(prefix) && editedMessage.endsWith(suffix) && editedMessage.length >= prefix.length + suffix.length) { | ||
| 6713 | return editedMessage.slice(prefix.length, editedMessage.length - suffix.length).trim(); | ||
| 6714 | } | ||
| 6715 | promptIndex = previousMessage.lastIndexOf(previousPrompt, promptIndex - 1); | ||
| 6716 | } | ||
| 6717 | } | ||
| 6718 | |||
| 6719 | return editedMessage.trim(); | ||
| 6720 | } | ||
| 6721 | |||
| 6722 | /** | ||
| 6723 | * Remembers the pre-edit message text. New image messages persist this information, | ||
| 6724 | * while this snapshot provides the same behavior for messages created before the fix. | ||
| 6725 | * @param {JQuery.ClickEvent} event Edit-button click event. | ||
| 6726 | */ | ||
| 6727 | function rememberImagePromptBeforeEdit(event) { | ||
| 6728 | const context = getContext(); | ||
| 6729 | const messageId = Number($(event.currentTarget).closest('.mes').attr('mesid')); | ||
| 6730 | const message = context.chat[messageId]; | ||
| 6731 | const mediaAttachment = getImagePromptMedia(message); | ||
| 6732 | if (!message || !mediaAttachment) { | ||
| 6733 | return; | ||
| 6734 | } | ||
| 6735 | |||
| 6736 | const messageText = String(message.mes ?? ''); | ||
| 6737 | const representedMedia = message.extra.media.findLast(item => item?.title && messageText.includes(item.title)); | ||
| 6738 | const prompt = message.extra?.sd_prompt ?? representedMedia?.title ?? mediaAttachment.title ?? message.extra?.title ?? ''; | ||
| 6739 | imagePromptEditSnapshots.set(message, { | ||
| 6740 | message: messageText, | ||
| 6741 | prompt: String(prompt), | ||
| 6742 | }); | ||
| 6743 | } | ||
| 6744 | |||
| 6745 | /** | ||
| 6746 | * Makes an edited generated-image message the prompt source for later image swipes. | ||
| 6747 | * @param {number} messageId Edited message index. | ||
| 6748 | */ | ||
| 6749 | function onImagePromptMessageEdited(messageId) { | ||
| 6750 | const context = getContext(); | ||
| 6751 | const message = context.chat[messageId]; | ||
| 6752 | const mediaAttachment = getImagePromptMedia(message); | ||
| 6753 | if (!message || !mediaAttachment) { | ||
| 6754 | return; | ||
| 6755 | } | ||
| 6756 | |||
| 6757 | const snapshot = imagePromptEditSnapshots.get(message) ?? ( | ||
| 6758 | typeof message.extra?.sd_prompt_message === 'string' | ||
| 6759 | ? { message: message.extra.sd_prompt_message, prompt: String(message.extra.sd_prompt ?? mediaAttachment.title ?? '') } | ||
| 6760 | : null | ||
| 6761 | ); | ||
| 6762 | imagePromptEditSnapshots.delete(message); | ||
| 6763 | if (!snapshot) { | ||
| 6764 | return; | ||
| 6765 | } | ||
| 6766 | |||
| 6767 | const editedMessage = String(message.mes ?? ''); | ||
| 6768 | if (editedMessage === snapshot.message) { | ||
| 6769 | return; | ||
| 6770 | } | ||
| 6771 | |||
| 6772 | const editedPrompt = extractEditedImagePrompt(snapshot.message, editedMessage, snapshot.prompt); | ||
| 6773 | message.extra.sd_prompt = editedPrompt; | ||
| 6774 | message.extra.sd_prompt_message = editedMessage; | ||
| 6775 | message.extra.sd_prompt_override = editedPrompt; | ||
| 6776 | } | ||
| 6777 | |||
| 6617 | /** | 6778 | /** |
| 6618 | * "Paintbrush" button handler to generate a new image for a message. | 6779 | * "Paintbrush" button handler to generate a new image for a message. |
| 6619 | * @param {JQuery<HTMLElement>} $icon The click target. | 6780 | * @param {JQuery<HTMLElement>} $icon The click target. |
| @@ -6761,7 +6922,7 @@ async function generateMediaSwipe(mediaAttachment, message, onStart, onComplete, | |||
| 6761 | 6922 | ||
| 6762 | try { | 6923 | try { |
| 6763 | const callback = (_a, _b, _c, _d, _e, _f, format) => { result.type = isVideo(format) ? MEDIA_TYPE.VIDEO : MEDIA_TYPE.IMAGE; }; | 6924 | const callback = (_a, _b, _c, _d, _e, _f, format) => { result.type = isVideo(format) ? MEDIA_TYPE.VIDEO : MEDIA_TYPE.IMAGE; }; |
| 6764 | const savedPrompt = mediaAttachment.title ?? message.extra.title ?? ''; | 6925 | const savedPrompt = message.extra.sd_prompt_override ?? mediaAttachment.title ?? message.extra.title ?? ''; |
| 6765 | const savedNegative = mediaAttachment.negative ?? message.extra.negative ?? ''; | 6926 | const savedNegative = mediaAttachment.negative ?? message.extra.negative ?? ''; |
| 6766 | const refineArgs = { | 6927 | const refineArgs = { |
| 6767 | negative: savedNegative, | 6928 | negative: savedNegative, |
| @@ -7438,6 +7599,7 @@ export async function init() { | |||
| 7438 | }); | 7599 | }); |
| 7439 | 7600 | ||
| 7440 | eventSource.on(event_types.CHAT_CHANGED, onChatChanged); | 7601 | eventSource.on(event_types.CHAT_CHANGED, onChatChanged); |
| 7602 | eventSource.on(event_types.MESSAGE_EDITED, onImagePromptMessageEdited); | ||
| 7441 | eventSource.on(event_types.IMAGE_SWIPED, onImageSwiped); | 7603 | eventSource.on(event_types.IMAGE_SWIPED, onImageSwiped); |
| 7442 | 7604 | ||
| 7443 | [event_types.SECRET_WRITTEN, event_types.SECRET_DELETED, event_types.SECRET_ROTATED].forEach(event => { | 7605 | [event_types.SECRET_WRITTEN, event_types.SECRET_DELETED, event_types.SECRET_ROTATED].forEach(event => { |
| @@ -100,9 +100,9 @@ | |||
| 100 | <span data-i18n="Copy ComfyUI URL">Copy ComfyUI URL</span> | 100 | <span data-i18n="Copy ComfyUI URL">Copy ComfyUI URL</span> |
| 101 | </div> | 101 | </div> |
| 102 | </div> | 102 | </div> |
| 103 | <small data-i18n="sd_runpod_small">While a SillyTavern tab is open, a keepalive ping keeps the pod alive; it shuts down 15 minutes after the last activity either way. Red = off, orange = starting/downloading models, green = ready.</small> | 103 | <small data-i18n="sd_runpod_small">The pod starts only when you press Warm up. A SillyTavern tab keeps an already-ready pod alive; it shuts down 15 minutes after the last activity. Image requests fail over instead of starting a stopped pod. Red = off, orange = starting/downloading models, green = ready.</small> |
| 104 | <h5 data-i18n="Model catalog">Model catalog</h5> | 104 | <h5 data-i18n="Model catalog">Model catalog</h5> |
| 105 | <small data-i18n="sd_runpod_models_small">When the ComfyUI URL points at the proxy, the Model dropdown lists these entries instead of asking the backend. The pod downloads exactly the selected entry's files at boot; changing the selection re-provisions with the new model. Filename = what "%model%" resolves to; downloads = one "subpath url" per line, including companion files (text encoder, VAE, LoRA).</small> | 105 | <small data-i18n="sd_runpod_models_small">When the ComfyUI URL points at the proxy, the Model dropdown lists these entries instead of asking the backend. Press Warm up to apply the selected model and download its files; changing the selection or generating an image never starts or provisions a stopped pod. Filename = what "%model%" resolves to; downloads = one "subpath url" per line, including companion files (text encoder, VAE, LoRA).</small> |
| 106 | <div id="sd_runpod_models_list" class="flex-container flexFlowColumn marginTopBot5"></div> | 106 | <div id="sd_runpod_models_list" class="flex-container flexFlowColumn marginTopBot5"></div> |
| 107 | <div class="flex-container marginTopBot5"> | 107 | <div class="flex-container marginTopBot5"> |
| 108 | <div id="sd_runpod_models_add" class="menu_button menu_button_icon"> | 108 | <div id="sd_runpod_models_add" class="menu_button menu_button_icon"> |