#!/usr/bin/env python3 from __future__ import annotations import argparse from datetime import datetime from pathlib import Path import gradio as gr from huggingface_hub import hf_hub_download from irodori_tts.inference_runtime import ( RuntimeKey, SamplingRequest, clear_cached_runtime, default_runtime_device, get_cached_runtime, list_available_runtime_devices, list_available_runtime_precisions, save_wav, ) FIXED_SECONDS = 30.0 MAX_GRADIO_CANDIDATES = 32 GRADIO_AUDIO_COLS_PER_ROW = 8 def _default_checkpoint() -> str: candidates = sorted( [ *Path(".").glob("**/checkpoint_*.pt"), *Path(".").glob("**/checkpoint_*.safetensors"), ] ) if not candidates: return "Aratako/Irodori-TTS-500M" return str(candidates[-1]) def _default_model_device() -> str: return default_runtime_device() def _default_codec_device() -> str: return default_runtime_device() def _precision_choices_for_device(device: str) -> list[str]: return list_available_runtime_precisions(device) def _on_model_device_change(device: str) -> gr.Dropdown: choices = _precision_choices_for_device(device) return gr.Dropdown(choices=choices, value=choices[0]) def _on_codec_device_change(device: str) -> gr.Dropdown: choices = _precision_choices_for_device(device) return gr.Dropdown(choices=choices, value=choices[0]) def _parse_optional_float(raw: str | None, label: str) -> float | None: if raw is None: return None text = str(raw).strip() if text == "" or text.lower() == "none": return None try: return float(text) except ValueError as exc: raise ValueError(f"{label} must be a float or blank.") from exc def _parse_optional_int(raw: str | None, label: str) -> int | None: if raw is None: return None text = str(raw).strip() if text == "" or text.lower() == "none": return None try: return int(text) except ValueError as exc: raise ValueError(f"{label} must be an int or blank.") from exc def _format_timings(stage_timings: list[tuple[str, float]], total_to_decode: float) -> str: lines = [ "[timing] ---- request ----", *[f"[timing] {name}: {sec * 1000.0:.1f} ms" for name, sec in stage_timings], f"[timing] total_to_decode: {total_to_decode:.3f} s", ] return "\n".join(lines) def _resolve_ref_wav(uploaded_audio: str | None) -> str | None: if uploaded_audio is not None and str(uploaded_audio).strip() != "": return str(uploaded_audio) return None def _resolve_checkpoint_path(raw_checkpoint: str) -> str: checkpoint = str(raw_checkpoint).strip() if checkpoint == "": raise ValueError("checkpoint is required.") suffix = Path(checkpoint).suffix.lower() if suffix in {".pt", ".safetensors"}: return checkpoint resolved = hf_hub_download(repo_id=checkpoint, filename="model.safetensors") print(f"[gradio] checkpoint: hf://{checkpoint} -> {resolved}", flush=True) return str(resolved) def _build_runtime_key( checkpoint: str, model_device: str, model_precision: str, codec_device: str, codec_precision: str, enable_watermark: bool, ) -> RuntimeKey: checkpoint_path = _resolve_checkpoint_path(checkpoint) return RuntimeKey( checkpoint=checkpoint_path, model_device=str(model_device), codec_repo="facebook/dacvae-watermarked", model_precision=str(model_precision), codec_device=str(codec_device), codec_precision=str(codec_precision), enable_watermark=bool(enable_watermark), compile_model=False, compile_dynamic=False, ) def _load_model( checkpoint: str, model_device: str, model_precision: str, codec_device: str, codec_precision: str, enable_watermark: bool, ) -> str: runtime_key = _build_runtime_key( checkpoint=checkpoint, model_device=model_device, model_precision=model_precision, codec_device=codec_device, codec_precision=codec_precision, enable_watermark=enable_watermark, ) _, reloaded = get_cached_runtime(runtime_key) if reloaded: status = "loaded model into memory" else: status = "model already loaded; reused existing runtime" return ( f"{status}\n" f"checkpoint: {runtime_key.checkpoint}\n" f"model_device: {runtime_key.model_device}\n" f"model_precision: {runtime_key.model_precision}\n" f"codec_device: {runtime_key.codec_device}\n" f"codec_precision: {runtime_key.codec_precision}" ) def _run_generation( checkpoint: str, model_device: str, model_precision: str, codec_device: str, codec_precision: str, enable_watermark: bool, text: str, uploaded_audio: str | None, num_steps: int, num_candidates: int, seed_raw: str, cfg_guidance_mode: str, cfg_scale_text: float, cfg_scale_speaker: float, cfg_scale_raw: str, cfg_min_t: float, cfg_max_t: float, context_kv_cache: bool, truncation_factor_raw: str, rescale_k_raw: str, rescale_sigma_raw: str, normalize_ref_audio: bool, speaker_kv_scale_raw: str, speaker_kv_min_t_raw: str, speaker_kv_max_layers_raw: str, ) -> tuple[object, ...]: def stdout_log(msg: str) -> None: print(msg, flush=True) runtime_key = _build_runtime_key( checkpoint=checkpoint, model_device=model_device, model_precision=model_precision, codec_device=codec_device, codec_precision=codec_precision, enable_watermark=enable_watermark, ) if str(text).strip() == "": raise ValueError("text is required.") requested_candidates = int(num_candidates) if requested_candidates <= 0: raise ValueError("num_candidates must be >= 1.") if requested_candidates > MAX_GRADIO_CANDIDATES: raise ValueError(f"num_candidates must be <= {MAX_GRADIO_CANDIDATES}.") cfg_scale = _parse_optional_float(cfg_scale_raw, "cfg_scale") truncation_factor = _parse_optional_float(truncation_factor_raw, "truncation_factor") rescale_k = _parse_optional_float(rescale_k_raw, "rescale_k") rescale_sigma = _parse_optional_float(rescale_sigma_raw, "rescale_sigma") speaker_kv_scale = _parse_optional_float(speaker_kv_scale_raw, "speaker_kv_scale") speaker_kv_min_t = _parse_optional_float(speaker_kv_min_t_raw, "speaker_kv_min_t") speaker_kv_max_layers = _parse_optional_int(speaker_kv_max_layers_raw, "speaker_kv_max_layers") seed = _parse_optional_int(seed_raw, "seed") ref_wav = _resolve_ref_wav(uploaded_audio=uploaded_audio) no_ref = ref_wav is None ref_normalize_db = -16.0 if bool(normalize_ref_audio) else None ref_ensure_max = bool(normalize_ref_audio) runtime, reloaded = get_cached_runtime(runtime_key) stdout_log(f"[gradio] runtime: {'reloaded' if reloaded else 'reused'}") stdout_log( ( "[gradio] request: model_device={} model_precision={} codec_device={} codec_precision={} " "watermark={} mode={} seconds={} steps={} seed={} no_ref={} candidates={}" ).format( model_device, model_precision, codec_device, codec_precision, enable_watermark, cfg_guidance_mode, FIXED_SECONDS, num_steps, "random" if seed is None else seed, no_ref, requested_candidates, ) ) result = runtime.synthesize( SamplingRequest( text=str(text), ref_wav=ref_wav, ref_latent=None, no_ref=bool(no_ref), ref_normalize_db=ref_normalize_db, ref_ensure_max=bool(ref_ensure_max), num_candidates=requested_candidates, decode_mode="sequential", seconds=FIXED_SECONDS, max_ref_seconds=30.0, max_text_len=None, num_steps=int(num_steps), seed=None if seed is None else int(seed), cfg_guidance_mode=str(cfg_guidance_mode), cfg_scale_text=float(cfg_scale_text), cfg_scale_speaker=float(cfg_scale_speaker), cfg_scale=cfg_scale, cfg_min_t=float(cfg_min_t), cfg_max_t=float(cfg_max_t), truncation_factor=truncation_factor, rescale_k=rescale_k, rescale_sigma=rescale_sigma, context_kv_cache=bool(context_kv_cache), speaker_kv_scale=speaker_kv_scale, speaker_kv_min_t=speaker_kv_min_t, speaker_kv_max_layers=speaker_kv_max_layers, trim_tail=True, ), log_fn=stdout_log, ) out_dir = Path("gradio_outputs") out_dir.mkdir(parents=True, exist_ok=True) stamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") out_paths: list[str] = [] for i, audio in enumerate(result.audios, start=1): out_path = save_wav( out_dir / f"sample_{stamp}_{i:03d}.wav", audio.float(), result.sample_rate, ) out_paths.append(str(out_path)) runtime_msg = "runtime: reloaded" if reloaded else "runtime: reused" detail_lines = [ runtime_msg, f"seed_used: {result.used_seed}", f"candidates: {len(result.audios)}", *[f"saved[{i}]: {path}" for i, path in enumerate(out_paths, start=1)], *result.messages, ] detail_text = "\n".join(detail_lines) timing_text = _format_timings(result.stage_timings, result.total_to_decode) stdout_log(f"[gradio] saved {len(out_paths)} candidates") audio_updates: list[object] = [] for i in range(MAX_GRADIO_CANDIDATES): if i < len(out_paths): audio_updates.append(gr.update(value=out_paths[i], visible=True)) else: audio_updates.append(gr.update(value=None, visible=False)) return tuple([*audio_updates, detail_text, timing_text]) def _clear_runtime_cache() -> str: clear_cached_runtime() return "cleared loaded model from memory" def build_ui() -> gr.Blocks: default_checkpoint = _default_checkpoint() default_model_device = _default_model_device() default_codec_device = _default_codec_device() device_choices = list_available_runtime_devices() model_precision_choices = _precision_choices_for_device(default_model_device) codec_precision_choices = _precision_choices_for_device(default_codec_device) with gr.Blocks(title="Irodori-TTS Gradio") as demo: gr.Markdown("# Irodori-TTS Inference (Cached Runtime)") gr.Markdown( "When settings are unchanged, runtime is reused and only sampling/decoding runs." ) with gr.Row(): checkpoint = gr.Textbox( label="Checkpoint (.pt/.safetensors or HF repo id)", value=default_checkpoint, scale=4, ) model_device = gr.Dropdown( label="Model Device", choices=device_choices, value=default_model_device, scale=1, ) model_precision = gr.Dropdown( label="Model Precision", choices=model_precision_choices, value=model_precision_choices[0], scale=1, ) codec_device = gr.Dropdown( label="Codec Device", choices=device_choices, value=default_codec_device, scale=1, ) codec_precision = gr.Dropdown( label="Codec Precision", choices=codec_precision_choices, value=codec_precision_choices[0], scale=1, ) enable_watermark = gr.Checkbox(label="Enable Watermark", value=False, scale=1) with gr.Row(): load_model_btn = gr.Button("Load Model") clear_cache_btn = gr.Button("Unload Model") clear_cache_msg = gr.Textbox(label="Model Status", interactive=False) text = gr.Textbox(label="Text", lines=4) with gr.Row(): uploaded_audio = gr.Audio( label="Reference Audio Upload (optional, blank = no-reference mode)", type="filepath", scale=4, ) with gr.Column(scale=1): normalize_ref_audio = gr.Checkbox( label="Normalize Ref Audio (-16 LUFS + ensure max)", value=True, ) with gr.Accordion("Sampling", open=True): with gr.Row(): num_steps = gr.Slider(label="Num Steps", minimum=1, maximum=120, value=40, step=1) num_candidates = gr.Slider( label="Num Candidates", minimum=1, maximum=MAX_GRADIO_CANDIDATES, value=1, step=1, ) seed_raw = gr.Textbox(label="Seed (blank=random)", value="") with gr.Row(): cfg_guidance_mode = gr.Dropdown( label="CFG Guidance Mode", choices=["independent", "joint", "alternating"], value="independent", ) cfg_scale_text = gr.Slider( label="CFG Scale Text", minimum=0.0, maximum=10.0, value=3.0, step=0.1, ) cfg_scale_speaker = gr.Slider( label="CFG Scale Speaker", minimum=0.0, maximum=10.0, value=5.0, step=0.1, ) with gr.Accordion("Advanced (Optional)", open=False): cfg_scale_raw = gr.Textbox(label="CFG Scale Override (optional)", value="") with gr.Row(): cfg_min_t = gr.Number(label="CFG Min t", value=0.5) cfg_max_t = gr.Number(label="CFG Max t", value=1.0) context_kv_cache = gr.Checkbox(label="Context KV Cache", value=True) with gr.Row(): truncation_factor_raw = gr.Textbox(label="Truncation Factor (optional)", value="") rescale_k_raw = gr.Textbox(label="Rescale k (optional)", value="") rescale_sigma_raw = gr.Textbox(label="Rescale sigma (optional)", value="") with gr.Row(): speaker_kv_scale_raw = gr.Textbox(label="Speaker KV Scale (optional)", value="") speaker_kv_min_t_raw = gr.Textbox(label="Speaker KV Min t (optional)", value="0.9") speaker_kv_max_layers_raw = gr.Textbox( label="Speaker KV Max Layers (optional)", value="" ) generate_btn = gr.Button("Generate", variant="primary") out_audios: list[gr.Audio] = [] num_rows = (MAX_GRADIO_CANDIDATES + GRADIO_AUDIO_COLS_PER_ROW - 1) // GRADIO_AUDIO_COLS_PER_ROW with gr.Column(): for row_idx in range(num_rows): with gr.Row(): for col_idx in range(GRADIO_AUDIO_COLS_PER_ROW): i = row_idx * GRADIO_AUDIO_COLS_PER_ROW + col_idx if i >= MAX_GRADIO_CANDIDATES: break out_audios.append( gr.Audio( label=f"Generated Audio {i + 1}", type="filepath", interactive=False, visible=(i == 0), min_width=160, ) ) out_log = gr.Textbox(label="Run Log", lines=8) out_timing = gr.Textbox(label="Timing", lines=8) generate_btn.click( _run_generation, inputs=[ checkpoint, model_device, model_precision, codec_device, codec_precision, enable_watermark, text, uploaded_audio, num_steps, num_candidates, seed_raw, cfg_guidance_mode, cfg_scale_text, cfg_scale_speaker, cfg_scale_raw, cfg_min_t, cfg_max_t, context_kv_cache, truncation_factor_raw, rescale_k_raw, rescale_sigma_raw, normalize_ref_audio, speaker_kv_scale_raw, speaker_kv_min_t_raw, speaker_kv_max_layers_raw, ], outputs=[*out_audios, out_log, out_timing], ) model_device.change( _on_model_device_change, inputs=[model_device], outputs=[model_precision] ) codec_device.change( _on_codec_device_change, inputs=[codec_device], outputs=[codec_precision] ) load_model_btn.click( _load_model, inputs=[ checkpoint, model_device, model_precision, codec_device, codec_precision, enable_watermark, ], outputs=[clear_cache_msg], ) clear_cache_btn.click(_clear_runtime_cache, outputs=[clear_cache_msg]) return demo def main() -> None: parser = argparse.ArgumentParser(description="Gradio app for Irodori-TTS with cached runtime.") parser.add_argument("--server-name", default="127.0.0.1") parser.add_argument("--server-port", type=int, default=7860) parser.add_argument("--share", action="store_true") parser.add_argument("--debug", action="store_true") args = parser.parse_args() demo = build_ui() demo.queue(default_concurrency_limit=1) demo.launch( server_name=args.server_name, server_port=args.server_port, share=bool(args.share), debug=bool(args.debug), ) if __name__ == "__main__": main()