""" Gradio app for TADA inference. Usage: pip install hume-tada python app.py # or with hot reload + share link: GRADIO_SHARE=1 gradio app.py """ import dataclasses import html import json import logging import os import shutil import tempfile import time import torch import torchaudio import gradio as gr 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 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Preset samples & transcripts # --------------------------------------------------------------------------- _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, "Arabic": "ar", "Chinese": "ch", "German": "de", "Spanish": "es", "French": "fr", "Italian": "it", "Japanese": "ja", "Polish": "pl", "Portuguese": "pt", } _CODE_TO_LANG_DIR = {None: "en", "ar": "ar", "ch": "ch", "de": "de", "es": "es", "fr": "fr", "it": "it", "ja": "ja", "pl": "pl", "pt": "pt"} def _discover_preset_samples(lang_code: str | None = None) -> dict[str, str]: """Return {display_name: absolute_path} for audio files in the language's samples/ subdir.""" 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 preset transcripts from synth_transcripts.json in the language's sample dir.""" 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 {} # Initialize with English samples _PRESET_SAMPLES = _discover_preset_samples(None) _PRESET_TRANSCRIPTS = _load_preset_transcripts(None) logger.info("Discovered %d preset audio samples, %d transcripts", len(_PRESET_SAMPLES), len(_PRESET_TRANSCRIPTS)) # --------------------------------------------------------------------------- # Global model state # --------------------------------------------------------------------------- _MODEL_CHOICES = ["HumeAI/tada-1b", "HumeAI/tada-3b-ml"] _DEFAULT_MODEL = "HumeAI/tada-3b-ml" _encoder_cache: dict[str | None, Encoder] = {} _model: TadaForCausalLM | None = None _current_model_name: str = "" _current_language: str | None = None _device = "cuda" def _move_encoder_output(output: EncoderOutput, device: str) -> EncoderOutput: """Move all tensor fields of an EncoderOutput to the given 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_encoder(language_code: str | None = None) -> Encoder: """Get or create an Encoder for the given language, with caching.""" if language_code not in _encoder_cache: _encoder_cache[language_code] = Encoder.from_pretrained( "HumeAI/tada-codec", language=language_code ).to(_device) return _encoder_cache[language_code] def _get_device_info() -> str: 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 (ZeroGPU provides GPU during inference)" def load_models(model_name: str = _DEFAULT_MODEL) -> str: """Load encoder and TADA model. Returns a status string.""" global _model, _current_model_name if _model is not None and _current_model_name == model_name: return f"Loaded: {model_name} on {_get_device_info()}" if _model is not None: del _model _model = None get_encoder(_current_language) logger.info("Loading %s ...", model_name) _model = TadaForCausalLM.from_pretrained(model_name) _current_model_name = model_name status = f"Loaded: {model_name} on {_get_device_info()}" logger.info(status) return status # --------------------------------------------------------------------------- # Core inference helpers # --------------------------------------------------------------------------- def _encode_prompt(audio_path: str | None, language_code: str | None = None) -> EncoderOutput: """Encode an audio file into an EncoderOutput prompt (or return an empty one).""" if audio_path is None or audio_path == "": return EncoderOutput.empty(_device) encoder = get_encoder(language_code) audio, sample_rate = torchaudio.load(audio_path) audio = audio.mean(dim=0, keepdim=True) # mono audio = audio / audio.abs().max().clamp(min=1e-8) * 0.95 audio = audio.to(_device) prompt = encoder(audio, sample_rate=sample_rate) return prompt def _decode_tokens_individually(tokenizer, token_ids: list[int]) -> list[str]: """Decode a list of token IDs into per-token strings, handling multi-byte characters.""" 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) -> str: """Build an HTML string: dots in grey, tokens as bold coloured spans.""" if prompt.text_tokens is None or prompt.token_positions is None: return "" encoder = get_encoder(language_code) 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( f'{escaped}' ) prev_pos = pos + 1 body = "".join(parts) return ( f'
' f'
{header}
' f"{body}
" ) @gpu_decorator @torch.inference_mode() def process_prompt(audio_path: str | None, language: str = "English") -> tuple[str, EncoderOutput | None]: """Encode the voice prompt and return (alignment_html, prompt_on_cpu).""" global _current_language language_code = LANGUAGE_MAP.get(language) _current_language = language_code _encoder = get_encoder(language_code) _encoder.to(_device) if audio_path is None or audio_path == "": return "No audio provided (zero-shot mode).", None try: prompt = _encode_prompt(audio_path, language_code) alignment_html = _format_token_alignment(prompt, language_code) # Move to CPU for gr.State serialization (ZeroGPU compatibility) prompt_cpu = _move_encoder_output(prompt, "cpu") return alignment_html, prompt_cpu except Exception as e: logger.exception("Prompt processing failed") raise gr.Error(f"Prompt processing failed: {e}") def _decode_byte_tokens(raw_tokens: list[str]) -> list[str]: """Decode GPT-2 byte-level token strings into proper Unicode per-token labels.""" 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 string from step_logs: dots for n_frames_before, tokens highlighted.""" 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 | {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( f'{escaped}' ) body = "".join(parts) return ( f'
' f'
{header}
' f"{body}
" ) @gpu_decorator(duration=120) @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, ) -> tuple[str | None, str]: """Run TADA generation using the provided prompt and return (wav_path, alignment_html).""" if _model is None: raise gr.Error("Models are not loaded. Click 'Load Model' first.") if cached_prompt is None: raise gr.Error("Please upload audio and click 'Process Prompt' first.") _model.to(_device) _model.decoder.to(_device) try: prompt = _move_encoder_output(cached_prompt, _device) logger.info("Generating speech for text: %s", text) # speed_up_factor: 0 means disabled (None) suf = float(speed_up_factor) if speed_up_factor > 0 else None t0 = time.time() 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=suf, 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 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 # Filter to only generated steps (exclude prompted/prefilled steps) all_logs = output.step_logs or [] generated_logs = [e for e in all_logs if e.get("acoustic_feat_src") not in ("prompted", None)] generated_html = _format_step_logs(generated_logs, audio_duration, wall_time) return tmp_path, generated_html except gr.Error: raise except Exception as e: logger.exception("Generation failed") raise gr.Error(f"Generation failed: {e}") # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- def build_ui() -> gr.Blocks: with gr.Blocks( title="TADA Inference", css=( ".gradio-container { max-width: 1400px !important; width: 100% !important; margin: auto !important; } " ".compact-audio { min-height: 0 !important; } " ".compact-audio audio { height: 36px !important; } " ), ) as demo: gr.Markdown("# TADA - Text-Acoustic Dual Alignment LLM") prompt_state = gr.State(value=None) with gr.Row(equal_height=False): with gr.Column(scale=1): with gr.Row(): model_dropdown = gr.Dropdown( choices=_MODEL_CHOICES, value=_current_model_name or _DEFAULT_MODEL, label="Model", scale=3, ) load_btn = gr.Button("Load Model", scale=1) load_status = gr.Textbox(label="Model Status", interactive=False, show_label=False) language_dd = gr.Dropdown( choices=list(LANGUAGE_MAP.keys()), value="English", label="Language", info="Selects the aligner for prompt encoding", ) with gr.Accordion("Text Settings", open=False): num_extra_steps = gr.Slider( minimum=0, maximum=200, value=0, step=1, label="Text Tokens to 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="0 = disabled. 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.Accordion("Acoustic Settings", open=False): acoustic_cfg_scale = gr.Slider( minimum=1.0, maximum=3.0, value=1.6, step=0.1, label="Acoustic CFG Scale", ) duration_cfg_scale = gr.Slider( minimum=1.0, maximum=3.0, value=1.0, step=0.1, label="Duration CFG Scale", ) negative_condition_source = gr.Dropdown( choices=["negative_step_output", "prompt", "zero"], value="negative_step_output", label="Negative Condition Source", ) noise_temperature = gr.Slider( minimum=0.4, maximum=1.2, value=0.9, step=0.1, label="Noise Temperature", ) num_flow_matching_steps = gr.Slider( minimum=5, maximum=50, value=20, step=5, label="Flow Matching Steps", ) speed_up_factor = gr.Slider( minimum=0.0, maximum=3.0, value=0.0, step=0.1, label="Speed Up Factor", info="0 = disabled (natural duration). >0 scales speech speed.", ) num_acoustic_candidates = gr.Slider( minimum=1, maximum=16, value=1, step=1, label="Acoustic Candidates", info="Number of candidates to generate and rank.", ) scorer_dropdown = gr.Dropdown( choices=["likelihood", "spkr_verification", "duration_median"], value="likelihood", label="Scorer", info="How 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="Weight for spkr_verification scorer.", ) with gr.Column(scale=2): preset_choices = ["None (zero-shot)"] + list(_PRESET_SAMPLES.keys()) _default_voice = "fb_ears_emo_amazement_freeform.wav" preset_dropdown = gr.Dropdown( choices=preset_choices, value=_default_voice if _default_voice in _PRESET_SAMPLES else "None (zero-shot)", label="Voice Prompt", info="Pick a preset or upload / record your own", ) _default_voice_path = _PRESET_SAMPLES.get(_default_voice) audio_input = gr.Audio( label="Prompt Preview", type="filepath", sources=["upload", "microphone"], value=_default_voice_path, elem_classes=["compact-audio"], ) def _on_preset_selected(choice: str) -> str | None: if choice == "None (zero-shot)": return None path = _PRESET_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 preset_dropdown.change( fn=_on_preset_selected, inputs=[preset_dropdown], outputs=[audio_input], ) def _on_language_changed(language: str): """Update preset samples and transcripts when language changes.""" lang_code = LANGUAGE_MAP.get(language) samples = _discover_preset_samples(lang_code) new_preset_choices = ["None (zero-shot)"] + list(samples.keys()) global _PRESET_SAMPLES, _PRESET_TRANSCRIPTS _PRESET_SAMPLES = samples _PRESET_TRANSCRIPTS = _load_preset_transcripts(lang_code) new_transcript_choices = ["(custom)"] + list(_PRESET_TRANSCRIPTS.keys()) first_sample = new_preset_choices[1] if len(new_preset_choices) > 1 else "None (zero-shot)" return ( gr.update(choices=new_preset_choices, value=first_sample), gr.update(choices=new_transcript_choices, value="(custom)"), ) process_prompt_btn = gr.Button("Process Prompt", variant="secondary", size="sm") with gr.Accordion("Token Alignment", open=True): prompt_alignment = gr.HTML(value="Upload audio and click Process Prompt before generating.") with gr.Column(scale=2): _default_transcript = "emo_interest_sentences" transcript_choices = ["(custom)"] + list(_PRESET_TRANSCRIPTS.keys()) transcript_dropdown = gr.Dropdown( choices=transcript_choices, value=_default_transcript if _default_transcript in _PRESET_TRANSCRIPTS else "(custom)", label="Transcript", info="Pick a preset or type your own below", ) text_input = gr.Textbox( label="Text to Speak", placeholder="Type what you want the model to say ...", autoscroll=False, max_lines=20, value=_PRESET_TRANSCRIPTS.get(_default_transcript, ""), ) def _on_transcript_selected(choice: str) -> str: if choice == "(custom)": return "" return _PRESET_TRANSCRIPTS.get(choice, "") transcript_dropdown.change( fn=_on_transcript_selected, inputs=[transcript_dropdown], outputs=[text_input], ) generate_btn = gr.Button("Generate", variant="primary", size="lg") # --- Wire language change to update presets + re-process prompt --- language_dd.change( fn=_on_language_changed, inputs=[language_dd], outputs=[preset_dropdown, transcript_dropdown], ) # --- Shared chain: show "Processing..." -> encode -> restore button --- def _wire_process_prompt(event): """Chain process_prompt onto any event.""" event.then( fn=lambda: (gr.update(value="Processing...", interactive=False), ""), inputs=[], outputs=[process_prompt_btn, prompt_alignment], ).then( fn=process_prompt, inputs=[audio_input, language_dd], outputs=[prompt_alignment, prompt_state], ).then( fn=lambda: gr.update(value="Process Prompt", interactive=True), inputs=[], outputs=[process_prompt_btn], ) # Manual click _wire_process_prompt(process_prompt_btn.click(fn=lambda: None, inputs=[], outputs=[])) # Load model (no auto-process; user must click Process Prompt) load_btn.click( fn=lambda: (gr.update(interactive=False), "Loading model..."), inputs=[], outputs=[load_btn, load_status], ).then( fn=load_models, inputs=[model_dropdown], outputs=[load_status], ).then( fn=lambda: gr.update(interactive=True), inputs=[], outputs=[load_btn], ) # --- Output --- audio_output = gr.Audio(label="Generated Audio") with gr.Accordion("Generated Alignment", open=False): generated_text_display = gr.HTML(value="Generate speech to see the alignment") # Wire up generate button 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=generate_speech, inputs=all_inputs, outputs=[audio_output, generated_text_display], ) return demo # --------------------------------------------------------------------------- # Entry-point # --------------------------------------------------------------------------- _share = os.environ.get("GRADIO_SHARE", "").lower() in ("1", "true", "yes") _port = int(os.environ.get("GRADIO_PORT", "7860")) # Auto-load models on startup load_models() # `demo` at module scope so the `gradio` CLI / HF Spaces can discover it. demo = build_ui() if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="TADA Inference Gradio App") parser.add_argument("--share", action="store_true", default=_share, help="Create a public Gradio share link") parser.add_argument("--port", type=int, default=_port, help="Server port (default: 7860)") args = parser.parse_args() demo.launch(server_name="0.0.0.0", server_port=args.port, share=args.share, allowed_paths=[_SAMPLES_DIR]) else: demo.launch(server_name="0.0.0.0", server_port=_port, share=_share, allowed_paths=[_SAMPLES_DIR])