文章目录
- 一、量化加速sdxl模型
- 1)涉及模型
- 2)环境安装
- 3)转换模型safetensor to pytorch文件
- 4)tensorRT的环境准备
- (1)下载tensorRT 10.10
- (2)下载cuda
一、量化加速sdxl模型
1)涉及模型
- 涉及模型
waiNSFWIllustrious_v120.safetensors
anyloraCheckpoint_bakedvaeBlessedFp16.safetensors
2)环境安装
torch>=1.9.0
safetensors>=0.3.0
numpy==2.2.6
3)转换模型safetensor to pytorch文件
- 转换指令
python3 safetensor_to_pytorch_converter.py
- 代码
#!/usr/bin/env python3
"""
Safetensor to PyTorch Converter
支持将safetensor格式的模型文件转换为PyTorch格式
"""import os
import sys
import argparse
from pathlib import Path
import torch
from safetensors.torch import load_file, save_file
import logging# 设置日志
logging.basicConfig(level=logging.INFO,format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)def convert_safetensor_to_pytorch(safetensor_path, output_path=None):"""将safetensor文件转换为PyTorch格式Args:safetensor_path (str): safetensor文件路径output_path (str, optional): 输出文件路径,如果为None则自动生成Returns:str: 输出文件路径"""try:# 检查输入文件是否存在if not os.path.exists(safetensor_path):raise FileNotFoundError(f"文件不存在: {safetensor_path}")# 如果未指定输出路径,自动生成if output_path is None:input_path = Path(safetensor_path)output_path = input_path.with_suffix('.pth')logger.info(f"正在加载safetensor文件: {safetensor_path}")# 加载safetensor文件state_dict = load_file(safetensor_path)logger.info(f"模型包含 {len(state_dict)} 个参数")# 保存为PyTorch格式torch.save(state_dict, output_path)logger.info(f"转换完成,保存到: {output_path}")# 显示文件大小信息input_size = os.path.getsize(safetensor_path) / (1024 * 1024) # MBoutput_size = os.path.getsize(output_path) / (1024 * 1024) # MBlogger.info(f"输入文件大小: {input_size:.2f} MB")logger.info(f"输出文件大小: {output_size:.2f} MB")return output_pathexcept Exception as e:logger.error(f"转换失败: {str(e)}")raisedef convert_specific_models():"""转换指定的两个模型文件""" models_to_convert = ["/home/fun/tensor_pic/safe_models/waiNSFWIllustrious_v120.safetensors","/home/fun/tensor_pic/safe_models/anyloraCheckpoint_bakedvaeBlessedFp16.safetensors"]converted_files = []for model_name in models_to_convert:if os.path.exists(model_name):logger.info(f"开始转换: {model_name}")try:output_file = convert_safetensor_to_pytorch(model_name)converted_files.append(output_file)logger.info(f"✅ {model_name} 转换成功")except Exception as e:logger.error(f"❌ {model_name} 转换失败: {str(e)}")else:logger.warning(f"⚠️ 文件不存在: {model_name}")return converted_filesdef main():parser = argparse.ArgumentParser(description='将safetensor文件转换为PyTorch格式')parser.add_argument('--input', '-i', type=str, help='输入的safetensor文件路径')parser.add_argument('--output', '-o', type=str, help='输出的PyTorch文件路径')parser.add_argument('--convert-specific', action='store_true', help='转换指定的两个模型文件')args = parser.parse_args()try:if args.convert_specific:logger.info("开始转换指定的模型文件...")converted_files = convert_specific_models()if converted_files:logger.info(f"成功转换 {len(converted_files)} 个文件:")for file in converted_files:logger.info(f" - {file}")else:logger.warning("没有成功转换任何文件")elif args.input:# 转换单个文件output_file = convert_safetensor_to_pytorch(args.input, args.output)logger.info(f"转换完成: {output_file}")else:# 默认转换指定模型logger.info("未指定参数,开始转换指定的模型文件...")convert_specific_models()except KeyboardInterrupt:logger.info("用户中断操作")sys.exit(1)except Exception as e:logger.error(f"程序执行失败: {str(e)}")sys.exit(1)if __name__ == "__main__":main()
- 输出
(tensor_sd_xl) fun@fun:~/tensor_pic/safe2pytorch_tool$ python3 safetensor_to_pytorch_converter.py
2025-07-31 13:02:27,927 - INFO - 未指定参数,开始转换指定的模型文件...
2025-07-31 13:02:27,927 - INFO - 开始转换: /home/fun/tensor_pic/safe_models/waiNSFWIllustrious_v120.safetensors
2025-07-31 13:02:27,927 - INFO - 正在加载safetensor文件: /home/fun/tensor_pic/safe_models/waiNSFWIllustrious_v120.safetensors
2025-07-31 13:02:27,945 - INFO - 模型包含 2515 个参数
2025-07-31 13:02:31,083 - INFO - 转换完成,保存到: /home/fun/tensor_pic/safe_models/waiNSFWIllustrious_v120.pth
2025-07-31 13:02:31,084 - INFO - 输入文件大小: 6616.63 MB
2025-07-31 13:02:31,084 - INFO - 输出文件大小: 6617.17 MB
2025-07-31 13:02:31,084 - INFO - ✅ /home/fun/tensor_pic/safe_models/waiNSFWIllustrious_v120.safetensors 转换成功
2025-07-31 13:02:31,084 - INFO - 开始转换: /home/fun/tensor_pic/safe_models/anyloraCheckpoint_bakedvaeBlessedFp16.safetensors
2025-07-31 13:02:31,085 - INFO - 正在加载safetensor文件: /home/fun/tensor_pic/safe_models/anyloraCheckpoint_bakedvaeBlessedFp16.safetensors
2025-07-31 13:02:31,092 - INFO - 模型包含 1143 个参数
2025-07-31 13:02:32,022 - INFO - 转换完成,保存到: /home/fun/tensor_pic/safe_models/anyloraCheckpoint_bakedvaeBlessedFp16.pth
2025-07-31 13:02:32,022 - INFO - 输入文件大小: 2033.85 MB
2025-07-31 13:02:32,022 - INFO - 输出文件大小: 2034.11 MB
2025-07-31 13:02:32,023 - INFO - ✅ /home/fun/tensor_pic/safe_models/anyloraCheckpoint_bakedvaeBlessedFp16.safetensors 转换成功
(tensor_sd_xl) fun@fun:~/tensor_pic/safe2pytorch_tool$
4)tensorRT的环境准备
(1)下载tensorRT 10.10
[链接]
命令安装这个包
1、安装包
sudo dpkg -i nv-tensorrt-local-repo-ubuntu2204-10.10.0-cuda-12.9_1.0-1_amd64.deb
2、查看 TensorRT 包是否可用:如果能看到 libnvinfer、libnvinfer-dev 等相关包,说明源添加成功
apt-cache search tensorrt
如果没有输出,说明还没安装。
3、安装 TensorRT:
sudo apt update
sudo apt install libnvinfer-dev libnvinfer-plugin-dev python3-tensorrt
4、安装完成后再用第2步命令检查
apt-cache search tensorrt
验证 Python TensorRT 是否可用
python3 -c "import tensorrt; print(tensorrt.__version__)"
如果能输出版本号,说明 Python TensorRT 安装成功。
(2)下载cuda
12.9 cuda链接
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/12.9.1/local_installers/cuda-repo-ubuntu2204-12-9-local_12.9.1-575.57.08-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-12-9-local_12.9.1-575.57.08-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-12-9-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda-toolkit-12-9