export_onnx.py

import argparse
import logging
import os
from pathlib import Path

import warnings
import torch
from sharp.models import PredictorParams, create_predictor

LOGGER = logging.getLogger(__name__)

# 默认权重文件名
DEFAULT_CHECKPOINT_NAME = "sharp_2572gikvuh.pt"
MODEL_URL = f"https://ml-site.cdn-apple.com/models/sharp/{DEFAULT_CHECKPOINT_NAME}"

def verify_model_precision(model_path: Path):
    """
    静态检测 ONNX 模型的权重数据类型分布，用于验证量化是否生效。
    """
    try:
        import onnx
        from collections import Counter
        # 尝试加载模型，注意大模型可能带 .data 后缀
        model = onnx.load(str(model_path))
        # 统计所有初始化器（权重）的数据类型
        # 在 ONNX 中，FLOAT=1, UINT8=2, INT8=3, FLOAT16=10, UINT4=27, INT4=28, FLOAT4=29
        dtypes = [onnx.TensorProto.DataType.Name(i.data_type) for i in model.graph.initializer]
        counter = Counter(dtypes)
        
        LOGGER.info(f"--- 权重类型分布报告 [{model_path.name}] ---")
        for dtype, count in counter.items():
            LOGGER.info(f" - {dtype}: {count} 个张量")
        
        # 计算文件大小并打印
        file_size = os.path.getsize(model_path)
        LOGGER.info(f" [文件信息] 路径: {model_path}")
        LOGGER.info(f" [文件信息] 体积: {file_size / (1024*1024):.2f} MB")
        LOGGER.info("------------------------------------------------")
    except ImportError:
        LOGGER.warning("未安装 'onnx' 模块，跳过权重精度验证。")
    except Exception as e:
        LOGGER.warning(f"验证模型精度时发生异常: {e}")

def export_onnx(output_path: Path, checkpoint_path: str = None, internal_shape: int = 1536, is_fp16: bool = False):
    device = torch.device("cpu")
    
    LOGGER.info("正在创建预测器模型...")
    predictor = create_predictor(PredictorParams())
    
    # FP16 处理：将网络参数转为半精度，体积减半
    if is_fp16:
        LOGGER.info("正在将模型转换为 FP16 (半精度) 以直接将其缩减进单一 < 2GB 的文件中...")
        predictor = predictor.half()
    
    if checkpoint_path is None:
        torch_hub_checkpoints = Path(torch.hub.get_dir()) / "checkpoints"
        search_paths = [
            Path(DEFAULT_CHECKPOINT_NAME),
            Path("data") / DEFAULT_CHECKPOINT_NAME,
            torch_hub_checkpoints / DEFAULT_CHECKPOINT_NAME,
        ]
        
        found_path = None
        for p in search_paths:
            if p.exists():
                found_path = p
                break
        
        if found_path:
            LOGGER.info(f"在本地找到模型权重: {found_path}，加载本地权重...")
            state_dict = torch.load(found_path, map_location="cpu", weights_only=True)
        else:
            state_dict = torch.hub.load_state_dict_from_url(MODEL_URL, progress=True, map_location="cpu")
    else:
        state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
        
    predictor.load_state_dict(state_dict)
    predictor.eval()
    predictor.to(device)
    
    dummy_image = torch.randn(1, 3, internal_shape, internal_shape, device=device)
    dummy_disparity = torch.tensor([1.0], device=device)
    
    if is_fp16:
        dummy_image = dummy_image.half()
        dummy_disparity = dummy_disparity.half()
        
    dummy_inputs = (dummy_image, dummy_disparity)
    
    LOGGER.info("正在导出为 ONNX 格式...")
    torch.onnx.export(
        predictor,
        dummy_inputs,
        str(output_path),
        export_params=True,
        opset_version=19,
        do_constant_folding=True,
        input_names=['image', 'disparity_factor'],
        output_names=['mean_vectors', 'singular_values', 'quaternions', 'colors', 'opacities'],
        keep_initializers_as_inputs=False,
        dynamic_axes={k: {0: 'batch_size'} for k in ['image', 'disparity_factor', 'mean_vectors', 'singular_values', 'quaternions', 'colors', 'opacities']}
    )
    LOGGER.info(f"成功将全精度 ONNX 模型导出至: {output_path}")

def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
    parser = argparse.ArgumentParser(description="将 ML-SHARP 模型导出为多种精度的 ONNX 格式。")
    parser.add_argument("--output", type=str, default="sharp.onnx", help="导出的基础 ONNX 模型保存路径。")
    parser.add_argument("--checkpoint", type=str, default=None, help="本地 .pt 权重文件路径。")
    parser.add_argument("--fp16", action="store_true", help="启用 FP16 半精度导出。")
    parser.add_argument("--int8", action="store_true", help="应用动态 INT8 量化。")
    parser.add_argument("--int4", action="store_true", help="应用真正的 INT4 重量级量化 (WOQ)。")
    args = parser.parse_args()

    # 如果启用了 FP16 模式且用户未指定 output，则重定向默认输出名
    if args.fp16 and args.output == "sharp.onnx":
        args.output = "sharp_fp16.onnx"
        
    output_path = Path(args.output)
    is_quant_requested = args.int8 or args.int4
    
    if is_quant_requested and output_path.exists():
        LOGGER.info(f"检测到基础模型 '{output_path}' 已存在，跳过导出并开始验证。")
        verify_model_precision(output_path)
    else:
        export_onnx(output_path, args.checkpoint, is_fp16=args.fp16)
        verify_model_precision(output_path)

    if is_quant_requested:
        try:
            from onnxruntime.quantization import quantize_dynamic, QuantType
            import onnx
            from onnx import shape_inference
        except ImportError:
            LOGGER.error("缺失依赖项：请安装 onnx 和 onnxruntime。")
            return

        preprocessed_path = output_path.with_name(f"{output_path.stem}_pre.onnx")
        if not preprocessed_path.exists():
            LOGGER.info("正在将基础模型转换为外部数据格式 (已跳过导致崩溃的形状推导模式)...")
            try:
                model_raw = onnx.load(str(output_path))
                onnx.save(
                    model_raw, 
                    str(preprocessed_path),
                    save_as_external_data=True,
                    all_tensors_to_one_file=True,
                    location=preprocessed_path.name + ".data",
                    size_threshold=1024,
                    convert_attribute=True
                )
                model_source = preprocessed_path
            except Exception as e:
                LOGGER.warning(f"预分流模型保存失败: {e}")
                model_source = output_path
        else:
            model_source = preprocessed_path

        tasks = []
        if args.int8: tasks.append("int8")
        if args.int4: tasks.append("int4")

        # 1. INT8 量化
        int8_path = output_path.with_name(f"{output_path.stem}_int8.onnx")
        if "int8" in tasks:
            LOGGER.info("正在应用动态 INT8 量化...")
            quantize_dynamic(str(model_source), str(int8_path), weight_type=QuantType.QUInt8)
            verify_model_precision(int8_path)

        # 2. INT4 重量级量化 (WOQ)
        if "int4" in tasks:
            LOGGER.info("正在准备进行真正的 INT4 权重压缩 (WOQ)...")
            quantized_output = output_path.with_name(f"{output_path.stem}_int4.onnx")
            
            import onnxruntime.quantization as oq
            import importlib
            import pkgutil

            # 优先尝试新版 Unified API
            quantize_func = getattr(oq, "quantize", None)
            WeightOnlyConfig = getattr(oq, "MatMulWeightOnlyQuantConfig", None)
            
            if quantize_func and WeightOnlyConfig:
                LOGGER.info("[INT4] 使用 MatMulWeightOnlyQuantConfig 执行量化...")
                try:
                    model = onnx.load(str(model_source))
                    config = WeightOnlyConfig(weight_type=getattr(QuantType, 'QInt4', QuantType.QUInt8))
                    quantize_func(model, str(quantized_output), config)
                    LOGGER.info(f"[INT4] Unified API 导出成功!")
                    verify_model_precision(quantized_output)
                    return
                except Exception as e:
                    LOGGER.warning(f"[INT4] Unified API 失败: {e}")

            # Fallback 1: 尝试 Bnb4Quantizer
            Bnb4Quantizer = None
            try:
                from onnxruntime.quantization.matmul_bnb4_quantizer import MatMulBnb4Quantizer
                Bnb4Quantizer = MatMulBnb4Quantizer
            except: pass

            if Bnb4Quantizer:
                LOGGER.info("[INT4] 使用 MatMulBnb4Quantizer (BnB 4-bit) 执行量化...")
                try:
                    import inspect
                    model = onnx.load(str(model_source))
                    sig = inspect.signature(Bnb4Quantizer.__init__)
                    params = sig.parameters
                    kwargs = {"model": model, "block_size": 32}
                    if "quant_type" in params:
                        kwargs["quant_type"] = 1
                    
                    q = Bnb4Quantizer(**kwargs)
                    q.process()
                    m = q.model.model if hasattr(q, 'model') and hasattr(q.model, 'model') else (q.model if hasattr(q, 'model') else q)
                    LOGGER.info(f"正在保存单体 INT4 模型...")
                    onnx.save(m, str(quantized_output))
                    LOGGER.info(f"[INT4] BnB 导出成功!")
                    verify_model_precision(quantized_output)
                    return
                except Exception as e:
                    LOGGER.warning(f"[INT4] BnB 失败: {e}")

            # Fallback 2: 暴力搜索兼容旧版 Class API
            def search_legacy(pkg):
                for attr in dir(pkg):
                    low = attr.lower()
                    if ("quantiz" in low or "quantis" in low) and ("4bit" in low or "nbit" in low or "weight" in low):
                        return getattr(pkg, attr), pkg.__name__
                if hasattr(pkg, "__path__"):
                    for _, modname, _ in pkgutil.walk_packages(pkg.__path__, pkg.__name__ + "."):
                        try:
                            mod = importlib.import_module(modname)
                            res, src = search_legacy(mod)
                            if res: return res, src
                        except: continue
                return None, None

            LegacyQuantizer, src = search_legacy(oq)
            if LegacyQuantizer and (isinstance(LegacyQuantizer, type) or callable(LegacyQuantizer)):
                LOGGER.info(f"[INT4] 找到量化类/函数: {src}.{LegacyQuantizer.__name__}")
                try:
                    model = onnx.load(str(model_source))
                    if "Quantizer" in LegacyQuantizer.__name__:
                        q = LegacyQuantizer(model, block_size=32, is_symmetric=True)
                        q.process()
                        m = q.model.model if hasattr(q, 'model') and hasattr(q.model, 'model') else (q.model if hasattr(q, 'model') else q)
                        onnx.save(m, str(quantized_output))
                    else:
                        LegacyQuantizer(model, str(quantized_output), weight_type=QuantType.QInt4)
                    LOGGER.info(f"[INT4] 导出成功!")
                    verify_model_precision(quantized_output)
                except Exception as e:
                    LOGGER.error(f"[INT4] 执行失败: {e}")
            else:
                LOGGER.error(f"[INT4] 未找到有效量化器。所有属性: {dir(oq)}")

if __name__ == "__main__":
    main()