#!/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()