generate-image.ts

#!/usr/bin/env node

import { $ } from 'bun'

const COMFYUI_HOST = process.env.COMFYUI_HOST || '192.168.1.212'
const COMFYUI_PORT = process.env.COMFYUI_PORT || '8188'
const WORKFLOW_FILE = process.env.WORKFLOW_FILE || '/home/grishberg/ComfyUI/workflows/flux2_txt2img.json'

async function checkComfyUI() {
  try {
    await fetch(`http://${COMFYUI_HOST}:${COMFYUI_PORT}/`)
    return true
  } catch {
    return false
  }
}

async function queuePrompt(prompt: string, seed?: number) {
  const workflow = await Bun.file(WORKFLOW_FILE).text()
  
  let updatedWorkflow = workflow.replace(
    /"text":\s*""/,
    `"text": ${JSON.stringify(prompt)}`
  )
  
  const parsed = JSON.parse(updatedWorkflow)
  if (seed !== undefined) {
    parsed['25']['inputs']['noise_seed'] = seed
  }
  
  const response = await fetch(`http://${COMFYUI_HOST}:${COMFYUI_PORT}/prompt`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ prompt: parsed })
  })
  
  const data = await response.json()
  return data.prompt_id
}

async function waitForCompletion(promptId: string) {
  const startTime = Date.now()
  const timeout = 5 * 60 * 1000 // 5 minutes
  
  while (Date.now() - startTime < timeout) {
    const queueResponse = await fetch(`http://${COMFYUI_HOST}:${COMFYUI_PORT}/queue`)
    const queueData = await queueResponse.json()
    
    if (queueData.queue_running?.some((r: any[]) => r.includes(promptId))) {
      await new Promise(resolve => setTimeout(resolve, 5000))
      continue
    }
    
    const historyResponse = await fetch(`http://${COMFYUI_HOST}:${COMFYUI_PORT}/history/${promptId}`)
    const history = await historyResponse.json()
    
    if (history && Object.keys(history).length > 0) {
      return history[promptId]
    }
    
    await new Promise(resolve => setTimeout(resolve, 5000))
  }
  
  throw new Error('Generation timed out after 5 minutes')
}

async function findOutputImage() {
  const outputDir = `/home/grishberg/ComfyUI/output`
  const files = await Bun.file(outputDir).list()
  const fluxFiles = files
    .filter(f => f.name.startsWith('flux2_gen_') && f.name.endsWith('.png'))
    .sort((a, b) => b.mtime.getTime() - a.mtime.getTime())
  
  return fluxFiles[0]?.name ? `${outputDir}/${fluxFiles[0].name}` : null
}

export async function generateImage(prompt: string, output?: string, seed?: number) {
  if (!(await checkComfyUI())) {
    throw new Error(`ComfyUI is not running at ${COMFYUI_HOST}:${COMFYUI_PORT}`)
  }
  
  console.log(`Generating image... (seed: ${seed || 'random'})`)
  
  const promptId = await queuePrompt(prompt, seed)
  console.log(`Queue: ${promptId}`)
  
  const result = await waitForCompletion(promptId)
  
  const outputPath = await findOutputImage()
  
  if (!outputPath) {
    throw new Error('Output image not found')
  }
  
  if (output) {
    await Bun.file(output).write(Bun.file(outputPath).bytes())
    console.log(`Saved to: ${output}`)
  } else {
    console.log(`Output: ${outputPath}`)
  }
  
  return outputPath
}

if (require.main === module) {
  const args = process.argv.slice(2)
  
  if (args.length < 2) {
    console.error('Usage: generate-image <prompt> <output.png> [seed]')
    process.exit(1)
  }
  
  const [prompt, output, seedStr] = args
  const seed = seedStr ? parseInt(seedStr) : undefined
  
  generateImage(prompt, output, seed).catch(err => {
    console.error('Error:', err.message)
    process.exit(1)
  })
}