Update app.py

app.py

@@ -638,6 +638,14 @@ from langchain.chains.conversation.memory import ConversationBufferWindowMemory
 from langchain.agents import Tool, initialize_agent
 from huggingface_hub import login
 
+from pydub import AudioSegment
+import io
+import math
+from threading import Thread
+from queue import Queue
+from transformers.generation.streamers import BaseStreamer
+import numpy as np
+
 # Check if the token is already set in the environment variables
 hf_token = os.getenv("HF_TOKEN")
 
@@ -952,7 +960,7 @@ def fetch_local_news():
     api_key = os.environ['SERP_API']
     url = f'https://serpapi.com/search.json?engine=google_news&q=birmingham headline&api_key={api_key}'
     response = requests.get(url)
-    if response.status_code == 200:
+    if response.status_code == 200):
         results = response.json().get("news_results", [])
         news_html = """
         <h2 style="font-family: 'Georgia', serif; color: #ff0000; background-color: #f8f8f8; padding: 10px; border-radius: 10px;">Birmingham Today</h2>
@@ -1111,44 +1119,126 @@ def generate_audio_elevenlabs(text):
         return None
 
 # Changes start here
-# Function to chunk the text
-def chunk_text(text, max_length=200):
-    words = text.split()
-    chunks = []
-    current_chunk = []
-    current_length = 0
-
-    for word in words:
-        if current_length + len(word) + 1 > max_length:
-            chunks.append(" ".join(current_chunk))
-            current_chunk = []
-            current_length = 0
-        current_chunk.append(word)
-        current_length += len(word) + 1
-
-    if current_chunk:
-        chunks.append(" ".join(current_chunk))
-
-    return chunks
-
-# Function to process each chunk
-def process_chunk(chunk, model, tokenizer, device):
-    description = "A female speaker with a slightly low-pitched voice delivers her words quite expressively, in a very confined sounding environment with clear audio quality. She speaks very fast."
+class ParlerTTSStreamer(BaseStreamer):
+    def __init__(
+        self,
+        model: ParlerTTSForConditionalGeneration,
+        device: Optional[str] = None,
+        play_steps: Optional[int] = 10,
+        stride: Optional[int] = None,
+        timeout: Optional[float] = None,
+    ):
+        self.decoder = model.decoder
+        self.audio_encoder = model.audio_encoder
+        self.generation_config = model.generation_config
+        self.device = device if device is not None else model.device
+
+        self.play_steps = play_steps
+        if stride is not None:
+            self.stride = stride
+        else:
+            hop_length = math.floor(self.audio_encoder.config.sampling_rate / self.audio_encoder.config.frame_rate)
+            self.stride = hop_length * (play_steps - self.decoder.num_codebooks) // 6
+        self.token_cache = None
+        self.to_yield = 0
+
+        self.audio_queue = Queue()
+        self.stop_signal = None
+        self.timeout = timeout
+
+    def apply_delay_pattern_mask(self, input_ids):
+        _, delay_pattern_mask = self.decoder.build_delay_pattern_mask(
+            input_ids[:, :1],
+            bos_token_id=self.generation_config.bos_token_id,
+            pad_token_id=self.generation_config.decoder_start_token_id,
+            max_length=input_ids.shape[-1],
+        )
+        input_ids = self.decoder.apply_delay_pattern_mask(input_ids, delay_pattern_mask)
 
-    input_ids = tokenizer(description, return_tensors="pt").input_ids.to(device)
-    prompt_input_ids = tokenizer(chunk, return_tensors="pt").input_ids.to(device)
+        mask = (delay_pattern_mask != self.generation_config.bos_token_id) & (delay_pattern_mask != self.generation_config.pad_token_id)
+        input_ids = input_ids[mask].reshape(1, self.decoder.num_codebooks, -1)
+        input_ids = input_ids[None, ...]
 
-    generation = model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids)
-    audio_arr = generation.cpu().numpy().squeeze()
+        input_ids = input_ids.to(self.audio_encoder.device)
 
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
-        sf.write(f.name, audio_arr, model.config.sampling_rate)
-        temp_audio_path = f.name
+        decode_sequentially = (
+            self.generation_config.bos_token_id in input_ids
+            or self.generation_config.pad_token_id in input_ids
+            or self.generation_config.eos_token_id in input_ids
+        )
+        if not decode_sequentially:
+            output_values = self.audio_encoder.decode(
+                input_ids,
+                audio_scales=[None],
+            )
+        else:
+            sample = input_ids[:, 0]
+            sample_mask = (sample >= self.audio_encoder.config.codebook_size).sum(dim=(0, 1)) == 0
+            sample = sample[:, :, sample_mask]
+            output_values = self.audio_encoder.decode(sample[None, ...], [None])
+
+        audio_values = output_values.audio_values[0, 0]
+        return audio_values.cpu().float().numpy()
+
+    def put(self, value):
+        batch_size = value.shape[0] // self.decoder.num_codebooks
+        if batch_size > 1:
+            raise ValueError("ParlerTTSStreamer only supports batch size 1")
+
+        if self.token_cache is None:
+            self.token_cache = value
+        else:
+            self.token_cache = torch.concatenate([self.token_cache, value[:, None]], dim=-1)
+
+        if self.token_cache.shape[-1] % self.play_steps == 0:
+            audio_values = self.apply_delay_pattern_mask(self.token_cache)
+            self.on_finalized_audio(audio_values[self.to_yield : -self.stride])
+            self.to_yield += len(audio_values) - self.to_yield - self.stride
+
+    def end(self):
+        if self.token_cache is not None:
+            audio_values = self.apply_delay_pattern_mask(self.token_cache)
+        else:
+            audio_values = np.zeros(self.to_yield)
 
-    logging.debug(f"Audio saved to {temp_audio_path}")
-    return temp_audio_path
+        self.on_finalized_audio(audio_values[self.to_yield :], stream_end=True)
+
+    def on_finalized_audio(self, audio: np.ndarray, stream_end: bool = False):
+        self.audio_queue.put(audio, timeout=self.timeout)
+        if stream_end:
+            self.audio_queue.put(self.stop_signal, timeout=self.timeout)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        value = self.audio_queue.get(timeout=self.timeout)
+        if not isinstance(value, np.ndarray) and value == self.stop_signal:
+            raise StopIteration()
+        else:
+            return value
+
+def numpy_to_mp3(audio_array, sampling_rate):
+    if np.issubdtype(audio_array.dtype, np.floating):
+        max_val = np.max(np.abs(audio_array))
+        audio_array = (audio_array / max_val) * 32767
+        audio_array = audio_array.astype(np.int16)
+
+    audio_segment = AudioSegment(
+        audio_array.tobytes(),
+        frame_rate=sampling_rate,
+        sample_width=audio_array.dtype.itemsize,
+        channels=1
+    )
+
+    mp3_io = io.BytesIO()
+    audio_segment.export(mp3_io, format="mp3", bitrate="320k")
+
+    mp3_bytes = mp3_io.getvalue()
+    mp3_io.close()
+
+    return mp3_bytes
 
-# Function to generate audio using Parler TTS
 def generate_audio_parler_tts(text):
     model_id = 'parler-tts/parler_tts_mini_v0.1'
     device = "cuda:0" if torch.cuda.is_available() else "cpu"
@@ -1161,28 +1251,38 @@ def generate_audio_parler_tts(text):
         model = ParlerTTSForConditionalGeneration.from_pretrained(model_id).to(device)
 
     tokenizer = AutoTokenizer.from_pretrained(model_id)
+    sampling_rate = model.audio_encoder.config.sampling_rate
+    frame_rate = model.audio_encoder.config.frame_rate
 
-    text_chunks = chunk_text(text)
-    audio_paths = []
+    description = "A female speaker with a slightly low-pitched voice delivers her words quite expressively, in a very confined sounding environment with clear audio quality. She speaks very fast."
+    play_steps_in_s = 2.0
+    play_steps = int(frame_rate * play_steps_in_s)
 
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        future_to_chunk = {executor.submit(process_chunk, chunk, model, tokenizer, device): chunk for chunk in text_chunks}
+    streamer = ParlerTTSStreamer(model, device=device, play_steps=play_steps)
 
-        for future in concurrent.futures.as_completed(future_to_chunk):
-            try:
-                audio_path = future.result()
-                audio_paths.append(audio_path)
-            except Exception as e:
-                logging.error(f"Error processing chunk: {e}")
+    inputs = tokenizer(description, return_tensors="pt").to(device)
+    prompt = tokenizer(text, return_tensors="pt").to(device)
+
+    generation_kwargs = dict(
+        input_ids=inputs.input_ids,
+        prompt_input_ids=prompt.input_ids,
+        streamer=streamer,
+        do_sample=True,
+        temperature=1.0,
+        min_new_tokens=10,
+    )
+
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
 
     combined_audio_path = tempfile.NamedTemporaryFile(delete=False, suffix=".wav").name
     combined_audio = []
 
-    for audio_path in audio_paths:
-        data, samplerate = sf.read(audio_path)
-        combined_audio.extend(data)
+    for new_audio in streamer:
+        print(f"Sample of length: {round(new_audio.shape[0] / sampling_rate, 2)} seconds")
+        combined_audio.extend(new_audio)
 
-    sf.write(combined_audio_path, combined_audio, samplerate)
+    sf.write(combined_audio_path, combined_audio, sampling_rate)
 
     logging.debug(f"Combined audio saved to {combined_audio_path}")
     return combined_audio_path
@@ -1256,3 +1356,4 @@ with gr.Blocks(theme='Pijush2023/scikit-learn-pijush') as demo:
      
 demo.queue()
 demo.launch(share=True)
+