File size: 2,305 Bytes
ce3c376 | 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 85 86 87 88 89 | #!/usr/bin/env python3
from __future__ import annotations
import argparse
import json
from pathlib import Path
MODEL_REQUIREMENTS = {
"qwen_image": [
"model_index.json",
"transformer",
"text_encoder",
"tokenizer",
"vae",
],
"cosyvoice": [
"cosyvoice3.yaml",
"llm.pt",
"flow.pt",
"hift.pt",
],
"wan_i2v": [
"configuration.json",
"Wan2.1_VAE.pth",
"high_noise_model",
"low_noise_model",
],
"qwen35_27b": [
"config.json",
"tokenizer.json",
],
}
def missing_index_shards(root: Path) -> list[str]:
missing: list[str] = []
for index_file in root.rglob("*.safetensors.index.json"):
data = json.loads(index_file.read_text())
weight_map = data.get("weight_map", {})
for shard in sorted(set(weight_map.values())):
shard_path = index_file.parent / shard
if not shard_path.exists():
missing.append(str(shard_path))
return missing
def check_model(name: str, root: Path) -> list[str]:
missing: list[str] = []
for item in MODEL_REQUIREMENTS[name]:
if not (root / item).exists():
missing.append(str(root / item))
if root.exists():
missing.extend(missing_index_shards(root))
return missing
def main() -> None:
parser = argparse.ArgumentParser(description="Validate VGen model directories.")
parser.add_argument("--qwen-image", required=True)
parser.add_argument("--cosyvoice", required=True)
parser.add_argument("--wan-i2v", required=True)
parser.add_argument("--qwen35-27b", required=True)
args = parser.parse_args()
paths = {
"qwen_image": Path(args.qwen_image),
"cosyvoice": Path(args.cosyvoice),
"wan_i2v": Path(args.wan_i2v),
"qwen35_27b": Path(args.qwen35_27b),
}
failed = False
for name, path in paths.items():
missing = check_model(name, path)
if missing:
failed = True
print(f"[FAIL] {name}: {path}")
for item in missing:
print(f" missing: {item}")
else:
print(f"[ OK ] {name}: {path}")
raise SystemExit(1 if failed else 0)
if __name__ == "__main__":
main()
|