aigen/server/utils/imageComfy.ts

let imageComfyHostOverride = ''

export function setImageComfyHostOverride(url: string) {
  imageComfyHostOverride = String(url || '').replace(/\/$/, '')
}

export function getImageComfyHost() {
  if (imageComfyHostOverride) return imageComfyHostOverride
  const config = useRuntimeConfig()
  let host = String(config.imageComfyHost || process.env.IMAGE_COMFY_HOST || '').trim().replace(/\/$/, '')
  const port = String(config.imageComfyPort || process.env.IMAGE_COMFY_PORT || '').trim()
  if (!host) {
    throw createError({ statusCode: 500, statusMessage: 'IMAGE_COMFY_HOST is not configured' })
  }
  if (!/^https?:\/\//i.test(host)) host = `http://${host}`
  try {
    const url = new URL(host)
    if (port && !url.port) url.port = port
    return url.origin
  } catch {
    return port ? `${host}:${port}` : host
  }
}

export function imageComfyConfigured() {
  const config = useRuntimeConfig()
  return Boolean(String(config.imageComfyHost || process.env.IMAGE_COMFY_HOST || '').trim())
}

export async function imageComfyFetch(path: string, init?: RequestInit) {
  const host = getImageComfyHost()
  try {
    return await fetch(`${host}${path}`, init)
  } catch (error) {
    throw createError({
      statusCode: 502,
      statusMessage: `Image ComfyUI host unreachable (${host})`,
      data: { cause: error instanceof Error ? error.message : String(error) }
    })
  }
}

export function imageComfyWsUrl(clientId: string) {
  return `${getImageComfyHost().replace(/^http/, 'ws')}/ws?clientId=${encodeURIComponent(clientId)}`
}

function imageInputFilename(original: string, jobId?: string) {
  const raw = String(original || 'still.png')
  const dot = raw.lastIndexOf('.')
  const ext = (dot >= 0 ? raw.slice(dot) : '.png').replace(/[^.a-zA-Z0-9]/g, '') || '.png'
  const base = (dot >= 0 ? raw.slice(0, dot) : raw).replace(/[^a-zA-Z0-9._-]+/g, '_').slice(0, 48) || 'still'
  const id = String(jobId || crypto.randomUUID()).replace(/-/g, '').slice(0, 10)
  return `aigen_edit_${id}_${base}${ext}`
}

export async function uploadImageEdit(file: { filename: string; data: Buffer; type?: string }, jobId?: string) {
  const body = new FormData()
  const blob = new Blob([new Uint8Array(file.data)], { type: file.type || 'application/octet-stream' })
  const filename = imageInputFilename(file.filename, jobId)
  body.append('image', blob, filename)
  body.append('overwrite', 'true')
  body.append('type', 'input')
  const res = await imageComfyFetch('/upload/image', { method: 'POST', body })
  if (!res.ok) {
    throw createError({ statusCode: 502, statusMessage: `Image upload failed (${res.status})` })
  }
  const uploaded = (await res.json()) as { name: string; subfolder?: string; type?: string }
  return {
    name: uploaded.name || filename,
    subfolder: uploaded.subfolder || '',
    type: uploaded.type || 'input'
  }
}

export async function queueImagePrompt(graph: unknown, clientId: string) {
  const res = await imageComfyFetch('/prompt', {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ prompt: graph, client_id: clientId })
  })
  const payload = await res.json().catch(() => ({}))
  if (!res.ok) {
    const message = (payload as { error?: { message?: string } }).error?.message
      || (payload as { node_errors?: unknown }).node_errors
      || `Queue failed (${res.status})`
    throw createError({ statusCode: 502, statusMessage: String(message), data: payload })
  }
  return payload as { prompt_id: string; number?: number }
}

export async function interruptImageComfy() {
  if (!imageComfyConfigured()) return false
  const res = await imageComfyFetch('/interrupt', { method: 'POST' })
  return res.ok
}

export async function fetchImageHistory(promptId: string) {
  const res = await imageComfyFetch(`/history/${encodeURIComponent(promptId)}`)
  if (!res.ok) return null
  return (await res.json()) as Record<string, unknown>
}

function isStillFile(item: { filename?: string } | null | undefined) {
  if (!item?.filename) return false
  return /\.(png|jpe?g|webp)$/i.test(item.filename)
}

function findStill(value: unknown, depth = 0): { filename: string; subfolder: string; type: string } | null {
  if (!value || typeof value !== 'object' || depth > 8) return null
  if (Array.isArray(value)) {
    const files = value.filter((item): item is { filename?: string; subfolder?: string; type?: string } => Boolean(item && typeof item === 'object'))
    const match = files.find(isStillFile)
    if (match?.filename) {
      return {
        filename: String(match.filename),
        subfolder: String(match.subfolder || ''),
        type: String(match.type || 'output')
      }
    }
    for (const item of value) {
      const nested = findStill(item, depth + 1)
      if (nested) return nested
    }
    return null
  }
  const record = value as { filename?: string; subfolder?: string; type?: string; images?: unknown }
  if (record.filename && isStillFile(record)) {
    return {
      filename: String(record.filename),
      subfolder: String(record.subfolder || ''),
      type: String(record.type || 'output')
    }
  }
  for (const nested of Object.values(value as Record<string, unknown>)) {
    const found = findStill(nested, depth + 1)
    if (found) return found
  }
  return null
}

export function extractEditedImage(history: Record<string, unknown> | null, promptId: string) {
  if (!history) return null
  const wrapped = history[promptId] as { outputs?: Record<string, unknown> } | undefined
  if (wrapped) return findStill(wrapped.outputs || {}) || findStill(wrapped)
  if ((history as { outputs?: unknown }).outputs) {
    return findStill((history as { outputs?: unknown }).outputs) || findStill(history)
  }
  return findStill(history)
}

export async function downloadEditedImage(image: { filename: string; subfolder: string; type: string }) {
  const params = new URLSearchParams({
    filename: image.filename,
    subfolder: image.subfolder || '',
    type: image.type || 'output'
  })
  const res = await imageComfyFetch(`/view?${params.toString()}`)
  if (!res.ok) {
    throw createError({ statusCode: 502, statusMessage: `Failed to fetch edited still (${res.status})` })
  }
  return Buffer.from(await res.arrayBuffer())
}

function historyError(history: Record<string, unknown> | null, promptId: string) {
  const entry = history?.[promptId] as { status?: { status_str?: string; completed?: boolean } } | undefined
  const status = entry?.status?.status_str
  if (status === 'interrupted') return 'Job interrupted.'
  if (status === 'error') return 'Image ComfyUI reported an execution error'
  if (entry?.status?.completed && !extractEditedImage(history, promptId)) {
    return 'Image ComfyUI finished without an output still'
  }
  return null
}

export async function waitForImageEdit(opts: {
  promptId: string
  clientId: string
  timeoutMs?: number
  onProgress?: (event: { message: string; progress: number; step?: number; maxStep?: number; node?: string | null }) => void
  isCancelled?: () => boolean
}) {
  const timeoutMs = opts.timeoutMs || 180_000
  const started = Date.now()
  let settled = false
  let lastError: string | null = null
  let result: { filename: string; subfolder: string; type: string } | null = null

  const finish = (image: { filename: string; subfolder: string; type: string } | null, error?: string) => {
    if (settled) return
    settled = true
    lastError = error || null
    result = image
  }

  let ws: WebSocket | null = null
  try {
    ws = new WebSocket(imageComfyWsUrl(opts.clientId))
    ws.addEventListener('message', (event) => {
      const payload = JSON.parse(String(event.data || '{}')) as {
        type?: string
        data?: { value?: number; max?: number; node?: string | null; prompt_id?: string }
      }
      if (payload.type === 'progress' && payload.data) {
        const max = Math.max(1, Number(payload.data.max || 4))
        const value = Number(payload.data.value || 0)
        opts.onProgress?.({
          message: 'Sampling edit...',
          progress: Math.min(90, 20 + Math.round((value / max) * 60)),
          step: value,
          maxStep: max,
          node: payload.data.node || '8'
        })
      }
      if (payload.type === 'executing' && payload.data?.node) {
        opts.onProgress?.({
          message: payload.data.node === '8' ? 'Sampling edit...' : 'Running Qwen Edit...',
          progress: payload.data.node === '10' ? 92 : 30,
          node: payload.data.node
        })
      }
      if (payload.type === 'execution_error') {
        finish(null, 'Image ComfyUI reported an execution error')
      }
    })
  } catch {
    ws = null
  }

  while (!settled && Date.now() - started < timeoutMs) {
    if (opts.isCancelled?.()) {
      finish(null, 'Job interrupted.')
      break
    }
    const history = await fetchImageHistory(opts.promptId)
    const image = extractEditedImage(history, opts.promptId)
    if (image) {
      finish(image)
      break
    }
    const error = historyError(history, opts.promptId)
    if (error) {
      finish(null, error)
      break
    }
    await new Promise(resolve => setTimeout(resolve, 700))
  }

  try { ws?.close() } catch { /* ignore */ }
  if (result) return result
  throw createError({
    statusCode: 502,
    statusMessage: lastError || 'Image edit timed out waiting for Qwen Edit to finish'
  })
}