Coordinate shared GPU access across studio instances

server/utils/comfyLifecycle.ts

@@ -1,3 +1,4 @@
+import { assertSharedGpu, sharedGpuHeaders, withSharedGpuStart } from '~/server/utils/sharedGpu'
 import { execFile, spawn } from 'node:child_process'
 import { promisify } from 'node:util'
 import { getComfyHost, setComfyHostOverride, comfyConfigured } from '~/server/utils/comfy'
@@ -85,13 +86,15 @@ export async function fetchQueue() {
 async function controlRequest(path: string, method = 'GET', timeoutMs = 5000, body?: Record<string, unknown>) {
   const { controlUrl, controlToken } = settings()
   if (!controlUrl) return null
+  if (method !== 'GET') await assertSharedGpu()
   try {
     const res = await fetch(`${controlUrl}${path}`, {
       method,
       headers: {
         Accept: 'application/json',
         ...(body ? { 'Content-Type': 'application/json' } : {}),
-        ...(controlToken ? { Authorization: `Bearer ${controlToken}` } : {})
+        ...(controlToken ? { Authorization: `Bearer ${controlToken}` } : {}),
+        ...(method !== 'GET' ? sharedGpuHeaders() : {})
       },
       ...(body ? { body: JSON.stringify(body) } : {}),
       signal: AbortSignal.timeout(timeoutMs)
@@ -344,7 +347,7 @@ async function waitWhileBusy(onStatus: StatusFn) {
   const started = Date.now()
   while (Date.now() - started < busyWaitMs) {
     const queue = await fetchQueue()
-    if (queue.running === 0) return
+    if (queue.ok && queue.running === 0 && queue.pending === 0) return
     onStatus({
       state: 'busy',
       message: `ComfyUI is busy with another prompt (${queue.running} running, ${queue.pending} queued). Waiting up to 3 minutes, then this job will be saved.`,
@@ -488,7 +491,10 @@ async function ensureUnlocked(onStatus: StatusFn, skipBusyWait = false) {
 
 export function ensureComfyReady(onStatus: StatusFn = () => undefined, options?: { skipBusyWait?: boolean }) {
   const skipBusyWait = options?.skipBusyWait === true
-  const run = gate.then(() => ensureUnlocked(onStatus, skipBusyWait))
+  const run = gate.then(() => withSharedGpuStart(
+    () => ensureUnlocked(onStatus, skipBusyWait),
+    async () => { await assertSharedGpu() }
+  ))
   gate = run.then(() => undefined, () => undefined)
   return run
 }