Coordinate shared GPU access across studio instances

server/utils/imageComfy.ts

@@ -1,3 +1,4 @@
+import { assertSharedGpu, sharedGpuHeaders, sharedGpuConfigured } from '~/server/utils/sharedGpu'
 import { createHash } from 'node:crypto'
 import { AsyncLocalStorage } from 'node:async_hooks'
 import { getComfyHost, viaAgentFrontDoor } from '~/server/utils/comfy'
@@ -61,6 +62,7 @@ export function getSidecarImageHost() {
 }
 
 export function getImageComfyHost() {
+  if (sharedGpuConfigured()) return getComfyHost()
   const fromAls = imageHostAls.getStore()
   if (fromAls) return viaAgentFrontDoor(fromAls) || fromAls
   if (imageComfyHostOverride) return viaAgentFrontDoor(imageComfyHostOverride) || imageComfyHostOverride
@@ -80,6 +82,12 @@ export async function imageComfyFetch(path: string, init?: RequestInit) {
   if (!host) {
     throw createError({ statusCode: 500, statusMessage: 'No image ComfyUI host is configured' })
   }
+  if (init?.method && !['GET', 'HEAD'].includes(init.method.toUpperCase())) {
+    await assertSharedGpu()
+    const headers = new Headers(init.headers)
+    for (const [key, value] of Object.entries(sharedGpuHeaders())) headers.set(key, value)
+    init = { ...init, headers }
+  }
   try {
     return await fetch(`${host}${path}`, init)
   } catch (error) {
@@ -390,7 +398,8 @@ async function purgeImageOnDesktop(opts: {
       headers: {
         'Content-Type': 'application/json',
         Accept: 'application/json',
-        ...(token ? { Authorization: `Bearer ${token}` } : {})
+        ...(token ? { Authorization: `Bearer ${token}` } : {}),
+        ...sharedGpuHeaders()
       },
       body: JSON.stringify({
         imageName: imageNames[0] || '',