import torch import random import re import yaml import spacy import warnings import noisereduce as nr def generate_false_start(word): """Creates progressive phonetic false starts optimized for TTS.""" clean_word = word.strip() if len(clean_word) <= 2: return f"{clean_word[0]}... {word}" stutters = [] num_stutters = random.choices([1, 2], weights=[0.7, 0.3])[0] if num_stutters == 1: cut = random.randint(1, max(1, len(clean_word) - 2)) stutters.append(f"{clean_word[:cut]}...") else: cut1 = max(1, len(clean_word) // 3) cut2 = max(cut1 + 1, int(len(clean_word) * 0.75)) stutters.append(f"{clean_word[:cut1]}...") if cut2 < len(clean_word): stutters.append(f"{clean_word[:cut2]}...") return " ".join(stutters) + f" {word}" nlp = None def inject_aphasia_cues(text, severity=0.3): """Injects fillers, repetitions, and false starts into a sentence.""" global nlp if nlp is None: try: nlp = spacy.load("en_core_web_sm") except OSError: raise RuntimeError( "\n" + "="*60 + "\n" "🚨 Missing Dependency: spaCy English Language Model\n" "To use apply_aphasia=True, you must download the NLP model first.\n" "Please run the following command in your terminal or Jupyter Notebook:\n\n" " python -m spacy download en_core_web_sm\n\n" + "="*60 ) doc = nlp(text) fillers = ["um...", "uh,", "uh...", "um,"] manipulated_words = [] for token in doc: is_content_word = token.pos_ in ["NOUN", "VERB", "ADJ", "PROPN"] if is_content_word and random.random() < severity: stutter_type = random.choices( ["filler", "repetition", "false_start"], weights=[0.4, 0.2, 0.4] )[0] if stutter_type == "false_start": manipulated_words.append(generate_false_start(token.text_with_ws)) continue elif stutter_type == "repetition" and token.i > 0 and doc[token.i - 1].pos_ != "PUNCT": prev_word = doc[token.i - 1].text manipulated_words.append(f"... {prev_word} ") elif stutter_type == "filler": filler = random.choice(fillers) manipulated_words.append(f"{filler} ") manipulated_words.append(token.text_with_ws) final_text = "".join(manipulated_words) final_text = re.sub(r'\s+([?.!,])', r'\1', final_text) return re.sub(r'\s{2,}', ' ', final_text).strip() class AphaVoicePipeline: def __init__(self, config_path="config.yaml", model_path="aphavoice_model.ckpt", vocoder_path="vocoder"): warnings.filterwarnings("ignore", category=FutureWarning) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 1. Load Config with open(config_path, "r") as f: self.config = yaml.safe_load(f) self.mel_mean = self.config['data']['mel_mean'] self.mel_std = self.config['data']['mel_std'] # 2. Load the Matcha Model using PyTorch Lightning's loader! from matcha.models.matcha_tts import MatchaTTS self.model = MatchaTTS.load_from_checkpoint( model_path, map_location=self.device, strict=False, n_vocab=self.config['model']['n_vocab'], n_spks=self.config['model']['n_spks'], spk_emb_dim=self.config['model']['spk_emb_dim'], weights_only=False ).to(self.device) self.model.mel_mean = torch.full((1, 80, 1), self.mel_mean).to(self.device) self.model.mel_std = torch.full((1, 80, 1), self.mel_std).to(self.device) self.model.eval() # 3. Load the Vocoder using your exact original logic! from matcha.hifigan.models import Generator as HiFiGAN from matcha.hifigan.config import v1 class AttrDict(dict): def __init__(self, *args, **kwargs): super(AttrDict, self).__init__(*args, **kwargs) self.__dict__ = self h = AttrDict(v1) self.vocoder = HiFiGAN(h).to(self.device) state_dict = torch.load(vocoder_path, map_location=self.device, weights_only=False) if 'generator' in state_dict: self.vocoder.load_state_dict(state_dict['generator']) else: self.vocoder.load_state_dict(state_dict) self.vocoder.eval() self.vocoder.remove_weight_norm() def synthesise(self, text, speaker_id=0, apply_aphasia=True, severity=0.3, length_scale=1.0): """ The main generation function incorporating all custom post-processing and denoising. """ from matcha.text import text_to_sequence from matcha.utils.utils import intersperse # Step 1: Apply textual dysfluency if toggled if apply_aphasia: text = inject_aphasia_cues(text, severity=severity) with torch.no_grad(): # Step 2: Convert text to sequence & handle tuple return sequence_output = text_to_sequence(text, [self.config['data']['cleaner']]) if isinstance(sequence_output, tuple): sequence = sequence_output[0] else: sequence = sequence_output # Interpserse adds blank tokens (0) between phonemes for better rhythm sequence = intersperse(sequence, 0) x = torch.tensor(sequence, dtype=torch.long, device=self.device).unsqueeze(0) x_lengths = torch.tensor([x.size(-1)], dtype=torch.long, device=self.device) # Step 3: Handle multi-speaker ID spk_tensor = torch.tensor([speaker_id], dtype=torch.long, device=self.device) # Step 4: Generate Mel-Spectrogram output = self.model.synthesise( x, x_lengths, n_timesteps=self.config['inference']['n_timesteps'], temperature=self.config['inference']['temperature'], spks=spk_tensor, length_scale=length_scale ) mel = output['mel'] # Step 5: 4D Tensor Fix (Prevents shape mismatch in vocoder) if mel.ndim == 4: if mel.shape[1] == mel.shape[2] == 80: mel = mel[:, 0, :, :] else: mel = mel.squeeze(1) # Step 6: Vocoder to Audio audio = self.vocoder(mel).squeeze() # Step 7: Peak Normalization audio = audio / (torch.max(torch.abs(audio)) + 1e-6) audio_np = audio.cpu().numpy() # Step 8: Spectral Denoising clean_audio_np = nr.reduce_noise( y=audio_np, sr=self.config['data']['sample_rate'], prop_decrease=0.20, n_fft=1024, win_length=1024 ) # Return sample rate, audio array (for Gradio player), and final manipulated text return self.config['data']['sample_rate'], clean_audio_np, text