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 validate_text(self, text: str, language: str = "English") -> dict:
        """Validate text for TTS."""
        try:
            response = requests.post(
                f"{self.api_url}/v1/text/validate",
                json={"text": text, "language": language},
                timeout=5,
            )
            return response.json() if response.status_code == 200 else {"valid": False, "error": response.text}
        except Exception as e:
            return {"valid": False, "error": str(e)}
    
    def get_usage(self) -> dict:
        """Get API usage statistics."""
        try:
            response = requests.get(f"{self.api_url}/v1/usage", timeout=5)
            return response.json() if response.status_code == 200 else {"error": "Unable to fetch usage"}
        except Exception as e:
            return {"error": str(e)}
    
    def get_quota(self) -> dict:
        """Get API quota and rate limits."""
        try:
            response = requests.get(f"{self.api_url}/v1/quota", timeout=5)
            return response.json() if response.status_code == 200 else {"error": "Unable to fetch quota"}
        except Exception as e:
            return {"error": str(e)}
    
    def list_voices(self) -> dict:
        """List all custom voices."""
        try:
            response = requests.get(f"{self.api_url}/v1/voices", timeout=5)
            return response.json() if response.status_code == 200 else {"data": []}
        except Exception as e:
            return {"data": [], "error": str(e)}
    
    def delete_voice(self, voice_id: str) -> dict:
        """Delete a custom voice."""
        try:
            response = requests.delete(f"{self.api_url}/v1/voices/{voice_id}", timeout=5)
            return response.json() if response.status_code == 200 else {"error": "Unable to delete voice"}
        except Exception as e:
            return {"error": str(e)}
    
    def get_model_config(self, model_id: str) -> dict:
        """Get model configuration."""
        try:
            response = requests.get(
                f"{self.api_url}/v1/models/{model_id}/config",
                timeout=5,
            )
            return response.json() if response.status_code == 200 else {"error": "Unable to fetch config"}
        except Exception as e:
            return {"error": str(e)}
    
    def get_model_languages(self, model_id: str) -> dict:
        """Get supported languages for a model."""
        try:
            response = requests.get(
                f"{self.api_url}/v1/models/{model_id}/languages",
                timeout=5,
            )
            return response.json() if response.status_code == 200 else {"languages": []}
        except Exception as e:
            return {"languages": [], "error": str(e)}
    
    def create_batch(self, items: list) -> dict:
        """Create a batch job."""
        try:
            response = requests.post(
                f"{self.api_url}/v1/batch/create",
                json={"items": items},
                timeout=5,
            )
            return response.json() if response.status_code == 200 else {"error": "Unable to create batch"}
        except Exception as e:
            return {"error": str(e)}
    
    def get_batch_status(self, job_id: str) -> dict:
        """Get batch job status."""
        try:
            response = requests.get(f"{self.api_url}/v1/batch/{job_id}", timeout=5)
            return response.json() if response.status_code == 200 else {"error": "Unable to fetch status"}
        except Exception as e:
            return {"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 validate_text_fn(text: str, language: str) -> str:
    """Validate text and return analysis."""
    if not text.strip():
        return "❌ Error: Please enter some text"
    
    validation = client.validate_text(text, language)
    
    if not validation.get("valid"):
        return f"❌ Invalid: {validation.get('error', 'Unknown error')}"
    
    result = f"""✅ Valid for TTS

📊 Analysis:
- Characters: {validation.get('character_count', 0)}
- Estimated Duration: {validation.get('estimated_duration', 0):.2f}s
- Language: {validation.get('language', language)}
"""
    
    if validation.get('warnings'):
        result += f"\n⚠️ Warnings:\n"
        for warning in validation['warnings']:
            result += f"  • {warning}\n"
    
    return result


def create_batch_fn(batch_text: str) -> str:
    """Create a batch job from multiple lines of text."""
    if not batch_text.strip():
        return "❌ Error: Please enter at least one line of text"
    
    lines = [line.strip() for line in batch_text.split('\n') if line.strip()]
    
    if len(lines) > 100:
        return f"❌ Error: Maximum 100 items per batch (you have {len(lines)})"
    
    items = [{"text": line} for line in lines]
    result = client.create_batch(items)
    
    if "error" in result:
        return f"❌ Error: {result['error']}"
    
    return f"""✅ Batch Job Created

📋 Job ID: {result.get('id', 'Unknown')}
Status: {result.get('status', 'Unknown')}
Items: {len(lines)}

Created: {result.get('created_at', 'Unknown')}

Check status using the job ID above.
"""


def get_batch_status_fn(job_id: str) -> str:
    """Get batch job status."""
    if not job_id.strip():
        return "❌ Error: Please enter a job ID"
    
    result = client.get_batch_status(job_id)
    
    if "error" in result:
        return f"❌ Error: {result['error']}"
    
    counts = result.get('request_counts', {})
    
    return f"""📋 Batch Job Status

Job ID: {result.get('id', 'Unknown')}
Status: {result.get('status', 'Unknown').upper()}

Progress:
  • Total: {counts.get('total', 0)}
  • Completed: {counts.get('completed', 0)}
  • Processing: {counts.get('processing', 0)}
  • Failed: {counts.get('failed', 0)}

Created: {result.get('created_at', 'Unknown')}
Updated: {result.get('updated_at', 'Unknown')}
"""


def get_usage_fn() -> str:
    """Get API usage statistics."""
    usage = client.get_usage()
    
    if "error" in usage:
        return f"❌ Error: {usage['error']}"
    
    result = f"""📊 API Usage Statistics

Total Requests: {usage.get('requests_made', 0)}
Audio Generated: {usage.get('audio_generated_minutes', 0):.2f} minutes

Requests by Model:
"""
    for model, count in usage.get('requests_by_model', {}).items():
        result += f"  • {model}: {count}\n"
    
    result += "\nRequests by Language:\n"
    for lang, count in usage.get('requests_by_language', {}).items():
        result += f"  • {lang}: {count}\n"
    
    return result


def get_quota_fn() -> str:
    """Get API quota information."""
    quota = client.get_quota()
    
    if "error" in quota:
        return f"❌ Error: {quota['error']}"
    
    return f"""📈 API Quota & Limits

Rate Limits:
  • Requests/minute: {quota.get('requests_per_minute', 'Unknown')}
  • Remaining this minute: {quota.get('remaining_requests', 'Unknown')}
  • Concurrent requests: {quota.get('concurrent_requests', 'Unknown')}

Restrictions:
  • Max text length: {quota.get('max_text_length', 'Unknown')} characters
  • Max batch size: {quota.get('max_batch_size', 'Unknown')} items
"""


def list_voices_fn() -> str:
    """List all custom voices."""
    result = client.list_voices()
    
    if "error" in result:
        return f"❌ Error: {result['error']}"
    
    voices = result.get('data', [])
    
    if not voices:
        return "No custom voices found.\n\nUpload a voice using the API endpoint:\nPOST /v1/voices/upload"
    
    output = f"🎭 Custom Voices ({len(voices)})\n\n"
    for voice in voices:
        output += f"Name: {voice.get('name', 'Unknown')}\n"
        output += f"  ID: {voice.get('id', 'Unknown')}\n"
        output += f"  Language: {voice.get('language', 'Unknown')}\n"
        output += f"  Created: {voice.get('created_at', 'Unknown')}\n\n"
    
    return output


def delete_voice_fn(voice_id: str) -> str:
    """Delete a custom voice."""
    if not voice_id.strip():
        return "❌ Error: Please enter a voice ID"
    
    result = client.delete_voice(voice_id)
    
    if "error" in result:
        return f"❌ Error: {result['error']}"
    
    return f"✅ Voice deleted: {voice_id}"


def get_model_info_fn(model_id: str) -> str:
    """Get model configuration and languages."""
    config = client.get_model_config(model_id)
    languages = client.get_model_languages(model_id)
    
    if "error" in config:
        return f"❌ Error: {config['error']}"
    
    result = f"""🤖 Model Configuration

Model: {config.get('id', 'Unknown')}
Size: {config.get('size', 'Unknown')}
Tokenizer: {config.get('tokenizer_type', 'Unknown')}

Max Text Length: {config.get('max_text_length', 'Unknown')} characters
Output Sample Rate: {config.get('output_sample_rate', 'Unknown')} Hz

Features:
  • Streaming: {'✅' if config.get('streaming_supported') else '❌'}
  • Voice Cloning: {'✅' if config.get('voice_cloning_supported') else '❌'}

Recommended Parameters:
"""
    
    for param, value in config.get('recommended_parameters', {}).items():
        result += f"  • {param}: {value}\n"
    
    result += "\nSupported Languages:\n"
    for lang in config.get('languages', []):
        result += f"  • {lang}\n"
    
    if languages.get('languages'):
        result += "\nLanguage Details:\n"
        for lang in languages['languages']:
            result += f"  • {lang.get('name')} ({lang.get('code')})\n"
    
    return result


def create_demo():
    """Create and return the Gradio demo interface."""
    
    with gr.Blocks(
        title="Qwen3-TTS API Demo",
        theme='Nymbo/Nymbo_Theme',
    ) as demo:
        gr.Markdown("""
        # 🎙️ Qwen3-TTS OpenAI-like API Demo
        
        Comprehensive demo showcasing the full **Qwen3-TTS API** with:
        - 🎵 Real-time streaming audio generation
        - 🎭 Voice cloning and management
        - 📊 Batch processing & usage tracking
        - 🌍 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")
            with gr.Row():
                status_button = gr.Button("Check API Status", variant="primary", scale=1)
                status_output = gr.Textbox(
                    label="Status",
                    interactive=False,
                    scale=4,
                )
            status_button.click(
                fn=check_api_status,
                outputs=status_output,
            )
        
        # Tab interface for different features
        with gr.Tabs():
            
            # Tab 1: Text-to-Speech
            with gr.TabItem("🎵 Text-to-Speech"):
                with gr.Group():
                    gr.Markdown("### Convert text to speech with optional streaming and voice cloning")
                    
                    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)",
                            )
                    
                    synthesize_button = gr.Button(
                        "Generate Speech 🎵",
                        variant="primary",
                        size="lg",
                    )
                    
                    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,
                        )
                    
                    synthesize_button.click(
                        fn=synthesize_speech,
                        inputs=[text_input, language_select, use_streaming, use_voice_clone],
                        outputs=[audio_output, status_output_gen],
                    )
            
            # Tab 2: Text Validation
            with gr.TabItem("✅ Text Validation"):
                with gr.Group():
                    gr.Markdown("### Validate text and get synthesis estimates")
                    
                    validate_text_input = gr.Textbox(
                        label="Text to Validate",
                        placeholder="Enter text to check...",
                        lines=4,
                    )
                    
                    validate_language = gr.Dropdown(
                        choices=["English", "Russian", "Chinese", "Japanese", "Korean"],
                        value="English",
                        label="Language",
                    )
                    
                    validate_button = gr.Button("Validate Text", variant="primary")
                    validate_output = gr.Textbox(
                        label="Validation Result",
                        interactive=False,
                        lines=8,
                    )
                    
                    validate_button.click(
                        fn=validate_text_fn,
                        inputs=[validate_text_input, validate_language],
                        outputs=validate_output,
                    )
            
            # Tab 3: Batch Processing
            with gr.TabItem("📦 Batch Processing"):
                with gr.Group():
                    gr.Markdown("### Submit multiple texts for batch processing")
                    
                    with gr.Row():
                        with gr.Column():
                            batch_text_input = gr.Textbox(
                                label="Texts (one per line)",
                                placeholder="Line 1: First text\nLine 2: Second text\nLine 3: Third text\n...",
                                lines=6,
                            )
                            create_batch_button = gr.Button("Create Batch Job", variant="primary")
                            batch_create_output = gr.Textbox(
                                label="Creation Result",
                                interactive=False,
                                lines=8,
                            )
                        
                        with gr.Column():
                            batch_id_input = gr.Textbox(
                                label="Job ID",
                                placeholder="Enter batch job ID to check status...",
                            )
                            check_batch_button = gr.Button("Check Batch Status", variant="primary")
                            batch_status_output = gr.Textbox(
                                label="Batch Status",
                                interactive=False,
                                lines=8,
                            )
                    
                    create_batch_button.click(
                        fn=create_batch_fn,
                        inputs=batch_text_input,
                        outputs=batch_create_output,
                    )
                    
                    check_batch_button.click(
                        fn=get_batch_status_fn,
                        inputs=batch_id_input,
                        outputs=batch_status_output,
                    )
            
            # Tab 4: Voice Management
            with gr.TabItem("🎭 Voice Management"):
                with gr.Group():
                    gr.Markdown("### Manage custom voices for voice cloning")
                    
                    with gr.Row():
                        with gr.Column():
                            list_voices_button = gr.Button("List All Voices", variant="primary")
                            voices_output = gr.Textbox(
                                label="Custom Voices",
                                interactive=False,
                                lines=10,
                            )
                        
                        with gr.Column():
                            delete_voice_id = gr.Textbox(
                                label="Voice ID to Delete",
                                placeholder="Enter voice ID...",
                            )
                            delete_voice_button = gr.Button("Delete Voice", variant="stop")
                            delete_voice_output = gr.Textbox(
                                label="Delete Result",
                                interactive=False,
                            )
                    
                    gr.Markdown("""
                    **To upload a custom voice, use the API endpoint:**
                    ```
                    POST /v1/voices/upload
                    - name: Voice name
                    - language: Language code
                    - audio: WAV audio file (5-30s)
                    - ref_text: Transcription of audio
                    ```
                    """)
                    
                    list_voices_button.click(
                        fn=list_voices_fn,
                        outputs=voices_output,
                    )
                    
                    delete_voice_button.click(
                        fn=delete_voice_fn,
                        inputs=delete_voice_id,
                        outputs=delete_voice_output,
                    )
            
            # Tab 5: Model Information
            with gr.TabItem("🤖 Model Info"):
                with gr.Group():
                    gr.Markdown("### View model configuration and capabilities")
                    
                    model_select = gr.Dropdown(
                        choices=["Qwen/Qwen3-TTS-12Hz-1.7B-Base"],
                        value="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
                        label="Model",
                    )
                    
                    model_info_button = gr.Button("Get Model Info", variant="primary")
                    model_info_output = gr.Textbox(
                        label="Model Configuration",
                        interactive=False,
                        lines=15,
                    )
                    
                    model_info_button.click(
                        fn=get_model_info_fn,
                        inputs=model_select,
                        outputs=model_info_output,
                    )
            
            # Tab 6: Usage & Quota
            with gr.TabItem("📊 Usage & Quota"):
                with gr.Group():
                    gr.Markdown("### Monitor API usage and rate limits")
                    
                    with gr.Row():
                        usage_button = gr.Button("Get Usage Stats", variant="primary")
                        quota_button = gr.Button("Get Quota Info", variant="primary")
                    
                    with gr.Row():
                        usage_output = gr.Textbox(
                            label="Usage Statistics",
                            interactive=False,
                            lines=12,
                        )
                        
                        quota_output = gr.Textbox(
                            label="Quota & Rate Limits",
                            interactive=False,
                            lines=12,
                        )
                    
                    usage_button.click(
                        fn=get_usage_fn,
                        outputs=usage_output,
                    )
                    
                    quota_button.click(
                        fn=get_quota_fn,
                        outputs=quota_output,
                    )
        
        # Footer with links
        with gr.Group():
            gr.Markdown("""
            ---
            
            ### 📚 Documentation & Resources
            
            - **Interactive API Docs**: [Swagger UI](http://localhost:8000/docs)
            - **Extended Endpoints**: [Documentation](https://github.com/Qwen/Qwen3-TTS/blob/main/api/EXTENDED_ENDPOINTS.md)
            - **Code Examples**: [Python Examples](https://github.com/Qwen/Qwen3-TTS/blob/main/api/examples.py)
            - **API Quick Reference**: [Quick Reference](https://github.com/Qwen/Qwen3-TTS/blob/main/API_QUICK_REFERENCE.md)
            """)
    
    return demo


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