scripts/export.py

#!/usr/bin/env python3
"""
模型导出脚本
用于导出训练好的模型为不同格式
"""

import os
import sys
import argparse
import yaml
import torch
import torch.nn as nn

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(__file__)))

from src.models.unet_light import UNetLight
from src.models.diffusion import DiffusionProcess
from src.inference.optimization import ModelOptimizer, ONNXExporter, optimize_model_for_p4


def load_config(config_path: str) -> dict:
    """加载配置文件"""
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    return config


def load_model(checkpoint_path: str, config: dict, device: torch.device) -> nn.Module:
    """加载模型"""
    # 加载模型配置
    model_config_path = config.get('model_config', 'configs/model/unet_light.yaml')
    model_config = load_config(model_config_path)
    
    # 创建模型
    model = UNetLight(model_config)
    
    # 加载检查点
    print(f"加载检查点: {checkpoint_path}")
    checkpoint = torch.load(checkpoint_path, map_location='cpu')
    
    # 加载模型权重
    if 'model_state_dict' in checkpoint:
        model.load_state_dict(checkpoint['model_state_dict'])
    elif 'state_dict' in checkpoint:
        model.load_state_dict(checkpoint['state_dict'])
    else:
        model.load_state_dict(checkpoint)
    
    # 移动到设备
    model = model.to(device)
    model.eval()
    
    print(f"模型加载完成")
    
    return model


def export_torchscript(model: nn.Module, output_path: str):
    """导出为TorchScript格式"""
    print(f"导出为TorchScript: {output_path}")
    
    # 创建示例输入
    example_input = torch.randn(1, model.in_channels, 64, 64)
    example_timestep = torch.tensor([500])
    example_context = torch.randn(1, 77, 768)
    
    # 跟踪模型
    traced_model = torch.jit.trace(
        model,
        (example_input, example_timestep, example_context),
        check_trace=False
    )
    
    # 保存
    traced_model.save(output_path)
    print(f"TorchScript模型已保存: {output_path}")
    
    return traced_model


def export_onnx(model: nn.Module, output_path: str, opset_version: int = 14):
    """导出为ONNX格式"""
    print(f"导出为ONNX: {output_path}")
    
    # 创建示例输入
    example_input = torch.randn(1, model.in_channels, 64, 64)
    example_timestep = torch.tensor([500])
    example_context = torch.randn(1, 77, 768)
    
    # 设置动态轴
    dynamic_axes = {
        'input': {0: 'batch_size'},
        'timestep': {0: 'batch_size'},
        'context': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
    
    # 导出
    torch.onnx.export(
        model,
        (example_input, example_timestep, example_context),
        output_path,
        input_names=['input', 'timestep', 'context'],
        output_names=['output'],
        dynamic_axes=dynamic_axes,
        opset_version=opset_version,
        do_constant_folding=True,
        verbose=False
    )
    
    print(f"ONNX模型已保存: {output_path}")
    
    # 验证ONNX模型
    import onnx
    onnx_model = onnx.load(output_path)
    onnx.checker.check_model(onnx_model)
    print("ONNX模型验证成功")


def export_safetensors(model: nn.Module, output_path: str):
    """导出为safetensors格式"""
    try:
        from safetensors.torch import save_file
        
        # 转换为safetensors格式
        state_dict = model.state_dict()
        save_file(state_dict, output_path)
        
        print(f"Safetensors模型已保存: {output_path}")
    except ImportError:
        print("safetensors未安装，跳过safetensors导出")
        print("安装: pip install safetensors")


def optimize_and_export(
    checkpoint_path: str,
    output_dir: str,
    formats: list = ['torchscript', 'onnx', 'safetensors'],
    optimize_for_p4: bool = True
):
    """优化并导出模型"""
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 加载配置
    config_path = "configs/training/p4_optimized.yaml"
    config = load_config(config_path)
    
    # 设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 加载模型
    model = load_model(checkpoint_path, config, device)
    
    # 优化模型（针对P4）
    if optimize_for_p4:
        print("优化模型（针对P4）...")
        model = optimize_model_for_p4(model)
    
    # 获取模型信息
    total_params = sum(p.numel() for p in model.parameters())
    model_size_mb = total_params * 4 / 1024**2  # fp32
    
    print(f"\n模型信息:")
    print(f"  参数量: {total_params:,}")
    print(f"  模型大小: {model_size_mb:.2f} MB (fp32)")
    
    # 导出为不同格式
    base_name = os.path.splitext(os.path.basename(checkpoint_path))[0]
    
    for fmt in formats:
        if fmt == 'torchscript':
            output_path = os.path.join(output_dir, f"{base_name}.torchscript.pt")
            export_torchscript(model, output_path)
        
        elif fmt == 'onnx':
            output_path = os.path.join(output_dir, f"{base_name}.onnx")
            export_onnx(model, output_path)
        
        elif fmt == 'safetensors':
            output_path = os.path.join(output_dir, f"{base_name}.safetensors")
            export_safetensors(model, output_path)
        
        elif fmt == 'pth':
            output_path = os.path.join(output_dir, f"{base_name}.pth")
            torch.save(model.state_dict(), output_path)
            print(f"PyTorch模型已保存: {output_path}")
        
        else:
            print(f"未知的格式: {fmt}")
    
    print(f"\n所有模型已导出到: {output_dir}")


def create_lite_version(model: nn.Module, reduction_factor: float = 0.5) -> nn.Module:
    """创建轻量版本（通过减少通道数）"""
    # 注意：这是一个示例，需要根据实际模型结构调整
    print(f"创建轻量版本，减少因子: {reduction_factor}")
    
    # 这里应该实现具体的轻量化逻辑
    # 例如，减少UNet的通道数
    
    return model


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="导出Lumina模型")
    
    parser.add_argument(
        "--checkpoint",
        type=str,
        required=True,
        help="模型检查点路径"
    )
    
    parser.add_argument(
        "--output-dir",
        type=str,
        default="./exported_models",
        help="输出目录"
    )
    
    parser.add_argument(
        "--formats",
        type=str,
        nargs="+",
        default=['torchscript', 'onnx'],
        choices=['torchscript', 'onnx', 'safetensors', 'pth'],
        help="导出格式"
    )
    
    parser.add_argument(
        "--optimize",
        action="store_true",
        help="优化模型（针对P4）"
    )
    
    parser.add_argument(
        "--lite",
        action="store_true",
        help="创建轻量版本"
    )
    
    parser.add_argument(
        "--lite-factor",
        type=float,
        default=0.5,
        help="轻量化减少因子"
    )
    
    args = parser.parse_args()
    
    # 检查输入文件
    if not os.path.exists(args.checkpoint):
        print(f"错误: 检查点文件不存在: {args.checkpoint}")
        return
    
    # 优化并导出
    optimize_and_export(
        checkpoint_path=args.checkpoint,
        output_dir=args.output_dir,
        formats=args.formats,
        optimize_for_p4=args.optimize
    )


if __name__ == "__main__":
    main()