update

BATCH_IMPLEMENTATION_PLAN.md

@@ -1,58 +0,0 @@
-# Plan for Implementing High-Performance Batch Processing
-
-This document outlines the necessary code modifications to implement a high-performance batch processing mode that can be toggled by the "Use VADER" checkbox in the GUI.
-
-The goal is to create two distinct modes:
-- **VADER On (Nuanced Mode):** Slower, processes chunks one-by-one with unique TTS parameters for nuanced delivery.
-- **VADER Off (Batch Mode):** Significantly faster, processes chunks in batches with a single set of TTS parameters.
-
----
-
-## 1. File to Modify: `src/chatterbox/tts.py`
-
-*   **Purpose:** To enable the core TTS model to handle batches of text.
-*   **Changes Needed:**
-    *   A new method, `generate_batch(self, texts: list, **tts_params)`, needs to be created within the `ChatterboxTTS` class.
-    *   This method must perform the following steps:
-        1.  Accept a list of text strings (`texts`).
-        2.  Tokenize each text string in the list.
-        3.  Pad the tokenized sequences to ensure they all have the same length, creating a single batch tensor. `torch.nn.utils.rnn.pad_sequence` is suitable for this.
-        4.  Feed the complete batch tensor to the underlying model (`self.t3.inference` and `self.s3gen.inference`).
-        5.  Return a list of the resulting audio waveforms.
-
----
-
-## 2. File to Modify: `modules/tts_engine.py`
-
-*   **Purpose:** To orchestrate the new batching workflow and choose the processing mode.
-*   **Changes Needed:**
-
-    ### a. Create a New Worker Function
-    *   Add a new function: `process_batch(batch_of_chunks, model, ...)`
-    *   This function will:
-        1.  Accept a list of chunk objects (e.g., a batch of 16).
-        2.  Extract the text from each chunk into a simple list.
-        3.  Call the new `model.generate_batch()` with the list of texts and the shared TTS parameters.
-        4.  Receive a list of audio waveforms back.
-        5.  Loop through the audio waves, apply the existing silence trimming and padding logic to each one, and save them to their respective `chunk_...wav` files.
-
-    ### b. Modify the Main `process_book_folder` Function
-    *   Locate the `use_vader` flag which is determined from the GUI options.
-    *   Wrap the core processing loop in an `if/else` block based on this flag.
-    *   **`if use_vader:` (Nuanced Mode):**
-        *   Keep the existing code that iterates through chunks one-by-one and submits them to the `process_one_chunk` function.
-    *   **`else:` (Batch Mode):**
-        *   Add the new logic here.
-        *   Group the `all_chunks` list into fixed-size batches based on `TTS_BATCH_SIZE` from the config.
-        *   Use the existing `ThreadPoolExecutor` to submit these new **batches** to the new `process_batch` worker function.
-
----
-
-## 3. Files to Modify: `config/config.py` and `chatterbox_gui.py`
-
-*   **Purpose:** To provide user control over the batch size for performance tuning.
-*   **Changes Needed:**
-    *   **In `config/config.py`:**
-        *   Add a new configuration variable: `TTS_BATCH_SIZE = 16` (or another sensible default).
-    *   **In `chatterbox_gui.py`:**
-        *   On the "Config" tab, add a new `QSpinBox` (numeric input field) that is linked to the `TTS_BATCH_SIZE` variable. This will allow the user to change the batch size without editing code.

HF_Deploy/.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text

HF_Deploy/.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

HF_Deploy/README.md

@@ -0,0 +1,14 @@
+---
+title: ChatterboxTTS DNXS Spokenword
+emoji: 🌖
+colorFrom: blue
+colorTo: red
+sdk: gradio
+sdk_version: 5.39.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+short_description: 'ChatterboxTTS Gradio interface for custom workflow. '
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

HF_Deploy/Text_Input/Goliath/test1.txt

@@ -0,0 +1,7 @@
+My dear fellow, I spent a considerable length of time in the Peninsula.  I was with a British rifle brigade when I met Sir Arthur Wellesley.  And I was a prisoner of the French at Salamanca - 18 12 I think it was.  I always find it’s best to see both sides of both sides, if you see what I mean.
+
+I didn’t really think you approved of war sir, said Benton sadly.
+
+The Doctor turned his attention back to the twisting country lane.  He sighed as he changed gear for another sharp corner.  Sometimes it’s inevitable, he noted with genuine sadness.  I’m a man of peace, but I seem to spend much of my time caught up in conflict.  The central paradox of my life, perhaps.
+
+Benton leant back in the seat.  What’s the central paradox of mine? he asked, fascinated.

HF_Deploy/Text_Input/README.md

@@ -0,0 +1,40 @@
+# Text Input Directory
+
+Place your book text files here for audiobook generation.
+
+## Directory Structure
+Create a subdirectory for each book:
+```
+Text_Input/
+├── Book Name 1/
+│   ├── book.txt          # Main text file
+│   ├── cover.jpg         # Book cover image (optional)
+│   └── book.nfo          # Metadata file (optional)
+├── Book Name 2/
+│   ├── another_book.txt
+│   └── cover.png
+└── ...
+```
+
+## Text File Requirements
+- **Format**: Plain text (.txt) files
+- **Encoding**: UTF-8
+- **Content**: Clean text without excessive formatting
+- **Structure**: Use paragraph breaks for natural speech flow
+
+## Optional Files
+- **cover.jpg/png**: Book cover image for M4B metadata
+- **book.nfo**: XML metadata file with book information (title, author, etc.)
+
+## Text Preparation Tips
+- Remove table of contents, page numbers, headers/footers
+- Keep chapter headings (e.g., "Chapter 1")
+- Use proper punctuation for natural speech
+- Remove excessive line breaks or formatting
+- Ensure UTF-8 encoding for special characters
+
+## Processing
+1. Add your book directory to Text_Input/
+2. Run the main program and select your book
+3. The system will chunk the text and generate JSON metadata
+4. Use the generated chunks for TTS audiobook creation

HF_Deploy/Text_Input/test

@@ -0,0 +1,20 @@
+She stood alone in the hallway. The lights flickered overhead. "I don't like this," she whispered. "Too quiet. Too cold."
+
+
+
+***
+
+Chapter 1
+
+A crash echoed from somewhere far off.
+He turned. "Was that you?"
+
+"No," she said. "It wasn't me."
+
+---
+
+They moved cautiously down the corridor. Every step sounded like thunder. Each shadow seemed to breathe.
+
+Chapter 2
+
+Something moved behind the curtain.

HF_Deploy/config/config.py

@@ -0,0 +1,159 @@
+"""
+GenTTS Configuration Module
+Central location for all settings, paths, and feature toggles
+"""
+
+import os
+from pathlib import Path
+
+# ============================================================================
+# CORE DIRECTORIES
+# ============================================================================
+TEXT_INPUT_ROOT = Path("Text_Input")
+AUDIOBOOK_ROOT = Path("Audiobook")
+VOICE_SAMPLES_DIR = Path("Voice_Samples")
+
+# ============================================================================
+# TEXT PROCESSING SETTINGS
+# ============================================================================
+MAX_CHUNK_WORDS = 28
+MIN_CHUNK_WORDS = 4
+
+# ============================================================================
+# WORKER AND PERFORMANCE SETTINGS
+# ============================================================================
+MAX_WORKERS = 2                       # Keep at 2 - GPU utilization already high
+TEST_MAX_WORKERS = 6                  # For experimentation
+USE_DYNAMIC_WORKERS = False           # Toggle for testing
+VRAM_SAFETY_THRESHOLD = 6.5           # GB
+
+# ============================================================================
+# AUDIO QUALITY SETTINGS
+# ============================================================================
+ENABLE_MID_DROP_CHECK = False
+ENABLE_ASR = False
+ASR_WORKERS = 4                       # Parallel ASR on CPU threads
+
+# ============================================================================
+# TTS HUM DETECTION SETTINGS
+# ============================================================================
+ENABLE_HUM_DETECTION = False          # Disabled for speed (re-enable if quality issues)
+HUM_FREQ_MIN = 50                     # Hz - Lower frequency bound for hum detection
+HUM_FREQ_MAX = 200                    # Hz - Upper frequency bound for hum detection
+HUM_ENERGY_THRESHOLD = 0.3            # Ratio of hum energy to total energy (0.1-0.5 range)
+HUM_STEADY_THRESHOLD = 0.6            # Ratio of segments with steady amplitude (0.5-0.8 range)
+HUM_AMPLITUDE_MIN = 0.005             # Minimum RMS for steady hum detection
+HUM_AMPLITUDE_MAX = 0.1               # Maximum RMS for steady hum detection
+
+# ============================================================================
+# AUDIO TRIMMING SETTINGS
+# ============================================================================
+ENABLE_AUDIO_TRIMMING = True          # Enable automatic audio trimming after TTS
+SPEECH_ENDPOINT_THRESHOLD = 0.005     # RMS threshold to detect end of speech (more aggressive)
+TRIMMING_BUFFER_MS = 50               # Small buffer after detected speech endpoint
+
+# ============================================================================
+# SILENCE DURATION SETTINGS (milliseconds)
+# ============================================================================
+SILENCE_CHAPTER_START = 500           # Half second for chapter beginnings
+SILENCE_CHAPTER_END = 800             # Longer pause before new chapter
+SILENCE_SECTION_BREAK = 600           # Section transitions
+SILENCE_PARAGRAPH_END = 300           # Standard paragraph breaks
+
+# Punctuation-specific silence settings (milliseconds)
+SILENCE_COMMA = 150                   # Brief pause after commas
+SILENCE_SEMICOLON = 250               # Medium pause after semicolons
+SILENCE_COLON = 300                   # Pause after colons
+SILENCE_PERIOD = 400                  # Sentence end pause
+SILENCE_QUESTION_MARK = 450           # Question pause (slightly longer)
+SILENCE_EXCLAMATION = 400             # Exclamation pause
+SILENCE_DASH = 200                    # Em dash pause
+SILENCE_ELLIPSIS = 350                # Ellipsis pause (suspense)
+SILENCE_QUOTE_END = 250               # End of quoted speech
+
+# Chunk-level silence settings
+ENABLE_CHUNK_END_SILENCE = True       # Add silence to end of every chunk
+CHUNK_END_SILENCE_MS = 200            # Default silence at end of each chunk
+
+# Content boundary silence settings (milliseconds)
+SILENCE_PARAGRAPH_FALLBACK = 500      # Original paragraph logic fallback
+
+# ============================================================================
+# AUDIO NORMALIZATION SETTINGS
+# ============================================================================
+ENABLE_NORMALIZATION = True           # Global ON/OFF switch for normalization
+NORMALIZATION_TYPE = "peak"           # Options: "loudness", "peak", "simple", "none"
+TARGET_LUFS = -16                     # Target loudness (LUFS) for broadcast standard
+TARGET_PEAK_DB = -1.5                 # Target peak level (dB) to prevent clipping
+TARGET_LRA = 11                       # Target loudness range for consistency
+
+# ============================================================================
+# AUDIO PLAYBACK SPEED SETTINGS
+# ============================================================================
+ATEMPO_SPEED = 0.95                    # Playback speed multiplier (0.5-2.0 range, 1.0 = normal speed)
+
+# ============================================================================
+# ENVIRONMENT SETUP
+# ============================================================================
+os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
+os.environ["TRANSFORMERS_NO_PROGRESS_BAR"] = "1"
+os.environ["HF_TRANSFORMERS_NO_TQDM"] = "1"
+os.environ["TORCH_HUB_DIR"] = "/tmp/torch_hub_silent"
+
+# ============================================================================
+# COLOR CODES FOR TERMINAL OUTPUT
+# ============================================================================
+RESET = "\033[0m"
+BOLD = "\033[1m"
+RED = "\033[91m"
+GREEN = "\033[92m"
+YELLOW = "\033[93m"
+CYAN = "\033[96m"
+
+# ============================================================================
+# TTS MODEL PARAMETERS (DEFAULTS)
+# ============================================================================
+DEFAULT_EXAGGERATION = 0.4            # Emotion intensity (0.0-2.0 range)
+DEFAULT_CFG_WEIGHT = 0.5             # Faithfulness to text (0.0-1.0 range)
+DEFAULT_TEMPERATURE = 0.9         # Randomness/creativity (0.0-1.0 range)
+
+# ============================================================================
+# VADER SENTIMENT TO TTS PARAMETER MAPPING
+# ============================================================================
+# These settings control how VADER sentiment analysis dynamically adjusts TTS parameters.
+# The formula used is: new_param = base_param + (compound_score * sensitivity)
+# The result is then clamped within the defined MIN/MAX range.
+
+# --- Base TTS Parameters (used as the starting point) ---
+# These are the same as the main defaults, but listed here for clarity.
+BASE_EXAGGERATION = DEFAULT_EXAGGERATION  # Default: 1.0
+BASE_CFG_WEIGHT = DEFAULT_CFG_WEIGHT      # Default: 0.7
+BASE_TEMPERATURE = DEFAULT_TEMPERATURE    # Default: 0.7
+
+# --- Sensitivity ---
+# How much VADER's compound score affects each parameter.
+# Higher values mean more dramatic changes based on sentiment.
+VADER_EXAGGERATION_SENSITIVITY = 0.5  # e.g., compound of 0.8 -> 1.0 + (0.8 * 0.5) = 1.4
+VADER_CFG_WEIGHT_SENSITIVITY = -0.2 # Negative: more emotional text is less strict
+VADER_TEMPERATURE_SENSITIVITY = 0.15  # More emotional text gets slightly more creative
+
+# --- Min/Max Clamps ---
+# Hard limits to prevent extreme, undesirable audio artifacts.
+TTS_PARAM_MIN_EXAGGERATION = 0.1
+TTS_PARAM_MAX_EXAGGERATION = 2.0
+TTS_PARAM_MIN_CFG_WEIGHT = 0.1
+TTS_PARAM_MAX_CFG_WEIGHT = 1.0
+
+TTS_PARAM_MIN_TEMPERATURE = 0.1
+TTS_PARAM_MAX_TEMPERATURE = 5.0
+
+# ============================================================================
+# BATCH PROCESSING SETTINGS
+# ============================================================================
+BATCH_SIZE = 250                      # Larger batches for better speed (monitor VRAM)
+CLEANUP_INTERVAL = 500                # Deep cleanup every N chunks (reduced frequency for speed)
+
+# ============================================================================
+# FEATURE TOGGLES
+# ============================================================================
+shutdown_requested = False           # Global shutdown flag

HF_Deploy/gradio_main_interface.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python3
+"""
+ChatterboxTTS DNXS-Spokneword Gradio Main Interface
+Modular web interface with separate tab modules
+"""
+
+import gradio as gr
+import sys
+import os
+from pathlib import Path
+
+# Add the current directory to Python path for imports
+sys.path.append(str(Path(__file__).parent))
+
+# Import tab modules
+try:
+    from gradio_tabs.tab1_convert_book import create_convert_book_tab
+    TAB1_AVAILABLE = True
+except ImportError as e:
+    print(f"⚠️  Tab 1 not available: {e}")
+    TAB1_AVAILABLE = False
+
+try:
+    from gradio_tabs.tab6_settings import create_settings_tab_interface
+    TAB6_AVAILABLE = True
+except ImportError as e:
+    print(f"⚠️  Tab 6 (Settings) not available: {e}")
+    TAB6_AVAILABLE = False
+
+def create_placeholder_tab(tab_name, tab_number):
+    """Create a placeholder tab for future implementation"""
+    with gr.Column():
+        gr.Markdown(f"# 🚧 {tab_name}")
+        gr.Markdown(f"*Tab {tab_number} - Coming Soon*")
+        gr.Markdown("This tab will be implemented in a future update.")
+
+        gr.Button("Placeholder Button", interactive=False)
+
+def create_main_interface():
+    """Create the main ChatterboxTTS Gradio interface with all tabs"""
+
+    with gr.Blocks(
+        title="ChatterboxTTS - Complete Interface",
+        theme=gr.themes.Soft(),
+        css="""
+        .gradio-container {
+            max-width: 1200px !important;
+        }
+        """
+    ) as demo:
+
+        # Header
+        gr.Markdown("""
+        # 🎤 ChatterboxTTS - Complete Web Interface
+        *Modular audiobook generation system with advanced TTS capabilities*
+        """)
+
+        # Tab interface
+        with gr.Tabs():
+            # Tab 1: Convert Book (Working)
+            if TAB1_AVAILABLE:
+                with gr.Tab("1. Convert Book"):
+                    create_convert_book_tab()
+            else:
+                with gr.Tab("1. Convert Book"):
+                    create_placeholder_tab("Convert Book", 1)
+
+            # Tab 2-10: Placeholders for now
+            with gr.Tab("2. File Management"):
+                create_placeholder_tab("File Management", 2)
+
+            with gr.Tab("3. Voice Analysis"):
+                create_placeholder_tab("Voice Analysis", 3)
+
+            with gr.Tab("4. Batch Processing"):
+                create_placeholder_tab("Batch Processing", 4)
+
+            with gr.Tab("5. Audio Tools"):
+                create_placeholder_tab("Audio Tools", 5)
+
+            # Tab 6: Settings (Working)
+            if TAB6_AVAILABLE:
+                with gr.Tab("6. Settings"):
+                    create_settings_tab_interface()
+            else:
+                with gr.Tab("6. Settings"):
+                    create_placeholder_tab("Settings", 6)
+
+            with gr.Tab("7. Chunk Tools"):
+                create_placeholder_tab("Chunk Tools", 7)
+
+            with gr.Tab("8. Voice Training"):
+                create_placeholder_tab("Voice Training", 8)
+
+            with gr.Tab("9. System Monitor"):
+                create_placeholder_tab("System Monitor", 9)
+
+            with gr.Tab("10. About"):
+                create_placeholder_tab("About", 10)
+
+        # Footer
+        gr.Markdown("""
+        ---
+        *ChatterboxTTS Gradio Interface - Modular Design*
+        Each tab is a separate module for easy maintenance and development.
+        """)
+
+    return demo
+
+def launch_interface():
+    """Launch the main interface"""
+    print("🚀 ChatterboxTTS - Starting Main Interface")
+    print("📊 Tab Status:")
+    print(f"   Tab 1 (Convert Book): {'✅ Available' if TAB1_AVAILABLE else '❌ Not Available'}")
+    print("   Tabs 2-10: 🚧 Placeholder (Coming Soon)")
+    print("-" * 50)
+
+    demo = create_main_interface()
+
+    # Launch configuration
+    launch_kwargs = {
+        'server_name': '0.0.0.0',
+        'server_port': 7860,
+        'show_error': True,
+        'quiet': False
+    }
+
+    # Detect cloud environments
+    if os.getenv("RUNPOD_POD_ID"):
+        print("☁️  RunPod deployment detected")
+        launch_kwargs['share'] = True
+    elif os.getenv("COLAB_GPU"):
+        print("☁️  Google Colab detected")
+        launch_kwargs['share'] = True
+    else:
+        print("💻 Local deployment")
+        launch_kwargs['share'] = False
+
+    print(f"🌐 Interface will be available at: http://localhost:{launch_kwargs['server_port']}")
+
+    try:
+        demo.launch(**launch_kwargs)
+    except Exception as e:
+        print(f"❌ Error launching interface: {e}")
+        raise
+
+if __name__ == "__main__":
+    launch_interface()

HF_Deploy/gradio_tabs/__init__.py

@@ -0,0 +1,7 @@
+"""
+ChatterboxTTS Gradio Tabs Package
+Modular tab system for the web interface
+"""
+
+# Make this directory a Python package
+__version__ = "1.0.0"

HF_Deploy/gradio_tabs/tab1_convert_book.py

@@ -0,0 +1,1173 @@
+#!/usr/bin/env python3
+"""
+Gradio Tab 1: Convert Book
+Exact replica of PyQt5 GUI Tab 1 functionality
+"""
+
+import gradio as gr
+import os
+import sys
+import threading
+import subprocess
+import tempfile
+import json
+import warnings
+import re
+import time
+from pathlib import Path
+from typing import List, Dict, Any, Optional, Tuple
+
+# Suppress CUDA deprecation warnings
+warnings.filterwarnings("ignore", category=FutureWarning, message=".*torch.backends.cuda.sdp_kernel.*")
+warnings.filterwarnings("ignore", category=FutureWarning, message=".*sdp_kernel.*")
+
+# Import ChatterboxTTS modules and ensure all config variables are available
+# First set defaults, then try to import from config
+DEFAULT_EXAGGERATION = 0.4
+DEFAULT_CFG_WEIGHT = 0.5
+DEFAULT_TEMPERATURE = 0.9
+TTS_PARAM_MIN_EXAGGERATION = 0.0
+TTS_PARAM_MAX_EXAGGERATION = 2.0
+TTS_PARAM_MIN_CFG_WEIGHT = 0.0
+TTS_PARAM_MAX_CFG_WEIGHT = 1.0
+TTS_PARAM_MIN_TEMPERATURE = 0.0
+TTS_PARAM_MAX_TEMPERATURE = 5.0
+ENABLE_REGENERATION_LOOP = True
+MAX_REGENERATION_ATTEMPTS = 3
+QUALITY_THRESHOLD = 0.7
+ENABLE_SENTIMENT_SMOOTHING = True
+SENTIMENT_SMOOTHING_WINDOW = 3
+SENTIMENT_SMOOTHING_METHOD = "rolling"
+ENABLE_MFCC_VALIDATION = False
+ENABLE_OUTPUT_VALIDATION = False
+SPECTRAL_ANOMALY_THRESHOLD = 0.8
+OUTPUT_VALIDATION_THRESHOLD = 0.85
+
+# Try to import config and override defaults if available
+try:
+    from config.config import *
+    CONFIG_AVAILABLE = True
+    print("✅ Config loaded successfully")
+except ImportError:
+    print("⚠️  Config not available - using defaults")
+    CONFIG_AVAILABLE = False
+
+# Import the actual conversion functions from GUI
+try:
+    # We need to import the actual conversion logic
+    import importlib.util
+    gui_spec = importlib.util.spec_from_file_location("chatterbox_gui", "chatterbox_gui.py")
+    gui_module = importlib.util.module_from_spec(gui_spec)
+    # We'll access the GUI's conversion methods
+    GUI_AVAILABLE = True
+except Exception as e:
+    print(f"⚠️  GUI module not available: {e}")
+    GUI_AVAILABLE = False
+
+# Global state for conversion with enhanced stats
+conversion_state = {
+    'running': False,
+    'progress': 0,
+    'status': 'Ready',
+    'thread': None,
+    'realtime_factor': '--',
+    'vram_usage': '-- GB',
+    'current_chunk': '--',
+    'eta': '--',
+    'elapsed': '--'
+}
+
+def parse_progress_stats(output_line):
+    """Parse progress statistics from TTS engine output"""
+    # Look for progress pattern: "🌀 Chunk 2/13 | ⏱ Elapsed: 0:01:31 | ETA: 0:09:54 | Remaining: 0:08:23 | Realtime: 0.11x | VRAM: 3.3GB"
+    progress_pattern = r'🌀 Chunk (\d+)/(\d+).*?Realtime: ([\d.]+)x.*?VRAM: ([\d.]+)GB'
+    match = re.search(progress_pattern, output_line)
+
+    if match:
+        current_chunk = int(match.group(1))
+        total_chunks = int(match.group(2))
+        realtime_factor = f"{match.group(3)}x"
+        vram_usage = f"{match.group(4)} GB"
+
+        # Update global state
+        conversion_state['current_chunk'] = f"{current_chunk}/{total_chunks}"
+        conversion_state['realtime_factor'] = realtime_factor
+        conversion_state['vram_usage'] = vram_usage
+        conversion_state['progress'] = int((current_chunk / total_chunks) * 100) if total_chunks > 0 else 0
+
+        print(f"📊 Stats Updated: Chunk {current_chunk}/{total_chunks}, {realtime_factor}, {vram_usage}")
+        return True
+    else:
+        # Try alternative patterns in case the format is different
+        alt_pattern = r'Chunk (\d+)/(\d+).*?Realtime: ([\d.]+)x.*?VRAM: ([\d.]+)GB'
+        alt_match = re.search(alt_pattern, output_line)
+        if alt_match:
+            current_chunk = int(alt_match.group(1))
+            total_chunks = int(alt_match.group(2))
+            realtime_factor = f"{alt_match.group(3)}x"
+            vram_usage = f"{alt_match.group(4)} GB"
+
+            conversion_state['current_chunk'] = f"{current_chunk}/{total_chunks}"
+            conversion_state['realtime_factor'] = realtime_factor
+            conversion_state['vram_usage'] = vram_usage
+            conversion_state['progress'] = int((current_chunk / total_chunks) * 100) if total_chunks > 0 else 0
+
+            print(f"📊 Stats Updated: Chunk {current_chunk}/{total_chunks}, {realtime_factor}, {vram_usage}")
+            return True
+
+    return False
+
+def get_progress_stats():
+    """Get current progress statistics for UI update"""
+    return (
+        conversion_state['realtime_factor'],
+        conversion_state['vram_usage'],
+        conversion_state['current_chunk'],
+        conversion_state['progress']
+    )
+
+def get_book_folders():
+    """Get available book folders from Text_Input directory"""
+    text_input_dir = Path("Text_Input")
+    if not text_input_dir.exists():
+        return []
+
+    folders = []
+    for item in text_input_dir.iterdir():
+        if item.is_dir():
+            folders.append(item.name)  # Show only folder name, not full path
+
+    return sorted(folders)
+
+def get_text_files_in_folder(folder_name):
+    """Get text files in selected book folder"""
+    if not folder_name:
+        return []
+
+    # Build full path from folder name
+    folder = Path("Text_Input") / folder_name
+    if not folder.exists():
+        return []
+
+    text_files = []
+    for file in folder.glob("*.txt"):
+        text_files.append(file.name)
+
+    return sorted(text_files)
+
+def get_voice_samples():
+    """Get available voice samples from Voice_Samples directory"""
+    voice_dir = Path("Voice_Samples")
+    if not voice_dir.exists():
+        return []
+
+    voices = []
+    for file in voice_dir.glob("*.wav"):
+        voices.append(file.name)  # Show only filename, not full path
+
+    return sorted(voices)
+
+def find_generated_audiobook(book_folder_path, voice_sample_path):
+    """Find the generated audiobook files"""
+    try:
+        book_folder = Path(book_folder_path)
+        voice_file = Path(voice_sample_path)
+        voice_name = voice_file.stem
+
+        # Look in Output/ directory first (final audiobooks)
+        output_dir = Path("Output")
+        if output_dir.exists():
+            # Look for M4B files with voice name
+            for m4b_file in output_dir.glob(f"*[{voice_name}]*.m4b"):
+                if m4b_file.exists():
+                    return str(m4b_file), "M4B audiobook"
+
+            # Look for WAV files with voice name
+            for wav_file in output_dir.glob(f"*[{voice_name}]*.wav"):
+                if wav_file.exists():
+                    return str(wav_file), "WAV audiobook"
+
+        # Look in Audiobook/ directory (processing output)
+        audiobook_dir = Path("Audiobook") / book_folder.name
+        if audiobook_dir.exists():
+            # Look for M4B files
+            for m4b_file in audiobook_dir.glob(f"*[{voice_name}]*.m4b"):
+                if m4b_file.exists():
+                    return str(m4b_file), "M4B audiobook"
+
+            # Look for WAV files
+            for wav_file in audiobook_dir.glob(f"*[{voice_name}]*.wav"):
+                if wav_file.exists():
+                    return str(wav_file), "WAV audiobook"
+
+            # Look for combined files
+            for combined_file in audiobook_dir.glob("*_combined.*"):
+                if combined_file.suffix in ['.wav', '.m4b', '.mp3']:
+                    return str(combined_file), f"{combined_file.suffix.upper()[1:]} combined audiobook"
+
+        return None, "No audiobook found"
+
+    except Exception as e:
+        print(f"Error finding audiobook: {e}")
+        return None, f"Error: {str(e)}"
+
+def run_book_conversion(book_path, text_file_path, voice_path, tts_params, quality_params, config_params):
+    """Run the actual book conversion - Direct call to TTS engine with progress monitoring"""
+    try:
+        # Import the real TTS engine function directly (avoid interface.py)
+        from modules.tts_engine import process_book_folder
+
+        # Extract enable_asr from tts_params (matching GUI exactly)
+        enable_asr = tts_params.get('enable_asr', False)
+
+        print(f"🚀 Starting book conversion with GUI parameters")
+        print(f"📖 Book: {book_path}")
+        print(f"📄 Text file: {text_file_path}")
+        print(f"🎤 Voice: {voice_path}")
+        print(f"🎛️ TTS Params: {tts_params}")
+        print(f"🔬 Quality Params: {quality_params}")
+        print(f"⚙️ Config Params: {config_params}")
+
+        # Set up progress callback function
+        def progress_callback(current_chunk, total_chunks, realtime_factor, vram_usage):
+            """Callback function to update progress from TTS engine"""
+            conversion_state['current_chunk'] = f"{current_chunk}/{total_chunks}"
+            conversion_state['realtime_factor'] = f"{realtime_factor}x"
+            conversion_state['vram_usage'] = f"{vram_usage} GB"
+            conversion_state['progress'] = int((current_chunk / total_chunks) * 100) if total_chunks > 0 else 0
+            print(f"📊 Progress: {current_chunk}/{total_chunks} ({conversion_state['progress']}%) - {realtime_factor}x - {vram_usage}GB")
+
+        # Add progress callback to config params
+        config_params['progress_callback'] = progress_callback
+
+        # Convert string paths to Path objects (required by TTS engine)
+        book_dir_path = Path(book_path)
+        voice_path_obj = Path(voice_path)
+
+        # Auto-detect device with fallback to CPU
+        import torch
+        if torch.cuda.is_available():
+            device = "cuda"
+            print("✅ Using CUDA GPU for processing")
+        else:
+            device = "cpu"
+            print("💻 Using CPU for processing (no GPU available)")
+        
+        # Direct call to TTS engine (function only accepts: book_dir, voice_path, tts_params, device, skip_cleanup)
+        result = process_book_folder(
+            book_dir=book_dir_path,
+            voice_path=voice_path_obj,
+            tts_params=tts_params,
+            device=device,
+            skip_cleanup=False
+        )
+
+        print(f"✅ Conversion completed successfully")
+        return {'success': True, 'result': result}
+
+    except Exception as e:
+        print(f"❌ Conversion failed: {e}")
+        import traceback
+        traceback.print_exc()
+        return {'success': False, 'error': str(e)}
+
+def regenerate_m4b_file(selected_m4b, playback_speed):
+    """Regenerate M4B file with new playback speed"""
+    if not selected_m4b:
+        return "❌ Please select an M4B file first", None
+
+    try:
+        print(f"🔄 Regenerating M4B: {selected_m4b} at {playback_speed}x speed")
+
+        # Import M4B regeneration tools
+        from tools.combine_only import apply_playback_speed_to_m4b
+
+        # Find the M4B file path
+        audiobook_root = Path("Audiobook")
+        m4b_path = None
+
+        for book_dir in audiobook_root.iterdir():
+            if book_dir.is_dir():
+                for m4b_file in book_dir.glob("*.m4b"):
+                    if m4b_file.name == selected_m4b:
+                        m4b_path = m4b_file
+                        break
+                if m4b_path:
+                    break
+
+        if not m4b_path:
+            return "❌ M4B file not found", None
+
+        # Create new filename with speed suffix
+        speed_suffix = f"_speed{playback_speed}x".replace(".", "p")
+        new_name = m4b_path.stem + speed_suffix + ".m4b"
+        output_path = m4b_path.parent / new_name
+
+        # Apply speed change
+        success = apply_playback_speed_to_m4b(str(m4b_path), str(output_path), playback_speed)
+
+        if success:
+            return f"✅ Regenerated M4B at {playback_speed}x speed: {new_name}", str(output_path)
+        else:
+            return "❌ Failed to regenerate M4B", None
+
+    except Exception as e:
+        print(f"❌ M4B regeneration failed: {e}")
+        return f"❌ Error: {str(e)}", None
+
+def create_convert_book_tab():
+    """Create Tab 1: Convert Book with all GUI functionality"""
+
+    with gr.Column():
+        gr.Markdown("# 🚀 Convert Book")
+        gr.Markdown("*Main TTS conversion functionality - matches GUI Tab 1*")
+
+        # Main Content Layout
+        with gr.Row():
+            # Left Column - File Uploads
+            with gr.Column(scale=2):
+                gr.Markdown("### 📚 Book Selection")
+
+                # Book text file upload only
+                text_file_upload = gr.File(
+                    label="📚 Upload Book Text File",
+                    file_types=[".txt"],
+                    file_count="single",
+                    interactive=True
+                )
+
+                gr.Markdown("### 🎤 Voice Selection")
+
+                # Single voice upload with integrated playback
+                voice_file_upload = gr.File(
+                    label="🎤 Upload Voice Sample",
+                    file_types=[".wav", ".mp3", ".m4a"],
+                    file_count="single",
+                    interactive=True
+                )
+
+                # Voice sample player (becomes active after upload)
+                voice_audio = gr.Audio(
+                    label="Voice Sample Preview",
+                    interactive=False,
+                    show_download_button=False,
+                    visible=False
+                )
+
+            # Right Column - All Settings
+            with gr.Column(scale=1):
+                gr.Markdown("### ⚙️ Quick Settings")
+
+                # VADER and ASR
+                vader_enabled = gr.Checkbox(
+                    label="Use VADER sentiment analysis",
+                    value=True,
+                    info="Adjust TTS params per chunk based on emotion"
+                )
+
+                # ASR System with intelligent model selection
+                with gr.Row():
+                    asr_enabled = gr.Checkbox(
+                        label="🎤 Enable ASR validation",
+                        value=False,
+                        info="Smart quality control with automatic model selection"
+                    )
+
+                # ASR Configuration (initially hidden)
+                with gr.Column(visible=False) as asr_config_group:
+                    gr.Markdown("#### 🔍 ASR Configuration")
+
+                    # System analysis display
+                    system_analysis = gr.Textbox(
+                        label="System Analysis",
+                        value="Click 'Analyze System' to detect capabilities",
+                        lines=3,
+                        interactive=False
+                    )
+
+                    analyze_system_btn = gr.Button(
+                        "🔍 Analyze System",
+                        size="sm",
+                        variant="secondary"
+                    )
+
+                    # ASR Level Selection
+                    asr_level = gr.Radio(
+                        label="ASR Quality Level",
+                        choices=[
+                            ("🟢 SAFE - Fast processing, basic accuracy", "safe"),
+                            ("🟡 MODERATE - Balanced speed/accuracy (recommended)", "moderate"),
+                            ("🔴 INSANE - Best accuracy, may stress system", "insane")
+                        ],
+                        value="moderate",
+                        info="Automatically selects best models for your system"
+                    )
+
+                    # Selected models display
+                    selected_models = gr.Textbox(
+                        label="Selected ASR Models",
+                        value="Select level to see model configuration",
+                        lines=2,
+                        interactive=False
+                    )
+
+                # Batch processing
+                add_to_batch = gr.Checkbox(
+                    label="📦 Add to batch queue",
+                    value=False,
+                    info="Queue for batch processing"
+                )
+
+                gr.Markdown("### 🔄 Regeneration Settings")
+
+                regeneration_enabled = gr.Checkbox(
+                    label="Enable automatic chunk regeneration",
+                    value=ENABLE_REGENERATION_LOOP,
+                    info="Retry failed chunks automatically"
+                )
+
+                max_attempts = gr.Slider(
+                    label="Max Attempts",
+                    minimum=1, maximum=10, step=1,
+                    value=MAX_REGENERATION_ATTEMPTS
+                )
+
+                quality_threshold = gr.Slider(
+                    label="Quality Threshold",
+                    minimum=0.1, maximum=1.0, step=0.05,
+                    value=QUALITY_THRESHOLD
+                )
+
+                gr.Markdown("### 📊 Sentiment Smoothing")
+
+                sentiment_smoothing = gr.Checkbox(
+                    label="Enable sentiment smoothing",
+                    value=ENABLE_SENTIMENT_SMOOTHING,
+                    info="Smooth emotional transitions"
+                )
+
+                smoothing_window = gr.Slider(
+                    label="Window Size",
+                    minimum=1, maximum=10, step=1,
+                    value=SENTIMENT_SMOOTHING_WINDOW
+                )
+
+                smoothing_method = gr.Dropdown(
+                    label="Smoothing Method",
+                    choices=["rolling", "exp_decay"],
+                    value=SENTIMENT_SMOOTHING_METHOD
+                )
+
+                gr.Markdown("### 🔍 Advanced Detection")
+
+                mfcc_validation = gr.Checkbox(
+                    label="MFCC spectral analysis",
+                    value=ENABLE_MFCC_VALIDATION,
+                    info="Advanced audio quality detection"
+                )
+
+                output_validation = gr.Checkbox(
+                    label="Output validation",
+                    value=ENABLE_OUTPUT_VALIDATION,
+                    info="Quality control clearinghouse for enabled checks"
+                )
+
+                spectral_threshold = gr.Slider(
+                    label="Spectral Threshold",
+                    minimum=0.1, maximum=1.0, step=0.05,
+                    value=SPECTRAL_ANOMALY_THRESHOLD
+                )
+
+                output_threshold = gr.Slider(
+                    label="Output Threshold",
+                    minimum=0.1, maximum=1.0, step=0.05,
+                    value=OUTPUT_VALIDATION_THRESHOLD
+                )
+
+
+        # TTS Parameters
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### 🎛️ TTS Parameters")
+
+                exaggeration = gr.Slider(
+                    label="Exaggeration",
+                    minimum=TTS_PARAM_MIN_EXAGGERATION,
+                    maximum=TTS_PARAM_MAX_EXAGGERATION,
+                    step=0.1,
+                    value=DEFAULT_EXAGGERATION,
+                    info="Emotional intensity"
+                )
+
+                cfg_weight = gr.Slider(
+                    label="CFG Weight",
+                    minimum=TTS_PARAM_MIN_CFG_WEIGHT,
+                    maximum=TTS_PARAM_MAX_CFG_WEIGHT,
+                    step=0.1,
+                    value=DEFAULT_CFG_WEIGHT,
+                    info="Text faithfulness"
+                )
+
+                temperature = gr.Slider(
+                    label="Temperature",
+                    minimum=TTS_PARAM_MIN_TEMPERATURE,
+                    maximum=TTS_PARAM_MAX_TEMPERATURE,
+                    step=0.1,
+                    value=DEFAULT_TEMPERATURE,
+                    info="Creativity/randomness"
+                )
+
+            with gr.Column():
+                gr.Markdown("### ⚡ Advanced Sampling")
+
+                min_p = gr.Slider(
+                    label="Min-P",
+                    minimum=0.0, maximum=0.5, step=0.01,
+                    value=0.05,
+                    info="Minimum probability threshold"
+                )
+
+                top_p = gr.Slider(
+                    label="Top-P",
+                    minimum=0.5, maximum=1.0, step=0.1,
+                    value=1.0,
+                    info="Nucleus sampling"
+                )
+
+                repetition_penalty = gr.Slider(
+                    label="Repetition Penalty",
+                    minimum=1.0, maximum=3.0, step=0.1,
+                    value=2.0,
+                    info="Reduce repetition"
+                )
+
+                gr.Markdown("### ⚙️ Performance Settings")
+                
+                max_workers = gr.Number(
+                    label="Max Workers",
+                    minimum=1, maximum=8, step=1,
+                    value=2,
+                    info="⚠️ Only increase above 2 if CPU/GPU utilization < 70%"
+                )
+
+        # Action Buttons and Status
+        with gr.Row():
+            with gr.Column(scale=2):
+                convert_btn = gr.Button(
+                    "🚀 Start Conversion",
+                    variant="primary",
+                    size="lg",
+                    interactive=True
+                )
+
+                # Status Display
+                status_display = gr.Textbox(
+                    label="Status",
+                    value="⏸ Ready",
+                    interactive=False,
+                    lines=1
+                )
+
+                progress_display = gr.Number(
+                    label="Progress %",
+                    value=0,
+                    interactive=False,
+                    precision=0
+                )
+
+            with gr.Column(scale=1):
+                gr.Markdown("### 📊 Processing Stats")
+
+                realtime_factor = gr.Textbox(
+                    label="Realtime Factor",
+                    value="--",
+                    interactive=False
+                )
+
+                vram_usage = gr.Textbox(
+                    label="VRAM Usage",
+                    value="-- GB",
+                    interactive=False
+                )
+
+                current_chunk = gr.Textbox(
+                    label="Current Chunk",
+                    value="--",
+                    interactive=False
+                )
+
+        # Regenerate M4B Section (moved above audiobook player)
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### 🔄 Regenerate M4B")
+
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        m4b_file_selector = gr.Dropdown(
+                            label="Select M4B File to Regenerate",
+                            choices=[],
+                            value=None,
+                            interactive=True,
+                            info="Choose from generated audiobook files"
+                        )
+
+                    with gr.Column(scale=1):
+                        playback_speed = gr.Slider(
+                            label="Playback Speed",
+                            minimum=0.5,
+                            maximum=2.0,
+                            step=0.1,
+                            value=1.0,
+                            info="Speed adjustment for regeneration"
+                        )
+
+                regenerate_m4b_btn = gr.Button(
+                    "🔄 Regenerate M4B",
+                    variant="secondary",
+                    size="lg"
+                )
+
+        # Generated Audiobook Player (simplified, play-only)
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### 🎧 Generated Audiobook Player")
+
+                # Audiobook file selector dropdown
+                audiobook_selector = gr.Dropdown(
+                    label="Select Audiobook",
+                    choices=[],
+                    value=None,
+                    interactive=True,
+                    info="Choose from session audiobooks"
+                )
+
+                # Main audio player - play only, no upload
+                audio_player = gr.Audio(
+                    label="Audiobook Player",
+                    value=None,
+                    interactive=False,
+                    show_download_button=True,
+                    show_share_button=False,
+                    waveform_options=gr.WaveformOptions(
+                        show_controls=True,
+                        show_recording_waveform=False,
+                        skip_length=10
+                    )
+                )
+
+    # Event Handlers
+    def handle_voice_upload(voice_file):
+        """Handle voice file upload and show player"""
+        if voice_file is None:
+            return gr.update(value=None, visible=False)
+
+        # Show the voice player with uploaded file
+        return gr.update(value=voice_file, visible=True)
+
+    def get_session_audiobooks():
+        """Get list of M4B files from current session, sorted by creation time (newest first)"""
+        audiobooks = []
+
+        # Look in Audiobook directory for M4B files
+        audiobook_root = Path("Audiobook")
+        if audiobook_root.exists():
+            for book_dir in audiobook_root.iterdir():
+                if book_dir.is_dir():
+                    # Look for M4B files in book directory
+                    for m4b_file in book_dir.glob("*.m4b"):
+                        # Get creation time for sorting
+                        creation_time = m4b_file.stat().st_mtime
+                        audiobooks.append((str(m4b_file), m4b_file.name, creation_time))
+
+        # Also check Output directory
+        output_root = Path("Output")
+        if output_root.exists():
+            for m4b_file in output_root.glob("*.m4b"):
+                creation_time = m4b_file.stat().st_mtime
+                audiobooks.append((str(m4b_file), m4b_file.name, creation_time))
+
+        # Sort by creation time (newest first)
+        audiobooks.sort(key=lambda x: x[2], reverse=True)
+
+        # Return just path and name (drop creation time)
+        return [(ab[0], ab[1]) for ab in audiobooks]
+
+    def update_audiobook_dropdowns(latest_file=None):
+        """Update audiobook dropdowns - after conversion both show latest, after regeneration only playback updates"""
+        audiobooks = get_session_audiobooks()
+        choices = [ab[1] for ab in audiobooks]  # Just filenames for display
+
+        # Determine what to set as selected
+        if latest_file:
+            # Use specific file if provided
+            selected_file = latest_file
+        elif choices:
+            # Default to newest file (first in sorted list)
+            selected_file = choices[0]
+        else:
+            selected_file = None
+
+        return (
+            gr.update(choices=choices, value=selected_file),  # audiobook_selector (playback)
+            gr.update(choices=choices, value=selected_file)   # m4b_file_selector (regeneration source)
+        )
+
+    def update_audiobook_dropdowns_after_conversion():
+        """Update both dropdowns to show the newest generated file after conversion"""
+        return update_audiobook_dropdowns()
+
+    def update_playback_only(new_file_name):
+        """Update only the playback dropdown after regeneration"""
+        audiobooks = get_session_audiobooks()
+        choices = [ab[1] for ab in audiobooks]
+
+        return (
+            gr.update(choices=choices, value=new_file_name),  # audiobook_selector (playback) - new file
+            gr.update()  # m4b_file_selector (regeneration) - no change
+        )
+
+    def load_selected_audiobook(selected_audiobook):
+        """Load selected audiobook into player"""
+        if not selected_audiobook:
+            return None
+
+        # Find the full path for the selected audiobook
+        audiobooks = get_session_audiobooks()
+        for full_path, filename in audiobooks:
+            if filename == selected_audiobook:
+                return full_path
+
+        return None
+
+    def handle_asr_toggle(asr_enabled_val):
+        """Show/hide ASR configuration when ASR is toggled"""
+        return gr.update(visible=asr_enabled_val)
+
+    def analyze_system():
+        """Analyze system capabilities and return summary"""
+        try:
+            from modules.system_detector import get_system_profile, print_system_summary, categorize_system
+
+            profile = get_system_profile()
+            categories = categorize_system(profile)
+
+            summary = f"🖥️ System Profile:\n"
+            summary += f"VRAM: {profile['gpu']['total_mb']:,}MB total, {profile['available_vram_after_tts']:,}MB available after TTS ({categories['vram']} class)\n"
+            summary += f"RAM: {profile['ram']['total_mb']:,}MB total, {profile['ram']['available_mb']:,}MB available ({categories['ram']} class)\n"
+            summary += f"CPU: {profile['cpu_cores']} cores ({categories['cpu']} class)"
+
+            if not profile['has_gpu']:
+                summary += f"\n⚠️ No CUDA GPU detected - ASR will run on CPU only"
+
+            return summary
+
+        except Exception as e:
+            return f"❌ Error analyzing system: {str(e)}"
+
+    def update_asr_models(asr_level_val):
+        """Update ASR model display based on selected level"""
+        try:
+            from modules.system_detector import get_system_profile, recommend_asr_models
+
+            profile = get_system_profile()
+            recommendations = recommend_asr_models(profile)
+
+            if asr_level_val not in recommendations:
+                return "❌ Invalid ASR level selected"
+
+            config = recommendations[asr_level_val]
+            primary = config['primary']
+            fallback = config['fallback']
+
+            result = f"Primary: {primary['model']} on {primary['device'].upper()}\n"
+            result += f"Fallback: {fallback['model']} on {fallback['device'].upper()}"
+
+            if asr_level_val == 'insane':
+                result += f"\n⚠️ WARNING: INSANE mode may cause memory pressure"
+
+            return result
+
+        except Exception as e:
+            return f"❌ Error getting models: {str(e)}"
+
+    def start_conversion(text_file_upload, voice_file_upload,
+                        vader_val, asr_val, asr_level_val, add_to_batch_val,
+                        regen_enabled_val, max_attempts_val, quality_thresh_val,
+                        sentiment_smooth_val, smooth_window_val, smooth_method_val,
+                        mfcc_val, output_val, spectral_thresh_val, output_thresh_val,
+                        exag_val, cfg_val, temp_val, min_p_val, top_p_val, rep_penalty_val,
+                        max_workers_val):
+        """Start the actual book conversion - file upload version"""
+
+        # Validation
+        if not text_file_upload:
+            return "❌ Please upload a text file", 0, None, None
+        if not voice_file_upload:
+            return "❌ Please upload a voice sample", 0, None, None
+
+        # Check if already running
+        if conversion_state['running']:
+            return "⚠️ Conversion already in progress", conversion_state['progress'], None, None
+
+        try:
+            # Create temporary book structure from uploads
+            import tempfile
+            import shutil
+            from datetime import datetime
+
+            # Generate unique book name from text file
+            text_filename = Path(text_file_upload).name
+            book_name = text_filename.replace('.txt', '').replace(' ', '_')
+            timestamp = datetime.now().strftime("%H%M%S")
+            unique_book_name = f"{book_name}_{timestamp}"
+
+            # Create directory structure
+            text_input_dir = Path("Text_Input")
+            text_input_dir.mkdir(exist_ok=True)
+
+            book_dir = text_input_dir / unique_book_name
+            book_dir.mkdir(exist_ok=True)
+
+            # Copy uploaded files to expected locations
+            text_dest = book_dir / f"{unique_book_name}.txt"
+            shutil.copy2(text_file_upload, text_dest)
+
+            voice_samples_dir = Path("Voice_Samples")
+            voice_samples_dir.mkdir(exist_ok=True)
+
+            voice_filename = Path(voice_file_upload).name
+            voice_dest = voice_samples_dir / voice_filename
+            shutil.copy2(voice_file_upload, voice_dest)
+
+            print(f"📁 Created book structure: {book_dir}")
+            print(f"📄 Text file: {text_dest}")
+            print(f"🎤 Voice file: {voice_dest}")
+
+        except Exception as e:
+            return f"❌ Error setting up files: {e}", 0, None, None
+
+        # Build ASR configuration first
+        asr_config = {'enabled': False}
+        if asr_val:
+            try:
+                from modules.system_detector import get_system_profile, recommend_asr_models
+                profile = get_system_profile()
+                recommendations = recommend_asr_models(profile)
+
+                if asr_level_val in recommendations:
+                    selected_config = recommendations[asr_level_val]
+                    primary = selected_config['primary']
+                    fallback = selected_config['fallback']
+
+                    asr_config = {
+                        'enabled': True,
+                        'level': asr_level_val,
+                        'primary_model': primary['model'],
+                        'primary_device': primary['device'],
+                        'fallback_model': fallback['model'],
+                        'fallback_device': fallback['device']
+                    }
+            except Exception as e:
+                print(f"⚠️ Error configuring ASR: {e}")
+                asr_config = {'enabled': False}
+
+        # Prepare parameters (matching GUI structure exactly)
+        tts_params = {
+            'exaggeration': exag_val,
+            'cfg_weight': cfg_val,
+            'temperature': temp_val,
+            'min_p': min_p_val,
+            'top_p': top_p_val,
+            'repetition_penalty': rep_penalty_val,
+            'enable_asr': asr_config.get('enabled', False),  # Match GUI pattern
+            'max_workers': int(max_workers_val)  # User-defined worker count
+        }
+
+        quality_params = {
+            'regeneration_enabled': regen_enabled_val,
+            'max_attempts': max_attempts_val,
+            'quality_threshold': quality_thresh_val,
+            'sentiment_smoothing': sentiment_smooth_val,
+            'smoothing_window': smooth_window_val,
+            'smoothing_method': smooth_method_val,
+            'mfcc_validation': mfcc_val,
+            'output_validation': output_val,
+            'spectral_threshold': spectral_thresh_val,
+            'output_threshold': output_thresh_val
+        }
+
+        config_params = {
+            'vader_enabled': vader_val,
+            'asr_enabled': asr_val,
+            'asr_config': asr_config,
+            'add_to_batch': add_to_batch_val
+        }
+
+        # Set conversion state
+        conversion_state['running'] = True
+        conversion_state['progress'] = 0
+        conversion_state['status'] = 'Starting conversion...'
+        conversion_state['current_book'] = book_dir.name  # Track current book
+
+        try:
+            # Run conversion using the modular backend in a separate thread
+            import threading
+
+            def run_conversion_thread():
+                try:
+                    result = run_book_conversion(
+                        str(book_dir), str(text_dest), str(voice_dest),
+                        tts_params, quality_params, config_params
+                    )
+
+                    if result['success']:
+                        conversion_state['status'] = '🎉 CONVERSION COMPLETE! M4B audiobook ready for playback.'
+                        conversion_state['progress'] = 100
+                        conversion_state['auto_refresh_needed'] = True  # Flag for auto-refresh
+                    else:
+                        conversion_state['status'] = f"❌ Conversion failed: {result.get('error', 'Unknown error')}"
+                        conversion_state['progress'] = 0
+
+                except Exception as e:
+                    conversion_state['status'] = f"❌ Error: {str(e)}"
+                    conversion_state['progress'] = 0
+                finally:
+                    conversion_state['running'] = False
+
+            # Start conversion thread
+            thread = threading.Thread(target=run_conversion_thread)
+            thread.start()
+
+            # Return immediate response - user will need to refresh to see final results
+            return (
+                "🚀 Conversion started in background...",
+                5,  # Initial progress
+                None,
+                gr.update(),
+                gr.update()
+            )
+
+        except Exception as e:
+            conversion_state['status'] = f"❌ Error: {str(e)}"
+            return conversion_state['status'], 0, None, gr.update(), gr.update()
+        finally:
+            conversion_state['running'] = False
+
+
+    # Connect event handlers
+
+    # ASR event handlers
+    asr_enabled.change(
+        handle_asr_toggle,
+        inputs=[asr_enabled],
+        outputs=[asr_config_group]
+    )
+
+    analyze_system_btn.click(
+        analyze_system,
+        inputs=[],
+        outputs=[system_analysis]
+    )
+
+    asr_level.change(
+        update_asr_models,
+        inputs=[asr_level],
+        outputs=[selected_models]
+    )
+
+    # Voice upload handler
+    voice_file_upload.change(
+        handle_voice_upload,
+        inputs=[voice_file_upload],
+        outputs=[voice_audio]
+    )
+
+    # Main conversion handler
+    convert_btn.click(
+        start_conversion,
+        inputs=[
+            text_file_upload, voice_file_upload,
+            vader_enabled, asr_enabled, asr_level, add_to_batch,
+            regeneration_enabled, max_attempts, quality_threshold,
+            sentiment_smoothing, smoothing_window, smoothing_method,
+            mfcc_validation, output_validation, spectral_threshold, output_threshold,
+            exaggeration, cfg_weight, temperature, min_p, top_p, repetition_penalty,
+            max_workers
+        ],
+        outputs=[status_display, progress_display, audio_player, audiobook_selector, m4b_file_selector]
+    )
+
+    # Audiobook selector handler
+    audiobook_selector.change(
+        load_selected_audiobook,
+        inputs=[audiobook_selector],
+        outputs=[audio_player]
+    )
+
+    # M4B regeneration handler
+    def handle_m4b_regeneration(selected_m4b, speed):
+        """Handle M4B regeneration and update player"""
+        status_msg, new_m4b_path = regenerate_m4b_file(selected_m4b, speed)
+
+        if new_m4b_path:
+            # Load the new M4B in the player
+            new_file_name = Path(new_m4b_path).name
+            new_audio = load_selected_audiobook(new_file_name)
+            # Update only playback dropdown, keep regeneration dropdown on source file
+            audiobook_choices, m4b_choices = update_playback_only(new_file_name)
+            return status_msg, new_audio, audiobook_choices, m4b_choices
+        else:
+            return status_msg, None, gr.update(), gr.update()
+
+    regenerate_m4b_btn.click(
+        handle_m4b_regeneration,
+        inputs=[m4b_file_selector, playback_speed],
+        outputs=[status_display, audio_player, audiobook_selector, m4b_file_selector]
+    )
+
+    # Progress monitoring with file-based approach
+    def get_current_stats():
+        """Get current progress statistics by monitoring output files"""
+        try:
+            if conversion_state['running']:
+                # Look for generated audio chunks to estimate progress
+                book_name = conversion_state.get('current_book', 'unknown')
+                audiobook_root = Path("Audiobook") / book_name / "TTS" / "audio_chunks"
+
+                if audiobook_root.exists():
+                    chunk_files = list(audiobook_root.glob("chunk_*.wav"))
+                    current_chunks = len(chunk_files)
+
+                    # Try to estimate total from JSON if available
+                    json_path = Path("Text_Input") / f"{book_name}_chunks.json"
+                    total_chunks = 0
+                    if json_path.exists():
+                        import json
+                        with open(json_path, 'r') as f:
+                            data = json.load(f)
+                            total_chunks = len(data)
+
+                    if total_chunks > 0:
+                        progress = int((current_chunks / total_chunks) * 100)
+                        conversion_state['progress'] = progress
+                        conversion_state['current_chunk'] = f"{current_chunks}/{total_chunks}"
+
+                        return (
+                            conversion_state.get('realtime_factor', '--'),
+                            conversion_state.get('vram_usage', '-- GB'),
+                            f"{current_chunks}/{total_chunks}",
+                            progress
+                        )
+
+            return (
+                conversion_state.get('realtime_factor', '--'),
+                conversion_state.get('vram_usage', '-- GB'),
+                conversion_state.get('current_chunk', '--'),
+                conversion_state.get('progress', 0)
+            )
+        except Exception as e:
+            print(f"Error getting stats: {e}")
+            return "--", "-- GB", "--", conversion_state.get('progress', 0)
+
+    def auto_check_completion():
+        """Automatically check for completion and refresh interface"""
+        # First get current stats
+        stats = get_current_stats()
+        
+        # Check if conversion just completed and needs auto-refresh
+        if (not conversion_state['running'] and 
+            conversion_state['progress'] == 100 and 
+            conversion_state.get('auto_refresh_needed', False)):
+            
+            # Clear the auto-refresh flag
+            conversion_state['auto_refresh_needed'] = False
+            print("🎉 Auto-detected completion! Refreshing interface...")
+            
+            # Get completion results
+            status, progress, audio, audiobook_choices, m4b_choices = get_status_and_results()
+            
+            # Return combined stats + completion results
+            return (
+                stats[0],  # realtime_factor
+                stats[1],  # vram_usage  
+                stats[2],  # current_chunk
+                100,       # progress (completed)
+                status,    # completion status
+                audio,     # audio player
+                audiobook_choices,  # audiobook dropdown
+                m4b_choices        # m4b dropdown
+            )
+        else:
+            # Return stats + current status (no completion)
+            return (
+                stats[0],  # realtime_factor
+                stats[1],  # vram_usage
+                stats[2],  # current_chunk  
+                stats[3],  # progress
+                conversion_state.get('status', '⏸ Ready'),  # current status
+                gr.update(),  # no audio update
+                gr.update(),  # no audiobook update
+                gr.update()   # no m4b update
+            )
+
+    def get_status_and_results():
+        """Get conversion status and results after completion"""
+        if not conversion_state['running'] and conversion_state['progress'] == 100:
+            # Conversion completed, update dropdowns
+            audiobook_choices, m4b_choices = update_audiobook_dropdowns_after_conversion()
+            latest_audiobook = None
+            if audiobook_choices['choices']:
+                latest_audiobook = load_selected_audiobook(audiobook_choices['choices'][0])
+
+            return (
+                conversion_state['status'],
+                conversion_state['progress'],
+                latest_audiobook,
+                audiobook_choices,
+                m4b_choices
+            )
+        else:
+            return (
+                conversion_state['status'],
+                conversion_state['progress'],
+                None,
+                gr.update(),
+                gr.update()
+            )
+
+    # Create refresh buttons
+    with gr.Row():
+        refresh_stats_btn = gr.Button("🔄 Refresh Stats", size="sm", variant="secondary")
+        check_completion_btn = gr.Button("📋 Check Completion", size="sm", variant="secondary")
+    
+    # Auto-refresh timer (checks every 5 seconds during conversion)
+    auto_timer = gr.Timer(5.0)  # 5 second interval
+
+    refresh_stats_btn.click(
+        auto_check_completion,
+        outputs=[realtime_factor, vram_usage, current_chunk, progress_display, status_display, audio_player, audiobook_selector, m4b_file_selector]
+    )
+
+    check_completion_btn.click(
+        get_status_and_results,
+        outputs=[status_display, progress_display, audio_player, audiobook_selector, m4b_file_selector]
+    )
+    
+    # Auto-timer for progress monitoring and completion detection
+    auto_timer.tick(
+        auto_check_completion,
+        outputs=[realtime_factor, vram_usage, current_chunk, progress_display, status_display, audio_player, audiobook_selector, m4b_file_selector]
+    )
+
+    return {
+        'convert_button': convert_btn,
+        'status_display': status_display,
+        'progress': progress_display
+    }
+
+if __name__ == "__main__":
+    # Test the tab
+    with gr.Blocks() as demo:
+        create_convert_book_tab()
+
+    demo.launch()

HF_Deploy/modules/asr_manager.py

@@ -0,0 +1,233 @@
+"""
+ASR Manager Module
+Centralized ASR model loading with adaptive GPU/CPU fallback and real-time VRAM monitoring
+"""
+
+import torch
+import logging
+from pathlib import Path
+from config.config import DEFAULT_ASR_MODEL, ASR_MODEL_VRAM_MB, ASR_MODEL_RAM_MB
+
+def get_real_time_vram_status():
+    """Get current GPU memory usage in real-time"""
+    try:
+        if torch.cuda.is_available():
+            gpu_count = torch.cuda.device_count()
+            if gpu_count > 0:
+                # Use first GPU
+                total_vram = torch.cuda.get_device_properties(0).total_memory
+                allocated_vram = torch.cuda.memory_allocated(0)
+                reserved_vram = torch.cuda.memory_reserved(0)
+                available_vram = total_vram - allocated_vram
+                
+                return {
+                    'total_mb': total_vram // 1024 // 1024,
+                    'allocated_mb': allocated_vram // 1024 // 1024,
+                    'reserved_mb': reserved_vram // 1024 // 1024,
+                    'available_mb': available_vram // 1024 // 1024,
+                    'has_gpu': True
+                }
+    except Exception as e:
+        logging.warning(f"Failed to get real-time VRAM status: {e}")
+    
+    return {
+        'total_mb': 0,
+        'allocated_mb': 0, 
+        'reserved_mb': 0,
+        'available_mb': 0,
+        'has_gpu': False
+    }
+
+def calculate_available_vram_for_asr(safety_buffer_mb=500):
+    """Calculate VRAM available for ASR with safety buffer"""
+    vram_status = get_real_time_vram_status()
+    
+    if not vram_status['has_gpu']:
+        return 0
+    
+    # Available VRAM minus safety buffer for stability
+    available_with_buffer = max(0, vram_status['available_mb'] - safety_buffer_mb)
+    
+    return available_with_buffer
+
+def can_model_fit_gpu(model_name, available_vram_mb):
+    """Check if a specific ASR model can fit in available VRAM"""
+    required_vram = ASR_MODEL_VRAM_MB.get(model_name, 0)
+    return available_vram_mb >= required_vram
+
+def try_load_model_with_fallback(model_name, primary_device, fallback_device="cpu"):
+    """Try to load model on primary device, fallback to secondary if it fails"""
+    import whisper
+    
+    # Convert device names for whisper compatibility
+    def convert_device_name(device):
+        if device.lower() == "gpu":
+            return "cuda"
+        return device.lower()
+    
+    primary_device_whisper = convert_device_name(primary_device)
+    fallback_device_whisper = convert_device_name(fallback_device)
+    
+    try:
+        print(f"🎯 Attempting to load {model_name} on {primary_device.upper()}")
+        model = whisper.load_model(model_name, device=primary_device_whisper)
+        print(f"✅ Successfully loaded {model_name} on {primary_device.upper()}")
+        return model, primary_device
+        
+    except Exception as e:
+        print(f"⚠️ {model_name} failed on {primary_device} ({str(e)[:50]}...)")
+        
+        if fallback_device_whisper != primary_device_whisper:
+            try:
+                print(f"🔄 Trying {model_name} on {fallback_device.upper()}")
+                model = whisper.load_model(model_name, device=fallback_device_whisper)
+                print(f"✅ Successfully loaded {model_name} on {fallback_device.upper()}")
+                return model, fallback_device
+                
+            except Exception as fallback_e:
+                print(f"❌ {model_name} also failed on {fallback_device} ({str(fallback_e)[:50]}...)")
+        
+        # Both failed
+        raise Exception(f"Model {model_name} failed on both {primary_device} and {fallback_device}")
+
+def load_asr_model_adaptive(asr_config=None):
+    """
+    Adaptive ASR model loading with real-time VRAM checking and intelligent fallback
+    
+    Args:
+        asr_config: ASR configuration dict from interfaces (None for GUI fallback)
+        
+    Returns:
+        tuple: (asr_model, actual_device_used) or (None, None) if all loading fails
+    """
+    print(f"🔍 Starting adaptive ASR model loading...")
+    
+    # Get current VRAM status
+    vram_status = get_real_time_vram_status()
+    available_vram = calculate_available_vram_for_asr()
+    
+    print(f"🖥️ Real-time VRAM status:")
+    print(f"   Total: {vram_status['total_mb']:,}MB")
+    print(f"   Allocated: {vram_status['allocated_mb']:,}MB") 
+    print(f"   Available for ASR: {available_vram:,}MB (with 500MB safety buffer)")
+    
+    # Determine what models to try based on config
+    if asr_config and asr_config.get('enabled') and 'primary_model' in asr_config:
+        # Intelligent selection from CLI/Gradio
+        primary_model = asr_config['primary_model']
+        primary_device = asr_config['primary_device']
+        fallback_model = asr_config['fallback_model'] 
+        fallback_device = asr_config['fallback_device']
+        
+        print(f"🧠 Using intelligent ASR config:")
+        print(f"   Primary: {primary_model} on {primary_device.upper()}")
+        print(f"   Fallback: {fallback_model} on {fallback_device.upper()}")
+        
+        # Real-time VRAM check for primary model
+        if primary_device.lower() == 'gpu':
+            if not vram_status['has_gpu']:
+                print(f"⚠️ No GPU available, forcing CPU mode")
+                primary_device = 'cpu'
+            elif not can_model_fit_gpu(primary_model, available_vram):
+                required = ASR_MODEL_VRAM_MB.get(primary_model, 0)
+                print(f"⚠️ Insufficient VRAM for {primary_model} (need {required}MB, have {available_vram}MB)")
+                print(f"🔄 Switching primary to CPU")
+                primary_device = 'cpu'
+        
+        # Try primary model
+        try:
+            return try_load_model_with_fallback(primary_model, primary_device, primary_device)
+        except:
+            # Primary failed, try fallback model
+            print(f"🔄 Primary model failed, trying fallback configuration...")
+            
+            # Real-time VRAM check for fallback model  
+            if fallback_device.lower() == 'gpu':
+                if not vram_status['has_gpu']:
+                    print(f"⚠️ No GPU available for fallback, using CPU")
+                    fallback_device = 'cpu'
+                elif not can_model_fit_gpu(fallback_model, available_vram):
+                    required = ASR_MODEL_VRAM_MB.get(fallback_model, 0)
+                    print(f"⚠️ Insufficient VRAM for fallback {fallback_model} (need {required}MB, have {available_vram}MB)")
+                    fallback_device = 'cpu'
+            
+            try:
+                return try_load_model_with_fallback(fallback_model, fallback_device, 'cpu')
+            except:
+                print(f"❌ Both configured models failed!")
+    
+    else:
+        # Fallback mode for GUI or missing config
+        print(f"🔧 Using fallback mode: {DEFAULT_ASR_MODEL}")
+    
+    # Last resort: try default model with adaptive device selection
+    print(f"🆘 Last resort: trying {DEFAULT_ASR_MODEL} with adaptive device selection")
+    
+    # Choose device based on real-time VRAM availability
+    if vram_status['has_gpu'] and can_model_fit_gpu(DEFAULT_ASR_MODEL, available_vram):
+        device = 'cuda'  # Use cuda directly for whisper
+        device_display = 'GPU'
+        print(f"✅ Using GPU for {DEFAULT_ASR_MODEL}")
+    else:
+        device = 'cpu'  
+        device_display = 'CPU'
+        print(f"🔄 Using CPU for {DEFAULT_ASR_MODEL}")
+    
+    try:
+        import whisper
+        model = whisper.load_model(DEFAULT_ASR_MODEL, device=device)
+        print(f"✅ Successfully loaded {DEFAULT_ASR_MODEL} on {device_display}")
+        return model, device_display.lower()
+    except Exception as e:
+        print(f"❌ Critical failure: Could not load {DEFAULT_ASR_MODEL} on {device}: {e}")
+        
+        # Ultimate fallback to CPU if GPU failed
+        if device == 'cuda':
+            try:
+                print(f"🆘 Ultimate fallback: {DEFAULT_ASR_MODEL} on CPU")
+                model = whisper.load_model(DEFAULT_ASR_MODEL, device='cpu')
+                print(f"✅ Successfully loaded {DEFAULT_ASR_MODEL} on CPU")
+                return model, 'cpu'
+            except Exception as cpu_e:
+                print(f"💀 Complete failure: {cpu_e}")
+        
+        return None, None
+
+def cleanup_asr_model(asr_model):
+    """Clean up ASR model to free memory"""
+    if asr_model is not None:
+        try:
+            del asr_model
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            print(f"🧹 ASR model cleaned up")
+        except Exception as e:
+            logging.warning(f"Failed to cleanup ASR model: {e}")
+
+def get_asr_memory_info():
+    """Get memory information for ASR debugging"""
+    vram_status = get_real_time_vram_status()
+    available_vram = calculate_available_vram_for_asr()
+    
+    info = {
+        'vram_total_mb': vram_status['total_mb'],
+        'vram_allocated_mb': vram_status['allocated_mb'],
+        'vram_available_for_asr_mb': available_vram,
+        'has_gpu': vram_status['has_gpu']
+    }
+    
+    return info
+
+if __name__ == "__main__":
+    # Test the adaptive loading
+    print("Testing ASR Manager...")
+    info = get_asr_memory_info()
+    print(f"Memory info: {info}")
+    
+    # Test adaptive loading
+    model, device = load_asr_model_adaptive()
+    if model:
+        print(f"Test successful: Model loaded on {device}")
+        cleanup_asr_model(model)
+    else:
+        print("Test failed: No model loaded")

HF_Deploy/modules/audio_processor.py

@@ -0,0 +1,569 @@
+"""
+Audio Processing Module
+Handles audio validation, effects, cleanup, and quality control
+"""
+
+import numpy as np
+import soundfile as sf
+import logging
+import shutil
+import re
+import time
+from pathlib import Path
+from pydub import AudioSegment, silence
+from config.config import *
+
+# ============================================================================
+# AUDIO QUALITY DETECTION
+# ============================================================================
+
+def check_audio_health(wav_path):
+    """Enhanced audio health checking"""
+    data, samplerate = sf.read(str(wav_path))
+    if len(data.shape) > 1:
+        data = data[:, 0]  # mono only
+
+    clipping = np.mean(np.abs(data) > 0.98)
+    silence_ratio = np.mean(np.abs(data) < 1e-4)
+    rms = np.sqrt(np.mean(data**2))
+    mean_abs = np.mean(np.abs(data))
+    flatness = mean_abs / (rms + 1e-8)
+
+    return {
+        "clipping_ratio": round(clipping, 4),
+        "silence_ratio": round(silence_ratio, 4),
+        "flatness": round(flatness, 4),
+    }
+
+def detect_tts_hum_artifact(wav_path):
+    """
+    Detect low-frequency TTS confusion hum using configurable parameters
+    """
+    if not ENABLE_HUM_DETECTION:
+        return False, {}
+
+    data, sr = sf.read(str(wav_path))
+    if data.ndim > 1:
+        data = data[:, 0]  # Mono
+
+    # FFT analysis for frequency content
+    fft = np.fft.rfft(data)
+    freqs = np.fft.rfftfreq(len(data), 1/sr)
+
+    # Focus on hum frequency range (configurable at top of file)
+    hum_mask = (freqs >= HUM_FREQ_MIN) & (freqs <= HUM_FREQ_MAX)
+    hum_energy = np.sum(np.abs(fft[hum_mask]))
+    total_energy = np.sum(np.abs(fft))
+
+    # Check for sustained low-level amplitude (steady hum characteristic)
+    segment_size = sr // 4  # 250ms segments
+    segments = [data[i:i+segment_size] for i in range(0, len(data)-segment_size, segment_size)]
+
+    steady_segments = 0
+    for segment in segments:
+        rms = np.sqrt(np.mean(segment**2))
+        if HUM_AMPLITUDE_MIN < rms < HUM_AMPLITUDE_MAX:
+            steady_segments += 1
+
+    # Calculate hum indicators using configurable thresholds
+    hum_ratio = hum_energy / (total_energy + 1e-10)
+    steady_ratio = steady_segments / len(segments) if segments else 0
+
+    # Detection logic using configurable thresholds
+    has_hum = (hum_ratio > HUM_ENERGY_THRESHOLD) and (steady_ratio > HUM_STEADY_THRESHOLD)
+
+    if has_hum:
+        logging.info(f"🔍 TTS hum detected: {wav_path.name}")
+        logging.info(f"   Frequency range: {HUM_FREQ_MIN}-{HUM_FREQ_MAX}Hz")
+        logging.info(f"   Hum energy ratio: {hum_ratio:.3f} (threshold: {HUM_ENERGY_THRESHOLD})")
+        logging.info(f"   Steady segments: {steady_ratio:.3f} (threshold: {HUM_STEADY_THRESHOLD})")
+
+    return has_hum, {
+        "hum_ratio": hum_ratio,
+        "steady_ratio": steady_ratio,
+        "freq_range": f"{HUM_FREQ_MIN}-{HUM_FREQ_MAX}Hz"
+    }
+
+def smart_audio_validation(wav_path):
+    """Comprehensive audio validation with intelligent responses"""
+    # Standard health check
+    health = check_audio_health(wav_path)
+
+    # TTS hum detection (if enabled)
+    has_hum, hum_metrics = detect_tts_hum_artifact(wav_path)
+
+    # Decision matrix
+    if health["clipping_ratio"] > 0.05:
+        return handle_problematic_chunks(wav_path, "clipping", health)
+    elif health["flatness"] > 0.9:
+        return handle_problematic_chunks(wav_path, "corrupted", health)
+    elif has_hum:
+        return handle_problematic_chunks(wav_path, "tts_hum", hum_metrics)
+    else:
+        return wav_path  # Passed all checks
+
+def has_mid_energy_drop(wav_tensor, sr, window_ms=250, threshold_ratio=None):
+    """Detect mid-chunk energy drops"""
+    wav = wav_tensor.squeeze().numpy()
+    win_samples = int(sr * window_ms / 1000)
+    segments = [wav[i:i+win_samples] for i in range(0, len(wav) - win_samples, win_samples)]
+
+    rms_vals = [np.sqrt(np.mean(seg**2)) for seg in segments]
+    rms_avg = np.mean(rms_vals)
+    dynamic_thresh = threshold_ratio or max(0.02, 0.1 if rms_avg < 0.01 else 0.2)
+
+    drop_sequence = 0
+    consecutive_required = 2
+
+    for i, rms in enumerate(rms_vals):
+        if i < 3:
+            continue
+        if rms < rms_avg * dynamic_thresh:
+            drop_sequence += 1
+            if drop_sequence >= consecutive_required:
+                return True
+        else:
+            drop_sequence = 0
+
+    return False
+
+# ============================================================================
+# PROBLEMATIC CHUNK HANDLING
+# ============================================================================
+
+def handle_problematic_chunks(wav_path, issue_type, metrics):
+    """Handle chunks with audio issues - quarantine for review"""
+    quarantine_dir = wav_path.parent / "quarantine"
+    quarantine_dir.mkdir(exist_ok=True)
+
+    # Move to quarantine with descriptive name
+    quarantine_path = quarantine_dir / f"{wav_path.stem}_{issue_type}.wav"
+    shutil.move(str(wav_path), str(quarantine_path))
+
+    # Log for user review
+    logging.warning(f"🚨 Quarantined {issue_type}: {wav_path.name} → {quarantine_path.name}")
+    logging.warning(f"   Metrics: {metrics}")
+
+    return quarantine_path
+
+def pause_for_chunk_review(quarantine_dir):
+    """Pause processing to allow manual chunk review/editing with proper workflow"""
+    quarantined_files = list(quarantine_dir.glob("*.wav"))
+
+    if not quarantined_files:
+        return  # No quarantined files, continue normally
+
+    print(f"\n⚠️ {len(quarantined_files)} chunks quarantined in: {quarantine_dir}")
+    print("\nQuarantined chunks:")
+    for qfile in quarantined_files:
+        print(f"   📁 {qfile.name}")
+
+    print("\n🔧 Options:")
+    print("1. Continue processing (use quarantined chunks as-is)")
+    print("2. Pause to manually review/edit chunks")
+
+    while True:
+        choice = input("\nEnter choice [1/2]: ").strip()
+        if choice in ['1', '2']:
+            break
+        print("❌ Invalid choice. Please enter 1 or 2.")
+
+    if choice == "2":
+        print(f"\n🛑 Processing paused for manual review.")
+        print(f"📂 Quarantined chunks are in: {quarantine_dir}")
+        print("\n📝 Instructions:")
+        print("   1. Edit the audio files in the quarantine folder")
+        print("   2. Keep the original filenames (chunk numbering intact)")
+        print("   3. Leave edited files IN the quarantine folder")
+        print("   4. Press Enter below to continue processing")
+
+        input("\n⏸️  Press Enter when you've finished editing...")
+
+        # Verify files still exist after user editing
+        edited_files = list(quarantine_dir.glob("*.wav"))
+        if not edited_files:
+            print("⚠️ No files found in quarantine folder after editing!")
+            return
+
+        print(f"✅ Found {len(edited_files)} edited files, continuing...")
+
+    # Move all chunks back to main audio folder (whether edited or not)
+    moved_count = 0
+    for qfile in quarantine_dir.glob("*.wav"):
+        # Extract original chunk name from quarantine filename - FIXED LINE:
+        original_name = re.sub(r'_(clipping|corrupted|tts_hum)$', '', qfile.stem) + ".wav"
+        main_path = qfile.parent.parent / original_name
+
+        try:
+            shutil.move(str(qfile), str(main_path))
+            moved_count += 1
+            print(f"↩️ Restored: {original_name}")
+        except Exception as e:
+            logging.error(f"❌ Failed to restore {qfile.name}: {e}")
+
+    print(f"\n✅ Restored {moved_count} chunks to main audio folder")
+
+    # Clean up empty quarantine directory
+    if not any(quarantine_dir.iterdir()):
+        quarantine_dir.rmdir()
+
+    return moved_count
+
+# ============================================================================
+# AUDIO EFFECTS AND PROCESSING
+# ============================================================================
+
+def detect_end_artifact(wav_path, window_ms=100):
+    """Enhanced artifact detection"""
+    data, sr = sf.read(str(wav_path))
+    if data.ndim > 1:
+        data = data[:, 0]
+
+    win_samples = int(window_ms / 1000 * sr)
+    if len(data) < win_samples * 2:
+        return False
+
+    end = data[-win_samples:]
+    middle = data[len(data)//2 : len(data)//2 + win_samples]
+
+    rms_end = np.sqrt(np.mean(end**2))
+    rms_mid = np.sqrt(np.mean(middle**2)) + 1e-10
+    rms_ratio = rms_end / rms_mid
+
+    zcr = np.mean(np.diff(np.sign(end)) != 0)
+
+    fft = np.fft.rfft(end)
+    freqs = np.fft.rfftfreq(len(end), 1/sr)
+    low_band = fft[freqs < 150]
+    low_energy = np.sum(np.abs(low_band)) / (np.sum(np.abs(fft)) + 1e-10)
+
+    logging.info(f"{GREEN}[DEBUG]{RESET} Artifact metrics - {YELLOW}RMS ratio: {rms_ratio:.3f}{RESET}, "
+                f"{GREEN}ZCR: {zcr:.3f}{RESET}, {CYAN}LowEnergy: {low_energy:.3f}{RESET}")
+
+    return rms_ratio > 0.6 or zcr > 0.2 or low_energy > 0.4
+
+def find_end_of_speech(wav_path, sr=16000):
+    """Find end of speech using Silero VAD"""
+    import torch
+    import os
+
+    # Set environment variables to suppress PyTorch Hub verbosity
+    old_vars = {}
+    suppress_vars = {
+        'TORCH_HUB_VERBOSE': '0',
+        'PYTHONWARNINGS': 'ignore',
+        'TF_CPP_MIN_LOG_LEVEL': '3'
+    }
+
+    # Save old values and set new ones
+    for key, value in suppress_vars.items():
+        old_vars[key] = os.environ.get(key)
+        os.environ[key] = value
+
+    # Temporarily disable logging for this operation
+    old_level = logging.getLogger().level
+    logging.getLogger().setLevel(logging.ERROR)
+
+    try:
+        model, utils = torch.hub.load(
+            repo_or_dir='snakers4/silero-vad',
+            model='silero_vad',
+            force_reload=False,
+            verbose=False
+        )
+        (get_speech_timestamps, _, read_audio, _, _) = utils
+
+        wav = read_audio(str(wav_path), sampling_rate=sr)
+        speech_segments = get_speech_timestamps(wav, model, sampling_rate=sr)
+
+        if not speech_segments:
+            return None
+
+        last_seg_end = speech_segments[-1]['end']
+        return int(last_seg_end * 1000 / sr)
+
+    finally:
+        # Restore everything
+        logging.getLogger().setLevel(old_level)
+        for key, old_value in old_vars.items():
+            if old_value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = old_value
+
+def fade_out_wav(wav_path, output_path=None, fade_ms=20):
+    """Apply fade-out to audio"""
+    data, sr = sf.read(str(wav_path))
+    if data.ndim > 1:
+        data = data[:, 0]
+
+    fade_samples = int(sr * fade_ms / 1000)
+    if len(data) < fade_samples:
+        return
+
+    debug_path = wav_path.parent / f"{wav_path.stem}_pre_fade.wav"
+    sf.write(str(debug_path), data, sr)
+
+    fade_curve = np.linspace(1.0, 0.0, fade_samples)
+    data[-fade_samples:] *= fade_curve
+
+    sf.write(str(output_path or wav_path), data, sr)
+
+def apply_smart_fade(wav_path):
+    """Apply smart fade with artifact detection"""
+    eos_ms = find_end_of_speech(wav_path)
+
+    if detect_end_artifact(wav_path):
+        fade_out_wav(wav_path)
+
+def apply_smart_fade_memory(audio_segment):
+    """Apply smart fade with artifact detection - in memory version"""
+    # For now, apply a gentle fade to all audio to prevent clicks
+    # TODO: Add proper artifact detection for memory processing
+    return audio_segment.fade_out(50)  # 50ms fade out
+
+def smart_audio_validation_memory(audio_segment, sample_rate):
+    """Enhanced audio validation in memory - returns (audio, is_quarantined)"""
+    # Basic validation - can be enhanced with hum detection later
+    # For now, just return the audio as-is
+    is_quarantined = False
+    
+    # Could add memory-based hum detection here
+    # is_quarantined = detect_hum_memory(audio_segment, sample_rate)
+    
+    return audio_segment, is_quarantined
+
+def add_contextual_silence_memory(audio_segment, boundary_type):
+    """Add appropriate silence based on content boundary type - in memory"""
+    from pydub import AudioSegment
+    from config.config import (
+        SILENCE_CHAPTER_START, SILENCE_CHAPTER_END, SILENCE_SECTION_BREAK, SILENCE_PARAGRAPH_END,
+        SILENCE_COMMA, SILENCE_SEMICOLON, SILENCE_COLON, SILENCE_PERIOD, SILENCE_QUESTION_MARK,
+        SILENCE_EXCLAMATION, SILENCE_DASH, SILENCE_ELLIPSIS, SILENCE_QUOTE_END
+    )
+    
+    silence_durations = {
+        # Structural boundaries
+        "chapter_start": SILENCE_CHAPTER_START,
+        "chapter_end": SILENCE_CHAPTER_END,
+        "section_break": SILENCE_SECTION_BREAK,
+        "paragraph_end": SILENCE_PARAGRAPH_END,
+        # Punctuation boundaries
+        "comma": SILENCE_COMMA,
+        "semicolon": SILENCE_SEMICOLON,
+        "colon": SILENCE_COLON,
+        "period": SILENCE_PERIOD,
+        "question_mark": SILENCE_QUESTION_MARK,
+        "exclamation": SILENCE_EXCLAMATION,
+        "dash": SILENCE_DASH,
+        "ellipsis": SILENCE_ELLIPSIS,
+        "quote_end": SILENCE_QUOTE_END,
+    }
+    
+    if boundary_type in silence_durations:
+        duration = silence_durations[boundary_type]
+        silence_segment = AudioSegment.silent(duration=duration)
+        return audio_segment + silence_segment
+    
+    return audio_segment
+
+def smart_fade_out(wav_path, silence_thresh_db=-40, min_silence_len=300):
+    """Smart fade-out for natural audio endings"""
+    audio = AudioSegment.from_wav(wav_path)
+    tail_window_ms = 2000
+
+    if len(audio) < tail_window_ms:
+        logging.info(f"⚠️ {YELLOW}Skipping fade: {wav_path.name} too short ({len(audio)}ms < {tail_window_ms}ms){RESET}")
+        return
+
+    tail = audio[-tail_window_ms:]
+    silent_ranges = silence.detect_silence(tail, min_silence_len=min_silence_len, silence_thresh=silence_thresh_db)
+
+    min_tail_energy = max(tail.get_array_of_samples())
+    if not silent_ranges or min_tail_energy > audio.max_possible_amplitude * 0.1:
+        logging.info(f"✅ {GREEN}No fade needed for {wav_path.name} (no valid trailing silence){RESET}")
+        return
+
+    fade_start_ms = silent_ranges[0][0]
+    fade_length_ms = tail_window_ms - fade_start_ms
+
+    if fade_length_ms < 100:
+        logging.info(f"✅ {GREEN}No fade needed for {wav_path.name} (fade too short: {fade_length_ms}ms){RESET}")
+        return
+
+    fade_start_point = silent_ranges[0][0]
+    logging.info(f"⚠️ {RED}Fading tail of {wav_path.name} from {fade_start_point}ms to end{RESET}")
+    faded = audio[:fade_start_point] + audio[fade_start_point:].fade_out(duration=fade_length_ms)
+    faded.export(wav_path, format="wav")
+
+# ============================================================================
+# AUDIO TRIMMING
+# ============================================================================
+
+def trim_audio_endpoint(audio_segment, threshold=None, buffer_ms=None):
+    """
+    Trim audio to the detected end of speech using RMS energy analysis.
+    
+    Args:
+        audio_segment: pydub AudioSegment object
+        threshold: RMS threshold for speech detection (from config if None)
+        buffer_ms: Buffer to add after detected endpoint (from config if None)
+    
+    Returns:
+        Trimmed AudioSegment
+    """
+    if threshold is None:
+        threshold = SPEECH_ENDPOINT_THRESHOLD
+    if buffer_ms is None:
+        buffer_ms = TRIMMING_BUFFER_MS
+    
+    # Convert to numpy array for analysis
+    samples = np.array(audio_segment.get_array_of_samples())
+    if audio_segment.channels == 2:
+        samples = samples.reshape((-1, 2)).mean(axis=1)
+    
+    # Normalize samples
+    samples = samples.astype(np.float32) / audio_segment.max_possible_amplitude
+    
+    # Calculate RMS in sliding windows (50ms windows)
+    window_size = int(0.05 * audio_segment.frame_rate)  # 50ms
+    rms_values = []
+    
+    for i in range(0, len(samples) - window_size, window_size // 2):
+        window = samples[i:i + window_size]
+        rms = np.sqrt(np.mean(window ** 2))
+        rms_values.append(rms)
+    
+    # Find actual end of speech using energy decay detection
+    speech_end_idx = 0  # Default to beginning if no speech found
+    
+    # Look for a significant and sustained drop in energy
+    # Scan backwards to find where energy consistently stays above a higher threshold
+    strong_speech_threshold = threshold * 3  # 3x threshold for "real" speech
+    
+    for i in range(len(rms_values) - 1, -1, -1):
+        if rms_values[i] > strong_speech_threshold:
+            # Found strong speech, check if it's sustained
+            # Look forward to see if energy drops and stays low
+            sustained_speech = True
+            windows_ahead = min(10, len(rms_values) - i)  # Look ahead up to 10 windows (250ms)
+            
+            # Check if most of the next windows have reasonable speech levels
+            speech_count = 0
+            for j in range(i, min(i + windows_ahead, len(rms_values))):
+                if rms_values[j] > threshold:
+                    speech_count += 1
+            
+            # If this looks like the end of sustained speech content
+            if speech_count >= max(1, windows_ahead * 0.3):  # At least 30% speech in next windows
+                speech_end_idx = i
+                break
+    
+    # If no strong speech found, fall back to simple threshold method but be conservative
+    if speech_end_idx == 0:
+        for i in range(len(rms_values) - 1, -1, -1):
+            if rms_values[i] > threshold * 2:  # Use 2x threshold for fallback
+                speech_end_idx = i
+                break
+    
+    # Convert back to milliseconds and add buffer
+    # Convert window index to sample position, then to milliseconds
+    sample_position = speech_end_idx * (window_size // 2)
+    speech_end_ms = int(sample_position * 1000 / audio_segment.frame_rate)
+    trim_point_ms = min(speech_end_ms + buffer_ms, len(audio_segment))
+    
+    return audio_segment[:trim_point_ms]
+
+def process_audio_with_trimming_and_silence(audio_segment, boundary_type, enable_trimming=None):
+    """
+    Complete audio processing: trim to speech endpoint + add punctuation-based silence.
+    
+    Args:
+        audio_segment: pydub AudioSegment object
+        boundary_type: Boundary type from text processing
+        enable_trimming: Whether to trim audio (from config if None)
+    
+    Returns:
+        Processed AudioSegment with trimming and appropriate silence
+    """
+    if enable_trimming is None:
+        enable_trimming = ENABLE_AUDIO_TRIMMING
+    
+    processed_audio = audio_segment
+    
+    # Step 1: Trim to speech endpoint if enabled
+    if enable_trimming:
+        processed_audio = trim_audio_endpoint(processed_audio)
+    
+    # Step 2: Add punctuation-appropriate silence
+    processed_audio = add_contextual_silence_memory(processed_audio, boundary_type)
+    
+    return processed_audio
+
+# ============================================================================
+# SILENCE AND CONTEXTUAL AUDIO
+# ============================================================================
+
+def add_contextual_silence(wav_path, boundary_type):
+    """Add appropriate silence based on content boundary type"""
+    silence_durations = {
+        # Structural boundaries
+        "chapter_start": SILENCE_CHAPTER_START,
+        "chapter_end": SILENCE_CHAPTER_END,
+        "section_break": SILENCE_SECTION_BREAK,
+        "paragraph_end": SILENCE_PARAGRAPH_END,
+        # Punctuation boundaries
+        "comma": SILENCE_COMMA,
+        "semicolon": SILENCE_SEMICOLON,
+        "colon": SILENCE_COLON,
+        "period": SILENCE_PERIOD,
+        "question_mark": SILENCE_QUESTION_MARK,
+        "exclamation": SILENCE_EXCLAMATION,
+        "dash": SILENCE_DASH,
+        "ellipsis": SILENCE_ELLIPSIS,
+        "quote_end": SILENCE_QUOTE_END,
+    }
+
+    if boundary_type in silence_durations:
+        duration = silence_durations[boundary_type]
+        audio = AudioSegment.from_wav(wav_path)
+        silence_segment = AudioSegment.silent(duration=duration)
+        extended_audio = audio + silence_segment
+        extended_audio.export(wav_path, format="wav")
+
+        logging.info(f"🔇 Added {duration}ms silence for {boundary_type}: {wav_path.name}")
+
+def add_chunk_end_silence(wav_path):
+    """Add configurable silence to end of chunk if enabled"""
+    if not ENABLE_CHUNK_END_SILENCE or CHUNK_END_SILENCE_MS <= 0:
+        return
+
+    try:
+        audio = AudioSegment.from_wav(wav_path)
+        silence_segment = AudioSegment.silent(duration=CHUNK_END_SILENCE_MS)
+        audio_with_silence = audio + silence_segment
+        audio_with_silence.export(wav_path, format="wav")
+        logging.info(f"➕ Added {CHUNK_END_SILENCE_MS}ms end silence to {wav_path.name}")
+    except Exception as e:
+        logging.warning(f"⚠️ Failed to add end silence to {wav_path.name}: {e}")
+
+# ============================================================================
+# AUDIO UTILITY FUNCTIONS
+# ============================================================================
+
+def get_wav_duration(wav_path):
+    """Get WAV file duration"""
+    import wave
+    with wave.open(str(wav_path), 'rb') as wf:
+        frames = wf.getnframes()
+        rate = wf.getframerate()
+        return frames / float(rate)
+
+def get_chunk_audio_duration(wav_path):
+    """Get actual audio duration from WAV file"""
+    try:
+        data, sr = sf.read(str(wav_path))
+        return len(data) / sr
+    except:
+        # Fallback to wave module
+        return get_wav_duration(wav_path)

HF_Deploy/modules/batch_processor.py

@@ -0,0 +1,31 @@
+"""
+Batch Processing Module
+Handles multi-book batch processing operations
+"""
+
+import torch
+from modules.tts_engine import process_book_folder
+
+def pipeline_book_processing(book_queue):
+    """Process multiple books in sequence"""
+    completed_books = []
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    for book_info in book_queue:
+        book_dir = book_info['book_dir']
+        voice_path = book_info['voice_path'] 
+        tts_params = book_info['tts_params']
+        
+        print(f"\n🎯 Processing: {book_dir.name}")
+        
+        try:
+            result = process_book_folder(book_dir, voice_path, tts_params, device)
+            if result[0]:  # Check if final_m4b_path exists
+                completed_books.append(book_info)
+                print(f"✅ Completed: {book_dir.name}")
+            else:
+                print(f"❌ Failed: {book_dir.name}")
+        except Exception as e:
+            print(f"❌ Error processing {book_dir.name}: {e}")
+    
+    return completed_books

HF_Deploy/modules/file_manager.py

@@ -0,0 +1,431 @@
+"""
+File Manager Module
+Handles I/O operations, M4B conversion, metadata, and FFmpeg operations
+"""
+
+import subprocess
+import soundfile as sf
+import os
+import re
+import time
+import logging
+from pathlib import Path
+from config.config import *
+
+# ============================================================================
+# VOICE SAMPLE MANAGEMENT
+# ============================================================================
+
+def list_voice_samples():
+    """List available voice samples"""
+    return sorted(VOICE_SAMPLES_DIR.glob("*.wav"))
+
+def ensure_voice_sample_compatibility(input_path, output_dir=None):
+    """Ensure voice sample is compatible with TTS (24kHz mono)"""
+    input_path = str(input_path)
+    ext = os.path.splitext(input_path)[1].lower()
+    basename = os.path.splitext(os.path.basename(input_path))[0]
+    output_dir = output_dir or os.path.dirname(input_path)
+    output_path = os.path.join(output_dir, basename + "_ttsready.wav")
+
+    try:
+        info = sf.info(input_path)
+        if (ext == '.wav' and info.samplerate == 24000 and info.channels == 1):
+            return input_path
+    except Exception:
+        pass
+
+    cmd = [
+        "ffmpeg", "-y",
+        "-i", input_path,
+        "-ar", "24000",
+        "-ac", "1",
+        output_path
+    ]
+    subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+    return output_path
+
+# ============================================================================
+# FFMPEG OPERATIONS
+# ============================================================================
+
+def run_ffmpeg(cmd):
+    """Run FFmpeg command with error handling"""
+    try:
+        subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+    except subprocess.CalledProcessError as e:
+        logging.info(f"FFmpeg command failed: {' '.join(cmd)}")
+        logging.info(f"Error: {e}")
+        subprocess.run(cmd)
+        raise
+
+# ============================================================================
+# M4B CONVERSION WITH NORMALIZATION
+# ============================================================================
+
+def convert_to_m4b_with_peak_normalization(wav_path, temp_m4b_path, target_db=-3.0):
+    """Convert WAV to M4B with peak normalization"""
+    print("🚀 Converting to m4b with peak normalization...")
+
+    # Build audio filter chain
+    audio_filters = [f"loudnorm=I=-16:TP={target_db}:LRA=11"]
+    if ATEMPO_SPEED != 1.0:
+        audio_filters.append(f"atempo={ATEMPO_SPEED}")
+    
+    cmd = [
+        "ffmpeg", "-y",
+        "-i", str(wav_path),
+        "-af", ",".join(audio_filters),
+        "-c:a", "aac",
+        str(temp_m4b_path)
+    ]
+
+    start_time = time.time()
+    process = subprocess.Popen(cmd, stderr=subprocess.PIPE, text=True)
+
+    audio_secs = 0.0
+    for line in process.stderr:
+        match = re.search(r"time=(\d{2}):(\d{2}):(\d{2})\.(\d{2})", line)
+        if match:
+            h, m, s, ms = map(int, match.groups())
+            audio_secs = h * 3600 + m * 60 + s + ms / 100
+            elapsed = time.time() - start_time
+            factor = audio_secs / elapsed if elapsed > 0 else 0.0
+            print(f"📼 FFmpeg (normalizing): {match.group(0)} | {factor:.2f}x realtime", end='\r')
+
+    process.wait()
+    print("\n✅ Conversion with normalization complete.")
+
+def convert_to_m4b_with_loudness_normalization(wav_path, temp_m4b_path):
+    """Convert WAV to M4B with two-pass loudness normalization"""
+    import json
+
+    print("🚀 Converting to m4b with loudness normalization...")
+
+    # Step 1: Analyze audio loudness
+    print("📊 Analyzing audio loudness...")
+    analyze_cmd = [
+        "ffmpeg", "-y",
+        "-i", str(wav_path),
+        "-af", "loudnorm=I=-16:TP=-1.5:LRA=11:print_format=json",
+        "-f", "null", "-"
+    ]
+
+    result = subprocess.run(analyze_cmd, capture_output=True, text=True)
+
+    # Extract loudness measurements from stderr
+    loudness_data = None
+    for line in result.stderr.split('\n'):
+        if line.strip().startswith('{'):
+            try:
+                loudness_data = json.loads(line.strip())
+                break
+            except:
+                continue
+
+    if not loudness_data:
+        print("⚠️ Could not analyze loudness, falling back to single-pass...")
+        return convert_to_m4b_with_peak_normalization(wav_path, temp_m4b_path)
+
+    # Step 2: Apply normalization with measured values
+    print("🔧 Applying normalization...")
+    
+    # Build audio filter chain
+    audio_filters = [f"loudnorm=I=-16:TP=-1.5:LRA=11:measured_I={loudness_data['input_i']}:measured_LRA={loudness_data['input_lra']}:measured_TP={loudness_data['input_tp']}:measured_thresh={loudness_data['input_thresh']}:offset={loudness_data['target_offset']}:linear=true:print_format=summary"]
+    if ATEMPO_SPEED != 1.0:
+        audio_filters.append(f"atempo={ATEMPO_SPEED}")
+    
+    cmd = [
+        "ffmpeg", "-y",
+        "-i", str(wav_path),
+        "-af", ",".join(audio_filters),
+        "-c:a", "aac",
+        str(temp_m4b_path)
+    ]
+
+    start_time = time.time()
+    process = subprocess.Popen(cmd, stderr=subprocess.PIPE, text=True)
+
+    audio_secs = 0.0
+    for line in process.stderr:
+        match = re.search(r"time=(\d{2}):(\d{2}):(\d{2})\.(\d{2})", line)
+        if match:
+            h, m, s, ms = map(int, match.groups())
+            audio_secs = h * 3600 + m * 60 + s + ms / 100
+            elapsed = time.time() - start_time
+            factor = audio_secs / elapsed if elapsed > 0 else 0.0
+            print(f"📼 FFmpeg (normalizing): {match.group(0)} | {factor:.2f}x realtime", end='\r')
+
+    process.wait()
+    print("\n✅ Two-pass normalization complete.")
+
+def convert_to_m4b_with_simple_normalization(wav_path, temp_m4b_path, target_db=-6.0):
+    """Convert WAV to M4B with simple peak normalization"""
+    print("🚀 Converting to m4b with simple normalization...")
+
+    # Build audio filter chain
+    audio_filters = [f"volume={target_db}dB"]
+    if ATEMPO_SPEED != 1.0:
+        audio_filters.append(f"atempo={ATEMPO_SPEED}")
+
+    cmd = [
+        "ffmpeg", "-y",
+        "-i", str(wav_path),
+        "-af", ",".join(audio_filters),
+        "-c:a", "aac",
+        str(temp_m4b_path)
+    ]
+
+    start_time = time.time()
+    process = subprocess.Popen(cmd, stderr=subprocess.PIPE, text=True)
+
+    audio_secs = 0.0
+    for line in process.stderr:
+        match = re.search(r"time=(\d{2}):(\d{2}):(\d{2})\.(\d{2})", line)
+        if match:
+            h, m, s, ms = map(int, match.groups())
+            audio_secs = h * 3600 + m * 60 + s + ms / 100
+            elapsed = time.time() - start_time
+            factor = audio_secs / elapsed if elapsed > 0 else 0.0
+            print(f"📼 FFmpeg (normalizing): {match.group(0)} | {factor:.2f}x realtime", end='\r')
+
+    process.wait()
+    print("\n✅ Simple normalization complete.")
+
+def convert_to_m4b(wav_path, temp_m4b_path):
+    """Convert WAV to M4B with configurable normalization"""
+    if not ENABLE_NORMALIZATION or NORMALIZATION_TYPE == "none":
+        # Original function without normalization
+        print("🚀 Converting to m4b...")
+
+        # Build audio filter for atempo if needed
+        audio_filter = []
+        if ATEMPO_SPEED != 1.0:
+            audio_filter = ["-filter:a", f"atempo={ATEMPO_SPEED}"]
+
+        cmd = [
+            "ffmpeg", "-y",
+            "-i", str(wav_path)
+        ] + audio_filter + [
+            "-c:a", "aac",
+            str(temp_m4b_path)
+        ]
+
+    elif NORMALIZATION_TYPE == "loudness":
+        # EBU R128 loudness normalization (recommended for audiobooks)
+        return convert_to_m4b_with_loudness_normalization(wav_path, temp_m4b_path)
+
+    elif NORMALIZATION_TYPE == "peak":
+        # Peak normalization
+        return convert_to_m4b_with_peak_normalization(wav_path, temp_m4b_path, TARGET_PEAK_DB)
+
+    elif NORMALIZATION_TYPE == "simple":
+        # Simple volume adjustment
+        return convert_to_m4b_with_simple_normalization(wav_path, temp_m4b_path, TARGET_PEAK_DB)
+
+    else:
+        # Fallback to no normalization
+        # Build audio filter for atempo if needed
+        audio_filter = []
+        if ATEMPO_SPEED != 1.0:
+            audio_filter = ["-filter:a", f"atempo={ATEMPO_SPEED}"]
+
+        cmd = [
+            "ffmpeg", "-y",
+            "-i", str(wav_path)
+        ] + audio_filter + [
+            "-c:a", "aac",
+            str(temp_m4b_path)
+        ]
+
+    # Run the conversion (if not handled by specialized functions above)
+    start_time = time.time()
+    process = subprocess.Popen(cmd, stderr=subprocess.PIPE, text=True)
+
+    audio_secs = 0.0
+    for line in process.stderr:
+        match = re.search(r"time=(\d{2}):(\d{2}):(\d{2})\.(\d{2})", line)
+        if match:
+            h, m, s, ms = map(int, match.groups())
+            audio_secs = h * 3600 + m * 60 + s + ms / 100
+            elapsed = time.time() - start_time
+            factor = audio_secs / elapsed if elapsed > 0 else 0.0
+            print(f"📼 FFmpeg: {match.group(0)} | {factor:.2f}x realtime", end='\r')
+
+    process.wait()
+    print("\n✅ Conversion complete.")
+
+def add_metadata_to_m4b(temp_m4b_path, final_m4b_path, cover_path=None, nfo_path=None):
+    """Add metadata and cover to M4B"""
+    cmd = ["ffmpeg", "-y", "-i", str(temp_m4b_path)]
+
+    if cover_path and cover_path.exists():
+        cmd.extend(["-i", str(cover_path), "-map", "0", "-map", "1", "-c", "copy", "-disposition:v:0", "attached_pic"])
+    else:
+        cmd.extend(["-map", "0", "-c", "copy"])
+
+    if nfo_path and nfo_path.exists():
+        with open(nfo_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                if ':' in line:
+                    key, val = line.strip().split(':', 1)
+                    cmd.extend(["-metadata", f"{key.strip()}={val.strip()}"])
+
+    cmd.append(str(final_m4b_path))
+    run_ffmpeg(cmd)
+    temp_m4b_path.unlink(missing_ok=True)
+
+# ============================================================================
+# FILE UTILITIES
+# ============================================================================
+
+def chunk_sort_key(f):
+    """Extracts the chunk number for natural sorting"""
+    m = re.match(r"chunk_(\d+)\.wav", f.name)
+    return int(m.group(1)) if m else 0
+
+def create_concat_file(chunk_paths, output_path):
+    """Create FFmpeg concat file for audio chunks"""
+    with open(output_path, 'w') as f:
+        for p in chunk_paths:
+            # Use absolute path to ensure FFmpeg can find the files
+            f.write(f"file '{str(p.resolve())}'\n")
+
+    logging.info(f"concat.txt written with {len(chunk_paths)} chunks.")
+    return output_path
+
+def cleanup_temp_files(directory, patterns):
+    """Clean up temporary files matching patterns"""
+    files_cleaned = 0
+    for pattern in patterns:
+        for temp_file in directory.glob(pattern):
+            temp_file.unlink(missing_ok=True)
+            files_cleaned += 1
+
+    return files_cleaned
+
+# ============================================================================
+# DIRECTORY MANAGEMENT
+# ============================================================================
+
+def setup_book_directories(book_dir):
+    """Set up directory structure for book processing"""
+    basename = book_dir.name
+    output_root = AUDIOBOOK_ROOT / basename
+    tts_dir = output_root / "TTS"
+    text_chunks_dir = tts_dir / "text_chunks"
+    audio_chunks_dir = tts_dir / "audio_chunks"
+
+    # Create directories
+    for d in [output_root, tts_dir, text_chunks_dir, audio_chunks_dir]:
+        d.mkdir(parents=True, exist_ok=True)
+
+    return output_root, tts_dir, text_chunks_dir, audio_chunks_dir
+
+def find_book_files(book_dir):
+    """Find text files, cover, and metadata for a book"""
+    text_files = sorted(book_dir.glob("*.txt"))
+    nfo_file = book_dir / "book.nfo"
+    cover_jpg = book_dir / "cover.jpg"
+    cover_png = book_dir / "cover.png"
+    cover_file = cover_jpg if cover_jpg.exists() else cover_png if cover_png.exists() else None
+
+    return {
+        'text': text_files[0] if text_files else None,
+        'cover': cover_file,
+        'nfo': nfo_file if nfo_file.exists() else None
+    }
+
+# ============================================================================
+# AUDIO FILE OPERATIONS
+# ============================================================================
+
+def combine_audio_chunks(chunk_paths, output_path):
+    """Combine audio chunks into single file using FFmpeg"""
+    concat_list_path = output_path.parent / "concat.txt"
+    create_concat_file(chunk_paths, concat_list_path)
+
+    run_ffmpeg([
+        "ffmpeg", "-y", "-f", "concat", "-safe", "0",
+        "-i", str(concat_list_path.resolve()),
+        "-c", "copy", str(output_path.resolve())
+    ])
+
+    return output_path
+
+def get_audio_files_in_directory(directory, pattern="chunk_*.wav"):
+    """Get sorted list of audio files matching pattern"""
+    chunk_paths = sorted([f for f in directory.glob(pattern)
+                         if re.fullmatch(r'chunk_\d{3,}\.wav', f.name)],
+                        key=chunk_sort_key)
+    return chunk_paths
+
+# ============================================================================
+# VALIDATION AND VERIFICATION
+# ============================================================================
+
+def verify_audio_file(wav_path):
+    """Verify audio file is valid and readable"""
+    try:
+        info = sf.info(str(wav_path))
+        return info.frames > 0 and info.samplerate > 0
+    except Exception as e:
+        logging.error(f"Invalid audio file {wav_path}: {e}")
+        return False
+
+def verify_chunk_completeness(audio_chunks_dir, expected_count):
+    """Verify all expected chunks exist and are valid"""
+    missing_chunks = []
+    invalid_chunks = []
+
+    for i in range(1, expected_count + 1):
+        chunk_path = audio_chunks_dir / f"chunk_{i:05}.wav"
+
+        if not chunk_path.exists():
+            missing_chunks.append(i)
+        elif not verify_audio_file(chunk_path):
+            invalid_chunks.append(i)
+
+    return missing_chunks, invalid_chunks
+
+# ============================================================================
+# EXPORT AND IMPORT FUNCTIONS
+# ============================================================================
+
+def export_processing_log(output_dir, processing_info):
+    """Export comprehensive processing log"""
+    log_path = output_dir / "processing_complete.log"
+
+    with open(log_path, 'w', encoding='utf-8') as f:
+        f.write("GenTTS Processing Complete\n")
+        f.write("=" * 50 + "\n\n")
+
+        for key, value in processing_info.items():
+            f.write(f"{key}: {value}\n")
+
+    return log_path
+
+def save_chunk_info(text_chunks_dir, chunks_info):
+    """Save chunk information for debugging/resume"""
+    info_path = text_chunks_dir / "chunks_info.json"
+
+    import json
+    with open(info_path, 'w', encoding='utf-8') as f:
+        json.dump(chunks_info, f, indent=2, ensure_ascii=False)
+
+    return info_path
+
+def load_chunk_info(text_chunks_dir):
+    """Load chunk information if available"""
+    info_path = text_chunks_dir / "chunks_info.json"
+
+    if not info_path.exists():
+        return None
+
+    import json
+    try:
+        with open(info_path, 'r', encoding='utf-8') as f:
+            return json.load(f)
+    except Exception as e:
+        logging.warning(f"Could not load chunk info: {e}")
+        return None

HF_Deploy/modules/gui_json_generator.py

@@ -0,0 +1,217 @@
+#!/usr/bin/env python3
+"""
+GUI JSON Audio Generation Module
+
+This module provides JSON-to-audiobook generation specifically for GUI use.
+It's based on utils/generate_from_json.py but adapted for GUI integration.
+"""
+
+import torch
+from pathlib import Path
+import sys
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import time
+from datetime import timedelta
+
+# Add project root to path to allow module imports
+project_root = Path(__file__).parent.parent
+sys.path.append(str(project_root))
+
+from config.config import *
+from modules.tts_engine import load_optimized_model, process_one_chunk
+from modules.file_manager import setup_book_directories, list_voice_samples, ensure_voice_sample_compatibility
+from wrapper.chunk_loader import load_chunks
+from src.chatterbox.tts import punc_norm
+from modules.progress_tracker import log_chunk_progress, log_run
+from tools.combine_only import combine_audio_for_book
+
+
+def generate_audiobook_from_json(json_path, voice_name, temp_setting=None):
+    """
+    Generate complete audiobook from JSON chunks file.
+    
+    Args:
+        json_path (str): Path to the JSON chunks file
+        voice_name (str): Name of the voice to use (without .wav extension)
+        temp_setting (float, optional): Temperature override for TTS
+        
+    Returns:
+        tuple: (success: bool, message: str, audiobook_path: str or None)
+    """
+    try:
+        print(f"🎵 GUI JSON Generator: Starting audiobook generation")
+        print(f"📄 JSON file: {json_path}")
+        print(f"🎤 Voice: {voice_name}")
+        if temp_setting:
+            print(f"🌡️ Temperature override: {temp_setting}")
+        
+        # Determine book name from JSON path
+        json_file = Path(json_path)
+        
+        # Try to extract book name from path structure
+        if 'Audiobook' in json_file.parts:
+            audiobook_index = json_file.parts.index('Audiobook')
+            if audiobook_index + 1 < len(json_file.parts):
+                book_name = json_file.parts[audiobook_index + 1]
+                print(f"📚 Detected book name from path: {book_name}")
+            else:
+                raise Exception("Cannot determine book name from Audiobook path")
+        elif json_file.stem.endswith('_chunks'):
+            book_name = json_file.stem.replace('_chunks', '')
+            print(f"📚 Detected book name from filename: {book_name}")
+        else:
+            book_name = json_file.stem
+            print(f"📚 Using filename as book name: {book_name}")
+
+        # Load JSON chunks (READ ONLY - never modify the original)
+        print(f"📖 Loading chunks from: {json_path}")
+        all_chunks = load_chunks(str(json_path))
+        print(f"✅ Found {len(all_chunks)} chunks.")
+
+        # Find voice file
+        voice_files = list_voice_samples()
+        voice_path = None
+        for voice_file in voice_files:
+            if voice_file.stem == voice_name:
+                voice_path = voice_file
+                break
+        
+        if not voice_path:
+            available_voices = [vf.stem for vf in voice_files]
+            return False, f"Voice '{voice_name}' not found. Available: {available_voices}", None
+
+        # Ensure voice compatibility
+        voice_path = ensure_voice_sample_compatibility(voice_path)
+        if isinstance(voice_path, str):
+            voice_path = Path(voice_path)
+        
+        print(f"🎤 Using voice: {voice_path.name}")
+
+        # Setup device
+        if torch.cuda.is_available():
+            device = "cuda"
+        elif torch.backends.mps.is_available():
+            device = "mps"
+        else:
+            device = "cpu"
+        
+        print(f"🚀 Using device: {device}")
+
+        # Load TTS model
+        print(f"🤖 Loading TTS model...")
+        model = load_optimized_model(device)
+        
+        # Prepare voice conditionals
+        print(f"🎤 Preparing voice conditionals...")
+        model.prepare_conditionals(voice_path)
+
+        # Setup output directories
+        output_root = AUDIOBOOK_ROOT / book_name
+        tts_dir = output_root / "TTS"
+        text_chunks_dir = tts_dir / "text_chunks"
+        audio_chunks_dir = tts_dir / "audio_chunks"
+        
+        # Create directories
+        for dir_path in [output_root, tts_dir, text_chunks_dir, audio_chunks_dir]:
+            dir_path.mkdir(parents=True, exist_ok=True)
+
+        # Clean existing audio chunks
+        print("🧹 Clearing old audio chunks...")
+        for wav_file in audio_chunks_dir.glob("*.wav"):
+            wav_file.unlink()
+
+        # Process chunks
+        start_time = time.time()
+        total_chunks = len(all_chunks)
+        log_path = output_root / "gui_json_generation.log"
+        
+        print(f"🔄 Generating {total_chunks} audio chunks...")
+
+        with ThreadPoolExecutor(max_workers=2) as executor:
+            futures = []
+            for i, chunk_data in enumerate(all_chunks):
+                # Use chunk's TTS params, with temperature override if provided
+                chunk_tts_params = chunk_data.get("tts_params", {}).copy()
+                if temp_setting is not None:
+                    chunk_tts_params["temperature"] = temp_setting
+                
+                # Ensure required TTS params exist
+                chunk_tts_params.setdefault("exaggeration", DEFAULT_EXAGGERATION)
+                chunk_tts_params.setdefault("cfg_weight", DEFAULT_CFG_WEIGHT)
+                chunk_tts_params.setdefault("temperature", DEFAULT_TEMPERATURE)
+
+                future = executor.submit(
+                    process_one_chunk,
+                    i, chunk_data['text'], text_chunks_dir, audio_chunks_dir,
+                    voice_path, chunk_tts_params, start_time, total_chunks,
+                    punc_norm, book_name, log_run, log_path, device,
+                    model, None, all_chunks, chunk_data.get('boundary_type', 'none')
+                )
+                futures.append(future)
+
+            # Wait for all chunks to complete
+            completed_chunks = 0
+            for future in as_completed(futures):
+                try:
+                    result = future.result()
+                    if result:
+                        idx, _ = result
+                        completed_chunks += 1
+                        log_chunk_progress(idx, total_chunks, start_time, 0)
+                        print(f"✅ Completed chunk {completed_chunks}/{total_chunks}")
+                except Exception as e:
+                    print(f"❌ Error processing chunk: {e}")
+
+        elapsed_time = time.time() - start_time
+        print(f"✅ Audio generation complete in {timedelta(seconds=int(elapsed_time))}")
+        print(f"🔊 Audio chunks generated in: {audio_chunks_dir}")
+
+        # Combine chunks into final audiobook
+        print("🔗 Combining audio chunks into final audiobook...")
+        try:
+            success = combine_audio_for_book(str(output_root), voice_name)
+            if success:
+                # Look for the created audiobook file with voice name
+                final_m4b = output_root / f"{book_name} [{voice_name}].m4b"
+                if final_m4b.exists():
+                    print(f"🎉 Audiobook created: {final_m4b.name}")
+                    return True, "Audiobook generation completed successfully", str(final_m4b)
+                else:
+                    return False, "Combine succeeded but final audiobook file not found", None
+            else:
+                return False, "Failed to combine audio chunks", None
+        except Exception as e:
+            return False, f"Error combining audio chunks: {e}", None
+
+    except Exception as e:
+        error_msg = f"JSON generation error: {e}"
+        print(f"❌ {error_msg}")
+        return False, error_msg, None
+
+
+def get_book_name_from_json_path(json_path):
+    """
+    Extract book name from JSON file path.
+    
+    Args:
+        json_path (str): Path to JSON file
+        
+    Returns:
+        str: Detected book name
+    """
+    json_file = Path(json_path)
+    
+    if 'Audiobook' in json_file.parts:
+        audiobook_index = json_file.parts.index('Audiobook')
+        if audiobook_index + 1 < len(json_file.parts):
+            return json_file.parts[audiobook_index + 1]
+    
+    if json_file.stem.endswith('_chunks'):
+        return json_file.stem.replace('_chunks', '')
+    
+    return json_file.stem
+
+
+if __name__ == "__main__":
+    # CLI compatibility for testing
+    print("GUI JSON Generator - use from GUI or import as module")

HF_Deploy/modules/path_manager.py

@@ -0,0 +1,19 @@
+from pathlib import Path
+from config.config import AUDIOBOOK_ROOT
+
+def get_book_paths(book_name):
+    """Return standardized paths for a given book name"""
+    base = AUDIOBOOK_ROOT / book_name
+    tts_dir = base / "TTS"
+    return {
+        "book_folder": base,
+        "tts_dir": tts_dir,
+        "text_chunks": tts_dir / "text_chunks",
+        "audio_chunks": tts_dir / "audio_chunks",
+        "combined_wav": base / f"{book_name}.wav",
+        "final_m4b": base / f"{book_name}.m4b",
+        "concat_list": tts_dir / "audio_chunks" / "concat.txt",
+        "quarantine": tts_dir / "audio_chunks" / "quarantine",
+        "run_log": base / "run.log",
+        "chunk_log": base / "chunk_validation.log"
+    }

HF_Deploy/modules/progress_tracker.py

@@ -0,0 +1,306 @@
+"""
+Progress Tracker Module
+Handles progress display, VRAM monitoring, logging systems, and performance tracking
+"""
+
+import time
+import sys
+import logging
+from datetime import timedelta
+from pathlib import Path
+from config.config import *
+
+# ============================================================================
+# LOGGING SETUP
+# ============================================================================
+
+def setup_logging(log_dir):
+    """Setup logging configuration"""
+    log_file = log_dir / "chunk_validation.log"
+
+    # Clear existing log
+    open(log_file, 'w').close()
+
+    logging.basicConfig(
+        filename=str(log_file),
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s",
+        filemode='w'  # Overwrite existing log
+    )
+
+    # Also log to console for important messages
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.WARNING)
+    formatter = logging.Formatter('%(levelname)s - %(message)s')
+    console_handler.setFormatter(formatter)
+    logging.getLogger().addHandler(console_handler)
+
+def log_console(message, color=None):
+    """Log to both console and file with optional color"""
+    color_codes = {
+        "RED": RED, "GREEN": GREEN, "YELLOW": YELLOW,
+        "CYAN": CYAN, "BOLD": BOLD, "RESET": RESET
+    }
+
+    prefix = color_codes.get(color, "")
+    suffix = RESET if color else ""
+
+    print(f"{prefix}{message}{suffix}")
+    logging.info(message)
+
+def log_run(message, log_path):
+    """Log to run file"""
+    with open(log_path, "a", encoding="utf-8") as logf:
+        logf.write(message + "\n")
+
+# ============================================================================
+# PROGRESS TRACKING
+# ============================================================================
+
+def log_chunk_progress(i, total_chunks, start_time, total_audio_duration=0.0):
+    """Enhanced progress logging with accurate realtime factor"""
+    elapsed = time.time() - start_time
+    avg_time = elapsed / (i + 1)
+    eta = avg_time * total_chunks
+    remaining = eta - elapsed
+
+    def fmt(seconds):
+        return str(timedelta(seconds=int(seconds)))
+
+    # Show VRAM usage in progress
+    allocated, _ = monitor_vram_usage("chunk_progress")
+
+    # Calculate ACCURATE realtime factor using actual audio duration
+    if total_audio_duration > 0 and elapsed > 0:
+        actual_realtime = total_audio_duration / elapsed
+        realtime_str = f"{GREEN}{actual_realtime:.2f}x{RESET}"
+        audio_str = f" | Audio: {GREEN}{fmt(total_audio_duration)}{RESET}"
+    else:
+        actual_realtime = 0.0  # Default value when calculating
+        realtime_str = f"{YELLOW}Calculating...{RESET}"
+        audio_str = ""
+
+    # Force immediate output with explicit flushing
+    progress_msg = (f"\n🌀 Chunk {i+1}/{total_chunks} | ⏱ Elapsed: {CYAN}{fmt(elapsed)}{RESET} | "
+                   f"ETA: {CYAN}{fmt(eta)}{RESET} | Remaining: {YELLOW}{fmt(remaining)}{RESET} | "
+                   f"Realtime: {realtime_str} | VRAM: {GREEN}{allocated:.1f}GB{RESET}{audio_str}")
+
+    print(progress_msg)
+    sys.stdout.flush()  # Force immediate output
+    
+    # Create clean status message for GUI (without ANSI color codes)
+    realtime_display = f"{actual_realtime:.2f}x" if actual_realtime > 0 else "Calculating..."
+    clean_status = (f"Elapsed: {fmt(elapsed)} | ETA: {fmt(eta)} | Remaining: {fmt(remaining)} | "
+                   f"Realtime: {realtime_display} | VRAM: {allocated:.1f}GB" +
+                   (f" | Audio: {fmt(total_audio_duration)}" if total_audio_duration > 0 else ""))
+    
+    # Emit status to GUI if callback is available
+    if hasattr(log_chunk_progress, '_status_callback') and log_chunk_progress._status_callback:
+        log_chunk_progress._status_callback(clean_status)
+
+    # Also log to file for debugging
+    realtime_log = f"{actual_realtime:.2f}x" if actual_realtime > 0 else "N/A"
+    logging.info(f"Progress: Chunk {i+1}/{total_chunks}, Elapsed: {fmt(elapsed)}, "
+                f"ETA: {fmt(eta)}, Realtime: {realtime_log}, "
+                f"Audio Duration: {fmt(total_audio_duration)}, VRAM: {allocated:.1f}GB")
+
+def display_batch_progress(batch_start, batch_end, total_chunks):
+    """Display batch processing progress"""
+    batch_progress = (batch_end / total_chunks) * 100
+    print(f"\n📊 Batch Progress: {batch_start+1}-{batch_end}/{total_chunks} ({batch_progress:.1f}%)")
+
+def display_final_summary(elapsed_time, audio_duration, chunk_count, realtime_factor):
+    """Display final processing summary"""
+    elapsed_td = timedelta(seconds=int(elapsed_time))
+    audio_td = timedelta(seconds=int(audio_duration))
+
+    print(f"\n🎉 {GREEN}Processing Complete!{RESET}")
+    print(f"📊 Final Statistics:")
+    print(f"   ⏱️ Processing Time: {CYAN}{elapsed_td}{RESET}")
+    print(f"   🎵 Audio Duration: {GREEN}{audio_td}{RESET}")
+    print(f"   📦 Total Chunks: {YELLOW}{chunk_count}{RESET}")
+    print(f"   🚀 Realtime Factor: {BOLD}{realtime_factor:.2f}x{RESET}")
+    print(f"   💾 Memory Efficiency: {GREEN}Optimized{RESET}")
+
+# ============================================================================
+# VRAM AND PERFORMANCE MONITORING
+# ============================================================================
+
+def monitor_vram_usage(operation_name=""):
+    """Real-time VRAM monitoring with threshold warnings"""
+    import torch
+
+    if not torch.cuda.is_available():
+        return 0, 0
+
+    allocated = torch.cuda.memory_allocated() / 1024**3
+    reserved = torch.cuda.memory_reserved() / 1024**3
+
+    if allocated > VRAM_SAFETY_THRESHOLD:
+        logging.warning(f"⚠️ High VRAM usage during {operation_name}: {allocated:.1f}GB allocated, {reserved:.1f}GB reserved")
+        # Trigger memory optimization if available
+        optimize_memory_if_needed()
+
+    return allocated, reserved
+
+def monitor_gpu_utilization():
+    """Monitor GPU utilization if pynvml is available"""
+    try:
+        import pynvml
+        pynvml.nvmlInit()
+        handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+        util = pynvml.nvmlDeviceGetUtilizationRates(handle)
+        temp = pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
+
+        return {
+            "gpu_util": util.gpu,
+            "memory_util": util.memory,
+            "temperature": temp
+        }
+    except:
+        return {"gpu_util": "N/A", "memory_util": "N/A", "temperature": "N/A"}
+
+def optimize_memory_if_needed():
+    """Trigger memory optimization when thresholds are exceeded"""
+    try:
+        # Try to use the enhanced CUDA memory optimization if available
+        from modules.tts_engine import optimize_cuda_memory_usage
+        optimize_cuda_memory_usage()
+    except ImportError:
+        # Fallback to basic optimization
+        import torch
+        import gc
+        torch.cuda.empty_cache()
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.ipc_collect()
+
+def display_system_info():
+    """Display system information at startup"""
+    import torch
+
+    print(f"\n🖥️ {CYAN}System Information:{RESET}")
+
+    # CUDA info
+    if torch.cuda.is_available():
+        gpu_name = torch.cuda.get_device_name(0)
+        total_vram = torch.cuda.get_device_properties(0).total_memory / 1024**3
+        print(f"   GPU: {GREEN}{gpu_name}{RESET}")
+        print(f"   VRAM: {GREEN}{total_vram:.1f}GB{RESET}")
+        print(f"   CUDA Version: {GREEN}{torch.version.cuda}{RESET}")
+    else:
+        print(f"   GPU: {RED}Not Available{RESET}")
+
+    # Memory threshold
+    print(f"   VRAM Safety Threshold: {YELLOW}{VRAM_SAFETY_THRESHOLD}GB{RESET}")
+
+    # Worker configuration
+    print(f"   Max Workers: {YELLOW}{MAX_WORKERS}{RESET}")
+    print(f"   Dynamic Workers: {YELLOW}{USE_DYNAMIC_WORKERS}{RESET}")
+
+# ============================================================================
+# PERFORMANCE TRACKING
+# ============================================================================
+
+class PerformanceTracker:
+    """Track performance metrics throughout processing"""
+
+    def __init__(self):
+        self.start_time = time.time()
+        self.chunk_times = []
+        self.vram_usage = []
+        self.batch_times = []
+
+    def log_chunk_completion(self, chunk_index, audio_duration):
+        """Log individual chunk completion"""
+        current_time = time.time()
+        chunk_time = current_time - (self.start_time + sum(self.chunk_times))
+
+        self.chunk_times.append(chunk_time)
+
+        # Track VRAM
+        allocated, reserved = monitor_vram_usage()
+        self.vram_usage.append((chunk_index, allocated, reserved))
+
+    def log_batch_completion(self, batch_size):
+        """Log batch completion"""
+        if len(self.chunk_times) >= batch_size:
+            batch_time = sum(self.chunk_times[-batch_size:])
+            self.batch_times.append(batch_time)
+
+    def get_performance_summary(self):
+        """Get comprehensive performance summary"""
+        total_time = time.time() - self.start_time
+        avg_chunk_time = sum(self.chunk_times) / len(self.chunk_times) if self.chunk_times else 0
+
+        vram_peak = max([usage[1] for usage in self.vram_usage]) if self.vram_usage else 0
+        vram_avg = sum([usage[1] for usage in self.vram_usage]) / len(self.vram_usage) if self.vram_usage else 0
+
+        return {
+            "total_time": total_time,
+            "avg_chunk_time": avg_chunk_time,
+            "total_chunks": len(self.chunk_times),
+            "vram_peak": vram_peak,
+            "vram_average": vram_avg,
+            "batch_count": len(self.batch_times)
+        }
+
+# ============================================================================
+# ERROR AND WARNING TRACKING
+# ============================================================================
+
+def log_processing_error(chunk_id, error_message, error_type="GENERAL"):
+    """Log processing errors with categorization"""
+    timestamp = time.strftime('%Y-%m-%d %H:%M:%S')
+    error_log = f"[{timestamp}] {error_type} ERROR - Chunk {chunk_id}: {error_message}"
+
+    logging.error(error_log)
+    print(f"{RED}❌ Error in chunk {chunk_id}: {error_message}{RESET}")
+
+def log_processing_warning(chunk_id, warning_message, warning_type="GENERAL"):
+    """Log processing warnings with categorization"""
+    timestamp = time.strftime('%Y-%m-%d %H:%M:%S')
+    warning_log = f"[{timestamp}] {warning_type} WARNING - Chunk {chunk_id}: {warning_message}"
+
+    logging.warning(warning_log)
+    print(f"{YELLOW}⚠️ Warning in chunk {chunk_id}: {warning_message}{RESET}")
+
+# ============================================================================
+# REAL-TIME STATUS DISPLAY
+# ============================================================================
+
+def create_status_line(current_chunk, total_chunks, elapsed_time, realtime_factor, vram_usage):
+    """Create a single-line status for real-time updates"""
+    progress_percent = (current_chunk / total_chunks) * 100
+    elapsed_str = str(timedelta(seconds=int(elapsed_time)))
+
+    status = (f"🔄 {current_chunk}/{total_chunks} ({progress_percent:.1f}%) | "
+             f"⏱️ {elapsed_str} | 🚀 {realtime_factor:.2f}x | 💾 {vram_usage:.1f}GB")
+
+    return status
+
+def update_status_line(status_message):
+    """Update status line in place"""
+    print(f"\r{status_message}", end='', flush=True)
+
+# ============================================================================
+# EXPORT FUNCTIONS
+# ============================================================================
+
+def export_performance_report(output_dir, performance_data):
+    """Export detailed performance report"""
+    report_path = output_dir / "performance_report.txt"
+
+    with open(report_path, 'w', encoding='utf-8') as f:
+        f.write("GenTTS Performance Report\n")
+        f.write("=" * 50 + "\n\n")
+
+        f.write(f"Processing Summary:\n")
+        f.write(f"  Total Processing Time: {timedelta(seconds=int(performance_data['total_time']))}\n")
+        f.write(f"  Average Chunk Time: {performance_data['avg_chunk_time']:.2f}s\n")
+        f.write(f"  Total Chunks Processed: {performance_data['total_chunks']}\n")
+        f.write(f"  Peak VRAM Usage: {performance_data['vram_peak']:.2f}GB\n")
+        f.write(f"  Average VRAM Usage: {performance_data['vram_average']:.2f}GB\n")
+        f.write(f"  Batch Count: {performance_data['batch_count']}\n")
+
+    return report_path

HF_Deploy/modules/resume_handler.py

@@ -0,0 +1,596 @@
+"""
+Resume Handler Module
+Handles resume functionality for interrupted processing
+"""
+
+import torch
+import time
+import logging
+from datetime import timedelta
+from pathlib import Path
+
+from config.config import *
+from modules.text_processor import smart_punctuate, sentence_chunk_text
+from modules.file_manager import (
+    setup_book_directories, find_book_files, list_voice_samples,
+    ensure_voice_sample_compatibility, get_audio_files_in_directory,
+    combine_audio_chunks, convert_to_m4b, add_metadata_to_m4b
+)
+from modules.audio_processor import get_chunk_audio_duration, pause_for_chunk_review
+from modules.progress_tracker import setup_logging, log_chunk_progress, log_run
+
+def analyze_existing_chunks(audio_chunks_dir):
+    """Analyze existing chunks to determine resume point"""
+    if not audio_chunks_dir.exists():
+        return 0, []
+
+    chunk_paths = get_audio_files_in_directory(audio_chunks_dir)
+
+    if not chunk_paths:
+        return 0, []
+
+    # Find the highest chunk number
+    chunk_numbers = []
+    for chunk_path in chunk_paths:
+        import re
+        match = re.match(r"chunk_(\d+)\.wav", chunk_path.name)
+        if match:
+            chunk_numbers.append(int(match.group(1)))
+
+    if not chunk_numbers:
+        return 0, []
+
+    chunk_numbers.sort()
+    last_chunk_number = max(chunk_numbers)
+
+    # Check for gaps in sequence
+    missing_chunks = []
+    for i in range(1, last_chunk_number + 1):
+        if i not in chunk_numbers:
+            missing_chunks.append(i)
+
+    print(f"📊 Existing chunks analysis:")
+    print(f"   Total chunks found: {GREEN}{len(chunk_numbers)}{RESET}")
+    print(f"   Highest chunk number: {GREEN}{last_chunk_number}{RESET}")
+    if missing_chunks:
+        print(f"   Missing chunks: {YELLOW}{len(missing_chunks)}{RESET}")
+        if len(missing_chunks) <= 10:
+            print(f"   Missing: {missing_chunks}")
+        else:
+            print(f"   Missing: {missing_chunks[:10]}... (+{len(missing_chunks)-10} more)")
+
+    return last_chunk_number, missing_chunks
+
+def suggest_resume_point(last_chunk, missing_chunks):
+    """Suggest optimal resume point based on existing chunks"""
+    if not missing_chunks:
+        # No gaps, can resume from next chunk
+        return last_chunk + 1
+
+    # If there are missing chunks, suggest resuming from first missing
+    first_missing = min(missing_chunks)
+
+    print(f"\n💡 Resume suggestions:")
+    print(f"   Resume from chunk {GREEN}{last_chunk + 1}{RESET} (continue from last)")
+    print(f"   Resume from chunk {YELLOW}{first_missing}{RESET} (fill gaps first)")
+
+    return first_missing
+
+def validate_resume_point(start_chunk, total_expected_chunks):
+    """Validate that resume point makes sense"""
+    if start_chunk < 1:
+        print(f"{RED}❌ Invalid resume point: {start_chunk}. Must be >= 1{RESET}")
+        return False
+
+    if start_chunk > total_expected_chunks:
+        print(f"{RED}❌ Resume point {start_chunk} exceeds expected total chunks {total_expected_chunks}{RESET}")
+        return False
+
+    return True
+
+def process_book_folder_resume(book_dir, voice_path, tts_params, device, start_chunk=1):
+    """Enhanced book processing with resume capability"""
+    from modules.tts_engine import process_one_chunk, load_optimized_model, get_optimal_workers
+    from src.chatterbox.tts import punc_norm
+    from concurrent.futures import ThreadPoolExecutor, as_completed
+
+    # Setup directories
+    output_root, tts_dir, text_chunks_dir, audio_chunks_dir = setup_book_directories(book_dir)
+
+    # Find book files
+    book_files = find_book_files(book_dir)
+    text_file = book_files['text']
+    cover_file = book_files['cover']
+    nfo_file = book_files['nfo']
+
+    if not text_file:
+        logging.info(f"[{book_dir.name}] ERROR: No .txt files found in the book folder.")
+        return None, None, []
+    
+    text_files = [text_file]  # Convert to list for compatibility
+
+    # IMPORTANT: Don't delete existing directories if resuming
+    print(f"🔍 DEBUG: start_chunk = {start_chunk}")
+    if start_chunk == 1:
+        print(f"⚠️ WARNING: start_chunk is 1 - this will clear existing chunks!")
+        print(f"📁 About to clear: {audio_chunks_dir}")
+        
+        # Only clear on fresh start
+        import shutil
+        for d in [text_chunks_dir, audio_chunks_dir]:
+            if d.exists() and d.is_dir():
+                print(f"🗑️ CLEARING DIRECTORY: {d}")
+                shutil.rmtree(d)
+
+        for d in [output_root, tts_dir, text_chunks_dir, audio_chunks_dir]:
+            d.mkdir(parents=True, exist_ok=True)
+    else:
+        print(f"✅ RESUME MODE: Preserving existing chunks in {audio_chunks_dir}")
+        # Ensure directories exist for resume
+        for d in [output_root, tts_dir, text_chunks_dir, audio_chunks_dir]:
+            d.mkdir(parents=True, exist_ok=True)
+
+    setup_logging(output_root)
+
+    # Load existing chunks from JSON (resume should use preprocessed data)
+    from modules.tts_engine import find_chunks_json_file
+    
+    json_file = find_chunks_json_file(book_dir.name)
+    if json_file:
+        print(f"📖 Loading preprocessed chunks from: {json_file.name}")
+        from wrapper.chunk_loader import load_chunks
+        all_chunks = load_chunks(str(json_file))
+        print(f"✅ Loaded {len(all_chunks)} chunks with metadata")
+    else:
+        print(f"❌ No preprocessed chunks found for {book_dir.name}")
+        print(f"💡 Use Option 1 to process this book from the beginning first.")
+        return None, None, []
+
+    # Validate resume point
+    if not validate_resume_point(start_chunk, len(all_chunks)):
+        return None, None, []
+
+    # Filter chunks to process (resume logic)
+    if start_chunk > 1:
+        print(f"🔄 Resuming from chunk {start_chunk}")
+        print(f"📊 Skipping chunks 1-{start_chunk-1} (already completed)")
+
+        # Check which chunks already exist
+        existing_chunks = []
+        for i in range(start_chunk-1):
+            chunk_path = audio_chunks_dir / f"chunk_{i+1:05}.wav"
+            if chunk_path.exists():
+                existing_chunks.append(i+1)
+
+        print(f"✅ Found {len(existing_chunks)} existing chunks")
+
+        # Only process remaining chunks
+        chunks_to_process = all_chunks[start_chunk-1:]
+        chunk_offset = start_chunk - 1
+    else:
+        chunks_to_process = all_chunks
+        chunk_offset = 0
+
+    run_log_lines = [
+        f"\n===== RESUME Processing: {book_dir.name} =====",
+        f"Voice: {voice_path.name}",
+        f"Started: {time.strftime('%Y-%m-%d %H:%M:%S')}",
+        f"Resume from chunk: {start_chunk}",
+        f"Text files processed: {len(text_files)}",
+        f"Total chunks generated: {len(all_chunks)}",
+        f"Chunks to process: {len(chunks_to_process)}"
+    ]
+
+    # Write initial run info immediately
+    initial_log = run_log_lines + [
+        f"--- Generation Settings ---",
+        f"Batch Processing: Enabled ({BATCH_SIZE} chunks per batch)",
+        f"ASR Enabled: {ENABLE_ASR}",
+        f"Hum Detection: {ENABLE_HUM_DETECTION}",
+        f"Dynamic Workers: {USE_DYNAMIC_WORKERS}",
+        f"Voice used: {voice_path.name}",
+        f"Exaggeration: {tts_params['exaggeration']}",
+        f"CFG weight: {tts_params['cfg_weight']}",
+        f"Temperature: {tts_params['temperature']}",
+        f"Processing Status: IN PROGRESS...",
+        f"="*50
+    ]
+
+    log_run("\n".join(initial_log), output_root / "run.log")
+    print(f"📝 Initial run info written to: {output_root / 'run.log'}")
+
+    start_time = time.time()
+    total_chunks = len(all_chunks)
+    remaining_chunks = len(chunks_to_process)
+    log_path = output_root / "chunk_validation.log"
+
+    # Calculate existing audio duration for accurate progress
+    total_audio_duration = 0.0
+    if start_chunk > 1:
+        print("📊 Calculating existing audio duration...")
+        for i in range(start_chunk-1):
+            chunk_path = audio_chunks_dir / f"chunk_{i+1:05}.wav"
+            if chunk_path.exists():
+                total_audio_duration += get_chunk_audio_duration(chunk_path)
+        print(f"📊 Existing audio: {timedelta(seconds=int(total_audio_duration))}")
+
+    # Initialize performance optimizations
+    from modules.tts_engine import detect_deployment_environment, enable_gpu_persistence_mode
+    deployment_env = detect_deployment_environment()
+    print(f"🌍 Deployment environment: {deployment_env}")
+    
+    # Enable GPU persistence mode for better performance
+    gpu_persistence_enabled = enable_gpu_persistence_mode()
+
+    # Batch processing for remaining chunks
+    print(f"📊 Processing {remaining_chunks} remaining chunks in batches of {BATCH_SIZE}")
+
+    all_results = []
+
+    for batch_start in range(0, remaining_chunks, BATCH_SIZE):
+        batch_end = min(batch_start + BATCH_SIZE, remaining_chunks)
+        batch_chunks = chunks_to_process[batch_start:batch_end]
+
+        actual_start_chunk = chunk_offset + batch_start + 1
+        actual_end_chunk = chunk_offset + batch_end
+
+        print(f"\n🔄 Processing batch: chunks {actual_start_chunk}-{actual_end_chunk}")
+
+        # Fresh model for each batch
+        model = load_optimized_model(device)
+        compatible_voice = ensure_voice_sample_compatibility(voice_path, output_dir=tts_dir)
+        
+        # Pre-warm model to eliminate first chunk quality variations
+        from modules.tts_engine import prewarm_model_with_voice
+        model = prewarm_model_with_voice(model, compatible_voice, tts_params)
+
+        # Load ASR model once per batch if needed using adaptive manager
+        asr_model = None
+        asr_device_used = None
+        if ENABLE_ASR:
+            from modules.asr_manager import load_asr_model_adaptive
+            print(f"🎤 Loading ASR model for resume mode...")
+            # Resume mode uses fallback config (no intelligent selection)
+            asr_model, asr_device_used = load_asr_model_adaptive()
+
+        futures = []
+        batch_results = []
+
+        # Dynamic worker allocation
+        optimal_workers = get_optimal_workers()
+        print(f"🔧 Using {optimal_workers} workers for batch {actual_start_chunk}-{actual_end_chunk}")
+
+        # Try producer-consumer pipeline first (Phase 4 optimization)
+        batch_results = []
+        if ENABLE_PRODUCER_CONSUMER_PIPELINE:
+            try:
+                print(f"🚀 Attempting producer-consumer pipeline for resume batch {actual_start_chunk}-{actual_end_chunk}")
+                from modules.tts_engine import process_chunks_with_pipeline
+                pipeline_results = process_chunks_with_pipeline(
+                    all_chunks, batch_chunks, chunk_offset, text_chunks_dir, audio_chunks_dir,
+                    voice_path, tts_params, start_time, total_chunks, punc_norm, book_dir.name,
+                    log_run, log_path, device, model, asr_model, True, optimal_workers,  # asr_enabled=True for resume
+                    total_audio_duration  # Pass accumulated duration for proper ETA calculation
+                )
+                
+                # Handle tuple return from pipeline
+                if isinstance(pipeline_results, tuple) and len(pipeline_results) == 2:
+                    batch_results, batch_audio_duration = pipeline_results
+                    total_audio_duration += batch_audio_duration
+                else:
+                    # Fallback for old return format
+                    batch_results = pipeline_results
+                
+                if batch_results:
+                    print(f"✅ Producer-consumer pipeline completed resume batch: {len(batch_results)} chunks")
+                    # Pipeline already handled progress logging internally
+                
+            except Exception as e:
+                logging.error(f"❌ Producer-consumer pipeline failed in resume: {e}")
+                if not ENABLE_PIPELINE_FALLBACK:
+                    raise
+                batch_results = []  # Clear failed results
+        
+        # Fallback to original sequential processing if pipeline disabled or failed
+        if not batch_results:
+            print(f"🔄 Using sequential processing fallback for resume batch {actual_start_chunk}-{actual_end_chunk}")
+            futures = []
+
+            with ThreadPoolExecutor(max_workers=optimal_workers) as executor:
+                for i, chunk_data in enumerate(batch_chunks):
+                    global_chunk_index = chunk_offset + i
+
+                    # Check for shutdown request
+                    if shutdown_requested:
+                        print(f"\n⏹️ {YELLOW}Stopping submission of new chunks...{RESET}")
+                        break
+
+                    chunk = chunk_data["text"]
+                    all_chunk_texts = [cd["text"] for cd in all_chunks]
+                    boundary_type = chunk_data.get("boundary_type", "none")
+
+                    futures.append(executor.submit(
+                        process_one_chunk,
+                        global_chunk_index, chunk, text_chunks_dir, audio_chunks_dir,
+                        voice_path, tts_params, start_time, total_chunks,
+                        punc_norm, book_dir.name, log_run, log_path, device,
+                        model, asr_model, all_chunk_texts, boundary_type
+                    ))
+
+                # Wait for batch to complete
+                print(f"🔄 {CYAN}Waiting for batch {actual_start_chunk}-{actual_end_chunk} to complete...{RESET}")
+                completed_count = 0
+
+                for fut in as_completed(futures):
+                    try:
+                        idx, wav_path = fut.result()
+                        if wav_path and wav_path.exists():
+                            # Measure actual audio duration for this chunk
+                            chunk_duration = get_chunk_audio_duration(wav_path)
+                            total_audio_duration += chunk_duration
+                            batch_results.append((idx, wav_path))
+
+                            # Update progress every 10 chunks within batch
+                            completed_count += 1
+                            if completed_count % 10 == 0:
+                                current_chunk = chunk_offset + completed_count
+                                log_chunk_progress(current_chunk - 1, total_chunks, start_time, total_audio_duration)
+
+                    except Exception as e:
+                        logging.error(f"Future failed in batch: {e}")
+
+        # Clean up model after batch
+        print(f"🧹 Cleaning up after batch {actual_start_chunk}-{actual_end_chunk}")
+        del model
+        if asr_model:
+            from modules.asr_manager import cleanup_asr_model
+            cleanup_asr_model(asr_model)
+        torch.cuda.empty_cache()
+        import gc
+        gc.collect()
+        time.sleep(2)
+
+        all_results.extend(batch_results)
+        print(f"✅ Batch {actual_start_chunk}-{actual_end_chunk} completed ({len(batch_results)} chunks)")
+
+    # Final processing - combine ALL chunks (existing + new)
+    quarantine_dir = audio_chunks_dir / "quarantine"
+    pause_for_chunk_review(quarantine_dir)
+
+    # Collect ALL chunk paths (both existing and newly created)
+    chunk_paths = []
+    for i in range(total_chunks):
+        chunk_path = audio_chunks_dir / f"chunk_{i+1:05}.wav"
+        if chunk_path.exists():
+            chunk_paths.append(chunk_path)
+        else:
+            logging.warning(f"Missing chunk file: chunk_{i+1:05}.wav")
+
+    if not chunk_paths:
+        logging.info(f"{RED}❌ No valid audio chunks found. Skipping concatenation and conversion.{RESET}")
+        return None, None, []
+
+    print(f"📊 Found {len(chunk_paths)} total chunks for final audiobook")
+
+    # Calculate timing
+    elapsed_total = time.time() - start_time
+    elapsed_td = timedelta(seconds=int(elapsed_total))
+
+    # Get total audio duration from ALL chunks
+    total_audio_duration_final = sum(get_chunk_audio_duration(chunk_path) for chunk_path in chunk_paths)
+    audio_duration_td = timedelta(seconds=int(total_audio_duration_final))
+    realtime_factor = total_audio_duration_final / elapsed_total if elapsed_total > 0 else 0.0
+
+    print(f"\n⏱️ Resume Processing Complete:")
+    print(f"   Elapsed Time: {CYAN}{str(elapsed_td)}{RESET}")
+    print(f"   Audio Duration: {GREEN}{str(audio_duration_td)}{RESET}")
+    print(f"   Realtime Factor: {YELLOW}{realtime_factor:.2f}x{RESET}")
+
+    # Combine audio
+    combined_wav_path = output_root / f"{book_dir.name} [{voice_path.stem}].wav"
+    print("\n💾 Saving WAV file...")
+    combine_audio_chunks(chunk_paths, combined_wav_path)
+
+    # M4B conversion
+    temp_m4b_path = output_root / "output.m4b"
+    final_m4b_path = output_root / f"{book_dir.name}[{voice_path.stem}].m4b"
+    convert_to_m4b(combined_wav_path, temp_m4b_path)
+    add_metadata_to_m4b(temp_m4b_path, final_m4b_path, cover_file, nfo_file)
+
+    logging.info(f"Audiobook created: {final_m4b_path}")
+
+    # Append final completion info
+    completion_log = [
+        f"\n--- Resume Processing Complete ---",
+        f"Completed: {time.strftime('%Y-%m-%d %H:%M:%S')}",
+        f"Processing Time: {str(elapsed_td)}",
+        f"Audio Duration: {str(audio_duration_td)}",
+        f"Realtime Factor: {realtime_factor:.2f}x",
+        f"Total Chunks: {len(chunk_paths)}",
+        f"Combined WAV: {combined_wav_path}",
+        f"Final M4B: {final_m4b_path}"
+    ]
+
+    # Append to existing log
+    log_run("\n".join(completion_log), output_root / "run.log")
+    print(f"📝 Final completion info appended to: {output_root / 'run.log'}")
+
+    return final_m4b_path, combined_wav_path, run_log_lines
+
+def resume_book_from_chunk(start_chunk):
+    """Interactive resume function for stuck book"""
+    print(f"\n🔄 Resume Book Processing from Chunk {start_chunk}")
+    print("=" * 50)
+
+    # Show available books from Audiobook directory (books that have started processing)
+    audiobook_root = Path(AUDIOBOOK_ROOT)
+    if not audiobook_root.exists():
+        print(f"{RED}No audiobook directory found at {AUDIOBOOK_ROOT}.{RESET}")
+        return None
+        
+    book_dirs = sorted([d for d in audiobook_root.iterdir() if d.is_dir() and d.name != "Audio_Revisions"])
+    if not book_dirs:
+        print(f"{RED}No books found in {AUDIOBOOK_ROOT}/ - no books have started processing.{RESET}")
+        print(f"💡 Use Option 1 to start processing a new book first.")
+        return None
+
+    print("Available books (in progress or completed):")
+    for i, book_dir in enumerate(book_dirs):
+        # All books in Audiobook/ should have processing data
+        audio_chunks_dir = book_dir / "TTS" / "audio_chunks"
+        if audio_chunks_dir.exists():
+            last_chunk, missing = analyze_existing_chunks(audio_chunks_dir)
+            if missing:
+                status = f"(last chunk: {last_chunk}, {len(missing)} missing)"
+            else:
+                status = f"(completed: {last_chunk} chunks)"
+        else:
+            status = "(processing started but no chunks yet)"
+
+        print(f"  [{i}] {book_dir.name} {status}")
+
+    while True:
+        try:
+            book_idx = int(input("Select book index: "))
+            if 0 <= book_idx < len(book_dirs):
+                audiobook_dir = book_dirs[book_idx]
+                # Find corresponding Text_Input directory
+                text_input_book_dir = TEXT_INPUT_ROOT / audiobook_dir.name
+                if text_input_book_dir.exists():
+                    book_dir = text_input_book_dir
+                else:
+                    print(f"❌ Text_Input directory not found for {audiobook_dir.name}")
+                    print(f"💡 The original book files may have been moved or deleted.")
+                    continue
+                break
+        except Exception:
+            pass
+        print("Invalid selection. Try again.")
+
+    # Analyze existing chunks for selected book
+    audiobook_dir = AUDIOBOOK_ROOT / book_dir.name
+    if audiobook_dir.exists():
+        audio_chunks_dir = audiobook_dir / "TTS" / "audio_chunks"
+        if audio_chunks_dir.exists():
+            last_chunk, missing = analyze_existing_chunks(audio_chunks_dir)
+            suggested_resume = suggest_resume_point(last_chunk, missing)
+
+            print(f"\nSuggested resume point: {GREEN}{suggested_resume}{RESET}")
+
+            # Allow user to override
+            user_input = input(f"Resume from chunk [{suggested_resume}]: ").strip()
+            if user_input:
+                try:
+                    start_chunk = int(user_input)
+                except ValueError:
+                    print(f"Invalid input, using suggested: {suggested_resume}")
+                    start_chunk = suggested_resume
+            else:
+                start_chunk = suggested_resume
+
+    # Show available voices
+    voice_files = list_voice_samples()
+    if not voice_files:
+        print(f"{RED}No voice samples found.{RESET}")
+        return None
+
+    print("\nAvailable voices:")
+    for i, voice in enumerate(voice_files):
+        print(f"  [{i}] {voice.name}")
+
+    while True:
+        try:
+            voice_idx = int(input("Select voice index: "))
+            if 0 <= voice_idx < len(voice_files):
+                voice_path = voice_files[voice_idx]
+                break
+        except Exception:
+            pass
+        print("Invalid selection. Try again.")
+
+    # Get TTS parameters
+    def prompt_float(prompt, default):
+        val = input(f"{prompt} [{default}]: ").strip()
+        return float(val) if val else default
+
+    exaggeration = prompt_float("Enter exaggeration (emotion intensity)", DEFAULT_EXAGGERATION)
+    cfg_weight = prompt_float("Enter cfg_weight (faithfulness to text)", DEFAULT_CFG_WEIGHT)
+    temperature = prompt_float("Enter temperature (randomness)", DEFAULT_TEMPERATURE)
+
+    tts_params = dict(exaggeration=exaggeration, cfg_weight=cfg_weight, temperature=temperature)
+
+    # Determine device with proper validation
+    from modules.tts_engine import get_best_available_device
+    device = get_best_available_device()
+
+    print(f"\n🚀 Resuming {book_dir.name} from chunk {start_chunk}")
+    print(f"🎤 Voice: {voice_path.name}")
+    print(f"⚙️ Parameters: {tts_params}")
+
+    # Process with resume
+    return process_book_folder_resume(book_dir, voice_path, tts_params, device, start_chunk)
+
+def find_incomplete_books():
+    """Find books that appear to be incomplete"""
+    incomplete_books = []
+
+    for book_dir in TEXT_INPUT_ROOT.iterdir():
+        if not book_dir.is_dir():
+            continue
+
+        audiobook_dir = AUDIOBOOK_ROOT / book_dir.name
+        if not audiobook_dir.exists():
+            continue
+
+        audio_chunks_dir = audiobook_dir / "TTS" / "audio_chunks"
+        if not audio_chunks_dir.exists():
+            continue
+
+        # Check if there's a final M4B
+        m4b_files = list(audiobook_dir.glob("*.m4b"))
+        wav_files = list(audiobook_dir.glob("*.wav"))
+
+        if not m4b_files and not wav_files:
+            # No final output, likely incomplete
+            last_chunk, missing = analyze_existing_chunks(audio_chunks_dir)
+            if last_chunk > 0:
+                incomplete_books.append({
+                    "name": book_dir.name,
+                    "last_chunk": last_chunk,
+                    "missing_chunks": len(missing),
+                    "path": book_dir
+                })
+
+    return incomplete_books
+
+def auto_resume_incomplete():
+    """Automatically suggest resume for incomplete books"""
+    incomplete = find_incomplete_books()
+
+    if not incomplete:
+        print(f"{GREEN}✅ No incomplete books found!{RESET}")
+        return
+
+    print(f"{YELLOW}📋 Found {len(incomplete)} incomplete books:{RESET}")
+    for i, book in enumerate(incomplete):
+        print(f"  [{i}] {book['name']} (last chunk: {book['last_chunk']}, missing: {book['missing_chunks']})")
+
+    choice = input(f"\nSelect book to resume [0-{len(incomplete)-1}] or 'q' to quit: ").strip()
+
+    if choice.lower() == 'q':
+        return
+
+    try:
+        idx = int(choice)
+        if 0 <= idx < len(incomplete):
+            selected_book = incomplete[idx]
+            suggested_resume = selected_book['last_chunk'] + 1
+
+            print(f"\n🎯 Selected: {selected_book['name']}")
+            print(f"💡 Suggested resume point: chunk {suggested_resume}")
+
+            return resume_book_from_chunk(suggested_resume)
+    except ValueError:
+        print("Invalid selection.")
+
+    return None

HF_Deploy/modules/system_detector.py

@@ -0,0 +1,231 @@
+"""
+System Resource Detection Module
+Detects VRAM, RAM, CPU cores and recommends appropriate ASR models
+"""
+
+import psutil
+import torch
+import os
+import sys
+from pathlib import Path
+
+# Add project root to path for imports
+if __name__ == "__main__":
+    sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from config.config import ASR_MODEL_VRAM_MB, ASR_MODEL_RAM_MB
+
+def get_gpu_memory():
+    """Get total and available GPU memory in MB"""
+    try:
+        if torch.cuda.is_available():
+            gpu_count = torch.cuda.device_count()
+            if gpu_count > 0:
+                # Use first GPU
+                total_vram = torch.cuda.get_device_properties(0).total_memory
+                allocated_vram = torch.cuda.memory_allocated(0)
+                available_vram = total_vram - allocated_vram
+                
+                return {
+                    'total_mb': total_vram // 1024 // 1024,
+                    'available_mb': available_vram // 1024 // 1024,
+                    'allocated_mb': allocated_vram // 1024 // 1024
+                }
+    except:
+        pass
+    
+    return {'total_mb': 0, 'available_mb': 0, 'allocated_mb': 0}
+
+def get_system_memory():
+    """Get total and available system RAM in MB"""
+    try:
+        memory = psutil.virtual_memory()
+        return {
+            'total_mb': memory.total // 1024 // 1024,
+            'available_mb': memory.available // 1024 // 1024,
+            'used_mb': memory.used // 1024 // 1024
+        }
+    except:
+        return {'total_mb': 0, 'available_mb': 0, 'used_mb': 0}
+
+def get_cpu_cores():
+    """Get number of CPU cores"""
+    try:
+        return psutil.cpu_count(logical=False) or psutil.cpu_count()
+    except:
+        return 1
+
+def estimate_tts_vram_usage():
+    """Estimate VRAM usage by ChatterboxTTS (updated based on real usage)"""
+    return 5500  # 5.5GB in MB (was 7GB, adjusted based on actual 3.5GB usage + buffer)
+
+def get_system_profile():
+    """Get complete system resource profile"""
+    gpu_info = get_gpu_memory()
+    ram_info = get_system_memory()
+    cpu_cores = get_cpu_cores()
+    
+    # Estimate available resources after TTS loading
+    tts_vram_estimate = estimate_tts_vram_usage()
+    available_vram_after_tts = max(0, gpu_info['available_mb'] - tts_vram_estimate)
+    
+    return {
+        'gpu': gpu_info,
+        'ram': ram_info,
+        'cpu_cores': cpu_cores,
+        'available_vram_after_tts': available_vram_after_tts,
+        'has_gpu': gpu_info['total_mb'] > 0
+    }
+
+def categorize_system(profile):
+    """Categorize system capabilities"""
+    gpu_total = profile['gpu']['total_mb']
+    ram_total = profile['ram']['total_mb']
+    cpu_cores = profile['cpu_cores']
+    
+    # VRAM categories
+    if gpu_total < 4000:
+        vram_category = "low"
+    elif gpu_total <= 12000:
+        vram_category = "medium"
+    else:
+        vram_category = "high"
+    
+    # RAM categories  
+    if ram_total < 16000:
+        ram_category = "low"
+    elif ram_total <= 64000:
+        ram_category = "medium"
+    else:
+        ram_category = "high"
+    
+    # CPU categories
+    if cpu_cores < 6:
+        cpu_category = "low"
+    elif cpu_cores <= 16:
+        cpu_category = "medium"
+    else:
+        cpu_category = "high"
+    
+    return {
+        'vram': vram_category,
+        'ram': ram_category,
+        'cpu': cpu_category
+    }
+
+def get_safe_asr_models(profile):
+    """Get ASR models that can safely run on GPU with available VRAM"""
+    available_vram = profile['available_vram_after_tts']
+    safe_models = []
+    
+    for model, vram_req in ASR_MODEL_VRAM_MB.items():
+        if vram_req <= available_vram:
+            safe_models.append(model)
+    
+    return safe_models
+
+def get_safe_cpu_models(profile):
+    """Get ASR models that can safely run on CPU with available RAM"""
+    available_ram = profile['ram']['available_mb']
+    safe_models = []
+    
+    for model, ram_req in ASR_MODEL_RAM_MB.items():
+        if ram_req <= available_ram:
+            safe_models.append(model)
+    
+    return safe_models
+
+def recommend_asr_models(profile):
+    """Recommend Safe/Moderate/Insane ASR model configurations"""
+    categories = categorize_system(profile)
+    safe_gpu_models = get_safe_asr_models(profile)
+    safe_cpu_models = get_safe_cpu_models(profile)
+    
+    recommendations = {}
+    
+    # Model priority order (best to worst)
+    model_priority = ["large-v3", "large", "large-v2", "medium", "small", "base", "tiny"]
+    
+    # Safe: Conservative choice
+    safe_gpu = None
+    safe_cpu = None
+    
+    for model in reversed(model_priority):  # Start from smallest
+        if model in safe_gpu_models and not safe_gpu:
+            safe_gpu = model
+        if model in safe_cpu_models and not safe_cpu:
+            safe_cpu = model
+        if safe_gpu and safe_cpu:
+            break
+    
+    # Moderate: Balanced choice  
+    moderate_gpu = None
+    moderate_cpu = None
+    
+    # Try to get a model 1-2 steps up from safe
+    safe_idx = model_priority.index(safe_gpu) if safe_gpu else len(model_priority)
+    moderate_idx = max(0, safe_idx - 2)
+    
+    for i in range(moderate_idx, len(model_priority)):
+        model = model_priority[i]
+        if model in safe_gpu_models and not moderate_gpu:
+            moderate_gpu = model
+        if model in safe_cpu_models and not moderate_cpu:
+            moderate_cpu = model
+        if moderate_gpu and moderate_cpu:
+            break
+    
+    # Insane: Push the limits (best available models)
+    insane_gpu = None
+    insane_cpu = None
+    
+    # Get the best (largest) models that are safe
+    for model in model_priority:  # Start from best
+        if model in safe_gpu_models and not insane_gpu:
+            insane_gpu = model
+        if model in safe_cpu_models and not insane_cpu:
+            insane_cpu = model
+        if insane_gpu and insane_cpu:
+            break
+    
+    # Build recommendations
+    recommendations['safe'] = {
+        'primary': {'model': safe_gpu or safe_cpu, 'device': 'gpu' if safe_gpu else 'cpu'},
+        'fallback': {'model': safe_cpu, 'device': 'cpu'}
+    }
+    
+    recommendations['moderate'] = {
+        'primary': {'model': moderate_gpu or moderate_cpu, 'device': 'gpu' if moderate_gpu else 'cpu'},
+        'fallback': {'model': moderate_cpu, 'device': 'cpu'}
+    }
+    
+    recommendations['insane'] = {
+        'primary': {'model': insane_gpu or insane_cpu, 'device': 'gpu' if insane_gpu else 'cpu'},
+        'fallback': {'model': insane_cpu, 'device': 'cpu'}
+    }
+    
+    return recommendations
+
+def print_system_summary(profile):
+    """Print a human-readable system summary"""
+    categories = categorize_system(profile)
+    
+    print(f"🖥️ System Profile:")
+    print(f"   VRAM: {profile['gpu']['total_mb']:,}MB total, {profile['available_vram_after_tts']:,}MB available after TTS ({categories['vram']} class)")
+    print(f"   RAM:  {profile['ram']['total_mb']:,}MB total, {profile['ram']['available_mb']:,}MB available ({categories['ram']} class)")
+    print(f"   CPU:  {profile['cpu_cores']} cores ({categories['cpu']} class)")
+    
+    if not profile['has_gpu']:
+        print(f"   ⚠️ No CUDA GPU detected - ASR will run on CPU only")
+
+if __name__ == "__main__":
+    # Test the detection
+    profile = get_system_profile()
+    print_system_summary(profile)
+    
+    recommendations = recommend_asr_models(profile)
+    print(f"\nASR Model Recommendations:")
+    for level, config in recommendations.items():
+        primary = config['primary']
+        fallback = config['fallback']
+        print(f"🟢 {level.upper()}: {primary['model']} ({primary['device']}) + {fallback['model']} (cpu fallback)")

HF_Deploy/modules/text_processor.py

@@ -0,0 +1,745 @@
+"""
+Text Processing Module
+Handles text chunking, abbreviations, and preprocessing for TTS
+"""
+
+import re
+import logging
+from pathlib import Path
+from config.config import MAX_CHUNK_WORDS, MIN_CHUNK_WORDS, YELLOW, RESET
+
+
+
+# ============================================================================
+# ABBREVIATION REPLACEMENT SYSTEM
+# ============================================================================
+
+def load_abbreviations(file_path="utils/abbreviations.txt"):
+    """Load abbreviation replacements from external file"""
+    replacements = {}
+    abbrev_file = Path(file_path)
+
+    if not abbrev_file.exists():
+        print(f"⚠️ {YELLOW}Abbreviations file not found: {file_path}{RESET}")
+        print(f"📝 Creating sample file...")
+        create_sample_abbreviations_file(abbrev_file)
+        return replacements
+
+    try:
+        with open(abbrev_file, 'r', encoding='utf-8') as f:
+            for line_num, line in enumerate(f, 1):
+                line = line.strip()
+
+                # Skip empty lines and comments
+                if not line or line.startswith('#'):
+                    continue
+
+                # Parse "abbrev -> replacement" format
+                if ' -> ' in line:
+                    abbrev, replacement = line.split(' -> ', 1)
+                    replacements[abbrev.strip()] = replacement.strip()
+                else:
+                    print(f"⚠️ Invalid format on line {line_num}: {line}")
+
+        print(f"✅ Loaded {len(replacements)} abbreviation replacements from {file_path}")
+
+    except Exception as e:
+        print(f"❌ Error loading abbreviations: {e}")
+
+    return replacements
+
+def create_sample_abbreviations_file(file_path):
+    """Create a sample abbreviations file with common replacements"""
+    sample_content = """# Abbreviation Replacements for TTS
+# Format: abbreviation -> replacement
+# Lines starting with # are comments
+
+# Common titles and abbreviations
+Dr. -> Doctor
+Mr. -> Mister
+Mrs. -> Missus
+Ms. -> Miss
+Prof. -> Professor
+Rev. -> Reverend
+Lt. -> Lieutenant
+Capt. -> Captain
+Gen. -> General
+Col. -> Colonel
+Jr. -> Junior
+Sr. -> Senior
+
+# Political and organizations
+M.P. -> MP
+U.S. -> US
+U.K. -> UK
+U.N. -> UN
+F.B.I. -> FBI
+C.I.A. -> CIA
+N.A.S.A. -> NASA
+
+# Common abbreviations
+etc. -> et cetera
+vs. -> versus
+e.g. -> for example
+i.e. -> that is
+Inc. -> Incorporated
+Corp. -> Corporation
+Ltd. -> Limited
+Co. -> Company
+
+# Numbers and ordinals
+1st -> first
+2nd -> second
+3rd -> third
+4th -> fourth
+5th -> fifth
+10th -> tenth
+20th -> twentieth
+21st -> twenty-first
+30th -> thirtieth
+40th -> fortieth
+50th -> fiftieth
+60th -> sixtieth
+70th -> seventieth
+80th -> eightieth
+90th -> ninetieth
+100th -> one hundredth
+
+# Time abbreviations
+a.m. -> AM
+p.m. -> PM
+A.M. -> AM
+P.M. -> PM
+"""
+
+    try:
+        with open(file_path, 'w', encoding='utf-8') as f:
+            f.write(sample_content)
+        print(f"📝 Created sample abbreviations file: {file_path}")
+        print(f"💡 Edit this file to add your own replacements!")
+    except Exception as e:
+        print(f"❌ Error creating sample file: {e}")
+
+def preprocess_abbreviations(text, replacements):
+    """Replace abbreviations with TTS-friendly versions"""
+    if not replacements:
+        return text
+
+    original_text = text
+    replacements_made = 0
+
+    # Apply replacements (order matters for overlapping patterns)
+    for abbrev, replacement in replacements.items():
+        if abbrev in text:
+            text = text.replace(abbrev, replacement)
+            replacements_made += 1
+
+    if replacements_made > 0:
+        logging.info(f"📝 Applied {replacements_made} abbreviation replacements")
+
+    return text
+
+# ============================================================================
+# TEXT PREPROCESSING AND CHUNKING
+# ============================================================================
+
+def smart_punctuate(text):
+    """
+    Enhanced punctuation normalization with abbreviation replacement.
+    
+    PROCESSING REQUIREMENTS:
+    - Load and apply abbreviation replacements (Dr. -> Doctor, etc.)
+    - Add periods to lines that don't end with punctuation
+    - Replace Unicode smart quotes with ASCII quotes (", ')
+    - Remove problematic formatting (bold markdown, underlines)
+    - Preserve paragraph breaks (empty lines)
+    
+    This prepares text for consistent TTS processing.
+    """
+
+    # Load abbreviations and apply them
+    abbreviation_replacements = load_abbreviations()
+    text = preprocess_abbreviations(text, abbreviation_replacements)
+
+    # Then continue with existing punctuation logic
+    lines = text.splitlines()
+    out = []
+
+    for l in lines:
+        stripped = l.strip()
+
+        # Preserve empty lines (paragraph breaks)
+        if not stripped:
+            out.append("")  # Keep the blank line
+        # Process non-empty lines
+        elif not re.search(r'[.!?]$', stripped) and not re.search(r'[.!?]["\']$', stripped):
+            out.append(stripped + ".")
+        else:
+            out.append(stripped)
+
+    result = "\n".join(out)
+
+    # Enhanced text preprocessing - replace curly quotes with straight quotes
+    result = result.replace('\u201c', '"').replace('\u201d', '"')  # Replace smart double quotes " "
+    result = result.replace('\u2018', "'").replace('\u2019', "'")  # Replace smart single quotes ' '
+
+    # Remove problematic formatting
+    result = re.sub(r'\*\*([^*]+)\*\*', r'\1', result)  # Remove bold markdown
+    result = re.sub(r'_{2,}', '', result)  # Remove underlines
+    
+    # Fix any escaped quotes that might appear in the text
+    result = result.replace('\\"', '"').replace("\\'", "'")
+    
+    # Additional quote normalization to prevent recurring dialogue corruption
+    result = re.sub(r'(["\'])\s*,\s*(["\'])', r'\1, \2', result)  # Fix quote spacing around commas
+    result = re.sub(r'(["\'])\s*\.\s*(["\'])', r'\1. \2', result)  # Fix quote spacing around periods
+    result = re.sub(r'(["\'])\s*([,.])\s*(["\'])\s*([,.])', r'\1\2 \3', result)  # Remove duplicate punctuation
+    
+    # Debug logging for dialogue patterns
+    if '"' in result and ('replied' in result or 'said' in result):
+        print(f"🗣️ DEBUG: Dialogue detected in smart_punctuate: {result[:100]}...")
+
+    return result
+
+def fix_short_sentence_artifacts(chunk_text):
+    """
+    Fix multiple short sentences that cause TTS errors.
+    Example: "Yes. No. Maybe." → "Yes, no, maybe."
+             "Right." → "Right," (if it's a single-word chunk)
+    """
+    # Handle full chunk that is just one short sentence
+    words = chunk_text.strip().split()
+    if len(words) == 1 and chunk_text.strip().endswith('.'):
+        return chunk_text.strip()[:-1] + ','  # Replace period with comma
+
+    parts = re.split(r'([.!?])', chunk_text.strip())
+    if len(parts) < 2:
+        return chunk_text  # nothing to fix
+
+    # Reconstruct sentence-punctuation pairs
+    sentences = []
+    for i in range(0, len(parts)-1, 2):
+        sentence = parts[i].strip()
+        punct = parts[i+1]
+        if sentence:
+            word_count = len(sentence.split())
+            sentences.append((sentence, punct, word_count))
+
+    # Handle multiple short sentences
+    short_count = sum(1 for _, _, wc in sentences if wc <= 3)
+
+    if short_count >= 2 and len(sentences) >= 2:
+        merged = ", ".join(s for s, _, _ in sentences) + "."
+        return merged
+
+    # Handle case where first sentence is a single word
+    if len(sentences) >= 2 and sentences[0][2] == 1 and sentences[0][1] == ".":
+        # Replace period with comma
+        first, second = sentences[0][0], sentences[1][0]
+        rest = " ".join(s for s, _, _ in sentences[2:])
+        new_text = f"{first}, {second}"
+        if rest:
+            new_text += " " + rest
+        return new_text
+
+    return chunk_text
+
+def _is_apostrophe(text, pos):
+    """Check if a single quote at position pos is likely an apostrophe (not speech quote)"""
+    if pos == 0 or pos >= len(text) - 1:
+        return False
+    
+    # Check characters before and after
+    before = text[pos - 1] if pos > 0 else ' '
+    after = text[pos + 1] if pos < len(text) - 1 else ' '
+    
+    # It's likely an apostrophe if:
+    # 1. Preceded and followed by letters (contractions like "don't", possessives like "John's")
+    # 2. Or preceded by letter and followed by 's' or 't' (common contractions)
+    if before.isalpha() and after.isalpha():
+        return True
+    if before.isalpha() and after in 's':
+        return True
+    
+    return False
+
+def sentence_chunk_text(text, max_words=MAX_CHUNK_WORDS, min_words=MIN_CHUNK_WORDS):
+    """
+    Simple and reliable text chunking that follows the exact rules:
+    
+    TEXT CHUNKING RULES:
+    1. Break at sentence boundaries (. ! ?) first (highest priority)
+    2. If sentence > max_words, break at punctuation working backwards (; — , in that order)
+    3. If no punctuation available, preserve sentence intact to maintain coherence
+    4. Ensure all chunks meet min_words requirement by combining small chunks
+    
+    PUNCTUATION HIERARCHY (for breaking long sentences):
+    1. . ! ? (sentence boundaries) - handled at sentence level
+    2. ; (semicolon) - major pause
+    3. — – (dashes) - major pause  
+    4. , (comma) - minor pause
+    5. Preserve overlong sentences if no punctuation available
+    """
+    import re
+    
+    # Process text paragraph by paragraph to preserve structure
+    paragraphs = text.split('\n\n')
+    all_final_chunks = []
+    
+    for paragraph in paragraphs:
+        paragraph = paragraph.strip()
+        if not paragraph:
+            continue
+            
+        # Check if this is a chapter/section header
+        para_lower = paragraph.lower().strip()
+        is_chapter_header = (
+            any(word in para_lower for word in ['chapter', 'section', 'part', 'prologue', 'epilogue']) and
+            len(paragraph.split()) <= 10
+        )
+        
+        if is_chapter_header:
+            # Chapter headers are their own chunks and always paragraph ends
+            all_final_chunks.append((paragraph, True))
+            continue
+        
+        # Split into sentences using periods, exclamation marks, question marks
+        # This avoids the complex quote detection that was causing problems
+        sentences = re.split(r'([.!?])\s+', paragraph.strip())
+        
+        # Reconstruct sentences with their punctuation
+        reconstructed_sentences = []
+        for i in range(0, len(sentences) - 1, 2):
+            sentence = sentences[i].strip()
+            if i + 1 < len(sentences):
+                punct = sentences[i + 1]
+                sentence += punct
+            if sentence:
+                reconstructed_sentences.append(sentence)
+        
+        # Handle any remaining text (no ending punctuation)
+        if sentences and sentences[-1].strip():
+            last_part = sentences[-1].strip()
+            if last_part and not last_part in '.!?':
+                reconstructed_sentences.append(last_part)
+        
+        # Process each sentence
+        paragraph_chunks = []
+        for sent_idx, sentence in enumerate(reconstructed_sentences):
+            is_last_sentence = (sent_idx == len(reconstructed_sentences) - 1)
+            words = sentence.split()
+            
+            if len(words) <= max_words:
+                # Sentence fits, use as-is
+                paragraph_chunks.append((sentence.strip(), is_last_sentence))
+            else:
+                # Sentence too long, break it using punctuation
+                broken_chunks = _break_long_sentence_simple(sentence, max_words)
+                # Only mark the last broken chunk as sentence end
+                for i, chunk in enumerate(broken_chunks):
+                    is_chunk_end = (is_last_sentence and i == len(broken_chunks) - 1)
+                    paragraph_chunks.append((chunk.strip(), is_chunk_end))
+        
+        all_final_chunks.extend(paragraph_chunks)
+    
+    # Combine small chunks that don't meet min_words requirement
+    combined_chunks = _combine_small_chunks(all_final_chunks, min_words, max_words)
+    
+    return combined_chunks
+
+def _break_long_sentence_simple(sentence, max_words):
+    """Break a long sentence at punctuation marks, working backwards"""
+    import re
+    
+    # Punctuation patterns in priority order
+    patterns = [
+        r';\s*',      # semicolon + optional space
+        r'—\s*',      # em dash + optional space  
+        r'–\s*',      # en dash + optional space
+        r',\s*',      # comma + optional space
+    ]
+    
+    chunks = []
+    remaining = sentence.strip()
+    
+    while remaining:
+        words = remaining.split()
+        if len(words) <= max_words:
+            chunks.append(remaining)
+            break
+        
+        # Find best break point working backwards
+        best_break = -1
+        
+        # Try each punctuation pattern
+        for pattern in patterns:
+            matches = list(re.finditer(pattern, remaining))
+            if matches:
+                # Find rightmost match that results in chunk <= max_words
+                for match in reversed(matches):
+                    test_chunk = remaining[:match.end()].strip()
+                    if len(test_chunk.split()) <= max_words:
+                        best_break = match.end()
+                        break
+                if best_break != -1:
+                    break
+        
+        if best_break != -1:
+            # Found good break point
+            chunk = remaining[:best_break].strip()
+            chunks.append(chunk)
+            remaining = remaining[best_break:].strip()
+        else:
+            # No punctuation found - preserve sentence coherence by keeping it intact
+            # This prevents splitting sentences with potentially different sentiment
+            chunks.append(remaining)
+            break
+    
+    return chunks
+
+def _combine_small_chunks(chunks, min_words, max_words):
+    """Combine chunks that are too small"""
+    combined = []
+    current_chunk = ""
+    current_is_para_end = False
+    
+    for chunk_text, is_para_end in chunks:
+        chunk_words = len(chunk_text.split())
+        current_words = len(current_chunk.split()) if current_chunk else 0
+        
+        if not current_chunk:
+            # First chunk
+            current_chunk = chunk_text
+            current_is_para_end = is_para_end
+        elif current_words + chunk_words <= max_words:
+            # Can combine
+            current_chunk = current_chunk + " " + chunk_text
+            current_is_para_end = is_para_end  # Use the latest para_end flag
+        else:
+            # Can't combine, flush current and start new
+            if current_words >= min_words:
+                combined.append((current_chunk, current_is_para_end))
+                current_chunk = chunk_text
+                current_is_para_end = is_para_end
+            else:
+                # Current chunk too small, force combine anyway
+                current_chunk = current_chunk + " " + chunk_text
+                current_is_para_end = is_para_end
+    
+    # Handle remaining chunk
+    if current_chunk:
+        combined.append((current_chunk, current_is_para_end))
+    
+    return combined
+
+def break_long_sentence_backwards(sentence, max_words, min_words):
+    """
+    Break a long sentence working backwards from the end to find natural punctuation.
+    
+    ALGORITHM:
+    1. Start from sentence end, work backwards to find punctuation within max_words
+    2. Break at the latest (rightmost) punctuation that keeps chunk <= max_words
+    3. This preserves natural pauses and speech rhythm
+    4. Continue processing remaining text normally
+    
+    PUNCTUATION HIERARCHY (in order of preference):
+    1. . ! ? (sentence boundaries) - highest priority
+    2. ; (semicolon) - major pause
+    3. — (em dash) - major pause  
+    4. , (comma) - minor pause
+    5. Force break at word limit (last resort)
+    """
+    
+    # Punctuation patterns to search for (in order of preference)
+    punctuation_patterns = [
+        r'[.!?]\s+',  # sentence boundaries + required space (highest priority)
+        r';\s*',      # semicolon + optional space
+        r'—\s*',      # em dash + optional space
+        r'–\s*',      # en dash + optional space
+        r',\s*',      # comma + optional space
+    ]
+    
+    chunks = []
+    remaining_text = sentence.strip()
+    
+    while remaining_text:
+        words = remaining_text.split()
+        
+        if len(words) <= max_words:
+            # Remaining text fits within limit
+            chunks.append(remaining_text.strip())
+            break
+            
+        # Text exceeds max_words - find backwards break point
+        # Search for punctuation within the current 'remaining_text' up to max_words
+        # We need to find the *last* punctuation mark that results in a chunk <= max_words
+        best_break_index = -1 # Index in 'words' list
+        best_break_pos_in_text = -1 # Character position in 'remaining_text'
+
+        # Iterate backwards from max_words down to min_words (or 1 if min_words is very small)
+        # to find the latest punctuation that keeps the chunk within limits.
+        for i in range(min(max_words, len(words)) -1, 0, -1):
+            sub_text = " ".join(words[:i+1]) # Text up to current word
+            
+            found_punctuation = False
+            for pattern in punctuation_patterns:
+                matches = list(re.finditer(pattern, sub_text))
+                if matches:
+                    # Take the rightmost match in this sub_text
+                    last_match = matches[-1]
+                    # Ensure the break is within the max_words limit
+                    if len(sub_text[:last_match.end()].split()) <= max_words:
+                        best_break_index = i # Store word index
+                        best_break_pos_in_text = last_match.end() # Store char position
+                        found_punctuation = True
+                        break # Found a good break for this sub_text, move to next i
+            if found_punctuation:
+                break # Found the best break for the overall chunk, exit outer loop
+
+        if best_break_pos_in_text != -1:
+            # Found punctuation - break after it, keeping punctuation with preceding text
+            chunk_text = remaining_text[:best_break_pos_in_text].strip()
+            chunks.append(chunk_text)
+            remaining_text = remaining_text[best_break_pos_in_text:].strip()
+        else:
+            # No punctuation found within the desired range - keep sentence intact
+            # This preserves sentence coherence over word count limits
+            chunks.append(remaining_text.strip())
+            break
+    
+    return chunks
+
+# ============================================================================
+# CONTENT BOUNDARY DETECTION
+# ============================================================================
+
+def detect_punctuation_boundary(chunk_text):
+    """
+    Detect the ending punctuation of a text chunk for precise silence insertion.
+    
+    Returns specific punctuation boundary types:
+    - "comma" -> Brief pause after commas
+    - "semicolon" -> Medium pause after semicolons  
+    - "colon" -> Pause after colons
+    - "period" -> Sentence end pause
+    - "question_mark" -> Question pause
+    - "exclamation" -> Exclamation pause
+    - "dash" -> Em dash pause
+    - "ellipsis" -> Ellipsis pause (suspense)
+    - "quote_end" -> End of quoted speech
+    - None -> No specific punctuation detected
+    """
+    # Strip whitespace and newlines for accurate detection
+    text = chunk_text.strip()
+    
+    if not text:
+        return None
+    
+    # Check ending punctuation patterns (in order of specificity)
+    if text.endswith('...'):
+        return "ellipsis"
+    elif text.endswith('"') or text.endswith("'"):
+        return "quote_end"
+    elif text.endswith('!'):
+        return "exclamation"
+    elif text.endswith('?'):
+        return "question_mark"
+    elif text.endswith('.'):
+        return "period"
+    elif text.endswith(':'):
+        return "colon"
+    elif text.endswith(';'):
+        return "semicolon"
+    elif text.endswith(','):
+        return "comma"
+    elif text.endswith('—') or text.endswith('–'):
+        return "dash"
+    
+    return None
+
+def detect_content_boundaries(chunk_text, chunk_index, all_chunks, is_paragraph_end=False):
+    """
+    Detect chapter breaks and paragraph endings for appropriate silence insertion.
+    Now enhanced with punctuation-specific boundary detection.
+    
+    BOUNDARY DETECTION REQUIREMENTS:
+    - Chapter start: "Chapter N", "Ch. N", "I.", "1." patterns
+    - Chapter end: Next chunk is a chapter start
+    - Section break: Multiple asterisks, hashes, or em-dashes
+    - Paragraph end: Detected via chunking process flag or content analysis
+    - Punctuation: Specific ending punctuation for precise silence timing
+    
+    Returns boundary_type for silence insertion:
+    - "chapter_start" -> Long pause before chapter
+    - "chapter_end" -> Long pause after chapter
+    - "section_break" -> Medium pause for section breaks  
+    - "paragraph_end" -> Short pause for paragraph breaks
+    - Punctuation types: "comma", "period", "question_mark", etc.
+    - None -> No special boundary detected
+    """
+    boundary_type = None
+
+    # Chapter detection (flexible patterns)
+    chapter_patterns = [
+        r'^(Chapter \d+|CHAPTER \d+)',
+        r'^(Ch\. \d+|CH\. \d+)',
+        r'^\d+\.',  # Simple "1." numbering
+        r'^[IVX]+\.',  # Roman numerals "I.", "II.", etc.
+    ]
+
+    for pattern in chapter_patterns:
+        if re.search(pattern, chunk_text.strip(), re.MULTILINE):
+            boundary_type = "chapter_start"
+            break
+
+    # Look ahead for chapter start (current chunk ends chapter)
+    if chunk_index + 1 < len(all_chunks):
+        next_chunk = all_chunks[chunk_index + 1]
+        for pattern in chapter_patterns:
+            if re.search(pattern, next_chunk.strip()):
+                boundary_type = "chapter_end"
+                break
+
+    # Section breaks (asterisks, multiple line breaks)
+    if re.search(r'\*{3,}|\#{3,}|—{3,}', chunk_text):
+        boundary_type = "section_break"
+
+    # Paragraph ending detection
+    # Use the is_paragraph_end flag from chunking process since newlines are stripped
+    if is_paragraph_end and boundary_type is None:
+        boundary_type = "paragraph_end"
+
+    # If no major structural boundary found, check punctuation
+    if boundary_type is None:
+        boundary_type = detect_punctuation_boundary(chunk_text)
+
+    return boundary_type
+
+def _split_long_dialogue(sentence, max_words, recursion_depth=0):
+    """
+    Split long dialogue sections that exceed word limits.
+    Tries to break at natural points: attribution, internal punctuation, then word boundaries.
+    """
+    # Prevent infinite recursion
+    if recursion_depth > 3:
+        # Force word boundary split if recursion gets too deep
+        words = sentence.split()
+        sentences = []
+        start = 0
+        while start < len(words):
+            end = min(start + max_words, len(words))
+            chunk_words = words[start:end]
+            sentences.append(' '.join(chunk_words))
+            start = end
+        return sentences
+    
+    words = sentence.split()
+    if len(words) <= max_words:
+        return [sentence]
+    
+    sentences = []
+    
+    # Strategy 1: Break at dialogue attribution (he said, she replied, etc.)
+    attribution_pattern = r'(\s+(?:he|she|I|they|[A-Z][a-z]+)\s+(?:said|replied|asked|shouted|whispered|continued|added|interrupted)[^.!?]*?[.!?]?\s*)'
+    attribution_matches = list(re.finditer(attribution_pattern, sentence, re.IGNORECASE))
+    
+    if attribution_matches:
+        start = 0
+        for match in attribution_matches:
+            # Check if breaking here keeps chunks under limit
+            before_attr = sentence[start:match.end()].strip()
+            if before_attr and len(before_attr.split()) <= max_words:
+                sentences.append(before_attr)
+                start = match.end()
+        
+        # Add remaining text
+        if start < len(sentence):
+            remaining = sentence[start:].strip()
+            if remaining:
+                if len(remaining.split()) > max_words:
+                    # Recursively split if still too long, but with depth tracking
+                    sentences.extend(_split_long_dialogue(remaining, max_words, recursion_depth + 1))
+                else:
+                    sentences.append(remaining)
+        
+        if sentences:  # If we successfully split, return result
+            return sentences
+    
+    # Strategy 2: Break at internal punctuation (commas, semicolons within quotes)
+    punct_pattern = r'([,;:]\s+)'
+    parts = re.split(punct_pattern, sentence)
+    
+    current_chunk = ""
+    sentences = []
+    for i, part in enumerate(parts):
+        test_chunk = current_chunk + part
+        if len(test_chunk.split()) > max_words and current_chunk:
+            sentences.append(current_chunk.strip())
+            current_chunk = part
+        else:
+            current_chunk = test_chunk
+    
+    if current_chunk.strip():
+        sentences.append(current_chunk.strip())
+    
+    # Check if any resulting chunk is still too long and needs further splitting
+    final_sentences = []
+    for chunk in sentences:
+        if len(chunk.split()) > max_words:
+            # Split oversized chunks using word boundaries
+            chunk_words = chunk.split()
+            start = 0
+            while start < len(chunk_words):
+                end = min(start + max_words, len(chunk_words))
+                sub_chunk_words = chunk_words[start:end]
+                final_sentences.append(' '.join(sub_chunk_words))
+                start = end
+        else:
+            final_sentences.append(chunk)
+    
+    if len(final_sentences) > 1:  # If we successfully split, return result
+        return final_sentences
+    
+    # Strategy 3: Force break at word boundaries (guaranteed to work)
+    sentences = []
+    start = 0
+    while start < len(words):
+        end = min(start + max_words, len(words))
+        chunk_words = words[start:end]
+        sentences.append(' '.join(chunk_words))
+        start = end
+    
+    return sentences
+
+# ============================================================================
+# UTILITY FUNCTIONS
+# ============================================================================
+
+def reload_abbreviations():
+    """Reload abbreviations from file (useful for testing changes)"""
+    return load_abbreviations()
+
+def test_abbreviations(test_text="Dr. Smith met with the M.P. at 3:30 p.m. on the 21st."):
+    """Test abbreviation replacements on sample text"""
+    abbreviation_replacements = load_abbreviations()
+    print(f"Original: {test_text}")
+    processed = preprocess_abbreviations(test_text, abbreviation_replacements)
+    print(f"Processed: {processed}")
+    return processed
+
+def test_chunking(test_text=None, max_words=20, min_words=4):
+    """Test the enhanced chunking with sample or custom text"""
+    if test_text is None:
+        test_text = '''Though perfectly worldly-wise, and able, as she expressed it, to take care of herself, there was yet something curiously ingenuous in her single-minded attitude towards life, and her whole-hearted determination to "make good." This glimpse of a world unknown to me was not without its charm, and I enjoyed seeing her vivid little face light up as she talked.'''
+
+    chunks = sentence_chunk_text(test_text, max_words=max_words, min_words=min_words)
+
+    print("Enhanced Chunking Results:")
+    for i, (chunk, is_para) in enumerate(chunks):
+        word_count = len(chunk.split())
+        print(f"Chunk {i+1} ({word_count} words): {chunk}")
+        if word_count > max_words:
+            print(f"  ✅ Over {max_words} words but complete sentence (follows punctuation rules)")
+        print()
+
+    return chunks

HF_Deploy/modules/tts_engine.py

@@ -0,0 +1,710 @@
+"""
+TTS Engine Module
+Handles ChatterboxTTS interface, model loading, and chunk processing coordination
+"""
+
+import torch
+import gc
+import time
+import logging
+import shutil
+import sys
+from datetime import timedelta
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+import torchaudio as ta
+
+from config.config import *
+from modules.text_processor import smart_punctuate, sentence_chunk_text, detect_content_boundaries
+
+def find_chunks_json_file(book_name):
+    """Find the corresponding chunks JSON file for a book"""
+    from config.config import AUDIOBOOK_ROOT
+    
+    # Look in the TTS processing directory
+    tts_chunks_dir = AUDIOBOOK_ROOT / book_name / "TTS" / "text_chunks"
+    json_path = tts_chunks_dir / "chunks_info.json"
+    
+    if json_path.exists():
+        return json_path
+    
+    # Also check old Text_Input location for backwards compatibility
+    text_input_dir = Path("Text_Input")
+    possible_names = [
+        f"{book_name}_chunks.json",
+        f"{book_name.lower()}_chunks.json",
+        f"{book_name.replace(' ', '_')}_chunks.json"
+    ]
+    
+    for name in possible_names:
+        old_json_path = text_input_dir / name
+        if old_json_path.exists():
+            return old_json_path
+    
+    return None
+from modules.audio_processor import (
+    smart_audio_validation, apply_smart_fade, add_chunk_end_silence,
+    add_contextual_silence, pause_for_chunk_review, get_chunk_audio_duration,
+    has_mid_energy_drop, apply_smart_fade_memory, smart_audio_validation_memory
+)
+from modules.file_manager import (
+    setup_book_directories, find_book_files, ensure_voice_sample_compatibility,
+    combine_audio_chunks, get_audio_files_in_directory, convert_to_m4b, add_metadata_to_m4b
+)
+from modules.progress_tracker import setup_logging, log_chunk_progress, log_run
+
+# ============================================================================
+# MEMORY AND MODEL MANAGEMENT
+# ============================================================================
+
+def monitor_gpu_activity(operation_name):
+    """Lightweight GPU monitoring for high-speed processing"""
+    # Disabled expensive pynvml queries to free up GPU cycles
+    if torch.cuda.is_available():
+        allocated = torch.cuda.memory_allocated() / 1024**3
+        # Skip GPU utilization queries during production runs
+        return allocated, 0
+    return 0, 0
+
+def optimize_memory_usage():
+    """Aggressive memory management for 8GB VRAM"""
+    torch.cuda.empty_cache()
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.ipc_collect()
+
+def monitor_vram_usage(operation_name=""):
+    """Real-time VRAM monitoring"""
+    if torch.cuda.is_available():
+        allocated = torch.cuda.memory_allocated() / 1024**3
+        reserved = torch.cuda.memory_reserved() / 1024**3
+
+        if allocated > VRAM_SAFETY_THRESHOLD:
+            logging.warning(f"⚠️ High VRAM usage during {operation_name}: {allocated:.1f}GB allocated, {reserved:.1f}GB reserved")
+            optimize_memory_usage()
+
+        return allocated, reserved
+    return 0, 0
+
+def get_optimal_workers(user_max_workers=None):
+    """Dynamic worker allocation based on device type and resources"""
+    # Check for user override first
+    if user_max_workers is not None:
+        print(f"👤 Using user-defined workers: {user_max_workers}")
+        return int(user_max_workers)
+        
+    if not USE_DYNAMIC_WORKERS:
+        return MAX_WORKERS
+
+    # CPU-based worker calculation
+    if not torch.cuda.is_available():
+        import psutil
+        cpu_cores = psutil.cpu_count(logical=False)  # Physical cores
+        available_memory = psutil.virtual_memory().available / 1024**3  # GB
+        
+        # Each TTS model instance needs ~2-3GB RAM
+        # Conservative estimation: allow 1 worker per 4GB available RAM
+        memory_limited_workers = max(1, int(available_memory / 4))
+        
+        # CPU-based calculation: use 50% of physical cores for intensive TTS work
+        cpu_limited_workers = max(1, int(cpu_cores * 0.5))
+        
+        optimal_workers = min(memory_limited_workers, cpu_limited_workers, MAX_WORKERS)
+        print(f"💻 CPU mode: {cpu_cores} cores, {available_memory:.1f}GB RAM → {optimal_workers} workers")
+        return optimal_workers
+    
+    # GPU-based worker calculation (existing logic)
+    allocated_vram = torch.cuda.memory_allocated() / 1024**3
+
+    if allocated_vram < 5.0:
+        return min(TEST_MAX_WORKERS, MAX_WORKERS)
+    elif allocated_vram < VRAM_SAFETY_THRESHOLD:
+        return min(2, MAX_WORKERS)
+    else:
+        return 1
+
+def load_optimized_model(device):
+    """Load TTS model with memory optimizations and device detection"""
+    from chatterbox.tts import ChatterboxTTS
+    
+    # Detect available device if not specified or if CUDA not available
+    if device == "cuda" and not torch.cuda.is_available():
+        print("⚠️  CUDA not available, falling back to CPU")
+        device = "cpu"
+    elif device == "auto":
+        if torch.cuda.is_available():
+            device = "cuda"
+            print("✅ CUDA detected, using GPU")
+        else:
+            device = "cpu"
+            print("💻 No GPU detected, using CPU")
+    
+    print(f"🔧 Loading ChatterboxTTS model on device: {device}")
+
+    try:
+        # Load model (ChatterboxTTS.from_pretrained doesn't support torch_dtype parameter)
+        model = ChatterboxTTS.from_pretrained(device=device)
+        logging.info(f"✅ Loaded ChatterboxTTS model on {device}")
+    except Exception as e:
+        print(f"❌ Failed to load model on {device}: {e}")
+        if device == "cuda":
+            print("🔄 Retrying with CPU...")
+            try:
+                model = ChatterboxTTS.from_pretrained(device="cpu")
+                logging.info("✅ Loaded model on CPU (GPU failed)")
+                device = "cpu"
+            except Exception as e2:
+                print(f"❌ Failed to load model on CPU: {e2}")
+                raise e2
+        else:
+            raise e
+
+    # Only apply eval() and benchmark if the model has these attributes
+    if hasattr(model, 'eval'):
+        model.eval()
+
+    # Set CUDNN benchmark for performance (if available)
+    if torch.backends.cudnn.is_available():
+        torch.backends.cudnn.benchmark = True
+
+    return model
+
+# ============================================================================
+# CHUNK PROCESSING
+# ============================================================================
+
+def patch_alignment_layer(tfmr, alignment_layer_idx=12):
+    """Patch alignment layer to avoid recursion"""
+    from types import MethodType
+    target_layer = tfmr.layers[alignment_layer_idx].self_attn
+    original_forward = target_layer.forward
+
+    def patched_forward(self, *args, **kwargs):
+        kwargs['output_attentions'] = True
+        return original_forward(*args, **kwargs)
+
+    target_layer.forward = MethodType(patched_forward, target_layer)
+
+def process_one_chunk(
+    i, chunk, text_chunks_dir, audio_chunks_dir,
+    voice_path, tts_params, start_time, total_chunks,
+    punc_norm, basename, log_run_func, log_path, device,
+    model, asr_model, all_chunks, boundary_type="none"
+):
+    """Enhanced chunk processing with quality control, contextual silence, and deep cleanup"""
+    import difflib
+    from pydub import AudioSegment
+
+    chunk_id_str = f"{i+1:05}"
+    chunk_path = text_chunks_dir / f"chunk_{chunk_id_str}.txt"
+    with open(chunk_path, 'w', encoding='utf-8') as cf:
+        cf.write(chunk)
+
+    chunk_audio_path = audio_chunks_dir / f"chunk_{chunk_id_str}.wav"
+
+    # ============================================================================
+    # ENHANCED PERIODIC DEEP CLEANUP
+    # ============================================================================
+    cleanup_interval = CLEANUP_INTERVAL
+
+    # Skip cleanup on model reinitialization chunks to avoid conflicts
+    if (i + 1) % cleanup_interval == 0 and (i + 1) % BATCH_SIZE != 0:
+        print(f"\n🧹 {YELLOW}DEEP CLEANUP at chunk {i+1}/{total_chunks}...{RESET}")
+
+        # Enhanced VRAM monitoring before cleanup
+        allocated_before = torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0
+        reserved_before = torch.cuda.memory_reserved() / 1024**3 if torch.cuda.is_available() else 0
+
+        print(f"   Before: VRAM Allocated: {allocated_before:.1f}GB | Reserved: {reserved_before:.1f}GB")
+
+        # Bulk temp file cleanup
+        print("   🗑️ Cleaning bulk temporary files...")
+        temp_patterns = ["*_try*.wav", "*_pre.wav", "*_fade*.wav", "*_debug*.wav", "*_temp*.wav", "*_backup*.wav"]
+        total_temp_files = 0
+        for pattern in temp_patterns:
+            temp_files = list(audio_chunks_dir.glob(pattern))
+            for temp_file in temp_files:
+                temp_file.unlink(missing_ok=True)
+            total_temp_files += len(temp_files)
+
+        if total_temp_files > 0:
+            print(f"   🗑️ Removed {total_temp_files} temporary audio files")
+
+        # Aggressive CUDA context reset
+        print("   🔄 Performing aggressive CUDA context reset...")
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+        # Force CUDA context reset
+        if hasattr(torch.cuda, 'reset_peak_memory_stats'):
+            torch.cuda.reset_peak_memory_stats()
+        if hasattr(torch._C, '_cuda_clearCublasWorkspaces'):
+            torch._C._cuda_clearCublasWorkspaces()
+
+        # Force garbage collection multiple times
+        for _ in range(3):
+            gc.collect()
+
+        # Clear model cache if it has one
+        if hasattr(model, 'clear_cache'):
+            model.clear_cache()
+        elif hasattr(model, 'reset_states'):
+            model.reset_states()
+
+        # Brief pause to let GPU settle
+        time.sleep(1.0)
+
+        # Monitor after cleanup
+        allocated_after = torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0
+        reserved_after = torch.cuda.memory_reserved() / 1024**3 if torch.cuda.is_available() else 0
+
+        print(f"   After:  VRAM Allocated: {allocated_after:.1f}GB | Reserved: {reserved_after:.1f}GB")
+        print(f"   Freed:  {allocated_before - allocated_after:.1f}GB allocated, {reserved_before - reserved_after:.1f}GB reserved")
+        print(f"🧹 {GREEN}Deep cleanup complete!{RESET}\n")
+
+    best_sim, best_asr_text = -1, ""
+    wav_path_active = None
+    attempt_paths = []
+    mid_drop_retries = 0
+    max_mid_drop_retries = 2
+
+    for attempt_num in range(1, 3):
+        logging.info(f"🔁 Starting TTS for chunk {chunk_id_str}, attempt {attempt_num}")
+        try:
+            tts_args = {k: v for k, v in tts_params.items() if k != "max_workers"}
+
+            # monitor_gpu_activity(f"Before TTS chunk_{chunk_id_str}")  # Disabled for speed
+            with torch.no_grad():
+                wav = model.generate(chunk, **tts_args).detach().cpu()
+            # monitor_gpu_activity(f"After TTS chunk_{chunk_id_str}")  # Disabled for speed
+
+            if wav.dim() == 1:
+                wav = wav.unsqueeze(0)
+
+            # Retry if mid-energy drop is enabled and detected (check in memory)
+            if ENABLE_MID_DROP_CHECK and has_mid_energy_drop(wav, model.sr):
+                mid_drop_retries += 1
+                if mid_drop_retries >= max_mid_drop_retries:
+                    logging.info(f"⚠️ Mid-drop retry limit reached for {chunk_id_str}. Accepting audio.")
+                else:
+                    logging.info(f"⚠️ Mid-chunk noise detected in {chunk_id_str}. Retrying...")
+                    continue
+
+            # Convert tensor to AudioSegment for in-memory processing
+            import io
+            import soundfile as sf
+            from pydub import AudioSegment
+            
+            # Convert wav tensor to AudioSegment (in memory)
+            wav_np = wav.squeeze().numpy()
+            with io.BytesIO() as wav_buffer:
+                sf.write(wav_buffer, wav_np, model.sr, format='wav')
+                wav_buffer.seek(0)
+                audio_segment = AudioSegment.from_wav(wav_buffer)
+            
+            # Smart fade removed - replaced by precise audio trimming
+            # Audio health validation disabled for speed
+            
+            # Note: Audio trimming will handle end-of-speech cleanup more precisely
+
+            # ASR validation (memory-based processing) - check user setting first
+            enable_asr_user = tts_params.get('enable_asr', False)
+            if (enable_asr_user or ENABLE_ASR) and asr_model is not None:
+                from modules.audio_processor import asr_f1_score
+                import io
+                import soundfile as sf
+                # monitor_gpu_activity(f"Before ASR chunk_{chunk_id_str}")  # Disabled for speed
+                try:
+                    # Process ASR completely in memory - no disk writes
+                    # Convert AudioSegment to numpy array for ASR
+                    samples = np.array(audio_segment.get_array_of_samples())
+                    if audio_segment.channels == 2:
+                        samples = samples.reshape((-1, 2)).mean(axis=1)
+                    
+                    # Normalize to float32 for ASR model
+                    audio_np = samples.astype(np.float32) / audio_segment.max_possible_amplitude
+                    
+                    # Use ASR model directly on numpy array (if supported)
+                    # Note: This depends on the ASR model's input capabilities
+                    result = asr_model.transcribe(audio_np)
+                    
+                    if not isinstance(result, dict) or "text" not in result:
+                        raise ValueError(f"Invalid ASR result type: {type(result)}")
+
+                    asr_text = result.get("text", "").strip()
+                    sim_ratio = asr_f1_score(punc_norm(chunk), asr_text)
+
+                except Exception as e:
+                    print(f"❌ ASR failed for {chunk_id_str}: {e}")
+                    log_run_func(f"ASR VALIDATION FAILED - Chunk {chunk_id_str}:\nExpected:\n{chunk}\nActual:\n<ASR Failure: {e}>\nSimilarity: -1.000\n" + "="*50, log_path)
+                    sim_ratio = -1.0
+                    continue
+
+                logging.info(f"ASR similarity for chunk {chunk_id_str}: {sim_ratio:.3f}")
+                if sim_ratio < 0.7:
+                    continue
+
+                # Track best valid match
+                best_sim = sim_ratio
+                best_asr_text = asr_text
+                # monitor_gpu_activity(f"After ASR chunk_{chunk_id_str}")  # Disabled for speed
+
+            # Success - we have processed audio in memory
+            final_audio = audio_segment
+            break
+
+        except Exception as e:
+            import traceback
+            logging.error(f"Exception during TTS attempt {attempt_num} for chunk {chunk_id_str}: {e}")
+            traceback.print_exc()
+            continue
+
+    if 'final_audio' not in locals():
+        logging.info(f"❌ Chunk {chunk_id_str} failed all attempts.")
+        return None, None
+
+    # Apply trimming and contextual silence in memory before final save
+    from modules.audio_processor import process_audio_with_trimming_and_silence
+    
+    if boundary_type and boundary_type != "none":
+        final_audio = process_audio_with_trimming_and_silence(final_audio, boundary_type)
+        print(f"🔇 Added {boundary_type} silence to chunk {i+1:05}")
+    else:
+        # Apply trimming even without boundary type if enabled
+        if ENABLE_AUDIO_TRIMMING:
+            from modules.audio_processor import trim_audio_endpoint
+            final_audio = trim_audio_endpoint(final_audio)
+
+    # Note: ENABLE_CHUNK_END_SILENCE is now handled by punctuation-specific silence
+    # The new system provides more precise silence based on actual punctuation
+
+    # Final save - only disk write in entire process
+    final_path = audio_chunks_dir / f"chunk_{chunk_id_str}.wav"
+    final_audio.export(final_path, format="wav")
+    logging.info(f"✅ Saved final chunk: {final_path.name}")
+
+    # No intermediate file cleanup needed - all processing done in memory
+
+    # Log details - only log ASR failures
+    asr_active = enable_asr_user or ENABLE_ASR
+    if asr_active and best_sim < 0.8:
+        log_run_func(f"ASR VALIDATION FAILED - Chunk {chunk_id_str}:\nExpected:\n{chunk}\nActual:\n{best_asr_text}\nSimilarity: {best_sim:.3f}\n" + "="*50, log_path)
+    elif not asr_active:
+        log_run_func(f"Chunk {chunk_id_str}: Original text: {chunk}", log_path)
+
+    # Silence already added in memory above - no disk processing needed
+
+    # Enhanced regular cleanup (every chunk)
+    del wav
+    optimize_memory_usage()
+
+    # Additional per-chunk cleanup for long runs
+    if (i + 1) % 50 == 0:
+        torch.cuda.empty_cache()
+        gc.collect()
+
+    return i, final_path
+
+# ============================================================================
+# MAIN BOOK PROCESSING FUNCTION
+# ============================================================================
+
+from vaderSentiment.vaderSentiment import SentimentIntensityAnalyzer
+from wrapper.chunk_loader import save_chunks
+
+def generate_enriched_chunks(text_file, output_dir, user_tts_params=None):
+    """Reads a text file, performs VADER sentiment analysis, and returns enriched chunks."""
+    analyzer = SentimentIntensityAnalyzer()
+    
+    raw_text = text_file.read_text(encoding='utf-8')
+    cleaned = smart_punctuate(raw_text)
+    chunks = sentence_chunk_text(cleaned)
+
+    # Use user-provided parameters as base, or fall back to config defaults
+    if user_tts_params:
+        base_exaggeration = user_tts_params.get('exaggeration', BASE_EXAGGERATION)
+        base_cfg_weight = user_tts_params.get('cfg_weight', BASE_CFG_WEIGHT)
+        base_temperature = user_tts_params.get('temperature', BASE_TEMPERATURE)
+    else:
+        base_exaggeration = BASE_EXAGGERATION
+        base_cfg_weight = BASE_CFG_WEIGHT
+        base_temperature = BASE_TEMPERATURE
+
+    enriched = []
+    chunk_texts = [chunk_text for chunk_text, _ in chunks]
+    
+    for i, (chunk_text, is_para_end) in enumerate(chunks):
+        sentiment_scores = analyzer.polarity_scores(chunk_text)
+        compound_score = sentiment_scores['compound']
+
+        exaggeration = base_exaggeration + (compound_score * VADER_EXAGGERATION_SENSITIVITY)
+        cfg_weight = base_cfg_weight + (compound_score * VADER_CFG_WEIGHT_SENSITIVITY)
+        temperature = base_temperature + (compound_score * VADER_TEMPERATURE_SENSITIVITY)
+
+        # Clamp values to defined min/max
+        exaggeration = round(max(TTS_PARAM_MIN_EXAGGERATION, min(exaggeration, TTS_PARAM_MAX_EXAGGERATION)), 2)
+        cfg_weight = round(max(TTS_PARAM_MIN_CFG_WEIGHT, min(cfg_weight, TTS_PARAM_MAX_CFG_WEIGHT)), 2)
+        temperature = round(max(TTS_PARAM_MIN_TEMPERATURE, min(temperature, TTS_PARAM_MAX_TEMPERATURE)), 2)
+
+        boundary_type = detect_content_boundaries(chunk_text, i, chunk_texts, is_para_end)
+        
+        enriched.append({
+            "index": i,
+            "text": chunk_text,
+            "word_count": len(chunk_text.split()),
+            "boundary_type": boundary_type if boundary_type else "none",
+            "sentiment_compound": compound_score,
+            "tts_params": {
+                "exaggeration": exaggeration,
+                "cfg_weight": cfg_weight,
+                "temperature": temperature
+            }
+        })
+
+    output_json_path = output_dir / "chunks_info.json"
+    save_chunks(output_json_path, enriched)
+    return enriched
+
+def process_book_folder(book_dir, voice_path, tts_params, device, skip_cleanup=False):
+    """Enhanced book processing with batch processing to prevent hangs"""
+    print(f"🔍 DEBUG: Entering process_book_folder with book_dir='{book_dir}', voice_path='{voice_path}'")
+    
+    from chatterbox.tts import punc_norm
+    print(f"🔍 DEBUG: Successfully imported punc_norm")
+
+    # Setup directories
+    print(f"🔍 DEBUG: Calling setup_book_directories...")
+    output_root, tts_dir, text_chunks_dir, audio_chunks_dir = setup_book_directories(book_dir)
+    print(f"🔍 DEBUG: Directory setup complete")
+    
+    # Clean previous processing files (but skip for resume operations)
+    if skip_cleanup:
+        print(f"🔄 RESUME MODE: Skipping cleanup to preserve existing chunks")
+        print(f"📁 Preserving: {text_chunks_dir}, {audio_chunks_dir}")
+    else:
+        print(f"🧹 FRESH PROCESSING: Cleaning previous processing files...")
+        import glob
+        
+        # Clear text chunks  
+        for txt_file in text_chunks_dir.glob("*.txt"):
+            txt_file.unlink(missing_ok=True)
+        for json_file in text_chunks_dir.glob("*.json"):
+            json_file.unlink(missing_ok=True)
+            
+        # Clear audio chunks
+        for wav_file in audio_chunks_dir.glob("*.wav"):
+            wav_file.unlink(missing_ok=True)
+            
+        # Clear logs
+        for log_file in output_root.glob("*.log"):
+            log_file.unlink(missing_ok=True)
+            
+        print(f"✅ Cleanup complete")
+
+    # Find book files
+    print(f"🔍 DEBUG: Calling find_book_files...")
+    book_files = find_book_files(book_dir)
+    text_files = [book_files['text']] if book_files['text'] else []
+    cover_file = book_files['cover']
+    nfo_file = book_files['nfo']
+    print(f"🔍 DEBUG: Found text files: {text_files}")
+
+    if not text_files:
+        logging.info(f"[{book_dir.name}] ERROR: No .txt files found in the book folder.")
+        return None, None, []
+
+    setup_logging(output_root)
+
+    # Generate enriched chunks with VADER analysis using user parameters
+    all_chunks = generate_enriched_chunks(text_files[0], text_chunks_dir, tts_params)
+
+    # Create run_log_lines
+    print(f"🔍 DEBUG: Creating run_log_lines...")
+    print(f"🔍 DEBUG: voice_path type: {type(voice_path)}, value: {voice_path}")
+    
+    # Extract voice name for logging
+    voice_name_for_log = voice_path.stem if hasattr(voice_path, 'stem') else Path(voice_path).stem
+    
+    run_log_lines = [
+        f"\n===== Processing: {book_dir.name} =====",
+        f"Voice: {voice_name_for_log}",
+        f"Started: {time.strftime('%Y-%m-%d %H:%M:%S')}",
+        f"Text files processed: {len(text_files)}",
+        f"Total chunks generated: {len(all_chunks)}"
+    ]
+
+    start_time = time.time()
+    total_chunks = len(all_chunks)
+    log_path = output_root / "chunk_validation.log"
+    total_audio_duration = 0.0
+
+    # Batch processing
+    print(f"📊 Processing {total_chunks} chunks in batches of {BATCH_SIZE}")
+
+    all_results = []
+
+    for batch_start in range(0, total_chunks, BATCH_SIZE):
+        batch_end = min(batch_start + BATCH_SIZE, total_chunks)
+        batch_chunks = all_chunks[batch_start:batch_end]
+
+        print(f"\n🔄 Processing batch: chunks {batch_start+1}-{batch_end}")
+
+        # Fresh model for each batch
+        model = load_optimized_model(device)
+        compatible_voice = ensure_voice_sample_compatibility(voice_path, output_dir=tts_dir)
+        model.prepare_conditionals(compatible_voice)
+
+        # Load ASR model once per batch if needed (check user settings first, then global config)
+        asr_model = None
+        enable_asr_user = tts_params.get('enable_asr', False)
+        if enable_asr_user or ENABLE_ASR:
+            import whisper
+            print(f"🎤 Loading Whisper ASR model for batch... (user setting: {enable_asr_user})")
+            # Use same device as TTS model, with fallback to CPU
+            asr_device = device if torch.cuda.is_available() and device == "cuda" else "cpu"
+            print(f"🎤 Loading ASR model on device: {asr_device}")
+            asr_model = whisper.load_model("base", device=asr_device)
+
+        futures = []
+        batch_results = []
+
+        # Dynamic worker allocation
+        user_max_workers = tts_params.get('max_workers', None)
+        optimal_workers = get_optimal_workers(user_max_workers)
+        print(f"🔧 Using {optimal_workers} workers for batch {batch_start+1}-{batch_end}")
+
+        with ThreadPoolExecutor(max_workers=optimal_workers) as executor:
+            for i, chunk_data in enumerate(batch_chunks):
+                global_chunk_index = batch_start + i
+
+                # Check for shutdown request
+                if shutdown_requested:
+                    print(f"\n⏹️ {YELLOW}Stopping submission of new chunks...{RESET}")
+                    break
+
+                # Handle both dictionary and tuple formats for chunk data
+                if isinstance(chunk_data, dict):
+                    chunk = chunk_data["text"]
+                    boundary_type = chunk_data.get("boundary_type", "none")
+                    # Use chunk-specific TTS params if available, otherwise fall back to global
+                    chunk_tts_params = chunk_data.get("tts_params", tts_params)
+                else:
+                    # Handle old tuple format (text, is_para_end) - convert to boundary_type
+                    chunk = chunk_data[0] if len(chunk_data) > 0 else str(chunk_data)
+                    # Convert old is_paragraph_end to boundary_type
+                    is_old_para_end = chunk_data[1] if len(chunk_data) > 1 else False
+                    boundary_type = "paragraph_end" if is_old_para_end else "none"
+                    chunk_tts_params = tts_params # Fallback for old format
+
+                # Handle both dictionary and tuple formats for backward compatibility  
+                all_chunk_texts = []
+                for cd in all_chunks:
+                    if isinstance(cd, dict):
+                        all_chunk_texts.append(cd["text"])
+                    else:
+                        # Handle old tuple format (text, is_para_end)
+                        all_chunk_texts.append(cd[0] if len(cd) > 0 else str(cd))
+
+                futures.append(executor.submit(
+                    process_one_chunk,
+                    global_chunk_index, chunk, text_chunks_dir, audio_chunks_dir,
+                    voice_path, chunk_tts_params, start_time, total_chunks,
+                    punc_norm, book_dir.name, log_run, log_path, device,
+                    model, asr_model, all_chunk_texts, boundary_type
+                ))
+
+            # Wait for batch to complete
+            print(f"🔄 {CYAN}Waiting for batch {batch_start+1}-{batch_end} to complete...{RESET}")
+            completed_count = 0
+
+            for fut in as_completed(futures):
+                try:
+                    idx, wav_path = fut.result()
+                    if wav_path and wav_path.exists():
+                        # Measure actual audio duration for this chunk
+                        chunk_duration = get_chunk_audio_duration(wav_path)
+                        total_audio_duration += chunk_duration
+                        batch_results.append((idx, wav_path))
+
+                        # Update progress every 10 chunks within batch
+                        completed_count += 1
+                        if completed_count % 10 == 0:
+                            log_chunk_progress(batch_start + completed_count - 1, total_chunks, start_time, total_audio_duration)
+
+                except Exception as e:
+                    logging.error(f"Future failed in batch: {e}")
+
+        # Clean up model after batch
+        print(f"🧹 Cleaning up after batch {batch_start+1}-{batch_end}")
+        del model
+        if asr_model:
+            del asr_model
+        torch.cuda.empty_cache()
+        gc.collect()
+        time.sleep(2)
+
+        all_results.extend(batch_results)
+        print(f"✅ Batch {batch_start+1}-{batch_end} completed ({len(batch_results)} chunks)")
+
+    # Final processing
+    quarantine_dir = audio_chunks_dir / "quarantine"
+    pause_for_chunk_review(quarantine_dir)
+
+    # Collect final chunk paths
+    chunk_paths = get_audio_files_in_directory(audio_chunks_dir)
+
+    if not chunk_paths:
+        logging.info(f"{RED}❌ No valid audio chunks found. Skipping concatenation and conversion.{RESET}")
+        return None, None, []
+
+    # Calculate timing
+    elapsed_total = time.time() - start_time
+    elapsed_td = timedelta(seconds=int(elapsed_total))
+
+    total_audio_duration_final = sum(get_chunk_audio_duration(chunk_path) for chunk_path in chunk_paths)
+    audio_duration_td = timedelta(seconds=int(total_audio_duration_final))
+    realtime_factor = total_audio_duration_final / elapsed_total if elapsed_total > 0 else 0.0
+
+    print(f"\n⏱️ TTS Processing Complete:")
+    print(f"   Elapsed Time: {CYAN}{str(elapsed_td)}{RESET}")
+    print(f"   Audio Duration: {GREEN}{str(audio_duration_td)}{RESET}")
+    print(f"   Realtime Factor: {YELLOW}{realtime_factor:.2f}x{RESET}")
+
+    # Combine audio
+    voice_name = voice_path.stem if hasattr(voice_path, 'stem') else Path(voice_path).stem
+    combined_wav_path = output_root / f"{book_dir.name} [{voice_name}].wav"
+    print("\n💾 Saving WAV file...")
+    combine_audio_chunks(chunk_paths, combined_wav_path)
+
+    # M4B conversion with normalization
+    temp_m4b_path = output_root / "output.m4b"
+    final_m4b_path = output_root / f"{book_dir.name}[{voice_name}].m4b"
+    convert_to_m4b(combined_wav_path, temp_m4b_path)
+    add_metadata_to_m4b(temp_m4b_path, final_m4b_path, cover_file, nfo_file)
+
+    logging.info(f"Audiobook created: {final_m4b_path}")
+
+    # Add final info to run log
+    run_log_lines.extend([
+        f"Combined WAV: {combined_wav_path}",
+        "--- Generation Settings ---",
+        f"Batch Processing: Enabled ({BATCH_SIZE} chunks per batch)",
+        f"ASR Enabled: {enable_asr_user or ENABLE_ASR} (user: {enable_asr_user}, global: {ENABLE_ASR})",
+        f"Hum Detection: {ENABLE_HUM_DETECTION}",
+        f"Dynamic Workers: {USE_DYNAMIC_WORKERS}",
+        f"Voice used: {voice_name}",
+        f"Exaggeration: {tts_params['exaggeration']}",
+        f"CFG weight: {tts_params['cfg_weight']}",
+        f"Temperature: {tts_params['temperature']}",
+        f"Processing Time: {str(elapsed_td)}",
+        f"Audio Duration: {str(audio_duration_td)}",
+        f"Realtime Factor: {realtime_factor:.2f}x",
+        f"Total Chunks: {len(chunk_paths)}"
+    ])
+
+    # Write the run log
+    log_run("\n".join(run_log_lines), output_root / "run.log")
+    print(f"📝 Run log written to: {output_root / 'run.log'}")
+
+    return final_m4b_path, combined_wav_path, run_log_lines

HF_Deploy/modules/voice_detector.py

@@ -0,0 +1,240 @@
+"""
+Voice Detection Module
+Handles voice detection from multiple sources: JSON metadata, log files, filenames
+"""
+
+import re
+import json
+from pathlib import Path
+from config.config import AUDIOBOOK_ROOT
+from modules.file_manager import list_voice_samples
+
+
+def get_likely_voices_for_book(book_name, chunks_json_path=None):
+    """
+    Get the most likely voice candidates for a book using the 3 detection methods:
+    1. JSON metadata/comments (if available)
+    2. run.log file 
+    3. Generated audiobook filenames (may return multiple)
+    
+    Returns: list of (voice_name, voice_path, detection_method) tuples
+    """
+    print(f"🔍 Finding likely voices for book: {book_name}")
+    likely_voices = []
+    
+    # Method 1: Check JSON metadata and comments
+    if chunks_json_path:
+        voice_from_json = get_voice_from_json(chunks_json_path)
+        if voice_from_json:
+            voice_path = find_voice_file_by_name(voice_from_json)
+            if voice_path:
+                likely_voices.append((voice_from_json, voice_path, "json_metadata"))
+                print(f"✅ Voice found in JSON: {voice_from_json}")
+    
+    # Method 2: Check run.log file
+    voice_from_log = get_voice_from_log(book_name)
+    if voice_from_log:
+        voice_path = find_voice_file_by_name(voice_from_log)
+        if voice_path:
+            # Avoid duplicates
+            if not any(v[0] == voice_from_log for v in likely_voices):
+                likely_voices.append((voice_from_log, voice_path, "run_log"))
+                print(f"✅ Voice found in run.log: {voice_from_log}")
+    
+    # Method 3: Check generated filename patterns (may find multiple)
+    voices_from_files = get_voices_from_filenames(book_name)
+    for voice_name in voices_from_files:
+        voice_path = find_voice_file_by_name(voice_name)
+        if voice_path:
+            # Avoid duplicates
+            if not any(v[0] == voice_name for v in likely_voices):
+                likely_voices.append((voice_name, voice_path, "filename_pattern"))
+                print(f"✅ Voice found in filename: {voice_name}")
+    
+    if not likely_voices:
+        print(f"⚠️ No likely voices detected for {book_name}")
+    else:
+        print(f"📋 Found {len(likely_voices)} likely voice candidates")
+    
+    return likely_voices
+
+def detect_voice_for_book(book_name, chunks_json_path=None):
+    """
+    Detect the most likely voice for a book (returns first candidate)
+    For backwards compatibility with existing code
+    """
+    likely_voices = get_likely_voices_for_book(book_name, chunks_json_path)
+    if likely_voices:
+        return likely_voices[0]  # Return the first (most likely) candidate
+    return None, None, "not_found"
+
+
+def get_voice_from_json(json_path):
+    """Extract voice information from JSON metadata"""
+    try:
+        with open(json_path, 'r', encoding='utf-8') as f:
+            content = f.read()
+        
+        # Check for voice metadata in JSON
+        if '"voice_used":' in content:
+            data = json.loads(content)
+            if isinstance(data, dict) and 'voice_used' in data:
+                return data['voice_used']
+            elif isinstance(data, list) and data and 'voice_used' in data[0]:
+                return data[0]['voice_used']
+        
+        # Check for voice as comment in JSON (fallback option)
+        voice_comment_match = re.search(r'//\s*voice:\s*([^\n]+)', content, re.IGNORECASE)
+        if voice_comment_match:
+            return voice_comment_match.group(1).strip()
+            
+    except Exception as e:
+        print(f"⚠️ Error reading JSON for voice info: {e}")
+    
+    return None
+
+
+def get_voice_from_log(book_name):
+    """Extract voice information from run.log file"""
+    audiobook_root = Path(AUDIOBOOK_ROOT)
+    log_file = audiobook_root / book_name / "run.log"
+    
+    if log_file.exists():
+        try:
+            with open(log_file, 'r', encoding='utf-8') as f:
+                for line in f:
+                    line = line.strip()
+                    if line.startswith("Voice: ") or line.startswith("Voice used: "):
+                        voice_name = line.split(": ", 1)[1].strip()
+                        return voice_name
+        except Exception as e:
+            print(f"⚠️ Error reading run log: {e}")
+    
+    return None
+
+
+def get_voices_from_filenames(book_name):
+    """Extract voice names from existing audiobook filename patterns (may return multiple)"""
+    audiobook_root = Path(AUDIOBOOK_ROOT)
+    book_dir = audiobook_root / book_name
+    
+    if not book_dir.exists():
+        return []
+    
+    found_voices = []
+    
+    # Look for WAV files with voice pattern: BookName [VoiceName].wav
+    for wav_file in book_dir.glob("*.wav"):
+        match = re.search(r'\[([^\]]+)\]\.wav$', wav_file.name)
+        if match:
+            voice_name = match.group(1)
+            if voice_name not in found_voices:
+                found_voices.append(voice_name)
+    
+    # Look for M4B files with voice pattern: BookName[VoiceName].m4b  
+    for m4b_file in book_dir.glob("*.m4b"):
+        match = re.search(r'\[([^\]]+)\]\.m4b$', m4b_file.name)
+        if match:
+            voice_name = match.group(1)
+            if voice_name not in found_voices:
+                found_voices.append(voice_name)
+    
+    return found_voices
+
+def get_voice_from_filename(book_name):
+    """Extract voice name from existing audiobook filename patterns (backwards compatibility)"""
+    voices = get_voices_from_filenames(book_name)
+    return voices[0] if voices else None
+
+
+def find_voice_file_by_name(voice_name):
+    """Find voice file by name in Voice_Samples directory"""
+    voice_files = list_voice_samples()
+    
+    # Exact match first
+    for voice_file in voice_files:
+        if voice_file.stem == voice_name:
+            return voice_file
+    
+    # Partial match (case insensitive)
+    voice_name_lower = voice_name.lower()
+    for voice_file in voice_files:
+        if voice_name_lower in voice_file.stem.lower():
+            return voice_file
+    
+    return None
+
+
+
+
+def add_voice_to_json(json_path, voice_name, method="metadata"):
+    """
+    Add voice information to JSON file
+    
+    method options:
+    - "metadata": Add as top-level metadata
+    - "comment": Add as comment that doesn't affect parsing
+    """
+    try:
+        with open(json_path, 'r', encoding='utf-8') as f:
+            content = f.read()
+        
+        if method == "metadata":
+            # Add voice as metadata to JSON structure
+            data = json.loads(content)
+            
+            if isinstance(data, list):
+                # For list format, add metadata as first element or update existing
+                if data and isinstance(data[0], dict) and not any(key.startswith('text') for key in data[0].keys()):
+                    # First element is already metadata
+                    data[0]['voice_used'] = voice_name
+                else:
+                    # Insert metadata as first element
+                    metadata = {"voice_used": voice_name, "_metadata": True}
+                    data.insert(0, metadata)
+            elif isinstance(data, dict):
+                # For dict format, add to top level
+                data['voice_used'] = voice_name
+            
+            # Save updated JSON
+            with open(json_path, 'w', encoding='utf-8') as f:
+                json.dump(data, f, indent=2, ensure_ascii=False)
+                
+        elif method == "comment":
+            # Add voice as comment at the top of file
+            voice_comment = f"// voice: {voice_name}\n"
+            
+            if not content.startswith("// voice:"):
+                content = voice_comment + content
+                with open(json_path, 'w', encoding='utf-8') as f:
+                    f.write(content)
+        
+        print(f"✅ Added voice '{voice_name}' to {json_path.name} using {method} method")
+        return True
+        
+    except Exception as e:
+        print(f"❌ Error adding voice to JSON: {e}")
+        return False
+
+
+def remove_voice_comment_from_json(json_path):
+    """Remove voice comment from JSON file for clean processing"""
+    try:
+        with open(json_path, 'r', encoding='utf-8') as f:
+            content = f.read()
+        
+        # Remove voice comment lines
+        lines = content.split('\n')
+        filtered_lines = [line for line in lines if not line.strip().startswith('// voice:')]
+        
+        if len(filtered_lines) != len(lines):
+            # Comments were removed, save cleaned version
+            cleaned_content = '\n'.join(filtered_lines)
+            with open(json_path, 'w', encoding='utf-8') as f:
+                f.write(cleaned_content)
+            return True
+            
+    except Exception as e:
+        print(f"⚠️ Error cleaning JSON comments: {e}")
+    
+    return False

HF_Deploy/requirements.txt

@@ -0,0 +1,56 @@
+# ChatterboxTTS HuggingFace Spaces Requirements
+# Optimized for HF Spaces environment with flexible versions
+
+# Core ML and TTS - Essential (pinned versions for fast builds)
+torch==2.6.0
+torchaudio==2.6.0
+transformers==4.46.3
+huggingface_hub>=0.15.0
+safetensors>=0.3.0
+
+# Audio processing - Required
+soundfile>=0.12.0
+librosa>=0.9.0
+pydub>=0.25.0
+audioread>=3.0.0
+
+# ASR System - Intelligent ASR with fallback
+openai-whisper>=20231117
+
+# System monitoring and resource detection
+psutil>=5.8.0
+pynvml>=11.0.0
+
+# Core scientific computing (pinned for fast builds)
+numpy==2.2.0
+scipy>=1.7.0
+
+# Text processing
+regex>=2023.0.0
+vaderSentiment>=3.3.0
+
+# Web interface - Gradio (let HF manage version)
+gradio>=4.0.0
+
+# Progress and logging
+tqdm>=4.60.0
+
+# File handling
+pathlib2>=2.3.0
+
+# Configuration and utilities
+python-dotenv>=1.0.0
+
+# Optional utilities
+requests>=2.25.0
+packaging>=21.0
+
+# Core ChatterboxTTS model dependencies
+chatterbox-tts>=0.1.2
+resemble-perth>=1.0.1
+omegaconf>=2.3.0
+einops>=0.6.0
+diffusers>=0.21.0
+tokenizers>=0.13.0
+conformer>=0.3.0
+s3tokenizer==0.2.0

HF_Deploy/src/chatterbox/__init__.py

@@ -0,0 +1,2 @@
+from .tts import ChatterboxTTS
+from .vc import ChatterboxVC

HF_Deploy/src/chatterbox/models/s3gen/__init__.py

@@ -0,0 +1,2 @@
+from .s3gen import S3Token2Wav as S3Gen
+from .const import S3GEN_SR

HF_Deploy/src/chatterbox/models/s3gen/const.py

@@ -0,0 +1 @@
+S3GEN_SR = 24000

HF_Deploy/src/chatterbox/models/s3gen/decoder.py

@@ -0,0 +1,317 @@
+# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import pack, rearrange, repeat
+
+from .utils.mask import add_optional_chunk_mask
+from .matcha.decoder import SinusoidalPosEmb, Block1D, ResnetBlock1D, Downsample1D, \
+    TimestepEmbedding, Upsample1D
+from .matcha.transformer import BasicTransformerBlock
+
+
+def mask_to_bias(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    assert mask.dtype == torch.bool
+    assert dtype in [torch.float32, torch.bfloat16, torch.float16]
+    mask = mask.to(dtype)
+    # attention mask bias
+    # NOTE(Mddct): torch.finfo jit issues
+    #     chunk_masks = (1.0 - chunk_masks) * torch.finfo(dtype).min
+    mask = (1.0 - mask) * -1.0e+10
+    return mask
+
+
+
+class Transpose(torch.nn.Module):
+    def __init__(self, dim0: int, dim1: int):
+        super().__init__()
+        self.dim0 = dim0
+        self.dim1 = dim1
+
+    def forward(self, x: torch.Tensor):
+        x = torch.transpose(x, self.dim0, self.dim1)
+        return x
+
+
+class CausalBlock1D(Block1D):
+    def __init__(self, dim: int, dim_out: int):
+        super(CausalBlock1D, self).__init__(dim, dim_out)
+        self.block = torch.nn.Sequential(
+            CausalConv1d(dim, dim_out, 3),
+            Transpose(1, 2),
+            nn.LayerNorm(dim_out),
+            Transpose(1, 2),
+            nn.Mish(),
+        )
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor):
+        output = self.block(x * mask)
+        return output * mask
+
+
+class CausalResnetBlock1D(ResnetBlock1D):
+    def __init__(self, dim: int, dim_out: int, time_emb_dim: int, groups: int = 8):
+        super(CausalResnetBlock1D, self).__init__(dim, dim_out, time_emb_dim, groups)
+        self.block1 = CausalBlock1D(dim, dim_out)
+        self.block2 = CausalBlock1D(dim_out, dim_out)
+
+
+class CausalConv1d(torch.nn.Conv1d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: str = 'zeros',
+        device=None,
+        dtype=None
+    ) -> None:
+        super(CausalConv1d, self).__init__(in_channels, out_channels,
+                                           kernel_size, stride,
+                                           padding=0, dilation=dilation,
+                                           groups=groups, bias=bias,
+                                           padding_mode=padding_mode,
+                                           device=device, dtype=dtype)
+        assert stride == 1
+        self.causal_padding = (kernel_size - 1, 0)
+
+    def forward(self, x: torch.Tensor):
+        x = F.pad(x, self.causal_padding)
+        x = super(CausalConv1d, self).forward(x)
+        return x
+
+
+class ConditionalDecoder(nn.Module):
+    def __init__(
+        self,
+        in_channels=320,
+        out_channels=80,
+        causal=True,
+        channels=[256],
+        dropout=0.0,
+        attention_head_dim=64,
+        n_blocks=4,
+        num_mid_blocks=12,
+        num_heads=8,
+        act_fn="gelu",
+    ):
+        """
+        This decoder requires an input with the same shape of the target. So, if your text content
+        is shorter or longer than the outputs, please re-sampling it before feeding to the decoder.
+        """
+        super().__init__()
+        channels = tuple(channels)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.causal = causal
+        self.time_embeddings = SinusoidalPosEmb(in_channels)
+        time_embed_dim = channels[0] * 4
+        self.time_mlp = TimestepEmbedding(
+            in_channels=in_channels,
+            time_embed_dim=time_embed_dim,
+            act_fn="silu",
+        )
+        self.down_blocks = nn.ModuleList([])
+        self.mid_blocks = nn.ModuleList([])
+        self.up_blocks = nn.ModuleList([])
+
+        # NOTE jrm: `static_chunk_size` is missing?
+        self.static_chunk_size = 0
+
+        output_channel = in_channels
+        for i in range(len(channels)):  # pylint: disable=consider-using-enumerate
+            input_channel = output_channel
+            output_channel = channels[i]
+            is_last = i == len(channels) - 1
+            resnet = CausalResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) if self.causal else \
+                ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim)
+            transformer_blocks = nn.ModuleList(
+                [
+                    BasicTransformerBlock(
+                        dim=output_channel,
+                        num_attention_heads=num_heads,
+                        attention_head_dim=attention_head_dim,
+                        dropout=dropout,
+                        activation_fn=act_fn,
+                    )
+                    for _ in range(n_blocks)
+                ]
+            )
+            downsample = (
+                Downsample1D(output_channel) if not is_last else
+                CausalConv1d(output_channel, output_channel, 3) if self.causal else nn.Conv1d(output_channel, output_channel, 3, padding=1)
+            )
+            self.down_blocks.append(nn.ModuleList([resnet, transformer_blocks, downsample]))
+
+        for _ in range(num_mid_blocks):
+            input_channel = channels[-1]
+            out_channels = channels[-1]
+            resnet = CausalResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) if self.causal else \
+                ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim)
+
+            transformer_blocks = nn.ModuleList(
+                [
+                    BasicTransformerBlock(
+                        dim=output_channel,
+                        num_attention_heads=num_heads,
+                        attention_head_dim=attention_head_dim,
+                        dropout=dropout,
+                        activation_fn=act_fn,
+                    )
+                    for _ in range(n_blocks)
+                ]
+            )
+
+            self.mid_blocks.append(nn.ModuleList([resnet, transformer_blocks]))
+
+        channels = channels[::-1] + (channels[0],)
+        for i in range(len(channels) - 1):
+            input_channel = channels[i] * 2
+            output_channel = channels[i + 1]
+            is_last = i == len(channels) - 2
+            resnet = CausalResnetBlock1D(
+                dim=input_channel,
+                dim_out=output_channel,
+                time_emb_dim=time_embed_dim,
+            ) if self.causal else ResnetBlock1D(
+                dim=input_channel,
+                dim_out=output_channel,
+                time_emb_dim=time_embed_dim,
+            )
+            transformer_blocks = nn.ModuleList(
+                [
+                    BasicTransformerBlock(
+                        dim=output_channel,
+                        num_attention_heads=num_heads,
+                        attention_head_dim=attention_head_dim,
+                        dropout=dropout,
+                        activation_fn=act_fn,
+                    )
+                    for _ in range(n_blocks)
+                ]
+            )
+            upsample = (
+                Upsample1D(output_channel, use_conv_transpose=True)
+                if not is_last
+                else CausalConv1d(output_channel, output_channel, 3) if self.causal else nn.Conv1d(output_channel, output_channel, 3, padding=1)
+            )
+            self.up_blocks.append(nn.ModuleList([resnet, transformer_blocks, upsample]))
+        self.final_block = CausalBlock1D(channels[-1], channels[-1]) if self.causal else Block1D(channels[-1], channels[-1])
+        self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d):
+                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.GroupNorm):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, x, mask, mu, t, spks=None, cond=None):
+        """Forward pass of the UNet1DConditional model.
+
+        Args:
+            x (torch.Tensor): shape (batch_size, in_channels, time)
+            mask (_type_): shape (batch_size, 1, time)
+            t (_type_): shape (batch_size)
+            spks (_type_, optional): shape: (batch_size, condition_channels). Defaults to None.
+            cond (_type_, optional): placeholder for future use. Defaults to None.
+
+        Raises:
+            ValueError: _description_
+            ValueError: _description_
+
+        Returns:
+            _type_: _description_
+        """
+
+        t = self.time_embeddings(t).to(t.dtype)
+        t = self.time_mlp(t)
+
+        x = pack([x, mu], "b * t")[0]
+
+        if spks is not None:
+            spks = repeat(spks, "b c -> b c t", t=x.shape[-1])
+            x = pack([x, spks], "b * t")[0]
+        if cond is not None:
+            x = pack([x, cond], "b * t")[0]
+
+        hiddens = []
+        masks = [mask]
+        for resnet, transformer_blocks, downsample in self.down_blocks:
+            mask_down = masks[-1]
+            x = resnet(x, mask_down, t)
+            x = rearrange(x, "b c t -> b t c").contiguous()
+            # attn_mask = torch.matmul(mask_down.transpose(1, 2).contiguous(), mask_down)
+            attn_mask = add_optional_chunk_mask(x, mask_down.bool(), False, False, 0, self.static_chunk_size, -1)
+            attn_mask = mask_to_bias(attn_mask == 1, x.dtype)
+            for transformer_block in transformer_blocks:
+                x = transformer_block(
+                    hidden_states=x,
+                    attention_mask=attn_mask,
+                    timestep=t,
+                )
+            x = rearrange(x, "b t c -> b c t").contiguous()
+            hiddens.append(x)  # Save hidden states for skip connections
+            x = downsample(x * mask_down)
+            masks.append(mask_down[:, :, ::2])
+        masks = masks[:-1]
+        mask_mid = masks[-1]
+
+        for resnet, transformer_blocks in self.mid_blocks:
+            x = resnet(x, mask_mid, t)
+            x = rearrange(x, "b c t -> b t c").contiguous()
+            # attn_mask = torch.matmul(mask_mid.transpose(1, 2).contiguous(), mask_mid)
+            attn_mask = add_optional_chunk_mask(x, mask_mid.bool(), False, False, 0, self.static_chunk_size, -1)
+            attn_mask = mask_to_bias(attn_mask == 1, x.dtype)
+            for transformer_block in transformer_blocks:
+                x = transformer_block(
+                    hidden_states=x,
+                    attention_mask=attn_mask,
+                    timestep=t,
+                )
+            x = rearrange(x, "b t c -> b c t").contiguous()
+
+        for resnet, transformer_blocks, upsample in self.up_blocks:
+            mask_up = masks.pop()
+            skip = hiddens.pop()
+            x = pack([x[:, :, :skip.shape[-1]], skip], "b * t")[0]
+            x = resnet(x, mask_up, t)
+            x = rearrange(x, "b c t -> b t c").contiguous()
+            # attn_mask = torch.matmul(mask_up.transpose(1, 2).contiguous(), mask_up)
+            attn_mask = add_optional_chunk_mask(x, mask_up.bool(), False, False, 0, self.static_chunk_size, -1)
+            attn_mask = mask_to_bias(attn_mask == 1, x.dtype)
+            for transformer_block in transformer_blocks:
+                x = transformer_block(
+                    hidden_states=x,
+                    attention_mask=attn_mask,
+                    timestep=t,
+                )
+            x = rearrange(x, "b t c -> b c t").contiguous()
+            x = upsample(x * mask_up)
+        x = self.final_block(x, mask_up)
+        output = self.final_proj(x * mask_up)
+        return output * mask

HF_Deploy/src/chatterbox/models/s3gen/f0_predictor.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+import torch.nn as nn
+from torch.nn.utils.parametrizations import weight_norm
+
+
+class ConvRNNF0Predictor(nn.Module):
+    def __init__(self,
+                 num_class: int = 1,
+                 in_channels: int = 80,
+                 cond_channels: int = 512
+                 ):
+        super().__init__()
+
+        self.num_class = num_class
+        self.condnet = nn.Sequential(
+            weight_norm(
+                nn.Conv1d(in_channels, cond_channels, kernel_size=3, padding=1)
+            ),
+            nn.ELU(),
+            weight_norm(
+                nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)
+            ),
+            nn.ELU(),
+            weight_norm(
+                nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)
+            ),
+            nn.ELU(),
+            weight_norm(
+                nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)
+            ),
+            nn.ELU(),
+            weight_norm(
+                nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)
+            ),
+            nn.ELU(),
+        )
+        self.classifier = nn.Linear(in_features=cond_channels, out_features=self.num_class)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.condnet(x)
+        x = x.transpose(1, 2)
+        return torch.abs(self.classifier(x).squeeze(-1))

HF_Deploy/src/chatterbox/models/s3gen/flow.py

@@ -0,0 +1,242 @@
+# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import logging
+import random
+from typing import Dict, Optional
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from omegaconf import DictConfig
+from .utils.mask import make_pad_mask
+
+
+class MaskedDiffWithXvec(torch.nn.Module):
+    def __init__(self,
+                 input_size: int = 512,
+                 output_size: int = 80,
+                 spk_embed_dim: int = 192,
+                 output_type: str = "mel",
+                 vocab_size: int = 4096,
+                 input_frame_rate: int = 50,
+                 only_mask_loss: bool = True,
+                 encoder: torch.nn.Module = None,
+                 length_regulator: torch.nn.Module = None,
+                 decoder: torch.nn.Module = None,
+                 decoder_conf: Dict = {'in_channels': 240, 'out_channel': 80, 'spk_emb_dim': 80, 'n_spks': 1,
+                                       'cfm_params': DictConfig({'sigma_min': 1e-06, 'solver': 'euler', 't_scheduler': 'cosine',
+                                                                 'training_cfg_rate': 0.2, 'inference_cfg_rate': 0.7, 'reg_loss_type': 'l1'}),
+                                       'decoder_params': {'channels': [256, 256], 'dropout': 0.0, 'attention_head_dim': 64,
+                                                          'n_blocks': 4, 'num_mid_blocks': 12, 'num_heads': 8, 'act_fn': 'gelu'}},
+                 mel_feat_conf: Dict = {'n_fft': 1024, 'num_mels': 80, 'sampling_rate': 22050,
+                                        'hop_size': 256, 'win_size': 1024, 'fmin': 0, 'fmax': 8000}):
+        super().__init__()
+        self.input_size = input_size
+        self.output_size = output_size
+        self.decoder_conf = decoder_conf
+        self.mel_feat_conf = mel_feat_conf
+        self.vocab_size = vocab_size
+        self.output_type = output_type
+        self.input_frame_rate = input_frame_rate
+        logging.info(f"input frame rate={self.input_frame_rate}")
+        self.input_embedding = nn.Embedding(vocab_size, input_size)
+        self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, output_size)
+        self.encoder = encoder
+        self.encoder_proj = torch.nn.Linear(self.encoder.output_size(), output_size)
+        self.decoder = decoder
+        self.length_regulator = length_regulator
+        self.only_mask_loss = only_mask_loss
+
+    def forward(
+            self,
+            batch: dict,
+            device: torch.device,
+    ) -> Dict[str, Optional[torch.Tensor]]:
+        token = batch['speech_token'].to(device)
+        token_len = batch['speech_token_len'].to(device)
+        feat = batch['speech_feat'].to(device)
+        feat_len = batch['speech_feat_len'].to(device)
+        embedding = batch['embedding'].to(device)
+
+        # xvec projection
+        embedding = F.normalize(embedding, dim=1)
+        embedding = self.spk_embed_affine_layer(embedding)
+
+        # concat text and prompt_text
+        mask = (~make_pad_mask(token_len)).float().unsqueeze(-1).to(device)
+        token = self.input_embedding(torch.clamp(token, min=0)) * mask
+
+        # text encode
+        h, h_lengths = self.encoder(token, token_len)
+        h = self.encoder_proj(h)
+        h, h_lengths = self.length_regulator(h, feat_len)
+
+        # get conditions
+        conds = torch.zeros(feat.shape, device=token.device)
+        for i, j in enumerate(feat_len):
+            if random.random() < 0.5:
+                continue
+            index = random.randint(0, int(0.3 * j))
+            conds[i, :index] = feat[i, :index]
+        conds = conds.transpose(1, 2)
+
+        mask = (~make_pad_mask(feat_len)).to(h)
+        feat = F.interpolate(feat.unsqueeze(dim=1), size=h.shape[1:], mode="nearest").squeeze(dim=1)
+        loss, _ = self.decoder.compute_loss(
+            feat.transpose(1, 2).contiguous(),
+            mask.unsqueeze(1),
+            h.transpose(1, 2).contiguous(),
+            embedding,
+            cond=conds
+        )
+        return {'loss': loss}
+
+    @torch.inference_mode()
+    def inference(self,
+                  token,
+                  token_len,
+                  prompt_token,
+                  prompt_token_len,
+                  prompt_feat,
+                  prompt_feat_len,
+                  embedding,
+                  flow_cache):
+        if self.fp16 is True:
+            prompt_feat = prompt_feat.half()
+            embedding = embedding.half()
+
+        assert token.shape[0] == 1
+        # xvec projection
+        embedding = F.normalize(embedding, dim=1)
+        embedding = self.spk_embed_affine_layer(embedding)
+
+        # concat text and prompt_text
+        token_len1, token_len2 = prompt_token.shape[1], token.shape[1]
+        token, token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + token_len
+        mask = (~make_pad_mask(token_len)).unsqueeze(-1).to(embedding)
+        token = self.input_embedding(torch.clamp(token, min=0)) * mask
+
+        # text encode
+        h, h_lengths = self.encoder(token, token_len)
+        h = self.encoder_proj(h)
+        mel_len1, mel_len2 = prompt_feat.shape[1], int(token_len2 / self.input_frame_rate * 22050 / 256)
+        h, h_lengths = self.length_regulator.inference(h[:, :token_len1], h[:, token_len1:], mel_len1, mel_len2, self.input_frame_rate)
+
+        # get conditions
+        conds = torch.zeros([1, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype)
+        conds[:, :mel_len1] = prompt_feat
+        conds = conds.transpose(1, 2)
+
+        mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h)
+        feat, flow_cache = self.decoder(
+            mu=h.transpose(1, 2).contiguous(),
+            mask=mask.unsqueeze(1),
+            spks=embedding,
+            cond=conds,
+            n_timesteps=10,
+            prompt_len=mel_len1,
+            flow_cache=flow_cache
+        )
+        feat = feat[:, :, mel_len1:]
+        assert feat.shape[2] == mel_len2
+        return feat.float(), flow_cache
+
+
+class CausalMaskedDiffWithXvec(torch.nn.Module):
+    def __init__(self,
+                 input_size: int = 512,
+                 output_size: int = 80,
+                 spk_embed_dim: int = 192,
+                 output_type: str = "mel",
+                 vocab_size: int = 6561,
+                 input_frame_rate: int = 25,
+                 only_mask_loss: bool = True,
+                 token_mel_ratio: int = 2,
+                 pre_lookahead_len: int = 3,
+                 encoder: torch.nn.Module = None,
+                 decoder: torch.nn.Module = None,
+                 decoder_conf: Dict = {'in_channels': 240, 'out_channel': 80, 'spk_emb_dim': 80, 'n_spks': 1,
+                                       'cfm_params': DictConfig({'sigma_min': 1e-06, 'solver': 'euler', 't_scheduler': 'cosine',
+                                                                 'training_cfg_rate': 0.2, 'inference_cfg_rate': 0.7, 'reg_loss_type': 'l1'}),
+                                       'decoder_params': {'channels': [256, 256], 'dropout': 0.0, 'attention_head_dim': 64,
+                                                          'n_blocks': 4, 'num_mid_blocks': 12, 'num_heads': 8, 'act_fn': 'gelu'}},
+                 mel_feat_conf: Dict = {'n_fft': 1024, 'num_mels': 80, 'sampling_rate': 22050,
+                                        'hop_size': 256, 'win_size': 1024, 'fmin': 0, 'fmax': 8000}):
+        super().__init__()
+        self.input_size = input_size
+        self.output_size = output_size
+        self.decoder_conf = decoder_conf
+        self.mel_feat_conf = mel_feat_conf
+        self.vocab_size = vocab_size
+        self.output_type = output_type
+        self.input_frame_rate = input_frame_rate
+        logging.info(f"input frame rate={self.input_frame_rate}")
+        self.input_embedding = nn.Embedding(vocab_size, input_size)
+        self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, output_size)
+        self.encoder = encoder
+        self.encoder_proj = torch.nn.Linear(self.encoder.output_size(), output_size)
+        self.decoder = decoder
+        self.only_mask_loss = only_mask_loss
+        self.token_mel_ratio = token_mel_ratio
+        self.pre_lookahead_len = pre_lookahead_len
+
+        # FIXME: this was missing - just putting it in as false
+        self.fp16 = False
+
+    @torch.inference_mode()
+    def inference(self,
+                  token,
+                  token_len,
+                  prompt_token,
+                  prompt_token_len,
+                  prompt_feat,
+                  prompt_feat_len,
+                  embedding,
+                  finalize):
+        if self.fp16 is True:
+            prompt_feat = prompt_feat.half()
+            embedding = embedding.half()
+
+        assert token.shape[0] == 1
+        # xvec projection
+        embedding = F.normalize(embedding, dim=1)
+        embedding = self.spk_embed_affine_layer(embedding)
+
+        # concat text and prompt_text
+        token, token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + token_len
+        mask = (~make_pad_mask(token_len)).unsqueeze(-1).to(embedding)
+        token = self.input_embedding(torch.clamp(token, min=0)) * mask
+
+        # text encode
+        h, h_lengths = self.encoder(token, token_len)
+        if finalize is False:
+            h = h[:, :-self.pre_lookahead_len * self.token_mel_ratio]
+        mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1]
+        h = self.encoder_proj(h)
+
+        # get conditions
+        conds = torch.zeros([1, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype)
+        conds[:, :mel_len1] = prompt_feat
+        conds = conds.transpose(1, 2)
+
+        mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h)
+        feat, _ = self.decoder(
+            mu=h.transpose(1, 2).contiguous(),
+            mask=mask.unsqueeze(1),
+            spks=embedding,
+            cond=conds,
+            n_timesteps=10
+        )
+        feat = feat[:, :, mel_len1:]
+        assert feat.shape[2] == mel_len2
+        return feat.float(), None  # NOTE jrm: why are they returning None here?

HF_Deploy/src/chatterbox/models/s3gen/flow_matching.py

@@ -0,0 +1,228 @@
+# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import threading
+import torch
+import torch.nn.functional as F
+from .matcha.flow_matching import BASECFM
+from omegaconf import OmegaConf
+
+
+CFM_PARAMS = OmegaConf.create({
+    "sigma_min": 1e-06,
+    "solver": "euler",
+    "t_scheduler": "cosine",
+    "training_cfg_rate": 0.2,
+    "inference_cfg_rate": 0.7,
+    "reg_loss_type": "l1"
+})
+
+
+class ConditionalCFM(BASECFM):
+    def __init__(self, in_channels, cfm_params, n_spks=1, spk_emb_dim=64, estimator: torch.nn.Module = None):
+        super().__init__(
+            n_feats=in_channels,
+            cfm_params=cfm_params,
+            n_spks=n_spks,
+            spk_emb_dim=spk_emb_dim,
+        )
+        self.t_scheduler = cfm_params.t_scheduler
+        self.training_cfg_rate = cfm_params.training_cfg_rate
+        self.inference_cfg_rate = cfm_params.inference_cfg_rate
+        in_channels = in_channels + (spk_emb_dim if n_spks > 0 else 0)
+        # Just change the architecture of the estimator here
+        self.estimator = estimator
+        self.lock = threading.Lock()
+
+    @torch.inference_mode()
+    def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, prompt_len=0, flow_cache=torch.zeros(1, 80, 0, 2)):
+        """Forward diffusion
+
+        Args:
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): output_mask
+                shape: (batch_size, 1, mel_timesteps)
+            n_timesteps (int): number of diffusion steps
+            temperature (float, optional): temperature for scaling noise. Defaults to 1.0.
+            spks (torch.Tensor, optional): speaker ids. Defaults to None.
+                shape: (batch_size, spk_emb_dim)
+            cond: Not used but kept for future purposes
+
+        Returns:
+            sample: generated mel-spectrogram
+                shape: (batch_size, n_feats, mel_timesteps)
+        """
+
+        z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature
+        cache_size = flow_cache.shape[2]
+        # fix prompt and overlap part mu and z
+        if cache_size != 0:
+            z[:, :, :cache_size] = flow_cache[:, :, :, 0]
+            mu[:, :, :cache_size] = flow_cache[:, :, :, 1]
+        z_cache = torch.concat([z[:, :, :prompt_len], z[:, :, -34:]], dim=2)
+        mu_cache = torch.concat([mu[:, :, :prompt_len], mu[:, :, -34:]], dim=2)
+        flow_cache = torch.stack([z_cache, mu_cache], dim=-1)
+
+        t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype)
+        if self.t_scheduler == 'cosine':
+            t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
+        return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), flow_cache
+
+    def solve_euler(self, x, t_span, mu, mask, spks, cond):
+        """
+        Fixed euler solver for ODEs.
+        Args:
+            x (torch.Tensor): random noise
+            t_span (torch.Tensor): n_timesteps interpolated
+                shape: (n_timesteps + 1,)
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): output_mask
+                shape: (batch_size, 1, mel_timesteps)
+            spks (torch.Tensor, optional): speaker ids. Defaults to None.
+                shape: (batch_size, spk_emb_dim)
+            cond: Not used but kept for future purposes
+        """
+        t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]
+        t = t.unsqueeze(dim=0)
+
+        # I am storing this because I can later plot it by putting a debugger here and saving it to a file
+        # Or in future might add like a return_all_steps flag
+        sol = []
+
+        # Do not use concat, it may cause memory format changed and trt infer with wrong results!
+        x_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=x.dtype)
+        mask_in = torch.zeros([2, 1, x.size(2)], device=x.device, dtype=x.dtype)
+        mu_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=x.dtype)
+        t_in = torch.zeros([2], device=x.device, dtype=x.dtype)
+        spks_in = torch.zeros([2, 80], device=x.device, dtype=x.dtype)
+        cond_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=x.dtype)
+        for step in range(1, len(t_span)):
+            # Classifier-Free Guidance inference introduced in VoiceBox
+            x_in[:] = x
+            mask_in[:] = mask
+            mu_in[0] = mu
+            t_in[:] = t.unsqueeze(0)
+            spks_in[0] = spks
+            cond_in[0] = cond
+            dphi_dt = self.forward_estimator(
+                x_in, mask_in,
+                mu_in, t_in,
+                spks_in,
+                cond_in
+            )
+            dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [x.size(0), x.size(0)], dim=0)
+            dphi_dt = ((1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt)
+            x = x + dt * dphi_dt
+            t = t + dt
+            sol.append(x)
+            if step < len(t_span) - 1:
+                dt = t_span[step + 1] - t
+
+        return sol[-1].float()
+
+    def forward_estimator(self, x, mask, mu, t, spks, cond):
+        if isinstance(self.estimator, torch.nn.Module):
+            return self.estimator.forward(x, mask, mu, t, spks, cond)
+        else:
+            with self.lock:
+                self.estimator.set_input_shape('x', (2, 80, x.size(2)))
+                self.estimator.set_input_shape('mask', (2, 1, x.size(2)))
+                self.estimator.set_input_shape('mu', (2, 80, x.size(2)))
+                self.estimator.set_input_shape('t', (2,))
+                self.estimator.set_input_shape('spks', (2, 80))
+                self.estimator.set_input_shape('cond', (2, 80, x.size(2)))
+                # run trt engine
+                self.estimator.execute_v2([x.contiguous().data_ptr(),
+                                           mask.contiguous().data_ptr(),
+                                           mu.contiguous().data_ptr(),
+                                           t.contiguous().data_ptr(),
+                                           spks.contiguous().data_ptr(),
+                                           cond.contiguous().data_ptr(),
+                                           x.data_ptr()])
+            return x
+
+    def compute_loss(self, x1, mask, mu, spks=None, cond=None):
+        """Computes diffusion loss
+
+        Args:
+            x1 (torch.Tensor): Target
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): target mask
+                shape: (batch_size, 1, mel_timesteps)
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            spks (torch.Tensor, optional): speaker embedding. Defaults to None.
+                shape: (batch_size, spk_emb_dim)
+
+        Returns:
+            loss: conditional flow matching loss
+            y: conditional flow
+                shape: (batch_size, n_feats, mel_timesteps)
+        """
+        b, _, t = mu.shape
+
+        # random timestep
+        t = torch.rand([b, 1, 1], device=mu.device, dtype=mu.dtype)
+        if self.t_scheduler == 'cosine':
+            t = 1 - torch.cos(t * 0.5 * torch.pi)
+        # sample noise p(x_0)
+        z = torch.randn_like(x1)
+
+        y = (1 - (1 - self.sigma_min) * t) * z + t * x1
+        u = x1 - (1 - self.sigma_min) * z
+
+        # during training, we randomly drop condition to trade off mode coverage and sample fidelity
+        if self.training_cfg_rate > 0:
+            cfg_mask = torch.rand(b, device=x1.device) > self.training_cfg_rate
+            mu = mu * cfg_mask.view(-1, 1, 1)
+            spks = spks * cfg_mask.view(-1, 1)
+            cond = cond * cfg_mask.view(-1, 1, 1)
+
+        pred = self.estimator(y, mask, mu, t.squeeze(), spks, cond)
+        loss = F.mse_loss(pred * mask, u * mask, reduction="sum") / (torch.sum(mask) * u.shape[1])
+        return loss, y
+
+
+class CausalConditionalCFM(ConditionalCFM):
+    def __init__(self, in_channels=240, cfm_params=CFM_PARAMS, n_spks=1, spk_emb_dim=80, estimator=None):
+        super().__init__(in_channels, cfm_params, n_spks, spk_emb_dim, estimator)
+        self.rand_noise = torch.randn([1, 80, 50 * 300])
+
+    @torch.inference_mode()
+    def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None):
+        """Forward diffusion
+
+        Args:
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): output_mask
+                shape: (batch_size, 1, mel_timesteps)
+            n_timesteps (int): number of diffusion steps
+            temperature (float, optional): temperature for scaling noise. Defaults to 1.0.
+            spks (torch.Tensor, optional): speaker ids. Defaults to None.
+                shape: (batch_size, spk_emb_dim)
+            cond: Not used but kept for future purposes
+
+        Returns:
+            sample: generated mel-spectrogram
+                shape: (batch_size, n_feats, mel_timesteps)
+        """
+
+        z = self.rand_noise[:, :, :mu.size(2)].to(mu.device).to(mu.dtype) * temperature
+        # fix prompt and overlap part mu and z
+        t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype)
+        if self.t_scheduler == 'cosine':
+            t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
+        return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None

HF_Deploy/src/chatterbox/models/s3gen/hifigan.py

@@ -0,0 +1,474 @@
+# jrm: adapted from CosyVoice/cosyvoice/hifigan/generator.py
+#      most modules should be reusable, but I found their SineGen changed a git.
+
+# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""HIFI-GAN"""
+
+from typing import Dict, Optional, List
+import numpy as np
+from scipy.signal import get_window
+import torch
+import torch.nn.functional as F
+from torch.nn import Conv1d
+from torch.nn import ConvTranspose1d
+from torch.nn.utils import remove_weight_norm
+from torch.nn.utils.parametrizations import weight_norm
+from torch.distributions.uniform import Uniform
+from torch import nn, sin, pow
+from torch.nn import Parameter
+
+
+class Snake(nn.Module):
+    '''
+    Implementation of a sine-based periodic activation function
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter
+    References:
+        - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snake(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    '''
+    def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):
+        '''
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha: trainable parameter
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            alpha will be trained along with the rest of your model.
+        '''
+        super(Snake, self).__init__()
+        self.in_features = in_features
+
+        # initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale: # log scale alphas initialized to zeros
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+        else: # linear scale alphas initialized to ones
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        '''
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        Snake ∶= x + 1/a * sin^2 (xa)
+        '''
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+        x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x
+
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+"""hifigan based generator implementation.
+
+This code is modified from https://github.com/jik876/hifi-gan
+ ,https://github.com/kan-bayashi/ParallelWaveGAN and
+ https://github.com/NVIDIA/BigVGAN
+
+"""
+
+
+class ResBlock(torch.nn.Module):
+    """Residual block module in HiFiGAN/BigVGAN."""
+    def __init__(
+        self,
+        channels: int = 512,
+        kernel_size: int = 3,
+        dilations: List[int] = [1, 3, 5],
+    ):
+        super(ResBlock, self).__init__()
+        self.convs1 = nn.ModuleList()
+        self.convs2 = nn.ModuleList()
+
+        for dilation in dilations:
+            self.convs1.append(
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation,
+                        padding=get_padding(kernel_size, dilation)
+                    )
+                )
+            )
+            self.convs2.append(
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1)
+                    )
+                )
+            )
+        self.convs1.apply(init_weights)
+        self.convs2.apply(init_weights)
+        self.activations1 = nn.ModuleList([
+            Snake(channels, alpha_logscale=False)
+            for _ in range(len(self.convs1))
+        ])
+        self.activations2 = nn.ModuleList([
+            Snake(channels, alpha_logscale=False)
+            for _ in range(len(self.convs2))
+        ])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for idx in range(len(self.convs1)):
+            xt = self.activations1[idx](x)
+            xt = self.convs1[idx](xt)
+            xt = self.activations2[idx](xt)
+            xt = self.convs2[idx](xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for idx in range(len(self.convs1)):
+            remove_weight_norm(self.convs1[idx])
+            remove_weight_norm(self.convs2[idx])
+
+
+class SineGen(torch.nn.Module):
+    """ Definition of sine generator
+    SineGen(samp_rate, harmonic_num = 0,
+            sine_amp = 0.1, noise_std = 0.003,
+            voiced_threshold = 0,
+            flag_for_pulse=False)
+    samp_rate: sampling rate in Hz
+    harmonic_num: number of harmonic overtones (default 0)
+    sine_amp: amplitude of sine-wavefrom (default 0.1)
+    noise_std: std of Gaussian noise (default 0.003)
+    voiced_thoreshold: F0 threshold for U/V classification (default 0)
+    flag_for_pulse: this SinGen is used inside PulseGen (default False)
+    Note: when flag_for_pulse is True, the first time step of a voiced
+        segment is always sin(np.pi) or cos(0)
+    """
+
+    def __init__(self, samp_rate, harmonic_num=0,
+                 sine_amp=0.1, noise_std=0.003,
+                 voiced_threshold=0):
+        super(SineGen, self).__init__()
+        self.sine_amp = sine_amp
+        self.noise_std = noise_std
+        self.harmonic_num = harmonic_num
+        self.sampling_rate = samp_rate
+        self.voiced_threshold = voiced_threshold
+
+    def _f02uv(self, f0):
+        # generate uv signal
+        uv = (f0 > self.voiced_threshold).type(torch.float32)
+        return uv
+
+    @torch.no_grad()
+    def forward(self, f0):
+        """
+        :param f0: [B, 1, sample_len], Hz
+        :return: [B, 1, sample_len]
+        """
+
+        F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to(f0.device)
+        for i in range(self.harmonic_num + 1):
+            F_mat[:, i: i + 1, :] = f0 * (i + 1) / self.sampling_rate
+
+        theta_mat = 2 * np.pi * (torch.cumsum(F_mat, dim=-1) % 1)
+        u_dist = Uniform(low=-np.pi, high=np.pi)
+        phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device)
+        phase_vec[:, 0, :] = 0
+
+        # generate sine waveforms
+        sine_waves = self.sine_amp * torch.sin(theta_mat + phase_vec)
+
+        # generate uv signal
+        uv = self._f02uv(f0)
+
+        # noise: for unvoiced should be similar to sine_amp
+        #        std = self.sine_amp/3 -> max value ~ self.sine_amp
+        # .       for voiced regions is self.noise_std
+        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
+        noise = noise_amp * torch.randn_like(sine_waves)
+
+        # first: set the unvoiced part to 0 by uv
+        # then: additive noise
+        sine_waves = sine_waves * uv + noise
+        return sine_waves, uv, noise
+
+
+class SourceModuleHnNSF(torch.nn.Module):
+    """ SourceModule for hn-nsf
+    SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0)
+    sampling_rate: sampling_rate in Hz
+    harmonic_num: number of harmonic above F0 (default: 0)
+    sine_amp: amplitude of sine source signal (default: 0.1)
+    add_noise_std: std of additive Gaussian noise (default: 0.003)
+        note that amplitude of noise in unvoiced is decided
+        by sine_amp
+    voiced_threshold: threhold to set U/V given F0 (default: 0)
+    Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+    F0_sampled (batchsize, length, 1)
+    Sine_source (batchsize, length, 1)
+    noise_source (batchsize, length 1)
+    uv (batchsize, length, 1)
+    """
+
+    def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0):
+        super(SourceModuleHnNSF, self).__init__()
+
+        self.sine_amp = sine_amp
+        self.noise_std = add_noise_std
+
+        # to produce sine waveforms
+        self.l_sin_gen = SineGen(sampling_rate, harmonic_num,
+                                 sine_amp, add_noise_std, voiced_threshod)
+
+        # to merge source harmonics into a single excitation
+        self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
+        self.l_tanh = torch.nn.Tanh()
+
+    def forward(self, x):
+        """
+        Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+        F0_sampled (batchsize, length, 1)
+        Sine_source (batchsize, length, 1)
+        noise_source (batchsize, length 1)
+        """
+        # source for harmonic branch
+        with torch.no_grad():
+            sine_wavs, uv, _ = self.l_sin_gen(x.transpose(1, 2))
+            sine_wavs = sine_wavs.transpose(1, 2)
+            uv = uv.transpose(1, 2)
+        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
+
+        # source for noise branch, in the same shape as uv
+        noise = torch.randn_like(uv) * self.sine_amp / 3
+        return sine_merge, noise, uv
+
+
+class HiFTGenerator(nn.Module):
+    """
+    HiFTNet Generator: Neural Source Filter + ISTFTNet
+    https://arxiv.org/abs/2309.09493
+    """
+    def __init__(
+            self,
+            in_channels: int = 80,
+            base_channels: int = 512,
+            nb_harmonics: int = 8,
+            sampling_rate: int = 22050,
+            nsf_alpha: float = 0.1,
+            nsf_sigma: float = 0.003,
+            nsf_voiced_threshold: float = 10,
+            upsample_rates: List[int] = [8, 8],
+            upsample_kernel_sizes: List[int] = [16, 16],
+            istft_params: Dict[str, int] = {"n_fft": 16, "hop_len": 4},
+            resblock_kernel_sizes: List[int] = [3, 7, 11],
+            resblock_dilation_sizes: List[List[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+            source_resblock_kernel_sizes: List[int] = [7, 11],
+            source_resblock_dilation_sizes: List[List[int]] = [[1, 3, 5], [1, 3, 5]],
+            lrelu_slope: float = 0.1,
+            audio_limit: float = 0.99,
+            f0_predictor: torch.nn.Module = None,
+    ):
+        super(HiFTGenerator, self).__init__()
+
+        self.out_channels = 1
+        self.nb_harmonics = nb_harmonics
+        self.sampling_rate = sampling_rate
+        self.istft_params = istft_params
+        self.lrelu_slope = lrelu_slope
+        self.audio_limit = audio_limit
+
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.m_source = SourceModuleHnNSF(
+            sampling_rate=sampling_rate,
+            upsample_scale=np.prod(upsample_rates) * istft_params["hop_len"],
+            harmonic_num=nb_harmonics,
+            sine_amp=nsf_alpha,
+            add_noise_std=nsf_sigma,
+            voiced_threshod=nsf_voiced_threshold)
+        self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * istft_params["hop_len"])
+
+        self.conv_pre = weight_norm(
+            Conv1d(in_channels, base_channels, 7, 1, padding=3)
+        )
+
+        # Up
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        base_channels // (2**i),
+                        base_channels // (2**(i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+
+        # Down
+        self.source_downs = nn.ModuleList()
+        self.source_resblocks = nn.ModuleList()
+        downsample_rates = [1] + upsample_rates[::-1][:-1]
+        downsample_cum_rates = np.cumprod(downsample_rates)
+        for i, (u, k, d) in enumerate(zip(downsample_cum_rates[::-1], source_resblock_kernel_sizes, source_resblock_dilation_sizes)):
+            if u == 1:
+                self.source_downs.append(
+                    Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), 1, 1)
+                )
+            else:
+                self.source_downs.append(
+                    Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), u * 2, u, padding=(u // 2))
+                )
+
+            self.source_resblocks.append(
+                ResBlock(base_channels // (2 ** (i + 1)), k, d)
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = base_channels // (2**(i + 1))
+            for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
+                self.resblocks.append(ResBlock(ch, k, d))
+
+        self.conv_post = weight_norm(Conv1d(ch, istft_params["n_fft"] + 2, 7, 1, padding=3))
+        self.ups.apply(init_weights)
+        self.conv_post.apply(init_weights)
+        self.reflection_pad = nn.ReflectionPad1d((1, 0))
+        self.stft_window = torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32))
+        self.f0_predictor = f0_predictor
+
+    def remove_weight_norm(self):
+        print('Removing weight norm...')
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+        remove_weight_norm(self.conv_pre)
+        remove_weight_norm(self.conv_post)
+        self.m_source.remove_weight_norm()
+        for l in self.source_downs:
+            remove_weight_norm(l)
+        for l in self.source_resblocks:
+            l.remove_weight_norm()
+
+    def _stft(self, x):
+        spec = torch.stft(
+            x,
+            self.istft_params["n_fft"], self.istft_params["hop_len"], self.istft_params["n_fft"], window=self.stft_window.to(x.device),
+            return_complex=True)
+        spec = torch.view_as_real(spec)  # [B, F, TT, 2]
+        return spec[..., 0], spec[..., 1]
+
+    def _istft(self, magnitude, phase):
+        magnitude = torch.clip(magnitude, max=1e2)
+        real = magnitude * torch.cos(phase)
+        img = magnitude * torch.sin(phase)
+        inverse_transform = torch.istft(torch.complex(real, img), self.istft_params["n_fft"], self.istft_params["hop_len"],
+                                        self.istft_params["n_fft"], window=self.stft_window.to(magnitude.device))
+        return inverse_transform
+
+    def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor:
+        s_stft_real, s_stft_imag = self._stft(s.squeeze(1))
+        s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1)
+
+        x = self.conv_pre(x)
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, self.lrelu_slope)
+            x = self.ups[i](x)
+
+            if i == self.num_upsamples - 1:
+                x = self.reflection_pad(x)
+
+            # fusion
+            si = self.source_downs[i](s_stft)
+            si = self.source_resblocks[i](si)
+            x = x + si
+
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        magnitude = torch.exp(x[:, :self.istft_params["n_fft"] // 2 + 1, :])
+        phase = torch.sin(x[:, self.istft_params["n_fft"] // 2 + 1:, :])  # actually, sin is redundancy
+
+        x = self._istft(magnitude, phase)
+        x = torch.clamp(x, -self.audio_limit, self.audio_limit)
+        return x
+
+    def forward(
+            self,
+            batch: dict,
+            device: torch.device,
+    ) -> Dict[str, Optional[torch.Tensor]]:
+        speech_feat = batch['speech_feat'].transpose(1, 2).to(device)
+        # mel->f0
+        f0 = self.f0_predictor(speech_feat)
+        # f0->source
+        s = self.f0_upsamp(f0[:, None]).transpose(1, 2)  # bs,n,t
+        s, _, _ = self.m_source(s)
+        s = s.transpose(1, 2)
+        # mel+source->speech
+        generated_speech = self.decode(x=speech_feat, s=s)
+        return generated_speech, f0
+
+    @torch.inference_mode()
+    def inference(self, speech_feat: torch.Tensor, cache_source: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor:
+        # mel->f0
+        f0 = self.f0_predictor(speech_feat)
+        # f0->source
+        s = self.f0_upsamp(f0[:, None]).transpose(1, 2)  # bs,n,t
+        s, _, _ = self.m_source(s)
+        s = s.transpose(1, 2)
+        # use cache_source to avoid glitch
+        if cache_source.shape[2] != 0:
+            s[:, :, :cache_source.shape[2]] = cache_source
+        generated_speech = self.decode(x=speech_feat, s=s)
+        return generated_speech, s

HF_Deploy/src/chatterbox/models/s3gen/matcha/decoder.py

@@ -0,0 +1,443 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from conformer import ConformerBlock
+from diffusers.models.activations import get_activation
+from einops import pack, rearrange, repeat
+
+from .transformer import BasicTransformerBlock
+
+
+class SinusoidalPosEmb(torch.nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+        assert self.dim % 2 == 0, "SinusoidalPosEmb requires dim to be even"
+
+    def forward(self, x, scale=1000):
+        if x.ndim < 1:
+            x = x.unsqueeze(0)
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
+        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class Block1D(torch.nn.Module):
+    def __init__(self, dim, dim_out, groups=8):
+        super().__init__()
+        self.block = torch.nn.Sequential(
+            torch.nn.Conv1d(dim, dim_out, 3, padding=1),
+            torch.nn.GroupNorm(groups, dim_out),
+            nn.Mish(),
+        )
+
+    def forward(self, x, mask):
+        output = self.block(x * mask)
+        return output * mask
+
+
+class ResnetBlock1D(torch.nn.Module):
+    def __init__(self, dim, dim_out, time_emb_dim, groups=8):
+        super().__init__()
+        self.mlp = torch.nn.Sequential(nn.Mish(), torch.nn.Linear(time_emb_dim, dim_out))
+
+        self.block1 = Block1D(dim, dim_out, groups=groups)
+        self.block2 = Block1D(dim_out, dim_out, groups=groups)
+
+        self.res_conv = torch.nn.Conv1d(dim, dim_out, 1)
+
+    def forward(self, x, mask, time_emb):
+        h = self.block1(x, mask)
+        h += self.mlp(time_emb).unsqueeze(-1)
+        h = self.block2(h, mask)
+        output = h + self.res_conv(x * mask)
+        return output
+
+
+class Downsample1D(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.conv = torch.nn.Conv1d(dim, dim, 3, 2, 1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        time_embed_dim: int,
+        act_fn: str = "silu",
+        out_dim: int = None,
+        post_act_fn: Optional[str] = None,
+        cond_proj_dim=None,
+    ):
+        super().__init__()
+
+        self.linear_1 = nn.Linear(in_channels, time_embed_dim)
+
+        if cond_proj_dim is not None:
+            self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)
+        else:
+            self.cond_proj = None
+
+        self.act = get_activation(act_fn)
+
+        if out_dim is not None:
+            time_embed_dim_out = out_dim
+        else:
+            time_embed_dim_out = time_embed_dim
+        self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out)
+
+        if post_act_fn is None:
+            self.post_act = None
+        else:
+            self.post_act = get_activation(post_act_fn)
+
+    def forward(self, sample, condition=None):
+        if condition is not None:
+            sample = sample + self.cond_proj(condition)
+        sample = self.linear_1(sample)
+
+        if self.act is not None:
+            sample = self.act(sample)
+
+        sample = self.linear_2(sample)
+
+        if self.post_act is not None:
+            sample = self.post_act(sample)
+        return sample
+
+
+class Upsample1D(nn.Module):
+    """A 1D upsampling layer with an optional convolution.
+
+    Parameters:
+        channels (`int`):
+            number of channels in the inputs and outputs.
+        use_conv (`bool`, default `False`):
+            option to use a convolution.
+        use_conv_transpose (`bool`, default `False`):
+            option to use a convolution transpose.
+        out_channels (`int`, optional):
+            number of output channels. Defaults to `channels`.
+    """
+
+    def __init__(self, channels, use_conv=False, use_conv_transpose=True, out_channels=None, name="conv"):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_conv_transpose = use_conv_transpose
+        self.name = name
+
+        self.conv = None
+        if use_conv_transpose:
+            self.conv = nn.ConvTranspose1d(channels, self.out_channels, 4, 2, 1)
+        elif use_conv:
+            self.conv = nn.Conv1d(self.channels, self.out_channels, 3, padding=1)
+
+    def forward(self, inputs):
+        assert inputs.shape[1] == self.channels
+        if self.use_conv_transpose:
+            return self.conv(inputs)
+
+        outputs = F.interpolate(inputs, scale_factor=2.0, mode="nearest")
+
+        if self.use_conv:
+            outputs = self.conv(outputs)
+
+        return outputs
+
+
+class ConformerWrapper(ConformerBlock):
+    def __init__(  # pylint: disable=useless-super-delegation
+        self,
+        *,
+        dim,
+        dim_head=64,
+        heads=8,
+        ff_mult=4,
+        conv_expansion_factor=2,
+        conv_kernel_size=31,
+        attn_dropout=0,
+        ff_dropout=0,
+        conv_dropout=0,
+        conv_causal=False,
+    ):
+        super().__init__(
+            dim=dim,
+            dim_head=dim_head,
+            heads=heads,
+            ff_mult=ff_mult,
+            conv_expansion_factor=conv_expansion_factor,
+            conv_kernel_size=conv_kernel_size,
+            attn_dropout=attn_dropout,
+            ff_dropout=ff_dropout,
+            conv_dropout=conv_dropout,
+            conv_causal=conv_causal,
+        )
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        timestep=None,
+    ):
+        return super().forward(x=hidden_states, mask=attention_mask.bool())
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        channels=(256, 256),
+        dropout=0.05,
+        attention_head_dim=64,
+        n_blocks=1,
+        num_mid_blocks=2,
+        num_heads=4,
+        act_fn="snake",
+        down_block_type="transformer",
+        mid_block_type="transformer",
+        up_block_type="transformer",
+    ):
+        super().__init__()
+        channels = tuple(channels)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.time_embeddings = SinusoidalPosEmb(in_channels)
+        time_embed_dim = channels[0] * 4
+        self.time_mlp = TimestepEmbedding(
+            in_channels=in_channels,
+            time_embed_dim=time_embed_dim,
+            act_fn="silu",
+        )
+
+        self.down_blocks = nn.ModuleList([])
+        self.mid_blocks = nn.ModuleList([])
+        self.up_blocks = nn.ModuleList([])
+
+        output_channel = in_channels
+        for i in range(len(channels)):  # pylint: disable=consider-using-enumerate
+            input_channel = output_channel
+            output_channel = channels[i]
+            is_last = i == len(channels) - 1
+            resnet = ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim)
+            transformer_blocks = nn.ModuleList(
+                [
+                    self.get_block(
+                        down_block_type,
+                        output_channel,
+                        attention_head_dim,
+                        num_heads,
+                        dropout,
+                        act_fn,
+                    )
+                    for _ in range(n_blocks)
+                ]
+            )
+            downsample = (
+                Downsample1D(output_channel) if not is_last else nn.Conv1d(output_channel, output_channel, 3, padding=1)
+            )
+
+            self.down_blocks.append(nn.ModuleList([resnet, transformer_blocks, downsample]))
+
+        for i in range(num_mid_blocks):
+            input_channel = channels[-1]
+            out_channels = channels[-1]
+
+            resnet = ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim)
+
+            transformer_blocks = nn.ModuleList(
+                [
+                    self.get_block(
+                        mid_block_type,
+                        output_channel,
+                        attention_head_dim,
+                        num_heads,
+                        dropout,
+                        act_fn,
+                    )
+                    for _ in range(n_blocks)
+                ]
+            )
+
+            self.mid_blocks.append(nn.ModuleList([resnet, transformer_blocks]))
+
+        channels = channels[::-1] + (channels[0],)
+        for i in range(len(channels) - 1):
+            input_channel = channels[i]
+            output_channel = channels[i + 1]
+            is_last = i == len(channels) - 2
+
+            resnet = ResnetBlock1D(
+                dim=2 * input_channel,
+                dim_out=output_channel,
+                time_emb_dim=time_embed_dim,
+            )
+            transformer_blocks = nn.ModuleList(
+                [
+                    self.get_block(
+                        up_block_type,
+                        output_channel,
+                        attention_head_dim,
+                        num_heads,
+                        dropout,
+                        act_fn,
+                    )
+                    for _ in range(n_blocks)
+                ]
+            )
+            upsample = (
+                Upsample1D(output_channel, use_conv_transpose=True)
+                if not is_last
+                else nn.Conv1d(output_channel, output_channel, 3, padding=1)
+            )
+
+            self.up_blocks.append(nn.ModuleList([resnet, transformer_blocks, upsample]))
+
+        self.final_block = Block1D(channels[-1], channels[-1])
+        self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1)
+
+        self.initialize_weights()
+        # nn.init.normal_(self.final_proj.weight)
+
+    @staticmethod
+    def get_block(block_type, dim, attention_head_dim, num_heads, dropout, act_fn):
+        if block_type == "conformer":
+            block = ConformerWrapper(
+                dim=dim,
+                dim_head=attention_head_dim,
+                heads=num_heads,
+                ff_mult=1,
+                conv_expansion_factor=2,
+                ff_dropout=dropout,
+                attn_dropout=dropout,
+                conv_dropout=dropout,
+                conv_kernel_size=31,
+            )
+        elif block_type == "transformer":
+            block = BasicTransformerBlock(
+                dim=dim,
+                num_attention_heads=num_heads,
+                attention_head_dim=attention_head_dim,
+                dropout=dropout,
+                activation_fn=act_fn,
+            )
+        else:
+            raise ValueError(f"Unknown block type {block_type}")
+
+        return block
+
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d):
+                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
+
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+            elif isinstance(m, nn.GroupNorm):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+            elif isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
+
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, x, mask, mu, t, spks=None, cond=None):
+        """Forward pass of the UNet1DConditional model.
+
+        Args:
+            x (torch.Tensor): shape (batch_size, in_channels, time)
+            mask (_type_): shape (batch_size, 1, time)
+            t (_type_): shape (batch_size)
+            spks (_type_, optional): shape: (batch_size, condition_channels). Defaults to None.
+            cond (_type_, optional): placeholder for future use. Defaults to None.
+
+        Raises:
+            ValueError: _description_
+            ValueError: _description_
+
+        Returns:
+            _type_: _description_
+        """
+
+        t = self.time_embeddings(t)
+        t = self.time_mlp(t)
+
+        x = pack([x, mu], "b * t")[0]
+
+        if spks is not None:
+            spks = repeat(spks, "b c -> b c t", t=x.shape[-1])
+            x = pack([x, spks], "b * t")[0]
+
+        hiddens = []
+        masks = [mask]
+        for resnet, transformer_blocks, downsample in self.down_blocks:
+            mask_down = masks[-1]
+            x = resnet(x, mask_down, t)
+            x = rearrange(x, "b c t -> b t c")
+            mask_down = rearrange(mask_down, "b 1 t -> b t")
+            for transformer_block in transformer_blocks:
+                x = transformer_block(
+                    hidden_states=x,
+                    attention_mask=mask_down,
+                    timestep=t,
+                )
+            x = rearrange(x, "b t c -> b c t")
+            mask_down = rearrange(mask_down, "b t -> b 1 t")
+            hiddens.append(x)  # Save hidden states for skip connections
+            x = downsample(x * mask_down)
+            masks.append(mask_down[:, :, ::2])
+
+        masks = masks[:-1]
+        mask_mid = masks[-1]
+
+        for resnet, transformer_blocks in self.mid_blocks:
+            x = resnet(x, mask_mid, t)
+            x = rearrange(x, "b c t -> b t c")
+            mask_mid = rearrange(mask_mid, "b 1 t -> b t")
+            for transformer_block in transformer_blocks:
+                x = transformer_block(
+                    hidden_states=x,
+                    attention_mask=mask_mid,
+                    timestep=t,
+                )
+            x = rearrange(x, "b t c -> b c t")
+            mask_mid = rearrange(mask_mid, "b t -> b 1 t")
+
+        for resnet, transformer_blocks, upsample in self.up_blocks:
+            mask_up = masks.pop()
+            x = resnet(pack([x, hiddens.pop()], "b * t")[0], mask_up, t)
+            x = rearrange(x, "b c t -> b t c")
+            mask_up = rearrange(mask_up, "b 1 t -> b t")
+            for transformer_block in transformer_blocks:
+                x = transformer_block(
+                    hidden_states=x,
+                    attention_mask=mask_up,
+                    timestep=t,
+                )
+            x = rearrange(x, "b t c -> b c t")
+            mask_up = rearrange(mask_up, "b t -> b 1 t")
+            x = upsample(x * mask_up)
+
+        x = self.final_block(x, mask_up)
+        output = self.final_proj(x * mask_up)
+
+        return output * mask

HF_Deploy/src/chatterbox/models/s3gen/matcha/flow_matching.py

@@ -0,0 +1,129 @@
+from abc import ABC
+
+import torch
+import torch.nn.functional as F
+
+from .decoder import Decoder
+
+
+class BASECFM(torch.nn.Module, ABC):
+    def __init__(
+        self,
+        n_feats,
+        cfm_params,
+        n_spks=1,
+        spk_emb_dim=128,
+    ):
+        super().__init__()
+        self.n_feats = n_feats
+        self.n_spks = n_spks
+        self.spk_emb_dim = spk_emb_dim
+        self.solver = cfm_params.solver
+        if hasattr(cfm_params, "sigma_min"):
+            self.sigma_min = cfm_params.sigma_min
+        else:
+            self.sigma_min = 1e-4
+
+        self.estimator = None
+
+    @torch.inference_mode()
+    def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None):
+        """Forward diffusion
+
+        Args:
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): output_mask
+                shape: (batch_size, 1, mel_timesteps)
+            n_timesteps (int): number of diffusion steps
+            temperature (float, optional): temperature for scaling noise. Defaults to 1.0.
+            spks (torch.Tensor, optional): speaker ids. Defaults to None.
+                shape: (batch_size, spk_emb_dim)
+            cond: Not used but kept for future purposes
+
+        Returns:
+            sample: generated mel-spectrogram
+                shape: (batch_size, n_feats, mel_timesteps)
+        """
+        z = torch.randn_like(mu) * temperature
+        t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device)
+        return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond)
+
+    def solve_euler(self, x, t_span, mu, mask, spks, cond):
+        """
+        Fixed euler solver for ODEs.
+        Args:
+            x (torch.Tensor): random noise
+            t_span (torch.Tensor): n_timesteps interpolated
+                shape: (n_timesteps + 1,)
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): output_mask
+                shape: (batch_size, 1, mel_timesteps)
+            spks (torch.Tensor, optional): speaker ids. Defaults to None.
+                shape: (batch_size, spk_emb_dim)
+            cond: Not used but kept for future purposes
+        """
+        t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]
+
+        # I am storing this because I can later plot it by putting a debugger here and saving it to a file
+        # Or in future might add like a return_all_steps flag
+        sol = []
+
+        for step in range(1, len(t_span)):
+            dphi_dt = self.estimator(x, mask, mu, t, spks, cond)
+
+            x = x + dt * dphi_dt
+            t = t + dt
+            sol.append(x)
+            if step < len(t_span) - 1:
+                dt = t_span[step + 1] - t
+
+        return sol[-1]
+
+    def compute_loss(self, x1, mask, mu, spks=None, cond=None):
+        """Computes diffusion loss
+
+        Args:
+            x1 (torch.Tensor): Target
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): target mask
+                shape: (batch_size, 1, mel_timesteps)
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            spks (torch.Tensor, optional): speaker embedding. Defaults to None.
+                shape: (batch_size, spk_emb_dim)
+
+        Returns:
+            loss: conditional flow matching loss
+            y: conditional flow
+                shape: (batch_size, n_feats, mel_timesteps)
+        """
+        b, _, t = mu.shape
+
+        # random timestep
+        t = torch.rand([b, 1, 1], device=mu.device, dtype=mu.dtype)
+        # sample noise p(x_0)
+        z = torch.randn_like(x1)
+
+        y = (1 - (1 - self.sigma_min) * t) * z + t * x1
+        u = x1 - (1 - self.sigma_min) * z
+
+        loss = F.mse_loss(self.estimator(y, mask, mu, t.squeeze(), spks), u, reduction="sum") / (
+            torch.sum(mask) * u.shape[1]
+        )
+        return loss, y
+
+
+class CFM(BASECFM):
+    def __init__(self, in_channels, out_channel, cfm_params, decoder_params, n_spks=1, spk_emb_dim=64):
+        super().__init__(
+            n_feats=in_channels,
+            cfm_params=cfm_params,
+            n_spks=n_spks,
+            spk_emb_dim=spk_emb_dim,
+        )
+
+        in_channels = in_channels + (spk_emb_dim if n_spks > 1 else 0)
+        # Just change the architecture of the estimator here
+        self.estimator = Decoder(in_channels=in_channels, out_channels=out_channel, **decoder_params)

HF_Deploy/src/chatterbox/models/s3gen/matcha/text_encoder.py

@@ -0,0 +1,413 @@
+""" from https://github.com/jaywalnut310/glow-tts """
+
+import math
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+
+def sequence_mask(length, max_length=None):
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels, eps=1e-4):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = torch.nn.Parameter(torch.ones(channels))
+        self.beta = torch.nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        n_dims = len(x.shape)
+        mean = torch.mean(x, 1, keepdim=True)
+        variance = torch.mean((x - mean) ** 2, 1, keepdim=True)
+
+        x = (x - mean) * torch.rsqrt(variance + self.eps)
+
+        shape = [1, -1] + [1] * (n_dims - 2)
+        x = x * self.gamma.view(*shape) + self.beta.view(*shape)
+        return x
+
+
+class ConvReluNorm(nn.Module):
+    def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+
+        self.conv_layers = torch.nn.ModuleList()
+        self.norm_layers = torch.nn.ModuleList()
+        self.conv_layers.append(torch.nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size // 2))
+        self.norm_layers.append(LayerNorm(hidden_channels))
+        self.relu_drop = torch.nn.Sequential(torch.nn.ReLU(), torch.nn.Dropout(p_dropout))
+        for _ in range(n_layers - 1):
+            self.conv_layers.append(
+                torch.nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size // 2)
+            )
+            self.norm_layers.append(LayerNorm(hidden_channels))
+        self.proj = torch.nn.Conv1d(hidden_channels, out_channels, 1)
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask):
+        x_org = x
+        for i in range(self.n_layers):
+            x = self.conv_layers[i](x * x_mask)
+            x = self.norm_layers[i](x)
+            x = self.relu_drop(x)
+        x = x_org + self.proj(x)
+        return x * x_mask
+
+
+class DurationPredictor(nn.Module):
+    def __init__(self, in_channels, filter_channels, kernel_size, p_dropout):
+        super().__init__()
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.p_dropout = p_dropout
+
+        self.drop = torch.nn.Dropout(p_dropout)
+        self.conv_1 = torch.nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size // 2)
+        self.norm_1 = LayerNorm(filter_channels)
+        self.conv_2 = torch.nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size // 2)
+        self.norm_2 = LayerNorm(filter_channels)
+        self.proj = torch.nn.Conv1d(filter_channels, 1, 1)
+
+    def forward(self, x, x_mask):
+        x = self.conv_1(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_1(x)
+        x = self.drop(x)
+        x = self.conv_2(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_2(x)
+        x = self.drop(x)
+        x = self.proj(x * x_mask)
+        return x * x_mask
+
+
+class RotaryPositionalEmbeddings(nn.Module):
+    """
+    ## RoPE module
+
+    Rotary encoding transforms pairs of features by rotating in the 2D plane.
+    That is, it organizes the $d$ features as $\frac{d}{2}$ pairs.
+    Each pair can be considered a coordinate in a 2D plane, and the encoding will rotate it
+    by an angle depending on the position of the token.
+    """
+
+    def __init__(self, d: int, base: int = 10_000):
+        r"""
+        * `d` is the number of features $d$
+        * `base` is the constant used for calculating $\Theta$
+        """
+        super().__init__()
+
+        self.base = base
+        self.d = int(d)
+        self.cos_cached = None
+        self.sin_cached = None
+
+    def _build_cache(self, x: torch.Tensor):
+        r"""
+        Cache $\cos$ and $\sin$ values
+        """
+        # Return if cache is already built
+        if self.cos_cached is not None and x.shape[0] <= self.cos_cached.shape[0]:
+            return
+
+        # Get sequence length
+        seq_len = x.shape[0]
+
+        # $\Theta = {\theta_i = 10000^{-\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
+        theta = 1.0 / (self.base ** (torch.arange(0, self.d, 2).float() / self.d)).to(x.device)
+
+        # Create position indexes `[0, 1, ..., seq_len - 1]`
+        seq_idx = torch.arange(seq_len, device=x.device).float().to(x.device)
+
+        # Calculate the product of position index and $\theta_i$
+        idx_theta = torch.einsum("n,d->nd", seq_idx, theta)
+
+        # Concatenate so that for row $m$ we have
+        # $[m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}, m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}]$
+        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=1)
+
+        # Cache them
+        self.cos_cached = idx_theta2.cos()[:, None, None, :]
+        self.sin_cached = idx_theta2.sin()[:, None, None, :]
+
+    def _neg_half(self, x: torch.Tensor):
+        # $\frac{d}{2}$
+        d_2 = self.d // 2
+
+        # Calculate $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., x^{(\frac{d}{2})}]$
+        return torch.cat([-x[:, :, :, d_2:], x[:, :, :, :d_2]], dim=-1)
+
+    def forward(self, x: torch.Tensor):
+        """
+        * `x` is the Tensor at the head of a key or a query with shape `[seq_len, batch_size, n_heads, d]`
+        """
+        # Cache $\cos$ and $\sin$ values
+        x = rearrange(x, "b h t d -> t b h d")
+
+        self._build_cache(x)
+
+        # Split the features, we can choose to apply rotary embeddings only to a partial set of features.
+        x_rope, x_pass = x[..., : self.d], x[..., self.d :]
+
+        # Calculate
+        # $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., x^{(\frac{d}{2})}]$
+        neg_half_x = self._neg_half(x_rope)
+
+        x_rope = (x_rope * self.cos_cached[: x.shape[0]]) + (neg_half_x * self.sin_cached[: x.shape[0]])
+
+        return rearrange(torch.cat((x_rope, x_pass), dim=-1), "t b h d -> b h t d")
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels,
+        out_channels,
+        n_heads,
+        heads_share=True,
+        p_dropout=0.0,
+        proximal_bias=False,
+        proximal_init=False,
+    ):
+        super().__init__()
+        assert channels % n_heads == 0
+
+        self.channels = channels
+        self.out_channels = out_channels
+        self.n_heads = n_heads
+        self.heads_share = heads_share
+        self.proximal_bias = proximal_bias
+        self.p_dropout = p_dropout
+        self.attn = None
+
+        self.k_channels = channels // n_heads
+        self.conv_q = torch.nn.Conv1d(channels, channels, 1)
+        self.conv_k = torch.nn.Conv1d(channels, channels, 1)
+        self.conv_v = torch.nn.Conv1d(channels, channels, 1)
+
+        # from https://nn.labml.ai/transformers/rope/index.html
+        self.query_rotary_pe = RotaryPositionalEmbeddings(self.k_channels * 0.5)
+        self.key_rotary_pe = RotaryPositionalEmbeddings(self.k_channels * 0.5)
+
+        self.conv_o = torch.nn.Conv1d(channels, out_channels, 1)
+        self.drop = torch.nn.Dropout(p_dropout)
+
+        torch.nn.init.xavier_uniform_(self.conv_q.weight)
+        torch.nn.init.xavier_uniform_(self.conv_k.weight)
+        if proximal_init:
+            self.conv_k.weight.data.copy_(self.conv_q.weight.data)
+            self.conv_k.bias.data.copy_(self.conv_q.bias.data)
+        torch.nn.init.xavier_uniform_(self.conv_v.weight)
+
+    def forward(self, x, c, attn_mask=None):
+        q = self.conv_q(x)
+        k = self.conv_k(c)
+        v = self.conv_v(c)
+
+        x, self.attn = self.attention(q, k, v, mask=attn_mask)
+
+        x = self.conv_o(x)
+        return x
+
+    def attention(self, query, key, value, mask=None):
+        b, d, t_s, t_t = (*key.size(), query.size(2))
+        query = rearrange(query, "b (h c) t-> b h t c", h=self.n_heads)
+        key = rearrange(key, "b (h c) t-> b h t c", h=self.n_heads)
+        value = rearrange(value, "b (h c) t-> b h t c", h=self.n_heads)
+
+        query = self.query_rotary_pe(query)
+        key = self.key_rotary_pe(key)
+
+        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.k_channels)
+
+        if self.proximal_bias:
+            assert t_s == t_t, "Proximal bias is only available for self-attention."
+            scores = scores + self._attention_bias_proximal(t_s).to(device=scores.device, dtype=scores.dtype)
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e4)
+        p_attn = torch.nn.functional.softmax(scores, dim=-1)
+        p_attn = self.drop(p_attn)
+        output = torch.matmul(p_attn, value)
+        output = output.transpose(2, 3).contiguous().view(b, d, t_t)
+        return output, p_attn
+
+    @staticmethod
+    def _attention_bias_proximal(length):
+        r = torch.arange(length, dtype=torch.float32)
+        diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+        return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
+
+
+class FFN(nn.Module):
+    def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0.0):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+
+        self.conv_1 = torch.nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size // 2)
+        self.conv_2 = torch.nn.Conv1d(filter_channels, out_channels, kernel_size, padding=kernel_size // 2)
+        self.drop = torch.nn.Dropout(p_dropout)
+
+    def forward(self, x, x_mask):
+        x = self.conv_1(x * x_mask)
+        x = torch.relu(x)
+        x = self.drop(x)
+        x = self.conv_2(x * x_mask)
+        return x * x_mask
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        **kwargs,
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+
+        self.drop = torch.nn.Dropout(p_dropout)
+        self.attn_layers = torch.nn.ModuleList()
+        self.norm_layers_1 = torch.nn.ModuleList()
+        self.ffn_layers = torch.nn.ModuleList()
+        self.norm_layers_2 = torch.nn.ModuleList()
+        for _ in range(self.n_layers):
+            self.attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout))
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask):
+        attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        for i in range(self.n_layers):
+            x = x * x_mask
+            y = self.attn_layers[i](x, x, attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class TextEncoder(nn.Module):
+    def __init__(
+        self,
+        encoder_type,
+        encoder_params,
+        duration_predictor_params,
+        n_vocab,
+        n_spks=1,
+        spk_emb_dim=128,
+    ):
+        super().__init__()
+        self.encoder_type = encoder_type
+        self.n_vocab = n_vocab
+        self.n_feats = encoder_params.n_feats
+        self.n_channels = encoder_params.n_channels
+        self.spk_emb_dim = spk_emb_dim
+        self.n_spks = n_spks
+
+        self.emb = torch.nn.Embedding(n_vocab, self.n_channels)
+        torch.nn.init.normal_(self.emb.weight, 0.0, self.n_channels**-0.5)
+
+        if encoder_params.prenet:
+            self.prenet = ConvReluNorm(
+                self.n_channels,
+                self.n_channels,
+                self.n_channels,
+                kernel_size=5,
+                n_layers=3,
+                p_dropout=0.5,
+            )
+        else:
+            self.prenet = lambda x, x_mask: x
+
+        self.encoder = Encoder(
+            encoder_params.n_channels + (spk_emb_dim if n_spks > 1 else 0),
+            encoder_params.filter_channels,
+            encoder_params.n_heads,
+            encoder_params.n_layers,
+            encoder_params.kernel_size,
+            encoder_params.p_dropout,
+        )
+
+        self.proj_m = torch.nn.Conv1d(self.n_channels + (spk_emb_dim if n_spks > 1 else 0), self.n_feats, 1)
+        self.proj_w = DurationPredictor(
+            self.n_channels + (spk_emb_dim if n_spks > 1 else 0),
+            duration_predictor_params.filter_channels_dp,
+            duration_predictor_params.kernel_size,
+            duration_predictor_params.p_dropout,
+        )
+
+    def forward(self, x, x_lengths, spks=None):
+        """Run forward pass to the transformer based encoder and duration predictor
+
+        Args:
+            x (torch.Tensor): text input
+                shape: (batch_size, max_text_length)
+            x_lengths (torch.Tensor): text input lengths
+                shape: (batch_size,)
+            spks (torch.Tensor, optional): speaker ids. Defaults to None.
+                shape: (batch_size,)
+
+        Returns:
+            mu (torch.Tensor): average output of the encoder
+                shape: (batch_size, n_feats, max_text_length)
+            logw (torch.Tensor): log duration predicted by the duration predictor
+                shape: (batch_size, 1, max_text_length)
+            x_mask (torch.Tensor): mask for the text input
+                shape: (batch_size, 1, max_text_length)
+        """
+        x = self.emb(x) * math.sqrt(self.n_channels)
+        x = torch.transpose(x, 1, -1)
+        x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
+
+        x = self.prenet(x, x_mask)
+        if self.n_spks > 1:
+            x = torch.cat([x, spks.unsqueeze(-1).repeat(1, 1, x.shape[-1])], dim=1)
+        x = self.encoder(x, x_mask)
+        mu = self.proj_m(x) * x_mask
+
+        x_dp = torch.detach(x)
+        logw = self.proj_w(x_dp, x_mask)
+
+        return mu, logw, x_mask

HF_Deploy/src/chatterbox/models/s3gen/matcha/transformer.py

@@ -0,0 +1,316 @@
+from typing import Any, Dict, Optional
+
+import torch
+import torch.nn as nn
+from diffusers.models.attention import (
+    GEGLU,
+    GELU,
+    AdaLayerNorm,
+    AdaLayerNormZero,
+    ApproximateGELU,
+)
+from diffusers.models.attention_processor import Attention
+from diffusers.models.lora import LoRACompatibleLinear
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+
+
+class SnakeBeta(nn.Module):
+    """
+    A modified Snake function which uses separate parameters for the magnitude of the periodic components
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter that controls frequency
+        - beta - trainable parameter that controls magnitude
+    References:
+        - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snakebeta(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    """
+
+    def __init__(self, in_features, out_features, alpha=1.0, alpha_trainable=True, alpha_logscale=True):
+        """
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha - trainable parameter that controls frequency
+            - beta - trainable parameter that controls magnitude
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            beta is initialized to 1 by default, higher values = higher-magnitude.
+            alpha will be trained along with the rest of your model.
+        """
+        super().__init__()
+        self.in_features = out_features if isinstance(out_features, list) else [out_features]
+        self.proj = LoRACompatibleLinear(in_features, out_features)
+
+        # initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # log scale alphas initialized to zeros
+            self.alpha = nn.Parameter(torch.zeros(self.in_features) * alpha)
+            self.beta = nn.Parameter(torch.zeros(self.in_features) * alpha)
+        else:  # linear scale alphas initialized to ones
+            self.alpha = nn.Parameter(torch.ones(self.in_features) * alpha)
+            self.beta = nn.Parameter(torch.ones(self.in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+        self.beta.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        """
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        SnakeBeta ∶= x + 1/b * sin^2 (xa)
+        """
+        x = self.proj(x)
+        if self.alpha_logscale:
+            alpha = torch.exp(self.alpha)
+            beta = torch.exp(self.beta)
+        else:
+            alpha = self.alpha
+            beta = self.beta
+
+        x = x + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(torch.sin(x * alpha), 2)
+
+        return x
+
+
+class FeedForward(nn.Module):
+    r"""
+    A feed-forward layer.
+
+    Parameters:
+        dim (`int`): The number of channels in the input.
+        dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
+        mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        mult: int = 4,
+        dropout: float = 0.0,
+        activation_fn: str = "geglu",
+        final_dropout: bool = False,
+    ):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        if activation_fn == "gelu":
+            act_fn = GELU(dim, inner_dim)
+        if activation_fn == "gelu-approximate":
+            act_fn = GELU(dim, inner_dim, approximate="tanh")
+        elif activation_fn == "geglu":
+            act_fn = GEGLU(dim, inner_dim)
+        elif activation_fn == "geglu-approximate":
+            act_fn = ApproximateGELU(dim, inner_dim)
+        elif activation_fn == "snakebeta":
+            act_fn = SnakeBeta(dim, inner_dim)
+
+        self.net = nn.ModuleList([])
+        # project in
+        self.net.append(act_fn)
+        # project dropout
+        self.net.append(nn.Dropout(dropout))
+        # project out
+        self.net.append(LoRACompatibleLinear(inner_dim, dim_out))
+        # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
+        if final_dropout:
+            self.net.append(nn.Dropout(dropout))
+
+    def forward(self, hidden_states):
+        for module in self.net:
+            hidden_states = module(hidden_states)
+        return hidden_states
+
+
+@maybe_allow_in_graph
+class BasicTransformerBlock(nn.Module):
+    r"""
+    A basic Transformer block.
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
+        only_cross_attention (`bool`, *optional*):
+            Whether to use only cross-attention layers. In this case two cross attention layers are used.
+        double_self_attention (`bool`, *optional*):
+            Whether to use two self-attention layers. In this case no cross attention layers are used.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        num_embeds_ada_norm (:
+            obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
+        attention_bias (:
+            obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        dropout=0.0,
+        cross_attention_dim: Optional[int] = None,
+        activation_fn: str = "geglu",
+        num_embeds_ada_norm: Optional[int] = None,
+        attention_bias: bool = False,
+        only_cross_attention: bool = False,
+        double_self_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_elementwise_affine: bool = True,
+        norm_type: str = "layer_norm",
+        final_dropout: bool = False,
+    ):
+        super().__init__()
+        self.only_cross_attention = only_cross_attention
+
+        self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
+        self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm"
+
+        if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
+            raise ValueError(
+                f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"
+                f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."
+            )
+
+        # Define 3 blocks. Each block has its own normalization layer.
+        # 1. Self-Attn
+        if self.use_ada_layer_norm:
+            self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)
+        elif self.use_ada_layer_norm_zero:
+            self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
+        else:
+            self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+        self.attn1 = Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            cross_attention_dim=cross_attention_dim if only_cross_attention else None,
+            upcast_attention=upcast_attention,
+        )
+
+        # 2. Cross-Attn
+        if cross_attention_dim is not None or double_self_attention:
+            # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
+            # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
+            # the second cross attention block.
+            self.norm2 = (
+                AdaLayerNorm(dim, num_embeds_ada_norm)
+                if self.use_ada_layer_norm
+                else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+            )
+            self.attn2 = Attention(
+                query_dim=dim,
+                cross_attention_dim=cross_attention_dim if not double_self_attention else None,
+                heads=num_attention_heads,
+                dim_head=attention_head_dim,
+                dropout=dropout,
+                bias=attention_bias,
+                upcast_attention=upcast_attention,
+                # scale_qk=False, # uncomment this to not to use flash attention
+            )  # is self-attn if encoder_hidden_states is none
+        else:
+            self.norm2 = None
+            self.attn2 = None
+
+        # 3. Feed-forward
+        self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+        self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout)
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int):
+        # Sets chunk feed-forward
+        self._chunk_size = chunk_size
+        self._chunk_dim = dim
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        timestep: Optional[torch.LongTensor] = None,
+        cross_attention_kwargs: Dict[str, Any] = None,
+        class_labels: Optional[torch.LongTensor] = None,
+    ):
+        # Notice that normalization is always applied before the real computation in the following blocks.
+        # 1. Self-Attention
+        if self.use_ada_layer_norm:
+            norm_hidden_states = self.norm1(hidden_states, timestep)
+        elif self.use_ada_layer_norm_zero:
+            norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+                hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
+            )
+        else:
+            norm_hidden_states = self.norm1(hidden_states)
+
+        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
+
+        attn_output = self.attn1(
+            norm_hidden_states,
+            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
+            attention_mask=encoder_attention_mask if self.only_cross_attention else attention_mask,
+            **cross_attention_kwargs,
+        )
+        if self.use_ada_layer_norm_zero:
+            attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = attn_output + hidden_states
+
+        # 2. Cross-Attention
+        if self.attn2 is not None:
+            norm_hidden_states = (
+                self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
+            )
+
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+                **cross_attention_kwargs,
+            )
+            hidden_states = attn_output + hidden_states
+
+        # 3. Feed-forward
+        norm_hidden_states = self.norm3(hidden_states)
+
+        if self.use_ada_layer_norm_zero:
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
+                raise ValueError(
+                    f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
+                )
+
+            num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
+            ff_output = torch.cat(
+                [self.ff(hid_slice) for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim)],
+                dim=self._chunk_dim,
+            )
+        else:
+            ff_output = self.ff(norm_hidden_states)
+
+        if self.use_ada_layer_norm_zero:
+            ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+        hidden_states = ff_output + hidden_states
+
+        return hidden_states

HF_Deploy/src/chatterbox/models/s3gen/s3gen.py

@@ -0,0 +1,305 @@
+# Modified from CosyVoice https://github.com/FunAudioLLM/CosyVoice
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+import numpy as np
+import torch
+import torchaudio as ta
+from functools import lru_cache
+from typing import Optional
+from omegaconf import DictConfig
+
+from ..s3tokenizer import S3_SR, SPEECH_VOCAB_SIZE, S3Tokenizer
+from .const import S3GEN_SR
+from .flow import CausalMaskedDiffWithXvec
+from .xvector import CAMPPlus
+from .utils.mel import mel_spectrogram
+from .f0_predictor import ConvRNNF0Predictor
+from .hifigan import HiFTGenerator
+from .transformer.upsample_encoder import UpsampleConformerEncoder
+from .flow_matching import CausalConditionalCFM
+from .decoder import ConditionalDecoder
+
+
+def drop_invalid_tokens(x):
+    assert len(x.shape) <= 2 and x.shape[0] == 1, "only batch size of one allowed for now"
+    return x[x < SPEECH_VOCAB_SIZE]
+
+
+# TODO: global resampler cache
+@lru_cache(100)
+def get_resampler(src_sr, dst_sr, device):
+    return ta.transforms.Resample(src_sr, dst_sr).to(device)
+
+
+class S3Token2Mel(torch.nn.Module):
+    """
+    CosyVoice2's CFM decoder maps S3 speech tokens to mel-spectrograms.
+
+    TODO: make these modules configurable?
+    """
+    def __init__(self):
+        super().__init__()
+        self.tokenizer = S3Tokenizer("speech_tokenizer_v2_25hz")
+        self.mel_extractor = mel_spectrogram # TODO: make it a torch module?
+        self.speaker_encoder = CAMPPlus()  # use default args
+
+        encoder = UpsampleConformerEncoder(
+            output_size=512,
+            attention_heads=8,
+            linear_units=2048,
+            num_blocks=6,
+            dropout_rate=0.1,
+            positional_dropout_rate=0.1,
+            attention_dropout_rate=0.1,
+            normalize_before=True,
+            input_layer='linear',
+            pos_enc_layer_type='rel_pos_espnet',
+            selfattention_layer_type='rel_selfattn',
+            input_size=512,
+            use_cnn_module=False,
+            macaron_style=False,
+        )
+
+        estimator = ConditionalDecoder(
+            in_channels=320,
+            out_channels=80,
+            causal=True,
+            channels=[256],
+            dropout=0.0,
+            attention_head_dim=64,
+            n_blocks=4,
+            num_mid_blocks=12,
+            num_heads=8,
+            act_fn='gelu',
+        )
+        cfm_params = DictConfig({
+            "sigma_min": 1e-06,
+            "solver": 'euler',
+            "t_scheduler": 'cosine',
+            "training_cfg_rate": 0.2,
+            "inference_cfg_rate": 0.7,
+            "reg_loss_type": 'l1',
+        })
+        decoder = CausalConditionalCFM(
+            spk_emb_dim=80,
+            cfm_params=cfm_params,
+            estimator=estimator,
+        )
+
+        self.flow = CausalMaskedDiffWithXvec(
+            encoder=encoder,
+            decoder=decoder
+        )
+
+        self.resamplers = {}
+
+    @property
+    def device(self):
+        params = self.tokenizer.parameters()
+        return next(params).device
+
+    def embed_ref(
+        self,
+        ref_wav: torch.Tensor,
+        ref_sr: int,
+        device="auto",
+        ref_fade_out=True,
+    ):
+        device = self.device if device == "auto" else device
+        if isinstance(ref_wav, np.ndarray):
+            ref_wav = torch.from_numpy(ref_wav).float()
+
+        if ref_wav.device != device:
+            ref_wav = ref_wav.to(device)
+
+        if len(ref_wav.shape) == 1:
+            ref_wav = ref_wav.unsqueeze(0)  # (B, L)
+
+        if ref_wav.size(1) > 10 * ref_sr:
+            print("WARNING: cosydec received ref longer than 10s")
+
+        ref_wav_24 = ref_wav
+        if ref_sr != S3GEN_SR:
+            ref_wav_24 = get_resampler(ref_sr, S3GEN_SR, device)(ref_wav)
+
+        ref_mels_24 = self.mel_extractor(ref_wav_24).transpose(1, 2).to(device)
+        ref_mels_24_len = None
+
+        # Resample to 16kHz
+        ref_wav_16 = get_resampler(ref_sr, S3_SR, device)(ref_wav).to(device)
+
+        # Speaker embedding
+        ref_x_vector = self.speaker_encoder.inference(ref_wav_16)
+
+        # Tokenize 16khz reference
+        ref_speech_tokens, ref_speech_token_lens = self.tokenizer(ref_wav_16)
+
+        # Make sure mel_len = 2 * stoken_len (happens when the input is not padded to multiple of 40ms)
+        if ref_mels_24.shape[1] != 2 * ref_speech_tokens.shape[1]:
+            logging.warning(
+                "Reference mel length is not equal to 2 * reference token length.\n"
+            )
+            ref_speech_tokens = ref_speech_tokens[:, :ref_mels_24.shape[1] // 2]
+            ref_speech_token_lens[0] = ref_speech_tokens.shape[1]
+
+        return dict(
+            prompt_token=ref_speech_tokens.to(device),
+            prompt_token_len=ref_speech_token_lens,
+            prompt_feat=ref_mels_24,
+            prompt_feat_len=ref_mels_24_len,
+            embedding=ref_x_vector,
+        )
+
+    def forward(
+        self,
+        speech_tokens: torch.LongTensor,
+        # locally-computed ref embedding (mutex with ref_dict)
+        ref_wav: Optional[torch.Tensor],
+        ref_sr: Optional[int],
+        # pre-computed ref embedding (prod API)
+        ref_dict: Optional[dict] = None,
+        finalize: bool = False,
+    ):
+        """
+        Generate waveforms from S3 speech tokens and a reference waveform, which the speaker timbre is inferred from.
+
+        NOTE:
+        - The speaker encoder accepts 16 kHz waveform.
+        - S3TokenizerV2 accepts 16 kHz waveform.
+        - The mel-spectrogram for the reference assumes 24 kHz input signal.
+        - This function is designed for batch_size=1 only.
+
+        Args
+        ----
+        - `speech_tokens`: S3 speech tokens [B=1, T]
+        - `ref_wav`: reference waveform (`torch.Tensor` with shape=[B=1, T])
+        - `ref_sr`: reference sample rate
+        - `finalize`: whether streaming is finished or not. Note that if False, the last 3 tokens will be ignored.
+        """
+        assert (ref_wav is None) ^ (ref_dict is None), f"Must provide exactly one of ref_wav or ref_dict (got {ref_wav} and {ref_dict})"
+
+        if ref_dict is None:
+            ref_dict = self.embed_ref(ref_wav, ref_sr)
+        else:
+            # type/device casting (all values will be numpy if it's from a prod API call)
+            for rk in list(ref_dict):
+                if isinstance(ref_dict[rk], np.ndarray):
+                    ref_dict[rk] = torch.from_numpy(ref_dict[rk])
+                if torch.is_tensor(ref_dict[rk]):
+                    ref_dict[rk] = ref_dict[rk].to(self.device)
+
+        if len(speech_tokens.shape) == 1:
+            speech_tokens = speech_tokens.unsqueeze(0)
+
+        # assert speech_tokens.shape[0] == 1, "only batch size of one allowed for now"
+        speech_token_lens = torch.LongTensor([speech_tokens.size(1)]).to(self.device)
+
+        output_mels, _ = self.flow.inference(
+            token=speech_tokens,
+            token_len=speech_token_lens,
+            finalize=finalize,
+            **ref_dict,
+        )
+        return output_mels
+
+
+class S3Token2Wav(S3Token2Mel):
+    """
+    The decoder of CosyVoice2 is a concat of token-to-mel (CFM) and a mel-to-waveform (HiFiGAN) modules.
+
+    TODO: make these modules configurable?
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        f0_predictor = ConvRNNF0Predictor()
+        self.mel2wav = HiFTGenerator(
+            sampling_rate=S3GEN_SR,
+            upsample_rates=[8, 5, 3],
+            upsample_kernel_sizes=[16, 11, 7],
+            source_resblock_kernel_sizes=[7, 7, 11],
+            source_resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+            f0_predictor=f0_predictor,
+        )
+
+        # silence out a few ms and fade audio in to reduce artifacts
+        n_trim = S3GEN_SR // 50  # 20ms = half of a frame
+        trim_fade = torch.zeros(2 * n_trim)
+        trim_fade[n_trim:] = (torch.cos(torch.linspace(torch.pi, 0, n_trim)) + 1) / 2
+        self.register_buffer("trim_fade", trim_fade, persistent=False) # (buffers get automatic device casting)
+
+    def forward(
+        self,
+        speech_tokens,
+        # locally-computed ref embedding (mutex with ref_dict)
+        ref_wav: Optional[torch.Tensor],
+        ref_sr: Optional[int],
+        # pre-computed ref embedding (prod API)
+        ref_dict: Optional[dict] = None,
+        finalize: bool = False
+    ):
+        output_mels = super().forward(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize)
+
+        # TODO jrm: ignoring the speed control (mel interpolation) and the HiFTGAN caching mechanisms for now.
+        hift_cache_source = torch.zeros(1, 1, 0).to(self.device)
+
+        output_wavs, *_ = self.mel2wav.inference(speech_feat=output_mels, cache_source=hift_cache_source)
+
+        if not self.training:
+            # NOTE: ad-hoc method to reduce "spillover" from the reference clip.
+            output_wavs[:, :len(self.trim_fade)] *= self.trim_fade
+
+        return output_wavs
+
+    @torch.inference_mode()
+    def flow_inference(
+        self,
+        speech_tokens,
+        # locally-computed ref embedding (mutex with ref_dict)
+        ref_wav: Optional[torch.Tensor] = None,
+        ref_sr: Optional[int] = None,
+        # pre-computed ref embedding (prod API)
+        ref_dict: Optional[dict] = None,
+        finalize: bool = False,
+    ):
+        return super().forward(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize)
+
+    @torch.inference_mode()
+    def hift_inference(self, speech_feat, cache_source: torch.Tensor = None):
+        if cache_source is None:
+            cache_source = torch.zeros(1, 1, 0).to(self.device)
+        return self.mel2wav.inference(speech_feat=speech_feat, cache_source=cache_source)
+
+    @torch.inference_mode()
+    def inference(
+        self,
+        speech_tokens,
+        # locally-computed ref embedding (mutex with ref_dict)
+        ref_wav: Optional[torch.Tensor] = None,
+        ref_sr: Optional[int] = None,
+        # pre-computed ref embedding (prod API)
+        ref_dict: Optional[dict] = None,
+        cache_source: torch.Tensor = None, # NOTE: this arg is for streaming, it can probably be removed here
+        finalize: bool = True,
+    ):
+        output_mels = self.flow_inference(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize)
+        output_wavs, output_sources = self.hift_inference(output_mels, cache_source)
+
+        # NOTE: ad-hoc method to reduce "spillover" from the reference clip.
+        output_wavs[:, :len(self.trim_fade)] *= self.trim_fade
+
+        return output_wavs, output_sources

HF_Deploy/src/chatterbox/models/s3gen/transformer/activation.py

@@ -0,0 +1,84 @@
+# Copyright (c) 2020 Johns Hopkins University (Shinji Watanabe)
+#               2020 Northwestern Polytechnical University (Pengcheng Guo)
+#               2020 Mobvoi Inc (Binbin Zhang)
+#               2024 Alibaba Inc (Xiang Lyu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Swish() activation function for Conformer."""
+
+import torch
+from torch import nn, sin, pow
+from torch.nn import Parameter
+
+
+class Swish(torch.nn.Module):
+    """Construct an Swish object."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Return Swish activation function."""
+        return x * torch.sigmoid(x)
+
+
+# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license.
+#   LICENSE is in incl_licenses directory.
+class Snake(nn.Module):
+    '''
+    Implementation of a sine-based periodic activation function
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter
+    References:
+        - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snake(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    '''
+    def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):
+        '''
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha: trainable parameter
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            alpha will be trained along with the rest of your model.
+        '''
+        super(Snake, self).__init__()
+        self.in_features = in_features
+
+        # initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # log scale alphas initialized to zeros
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+        else:  # linear scale alphas initialized to ones
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        '''
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        Snake ∶= x + 1/a * sin^2 (xa)
+        '''
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+        x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x

HF_Deploy/src/chatterbox/models/s3gen/transformer/attention.py

@@ -0,0 +1,330 @@
+# Copyright (c) 2019 Shigeki Karita
+#               2020 Mobvoi Inc (Binbin Zhang)
+#               2022 Xingchen Song (sxc19@mails.tsinghua.edu.cn)
+#               2024 Alibaba Inc (Xiang Lyu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Multi-Head Attention layer definition."""
+
+import math
+from typing import Tuple
+
+import torch
+from torch import nn
+
+
+class MultiHeadedAttention(nn.Module):
+    """Multi-Head Attention layer.
+
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self,
+                 n_head: int,
+                 n_feat: int,
+                 dropout_rate: float,
+                 key_bias: bool = True):
+        """Construct an MultiHeadedAttention object."""
+        super().__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat, bias=key_bias)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward_qkv(
+        self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Transform query, key and value.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+
+        Returns:
+            torch.Tensor: Transformed query tensor, size
+                (#batch, n_head, time1, d_k).
+            torch.Tensor: Transformed key tensor, size
+                (#batch, n_head, time2, d_k).
+            torch.Tensor: Transformed value tensor, size
+                (#batch, n_head, time2, d_k).
+
+        """
+        n_batch = query.size(0)
+        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
+        q = q.transpose(1, 2)  # (batch, head, time1, d_k)
+        k = k.transpose(1, 2)  # (batch, head, time2, d_k)
+        v = v.transpose(1, 2)  # (batch, head, time2, d_k)
+
+        return q, k, v
+
+    def forward_attention(
+        self,
+        value: torch.Tensor,
+        scores: torch.Tensor,
+        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
+    ) -> torch.Tensor:
+        """Compute attention context vector.
+
+        Args:
+            value (torch.Tensor): Transformed value, size
+                (#batch, n_head, time2, d_k).
+            scores (torch.Tensor): Attention score, size
+                (#batch, n_head, time1, time2).
+            mask (torch.Tensor): Mask, size (#batch, 1, time2) or
+                (#batch, time1, time2), (0, 0, 0) means fake mask.
+
+        Returns:
+            torch.Tensor: Transformed value (#batch, time1, d_model)
+                weighted by the attention score (#batch, time1, time2).
+
+        """
+        n_batch = value.size(0)
+        # NOTE(xcsong): When will `if mask.size(2) > 0` be True?
+        #   1. onnx(16/4) [WHY? Because we feed real cache & real mask for the
+        #           1st chunk to ease the onnx export.]
+        #   2. pytorch training
+        if mask.size(2) > 0:  # time2 > 0
+            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
+            # For last chunk, time2 might be larger than scores.size(-1)
+            mask = mask[:, :, :, :scores.size(-1)]  # (batch, 1, *, time2)
+            scores = scores.masked_fill(mask, -float('inf'))
+            attn = torch.softmax(scores, dim=-1).masked_fill(
+                mask, 0.0)  # (batch, head, time1, time2)
+        # NOTE(xcsong): When will `if mask.size(2) > 0` be False?
+        #   1. onnx(16/-1, -1/-1, 16/0)
+        #   2. jit (16/-1, -1/-1, 16/0, 16/4)
+        else:
+            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+
+        p_attn = self.dropout(attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = (x.transpose(1, 2).contiguous().view(n_batch, -1,
+                                                 self.h * self.d_k)
+             )  # (batch, time1, d_model)
+
+        return self.linear_out(x)  # (batch, time1, d_model)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        pos_emb: torch.Tensor = torch.empty(0),
+        cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Compute scaled dot product attention.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2).
+                1.When applying cross attention between decoder and encoder,
+                the batch padding mask for input is in (#batch, 1, T) shape.
+                2.When applying self attention of encoder,
+                the mask is in (#batch, T, T)  shape.
+                3.When applying self attention of decoder,
+                the mask is in (#batch, L, L)  shape.
+                4.If the different position in decoder see different block
+                of the encoder, such as Mocha, the passed in mask could be
+                in (#batch, L, T) shape. But there is no such case in current
+                CosyVoice.
+            cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2),
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == size`
+
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+            torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2)
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == size`
+
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+
+        # NOTE(xcsong):
+        #   when export onnx model, for 1st chunk, we feed
+        #       cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode)
+        #       or cache(1, head, real_cache_t, d_k * 2) (16/4 mode).
+        #       In all modes, `if cache.size(0) > 0` will alwayse be `True`
+        #       and we will always do splitting and
+        #       concatnation(this will simplify onnx export). Note that
+        #       it's OK to concat & split zero-shaped tensors(see code below).
+        #   when export jit  model, for 1st chunk, we always feed
+        #       cache(0, 0, 0, 0) since jit supports dynamic if-branch.
+        # >>> a = torch.ones((1, 2, 0, 4))
+        # >>> b = torch.ones((1, 2, 3, 4))
+        # >>> c = torch.cat((a, b), dim=2)
+        # >>> torch.equal(b, c)        # True
+        # >>> d = torch.split(a, 2, dim=-1)
+        # >>> torch.equal(d[0], d[1])  # True
+        if cache.size(0) > 0:
+            key_cache, value_cache = torch.split(cache,
+                                                 cache.size(-1) // 2,
+                                                 dim=-1)
+            k = torch.cat([key_cache, k], dim=2)
+            v = torch.cat([value_cache, v], dim=2)
+        # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's
+        #   non-trivial to calculate `next_cache_start` here.
+        new_cache = torch.cat((k, v), dim=-1)
+
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+        return self.forward_attention(v, scores, mask), new_cache
+
+
+class RelPositionMultiHeadedAttention(MultiHeadedAttention):
+    """Multi-Head Attention layer with relative position encoding.
+    Paper: https://arxiv.org/abs/1901.02860
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+    """
+
+    def __init__(self,
+                 n_head: int,
+                 n_feat: int,
+                 dropout_rate: float,
+                 key_bias: bool = True):
+        """Construct an RelPositionMultiHeadedAttention object."""
+        super().__init__(n_head, n_feat, dropout_rate, key_bias)
+        # linear transformation for positional encoding
+        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
+        # these two learnable bias are used in matrix c and matrix d
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k))
+        self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k))
+        torch.nn.init.xavier_uniform_(self.pos_bias_u)
+        torch.nn.init.xavier_uniform_(self.pos_bias_v)
+
+    def rel_shift(self, x: torch.Tensor) -> torch.Tensor:
+        """Compute relative positional encoding.
+
+        Args:
+            x (torch.Tensor): Input tensor (batch, head, time1, 2*time1-1).
+            time1 means the length of query vector.
+
+        Returns:
+            torch.Tensor: Output tensor.
+
+        """
+        zero_pad = torch.zeros((x.size()[0], x.size()[1], x.size()[2], 1),
+                               device=x.device,
+                               dtype=x.dtype)
+        x_padded = torch.cat([zero_pad, x], dim=-1)
+
+        x_padded = x_padded.view(x.size()[0],
+                                 x.size()[1],
+                                 x.size(3) + 1, x.size(2))
+        x = x_padded[:, :, 1:].view_as(x)[
+            :, :, :, : x.size(-1) // 2 + 1
+        ]  # only keep the positions from 0 to time2
+        return x
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        pos_emb: torch.Tensor = torch.empty(0),
+        cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Compute 'Scaled Dot Product Attention' with rel. positional encoding.
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2), (0, 0, 0) means fake mask.
+            pos_emb (torch.Tensor): Positional embedding tensor
+                (#batch, time2, size).
+            cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2),
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == size`
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+            torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2)
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == size`
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+        q = q.transpose(1, 2)  # (batch, time1, head, d_k)
+
+        # NOTE(xcsong):
+        #   when export onnx model, for 1st chunk, we feed
+        #       cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode)
+        #       or cache(1, head, real_cache_t, d_k * 2) (16/4 mode).
+        #       In all modes, `if cache.size(0) > 0` will alwayse be `True`
+        #       and we will always do splitting and
+        #       concatnation(this will simplify onnx export). Note that
+        #       it's OK to concat & split zero-shaped tensors(see code below).
+        #   when export jit  model, for 1st chunk, we always feed
+        #       cache(0, 0, 0, 0) since jit supports dynamic if-branch.
+        # >>> a = torch.ones((1, 2, 0, 4))
+        # >>> b = torch.ones((1, 2, 3, 4))
+        # >>> c = torch.cat((a, b), dim=2)
+        # >>> torch.equal(b, c)        # True
+        # >>> d = torch.split(a, 2, dim=-1)
+        # >>> torch.equal(d[0], d[1])  # True
+        if cache.size(0) > 0:
+            key_cache, value_cache = torch.split(cache,
+                                                 cache.size(-1) // 2,
+                                                 dim=-1)
+            k = torch.cat([key_cache, k], dim=2)
+            v = torch.cat([value_cache, v], dim=2)
+        # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's
+        #   non-trivial to calculate `next_cache_start` here.
+        new_cache = torch.cat((k, v), dim=-1)
+
+        n_batch_pos = pos_emb.size(0)
+        p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)
+        p = p.transpose(1, 2)  # (batch, head, time1, d_k)
+
+        # (batch, head, time1, d_k)
+        q_with_bias_u = (q + self.pos_bias_u.to(q.device)).transpose(1, 2)
+        # (batch, head, time1, d_k)
+        q_with_bias_v = (q + self.pos_bias_v.to(q.device)).transpose(1, 2)
+
+        # compute attention score
+        # first compute matrix a and matrix c
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        # (batch, head, time1, time2)
+        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
+
+        # compute matrix b and matrix d
+        # (batch, head, time1, time2)
+        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
+        # NOTE(Xiang Lyu): Keep rel_shift since espnet rel_pos_emb is used
+        if matrix_ac.shape != matrix_bd.shape:
+            matrix_bd = self.rel_shift(matrix_bd)
+
+        scores = (matrix_ac + matrix_bd) / math.sqrt(
+            self.d_k)  # (batch, head, time1, time2)
+
+        return self.forward_attention(v, scores, mask), new_cache

HF_Deploy/src/chatterbox/models/s3gen/transformer/convolution.py

@@ -0,0 +1,145 @@
+# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Di Wu)
+#               2024 Alibaba Inc (Xiang Lyu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Modified from ESPnet(https://github.com/espnet/espnet)
+"""ConvolutionModule definition."""
+
+from typing import Tuple
+
+import torch
+from torch import nn
+
+
+class ConvolutionModule(nn.Module):
+    """ConvolutionModule in Conformer model."""
+
+    def __init__(self,
+                 channels: int,
+                 kernel_size: int = 15,
+                 activation: nn.Module = nn.ReLU(),
+                 norm: str = "batch_norm",
+                 causal: bool = False,
+                 bias: bool = True):
+        """Construct an ConvolutionModule object.
+        Args:
+            channels (int): The number of channels of conv layers.
+            kernel_size (int): Kernel size of conv layers.
+            causal (int): Whether use causal convolution or not
+        """
+        super().__init__()
+
+        self.pointwise_conv1 = nn.Conv1d(
+            channels,
+            2 * channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=bias,
+        )
+        # self.lorder is used to distinguish if it's a causal convolution,
+        # if self.lorder > 0: it's a causal convolution, the input will be
+        #    padded with self.lorder frames on the left in forward.
+        # else: it's a symmetrical convolution
+        if causal:
+            padding = 0
+            self.lorder = kernel_size - 1
+        else:
+            # kernel_size should be an odd number for none causal convolution
+            assert (kernel_size - 1) % 2 == 0
+            padding = (kernel_size - 1) // 2
+            self.lorder = 0
+        self.depthwise_conv = nn.Conv1d(
+            channels,
+            channels,
+            kernel_size,
+            stride=1,
+            padding=padding,
+            groups=channels,
+            bias=bias,
+        )
+
+        assert norm in ['batch_norm', 'layer_norm']
+        if norm == "batch_norm":
+            self.use_layer_norm = False
+            self.norm = nn.BatchNorm1d(channels)
+        else:
+            self.use_layer_norm = True
+            self.norm = nn.LayerNorm(channels)
+
+        self.pointwise_conv2 = nn.Conv1d(
+            channels,
+            channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=bias,
+        )
+        self.activation = activation
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        cache: torch.Tensor = torch.zeros((0, 0, 0)),
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Compute convolution module.
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, channels).
+            mask_pad (torch.Tensor): used for batch padding (#batch, 1, time),
+                (0, 0, 0) means fake mask.
+            cache (torch.Tensor): left context cache, it is only
+                used in causal convolution (#batch, channels, cache_t),
+                (0, 0, 0) meas fake cache.
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, channels).
+        """
+        # exchange the temporal dimension and the feature dimension
+        x = x.transpose(1, 2)  # (#batch, channels, time)
+
+        # mask batch padding
+        if mask_pad.size(2) > 0:  # time > 0
+            x.masked_fill_(~mask_pad, 0.0)
+
+        if self.lorder > 0:
+            if cache.size(2) == 0:  # cache_t == 0
+                x = nn.functional.pad(x, (self.lorder, 0), 'constant', 0.0)
+            else:
+                assert cache.size(0) == x.size(0)  # equal batch
+                assert cache.size(1) == x.size(1)  # equal channel
+                x = torch.cat((cache, x), dim=2)
+            assert (x.size(2) > self.lorder)
+            new_cache = x[:, :, -self.lorder:]
+        else:
+            # It's better we just return None if no cache is required,
+            # However, for JIT export, here we just fake one tensor instead of
+            # None.
+            new_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device)
+
+        # GLU mechanism
+        x = self.pointwise_conv1(x)  # (batch, 2*channel, dim)
+        x = nn.functional.glu(x, dim=1)  # (batch, channel, dim)
+
+        # 1D Depthwise Conv
+        x = self.depthwise_conv(x)
+        if self.use_layer_norm:
+            x = x.transpose(1, 2)
+        x = self.activation(self.norm(x))
+        if self.use_layer_norm:
+            x = x.transpose(1, 2)
+        x = self.pointwise_conv2(x)
+        # mask batch padding
+        if mask_pad.size(2) > 0:  # time > 0
+            x.masked_fill_(~mask_pad, 0.0)
+
+        return x.transpose(1, 2), new_cache

HF_Deploy/src/chatterbox/models/s3gen/transformer/embedding.py

@@ -0,0 +1,294 @@
+# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Di Wu)
+#               2024 Alibaba Inc (Xiang Lyu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Modified from ESPnet(https://github.com/espnet/espnet)
+"""Positonal Encoding Module."""
+
+import math
+from typing import Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+
+class PositionalEncoding(torch.nn.Module):
+    """Positional encoding.
+
+    :param int d_model: embedding dim
+    :param float dropout_rate: dropout rate
+    :param int max_len: maximum input length
+
+    PE(pos, 2i)   = sin(pos/(10000^(2i/dmodel)))
+    PE(pos, 2i+1) = cos(pos/(10000^(2i/dmodel)))
+    """
+
+    def __init__(self,
+                 d_model: int,
+                 dropout_rate: float,
+                 max_len: int = 5000,
+                 reverse: bool = False):
+        """Construct an PositionalEncoding object."""
+        super().__init__()
+        self.d_model = d_model
+        self.xscale = math.sqrt(self.d_model)
+        self.dropout = torch.nn.Dropout(p=dropout_rate)
+        self.max_len = max_len
+
+        self.pe = torch.zeros(self.max_len, self.d_model)
+        position = torch.arange(0, self.max_len,
+                                dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, self.d_model, 2, dtype=torch.float32) *
+            -(math.log(10000.0) / self.d_model))
+        self.pe[:, 0::2] = torch.sin(position * div_term)
+        self.pe[:, 1::2] = torch.cos(position * div_term)
+        self.pe = self.pe.unsqueeze(0)
+
+    def forward(self,
+                x: torch.Tensor,
+                offset: Union[int, torch.Tensor] = 0) \
+            -> Tuple[torch.Tensor, torch.Tensor]:
+        """Add positional encoding.
+
+        Args:
+            x (torch.Tensor): Input. Its shape is (batch, time, ...)
+            offset (int, torch.tensor): position offset
+
+        Returns:
+            torch.Tensor: Encoded tensor. Its shape is (batch, time, ...)
+            torch.Tensor: for compatibility to RelPositionalEncoding
+        """
+
+        self.pe = self.pe.to(x.device)
+        pos_emb = self.position_encoding(offset, x.size(1), False)
+        x = x * self.xscale + pos_emb
+        return self.dropout(x), self.dropout(pos_emb)
+
+    def position_encoding(self,
+                          offset: Union[int, torch.Tensor],
+                          size: int,
+                          apply_dropout: bool = True) -> torch.Tensor:
+        """ For getting encoding in a streaming fashion
+
+        Attention!!!!!
+        we apply dropout only once at the whole utterance level in a none
+        streaming way, but will call this function several times with
+        increasing input size in a streaming scenario, so the dropout will
+        be applied several times.
+
+        Args:
+            offset (int or torch.tensor): start offset
+            size (int): required size of position encoding
+
+        Returns:
+            torch.Tensor: Corresponding encoding
+        """
+        # How to subscript a Union type:
+        #   https://github.com/pytorch/pytorch/issues/69434
+        if isinstance(offset, int):
+            assert offset + size <= self.max_len
+            pos_emb = self.pe[:, offset:offset + size]
+        elif isinstance(offset, torch.Tensor) and offset.dim() == 0:  # scalar
+            assert offset + size <= self.max_len
+            pos_emb = self.pe[:, offset:offset + size]
+        else:  # for batched streaming decoding on GPU
+            assert torch.max(offset) + size <= self.max_len
+            index = offset.unsqueeze(1) + \
+                torch.arange(0, size).to(offset.device)  # B X T
+            flag = index > 0
+            # remove negative offset
+            index = index * flag
+            pos_emb = F.embedding(index, self.pe[0])  # B X T X d_model
+
+        if apply_dropout:
+            pos_emb = self.dropout(pos_emb)
+        return pos_emb
+
+
+class RelPositionalEncoding(PositionalEncoding):
+    """Relative positional encoding module.
+    See : Appendix B in https://arxiv.org/abs/1901.02860
+    Args:
+        d_model (int): Embedding dimension.
+        dropout_rate (float): Dropout rate.
+        max_len (int): Maximum input length.
+    """
+
+    def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000):
+        """Initialize class."""
+        super().__init__(d_model, dropout_rate, max_len, reverse=True)
+
+    def forward(self,
+                x: torch.Tensor,
+                offset: Union[int, torch.Tensor] = 0) \
+            -> Tuple[torch.Tensor, torch.Tensor]:
+        """Compute positional encoding.
+        Args:
+            x (torch.Tensor): Input tensor (batch, time, `*`).
+        Returns:
+            torch.Tensor: Encoded tensor (batch, time, `*`).
+            torch.Tensor: Positional embedding tensor (1, time, `*`).
+        """
+        self.pe = self.pe.to(x.device)
+        x = x * self.xscale
+        pos_emb = self.position_encoding(offset, x.size(1), False)
+        return self.dropout(x), self.dropout(pos_emb)
+
+
+class WhisperPositionalEncoding(PositionalEncoding):
+    """ Sinusoids position encoding used in openai-whisper.encoder
+    """
+
+    def __init__(self, d_model: int, dropout_rate: float, max_len: int = 1500):
+        super().__init__(d_model, dropout_rate, max_len)
+        self.xscale = 1.0
+        log_timescale_increment = np.log(10000) / (d_model // 2 - 1)
+        inv_timescales = torch.exp(-log_timescale_increment *
+                                   torch.arange(d_model // 2))
+        scaled_time = torch.arange(max_len)[:, np.newaxis] * \
+            inv_timescales[np.newaxis, :]
+        pe = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
+        delattr(self, "pe")
+        self.register_buffer("pe", pe.unsqueeze(0))
+
+
+class LearnablePositionalEncoding(PositionalEncoding):
+    """ Learnable position encoding used in openai-whisper.decoder
+    """
+
+    def __init__(self, d_model: int, dropout_rate: float, max_len: int = 448):
+        super().__init__(d_model, dropout_rate, max_len)
+        # NOTE(xcsong): overwrite self.pe & self.xscale
+        self.pe = torch.nn.Parameter(torch.empty(1, max_len, d_model))
+        self.xscale = 1.0
+
+
+class NoPositionalEncoding(torch.nn.Module):
+    """ No position encoding
+    """
+
+    def __init__(self, d_model: int, dropout_rate: float):
+        super().__init__()
+        self.d_model = d_model
+        self.dropout = torch.nn.Dropout(p=dropout_rate)
+
+    def forward(self,
+                x: torch.Tensor,
+                offset: Union[int, torch.Tensor] = 0) \
+            -> Tuple[torch.Tensor, torch.Tensor]:
+        """ Just return zero vector for interface compatibility
+        """
+        pos_emb = torch.zeros(1, x.size(1), self.d_model).to(x.device)
+        return self.dropout(x), pos_emb
+
+    def position_encoding(self, offset: Union[int, torch.Tensor],
+                          size: int) -> torch.Tensor:
+        return torch.zeros(1, size, self.d_model)
+
+
+class EspnetRelPositionalEncoding(torch.nn.Module):
+    """Relative positional encoding module (new implementation).
+
+    Details can be found in https://github.com/espnet/espnet/pull/2816.
+
+    See : Appendix B in https://arxiv.org/abs/1901.02860
+
+    Args:
+        d_model (int): Embedding dimension.
+        dropout_rate (float): Dropout rate.
+        max_len (int): Maximum input length.
+
+    """
+
+    def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000):
+        """Construct an PositionalEncoding object."""
+        super(EspnetRelPositionalEncoding, self).__init__()
+        self.d_model = d_model
+        self.xscale = math.sqrt(self.d_model)
+        self.dropout = torch.nn.Dropout(p=dropout_rate)
+        self.pe = None
+        self.extend_pe(torch.tensor(0.0).expand(1, max_len))
+
+    def extend_pe(self, x: torch.Tensor):
+        """Reset the positional encodings."""
+        if self.pe is not None:
+            # self.pe contains both positive and negative parts
+            # the length of self.pe is 2 * input_len - 1
+            if self.pe.size(1) >= x.size(1) * 2 - 1:
+                if self.pe.dtype != x.dtype or self.pe.device != x.device:
+                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+                return
+        # Suppose `i` means to the position of query vecotr and `j` means the
+        # position of key vector. We use position relative positions when keys
+        # are to the left (i>j) and negative relative positions otherwise (i<j).
+        pe_positive = torch.zeros(x.size(1), self.d_model)
+        pe_negative = torch.zeros(x.size(1), self.d_model)
+        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, self.d_model, 2, dtype=torch.float32)
+            * -(math.log(10000.0) / self.d_model)
+        )
+        pe_positive[:, 0::2] = torch.sin(position * div_term)
+        pe_positive[:, 1::2] = torch.cos(position * div_term)
+        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)
+        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)
+
+        # Reserve the order of positive indices and concat both positive and
+        # negative indices. This is used to support the shifting trick
+        # as in https://arxiv.org/abs/1901.02860
+        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)
+        pe_negative = pe_negative[1:].unsqueeze(0)
+        pe = torch.cat([pe_positive, pe_negative], dim=1)
+        self.pe = pe.to(device=x.device, dtype=x.dtype)
+
+    def forward(self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0) \
+            -> Tuple[torch.Tensor, torch.Tensor]:
+        """Add positional encoding.
+
+        Args:
+            x (torch.Tensor): Input tensor (batch, time, `*`).
+
+        Returns:
+            torch.Tensor: Encoded tensor (batch, time, `*`).
+
+        """
+        self.extend_pe(x)
+        x = x * self.xscale
+        pos_emb = self.position_encoding(size=x.size(1), offset=offset)
+        return self.dropout(x), self.dropout(pos_emb)
+
+    def position_encoding(self,
+                          offset: Union[int, torch.Tensor],
+                          size: int) -> torch.Tensor:
+        """ For getting encoding in a streaming fashion
+
+        Attention!!!!!
+        we apply dropout only once at the whole utterance level in a none
+        streaming way, but will call this function several times with
+        increasing input size in a streaming scenario, so the dropout will
+        be applied several times.
+
+        Args:
+            offset (int or torch.tensor): start offset
+            size (int): required size of position encoding
+
+        Returns:
+            torch.Tensor: Corresponding encoding
+        """
+        pos_emb = self.pe[
+            :,
+            self.pe.size(1) // 2 - size + 1: self.pe.size(1) // 2 + size,
+        ]
+        return pos_emb

HF_Deploy/src/chatterbox/models/s3gen/transformer/encoder_layer.py

@@ -0,0 +1,236 @@
+# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu)
+#               2022 Xingchen Song (sxc19@mails.tsinghua.edu.cn)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Modified from ESPnet(https://github.com/espnet/espnet)
+"""Encoder self-attention layer definition."""
+
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+
+
+class TransformerEncoderLayer(nn.Module):
+    """Encoder layer module.
+
+    Args:
+        size (int): Input dimension.
+        self_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` or `RelPositionMultiHeadedAttention`
+            instance can be used as the argument.
+        feed_forward (torch.nn.Module): Feed-forward module instance.
+            `PositionwiseFeedForward`, instance can be used as the argument.
+        dropout_rate (float): Dropout rate.
+        normalize_before (bool):
+            True: use layer_norm before each sub-block.
+            False: to use layer_norm after each sub-block.
+    """
+
+    def __init__(
+        self,
+        size: int,
+        self_attn: torch.nn.Module,
+        feed_forward: torch.nn.Module,
+        dropout_rate: float,
+        normalize_before: bool = True,
+    ):
+        """Construct an EncoderLayer object."""
+        super().__init__()
+        self.self_attn = self_attn
+        self.feed_forward = feed_forward
+        self.norm1 = nn.LayerNorm(size, eps=1e-12)
+        self.norm2 = nn.LayerNorm(size, eps=1e-12)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.size = size
+        self.normalize_before = normalize_before
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor,
+        pos_emb: torch.Tensor,
+        mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        att_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+        cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Compute encoded features.
+
+        Args:
+            x (torch.Tensor): (#batch, time, size)
+            mask (torch.Tensor): Mask tensor for the input (#batch, time，time),
+                (0, 0, 0) means fake mask.
+            pos_emb (torch.Tensor): just for interface compatibility
+                to ConformerEncoderLayer
+            mask_pad (torch.Tensor): does not used in transformer layer,
+                just for unified api with conformer.
+            att_cache (torch.Tensor): Cache tensor of the KEY & VALUE
+                (#batch=1, head, cache_t1, d_k * 2), head * d_k == size.
+            cnn_cache (torch.Tensor): Convolution cache in conformer layer
+                (#batch=1, size, cache_t2), not used here, it's for interface
+                compatibility to ConformerEncoderLayer.
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time, time).
+            torch.Tensor: att_cache tensor,
+                (#batch=1, head, cache_t1 + time, d_k * 2).
+            torch.Tensor: cnn_cahce tensor (#batch=1, size, cache_t2).
+
+        """
+        residual = x
+        if self.normalize_before:
+            x = self.norm1(x)
+        x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb=pos_emb, cache=att_cache)
+        x = residual + self.dropout(x_att)
+        if not self.normalize_before:
+            x = self.norm1(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm2(x)
+        x = residual + self.dropout(self.feed_forward(x))
+        if not self.normalize_before:
+            x = self.norm2(x)
+
+        fake_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device)
+        return x, mask, new_att_cache, fake_cnn_cache
+
+
+class ConformerEncoderLayer(nn.Module):
+    """Encoder layer module.
+    Args:
+        size (int): Input dimension.
+        self_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` or `RelPositionMultiHeadedAttention`
+            instance can be used as the argument.
+        feed_forward (torch.nn.Module): Feed-forward module instance.
+            `PositionwiseFeedForward` instance can be used as the argument.
+        feed_forward_macaron (torch.nn.Module): Additional feed-forward module
+             instance.
+            `PositionwiseFeedForward` instance can be used as the argument.
+        conv_module (torch.nn.Module): Convolution module instance.
+            `ConvlutionModule` instance can be used as the argument.
+        dropout_rate (float): Dropout rate.
+        normalize_before (bool):
+            True: use layer_norm before each sub-block.
+            False: use layer_norm after each sub-block.
+    """
+
+    def __init__(
+        self,
+        size: int,
+        self_attn: torch.nn.Module,
+        feed_forward: Optional[nn.Module] = None,
+        feed_forward_macaron: Optional[nn.Module] = None,
+        conv_module: Optional[nn.Module] = None,
+        dropout_rate: float = 0.1,
+        normalize_before: bool = True,
+    ):
+        """Construct an EncoderLayer object."""
+        super().__init__()
+        self.self_attn = self_attn
+        self.feed_forward = feed_forward
+        self.feed_forward_macaron = feed_forward_macaron
+        self.conv_module = conv_module
+        self.norm_ff = nn.LayerNorm(size, eps=1e-12)  # for the FNN module
+        self.norm_mha = nn.LayerNorm(size, eps=1e-12)  # for the MHA module
+        if feed_forward_macaron is not None:
+            self.norm_ff_macaron = nn.LayerNorm(size, eps=1e-12)
+            self.ff_scale = 0.5
+        else:
+            self.ff_scale = 1.0
+        if self.conv_module is not None:
+            self.norm_conv = nn.LayerNorm(size, eps=1e-12)  # for the CNN module
+            self.norm_final = nn.LayerNorm(
+                size, eps=1e-12)  # for the final output of the block
+        self.dropout = nn.Dropout(dropout_rate)
+        self.size = size
+        self.normalize_before = normalize_before
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor,
+        pos_emb: torch.Tensor,
+        mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        att_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+        cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Compute encoded features.
+
+        Args:
+            x (torch.Tensor): (#batch, time, size)
+            mask (torch.Tensor): Mask tensor for the input (#batch, time，time),
+                (0, 0, 0) means fake mask.
+            pos_emb (torch.Tensor): positional encoding, must not be None
+                for ConformerEncoderLayer.
+            mask_pad (torch.Tensor): batch padding mask used for conv module.
+                (#batch, 1，time), (0, 0, 0) means fake mask.
+            att_cache (torch.Tensor): Cache tensor of the KEY & VALUE
+                (#batch=1, head, cache_t1, d_k * 2), head * d_k == size.
+            cnn_cache (torch.Tensor): Convolution cache in conformer layer
+                (#batch=1, size, cache_t2)
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time, time).
+            torch.Tensor: att_cache tensor,
+                (#batch=1, head, cache_t1 + time, d_k * 2).
+            torch.Tensor: cnn_cahce tensor (#batch, size, cache_t2).
+        """
+
+        # whether to use macaron style
+        if self.feed_forward_macaron is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_ff_macaron(x)
+            x = residual + self.ff_scale * self.dropout(
+                self.feed_forward_macaron(x))
+            if not self.normalize_before:
+                x = self.norm_ff_macaron(x)
+
+        # multi-headed self-attention module
+        residual = x
+        if self.normalize_before:
+            x = self.norm_mha(x)
+        x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb,
+                                              att_cache)
+        x = residual + self.dropout(x_att)
+        if not self.normalize_before:
+            x = self.norm_mha(x)
+
+        # convolution module
+        # Fake new cnn cache here, and then change it in conv_module
+        new_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device)
+        if self.conv_module is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_conv(x)
+            x, new_cnn_cache = self.conv_module(x, mask_pad, cnn_cache)
+            x = residual + self.dropout(x)
+
+            if not self.normalize_before:
+                x = self.norm_conv(x)
+
+        # feed forward module
+        residual = x
+        if self.normalize_before:
+            x = self.norm_ff(x)
+
+        x = residual + self.ff_scale * self.dropout(self.feed_forward(x))
+        if not self.normalize_before:
+            x = self.norm_ff(x)
+
+        if self.conv_module is not None:
+            x = self.norm_final(x)
+
+        return x, mask, new_att_cache, new_cnn_cache

HF_Deploy/src/chatterbox/models/s3gen/transformer/positionwise_feed_forward.py

@@ -0,0 +1,115 @@
+# Copyright (c) 2019 Shigeki Karita
+#               2020 Mobvoi Inc (Binbin Zhang)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Positionwise feed forward layer definition."""
+
+import torch
+
+
+class PositionwiseFeedForward(torch.nn.Module):
+    """Positionwise feed forward layer.
+
+    FeedForward are appied on each position of the sequence.
+    The output dim is same with the input dim.
+
+    Args:
+        idim (int): Input dimenstion.
+        hidden_units (int): The number of hidden units.
+        dropout_rate (float): Dropout rate.
+        activation (torch.nn.Module): Activation function
+    """
+
+    def __init__(
+            self,
+            idim: int,
+            hidden_units: int,
+            dropout_rate: float,
+            activation: torch.nn.Module = torch.nn.ReLU(),
+    ):
+        """Construct a PositionwiseFeedForward object."""
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = torch.nn.Linear(idim, hidden_units)
+        self.activation = activation
+        self.dropout = torch.nn.Dropout(dropout_rate)
+        self.w_2 = torch.nn.Linear(hidden_units, idim)
+
+    def forward(self, xs: torch.Tensor) -> torch.Tensor:
+        """Forward function.
+
+        Args:
+            xs: input tensor (B, L, D)
+        Returns:
+            output tensor, (B, L, D)
+        """
+        return self.w_2(self.dropout(self.activation(self.w_1(xs))))
+
+
+class MoEFFNLayer(torch.nn.Module):
+    """
+    Mixture of expert with Positionwise feed forward layer
+    See also figure 1 in https://arxiv.org/pdf/2305.15663.pdf
+    The output dim is same with the input dim.
+
+    Modified from https://github.com/Lightning-AI/lit-gpt/pull/823
+                  https://github.com/mistralai/mistral-src/blob/b46d6/moe_one_file_ref.py#L203-L219
+    Args:
+        n_expert: number of expert.
+        n_expert_per_token: The actual number of experts used for each frame
+        idim (int): Input dimenstion.
+        hidden_units (int): The number of hidden units.
+        dropout_rate (float): Dropout rate.
+        activation (torch.nn.Module): Activation function
+    """
+
+    def __init__(
+            self,
+            n_expert: int,
+            n_expert_per_token: int,
+            idim: int,
+            hidden_units: int,
+            dropout_rate: float,
+            activation: torch.nn.Module = torch.nn.ReLU(),
+    ):
+        super(MoEFFNLayer, self).__init__()
+        self.gate = torch.nn.Linear(idim, n_expert, bias=False)
+        self.experts = torch.nn.ModuleList(
+            PositionwiseFeedForward(idim, hidden_units, dropout_rate,
+                                    activation) for _ in range(n_expert))
+        self.n_expert_per_token = n_expert_per_token
+
+    def forward(self, xs: torch.Tensor) -> torch.Tensor:
+        """Foward function.
+        Args:
+            xs: input tensor (B, L, D)
+        Returns:
+            output tensor, (B, L, D)
+
+        """
+        B, L, D = xs.size(
+        )  # batch size, sequence length, embedding dimension (idim)
+        xs = xs.view(-1, D)  # (B*L, D)
+        router = self.gate(xs)  # (B*L, n_expert)
+        logits, indices = torch.topk(
+            router, self.n_expert_per_token
+        )  # probs:(B*L, n_expert), indices: (B*L, n_expert)
+        weights = torch.nn.functional.softmax(
+            logits, dim=1,
+            dtype=torch.float).to(dtype=xs.dtype)  # (B*L, n_expert_per_token)
+        output = torch.zeros_like(xs)  # (B*L, D)
+        for i, expert in enumerate(self.experts):
+            mask = indices == i
+            batch_idx, ith_expert = torch.where(mask)
+            output[batch_idx] += weights[batch_idx, ith_expert, None] * expert(
+                xs[batch_idx])
+        return output.view(B, L, D)

HF_Deploy/src/chatterbox/models/s3gen/transformer/subsampling.py

@@ -0,0 +1,383 @@
+# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu)
+#               2024 Alibaba Inc (Xiang Lyu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Modified from ESPnet(https://github.com/espnet/espnet)
+"""Subsampling layer definition."""
+
+from typing import Tuple, Union
+
+import torch
+
+
+class BaseSubsampling(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.right_context = 0
+        self.subsampling_rate = 1
+
+    def position_encoding(self, offset: Union[int, torch.Tensor],
+                          size: int) -> torch.Tensor:
+        return self.pos_enc.position_encoding(offset, size)
+
+
+class EmbedinigNoSubsampling(BaseSubsampling):
+    """Embedding input without subsampling
+    """
+
+    def __init__(self, idim: int, odim: int, dropout_rate: float,
+                 pos_enc_class: torch.nn.Module):
+        super().__init__()
+        self.embed = torch.nn.Embedding(idim, odim)
+        self.pos_enc = pos_enc_class
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        offset: Union[int, torch.Tensor] = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Input x.
+
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: linear input tensor (#batch, time', odim),
+                where time' = time .
+            torch.Tensor: linear input mask (#batch, 1, time'),
+                where time' = time .
+
+        """
+        x = self.embed(x)
+        x, pos_emb = self.pos_enc(x, offset)
+        return x, pos_emb, x_mask
+
+
+class LinearNoSubsampling(BaseSubsampling):
+    """Linear transform the input without subsampling
+
+    Args:
+        idim (int): Input dimension.
+        odim (int): Output dimension.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, idim: int, odim: int, dropout_rate: float,
+                 pos_enc_class: torch.nn.Module):
+        """Construct an linear object."""
+        super().__init__()
+        self.out = torch.nn.Sequential(
+            torch.nn.Linear(idim, odim),
+            torch.nn.LayerNorm(odim, eps=1e-5),
+            torch.nn.Dropout(dropout_rate),
+        )
+        self.pos_enc = pos_enc_class
+        self.right_context = 0
+        self.subsampling_rate = 1
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        offset: Union[int, torch.Tensor] = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Input x.
+
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: linear input tensor (#batch, time', odim),
+                where time' = time .
+            torch.Tensor: linear input mask (#batch, 1, time'),
+                where time' = time .
+
+        """
+        x = self.out(x)
+        x, pos_emb = self.pos_enc(x, offset)
+        return x, pos_emb, x_mask
+
+
+class Conv1dSubsampling2(BaseSubsampling):
+    """Convolutional 1D subsampling (to 1/2 length).
+       It is designed for Whisper, ref:
+       https://github.com/openai/whisper/blob/main/whisper/model.py
+
+    Args:
+        idim (int): Input dimension.
+        odim (int): Output dimension.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, idim: int, odim: int, dropout_rate: float,
+                 pos_enc_class: torch.nn.Module):
+        """Construct an Conv1dSubsampling2 object."""
+        super().__init__()
+        self.conv = torch.nn.Sequential(
+            torch.nn.Conv1d(idim, odim, kernel_size=3, padding=1),
+            torch.nn.GELU(),
+            torch.nn.Conv1d(odim, odim, kernel_size=3, stride=2, padding=1),
+            torch.nn.GELU(),
+        )
+        self.pos_enc = pos_enc_class
+        # The right context for every conv layer is computed by:
+        # (kernel_size - 1) * frame_rate_of_this_layer
+        self.subsampling_rate = 2
+        # 4 = (3 - 1) * 1 + (3 - 1) * 1
+        self.right_context = 4
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        offset: Union[int, torch.Tensor] = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Subsample x.
+
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: Subsampled tensor (#batch, time', odim),
+                where time' = time // 2.
+            torch.Tensor: Subsampled mask (#batch, 1, time'),
+                where time' = time // 2.
+            torch.Tensor: positional encoding
+
+        """
+        time = x.size(1)
+        x = x.transpose(1, 2)  # (b, f, t)
+        x = self.conv(x)
+        x = x.transpose(1, 2)  # (b, t, f)
+        x, pos_emb = self.pos_enc(x, offset)
+        return x, pos_emb, x_mask[:, :, (time + 1) % 2::2]
+
+
+class Conv2dSubsampling4(BaseSubsampling):
+    """Convolutional 2D subsampling (to 1/4 length).
+
+    Args:
+        idim (int): Input dimension.
+        odim (int): Output dimension.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, idim: int, odim: int, dropout_rate: float,
+                 pos_enc_class: torch.nn.Module):
+        """Construct an Conv2dSubsampling4 object."""
+        super().__init__()
+        self.conv = torch.nn.Sequential(
+            torch.nn.Conv2d(1, odim, 3, 2),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(odim, odim, 3, 2),
+            torch.nn.ReLU(),
+        )
+        self.out = torch.nn.Sequential(
+            torch.nn.Linear(odim * (((idim - 1) // 2 - 1) // 2), odim))
+        self.pos_enc = pos_enc_class
+        # The right context for every conv layer is computed by:
+        # (kernel_size - 1) * frame_rate_of_this_layer
+        self.subsampling_rate = 4
+        # 6 = (3 - 1) * 1 + (3 - 1) * 2
+        self.right_context = 6
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        offset: Union[int, torch.Tensor] = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Subsample x.
+
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: Subsampled tensor (#batch, time', odim),
+                where time' = time // 4.
+            torch.Tensor: Subsampled mask (#batch, 1, time'),
+                where time' = time // 4.
+            torch.Tensor: positional encoding
+
+        """
+        x = x.unsqueeze(1)  # (b, c=1, t, f)
+        x = self.conv(x)
+        b, c, t, f = x.size()
+        x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f))
+        x, pos_emb = self.pos_enc(x, offset)
+        return x, pos_emb, x_mask[:, :, 2::2][:, :, 2::2]
+
+
+class Conv2dSubsampling6(BaseSubsampling):
+    """Convolutional 2D subsampling (to 1/6 length).
+    Args:
+        idim (int): Input dimension.
+        odim (int): Output dimension.
+        dropout_rate (float): Dropout rate.
+        pos_enc (torch.nn.Module): Custom position encoding layer.
+    """
+
+    def __init__(self, idim: int, odim: int, dropout_rate: float,
+                 pos_enc_class: torch.nn.Module):
+        """Construct an Conv2dSubsampling6 object."""
+        super().__init__()
+        self.conv = torch.nn.Sequential(
+            torch.nn.Conv2d(1, odim, 3, 2),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(odim, odim, 5, 3),
+            torch.nn.ReLU(),
+        )
+        self.linear = torch.nn.Linear(odim * (((idim - 1) // 2 - 2) // 3),
+                                      odim)
+        self.pos_enc = pos_enc_class
+        # 10 = (3 - 1) * 1 + (5 - 1) * 2
+        self.subsampling_rate = 6
+        self.right_context = 10
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        offset: Union[int, torch.Tensor] = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Subsample x.
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: Subsampled tensor (#batch, time', odim),
+                where time' = time // 6.
+            torch.Tensor: Subsampled mask (#batch, 1, time'),
+                where time' = time // 6.
+            torch.Tensor: positional encoding
+        """
+        x = x.unsqueeze(1)  # (b, c, t, f)
+        x = self.conv(x)
+        b, c, t, f = x.size()
+        x = self.linear(x.transpose(1, 2).contiguous().view(b, t, c * f))
+        x, pos_emb = self.pos_enc(x, offset)
+        return x, pos_emb, x_mask[:, :, 2::2][:, :, 4::3]
+
+
+class Conv2dSubsampling8(BaseSubsampling):
+    """Convolutional 2D subsampling (to 1/8 length).
+
+    Args:
+        idim (int): Input dimension.
+        odim (int): Output dimension.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, idim: int, odim: int, dropout_rate: float,
+                 pos_enc_class: torch.nn.Module):
+        """Construct an Conv2dSubsampling8 object."""
+        super().__init__()
+        self.conv = torch.nn.Sequential(
+            torch.nn.Conv2d(1, odim, 3, 2),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(odim, odim, 3, 2),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(odim, odim, 3, 2),
+            torch.nn.ReLU(),
+        )
+        self.linear = torch.nn.Linear(
+            odim * ((((idim - 1) // 2 - 1) // 2 - 1) // 2), odim)
+        self.pos_enc = pos_enc_class
+        self.subsampling_rate = 8
+        # 14 = (3 - 1) * 1 + (3 - 1) * 2 + (3 - 1) * 4
+        self.right_context = 14
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        offset: Union[int, torch.Tensor] = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Subsample x.
+
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: Subsampled tensor (#batch, time', odim),
+                where time' = time // 8.
+            torch.Tensor: Subsampled mask (#batch, 1, time'),
+                where time' = time // 8.
+            torch.Tensor: positional encoding
+        """
+        x = x.unsqueeze(1)  # (b, c, t, f)
+        x = self.conv(x)
+        b, c, t, f = x.size()
+        x = self.linear(x.transpose(1, 2).contiguous().view(b, t, c * f))
+        x, pos_emb = self.pos_enc(x, offset)
+        return x, pos_emb, x_mask[:, :, 2::2][:, :, 2::2][:, :, 2::2]
+
+
+class LegacyLinearNoSubsampling(BaseSubsampling):
+    """Linear transform the input without subsampling
+
+    Args:
+        idim (int): Input dimension.
+        odim (int): Output dimension.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, idim: int, odim: int, dropout_rate: float,
+                 pos_enc_class: torch.nn.Module):
+        """Construct an linear object."""
+        super().__init__()
+        self.out = torch.nn.Sequential(
+            torch.nn.Linear(idim, odim),
+            torch.nn.LayerNorm(odim, eps=1e-5),
+            torch.nn.Dropout(dropout_rate),
+            torch.nn.ReLU(),
+        )
+        self.pos_enc = pos_enc_class
+        self.right_context = 0
+        self.subsampling_rate = 1
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        offset: Union[int, torch.Tensor] = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Input x.
+
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: linear input tensor (#batch, time', odim),
+                where time' = time .
+            torch.Tensor: linear input mask (#batch, 1, time'),
+                where time' = time .
+
+        """
+        x = self.out(x)
+        x, pos_emb = self.pos_enc(x, offset)
+        return x, pos_emb, x_mask

HF_Deploy/src/chatterbox/models/s3gen/transformer/upsample_encoder.py

@@ -0,0 +1,318 @@
+# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu)
+#               2022 Xingchen Song (sxc19@mails.tsinghua.edu.cn)
+#               2024 Alibaba Inc (Xiang Lyu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Modified from ESPnet(https://github.com/espnet/espnet)
+"""Encoder definition."""
+from typing import Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .convolution import ConvolutionModule
+from .encoder_layer import ConformerEncoderLayer
+from .positionwise_feed_forward import PositionwiseFeedForward
+from ..utils.class_utils import (
+    COSYVOICE_EMB_CLASSES,
+    COSYVOICE_SUBSAMPLE_CLASSES,
+    COSYVOICE_ATTENTION_CLASSES,
+    COSYVOICE_ACTIVATION_CLASSES,
+)
+from ..utils.mask import make_pad_mask
+from ..utils.mask import add_optional_chunk_mask
+
+
+class Upsample1D(nn.Module):
+    """A 1D upsampling layer with an optional convolution.
+
+    Parameters:
+        channels (`int`):
+            number of channels in the inputs and outputs.
+        use_conv (`bool`, default `False`):
+            option to use a convolution.
+        use_conv_transpose (`bool`, default `False`):
+            option to use a convolution transpose.
+        out_channels (`int`, optional):
+            number of output channels. Defaults to `channels`.
+    """
+
+    def __init__(self, channels: int, out_channels: int, stride: int = 2):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels
+        self.stride = stride
+        # In this mode, first repeat interpolate, than conv with stride=1
+        self.conv = nn.Conv1d(self.channels, self.out_channels, stride * 2 + 1, stride=1, padding=0)
+
+    def forward(self, inputs: torch.Tensor, input_lengths: torch.Tensor):
+        outputs = F.interpolate(inputs, scale_factor=float(self.stride), mode="nearest")
+        outputs = F.pad(outputs, (self.stride * 2, 0), value=0.0)
+        outputs = self.conv(outputs)
+        return outputs, input_lengths * self.stride
+
+
+class PreLookaheadLayer(nn.Module):
+    def __init__(self, channels: int, pre_lookahead_len: int = 1):
+        super().__init__()
+        self.channels = channels
+        self.pre_lookahead_len = pre_lookahead_len
+        self.conv1 = nn.Conv1d(
+            channels, channels,
+            kernel_size=pre_lookahead_len + 1,
+            stride=1, padding=0,
+        )
+        self.conv2 = nn.Conv1d(
+            channels, channels,
+            kernel_size=3, stride=1, padding=0,
+        )
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """
+        inputs: (batch_size, seq_len, channels)
+        """
+        outputs = inputs.transpose(1, 2).contiguous()
+        # look ahead
+        outputs = F.pad(outputs, (0, self.pre_lookahead_len), mode='constant', value=0.0)
+        outputs = F.leaky_relu(self.conv1(outputs))
+        # outputs
+        outputs = F.pad(outputs, (2, 0), mode='constant', value=0.0)
+        outputs = self.conv2(outputs)
+        outputs = outputs.transpose(1, 2).contiguous()
+
+        # residual connection
+        outputs = outputs + inputs
+        return outputs
+
+
+class UpsampleConformerEncoder(torch.nn.Module):
+
+    def __init__(
+        self,
+        input_size: int = 512,
+        output_size: int = 512,
+        attention_heads: int = 8,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        attention_dropout_rate: float = 0.1,
+        input_layer: str = "linear",
+        pos_enc_layer_type: str = "rel_pos_espnet",
+        normalize_before: bool = True,
+        static_chunk_size: int = 0,
+        use_dynamic_chunk: bool = False,
+        global_cmvn: torch.nn.Module = None,
+        use_dynamic_left_chunk: bool = False,
+        positionwise_conv_kernel_size: int = 1,
+        macaron_style: bool = False,
+        selfattention_layer_type: str = "rel_selfattn",
+        activation_type: str = "swish",
+        use_cnn_module: bool = False,
+        cnn_module_kernel: int = 15,
+        causal: bool = False,
+        cnn_module_norm: str = "batch_norm",
+        key_bias: bool = True,
+        gradient_checkpointing: bool = False,
+    ):
+        """
+        Args:
+            input_size (int): input dim
+            output_size (int): dimension of attention
+            attention_heads (int): the number of heads of multi head attention
+            linear_units (int): the hidden units number of position-wise feed
+                forward
+            num_blocks (int): the number of decoder blocks
+            dropout_rate (float): dropout rate
+            attention_dropout_rate (float): dropout rate in attention
+            positional_dropout_rate (float): dropout rate after adding
+                positional encoding
+            input_layer (str): input layer type.
+                optional [linear, conv2d, conv2d6, conv2d8]
+            pos_enc_layer_type (str): Encoder positional encoding layer type.
+                opitonal [abs_pos, scaled_abs_pos, rel_pos, no_pos]
+            normalize_before (bool):
+                True: use layer_norm before each sub-block of a layer.
+                False: use layer_norm after each sub-block of a layer.
+            static_chunk_size (int): chunk size for static chunk training and
+                decoding
+            use_dynamic_chunk (bool): whether use dynamic chunk size for
+                training or not, You can only use fixed chunk(chunk_size > 0)
+                or dyanmic chunk size(use_dynamic_chunk = True)
+            global_cmvn (Optional[torch.nn.Module]): Optional GlobalCMVN module
+            use_dynamic_left_chunk (bool): whether use dynamic left chunk in
+                dynamic chunk training
+            key_bias: whether use bias in attention.linear_k, False for whisper models.
+            gradient_checkpointing: rerunning a forward-pass segment for each
+                checkpointed segment during backward.
+        """
+        super().__init__()
+        self._output_size = output_size
+
+        self.global_cmvn = global_cmvn
+        self.embed = COSYVOICE_SUBSAMPLE_CLASSES[input_layer](
+            input_size,
+            output_size,
+            dropout_rate,
+            COSYVOICE_EMB_CLASSES[pos_enc_layer_type](output_size,
+                                                      positional_dropout_rate),
+        )
+
+        self.normalize_before = normalize_before
+        self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5)
+        self.static_chunk_size = static_chunk_size
+        self.use_dynamic_chunk = use_dynamic_chunk
+        self.use_dynamic_left_chunk = use_dynamic_left_chunk
+        self.gradient_checkpointing = gradient_checkpointing
+        activation = COSYVOICE_ACTIVATION_CLASSES[activation_type]()
+        # self-attention module definition
+        encoder_selfattn_layer_args = (
+            attention_heads,
+            output_size,
+            attention_dropout_rate,
+            key_bias,
+        )
+        # feed-forward module definition
+        positionwise_layer_args = (
+            output_size,
+            linear_units,
+            dropout_rate,
+            activation,
+        )
+        # convolution module definition
+        convolution_layer_args = (output_size, cnn_module_kernel, activation,
+                                  cnn_module_norm, causal)
+        self.pre_lookahead_layer = PreLookaheadLayer(channels=512, pre_lookahead_len=3)
+        self.encoders = torch.nn.ModuleList([
+            ConformerEncoderLayer(
+                output_size,
+                COSYVOICE_ATTENTION_CLASSES[selfattention_layer_type](
+                    *encoder_selfattn_layer_args),
+                PositionwiseFeedForward(*positionwise_layer_args),
+                PositionwiseFeedForward(
+                    *positionwise_layer_args) if macaron_style else None,
+                ConvolutionModule(
+                    *convolution_layer_args) if use_cnn_module else None,
+                dropout_rate,
+                normalize_before,
+            ) for _ in range(num_blocks)
+        ])
+        self.up_layer = Upsample1D(channels=512, out_channels=512, stride=2)
+        self.up_embed = COSYVOICE_SUBSAMPLE_CLASSES[input_layer](
+            input_size,
+            output_size,
+            dropout_rate,
+            COSYVOICE_EMB_CLASSES[pos_enc_layer_type](output_size,
+                                                      positional_dropout_rate),
+        )
+        self.up_encoders = torch.nn.ModuleList([
+            ConformerEncoderLayer(
+                output_size,
+                COSYVOICE_ATTENTION_CLASSES[selfattention_layer_type](
+                    *encoder_selfattn_layer_args),
+                PositionwiseFeedForward(*positionwise_layer_args),
+                PositionwiseFeedForward(
+                    *positionwise_layer_args) if macaron_style else None,
+                ConvolutionModule(
+                    *convolution_layer_args) if use_cnn_module else None,
+                dropout_rate,
+                normalize_before,
+            ) for _ in range(4)
+        ])
+
+    def output_size(self) -> int:
+        return self._output_size
+
+    def forward(
+        self,
+        xs: torch.Tensor,
+        xs_lens: torch.Tensor,
+        decoding_chunk_size: int = 0,
+        num_decoding_left_chunks: int = -1,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Embed positions in tensor.
+
+        Args:
+            xs: padded input tensor (B, T, D)
+            xs_lens: input length (B)
+            decoding_chunk_size: decoding chunk size for dynamic chunk
+                0: default for training, use random dynamic chunk.
+                <0: for decoding, use full chunk.
+                >0: for decoding, use fixed chunk size as set.
+            num_decoding_left_chunks: number of left chunks, this is for decoding,
+            the chunk size is decoding_chunk_size.
+                >=0: use num_decoding_left_chunks
+                <0: use all left chunks
+        Returns:
+            encoder output tensor xs, and subsampled masks
+            xs: padded output tensor (B, T' ~= T/subsample_rate, D)
+            masks: torch.Tensor batch padding mask after subsample
+                (B, 1, T' ~= T/subsample_rate)
+        NOTE(xcsong):
+            We pass the `__call__` method of the modules instead of `forward` to the
+            checkpointing API because `__call__` attaches all the hooks of the module.
+            https://discuss.pytorch.org/t/any-different-between-model-input-and-model-forward-input/3690/2
+        """
+        T = xs.size(1)
+        masks = ~make_pad_mask(xs_lens, T).unsqueeze(1)  # (B, 1, T)
+        if self.global_cmvn is not None:
+            xs = self.global_cmvn(xs)
+        xs, pos_emb, masks = self.embed(xs, masks)
+        mask_pad = masks  # (B, 1, T/subsample_rate)
+        chunk_masks = add_optional_chunk_mask(xs, masks,
+                                              self.use_dynamic_chunk,
+                                              self.use_dynamic_left_chunk,
+                                              decoding_chunk_size,
+                                              self.static_chunk_size,
+                                              num_decoding_left_chunks)
+        # lookahead + conformer encoder
+        xs = self.pre_lookahead_layer(xs)
+        xs = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad)
+
+        # upsample + conformer encoder
+        xs = xs.transpose(1, 2).contiguous()
+        xs, xs_lens = self.up_layer(xs, xs_lens)
+        xs = xs.transpose(1, 2).contiguous()
+        T = xs.size(1)
+        masks = ~make_pad_mask(xs_lens, T).unsqueeze(1)  # (B, 1, T)
+        xs, pos_emb, masks = self.up_embed(xs, masks)
+        mask_pad = masks  # (B, 1, T/subsample_rate)
+        chunk_masks = add_optional_chunk_mask(xs, masks,
+                                              self.use_dynamic_chunk,
+                                              self.use_dynamic_left_chunk,
+                                              decoding_chunk_size,
+                                              self.static_chunk_size * self.up_layer.stride,
+                                              num_decoding_left_chunks)
+        xs = self.forward_up_layers(xs, chunk_masks, pos_emb, mask_pad)
+
+        if self.normalize_before:
+            xs = self.after_norm(xs)
+        # Here we assume the mask is not changed in encoder layers, so just
+        # return the masks before encoder layers, and the masks will be used
+        # for cross attention with decoder later
+        return xs, masks
+
+    def forward_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor,
+                       pos_emb: torch.Tensor,
+                       mask_pad: torch.Tensor) -> torch.Tensor:
+        for layer in self.encoders:
+            xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad)
+        return xs
+
+    def forward_up_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor,
+                          pos_emb: torch.Tensor,
+                          mask_pad: torch.Tensor) -> torch.Tensor:
+        for layer in self.up_encoders:
+            xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad)
+        return xs

HF_Deploy/src/chatterbox/models/s3gen/utils/class_utils.py

@@ -0,0 +1,71 @@
+# Copyright [2023-11-28] <sxc19@mails.tsinghua.edu.cn, Xingchen Song>
+#            2024 Alibaba Inc (authors: Xiang Lyu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+
+from ..transformer.activation import Swish
+from ..transformer.subsampling import (
+    LinearNoSubsampling,
+    EmbedinigNoSubsampling,
+    Conv1dSubsampling2,
+    Conv2dSubsampling4,
+    Conv2dSubsampling6,
+    Conv2dSubsampling8,
+)
+from ..transformer.embedding import (
+    PositionalEncoding,
+    RelPositionalEncoding,
+    WhisperPositionalEncoding,
+    LearnablePositionalEncoding,
+    NoPositionalEncoding)
+from ..transformer.attention import (MultiHeadedAttention,
+    RelPositionMultiHeadedAttention)
+from ..transformer.embedding import EspnetRelPositionalEncoding
+from ..transformer.subsampling import LegacyLinearNoSubsampling
+
+
+COSYVOICE_ACTIVATION_CLASSES = {
+    "hardtanh": torch.nn.Hardtanh,
+    "tanh": torch.nn.Tanh,
+    "relu": torch.nn.ReLU,
+    "selu": torch.nn.SELU,
+    "swish": getattr(torch.nn, "SiLU", Swish),
+    "gelu": torch.nn.GELU,
+}
+
+COSYVOICE_SUBSAMPLE_CLASSES = {
+    "linear": LinearNoSubsampling,
+    "linear_legacy": LegacyLinearNoSubsampling,
+    "embed": EmbedinigNoSubsampling,
+    "conv1d2": Conv1dSubsampling2,
+    "conv2d": Conv2dSubsampling4,
+    "conv2d6": Conv2dSubsampling6,
+    "conv2d8": Conv2dSubsampling8,
+    'paraformer_dummy': torch.nn.Identity
+}
+
+COSYVOICE_EMB_CLASSES = {
+    "embed": PositionalEncoding,
+    "abs_pos": PositionalEncoding,
+    "rel_pos": RelPositionalEncoding,
+    "rel_pos_espnet": EspnetRelPositionalEncoding,
+    "no_pos": NoPositionalEncoding,
+    "abs_pos_whisper": WhisperPositionalEncoding,
+    "embed_learnable_pe": LearnablePositionalEncoding,
+}
+
+COSYVOICE_ATTENTION_CLASSES = {
+    "selfattn": MultiHeadedAttention,
+    "rel_selfattn": RelPositionMultiHeadedAttention,
+}