Qwen3-TTS-streaming/api/gradio_demo.py

"""
Gradio demo for Qwen3-TTS API.
Provides a user-friendly interface for testing the streaming TTS API.
"""

import requests
import base64
import io
import time
from pathlib import Path

import gradio as gr
import soundfile as sf
import numpy as np


class TTSAPIClient:
    """Client for interacting with Qwen3-TTS API."""
    
    def __init__(self, api_url: str = "http://localhost:8000"):
        self.api_url = api_url
    
    def health_check(self) -> dict:
        """Check API health status."""
        try:
            response = requests.get(f"{self.api_url}/v1/health", timeout=5)
            return response.json() if response.status_code == 200 else {"status": "unavailable"}
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def list_models(self) -> list:
        """Get list of available models."""
        try:
            response = requests.get(f"{self.api_url}/v1/models", timeout=5)
            if response.status_code == 200:
                data = response.json()
                return [m["id"] for m in data.get("data", [])]
            return []
        except Exception as e:
            return f"Error: {str(e)}"
    
    def generate_speech(
        self,
        text: str,
        language: str = "English",
        voice_clone: bool = False,
    ) -> tuple:
        """
        Generate speech using TTS API.
        
        Returns:
            Tuple of (audio_data, sample_rate) or (None, error_message)
        """
        try:
            payload = {
                "text": text,
                "language": language,
                "voice_clone_mode": "reference_audio" if voice_clone else "disabled",
            }
            
            response = requests.post(
                f"{self.api_url}/v1/audio/speech",
                json=payload,
                timeout=60,
            )
            
            if response.status_code == 200:
                data = response.json()
                # Decode base64 audio
                audio_bytes = base64.b64decode(data["audio_base64"])
                audio, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32")
                return (sr, audio), None
            else:
                error = response.json().get("error", "Unknown error")
                return None, f"Error: {error}"
        
        except requests.exceptions.Timeout:
            return None, "Request timeout - model may be loading or busy"
        except Exception as e:
            return None, f"Error: {str(e)}"
    
    def stream_speech(
        self,
        text: str,
        language: str = "English",
        voice_clone: bool = False,
    ) -> tuple:
        """
        Stream speech from TTS API with chunks.
        
        Yields audio chunks to stream in near real-time.
        """
        try:
            payload = {
                "text": text,
                "language": language,
                "voice_clone_mode": "reference_audio" if voice_clone else "disabled",
                "stream_options": {
                    "emit_every_frames": 8,
                    "decode_window_frames": 80,
                    "overlap_samples": 512,
                }
            }
            
            response = requests.post(
                f"{self.api_url}/v1/audio/speech/stream",
                json=payload,
                stream=True,
                timeout=120,
            )
            
            if response.status_code != 200:
                error = response.json().get("error", "Unknown error")
                return None, f"Error: {error}"
            
            chunks = []
            sample_rate = None
            chunk_count = 0
            
            # Read streaming chunks with length prefix (4 bytes big-endian)
            buffer = b""
            for chunk in response.iter_content(chunk_size=65536):
                if chunk:
                    buffer += chunk
                    
                    # Process complete frames
                    while len(buffer) >= 4:
                        frame_len = int.from_bytes(buffer[:4], byteorder='big')
                        
                        if len(buffer) < 4 + frame_len:
                            break  # Need more data
                        
                        frame_data = buffer[4:4 + frame_len]
                        buffer = buffer[4 + frame_len:]
                        
                        try:
                            audio, sample_rate = sf.read(
                                io.BytesIO(frame_data),
                                dtype="float32"
                            )
                            chunks.append(audio)
                            chunk_count += 1
                        except Exception as e:
                            print(f"Error decoding chunk: {e}")
            
            if chunks and sample_rate:
                final_audio = np.concatenate(chunks)
                return (sample_rate, final_audio), None
            else:
                return None, "No audio data received"
        
        except requests.exceptions.Timeout:
            return None, "Request timeout - stream took too long"
        except Exception as e:
            return None, f"Streaming error: {str(e)}"


# Initialize API client
client = TTSAPIClient()


def check_api_status():
    """Check API health and return status message."""
    health = client.health_check()
    if health.get("status") == "healthy":
        return f"✅ API Ready\nModel: {health.get('model', 'Unknown')}\nDevice: {health.get('device', 'Unknown')}"
    else:
        return f"❌ API Unavailable\nMake sure the API server is running on http://localhost:8000"


def synthesize_speech(text: str, language: str, use_streaming: bool, use_voice_clone: bool) -> tuple:
    """
    Synthesize speech using the selected options.
    
    Returns:
        Tuple of (audio_output, status_message)
    """
    if not text.strip():
        return None, "❌ Error: Please enter some text"
    
    # Check API status
    health = client.health_check()
    if health.get("status") != "healthy":
        return None, "❌ Error: API is not available. Please start the API server."
    
    status_update = f"🔄 Generating speech ({len(text)} characters)...\n"
    status_update += f"Language: {language}\n"
    status_update += f"Voice Clone: {'Yes' if use_voice_clone else 'No'}\n"
    status_update += f"Streaming: {'Yes' if use_streaming else 'No'}\n\n"
    
    start_time = time.time()
    
    if use_streaming:
        result, error = client.stream_speech(text, language, use_voice_clone)
    else:
        result, error = client.generate_speech(text, language, use_voice_clone)
    
    elapsed = time.time() - start_time
    
    if error:
        return None, f"❌ {error}"
    
    if result:
        sample_rate, audio = result
        duration = len(audio) / sample_rate
        
        status_update += f"✅ Success!\n"
        status_update += f"Duration: {duration:.2f}s\n"
        status_update += f"Generation time: {elapsed:.2f}s\n"
        status_update += f"Sample rate: {sample_rate} Hz\n"
        
        if use_streaming:
            status_update += f"(Streaming optimized)\n"
        
        return (sample_rate, audio), status_update
    
    return None, "❌ Error: No audio generated"


def create_demo():
    """Create and return the Gradio demo interface."""
    
    with gr.Blocks(
        title="Qwen3-TTS API Demo",
        theme=gr.themes.Soft(),
    ) as demo:
        gr.Markdown("""
        # 🎙️ Qwen3-TTS Streaming API Demo
        
        This demo showcases the **Qwen3-TTS OpenAI-like API** with support for:
        - 🎵 Real-time streaming audio generation
        - 🎭 Voice cloning (when reference audio is provided)
        - 🌍 Multiple language support
        
        **Note:** Make sure the API server is running at `http://localhost:8000`
        """)
        
        # API Status Section
        with gr.Group():
            gr.Markdown("### API Status")
            status_button = gr.Button("Check API Status", variant="primary")
            status_output = gr.Textbox(
                label="Status",
                interactive=False,
                lines=3,
            )
            status_button.click(
                fn=check_api_status,
                outputs=status_output,
            )
        
        # Main Input Section
        with gr.Group():
            gr.Markdown("### Text to Speech Settings")
            
            with gr.Row():
                text_input = gr.Textbox(
                    label="Text to Synthesize",
                    placeholder="Enter the text you want to convert to speech...",
                    lines=4,
                    scale=3,
                )
                
                with gr.Column(scale=1):
                    language_select = gr.Dropdown(
                        choices=["English", "Russian", "Chinese", "Japanese", "Korean", "Auto"],
                        value="English",
                        label="Language",
                    )
                    
                    use_streaming = gr.Checkbox(
                        label="Use Streaming",
                        value=True,
                        info="Stream audio for faster first chunk",
                    )
                    
                    use_voice_clone = gr.Checkbox(
                        label="Voice Cloning",
                        value=False,
                        info="Use reference voice (if available)",
                    )
            
            # Synthesis Button
            with gr.Row():
                synthesize_button = gr.Button(
                    "Generate Speech 🎵",
                    variant="primary",
                    size="lg",
                )
        
        # Output Section
        with gr.Group():
            gr.Markdown("### Audio Output")
            
            with gr.Row():
                audio_output = gr.Audio(
                    label="Generated Audio",
                    type="numpy",
                    interactive=False,
                )
                
                status_output_gen = gr.Textbox(
                    label="Generation Status",
                    interactive=False,
                    lines=6,
                )
            
            # Connect synthesis button
            synthesize_button.click(
                fn=synthesize_speech,
                inputs=[text_input, language_select, use_streaming, use_voice_clone],
                outputs=[audio_output, status_output_gen],
            )
        
        # Info Section
        with gr.Group():
            gr.Markdown("""
            ### ℹ️ Information
            
            **Voice Cloning:**
            - To enable voice cloning, place `ref_audio.wav` and `ref_text.txt` in `assets/voice_cloning/`
            - `ref_audio.wav`: Reference audio file in WAV format
            - `ref_text.txt`: Transcription of the reference audio
            
            **API Endpoints:**
            - `GET /v1/health` - Health check
            - `GET /v1/models` - List available models
            - `POST /v1/audio/speech` - Generate speech
            - `POST /v1/audio/speech/stream` - Stream speech
            
            **Documentation:** Visit `http://localhost:8000/docs` for API documentation
            """)
    
    return demo


if __name__ == "__main__":
    demo = create_demo()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
    )