# /// script
# requires-python = ">=3.14"
# dependencies = [
#     "safetensors>=0.7.0",
#     "torch>=2.9.1",
#     "packaging",
#     "huggingface_hub",
#     "numpy",
# ]
# ///

import os
import shutil
from pathlib import Path
from safetensors.torch import load_file, save_file
import torch
from huggingface_hub import HfApi
import tempfile

def load_and_zero_out_voice_cloning_tensors(model_path: Path):
    print("Converting to bf16 tensors in", model_path)

    tensors = load_file(model_path)
    for key, tensor in tensors.items():
        if tensor.dtype == torch.float32:
            tensors[key] = tensor.to(torch.bfloat16)
    save_file(tensors, model_path)

current_repository = Path(__file__).parent


# We look for all the files named model.safetensors and zero out the voice cloning tensors in them
for file_path in current_repository.rglob("model.safetensors"):
    load_and_zero_out_voice_cloning_tensors(file_path)