aigen/server/utils/extensionFrame.ts

import { videoSourcePaths } from './videoSources'
import { copyFileSync, existsSync, statSync, writeFileSync } from 'node:fs'
import refineTemplate from '../assets/extension_frame_refine.json'
import { clipAnchorFramePath, clipVideoPath, downloadComfyImage } from '~/server/utils/library'
import { extractLastFrame } from '~/server/utils/ffmpeg'
import { comfyFilenamePrefix, fetchHistory, queuePrompt, uploadImage } from '~/server/utils/comfy'
import { ensureComfyLoraNames, resolveGraphLoraNames } from '~/server/utils/loras'
import { extractEditedImage } from '~/server/utils/imageComfy'
import {
  clampRefinementDenoise,
  REFINEMENT_DENOISE_DEFAULT
} from '~/utils/extensionFrame'

type WorkflowNode = { class_type: string; inputs: Record<string, unknown>; _meta?: { title?: string } }
type WorkflowGraph = Record<string, WorkflowNode>

const MIN_ANCHOR_BYTES = 64

export function isValidAnchorFile(path: string) {
  try {
    return existsSync(path) && statSync(path).size >= MIN_ANCHOR_BYTES
  } catch {
    return false
  }
}

export async function resolveExtensionHandoffFrame(opts: {
  ownerKey: string
  sourceClipId?: string
  sourceVideoPath: string
  destPath: string
}) {
  const anchorPath = opts.sourceClipId
    ? clipAnchorFramePath(opts.ownerKey, opts.sourceClipId)
    : ''
  if (anchorPath && isValidAnchorFile(anchorPath)) {
    copyFileSync(anchorPath, opts.destPath)
    return { path: opts.destPath, source: 'anchor' as const }
  }
  const originalSources = videoSourcePaths(opts.sourceVideoPath)
  await extractLastFrame(originalSources[originalSources.length - 1], opts.destPath)
  if (!isValidAnchorFile(opts.destPath)) {
    throw new Error('Could not extract the last frame of the clip being extended')
  }
  return { path: opts.destPath, source: 'extract' as const }
}

export async function persistClipAnchorFrame(opts: {
  ownerKey: string
  clipId: string
  fromPath?: string
  fromBuffer?: Buffer
  videoPath?: string
}) {
  const dest = clipAnchorFramePath(opts.ownerKey, opts.clipId)
  if (opts.fromBuffer?.length && opts.fromBuffer.length >= MIN_ANCHOR_BYTES) {
    writeFileSync(dest, opts.fromBuffer)
    return dest
  }
  if (opts.fromPath && isValidAnchorFile(opts.fromPath)) {
    copyFileSync(opts.fromPath, dest)
    return dest
  }
  const video = opts.videoPath || clipVideoPath(opts.ownerKey, opts.clipId)
  if (!existsSync(video)) return ''
  const originalSources = videoSourcePaths(video)
  await extractLastFrame(originalSources[originalSources.length - 1], dest)
  return isValidAnchorFile(dest) ? dest : ''
}

export function buildRefineGraph(opts: {
  imageName: string
  denoise: number
  seed: number
  filenamePrefix: string
}) {
  const graph = structuredClone(refineTemplate) as WorkflowGraph
  if (graph['1']?.inputs) graph['1'].inputs.image = opts.imageName
  if (graph['13']?.inputs) graph['13'].inputs.noise_seed = opts.seed
  if (graph['15']?.inputs) graph['15'].inputs.denoise = opts.denoise
  if (graph['20']?.inputs) graph['20'].inputs.filename_prefix = opts.filenamePrefix
  return graph
}

async function waitForRefineImage(promptId: string, timeoutMs = 120_000) {
  const started = Date.now()
  while (Date.now() - started < timeoutMs) {
    const history = await fetchHistory(promptId)
    const image = extractEditedImage(history, promptId)
    if (image) return image
    const entry = history?.[promptId] as { status?: { status_str?: string } } | undefined
    const status = entry?.status?.status_str
    if (status === 'error' || status === 'interrupted') {
      throw new Error(status === 'interrupted' ? 'Frame refine interrupted' : 'Frame refine failed on ComfyUI')
    }
    await new Promise(resolve => setTimeout(resolve, 700))
  }
  throw new Error('Frame refine timed out waiting for ComfyUI')
}

/**
 * Low-denoise Flux cleanup before the video I2V start frame.
 * Runs as its own Comfy prompt so MiniMax does not share VRAM with Klein.
 * When refine is off, callers skip this and upload the PNG directly (fast path).
 */
export async function refineExtensionHandoffFrame(opts: {
  frame: Buffer
  denoise?: number
  seed?: number
  clientId: string
  jobId: string
  onStatus?: (message: string) => void
  onQueued?: (promptId: string) => void
}) {
  const denoise = clampRefinementDenoise(opts.denoise, REFINEMENT_DENOISE_DEFAULT)
  opts.onStatus?.('Refining extension handoff frame...')
  const uploaded = await uploadImage({
    filename: 'last_frame.png',
    data: opts.frame,
    type: 'image/png'
  }, opts.jobId)
  const prefix = `${comfyFilenamePrefix()}_extrefine`
  const graph = buildRefineGraph({
    imageName: uploaded.name,
    denoise,
    seed: opts.seed ?? Math.floor(Math.random() * 2_147_483_647),
    filenamePrefix: `aigen/extension_refine/${prefix}`
  })
  await ensureComfyLoraNames('image')
  resolveGraphLoraNames(graph, 'image')
  const queued = await queuePrompt(graph, opts.clientId)
  opts.onQueued?.(queued.prompt_id)
  const image = await waitForRefineImage(queued.prompt_id)
  const buffer = await downloadComfyImage(image)
  if (!buffer.length || buffer.length < MIN_ANCHOR_BYTES) {
    throw new Error('Frame refine finished without an output still')
  }
  return buffer
}

export function extractAnchorImageFromHistory(history: Record<string, unknown> | null, promptId: string) {
  if (!history || !promptId) return null
  const entry = history[promptId] as { outputs?: Record<string, unknown> } | undefined
  const outputs = entry?.outputs
  if (!outputs || typeof outputs !== 'object') return null
  const anchor = outputs.aigen_anchor_save as { images?: Array<{ filename?: string; subfolder?: string; type?: string }> } | undefined
  const file = anchor?.images?.[0]
  if (file?.filename) {
    return {
      filename: file.filename,
      subfolder: file.subfolder || '',
      type: file.type || 'output'
    }
  }
  return null
}

export async function downloadAnchorImage(image: { filename: string; subfolder: string; type: string }) {
  return downloadComfyImage(image)
}