app.py

#!/usr/bin/env python3
"""
ZipVoice Gradio Web Interface for HuggingFace Spaces
Updated for Gradio 5.47.0 compatibility
"""

import os
import sys
import tempfile
import gradio as gr
import torch
from pathlib import Path
import spaces
import whisper

# Add current directory to Python path for local zipvoice package
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# Import ZipVoice components
from zipvoice.models.zipvoice import ZipVoice
from zipvoice.models.zipvoice_distill import ZipVoiceDistill
from zipvoice.tokenizer.tokenizer import EmiliaTokenizer
from zipvoice.utils.checkpoint import load_checkpoint
from zipvoice.utils.feature import VocosFbank
from zipvoice.bin.infer_zipvoice import generate_sentence
from lhotse.utils import fix_random_seed

# Global variables for caching models
_models_cache = {}
_tokenizer_cache = None
_vocoder_cache = None
_feature_extractor_cache = None


def load_models_and_components(model_name: str):
    """Load and cache models, tokenizer, vocoder, and feature extractor."""
    global _models_cache, _tokenizer_cache, _vocoder_cache, _feature_extractor_cache

    # Set device (GPU if available for Spaces GPU acceleration)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if model_name not in _models_cache:
        print(f"Loading {model_name} model...")

        # Model directory mapping
        model_dir_map = {
            "zipvoice": "zipvoice",
            "zipvoice_distill": "zipvoice_distill",
        }

        huggingface_repo = "k2-fsa/ZipVoice"

        # Download model files from HuggingFace
        from huggingface_hub import hf_hub_download

        model_ckpt = hf_hub_download(
            huggingface_repo, filename=f"{model_dir_map[model_name]}/model.pt"
        )
        model_config_path = hf_hub_download(
            huggingface_repo, filename=f"{model_dir_map[model_name]}/model.json"
        )
        token_file = hf_hub_download(
            huggingface_repo, filename=f"{model_dir_map[model_name]}/tokens.txt"
        )

        # Load tokenizer (cache it)
        if _tokenizer_cache is None:
            _tokenizer_cache = EmiliaTokenizer(token_file=token_file)
        tokenizer = _tokenizer_cache
        tokenizer_config = {"vocab_size": tokenizer.vocab_size, "pad_id": tokenizer.pad_id}

        # Load model configuration
        import json
        with open(model_config_path, "r") as f:
            model_config = json.load(f)

        # Create model
        if model_name == "zipvoice":
            model = ZipVoice(**model_config["model"], **tokenizer_config)
        else:
            model = ZipVoiceDistill(**model_config["model"], **tokenizer_config)

        # Load model weights
        load_checkpoint(filename=model_ckpt, model=model, strict=True)
        model = model.to(device)
        model.eval()

        _models_cache[model_name] = model

    # Load vocoder (cache it)
    if _vocoder_cache is None:
        from vocos import Vocos
        _vocoder_cache = Vocos.from_pretrained("charactr/vocos-mel-24khz")
        _vocoder_cache = _vocoder_cache.to(device)
        _vocoder_cache.eval()

    # Load feature extractor (cache it)
    if _feature_extractor_cache is None:
        _feature_extractor_cache = VocosFbank()

    return (_models_cache[model_name], _tokenizer_cache,
            _vocoder_cache, _feature_extractor_cache,
            model_config["feature"]["sampling_rate"])


@spaces.GPU
def transcribe_audio_whisper(audio_file):
    """Transcribe audio file using Whisper."""
    if audio_file is None:
        return "Error: Please upload an audio file first."

    try:
        # Load Whisper model (will be done on GPU)
        model = whisper.load_model("small")

        # Save uploaded audio to temporary file for processing
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio:
            temp_audio_path = temp_audio.name
            with open(temp_audio_path, "wb") as f:
                f.write(audio_file)

        # Transcribe the audio
        result = model.transcribe(temp_audio_path)

        # Clean up temporary file
        os.unlink(temp_audio_path)

        return result["text"].strip()

    except Exception as e:
        return f"Error during transcription: {str(e)}"


@spaces.GPU
def synthesize_speech_gradio(
    text: str,
    prompt_audio_file,
    prompt_text: str,
    model_name: str,
    speed: float
):
    """Synthesize speech using ZipVoice for Gradio interface."""
    if not text.strip():
        return None, "Error: Please enter text to synthesize."

    if prompt_audio_file is None:
        return None, "Error: Please upload a prompt audio file."

    if not prompt_text.strip():
        return None, "Error: Please enter the transcription of the prompt audio."

    try:
        # Set random seed for reproducibility
        fix_random_seed(666)

        # Load models and components
        model, tokenizer, vocoder, feature_extractor, sampling_rate = load_models_and_components(model_name)

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Save uploaded audio to temporary file
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio:
            temp_audio_path = temp_audio.name
            with open(temp_audio_path, "wb") as f:
                f.write(prompt_audio_file)

        # Create temporary output file
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_output:
            output_path = temp_output.name

        print(f"Synthesizing: '{text}' using {model_name}")
        print(f"Prompt: {prompt_text}")
        print(f"Speed: {speed}")

        # Generate speech
        with torch.inference_mode():
            metrics = generate_sentence(
                save_path=output_path,
                prompt_text=prompt_text,
                prompt_wav=temp_audio_path,
                text=text,
                model=model,
                vocoder=vocoder,
                tokenizer=tokenizer,
                feature_extractor=feature_extractor,
                device=device,
                num_step=16 if model_name == "zipvoice" else 8,
                guidance_scale=1.0 if model_name == "zipvoice" else 3.0,
                speed=speed,
                t_shift=0.5,
                target_rms=0.1,
                feat_scale=0.1,
                sampling_rate=sampling_rate,
                max_duration=100,
                remove_long_sil=False,
            )

        # Read the generated audio file
        with open(output_path, "rb") as f:
            audio_data = f.read()

        # Clean up temporary files
        os.unlink(temp_audio_path)
        os.unlink(output_path)

        success_msg = f"Synthesis completed! Duration: {metrics['wav_seconds']:.2f}s, RTF: {metrics['rtf']:.2f}"
        return audio_data, success_msg

    except Exception as e:
        error_msg = f"Error during synthesis: {str(e)}"
        print(error_msg)
        return None, error_msg


def create_gradio_interface():
    """Create the Gradio web interface."""

    # Enhanced CSS for modern UI/UX
    css = """
    .gradio-container {
        max-width: 1400px;
        margin: auto;
        font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
    }
    .title {
        text-align: center;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        font-size: 3.5em;
        font-weight: 800;
        margin-bottom: 0.5em;
        letter-spacing: -0.02em;
    }
    .subtitle {
        text-align: center;
        color: #64748b;
        font-size: 1.3em;
        margin-bottom: 2.5em;
        font-weight: 300;
    }
    .step-card {
        background: linear-gradient(145deg, #f8fafc, #e2e8f0);
        border: 1px solid #cbd5e1;
        border-radius: 16px;
        padding: 1.5em;
        margin: 1em 0;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
        transition: all 0.3s ease;
    }
    .step-card:hover {
        transform: translateY(-2px);
        box-shadow: 0 8px 25px -5px rgba(0, 0, 0, 0.1);
    }
    .step-number {
        background: linear-gradient(135deg, #667eea, #764ba2);
        color: white;
        width: 32px;
        height: 32px;
        border-radius: 50%;
        display: inline-flex;
        align-items: center;
        justify-content: center;
        font-weight: bold;
        font-size: 0.9em;
        margin-right: 12px;
    }
    .feature-grid {
        display: grid;
        grid-template-columns: repeat(auto-fit, minmax(300px, 1fr));
        gap: 1.5em;
        margin: 2em 0;
    }
    .feature-card {
        background: white;
        border: 1px solid #e2e8f0;
        border-radius: 12px;
        padding: 1.5em;
        box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
        transition: all 0.3s ease;
    }
    .feature-card:hover {
        border-color: #667eea;
        box-shadow: 0 8px 25px rgba(102, 126, 234, 0.1);
    }
    .btn-primary {
        background: linear-gradient(135deg, #667eea, #764ba2) !important;
        border: none !important;
        color: white !important;
        font-weight: 600 !important;
        transition: all 0.3s ease !important;
    }
    .btn-primary:hover {
        transform: translateY(-1px) !important;
        box-shadow: 0 8px 25px rgba(102, 126, 234, 0.3) !important;
    }
    .output-section {
        background: linear-gradient(145deg, #f1f5f9, #e2e8f0);
        border-radius: 16px;
        padding: 2em;
        margin-top: 1em;
    }
    .example-card {
        background: white;
        border: 1px solid #e2e8f0;
        border-radius: 8px;
        padding: 1em;
        margin: 0.5em 0;
        transition: all 0.2s ease;
    }
    .example-card:hover {
        border-color: #667eea;
        background: #fafbfc;
    }
    """

    with gr.Blocks(title="ZipVoice - Zero-Shot Text-to-Speech", css=css) as interface:

        gr.HTML("""
        <div class="title">🎵 ZipVoice</div>
        <div class="subtitle">Fast and High-Quality Zero-Shot Text-to-Speech with Flow Matching</div>

        <div style="background: #f8fafc; border: 1px solid #e2e8f0; border-radius: 8px; padding: 1.5em; margin: 1em 0; font-size: 0.9em;">
            <h3 style="margin-top: 0; color: #1e293b;">📖 How to Use / 使用說明</h3>

            <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 2em; margin-top: 1em;">
                <div>
                    <h4 style="color: #2563eb; margin-bottom: 0.5em;">English / 英文</h4>
                    <ol style="margin: 0; padding-left: 1.2em; line-height: 1.6;">
                        <li><b>Upload Audio:</b> Choose a short audio clip (1-3 seconds) of the voice you want to clone</li>
                        <li><b>Transcribe:</b> Click "🎤 Transcribe Audio" to get automatic transcription</li>
                        <li><b>Enter Text:</b> Type the text you want to convert to speech</li>
                        <li><b>Choose Model:</b> Select ZipVoice (better quality) or ZipVoice Distill (faster)</li>
                        <li><b>Adjust Speed:</b> Modify speech speed (0.5 = slower, 2.0 = faster)</li>
                        <li><b>Generate:</b> Click "🎵 Generate Speech" to create your audio</li>
                    </ol>
                    <p style="margin-top: 1em; color: #64748b;"><b>Tips:</b> Use clear audio with minimal background noise for best results.</p>
                </div>

                <div>
                    <h4 style="color: #2563eb; margin-bottom: 0.5em;">繁體中文 / Traditional Chinese</h4>
                    <ol style="margin: 0; padding-left: 1.2em; line-height: 1.6;">
                        <li><b>上傳音訊：</b>選擇一個簡短的音訊片段（1-3秒）作為要克隆的聲音</li>
                        <li><b>轉錄音訊：</b>點選「🎤 Transcribe Audio」按鈕進行自動轉錄，或自行輸入音訊片段的文字</li>
                        <li><b>輸入文字：</b>輸入您要轉換成語音的文字</li>
                        <li><b>選擇模型：</b>選擇 ZipVoice（品質較好）或 ZipVoice Distill（速度較快）</li>
                        <li><b>調整速度：</b>修改語音速度（0.5 = 較慢，2.0 = 較快）</li>
                        <li><b>生成語音：</b>點選「🎵 Generate Speech」生成音訊</li>
                    </ol>
                    <p style="margin-top: 1em; color: #64748b;"><b>提示：</b>使用清晰且背景噪音少的音頻以獲得最佳效果。</p>
                </div>
            </div>
        </div>
        """)

        with gr.Row():
            with gr.Column(scale=2):
                text_input = gr.Textbox(
                    label="Text to Synthesize",
                    placeholder="Enter the text you want to convert to speech...",
                    lines=3,
                    value="這是一則語音測試"
                )

                with gr.Row():
                    model_dropdown = gr.Dropdown(
                        choices=["zipvoice", "zipvoice_distill"],
                        value="zipvoice",
                        label="Model"
                    )

                    speed_slider = gr.Slider(
                        minimum=0.5,
                        maximum=2.0,
                        value=1.0,
                        step=0.1,
                        label="Speed"
                    )

                prompt_audio = gr.File(
                    label="Prompt Audio",
                    file_types=["audio"],
                    type="binary"
                )

                prompt_text = gr.Textbox(
                    label="Prompt Transcription",
                    placeholder="Enter the exact transcription of the prompt audio...",
                    lines=2
                )

                transcribe_btn = gr.Button(
                    "🎤 Transcribe Audio",
                    variant="secondary",
                    size="sm"
                )

                generate_btn = gr.Button(
                    "🎵 Generate Speech",
                    variant="primary",
                    size="lg"
                )

            with gr.Column(scale=1):
                output_audio = gr.Audio(
                    label="Generated Speech",
                    type="filepath"
                )

                status_text = gr.Textbox(
                    label="Status",
                    interactive=False,
                    lines=3
                )

                gr.Examples(
                    examples=[
                        ["I have a dream that one day this nation will rise up and live out the true meaning of its creed.", "jfk.wav", "ask not what your country can do for you, ask what you can do for your country", "zipvoice", 1.0],
                        ["今天天氣真好，我們去公園散步吧！", "jfk.wav", "ask not what your country can do for you, ask what you can do for your country", "zipvoice", 1.0],
                        ["The quick brown fox jumps over the lazy dog.", "jfk.wav", "ask not what your country can do for you, ask what you can do for your country", "zipvoice_distill", 1.2],
                    ],
                    inputs=[text_input, prompt_audio, prompt_text, model_dropdown, speed_slider],
                    label="Quick Examples"
                )

        # Event handling
        transcribe_btn.click(
            fn=transcribe_audio_whisper,
            inputs=[prompt_audio],
            outputs=[prompt_text]
        )

        generate_btn.click(
            fn=synthesize_speech_gradio,
            inputs=[text_input, prompt_audio, prompt_text, model_dropdown, speed_slider],
            outputs=[output_audio, status_text]
        )

        # Footer
        gr.HTML("""
        <div style="text-align: center; margin-top: 2em; color: #64748b; font-size: 0.9em;">
            <p>Powered by <a href="https://github.com/k2-fsa/ZipVoice" target="_blank">ZipVoice</a> |
            Built with <a href="https://gradio.app" target="_blank">Gradio</a></p>
            <p>Upload a short audio clip as prompt, and ZipVoice will synthesize speech in that voice style!</p>
        </div>
        """)

    return interface


if __name__ == "__main__":
    # Create and launch the interface
    interface = create_gradio_interface()
    interface.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        show_error=True
    )