Add optional reference photo conditioning to Flux generation

server/api/v2/generate.post.ts

@@ -183,9 +183,10 @@ export default defineEventHandler(async (event) => {
   if (mode === 'iterate' && !imageA && (uploadedA || String(fields.image_a || fields.stillId || '').trim())) {
     throw createError({ statusCode: 400, statusMessage: 'Could not load still A for iterate' })
   }
-  const imageB = mode === 'refine' || mode === 'generate' || (mode === 'iterate' && !imageA)
+  const referenceGenerate = (mode === 'generate' || (mode === 'iterate' && !imageA)) && parseImageV2Engine(fields.engine, 'flux') === 'flux'
+  const imageB = mode === 'refine' || (!referenceGenerate && (mode === 'generate' || (mode === 'iterate' && !imageA)))
     ? null
-    : await resolveImageRef(ownerKey, fields.image_b, uploadedB).catch(() => null)
+    : await resolveImageRef(ownerKey, fields.image_b, uploadedB)
   const mask = mode === 'refine' ? await resolveImageRef(ownerKey, fields.mask, uploadedMask) : null
   const graphMode = mode === 'iterate'
     ? (imageA ? (imageB ? 'compose' as const : 'edit' as const) : 'generate' as const)

server/utils/fluxReference.ts

@@ -0,0 +1,12 @@
+/** Reference conditioning only: the sampler still starts from its empty generated latent. */
+export function addFluxGenerationReference(graph: Record<string, any>, filename: string) {
+  if (!filename) return
+  graph['reference:load'] = { class_type: 'LoadImage', inputs: { image: filename }, _meta: { title: 'Load reference photo' } }
+  graph['reference:scale'] = { class_type: 'ImageScaleToTotalPixels', inputs: { image: ['reference:load', 0], upscale_method: 'lanczos', megapixels: 1, resolution_steps: 1 }, _meta: { title: 'Scale reference photo' } }
+  graph['reference:encode'] = { class_type: 'VAEEncode', inputs: { pixels: ['reference:scale', 0], vae: ['6', 0] }, _meta: { title: 'Encode reference photo' } }
+  for (const [name, prompt] of [['positive', '9'], ['negative', '10']]) {
+    const id = `reference:${name}`
+    graph[id] = { class_type: 'ReferenceLatent', inputs: { conditioning: [prompt, 0], latent: ['reference:encode', 0] }, _meta: { title: 'Reference photo guidance' } }
+    graph['18'].inputs[name] = [id, 0]
+  }
+}

server/utils/imageChainV2.ts

@@ -118,9 +118,9 @@ export async function runEditV2(job: Job, params: EditV2RunParams) {
       let uploadedMask: { name: string } | null = null
 
       await ensureComfyLoraNames('image')
-      if (!generate && params.megapixels > 0) await assertImageScaleToTotalPixelsNode()
+      if ((!generate && params.megapixels > 0) || (generate && params.reference)) await assertImageScaleToTotalPixelsNode()
       const krea = engine === 'krea' ? await resolveKreaGenerateAssets() : null
-      if (!generate && graphMode === 'compose' && params.reference) {
+      if ((graphMode === 'compose' || (generate && engine === 'flux')) && params.reference) {
         uploadedRef = await uploadImage({
           ...params.reference,
           filename: `ref_${params.reference.filename || 'image_b.png'}`

server/utils/imageWorkflowV2.ts

@@ -1,5 +1,6 @@
 import editTemplate from '../assets/klein_v2_edit.json'
 import composeTemplate from '../assets/klein_v2_compose.json'
+import { addFluxGenerationReference } from './fluxReference'
 import refineTemplate from '../assets/klein_v2_refine.json'
 import generateTemplate from '../assets/klein_v2_generate.json'
 import kreaGenerateTemplate from '../assets/krea_v2_generate.json'
@@ -304,7 +305,7 @@ export function assertImageV2Graph(graph: WorkflowGraph, mode: ImageV2Mode, imag
   } else if (engine === 'krea') {
     assertKreaItiGraph(graph, mode)
   } else if (mode === 'generate') {
-    if (loaders.length) {
+    if (loaders.length && !(imageBName && loaders.length === 1 && graph['reference:load']?.inputs.image === imageBName)) {
       throw createError({
         statusCode: 500,
         statusMessage: 'Generate graph has a required LoadImage. Refusing to run an edit fallback.'
@@ -545,10 +546,11 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) {
   if (!generate) applyScaleMegapixels(graph, params.megapixels)
 
   const userLoras = applyImageV2UserLoras(graph, imageV2UserStack(params, 'flux'), 'flux')
+  if (generate && params.imageBName) addFluxGenerationReference(graph, params.imageBName)
   const guider = Object.values(graph).find(node => node.class_type === 'CFGGuider' || node.class_type === 'BasicGuider')
   const modelFeed = Array.isArray(guider?.inputs.model) ? String(guider.inputs.model[0] || '') : ''
 
-  assertImageV2Graph(graph, params.mode, generate || refine ? undefined : params.imageBName, engine)
+  assertImageV2Graph(graph, params.mode, refine ? undefined : params.imageBName, engine)
 
   const loaders = loadImageNames(graph)
   console.log(JSON.stringify({

server/utils/studioQueue.ts

@@ -1149,6 +1149,9 @@ async function startStudioEditJob(item: StudioJob) {
       }
       : null
     const refId = payload.referenceStillId
+    if (refId && !existsSync(stillPath(item.ownerKey, refId))) {
+      throw new Error('The reference photo is missing from the library')
+    }
     const reference = refId && existsSync(stillPath(item.ownerKey, refId))
       ? {
         filename: payload.referenceStillFilename || 'image2.png',
@@ -1230,7 +1233,7 @@ async function startStudioEditJob(item: StudioJob) {
         engine: v2Engine,
         task: payload.v2Task || (graphMode === 'generate' ? 't2i' : 'scene'),
         image: graphMode === 'generate' ? null : image,
-        reference: graphMode === 'compose' ? reference : null,
+        reference: graphMode === 'compose' || (graphMode === 'generate' && v2Engine === 'flux') ? reference : null,
         mask,
         prompt: payload.prompt,
         negative: payload.negative || '',