Improve generation limits and cancellation

README.md

@@ -20,7 +20,10 @@ The model selector defaults to the 30B-A3B and can switch to 2B model.
 
 The demo includes **audio understanding, speech recognition, speech translation, text reasoning, text-to-speech, and speech-to-speech**.
 
-Reasoning has no separate token cap by default for most tasks. Speech-to-speech defaults to a 1,536-token reasoning budget so the model has room for its final answer. `Max new tokens` sets a total cap shared by reasoning and the final answer.
+* Audio inputs support up to 15 minutes.
+* Text generation defaults to 1,024 tokens and supports up to 4,096. `Max new tokens` sets a total cap shared by reasoning and the final answer. Reasoning has no separate token cap by default for most tasks.
+* TTS uses a separate 256-token speech budget.
+* Speech-to-speech defaults to a 1,536-token reasoning budget so the model has room for its final answer.
 
 The hosted ZeroGPU runtime uses a prebuilt `mamba-ssm==2.3.2.post1` Blackwell wheel.
 

app.py

@@ -70,10 +70,10 @@ MODEL_2B_NAME = "Nemotron-Labs-Audex-2B"
 
 # Shared generation configuration
 SAMPLE_RATE = 16000
-MAX_AUDIO_DURATION_SECONDS = float(os.environ.get("AUDEX_MAX_AUDIO_SECONDS", "120"))
-MAX_NEW_TOKENS = int(os.environ.get("AUDEX_MAX_NEW_TOKENS", "2048"))
+MAX_AUDIO_DURATION_SECONDS = float(os.environ.get("AUDEX_MAX_AUDIO_SECONDS", "900"))
+MAX_NEW_TOKENS = int(os.environ.get("AUDEX_MAX_NEW_TOKENS", "4096"))
 DEFAULT_MAX_NEW_TOKENS = min(
-    int(os.environ.get("AUDEX_DEFAULT_MAX_NEW_TOKENS", "256")),
+    int(os.environ.get("AUDEX_DEFAULT_MAX_NEW_TOKENS", "1024")),
     MAX_NEW_TOKENS,
 )
 DEFAULT_REASONING_BUDGET = max(
@@ -146,9 +146,12 @@ TTS_EXAMPLE_TEXT = (
 
 # Task-specific limits and feature switches
 TTS_MAX_NEW_TOKENS = min(
-    int(os.environ.get("AUDEX_TTS_MAX_NEW_TOKENS", "192")),
+    int(os.environ.get("AUDEX_TTS_MAX_NEW_TOKENS", "256")),
     MAX_NEW_TOKENS,
 )
+TTS_30B_GPU_DURATION_SECONDS = int(
+    os.environ.get("AUDEX_TTS_30B_GPU_DURATION_SECONDS", "90")
+)
 S2S_TEXT_MAX_NEW_TOKENS = int(os.environ.get("AUDEX_S2S_TEXT_MAX_NEW_TOKENS", "2048"))
 S2S_REASONING_BUDGET = int(os.environ.get("AUDEX_S2S_REASONING_BUDGET", "1536"))
 S2S_GPU_DURATION_SECONDS = int(os.environ.get("AUDEX_S2S_GPU_DURATION_SECONDS", "120"))
@@ -833,7 +836,7 @@ def _stream_model_generate(
         )
     finally:
         cancel_event.set()
-        thread.join(timeout=10)
+        thread.join()
 
 
 def _generate(
@@ -1336,6 +1339,7 @@ def _estimate_unified(
     reasoning: bool,
     reasoning_budget: int,
     max_new_tokens: int,
+    tts_max_new_tokens: int,
     temperature: float,
     top_p: float,
     guidance_scale: float,
@@ -1344,9 +1348,11 @@ def _estimate_unified(
 ) -> int:
     if task == S2S_TASK:
         return S2S_GPU_DURATION_SECONDS
-    if task in {TEXT_TASK, TTS_TASK}:
-        if task == TTS_TASK and int(max_new_tokens) > TTS_MAX_NEW_TOKENS:
-            return 10
+    if task == TTS_TASK:
+        if model_name == DEFAULT_MODEL_NAME:
+            return TTS_30B_GPU_DURATION_SECONDS
+        return MAX_GPU_DURATION_SECONDS
+    if task == TEXT_TASK:
         return _estimate_text(
             prompt,
             reasoning,
@@ -1374,6 +1380,7 @@ def run_unified(
     reasoning: bool,
     reasoning_budget: int,
     max_new_tokens: int,
+    tts_max_new_tokens: int,
     temperature: float,
     top_p: float,
     guidance_scale: float,
@@ -1402,7 +1409,7 @@ def run_unified(
         for status, wav_path, player in _generate_tts(
             runtime,
             prompt,
-            int(max_new_tokens),
+            int(tts_max_new_tokens),
             temperature,
             top_p,
             guidance_scale=float(guidance_scale),
@@ -1449,7 +1456,6 @@ def unified_task_defaults(
     has_speech_output = settings["modality"] in {"tts", "s2s"}
     reasoning = bool(settings["reasoning"])
     default_max_new_tokens = int(settings.get("max_new_tokens", DEFAULT_MAX_NEW_TOKENS))
-    max_new_tokens_limit = int(settings.get("max_new_tokens_limit", MAX_NEW_TOKENS))
     reasoning_budget_limit = _max_reasoning_budget(default_max_new_tokens)
     return (
         gr.update(visible=is_audio),
@@ -1479,9 +1485,10 @@ def unified_task_defaults(
         float(settings["temperature"]),
         float(settings["top_p"]),
         gr.update(
-            value=TTS_MAX_NEW_TOKENS if is_tts else default_max_new_tokens,
-            maximum=TTS_MAX_NEW_TOKENS if is_tts else max_new_tokens_limit,
+            value=default_max_new_tokens,
+            visible=not is_tts,
         ),
+        gr.update(value=TTS_MAX_NEW_TOKENS, visible=is_tts),
         gr.update(
             value=float(settings.get("guidance_scale", 2.0)),
             visible=has_speech_output,
@@ -1577,6 +1584,17 @@ with gr.Blocks(title="Nemotron-Labs-Audex") as demo:
                         "If reached, the displayed output is incomplete."
                     ),
                 )
+                unified_tts_max_new_tokens = gr.Slider(
+                    16,
+                    TTS_MAX_NEW_TOKENS,
+                    value=TTS_MAX_NEW_TOKENS,
+                    step=16,
+                    label="Max speech tokens",
+                    info=(
+                        "Speech-token cap. If reached, generated audio may be incomplete."
+                    ),
+                    visible=False,
+                )
                 unified_temperature = gr.Slider(
                     0.1,
                     1.5,
@@ -1638,6 +1656,7 @@ with gr.Blocks(title="Nemotron-Labs-Audex") as demo:
             unified_temperature,
             unified_top_p,
             unified_max_new_tokens,
+            unified_tts_max_new_tokens,
             unified_guidance_scale,
             unified_answer_out,
             unified_thinking_out,
@@ -1669,6 +1688,7 @@ with gr.Blocks(title="Nemotron-Labs-Audex") as demo:
             unified_reasoning,
             unified_reasoning_budget,
             unified_max_new_tokens,
+            unified_tts_max_new_tokens,
             unified_temperature,
             unified_top_p,
             unified_guidance_scale,

run_local.sh

@@ -35,9 +35,9 @@ export AUDEX_TORCH_ARCH="${AUDEX_TORCH_ARCH:-$detected_arch}"
 export AUDEX_WHEELS_DIR="${AUDEX_WHEELS_DIR:-$APP_DIR/.local/wheels/$runtime_tag}"
 export AUDEX_BUILD_DIR="${AUDEX_BUILD_DIR:-$APP_DIR/.local/build/$runtime_tag}"
 export AUDEX_MAX_AUDIO_SECONDS="${AUDEX_MAX_AUDIO_SECONDS:-900}"
-export AUDEX_MAX_NEW_TOKENS="${AUDEX_MAX_NEW_TOKENS:-2048}"
+export AUDEX_MAX_NEW_TOKENS="${AUDEX_MAX_NEW_TOKENS:-4096}"
 export AUDEX_DEFAULT_MAX_NEW_TOKENS="${AUDEX_DEFAULT_MAX_NEW_TOKENS:-1024}"
-export AUDEX_TTS_MAX_NEW_TOKENS="${AUDEX_TTS_MAX_NEW_TOKENS:-1024}"
+export AUDEX_TTS_MAX_NEW_TOKENS="${AUDEX_TTS_MAX_NEW_TOKENS:-256}"
 export AUDEX_S2S_TEXT_MAX_NEW_TOKENS="${AUDEX_S2S_TEXT_MAX_NEW_TOKENS:-2048}"
 export AUDEX_S2S_TTS_MAX_NEW_TOKENS="${AUDEX_S2S_TTS_MAX_NEW_TOKENS:-2400}"
 export GRADIO_SERVER_NAME="${GRADIO_SERVER_NAME:-0.0.0.0}"

tts_utils.py

@@ -298,7 +298,7 @@ def stream_tts(
         )
     finally:
         cancel_event.set()
-        thread.join(timeout=10)
+        thread.join()
 
 
 def encode_pcm_chunk(pcm: np.ndarray) -> str: