Upload convert_precision.py

convert_precision.py

@@ -1,56 +1,60 @@
-import torch
-import os
-
-# --- 配置 ---
-# 1. 设置目标精度: 'fp16' 或 'bf16'
-TARGET_PRECISION = 'bf16' 
-
-# 2. 设置原始32位训练检查点文件的路径
-fp32_checkpoint_path = r"E:\comfyui\ComfyUI-aki-v1.3\models\SDMatte\SDMatte_plus.pth"
-
-# --------------------------------------------------------------------
-
-# 自动生成输出文件名
-output_filename = fp32_checkpoint_path.replace('.pth', f'_{TARGET_PRECISION}_inference.pth')
-
-if not os.path.exists(fp32_checkpoint_path):
-    print(f"[错误] 文件不存在: {fp32_checkpoint_path}")
-else:
-    try:
-        print(f"--- 开始处理训练检查点: {fp32_checkpoint_path} ---")
-        full_checkpoint = torch.load(fp32_checkpoint_path, map_location="cpu", weights_only=False)
-
-        # 检查 'model' 键是否存在，这是包含我们所需权重的部分
-        if 'model' in full_checkpoint:
-            # 明确提取出模型的 state_dict
-            state_dict = full_checkpoint['model']
-            print("成功提取到 'model' 键中的权重字典。")
-        else:
-            # 如果没有 'model' 键，则假定整个文件就是 state_dict
-            print("[警告] 未在顶层找到 'model' 键，将尝试转换整个文件。")
-            state_dict = full_checkpoint
-        
-        print(f"开始将权重转换为 {TARGET_PRECISION} ...")
-        
-        target_dtype = torch.float16 if TARGET_PRECISION == 'fp16' else torch.bfloat16
-
-        # 遍历权重字典中的每一项并进行转换
-        for key in state_dict:
-            if isinstance(state_dict[key], torch.Tensor) and state_dict[key].is_floating_point():
-                state_dict[key] = state_dict[key].to(target_dtype)
-
-        print(f"正在保存纯推理模型到: {output_filename} ...")
-        # 直接保存处理后的 state_dict，不包含任何训练相关的附加信息
-        torch.save(state_dict, output_filename)
-        
-        # 打印大小对比
-        original_size_gb = os.path.getsize(fp32_checkpoint_path) / (1024**3)
-        final_size_gb = os.path.getsize(output_filename) / (1024**3)
-
-        print("\n--- 转换成功 ---")
-        print(f"原始训练检查点大小: {original_size_gb:.2f} GB")
-        print(f"生成的纯推理模型大小 ({TARGET_PRECISION.upper()}): {final_size_gb:.2f} GB")
-        print("说明: 新文件只包含用于推理的核心模型权重，已移除训练相关的优化器状态。")
-
-    except Exception as e:
+import torch
+import os
+
+# --- 配置 ---
+# 1. 设置目标精度: 'fp32', 'fp16', 或 'bf16'
+#    - 'fp32': 移除训练数据，保留32位全精度。
+#    - 'fp16': 转换为16位半精度，体积最小。
+#    - 'bf16': 转换为bfloat16精度，推荐RTX 30系及以上GPU。
+TARGET_PRECISION = 'fp32' 
+
+# 2. 设置原始32位训练检查点文件的路径,例如"E:\comfyui\ComfyUI-aki-v1.3\models\SDMatte\1\SDMatte_plus.pth"
+fp32_checkpoint_path = r"E:\comfyui\ComfyUI-aki-v1.3\models\SDMatte\1\SDMatte_plus.pth"
+
+# --------------------------------------------------------------------
+
+# 自动生成输出文件名
+output_filename = fp32_checkpoint_path.replace('.pth', f'_{TARGET_PRECISION}_inference.pth')
+
+if not os.path.exists(fp32_checkpoint_path):
+    print(f"[错误] 文件不存在: {fp32_checkpoint_path}")
+else:
+    try:
+        print(f"--- 开始处理训练检查点: {fp32_checkpoint_path} ---")
+        full_checkpoint = torch.load(fp32_checkpoint_path, map_location="cpu", weights_only=False)
+
+        if 'model' in full_checkpoint:
+            state_dict = full_checkpoint['model']
+            print("成功提取到 'model' 键中的权重字典。")
+        else:
+            print("[警告] 未在顶层找到 'model' 键，将尝试转换整个文件。")
+            state_dict = full_checkpoint
+        
+        print(f"目标输出精度: {TARGET_PRECISION}")
+
+        # --- MODIFICATION START ---
+        # 仅当目标精度不是 'fp32' 时，才执行类型转换
+        if TARGET_PRECISION != 'fp32':
+            print(f"开始将权重转换为 {TARGET_PRECISION} ...")
+            target_dtype = torch.float16 if TARGET_PRECISION == 'fp16' else torch.bfloat16
+
+            for key in state_dict:
+                if isinstance(state_dict[key], torch.Tensor) and state_dict[key].is_floating_point():
+                    state_dict[key] = state_dict[key].to(target_dtype)
+        else:
+            print("保留原始 FP32 精度，仅剥离训练数据。")
+        # --- MODIFICATION END ---
+
+        print(f"正在保存纯推理模型到: {output_filename} ...")
+        torch.save(state_dict, output_filename)
+        
+        original_size_gb = os.path.getsize(fp32_checkpoint_path) / (1024**3)
+        final_size_gb = os.path.getsize(output_filename) / (1024**3)
+
+        print("\n--- 转换成功 ---")
+        print(f"原始训练检查点大小: {original_size_gb:.2f} GB")
+        print(f"生成的纯推理模型大小 ({TARGET_PRECISION.upper()}): {final_size_gb:.2f} GB")
+        print("说明: 新文件只包含用于推理的核心模型权重，已移除训练相关的优化器状态。")
+
+    except Exception as e:
         print(f"\n[错误] 处理过程中发生错误: {e}")