# /// 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("Zeroing out voice cloning tensors in", model_path)

    tensors = load_file(model_path)
    new_tensors = {}
    for key, tensor in tensors.items():
        if key.startswith("mimi.encoder"):
            new_tensors[key] = torch.zeros_like(tensor)
        else:
            new_tensors[key] = tensor
    save_file(new_tensors, model_path)

current_repository = Path(__file__).parent

with tempfile.TemporaryDirectory() as destination_dir:
    destination_dir = Path(destination_dir)
    destination_dir.rmdir()

    shutil.copytree(current_repository, destination_dir)

    # legacy name
    load_and_zero_out_voice_cloning_tensors(destination_dir / "tts_b6369a24.safetensors")

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

    api = HfApi()
    api.upload_folder(
        folder_path=destination_dir,
        repo_id="kyutai/pocket-tts-without-voice-cloning",
        delete_patterns="*",
    )