towsty/aigen
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Add optional reference photo conditioning to Flux generation

Browse Source
This commit is contained in:
Towsty
2026-09-06 23:05:36 -05:00
parent bb91b4fd3f
commit 53bb58bbd5
7 changed files with 66 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pages/index.vue
+9 -8
View File
@@ -182,9 +182,9 @@
</label>
</div>
<div v-if="studioMode === 'edit' || (studioMode === 'editv2' && (v2Mode === 'compose' || (v2Mode === 'iterate' && file)))" class="rounded-2xl border border-dashed border-white/15 bg-zinc-950/50 p-4">
<div v-if="fluxTextReference || studioMode === 'edit' || (studioMode === 'editv2' && (v2Mode === 'compose' || (v2Mode === 'iterate' && file)))" class="rounded-2xl border border-dashed border-white/15 bg-zinc-950/50 p-4">
<div class="mb-2 flex items-center justify-between gap-2">
<p class="text-sm font-medium text-zinc-300">{{ studioMode === 'editv2' && v2Mode === 'compose' ? 'Still B · required' : studioMode === 'editv2' ? 'Still B · optional' : 'Image 2 · optional' }}</p>
<p class="text-sm font-medium text-zinc-300">{{ fluxTextReference ? 'Reference photo · optional' : studioMode === 'editv2' && v2Mode === 'compose' ? 'Still B · required' : studioMode === 'editv2' ? 'Still B · optional' : 'Image 2 · optional' }}</p>
<button
v-if="editRefFile"
type="button"
@@ -213,11 +213,11 @@
class="max-h-40 w-full object-contain bg-black/40"
>
<span v-else class="flex min-h-28 items-center justify-center px-3 py-6 text-center text-sm text-zinc-500">
{{ studioMode === 'editv2' ? stillBHint : 'Apply this still onto image 1' }}
{{ fluxTextReference ? 'Choose a reference photo from your library or upload one' : studioMode === 'editv2' ? stillBHint : 'Apply this still onto image 1' }}
</span>
</button>
<p class="mt-2 text-xs text-zinc-500">
{{ 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<ImageV2Mode>({
set: mode => { selectedV2Mode.value = mode }
})
const v2Engine = ref<ImageV2Engine>('flux')
const fluxTextReference = computed(() => studioMode.value === 'editv2' && v2Engine.value === 'flux' && (v2Mode.value === 'generate' || (v2Mode.value === 'iterate' && !file.value)))
const v2Task = ref<ImageV2Task>('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<Blob>(`/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()
server/api/v2/generate.post.ts
+3 -2
View File
@@ -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
+12
View File
@@ -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
+2 -2
View File
@@ -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
+4 -2
View File
@@ -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
+4 -1
View File
@@ -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 || '',
tests/flux-reference.test.mjs
+32
View File
@@ -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)
})