import io import os import spaces # MUST come before torch / any CUDA-touching import import gradio as gr import numpy as np import torch from irodori_tts.gradio_emoji_palette import EMOJI_PALETTE_CSS, build_emoji_palette from irodori_tts.inference_runtime import ( InferenceRuntime, RuntimeKey, SamplingRequest, download_hf_checkpoint, ) # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- MODEL_REPO = os.environ.get("MODEL_REPO", "Aratako/Irodori-TTS-v4.1-Small") CODEC_REPO = "Aratako/Semantic-DACVAE-Japanese-32dim" MAX_GRADIO_CANDIDATES = int(os.environ.get("MAX_GRADIO_CANDIDATES", "8")) GRADIO_AUDIO_COLS_PER_ROW = 4 CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ # Global state _runtime: InferenceRuntime | None = None # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- 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 _on_t_schedule_mode_change(mode: str) -> object: return gr.update(interactive=str(mode).strip().lower() == "sway") def _coerce_gradio_file_path(value: object) -> str | None: if value is None: return None if isinstance(value, str): text = value.strip() return text or None if isinstance(value, dict): for key in ("path", "name"): candidate = value.get(key) if candidate is not None and str(candidate).strip(): return str(candidate) return None candidate = getattr(value, "name", None) if candidate is not None and str(candidate).strip(): return str(candidate) text = str(value).strip() return text or None def _resolve_ref_wavs(uploaded_audios: object) -> list[str]: if uploaded_audios is None: return [] values = uploaded_audios if isinstance(uploaded_audios, (list, tuple)) else [uploaded_audios] paths = [_coerce_gradio_file_path(value) for value in values] return [path for path in paths if path is not None] # --------------------------------------------------------------------------- # Model Loading (module scope, eager .to("cuda") for ZeroGPU) # --------------------------------------------------------------------------- def _load_models(): global _runtime if _runtime is not None: return print(f"[Info] Downloading checkpoint and tokenizer from {MODEL_REPO}...", flush=True) checkpoint_path = download_hf_checkpoint(MODEL_REPO) print(f"[Info] Checkpoint path: {checkpoint_path}", flush=True) device = "cuda" precision = "bf16" key = RuntimeKey( checkpoint=checkpoint_path, model_device=device, codec_repo=CODEC_REPO, model_precision=precision, codec_device=device, codec_precision=precision, ) print("[Info] Building runtime (this may take a while)...", flush=True) _runtime = InferenceRuntime.from_key(key) print("[Info] All models loaded successfully.", flush=True) # Load models at startup (ZeroGPU intercepts .to("cuda")) _load_models() _quantized = False # --------------------------------------------------------------------------- # GPU-decorated Inference # --------------------------------------------------------------------------- @spaces.GPU(duration=120) def run_inference_gpu( text: str, caption: str = "", uploaded_audios: object = None, num_steps: int = 40, num_candidates: int = 1, seed_raw: str = "", seconds_raw: str = "", duration_scale: float = 1.0, t_schedule_mode: str = "linear", sway_coeff: float = -1.0, cfg_guidance_mode: str = "independent", cfg_scale_text: float = 3.0, cfg_scale_caption: float = 4.0, cfg_scale_speaker: float = 5.0, cfg_scale_raw: str = "", cfg_min_t: float = 0.5, cfg_max_t: float = 1.0, context_kv_cache: bool = True, max_text_len_raw: str = "", max_caption_len_raw: str = "", truncation_factor_raw: str = "", rescale_k_raw: str = "", rescale_sigma_raw: str = "", speaker_kv_scale_raw: str = "", ) -> tuple[list[tuple[int, np.ndarray]], str]: """Generate Japanese speech from text using Irodori-TTS-v4.1-Small (INT8 quantized). The base bf16 model is loaded at startup; INT8 weight-only quantization is applied on the first GPU call to reduce VRAM usage while preserving quality. Args: text: Japanese text to synthesize. caption: Optional style/emotion prompt. uploaded_audios: Optional reference audio files for voice cloning. num_steps: Number of diffusion sampling steps. num_candidates: Number of audio candidates to generate. seed_raw: Random seed (blank = random). seconds_raw: Manual output duration in seconds (blank = auto). duration_scale: Scale factor for predicted duration. t_schedule_mode: Timestep schedule mode ("linear" or "sway"). sway_coeff: Sway sampling coefficient. cfg_guidance_mode: CFG guidance mode. cfg_scale_text: CFG scale for text conditioning. cfg_scale_caption: CFG scale for caption conditioning. cfg_scale_speaker: CFG scale for speaker conditioning. cfg_scale_raw: Optional override for all CFG scales. cfg_min_t: Minimum timestep for CFG. cfg_max_t: Maximum timestep for CFG. context_kv_cache: Whether to use context KV cache. max_text_len_raw: Optional max text token length. max_caption_len_raw: Optional max caption token length. truncation_factor_raw: Optional noise truncation factor. rescale_k_raw: Optional temporal rescale k. rescale_sigma_raw: Optional temporal rescale sigma. speaker_kv_scale_raw: Optional speaker KV scale factor. Returns: A tuple of (list of (sample_rate, waveform) pairs, log text). """ if _runtime is None: _load_models() global _quantized if not _quantized: from irodori_tts.quantization import quantize_model quantized_count = quantize_model( _runtime.model, quantization_type="int8_weight_only", profile="core" ) print(f"[Info] Quantized {quantized_count} layers to INT8 weight-only.", flush=True) _quantized = True log_buffer = io.StringIO() def stdout_log(msg: str) -> None: print(msg, flush=True) log_buffer.write(msg + "\n") text_value = "" if text is None else str(text).strip() if not text_value: raise gr.Error("Please enter text to synthesize.") caption_value = "" if caption is None else str(caption).strip() cfg_scale = _parse_optional_float(cfg_scale_raw, "cfg_scale") max_text_len = _parse_optional_int(max_text_len_raw, "max_text_len") max_caption_len = _parse_optional_int(max_caption_len_raw, "max_caption_len") 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") seed = _parse_optional_int(seed_raw, "seed") manual_seconds = _parse_optional_float(seconds_raw, "seconds") if num_candidates is None or (isinstance(num_candidates, str) and num_candidates.strip() == ""): num_candidates = 1 requested_candidates = int(num_candidates) if requested_candidates <= 0: raise gr.Error("num_candidates must be >= 1.") if requested_candidates > MAX_GRADIO_CANDIDATES: raise gr.Error(f"num_candidates must be <= {MAX_GRADIO_CANDIDATES}.") ref_wavs = _resolve_ref_wavs(uploaded_audios) no_ref = not ref_wavs if _runtime is not None and not _runtime.model_cfg.use_speaker_condition_resolved: ref_wavs = [] no_ref = True stdout_log( ( "[Info] request: mode={} seconds={} duration_scale={} " "schedule={} sway_coeff={} steps={} seed={} caption={} no_ref={} candidates={}" ).format( cfg_guidance_mode, "auto" if manual_seconds is None else manual_seconds, float(duration_scale), t_schedule_mode, float(sway_coeff), int(num_steps), "random" if seed is None else seed, "on" if caption_value else "off", no_ref, requested_candidates, ) ) result = _runtime.synthesize( SamplingRequest( text=text_value, caption=caption_value or None, ref_wav=None, ref_wavs=ref_wavs or None, ref_latent=None, no_ref=bool(no_ref), ref_normalize_db=-16.0, ref_ensure_max=True, num_candidates=requested_candidates, decode_mode="sequential", seconds=manual_seconds, duration_scale=float(duration_scale), max_ref_seconds=None, max_text_len=max_text_len, max_caption_len=max_caption_len, 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_caption=float(cfg_scale_caption), cfg_scale_speaker=0.0 if no_ref else 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=None if no_ref else speaker_kv_scale, speaker_kv_min_t=None, speaker_kv_max_layers=None, t_schedule_mode=str(t_schedule_mode), sway_coeff=float(sway_coeff), trim_tail=True, ), log_fn=stdout_log, ) sample_rate = result.sample_rate audio_results: list[tuple[int, np.ndarray]] = [] for audio in result.audios: waveform = audio.squeeze(0).float().numpy() audio_results.append((sample_rate, waveform)) stdout_log(f"[Info] seed_used: {result.used_seed}") stdout_log(f"[Info] candidates: {len(result.audios)}") return audio_results, log_buffer.getvalue() # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- def build_demo(): MODEL_LINK = "https://huggingface.co/Aratako/Irodori-TTS-v4.1-Small-Quantized" BASE_MODEL_LINK = "https://huggingface.co/Aratako/Irodori-TTS-v4.1-Small" GITHUB_REPO = "https://github.com/Aratako/Irodori-TTS" title = "# Irodori-TTS-v4.1-Small (INT8 Quantized) Demo" description = f"""\ [Quantized Model]({MODEL_LINK}) | [Base Model]({BASE_MODEL_LINK}) | [GitHub]({GITHUB_REPO}) Flow-matching based Japanese TTS model (approximately 766M parameters), \ loaded from the **INT8 weight-only quantized** variant for reduced memory footprint. \ Generates speech from text, optional reference audio, and optional style caption. - **Reference audio**: Optional. One or more clips can be concatenated in the displayed order, up to the model's 120-second limit. - **Caption**: Optional style prompt for emotion, tone, speaking style, or acoustic scene. - **Duration**: By default, v4.1-Small predicts the output duration automatically. \ Use Duration Scale for small adjustments or Seconds for exact manual control. """ with gr.Blocks() as demo: with gr.Column(elem_id="col-container"): gr.Markdown(title) gr.Markdown(description) text = gr.Textbox( label="Text", lines=4, elem_id="irodori-voicedesign-text-input", placeholder="合成したいテキストを入力してください", ) build_emoji_palette(text, open=False) caption = gr.Textbox( label="Caption / Style Prompt (optional)", lines=3, placeholder="e.g. 優しく、ゆっくりと", ) uploaded_audios = gr.File( label=( "Reference Audio Uploads (optional; concatenated in displayed order, " "blank = no-reference mode)" ), type="filepath", file_count="multiple", file_types=["audio"], allow_reordering=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="", ) seconds_raw = gr.Textbox( label="Seconds (blank=auto)", value="", ) duration_scale = gr.Slider( label="Duration Scale", minimum=0.5, maximum=1.5, value=1.0, step=0.01, ) with gr.Row(): t_schedule_mode = gr.Dropdown( label="Time Schedule", choices=["linear", "sway"], value="linear", ) sway_coeff = gr.Slider( label="Sway Coeff", minimum=-1.0, maximum=1.5, value=-1.0, step=0.1, interactive=False, ) 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_caption = gr.Slider( label="CFG Scale Caption", minimum=0.0, maximum=10.0, value=4.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) speaker_kv_scale_raw = gr.Textbox(label="Speaker KV Scale (optional)", value="") with gr.Row(): max_text_len_raw = gr.Textbox(label="Max Text Len (optional)", value="") max_caption_len_raw = gr.Textbox(label="Max Caption Len (optional)", value="") 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="") 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="numpy", visible=(i == 0), ) ) out_log = gr.Textbox(label="Run Log", lines=6) gr.Examples( examples=[ ["こんにちは、これは音声合成のテストです。"], ["今日もいい天気ですね。散歩に行きたくなります。"], ["私は人工知能のアシスタントです。何かお手伝いできることはありますか？"], ], inputs=[text], fn=run_inference_gpu, outputs=[*out_audios, out_log], cache_examples=True, cache_mode="lazy", ) def gradio_inference( text_val, caption_val, uploaded_audios_val, num_steps_val, num_candidates_val, seed_raw_val, seconds_raw_val, duration_scale_val, t_schedule_mode_val, sway_coeff_val, cfg_guidance_mode_val, cfg_scale_text_val, cfg_scale_caption_val, cfg_scale_speaker_val, cfg_scale_raw_val, cfg_min_t_val, cfg_max_t_val, context_kv_cache_val, max_text_len_raw_val, max_caption_len_raw_val, truncation_factor_raw_val, rescale_k_raw_val, rescale_sigma_raw_val, speaker_kv_scale_raw_val, ): try: audio_results, log_text = run_inference_gpu( text=text_val, caption=caption_val, uploaded_audios=uploaded_audios_val, num_steps=num_steps_val, num_candidates=num_candidates_val, seed_raw=seed_raw_val, seconds_raw=seconds_raw_val, duration_scale=duration_scale_val, t_schedule_mode=t_schedule_mode_val, sway_coeff=sway_coeff_val, cfg_guidance_mode=cfg_guidance_mode_val, cfg_scale_text=cfg_scale_text_val, cfg_scale_caption=cfg_scale_caption_val, cfg_scale_speaker=cfg_scale_speaker_val, cfg_scale_raw=cfg_scale_raw_val, cfg_min_t=cfg_min_t_val, cfg_max_t=cfg_max_t_val, context_kv_cache=context_kv_cache_val, max_text_len_raw=max_text_len_raw_val, max_caption_len_raw=max_caption_len_raw_val, truncation_factor_raw=truncation_factor_raw_val, rescale_k_raw=rescale_k_raw_val, rescale_sigma_raw=rescale_sigma_raw_val, speaker_kv_scale_raw=speaker_kv_scale_raw_val, ) audio_updates: list[object] = [] for i in range(MAX_GRADIO_CANDIDATES): if i < len(audio_results): audio_updates.append(gr.update(value=audio_results[i], visible=True)) else: audio_updates.append(gr.update(value=None, visible=False)) return (*audio_updates, log_text) except Exception as e: raise gr.Error(str(e)) from e generate_btn.click( fn=gradio_inference, inputs=[ text, caption, uploaded_audios, num_steps, num_candidates, seed_raw, seconds_raw, duration_scale, t_schedule_mode, sway_coeff, cfg_guidance_mode, cfg_scale_text, cfg_scale_caption, cfg_scale_speaker, cfg_scale_raw, cfg_min_t, cfg_max_t, context_kv_cache, max_text_len_raw, max_caption_len_raw, truncation_factor_raw, rescale_k_raw, rescale_sigma_raw, speaker_kv_scale_raw, ], outputs=[*out_audios, out_log], api_name="generate", ) t_schedule_mode.change( _on_t_schedule_mode_change, inputs=[t_schedule_mode], outputs=[sway_coeff] ) return demo if __name__ == "__main__": demo = build_demo() demo.queue(default_concurrency_limit=1) demo.launch(theme=gr.themes.Citrus(), css=EMOJI_PALETTE_CSS + CSS, mcp_server=True)