Upload 2 files

null_embedding.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:598bbf18a1568585795933f063f6d2bb71ae6b92715b69425690698f7282494e
+size 316607

prune_model.py

@@ -0,0 +1,94 @@
+import torch
+import os
+from collections import OrderedDict
+
+def prune_and_convert_model():
+    """
+    交互式脚本，用于加载PyTorch检查点，移除trainer状态和text_encoder权重，
+    并根据用户选择的目标精度（fp32, fp16, bf16）保存。
+    """
+    # --- 交互式输入 ---
+    source_path = input("请输入源模型文件的完整路径（例如 D:\\models\\SDMatte_plus.pth）：\n> ")
+    
+    if not os.path.exists(source_path):
+        print(f"\n[错误] 文件未找到: {source_path}")
+        return
+
+    while True:
+        target_precision = input("\n请选择目标精度 (输入 fp32, fp16, 或 bf16):\n> ").lower()
+        if target_precision in ['fp32', 'fp16', 'bf16']:
+            break
+        else:
+            print("[错误] 无效输入，请输入 'fp32', 'fp16', 或 'bf16' 中的一个。")
+
+    # --- 核心逻辑 ---
+    # 自动生成输出文件名
+    base, ext = os.path.splitext(source_path)
+    output_path = f"{base}_pruned_{target_precision}{ext}"
+    
+    print(f"\n--- 开始处理检查点: {os.path.basename(source_path)} ---")
+    print(f"目标精度: {target_precision.upper()}")
+    print(f"输出文件将保存为: {os.path.basename(output_path)}")
+
+    try:
+        print("\n正在加载原始权重文件...")
+        full_checkpoint = torch.load(source_path, map_location="cpu", weights_only=False)
+        
+        state_dict = None
+        if isinstance(full_checkpoint, dict):
+            if 'model' in full_checkpoint:
+                state_dict = full_checkpoint['model']
+                print("检测到训练检查点，已从 'model' 键中提取权重。")
+            elif 'state_dict' in full_checkpoint:
+                 state_dict = full_checkpoint['state_dict']
+                 print("检测到训练检查点，已从 'state_dict' 键中提取权重。")
+            else:
+                state_dict = full_checkpoint
+                print("文件似乎是一个纯权重字典，将直接处理。")
+        else:
+             print("[错误] 文件格式无法识别。")
+             return
+
+        pruned_state_dict = OrderedDict()
+        removed_count = 0
+        kept_count = 0
+        
+        print("正在移除 'text_encoder.' 相关权重...")
+        for key, value in state_dict.items():
+            if not key.startswith("text_encoder."):
+                pruned_state_dict[key] = value
+                kept_count += 1
+            else:
+                removed_count += 1
+        
+        print(f"移除 {removed_count} 个 text_encoder 权重，保留 {kept_count} 个。")
+        
+        # 根据用户选择，进行精度转换
+        if target_precision != 'fp32':
+            print(f"开始将保留的权重转换为 {target_precision} ...")
+            target_dtype = torch.float16 if target_precision == 'fp16' else torch.bfloat16
+            
+            for key, value in pruned_state_dict.items():
+                if isinstance(value, torch.Tensor) and value.is_floating_point():
+                    pruned_state_dict[key] = value.to(target_dtype)
+        else:
+            print("保留原始 FP32 精度。")
+
+        print(f"正在保存处理后的模型到: {output_path}")
+        torch.save(pruned_state_dict, output_path)
+
+        source_size = os.path.getsize(source_path) / (1024**2)
+        pruned_size = os.path.getsize(output_path) / (1024**2)
+
+        print("\n--- 成功 ---")
+        print(f"原始文件大小: {source_size:.2f} MB")
+        print(f"裁剪后文件大小 ({target_precision.upper()}): {pruned_size:.2f} MB")
+        print(f"总共节省空间: {source_size - pruned_size:.2f} MB")
+
+    except Exception as e:
+        print(f"\n[错误] 处理过程中发生错误: {e}")
+        import traceback
+        traceback.print_exc()
+
+if __name__ == "__main__":
+    prune_and_convert_model()