Upload 8 files

README.md

@@ -1,6 +1,6 @@
 ---
 title: MioTTS 0.1B Demo
-emoji: 👀
+emoji: 📈
 colorFrom: pink
 colorTo: gray
 sdk: gradio
@@ -8,6 +8,11 @@ sdk_version: 6.5.1
 python_version: '3.12'
 app_file: app.py
 pinned: false
+license: mit
+short_description: TTS demo for MioTTS-0.1B
+models:
+  - Aratako/MioTTS-0.1B
+  - Aratako/MioCodec-25Hz-24kHz
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,444 @@
+import os
+import random
+import re
+from typing import Optional
+
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from miocodec import MioCodecModel
+
+from text import normalize_text
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+MODEL_REPO = os.environ.get("MODEL_REPO", "Aratako/MioTTS-0.1B")
+CODEC_REPO = os.environ.get("CODEC_REPO", "Aratako/MioCodec-25Hz-24kHz")
+
+# Global variables for lazy loading
+_model = None
+_tokenizer = None
+_codec = None
+
+# Presets directory
+PRESETS_DIR = "presets"
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+TOKEN_PATTERN = re.compile(r"<\|s_(\d+)\|>")
+
+
+def seed_everything(seed: Optional[int]) -> int:
+    if seed is None:
+        seed = random.SystemRandom().randint(0, 2**31 - 1)
+        print(f"[Info] No seed provided; using random seed {seed}")
+
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    return seed
+
+
+def parse_speech_tokens(text: str) -> list[int]:
+    tokens = [int(value) for value in TOKEN_PATTERN.findall(text)]
+    if not tokens:
+        raise ValueError("No speech tokens found in LLM output.")
+    return tokens
+
+
+# ---------------------------------------------------------------------------
+# Model Loading
+# ---------------------------------------------------------------------------
+
+def load_models():
+    global _model, _tokenizer, _codec
+
+    if _model is not None:
+        return
+
+    print(f"[Info] Loading LLM from {MODEL_REPO}...")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    _tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
+    _model = AutoModelForCausalLM.from_pretrained(
+        MODEL_REPO,
+        torch_dtype=torch.bfloat16,
+    ).to(device)
+    _model.eval()
+
+    print(f"[Info] Loading codec from {CODEC_REPO}...")
+    _codec = MioCodecModel.from_pretrained(CODEC_REPO)
+    _codec = _codec.eval().to(device)
+
+    print("[Info] Models loaded successfully.")
+
+
+def get_preset_list() -> list[str]:
+    if not os.path.exists(PRESETS_DIR):
+        return []
+    presets = []
+    for f in os.listdir(PRESETS_DIR):
+        if f.endswith(".pt"):
+            presets.append(f[:-3])
+    return sorted(presets)
+
+
+def load_preset_embedding(preset_id: str) -> torch.Tensor:
+    path = os.path.join(PRESETS_DIR, f"{preset_id}.pt")
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"Preset '{preset_id}' not found.")
+    embedding = torch.load(path, map_location="cpu", weights_only=True)
+    if isinstance(embedding, dict):
+        embedding = embedding.get("global_embedding", embedding)
+    return embedding.squeeze()
+
+
+# ---------------------------------------------------------------------------
+# GPU-decorated Inference Functions
+# ---------------------------------------------------------------------------
+
+@spaces.GPU(duration=120)
+def run_inference_gpu(
+    target_text: str,
+    reference_mode: str,
+    reference_audio: Optional[tuple[int, np.ndarray]],
+    preset_id: Optional[str],
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    repetition_penalty: float,
+    max_tokens: int,
+    seed: Optional[int],
+    num_samples: int = 1,
+) -> list[tuple[int, np.ndarray]]:
+    load_models()
+
+    used_seed = seed_everything(None if seed is None else int(seed))
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Normalize text
+    normalized_text = normalize_text(target_text)
+    print(f"[Info] Normalized text: {normalized_text}")
+
+    # Prepare reference
+    reference_waveform = None
+    global_embedding = None
+
+    if reference_mode == "upload" and reference_audio is not None:
+        sr, audio = reference_audio
+        # Convert to tensor
+        if audio.ndim == 1:
+            audio_tensor = torch.from_numpy(audio).float()
+        else:
+            audio_tensor = torch.from_numpy(audio.mean(axis=1)).float()
+
+        # Resample if needed
+        codec_sr = _codec.config.sample_rate
+        if sr != codec_sr:
+            import torchaudio
+            audio_tensor = audio_tensor.unsqueeze(0)
+            resampler = torchaudio.transforms.Resample(sr, codec_sr)
+            audio_tensor = resampler(audio_tensor).squeeze(0)
+
+        # Trim to max 20 seconds
+        max_samples = int(codec_sr * 20)
+        if audio_tensor.shape[0] > max_samples:
+            audio_tensor = audio_tensor[:max_samples]
+            print(f"[Info] Reference audio trimmed to 20 seconds")
+
+        reference_waveform = audio_tensor.to(device)
+    elif reference_mode == "preset" and preset_id:
+        global_embedding = load_preset_embedding(preset_id).to(device)
+    else:
+        raise ValueError("Either reference audio or preset must be provided.")
+
+    # Tokenize input
+    messages = [{"role": "user", "content": normalized_text}]
+    input_text = _tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    inputs = _tokenizer(input_text, return_tensors="pt").to(device)
+    # Remove token_type_ids if present (not used by this model)
+    inputs.pop("token_type_ids", None)
+
+    # Generate (batch)
+    with torch.no_grad():
+        outputs = _model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            repetition_penalty=repetition_penalty,
+            do_sample=True,
+            pad_token_id=_tokenizer.eos_token_id,
+            num_return_sequences=num_samples,
+        )
+
+    # Parse all generated sequences
+    tokens_list = []
+    for i in range(outputs.shape[0]):
+        generated_text = _tokenizer.decode(outputs[i], skip_special_tokens=False)
+        generated_part = generated_text[len(input_text):]
+        try:
+            speech_tokens = parse_speech_tokens(generated_part)
+            tokens_list.append(speech_tokens)
+        except ValueError as e:
+            print(f"[Warning] Sample {i + 1}: {e}")
+
+    if not tokens_list:
+        raise ValueError("No valid speech tokens generated.")
+
+    # Decode audio (batch)
+    results = []
+    sample_rate = _codec.config.sample_rate
+
+    with torch.no_grad():
+        # Prepare batch tokens
+        max_len = max(len(t) for t in tokens_list)
+        batch_tokens = torch.zeros((len(tokens_list), max_len), dtype=torch.long, device=device)
+        content_lengths = []
+        for i, tokens in enumerate(tokens_list):
+            batch_tokens[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long)
+            content_lengths.append(len(tokens))
+
+        # Get global embeddings
+        if reference_waveform is not None:
+            # Extract global embedding from reference waveform
+            ref_features = _codec.encode(reference_waveform, return_content=False, return_global=True)
+            global_embeddings = ref_features.global_embedding.unsqueeze(0).expand(len(tokens_list), -1)
+        else:
+            global_embeddings = global_embedding.unsqueeze(0).expand(len(tokens_list), -1)
+
+        # Batch decode
+        audio_batch, audio_lengths = _codec.decode_batch(
+            global_embeddings=global_embeddings,
+            content_token_indices=batch_tokens,
+            content_lengths=content_lengths,
+        )
+        for i in range(len(tokens_list)):
+            audio_len = int(audio_lengths[i])
+            audio_np = audio_batch[i, :audio_len].cpu().numpy()
+            results.append((sample_rate, audio_np))
+
+    print(f"[Info] Seed used: {used_seed}")
+    return results
+
+
+# Load models at startup
+load_models()
+
+
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+
+MAX_NUM_SAMPLES = 32
+
+
+def gradio_inference(
+    target_text: str,
+    reference_mode: str,
+    reference_audio: Optional[tuple[int, np.ndarray]],
+    preset_id: Optional[str],
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    repetition_penalty: float,
+    max_tokens: int,
+    seed: str,
+    num_samples: int,
+):
+    if not target_text.strip():
+        outputs = [gr.update(value=None, visible=False) for _ in range(MAX_NUM_SAMPLES)]
+        return outputs
+
+    seed_val = None
+    if seed.strip() not in {"", "None", "none"}:
+        seed_val = int(float(seed))
+
+    try:
+        results = run_inference_gpu(
+            target_text=target_text,
+            reference_mode=reference_mode,
+            reference_audio=reference_audio,
+            preset_id=preset_id,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            repetition_penalty=repetition_penalty,
+            max_tokens=max_tokens,
+            seed=seed_val,
+            num_samples=int(num_samples),
+        )
+    except Exception as e:
+        print(f"[Error] {e}")
+        raise gr.Error(str(e))
+
+    outputs = []
+    for i in range(MAX_NUM_SAMPLES):
+        if i < len(results):
+            outputs.append(gr.update(value=results[i], visible=True))
+        else:
+            outputs.append(gr.update(value=None, visible=False))
+    return outputs
+
+
+def build_demo():
+    presets = get_preset_list()
+
+    MODEL_LINK = f"https://huggingface.co/{MODEL_REPO}"
+    GITHUB_REPO = "https://github.com/Aratako/MioTTS-Inference"
+
+    title = "# MioTTS-0.1B Demo"
+    description = f"""
+- **Model**: [{MODEL_REPO}]({MODEL_LINK})
+- For faster and more efficient inference, see [MioTTS-Inference]({GITHUB_REPO})
+
+**Usage:**
+- Select a preset voice OR upload your own reference audio (max 20 seconds)
+- Enter text to synthesize
+- Adjust generation parameters as needed
+"""
+
+    with gr.Blocks() as demo:
+        gr.Markdown(title)
+        gr.Markdown(description)
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                reference_mode = gr.Radio(
+                    choices=["preset", "upload"],
+                    value="preset",
+                    label="Reference Mode",
+                )
+                preset_id = gr.Dropdown(
+                    choices=presets,
+                    value=presets[0] if presets else None,
+                    label="Preset Voice",
+                    allow_custom_value=False,
+                    visible=True,
+                )
+                reference_audio = gr.Audio(
+                    label="Reference Audio",
+                    type="numpy",
+                    visible=False,
+                )
+
+        def update_reference_visibility(mode):
+            if mode == "preset":
+                return gr.update(visible=True), gr.update(visible=False)
+            else:
+                return gr.update(visible=False), gr.update(visible=True)
+
+        reference_mode.change(
+            fn=update_reference_visibility,
+            inputs=[reference_mode],
+            outputs=[preset_id, reference_audio],
+        )
+
+        target_text = gr.Textbox(
+            label="Text to Synthesize",
+            value="",
+            placeholder="Enter text to synthesize",
+            lines=3,
+        )
+
+        with gr.Row():
+            seed_box = gr.Textbox(
+                label="Seed (optional)",
+                value="",
+                placeholder="Leave blank for random",
+            )
+            num_samples = gr.Slider(
+                label="Number of Samples",
+                minimum=1,
+                maximum=MAX_NUM_SAMPLES,
+                step=1,
+                value=1,
+            )
+
+        with gr.Row():
+            temperature = gr.Slider(
+                label="Temperature", minimum=0.1, maximum=1.5, step=0.05, value=0.8
+            )
+            top_p = gr.Slider(
+                label="Top-p", minimum=0.1, maximum=1.0, step=0.05, value=1.0
+            )
+            top_k = gr.Slider(
+                label="Top-k", minimum=0, maximum=100, step=1, value=50
+            )
+
+        with gr.Row():
+            repetition_penalty = gr.Slider(
+                label="Repetition Penalty",
+                minimum=1.0,
+                maximum=1.5,
+                step=0.05,
+                value=1.0,
+            )
+            max_tokens = gr.Slider(
+                label="Max Tokens",
+                minimum=100,
+                maximum=1000,
+                step=50,
+                value=700,
+            )
+
+        generate_button = gr.Button("Generate", variant="primary")
+
+        # Output audio components
+        output_audios = []
+        cols_per_row = 4
+        num_rows = (MAX_NUM_SAMPLES + cols_per_row - 1) // cols_per_row
+        with gr.Column():
+            for row_idx in range(num_rows):
+                with gr.Row():
+                    for col_idx in range(cols_per_row):
+                        i = row_idx * cols_per_row + col_idx
+                        if i >= MAX_NUM_SAMPLES:
+                            break
+                        audio = gr.Audio(
+                            label=f"Sample #{i+1}",
+                            type="numpy",
+                            interactive=False,
+                            visible=(i == 0),
+                        )
+                        output_audios.append(audio)
+
+        generate_button.click(
+            fn=gradio_inference,
+            inputs=[
+                target_text,
+                reference_mode,
+                reference_audio,
+                preset_id,
+                temperature,
+                top_p,
+                top_k,
+                repetition_penalty,
+                max_tokens,
+                seed_box,
+                num_samples,
+            ],
+            outputs=output_audios,
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    demo = build_demo()
+    demo.launch()

presets/en_female.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a386d2ee0b48036586fc322da0ccf4b88f585ef9162e0371a9999637ebd7645
+size 1997

presets/en_male.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:33a37cffd19795491707edcd760b5f9ecf9da1354296fbca6b7add25f2de42d1
+size 2096

presets/jp_female.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b9fb6483be458d81b7edcb8edc49487a781ea2a344495f964bfd0d463d560dba
+size 2103

presets/jp_male.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb6667eff5e7a1a80314dd2d52b3a818b7cc00d54a3b659e1b21f2423581cb82
+size 2096

requirements.txt

@@ -0,0 +1,8 @@
+torch
+torchaudio
+transformers<5
+accelerate
+gradio
+soundfile
+numpy
+miocodec @ git+https://github.com/Aratako/MioCodec@main

text.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import re
+
+REPLACE_MAP: dict[str, str] = {
+    r"\t": "",
+    r"\[n\]": "",
+    r"　": "",
+    r"[;▼♀♂《》≪≫①②③④⑤⑥]": "",
+    r"[\u02d7\u2010-\u2015\u2043\u2212\u23af\u23e4\u2500\u2501\u2e3a\u2e3b]": "",
+    r"[\uff5e\u301C]": "ー",
+    r"？": "?",
+    r"！": "!",
+    r"[●◯〇]": "○",
+    r"♥": "♡",
+}
+
+FULLWIDTH_ALPHA_TO_HALFWIDTH = str.maketrans(
+    {
+        chr(full): chr(half)
+        for full, half in zip(
+            list(range(0xFF21, 0xFF3B)) + list(range(0xFF41, 0xFF5B)),
+            list(range(0x41, 0x5B)) + list(range(0x61, 0x7B)),
+            strict=True,
+        )
+    }
+)
+_HALFWIDTH_KATAKANA_CHARS = "ｦｧｨｩｪｫｬｭｮｯｰｱｲｳｴｵｶｷｸｹｺｻｼｽｾｿﾀﾁﾂﾃﾄﾅﾆﾇﾈﾉﾊﾋﾌﾍﾎﾏﾐﾑﾒﾓﾔﾕﾖﾗﾘﾙﾚﾛﾜﾝ"
+_FULLWIDTH_KATAKANA_CHARS = "ヲァィゥェォャュョッーアイウエオカキクケコサシスセソタチツテトナニヌネノハヒフヘホマミムメモヤユヨラリルレロワン"
+HALFWIDTH_KATAKANA_TO_FULLWIDTH = str.maketrans(
+    _HALFWIDTH_KATAKANA_CHARS, _FULLWIDTH_KATAKANA_CHARS
+)
+FULLWIDTH_DIGITS_TO_HALFWIDTH = str.maketrans(
+    {
+        chr(full): chr(half)
+        for full, half in zip(range(0xFF10, 0xFF1A), range(0x30, 0x3A), strict=True)
+    }
+)
+
+
+def normalize_text(text: str) -> str:
+    """Normalize text for TTS."""
+    for pattern, replacement in REPLACE_MAP.items():
+        text = re.sub(pattern, replacement, text)
+
+    text = text.translate(FULLWIDTH_ALPHA_TO_HALFWIDTH)
+    text = text.translate(FULLWIDTH_DIGITS_TO_HALFWIDTH)
+    text = text.translate(HALFWIDTH_KATAKANA_TO_FULLWIDTH)
+
+    text = re.sub(r"…{3,}", "……", text)
+
+    if text.startswith("「") and text.endswith("」"):
+        text = text[1:-1]
+    if text.startswith("『") and text.endswith("』"):
+        text = text[1:-1]
+    if text.startswith("（") and text.endswith("）"):
+        text = text[1:-1]
+    if text.startswith("【") and text.endswith("】"):
+        text = text[1:-1]
+    if text.startswith("(") and text.endswith(")"):
+        text = text[1:-1]
+
+    if text.endswith("。"):
+        text = text.rstrip("。")
+
+    return text