yo

.gitignore

@@ -0,0 +1,77 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Environment
+.env
+.venv
+env/
+venv/
+ENV/
+
+# Model & data files
+models/
+tokenizers/
+voice_memories/
+voice_samples/
+cloned_voices/
+*.pt
+*.pth
+*.bin
+*.ckpt
+
+# Audio files
+*.wav
+*.mp3
+*.flac
+*.opus
+*.aac
+
+# Logs
+logs/
+*.log
+log.txt
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+.DS_Store
+
+# Temp files
+tmp/
+temp/
+.temp/
+.~*
+
+# Project specific
+.cache/
+.neptune/
+MANIFEST.in
+.history/
+.mypy_cache/
+.pytest_cache/
+__pycache__/
+.ipynb_checkpoints/
+
+# Ignore Hugging Face credentials
+.huggingface/

Dockerfile

@@ -0,0 +1,92 @@
+# Multi-stage build for authenticated model downloads
+FROM python:3.10-slim AS model-downloader
+# Install huggingface-cli
+RUN pip install huggingface_hub
+# Set working directory
+WORKDIR /model-downloader
+# Create directory for downloaded models
+RUN mkdir -p /model-downloader/models
+# This will run when building the image
+# You'll need to pass your Hugging Face token at build time
+ARG HF_TOKEN
+ENV HF_TOKEN=${HF_TOKEN}
+# Login and download model
+RUN if [ -n "$HF_TOKEN" ]; then \
+    huggingface-cli login --token ${HF_TOKEN}; \
+    huggingface-cli download sesame/csm-1b ckpt.pt --local-dir /model-downloader/models; \
+    else echo "No HF_TOKEN provided, model download will be skipped"; fi
+
+# Now for the main application stage
+FROM nvidia/cuda:12.4.0-base-ubuntu22.04
+# Set environment variables
+ENV PYTHONFAULTHANDLER=1 \
+    PYTHONUNBUFFERED=1 \
+    PYTHONHASHSEED=random \
+    PIP_NO_CACHE_DIR=1 \
+    PIP_DISABLE_PIP_VERSION_CHECK=1 \
+    PIP_DEFAULT_TIMEOUT=100 \
+    NVIDIA_VISIBLE_DEVICES=all \
+    NVIDIA_DRIVER_CAPABILITIES=compute,utility \
+    TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6" \
+    TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3 \
+    python3-pip \
+    python3-dev \
+    ffmpeg \
+    git \
+    build-essential \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set working directory
+WORKDIR /app
+
+# Copy requirements first for better caching
+COPY requirements.txt .
+
+# Create and set up persistent directories with proper permissions
+RUN mkdir -p /app/static /app/models /app/voice_memories /app/voice_references \
+    /app/voice_profiles /app/cloned_voices /app/audio_cache /app/tokenizers /app/logs && \
+    chmod -R 777 /app/voice_references /app/voice_profiles /app/voice_memories \
+    /app/cloned_voices /app/audio_cache /app/static /app/logs /app/tokenizers /app/models
+
+# Copy static files
+COPY ./static /app/static
+
+# Install Python dependencies
+RUN pip3 install --no-cache-dir --upgrade pip && \
+    pip3 install torch torchaudio numpy
+
+# Install torchao from source
+RUN pip3 install git+https://github.com/pytorch/ao.git
+
+# Install torchtune from source with specific branch for latest features
+RUN git clone https://github.com/pytorch/torchtune.git /tmp/torchtune && \
+    cd /tmp/torchtune && \
+    # Try to use the main branch, which should have llama3_2
+    git checkout main && \
+    pip install -e .
+
+# Install remaining dependencies
+RUN pip3 install -r requirements.txt
+
+# Install additional dependencies for streaming and voice cloning
+RUN pip3 install yt-dlp openai-whisper
+
+# Copy application code
+COPY ./app /app/app
+
+# Copy downloaded model from the model-downloader stage
+COPY --from=model-downloader /model-downloader/models /app/models
+
+# Show available models in torchtune
+RUN python3 -c "import torchtune.models; print('Available models in torchtune:', dir(torchtune.models))"
+
+# Expose port
+EXPOSE 8000
+
+# Command to run the application
+CMD ["python3", "-m", "app.main"]

README.md

@@ -1,10 +1,475 @@
 ---
-title: Sesame Openai
-emoji: ⚡
-colorFrom: red
-colorTo: gray
+title: CSM-1B TTS Interface  # Choose a descriptive title
+emoji: 🔊                 # Choose an emoji (e.g., 🎤, 🔊, ✨)
+colorFrom: blue           # Optional: Start color for card gradient
+colorTo: indigo         # Optional: End color for card gradient
 sdk: docker
+sdk_version: "28.0.1"     # <-- IMPORTANT: Check your requirements.txt for the exact gradio version you installed! Update this value.
+app_file: app/main.py   # <-- IMPORTANT: Use the EXACT filename of your main Gradio script (the one with `demo.launch()`)
 pinned: false
+# Optional: Add other configurations like python_version if needed
+# python_version: "3.10"
+# Optional: Specify hardware if needed (e.g., for GPU)
+# hardware: cpu-upgrade # or gpu-small, gpu-a10g-small etc. Check HF pricing/docs
+# Optional: Specify secrets needed (like HF_TOKEN if model download needs it)
+# secrets:
+#   - HF_TOKEN
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+
+
+
+# CSM-1B TTS API
+
+An OpenAI-compatible Text-to-Speech API that harnesses the power of Sesame's Conversational Speech Model (CSM-1B). This API allows you to generate high-quality speech from text using a variety of consistent voices, compatible with systems like OpenWebUI, ChatBot UI, and any platform that supports the OpenAI TTS API format.
+
+## Features
+
+- **OpenAI API Compatibility**: Drop-in replacement for OpenAI's TTS API
+- **Multiple Voices**: Six distinct voices (alloy, echo, fable, onyx, nova, shimmer)
+- **Voice Consistency**: Maintains consistent voice characteristics across multiple requests
+- **Voice Cloning**: Clone your own voice from audio samples
+- **Conversational Context**: Supports conversational context for improved naturalness
+- **Multiple Audio Formats**: Supports MP3, OPUS, AAC, FLAC, and WAV
+- **Speed Control**: Adjustable speech speed
+- **CUDA Acceleration**: GPU support for faster generation
+- **Web UI**: Simple interface for voice cloning and speech generation
+
+## Getting Started
+
+### Prerequisites
+
+- Docker and Docker Compose
+- NVIDIA GPU with CUDA support (recommended)
+- Hugging Face account with access to `sesame/csm-1b` model
+
+### Installation
+
+1. Clone this repository:
+```bash
+git clone https://github.com/phildougherty/sesame_csm_openai
+cd sesame_csm_openai
+```
+
+2. Create a `.env` file in the /app folder with your Hugging Face token:
+```
+HF_TOKEN=your_hugging_face_token_here
+```
+
+3. Build and start the container:
+```bash
+docker compose up -d --build
+```
+
+The server will start on port 8000. First startup may take some time as it downloads the model files.
+
+## Hugging Face Configuration (ONLY NEEDED TO ACCEPT TERMS/DOWNLOAD MODEL)
+
+This API requires access to the `sesame/csm-1b` model on Hugging Face:
+
+1. Create a Hugging Face account if you don't have one: [https://huggingface.co/join](https://huggingface.co/join)
+2. Accept the model license at [https://huggingface.co/sesame/csm-1b](https://huggingface.co/sesame/csm-1b)
+3. Generate an access token at [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)
+4. Use this token in your `.env` file or pass it directly when building the container:
+
+```bash
+HF_TOKEN=your_token docker compose up -d --build
+```
+
+### Required Models
+
+The API uses the following models which are downloaded automatically:
+
+- **CSM-1B**: The main speech generation model from Sesame
+- **Mimi**: Audio codec for high-quality audio generation
+- **Llama Tokenizer**: Uses the unsloth/Llama-3.2-1B tokenizer for text processing
+
+## Multi-GPU Support
+
+The CSM-1B model can be distributed across multiple GPUs to handle larger models or improve performance. To enable multi-GPU support, set the `CSM_DEVICE_MAP` environment variable:
+
+```bash
+# Automatic device mapping (recommended)
+CSM_DEVICE_MAP=auto docker compose up -d
+
+# Balanced distribution of layers across GPUs
+CSM_DEVICE_MAP=balanced docker compose up -d
+
+# Sequential distribution (backbone on first GPUs, decoder on remaining)
+CSM_DEVICE_MAP=sequential docker compose up -d
+
+## Voice Cloning Guide
+
+The CSM-1B TTS API comes with powerful voice cloning capabilities that allow you to create custom voices from audio samples. Here's how to use this feature:
+
+### Method 1: Using the Web Interface
+
+1. Access the voice cloning UI by navigating to `http://your-server-ip:8000/voice-cloning` in your browser.
+
+2. **Clone a Voice**:
+   - Go to the "Clone Voice" tab
+   - Enter a name for your voice
+   - Upload an audio sample (2-3 minutes of clear speech works best)
+   - Optionally provide a transcript of the audio for better results
+   - Click "Clone Voice"
+
+3. **View Your Voices**:
+   - Navigate to the "My Voices" tab to see all your cloned voices
+   - You can preview or delete voices from this tab
+
+4. **Generate Speech**:
+   - Go to the "Generate Speech" tab
+   - Select one of your cloned voices
+   - Enter the text you want to synthesize
+   - Adjust the temperature slider if needed (lower for more consistent results)
+   - Click "Generate Speech" and listen to the result
+
+### Method 2: Using the API
+
+1. **Clone a Voice**:
+```bash
+curl -X POST http://localhost:8000/v1/voice-cloning/clone \
+  -F "name=My Voice" \
+  -F "audio_file=@path/to/your/voice_sample.mp3" \
+  -F "transcript=Optional transcript of the audio sample" \
+  -F "description=A description of this voice"
+```
+
+2. **List Available Cloned Voices**:
+```bash
+curl -X GET http://localhost:8000/v1/voice-cloning/voices
+```
+
+3. **Generate Speech with a Cloned Voice**:
+```bash
+curl -X POST http://localhost:8000/v1/voice-cloning/generate \
+  -H "Content-Type: application/json" \
+  -d '{
+    "voice_id": "1234567890_my_voice",
+    "text": "This is my cloned voice speaking.",
+    "temperature": 0.7
+  }' \
+  --output cloned_speech.mp3
+```
+
+4. **Generate a Voice Preview**:
+```bash
+curl -X POST http://localhost:8000/v1/voice-cloning/voices/1234567890_my_voice/preview \
+  --output voice_preview.mp3
+```
+
+5. **Delete a Cloned Voice**:
+```bash
+curl -X DELETE http://localhost:8000/v1/voice-cloning/voices/1234567890_my_voice
+```
+
+### Voice Cloning Best Practices
+
+For the best voice cloning results:
+
+1. **Use High-Quality Audio**: Record in a quiet environment with minimal background noise and echo.
+
+2. **Provide Sufficient Length**: 2-3 minutes of speech provides better results than shorter samples.
+
+3. **Clear, Natural Speech**: Speak naturally at a moderate pace with clear pronunciation.
+
+4. **Include Various Intonations**: Sample should contain different sentence types (statements, questions) for better expressiveness.
+
+5. **Add a Transcript**: While optional, providing an accurate transcript of your recording helps the model better capture your voice characteristics.
+
+6. **Adjust Temperature**: For more consistent results, use lower temperature values (0.6-0.7). For more expressiveness, use higher values (0.7-0.9).
+
+7. **Try Multiple Samples**: If you're not satisfied with the results, try recording a different sample or adjusting the speaking style.
+
+### Using Cloned Voices with the Standard TTS Endpoint
+
+Cloned voices are automatically available through the standard OpenAI-compatible endpoint. Simply use the voice ID or name as the `voice` parameter:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "csm-1b",
+    "input": "This is my cloned voice speaking through the standard endpoint.",
+    "voice": "1234567890_my_voice",
+    "response_format": "mp3"
+  }' \
+  --output cloned_speech.mp3
+```
+
+## YouTube Voice Cloning 
+
+The CSM-1B TTS API now includes the ability to clone voices directly from YouTube videos. This feature allows you to extract voice characteristics from any YouTube content and create custom TTS voices without needing to download or prepare audio samples yourself.
+
+## How to Clone a Voice from YouTube
+
+### API Endpoint
+
+```
+POST /v1/audio/speech/voice-cloning/youtube
+```
+
+Parameters:
+- `youtube_url`: URL of the YouTube video
+- `voice_name`: Name for the cloned voice
+- `start_time` (optional): Start time in seconds (default: 0)
+- `duration` (optional): Duration to extract in seconds (default: 180)
+- `description` (optional): Description of the voice
+
+Example request:
+```json
+{
+  "youtube_url": "https://www.youtube.com/watch?v=dQw4w9WgXcQ",
+  "voice_name": "rick_astley",
+  "start_time": 30,
+  "duration": 60,
+  "description": "Never gonna give you up"
+}
+```
+
+Response:
+```json
+{
+  "voice_id": "1710805983_rick_astley",
+  "name": "rick_astley",
+  "description": "Never gonna give you up",
+  "created_at": "2025-03-18T22:53:03Z",
+  "audio_duration": 60.0,
+  "sample_count": 1440000
+}
+```
+
+## How It Works
+
+1. The system downloads the audio from the specified YouTube video
+2. It extracts the specified segment (start time and duration)
+3. Whisper ASR generates a transcript of the audio for better voice matching
+4. The audio is processed to remove noise and silence
+5. The voice is cloned and made available for TTS generation
+
+## Best Practices for YouTube Voice Cloning
+
+For optimal results:
+
+1. **Choose Clear Speech Segments**
+   - Select portions of the video with clear, uninterrupted speech
+   - Avoid segments with background music, sound effects, or multiple speakers
+
+2. **Optimal Duration**
+   - 30-60 seconds of clean speech typically provides the best results
+   - Longer isn't always better - quality matters more than quantity
+
+3. **Specify Time Ranges Precisely**
+   - Use `start_time` and `duration` to target the exact speech segment
+   - Preview the segment in YouTube before cloning to ensure it's suitable
+
+4. **Consider Audio Quality**
+   - Higher quality videos generally produce better voice clones
+   - Interviews, vlogs, and speeches often work better than highly produced content
+
+## Limitations
+
+- YouTube videos with heavy background music may result in lower quality voice clones
+- Very noisy or low-quality audio sources will produce less accurate voice clones
+- The system works best with natural speech rather than singing or exaggerated voices
+- Copyright restrictions apply - only clone voices you have permission to use
+
+## Example Use Cases
+
+- Create a voice clone of a public figure for educational content
+- Clone your own YouTube voice for consistent TTS across your applications
+- Create voice clones from historical speeches or interviews (public domain)
+- Develop custom voices for creative projects with proper permissions
+
+## Ethical Considerations
+
+Please use YouTube voice cloning responsibly:
+- Only clone voices from content you have permission to use
+- Respect copyright and intellectual property rights
+- Clearly disclose when using AI-generated or cloned voices
+- Do not use cloned voices for impersonation, deception, or harmful content
+
+## How the Voices Work
+
+Unlike traditional TTS systems with pre-trained voice models, CSM-1B works differently:
+
+- The base CSM-1B model is capable of producing a wide variety of voices but doesn't have fixed voice identities
+- This API creates consistent voices by using acoustic "seed" samples for each named voice
+- When you specify a voice (e.g., "alloy"), the API uses a consistent acoustic seed and speaker ID
+- The most recent generated audio becomes the new reference for that voice, maintaining voice consistency
+- Each voice has unique tonal qualities:
+  - **alloy**: Balanced mid-tones with natural inflection
+  - **echo**: Resonant with slight reverberance
+  - **fable**: Brighter with higher pitch
+  - **onyx**: Deep and resonant
+  - **nova**: Warm and smooth
+  - **shimmer**: Light and airy with higher frequencies
+
+The voice system can be extended with your own voice samples by using the voice cloning feature.
+
+## API Usage
+
+### Basic Usage
+
+Generate speech with a POST request to `/v1/audio/speech`:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "csm-1b",
+    "input": "Hello, this is a test of the CSM text to speech system.",
+    "voice": "alloy",
+    "response_format": "mp3"
+  }' \
+  --output speech.mp3
+```
+
+### Available Endpoints
+
+#### Standard TTS Endpoints
+- `GET /v1/audio/models` - List available models
+- `GET /v1/audio/voices` - List available voices (including cloned voices)
+- `GET /v1/audio/speech/response-formats` - List available response formats
+- `POST /v1/audio/speech` - Generate speech from text
+- `POST /api/v1/audio/conversation` - Advanced endpoint for conversational speech
+
+#### Voice Cloning Endpoints
+- `POST /v1/voice-cloning/clone` - Clone a new voice from an audio sample
+- `GET /v1/voice-cloning/voices` - List all cloned voices
+- `POST /v1/voice-cloning/generate` - Generate speech with a cloned voice
+- `POST /v1/voice-cloning/voices/{voice_id}/preview` - Generate a preview of a cloned voice
+- `DELETE /v1/voice-cloning/voices/{voice_id}` - Delete a cloned voice
+
+### Request Parameters
+
+#### Standard TTS
+| Parameter | Description | Type | Default |
+|-----------|-------------|------|---------|
+| `model` | Model ID to use | string | "csm-1b" |
+| `input` | The text to convert to speech | string | Required |
+| `voice` | The voice to use (standard or cloned voice ID) | string | "alloy" |
+| `response_format` | Audio format | string | "mp3" |
+| `speed` | Speech speed multiplier | float | 1.0 |
+| `temperature` | Sampling temperature | float | 0.8 |
+| `max_audio_length_ms` | Maximum audio length in ms | integer | 90000 |
+
+#### Voice Cloning
+| Parameter | Description | Type | Default |
+|-----------|-------------|------|---------|
+| `name` | Name for the cloned voice | string | Required |
+| `audio_file` | Audio sample file | file | Required |
+| `transcript` | Transcript of the audio | string | Optional |
+| `description` | Description of the voice | string | Optional |
+
+### Available Voices
+
+- `alloy` - Balanced and natural
+- `echo` - Resonant
+- `fable` - Bright and higher-pitched
+- `onyx` - Deep and resonant
+- `nova` - Warm and smooth
+- `shimmer` - Light and airy
+- `[cloned voice ID]` - Any voice you've cloned using the voice cloning feature
+
+### Response Formats
+
+- `mp3` - MP3 audio format
+- `opus` - Opus audio format
+- `aac` - AAC audio format
+- `flac` - FLAC audio format
+- `wav` - WAV audio format
+
+## Integration with OpenWebUI
+
+OpenWebUI is a popular open-source UI for AI models that supports custom TTS endpoints. Here's how to integrate the CSM-1B TTS API:
+
+1. Access your OpenWebUI settings
+2. Navigate to the TTS settings section
+3. Select "Custom TTS Endpoint"
+4. Enter your CSM-1B TTS API URL: `http://your-server-ip:8000/v1/audio/speech`
+5. Use the API Key field to add any authentication if you've configured it (not required by default)
+6. Test the connection
+7. Save your settings
+
+Once configured, OpenWebUI will use your CSM-1B TTS API for all text-to-speech conversion, producing high-quality speech with the selected voice.
+
+### Using Cloned Voices with OpenWebUI
+
+Your cloned voices will automatically appear in OpenWebUI's voice selector. Simply choose your cloned voice from the dropdown menu in the TTS settings or chat interface.
+
+## Advanced Usage
+
+### Conversational Context
+
+For more natural-sounding speech in a conversation, you can use the conversation endpoint:
+
+```bash
+curl -X POST http://localhost:8000/api/v1/audio/conversation \
+  -H "Content-Type: application/json" \
+  -d '{
+    "text": "Nice to meet you too!",
+    "speaker_id": 0,
+    "context": [
+      {
+        "speaker": 1,
+        "text": "Hello, nice to meet you.",
+        "audio": "BASE64_ENCODED_AUDIO"
+      }
+    ]
+  }' \
+  --output response.wav
+```
+
+This allows the model to take into account the previous utterances for more contextually appropriate speech.
+
+### Model Parameters
+
+For fine-grained control, you can adjust:
+
+- `temperature` (0.0-1.0): Higher values produce more variation but may be less stable
+- `topk` (1-100): Controls diversity of generated speech
+- `max_audio_length_ms`: Maximum length of generated audio in milliseconds
+- `voice_consistency` (0.0-1.0): How strongly to maintain voice characteristics across segments
+
+## Troubleshooting
+
+### API Returns 503 Service Unavailable
+
+- Verify your Hugging Face token has access to `sesame/csm-1b`
+- Check if the model downloaded successfully in the logs
+- Ensure you have enough GPU memory (at least 8GB recommended)
+
+### Audio Quality Issues
+
+- Try different voices - some may work better for your specific text
+- Adjust temperature (lower for more stable output)
+- For longer texts, the API automatically splits into smaller chunks for better quality
+- For cloned voices, try recording a cleaner audio sample
+
+### Voice Cloning Issues
+
+- **Poor Voice Quality**: Try recording in a quieter environment with less background noise
+- **Inconsistent Voice**: Provide a longer and more varied audio sample (2-3 minutes)
+- **Accent Issues**: Make sure your sample contains similar words/sounds to what you'll be generating
+- **Low Volume**: The sample is normalized automatically, but ensure it's not too quiet or distorted
+
+### Voice Inconsistency
+
+- The API maintains voice consistency across separate requests
+- However, very long pauses between requests may result in voice drift
+- For critical applications, consider using the same seed audio
+
+## License
+
+This project is released under the MIT License. The CSM-1B model is subject to its own license terms defined by Sesame.
+
+## Acknowledgments
+
+- [Sesame](https://www.sesame.com) for releasing the CSM-1B model
+- This project is not affiliated with or endorsed by Sesame or OpenAI
+
+---
+
+Happy speech generating!

app/.dockerignore

@@ -0,0 +1,94 @@
+# Version control
+.git
+.gitignore
+.github
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Environment
+.env
+.venv
+env/
+venv/
+ENV/
+
+# Docker related
+Dockerfile*
+docker-compose*
+.docker
+.dockerignore
+
+# Documentation
+README.md
+CHANGELOG.md
+LICENSE
+docs/
+*.md
+examples/
+tests/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+.DS_Store
+
+# Logs and temp files
+logs/
+*.log
+log.txt
+tmp/
+temp/
+.temp/
+.~*
+
+# Model files - we'll download these in the container
+models/
+tokenizers/
+voice_memories/
+voice_samples/
+*.pt
+*.pth
+*.bin
+*.ckpt
+
+# Audio files
+*.wav
+*.mp3
+*.flac
+*.opus
+*.aac
+
+# Project specific
+.cache/
+.neptune/
+MANIFEST.in
+.history/
+.mypy_cache/
+.pytest_cache/
+__pycache__/
+.ipynb_checkpoints/
+
+# Ignore Hugging Face credentials except specified files
+.huggingface/

app/api/init.py

@@ -0,0 +1 @@
+"""API package for CSM-1B."""

app/api/routes.py

@@ -0,0 +1,1048 @@
+"""
+API routes for the CSM-1B TTS API.
+"""
+import os
+import io
+import base64
+import time
+import tempfile
+import logging
+import asyncio
+from enum import Enum
+from typing import Dict, List, Optional, Any, Union
+import torch
+import torchaudio
+import numpy as np
+from fastapi import APIRouter, Request, HTTPException, BackgroundTasks, Body, Response, Query
+from fastapi.responses import StreamingResponse
+from app.api.schemas import SpeechRequest, ResponseFormat, Voice
+from app.models import Segment
+from app.api.streaming import AudioChunker
+from app.prompt_engineering import split_into_segments
+
+# Set up logging
+logger = logging.getLogger(__name__)
+router = APIRouter()
+
+# Mapping of response_format to MIME types
+MIME_TYPES = {
+    "mp3": "audio/mpeg",
+    "opus": "audio/opus",
+    "aac": "audio/aac",
+    "flac": "audio/flac",
+    "wav": "audio/wav",
+}
+
+def get_speaker_id(app_state, voice):
+    """Helper function to get speaker ID from voice name or ID"""
+    if hasattr(app_state, "voice_speaker_map") and voice in app_state.voice_speaker_map:
+        return app_state.voice_speaker_map[voice]
+        
+    # Standard voices mapping
+    voice_to_speaker = {"alloy": 0, "echo": 1, "fable": 2, "onyx": 3, "nova": 4, "shimmer": 5}
+    
+    if voice in voice_to_speaker:
+        return voice_to_speaker[voice]
+    
+    # Try parsing as integer
+    try:
+        speaker_id = int(voice)
+        if 0 <= speaker_id < 6:
+            return speaker_id
+    except (ValueError, TypeError):
+        pass
+    
+    # Check cloned voices if the voice cloner exists
+    if hasattr(app_state, "voice_cloner") and app_state.voice_cloner is not None:
+        # Check by ID
+        if voice in app_state.voice_cloner.cloned_voices:
+            return app_state.voice_cloner.cloned_voices[voice].speaker_id
+            
+        # Check by name
+        for v_id, v_info in app_state.voice_cloner.cloned_voices.items():
+            if v_info.name.lower() == voice.lower():
+                return v_info.speaker_id
+    
+    # Default to alloy
+    return 0
+
+@router.post("/audio/speech", tags=["Audio"], response_class=Response)
+async def generate_speech(
+    request: Request,
+    speech_request: SpeechRequest,
+):
+    """
+    Generate audio of text being spoken by a realistic voice.
+    
+    This endpoint is compatible with the OpenAI TTS API.
+    
+    For streaming responses, use `/v1/audio/speech/streaming` instead.
+    """
+    # Check if model is available
+    if not hasattr(request.app.state, "generator") or request.app.state.generator is None:
+        raise HTTPException(status_code=503, detail="TTS model not available")
+    
+    # Set default values
+    model = speech_request.model
+    voice = speech_request.voice
+    input_text = speech_request.input
+    response_format = speech_request.response_format
+    speed = speech_request.speed
+    temperature = speech_request.temperature
+    max_audio_length_ms = speech_request.max_audio_length_ms
+    
+    # Log request details
+    logger.info(f"TTS request: text length={len(input_text)}, voice={voice}, format={response_format}")
+    
+    try:
+        # Get speaker ID for the voice
+        speaker_id = get_speaker_id(request.app.state, voice)
+        if speaker_id is None:
+            raise HTTPException(status_code=400, detail=f"Voice '{voice}' not found")
+        
+        # Check if this is a cloned voice
+        voice_info = None
+        cloned_voice_id = None
+        
+        if hasattr(request.app.state, "get_voice_info"):
+            voice_info = request.app.state.get_voice_info(voice)
+            if voice_info and voice_info["type"] == "cloned":
+                cloned_voice_id = voice_info["voice_id"]
+                
+        # Generate audio based on whether it's a standard or cloned voice
+        if cloned_voice_id is not None and hasattr(request.app.state, "voice_cloner"):
+            # Generate speech with cloned voice
+            logger.info(f"Generating speech with cloned voice ID: {cloned_voice_id}")
+            try:
+                voice_cloner = request.app.state.voice_cloner
+                audio = voice_cloner.generate_speech(
+                    text=input_text,
+                    voice_id=cloned_voice_id,
+                    temperature=temperature,
+                    topk=speech_request.topk or 30,
+                    max_audio_length_ms=max_audio_length_ms
+                )
+                sample_rate = request.app.state.sample_rate
+                logger.info(f"Generated speech with cloned voice, length: {len(audio)/sample_rate:.2f}s")
+            except Exception as e:
+                logger.error(f"Error generating speech with cloned voice: {e}", exc_info=True)
+                raise HTTPException(
+                    status_code=500, 
+                    detail=f"Failed to generate speech with cloned voice: {str(e)}"
+                )
+        else:
+            # Generate speech with standard voice
+            # Use voice context from memory if enabled
+            if hasattr(request.app.state, "voice_memory_enabled") and request.app.state.voice_memory_enabled:
+                from app.voice_memory import get_voice_context
+                context = get_voice_context(voice, torch.device(request.app.state.device))
+            else:
+                context = []
+            
+            # Apply optional text enhancement for better voice consistency
+            enhanced_text = input_text
+            if hasattr(request.app.state, "prompt_templates"):
+                from app.prompt_engineering import format_text_for_voice
+                enhanced_text = format_text_for_voice(input_text, voice)
+            
+            # Generate audio
+            audio = request.app.state.generator.generate(
+                text=enhanced_text,
+                speaker=speaker_id,
+                context=context,
+                temperature=temperature,
+                topk=speech_request.topk or 50,
+                max_audio_length_ms=max_audio_length_ms
+            )
+            sample_rate = request.app.state.sample_rate
+            
+            # Process audio for better quality
+            if hasattr(request.app.state, "voice_enhancement_enabled") and request.app.state.voice_enhancement_enabled:
+                from app.voice_enhancement import process_generated_audio
+                audio = process_generated_audio(
+                    audio=audio,
+                    voice_name=voice,
+                    sample_rate=sample_rate,
+                    text=input_text
+                )
+            
+            # Update voice memory if enabled
+            if hasattr(request.app.state, "voice_memory_enabled") and request.app.state.voice_memory_enabled:
+                from app.voice_memory import update_voice_memory
+                update_voice_memory(voice, audio, input_text)
+        
+        # Handle speed adjustments if not 1.0
+        if speed != 1.0 and speed > 0:
+            try:
+                # Adjust speed using torchaudio
+                effects = [
+                    ["tempo", str(speed)]
+                ]
+                audio_cpu = audio.cpu()
+                adjusted_audio, _ = torchaudio.sox_effects.apply_effects_tensor(
+                    audio_cpu.unsqueeze(0), 
+                    sample_rate, 
+                    effects
+                )
+                audio = adjusted_audio.squeeze(0)
+                logger.info(f"Adjusted speech speed to {speed}x")
+            except Exception as e:
+                logger.warning(f"Failed to adjust speech speed: {e}")
+        
+        # Format the audio according to the requested format
+        response_data, content_type = await format_audio(
+            audio, 
+            response_format, 
+            sample_rate, 
+            request.app.state
+        )
+        
+        # Create and return the response
+        return Response(
+            content=response_data,
+            media_type=content_type,
+            headers={"Content-Disposition": f"attachment; filename=speech.{response_format}"}
+        )
+                
+    except Exception as e:
+        logger.error(f"Error in text_to_speech: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+
+@router.post("/audio/speech/stream", tags=["Audio"])
+async def stream_speech(request: Request, speech_request: SpeechRequest):
+    """Stream audio in real-time as it's being generated."""
+    # Check if model is loaded
+    if not hasattr(request.app.state, "generator") or request.app.state.generator is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+
+    # Get request parameters
+    input_text = speech_request.input
+    voice = speech_request.voice
+    response_format = speech_request.response_format
+    temperature = speech_request.temperature
+    
+    logger.info(f"Real-time streaming speech from text ({len(input_text)} chars) with voice '{voice}'")
+    
+    # Get speaker ID for the voice
+    speaker_id = get_speaker_id(request.app.state, voice)
+    if speaker_id is None:
+        raise HTTPException(status_code=400, detail=f"Voice '{voice}' not found")
+
+    # Split text into very small segments for incremental generation
+    text_segments = split_into_segments(input_text, max_chars=50)  # Smaller segments for faster first response
+    logger.info(f"Split text into {len(text_segments)} segments")
+
+    # Create media type based on format
+    media_type = {
+        "mp3": "audio/mpeg",
+        "opus": "audio/opus",
+        "aac": "audio/aac",
+        "flac": "audio/flac",
+        "wav": "audio/wav",
+    }.get(response_format, "audio/mpeg")
+    
+    # For streaming, WAV works best
+    streaming_format = "wav"
+    
+    # Set up WAV header for streaming
+    sample_rate = request.app.state.sample_rate
+    
+    async def generate_streaming_audio():
+        # Get context for the voice
+        if hasattr(request.app.state, "voice_cloning_enabled") and request.app.state.voice_cloning_enabled:
+            voice_info = request.app.state.get_voice_info(voice)
+            if voice_info and voice_info["type"] == "cloned":
+                # Use cloned voice context
+                voice_cloner = request.app.state.voice_cloner
+                context = voice_cloner.get_voice_context(voice_info["voice_id"])
+            else:
+                # Standard voice
+                from app.voice_enhancement import get_voice_segments
+                context = get_voice_segments(voice, request.app.state.device)
+        else:
+            # Standard voice
+            from app.voice_enhancement import get_voice_segments
+            context = get_voice_segments(voice, request.app.state.device)
+
+        # Send WAV header immediately
+        if streaming_format == "wav":
+            # Create a WAV header for 16-bit mono audio
+            header = bytes()
+            # RIFF header
+            header += b'RIFF'
+            header += b'\x00\x00\x00\x00'  # Placeholder for file size
+            header += b'WAVE'
+            # Format chunk
+            header += b'fmt '
+            header += (16).to_bytes(4, 'little')  # Format chunk size
+            header += (1).to_bytes(2, 'little')   # PCM format
+            header += (1).to_bytes(2, 'little')   # Mono channel
+            header += (sample_rate).to_bytes(4, 'little')  # Sample rate
+            header += (sample_rate * 2).to_bytes(4, 'little')  # Byte rate
+            header += (2).to_bytes(2, 'little')   # Block align
+            header += (16).to_bytes(2, 'little')  # Bits per sample
+            # Data chunk
+            header += b'data'
+            header += b'\x00\x00\x00\x00'  # Placeholder for data size
+            yield header
+
+        # Process each segment and stream immediately
+        for i, segment_text in enumerate(text_segments):
+            try:
+                logger.info(f"Generating segment {i+1}/{len(text_segments)}")
+                
+                # For cloned voices, use the voice cloner
+                if hasattr(request.app.state, "voice_cloning_enabled") and request.app.state.voice_cloning_enabled:
+                    voice_info = request.app.state.get_voice_info(voice)
+                    if voice_info and voice_info["type"] == "cloned":
+                        # Use cloned voice
+                        voice_cloner = request.app.state.voice_cloner
+                        segment_audio = await asyncio.to_thread(
+                            voice_cloner.generate_speech,
+                            segment_text,
+                            voice_info["voice_id"],
+                            temperature=temperature,
+                            topk=30,
+                            max_audio_length_ms=2000  # Keep it very short for fast generation
+                        )
+                    else:
+                        # Use standard voice with generator
+                        segment_audio = await asyncio.to_thread(
+                            request.app.state.generator.generate,
+                            segment_text,
+                            speaker_id,
+                            context,
+                            max_audio_length_ms=2000,  # Short for quicker generation
+                            temperature=temperature
+                        )
+                else:
+                    # Use standard voice with generator
+                    segment_audio = await asyncio.to_thread(
+                        request.app.state.generator.generate,
+                        segment_text,
+                        speaker_id,
+                        context,
+                        max_audio_length_ms=2000,  # Short for quicker generation
+                        temperature=temperature
+                    )
+                
+                # Skip empty or problematic audio
+                if segment_audio is None or segment_audio.numel() == 0:
+                    logger.warning(f"Empty audio for segment {i+1}")
+                    continue
+                    
+                # Convert to bytes and stream immediately
+                buf = io.BytesIO()
+                audio_to_save = segment_audio.unsqueeze(0) if len(segment_audio.shape) == 1 else segment_audio
+                torchaudio.save(buf, audio_to_save.cpu(), sample_rate, format=streaming_format)
+                buf.seek(0)
+                
+                # For WAV format, skip the header for all segments after the first
+                if streaming_format == "wav" and i > 0:
+                    buf.seek(44)  # Skip WAV header
+                    
+                segment_bytes = buf.read()
+                yield segment_bytes
+                
+                # Update context with this segment for next generation
+                context = [
+                    Segment(
+                        text=segment_text,
+                        speaker=speaker_id,
+                        audio=segment_audio
+                    )
+                ]
+                
+            except Exception as e:
+                logger.error(f"Error generating segment {i+1}: {e}")
+                # Continue to next segment
+        
+    # Return the streaming response
+    return StreamingResponse(
+        generate_streaming_audio(),
+        media_type=media_type,
+        headers={
+            "X-Accel-Buffering": "no",  # Prevent buffering in nginx
+            "Cache-Control": "no-cache, no-store, must-revalidate",
+            "Connection": "keep-alive",
+            "Transfer-Encoding": "chunked"
+        }
+    )
+
+@router.post("/audio/speech/streaming", tags=["Audio"])
+async def openai_stream_speech(
+    request: Request,
+    speech_request: SpeechRequest,
+):
+    """
+    Stream audio in OpenAI-compatible streaming format.
+    
+    This endpoint is compatible with the OpenAI streaming TTS API.
+    """
+    # Use the same logic as the stream_speech endpoint but with a different name
+    # to maintain the OpenAI API naming convention
+    return await stream_speech(request, speech_request)
+
+async def format_audio(audio, response_format, sample_rate, app_state):
+    """
+    Format audio according to requested format.
+    
+    Args:
+        audio: Audio tensor from TTS generation
+        response_format: Format as string or enum ('mp3', 'opus', 'aac', 'flac', 'wav')
+        sample_rate: Sample rate of the audio
+        app_state: FastAPI app state with config and cache settings
+    
+    Returns:
+        Tuple of (response_data, content_type)
+    """
+    import io
+    import torch
+    import torchaudio
+    import tempfile
+    import os
+    import hashlib
+    import time
+    
+    # Handle enum or string for response_format
+    if hasattr(response_format, 'value'):
+        response_format = response_format.value
+    
+    # Normalize response_format to lowercase
+    response_format = str(response_format).lower()
+    
+    # Map formats to content types
+    format_to_content_type = {
+        'mp3': 'audio/mpeg',
+        'opus': 'audio/opus',
+        'aac': 'audio/aac',
+        'flac': 'audio/flac',
+        'wav': 'audio/wav'
+    }
+    
+    # Ensure response format is supported
+    if response_format not in format_to_content_type:
+        logger.warning(f"Unsupported format: {response_format}, defaulting to mp3")
+        response_format = 'mp3'
+    
+    # Generate a cache key based on audio content and format
+    cache_enabled = getattr(app_state, "audio_cache_enabled", False)
+    cache_key = None
+    
+    if cache_enabled:
+        # Generate a hash of the audio tensor for caching
+        audio_hash = hashlib.md5(audio.cpu().numpy().tobytes()).hexdigest()
+        cache_key = f"{audio_hash}_{response_format}"
+        cache_dir = getattr(app_state, "audio_cache_dir", "/app/audio_cache")
+        os.makedirs(cache_dir, exist_ok=True)
+        cache_path = os.path.join(cache_dir, f"{cache_key}")
+        
+        # Check if we have a cache hit
+        if os.path.exists(cache_path):
+            try:
+                with open(cache_path, "rb") as f:
+                    cached_data = f.read()
+                logger.info(f"Cache hit for {response_format} audio")
+                return cached_data, format_to_content_type[response_format]
+            except Exception as e:
+                logger.warning(f"Error reading from cache: {e}")
+    
+    # Process audio to the required format
+    start_time = time.time()
+    
+    # Move audio to CPU before saving
+    audio_cpu = audio.cpu()
+    
+    # Use a temporary file for format conversion
+    with tempfile.NamedTemporaryFile(suffix=f".{response_format}", delete=False) as temp_file:
+        temp_path = temp_file.name
+        try:
+            if response_format == 'wav':
+                # Direct save for WAV
+                torchaudio.save(temp_path, audio_cpu.unsqueeze(0), sample_rate)
+            else:
+                # For other formats, first save as WAV then convert
+                wav_path = f"{temp_path}.wav"
+                torchaudio.save(wav_path, audio_cpu.unsqueeze(0), sample_rate)
+                
+                # Use ffmpeg via torchaudio for conversion
+                if hasattr(torchaudio.backend, 'sox_io_backend'):  # New torchaudio structure
+                    if response_format == 'mp3':
+                        # For MP3, use higher quality
+                        sox_effects = torchaudio.sox_effects.SoxEffectsChain()
+                        sox_effects.set_input_file(wav_path)
+                        sox_effects.append_effect_to_chain(["rate", f"{sample_rate}"])
+                        # Higher bitrate for better quality
+                        sox_effects.append_effect_to_chain(["gain", "-n"])  # Normalize
+                        out, _ = sox_effects.sox_build_flow_effects()
+                        torchaudio.save(temp_path, out, sample_rate, format="mp3", compression=128)
+                    elif response_format == 'opus':
+                        # Use ffmpeg for opus through a system call
+                        import subprocess
+                        subprocess.run([
+                            "ffmpeg", "-i", wav_path, "-c:a", "libopus", 
+                            "-b:a", "64k", "-vbr", "on", temp_path,
+                            "-y", "-loglevel", "error"
+                        ], check=True)
+                    elif response_format == 'aac':
+                        # Use ffmpeg for AAC through a system call
+                        import subprocess
+                        subprocess.run([
+                            "ffmpeg", "-i", wav_path, "-c:a", "aac", 
+                            "-b:a", "128k", temp_path,
+                            "-y", "-loglevel", "error"
+                        ], check=True)
+                    elif response_format == 'flac':
+                        torchaudio.save(temp_path, audio_cpu.unsqueeze(0), sample_rate, format="flac")
+                else:
+                    # Fallback using external command
+                    import subprocess
+                    if response_format == 'mp3':
+                        subprocess.run([
+                            "ffmpeg", "-i", wav_path, "-codec:a", "libmp3lame", 
+                            "-qscale:a", "2", temp_path,
+                            "-y", "-loglevel", "error"
+                        ], check=True)
+                    elif response_format == 'opus':
+                        subprocess.run([
+                            "ffmpeg", "-i", wav_path, "-c:a", "libopus", 
+                            "-b:a", "64k", "-vbr", "on", temp_path,
+                            "-y", "-loglevel", "error"
+                        ], check=True)
+                    elif response_format == 'aac':
+                        subprocess.run([
+                            "ffmpeg", "-i", wav_path, "-c:a", "aac", 
+                            "-b:a", "128k", temp_path,
+                            "-y", "-loglevel", "error"
+                        ], check=True)
+                    elif response_format == 'flac':
+                        subprocess.run([
+                            "ffmpeg", "-i", wav_path, "-c:a", "flac", temp_path,
+                            "-y", "-loglevel", "error"
+                        ], check=True)
+                
+                # Clean up the temporary WAV file
+                try:
+                    os.unlink(wav_path)
+                except:
+                    pass
+            
+            # Read the processed audio file
+            with open(temp_path, "rb") as f:
+                response_data = f.read()
+            
+            # Store in cache if enabled
+            if cache_enabled and cache_key:
+                try:
+                    cache_path = os.path.join(getattr(app_state, "audio_cache_dir", "/app/audio_cache"), f"{cache_key}")
+                    with open(cache_path, "wb") as f:
+                        f.write(response_data)
+                    logger.debug(f"Cached {response_format} audio with key: {cache_key}")
+                except Exception as e:
+                    logger.warning(f"Error writing to cache: {e}")
+            
+            # Log processing time
+            processing_time = time.time() - start_time
+            logger.info(f"Processed audio to {response_format} in {processing_time:.3f}s")
+            
+            return response_data, format_to_content_type[response_format]
+        
+        except Exception as e:
+            logger.error(f"Error converting audio to {response_format}: {e}")
+            # Fallback to WAV if conversion fails
+            try:
+                wav_path = f"{temp_path}.wav"
+                torchaudio.save(wav_path, audio_cpu.unsqueeze(0), sample_rate)
+                with open(wav_path, "rb") as f:
+                    response_data = f.read()
+                os.unlink(wav_path)
+                return response_data, "audio/wav"
+            except Exception as fallback_error:
+                logger.error(f"Fallback to WAV also failed: {fallback_error}")
+                raise RuntimeError(f"Failed to generate audio in any format: {str(e)}")
+        
+        finally:
+            # Clean up the temporary file
+            try:
+                os.unlink(temp_path)
+            except:
+                pass
+
+@router.post("/audio/conversation", tags=["Conversation API"])
+async def conversation_to_speech(
+    request: Request,
+    text: str = Body(..., description="Text to convert to speech"),
+    speaker_id: int = Body(0, description="Speaker ID"),
+    context: List[Dict] = Body([], description="Context segments with speaker, text, and audio path"),
+):
+    """
+    Custom endpoint for conversational TTS using CSM-1B.
+    
+    This is not part of the OpenAI API but provides the unique conversational
+    capability of the CSM model.
+    """
+    # Get generator from app state
+    generator = request.app.state.generator
+    
+    # Validate model availability
+    if generator is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    try:
+        segments = []
+        
+        # Process context if provided
+        for ctx in context:
+            if 'speaker' not in ctx or 'text' not in ctx or 'audio' not in ctx:
+                continue
+                
+            # Audio should be base64-encoded
+            audio_data = base64.b64decode(ctx['audio'])
+            audio_file = io.BytesIO(audio_data)
+            
+            # Save to temporary file for torchaudio
+            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp:
+                temp.write(audio_file.read())
+                temp_path = temp.name
+            
+            # Load audio
+            audio_tensor, sample_rate = torchaudio.load(temp_path)
+            audio_tensor = torchaudio.functional.resample(
+                audio_tensor.squeeze(0), 
+                orig_freq=sample_rate, 
+                new_freq=generator.sample_rate
+            )
+            
+            # Clean up
+            os.unlink(temp_path)
+            
+            # Create segment
+            segments.append(
+                Segment(
+                    speaker=ctx['speaker'],
+                    text=ctx['text'],
+                    audio=audio_tensor
+                )
+            )
+            
+        logger.info(f"Conversation request: '{text}' with {len(segments)} context segments")
+        
+        # Format the text for better voice consistency
+        from app.prompt_engineering import format_text_for_voice
+        
+        # Determine voice name from speaker_id
+        voice_names = ["alloy", "echo", "fable", "onyx", "nova", "shimmer"]
+        voice_name = voice_names[speaker_id] if 0 <= speaker_id < len(voice_names) else "alloy"
+        
+        formatted_text = format_text_for_voice(text, voice_name)
+        
+        # Generate audio with context
+        audio = generator.generate(
+            text=formatted_text,
+            speaker=speaker_id,
+            context=segments,
+            max_audio_length_ms=20000,  # 20 seconds
+            temperature=0.7,  # Lower temperature for more stable output
+            topk=40,
+        )
+        
+        # Process audio for better quality
+        from app.voice_enhancement import process_generated_audio
+        
+        processed_audio = process_generated_audio(
+            audio, 
+            voice_name,
+            generator.sample_rate,
+            text
+        )
+        
+        # Save to temporary file
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp:
+            temp_path = temp.name
+        
+        # Save audio
+        torchaudio.save(temp_path, processed_audio.unsqueeze(0).cpu(), generator.sample_rate)
+        
+        # Return audio file
+        def iterfile():
+            with open(temp_path, 'rb') as f:
+                yield from f
+            # Clean up
+            if os.path.exists(temp_path):
+                os.unlink(temp_path)
+        
+        logger.info(f"Generated conversation response, duration: {processed_audio.shape[0]/generator.sample_rate:.2f}s")
+        
+        return StreamingResponse(
+            iterfile(),
+            media_type="audio/wav",
+            headers={'Content-Disposition': 'attachment; filename="speech.wav"'}
+        )
+    
+    except Exception as e:
+        import traceback
+        error_trace = traceback.format_exc()
+        logger.error(f"Conversation speech generation failed: {str(e)}\n{error_trace}")
+        raise HTTPException(status_code=500, detail=f"Conversation speech generation failed: {str(e)}")
+
+@router.get("/audio/voices", tags=["Audio"])
+async def list_voices(request: Request):
+    """
+    List available voices in a format compatible with OpenAI and OpenWebUI.
+    """
+    # Use app state's get_all_voices function if available
+    if hasattr(request.app.state, "get_all_voices"):
+        voices = request.app.state.get_all_voices()
+        logger.info(f"Listing {len(voices)} voices")
+        return {"voices": voices}
+    
+    # Fallback to standard voices if necessary
+    standard_voices = [
+        {"voice_id": "alloy", "name": "Alloy"},
+        {"voice_id": "echo", "name": "Echo"},
+        {"voice_id": "fable", "name": "Fable"},
+        {"voice_id": "onyx", "name": "Onyx"},
+        {"voice_id": "nova", "name": "Nova"},
+        {"voice_id": "shimmer", "name": "Shimmer"}
+    ]
+    
+    # Add cloned voices if available
+    if hasattr(request.app.state, "voice_cloner") and request.app.state.voice_cloner is not None:
+        cloned_voices = request.app.state.voice_cloner.list_voices()
+        for voice in cloned_voices:
+            standard_voices.append({
+                "voice_id": voice.id,  # This has to be specifically voice_id
+                "name": voice.name     # This has to be specifically name
+            })
+    
+    logger.info(f"Listing {len(standard_voices)} voices")
+    return {"voices": standard_voices}
+
+# Add OpenAI-compatible models list endpoint
+@router.get("/audio/models", tags=["Audio"], summary="List available audio models")
+async def list_models():
+    """
+    OpenAI compatible endpoint that returns a list of available audio models.
+    """
+    models = [
+        {
+            "id": "csm-1b",
+            "name": "CSM-1B",
+            "description": "Conversational Speech Model 1B from Sesame",
+            "created": 1716019200,  # March 13, 2025 (from the example)
+            "object": "audio",
+            "owned_by": "sesame",
+            "capabilities": {
+                "tts": True,
+                "voice_generation": True,
+                "voice_cloning": hasattr(router.app, "voice_cloner"),
+                "streaming": True
+            },
+            "max_input_length": 4096,
+            "price": {"text-to-speech": 0.00}
+        },
+        {
+            "id": "tts-1",
+            "name": "CSM-1B (Compatibility Mode)",
+            "description": "CSM-1B with OpenAI TTS-1 compatibility",
+            "created": 1716019200,
+            "object": "audio",
+            "owned_by": "sesame",
+            "capabilities": {
+                "tts": True,
+                "voice_generation": True,
+                "streaming": True
+            },
+            "max_input_length": 4096,
+            "price": {"text-to-speech": 0.00}
+        },
+        {
+            "id": "tts-1-hd",
+            "name": "CSM-1B (HD Mode)",
+            "description": "CSM-1B with higher quality settings",
+            "created": 1716019200,
+            "object": "audio",
+            "owned_by": "sesame",
+            "capabilities": {
+                "tts": True,
+                "voice_generation": True,
+                "streaming": True
+            },
+            "max_input_length": 4096,
+            "price": {"text-to-speech": 0.00}
+        }
+    ]
+    
+    return {"data": models, "object": "list"}
+
+# Response format options endpoint
+@router.get("/audio/speech/response-formats", tags=["Audio"], summary="List available response formats")
+async def list_response_formats():
+    """List available response formats for speech synthesis."""
+    formats = [
+        {"name": "mp3", "content_type": "audio/mpeg"},
+        {"name": "opus", "content_type": "audio/opus"},
+        {"name": "aac", "content_type": "audio/aac"},
+        {"name": "flac", "content_type": "audio/flac"},
+        {"name": "wav", "content_type": "audio/wav"}
+    ]
+    
+    return {"response_formats": formats}
+
+# Streaming format options endpoint
+@router.get("/audio/speech/stream-formats", tags=["Audio"], summary="List available streaming formats")
+async def list_stream_formats():
+    """List available streaming formats for TTS."""
+    return {
+        "stream_formats": [
+            {
+                "format": "mp3", 
+                "content_type": "audio/mpeg",
+                "description": "MP3 audio format (streaming)"
+            },
+            {
+                "format": "opus", 
+                "content_type": "audio/opus",
+                "description": "Opus audio format (streaming)"
+            },
+            {
+                "format": "aac", 
+                "content_type": "audio/aac",
+                "description": "AAC audio format (streaming)"
+            },
+            {
+                "format": "flac", 
+                "content_type": "audio/flac",
+                "description": "FLAC audio format (streaming)"
+            },
+            {
+                "format": "wav", 
+                "content_type": "audio/wav",
+                "description": "WAV audio format (streaming)"
+            }
+        ]
+    }
+
+# Simple test endpoint
+@router.get("/test", tags=["Utility"], summary="Test endpoint")
+async def test_endpoint():
+    """Simple test endpoint that returns a successful response."""
+    return {"status": "ok", "message": "API is working"}
+
+# Debug endpoint
+@router.get("/debug", tags=["Utility"], summary="Debug endpoint")
+async def debug_info(request: Request):
+    """Get debug information about the API."""
+    generator = request.app.state.generator
+    
+    # Basic info
+    debug_info = {
+        "model_loaded": generator is not None,
+        "device": generator.device if generator is not None else None,
+        "sample_rate": generator.sample_rate if generator is not None else None,
+    }
+    
+    # Add voice enhancement info if available
+    try:
+        from app.voice_enhancement import VOICE_PROFILES
+        voice_info = {}
+        for name, profile in VOICE_PROFILES.items():
+            voice_info[name] = {
+                "pitch_range": f"{profile.pitch_range[0]}-{profile.pitch_range[1]}Hz",
+                "timbre": profile.timbre,
+                "ref_segments": len(profile.reference_segments),
+            }
+        debug_info["voice_profiles"] = voice_info
+    except ImportError:
+        debug_info["voice_profiles"] = "Not available"
+        
+    # Add voice cloning info if available
+    if hasattr(request.app.state, "voice_cloner"):
+        voice_cloner = request.app.state.voice_cloner
+        debug_info["voice_cloning"] = {
+            "enabled": True,
+            "cloned_voices_count": len(voice_cloner.list_voices()),
+            "cloned_voices": [v.name for v in voice_cloner.list_voices()]
+        }
+    else:
+        debug_info["voice_cloning"] = {"enabled": False}
+    
+    # Add streaming info
+    debug_info["streaming"] = {"enabled": True}
+    
+    # Add memory usage info for CUDA
+    if torch.cuda.is_available():
+        debug_info["cuda"] = {
+            "allocated_memory_gb": torch.cuda.memory_allocated() / 1e9,
+            "reserved_memory_gb": torch.cuda.memory_reserved() / 1e9,
+            "max_memory_gb": torch.cuda.get_device_properties(0).total_memory / 1e9,
+        }
+    
+    return debug_info
+
+@router.get("/voice-management/info", tags=["Voice Management"])
+async def get_voice_storage_info(request: Request):
+    """Get information about voice storage usage and status."""
+    from app.utils.voice_manager import get_voice_storage_info
+    return get_voice_storage_info()
+
+@router.post("/voice-management/backup", tags=["Voice Management"])
+async def create_voice_backup(request: Request):
+    """Create a backup of all voice data."""
+    from app.utils.voice_manager import backup_voice_data
+    backup_path = backup_voice_data()
+    return {"status": "success", "backup_path": backup_path}
+
+@router.post("/voice-management/reset-voices", tags=["Voice Management"])
+async def reset_voices(request: Request):
+    """Reset voices to their default state."""
+    from app.utils.voice_manager import restore_default_voices
+    backup_path = restore_default_voices()
+    return {"status": "success", "backup_path": backup_path, "message": "Voices reset to default state"}
+
+@router.get("/voice-management/verify-references", tags=["Voice Management"])
+async def verify_references(request: Request):
+    """Check if voice references are complete and valid."""
+    from app.utils.voice_manager import verify_voice_references
+    return verify_voice_references()
+
+# Voice diagnostics endpoint
+@router.get("/debug/voices", tags=["Debug"], summary="Voice diagnostics")
+async def voice_diagnostics():
+    """Get diagnostic information about voice references."""
+    try:
+        from app.voice_enhancement import VOICE_PROFILES
+        
+        diagnostics = {}
+        for name, profile in VOICE_PROFILES.items():
+            ref_info = []
+            for i, ref in enumerate(profile.reference_segments):
+                if ref is not None:
+                    duration = ref.shape[0] / 24000  # Assume 24kHz
+                    ref_info.append({
+                        "index": i,
+                        "duration_seconds": f"{duration:.2f}",
+                        "samples": ref.shape[0],
+                        "min": float(ref.min()),
+                        "max": float(ref.max()),
+                        "rms": float(torch.sqrt(torch.mean(ref ** 2))),
+                    })
+            
+            diagnostics[name] = {
+                "speaker_id": profile.speaker_id,
+                "pitch_range": f"{profile.pitch_range[0]}-{profile.pitch_range[1]}Hz",
+                "references": ref_info,
+                "reference_count": len(ref_info),
+            }
+        
+        return {"diagnostics": diagnostics}
+    except ImportError:
+        return {"error": "Voice enhancement module not available"}
+
+# Specialized debugging endpoint for speech generation
+@router.post("/debug/speech", tags=["Debug"], summary="Debug speech generation")
+async def debug_speech(
+    request: Request,
+    text: str = Body(..., embed=True),
+    voice: str = Body("alloy", embed=True),
+    use_enhancement: bool = Body(True, embed=True)
+):
+    """Debug endpoint for speech generation with enhancement options."""
+    generator = request.app.state.generator
+    
+    if generator is None:
+        return {"error": "Model not loaded"}
+    
+    try:
+        # Convert voice name to speaker ID
+        voice_map = {
+            "alloy": 0, 
+            "echo": 1, 
+            "fable": 2, 
+            "onyx": 3, 
+            "nova": 4, 
+            "shimmer": 5
+        }
+        speaker = voice_map.get(voice, 0)
+        
+        # Format text if using enhancement
+        if use_enhancement:
+            from app.prompt_engineering import format_text_for_voice
+            formatted_text = format_text_for_voice(text, voice)
+            logger.info(f"Using formatted text: {formatted_text}")
+        else:
+            formatted_text = text
+            
+        # Get context if using enhancement
+        if use_enhancement:
+            from app.voice_enhancement import get_voice_segments
+            context = get_voice_segments(voice, generator.device)
+            logger.info(f"Using {len(context)} context segments")
+        else:
+            context = []
+            
+        # Generate audio
+        start_time = time.time()
+        audio = generator.generate(
+            text=formatted_text,
+            speaker=speaker,
+            context=context,
+            max_audio_length_ms=10000,  # 10 seconds
+            temperature=0.7 if use_enhancement else 0.9,
+            topk=40 if use_enhancement else 50,
+        )
+        generation_time = time.time() - start_time
+        
+        # Process audio if using enhancement
+        if use_enhancement:
+            from app.voice_enhancement import process_generated_audio
+            start_time = time.time()
+            processed_audio = process_generated_audio(audio, voice, generator.sample_rate, text)
+            processing_time = time.time() - start_time
+        else:
+            processed_audio = audio
+            processing_time = 0
+        
+        # Save to temporary WAV file
+        temp_path = f"/tmp/debug_speech_{voice}_{int(time.time())}.wav"
+        torchaudio.save(temp_path, processed_audio.unsqueeze(0).cpu(), generator.sample_rate)
+        
+        # Also save original if enhanced
+        if use_enhancement:
+            orig_path = f"/tmp/debug_speech_{voice}_original_{int(time.time())}.wav"
+            torchaudio.save(orig_path, audio.unsqueeze(0).cpu(), generator.sample_rate)
+        else:
+            orig_path = temp_path
+            
+        # Calculate audio metrics
+        duration = processed_audio.shape[0] / generator.sample_rate
+        rms = float(torch.sqrt(torch.mean(processed_audio ** 2)))
+        peak = float(processed_audio.abs().max())
+        
+        return {
+            "status": "success",
+            "message": f"Audio generated successfully and saved to {temp_path}",
+            "audio": {
+                "duration_seconds": f"{duration:.2f}",
+                "samples": processed_audio.shape[0],
+                "sample_rate": generator.sample_rate,
+                "rms_level": f"{rms:.3f}",
+                "peak_level": f"{peak:.3f}",
+            },
+            "processing": {
+                "enhancement_used": use_enhancement,
+                "generation_time_seconds": f"{generation_time:.3f}",
+                "processing_time_seconds": f"{processing_time:.3f}",
+                "original_path": orig_path,
+                "processed_path": temp_path,
+            }
+        }
+    except Exception as e:
+        import traceback
+        error_trace = traceback.format_exc()
+        logger.error(f"Debug speech generation failed: {e}\n{error_trace}")
+        return {
+            "status": "error",
+            "message": str(e),
+            "traceback": error_trace
+        }

app/api/schemas.py

@@ -0,0 +1,41 @@
+# app/api/schemas.py
+from enum import Enum
+from typing import Optional, List, Dict, Any, Union
+from pydantic import BaseModel, Field
+
+# Voice options as a non-restrictive string
+class Voice(str):
+    """Voice options for CSM model - allowing any string value"""
+    pass
+
+class ResponseFormat(str, Enum):
+    mp3 = "mp3"
+    opus = "opus"
+    aac = "aac"
+    flac = "flac"
+    wav = "wav"
+
+# Create SpeechRequest for compatibility with our new code
+class SpeechRequest(BaseModel):
+    model: Optional[str] = Field("csm-1b", description="The TTS model to use")
+    input: str = Field(..., description="The text to generate audio for")
+    voice: Optional[str] = Field("alloy", description="The voice to use for generation")
+    response_format: Optional[ResponseFormat] = Field(ResponseFormat.mp3, description="The format of the audio response")
+    speed: Optional[float] = Field(1.0, description="The speed of the audio", ge=0.25, le=4.0)
+    # CSM-specific parameters
+    max_audio_length_ms: Optional[float] = Field(90000, description="Maximum audio length in milliseconds")
+    temperature: Optional[float] = Field(0.9, description="Sampling temperature", ge=0.0, le=2.0)
+    topk: Optional[int] = Field(50, description="Top-k for sampling", ge=1, le=100)
+    
+    class Config:
+        populate_by_name = True
+        extra = "ignore"  # Allow extra fields without error
+
+# Maintain TTSRequest for backward compatibility
+class TTSRequest(SpeechRequest):
+    """Legacy alias for SpeechRequest for backward compatibility"""
+    pass
+
+class TTSResponse(BaseModel):
+    """Only used for API documentation"""
+    pass

app/api/streaming.py

@@ -0,0 +1,315 @@
+"""Streaming support for the TTS API."""
+import asyncio
+import io
+import logging
+import time
+from typing import AsyncGenerator, Optional, List
+import torch
+import torchaudio
+from fastapi import APIRouter, Request, HTTPException
+from fastapi.responses import StreamingResponse
+from app.api.schemas import SpeechRequest, ResponseFormat
+from app.prompt_engineering import split_into_segments
+from app.models import Segment
+
+logger = logging.getLogger(__name__)
+router = APIRouter()
+
+class AudioChunker:
+    """Handle audio chunking for streaming responses."""
+    def __init__(self,
+                 sample_rate: int,
+                 format: str = "mp3",
+                 chunk_size_ms: int = 200):  # Smaller chunks for better streaming
+        """
+        Initialize audio chunker.
+        Args:
+            sample_rate: Audio sample rate in Hz
+            format: Output audio format (mp3, opus, etc.)
+            chunk_size_ms: Size of each chunk in milliseconds
+        """
+        self.sample_rate = sample_rate
+        self.format = format.lower()
+        self.chunk_size_samples = int(sample_rate * (chunk_size_ms / 1000))
+        logger.info(f"Audio chunker initialized with {chunk_size_ms}ms chunks ({self.chunk_size_samples} samples)")
+
+    async def chunk_audio(self,
+                    audio: torch.Tensor,
+                    delay_ms: int = 0) -> AsyncGenerator[bytes, None]:
+        """
+        Convert audio tensor to streaming chunks.
+        Args:
+            audio: Audio tensor to stream
+            delay_ms: Artificial delay between chunks (for testing)
+        Yields:
+            Audio chunks as bytes
+        """
+        # Ensure audio is on CPU
+        if audio.is_cuda:
+            audio = audio.cpu()
+        # Calculate number of chunks
+        num_samples = audio.shape[0]
+        num_chunks = (num_samples + self.chunk_size_samples - 1) // self.chunk_size_samples
+        logger.info(f"Streaming {num_samples} samples as {num_chunks} chunks")
+        for i in range(num_chunks):
+            start_idx = i * self.chunk_size_samples
+            end_idx = min(start_idx + self.chunk_size_samples, num_samples)
+            # Extract chunk
+            chunk = audio[start_idx:end_idx]
+            # Convert to bytes in requested format
+            chunk_bytes = await self._format_chunk(chunk)
+            # Add artificial delay if requested (for testing)
+            if delay_ms > 0:
+                await asyncio.sleep(delay_ms / 1000)
+            yield chunk_bytes
+
+    async def _format_chunk(self, chunk: torch.Tensor) -> bytes:
+        """Convert audio chunk to bytes in the specified format."""
+        buf = io.BytesIO()
+        # Ensure chunk is 1D and on CPU
+        if len(chunk.shape) == 1:
+            chunk = chunk.unsqueeze(0)  # Add channel dimension
+        # Ensure chunk is on CPU
+        if chunk.is_cuda:
+            chunk = chunk.cpu()
+        # Save to buffer in specified format
+        if self.format == "mp3":
+            torchaudio.save(buf, chunk, self.sample_rate, format="mp3")
+        elif self.format == "opus":
+            torchaudio.save(buf, chunk, self.sample_rate, format="opus")
+        elif self.format == "aac":
+            torchaudio.save(buf, chunk, self.sample_rate, format="aac")
+        elif self.format == "flac":
+            torchaudio.save(buf, chunk, self.sample_rate, format="flac")
+        elif self.format == "wav":
+            torchaudio.save(buf, chunk, self.sample_rate, format="wav")
+        else:
+            # Default to mp3
+            torchaudio.save(buf, chunk, self.sample_rate, format="mp3")
+        # Get bytes from buffer
+        buf.seek(0)
+        return buf.read()
+
+# Helper function to get speaker ID for a voice
+def get_speaker_id(app_state, voice):
+    """Helper function to get speaker ID from voice name or ID"""
+    if hasattr(app_state, "voice_speaker_map") and voice in app_state.voice_speaker_map:
+        return app_state.voice_speaker_map[voice]
+    # Standard voices mapping
+    voice_to_speaker = {"alloy": 0, "echo": 1, "fable": 2, "onyx": 3, "nova": 4, "shimmer": 5}
+    if voice in voice_to_speaker:
+        return voice_to_speaker[voice]
+    # Try parsing as integer
+    try:
+        speaker_id = int(voice)
+        if 0 <= speaker_id < 6:
+            return speaker_id
+    except (ValueError, TypeError):
+        pass
+    # Check cloned voices if the voice cloner exists
+    if hasattr(app_state, "voice_cloner") and app_state.voice_cloner is not None:
+        # Check by ID
+        if voice in app_state.voice_cloner.cloned_voices:
+            return app_state.voice_cloner.cloned_voices[voice].speaker_id
+        # Check by name
+        for v_id, v_info in app_state.voice_cloner.cloned_voices.items():
+            if v_info.name.lower() == voice.lower():
+                return v_info.speaker_id
+    # Default to alloy
+    return 0
+
+@router.post("/audio/speech/stream", tags=["Audio"])
+async def stream_speech(
+    request: Request,
+    speech_request: SpeechRequest,
+):
+    """
+    Stream audio of text being spoken by a realistic voice.
+    This endpoint provides an OpenAI-compatible streaming interface for TTS.
+    """
+    # Check if model is loaded
+    if not hasattr(request.app.state, "generator") or request.app.state.generator is None:
+        raise HTTPException(
+            status_code=503,
+            detail="Model not loaded. Please try again later."
+        )
+
+    # Get request parameters
+    model = speech_request.model
+    input_text = speech_request.input
+    voice = speech_request.voice
+    response_format = speech_request.response_format
+    speed = speech_request.speed
+    temperature = speech_request.temperature
+    max_audio_length_ms = speech_request.max_audio_length_ms
+
+    # Log the request
+    logger.info(f"Real-time streaming speech from text ({len(input_text)} chars) with voice '{voice}'")
+
+    # Check if text is empty
+    if not input_text or len(input_text.strip()) == 0:
+        raise HTTPException(
+            status_code=400,
+            detail="Input text cannot be empty"
+        )
+
+    # Get speaker ID for the voice
+    speaker_id = get_speaker_id(request.app.state, voice)
+    if speaker_id is None:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Voice '{voice}' not found. Available voices: {request.app.state.available_voices}"
+        )
+
+    try:
+        # Create media type based on format
+        media_type = {
+            "mp3": "audio/mpeg",
+            "opus": "audio/opus",
+            "aac": "audio/aac",
+            "flac": "audio/flac",
+            "wav": "audio/wav",
+        }.get(response_format, "audio/mpeg")
+
+        # Create the chunker for streaming
+        sample_rate = request.app.state.sample_rate
+        chunker = AudioChunker(sample_rate, response_format)
+
+        # Split text into segments using the imported function
+        from app.prompt_engineering import split_into_segments
+        text_segments = split_into_segments(input_text, max_chars=50)  # Smaller segments for faster first response
+        
+        logger.info(f"Split text into {len(text_segments)} segments for incremental streaming")
+
+        async def generate_streaming_audio():
+            # Check for cloned voice
+            voice_info = None
+            from_cloned_voice = False
+            
+            if hasattr(request.app.state, "voice_cloning_enabled") and request.app.state.voice_cloning_enabled:
+                voice_info = request.app.state.get_voice_info(voice)
+                from_cloned_voice = voice_info and voice_info["type"] == "cloned"
+                if from_cloned_voice:
+                    # Use cloned voice context for first segment
+                    voice_cloner = request.app.state.voice_cloner
+                    context = voice_cloner.get_voice_context(voice_info["voice_id"])
+                else:
+                    # Use standard voice context
+                    from app.voice_enhancement import get_voice_segments
+                    context = get_voice_segments(voice, request.app.state.device)
+            else:
+                # Use standard voice context
+                from app.voice_enhancement import get_voice_segments
+                context = get_voice_segments(voice, request.app.state.device)
+
+            # Send an empty chunk to initialize the connection
+            yield b''
+
+            # Process each text segment incrementally and stream in real time
+            for i, segment_text in enumerate(text_segments):
+                try:
+                    logger.info(f"Generating segment {i+1}/{len(text_segments)}")
+                    
+                    # Generate audio for this segment - use async to avoid blocking
+                    if from_cloned_voice:
+                        # Generate with cloned voice
+                        voice_cloner = request.app.state.voice_cloner
+                        
+                        # Convert to asynchronous with asyncio.to_thread
+                        segment_audio = await asyncio.to_thread(
+                            voice_cloner.generate_speech,
+                            segment_text,
+                            voice_info["voice_id"],
+                            temperature=temperature,
+                            topk=30,
+                            max_audio_length_ms=2000  # Keep segments short for streaming
+                        )
+                    else:
+                        # Use standard voice with generator
+                        segment_audio = await asyncio.to_thread(
+                            request.app.state.generator.generate,
+                            segment_text,
+                            speaker_id,
+                            context,
+                            max_audio_length_ms=2000,  # Short for quicker generation
+                            temperature=temperature
+                        )
+                    
+                    # Process audio quality for this segment
+                    if hasattr(request.app.state, "voice_enhancement_enabled") and request.app.state.voice_enhancement_enabled:
+                        from app.voice_enhancement import process_generated_audio
+                        segment_audio = process_generated_audio(
+                            audio=segment_audio,
+                            voice_name=voice,
+                            sample_rate=sample_rate,
+                            text=segment_text
+                        )
+                        
+                    # Handle speed adjustment
+                    if speed != 1.0 and speed > 0:
+                        try:
+                            # Adjust speed using torchaudio
+                            effects = [["tempo", str(speed)]]
+                            audio_cpu = segment_audio.cpu()
+                            adjusted_audio, _= torchaudio.sox_effects.apply_effects_tensor(
+                                audio_cpu.unsqueeze(0),
+                                sample_rate,
+                                effects
+                            )
+                            segment_audio = adjusted_audio.squeeze(0)
+                        except Exception as e:
+                            logger.warning(f"Failed to adjust speech speed: {e}")
+                            
+                    # Convert this segment to bytes and stream immediately
+                    buf = io.BytesIO()
+                    audio_to_save = segment_audio.unsqueeze(0) if len(segment_audio.shape) == 1 else segment_audio
+                    torchaudio.save(buf, audio_to_save.cpu(), sample_rate, format=response_format)
+                    buf.seek(0)
+                    segment_bytes = buf.read()
+                    
+                    # Stream this segment immediately
+                    yield segment_bytes
+                    
+                    # Update context with this segment for next generation
+                    context = [
+                        Segment(
+                            text=segment_text,
+                            speaker=speaker_id,
+                            audio=segment_audio
+                        )
+                    ]
+                    
+                except Exception as e:
+                    logger.error(f"Error generating segment {i+1}: {e}")
+                    # Try to continue with next segment
+                    
+        # Return streaming response
+        return StreamingResponse(
+            generate_streaming_audio(),
+            media_type=media_type,
+            headers={
+                "Content-Disposition": f'attachment; filename="speech.{response_format}"',
+                "X-Accel-Buffering": "no",  # Prevent buffering in nginx
+                "Cache-Control": "no-cache, no-store, must-revalidate",  # Prevent caching
+                "Pragma": "no-cache",
+                "Expires": "0",
+                "Connection": "keep-alive",
+                "Transfer-Encoding": "chunked"
+            }
+        )
+    except Exception as e:
+        logger.error(f"Error in stream_speech: {e}")
+        raise HTTPException(status_code=500, detail=f"Error generating speech: {str(e)}")
+    
+@router.post("/audio/speech/streaming", tags=["Audio"])
+async def openai_stream_speech(
+    request: Request,
+    speech_request: SpeechRequest,
+):
+    """
+    Stream audio in OpenAI-compatible streaming format.
+    This endpoint is compatible with the OpenAI streaming TTS API.
+    """
+    # Use the same logic as the stream_speech endpoint but with a different name
+    # to maintain the OpenAI API naming convention
+    return await stream_speech(request, speech_request)

app/api/utils.py

@@ -0,0 +1,25 @@
+def get_all_voices(app_state):
+    """
+    Get all available voices including standard and cloned voices.
+    Returns them in a format compatible with OpenAI's API.
+    """
+    # Standard voices
+    voices = [
+        {"voice_id": "alloy", "name": "Alloy"},
+        {"voice_id": "echo", "name": "Echo"},
+        {"voice_id": "fable", "name": "Fable"},
+        {"voice_id": "onyx", "name": "Onyx"},
+        {"voice_id": "nova", "name": "Nova"},
+        {"voice_id": "shimmer", "name": "Shimmer"}
+    ]
+    
+    # Add cloned voices if available
+    if hasattr(app_state, "voice_cloner") and app_state.voice_cloner is not None:
+        cloned_voices = app_state.voice_cloner.list_voices()
+        for voice in cloned_voices:
+            voices.append({
+                "voice_id": voice.id,
+                "name": voice.name
+            })
+    
+    return voices

app/api/voice_cloning_routes.py

@@ -0,0 +1,289 @@
+"""Voice cloning API routes for CSM-1B TTS API."""
+import os
+import io
+import time
+import tempfile
+from typing import Dict, List, Optional, Any
+import torch
+import torchaudio
+from fastapi import APIRouter, Request, Response, HTTPException, UploadFile, File, Form, Body
+from fastapi.responses import StreamingResponse, JSONResponse
+from app.voice_cloning import ClonedVoice
+
+# Create router
+router = APIRouter(prefix="/voice-cloning", tags=["Voice Cloning"])
+
+@router.post("/clone", summary="Clone a new voice")
+async def clone_voice(
+    request: Request,
+    audio_file: UploadFile = File(...),
+    name: str = Form(...),
+    transcript: Optional[str] = Form(None),
+    description: Optional[str] = Form(None)
+):
+    """
+    Clone a new voice from an audio file.
+    
+    - **audio_file**: Audio file with the voice to clone (MP3, WAV, etc.)
+    - **name**: Name for the cloned voice
+    - **transcript**: Optional transcript of the audio
+    - **description**: Optional description of the voice
+    """
+    if not hasattr(request.app.state, "voice_cloner"):
+        raise HTTPException(status_code=503, detail="Voice cloning service not available")
+    
+    voice_cloner = request.app.state.voice_cloner
+    
+    try:
+        voice = await voice_cloner.clone_voice(
+            audio_file=audio_file,
+            voice_name=name,
+            transcript=transcript,
+            description=description
+        )
+        
+        return {
+            "status": "success",
+            "message": "Voice cloned successfully",
+            "voice": voice
+        }
+    except Exception as e:
+        import traceback
+        error_trace = traceback.format_exc()
+        request.app.state.logger.error(f"Voice cloning failed: {e}\n{error_trace}")
+        raise HTTPException(status_code=500, detail=f"Voice cloning failed: {str(e)}")
+
+@router.get("/voices", summary="List cloned voices")
+async def list_voices(request: Request):
+    """List all cloned voices available in the system."""
+    if not hasattr(request.app.state, "voice_cloner"):
+        raise HTTPException(status_code=503, detail="Voice cloning service not available")
+    
+    voice_cloner = request.app.state.voice_cloner
+    voices = voice_cloner.list_voices()
+    
+    return {
+        "voices": voices
+    }
+
+@router.delete("/voices/{voice_id}", summary="Delete a cloned voice")
+async def delete_voice(request: Request, voice_id: str):
+    """Delete a cloned voice by ID."""
+    if not hasattr(request.app.state, "voice_cloner"):
+        raise HTTPException(status_code=503, detail="Voice cloning service not available")
+    
+    voice_cloner = request.app.state.voice_cloner
+    success = voice_cloner.delete_voice(voice_id)
+    
+    if not success:
+        raise HTTPException(status_code=404, detail=f"Voice with ID {voice_id} not found")
+    
+    return {
+        "status": "success",
+        "message": f"Voice {voice_id} deleted successfully"
+    }
+
+@router.post("/clone-from-youtube", summary="Clone a voice from YouTube")
+async def clone_voice_from_youtube(
+    request: Request,
+    data: dict = Body(...)  # Use a single body parameter to avoid conflicts
+):
+    """
+    Clone a voice from a YouTube video.
+    
+    - **youtube_url**: URL of the YouTube video
+    - **voice_name**: Name for the cloned voice
+    - **start_time**: Start time in seconds (default: 0)
+    - **duration**: Duration to extract in seconds (default: 180)
+    - **description**: Optional description of the voice
+    """
+    if not hasattr(request.app.state, "voice_cloner"):
+        raise HTTPException(status_code=503, detail="Voice cloning service not available")
+    
+    voice_cloner = request.app.state.voice_cloner
+    
+    # Extract parameters from the request body
+    youtube_url = data.get("youtube_url")
+    voice_name = data.get("voice_name")
+    start_time = data.get("start_time", 0)
+    duration = data.get("duration", 180)
+    description = data.get("description")
+    
+    # Validate required parameters
+    if not youtube_url or not voice_name:
+        raise HTTPException(status_code=400, detail="Missing required parameters: youtube_url and voice_name")
+    
+    try:
+        voice = await voice_cloner.clone_voice_from_youtube(
+            youtube_url=youtube_url,
+            voice_name=voice_name,
+            start_time=start_time,
+            duration=duration,
+            description=description
+        )
+        
+        return {
+            "status": "success",
+            "message": "Voice cloned successfully from YouTube",
+            "voice": voice
+        }
+    except Exception as e:
+        import traceback
+        error_trace = traceback.format_exc()
+        request.app.state.logger.error(f"YouTube voice cloning failed: {e}\n{error_trace}")
+        raise HTTPException(status_code=500, detail=f"YouTube voice cloning failed: {str(e)}")
+    
+@router.post("/generate", summary="Generate speech with cloned voice")
+async def generate_speech(
+    request: Request,
+    voice_id: str = Body(..., embed=True),
+    text: str = Body(..., embed=True),
+    temperature: float = Body(0.65, embed=True),
+    response_format: str = Body("mp3", embed=True)
+):
+    """
+    Generate speech using a cloned voice.
+    
+    - **voice_id**: ID of the cloned voice to use
+    - **text**: Text to synthesize
+    - **temperature**: Sampling temperature (lower = more stable, higher = more varied)
+    - **response_format**: Audio format (mp3, wav, etc.)
+    """
+    if not hasattr(request.app.state, "voice_cloner"):
+        raise HTTPException(status_code=503, detail="Voice cloning service not available")
+    
+    voice_cloner = request.app.state.voice_cloner
+    
+    # Validate voice ID
+    if voice_id not in voice_cloner.cloned_voices:
+        raise HTTPException(status_code=404, detail=f"Voice with ID {voice_id} not found")
+    
+    # MIME type mapping
+    mime_types = {
+        "mp3": "audio/mpeg",
+        "wav": "audio/wav",
+        "ogg": "audio/ogg",
+        "flac": "audio/flac",
+        "m4a": "audio/mp4",
+    }
+    
+    # Set default if format not specified
+    if response_format not in mime_types:
+        response_format = "mp3"
+    
+    try:
+        # Generate speech with the cloned voice - IMPORTANT: This is a synchronous operation
+        # Remove the await keyword here
+        audio = voice_cloner.generate_speech(
+            text=text,
+            voice_id=voice_id,
+            temperature=temperature
+        )
+        
+        # Create temporary file for audio conversion
+        with tempfile.NamedTemporaryFile(suffix=f".{response_format}", delete=False) as temp_file:
+            temp_path = temp_file.name
+            
+        # Save to WAV first (direct format for torchaudio)
+        wav_path = f"{temp_path}.wav"
+        torchaudio.save(wav_path, audio.unsqueeze(0).cpu(), voice_cloner.sample_rate)
+        
+        # Convert to requested format
+        import ffmpeg
+        
+        if response_format == "mp3":
+            (
+                ffmpeg.input(wav_path)
+                .output(temp_path, format='mp3', audio_bitrate='128k')
+                .run(quiet=True, overwrite_output=True)
+            )
+        elif response_format == "ogg":
+            (
+                ffmpeg.input(wav_path)
+                .output(temp_path, format='ogg')
+                .run(quiet=True, overwrite_output=True)
+            )
+        elif response_format == "flac":
+            (
+                ffmpeg.input(wav_path)
+                .output(temp_path, format='flac')
+                .run(quiet=True, overwrite_output=True)
+            )
+        elif response_format == "m4a":
+            (
+                ffmpeg.input(wav_path)
+                .output(temp_path, format='mp4')
+                .run(quiet=True, overwrite_output=True)
+            )
+        else:  # wav
+            temp_path = wav_path
+        
+        # Clean up the temporary WAV file if we created a different format
+        if temp_path != wav_path and os.path.exists(wav_path):
+            os.unlink(wav_path)
+        
+        # Return audio file as response
+        def iterfile():
+            with open(temp_path, 'rb') as f:
+                yield from f
+            # Clean up temp file after streaming
+            if os.path.exists(temp_path):
+                os.unlink(temp_path)
+        
+        return StreamingResponse(
+            iterfile(),
+            media_type=mime_types.get(response_format, "application/octet-stream"),
+            headers={'Content-Disposition': f'attachment; filename="speech.{response_format}"'}
+        )
+    
+    except Exception as e:
+        import traceback
+        error_trace = traceback.format_exc()
+        request.app.state.logger.error(f"Speech generation failed: {e}\n{error_trace}")
+        raise HTTPException(status_code=500, detail=f"Speech generation failed: {str(e)}")
+
+@router.post("/voices/{voice_id}/preview", summary="Generate a preview of a cloned voice")
+async def generate_preview(
+    request: Request,
+    voice_id: str,
+    text: Optional[str] = Body("This is a preview of my cloned voice.", embed=True),
+    response_format: str = Body("mp3", embed=True)
+):
+    """
+    Generate a preview of a cloned voice with a standard text.
+    
+    - **voice_id**: ID of the cloned voice to use
+    - **text**: Optional custom text for the preview
+    - **response_format**: Audio format (mp3, wav, etc.)
+    """
+    # Use the generate_speech endpoint with a standard text
+    return await generate_speech(
+        request=request,
+        voice_id=voice_id,
+        text=text,
+        temperature=0.7,
+        response_format=response_format
+    )
+
+@router.get("/openai-compatible-voices", summary="List cloned voices in OpenAI format")
+async def list_voices_openai_format(request: Request):
+    """List all cloned voices in OpenAI-compatible format."""
+    if not hasattr(request.app.state, "voice_cloner"):
+        raise HTTPException(status_code=503, detail="Voice cloning service not available")
+    
+    voice_cloner = request.app.state.voice_cloner
+    voices = voice_cloner.list_voices()
+    
+    # Format voices in OpenAI-compatible format
+    openai_voices = []
+    for voice in voices:
+        openai_voices.append({
+            "voice_id": voice.id,
+            "name": voice.name,
+            "preview_url": f"/v1/voice-cloning/voices/{voice.id}/preview",
+            "description": voice.description or f"Cloned voice: {voice.name}",
+            "languages": [{"language_code": "en", "name": "English"}],
+            "cloned": True
+        })
+    
+    return {"voices": openai_voices}

app/audio_processing.py

@@ -0,0 +1,230 @@
+"""Audio processing utilities for CSM-1B TTS API."""
+import logging
+import numpy as np
+import torch
+from scipy import signal
+
+logger = logging.getLogger(__name__)
+
+def remove_long_silences(
+    audio: torch.Tensor, 
+    sample_rate: int,
+    min_speech_energy: float = 0.015,
+    max_silence_sec: float = 0.4,
+    keep_silence_sec: float = 0.1,
+) -> torch.Tensor:
+    """
+    Remove uncomfortably long silences from audio while preserving natural pauses.
+    
+    Args:
+        audio: Audio tensor
+        sample_rate: Sample rate in Hz
+        min_speech_energy: Minimum RMS energy to consider as speech
+        max_silence_sec: Maximum silence duration to keep in seconds
+        keep_silence_sec: Amount of silence to keep at speech boundaries
+        
+    Returns:
+        Audio with long silences removed
+    """
+    # Convert to numpy for processing
+    audio_np = audio.cpu().numpy()
+    
+    # Calculate frame size and hop length
+    frame_size = int(0.02 * sample_rate)  # 20ms frames
+    hop_length = int(0.01 * sample_rate)  # 10ms hop
+    
+    # Compute frame energy
+    frames = []
+    for i in range(0, len(audio_np) - frame_size + 1, hop_length):
+        frames.append(audio_np[i:i+frame_size])
+    
+    if len(frames) < 2:  # If audio is too short for analysis
+        return audio
+        
+    frames = np.array(frames)
+    # Root mean square energy
+    frame_energy = np.sqrt(np.mean(frames**2, axis=1))
+    
+    # Adaptive threshold based on audio content
+    # Uses a percentile to adapt to different audio characteristics
+    energy_threshold = max(
+        min_speech_energy,  # Minimum threshold
+        np.percentile(frame_energy, 10)  # Adapt to audio
+    )
+    
+    # Identify speech frames
+    is_speech = frame_energy > energy_threshold
+    
+    # Convert frame indices to sample indices considering overlapping frames
+    speech_segments = []
+    in_speech = False
+    speech_start = 0
+    
+    for i in range(len(is_speech)):
+        if is_speech[i] and not in_speech:
+            # Start of speech
+            in_speech = True
+            # Calculate start sample including keep_silence
+            speech_start = max(0, i * hop_length - int(keep_silence_sec * sample_rate))
+            
+        elif not is_speech[i] and in_speech:
+            # Potential end of speech, look ahead to check if silence continues
+            silence_length = 0
+            for j in range(i, min(len(is_speech), i + int(max_silence_sec * sample_rate / hop_length))):
+                if not is_speech[j]:
+                    silence_length += 1
+                else:
+                    break
+                    
+            if silence_length * hop_length >= max_silence_sec * sample_rate:
+                # End of speech, long enough silence detected
+                in_speech = False
+                # Calculate end sample including keep_silence
+                speech_end = min(len(audio_np), i * hop_length + int(keep_silence_sec * sample_rate))
+                speech_segments.append((speech_start, speech_end))
+    
+    # Handle the case where audio ends during speech
+    if in_speech:
+        speech_segments.append((speech_start, len(audio_np)))
+    
+    if not speech_segments:
+        logger.warning("No speech segments detected, returning original audio")
+        return audio
+    
+    # Combine speech segments with controlled silence durations
+    result = []
+    
+    # Add initial silence if the first segment doesn't start at the beginning
+    if speech_segments[0][0] > 0:
+        # Add a short leading silence (100ms)
+        silence_samples = min(int(0.1 * sample_rate), speech_segments[0][0])
+        if silence_samples > 0:
+            result.append(audio_np[speech_segments[0][0] - silence_samples:speech_segments[0][0]])
+    
+    # Process each speech segment
+    for i, (start, end) in enumerate(speech_segments):
+        # Add this speech segment
+        result.append(audio_np[start:end])
+        
+        # Add a controlled silence between segments
+        if i < len(speech_segments) - 1:
+            next_start = speech_segments[i+1][0]
+            # Calculate available silence duration
+            available_silence = next_start - end
+            
+            if available_silence > 0:
+                # Use either the actual silence (if shorter than max) or the max allowed
+                silence_duration = min(available_silence, int(max_silence_sec * sample_rate))
+                # Take the first portion of the silence - usually cleaner
+                result.append(audio_np[end:end + silence_duration])
+    
+    # Combine all parts
+    processed_audio = np.concatenate(result)
+    
+    # Log the results
+    original_duration = len(audio_np) / sample_rate
+    processed_duration = len(processed_audio) / sample_rate
+    logger.info(f"Silence removal: {original_duration:.2f}s -> {processed_duration:.2f}s ({processed_duration/original_duration*100:.1f}%)")
+    
+    # Return as tensor with original device and dtype
+    return torch.tensor(processed_audio, device=audio.device, dtype=audio.dtype)
+
+def create_high_shelf_filter(audio, sample_rate, frequency=4000, gain_db=3.0):
+    """
+    Create a high shelf filter to boost frequencies above the given frequency.
+    
+    Args:
+        audio: Audio numpy array
+        sample_rate: Sample rate in Hz
+        frequency: Shelf frequency in Hz
+        gain_db: Gain in dB for frequencies above the shelf
+        
+    Returns:
+        Filtered audio
+    """
+    # Convert gain from dB to linear
+    gain = 10 ** (gain_db / 20.0)
+    
+    # Normalized frequency (0 to 1, where 1 is Nyquist frequency)
+    normalized_freq = 2.0 * frequency / sample_rate
+    
+    # Design a high-shelf biquad filter
+    # This is a standard second-order section (SOS) implementation
+    b0 = gain
+    b1 = 0
+    b2 = 0
+    a0 = 1
+    a1 = 0
+    a2 = 0
+    
+    # Simple first-order high-shelf filter
+    alpha = np.sin(np.pi * normalized_freq) / 2 * np.sqrt((gain + 1/gain) * (1/0.5 - 1) + 2)
+    cos_w0 = np.cos(np.pi * normalized_freq)
+    
+    b0 = gain * ((gain + 1) + (gain - 1) * cos_w0 + 2 * np.sqrt(gain) * alpha)
+    b1 = -2 * gain * ((gain - 1) + (gain + 1) * cos_w0)
+    b2 = gain * ((gain + 1) + (gain - 1) * cos_w0 - 2 * np.sqrt(gain) * alpha)
+    a0 = (gain + 1) - (gain - 1) * cos_w0 + 2 * np.sqrt(gain) * alpha
+    a1 = 2 * ((gain - 1) - (gain + 1) * cos_w0)
+    a2 = (gain + 1) - (gain - 1) * cos_w0 - 2 * np.sqrt(gain) * alpha
+    
+    # Normalize coefficients
+    b = np.array([b0, b1, b2]) / a0
+    a = np.array([1.0, a1/a0, a2/a0])
+    
+    # Apply the filter
+    return signal.lfilter(b, a, audio)
+
+def enhance_audio_quality(audio: torch.Tensor, sample_rate: int) -> torch.Tensor:
+    """
+    Enhance audio quality by applying various processing techniques.
+    
+    Args:
+        audio: Audio tensor
+        sample_rate: Sample rate in Hz
+        
+    Returns:
+        Enhanced audio tensor
+    """
+    try:
+        audio_np = audio.cpu().numpy()
+        
+        # Remove DC offset
+        audio_np = audio_np - np.mean(audio_np)
+        
+        # Apply light compression to improve perceived loudness
+        # Compress by reducing peaks and increasing quieter parts slightly
+        threshold = 0.5
+        ratio = 1.5
+        attack = 0.01
+        release = 0.1
+        
+        # Simple compression algorithm
+        gain = np.ones_like(audio_np)
+        for i in range(1, len(audio_np)):
+            level = abs(audio_np[i])
+            if level > threshold:
+                gain[i] = threshold + (level - threshold) / ratio
+                gain[i] = gain[i] / level if level > 0 else 1.0
+            else:
+                gain[i] = 1.0
+            
+            # Smooth gain changes
+            gain[i] = gain[i-1] + (gain[i] - gain[i-1]) * (attack if gain[i] < gain[i-1] else release)
+        
+        audio_np = audio_np * gain
+        
+        # Apply high-shelf filter to enhance speech clarity
+        # Boost frequencies above 4000 Hz by 3 dB
+        audio_np = create_high_shelf_filter(audio_np, sample_rate, frequency=4000, gain_db=3.0)
+        
+        # Normalize to prevent clipping
+        max_val = np.max(np.abs(audio_np))
+        if max_val > 0:
+            audio_np = audio_np * 0.95 / max_val
+        
+        return torch.tensor(audio_np, device=audio.device, dtype=audio.dtype)
+        
+    except Exception as e:
+        logger.warning(f"Audio quality enhancement failed: {e}")
+        return audio

app/custom_transformer.py

@@ -0,0 +1,339 @@
+"""Custom transformer implementation for fallback."""
+import torch
+import torch.nn as nn
+import math
+import logging
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization."""
+    
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        # Calculate RMS
+        rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return self.weight * rms * x
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary positional embedding."""
+    
+    def __init__(self, dim, max_seq_len=2048, base=10000):
+        super().__init__()
+        self.dim = dim
+        self.max_seq_len = max_seq_len
+        self.base = base
+        
+        # Generate frequency tensor
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        
+        # Generate cos and sin cache
+        self._update_cos_sin_cache(max_seq_len)
+        
+    def _update_cos_sin_cache(self, max_seq_len):
+        """Update the cache of cos and sin values."""
+        self.max_seq_len = max_seq_len
+        t = torch.arange(max_seq_len, device=self.inv_freq.device)
+        
+        # Compute cos and sin at each position
+        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+        cos = freqs.cos()
+        sin = freqs.sin()
+        
+        self.register_buffer("cos_cache", cos, persistent=False)
+        self.register_buffer("sin_cache", sin, persistent=False)
+        
+    def forward(self, x, seq_len=None, pos=None):
+        # Get appropriate parts of the cache
+        if pos is not None:
+            # Handle arbitrary positions
+            cos = self.cos_cache[pos]
+            sin = self.sin_cache[pos]
+        else:
+            # Handle sequential positions
+            seq_len = x.shape[1] if seq_len is None else seq_len
+            cos = self.cos_cache[:seq_len]
+            sin = self.sin_cache[:seq_len]
+            
+        return cos, sin
+
+
+def rotate_half(x):
+    """Rotate half the dimensions of the input."""
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None):
+    """Apply rotary position embedding to q and k."""
+    if position_ids is not None:
+        # Handle arbitrary positions
+        cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
+        sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
+    else:
+        # Handle sequential positions
+        cos = cos.unsqueeze(0).unsqueeze(0)  # [1, 1, seq_len, dim]
+        sin = sin.unsqueeze(0).unsqueeze(0)  # [1, 1, seq_len, dim]
+    
+    # Apply rotation
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    
+    return q_embed, k_embed
+
+
+class CustomAttention(nn.Module):
+    """Multi-head attention with support for KV caching."""
+    
+    def __init__(self, dim, num_heads, num_kv_heads=None, dropout=0.0):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads or num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        
+        # Attention projections
+        self.q_proj = nn.Linear(dim, num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(dim, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(num_heads * self.head_dim, dim, bias=False)
+        
+        # Rotary embedding
+        self.rope = RotaryEmbedding(self.head_dim)
+        
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+        
+    def _repeat_kv(self, x):
+        """Repeat KV heads to match the number of query heads."""
+        if self.num_kv_heads == self.num_heads:
+            return x
+        
+        b, s, n_kv_head, head_dim = x.shape
+        
+        # Repeat the KV heads to match the number of query heads
+        repeats = self.num_heads // self.num_kv_heads
+        x = x.repeat_interleave(repeats, dim=2)
+        
+        return x
+        
+    def forward(self, x, mask=None, input_pos=None, kv_cache=None):
+        batch_size, seq_len, _ = x.shape
+        
+        # Project to q, k, v
+        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)  # [b, nh, s, hd]
+        k = self.k_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)  # [b, nkh, s, hd]
+        v = self.v_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)  # [b, nkh, s, hd]
+        
+        # Apply rotary embeddings
+        cos, sin = self.rope.forward(x, seq_len=seq_len, pos=input_pos)
+        q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids=input_pos)
+        
+        # Handle KV cache
+        if kv_cache is not None:
+            k_cache, v_cache = kv_cache
+            
+            if input_pos is not None:
+                # Update cache at specific positions
+                k_cache.index_copy_(2, input_pos, k)
+                v_cache.index_copy_(2, input_pos, v)
+                
+                # Use the entire cache
+                k, v = k_cache, v_cache
+        
+        # Repeat KV if needed
+        k = self._repeat_kv(k)
+        v = self._repeat_kv(v)
+        
+        # Calculate attention scores
+        attention_scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+        
+        # Apply mask if provided
+        if mask is not None:
+            attention_scores = attention_scores.masked_fill(mask == 0, -10000.0)
+            
+        # Apply softmax and dropout
+        attention_probs = self.dropout(torch.softmax(attention_scores, dim=-1))
+        
+        # Get context vector
+        context = torch.matmul(attention_probs, v)
+        
+        # Reshape and project back
+        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
+        output = self.o_proj(context)
+        
+        return output
+
+
+class FeedForward(nn.Module):
+    """Feed-forward network with GELU activation."""
+    
+    def __init__(self, dim, hidden_dim, dropout=0.0):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+        self.act = nn.GELU()
+        
+    def forward(self, x):
+        x = self.w1(x)
+        x = self.act(x)
+        x = self.dropout(x)
+        x = self.w2(x)
+        return x
+
+
+class TransformerLayer(nn.Module):
+    """A single transformer layer."""
+    
+    def __init__(
+        self, 
+        dim, 
+        num_heads, 
+        num_kv_heads=None,
+        ffn_dim=None, 
+        dropout=0.0, 
+        norm_eps=1e-5
+    ):
+        super().__init__()
+        self.norm1 = RMSNorm(dim, eps=norm_eps)
+        self.attn = CustomAttention(dim, num_heads, num_kv_heads, dropout)
+        self.norm2 = RMSNorm(dim, eps=norm_eps)
+        self.ffn = FeedForward(
+            dim, 
+            ffn_dim or 4 * dim, 
+            dropout
+        )
+        
+    def forward(self, x, mask=None, input_pos=None, kv_cache=None):
+        # Self-attention with residual
+        h = self.norm1(x)
+        h = self.attn(h, mask=mask, input_pos=input_pos, kv_cache=kv_cache)
+        x = x + h
+        
+        # FFN with residual
+        h = self.norm2(x)
+        h = self.ffn(h)
+        x = x + h
+        
+        return x
+
+
+class CustomTransformerDecoder(nn.Module):
+    """Custom transformer decoder that mimics Llama architecture."""
+    
+    def __init__(
+        self,
+        vocab_size,
+        num_layers,
+        num_heads,
+        num_kv_heads,
+        embed_dim,
+        max_seq_len,
+        intermediate_dim,
+        attn_dropout=0.0,
+        norm_eps=1e-5,
+        rope_base=10000,
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.max_seq_len = max_seq_len
+        self.embed_dim = embed_dim
+        
+        # Token embeddings
+        self.tok_embeddings = nn.Embedding(vocab_size, embed_dim)
+        
+        # Transformer layers
+        self.layers = nn.ModuleList([
+            TransformerLayer(
+                embed_dim,
+                num_heads,
+                num_kv_heads,
+                intermediate_dim,
+                attn_dropout,
+                norm_eps
+            )
+            for _ in range(num_layers)
+        ])
+        
+        # Final normalization and output projection
+        self.norm = RMSNorm(embed_dim, eps=norm_eps)
+        self.output = nn.Linear(embed_dim, vocab_size, bias=False)
+        
+        # Initialize the KV cache
+        self._kv_cache = None
+        self._has_cache = False
+        
+        logger.info(f"Initialized CustomTransformerDecoder with {num_layers} layers, {num_heads} heads, {embed_dim} dim")
+        
+    def setup_caches(self, batch_size, dtype, decoder_max_seq_len=None):
+        """Set up KV caches for inference."""
+        max_seq_len = decoder_max_seq_len or self.max_seq_len
+        device = next(self.parameters()).device
+        
+        self._kv_cache = []
+        for i, layer in enumerate(self.layers):
+            # Create a KV cache for each layer
+            k_cache = torch.zeros(
+                batch_size,
+                layer.attn.num_kv_heads,
+                max_seq_len,
+                layer.attn.head_dim,
+                device=device,
+                dtype=dtype
+            )
+            v_cache = torch.zeros(
+                batch_size,
+                layer.attn.num_kv_heads,
+                max_seq_len,
+                layer.attn.head_dim,
+                device=device,
+                dtype=dtype
+            )
+            self._kv_cache.append((k_cache, v_cache))
+        
+        self._has_cache = True
+        logger.info(f"KV caches set up for {batch_size} batches, {max_seq_len} seq length")
+        
+    def caches_are_enabled(self):
+        """Check if caches are enabled."""
+        return self._has_cache
+        
+    def reset_caches(self):
+        """Reset the KV cache to zeros."""
+        if self._has_cache and self._kv_cache:
+            for k_cache, v_cache in self._kv_cache:
+                k_cache.zero_()
+                v_cache.zero_()
+    
+    def forward(self, x, mask=None, input_pos=None):
+        batch_size, seq_len = x.shape[:2]
+        
+        # Apply embedding if input is token IDs
+        if x.dim() == 2:
+            x = self.tok_embeddings(x)
+        
+        # Apply transformer layers
+        for i, layer in enumerate(self.layers):
+            layer_cache = self._kv_cache[i] if self._has_cache else None
+            x = layer(x, mask=mask, input_pos=input_pos, kv_cache=layer_cache)
+        
+        # Apply final norm
+        x = self.norm(x)
+        
+        # Skip output projection if using Identity
+        if isinstance(self.output, nn.Identity):
+            return x
+        
+        # Apply output projection
+        logits = self.output(x)
+        
+        return logits

app/download_model.py

@@ -0,0 +1,28 @@
+"""Download CSM-1B model from Hugging Face."""
+
+import os
+import argparse
+from huggingface_hub import hf_hub_download
+
+def download_model(output_dir="models"):
+    """Download CSM-1B model from Hugging Face."""
+    print("Downloading CSM-1B model...")
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Download model
+    model_path = hf_hub_download(
+        repo_id="sesame/csm-1b",
+        filename="ckpt.pt",
+        local_dir=output_dir,
+        local_dir_use_symlinks=False
+    )
+    
+    print(f"Model downloaded to {model_path}")
+    return model_path
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Download CSM-1B model")
+    parser.add_argument("--output", type=str, default="models", help="Output directory")
+    args = parser.parse_args()
+    
+    download_model(args.output)

app/generator.py

@@ -0,0 +1,834 @@
+# Updated generator.py with proper function order
+from dataclasses import dataclass
+from typing import List, Tuple
+import torch
+import torchaudio
+import logging
+import os
+from huggingface_hub import hf_hub_download
+from transformers import AutoTokenizer
+from tokenizers.processors import TemplateProcessing
+from app.models import Segment
+from app.text_normalizer import clean_text_for_tts
+from app.text_normalizer import TextNormalizer
+
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+# Import the CSM watermarking code
+try:
+    from app.watermarking import CSM_1B_GH_WATERMARK, load_watermarker, watermark
+except ImportError:
+    # Define stubs for watermarking if the module is not available
+    CSM_1B_GH_WATERMARK = "CSM1B"
+    def load_watermarker(device="cpu"):
+        return None
+    def watermark(watermarker, audio, sample_rate, key):
+        return audio, sample_rate
+
+def load_llama3_tokenizer():
+    """
+    Load tokenizer for Llama 3.2, using unsloth's open version
+    instead of the gated meta-llama version.
+    """
+    try:
+        # Use the unsloth version which is not gated
+        tokenizer_name = "unsloth/Llama-3.2-1B"
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+        bos = tokenizer.bos_token
+        eos = tokenizer.eos_token
+        tokenizer._tokenizer.post_processor = TemplateProcessing(
+            single=f"{bos}:0 $A:0 {eos}:0",
+            pair=f"{bos}:0 $A:0 {eos}:0 {bos}:1 $B:1 {eos}:1",
+            special_tokens=[(f"{bos}", tokenizer.bos_token_id), (f"{eos}", tokenizer.eos_token_id)],
+        )
+        logger.info("Successfully loaded tokenizer from unsloth/Llama-3.2-1B")
+        return tokenizer
+    except Exception as e:
+        logger.error(f"Error loading tokenizer from unsloth: {e}")
+        # Fallback to a simpler tokenizer if needed
+        try:
+            from transformers import GPT2Tokenizer
+            logger.warning("Falling back to GPT2Tokenizer")
+            tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+            tokenizer.pad_token = tokenizer.eos_token
+            return tokenizer
+        except Exception as fallback_e:
+            logger.error(f"Fallback tokenizer also failed: {fallback_e}")
+            raise RuntimeError("Could not load any suitable tokenizer")
+
+class Generator:
+    """Generator class for CSM-1B model."""
+    def __init__(self, model):
+        """Initialize generator with model."""
+        self._model = model
+        self._model.setup_caches(1)
+        self._text_tokenizer = load_llama3_tokenizer()
+        device = next(model.parameters()).device
+        # Load Mimi codec for audio tokenization
+        try:
+            logger.info("Loading Mimi audio codec...")
+            from huggingface_hub import hf_hub_download
+            # First try to import from moshi
+            try:
+                from moshi.models import loaders
+                DEFAULT_REPO = loaders.DEFAULT_REPO
+                MIMI_NAME = loaders.MIMI_NAME
+                get_mimi = loaders.get_mimi
+            except ImportError:
+                logger.warning("moshi.models.loaders not found, using fallback")
+                # Fallback values if moshi.models.loaders is not available
+                DEFAULT_REPO = "kyutai/mimi"
+                MIMI_NAME = "mimi-december.pt"
+                # Fallback function to load mimi
+                def get_mimi(checkpoint_path, device):
+                    from moshi.models.vqvae_model import MiMiModule
+                    checkpoint = torch.load(checkpoint_path, map_location=device)
+                    model = MiMiModule.init_from_checkpoint(checkpoint, device=device)
+                    return model
+            mimi_weight = hf_hub_download(DEFAULT_REPO, MIMI_NAME)
+            mimi = get_mimi(mimi_weight, device=device)
+            mimi.set_num_codebooks(32)
+            self._audio_tokenizer = mimi
+            self.sample_rate = mimi.sample_rate
+            logger.info(f"Mimi codec loaded successfully with sample rate {self.sample_rate}")
+        except Exception as e:
+            logger.error(f"Error loading Mimi codec: {e}")
+            self._audio_tokenizer = None
+            self.sample_rate = 24000  # Default sample rate
+            logger.warning(f"Using fallback sample rate: {self.sample_rate}")
+            raise RuntimeError(f"Failed to load Mimi codec: {e}")
+        try:
+            self._watermarker = load_watermarker(device=device)
+            logger.info("Watermarker loaded successfully")
+        except Exception as e:
+            logger.warning(f"Error loading watermarker: {e}. Watermarking will be disabled.")
+            self._watermarker = None
+            
+        self.device = device
+        # Optimize for CUDA throughput
+        if torch.cuda.is_available():
+            torch.backends.cudnn.benchmark = True
+            torch.cuda.empty_cache()
+            logger.info("CUDA optimizations enabled")
+            
+    def _tokenize_text_segment(self, text: str, speaker: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Tokenize a text segment."""
+        frame_tokens = []
+        frame_masks = []
+        # Strip any voice instructions in square brackets to avoid them being read out
+        text = self._clean_text_input(text)
+        text_tokens = self._text_tokenizer.encode(f"[{speaker}]{text}")
+        text_frame = torch.zeros(len(text_tokens), 33).long()
+        text_frame_mask = torch.zeros(len(text_tokens), 33).bool()
+        text_frame[:, -1] = torch.tensor(text_tokens)
+        text_frame_mask[:, -1] = True
+        frame_tokens.append(text_frame.to(self.device))
+        frame_masks.append(text_frame_mask.to(self.device))
+        return torch.cat(frame_tokens, dim=0), torch.cat(frame_masks, dim=0)
+    
+    def _clean_text_input(self, text: str) -> str:
+        """Clean and normalize text for TTS."""
+        return clean_text_for_tts(text)
+    
+    def _tokenize_audio(self, audio: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Tokenize audio."""
+        if self._audio_tokenizer is None:
+            raise RuntimeError("Audio tokenizer not initialized")
+        frame_tokens = []
+        frame_masks = []
+        # (K, T)
+        audio = audio.to(self.device)
+        audio_tokens = self._audio_tokenizer.encode(audio.unsqueeze(0).unsqueeze(0))[0]
+        # add EOS frame
+        eos_frame = torch.zeros(audio_tokens.size(0), 1).to(self.device)
+        audio_tokens = torch.cat([audio_tokens, eos_frame], dim=1)
+        audio_frame = torch.zeros(audio_tokens.size(1), 33).long().to(self.device)
+        audio_frame_mask = torch.zeros(audio_tokens.size(1), 33).bool().to(self.device)
+        audio_frame[:, :-1] = audio_tokens.transpose(0, 1)
+        audio_frame_mask[:, :-1] = True
+        frame_tokens.append(audio_frame)
+        frame_masks.append(audio_frame_mask)
+        return torch.cat(frame_tokens, dim=0), torch.cat(frame_masks, dim=0)
+    
+    def _tokenize_segment(self, segment: Segment) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Tokenize a segment of text and audio."""
+        text_tokens, text_masks = self._tokenize_text_segment(segment.text, segment.speaker)
+        audio_tokens, audio_masks = self._tokenize_audio(segment.audio)
+        return torch.cat([text_tokens, audio_tokens], dim=0), torch.cat([text_masks, audio_masks], dim=0)
+
+    def generate_quick(
+        self,
+        text: str,
+        speaker: int,
+        context: List[Segment],
+        max_audio_length_ms: float = 2000,  # Short for quick generation
+        temperature: float = 0.7,  # Lower for more predictable output
+        topk: int = 20,  # Lower for faster beam selection
+    ) -> torch.Tensor:
+        """Generate audio quickly for real-time streaming."""
+        # Similar to generate() but optimized for speed
+        self._model.reset_caches()
+        
+        # Convert max_audio_length_ms to frames - limit for faster generation
+        max_audio_frames = min(int(max_audio_length_ms / 80), 128)  # Smaller limit
+        
+        # Process text
+        cleaned_text = clean_text_for_tts(text)
+        
+        # Prepare tokens
+        tokens, tokens_mask = [], []
+        # Add context segments (limited to 1 for speed)
+        if context:
+            segment_tokens, segment_tokens_mask = self._tokenize_segment(context[0])
+            tokens.append(segment_tokens)
+            tokens_mask.append(segment_tokens_mask)
+        # Add text tokens
+        gen_segment_tokens, gen_segment_tokens_mask = self._tokenize_text_segment(cleaned_text, speaker)
+        tokens.append(gen_segment_tokens)
+        tokens_mask.append(gen_segment_tokens_mask)
+        
+        prompt_tokens = torch.cat(tokens, dim=0).long().to(self.device)
+        prompt_tokens_mask = torch.cat(tokens_mask, dim=0).bool().to(self.device)
+        
+        # Generate with larger batch size for fewer iterations
+        curr_tokens = prompt_tokens.unsqueeze(0)
+        curr_tokens_mask = prompt_tokens_mask.unsqueeze(0)
+        curr_pos = torch.arange(0, prompt_tokens.size(0)).unsqueeze(0).long().to(self.device)
+        
+        # Use larger batch size
+        batch_size = 64  # Generate more frames at once
+        all_samples = []
+        for start_idx in range(0, max_audio_frames, batch_size):
+            end_idx = min(start_idx + batch_size, max_audio_frames)
+            batch_frames = end_idx - start_idx
+            samples_batch = []
+            for i in range(batch_frames):
+                sample = self._model.generate_frame(curr_tokens, curr_tokens_mask, curr_pos, temperature, topk)
+                samples_batch.append(sample)
+                if torch.all(sample == 0):
+                    break
+                curr_tokens = torch.cat([sample, torch.zeros(1, 1).long().to(self.device)], dim=1).unsqueeze(1)
+                curr_tokens_mask = torch.cat(
+                    [torch.ones_like(sample).bool(), torch.zeros(1, 1).bool().to(self.device)], dim=1
+                ).unsqueeze(1)
+                curr_pos = curr_pos[:, -1:] + 1
+            all_samples.extend(samples_batch)
+            if len(samples_batch) < batch_frames:
+                break
+        
+        if not all_samples:
+            return torch.zeros(10, device=self.device)  # Return short empty audio
+            
+        # Decode audio
+        audio = self._audio_tokenizer.decode(torch.stack(all_samples).permute(1, 2, 0)).squeeze(0).squeeze(0)
+        return audio
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        text: str,
+        speaker: int,
+        context: List[Segment],
+        max_audio_length_ms: float = 90_000,
+        temperature: float = 0.9,
+        topk: int = 50,
+    ) -> torch.Tensor:
+        """Generate audio from text."""
+        if self._audio_tokenizer is None:
+            raise RuntimeError("Audio tokenizer not initialized")
+        
+        # Start timing
+        start_time = torch.cuda.Event(enable_timing=True)
+        end_time = torch.cuda.Event(enable_timing=True)
+        start_time.record()
+            
+        self._model.reset_caches()
+        
+        # Convert max_audio_length_ms to frames - this controls the maximum generation length
+        max_audio_frames = min(int(max_audio_length_ms / 80), 1024)  # Limit to reasonable size
+        max_seq_len = 2048 - max_audio_frames
+        
+        # Check if text is long and should be split
+        if len(text) > 200:
+            logger.info(f"Long text detected ({len(text)} chars), processing in segments")
+            sentences = TextNormalizer.split_into_sentences(text)
+            logger.info(f"Split into {len(sentences)} segments")
+            
+            # Process sentences individually and concatenate the results
+            all_audio_segments = []
+            
+            # Use the first sentence to establish voice
+            first_sentence = sentences[0]
+            cleaned_text = clean_text_for_tts(first_sentence)
+            
+            # Generate the first segment
+            tokens, tokens_mask = [], []
+            
+            # Add context segments for the first sentence
+            for segment in context:
+                segment_tokens, segment_tokens_mask = self._tokenize_segment(segment)
+                tokens.append(segment_tokens)
+                tokens_mask.append(segment_tokens_mask)
+            
+            # Add first sentence tokens
+            gen_segment_tokens, gen_segment_tokens_mask = self._tokenize_text_segment(cleaned_text, speaker)
+            tokens.append(gen_segment_tokens)
+            tokens_mask.append(gen_segment_tokens_mask)
+            
+            prompt_tokens = torch.cat(tokens, dim=0).long().to(self.device)
+            prompt_tokens_mask = torch.cat(tokens_mask, dim=0).bool().to(self.device)
+            
+            # Check context size and truncate if needed
+            if prompt_tokens.size(0) >= max_seq_len:
+                logger.warning(f"Inputs too long ({prompt_tokens.size(0)} tokens), truncating to {max_seq_len - 50}")
+                prompt_tokens = prompt_tokens[-max_seq_len+50:]
+                prompt_tokens_mask = prompt_tokens_mask[-max_seq_len+50:]
+            
+            # Generate first sentence audio
+            curr_tokens = prompt_tokens.unsqueeze(0)
+            curr_tokens_mask = prompt_tokens_mask.unsqueeze(0)
+            curr_pos = torch.arange(0, prompt_tokens.size(0)).unsqueeze(0).long().to(self.device)
+            
+            # Generate first segment
+            first_segment_samples = []
+            for start_idx in range(0, max_audio_frames, 32):
+                end_idx = min(start_idx + 32, max_audio_frames)
+                batch_frames = end_idx - start_idx
+                samples_batch = []
+                
+                for i in range(batch_frames):
+                    sample = self._model.generate_frame(curr_tokens, curr_tokens_mask, curr_pos, temperature, topk)
+                    samples_batch.append(sample)
+                    
+                    if torch.all(sample == 0):
+                        break
+                    
+                    curr_tokens = torch.cat([sample, torch.zeros(1, 1).long().to(self.device)], dim=1).unsqueeze(1)
+                    curr_tokens_mask = torch.cat(
+                        [torch.ones_like(sample).bool(), torch.zeros(1, 1).bool().to(self.device)], dim=1
+                    ).unsqueeze(1)
+                    curr_pos = curr_pos[:, -1:] + 1
+                
+                first_segment_samples.extend(samples_batch)
+                
+                if len(samples_batch) < batch_frames:
+                    break
+            
+            if not first_segment_samples:
+                raise RuntimeError("No audio generated for first segment")
+            
+            # Decode first segment
+            first_segment_audio = self._audio_tokenizer.decode(
+                torch.stack(first_segment_samples).permute(1, 2, 0)
+            ).squeeze(0).squeeze(0)
+            
+            all_audio_segments.append(first_segment_audio)
+            
+            # Now process remaining sentences using the first as context
+            for i, sentence in enumerate(sentences[1:], 1):
+                logger.info(f"Generating segment {i+1}/{len(sentences)}")
+                cleaned_text = clean_text_for_tts(sentence)
+                
+                # Create a context segment from the previous generation
+                prev_segment = Segment(
+                    speaker=speaker,
+                    text=sentences[i-1],
+                    audio=all_audio_segments[-1]
+                )
+                
+                # Generate with this segment as context
+                segment_tokens, segment_tokens_mask = [], []
+                segment_tokens.append(self._tokenize_segment(prev_segment)[0])
+                segment_tokens_mask.append(self._tokenize_segment(prev_segment)[1])
+                
+                # Add current segment tokens
+                current_tokens, current_tokens_mask = self._tokenize_text_segment(cleaned_text, speaker)
+                segment_tokens.append(current_tokens)
+                segment_tokens_mask.append(current_tokens_mask)
+                
+                segment_prompt_tokens = torch.cat(segment_tokens, dim=0).long().to(self.device)
+                segment_prompt_tokens_mask = torch.cat(segment_tokens_mask, dim=0).bool().to(self.device)
+                
+                # Check length and truncate if needed
+                if segment_prompt_tokens.size(0) >= max_seq_len:
+                    segment_prompt_tokens = segment_prompt_tokens[-max_seq_len+50:]
+                    segment_prompt_tokens_mask = segment_prompt_tokens_mask[-max_seq_len+50:]
+                
+                # Generate audio for this segment
+                curr_tokens = segment_prompt_tokens.unsqueeze(0)
+                curr_tokens_mask = segment_prompt_tokens_mask.unsqueeze(0)
+                curr_pos = torch.arange(0, segment_prompt_tokens.size(0)).unsqueeze(0).long().to(self.device)
+                
+                # Generate segment
+                segment_samples = []
+                for start_idx in range(0, max_audio_frames, 32):
+                    end_idx = min(start_idx + 32, max_audio_frames)
+                    batch_frames = end_idx - start_idx
+                    samples_batch = []
+                    
+                    for i in range(batch_frames):
+                        sample = self._model.generate_frame(curr_tokens, curr_tokens_mask, curr_pos, temperature, topk)
+                        samples_batch.append(sample)
+                        
+                        if torch.all(sample == 0):
+                            break
+                        
+                        curr_tokens = torch.cat([sample, torch.zeros(1, 1).long().to(self.device)], dim=1).unsqueeze(1)
+                        curr_tokens_mask = torch.cat(
+                            [torch.ones_like(sample).bool(), torch.zeros(1, 1).bool().to(self.device)], dim=1
+                        ).unsqueeze(1)
+                        curr_pos = curr_pos[:, -1:] + 1
+                    
+                    segment_samples.extend(samples_batch)
+                    
+                    if len(samples_batch) < batch_frames:
+                        break
+                
+                if not segment_samples:
+                    logger.warning(f"No audio generated for segment {i+1}")
+                    continue
+                
+                # Decode segment
+                segment_audio = self._audio_tokenizer.decode(
+                    torch.stack(segment_samples).permute(1, 2, 0)
+                ).squeeze(0).squeeze(0)
+                
+                all_audio_segments.append(segment_audio)
+            
+            # Combine all segments with small pauses
+            pause_samples = int(0.3 * self.sample_rate)  # 300ms pause
+            pause = torch.zeros(pause_samples, device=self.device)
+            
+            audio_parts = []
+            for i, segment_audio in enumerate(all_audio_segments):
+                audio_parts.append(segment_audio)
+                if i < len(all_audio_segments) - 1:
+                    audio_parts.append(pause)
+            
+            audio = torch.cat(audio_parts)
+            logger.info(f"Combined {len(all_audio_segments)} segments into final audio")
+        
+        else:
+            # For shorter text, standard processing
+            tokens, tokens_mask = [], []
+            
+            # Add context segments
+            for segment in context:
+                segment_tokens, segment_tokens_mask = self._tokenize_segment(segment)
+                tokens.append(segment_tokens)
+                tokens_mask.append(segment_tokens_mask)
+            
+            # Process text
+            cleaned_text = clean_text_for_tts(text)
+            gen_segment_tokens, gen_segment_tokens_mask = self._tokenize_text_segment(cleaned_text, speaker)
+            tokens.append(gen_segment_tokens)
+            tokens_mask.append(gen_segment_tokens_mask)
+            
+            prompt_tokens = torch.cat(tokens, dim=0).long().to(self.device)
+            prompt_tokens_mask = torch.cat(tokens_mask, dim=0).bool().to(self.device)
+            
+            # Check context size
+            if prompt_tokens.size(0) >= max_seq_len:
+                logger.warning(f"Inputs too long ({prompt_tokens.size(0)} tokens), truncating to {max_seq_len - 50}")
+                prompt_tokens = prompt_tokens[-max_seq_len+50:]
+                prompt_tokens_mask = prompt_tokens_mask[-max_seq_len+50:]
+            
+            # Generate audio - optimized batch generation
+            curr_tokens = prompt_tokens.unsqueeze(0)
+            curr_tokens_mask = prompt_tokens_mask.unsqueeze(0)
+            curr_pos = torch.arange(0, prompt_tokens.size(0)).unsqueeze(0).long().to(self.device)
+            
+            # Using optimized batch generation
+            batch_size = 32  # Generate this many frames at once 
+            all_samples = []
+            
+            for start_idx in range(0, max_audio_frames, batch_size):
+                end_idx = min(start_idx + batch_size, max_audio_frames)
+                batch_frames = end_idx - start_idx
+                
+                samples_batch = []
+                
+                for i in range(batch_frames):
+                    sample = self._model.generate_frame(curr_tokens, curr_tokens_mask, curr_pos, temperature, topk)
+                    samples_batch.append(sample)
+                    
+                    if torch.all(sample == 0):
+                        break
+                    
+                    curr_tokens = torch.cat([sample, torch.zeros(1, 1).long().to(self.device)], dim=1).unsqueeze(1)
+                    curr_tokens_mask = torch.cat(
+                        [torch.ones_like(sample).bool(), torch.zeros(1, 1).bool().to(self.device)], dim=1
+                    ).unsqueeze(1)
+                    curr_pos = curr_pos[:, -1:] + 1
+                
+                all_samples.extend(samples_batch)
+                
+                if len(samples_batch) < batch_frames:
+                    logger.info(f"Early stopping at frame {start_idx + len(samples_batch)}/{max_audio_frames}")
+                    break
+            
+            if not all_samples:
+                raise RuntimeError("No audio generated - model produced empty output")
+            
+            # Decode audio
+            audio = self._audio_tokenizer.decode(torch.stack(all_samples).permute(1, 2, 0)).squeeze(0).squeeze(0)
+        
+        # Apply watermark
+        if self._watermarker is not None:
+            try:
+                audio, wm_sample_rate = watermark(self._watermarker, audio, self.sample_rate, CSM_1B_GH_WATERMARK)
+                audio = torchaudio.functional.resample(audio, orig_freq=wm_sample_rate, new_freq=self.sample_rate)
+            except Exception as e:
+                logger.warning(f"Error applying watermark: {e}. Continuing without watermark.")
+        
+        # Record execution time
+        end_time.record()
+        torch.cuda.synchronize()
+        execution_ms = start_time.elapsed_time(end_time)
+        audio_length_ms = (audio.shape[0] / self.sample_rate) * 1000
+        
+        # Calculate real-time factor (RTF)
+        rtf = execution_ms / audio_length_ms
+        logger.info(f"Audio generated in {execution_ms:.2f}ms, length: {audio_length_ms:.2f}ms, RTF: {rtf:.2f}x")
+        
+        return audio
+
+# Define helper functions for multi-GPU support
+def _manual_device_map(model, state_dict, strategy="balanced"):
+    """Apply manual device mapping for multi-GPU setups.
+    
+    Args:
+        model: The model to map
+        state_dict: Model state dict
+        strategy: Mapping strategy ('balanced', 'sequential')
+        
+    Returns:
+        Model with weights distributed across GPUs
+    """
+    num_gpus = torch.cuda.device_count()
+    if num_gpus <= 1:
+        # No need for mapping with single GPU
+        model.load_state_dict(state_dict)
+        model = model.to("cuda")
+        return model
+    
+    logger.info(f"Applying manual {strategy} device mapping across {num_gpus} GPUs")
+    
+    # Get all layer names from state dict
+    layer_names = [name for name in state_dict.keys() if "layers" in name]
+    backbone_layers = [name for name in layer_names if "backbone.layers" in name]
+    decoder_layers = [name for name in layer_names if "decoder.layers" in name]
+    
+    # Count number of backbone and decoder layers
+    backbone_layer_indices = set()
+    for name in backbone_layers:
+        parts = name.split('.')
+        if len(parts) > 2:
+            try:
+                backbone_layer_indices.add(int(parts[2]))
+            except ValueError:
+                pass
+    
+    decoder_layer_indices = set()
+    for name in decoder_layers:
+        parts = name.split('.')
+        if len(parts) > 2:
+            try:
+                decoder_layer_indices.add(int(parts[2]))
+            except ValueError:
+                pass
+    
+    num_backbone_layers = len(backbone_layer_indices)
+    num_decoder_layers = len(decoder_layer_indices)
+    
+    # Create device map
+    device_map = {}
+    
+    if strategy == "balanced":
+        # Distribute layers evenly across GPUs
+        layers_per_gpu = (num_backbone_layers + num_decoder_layers) // num_gpus
+        remainder = (num_backbone_layers + num_decoder_layers) % num_gpus
+        
+        # Assign backbone layers
+        for i in backbone_layer_indices:
+            gpu_idx = min(i // layers_per_gpu, num_gpus - 1)
+            device_map[f"backbone.layers.{i}"] = f"cuda:{gpu_idx}"
+        
+        # Assign decoder layers
+        for i in decoder_layer_indices:
+            gpu_idx = min((i + num_backbone_layers) // layers_per_gpu, num_gpus - 1)
+            device_map[f"decoder.layers.{i}"] = f"cuda:{gpu_idx}"
+            
+    elif strategy == "sequential":
+        # Fill each GPU sequentially
+        # Backbone layers on first GPU(s)
+        backbone_per_gpu = max(1, num_backbone_layers // ((num_gpus + 1) // 2))
+        for i in backbone_layer_indices:
+            gpu_idx = min(i // backbone_per_gpu, (num_gpus + 1) // 2 - 1)
+            device_map[f"backbone.layers.{i}"] = f"cuda:{gpu_idx}"
+        
+        # Decoder layers on remaining GPU(s)
+        decoder_per_gpu = max(1, num_decoder_layers // (num_gpus - (num_gpus + 1) // 2 + 1))
+        for i in decoder_layer_indices:
+            gpu_idx = min(i // decoder_per_gpu + (num_gpus + 1) // 2 - 1, num_gpus - 1)
+            device_map[f"decoder.layers.{i}"] = f"cuda:{gpu_idx}"
+    
+    # Assign embeddings and other components
+    device_map["text_embeddings"] = "cuda:0"
+    device_map["audio_embeddings"] = "cuda:0"
+    device_map["projection"] = "cuda:0"
+    device_map["codebook0_head"] = "cuda:0"
+    device_map["audio_head"] = "cuda:0"
+    
+    # Load state dict with device mapping
+    model.load_state_dict(state_dict)
+    
+    # Move model parts to assigned devices
+    for name, device in device_map.items():
+        if "backbone.layers" in name:
+            layer_idx = int(name.split('.')[-1])
+            if hasattr(model.backbone, 'layers') and layer_idx < len(model.backbone.layers):
+                model.backbone.layers[layer_idx] = model.backbone.layers[layer_idx].to(device)
+        elif "decoder.layers" in name:
+            layer_idx = int(name.split('.')[-1])
+            if hasattr(model.decoder, 'layers') and layer_idx < len(model.decoder.layers):
+                model.decoder.layers[layer_idx] = model.decoder.layers[layer_idx].to(device)
+        elif hasattr(model, name):
+            setattr(model, name, getattr(model, name).to(device))
+    
+    logger.info(f"Model distributed across GPUs with {strategy} strategy")
+    return model
+
+def load_csm_1b(ckpt_path: str = "ckpt.pt", device: str = "cuda", device_map: str = None) -> Generator:
+    """Load CSM-1B model and create generator with performance optimizations.
+    
+    Args:
+        ckpt_path: Path to model checkpoint
+        device: Device to load model on ('cuda', 'cpu', or specific CUDA device)
+        device_map: Optional device mapping strategy ('auto', 'balanced', 'sequential', or None)
+    
+    Returns:
+        Generator instance with optimized settings
+    """
+    try:
+        # Import models module for CSM
+        from app.torchtune_models import Model, ModelArgs
+        
+        # Create model
+        model_args = ModelArgs(
+            backbone_flavor="llama-1B",
+            decoder_flavor="llama-100M",
+            text_vocab_size=128256,
+            audio_vocab_size=2051,
+            audio_num_codebooks=32,
+        )
+        
+        # Load model
+        logger.info(f"Loading CSM-1B model from {ckpt_path} with device={device}, device_map={device_map}")
+        
+        # Check for CUDA availability
+        cuda_available = device == "cuda" and torch.cuda.is_available()
+        
+        # Set up torch for optimized inference
+        if cuda_available:
+            # Check if we should enable TF32 (faster but slightly less precise)
+            enable_tf32 = os.environ.get("ENABLE_TF32", "true").lower() == "true"
+            if enable_tf32:
+                logger.info("Enabling TF32 for faster matrix multiplications")
+                torch.backends.cuda.matmul.allow_tf32 = True
+                torch.backends.cudnn.allow_tf32 = True
+            
+            # Check for available precision modes
+            use_bfloat16 = torch.cuda.is_bf16_supported()
+            use_float16 = not use_bfloat16 and torch.cuda.is_available()  # Fallback to float16
+            
+            if use_bfloat16:
+                dtype = torch.bfloat16
+                logger.info("Using bfloat16 precision for faster inference")
+            elif use_float16:
+                dtype = torch.float16
+                logger.info("Using float16 precision for faster inference")
+            else:
+                dtype = torch.float32
+                logger.info("Using float32 precision (mixed precision not available)")
+            
+            # Enable Flash Attention if available
+            try:
+                import flash_attn
+                if os.environ.get("ENABLE_FLASH_ATTN", "true").lower() == "true":
+                    logger.info("Flash Attention detected - enabling for faster attention")
+                    os.environ["PYTORCH_FLASH_ATTENTION_ENABLED"] = "1"
+            except ImportError:
+                logger.info("Flash Attention not available (install flash-attn for faster inference)")
+        else:
+            # CPU-only mode
+            dtype = torch.float32
+            logger.info("Using CPU mode with float32 precision")
+        
+        # Check for quantization
+        enable_quantization = os.environ.get("ENABLE_QUANTIZATION", "false").lower() == "true"
+        is_quantized = False
+        
+        # Check for multi-GPU setup
+        if device_map and torch.cuda.device_count() > 1:
+            logger.info(f"Using device_map={device_map} across {torch.cuda.device_count()} GPUs")
+            
+            # Create model with device map
+            model = Model(model_args)
+            
+            # Load state dict
+            state_dict = torch.load(ckpt_path, map_location='cpu')
+            
+            # Try quantization before device mapping if enabled
+            if enable_quantization and cuda_available:
+                try:
+                    from bitsandbytes.nn import Linear8bitLt
+                    
+                    def replace_with_8bit(model):
+                        """Replace linear layers with 8-bit quantized versions"""
+                        for name, module in model.named_modules():
+                            if isinstance(module, torch.nn.Linear) and module.out_features > 256:
+                                parent_name = name.rsplit('.', 1)[0] if '.' in name else ''
+                                parent = model
+                                if parent_name:
+                                    for attr in parent_name.split('.'):
+                                        parent = getattr(parent, attr)
+                                child_name = name.rsplit('.', 1)[1] if '.' in name else name
+                                setattr(parent, child_name, Linear8bitLt.from_float(module))
+                        return model
+                    
+                    logger.info("Applying 8-bit quantization to linear layers")
+                    model = replace_with_8bit(model)
+                    is_quantized = True
+                except ImportError:
+                    logger.warning("bitsandbytes not available, skipping quantization")
+            
+            # Apply device mapping
+            if device_map == "auto":
+                # Use accelerate for automatic device mapping
+                try:
+                    from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+                    
+                    # Initialize empty model
+                    with init_empty_weights():
+                        empty_model = Model(model_args)
+                    
+                    # Load and dispatch model across GPUs
+                    model = load_checkpoint_and_dispatch(
+                        empty_model, 
+                        ckpt_path, 
+                        device_map="auto",
+                        no_split_module_classes=["TransformerLayer"],
+                        # Offload CPU if very large model
+                        offload_folder="offload" if os.environ.get("OFFLOAD_TO_CPU", "false").lower() == "true" else None
+                    )
+                    logger.info("Model loaded with automatic device mapping")
+                except ImportError:
+                    logger.warning("accelerate package not found, falling back to manual device mapping")
+                    model = _manual_device_map(model, state_dict, "balanced")
+                except Exception as mapping_error:
+                    logger.error(f"Auto device mapping failed: {mapping_error}, falling back to manual")
+                    model = _manual_device_map(model, state_dict, "balanced")
+            else:
+                # Manual device mapping
+                model = _manual_device_map(model, state_dict, device_map or "balanced")
+        else:
+            # Single GPU or CPU setup
+            
+            # Try quantization before loading if enabled (GPU only)
+            if enable_quantization and cuda_available and not is_quantized:
+                try:
+                    # First load to CPU for quantization
+                    model = Model(model_args).to("cpu")
+                    state_dict = torch.load(ckpt_path, map_location="cpu")
+                    model.load_state_dict(state_dict)
+                    
+                    from bitsandbytes.nn import Linear8bitLt
+                    
+                    def replace_with_8bit(model):
+                        """Replace linear layers with 8-bit quantized versions"""
+                        for name, module in model.named_modules():
+                            if isinstance(module, torch.nn.Linear) and module.out_features > 256:
+                                parent_name = name.rsplit('.', 1)[0] if '.' in name else ''
+                                parent = model
+                                if parent_name:
+                                    for attr in parent_name.split('.'):
+                                        parent = getattr(parent, attr)
+                                child_name = name.rsplit('.', 1)[1] if '.' in name else name
+                                setattr(parent, child_name, Linear8bitLt.from_float(module))
+                        return model
+                    
+                    logger.info("Applying 8-bit quantization to linear layers")
+                    model = replace_with_8bit(model)
+                    model = model.to(device=device)
+                    is_quantized = True
+                except ImportError:
+                    logger.warning("bitsandbytes not available, loading without quantization")
+                    # Load the standard way
+                    model = Model(model_args).to(device=device, dtype=dtype)
+                    state_dict = torch.load(ckpt_path, map_location=device)
+                    model.load_state_dict(state_dict)
+                except Exception as quant_error:
+                    logger.error(f"Quantization failed: {quant_error}, loading without quantization")
+                    # Load the standard way
+                    model = Model(model_args).to(device=device, dtype=dtype)
+                    state_dict = torch.load(ckpt_path, map_location=device)
+                    model.load_state_dict(state_dict)
+            else:
+                # Standard load without quantization
+                model = Model(model_args).to(device=device, dtype=dtype)
+                state_dict = torch.load(ckpt_path, map_location=device)
+                model.load_state_dict(state_dict)
+        
+        # Apply torch.compile if available (PyTorch 2.0+)
+        compile_mode = os.environ.get("TORCH_COMPILE_MODE", "none")
+        if hasattr(torch, 'compile') and compile_mode != "none" and cuda_available:
+            try:
+                logger.info(f"Using torch.compile with mode '{compile_mode}' for faster inference")
+                if compile_mode == "default":
+                    model = torch.compile(model)
+                else:
+                    model = torch.compile(model, mode=compile_mode)
+            except Exception as compile_error:
+                logger.warning(f"Torch compile failed (requires PyTorch 2.0+): {compile_error}")
+        
+        # Try to optimize CUDA graphs for faster inference (advanced)
+        use_cuda_graphs = os.environ.get("USE_CUDA_GRAPHS", "false").lower() == "true"
+        if use_cuda_graphs and cuda_available and hasattr(torch.cuda, 'CUDAGraph'):
+            try:
+                logger.info("Setting up CUDA graphs for repeated inference patterns")
+                # This requires custom integration inside the model's forward method
+                # Just flagging that CUDA graphs should be used
+                model.use_cuda_graphs = True
+            except Exception as cuda_graph_error:
+                logger.warning(f"CUDA graphs setup failed: {cuda_graph_error}")
+                model.use_cuda_graphs = False
+        
+        # Set optimal settings for CUDA context
+        if cuda_available:
+            # Set benchmark mode for hardware-specific optimizations
+            torch.backends.cudnn.benchmark = True
+            # Clean up CUDA cache before creating generator
+            torch.cuda.empty_cache()
+            # Ensure all CUDA work is completed to avoid launch delays
+            torch.cuda.synchronize()
+        
+        # Create generator
+        logger.info("Creating generator with optimized settings")
+        generator = Generator(model)
+        
+        # Log memory usage if on CUDA
+        if cuda_available:
+            memory_allocated = torch.cuda.memory_allocated() / (1024**3)
+            memory_reserved = torch.cuda.memory_reserved() / (1024**3)
+            logger.info(f"Model loaded, CUDA memory: {memory_allocated:.2f}GB allocated, {memory_reserved:.2f}GB reserved")
+        
+        logger.info(f"Generator created successfully: precision={dtype}, quantized={is_quantized}")
+        return generator
+    except Exception as e:
+        logger.error(f"Failed to load CSM-1B model: {e}")
+        import traceback
+        logger.error(traceback.format_exc())
+        raise

app/main.py

@@ -0,0 +1,647 @@
+"""
+CSM-1B TTS API main application.
+Provides an OpenAI-compatible API for the CSM-1B text-to-speech model.
+"""
+import os
+import time
+import tempfile
+import logging
+from logging.handlers import RotatingFileHandler
+import traceback
+import asyncio
+import glob
+import torch
+import uvicorn
+from contextlib import asynccontextmanager
+from fastapi import FastAPI, Depends, HTTPException, Request, Response
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import RedirectResponse, FileResponse
+from fastapi.staticfiles import StaticFiles
+from app.api.routes import router as api_router
+
+# Setup logging
+os.makedirs("logs", exist_ok=True)
+log_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+
+# Console handler
+console_handler = logging.StreamHandler()
+console_handler.setFormatter(logging.Formatter(log_format))
+
+# File handler
+file_handler = RotatingFileHandler(
+    "logs/csm_tts_api.log", 
+    maxBytes=10*1024*1024,  # 10MB
+    backupCount=5
+)
+file_handler.setFormatter(logging.Formatter(log_format))
+
+# Configure root logger
+logging.basicConfig(
+    level=logging.INFO,
+    format=log_format,
+    handlers=[console_handler, file_handler]
+)
+logger = logging.getLogger(__name__)
+logger.info("Starting CSM-1B TTS API")
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Application lifespan manager for startup and shutdown events."""
+    # STARTUP EVENT
+    logger.info("Starting application initialization")
+    app.state.startup_time = time.time()
+    app.state.generator = None  # Will be populated later if model loads
+    app.state.logger = logger  # Make logger available to routes
+    
+    # Create necessary directories - use persistent locations
+    os.makedirs("/app/models", exist_ok=True)
+    os.makedirs("/app/tokenizers", exist_ok=True)
+    os.makedirs("/app/voice_memories", exist_ok=True)
+    os.makedirs("/app/voice_references", exist_ok=True)
+    os.makedirs("/app/voice_profiles", exist_ok=True)
+    os.makedirs("/app/cloned_voices", exist_ok=True)
+    os.makedirs("/app/audio_cache", exist_ok=True)
+    os.makedirs("/app/static", exist_ok=True)
+    
+    # Set tokenizer cache
+    try:
+        os.environ["TRANSFORMERS_CACHE"] = "/app/tokenizers"
+        logger.info(f"Set tokenizer cache to: {os.environ['TRANSFORMERS_CACHE']}")
+    except Exception as e:
+        logger.error(f"Error setting tokenizer cache: {e}")
+    
+    # Install additional dependencies if needed
+    try:
+        import scipy
+        import soundfile
+        logger.info("Audio processing dependencies available")
+    except ImportError as e:
+        logger.warning(f"Audio processing dependency missing: {e}. Some audio enhancements may not work.")
+        logger.warning("Consider installing: pip install scipy soundfile")
+    
+    # Check CUDA availability
+    cuda_available = torch.cuda.is_available()
+    if cuda_available:
+        device_count = torch.cuda.device_count()
+        device_name = torch.cuda.get_device_name(0) if device_count > 0 else "unknown"
+        logger.info(f"CUDA is available: {device_count} device(s). Using {device_name}")
+        # Report CUDA memory
+        if hasattr(torch.cuda, 'get_device_properties'):
+            total_memory = torch.cuda.get_device_properties(0).total_memory
+            logger.info(f"Total CUDA memory: {total_memory / (1024**3):.2f} GB")
+    else:
+        logger.warning("CUDA is not available. Using CPU (this will be slow)")
+    
+    # Determine device and device mapping
+    device = "cuda" if cuda_available else "cpu"
+    device_map = os.environ.get("CSM_DEVICE_MAP", None)  # Options: "auto", "balanced", "sequential"
+    if device_map and cuda_available:
+        if torch.cuda.device_count() > 1:
+            logger.info(f"Using device mapping strategy: {device_map} across {torch.cuda.device_count()} GPUs")
+        else:
+            logger.info("Device mapping requested but only one GPU available, ignoring device_map")
+            device_map = None
+    else:
+        device_map = None
+    
+    logger.info(f"Using device: {device}")
+    app.state.device = device
+    app.state.device_map = device_map
+    
+    # Check if model file exists
+    model_path = os.path.join("/app/models", "ckpt.pt")
+    if not os.path.exists(model_path):
+        # Try to download at runtime if not present
+        logger.info("Model not found. Attempting to download...")
+        try:
+            from huggingface_hub import hf_hub_download, login
+            # Check for token in environment
+            hf_token = os.environ.get("HF_TOKEN")
+            if hf_token:
+                logger.info("Logging in to Hugging Face using provided token")
+                login(token=hf_token)
+            logger.info("Downloading CSM-1B model from Hugging Face...")
+            download_start = time.time()
+            model_path = hf_hub_download(
+                repo_id="sesame/csm-1b", 
+                filename="ckpt.pt", 
+                local_dir="/app/models"
+            )
+            download_time = time.time() - download_start
+            logger.info(f"Model downloaded to {model_path} in {download_time:.2f} seconds")
+        except Exception as e:
+            error_stack = traceback.format_exc()
+            logger.error(f"Error downloading model: {str(e)}\n{error_stack}")
+            logger.error("Please build the image with HF_TOKEN to download the model")
+            logger.error("Starting without model - API will return 503 Service Unavailable")
+    else:
+        logger.info(f"Found existing model at {model_path}")
+        logger.info(f"Model size: {os.path.getsize(model_path) / (1024 * 1024):.2f} MB")
+    
+    # Load the model
+    try:
+        logger.info("Loading CSM-1B model...")
+        load_start = time.time()
+        from app.generator import load_csm_1b
+        app.state.generator = load_csm_1b(model_path, device, device_map)
+        load_time = time.time() - load_start
+        logger.info(f"Model loaded successfully in {load_time:.2f} seconds")
+        
+        # Store sample rate in app state
+        app.state.sample_rate = app.state.generator.sample_rate
+        logger.info(f"Model sample rate: {app.state.sample_rate} Hz")
+        
+        # Initialize voice enhancement system (this will create proper voice profiles)
+        logger.info("Initializing voice enhancement system...")
+        try:
+            from app.voice_enhancement import initialize_voice_profiles, save_voice_profiles
+            initialize_voice_profiles()
+            app.state.voice_enhancement_enabled = True
+            logger.info("Voice profiles initialized successfully")
+        except Exception as e:
+            error_stack = traceback.format_exc()
+            logger.error(f"Error initializing voice profiles: {str(e)}\n{error_stack}")
+            logger.warning("Voice enhancement features will be limited")
+            app.state.voice_enhancement_enabled = False
+        
+        # Initialize voice memory system for consistent generation
+        logger.info("Initializing voice memory system...")
+        try:
+            from app.voice_memory import initialize_voices
+            initialize_voices(app.state.sample_rate)
+            app.state.voice_memory_enabled = True
+            logger.info("Voice memory system initialized")
+        except Exception as e:
+            logger.warning(f"Error initializing voice memory: {e}")
+            app.state.voice_memory_enabled = False
+        
+        # Initialize voice cloning system
+        try:
+            logger.info("Initializing voice cloning system...")
+            from app.voice_cloning import VoiceCloner, CLONED_VOICES_DIR
+            # Update the cloned voices directory to use the persistent volume
+            app.state.cloned_voices_dir = "/app/cloned_voices"  # Store path in app state for access
+            os.makedirs(app.state.cloned_voices_dir, exist_ok=True)
+            CLONED_VOICES_DIR = app.state.cloned_voices_dir  # Update the module constant
+            
+            # Initialize the voice cloner with proper device
+            app.state.voice_cloner = VoiceCloner(app.state.generator, device=device)
+            
+            # Make sure existing voices are loaded
+            app.state.voice_cloner._load_existing_voices()
+            
+            # Log the available voices
+            cloned_voices = app.state.voice_cloner.list_voices()
+            logger.info(f"Voice cloning system initialized with {len(cloned_voices)} existing voices")
+            for voice in cloned_voices:
+                logger.info(f"  - {voice.name} (ID: {voice.id}, Speaker ID: {voice.speaker_id})")
+            
+            # Flag for voice cloning availability
+            app.state.voice_cloning_enabled = True
+        except Exception as e:
+            error_stack = traceback.format_exc()
+            logger.error(f"Error initializing voice cloning: {e}\n{error_stack}")
+            logger.warning("Voice cloning features will not be available")
+            app.state.voice_cloning_enabled = False
+        
+        # Create prompt templates for consistent generation
+        logger.info("Setting up prompt engineering templates...")
+        try:
+            from app.prompt_engineering import initialize_templates
+            app.state.prompt_templates = initialize_templates()
+            logger.info("Prompt templates initialized")
+        except Exception as e:
+            error_stack = traceback.format_exc()
+            logger.error(f"Error initializing prompt templates: {e}\n{error_stack}")
+            logger.warning("Voice consistency features will be limited")
+        
+        # Generate voice reference samples (runs in background to avoid blocking startup)
+        async def generate_samples_async():
+            try:
+                logger.info("Starting voice reference generation (background task)...")
+                from app.voice_enhancement import create_voice_segments
+                create_voice_segments(app.state)
+                logger.info("Voice reference generation completed")
+            except Exception as e:
+                error_stack = traceback.format_exc()
+                logger.error(f"Error in voice reference generation: {str(e)}\n{error_stack}")
+        
+        # Start as a background task
+        asyncio.create_task(generate_samples_async())
+        
+        # Initialize voice cache for all voices (standard + cloned)
+        app.state.voice_cache = {}
+        
+        # Add standard voices
+        standard_voices = ["alloy", "echo", "fable", "onyx", "nova", "shimmer"]
+        for voice in standard_voices:
+            app.state.voice_cache[voice] = []
+        
+        # Add cloned voices to cache if they exist
+        if app.state.voice_cloning_enabled and hasattr(app.state, "voice_cloner"):
+            for voice in app.state.voice_cloner.list_voices():
+                app.state.voice_cache[voice.id] = []
+                # Also add by name for more flexible lookup
+                app.state.voice_cache[voice.name] = []
+        
+        # Create mapping from voice name/id to speaker_id for easy lookup
+        app.state.voice_speaker_map = {
+            "alloy": 0, "echo": 1, "fable": 2, "onyx": 3, "nova": 4, "shimmer": 5
+        }
+        
+        # Add cloned voices to the speaker map
+        if app.state.voice_cloning_enabled and hasattr(app.state, "voice_cloner"):
+            for voice in app.state.voice_cloner.list_voices():
+                app.state.voice_speaker_map[voice.id] = voice.speaker_id
+                app.state.voice_speaker_map[voice.name] = voice.speaker_id
+                app.state.voice_speaker_map[str(voice.speaker_id)] = voice.speaker_id
+        
+        # Compile voice information for API
+        app.state.available_voices = standard_voices.copy()
+        if app.state.voice_cloning_enabled and hasattr(app.state, "voice_cloner"):
+            for voice in app.state.voice_cloner.list_voices():
+                app.state.available_voices.append(voice.id)
+                app.state.available_voices.append(voice.name)
+        
+        # Store model information for API endpoints
+        app.state.model_info = {
+            "name": "CSM-1B",
+            "device": device,
+            "device_map": device_map,
+            "sample_rate": app.state.sample_rate,
+            "standard_voices": standard_voices,
+            "cloned_voices": [v.id for v in app.state.voice_cloner.list_voices()] if app.state.voice_cloning_enabled else [],
+            "voice_enhancement_enabled": app.state.voice_enhancement_enabled,
+            "voice_memory_enabled": app.state.voice_memory_enabled,
+            "voice_cloning_enabled": app.state.voice_cloning_enabled,
+            "streaming_enabled": True
+        }
+        
+        # Create a function to access all voices in a standardized format
+        def get_all_available_voices():
+            """Helper function to get all available voices for API endpoints"""
+            # Standard voices with fixed descriptions
+            all_voices = [
+                {"voice_id": "alloy", "name": "Alloy", "description": "Balanced and natural"},
+                {"voice_id": "echo", "name": "Echo", "description": "Resonant and deeper"},
+                {"voice_id": "fable", "name": "Fable", "description": "Bright and higher-pitched"},
+                {"voice_id": "onyx", "name": "Onyx", "description": "Deep and authoritative"},
+                {"voice_id": "nova", "name": "Nova", "description": "Warm and smooth"},
+                {"voice_id": "shimmer", "name": "Shimmer", "description": "Light and airy"}
+            ]
+            
+            # Add cloned voices if available
+            if app.state.voice_cloning_enabled and hasattr(app.state, "voice_cloner"):
+                for voice in app.state.voice_cloner.list_voices():
+                    all_voices.append({
+                        "voice_id": voice.id,
+                        "name": voice.name,
+                        "description": voice.description or f"Cloned voice: {voice.name}"
+                    })
+            
+            return all_voices
+        
+        app.state.get_all_voices = get_all_available_voices
+        
+        # Add helper function to lookup voice info
+        def get_voice_info(voice_identifier):
+            """Look up voice information based on name, ID, or speaker_id"""
+            # Check standard voices
+            if voice_identifier in standard_voices:
+                return {
+                    "type": "standard",
+                    "voice_id": voice_identifier,
+                    "name": voice_identifier,
+                    "speaker_id": standard_voices.index(voice_identifier)
+                }
+            
+            # Look for cloned voice
+            if not app.state.voice_cloning_enabled or not hasattr(app.state, "voice_cloner"):
+                return None
+                
+            # Check by ID
+            if voice_identifier in app.state.voice_cloner.cloned_voices:
+                voice = app.state.voice_cloner.cloned_voices[voice_identifier]
+                return {
+                    "type": "cloned",
+                    "voice_id": voice.id,
+                    "name": voice.name,
+                    "speaker_id": voice.speaker_id
+                }
+                
+            # Check by name
+            for v_id, voice in app.state.voice_cloner.cloned_voices.items():
+                if voice.name == voice_identifier:
+                    return {
+                        "type": "cloned",
+                        "voice_id": voice.id,
+                        "name": voice.name,
+                        "speaker_id": voice.speaker_id
+                    }
+                    
+            # Check by speaker_id (string representation)
+            try:
+                speaker_id = int(voice_identifier)
+                # Check if any cloned voice has this speaker_id
+                for v_id, voice in app.state.voice_cloner.cloned_voices.items():
+                    if voice.speaker_id == speaker_id:
+                        return {
+                            "type": "cloned", 
+                            "voice_id": voice.id,
+                            "name": voice.name,
+                            "speaker_id": speaker_id
+                        }
+            except (ValueError, TypeError):
+                pass
+                
+            # No match found
+            return None
+            
+        app.state.get_voice_info = get_voice_info
+        
+        # Set up audio cache
+        app.state.audio_cache_enabled = os.environ.get("ENABLE_AUDIO_CACHE", "true").lower() == "true"
+        if app.state.audio_cache_enabled:
+            app.state.audio_cache_dir = "/app/audio_cache"
+            logger.info(f"Audio cache enabled, cache dir: {app.state.audio_cache_dir}")
+        
+        # Log GPU utilization after model loading
+        if cuda_available:
+            memory_allocated = torch.cuda.memory_allocated() / (1024**3)
+            memory_reserved = torch.cuda.memory_reserved() / (1024**3)
+            logger.info(f"GPU memory: {memory_allocated:.2f} GB allocated, {memory_reserved:.2f} GB reserved")
+            
+            if torch.cuda.device_count() > 1 and device_map:
+                logger.info("Multi-GPU setup active with the following memory usage:")
+                for i in range(torch.cuda.device_count()):
+                    memory_allocated = torch.cuda.memory_allocated(i) / (1024**3)
+                    memory_reserved = torch.cuda.memory_reserved(i) / (1024**3)
+                    logger.info(f"GPU {i}: {memory_allocated:.2f} GB allocated, {memory_reserved:.2f} GB reserved")
+        
+        # Set up scheduled tasks
+        try:
+            # Create a background task for periodic voice profile backup
+            async def periodic_voice_profile_backup(interval_hours=6):
+                """Periodically save voice profiles to persistent storage."""
+                while True:
+                    try:
+                        # Wait for the specified interval
+                        await asyncio.sleep(interval_hours * 3600)
+                        
+                        # Log the backup
+                        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
+                        logger.info(f"Scheduled voice profile backup started at {timestamp}")
+                        
+                        # Save voice profiles
+                        if hasattr(app.state, "voice_enhancement_enabled") and app.state.voice_enhancement_enabled:
+                            from app.voice_enhancement import save_voice_profiles
+                            save_voice_profiles()
+                            logger.info("Voice profiles saved successfully")
+                            
+                        # Save voice memories
+                        if hasattr(app.state, "voice_memory_enabled") and app.state.voice_memory_enabled:
+                            from app.voice_memory import VOICE_MEMORIES
+                            for voice_name, memory in VOICE_MEMORIES.items():
+                                memory.save()
+                            logger.info("Voice memories saved successfully")
+                            
+                    except Exception as e:
+                        logger.error(f"Error in periodic voice profile backup: {e}")
+            
+            # Start the scheduled task
+            asyncio.create_task(periodic_voice_profile_backup(interval_hours=6))
+            logger.info("Started scheduled voice profile backup task")
+            
+        except Exception as e:
+            logger.warning(f"Failed to set up scheduled tasks: {e}")
+        
+        logger.info(f"CSM-1B TTS API is ready on {device} with sample rate {app.state.sample_rate}")
+        logger.info(f"Standard voices: {standard_voices}")
+        cloned_count = len(app.state.voice_cloner.list_voices()) if app.state.voice_cloning_enabled else 0
+        logger.info(f"Cloned voices: {cloned_count}")
+        
+    except Exception as e:
+        error_stack = traceback.format_exc()
+        logger.error(f"Error loading model: {str(e)}\n{error_stack}")
+        app.state.generator = None
+    
+    # Calculate total startup time
+    startup_time = time.time() - app.state.startup_time
+    logger.info(f"Application startup completed in {startup_time:.2f} seconds")
+    
+    yield  # This is where the application runs
+    
+    # SHUTDOWN EVENT
+    logger.info("Application shutdown initiated")
+    
+    # Clean up model resources
+    if hasattr(app.state, "generator") and app.state.generator is not None:
+        try:
+            # Clean up CUDA memory if available
+            if torch.cuda.is_available():
+                logger.info("Clearing CUDA cache")
+                torch.cuda.empty_cache()
+                torch.cuda.synchronize()
+        except Exception as e:
+            logger.error(f"Error during CUDA cleanup: {e}")
+    
+    # Save voice profiles if they've been updated
+    try:
+        if hasattr(app.state, "voice_enhancement_enabled") and app.state.voice_enhancement_enabled:
+            from app.voice_enhancement import save_voice_profiles
+            logger.info("Saving voice profiles...")
+            save_voice_profiles()
+            logger.info("Voice profiles saved successfully")
+    except Exception as e:
+        logger.error(f"Error saving voice profiles: {e}")
+    
+    # Save voice memories if they've been updated
+    try:
+        if hasattr(app.state, "voice_memory_enabled") and app.state.voice_memory_enabled:
+            from app.voice_memory import VOICE_MEMORIES
+            logger.info("Saving voice memories...")
+            for voice_name, memory in VOICE_MEMORIES.items():
+                memory.save()
+            logger.info("Voice memories saved successfully")
+    except Exception as e:
+        logger.error(f"Error saving voice memories: {e}")
+    
+    # Clean up any temporary files
+    try:
+        for temp_file in glob.glob(os.path.join(tempfile.gettempdir(), "csm_tts_*")):
+            try:
+                os.remove(temp_file)
+                logger.info(f"Removed temporary file: {temp_file}")
+            except:
+                pass
+    except Exception as e:
+        logger.warning(f"Error cleaning up temporary files: {e}")
+    
+    logger.info("Application shutdown complete")
+    
+# Initialize FastAPI app
+app = FastAPI(
+    title="CSM-1B TTS API",
+    description="OpenAI-compatible TTS API using the CSM-1B model from Sesame",
+    version="1.0.0",
+    lifespan=lifespan
+)
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Create static and other required directories
+os.makedirs("/app/static", exist_ok=True)
+os.makedirs("/app/cloned_voices", exist_ok=True)
+
+# Mount the static files directory
+app.mount("/static", StaticFiles(directory="/app/static"), name="static")
+
+# Include routers
+app.include_router(api_router, prefix="/api/v1")
+
+# Add OpenAI compatible route
+app.include_router(api_router, prefix="/v1")
+
+# Add voice cloning routes
+from app.api.voice_cloning_routes import router as voice_cloning_router
+app.include_router(voice_cloning_router, prefix="/api/v1")
+app.include_router(voice_cloning_router, prefix="/v1")
+
+# Add streaming routes
+from app.api.streaming import router as streaming_router
+app.include_router(streaming_router, prefix="/api/v1")
+app.include_router(streaming_router, prefix="/v1")
+
+# Middleware for request timing
+@app.middleware("http")
+async def add_process_time_header(request: Request, call_next):
+    """Middleware to track request processing time."""
+    start_time = time.time()
+    response = await call_next(request)
+    process_time = time.time() - start_time
+    response.headers["X-Process-Time"] = str(process_time)
+    logger.debug(f"Request to {request.url.path} processed in {process_time:.3f} seconds")
+    return response
+
+# Health check endpoint
+@app.get("/health", include_in_schema=False)
+async def health_check(request: Request):
+    """Health check endpoint that returns the status of the API."""
+    model_status = "healthy" if hasattr(request.app.state, "generator") and request.app.state.generator is not None else "unhealthy"
+    uptime = time.time() - getattr(request.app.state, "startup_time", time.time())
+
+    # Get voice information
+    standard_voices = ["alloy", "echo", "fable", "onyx", "nova", "shimmer"]
+    cloned_voices = []
+    
+    if hasattr(request.app.state, "voice_cloner") and request.app.state.voice_cloner is not None:
+        cloned_voices = [
+            {"id": v.id, "name": v.name, "speaker_id": v.speaker_id}
+            for v in request.app.state.voice_cloner.list_voices()
+        ]
+    
+    # Get CUDA memory stats if available
+    cuda_stats = None
+    if torch.cuda.is_available():
+        cuda_stats = {
+            "allocated_gb": torch.cuda.memory_allocated() / (1024**3),
+            "reserved_gb": torch.cuda.memory_reserved() / (1024**3)
+        }
+    
+    return {
+        "status": model_status,
+        "uptime": f"{uptime:.2f} seconds",
+        "device": getattr(request.app.state, "device", "unknown"),
+        "model": "CSM-1B",
+        "standard_voices": standard_voices,
+        "cloned_voices": cloned_voices,
+        "cloned_voices_count": len(cloned_voices),
+        "sample_rate": getattr(request.app.state, "sample_rate", 0),
+        "enhancements": "enabled" if hasattr(request.app.state, "model_info") and 
+                      request.app.state.model_info.get("voice_enhancement_enabled", False) else "disabled",
+        "streaming": "enabled",
+        "cuda": cuda_stats,
+        "version": "1.0.0"
+    }
+
+# Version endpoint
+@app.get("/version", include_in_schema=False)
+async def version():
+    """Version endpoint that returns API version information."""
+    return {
+        "api_version": "1.0.0",
+        "model_version": "CSM-1B",
+        "compatible_with": "OpenAI TTS v1",
+        "enhancements": "voice consistency and audio quality v1.0",
+        "voice_cloning": "enabled" if hasattr(app.state, "voice_cloner") else "disabled",
+        "streaming": "enabled"
+    }
+
+# Voice cloning UI endpoint
+@app.get("/voice-cloning", include_in_schema=False)
+async def voice_cloning_ui():
+    """Voice cloning UI endpoint."""
+    return FileResponse("/app/static/voice-cloning.html")
+
+# Streaming demo endpoint
+@app.get("/streaming-demo", include_in_schema=False)
+async def streaming_demo():
+    """Streaming TTS demo endpoint."""
+    return FileResponse("/app/static/streaming-demo.html")
+
+@app.get("/", include_in_schema=False)
+async def root():
+    """Root endpoint that redirects to docs."""
+    logger.debug("Root endpoint accessed, redirecting to docs")
+    return RedirectResponse(url="/docs")
+
+if __name__ == "__main__":
+    # Get port from environment or use default
+    port = int(os.environ.get("PORT", 8000))
+    
+    # Development mode flag
+    dev_mode = os.environ.get("DEV_MODE", "false").lower() == "true"
+    
+    # Log level (default to INFO, but can be overridden)
+    log_level = os.environ.get("LOG_LEVEL", "INFO").upper()
+    logging.getLogger().setLevel(log_level)
+    
+    # Check for audio enhancement and voice cloning flags
+    enable_enhancements = os.environ.get("ENABLE_ENHANCEMENTS", "true").lower() == "true"
+    enable_voice_cloning = os.environ.get("ENABLE_VOICE_CLONING", "true").lower() == "true"
+    
+    if not enable_enhancements:
+        logger.warning("Voice enhancements disabled by environment variable")
+    if not enable_voice_cloning:
+        logger.warning("Voice cloning disabled by environment variable")
+    
+    logger.info(f"Voice enhancements: {'enabled' if enable_enhancements else 'disabled'}")
+    logger.info(f"Voice cloning: {'enabled' if enable_voice_cloning else 'disabled'}")
+    logger.info(f"Streaming: enabled")
+    logger.info(f"Log level: {log_level}")
+    
+    if dev_mode:
+        logger.info(f"Running in development mode with auto-reload enabled on port {port}")
+        uvicorn.run(
+            "app.main:app", 
+            host="0.0.0.0", 
+            port=port, 
+            reload=True, 
+            log_level=log_level.lower()
+        )
+    else:
+        logger.info(f"Running in production mode on port {port}")
+        uvicorn.run(
+            "app.main:app", 
+            host="0.0.0.0", 
+            port=port, 
+            reload=False, 
+            log_level=log_level.lower()
+        )

app/models.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+from typing import List, Tuple
+
+import torch
+
+
+@dataclass
+class Segment:
+    """A segment of speech with text, speaker, and audio."""
+    speaker: int
+    text: str
+    # (num_samples,), sample_rate = 24_000
+    audio: torch.Tensor

app/prompt_engineering.py

@@ -0,0 +1,129 @@
+"""Prompt engineering for consistent voice generation."""
+import re
+import random
+from typing import List, Dict, Optional
+import logging
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+# Voice style descriptors for consistent prompting
+VOICE_STYLES = {
+    "alloy": {
+        "adjectives": ["balanced", "natural", "clear", "articulate", "neutral", "conversational"],
+        "characteristics": ["medium pitch", "even pacing", "neutral tone", "balanced resonance"],
+        "speaking_style": "conversational and balanced"
+    },
+    "echo": {
+        "adjectives": ["resonant", "deep", "reverberant", "rich", "sonorous", "full"],
+        "characteristics": ["lower pitch", "deliberate pacing", "resonant tone", "deeper timbre"],
+        "speaking_style": "rich and resonant"
+    },
+    "fable": {
+        "adjectives": ["bright", "light", "clear", "energetic", "articulate", "animated"],
+        "characteristics": ["higher pitch", "lively pacing", "bright tone", "clear articulation"],
+        "speaking_style": "bright and energetic"
+    },
+    "onyx": {
+        "adjectives": ["deep", "authoritative", "powerful", "commanding", "strong", "resolute"],
+        "characteristics": ["low pitch", "measured pacing", "authoritative tone", "strong projection"],
+        "speaking_style": "deep and authoritative"
+    },
+    "nova": {
+        "adjectives": ["warm", "pleasant", "smooth", "harmonious", "gentle", "comforting"],
+        "characteristics": ["medium pitch", "smooth pacing", "warm tone", "pleasant timbre"],
+        "speaking_style": "warm and smooth"
+    },
+    "shimmer": {
+        "adjectives": ["light", "airy", "bright", "crystalline", "delicate", "expressive"],
+        "characteristics": ["higher pitch", "quick pacing", "light tone", "bright timbre"],
+        "speaking_style": "light and expressive"
+    },
+    "custom": {
+        "adjectives": ["clear", "distinct", "authentic", "natural", "personalized", "unique"],
+        "characteristics": ["natural rhythm", "authentic tone", "personal inflection", "distinctive sound"],
+        "speaking_style": "authentic and natural"
+    }
+}
+
+def initialize_templates():
+    """Initialize prompt templates - placeholder for any future setup."""
+    logger.info("Prompt templates initialized")
+    return VOICE_STYLES
+
+def split_into_segments(text: str, max_chars: int = 150) -> List[str]:
+    """Split text into optimal segments for better generation.
+    Args:
+        text: Text to split
+        max_chars: Maximum characters per segment
+    Returns:
+        List of text segments
+    """
+    # Handle empty or very short text
+    if not text or len(text) <= max_chars:
+        return [text]
+    
+    # Split by sentences first
+    sentences = re.split(r'(?<=[.!?])\s+', text)
+    
+    # Initialize segments
+    segments = []
+    current_segment = ""
+    
+    for sentence in sentences:
+        # If adding this sentence would exceed max_chars
+        if len(current_segment) + len(sentence) > max_chars:
+            # If current segment is not empty, add it to segments
+            if current_segment:
+                segments.append(current_segment.strip())
+                current_segment = ""
+                
+            # If this sentence alone exceeds max_chars, split it by phrases
+            if len(sentence) > max_chars:
+                phrases = re.split(r'(?<=[,;:])\s+', sentence)
+                for phrase in phrases:
+                    if len(phrase) > max_chars:
+                        # Split long phrases into chunks
+                        words = phrase.split()
+                        chunk = ""
+                        for word in words:
+                            if len(chunk) + len(word) + 1 <= max_chars:
+                                chunk += " " + word if chunk else word
+                            else:
+                                segments.append(chunk.strip())
+                                chunk = word
+                        if chunk:
+                            segments.append(chunk.strip())
+                    else:
+                        if len(current_segment) + len(phrase) <= max_chars:
+                            current_segment += " " + phrase if current_segment else phrase
+                        else:
+                            segments.append(current_segment.strip())
+                            current_segment = phrase
+            else:
+                current_segment = sentence
+        else:
+            current_segment += " " + sentence if current_segment else sentence
+            
+    # Add the last segment
+    if current_segment:
+        segments.append(current_segment.strip())
+        
+    logger.info(f"Split text into {len(segments)} segments")
+    return segments
+
+def format_text_for_voice(text: str, voice_name: str, segment_index: int = 0, total_segments: int = 1) -> str:
+    """Format text with voice characteristics for more consistent generation.
+    Args:
+        text: Text to format
+        voice_name: Name of the voice
+        segment_index: Index of this segment (for multi-segment texts)
+        total_segments: Total number of segments
+    Returns:
+        Formatted text optimized for consistent voice generation
+    """
+    # IMPORTANT: We no longer add voice instructions in brackets since CSM reads them aloud
+    # Instead, we're using speaker IDs to control voice identity which is what the model expects
+    
+    # Just return the unmodified text - the Generator class will handle proper formatting
+    return text

app/text_normalizer.py

@@ -0,0 +1,194 @@
+"""
+Text normalization and cleaning utilities for CSM-1B TTS system.
+Handles common issues like contractions, numbers, and special characters.
+"""
+import re
+import logging
+
+logger = logging.getLogger(__name__)
+
+class TextNormalizer:
+    """Text normalization utilities for TTS."""
+    
+    # Common English contractions mapping
+    CONTRACTIONS = {
+        "don't": "dont",
+        "won't": "wont",
+        "can't": "cant",
+        "isn't": "isnt",
+        "he's": "hes",
+        "she's": "shes",
+        "they're": "theyre",
+        "we're": "were",
+        "you're": "youre",
+        "that's": "thats",
+        "it's": "its",
+        "what's": "whats",
+        "let's": "lets",
+        "who's": "whos",
+        "how's": "hows",
+        "where's": "wheres",
+        "there's": "theres",
+        "wouldn't": "wouldnt",
+        "shouldn't": "shouldnt",
+        "couldn't": "couldnt",
+        "hasn't": "hasnt",
+        "haven't": "havent",
+        "hadn't": "hadnt",
+        "didn't": "didnt",
+        "i'm": "im",
+        "i've": "ive",
+        "i'd": "id",
+        "i'll": "ill",
+        "you've": "youve",
+        "you'll": "youll",
+        "you'd": "youd",
+        "we've": "weve",
+        "we'll": "well",
+        "we'd": "wed",
+        "they've": "theyve",
+        "they'll": "theyll",
+        "they'd": "theyd",
+        "aren't": "arent",
+        "weren't": "werent",
+        "wasn't": "wasnt",
+    }
+    
+    # Common abbreviations to expand
+    ABBREVIATIONS = {
+        "Mr.": "Mister",
+        "Mrs.": "Misses",
+        "Dr.": "Doctor",
+        "Prof.": "Professor",
+        "St.": "Street",
+        "Rd.": "Road",
+        "Ave.": "Avenue",
+        "vs.": "versus",
+        "etc.": "etcetera",
+        "e.g.": "for example",
+        "i.e.": "that is",
+        "approx.": "approximately",
+    }
+    
+    # Simple number words for common numbers
+    NUMBER_WORDS = {
+        "0": "zero",
+        "1": "one",
+        "2": "two",
+        "3": "three",
+        "4": "four",
+        "5": "five",
+        "6": "six",
+        "7": "seven",
+        "8": "eight",
+        "9": "nine",
+        "10": "ten",
+        "11": "eleven",
+        "12": "twelve",
+        "13": "thirteen",
+        "14": "fourteen",
+        "15": "fifteen",
+        "16": "sixteen",
+        "17": "seventeen",
+        "18": "eighteen",
+        "19": "nineteen",
+        "20": "twenty",
+        "30": "thirty",
+        "40": "forty",
+        "50": "fifty",
+        "60": "sixty",
+        "70": "seventy",
+        "80": "eighty",
+        "90": "ninety",
+        "100": "one hundred",
+        "1000": "one thousand",
+        "1000000": "one million",
+        "1000000000": "one billion",
+    }
+    
+    @classmethod
+    def normalize_text(cls, text: str) -> str:
+        """
+        Normalize text for TTS: handle contractions, punctuation, and special cases.
+        
+        Args:
+            text: Input text to normalize
+            
+        Returns:
+            Normalized text ready for TTS
+        """
+        if not text:
+            return text
+            
+        # Log original text for debugging
+        logger.debug(f"Normalizing text: '{text}'")
+        
+        # Remove voice instructions in square brackets
+        text = re.sub(r'\[.*?\]', '', text)
+        
+        # Handle contractions - preserving case sensitivity
+        for contraction, replacement in cls.CONTRACTIONS.items():
+            # Case insensitive replacement
+            text = re.sub(r'\b' + re.escape(contraction) + r'\b', replacement, text, flags=re.IGNORECASE)
+        
+        # Expand common abbreviations
+        for abbr, expanded in cls.ABBREVIATIONS.items():
+            text = text.replace(abbr, expanded)
+        
+        # Handle numbers - only convert standalone numbers
+        def replace_number(match):
+            number = match.group(0)
+            if number in cls.NUMBER_WORDS:
+                return cls.NUMBER_WORDS[number]
+            return number
+            
+        text = re.sub(r'\b\d+\b', replace_number, text)
+        
+        # Replace problematic symbols
+        text = text.replace("&", " and ")
+        text = text.replace("%", " percent ")
+        text = text.replace("@", " at ")
+        text = text.replace("#", " number ")
+        text = text.replace("$", " dollar ")
+        text = text.replace("€", " euro ")
+        text = text.replace("£", " pound ")
+        text = text.replace("¥", " yen ")
+        
+        # Handle dates in MM/DD/YYYY format
+        text = re.sub(r'\b(\d{1,2})/(\d{1,2})/(\d{4})\b', r'\1 \2 \3', text)
+        
+        # Fix excessive spaces
+        text = re.sub(r'\s+', ' ', text).strip()
+        
+        # Ensure sentence ends with punctuation
+        if not text[-1] in ['.', '!', '?', ';', ':', ',']:
+            text = text + '.'
+            
+        logger.debug(f"Normalized text: '{text}'")
+        return text
+    
+    @classmethod
+    def split_into_sentences(cls, text: str) -> list:
+        """
+        Split text into sentences for better TTS performance.
+        
+        Args:
+            text: Input text to split
+            
+        Returns:
+            List of sentences
+        """
+        # Normalize first
+        text = cls.normalize_text(text)
+        
+        # Split on sentence boundaries
+        sentences = re.split(r'(?<=[.!?])\s+', text)
+        
+        # Remove empty sentences
+        sentences = [s for s in sentences if s.strip()]
+        
+        return sentences
+
+def clean_text_for_tts(text: str) -> str:
+    """Clean and normalize text for TTS processing."""
+    return TextNormalizer.normalize_text(text)

app/torchtune_models.py

@@ -0,0 +1,282 @@
+"""Torchtune models for CSM-1B."""
+import logging
+from dataclasses import dataclass
+import torch
+import torch.nn as nn
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+# First, try to import llama3_2 from torchtune directly
+try:
+    import torchtune
+    logger.info(f"Torchtune version: {getattr(torchtune, '__version__', 'unknown')}")
+    
+    # Print available modules in torchtune.models
+    try:
+        import torchtune.models
+        logger.info(f"Available modules in torchtune.models: {dir(torchtune.models)}")
+    except Exception as e:
+        logger.error(f"Error inspecting torchtune.models: {e}")
+    
+    # Try to import llama3_2 model
+    try:
+        from torchtune.models.llama3_2 import llama3_2
+        logger.info("Successfully imported llama3_2 from torchtune")
+    except ImportError as e:
+        logger.warning(f"Could not import llama3_2: {e}")
+        # Try to import regular llama as fallback
+        try:
+            from torchtune.models.llama import llama
+            logger.info("Using llama from torchtune.models.llama as fallback")
+            llama3_2 = llama  # Alias llama as llama3_2
+        except ImportError:
+            logger.error("Could not import llama model either. Will use custom implementation.")
+            llama3_2 = None
+except ImportError as e:
+    logger.error(f"Torchtune not available: {e}")
+    torchtune = None
+    llama3_2 = None
+
+
+# Define our own model implementations as fallbacks
+def llama3_2_1B_custom():
+    """Create a Llama 3.2 1B model."""
+    from app.custom_transformer import CustomTransformerDecoder
+    return CustomTransformerDecoder(
+        vocab_size=128_256,
+        num_layers=16,
+        num_heads=32,
+        num_kv_heads=8,
+        embed_dim=2048,
+        max_seq_len=2048,
+        intermediate_dim=8192,
+        attn_dropout=0.0,
+        norm_eps=1e-5,
+    )
+
+
+def llama3_2_100M_custom():
+    """Create a Llama 3.2 100M model."""
+    from app.custom_transformer import CustomTransformerDecoder
+    return CustomTransformerDecoder(
+        vocab_size=128_256,
+        num_layers=4,
+        num_heads=8,
+        num_kv_heads=2,
+        embed_dim=1024,
+        max_seq_len=2048,
+        intermediate_dim=8192,
+        attn_dropout=0.0,
+        norm_eps=1e-5,
+    )
+
+
+# Setup fallback to our own implementations if needed
+if llama3_2 is None:
+    logger.warning("Using custom implementations for Llama models")
+    FLAVORS = {
+        "llama-1B": llama3_2_1B_custom,
+        "llama-100M": llama3_2_100M_custom,
+    }
+else:
+    logger.info("Using torchtune implementations for Llama models")
+    FLAVORS = {
+        "llama-1B": lambda: llama3_2(
+            vocab_size=128_256,
+            num_layers=16,
+            num_heads=32,
+            num_kv_heads=8,
+            embed_dim=2048,
+            max_seq_len=2048,
+            intermediate_dim=8192,
+            attn_dropout=0.0,
+            norm_eps=1e-5,
+            rope_base=500_000,
+            scale_factor=32,
+        ),
+        "llama-100M": lambda: llama3_2(
+            vocab_size=128_256,
+            num_layers=4,
+            num_heads=8,
+            num_kv_heads=2,
+            embed_dim=1024,
+            max_seq_len=2048,
+            intermediate_dim=8192,
+            attn_dropout=0.0,
+            norm_eps=1e-5,
+            rope_base=500_000,
+            scale_factor=32,
+        ),
+    }
+
+
+def _prepare_transformer(model):
+    """Prepare transformer for use."""
+    embed_dim = model.tok_embeddings.embedding_dim
+    model.tok_embeddings = nn.Identity()
+    model.output = nn.Identity()
+    return model, embed_dim
+
+
+def _create_causal_mask(seq_len: int, device: torch.device):
+    """Create causal mask."""
+    return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))
+
+
+def _index_causal_mask(mask: torch.Tensor, input_pos: torch.Tensor):
+    """Index causal mask.
+    
+    Args:
+        mask: (max_seq_len, max_seq_len)
+        input_pos: (batch_size, seq_len)
+    
+    Returns:
+        (batch_size, seq_len, max_seq_len)
+    """
+    r = mask[input_pos, :]
+    return r
+
+
+def _multinomial_sample_one_no_sync(probs):
+    """Do multinomial sampling without a cuda synchronization."""
+    q = torch.empty_like(probs).exponential_(1)
+    return torch.argmax(probs / q, dim=-1, keepdim=True).to(dtype=torch.int)
+
+
+def sample_topk(logits: torch.Tensor, topk: int, temperature: float):
+    """Sample from top-k logits."""
+    logits = logits / temperature
+    filter_value: float = -float("Inf")
+    indices_to_remove = logits < torch.topk(logits, topk)[0][..., -1, None]
+    scores_processed = logits.masked_fill(indices_to_remove, filter_value)
+    scores_processed = torch.nn.functional.log_softmax(scores_processed, dim=-1)
+    probs = torch.nn.functional.softmax(scores_processed, dim=-1)
+    sample_token = _multinomial_sample_one_no_sync(probs)
+    return sample_token
+
+
+@dataclass
+class ModelArgs:
+    """Model arguments."""
+    backbone_flavor: str
+    decoder_flavor: str
+    text_vocab_size: int
+    audio_vocab_size: int
+    audio_num_codebooks: int
+
+
+class Model(nn.Module):
+    """CSM-1B model."""
+    
+    def __init__(self, args: ModelArgs):
+        """Initialize model."""
+        super().__init__()
+        self.args = args
+        logger.info(f"Creating model with backbone: {args.backbone_flavor}, decoder: {args.decoder_flavor}")
+        
+        # Load backbone and decoder
+        self.backbone, backbone_dim = _prepare_transformer(FLAVORS[args.backbone_flavor]())
+        self.decoder, decoder_dim = _prepare_transformer(FLAVORS[args.decoder_flavor]())
+        
+        # Embeddings
+        self.text_embeddings = nn.Embedding(args.text_vocab_size, backbone_dim)
+        self.audio_embeddings = nn.Embedding(args.audio_vocab_size * args.audio_num_codebooks, backbone_dim)
+        
+        # Projection and heads
+        self.projection = nn.Linear(backbone_dim, decoder_dim, bias=False)
+        self.codebook0_head = nn.Linear(backbone_dim, args.audio_vocab_size, bias=False)
+        self.audio_head = nn.Parameter(torch.empty(args.audio_num_codebooks - 1, decoder_dim, args.audio_vocab_size))
+        
+        # Initialize audio head
+        nn.init.normal_(self.audio_head, mean=0.0, std=0.02)
+    
+    def setup_caches(self, max_batch_size: int) -> torch.Tensor:
+        """Setup KV caches and return a causal mask."""
+        dtype = next(self.parameters()).dtype
+        device = next(self.parameters()).device
+        
+        with device:
+            self.backbone.setup_caches(max_batch_size, dtype)
+            self.decoder.setup_caches(max_batch_size, dtype, decoder_max_seq_len=self.args.audio_num_codebooks)
+        
+        self.register_buffer("backbone_causal_mask", _create_causal_mask(self.backbone.max_seq_len, device))
+        self.register_buffer("decoder_causal_mask", _create_causal_mask(self.args.audio_num_codebooks, device))
+    
+    def generate_frame(
+        self,
+        tokens: torch.Tensor,
+        tokens_mask: torch.Tensor,
+        input_pos: torch.Tensor,
+        temperature: float,
+        topk: int,
+    ) -> torch.Tensor:
+        """Generate a frame of audio tokens.
+        
+        Args:
+            tokens: (batch_size, seq_len, audio_num_codebooks+1)
+            tokens_mask: (batch_size, seq_len, audio_num_codebooks+1)
+            input_pos: (batch_size, seq_len) positions for each token
+        
+        Returns:
+            (batch_size, audio_num_codebooks) sampled tokens
+        """
+        dtype = next(self.parameters()).dtype
+        b, s = tokens.size()[:2]
+        
+        assert self.backbone.caches_are_enabled(), "backbone caches are not enabled"
+        
+        curr_backbone_mask = _index_causal_mask(self.backbone_causal_mask, input_pos)
+        embeds = self._embed_tokens(tokens)
+        masked_embeds = embeds * tokens_mask.unsqueeze(-1)
+        h = masked_embeds.sum(dim=2)
+        
+        h = self.backbone(h, input_pos=input_pos, mask=curr_backbone_mask).to(dtype=dtype)
+        
+        last_h = h[:, -1, :]
+        c0_logits = self.codebook0_head(last_h)
+        c0_sample = sample_topk(c0_logits, topk, temperature)
+        c0_embed = self._embed_audio(0, c0_sample)
+        
+        curr_h = torch.cat([last_h.unsqueeze(1), c0_embed], dim=1)
+        curr_sample = c0_sample.clone()
+        curr_pos = torch.arange(0, curr_h.size(1), device=curr_h.device).unsqueeze(0).repeat(curr_h.size(0), 1)
+        
+        # Decoder caches must be reset every frame.
+        self.decoder.reset_caches()
+        
+        for i in range(1, self.args.audio_num_codebooks):
+            curr_decoder_mask = _index_causal_mask(self.decoder_causal_mask, curr_pos)
+            decoder_h = self.decoder(self.projection(curr_h), input_pos=curr_pos, mask=curr_decoder_mask).to(
+                dtype=dtype
+            )
+            ci_logits = torch.mm(decoder_h[:, -1, :], self.audio_head[i - 1])
+            ci_sample = sample_topk(ci_logits, topk, temperature)
+            ci_embed = self._embed_audio(i, ci_sample)
+            
+            curr_h = ci_embed
+            curr_sample = torch.cat([curr_sample, ci_sample], dim=1)
+            curr_pos = curr_pos[:, -1:] + 1
+        
+        return curr_sample
+    
+    def reset_caches(self):
+        """Reset KV caches."""
+        self.backbone.reset_caches()
+        self.decoder.reset_caches()
+    
+    def _embed_audio(self, codebook: int, tokens: torch.Tensor) -> torch.Tensor:
+        """Embed audio tokens."""
+        return self.audio_embeddings(tokens + codebook * self.args.audio_vocab_size)
+    
+    def _embed_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
+        """Embed tokens."""
+        text_embeds = self.text_embeddings(tokens[:, :, -1]).unsqueeze(-2)
+        audio_tokens = tokens[:, :, :-1] + (
+            self.args.audio_vocab_size * torch.arange(self.args.audio_num_codebooks, device=tokens.device)
+        )
+        audio_embeds = self.audio_embeddings(audio_tokens.view(-1)).reshape(
+            tokens.size(0), tokens.size(1), self.args.audio_num_codebooks, -1
+        )
+        
+        return torch.cat([audio_embeds, text_embeds], dim=-2)

app/utils/audio_utils.py

@@ -0,0 +1,64 @@
+"""Audio utilities for CSM-1B API."""
+
+import io
+import tempfile
+from typing import Optional
+import os
+
+import torch
+import torchaudio
+import ffmpeg
+
+
+def convert_audio_format(
+    audio_tensor: torch.Tensor,
+    sample_rate: int,
+    format: str = "mp3",
+    bit_rate: Optional[str] = "128k",
+) -> bytes:
+    """Convert audio tensor to specified format.
+    
+    Args:
+        audio_tensor: Audio tensor (channels, samples)
+        sample_rate: Sample rate
+        format: Output format (mp3, opus, aac, flac, wav)
+        bit_rate: Bit rate for lossy formats
+    
+    Returns:
+        Audio bytes in specified format
+    """
+    # Create temporary files
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_wav:
+        wav_path = temp_wav.name
+    
+    temp_out = tempfile.NamedTemporaryFile(suffix=f".{format}", delete=False)
+    out_path = temp_out.name
+    temp_out.close()
+    
+    try:
+        # Save as WAV first (native format for torchaudio)
+        torchaudio.save(wav_path, audio_tensor.unsqueeze(0) if audio_tensor.dim() == 1 else audio_tensor, 
+                        sample_rate)
+        
+        # Convert to desired format using ffmpeg
+        if format == "mp3":
+            ffmpeg.input(wav_path).output(out_path, format=format, audio_bitrate=bit_rate).run(quiet=True)
+        elif format in ["opus", "aac"]:
+            ffmpeg.input(wav_path).output(out_path, format=format).run(quiet=True)
+        elif format == "flac":
+            ffmpeg.input(wav_path).output(out_path, format=format).run(quiet=True)
+        elif format == "wav":
+            # Already saved as WAV
+            pass
+        
+        # Read the output file
+        with open(out_path if format != "wav" else wav_path, "rb") as f:
+            audio_bytes = f.read()
+        
+        return audio_bytes
+    
+    finally:
+        # Clean up temporary files
+        for path in [wav_path, out_path]:
+            if os.path.exists(path):
+                os.unlink(path)

app/utils/init.py

@@ -0,0 +1 @@
+"""Utilities for CSM-1B API."""

app/utils/scheduled_tasks.py

@@ -0,0 +1,42 @@
+"""Scheduled tasks for the TTS API."""
+import asyncio
+import logging
+import time
+from datetime import datetime
+
+logger = logging.getLogger(__name__)
+
+async def periodic_voice_profile_backup(app_state, interval_hours=6):
+    """
+    Periodically save voice profiles to persistent storage.
+    
+    Args:
+        app_state: The application state object
+        interval_hours: Backup interval in hours
+    """
+    while True:
+        try:
+            # Wait for the specified interval
+            await asyncio.sleep(interval_hours * 3600)
+            
+            # Log the backup
+            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            logger.info(f"Scheduled voice profile backup started at {timestamp}")
+            
+            # Save voice profiles
+            if hasattr(app_state, "voice_enhancement_enabled") and app_state.voice_enhancement_enabled:
+                from app.voice_enhancement import save_voice_profiles
+                save_voice_profiles()
+                logger.info("Voice profiles saved successfully")
+                
+            # Save voice memories
+            if hasattr(app_state, "voice_memory_enabled") and app_state.voice_memory_enabled:
+                for voice_name in app_state.voice_cache:
+                    if voice_name in ["alloy", "echo", "fable", "onyx", "nova", "shimmer"]:
+                        from app.voice_memory import VOICE_MEMORIES
+                        if voice_name in VOICE_MEMORIES:
+                            VOICE_MEMORIES[voice_name].save()
+                logger.info("Voice memories saved successfully")
+                
+        except Exception as e:
+            logger.error(f"Error in periodic voice profile backup: {e}")

app/utils/voice_manager.py

@@ -0,0 +1,132 @@
+"""Utility functions for managing voice references and profiles."""
+import os
+import logging
+import torch
+import torchaudio
+import shutil
+from typing import Dict, List, Optional, Any
+
+logger = logging.getLogger(__name__)
+
+# Define persistent paths
+VOICE_REFERENCES_DIR = "/app/voice_references"
+VOICE_PROFILES_DIR = "/app/voice_profiles"
+VOICE_MEMORIES_DIR = "/app/voice_memories"
+
+# Ensure directories exist
+os.makedirs(VOICE_REFERENCES_DIR, exist_ok=True)
+os.makedirs(VOICE_PROFILES_DIR, exist_ok=True)
+os.makedirs(VOICE_MEMORIES_DIR, exist_ok=True)
+
+def backup_voice_data(backup_dir: str = "/app/voice_backups"):
+    """Create a backup of all voice data."""
+    os.makedirs(backup_dir, exist_ok=True)
+    timestamp = torch.datetime.now().strftime("%Y%m%d_%H%M%S")
+    backup_path = os.path.join(backup_dir, f"voice_backup_{timestamp}")
+    os.makedirs(backup_path, exist_ok=True)
+    
+    # Backup voice references
+    if os.path.exists(VOICE_REFERENCES_DIR):
+        refs_backup = os.path.join(backup_path, "voice_references")
+        shutil.copytree(VOICE_REFERENCES_DIR, refs_backup)
+        
+    # Backup voice profiles
+    if os.path.exists(VOICE_PROFILES_DIR):
+        profiles_backup = os.path.join(backup_path, "voice_profiles")
+        shutil.copytree(VOICE_PROFILES_DIR, profiles_backup)
+        
+    # Backup voice memories
+    if os.path.exists(VOICE_MEMORIES_DIR):
+        memories_backup = os.path.join(backup_path, "voice_memories")
+        shutil.copytree(VOICE_MEMORIES_DIR, memories_backup)
+        
+    logger.info(f"Voice data backed up to {backup_path}")
+    return backup_path
+
+def restore_default_voices():
+    """Reset voices to their default state by removing existing voice data."""
+    for voice_dir in [VOICE_REFERENCES_DIR, VOICE_PROFILES_DIR, VOICE_MEMORIES_DIR]:
+        if os.path.exists(voice_dir):
+            # Create a backup before deleting
+            backup_path = backup_voice_data()
+            
+            # Remove existing data
+            for item in os.listdir(voice_dir):
+                item_path = os.path.join(voice_dir, item)
+                if os.path.isdir(item_path):
+                    shutil.rmtree(item_path)
+                else:
+                    os.remove(item_path)
+                    
+            logger.info(f"Removed existing voice data from {voice_dir}")
+            
+    logger.info(f"Voices reset to default state (backup created at {backup_path})")
+    return backup_path
+
+def verify_voice_references():
+    """Check if voice references are complete and valid."""
+    standard_voices = ["alloy", "echo", "fable", "onyx", "nova", "shimmer"]
+    missing_voices = []
+    
+    for voice in standard_voices:
+        voice_dir = os.path.join(VOICE_REFERENCES_DIR, voice)
+        # Check if directory exists and contains reference files
+        if not os.path.exists(voice_dir) or len(os.listdir(voice_dir)) == 0:
+            missing_voices.append(voice)
+            
+    return {
+        "complete": len(missing_voices) == 0,
+        "missing_voices": missing_voices,
+        "references_dir": VOICE_REFERENCES_DIR
+    }
+
+def get_voice_storage_info() -> Dict[str, Any]:
+    """Get information about voice storage usage and status."""
+    result = {
+        "voice_references": {
+            "path": VOICE_REFERENCES_DIR,
+            "exists": os.path.exists(VOICE_REFERENCES_DIR),
+            "voices": [],
+            "total_size_mb": 0
+        },
+        "voice_profiles": {
+            "path": VOICE_PROFILES_DIR,
+            "exists": os.path.exists(VOICE_PROFILES_DIR),
+            "file_count": 0,
+            "total_size_mb": 0
+        },
+        "voice_memories": {
+            "path": VOICE_MEMORIES_DIR,
+            "exists": os.path.exists(VOICE_MEMORIES_DIR),
+            "voices": [],
+            "total_size_mb": 0
+        }
+    }
+    
+    # Get voice references info
+    if result["voice_references"]["exists"]:
+        for voice in os.listdir(VOICE_REFERENCES_DIR):
+            voice_dir = os.path.join(VOICE_REFERENCES_DIR, voice)
+            if os.path.isdir(voice_dir):
+                file_count = len([f for f in os.listdir(voice_dir) if f.endswith('.wav')])
+                dir_size = sum(os.path.getsize(os.path.join(voice_dir, f)) for f in os.listdir(voice_dir) if os.path.isfile(os.path.join(voice_dir, f)))
+                result["voice_references"]["voices"].append({
+                    "name": voice,
+                    "file_count": file_count,
+                    "size_mb": dir_size / (1024 * 1024)
+                })
+                result["voice_references"]["total_size_mb"] += dir_size / (1024 * 1024)
+                
+    # Get voice profiles info
+    if result["voice_profiles"]["exists"]:
+        files = [f for f in os.listdir(VOICE_PROFILES_DIR) if os.path.isfile(os.path.join(VOICE_PROFILES_DIR, f))]
+        result["voice_profiles"]["file_count"] = len(files)
+        result["voice_profiles"]["total_size_mb"] = sum(os.path.getsize(os.path.join(VOICE_PROFILES_DIR, f)) for f in files) / (1024 * 1024)
+        
+    # Get voice memories info
+    if result["voice_memories"]["exists"]:
+        files = [f for f in os.listdir(VOICE_MEMORIES_DIR) if os.path.isfile(os.path.join(VOICE_MEMORIES_DIR, f))]
+        result["voice_memories"]["voices"] = [f.replace('.pt', '') for f in files if f.endswith('.pt')]
+        result["voice_memories"]["total_size_mb"] = sum(os.path.getsize(os.path.join(VOICE_MEMORIES_DIR, f)) for f in files) / (1024 * 1024)
+        
+    return result

app/voice_cloning.py

@@ -0,0 +1,705 @@
+"""
+Voice cloning module for CSM-1B TTS API.
+
+This module provides functionality to clone voices from audio samples,
+with advanced audio preprocessing and voice adaptation techniques.
+"""
+import os
+import io
+import time
+import tempfile
+import logging
+import asyncio
+import yt_dlp
+import whisper
+from typing import Dict, List, Optional, Union, Tuple, BinaryIO
+from pathlib import Path
+
+import numpy as np
+import torch
+import torchaudio
+from pydantic import BaseModel
+from fastapi import UploadFile
+
+from app.models import Segment
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+# Directory for storing cloned voice data
+CLONED_VOICES_DIR = "/app/cloned_voices"
+os.makedirs(CLONED_VOICES_DIR, exist_ok=True)
+
+class ClonedVoice(BaseModel):
+    """Model representing a cloned voice."""
+    id: str
+    name: str
+    created_at: float
+    speaker_id: int
+    description: Optional[str] = None
+    audio_duration: float
+    sample_count: int
+
+
+class VoiceCloner:
+    """Voice cloning utility for CSM-1B model."""
+
+    def __init__(self, generator, device="cuda"):
+        """Initialize the voice cloner with a generator instance."""
+        self.generator = generator
+        self.device = device
+        self.sample_rate = generator.sample_rate
+        self.cloned_voices = self._load_existing_voices()
+        logger.info(f"Voice cloner initialized with {len(self.cloned_voices)} existing voices")
+
+    def _load_existing_voices(self) -> Dict[str, ClonedVoice]:
+        """Load existing cloned voices from disk."""
+        voices = {}
+        if not os.path.exists(CLONED_VOICES_DIR):
+            return voices
+
+        for voice_dir in os.listdir(CLONED_VOICES_DIR):
+            voice_path = os.path.join(CLONED_VOICES_DIR, voice_dir)
+            if not os.path.isdir(voice_path):
+                continue
+
+            info_path = os.path.join(voice_path, "info.json")
+            if os.path.exists(info_path):
+                try:
+                    import json
+                    with open(info_path, "r") as f:
+                        voice_info = json.load(f)
+                        voices[voice_dir] = ClonedVoice(**voice_info)
+                        logger.info(f"Loaded cloned voice: {voice_dir}")
+                except Exception as e:
+                    logger.error(f"Error loading voice {voice_dir}: {e}")
+
+        return voices
+
+    async def process_audio_file(
+        self, 
+        file: Union[UploadFile, BinaryIO, str],
+        transcript: Optional[str] = None
+    ) -> Tuple[torch.Tensor, Optional[str], float]:
+        """
+        Process an audio file for voice cloning.
+        
+        Args:
+            file: The audio file (UploadFile, file-like object, or path)
+            transcript: Optional transcript of the audio
+            
+        Returns:
+            Tuple of (processed_audio, transcript, duration_seconds)
+        """
+        temp_path = None
+        
+        try:
+            # Handle different input types
+            if isinstance(file, str):
+                # It's a file path
+                audio_path = file
+                logger.info(f"Processing audio from file path: {audio_path}")
+            else:
+                # Create a temporary file
+                temp_fd, temp_path = tempfile.mkstemp(suffix=".wav")
+                os.close(temp_fd)  # Close the file descriptor
+                
+                if isinstance(file, UploadFile):
+                    # It's a FastAPI UploadFile
+                    logger.info("Processing audio from UploadFile")
+                    contents = await file.read()
+                    with open(temp_path, "wb") as f:
+                        f.write(contents)
+                elif hasattr(file, 'read'):
+                    # It's a file-like object - check if it's async
+                    logger.info("Processing audio from file-like object")
+                    if asyncio.iscoroutinefunction(file.read):
+                        # It's an async read method
+                        contents = await file.read()
+                    else:
+                        # It's a sync read method
+                        contents = file.read()
+                        
+                    with open(temp_path, "wb") as f:
+                        f.write(contents)
+                else:
+                    raise ValueError(f"Unsupported file type: {type(file)}")
+                
+                audio_path = temp_path
+                logger.info(f"Saved uploaded audio to temporary file: {audio_path}")
+
+            # Load audio
+            logger.info(f"Loading audio from {audio_path}")
+            audio, sr = torchaudio.load(audio_path)
+            
+            # Convert to mono if stereo
+            if audio.shape[0] > 1:
+                logger.info(f"Converting {audio.shape[0]} channels to mono")
+                audio = torch.mean(audio, dim=0, keepdim=True)
+            
+            # Remove first dimension if it's 1
+            if audio.shape[0] == 1:
+                audio = audio.squeeze(0)
+            
+            # Resample if necessary
+            if sr != self.sample_rate:
+                logger.info(f"Resampling from {sr}Hz to {self.sample_rate}Hz")
+                audio = torchaudio.functional.resample(
+                    audio, orig_freq=sr, new_freq=self.sample_rate
+                )
+            
+            # Get audio duration
+            duration_seconds = len(audio) / self.sample_rate
+            
+            # Process audio for better quality
+            logger.info(f"Preprocessing audio for quality enhancement")
+            processed_audio = self._preprocess_audio(audio)
+            processed_duration = len(processed_audio) / self.sample_rate
+            
+            logger.info(
+                f"Processed audio: original duration={duration_seconds:.2f}s, "
+                f"processed duration={processed_duration:.2f}s"
+            )
+            
+            return processed_audio, transcript, duration_seconds
+                
+        except Exception as e:
+            logger.error(f"Error processing audio: {e}", exc_info=True)
+            raise RuntimeError(f"Failed to process audio file: {e}")
+            
+        finally:
+            # Clean up temp file if we created one
+            if temp_path and os.path.exists(temp_path):
+                try:
+                    os.unlink(temp_path)
+                    logger.debug(f"Deleted temporary file {temp_path}")
+                except Exception as e:
+                    logger.warning(f"Failed to delete temporary file {temp_path}: {e}")
+
+    def _preprocess_audio(self, audio: torch.Tensor) -> torch.Tensor:
+        """
+        Preprocess audio for better voice cloning quality.
+        
+        Args:
+            audio: Raw audio tensor
+            
+        Returns:
+            Processed audio tensor
+        """
+        # Normalize volume
+        if torch.max(torch.abs(audio)) > 0:
+            audio = audio / torch.max(torch.abs(audio))
+        
+        # Remove silence with dynamic threshold
+        audio = self._remove_silence(audio, threshold=0.02)  # Slightly higher threshold to remove more noise
+        
+        # Remove DC offset (very low frequency noise)
+        audio = audio - torch.mean(audio)
+        
+        # Apply simple noise reduction
+        # This filters out very high frequencies that might contain noise
+        try:
+            audio_np = audio.cpu().numpy()
+            from scipy import signal
+            
+            # Apply a bandpass filter to focus on speech frequencies (80Hz - 8000Hz)
+            sos = signal.butter(3, [80, 8000], 'bandpass', fs=self.sample_rate, output='sos')
+            filtered = signal.sosfilt(sos, audio_np)
+            
+            # Normalize the filtered audio
+            filtered = filtered / (np.max(np.abs(filtered)) + 1e-8)
+            
+            # Convert back to torch tensor
+            audio = torch.tensor(filtered, device=audio.device)
+        except Exception as e:
+            logger.warning(f"Advanced audio filtering failed, using basic processing: {e}")
+        
+        # Ensure audio has correct amplitude
+        audio = audio * 0.9  # Slightly reduce volume to prevent clipping
+        
+        return audio
+
+    def _remove_silence(
+        self, 
+        audio: torch.Tensor, 
+        threshold: float = 0.015, 
+        min_silence_duration: float = 0.2
+    ) -> torch.Tensor:
+        """
+        Remove silence from audio while preserving speech rhythm.
+        
+        Args:
+            audio: Input audio tensor
+            threshold: Energy threshold for silence detection
+            min_silence_duration: Minimum silence duration in seconds
+            
+        Returns:
+            Audio with silence removed
+        """
+        # Convert to numpy for easier processing
+        audio_np = audio.cpu().numpy()
+        
+        # Calculate energy
+        energy = np.abs(audio_np)
+        
+        # Find regions above threshold (speech)
+        is_speech = energy > threshold
+        
+        # Convert min_silence_duration to samples
+        min_silence_samples = int(min_silence_duration * self.sample_rate)
+        
+        # Find speech segments
+        speech_segments = []
+        in_speech = False
+        speech_start = 0
+        
+        for i in range(len(is_speech)):
+            if is_speech[i] and not in_speech:
+                # Start of speech segment
+                in_speech = True
+                speech_start = i
+            elif not is_speech[i] and in_speech:
+                # Potential end of speech segment
+                # Only end if silence is long enough
+                silence_count = 0
+                for j in range(i, min(len(is_speech), i + min_silence_samples)):
+                    if not is_speech[j]:
+                        silence_count += 1
+                    else:
+                        break
+                
+                if silence_count >= min_silence_samples:
+                    # End of speech segment
+                    in_speech = False
+                    speech_segments.append((speech_start, i))
+        
+        # Handle case where audio ends during speech
+        if in_speech:
+            speech_segments.append((speech_start, len(is_speech)))
+        
+        # If no speech segments found, return original audio
+        if not speech_segments:
+            logger.warning("No speech segments detected, returning original audio")
+            return audio
+        
+        # Add small buffer around segments
+        buffer_samples = int(0.05 * self.sample_rate)  # 50ms buffer
+        processed_segments = []
+        
+        for start, end in speech_segments:
+            buffered_start = max(0, start - buffer_samples)
+            buffered_end = min(len(audio_np), end + buffer_samples)
+            processed_segments.append(audio_np[buffered_start:buffered_end])
+        
+        # Concatenate all segments with small pauses between them
+        small_pause = np.zeros(int(0.15 * self.sample_rate))  # 150ms pause
+        result = processed_segments[0]
+        
+        for segment in processed_segments[1:]:
+            result = np.concatenate([result, small_pause, segment])
+        
+        return torch.tensor(result, device=audio.device)
+
+    def _enhance_speech(self, audio: torch.Tensor) -> torch.Tensor:
+        """Enhance speech quality for better cloning results."""
+        # This is a placeholder for more advanced speech enhancement
+        # In a production implementation, you could add:
+        # - Noise reduction
+        # - Equalization for speech frequencies
+        # - Gentle compression for better dynamics
+        return audio
+
+    async def clone_voice(
+        self,
+        audio_file: Union[UploadFile, BinaryIO, str],
+        voice_name: str,
+        transcript: Optional[str] = None,
+        description: Optional[str] = None,
+        speaker_id: Optional[int] = None  # Make this optional
+    ) -> ClonedVoice:
+        """
+        Clone a voice from an audio file.
+        
+        Args:
+            audio_file: Audio file with the voice to clone
+            voice_name: Name for the cloned voice
+            transcript: Transcript of the audio (optional)
+            description: Description of the voice (optional)
+            speaker_id: Speaker ID to use (default: auto-assigned)
+            
+        Returns:
+            ClonedVoice object with voice information
+        """
+        logger.info(f"Cloning new voice '{voice_name}' from audio file")
+        
+        # Process the audio file
+        processed_audio, provided_transcript, duration = await self.process_audio_file(
+            audio_file, transcript
+        )
+        
+        # Use a better speaker ID assignment - use a small number similar to the built-in voices
+        # This prevents issues with the speaker ID being interpreted as speech
+        if speaker_id is None:
+            # Use a number between 10-20 to avoid conflicts with built-in voices (0-5)
+            # but not too large like 999 which might cause issues
+            existing_ids = [v.speaker_id for v in self.cloned_voices.values()]
+            for potential_id in range(10, 20):
+                if potential_id not in existing_ids:
+                    speaker_id = potential_id
+                    break
+            else:
+                # If all IDs in range are taken, use a fallback
+                speaker_id = 10
+        
+        # Generate a unique ID for the voice
+        voice_id = f"{int(time.time())}_{voice_name.lower().replace(' ', '_')}"
+        
+        # Create directory for the voice
+        voice_dir = os.path.join(CLONED_VOICES_DIR, voice_id)
+        os.makedirs(voice_dir, exist_ok=True)
+        
+        # Save the processed audio
+        audio_path = os.path.join(voice_dir, "reference.wav")
+        torchaudio.save(audio_path, processed_audio.unsqueeze(0).cpu(), self.sample_rate)
+        
+        # Save the transcript if provided
+        if provided_transcript:
+            transcript_path = os.path.join(voice_dir, "transcript.txt")
+            with open(transcript_path, "w") as f:
+                f.write(provided_transcript)
+        
+        # Create and save voice info
+        voice_info = ClonedVoice(
+            id=voice_id,
+            name=voice_name,
+            created_at=time.time(),
+            speaker_id=speaker_id,
+            description=description,
+            audio_duration=duration,
+            sample_count=len(processed_audio)
+        )
+        
+        # Save voice info as JSON
+        import json
+        with open(os.path.join(voice_dir, "info.json"), "w") as f:
+            f.write(json.dumps(voice_info.dict()))
+        
+        # Add to cloned voices dictionary
+        self.cloned_voices[voice_id] = voice_info
+        
+        logger.info(f"Voice '{voice_name}' cloned successfully with ID: {voice_id} and speaker_id: {speaker_id}")
+        
+        return voice_info
+
+    def get_voice_context(self, voice_id: str) -> List[Segment]:
+        """
+        Get context segments for a cloned voice.
+        
+        Args:
+            voice_id: ID of the cloned voice
+            
+        Returns:
+            List of context segments for the voice
+        """
+        if voice_id not in self.cloned_voices:
+            logger.warning(f"Voice ID {voice_id} not found")
+            return []
+        
+        voice = self.cloned_voices[voice_id]
+        voice_dir = os.path.join(CLONED_VOICES_DIR, voice_id)
+        audio_path = os.path.join(voice_dir, "reference.wav")
+        
+        if not os.path.exists(audio_path):
+            logger.error(f"Audio file for voice {voice_id} not found at {audio_path}")
+            return []
+        
+        try:
+            # Load the audio
+            audio, sr = torchaudio.load(audio_path)
+            audio = audio.squeeze(0)
+            
+            # Resample if necessary
+            if sr != self.sample_rate:
+                audio = torchaudio.functional.resample(
+                    audio, orig_freq=sr, new_freq=self.sample_rate
+                )
+            
+            # Trim to a maximum of 5 seconds to avoid sequence length issues
+            # This is a balance between voice quality and model limitations
+            max_samples = 5 * self.sample_rate  # 5 seconds
+            if audio.shape[0] > max_samples:
+                logger.info(f"Trimming voice sample from {audio.shape[0]} to {max_samples} samples")
+                # Take from beginning for better voice characteristics 
+                audio = audio[:max_samples]
+            
+            # Load transcript if available
+            transcript_path = os.path.join(voice_dir, "transcript.txt")
+            transcript = ""
+            if os.path.exists(transcript_path):
+                with open(transcript_path, "r") as f:
+                    full_transcript = f.read()
+                    # Take a portion of transcript that roughly matches our audio portion
+                    words = full_transcript.split()
+                    # Estimate 3 words per second as a rough average
+                    word_count = min(len(words), int(5 * 3))  # 5 seconds * 3 words/second
+                    transcript = " ".join(words[:word_count])
+            else:
+                transcript = f"Voice sample for {voice.name}"
+            
+            # Create context segment
+            segment = Segment(
+                text=transcript,
+                speaker=voice.speaker_id,
+                audio=audio.to(self.device)
+            )
+            
+            logger.info(f"Created voice context segment with {audio.shape[0]/self.sample_rate:.1f}s audio")
+            return [segment]
+            
+        except Exception as e:
+            logger.error(f"Error getting voice context for {voice_id}: {e}")
+            return []
+        
+    def list_voices(self) -> List[ClonedVoice]:
+        """List all available cloned voices."""
+        return list(self.cloned_voices.values())
+
+    def delete_voice(self, voice_id: str) -> bool:
+        """
+        Delete a cloned voice.
+        
+        Args:
+            voice_id: ID of the voice to delete
+            
+        Returns:
+            True if successful, False otherwise
+        """
+        if voice_id not in self.cloned_voices:
+            return False
+        
+        voice_dir = os.path.join(CLONED_VOICES_DIR, voice_id)
+        if os.path.exists(voice_dir):
+            try:
+                import shutil
+                shutil.rmtree(voice_dir)
+                del self.cloned_voices[voice_id]
+                return True
+            except Exception as e:
+                logger.error(f"Error deleting voice {voice_id}: {e}")
+                return False
+        
+        return False
+
+    async def clone_voice_from_youtube(
+        self,  # Don't forget the self parameter for class methods
+        youtube_url: str,
+        voice_name: str,
+        start_time: int = 0,
+        duration: int = 180,
+        description: str = None
+    ) -> ClonedVoice:
+        """
+        Clone a voice from a YouTube video.
+        
+        Args:
+            youtube_url: URL of the YouTube video
+            voice_name: Name for the cloned voice
+            start_time: Start time in seconds
+            duration: Duration to extract in seconds
+            description: Optional description of the voice
+            
+        Returns:
+            ClonedVoice object with voice information
+        """
+        logger.info(f"Cloning voice '{voice_name}' from YouTube: {youtube_url}")
+        
+        # Create temporary directory for processing
+        with tempfile.TemporaryDirectory() as temp_dir:
+            # Step 1: Download audio from YouTube
+            audio_path = await self._download_youtube_audio(youtube_url, temp_dir, start_time, duration)
+            
+            # Step 2: Generate transcript using Whisper
+            transcript = await self._generate_transcript(audio_path)
+            
+            # Step 3: Clone the voice using the extracted audio and transcript
+            voice = await self.clone_voice(
+                audio_file=audio_path,
+                voice_name=voice_name,
+                transcript=transcript,
+                description=description or f"Voice cloned from YouTube: {youtube_url}"
+            )
+            
+            return voice
+
+    async def _download_youtube_audio(
+        self,  # Don't forget the self parameter
+        url: str, 
+        output_dir: str, 
+        start_time: int = 0, 
+        duration: int = 180
+    ) -> str:
+        """
+        Download audio from a YouTube video.
+        
+        Args:
+            url: YouTube URL
+            output_dir: Directory to save the audio
+            start_time: Start time in seconds
+            duration: Duration to extract in seconds
+            
+        Returns:
+            Path to the downloaded audio file
+        """
+        output_path = os.path.join(output_dir, "youtube_audio.wav")
+        
+        # Configure yt-dlp options
+        ydl_opts = {
+            'format': 'bestaudio/best',
+            'postprocessors': [{
+                'key': 'FFmpegExtractAudio',
+                'preferredcodec': 'wav',
+                'preferredquality': '192',
+            }],
+            'outtmpl': output_path.replace(".wav", ""),
+            'quiet': True,
+            'no_warnings': True
+        }
+        
+        # Download the video
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            ydl.download([url])
+        
+        # Trim the audio to the specified segment
+        if start_time > 0 or duration < float('inf'):
+            import ffmpeg
+            trimmed_path = os.path.join(output_dir, "trimmed_audio.wav")
+            
+            # Use ffmpeg to trim the audio
+            (
+                ffmpeg.input(output_path)
+                .audio
+                .filter('atrim', start=start_time, duration=duration)
+                .output(trimmed_path)
+                .run(quiet=True, overwrite_output=True)
+            )
+            
+            return trimmed_path
+        
+        return output_path
+
+    async def _generate_transcript(self, audio_path: str) -> str:
+        """
+        Generate transcript from audio using Whisper.
+        
+        Args:
+            audio_path: Path to the audio file
+            
+        Returns:
+            Transcript text
+        """
+        # Load Whisper model (use small model for faster processing)
+        model = whisper.load_model("small")
+        
+        # Transcribe the audio
+        result = model.transcribe(audio_path)
+        
+        return result["text"]
+
+    def generate_speech(
+        self,
+        text: str,
+        voice_id: str,
+        temperature: float = 0.65,
+        topk: int = 30,
+        max_audio_length_ms: int = 15000
+    ) -> torch.Tensor:
+        """
+        Generate speech with a cloned voice.
+        Args:
+            text: Text to synthesize
+            voice_id: ID of the cloned voice to use
+            temperature: Sampling temperature (lower = more stable, higher = more varied)
+            topk: Top-k sampling parameter
+            max_audio_length_ms: Maximum audio length in milliseconds
+        Returns:
+            Generated audio tensor
+        """
+        # Remove any async/await keywords - this is a synchronous function
+        if voice_id not in self.cloned_voices:
+            raise ValueError(f"Voice ID {voice_id} not found")
+        voice = self.cloned_voices[voice_id]
+        context = self.get_voice_context(voice_id)
+        if not context:
+            raise ValueError(f"Could not get context for voice {voice_id}")
+        # Preprocess text for better pronunciation
+        processed_text = self._preprocess_text(text)
+        logger.info(f"Generating speech with voice '{voice.name}' (ID: {voice_id}, speaker: {voice.speaker_id})")
+        try:
+            # Check if text is too long and should be split
+            if len(processed_text) > 200:
+                logger.info(f"Text is long ({len(processed_text)} chars), splitting for better quality")
+                from app.prompt_engineering import split_into_segments
+                # Split text into manageable segments
+                segments = split_into_segments(processed_text, max_chars=150)
+                logger.info(f"Split text into {len(segments)} segments")
+                all_audio_chunks = []
+                # Process each segment
+                for i, segment_text in enumerate(segments):
+                    logger.info(f"Generating segment {i+1}/{len(segments)}")
+                    # Generate this segment - using plain text without formatting
+                    segment_audio = self.generator.generate(
+                        text=segment_text,  # Use plain text, no formatting
+                        speaker=voice.speaker_id,
+                        context=context,
+                        max_audio_length_ms=min(max_audio_length_ms, 10000),
+                        temperature=temperature,
+                        topk=topk,
+                    )
+                    all_audio_chunks.append(segment_audio)
+                    # Use this segment as context for the next one for consistency
+                    if i < len(segments) - 1:
+                        context = [
+                            Segment(
+                                text=segment_text,
+                                speaker=voice.speaker_id,
+                                audio=segment_audio
+                            )
+                        ]
+                # Combine chunks with small silence between them
+                if len(all_audio_chunks) == 1:
+                    audio = all_audio_chunks[0]
+                else:
+                    silence_samples = int(0.1 * self.sample_rate)  # 100ms silence
+                    silence = torch.zeros(silence_samples, device=all_audio_chunks[0].device)
+                    # Join segments with silence
+                    audio_parts = []
+                    for i, chunk in enumerate(all_audio_chunks):
+                        audio_parts.append(chunk)
+                        if i < len(all_audio_chunks) - 1:  # Don't add silence after the last chunk
+                            audio_parts.append(silence)
+                    # Concatenate all parts
+                    audio = torch.cat(audio_parts)
+                return audio
+            else:
+                # For short text, generate directly - using plain text without formatting
+                audio = self.generator.generate(
+                    text=processed_text,  # Use plain text, no formatting
+                    speaker=voice.speaker_id,
+                    context=context,
+                    max_audio_length_ms=max_audio_length_ms,
+                    temperature=temperature,
+                    topk=topk,
+                )
+                return audio
+        except Exception as e:
+            logger.error(f"Error generating speech with voice {voice_id}: {e}")
+            raise
+            
+    def _preprocess_text(self, text: str) -> str:
+        """Preprocess text for better pronunciation and voice cloning."""
+        # Make sure text ends with punctuation for better phrasing
+        text = text.strip()
+        if not text.endswith(('.', '?', '!', ';')):
+            text = text + '.'
+        
+        return text

app/voice_embeddings.py

@@ -0,0 +1,85 @@
+"""Voice embeddings for consistent voice generation."""
+import os
+import torch
+import torchaudio
+import numpy as np
+from typing import Dict
+
+# Path to store voice samples
+VOICE_SAMPLES_DIR = os.path.join(os.path.dirname(__file__), "voice_samples")
+os.makedirs(VOICE_SAMPLES_DIR, exist_ok=True)
+
+# Dictionary to store voice embeddings/samples
+VOICE_DICT: Dict[str, torch.Tensor] = {}
+
+
+def initialize_voices(sample_rate: int = 24000):
+    """Initialize voice dictionary with consistent samples."""
+    # Generate consistent seed audio for each voice
+    for voice_id in range(6):
+        voice_name = ["alloy", "echo", "fable", "onyx", "nova", "shimmer"][voice_id]
+        
+        # Create deterministic audio sample for each voice
+        np.random.seed(voice_id + 42)  # Use a fixed seed based on voice ID
+        
+        # Generate 1 second of "seed" audio with deterministic characteristics
+        # This differs per voice but remains consistent across runs
+        duration = 1.0  # seconds
+        t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
+        
+        # Create a distinctive waveform for each voice
+        if voice_id == 0:  # alloy - rich mid tones
+            freq1, freq2 = 220, 440
+            audio = 0.5 * np.sin(2 * np.pi * freq1 * t) + 0.3 * np.sin(2 * np.pi * freq2 * t)
+        elif voice_id == 1:  # echo - reverberant
+            freq = 330
+            audio = np.sin(2 * np.pi * freq * t) * np.exp(-t * 3)
+        elif voice_id == 2:  # fable - bright, higher pitch
+            freq = 523
+            audio = 0.7 * np.sin(2 * np.pi * freq * t)
+        elif voice_id == 3:  # onyx - deep and resonant
+            freq = 165
+            audio = 0.8 * np.sin(2 * np.pi * freq * t)
+        elif voice_id == 4:  # nova - warm and smooth
+            freq1, freq2 = 392, 196
+            audio = 0.4 * np.sin(2 * np.pi * freq1 * t) + 0.4 * np.sin(2 * np.pi * freq2 * t)
+        else:  # shimmer - airy and light
+            freq1, freq2, freq3 = 587, 880, 1174
+            audio = 0.3 * np.sin(2 * np.pi * freq1 * t) + 0.2 * np.sin(2 * np.pi * freq2 * t) + 0.1 * np.sin(2 * np.pi * freq3 * t)
+            
+        # Normalize
+        audio = audio / np.max(np.abs(audio))
+        
+        # Convert to tensor
+        audio_tensor = torch.tensor(audio, dtype=torch.float32)
+        
+        # Store the audio tensor
+        VOICE_DICT[voice_name] = audio_tensor
+        
+        # Save as wav for reference
+        save_path = os.path.join(VOICE_SAMPLES_DIR, f"{voice_name}_seed.wav")
+        torchaudio.save(save_path, audio_tensor.unsqueeze(0), sample_rate)
+        
+        print(f"Initialized voice seed for {voice_name}")
+
+
+def get_voice_sample(voice_name: str) -> torch.Tensor:
+    """Get the voice sample for a given voice name."""
+    if not VOICE_DICT:
+        initialize_voices()
+    
+    if voice_name in VOICE_DICT:
+        return VOICE_DICT[voice_name]
+        
+    # Default to alloy if voice not found
+    print(f"Voice {voice_name} not found, defaulting to alloy")
+    return VOICE_DICT["alloy"]
+
+
+def update_voice_sample(voice_name: str, audio: torch.Tensor):
+    """Update the voice sample with recently generated audio."""
+    # Only update if we've already initialized
+    if VOICE_DICT:
+        # Take the last second of audio (or whatever is available)
+        sample_length = min(24000, audio.shape[0])
+        VOICE_DICT[voice_name] = audio[-sample_length:].detach().cpu()

app/voice_enhancement.py

@@ -0,0 +1,581 @@
+"""Advanced voice enhancement and consistency system for CSM-1B."""
+import os
+import torch
+import torchaudio
+import numpy as np
+import soundfile as sf
+from typing import Dict, List, Optional, Tuple
+import logging
+from dataclasses import dataclass
+from scipy import signal
+
+# Setup logging
+logger = logging.getLogger(__name__)
+
+# Define persistent paths
+VOICE_REFERENCES_DIR = "/app/voice_references"
+VOICE_PROFILES_DIR = "/app/voice_profiles"
+
+# Ensure directories exist
+os.makedirs(VOICE_REFERENCES_DIR, exist_ok=True)
+os.makedirs(VOICE_PROFILES_DIR, exist_ok=True)
+
+@dataclass
+class VoiceProfile:
+    """Detailed voice profile with acoustic characteristics."""
+    name: str
+    speaker_id: int
+    # Acoustic parameters
+    pitch_range: Tuple[float, float]  # Min/max pitch in Hz
+    intensity_range: Tuple[float, float]  # Min/max intensity (volume)
+    spectral_tilt: float  # Brightness vs. darkness
+    prosody_pattern: str  # Pattern of intonation and rhythm
+    speech_rate: float  # Relative speech rate (1.0 = normal)
+    formant_shift: float  # Formant frequency shift (1.0 = no shift)
+    # Reference audio
+    reference_segments: List[torch.Tensor]
+    # Normalization parameters
+    target_rms: float = 0.2
+    target_peak: float = 0.95
+    
+    def get_enhancement_params(self) -> Dict:
+        """Get parameters for enhancing generated audio."""
+        return {
+            "target_rms": self.target_rms,
+            "target_peak": self.target_peak,
+            "pitch_range": self.pitch_range,
+            "formant_shift": self.formant_shift,
+            "speech_rate": self.speech_rate,
+            "spectral_tilt": self.spectral_tilt
+        }
+
+# Voice profiles with carefully tuned parameters
+VOICE_PROFILES = {
+    "alloy": VoiceProfile(
+        name="alloy",
+        speaker_id=0,
+        pitch_range=(85, 180),  # Hz - balanced range
+        intensity_range=(0.15, 0.3),  # moderate intensity
+        spectral_tilt=0.0,  # neutral tilt
+        prosody_pattern="balanced",
+        speech_rate=1.0,  # normal rate
+        formant_shift=1.0,  # no shift
+        reference_segments=[],
+        target_rms=0.2,
+        target_peak=0.95
+    ),
+    "echo": VoiceProfile(
+        name="echo",
+        speaker_id=1,
+        pitch_range=(75, 165),  # Hz - lower, resonant 
+        intensity_range=(0.2, 0.35),  # slightly stronger
+        spectral_tilt=-0.2,  # more low frequencies
+        prosody_pattern="deliberate",
+        speech_rate=0.95,  # slightly slower
+        formant_shift=0.95,  # slightly lower formants
+        reference_segments=[],
+        target_rms=0.22,  # slightly louder
+        target_peak=0.95
+    ),
+    "fable": VoiceProfile(
+        name="fable",
+        speaker_id=2,
+        pitch_range=(120, 250),  # Hz - higher range
+        intensity_range=(0.15, 0.28),  # moderate intensity
+        spectral_tilt=0.2,  # more high frequencies
+        prosody_pattern="animated",
+        speech_rate=1.05,  # slightly faster
+        formant_shift=1.05,  # slightly higher formants
+        reference_segments=[],
+        target_rms=0.19,
+        target_peak=0.95
+    ),
+    "onyx": VoiceProfile(
+        name="onyx",
+        speaker_id=3,
+        pitch_range=(65, 150),  # Hz - deeper range
+        intensity_range=(0.18, 0.32),  # moderate-strong
+        spectral_tilt=-0.3,  # more low frequencies
+        prosody_pattern="authoritative",
+        speech_rate=0.93,  # slightly slower
+        formant_shift=0.9,  # lower formants
+        reference_segments=[],
+        target_rms=0.23,  # stronger
+        target_peak=0.95
+    ),
+    "nova": VoiceProfile(
+        name="nova",
+        speaker_id=4,
+        pitch_range=(90, 200),  # Hz - warm midrange
+        intensity_range=(0.15, 0.27),  # moderate
+        spectral_tilt=-0.1,  # slightly warm
+        prosody_pattern="flowing",
+        speech_rate=1.0,  # normal rate
+        formant_shift=1.0,  # no shift
+        reference_segments=[],
+        target_rms=0.2,
+        target_peak=0.95
+    ),
+    "shimmer": VoiceProfile(
+        name="shimmer",
+        speaker_id=5,
+        pitch_range=(140, 280),  # Hz - brighter, higher
+        intensity_range=(0.15, 0.25),  # moderate-light
+        spectral_tilt=0.3,  # more high frequencies
+        prosody_pattern="light",
+        speech_rate=1.07,  # slightly faster
+        formant_shift=1.1,  # higher formants
+        reference_segments=[],
+        target_rms=0.18,  # slightly softer
+        target_peak=0.95
+    )
+}
+
+# Voice-specific prompt templates - crafted to establish voice identity clearly
+VOICE_PROMPTS = {
+    "alloy": [
+        "Hello, I'm Alloy. I speak with a balanced, natural tone that's easy to understand.",
+        "This is Alloy speaking. My voice is designed to be clear and conversational.",
+        "Alloy here - I have a neutral, friendly voice with balanced tone qualities."
+    ],
+    "echo": [
+        "Hello, I'm Echo. I speak with a resonant, deeper voice that carries well.",
+        "This is Echo speaking. My voice has a rich, resonant quality with depth.",
+        "Echo here - My voice is characterized by its warm, resonant tones."
+    ],
+    "fable": [
+        "Hello, I'm Fable. I speak with a bright, higher-pitched voice that's full of energy.",
+        "This is Fable speaking. My voice is characterized by its clear, bright quality.",
+        "Fable here - My voice is light, articulate, and slightly higher-pitched."
+    ],
+    "onyx": [
+        "Hello, I'm Onyx. I speak with a deep, authoritative voice that commands attention.",
+        "This is Onyx speaking. My voice has a powerful, deep quality with gravitas.",
+        "Onyx here - My voice is characterized by its depth and commanding presence."
+    ],
+    "nova": [
+        "Hello, I'm Nova. I speak with a warm, pleasant mid-range voice that's easy to listen to.",
+        "This is Nova speaking. My voice has a smooth, harmonious quality.",
+        "Nova here - My voice is characterized by its warm, friendly mid-tones."
+    ],
+    "shimmer": [
+        "Hello, I'm Shimmer. I speak with a light, bright voice that's expressive and clear.",
+        "This is Shimmer speaking. My voice has an airy, higher-pitched quality.",
+        "Shimmer here - My voice is characterized by its bright, crystalline tones."
+    ]
+}
+
+def initialize_voice_profiles():
+    """Initialize voice profiles with default settings.
+    
+    This function loads existing voice profiles from disk if available,
+    or initializes them with default settings.
+    """
+    global VOICE_PROFILES
+    
+    # Try to load existing profiles from persistent storage
+    profile_path = os.path.join(VOICE_PROFILES_DIR, "voice_profiles.pt")
+    
+    if os.path.exists(profile_path):
+        try:
+            logger.info(f"Loading voice profiles from {profile_path}")
+            saved_profiles = torch.load(profile_path)
+            
+            # Update existing profiles with saved data
+            for name, data in saved_profiles.items():
+                if name in VOICE_PROFILES:
+                    VOICE_PROFILES[name].reference_segments = [
+                        seg.to(torch.device("cpu")) for seg in data.get('reference_segments', [])
+                    ]
+            
+            logger.info(f"Loaded voice profiles for {len(saved_profiles)} voices")
+        except Exception as e:
+            logger.error(f"Error loading voice profiles: {e}")
+            logger.info("Using default voice profiles")
+    else:
+        logger.info("No saved voice profiles found, using defaults")
+    
+    # Ensure all voices have at least empty reference segments
+    for name, profile in VOICE_PROFILES.items():
+        if not hasattr(profile, 'reference_segments'):
+            profile.reference_segments = []
+    
+    logger.info(f"Voice profiles initialized for {len(VOICE_PROFILES)} voices")
+    return VOICE_PROFILES
+
+def normalize_audio(audio: torch.Tensor, target_rms: float = 0.2, target_peak: float = 0.95) -> torch.Tensor:
+    """Apply professional-grade normalization to audio.
+    
+    Args:
+        audio: Audio tensor
+        target_rms: Target RMS level for normalization
+        target_peak: Target peak level for limiting
+        
+    Returns:
+        Normalized audio tensor
+    """
+    # Ensure audio is on CPU for processing
+    audio_cpu = audio.detach().cpu()
+    
+    # Handle silent audio
+    if audio_cpu.abs().max() < 1e-6:
+        logger.warning("Audio is nearly silent, returning original")
+        return audio
+    
+    # Calculate current RMS
+    current_rms = torch.sqrt(torch.mean(audio_cpu ** 2))
+    
+    # Apply RMS normalization
+    if current_rms > 0:
+        gain = target_rms / current_rms
+        normalized = audio_cpu * gain
+    else:
+        normalized = audio_cpu
+        
+    # Apply peak limiting
+    current_peak = normalized.abs().max()
+    if current_peak > target_peak:
+        normalized = normalized * (target_peak / current_peak)
+    
+    # Return to original device
+    return normalized.to(audio.device)
+
+def apply_anti_muffling(audio: torch.Tensor, sample_rate: int, clarity_boost: float = 1.2) -> torch.Tensor:
+    """Apply anti-muffling to improve clarity.
+    
+    Args:
+        audio: Audio tensor
+        sample_rate: Audio sample rate
+        clarity_boost: Amount of high frequency boost (1.0 = no boost)
+        
+    Returns:
+        Processed audio tensor
+    """
+    # Convert to numpy for filtering
+    audio_np = audio.detach().cpu().numpy()
+    
+    try:
+        # Design a high shelf filter to boost high frequencies
+        # Use a standard high-shelf filter that's supported by scipy.signal
+        # We'll use a second-order Butterworth high-pass filter as an alternative
+        cutoff = 2000  # Hz
+        b, a = signal.butter(2, cutoff/(sample_rate/2), btype='high', analog=False)
+        
+        # Apply the filter with the clarity boost gain
+        boosted = signal.filtfilt(b, a, audio_np, axis=0) * clarity_boost
+        
+        # Mix with original to maintain some warmth
+        mix_ratio = 0.7  # 70% processed, 30% original
+        processed = mix_ratio * boosted + (1-mix_ratio) * audio_np
+        
+    except Exception as e:
+        logger.warning(f"Audio enhancement failed, using original: {e}")
+        # Return original audio if enhancement fails
+        return audio
+    
+    # Convert back to tensor on original device
+    return torch.tensor(processed, dtype=audio.dtype, device=audio.device)
+
+def enhance_audio(audio: torch.Tensor, sample_rate: int, voice_profile: VoiceProfile) -> torch.Tensor:
+    """Apply comprehensive audio enhancement based on voice profile.
+    
+    Args:
+        audio: Audio tensor
+        sample_rate: Audio sample rate
+        voice_profile: Voice profile containing enhancement parameters
+        
+    Returns:
+        Enhanced audio tensor
+    """
+    if audio is None or audio.numel() == 0:
+        logger.error("Cannot enhance empty audio")
+        return audio
+        
+    try:
+        # Step 1: Normalize audio levels
+        params = voice_profile.get_enhancement_params()
+        normalized = normalize_audio(
+            audio, 
+            target_rms=params["target_rms"],
+            target_peak=params["target_peak"]
+        )
+        
+        # Step 2: Apply anti-muffling based on spectral tilt
+        # Positive tilt means brighter voice so less clarity boost needed
+        clarity_boost = 1.0 + max(0, -params["spectral_tilt"]) * 0.5
+        clarified = apply_anti_muffling(
+            normalized,
+            sample_rate,
+            clarity_boost=clarity_boost
+        )
+        
+        # Log the enhancement
+        logger.debug(
+            f"Enhanced audio for {voice_profile.name}: "
+            f"RMS: {audio.pow(2).mean().sqrt().item():.3f}->{clarified.pow(2).mean().sqrt().item():.3f}, "
+            f"Peak: {audio.abs().max().item():.3f}->{clarified.abs().max().item():.3f}"
+        )
+        
+        return clarified
+        
+    except Exception as e:
+        logger.error(f"Error in audio enhancement: {e}")
+        return audio  # Return original audio if enhancement fails
+
+def validate_generated_audio(
+    audio: torch.Tensor, 
+    voice_name: str, 
+    sample_rate: int,
+    min_expected_duration: float = 0.5
+) -> Tuple[bool, torch.Tensor, str]:
+    """Validate and fix generated audio.
+    
+    Args:
+        audio: Audio tensor to validate
+        voice_name: Name of the voice used
+        sample_rate: Audio sample rate
+        min_expected_duration: Minimum expected duration in seconds
+        
+    Returns:
+        Tuple of (is_valid, fixed_audio, message)
+    """
+    if audio is None:
+        return False, torch.zeros(1), "Audio is None"
+    
+    # Check for NaN values
+    if torch.isnan(audio).any():
+        logger.warning(f"Audio for {voice_name} contains NaN values, replacing with zeros")
+        audio = torch.where(torch.isnan(audio), torch.zeros_like(audio), audio)
+    
+    # Check audio duration
+    duration = audio.shape[0] / sample_rate
+    if duration < min_expected_duration:
+        logger.warning(f"Audio for {voice_name} is too short ({duration:.2f}s < {min_expected_duration}s)")
+        return False, audio, f"Audio too short: {duration:.2f}s"
+    
+    # Check for silent sections - this can indicate generation problems
+    rms = torch.sqrt(torch.mean(audio ** 2))
+    if rms < 0.01:  # Very low RMS indicates near silence
+        logger.warning(f"Audio for {voice_name} is nearly silent (RMS: {rms:.6f})")
+        return False, audio, f"Audio nearly silent: RMS = {rms:.6f}"
+    
+    # Check if audio suddenly cuts off - this detects premature stopping
+    # Calculate RMS in the last 100ms
+    last_samples = int(0.1 * sample_rate)
+    if audio.shape[0] > last_samples:
+        end_rms = torch.sqrt(torch.mean(audio[-last_samples:] ** 2))
+        if end_rms > 0.1:  # High RMS at the end suggests an abrupt cutoff
+            logger.warning(f"Audio for {voice_name} may have cut off prematurely (end RMS: {end_rms:.3f})")
+            return True, audio, "Audio may have cut off prematurely"
+    
+    return True, audio, "Audio validation passed"
+
+def create_voice_segments(app_state, regenerate: bool = False):
+    """Create high-quality voice reference segments.
+    
+    Args:
+        app_state: Application state containing generator
+        regenerate: Whether to regenerate existing references
+    """
+    generator = app_state.generator
+    if not generator:
+        logger.error("Cannot create voice segments: generator not available")
+        return
+    
+    # Use persistent directory for voice reference segments
+    os.makedirs(VOICE_REFERENCES_DIR, exist_ok=True)
+    
+    for voice_name, profile in VOICE_PROFILES.items():
+        voice_dir = os.path.join(VOICE_REFERENCES_DIR, voice_name)
+        os.makedirs(voice_dir, exist_ok=True)
+        
+        # Check if we already have references
+        if not regenerate and profile.reference_segments:
+            logger.info(f"Voice {voice_name} already has {len(profile.reference_segments)} reference segments")
+            continue
+        
+        # Get prompts for this voice
+        prompts = VOICE_PROMPTS[voice_name]
+        
+        # Generate reference segments
+        logger.info(f"Generating reference segments for voice: {voice_name}")
+        reference_segments = []
+        
+        for i, prompt in enumerate(prompts):
+            ref_path = os.path.join(voice_dir, f"{voice_name}_ref_{i}.wav")
+            
+            # Skip if file exists and we're not regenerating
+            if not regenerate and os.path.exists(ref_path):
+                try:
+                    # Load existing reference
+                    audio_tensor, sr = torchaudio.load(ref_path)
+                    if sr != generator.sample_rate:
+                        audio_tensor = torchaudio.functional.resample(
+                            audio_tensor.squeeze(0), orig_freq=sr, new_freq=generator.sample_rate
+                        )
+                    else:
+                        audio_tensor = audio_tensor.squeeze(0)
+                    reference_segments.append(audio_tensor.to(generator.device))
+                    logger.info(f"Loaded existing reference {i+1}/{len(prompts)} for {voice_name}")
+                    continue
+                except Exception as e:
+                    logger.warning(f"Failed to load existing reference {i+1} for {voice_name}: {e}")
+            
+            try:
+                # Use a lower temperature for more stability in reference samples
+                logger.info(f"Generating reference {i+1}/{len(prompts)} for {voice_name}: '{prompt}'")
+                
+                # We want references to be as clean as possible
+                audio = generator.generate(
+                    text=prompt,
+                    speaker=profile.speaker_id,
+                    context=[],  # No context for initial samples to prevent voice bleed
+                    max_audio_length_ms=6000,  # Shorter for more control
+                    temperature=0.7,  # Lower temperature for more stability
+                    topk=30,  # More focused sampling
+                )
+                
+                # Validate and enhance the audio
+                is_valid, audio, message = validate_generated_audio(
+                    audio, voice_name, generator.sample_rate
+                )
+                
+                if is_valid:
+                    # Enhance the audio
+                    audio = enhance_audio(audio, generator.sample_rate, profile)
+                    
+                    # Save the reference to persistent storage
+                    torchaudio.save(ref_path, audio.unsqueeze(0).cpu(), generator.sample_rate)
+                    reference_segments.append(audio)
+                    logger.info(f"Generated reference {i+1} for {voice_name}: {message}")
+                else:
+                    logger.warning(f"Invalid reference for {voice_name}: {message}")
+                    # Try again with different settings if invalid
+                    if i < len(prompts) - 1:
+                        logger.info(f"Trying again with next prompt")
+                        continue
+            
+            except Exception as e:
+                logger.error(f"Error generating reference for {voice_name}: {e}")
+        
+        # Update the voice profile with references
+        if reference_segments:
+            VOICE_PROFILES[voice_name].reference_segments = reference_segments
+            logger.info(f"Updated {voice_name} with {len(reference_segments)} reference segments")
+    
+    # Save the updated profiles to persistent storage
+    save_voice_profiles()
+
+def get_voice_segments(voice_name: str, device: torch.device) -> List:
+    """Get context segments for a given voice.
+    
+    Args:
+        voice_name: Name of the voice to use
+        device: Device to place tensors on
+        
+    Returns:
+        List of context segments
+    """
+    from app.models import Segment
+    
+    if voice_name not in VOICE_PROFILES:
+        logger.warning(f"Voice {voice_name} not found, defaulting to alloy")
+        voice_name = "alloy"
+    
+    profile = VOICE_PROFILES[voice_name]
+    
+    # If we don't have reference segments yet, create them
+    if not profile.reference_segments:
+        try:
+            # Try to load from disk - use persistent storage
+            voice_dir = os.path.join(VOICE_REFERENCES_DIR, voice_name)
+            
+            if os.path.exists(voice_dir):
+                reference_segments = []
+                prompts = VOICE_PROMPTS[voice_name]
+                
+                for i, prompt in enumerate(prompts):
+                    ref_path = os.path.join(voice_dir, f"{voice_name}_ref_{i}.wav")
+                    if os.path.exists(ref_path):
+                        audio_tensor, sr = torchaudio.load(ref_path)
+                        audio_tensor = audio_tensor.squeeze(0)
+                        reference_segments.append(audio_tensor)
+                
+                if reference_segments:
+                    profile.reference_segments = reference_segments
+                    logger.info(f"Loaded {len(reference_segments)} reference segments for {voice_name}")
+        except Exception as e:
+            logger.error(f"Error loading reference segments for {voice_name}: {e}")
+    
+    # Create context segments from references
+    context = []
+    if profile.reference_segments:
+        for i, ref_audio in enumerate(profile.reference_segments):
+            # Use corresponding prompt if available, otherwise use a generic one
+            text = VOICE_PROMPTS[voice_name][i] if i < len(VOICE_PROMPTS[voice_name]) else f"Voice reference for {voice_name}"
+            
+            context.append(
+                Segment(
+                    speaker=profile.speaker_id,
+                    text=text,
+                    audio=ref_audio.to(device)
+                )
+            )
+    
+    logger.info(f"Returning {len(context)} context segments for {voice_name}")
+    return context
+
+def save_voice_profiles():
+    """Save voice profiles to persistent storage."""
+    os.makedirs(VOICE_PROFILES_DIR, exist_ok=True)
+    profile_path = os.path.join(VOICE_PROFILES_DIR, "voice_profiles.pt")
+    
+    # Create a serializable version of the profiles
+    serializable_profiles = {}
+    for name, profile in VOICE_PROFILES.items():
+        serializable_profiles[name] = {
+            'reference_segments': [seg.cpu() for seg in profile.reference_segments]
+        }
+    
+    # Save to persistent storage
+    torch.save(serializable_profiles, profile_path)
+    logger.info(f"Saved voice profiles to {profile_path}")
+    
+def process_generated_audio(
+    audio: torch.Tensor, 
+    voice_name: str, 
+    sample_rate: int,
+    text: str
+) -> torch.Tensor:
+    """Process generated audio for consistency and quality.
+    
+    Args:
+        audio: Audio tensor
+        voice_name: Name of voice used
+        sample_rate: Audio sample rate  
+        text: Text that was spoken
+        
+    Returns:
+        Processed audio tensor
+    """
+    # Validate the audio
+    is_valid, audio, message = validate_generated_audio(audio, voice_name, sample_rate)
+    if not is_valid:
+        logger.warning(f"Generated audio validation issue: {message}")
+    
+    # Get voice profile for enhancement
+    profile = VOICE_PROFILES.get(voice_name, VOICE_PROFILES["alloy"])
+    
+    # Enhance the audio based on voice profile
+    enhanced = enhance_audio(audio, sample_rate, profile)
+    
+    # Log the enhancement
+    original_duration = audio.shape[0] / sample_rate
+    enhanced_duration = enhanced.shape[0] / sample_rate
+    logger.info(
+        f"Processed audio for '{voice_name}': "
+        f"Duration: {original_duration:.2f}s->{enhanced_duration:.2f}s, "
+        f"RMS: {audio.pow(2).mean().sqrt().item():.3f}->{enhanced.pow(2).mean().sqrt().item():.3f}"
+    )
+    
+    return enhanced

app/voice_memory.py

@@ -0,0 +1,474 @@
+"""Advanced voice memory system for consistent voice generation."""
+import os
+import torch
+import torchaudio
+import numpy as np
+import random
+import logging
+from typing import Dict, List, Optional
+from dataclasses import dataclass
+from app.models import Segment
+
+# Setup logging
+logger = logging.getLogger(__name__)
+
+# Path to store voice memories - use persistent location
+VOICE_MEMORIES_DIR = "/app/voice_memories"
+os.makedirs(VOICE_MEMORIES_DIR, exist_ok=True)
+
+@dataclass
+class VoiceMemory:
+    """Store voice characteristics for consistent generation."""
+    name: str  # Voice name (alloy, echo, etc.)
+    speaker_id: int  # Speaker ID (0-5)
+    # Store multiple audio segments for context
+    audio_segments: List[torch.Tensor]  
+    # Store text prompts that produced good results
+    text_segments: List[str]
+    # Base characteristics for this voice
+    pitch_base: float  # Base pitch characteristic (Hz)
+    timbre: str  # Voice quality descriptor
+    
+    def get_context_segments(self, device: torch.device, max_segments: int = 2) -> List[Segment]:
+        """Get context segments for this voice."""
+        if not self.audio_segments:
+            return []
+            
+        # Select a limited number of segments to avoid context overflow
+        num_segments = min(len(self.audio_segments), max_segments)
+        indices = list(range(len(self.audio_segments)))
+        random.shuffle(indices)
+        selected_indices = indices[:num_segments]
+        
+        segments = []
+        for i in selected_indices:
+            segments.append(
+                Segment(
+                    speaker=self.speaker_id,
+                    text=self.text_segments[i] if i < len(self.text_segments) else f"Voice sample {i}",
+                    audio=self.audio_segments[i].to(device)
+                )
+            )
+            
+        return segments
+    
+    def update_with_new_audio(self, audio: torch.Tensor, text: str, max_stored: int = 5):
+        """Update voice memory with newly generated audio."""
+        # Add new audio and text
+        self.audio_segments.append(audio.detach().cpu())
+        self.text_segments.append(text)
+        
+        # Keep only the most recent segments
+        if len(self.audio_segments) > max_stored:
+            self.audio_segments = self.audio_segments[-max_stored:]
+            self.text_segments = self.text_segments[-max_stored:]
+            
+    def save(self):
+        """Save voice memory to persistent storage."""
+        data = {
+            "name": self.name,
+            "speaker_id": self.speaker_id,
+            "audio_segments": self.audio_segments,
+            "text_segments": self.text_segments,
+            "pitch_base": self.pitch_base,
+            "timbre": self.timbre
+        }
+        
+        # Save to the persistent directory
+        save_path = os.path.join(VOICE_MEMORIES_DIR, f"{self.name}.pt")
+        try:
+            torch.save(data, save_path)
+            logger.info(f"Saved voice memory for {self.name} to {save_path}")
+        except Exception as e:
+            logger.error(f"Error saving voice memory for {self.name}: {e}")
+        
+    @classmethod
+    def load(cls, name: str) -> Optional['VoiceMemory']:
+        """Load voice memory from persistent storage."""
+        path = os.path.join(VOICE_MEMORIES_DIR, f"{name}.pt")
+        if not os.path.exists(path):
+            logger.info(f"No saved voice memory found for {name} at {path}")
+            return None
+            
+        try:
+            data = torch.load(path)
+            return cls(
+                name=data["name"],
+                speaker_id=data["speaker_id"],
+                audio_segments=data["audio_segments"],
+                text_segments=data["text_segments"],
+                pitch_base=data["pitch_base"],
+                timbre=data["timbre"]
+            )
+        except Exception as e:
+            logger.error(f"Error loading voice memory for {name}: {e}")
+            return None
+
+# Dictionary of voice memories
+VOICE_MEMORIES: Dict[str, VoiceMemory] = {}
+
+# Voice characteristics
+VOICE_CHARACTERISTICS = {
+    "alloy": {"pitch": 220.0, "timbre": "balanced", "description": "A balanced, natural voice with medium pitch"},
+    "echo": {"pitch": 330.0, "timbre": "resonant", "description": "A resonant voice with a reverberant quality"},
+    "fable": {"pitch": 523.0, "timbre": "bright", "description": "A bright, higher-pitched voice with clear articulation"},
+    "onyx": {"pitch": 165.0, "timbre": "deep", "description": "A deep, authoritative voice with lower pitch"},
+    "nova": {"pitch": 392.0, "timbre": "warm", "description": "A warm, smooth voice with pleasant midrange tone"},
+    "shimmer": {"pitch": 587.0, "timbre": "light", "description": "A light, airy voice with higher frequencies"}
+}
+
+# Voice intro texts - carefully crafted to capture voice characteristics
+VOICE_INTROS = {
+    "alloy": [
+        "Hello, I'm Alloy. My voice is designed to be clear and balanced.",
+        "This is the Alloy voice. I aim to sound natural and easy to understand.",
+        "Welcome, I'm the voice known as Alloy. I have a balanced, medium-range tone."
+    ],
+    "echo": [
+        "Hello, I'm Echo. My voice has a rich, resonant quality.",
+        "This is the Echo voice. Notice my distinctive resonance and depth.",
+        "Welcome, I'm the voice known as Echo. My tone is designed to resonate clearly."
+    ],
+    "fable": [
+        "Hello, I'm Fable. My voice is bright and articulate.",
+        "This is the Fable voice. I have a higher pitch with clear pronunciation.",
+        "Welcome, I'm the voice known as Fable. I speak with a bright, energetic tone."
+    ],
+    "onyx": [
+        "Hello, I'm Onyx. My voice is deep and authoritative.",
+        "This is the Onyx voice. I speak with a lower pitch and commanding presence.",
+        "Welcome, I'm the voice known as Onyx. My tone is deep and resonant."
+    ],
+    "nova": [
+        "Hello, I'm Nova. My voice is warm and harmonious.",
+        "This is the Nova voice. I have a smooth, pleasant mid-range quality.",
+        "Welcome, I'm the voice known as Nova. I speak with a warm, friendly tone."
+    ],
+    "shimmer": [
+        "Hello, I'm Shimmer. My voice is light and expressive.",
+        "This is the Shimmer voice. I have a higher-pitched, airy quality.",
+        "Welcome, I'm the voice known as Shimmer. My tone is bright and crisp."
+    ]
+}
+
+def initialize_voices(sample_rate: int = 24000):
+    """Initialize voice memories with consistent base samples."""
+    global VOICE_MEMORIES
+    
+    # Check if persistent directory exists, create if needed
+    os.makedirs(VOICE_MEMORIES_DIR, exist_ok=True)
+    logger.info(f"Using voice memories directory: {VOICE_MEMORIES_DIR}")
+    
+    # First try to load existing memories from persistent storage
+    for voice_name in ["alloy", "echo", "fable", "onyx", "nova", "shimmer"]:
+        memory = VoiceMemory.load(voice_name)
+        if memory:
+            VOICE_MEMORIES[voice_name] = memory
+            logger.info(f"Loaded existing voice memory for {voice_name} with {len(memory.audio_segments)} segments")
+            continue
+            
+        # If no memory exists, create a new one
+        speaker_id = ["alloy", "echo", "fable", "onyx", "nova", "shimmer"].index(voice_name)
+        characteristics = VOICE_CHARACTERISTICS[voice_name]
+        
+        # Create deterministic seed audio
+        np.random.seed(speaker_id + 42)
+        duration = 1.0  # seconds
+        t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
+        
+        # Create characteristic waveform
+        pitch = characteristics["pitch"]
+        if voice_name == "alloy":
+            audio = 0.5 * np.sin(2 * np.pi * pitch * t) + 0.3 * np.sin(2 * np.pi * pitch * 2 * t)
+        elif voice_name == "echo":
+            audio = np.sin(2 * np.pi * pitch * t) * np.exp(-t * 3)
+        elif voice_name == "fable":
+            audio = 0.7 * np.sin(2 * np.pi * pitch * t)
+        elif voice_name == "onyx":
+            audio = 0.8 * np.sin(2 * np.pi * pitch * t) + 0.1 * np.sin(2 * np.pi * pitch * 0.5 * t)
+        elif voice_name == "nova":
+            audio = 0.4 * np.sin(2 * np.pi * pitch * t) + 0.4 * np.sin(2 * np.pi * pitch * 0.5 * t)
+        else:  # shimmer
+            audio = 0.3 * np.sin(2 * np.pi * pitch * t) + 0.2 * np.sin(2 * np.pi * pitch * 1.5 * t) + 0.1 * np.sin(2 * np.pi * pitch * 2 * t)
+            
+        # Normalize
+        audio = audio / np.max(np.abs(audio))
+        
+        # Convert to tensor
+        audio_tensor = torch.tensor(audio, dtype=torch.float32)
+        
+        # Create voice memory
+        memory = VoiceMemory(
+            name=voice_name,
+            speaker_id=speaker_id,
+            audio_segments=[audio_tensor],
+            text_segments=[f"This is the voice of {voice_name}"],
+            pitch_base=characteristics["pitch"],
+            timbre=characteristics["timbre"]
+        )
+        
+        # Save the voice memory to persistent storage
+        memory.save()
+        
+        # Store in dictionary
+        VOICE_MEMORIES[voice_name] = memory
+        
+        # Save as wav for reference
+        save_path = os.path.join(VOICE_MEMORIES_DIR, f"{voice_name}_seed.wav")
+        torchaudio.save(save_path, audio_tensor.unsqueeze(0), sample_rate)
+        
+        logger.info(f"Initialized new voice memory for {voice_name}")
+
+def get_voice_context(voice_name: str, device: torch.device, max_segments: int = 2) -> List[Segment]:
+    """Get context segments for a given voice."""
+    if not VOICE_MEMORIES:
+        initialize_voices()
+        
+    if voice_name in VOICE_MEMORIES:
+        return VOICE_MEMORIES[voice_name].get_context_segments(device, max_segments=max_segments)
+        
+    # Default to alloy if voice not found
+    logger.warning(f"Voice {voice_name} not found, defaulting to alloy")
+    return VOICE_MEMORIES["alloy"].get_context_segments(device, max_segments=max_segments)
+
+def update_voice_memory(voice_name: str, audio: torch.Tensor, text: str):
+    """Update voice memory with newly generated audio and save to persistent storage."""
+    if not VOICE_MEMORIES:
+        initialize_voices()
+        
+    if voice_name in VOICE_MEMORIES:
+        VOICE_MEMORIES[voice_name].update_with_new_audio(audio, text)
+        VOICE_MEMORIES[voice_name].save()
+        logger.info(f"Updated voice memory for {voice_name}, now has {len(VOICE_MEMORIES[voice_name].audio_segments)} segments")
+
+def generate_voice_samples(app_state):
+    """Generate high-quality voice samples for each voice.
+    
+    Args:
+        app_state: The FastAPI app state containing the generator
+    """
+    generator = app_state.generator
+    if not generator:
+        logger.error("Cannot generate voice samples: generator not available")
+        return
+    
+    logger.info("Beginning voice sample generation...")
+    
+    # Ensure persistent directory exists
+    os.makedirs(VOICE_MEMORIES_DIR, exist_ok=True)
+    
+    for voice_name in ["alloy", "echo", "fable", "onyx", "nova", "shimmer"]:
+        speaker_id = ["alloy", "echo", "fable", "onyx", "nova", "shimmer"].index(voice_name)
+        
+        # Get multiple sample texts for this voice
+        sample_texts = VOICE_INTROS[voice_name]
+        
+        # Generate a collection of samples for this voice
+        logger.info(f"Generating samples for voice: {voice_name}")
+        audio_segments = []
+        text_segments = []
+        
+        for i, sample_text in enumerate(sample_texts):
+            try:
+                # Check if we already have a sample
+                sample_path = os.path.join(VOICE_MEMORIES_DIR, f"{voice_name}_sample_{i}.wav")
+                if os.path.exists(sample_path):
+                    logger.info(f"Found existing sample {i+1} for {voice_name}, loading from {sample_path}")
+                    audio_tensor, sr = torchaudio.load(sample_path)
+                    if sr != generator.sample_rate:
+                        audio_tensor = torchaudio.functional.resample(
+                            audio_tensor.squeeze(0), orig_freq=sr, new_freq=generator.sample_rate
+                        )
+                    else:
+                        audio_tensor = audio_tensor.squeeze(0)
+                    audio_segments.append(audio_tensor)
+                    text_segments.append(sample_text)
+                    continue
+                    
+                # Generate without context first for seed samples
+                logger.info(f"Generating sample {i+1}/{len(sample_texts)} for {voice_name}: '{sample_text}'")
+                
+                # Use a lower temperature for more stable output
+                audio = generator.generate(
+                    text=sample_text,
+                    speaker=speaker_id,
+                    context=[],  # No context for initial samples
+                    max_audio_length_ms=10000,
+                    temperature=0.7,  # Lower temperature for more stable output
+                    topk=30,
+                )
+                
+                # Save this segment
+                audio_segments.append(audio.detach().cpu())
+                text_segments.append(sample_text)
+                
+                # Save as WAV for reference to persistent storage
+                torchaudio.save(sample_path, audio.unsqueeze(0).cpu(), generator.sample_rate)
+                
+                logger.info(f"Generated sample {i+1} for {voice_name}, length: {audio.shape[0]/generator.sample_rate:.2f}s")
+            
+            except Exception as e:
+                logger.error(f"Error generating sample {i+1} for {voice_name}: {e}")
+        
+        # Use the generated samples to update the voice memory
+        if voice_name in VOICE_MEMORIES and audio_segments:
+            # Replace existing samples with these high quality ones
+            VOICE_MEMORIES[voice_name].audio_segments = audio_segments
+            VOICE_MEMORIES[voice_name].text_segments = text_segments
+            VOICE_MEMORIES[voice_name].save()
+            
+            logger.info(f"Updated voice memory for {voice_name} with {len(audio_segments)} high-quality samples")
+        
+        # Now generate a second pass with context from these samples
+        if len(audio_segments) >= 2:
+            try:
+                # Check if we already have a character sample
+                character_path = os.path.join(VOICE_MEMORIES_DIR, f"{voice_name}_character.wav")
+                if os.path.exists(character_path):
+                    logger.info(f"Found existing character sample for {voice_name}, loading from {character_path}")
+                    audio_tensor, sr = torchaudio.load(character_path)
+                    if sr != generator.sample_rate:
+                        audio_tensor = torchaudio.functional.resample(
+                            audio_tensor.squeeze(0), orig_freq=sr, new_freq=generator.sample_rate
+                        )
+                    else:
+                        audio_tensor = audio_tensor.squeeze(0)
+                
+                    character_sample_text = f"I'm the voice assistant known as {voice_name}. I'm designed to have a distinctive voice that you can easily recognize."
+                    VOICE_MEMORIES[voice_name].audio_segments.append(audio_tensor)
+                    VOICE_MEMORIES[voice_name].text_segments.append(character_sample_text)
+                    VOICE_MEMORIES[voice_name].save()
+                    continue
+                
+                # Get intro and conclusion prompts that build voice consistency
+                context = [
+                    Segment(
+                        speaker=speaker_id,
+                        text=text_segments[0],
+                        audio=audio_segments[0].to(generator.device)
+                    )
+                ]
+                
+                # Create a longer sample with the voice characteristics now established
+                character_sample_text = f"I'm the voice assistant known as {voice_name}. I'm designed to have a distinctive voice that you can easily recognize. My speech patterns and tone should remain consistent throughout our conversation."
+                
+                logger.info(f"Generating character sample for {voice_name} with context")
+                character_audio = generator.generate(
+                    text=character_sample_text,
+                    speaker=speaker_id,
+                    context=context,
+                    max_audio_length_ms=15000,
+                    temperature=0.7, 
+                    topk=30,
+                )
+                
+                # Save this comprehensive character sample to persistent storage
+                torchaudio.save(character_path, character_audio.unsqueeze(0).cpu(), generator.sample_rate)
+                
+                # Add this to the memory as well
+                VOICE_MEMORIES[voice_name].audio_segments.append(character_audio.detach().cpu()) 
+                VOICE_MEMORIES[voice_name].text_segments.append(character_sample_text)
+                VOICE_MEMORIES[voice_name].save()
+                
+                logger.info(f"Generated character sample for {voice_name}, length: {character_audio.shape[0]/generator.sample_rate:.2f}s")
+                
+            except Exception as e:
+                logger.error(f"Error generating character sample for {voice_name}: {e}")
+
+def create_custom_voice(
+    app_state,
+    name: str, 
+    initial_text: str, 
+    speaker_id: int = 0,
+    pitch: Optional[float] = None,
+    timbre: str = "custom"
+) -> Dict:
+    """Create a new custom voice.
+    
+    Args:
+        app_state: The FastAPI app state containing the generator
+        name: Name for the new voice
+        initial_text: Text for the initial voice sample
+        speaker_id: Base speaker ID (0-5)
+        pitch: Base pitch in Hz (optional)
+        timbre: Voice quality descriptor
+        
+    Returns:
+        Dict with creation status and voice info
+    """
+    generator = app_state.generator
+    if not generator:
+        return {"status": "error", "message": "Generator not available"}
+    
+    # Check if voice already exists
+    if not VOICE_MEMORIES:
+        initialize_voices()
+        
+    if name in VOICE_MEMORIES:
+        return {"status": "error", "message": f"Voice '{name}' already exists"}
+    
+    # Generate a voice sample
+    try:
+        logger.info(f"Creating custom voice '{name}' with text: '{initial_text}'")
+        
+        audio = generator.generate(
+            text=initial_text,
+            speaker=speaker_id,
+            context=[],
+            max_audio_length_ms=10000,
+            temperature=0.7,
+        )
+        
+        # Determine base pitch if not provided
+        if pitch is None:
+            if speaker_id == 0:  # alloy
+                pitch = 220.0
+            elif speaker_id == 1:  # echo
+                pitch = 330.0
+            elif speaker_id == 2:  # fable
+                pitch = 523.0
+            elif speaker_id == 3:  # onyx
+                pitch = 165.0
+            elif speaker_id == 4:  # nova
+                pitch = 392.0
+            else:  # shimmer
+                pitch = 587.0
+        
+        # Create a new voice memory
+        memory = VoiceMemory(
+            name=name,
+            speaker_id=speaker_id,
+            audio_segments=[audio.detach().cpu()],
+            text_segments=[initial_text],
+            pitch_base=pitch,
+            timbre=timbre
+        )
+        
+        # Save the voice memory to persistent storage
+        memory.save()
+        VOICE_MEMORIES[name] = memory
+        
+        # Save sample as WAV for reference to persistent storage
+        sample_path = os.path.join(VOICE_MEMORIES_DIR, f"{name}_sample.wav")
+        torchaudio.save(sample_path, audio.unsqueeze(0).cpu(), generator.sample_rate)
+        
+        logger.info(f"Created custom voice '{name}' successfully")
+        
+        return {
+            "status": "success",
+            "message": f"Voice '{name}' created successfully",
+            "voice": {
+                "name": name,
+                "speaker_id": speaker_id,
+                "pitch": pitch,
+                "timbre": timbre,
+                "sample_length_seconds": audio.shape[0] / generator.sample_rate
+            }
+        }
+    
+    except Exception as e:
+        logger.error(f"Error creating custom voice '{name}': {e}")
+        return {
+            "status": "error",
+            "message": f"Error creating voice: {str(e)}"
+        }

app/watermarking.py

@@ -0,0 +1,22 @@
+"""Stub for watermarking module.
+
+The original CSM code has a watermarking module that adds
+an imperceptible watermark to generated audio.
+"""
+
+# Watermark key used by CSM
+CSM_1B_GH_WATERMARK = "CSM1B@GitHub"
+
+def load_watermarker(device="cpu"):
+    """Stub for watermarker loading.
+    
+    In a real implementation, this would load the actual watermarker.
+    """
+    return None
+
+def watermark(watermarker, audio, sample_rate, key):
+    """Stub for watermarking function.
+    
+    In a real implementation, this would add the watermark.
+    """
+    return audio, sample_rate

docker-compose.yml

@@ -0,0 +1,23 @@
+services:
+  csm-api:
+    build:
+      context: .
+      dockerfile: Dockerfile
+      args:
+        - HF_TOKEN=${HF_TOKEN}
+    ports:
+      - "8000:8000"
+    volumes:
+      - ./models:/app/models
+      - ./cloned_voices:/app/cloned_voices
+      - ./voice_references:/app/voice_references  
+      - ./voice_profiles:/app/voice_profiles    
+    environment:
+      - HF_TOKEN=${HF_TOKEN}  
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]

requirements.txt

@@ -0,0 +1,24 @@
+fastapi>=0.100.0
+uvicorn>=0.24.0
+pydantic>=2.4.0
+python-multipart>=0.0.6
+huggingface_hub>=0.20.0
+torch>=2.0.0
+torchaudio>=2.0.0
+transformers>=4.35.0
+tokenizers>=0.14.0
+sentencepiece>=0.1.99
+triton>=2.1.0; platform_system != "Windows"
+triton-windows>=2.1.0; platform_system == "Windows"
+torchao>=0.1.0
+torchtune>=0.5.0
+numpy>=1.25.0
+ffmpeg-python>=0.2.0
+moshi>=0.1.0
+soundfile>=0.12.1
+scipy>=1.10.0
+librosa>=0.10.0
+yt-dlp>=2023.3.4
+openai-whisper>=20230314
+ffmpeg-python>=0.2.0
+accelerate>=0.20.0

static/voice-cloning.html

@@ -0,0 +1,1385 @@
+<!DOCTYPE html>
+<html lang="en" data-theme="light">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0, maximum-scale=1.0, user-scalable=no">
+    <title>Voice Cloning</title>
+    <style>
+        :root {
+            /* Base colors */
+            --bg-color: #ffffff;
+            --text-color: #333333;
+            --card-bg: #ffffff;
+            --card-border: #e2e8f0;
+            --card-shadow: rgba(0,0,0,0.1);
+            
+            /* Primary colors */
+            --primary-color: #4f46e5;
+            --primary-hover: #4338ca;
+            --primary-light: rgba(79, 70, 229, 0.1);
+            
+            /* Secondary colors */
+            --secondary-color: #6b7280;
+            --secondary-hover: #4b5563;
+            
+            /* Accent colors */
+            --success-color: #10b981;
+            --success-hover: #059669;
+            --danger-color: #ef4444;
+            --danger-hover: #dc2626;
+            --warning-color: #f59e0b;
+            --warning-hover: #d97706;
+            
+            /* Input elements */
+            --input-bg: #ffffff;
+            --input-border: #d1d5db;
+            --input-focus-border: #4f46e5;
+            --input-focus-shadow: rgba(79, 70, 229, 0.2);
+            
+            /* Voice cards */
+            --voice-card-bg: #f9fafb;
+            --voice-card-border: #e5e7eb;
+            --voice-card-shadow: rgba(0,0,0,0.05);
+            
+            /* Tabs */
+            --tab-border: #e5e7eb;
+            --tab-text: #4b5563;
+            --tab-active: #4f46e5;
+            --tab-active-bg: rgba(79, 70, 229, 0.1);
+            
+            /* Toggle */
+            --toggle-bg: #e5e7eb;
+            --toggle-active: #4f46e5;
+            --toggle-circle: #ffffff;
+            
+            /* Status indicators */
+            --status-success-bg: #dcfce7;
+            --status-success-text: #166534;
+            --status-error-bg: #fee2e2;
+            --status-error-text: #b91c1c;
+            --status-warning-bg: #fff7ed;
+            --status-warning-text: #c2410c;
+            --status-info-bg: #eff6ff;
+            --status-info-text: #1e40af;
+        }
+        
+        [data-theme="dark"] {
+            /* Base colors */
+            --bg-color: #111827;
+            --text-color: #f3f4f6;
+            --card-bg: #1f2937;
+            --card-border: #374151;
+            --card-shadow: rgba(0,0,0,0.3);
+            
+            /* Primary colors */
+            --primary-color: #6366f1;
+            --primary-hover: #4f46e5;
+            --primary-light: rgba(99, 102, 241, 0.2);
+            
+            /* Secondary colors */
+            --secondary-color: #9ca3af;
+            --secondary-hover: #6b7280;
+            
+            /* Accent colors - slightly brighter for dark theme */
+            --success-color: #34d399;
+            --success-hover: #10b981;
+            --danger-color: #f87171;
+            --danger-hover: #ef4444;
+            --warning-color: #fbbf24;
+            --warning-hover: #f59e0b;
+            
+            /* Input elements */
+            --input-bg: #374151;
+            --input-border: #4b5563;
+            --input-focus-border: #6366f1;
+            --input-focus-shadow: rgba(99, 102, 241, 0.3);
+            
+            /* Voice cards */
+            --voice-card-bg: #1f2937;
+            --voice-card-border: #374151;
+            --voice-card-shadow: rgba(0,0,0,0.2);
+            
+            /* Tabs */
+            --tab-border: #374151;
+            --tab-text: #9ca3af;
+            --tab-active: #6366f1;
+            --tab-active-bg: rgba(99, 102, 241, 0.2);
+            
+            /* Toggle */
+            --toggle-bg: #4b5563;
+            --toggle-active: #6366f1;
+            --toggle-circle: #e5e7eb;
+            
+            /* Status indicators */
+            --status-success-bg: rgba(16, 185, 129, 0.2);
+            --status-success-text: #34d399;
+            --status-error-bg: rgba(239, 68, 68, 0.2);
+            --status-error-text: #f87171;
+            --status-warning-bg: rgba(245, 158, 11, 0.2);
+            --status-warning-text: #fbbf24;
+            --status-info-bg: rgba(59, 130, 246, 0.2);
+            --status-info-text: #60a5fa;
+        }
+        
+        * {
+            box-sizing: border-box;
+        }
+        
+        html, body {
+            margin: 0;
+            padding: 0;
+            overflow-x: hidden;
+            width: 100%;
+        }
+        
+        body {
+            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, 'Open Sans', 'Helvetica Neue', sans-serif;
+            max-width: 100%;
+            padding: 16px;
+            line-height: 1.6;
+            background-color: var(--bg-color);
+            color: var(--text-color);
+            transition: background-color 0.3s, color 0.3s;
+        }
+        
+        .container {
+            width: 100%;
+            max-width: 800px;
+            margin: 0 auto;
+            padding: 0 8px;
+        }
+        
+        h1, h2, h3 {
+            color: var(--text-color);
+            margin-top: 0;
+        }
+        
+        h1 {
+            font-size: 1.8rem;
+            margin-bottom: 1rem;
+            text-align: center;
+        }
+        
+        h2 {
+            font-size: 1.4rem;
+            margin-bottom: 1rem;
+        }
+        
+        .card {
+            border: 1px solid var(--card-border);
+            border-radius: 12px;
+            padding: 20px;
+            margin-bottom: 20px;
+            box-shadow: 0 4px 6px var(--card-shadow);
+            background-color: var(--card-bg);
+            transition: box-shadow 0.3s ease;
+        }
+        
+        .card:hover {
+            box-shadow: 0 6px 12px var(--card-shadow);
+        }
+        
+        .form-group {
+            margin-bottom: 20px;
+        }
+        
+        label {
+            display: block;
+            margin-bottom: 6px;
+            font-weight: 500;
+            font-size: 0.95rem;
+        }
+        
+        input, textarea, select {
+            width: 100%;
+            padding: 10px 12px;
+            border: 1px solid var(--input-border);
+            border-radius: 8px;
+            background-color: var(--input-bg);
+            color: var(--text-color);
+            font-size: 1rem;
+            transition: border-color 0.3s, box-shadow 0.3s;
+        }
+        
+        input:focus, textarea:focus, select:focus {
+            outline: none;
+            border-color: var(--input-focus-border);
+            box-shadow: 0 0 0 3px var(--input-focus-shadow);
+        }
+        
+        /* File input styling */
+        input[type="file"] {
+            padding: 8px;
+            background-color: var(--input-bg);
+            border: 1px dashed var(--input-border);
+            border-radius: 8px;
+            cursor: pointer;
+        }
+        
+        input[type="file"]:hover {
+            border-color: var(--primary-color);
+        }
+        
+        button {
+            background-color: var(--primary-color);
+            color: white;
+            border: none;
+            border-radius: 8px;
+            padding: 12px 20px;
+            cursor: pointer;
+            font-weight: 600;
+            font-size: 1rem;
+            transition: background-color 0.3s, transform 0.1s;
+            width: 100%;
+        }
+        
+        button:hover {
+            background-color: var(--primary-hover);
+        }
+        
+        button:active {
+            transform: translateY(1px);
+        }
+        
+        button:disabled {
+            opacity: 0.7;
+            cursor: not-allowed;
+        }
+        
+        .btn-row {
+            display: flex;
+            gap: 10px;
+        }
+        
+        .btn-row button {
+            flex: 1;
+        }
+        
+        .voice-list {
+            display: grid;
+            grid-template-columns: repeat(auto-fill, minmax(280px, 1fr));
+            gap: 16px;
+        }
+        
+        .voice-card {
+            border: 1px solid var(--voice-card-border);
+            border-radius: 10px;
+            padding: 16px;
+            background-color: var(--voice-card-bg);
+            box-shadow: 0 2px 6px var(--voice-card-shadow);
+            transition: transform 0.2s ease, box-shadow 0.2s ease;
+        }
+        
+        .voice-card:hover {
+            transform: translateY(-3px);
+            box-shadow: 0 6px 12px var(--voice-card-shadow);
+        }
+        
+        .controls {
+            display: flex;
+            gap: 8px;
+            margin-top: 12px;
+        }
+        
+        .voice-name {
+            font-weight: 600;
+            font-size: 18px;
+            margin: 0 0 8px 0;
+            color: var(--primary-color);
+        }
+        
+        .btn-danger {
+            background-color: var(--danger-color);
+        }
+        
+        .btn-danger:hover {
+            background-color: var(--danger-hover);
+        }
+        
+        .btn-secondary {
+            background-color: var(--secondary-color);
+        }
+        
+        .btn-secondary:hover {
+            background-color: var(--secondary-hover);
+        }
+        
+        #audio-preview {
+            margin-top: 20px;
+            width: 100%;
+            border-radius: 8px;
+            background-color: var(--card-bg);
+        }
+        
+        /* Status indicators */
+        .status-indicator {
+            padding: 12px 16px;
+            margin: 16px 0;
+            border-radius: 8px;
+            font-weight: 500;
+            display: flex;
+            align-items: center;
+            opacity: 0;
+            transition: opacity 0.3s ease;
+            max-height: 0;
+            overflow: hidden;
+        }
+        
+        .status-indicator.show {
+            opacity: 1;
+            max-height: 100px;
+        }
+        
+        .status-indicator.success {
+            background-color: var(--status-success-bg);
+            color: var(--status-success-text);
+        }
+        
+        .status-indicator.error {
+            background-color: var(--status-error-bg);
+            color: var(--status-error-text);
+        }
+        
+        .status-indicator.warning {
+            background-color: var(--status-warning-bg);
+            color: var(--status-warning-text);
+        }
+        
+        .status-indicator.info {
+            background-color: var(--status-info-bg);
+            color: var(--status-info-text);
+        }
+        
+        .status-indicator svg {
+            margin-right: 8px;
+            flex-shrink: 0;
+        }
+        
+        /* Tabs */
+        .tabs {
+            display: flex;
+            margin-bottom: 20px;
+            border-bottom: 1px solid var(--tab-border);
+            overflow-x: auto;
+            -webkit-overflow-scrolling: touch;
+            scrollbar-width: none; /* Hide scrollbar for Firefox */
+        }
+        
+        .tabs::-webkit-scrollbar {
+            display: none; /* Hide scrollbar for Chrome/Safari */
+        }
+        
+        .tabs button {
+            background-color: transparent;
+            color: var(--tab-text);
+            border: none;
+            padding: 12px 16px;
+            margin-right: 8px;
+            cursor: pointer;
+            position: relative;
+            font-weight: 600;
+            border-radius: 8px 8px 0 0;
+            white-space: nowrap;
+            width: auto;
+            flex-shrink: 0;
+        }
+        
+        .tabs button.active {
+            color: var(--tab-active);
+            background-color: var(--tab-active-bg);
+        }
+        
+        .tabs button.active::after {
+            content: '';
+            position: absolute;
+            bottom: -1px;
+            left: 0;
+            right: 0;
+            height: 2px;
+            background-color: var(--tab-active);
+        }
+        
+        .tab-content {
+            display: none;
+        }
+        
+        .tab-content.active {
+            display: block;
+        }
+        
+        /* Theme toggle switch */
+        .theme-switch-wrapper {
+            display: flex;
+            align-items: center;
+            justify-content: flex-end;
+            margin-bottom: 16px;
+        }
+        
+        .theme-switch {
+            display: inline-block;
+            height: 28px;
+            position: relative;
+            width: 54px;
+        }
+        
+        .theme-switch input {
+            display: none;
+        }
+        
+        .slider {
+            background-color: var(--toggle-bg);
+            bottom: 0;
+            cursor: pointer;
+            left: 0;
+            position: absolute;
+            right: 0;
+            top: 0;
+            transition: .4s;
+            border-radius: 28px;
+        }
+        
+        .slider:before {
+            background-color: var(--toggle-circle);
+            bottom: 4px;
+            content: "";
+            height: 20px;
+            left: 4px;
+            position: absolute;
+            transition: .4s;
+            width: 20px;
+            border-radius: 50%;
+            box-shadow: 0 1px 3px rgba(0, 0, 0, 0.2);
+        }
+        
+        input:checked + .slider {
+            background-color: var(--toggle-active);
+        }
+        
+        input:checked + .slider:before {
+            transform: translateX(26px);
+        }
+        
+        .theme-icon {
+            width: 16px;
+            height: 16px;
+            display: inline-block;
+            margin: 0 8px;
+            font-size: 16px;
+            line-height: 1;
+        }
+        
+        /* Progress bar */
+        .progress-bar {
+            width: 100%;
+            height: 12px;
+            background-color: var(--input-bg);
+            border-radius: 6px;
+            margin-top: 12px;
+            overflow: hidden;
+            box-shadow: inset 0 1px 3px var(--card-shadow);
+        }
+        
+        .progress-fill {
+            height: 100%;
+            background: linear-gradient(to right, var(--primary-color), var(--primary-hover));
+            width: 0%;
+            transition: width 0.5s ease-in-out;
+            border-radius: 6px;
+            position: relative;
+        }
+        
+        .progress-fill::after {
+            content: '';
+            position: absolute;
+            top: 0;
+            left: 0;
+            right: 0;
+            bottom: 0;
+            background: linear-gradient(
+                -45deg,
+                rgba(255, 255, 255, 0.2) 25%,
+                transparent 25%,
+                transparent 50%,
+                rgba(255, 255, 255, 0.2) 50%,
+                rgba(255, 255, 255, 0.2) 75%,
+                transparent 75%
+            );
+            background-size: 16px 16px;
+            animation: progress-animation 1s linear infinite;
+            border-radius: 6px;
+        }
+        
+        @keyframes progress-animation {
+            0% {
+                background-position: 0 0;
+            }
+            100% {
+                background-position: 16px 0;
+            }
+        }
+        
+        /* Divider */
+        .divider {
+            margin: 24px 0;
+            border-top: 1px solid var(--card-border);
+            position: relative;
+        }
+        
+        .divider-text {
+            position: absolute;
+            top: -10px;
+            left: 50%;
+            transform: translateX(-50%);
+            background-color: var(--bg-color);
+            padding: 0 12px;
+            color: var(--secondary-color);
+            font-size: 0.9rem;
+        }
+        
+        /* Small text */
+        small {
+            color: var(--secondary-color);
+            display: block;
+            margin-top: 6px;
+            font-size: 0.85rem;
+        }
+        
+        /* Range slider styling */
+        input[type="range"] {
+            -webkit-appearance: none;
+            height: 8px;
+            border-radius: 4px;
+            background: var(--input-bg);
+            outline: none;
+            padding: 0;
+            margin: 10px 0;
+        }
+        
+        input[type="range"]::-webkit-slider-thumb {
+            -webkit-appearance: none;
+            appearance: none;
+            width: 20px;
+            height: 20px;
+            border-radius: 50%;
+            background: var(--primary-color);
+            cursor: pointer;
+            border: 2px solid white;
+            box-shadow: 0 1px 3px rgba(0, 0, 0, 0.2);
+        }
+        
+        input[type="range"]::-moz-range-thumb {
+            width: 20px;
+            height: 20px;
+            border-radius: 50%;
+            background: var(--primary-color);
+            cursor: pointer;
+            border: 2px solid white;
+            box-shadow: 0 1px 3px rgba(0, 0, 0, 0.2);
+        }
+        
+        input[type="range"]::-ms-thumb {
+            width: 20px;
+            height: 20px;
+            border-radius: 50%;
+            background: var(--primary-color);
+            cursor: pointer;
+            border: 2px solid white;
+            box-shadow: 0 1px 3px rgba(0, 0, 0, 0.2);
+        }
+        
+        #temperature-value {
+            display: inline-block;
+            width: 40px;
+            text-align: center;
+            font-weight: 600;
+            color: var(--primary-color);
+        }
+        
+        /* Toast notifications */
+        .toast-container {
+            position: fixed;
+            bottom: 20px;
+            right: 20px;
+            z-index: 1000;
+            max-width: 100%;
+            width: 300px;
+        }
+        
+        .toast {
+            padding: 12px 16px;
+            margin-bottom: 12px;
+            border-radius: 8px;
+            color: white;
+            font-weight: 500;
+            box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);
+            display: flex;
+            align-items: center;
+            justify-content: space-between;
+            animation: toast-in 0.3s ease forwards;
+            max-width: 100%;
+        }
+        
+        .toast.success {
+            background-color: var(--success-color);
+        }
+        
+        .toast.error {
+            background-color: var(--danger-color);
+        }
+        
+        .toast.warning {
+            background-color: var(--warning-color);
+        }
+        
+        .toast.info {
+            background-color: var(--primary-color);
+        }
+        
+        .toast-close {
+            background: none;
+            border: none;
+            color: white;
+            font-size: 18px;
+            cursor: pointer;
+            opacity: 0.8;
+            width: auto;
+            padding: 0 0 0 12px;
+        }
+        
+        .toast-close:hover {
+            opacity: 1;
+            background: none;
+        }
+        
+        @keyframes toast-in {
+            from {
+                transform: translateX(100%);
+                opacity: 0;
+            }
+            to {
+                transform: translateX(0);
+                opacity: 1;
+            }
+        }
+        
+        /* Loading spinner */
+        .spinner {
+            display: inline-block;
+            width: 20px;
+            height: 20px;
+            margin-right: 8px;
+            border: 3px solid rgba(255, 255, 255, 0.3);
+            border-radius: 50%;
+            border-top-color: white;
+            animation: spin 1s ease-in-out infinite;
+        }
+        
+        @keyframes spin {
+            to { transform: rotate(360deg); }
+        }
+        
+        /* Mobile responsiveness */
+        @media (max-width: 640px) {
+            body {
+                padding: 12px 8px;
+            }
+            
+            h1 {
+                font-size: 1.6rem;
+            }
+            
+            .card {
+                padding: 16px;
+            }
+            
+            .voice-list {
+                grid-template-columns: 1fr;
+            }
+            
+            .controls {
+                flex-direction: column;
+            }
+            
+            .controls button {
+                width: 100%;
+            }
+            
+            .tabs button {
+                padding: 8px 12px;
+                font-size: 0.9rem;
+            }
+        }
+        /* Checkbox styling */
+        .checkbox-group {
+            margin-bottom: 20px;
+        }
+        .checkbox-label {
+            display: flex;
+            align-items: center;
+            cursor: pointer;
+            font-weight: 500;
+            font-size: 0.95rem;
+        }
+        input[type="checkbox"] {
+            width: auto;
+            margin-right: 10px;
+            height: 18px;
+            width: 18px;
+            cursor: pointer;
+            accent-color: var(--primary-color);
+        }
+        .checkbox-text {
+            position: relative;
+            top: 1px;
+        }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <div class="theme-switch-wrapper">
+            <span class="theme-icon">☀️</span>
+            <label class="theme-switch" for="checkbox">
+                <input type="checkbox" id="checkbox" />
+                <div class="slider"></div>
+            </label>
+            <span class="theme-icon">🌙</span>
+        </div>
+        
+        <h1>Voice Cloning</h1>
+        
+        <div class="tabs">
+            <button id="tab-clone" class="active">Clone Voice</button>
+            <button id="tab-voices">My Voices</button>
+            <button id="tab-generate">Generate Speech</button>
+        </div>
+        
+        <!-- Status indicator -->
+        <div id="status-message" class="status-indicator">
+            <!-- Content will be dynamically inserted -->
+        </div>
+        
+        <div id="clone-tab" class="tab-content active">
+            <div class="card">
+                <h2>Clone a New Voice</h2>
+                <form id="clone-form">
+                    <div class="form-group">
+                        <label for="voice-name">Voice Name</label>
+                        <input type="text" id="voice-name" name="name" required placeholder="e.g. My Voice">
+                    </div>
+                    <div class="form-group">
+                        <label for="audio-file">Voice Sample (2-3 minute audio recording)</label>
+                        <input type="file" id="audio-file" name="audio_file" required accept="audio/*">
+                        <small>For best results, provide a clear recording with minimal background noise.</small>
+                    </div>
+                    <div class="form-group">
+                        <label for="transcript">Transcript (Optional)</label>
+                        <textarea id="transcript" name="transcript" rows="4" placeholder="Exact transcript of your audio sample..."></textarea>
+                        <small>Adding a transcript helps improve voice accuracy.</small>
+                    </div>
+                    <div class="form-group">
+                        <label for="description">Description (Optional)</label>
+                        <input type="text" id="description" name="description" placeholder="A description of this voice">
+                    </div>
+                    <button type="submit">Clone Voice</button>
+                </form>
+            </div>
+            
+            <div class="divider">
+                <span class="divider-text">OR</span>
+            </div>
+            
+            <div class="card">
+                <h2>Clone Voice from YouTube</h2>
+                <form id="youtube-clone-form">
+                    <div class="form-group">
+                        <label for="youtube-url">YouTube URL</label>
+                        <input type="url" id="youtube-url" name="youtube_url" required placeholder="https://www.youtube.com/watch?v=...">
+                    </div>
+                    <div class="form-group">
+                        <label for="youtube-voice-name">Voice Name</label>
+                        <input type="text" id="youtube-voice-name" name="voice_name" required placeholder="e.g. YouTube Voice">
+                    </div>
+                    <div class="form-group">
+                        <label for="start-time">Start Time (seconds)</label>
+                        <input type="number" id="start-time" name="start_time" min="0" value="0">
+                    </div>
+                    <div class="form-group">
+                        <label for="duration">Duration (seconds)</label>
+                        <input type="number" id="duration" name="duration" min="10" max="600" value="180">
+                        <small>Recommended: 2-3 minutes of clear speech</small>
+                    </div>
+                    <div class="form-group">
+                        <label for="youtube-description">Description (Optional)</label>
+                        <input type="text" id="youtube-description" name="description" placeholder="A description of this voice">
+                    </div>
+                    <button type="submit">Clone from YouTube</button>
+                </form>
+                <div id="youtube-progress" style="display: none; margin-top: 16px;">
+                    <p>Processing YouTube video... <span id="progress-status">Downloading</span></p>
+                    <div class="progress-bar">
+                        <div class="progress-fill"></div>
+                    </div>
+                </div>
+            </div>
+        </div>
+        
+        <div id="voices-tab" class="tab-content">
+            <h2>My Cloned Voices</h2>
+            <div id="voice-list" class="voice-list">
+                <!-- Voice cards will be added here -->
+            </div>
+        </div>
+        
+        <div id="generate-tab" class="tab-content">
+            <div class="card">
+                <h2>Generate Speech with Cloned Voice</h2>
+                <form id="generate-form">
+                    <div class="form-group">
+                        <label for="voice-select">Select Voice</label>
+                        <select id="voice-select" name="voice" required>
+                            <option value="">Select a voice</option>
+                            <!-- Voice options will be added here -->
+                        </select>
+                    </div>
+                    <div class="form-group">
+                        <label for="generate-text">Text to Speak</label>
+                        <textarea id="generate-text" name="text" rows="4" required placeholder="Enter text to synthesize with the selected voice..."></textarea>
+                    </div>
+                    <div class="form-group">
+                        <label>Temperature: <span id="temperature-value">0.7</span></label>
+                        <input type="range" id="temperature" name="temperature" min="0.5" max="1.0" step="0.05" value="0.7">
+                        <small>Lower values (0.5-0.7) produce more consistent speech, higher values (0.8-1.0) produce more varied speech.</small>
+                    </div>
+                    <div class="form-group checkbox-group">
+                        <label class="checkbox-label">
+                            <input type="checkbox" id="use-streaming" name="use_streaming">
+                            <span class="checkbox-text">Use streaming mode</span>
+                        </label>
+                        <small>Stream audio as it's generated for faster start and lower latency.</small>
+                    </div>
+                    <button type="submit">Generate Speech</button>
+                </form>
+                <audio id="audio-preview" controls style="display: none;"></audio>
+            </div>
+        </div>
+    </div>
+    
+    <!-- Toast notifications container -->
+    <div class="toast-container" id="toast-container"></div>
+    
+    <script>
+        // Theme toggle functionality
+        const toggleSwitch = document.querySelector('#checkbox');
+        const html = document.querySelector('html');
+        
+        // Check for saved theme preference or use system preference
+        function getThemePreference() {
+            const savedTheme = localStorage.getItem('theme');
+            if (savedTheme) {
+                return savedTheme;
+            }
+            // Check if system prefers dark mode
+            return window.matchMedia('(prefers-color-scheme: dark)').matches ? 'dark' : 'light';
+        }
+        
+        // Apply the theme
+        function setTheme(theme) {
+            html.setAttribute('data-theme', theme);
+            localStorage.setItem('theme', theme);
+            toggleSwitch.checked = theme === 'dark';
+        }
+        
+        // Initialize theme
+        setTheme(getThemePreference());
+        
+        // Listen for toggle changes
+        toggleSwitch.addEventListener('change', function(e) {
+            if (e.target.checked) {
+                setTheme('dark');
+            } else {
+                setTheme('light');
+            }
+        });
+        
+        // Toast notification system
+        function showToast(message, type = 'info', duration = 5000) {
+            const container = document.getElementById('toast-container');
+            const toast = document.createElement('div');
+            toast.className = `toast ${type}`;
+            toast.innerHTML = `
+                <span>${message}</span>
+                <button class="toast-close">&times;</button>
+            `;
+            container.appendChild(toast);
+            
+            // Auto remove after duration
+            const timeout = setTimeout(() => {
+                toast.style.opacity = '0';
+                setTimeout(() => {
+                    container.removeChild(toast);
+                }, 300);
+            }, duration);
+            
+            // Manual close
+            toast.querySelector('.toast-close').addEventListener('click', () => {
+                clearTimeout(timeout);
+                toast.style.opacity = '0';
+                setTimeout(() => {
+                    container.removeChild(toast);
+                }, 300);
+            });
+        }
+        
+        // Status indicator functions
+        function showStatus(message, type) {
+            const statusElem = document.getElementById('status-message');
+            statusElem.className = `status-indicator ${type} show`;
+            
+            let icon = '';
+            switch(type) {
+                case 'success':
+                    icon = '<svg width="20" height="20" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M9 12l2 2 4-4M21 12a9 9 0 11-18 0 9 9 0 0118 0z" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"/></svg>';
+                    break;
+                case 'error':
+                    icon = '<svg width="20" height="20" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M10 14l2-2m0 0l2-2m-2 2l-2-2m2 2l2 2m7-2a9 9 0 11-18 0 9 9 0 0118 0z" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"/></svg>';
+                    break;
+                case 'warning':
+                    icon = '<svg width="20" height="20" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M12 9v2m0 4h.01m-6.938 4h13.856c1.54 0 2.502-1.667 1.732-3L13.732 4c-.77-1.333-2.694-1.333-3.464 0L3.34 16c-.77 1.333.192 3 1.732 3z" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"/></svg>';
+                    break;
+                case 'info':
+                    icon = '<svg width="20" height="20" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M13 16h-1v-4h-1m1-4h.01M21 12a9 9 0 11-18 0 9 9 0 0118 0z" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"/></svg>';
+                    break;
+            }
+            
+            statusElem.innerHTML = icon + message;
+            
+            // Auto-hide after 5 seconds
+            setTimeout(() => {
+                statusElem.className = 'status-indicator';
+            }, 5000);
+        }
+        
+        function hideStatus() {
+            const statusElem = document.getElementById('status-message');
+            statusElem.className = 'status-indicator';
+        }
+        
+        // Tab functionality
+        const tabs = document.querySelectorAll('.tabs button');
+        const tabContents = document.querySelectorAll('.tab-content');
+        
+        tabs.forEach(tab => {
+            tab.addEventListener('click', () => {
+                // Remove active class from all tabs
+                tabs.forEach(t => t.classList.remove('active'));
+                tabContents.forEach(tc => tc.classList.remove('active'));
+                
+                // Add active class to clicked tab
+                tab.classList.add('active');
+                
+                // Show corresponding tab content
+                const tabId = tab.id.replace('tab-', '');
+                document.getElementById(`${tabId}-tab`).classList.add('active');
+                
+                // Hide any status messages when changing tabs
+                hideStatus();
+            });
+        });
+        
+        // Temperature slider
+        const temperatureSlider = document.getElementById('temperature');
+        const temperatureValue = document.getElementById('temperature-value');
+        
+        temperatureSlider.addEventListener('input', () => {
+            temperatureValue.textContent = temperatureSlider.value;
+        });
+        
+        // Load voices
+        async function loadVoices() {
+            try {
+                const response = await fetch('/v1/voice-cloning/voices');
+                const data = await response.json();
+                const voiceList = document.getElementById('voice-list');
+                const voiceSelect = document.getElementById('voice-select');
+                
+                // Clear existing content
+                voiceList.innerHTML = '';
+                
+                // Clear voice select options but keep the first one
+                while (voiceSelect.options.length > 1) {
+                    voiceSelect.remove(1);
+                }
+                
+                if (data.voices && data.voices.length > 0) {
+                    data.voices.forEach(voice => {
+                        // Add to voice list
+                        const voiceCard = document.createElement('div');
+                        voiceCard.className = 'voice-card';
+                        voiceCard.innerHTML = `
+                            <h3 class="voice-name">${voice.name}</h3>
+                            <p>${voice.description || 'No description'}</p>
+                            <p>Created: ${new Date(voice.created_at * 1000).toLocaleString()}</p>
+                            <div class="controls">
+                                <button class="btn-secondary preview-voice" data-id="${voice.id}">Preview</button>
+                                <button class="btn-danger delete-voice" data-id="${voice.id}">Delete</button>
+                            </div>
+                        `;
+                        voiceList.appendChild(voiceCard);
+                        
+                        // Add to voice select
+                        const option = document.createElement('option');
+                        option.value = voice.id;
+                        option.textContent = voice.name;
+                        voiceSelect.appendChild(option);
+                    });
+                    
+                    // Add event listeners for preview and delete buttons
+                    document.querySelectorAll('.preview-voice').forEach(button => {
+                        button.addEventListener('click', previewVoice);
+                    });
+                    document.querySelectorAll('.delete-voice').forEach(button => {
+                        button.addEventListener('click', deleteVoice);
+                    });
+                    
+                    showStatus(`Loaded ${data.voices.length} voices successfully`, 'success');
+                } else {
+                    voiceList.innerHTML = '<p>No cloned voices yet. Create one in the "Clone Voice" tab.</p>';
+                }
+            } catch (error) {
+                console.error('Error loading voices:', error);
+                showStatus('Failed to load voices', 'error');
+            }
+        }
+        
+        // Preview voice
+        async function previewVoice(event) {
+            const button = event.target;
+            const originalText = button.textContent;
+            button.disabled = true;
+            button.innerHTML = '<div class="spinner"></div> Loading...';
+            
+            const voiceId = button.dataset.id;
+            const audioPreview = document.getElementById('audio-preview');
+            
+            try {
+                const response = await fetch(`/v1/voice-cloning/voices/${voiceId}/preview`, {
+                    method: 'POST',
+                    headers: {
+                        'Content-Type': 'application/json'
+                    },
+                    body: JSON.stringify({
+                        text: "This is a preview of my cloned voice. I hope you like how it sounds!"
+                    })
+                });
+                
+                if (response.ok) {
+                    const blob = await response.blob();
+                    const url = URL.createObjectURL(blob);
+                    audioPreview.src = url;
+                    audioPreview.style.display = 'block';
+                    
+                    // Switch to the generate tab
+                    document.getElementById('tab-generate').click();
+                    
+                    // Set the voice in the select
+                    document.getElementById('voice-select').value = voiceId;
+                    audioPreview.play();
+                    
+                    showToast('Voice preview loaded', 'success');
+                } else {
+                    showToast('Failed to preview voice', 'error');
+                }
+            } catch (error) {
+                console.error('Error previewing voice:', error);
+                showToast('Error previewing voice', 'error');
+            } finally {
+                button.disabled = false;
+                button.textContent = originalText;
+            }
+        }
+        
+        // Delete voice
+        async function deleteVoice(event) {
+            if (!confirm('Are you sure you want to delete this voice? This cannot be undone.')) {
+                return;
+            }
+            
+            const button = event.target;
+            const originalText = button.textContent;
+            button.disabled = true;
+            button.innerHTML = '<div class="spinner"></div> Deleting...';
+            
+            const voiceId = button.dataset.id;
+            
+            try {
+                const response = await fetch(`/v1/voice-cloning/voices/${voiceId}`, {
+                    method: 'DELETE'
+                });
+                
+                if (response.ok) {
+                    showToast('Voice deleted successfully', 'success');
+                    loadVoices();
+                } else {
+                    showToast('Failed to delete voice', 'error');
+                }
+            } catch (error) {
+                console.error('Error deleting voice:', error);
+                showToast('Error deleting voice', 'error');
+            } finally {
+                button.disabled = false;
+                button.textContent = originalText;
+            }
+        }
+        
+        // Clone voice form submission
+        document.getElementById('clone-form').addEventListener('submit', async (event) => {
+            event.preventDefault();
+            const formData = new FormData(event.target);
+            const submitButton = event.target.querySelector('button[type="submit"]');
+            const originalText = submitButton.textContent;
+            
+            submitButton.disabled = true;
+            submitButton.innerHTML = '<div class="spinner"></div> Cloning Voice...';
+            showStatus('Processing your audio sample...', 'info');
+            
+            try {
+                const response = await fetch('/v1/voice-cloning/clone', {
+                    method: 'POST',
+                    body: formData
+                });
+                
+                if (response.ok) {
+                    const result = await response.json();
+                    showStatus('Voice cloned successfully!', 'success');
+                    showToast('Voice cloned successfully!', 'success');
+                    event.target.reset();
+                    
+                    // Switch to the voices tab
+                    document.getElementById('tab-voices').click();
+                    loadVoices();
+                } else {
+                    const error = await response.json();
+                    showStatus(`Failed to clone voice: ${error.detail}`, 'error');
+                    showToast('Failed to clone voice', 'error');
+                }
+            } catch (error) {
+                console.error('Error cloning voice:', error);
+                showStatus('Error processing your request', 'error');
+                showToast('Error cloning voice', 'error');
+            } finally {
+                submitButton.disabled = false;
+                submitButton.textContent = originalText;
+            }
+        });
+        
+        // YouTube voice cloning form submission
+        document.getElementById('youtube-clone-form').addEventListener('submit', async (event) => {
+            event.preventDefault();
+            const formData = new FormData(event.target);
+            const youtubeUrl = formData.get('youtube_url');
+            const voiceName = formData.get('voice_name');
+            const startTime = parseInt(formData.get('start_time'));
+            const duration = parseInt(formData.get('duration'));
+            const description = formData.get('description');
+            const progressDiv = document.getElementById('youtube-progress');
+            const progressStatus = document.getElementById('progress-status');
+            const progressFill = document.querySelector('.progress-fill');
+            const submitButton = event.target.querySelector('button[type="submit"]');
+            const originalText = submitButton.textContent;
+            
+            submitButton.disabled = true;
+            submitButton.innerHTML = '<div class="spinner"></div> Processing...';
+            showStatus('Starting YouTube download...', 'info');
+            
+            // Show progress bar
+            progressDiv.style.display = 'block';
+            progressFill.style.width = '10%';
+            progressStatus.textContent = 'Downloading audio...';
+            
+            // Simulate progress updates (since we can't get real-time updates easily)
+            let progress = 10;
+            const progressInterval = setInterval(() => {
+                if (progress < 90) {
+                    progress += 5;
+                    progressFill.style.width = `${progress}%`;
+                    if (progress > 30 && progress < 60) {
+                        progressStatus.textContent = 'Generating transcript...';
+                    } else if (progress >= 60) {
+                        progressStatus.textContent = 'Cloning voice...';
+                    }
+                }
+            }, 1000);
+            
+            try {
+                const response = await fetch('/v1/voice-cloning/clone-from-youtube', {
+                    method: 'POST',
+                    headers: {
+                        'Content-Type': 'application/json'
+                    },
+                    body: JSON.stringify({
+                        youtube_url: youtubeUrl,
+                        voice_name: voiceName,
+                        start_time: startTime,
+                        duration: duration,
+                        description: description
+                    })
+                });
+                
+                clearInterval(progressInterval);
+                
+                if (response.ok) {
+                    progressFill.style.width = '100%';
+                    progressStatus.textContent = 'Complete!';
+                    
+                    const result = await response.json();
+                    showStatus('Voice cloned successfully from YouTube!', 'success');
+                    showToast('Voice cloned from YouTube!', 'success');
+                    
+                    event.target.reset();
+                    
+                    // Switch to the voices tab
+                    document.getElementById('tab-voices').click();
+                    loadVoices();
+                } else {
+                    const error = await response.json();
+                    showStatus(`Failed to clone voice from YouTube: ${error.detail}`, 'error');
+                    showToast('Failed to clone voice from YouTube', 'error');
+                    progressDiv.style.display = 'none';
+                }
+            } catch (error) {
+                console.error('Error cloning voice from YouTube:', error);
+                showStatus('Error processing YouTube video', 'error');
+                showToast('Error cloning voice from YouTube', 'error');
+                progressDiv.style.display = 'none';
+            } finally {
+                clearInterval(progressInterval);
+                submitButton.disabled = false;
+                submitButton.textContent = originalText;
+            }
+        });
+        
+        // Generate speech form submission
+        document.getElementById('generate-form').addEventListener('submit', async (event) => {
+            event.preventDefault();
+            const formData = new FormData(event.target);
+            const voiceId = formData.get('voice');
+            const text = formData.get('text');
+            const temperature = formData.get('temperature');
+            const useStreaming = formData.get('use_streaming') === 'on';
+            
+            if (!voiceId) {
+                showToast('Please select a voice', 'warning');
+                return;
+            }
+            
+            const submitButton = event.target.querySelector('button[type="submit"]');
+            const originalText = submitButton.textContent;
+            submitButton.disabled = true;
+            submitButton.innerHTML = '<div class="spinner"></div> Generating...';
+            showStatus(useStreaming ? 'Streaming speech...' : 'Generating speech...', 'info');
+            
+            try {
+                const audioPreview = document.getElementById('audio-preview');
+                
+                if (useStreaming) {
+                    // For streaming, we need to handle the response differently to play audio as it arrives
+                    try {
+                        // Reset audio element
+                        audioPreview.style.display = 'block';
+                        
+                        // Prepare the request
+                        const requestOptions = {
+                            method: 'POST',
+                            headers: {
+                                'Content-Type': 'application/json'
+                            },
+                            body: JSON.stringify({
+                                model: "csm-1b",
+                                input: text,
+                                voice: voiceId,
+                                response_format: "mp3",
+                                temperature: parseFloat(temperature),
+                                speed: 1.0
+                            })
+                        };
+                        
+                        // Create a unique URL for this request to avoid caching issues
+                        const timestamp = new Date().getTime();
+                        const streamingUrl = `/v1/audio/speech/streaming?t=${timestamp}`;
+                        
+                        // Fetch from streaming endpoint
+                        const response = await fetch(streamingUrl, requestOptions);
+                        
+                        if (response.ok) {
+                            // Create a blob URL for immediate playback
+                            const blob = await response.blob();
+                            const url = URL.createObjectURL(blob);
+                            
+                            // Set the audio source and play immediately
+                            audioPreview.src = url;
+                            audioPreview.autoplay = true;
+                            
+                            // Event listeners for success/failure
+                            audioPreview.onplay = () => {
+                                showStatus('Speech streamed successfully', 'success');
+                                showToast('Speech streaming playback started', 'success');
+                            };
+                            
+                            audioPreview.onerror = (e) => {
+                                console.error('Audio playback error:', e);
+                                showStatus('Error playing streamed audio', 'error');
+                                showToast('Streaming playback error', 'error');
+                            };
+                        } else {
+                            const error = await response.json();
+                            showStatus(`Failed to stream speech: ${error.detail || 'Unknown error'}`, 'error');
+                            showToast('Failed to stream speech', 'error');
+                        }
+                    } catch (error) {
+                        console.error('Streaming error:', error);
+                        showStatus(`Error streaming speech: ${error.message}`, 'error');
+                        showToast('Error streaming speech', 'error');
+                    }
+                } else {
+                    // Non-streaming uses the original endpoint
+                    const response = await fetch('/v1/voice-cloning/generate', {
+                        method: 'POST',
+                        headers: {
+                            'Content-Type': 'application/json'
+                        },
+                        body: JSON.stringify({
+                            voice_id: voiceId,
+                            text: text,
+                            temperature: parseFloat(temperature)
+                        })
+                    });
+                    
+                    if (response.ok) {
+                        const blob = await response.blob();
+                        const url = URL.createObjectURL(blob);
+                        audioPreview.src = url;
+                        audioPreview.style.display = 'block';
+                        audioPreview.play();
+                        showStatus('Speech generated successfully', 'success');
+                        showToast('Speech generated successfully', 'success');
+                    } else {
+                        const error = await response.json();
+                        showStatus(`Failed to generate speech: ${error.detail}`, 'error');
+                        showToast('Failed to generate speech', 'error');
+                    }
+                }
+            } catch (error) {
+                console.error('Error generating speech:', error);
+                showStatus('Error generating speech', 'error');
+                showToast('Error generating speech', 'error');
+            } finally {
+                submitButton.disabled = false;
+                submitButton.textContent = originalText;
+            }
+        });
+
+        // Load voices on page load
+        loadVoices();
+    </script>
+</body>
+</html>