PyPI - pocket-tts - Versions diffs - 1.0.2__py3-none-any.whl - Mend

pocket-tts 1.0.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (38) hide show

pocket_tts/__init__.py +16 -0
pocket_tts/__main__.py +6 -0
pocket_tts/conditioners/__init__.py +0 -0
pocket_tts/conditioners/base.py +38 -0
pocket_tts/conditioners/text.py +61 -0
pocket_tts/config/b6369a24.yaml +57 -0
pocket_tts/data/__init__.py +2 -0
pocket_tts/data/audio.py +144 -0
pocket_tts/data/audio_utils.py +28 -0
pocket_tts/default_parameters.py +7 -0
pocket_tts/main.py +262 -0
pocket_tts/models/__init__.py +3 -0
pocket_tts/models/flow_lm.py +208 -0
pocket_tts/models/mimi.py +111 -0
pocket_tts/models/tts_model.py +782 -0
pocket_tts/modules/__init__.py +1 -0
pocket_tts/modules/conv.py +161 -0
pocket_tts/modules/dummy_quantizer.py +18 -0
pocket_tts/modules/layer_scale.py +11 -0
pocket_tts/modules/mimi_transformer.py +285 -0
pocket_tts/modules/mlp.py +215 -0
pocket_tts/modules/resample.py +46 -0
pocket_tts/modules/rope.py +74 -0
pocket_tts/modules/seanet.py +180 -0
pocket_tts/modules/stateful_module.py +45 -0
pocket_tts/modules/transformer.py +124 -0
pocket_tts/static/index.html +374 -0
pocket_tts/utils/__init__.py +1 -0
pocket_tts/utils/config.py +122 -0
pocket_tts/utils/debugging.py +26 -0
pocket_tts/utils/logging_utils.py +41 -0
pocket_tts/utils/utils.py +103 -0
pocket_tts/utils/weights_loading.py +35 -0
pocket_tts-1.0.2.dist-info/METADATA +174 -0
pocket_tts-1.0.2.dist-info/RECORD +38 -0
pocket_tts-1.0.2.dist-info/WHEEL +4 -0
pocket_tts-1.0.2.dist-info/entry_points.txt +2 -0
pocket_tts-1.0.2.dist-info/licenses/LICENSE +23 -0

pocket_tts/static/index.html ADDED Viewed

@@ -0,0 +1,374 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Pocket TTS - Streaming</title>
+    <script src="https://cdn.tailwindcss.com"></script>
+    <style>
+        body {
+            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+        }
+        .spinner {
+            animation: spin 1s linear infinite;
+        }
+        @keyframes spin {
+            from { transform: rotate(0deg); }
+            to { transform: rotate(360deg); }
+        }
+    </style>
+</head>
+<body class="bg-gray-50 min-h-screen">
+    <div class="max-w-xl mx-auto p-4 space-y-4">
+        <h1 class="text-xl font-bold text-center">Pocket TTS</h1>
+        <div class="bg-white p-4 rounded-lg shadow">
+            <textarea
+                id="text-input"
+                class="w-full border rounded p-2"
+                placeholder="Enter your text here..."
+                rows="4"
+            >Hello world. I am Kyutai's Pocket TTS. I'm fast enough to run on small CPUs. I hope you'll like me.</textarea>
+        </div>
+        <div class="bg-white p-4 rounded-lg shadow">
+            <label for="voice-url-input" class="block text-sm font-medium text-gray-700 mb-2">
+                Optional voice URL (leave empty to use default voice):
+            </label>
+            <input
+                type="text"
+                id="voice-url-input"
+                class="w-full border rounded p-2"
+                placeholder="hf://kyutai/tts-voices/alba-mackenna/casual.wav"
+                value="alba"
+            />
+            <p class="text-xs text-gray-500 mt-1">
+                Supports: http://, https://, or hf:// URLs.<br>
+                You can also use predefined voices:<br>
+                "alba", "marius", "javert", "jean", "fantine", "cosette", "eponine", "azelma".
+            </p>
+        </div>
+        <div class="bg-white p-4 rounded-lg shadow">
+            <label for="voice-wav-input" class="block text-sm font-medium text-gray-700 mb-2">
+                Or upload an audio file for voice cloning:
+            </label>
+            <input
+                type="file"
+                id="voice-wav-input"
+                class="w-full border rounded p-2"
+                accept=".wav,.mp3,.flac,.ogg,.m4a,audio/*"
+            />
+            <p class="text-xs text-gray-500 mt-1">
+                Upload an audio file (WAV, MP3, FLAC, etc.) to use as voice reference. Takes precedence over voice URL.
+            </p>
+        </div>
+        <button
+            id="generate-btn"
+            class="w-full px-4 py-2 rounded text-white bg-blue-600 hover:bg-blue-700 disabled:bg-gray-400 disabled:cursor-not-allowed"
+        >
+            <span id="generate-text">Generate audio</span>
+        </button>
+        <div id="status" class="hidden">
+            <div class="flex items-center space-x-2">
+                <div id="status-spinner" class="spinner w-4 h-4 border-2 border-blue-600 border-t-transparent rounded-full"></div>
+                <span id="status-text"></span>
+            </div>
+        </div>
+        <div id="audio-section" class="hidden">
+            <div class="bg-white p-4 rounded-lg shadow">
+                <h3 class="text-lg font-semibold mb-2">Generated Audio</h3>
+                <audio id="audio" controls class="w-full"></audio>
+                <div class="mt-2">
+                    <a id="download-link" class="text-blue-600 text-sm hover:underline">Download</a>
+                </div>
+            </div>
+        </div>
+    </div>
+    <script>
+        class StreamingWavPlayer {
+            constructor() {
+                this.audioContext = new (window.AudioContext || window.webkitAudioContext)();
+                this.sampleRate = 0;
+                this.numChannels = 0;
+                this.headerParsed = false;
+                this.headerBuffer = new Uint8Array(44);
+                this.headerBytesReceived = 0;
+                this.nextStartTime = 0;
+                this.isPlaying = false;
+                this.minBufferSize = 16384;
+                this.pcmData = new Uint8Array(0);
+            }
+            parseWavHeader(header) {
+                const view = new DataView(header.buffer);
+                const riff = String.fromCharCode.apply(null, Array.from(header.slice(0, 4)));
+                const wave = String.fromCharCode.apply(null, Array.from(header.slice(8, 12)));
+                if (riff !== 'RIFF' || wave !== 'WAVE') {
+                    throw new Error('Invalid WAV file');
+                }
+                this.numChannels = view.getUint16(22, true);
+                this.sampleRate = view.getUint32(24, true);
+                const bitsPerSample = view.getUint16(34, true);
+                console.log(`WAV Format: ${this.sampleRate}Hz, ${this.numChannels} channels, ${bitsPerSample} bits`);
+                this.headerParsed = true;
+            }
+            appendPcmData(newData) {
+                const newBuffer = new Uint8Array(this.pcmData.length + newData.length);
+                newBuffer.set(this.pcmData);
+                newBuffer.set(newData, this.pcmData.length);
+                this.pcmData = newBuffer;
+            }
+            async tryPlayBuffer() {
+                if (!this.headerParsed || this.pcmData.length < this.minBufferSize) {
+                    return;
+                }
+                const bytesPerSample = this.numChannels * 2;
+                const samplesToPlay = Math.floor(this.pcmData.length / bytesPerSample);
+                const bytesToPlay = samplesToPlay * bytesPerSample;
+                if (bytesToPlay === 0) return;
+                const dataToPlay = this.pcmData.slice(0, bytesToPlay);
+                this.pcmData = this.pcmData.slice(bytesToPlay);
+                const audioBuffer = this.audioContext.createBuffer(
+                    this.numChannels,
+                    samplesToPlay,
+                    this.sampleRate
+                );
+                const int16Data = new Int16Array(dataToPlay.buffer, dataToPlay.byteOffset, samplesToPlay * this.numChannels);
+                for (let channel = 0; channel < this.numChannels; channel++) {
+                    const channelData = audioBuffer.getChannelData(channel);
+                    for (let i = 0; i < samplesToPlay; i++) {
+                        channelData[i] = int16Data[i * this.numChannels + channel] / 32768;
+                    }
+                }
+                const source = this.audioContext.createBufferSource();
+                source.buffer = audioBuffer;
+                source.connect(this.audioContext.destination);
+                const currentTime = this.audioContext.currentTime;
+                const startTime = Math.max(currentTime, this.nextStartTime);
+                source.start(startTime);
+                // Track first audio playback
+                if (!this.firstAudioPlayed && window.firstAudioCallback) {
+                    this.firstAudioPlayed = true;
+                    window.firstAudioCallback();
+                }
+                this.nextStartTime = startTime + audioBuffer.duration;
+                this.isPlaying = true;
+                if (this.pcmData.length >= this.minBufferSize) {
+                    setTimeout(() => this.tryPlayBuffer(), 10);
+                }
+            }
+            addChunk(chunk) {
+                if (!this.headerParsed) {
+                    const headerBytesNeeded = 44 - this.headerBytesReceived;
+                    const bytesToCopy = Math.min(headerBytesNeeded, chunk.length);
+                    this.headerBuffer.set(
+                        chunk.slice(0, bytesToCopy),
+                        this.headerBytesReceived
+                    );
+                    this.headerBytesReceived += bytesToCopy;
+                    if (this.headerBytesReceived >= 44) {
+                        this.parseWavHeader(this.headerBuffer);
+                        if (chunk.length > bytesToCopy) {
+                            this.appendPcmData(chunk.slice(bytesToCopy));
+                        }
+                    }
+                } else {
+                    this.appendPcmData(chunk);
+                }
+                this.tryPlayBuffer();
+            }
+            stop() {
+                this.audioContext.close();
+                this.isPlaying = false;
+            }
+        }
+        // Application state
+        let streamingPlayer = null;
+        let currentAudioBlob = null;
+        // DOM elements
+        const textInput = document.getElementById('text-input');
+        const generateBtn = document.getElementById('generate-btn');
+        const generateText = document.getElementById('generate-text');
+        const status = document.getElementById('status');
+        const statusText = document.getElementById('status-text');
+        const statusSpinner = document.getElementById('status-spinner');
+        const audioSection = document.getElementById('audio-section');
+        const audio = document.getElementById('audio');
+        const downloadLink = document.getElementById('download-link');
+        // Event listeners
+        generateBtn.addEventListener('click', generateAudio);
+        function showStatus(message, isLoading = false) {
+            statusText.textContent = message;
+            statusSpinner.style.display = isLoading ? 'block' : 'none';
+            status.classList.remove('hidden');
+        }
+        function hideStatus() {
+            status.classList.add('hidden');
+        }
+        async function generateAudio() {
+            const text = textInput.value.trim();
+            if (!text) {
+                showStatus('Please enter some text to generate speech.', false);
+                setTimeout(hideStatus, 3000);
+                return;
+            }
+            // Stop any currently playing audio
+            if (streamingPlayer) {
+                streamingPlayer.stop();
+                streamingPlayer = null;
+            }
+            // Track timing
+            const startTime = performance.now();
+            let firstAudioTime = null;
+            // Set callback for first audio
+            window.firstAudioCallback = () => {
+                if (!firstAudioTime) {
+                    firstAudioTime = performance.now();
+                    const timeToFirstAudio = ((firstAudioTime - startTime) / 1000).toFixed(2);
+                    showStatus(`First audio in ${timeToFirstAudio}s...`, true);
+                }
+            };
+            // Update UI
+            generateBtn.disabled = true;
+            generateText.textContent = 'Generating...';
+            showStatus('Generating speech...', true);
+            audioSection.classList.add('hidden');
+            try {
+                const formData = new FormData();
+                formData.append('text', text);
+                // Add voice URL if provided (only if no WAV file is uploaded)
+                const voiceUrl = document.getElementById('voice-url-input').value.trim();
+                const voiceWavFile = document.getElementById('voice-wav-input').files[0];
+                if (voiceWavFile) {
+                    // If WAV file is uploaded, only use the WAV file (ignore voice URL)
+                    formData.append('voice_wav', voiceWavFile);
+                } else if (voiceUrl) {
+                    // Only use voice URL if no WAV file is uploaded
+                    formData.append('voice_url', voiceUrl);
+                }
+                const response = await fetch('/tts', {
+                    method: 'POST',
+                    body: formData
+                });
+                if (!response.ok) {
+                    throw new Error(`Server error: ${response.status}`);
+                }
+                // Clone the response for both streaming and blob collection
+                const responseForPlayback = response.clone();
+                const responseForHistory = response.clone();
+                // Start streaming playback
+                const reader = responseForPlayback.body.getReader();
+                streamingPlayer = new StreamingWavPlayer();
+                const processStream = async () => {
+                    try {
+                        while (true) {
+                            const { done, value } = await reader.read();
+                            if (done) break;
+                            if (value) {
+                                streamingPlayer.addChunk(value);
+                            }
+                        }
+                    } catch (e) {
+                        console.error('Error processing stream:', e);
+                    }
+                };
+                // Start processing the stream
+                processStream();
+                // Also collect the full blob for download/audio element
+                const blob = await responseForHistory.blob();
+                const totalTime = ((performance.now() - startTime) / 1000).toFixed(2);
+                currentAudioBlob = blob;
+                // Update UI with audio
+                const audioUrl = URL.createObjectURL(blob);
+                audio.src = audioUrl;
+                // Wait for audio metadata to load to get duration
+                audio.addEventListener('loadedmetadata', () => {
+                    const audioDuration = audio.duration;
+                    const speedRatio = (audioDuration / parseFloat(totalTime)).toFixed(1);
+                    const timeToFirst = firstAudioTime ? ((firstAudioTime - startTime) / 1000).toFixed(2) : 'N/A';
+                    showStatus(
+                        `✨ First audio: ${timeToFirst}s | Total: ${totalTime}s | ${speedRatio}x faster than real-time`,
+                        false
+                    );
+                });
+                downloadLink.href = audioUrl;
+                downloadLink.download = `tts-audio.wav`;
+                audioSection.classList.remove('hidden');
+            } catch (error) {
+                console.error('Error generating audio:', error);
+                showStatus(`Error: ${error.message}`, false);
+                setTimeout(hideStatus, 3000);
+            } finally {
+                // Re-enable button
+                generateBtn.disabled = false;
+                generateText.textContent = 'Generate Audio';
+            }
+        }
+        // Focus on text input when page loads
+        window.addEventListener('load', () => {
+            textInput.focus();
+        });
+    </script>
+</body>
+</html>

pocket_tts/utils/__init__.py ADDED Viewed

	@@ -0,0 +1 @@
1	+ """Utilities."""

pocket_tts/utils/config.py ADDED Viewed

@@ -0,0 +1,122 @@
+"""Configuration models for loading YAML config files."""
+from pathlib import Path
+import yaml
+from pydantic import BaseModel, ConfigDict
+class StrictModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+# Flow configuration
+class FlowConfig(StrictModel):
+    dim: int
+    depth: int
+# Transformer configuration for FlowLM
+class FlowLMTransformerConfig(StrictModel):
+    hidden_scale: int
+    max_period: int
+    d_model: int
+    num_heads: int
+    num_layers: int
+class LookupTable(StrictModel):
+    dim: int
+    n_bins: int
+    tokenizer: str
+    tokenizer_path: str
+# Root configuration
+class FlowLMConfig(StrictModel):
+    """Root configuration model for YAML config files."""
+    dtype: str
+    # Nested configurations
+    flow: FlowConfig
+    transformer: FlowLMTransformerConfig
+    # conditioning
+    lookup_table: LookupTable
+    weights_path: str | None = None
+# SEANet configuration
+class SEANetConfig(StrictModel):
+    dimension: int
+    channels: int
+    n_filters: int
+    n_residual_layers: int
+    ratios: list[int]
+    kernel_size: int
+    residual_kernel_size: int
+    last_kernel_size: int
+    dilation_base: int
+    pad_mode: str
+    compress: int
+# Transformer configuration for Mimi
+class MimiTransformerConfig(StrictModel):
+    d_model: int
+    input_dimension: int
+    output_dimensions: tuple[int, ...]
+    num_heads: int
+    num_layers: int
+    layer_scale: float
+    context: int
+    max_period: float = 10000.0
+    dim_feedforward: int
+# Quantizer configuration
+class QuantizerConfig(StrictModel):
+    dimension: int
+    output_dimension: int
+# Root configuration
+class MimiConfig(StrictModel):
+    """Root configuration model for Mimi YAML config files."""
+    dtype: str
+    # Sample rate and channels
+    sample_rate: int
+    channels: int
+    frame_rate: float
+    # SEANet configurations
+    seanet: SEANetConfig
+    # Transformer
+    transformer: MimiTransformerConfig
+    # Quantizer
+    quantizer: QuantizerConfig
+    weights_path: str | None = None
+class Config(StrictModel):
+    flow_lm: FlowLMConfig
+    mimi: MimiConfig
+    weights_path: str | None = None
+    weights_path_without_voice_cloning: str | None = None
+def load_config(yaml_path: str | Path) -> Config:
+    yaml_path = Path(yaml_path)
+    if not yaml_path.exists():
+        raise FileNotFoundError(f"Config file not found: {yaml_path}")
+    with open(yaml_path, "r") as f:
+        config_dict = yaml.safe_load(f)
+    return Config(**config_dict)

pocket_tts/utils/debugging.py ADDED Viewed

@@ -0,0 +1,26 @@
+import torch
+from torch.utils._python_dispatch import TorchDispatchMode
+def to_str(obj):
+    if isinstance(obj, (torch.Tensor, torch.nn.Parameter)):
+        return f"T(s={list(obj.shape)})"
+    elif isinstance(obj, (list, tuple)):
+        return "[" + ", ".join(to_str(o) for o in obj) + "]"
+    elif isinstance(obj, dict):
+        return "{" + ", ".join(f"{to_str(k)}: {to_str(v)}" for k, v in obj.items()) + "}"
+    else:
+        return str(obj)
+class LoggingMode(TorchDispatchMode):
+    """Useful to check implementation differences."""
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        output = func(*args, **kwargs or {})
+        print(
+            f"Aten function called: {func}, args: "
+            f"{to_str(args)}, kwargs: {to_str(kwargs)} -> "
+            f"output: {to_str(output)}"
+        )
+        return output

pocket_tts/utils/logging_utils.py ADDED Viewed

@@ -0,0 +1,41 @@
+import logging
+from contextlib import contextmanager
+class PocketTTSFilter(logging.Filter):
+    def filter(self, record):
+        return record.name.startswith("pocket_tts")
+@contextmanager
+def enable_logging(library_name, level):
+    # Get the specific logger and its parent
+    logger = logging.getLogger(library_name)
+    parent_logger = logging.getLogger("pocket_tts")
+    # Store original configuration
+    old_level = logger.level
+    old_parent_level = parent_logger.level
+    old_handlers = parent_logger.handlers.copy()
+    # Configure logging format for pocket_tts logger
+    parent_logger.setLevel(level)
+    # Clear existing handlers and add our custom formatter with filter
+    parent_logger.handlers.clear()
+    handler = logging.StreamHandler()
+    formatter = logging.Formatter("%(levelname)s: %(message)s")
+    handler.setFormatter(formatter)
+    handler.addFilter(PocketTTSFilter())
+    parent_logger.addHandler(handler)
+    parent_logger.propagate = False
+    try:
+        yield logger
+    finally:
+        # Restore original configuration
+        logger.setLevel(old_level)
+        parent_logger.setLevel(old_parent_level)
+        parent_logger.handlers.clear()
+        for h in old_handlers:
+            parent_logger.addHandler(h)

pocket_tts/utils/utils.py ADDED Viewed

@@ -0,0 +1,103 @@
+import hashlib
+import logging
+import time
+from pathlib import Path
+import requests
+import safetensors.torch
+import torch
+from huggingface_hub import hf_hub_download
+from torch import nn
+PROJECT_ROOT = Path(__file__).parent.parent.parent
+_voices_names = ["alba", "marius", "javert", "jean", "fantine", "cosette", "eponine", "azelma"]
+PREDEFINED_VOICES = {
+    # don't forget to change this
+    x: f"hf://kyutai/pocket-tts-without-voice-cloning/embeddings/{x}.safetensors@d4fdd22ae8c8e1cb3634e150ebeff1dab2d16df3"
+    for x in _voices_names
+}
+def make_cache_directory() -> Path:
+    cache_dir = Path.home() / ".cache" / "pocket_tts"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+def print_nb_parameters(model: nn.Module, model_name: str):
+    logger = logging.getLogger(__name__)
+    state_dict = model.state_dict()
+    total = 0
+    for key, value in state_dict.items():
+        logger.info("%s: %,d", key, value.numel())
+        total += value.numel()
+    logger.info("Total number of parameters in %s: %,d", model_name, total)
+def size_of_dict(state_dict: dict) -> int:
+    total_size = 0
+    for value in state_dict.values():
+        if isinstance(value, torch.Tensor):
+            total_size += value.numel() * value.element_size()
+        elif isinstance(value, dict):
+            total_size += size_of_dict(value)
+    return total_size
+class display_execution_time:
+    def __init__(self, task_name: str, print_output: bool = True):
+        self.task_name = task_name
+        self.print_output = print_output
+        self.start_time = None
+        self.elapsed_time_ms = None
+        self.logger = logging.getLogger(__name__)
+    def __enter__(self):
+        self.start_time = time.monotonic()
+        return self
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        end_time = time.monotonic()
+        self.elapsed_time_ms = int((end_time - self.start_time) * 1000)
+        if self.print_output:
+            self.logger.info("%s took %d ms", self.task_name, self.elapsed_time_ms)
+        return False  # Don't suppress exceptions
+def download_if_necessary(file_path: str) -> Path:
+    if file_path.startswith("http://") or file_path.startswith("https://"):
+        cache_dir = make_cache_directory()
+        cached_file = cache_dir / (
+            hashlib.sha256(file_path.encode()).hexdigest() + "." + file_path.split(".")[-1]
+        )
+        if not cached_file.exists():
+            response = requests.get(file_path)
+            response.raise_for_status()
+            with open(cached_file, "wb") as f:
+                f.write(response.content)
+        return cached_file
+    elif file_path.startswith("hf://"):
+        file_path = file_path.removeprefix("hf://")
+        splitted = file_path.split("/")
+        repo_id = "/".join(splitted[:2])
+        filename = "/".join(splitted[2:])
+        if "@" in filename:
+            filename, revision = filename.split("@")
+        else:
+            revision = None
+        cached_file = hf_hub_download(repo_id=repo_id, filename=filename, revision=revision)
+        return Path(cached_file)
+    else:
+        return Path(file_path)
+def load_predefined_voice(voice_name: str) -> torch.Tensor:
+    if voice_name not in PREDEFINED_VOICES:
+        raise ValueError(
+            f"Predefined voice '{voice_name}' not found"
+            f", available voices are {list(PREDEFINED_VOICES)}."
+        )
+    voice_file = download_if_necessary(PREDEFINED_VOICES[voice_name])
+    # There is only one tensor in the file.
+    return safetensors.torch.load_file(voice_file)["audio_prompt"]