generate_ai_model_hash.py

def generate_hash_files(dir):
    import os
    for root, _, filenames in os.walk(dir):
        for filename in filenames:
            if filename.endswith(".safetensors") or filename.endswith(".ckpt") or filename.endswith(".pt"):
                fpath = os.path.join(root, filename)
                if os.path.exists(fpath + ".sha256"):
                    print("Hash file already exists for", fpath, "- skipping")
                else:
                    print("Generating hash for", fpath)
                    fhash = calculate_hash(fpath)
                    with open(fpath + ".sha256", "wt") as f:
                        f.write(fhash)
                        print("Hash written to", fpath + ".sha256")



def calculate_hash(filepath, chunksize=1024*1024):
    import hashlib
    sha256_hash = hashlib.sha256()
    with open(filepath, "rb") as f:
        while chunk := f.read(chunksize):
            sha256_hash.update(chunk)
    return sha256_hash.hexdigest()

if __name__ == "__main__":
    import sys
    if len(sys.argv) != 2:
        print("Usage: python generate_model_hashes.py <directory>")
        sys.exit(1)
    generate_hash_files(sys.argv[1])