Fix cold start. Add SharedGenerationState

app.py

@@ -24,6 +24,7 @@ from voxtream.utils.app import (
     CUSTOM_CSS,
     AppConfig,
     GenerationControl,
+    SharedGenerationState,
     SpeakingRateState,
     VisualizationState,
     build_low_latency_audio_head,
@@ -87,6 +88,7 @@ def demo_app(
     synthesize_fn,
     speaking_rate_state: SpeakingRateState,
     generation_control: GenerationControl,
+    shared_generation_state: SharedGenerationState,
 ):
     with gr.Blocks(
         css=CUSTOM_CSS,
@@ -175,13 +177,14 @@ def demo_app(
         text_progress = gr.HTML(
             render_text_progress(app_config, None), elem_id="text-progress-container"
         )
+        generation_session = gr.State("")
 
         def validate_inputs(audio, ttext):
             if not audio:
                 return gr.update(
                     visible=True, value="⚠️ Please provide a prompt audio."
                 ), gr.update(interactive=False)
-            if not ttext.strip():
+            if not ttext or not ttext.strip():
                 return gr.update(
                     visible=True, value="⚠️ Please provide target text."
                 ), gr.update(interactive=False)
@@ -197,10 +200,22 @@ def demo_app(
             inputs=prompt_enhancement,
             outputs=prompt_enhancement_msg,
         )
+
+        def update_speaking_rate(value, session_id):
+            speaking_rate_state.update(value)
+            shared_generation_state.update_speaking_rate(session_id, value)
+
+        speaking_rate_control.input(
+            fn=update_speaking_rate,
+            inputs=[speaking_rate_control, generation_session],
+            queue=False,
+            show_progress="hidden",
+        )
         speaking_rate_control.release(
-            fn=lambda value: speaking_rate_state.update(value),
-            inputs=speaking_rate_control,
+            fn=update_speaking_rate,
+            inputs=[speaking_rate_control, generation_session],
             queue=False,
+            show_progress="hidden",
         )
 
         for inp in [prompt_audio, target_text]:
@@ -211,6 +226,7 @@ def demo_app(
             )
 
         def prepare_generation(speaking_rate, enable_rate):
+            session_id = shared_generation_state.create(speaking_rate)
             generation_control.start()
             speaking_rate_state.start(speaking_rate)
             return (
@@ -218,8 +234,9 @@ def demo_app(
                 gr.update(interactive=False),
                 empty_rate_plot(app_config, show_target=enable_rate),
                 render_text_progress(app_config, None),
-                render_audio_stream(app_config, session_id=uuid.uuid4().hex),
+                render_audio_stream(app_config, session_id=session_id),
                 *generation_button_updates(running=True),
+                session_id,
             )
 
         submit_btn.click(
@@ -234,6 +251,7 @@ def demo_app(
                 pause_btn,
                 resume_btn,
                 stop_btn,
+                generation_session,
             ],
             show_progress="hidden",
         ).then(
@@ -246,6 +264,7 @@ def demo_app(
                 streaming_input,
                 speaking_rate_control,
                 enable_speaking_rate,
+                generation_session,
             ],
             outputs=[
                 output_audio,
@@ -256,25 +275,29 @@ def demo_app(
                 pause_btn,
                 resume_btn,
                 stop_btn,
+                generation_session,
             ],
         )
 
-        def pause_generation():
+        def pause_generation(session_id):
             generation_control.pause()
+            shared_generation_state.pause(session_id)
             return generation_button_updates(running=True, paused=True)
 
-        def resume_generation():
+        def resume_generation(session_id):
             generation_control.resume()
+            shared_generation_state.resume(session_id)
             return generation_button_updates(running=True)
 
-        def stop_generation():
+        def stop_generation(session_id):
             generation_control.stop()
             speaking_rate_state.stop()
+            shared_generation_state.stop(session_id)
             return generation_button_updates(running=False)
 
         pause_btn.click(
             fn=pause_generation,
-            inputs=[],
+            inputs=generation_session,
             outputs=[pause_btn, resume_btn, stop_btn],
             js=(
                 "() => { if (window.voxtreamLowLatencyAudio) { "
@@ -284,7 +307,7 @@ def demo_app(
         )
         resume_btn.click(
             fn=resume_generation,
-            inputs=[],
+            inputs=generation_session,
             outputs=[pause_btn, resume_btn, stop_btn],
             js=(
                 "() => { if (window.voxtreamLowLatencyAudio) { "
@@ -294,7 +317,7 @@ def demo_app(
         )
         stop_btn.click(
             fn=stop_generation,
-            inputs=[],
+            inputs=generation_session,
             outputs=[pause_btn, resume_btn, stop_btn],
             js=(
                 "() => { if (window.voxtreamLowLatencyAudio) { "
@@ -303,8 +326,11 @@ def demo_app(
             queue=False,
         )
 
-        clear_btn.click(
-            fn=lambda: (
+        def clear_generation(session_id):
+            generation_control.stop()
+            speaking_rate_state.stop()
+            shared_generation_state.stop(session_id)
+            return (
                 gr.update(value=None),
                 gr.update(value=""),
                 gr.update(value=None, visible=False),
@@ -315,8 +341,12 @@ def demo_app(
                 render_text_progress(app_config, None),
                 render_audio_stream(app_config, session_id=uuid.uuid4().hex),
                 *generation_button_updates(running=False),
-            ),
-            inputs=[],
+                "",
+            )
+
+        clear_btn.click(
+            fn=clear_generation,
+            inputs=generation_session,
             outputs=[
                 prompt_audio,
                 target_text,
@@ -330,6 +360,7 @@ def demo_app(
                 pause_btn,
                 resume_btn,
                 stop_btn,
+                generation_session,
             ],
         )
 
@@ -354,19 +385,21 @@ def demo_app(
                 pause_btn,
                 resume_btn,
                 stop_btn,
+                generation_session,
             ],
             fn=synthesize_fn,
             cache_examples=False,
         )
 
         ex.dataset.click(
-            fn=lambda: clear_outputs(app_config),
+            fn=lambda: (*clear_outputs(app_config), ""),
             inputs=[],
             outputs=[
                 output_audio,
                 rate_plot,
                 text_progress,
                 stream_audio,
+                generation_session,
             ],
             queue=False,
         ).then(
@@ -376,7 +409,7 @@ def demo_app(
             queue=False,
         )
 
-    demo.launch()
+    demo.queue(default_concurrency_limit=1).launch()
 
 
 def main():
@@ -435,6 +468,7 @@ def main():
     speech_generator = SpeechGenerator(config, spk_rate_config)
     speaking_rate_state = SpeakingRateState(app_config.speaking_rate_default)
     generation_control = GenerationControl()
+    shared_generation_state = SharedGenerationState()
     chunk_size = int(config.mimi_sr * app_config.min_chunk_sec)
 
     @spaces.GPU
@@ -446,13 +480,18 @@ def main():
         streaming_input,
         speaking_rate_control,
         enable_speaking_rate=True,
+        generation_session_id="",
     ):
-        stream_session_id = uuid.uuid4().hex
+        control_session_id = generation_session_id or shared_generation_state.create(
+            speaking_rate_control
+        )
+        stream_session_id = control_session_id or uuid.uuid4().hex
         stream_seq = 0
 
         if not prompt_audio_path or not target_text:
             speaking_rate_state.stop()
             generation_control.finish()
+            shared_generation_state.finish(control_session_id)
             yield (
                 gr.update(value=None, visible=False),
                 gr.update(interactive=True),
@@ -460,13 +499,38 @@ def main():
                 render_text_progress(app_config, None),
                 render_audio_stream(app_config, session_id=stream_session_id),
                 *generation_button_updates(running=False),
+                control_session_id,
+            )
+            return
+
+        if shared_generation_state.is_stopped(control_session_id):
+            speaking_rate_state.stop()
+            generation_control.finish()
+            shared_generation_state.finish(control_session_id)
+            yield (
+                gr.update(value=None, visible=False),
+                gr.update(interactive=True),
+                empty_rate_plot(app_config, show_target=enable_speaking_rate),
+                render_text_progress(app_config, None),
+                render_audio_stream(
+                    app_config,
+                    session_id=stream_session_id,
+                    active=False,
+                    final=True,
+                ),
+                *generation_button_updates(running=False),
+                control_session_id,
             )
             return
 
         ensure_generator_on_cuda(speech_generator)
         speaking_rate_state.ensure_started(speaking_rate_control)
         speaking_rate_gen = (
-            speaking_rate_state.values() if enable_speaking_rate else None
+            shared_generation_state.speaking_rate_values(
+                control_session_id, speaking_rate_control
+            )
+            if enable_speaking_rate
+            else None
         )
         text_metadata = build_text_progress_metadata(
             target_text,
@@ -511,14 +575,14 @@ def main():
 
         stream_iter = iter(stream)
         while True:
-            if not generation_control.wait_if_paused():
+            if not shared_generation_state.wait_if_paused(control_session_id):
                 stopped = True
                 break
             try:
                 frame, _, progress = next(stream_iter)
             except StopIteration:
                 break
-            if generation_control.is_stopped():
+            if shared_generation_state.is_stopped(control_session_id):
                 stopped = True
                 break
 
@@ -528,7 +592,7 @@ def main():
             plot_update, text_update = visualization.update(progress)
 
             if buffer_len >= chunk_size:
-                if generation_control.is_stopped():
+                if shared_generation_state.is_stopped(control_session_id):
                     stopped = True
                     break
                 audio = np.concatenate(buffer)
@@ -547,14 +611,16 @@ def main():
                         active=True,
                     ),
                     *generation_button_updates(
-                        running=True, paused=generation_control.is_paused()
+                        running=True,
+                        paused=shared_generation_state.is_paused(control_session_id),
                     ),
+                    control_session_id,
                 )
 
                 buffer = []
                 buffer_len = 0
 
-        stopped = stopped or generation_control.is_stopped()
+        stopped = stopped or shared_generation_state.is_stopped(control_session_id)
         if stopped and hasattr(stream, "close"):
             stream.close()
         final_text = visualization.final_text()
@@ -580,8 +646,10 @@ def main():
                     active=True,
                 ),
                 *generation_button_updates(
-                    running=True, paused=generation_control.is_paused()
+                    running=True,
+                    paused=shared_generation_state.is_paused(control_session_id),
                 ),
+                control_session_id,
             )
 
         if len(total_buffer) > 0:
@@ -598,6 +666,7 @@ def main():
 
             speaking_rate_state.stop()
             generation_control.finish()
+            shared_generation_state.finish(control_session_id)
             yield (
                 gr.update(value=file_path, visible=True),
                 gr.update(interactive=True),
@@ -612,10 +681,12 @@ def main():
                     final=True,
                 ),
                 *generation_button_updates(running=False),
+                control_session_id,
             )
         else:
             speaking_rate_state.stop()
             generation_control.finish()
+            shared_generation_state.finish(control_session_id)
             yield (
                 gr.update(value=None, visible=False),
                 gr.update(interactive=True),
@@ -630,6 +701,7 @@ def main():
                     final=True,
                 ),
                 *generation_button_updates(running=False),
+                control_session_id,
             )
 
     demo_app(
@@ -639,6 +711,7 @@ def main():
         synthesize_fn,
         speaking_rate_state,
         generation_control,
+        shared_generation_state,
     )
 
 

requirements.txt

@@ -1 +1 @@
-voxtream==0.2.2
+voxtream==0.2.3