Harden unified audio generation

README.md

@@ -27,7 +27,7 @@ The hosted ZeroGPU runtime uses a prebuilt `mamba-ssm==2.3.2.post1` Blackwell wh
 ## Run locally
 
 ```bash
-git clone https://huggingface.co/spaces/L0SG/Nemotron-Labs-Audex
+git clone https://huggingface.co/spaces/nvidia/Nemotron-Labs-Audex
 cd Nemotron-Labs-Audex
 bash setup_local.sh
 bash run_local.sh

app.py

@@ -53,7 +53,7 @@ DEFAULT_WHEELS_DIR = APP_DIR / ("wheels" if IS_ZERO_GPU else ".local/wheels")
 DEFAULT_BUILD_DIR = APP_DIR / ("_wheelout" if IS_ZERO_GPU else ".local/build")
 WHEELS_DIR = Path(os.environ.get("AUDEX_WHEELS_DIR", str(DEFAULT_WHEELS_DIR)))
 BUILD_DIR = Path(os.environ.get("AUDEX_BUILD_DIR", str(DEFAULT_BUILD_DIR)))
-REPO_ID = os.environ.get("SPACE_ID", "L0SG/Nemotron-Labs-Audex")
+REPO_ID = os.environ.get("SPACE_ID", "nvidia/Nemotron-Labs-Audex")
 CAUSAL_CONV1D_VERSION = "1.6.2.post1"
 MAMBA_SSM_VERSION = "2.3.2.post1"
 CAUSAL_CONV1D_SPEC = f"causal-conv1d=={CAUSAL_CONV1D_VERSION}"
@@ -89,6 +89,9 @@ MAX_TOKEN_WARNING = (
     "⚠️ Maximum new-token limit reached; this output is incomplete. "
     "Increase Max new tokens and run again."
 )
+TTS_TOKEN_WARNING = (
+    "⚠️ Maximum speech-token limit reached; generated audio may be incomplete."
+)
 TEXT_VOCAB_SIZE = 131072
 TEXT_SEED = 100
 TEXT_SYSTEM_PROMPT = (
@@ -1098,6 +1101,7 @@ def _generate_tts(
     chunks: list[np.ndarray] = []
     first_token_seconds: float | None = None
     total_tokens = 0
+    truncated = False
     request_tag = hashlib.sha256(text.encode()).hexdigest()[:8]
     started_at = time.perf_counter()
     torch.cuda.reset_peak_memory_stats()
@@ -1110,6 +1114,7 @@ def _generate_tts(
     try:
         for segment_index, segment in enumerate(segments, start=1):
             segment_tokens = 0
+            segment_truncated = False
             for event in stream_tts(
                 model=runtime.model,
                 tokenizer=runtime.tokenizer,
@@ -1122,6 +1127,8 @@ def _generate_tts(
                 guidance_scale=guidance_scale,
             ):
                 segment_tokens = event.token_count
+                if event.done:
+                    segment_truncated = event.truncated
                 if event.token_count and first_token_seconds is None:
                     first_token_seconds = time.perf_counter() - started_at
                 if event.pcm is not None:
@@ -1144,6 +1151,9 @@ def _generate_tts(
                     )
 
             total_tokens += segment_tokens
+            if segment_truncated:
+                truncated = True
+                break
             if segment_index < len(segments):
                 chunks.append(
                     np.zeros(
@@ -1164,11 +1174,17 @@ def _generate_tts(
             f"[gpu] tts_request={request_tag} tts_tokens={total_tokens} "
             f"ttfc_seconds={first_token_seconds:.2f} "
             f"tokens_per_second={token_rate:.2f} audio_seconds={waveform.size / SAMPLE_RATE:.2f} "
-            f"elapsed_seconds={elapsed_seconds:.2f} peak_memory_gib={peak_gib:.2f}",
+            f"elapsed_seconds={elapsed_seconds:.2f} peak_memory_gib={peak_gib:.2f} "
+            f"truncated={truncated}",
             flush=True,
         )
+        status = (
+            f"{TTS_TOKEN_WARNING} · {waveform.size / SAMPLE_RATE:.1f}s audio"
+            if truncated
+            else f"Complete · {waveform.size / SAMPLE_RATE:.1f}s audio"
+        )
         yield (
-            f"Complete · {waveform.size / SAMPLE_RATE:.1f}s audio",
+            status,
             wav_path,
             (
                 player_value(sequence, token_count=total_tokens, done=True)
@@ -1191,11 +1207,13 @@ def _split_sentence_segments(text: str) -> list[str]:
 
 
 def _clean_transcription(text: str) -> str:
-    first_quote = text.find("'")
-    last_quote = text.rfind("'")
-    if first_quote != -1 and last_quote > first_quote:
-        return text[first_quote + 1 : last_quote].strip()
-    return text.strip()
+    text = text.strip()
+    quoted = re.fullmatch(
+        r"(?is)(?:(?:the\s+)?(?:transcription|transcript)"
+        r"(?:\s+(?:is|reads))?\s*:?\s*)?(['\"])(.*)\1[.!]?",
+        text,
+    )
+    return quoted.group(2).strip() if quoted else text
 
 
 def _compose_s2s_input(response_prompt: str, transcript: str) -> str:
@@ -1267,7 +1285,7 @@ def _generate_s2s(
         raise gr.Error("Speech-to-speech response generation produced no final answer.")
     yield answer, thinking, gr.skip(), gr.skip()
 
-    for _, wav_path, player in _generate_tts(
+    for status, wav_path, player in _generate_tts(
         runtime,
         answer,
         S2S_TTS_MAX_NEW_TOKENS,
@@ -1278,7 +1296,12 @@ def _generate_s2s(
         token_limit=S2S_TTS_MAX_NEW_TOKENS,
         segment_sentences=True,
     ):
-        yield answer, thinking, wav_path or gr.skip(), player
+        visible_answer = (
+            f"{answer}\n\n{TTS_TOKEN_WARNING}"
+            if TTS_TOKEN_WARNING in status
+            else answer
+        )
+        yield visible_answer, thinking, wav_path or gr.skip(), player
 
 
 @spaces.GPU(duration=_estimate_text, size="xlarge")
@@ -1354,7 +1377,7 @@ def run_unified(
     guidance_scale: float,
 ) -> Iterator[tuple[object, object, object, object]]:
     """Run a selected Audex model for audio, text, or speech generation."""
-    yield "", "", gr.skip(), gr.skip()
+    yield "", "", None, gr.skip()
     runtime = _get_runtime(model_name)
     settings = TASK_SPECS[task]
     if settings["modality"] == "s2s":
@@ -1400,7 +1423,7 @@ def run_unified(
         yield "Please provide an audio input.", "", gr.skip(), gr.skip()
         return
 
-    instruction = (prompt.strip() or str(settings["template"])).replace("<sound>", "").strip()
+    instruction = prompt.strip() or str(settings["template"])
     for answer, thinking in _generate(
         runtime,
         audio,

audio_utils.py

@@ -136,10 +136,12 @@ def build_prompt_template(
         raise ValueError(f"Unknown prompt repetition mode: {prompt_repitition}")
     if prompt_repitition == "repetition":
         prompt = f"{prompt}\n{prompt}"
+    if SOUND_PLACEHOLDER not in prompt:
+        prompt = f"{SOUND_PLACEHOLDER}\n{prompt}"
 
     if reasoning:
-        return f"<|im_start|>user\n<sound>\n{prompt}<|im_end|>\n<|im_start|>assistant\n<think>\n"
-    return f"<|im_start|>user\n<sound>\n{prompt}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+        return f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n<think>\n"
+    return f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n<think></think>"
 
 
 def expand_sound_placeholder(prompt: str, num_embeddings: int) -> str:

tts_utils.py

@@ -41,6 +41,7 @@ class TTSStreamEvent:
     token_count: int
     elapsed_seconds: float
     done: bool = False
+    truncated: bool = False
 
 
 class SpeechTokenLogitsProcessor(LogitsProcessor):
@@ -256,11 +257,11 @@ def stream_tts(
     thread = Thread(target=generate, daemon=True)
     thread.start()
     token_count = 0
-    completed = False
+    terminal_token_seen = False
     try:
         for token_id in streamer:
             if token_id == token_map.end or token_id == token_map.eos:
-                completed = True
+                terminal_token_seen = True
                 break
             if not token_map.codec_start <= token_id <= token_map.codec_end:
                 raise RuntimeError(f"Unexpected token in TTS output: {token_id}")
@@ -279,24 +280,22 @@ def stream_tts(
                     token_count=token_count,
                     elapsed_seconds=time.perf_counter() - started_at,
                 )
-        else:
-            completed = True
 
         if generation_error:
             raise generation_error[0]
-        if completed:
-            for _, pcm in session.flush():
-                yield TTSStreamEvent(
-                    pcm=np.asarray(pcm, dtype=np.float32),
-                    token_count=token_count,
-                    elapsed_seconds=time.perf_counter() - started_at,
-                )
+        for _, pcm in session.flush():
             yield TTSStreamEvent(
-                pcm=None,
+                pcm=np.asarray(pcm, dtype=np.float32),
                 token_count=token_count,
                 elapsed_seconds=time.perf_counter() - started_at,
-                done=True,
             )
+        yield TTSStreamEvent(
+            pcm=None,
+            token_count=token_count,
+            elapsed_seconds=time.perf_counter() - started_at,
+            done=True,
+            truncated=not terminal_token_seen,
+        )
     finally:
         cancel_event.set()
         thread.join(timeout=10)