def _patch_asyncio_event_loop_del(): """ Patch a noisy asyncio teardown issue sometimes seen in Spaces environments. In some runtime/container combinations, Python may try to close an already invalid file descriptor when the event loop is garbage-collected. We silence only that specific harmless case. """ try: import asyncio.base_events as base_events original_del = getattr(base_events.BaseEventLoop, "__del__", None) if original_del is None: return def patched_del(self): try: original_del(self) except ValueError as e: if "Invalid file descriptor" not in str(e): raise base_events.BaseEventLoop.__del__ = patched_del except Exception: pass _patch_asyncio_event_loop_del() """ Gradio app for TADA inference with Hugging Face OAuth login. Usage on Hugging Face Spaces: - Add to README.md front matter: hf_oauth: true hf_oauth_scopes: - gated-repos This app expects users to sign in with Hugging Face before loading a gated model. Design goals: - Work on fixed GPU environments - Work on ZeroGPU Spaces - Keep the encoder on CPU by default to reduce VRAM pressure - Move the TADA model to the runtime device only when needed """ import dataclasses import html import json import logging import os import shutil import tempfile import time from typing import Optional import gradio as gr import torch import torchaudio try: import spaces gpu_decorator = spaces.GPU except ImportError: gpu_decorator = lambda fn=None, **kw: fn if fn else (lambda f: f) from tada.modules.encoder import Encoder, EncoderOutput # noqa: E402 from tada.modules.tada import InferenceOptions, TadaForCausalLM # noqa: E402 from tada.utils.text import normalize_text as normalize_text_fn # noqa: E402 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- _SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) _SAMPLES_DIR = os.path.join(_SCRIPT_DIR, "samples") _AUDIO_EXTENSIONS = (".wav", ".mp3", ".flac") LANGUAGE_MAP = { "English": None, "German": "de", "Japanese": "ja", } _CODE_TO_LANG_DIR = {None: "en", "de": "de", "ja": "ja"} _MODEL_CHOICES = ["HumeAI/tada-1b", "HumeAI/tada-3b-ml"] _DEFAULT_MODEL = "HumeAI/tada-3b-ml" _MULTILINGUAL_MODELS = {"HumeAI/tada-3b-ml"} _GATED_REPO_ID = "meta-llama/Llama-3.2-1B" _GATED_REPO_URL = f"https://huggingface.co/{_GATED_REPO_ID}" # Keep encoder on CPU by default. This is safer for memory-constrained setups, # especially because the encoder may lazily load ASR/alignment subcomponents. _DEFAULT_ENCODER_DEVICE = os.environ.get("TADA_ENCODER_DEVICE", "cpu").strip().lower() # By default, move the model back to CPU after generation. This is helpful for # ZeroGPU and conservative for fixed GPU setups. Override if needed. _OFFLOAD_MODEL_AFTER_INFERENCE = os.environ.get( "TADA_OFFLOAD_MODEL_AFTER_INFERENCE", "true" ).strip().lower() in {"1", "true", "yes", "on"} if "PYTORCH_CUDA_ALLOC_CONF" not in os.environ: os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # --------------------------------------------------------------------------- # Runtime device helpers # --------------------------------------------------------------------------- def _get_runtime_model_device() -> str: """ Resolve the model device at call time. This matters for ZeroGPU: GPU availability may change dynamically between app startup and the execution of a @spaces.GPU-decorated function. """ if torch.cuda.is_available(): return "cuda" if torch.backends.mps.is_available(): return "mps" return "cpu" def _get_encoder_device() -> str: """ Resolve encoder device. Allowed: - cpu - cuda - mps - auto """ requested = _DEFAULT_ENCODER_DEVICE if requested == "auto": return _get_runtime_model_device() if requested == "cuda" and not torch.cuda.is_available(): return "cpu" if requested == "mps" and not torch.backends.mps.is_available(): return "cpu" if requested not in {"cpu", "cuda", "mps"}: return "cpu" return requested _ENCODER_DEVICE = _get_encoder_device() # --------------------------------------------------------------------------- # Sample discovery and transcripts # --------------------------------------------------------------------------- def _discover_preset_samples(lang_code: str | None = None) -> dict[str, str]: """Return {display_name: absolute_path} for audio files in samples//.""" presets: dict[str, str] = {} lang_dir = _CODE_TO_LANG_DIR.get(lang_code, "en") search_dir = os.path.join(_SAMPLES_DIR, lang_dir) if not os.path.isdir(search_dir): return presets for fname in sorted(os.listdir(search_dir)): if fname.lower().endswith(_AUDIO_EXTENSIONS): presets[fname] = os.path.join(search_dir, fname) return presets def _load_preset_transcripts(lang_code: str | None = None) -> dict[str, str]: """Load synth_transcripts.json for the selected language.""" lang_dir = _CODE_TO_LANG_DIR.get(lang_code, "en") candidate = os.path.join(_SAMPLES_DIR, lang_dir, "synth_transcripts.json") if os.path.isfile(candidate): with open(candidate) as f: return json.load(f) return {} def _load_prompt_transcripts(lang_code: str | None = None) -> dict[str, str]: """ Load prompt_transcripts.json for the selected language. For preset reference audios, these transcripts let us bypass built-in ASR and provide stable alignment text directly to the encoder. """ lang_dir = _CODE_TO_LANG_DIR.get(lang_code, "en") candidate = os.path.join(_SAMPLES_DIR, lang_dir, "prompt_transcripts.json") if os.path.isfile(candidate): with open(candidate) as f: return json.load(f) return {} def _language_choices_for_model(model_name: str) -> list[str]: """Return language options supported by the selected model.""" if model_name in _MULTILINGUAL_MODELS: return list(LANGUAGE_MAP.keys()) return ["English"] def _get_language_assets(language: str) -> tuple[dict[str, str], dict[str, str], dict[str, str]]: """Return (samples, synth_transcripts, prompt_transcripts) for a language.""" lang_code = LANGUAGE_MAP.get(language) return ( _discover_preset_samples(lang_code), _load_preset_transcripts(lang_code), _load_prompt_transcripts(lang_code), ) _PRESET_SAMPLES, _PRESET_TRANSCRIPTS, _PROMPT_TRANSCRIPTS = _get_language_assets("English") logger.info( "Discovered %d preset audio samples, %d transcripts", len(_PRESET_SAMPLES), len(_PRESET_TRANSCRIPTS), ) # --------------------------------------------------------------------------- # Global state # --------------------------------------------------------------------------- _encoder_cache: dict[str | None, Encoder] = {} _model: Optional[TadaForCausalLM] = None _current_model_name: str = "" _current_language: str | None = None _hf_token: str | None = None # --------------------------------------------------------------------------- # Utility helpers # --------------------------------------------------------------------------- def _move_encoder_output(output: EncoderOutput, device: str) -> EncoderOutput: """Move all tensor fields inside an EncoderOutput to the target device.""" kwargs = {} for f in dataclasses.fields(output): val = getattr(output, f.name) if isinstance(val, torch.Tensor): kwargs[f.name] = val.to(device) else: kwargs[f.name] = val return EncoderOutput(**kwargs) def _get_device_info() -> str: """Return a human-readable description of the currently available runtime device.""" if torch.cuda.is_available(): names = [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())] return f"CUDA - {', '.join(names)}" if torch.backends.mps.is_available(): return "MPS (Apple Silicon)" return "CPU" def _status_md(text: str) -> str: """Format a short status message for markdown display.""" return f"**Status:** {text}" def _build_prompt_help_html(message: str) -> str: """Simple helper text block used for short statuses.""" return ( '
' f"{html.escape(message)}" "
" ) def _build_info_panel_html( title: str, paragraphs: list[str], bullets: list[str] | None = None, eli5: str | None = None, ) -> str: """ Build a richer explanatory panel for empty token/result boxes. """ paragraph_html = "".join( f"

{html.escape(paragraph)}

" for paragraph in paragraphs ) bullets_html = "" if bullets: bullets_html = ( "" ) eli5_html = "" if eli5: eli5_html = ( "
" "ELI5: " f"{html.escape(eli5)}" "
" ) return ( "
" f"
{html.escape(title)}
" f"{paragraph_html}" f"{bullets_html}" f"{eli5_html}" "
" ) def _reference_info_panel() -> str: """Placeholder/explainer for the reference alignment box.""" return _build_info_panel_html( "How TADA prepares a reference voice", [ "Start by choosing how to provide the reference voice: select a preset sample, upload your own recording, or record one live with the microphone.", "When you click Prepare Reference Voice, TADA encodes the reference audio and aligns it to text-like token positions. The goal is not just to store a voice example, but to extract speaker, style, and timing cues that guide generation.", "This preparation stage turns raw audio into a structured acoustic prompt. Instead of using the reference only as a loose example, TADA builds a representation that can stay synchronized with text during later generation.", "Once preparation is complete, this panel will display the reference alignment trace derived from the audio.", ], bullets=[ "Highlighted tokens show text units aligned with the reference audio.", "Grey dots indicate timing positions or gaps between aligned tokens.", "The header summarizes token count and reference audio duration.", ], eli5=( "Imagine teaching the model how someone speaks before asking it to talk. " "The model listens to the reference voice, figures out where words would fall in time, " "and stores those timing and style patterns so it can speak in a similar way later." ), ) def _generation_info_panel() -> str: """Placeholder/explainer for the generation trace box.""" return _build_info_panel_html( "How TADA generates speech", [ "Enter text manually or choose a preset, then click Generate Speech to create audio conditioned on the prepared reference voice.", "TADA generates speech using a synchronized text-acoustic process. Instead of letting text and audio drift apart into two loosely connected streams, the model keeps linguistic content and acoustic structure aligned during generation.", "This synchronized design is the key idea behind TADA. Each step helps decide both what is being said and how it should sound, which makes the process easier to inspect in a research demo like this one.", "After generation, this panel will display the text-acoustic trace used to produce the final audio.", ], bullets=[ "Highlighted tokens show generated text units involved in the output.", "Grey dots represent acoustic timing steps or spacing between token events.", "The header reports generation length, audio duration, wall time, and real-time factor.", ], eli5=( "Think of the model reading and speaking at the same time. " "At each step, it decides both the next word and how that word should sound. " "Because the two stay synchronized, the speech can come out faster and more consistently." ), ) def _gated_repo_access_text() -> str: """Return a user-facing message explaining gated repo access.""" return ( f"Your Hugging Face account may not yet have access to the gated repo " f"{_GATED_REPO_ID}. Request or verify gated-repo access here: {_GATED_REPO_URL}" ) def _is_gated_repo_error(exc: Exception) -> bool: """Best-effort detection for gated repo access failures.""" text = str(exc).lower() return ( "gated repo" in text or "cannot access gated repo" in text or "access to model" in text or "401 client error" in text or _GATED_REPO_ID.lower() in text ) def _is_oom_error(exc: Exception) -> bool: """Best-effort detection for memory errors.""" if isinstance(exc, torch.OutOfMemoryError): return True text = str(exc).lower() return "out of memory" in text or "cuda out of memory" in text def _clear_cuda_memory(): """ Best-effort CUDA cleanup. Important for ZeroGPU: do not force CUDA initialization from the main process. """ try: if not torch.cuda.is_available(): return torch.cuda.empty_cache() except Exception: pass def _move_model_to_device(device: str): """Move the model and decoder to the requested runtime device.""" global _model if _model is None: raise gr.Error("Model is not loaded yet.") _model.to(device) _model.decoder.to(device) def _offload_model_if_needed(): """ Optionally move the model back to CPU after inference. """ global _model if _model is None or not _OFFLOAD_MODEL_AFTER_INFERENCE: return try: _model.to("cpu") _clear_cuda_memory() except Exception: logger.exception("Failed to offload model back to CPU") # --------------------------------------------------------------------------- # Hugging Face auth helpers # --------------------------------------------------------------------------- def _hf_auth_status(profile: gr.OAuthProfile | None) -> str: """Return a user-facing Hugging Face auth status string.""" if profile is None: return ( "**Hugging Face:** not signed in \n" f"To use this demo, sign in and make sure your account has gated-repo access to " f"[{_GATED_REPO_ID}]({_GATED_REPO_URL})." ) return ( f"**Hugging Face:** signed in as `{profile.username}` \n" f"If loading fails, verify gated-repo access to [{_GATED_REPO_ID}]({_GATED_REPO_URL})." ) def _ensure_hf_auth(oauth_token: gr.OAuthToken | None = None) -> str: """ Register the user OAuth token for downstream huggingface_hub / transformers calls. """ global _hf_token token_value = None if oauth_token is not None and getattr(oauth_token, "token", None): token_value = oauth_token.token elif _hf_token: token_value = _hf_token if not token_value: raise gr.Error( "Please sign in with Hugging Face first. " + _gated_repo_access_text() ) if token_value != _hf_token: logger.info("Registering user OAuth token for downstream HF libraries") _hf_token = token_value os.environ["HF_TOKEN"] = token_value return token_value # --------------------------------------------------------------------------- # Encoder/model loading # --------------------------------------------------------------------------- def get_encoder( language_code: str | None = None, oauth_token: gr.OAuthToken | None = None, ) -> Encoder: """ Get or create an encoder for the given language. """ _ensure_hf_auth(oauth_token) if language_code not in _encoder_cache: logger.info("Loading encoder for language=%s on %s", language_code, _ENCODER_DEVICE) _encoder_cache[language_code] = Encoder.from_pretrained( "HumeAI/tada-codec", language=language_code, ).to(_ENCODER_DEVICE) return _encoder_cache[language_code] def load_models( model_name: str = _DEFAULT_MODEL, language: str = "English", oauth_token: gr.OAuthToken | None = None, progress: gr.Progress = gr.Progress(track_tqdm=True), ) -> str: """ Load the encoder and TADA model. Important design choice: - The encoder lives on its configured device (CPU by default). - The TADA model is loaded and staged on CPU. - During generation, the model is moved to the runtime device on demand. """ global _model, _current_model_name, _current_language if oauth_token is None or not getattr(oauth_token, "token", None): raise gr.Error( "Please sign in with Hugging Face first. " + _gated_repo_access_text() ) language_code = LANGUAGE_MAP.get(language) _current_language = language_code try: progress(0.05, desc="Authenticating with Hugging Face") _ensure_hf_auth(oauth_token) if _model is not None and _current_model_name == model_name: progress(0.25, desc="Checking encoder") get_encoder(language_code, oauth_token=oauth_token) progress(1.0, desc="Model already loaded") return ( f"Loaded: {model_name} | runtime device {_get_runtime_model_device().upper()} | " f"encoder on {_ENCODER_DEVICE.upper()}" ) if _model is not None: progress(0.15, desc="Clearing previous model from memory") del _model _model = None progress(0.35, desc="Loading encoder") get_encoder(language_code, oauth_token=oauth_token) progress(0.55, desc="Loading TADA model weights") _model = TadaForCausalLM.from_pretrained( model_name, token=oauth_token.token, ) # Keep model on CPU here. Runtime GPU moves happen only in decorated functions. _model.to("cpu") progress(0.90, desc="Finalizing model setup") _current_model_name = model_name status = ( f"Loaded: {model_name} | model staged on CPU, runtime device {_get_runtime_model_device().upper()} | " f"encoder on {_ENCODER_DEVICE.upper()}" ) logger.info(status) progress(1.0, desc="Model ready") return status except Exception as e: logger.exception("Model loading failed") if _is_gated_repo_error(e): raise gr.Error( "Failed to load the model because this Hugging Face account does not yet " f"have access to {_GATED_REPO_ID}. " f"Request or verify gated-repo access here: {_GATED_REPO_URL}" ) from e if _is_oom_error(e): raise gr.Error( "Model loading failed due to insufficient memory. " "Try a larger Space, a smaller model, or restart the Space." ) from e raise gr.Error(f"Failed to load the model: {e}") from e # --------------------------------------------------------------------------- # Reference encoding and alignment # --------------------------------------------------------------------------- def _encode_prompt( audio_path: str | None, language_code: str | None = None, prompt_text: str | None = None, oauth_token: gr.OAuthToken | None = None, ) -> EncoderOutput: """ Encode a reference audio file into an EncoderOutput prompt. """ if audio_path is None or audio_path == "": return EncoderOutput.empty(_ENCODER_DEVICE) encoder = get_encoder(language_code, oauth_token=oauth_token) audio, sample_rate = torchaudio.load(audio_path) audio = audio.mean(dim=0, keepdim=True) audio = audio / audio.abs().max().clamp(min=1e-8) * 0.95 audio = audio.to(_ENCODER_DEVICE) text_kwarg = [prompt_text] if prompt_text else None prompt = encoder(audio, text=text_kwarg, sample_rate=sample_rate) return prompt def _decode_tokens_individually(tokenizer, token_ids: list[int]) -> list[str]: """Decode token IDs into user-readable per-token strings.""" labels: list[str] = [] for i in range(len(token_ids)): prefix = tokenizer.decode(token_ids[:i], skip_special_tokens=True) full = tokenizer.decode(token_ids[: i + 1], skip_special_tokens=True) token_str = full[len(prefix) :] labels.append(token_str) return labels def _format_token_alignment( prompt: EncoderOutput, language_code: str | None = None, oauth_token: gr.OAuthToken | None = None, ) -> str: """ Build an HTML representation of reference alignment. """ if prompt.text_tokens is None or prompt.token_positions is None: return "" encoder = get_encoder(language_code, oauth_token=oauth_token) tokenizer = encoder.tokenizer n_tokens = ( int(prompt.text_tokens_len[0].item()) if prompt.text_tokens_len is not None else prompt.text_tokens.shape[1] ) token_ids = prompt.text_tokens[0, :n_tokens].cpu().tolist() positions = prompt.token_positions[0, :n_tokens].cpu().long().tolist() labels = _decode_tokens_individually(tokenizer, token_ids) audio_dur = prompt.audio.shape[-1] / prompt.sample_rate if prompt.audio.numel() > 0 else 0.0 header = f"{n_tokens} tokens | {audio_dur:.2f}s audio" parts: list[str] = [] prev_pos = 0 for pos, label in zip(positions, labels): gap = max(0, pos - prev_pos) if gap > 0: parts.append(f'{"." * gap}') escaped = html.escape(label) parts.append( '{escaped}' ) prev_pos = pos + 1 body = "".join(parts) return ( '
' f'
{header}
' f"{body}
" ) @gpu_decorator @torch.inference_mode() def process_prompt( audio_path: str | None, language: str = "English", oauth_token: gr.OAuthToken | None = None, progress: gr.Progress = gr.Progress(track_tqdm=True), ) -> tuple[str, EncoderOutput]: """ Prepare the reference voice and return (alignment_html, prompt_on_cpu). """ global _current_language language_code = LANGUAGE_MAP.get(language) _current_language = language_code try: progress(0.05, desc="Authenticating") _ensure_hf_auth(oauth_token) progress(0.15, desc="Loading reference encoder") encoder = get_encoder(language_code, oauth_token=oauth_token) encoder.to(_ENCODER_DEVICE) if audio_path is None or audio_path == "": progress(0.90, desc="Preparing zero-shot prompt") prompt = EncoderOutput.empty(_ENCODER_DEVICE) prompt_cpu = _move_encoder_output(prompt, "cpu") progress(1.0, desc="Reference ready") return ( _build_prompt_help_html( "No reference audio provided. Zero-shot mode is ready. " "In this mode, generation relies on text alone instead of a prepared acoustic reference." ), prompt_cpu, ) progress(0.35, desc="Loading reference audio") prompt_text = None audio_fname = os.path.basename(audio_path) for key in (audio_fname, audio_fname.replace("tada_preset_", "")): if key in _PROMPT_TRANSCRIPTS: prompt_text = _PROMPT_TRANSCRIPTS[key] break progress(0.55, desc="Preparing transcript and alignment") prompt = _encode_prompt( audio_path, language_code, prompt_text=prompt_text, oauth_token=oauth_token, ) progress(0.85, desc="Formatting reference analysis") alignment_html = _format_token_alignment( prompt, language_code, oauth_token=oauth_token, ) prompt_cpu = _move_encoder_output(prompt, "cpu") if not alignment_html: alignment_html = _build_prompt_help_html("Reference prepared successfully.") progress(1.0, desc="Reference ready") return alignment_html, prompt_cpu except Exception as e: logger.exception("Prompt processing failed") if _is_gated_repo_error(e): raise gr.Error( "Prompt preparation failed because this Hugging Face account does not yet " f"have access to {_GATED_REPO_ID}. " f"Request or verify access here: {_GATED_REPO_URL}" ) from e if _is_oom_error(e): raise gr.Error( "Prompt preparation ran out of memory. " "Try a larger GPU, restart the Space, or use a preset/reference audio with a known transcript." ) from e raise gr.Error(f"Prompt preparation failed: {e}") from e # --------------------------------------------------------------------------- # Generation and generation trace formatting # --------------------------------------------------------------------------- def _decode_byte_tokens(raw_tokens: list[str]) -> list[str]: """Decode byte-level model tokens into readable strings.""" if not raw_tokens or _model is None: return raw_tokens try: tokenizer = _model.tokenizer token_ids = tokenizer.convert_tokens_to_ids(raw_tokens) return _decode_tokens_individually(tokenizer, token_ids) except Exception: return [t.replace("\u0120", " ") for t in raw_tokens] def _format_step_logs(step_logs: list[dict], audio_duration: float, wall_time: float) -> str: """ Build an HTML view of generation step logs. """ if not step_logs: return "" n_tokens = len(step_logs) total_frames = sum(entry.get("n_frames_before", 0) for entry in step_logs) rtf = wall_time / audio_duration if audio_duration > 0 else float("inf") header = ( f"{n_tokens} steps | {audio_duration:.1f}s audio | {total_frames} frames | " f"{wall_time:.1f}s wall | RTF {rtf:.2f}" ) raw_tokens = [entry.get("token", "") for entry in step_logs] labels = _decode_byte_tokens(raw_tokens) parts: list[str] = [] for entry, label in zip(step_logs, labels): n_frames = entry.get("n_frames_before", 0) if n_frames > 0: parts.append(f'{"." * n_frames}') escaped = html.escape(label) parts.append( '{escaped}' ) body = "".join(parts) return ( '
' f'
{header}
' f"{body}
" ) @gpu_decorator @torch.inference_mode() def generate_speech( text: str, num_extra_steps: float = 0, noise_temperature: float = 0.9, acoustic_cfg_scale: float = 2.0, duration_cfg_scale: float = 2.0, num_flow_matching_steps: float = 20, negative_condition_source: str = "negative_step_output", text_only_logit_scale: float = 0.0, num_acoustic_candidates: float = 1, scorer: str = "likelihood", spkr_verification_weight: float = 1.0, speed_up_factor: float = 0.0, normalize_text: bool = True, cached_prompt: EncoderOutput | None = None, progress: gr.Progress = gr.Progress(track_tqdm=True), ) -> tuple[str | None, str]: """ Generate speech from text using the prepared reference prompt. """ if _model is None: raise gr.Error("Models are not loaded yet. Sign in and click 'Load Model' first.") if not text or not text.strip(): raise gr.Error("Please enter some text to synthesize.") runtime_device = _get_runtime_model_device() try: progress(0.05, desc="Preparing generation") _move_model_to_device(runtime_device) if cached_prompt is None: prompt = EncoderOutput.empty(runtime_device) else: prompt = _move_encoder_output(cached_prompt, runtime_device) logger.info("Generating speech for text: %s", text) speed = float(speed_up_factor) if speed_up_factor > 0 else None progress(0.20, desc=f"Using device: {runtime_device.upper()}") t0 = time.time() progress(0.40, desc="Running speech generation") output = _model.generate( prompt=prompt, text=text, num_transition_steps=0, num_extra_steps=int(num_extra_steps), normalize_text=normalize_text, inference_options=InferenceOptions( acoustic_cfg_scale=float(acoustic_cfg_scale), duration_cfg_scale=float(duration_cfg_scale), num_flow_matching_steps=int(num_flow_matching_steps), noise_temperature=float(noise_temperature), speed_up_factor=speed, time_schedule="logsnr", negative_condition_source=negative_condition_source, text_only_logit_scale=float(text_only_logit_scale), num_acoustic_candidates=int(num_acoustic_candidates), scorer=scorer, spkr_verification_weight=float(spkr_verification_weight), ), system_prompt="", ) wall_time = time.time() - t0 progress(0.85, desc="Saving generated audio") wav = output.audio[0].detach().cpu().float() if wav.dim() == 1: wav = wav.unsqueeze(0) tmp_path = os.path.join(tempfile.gettempdir(), f"tada_output_{id(output)}.wav") torchaudio.save(tmp_path, wav, 24_000) audio_duration = wav.shape[-1] / 24_000 progress(0.92, desc="Formatting generation trace") all_logs = output.step_logs or [] if _model is not None and text and output.input_text_ids is not None: input_ids = output.input_text_ids[0] seq_len = input_ids.shape[0] n_eos = _model.config.shift_acoustic normalized = normalize_text_fn(text) if normalize_text else text n_text_tokens = len(_model.tokenizer.encode(normalized, add_special_tokens=False)) text_end = seq_len - n_eos text_start = text_end - n_text_tokens log_by_step = {e["step"]: e for e in all_logs} text_logs = [] for step in range(text_start, text_end): if step in log_by_step: text_logs.append(log_by_step[step]) else: token_id = input_ids[step].item() token_str = _model.tokenizer.convert_ids_to_tokens([token_id])[0] text_logs.append( { "step": step, "token": token_str, "n_frames_before": 0, "n_frames_after": 0, "n_frames_src": "prefilled", "acoustic_mask": 1, "acoustic_feat_src": "prefilled", "acoustic_feat_norm": 0.0, } ) generated_logs = text_logs else: generated_logs = all_logs generated_html = _format_step_logs(generated_logs, audio_duration, wall_time) if not generated_html: generated_html = _build_prompt_help_html("Generation completed.") progress(1.0, desc="Generation complete") return tmp_path, generated_html except gr.Error: raise except Exception as e: logger.exception("Generation failed") if _is_oom_error(e): raise gr.Error( "Generation ran out of GPU memory. Try a larger GPU, a shorter text, or restart the Space." ) from e raise gr.Error(f"Generation failed: {e}") from e finally: _offload_model_if_needed() # --------------------------------------------------------------------------- # UI helper callbacks # --------------------------------------------------------------------------- def _default_text_for_language(language: str) -> tuple[list[str], str, str]: """Return (choices, default_key, default_text) for synth transcript presets.""" _, synth_transcripts, _ = _get_language_assets(language) transcript_choices = ["(custom)"] + list(synth_transcripts.keys()) default_key = transcript_choices[1] if len(transcript_choices) > 1 else "(custom)" default_text = synth_transcripts.get(default_key, "") if default_key != "(custom)" else "" return transcript_choices, default_key, default_text def _default_voice_for_language(language: str) -> tuple[list[str], str, str | None]: """Return (choices, default_choice, default_audio_path) for reference voice presets.""" samples, _, _ = _get_language_assets(language) preset_choices = ["None (zero-shot)"] + list(samples.keys()) default_choice = preset_choices[1] if len(preset_choices) > 1 else "None (zero-shot)" if default_choice == "None (zero-shot)": return preset_choices, default_choice, None return preset_choices, default_choice, samples.get(default_choice) def _on_model_selected(model_name: str): """Update the language dropdown when the selected model changes.""" choices = _language_choices_for_model(model_name) value = "English" if "English" in choices else choices[0] return gr.update(choices=choices, value=value) def _on_preset_selected(choice: str, language: str) -> str | None: """ Resolve a selected preset into a temporary file path. """ samples, _, _ = _get_language_assets(language) if choice == "None (zero-shot)": return None path = samples.get(choice) if path is None: return None tmp_path = os.path.join(tempfile.gettempdir(), f"tada_preset_{choice}") shutil.copy2(path, tmp_path) return tmp_path def _on_transcript_selected(choice: str, language: str) -> str: """Populate the text box from a transcript preset.""" _, synth_transcripts, _ = _get_language_assets(language) if choice == "(custom)": return "" return synth_transcripts.get(choice, "") def _on_language_changed(language: str): """ Refresh all language-dependent assets. """ global _PRESET_SAMPLES, _PRESET_TRANSCRIPTS, _PROMPT_TRANSCRIPTS gr.Info("Language changed. Please prepare the reference voice again.") _PRESET_SAMPLES, _PRESET_TRANSCRIPTS, _PROMPT_TRANSCRIPTS = _get_language_assets(language) voice_choices, voice_default, voice_path = _default_voice_for_language(language) transcript_choices, transcript_default, transcript_text = _default_text_for_language(language) return ( gr.update(choices=voice_choices, value=voice_default), voice_path, gr.update(choices=transcript_choices, value=transcript_default), transcript_text, None, gr.update(interactive=False), _status_md("Language changed. Prepare the reference voice again."), _reference_info_panel(), ) def _invalidate_prompt_state(message: str, notify: bool = False): """Clear the cached prompt and disable generation until reference is re-prepared.""" if notify: gr.Info(message) return ( None, gr.update(interactive=False), _status_md(message), _reference_info_panel(), ) def _before_prepare(): """UI updates right before preparing the reference voice.""" gr.Info("Preparing the reference voice. This extracts the acoustic prompt used for generation.") return ( gr.update(value="Preparing reference...", interactive=False), _status_md("Preparing reference voice..."), ) def _after_prepare_success(): """UI updates after successful reference preparation.""" gr.Info("Reference voice is ready. You can now generate speech.") return ( gr.update(value="Prepare Reference Voice", interactive=True), gr.update(interactive=True), _status_md("Reference ready. You can now generate speech."), ) def _before_load_model(): """UI updates right before model loading.""" gr.Info("Loading model on CPU first. On ZeroGPU, the GPU is only used during inference.") return ( gr.update(value="Loading model...", interactive=False), "Loading model...", ) def _after_load_model(): """UI updates after successful model loading.""" gr.Info("Model loaded. Next, prepare a reference voice.") return ( gr.update(value="Load Model", interactive=True), gr.update(interactive=True), _status_md("Model loaded. You can now prepare a reference voice."), ) def _before_generate(): """UI updates right before generation.""" gr.Info("Generating speech from your text and prepared reference voice...") return _status_md("Generating speech...") def _after_generate(): """UI updates after successful generation.""" gr.Info("Generation complete. You can listen to the audio and inspect the trace.") return _status_md("Generation completed.") # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- def _top_links_html() -> str: paper_url = "https://huggingface.co/papers/2602.23068" duplicate_url = "https://huggingface.co/spaces/fffiloni/tada-dual-alignment-tts-demo?duplicate=true" profile_url = "https://huggingface.co/fffiloni" pro_url = "https://huggingface.co/pro" return f"""
Paper page Duplicate this Space Follow me on HF Subscribe to PRO
""" def build_ui() -> gr.Blocks: with gr.Blocks(title="TADA Inference") as demo: prompt_state = gr.State(value=None) with gr.Column(elem_classes=["hero-header"]): gr.Markdown( "# TADA: A Generative Framework for Speech Modeling via Text-Acoustic Dual Alignment" ) gr.Markdown( "A unified speech-language model that synchronizes speech and text into a single, cohesive stream via 1:1 alignment.", elem_classes=["hero-subtitle"], ) gr.HTML(_top_links_html()) with gr.Column(elem_classes=["section-card"]): with gr.Row(equal_height=True): with gr.Column(scale=4, elem_classes=["split-left"]): gr.Markdown("## 1. Model & Generation Settings", elem_classes=["section-title"]) gr.Markdown( "Sign in, load the model, and adjust how TADA turns text plus a reference voice into synchronized speech generation. " "These controls influence how strongly the model follows linguistic content, acoustic structure, timing, and candidate ranking during synthesis.", elem_classes=["section-subtitle"], ) gr.LoginButton("Sign in with Hugging Face") hf_auth_status = gr.Markdown("**Hugging Face:** not signed in") with gr.Row(): model_dropdown = gr.Dropdown( choices=_MODEL_CHOICES, value=_DEFAULT_MODEL, label="Model", scale=4, ) load_btn = gr.Button("Load Model", scale=1) model_status = gr.Textbox( label="Model Status", value="Not loaded", interactive=False, ) gr.Markdown("### Generation Settings") gr.Markdown( "These settings control how TADA combines text and acoustics during speech generation." ) with gr.Tabs(): with gr.Tab("Text Controls"): num_extra_steps = gr.Slider( minimum=0, maximum=200, value=0, step=1, label="Text Tokens to Generate", info="Number of additional text tokens the model may generate.", ) text_only_logit_scale = gr.Slider( minimum=0.0, maximum=5.0, value=0.0, step=0.1, label="Text-Only Logit Scale", info="Blends text-only logits with audio-conditioned logits.", ) normalize_text_cb = gr.Checkbox( value=True, label="Normalize Text", info="Apply text normalization before generation.", ) with gr.Tab("Acoustic · Sampling"): noise_temperature = gr.Slider( minimum=0.4, maximum=1.2, value=0.9, step=0.1, label="Noise Temperature", info="Controls acoustic randomness.", ) num_flow_matching_steps = gr.Slider( minimum=5, maximum=50, value=20, step=5, label="Flow Matching Steps", info="Higher values can improve quality at the cost of speed.", ) speed_up_factor = gr.Slider( minimum=0.0, maximum=3.0, value=0.0, step=0.1, label="Speed Up Factor", info="0 keeps natural duration. Higher values speed up speech.", ) with gr.Tab("Acoustic · Guidance"): acoustic_cfg_scale = gr.Slider( minimum=1.0, maximum=3.0, value=1.6, step=0.1, label="Acoustic CFG Scale", info="Controls acoustic conditioning strength.", ) duration_cfg_scale = gr.Slider( minimum=1.0, maximum=3.0, value=1.0, step=0.1, label="Duration CFG Scale", info="Controls duration guidance strength.", ) negative_condition_source = gr.Dropdown( choices=["negative_step_output", "prompt", "zero"], value="negative_step_output", label="Negative Condition Source", info="Source used for negative conditioning.", ) with gr.Tab("Acoustic · Ranking"): num_acoustic_candidates = gr.Slider( minimum=1, maximum=16, value=1, step=1, label="Acoustic Candidates", info="Number of acoustic candidates to generate and rank.", ) scorer_dropdown = gr.Dropdown( choices=["likelihood", "spkr_verification", "duration_median"], value="likelihood", label="Scorer", info="Method used to rank acoustic candidates.", ) spkr_verification_weight = gr.Slider( minimum=0.0, maximum=5.0, value=1.0, step=0.1, label="Speaker Verification Weight", info="Used only for the speaker-verification scorer.", ) with gr.Column(scale=6, elem_classes=["split-right"]): gr.Markdown("## 2. Reference Voice", elem_classes=["section-title"]) gr.Markdown( "Choose how to provide the reference voice: use a preset sample, upload your own recording, or record one live. " "TADA encodes this audio to guide speaker and style during generation.", elem_classes=["section-subtitle"], ) with gr.Row(): language_dd = gr.Dropdown( choices=list(LANGUAGE_MAP.keys()), value="English", label="Language", info="Used to encode and align reference audio.", scale=4, ) voice_choices, voice_default, voice_path = _default_voice_for_language("English") preset_dropdown = gr.Dropdown( choices=voice_choices, value=voice_default, label="Reference Voice Preset", info="Choose a sample voice or switch to zero-shot mode.", scale=6, ) gr.Markdown( "
Upload or record a short voice sample, or use the selected preset.
" ) audio_input = gr.Audio( label="Reference Audio", type="filepath", sources=["upload", "microphone"], value=voice_path, elem_classes=["compact-audio"], ) process_prompt_btn = gr.Button( "Prepare Reference Voice", variant="secondary", interactive=False, ) gr.Markdown("### Reference Text-Acoustic Alignment") with gr.Column(elem_classes=["token-box-small"]): prompt_alignment = gr.HTML( value=_reference_info_panel() ) with gr.Column(elem_classes=["section-card"]): gr.Markdown("## 3. Generate Speech", elem_classes=["section-title"]) gr.Markdown( "After loading the model and preparing a reference voice, choose a preset text or write your own to generate new speech conditioned on that reference.", elem_classes=["section-subtitle"], ) transcript_choices, transcript_default, transcript_text = _default_text_for_language("English") with gr.Row(equal_height=False): with gr.Column(scale=4): transcript_dropdown = gr.Dropdown( choices=transcript_choices, value=transcript_default, label="Text Preset", info="Pick a preset or switch to custom text.", ) text_input = gr.Textbox( label="Text to Synthesize", placeholder="Type the text you want the model to say...", autoscroll=False, max_lines=12, value=transcript_text, ) generate_btn = gr.Button( "Generate Speech", variant="primary", size="lg", interactive=False, ) gr.Markdown("### Generated Audio") audio_output = gr.Audio(label=None) app_status = gr.Markdown( _status_md("Sign in, load a model, then prepare a reference voice.") ) with gr.Column(scale=6): gr.Markdown("### Text-Acoustic Generation Trace") with gr.Column(elem_classes=["token-box-large"]): generated_text_display = gr.HTML( value=_generation_info_panel() ) demo.load( fn=_hf_auth_status, inputs=None, outputs=[hf_auth_status], ) model_dropdown.change( fn=_on_model_selected, inputs=[model_dropdown], outputs=[language_dd], ) load_btn.click( fn=_before_load_model, inputs=[], outputs=[load_btn, model_status], ).then( fn=load_models, inputs=[model_dropdown, language_dd], outputs=[model_status], ).then( fn=_after_load_model, inputs=[], outputs=[load_btn, process_prompt_btn, app_status], ) language_dd.change( fn=_on_language_changed, inputs=[language_dd], outputs=[ preset_dropdown, audio_input, transcript_dropdown, text_input, prompt_state, generate_btn, app_status, prompt_alignment, ], ) preset_dropdown.change( fn=_on_preset_selected, inputs=[preset_dropdown, language_dd], outputs=[audio_input], ).then( fn=lambda: _invalidate_prompt_state( "Reference voice changed. Please prepare it again before generating.", notify=True, ), inputs=[], outputs=[prompt_state, generate_btn, app_status, prompt_alignment], ) audio_input.change( fn=lambda: _invalidate_prompt_state( "Reference audio changed. Please prepare it again before generating.", notify=False, ), inputs=[], outputs=[prompt_state, generate_btn, app_status, prompt_alignment], ) transcript_dropdown.change( fn=_on_transcript_selected, inputs=[transcript_dropdown, language_dd], outputs=[text_input], ) process_prompt_btn.click( fn=_before_prepare, inputs=[], outputs=[process_prompt_btn, app_status], ).then( fn=process_prompt, inputs=[audio_input, language_dd], outputs=[prompt_alignment, prompt_state], ).then( fn=_after_prepare_success, inputs=[], outputs=[process_prompt_btn, generate_btn, app_status], ) all_inputs = [ text_input, num_extra_steps, noise_temperature, acoustic_cfg_scale, duration_cfg_scale, num_flow_matching_steps, negative_condition_source, text_only_logit_scale, num_acoustic_candidates, scorer_dropdown, spkr_verification_weight, speed_up_factor, normalize_text_cb, prompt_state, ] generate_btn.click( fn=_before_generate, inputs=[], outputs=[app_status], ).then( fn=generate_speech, inputs=all_inputs, outputs=[audio_output, generated_text_display], ).then( fn=_after_generate, inputs=[], outputs=[app_status], ) return demo # --------------------------------------------------------------------------- # App setup # --------------------------------------------------------------------------- demo = build_ui() CSS = """ .gradio-container { max-width: 1280px !important; width: 100% !important; margin: auto !important; } .hero-header { gap: 0 !important; margin-bottom: 12px !important; } .hero-header h1 { margin: 0 0 4px 0 !important; line-height: 1.15 !important; } .section-card { border: 1px solid var(--border-color-primary); border-radius: 16px; padding: 18px; margin-bottom: 18px; background: var(--block-background-fill); } .split-left { padding-right: 16px; border-right: 1px solid var(--border-color-primary); } .split-right { padding-left: 16px; } .compact-audio { min-height: 0 !important; } .compact-audio audio { height: 40px !important; } .section-title { margin-bottom: 4px !important; } .section-subtitle { color: var(--body-text-color-subdued); margin-bottom: 14px !important; } .hero-subtitle { color: var(--body-text-color-subdued); margin-top: 0 !important; margin-bottom: 4px !important; font-size: 1.05rem; } .hf-badges { display: flex; gap: 4px; flex-wrap: wrap; align-items: center; justify-content: flex-start; margin: 0 !important; padding: 0 !important; } .hf-badges a { display: inline-flex; margin: 0 !important; padding: 0 !important; line-height: 0; } .hf-badges img { display: block; margin: 0 !important; } .helper-box { margin-top: 8px; color: var(--body-text-color-subdued); } .token-box-small { border: 1px solid var(--border-color-primary); border-radius: 12px; padding: 12px; background: var(--input-background-fill); overflow-x: auto; overflow-y: auto; max-width: 100%; box-sizing: border-box; min-height: 300px; max-height: 300px; } .token-box-large { border: 1px solid var(--border-color-primary); border-radius: 12px; padding: 12px; background: var(--input-background-fill); overflow-x: auto; overflow-y: auto; max-width: 100%; box-sizing: border-box; min-height: 620px; max-height: 620px; } .token-box-small *, .token-box-large * { max-width: 100%; box-sizing: border-box; } """ if __name__ == "__main__": demo.launch( server_name="0.0.0.0", server_port=7860, ssr_mode=False, allowed_paths=[_SAMPLES_DIR], css=CSS, )