Add optional reference photo conditioning to Flux generation

2026-09-06 23:05:36 -05:00
parent bb91b4fd3f
commit 53bb58bbd5
7 changed files with 66 additions and 15 deletions
@@ -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)
+})