SadTalker Models & Checkpoints: Verified Files, Sizes, Folder Structure & Missing Model Fixes

#!/usr/bin/env python3
"""
SadTalker Checkpoint Verifier
Scans checkpoints/ and gfpgan/weights/ to ensure all required model weights
exist, are non-empty, and match expected byte size bounds.
"""
import os
import sys

MODELS = [
    # Category, Relative Path, Minimum Bytes, Description, Required?
    ("Core 256", "checkpoints/SadTalker_V0.0.2_256.safetensors", 700_000_000, "Main 256px Safetensors Generator", True),
    ("Core 256", "checkpoints/mapping_00109-model.pth.tar", 150_000_000, "256px ExpNet Acoustic Mapping Model", True),
    ("Core 3DMM", "checkpoints/BFM_Fitting/similarity_Lm3D_all.mat", 3_000, "3DMM Landmark Similarity Matrix", True),
    ("Core 3DMM", "checkpoints/BFM_Fitting/BFM_model_front.mat", 20_000_000, "Basel Face Model Front Mesh Basis", True),
    ("Optional 512", "checkpoints/SadTalker_V0.0.2_512.safetensors", 700_000_000, "High-Res 512px Safetensors Generator", False),
    ("Optional 512", "checkpoints/mapping_00229-model.pth.tar", 150_000_000, "512px ExpNet Acoustic Mapping Model", False),
    ("Face Enhancer", "gfpgan/weights/GFPGANv1.4.pth", 340_000_000, "GFPGAN Face Restoration Generator", False),
    ("Face Enhancer", "gfpgan/weights/alignment_WFLW_4HG.pth", 190_000_000, "Facexlib 98-point Landmark Model", False),
    ("Face Enhancer", "gfpgan/weights/detection_Resnet50_Final.pth", 100_000_000, "RetinaFace Face Detection Backbone", False),
    ("Face Enhancer", "gfpgan/weights/parsing_parsenet.pth", 80_000_000, "BiSeNet Face Parsing Network", False),
]

def format_size(bytes_val):
    for unit in ['B', 'KB', 'MB', 'GB']:
        if bytes_val < 1024.0:
            return f"{bytes_val:.2f} {unit}"
        bytes_val /= 1024.0
    return f"{bytes_val:.2f} TB"

def main():
    print("=" * 68)
    print(" SadTalker Checkpoint Integrity & Hierarchy Verifier")
    print("=" * 68)
    
    missing_required = 0
    missing_optional = 0
    found_count = 0

    for category, rel_path, min_size, desc, is_required in MODELS:
        req_tag = "[REQUIRED]" if is_required else "[OPTIONAL]"
        if not os.path.exists(rel_path):
            if is_required:
                print(f"\033[91m✗ MISSING {req_tag}:\033[0m {rel_path} ({desc})")
                missing_required += 1
            else:
                print(f"\033[93m! NOT FOUND {req_tag}:\033[0m {rel_path} ({desc})")
                missing_optional += 1
            continue

        size = os.path.getsize(rel_path)
        if size < min_size:
            print(f"\033[91m✗ CORRUPTED {req_tag}:\033[0m {rel_path} -> Only {format_size(size)} (Expected > {format_size(min_size)})")
            if is_required:
                missing_required += 1
        else:
            print(f"\033[92m✓ OK {req_tag}:\033[0m {rel_path} [{format_size(size)}]")
            found_count += 1

    print("-" * 68)
    if missing_required == 0:
        print(f"\033[92m[SUCCESS] All critical SadTalker checkpoints are verified and ready!\033[0m")
        if missing_optional > 0:
            print(f"Note: {missing_optional} optional checkpoints (512 mode or GFPGAN) are missing.")
        sys.exit(0)
    else:
        print(f"\033[91m[FAILURE] Missing {missing_required} required checkpoints.\033[0m")
        print("SadTalker will throw a FileNotFoundError or EOF error during inference.")
        print("Consult the guide at https://sadtalker.ai/sadtalker-models-checkpoints for 1-click fixes.")
        sys.exit(1)

if __name__ == "__main__":
    main()