From 6226c55c77bb203e97e46acf6ac0164edf8cefe1 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 26 Mar 2026 16:34:27 +0100 Subject: [PATCH 01/11] Fix(Flux2): Correct guidance_embed, add guidance support for Klein 9B Base, and fix metadata recall Klein 4B and 9B (distilled) have guidance_embeds=False, while Klein 9B Base (undistilled) has guidance_embeds=True. This commit: - Sets guidance_embed=False for Klein 4B/9B and adds Klein9BBase with True - Adds guidance parameter to Flux2DenoiseInvocation (used by Klein 9B Base) - Passes real guidance value instead of hardcoded 1.0 in flux2/denoise.py - Hides guidance slider for distilled Klein models, shows it for Klein 9B Base - Shows Flux scheduler dropdown for all Flux2 Klein models - Passes scheduler to Flux2 denoise node and saves it in metadata - Adds KleinVAEModel and KleinQwen3EncoderModel to recall parameters panel --- invokeai/app/invocations/flux2_denoise.py | 8 +++++++ invokeai/backend/flux/util.py | 17 +++++++++++++- invokeai/backend/flux2/denoise.py | 22 ++++++++++++------- .../ImageMetadataActions.tsx | 2 ++ .../util/graph/generation/buildFLUXGraph.ts | 3 +++ .../GenerationSettingsAccordion.tsx | 7 ++++-- .../frontend/web/src/services/api/schema.ts | 6 +++++ 7 files changed, 54 insertions(+), 11 deletions(-) diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index c387a72790e..268652ca090 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/invokeai/app/invocations/flux2_denoise.py @@ -101,6 +101,13 @@ class Flux2DenoiseInvocation(BaseInvocation): description="Negative conditioning tensor. Can be None if cfg_scale is 1.0.", input=Input.Connection, ) + guidance: float = InputField( + default=4.0, + ge=0, + le=20, + description="The guidance strength. Only used by undistilled models (Klein 9B Base). " + "Ignored by distilled models (Klein 4B, Klein 9B).", + ) cfg_scale: float = InputField( default=1.0, description=FieldDescriptions.cfg_scale, @@ -467,6 +474,7 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: txt_ids=txt_ids, timesteps=timesteps, step_callback=self._build_step_callback(context), + guidance=self.guidance, cfg_scale=cfg_scale_list, neg_txt=neg_txt, neg_txt_ids=neg_txt_ids, diff --git a/invokeai/backend/flux/util.py b/invokeai/backend/flux/util.py index da6590c7573..81f8caf46a1 100644 --- a/invokeai/backend/flux/util.py +++ b/invokeai/backend/flux/util.py @@ -133,11 +133,26 @@ def get_flux_ae_params() -> AutoEncoderParams: axes_dim=[16, 56, 56], theta=10_000, qkv_bias=True, - guidance_embed=True, + guidance_embed=False, ), # Flux2 Klein 9B uses Qwen3 8B text encoder with stacked embeddings from layers [9, 18, 27] # The context_in_dim is 3 * hidden_size of Qwen3 (3 * 4096 = 12288) Flux2VariantType.Klein9B: FluxParams( + in_channels=64, + vec_in_dim=4096, # Qwen3-8B hidden size (used for pooled output) + context_in_dim=12288, # 3 layers * 4096 = 12288 for Qwen3-8B + hidden_size=3072, + mlp_ratio=4.0, + num_heads=24, + depth=19, + depth_single_blocks=38, + axes_dim=[16, 56, 56], + theta=10_000, + qkv_bias=True, + guidance_embed=False, + ), + # Flux2 Klein 9B Base is the undistilled foundation model with guidance_embeds=True + Flux2VariantType.Klein9BBase: FluxParams( in_channels=64, vec_in_dim=4096, # Qwen3-8B hidden size (used for pooled output) context_in_dim=12288, # 3 layers * 4096 = 12288 for Qwen3-8B diff --git a/invokeai/backend/flux2/denoise.py b/invokeai/backend/flux2/denoise.py index 7b5bd6194e0..47a1af68023 100644 --- a/invokeai/backend/flux2/denoise.py +++ b/invokeai/backend/flux2/denoise.py @@ -26,6 +26,7 @@ def denoise( # sampling parameters timesteps: list[float], step_callback: Callable[[PipelineIntermediateState], None], + guidance: float, cfg_scale: list[float], # Negative conditioning for CFG neg_txt: torch.Tensor | None = None, @@ -45,7 +46,9 @@ def denoise( This is a simplified denoise function for FLUX.2 Klein models that uses the diffusers Flux2Transformer2DModel interface. - Note: FLUX.2 Klein has guidance_embeds=False, so no guidance parameter is used. + Distilled models (Klein 4B, Klein 9B) have guidance_embeds=False, so the guidance + value is passed but ignored by the model. Undistilled models (Klein 9B Base) have + guidance_embeds=True and use the guidance value for generation. CFG is applied externally using negative conditioning when cfg_scale != 1.0. Args: @@ -56,6 +59,8 @@ def denoise( txt_ids: Text position IDs tensor. timesteps: List of timesteps for denoising schedule (linear sigmas from 1.0 to 1/n). step_callback: Callback function for progress updates. + guidance: Guidance strength. Used by undistilled models (Klein 9B Base), + ignored by distilled models (Klein 4B, Klein 9B). cfg_scale: List of CFG scale values per step. neg_txt: Negative text embeddings for CFG (optional). neg_txt_ids: Negative text position IDs (optional). @@ -76,9 +81,10 @@ def denoise( img = torch.cat([img, img_cond_seq], dim=1) img_ids = torch.cat([img_ids, img_cond_seq_ids], dim=1) - # Klein has guidance_embeds=False, but the transformer forward() still requires a guidance tensor - # We pass a dummy value (1.0) since it won't affect the output when guidance_embeds=False - guidance = torch.full((img.shape[0],), 1.0, device=img.device, dtype=img.dtype) + # The transformer forward() requires a guidance tensor. + # For distilled models (guidance_embeds=False), this value is ignored by the model. + # For undistilled models (Klein 9B Base, guidance_embeds=True), it controls guidance strength. + guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype) # Use scheduler if provided use_scheduler = scheduler is not None @@ -121,7 +127,7 @@ def denoise( timestep=t_vec, img_ids=img_ids, txt_ids=txt_ids, - guidance=guidance, + guidance=guidance_vec, return_dict=False, ) @@ -141,7 +147,7 @@ def denoise( timestep=t_vec, img_ids=img_ids, txt_ids=neg_txt_ids if neg_txt_ids is not None else txt_ids, - guidance=guidance, + guidance=guidance_vec, return_dict=False, ) @@ -222,7 +228,7 @@ def denoise( timestep=t_vec, img_ids=img_ids, txt_ids=txt_ids, - guidance=guidance, + guidance=guidance_vec, return_dict=False, ) @@ -242,7 +248,7 @@ def denoise( timestep=t_vec, img_ids=img_ids, txt_ids=neg_txt_ids if neg_txt_ids is not None else txt_ids, - guidance=guidance, + guidance=guidance_vec, return_dict=False, ) diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageMetadataViewer/ImageMetadataActions.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageMetadataViewer/ImageMetadataActions.tsx index 8d04aad1013..8569d4862a8 100644 --- a/invokeai/frontend/web/src/features/gallery/components/ImageMetadataViewer/ImageMetadataActions.tsx +++ b/invokeai/frontend/web/src/features/gallery/components/ImageMetadataViewer/ImageMetadataActions.tsx @@ -59,6 +59,8 @@ export const ImageMetadataActions = memo((props: Props) => { + + ); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts index ba27e5dbf6e..30c7fb7c4d8 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts @@ -160,6 +160,8 @@ export const buildFLUXGraph = async (arg: GraphBuilderArg): Promise = { model: Graph.getModelMetadataField(model), steps, + scheduler: fluxScheduler, }; if (kleinVaeModel) { flux2Metadata.vae = kleinVaeModel; diff --git a/invokeai/frontend/web/src/features/settingsAccordions/components/GenerationSettingsAccordion/GenerationSettingsAccordion.tsx b/invokeai/frontend/web/src/features/settingsAccordions/components/GenerationSettingsAccordion/GenerationSettingsAccordion.tsx index ffdbc4ce778..2638141ec68 100644 --- a/invokeai/frontend/web/src/features/settingsAccordions/components/GenerationSettingsAccordion/GenerationSettingsAccordion.tsx +++ b/invokeai/frontend/web/src/features/settingsAccordions/components/GenerationSettingsAccordion/GenerationSettingsAccordion.tsx @@ -83,10 +83,13 @@ export const GenerationSettingsAccordion = memo(() => { {!isFLUX && !isFlux2 && !isSD3 && !isCogView4 && !isZImage && } - {isFLUX && } + {(isFLUX || isFlux2) && } {isZImage && } - {(isFLUX || isFlux2) && modelConfig && !isFluxFillMainModelModelConfig(modelConfig) && } + {isFLUX && modelConfig && !isFluxFillMainModelModelConfig(modelConfig) && } + {isFlux2 && modelConfig && 'variant' in modelConfig && modelConfig.variant === 'klein_9b_base' && ( + + )} {!isFLUX && !isFlux2 && } {isFLUX && } {isFLUX && fluxDypePreset === 'manual' && } diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 7ae0c77bdf9..7408aa440f3 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -8872,6 +8872,12 @@ export type components = { * @default null */ negative_text_conditioning?: components["schemas"]["FluxConditioningField"] | null; + /** + * Guidance + * @description The guidance strength. Only used by undistilled models (Klein 9B Base). Ignored by distilled models (Klein 4B, Klein 9B). + * @default 4 + */ + guidance?: number; /** * CFG Scale * @description Classifier-Free Guidance scale From edbc705a64bce3dd7c4f65d48f6b4ef5213341da Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 8 Apr 2026 01:17:06 +0200 Subject: [PATCH 02/11] test(flux2): cover Klein guidance gating, scheduler metadata, and recall dedupe Add a mock-based harness for buildFLUXGraph that locks in the FLUX.2 orchestration: guidance is written to metadata and the flux2_denoise node only for klein_9b_base, distilled variants (klein_9b, klein_4b) omit it, the FLUX scheduler is persisted into both metadata and the denoise node, and separately selected Klein VAE / Qwen3 encoder land in metadata. Add parsing tests for the metadata recall handlers: KleinVAEModel and KleinQwen3EncoderModel only fire when the current main model is FLUX.2, and the generic VAEModel handler now bails out for flux2 / z-image so the metadata viewer no longer renders duplicate VAE rows next to the dedicated Klein / Z-Image handlers. --- .../src/features/metadata/parsing.test.tsx | 131 +++++++++++ .../web/src/features/metadata/parsing.tsx | 3 + .../graph/generation/buildFLUXGraph.test.ts | 205 ++++++++++++++++++ .../util/graph/generation/buildFLUXGraph.ts | 3 + 4 files changed, 342 insertions(+) create mode 100644 invokeai/frontend/web/src/features/metadata/parsing.test.tsx create mode 100644 invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.test.ts diff --git a/invokeai/frontend/web/src/features/metadata/parsing.test.tsx b/invokeai/frontend/web/src/features/metadata/parsing.test.tsx new file mode 100644 index 00000000000..b029f04b010 --- /dev/null +++ b/invokeai/frontend/web/src/features/metadata/parsing.test.tsx @@ -0,0 +1,131 @@ +import type { AppStore } from 'app/store/store'; +import type * as paramsSliceModule from 'features/controlLayers/store/paramsSlice'; +import { ImageMetadataHandlers } from 'features/metadata/parsing'; +import type * as modelsApiModule from 'services/api/endpoints/models'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +// --------------------------------------------------------------------------- +// Module mocks +// +// We are testing only the *gating* logic of the model-related metadata +// handlers (`VAEModel`, `KleinVAEModel`, `KleinQwen3EncoderModel`). The actual +// model lookup goes through `parseModelIdentifier`, which dispatches RTK +// Query thunks. We stub the models endpoint so that any lookup resolves to a +// canned model identifier — the parse step then succeeds and the assertions +// inside each handler become observable. +// --------------------------------------------------------------------------- + +let currentBase: string | null = 'flux2'; + +vi.mock('features/controlLayers/store/paramsSlice', async (importOriginal) => { + const mod = await importOriginal(); + return { ...mod, selectBase: () => currentBase }; +}); + +const fakeModel = (type: 'vae' | 'qwen3_encoder', base: string) => ({ + key: `${type}-key`, + hash: 'hash', + name: `Some ${type}`, + base, + type, +}); + +let nextResolved: ReturnType = fakeModel('vae', 'flux2'); + +vi.mock('services/api/endpoints/models', async (importOriginal) => { + const mod = await importOriginal(); + return { + ...mod, + modelsApi: { + ...mod.modelsApi, + endpoints: { + ...mod.modelsApi.endpoints, + getModelConfig: { initiate: (key: string) => ({ type: 'rtkq/initiate', key }) }, + }, + }, + }; +}); + +const makeStore = (): AppStore => + ({ + dispatch: vi.fn(() => ({ + unwrap: () => Promise.resolve(nextResolved), + })), + getState: () => ({}), + }) as unknown as AppStore; + +beforeEach(() => { + currentBase = 'flux2'; + nextResolved = fakeModel('vae', 'flux2'); +}); + +describe('ImageMetadataHandlers — Klein recall gating', () => { + describe('KleinVAEModel', () => { + it('parses metadata.vae when the current main model is FLUX.2 Klein', async () => { + currentBase = 'flux2'; + nextResolved = fakeModel('vae', 'flux2'); + const store = makeStore(); + + const parsed = await ImageMetadataHandlers.KleinVAEModel.parse({ vae: nextResolved }, store); + + expect(parsed.key).toBe('vae-key'); + expect(parsed.type).toBe('vae'); + }); + + it('rejects parsing when the current main model is not FLUX.2 Klein', async () => { + currentBase = 'sdxl'; + nextResolved = fakeModel('vae', 'flux2'); + const store = makeStore(); + + await expect(ImageMetadataHandlers.KleinVAEModel.parse({ vae: nextResolved }, store)).rejects.toThrow(); + }); + }); + + describe('KleinQwen3EncoderModel', () => { + it('parses metadata.qwen3_encoder when the current main model is FLUX.2 Klein', async () => { + currentBase = 'flux2'; + nextResolved = fakeModel('qwen3_encoder', 'flux2'); + const store = makeStore(); + + const parsed = await ImageMetadataHandlers.KleinQwen3EncoderModel.parse( + { qwen3_encoder: nextResolved }, + store + ); + + expect(parsed.key).toBe('qwen3_encoder-key'); + expect(parsed.type).toBe('qwen3_encoder'); + }); + + it('rejects parsing when the current main model is not FLUX.2 Klein', async () => { + currentBase = 'sdxl'; + nextResolved = fakeModel('qwen3_encoder', 'flux2'); + const store = makeStore(); + + await expect( + ImageMetadataHandlers.KleinQwen3EncoderModel.parse({ qwen3_encoder: nextResolved }, store) + ).rejects.toThrow(); + }); + }); + + describe('VAEModel (generic)', () => { + // The generic VAEModel handler must NOT also fire for FLUX.2 / Z-Image + // images, otherwise the metadata viewer renders duplicate VAE rows next + // to the dedicated KleinVAEModel / ZImageVAEModel handlers. + it.each(['flux2', 'z-image'])('rejects parsing when current base is %s', async (base) => { + currentBase = base; + nextResolved = fakeModel('vae', base); + const store = makeStore(); + + await expect(ImageMetadataHandlers.VAEModel.parse({ vae: nextResolved }, store)).rejects.toThrow(); + }); + + it('parses successfully for non-Klein, non-Z-Image bases', async () => { + currentBase = 'sdxl'; + nextResolved = fakeModel('vae', 'sdxl'); + const store = makeStore(); + + const parsed = await ImageMetadataHandlers.VAEModel.parse({ vae: nextResolved }, store); + expect(parsed.key).toBe('vae-key'); + }); + }); +}); diff --git a/invokeai/frontend/web/src/features/metadata/parsing.tsx b/invokeai/frontend/web/src/features/metadata/parsing.tsx index 7d1d511a3c2..7e415af611a 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.tsx +++ b/invokeai/frontend/web/src/features/metadata/parsing.tsx @@ -844,6 +844,9 @@ const VAEModel: SingleMetadataHandler = { const parsed = await parseModelIdentifier(raw, store, 'vae'); assert(parsed.type === 'vae'); assert(isCompatibleWithMainModel(parsed, store)); + // Z-Image and FLUX.2 Klein have dedicated VAE handlers; avoid rendering a duplicate row. + const base = selectBase(store.getState()); + assert(base !== 'z-image' && base !== 'flux2', 'VAEModel handler does not apply to Z-Image or FLUX.2 Klein'); return Promise.resolve(parsed); }, recall: (value, store) => { diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.test.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.test.ts new file mode 100644 index 00000000000..a8cdcadbf72 --- /dev/null +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.test.ts @@ -0,0 +1,205 @@ +import { buildFLUXGraph } from 'features/nodes/util/graph/generation/buildFLUXGraph'; +import type { Graph } from 'features/nodes/util/graph/generation/Graph'; +import type { GraphBuilderArg } from 'features/nodes/util/graph/types'; +import type { Invocation } from 'services/api/types'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +// --------------------------------------------------------------------------- +// Module mocks +// +// `buildFLUXGraph` pulls in a large slice of the app: redux selectors, every +// `add*` helper, validators, the canvas manager, etc. The function itself only +// orchestrates these; the unit under test here is the orchestration logic +// (variant-gated guidance, scheduler propagation, metadata persistence). So we +// stub out every collaborator and assert against the resulting `Graph` object. +// --------------------------------------------------------------------------- + +const mockState = { + // buildFLUXGraph reads `state.system.shouldUse{NSFWChecker,Watermarker}` directly, + // every other access is funneled through the mocked selectors below. + system: { shouldUseNSFWChecker: false, shouldUseWatermarker: false }, +} as unknown as Parameters[0]['state']; + +const mockParams = { + guidance: 3.5, + steps: 28, + fluxScheduler: 'euler' as const, + fluxDypePreset: 'off' as const, + fluxDypeScale: 1, + fluxDypeExponent: 1, + fluxVAE: null, + t5EncoderModel: null, + clipEmbedModel: null, +}; + +let currentModel: { key: string; hash: string; name: string; base: string; type: string; variant?: string } | null; +let currentKleinVae: { key: string; hash: string; name: string; base: string; type: string } | null; +let currentKleinQwen3: { key: string; hash: string; name: string; base: string; type: string } | null; + +vi.mock('features/controlLayers/store/paramsSlice', () => ({ + selectMainModelConfig: () => currentModel, + selectParamsSlice: () => mockParams, + selectKleinVaeModel: () => currentKleinVae, + selectKleinQwen3EncoderModel: () => currentKleinQwen3, +})); + +vi.mock('features/controlLayers/store/refImagesSlice', () => ({ + selectRefImagesSlice: () => ({ entities: [] }), +})); + +vi.mock('features/controlLayers/store/selectors', () => ({ + selectCanvasSlice: () => ({ + bbox: { rect: { x: 0, y: 0, width: 1024, height: 1024 } }, + controlLayers: { entities: [] }, + regionalGuidance: { entities: [] }, + }), + selectCanvasMetadata: () => ({}), +})); + +vi.mock('features/controlLayers/store/types', () => ({ + isFlux2ReferenceImageConfig: () => false, + isFluxKontextReferenceImageConfig: () => false, +})); + +vi.mock('features/controlLayers/store/validators', () => ({ + getGlobalReferenceImageWarnings: () => [], +})); + +vi.mock('features/ui/store/uiSelectors', () => ({ + selectActiveTab: () => 'generate', +})); + +vi.mock('features/nodes/util/graph/graphBuilderUtils', () => ({ + selectCanvasOutputFields: () => ({}), +})); + +// Helper add* functions: each test cares only that the FLUX.2 orchestration +// path produces the right metadata + denoise inputs. The actual node graph +// produced by these helpers is irrelevant here. +vi.mock('features/nodes/util/graph/generation/addTextToImage', () => ({ + addTextToImage: ({ l2i }: { l2i: Invocation<'flux2_vae_decode'> }) => l2i, +})); +vi.mock('features/nodes/util/graph/generation/addImageToImage', () => ({ + addImageToImage: vi.fn(), +})); +vi.mock('features/nodes/util/graph/generation/addInpaint', () => ({ addInpaint: vi.fn() })); +vi.mock('features/nodes/util/graph/generation/addOutpaint', () => ({ addOutpaint: vi.fn() })); +vi.mock('features/nodes/util/graph/generation/addNSFWChecker', () => ({ addNSFWChecker: vi.fn() })); +vi.mock('features/nodes/util/graph/generation/addWatermarker', () => ({ addWatermarker: vi.fn() })); +vi.mock('features/nodes/util/graph/generation/addRegions', () => ({ addRegions: vi.fn(() => []) })); +vi.mock('features/nodes/util/graph/generation/addFLUXLoRAs', () => ({ addFLUXLoRAs: vi.fn() })); +vi.mock('features/nodes/util/graph/generation/addFlux2KleinLoRAs', () => ({ addFlux2KleinLoRAs: vi.fn() })); +vi.mock('features/nodes/util/graph/generation/addFLUXFill', () => ({ addFLUXFill: vi.fn() })); +vi.mock('features/nodes/util/graph/generation/addFLUXRedux', () => ({ + addFLUXReduxes: () => ({ addedFLUXReduxes: 0 }), +})); +vi.mock('features/nodes/util/graph/generation/addControlAdapters', () => ({ + addControlNets: vi.fn(() => Promise.resolve({ addedControlNets: 0 })), + addControlLoRA: vi.fn(), +})); +vi.mock('features/nodes/util/graph/generation/addIPAdapters', () => ({ + addIPAdapters: () => ({ addedIPAdapters: 0 }), +})); + +// --------------------------------------------------------------------------- +// Test harness +// --------------------------------------------------------------------------- + +const makeFlux2Model = (variant: string) => ({ + key: `flux2-${variant}`, + hash: 'hash', + name: `FLUX.2 Klein ${variant}`, + base: 'flux2', + type: 'main', + variant, +}); + +const buildArg = (): GraphBuilderArg => + ({ + generationMode: 'txt2img', + state: mockState, + manager: null, + }) as unknown as GraphBuilderArg; + +const findFlux2Denoise = (g: Graph): Invocation<'flux2_denoise'> | undefined => { + // The Graph object stores nodes on `_graph.nodes` keyed by id. + const nodes = (g as unknown as { _graph: { nodes: Record } })._graph.nodes; + return Object.values(nodes).find((n) => n.type === 'flux2_denoise') as Invocation<'flux2_denoise'> | undefined; +}; + +const getMetadata = (g: Graph): Record => + (g as unknown as { getMetadataNode: () => Record }).getMetadataNode(); + +beforeEach(() => { + currentModel = null; + currentKleinVae = null; + currentKleinQwen3 = null; +}); + +describe('buildFLUXGraph (FLUX.2 Klein)', () => { + describe('guidance gating', () => { + it('writes guidance into metadata and the denoise node for klein_9b_base', async () => { + currentModel = makeFlux2Model('klein_9b_base'); + + const { g } = await buildFLUXGraph(buildArg()); + + const metadata = getMetadata(g); + expect(metadata.guidance).toBe(mockParams.guidance); + + const denoise = findFlux2Denoise(g); + expect(denoise).toBeDefined(); + expect(denoise?.guidance).toBe(mockParams.guidance); + }); + + it.each(['klein_9b', 'klein_4b'])( + 'omits guidance from metadata and denoise for distilled variant %s', + async (variant) => { + currentModel = makeFlux2Model(variant); + + const { g } = await buildFLUXGraph(buildArg()); + + const metadata = getMetadata(g); + expect(metadata.guidance).toBeUndefined(); + + const denoise = findFlux2Denoise(g); + expect(denoise).toBeDefined(); + expect(denoise?.guidance).toBeUndefined(); + } + ); + }); + + describe('scheduler persistence', () => { + it('writes the FLUX scheduler into metadata and the denoise node for FLUX.2', async () => { + currentModel = makeFlux2Model('klein_9b_base'); + + const { g } = await buildFLUXGraph(buildArg()); + + expect(getMetadata(g).scheduler).toBe(mockParams.fluxScheduler); + expect(findFlux2Denoise(g)?.scheduler).toBe(mockParams.fluxScheduler); + }); + }); + + describe('Klein VAE / Qwen3 metadata', () => { + it('persists separately selected Klein VAE and Qwen3 encoder into metadata', async () => { + currentModel = makeFlux2Model('klein_9b_base'); + currentKleinVae = { key: 'vae-1', hash: 'h', name: 'Klein VAE', base: 'flux2', type: 'vae' }; + currentKleinQwen3 = { key: 'q3-1', hash: 'h', name: 'Qwen3', base: 'flux2', type: 'qwen3_encoder' }; + + const { g } = await buildFLUXGraph(buildArg()); + + const metadata = getMetadata(g); + expect(metadata.vae).toEqual(currentKleinVae); + expect(metadata.qwen3_encoder).toEqual(currentKleinQwen3); + }); + + it('omits vae / qwen3_encoder when none are selected', async () => { + currentModel = makeFlux2Model('klein_9b_base'); + + const { g } = await buildFLUXGraph(buildArg()); + + const metadata = getMetadata(g); + expect(metadata.vae).toBeUndefined(); + expect(metadata.qwen3_encoder).toBeUndefined(); + }); + }); +}); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts index 30c7fb7c4d8..d03924d4979 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts @@ -232,6 +232,9 @@ export const buildFLUXGraph = async (arg: GraphBuilderArg): Promise Date: Wed, 8 Apr 2026 01:26:01 +0200 Subject: [PATCH 03/11] Chore pnpm fix --- invokeai/frontend/web/src/features/metadata/parsing.test.tsx | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/invokeai/frontend/web/src/features/metadata/parsing.test.tsx b/invokeai/frontend/web/src/features/metadata/parsing.test.tsx index b029f04b010..01e33ee2dbb 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.test.tsx +++ b/invokeai/frontend/web/src/features/metadata/parsing.test.tsx @@ -87,10 +87,7 @@ describe('ImageMetadataHandlers — Klein recall gating', () => { nextResolved = fakeModel('qwen3_encoder', 'flux2'); const store = makeStore(); - const parsed = await ImageMetadataHandlers.KleinQwen3EncoderModel.parse( - { qwen3_encoder: nextResolved }, - store - ); + const parsed = await ImageMetadataHandlers.KleinQwen3EncoderModel.parse({ qwen3_encoder: nextResolved }, store); expect(parsed.key).toBe('qwen3_encoder-key'); expect(parsed.type).toBe('qwen3_encoder'); From f5866d081ab56e2247d4748e31d659238075ce1d Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 9 Apr 2026 18:43:29 +0200 Subject: [PATCH 04/11] Update version to 1.5.0 in flux2_denoise.py --- invokeai/app/invocations/flux2_denoise.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index 268652ca090..0a5d854f534 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/invokeai/app/invocations/flux2_denoise.py @@ -54,7 +54,7 @@ title="FLUX2 Denoise", tags=["image", "flux", "flux2", "klein", "denoise"], category="image", - version="1.4.0", + version="1.5.0", classification=Classification.Prototype, ) class Flux2DenoiseInvocation(BaseInvocation): From 4d9e2aac7082712484f556753047d4bc9167c30c Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 9 Apr 2026 18:46:52 +0200 Subject: [PATCH 05/11] Update condition for rendering ParamFluxScheduler --- .../GenerationSettingsAccordion/GenerationSettingsAccordion.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/frontend/web/src/features/settingsAccordions/components/GenerationSettingsAccordion/GenerationSettingsAccordion.tsx b/invokeai/frontend/web/src/features/settingsAccordions/components/GenerationSettingsAccordion/GenerationSettingsAccordion.tsx index af24828e1fd..0695276febf 100644 --- a/invokeai/frontend/web/src/features/settingsAccordions/components/GenerationSettingsAccordion/GenerationSettingsAccordion.tsx +++ b/invokeai/frontend/web/src/features/settingsAccordions/components/GenerationSettingsAccordion/GenerationSettingsAccordion.tsx @@ -86,7 +86,7 @@ export const GenerationSettingsAccordion = memo(() => { {!isFLUX && !isFlux2 && !isSD3 && !isCogView4 && !isZImage && !isAnima && } - {isFLUX && } + {(isFLUX || isFlux2) && } {isZImage && } {isAnima && } From 6cf6368608854cb435e5fd16fee80c412210c1c4 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Tue, 21 Apr 2026 23:03:04 +0200 Subject: [PATCH 06/11] feat(flux2): add Klein4BBase variant for FLUX.2 Klein Base 4B models Recognize FLUX.2-klein-base-4B on import via filename heuristic. The variant shares Klein4B's architecture (Qwen3-4B encoder, context_in_dim=7680) and reports guidance_embeds=False in its HF config, consistent with Klein 9B Base. UI behavior stays identical to distilled Klein4B until CFG support is wired up in a follow-up. --- .../invocations/flux2_klein_model_loader.py | 4 +-- invokeai/backend/flux/util.py | 17 +++++++++++++ .../backend/model_manager/configs/main.py | 25 ++++++++++++------- invokeai/backend/model_manager/taxonomy.py | 5 +++- invokeai/frontend/web/openapi.json | 2 +- .../web/src/features/modelManagerV2/models.ts | 1 + .../web/src/features/nodes/types/common.ts | 2 +- .../Advanced/ParamFlux2KleinModelSelect.tsx | 1 + .../frontend/web/src/services/api/schema.ts | 2 +- 9 files changed, 44 insertions(+), 15 deletions(-) diff --git a/invokeai/app/invocations/flux2_klein_model_loader.py b/invokeai/app/invocations/flux2_klein_model_loader.py index f39e7688f3e..2091fd380d7 100644 --- a/invokeai/app/invocations/flux2_klein_model_loader.py +++ b/invokeai/app/invocations/flux2_klein_model_loader.py @@ -207,9 +207,9 @@ def _validate_qwen3_encoder_variant(self, context: InvocationContext, main_confi flux2_variant = main_config.variant # Validate the variants match - # Klein4B requires Qwen3_4B, Klein9B/Klein9BBase requires Qwen3_8B + # Klein4B/Klein4BBase requires Qwen3_4B, Klein9B/Klein9BBase requires Qwen3_8B expected_qwen3_variant = None - if flux2_variant == Flux2VariantType.Klein4B: + if flux2_variant in (Flux2VariantType.Klein4B, Flux2VariantType.Klein4BBase): expected_qwen3_variant = Qwen3VariantType.Qwen3_4B elif flux2_variant in (Flux2VariantType.Klein9B, Flux2VariantType.Klein9BBase): expected_qwen3_variant = Qwen3VariantType.Qwen3_8B diff --git a/invokeai/backend/flux/util.py b/invokeai/backend/flux/util.py index 81f8caf46a1..1d696ccdd64 100644 --- a/invokeai/backend/flux/util.py +++ b/invokeai/backend/flux/util.py @@ -135,6 +135,23 @@ def get_flux_ae_params() -> AutoEncoderParams: qkv_bias=True, guidance_embed=False, ), + # Flux2 Klein 4B Base is the undistilled foundation model. It shares the same + # architecture as Klein 4B (distilled) and reports guidance_embeds=False in its + # HF transformer config - classical CFG (external negative pass) is the guidance mechanism. + Flux2VariantType.Klein4BBase: FluxParams( + in_channels=64, + vec_in_dim=2560, # Qwen3-4B hidden size (used for pooled output) + context_in_dim=7680, # 3 layers * 2560 = 7680 for Qwen3-4B + hidden_size=3072, + mlp_ratio=4.0, + num_heads=24, + depth=19, + depth_single_blocks=38, + axes_dim=[16, 56, 56], + theta=10_000, + qkv_bias=True, + guidance_embed=False, + ), # Flux2 Klein 9B uses Qwen3 8B text encoder with stacked embeddings from layers [9, 18, 27] # The context_in_dim is 3 * hidden_size of Qwen3 (3 * 4096 = 12288) Flux2VariantType.Klein9B: FluxParams( diff --git a/invokeai/backend/model_manager/configs/main.py b/invokeai/backend/model_manager/configs/main.py index 1be349f3942..35f5e43823c 100644 --- a/invokeai/backend/model_manager/configs/main.py +++ b/invokeai/backend/model_manager/configs/main.py @@ -81,8 +81,8 @@ def from_base( return cls(steps=35, cfg_scale=4.5, width=1024, height=1024) case BaseModelType.Flux2: # Different defaults based on variant - if variant == Flux2VariantType.Klein9BBase: - # Undistilled base model needs more steps + if variant in (Flux2VariantType.Klein4BBase, Flux2VariantType.Klein9BBase): + # Undistilled base models need more steps return cls(steps=28, cfg_scale=1.0, width=1024, height=1024) else: # Distilled models (Klein 4B, Klein 9B) use fewer steps @@ -386,6 +386,7 @@ def _get_flux2_variant(state_dict: dict[str | int, Any]) -> Flux2VariantType | N # Default to Klein9B - callers use filename heuristics to detect Klein9BBase return Flux2VariantType.Klein9B elif context_in_dim == KLEIN_4B_CONTEXT_DIM: + # Default to Klein4B - callers use filename heuristics to detect Klein4BBase return Flux2VariantType.Klein4B elif context_in_dim > 4096: # Unknown FLUX.2 variant, default to 4B @@ -570,10 +571,12 @@ def _get_variant_or_raise(cls, mod: ModelOnDisk) -> Flux2VariantType: if variant is None: raise NotAMatchError("unable to determine FLUX.2 model variant from state dict") - # Klein 9B Base and Klein 9B have identical architectures. - # Use filename heuristic to detect the Base (undistilled) variant. + # Base (undistilled) and distilled variants share identical architectures. + # Use filename heuristic to detect the Base variant. if variant == Flux2VariantType.Klein9B and _filename_suggests_base(mod.name): return Flux2VariantType.Klein9BBase + if variant == Flux2VariantType.Klein4B and _filename_suggests_base(mod.name): + return Flux2VariantType.Klein4BBase return variant @@ -742,10 +745,12 @@ def _get_variant_or_raise(cls, mod: ModelOnDisk) -> Flux2VariantType: if variant is None: raise NotAMatchError("unable to determine FLUX.2 model variant from state dict") - # Klein 9B Base and Klein 9B have identical architectures. - # Use filename heuristic to detect the Base (undistilled) variant. + # Base (undistilled) and distilled variants share identical architectures. + # Use filename heuristic to detect the Base variant. if variant == Flux2VariantType.Klein9B and _filename_suggests_base(mod.name): return Flux2VariantType.Klein9BBase + if variant == Flux2VariantType.Klein4B and _filename_suggests_base(mod.name): + return Flux2VariantType.Klein4BBase return variant @@ -853,11 +858,11 @@ def _get_variant_or_raise(cls, mod: ModelOnDisk) -> Flux2VariantType: """Determine the FLUX.2 variant from the transformer config. FLUX.2 Klein uses Qwen3 text encoder with larger joint_attention_dim: - - Klein 4B: joint_attention_dim = 7680 (3×Qwen3-4B hidden size) + - Klein 4B/4B Base: joint_attention_dim = 7680 (3×Qwen3-4B hidden size) - Klein 9B/9B Base: joint_attention_dim = 12288 (3×Qwen3-8B hidden size) - Klein 9B (distilled) and Klein 9B Base (undistilled) have identical architectures - and both have guidance_embeds=False. We use a filename heuristic to detect Base models. + Distilled and Base variants share identical architectures and both have + guidance_embeds=False. We use a filename heuristic to detect Base models. """ KLEIN_4B_CONTEXT_DIM = 7680 # 3 × 2560 KLEIN_9B_CONTEXT_DIM = 12288 # 3 × 4096 @@ -872,6 +877,8 @@ def _get_variant_or_raise(cls, mod: ModelOnDisk) -> Flux2VariantType: return Flux2VariantType.Klein9BBase return Flux2VariantType.Klein9B elif joint_attention_dim == KLEIN_4B_CONTEXT_DIM: + if _filename_suggests_base(mod.name): + return Flux2VariantType.Klein4BBase return Flux2VariantType.Klein4B elif joint_attention_dim > 4096: # Unknown FLUX.2 variant, default to 4B diff --git a/invokeai/backend/model_manager/taxonomy.py b/invokeai/backend/model_manager/taxonomy.py index b2b55ebd3fc..776ccadc6c2 100644 --- a/invokeai/backend/model_manager/taxonomy.py +++ b/invokeai/backend/model_manager/taxonomy.py @@ -131,7 +131,10 @@ class Flux2VariantType(str, Enum): """FLUX.2 model variants.""" Klein4B = "klein_4b" - """Flux2 Klein 4B variant using Qwen3 4B text encoder.""" + """Flux2 Klein 4B variant using Qwen3 4B text encoder (distilled).""" + + Klein4BBase = "klein_4b_base" + """Flux2 Klein 4B Base variant - undistilled foundation model using Qwen3 4B text encoder.""" Klein9B = "klein_9b" """Flux2 Klein 9B variant using Qwen3 8B text encoder (distilled).""" diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 19e5a3a68e9..6bb5aa662f0 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -19924,7 +19924,7 @@ }, "Flux2VariantType": { "type": "string", - "enum": ["klein_4b", "klein_9b", "klein_9b_base"], + "enum": ["klein_4b", "klein_4b_base", "klein_9b", "klein_9b_base"], "title": "Flux2VariantType", "description": "FLUX.2 model variants." }, diff --git a/invokeai/frontend/web/src/features/modelManagerV2/models.ts b/invokeai/frontend/web/src/features/modelManagerV2/models.ts index 9cc4ed24d9b..3fcfb3f21c5 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/models.ts +++ b/invokeai/frontend/web/src/features/modelManagerV2/models.ts @@ -228,6 +228,7 @@ export const MODEL_VARIANT_TO_LONG_NAME: Record = { dev_fill: 'FLUX Dev - Fill', schnell: 'FLUX Schnell', klein_4b: 'FLUX.2 Klein 4B', + klein_4b_base: 'FLUX.2 Klein 4B Base', klein_9b: 'FLUX.2 Klein 9B', klein_9b_base: 'FLUX.2 Klein 9B Base', turbo: 'Z-Image Turbo', diff --git a/invokeai/frontend/web/src/features/nodes/types/common.ts b/invokeai/frontend/web/src/features/nodes/types/common.ts index 75c3415cefb..8cdc9dc01c9 100644 --- a/invokeai/frontend/web/src/features/nodes/types/common.ts +++ b/invokeai/frontend/web/src/features/nodes/types/common.ts @@ -157,7 +157,7 @@ export const zSubModelType = z.enum([ export const zClipVariantType = z.enum(['large', 'gigantic']); export const zModelVariantType = z.enum(['normal', 'inpaint', 'depth']); export const zFluxVariantType = z.enum(['dev', 'dev_fill', 'schnell']); -export const zFlux2VariantType = z.enum(['klein_4b', 'klein_9b', 'klein_9b_base']); +export const zFlux2VariantType = z.enum(['klein_4b', 'klein_4b_base', 'klein_9b', 'klein_9b_base']); export const zZImageVariantType = z.enum(['turbo', 'zbase']); const zQwenImageVariantType = z.enum(['generate', 'edit']); export const zQwen3VariantType = z.enum(['qwen3_4b', 'qwen3_8b', 'qwen3_06b']); diff --git a/invokeai/frontend/web/src/features/parameters/components/Advanced/ParamFlux2KleinModelSelect.tsx b/invokeai/frontend/web/src/features/parameters/components/Advanced/ParamFlux2KleinModelSelect.tsx index 1d652a16458..6c0669acacc 100644 --- a/invokeai/frontend/web/src/features/parameters/components/Advanced/ParamFlux2KleinModelSelect.tsx +++ b/invokeai/frontend/web/src/features/parameters/components/Advanced/ParamFlux2KleinModelSelect.tsx @@ -64,6 +64,7 @@ ParamFlux2KleinVaeModelSelect.displayName = 'ParamFlux2KleinVaeModelSelect'; */ const KLEIN_TO_QWEN3_VARIANT_MAP: Record = { klein_4b: 'qwen3_4b', + klein_4b_base: 'qwen3_4b', klein_9b: 'qwen3_8b', klein_9b_base: 'qwen3_8b', }; diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 553a269d16b..c8a4ac969f0 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -10236,7 +10236,7 @@ export type components = { * @description FLUX.2 model variants. * @enum {string} */ - Flux2VariantType: "klein_4b" | "klein_9b" | "klein_9b_base"; + Flux2VariantType: "klein_4b" | "klein_4b_base" | "klein_9b" | "klein_9b_base"; /** * FluxConditioningCollectionOutput * @description Base class for nodes that output a collection of conditioning tensors From 7de35e3d87149780ecd7a02b3fec76bfa5917222 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 22 Apr 2026 00:18:48 +0200 Subject: [PATCH 07/11] Change Wrong Comment --- invokeai/backend/model_manager/configs/main.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/invokeai/backend/model_manager/configs/main.py b/invokeai/backend/model_manager/configs/main.py index 16620df4805..da5bc5eed36 100644 --- a/invokeai/backend/model_manager/configs/main.py +++ b/invokeai/backend/model_manager/configs/main.py @@ -864,8 +864,7 @@ def _get_variant_or_raise(cls, mod: ModelOnDisk) -> Flux2VariantType: - Klein 4B/4B Base: joint_attention_dim = 7680 (3×Qwen3-4B hidden size) - Klein 9B/9B Base: joint_attention_dim = 12288 (3×Qwen3-8B hidden size) - Distilled and Base variants share identical architectures and both have - guidance_embeds=False. We use a filename heuristic to detect Base models. + Distilled and Base variants share identical architectures. We use a filename heuristic to detect Base models. """ KLEIN_4B_CONTEXT_DIM = 7680 # 3 × 2560 KLEIN_9B_CONTEXT_DIM = 12288 # 3 × 4096 From f6ecc1a3c3f7d93d095f6483c9b481ed3155e20f Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 23 Apr 2026 01:14:39 +0200 Subject: [PATCH 08/11] refactor(flux2): remove inert guidance UI/metadata for FLUX.2 Klein All current FLUX.2 Klein variants (4B, 4B Base, 9B, 9B Base) report guidance_embeds=false in their HF transformer config (or have zeroed projection weights), so the guidance scalar has no effect on output. The linear UI previously exposed a guidance slider for klein_9b_base and wrote the value into metadata, which misled users into thinking it was steering generation. --- invokeai/app/invocations/flux2_denoise.py | 5 +++-- invokeai/backend/flux/util.py | 6 ++++-- invokeai/backend/flux2/denoise.py | 17 +++++++++-------- .../util/graph/generation/buildFLUXGraph.ts | 4 ---- .../GenerationSettingsAccordion.tsx | 5 ----- 5 files changed, 16 insertions(+), 21 deletions(-) diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index 0a5d854f534..d4239e41420 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/invokeai/app/invocations/flux2_denoise.py @@ -105,8 +105,9 @@ class Flux2DenoiseInvocation(BaseInvocation): default=4.0, ge=0, le=20, - description="The guidance strength. Only used by undistilled models (Klein 9B Base). " - "Ignored by distilled models (Klein 4B, Klein 9B).", + description="Guidance strength for distilled guidance-embedding models. " + "Inert for all current FLUX.2 Klein variants (their guidance_embeds weights are absent/zero); " + "kept for node-graph compatibility and future guidance-embedded models.", ) cfg_scale: float = InputField( default=1.0, diff --git a/invokeai/backend/flux/util.py b/invokeai/backend/flux/util.py index 1d696ccdd64..81b10a913ac 100644 --- a/invokeai/backend/flux/util.py +++ b/invokeai/backend/flux/util.py @@ -168,7 +168,9 @@ def get_flux_ae_params() -> AutoEncoderParams: qkv_bias=True, guidance_embed=False, ), - # Flux2 Klein 9B Base is the undistilled foundation model with guidance_embeds=True + # Flux2 Klein 9B Base is the undistilled foundation model. It shares the same + # architecture as Klein 9B (distilled) and reports guidance_embeds=False in its + # HF transformer config - the guidance scalar is inert for all Klein variants. Flux2VariantType.Klein9BBase: FluxParams( in_channels=64, vec_in_dim=4096, # Qwen3-8B hidden size (used for pooled output) @@ -181,7 +183,7 @@ def get_flux_ae_params() -> AutoEncoderParams: axes_dim=[16, 56, 56], theta=10_000, qkv_bias=True, - guidance_embed=True, + guidance_embed=False, ), } diff --git a/invokeai/backend/flux2/denoise.py b/invokeai/backend/flux2/denoise.py index 47a1af68023..b4438094f7b 100644 --- a/invokeai/backend/flux2/denoise.py +++ b/invokeai/backend/flux2/denoise.py @@ -46,9 +46,10 @@ def denoise( This is a simplified denoise function for FLUX.2 Klein models that uses the diffusers Flux2Transformer2DModel interface. - Distilled models (Klein 4B, Klein 9B) have guidance_embeds=False, so the guidance - value is passed but ignored by the model. Undistilled models (Klein 9B Base) have - guidance_embeds=True and use the guidance value for generation. + All current FLUX.2 Klein variants (4B, 4B Base, 9B, 9B Base) have guidance_embeds=False + in their HF transformer config (or absent/zeroed projection weights), so the guidance + value is passed but effectively ignored by the model. The argument is retained for + node-graph compatibility and future variants that may ship trained guidance projections. CFG is applied externally using negative conditioning when cfg_scale != 1.0. Args: @@ -59,8 +60,8 @@ def denoise( txt_ids: Text position IDs tensor. timesteps: List of timesteps for denoising schedule (linear sigmas from 1.0 to 1/n). step_callback: Callback function for progress updates. - guidance: Guidance strength. Used by undistilled models (Klein 9B Base), - ignored by distilled models (Klein 4B, Klein 9B). + guidance: Guidance strength. Inert for all current FLUX.2 Klein variants + (their guidance_embeds projection weights are absent/zero). cfg_scale: List of CFG scale values per step. neg_txt: Negative text embeddings for CFG (optional). neg_txt_ids: Negative text position IDs (optional). @@ -81,9 +82,9 @@ def denoise( img = torch.cat([img, img_cond_seq], dim=1) img_ids = torch.cat([img_ids, img_cond_seq_ids], dim=1) - # The transformer forward() requires a guidance tensor. - # For distilled models (guidance_embeds=False), this value is ignored by the model. - # For undistilled models (Klein 9B Base, guidance_embeds=True), it controls guidance strength. + # The transformer forward() requires a guidance tensor even when guidance_embeds=False, + # because the Flux2TimestepGuidanceEmbeddings forward signature takes it unconditionally. + # All current Klein variants have guidance_embeds=False, so the value is ignored internally. guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype) # Use scheduler if provided diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts index da64f9a31b8..dafcd9310ec 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts @@ -181,7 +181,6 @@ export const buildFLUXGraph = async (arg: GraphBuilderArg): Promise { {!isExternal && isFLUX && modelConfig && !isFluxFillMainModelModelConfig(modelConfig) && ( )} - {!isExternal && - isFlux2 && - modelConfig && - 'variant' in modelConfig && - modelConfig.variant === 'klein_9b_base' && } {!isExternal && !isFLUX && !isFlux2 && } {!isExternal && isZImage && } {!isExternal && isQwenImage && } From af0bd5041c6aeba122e9f199315c7687259c7560 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 23 Apr 2026 01:15:13 +0200 Subject: [PATCH 09/11] Chore typegen --- invokeai/frontend/web/src/services/api/schema.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index ab374d9474b..d793f2cc6ef 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -9802,7 +9802,7 @@ export type components = { negative_text_conditioning?: components["schemas"]["FluxConditioningField"] | null; /** * Guidance - * @description The guidance strength. Only used by undistilled models (Klein 9B Base). Ignored by distilled models (Klein 4B, Klein 9B). + * @description Guidance strength for distilled guidance-embedding models. Inert for all current FLUX.2 Klein variants (their guidance_embeds weights are absent/zero); kept for node-graph compatibility and future guidance-embedded models. * @default 4 */ guidance?: number; From 3f742d35dd7c6e4dfbf65421ab12309671dda3f9 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 23 Apr 2026 05:56:47 +0200 Subject: [PATCH 10/11] fix test --- .../graph/generation/buildFLUXGraph.test.ts | 20 +++++-------------- 1 file changed, 5 insertions(+), 15 deletions(-) diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.test.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.test.ts index 91e76a925bb..5b9f3d0a468 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.test.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.test.ts @@ -246,21 +246,11 @@ afterEach(resetState); describe('buildFLUXGraph (FLUX.2 Klein)', () => { describe('guidance gating', () => { - it('writes guidance into metadata and the denoise node for klein_9b_base', async () => { - currentModel = makeFlux2Model('klein_9b_base'); - - const { g } = await buildFLUXGraph(buildGraphArg()); - - const metadata = getMetadata(g); - expect(metadata.guidance).toBe(mockParams.guidance); - - const denoise = findFlux2Denoise(g); - expect(denoise).toBeDefined(); - expect(denoise?.guidance).toBe(mockParams.guidance); - }); - - it.each(['klein_9b', 'klein_4b'])( - 'omits guidance from metadata and denoise for distilled variant %s', + // guidance_embeds is inert for all current FLUX.2 Klein variants (weights are + // absent or zeroed), so the linear UI does not expose it and the graph builder + // must not write it into the denoise node or metadata. + it.each(['klein_9b_base', 'klein_9b', 'klein_4b_base', 'klein_4b'])( + 'omits guidance from metadata and denoise for variant %s', async (variant) => { currentModel = makeFlux2Model(variant); From 7612eb18366227cc25ac70fa84316340ca183cd9 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 23 Apr 2026 16:14:47 +0200 Subject: [PATCH 11/11] fix(flux2): skip Guidance metadata recall for legacy FLUX.2 images MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The generic Guidance metadata handler unconditionally parsed `metadata.guidance` and dispatched `setGuidance(value)` into the shared params slice. For images generated before the Klein guidance cleanup, this still fired — silently writing a stale guidance value into the global state, which then leaked back into FLUX.1 on model switch. Gate the handler on `metadata.model.base`: reject parsing when the image was generated with a FLUX.2 model. The handler is then skipped for both display and recall on legacy FLUX.2 metadata, matching the "silently ignored" contract stated in the PR. - parsing.tsx: check metadata.model.base in Guidance.parse() - parsing.test.tsx: three new cases covering FLUX.2 gating, FLUX.1 pass-through, and back-compat for metadata without a model field --- .../src/features/metadata/parsing.test.tsx | 46 +++++++++++++++++++ .../web/src/features/metadata/parsing.tsx | 9 ++++ 2 files changed, 55 insertions(+) diff --git a/invokeai/frontend/web/src/features/metadata/parsing.test.tsx b/invokeai/frontend/web/src/features/metadata/parsing.test.tsx index 01e33ee2dbb..bb295303273 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.test.tsx +++ b/invokeai/frontend/web/src/features/metadata/parsing.test.tsx @@ -125,4 +125,50 @@ describe('ImageMetadataHandlers — Klein recall gating', () => { expect(parsed.key).toBe('vae-key'); }); }); + + describe('Guidance (legacy FLUX.2 gating)', () => { + // Prior to the Klein guidance cleanup, FLUX.2 images wrote a `guidance` + // field into metadata. The guidance scalar is inert for all current Klein + // variants, so legacy values must not be recalled into the shared guidance + // state — otherwise they leak back into FLUX.1 when the user switches + // models. + it('rejects parsing when the image was generated with a FLUX.2 model', async () => { + const store = makeStore(); + + await expect( + Promise.resolve().then(() => + ImageMetadataHandlers.Guidance.parse( + { + model: { key: 'k', hash: 'h', name: 'Klein 9B Base', base: 'flux2', type: 'main' }, + guidance: 3.5, + }, + store + ) + ) + ).rejects.toThrow(); + }); + + it('parses successfully for FLUX.1 metadata', async () => { + const store = makeStore(); + + const parsed = await ImageMetadataHandlers.Guidance.parse( + { + model: { key: 'k', hash: 'h', name: 'FLUX Dev', base: 'flux', type: 'main' }, + guidance: 3.5, + }, + store + ); + + expect(parsed).toBe(3.5); + }); + + it('parses successfully when no model metadata is present', async () => { + // Metadata without a model field should still parse (back-compat for + // images where only scalar params were saved). + const store = makeStore(); + + const parsed = await ImageMetadataHandlers.Guidance.parse({ guidance: 3.5 }, store); + expect(parsed).toBe(3.5); + }); + }); }); diff --git a/invokeai/frontend/web/src/features/metadata/parsing.tsx b/invokeai/frontend/web/src/features/metadata/parsing.tsx index 8532c478fc4..cf55f378106 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.tsx +++ b/invokeai/frontend/web/src/features/metadata/parsing.tsx @@ -379,6 +379,15 @@ const Guidance: SingleMetadataHandler = { [SingleMetadataKey]: true, type: 'Guidance', parse: (metadata, _store) => { + // Legacy FLUX.2 images may still carry a `guidance` field, but guidance_embeds + // is inert for all current Klein variants. Reject parsing for FLUX.2 metadata + // so the handler is skipped on both display and recall - avoids leaking a stale + // value into the shared guidance param (which is still used by FLUX.1). + const rawModel = getProperty(metadata, 'model'); + const modelBase = (rawModel as { base?: unknown } | undefined)?.base; + if (modelBase === 'flux2') { + throw new Error('Guidance is not used for FLUX.2 Klein models.'); + } const raw = getProperty(metadata, 'guidance'); const parsed = zParameterGuidance.parse(raw); return Promise.resolve(parsed);