文章目錄
- 一、量化加速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