flux2distill/model_utils.py

"""Model loading, param-group split (Muon vs AdamW by name), and param accounting.

Verified klein-4B layout (transformer/config.json):
  transformer_blocks: 5 x Flux2TransformerBlock (245.4M each)
  single_transformer_blocks: 20 x Flux2SingleTransformerBlock (122.68M each)
  d=3072, heads=24, head_dim=128, mlp_ratio=3, joint_attention_dim=7680,
  guidance_embeds=false (CFG-free), is_distilled=true.
"""
from __future__ import annotations

import torch

# Top-level modules routed to AdamW (embedders / modulation / output / norms).
# Everything else (2D attention/MLP weights inside blocks + surrogate A/B) -> Muon.
ADAMW_PREFIXES = (
    "pos_embed",
    "time_guidance_embed",
    "double_stream_modulation_img",
    "double_stream_modulation_txt",
    "single_stream_modulation",
    "x_embedder",
    "context_embedder",
    "norm_out",
    "proj_out",
)

DTYPES = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}


def load_transformer(local_dir="models/klein-4b/transformer", dtype="bfloat16", device="cpu"):
    from diffusers import Flux2Transformer2DModel
    tf = Flux2Transformer2DModel.from_pretrained(local_dir, torch_dtype=DTYPES[dtype])
    return tf.to(device).eval()


def load_pipeline(local_dir="models/klein-4b", dtype="bfloat16", device="cuda"):
    from diffusers import Flux2KleinPipeline
    pipe = Flux2KleinPipeline.from_pretrained(local_dir, torch_dtype=DTYPES[dtype])
    pipe = pipe.to(device)
    return pipe


def build_param_groups(transformer, lr_muon=0.02, lr_adamw=2e-4, weight_decay=0.01):
    """Split params: Muon on 2D block weights, AdamW on embedders/modulation/norms/biases.

    Returns (muon_params, adamw_params, info). Route 1D params (norms/biases) anywhere
    to AdamW since Muon requires 2D updates.
    """
    muon, adamw = [], []
    for name, p in transformer.named_parameters():
        if not p.requires_grad:
            continue
        is_adamw_name = any(name.startswith(pre) or f".{pre}" in name for pre in ADAMW_PREFIXES)
        if is_adamw_name or p.ndim != 2:
            adamw.append(p)
        else:
            muon.append(p)
    info = {
        "muon_tensors": len(muon), "muon_params": sum(p.numel() for p in muon),
        "adamw_tensors": len(adamw), "adamw_params": sum(p.numel() for p in adamw),
    }
    return muon, adamw, info


def param_summary(transformer):
    total = sum(p.numel() for p in transformer.parameters())
    n_double = len(transformer.transformer_blocks)
    singles = transformer.single_transformer_blocks
    from .surrogate import LowRankResidualSurrogate, LinearAttentionSurrogate
    n_surrogate = sum(isinstance(b, (LowRankResidualSurrogate, LinearAttentionSurrogate)) for b in singles)
    n_full_single = len(singles) - n_surrogate
    return {
        "total_params": total,
        "total_B": total / 1e9,
        "n_double": n_double,
        "n_full_single": n_full_single,
        "n_surrogate": n_surrogate,
    }