Add inference steps and seed controls to Gradio demo (#40)

demo/gradio_demo.py

@@ -19,6 +19,7 @@ import soundfile as sf
 import torch
 import os
 import traceback
+import re
 
 from vibevoice.modular.configuration_vibevoice import VibeVoiceConfig
 from vibevoice.modular.modeling_vibevoice_inference import VibeVoiceForConditionalGenerationInference
@@ -210,6 +211,8 @@ class VibeVoiceDemo:
                                  speaker_3: str = None,
                                  speaker_4: str = None,
                                  cfg_scale: float = 1.3,
+                                 inference_steps: Optional[int] = None,
+                                 seed: Optional[int] = None,
                                  disable_voice_cloning: bool = False) -> Iterator[tuple]:
         try:
             
@@ -240,9 +243,37 @@ class VibeVoiceDemo:
             
             voice_cloning_enabled = not disable_voice_cloning
 
+            # Resolve per-run parameters
+            resolved_inference_steps = int(inference_steps) if inference_steps is not None else int(self.inference_steps)
+            resolved_seed: Optional[int]
+            if seed is None:
+                resolved_seed = None
+            else:
+                # Gradio Number may come through as float
+                resolved_seed = int(seed)
+                if resolved_seed < 0:
+                    resolved_seed = None
+
+            def _sanitize_filename_part(value: str, max_len: int = 60) -> str:
+                value = (value or "").strip()
+                value = re.sub(r"\s+", "_", value)
+                value = re.sub(r"[^A-Za-z0-9_\-\+\.]+", "", value)
+                return value[:max_len] if len(value) > max_len else value
+
+            def _write_complete_wav(audio_int16: np.ndarray, sample_rate: int = 24000) -> str:
+                ts = datetime.now().strftime("%Y%m%d%H%M%S")
+                speakers_part = _sanitize_filename_part("-".join(selected_speakers) if selected_speakers else "speaker")
+                cfg_part = _sanitize_filename_part(f"{cfg_scale:.2f}")
+                seed_part = str(resolved_seed) if resolved_seed is not None else "rand"
+                # Requested pattern: yyyyMMddhhmmss_[speaker]_[cfg_scale]_[inference_step][seed].wav
+                basename = f"{ts}_{speakers_part}_{cfg_part}_{resolved_inference_steps}_{seed_part}.wav"
+                out_path = os.path.join(tempfile.gettempdir(), basename)
+                sf.write(out_path, np.asarray(audio_int16, dtype=np.int16), sample_rate, subtype="PCM_16")
+                return out_path
+
             # Build initial log
             log = f"🎙️ Generating podcast with {num_speakers} speakers\n"
-            log += f"📊 Parameters: CFG Scale={cfg_scale}, Inference Steps={self.inference_steps}\n"
+            log += f"📊 Parameters: CFG Scale={cfg_scale}, Inference Steps={resolved_inference_steps}, Seed={(resolved_seed if resolved_seed is not None else 'random')}\n"
             log += f"🎭 Speakers: {', '.join(selected_speakers)}\n"
             log += f"🔊 Voice cloning: {'Enabled' if voice_cloning_enabled else 'Disabled'}\n"
             if self.loaded_adapter_root:
@@ -331,7 +362,7 @@ class VibeVoiceDemo:
             # Start generation in a separate thread
             generation_thread = threading.Thread(
                 target=self._generate_with_streamer,
-                args=(inputs, cfg_scale, audio_streamer, voice_cloning_enabled)
+                args=(inputs, cfg_scale, audio_streamer, voice_cloning_enabled, resolved_inference_steps, resolved_seed, target_device)
             )
             generation_thread.start()
             
@@ -465,8 +496,9 @@ class VibeVoiceDemo:
                 final_log += "✨ Generation successful! Complete audio is ready.\n"
                 final_log += "💡 Not satisfied? You can regenerate or adjust the CFG scale for different results."
                 
-                # Yield the complete audio
+                # Yield the complete audio as a filepath so download uses our filename
-                yield None, (sample_rate, complete_audio), final_log, gr.update(visible=False)
+                complete_wav_path = _write_complete_wav(complete_audio, sample_rate=sample_rate)
+                yield None, complete_wav_path, final_log, gr.update(visible=False)
                 return
             
             if not has_received_chunks:
@@ -490,8 +522,9 @@ class VibeVoiceDemo:
                 final_log += "✨ Generation successful! Complete audio is ready in the 'Complete Audio' tab.\n"
                 final_log += "💡 Not satisfied? You can regenerate or adjust the CFG scale for different results."
                 
-                # Final yield: Clear streaming audio and provide complete audio
+                # Final yield: Clear streaming audio and provide complete audio as filepath
-                yield None, (sample_rate, complete_audio), final_log, gr.update(visible=False)
+                complete_wav_path = _write_complete_wav(complete_audio, sample_rate=sample_rate)
+                yield None, complete_wav_path, final_log, gr.update(visible=False)
             else:
                 final_log = log + "❌ No audio was generated."
                 yield None, None, final_log, gr.update(visible=False)
@@ -513,7 +546,16 @@ class VibeVoiceDemo:
             traceback.print_exc()
             yield None, None, error_msg, gr.update(visible=False)
     
-    def _generate_with_streamer(self, inputs, cfg_scale, audio_streamer, voice_cloning_enabled: bool):
+    def _generate_with_streamer(
+        self,
+        inputs,
+        cfg_scale,
+        audio_streamer,
+        voice_cloning_enabled: bool,
+        inference_steps: int,
+        seed: Optional[int],
+        target_device: str,
+    ):
         """Helper method to run generation with streamer in a separate thread."""
         try:
             # Check for stop signal before starting generation
@@ -521,10 +563,28 @@ class VibeVoiceDemo:
                 audio_streamer.end()
                 return
 
+            # Apply per-run DDPM steps
+            try:
+                self.model.set_ddpm_inference_steps(num_steps=int(inference_steps))
+            except Exception as e:
+                print(f"Warning: failed to set inference steps ({inference_steps}): {e}")
+                
             # Define a stop check function that can be called from generate
             def check_stop_generation():
                 return self.stop_generation
 
+            generator = None
+            if seed is not None:
+                try:
+                    if target_device == "cuda":
+                        generator = torch.Generator(device="cuda")
+                    else:
+                        generator = torch.Generator()
+                    generator.manual_seed(int(seed))
+                except Exception as e:
+                    print(f"Warning: failed to create seeded generator (seed={seed}, device={target_device}): {e}")
+                    generator = None
+                
             outputs = self.model.generate(
                 **inputs,
                 max_new_tokens=None,
@@ -533,6 +593,7 @@ class VibeVoiceDemo:
                 generation_config={
                     'do_sample': False,
                 },
+                generator=generator,
                 audio_streamer=audio_streamer,
                 stop_check_fn=check_stop_generation,  # Pass the stop check function
                 verbose=False,  # Disable verbose in streaming mode
@@ -687,13 +748,26 @@ def create_demo_interface(demo_instance: VibeVoiceDemo):
                 with gr.Accordion("Generation Parameters", open=False):
                     cfg_scale = gr.Slider(
                         minimum=1.0,
-                        maximum=2.0,
+                        maximum=12.0,
                         value=1.3,
                         step=0.05,
                         label="CFG Scale (Guidance Strength)",
                         # info="Higher values increase adherence to text",
                         elem_classes="slider-container"
                     )
+                    inference_steps = gr.Slider(
+                        minimum=1,
+                        maximum=50,
+                        value=int(getattr(demo_instance, "inference_steps", 10)),
+                        step=1,
+                        label="Inference Steps",
+                        elem_classes="slider-container",
+                    )
+                    seed = gr.Number(
+                        value=42,
+                        precision=0,
+                        label="Seed (-1 = random)",
+                    )
                     disable_voice_cloning = gr.Checkbox(
                         value=False,
                         label="Disable voice cloning (skip conditioning voice prompts)",
@@ -782,7 +856,7 @@ Or paste text directly and it will auto-assign speakers.""",
                 # Complete audio output (non-streaming)
                 complete_audio_output = gr.Audio(
                     label="Complete Podcast (Download after generation)",
-                    type="numpy",
+                    type="filepath",
                     elem_classes="audio-output complete-audio-section",
                     streaming=False,  # Non-streaming mode
                     autoplay=False,
@@ -817,7 +891,7 @@ Or paste text directly and it will auto-assign speakers.""",
         )
         
         # Main generation function with streaming
-        def generate_podcast_wrapper(num_speakers, script, speaker_1, speaker_2, speaker_3, speaker_4, cfg_scale, disable_voice_cloning):
+        def generate_podcast_wrapper(num_speakers, script, speaker_1, speaker_2, speaker_3, speaker_4, cfg_scale, inference_steps, seed, disable_voice_cloning):
             """Wrapper function to handle the streaming generation call."""
             try:
                 speakers = [speaker_1, speaker_2, speaker_3, speaker_4]
@@ -836,6 +910,8 @@ Or paste text directly and it will auto-assign speakers.""",
                     speaker_3=speakers[2],
                     speaker_4=speakers[3],
                     cfg_scale=cfg_scale,
+                    inference_steps=inference_steps,
+                    seed=seed,
                     disable_voice_cloning=disable_voice_cloning
                 ):
                     final_log = log
@@ -884,7 +960,7 @@ Or paste text directly and it will auto-assign speakers.""",
             queue=False
         ).then(
             fn=generate_podcast_wrapper,
-            inputs=[num_speakers, script_input] + speaker_selections + [cfg_scale, disable_voice_cloning],
+            inputs=[num_speakers, script_input] + speaker_selections + [cfg_scale, inference_steps, seed, disable_voice_cloning],
             outputs=[audio_output, complete_audio_output, log_output, streaming_status, generate_btn, stop_btn],
             queue=True  # Enable Gradio's built-in queue
         )