79 lines
3.2 KiB
Python
79 lines
3.2 KiB
Python
import torch
|
|
import os
|
|
from safetensors import safe_open
|
|
|
|
def quick_model_check():
|
|
"""Quick check of both models to understand their structure"""
|
|
|
|
# Model paths
|
|
pytorch_model_path = "/mnt/sharefs/users/hao.zhang/DMD/wan_bidirectional_dmd_from_scratch/2025-06-20-08-17-06.607828_seed1024/checkpoint_model_004800/model.pt"
|
|
safetensors_dir = "../Wan-AI/Wan2.1-T2V-1.3B-Diffusers/transformer"
|
|
|
|
print("Quick Model Structure Check")
|
|
print("=" * 40)
|
|
|
|
# Check PyTorch model
|
|
print(f"\n1. PyTorch Model: {pytorch_model_path}")
|
|
if os.path.exists(pytorch_model_path):
|
|
try:
|
|
checkpoint = torch.load(pytorch_model_path, map_location='cpu')
|
|
print(f" File size: {os.path.getsize(pytorch_model_path) / (1024**3):.2f} GB")
|
|
print(f" Checkpoint type: {type(checkpoint)}")
|
|
|
|
if isinstance(checkpoint, dict):
|
|
print(f" Checkpoint keys: {list(checkpoint.keys())}")
|
|
if 'state_dict' in checkpoint:
|
|
state_dict = checkpoint['state_dict']
|
|
elif 'model' in checkpoint:
|
|
state_dict = checkpoint['model']
|
|
else:
|
|
state_dict = checkpoint
|
|
else:
|
|
state_dict = checkpoint
|
|
|
|
print(f" Parameters: {len(state_dict)}")
|
|
total_params = sum(tensor.numel() for tensor in state_dict.values())
|
|
print(f" Total params: {total_params:,}")
|
|
|
|
# Show first few keys
|
|
print(f" Sample keys:")
|
|
for i, key in enumerate(list(state_dict.keys())[:5]):
|
|
tensor = state_dict[key]
|
|
print(f" {i+1}. {key}: {tensor.shape}, {tensor.dtype}")
|
|
|
|
except Exception as e:
|
|
print(f" Error: {e}")
|
|
else:
|
|
print(f" File not found")
|
|
|
|
# Check safetensors directory
|
|
print(f"\n2. Safetensors Directory: {safetensors_dir}")
|
|
if os.path.exists(safetensors_dir):
|
|
safetensors_files = []
|
|
for file in os.listdir(safetensors_dir):
|
|
if file.endswith('.safetensors'):
|
|
file_path = os.path.join(safetensors_dir, file)
|
|
safetensors_files.append((file, file_path))
|
|
|
|
print(f" Found {len(safetensors_files)} safetensors files:")
|
|
for file_name, file_path in safetensors_files:
|
|
file_size = os.path.getsize(file_path) / (1024**3)
|
|
print(f" - {file_name}: {file_size:.2f} GB")
|
|
|
|
# Quick load to check structure
|
|
try:
|
|
with safe_open(file_path, framework="pt", device="cpu") as f:
|
|
keys = list(f.keys())
|
|
print(f" Parameters: {len(keys)}")
|
|
if keys:
|
|
# Get first tensor to check shape
|
|
first_key = keys[0]
|
|
first_tensor = f.get_tensor(first_key)
|
|
print(f" Sample: {first_key}: {first_tensor.shape}, {first_tensor.dtype}")
|
|
except Exception as e:
|
|
print(f" Error loading: {e}")
|
|
else:
|
|
print(f" Directory not found")
|
|
|
|
if __name__ == "__main__":
|
|
quick_model_check() |