Preserve extension quality with original segments and reliable refinement

server/utils/pending.ts

@@ -1,3 +1,4 @@
+import { videoSourcePaths } from './videoSources'
 import { existsSync, mkdirSync, readdirSync, readFileSync, renameSync, rmSync, writeFileSync } from 'node:fs'
 import { join } from 'node:path'
 import type { Job } from '~/server/utils/jobs'
@@ -295,6 +296,11 @@ export async function completePendingIfReady(pending: PendingJob) {
   const video = extractVideo(history, pending.promptId)
   if (!video) return null
   let buffer = await downloadComfyVideo(video)
+  const originalSegment = buffer
+  const parentVideo = pending.parentClipId ? clipVideoPath(pending.ownerKey, pending.parentClipId) : ''
+  const sourceSegments = pending.extendPart1Path
+    ? videoSourcePaths(parentVideo && existsSync(parentVideo) ? parentVideo : pending.extendPart1Path)
+    : []
   if (inFlight()) return null
   const afterDownload = alreadySaved()
   if (afterDownload) return afterDownload
@@ -313,13 +319,13 @@ export async function completePendingIfReady(pending: PendingJob) {
       buffer = await stitchExtension({
         part1Path: pending.extendPart1Path,
         part2: buffer,
+        sourcePaths: sourceSegments,
         tmpDir: pending.extendTmpDir
       })
     } catch (error) {
       removeExtendTemp(pending.extendTmpDir)
       throw error
     }
-    removeExtendTemp(pending.extendTmpDir)
   }
   const afterStitch = alreadySaved()
   if (afterStitch) return afterStitch
@@ -349,9 +355,12 @@ export async function completePendingIfReady(pending: PendingJob) {
     permanenceRefs: mergePermanenceRefs(pending.permanenceRefs, pending.shotPermanenceRefs?.[pending.chainIndex || 0]),
     loraName: pending.loraName,
     loraStack: pending.loraStack,
-    segmentFirstFrame
+    segmentFirstFrame,
+    originalSegment,
+    sourceSegments
   })
-  if (pending.saveLosslessAnchor === true) {
+  if (pending.extendTmpDir) removeExtendTemp(pending.extendTmpDir)
+  if (pending.saveLosslessAnchor !== false) {
     try {
       const { persistClipAnchorFrame, extractAnchorImageFromHistory, downloadAnchorImage } = await import('~/server/utils/extensionFrame')
       let saved = false