inference.py

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