feat: add diffusers format support for FLUX.2-klein#1369
Open
ComputerPers wants to merge 9 commits intoleejet:masterfrom
Open
feat: add diffusers format support for FLUX.2-klein#1369ComputerPers wants to merge 9 commits intoleejet:masterfrom
ComputerPers wants to merge 9 commits intoleejet:masterfrom
Conversation
…ion() Add checks for diffusers-style tensor names alongside existing BFL-format checks: - double_stream_modulation_img.linear.weight (.linear. vs .lin.) - single_transformer_blocks.47. (diffusers block naming) This enables get_sd_version() to correctly identify FLUX.2 and FLUX.2-klein models loaded from diffusers-format safetensors checkpoints.
…LUX.2 diffusers support Implement diffusers→BFL tensor name mapping for FLUX.2 architecture: - Shared modulation: .linear. → .lin. - Time embedding: time_guidance_embed → time_in - Double blocks: transformer_blocks → double_blocks (with SwiGLU MLP) - Single blocks: single_transformer_blocks → single_blocks (fused qkv) - Final layer mapping Block counts are parametric (not hardcoded) for future FLUX.2 variants. Route FLUX.2 versions to new function in convert_diffusion_model_name(). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Add dst_offset to TensorStorage for partial tensor loads. Detect split tensor parts (.weight.1, .weight.2) and compute byte offsets to fuse them into the expected combined tensor. Handle adaLN modulation half-swap (diffusers [shift,scale] → BFL [scale,shift]). Add SD_DUMP_CHECKSUMS env var for tensor comparison debugging. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
41 tests covering all tensor categories: embedders, shared modulations, double blocks (q/k/v split, proj, norms, MLP SwiGLU), single blocks (fused qkv, proj, norms), final layer. Both VERSION_FLUX2 and VERSION_FLUX2_KLEIN variants tested. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Add init_from_safetensors_index() to load multi-shard safetensors files by reading model.safetensors.index.json and loading each shard. Update init_from_file() to auto-detect index.json when .safetensors file is missing. Update init_from_diffusers_file() to support transformer/ directory (DiT models like FLUX, SD3) alongside unet/. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Read model_index.json to detect LLM-based text encoders (Qwen, Llama) and use text_encoders.llm. prefix instead of te. for correct loading. Support transformer/ directory in init_from_diffusers_file() for DiT models (FLUX, SD3). Enables single-directory loading: sd-cli --diffusion-model <diffusers-dir>/ Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
|
Not very useful if you ask me. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This adds support for loading FLUX.2-klein models in HuggingFace diffusers format — users can load models directly from
~/.cache/huggingface/hub/without converting to GGUF or BFL first.Also fixes #699 — the split-fuse qkv fix is generic and resolves diffusers format loading for SD3, FLUX.1 and Z-Image as well.
Along the way I fixed two issues that affect all diffusers-format models, not just FLUX.2:
Split-fuse qkv loading. The
.weight.Nsuffix convention inname_conversion.cppwas already used for SD3, FLUX.1 and Z-Image to split q/k/v into separate tensors, butload_tensors()never handled the split parts — they were silently dropped as "unknown tensor" and the base tensor had a shape mismatch. This PR addsdst_offsetsupport so split parts get written into the correct positions of the fused destination tensor. This fixes diffusers loading for all models that use split qkv, not just FLUX.2.Multi-shard safetensors. Large HF models are often sharded into multiple
.safetensorsfiles with amodel.safetensors.index.jsonindex. Addedinit_from_safetensors_index()that reads the index and loads each shard. This enables loading sharded text encoders (e.g. FLUX.2-klein-4B Qwen3 TE is 2 shards, FLUX.1-dev T5 is 2 shards).Other changes:
init_from_diffusers_file()now triestransformer/beforeunet/for DiT modelsmodel_index.jsonto auto-detect LLM text encoders (Qwen, Llama) and use the correct prefixBefore / After
Without this PR, loading FLUX.2-klein-4B diffusers format fails:
With this PR, both work:
Testing
Tested on Apple Silicon (M3 Max, 36 GB):
--diffusion-model <dir>/): auto-detected Qwen3 TE, generated OKBFL (reference) vs Diffusers (this PR) — FLUX.2-klein-4B, 1024x1024, seed 42, 4 steps, prompt "gandalf you shall not pass":