Spaces:
Running on Zero
Running on Zero
| """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 | |
| 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 | |