{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "880da911",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "053f4848",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/chatbot/miniconda3/envs/f5-tts/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "`torch_dtype` is deprecated! Use `dtype` instead!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Flash Attention 2 enabled.\n",
      "Loading base model...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading weights: 100%|██████████| 290/290 [00:00<00:00, 13434.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving base model to cuda...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model and adapter loaded successfully on GPU!\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from peft import PeftModel\n",
    "import torch\n",
    "import os\n",
    "\n",
    "use_flash = True\n",
    "base_model_path = \"outputs/ViSparkTTS\"\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "load_kwargs = {\n",
    "    \"pretrained_model_name_or_path\": f\"{base_model_path}/LLM\",\n",
    "    \"torch_dtype\": torch.bfloat16,\n",
    "    \"use_cache\": True,\n",
    "}\n",
    "\n",
    "if use_flash and torch.cuda.is_available():\n",
    "    try:\n",
    "        import flash_attn\n",
    "        load_kwargs[\"attn_implementation\"] = \"flash_attention_2\"\n",
    "        print(\"Flash Attention 2 enabled.\")\n",
    "    except ImportError:\n",
    "        print(\"flash_attn not installed, using standard attention.\")\n",
    "\n",
    "# 1. Nạp base model\n",
    "print(\"Loading base model...\")\n",
    "model = AutoModelForCausalLM.from_pretrained(**load_kwargs)\n",
    "\n",
    "print(f\"Moving base model to {device}...\")\n",
    "model = model.to(device)\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(f\"{base_model_path}/LLM\")\n",
    "\n",
    "# 5. Đưa về chế độ eval\n",
    "model = model.eval()\n",
    "print(\"Model and adapter loaded successfully on GPU!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "420d9838",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tok_str(ids, prefix):\n",
    "    return \"\".join(f\"<|bicodec_{prefix}_{i}|>\" for i in ids.cpu().numpy().flatten())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3d122c03",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchaudio.transforms as T\n",
    "\n",
    "def load_reference_audio(audio_path, audio_tokenizer):\n",
    "    import soundfile as sf\n",
    "    arr, sr = sf.read(audio_path)\n",
    "    if len(arr.shape) > 1:\n",
    "        arr = arr.mean(axis=1)\n",
    "    target_sr = audio_tokenizer.config[\"sample_rate\"]\n",
    "    if sr != target_sr:\n",
    "        resampler = T.Resample(orig_freq=sr, new_freq=target_sr)\n",
    "        arr = resampler(torch.from_numpy(arr).float()).numpy()\n",
    "    # if audio_tokenizer.config.get(\"volume_normalize\"):\n",
    "    #     arr = audio_volume_normalize(arr)\n",
    "    ref_np = audio_tokenizer.get_ref_clip(arr)\n",
    "    wav_t = torch.from_numpy(arr).unsqueeze(0).float().to(audio_tokenizer.device)\n",
    "    ref_t = torch.from_numpy(ref_np).unsqueeze(0).float().to(audio_tokenizer.device)\n",
    "    feat = audio_tokenizer.extract_features(wav_t)\n",
    "    feat = feat.to(audio_tokenizer.device)\n",
    "    sem_ids, glo_ids = audio_tokenizer.model.tokenize({\n",
    "        \"wav\": wav_t, \"ref_wav\": ref_t, \"feat\": feat,\n",
    "    })\n",
    "    return glo_ids, sem_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6250135b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sparktts.models.bicodec import BiCodec\n",
    "from sparktts.utils.file import load_config\n",
    "import numpy as np\n",
    "\n",
    "class FastAudioTokenizer:\n",
    "    def __init__(self, model_dir, device=\"cuda\"):\n",
    "        self.device = device\n",
    "        self.model_dir = model_dir\n",
    "        self.config = load_config(f\"{model_dir}/config.yaml\")\n",
    "        self.model = BiCodec.load_from_checkpoint(f\"{model_dir}/BiCodec\").to(device)\n",
    "        # Wav2Vec2 only used once for reference audio — keep on CPU to save GPU memory\n",
    "        from transformers import Wav2Vec2Model, Wav2Vec2FeatureExtractor as _WF\n",
    "        self.processor = _WF.from_pretrained(f\"{model_dir}/wav2vec2-large-xlsr-53\")\n",
    "        self.feature_extractor = Wav2Vec2Model.from_pretrained(f\"{model_dir}/wav2vec2-large-xlsr-53\")\n",
    "        self.feature_extractor.config.output_hidden_states = True\n",
    "        self.feature_extractor.to(\"cpu\")\n",
    "\n",
    "    @torch.inference_mode()\n",
    "    def extract_features(self, wavs):\n",
    "        if wavs.shape[0] != 1:\n",
    "            raise ValueError(f\"Expected batch size 1, got {wavs.shape}\")\n",
    "        wav_np = wavs.squeeze(0).cpu().numpy()\n",
    "        processed = self.processor(wav_np, sampling_rate=16000, return_tensors=\"pt\", padding=True)\n",
    "        out = self.feature_extractor(processed.input_values.to(\"cpu\"))\n",
    "        hs = out.hidden_states\n",
    "        feats = (hs[11] + hs[14] + hs[16]) / 3\n",
    "        return feats\n",
    "\n",
    "    def get_ref_clip(self, wav):\n",
    "        ref_len = (\n",
    "            int(self.config[\"sample_rate\"] * self.config[\"ref_segment_duration\"])\n",
    "            // self.config[\"latent_hop_length\"]\n",
    "            * self.config[\"latent_hop_length\"]\n",
    "        )\n",
    "        wav_len = len(wav)\n",
    "        if ref_len > wav_len:\n",
    "            wav = np.tile(wav, ref_len // wav_len + 1)\n",
    "        return wav[:ref_len]\n",
    "\n",
    "    @torch.inference_mode()\n",
    "    def detokenize(self, global_tokens, semantic_tokens):\n",
    "        global_tokens = global_tokens.unsqueeze(1)\n",
    "        wav_rec = self.model.detokenize(semantic_tokens, global_tokens)\n",
    "        return wav_rec.detach().squeeze().cpu().numpy()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "89fd9472",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/chatbot/miniconda3/envs/f5-tts/lib/python3.12/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.\n",
      "  WeightNorm.apply(module, name, dim)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Missing tensor: mel_transformer.spectrogram.window\n",
      "Missing tensor: mel_transformer.mel_scale.fb\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading weights: 100%|██████████| 422/422 [00:00<00:00, 86156.36it/s]\n",
      "\u001b[1mWav2Vec2Model LOAD REPORT\u001b[0m from: outputs/ViSparkTTS/wav2vec2-large-xlsr-53\n",
      "Key                          | Status     |  | \n",
      "-----------------------------+------------+--+-\n",
      "project_hid.weight           | UNEXPECTED |  | \n",
      "quantizer.codevectors        | UNEXPECTED |  | \n",
      "project_hid.bias             | UNEXPECTED |  | \n",
      "project_q.bias               | UNEXPECTED |  | \n",
      "quantizer.weight_proj.bias   | UNEXPECTED |  | \n",
      "quantizer.weight_proj.weight | UNEXPECTED |  | \n",
      "project_q.weight             | UNEXPECTED |  | \n",
      "\n",
      "Notes:\n",
      "- UNEXPECTED:\tcan be ignored when loading from different task/architecture; not ok if you expect identical arch.\n"
     ]
    }
   ],
   "source": [
    "audio_tokenizer = FastAudioTokenizer(\"outputs/ViSparkTTS\", device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bf88e3d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tok_str(ids, prefix):\n",
    "    return \"\".join(f\"<|bicodec_{prefix}_{i}|>\" for i in ids.cpu().numpy().flatten())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fe1e08ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "ref_glo_ids, _ = load_reference_audio(\"/mnt/data/data_create/data/data_20260609_135252/audio_segments/3_thứ_xiềng_xích_cần_gỡ_bỏ_trong_đời__Sách_Of_Human_Bondage_-_Kiếp_Người_-_NBCYduRAYZs_seg_005.mp3\", audio_tokenizer)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0cc1eeb0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['AI-Artificial Intelligence, hay trí tuệ nhân tạo, là lĩnh vực của khoa học máy tính tập trung vào việc tạo ra các hệ thống có khả năng thực hiện những nhiệm vụ vốn đòi hỏi trí thông minh của con người.',\n",
       " 'Những nhiệm vụ này bao gồm học hỏi từ dữ liệu, nhận diện hình ảnh và giọng nói, hiểu và tạo ngôn ngữ tự nhiên, đưa ra quyết định cũng như giải quyết vấn đề.',\n",
       " 'Thay vì chỉ thực hiện các lệnh được lập trình sẵn, các hệ thống AI hiện đại có thể tự cải thiện hiệu suất thông qua quá trình học từ dữ liệu.',\n",
       " 'Nhờ sự phát triển mạnh mẽ của sức mạnh tính toán, dữ liệu lớn và các mô hình học sâu Deep Learning, AI ngày càng được ứng dụng rộng rãi trong nhiều lĩnh vực như y tế, giáo dục, tài chính, sản xuất, giao thông và giải trí, góp phần tự động hóa công việc, nâng cao hiệu quả và hỗ trợ con người giải quyết các bài toán phức tạp.']"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import re\n",
    "\n",
    "texts = \"\"\"AI-Artificial Intelligence, hay trí tuệ nhân tạo, là lĩnh vực của khoa học máy tính tập trung vào việc tạo ra các hệ thống có khả năng thực hiện những nhiệm vụ vốn đòi hỏi trí thông minh của con người. Những nhiệm vụ này bao gồm học hỏi từ dữ liệu, nhận diện hình ảnh và giọng nói, hiểu và tạo ngôn ngữ tự nhiên, đưa ra quyết định cũng như giải quyết vấn đề. Thay vì chỉ thực hiện các lệnh được lập trình sẵn, các hệ thống AI hiện đại có thể tự cải thiện hiệu suất thông qua quá trình học từ dữ liệu. Nhờ sự phát triển mạnh mẽ của sức mạnh tính toán, dữ liệu lớn và các mô hình học sâu Deep Learning, AI ngày càng được ứng dụng rộng rãi trong nhiều lĩnh vực như y tế, giáo dục, tài chính, sản xuất, giao thông và giải trí, góp phần tự động hóa công việc, nâng cao hiệu quả và hỗ trợ con người giải quyết các bài toán phức tạp.\"\"\"\n",
    "sentences = re.split(r'(?<=[.!?])\\s+', texts.strip())\n",
    "sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "44f1809e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating token sequence...\n",
      "Token sequence generated.\n",
      "Thời gian: 7.4367570877075195\n",
      "Generating token sequence...\n",
      "Token sequence generated.\n",
      "Thời gian: 5.476766109466553\n",
      "Generating token sequence...\n",
      "Token sequence generated.\n",
      "Thời gian: 4.752292633056641\n",
      "Generating token sequence...\n",
      "Token sequence generated.\n",
      "Thời gian: 10.559756755828857\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "for text in sentences:\n",
    "    prompt = \"\".join([\n",
    "        \"<|task_tts|><|start_content|>\", text, \"<|end_content|>\",\n",
    "        \"<|start_global_token|>\", tok_str(ref_glo_ids, \"global\"),\n",
    "        \"<|end_global_token|>\",\n",
    "    ])\n",
    "    model_inputs = tokenizer([prompt], return_tensors=\"pt\", padding=True, truncation=True).to(device)\n",
    "    print(\"Generating token sequence...\")\n",
    "    start = time.time()\n",
    "    generated_ids = model.generate(\n",
    "        **model_inputs,\n",
    "        max_new_tokens = 2048, # Limit generation length\n",
    "        do_sample = True,\n",
    "        temperature = 0.8,\n",
    "        top_k = 50,\n",
    "        top_p = 1.0,\n",
    "        eos_token_id = tokenizer.eos_token_id, # Stop token\n",
    "        pad_token_id = tokenizer.pad_token_id # Use models pad token id\n",
    "    )\n",
    "    print(\"Token sequence generated.\")\n",
    "    print(f\"Thời gian: {time.time() - start}\")\n",
    "    import re\n",
    "    predicts_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=False)[0]\n",
    "    semantic_matches = re.findall(r\"bicodec_semantic_(\\d+)\", predicts_text)\n",
    "\n",
    "    pred_semantic_ids = torch.tensor([int(t) for t in semantic_matches]).long().unsqueeze(0)\n",
    "\n",
    "    global_matches = re.findall(r\"bicodec_global_(\\d+)\", predicts_text)\n",
    "\n",
    "    pred_global_ids = torch.tensor([int(t) for t in global_matches]).long().unsqueeze(0)\n",
    "    pred_global_ids = pred_global_ids.unsqueeze(0)\n",
    "    audio_tokenizer.device = device\n",
    "    audio_tokenizer.model.to(device)\n",
    "\n",
    "    wav_np = audio_tokenizer.detokenize(\n",
    "        pred_global_ids.to(device).squeeze(0),\n",
    "        pred_semantic_ids.to(device),\n",
    "    )\n",
    "    import soundfile as sf\n",
    "    sf.write(\"output.wav\", wav_np, 16000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4395c45c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[165137, 165146,  79759,  76294, 128277, 128362, 128531, 128818, 137125,\n",
       "          59735, 128738, 128818, 128439, 130051,     11, 134897,  88057, 128465,\n",
       "          47742,  78228, 130179, 128338, 128292, 131307,  18183,  20909,     11,\n",
       "          15235, 128300, 129357,  63478, 128740, 128284, 129922, 142423,  69086,\n",
       "         128271, 132122, 128978, 128260,    379, 128384,     11, 129023, 129677,\n",
       "             11, 128556, 128311,     11,  85636, 128415,     11,    342,  22516,\n",
       "          92430,  47742, 128486, 128597,     11, 131072, 128463, 128377, 128285,\n",
       "         128585,  62029, 128289,     11, 130314,   2162,     78, 128423, 128432,\n",
       "          47742, 129481, 128886,    390, 128251, 128486, 128784,  78228, 128601,\n",
       "         130051, 142244, 138628,     13, 165152, 165150, 152540, 154977, 153745,\n",
       "         155709, 154095, 151924, 152852, 155149, 153679, 155569, 152561, 155231,\n",
       "         152623, 155038, 155160, 152114, 154539, 153530, 153606, 155081, 151665,\n",
       "         155518, 154616, 154585, 155316, 154605, 155198, 153404, 152531, 152189,\n",
       "         155069, 151734, 165156, 165151, 161183, 163406, 156108, 157172, 159747,\n",
       "         160274, 156508, 163943, 161655, 160188, 157438, 159697, 160962, 161922,\n",
       "         157329, 163055, 163565, 157884, 160847, 162349, 162166, 159898, 162137,\n",
       "         155952, 162329, 158860, 158368, 163320, 163698, 158866, 161754, 156045,\n",
       "         161781, 162643, 159845, 156491, 159098, 161443, 157418, 160205, 159677,\n",
       "         157195, 159060, 156294, 162855, 161141, 160884, 162908, 163591, 163721,\n",
       "         159985, 161505, 163400, 157075, 158176, 163852, 156724, 155960, 157134,\n",
       "         156730, 157100, 157719, 160138, 156053, 163082, 162112, 163222, 162708,\n",
       "         163951, 159324, 162149, 162244, 163166, 161344, 163295, 160045, 158717,\n",
       "         159497, 160215, 159464, 157545, 159952, 161005, 159927, 156467, 163005,\n",
       "         157756, 157559, 163434, 161419, 161814, 159068, 159545, 159537, 160482,\n",
       "         158111, 155914, 161622, 162466, 163557, 162132, 162480, 159966, 156904,\n",
       "         160164, 159862, 158045, 157168, 163705, 156507, 157418, 163150, 162201,\n",
       "         159854, 162424, 162622, 163095, 156420, 159733, 155826, 159546, 160848,\n",
       "         163520, 163801, 160157, 163907, 156737, 163046, 163649, 156372, 163141,\n",
       "         158528, 160423, 156648, 161102, 163150, 162889, 161244, 162295, 160403,\n",
       "         162932, 163151, 163565, 160506, 155962, 163325, 160052, 155847, 159080,\n",
       "         157075, 162074, 159234, 158195, 161208, 160646, 158204, 158771, 163784,\n",
       "         161043, 159789, 160898, 157406, 157634, 163185, 159604, 157859, 159819,\n",
       "         156740, 162216, 159187, 159782, 162770, 159698, 161584, 158792, 162061,\n",
       "         160414, 162513, 161488, 161657, 158845, 162265, 163831, 156452, 159990,\n",
       "         162662, 159343, 163656, 160513, 163855, 163545, 161231, 161768, 159900,\n",
       "         161432, 162328, 160985, 162328, 163727, 158142, 163028, 161832, 157185,\n",
       "         156924, 163678, 158274, 157577, 156944, 161810, 156353, 159171, 156388,\n",
       "         158483, 162603, 162159, 161306, 160622, 160755, 158437, 158437, 157435,\n",
       "         156491, 163514, 158383, 156499, 158873, 157740, 155867, 162259, 157845,\n",
       "         155973, 162236, 163782, 162715, 163906, 158965, 160661, 163530, 163219,\n",
       "         161738, 156337, 157972, 158792, 161548, 161805, 163637, 160907, 156925,\n",
       "         158766, 161713, 162490, 161189, 162721, 158641, 158350, 156773, 158579,\n",
       "         160833, 159815, 157764, 160912, 158101, 158716, 162686, 161412, 156938,\n",
       "         159160, 163721, 163922, 156388, 157269, 162874, 162874, 161258, 160213,\n",
       "         160418, 161376, 159695, 156193, 156256, 157541, 158572, 162386, 161397,\n",
       "         160905, 162927, 157484, 158890, 157012, 156608, 161257, 155876, 163827,\n",
       "         156298, 157787, 161126, 163343, 158296, 155928, 162380, 157533, 162002,\n",
       "         163124, 156189, 162390, 161335, 158773, 157377, 163346, 159620, 161849,\n",
       "         159265, 158042, 163777, 158735, 163733, 159987, 162081, 159976, 157742,\n",
       "         158104, 161394, 161174, 155917, 162853, 160739, 158496, 156575, 157255,\n",
       "         160458, 158161, 160103, 160585, 156197, 156881, 163372, 157876, 156914,\n",
       "         161033, 156623, 163401, 159018, 163484, 163768, 161356, 157356, 156116,\n",
       "         159824, 159158, 163647, 161644, 157942, 163125, 159093, 157775, 158567,\n",
       "         161205, 155994, 159083, 161406, 162908, 156980, 162425, 157422, 157812,\n",
       "         163513, 159634, 161329, 156187, 161973, 156361, 156145, 156136, 159875,\n",
       "         161250, 158473, 163208, 162969, 161974, 163567, 157469, 158768, 155916,\n",
       "         158267, 158663, 156982, 156856, 160410, 162978, 157737, 159002, 157011,\n",
       "         159114, 157975, 161974, 162809, 158845, 158641, 157465, 163443, 163740,\n",
       "         157985, 156421, 156039, 156559, 162002, 157471, 155807, 163721, 157562,\n",
       "         156761, 159608, 156517, 161510, 160470, 159108, 156089, 161799, 161880,\n",
       "         157832, 160562, 160730, 160747, 160424, 155874, 158572, 157522, 160590,\n",
       "         161429, 161853, 157486, 155795, 159328, 157459, 159532, 158455, 163600,\n",
       "         163514, 160637, 159233, 162269, 162706, 162607, 156981, 162798, 161476,\n",
       "         158157, 158378, 158249, 162037, 156016, 157037, 158740, 158539, 163241,\n",
       "         163907, 158196, 157160, 159884, 161310, 155880, 162660, 156522, 160279,\n",
       "         157956, 163054, 159612, 158775, 158695, 162490, 157567, 160666, 160472,\n",
       "         158862, 156720, 156469, 156168, 158383, 156972, 159716, 163436, 158167,\n",
       "         156856, 159858, 161586, 156033, 159081, 158564, 159367, 163561, 161923,\n",
       "         157229, 156313, 162867, 158728, 157292, 162714, 163478, 156274, 158350,\n",
       "         161517, 159355, 159194, 161766, 156837, 158563, 157404, 158928, 156137,\n",
       "         158565, 161066, 161250, 156239, 159327, 163410, 159604, 157076, 163156,\n",
       "         157964, 163563, 161879, 156439, 159052, 157269, 159617, 156033, 162110,\n",
       "         157750, 156041, 162065, 158512, 163283, 156414, 161068, 161614, 163414,\n",
       "         155966, 160038, 163226, 160007, 161952, 159165, 156922, 156688, 156327,\n",
       "         157473, 159286, 157076, 160176, 158531, 163866, 160666, 159828, 163907,\n",
       "         162162, 157122, 162382, 158097, 159691, 158499, 162803, 160701, 163145,\n",
       "         159860, 156842, 160822, 156160, 158768, 160910, 159449, 162739, 160179,\n",
       "         156298, 157817, 160592, 161057, 156877, 161208, 158115, 156640, 156417,\n",
       "         161949, 163393, 160093, 158398, 158613, 157881, 160154, 157742, 160199,\n",
       "         159077, 157533, 160580, 163449, 156503, 158469, 162076, 161117, 161687,\n",
       "         163628, 157360, 163491, 156126, 158641, 161916, 158740, 158306, 161712,\n",
       "         160201, 157947, 158678, 158849, 156137, 162708, 160270, 157276, 157330,\n",
       "         158845, 161250, 161803, 156906, 161057, 161318, 163090, 157855, 159008,\n",
       "         162941, 162059, 159081, 162641, 156393, 160894, 162412, 158047, 161598,\n",
       "         157501, 156060, 163844, 159779, 158874, 158062, 157535, 156007, 161720,\n",
       "         158300, 162541, 160303, 159967, 161713, 156485, 160480, 160294, 161039,\n",
       "         155838, 157197, 162588, 159049, 157385, 161433, 162520, 161980, 159917,\n",
       "         160979, 156234, 159717, 162639, 157373, 161310, 163244, 159620, 159077,\n",
       "         157022, 157528, 157646, 162358, 155823, 160193, 160854, 156193, 156618,\n",
       "         157868, 162074, 159782, 163396, 161623, 157175, 160027, 163548, 156595,\n",
       "         163398, 159130, 158469, 158109, 157645, 156517, 159823, 156674, 161376,\n",
       "         160885, 163585, 163521, 158449, 156457, 161765, 159288, 162833, 160186,\n",
       "         157683, 163104, 161640, 157648, 157888, 163688, 160069, 160509, 158408,\n",
       "         159532, 161764, 163587, 160686, 158788, 160324, 157333, 157935, 162778,\n",
       "         157893, 159531, 160356, 161639, 160019, 158049, 161637, 160054, 161342,\n",
       "         163723, 157771, 159726, 161556, 156414, 162366, 159440, 163406, 158186,\n",
       "         160393, 163885, 158826, 161086, 160365, 156583, 158647, 159308, 162298,\n",
       "         158576, 155889, 162758, 162664, 156428, 163043, 157768, 163446, 160511,\n",
       "         156522, 161673, 158588, 162701, 157156, 160574, 162711, 155952, 159142,\n",
       "         156023, 163358, 157922, 159704, 162446, 160249, 157404, 161433, 163391,\n",
       "         156414, 155900, 161697, 163513, 158587, 163331, 161262, 158519, 157643,\n",
       "         160804, 161436, 156783, 159917, 160078, 162051, 157072, 157329, 156367,\n",
       "         159901, 155811, 158613, 162092, 160154, 157413, 160382, 162229, 159086,\n",
       "         157913, 160010, 155784, 159634, 156090, 156639, 158023, 161008, 156873,\n",
       "         158102, 156941, 159192, 162567, 156414, 159001, 163072, 162352, 162743,\n",
       "         163146, 156619, 161002, 161961, 159894, 157893, 160373, 160074, 162189,\n",
       "         157784, 160062, 162366, 161529, 161949, 157608, 161405, 162382, 156747,\n",
       "         160199, 160445, 163132, 162363, 158649, 160238, 156882, 161814, 163103,\n",
       "         156120, 156752, 162762, 158104, 159608, 162366, 163390, 159982, 160682,\n",
       "         157807, 156011, 162965, 157311, 160523, 161538, 158298, 158210, 156434,\n",
       "         156251, 159913, 161490, 157601, 162664, 125123, 151645]],\n",
       "       device='cuda:0')"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "generated_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6b7622fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "predicts_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=False)[0]\n",
    "semantic_matches = re.findall(r\"bicodec_semantic_(\\d+)\", predicts_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4def7db9",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "pred_semantic_ids = torch.tensor([int(t) for t in semantic_matches]).long().unsqueeze(0)\n",
    "\n",
    "global_matches = re.findall(r\"bicodec_global_(\\d+)\", predicts_text)\n",
    "\n",
    "pred_global_ids = torch.tensor([int(t) for t in global_matches]).long().unsqueeze(0)\n",
    "pred_global_ids = pred_global_ids.unsqueeze(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4a13d865",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_tokenizer.device = device\n",
    "audio_tokenizer.model.to(device)\n",
    "\n",
    "wav_np = audio_tokenizer.detokenize(\n",
    "    pred_global_ids.to(device).squeeze(0),\n",
    "    pred_semantic_ids.to(device),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5cc4d164",
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(\"Merging...\")\n",
    "# model = model.merge_and_unload()\n",
    "\n",
    "# print(\"Saving merged model...\")\n",
    "# model.save_pretrained(\n",
    "#     \"/mnt/data/data_create/train_or_fineturn/TTS/outputs/merged_model\",\n",
    "#     safe_serialization=True,\n",
    "# )\n",
    "\n",
    "# tokenizer.save_pretrained(\"/mnt/data/data_create/train_or_fineturn/TTS/outputs/merged_model\")\n",
    "\n",
    "# print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b3f8417",
   "metadata": {},
   "outputs": [],
   "source": [
    "import soundfile as sf\n",
    "sf.write(\"output.wav\", wav_np, 16000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2d1412b",
   "metadata": {},
   "source": [
    "# Push Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb2ff645",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/dolphinnlp/ViSparkTTS/commit/2e51d8230dcb13214dfd1c353c73b5a450c03643', commit_message='Upload folder using huggingface_hub', commit_description='', oid='2e51d8230dcb13214dfd1c353c73b5a450c03643', pr_url=None, repo_url=RepoUrl('https://huggingface.co/dolphinnlp/ViSparkTTS', endpoint='https://huggingface.co', repo_type='model', repo_id='dolphinnlp/ViSparkTTS'), pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from huggingface_hub import HfApi, create_repo, login\n",
    "\n",
    "repo_id = \"dolphinnlp/ViSparkTTS\"\n",
    "\n",
    "create_repo(repo_id, exist_ok=True)\n",
    "\n",
    "api = HfApi()\n",
    "api.upload_folder(\n",
    "    folder_path=\"/mnt/data/data_create/train_or_fineturn/TTS/outputs/ViSparkTTS\",\n",
    "    repo_id=repo_id,\n",
    "    repo_type=\"model\",\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "f5-tts",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}