aigen/server/utils/extendChain.ts

import { copyFileSync, existsSync, readFileSync, writeFileSync } from 'node:fs'
import { join } from 'node:path'
import { composeShotPrompt } from '~/utils/globalLocks'
import { composePromptParts } from '~/utils/promptParts'
import { parsePostedLoraStack, listStudioLoras, persistLoraFields, ensureComfyLoraNames } from '~/server/utils/loras'
import { readLoraStack } from '~/utils/loras'
import { clampVideoCfg } from '~/utils/generationPresets'
import { defaultVideoSteps, isLtxWorkflow, LTX_DISABLED_MESSAGE, ltxWorkflowEnabled, parseVideoWorkflow } from '~/utils/videoModels'
import { clampVideoDuration, clipPartDuration } from '~/utils/videoDuration'
import { createJob, emitJob, type Job } from '~/server/utils/jobs'
import {
  clipTitle,
  clipVideoPath,
  extendTempDir,
  getClip,
  nextFamilyPartName,
  removeExtendTemp
} from '~/server/utils/library'
import { probeHasAudio } from '~/server/utils/ffmpeg'
import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
import { waitForComfySocket, watchComfyJob } from '~/server/utils/watch'
import { comfyFilenamePrefix, queuePrompt, uploadImage } from '~/server/utils/comfy'
import { buildWorkflow } from '~/server/utils/workflow'
import { writePendingJob } from '~/server/utils/pending'
import {
  refineExtensionHandoffFrame,
  resolveExtensionHandoffFrame
} from '~/server/utils/extensionFrame'
import {
  clampRefinementDenoise,
  parseRefineExtensionFrame,
  parseSaveLosslessAnchor,
  REFINEMENT_DENOISE_DEFAULT
} from '~/utils/extensionFrame'

const SAMPLERS = new Set(['res_multistep', 'euler', 'dpmpp_2m'])
const SCHEDULERS = new Set(['simple', 'ddim_uniform', 'sgm_uniform'])

export function parseExtendDuration(raw: unknown, fallback = 5) {
  return clampVideoDuration(raw, fallback)
}

function frameLength(seconds: number, fps: number) {
  return Math.max(5, Math.floor(seconds * fps))
}

function parseFps(raw: unknown) {
  const fps = Number(raw)
  return fps === 12 || fps === 30 || fps === 24 ? fps : 24
}

function parseCfg(raw: unknown, turbo: boolean) {
  return clampVideoCfg(raw, turbo ? 1.5 : 4)
}

function parseSampler(raw: unknown) {
  const value = String(raw || '')
  return SAMPLERS.has(value) ? value : 'res_multistep'
}

function parseScheduler(raw: unknown) {
  const value = String(raw || '')
  return SCHEDULERS.has(value) ? value : 'simple'
}

export async function beginExtendFromClip(params: {
  ownerKey: string
  clipId: string
  prompt: string
  promptPre?: string
  promptPost?: string
  duration?: number
  loraStack?: unknown
  folderLocked?: boolean
  name?: string
  refineExtensionFrame?: boolean
  saveLosslessAnchor?: boolean
  refinementDenoise?: number
}) {
  const source = getClip(params.ownerKey, params.clipId)
  const sourcePath = clipVideoPath(params.ownerKey, source.id)
  if (!existsSync(sourcePath)) {
    throw createError({ statusCode: 404, statusMessage: 'Source video is missing' })
  }
  const prompt = String(params.prompt || '').trim()
  const parent = source.parentClipId
    ? (() => { try { return getClip(params.ownerKey, source.parentClipId!) } catch { return null } })()
    : null
  const inherited = clipPartDuration(source, parent)
  const durationSeconds = params.duration == null
    ? inherited
    : parseExtendDuration(params.duration, inherited)
  const promptPre = params.promptPre != null
    ? String(params.promptPre).trim()
    : String(source.promptPre || '').trim()
  const promptPost = params.promptPost != null
    ? String(params.promptPost).trim()
    : String(source.promptPost || '').trim()
  if (!composePromptParts(promptPre, prompt, promptPost)) {
    throw createError({ statusCode: 400, statusMessage: 'An extension prompt is required' })
  }
  const composedPrompt = composeShotPrompt({
    globalLocks: source.globalLocks,
    prompt: composePromptParts(promptPre, prompt, promptPost),
    shotIndex: (source.chainIndex || 0) + 1,
    familyRefs: source.permanenceRefs
  })
  const workflow = parseVideoWorkflow(source.workflow)
  if (isLtxWorkflow(workflow) && !ltxWorkflowEnabled()) {
    throw createError({ statusCode: 400, statusMessage: LTX_DISABLED_MESSAGE })
  }
  const turbo = source.turbo !== false
  const fps = parseFps(source.fps)
  const cfg = parseCfg(source.cfg, turbo)
  const samplerName = parseSampler(source.samplerName)
  const scheduler = parseScheduler(source.scheduler)
  const steps = Number(source.steps || defaultVideoSteps(turbo, workflow))
  const seed = Math.floor(Math.random() * 2_147_483_647)
  await listStudioLoras()
  const postedStack = parsePostedLoraStack(params.loraStack, 'video')
  const loraStack = postedStack.length ? postedStack : parsePostedLoraStack(readLoraStack(source), 'video')
  const loraFields = persistLoraFields(loraStack)
  const length = frameLength(durationSeconds, fps)
  const width = source.width || 1344
  const height = source.height || 768
  const clipName = String(params.name || '').trim() || nextFamilyPartName(params.ownerKey, source)
  const refineExtensionFrame = parseRefineExtensionFrame(params.refineExtensionFrame, true)
  const saveLosslessAnchor = parseSaveLosslessAnchor(params.saveLosslessAnchor, true)
  const refinementDenoise = clampRefinementDenoise(params.refinementDenoise, REFINEMENT_DENOISE_DEFAULT)

  const job = createJob()
  job.kind = 'video'
  job.maxStep = steps
  job.hideThumbnail = source.hideThumbnail
  const tmpDir = extendTempDir(params.ownerKey, job.id)
  const part1Path = join(tmpDir, 'part1.mp4')
  const framePath = join(tmpDir, 'last_frame.png')
  job.library = {
    ownerKey: params.ownerKey,
    folderId: source.folderId,
    hideThumbnail: source.hideThumbnail,
    folderLocked: params.folderLocked,
    name: clipName,
    prompt,
    promptMid: prompt,
    promptPre: promptPre || undefined,
    promptPost: promptPost || undefined,
    aspect: source.aspect || '16:9',
    width,
    height,
    steps,
    turbo,
    seed,
    cfg,
    fps,
    samplerName,
    scheduler,
    duration: durationSeconds,
    extendTmpDir: tmpDir,
    extendPart1Path: part1Path,
    extendSourceClipId: source.id,
    familyId: source.familyId || crypto.randomUUID(),
    parentClipId: source.id,
    chainIndex: (source.chainIndex || 0) + 1,
    workflow,
    globalLocks: source.globalLocks,
    permanenceRefs: source.permanenceRefs,
    refineExtensionFrame,
    saveLosslessAnchor,
    refinementDenoise,
    ...loraFields
  }

  emitJob(job, { type: 'status', message: 'Extracting last frame & stitching extension...', progress: 2 })

  void runExtend(job, {
    sourcePath,
    part1Path,
    framePath,
    prompt: composedPrompt,
    width,
    height,
    steps,
    seed,
    turbo,
    length,
    cfg,
    fps,
    samplerName,
    scheduler,
    refineExtensionFrame,
    saveLosslessAnchor,
    refinementDenoise,
    ...loraFields
  }).catch((error) => {
    removeExtendTemp(job.library?.extendTmpDir)
    const message = error instanceof Error ? error.message : String(error)
    job.status = 'error'
    job.error = message
    emitJob(job, { type: 'error', error: message, message })
  })

  return job
}

async function runExtend(
  job: Job,
  params: {
    sourcePath: string
    part1Path: string
    framePath: string
    prompt: string
    width: number
    height: number
    steps: number
    seed: number
    turbo: boolean
    length: number
    cfg: number
    fps: number
    samplerName: string
    scheduler: string
    refineExtensionFrame?: boolean
    saveLosslessAnchor?: boolean
    refinementDenoise?: number
    loraName?: string
    loraStack?: import('~/utils/loras').LoraStackItem[]
  }
) {
  copyFileSync(params.sourcePath, params.part1Path)
  const handoff = await resolveExtensionHandoffFrame({
    ownerKey: job.library!.ownerKey,
    sourceClipId: job.library?.extendSourceClipId,
    sourceVideoPath: params.sourcePath,
    destPath: params.framePath
  })
  let frame = readFileSync(handoff.path)
  if (!frame.length || frame.length < 64) {
    throw new Error('Could not extract the last frame of the clip being extended')
  }
  if (job.library) job.library.thumb = undefined
  const sound = await probeHasAudio(params.part1Path)

  await ensureComfyReady((status) => {
    emitJob(job, {
      type: status.state === 'busy' ? 'busy' : 'status',
      message: status.message,
      progress: status.state === 'online' ? Math.max(job.progress, 3) : Math.max(job.progress, 2),
      busy: status.state === 'busy',
      queueRunning: status.queueRunning,
      queuePending: status.queuePending
    })
  })

  if (params.refineExtensionFrame !== false) {
    emitJob(job, {
      type: 'status',
      message: handoff.source === 'anchor'
        ? 'Refining lossless anchor frame...'
        : 'Refining extracted handoff frame...',
      progress: 3
    })
    try {
      frame = await refineExtensionHandoffFrame({
        frame,
        denoise: params.refinementDenoise,
        seed: params.seed,
        clientId: job.clientId,
        jobId: job.id,
        onQueued: (promptId) => { job.promptId = promptId },
        onStatus: (message) => emitJob(job, { type: 'status', message, progress: 4 })
      })
      writeFileSync(params.framePath, frame)
    } catch (error) {
      throw new Error(`Extension frame refinement failed: ${error instanceof Error ? error.message : String(error)}. Retry after fixing refinement, or explicitly turn it off.`)
    }
  }

  if (job.status === 'cancelled' || job.status === 'error') throw new Error('Extension stopped before video generation')
  job.promptId = undefined
  const done = watchComfyJob(job)
  job.status = 'uploading'
  emitJob(job, { type: 'status', message: 'Uploading last frame to ComfyUI...', progress: 5 })
  const uploaded = await uploadImage({ filename: 'last_frame.png', data: frame, type: 'image/png' }, job.id)
  if (job.library) {
    job.library.imageName = uploaded.name
    job.library.imageSubfolder = uploaded.subfolder
  }
  emitJob(job, { type: 'status', message: 'Queueing extension on MiniMax H3...', progress: 6 })
  await waitForComfySocket(job, 4000)
  await ensureComfyLoraNames('video')

  const graph = buildWorkflow({
    prompt: params.prompt,
    imageName: uploaded.name,
    width: params.width,
    height: params.height,
    steps: params.steps,
    seed: params.seed,
    turbo: params.turbo,
    length: params.length,
    cfg: params.cfg,
    fps: params.fps,
    samplerName: params.samplerName,
    scheduler: params.scheduler,
    filenamePrefix: comfyFilenamePrefix(),
    sound,
    workflow: job.library?.workflow,
    duration: job.library?.duration,
    loraStack: params.loraStack || job.library?.loraStack || params.loraName || job.library?.loraName,
    saveLosslessAnchor: (params.saveLosslessAnchor ?? job.library?.saveLosslessAnchor) !== false
  })

  const queued = await queuePrompt(graph, job.clientId)
  job.promptId = queued.prompt_id
  job.status = 'running'
  if (job.library && job.promptId) {
    writePendingJob({
      jobId: job.id,
      promptId: job.promptId,
      clientId: job.clientId,
      ownerKey: job.library.ownerKey,
      folderId: job.library.folderId,
      hideThumbnail: job.library.hideThumbnail,
      folderLocked: job.library.folderLocked,
      name: job.library.name,
      prompt: job.library.prompt,
      aspect: job.library.aspect,
      width: job.library.width,
      height: job.library.height,
      steps: job.library.steps,
      turbo: job.library.turbo,
      seed: job.library.seed,
      startedAt: job.startedAt,
      imageName: job.library.imageName,
      imageSubfolder: job.library.imageSubfolder,
      extendTmpDir: job.library.extendTmpDir,
      extendPart1Path: job.library.extendPart1Path,
      familyId: job.library.familyId,
      parentClipId: job.library.parentClipId,
      chainIndex: job.library.chainIndex,
      stillId: job.library.stillId,
      sound: job.library.sound,
      workflow: job.library.workflow,
      duration: job.library.duration,
      globalLocks: job.library.globalLocks,
      permanenceRefs: job.library.permanenceRefs,
      loraName: job.library.loraName,
      loraStack: job.library.loraStack
    })
  }
  emitJob(job, { type: 'status', message: 'Job queued on ComfyUI', progress: 8 })
  await done
}