Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,137 @@
+.vscode
+
+# Pylance
+pyrightconfig.json
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+src/**/__pycache__/
+src/**/*.pyc
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyderworkspace
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+syn_out/
+checkpoints/
+.gradio
+
+# Ignore generated sample .wav files
+**/*.wav
+
+# User data directories - keep repository clean
+venv/
+audiobook_projects/
+speakers/
+source/
+audiobook_config.json
+
+# Development and archive directories
+archive/
+development/
+
+# OS generated files
+.DS_Store
+.DS_Store?
+._*
+.Spotlight-V100
+.Trashes
+ehthumbs.db
+Thumbs.db

AUDIOBOOK_FEATURES.md

@@ -0,0 +1,169 @@
+# 🎧 Chatterbox TTS - Audiobook Edition Features
+
+## 🚀 New Voice Management System
+
+The Audiobook Edition adds powerful voice management capabilities perfect for creating consistent character voices across your audiobook projects.
+
+## ✨ Key Features
+
+### 📚 Voice Library Tab
+- **Organized Voice Storage**: Keep all your character voices in one place
+- **Custom Voice Profiles**: Save voice settings with names, descriptions, and reference audio
+- **Easy Voice Selection**: Quick dropdown to switch between saved voices
+- **Voice Testing**: Test voices before saving or using them
+
+### 🎭 Character Voice Management
+- **Voice Profiles**: Each voice includes:
+  - Voice name (for file organization)
+  - Display name (human-readable)
+  - Description (character notes)
+  - Reference audio file
+  - Optimized settings (exaggeration, CFG/pace, temperature)
+
+### 🎙️ Voice Testing & Configuration
+- **Live Testing**: Test voice settings with custom text
+- **Parameter Tuning**: Fine-tune exaggeration, CFG/pace, and temperature
+- **Instant Feedback**: Hear changes immediately
+- **Save Optimized Settings**: Store perfect settings for each character
+
+## 🛠️ How to Use
+
+### 1. Launch the Audiobook Edition
+```bash
+# Use the audiobook launcher
+launch_audiobook.bat
+```
+
+### 2. Set Up Your Voice Library
+1. Go to the **"📚 Voice Library"** tab
+2. Set your voice library folder path (default: `voice_library`)
+3. Click **"📁 Update Library Path"**
+
+### 3. Create a Voice Profile
+1. **Upload Reference Audio**: Upload 10-30 seconds of clear speech
+2. **Configure Settings**:
+   - **Exaggeration**: 0.3-0.7 for most voices
+   - **CFG/Pace**: Lower = slower, more deliberate
+   - **Temperature**: Higher = more variation
+3. **Test the Voice**: Use the test text to hear how it sounds
+4. **Save Profile**: Give it a name and description, then save
+
+### 4. Use Saved Voices
+1. **Select Voice**: Choose from dropdown in Voice Library
+2. **Load Voice**: Click "📥 Load Voice" to load settings
+3. **Generate Speech**: Switch to TTS tab and generate with loaded voice
+
+## 📁 Voice Library Structure
+
+```
+voice_library/
+├── narrator_male_deep/
+│   ├── config.json          # Voice settings
+│   └── reference.wav        # Reference audio
+├── character_female_young/
+│   ├── config.json
+│   └── reference.mp3
+└── villain_gravelly/
+    ├── config.json
+    └── reference.wav
+```
+
+## 🎯 Audiobook Workflow
+
+### Step 1: Character Planning
+- List all characters in your audiobook
+- Gather reference audio for each (record or find samples)
+- Plan voice characteristics (age, personality, accent)
+
+### Step 2: Voice Creation
+- Create a voice profile for each character
+- Test and refine settings for consistency
+- Save with descriptive names (e.g., "Harry_confident", "Hermione_intelligent")
+
+### Step 3: Production
+- Load character voice before generating their dialogue
+- Use consistent settings throughout the book
+- Test voice regularly to maintain quality
+
+### Step 4: Quality Control
+- Use the same test phrase for all characters
+- Ensure voices are distinguishable
+- Adjust settings if characters sound too similar
+
+## 💡 Pro Tips
+
+### Voice Creation
+- **Reference Audio**: Use clean, noise-free recordings
+- **Length**: 10-30 seconds is optimal
+- **Content**: Natural speech, not overly dramatic
+- **Quality**: Higher quality audio = better cloning
+
+### Settings Optimization
+- **Exaggeration**:
+  - 0.3-0.5: Subtle, natural voices
+  - 0.5-0.7: Standard character voices
+  - 0.7-1.0: Dramatic or distinctive voices
+  
+- **CFG/Pace**:
+  - 0.3-0.4: Slow, deliberate (elderly, wise characters)
+  - 0.5: Standard pace
+  - 0.6-0.8: Faster pace (young, energetic characters)
+
+- **Temperature**:
+  - 0.5-0.8: Consistent delivery
+  - 0.8-1.2: More natural variation
+  - 1.2+: Creative but less predictable
+
+### Organization
+- **Naming Convention**: Use descriptive names (character_trait_type)
+- **Descriptions**: Include character details and usage notes
+- **Backup**: Keep your voice_library folder backed up
+- **Version Control**: Save multiple versions for different emotions
+
+## 🔧 Advanced Features
+
+### Voice Library Management
+- **Import/Export**: Copy voice_library folder between projects
+- **Sharing**: Share voice profiles with other audiobook creators
+- **Backup**: Regular backups of your voice library
+- **Organization**: Folder structure for different projects
+
+### Batch Processing (Future)
+- Process entire chapters with character voice switching
+- Automatic voice detection based on speaker tags
+- Export management for audiobook production
+
+## 🎵 Example Character Voices
+
+### Narrator
+- **Settings**: Exaggeration 0.4, CFG 0.5, Temp 0.7
+- **Description**: Clear, neutral, professional tone
+- **Use**: Chapter narration, scene descriptions
+
+### Hero Character
+- **Settings**: Exaggeration 0.6, CFG 0.6, Temp 0.8
+- **Description**: Confident, determined, slightly higher energy
+- **Use**: Main character dialogue
+
+### Wise Mentor
+- **Settings**: Exaggeration 0.3, CFG 0.3, Temp 0.6
+- **Description**: Slow, deliberate, thoughtful delivery
+- **Use**: Advisor character, important wisdom
+
+### Comic Relief
+- **Settings**: Exaggeration 0.8, CFG 0.7, Temp 1.0
+- **Description**: Energetic, expressive, variable delivery
+- **Use**: Funny sidekick, lighthearted moments
+
+## 🛡️ Best Practices
+
+1. **Consistency**: Always use the same voice profile for each character
+2. **Testing**: Test voices regularly during production
+3. **Backup**: Keep voice profiles backed up
+4. **Documentation**: Maintain character voice notes
+5. **Quality**: Use high-quality reference audio
+6. **Organization**: Use clear naming conventions
+
+---
+
+**Ready to create amazing audiobooks with consistent character voices? Launch the Audiobook Edition and start building your voice library! 🎧✨** 

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Resemble AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,12 +1,368 @@
----
-title: Https Huggingface Co Spaces Jujutechnology Ebook2audiobook
-emoji: 📉
-colorFrom: pink
-colorTo: blue
-sdk: gradio
-sdk_version: 5.34.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🎧 Chatterbox Audiobook Generator
+
+**This is a work in progress. You can consider this a pre-launch repo at the moment, but if you find bugs, please put them in the issues area. Thank you.**
+**Transform your text into high-quality audiobooks with advanced TTS models, voice cloning, and professional volume normalization.**
+
+## 🚀 Quick Start
+
+### 1. Install Dependencies
+```bash
+./install-audiobook.bat
+```
+
+### 2. Launch the Application
+```bash
+./launch_audiobook.bat
+```
+
+### 3. CUDA Issue Fix (If Needed)
+If you encounter CUDA assertion errors during generation, install the patched version:
+```bash
+# Activate your virtual environment first
+venv\Scripts\activate.bat
+
+# Install the CUDA-fixed version
+pip install --force-reinstall --no-cache-dir "chatterbox-tts @ git+https://github.com/fakerybakery/better-chatterbox@fix-cuda-issue"
+```
+
+The web interface will open automatically in your browser at `http://localhost:7860`
+
+---
+
+## ✨ Features
+
+### 📚 **Audiobook Creation**
+- **Single Voice**: Generate entire audiobooks with one consistent voice
+- **Multi-Voice**: Create dynamic audiobooks with multiple characters
+- **Custom Voices**: Clone voices from audio samples for personalized narration
+- **Professional Volume Normalization**: Ensure consistent audio levels across all voices
+- **📋 Text Queuing System** ⭐ *NEW*: Upload books in any size chapters and generate continuously
+- **🔄 Chunk-Based Processing** ⭐ *NEW*: Improved reliability for longer text generations
+
+### 🎵 **Audio Processing**
+- **Smart Cleanup**: Remove unwanted silence and audio artifacts
+- **Volume Normalization**: Professional-grade volume balancing for all voices
+- **Real-time Audio Analysis**: Live volume level monitoring and feedback
+- **Preview System**: Test settings before applying to entire projects
+- **Batch Processing**: Process multiple projects efficiently
+- **Quality Control**: Advanced audio optimization tools
+- **🎯 Enhanced Audio Quality** ⭐ *NEW*: Improved P-top and minimum P parameters for better voice generation
+
+### 🎭 **Voice Management**
+- **Voice Library**: Organize and manage your voice collection
+- **Voice Cloning**: Create custom voices from audio samples
+- **Volume Settings**: Configure target volume levels for each voice
+- **Professional Presets**: Industry-standard volume levels (audiobook, podcast, broadcast)
+- **Character Assignment**: Map specific voices to story characters
+
+### 📊 **Volume Normalization System** ⭐ *NEW*
+- **Professional Standards**: Audiobook (-18 dB), Podcast (-16 dB), Broadcast (-23 dB) presets
+- **Consistent Character Voices**: All characters maintain the same volume level
+- **Real-time Analysis**: Color-coded volume status with RMS and peak level display
+- **Retroactive Normalization**: Apply volume settings to existing voice projects
+- **Multi-Voice Support**: Batch normalize all voices in multi-character audiobooks
+- **Soft Limiting**: Intelligent audio limiting to prevent distortion
+
+### 📖 **Text Processing**
+- **Chapter Support**: Automatic chapter detection and organization
+- **Multi-Voice Parsing**: Parse character dialogue automatically
+- **Text Validation**: Ensure proper formatting before generation
+- **📋 Queue Management** ⭐ *NEW*: Batch process multiple text files sequentially
+- **🔇 Return Pause System** ⭐ *NEW*: Automatic pause insertion based on line breaks for natural speech flow
+
+---
+
+## 🎭 Custom Audiobook Processing Pipeline ⭐ *NEW*
+
+Our advanced text processing pipeline transforms your written content into natural-sounding audiobooks with intelligent pause placement and character flow management.
+
+### 🔇 **Return Pause System**
+
+**Automatic pause insertion based on your text formatting** - Every line break (`\n`) in your text automatically adds a 0.1-second pause to the generated audio, creating natural speech rhythms without manual intervention.
+
+#### **How It Works**
+- **Line Break Detection**: System automatically counts all line breaks in your text
+- **Pause Calculation**: Each return adds exactly 0.1 seconds of silence
+- **Accumulative Pauses**: Multiple consecutive line breaks create longer pauses
+- **Universal Support**: Works with single-voice, multi-voice, and batch processing
+
+#### **Example Text Formatting**
+```
+[Narrator] The sun was setting over the hills.
+
+[Character1] "We need to find shelter soon."
+
+[Character2] "I see a cave up ahead.
+Let's hurry before it gets dark."
+
+
+[Narrator] They rushed toward the cave, hearts pounding.
+```
+**Result**: Natural pauses between dialogue, emphasis pauses for dramatic effect, and smooth character transitions.
+
+### 📝 **Text Formatting Best Practices**
+
+#### **🎭 Multi-Voice Dialogue Structure**
+```
+[Character Name] Dialogue content here.
+
+[Another Character] Response content here.
+Multiple lines can be used for the same character.
+
+[Narrator] Descriptive text and scene setting.
+```
+
+#### **🎪 Natural Flow Techniques**
+- **Paragraph Breaks**: Use double line breaks for scene transitions
+- **Emphasis Pauses**: Add extra returns before important revelations
+- **Character Separation**: Single returns between different speakers
+- **Breathing Room**: Natural pauses for complex concepts or emotional moments
+
+#### **📖 Single Voice Formatting**
+```
+Chapter content flows naturally here.
+
+New paragraphs create natural pauses.
+
+
+Extended pauses can emphasize dramatic moments.
+
+Regular text continues with normal pacing.
+```
+
+### 🔄 **Processing Pipeline Features**
+
+#### **🧠 Intelligent Text Analysis**
+- **Line Break Preservation**: Maintains your formatting intentions throughout processing
+- **Character Assignment**: Automatically maps voice tags to selected voice profiles
+- **Chunk Optimization**: Breaks long texts into optimal segments while preserving pause timing
+- **Error Recovery**: Validates text and provides helpful formatting suggestions
+
+#### **⚡ Real-Time Processing**
+- **Live Feedback**: Console output shows exactly how many pauses are being added
+- **Debug Information**: Detailed logging of pause detection and application
+- **Progress Tracking**: Monitor pause processing alongside audio generation
+- **Quality Assurance**: Automatic validation of pause placement
+
+#### **🎚️ Professional Output**
+- **Seamless Integration**: Pauses blend naturally with generated speech
+- **Volume Consistency**: Silence segments match the audio output specifications
+- **Format Compatibility**: Works with all supported audio formats and quality settings
+- **Project Preservation**: Pause information saved in project metadata for regeneration
+
+### 💡 **Pro Tips for Better Audiobooks**
+
+#### **🎯 Dialogue Formatting**
+- **Character Consistency**: Always use the same character name format `[Name]`
+- **Natural Breaks**: Place returns where a human reader would naturally pause
+- **Scene Transitions**: Use multiple returns (2-3) for major scene changes
+- **Emotional Beats**: Add single returns before/after emotional dialogue
+
+#### **📚 Chapter Structure**
+```
+Chapter 1: The Beginning
+
+Opening paragraph with scene setting.
+
+"Character dialogue with natural flow."
+
+Descriptive narrative continues.
+
+
+Major scene transition with extended pause.
+
+New section begins here.
+```
+
+#### **🎪 Advanced Techniques**
+- **Cliffhangers**: Use extended pauses before revealing crucial information
+- **Action Sequences**: Shorter, punchy sentences with minimal pauses for intensity
+- **Contemplative Moments**: Longer pauses for reflection and character development
+- **Comedic Timing**: Strategic pauses before punchlines or comedic reveals
+
+### 🔍 **Debug Output Examples**
+
+When generating your audiobook, watch for these helpful console messages:
+```
+🔇 Detected 15 line breaks → 1.5s total pause time
+🔇 Line breaks detected in [Character1]: +0.3s pause (from 3 returns)
+🔇 Chunk 2 (Narrator): Added 0.2s pause after speech
+```
+
+This real-time feedback helps you understand exactly how your formatting translates to audio timing.
+
+---
+
+## 🆕 Recent Improvements
+
+### 🎯 **Audio Quality Enhancements**
+We've significantly improved audio generation quality by optimizing the underlying TTS parameters:
+
+- **Enhanced P-top and Minimum P Settings**: Fine-tuned probability parameters for more natural speech patterns
+- **Reduced Audio Artifacts**: Better handling of pronunciation and intonation
+- **Improved Voice Consistency**: More stable voice characteristics across long generations
+- **Better Pronunciation**: Enhanced handling of complex words and names
+
+**📝 Note for Existing Users**: 
+- Older voice profiles will continue to work as before
+- To take advantage of the new audio quality improvements, consider re-creating voice profiles
+- Existing projects remain fully compatible
+
+### 📋 **Text Queuing System**
+Perfect for processing large books or multiple chapters:
+
+- **Batch Upload**: Upload multiple text files of any size
+- **Sequential Processing**: Automatically processes files one after another
+- **Progress Tracking**: Monitor generation progress across all queued items
+- **Flexible Chapter Sizes**: No restrictions on individual file length
+- **Unattended Generation**: Set up large projects and let them run automatically
+
+### 🔄 **Chunk-Based TTS System**
+Enhanced the core text-to-speech engine for better reliability:
+
+- **Background Chunking**: Automatically splits long texts into optimal chunks
+- **Memory Management**: Better handling of large text inputs
+- **Error Recovery**: Improved resilience during long generation sessions
+- **Consistent Quality**: Maintains voice quality across chunk boundaries
+- **Progress Feedback**: Real-time updates on generation progress
+
+---
+
+## 🎚️ Volume Normalization Guide
+
+### **Individual Voice Setup**
+1. Go to **Voice Library** tab
+2. Upload your voice sample and configure settings
+3. Set target volume level (default: -18 dB for audiobooks)
+4. Choose from professional presets or use custom levels
+5. Save voice profile with volume settings
+
+### **Multi-Voice Projects**
+1. Navigate to **Multi-Voice Audiobook Creation** tab
+2. Enable volume normalization for all voices
+3. Set target level for consistent character voices
+4. All characters will be automatically normalized during generation
+
+### **Text Queuing Workflow** ⭐ *NEW*
+1. Go to **Production Studio** tab
+2. Select "Batch Processing" mode
+3. Upload multiple text files (chapters, sections, etc.)
+4. Choose your voice and settings
+5. Start batch processing - files will generate sequentially
+6. Monitor progress and download completed audiobooks
+
+### **Professional Standards**
+- **📖 Audiobook Standard**: -18 dB RMS (recommended for most audiobooks)
+- **🎙️ Podcast Standard**: -16 dB RMS (for podcast-style content)
+- **🔇 Quiet/Comfortable**: -20 dB RMS (for quiet listening environments)
+- **🔊 Loud/Energetic**: -14 dB RMS (for dynamic, energetic content)
+- **📺 Broadcast Standard**: -23 dB RMS (for broadcast television standards)
+
+---
+
+## 📁 Project Structure
+
+```
+📦 Your Audiobook Projects
+├── 🎤 speakers/           # Voice library and samples
+├── 📚 audiobook_projects/ # Generated audiobooks
+├── 🔧 src/audiobook/      # Core processing modules
+└── 📄 Generated files...  # Audio chunks and final outputs
+```
+
+---
+
+## 🎯 Workflow
+
+1. **📝 Prepare Text**: Format your story with proper chapter breaks and strategic line breaks for natural pauses
+2. **🎤 Select Voices**: Choose or clone voices for your characters  
+3. **🎚️ Configure Volume**: Set professional volume levels and normalization
+4. **⚙️ Configure Settings**: Adjust quality, speed, and processing options
+5. **🎧 Generate Audio**: Create your audiobook with advanced TTS and automatic pause insertion
+6. **🧹 Clean & Optimize**: Use smart cleanup tools for perfect audio
+7. **📦 Export**: Get your finished audiobook ready for distribution
+
+### 🎭 **Enhanced Multi-Voice Workflow**
+1. **📝 Format Dialogue**: Use `[Character]` tags and strategic line breaks for natural flow
+2. **🔇 Add Return Pauses**: Place line breaks where you want natural speech pauses (0.1s each)
+3. **🎤 Assign Voices**: Map each character to their voice profile
+4. **⚡ Process with Intelligence**: Watch console output for pause detection feedback
+5. **🎧 Review & Adjust**: Listen to generated audio and refine formatting if needed
+
+### 📋 **Batch Processing Workflow** ⭐ *NEW*
+1. **📚 Organize Chapters**: Split your book into individual text files
+2. **📋 Queue Setup**: Upload all files to the batch processing system
+3. **🎤 Voice Selection**: Choose voice and configure settings once
+4. **🔄 Automated Generation**: Let the system process all files sequentially
+5. **📊 Monitor Progress**: Track completion status in real-time
+6. **📦 Collect Results**: Download all generated audiobook chapters
+
+---
+
+## 🛠️ Technical Requirements
+
+- **Python 3.8+**
+- **CUDA GPU** (recommended for faster processing)
+- **8GB+ RAM** (16GB recommended for large projects)
+- **Modern web browser** for the interface
+
+### 🔧 **CUDA Support**
+- CUDA compatibility issues have been resolved with updated dependencies
+- GPU acceleration is now stable for extended generation sessions
+- Fallback to CPU processing available if CUDA issues occur
+- **If you encounter CUDA assertion errors**: Use the patched version from the installation instructions above
+- The fix addresses PyTorch indexing issues that could cause crashes during audio generation
+
+---
+
+## ⚠️ Known Issues & Compatibility
+
+### **Multi-Voice Generation**
+- Short sentences or sections may occasionally cause issues during multi-voice generation
+- This is a limitation of the underlying TTS models rather than the implementation
+- **Workaround**: Use longer, more detailed sentences for better stability
+- Single-voice generation is not affected by this issue
+
+### **Voice Profile Compatibility**
+- **Existing Voices**: All older voice profiles remain fully functional
+- **New Features**: To benefit from improved audio quality, consider re-creating voice profiles
+- **Project Compatibility**: Existing audiobook projects work without modification
+- **Regeneration**: Individual chunks can be regenerated with improved quality settings
+
+### **Batch Processing Considerations**
+- Large batch jobs may take significant time depending on text length and hardware
+- Monitor system resources during extended batch processing sessions
+- Consider processing very large books in smaller batches for better control
+
+---
+
+## 📋 Supported Formats
+
+### Input
+- **Text**: `.txt`, `.md`, formatted stories and scripts
+- **Audio Samples**: `.wav`, `.mp3`, `.flac` for voice cloning
+- **Batch Files**: Multiple text files for queue processing
+
+### Output
+- **Audio**: High-quality `.wav` files with professional volume levels
+- **Projects**: Organized folder structure with chapters
+- **Exports**: Ready-to-use audiobook files
+- **Batch Results**: Multiple completed audiobooks from queue processing
+
+---
+
+## 🆘 Support
+
+- **Features Guide**: See `AUDIOBOOK_FEATURES.md` for detailed capabilities
+- **Development Notes**: Check `development/` folder for technical details
+- **Issues**: Report problems via GitHub issues
+
+---
+
+## 📄 License
+
+This project is licensed under the terms specified in `LICENSE`.
+
+---
+
+**🎉 Ready to create amazing audiobooks with professional volume levels and enhanced audio quality? Run `./launch_audiobook.bat` and start generating!** 

VOICE_LIBRARY_ENHANCEMENT_COMPLETE.md

@@ -0,0 +1,99 @@
+# ✅ Voice Library Enhancement Complete
+
+## 🎯 **Problem Solved**
+The Voice Library UI was missing **advanced TTS parameters** (Min-P, Top-P, Repetition Penalty) that were available in the backend but not exposed to users.
+
+## 🛠️ **Changes Made**
+
+### 1. **Enhanced Voice Profile Storage** ⚙️
+- Updated `save_voice_profile()` function to accept and store:
+  - **Min-P** (default: 0.05) - Minimum probability threshold
+  - **Top-P** (default: 1.0) - Nucleus sampling threshold  
+  - **Repetition Penalty** (default: 1.2) - Token repetition control
+- Incremented version to **v2.1** for backward compatibility
+- Enhanced status messages to show advanced settings
+
+### 2. **Enhanced Voice Profile Loading** 📥
+- Updated `load_voice_profile()` function to return new parameters
+- Added backward compatibility - old voice profiles get sensible defaults
+- Enhanced status messages to show profile version
+
+### 3. **New Voice Library UI Controls** 🎛️
+Added **"Advanced Voice Parameters"** section in Voice Library tab:
+```
+🎛️ Advanced Voice Parameters
+├── Min-P (0.01-0.5) - "Minimum probability threshold for token selection (lower = more diverse)"
+├── Top-P (0.1-1.0) - "Nucleus sampling threshold (lower = more focused)"  
+└── Repetition Penalty (1.0-2.0) - "Penalty for repeating tokens (higher = less repetition)"
+```
+
+### 4. **Enhanced TTS Generation** 🎵
+- Updated core `generate()` function to accept new parameters
+- Updated `generate_with_cpu_fallback()` function for fallback mode
+- Updated `generate_with_retry()` function for robust generation
+- All TTS calls now use voice-specific advanced parameters
+
+### 5. **Enhanced Voice Configuration** 📋
+- Updated `get_voice_config()` function to include new parameters
+- All audiobook generation now uses saved voice settings
+- Backward compatibility maintained for existing voices
+
+### 6. **UI Integration** 🔗
+- **Save Button**: Now includes all 3 new parameters in voice profiles
+- **Load Button**: Populates all UI sliders with saved values
+- **Test Button**: Uses advanced parameters for voice testing
+
+## 🎮 **User Experience**
+
+### **Before** ❌
+- Only basic parameters: Exaggeration, CFG/Pace, Temperature
+- Advanced TTS controls were hidden and inaccessible
+- All voices used default Min-P/Top-P/Rep-Penalty values
+
+### **After** ✅  
+- **Full control** over TTS generation parameters
+- **Professional voice tuning** with industry-standard controls
+- **Per-voice customization** - each voice can have unique settings
+- **Backward compatibility** - existing voices continue working
+- **Enhanced voice testing** with all parameters
+
+## 📊 **Technical Benefits**
+
+### **Voice Quality Control** 🎭
+- **Min-P**: Fine-tune creativity vs consistency
+- **Top-P**: Control focus vs diversity in voice generation
+- **Repetition Penalty**: Eliminate unwanted voice repetitions
+
+### **Professional Workflow** 🎯
+- Voice artists can now fine-tune voices like professional TTS systems
+- Each character voice can have unique personality parameters
+- Better control over audiobook consistency and quality
+
+### **Future-Proof Architecture** 🚀
+- Versioned voice profiles (v2.1) support new features
+- Clean parameter passing through all generation functions  
+- Ready for additional TTS parameters in future updates
+
+## 🧪 **Testing Recommendations**
+
+1. **Create New Voice**: Test all advanced parameters
+2. **Load Old Voice**: Verify backward compatibility  
+3. **Generate Audio**: Confirm parameters affect output quality
+4. **Multi-Voice**: Test advanced parameters in character dialogue
+5. **Volume + Advanced**: Test combined normalization + advanced settings
+
+## ✨ **What Users See Now**
+
+When saving a voice, users get confirmation like:
+```
+✅ Voice profile 'Deep Male Narrator' saved successfully!
+📊 Audio normalized from -12.3 dB to -18.0 dB  
+🎛️ Advanced settings: Min-P=0.03, Top-P=0.9, Rep. Penalty=1.3
+```
+
+When loading a voice profile, version info is shown:
+```
+✅ Loaded voice profile: Deep Male Narrator (v2.1)
+```
+
+**The Voice Library now provides complete professional-grade TTS control!** 🎉 

install-audiobook.bat

@@ -0,0 +1,151 @@
+d@echo off
+echo ========================================
+echo   Chatterbox TTS - Installation Setup
+echo ========================================
+echo.
+echo This will install Chatterbox TTS in a virtual environment
+echo to keep it isolated from other Python projects.
+echo.
+echo Requirements:
+echo - Python 3.10 or higher
+echo - NVIDIA GPU with CUDA support (recommended)
+echo - Git (if you want to pull updates)
+echo.
+echo Current directory: %CD%
+echo.
+pause
+
+echo.
+echo [1/9] Checking Python installation...
+python --version
+if %errorlevel% neq 0 (
+    echo ERROR: Python is not installed or not in PATH
+    echo Please install Python 3.10+ from https://python.org
+    pause
+    exit /b 1
+)
+
+echo.
+echo [2/9] Checking if we're in the correct directory...
+if not exist "pyproject.toml" (
+    echo ERROR: pyproject.toml not found!
+    echo Please make sure you're running this from the chatterbox repository root.
+    echo Expected files: pyproject.toml, gradio_tts_app.py, src/chatterbox/
+    pause
+    exit /b 1
+)
+
+if not exist "src\chatterbox" (
+    echo ERROR: src\chatterbox directory not found!
+    echo Please make sure you're in the correct chatterbox repository.
+    pause
+    exit /b 1
+)
+
+echo Repository structure verified ✓
+
+echo.
+echo [3/9] Creating virtual environment...
+if exist "venv" (
+    echo Virtual environment already exists. Removing old one...
+    rmdir /s /q venv
+)
+python -m venv venv
+
+echo.
+echo [4/9] Activating virtual environment...
+call venv\Scripts\activate.bat
+
+echo.
+echo [5/9] Upgrading pip...
+python -m pip install --upgrade pip
+
+echo.
+echo [6/9] Installing compatible PyTorch with CUDA support...
+echo This may take a while (downloading ~2.5GB)...
+echo Installing PyTorch 2.4.1 + torchvision 0.19.1 (compatible versions)...
+pip install torch==2.4.1+cu121 torchvision==0.19.1+cu121 torchaudio==2.4.1+cu121 --index-url https://download.pytorch.org/whl/cu121
+
+echo.
+echo [7/9] Installing Chatterbox TTS and dependencies...
+pip install -e .
+pip install gradio
+
+echo.
+echo [8/9] Installing and configuring pydantic (tested version)...
+echo Uninstalling any existing pydantic versions...
+pip uninstall pydantic -y
+echo Installing pydantic version 2.10.6 (tested and verified)...
+pip install pydantic==2.10.6
+echo Verifying pydantic installation...
+pip show pydantic | findstr /C:"Version: 2.10.6"
+if %errorlevel% neq 0 (
+    echo WARNING: Pydantic 2.10.6 installation may have issues.
+    echo Attempting alternative installation...
+    pip install pydantic==2.10.6 --force-reinstall
+)
+
+echo Installing numpy (compatible version)...
+pip install numpy==1.26.0 --force-reinstall
+
+echo.
+echo [9/9] Testing installation...
+echo Testing PyTorch and CUDA...
+python -c "import torch; print('PyTorch version:', torch.__version__); print('CUDA available:', torch.cuda.is_available())"
+
+if %errorlevel% neq 0 (
+    echo WARNING: PyTorch test failed. Trying to fix torchvision compatibility...
+    pip uninstall torchvision -y
+    pip install torchvision==0.19.1+cu121 --index-url https://download.pytorch.org/whl/cu121 --force-reinstall
+    echo Retesting...
+    python -c "import torch; print('PyTorch version:', torch.__version__); print('CUDA available:', torch.cuda.is_available())"
+)
+
+echo.
+echo Testing Chatterbox import...
+python -c "from chatterbox.tts import ChatterboxTTS; print('Chatterbox TTS imported successfully!')"
+
+if %errorlevel% neq 0 (
+    echo WARNING: Chatterbox import failed. This might be a dependency issue.
+    echo The installation will continue, but you may need to troubleshoot.
+    echo Common fixes:
+    echo 1. Run install.bat again
+    echo 2. Check NVIDIA drivers are up to date
+    echo 3. Restart your computer
+)
+
+echo.
+echo Testing pydantic compatibility...
+python -c "import pydantic; print('Pydantic version:', pydantic.__version__)"
+
+echo.
+echo ========================================
+echo        Installation Complete!
+echo ========================================
+echo.
+echo Virtual environment created at: %CD%\venv
+echo.
+
+echo Final system check...
+python -c "import torch; print('GPU:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'None')"
+
+echo.
+echo ========================================
+echo           Ready for Audiobooks!
+echo ========================================
+echo.
+echo To start Chatterbox TTS:
+echo 1. Run launch_audiobook.bat (recommended)
+echo 2. Or manually: venv\Scripts\activate.bat then python gradio_tts_app_audiobook.py
+echo.
+echo Perfect for:
+echo - Voice cloning for audiobook narration
+echo - Multiple character voices
+echo - Consistent voice quality across chapters
+echo - Professional audiobook production
+echo.
+echo Note: If you encounter pydantic compatibility issues later,
+echo you can run update.bat to specifically update pydantic.
+echo.
+echo Installation finished successfully!
+pause 

launch_audiobook.bat

@@ -0,0 +1,52 @@
+@echo off
+setlocal
+
+rem Performance and Debugging Section
+rem =================================
+rem Enable CUDA_LAUNCH_BLOCKING for detailed error reports, but it hurts performance.
+rem set "CUDA_LAUNCH_BLOCKING=1"
+rem set "TORCH_USE_CUDA_DSA=1"
+
+echo Checking for virtual environment...
+if not exist "venv\Scripts\activate.bat" (
+    echo ERROR: Virtual environment not found!
+    echo Please run install.bat first to set up the environment.
+    echo.
+    echo Make sure you're in the chatterbox repository directory.
+    pause
+    exit /b 1
+)
+
+echo Checking repository structure...
+if not exist "gradio_tts_app_audiobook.py" (
+    echo ERROR: gradio_tts_app_audiobook.py not found!
+    echo Please make sure you're in the chatterbox repository root.
+    pause
+    exit /b 1
+)
+
+echo Activating virtual environment...
+call venv\Scripts\activate.bat
+
+echo.
+echo Starting Chatterbox TTS Audiobook Edition...
+echo Features: Voice Library, Character Management, Audiobook Tools
+echo Audio Cleaning Available in "Clean Samples" Tab
+echo This may take a moment to load the models...
+echo.
+echo Current directory: %CD%
+echo Python environment: %VIRTUAL_ENV%
+echo Voice library will be created at: %CD%\voice_library
+echo.
+
+python gradio_tts_app_audiobook.py
+
+echo.
+echo Chatterbox TTS Audiobook Edition has stopped.
+echo Deactivating virtual environment...
+deactivate
+echo.
+echo Thanks for using Chatterbox TTS Audiobook Edition! 🎧✨
+echo Your voice profiles are saved in the voice_library folder.
+echo Audio cleaning features are in the "Clean Samples" tab!
+pause 

pyproject.toml

@@ -0,0 +1,38 @@
+[project]
+name = "chatterbox-tts"
+version = "0.1.1"
+description = "Chatterbox: Open Source TTS and Voice Conversion by Resemble AI"
+readme = "README.md"
+requires-python = ">=3.8"
+license = {file = "LICENSE"}
+authors = [
+    {name = "resemble-ai", email = "[email protected]"}
+]
+dependencies = [
+    "numpy==1.26.0",
+    "resampy==0.4.3",
+    "librosa==0.10.0",
+    "s3tokenizer",
+    "torch==2.4.1",
+    "torchaudio==2.4.1",
+    "transformers==4.46.3",
+    "diffusers==0.29.0",
+    "resemble-perth==1.0.1",
+    "omegaconf==2.3.0",
+    "conformer==0.3.2",
+    "spacy>=3.4.0",
+]
+
+[project.optional-dependencies]
+advanced-nlp = ["spacy[en_core_web_sm]>=3.4.0"]
+
+[project.urls]
+Homepage = "https://github.com/resemble-ai/chatterbox"
+Repository = "https://github.com/resemble-ai/chatterbox"
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["src"]

simple_batch_demo.py

@@ -0,0 +1,146 @@
+"""
+Simple Batch Processing UI Demo for ChatterBox Audiobook
+
+This file shows exactly what needs to be added to the main interface
+to enable the batch processing functionality.
+"""
+
+import gradio as gr
+
+# Simulated functions (these already exist in your main file)
+def load_text_files_batch(file_paths):
+    if not file_paths:
+        return [], "No files uploaded"
+    return [{"filename": f"file_{i}.txt", "content": "sample", "words": 100} for i in range(len(file_paths))], f"Loaded {len(file_paths)} files"
+
+def validate_batch_audiobook_input(file_list, voice, project_name):
+    if not file_list:
+        return gr.Button(interactive=False), "❌ No files loaded", None
+    if not voice:
+        return gr.Button(interactive=False), "❌ Select a voice", None  
+    if not project_name:
+        return gr.Button(interactive=False), "❌ Enter project name", None
+    return gr.Button(interactive=True), f"✅ Ready to process {len(file_list)} files", None
+
+def create_batch_audiobook(model, file_list, voice_lib, voice, project_name, norm, level):
+    return None, f"✅ Batch processing complete! Created {len(file_list)} audiobooks with names {project_name}-1, {project_name}-2, etc."
+
+def demo_interface():
+    with gr.Blocks(title="Batch Processing Demo") as demo:
+        gr.HTML("""
+        <h1>🎵 Batch Processing Demo</h1>
+        <p>This shows the UI components that need to be added to your main interface.</p>
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                # Upload Mode Selection
+                upload_mode = gr.Radio(
+                    choices=[("Single File", "single"), ("Batch Processing", "batch")],
+                    value="single",
+                    label="📋 Upload Mode",
+                    info="Switch between single file and batch processing"
+                )
+                
+                # Single file upload (visible by default)
+                with gr.Group(visible=True) as single_group:
+                    single_file = gr.File(
+                        label="📄 Upload Single Text File",
+                        file_types=[".txt"],
+                        type="filepath"
+                    )
+                    single_status = gr.HTML("📄 Single file mode")
+                
+                # Batch file upload (hidden by default)  
+                with gr.Group(visible=False) as batch_group:
+                    batch_files = gr.File(
+                        label="📚 Upload Multiple Text Files",
+                        file_types=[".txt"],
+                        file_count="multiple",
+                        type="filepath"
+                    )
+                    load_batch_btn = gr.Button("📂 Load Batch Files")
+                    batch_status = gr.HTML("📚 Batch processing mode")
+                
+                # Voice and project settings
+                voice_dropdown = gr.Dropdown(
+                    choices=["Voice 1", "Voice 2", "Voice 3"],
+                    label="Select Voice",
+                    value=None
+                )
+                
+                project_name = gr.Textbox(
+                    label="Project Name",
+                    placeholder="my_audiobook"
+                )
+                
+                # Batch file list state
+                batch_file_list = gr.State([])
+                
+            with gr.Column():
+                # Processing buttons
+                validate_batch_btn = gr.Button("🔍 Validate Batch", variant="secondary")
+                process_batch_btn = gr.Button("🎵 Create Batch Audiobooks", variant="primary", interactive=False)
+                
+                # Status and output
+                processing_status = gr.HTML("Ready for batch processing")
+                output_audio = gr.Audio(label="Preview (last created audiobook)", visible=False)
+        
+        # Event handlers
+        def toggle_upload_mode(mode):
+            if mode == "single":
+                return gr.Group(visible=True), gr.Group(visible=False)
+            else:
+                return gr.Group(visible=False), gr.Group(visible=True)
+        
+        upload_mode.change(
+            fn=toggle_upload_mode,
+            inputs=[upload_mode],
+            outputs=[single_group, batch_group]
+        )
+        
+        load_batch_btn.click(
+            fn=load_text_files_batch,
+            inputs=[batch_files],
+            outputs=[batch_file_list, batch_status]
+        )
+        
+        validate_batch_btn.click(
+            fn=validate_batch_audiobook_input,
+            inputs=[batch_file_list, voice_dropdown, project_name],
+            outputs=[process_batch_btn, processing_status, gr.State()]
+        )
+        
+        process_batch_btn.click(
+            fn=create_batch_audiobook,
+            inputs=[gr.State(None), batch_file_list, gr.State(""), voice_dropdown, project_name, gr.State(True), gr.State(-18)],
+            outputs=[output_audio, processing_status]
+        )
+        
+        gr.HTML("""
+        <div style="margin-top: 20px; padding: 15px; border: 1px solid #ddd; border-radius: 5px;">
+            <h3>📋 To Add This to Your Main Interface:</h3>
+            <ol>
+                <li>Replace the simple file upload section with the Upload Mode selection</li>
+                <li>Add the single and batch upload groups</li>
+                <li>Add the batch processing buttons</li>
+                <li>Wire up the event handlers</li>
+                <li>Add the batch_file_list State component</li>
+            </ol>
+            <p><strong>Key Components Needed:</strong></p>
+            <ul>
+                <li>upload_mode (Radio)</li>
+                <li>single_upload_group and batch_upload_group (Group)</li>
+                <li>batch_files (File with file_count="multiple")</li>
+                <li>load_batch_btn (Button)</li>
+                <li>validate_batch_btn and process_batch_btn (Buttons)</li>
+                <li>batch_file_list (State)</li>
+            </ul>
+        </div>
+        """)
+    
+    return demo
+
+if __name__ == "__main__":
+    demo = demo_interface()
+    demo.launch() 

src/audiobook/__init__.py

@@ -0,0 +1,8 @@
+"""
+ChatterBox Audiobook Generator
+
+A modular audiobook creation system using TTS technology.
+"""
+
+__version__ = "1.0.0"
+__author__ = "ChatterBox Team" 

src/audiobook/audio_processing.py

@@ -0,0 +1,480 @@
+"""
+Audio processing utilities for audiobook generation.
+
+Handles audio saving, combining, trimming, and file operations.
+"""
+
+import os
+import wave
+import numpy as np
+from pathlib import Path
+from typing import List, Tuple, Optional, Any
+import tempfile
+import shutil
+
+# Optional audio processing imports
+try:
+    import librosa
+    import soundfile as sf
+    AUDIO_PROCESSING_AVAILABLE = True
+except ImportError:
+    AUDIO_PROCESSING_AVAILABLE = False
+
+
+def save_audio_chunks(
+    audio_chunks: List[np.ndarray], 
+    sample_rate: int, 
+    project_name: str, 
+    output_dir: str = "audiobook_projects"
+) -> List[str]:
+    """Save audio chunks as numbered WAV files.
+    
+    Args:
+        audio_chunks: List of audio arrays
+        sample_rate: Audio sample rate
+        project_name: Name of the project
+        output_dir: Output directory for projects
+        
+    Returns:
+        List of saved file paths
+    """
+    if not project_name.strip():
+        project_name = "untitled_audiobook"
+    
+    # Sanitize project name
+    safe_project_name = "".join(c for c in project_name if c.isalnum() or c in (' ', '-', '_')).rstrip()
+    safe_project_name = safe_project_name.replace(' ', '_')
+    
+    # Create output directory
+    project_dir = os.path.join(output_dir, safe_project_name)
+    os.makedirs(project_dir, exist_ok=True)
+    
+    saved_files = []
+    
+    for i, audio_chunk in enumerate(audio_chunks, 1):
+        filename = f"{safe_project_name}_{i:03d}.wav"
+        filepath = os.path.join(project_dir, filename)
+        
+        # Save as WAV file
+        with wave.open(filepath, 'wb') as wav_file:
+            wav_file.setnchannels(1)  # Mono
+            wav_file.setsampwidth(2)  # 16-bit
+            wav_file.setframerate(sample_rate)
+            
+            # Convert float32 to int16
+            audio_int16 = (audio_chunk * 32767).astype(np.int16)
+            wav_file.writeframes(audio_int16.tobytes())
+        
+        saved_files.append(filepath)
+    
+    return saved_files
+
+
+def combine_audio_files(file_paths: List[str], output_path: str, output_format: str = "wav") -> str:
+    """Combine multiple audio files into a single file.
+    
+    Args:
+        file_paths: List of audio file paths to combine
+        output_path: Output file path
+        output_format: Output format (wav or mp3)
+        
+    Returns:
+        Success message or error
+    """
+    if not file_paths:
+        return "❌ No audio files to combine"
+    
+    try:
+        # Read all audio files and combine
+        combined_audio = []
+        sample_rate = None
+        
+        for file_path in file_paths:
+            if not os.path.exists(file_path):
+                continue
+                
+            with wave.open(file_path, 'rb') as wav_file:
+                if sample_rate is None:
+                    sample_rate = wav_file.getframerate()
+                
+                frames = wav_file.readframes(wav_file.getnframes())
+                audio_data = np.frombuffer(frames, dtype=np.int16)
+                combined_audio.append(audio_data)
+        
+        if not combined_audio:
+            return "❌ No valid audio files found"
+        
+        # Concatenate all audio
+        final_audio = np.concatenate(combined_audio)
+        
+        # Save combined audio
+        if output_format.lower() == "wav":
+            with wave.open(output_path, 'wb') as wav_file:
+                wav_file.setnchannels(1)
+                wav_file.setsampwidth(2)
+                wav_file.setframerate(sample_rate)
+                wav_file.writeframes(final_audio.tobytes())
+        else:
+            # For MP3, we'd need additional dependencies like pydub
+            return "❌ MP3 export not implemented yet"
+        
+        return f"✅ Combined {len(file_paths)} files into {output_path}"
+        
+    except Exception as e:
+        return f"❌ Error combining audio files: {str(e)}"
+
+
+def save_trimmed_audio(audio_data: Any, original_file_path: str, chunk_num: int) -> Tuple[str, str]:
+    """Save trimmed audio data to a new file.
+    
+    Args:
+        audio_data: Audio data from Gradio component
+        original_file_path: Original audio file path
+        chunk_num: Chunk number
+        
+    Returns:
+        tuple: (success_message, file_path)
+    """
+    if audio_data is None:
+        return "❌ No audio data provided", ""
+    
+    try:
+        # Extract sample rate and audio array from gradio format
+        if isinstance(audio_data, tuple) and len(audio_data) == 2:
+            sample_rate, audio_array = audio_data
+        else:
+            return "❌ Invalid audio data format", ""
+        
+        # Create temporary file path
+        base_dir = os.path.dirname(original_file_path)
+        base_name = os.path.splitext(os.path.basename(original_file_path))[0]
+        trimmed_path = os.path.join(base_dir, f"{base_name}_trimmed.wav")
+        
+        # Save trimmed audio
+        with wave.open(trimmed_path, 'wb') as wav_file:
+            wav_file.setnchannels(1)
+            wav_file.setsampwidth(2)
+            wav_file.setframerate(sample_rate)
+            
+            # Convert to int16 if needed
+            if audio_array.dtype != np.int16:
+                if audio_array.dtype == np.float32 or audio_array.dtype == np.float64:
+                    audio_array = (audio_array * 32767).astype(np.int16)
+                else:
+                    audio_array = audio_array.astype(np.int16)
+            
+            wav_file.writeframes(audio_array.tobytes())
+        
+        return f"✅ Trimmed audio saved for chunk {chunk_num}", trimmed_path
+        
+    except Exception as e:
+        return f"❌ Error saving trimmed audio: {str(e)}", ""
+
+
+def extract_audio_segment(
+    audio_data: Any, 
+    start_time: Optional[float] = None, 
+    end_time: Optional[float] = None
+) -> Tuple[str, Any]:
+    """Extract a segment from audio data based on time stamps.
+    
+    Args:
+        audio_data: Audio data tuple (sample_rate, audio_array)
+        start_time: Start time in seconds
+        end_time: End time in seconds
+        
+    Returns:
+        tuple: (status_message, extracted_audio_data)
+    """
+    if audio_data is None:
+        return "❌ No audio data provided", None
+    
+    try:
+        if isinstance(audio_data, tuple) and len(audio_data) == 2:
+            sample_rate, audio_array = audio_data
+        else:
+            return "❌ Invalid audio data format", None
+        
+        if start_time is None and end_time is None:
+            return "❌ Please specify start time or end time", None
+        
+        # Convert time to sample indices
+        start_sample = int(start_time * sample_rate) if start_time is not None else 0
+        end_sample = int(end_time * sample_rate) if end_time is not None else len(audio_array)
+        
+        # Validate bounds
+        start_sample = max(0, start_sample)
+        end_sample = min(len(audio_array), end_sample)
+        
+        if start_sample >= end_sample:
+            return "❌ Invalid time range", None
+        
+        # Extract segment
+        extracted_audio = audio_array[start_sample:end_sample]
+        
+        return "✅ Audio segment extracted", (sample_rate, extracted_audio)
+        
+    except Exception as e:
+        return f"❌ Error extracting audio segment: {str(e)}", None
+
+
+def handle_audio_trimming(audio_data: Any) -> Tuple[str, Any]:
+    """Handle audio trimming from Gradio component.
+    
+    Args:
+        audio_data: Audio data from Gradio component
+        
+    Returns:
+        tuple: (status_message, processed_audio_data)
+    """
+    if audio_data is None:
+        return "No audio data", None
+    
+    try:
+        # Process audio data from Gradio
+        if isinstance(audio_data, tuple) and len(audio_data) == 2:
+            sample_rate, audio_array = audio_data
+            
+            # Validate audio array
+            if audio_array is None or len(audio_array) == 0:
+                return "Empty audio data", None
+            
+            return "Audio ready for processing", audio_data
+        else:
+            return "Invalid audio format", None
+            
+    except Exception as e:
+        return f"Error processing audio: {str(e)}", None
+
+
+def cleanup_temp_files(file_paths: List[str]) -> None:
+    """Clean up temporary files.
+    
+    Args:
+        file_paths: List of file paths to delete
+    """
+    for file_path in file_paths:
+        try:
+            if os.path.exists(file_path):
+                os.remove(file_path)
+        except Exception as e:
+            print(f"Warning: Could not delete {file_path}: {e}")
+
+
+def analyze_audio_quality(file_path: str) -> dict:
+    """Analyze audio file quality metrics.
+    
+    Args:
+        file_path: Path to audio file
+        
+    Returns:
+        Dictionary with quality metrics
+    """
+    try:
+        if AUDIO_PROCESSING_AVAILABLE:
+            # Use librosa for more detailed analysis
+            y, sr = librosa.load(file_path, sr=None)
+            
+            # Calculate advanced metrics
+            rms = np.sqrt(np.mean(y**2))
+            peak = np.max(np.abs(y))
+            duration = len(y) / sr
+            
+            # Calculate spectral centroid (brightness)
+            spectral_centroids = librosa.feature.spectral_centroid(y=y, sr=sr)[0]
+            spectral_centroid_mean = np.mean(spectral_centroids)
+            
+            # Calculate zero crossing rate (useful for speech analysis)
+            zcr = librosa.feature.zero_crossing_rate(y)[0]
+            zcr_mean = np.mean(zcr)
+            
+            return {
+                'duration': duration,
+                'sample_rate': sr,
+                'rms_level': float(rms),
+                'peak_level': float(peak),
+                'dynamic_range': float(peak / (rms + 1e-6)),
+                'spectral_centroid': float(spectral_centroid_mean),
+                'zero_crossing_rate': float(zcr_mean),
+                'has_advanced_analysis': True
+            }
+        else:
+            # Fallback to basic wave analysis
+            with wave.open(file_path, 'rb') as wav_file:
+                sample_rate = wav_file.getframerate()
+                n_frames = wav_file.getnframes()
+                duration = n_frames / sample_rate
+                
+                frames = wav_file.readframes(n_frames)
+                audio_data = np.frombuffer(frames, dtype=np.int16)
+                
+                # Normalize to float
+                audio_data = audio_data.astype(np.float32) / 32768.0
+                
+                # Calculate basic metrics
+                rms = np.sqrt(np.mean(audio_data**2))
+                peak = np.max(np.abs(audio_data))
+                
+                return {
+                    'duration': duration,
+                    'sample_rate': sample_rate,
+                    'rms_level': float(rms),
+                    'peak_level': float(peak),
+                    'dynamic_range': float(peak / (rms + 1e-6)),
+                    'has_advanced_analysis': False
+                }
+            
+    except Exception as e:
+        return {'error': str(e)}
+
+
+def auto_remove_silence(
+    file_path: str, 
+    silence_threshold: float = -50.0, 
+    min_silence_duration: float = 0.5
+) -> Tuple[str, str]:
+    """Automatically remove silence from audio file using advanced audio processing.
+    
+    Args:
+        file_path: Path to audio file
+        silence_threshold: Silence threshold in dB
+        min_silence_duration: Minimum silence duration to remove in seconds
+        
+    Returns:
+        tuple: (status_message, output_file_path)
+    """
+    if not AUDIO_PROCESSING_AVAILABLE:
+        # Fallback behavior - just copy the file with a warning
+        try:
+            output_path = file_path.replace('.wav', '_cleaned.wav')
+            shutil.copy2(file_path, output_path)
+            return "⚠️ Audio processing libraries not available. File copied without cleaning. Install librosa and soundfile for real audio processing.", output_path
+        except Exception as e:
+            return f"❌ Error copying file: {str(e)}", ""
+    
+    try:
+        # Load audio with librosa
+        y, sr = librosa.load(file_path, sr=None)
+        
+        if len(y) == 0:
+            return "❌ Audio file is empty", ""
+        
+        # Convert threshold from dB to amplitude
+        # silence_threshold is in dB (e.g., -50 dB)
+        threshold_amplitude = 10 ** (silence_threshold / 20)
+        
+        # Calculate frame length based on minimum silence duration
+        frame_length = int(min_silence_duration * sr)
+        hop_length = frame_length // 4  # 75% overlap
+        
+        # Calculate RMS energy for each frame
+        rms = librosa.feature.rms(y=y, frame_length=frame_length, hop_length=hop_length)[0]
+        
+        # Create time array for frames
+        times = librosa.frames_to_time(np.arange(len(rms)), sr=sr, hop_length=hop_length)
+        
+        # Find frames above threshold (non-silent)
+        non_silent_frames = rms > threshold_amplitude
+        
+        if not np.any(non_silent_frames):
+            return "❌ Entire audio file is below silence threshold", ""
+        
+        # Find continuous segments of non-silent audio
+        # Add padding to avoid cutting speech too close
+        padding_frames = max(1, int(0.1 * sr / hop_length))  # 100ms padding
+        
+        # Expand non-silent regions
+        expanded_mask = np.copy(non_silent_frames)
+        for i in range(len(non_silent_frames)):
+            if non_silent_frames[i]:
+                start_pad = max(0, i - padding_frames)
+                end_pad = min(len(expanded_mask), i + padding_frames + 1)
+                expanded_mask[start_pad:end_pad] = True
+        
+        # Convert frame indices back to sample indices
+        non_silent_samples = np.zeros(len(y), dtype=bool)
+        for i, is_voice in enumerate(expanded_mask):
+            if is_voice:
+                start_sample = int(times[i] * sr) if i < len(times) else len(y)
+                end_sample = int(times[i + 1] * sr) if i + 1 < len(times) else len(y)
+                start_sample = max(0, min(start_sample, len(y)))
+                end_sample = max(0, min(end_sample, len(y)))
+                non_silent_samples[start_sample:end_sample] = True
+        
+        # Extract non-silent audio
+        cleaned_audio = y[non_silent_samples]
+        
+        if len(cleaned_audio) == 0:
+            return "❌ No audio remaining after silence removal", ""
+        
+        # Save cleaned audio
+        output_path = file_path.replace('.wav', '_cleaned.wav')
+        sf.write(output_path, cleaned_audio, sr)
+        
+        # Calculate statistics
+        original_duration = len(y) / sr
+        cleaned_duration = len(cleaned_audio) / sr
+        removed_duration = original_duration - cleaned_duration
+        percentage_removed = (removed_duration / original_duration) * 100
+        
+        return (
+            f"✅ Silence removal completed! "
+            f"Removed {removed_duration:.2f}s ({percentage_removed:.1f}%) of silence. "
+            f"Final duration: {cleaned_duration:.2f}s",
+            output_path
+        )
+        
+    except Exception as e:
+        return f"❌ Error removing silence: {str(e)}", ""
+
+
+def normalize_audio_levels(
+    file_path: str,
+    target_lufs: float = -23.0,
+    peak_limit: float = -1.0
+) -> Tuple[str, str]:
+    """Normalize audio levels to broadcast standards.
+    
+    Args:
+        file_path: Path to audio file
+        target_lufs: Target loudness in LUFS (default: -23 for broadcast)
+        peak_limit: Peak limit in dB (default: -1.0)
+        
+    Returns:
+        tuple: (status_message, output_file_path)
+    """
+    if not AUDIO_PROCESSING_AVAILABLE:
+        return "❌ Audio processing libraries not available. Install librosa and soundfile.", ""
+    
+    try:
+        # Load audio
+        y, sr = librosa.load(file_path, sr=None)
+        
+        if len(y) == 0:
+            return "❌ Audio file is empty", ""
+        
+        # Simple peak normalization (more advanced LUFS would require pyloudnorm)
+        current_peak = np.max(np.abs(y))
+        target_peak = 10 ** (peak_limit / 20)  # Convert dB to linear
+        
+        if current_peak > 0:
+            # Normalize to target peak
+            normalized_audio = y * (target_peak / current_peak)
+        else:
+            normalized_audio = y
+        
+        # Save normalized audio
+        output_path = file_path.replace('.wav', '_normalized.wav')
+        sf.write(output_path, normalized_audio, sr)
+        
+        # Calculate gain applied
+        gain_db = 20 * np.log10(target_peak / current_peak) if current_peak > 0 else 0
+        
+        return (
+            f"✅ Audio normalized! Applied {gain_db:+.2f} dB gain. "
+            f"Peak level now at {peak_limit:.1f} dB.",
+            output_path
+        )
+        
+    except Exception as e:
+        return f"❌ Error normalizing audio: {str(e)}", "" 

src/audiobook/config.py

@@ -0,0 +1,72 @@
+"""
+Configuration management for the audiobook generator.
+
+Handles loading and saving of application configuration including voice library paths.
+"""
+
+import json
+import os
+from pathlib import Path
+
+
+# Default configuration values
+DEFAULT_VOICE_LIBRARY = "voice_library"
+CONFIG_FILE = "audiobook_config.json"
+
+
+def load_config() -> str:
+    """Load configuration including voice library path.
+    
+    Returns:
+        str: Path to the voice library directory
+    """
+    if os.path.exists(CONFIG_FILE):
+        try:
+            with open(CONFIG_FILE, 'r') as f:
+                config = json.load(f)
+            return config.get('voice_library_path', DEFAULT_VOICE_LIBRARY)
+        except Exception:
+            return DEFAULT_VOICE_LIBRARY
+    return DEFAULT_VOICE_LIBRARY
+
+
+def save_config(voice_library_path: str) -> str:
+    """Save configuration including voice library path.
+    
+    Args:
+        voice_library_path: Path to the voice library directory
+        
+    Returns:
+        str: Success or error message
+    """
+    config = {
+        'voice_library_path': voice_library_path,
+        'last_updated': str(Path().resolve())  # timestamp
+    }
+    try:
+        with open(CONFIG_FILE, 'w') as f:
+            json.dump(config, f, indent=2)
+        return f"✅ Configuration saved - Voice library path: {voice_library_path}"
+    except Exception as e:
+        return f"❌ Error saving configuration: {str(e)}"
+
+
+def update_voice_library_path(new_path: str) -> tuple[str, str]:
+    """Update the voice library path in configuration.
+    
+    Args:
+        new_path: New path to the voice library
+        
+    Returns:
+        tuple: (status_message, updated_path)
+    """
+    if not new_path.strip():
+        return "❌ Voice library path cannot be empty", ""
+    
+    # Create directory if it doesn't exist
+    try:
+        os.makedirs(new_path, exist_ok=True)
+        save_result = save_config(new_path)
+        return save_result, new_path
+    except Exception as e:
+        return f"❌ Error updating voice library path: {str(e)}", "" 

src/audiobook/models.py

@@ -0,0 +1,236 @@
+"""Model management and TTS operations for the audiobook generation system."""
+
+import torch
+import random
+import numpy as np
+from chatterbox.tts import ChatterboxTTS
+from typing import Any, Tuple, Optional
+
+
+# Global device setting - will be imported from main file
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MULTI_VOICE_DEVICE = "cpu"  # Force CPU for multi-voice processing
+
+
+def set_seed(seed: int) -> None:
+    """Set random seeds for reproducible generation.
+    
+    Args:
+        seed: Random seed value
+    """
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+
+
+def load_model() -> ChatterboxTTS:
+    """Load TTS model for the default device.
+    
+    Returns:
+        ChatterboxTTS: Loaded TTS model
+    """
+    model = ChatterboxTTS.from_pretrained(DEVICE)
+    return model
+
+
+def load_model_cpu() -> ChatterboxTTS:
+    """Load model specifically for CPU processing.
+    
+    Returns:
+        ChatterboxTTS: CPU-loaded TTS model
+    """
+    model = ChatterboxTTS.from_pretrained("cpu")
+    return model
+
+
+def clear_gpu_memory() -> None:
+    """Clear GPU memory cache to prevent CUDA errors."""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+
+def check_gpu_memory() -> str:
+    """Check current GPU memory usage.
+    
+    Returns:
+        str: GPU memory status information
+    """
+    if torch.cuda.is_available():
+        allocated = torch.cuda.memory_allocated()
+        cached = torch.cuda.memory_reserved()
+        return f"GPU Memory - Allocated: {allocated//1024//1024}MB, Cached: {cached//1024//1024}MB"
+    return "CUDA not available"
+
+
+def generate(
+    model: ChatterboxTTS, 
+    text: str, 
+    audio_prompt_path: str, 
+    exaggeration: float, 
+    temperature: float, 
+    seed_num: int, 
+    cfgw: float
+) -> Tuple[int, np.ndarray]:
+    """Generate audio from text using the TTS model.
+    
+    Args:
+        model: TTS model instance
+        text: Text to convert to speech
+        audio_prompt_path: Path to audio prompt file
+        exaggeration: Exaggeration parameter for generation
+        temperature: Temperature for generation randomness
+        seed_num: Random seed (0 for random)
+        cfgw: CFG weight parameter
+        
+    Returns:
+        tuple: (sample_rate, audio_array)
+    """
+    if model is None:
+        model = ChatterboxTTS.from_pretrained(DEVICE)
+
+    if seed_num != 0:
+        set_seed(int(seed_num))
+
+    wav = model.generate(
+        text,
+        audio_prompt_path=audio_prompt_path,
+        exaggeration=exaggeration,
+        temperature=temperature,
+        cfg_weight=cfgw,
+    )
+    return (model.sr, wav.squeeze(0).numpy())
+
+
+def generate_with_cpu_fallback(
+    model: ChatterboxTTS, 
+    text: str, 
+    audio_prompt_path: str, 
+    exaggeration: float, 
+    temperature: float, 
+    cfg_weight: float
+) -> Tuple[Any, str]:
+    """Generate audio with automatic CPU fallback for problematic CUDA errors.
+    
+    Args:
+        model: TTS model instance
+        text: Text to convert to speech
+        audio_prompt_path: Path to audio prompt file
+        exaggeration: Exaggeration parameter
+        temperature: Temperature parameter
+        cfg_weight: CFG weight parameter
+        
+    Returns:
+        tuple: (audio_wav, device_used)
+    """
+    # First try GPU if available
+    if DEVICE == "cuda":
+        try:
+            clear_gpu_memory()
+            wav = model.generate(
+                text,
+                audio_prompt_path=audio_prompt_path,
+                exaggeration=exaggeration,
+                temperature=temperature,
+                cfg_weight=cfg_weight,
+            )
+            return wav, "GPU"
+        except RuntimeError as e:
+            if ("srcIndex < srcSelectDimSize" in str(e) or 
+                "CUDA" in str(e) or 
+                "out of memory" in str(e).lower()):
+                
+                print(f"⚠️ CUDA error detected, falling back to CPU: {str(e)[:100]}...")
+                # Fall through to CPU mode
+            else:
+                raise e
+    
+    # CPU fallback or primary CPU mode
+    try:
+        # Load CPU model if needed
+        cpu_model = ChatterboxTTS.from_pretrained("cpu")
+        wav = cpu_model.generate(
+            text,
+            audio_prompt_path=audio_prompt_path,
+            exaggeration=exaggeration,
+            temperature=temperature,
+            cfg_weight=cfg_weight,
+        )
+        return wav, "CPU"
+    except Exception as e:
+        raise RuntimeError(f"Both GPU and CPU generation failed: {str(e)}")
+
+
+def generate_with_retry(
+    model: ChatterboxTTS, 
+    text: str, 
+    audio_prompt_path: str, 
+    exaggeration: float, 
+    temperature: float, 
+    cfg_weight: float, 
+    max_retries: int = 3
+) -> Tuple[Any, str]:
+    """Generate audio with retry mechanism for robustness.
+    
+    Args:
+        model: TTS model instance
+        text: Text to convert to speech
+        audio_prompt_path: Path to audio prompt file
+        exaggeration: Exaggeration parameter
+        temperature: Temperature parameter
+        cfg_weight: CFG weight parameter
+        max_retries: Maximum number of retry attempts
+        
+    Returns:
+        tuple: (audio_wav, device_used)
+    """
+    last_error = None
+    
+    for attempt in range(max_retries):
+        try:
+            return generate_with_cpu_fallback(
+                model, text, audio_prompt_path, exaggeration, temperature, cfg_weight
+            )
+        except Exception as e:
+            last_error = e
+            print(f"Attempt {attempt + 1}/{max_retries} failed: {str(e)}")
+            if attempt < max_retries - 1:
+                clear_gpu_memory()
+    
+    raise RuntimeError(f"All {max_retries} attempts failed. Last error: {last_error}")
+
+
+def force_cpu_processing() -> bool:
+    """Check if we should force CPU processing for stability.
+    
+    Returns:
+        bool: True if CPU processing should be forced
+    """
+    # For multi-voice, always use CPU to avoid CUDA indexing issues
+    return True
+
+
+def get_model_device_str(model_obj: Optional[ChatterboxTTS]) -> str:
+    """Get the device string for a model object.
+    
+    Args:
+        model_obj: TTS model instance
+        
+    Returns:
+        str: Device information string
+    """
+    if model_obj is None:
+        return "No model loaded"
+    
+    try:
+        # Try to access model device info
+        if hasattr(model_obj, 'device'):
+            return f"Model device: {model_obj.device}"
+        elif hasattr(model_obj, 'model') and hasattr(model_obj.model, 'device'):
+            return f"Model device: {model_obj.model.device}"
+        else:
+            return "Device info unavailable"
+    except Exception as e:
+        return f"Error getting device info: {str(e)}" 

src/audiobook/processing.py

@@ -0,0 +1,928 @@
+"""
+Text processing utilities for audiobook generation.
+
+Handles text chunking, validation, multi-voice parsing, and text cleanup.
+"""
+
+import re
+import os
+import wave
+import numpy as np
+from pathlib import Path
+from typing import List, Dict, Tuple, Any
+
+
+def chunk_text_by_sentences(text: str, max_words: int = 50) -> List[str]:
+    """Split text into chunks, breaking at sentence boundaries after reaching max_words.
+    
+    Args:
+        text: Input text to chunk
+        max_words: Maximum words per chunk
+        
+    Returns:
+        List of text chunks
+    """
+    # Split text into sentences using regex to handle multiple punctuation marks
+    sentences = re.split(r'([.!?]+\s*)', text)
+    
+    chunks = []
+    current_chunk = ""
+    current_word_count = 0
+    
+    i = 0
+    while i < len(sentences):
+        sentence = sentences[i].strip()
+        if not sentence:
+            i += 1
+            continue
+            
+        # Add punctuation if it exists
+        if i + 1 < len(sentences) and re.match(r'[.!?]+\s*', sentences[i + 1]):
+            sentence += sentences[i + 1]
+            i += 2
+        else:
+            i += 1
+        
+        sentence_words = len(sentence.split())
+        
+        # If adding this sentence would exceed max_words, start new chunk
+        if current_word_count > 0 and current_word_count + sentence_words > max_words:
+            if current_chunk.strip():
+                chunks.append(current_chunk.strip())
+            current_chunk = sentence
+            current_word_count = sentence_words
+        else:
+            current_chunk += " " + sentence if current_chunk else sentence
+            current_word_count += sentence_words
+    
+    # Add the last chunk if it exists
+    if current_chunk.strip():
+        chunks.append(current_chunk.strip())
+    
+    return chunks
+
+
+def adaptive_chunk_text(text: str, max_words: int = 50, reduce_on_error: bool = True) -> List[str]:
+    """Adaptively chunk text with error handling.
+    
+    Args:
+        text: Input text to chunk
+        max_words: Maximum words per chunk
+        reduce_on_error: Whether to reduce chunk size on errors
+        
+    Returns:
+        List of text chunks
+    """
+    return chunk_text_by_sentences(text, max_words)
+
+
+def load_text_file(file_path: str) -> Tuple[str, str]:
+    """Load text content from a file with encoding detection.
+    
+    Args:
+        file_path: Path to the text file
+        
+    Returns:
+        tuple: (text_content, status_message)
+    """
+    if not file_path:
+        return "", "No file selected"
+    
+    try:
+        # Try UTF-8 first
+        try:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+        except UnicodeDecodeError:
+            # Fallback to latin-1 for older files
+            with open(file_path, 'r', encoding='latin-1') as f:
+                content = f.read()
+        
+        if not content.strip():
+            return "", "File is empty"
+        
+        return content.strip(), f"✅ Loaded {len(content.split())} words from file"
+    
+    except FileNotFoundError:
+        return "", "❌ File not found"
+    except Exception as e:
+        return "", f"❌ Error reading file: {str(e)}"
+
+
+def validate_audiobook_input(text_content: str, selected_voice: str, project_name: str) -> Tuple[bool, str]:
+    """Validate input for single-voice audiobook creation.
+    
+    Args:
+        text_content: Text to validate
+        selected_voice: Selected voice name
+        project_name: Project name
+        
+    Returns:
+        tuple: (is_valid, error_message)
+    """
+    if not text_content or not text_content.strip():
+        return False, "❌ Please provide text content or upload a text file"
+    
+    if not selected_voice:
+        return False, "❌ Please select a voice"
+    
+    if not project_name or not project_name.strip():
+        return False, "❌ Please provide a project name"
+    
+    word_count = len(text_content.split())
+    if word_count < 10:
+        return False, "❌ Text content too short (minimum 10 words)"
+    
+    if word_count > 50000:
+        return False, "❌ Text content too long (maximum 50,000 words for performance)"
+    
+    return True, ""
+
+
+def parse_multi_voice_text(text: str) -> List[Dict[str, str]]:
+    """Parse text with multi-voice format markers.
+    
+    Expected format:
+    [CHARACTER_NAME] dialogue text (no colon needed, tag and dialogue can be on the same line)
+    
+    Args:
+        text: Input text with character markers
+        
+    Returns:
+        List of segments with character and text, e.g.,
+        [{'character': 'Character1', 'text': 'Dialogue for character 1.'}, ...]
+    """
+    segments = []
+    
+    # Regex to find [CharacterName] tags
+    # It captures the character name and the text that follows until the next tag or end of string
+    # Using re.split to capture text between tags and the tags themselves
+    parts = re.split(r'(\[[^\]]+\])', text)
+    
+    current_character = None
+    buffer = ""
+    
+    for part in parts:
+        if not part:
+            continue # Skip empty parts that can result from re.split
+            
+        part_stripped = part.strip()
+        if re.match(r'^\[[^\]]+\]$', part_stripped): # It's a character tag
+            if current_character and buffer.strip():
+                segments.append({
+                    'character': current_character,
+                    'text': buffer.strip()
+                })
+            current_character = part_stripped[1:-1] # Remove brackets
+            buffer = ""
+        else: # It's text content
+            if current_character is None and part_stripped: # Text before any character tag
+                # Assign to a default "Narrator" if no character tag precedes it.
+                # This can be adjusted based on desired behavior for untagged leading text.
+                segments.append({
+                    'character': "Narrator", # Or None, if untagged leading text should be handled differently
+                    'text': part_stripped
+                })
+                buffer = "" # Clear buffer as this part is processed
+            elif current_character:
+                buffer += part # Append to current character's text buffer
+            # If no current_character and it's not leading text, this part might be ignored
+            # or could be appended to a default narrator if strict tagging isn't enforced.
+            # Current logic: only appends if current_character is set.
+                
+    # Add any remaining text in the buffer for the last character
+    if current_character and buffer.strip():
+        segments.append({
+            'character': current_character,
+            'text': buffer.strip()
+        })
+    elif not current_character and buffer.strip() and not segments: # Only if it's the *only* content
+        # If the entire text has no tags, assign it all to Narrator
+        segments.append({
+            'character': "Narrator",
+            'text': buffer.strip()
+        })
+        
+    # Filter out any segments where the text is empty after stripping
+    final_segments = [seg for seg in segments if seg['text']]
+    
+    # Debug: Print parsed segments by the module
+    # print("Parsed Segments by text_processing.py module:", final_segments)
+    return final_segments
+
+
+def clean_character_name_from_text(text: str, voice_name: str) -> str:
+    """Clean character name markers from text.
+       The new parse_multi_voice_text in this module should handle name/dialogue separation.
+       This function primarily acts as a pass-through or for minor cleanup.
+    
+    Args:
+        text: Text that may contain character markers
+        voice_name: Voice name (largely ignored by this simplified version)
+        
+    Returns:
+        Cleaned text
+    """
+    # The parsing logic should have already separated the character name.
+    # This function just ensures the text is stripped.
+    return text.strip()
+
+
+def chunk_multi_voice_segments(segments: List[Dict[str, str]], max_words: int = 50) -> List[Dict[str, str]]:
+    """Chunk multi-voice segments while preserving character assignments.
+    
+    Args:
+        segments: List of character segments
+        max_words: Maximum words per chunk
+        
+    Returns:
+        List of chunked segments with character assignments
+    """
+    chunked_segments = []
+    
+    for segment in segments:
+        character = segment['character']
+        text = segment['text']
+        
+        # Chunk the text for this character
+        text_chunks = chunk_text_by_sentences(text, max_words)
+        
+        # Create segment for each chunk
+        for chunk in text_chunks:
+            chunked_segments.append({
+                'character': character,
+                'text': chunk
+            })
+    
+    return chunked_segments
+
+
+def validate_multi_voice_text(text_content: str, voice_library_path: str) -> Tuple[bool, str, List[str]]:
+    """Validate multi-voice text format and extract characters.
+    
+    Args:
+        text_content: Text to validate
+        voice_library_path: Path to voice library
+        
+    Returns:
+        tuple: (is_valid, error_message, character_list)
+    """
+    if not text_content or not text_content.strip():
+        return False, "❌ Please provide text content", []
+    
+    # Parse segments to extract characters
+    segments = parse_multi_voice_text(text_content)
+    
+    if not segments:
+        return False, "❌ No valid character segments found. Use format: [CHARACTER_NAME]: dialogue", []
+    
+    # Extract unique characters
+    characters = list(set(segment['character'] for segment in segments))
+    
+    if len(characters) < 2:
+        return False, "❌ Multi-voice requires at least 2 different characters", characters
+    
+    if len(characters) > 6:
+        return False, "❌ Too many characters (maximum 6 for performance)", characters
+    
+    # Check if we have enough text
+    total_words = sum(len(segment['text'].split()) for segment in segments)
+    if total_words < 20:
+        return False, "❌ Not enough text content (minimum 20 words)", characters
+    
+    return True, "", characters
+
+
+def validate_multi_audiobook_input(text_content: str, voice_library_path: str, project_name: str) -> Tuple[bool, str]:
+    """Validate input for multi-voice audiobook creation.
+    
+    Args:
+        text_content: Text to validate
+        voice_library_path: Path to voice library
+        project_name: Project name
+        
+    Returns:
+        tuple: (is_valid, error_message)
+    """
+    if not project_name or not project_name.strip():
+        return False, "❌ Please provide a project name"
+    
+    is_valid, error_msg, _ = validate_multi_voice_text(text_content, voice_library_path)
+    return is_valid, error_msg
+
+
+def analyze_multi_voice_text(text_content: str, voice_library_path: str) -> Tuple[bool, str, Dict[str, int]]:
+    """Analyze multi-voice text and return character statistics.
+    
+    Args:
+        text_content: Text to analyze
+        voice_library_path: Path to voice library
+        
+    Returns:
+        tuple: (is_valid, message, character_counts)
+    """
+    is_valid, error_msg, characters = validate_multi_voice_text(text_content, voice_library_path)
+    
+    if not is_valid:
+        return False, error_msg, {}
+    
+    # Parse segments and count words per character
+    segments = parse_multi_voice_text(text_content)
+    character_counts = {}
+    
+    for segment in segments:
+        character = segment['character']
+        word_count = len(segment['text'].split())
+        character_counts[character] = character_counts.get(character, 0) + word_count
+    
+    total_words = sum(character_counts.values())
+    message = f"✅ Found {len(characters)} characters with {total_words} total words"
+    
+    return True, message, character_counts
+
+
+def _filter_problematic_short_chunks(chunks: List[str], voice_assignments: Dict[str, str]) -> List[str]:
+    """Filter out problematic short chunks that might cause TTS issues.
+    
+    Args:
+        chunks: List of text chunks
+        voice_assignments: Character to voice mappings
+        
+    Returns:
+        Filtered list of chunks
+    """
+    filtered_chunks = []
+    min_length = 10  # Minimum character length
+    
+    for chunk in chunks:
+        # Skip very short chunks
+        if len(chunk.strip()) < min_length:
+            continue
+        
+        # Skip chunks that are just punctuation or whitespace
+        if not re.search(r'[a-zA-Z]', chunk):
+            continue
+        
+        filtered_chunks.append(chunk)
+    
+    return filtered_chunks
+
+
+# PHASE 4 REFACTOR: Adding audio processing functions to this module
+# Originally from gradio_tts_app_audiobook.py save_audio_chunks() function
+
+def save_audio_chunks(audio_chunks: List[np.ndarray], sample_rate: int, project_name: str, output_dir: str = "audiobook_projects") -> Tuple[List[str], str]:
+    """
+    Save audio chunks as numbered WAV files
+    
+    Args:
+        audio_chunks: List of audio numpy arrays
+        sample_rate: Sample rate for audio files
+        project_name: Name of the project
+        output_dir: Directory to save project files
+        
+    Returns:
+        tuple: (list of saved file paths, project directory path)
+    """
+    if not project_name.strip():
+        project_name = "untitled_audiobook"
+    
+    # Sanitize project name
+    safe_project_name = "".join(c for c in project_name if c.isalnum() or c in (' ', '-', '_')).rstrip()
+    safe_project_name = safe_project_name.replace(' ', '_')
+    
+    # Create output directory
+    project_dir = os.path.join(output_dir, safe_project_name)
+    os.makedirs(project_dir, exist_ok=True)
+    
+    saved_files = []
+    
+    for i, audio_chunk in enumerate(audio_chunks, 1):
+        filename = f"{safe_project_name}_{i:03d}.wav"
+        filepath = os.path.join(project_dir, filename)
+        
+        # Save as WAV file
+        with wave.open(filepath, 'wb') as wav_file:
+            wav_file.setnchannels(1)  # Mono
+            wav_file.setsampwidth(2)  # 16-bit
+            wav_file.setframerate(sample_rate)
+            
+            # Convert float32 to int16
+            audio_int16 = (audio_chunk * 32767).astype(np.int16)
+            wav_file.writeframes(audio_int16.tobytes())
+        
+        saved_files.append(filepath)
+    
+    return saved_files, project_dir
+
+
+# PHASE 4 REFACTOR: Adding extract_audio_segment function from gradio_tts_app_audiobook.py
+def extract_audio_segment(audio_data, start_time: float = None, end_time: float = None) -> tuple:
+    """Extract a segment from audio data.
+    
+    Args:
+        audio_data: Numpy array of audio data
+        start_time: Start time in seconds (None for beginning)
+        end_time: End time in seconds (None for end)
+        
+    Returns:
+        tuple: (status_message, extracted_audio_data)
+    """
+    try:
+        sample_rate = 24000  # Default sample rate
+        
+        if audio_data is None or len(audio_data) == 0:
+            return "❌ No audio data to extract from", None
+            
+        total_duration = len(audio_data) / sample_rate
+        
+        start_sample = int(start_time * sample_rate) if start_time else 0
+        end_sample = int(end_time * sample_rate) if end_time else len(audio_data)
+        
+        # Validate bounds
+        start_sample = max(0, min(start_sample, len(audio_data)))
+        end_sample = max(start_sample, min(end_sample, len(audio_data)))
+        
+        extracted_audio = audio_data[start_sample:end_sample]
+        
+        if len(extracted_audio) == 0:
+            return "❌ Invalid time range - no audio extracted", None
+            
+        extracted_duration = len(extracted_audio) / sample_rate
+        return f"✅ Extracted {extracted_duration:.2f}s of audio", extracted_audio
+        
+    except Exception as e:
+        return f"❌ Error extracting audio segment: {str(e)}", None
+
+
+def process_text_for_pauses(text: str, pause_duration: float = 0.1) -> tuple:
+    """Process text to count returns and calculate total pause time.
+    
+    Args:
+        text: Input text to process
+        pause_duration: Duration in seconds per line break (default 0.1)
+        
+    Returns:
+        tuple: (processed_text, return_count, total_pause_duration)
+    """
+    # Count line breaks (both \n and \r\n)
+    return_count = text.count('\n') + text.count('\r')
+    total_pause_duration = return_count * pause_duration
+    
+    # Clean up text for TTS (normalize line breaks but keep content)
+    processed_text = text.replace('\r\n', '\n').replace('\r', '\n')
+    # Replace multiple consecutive newlines with single space to avoid empty chunks
+    processed_text = re.sub(r'\n+', ' ', processed_text).strip()
+    
+    print(f"🔇 Detected {return_count} line breaks → {total_pause_duration:.1f}s total pause time")
+    
+    return processed_text, return_count, total_pause_duration
+
+
+def create_silence_audio(duration: float, sample_rate: int = 24000) -> np.ndarray:
+    """Create silence audio of specified duration.
+    
+    Args:
+        duration: Duration in seconds
+        sample_rate: Sample rate for the audio
+        
+    Returns:
+        numpy array of silence audio
+    """
+    num_samples = int(duration * sample_rate)
+    return np.zeros(num_samples, dtype=np.float32)
+
+
+def insert_pauses_between_chunks(audio_chunks: List[np.ndarray], 
+                                return_count: int, 
+                                sample_rate: int = 24000,
+                                pause_duration: float = 0.1) -> np.ndarray:
+    """Insert pauses between audio chunks based on return count.
+    
+    Args:
+        audio_chunks: List of audio chunk arrays
+        return_count: Number of returns detected in original text
+        sample_rate: Sample rate for audio
+        pause_duration: Duration per return in seconds
+        
+    Returns:
+        Combined audio with pauses inserted
+    """
+    if not audio_chunks:
+        return np.array([], dtype=np.float32)
+    
+    if return_count == 0:
+        # No pauses needed, just concatenate
+        return np.concatenate(audio_chunks)
+    
+    # Calculate how to distribute pauses
+    # For simplicity, we'll add all pause time at the end
+    # In a more sophisticated approach, we could distribute pauses throughout
+    total_pause_time = return_count * pause_duration
+    pause_audio = create_silence_audio(total_pause_time, sample_rate)
+    
+    print(f"🔇 Adding {total_pause_time:.1f}s pause ({return_count} returns × {pause_duration}s each)")
+    
+    # Concatenate audio chunks with pause at the end
+    combined_audio = np.concatenate(audio_chunks)
+    final_audio = np.concatenate([combined_audio, pause_audio])
+    
+    return final_audio
+
+
+def process_text_with_distributed_pauses(text: str, max_words: int = 50, 
+                                        pause_duration: float = 0.1) -> tuple:
+    """Process text and distribute pauses throughout chunks based on line breaks.
+    
+    Args:
+        text: Input text to process
+        max_words: Maximum words per chunk
+        pause_duration: Duration per line break in seconds
+        
+    Returns:
+        tuple: (chunks_with_pauses, total_return_count, total_pause_duration)
+    """
+    # First, process text to understand pause requirements
+    processed_text, return_count, total_pause_duration = process_text_for_pauses(text, pause_duration)
+    
+    # Split into lines to track where pauses should be
+    lines = text.split('\n')
+    chunks_with_pauses = []
+    
+    current_chunk = ""
+    current_word_count = 0
+    pauses_for_chunk = 0
+    
+    for i, line in enumerate(lines):
+        line = line.strip()
+        if not line:
+            pauses_for_chunk += 1  # Empty line counts as a pause
+            continue
+            
+        line_words = len(line.split())
+        
+        # If adding this line would exceed max_words, finalize current chunk
+        if current_word_count > 0 and current_word_count + line_words > max_words:
+            if current_chunk.strip():
+                chunks_with_pauses.append({
+                    'text': current_chunk.strip(),
+                    'pauses': pauses_for_chunk
+                })
+            current_chunk = line
+            current_word_count = line_words
+            pauses_for_chunk = 0
+        else:
+            current_chunk += " " + line if current_chunk else line
+            current_word_count += line_words
+        
+        # Add pause if not the last line
+        if i < len(lines) - 1:
+            pauses_for_chunk += 1
+    
+    # Add the last chunk if it exists
+    if current_chunk.strip():
+        chunks_with_pauses.append({
+            'text': current_chunk.strip(),
+            'pauses': pauses_for_chunk
+        })
+    
+    return chunks_with_pauses, return_count, total_pause_duration
+
+
+def map_line_breaks_to_chunks(original_text: str, chunks: List[str], pause_duration: float = 0.1) -> tuple:
+    """Map line breaks from original text to corresponding chunks.
+    
+    Args:
+        original_text: Original text with line breaks
+        chunks: List of text chunks created by sentence chunking
+        pause_duration: Duration per line break in seconds
+        
+    Returns:
+        tuple: (chunk_pause_map, total_pause_duration)
+            chunk_pause_map: Dict mapping chunk index to pause duration
+            total_pause_duration: Total pause time across all chunks
+    """
+    import re
+    
+    chunk_pause_map = {}
+    total_pause_duration = 0.0
+    
+    # Create a version of original text for matching (remove extra whitespace but keep structure)
+    normalized_original = re.sub(r'\s+', ' ', original_text.replace('\n', ' ')).strip()
+    
+    # Track position in original text
+    original_position = 0
+    
+    for chunk_idx, chunk in enumerate(chunks):
+        chunk_normalized = chunk.strip()
+        if not chunk_normalized:
+            continue
+            
+        # Find this chunk in the original text
+        chunk_start = normalized_original.find(chunk_normalized, original_position)
+        if chunk_start == -1:
+            # Fallback: try to find it without position constraint
+            chunk_start = normalized_original.find(chunk_normalized)
+        
+        if chunk_start == -1:
+            # Can't find chunk, no pauses for this one
+            continue
+            
+        chunk_end = chunk_start + len(chunk_normalized)
+        
+        # Count line breaks in the corresponding section of original text
+        # Map back to original text position
+        orig_text_section_start = 0
+        orig_text_section_end = len(original_text)
+        
+        # Find the corresponding section in original text
+        words_before = len(normalized_original[:chunk_start].split())
+        words_in_chunk = len(chunk_normalized.split())
+        
+        # Find the section in original text that corresponds to this chunk
+        original_words = original_text.split()
+        if words_before < len(original_words):
+            # Find the start position in original text
+            words_section = ' '.join(original_words[words_before:words_before + words_in_chunk])
+            section_start = original_text.find(words_section)
+            if section_start != -1:
+                section_end = section_start + len(words_section)
+                # Count line breaks in this section and the gap after it (until next chunk)
+                next_chunk_start = section_end
+                if chunk_idx < len(chunks) - 1:
+                    next_chunk_text = chunks[chunk_idx + 1].strip()
+                    next_chunk_pos = original_text.find(next_chunk_text, section_end)
+                    if next_chunk_pos != -1:
+                        next_chunk_start = next_chunk_pos
+                
+                # Count line breaks from end of current chunk to start of next chunk
+                gap_text = original_text[section_end:next_chunk_start]
+                line_breaks = gap_text.count('\n')
+                
+                if line_breaks > 0:
+                    pause_time = line_breaks * pause_duration
+                    chunk_pause_map[chunk_idx] = pause_time
+                    total_pause_duration += pause_time
+        
+        original_position = chunk_end
+    
+    return chunk_pause_map, total_pause_duration 
+
+
+def chunk_text_by_sentences_local(text, max_words=50):
+    """Local copy of sentence chunking to avoid circular imports."""
+    import re
+    
+    # Split into sentences using common sentence endings
+    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
+    
+    chunks = []
+    current_chunk = ""
+    current_word_count = 0
+    
+    for sentence in sentences:
+        if not sentence.strip():
+            continue
+            
+        sentence_words = len(sentence.split())
+        
+        # If adding this sentence would exceed max_words and we have content, start a new chunk
+        if current_word_count > 0 and current_word_count + sentence_words > max_words:
+            if current_chunk.strip():
+                chunks.append(current_chunk.strip())
+            current_chunk = sentence
+            current_word_count = sentence_words
+        else:
+            current_chunk += " " + sentence if current_chunk else sentence
+            current_word_count += sentence_words
+    
+    # Add the last chunk if it exists
+    if current_chunk.strip():
+        chunks.append(current_chunk.strip())
+    
+    return chunks
+
+def chunk_text_with_line_break_priority(text: str, max_words: int = 50, pause_duration: float = 0.1) -> tuple:
+    """Chunk text with line breaks taking priority over sentence breaks.
+    
+    This function first splits on line breaks, then applies sentence chunking
+    within each line break segment if needed.
+    
+    Args:
+        text: Input text with line breaks
+        max_words: Maximum words per chunk
+        pause_duration: Duration per line break in seconds
+        
+    Returns:
+        tuple: (chunks_with_pauses, total_pause_duration)
+            chunks_with_pauses: List of dicts with 'text' and 'pause_duration' keys
+            total_pause_duration: Total pause time across all chunks
+    """
+    import re
+    
+    chunks_with_pauses = []
+    total_pause_duration = 0.0
+    
+    # Split text by line breaks, keeping track of consecutive breaks
+    line_segments = re.split(r'(\n+)', text)
+    
+    for i, segment in enumerate(line_segments):
+        if not segment:
+            continue
+            
+        # Check if this segment is line breaks
+        if re.match(r'\n+', segment):
+            # Count the number of line breaks for pause calculation
+            line_break_count = segment.count('\n')
+            pause_time = line_break_count * pause_duration
+            
+            # Add pause to the previous chunk if it exists
+            if chunks_with_pauses:
+                chunks_with_pauses[-1]['pause_duration'] += pause_time
+                total_pause_duration += pause_time
+                print(f"🔇 Line breaks detected: +{pause_time:.1f}s pause (from {line_break_count} returns)")
+            continue
+        
+        # This is actual text content - chunk it by sentences if needed
+        text_content = segment.strip()
+        if not text_content:
+            continue
+            
+        # Apply sentence chunking to this segment
+        text_chunks = chunk_text_by_sentences_local(text_content, max_words)
+        
+        # Add these chunks with initial pause duration of 0
+        for chunk in text_chunks:
+            if chunk.strip():
+                chunks_with_pauses.append({
+                    'text': chunk.strip(),
+                    'pause_duration': 0.0
+                })
+    
+    return chunks_with_pauses, total_pause_duration 
+
+
+def parse_multi_voice_text_local(text):
+    """Local copy of multi-voice text parsing to avoid circular imports."""
+    import re
+    
+    # Pattern to match [CharacterName] at the beginning of lines
+    pattern = r'^\[([^\]]+)\]\s*(.*?)(?=^\[|\Z)'
+    matches = re.findall(pattern, text, re.MULTILINE | re.DOTALL)
+    
+    if not matches:
+        # If no voice tags found, treat as single narrator
+        return [("Narrator", text.strip())]
+    
+    segments = []
+    for character_name, content in matches:
+        # DON'T strip content to preserve line breaks for pause processing
+        # Only strip leading/trailing spaces, but preserve newlines
+        content = content.rstrip(' \t').lstrip(' \t')
+        if content:
+            segments.append((character_name.strip(), content))
+    
+    return segments
+
+def chunk_multi_voice_text_with_line_break_priority(text: str, max_words: int = 30, pause_duration: float = 0.1) -> tuple:
+    """Chunk multi-voice text with line breaks taking priority over sentence breaks.
+    
+    Args:
+        text: Input text with voice tags and line breaks
+        max_words: Maximum words per chunk
+        pause_duration: Duration per line break in seconds
+        
+    Returns:
+        tuple: (segments_with_pauses, total_pause_duration)
+            segments_with_pauses: List of dicts with 'voice', 'text', and 'pause_duration' keys
+            total_pause_duration: Total pause time across all segments
+    """
+    import re
+    
+    # Add debugging output for the input text
+    print(f"🔍 DEBUG: chunk_multi_voice_text_with_line_break_priority input:")
+    print(f"🔍 DEBUG: Input text length: {len(text)} characters")
+    print(f"🔍 DEBUG: Line breaks in input: {text.count(chr(10))} \\n chars, {text.count(chr(13))} \\r chars")
+    print(f"🔍 DEBUG: First 200 chars: {repr(text[:200])}")
+    
+    # NEW APPROACH: Process line breaks in the full text before voice parsing
+    # Split the entire text by voice segments while preserving line breaks
+    segments_with_pauses = []
+    total_pause_duration = 0.0
+    
+    # Find all voice segments with their positions, preserving everything in between
+    voice_pattern = r'(\[([^\]]+)\]\s*)'
+    split_parts = re.split(voice_pattern, text)
+    
+    print(f"🔍 DEBUG: Split text into {len(split_parts)} parts")
+    for i, part in enumerate(split_parts):
+        print(f"🔍 DEBUG: Part {i}: {repr(part[:50])}")
+    
+    current_voice = None
+    
+    i = 0
+    while i < len(split_parts):
+        part = split_parts[i]
+        
+        # Check if this part is a voice tag match
+        if i + 2 < len(split_parts) and re.match(r'\[([^\]]+)\]\s*', part):
+            # This is a voice tag, extract the voice name
+            current_voice = split_parts[i + 1]  # The captured voice name
+            print(f"🔍 DEBUG: Found voice tag: '{current_voice}'")
+            
+            # The content is in the next part after the voice tag and whitespace
+            content_part = split_parts[i + 2] if i + 2 < len(split_parts) else ""
+            
+            # Process the content with line break awareness
+            if content_part:
+                processed_segments = process_voice_content_with_line_breaks(
+                    current_voice, content_part, max_words, pause_duration
+                )
+                
+                for segment in processed_segments:
+                    segments_with_pauses.append(segment)
+                    total_pause_duration += segment['pause_duration']
+            
+            i += 3  # Skip voice tag, voice name, and content
+        else:
+            # This is content between voice tags or before first voice tag
+            if current_voice and part.strip():
+                # Content continuation for current voice
+                processed_segments = process_voice_content_with_line_breaks(
+                    current_voice, part, max_words, pause_duration
+                )
+                
+                for segment in processed_segments:
+                    segments_with_pauses.append(segment)
+                    total_pause_duration += segment['pause_duration']
+            elif not current_voice and part.strip():
+                # Content before any voice tag - treat as narrator
+                processed_segments = process_voice_content_with_line_breaks(
+                    "Narrator", part, max_words, pause_duration
+                )
+                
+                for segment in processed_segments:
+                    segments_with_pauses.append(segment)
+                    total_pause_duration += segment['pause_duration']
+            
+            i += 1
+    
+    print(f"🔍 DEBUG: Final result: {len(segments_with_pauses)} segments, {total_pause_duration:.1f}s total pause time")
+    
+    return segments_with_pauses, total_pause_duration
+
+
+def process_voice_content_with_line_breaks(voice_name: str, content: str, max_words: int, pause_duration: float) -> list:
+    """Process voice content while preserving line breaks for pauses."""
+    import re
+    
+    segments = []
+    
+    # Split content by line breaks, keeping the line breaks
+    line_segments = re.split(r'(\n+)', content)
+    
+    print(f"🔍 DEBUG: Processing voice '{voice_name}' content split into {len(line_segments)} line segments")
+    
+    for i, line_segment in enumerate(line_segments):
+        if not line_segment:
+            continue
+            
+        # Check if this segment is line breaks
+        if re.match(r'\n+', line_segment):
+            # Count the number of line breaks for pause calculation
+            line_break_count = line_segment.count('\n')
+            pause_time = line_break_count * pause_duration
+            
+            print(f"🔍 DEBUG: Found {line_break_count} line breaks, calculating {pause_time:.1f}s pause")
+            
+            # Add pause to the previous segment if it exists and has the same voice
+            if segments and segments[-1]['voice'] == voice_name:
+                segments[-1]['pause_duration'] += pause_time
+                print(f"🔇 Line breaks detected in [{voice_name}]: +{pause_time:.1f}s pause (from {line_break_count} returns)")
+            else:
+                print(f"🔍 DEBUG: No previous segment to add pause to, or voice mismatch")
+            continue
+        
+        # This is actual text content - chunk it by sentences if needed
+        text_content = line_segment.strip()
+        if not text_content:
+            continue
+            
+        print(f"🔍 DEBUG: Processing text content: '{text_content[:50]}...'")
+        
+        # Apply sentence chunking to this segment
+        text_chunks = chunk_text_by_sentences_local(text_content, max_words)
+        
+        print(f"🔍 DEBUG: chunk_text_by_sentences_local produced {len(text_chunks)} chunks")
+        
+        # Add these chunks with voice assignment and initial pause duration of 0
+        for chunk in text_chunks:
+            if chunk.strip():
+                segments.append({
+                    'voice': voice_name,
+                    'text': chunk.strip(),
+                    'pause_duration': 0.0
+                })
+                print(f"🔍 DEBUG: Added segment: voice='{voice_name}', text='{chunk.strip()[:30]}...', pause=0.0")
+    
+    return segments 

src/audiobook/project_management.py

@@ -0,0 +1,656 @@
+"""
+Project management utilities for audiobook generation.
+
+Handles project creation, loading, metadata, file organization, and project lifecycle.
+"""
+
+import os
+import json
+import shutil
+import time
+from pathlib import Path
+from typing import List, Dict, Tuple, Optional, Any
+from datetime import datetime
+
+from .text_processing import chunk_text_by_sentences, parse_multi_voice_text, chunk_multi_voice_segments
+from .audio_processing import save_audio_chunks, auto_remove_silence, normalize_audio_levels, analyze_audio_quality
+from .voice_management import load_voice_for_tts, get_voice_config
+from .models import generate_with_retry, load_model_cpu
+
+
+# Constants
+MAX_CHUNKS_FOR_INTERFACE = 100
+MAX_CHUNKS_FOR_AUTO_SAVE = 100
+
+
+def save_project_metadata(
+    project_dir: str,
+    project_name: str,
+    text_content: str,
+    voice_info: dict,
+    chunks: list,
+    project_type: str = "single_voice"
+) -> None:
+    """Save project metadata to JSON file.
+    
+    Args:
+        project_dir: Project directory path
+        project_name: Name of the project
+        text_content: Original text content
+        voice_info: Voice configuration information
+        chunks: List of text chunks
+        project_type: Type of project (single_voice or multi_voice)
+    """
+    metadata = {
+        'project_name': project_name,
+        'project_type': project_type,
+        'created_at': datetime.now().isoformat(),
+        'text_content': text_content,
+        'voice_info': voice_info,
+        'chunks': chunks,
+        'total_chunks': len(chunks),
+        'status': 'in_progress'
+    }
+    
+    metadata_path = os.path.join(project_dir, 'metadata.json')
+    with open(metadata_path, 'w', encoding='utf-8') as f:
+        json.dump(metadata, f, indent=2)
+
+
+def load_project_metadata(project_dir: str) -> dict:
+    """Load project metadata from JSON file.
+    
+    Args:
+        project_dir: Project directory path
+        
+    Returns:
+        Project metadata dictionary
+    """
+    metadata_path = os.path.join(project_dir, 'metadata.json')
+    if os.path.exists(metadata_path):
+        try:
+            with open(metadata_path, 'r', encoding='utf-8') as f:
+                return json.load(f)
+        except Exception as e:
+            print(f"Warning: Could not load metadata for {project_dir}: {e}")
+    return {}
+
+
+def get_existing_projects(output_dir: str = "audiobook_projects") -> List[Dict[str, Any]]:
+    """Get list of existing audiobook projects.
+    
+    Args:
+        output_dir: Directory containing projects
+        
+    Returns:
+        List of project information dictionaries
+    """
+    projects = []
+    
+    if not os.path.exists(output_dir):
+        return projects
+    
+    try:
+        for item in os.listdir(output_dir):
+            project_dir = os.path.join(output_dir, item)
+            if os.path.isdir(project_dir):
+                metadata = load_project_metadata(project_dir)
+                
+                if metadata:
+                    # Use metadata information
+                    project_info = {
+                        'name': metadata.get('project_name', item),
+                        'path': project_dir,
+                        'type': metadata.get('project_type', 'unknown'),
+                        'created_at': metadata.get('created_at', ''),
+                        'total_chunks': metadata.get('total_chunks', 0),
+                        'status': metadata.get('status', 'unknown')
+                    }
+                else:
+                    # Fallback to directory scanning
+                    audio_files = [f for f in os.listdir(project_dir) if f.endswith('.wav')]
+                    project_info = {
+                        'name': item,
+                        'path': project_dir,
+                        'type': 'legacy',
+                        'created_at': '',
+                        'total_chunks': len(audio_files),
+                        'status': 'completed' if audio_files else 'empty'
+                    }
+                
+                projects.append(project_info)
+                
+    except Exception as e:
+        print(f"Warning: Error scanning projects directory: {e}")
+    
+    # Sort by creation date (newest first)
+    def get_sort_key(project):
+        created_at = project.get('created_at', '')
+        if created_at:
+            try:
+                return datetime.fromisoformat(created_at)
+            except:
+                pass
+        return datetime.min
+    
+    projects.sort(key=get_sort_key, reverse=True)
+    return projects
+
+
+def get_project_choices() -> List[str]:
+    """Get project names for UI dropdowns.
+    
+    Returns:
+        List of project names
+    """
+    projects = get_existing_projects()
+    if not projects:
+        return ["No projects found"]
+    
+    # Format: "project_name (type - chunks)"
+    choices = []
+    for project in projects:
+        name = project['name']
+        project_type = project['type']
+        chunk_count = project['total_chunks']
+        formatted = f"{name} ({project_type} - {chunk_count} chunks)"
+        choices.append(formatted)
+    
+    return choices
+
+
+def load_project_for_regeneration(project_name: str) -> Tuple[str, str, str, str]:
+    """Load project data for regeneration interface.
+    
+    Args:
+        project_name: Name of the project to load
+        
+    Returns:
+        tuple: (text_content, voice_name, project_type, status_message)
+    """
+    if not project_name or project_name == "No projects found":
+        return "", "", "", "No project selected"
+    
+    # Extract actual project name from formatted string
+    actual_name = project_name.split(' (')[0] if ' (' in project_name else project_name
+    
+    projects = get_existing_projects()
+    project_info = None
+    
+    for project in projects:
+        if project['name'] == actual_name:
+            project_info = project
+            break
+    
+    if not project_info:
+        return "", "", "", f"❌ Project '{actual_name}' not found"
+    
+    # Load project metadata
+    metadata = load_project_metadata(project_info['path'])
+    
+    if not metadata:
+        return "", "", "", f"❌ Could not load project metadata for '{actual_name}'"
+    
+    text_content = metadata.get('text_content', '')
+    voice_info = metadata.get('voice_info', {})
+    project_type = metadata.get('project_type', 'single_voice')
+    
+    # Extract voice name based on project type
+    if project_type == 'single_voice':
+        voice_name = voice_info.get('voice_name', '')
+    else:
+        voice_name = 'Multi-voice project'
+    
+    return text_content, voice_name, project_type, f"✅ Loaded project '{actual_name}'"
+
+
+def create_audiobook(
+    model: Any,
+    text_content: str,
+    voice_library_path: str,
+    selected_voice: str,
+    project_name: str,
+    resume: bool = False,
+    autosave_interval: int = 10
+) -> Tuple[str, List[str], str]:
+    """Create a single-voice audiobook project.
+    
+    Args:
+        model: TTS model instance
+        text_content: Text to convert to audio
+        voice_library_path: Path to voice library
+        selected_voice: Name of selected voice
+        project_name: Name for the project
+        resume: Whether to resume existing project
+        autosave_interval: Chunks between auto-saves
+        
+    Returns:
+        tuple: (status_message, audio_file_paths, project_path)
+    """
+    if not model:
+        model = load_model_cpu()
+    
+    # Load voice configuration
+    audio_prompt_path, voice_config = load_voice_for_tts(voice_library_path, selected_voice)
+    
+    if not audio_prompt_path:
+        return f"❌ Could not load voice '{selected_voice}'", [], ""
+    
+    # Get voice parameters
+    exaggeration = voice_config.get('exaggeration', 1.0)
+    temperature = voice_config.get('temperature', 0.7)
+    cfg_weight = voice_config.get('cfg_weight', 1.0)
+    
+    # Chunk the text
+    chunks = chunk_text_by_sentences(text_content, max_words=50)
+    
+    # Create project directory
+    safe_project_name = "".join(c for c in project_name if c.isalnum() or c in (' ', '-', '_')).strip()
+    safe_project_name = safe_project_name.replace(' ', '_')
+    project_dir = os.path.join("audiobook_projects", safe_project_name)
+    os.makedirs(project_dir, exist_ok=True)
+    
+    # Save project metadata
+    voice_info = {
+        'voice_name': selected_voice,
+        'audio_prompt_path': audio_prompt_path,
+        'exaggeration': exaggeration,
+        'temperature': temperature,
+        'cfg_weight': cfg_weight
+    }
+    
+    save_project_metadata(project_dir, project_name, text_content, voice_info, chunks, "single_voice")
+    
+    # Generate audio for chunks
+    audio_chunks = []
+    generated_files = []
+    
+    try:
+        for i, chunk in enumerate(chunks):
+            print(f"Generating chunk {i+1}/{len(chunks)}")
+            
+            # Generate audio
+            wav, device_used = generate_with_retry(
+                model, chunk, audio_prompt_path, exaggeration, temperature, cfg_weight
+            )
+            
+            # Convert to numpy array if needed
+            if hasattr(wav, 'squeeze'):
+                audio_array = wav.squeeze(0).numpy()
+            else:
+                audio_array = wav
+            
+            audio_chunks.append(audio_array)
+            
+            # Auto-save periodically
+            if (i + 1) % autosave_interval == 0 or i == len(chunks) - 1:
+                # Save current batch
+                batch_files = save_audio_chunks(
+                    audio_chunks, model.sr, safe_project_name, "audiobook_projects"
+                )
+                generated_files.extend(batch_files)
+                audio_chunks = []  # Reset for next batch
+        
+        # Update metadata to completed
+        metadata = load_project_metadata(project_dir)
+        metadata['status'] = 'completed'
+        metadata['completed_at'] = datetime.now().isoformat()
+        
+        metadata_path = os.path.join(project_dir, 'metadata.json')
+        with open(metadata_path, 'w', encoding='utf-8') as f:
+            json.dump(metadata, f, indent=2)
+        
+        return f"✅ Audiobook '{project_name}' created successfully! Generated {len(chunks)} audio chunks.", generated_files, project_dir
+        
+    except Exception as e:
+        return f"❌ Error creating audiobook: {str(e)}", generated_files, project_dir
+
+
+def create_multi_voice_audiobook_with_assignments(
+    model: Any,
+    text_content: str,
+    voice_library_path: str,
+    project_name: str,
+    voice_assignments: Dict[str, str],
+    resume: bool = False,
+    autosave_interval: int = 10
+) -> Tuple[str, List[str], str]:
+    """Create a multi-voice audiobook project with character voice assignments.
+    
+    Args:
+        model: TTS model instance
+        text_content: Text with character markers
+        voice_library_path: Path to voice library
+        project_name: Name for the project
+        voice_assignments: Character to voice mappings
+        resume: Whether to resume existing project
+        autosave_interval: Chunks between auto-saves
+        
+    Returns:
+        tuple: (status_message, audio_file_paths, project_path)
+    """
+    if not model:
+        model = load_model_cpu()
+    
+    # Parse multi-voice text
+    segments = parse_multi_voice_text(text_content)
+    chunked_segments = chunk_multi_voice_segments(segments, max_words=50)
+    
+    # Create project directory
+    safe_project_name = "".join(c for c in project_name if c.isalnum() or c in (' ', '-', '_')).strip()
+    safe_project_name = safe_project_name.replace(' ', '_')
+    project_dir = os.path.join("audiobook_projects", safe_project_name)
+    os.makedirs(project_dir, exist_ok=True)
+    
+    # Save project metadata
+    voice_info = {
+        'voice_assignments': voice_assignments,
+        'characters': list(voice_assignments.keys())
+    }
+    
+    save_project_metadata(project_dir, project_name, text_content, voice_info, chunked_segments, "multi_voice")
+    
+    # Generate audio for segments
+    audio_chunks = []
+    generated_files = []
+    
+    try:
+        for i, segment in enumerate(chunked_segments):
+            character = segment['character']
+            text = segment['text']
+            
+            # Get assigned voice for character
+            assigned_voice = voice_assignments.get(character)
+            if not assigned_voice:
+                print(f"Warning: No voice assigned for character '{character}', skipping segment")
+                continue
+            
+            # Load voice configuration
+            audio_prompt_path, voice_config = load_voice_for_tts(voice_library_path, assigned_voice)
+            
+            if not audio_prompt_path:
+                print(f"Warning: Could not load voice '{assigned_voice}' for character '{character}'")
+                continue
+            
+            print(f"Generating segment {i+1}/{len(chunked_segments)} - {character}: {text[:50]}...")
+            
+            # Generate audio
+            wav, device_used = generate_with_retry(
+                model, text, audio_prompt_path,
+                voice_config.get('exaggeration', 1.0),
+                voice_config.get('temperature', 0.7),
+                voice_config.get('cfg_weight', 1.0)
+            )
+            
+            # Convert to numpy array if needed
+            if hasattr(wav, 'squeeze'):
+                audio_array = wav.squeeze(0).numpy()
+            else:
+                audio_array = wav
+            
+            audio_chunks.append(audio_array)
+            
+            # Auto-save periodically
+            if (i + 1) % autosave_interval == 0 or i == len(chunked_segments) - 1:
+                # Save current batch
+                batch_files = save_audio_chunks(
+                    audio_chunks, model.sr, safe_project_name, "audiobook_projects"
+                )
+                generated_files.extend(batch_files)
+                audio_chunks = []  # Reset for next batch
+        
+        # Update metadata to completed
+        metadata = load_project_metadata(project_dir)
+        metadata['status'] = 'completed'
+        metadata['completed_at'] = datetime.now().isoformat()
+        
+        metadata_path = os.path.join(project_dir, 'metadata.json')
+        with open(metadata_path, 'w', encoding='utf-8') as f:
+            json.dump(metadata, f, indent=2)
+        
+        return f"✅ Multi-voice audiobook '{project_name}' created successfully! Generated {len(chunked_segments)} audio segments.", generated_files, project_dir
+        
+    except Exception as e:
+        return f"❌ Error creating multi-voice audiobook: {str(e)}", generated_files, project_dir
+
+
+def cleanup_project_temp_files(project_name: str) -> str:
+    """Clean up temporary files for a project.
+    
+    Args:
+        project_name: Name of the project
+        
+    Returns:
+        Status message
+    """
+    if not project_name:
+        return "❌ No project specified"
+    
+    # Extract actual project name
+    actual_name = project_name.split(' (')[0] if ' (' in project_name else project_name
+    safe_name = "".join(c for c in actual_name if c.isalnum() or c in (' ', '-', '_')).strip()
+    safe_name = safe_name.replace(' ', '_')
+    
+    project_dir = os.path.join("audiobook_projects", safe_name)
+    
+    if not os.path.exists(project_dir):
+        return f"❌ Project directory not found: {safe_name}"
+    
+    try:
+        temp_files = []
+        for file in os.listdir(project_dir):
+            if 'temp' in file.lower() or 'trimmed' in file.lower():
+                temp_files.append(os.path.join(project_dir, file))
+        
+        for temp_file in temp_files:
+            if os.path.exists(temp_file):
+                os.remove(temp_file)
+        
+        return f"✅ Cleaned up {len(temp_files)} temporary files for project '{actual_name}'"
+        
+    except Exception as e:
+        return f"❌ Error cleaning up project files: {str(e)}"
+
+
+def auto_clean_project_audio(
+    project_name: str,
+    silence_threshold: float = -50.0,
+    min_silence_duration: float = 0.5
+) -> str:
+    """Automatically clean audio for all chunks in a project.
+    
+    Args:
+        project_name: Name of the project
+        silence_threshold: Silence threshold in dB
+        min_silence_duration: Minimum silence duration to remove
+        
+    Returns:
+        Status message
+    """
+    if not project_name:
+        return "❌ No project specified"
+    
+    # Extract actual project name
+    actual_name = project_name.split(' (')[0] if ' (' in project_name else project_name
+    safe_name = "".join(c for c in actual_name if c.isalnum() or c in (' ', '-', '_')).strip()
+    safe_name = safe_name.replace(' ', '_')
+    
+    project_dir = os.path.join("audiobook_projects", safe_name)
+    
+    if not os.path.exists(project_dir):
+        return f"❌ Project directory not found: {safe_name}"
+    
+    try:
+        # Get all WAV files in the project
+        audio_files = [f for f in os.listdir(project_dir) 
+                      if f.endswith('.wav') and not 'cleaned' in f.lower() and not 'temp' in f.lower()]
+        
+        if not audio_files:
+            return f"❌ No audio files found in project '{actual_name}'"
+        
+        cleaned_count = 0
+        failed_count = 0
+        total_time_saved = 0.0
+        
+        for audio_file in audio_files:
+            file_path = os.path.join(project_dir, audio_file)
+            
+            # Clean the audio
+            status_msg, cleaned_path = auto_remove_silence(
+                file_path, silence_threshold, min_silence_duration
+            )
+            
+            if "✅" in status_msg:
+                cleaned_count += 1
+                # Extract time saved from status message
+                if "Removed" in status_msg:
+                    try:
+                        # Parse "Removed X.XXs" from status message
+                        import re
+                        match = re.search(r'Removed (\d+\.?\d*)s', status_msg)
+                        if match:
+                            total_time_saved += float(match.group(1))
+                    except:
+                        pass
+            else:
+                failed_count += 1
+                print(f"Failed to clean {audio_file}: {status_msg}")
+        
+        if cleaned_count > 0:
+            return (
+                f"✅ Auto-cleaned {cleaned_count}/{len(audio_files)} audio files for project '{actual_name}'. "
+                f"Total silence removed: {total_time_saved:.2f}s. "
+                f"Failed: {failed_count}"
+            )
+        else:
+            return f"❌ Failed to clean any audio files for project '{actual_name}'"
+        
+    except Exception as e:
+        return f"❌ Error auto-cleaning project audio: {str(e)}"
+
+
+def analyze_project_audio_quality(project_name: str) -> str:
+    """Analyze audio quality for all chunks in a project.
+    
+    Args:
+        project_name: Name of the project
+        
+    Returns:
+        Analysis results
+    """
+    if not project_name:
+        return "❌ No project specified"
+    
+    # Extract actual project name
+    actual_name = project_name.split(' (')[0] if ' (' in project_name else project_name
+    safe_name = "".join(c for c in actual_name if c.isalnum() or c in (' ', '-', '_')).strip()
+    safe_name = safe_name.replace(' ', '_')
+    
+    project_dir = os.path.join("audiobook_projects", safe_name)
+    
+    if not os.path.exists(project_dir):
+        return f"❌ Project directory not found: {safe_name}"
+    
+    try:
+        # Get all WAV files in the project
+        audio_files = [f for f in os.listdir(project_dir) 
+                      if f.endswith('.wav') and not 'temp' in f.lower()]
+        
+        if not audio_files:
+            return f"❌ No audio files found in project '{actual_name}'"
+        
+        total_duration = 0.0
+        total_rms = 0.0
+        peak_levels = []
+        analyzed_count = 0
+        
+        for audio_file in audio_files:
+            file_path = os.path.join(project_dir, audio_file)
+            
+            # Analyze the audio
+            metrics = analyze_audio_quality(file_path)
+            
+            if 'error' not in metrics:
+                total_duration += metrics.get('duration', 0)
+                total_rms += metrics.get('rms_level', 0)
+                peak_levels.append(metrics.get('peak_level', 0))
+                analyzed_count += 1
+        
+        if analyzed_count > 0:
+            avg_rms = total_rms / analyzed_count
+            max_peak = max(peak_levels) if peak_levels else 0
+            avg_peak = sum(peak_levels) / len(peak_levels) if peak_levels else 0
+            
+            # Convert to dB
+            avg_rms_db = 20 * np.log10(avg_rms) if avg_rms > 0 else -np.inf
+            max_peak_db = 20 * np.log10(max_peak) if max_peak > 0 else -np.inf
+            avg_peak_db = 20 * np.log10(avg_peak) if avg_peak > 0 else -np.inf
+            
+            return (
+                f"📊 Audio Quality Analysis for '{actual_name}':\n"
+                f"• Files analyzed: {analyzed_count}/{len(audio_files)}\n"
+                f"• Total duration: {total_duration:.2f} seconds\n"
+                f"• Average RMS level: {avg_rms_db:.1f} dB\n"
+                f"• Average peak level: {avg_peak_db:.1f} dB\n"
+                f"• Maximum peak level: {max_peak_db:.1f} dB\n"
+                f"• Recommended: Keep peaks below -3 dB for headroom"
+            )
+        else:
+            return f"❌ Failed to analyze any audio files for project '{actual_name}'"
+        
+    except Exception as e:
+        return f"❌ Error analyzing project audio quality: {str(e)}"
+
+
+def get_project_chunks(project_name: str) -> List[Dict[str, Any]]:
+    """Get list of audio chunks for a project.
+    
+    Args:
+        project_name: Name of the project
+        
+    Returns:
+        List of chunk information dictionaries
+    """
+    if not project_name or project_name == "No projects found":
+        return []
+    
+    # Extract actual project name
+    actual_name = project_name.split(' (')[0] if ' (' in project_name else project_name
+    safe_name = "".join(c for c in actual_name if c.isalnum() or c in (' ', '-', '_')).strip()
+    safe_name = safe_name.replace(' ', '_')
+    
+    project_dir = os.path.join("audiobook_projects", safe_name)
+    
+    if not os.path.exists(project_dir):
+        return []
+    
+    try:
+        chunks = []
+        audio_files = [f for f in os.listdir(project_dir) if f.endswith('.wav') and not 'temp' in f.lower()]
+        
+        # Sort files by chunk number
+        def extract_chunk_num_from_filename(filename: str) -> int:
+            # Extract number from filename like "project_001.wav"
+            parts = filename.replace('.wav', '').split('_')
+            for part in reversed(parts):
+                if part.isdigit():
+                    return int(part)
+            return 0
+        
+        audio_files.sort(key=extract_chunk_num_from_filename)
+        
+        for i, filename in enumerate(audio_files):
+            file_path = os.path.join(project_dir, filename)
+            chunk_info = {
+                'chunk_num': i + 1,
+                'filename': filename,
+                'file_path': file_path,
+                'size': os.path.getsize(file_path) if os.path.exists(file_path) else 0
+            }
+            chunks.append(chunk_info)
+        
+        return chunks
+        
+    except Exception as e:
+        print(f"Error getting project chunks: {e}")
+        return [] 

src/audiobook/voice_management.py

@@ -0,0 +1,332 @@
+"""
+Voice management utilities for audiobook generation.
+
+Handles voice profile CRUD operations, voice library management, and voice selection.
+"""
+
+import os
+import json
+import shutil
+from pathlib import Path
+from typing import List, Dict, Tuple, Optional, Any
+
+
+def ensure_voice_library_exists(voice_library_path: str) -> None:
+    """Ensure the voice library directory exists.
+    
+    Args:
+        voice_library_path: Path to voice library directory
+    """
+    os.makedirs(voice_library_path, exist_ok=True)
+
+
+def get_voice_profiles(voice_library_path: str) -> List[Dict[str, Any]]:
+    """Get all voice profiles from the voice library.
+    
+    Args:
+        voice_library_path: Path to voice library directory
+        
+    Returns:
+        List of voice profile dictionaries
+    """
+    ensure_voice_library_exists(voice_library_path)
+    profiles = []
+    
+    try:
+        for item in os.listdir(voice_library_path):
+            profile_dir = os.path.join(voice_library_path, item)
+            if os.path.isdir(profile_dir):
+                config_file = os.path.join(profile_dir, "config.json")
+                if os.path.exists(config_file):
+                    try:
+                        with open(config_file, 'r', encoding='utf-8') as f:
+                            profile = json.load(f)
+                            profile['voice_name'] = item
+                            profiles.append(profile)
+                    except Exception as e:
+                        print(f"Warning: Could not load profile {item}: {e}")
+    except Exception as e:
+        print(f"Warning: Could not read voice library: {e}")
+    
+    return profiles
+
+
+def get_voice_choices(voice_library_path: str) -> List[str]:
+    """Get list of available voice names for UI dropdowns.
+    
+    Args:
+        voice_library_path: Path to voice library directory
+        
+    Returns:
+        List of voice names
+    """
+    profiles = get_voice_profiles(voice_library_path)
+    return [profile['voice_name'] for profile in profiles]
+
+
+def get_audiobook_voice_choices(voice_library_path: str) -> List[str]:
+    """Get voice choices formatted for audiobook interface.
+    
+    Args:
+        voice_library_path: Path to voice library directory
+        
+    Returns:
+        List of voice names with display formatting
+    """
+    choices = get_voice_choices(voice_library_path)
+    if not choices:
+        return ["No voices available - Please add voices first"]
+    return choices
+
+
+def get_voice_config(voice_library_path: str, voice_name: str) -> Dict[str, Any]:
+    """Get configuration for a specific voice.
+    
+    Args:
+        voice_library_path: Path to voice library directory
+        voice_name: Name of the voice
+        
+    Returns:
+        Voice configuration dictionary
+    """
+    profile_dir = os.path.join(voice_library_path, voice_name)
+    config_file = os.path.join(profile_dir, "config.json")
+    
+    default_config = {
+        'voice_name': voice_name,
+        'display_name': voice_name,
+        'description': '',
+        'exaggeration': 1.0,
+        'cfg_weight': 1.0,
+        'temperature': 0.7
+    }
+    
+    if os.path.exists(config_file):
+        try:
+            with open(config_file, 'r', encoding='utf-8') as f:
+                config = json.load(f)
+                return {**default_config, **config}
+        except Exception as e:
+            print(f"Warning: Could not load config for {voice_name}: {e}")
+    
+    return default_config
+
+
+def load_voice_for_tts(voice_library_path: str, voice_name: str) -> Tuple[Optional[str], Dict[str, Any]]:
+    """Load voice audio file and configuration for TTS generation.
+    
+    Args:
+        voice_library_path: Path to voice library directory
+        voice_name: Name of the voice to load
+        
+    Returns:
+        tuple: (audio_file_path, voice_config)
+    """
+    if not voice_name:
+        return None, {}
+    
+    profile_dir = os.path.join(voice_library_path, voice_name)
+    if not os.path.exists(profile_dir):
+        return None, {}
+    
+    # Look for audio file
+    audio_file = None
+    for ext in ['.wav', '.mp3', '.flac']:
+        potential_file = os.path.join(profile_dir, f"voice{ext}")
+        if os.path.exists(potential_file):
+            audio_file = potential_file
+            break
+    
+    # Get voice configuration
+    config = get_voice_config(voice_library_path, voice_name)
+    
+    return audio_file, config
+
+
+def save_voice_profile(
+    voice_library_path: str,
+    voice_name: str,
+    display_name: str,
+    description: str,
+    audio_file: Any,
+    exaggeration: float,
+    cfg_weight: float,
+    temperature: float
+) -> str:
+    """Save a new voice profile to the library.
+    
+    Args:
+        voice_library_path: Path to voice library directory
+        voice_name: Internal voice name (used for directory)
+        display_name: Display name for UI
+        description: Voice description
+        audio_file: Audio file data from Gradio
+        exaggeration: Exaggeration parameter
+        cfg_weight: CFG weight parameter
+        temperature: Temperature parameter
+        
+    Returns:
+        Status message
+    """
+    if not voice_name.strip():
+        return "❌ Voice name cannot be empty"
+    
+    # Sanitize voice name for directory
+    safe_voice_name = "".join(c for c in voice_name if c.isalnum() or c in (' ', '-', '_')).strip()
+    safe_voice_name = safe_voice_name.replace(' ', '_')
+    
+    if not safe_voice_name:
+        return "❌ Voice name contains only invalid characters"
+    
+    ensure_voice_library_exists(voice_library_path)
+    
+    profile_dir = os.path.join(voice_library_path, safe_voice_name)
+    os.makedirs(profile_dir, exist_ok=True)
+    
+    try:
+        # Save audio file
+        if audio_file is not None:
+            audio_path = os.path.join(profile_dir, "voice.wav")
+            if isinstance(audio_file, str):
+                # File path provided
+                shutil.copy2(audio_file, audio_path)
+            elif hasattr(audio_file, 'name'):
+                # Gradio file object
+                shutil.copy2(audio_file.name, audio_path)
+            else:
+                return "❌ Invalid audio file format"
+        
+        # Save configuration
+        config = {
+            'voice_name': safe_voice_name,
+            'display_name': display_name or safe_voice_name,
+            'description': description or '',
+            'exaggeration': float(exaggeration),
+            'cfg_weight': float(cfg_weight),
+            'temperature': float(temperature)
+        }
+        
+        config_path = os.path.join(profile_dir, "config.json")
+        with open(config_path, 'w', encoding='utf-8') as f:
+            json.dump(config, f, indent=2)
+        
+        return f"✅ Voice profile '{display_name}' saved successfully"
+        
+    except Exception as e:
+        return f"❌ Error saving voice profile: {str(e)}"
+
+
+def load_voice_profile(voice_library_path: str, voice_name: str) -> Tuple[str, str, str, float, float, float]:
+    """Load voice profile data for editing.
+    
+    Args:
+        voice_library_path: Path to voice library directory
+        voice_name: Name of voice to load
+        
+    Returns:
+        tuple: (display_name, description, audio_path, exaggeration, cfg_weight, temperature)
+    """
+    if not voice_name:
+        return "", "", "", 1.0, 1.0, 0.7
+    
+    config = get_voice_config(voice_library_path, voice_name)
+    audio_file, _ = load_voice_for_tts(voice_library_path, voice_name)
+    
+    return (
+        config.get('display_name', voice_name),
+        config.get('description', ''),
+        audio_file or "",
+        config.get('exaggeration', 1.0),
+        config.get('cfg_weight', 1.0),
+        config.get('temperature', 0.7)
+    )
+
+
+def delete_voice_profile(voice_library_path: str, voice_name: str) -> str:
+    """Delete a voice profile from the library.
+    
+    Args:
+        voice_library_path: Path to voice library directory
+        voice_name: Name of voice to delete
+        
+    Returns:
+        Status message
+    """
+    if not voice_name:
+        return "❌ No voice selected for deletion"
+    
+    profile_dir = os.path.join(voice_library_path, voice_name)
+    
+    if not os.path.exists(profile_dir):
+        return f"❌ Voice profile '{voice_name}' not found"
+    
+    try:
+        shutil.rmtree(profile_dir)
+        return f"✅ Voice profile '{voice_name}' deleted successfully"
+    except Exception as e:
+        return f"❌ Error deleting voice profile: {str(e)}"
+
+
+def refresh_voice_list(voice_library_path: str) -> List[str]:
+    """Refresh and return the current voice list.
+    
+    Args:
+        voice_library_path: Path to voice library directory
+        
+    Returns:
+        Updated list of voice names
+    """
+    return get_voice_choices(voice_library_path)
+
+
+def refresh_voice_choices(voice_library_path: str) -> List[str]:
+    """Refresh voice choices for regular dropdowns.
+    
+    Args:
+        voice_library_path: Path to voice library directory
+        
+    Returns:
+        Updated list of voice choices
+    """
+    return get_voice_choices(voice_library_path)
+
+
+def refresh_audiobook_voice_choices(voice_library_path: str) -> List[str]:
+    """Refresh voice choices for audiobook interface.
+    
+    Args:
+        voice_library_path: Path to voice library directory
+        
+    Returns:
+        Updated list of audiobook voice choices
+    """
+    return get_audiobook_voice_choices(voice_library_path)
+
+
+def create_assignment_interface_with_dropdowns(
+    voice_counts: Dict[str, int], 
+    voice_library_path: str
+) -> List[Any]:
+    """Create voice assignment interface components.
+    
+    Args:
+        voice_counts: Dictionary mapping character names to word counts
+        voice_library_path: Path to voice library directory
+        
+    Returns:
+        List of interface components
+    """
+    # This would typically return Gradio components
+    # For now, return character names and available voices
+    characters = list(voice_counts.keys())
+    available_voices = get_voice_choices(voice_library_path)
+    
+    # Return data that can be used to create dropdowns
+    return [
+        {
+            'character': char,
+            'word_count': voice_counts[char],
+            'available_voices': available_voices
+        }
+        for char in characters[:6]  # Limit to 6 characters
+    ] 

src/chatterbox/__init__.py

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

src/chatterbox/models/s3gen/__init__.py

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

src/chatterbox/models/s3gen/const.py

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

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

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))

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?

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

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

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

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)

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

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

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

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

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

@@ -0,0 +1,330 @@
+# Copyright (c) 2019 Shigeki Karita
+#               2020 Mobvoi Inc (Binbin Zhang)
+#               2022 Xingchen Song ([email protected])
+#               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

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

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

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

@@ -0,0 +1,236 @@
+# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu)
+#               2022 Xingchen Song ([email protected])
+#
+# 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

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)

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

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

@@ -0,0 +1,318 @@
+# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu)
+#               2022 Xingchen Song ([email protected])
+#               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

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

@@ -0,0 +1,71 @@
+# Copyright [2023-11-28] <[email protected], 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,
+}

src/chatterbox/models/s3gen/utils/mask.py

@@ -0,0 +1,193 @@
+# Copyright (c) 2019 Shigeki Karita
+#               2020 Mobvoi Inc (Binbin Zhang)
+#               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
+
+'''
+def subsequent_mask(
+        size: int,
+        device: torch.device = torch.device("cpu"),
+) -> torch.Tensor:
+    """Create mask for subsequent steps (size, size).
+
+    This mask is used only in decoder which works in an auto-regressive mode.
+    This means the current step could only do attention with its left steps.
+
+    In encoder, fully attention is used when streaming is not necessary and
+    the sequence is not long. In this  case, no attention mask is needed.
+
+    When streaming is need, chunk-based attention is used in encoder. See
+    subsequent_chunk_mask for the chunk-based attention mask.
+
+    Args:
+        size (int): size of mask
+        str device (str): "cpu" or "cuda" or torch.Tensor.device
+        dtype (torch.device): result dtype
+
+    Returns:
+        torch.Tensor: mask
+
+    Examples:
+        >>> subsequent_mask(3)
+        [[1, 0, 0],
+         [1, 1, 0],
+         [1, 1, 1]]
+    """
+    ret = torch.ones(size, size, device=device, dtype=torch.bool)
+    return torch.tril(ret)
+'''
+
+
+def subsequent_chunk_mask(
+        size: int,
+        chunk_size: int,
+        num_left_chunks: int = -1,
+        device: torch.device = torch.device("cpu"),
+) -> torch.Tensor:
+    """Create mask for subsequent steps (size, size) with chunk size,
+       this is for streaming encoder
+
+    Args:
+        size (int): size of mask
+        chunk_size (int): size of chunk
+        num_left_chunks (int): number of left chunks
+            <0: use full chunk
+            >=0: use num_left_chunks
+        device (torch.device): "cpu" or "cuda" or torch.Tensor.device
+
+    Returns:
+        torch.Tensor: mask
+
+    Examples:
+        >>> subsequent_chunk_mask(4, 2)
+        [[1, 1, 0, 0],
+         [1, 1, 0, 0],
+         [1, 1, 1, 1],
+         [1, 1, 1, 1]]
+    """
+    # NOTE this modified implementation meets onnx export requirements, but it doesn't support num_left_chunks
+    # actually this is not needed after we have inference cache implemented, will remove it later
+    pos_idx = torch.arange(size, device=device)
+    block_value = (torch.div(pos_idx, chunk_size, rounding_mode='trunc') + 1) * chunk_size
+    ret = pos_idx.unsqueeze(0) < block_value.unsqueeze(1)
+    return ret
+
+
+def add_optional_chunk_mask(xs: torch.Tensor,
+                            masks: torch.Tensor,
+                            use_dynamic_chunk: bool,
+                            use_dynamic_left_chunk: bool,
+                            decoding_chunk_size: int,
+                            static_chunk_size: int,
+                            num_decoding_left_chunks: int,
+                            enable_full_context: bool = True):
+    """ Apply optional mask for encoder.
+
+    Args:
+        xs (torch.Tensor): padded input, (B, L, D), L for max length
+        mask (torch.Tensor): mask for xs, (B, 1, L)
+        use_dynamic_chunk (bool): whether to use dynamic chunk or not
+        use_dynamic_left_chunk (bool): whether to use dynamic left chunk for
+            training.
+        decoding_chunk_size (int): decoding chunk size for dynamic chunk, it's
+            0: default for training, use random dynamic chunk.
+            <0: for decoding, use full chunk.
+            >0: for decoding, use fixed chunk size as set.
+        static_chunk_size (int): chunk size for static chunk training/decoding
+            if it's greater than 0, if use_dynamic_chunk is true,
+            this parameter will be ignored
+        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
+        enable_full_context (bool):
+            True: chunk size is either [1, 25] or full context(max_len)
+            False: chunk size ~ U[1, 25]
+
+    Returns:
+        torch.Tensor: chunk mask of the input xs.
+    """
+    # Whether to use chunk mask or not
+    if use_dynamic_chunk:
+        max_len = xs.size(1)
+        if decoding_chunk_size < 0:
+            chunk_size = max_len
+            num_left_chunks = -1
+        elif decoding_chunk_size > 0:
+            chunk_size = decoding_chunk_size
+            num_left_chunks = num_decoding_left_chunks
+        else:
+            # chunk size is either [1, 25] or full context(max_len).
+            # Since we use 4 times subsampling and allow up to 1s(100 frames)
+            # delay, the maximum frame is 100 / 4 = 25.
+            chunk_size = torch.randint(1, max_len, (1, )).item()
+            num_left_chunks = -1
+            if chunk_size > max_len // 2 and enable_full_context:
+                chunk_size = max_len
+            else:
+                chunk_size = chunk_size % 25 + 1
+                if use_dynamic_left_chunk:
+                    max_left_chunks = (max_len - 1) // chunk_size
+                    num_left_chunks = torch.randint(0, max_left_chunks,
+                                                    (1, )).item()
+        chunk_masks = subsequent_chunk_mask(xs.size(1), chunk_size,
+                                            num_left_chunks,
+                                            xs.device)  # (L, L)
+        chunk_masks = chunk_masks.unsqueeze(0)  # (1, L, L)
+        chunk_masks = masks & chunk_masks  # (B, L, L)
+    elif static_chunk_size > 0:
+        num_left_chunks = num_decoding_left_chunks
+        chunk_masks = subsequent_chunk_mask(xs.size(1), static_chunk_size,
+                                            num_left_chunks,
+                                            xs.device)  # (L, L)
+        chunk_masks = chunk_masks.unsqueeze(0)  # (1, L, L)
+        chunk_masks = masks & chunk_masks  # (B, L, L)
+    else:
+        chunk_masks = masks
+    assert chunk_masks.dtype == torch.bool
+    if (chunk_masks.sum(dim=-1) == 0).sum().item() != 0:
+        logging.warning('get chunk_masks all false at some timestep, force set to true, make sure they are masked in futuer computation!')
+        chunk_masks[chunk_masks.sum(dim=-1)==0] = True
+    return chunk_masks
+
+
+def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
+    """Make mask tensor containing indices of padded part.
+
+    See description of make_non_pad_mask.
+
+    Args:
+        lengths (torch.Tensor): Batch of lengths (B,).
+    Returns:
+        torch.Tensor: Mask tensor containing indices of padded part.
+
+    Examples:
+        >>> lengths = [5, 3, 2]
+        >>> make_pad_mask(lengths)
+        masks = [[0, 0, 0, 0 ,0],
+                 [0, 0, 0, 1, 1],
+                 [0, 0, 1, 1, 1]]
+    """
+    batch_size = lengths.size(0)
+    max_len = max_len if max_len > 0 else lengths.max().item()
+    seq_range = torch.arange(0,
+                             max_len,
+                             dtype=torch.int64,
+                             device=lengths.device)
+    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
+    seq_length_expand = lengths.unsqueeze(-1)
+    mask = seq_range_expand >= seq_length_expand
+    return mask

src/chatterbox/models/s3gen/utils/mel.py

@@ -0,0 +1,81 @@
+"""mel-spectrogram extraction in Matcha-TTS"""
+from librosa.filters import mel as librosa_mel_fn
+import torch
+import numpy as np
+
+
+# NOTE: they decalred these global vars
+mel_basis = {}
+hann_window = {}
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def spectral_normalize_torch(magnitudes):
+    output = dynamic_range_compression_torch(magnitudes)
+    return output
+
+"""
+feat_extractor: !name:matcha.utils.audio.mel_spectrogram
+    n_fft: 1920
+    num_mels: 80
+    sampling_rate: 24000
+    hop_size: 480
+    win_size: 1920
+    fmin: 0
+    fmax: 8000
+    center: False
+
+"""
+
+def mel_spectrogram(y, n_fft=1920, num_mels=80, sampling_rate=24000, hop_size=480, win_size=1920,
+                    fmin=0, fmax=8000, center=False):
+    """Copied from https://github.com/shivammehta25/Matcha-TTS/blob/main/matcha/utils/audio.py
+    Set default values according to Cosyvoice's config.
+    """
+
+    if isinstance(y, np.ndarray):
+        y = torch.tensor(y).float()
+
+    if len(y.shape) == 1:
+        y = y[None, ]
+
+    if torch.min(y) < -1.0:
+        print("min value is ", torch.min(y))
+    if torch.max(y) > 1.0:
+        print("max value is ", torch.max(y))
+
+    global mel_basis, hann_window  # pylint: disable=global-statement,global-variable-not-assigned
+    if f"{str(fmax)}_{str(y.device)}" not in mel_basis:
+        mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
+        mel_basis[str(fmax) + "_" + str(y.device)] = torch.from_numpy(mel).float().to(y.device)
+        hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device)
+
+    y = torch.nn.functional.pad(
+        y.unsqueeze(1), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), mode="reflect"
+    )
+    y = y.squeeze(1)
+
+    spec = torch.view_as_real(
+        torch.stft(
+            y,
+            n_fft,
+            hop_length=hop_size,
+            win_length=win_size,
+            window=hann_window[str(y.device)],
+            center=center,
+            pad_mode="reflect",
+            normalized=False,
+            onesided=True,
+            return_complex=True,
+        )
+    )
+
+    spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))
+
+    spec = torch.matmul(mel_basis[str(fmax) + "_" + str(y.device)], spec)
+    spec = spectral_normalize_torch(spec)
+
+    return spec

src/chatterbox/models/s3gen/xvector.py

@@ -0,0 +1,428 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
+
+
+from collections import OrderedDict
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint as cp
+import torchaudio.compliance.kaldi as Kaldi
+
+
+def pad_list(xs, pad_value):
+    """Perform padding for the list of tensors.
+
+    Args:
+        xs (List): List of Tensors [(T_1, `*`), (T_2, `*`), ..., (T_B, `*`)].
+        pad_value (float): Value for padding.
+
+    Returns:
+        Tensor: Padded tensor (B, Tmax, `*`).
+
+    Examples:
+        >>> x = [torch.ones(4), torch.ones(2), torch.ones(1)]
+        >>> x
+        [tensor([1., 1., 1., 1.]), tensor([1., 1.]), tensor([1.])]
+        >>> pad_list(x, 0)
+        tensor([[1., 1., 1., 1.],
+                [1., 1., 0., 0.],
+                [1., 0., 0., 0.]])
+
+    """
+    n_batch = len(xs)
+    max_len = max(x.size(0) for x in xs)
+    pad = xs[0].new(n_batch, max_len, *xs[0].size()[1:]).fill_(pad_value)
+
+    for i in range(n_batch):
+        pad[i, : xs[i].size(0)] = xs[i]
+
+    return pad
+
+
+def extract_feature(audio):
+    features = []
+    feature_times = []
+    feature_lengths = []
+    for au in audio:
+        feature = Kaldi.fbank(au.unsqueeze(0), num_mel_bins=80)
+        feature = feature - feature.mean(dim=0, keepdim=True)
+        features.append(feature)
+        feature_times.append(au.shape[0])
+        feature_lengths.append(feature.shape[0])
+    # padding for batch inference
+    features_padded = pad_list(features, pad_value=0)
+    # features = torch.cat(features)
+    return features_padded, feature_lengths, feature_times
+
+
+class BasicResBlock(torch.nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(BasicResBlock, self).__init__()
+        self.conv1 = torch.nn.Conv2d(
+            in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False
+        )
+        self.bn1 = torch.nn.BatchNorm2d(planes)
+        self.conv2 = torch.nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = torch.nn.BatchNorm2d(planes)
+
+        self.shortcut = torch.nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = torch.nn.Sequential(
+                torch.nn.Conv2d(
+                    in_planes,
+                    self.expansion * planes,
+                    kernel_size=1,
+                    stride=(stride, 1),
+                    bias=False,
+                ),
+                torch.nn.BatchNorm2d(self.expansion * planes),
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class FCM(torch.nn.Module):
+    def __init__(self, block=BasicResBlock, num_blocks=[2, 2], m_channels=32, feat_dim=80):
+        super(FCM, self).__init__()
+        self.in_planes = m_channels
+        self.conv1 = torch.nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = torch.nn.BatchNorm2d(m_channels)
+
+        self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2)
+        self.layer2 = self._make_layer(block, m_channels, num_blocks[0], stride=2)
+
+        self.conv2 = torch.nn.Conv2d(
+            m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False
+        )
+        self.bn2 = torch.nn.BatchNorm2d(m_channels)
+        self.out_channels = m_channels * (feat_dim // 8)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+        return torch.nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = x.unsqueeze(1)
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = F.relu(self.bn2(self.conv2(out)))
+
+        shape = out.shape
+        out = out.reshape(shape[0], shape[1] * shape[2], shape[3])
+        return out
+
+
+def get_nonlinear(config_str, channels):
+    nonlinear = torch.nn.Sequential()
+    for name in config_str.split("-"):
+        if name == "relu":
+            nonlinear.add_module("relu", torch.nn.ReLU(inplace=True))
+        elif name == "prelu":
+            nonlinear.add_module("prelu", torch.nn.PReLU(channels))
+        elif name == "batchnorm":
+            nonlinear.add_module("batchnorm", torch.nn.BatchNorm1d(channels))
+        elif name == "batchnorm_":
+            nonlinear.add_module("batchnorm", torch.nn.BatchNorm1d(channels, affine=False))
+        else:
+            raise ValueError("Unexpected module ({}).".format(name))
+    return nonlinear
+
+
+def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, eps=1e-2):
+    mean = x.mean(dim=dim)
+    std = x.std(dim=dim, unbiased=unbiased)
+    stats = torch.cat([mean, std], dim=-1)
+    if keepdim:
+        stats = stats.unsqueeze(dim=dim)
+    return stats
+
+
+class StatsPool(torch.nn.Module):
+    def forward(self, x):
+        return statistics_pooling(x)
+
+
+class TDNNLayer(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        bias=False,
+        config_str="batchnorm-relu",
+    ):
+        super(TDNNLayer, self).__init__()
+        if padding < 0:
+            assert (
+                kernel_size % 2 == 1
+            ), "Expect equal paddings, but got even kernel size ({})".format(kernel_size)
+            padding = (kernel_size - 1) // 2 * dilation
+        self.linear = torch.nn.Conv1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+        self.nonlinear = get_nonlinear(config_str, out_channels)
+
+    def forward(self, x):
+        x = self.linear(x)
+        x = self.nonlinear(x)
+        return x
+
+
+class CAMLayer(torch.nn.Module):
+    def __init__(
+        self, bn_channels, out_channels, kernel_size, stride, padding, dilation, bias, reduction=2
+    ):
+        super(CAMLayer, self).__init__()
+        self.linear_local = torch.nn.Conv1d(
+            bn_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+        self.linear1 = torch.nn.Conv1d(bn_channels, bn_channels // reduction, 1)
+        self.relu = torch.nn.ReLU(inplace=True)
+        self.linear2 = torch.nn.Conv1d(bn_channels // reduction, out_channels, 1)
+        self.sigmoid = torch.nn.Sigmoid()
+
+    def forward(self, x):
+        y = self.linear_local(x)
+        context = x.mean(-1, keepdim=True) + self.seg_pooling(x)
+        context = self.relu(self.linear1(context))
+        m = self.sigmoid(self.linear2(context))
+        return y * m
+
+    def seg_pooling(self, x, seg_len=100, stype="avg"):
+        if stype == "avg":
+            seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
+        elif stype == "max":
+            seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
+        else:
+            raise ValueError("Wrong segment pooling type.")
+        shape = seg.shape
+        seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1)
+        seg = seg[..., : x.shape[-1]]
+        return seg
+
+
+class CAMDenseTDNNLayer(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        bn_channels,
+        kernel_size,
+        stride=1,
+        dilation=1,
+        bias=False,
+        config_str="batchnorm-relu",
+        memory_efficient=False,
+    ):
+        super(CAMDenseTDNNLayer, self).__init__()
+        assert kernel_size % 2 == 1, "Expect equal paddings, but got even kernel size ({})".format(
+            kernel_size
+        )
+        padding = (kernel_size - 1) // 2 * dilation
+        self.memory_efficient = memory_efficient
+        self.nonlinear1 = get_nonlinear(config_str, in_channels)
+        self.linear1 = torch.nn.Conv1d(in_channels, bn_channels, 1, bias=False)
+        self.nonlinear2 = get_nonlinear(config_str, bn_channels)
+        self.cam_layer = CAMLayer(
+            bn_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+
+    def bn_function(self, x):
+        return self.linear1(self.nonlinear1(x))
+
+    def forward(self, x):
+        if self.training and self.memory_efficient:
+            x = cp.checkpoint(self.bn_function, x)
+        else:
+            x = self.bn_function(x)
+        x = self.cam_layer(self.nonlinear2(x))
+        return x
+
+
+class CAMDenseTDNNBlock(torch.nn.ModuleList):
+    def __init__(
+        self,
+        num_layers,
+        in_channels,
+        out_channels,
+        bn_channels,
+        kernel_size,
+        stride=1,
+        dilation=1,
+        bias=False,
+        config_str="batchnorm-relu",
+        memory_efficient=False,
+    ):
+        super(CAMDenseTDNNBlock, self).__init__()
+        for i in range(num_layers):
+            layer = CAMDenseTDNNLayer(
+                in_channels=in_channels + i * out_channels,
+                out_channels=out_channels,
+                bn_channels=bn_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                dilation=dilation,
+                bias=bias,
+                config_str=config_str,
+                memory_efficient=memory_efficient,
+            )
+            self.add_module("tdnnd%d" % (i + 1), layer)
+
+    def forward(self, x):
+        for layer in self:
+            x = torch.cat([x, layer(x)], dim=1)
+        return x
+
+
+class TransitLayer(torch.nn.Module):
+    def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"):
+        super(TransitLayer, self).__init__()
+        self.nonlinear = get_nonlinear(config_str, in_channels)
+        self.linear = torch.nn.Conv1d(in_channels, out_channels, 1, bias=bias)
+
+    def forward(self, x):
+        x = self.nonlinear(x)
+        x = self.linear(x)
+        return x
+
+
+class DenseLayer(torch.nn.Module):
+    def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"):
+        super(DenseLayer, self).__init__()
+        self.linear = torch.nn.Conv1d(in_channels, out_channels, 1, bias=bias)
+        self.nonlinear = get_nonlinear(config_str, out_channels)
+
+    def forward(self, x):
+        if len(x.shape) == 2:
+            x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1)
+        else:
+            x = self.linear(x)
+        x = self.nonlinear(x)
+        return x
+
+# @tables.register("model_classes", "CAMPPlus")
+class CAMPPlus(torch.nn.Module):
+    def __init__(
+        self,
+        feat_dim=80,
+        embedding_size=192,
+        growth_rate=32,
+        bn_size=4,
+        init_channels=128,
+        config_str="batchnorm-relu",
+        memory_efficient=True,
+        output_level="segment",
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.head = FCM(feat_dim=feat_dim)
+        channels = self.head.out_channels
+        self.output_level = output_level
+
+        self.xvector = torch.nn.Sequential(
+            OrderedDict(
+                [
+                    (
+                        "tdnn",
+                        TDNNLayer(
+                            channels,
+                            init_channels,
+                            5,
+                            stride=2,
+                            dilation=1,
+                            padding=-1,
+                            config_str=config_str,
+                        ),
+                    ),
+                ]
+            )
+        )
+        channels = init_channels
+        for i, (num_layers, kernel_size, dilation) in enumerate(
+            zip((12, 24, 16), (3, 3, 3), (1, 2, 2))
+        ):
+            block = CAMDenseTDNNBlock(
+                num_layers=num_layers,
+                in_channels=channels,
+                out_channels=growth_rate,
+                bn_channels=bn_size * growth_rate,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                config_str=config_str,
+                memory_efficient=memory_efficient,
+            )
+            self.xvector.add_module("block%d" % (i + 1), block)
+            channels = channels + num_layers * growth_rate
+            self.xvector.add_module(
+                "transit%d" % (i + 1),
+                TransitLayer(channels, channels // 2, bias=False, config_str=config_str),
+            )
+            channels //= 2
+
+        self.xvector.add_module("out_nonlinear", get_nonlinear(config_str, channels))
+
+        if self.output_level == "segment":
+            self.xvector.add_module("stats", StatsPool())
+            self.xvector.add_module(
+                "dense", DenseLayer(channels * 2, embedding_size, config_str="batchnorm_")
+            )
+        else:
+            assert (
+                self.output_level == "frame"
+            ), "`output_level` should be set to 'segment' or 'frame'. "
+
+        for m in self.modules():
+            if isinstance(m, (torch.nn.Conv1d, torch.nn.Linear)):
+                torch.nn.init.kaiming_normal_(m.weight.data)
+                if m.bias is not None:
+                    torch.nn.init.zeros_(m.bias)
+
+    def forward(self, x):
+        x = x.permute(0, 2, 1)  # (B,T,F) => (B,F,T)
+        x = self.head(x)
+        x = self.xvector(x)
+        if self.output_level == "frame":
+            x = x.transpose(1, 2)
+        return x
+
+    def inference(self, audio_list):
+        speech, speech_lengths, speech_times = extract_feature(audio_list)
+        results = self.forward(speech.to(torch.float32))
+        return results

src/chatterbox/models/s3tokenizer/__init__.py

@@ -0,0 +1,30 @@
+from .s3tokenizer import (
+    S3_SR,
+    S3_HOP,
+    S3_TOKEN_HOP,
+    S3_TOKEN_RATE,
+    SPEECH_VOCAB_SIZE,
+    S3Tokenizer,
+)
+
+
+SOS = SPEECH_VOCAB_SIZE
+EOS = SPEECH_VOCAB_SIZE + 1
+
+
+
+def drop_invalid_tokens(x):
+    """Drop SoS and EoS"""
+    assert len(x.shape) == 1 or (len(x.shape) == 2 and x.shape[0] == 1), "only batch size of one allowed for now"
+    if SOS in x:
+        s = (x == SOS).nonzero(as_tuple=True)[0].squeeze(0) + 1
+    else:
+        s = 0
+
+    if EOS in x:
+        e = (x == EOS).nonzero(as_tuple=True)[0].squeeze(0)
+    else:
+        e = None
+
+    x = x[s: e]
+    return x

src/chatterbox/models/s3tokenizer/s3tokenizer.py

@@ -0,0 +1,168 @@
+from typing import List, Tuple
+
+import numpy as np
+import librosa
+import torch
+import torch.nn.functional as F
+from s3tokenizer.utils import padding
+from s3tokenizer.model_v2 import (
+    S3TokenizerV2,
+    ModelConfig,
+)
+
+
+# Sampling rate of the inputs to S3TokenizerV2
+S3_SR = 16_000
+S3_HOP = 160  # 100 frames/sec
+S3_TOKEN_HOP = 640  # 25 tokens/sec
+S3_TOKEN_RATE = 25
+SPEECH_VOCAB_SIZE = 6561
+
+
+class S3Tokenizer(S3TokenizerV2):
+    """
+    s3tokenizer.S3TokenizerV2 with the following changes:
+    - a more integrated `forward`
+    - compute `log_mel_spectrogram` using `_mel_filters` and `window` in `register_buffers`
+    """
+
+    ignore_state_dict_missing = ("_mel_filters", "window")
+
+    def __init__(
+        self,
+        name: str="speech_tokenizer_v2_25hz",
+        config: ModelConfig = ModelConfig()
+    ):
+        super().__init__(name)
+
+        self.n_fft = 400
+        _mel_filters = librosa.filters.mel(
+            sr=S3_SR,
+            n_fft=self.n_fft,
+            n_mels=config.n_mels
+        )
+        self.register_buffer(
+            "_mel_filters",
+            torch.FloatTensor(_mel_filters),
+        )
+
+        self.register_buffer(
+            "window",
+            torch.hann_window(self.n_fft),
+        )
+
+    def pad(self, wavs, sr) -> List[torch.Tensor]:
+        """
+        Given a list of wavs with the same `sample_rate`, pad them so that the length is multiple of 40ms (S3 runs at 25 token/sec).
+        """
+        processed_wavs = []
+        for wav in wavs:
+            if isinstance(wav, np.ndarray):
+                wav = torch.from_numpy(wav)
+            if wav.dim() == 1:
+                wav = wav.unsqueeze(0)
+
+            n_tokens = (wav.shape[1] / sr) * S3_TOKEN_RATE
+            n_tokens = np.ceil(n_tokens)
+            intended_wav_len = n_tokens * (sr / S3_TOKEN_RATE)
+            intended_wav_len = int(intended_wav_len)
+            wav = torch.nn.functional.pad(
+                wav,
+                (0, intended_wav_len - wav.shape[-1]),
+                mode="constant",
+                value=0
+            )
+            processed_wavs.append(wav)
+        return processed_wavs
+
+    def _prepare_audio(self, wavs):
+        """Prepare a list of audios for s3tokenizer processing."""
+        processed_wavs = []
+        for wav in wavs:
+            if isinstance(wav, np.ndarray):
+                wav = torch.from_numpy(wav)
+            if wav.dim() == 1:
+                wav = wav.unsqueeze(0)
+
+            processed_wavs.append(wav)
+        return processed_wavs
+
+    @torch.no_grad()
+    def forward(
+        self,
+        wavs: torch.Tensor,
+        accelerator: 'Accelerator'=None,
+        max_len: int=None,
+    ) -> Tuple[torch.Tensor, torch.LongTensor]:
+        """
+        NOTE: mel-spec has a hop size of 160 points (100 frame/sec).
+        FIXME: this class inherits `nn.Module` but doesn't accept `torch.Tensor` and handles a list of wavs one by one, which is unexpected.
+
+        Args
+        ----
+        - `wavs`: 16 kHz speech audio
+        - `max_len` max length to truncate the output sequence to (25 token/sec).
+        NOTE: please pad the waveform if longer sequence is needed.
+        """
+        processed_wavs = self._prepare_audio(wavs)
+        mels, mel_lens = [], []
+        for wav in processed_wavs:
+            wav = wav.to(self.device)
+            mel = self.log_mel_spectrogram(wav)  # [B=1, F, T]
+            if max_len is not None:
+                mel = mel[..., :max_len * 4]  # num_mel_frames = 4 * num_tokens
+            mels.append(mel.squeeze(0))
+
+        mels, mel_lens = padding(mels)
+        if accelerator is None:
+            tokenizer = self
+        else:
+            tokenizer = accelerator.unwrap_model(self)
+
+        speech_tokens, speech_token_lens = tokenizer.quantize(mels, mel_lens.to(self.device))
+        return (
+            speech_tokens.long().detach(),
+            speech_token_lens.long().detach(),
+        )
+
+    def log_mel_spectrogram(
+        self,
+        audio: torch.Tensor,
+        padding: int = 0,
+    ):
+        """
+        Compute the log-Mel spectrogram of
+
+        Parameters
+        ----------
+        audio: torch.Tensor, shape = (*)
+            The path to audio or either a NumPy array or Tensor containing the
+            audio waveform in 16 kHz
+
+        padding: int
+            Number of zero samples to pad to the right
+
+        Returns
+        -------
+        torch.Tensor, shape = (128, n_frames)
+            A Tensor that contains the Mel spectrogram
+        """
+        if not torch.is_tensor(audio):
+            audio = torch.from_numpy(audio)
+
+        audio = audio.to(self.device)
+        if padding > 0:
+            audio = F.pad(audio, (0, padding))
+        stft = torch.stft(
+            audio, self.n_fft, S3_HOP,
+            window=self.window.to(self.device),
+            return_complex=True
+        )
+        magnitudes = stft[..., :-1].abs()**2
+
+        mel_spec = self._mel_filters.to(self.device) @ magnitudes
+
+        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+        log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
+        log_spec = (log_spec + 4.0) / 4.0
+        return log_spec

src/chatterbox/models/t3/__init__.py

@@ -0,0 +1 @@
+from .t3 import T3

src/chatterbox/models/t3/inference/alignment_stream_analyzer.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2025 Resemble AI
+# Author: John Meade, Jeremy Hsu
+# MIT License
+import logging
+import torch
+from dataclasses import dataclass
+from types import MethodType
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AlignmentAnalysisResult:
+    # was this frame detected as being part of a noisy beginning chunk with potential hallucinations?
+    false_start: bool
+    # was this frame detected as being part of a long tail with potential hallucinations?
+    long_tail: bool
+    # was this frame detected as repeating existing text content?
+    repetition: bool
+    # was the alignment position of this frame too far from the previous frame?
+    discontinuity: bool
+    # has inference reached the end of the text tokens? eg, this remains false if inference stops early
+    complete: bool
+    # approximate position in the text token sequence. Can be used for generating online timestamps.
+    position: int
+
+
+class AlignmentStreamAnalyzer:
+    def __init__(self, tfmr, queue, text_tokens_slice, alignment_layer_idx=9, eos_idx=0):
+        """
+        Some transformer TTS models implicitly solve text-speech alignment in one or more of their self-attention
+        activation maps. This module exploits this to perform online integrity checks which streaming.
+        A hook is injected into the specified attention layer, and heuristics are used to determine alignment
+        position, repetition, etc.
+
+        NOTE: currently requires no queues.
+        """
+        # self.queue = queue
+        self.text_tokens_slice = (i, j) = text_tokens_slice
+        self.eos_idx = eos_idx
+        self.alignment = torch.zeros(0, j-i)
+        # self.alignment_bin = torch.zeros(0, j-i)
+        self.curr_frame_pos = 0
+        self.text_position = 0
+
+        self.started = False
+        self.started_at = None
+
+        self.complete = False
+        self.completed_at = None
+
+        # Using `output_attentions=True` is incompatible with optimized attention kernels, so
+        # using it for all layers slows things down too much. We can apply it to just one layer
+        # by intercepting the kwargs and adding a forward hook (credit: jrm)
+        self.last_aligned_attn = None
+        self._add_attention_spy(tfmr, alignment_layer_idx)
+
+    def _add_attention_spy(self, tfmr, alignment_layer_idx):
+        """
+        Adds a forward hook to a specific attention layer to collect outputs.
+        Using `output_attentions=True` is incompatible with optimized attention kernels, so
+        using it for all layers slows things down too much.
+        (credit: jrm)
+        """
+
+        def attention_forward_hook(module, input, output):
+            """
+            See `LlamaAttention.forward`; the output is a 3-tuple: `attn_output, attn_weights, past_key_value`.
+            NOTE:
+            - When `output_attentions=True`, `LlamaSdpaAttention.forward` calls `LlamaAttention.forward`.
+            - `attn_output` has shape [B, H, T0, T0] for the 0th entry, and [B, H, 1, T0+i] for the rest i-th.
+            """
+            step_attention = output[1].cpu() # (B, 16, N, N)
+            self.last_aligned_attn = step_attention[0].mean(0) # (N, N)
+
+        target_layer = tfmr.layers[alignment_layer_idx].self_attn
+        hook_handle = target_layer.register_forward_hook(attention_forward_hook)
+
+        # Backup original forward
+        original_forward = target_layer.forward
+        def patched_forward(self, *args, **kwargs):
+            kwargs['output_attentions'] = True
+            return original_forward(*args, **kwargs)
+
+        # TODO: how to unpatch it?
+        target_layer.forward = MethodType(patched_forward, target_layer)
+
+    def step(self, logits):
+        """
+        Emits an AlignmentAnalysisResult into the output queue, and potentially modifies the logits to force an EOS.
+        """
+        # extract approximate alignment matrix chunk (1 frame at a time after the first chunk)
+        aligned_attn = self.last_aligned_attn # (N, N)
+        i, j = self.text_tokens_slice
+        if self.curr_frame_pos == 0:
+            # first chunk has conditioning info, text tokens, and BOS token
+            A_chunk = aligned_attn[j:, i:j].clone().cpu() # (T, S)
+        else:
+            # subsequent chunks have 1 frame due to KV-caching
+            A_chunk = aligned_attn[:, i:j].clone().cpu() # (1, S)
+
+        # TODO: monotonic masking; could have issue b/c spaces are often skipped.
+        A_chunk[:, self.curr_frame_pos + 1:] = 0
+
+
+        self.alignment = torch.cat((self.alignment, A_chunk), dim=0)
+
+        A = self.alignment
+        T, S = A.shape
+
+        # update position
+        cur_text_posn = A_chunk[-1].argmax()
+        discontinuity = not(-4 < cur_text_posn - self.text_position < 7) # NOTE: very lenient!
+        if not discontinuity:
+            self.text_position = cur_text_posn
+
+        # Hallucinations at the start of speech show up as activations at the bottom of the attention maps!
+        # To mitigate this, we just wait until there are no activations far off-diagonal in the last 2 tokens,
+        # and there are some strong activations in the first few tokens.
+        false_start = (not self.started) and (A[-2:, -2:].max() > 0.1 or A[:, :4].max() < 0.5)
+        self.started = not false_start
+        if self.started and self.started_at is None:
+            self.started_at = T
+
+        # Is generation likely complete?
+        self.complete = self.complete or self.text_position >= S - 3
+        if self.complete and self.completed_at is None:
+            self.completed_at = T
+
+        # NOTE: EOS rarely assigned activations, and second-last token is often punctuation, so use last 3 tokens.
+        # NOTE: due to the false-start behaviour, we need to make sure we skip activations for the first few tokens.
+        last_text_token_duration = A[15:, -3:].sum()
+
+        # Activations for the final token that last too long are likely hallucinations.
+        long_tail = self.complete and (A[self.completed_at:, -3:].sum(dim=0).max() >= 10) # 400ms
+
+        # If there are activations in previous tokens after generation has completed, assume this is a repetition error.
+        repetition = self.complete and (A[self.completed_at:, :-5].max(dim=1).values.sum() > 5)
+
+        # If a bad ending is detected, force emit EOS by modifying logits
+        # NOTE: this means logits may be inconsistent with latents!
+        if long_tail or repetition:
+            logger.warn(f"forcing EOS token, {long_tail=}, {repetition=}")
+            # (±2**15 is safe for all dtypes >= 16bit)
+            logits = -(2**15) * torch.ones_like(logits)
+            logits[..., self.eos_idx] = 2**15
+
+        # Suppress EoS to prevent early termination
+        if cur_text_posn < S - 3: # FIXME: arbitrary
+            logits[..., self.eos_idx] = -2**15
+
+        self.curr_frame_pos += 1
+        return logits

src/chatterbox/models/t3/inference/t3_hf_backend.py

@@ -0,0 +1,116 @@
+from typing import Optional
+
+import torch
+from torch import nn as nn
+from transformers import LlamaConfig, LlamaModel, LlamaPreTrainedModel, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+
+
+class T3HuggingfaceBackend(LlamaPreTrainedModel, GenerationMixin):
+    """
+    Override some HuggingFace interface methods so we can use the standard `generate` method with our
+    custom embedding / logit layers.
+
+    NOTE: need to extend "*PreTrainedModel" to avoid re-initializing weights!
+    """
+
+    def __init__(
+        self,
+        config: LlamaConfig,
+        llama: LlamaModel,
+        *,
+        speech_enc,
+        speech_head,
+        latents_queue=None,
+        logits_queue=None,
+        alignment_stream_analyzer: 'AlignmentStreamAnalyzer'=None,
+    ):
+        super().__init__(config)
+        self.model = llama
+        self.speech_enc = speech_enc
+        self.speech_head = speech_head
+        self._added_cond = False
+        self.alignment_stream_analyzer = alignment_stream_analyzer
+
+    @torch.inference_mode()
+    def prepare_inputs_for_generation(
+        self, input_ids: torch.Tensor, decoder_cond: torch.Tensor, use_cache: bool, past_key_values=None,
+        # This argument was introduced in some recent version of transformers (>=4.29.1)
+        cache_position=None
+    ):
+        """
+        This is a method used by huggingface's generate() method.
+        Overridden here to apply our custom speech token embedding layer.
+
+        :param input_ids: (B, S) int64 tensors of input tokens.
+        :param decoder_cond: (B, T, C) float32 tensor of conditioning (prefixed to <input_embeds>)
+        """
+
+        # Make use of the kv cache: only the last input ID is new, we trim away all the ones before
+        if not use_cache:
+            past_key_values = None
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+
+        # custom speech token embedding layer
+        inputs_embeds = self.speech_enc(input_ids)
+
+        # prefix decoder conditioning if applicable
+        if not self._added_cond:
+            assert past_key_values is not None # should be first step
+            if decoder_cond.size(0) != inputs_embeds.size(0):
+                decoder_cond = decoder_cond.expand(inputs_embeds.size(0), -1, -1)
+            inputs_embeds = torch.cat([decoder_cond, inputs_embeds], dim=1)
+            self._added_cond = True
+
+        return {
+            "inputs_embeds": inputs_embeds,
+            "past_key_values": past_key_values,
+            "use_cache": use_cache,
+        }
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        past_key_values: Optional[torch.Tensor]=None,
+        use_cache=True,
+        output_attentions=False,
+        output_hidden_states=True,
+        return_dict=True,
+    ):
+        """
+        This is a method used by huggingface's generate() method.
+        Overridden here to apply our custom layer norm and speech logit projection layers.
+
+        :param inputs_embeds: (B, S, C) float32 tensor of conditioning inputs. If past key values are given,
+        S should be 1.
+        """
+        is_large_input = inputs_embeds.size(1) != 1
+        has_cache = past_key_values is not None and len(past_key_values) > 0
+        assert not (is_large_input and has_cache)
+        assert return_dict
+        assert output_hidden_states
+
+        tfmr_out = self.model(
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+        )
+        hidden_states = tfmr_out.hidden_states[-1]  # (B, seq, dim)
+
+        logits = self.speech_head(hidden_states)
+        # assert inputs_embeds.size(0) == 1 # (disabled for CFG)
+
+        # NOTE: hallucination handler may modify logits to force emit an EOS token
+        # logits = self.alignment_stream_analyzer.step(logits)
+
+        return CausalLMOutputWithCrossAttentions(
+            logits=logits,
+            past_key_values=tfmr_out.past_key_values,
+            hidden_states=tfmr_out.hidden_states,
+            attentions=tfmr_out.attentions,
+        )

src/chatterbox/models/t3/llama_configs.py

@@ -0,0 +1,37 @@
+LLAMA_520M_CONFIG_DICT = dict(
+    # Arbitrary small number that won't cause problems when loading.
+    # These param are unused due to custom input layers.
+    vocab_size=8,
+    # default params needed for loading most pretrained 1B weights
+    max_position_embeddings=131072,
+    hidden_size=1024,
+    intermediate_size=4096,
+    num_hidden_layers=30,
+    num_attention_heads=16,
+    attn_implementation="sdpa",
+    head_dim=64,
+    tie_word_embeddings=False,
+    hidden_act="silu",
+    attention_bias=False,
+    attention_dropout=0.0,
+    initializer_range=0.02,
+    mlp_bias=False,
+    model_type="llama",
+    num_key_value_heads=16,
+    pretraining_tp=1,
+    rms_norm_eps=1e-05,
+    rope_scaling=dict(
+        factor=8.0,
+        high_freq_factor=4.0,
+        low_freq_factor=1.0,
+        original_max_position_embeddings=8192,
+        rope_type="llama3"
+    ),
+    rope_theta=500000.0,
+    torch_dtype="bfloat16",
+    use_cache=True,
+)
+
+LLAMA_CONFIGS = {
+    "Llama_520M": LLAMA_520M_CONFIG_DICT,
+}