Skip to content

fix(mm): detect FLUX.2 Klein 9B Base variant via filename heuristic#9011

Merged
lstein merged 2 commits intomainfrom
fix/flux2-klein-guidance-and-metadata
Apr 7, 2026
Merged

fix(mm): detect FLUX.2 Klein 9B Base variant via filename heuristic#9011
lstein merged 2 commits intomainfrom
fix/flux2-klein-guidance-and-metadata

Conversation

@Pfannkuchensack
Copy link
Copy Markdown
Collaborator

Summary

FLUX.2 Klein 9B Base (undistilled) and Klein 9B (distilled) have identical model architectures — same block counts, same tensor shapes, same keys. They cannot be distinguished from the state dict alone.

This PR adds a filename heuristic to detect the Base variant: if the model name contains "base" (e.g. flux-2-klein-base-9b-fp8.safetensors), it is classified as klein_9b_base instead of klein_9b.

Also fixes the diffusers format detection, which incorrectly relied on the guidance_embeds config field to distinguish the variants. Both variants have guidance_embeds=False.

QA Instructions

  1. Import a FLUX.2 Klein 9B Base checkpoint with "base" in the filename (e.g. flux-2-klein-base-9b-fp8.safetensors)
  2. Verify it is detected as klein_9b_base
  3. Import a FLUX.2 Klein 9B (distilled) checkpoint without "base" in the filename
  4. Verify it is detected as klein_9b
  5. Verify Klein 4B models are unaffected

Merge Plan

No special considerations.

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Tests added / updated (if applicable)
  • ❗Changes to a redux slice have a corresponding migration
  • Documentation added / updated (if applicable)
  • Updated What's New copy (if doing a release after this PR)

Klein 9B Base (undistilled) and Klein 9B (distilled) have identical
architectures and cannot be distinguished from the state dict alone.
Use a filename heuristic ("base" in the name) to detect the Base
variant for checkpoint, GGUF, and diffusers format models.

Also fixes the incorrect guidance_embeds-based detection for diffusers
format, since both variants have guidance_embeds=False.
@github-actions github-actions bot added python PRs that change python files backend PRs that change backend files labels Apr 2, 2026
@lstein lstein self-assigned this Apr 7, 2026
@lstein lstein added the v6.13.x label Apr 7, 2026
@lstein lstein moved this to 6.13.x Theme: MODELS in Invoke - Community Roadmap Apr 7, 2026
Copy link
Copy Markdown
Collaborator

@lstein lstein left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Works as advertised.

@lstein lstein enabled auto-merge (squash) April 7, 2026 02:23
@lstein lstein merged commit dbbf289 into main Apr 7, 2026
13 checks passed
@lstein lstein deleted the fix/flux2-klein-guidance-and-metadata branch April 7, 2026 02:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

backend PRs that change backend files python PRs that change python files v6.13.x

Projects

Status: 6.13.x Theme: MODELS

Development

Successfully merging this pull request may close these issues.

2 participants