Update app.py

app.py

@@ -1,5 +1,6 @@
 import queue
 import threading
+import spaces
 import os
 import gradio as gr
 from dia.model import Dia
@@ -12,7 +13,7 @@ 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="float16")
+model = Dia.from_pretrained("nari-labs/Dia-1.6B", compute_dtype="float32")
 
 # Queue for audio streaming
 audio_queue = queue.Queue()
@@ -43,7 +44,56 @@ def split_podcast_into_chunks(podcast_text, chunk_size=3):
 
     return chunks
 
-
+def postprocess_audio(output_audio_np, speed_factor: float=0.94):
+    """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
+
+    # --- Slow down audio ---
+    original_len = len(output_audio_np)
+    # Ensure speed_factor is positive and not excessively small/large to avoid issues
+    speed_factor = max(0.1, min(speed_factor, 5.0))
+    target_len = int(
+        original_len / speed_factor
+    )  # Target length based on speed_factor
+    if (
+        target_len != original_len and target_len > 0
+    ):  # Only interpolate if length changes and is valid
+        x_original = np.arange(original_len)
+        x_resampled = np.linspace(0, original_len - 1, target_len)
+        resampled_audio_np = np.interp(x_resampled, x_original, output_audio_np)
+        output_audio = (
+            output_sr,
+            resampled_audio_np.astype(np.float32),
+        )  # Use resampled audio
+        print(
+            f"Resampled audio from {original_len} to {target_len} samples for {speed_factor:.2f}x speed."
+        )
+    else:
+        output_audio = (
+            output_sr,
+            output_audio_np,
+        )  # Keep original if calculation fails or no change
+        print(f"Skipping audio speed adjustment (factor: {speed_factor:.2f}).")
+    # --- End slowdown ---
+
+    print(
+        f"Audio conversion successful. Final shape: {output_audio[1].shape}, Sample Rate: {output_sr}"
+    )
+
+    # Explicitly convert to int16 to prevent Gradio warning
+    if (
+        output_audio[1].dtype == np.float32
+        or output_audio[1].dtype == np.float64
+    ):
+        audio_for_gradio = np.clip(output_audio[1], -1.0, 1.0)
+        audio_for_gradio = (audio_for_gradio * 32767).astype(np.int16)
+        output_audio = (output_sr, audio_for_gradio)
+        print("Converted audio to int16 for Gradio output.")
+    return output_audio
+
+
+@spaces.GPU
 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
@@ -51,9 +101,15 @@ def process_audio_chunks(podcast_text):
         if stop_signal.is_set():
             break
         set_seed(42)
-        raw_audio = model.generate(chunk, use_torch_compile=True, verbose=False)
-        audio_chunk = np.array(raw_audio, dtype=np.float32)
-        audio_queue.put((sample_rate, audio_chunk))
+        raw_audio = model.generate(
+            chunk,
+            use_torch_compile=False,
+            verbose=False,
+            temperature=1.3
+            top_p=0.95
+        )
+        audio_chunk_np = np.array(raw_audio, dtype=np.float32)
+        audio_queue.put(postprocess_audio(audio_chunk_np))
 
     audio_queue.put(None)