"""MiniMax-H3's Qwen3-VL conditioner, streamed one decoder layer at a time from disk on the CPU.

H3 conditions on `hidden_states[50]` of a text-only presentation, so only `embed_tokens` and decoder layers 0-49 are
read, and no layer stays resident: each is loaded from its shard (the next one in a background thread), run, dropped.
For a text-only token list this is `get_qwen3vl_prompt_embeds` without the 66 GB model in memory.
"""
import json
from concurrent.futures import ThreadPoolExecutor

import torch
from safetensors import safe_open
from transformers import AutoConfig
from transformers.masking_utils import create_causal_mask
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextDecoderLayer, Qwen3VLTextRotaryEmbedding

LM = "model.language_model."


def load(te_dir, weight_map, prefix):
    out = {}
    for f in sorted({f for k, f in weight_map.items() if k.startswith(prefix)}):
        with safe_open(f"{te_dir}/{f}", "pt") as sf:
            for k in sf.keys():
                if k.startswith(prefix):
                    out[k[len(prefix):]] = sf.get_tensor(k)
    return out


@torch.no_grad()
def encode(te_dir, token_ids, layer=50):
    cfg = AutoConfig.from_pretrained(te_dir).text_config
    cfg._attn_implementation = "sdpa"
    wm = json.load(open(f"{te_dir}/model.safetensors.index.json"))["weight_map"]
    pool = ThreadPoolExecutor(1)
    nxt = pool.submit(load, te_dir, wm, f"{LM}layers.0.")
    ids = torch.tensor([token_ids])
    h = load(te_dir, wm, LM + "embed_tokens.")["weight"][ids]
    # text only: all three M-RoPE axes carry the same 0..T-1, as Qwen3VLModel builds them without vision
    pos = torch.arange(len(token_ids))[None]
    pe = Qwen3VLTextRotaryEmbedding(cfg)(h, pos[None].expand(3, 1, -1))
    mask = create_causal_mask(config=cfg, inputs_embeds=h, attention_mask=torch.ones_like(ids),
                              past_key_values=None, position_ids=pos)
    with torch.device("meta"):
        block = Qwen3VLTextDecoderLayer(cfg, 0)
    for i in range(layer):
        sd = nxt.result()
        if i + 1 < layer:
            nxt = pool.submit(load, te_dir, wm, f"{LM}layers.{i + 1}.")
        block.load_state_dict(sd, assign=True)
        h = block(h, position_embeddings=pe, attention_mask=mask, position_ids=pos)
    return h


if __name__ == "__main__":
    # python te_stream.py <text_encoder dir> <token ids .json> <out .pt>
    # app.py runs this as a subprocess: a fresh interpreter, so the app's own process never does the CPU work
    # (OpenMP state in a process that ZeroGPU later forks is asking for a hang) and the layers' memory dies with it.
    import sys

    te_dir, ids, out = sys.argv[1:4]
    torch.save(encode(te_dir, json.load(open(ids))).contiguous(), out)  # (1, T, 5120) bf16, as sample.py feeds it