fix(flux2): remove inert guidance UI, add Klein 4B Base variant, fix metadata recall (#8995)

* 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

* 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.

* Chore pnpm fix

* Update version to 1.5.0 in flux2_denoise.py

* Update condition for rendering ParamFluxScheduler

* 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.

* Change Wrong Comment

* 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.

* Chore typegen

* fix test

* fix(flux2): skip Guidance metadata recall for legacy FLUX.2 images

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

---------

Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com>
Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com>

Author: 3 people

invokeai/app/invocations/flux2_denoise.py

@@ -53,8 +53,8 @@
     "flux2_denoise",
     title="FLUX2 Denoise",
     tags=["image", "flux", "flux2", "klein", "denoise"],
-    category="latents",
-    version="1.4.0",
+    category="image",
+    version="1.5.0",
     classification=Classification.Prototype,
 )
 class Flux2DenoiseInvocation(BaseInvocation):
@@ -101,6 +101,14 @@ 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="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,
         description=FieldDescriptions.cfg_scale,
@@ -467,6 +475,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,

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

invokeai/backend/flux/util.py

@@ -133,7 +133,24 @@ 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 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)
@@ -149,7 +166,24 @@ 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 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)
+        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,
     ),
 }
 

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,10 @@ 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.
+    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:
@@ -56,6 +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. 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).
@@ -76,9 +82,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 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
     use_scheduler = scheduler is not None
@@ -121,7 +128,7 @@ def denoise(
                 timestep=t_vec,
                 img_ids=img_ids,
                 txt_ids=txt_ids,
-                guidance=guidance,
+                guidance=guidance_vec,
                 return_dict=False,
             )
 
@@ -141,7 +148,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 +229,7 @@ def denoise(
                 timestep=t_vec,
                 img_ids=img_ids,
                 txt_ids=txt_ids,
-                guidance=guidance,
+                guidance=guidance_vec,
                 return_dict=False,
             )
 
@@ -242,7 +249,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,
                 )
 

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
@@ -389,6 +389,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
@@ -573,10 +574,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
 
@@ -745,10 +748,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
 
@@ -856,11 +861,10 @@ 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. We use a filename heuristic to detect Base models.
         """
         KLEIN_4B_CONTEXT_DIM = 7680  # 3 × 2560
         KLEIN_9B_CONTEXT_DIM = 12288  # 3 × 4096
@@ -875,6 +879,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

invokeai/backend/model_manager/taxonomy.py

@@ -132,7 +132,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)."""

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."
       },

invokeai/frontend/web/src/features/gallery/components/ImageMetadataViewer/ImageMetadataActions.tsx

@@ -63,6 +63,8 @@ export const ImageMetadataActions = memo((props: Props) => {
       <SingleMetadataDatum metadata={metadata} handler={ImageMetadataHandlers.QwenImageShift} />
       <SingleMetadataDatum metadata={metadata} handler={ImageMetadataHandlers.CanvasLayers} />
       <CollectionMetadataDatum metadata={metadata} handler={ImageMetadataHandlers.RefImages} />
+      <SingleMetadataDatum metadata={metadata} handler={ImageMetadataHandlers.KleinVAEModel} />
+      <SingleMetadataDatum metadata={metadata} handler={ImageMetadataHandlers.KleinQwen3EncoderModel} />
       <CollectionMetadataDatum metadata={metadata} handler={ImageMetadataHandlers.LoRAs} />
     </Flex>
   );