File size: 3,175 Bytes
300ccd6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """Loader for the compressed FLUX.2-klein-9B text encoder.
Mixed per-layer FFN widths / head counts are not expressible in a stock Qwen3 config, so
this loader rebuilds the module shapes from pruning_metadata.json and then strict-loads
the weights. NEVER load this checkpoint with ignore_mismatched_sizes=True.
Usage:
from loading import load_text_encoder, load_pipeline
te = load_text_encoder(".") # or a downloaded repo dir
pipe = load_pipeline("black-forest-labs/FLUX.2-klein-9B", te)
image = pipe(prompt, text_encoder_out_layers=(9, 17, 25), num_inference_steps=4,
guidance_scale=1.0).images[0]
IMPORTANT: every pipeline call MUST pass text_encoder_out_layers=(9, 17, 25).
Without it the pipeline reads the 36-layer default taps and silently produces
wrong embeddings.
"""
import json
from pathlib import Path
import torch
import torch.nn as nn
from transformers import AutoConfig
from transformers.models.qwen3.modeling_qwen3 import Qwen3Model
TEXT_ENCODER_OUT_LAYERS = (9, 17, 25)
def _resize(linear, out_features=None, in_features=None):
new = nn.Linear(in_features or linear.in_features, out_features or linear.out_features,
bias=linear.bias is not None)
return new.to(dtype=linear.weight.dtype)
def load_text_encoder(model_dir, torch_dtype=torch.bfloat16):
model_dir = Path(model_dir)
meta = json.loads((model_dir / "pruning_metadata.json").read_text(encoding="utf-8"))
config = AutoConfig.from_pretrained(model_dir)
model = Qwen3Model(config).to(torch_dtype)
head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
qpg = config.num_attention_heads // config.num_key_value_heads
for idx, keep in meta["ffn_keep_by_exported_layer"].items():
mlp = model.layers[int(idx)].mlp
mlp.gate_proj = _resize(mlp.gate_proj, out_features=int(keep))
mlp.up_proj = _resize(mlp.up_proj, out_features=int(keep))
mlp.down_proj = _resize(mlp.down_proj, in_features=int(keep))
for idx, kept in meta["head_groups_by_exported_layer"].items():
attn = model.layers[int(idx)].self_attn
attn.q_proj = _resize(attn.q_proj, out_features=int(kept) * qpg * head_dim)
attn.k_proj = _resize(attn.k_proj, out_features=int(kept) * head_dim)
attn.v_proj = _resize(attn.v_proj, out_features=int(kept) * head_dim)
attn.o_proj = _resize(attn.o_proj, in_features=int(kept) * qpg * head_dim)
if hasattr(attn, "num_key_value_groups"):
attn.num_key_value_groups = qpg
if meta.get("final_norm_identity"):
model.norm = nn.Identity()
import safetensors.torch as st
state = {}
for shard in sorted(model_dir.glob("model*.safetensors")):
state.update(st.load_file(shard))
model.load_state_dict(state, strict=True)
model.eval()
return model
def load_pipeline(base_model_id, text_encoder, torch_dtype=torch.bfloat16, **kwargs):
from diffusers import Flux2KleinPipeline
pipe = Flux2KleinPipeline.from_pretrained(
base_model_id, text_encoder=text_encoder, torch_dtype=torch_dtype, **kwargs)
return pipe
|