242 lines
8.1 KiB
Python
242 lines
8.1 KiB
Python
import os
|
|
import json
|
|
import argparse
|
|
from pathlib import Path
|
|
from safetensors import safe_open
|
|
|
|
|
|
def validate_components(model_path):
|
|
"""Validate all model components exist."""
|
|
print("=" * 60)
|
|
print("VALIDATING COMPONENTS")
|
|
print("=" * 60)
|
|
|
|
components = {
|
|
"tokenizer": ["special_tokens_map.json", "tokenizer_config.json"],
|
|
"text_encoder": ["config.json", "model.safetensors.index.json"],
|
|
"vae": ["config.json", "diffusion_pytorch_model.safetensors"],
|
|
"scheduler": ["scheduler_config.json"],
|
|
"transformer": ["config.json", "diffusion_pytorch_model.safetensors.index.json"]
|
|
}
|
|
|
|
all_valid = True
|
|
for component, required_files in components.items():
|
|
component_path = os.path.join(model_path, component)
|
|
print(f"\n{component}:")
|
|
|
|
if not os.path.exists(component_path):
|
|
print(f" ✗ Directory not found")
|
|
all_valid = False
|
|
continue
|
|
|
|
for req_file in required_files:
|
|
file_path = os.path.join(component_path, req_file)
|
|
exists = os.path.exists(file_path)
|
|
symbol = "✓" if exists else "✗"
|
|
print(f" {symbol} {req_file}")
|
|
if not exists:
|
|
all_valid = False
|
|
|
|
return all_valid
|
|
|
|
|
|
def validate_dit_weights(dit_path):
|
|
"""Validate DiT weights structure."""
|
|
print("\n" + "=" * 60)
|
|
print("VALIDATING DiT WEIGHTS")
|
|
print("=" * 60)
|
|
|
|
# Load config
|
|
config_path = os.path.join(dit_path, "config.json")
|
|
with open(config_path) as f:
|
|
config = json.load(f)
|
|
|
|
hidden_size = config["hidden_size"]
|
|
depth = config["depth"]
|
|
num_heads = config["num_heads"]
|
|
|
|
print(f"\nArchitecture:")
|
|
print(f" - hidden_size: {hidden_size}")
|
|
print(f" - depth: {depth}")
|
|
print(f" - num_heads: {num_heads}")
|
|
|
|
# Load weight index
|
|
index_path = os.path.join(dit_path, "diffusion_pytorch_model.safetensors.index.json")
|
|
with open(index_path) as f:
|
|
index = json.load(f)
|
|
|
|
weight_map = index["weight_map"]
|
|
all_keys = list(weight_map.keys())
|
|
|
|
print(f"\nWeight statistics:")
|
|
print(f" - Total keys: {len(all_keys)}")
|
|
print(f" - Total size: {index['metadata']['total_size'] / 1e9:.2f} GB")
|
|
|
|
# Check structure
|
|
embedder_keys = [k for k in all_keys if 'embedder' in k]
|
|
block_keys = [k for k in all_keys if k.startswith('blocks.')]
|
|
final_keys = [k for k in all_keys if k.startswith('final_layer.')]
|
|
|
|
print(f"\nKey distribution:")
|
|
print(f" - Embedder layers: {len(embedder_keys)}")
|
|
print(f" - Transformer blocks: {len(block_keys)}")
|
|
print(f" - Final layer: {len(final_keys)}")
|
|
|
|
# Verify all blocks present
|
|
block_nums = set()
|
|
for key in block_keys:
|
|
if key.startswith('blocks.'):
|
|
block_num = int(key.split('.')[1])
|
|
block_nums.add(block_num)
|
|
|
|
expected_blocks = set(range(depth))
|
|
missing_blocks = expected_blocks - block_nums
|
|
|
|
if missing_blocks:
|
|
print(f"\n✗ Missing blocks: {sorted(missing_blocks)}")
|
|
return False
|
|
else:
|
|
print(f"\n✓ All {depth} blocks present (0-{depth-1})")
|
|
|
|
# Sample weights
|
|
first_shard = os.path.join(dit_path, "diffusion_pytorch_model-00001-of-00006.safetensors")
|
|
print(f"\nSampling weights from first shard:")
|
|
|
|
with safe_open(first_shard, framework="pt", device="cpu") as f:
|
|
sample_keys = [k for k in f.keys() if k in all_keys][:5]
|
|
for key in sample_keys:
|
|
tensor = f.get_tensor(key)
|
|
print(f" - {key}")
|
|
print(f" Shape: {tuple(tensor.shape)}, Dtype: {tensor.dtype}")
|
|
|
|
return True
|
|
|
|
|
|
def validate_shapes(dit_path):
|
|
"""Validate weight shapes match expected architecture."""
|
|
print("\n" + "=" * 60)
|
|
print("VALIDATING WEIGHT SHAPES")
|
|
print("=" * 60)
|
|
|
|
# Load config
|
|
config_path = os.path.join(dit_path, "config.json")
|
|
with open(config_path) as f:
|
|
config = json.load(f)
|
|
|
|
hidden_size = config["hidden_size"]
|
|
num_heads = config["num_heads"]
|
|
head_dim = hidden_size // num_heads
|
|
adaln_dim = config.get("adaln_tembed_dim", 512)
|
|
mlp_ratio = config.get("mlp_ratio", 4)
|
|
|
|
# Calculate FFN hidden_dim using SwiGLU formula from blocks.py
|
|
# hidden_dim = int(2 * (hidden_size * mlp_ratio) / 3)
|
|
# rounded to multiple_of=256
|
|
multiple_of = 256
|
|
ffn_hidden = int(2 * hidden_size * mlp_ratio / 3)
|
|
ffn_hidden = multiple_of * ((ffn_hidden + multiple_of - 1) // multiple_of)
|
|
|
|
# Expected shapes
|
|
expected = {
|
|
"x_embedder.proj.weight": (hidden_size, 16, 1, 2, 2),
|
|
"x_embedder.proj.bias": (hidden_size,),
|
|
"t_embedder.mlp.0.weight": (adaln_dim, 256),
|
|
"t_embedder.mlp.2.weight": (adaln_dim, adaln_dim),
|
|
"y_embedder.y_proj.0.weight": (hidden_size, 4096),
|
|
"blocks.0.attn.qkv.weight": (3 * hidden_size, hidden_size),
|
|
"blocks.0.attn.q_norm.weight": (head_dim,),
|
|
"blocks.0.attn.proj.weight": (hidden_size, hidden_size),
|
|
"blocks.0.cross_attn.q_linear.weight": (hidden_size, hidden_size),
|
|
"blocks.0.cross_attn.kv_linear.weight": (2 * hidden_size, hidden_size),
|
|
"blocks.0.ffn.w1.weight": (ffn_hidden, hidden_size),
|
|
"blocks.0.adaLN_modulation.1.weight": (6 * hidden_size, adaln_dim),
|
|
"final_layer.linear.weight": (64, hidden_size),
|
|
}
|
|
|
|
# Load and check
|
|
first_shard = os.path.join(dit_path, "diffusion_pytorch_model-00001-of-00006.safetensors")
|
|
|
|
all_valid = True
|
|
with safe_open(first_shard, framework="pt", device="cpu") as f:
|
|
for key, expected_shape in expected.items():
|
|
if key in f.keys():
|
|
tensor = f.get_tensor(key)
|
|
actual_shape = tuple(tensor.shape)
|
|
|
|
if actual_shape == expected_shape:
|
|
print(f"✓ {key}: {actual_shape}")
|
|
else:
|
|
print(f"✗ {key}: expected {expected_shape}, got {actual_shape}")
|
|
all_valid = False
|
|
|
|
return all_valid
|
|
|
|
|
|
def validate_model_index(model_path):
|
|
"""Validate model_index.json exists and is correct."""
|
|
print("\n" + "=" * 60)
|
|
print("VALIDATING MODEL INDEX")
|
|
print("=" * 60)
|
|
|
|
model_index_path = os.path.join(model_path, "model_index.json")
|
|
|
|
if not os.path.exists(model_index_path):
|
|
print("✗ model_index.json not found")
|
|
return False
|
|
|
|
with open(model_index_path) as f:
|
|
index = json.load(f)
|
|
|
|
required_keys = ["_class_name", "workload_type", "tokenizer", "text_encoder",
|
|
"vae", "scheduler", "transformer"]
|
|
|
|
all_valid = True
|
|
for key in required_keys:
|
|
if key in index:
|
|
print(f"✓ {key}: {index[key]}")
|
|
else:
|
|
print(f"✗ {key}: missing")
|
|
all_valid = False
|
|
|
|
return all_valid
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Validate LongCat weights for FastVideo")
|
|
parser.add_argument("--model-path", type=str, required=True,
|
|
help="Path to LongCat model directory")
|
|
parser.add_argument("--check-shapes", action="store_true",
|
|
help="Also validate weight shapes (slower)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
print(f"\nValidating: {args.model_path}\n")
|
|
|
|
# Run validations
|
|
components_valid = validate_components(args.model_path)
|
|
model_index_valid = validate_model_index(args.model_path)
|
|
dit_valid = validate_dit_weights(os.path.join(args.model_path, "transformer"))
|
|
|
|
if args.check_shapes:
|
|
shapes_valid = validate_shapes(os.path.join(args.model_path, "transformer"))
|
|
else:
|
|
shapes_valid = True
|
|
print("\nSkipping shape validation (use --check-shapes to enable)")
|
|
|
|
# Summary
|
|
print("\n" + "=" * 60)
|
|
print("VALIDATION SUMMARY")
|
|
print("=" * 60)
|
|
|
|
all_valid = components_valid and model_index_valid and dit_valid and shapes_valid
|
|
|
|
if all_valid:
|
|
print("✓ All validations passed!")
|
|
print("✓ Model ready for FastVideo")
|
|
else:
|
|
print("✗ Some validations failed")
|
|
print("✗ Please check errors above")
|
|
|
|
print("=" * 60)
|
|
|