Preserve extension quality with original segments and reliable refinement

server/utils/extensionFrame.ts

@@ -1,8 +1,10 @@
+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,
@@ -35,7 +37,8 @@ export async function resolveExtensionHandoffFrame(opts: {
     copyFileSync(anchorPath, opts.destPath)
     return { path: opts.destPath, source: 'anchor' as const }
   }
-  await extractLastFrame(opts.sourceVideoPath, opts.destPath)
+  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')
   }
@@ -60,11 +63,12 @@ export async function persistClipAnchorFrame(opts: {
   }
   const video = opts.videoPath || clipVideoPath(opts.ownerKey, opts.clipId)
   if (!existsSync(video)) return ''
-  await extractLastFrame(video, dest)
+  const originalSources = videoSourcePaths(video)
+  await extractLastFrame(originalSources[originalSources.length - 1], dest)
   return isValidAnchorFile(dest) ? dest : ''
 }
 
-function buildRefineGraph(opts: {
+export function buildRefineGraph(opts: {
   imageName: string
   denoise: number
   seed: number
@@ -95,7 +99,7 @@ async function waitForRefineImage(promptId: string, timeoutMs = 120_000) {
 }
 
 /**
- * Bilateral + low-denoise Flux cleanup before the video I2V start frame.
+ * 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).
  */
@@ -106,6 +110,7 @@ export async function refineExtensionHandoffFrame(opts: {
   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...')
@@ -121,7 +126,10 @@ export async function refineExtensionHandoffFrame(opts: {
     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) {
@@ -134,7 +142,7 @@ export function extractAnchorImageFromHistory(history: Record<string, unknown> |
   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 extractEditedImage(history, promptId)
+  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) {
@@ -144,7 +152,7 @@ export function extractAnchorImageFromHistory(history: Record<string, unknown> |
       type: file.type || 'output'
     }
   }
-  return extractEditedImage(history, promptId)
+  return null
 }
 
 export async function downloadAnchorImage(image: { filename: string; subfolder: string; type: string }) {