Working Nari labs code

app.py

@@ -7,13 +7,15 @@ from dia.model import Dia
 from huggingface_hub import InferenceClient
 import numpy as np
 from transformers import set_seed
+import io, soundfile as sf
+
 
 # Hardcoded podcast subject
 PODCAST_SUBJECT = "The future of AI and its impact on society"
 
 # Initialize the inference client
 client = InferenceClient("meta-llama/Llama-3.3-70B-Instruct", provider="cerebras", token=os.getenv("HF_TOKEN"))
-model = Dia.from_pretrained("nari-labs/Dia-1.6B", compute_dtype="float32")
+model = Dia.from_pretrained("nari-labs/Dia-1.6B", compute_dtype="bfloat16")
 
 # Queue for audio streaming
 audio_queue = queue.Queue()
@@ -36,15 +38,9 @@ Now go on, make 5 minutes of podcast.
 
 def split_podcast_into_chunks(podcast_text, chunk_size=3):
     lines = podcast_text.strip().split("\n")
-    chunks = []
-
-    for i in range(0, len(lines), chunk_size):
-        chunk = "\n".join(lines[i : i + chunk_size])
-        chunks.append(chunk)
+    return ["\n".join(lines[i : i + chunk_size]) for i in range(0, len(lines), chunk_size)]
 
-    return chunks
-
-def postprocess_audio(output_audio_np, speed_factor: float=0.94):
+def postprocess_audio(output_audio_np, speed_factor: float=0.8):
     """Taken from https://huggingface.co/spaces/nari-labs/Dia-1.6B/blob/main/app.py"""
     # Get sample rate from the loaded DAC model
     output_sr = 44100
@@ -98,6 +94,7 @@ def process_audio_chunks(podcast_text):
     chunks = split_podcast_into_chunks(podcast_text)
     sample_rate = 44100 # Modified from https://huggingface.co/spaces/nari-labs/Dia-1.6B/blob/main/app.py has 44100
     for chunk in chunks:
+        print(f"Processing chunk: {chunk}")
         if stop_signal.is_set():
             break
         set_seed(42)
@@ -117,26 +114,21 @@ def process_audio_chunks(podcast_text):
 def stream_audio_generator(podcast_text):
     """Creates a generator that yields audio chunks for streaming"""
     stop_signal.clear()
+    threading.Thread(target=process_audio_chunks, args=(podcast_text,)).start()
 
-    # Start audio generation in a separate thread
-    gen_thread = threading.Thread(target=process_audio_chunks, args=(podcast_text,))
-    gen_thread.start()
-
-    try:
-        while True:
-            # Get next chunk from queue
-            chunk = audio_queue.get()
-
-            # None signals end of generation
-            if chunk is None:
-                break
-
-            # Yield the audio chunk with sample rate
-            print(chunk)
-            yield chunk
-
-    except Exception as e:
-        print(f"Error in streaming: {e}")
+    while True:
+        chunk = audio_queue.get()
+        if chunk is None:
+            break
+        sr, data = chunk           # the tuple you produced earlier
+
+        # Encode the numpy array into a WAV blob
+        buf = io.BytesIO()
+        sf.write(buf, data.astype(np.float32) / 32768.0, sr, format="wav")
+        buf.seek(0)
+        buffer = buf.getvalue()
+        print("PRINTING BUFFER:", buffer)
+        yield buffer# <-- bytes, so the browser can play it
 
 
 def stop_generation():