export_script/a_export_to_onnx.py

"""
Export Qwen3-VL-Embedding model to ONNX files for inference.

Produces 2 ONNX files:
  1. Vision.onnx      — Vision encoder (fixed image resolution)
  2. Transformer.onnx — Transformer decoder layers (embedding mode, no KV cache)

Also saves rotary_params.npz, which contains:
  • inv_freq + mrope_section — for Python-side mRoPE computation
  • embed_weight             — token-embedding lookup table (replaces Embed.onnx)
  • image dimensions / config values

Usage:
    python 1_export_to_onnx.py
    python 1_export_to_onnx.py --model_path /path/to/model --output_dir /path/to/output
"""

import math
import os
import gc
import glob
import argparse
import torch
import onnx
import numpy as np
from transformers.models.qwen3_vl.modeling_qwen3_vl import (
    Qwen3VLVisionModel,
)
from qwen3_vl_embedding import Qwen3VLForEmbedding, MAX_TOTAL_PIXELS


def consolidate_external_data(onnx_path: str):
    """
    Re-save an ONNX model so all external weights live in one .onnx.data file
    instead of hundreds of per-tensor files.
    """
    out_dir  = os.path.dirname(onnx_path)
    base     = os.path.basename(onnx_path)
    data_rel = base + ".data"                      # sibling file: Transformer.onnx.data
    model    = onnx.load(onnx_path, load_external_data=True)

    # Delete ONNX per-tensor external files (e.g.
    # `model.model.language_model.layers.0.self_attn.qk_norm_weight`). They can
    # be named with dots, so we use an allow-list of known suffixes we want to
    # keep. Subdirectories (e.g. `tokenizer/`) are skipped via the isfile check.
    PRESERVE_SUFFIXES = (".onnx", ".onnx.data", ".npz", ".bin")
    for fname in list(os.listdir(out_dir)):
        fpath = os.path.join(out_dir, fname)
        if not os.path.isfile(fpath):
            continue
        if any(fname.endswith(sfx) for sfx in PRESERVE_SUFFIXES):
            continue
        try:
            os.remove(fpath)
        except OSError:
            pass

    # Remove the target .data file if it already exists (stale)
    target_data = os.path.join(out_dir, data_rel)
    if os.path.exists(target_data):
        os.remove(target_data)

    onnx.save_model(
        model, onnx_path,
        save_as_external_data=True,
        all_tensors_to_one_file=True,
        location=data_rel,
        size_threshold=1024,
        convert_attribute=False,
    )
    del model
    gc.collect()


# ══════════════════════════════════════════════════════════════════════════════
# Configuration — reads from the module's config.py
# ══════════════════════════════════════════════════════════════════════════════
IMG_SIZE = (768,768) # (h,w)
QWEN3VLE_MODEL_HF_PATH = os.getenv(
    "MODEL_QWEN3VLE_MODEL_HF_PATH", "/home/jordan/Research/Product-AI-mono/assets/model/Qwen3-VL-Embedding-2B"
)
QWEN3VLE_TRT_DIR_PATH = os.getenv(
    "MODEL_QWEN3VLE_TRT_DIR_PATH", "assets/model/Qwen3-VL-Embedding-2B-onnx"
)
TEMPORAL_SIZE=1

DEFAULT_MODEL_PATH = QWEN3VLE_MODEL_HF_PATH
DEFAULT_OUTPUT_DIR = QWEN3VLE_TRT_DIR_PATH

# Vision — fixed image resolution at export time.
# Must be multiples of (patch_size * merge_size), validated after model loading.
IMAGE_HEIGHT     = IMG_SIZE[0]
IMAGE_WIDTH      = IMG_SIZE[1]

MAX_SEQ_LEN      = MAX_TOTAL_PIXELS                           # 8192
OPSET            = 17


# ══════════════════════════════════════════════════════════════════════════════
# Token Embedding Module
# ══════════════════════════════════════════════════════════════════════════════
class LLM_EMBED(torch.nn.Module):
    """Extract the token-embedding layer and run it in float32."""

    def __init__(self, model):
        super().__init__()
        self.embed_tokens = model.model.language_model.embed_tokens.float()

    def forward(self, input_ids):
        return self.embed_tokens(input_ids)


# ══════════════════════════════════════════════════════════════════════════════
# Vision Encoder Module
# ══════════════════════════════════════════════════════════════════════════════
class LLM_VISION_EMBED(torch.nn.Module):
    """
    ONNX-exportable vision encoder with baked-in positional / rotary
    embeddings for a fixed image resolution.

    Takes the HF processor's pixel_values directly (already CLIP-normalised
    and patched) so the patch ordering, normalisation, and Conv3d behaviour
    are identical to the PyTorch model.

    Input : float32 pixel_values [total_patches, flatten_dim]
    Output: (deepstack_feature_0, …, vision_hidden_states)
    """

    def __init__(self, model, height_factor, width_factor, temporal_patches=1):
        super().__init__()
        visual        = model.model.visual
        vision_config = model.config.vision_config

        self.num_heads     = vision_config.num_heads
        self.head_dim      = vision_config.hidden_size // self.num_heads
        self.head_dim_half = self.head_dim // 2
        self.embed_dim     = visual.patch_embed.embed_dim
        self.merge_size    = visual.spatial_merge_size
        self.t_patches     = temporal_patches

        ms = self.merge_size
        self.grid_h = height_factor * ms
        self.grid_w = width_factor  * ms

        # ── pre-compute positional / rotary embeddings for the fixed grid ─
        grid_thw = torch.tensor(
            [[self.t_patches, self.grid_h, self.grid_w]], dtype=torch.int32
        )

        pos_embeds = Qwen3VLVisionModel.fast_pos_embed_interpolate(
            visual, grid_thw
        ).unsqueeze(0)                                  # [1, N, embed_dim]
        self.register_buffer("pos_embeds", pos_embeds)

        rot_pos_emb = (
            Qwen3VLVisionModel.rot_pos_emb(visual, grid_thw)
            .float()
            .unsqueeze(0).unsqueeze(0).unsqueeze(0)     # [1, 1, 1, N, dim/2]
        )
        cos = rot_pos_emb.cos()
        sin = rot_pos_emb.sin()
        self.register_buffer("rotary_cos", torch.cat([cos, cos], dim=-1))
        self.register_buffer("rotary_sin", torch.cat([-sin, sin], dim=-1))

        # ── ONNX-friendly GELU ──────────────────────────────────────────
        self._replace_gelu(visual)

        # ── fuse layer-norms into the subsequent linear layers ──────────
        scaling = self.head_dim ** -0.25
        for blk in visual.blocks:
            blk.attn.qkv.weight.data[: -self.embed_dim] *= scaling
            blk.attn.qkv.bias.data[: -self.embed_dim]   *= scaling
            self._fuse_norm(blk.norm1, blk.attn.qkv)
            self._fuse_norm(blk.norm2, blk.mlp.linear_fc1)

        for ds_layer in visual.deepstack_merger_list:
            self._fuse_norm(ds_layer.norm, ds_layer.linear_fc1)

        self._fuse_norm(visual.merger.norm, visual.merger.linear_fc1)

        self.visual = visual

    # ── helpers ──────────────────────────────────────────────────────────
    @staticmethod
    def _fuse_norm(norm, linear):
        norm_bias   = norm.bias.data
        norm_weight = norm.weight.data
        if linear.weight.shape[1] != norm_bias.shape[0]:
            repeat_factor = linear.weight.shape[1] // norm_bias.shape[0]
            norm_bias     = norm_bias.repeat(repeat_factor)
            norm_weight   = norm_weight.repeat(repeat_factor)
        linear.bias.data.add_(torch.matmul(linear.weight.data, norm_bias))
        linear.weight.data.mul_(norm_weight.unsqueeze(0))
        norm.elementwise_affine = False
        norm.weight = None
        norm.bias   = None

    @staticmethod
    def _replace_gelu(module):
        for name, child in module.named_children():
            if isinstance(child, torch.nn.GELU):
                setattr(module, name, torch.nn.GELU(approximate="tanh"))
            else:
                LLM_VISION_EMBED._replace_gelu(child)

    def _rotate_half(self, x, batch_size):
        x = x.view(2, batch_size, self.num_heads, -1, 2, self.head_dim_half)
        x = x.flip(-2)
        return x.view(2, batch_size, self.num_heads, -1, self.head_dim)

    def _spatial_merge(self, x, target_hidden_size):
        """Explicit spatial reshape: group merge_size×merge_size patches per frame."""
        ms = self.merge_size
        x = x.view(1, self.t_patches, self.grid_h, self.grid_w, -1)
        x = x.view(1, self.t_patches, self.grid_h // ms, ms, self.grid_w // ms, ms, -1)
        x = x.permute(0, 1, 2, 4, 3, 5, 6)
        return x.reshape(1, -1, target_hidden_size)

    # ── forward ──────────────────────────────────────────────────────────
    def forward(self, pixel_values):
        """
        pixel_values : float32 [total_patches, flatten_dim]
            Already CLIP-normalised, from the HF processor.
        """
        batch_size = 1

        # patch embedding (original HF module)
        hidden = self.visual.patch_embed(pixel_values)           # [N, embed_dim]
        hidden = hidden.unsqueeze(0)                             # [1, N, embed_dim]
        hidden = hidden + self.pos_embeds

        # transformer blocks
        deepstack_features = []
        ds_indices = self.visual.deepstack_visual_indexes
        ds_modules = self.visual.deepstack_merger_list

        for layer_num, blk in enumerate(self.visual.blocks):
            # --- self-attention (manual, ONNX-safe) ---
            h_norm = blk.norm1(hidden)
            qkv    = blk.attn.qkv(h_norm)
            qkv    = qkv.reshape(batch_size, -1, 3, self.num_heads, self.head_dim)
            qkv    = qkv.permute(2, 0, 3, 1, 4)          # [3, B, heads, seq, dim]
            qk, v  = qkv.split([2, 1], dim=0)

            qk_rot = (
                qk * self.rotary_cos
                + self._rotate_half(qk, batch_size) * self.rotary_sin
            )
            q_rot, k_rot = qk_rot.split([1, 1], dim=0)
            attn = torch.matmul(q_rot, k_rot.transpose(-1, -2))
            attn = torch.softmax(attn, dim=-1)
            attn = torch.matmul(attn, v)
            attn = attn.transpose(2, 3).reshape(
                batch_size, -1, blk.attn.proj.in_features
            )
            hidden = hidden + blk.attn.proj(attn)

            # --- feed-forward ---
            mlp_out = blk.mlp.linear_fc1(blk.norm2(hidden))
            mlp_out = blk.mlp.act_fn(mlp_out)
            mlp_out = blk.mlp.linear_fc2(mlp_out)
            hidden  = hidden + mlp_out

            # --- deepstack (explicit spatial merge) ---
            if layer_num in ds_indices:
                idx      = ds_indices.index(layer_num)
                ds_layer = ds_modules[idx]
                x_ds = self._spatial_merge(hidden, ds_layer.hidden_size)
                x_ds = ds_layer.norm(x_ds)
                x_ds = ds_layer.linear_fc1(x_ds)
                x_ds = ds_layer.act_fn(x_ds)
                x_ds = ds_layer.linear_fc2(x_ds)
                deepstack_features.append(x_ds)

        # merger (explicit spatial merge)
        hidden = self.visual.merger.norm(hidden)
        hidden = self._spatial_merge(hidden, self.visual.merger.hidden_size)
        hidden = self.visual.merger.linear_fc1(hidden)
        hidden = self.visual.merger.act_fn(hidden)
        hidden = self.visual.merger.linear_fc2(hidden)

        return *deepstack_features, hidden


# ══════════════════════════════════════════════════════════════════════════════
# Transformer Module (embedding mode — no KV cache, no lm_head)
# ══════════════════════════════════════════════════════════════════════════════
class LLM_MAIN_EMBED(torch.nn.Module):
    """
    All decoder layers in a single ONNX graph.

    Differences from the generation-oriented LLM_MAIN:
      • No KV-cache inputs / outputs  (single forward pass)
      • No lm_head projection         (we need hidden states, not logits)
      • Final RMSNorm is kept         (not fused into lm_head)
      • Outputs the full sequence      (not just last token)

    Input layout  (via *all_inputs):
        hidden_states, ds_feat_0 … ds_feat_N, rotary_cos, rotary_sin, attn_mask

    Output:
        last_hidden_state  [batch, seq_len, hidden_size]
    """

    def __init__(
        self,
        model,
        num_heads,
        num_key_value_heads,
        head_dim,
        num_layers,
        hidden_size,
        deepstack_features_len,
    ):
        super().__init__()
        self.model = model

        self.head_dim            = head_dim
        self.head_dim_half       = head_dim // 2
        self.num_heads           = num_heads
        self.num_key_value_heads = num_key_value_heads
        self.num_key_value_groups = num_heads // num_key_value_heads
        self.qk_heads            = num_heads + num_key_value_heads
        self.num_layers          = num_layers
        self.deepstack_features_len = deepstack_features_len

        # 3 trailing args after deepstack features: cos, sin, mask
        self._ds_offset = 3 + deepstack_features_len

        # ── save the final RMSNorm (before weight fusion touches layers) ─
        lang = model.model.language_model
        self.register_buffer(
            "final_norm_weight", lang.norm.weight.data.clone()
        )
        self.final_norm_eps = float(lang.norm.variance_epsilon)

        # ── ONNX-friendly GELU ──────────────────────────────────────────
        self._replace_gelu(model)

        # ── fuse weights (same optimisations as reference) ──────────────
        self._fuse_weights(hidden_size)

    # ══════════════════════════════════════════════════════════════════════
    # Weight fusion (runs once at init)
    # ══════════════════════════════════════════════════════════════════════
    def _fuse_weights(self, hidden_size):
        scale_factor   = self.head_dim ** -0.25
        norm_factor    = hidden_size ** 0.5
        norm_factor_qk = self.head_dim ** 0.5

        with torch.no_grad():
            for layer in self.model.model.language_model.layers:
                self._fuse_qkv(layer, scale_factor, norm_factor, norm_factor_qk)
                self._fuse_gate_up(layer, norm_factor)
            # NOTE: we do NOT fuse the final norm into lm_head (no lm_head).

    def _fuse_qkv(self, layer, scale_factor, norm_factor, norm_factor_qk):
        attn = layer.self_attn
        q, k, v = attn.q_proj, attn.k_proj, attn.v_proj

        in_f  = int(q.in_features)
        out_f = int(q.out_features + k.out_features + v.out_features)
        has_bias = any(p.bias is not None for p in (q, k, v))

        qkv = torch.nn.Linear(in_f, out_f, bias=has_bias)
        qkv.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0))
        if has_bias:
            def _b(p):
                return p.bias if p.bias is not None else torch.zeros(
                    p.out_features, dtype=qkv.weight.dtype
                )
            qkv.bias.copy_(torch.cat([_b(q), _b(k), _b(v)], dim=0))

        attn.q_out_features = int(q.out_features)
        attn.k_out_features = int(k.out_features)
        attn.v_out_features = int(v.out_features)
        del attn.q_proj, attn.k_proj, attn.v_proj

        # fuse QK norms + attention scaling
        combined_scale = scale_factor * norm_factor_qk
        attn.q_norm.weight.mul_(combined_scale)
        attn.k_norm.weight.mul_(combined_scale)
        q_norm_rep = attn.q_norm.weight.repeat(self.num_heads)
        k_norm_rep = attn.k_norm.weight.repeat(self.num_key_value_heads)
        attn.qk_norm_weight = torch.nn.Parameter(
            torch.cat([q_norm_rep, k_norm_rep], dim=0).view(
                1, 1, 1, -1, self.head_dim
            )
        )
        del attn.q_norm, attn.k_norm

        # absorb input LayerNorm into QKV
        input_norm_w = layer.input_layernorm.weight.unsqueeze(0) * norm_factor
        qkv.weight.mul_(input_norm_w)
        attn.qkv = qkv
        del layer.input_layernorm

    def _fuse_gate_up(self, layer, norm_factor):
        post_norm_w = layer.post_attention_layernorm.weight.unsqueeze(0) * norm_factor
        gate, up = layer.mlp.gate_proj, layer.mlp.up_proj
        gate_up = torch.nn.Linear(
            gate.in_features, gate.out_features + up.out_features, bias=False
        )
        gate_up.weight.copy_(
            torch.cat([gate.weight * post_norm_w, up.weight * post_norm_w], dim=0)
        )
        layer.mlp.gate_up_proj = gate_up
        del layer.mlp.gate_proj, layer.mlp.up_proj, layer.post_attention_layernorm

    # ══════════════════════════════════════════════════════════════════════
    # Utility
    # ══════════════════════════════════════════════════════════════════════
    @staticmethod
    def _replace_gelu(module):
        for name, child in module.named_children():
            if isinstance(child, torch.nn.GELU):
                setattr(module, name, torch.nn.GELU(approximate="tanh"))
            else:
                LLM_MAIN_EMBED._replace_gelu(child)

    def _rms_norm(self, x):
        """Sum-based RMS norm (weight is absorbed into the subsequent linear)."""
        return x * torch.rsqrt(x.square().sum(-1, keepdim=True))

    def _rotate_half(self, x, batch_size):
        x = x.view(batch_size, -1, 1, self.qk_heads, 2, self.head_dim_half)
        x = x.flip(-2)
        return x.view(batch_size, -1, 1, self.qk_heads, self.head_dim)

    # ══════════════════════════════════════════════════════════════════════
    # Forward
    # ══════════════════════════════════════════════════════════════════════
    def forward(self, *all_inputs):
        """
        all_inputs layout:
            hidden_states          [batch, seq, hidden]
            deepstack_feat_0 …     [1, seq, hidden]   (×deepstack_features_len)
            rotary_cos             [1, seq, 1, 1, head_dim]
            rotary_sin             [1, seq, 1, 1, head_dim]
            attention_mask         [1, 1, 1, seq, seq]
        """
        hidden_states  = all_inputs[0]
        rotary_cos     = all_inputs[-3]
        rotary_sin     = all_inputs[-2]
        attention_mask  = all_inputs[-1]
        batch_size     = hidden_states.shape[0]

        for i, layer in enumerate(self.model.model.language_model.layers):
            # ── self-attention ───────────────────────────────────────
            residual      = hidden_states
            hidden_states = self._rms_norm(hidden_states)

            qkv = layer.self_attn.qkv(hidden_states)
            qkv = qkv.reshape(
                batch_size, -1, 1,
                self.qk_heads + self.num_key_value_heads,
                self.head_dim,
            )
            qk, v = torch.split(
                qkv, [self.qk_heads, self.num_key_value_heads], dim=-2
            )

            qk     = self._rms_norm(qk) * layer.self_attn.qk_norm_weight
            qk_rot = (
                qk * rotary_cos
                + self._rotate_half(qk, batch_size) * rotary_sin
            )

            q, k = torch.split(
                qk_rot, [self.num_heads, self.num_key_value_heads], dim=-2
            )
            q = q.reshape(
                batch_size, -1,
                self.num_key_value_heads,
                self.num_key_value_groups,
                self.head_dim,
            )
            q = q.permute(0, 2, 3, 1, 4)      # [B, kv_h, g, seq, dim]
            k = k.permute(0, 3, 2, 4, 1)       # [B, kv_h, 1, dim, seq]
            v = v.transpose(1, 3)               # [B, kv_h, 1, seq, dim]

            # no KV cache — direct attention over the full sequence
            attn = torch.matmul(q, k) + attention_mask
            attn = torch.softmax(attn, dim=-1)
            attn = torch.matmul(attn, v)

            attn = attn.permute(0, 3, 1, 2, 4).reshape(
                batch_size, -1, layer.self_attn.o_proj.in_features
            )
            hidden_states = residual + layer.self_attn.o_proj(attn)

            # ── feed-forward ─────────────────────────────────────────
            residual      = hidden_states
            hidden_states = self._rms_norm(hidden_states)

            gate_up = layer.mlp.gate_up_proj(hidden_states)
            gate, up = torch.split(
                gate_up,
                [layer.mlp.down_proj.in_features, layer.mlp.down_proj.in_features],
                dim=-1,
            )
            hidden_states = residual + layer.mlp.down_proj(
                layer.mlp.act_fn(gate) * up
            )

            # ── deepstack feature injection ──────────────────────────
            if i < self.deepstack_features_len:
                hidden_states = all_inputs[1 + i] + hidden_states

        # ── final RMSNorm (standard, not fused) ─────────────────────────
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(
            variance + self.final_norm_eps
        )
        hidden_states = hidden_states * self.final_norm_weight

        return hidden_states


# ══════════════════════════════════════════════════════════════════════════════
# Export
# ══════════════════════════════════════════════════════════════════════════════
def main(model_path, output_dir):
    os.makedirs(output_dir, exist_ok=True)

    onnx_vision      = os.path.join(output_dir, "Vision.onnx")
    onnx_transformer = os.path.join(output_dir, "Transformer.onnx")
    rotary_params    = os.path.join(output_dir, "rotary_params.npz")

    # ── Save processor files (tokenizer + image processor config) ───────
    # Stored in a `tokenizer/` subdirectory so the output layout stays tidy
    # and the directory remains self-contained — no HF repo pull at inference.
    print("Saving processor files …")
    from transformers.models.qwen3_vl.processing_qwen3_vl import Qwen3VLProcessor
    tokenizer_dir = os.path.join(output_dir, "tokenizer")
    os.makedirs(tokenizer_dir, exist_ok=True)
    Qwen3VLProcessor.from_pretrained(model_path).save_pretrained(tokenizer_dir)
    print(f"  Processor saved to {tokenizer_dir}")

    print("Loading model …")
    model = Qwen3VLForEmbedding.from_pretrained(
        model_path, torch_dtype=torch.float32, device_map="cpu", low_cpu_mem_usage=True
    ).eval()

    # ── extract config ──────────────────────────────────────────────────
    text_cfg = model.config.text_config
    num_heads           = text_cfg.num_attention_heads
    num_key_value_heads = text_cfg.num_key_value_heads
    head_dim            = text_cfg.head_dim
    num_layers          = text_cfg.num_hidden_layers
    hidden_size         = text_cfg.hidden_size
    deepstack_features_len = len(model.model.visual.deepstack_visual_indexes)

    patch_size  = model.model.visual.patch_size
    merge_size  = model.model.visual.spatial_merge_size
    temporal_patch_size = model.model.visual.patch_embed.temporal_patch_size
    unit        = patch_size * merge_size

    # Compute or validate IMAGE_HEIGHT / IMAGE_WIDTH
    if IMAGE_HEIGHT is None:
        IMAGE_HEIGHT_val = 16 * unit   # sensible default
    else:
        IMAGE_HEIGHT_val = IMAGE_HEIGHT
    if IMAGE_WIDTH is None:
        IMAGE_WIDTH_val = 16 * unit
    else:
        IMAGE_WIDTH_val = IMAGE_WIDTH

    assert IMAGE_HEIGHT_val % unit == 0, (
        f"IMAGE_HEIGHT ({IMAGE_HEIGHT_val}) must be a multiple of "
        f"patch_size * merge_size = {unit}"
    )
    assert IMAGE_WIDTH_val % unit == 0, (
        f"IMAGE_WIDTH ({IMAGE_WIDTH_val}) must be a multiple of "
        f"patch_size * merge_size = {unit}"
    )

    HEIGHT_FACTOR     = IMAGE_HEIGHT_val // unit
    WIDTH_FACTOR      = IMAGE_WIDTH_val  // unit

    # Number of temporal patches after Conv3d with stride=temporal_patch_size.
    # The HF processor pads frames to multiples of temporal_patch_size.
    temporal_patches  = math.ceil(TEMPORAL_SIZE / temporal_patch_size)
    vision_embed_size = temporal_patches * HEIGHT_FACTOR * WIDTH_FACTOR

    print(f"  text  : layers={num_layers}  heads={num_heads}  kv_heads={num_key_value_heads}  "
          f"head_dim={head_dim}  hidden={hidden_size}")
    print(f"  vision: deepstack={deepstack_features_len}  patch={patch_size}  merge={merge_size}")
    print(f"  TEMPORAL_SIZE={TEMPORAL_SIZE}  temporal_patch_size={temporal_patch_size}"
          f"  →  temporal_patches={temporal_patches}")
    print(f"  image size: {IMAGE_HEIGHT_val}x{IMAGE_WIDTH_val}  →  {vision_embed_size} vision tokens")

    with torch.inference_mode():

        # ══════════════════════════════════════════════════════════════════
        # 1. Save rotary parameters + token embedding weight (to skip Embed.onnx)
        # ══════════════════════════════════════════════════════════════════
        rotary_emb = model.model.language_model.rotary_emb
        inv_freq      = rotary_emb.inv_freq.cpu().numpy()
        mrope_section = np.array(rotary_emb.mrope_section, dtype=np.int64)
        embed_weight = (
            model.model.language_model.embed_tokens.weight
            .detach().cpu().float().numpy()
        )
        np.savez(
            rotary_params,
            inv_freq=inv_freq,
            mrope_section=mrope_section,
            embed_weight=embed_weight,
            image_height=IMAGE_HEIGHT_val,
            image_width=IMAGE_WIDTH_val,
            height_factor=HEIGHT_FACTOR,
            width_factor=WIDTH_FACTOR,
            patch_size=patch_size,
            merge_size=merge_size,
            hidden_size=hidden_size,
            head_dim=head_dim,
            max_seq_len=MAX_SEQ_LEN,
            temporal_patch_size=temporal_patch_size,
            temporal_patches=temporal_patches,
            temporal_size=TEMPORAL_SIZE,
        )
        print(f"Saved rotary params → {rotary_params}")

        # ══════════════════════════════════════════════════════════════════
        # 2. Export Vision
        # ══════════════════════════════════════════════════════════════════
        print("Exporting Vision …")
        total_patches = temporal_patches * HEIGHT_FACTOR * merge_size * WIDTH_FACTOR * merge_size
        flatten_dim   = 3 * temporal_patch_size * patch_size * patch_size
        dummy_pixels  = torch.randn(total_patches, flatten_dim, dtype=torch.float32)

        vision_output_names = []
        vision_dynamic_axes = {}
        for i in range(deepstack_features_len):
            name = f"deepstack_feature_{i}"
            vision_output_names.append(name)
        vision_output_names.append("vision_hidden_states")

        torch.onnx.export(
            LLM_VISION_EMBED(model, HEIGHT_FACTOR, WIDTH_FACTOR, temporal_patches),
            (dummy_pixels,),
            onnx_vision,
            input_names=["pixel_values"],
            output_names=vision_output_names,
            dynamic_axes=vision_dynamic_axes if vision_dynamic_axes else None,
            opset_version=OPSET,
            dynamo=False,
        )
        del dummy_pixels
        gc.collect()
        print(f"  → {onnx_vision}")
        print("  Consolidating external weights …")
        consolidate_external_data(onnx_vision)

        # ══════════════════════════════════════════════════════════════════
        # 3. Export Transformer
        # ══════════════════════════════════════════════════════════════════
        print("Exporting Transformer …")

        # dummy sequence length for tracing
        dummy_seq   = 10 + vision_embed_size
        dummy_hidden = torch.ones((1, dummy_seq, hidden_size), dtype=torch.float32)
        dummy_ds     = torch.ones((1, dummy_seq, hidden_size), dtype=torch.float32)
        dummy_cos    = torch.zeros((1, dummy_seq, 1, 1, head_dim), dtype=torch.float32)
        dummy_sin    = torch.zeros((1, dummy_seq, 1, 1, head_dim), dtype=torch.float32)
        dummy_mask   = torch.zeros((1, 1, 1, dummy_seq, dummy_seq), dtype=torch.float32)

        all_inputs  = [dummy_hidden]
        input_names = ["hidden_states"]
        dynamic_axes = {
            "hidden_states":     {0: "batch", 1: "seq_len"},
            "last_hidden_state": {0: "batch", 1: "seq_len"},
            "rotary_cos":        {1: "seq_len"},
            "rotary_sin":        {1: "seq_len"},
            "attention_mask":    {3: "seq_len", 4: "seq_len"},
        }

        for i in range(deepstack_features_len):
            name = f"deepstack_features_{i}"
            input_names.append(name)
            all_inputs.append(dummy_ds)
            dynamic_axes[name] = {1: "seq_len"}

        all_inputs.extend([dummy_cos, dummy_sin, dummy_mask])
        input_names.extend(["rotary_cos", "rotary_sin", "attention_mask"])

        model_main = LLM_MAIN_EMBED(
            model, num_heads, num_key_value_heads, head_dim,
            num_layers, hidden_size, deepstack_features_len,
        )
        del model
        gc.collect()

        torch.onnx.export(
            model_main,
            tuple(all_inputs),
            onnx_transformer,
            input_names=input_names,
            output_names=["last_hidden_state"],
            dynamic_axes=dynamic_axes,
            opset_version=OPSET,
            dynamo=False,
        )
        del model_main, all_inputs
        gc.collect()
        print(f"  → {onnx_transformer}")
        print("  Consolidating external weights …")
        consolidate_external_data(onnx_transformer)

    print("Export complete.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Export Qwen3-VL-Embedding model to ONNX files"
    )
    parser.add_argument(
        "--model_path", type=str, default=DEFAULT_MODEL_PATH,
        help=f"Path to the HF model (default: {DEFAULT_MODEL_PATH})",
    )
    parser.add_argument(
        "--output_dir", type=str, default=DEFAULT_OUTPUT_DIR,
        help=f"Directory to save ONNX files (default: {DEFAULT_OUTPUT_DIR})",
    )
    args = parser.parse_args()

    main(args.model_path, args.output_dir)