File size: 2,443 Bytes
d477207 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | #!/usr/bin/env python3
"""
检查提取的VAE和UNet组件的精度和基本信息
"""
import os
import sys
from pathlib import Path
from safetensors.torch import load_file
def check_component_info(file_path):
"""检查组件文件的信息"""
if not Path(file_path).exists():
print(f"❌ 文件不存在: {file_path}")
return
print(f"\n📄 检查文件: {Path(file_path).name}")
try:
# 加载权重
state_dict = load_file(file_path)
# 统计信息
total_params = 0
dtypes = set()
for key, tensor in state_dict.items():
total_params += tensor.numel()
dtypes.add(str(tensor.dtype))
file_size = Path(file_path).stat().st_size / (1024**3) # GB
print(f" 📊 参数数量: {total_params:,}")
print(f" 💾 文件大小: {file_size:.2f} GB")
print(f" 🎯 数据精度: {', '.join(dtypes)}")
print(f" 🔧 参数种类: {len(state_dict)} 个张量")
# 显示前几个参数的键名
keys = list(state_dict.keys())[:5]
print(f" 🔑 示例参数: {', '.join(keys)}")
if len(state_dict) > 5:
print(f" ... 还有 {len(state_dict) - 5} 个参数")
except Exception as e:
print(f"❌ 检查失败: {e}")
def main():
# 设置路径
script_dir = Path(__file__).parent.absolute()
base_dir = script_dir.parent
extracted_dir = base_dir / "models" / "extracted_components"
print("🔍 检查提取的组件信息")
print("=" * 50)
# 检查各个组件文件
components = [
"waiNSFWIllustrious_v140_vae.safetensors",
"waiNSFWIllustrious_v140_unet.safetensors"
]
for component in components:
file_path = extracted_dir / component
check_component_info(file_path)
print(f"\n📂 检查目录: {extracted_dir}")
# 检查配置文件
config_files = [
"waiNSFWIllustrious_v140_vae_config.json",
"waiNSFWIllustrious_v140_unet_config.json"
]
print(f"\n📋 配置文件:")
for config_file in config_files:
config_path = extracted_dir / config_file
if config_path.exists():
print(f" ✅ {config_file}")
else:
print(f" ❌ {config_file} (缺失)")
if __name__ == "__main__":
main()
|