diff --git a/pages/index.vue b/pages/index.vue index f18e942..94bb75d 100644 --- a/pages/index.vue +++ b/pages/index.vue @@ -182,9 +182,9 @@ -
+
-

{{ studioMode === 'editv2' && v2Mode === 'compose' ? 'Still B · required' : studioMode === 'editv2' ? 'Still B · optional' : 'Image 2 · optional' }}

+

{{ fluxTextReference ? 'Reference photo · optional' : studioMode === 'editv2' && v2Mode === 'compose' ? 'Still B · required' : studioMode === 'editv2' ? 'Still B · optional' : 'Image 2 · optional' }}

- {{ studioMode === 'editv2' + {{ fluxTextReference ? 'Guides a new image from your text. This is not a starting image; your selected output size stays in effect. Describe which details to use from the photo.' : studioMode === 'editv2' ? (editRefFile ? stillBHint : 'Compose will reject this job until still B is loaded.') : (editRefFile ? 'Two-image branch on. Image 2 is applied onto image 1.' @@ -2716,6 +2716,7 @@ const v2Mode = computed({ set: mode => { selectedV2Mode.value = mode } }) const v2Engine = ref('flux') +const fluxTextReference = computed(() => studioMode.value === 'editv2' && v2Engine.value === 'flux' && (v2Mode.value === 'generate' || (v2Mode.value === 'iterate' && !file.value))) const v2Task = ref('scene') const v2Negative = ref('') const v2Steps = ref(IMAGE_V2_STEPS_DEFAULT) @@ -2923,7 +2924,7 @@ const REFS_STORE = 'aigen-permanence-refs' const pickerHeading = computed(() => { if (pickerKind.value === 'permanence') return 'Choose a still to name' if (pickerSlot.value === 'main') return studioMode.value === 'edit' || studioMode.value === 'editv2' ? (studioMode.value === 'editv2' ? 'Choose still A' : 'Choose image 1') : 'Choose an image' - if (pickerSlot.value === 'editRef') return studioMode.value === 'editv2' ? 'Choose still B' : 'Choose image 2' + if (pickerSlot.value === 'editRef') return fluxTextReference.value ? 'Choose a reference photo' : studioMode.value === 'editv2' ? 'Choose still B' : 'Choose image 2' return `Choose Picture ${Number(pickerSlot.value) + 2}` }) const permanencePicker = computed(() => pickerKind.value === 'permanence') @@ -5552,7 +5553,7 @@ async function restoreStudioEditJob(payload: { if (payload.hideInput && showPrivacyToggles.value) hideInputPreview.value = true readFile(new File([blob], stillFileName(stills.value.find(item => item.id === payload.stillId)) || payload.stillFilename || 'still.png', { type: blob.type || 'image/png' }), { stillId: payload.stillId }) } - if (v2 && payload.v2Mode === 'compose' && payload.referenceStillId) { + if (v2 && ['compose', 'generate', 'iterate'].includes(payload.v2Mode || '') && payload.referenceStillId) { await loadStillAsEditRef(payload.referenceStillId).catch(() => null) } else if (!(v2 && payload.v2Mode === 'compose')) { clearEditRef() @@ -6472,7 +6473,7 @@ async function useStillAsInput(still: LibraryStill) { const blob = await $fetch(`/api/library/stills/${still.id}`, { responseType: 'blob' }) if (still.hideInput && showPrivacyToggles.value) hideInputPreview.value = true readFile(new File([blob], stillFileName(still), { type: blob.type || 'image/png' }), { stillId: still.id }) - if (dest === 'imagev2' && still.settings?.mode === 'compose' && still.settings.referenceStillId) { + if (dest === 'imagev2' && ['compose', 'generate', 'iterate'].includes(still.settings?.mode || '') && still.settings?.referenceStillId) { try { await loadStillAsEditRef(still.settings.referenceStillId) } catch { @@ -7316,7 +7317,7 @@ async function editImageV2() { ? v2Task.value : 'scene') if (v2Mode.value !== 'generate' && file.value) body.append('image_a', file.value) - if ((v2Mode.value === 'compose' || v2Mode.value === 'iterate') && editRefFile.value) body.append('image_b', editRefFile.value) + if ((fluxTextReference.value || v2Mode.value === 'compose' || v2Mode.value === 'iterate') && editRefFile.value) body.append('image_b', editRefFile.value) if (v2Mode.value === 'refine' || (v2Engine.value === 'krea' && v2Mode.value !== 'generate' && !iterateGenerate)) { if (v2Mode.value === 'refine') { const maskBlob = await refinePainter.value?.exportPng() diff --git a/server/api/v2/generate.post.ts b/server/api/v2/generate.post.ts index dae7d4a..45cb51d 100644 --- a/server/api/v2/generate.post.ts +++ b/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) diff --git a/server/utils/fluxReference.ts b/server/utils/fluxReference.ts new file mode 100644 index 0000000..b1ed26a --- /dev/null +++ b/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, 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] + } +} diff --git a/server/utils/imageChainV2.ts b/server/utils/imageChainV2.ts index f8bdce0..a7a0b20 100644 --- a/server/utils/imageChainV2.ts +++ b/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'}` diff --git a/server/utils/imageWorkflowV2.ts b/server/utils/imageWorkflowV2.ts index 227a8b8..c6d1402 100644 --- a/server/utils/imageWorkflowV2.ts +++ b/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({ diff --git a/server/utils/studioQueue.ts b/server/utils/studioQueue.ts index b704766..8fb68b2 100644 --- a/server/utils/studioQueue.ts +++ b/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 || '', diff --git a/tests/flux-reference.test.mjs b/tests/flux-reference.test.mjs new file mode 100644 index 0000000..dd43308 --- /dev/null +++ b/tests/flux-reference.test.mjs @@ -0,0 +1,32 @@ +import test from 'node:test' +import assert from 'node:assert/strict' +import { readFileSync } from 'node:fs' +import ts from 'typescript' +const source = readFileSync(new URL('../server/utils/fluxReference.ts', import.meta.url), 'utf8') +const code = ts.transpileModule(source, { compilerOptions: { module: ts.ModuleKind.CommonJS } }).outputText +const api = {} +new Function('exports', code)(api) +const template = JSON.parse(readFileSync(new URL('../server/assets/klein_v2_generate.json', import.meta.url), 'utf8')) + +test('Flux reference conditions a new image without initializing its latent or changing output dimensions', () => { + const graph = structuredClone(template) + graph['14'].inputs.width = 1344 + graph['14'].inputs.height = 768 + const sampler = structuredClone(graph['19']) + api.addFluxGenerationReference(graph, 'reference.png') + assert.deepEqual(graph['19'], sampler) + assert.equal(graph['14'].class_type, 'EmptyFlux2LatentImage') + assert.equal(graph['14'].inputs.width, 1344) + assert.equal(graph['14'].inputs.height, 768) + assert.equal(graph['reference:load'].inputs.image, 'reference.png') + assert.deepEqual(graph['18'].inputs.positive, ['reference:positive', 0]) + assert.deepEqual(graph['reference:positive'].inputs.conditioning, ['9', 0]) + assert.deepEqual(graph['reference:negative'].inputs.conditioning, ['10', 0]) + assert.deepEqual(graph['reference:encode'].inputs.pixels, ['reference:scale', 0]) +}) + +test('Flux generation without a reference retains the text-only graph', () => { + const graph = structuredClone(template) + api.addFluxGenerationReference(graph, '') + assert.deepEqual(graph, template) +})