Fix: Peak audio normalization (-1 dB), robust speaker search, and tuned CFG for crystal-clear benchmark generation

inference.py

@@ -137,11 +137,23 @@ def synthesize_speech(
             speed=speed,
         )
 
-    # 5. Save output to standard 24kHz 16-bit PCM .wav with Peak Normalization
+    # 5. Save output to standard 24kHz 16-bit PCM .wav with DC Removal + Peak Normalization
     audio_flat = audio_tensor.squeeze().detach().cpu().float()
+
+    # Step A: DC Offset Removal (Removes massive electrical hum/bias)
+    audio_flat = audio_flat - audio_flat.mean()
+
+    # Step B: High-pass filtering (Remove sub-audible rumble below 50Hz)
+    try:
+        import torchaudio
+        audio_flat = torchaudio.functional.highpass_biquad(audio_flat.unsqueeze(0), 24000, cutoff_freq=50.0).squeeze(0)
+        audio_flat = audio_flat - audio_flat.mean()
+    except Exception:
+        pass
+
+    # Step C: Safe peak normalization to -1 dB (0.89) without amplifying DC bias
     max_amp = audio_flat.abs().max().item()
     if max_amp > 1e-5:
-        # Normalize to 0.89 (-1 dB) so audio is loud, clear, and distortion-free
         audio_flat = audio_flat * (0.89 / max_amp)
     audio_flat = audio_flat.clamp(-1.0, 1.0)
     audio_int16 = (audio_flat * 32767.0).to(torch.int16).numpy().tobytes()

models/viuai_tts.py

@@ -241,15 +241,25 @@ class ViuAITTS200M(nn.Module):
         # 1. Encode text
         text_latents = self.text_encoder(text_tokens)
 
-        # 2. Predict durations and calculate bounded total mel length (OOM Guard)
+        # 2. Predict durations with intelligent natural speech pacing safeguard
+        # (Eliminates the 0.49s audio collapse bug when duration predictor is under-converged)
         pred_log_dur, pred_pitch, pred_energy = self.prosody_engine(text_latents)
-        raw_dur = torch.round(torch.exp(pred_log_dur) - 1.0) / max(0.01, speed)
-        pred_dur = torch.clamp(raw_dur, min=1, max=50).long()
+        raw_dur = (torch.exp(pred_log_dur) - 1.0) / max(0.01, speed)
+        
+        # Human natural speech rate: ~12-15 chars/sec => ~7-10 frames per phoneme at 24kHz (hop 256)
+        mean_dur = raw_dur.mean().item()
+        if mean_dur < 5.5:
+            # Automatic pacing scale: ensures natural human speaking pace (~8.5 frames/token)
+            scale_factor = 8.5 / max(0.5, mean_dur)
+            raw_dur = raw_dur * scale_factor
+
+        # Enforce minimum clamp of 4 frames (~42ms) per token so characters never become 1-frame micro-clicks
+        pred_dur = torch.clamp(torch.round(raw_dur), min=4, max=60).long()
         total_mel_len = int(pred_dur.sum().item())
 
         # Guard against OOM: clamp total_mel_len
         cap = max_mel_len if max_mel_len is not None else self.max_mel_len
-        total_mel_len = min(max(total_mel_len, 10), cap)
+        total_mel_len = min(max(total_mel_len, 20), cap)
 
         # Text conditioning matches training representation exactly
         text_cond = text_latents
@@ -281,7 +291,8 @@ class ViuAITTS200M(nn.Module):
         # [B, T_mel, 80] -> [B, 80, T_mel]
         mel_out = generated_mel_seq.transpose(1, 2)
 
-        # 4. Synthesize raw waveform via Neural Vocoder
+        # 4. Synthesize raw waveform via Neural Vocoder with DC cleaning
         audio_out = self.vocoder(mel_out)
+        audio_out = audio_out - audio_out.mean(dim=-1, keepdim=True)
 
         return audio_out, mel_out

models/vocoder.py

@@ -99,7 +99,8 @@ class NeuralVocoder(nn.Module):
                 v = v.to(device)
                 v.eval()
                 object.__setattr__(self, "_vocos_instance", v)
-            except Exception:
+            except Exception as e:
+                print(f"  [!] Vocos load warning: {e}. Internal vocoder fallback active.")
                 object.__setattr__(self, "_vocos_instance", False)
         inst = getattr(self, "_vocos_instance", None)
         return inst if inst is not False else None
@@ -149,14 +150,26 @@ class NeuralVocoder(nn.Module):
                         if mel.shape[1] == 80:
                             proj = self._get_mel_80_to_100_proj(mel.device)
                             mel_in = torch.matmul(proj, mel)
+                            # Roll-off air-band channels (87 to 99) naturally from channel 86
+                            # Prevents zero-energy injection which log-mel treats as loud high-frequency noise!
+                            roll_off = torch.linspace(0.4, 3.5, 13, device=mel.device).unsqueeze(0).unsqueeze(-1)
+                            mel_in[:, 87:, :] = mel_in[:, 86:87, :].expand(-1, 13, -1) - roll_off
                         else:
                             mel_in = mel
                         wav = vocos.decode(mel_in)
                         if wav.ndim == 2:
                             wav = wav.unsqueeze(1)
+                        # Automatic DC offset subtraction and sub-bass filtering
+                        wav = wav - wav.mean(dim=-1, keepdim=True)
+                        try:
+                            import torchaudio
+                            wav = torchaudio.functional.highpass_biquad(wav, 24000, cutoff_freq=50.0)
+                            wav = wav - wav.mean(dim=-1, keepdim=True)
+                        except Exception:
+                            pass
                         return wav
-                except Exception:
-                    pass
+                except Exception as ex:
+                    print(f"  [!] Vocos decode exception: {ex}")
 
         # Internal Neural Vocoder fallback
         x = self.conv_pre(mel)
@@ -175,4 +188,7 @@ class NeuralVocoder(nn.Module):
         x = F.leaky_relu(x)
         x = self.conv_post(x)
         x = torch.tanh(x)
+
+        # DC offset subtraction on internal vocoder output
+        x = x - x.mean(dim=-1, keepdim=True)
         return x

train.py

@@ -410,29 +410,7 @@ class ViuAISpeechDataset(Dataset):
         T_text = len(tokens)
         T_mel = mel.shape[-1]
 
-        # 3. Dynamic Energy-Aware Duration Allocation
-        # Content characters receive more frames; vowels get high duration, spaces/punct get micro-pauses
-        token_weights = torch.ones(T_text, dtype=torch.float32)
-        for ti, ch_id in enumerate(tokens):
-            if ch_id in (0, 1, 2, 3):  # pad, unk, bos, eos
-                token_weights[ti] = 0.5
-            elif ch_id in (4, 5, 6, 7, 8, 9, 10, 11):  # punctuation/space
-                token_weights[ti] = 1.0
-            else:
-                token_weights[ti] = 2.5  # speech letters/Devanagari matras
-
-        token_weights = token_weights / token_weights.sum()
-        durations = torch.clamp(torch.round(token_weights * T_mel).long(), min=1)
-        diff = T_mel - int(durations.sum().item())
-        if diff > 0:
-            durations[0] += diff
-        elif diff < 0:
-            for _ in range(abs(diff)):
-                m_idx = torch.argmax(durations)
-                if durations[m_idx] > 1:
-                    durations[m_idx] -= 1
-
-        # Extract real acoustic pitch (F0) and RMS energy from speech waveform
+        # Extract real acoustic pitch (F0) and RMS energy from speech waveform first
         frame_pitch, frame_energy = extract_pitch_and_energy(wav, sample_rate=24000, hop_length=256)
         if frame_pitch.numel() > 0:
             text_pitch = F.interpolate(
@@ -448,6 +426,39 @@ class ViuAISpeechDataset(Dataset):
         else:
             text_energy = torch.zeros(T_text, dtype=torch.float32)
 
+        # 3. Dynamic Energy-Aware Duration Allocation
+        # Content characters receive more frames; vowels get high duration, spaces/punct get micro-pauses
+        token_weights = torch.ones(T_text, dtype=torch.float32)
+        for ti, ch_id in enumerate(tokens):
+            if ch_id in (0, 1, 2, 3):  # pad, unk, bos, eos
+                token_weights[ti] = 0.4
+            elif ch_id in (4, 5, 6, 7, 8, 9, 10, 11):  # punctuation/space
+                token_weights[ti] = 0.9
+            else:
+                token_weights[ti] = 2.4  # speech letters/Devanagari matras
+
+        # Modulate content weights with active speech energy ratio
+        if frame_energy.numel() > 0:
+            speech_ratio = (frame_energy > 0.15).float().mean().item()
+            token_weights[token_weights > 1.0] *= (1.0 + speech_ratio * 0.4)
+
+        token_weights = token_weights / token_weights.sum()
+        durations = torch.clamp(torch.round(token_weights * T_mel).long(), min=1)
+        diff = T_mel - int(durations.sum().item())
+
+        # Distribute residual evenly across content tokens (NEVER dumping all diff into token 0!)
+        if diff > 0:
+            sort_indices = torch.argsort(token_weights, descending=True)
+            for i in range(diff):
+                durations[sort_indices[i % T_text]] += 1
+        elif diff < 0:
+            sort_indices = torch.argsort(durations, descending=True)
+            for i in range(abs(diff)):
+                for idx in sort_indices:
+                    if durations[idx] > 1:
+                        durations[idx] -= 1
+                        break
+
         # 4. In-Context Reference Audio from SAME Speaker for Zero-Shot Voice Conditioning
         # Fast RAM-cached path: avoids reloading duplicate reference audio files from disk
         ref_mel = None
@@ -582,7 +593,16 @@ def run_golden_benchmark(model, dataset_dir: str, tag: str, output_dir: str = "c
 
             out_wav_path = os.path.join(save_dir, f"{file_tag}.wav")
             audio_flat = audio_out.squeeze().detach().cpu().float()
-            # Peak normalize to -1 dB (0.89) to eliminate faint/whispering audio
+            # Step A: DC Offset Removal (Removes massive electrical hum/bias)
+            audio_flat = audio_flat - audio_flat.mean()
+            # Step B: High-pass filtering (Remove sub-audible rumble below 50Hz)
+            try:
+                import torchaudio
+                audio_flat = torchaudio.functional.highpass_biquad(audio_flat.unsqueeze(0), 24000, cutoff_freq=50.0).squeeze(0)
+                audio_flat = audio_flat - audio_flat.mean()
+            except Exception:
+                pass
+            # Step C: Peak normalize to -1 dB (0.89) safely without amplifying DC noise
             max_amp = audio_flat.abs().max().item()
             if max_amp > 1e-5:
                 audio_flat = audio_flat * (0.89 / max_amp)
@@ -606,7 +626,7 @@ def run_golden_benchmark(model, dataset_dir: str, tag: str, output_dir: str = "c
 # ---------------------------------------------------------------------------
 def train(
     dataset_dir: str = None,
-    epochs: int = 60,
+    epochs: int = 200,
     batch_size: int = 0,  # 0/auto => GPU profile default
     lr: float = 1e-4,
     checkpoint_dir: str = "checkpoints",
@@ -622,6 +642,7 @@ def train(
     save_every: int = 5,
     save_every_steps: int = 0,
     log_every: int = 10,
+    train_vocoder: bool = False,
 ):
     is_ddp, rank, world, device = setup_distributed()
     is_main = (rank == 0)
@@ -838,7 +859,7 @@ def train(
                         mel_lengths=mel_lens,
                         cond_drop_prob=0.15,
                         audio_target=batch_audios,
-                        train_vocoder=False,  # Decoupled vocoder: 100% compute on acoustic DiT!
+                        train_vocoder=train_vocoder,
                     )
                     loss = losses["total_loss"] / accum
 
@@ -973,7 +994,7 @@ def train(
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="ViuAI_TTS_200M Master Training (SOTA Mode)")
     parser.add_argument("--data_dir", type=str, default=None, help="Path to dataset dir with metadata.csv")
-    parser.add_argument("--epochs", type=int, default=60, help="Number of full training epochs (default: 60)")
+    parser.add_argument("--epochs", type=int, default=200, help="Number of full training epochs (default: 200)")
     parser.add_argument("--batch_size", type=int, default=0, help="Per-GPU batch size (0/auto = GPU profile default)")
     parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate")
     parser.add_argument("--seed", type=int, default=42, help="Random seed")
@@ -989,6 +1010,7 @@ if __name__ == "__main__":
     parser.add_argument("--save_every", type=int, default=5, help="Save checkpoint every N epochs (default: 5)")
     parser.add_argument("--save_every_steps", type=int, default=0, help="Save checkpoint every N optimizer steps (0 = off)")
     parser.add_argument("--log_every", type=int, default=10, help="Log every N optimizer steps")
+    parser.add_argument("--train_vocoder", action="store_true", help="Enable Phase-2 neural vocoder fine-tuning alongside acoustic DiT")
     args = parser.parse_args()
 
     train(
@@ -1008,4 +1030,5 @@ if __name__ == "__main__":
         save_every=args.save_every,
         save_every_steps=args.save_every_steps,
         log_every=args.log_every,
+        train_vocoder=args.train_vocoder,
     )

train_colab.ipynb

@@ -1,129 +1,130 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# 🎙️ ViuAI_TTS_200M: GPU Training Notebook in Google Colab\n",
-    "\n",
-    "Welcome to the official training notebook for **ViuAI_TTS_200M** (~190.5M Parameters).\n",
-    "\n",
-    "### 🚀 Features:\n",
-    "* **Optimal Transport Flow Matching (OT-CFM)**\n",
-    "* **Diffusion Transformer (DiT)** with AdaLN-Zero\n",
-    "* **Human Prosody Engine** (Dynamic F0 Pitch + Energy + Duration)\n",
-    "* **Zero-Shot Human Voice Cloning** (ElevenLabs / Grok style)\n",
-    "* **24,000 Hz Studio Audio Generation**\n",
-    "\n",
-    "---\n",
-    "### ⚡ Step 0: Ensure GPU is Enabled\n",
-    "Click **Runtime** > **Change runtime type** > Select **T4 GPU** (or A100/L4 if Pro) > **Save**."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# 1. Check GPU\n",
-    "!nvidia-smi"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# 2. Clone ViuAI_TTS_200M Model Repository from Hugging Face\n",
-    "!git clone https://huggingface.co/ViuAI/ViuAI_TTS_200M\n",
-    "%cd ViuAI_TTS_200M\n",
-    "!ls -la"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# 3. Install Dependencies\n",
-    "!pip install -q huggingface_hub datasets"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# 4. Pull Dataset from Hugging Face (ViuAI/ViuAI_TTS_Dataset)\n",
-    "import os\n",
-    "from huggingface_hub import snapshot_download\n",
-    "\n",
-    "dataset_dir = \"../dataset\"\n",
-    "snapshot_download(\n",
-    "    repo_id=\"ViuAI/ViuAI_TTS_Dataset\",\n",
-    "    repo_type=\"dataset\",\n",
-    "    local_dir=dataset_dir\n",
-    ")\n",
-    "print(\"Dataset downloaded successfully to:\", dataset_dir)\n",
-    "!ls -la ../dataset"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# 5. Verify Model Architecture & Parameter Budget on GPU\n",
-    "!python verify_model.py"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# 6. Start Mixed Precision (AMP) GPU Training Loop\n",
-    "!python train.py"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# 7. Test Inference & Play Synthesized Audio Directly in Colab!\n",
-    "import IPython.display as ipd\n",
-    "\n",
-    "!python inference.py \\\n",
-    "  --text \"Namaste! Yeh ViuAI TTS 200M model ka live audio output hai jo GPU par train hua hai.\" \\\n",
-    "  --cfg_scale 2.0 \\\n",
-    "  --steps 15 \\\n",
-    "  --output \"colab_output.wav\"\n",
-    "\n",
-    "ipd.Audio(\"colab_output.wav\")"
-   ]
-  }
- ],
- "metadata": {
-  "accelerator": "GPU",
-  "colab": {
-   "provenance": []
-  },
-  "kernelspec": {
-   "display_name": "Python 3",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 0
-}
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# \ud83c\udf99\ufe0f ViuAI_TTS_200M: GPU Training Notebook in Google Colab\n",
+    "\n",
+    "Welcome to the official training notebook for **ViuAI_TTS_200M** (~190.5M Parameters).\n",
+    "\n",
+    "### \ud83d\ude80 Features:\n",
+    "* **Optimal Transport Flow Matching (OT-CFM)**\n",
+    "* **Diffusion Transformer (DiT)** with AdaLN-Zero\n",
+    "* **Human Prosody Engine** (Dynamic F0 Pitch + Energy + Duration)\n",
+    "* **Zero-Shot Human Voice Cloning** (ElevenLabs / Grok style)\n",
+    "* **24,000 Hz Studio Audio Generation**\n",
+    "\n",
+    "---\n",
+    "### \u26a1 Step 0: Ensure GPU is Enabled\n",
+    "Click **Runtime** > **Change runtime type** > Select **T4 GPU** (or A100/L4 if Pro) > **Save**."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 1. Check GPU\n",
+    "!nvidia-smi"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 2. Clone ViuAI_TTS_200M Model Repository from Hugging Face (skip heavy LFS checkpoints to save Colab disk)\n",
+    "!GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/ViuAI/ViuAI_TTS_200M\n",
+    "%cd ViuAI_TTS_200M\n",
+    "!ls -la"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 3. Install All Required Dependencies (Vocos, Torchaudio, SoundFile included)\n",
+    "!pip install -q vocos huggingface_hub datasets torchaudio soundfile"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 4. Pull Dataset from Hugging Face (ViuAI/ViuAI_TTS_Dataset)\n",
+    "import os\n",
+    "from huggingface_hub import snapshot_download\n",
+    "\n",
+    "dataset_dir = \"../dataset\"\n",
+    "snapshot_download(\n",
+    "    repo_id=\"ViuAI/ViuAI_TTS_Dataset\",\n",
+    "    repo_type=\"dataset\",\n",
+    "    local_dir=dataset_dir\n",
+    ")\n",
+    "print(\"Dataset downloaded successfully to:\", dataset_dir)\n",
+    "!ls -la ../dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 5. Verify Model Architecture & Parameter Budget on GPU\n",
+    "!python verify_model.py"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 6. Start Mixed Precision (AMP) GPU Training Loop (Auto-resumes from latest checkpoint)\n",
+    "!python train.py --epochs 200 --resume auto --train_vocoder --push_to_hub"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 7. Test Inference & Play Synthesized Audio Directly in Colab!\n",
+    "import IPython.display as ipd\n",
+    "\n",
+    "!python inference.py \\\n",
+    "  --checkpoint checkpoints/viuai_tts_latest.pt \\\n",
+    "  --text \"Namaste! Yeh ViuAI TTS 200M model ka live audio output hai jo GPU par train hua hai.\" \\\n",
+    "  --cfg_scale 1.6 \\\n",
+    "  --steps 20 \\\n",
+    "  --output \"colab_output.wav\"\n",
+    "\n",
+    "ipd.Audio(\"colab_output.wav\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}