"""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