import os
import argparse
import torch

def convert_pt_to_fp8(input_file, output_file):
    try:
        # 获取当前运行目录的绝对路径，锁定范围
        current_dir = os.getcwd()
        
        in_path = os.path.abspath(os.path.join(current_dir, input_file))
        out_path = os.path.abspath(os.path.join(current_dir, output_file))
        
        # 严格限制：只允许处理当前目录下的文件
        if os.path.dirname(in_path) != current_dir or os.path.dirname(out_path) != current_dir:
            print("❌ 错误：基于安全限制，本脚本只允许处理当前目录下的文件！")
            return

        if not os.path.exists(in_path):
            print(f"❌ 错误：找不到输入文件 '{input_file}'")
            return

        print(f"📦 正在加载 FP16 模型权重: {input_file} ...")
        # 视你的硬件情况，可将 map_location 设为 'cuda' 或 'cpu'
        state_dict = torch.load(in_path, map_location='cpu')

        print("⚡ 正在将浮点数转换为底层 FP8 (E4M3FN) 格式...")
        fp8_state_dict = {}
        
        # 判断载入的是直接的 OrderedDict 还是包裹在字典里的权重
        if isinstance(state_dict, dict) and 'model' in state_dict:
            target_dict = state_dict['model']
            is_wrapped = True
        else:
            target_dict = state_dict
            is_wrapped = False

        for k, v in target_dict.items():
            if isinstance(v, torch.Tensor):
                # 只有浮点数（FP32/FP16）需要转为 FP8，整数（如层索引、形状）原样保留
                if v.is_floating_point():
                    # 转换为 PyTorch 标准的 FP8 权重存储格式
                    fp8_state_dict[k] = v.to(dtype=torch.float8_e4m3fn)
                else:
                    fp8_state_dict[k] = v
            else:
                fp8_state_dict[k] = v

        # 保持原有的包裹外壳结构
        output_data = {'model': fp8_state_dict} if is_wrapped else fp8_state_dict

        print(f"💾 正在保存 FP8 模型至: {output_file} ...")
        torch.save(output_data, out_path)
        print(f"✅ 成功生成 FP8 权重文件！")

    except Exception as e:
        print(f"❌ 转换失败，错误原因: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="将 FP16/FP32 的 .pt 权重转换为 FP8(E4M3) 格式（仅限当前目录）")
    parser.add_argument("-i", "--input", required=True, help="输入的 .pt 权重文件名")
    parser.add_argument("-o", "--output", required=True, help="输出的 FP8 .pt 文件名")
    
    args = parser.parse_args()
    convert_pt_to_fp8(args.input, args.output)