Update util.py

util.py

@@ -1,4 +1,6 @@
 import torch
+import librosa
+import requests
 from nemo.collections.tts.models import AudioCodecModel
 from dataclasses import dataclass
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -59,7 +61,6 @@ class NemoAudioPlayer:
         if not (start_of_speech_flag and end_of_speech_flag):
             raise ValueError('Special speech tokens not found in output!')
         
-        print("Output validation passed - speech tokens found")
 
     def get_nano_codes(self, out_ids):
         """Extract nano codec tokens from model output"""
@@ -88,8 +89,6 @@ class NemoAudioPlayer:
 
         audio_codes = audio_codes.T.unsqueeze(0)
         len_ = torch.tensor([audio_codes.shape[-1]])
-        
-        print(f"Extracted audio codes shape: {audio_codes.shape}")
         return audio_codes, len_
 
     def get_text(self, out_ids):
@@ -107,8 +106,7 @@ class NemoAudioPlayer:
     def get_waveform(self, out_ids):
         """Convert model output to audio waveform"""
         out_ids = out_ids.flatten()
-        print("Starting waveform generation...")
-        
+
         # Validate output
         self.output_validation(out_ids)
         
@@ -116,15 +114,12 @@ class NemoAudioPlayer:
         audio_codes, len_ = self.get_nano_codes(out_ids)
         audio_codes, len_ = audio_codes.to(self.device), len_.to(self.device)
         
-        print("Decoding audio with NeMo codec...")
         with torch.inference_mode():
             reconstructed_audio, _ = self.nemo_codec_model.decode(
                 tokens=audio_codes, 
                 tokens_len=len_
             )
             output_audio = reconstructed_audio.cpu().detach().numpy().squeeze()
-
-        print(f"Generated audio shape: {output_audio.shape}")
         
         if self.text_tokenizer_name:
             text = self.get_text(out_ids)
@@ -175,18 +170,12 @@ class KaniModel:
         # Concatenate tokens
         modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1)
         attention_mask = torch.ones(1, modified_input_ids.shape[1], dtype=torch.int64)
-        
-        print(f"Input sequence length: {modified_input_ids.shape[1]}")
         return modified_input_ids, attention_mask
 
     def model_request(self, input_ids: torch.tensor, attention_mask: torch.tensor) -> torch.tensor:
         """Generate tokens using the model"""
         input_ids = input_ids.to(self.device)
         attention_mask = attention_mask.to(self.device)
-
-        print("Starting model generation...")
-        print(f"Generation parameters: max_tokens={self.conf.max_new_tokens}, "
-              f"temp={self.conf.temperature}, top_p={self.conf.top_p}")
         
         with torch.no_grad():
             generated_ids = self.model.generate(
@@ -201,14 +190,10 @@ class KaniModel:
                 eos_token_id=self.player.end_of_speech,
                 pad_token_id=self.tokenizer.pad_token_id if self.tokenizer.pad_token_id else self.tokenizer.eos_token_id
             )
-        
-        print(f"Generated sequence length: {generated_ids.shape[1]}")
         return generated_ids.to('cpu')
 
     def run_model(self, text: str):
-        """Complete pipeline: text -> tokens -> generation -> audio"""
-        print(f"Processing text: '{text[:50]}{'...' if len(text) > 50 else ''}'")
-        
+        """Complete pipeline: text -> tokens -> generation -> audio"""        
         # Prepare input
         input_ids, attention_mask = self.get_input_ids(text)
         
@@ -217,6 +202,50 @@ class KaniModel:
         
         # Convert to audio
         audio, _ = self.player.get_waveform(model_output)
-        
-        print("Text-to-speech generation completed successfully!")
-        return audio, text
+        return audio, text
+
+
+class Demo:
+    def __init__(self):
+        self.audio_dir = './audio_examples'
+        os.makedirs(self.audio_dir, exist_ok=True)
+        self.sentences = [
+            "You make my days brighter, and my wildest dreams feel like reality. How do you do that?",
+            "Anyway, um, so, um, tell me, tell me all about her. I mean, what's she like? Is she really, you know, pretty?",
+            "Great, and just a couple quick questions so we can match you with the right buyer. Is your home address still 330 East Charleston Road?",
+            "No, that does not make you a failure. No, sweetie, no. It just, uh, it just means that you're having a tough time...",
+            "Oh, yeah. I mean did you want to get a quick snack together or maybe something before you go?",
+            "I-- Oh, I am such an idiot sometimes. I'm so sorry. Um, I-I don't know where my head's at.",
+            "Got it. $300,000. I can definitely help you get a very good price for your property by selecting a realtor.",
+            "Holy fu- Oh my God! Don't you understand how dangerous it is, huh?"
+        ]
+        self.urls = [
+            'https://www.nineninesix.ai/examples/kani/1.wav',
+            'https://www.nineninesix.ai/examples/kani/2.wav',
+            'https://www.nineninesix.ai/examples/kani/5.wav',
+            'https://www.nineninesix.ai/examples/kani/6.wav',
+            'https://www.nineninesix.ai/examples/kani/3.wav',
+            'https://www.nineninesix.ai/examples/kani/7.wav',
+            'https://www.nineninesix.ai/examples/kani/4.wav',
+            'https://www.nineninesix.ai/examples/kani/8.wav'
+        ]
+    
+    def download_audio(self, url: str, filename: str):
+        filepath = os.path.join(self.audio_dir, filename)
+        if not os.path.exists(filepath):
+            r = requests.get(url)
+            r.raise_for_status()
+            with open(filepath, 'wb') as f:
+                f.write(r.content)
+        return filepath
+
+    def get_audio(self, filepath: str):
+        return librosa.load(filepath, sr=22050)
+
+    def __call__(self):
+        examples = {}
+        for idx, (sentence, url) in enumerate(zip(self.sentences, self.urls), start=1):
+            filename = f"{idx}.wav"
+            filepath = self.download_audio(url, filename)
+            examples[sentence] = self.get_audio(filepath)
+        return examples