Update app.py

app.py

@@ -73,9 +73,8 @@ class VibeVoiceDemo:
                          speaker_1: str = None, speaker_2: str = None,
                          speaker_3: str = None, speaker_4: str = None,
                          cfg_scale: float = 1.3):
-        """Final audio generation only (no streaming, runs fully on GPU)."""
+        """Final audio generation only (no streaming)."""
         self.is_generating = True
-        self.stop_generation = False
     
         if not script.strip():
             raise gr.Error("Please provide a script.")
@@ -83,18 +82,17 @@ class VibeVoiceDemo:
         if num_speakers < 1 or num_speakers > 4:
             raise gr.Error("Number of speakers must be 1–4.")
     
-        # Collect selected speakers
         selected = [speaker_1, speaker_2, speaker_3, speaker_4][:num_speakers]
         for i, sp in enumerate(selected):
             if not sp or sp not in self.available_voices:
                 raise gr.Error(f"Invalid speaker {i+1} selection.")
     
-        # Load voices into memory
+        # load voices
         voice_samples = [self.read_audio(self.available_voices[sp]) for sp in selected]
         if any(len(v) == 0 for v in voice_samples):
             raise gr.Error("Failed to load one or more voice samples.")
     
-        # Format script
+        # format script
         lines = script.strip().split("\n")
         formatted = []
         for i, line in enumerate(lines):
@@ -108,7 +106,7 @@ class VibeVoiceDemo:
                 formatted.append(f"Speaker {sp_id}: {line}")
         formatted_script = "\n".join(formatted)
     
-        # Prepare processor inputs
+        # processor input
         inputs = self.processor(
             text=[formatted_script],
             voice_samples=[voice_samples],
@@ -118,48 +116,39 @@ class VibeVoiceDemo:
         )
     
         start = time.time()
-        sample_rate = 24000
-        audio_streamer = AudioStreamer(batch_size=1)
-    
-        # Run generation fully on GPU
-        self.model.generate(
+        outputs = self.model.generate(
             **inputs,
             max_new_tokens=None,
             cfg_scale=cfg_scale,
             tokenizer=self.processor.tokenizer,
             generation_config={'do_sample': False},
-            audio_streamer=audio_streamer,
             verbose=False,
         )
     
-        # Collect all audio chunks
-        all_chunks = []
-        for audio_chunk in audio_streamer.get_stream(0):
-            if torch.is_tensor(audio_chunk):
-                audio_chunk = audio_chunk.float().cpu().numpy()
-            if audio_chunk.ndim > 1:
-                audio_chunk = audio_chunk.squeeze()
-            all_chunks.append(audio_chunk)
-    
-        if not all_chunks:
+        # --- FIX: pull from speech_outputs ---
+        if outputs.speech_outputs and outputs.speech_outputs[0] is not None:
+            audio = outputs.speech_outputs[0].cpu().numpy()
+        else:
             self.is_generating = False
             raise gr.Error("❌ No audio was generated by the model.")
     
-        complete_audio = np.concatenate(all_chunks)
-        audio16 = convert_to_16_bit_wav(complete_audio)
+        if audio.ndim > 1:
+            audio = audio.squeeze()
+    
+        sample_rate = 24000
     
         # Save automatically to disk
         os.makedirs("outputs", exist_ok=True)
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         file_path = os.path.join("outputs", f"podcast_{timestamp}.wav")
-        sf.write(file_path, complete_audio, sample_rate)
+        sf.write(file_path, audio, sample_rate)
         print(f"💾 Saved podcast to {file_path}")
     
-        total_dur = len(complete_audio) / sample_rate
+        total_dur = len(audio) / sample_rate
         log = f"✅ Generation complete in {time.time()-start:.1f}s, {total_dur:.1f}s audio\nSaved to {file_path}"
     
         self.is_generating = False
-        return (sample_rate, complete_audio), log
+        return (sample_rate, audio), log