VideoX-Fun/infer_split_axmodels.py

#!/usr/bin/env python3
"""串行执行切分后的 AXModel 子图以复现完整推理路径。"""
from __future__ import annotations

import argparse
import json
import logging
import sys
import types
from pathlib import Path
from typing import Dict, List

import numpy as np
import onnx
from axengine import InferenceSession

PROJECT_ROOT = Path(__file__).resolve().parent
SCRIPTS_DIR = PROJECT_ROOT / "scripts"
if SCRIPTS_DIR.as_posix() not in sys.path:
    sys.path.insert(0, SCRIPTS_DIR.as_posix())

from split_onnx_by_subconfigs import (  # type: ignore
    SubGraphSpec,
    build_graph_index,
    derive_interface,
    ordered_specs,
    sanitize,
    trace_nodes_between,
    untouched_components,
)


def _ensure_numpy_core_alias() -> None:
    """兼容 numpy>=2 保存、numpy<2 读取时的 pickle 模块路径差异。"""
    if "numpy._core" in sys.modules:
        return
    try:
        core_module = sys.modules.get("numpy.core") or __import__("numpy.core")
    except Exception:  # pragma: no cover - 极端环境
        return
    alias = types.ModuleType("numpy._core")
    alias.__dict__.update(core_module.__dict__)
    sys.modules["numpy._core"] = alias
    submods = {
        "multiarray": getattr(core_module, "multiarray", None),
        "umath": getattr(core_module, "umath", None),
        "numerictypes": getattr(core_module, "numerictypes", None),
        "_multiarray_umath": getattr(core_module, "_multiarray_umath", None),
    }
    for name, module in submods.items():
        if module is not None:
            sys.modules[f"numpy._core.{name}"] = module


def load_specs(config_path: Path, onnx_path: Path) -> List[SubGraphSpec]:
    with config_path.open("r", encoding="utf-8") as f:
        config = json.load(f)
    sub_configs = config.get("compiler", {}).get("sub_configs", [])
    if not sub_configs:
        raise ValueError("配置文件中缺少 compiler.sub_configs")

    model = onnx.load(onnx_path.as_posix())
    index = build_graph_index(model)

    specs: List[SubGraphSpec] = []
    covered_nodes = set()
    for idx, entry in enumerate(sub_configs):
        start = [name for name in entry.get("start_tensor_names", []) if name]
        end = [name for name in entry.get("end_tensor_names", []) if name]
        if not start or not end:
            raise ValueError(f"sub_config[{idx}] 缺少 tensor 名称")
        spec = SubGraphSpec(
            label=f"cfg_{idx:02d}",
            start=start,
            end=end,
            node_names=set(),
            source="config",
        )
        nodes = trace_nodes_between(spec, index)
        spec.node_names = nodes
        covered_nodes.update(nodes)
        specs.append(spec)

    leftovers = untouched_components(index.node_order, covered_nodes, index)
    for idx, component in enumerate(leftovers):
        start, end = derive_interface(component, index)
        if not end:
            continue
        spec = SubGraphSpec(
            label=f"auto_{idx:02d}",
            start=start,
            end=end,
            node_names=component,
            source="auto",
        )
        specs.append(spec)

    return ordered_specs(specs, index)


def expected_model_path(spec: SubGraphSpec, model_dir: Path) -> Path:
    head = sanitize(spec.start[0]) if spec.start else "const"
    tail = sanitize(spec.end[0]) if spec.end else "out"
    filename = f"{spec.label}_{head}_to_{tail}_{spec.source}.axmodel"
    path = model_dir / filename
    if not path.exists():
        raise FileNotFoundError(f"未找到模型文件: {path}")
    return path


def run_pipeline(
    specs: List[SubGraphSpec],
    model_dir: Path,
    feed_dict: Dict[str, np.ndarray],
) -> Dict[str, np.ndarray]:
    tensor_store: Dict[str, np.ndarray] = dict(feed_dict)
    for spec in specs:
        model_path = expected_model_path(spec, model_dir)
        session = InferenceSession(model_path.as_posix())
        inputs = {}
        missing = []
        for name in spec.start:
            value = tensor_store.get(name)
            if value is None:
                missing.append(name)
            else:
                inputs[name] = value
        if missing:
            raise KeyError(f"模型 {model_path.name} 缺少输入张量: {missing}")
        results = session.run(spec.end, inputs)
        if not isinstance(results, (list, tuple)):
            results = [results]
        for out_name, value in zip(spec.end, results):
            tensor_store[out_name] = value
        shapes = {out_name: tuple(value.shape) for out_name, value in zip(spec.end, results)}
        logging.info("完成 %s: %s", model_path.name, shapes)
    return tensor_store


def load_input_file(path: Path, single_name: str | None) -> Dict[str, np.ndarray]:
    suffix = path.suffix.lower()
    if suffix == ".npz":
        data = np.load(path, allow_pickle=False)
        return {key: np.asarray(data[key]) for key in data.files}

    if suffix == ".npy":
        _ensure_numpy_core_alias()
        blob = np.load(path, allow_pickle=True)
        if isinstance(blob, np.ndarray) and blob.dtype == object:
            obj = blob.item()
            if isinstance(obj, dict):
                return {str(k): np.asarray(v) for k, v in obj.items()}
            if single_name is None:
                raise ValueError("单一 .npy 输入需要通过 --single-input-name 指定张量名")
            return {single_name: np.asarray(obj)}
        if single_name is None:
            raise ValueError("单一 .npy 输入需要通过 --single-input-name 指定张量名")
        return {single_name: np.asarray(blob)}

    raise ValueError(f"暂不支持的输入文件格式: {path}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="串行执行切分后的 AXModel")
    parser.add_argument("--onnx", required=True, type=Path, help="原始未切分 ONNX")
    parser.add_argument("--config", required=True, type=Path, help="sub_config 配置文件")
    parser.add_argument("--model-dir", required=True, type=Path, help=".axmodel 所在目录")
    parser.add_argument(
        "--input-file",
        required=True,
        type=Path,
        help="包含原始输入张量的 npz 或 npy 文件",
    )
    parser.add_argument(
        "--single-input-name",
        type=str,
        help="当 --input-file 是单张量 .npy 文件时指定对应的输入名",
    )
    parser.add_argument("--dump-npz", type=Path, help="可选，将最终输出保存为 npz")
    parser.add_argument("--log", default="INFO", help="日志等级")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    logging.basicConfig(level=getattr(logging, args.log.upper(), logging.INFO))

    specs = load_specs(args.config, args.onnx)
    feed = load_input_file(args.input_file, args.single_input_name)
    store = run_pipeline(specs, args.model_dir, feed)

    final_outputs = {name: store[name] for name in specs[-1].end}
    for name, value in final_outputs.items():
        logging.info("最终输出 %s: shape=%s", name, value.shape)

    if args.dump_npz:
        np.savez(args.dump_npz, **final_outputs)
        logging.info("已保存输出到 %s", args.dump_npz)


if __name__ == "__main__":
    """
    python3 infer_split_axmodels.py \
        --onnx onnx-models/transformer.onnx \
        --config pulsar2_configs/transformers_subgraph.json \
        --model-dir compiled_slice_quant_onnx \
        --input-file onnx-calibration-no-controlnet/transformer_inputs_prompt000_step00.npy \
        --single-input-name timestep
    """
    main()