# /// 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

current_repository = Path(__file__).parent

with tempfile.TemporaryDirectory() as destination_dir:
    destination_dir = Path(destination_dir)
    if destination_dir.exists():
        shutil.rmtree(destination_dir)

    shutil.copytree(current_repository, destination_dir)

    model_name = "tts_b6369a24.safetensors"

    tensors = load_file(current_repository / model_name)

    new_tensors = {}

    for key, tensor in tensors.items():
        if key.startswith("mimi.encoder"):
            print("zeroing out", key)
            new_tensors[key] = torch.zeros_like(tensor)
        else:
            new_tensors[key] = tensor

    save_file(new_tensors, destination_dir / model_name)

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