Fix image prompt swipes and RunPod cold starts

abc1cb103bccdb271ac27a2e4ae81b744cd67c57

permissionBRICK <40219477+permissionBRICK@users.noreply.github.com>

Parents: f4309221b126179b42ec9664e0948610685fba54

.gitignore+2 -0 Viewed

 yarn.lock
 *.code-workspace
 test-results/
 +__pycache__/
 +*.py[cod]

homelab/README.md+1 -1 Viewed

 | File | Deployed at | Purpose |
 |---|---|---|
 | `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). |
 | `runpod-lazy-proxy.py` | `/mnt/user/appdata/runpod-lazy/proxy.py` (container `runpod-lazy`, :8189) | ComfyUI-compatible lazyon-demand proxy for a volumeless RunPod GPU pod: wake-on-generateonly with`/lazy/warmup` model-groupstarts detectionor provisions, while cold ComfyUI requests return 503 for provider failover. `/lazy/status|warmup|shutdown|ping|catalog` control API for theprovide SillyTavern warmupcontrols UI,and a NAS-side 15-min idle terminate. Model-set changes onrequested aby runningWarm podup 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. |
 | `runpod-pod.sh` | `/boot/config/runpod-pod.sh` | CLI for the runpod-lazy control API (`status|warmup|shutdown|ping|watch`). |
 Related: the RunPod worker image (ComfyUI + custom nodes + boot-time model

homelab/runpod-lazy-proxy.py+81 -66 Viewed

   GET  /lazy/status    -> {"state": red|orange|green, "model": ..., ...}
   POST /lazy/warmup    -> start pod for the active catalog entry
   POST /lazy/shutdown  -> terminate pod now
   POST /lazy/ping      -> keepalive (extendsfor idlea timer;ready pod (never starts a pod)
   POST /lazy/catalog   -> {"models": [...], "active": "<value>"} store catalog;
 -                          re-provision the pod if the active model changed
 +                          never starts or provisions a pod
 WakeStart semantics: aonly ComfyUIPOST /prompt arrivinglazy/warmup withmay noprovision readyor podchange matchesmodels. theComfyUI
 -workflow against catalog filenames and provisions for that entry (falls back
 +requests return 503 unless a pod is already ready, allowing SillyTavern's image
 -to the active entry). The idle timer lives here (NAS side): IDLE_SECONDS after
 +target fallback chain to continue without waking this target. The idle timer
 -the last activity signal the pod is terminated, frontend or no frontend.
 +lives here (NAS side): IDLE_SECONDS after the last activity signal the pod is
 +terminated, frontend or no frontend.
 """
 import json
 import os

 START_TIMEOUT = int(os.environ.get('START_TIMEOUT', '1500'))
 HF_TOKEN = os.environ.get('HF_TOKEN', '')
 COMFY_ARGS = os.environ.get('COMFY_ARGS', '--listen 0.0.0.0 --port 8188 --use-pytorch-cross-attention')
 -CACHE_FILE = os.environ.get('CACHE_FILE', '/app/system_stats.cache')
 CATALOG_FILE = os.environ.get('CATALOG_FILE', '/app/catalog.json')
 state = {

     'since': time.time(),
     'last': time.monotonic(),
     'ensuring': 0,      # active ensure_pod waiters (status skips probing then)
 +    'control_epoch': 0, # incremented by shutdown to cancel in-flight warmups
     'lock': threading.Lock(),
+}

         req = urllib.request.Request(pod_url(pod_id) + '/system_stats', headers={'User-Agent': UA})
         with urllib.request.urlopen(req, timeout=5) as resp:
             if resp.status == 200:
                 data = resp.read()
 -                with open(CACHE_FILE, 'wb') as f:
 -                    f.write(data)
                 return True
     except Exception:
         pass

     raise RuntimeError(f'could not create pod: {last_err}')
 def ensure_pod(values=None, wait=True, control_epoch=None):
     """Ensures a pod holding the catalog values (None = active selection).
     An existing pod gets missing models downloaded IN PLACE through the in-pod
     model manager; recreation only happens for old-image pods without one."""
 +    def check_cancelled():
 +        if control_epoch is not None and control_epoch != state['control_epoch']:
 +            raise RuntimeError('pod warmup cancelled')
++
     values = values or active_values()
     key = values_key(values)
     files = needed_files(values)
     dests = [f['dest'].lstrip('/') for f in files]
     created = False
     with state['lock']:
 +        check_cancelled()
         pod_id, have, gpu = (state['pod_id'], state['model'], state['gpu'])
         if not pod_id:
             pod_id, have, gpu = find_pod()
         if not pod_id:
             pod_id, gpu = _create_pod_with_retries(values)
             have, created = key, True
 -        state['pod_id'], state['model'], state['gpu'] = pod_id, have, gpu
 +        check_cancelled()
 -        state['last'] = time.monotonic()
 +        state.update({
 +            'pod_id': pod_id,
 +            'model': have,
 +            'gpu': gpu,
 +            'last': time.monotonic(),
 +            'phase': 'green' if state['phase'] == 'green' else 'orange',
 +        })
 -    if state['phase'] != 'green':
 -        state['phase'] = 'orange'
     if not wait:
         return pod_id
     # Freshly created pods boot with the right manifest; existing pods need an
     # /ensure pushed once the manager answers.
     ensured = created or not files
     no_manager_strikes = 0
 +    with state['lock']:
 +        check_cancelled()
         state['ensuring'] += 1
     try:
         deadline = time.monotonic() + START_TIMEOUT
         while time.monotonic() < deadline:
             # Booting/downloading counts as activity, else the reaper would
             # kill a warming pod mid-download.
 +            with state['lock']:
 +                check_cancelled()
 +                if state['pod_id'] != pod_id:
 +                    raise RuntimeError('pod replaced while waiting')
                 state['last'] = time.monotonic()
             if not ensured:
                 try:
                     mm_request(pod_id, 'POST', '/ensure', {'files': files, 'priority': True})
 +                except Exception:
 +                    pass  # manager still booting, or old image without one
 +                else:
                     ensured = True
 +                    with state['lock']:
 +                        check_cancelled()
                         state['model'] = values_key(key_values(state['model']) | set(values))
                     log(f'in-place ensure requested: {key}')
 -                except Exception:
 -                    pass  # manager still booting, or old image without one
             ready = upstream_ready(pod_id)
             models_ok = True
             if not created and files and ensured:

                 # Comfy is up but there is no manager (old image): last resort.
                 log(f'pod has models={state["model"]}, need {key} - recreating (no model manager)')
                 with state['lock']:
 +                    check_cancelled()
                     terminate(pod_id)
                     pod_id, gpu = _create_pod_with_retries(values)
 +                    check_cancelled()
                     state.update({'pod_id': pod_id, 'model': key, 'gpu': gpu, 'last': time.monotonic()})
                 created, ensured = True, True
                 continue
             if ready and models_ok:
 +                with state['lock']:
 +                    check_cancelled()
 +                    if state['pod_id'] != pod_id:
 +                        raise RuntimeError('pod replaced while waiting')
                     if state['phase'] != 'green':
                         state['phase'] = 'green'
                         state['since'] = time.time()

                         state['prefetch_pod'] = pod_id
                         threading.Thread(target=prefetch_rest, args=(pod_id,), daemon=True).start()
                 return pod_id
 +            with state['lock']:
 +                check_cancelled()
                 if state['phase'] == 'green':
                     state['phase'] = 'orange'  # in-place download in progress
                 if state['pod_id'] != pod_id:

             time.sleep(5)
         raise RuntimeError('pod not ready within START_TIMEOUT')
     finally:
 +        with state['lock']:
             state['ensuring'] -= 1
 -def match_workflow_models(workflow_bytes):
 -    """Returns all catalog values whose filenames appear in the workflow."""
 -    try:
 -        text = workflow_bytes.decode('utf-8', 'replace')
 -    except Exception:
 -        return []
 -    found = []
 -    for entry in load_catalog().get('models', []):
 -        value = entry.get('value')
 -        if value and value in text:
 -            found.append(value)
 -    return found
+-
+-
 def idle_reaper():
     while True:
         time.sleep(30)

 def status_body():
     pod_id = state['pod_id']
 +    previous_phase = state['phase']
     if state['ensuring'] > 0:
         # An ensure_pod waiter owns the phase; probing here could report green
         # mid-download or block the status reply on a dead pod's proxy URL.
         pass
     elif pod_id:
         if state['phase'] != 'green' and upstream_ready(pod_id):
             state['phase'] = 'green'
 +        else:
 +            found, have, gpu = find_pod()
 +            if found:
 +                state.update({'pod_id': found, 'model': have, 'gpu': gpu, 'phase': 'orange'})
 +            else:
 +                state.update({'pod_id': None, 'model': None, 'gpu': None, 'phase': 'red'})
 +        if state['phase'] == 'green' and previous_phase != 'green':
             state['since'] = time.time()
     else:
         found, have, gpu = find_pod()

         # ---- control API ----
         if path == '/lazy/status':
 +            if self.command != 'GET':
 +                return self._reply(405, b'{"error": "method not allowed"}')
             return self._reply(200, status_body())
 +        if path in ('/lazy/ping', '/lazy/warmup', '/lazy/shutdown', '/lazy/catalog') and self.command != 'POST':
 +            return self._reply(405, b'{"error": "method not allowed"}')
         if path == '/lazy/ping':
             if state['pod_id'] and state['phase'] == 'green':
                 state['last'] = time.monotonic()
             return self._reply(200, b'{"ok": true}')
         if path == '/lazy/warmup':
 -            threading.Thread(target=self._safe_ensure, args=(None,), daemon=True).start()
 +            with state['lock']:
 +                control_epoch = state['control_epoch']
                 if state['phase'] != 'green':
                     state['phase'] = 'orange'
 +            threading.Thread(target=self._safe_ensure, args=(None, control_epoch), daemon=True).start()
             return self._reply(200, status_body())
         if path == '/lazy/shutdown':
 +            # Serialize with pod creation/recreation. Incrementing the epoch
 +            # also cancels a warmup thread that has not entered its lock yet.
 +            with state['lock']:
 +                state['control_epoch'] += 1
                 pod_id = state['pod_id'] or find_pod()[0]
                 if pod_id:
                     terminate(pod_id)

                 catalog = {'models': payload.get('models', []), 'active': [a for a in (active or []) if a]}
                 save_catalog(catalog)
                 log(f"catalog updated: {len(catalog['models'])} models, active={catalog['active']}")
 -                # Fetch missing models proactively when the active set changed
 -                # under a live pod (in place; recreate only for old images).
 -                if state['pod_id'] and catalog['active'] and state['model'] and not (set(catalog['active']) <= key_values(state['model'])):
 -                    log('active model set changed - ensuring models on pod')
 -                    threading.Thread(target=self._safe_ensure, args=(catalog['active'],), daemon=True).start()
                 return self._reply(200, status_body())
             except Exception as err:
                 return self._reply(400, json.dumps({'error': str(err)}).encode())
         # ---- comfy proxying ----
 -        state['last'] = time.monotonic()
         pod_id = state['pod_id']
 -        ready = pod_id and upstream_ready(pod_id)
 +        if not pod_id:
 -        if not ready and self.command == 'GET' and path == '/system_stats':
 +            pod_id, have, gpu = find_pod()
             tryif pod_id:
 -                with open(CACHE_FILE, 'rb') as f:
 +                state.update({'pod_id': pod_id, 'model': have, 'gpu': gpu, 'phase': 'orange'})
 -                    return self._reply(200, f.read())
 +        ready = pod_id and state['ensuring'] == 0 and upstream_ready(pod_id)
 -            except FileNotFoundError:
 +        if not ready:
 -                return self._reply(200, json.dumps({'system': {'comfyui_version': 'cold'}, 'devices': []}).encode())
 +            return self._reply(503, json.dumps({
 -        if not ready or (self.command == 'POST' and path == '/prompt'):
 +                'error': 'RunPod pod is not ready; start it with SillyTavern\'s Warm up control.',
 -            values = match_workflow_models(body) if (self.command == 'POST' and path == '/prompt' and body) else None
 +            }).encode())
 -            try:
 +        if state['phase'] != 'green':
 -                pod_id = ensure_pod(values or None)
 +            state.update({'phase': 'green', 'since': time.time()})
 -            except Exception as err:
 +        state['last'] = time.monotonic()
 -                log('ensure_pod failed:', err)
 -                return self._reply(502, json.dumps({'error': str(err)}).encode())
         try:
             req = urllib.request.Request(pod_url(pod_id) + self.path, body, method=self.command)
             req.add_header('User-Agent', UA)

         finally:
             state['last'] = time.monotonic()
     def _safe_ensure(self, value, control_epoch):
         try:
             ensure_pod(value, control_epoch=control_epoch)
         except Exception as err:
             log('provision failed:', err)

 def main():
     threading.Thread(target=idle_reaper, daemon=True).start()
     # Adopt a pre-existing pod (e.g. after a proxy restart mid-without provisioning) soor
     # its readinesschanging waitit. Read- which alsoonly feedsstatus thepolls idlewill timerupdate -its keepsreadiness runningphase.
     pod_id, have, gpu = find_pod()
     if pod_id:
         state.update({'pod_id': pod_id, 'model': have, 'gpu': gpu, 'phase': 'orange', 'last': time.monotonic()})
         log(f'adopting existing pod {pod_id} (model={have})')
+-
 -        def _adopt():
 -            try:
 -                ensure_pod(have.split('+') if have else None)
 -            except Exception as err:
 -                log('adoption wait failed:', err)
 -        threading.Thread(target=_adopt, daemon=True).start()
     server = ThreadingHTTPServer(('0.0.0.0', LISTEN_PORT), Proxy)
     log(f'runpod-lazy v5v6 (in-placemanual modelwarmup switchingonly) on :{LISTEN_PORT} (pod "{POD_NAME}", DCs {DATACENTERS or "any"}, idle {IDLE_SECONDS}s)')
     server.serve_forever()

homelab/tests/test_runpod_lazy_proxy.py+161 -0 Viewed

 +import http.client
 +import importlib.util
 +import json
 +import os
 +import pathlib
 +import tempfile
 +import threading
 +import unittest
 +from unittest import mock
++
++
 +os.environ.setdefault('RUNPOD_KEY', 'test-key')
 +PROXY_PATH = pathlib.Path(__file__).parents[1] / 'runpod-lazy-proxy.py'
 +SPEC = importlib.util.spec_from_file_location('runpod_lazy_proxy', PROXY_PATH)
 +proxy = importlib.util.module_from_spec(SPEC)
 +SPEC.loader.exec_module(proxy)
++
++
 +class RunpodLazyProxyTest(unittest.TestCase):
 +    def setUp(self):
 +        self.temp_dir = tempfile.TemporaryDirectory()
 +        proxy.CATALOG_FILE = str(pathlib.Path(self.temp_dir.name) / 'catalog.json')
 +        proxy.state.update({
 +            'pod_id': None,
 +            'model': None,
 +            'phase': 'red',
 +            'gpu': None,
 +            'since': 0,
 +            'last': 100,
 +            'ensuring': 0,
 +            'control_epoch': 0,
 +        })
 +        self.find_pod = mock.patch.object(proxy, 'find_pod', return_value=(None, None, None))
 +        self.find_pod_mock = self.find_pod.start()
 +        self.server = proxy.ThreadingHTTPServer(('127.0.0.1', 0), proxy.Proxy)
 +        self.server_thread = threading.Thread(target=self.server.serve_forever, daemon=True)
 +        self.server_thread.start()
++
 +    def tearDown(self):
 +        self.server.shutdown()
 +        self.server.server_close()
 +        self.server_thread.join()
 +        self.find_pod.stop()
 +        self.temp_dir.cleanup()
++
 +    def request(self, method, path, payload=None):
 +        body = json.dumps(payload) if payload is not None else None
 +        headers = {'Content-Type': 'application/json'} if body is not None else {}
 +        connection = http.client.HTTPConnection('127.0.0.1', self.server.server_port, timeout=2)
 +        connection.request(method, path, body=body, headers=headers)
 +        response = connection.getresponse()
 +        data = response.read()
 +        connection.close()
 +        return response.status, json.loads(data or b'{}')
++
 +    def test_cold_comfy_requests_fail_without_provisioning(self):
 +        with mock.patch.object(proxy, 'ensure_pod') as ensure_pod:
 +            for method, path, payload in [
 +                ('GET', '/system_stats', None),
 +                ('GET', '/object_info', None),
 +                ('POST', '/prompt', {'prompt': {}}),
 +            ]:
 +                with self.subTest(path=path):
 +                    status, body = self.request(method, path, payload)
 +                    self.assertEqual(status, 503)
 +                    self.assertIn('Warm up', body['error'])
++
 +            ensure_pod.assert_not_called()
++
 +    def test_status_ping_and_catalog_do_not_provision_cold_pod(self):
 +        with mock.patch.object(proxy, 'ensure_pod') as ensure_pod:
 +            status, body = self.request('GET', '/lazy/status')
 +            self.assertEqual(status, 200)
 +            self.assertEqual(body['state'], 'red')
++
 +            status, _ = self.request('POST', '/lazy/ping')
 +            self.assertEqual(status, 200)
 +            self.assertEqual(proxy.state['last'], 100)
++
 +            catalog = {
 +                'models': [{'name': 'Test', 'value': 'test.safetensors', 'files': []}],
 +                'active': ['test.safetensors'],
 +            }
 +            status, _ = self.request('POST', '/lazy/catalog', catalog)
 +            self.assertEqual(status, 200)
 +            self.assertEqual(proxy.load_catalog(), catalog)
 +            ensure_pod.assert_not_called()
++
 +    def test_status_clears_a_stale_ready_pod_without_restarting_it(self):
 +        proxy.state.update({'pod_id': 'stopped-pod', 'phase': 'green', 'last': 100})
++
 +        with mock.patch.object(proxy, 'upstream_ready', return_value=False):
 +            status, body = self.request('GET', '/lazy/status')
++
 +        self.assertEqual(status, 200)
 +        self.assertEqual(body['state'], 'red')
 +        self.assertIsNone(body['pod_id'])
++
 +        status, _ = self.request('POST', '/lazy/ping')
 +        self.assertEqual(status, 200)
 +        self.assertEqual(proxy.state['last'], 100)
++
 +    def test_get_warmup_is_rejected_without_provisioning(self):
 +        with mock.patch.object(proxy, 'ensure_pod') as ensure_pod:
 +            status, body = self.request('GET', '/lazy/warmup')
++
 +            self.assertEqual(status, 405)
 +            self.assertEqual(body['error'], 'method not allowed')
 +            ensure_pod.assert_not_called()
++
 +    def test_warmup_is_the_explicit_provisioning_path(self):
 +        called = threading.Event()
 +        with mock.patch.object(proxy, 'ensure_pod') as ensure_pod:
 +            ensure_pod.side_effect = lambda _value, control_epoch: called.set()
++
 +            status, body = self.request('POST', '/lazy/warmup')
++
 +            self.assertEqual(status, 200)
 +            self.assertEqual(body['state'], 'orange')
 +            self.assertTrue(called.wait(1))
 +            ensure_pod.assert_called_once_with(None, control_epoch=0)
++
 +    def test_shutdown_cancels_pod_creation_in_progress(self):
 +        create_started = threading.Event()
 +        finish_create = threading.Event()
 +        shutdown_finished = threading.Event()
 +        shutdown_result = {}
++
 +        def create_pod(_values):
 +            create_started.set()
 +            self.assertTrue(finish_create.wait(2))
 +            return 'created-after-warmup', 'test-gpu'
++
 +        def shutdown_request():
 +            shutdown_result['response'] = self.request('POST', '/lazy/shutdown')
 +            shutdown_finished.set()
++
 +        with mock.patch.object(proxy, '_create_pod_with_retries', side_effect=create_pod), \
 +                mock.patch.object(proxy, 'terminate') as terminate:
 +            status, body = self.request('POST', '/lazy/warmup')
 +            self.assertEqual(status, 200)
 +            self.assertEqual(body['state'], 'orange')
 +            self.assertTrue(create_started.wait(1))
++
 +            shutdown_thread = threading.Thread(target=shutdown_request)
 +            shutdown_thread.start()
 +            self.assertFalse(shutdown_finished.wait(0.1))
++
 +            finish_create.set()
 +            shutdown_thread.join(2)
 +            self.assertFalse(shutdown_thread.is_alive())
++
 +            status, body = shutdown_result['response']
 +            self.assertEqual(status, 200)
 +            self.assertEqual(body['state'], 'red')
 +            self.assertIsNone(proxy.state['pod_id'])
 +            terminate.assert_called_once_with('created-after-warmup')
++
++
 +if __name__ == '__main__':
 +    unittest.main()

public/scripts/extensions/stable-diffusion/index.js+181 -19 Viewed

     switch (extension_settings.sd.source) {
         case sources.comfy:
             returnif (extension_settings.sd.comfy_type =!== comfyTypes.standard) {
 -                ? probe('/api/sd/comfy/ping', { url: extension_settings.sd.comfy_url })
 +                return true;
 -                : true;
 +            }
 +            if (isRunpodProxyUrl(extension_settings.sd.comfy_url)) {
 +                return isRunpodReady(SOURCE_PROBE_TIMEOUT_MS);
 +            }
 +            return probe('/api/sd/comfy/ping', { url: extension_settings.sd.comfy_url });
         case sources.auto:
         case sources.vlad:
         case sources.drawthings:

     renderRefImages();
     renderRunpodModels();
     setupRunpodLoops();
 -    pushRunpodCatalog();
     for (const style of extension_settings.sd.styles) {
         const option = document.createElement('option');

 /** Poll cadence for the pod status indicator (faster while it is starting). */
 const RUNPOD_POLL_IDLE_MS = 30000;
 const RUNPOD_POLL_BUSY_MS = 5000;
 /** Keepalive cadence: signalsfor "a SillyTavernpod tabthat isthe open"status toAPI thehas proxyconfirmed is ready. */
 const RUNPOD_PING_MS = 60000;
 let runpodStatusTimer = null;

+}
 /**
 + * Reads the lazy proxy status without starting or warming a pod.
 + * @param {number} timeout Request timeout in milliseconds.
 + * @returns {Promise<object>} Parsed /lazy/status response.
 + */
 +async function getRunpodStatus(timeout = 8000) {
 +    const url = getRunpodLazyUrl();
 +    if (!url) {
 +        throw new Error('RunPod lazy proxy URL is not configured.');
 +    }
 +    const result = await fetch(`${url}/lazy/status`, { signal: AbortSignal.timeout(timeout) });
 +    if (!result.ok) {
 +        throw new Error(`RunPod status returned ${result.status}.`);
 +    }
 +    return result.json();
 +}
++
 +/**
 + * Checks whether the lazy proxy already has a ready pod.
 + * @param {number} timeout Request timeout in milliseconds.
 + * @returns {Promise<boolean>} True only when /lazy/status reports green.
 + */
 +async function isRunpodReady(timeout = 8000) {
 +    try {
 +        return (await getRunpodStatus(timeout))?.state === 'green';
 +    } catch {
 +        return false;
 +    }
 +}
++
 +/**
  * Updates the status dot + text from a /lazy/status response (or an error).
  * @param {object|null} status Parsed status JSON, or null when unreachable.
  */

+    }
     let phase = runpodLastPhase;
     try {
 -        const result = await fetch(`${url}/lazy/status`, { signal: AbortSignal.timeout(8000) });
 -        if (!result.ok) {
 -            throw new Error(`status ${result.status}`);
 -        }
         runpodPollFailures = 0;
         phase = renderRunpodStatus(await result.jsongetRunpodStatus());
     } catch {
         // A single slow/failed poll (e.g. while the proxy is provisioning a pod)
         // must not flip the dot to red; only sustained unreachability does.

         return;
+    }
     try {
 +        if (action === 'warmup' && !(await pushRunpodCatalog())) {
 +            throw new Error('Could not sync the RunPod model catalog. The pod was not started.');
 +        }
         const result = await fetch(`${url}/lazy/${action}`, { method: 'POST', signal: AbortSignal.timeout(10000) });
 -        renderRunpodStatus(result.ok ? await result.json() : null);
 +        if (!result.ok) {
 +            throw new Error(`RunPod ${action} returned ${result.status}.`);
 +        }
 +        renderRunpodStatus(await result.json());
         if (action === 'warmup') {
             toastr.info(t`Pod warmup requested. Models will pre-download; the dot turns green when ready.`, t`Image Generation`);
+        }

 function runpodKeepalive() {
     const url = getRunpodLazyUrl();
     if (!url || runpodLastPhase !== 'green') {
         return;
+    }
 -    // Only extends the pod's idle timer; the proxy never starts a pod for a ping.
 +    // The proxy also rejects wake-on-ping; this phase guard avoids unnecessary
 +    // requests whenever the last read-only status poll says the pod is down.
     fetch(`${url}/lazy/ping`, { method: 'POST', signal: AbortSignal.timeout(5000) }).catch(() => { });
+}

     return [config.model, config.lora].filter(v => v && known.has(v));
+}
 -/** Pushes the model catalog + active selection to the proxy (fire-and-forget). */
 +/**
 + * Pushes the model catalog + active selection to the proxy.
 + * Catalog updates only store configuration; only /lazy/warmup may provision.
 + * @returns {Promise<boolean>} Whether the proxy accepted the catalog.
 + */
 async function pushRunpodCatalog() {
     const url = getRunpodLazyUrl();
     if (!url) {
         return false;
+    }
     const models = getRunpodCatalog().map(m => ({
         name: m.name || m.value,

         files: parseRunpodFiles(m.downloads),
     }));
     try {
         const result = await fetch(`${url}/lazy/catalog`, {
             method: 'POST',
             headers: { 'Content-Type': 'application/json' },
             signal: AbortSignal.timeout(8000),
             body: JSON.stringify({ models, active: getRunpodActiveModels() }),
         });
 +        if (!result.ok) {
 +            throw new Error(`catalog returned ${result.status}`);
 +        }
 +        return true;
     } catch (error) {
         console.warn('SD: runpod catalog push failed', error);
 +        return false;
+    }
+}

     rememberWorkflowPref('model', extension_settings.sd.model);
     saveSettingsDebounced();
 -    // Selecting a catalog model while pointed at the RunPod proxy kicks off the
 +    // Keep the proxy's stored selection current, but catalog updates never
 -    // download / re-provision for it right away.
 +    // download, provision, or start a pod. Warm up applies the selection.
     if (isRunpodProxyUrl(extension_settings.sd.comfy_url)) {
         pushRunpodCatalog();
+    }

     if (extension_settings.sd.comfy_type === comfyTypes.runpod_serverless) {
         return ['N/A'];
+    }
 +    if (isRunpodProxyUrl(extension_settings.sd.comfy_url)) {
 +        // Do not query /object_info through a cold lazy proxy on page load.
 +        return extension_settings.sd.sampler ? [extension_settings.sd.sampler] : [];
 +    }
     if (!extension_settings.sd.comfy_url) {
         return [];
+    }

     if (extension_settings.sd.comfy_type === comfyTypes.runpod_serverless) {
         return ['N/A'];
+    }
 +    if (isRunpodProxyUrl(extension_settings.sd.comfy_url)) {
 +        // Do not query /object_info through a cold lazy proxy on page load.
 +        return extension_settings.sd.scheduler ? [extension_settings.sd.scheduler] : [];
 +    }
     if (!extension_settings.sd.comfy_url) {
         return [];
+    }

  * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
  */
 async function generateComfyImage(prompt, negativePrompt, signal) {
 +    if (isRunpodProxyUrl(extension_settings.sd.comfy_url) && !(await isRunpodReady(SOURCE_PROBE_TIMEOUT_MS))) {
 +        throw new Error('RunPod on-demand pod is not ready. Start it with the Warm up control.');
 +    }
     const placeholders = [
         'model',
         'vae',

             media_display: MEDIA_DISPLAY.GALLERY,
             media_index: 0,
             inline_image: false,
 +            sd_prompt: prompt,
 +            sd_prompt_message: messageText,
         },
     };
     context.chat.push(message);

     });
     $(document).on('click', '.sd_message_gen', (e) => sdMessageButton($(e.currentTarget), { animate: false }));
 +    $(document).on('click', '.mes_edit', rememberImagePromptBeforeEdit);
     $(document).on('click touchend', function (e) {
         const target = $(e.target);

 /** @type {WeakMap<HTMLElement, AbortController>} */
 const buttonAbortControllers = new WeakMap();
 +/** @type {WeakMap<ChatMessage, { message: string, prompt: string }>} */
 +const imagePromptEditSnapshots = new WeakMap();
++
 +/**
 + * Gets the most recent generated media attachment whose prompt backs image swipes.
 + * @param {ChatMessage} message Message that owns the media gallery.
 + * @returns {MediaAttachment|null} Generated media attachment, if one exists.
 + */
 +function getImagePromptMedia(message) {
 +    const media = message?.extra?.media;
 +    if (!Array.isArray(media)) {
 +        return null;
 +    }
++
 +    return media.findLast(item => item?.source === MEDIA_SOURCE.GENERATED || item?.generation_type !== undefined) ?? null;
 +}
++
 +/**
 + * Extracts an edited prompt while preserving the surrounding generated-message template.
 + * Falls back to the entire edited message when the old prompt cannot be located, which
 + * also supports users who intentionally replace the template along with the prompt.
 + * @param {string} previousMessage Message text before editing.
 + * @param {string} editedMessage Message text after editing.
 + * @param {string} previousPrompt Prompt represented by the previous message.
 + * @returns {string} Prompt to use for subsequent image swipes.
 + */
 +function extractEditedImagePrompt(previousMessage, editedMessage, previousPrompt) {
 +    if (previousMessage === editedMessage) {
 +        return previousPrompt;
 +    }
++
 +    if (previousPrompt) {
 +        let promptIndex = previousMessage.lastIndexOf(previousPrompt);
 +        while (promptIndex >= 0) {
 +            const prefix = previousMessage.slice(0, promptIndex);
 +            const suffix = previousMessage.slice(promptIndex + previousPrompt.length);
 +            if (editedMessage.startsWith(prefix) && editedMessage.endsWith(suffix) && editedMessage.length >= prefix.length + suffix.length) {
 +                return editedMessage.slice(prefix.length, editedMessage.length - suffix.length).trim();
 +            }
 +            promptIndex = previousMessage.lastIndexOf(previousPrompt, promptIndex - 1);
 +        }
 +    }
++
 +    return editedMessage.trim();
 +}
++
 +/**
 + * Remembers the pre-edit message text. New image messages persist this information,
 + * while this snapshot provides the same behavior for messages created before the fix.
 + * @param {JQuery.ClickEvent} event Edit-button click event.
 + */
 +function rememberImagePromptBeforeEdit(event) {
 +    const context = getContext();
 +    const messageId = Number($(event.currentTarget).closest('.mes').attr('mesid'));
 +    const message = context.chat[messageId];
 +    const mediaAttachment = getImagePromptMedia(message);
 +    if (!message || !mediaAttachment) {
 +        return;
 +    }
++
 +    const messageText = String(message.mes ?? '');
 +    const representedMedia = message.extra.media.findLast(item => item?.title && messageText.includes(item.title));
 +    const prompt = message.extra?.sd_prompt ?? representedMedia?.title ?? mediaAttachment.title ?? message.extra?.title ?? '';
 +    imagePromptEditSnapshots.set(message, {
 +        message: messageText,
 +        prompt: String(prompt),
 +    });
 +}
++
 +/**
 + * Makes an edited generated-image message the prompt source for later image swipes.
 + * @param {number} messageId Edited message index.
 + */
 +function onImagePromptMessageEdited(messageId) {
 +    const context = getContext();
 +    const message = context.chat[messageId];
 +    const mediaAttachment = getImagePromptMedia(message);
 +    if (!message || !mediaAttachment) {
 +        return;
 +    }
++
 +    const snapshot = imagePromptEditSnapshots.get(message) ?? (
 +        typeof message.extra?.sd_prompt_message === 'string'
 +            ? { message: message.extra.sd_prompt_message, prompt: String(message.extra.sd_prompt ?? mediaAttachment.title ?? '') }
 +            : null
 +    );
 +    imagePromptEditSnapshots.delete(message);
 +    if (!snapshot) {
 +        return;
 +    }
++
 +    const editedMessage = String(message.mes ?? '');
 +    if (editedMessage === snapshot.message) {
 +        return;
 +    }
++
 +    const editedPrompt = extractEditedImagePrompt(snapshot.message, editedMessage, snapshot.prompt);
 +    message.extra.sd_prompt = editedPrompt;
 +    message.extra.sd_prompt_message = editedMessage;
 +    message.extra.sd_prompt_override = editedPrompt;
 +}
++
 /**
  * "Paintbrush" button handler to generate a new image for a message.
  * @param {JQuery<HTMLElement>} $icon The click target.

     try {
         const callback = (_a, _b, _c, _d, _e, _f, format) => { result.type = isVideo(format) ? MEDIA_TYPE.VIDEO : MEDIA_TYPE.IMAGE; };
         const savedPrompt = message.extra.sd_prompt_override ?? mediaAttachment.title ?? message.extra.title ?? '';
         const savedNegative = mediaAttachment.negative ?? message.extra.negative ?? '';
         const refineArgs = {
             negative: savedNegative,

     });
     eventSource.on(event_types.CHAT_CHANGED, onChatChanged);
 +    eventSource.on(event_types.MESSAGE_EDITED, onImagePromptMessageEdited);
     eventSource.on(event_types.IMAGE_SWIPED, onImageSwiped);
     [event_types.SECRET_WRITTEN, event_types.SECRET_DELETED, event_types.SECRET_ROTATED].forEach(event => {

public/scripts/extensions/stable-diffusion/settings.html+2 -2 Viewed

                     <span data-i18n="Copy ComfyUI URL">Copy ComfyUI URL</span>
                 </div>
             </div>
             <small data-i18n="sd_runpod_small">WhileThe apod SillyTavernstarts tabonly iswhen open,you apress keepaliveWarm pingup. A SillyTavern tab keeps thean already-ready pod alive; it shuts down 15 minutes after the last activity. eitherImage wayrequests fail over instead of starting a stopped pod. Red = off, orange = starting/downloading models, green = ready.</small>
             <h5 data-i18n="Model catalog">Model catalog</h5>
             <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. ThePress podWarm downloadsup exactlyto apply the selected entry'smodel filesand atdownload bootits files; changing the selection re-provisionsor withgenerating thean newimage modelnever starts or provisions a stopped pod. Filename = what &quot;%model%&quot; resolves to; downloads = one &quot;subpath url&quot; per line, including companion files (text encoder, VAE, LoRA).</small>
             <div id="sd_runpod_models_list" class="flex-container flexFlowColumn marginTopBot5"></div>
             <div class="flex-container marginTopBot5">
                 <div id="sd_runpod_models_add" class="menu_button menu_button_icon">