Purge identity refs from desktop Comfy and keep them in the server stills library.

Co-authored-by: Cursor <cursoragent@cursor.com>

scripts/comfy-host-agent.mjs

@@ -181,9 +181,12 @@ function readJson(req) {
 
 function purgeDesktopFiles(body) {
   const deleted = []
-  const imageName = String(body?.imageName || body?.filename || '')
+  const names = [
+    String(body?.imageName || body?.filename || ''),
+    ...(Array.isArray(body?.imageNames) ? body.imageNames : [])
+  ].map(name => String(name || '').trim()).filter(Boolean)
   const imageSub = String(body?.imageSubfolder || '')
-  if (imageName) {
+  for (const imageName of names) {
     for (const root of inputRoots()) {
       if (!existsSync(root)) continue
       const nested = safeFile(root, imageSub, imageName)
@@ -220,7 +223,7 @@ function purgeDesktopFiles(body) {
       }
     }
   }
-  console.log(JSON.stringify({ src: 'comfy-host-agent', event: 'purge', deleted: deleted.length, files: deleted.map((path) => basename(path)), imageName }))
+  console.log(JSON.stringify({ src: 'comfy-host-agent', event: 'purge', deleted: deleted.length, files: deleted.map((path) => basename(path)), imageNames: names }))
   return { ok: true, deleted }
 }
 

server/api/generate.post.ts

@@ -152,6 +152,20 @@ export default defineEventHandler(async (event) => {
       height,
       hideInput
     })
+  if (useIdentityRefs) {
+    for (const [index, ref] of referenceImages.entries()) {
+      if (!ref?.data?.length || isPipelineFrameFilename(ref.filename)) continue
+      await saveStill({
+        ownerKey,
+        folderId,
+        filename: ref.filename || `identity-ref-${index + 2}.png`,
+        data: ref.data,
+        width,
+        height,
+        hideInput
+      })
+    }
+  }
 
   const chainTotal = 1 + extensions.length
   const job = createJob()
@@ -332,6 +346,7 @@ async function queueMiniMax(
   if (job.library) {
     job.library.imageName = uploaded.name
     job.library.imageSubfolder = uploaded.subfolder
+    job.library.referenceImageNames = referenceNames.filter(Boolean)
   }
   emitChainJob(job, { type: 'status', message: queueing, progress: 6 })
   await waitForComfySocket(job, 4000)
@@ -390,7 +405,8 @@ async function queueMiniMax(
       parentClipId: job.library.parentClipId,
       chainIndex: job.library.chainIndex,
       stillId: job.library.stillId,
-      sound: job.library.sound
+      sound: job.library.sound,
+      referenceImageNames: job.library.referenceImageNames
     })
   }
   emitChainJob(job, { type: 'status', message: 'Job queued on ComfyUI', progress: 8 })

server/utils/comfy.ts

@@ -293,6 +293,7 @@ async function purgeOnDesktop(opts: {
   video?: { filename: string; subfolder: string; type: string }
   imageName?: string
   imageSubfolder?: string
+  extraImageNames?: string[]
 }) {
   const config = useRuntimeConfig()
   const controlUrl = String(config.comfyControlUrl || process.env.COMFY_CONTROL_URL || '').replace(/\/$/, '')
@@ -309,6 +310,7 @@ async function purgeOnDesktop(opts: {
       body: JSON.stringify({
         imageName: opts.imageName || '',
         imageSubfolder: opts.imageSubfolder || '',
+        imageNames: (opts.extraImageNames || []).filter(Boolean),
         video: opts.video || null
       }),
       signal: AbortSignal.timeout(8000)
@@ -322,6 +324,7 @@ export async function purgeComfyArtifacts(opts: {
   video?: { filename: string; subfolder: string; type: string }
   imageName?: string
   imageSubfolder?: string
+  extraImageNames?: string[]
   promptId?: string
 }) {
   if (!purgeComfyEnabled()) return
@@ -329,6 +332,10 @@ export async function purgeComfyArtifacts(opts: {
   const files: { filename: string; subfolder: string; type: string }[] = []
   if (opts.video?.filename) files.push(opts.video)
   if (opts.imageName) files.push({ filename: opts.imageName, subfolder: opts.imageSubfolder || '', type: 'input' })
+  for (const name of opts.extraImageNames || []) {
+    if (!name || name === opts.imageName) continue
+    files.push({ filename: name, subfolder: opts.imageSubfolder || '', type: 'input' })
+  }
   for (const file of files) {
     try {
       await comfyFetch('/aigen/purge', {

server/utils/jobs.ts

@@ -66,6 +66,7 @@ export interface Job {
     imageSubfolder?: string
     stillId?: string
     stillFilename?: string
+    referenceImageNames?: string[]
     duration?: number
     sound?: boolean
     draftId?: string

server/utils/pending.ts

@@ -27,6 +27,7 @@ export interface PendingJob {
   chainIndex?: number
   stillId?: string
   sound?: boolean
+  referenceImageNames?: string[]
 }
 
 function pendingRoot() {
@@ -148,7 +149,13 @@ export async function completePendingIfReady(pending: PendingJob) {
     stillId: pending.stillId,
     sound: pending.sound
   })
-  await purgeComfyArtifacts({ video, imageName: pending.imageName, imageSubfolder: pending.imageSubfolder, promptId: pending.promptId })
+  await purgeComfyArtifacts({
+    video,
+    imageName: pending.imageName,
+    imageSubfolder: pending.imageSubfolder,
+    extraImageNames: pending.referenceImageNames,
+    promptId: pending.promptId
+  })
   deletePendingJob(pending.jobId)
   return {
     type: 'complete' as const,

server/utils/watch.ts

@@ -132,6 +132,11 @@ export function watchComfyJob(job: Job, options: { persist?: boolean } = {}): Pr
       emitChainJob(job, { type: 'error', error: job.error, message: job.error })
       if (job.id) deletePendingJob(job.id)
       removeExtendTemp(job.library?.extendTmpDir)
+      await purgeComfyArtifacts({
+        imageName: job.library?.imageName,
+        imageSubfolder: job.library?.imageSubfolder,
+        extraImageNames: job.library?.referenceImageNames
+      })
       resolve()
     }
 
@@ -209,6 +214,7 @@ export function watchComfyJob(job: Job, options: { persist?: boolean } = {}): Pr
               video,
               imageName: job.library.imageName,
               imageSubfolder: job.library.imageSubfolder,
+              extraImageNames: job.library.referenceImageNames,
               promptId: job.promptId
             })
             if (!persist) {
@@ -405,6 +411,7 @@ export function ensurePendingWatch(pending: {
   chainIndex?: number
   stillId?: string
   sound?: boolean
+  referenceImageNames?: string[]
 }) {
   const existing = getJob(pending.jobId)
   if (existing) return existing
@@ -435,7 +442,8 @@ export function ensurePendingWatch(pending: {
       parentClipId: pending.parentClipId,
       chainIndex: pending.chainIndex,
       stillId: pending.stillId,
-      sound: pending.sound
+      sound: pending.sound,
+      referenceImageNames: pending.referenceImageNames
     }
   })
   if (!pendingWatches.has(job.id)) {