Deploy ACE-Step Custom Edition with bug fixes

.gitattributes

@@ -0,0 +1,3 @@
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,67 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+venv/
+env/
+ENV/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Application
+outputs/
+timelines/
+lora_training/prepared_data/
+lora_training/models/
+logs/
+models/
+*.wav
+*.mp3
+*.flac
+*.ogg
+
+# Config (keep example)
+config.yaml
+
+# Jupyter
+.ipynb_checkpoints/
+*.ipynb
+
+# Model cache
+.cache/
+huggingface/
+
+# Environment
+.env
+.env.local
+
+# Test outputs
+test_outputs/

.python-version

@@ -0,0 +1 @@
+python-3.11.0

DEPLOYMENT.md

@@ -0,0 +1,367 @@
+# HuggingFace Spaces Deployment Guide
+
+## Quick Deploy to HuggingFace Spaces
+
+### Prerequisites
+- HuggingFace account (create at https://huggingface.co/join)
+- Git installed on your machine
+- Git LFS installed (for large files)
+
+### Method 1: Web Upload (Easiest)
+
+1. **Create New Space**
+   - Go to https://huggingface.co/new-space
+   - Name: `ace-step-custom` (or your choice)
+   - License: MIT
+   - SDK: Gradio
+   - Hardware: A10G Small (or better)
+   - Click "Create Space"
+
+2. **Upload Files**
+   - Click "Files and versions" tab
+   - Click "Add file" → "Upload files"
+   - Upload all files from `d:\2025-vibe-coding\ACE-Step-Custom\`:
+     - `app.py`
+     - `requirements.txt`
+     - `config.yaml`
+     - `README.md` (with YAML frontmatter)
+     - `LICENSE`
+     - `.gitignore`
+     - Entire `src/` directory
+     - Entire `scripts/` directory
+   - Commit changes
+
+3. **Configure Space**
+   - Go to "Settings" tab
+   - Set Python version: 3.10
+   - Enable GPU: A10G Small (minimum) or A100 (recommended)
+   - Set timeout: 30 minutes (for long generations)
+   - Save settings
+
+4. **Wait for Build**
+   - Space will automatically build and deploy
+   - First build takes 5-10 minutes
+   - Model will download on first run (~7GB)
+
+### Method 2: Git Push (For Developers)
+
+1. **Create Space on HuggingFace**
+   - Go to https://huggingface.co/new-space
+   - Create space as above
+
+2. **Clone and Push**
+   ```powershell
+   # Navigate to project
+   cd d:\2025-vibe-coding\ACE-Step-Custom
+
+   # Initialize git (if not already)
+   git init
+   git add .
+   git commit -m "Initial commit"
+
+   # Add HuggingFace remote
+   git remote add hf https://huggingface.co/spaces/YOUR_USERNAME/ace-step-custom
+
+   # Push to HuggingFace
+   git push hf main
+   ```
+
+3. **Configure Git LFS for Large Files**
+   ```powershell
+   git lfs install
+   git lfs track "*.wav"
+   git lfs track "*.pth"
+   git lfs track "*.bin"
+   git lfs track "models/**"
+   git add .gitattributes
+   git commit -m "Add LFS tracking"
+   git push hf main
+   ```
+
+### Method 3: HuggingFace CLI (Fastest)
+
+1. **Install HuggingFace CLI**
+   ```powershell
+   pip install huggingface_hub
+   ```
+
+2. **Login**
+   ```powershell
+   huggingface-cli login
+   # Enter your HuggingFace token
+   ```
+
+3. **Create and Upload**
+   ```powershell
+   cd d:\2025-vibe-coding\ACE-Step-Custom
+   
+   # Create space
+   huggingface-cli repo create ace-step-custom --type space --space_sdk gradio
+   
+   # Upload files
+   huggingface-cli upload YOUR_USERNAME/ace-step-custom . --repo-type space
+   ```
+
+## Space Configuration
+
+### Hardware Recommendations
+
+| GPU | VRAM | Cost | Performance | Recommended For |
+|-----|------|------|-------------|-----------------|
+| CPU | - | Free | Very Slow | Testing only |
+| T4 Small | 16GB | ~$0.60/hr | Slow | Light testing |
+| **A10G Small** | **24GB** | **~$1.05/hr** | **Good** | **Recommended** |
+| A10G Large | 24GB | ~$3.15/hr | Good | Production |
+| A100 Large | 40GB | ~$4.13/hr | Excellent | Best quality |
+
+**Recommendation:** Start with A10G Small for testing, upgrade to A100 for production.
+
+### Environment Variables (Optional)
+
+In Space settings, you can add:
+
+```
+GRADIO_SERVER_NAME=0.0.0.0
+GRADIO_SERVER_PORT=7860
+HF_HOME=/data/huggingface
+TORCH_HOME=/data/torch
+```
+
+### Secrets (If Needed)
+
+For API keys or sensitive data:
+- Go to Space Settings → Repository secrets
+- Add secrets like `HF_TOKEN`, `API_KEY`, etc.
+- Access in code: `os.environ.get("SECRET_NAME")`
+
+## Post-Deployment Setup
+
+### First Launch
+
+1. **Wait for Model Download**
+   - First launch downloads ACE-Step model (~7GB)
+   - Takes 5-10 minutes depending on connection
+   - Model cached for subsequent runs
+
+2. **Test Basic Generation**
+   - Go to Tab 1 (Standard ACE-Step)
+   - Enter simple prompt: "Happy pop song"
+   - Set duration to 10 seconds
+   - Click Generate
+
+3. **Test Timeline**
+   - Go to Tab 2 (Timeline Workflow)
+   - Enter prompt and lyrics
+   - Set context length to 30s
+   - Generate first clip
+
+4. **Test LoRA Training**
+   - Go to Tab 3 (LoRA Training)
+   - Upload 2-3 test audio files
+   - Run quick training (2-3 epochs)
+
+### Monitoring
+
+**View Logs:**
+- Click "Logs" tab in your Space
+- Monitor for errors or warnings
+- Check GPU usage and memory
+
+**Performance Metrics:**
+- Generation time
+- Memory usage
+- Error rate
+- User feedback
+
+### Troubleshooting
+
+**Space Not Building:**
+- Check requirements.txt for conflicts
+- Verify Python 3.10 compatibility
+- Check logs for specific errors
+
+**Out of Memory:**
+- Upgrade to larger GPU
+- Reduce batch size in LoRA training
+- Limit generation duration
+
+**Model Not Loading:**
+- Check HuggingFace Hub access
+- Verify model ID in config.yaml
+- Check internet connectivity
+
+**Slow Performance:**
+- Upgrade GPU tier
+- Reduce concurrent users
+- Optimize generation parameters
+
+## Optimization Tips
+
+### Reduce Startup Time
+
+1. **Cache Models**
+   ```python
+   # In app.py, add before model loading:
+   os.environ["HF_HOME"] = "/data/huggingface"
+   ```
+
+2. **Preload on Startup**
+   - Models download on first run
+   - Cached for subsequent uses
+   - Consider pre-downloading to Space
+
+### Improve Response Time
+
+1. **Use Queuing**
+   - Gradio automatically queues requests
+   - Set `max_size` in `app.launch()`
+
+2. **Optimize Generation**
+   - Lower default duration
+   - Reduce sampling steps
+   - Use FP16 precision
+
+### Cost Optimization
+
+1. **Auto-Sleep**
+   - Space sleeps after inactivity
+   - Wakes on first request
+   - Configure in Space settings
+
+2. **Usage Limits**
+   - Set max concurrent users
+   - Limit generation duration
+   - Add rate limiting if needed
+
+## Going Live
+
+### Before Public Release
+
+- [ ] Test all three tabs thoroughly
+- [ ] Verify LoRA training works
+- [ ] Test with different prompts and styles
+- [ ] Check error handling
+- [ ] Review logs for issues
+- [ ] Test on mobile devices
+- [ ] Add usage examples
+- [ ] Create demo video
+
+### Public Space Settings
+
+1. **Enable Discussions**
+   - Let users report issues
+   - Gather feedback
+
+2. **Add Examples**
+   - Create example prompts
+   - Show best practices
+   - Include sample outputs
+
+3. **Update README**
+   - Clear usage instructions
+   - Feature highlights
+   - Limitations and known issues
+
+4. **Pin Space**
+   - Makes it discoverable
+   - Shows on your profile
+
+## Maintenance
+
+### Regular Updates
+
+```powershell
+# Update code
+cd d:\2025-vibe-coding\ACE-Step-Custom
+git add .
+git commit -m "Update description"
+git push hf main
+```
+
+### Monitor Usage
+
+- Check Space analytics
+- Review user feedback
+- Monitor error rates
+- Track popular features
+
+### Scaling
+
+**If Space Gets Popular:**
+1. Upgrade GPU tier
+2. Add request queuing
+3. Implement caching
+4. Consider duplicate Spaces for load balancing
+
+## Support & Community
+
+### Get Help
+
+- HuggingFace Forums: https://discuss.huggingface.co/
+- Discord: https://discord.gg/huggingface
+- Docs: https://huggingface.co/docs/hub/spaces
+
+### Share Your Space
+
+- Post on Twitter/X with #HuggingFace #ACEStep
+- Share in AI music communities
+- Add to your portfolio
+- Write blog post about it
+
+## Advanced Configuration
+
+### Custom Domain (Pro)
+
+HuggingFace Pro users can set custom domains:
+1. Go to Space settings
+2. Add custom domain
+3. Configure DNS
+
+### Persistent Storage
+
+For saving user data:
+```python
+import os
+PERSIST_DIR = os.environ.get("SPACE_ID", "local")
+# Save to /data/{SPACE_ID}/
+```
+
+### Analytics Integration
+
+Add Google Analytics or similar:
+```python
+# In app.py
+analytics_code = """
+<script async src="https://www.googletagmanager.com/gtag/js?id=YOUR-ID"></script>
+<script>
+  window.dataLayer = window.dataLayer || [];
+  function gtag(){dataLayer.push(arguments);}
+  gtag('js', new Date());
+  gtag('config', 'YOUR-ID');
+</script>
+"""
+```
+
+## Success Checklist
+
+Before announcing your Space:
+
+- ✅ All features working
+- ✅ Clear documentation
+- ✅ Example outputs included
+- ✅ Error handling robust
+- ✅ Performance optimized
+- ✅ Mobile-friendly UI
+- ✅ Clear limitations stated
+- ✅ License properly attributed
+- ✅ Usage guidelines clear
+- ✅ Contact/support info provided
+
+## Your Space URL
+
+After deployment, your Space will be available at:
+```
+https://huggingface.co/spaces/YOUR_USERNAME/ace-step-custom
+```
+
+Share it with the world! 🎵🚀

DEPLOYMENT_CHECKLIST.txt

@@ -0,0 +1,226 @@
+⚠️ IMPORTANT: Follow these steps in order ⚠️
+
+═══════════════════════════════════════════════════════════════
+🚀 HuggingFace Spaces Deployment - Step by Step
+═══════════════════════════════════════════════════════════════
+
+📋 PREREQUISITES
+═══════════════════════════════════════════════════════════════
+☐ HuggingFace account created: https://huggingface.co/join
+☐ HuggingFace token obtained: https://huggingface.co/settings/tokens
+   (Create new token with "write" access)
+☐ HuggingFace CLI installed (already done ✓)
+
+═══════════════════════════════════════════════════════════════
+🎯 DEPLOYMENT STEPS
+═══════════════════════════════════════════════════════════════
+
+Choose ONE method below:
+
+┌─────────────────────────────────────────────────────────────┐
+│ METHOD 1: AUTOMATED SCRIPT (EASIEST) ⭐                     │
+└─────────────────────────────────────────────────────────────┘
+
+1. Open PowerShell in this directory:
+   d:\2025-vibe-coding\ACE-Step-Custom
+
+2. Run the deployment script:
+   .\deploy_hf.bat
+
+3. Follow the prompts:
+   - Login with your HF token
+   - Enter Space name (e.g., "ace-step-custom")
+   - Wait for upload
+
+┌─────────────────────────────────────────────────────────────┐
+│ METHOD 2: MANUAL CLI (FOR DEVELOPERS)                       │
+└─────────────────────────────────────────────────────────────┘
+
+1. Login to HuggingFace:
+   huggingface-cli login
+   [Paste your token]
+
+2. Create the Space:
+   huggingface-cli repo create ace-step-custom --type space --space_sdk gradio
+
+3. Upload files:
+   huggingface-cli upload YOUR_USERNAME/ace-step-custom . --repo-type space
+
+┌─────────────────────────────────────────────────────────────┐
+│ METHOD 3: WEB INTERFACE (NO CLI NEEDED)                     │
+└─────────────────────────────────────────────────────────────┘
+
+1. Go to: https://huggingface.co/new-space
+
+2. Fill in Space details:
+   Name: ace-step-custom
+   License: MIT
+   SDK: Gradio
+   Hardware: A10G Small
+
+3. Click "Create Space"
+
+4. Click "Files and versions" → "Add file" → "Upload files"
+
+5. Upload these files/folders:
+   ✓ app.py
+   ✓ requirements.txt
+   ✓ config.yaml
+   ✓ README.md
+   ✓ LICENSE
+   ✓ .gitignore
+   ✓ src/ (entire folder)
+   ✓ scripts/ (entire folder)
+
+6. Commit changes
+
+═══════════════════════════════════════════════════════════════
+⚙️ POST-DEPLOYMENT CONFIGURATION
+═══════════════════════════════════════════════════════════════
+
+After upload, configure your Space:
+
+1. ☐ Go to your Space URL:
+   https://huggingface.co/spaces/YOUR_USERNAME/ace-step-custom
+
+2. ☐ Click "Settings" tab
+
+3. ☐ Configure Hardware:
+   - Select: "A10G Small" (24GB VRAM) - MINIMUM
+   - Or: "A100 Large" (40GB VRAM) - RECOMMENDED
+   - Click "Save"
+
+4. ☐ Set Python version: 3.10 (should be automatic)
+
+5. ☐ Set timeout: 30 minutes (optional, for long generations)
+
+6. ☐ Enable Discussions (optional, for user feedback)
+
+═══════════════════════════════════════════════════════════════
+⏱️ BUILD & TESTING
+═══════════════════════════════════════════════════════════════
+
+1. ☐ Wait for build to complete:
+   - Click "Logs" tab to monitor
+   - First build: 5-10 minutes
+   - Model download: ~7GB (first run only)
+
+2. ☐ Space will show "Running" when ready
+
+3. ☐ Test Tab 1 (Standard ACE-Step):
+   - Enter prompt: "Happy pop song with piano"
+   - Set duration: 10 seconds
+   - Click "Generate"
+   - Verify audio plays
+
+4. ☐ Test Tab 2 (Timeline Workflow):
+   - Enter prompt and lyrics
+   - Set context length: 30 seconds
+   - Click "Generate Clip"
+   - Verify timeline updates
+
+5. ☐ Test Tab 3 (LoRA Training):
+   - Upload 2-3 test audio files
+   - Set epochs to 2
+   - Click "Start Training"
+   - Verify progress updates
+
+═══════════════════════════════════════════════════════════════
+💰 COST MANAGEMENT
+═══════════════════════════════════════════════════════════════
+
+GPU Costs:
+- A10G Small (24GB): ~$1.05/hour ⭐ RECOMMENDED
+- A100 Large (40GB): ~$4.13/hour
+
+Auto-Sleep:
+✓ Space sleeps automatically after 48 hours of inactivity
+✓ Wakes up on first request (30-60 second startup)
+✓ No charges while sleeping
+
+Testing Budget:
+- Initial testing: ~$5-10
+- Active use: ~$10-50/month
+- Production: Scale as needed
+
+═══════════════════════════════════════════════════════════════
+🐛 TROUBLESHOOTING
+═══════════════════════════════════════════════════════════════
+
+Problem: Space won't start
+Solution: 
+- Check "Logs" tab for errors
+- Verify all files uploaded
+- Ensure README.md has YAML frontmatter
+
+Problem: Out of memory error
+Solution:
+- Upgrade to A100 Large
+- Reduce generation duration in UI
+- Lower batch size in LoRA training
+
+Problem: Slow generation
+Solution:
+- Verify GPU is enabled (not CPU)
+- Check Space isn't sleeping
+- Reduce sampling steps in config
+
+Problem: Model download fails
+Solution:
+- Check HuggingFace Hub status
+- Verify internet connectivity
+- Wait and retry
+
+═══════════════════════════════════════════════════════════════
+✅ SUCCESS CHECKLIST
+═══════════════════════════════════════════════════════════════
+
+Before announcing your Space:
+
+☐ All three tabs tested and working
+☐ Example generations added to README
+☐ Clear usage instructions visible
+☐ GPU enabled (A10G Small minimum)
+☐ Error handling tested
+☐ Mobile view checked
+☐ Discussions enabled
+☐ License properly displayed
+☐ Contact/support info added
+☐ Share link works
+
+═══════════════════════════════════════════════════════════════
+🎉 GO LIVE!
+═══════════════════════════════════════════════════════════════
+
+Your Space URL:
+https://huggingface.co/spaces/YOUR_USERNAME/ace-step-custom
+
+Share it:
+□ Twitter/X: "Just deployed ACE-Step 1.5 Custom on @huggingface! 🎵 
+   Check it out: [your-url] #AIMusic #HuggingFace #ACEStep"
+□ LinkedIn post
+□ Reddit (r/MachineLearning, r/artificial, r/WeAreTheMusicMakers)
+□ Discord communities
+□ Personal blog/portfolio
+
+═══════════════════════════════════════════════════════════════
+📚 ADDITIONAL RESOURCES
+═══════════════════════════════════════════════════════════════
+
+Documentation:
+- DEPLOY_QUICK.md - Quick reference
+- DEPLOYMENT.md - Complete guide
+- README.md - Project documentation
+
+Support:
+- HuggingFace Docs: https://huggingface.co/docs/hub/spaces
+- HuggingFace Discord: https://discord.gg/huggingface
+- GitHub Issues: [your-repo-url]
+
+═══════════════════════════════════════════════════════════════
+
+Ready to deploy? 🚀
+
+Run: .\deploy_hf.bat
+
+═══════════════════════════════════════════════════════════════

DEPLOY_QUICK.md

@@ -0,0 +1,126 @@
+# Quick Deployment to HuggingFace Spaces
+
+## Prerequisites
+✅ HuggingFace account: https://huggingface.co/join
+✅ HuggingFace token: https://huggingface.co/settings/tokens
+
+## Fastest Method (Windows)
+
+Run the deployment script:
+
+```powershell
+cd d:\2025-vibe-coding\ACE-Step-Custom
+.\deploy_hf.bat
+```
+
+The script will:
+1. Install HuggingFace CLI (if needed)
+2. Login to your account
+3. Create new Space
+4. Upload all files
+5. Provide your Space URL
+
+## Fastest Method (Linux/Mac)
+
+```bash
+cd /path/to/ACE-Step-Custom
+chmod +x deploy_hf.sh
+./deploy_hf.sh
+```
+
+## Manual Deployment (If Script Fails)
+
+### 1. Install HuggingFace CLI
+```powershell
+pip install huggingface_hub
+```
+
+### 2. Login
+```powershell
+huggingface-cli login
+```
+Enter your token from: https://huggingface.co/settings/tokens
+
+### 3. Create Space
+```powershell
+huggingface-cli repo create ace-step-custom --type space --space_sdk gradio
+```
+
+### 4. Upload Files
+```powershell
+cd d:\2025-vibe-coding\ACE-Step-Custom
+huggingface-cli upload YOUR_USERNAME/ace-step-custom . --repo-type space
+```
+
+Replace `YOUR_USERNAME` with your HuggingFace username.
+
+## After Upload
+
+### 1. Configure GPU
+- Go to your Space: https://huggingface.co/spaces/YOUR_USERNAME/ace-step-custom
+- Click "Settings" tab
+- Under "Hardware", select: **A10G Small** (recommended)
+- Click "Save"
+
+### 2. Wait for Build
+- Space will build automatically (5-10 minutes)
+- Check "Logs" tab for progress
+- Model downloads on first run (~7GB)
+
+### 3. Test Your Space
+1. Open Space URL
+2. Test Tab 1: Generate 10-second clip
+3. Test Tab 2: Generate timeline clip
+4. Test Tab 3: Upload test audio
+
+## Troubleshooting
+
+**Login Failed:**
+```powershell
+# Make sure you copied the full token
+huggingface-cli whoami  # Check if logged in
+```
+
+**Upload Failed:**
+```powershell
+# Try with explicit exclusions
+huggingface-cli upload YOUR_USERNAME/ace-step-custom . --repo-type space --exclude "*.pyc" --exclude "outputs/*" --exclude "__pycache__/*"
+```
+
+**Space Not Starting:**
+- Check "Logs" tab for errors
+- Verify requirements.txt is uploaded
+- Ensure README.md has correct YAML frontmatter
+
+**Out of Memory:**
+- Upgrade GPU in Settings
+- Start with A10G Small minimum
+
+## Your Space URL
+
+After deployment:
+```
+https://huggingface.co/spaces/YOUR_USERNAME/ace-step-custom
+```
+
+## Cost Estimate
+
+- **A10G Small (24GB):** ~$1.05/hour
+- **Auto-sleep:** Space sleeps when inactive (no charge)
+- **Testing:** Budget ~$5-10 for initial testing
+
+## Need Help?
+
+See full guide: [DEPLOYMENT.md](DEPLOYMENT.md)
+
+## Next Steps
+
+1. ✅ Deploy Space
+2. ✅ Test all features
+3. ✅ Enable Discussions in Settings
+4. ✅ Add example outputs to README
+5. ✅ Share your Space!
+
+---
+
+🎵 Happy testing! Your Space will be live in minutes! 🚀

Dockerfile

@@ -0,0 +1,35 @@
+FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
+
+# Set working directory
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    python3.10 \
+    python3-pip \
+    git \
+    ffmpeg \
+    libsndfile1 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip3 install --no-cache-dir -r requirements.txt
+
+# Copy application files
+COPY . .
+
+# Create necessary directories
+RUN mkdir -p outputs timelines lora_training logs models
+
+# Expose Gradio port
+EXPOSE 7860
+
+# Set environment variables
+ENV GRADIO_SERVER_NAME="0.0.0.0"
+ENV GRADIO_SERVER_PORT=7860
+
+# Run the application
+CMD ["python3", "app.py"]

LICENSE

@@ -0,0 +1,28 @@
+MIT License
+
+Copyright (c) 2026 Gamahea Development Team
+
+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.
+
+---
+
+This project uses ACE-Step, which is subject to its own license:
+https://github.com/ace-step/ACE-Step
+
+Please refer to the original ACE-Step repository for their licensing terms.

QUICKSTART.md

@@ -0,0 +1,115 @@
+# ACE-Step 1.5 Custom Edition - Quick Start Guide
+
+## Installation
+
+### Option 1: Local Setup
+
+1. **Clone the repository**
+```bash
+git clone https://github.com/yourusername/ace-step-custom.git
+cd ace-step-custom
+```
+
+2. **Create virtual environment**
+```bash
+python -m venv venv
+
+# On Windows:
+venv\Scripts\activate
+
+# On Linux/Mac:
+source venv/bin/activate
+```
+
+3. **Run setup**
+```bash
+python scripts/setup.py
+```
+
+4. **Download model**
+```bash
+python scripts/download_model.py
+```
+
+5. **Launch application**
+```bash
+python app.py
+```
+
+6. **Open browser to** `http://localhost:7860`
+
+### Option 2: HuggingFace Spaces
+
+1. Create new Space on HuggingFace
+2. Upload all project files
+3. Set Space configuration:
+   - SDK: `gradio`
+   - Python: `3.10`
+   - GPU: `A10G` (or better)
+4. Space will auto-deploy
+
+## Usage
+
+### Tab 1: Standard ACE-Step
+
+Standard interface with all original ACE-Step features:
+- Text-to-music generation
+- Variation generation
+- Repainting sections
+- Lyric editing
+
+### Tab 2: Timeline Workflow
+
+Advanced timeline-based generation:
+1. Enter prompt and lyrics
+2. Set context length (0-120s)
+3. Click "Generate" for 32s clips
+4. Clips auto-blend into timeline
+5. Use "Extend" to continue
+6. Use "Inpaint" to edit regions
+
+### Tab 3: LoRA Training
+
+Train custom models:
+1. Upload audio files (10+ recommended)
+2. Set training parameters
+3. Click "Start Training"
+4. Download trained model
+5. Use in Tab 1 or Tab 2
+
+## Tips
+
+- **First time:** Start with Standard tab to understand basics
+- **For longer songs:** Use Timeline tab with context length 30-60s
+- **For custom styles:** Train LoRA with 20+ similar audio files
+- **GPU recommended:** 8GB+ VRAM for best performance
+- **CPU mode:** Works but slower, use shorter durations
+
+## Troubleshooting
+
+### Out of Memory
+- Reduce batch size in LoRA training
+- Use shorter audio durations
+- Close other GPU applications
+
+### Poor Quality
+- Increase context length
+- Try different seeds
+- Adjust temperature (0.6-0.8 is usually good)
+
+### Blend Artifacts
+- Reduce lead-in/lead-out durations
+- Ensure consistent style across clips
+- Use lower context length for more variety
+
+## Support
+
+- GitHub Issues: [Report bugs here]
+- Documentation: See `docs/` directory
+- Examples: See `examples/` directory
+
+## Credits
+
+Based on ACE-Step by ACE Studio and Step Fun
+- Website: https://ace-step.github.io/
+- Paper: https://arxiv.org/abs/2506.00045

README.md

@@ -0,0 +1,73 @@
+---
+title: ACE-Step 1.5 Custom Edition
+emoji: 🎵
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+sdk_version: 4.0.0
+app_file: app.py
+pinned: false
+license: mit
+python_version: 3.11
+---
+
+# ACE-Step 1.5 Custom Edition
+
+A comprehensive music generation system built on ACE-Step 1.5, featuring:
+
+## 🌟 Features
+
+### 1. Standard ACE-Step Interface
+Full-featured standard ACE-Step 1.5 GUI with all original capabilities including:
+- Text-to-music generation with style control
+- Variation generation
+- Section repainting
+- Lyric editing
+
+### 2. Custom Timeline Workflow
+Advanced timeline-based generation system:
+- Generate 32-second clips with seamless blending
+- Adjustable context length (0-120 seconds) for style consistency
+- Master timeline with visual representation
+- Extend, inpaint, and remix capabilities
+- Automatic crossfading between clips
+
+### 3. LoRA Training Studio
+Complete training interface for custom models:
+- Upload and preprocess audio files
+- Configure training parameters
+- Train specialized models for voices, instruments, or styles
+- Download and reuse trained models
+- Continue training from existing LoRAs
+
+## 🚀 Quick Start
+
+1. **Standard Generation**: Use Tab 1 for traditional text-to-music
+2. **Timeline Creation**: Use Tab 2 to build longer songs with consistent style
+3. **Custom Training**: Use Tab 3 to create specialized models
+
+## 💡 Tips
+
+- Start with context length of 30-60s for best results
+- For custom voices, train LoRA with 20+ audio samples
+- Adjust temperature between 0.6-0.8 for quality vs creativity
+- Use "Extend" in Timeline mode to continue your song
+
+## 🎯 Use Cases
+
+- **Musicians**: Create backing tracks and song ideas
+- **Content Creators**: Generate royalty-free music for videos
+- **Game Developers**: Create adaptive game soundtracks
+- **AI Researchers**: Experiment with music generation and LoRA training
+
+## 📚 Documentation
+
+See the repository for full documentation and examples.
+
+## 🙏 Credits
+
+Built on top of [ACE-Step](https://ace-step.github.io/) by ACE Studio and Step Fun.
+
+## ⚠️ Note
+
+This is a custom implementation focusing on enhanced workflows and training capabilities. Generation quality depends on the base ACE-Step model and your usage patterns.

README_HF.md

@@ -0,0 +1,73 @@
+---
+title: ACE-Step 1.5 Custom Edition
+emoji: 🎵
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+sdk_version: 4.0.0
+app_file: app.py
+pinned: false
+license: mit
+python_version: 3.11
+---
+
+# ACE-Step 1.5 Custom Edition
+
+A comprehensive music generation system built on ACE-Step 1.5, featuring:
+
+## 🌟 Features
+
+### 1. Standard ACE-Step Interface
+Full-featured standard ACE-Step 1.5 GUI with all original capabilities including:
+- Text-to-music generation with style control
+- Variation generation
+- Section repainting
+- Lyric editing
+
+### 2. Custom Timeline Workflow
+Advanced timeline-based generation system:
+- Generate 32-second clips with seamless blending
+- Adjustable context length (0-120 seconds) for style consistency
+- Master timeline with visual representation
+- Extend, inpaint, and remix capabilities
+- Automatic crossfading between clips
+
+### 3. LoRA Training Studio
+Complete training interface for custom models:
+- Upload and preprocess audio files
+- Configure training parameters
+- Train specialized models for voices, instruments, or styles
+- Download and reuse trained models
+- Continue training from existing LoRAs
+
+## 🚀 Quick Start
+
+1. **Standard Generation**: Use Tab 1 for traditional text-to-music
+2. **Timeline Creation**: Use Tab 2 to build longer songs with consistent style
+3. **Custom Training**: Use Tab 3 to create specialized models
+
+## 💡 Tips
+
+- Start with context length of 30-60s for best results
+- For custom voices, train LoRA with 20+ audio samples
+- Adjust temperature between 0.6-0.8 for quality vs creativity
+- Use "Extend" in Timeline mode to continue your song
+
+## 🎯 Use Cases
+
+- **Musicians**: Create backing tracks and song ideas
+- **Content Creators**: Generate royalty-free music for videos
+- **Game Developers**: Create adaptive game soundtracks
+- **AI Researchers**: Experiment with music generation and LoRA training
+
+## 📚 Documentation
+
+See the repository for full documentation and examples.
+
+## 🙏 Credits
+
+Built on top of [ACE-Step](https://ace-step.github.io/) by ACE Studio and Step Fun.
+
+## ⚠️ Note
+
+This is a custom implementation focusing on enhanced workflows and training capabilities. Generation quality depends on the base ACE-Step model and your usage patterns.

README_PROJECT.md

@@ -0,0 +1,116 @@
+---
+title: ACE-Step 1.5 Custom Edition
+emoji: 🎵
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+sdk_version: 5.9.1
+app_file: app.py
+pinned: false
+license: mit
+python_version: "3.11"
+hardware: zero-gpu-medium
+---
+
+# ACE-Step 1.5 Custom Edition
+
+A fully-featured implementation of ACE-Step 1.5 with custom GUI and workflow capabilities for local use and HuggingFace Space deployment.
+
+## Features
+
+### 🎵 Three Main Interfaces
+
+1. **Standard ACE-Step GUI**: Full-featured standard ACE-Step 1.5 interface with all original capabilities
+2. **Custom Timeline Workflow**: Advanced timeline-based generation with:
+   - 32-second clip generation (2s lead-in + 28s main + 2s lead-out)
+   - Seamless clip blending for continuous music
+   - Context Length slider (0-120 seconds) for style guidance
+   - Master timeline with extend, inpaint, and remix capabilities
+3. **LoRA Training Studio**: Complete LoRA training interface with:
+   - Audio file upload and preprocessing
+   - Custom training configuration
+   - Model download/upload for continued training
+
+## Architecture
+
+- **Base Model**: ACE-Step v1.5 Turbo
+- **Framework**: Gradio 5.9.1, PyTorch
+- **Deployment**: Local execution + HuggingFace Spaces
+- **Audio Processing**: DiT + VAE + 5Hz Language Model
+
+## Installation
+
+### Local Setup
+
+```bash
+# Clone the repository
+git clone https://github.com/yourusername/ace-step-custom.git
+cd ace-step-custom
+
+# Create virtual environment
+python -m venv venv
+source venv/bin/activate  # On Windows: venv\Scripts\activate
+
+# Install dependencies
+pip install -r requirements.txt
+
+# Download ACE-Step model
+python scripts/download_model.py
+
+# Run the application
+python app.py
+```
+
+### HuggingFace Space Deployment
+
+1. Create a new Space on HuggingFace
+2. Upload all files to the Space
+3. Set Space to use GPU (recommended: H200 or A100)
+4. The app will automatically download models and start
+
+## Usage
+
+### Standard Mode
+Use the first tab for standard ACE-Step generation with all original features.
+
+### Timeline Mode
+1. Enter your prompt/lyrics
+2. Adjust Context Length (how far back to reference previous clips)
+3. Click "Generate" to create 32-second clips
+4. Clips automatically blend and add to timeline
+5. Use "Extend" to continue the song or other options for variations
+
+### LoRA Training
+1. Upload audio files for training
+2. Configure training parameters
+3. Train custom LoRA models
+4. Download and reuse for continued training
+
+## System Requirements
+
+### Minimum
+- GPU: 8GB VRAM (with optimizations)
+- RAM: 16GB
+- Storage: 20GB
+
+### Recommended
+- GPU: 16GB+ VRAM (A100, H200, or consumer GPUs)
+- RAM: 32GB
+- Storage: 50GB
+
+## Technical Details
+
+- **Audio Format**: 48kHz, stereo
+- **Generation Speed**: ~8 inference steps (turbo model)
+- **Context Window**: Up to 120 seconds for style guidance
+- **Blend Regions**: 2-second crossfade between clips
+
+## Credits
+
+Based on ACE-Step 1.5 by ACE Studio
+- GitHub: https://github.com/ace-step/ACE-Step-1.5
+- Original Demo: https://huggingface.co/spaces/ACE-Step/ACE-Step
+
+## License
+
+MIT License (see LICENSE file)

acestep/__init__.py

@@ -0,0 +1 @@
+"""ACE-Step package."""

acestep/acestep_v15_pipeline.py

@@ -0,0 +1,411 @@
+"""
+ACE-Step V1.5 Pipeline
+Handler wrapper connecting model and UI
+"""
+import os
+import sys
+
+# Load environment variables from .env file at most once per process to avoid
+# epoch-boundary stalls (e.g. on Windows when Gradio yields during training)
+_env_loaded = False  # module-level so we never reload .env in the same process
+try:
+    from dotenv import load_dotenv
+    if not _env_loaded:
+        _current_file = os.path.abspath(__file__)
+        _project_root = os.path.dirname(os.path.dirname(_current_file))
+        _env_path = os.path.join(_project_root, '.env')
+        _env_example_path = os.path.join(_project_root, '.env.example')
+        if os.path.exists(_env_path):
+            load_dotenv(_env_path)
+            print(f"Loaded configuration from {_env_path}")
+        elif os.path.exists(_env_example_path):
+            load_dotenv(_env_example_path)
+            print(f"Loaded configuration from {_env_example_path} (fallback)")
+        _env_loaded = True
+except ImportError:
+    # python-dotenv not installed, skip loading .env
+    pass
+
+# Clear proxy settings that may affect Gradio
+for proxy_var in ['http_proxy', 'https_proxy', 'HTTP_PROXY', 'HTTPS_PROXY', 'ALL_PROXY']:
+    os.environ.pop(proxy_var, None)
+
+try:
+    # When executed as a module: `python -m acestep.acestep_v15_pipeline`
+    from .handler import AceStepHandler
+    from .llm_inference import LLMHandler
+    from .dataset_handler import DatasetHandler
+    from .gradio_ui import create_gradio_interface
+    from .gpu_config import get_gpu_config, get_gpu_memory_gb, print_gpu_config_info, set_global_gpu_config, VRAM_16GB_MIN_GB
+except ImportError:
+    # When executed as a script: `python acestep/acestep_v15_pipeline.py`
+    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    if project_root not in sys.path:
+        sys.path.insert(0, project_root)
+    from acestep.handler import AceStepHandler
+    from acestep.llm_inference import LLMHandler
+    from acestep.dataset_handler import DatasetHandler
+    from acestep.gradio_ui import create_gradio_interface
+    from acestep.gpu_config import get_gpu_config, get_gpu_memory_gb, print_gpu_config_info, set_global_gpu_config, VRAM_16GB_MIN_GB
+
+
+def create_demo(init_params=None, language='en'):
+    """
+    Create Gradio demo interface
+    
+    Args:
+        init_params: Dictionary containing initialization parameters and state.
+                    If None, service will not be pre-initialized.
+                    Keys: 'pre_initialized' (bool), 'checkpoint', 'config_path', 'device',
+                          'init_llm', 'lm_model_path', 'backend', 'use_flash_attention',
+                          'offload_to_cpu', 'offload_dit_to_cpu', 'init_status',
+                          'dit_handler', 'llm_handler' (initialized handlers if pre-initialized),
+                          'language' (UI language code)
+        language: UI language code ('en', 'zh', 'ja', default: 'en')
+    
+    Returns:
+        Gradio Blocks instance
+    """
+    # Use pre-initialized handlers if available, otherwise create new ones
+    if init_params and init_params.get('pre_initialized') and 'dit_handler' in init_params:
+        dit_handler = init_params['dit_handler']
+        llm_handler = init_params['llm_handler']
+    else:
+        dit_handler = AceStepHandler()  # DiT handler
+        llm_handler = LLMHandler()      # LM handler
+    
+    dataset_handler = DatasetHandler()  # Dataset handler
+    
+    # Create Gradio interface with all handlers and initialization parameters
+    demo = create_gradio_interface(dit_handler, llm_handler, dataset_handler, init_params=init_params, language=language)
+    
+    return demo
+
+
+def main():
+    """Main entry function"""
+    import argparse
+    
+    # Detect GPU memory and get configuration
+    gpu_config = get_gpu_config()
+    set_global_gpu_config(gpu_config)  # Set global config for use across modules
+    
+    gpu_memory_gb = gpu_config.gpu_memory_gb
+    auto_offload = gpu_memory_gb > 0 and gpu_memory_gb < VRAM_16GB_MIN_GB
+    
+    # Print GPU configuration info
+    print(f"\n{'='*60}")
+    print("GPU Configuration Detected:")
+    print(f"{'='*60}")
+    print(f"  GPU Memory: {gpu_memory_gb:.2f} GB")
+    print(f"  Configuration Tier: {gpu_config.tier}")
+    print(f"  Max Duration (with LM): {gpu_config.max_duration_with_lm}s ({gpu_config.max_duration_with_lm // 60} min)")
+    print(f"  Max Duration (without LM): {gpu_config.max_duration_without_lm}s ({gpu_config.max_duration_without_lm // 60} min)")
+    print(f"  Max Batch Size (with LM): {gpu_config.max_batch_size_with_lm}")
+    print(f"  Max Batch Size (without LM): {gpu_config.max_batch_size_without_lm}")
+    print(f"  Default LM Init: {gpu_config.init_lm_default}")
+    print(f"  Available LM Models: {gpu_config.available_lm_models or 'None'}")
+    print(f"{'='*60}\n")
+    
+    if auto_offload:
+        print(f"Auto-enabling CPU offload (GPU < 16GB)")
+    elif gpu_memory_gb > 0:
+        print(f"CPU offload disabled by default (GPU >= 16GB)")
+    else:
+        print("No GPU detected, running on CPU")
+
+    # Define local outputs directory
+    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    output_dir = os.path.join(project_root, "gradio_outputs")
+    # Normalize path to use forward slashes for Gradio 6 compatibility on Windows
+    output_dir = output_dir.replace("\\", "/")
+    os.makedirs(output_dir, exist_ok=True)
+    print(f"Output directory: {output_dir}")
+
+    parser = argparse.ArgumentParser(description="Gradio Demo for ACE-Step V1.5")
+    parser.add_argument("--port", type=int, default=7860, help="Port to run the gradio server on")
+    parser.add_argument("--share", action="store_true", help="Create a public link")
+    parser.add_argument("--debug", action="store_true", help="Enable debug mode")
+    parser.add_argument("--server-name", type=str, default="127.0.0.1", help="Server name (default: 127.0.0.1, use 0.0.0.0 for all interfaces)")
+    parser.add_argument("--language", type=str, default="en", choices=["en", "zh", "he", "ja"], help="UI language: en (English), zh (中文), he (עברית), ja (日本語)")
+    parser.add_argument(
+        "--allowed-path",
+        action="append",
+        default=[],
+        help="Additional allowed file paths for Gradio (repeatable).",
+    )
+    
+    # Service mode argument
+    parser.add_argument("--service_mode", type=lambda x: x.lower() in ['true', '1', 'yes'], default=False, 
+                       help="Enable service mode (default: False). When enabled, uses preset models and restricts UI options.")
+    
+    # Service initialization arguments
+    parser.add_argument("--init_service", type=lambda x: x.lower() in ['true', '1', 'yes'], default=False, help="Initialize service on startup (default: False)")
+    parser.add_argument("--checkpoint", type=str, default=None, help="Checkpoint file path (optional, for display purposes)")
+    parser.add_argument("--config_path", type=str, default=None, help="Main model path (e.g., 'acestep-v15-turbo')")
+    parser.add_argument("--device", type=str, default="auto", choices=["auto", "cuda", "mps", "xpu", "cpu"], help="Processing device (default: auto)")
+    parser.add_argument("--init_llm", type=lambda x: x.lower() in ['true', '1', 'yes'], default=None, help="Initialize 5Hz LM (default: auto based on GPU memory)")
+    parser.add_argument("--lm_model_path", type=str, default=None, help="5Hz LM model path (e.g., 'acestep-5Hz-lm-0.6B')")
+    parser.add_argument("--backend", type=str, default="vllm", choices=["vllm", "pt", "mlx"], help="5Hz LM backend (default: vllm, use 'mlx' for native Apple Silicon acceleration)")
+    parser.add_argument("--use_flash_attention", type=lambda x: x.lower() in ['true', '1', 'yes'], default=None, help="Use flash attention (default: auto-detect)")
+    parser.add_argument("--offload_to_cpu", type=lambda x: x.lower() in ['true', '1', 'yes'], default=auto_offload, help=f"Offload models to CPU (default: {'True' if auto_offload else 'False'}, auto-detected based on GPU VRAM)")
+    parser.add_argument("--offload_dit_to_cpu", type=lambda x: x.lower() in ['true', '1', 'yes'], default=False, help="Offload DiT to CPU (default: False)")
+    parser.add_argument("--download-source", type=str, default=None, choices=["huggingface", "modelscope", "auto"], help="Preferred model download source (default: auto-detect based on network)")
+
+    # API mode argument
+    parser.add_argument("--enable-api", action="store_true", help="Enable API endpoints (default: False)")
+
+    # Authentication arguments
+    parser.add_argument("--auth-username", type=str, default=None, help="Username for Gradio authentication")
+    parser.add_argument("--auth-password", type=str, default=None, help="Password for Gradio authentication")
+    parser.add_argument("--api-key", type=str, default=None, help="API key for API endpoints authentication")
+
+    args = parser.parse_args()
+
+    # Enable API requires init_service
+    if args.enable_api:
+        args.init_service = True
+        # Load config from .env if not specified
+        if args.config_path is None:
+            args.config_path = os.environ.get("ACESTEP_CONFIG_PATH")
+        if args.lm_model_path is None:
+            args.lm_model_path = os.environ.get("ACESTEP_LM_MODEL_PATH")
+        if os.environ.get("ACESTEP_LM_BACKEND"):
+            args.backend = os.environ.get("ACESTEP_LM_BACKEND")
+
+    # Service mode defaults (can be configured via .env file)
+    if args.service_mode:
+        print("Service mode enabled - applying preset configurations...")
+        # Force init_service in service mode
+        args.init_service = True
+        # Default DiT model for service mode (from env or fallback)
+        if args.config_path is None:
+            args.config_path = os.environ.get(
+                "SERVICE_MODE_DIT_MODEL",
+                "acestep-v15-turbo-fix-inst-shift-dynamic"
+            )
+        # Default LM model for service mode (from env or fallback)
+        if args.lm_model_path is None:
+            args.lm_model_path = os.environ.get(
+                "SERVICE_MODE_LM_MODEL",
+                "acestep-5Hz-lm-1.7B-v4-fix"
+            )
+        # Backend for service mode (from env or fallback to vllm)
+        args.backend = os.environ.get("SERVICE_MODE_BACKEND", "vllm")
+        print(f"  DiT model: {args.config_path}")
+        print(f"  LM model: {args.lm_model_path}")
+        print(f"  Backend: {args.backend}")
+    
+    # Auto-enable CPU offload for tier6 GPUs (16-24GB) when using the 4B LM model
+    # The 4B LM (~8GB) + DiT (~4.7GB) + VAE + text encoder exceeds 16-20GB with activations
+    if not args.offload_to_cpu and args.lm_model_path and "4B" in args.lm_model_path:
+        if 0 < gpu_memory_gb <= 24:
+            args.offload_to_cpu = True
+            print(f"Auto-enabling CPU offload (4B LM model requires offloading on {gpu_memory_gb:.0f}GB GPU)")
+
+    try:
+        init_params = None
+        dit_handler = None
+        llm_handler = None
+
+        # If init_service is True, perform initialization before creating UI
+        if args.init_service:
+            print("Initializing service from command line...")
+            
+            # Create handler instances for initialization
+            dit_handler = AceStepHandler()
+            llm_handler = LLMHandler()
+            
+            # Auto-select config_path if not provided
+            if args.config_path is None:
+                available_models = dit_handler.get_available_acestep_v15_models()
+                if available_models:
+                    args.config_path = "acestep-v15-turbo" if "acestep-v15-turbo" in available_models else available_models[0]
+                    print(f"Auto-selected config_path: {args.config_path}")
+                else:
+                    print("Error: No available models found. Please specify --config_path", file=sys.stderr)
+                    sys.exit(1)
+            
+            # Get project root (same logic as in handler)
+            current_file = os.path.abspath(__file__)
+            project_root = os.path.dirname(os.path.dirname(current_file))
+            
+            # Determine flash attention setting
+            use_flash_attention = args.use_flash_attention
+            if use_flash_attention is None:
+                use_flash_attention = dit_handler.is_flash_attention_available(args.device)
+
+            # Determine download source preference
+            prefer_source = None
+            if args.download_source and args.download_source != "auto":
+                prefer_source = args.download_source
+                print(f"Using preferred download source: {prefer_source}")
+
+            # Initialize DiT handler
+            print(f"Initializing DiT model: {args.config_path} on {args.device}...")
+            init_status, enable_generate = dit_handler.initialize_service(
+                project_root=project_root,
+                config_path=args.config_path,
+                device=args.device,
+                use_flash_attention=use_flash_attention,
+                compile_model=False,
+                offload_to_cpu=args.offload_to_cpu,
+                offload_dit_to_cpu=args.offload_dit_to_cpu,
+                prefer_source=prefer_source
+            )
+            
+            if not enable_generate:
+                print(f"Error initializing DiT model: {init_status}", file=sys.stderr)
+                sys.exit(1)
+            
+            print(f"DiT model initialized successfully")
+            
+            # Initialize LM handler if requested
+            # Auto-determine init_llm based on GPU config if not explicitly set
+            if args.init_llm is None:
+                args.init_llm = gpu_config.init_lm_default
+                print(f"Auto-setting init_llm to {args.init_llm} based on GPU configuration")
+            
+            lm_status = ""
+            if args.init_llm:
+                if args.lm_model_path is None:
+                    # Try to get default LM model
+                    available_lm_models = llm_handler.get_available_5hz_lm_models()
+                    if available_lm_models:
+                        args.lm_model_path = available_lm_models[0]
+                        print(f"Using default LM model: {args.lm_model_path}")
+                    else:
+                        print("Warning: No LM models available, skipping LM initialization", file=sys.stderr)
+                        args.init_llm = False
+                
+                if args.init_llm and args.lm_model_path:
+                    checkpoint_dir = os.path.join(project_root, "checkpoints")
+                    print(f"Initializing 5Hz LM: {args.lm_model_path} on {args.device}...")
+                    lm_status, lm_success = llm_handler.initialize(
+                        checkpoint_dir=checkpoint_dir,
+                        lm_model_path=args.lm_model_path,
+                        backend=args.backend,
+                        device=args.device,
+                        offload_to_cpu=args.offload_to_cpu,
+                        dtype=None,
+                    )
+                    
+                    if lm_success:
+                        print(f"5Hz LM initialized successfully")
+                        init_status += f"\n{lm_status}"
+                    else:
+                        print(f"Warning: 5Hz LM initialization failed: {lm_status}", file=sys.stderr)
+                        init_status += f"\n{lm_status}"
+            
+            # Prepare initialization parameters for UI
+            init_params = {
+                'pre_initialized': True,
+                'service_mode': args.service_mode,
+                'checkpoint': args.checkpoint,
+                'config_path': args.config_path,
+                'device': args.device,
+                'init_llm': args.init_llm,
+                'lm_model_path': args.lm_model_path,
+                'backend': args.backend,
+                'use_flash_attention': use_flash_attention,
+                'offload_to_cpu': args.offload_to_cpu,
+                'offload_dit_to_cpu': args.offload_dit_to_cpu,
+                'init_status': init_status,
+                'enable_generate': enable_generate,
+                'dit_handler': dit_handler,
+                'llm_handler': llm_handler,
+                'language': args.language,
+                'gpu_config': gpu_config,  # Pass GPU config to UI
+                'output_dir': output_dir,  # Pass output dir to UI
+            }
+            
+            print("Service initialization completed successfully!")
+        
+        # Create and launch demo
+        print(f"Creating Gradio interface with language: {args.language}...")
+        
+        # If not using init_service, still pass gpu_config to init_params
+        if init_params is None:
+            init_params = {
+                'gpu_config': gpu_config,
+                'language': args.language,
+                'output_dir': output_dir,  # Pass output dir to UI
+            }
+        
+        demo = create_demo(init_params=init_params, language=args.language)
+        
+        # Enable queue for multi-user support
+        # This ensures proper request queuing and prevents concurrent generation conflicts
+        print("Enabling queue for multi-user support...")
+        demo.queue(
+            max_size=20,  # Maximum queue size (adjust based on your needs)
+            status_update_rate="auto",  # Update rate for queue status
+            default_concurrency_limit=1,  # Prevents VRAM saturation
+        )
+
+        print(f"Launching server on {args.server_name}:{args.port}...")
+
+        # Setup authentication if provided
+        auth = None
+        if args.auth_username and args.auth_password:
+            auth = (args.auth_username, args.auth_password)
+            print("Authentication enabled")
+
+        allowed_paths = [output_dir]
+        for p in args.allowed_path:
+            if p and p not in allowed_paths:
+                allowed_paths.append(p)
+
+        # Enable API endpoints if requested
+        if args.enable_api:
+            print("Enabling API endpoints...")
+            from acestep.gradio_ui.api_routes import setup_api_routes
+
+            # Launch Gradio first with prevent_thread_lock=True
+            demo.launch(
+                server_name=args.server_name,
+                server_port=args.port,
+                share=args.share,
+                debug=args.debug,
+                show_error=True,
+                prevent_thread_lock=True,  # Don't block, so we can add routes
+                inbrowser=False,
+                auth=auth,
+                allowed_paths=allowed_paths,  # include output_dir + user-provided
+            )
+
+            # Now add API routes to Gradio's FastAPI app (app is available after launch)
+            setup_api_routes(demo, dit_handler, llm_handler, api_key=args.api_key)
+
+            if args.api_key:
+                print("API authentication enabled")
+            print("API endpoints enabled: /health, /v1/models, /release_task, /query_result, /create_random_sample, /format_lyrics")
+
+            # Keep the main thread alive
+            try:
+                while True:
+                    import time
+                    time.sleep(1)
+            except KeyboardInterrupt:
+                print("\nShutting down...")
+        else:
+            demo.launch(
+                server_name=args.server_name,
+                server_port=args.port,
+                share=args.share,
+                debug=args.debug,
+                show_error=True,
+                prevent_thread_lock=False,
+                inbrowser=False,
+                auth=auth,
+                allowed_paths=allowed_paths,  # include output_dir + user-provided
+            )
+    except Exception as e:
+        print(f"Error launching Gradio: {e}", file=sys.stderr)
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

acestep/audio_utils.py

@@ -0,0 +1,354 @@
+"""
+Audio saving and transcoding utility module
+
+Independent audio file operations outside of handler, supporting:
+- Save audio tensor/numpy to files (default FLAC format, fast)
+- Format conversion (FLAC/WAV/MP3)
+- Batch processing
+"""
+
+import os
+import hashlib
+import json
+from pathlib import Path
+from typing import Union, Optional, List, Tuple
+import torch
+import numpy as np
+import torchaudio
+from loguru import logger
+
+
+class AudioSaver:
+    """Audio saving and transcoding utility class"""
+    
+    def __init__(self, default_format: str = "flac"):
+        """
+        Initialize audio saver
+        
+        Args:
+            default_format: Default save format ('flac', 'wav', 'mp3')
+        """
+        self.default_format = default_format.lower()
+        if self.default_format not in ["flac", "wav", "mp3"]:
+            logger.warning(f"Unsupported format {default_format}, using 'flac'")
+            self.default_format = "flac"
+    
+    def save_audio(
+        self,
+        audio_data: Union[torch.Tensor, np.ndarray],
+        output_path: Union[str, Path],
+        sample_rate: int = 48000,
+        format: Optional[str] = None,
+        channels_first: bool = True,
+    ) -> str:
+        """
+        Save audio data to file
+        
+        Args:
+            audio_data: Audio data, torch.Tensor [channels, samples] or numpy.ndarray
+            output_path: Output file path (extension can be omitted)
+            sample_rate: Sample rate
+            format: Audio format ('flac', 'wav', 'mp3'), defaults to default_format
+            channels_first: If True, tensor format is [channels, samples], else [samples, channels]
+        
+        Returns:
+            Actual saved file path
+        """
+        format = (format or self.default_format).lower()
+        if format not in ["flac", "wav", "mp3"]:
+            logger.warning(f"Unsupported format {format}, using {self.default_format}")
+            format = self.default_format
+        
+        # Ensure output path has correct extension
+        output_path = Path(output_path)
+        if output_path.suffix.lower() not in ['.flac', '.wav', '.mp3']:
+            output_path = output_path.with_suffix(f'.{format}')
+        
+        # Convert to torch tensor
+        if isinstance(audio_data, np.ndarray):
+            if channels_first:
+                # numpy [samples, channels] -> tensor [channels, samples]
+                audio_tensor = torch.from_numpy(audio_data.T).float()
+            else:
+                # numpy [samples, channels] -> tensor [samples, channels] -> [channels, samples]
+                audio_tensor = torch.from_numpy(audio_data).float()
+                if audio_tensor.dim() == 2 and audio_tensor.shape[0] < audio_tensor.shape[1]:
+                    audio_tensor = audio_tensor.T
+        else:
+            # torch tensor
+            audio_tensor = audio_data.cpu().float()
+            if not channels_first and audio_tensor.dim() == 2:
+                # [samples, channels] -> [channels, samples]
+                if audio_tensor.shape[0] > audio_tensor.shape[1]:
+                    audio_tensor = audio_tensor.T
+        
+        # Ensure memory is contiguous
+        audio_tensor = audio_tensor.contiguous()
+        
+        # Select backend and save
+        try:
+            if format == "mp3":
+                # MP3 uses ffmpeg backend
+                torchaudio.save(
+                    str(output_path),
+                    audio_tensor,
+                    sample_rate,
+                    channels_first=True,
+                    backend='ffmpeg',
+                )
+            elif format in ["flac", "wav"]:
+                # FLAC and WAV use soundfile backend (fastest)
+                torchaudio.save(
+                    str(output_path),
+                    audio_tensor,
+                    sample_rate,
+                    channels_first=True,
+                    backend='soundfile',
+                )
+            else:
+                # Other formats use default backend
+                torchaudio.save(
+                    str(output_path),
+                    audio_tensor,
+                    sample_rate,
+                    channels_first=True,
+                )
+            
+            logger.debug(f"[AudioSaver] Saved audio to {output_path} ({format}, {sample_rate}Hz)")
+            return str(output_path)
+            
+        except Exception as e:
+            try:
+                import soundfile as sf
+                audio_np = audio_tensor.transpose(0, 1).numpy()  # -> [samples, channels]
+                sf.write(str(output_path), audio_np, sample_rate, format=format.upper())
+                logger.debug(f"[AudioSaver] Fallback soundfile Saved audio to {output_path} ({format}, {sample_rate}Hz)")
+                return str(output_path)
+            except Exception as e:
+                logger.error(f"[AudioSaver] Failed to save audio: {e}")
+                raise
+    
+    def convert_audio(
+        self,
+        input_path: Union[str, Path],
+        output_path: Union[str, Path],
+        output_format: str,
+        remove_input: bool = False,
+    ) -> str:
+        """
+        Convert audio format
+        
+        Args:
+            input_path: Input audio file path
+            output_path: Output audio file path
+            output_format: Target format ('flac', 'wav', 'mp3')
+            remove_input: Whether to delete input file
+        
+        Returns:
+            Output file path
+        """
+        input_path = Path(input_path)
+        output_path = Path(output_path)
+        
+        if not input_path.exists():
+            raise FileNotFoundError(f"Input file not found: {input_path}")
+        
+        # Load audio
+        audio_tensor, sample_rate = torchaudio.load(str(input_path))
+        
+        # Save as new format
+        output_path = self.save_audio(
+            audio_tensor,
+            output_path,
+            sample_rate=sample_rate,
+            format=output_format,
+            channels_first=True
+        )
+        
+        # Delete input file if needed
+        if remove_input:
+            input_path.unlink()
+            logger.debug(f"[AudioSaver] Removed input file: {input_path}")
+        
+        return output_path
+    
+    def save_batch(
+        self,
+        audio_batch: Union[List[torch.Tensor], torch.Tensor],
+        output_dir: Union[str, Path],
+        file_prefix: str = "audio",
+        sample_rate: int = 48000,
+        format: Optional[str] = None,
+        channels_first: bool = True,
+    ) -> List[str]:
+        """
+        Save audio batch
+        
+        Args:
+            audio_batch: Audio batch, List[tensor] or tensor [batch, channels, samples]
+            output_dir: Output directory
+            file_prefix: File prefix
+            sample_rate: Sample rate
+            format: Audio format
+            channels_first: Tensor format flag
+        
+        Returns:
+            List of saved file paths
+        """
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Process batch
+        if isinstance(audio_batch, torch.Tensor) and audio_batch.dim() == 3:
+            # [batch, channels, samples]
+            audio_list = [audio_batch[i] for i in range(audio_batch.shape[0])]
+        elif isinstance(audio_batch, list):
+            audio_list = audio_batch
+        else:
+            audio_list = [audio_batch]
+        
+        saved_paths = []
+        for i, audio in enumerate(audio_list):
+            output_path = output_dir / f"{file_prefix}_{i:04d}"
+            saved_path = self.save_audio(
+                audio,
+                output_path,
+                sample_rate=sample_rate,
+                format=format,
+                channels_first=channels_first
+            )
+            saved_paths.append(saved_path)
+        
+        return saved_paths
+
+
+def get_audio_file_hash(audio_file) -> str:
+    """
+    Get hash identifier for an audio file.
+    
+    Args:
+        audio_file: Path to audio file (str) or file-like object
+    
+    Returns:
+        Hash string or empty string
+    """
+    if audio_file is None:
+        return ""
+    
+    try:
+        if isinstance(audio_file, str):
+            if os.path.exists(audio_file):
+                with open(audio_file, 'rb') as f:
+                    return hashlib.md5(f.read()).hexdigest()
+            return hashlib.md5(audio_file.encode('utf-8')).hexdigest()
+        elif hasattr(audio_file, 'name'):
+            return hashlib.md5(str(audio_file.name).encode('utf-8')).hexdigest()
+        return hashlib.md5(str(audio_file).encode('utf-8')).hexdigest()
+    except Exception:
+        return hashlib.md5(str(audio_file).encode('utf-8')).hexdigest()
+
+
+def generate_uuid_from_params(params_dict) -> str:
+    """
+    Generate deterministic UUID from generation parameters.
+    Same parameters will always generate the same UUID.
+    
+    Args:
+        params_dict: Dictionary of parameters
+    
+    Returns:
+        UUID string
+    """
+    
+    params_json = json.dumps(params_dict, sort_keys=True, ensure_ascii=False)
+    hash_obj = hashlib.sha256(params_json.encode('utf-8'))
+    hash_hex = hash_obj.hexdigest()
+    uuid_str = f"{hash_hex[0:8]}-{hash_hex[8:12]}-{hash_hex[12:16]}-{hash_hex[16:20]}-{hash_hex[20:32]}"
+    return uuid_str
+
+
+def generate_uuid_from_audio_data(
+    audio_data: Union[torch.Tensor, np.ndarray],
+    seed: Optional[int] = None
+) -> str:
+    """
+    Generate UUID from audio data (for caching/deduplication)
+    
+    Args:
+        audio_data: Audio data
+        seed: Optional seed value
+    
+    Returns:
+        UUID string
+    """
+    if isinstance(audio_data, torch.Tensor):
+        # Convert to numpy and calculate hash
+        audio_np = audio_data.cpu().numpy()
+    else:
+        audio_np = audio_data
+    
+    # Calculate data hash
+    data_hash = hashlib.md5(audio_np.tobytes()).hexdigest()
+    
+    if seed is not None:
+        combined = f"{data_hash}_{seed}"
+        return hashlib.md5(combined.encode()).hexdigest()
+    
+    return data_hash
+
+
+# Global default instance
+_default_saver = AudioSaver(default_format="flac")
+
+SILENT_RMS_THRESHOLD = 1e-5
+SILENT_PEAK_THRESHOLD = 1e-5
+
+
+def is_audio_silent(
+    audio_data: Union[torch.Tensor, np.ndarray],
+    rms_threshold: float = SILENT_RMS_THRESHOLD,
+    peak_threshold: float = SILENT_PEAK_THRESHOLD,
+    channels_first: bool = True,
+) -> Tuple[bool, float, float]:
+    """
+    Check if audio is silent or near-silent (e.g. zeroed conditioning output).
+    Returns (is_silent, rms, peak) where rms/peak are computed over the full signal.
+    """
+    if audio_data is None:
+        return True, 0.0, 0.0
+    if isinstance(audio_data, np.ndarray):
+        x = np.asarray(audio_data, dtype=np.float64).ravel()
+    else:
+        x = audio_data.cpu().float().numpy().ravel()
+    if x.size == 0:
+        return True, 0.0, 0.0
+    rms = float(np.sqrt(np.mean(x * x)))
+    peak = float(np.max(np.abs(x)))
+    is_silent = rms <= rms_threshold and peak <= peak_threshold
+    return is_silent, rms, peak
+
+
+def save_audio(
+    audio_data: Union[torch.Tensor, np.ndarray],
+    output_path: Union[str, Path],
+    sample_rate: int = 48000,
+    format: Optional[str] = None,
+    channels_first: bool = True,
+) -> str:
+    """
+    Convenience function: save audio (using default configuration)
+    
+    Args:
+        audio_data: Audio data
+        output_path: Output path
+        sample_rate: Sample rate
+        format: Format (default flac)
+        channels_first: Tensor format flag
+    
+    Returns:
+        Saved file path
+    """
+    return _default_saver.save_audio(
+        audio_data, output_path, sample_rate, format, channels_first
+    )
+

acestep/constants.py

@@ -0,0 +1,193 @@
+"""
+Constants for ACE-Step
+Centralized constants used across the codebase
+"""
+
+# ==============================================================================
+# Language Constants
+# ==============================================================================
+
+# Supported languages for vocal generation and language detection
+# Covers major world languages with good TTS support in the underlying model
+# 'unknown' is used when language cannot be determined automatically
+VALID_LANGUAGES = [
+    'ar', 'az', 'bg', 'bn', 'ca', 'cs', 'da', 'de', 'el', 'en',
+    'es', 'fa', 'fi', 'fr', 'he', 'hi', 'hr', 'ht', 'hu', 'id',
+    'is', 'it', 'ja', 'ko', 'la', 'lt', 'ms', 'ne', 'nl', 'no',
+    'pa', 'pl', 'pt', 'ro', 'ru', 'sa', 'sk', 'sr', 'sv', 'sw',
+    'ta', 'te', 'th', 'tl', 'tr', 'uk', 'ur', 'vi', 'yue', 'zh',
+    'unknown'
+]
+
+
+# ==============================================================================
+# Keyscale Constants
+# ==============================================================================
+
+# Musical note names using standard Western notation
+KEYSCALE_NOTES = ['A', 'B', 'C', 'D', 'E', 'F', 'G']
+
+# Supported accidentals: natural, ASCII sharp/flat, Unicode sharp/flat
+KEYSCALE_ACCIDENTALS = ['', '#', 'b', '♯', '♭']  # empty + ASCII sharp/flat + Unicode sharp/flat
+
+# Major and minor scale modes
+KEYSCALE_MODES = ['major', 'minor']
+
+# Generate all valid keyscales: 7 notes × 5 accidentals × 2 modes = 70 combinations
+# Examples: "C major", "F# minor", "B♭ major"
+VALID_KEYSCALES = set()
+for note in KEYSCALE_NOTES:
+    for acc in KEYSCALE_ACCIDENTALS:
+        for mode in KEYSCALE_MODES:
+            VALID_KEYSCALES.add(f"{note}{acc} {mode}")
+
+
+# ==============================================================================
+# Metadata Range Constants
+# ==============================================================================
+
+# BPM (Beats Per Minute) range - covers most musical styles
+# 30 BPM: Very slow ballads, ambient music
+# 300 BPM: Fast electronic dance music, extreme metal
+BPM_MIN = 30
+BPM_MAX = 300
+
+# Duration range (in seconds) - balances quality vs. computational cost
+# 10s: Short loops, musical excerpts
+# 600s: Full songs, extended compositions (10 minutes)
+DURATION_MIN = 10
+DURATION_MAX = 600
+
+# Valid time signatures - common musical meter patterns
+# 2: 2/4 time (marches, polka)
+# 3: 3/4 time (waltzes, ballads) 
+# 4: 4/4 time (most pop, rock, hip-hop)
+# 6: 6/8 time (compound time, folk dances)
+VALID_TIME_SIGNATURES = [2, 3, 4, 6]
+
+
+# ==============================================================================
+# Task Type Constants  
+# ==============================================================================
+
+# All supported generation tasks across different model variants
+TASK_TYPES = ["text2music", "repaint", "cover", "extract", "lego", "complete"]
+
+# Task types available for turbo models (optimized subset for speed)
+# - text2music: Generate from text descriptions
+# - repaint: Selective audio editing/regeneration  
+# - cover: Style transfer using reference audio
+TASK_TYPES_TURBO = ["text2music", "repaint", "cover"]
+
+# Task types available for base models (full feature set)
+# Additional tasks requiring more computational resources:
+# - extract: Separate individual tracks/stems from audio
+# - lego: Multi-track generation (add layers)
+# - complete: Automatic completion of partial audio
+TASK_TYPES_BASE = ["text2music", "repaint", "cover", "extract", "lego", "complete"]
+
+
+# ==============================================================================
+# Instruction Constants
+# ==============================================================================
+
+# Default instructions
+DEFAULT_DIT_INSTRUCTION = "Fill the audio semantic mask based on the given conditions:"
+DEFAULT_LM_INSTRUCTION = "Generate audio semantic tokens based on the given conditions:"
+DEFAULT_LM_UNDERSTAND_INSTRUCTION = "Understand the given musical conditions and describe the audio semantics accordingly:"
+DEFAULT_LM_INSPIRED_INSTRUCTION = "Expand the user's input into a more detailed and specific musical description:"
+DEFAULT_LM_REWRITE_INSTRUCTION = "Format the user's input into a more detailed and specific musical description:"
+
+# Instruction templates for each task type
+# Note: Some instructions use placeholders like {TRACK_NAME} or {TRACK_CLASSES}
+# These should be formatted using .format() or f-strings when used
+TASK_INSTRUCTIONS = {
+    "text2music": "Fill the audio semantic mask based on the given conditions:",
+    "repaint": "Repaint the mask area based on the given conditions:",
+    "cover": "Generate audio semantic tokens based on the given conditions:",
+    "extract": "Extract the {TRACK_NAME} track from the audio:",
+    "extract_default": "Extract the track from the audio:",
+    "lego": "Generate the {TRACK_NAME} track based on the audio context:",
+    "lego_default": "Generate the track based on the audio context:",
+    "complete": "Complete the input track with {TRACK_CLASSES}:",
+    "complete_default": "Complete the input track:",
+}
+
+
+# ==============================================================================
+# Track/Instrument Constants
+# ==============================================================================
+
+# Supported instrumental track types for multi-track generation and extraction
+# Organized by instrument families for logical grouping:
+# - Wind instruments: woodwinds, brass
+# - Electronic: fx (effects), synth (synthesizer)  
+# - String instruments: strings, guitar, bass
+# - Rhythm section: percussion, drums, keyboard
+# - Vocals: backing_vocals, vocals (lead vocals)
+TRACK_NAMES = [
+    "woodwinds", "brass", "fx", "synth", "strings", "percussion",
+    "keyboard", "guitar", "bass", "drums", "backing_vocals", "vocals"
+]
+
+# Template for SFT (Supervised Fine-Tuning) model prompts
+# Used to format inputs for the language model with instruction, caption, and metadata
+SFT_GEN_PROMPT = """# Instruction
+{}
+
+# Caption
+{}
+
+# Metas
+{}<|endoftext|>
+"""
+
+
+# ==============================================================================
+# GPU Memory Configuration Constants
+# ==============================================================================
+
+# GPU tier thresholds (in GB)
+GPU_TIER_THRESHOLDS = {
+    "tier1": 4,    # <= 4GB
+    "tier2": 6,    # 4-6GB
+    "tier3": 8,    # 6-8GB
+    "tier4": 12,   # 8-12GB
+    "tier5": 16,   # 12-16GB
+    "tier6": 24,   # 16-24GB
+    # "unlimited" for >= 24GB
+}
+
+# LM model memory requirements (in GB)
+LM_MODEL_MEMORY_GB = {
+    "0.6B": 3.0,
+    "1.7B": 8.0,
+    "4B": 12.0,
+}
+
+# LM model names mapping
+LM_MODEL_NAMES = {
+    "0.6B": "acestep-5Hz-lm-0.6B",
+    "1.7B": "acestep-5Hz-lm-1.7B",
+    "4B": "acestep-5Hz-lm-4B",
+}
+
+
+# ==============================================================================
+# Debug Constants
+# ==============================================================================
+
+# Tensor debug mode (values: "OFF" | "ON" | "VERBOSE")
+TENSOR_DEBUG_MODE = "OFF"
+
+# Placeholder debug switches for other main functionality (default "OFF")
+# Update names/usage as features adopt them.
+DEBUG_API_SERVER = "OFF"
+DEBUG_INFERENCE = "OFF"
+DEBUG_TRAINING = "OFF"
+DEBUG_DATASET = "OFF"
+DEBUG_AUDIO = "OFF"
+DEBUG_LLM = "OFF"
+DEBUG_UI = "OFF"
+DEBUG_MODEL_LOADING = "OFF"
+DEBUG_GPU = "OFF"

acestep/dataset_handler.py

@@ -0,0 +1,83 @@
+"""
+Dataset Handler Module
+
+Handles dataset import and exploration functionality for ACE-Step training.
+This module provides a placeholder implementation for dataset operations
+when the full training dataset dependencies are not available.
+
+Note: Full dataset functionality requires Text2MusicDataset which may not be
+included in the basic installation to reduce dependencies.
+"""
+from typing import Optional, Tuple, Any, Dict
+
+
+class DatasetHandler:
+    """
+    Dataset Handler for Dataset Explorer functionality.
+    
+    Provides interface for dataset import and exploration features in the Gradio UI.
+    When training dependencies are not available, returns appropriate fallback responses.
+    """
+    
+    def __init__(self):
+        """Initialize dataset handler with empty state"""
+        self.dataset = None
+        self.dataset_imported = False
+    
+    def import_dataset(self, dataset_type: str) -> str:
+        """
+        Import dataset (currently disabled in base installation)
+        
+        Args:
+            dataset_type: Type of dataset to import (e.g., "train", "test", "validation")
+            
+        Returns:
+            Status message indicating dataset import is disabled
+            
+        Note:
+            This is a placeholder implementation. Full dataset support requires:
+            - Text2MusicDataset dependency
+            - Training data files  
+            - Additional configuration
+        """
+        self.dataset_imported = False
+        return f"⚠️ Dataset import is currently disabled. Text2MusicDataset dependency not available."
+    
+    def get_item_data(self, *args, **kwargs) -> Tuple:
+        """
+        Get dataset item data (placeholder implementation)
+        
+        Args:
+            *args: Variable arguments (ignored in placeholder)
+            **kwargs: Keyword arguments (ignored in placeholder)
+            
+        Returns:
+            Tuple of placeholder values matching the expected return format:
+            (caption, lyrics, language, bpm, keyscale, ref_audio, src_audio, codes,
+             status_msg, instruction, duration, timesig, audio1, audio2, audio3, 
+             metadata, task_type)
+             
+        Note:
+            Returns empty/default values since dataset is not available.
+            Real implementation would return actual dataset samples.
+        """
+        return (
+            "",           # caption: empty string
+            "",           # lyrics: empty string  
+            "",           # language: empty string
+            "",           # bpm: empty string
+            "",           # keyscale: empty string
+            None,         # ref_audio: no audio file
+            None,         # src_audio: no audio file
+            None,         # codes: no audio codes
+            "❌ Dataset not available",  # status_msg: error indicator
+            "",           # instruction: empty string
+            0,            # duration: zero
+            "",           # timesig: empty string
+            None,         # audio1: no audio
+            None,         # audio2: no audio
+            None,         # audio3: no audio
+            {},           # metadata: empty dict
+            "text2music"  # task_type: default task
+        )
+

acestep/debug_utils.py

@@ -0,0 +1,122 @@
+"""
+Debug helpers (global).
+"""
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Optional, Callable, Union
+
+from acestep.constants import (
+    TENSOR_DEBUG_MODE,
+    DEBUG_API_SERVER,
+    DEBUG_INFERENCE,
+    DEBUG_TRAINING,
+    DEBUG_DATASET,
+    DEBUG_AUDIO,
+    DEBUG_LLM,
+    DEBUG_UI,
+    DEBUG_MODEL_LOADING,
+    DEBUG_GPU,
+)
+
+
+def _normalize_mode(mode: str) -> str:
+    return (mode or "").strip().upper()
+
+
+def is_debug_enabled(mode: str) -> bool:
+    return _normalize_mode(mode) != "OFF"
+
+
+def is_debug_verbose(mode: str) -> bool:
+    return _normalize_mode(mode) == "VERBOSE"
+
+
+def debug_log(message: Union[str, Callable[[], str]], *, mode: str = TENSOR_DEBUG_MODE, prefix: str = "debug") -> None:
+    """Emit a timestamped debug log line if the mode is enabled."""
+    if not is_debug_enabled(mode):
+        return
+    if callable(message):
+        message = message()
+    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
+    print(f"[{prefix}] {ts} {message}", flush=True)
+
+
+# Placeholder debug switches registry (for centralized access)
+DEBUG_SWITCHES = {
+    "tensor": TENSOR_DEBUG_MODE,
+    "api_server": DEBUG_API_SERVER,
+    "inference": DEBUG_INFERENCE,
+    "training": DEBUG_TRAINING,
+    "dataset": DEBUG_DATASET,
+    "audio": DEBUG_AUDIO,
+    "llm": DEBUG_LLM,
+    "ui": DEBUG_UI,
+    "model_loading": DEBUG_MODEL_LOADING,
+    "gpu": DEBUG_GPU,
+}
+
+
+def get_debug_mode(name: str, default: str = "OFF") -> str:
+    """Fetch a placeholder debug mode by name."""
+    return DEBUG_SWITCHES.get((name or "").strip().lower(), default)
+
+
+def debug_log_for(name: str, message: Union[str, Callable[[], str]], *, prefix: str | None = None) -> None:
+    """Emit a timestamped debug log for a named subsystem."""
+    mode = get_debug_mode(name)
+    debug_log(message, mode=mode, prefix=prefix or name)
+
+
+def debug_start_for(name: str, label: str) -> Optional[float]:
+    """Start timing for a named subsystem."""
+    mode = get_debug_mode(name)
+    return debug_start(label, mode=mode, prefix=name)
+
+
+def debug_end_for(name: str, label: str, start_ts: Optional[float]) -> None:
+    """End timing for a named subsystem."""
+    mode = get_debug_mode(name)
+    debug_end(label, start_ts, mode=mode, prefix=name)
+
+
+def debug_log_verbose_for(name: str, message: Union[str, Callable[[], str]], *, prefix: str | None = None) -> None:
+    """Emit a timestamped debug log only in VERBOSE mode for a named subsystem."""
+    mode = get_debug_mode(name)
+    if not is_debug_verbose(mode):
+        return
+    debug_log(message, mode=mode, prefix=prefix or name)
+
+
+def debug_start_verbose_for(name: str, label: str) -> Optional[float]:
+    """Start timing only in VERBOSE mode for a named subsystem."""
+    mode = get_debug_mode(name)
+    if not is_debug_verbose(mode):
+        return None
+    return debug_start(label, mode=mode, prefix=name)
+
+
+def debug_end_verbose_for(name: str, label: str, start_ts: Optional[float]) -> None:
+    """End timing only in VERBOSE mode for a named subsystem."""
+    mode = get_debug_mode(name)
+    if not is_debug_verbose(mode):
+        return
+    debug_end(label, start_ts, mode=mode, prefix=name)
+
+
+def debug_start(name: str, *, mode: str = TENSOR_DEBUG_MODE, prefix: str = "debug") -> Optional[float]:
+    """Return a start timestamp (perf counter) if enabled, otherwise None."""
+    if not is_debug_enabled(mode):
+        return None
+    debug_log(f"START {name}", mode=mode, prefix=prefix)
+    from time import perf_counter
+    return perf_counter()
+
+
+def debug_end(name: str, start_ts: Optional[float], *, mode: str = TENSOR_DEBUG_MODE, prefix: str = "debug") -> None:
+    """Emit an END log with elapsed ms if enabled and start_ts is present."""
+    if start_ts is None or not is_debug_enabled(mode):
+        return
+    from time import perf_counter
+    elapsed_ms = (perf_counter() - start_ts) * 1000.0
+    debug_log(f"END {name} ({elapsed_ms:.1f} ms)", mode=mode, prefix=prefix)

acestep/dit_alignment_score.py

@@ -0,0 +1,877 @@
+"""
+DiT Alignment Score Module
+
+This module provides lyrics-to-audio alignment using cross-attention matrices
+from DiT model for generating LRC timestamps.
+
+Refactored from lyrics_alignment_infos.py for integration with ACE-Step.
+"""
+import numba
+import torch
+import numpy as np
+import torch.nn.functional as F
+from dataclasses import dataclass, asdict
+from typing import List, Dict, Any, Optional, Tuple, Union
+
+
+# ================= Data Classes =================
+@dataclass
+class TokenTimestamp:
+    """Stores per-token timing information."""
+    token_id: int
+    text: str
+    start: float
+    end: float
+    probability: float
+
+
+@dataclass
+class SentenceTimestamp:
+    """Stores per-sentence timing information with token list."""
+    text: str
+    start: float
+    end: float
+    tokens: List[TokenTimestamp]
+    confidence: float
+
+
+# ================= DTW Algorithm (Numba Optimized) =================
+@numba.jit(nopython=True)
+def dtw_cpu(x: np.ndarray):
+    """
+    Dynamic Time Warping algorithm optimized with Numba.
+    
+    Args:
+        x: Cost matrix of shape [N, M]
+        
+    Returns:
+        Tuple of (text_indices, time_indices) arrays
+    """
+    N, M = x.shape
+    # Use float32 for memory efficiency
+    cost = np.ones((N + 1, M + 1), dtype=np.float32) * np.inf
+    trace = -np.ones((N + 1, M + 1), dtype=np.float32)
+    cost[0, 0] = 0
+
+    for j in range(1, M + 1):
+        for i in range(1, N + 1):
+            c0 = cost[i - 1, j - 1]
+            c1 = cost[i - 1, j]
+            c2 = cost[i, j - 1]
+
+            if c0 < c1 and c0 < c2:
+                c, t = c0, 0
+            elif c1 < c0 and c1 < c2:
+                c, t = c1, 1
+            else:
+                c, t = c2, 2
+
+            cost[i, j] = x[i - 1, j - 1] + c
+            trace[i, j] = t
+
+    return _backtrace(trace, N, M)
+
+
+@numba.jit(nopython=True)
+def _backtrace(trace: np.ndarray, N: int, M: int):
+    """
+    Optimized backtrace function for DTW.
+    
+    Args:
+        trace: Trace matrix of shape (N+1, M+1)
+        N, M: Original matrix dimensions
+        
+    Returns:
+        Path array of shape (2, path_len) - first row is text indices, second is time indices
+    """
+    # Boundary handling
+    trace[0, :] = 2
+    trace[:, 0] = 1
+    
+    # Pre-allocate array, max path length is N+M
+    max_path_len = N + M
+    path = np.zeros((2, max_path_len), dtype=np.int32)
+    
+    i, j = N, M
+    path_idx = max_path_len - 1
+    
+    while i > 0 or j > 0:
+        path[0, path_idx] = i - 1  # text index
+        path[1, path_idx] = j - 1  # time index
+        path_idx -= 1
+        
+        t = trace[i, j]
+        if t == 0:
+            i -= 1
+            j -= 1
+        elif t == 1:
+            i -= 1
+        elif t == 2:
+            j -= 1
+        else:
+            break
+    
+    actual_len = max_path_len - path_idx - 1
+    return path[:, path_idx + 1:max_path_len]
+
+
+# ================= Utility Functions =================
+def median_filter(x: torch.Tensor, filter_width: int) -> torch.Tensor:
+    """
+    Apply median filter to tensor.
+    
+    Args:
+        x: Input tensor
+        filter_width: Width of median filter
+        
+    Returns:
+        Filtered tensor
+    """
+    pad_width = filter_width // 2
+    if x.shape[-1] <= pad_width:
+        return x
+    if x.ndim == 2:
+        x = x[None, :]
+    x = F.pad(x, (filter_width // 2, filter_width // 2, 0, 0), mode="reflect")
+    result = x.unfold(-1, filter_width, 1).sort()[0][..., filter_width // 2]
+    if result.ndim > 2:
+        result = result.squeeze(0)
+    return result
+
+
+# ================= Main Aligner Class =================
+class MusicStampsAligner:
+    """
+    Aligner class for generating lyrics timestamps from cross-attention matrices.
+    
+    Uses bidirectional consensus denoising and DTW for alignment.
+    """
+    
+    def __init__(self, tokenizer):
+        """
+        Initialize the aligner.
+        
+        Args:
+            tokenizer: Text tokenizer for decoding tokens
+        """
+        self.tokenizer = tokenizer
+
+    def _apply_bidirectional_consensus(
+        self, 
+        weights_stack: torch.Tensor, 
+        violence_level: float, 
+        medfilt_width: int
+    ) -> tuple:
+        """
+        Core denoising logic using bidirectional consensus.
+        
+        Args:
+            weights_stack: Attention weights [Heads, Tokens, Frames]
+            violence_level: Denoising strength coefficient
+            medfilt_width: Median filter width
+            
+        Returns:
+            Tuple of (calc_matrix, energy_matrix) as numpy arrays
+        """
+        # A. Bidirectional Consensus
+        row_prob = F.softmax(weights_stack, dim=-1)  # Token -> Frame
+        col_prob = F.softmax(weights_stack, dim=-2)  # Frame -> Token
+        processed = row_prob * col_prob
+
+        # 1. Row suppression (kill horizontal crossing lines)
+        row_medians = torch.quantile(processed, 0.5, dim=-1, keepdim=True)
+        processed = processed - (violence_level * row_medians)
+        processed = torch.relu(processed)
+
+        # 2. Column suppression (kill vertical crossing lines)
+        col_medians = torch.quantile(processed, 0.5, dim=-2, keepdim=True)
+        processed = processed - (violence_level * col_medians)
+        processed = torch.relu(processed)
+
+        # C. Power sharpening
+        processed = processed ** 2
+
+        # Energy matrix for confidence
+        energy_matrix = processed.mean(dim=0).cpu().numpy()
+        
+        # D. Z-Score normalization
+        std, mean = torch.std_mean(processed, unbiased=False)
+        weights_processed = (processed - mean) / (std + 1e-9)
+
+        # E. Median filtering
+        weights_processed = median_filter(weights_processed, filter_width=medfilt_width)
+        calc_matrix = weights_processed.mean(dim=0).numpy()
+        
+        return calc_matrix, energy_matrix
+
+    def _preprocess_attention(
+        self, 
+        attention_matrix: torch.Tensor, 
+        custom_config: Dict[int, List[int]], 
+        violence_level: float, 
+        medfilt_width: int = 7
+    ) -> tuple:
+        """
+        Preprocess attention matrix for alignment.
+        
+        Args:
+            attention_matrix: Attention tensor [Layers, Heads, Tokens, Frames]
+            custom_config: Dict mapping layer indices to head indices
+            violence_level: Denoising strength
+            medfilt_width: Median filter width
+            
+        Returns:
+            Tuple of (calc_matrix, energy_matrix, visual_matrix)
+        """
+        if not isinstance(attention_matrix, torch.Tensor):
+            weights = torch.tensor(attention_matrix)
+        else:
+            weights = attention_matrix.clone()
+
+        weights = weights.cpu().float()
+
+        selected_tensors = []
+        for layer_idx, head_indices in custom_config.items():
+            for head_idx in head_indices:
+                if layer_idx < weights.shape[0] and head_idx < weights.shape[1]:
+                    head_matrix = weights[layer_idx, head_idx]
+                    selected_tensors.append(head_matrix)
+
+        if not selected_tensors:
+            return None, None, None
+
+        # Stack selected heads: [Heads, Tokens, Frames]
+        weights_stack = torch.stack(selected_tensors, dim=0)
+        visual_matrix = weights_stack.mean(dim=0).numpy()
+
+        calc_matrix, energy_matrix = self._apply_bidirectional_consensus(
+            weights_stack, violence_level, medfilt_width
+        )
+
+        return calc_matrix, energy_matrix, visual_matrix
+
+    def stamps_align_info(
+        self,
+        attention_matrix: torch.Tensor,
+        lyrics_tokens: List[int],
+        total_duration_seconds: float,
+        custom_config: Dict[int, List[int]],
+        return_matrices: bool = False,
+        violence_level: float = 2.0,
+        medfilt_width: int = 1
+    ) -> Dict[str, Any]:
+        """
+        Get alignment information from attention matrix.
+        
+        Args:
+            attention_matrix: Cross-attention tensor [Layers, Heads, Tokens, Frames]
+            lyrics_tokens: List of lyrics token IDs
+            total_duration_seconds: Total audio duration in seconds
+            custom_config: Dict mapping layer indices to head indices
+            return_matrices: Whether to return intermediate matrices
+            violence_level: Denoising strength
+            medfilt_width: Median filter width
+            
+        Returns:
+            Dict containing calc_matrix, lyrics_tokens, total_duration_seconds,
+            and optionally energy_matrix and vis_matrix
+        """
+        calc_matrix, energy_matrix, visual_matrix = self._preprocess_attention(
+            attention_matrix, custom_config, violence_level, medfilt_width
+        )
+        
+        if calc_matrix is None:
+            return {
+                "calc_matrix": None,
+                "lyrics_tokens": lyrics_tokens,
+                "total_duration_seconds": total_duration_seconds,
+                "error": "No valid attention heads found"
+            }
+        
+        return_dict = {
+            "calc_matrix": calc_matrix,
+            "lyrics_tokens": lyrics_tokens,
+            "total_duration_seconds": total_duration_seconds
+        }
+        
+        if return_matrices:
+            return_dict['energy_matrix'] = energy_matrix
+            return_dict['vis_matrix'] = visual_matrix
+
+        return return_dict
+
+    def _decode_tokens_incrementally(self, token_ids: List[int]) -> List[str]:
+        """
+        Decode tokens incrementally to properly handle multi-byte UTF-8 characters.
+        
+        For Chinese and other multi-byte characters, the tokenizer may split them
+        into multiple byte-level tokens. Decoding each token individually produces
+        invalid UTF-8 sequences (showing as �). This method uses byte-level comparison
+        to correctly track which characters each token contributes.
+        
+        Args:
+            token_ids: List of token IDs
+            
+        Returns:
+            List of decoded text for each token position
+        """
+        decoded_tokens = []
+        prev_bytes = b""
+        
+        for i in range(len(token_ids)):
+            # Decode tokens from start to current position
+            current_text = self.tokenizer.decode(token_ids[:i+1], skip_special_tokens=False)
+            current_bytes = current_text.encode('utf-8', errors='surrogatepass')
+            
+            # The contribution of current token is the new bytes added
+            if len(current_bytes) >= len(prev_bytes):
+                new_bytes = current_bytes[len(prev_bytes):]
+                # Try to decode the new bytes; if incomplete, use empty string
+                try:
+                    token_text = new_bytes.decode('utf-8')
+                except UnicodeDecodeError:
+                    # Incomplete UTF-8 sequence, this token doesn't complete a character
+                    token_text = ""
+            else:
+                # Edge case: current decode is shorter (shouldn't happen normally)
+                token_text = ""
+            
+            decoded_tokens.append(token_text)
+            prev_bytes = current_bytes
+        
+        return decoded_tokens
+
+    def token_timestamps(
+        self,
+        calc_matrix: np.ndarray,
+        lyrics_tokens: List[int],
+        total_duration_seconds: float
+    ) -> List[TokenTimestamp]:
+        """
+        Generate per-token timestamps using DTW.
+        
+        Args:
+            calc_matrix: Processed attention matrix [Tokens, Frames]
+            lyrics_tokens: List of token IDs
+            total_duration_seconds: Total audio duration
+            
+        Returns:
+            List of TokenTimestamp objects
+        """
+        n_frames = calc_matrix.shape[-1]
+        text_indices, time_indices = dtw_cpu(-calc_matrix.astype(np.float64))
+
+        seconds_per_frame = total_duration_seconds / n_frames
+        alignment_results = []
+        
+        # Use incremental decoding to properly handle multi-byte UTF-8 characters
+        decoded_tokens = self._decode_tokens_incrementally(lyrics_tokens)
+
+        for i in range(len(lyrics_tokens)):
+            mask = (text_indices == i)
+
+            if not np.any(mask):
+                start = alignment_results[-1].end if alignment_results else 0.0
+                end = start
+                token_conf = 0.0
+            else:
+                times = time_indices[mask] * seconds_per_frame
+                start = times[0]
+                end = times[-1]
+                token_conf = 0.0
+
+            if end < start:
+                end = start
+
+            alignment_results.append(TokenTimestamp(
+                token_id=lyrics_tokens[i],
+                text=decoded_tokens[i],
+                start=float(start),
+                end=float(end),
+                probability=token_conf
+            ))
+
+        return alignment_results
+
+    def _decode_sentence_from_tokens(self, tokens: List[TokenTimestamp]) -> str:
+        """
+        Decode a sentence by decoding all token IDs together.
+        This avoids UTF-8 encoding issues from joining individual token texts.
+        
+        Args:
+            tokens: List of TokenTimestamp objects
+            
+        Returns:
+            Properly decoded sentence text
+        """
+        token_ids = [t.token_id for t in tokens]
+        return self.tokenizer.decode(token_ids, skip_special_tokens=False)
+
+    def sentence_timestamps(
+        self, 
+        token_alignment: List[TokenTimestamp]
+    ) -> List[SentenceTimestamp]:
+        """
+        Group token timestamps into sentence timestamps.
+        
+        Args:
+            token_alignment: List of TokenTimestamp objects
+            
+        Returns:
+            List of SentenceTimestamp objects
+        """
+        results = []
+        current_tokens = []
+
+        for token in token_alignment:
+            current_tokens.append(token)
+
+            if '\n' in token.text:
+                # Decode all token IDs together to avoid UTF-8 issues
+                full_text = self._decode_sentence_from_tokens(current_tokens)
+
+                if full_text.strip():
+                    valid_scores = [t.probability for t in current_tokens if t.probability > 0]
+                    sent_conf = sum(valid_scores) / len(valid_scores) if valid_scores else 0.0
+
+                    results.append(SentenceTimestamp(
+                        text=full_text.strip(),
+                        start=round(current_tokens[0].start, 3),
+                        end=round(current_tokens[-1].end, 3),
+                        tokens=list(current_tokens),
+                        confidence=sent_conf
+                    ))
+
+                current_tokens = []
+
+        # Handle last sentence
+        if current_tokens:
+            # Decode all token IDs together to avoid UTF-8 issues
+            full_text = self._decode_sentence_from_tokens(current_tokens)
+            if full_text.strip():
+                valid_scores = [t.probability for t in current_tokens if t.probability > 0]
+                sent_conf = sum(valid_scores) / len(valid_scores) if valid_scores else 0.0
+
+                results.append(SentenceTimestamp(
+                    text=full_text.strip(),
+                    start=round(current_tokens[0].start, 3),
+                    end=round(current_tokens[-1].end, 3),
+                    tokens=list(current_tokens),
+                    confidence=sent_conf
+                ))
+
+        # Normalize confidence scores
+        if results:
+            all_scores = [s.confidence for s in results]
+            min_score = min(all_scores)
+            max_score = max(all_scores)
+            score_range = max_score - min_score
+
+            if score_range > 1e-9:
+                for s in results:
+                    normalized_score = (s.confidence - min_score) / score_range
+                    s.confidence = round(normalized_score, 2)
+            else:
+                for s in results:
+                    s.confidence = round(s.confidence, 2)
+
+        return results
+
+    def format_lrc(
+        self, 
+        sentence_timestamps: List[SentenceTimestamp],
+        include_end_time: bool = False
+    ) -> str:
+        """
+        Format sentence timestamps as LRC lyrics format.
+        
+        Args:
+            sentence_timestamps: List of SentenceTimestamp objects
+            include_end_time: Whether to include end time (enhanced LRC format)
+            
+        Returns:
+            LRC formatted string
+        """
+        lines = []
+        
+        for sentence in sentence_timestamps:
+            # Convert seconds to mm:ss.xx format
+            start_minutes = int(sentence.start // 60)
+            start_seconds = sentence.start % 60
+            
+            if include_end_time:
+                end_minutes = int(sentence.end // 60)
+                end_seconds = sentence.end % 60
+                timestamp = f"[{start_minutes:02d}:{start_seconds:05.2f}][{end_minutes:02d}:{end_seconds:05.2f}]"
+            else:
+                timestamp = f"[{start_minutes:02d}:{start_seconds:05.2f}]"
+            
+            # Clean the text (remove structural tags like [verse], [chorus])
+            text = sentence.text
+            
+            lines.append(f"{timestamp}{text}")
+        
+        return "\n".join(lines)
+
+    def get_timestamps_and_lrc(
+        self,
+        calc_matrix: np.ndarray,
+        lyrics_tokens: List[int],
+        total_duration_seconds: float
+    ) -> Dict[str, Any]:
+        """
+        Convenience method to get both timestamps and LRC in one call.
+        
+        Args:
+            calc_matrix: Processed attention matrix
+            lyrics_tokens: List of token IDs
+            total_duration_seconds: Total audio duration
+            
+        Returns:
+            Dict containing token_timestamps, sentence_timestamps, and lrc_text
+        """
+        token_stamps = self.token_timestamps(
+            calc_matrix=calc_matrix,
+            lyrics_tokens=lyrics_tokens,
+            total_duration_seconds=total_duration_seconds
+        )
+        
+        sentence_stamps = self.sentence_timestamps(token_stamps)
+        lrc_text = self.format_lrc(sentence_stamps)
+        
+        return {
+            "token_timestamps": token_stamps,
+            "sentence_timestamps": sentence_stamps,
+            "lrc_text": lrc_text
+        }
+
+
+class MusicLyricScorer:
+    """
+    Scorer class for evaluating lyrics-to-audio alignment quality.
+
+    Focuses on calculating alignment quality metrics (Coverage, Monotonicity, Confidence)
+    using tensor operations for potential differentiability or GPU acceleration.
+    """
+
+    def __init__(self, tokenizer: Any):
+        """
+        Initialize the aligner.
+
+        Args:
+            tokenizer: Tokenizer instance (must implement .decode()).
+        """
+        self.tokenizer = tokenizer
+
+    def _generate_token_type_mask(self, token_ids: List[int]) -> np.ndarray:
+        """
+        Generate a mask distinguishing lyrics (1) from structural tags (0).
+        Uses self.tokenizer to decode tokens.
+
+        Args:
+            token_ids: List of token IDs.
+
+        Returns:
+            Numpy array of shape [len(token_ids)] with 1 or 0.
+        """
+        decoded_tokens = [self.tokenizer.decode([tid]) for tid in token_ids]
+        mask = np.ones(len(token_ids), dtype=np.int32)
+        in_bracket = False
+
+        for i, token_str in enumerate(decoded_tokens):
+            if '[' in token_str:
+                in_bracket = True
+            if in_bracket:
+                mask[i] = 0
+            if ']' in token_str:
+                in_bracket = False
+                mask[i] = 0
+        return mask
+
+    def _preprocess_attention(
+            self,
+            attention_matrix: Union[torch.Tensor, np.ndarray],
+            custom_config: Dict[int, List[int]],
+            medfilt_width: int = 1
+    ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[torch.Tensor]]:
+        """
+        Extracts and normalizes the attention matrix.
+
+        Logic V4: Uses Min-Max normalization to highlight energy differences.
+
+        Args:
+            attention_matrix: Raw attention tensor [Layers, Heads, Tokens, Frames].
+            custom_config: Config mapping layers to heads.
+            medfilt_width: Width for median filtering.
+
+        Returns:
+            Tuple of (calc_matrix, energy_matrix, avg_weights_tensor).
+        """
+        # 1. Prepare Tensor
+        if not isinstance(attention_matrix, torch.Tensor):
+            weights = torch.tensor(attention_matrix)
+        else:
+            weights = attention_matrix.clone()
+        weights = weights.cpu().float()
+
+        # 2. Select Heads based on config
+        selected_tensors = []
+        for layer_idx, head_indices in custom_config.items():
+            for head_idx in head_indices:
+                if layer_idx < weights.shape[0] and head_idx < weights.shape[1]:
+                    selected_tensors.append(weights[layer_idx, head_idx])
+
+        if not selected_tensors:
+            return None, None, None
+
+        weights_stack = torch.stack(selected_tensors, dim=0)
+
+        # 3. Average Heads
+        avg_weights = weights_stack.mean(dim=0)  # [Tokens, Frames]
+
+        # 4. Preprocessing Logic
+        # Min-Max normalization preserving energy distribution
+        # Median filter is applied to the energy matrix
+        energy_tensor = median_filter(avg_weights, filter_width=medfilt_width)
+        energy_matrix = energy_tensor.numpy()
+
+        e_min, e_max = energy_matrix.min(), energy_matrix.max()
+
+        if e_max - e_min > 1e-9:
+            energy_matrix = (energy_matrix - e_min) / (e_max - e_min)
+        else:
+            energy_matrix = np.zeros_like(energy_matrix)
+
+        # Contrast enhancement for DTW pathfinding
+        # calc_matrix is used for pathfinding, energy_matrix for scoring
+        calc_matrix = energy_matrix ** 2
+
+        return calc_matrix, energy_matrix, avg_weights
+
+    def _compute_alignment_metrics(
+            self,
+            energy_matrix: torch.Tensor,
+            path_coords: torch.Tensor,
+            type_mask: torch.Tensor,
+            time_weight: float = 0.01,
+            overlap_frames: float = 9.0,
+            instrumental_weight: float = 1.0
+    ) -> Tuple[float, float, float]:
+        """
+        Core metric calculation logic using high-precision Tensor operations.
+
+        Args:
+            energy_matrix: Normalized energy [Rows, Cols].
+            path_coords: DTW path coordinates [Steps, 2].
+            type_mask: Token type mask [Rows] (1=Lyrics, 0=Tags).
+            time_weight: Minimum energy threshold for monotonicity.
+            overlap_frames: Allowed overlap for monotonicity check.
+            instrumental_weight: Weight for non-lyric tokens in confidence calc.
+
+        Returns:
+            Tuple of (coverage, monotonicity, confidence).
+        """
+        # Ensure high precision for internal calculation
+        energy_matrix = energy_matrix.to(dtype=torch.float64)
+        path_coords = path_coords.long()
+        type_mask = type_mask.long()
+
+        device = energy_matrix.device
+        rows, cols = energy_matrix.shape
+
+        is_lyrics_row = (type_mask == 1)
+
+        # ================= A. Coverage Score =================
+        # Ratio of lyric lines that have significant energy peak
+        row_max_energies = energy_matrix.max(dim=1).values
+        total_sung_rows = is_lyrics_row.sum().double()
+
+        coverage_threshold = 0.1
+        valid_sung_mask = is_lyrics_row & (row_max_energies > coverage_threshold)
+        valid_sung_rows = valid_sung_mask.sum().double()
+
+        if total_sung_rows > 0:
+            coverage_score = valid_sung_rows / total_sung_rows
+        else:
+            coverage_score = torch.tensor(1.0, device=device, dtype=torch.float64)
+
+        # ================= B. Monotonicity Score =================
+        # Check if the "center of mass" of lyric lines moves forward in time
+        col_indices = torch.arange(cols, device=device, dtype=torch.float64)
+
+        # Zero out low energy noise
+        weights = torch.where(
+            energy_matrix > time_weight,
+            energy_matrix,
+            torch.zeros_like(energy_matrix)
+        )
+
+        sum_w = weights.sum(dim=1)
+        sum_t = (weights * col_indices).sum(dim=1)
+
+        # Calculate centroids
+        centroids = torch.full((rows,), -1.0, device=device, dtype=torch.float64)
+        valid_w_mask = sum_w > 1e-9
+        centroids[valid_w_mask] = sum_t[valid_w_mask] / sum_w[valid_w_mask]
+
+        # Extract sequence of valid lyrics centroids
+        valid_sequence_mask = is_lyrics_row & (centroids >= 0)
+        sung_centroids = centroids[valid_sequence_mask]
+
+        cnt = sung_centroids.shape[0]
+        if cnt > 1:
+            curr_c = sung_centroids[:-1]
+            next_c = sung_centroids[1:]
+
+            # Check non-decreasing order with overlap tolerance
+            non_decreasing = (next_c >= (curr_c - overlap_frames)).double().sum()
+            pairs = torch.tensor(cnt - 1, device=device, dtype=torch.float64)
+            monotonicity_score = non_decreasing / pairs
+        else:
+            monotonicity_score = torch.tensor(1.0, device=device, dtype=torch.float64)
+
+        # ================= C. Path Confidence =================
+        # Average energy along the optimal path
+        if path_coords.shape[0] > 0:
+            p_rows = path_coords[:, 0]
+            p_cols = path_coords[:, 1]
+
+            path_energies = energy_matrix[p_rows, p_cols]
+            step_weights = torch.ones_like(path_energies)
+
+            # Lower weight for instrumental/tag steps
+            is_inst_step = (type_mask[p_rows] == 0)
+            step_weights[is_inst_step] = instrumental_weight
+
+            total_energy = (path_energies * step_weights).sum()
+            total_steps = step_weights.sum()
+
+            if total_steps > 0:
+                path_confidence = total_energy / total_steps
+            else:
+                path_confidence = torch.tensor(0.0, device=device, dtype=torch.float64)
+        else:
+            path_confidence = torch.tensor(0.0, device=device, dtype=torch.float64)
+
+        return coverage_score.item(), monotonicity_score.item(), path_confidence.item()
+
+    def lyrics_alignment_info(
+            self,
+            attention_matrix: Union[torch.Tensor, np.ndarray],
+            token_ids: List[int],
+            custom_config: Dict[int, List[int]],
+            return_matrices: bool = False,
+            medfilt_width: int = 1
+    ) -> Dict[str, Any]:
+        """
+        Generates alignment path and processed matrices.
+
+        Args:
+            attention_matrix: Input attention tensor.
+            token_ids: Corresponding token IDs.
+            custom_config: Layer/Head configuration.
+            return_matrices: If True, returns matrices in the output.
+            medfilt_width: Median filter width.
+
+        Returns:
+            Dict or AlignmentInfo object containing path and masks.
+        """
+        calc_matrix, energy_matrix, vis_matrix = self._preprocess_attention(
+            attention_matrix, custom_config, medfilt_width
+        )
+
+        if calc_matrix is None:
+            return {
+                "calc_matrix": None,
+                "error": "No valid attention heads found"
+            }
+
+        # 1. Generate Semantic Mask (1=Lyrics, 0=Tags)
+        # Uses self.tokenizer internally
+        type_mask = self._generate_token_type_mask(token_ids)
+
+        # Safety check for shape mismatch
+        if len(type_mask) != energy_matrix.shape[0]:
+            # Fallback to all lyrics if shapes don't align
+            type_mask = np.ones(energy_matrix.shape[0], dtype=np.int32)
+
+        # 2. DTW Pathfinding
+        # Using negative calc_matrix because DTW minimizes cost
+        text_indices, time_indices = dtw_cpu(-calc_matrix.astype(np.float32))
+        path_coords = np.stack([text_indices, time_indices], axis=1)
+
+        return_dict = {
+            "path_coords": path_coords,
+            "type_mask": type_mask,
+            "energy_matrix": energy_matrix
+        }
+        if return_matrices:
+            return_dict['calc_matrix'] = calc_matrix
+            return_dict['vis_matrix'] = vis_matrix
+
+        return return_dict
+
+    def calculate_score(
+            self,
+            energy_matrix: Union[torch.Tensor, np.ndarray],
+            type_mask: Union[torch.Tensor, np.ndarray],
+            path_coords: Union[torch.Tensor, np.ndarray],
+            time_weight: float = 0.01,
+            overlap_frames: float = 9.0,
+            instrumental_weight: float = 1.0
+    ) -> Dict[str, Any]:
+        """
+        Calculates the final alignment score based on pre-computed components.
+
+        Args:
+            energy_matrix: Processed energy matrix.
+            type_mask: Token type mask.
+            path_coords: DTW path coordinates.
+            time_weight: Minimum energy threshold for monotonicity.
+            overlap_frames: Allowed backward movement frames.
+            instrumental_weight: Weight for non-lyric path steps.
+
+        Returns:
+            AlignmentScore object containing individual metrics and final score.
+        """
+        # Ensure Inputs are Tensors on the correct device
+        if not isinstance(energy_matrix, torch.Tensor):
+            # Use available accelerator device; fallback to CPU if none
+            if torch.cuda.is_available():
+                _score_device = "cuda"
+            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+                _score_device = "mps"
+            else:
+                _score_device = "cpu"
+            energy_matrix = torch.tensor(energy_matrix, device=_score_device, dtype=torch.float32)
+
+        device = energy_matrix.device
+
+        if not isinstance(type_mask, torch.Tensor):
+            type_mask = torch.tensor(type_mask, device=device, dtype=torch.long)
+        else:
+            type_mask = type_mask.to(device=device, dtype=torch.long)
+
+        if not isinstance(path_coords, torch.Tensor):
+            path_coords = torch.tensor(path_coords, device=device, dtype=torch.long)
+        else:
+            path_coords = path_coords.to(device=device, dtype=torch.long)
+
+        # Compute Metrics
+        coverage, monotonicity, confidence = self._compute_alignment_metrics(
+            energy_matrix=energy_matrix,
+            path_coords=path_coords,
+            type_mask=type_mask,
+            time_weight=time_weight,
+            overlap_frames=overlap_frames,
+            instrumental_weight=instrumental_weight
+        )
+
+        # Final Score Calculation
+        # (Cov^2 * Mono^2 * Conf)
+        final_score = (coverage ** 2) * (monotonicity ** 2) * confidence
+        final_score = float(np.clip(final_score, 0.0, 1.0))
+
+        return {
+            "lyrics_score": round(final_score, 4)
+        }

acestep/gpu_config.py

@@ -0,0 +1,549 @@
+"""
+GPU Configuration Module
+Centralized GPU memory detection and adaptive configuration management
+
+    Debug Mode:
+        Set environment variable MAX_CUDA_VRAM to simulate different GPU memory sizes.
+        Example: MAX_CUDA_VRAM=8 python acestep  # Simulates 8GB GPU
+        
+        For MPS testing, use MAX_MPS_VRAM to simulate MPS memory.
+        Example: MAX_MPS_VRAM=16 python acestep  # Simulates 16GB MPS
+    
+    This is useful for testing GPU tier configurations on high-end hardware.
+"""
+
+import os
+import sys
+from dataclasses import dataclass
+from typing import Optional, List, Dict, Tuple
+from loguru import logger
+
+
+# Environment variable for debugging/testing different GPU memory configurations
+DEBUG_MAX_CUDA_VRAM_ENV = "MAX_CUDA_VRAM"
+DEBUG_MAX_MPS_VRAM_ENV = "MAX_MPS_VRAM"
+
+# Tolerance for 16GB detection: reported VRAM like 15.5GB is effectively 16GB hardware
+# Real-world 16GB GPUs often report 15.7-15.9GB due to system/driver reservations
+VRAM_16GB_TOLERANCE_GB = 0.5
+VRAM_16GB_MIN_GB = 16.0 - VRAM_16GB_TOLERANCE_GB  # treat as 16GB class if >= this
+
+# PyTorch installation URLs for diagnostics
+PYTORCH_CUDA_INSTALL_URL = "https://download.pytorch.org/whl/cu121"
+PYTORCH_ROCM_INSTALL_URL = "https://download.pytorch.org/whl/rocm6.0"
+
+
+@dataclass
+class GPUConfig:
+    """GPU configuration based on available memory"""
+    tier: str  # "tier1", "tier2", etc. or "unlimited"
+    gpu_memory_gb: float
+    
+    # Duration limits (in seconds)
+    max_duration_with_lm: int  # When LM is initialized
+    max_duration_without_lm: int  # When LM is not initialized
+    
+    # Batch size limits
+    max_batch_size_with_lm: int
+    max_batch_size_without_lm: int
+    
+    # LM configuration
+    init_lm_default: bool  # Whether to initialize LM by default
+    available_lm_models: List[str]  # Available LM models for this tier
+    
+    # LM memory allocation (GB) for each model size
+    lm_memory_gb: Dict[str, float]  # e.g., {"0.6B": 3, "1.7B": 8, "4B": 12}
+
+
+# GPU tier configurations
+GPU_TIER_CONFIGS = {
+    "tier1": {  # <= 4GB
+        "max_duration_with_lm": 180,  # 3 minutes
+        "max_duration_without_lm": 180,  # 3 minutes
+        "max_batch_size_with_lm": 1,
+        "max_batch_size_without_lm": 1,
+        "init_lm_default": False,
+        "available_lm_models": [],
+        "lm_memory_gb": {},
+    },
+    "tier2": {  # 4-6GB
+        "max_duration_with_lm": 360,  # 6 minutes
+        "max_duration_without_lm": 360,  # 6 minutes
+        "max_batch_size_with_lm": 1,
+        "max_batch_size_without_lm": 1,
+        "init_lm_default": False,
+        "available_lm_models": [],
+        "lm_memory_gb": {},
+    },
+    "tier3": {  # 6-8GB
+        "max_duration_with_lm": 240,  # 4 minutes with LM
+        "max_duration_without_lm": 360,  # 6 minutes without LM
+        "max_batch_size_with_lm": 1,
+        "max_batch_size_without_lm": 2,
+        "init_lm_default": False,  # Don't init by default due to limited memory
+        "available_lm_models": ["acestep-5Hz-lm-0.6B"],
+        "lm_memory_gb": {"0.6B": 3},
+    },
+    "tier4": {  # 8-12GB
+        "max_duration_with_lm": 240,  # 4 minutes with LM
+        "max_duration_without_lm": 360,  # 6 minutes without LM
+        "max_batch_size_with_lm": 2,
+        "max_batch_size_without_lm": 4,
+        "init_lm_default": False,  # Don't init by default
+        "available_lm_models": ["acestep-5Hz-lm-0.6B"],
+        "lm_memory_gb": {"0.6B": 3},
+    },
+    "tier5": {  # 12-16GB
+        "max_duration_with_lm": 240,  # 4 minutes with LM
+        "max_duration_without_lm": 360,  # 6 minutes without LM
+        "max_batch_size_with_lm": 2,
+        "max_batch_size_without_lm": 4,
+        "init_lm_default": True,
+        "available_lm_models": ["acestep-5Hz-lm-0.6B", "acestep-5Hz-lm-1.7B"],
+        "lm_memory_gb": {"0.6B": 3, "1.7B": 8},
+    },
+    "tier6": {  # 16-24GB
+        "max_duration_with_lm": 480,  # 8 minutes
+        "max_duration_without_lm": 480,  # 8 minutes
+        "max_batch_size_with_lm": 4,
+        "max_batch_size_without_lm": 8,
+        "init_lm_default": True,
+        "available_lm_models": ["acestep-5Hz-lm-0.6B", "acestep-5Hz-lm-1.7B", "acestep-5Hz-lm-4B"],
+        "lm_memory_gb": {"0.6B": 3, "1.7B": 8, "4B": 12},
+    },
+    "unlimited": {  # >= 24GB
+        "max_duration_with_lm": 600,  # 10 minutes (max supported)
+        "max_duration_without_lm": 600,  # 10 minutes
+        "max_batch_size_with_lm": 8,
+        "max_batch_size_without_lm": 8,
+        "init_lm_default": True,
+        "available_lm_models": ["acestep-5Hz-lm-0.6B", "acestep-5Hz-lm-1.7B", "acestep-5Hz-lm-4B"],
+        "lm_memory_gb": {"0.6B": 3, "1.7B": 8, "4B": 12},
+    },
+}
+
+
+def get_gpu_memory_gb() -> float:
+    """
+    Get GPU memory in GB. Returns 0 if no GPU is available.
+    
+    Debug Mode:
+        Set environment variable MAX_CUDA_VRAM to override the detected GPU memory.
+        Example: MAX_CUDA_VRAM=8 python acestep  # Simulates 8GB GPU
+        
+        For MPS testing, set MAX_MPS_VRAM to override MPS memory detection.
+        Example: MAX_MPS_VRAM=16 python acestep  # Simulates 16GB MPS
+        
+        This allows testing different GPU tier configurations on high-end hardware.
+    """
+    # Check for debug override first
+    debug_vram = os.environ.get(DEBUG_MAX_CUDA_VRAM_ENV)
+    if debug_vram is not None:
+        try:
+            simulated_gb = float(debug_vram)
+            logger.warning(f"⚠️ DEBUG MODE: Simulating GPU memory as {simulated_gb:.1f}GB (set via {DEBUG_MAX_CUDA_VRAM_ENV} environment variable)")
+            return simulated_gb
+        except ValueError:
+            logger.warning(f"Invalid {DEBUG_MAX_CUDA_VRAM_ENV} value: {debug_vram}, ignoring")
+    debug_mps_vram = os.environ.get(DEBUG_MAX_MPS_VRAM_ENV)
+    if debug_mps_vram is not None:
+        try:
+            simulated_gb = float(debug_mps_vram)
+            logger.warning(f"⚠️ DEBUG MODE: Simulating MPS memory as {simulated_gb:.1f}GB (set via {DEBUG_MAX_MPS_VRAM_ENV} environment variable)")
+            return simulated_gb
+        except ValueError:
+            logger.warning(f"Invalid {DEBUG_MAX_MPS_VRAM_ENV} value: {debug_mps_vram}, ignoring")
+    
+    try:
+        import torch
+        if torch.cuda.is_available():
+            # Get total memory of the first GPU in GB
+            total_memory = torch.cuda.get_device_properties(0).total_memory
+            memory_gb = total_memory / (1024**3)  # Convert bytes to GB
+            device_name = torch.cuda.get_device_name(0)
+            is_rocm = hasattr(torch.version, 'hip') and torch.version.hip is not None
+            if is_rocm:
+                logger.info(f"ROCm GPU detected: {device_name} ({memory_gb:.1f} GB, HIP {torch.version.hip})")
+            else:
+                logger.info(f"CUDA GPU detected: {device_name} ({memory_gb:.1f} GB)")
+            return memory_gb
+        elif hasattr(torch, 'xpu') and torch.xpu.is_available():
+            # Get total memory of the first XPU in GB
+            total_memory = torch.xpu.get_device_properties(0).total_memory
+            memory_gb = total_memory / (1024**3)  # Convert bytes to GB
+            return memory_gb
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            mps_module = getattr(torch, "mps", None)
+            try:
+                if mps_module is not None and hasattr(mps_module, "recommended_max_memory"):
+                    total_memory = mps_module.recommended_max_memory()
+                    memory_gb = total_memory / (1024**3)  # Convert bytes to GB
+                    return memory_gb
+                if mps_module is not None and hasattr(mps_module, "get_device_properties"):
+                    props = mps_module.get_device_properties(0)
+                    total_memory = getattr(props, "total_memory", None)
+                    if total_memory:
+                        memory_gb = total_memory / (1024**3)
+                        return memory_gb
+            except Exception as e:
+                logger.warning(f"Failed to detect MPS memory: {e}")
+
+            # Fallback: estimate from system unified memory (Apple Silicon shares CPU/GPU RAM)
+            try:
+                import subprocess
+                result = subprocess.run(
+                    ["sysctl", "-n", "hw.memsize"],
+                    capture_output=True, text=True, timeout=5
+                )
+                total_system_bytes = int(result.stdout.strip())
+                # MPS can use up to ~75% of unified memory for GPU workloads
+                memory_gb = (total_system_bytes / (1024**3)) * 0.75
+                return memory_gb
+            except Exception:
+                logger.warning(f"MPS available but total memory not exposed. Set {DEBUG_MAX_MPS_VRAM_ENV} to enable tiering.")
+                # Conservative fallback for M1/M2
+                return 8.0
+        else:
+            # No GPU detected - provide diagnostic information
+            _log_gpu_diagnostic_info(torch)
+            return 0
+    except Exception as e:
+        logger.warning(f"Failed to detect GPU memory: {e}")
+        return 0
+
+
+def _log_gpu_diagnostic_info(torch_module):
+    """
+    Log diagnostic information when GPU is not detected to help users troubleshoot.
+    
+    Args:
+        torch_module: The torch module to inspect for build information
+    """
+    logger.warning("=" * 80)
+    logger.warning("⚠️ GPU NOT DETECTED - DIAGNOSTIC INFORMATION")
+    logger.warning("=" * 80)
+    
+    # Check PyTorch build type
+    is_rocm_build = hasattr(torch_module.version, 'hip') and torch_module.version.hip is not None
+    is_cuda_build = hasattr(torch_module.version, 'cuda') and torch_module.version.cuda is not None
+    
+    if is_rocm_build:
+        logger.warning("✓ PyTorch ROCm build detected")
+        logger.warning(f"  HIP version: {torch_module.version.hip}")
+        logger.warning("")
+        logger.warning("❌ torch.cuda.is_available() returned False")
+        logger.warning("")
+        logger.warning("Common causes for AMD/ROCm GPUs:")
+        logger.warning("  1. ROCm drivers not installed or not properly configured")
+        logger.warning("  2. GPU not supported by installed ROCm version")
+        logger.warning("  3. Missing or incorrect HSA_OVERRIDE_GFX_VERSION environment variable")
+        logger.warning("  4. ROCm runtime libraries not in system path")
+        logger.warning("")
+        
+        # Check for common environment variables
+        hsa_override = os.environ.get('HSA_OVERRIDE_GFX_VERSION')
+        if hsa_override:
+            logger.warning(f"  HSA_OVERRIDE_GFX_VERSION is set to: {hsa_override}")
+        else:
+            logger.warning("  ⚠️ HSA_OVERRIDE_GFX_VERSION is not set")
+            logger.warning("     For RDNA3 GPUs (RX 7000 series, RX 9000 series):")
+            logger.warning("       - RX 7900 XT/XTX, RX 9070 XT: set HSA_OVERRIDE_GFX_VERSION=11.0.0")
+            logger.warning("       - RX 7800 XT, RX 7700 XT: set HSA_OVERRIDE_GFX_VERSION=11.0.1")
+            logger.warning("       - RX 7600: set HSA_OVERRIDE_GFX_VERSION=11.0.2")
+        
+        logger.warning("")
+        logger.warning("Troubleshooting steps:")
+        logger.warning("  1. Verify ROCm installation:")
+        logger.warning("     rocm-smi  # Should list your GPU")
+        logger.warning("  2. Check PyTorch ROCm build:")
+        logger.warning("     python -c \"import torch; print(f'ROCm: {torch.version.hip}')\"")
+        logger.warning("  3. Set HSA_OVERRIDE_GFX_VERSION for your GPU (see above)")
+        logger.warning("  4. On Windows: Use start_gradio_ui_rocm.bat which sets required env vars")
+        logger.warning("  5. See docs/en/ACE-Step1.5-Rocm-Manual-Linux.md for Linux setup")
+        logger.warning("  6. See requirements-rocm.txt for Windows ROCm setup instructions")
+        
+    elif is_cuda_build:
+        logger.warning("✓ PyTorch CUDA build detected")
+        logger.warning(f"  CUDA version: {torch_module.version.cuda}")
+        logger.warning("")
+        logger.warning("❌ torch.cuda.is_available() returned False")
+        logger.warning("")
+        logger.warning("Common causes for NVIDIA GPUs:")
+        logger.warning("  1. NVIDIA drivers not installed")
+        logger.warning("  2. CUDA runtime not installed or version mismatch")
+        logger.warning("  3. GPU not supported by installed CUDA version")
+        logger.warning("")
+        logger.warning("Troubleshooting steps:")
+        logger.warning("  1. Verify NVIDIA driver installation:")
+        logger.warning("     nvidia-smi  # Should list your GPU")
+        logger.warning("  2. Check CUDA version compatibility")
+        logger.warning("  3. Reinstall PyTorch with CUDA support:")
+        logger.warning(f"     pip install torch --index-url {PYTORCH_CUDA_INSTALL_URL}")
+        
+    else:
+        logger.warning("⚠️ PyTorch build type: CPU-only")
+        logger.warning("")
+        logger.warning("You have installed a CPU-only version of PyTorch!")
+        logger.warning("")
+        logger.warning("For NVIDIA GPUs:")
+        logger.warning(f"  pip install torch --index-url {PYTORCH_CUDA_INSTALL_URL}")
+        logger.warning("")
+        logger.warning("For AMD GPUs with ROCm:")
+        logger.warning("  Windows: See requirements-rocm.txt for detailed instructions")
+        logger.warning(f"  Linux: pip install torch --index-url {PYTORCH_ROCM_INSTALL_URL}")
+        logger.warning("")
+        logger.warning("For more information, see README.md section 'AMD / ROCm GPUs'")
+    
+    logger.warning("=" * 80)
+
+
+def get_gpu_tier(gpu_memory_gb: float) -> str:
+    """
+    Determine GPU tier based on available memory.
+    
+    Args:
+        gpu_memory_gb: GPU memory in GB
+        
+    Returns:
+        Tier string: "tier1", "tier2", "tier3", "tier4", "tier5", "tier6", or "unlimited"
+    """
+    if gpu_memory_gb <= 0:
+        # CPU mode - use tier1 limits
+        return "tier1"
+    elif gpu_memory_gb <= 4:
+        return "tier1"
+    elif gpu_memory_gb <= 6:
+        return "tier2"
+    elif gpu_memory_gb <= 8:
+        return "tier3"
+    elif gpu_memory_gb <= 12:
+        return "tier4"
+    elif gpu_memory_gb < VRAM_16GB_MIN_GB:
+        return "tier5"
+    elif gpu_memory_gb <= 24:
+        if gpu_memory_gb < 16.0:
+            logger.info(f"Detected {gpu_memory_gb:.2f}GB VRAM — treating as 16GB class GPU")
+        return "tier6"
+    else:
+        return "unlimited"
+
+
+def get_gpu_config(gpu_memory_gb: Optional[float] = None) -> GPUConfig:
+    """
+    Get GPU configuration based on detected or provided GPU memory.
+    
+    Args:
+        gpu_memory_gb: GPU memory in GB. If None, will be auto-detected.
+        
+    Returns:
+        GPUConfig object with all configuration parameters
+    """
+    if gpu_memory_gb is None:
+        gpu_memory_gb = get_gpu_memory_gb()
+    
+    tier = get_gpu_tier(gpu_memory_gb)
+    config = GPU_TIER_CONFIGS[tier]
+    
+    return GPUConfig(
+        tier=tier,
+        gpu_memory_gb=gpu_memory_gb,
+        max_duration_with_lm=config["max_duration_with_lm"],
+        max_duration_without_lm=config["max_duration_without_lm"],
+        max_batch_size_with_lm=config["max_batch_size_with_lm"],
+        max_batch_size_without_lm=config["max_batch_size_without_lm"],
+        init_lm_default=config["init_lm_default"],
+        available_lm_models=config["available_lm_models"],
+        lm_memory_gb=config["lm_memory_gb"],
+    )
+
+
+def get_lm_model_size(model_path: str) -> str:
+    """
+    Extract LM model size from model path.
+    
+    Args:
+        model_path: Model path string (e.g., "acestep-5Hz-lm-0.6B")
+        
+    Returns:
+        Model size string: "0.6B", "1.7B", or "4B"
+    """
+    if "0.6B" in model_path:
+        return "0.6B"
+    elif "1.7B" in model_path:
+        return "1.7B"
+    elif "4B" in model_path:
+        return "4B"
+    else:
+        # Default to smallest model assumption
+        return "0.6B"
+
+
+def get_lm_gpu_memory_ratio(model_path: str, total_gpu_memory_gb: float) -> Tuple[float, float]:
+    """
+    Calculate GPU memory utilization ratio for LM model.
+    
+    Args:
+        model_path: LM model path (e.g., "acestep-5Hz-lm-0.6B")
+        total_gpu_memory_gb: Total GPU memory in GB
+        
+    Returns:
+        Tuple of (gpu_memory_utilization_ratio, target_memory_gb)
+    """
+    model_size = get_lm_model_size(model_path)
+    
+    # Target memory allocation for each model size
+    target_memory = {
+        "0.6B": 3.0,
+        "1.7B": 8.0,
+        "4B": 12.0,
+    }
+    
+    target_gb = target_memory.get(model_size, 3.0)
+    
+    # For large GPUs (>=24GB), don't restrict memory too much
+    if total_gpu_memory_gb >= 24:
+        # Use a reasonable ratio that allows the model to run efficiently
+        ratio = min(0.9, max(0.2, target_gb / total_gpu_memory_gb))
+    else:
+        # For smaller GPUs, strictly limit memory usage
+        ratio = min(0.9, max(0.1, target_gb / total_gpu_memory_gb))
+    
+    return ratio, target_gb
+
+
+def check_duration_limit(
+    duration: float,
+    gpu_config: GPUConfig,
+    lm_initialized: bool
+) -> Tuple[bool, str]:
+    """
+    Check if requested duration is within limits for current GPU configuration.
+    
+    Args:
+        duration: Requested duration in seconds
+        gpu_config: Current GPU configuration
+        lm_initialized: Whether LM is initialized
+        
+    Returns:
+        Tuple of (is_valid, warning_message)
+    """
+    max_duration = gpu_config.max_duration_with_lm if lm_initialized else gpu_config.max_duration_without_lm
+    
+    if duration > max_duration:
+        warning_msg = (
+            f"⚠️ Requested duration ({duration:.0f}s) exceeds the limit for your GPU "
+            f"({gpu_config.gpu_memory_gb:.1f}GB). Maximum allowed: {max_duration}s "
+            f"({'with' if lm_initialized else 'without'} LM). "
+            f"Duration will be clamped to {max_duration}s."
+        )
+        return False, warning_msg
+    
+    return True, ""
+
+
+def check_batch_size_limit(
+    batch_size: int,
+    gpu_config: GPUConfig,
+    lm_initialized: bool
+) -> Tuple[bool, str]:
+    """
+    Check if requested batch size is within limits for current GPU configuration.
+    
+    Args:
+        batch_size: Requested batch size
+        gpu_config: Current GPU configuration
+        lm_initialized: Whether LM is initialized
+        
+    Returns:
+        Tuple of (is_valid, warning_message)
+    """
+    max_batch_size = gpu_config.max_batch_size_with_lm if lm_initialized else gpu_config.max_batch_size_without_lm
+    
+    if batch_size > max_batch_size:
+        warning_msg = (
+            f"⚠️ Requested batch size ({batch_size}) exceeds the limit for your GPU "
+            f"({gpu_config.gpu_memory_gb:.1f}GB). Maximum allowed: {max_batch_size} "
+            f"({'with' if lm_initialized else 'without'} LM). "
+            f"Batch size will be clamped to {max_batch_size}."
+        )
+        return False, warning_msg
+    
+    return True, ""
+
+
+def is_lm_model_supported(model_path: str, gpu_config: GPUConfig) -> Tuple[bool, str]:
+    """
+    Check if the specified LM model is supported for current GPU configuration.
+    
+    Args:
+        model_path: LM model path
+        gpu_config: Current GPU configuration
+        
+    Returns:
+        Tuple of (is_supported, warning_message)
+    """
+    if not gpu_config.available_lm_models:
+        return False, (
+            f"⚠️ Your GPU ({gpu_config.gpu_memory_gb:.1f}GB) does not have enough memory "
+            f"to run any LM model. Please disable LM initialization."
+        )
+    
+    model_size = get_lm_model_size(model_path)
+    
+    # Check if model size is in available models
+    for available_model in gpu_config.available_lm_models:
+        if model_size in available_model:
+            return True, ""
+    
+    return False, (
+        f"⚠️ LM model {model_path} ({model_size}) is not supported for your GPU "
+        f"({gpu_config.gpu_memory_gb:.1f}GB). Available models: {', '.join(gpu_config.available_lm_models)}"
+    )
+
+
+def get_recommended_lm_model(gpu_config: GPUConfig) -> Optional[str]:
+    """
+    Get recommended LM model for current GPU configuration.
+    
+    Args:
+        gpu_config: Current GPU configuration
+        
+    Returns:
+        Recommended LM model path, or None if LM is not supported
+    """
+    if not gpu_config.available_lm_models:
+        return None
+    
+    # Return the largest available model (last in the list)
+    return gpu_config.available_lm_models[-1]
+
+
+def print_gpu_config_info(gpu_config: GPUConfig):
+    """Print GPU configuration information for debugging."""
+    logger.info(f"GPU Configuration:")
+    logger.info(f"  - GPU Memory: {gpu_config.gpu_memory_gb:.1f} GB")
+    logger.info(f"  - Tier: {gpu_config.tier}")
+    logger.info(f"  - Max Duration (with LM): {gpu_config.max_duration_with_lm}s ({gpu_config.max_duration_with_lm // 60} min)")
+    logger.info(f"  - Max Duration (without LM): {gpu_config.max_duration_without_lm}s ({gpu_config.max_duration_without_lm // 60} min)")
+    logger.info(f"  - Max Batch Size (with LM): {gpu_config.max_batch_size_with_lm}")
+    logger.info(f"  - Max Batch Size (without LM): {gpu_config.max_batch_size_without_lm}")
+    logger.info(f"  - Init LM by Default: {gpu_config.init_lm_default}")
+    logger.info(f"  - Available LM Models: {gpu_config.available_lm_models or 'None'}")
+
+
+# Global GPU config instance (initialized lazily)
+_global_gpu_config: Optional[GPUConfig] = None
+
+
+def get_global_gpu_config() -> GPUConfig:
+    """Get the global GPU configuration, initializing if necessary."""
+    global _global_gpu_config
+    if _global_gpu_config is None:
+        _global_gpu_config = get_gpu_config()
+    return _global_gpu_config
+
+
+def set_global_gpu_config(config: GPUConfig):
+    """Set the global GPU configuration."""
+    global _global_gpu_config
+    _global_gpu_config = config

acestep/gradio_ui/__init__.py

@@ -0,0 +1 @@
+from acestep.gradio_ui.interfaces import create_gradio_interface

acestep/gradio_ui/api_routes.py

@@ -0,0 +1,564 @@
+"""
+Gradio API Routes Module
+Add API endpoints compatible with api_server.py and CustomAceStep to Gradio application
+"""
+import json
+import os
+import random
+import time
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
+
+from fastapi import APIRouter, HTTPException, Request, Depends, Header
+from fastapi.responses import FileResponse
+
+# Global results directory inside project root
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+DEFAULT_RESULTS_DIR = os.path.join(PROJECT_ROOT, "gradio_outputs").replace("\\", "/")
+os.makedirs(DEFAULT_RESULTS_DIR, exist_ok=True)
+
+# API Key storage (set via setup_api_routes)
+_api_key: Optional[str] = None
+
+
+def set_api_key(key: Optional[str]):
+    """Set the API key for authentication"""
+    global _api_key
+    _api_key = key
+
+
+def _wrap_response(data: Any, code: int = 200, error: Optional[str] = None) -> Dict[str, Any]:
+    """Wrap response data in standard format compatible with CustomAceStep."""
+    return {
+        "data": data,
+        "code": code,
+        "error": error,
+        "timestamp": int(time.time() * 1000),
+        "extra": None,
+    }
+
+
+def verify_token_from_request(body: dict, authorization: Optional[str] = None) -> Optional[str]:
+    """
+    Verify API key from request body (ai_token) or Authorization header.
+    Returns the token if valid, None if no auth required.
+    """
+    if _api_key is None:
+        return None  # No auth required
+
+    # Try ai_token from body first
+    ai_token = body.get("ai_token") if body else None
+    if ai_token:
+        if ai_token == _api_key:
+            return ai_token
+        raise HTTPException(status_code=401, detail="Invalid ai_token")
+
+    # Fallback to Authorization header
+    if authorization:
+        if authorization.startswith("Bearer "):
+            token = authorization[7:]
+        else:
+            token = authorization
+        if token == _api_key:
+            return token
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+    # No token provided but auth is required
+    raise HTTPException(status_code=401, detail="Missing ai_token or Authorization header")
+
+
+async def verify_api_key(authorization: Optional[str] = Header(None)):
+    """Verify API key from Authorization header (legacy, for non-body endpoints)"""
+    if _api_key is None:
+        return  # No auth required
+
+    if not authorization:
+        raise HTTPException(status_code=401, detail="Missing Authorization header")
+
+    # Support "Bearer <key>" format
+    if authorization.startswith("Bearer "):
+        token = authorization[7:]
+    else:
+        token = authorization
+
+    if token != _api_key:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+
+# Use diskcache to store results
+try:
+    import diskcache
+    _cache_dir = os.path.join(os.path.dirname(__file__), ".cache", "api_results")
+    os.makedirs(_cache_dir, exist_ok=True)
+    _result_cache = diskcache.Cache(_cache_dir)
+    DISKCACHE_AVAILABLE = True
+except ImportError:
+    _result_cache = {}
+    DISKCACHE_AVAILABLE = False
+
+RESULT_EXPIRE_SECONDS = 7 * 24 * 60 * 60  # 7 days expiration
+RESULT_KEY_PREFIX = "ace_step_v1.5_"
+
+# =============================================================================
+# Example Data for Random Sample
+# =============================================================================
+
+def _get_project_root() -> str:
+    """Get project root directory"""
+    return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+def _load_all_examples(sample_mode: str = "simple_mode") -> List[Dict[str, Any]]:
+    """Load all example JSON files from examples directory"""
+    project_root = _get_project_root()
+    if sample_mode == "simple_mode":
+        examples_dir = os.path.join(project_root, "examples", "simple_mode")
+    else:
+        examples_dir = os.path.join(project_root, "examples", "text2music")
+
+    if not os.path.isdir(examples_dir):
+        return []
+
+    all_examples = []
+    for filename in os.listdir(examples_dir):
+        if filename.endswith(".json"):
+            filepath = os.path.join(examples_dir, filename)
+            try:
+                with open(filepath, "r", encoding="utf-8") as f:
+                    data = json.load(f)
+                    if isinstance(data, list):
+                        all_examples.extend(data)
+                    elif isinstance(data, dict):
+                        all_examples.append(data)
+            except Exception:
+                pass
+    return all_examples
+
+
+# Pre-load example data
+SIMPLE_EXAMPLE_DATA = _load_all_examples("simple_mode")
+CUSTOM_EXAMPLE_DATA = _load_all_examples("custom_mode")
+
+
+def store_result(task_id: str, result: dict, status: str = "succeeded"):
+    """Store result to diskcache"""
+    data = {
+        "result": result,
+        "created_at": time.time(),
+        "status": status
+    }
+    key = f"{RESULT_KEY_PREFIX}{task_id}"
+    if DISKCACHE_AVAILABLE:
+        _result_cache.set(key, data, expire=RESULT_EXPIRE_SECONDS)
+    else:
+        _result_cache[key] = data
+
+
+def get_result(task_id: str) -> Optional[dict]:
+    """Get result from diskcache"""
+    key = f"{RESULT_KEY_PREFIX}{task_id}"
+    if DISKCACHE_AVAILABLE:
+        return _result_cache.get(key)
+    else:
+        return _result_cache.get(key)
+
+
+router = APIRouter()
+
+
+@router.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return _wrap_response({
+        "status": "ok",
+        "service": "ACE-Step Gradio API",
+        "version": "1.0",
+    })
+
+
+@router.get("/v1/models")
+async def list_models(request: Request, _: None = Depends(verify_api_key)):
+    """List available DiT models"""
+    dit_handler = request.app.state.dit_handler
+
+    models = []
+    if dit_handler and dit_handler.model is not None:
+        # Get current loaded model name
+        config_path = getattr(dit_handler, 'config_path', '') or ''
+        model_name = os.path.basename(config_path.rstrip("/\\")) if config_path else "unknown"
+        models.append({
+            "name": model_name,
+            "is_default": True,
+        })
+
+    return _wrap_response({
+        "models": models,
+        "default_model": models[0]["name"] if models else None,
+    })
+
+
+@router.get("/v1/audio")
+async def get_audio(path: str, _: None = Depends(verify_api_key)):
+    """Download audio file"""
+    # Security: Validate path is within allowed directory to prevent path traversal
+    resolved_path = os.path.realpath(path)
+    allowed_dir = os.path.realpath(DEFAULT_RESULTS_DIR)
+    if not resolved_path.startswith(allowed_dir + os.sep) and resolved_path != allowed_dir:
+        raise HTTPException(status_code=403, detail="Access denied: path outside allowed directory")
+    if not os.path.exists(resolved_path):
+        raise HTTPException(status_code=404, detail="Audio file not found")
+
+    ext = os.path.splitext(resolved_path)[1].lower()
+    media_types = {
+        ".mp3": "audio/mpeg",
+        ".wav": "audio/wav",
+        ".flac": "audio/flac",
+        ".ogg": "audio/ogg",
+    }
+    media_type = media_types.get(ext, "audio/mpeg")
+
+    return FileResponse(resolved_path, media_type=media_type)
+
+
+@router.post("/create_random_sample")
+async def create_random_sample(request: Request, authorization: Optional[str] = Header(None)):
+    """Get random sample parameters from pre-loaded example data"""
+    content_type = (request.headers.get("content-type") or "").lower()
+
+    if "json" in content_type:
+        body = await request.json()
+    else:
+        form = await request.form()
+        body = {k: v for k, v in form.items()}
+
+    verify_token_from_request(body, authorization)
+    sample_type = body.get("sample_type", "simple_mode") or "simple_mode"
+
+    if sample_type == "simple_mode":
+        example_data = SIMPLE_EXAMPLE_DATA
+    else:
+        example_data = CUSTOM_EXAMPLE_DATA
+
+    if not example_data:
+        return _wrap_response(None, code=500, error="No example data available")
+
+    random_example = random.choice(example_data)
+    return _wrap_response(random_example)
+
+
+@router.post("/query_result")
+async def query_result(request: Request, authorization: Optional[str] = Header(None)):
+    """Batch query task results"""
+    content_type = (request.headers.get("content-type") or "").lower()
+
+    if "json" in content_type:
+        body = await request.json()
+    else:
+        form = await request.form()
+        body = {k: v for k, v in form.items()}
+
+    verify_token_from_request(body, authorization)
+    task_ids = body.get("task_id_list", [])
+
+    if isinstance(task_ids, str):
+        try:
+            task_ids = json.loads(task_ids)
+        except Exception:
+            task_ids = []
+
+    results = []
+    for task_id in task_ids:
+        data = get_result(task_id)
+        if data and data.get("status") == "succeeded":
+            results.append({
+                "task_id": task_id,
+                "status": 1,
+                "result": json.dumps(data["result"], ensure_ascii=False)
+            })
+        else:
+            results.append({
+                "task_id": task_id,
+                "status": 0,
+                "result": "[]"
+            })
+
+    return _wrap_response(results)
+
+
+@router.post("/format_input")
+async def format_input(request: Request, authorization: Optional[str] = Header(None)):
+    """Format and enhance lyrics/caption via LLM"""
+    llm_handler = request.app.state.llm_handler
+
+    if not llm_handler or not llm_handler.llm_initialized:
+        return _wrap_response(None, code=500, error="LLM not initialized")
+
+    content_type = (request.headers.get("content-type") or "").lower()
+    if "json" in content_type:
+        body = await request.json()
+    else:
+        form = await request.form()
+        body = {k: v for k, v in form.items()}
+
+    verify_token_from_request(body, authorization)
+
+    caption = body.get("prompt", "") or ""
+    lyrics = body.get("lyrics", "") or ""
+    temperature = float(body.get("temperature", 0.85))
+
+    from acestep.inference import format_sample
+
+    try:
+        result = format_sample(
+            llm_handler=llm_handler,
+            caption=caption,
+            lyrics=lyrics,
+            temperature=temperature,
+            use_constrained_decoding=True,
+        )
+
+        if not result.success:
+            return _wrap_response(None, code=500, error=result.status_message)
+
+        return _wrap_response({
+            "caption": result.caption or caption,
+            "lyrics": result.lyrics or lyrics,
+            "bpm": result.bpm,
+            "key_scale": result.keyscale,
+            "time_signature": result.timesignature,
+            "duration": result.duration,
+            "vocal_language": result.language or "unknown",
+        })
+    except Exception as e:
+        return _wrap_response(None, code=500, error=str(e))
+
+
+@router.post("/release_task")
+async def release_task(request: Request, authorization: Optional[str] = Header(None)):
+    """Create music generation task"""
+    dit_handler = request.app.state.dit_handler
+    llm_handler = request.app.state.llm_handler
+
+    if not dit_handler or dit_handler.model is None:
+        raise HTTPException(status_code=500, detail="DiT model not initialized")
+
+    content_type = (request.headers.get("content-type") or "").lower()
+    if "json" in content_type:
+        body = await request.json()
+    else:
+        form = await request.form()
+        body = {k: v for k, v in form.items()}
+
+    verify_token_from_request(body, authorization)
+    task_id = str(uuid4())
+
+    from acestep.inference import generate_music, GenerationParams, GenerationConfig, create_sample, format_sample
+
+    # Parse param_obj if provided
+    param_obj = body.get("param_obj", {})
+    if isinstance(param_obj, str):
+        try:
+            param_obj = json.loads(param_obj)
+        except Exception:
+            param_obj = {}
+
+    # Helper to get param with aliases
+    def get_param(key, *aliases, default=None):
+        for k in [key] + list(aliases):
+            if k in body and body[k] is not None:
+                return body[k]
+            if k in param_obj and param_obj[k] is not None:
+                return param_obj[k]
+        return default
+
+    def to_bool(val, default=False):
+        if val is None:
+            return default
+        if isinstance(val, bool):
+            return val
+        if isinstance(val, str):
+            return val.lower() in ("true", "1", "yes")
+        return bool(val)
+
+    try:
+        # Get sample_mode and sample_query parameters
+        sample_mode = to_bool(get_param("sample_mode", "sampleMode"), False)
+        sample_query = get_param("sample_query", "sampleQuery", "description", "desc", default="") or ""
+        use_format = to_bool(get_param("use_format", "useFormat"), False)
+        has_sample_query = bool(sample_query and sample_query.strip())
+
+        # Get base parameters
+        caption = get_param("prompt", "caption", default="") or ""
+        lyrics = get_param("lyrics", default="") or ""
+        vocal_language = get_param("vocal_language", "language", default="en") or "en"
+        lm_temperature = float(get_param("lm_temperature", "temperature", default=0.85) or 0.85)
+
+        # Process sample_mode: use LLM to auto-generate caption/lyrics/metas
+        if sample_mode or has_sample_query:
+            if not llm_handler or not llm_handler.llm_initialized:
+                raise HTTPException(status_code=500, detail="sample_mode requires LLM to be initialized")
+
+            query = sample_query if has_sample_query else "NO USER INPUT"
+            sample_result = create_sample(
+                llm_handler=llm_handler,
+                query=query,
+                vocal_language=vocal_language if vocal_language not in ("en", "unknown", "") else None,
+                temperature=lm_temperature,
+            )
+
+            if not sample_result.success:
+                raise HTTPException(status_code=500, detail=sample_result.error or sample_result.status_message)
+
+            # Use generated values
+            caption = sample_result.caption or caption
+            lyrics = sample_result.lyrics or lyrics
+            # Override metas from sample result if available
+            sample_bpm = sample_result.bpm
+            sample_duration = sample_result.duration
+            sample_keyscale = sample_result.keyscale
+            sample_timesignature = sample_result.timesignature
+            sample_language = sample_result.language or vocal_language
+        else:
+            sample_bpm = None
+            sample_duration = None
+            sample_keyscale = None
+            sample_timesignature = None
+            sample_language = vocal_language
+
+        # Process use_format: enhance caption/lyrics via LLM
+        if use_format and not sample_mode and not has_sample_query:
+            if llm_handler and llm_handler.llm_initialized:
+                format_result = format_sample(
+                    llm_handler=llm_handler,
+                    caption=caption,
+                    lyrics=lyrics,
+                    temperature=lm_temperature,
+                )
+                if format_result.success:
+                    caption = format_result.caption or caption
+                    lyrics = format_result.lyrics or lyrics
+                    if format_result.bpm:
+                        sample_bpm = format_result.bpm
+                    if format_result.duration:
+                        sample_duration = format_result.duration
+                    if format_result.keyscale:
+                        sample_keyscale = format_result.keyscale
+                    if format_result.timesignature:
+                        sample_timesignature = format_result.timesignature
+                    if format_result.language:
+                        sample_language = format_result.language
+
+        # Build generation params with alias support
+        params = GenerationParams(
+            task_type=get_param("task_type", default="text2music"),
+            caption=caption,
+            lyrics=lyrics,
+            bpm=sample_bpm or get_param("bpm"),
+            keyscale=sample_keyscale or get_param("key_scale", "keyscale", "key", default=""),
+            timesignature=sample_timesignature or get_param("time_signature", "timesignature", default=""),
+            duration=sample_duration or get_param("audio_duration", "duration", default=-1),
+            vocal_language=sample_language,
+            inference_steps=get_param("inference_steps", default=8),
+            guidance_scale=float(get_param("guidance_scale", default=7.0) or 7.0),
+            seed=int(get_param("seed", default=-1) or -1),
+            thinking=to_bool(get_param("thinking"), False),
+            lm_temperature=lm_temperature,
+            lm_cfg_scale=float(get_param("lm_cfg_scale", default=2.0) or 2.0),
+            lm_negative_prompt=get_param("lm_negative_prompt", default="NO USER INPUT") or "NO USER INPUT",
+        )
+
+        config = GenerationConfig(
+            batch_size=get_param("batch_size", default=2),
+            use_random_seed=get_param("use_random_seed", default=True),
+            audio_format=get_param("audio_format", default="mp3"),
+        )
+
+        # Get output directory
+        save_dir = os.path.join(DEFAULT_RESULTS_DIR, f"api_{int(time.time())}").replace("\\", "/")
+        os.makedirs(save_dir, exist_ok=True)
+
+        # Call generation function
+        result = generate_music(
+            dit_handler=dit_handler,
+            llm_handler=llm_handler if llm_handler and llm_handler.llm_initialized else None,
+            params=params,
+            config=config,
+            save_dir=save_dir,
+        )
+
+        if not result.success:
+            raise HTTPException(status_code=500, detail=result.error or result.status_message)
+
+        # Extract audio paths
+        audio_paths = [a["path"] for a in result.audios if a.get("path")]
+
+        # Build result data with download URLs
+        from urllib.parse import urlencode
+        result_data = [{
+            "file": p,
+            "url": f"/v1/audio?{urlencode({'path': p})}",
+            "status": 1,
+            "create_time": int(time.time()),
+        } for p in audio_paths]
+
+        # Store result
+        store_result(task_id, result_data)
+
+        return _wrap_response({"task_id": task_id, "status": "succeeded"})
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+def setup_api_routes_to_app(app, dit_handler, llm_handler, api_key: Optional[str] = None):
+    """
+    Mount API routes to a FastAPI application (for use with gr.mount_gradio_app)
+
+    Args:
+        app: FastAPI application instance
+        dit_handler: DiT handler
+        llm_handler: LLM handler
+        api_key: Optional API key for authentication
+    """
+    set_api_key(api_key)
+    app.state.dit_handler = dit_handler
+    app.state.llm_handler = llm_handler
+    app.include_router(router)
+
+
+def setup_api_routes(demo, dit_handler, llm_handler, api_key: Optional[str] = None):
+    """
+    Mount API routes to Gradio application
+
+    Args:
+        demo: Gradio Blocks instance
+        dit_handler: DiT handler
+        llm_handler: LLM handler
+        api_key: Optional API key for authentication
+    """
+    set_api_key(api_key)
+    app = demo.app
+    app.state.dit_handler = dit_handler
+    app.state.llm_handler = llm_handler
+    app.include_router(router)
+    
+    # Override the /info endpoint to handle schema generation errors gracefully
+    from fastapi.responses import JSONResponse
+    
+    @app.get("/info")
+    async def custom_api_info():
+        """Custom API info endpoint with error handling for schema generation issues"""
+        try:
+            # Try to get the original API info
+            from gradio import utils
+            api_info = utils.safe_deepcopy(demo.get_api_info())
+            return JSONResponse(content=api_info)
+        except (TypeError, AttributeError, KeyError) as e:
+            # If schema generation fails, return a minimal response
+            return JSONResponse(content={
+                "error": "API schema generation not available",
+                "message": "Custom API routes are available at /health, /v1/models, /release_task, /query_result, /create_random_sample, /format_input",
+                "detail": str(e)
+            })

acestep/gradio_ui/events/__init__.py

@@ -0,0 +1,1254 @@
+"""
+Gradio UI Event Handlers Module
+Main entry point for setting up all event handlers
+"""
+import gradio as gr
+from typing import Optional
+from loguru import logger
+
+# Import handler modules
+from . import generation_handlers as gen_h
+from . import results_handlers as res_h
+from . import training_handlers as train_h
+from acestep.gradio_ui.i18n import t
+
+
+def setup_event_handlers(demo, dit_handler, llm_handler, dataset_handler, dataset_section, generation_section, results_section):
+    """Setup event handlers connecting UI components and business logic"""
+    
+    # ========== Dataset Handlers ==========
+    dataset_section["import_dataset_btn"].click(
+        fn=dataset_handler.import_dataset,
+        inputs=[dataset_section["dataset_type"]],
+        outputs=[dataset_section["data_status"]]
+    )
+    
+    # ========== Service Initialization ==========
+    generation_section["refresh_btn"].click(
+        fn=lambda: gen_h.refresh_checkpoints(dit_handler),
+        outputs=[generation_section["checkpoint_dropdown"]]
+    )
+    
+    generation_section["config_path"].change(
+        fn=gen_h.update_model_type_settings,
+        inputs=[generation_section["config_path"]],
+        outputs=[
+            generation_section["inference_steps"],
+            generation_section["guidance_scale"],
+            generation_section["use_adg"],
+            generation_section["shift"],
+            generation_section["cfg_interval_start"],
+            generation_section["cfg_interval_end"],
+            generation_section["task_type"],
+        ]
+    )
+    
+    generation_section["init_btn"].click(
+        fn=lambda *args: gen_h.init_service_wrapper(dit_handler, llm_handler, *args),
+        inputs=[
+            generation_section["checkpoint_dropdown"],
+            generation_section["config_path"],
+            generation_section["device"],
+            generation_section["init_llm_checkbox"],
+            generation_section["lm_model_path"],
+            generation_section["backend_dropdown"],
+            generation_section["use_flash_attention_checkbox"],
+            generation_section["offload_to_cpu_checkbox"],
+            generation_section["offload_dit_to_cpu_checkbox"],
+            generation_section["compile_model_checkbox"],
+            generation_section["quantization_checkbox"],
+        ],
+        outputs=[
+            generation_section["init_status"], 
+            generation_section["generate_btn"], 
+            generation_section["service_config_accordion"],
+            # Model type settings (updated based on actual loaded model)
+            generation_section["inference_steps"],
+            generation_section["guidance_scale"],
+            generation_section["use_adg"],
+            generation_section["shift"],
+            generation_section["cfg_interval_start"],
+            generation_section["cfg_interval_end"],
+            generation_section["task_type"],
+        ]
+    )
+    
+    # ========== LoRA Handlers ==========
+    generation_section["load_lora_btn"].click(
+        fn=dit_handler.load_lora,
+        inputs=[generation_section["lora_path"]],
+        outputs=[generation_section["lora_status"]]
+    ).then(
+        # Update checkbox to enabled state after loading
+        fn=lambda: gr.update(value=True),
+        outputs=[generation_section["use_lora_checkbox"]]
+    )
+    
+    generation_section["unload_lora_btn"].click(
+        fn=dit_handler.unload_lora,
+        outputs=[generation_section["lora_status"]]
+    ).then(
+        # Update checkbox to disabled state after unloading
+        fn=lambda: gr.update(value=False),
+        outputs=[generation_section["use_lora_checkbox"]]
+    )
+    
+    generation_section["use_lora_checkbox"].change(
+        fn=dit_handler.set_use_lora,
+        inputs=[generation_section["use_lora_checkbox"]],
+        outputs=[generation_section["lora_status"]]
+    )
+    
+    generation_section["lora_scale_slider"].change(
+        fn=dit_handler.set_lora_scale,
+        inputs=[generation_section["lora_scale_slider"]],
+        outputs=[generation_section["lora_status"]]
+    )
+    
+    # ========== UI Visibility Updates ==========
+    generation_section["init_llm_checkbox"].change(
+        fn=gen_h.update_negative_prompt_visibility,
+        inputs=[generation_section["init_llm_checkbox"]],
+        outputs=[generation_section["lm_negative_prompt"]]
+    )
+    
+    generation_section["init_llm_checkbox"].change(
+        fn=gen_h.update_audio_cover_strength_visibility,
+        inputs=[generation_section["task_type"], generation_section["init_llm_checkbox"], generation_section["reference_audio"]],
+        outputs=[generation_section["audio_cover_strength"]]
+    )
+    
+    generation_section["task_type"].change(
+        fn=gen_h.update_audio_cover_strength_visibility,
+        inputs=[generation_section["task_type"], generation_section["init_llm_checkbox"], generation_section["reference_audio"]],
+        outputs=[generation_section["audio_cover_strength"]]
+    )
+    
+    generation_section["reference_audio"].change(
+        fn=gen_h.update_audio_cover_strength_visibility,
+        inputs=[generation_section["task_type"], generation_section["init_llm_checkbox"], generation_section["reference_audio"]],
+        outputs=[generation_section["audio_cover_strength"]]
+    )
+    
+    generation_section["batch_size_input"].change(
+        fn=gen_h.update_audio_components_visibility,
+        inputs=[generation_section["batch_size_input"]],
+        outputs=[
+            results_section["audio_col_1"],
+            results_section["audio_col_2"],
+            results_section["audio_col_3"],
+            results_section["audio_col_4"],
+            results_section["audio_row_5_8"],
+            results_section["audio_col_5"],
+            results_section["audio_col_6"],
+            results_section["audio_col_7"],
+            results_section["audio_col_8"],
+        ]
+    )
+    
+    # ========== Audio Conversion ==========
+    generation_section["convert_src_to_codes_btn"].click(
+        fn=lambda src: gen_h.convert_src_audio_to_codes_wrapper(dit_handler, src),
+        inputs=[generation_section["src_audio"]],
+        outputs=[generation_section["text2music_audio_code_string"]]
+    )
+    
+    # ========== Instruction UI Updates ==========
+    for trigger in [generation_section["task_type"], generation_section["track_name"], generation_section["complete_track_classes"], generation_section["reference_audio"]]:
+        trigger.change(
+            fn=lambda *args: gen_h.update_instruction_ui(dit_handler, *args),
+            inputs=[
+                generation_section["task_type"],
+                generation_section["track_name"],
+                generation_section["complete_track_classes"],
+                generation_section["text2music_audio_code_string"],
+                generation_section["init_llm_checkbox"],
+                generation_section["reference_audio"],
+            ],
+            outputs=[
+                generation_section["instruction_display_gen"],
+                generation_section["track_name"],
+                generation_section["complete_track_classes"],
+                generation_section["audio_cover_strength"],
+                generation_section["repainting_group"],
+                generation_section["text2music_audio_codes_group"],
+            ]
+        )
+    
+    # ========== Sample/Transcribe Handlers ==========
+    # Load random example from ./examples/text2music directory
+    generation_section["sample_btn"].click(
+        fn=lambda task: gen_h.load_random_example(task, llm_handler) + (True,),
+        inputs=[
+            generation_section["task_type"],
+        ],
+        outputs=[
+            generation_section["captions"],
+            generation_section["lyrics"],
+            generation_section["think_checkbox"],
+            generation_section["bpm"],
+            generation_section["audio_duration"],
+            generation_section["key_scale"],
+            generation_section["vocal_language"],
+            generation_section["time_signature"],
+            results_section["is_format_caption_state"]
+        ]
+    )
+    
+    generation_section["text2music_audio_code_string"].change(
+        fn=gen_h.update_transcribe_button_text,
+        inputs=[generation_section["text2music_audio_code_string"]],
+        outputs=[generation_section["transcribe_btn"]]
+    )
+    
+    generation_section["transcribe_btn"].click(
+        fn=lambda codes, debug: gen_h.transcribe_audio_codes(llm_handler, codes, debug),
+        inputs=[
+            generation_section["text2music_audio_code_string"],
+            generation_section["constrained_decoding_debug"]
+        ],
+        outputs=[
+            results_section["status_output"],
+            generation_section["captions"],
+            generation_section["lyrics"],
+            generation_section["bpm"],
+            generation_section["audio_duration"],
+            generation_section["key_scale"],
+            generation_section["vocal_language"],
+            generation_section["time_signature"],
+            results_section["is_format_caption_state"]
+        ]
+    )
+    
+    # ========== Reset Format Caption Flag ==========
+    for trigger in [generation_section["captions"], generation_section["lyrics"], generation_section["bpm"],
+                    generation_section["key_scale"], generation_section["time_signature"],
+                    generation_section["vocal_language"], generation_section["audio_duration"]]:
+        trigger.change(
+            fn=gen_h.reset_format_caption_flag,
+            inputs=[],
+            outputs=[results_section["is_format_caption_state"]]
+        )
+    
+    # ========== Audio Uploads Accordion ==========
+    for trigger in [generation_section["reference_audio"], generation_section["src_audio"]]:
+        trigger.change(
+            fn=gen_h.update_audio_uploads_accordion,
+            inputs=[generation_section["reference_audio"], generation_section["src_audio"]],
+            outputs=[generation_section["audio_uploads_accordion"]]
+        )
+    
+    # ========== Instrumental Checkbox ==========
+    generation_section["instrumental_checkbox"].change(
+        fn=gen_h.handle_instrumental_checkbox,
+        inputs=[generation_section["instrumental_checkbox"], generation_section["lyrics"]],
+        outputs=[generation_section["lyrics"]]
+    )
+    
+    # ========== Format Button ==========
+    # Note: cfg_scale and negative_prompt are not supported in format mode
+    generation_section["format_btn"].click(
+        fn=lambda caption, lyrics, bpm, duration, key_scale, time_sig, temp, top_k, top_p, debug: gen_h.handle_format_sample(
+            llm_handler, caption, lyrics, bpm, duration, key_scale, time_sig, temp, top_k, top_p, debug
+        ),
+        inputs=[
+            generation_section["captions"],
+            generation_section["lyrics"],
+            generation_section["bpm"],
+            generation_section["audio_duration"],
+            generation_section["key_scale"],
+            generation_section["time_signature"],
+            generation_section["lm_temperature"],
+            generation_section["lm_top_k"],
+            generation_section["lm_top_p"],
+            generation_section["constrained_decoding_debug"],
+        ],
+        outputs=[
+            generation_section["captions"],
+            generation_section["lyrics"],
+            generation_section["bpm"],
+            generation_section["audio_duration"],
+            generation_section["key_scale"],
+            generation_section["vocal_language"],
+            generation_section["time_signature"],
+            results_section["is_format_caption_state"],
+            results_section["status_output"],
+        ]
+    )
+    
+    # ========== Simple/Custom Mode Toggle ==========
+    generation_section["generation_mode"].change(
+        fn=gen_h.handle_generation_mode_change,
+        inputs=[generation_section["generation_mode"]],
+        outputs=[
+            generation_section["simple_mode_group"],
+            generation_section["caption_accordion"],
+            generation_section["lyrics_accordion"],
+            generation_section["generate_btn"],
+            generation_section["simple_sample_created"],
+            generation_section["optional_params_accordion"],
+        ]
+    )
+    
+    # ========== Simple Mode Instrumental Checkbox ==========
+    # When instrumental is checked, disable vocal language and set to ["unknown"]
+    generation_section["simple_instrumental_checkbox"].change(
+        fn=gen_h.handle_simple_instrumental_change,
+        inputs=[generation_section["simple_instrumental_checkbox"]],
+        outputs=[generation_section["simple_vocal_language"]]
+    )
+    
+    # ========== Random Description Button ==========
+    generation_section["random_desc_btn"].click(
+        fn=gen_h.load_random_simple_description,
+        inputs=[],
+        outputs=[
+            generation_section["simple_query_input"],
+            generation_section["simple_instrumental_checkbox"],
+            generation_section["simple_vocal_language"],
+        ]
+    )
+    
+    # ========== Create Sample Button (Simple Mode) ==========
+    # Note: cfg_scale and negative_prompt are not supported in create_sample mode
+    generation_section["create_sample_btn"].click(
+        fn=lambda query, instrumental, vocal_lang, temp, top_k, top_p, debug: gen_h.handle_create_sample(
+            llm_handler, query, instrumental, vocal_lang, temp, top_k, top_p, debug
+        ),
+        inputs=[
+            generation_section["simple_query_input"],
+            generation_section["simple_instrumental_checkbox"],
+            generation_section["simple_vocal_language"],
+            generation_section["lm_temperature"],
+            generation_section["lm_top_k"],
+            generation_section["lm_top_p"],
+            generation_section["constrained_decoding_debug"],
+        ],
+        outputs=[
+            generation_section["captions"],
+            generation_section["lyrics"],
+            generation_section["bpm"],
+            generation_section["audio_duration"],
+            generation_section["key_scale"],
+            generation_section["vocal_language"],
+            generation_section["simple_vocal_language"],
+            generation_section["time_signature"],
+            generation_section["instrumental_checkbox"],
+            generation_section["caption_accordion"],
+            generation_section["lyrics_accordion"],
+            generation_section["generate_btn"],
+            generation_section["simple_sample_created"],
+            generation_section["think_checkbox"],
+            results_section["is_format_caption_state"],
+            results_section["status_output"],
+        ]
+    )
+    
+    # ========== Load/Save Metadata ==========
+    generation_section["load_file"].upload(
+        fn=lambda file_obj: gen_h.load_metadata(file_obj, llm_handler),
+        inputs=[generation_section["load_file"]],
+        outputs=[
+            generation_section["task_type"],
+            generation_section["captions"],
+            generation_section["lyrics"],
+            generation_section["vocal_language"],
+            generation_section["bpm"],
+            generation_section["key_scale"],
+            generation_section["time_signature"],
+            generation_section["audio_duration"],
+            generation_section["batch_size_input"],
+            generation_section["inference_steps"],
+            generation_section["guidance_scale"],
+            generation_section["seed"],
+            generation_section["random_seed_checkbox"],
+            generation_section["use_adg"],
+            generation_section["cfg_interval_start"],
+            generation_section["cfg_interval_end"],
+            generation_section["shift"],
+            generation_section["infer_method"],
+            generation_section["custom_timesteps"],
+            generation_section["audio_format"],
+            generation_section["lm_temperature"],
+            generation_section["lm_cfg_scale"],
+            generation_section["lm_top_k"],
+            generation_section["lm_top_p"],
+            generation_section["lm_negative_prompt"],
+            generation_section["use_cot_metas"],  # Added: use_cot_metas
+            generation_section["use_cot_caption"],
+            generation_section["use_cot_language"],
+            generation_section["audio_cover_strength"],
+            generation_section["think_checkbox"],
+            generation_section["text2music_audio_code_string"],
+            generation_section["repainting_start"],
+            generation_section["repainting_end"],
+            generation_section["track_name"],
+            generation_section["complete_track_classes"],
+            generation_section["instrumental_checkbox"],  # Added: instrumental_checkbox
+            results_section["is_format_caption_state"]
+        ]
+    )
+    
+    # Save buttons for all 8 audio outputs
+    download_existing_js = """(current_audio, batch_files) => {
+    // Debug: print what the input actually is
+    console.log("👉 [Debug] Current Audio Input:", current_audio);
+    
+    // 1. Safety check
+    if (!current_audio) {
+        console.warn("⚠️ No audio selected or audio is empty.");
+        return;
+    }
+    if (!batch_files || !Array.isArray(batch_files)) {
+        console.warn("⚠️ Batch file list is empty/not ready.");
+        return;
+    }
+
+    // 2. Smartly extract path string
+    let pathString = "";
+    
+    if (typeof current_audio === "string") {
+        // Case A: direct path string received
+        pathString = current_audio;
+    } else if (typeof current_audio === "object") {
+        // Case B: an object is received, try common properties
+        // Gradio file objects usually have path, url, or name
+        pathString = current_audio.path || current_audio.name || current_audio.url || "";
+    }
+
+    if (!pathString) {
+        console.error("❌ Error: Could not extract a valid path string from input.", current_audio);
+        return;
+    }
+
+    // 3. Extract Key (UUID)
+    // Path could be /tmp/.../uuid.mp3 or url like /file=.../uuid.mp3
+    let filename = pathString.split(/[\\\\/]/).pop(); // get the filename
+    let key = filename.split('.')[0]; // get UUID without extension
+
+    console.log(`🔑 Key extracted: ${key}`);
+
+    // 4. Find matching file(s) in the list
+    let targets = batch_files.filter(f => {
+        // Also extract names from batch_files objects
+        // f usually contains name (backend path) and orig_name (download name)
+        const fPath = f.name || f.path || ""; 
+        return fPath.includes(key);
+    });
+
+    if (targets.length === 0) {
+        console.warn("❌ No matching files found in batch list for key:", key);
+        alert("Batch list does not contain this file yet. Please wait for generation to finish.");
+        return;
+    }
+
+    // 5. Trigger download(s)
+    console.log(`🎯 Found ${targets.length} files to download.`);
+    targets.forEach((f, index) => {
+        setTimeout(() => {
+            const a = document.createElement('a');
+            // Prefer url (frontend-accessible link), otherwise try data
+            a.href = f.url || f.data; 
+            a.download = f.orig_name || "download";
+            a.style.display = 'none';
+            document.body.appendChild(a);
+            a.click();
+            document.body.removeChild(a);
+        }, index * 1000); // 300ms interval to avoid browser blocking
+    });
+}
+"""
+    for btn_idx in range(1, 9):
+        results_section[f"save_btn_{btn_idx}"].click(
+            fn=None,
+            inputs=[
+                results_section[f"generated_audio_{btn_idx}"],
+                results_section["generated_audio_batch"],
+            ],
+        js=download_existing_js  # Run the above JS
+    )
+    # ========== Send to SRC Handlers ==========
+    for btn_idx in range(1, 9):
+        results_section[f"send_to_src_btn_{btn_idx}"].click(
+            fn=res_h.send_audio_to_src_with_metadata,
+            inputs=[
+                results_section[f"generated_audio_{btn_idx}"],
+                results_section["lm_metadata_state"]
+            ],
+            outputs=[
+                generation_section["src_audio"],
+                generation_section["bpm"],
+                generation_section["captions"],
+                generation_section["lyrics"],
+                generation_section["audio_duration"],
+                generation_section["key_scale"],
+                generation_section["vocal_language"],
+                generation_section["time_signature"],
+                results_section["is_format_caption_state"]
+            ]
+        )
+    
+    # ========== Score Calculation Handlers ==========
+    # Use default argument to capture btn_idx value at definition time (Python closure fix)
+    def make_score_handler(idx):
+        return lambda scale, batch_idx, queue: res_h.calculate_score_handler_with_selection(
+            dit_handler, llm_handler, idx, scale, batch_idx, queue
+        )
+    
+    for btn_idx in range(1, 9):
+        results_section[f"score_btn_{btn_idx}"].click(
+            fn=make_score_handler(btn_idx),
+            inputs=[
+                generation_section["score_scale"],
+                results_section["current_batch_index"],
+                results_section["batch_queue"],
+            ],
+            outputs=[
+                results_section[f"score_display_{btn_idx}"],
+                results_section[f"details_accordion_{btn_idx}"],
+                results_section["batch_queue"]
+            ]
+        )
+    
+    # ========== LRC Timestamp Handlers ==========
+    # Use default argument to capture btn_idx value at definition time (Python closure fix)
+    def make_lrc_handler(idx):
+        return lambda batch_idx, queue, vocal_lang, infer_steps: res_h.generate_lrc_handler(
+            dit_handler, idx, batch_idx, queue, vocal_lang, infer_steps
+        )
+    
+    for btn_idx in range(1, 9):
+        results_section[f"lrc_btn_{btn_idx}"].click(
+            fn=make_lrc_handler(btn_idx),
+            inputs=[
+                results_section["current_batch_index"],
+                results_section["batch_queue"],
+                generation_section["vocal_language"],
+                generation_section["inference_steps"],
+            ],
+            outputs=[
+                results_section[f"lrc_display_{btn_idx}"],
+                results_section[f"details_accordion_{btn_idx}"],
+                # NOTE: Removed generated_audio output!
+                # Audio subtitles are now updated via lrc_display.change() event.
+                results_section["batch_queue"]
+            ]
+        )
+    
+    def generation_wrapper(*args):
+        yield from res_h.generate_with_batch_management(dit_handler, llm_handler, *args)
+    # ========== Generation Handler ==========
+    generation_section["generate_btn"].click(
+        fn=res_h.clear_audio_outputs_for_new_generation,
+        outputs=[
+            results_section["generated_audio_1"],
+            results_section["generated_audio_2"],
+            results_section["generated_audio_3"],
+            results_section["generated_audio_4"],
+            results_section["generated_audio_5"],
+            results_section["generated_audio_6"],
+            results_section["generated_audio_7"],
+            results_section["generated_audio_8"],
+            results_section["generated_audio_batch"],
+        ],
+    ).then(
+        fn=generation_wrapper,
+        inputs=[
+            generation_section["captions"],
+            generation_section["lyrics"],
+            generation_section["bpm"],
+            generation_section["key_scale"],
+            generation_section["time_signature"],
+            generation_section["vocal_language"],
+            generation_section["inference_steps"],
+            generation_section["guidance_scale"],
+            generation_section["random_seed_checkbox"],
+            generation_section["seed"],
+            generation_section["reference_audio"],
+            generation_section["audio_duration"],
+            generation_section["batch_size_input"],
+            generation_section["src_audio"],
+            generation_section["text2music_audio_code_string"],
+            generation_section["repainting_start"],
+            generation_section["repainting_end"],
+            generation_section["instruction_display_gen"],
+            generation_section["audio_cover_strength"],
+            generation_section["task_type"],
+            generation_section["use_adg"],
+            generation_section["cfg_interval_start"],
+            generation_section["cfg_interval_end"],
+            generation_section["shift"],
+            generation_section["infer_method"],
+            generation_section["custom_timesteps"],
+            generation_section["audio_format"],
+            generation_section["lm_temperature"],
+            generation_section["think_checkbox"],
+            generation_section["lm_cfg_scale"],
+            generation_section["lm_top_k"],
+            generation_section["lm_top_p"],
+            generation_section["lm_negative_prompt"],
+            generation_section["use_cot_metas"],
+            generation_section["use_cot_caption"],
+            generation_section["use_cot_language"],
+            results_section["is_format_caption_state"],
+            generation_section["constrained_decoding_debug"],
+            generation_section["allow_lm_batch"],
+            generation_section["auto_score"],
+            generation_section["auto_lrc"],
+            generation_section["score_scale"],
+            generation_section["lm_batch_chunk_size"],
+            generation_section["track_name"],
+            generation_section["complete_track_classes"],
+            generation_section["autogen_checkbox"],
+            results_section["current_batch_index"],
+            results_section["total_batches"],
+            results_section["batch_queue"],
+            results_section["generation_params_state"],
+        ],
+        outputs=[
+            results_section["generated_audio_1"],
+            results_section["generated_audio_2"],
+            results_section["generated_audio_3"],
+            results_section["generated_audio_4"],
+            results_section["generated_audio_5"],
+            results_section["generated_audio_6"],
+            results_section["generated_audio_7"],
+            results_section["generated_audio_8"],
+            results_section["generated_audio_batch"],
+            results_section["generation_info"],
+            results_section["status_output"],
+            generation_section["seed"],
+            results_section["score_display_1"],
+            results_section["score_display_2"],
+            results_section["score_display_3"],
+            results_section["score_display_4"],
+            results_section["score_display_5"],
+            results_section["score_display_6"],
+            results_section["score_display_7"],
+            results_section["score_display_8"],
+            results_section["codes_display_1"],
+            results_section["codes_display_2"],
+            results_section["codes_display_3"],
+            results_section["codes_display_4"],
+            results_section["codes_display_5"],
+            results_section["codes_display_6"],
+            results_section["codes_display_7"],
+            results_section["codes_display_8"],
+            results_section["details_accordion_1"],
+            results_section["details_accordion_2"],
+            results_section["details_accordion_3"],
+            results_section["details_accordion_4"],
+            results_section["details_accordion_5"],
+            results_section["details_accordion_6"],
+            results_section["details_accordion_7"],
+            results_section["details_accordion_8"],
+            results_section["lrc_display_1"],
+            results_section["lrc_display_2"],
+            results_section["lrc_display_3"],
+            results_section["lrc_display_4"],
+            results_section["lrc_display_5"],
+            results_section["lrc_display_6"],
+            results_section["lrc_display_7"],
+            results_section["lrc_display_8"],
+            results_section["lm_metadata_state"],
+            results_section["is_format_caption_state"],
+            results_section["current_batch_index"],
+            results_section["total_batches"],
+            results_section["batch_queue"],
+            results_section["generation_params_state"],
+            results_section["batch_indicator"],
+            results_section["prev_batch_btn"],
+            results_section["next_batch_btn"],
+            results_section["next_batch_status"],
+            results_section["restore_params_btn"],
+        ],
+    ).then(
+        fn=lambda *args: res_h.generate_next_batch_background(dit_handler, llm_handler, *args),
+        inputs=[
+            generation_section["autogen_checkbox"],
+            results_section["generation_params_state"],
+            results_section["current_batch_index"],
+            results_section["total_batches"],
+            results_section["batch_queue"],
+            results_section["is_format_caption_state"],
+        ],
+        outputs=[
+            results_section["batch_queue"],
+            results_section["total_batches"],
+            results_section["next_batch_status"],
+            results_section["next_batch_btn"],
+        ]
+    )
+    
+    # ========== Batch Navigation Handlers ==========
+    results_section["prev_batch_btn"].click(
+        fn=res_h.navigate_to_previous_batch,
+        inputs=[
+            results_section["current_batch_index"],
+            results_section["batch_queue"],
+        ],
+        outputs=[
+            results_section["generated_audio_1"],
+            results_section["generated_audio_2"],
+            results_section["generated_audio_3"],
+            results_section["generated_audio_4"],
+            results_section["generated_audio_5"],
+            results_section["generated_audio_6"],
+            results_section["generated_audio_7"],
+            results_section["generated_audio_8"],
+            results_section["generated_audio_batch"],
+            results_section["generation_info"],
+            results_section["current_batch_index"],
+            results_section["batch_indicator"],
+            results_section["prev_batch_btn"],
+            results_section["next_batch_btn"],
+            results_section["status_output"],
+            results_section["score_display_1"],
+            results_section["score_display_2"],
+            results_section["score_display_3"],
+            results_section["score_display_4"],
+            results_section["score_display_5"],
+            results_section["score_display_6"],
+            results_section["score_display_7"],
+            results_section["score_display_8"],
+            results_section["codes_display_1"],
+            results_section["codes_display_2"],
+            results_section["codes_display_3"],
+            results_section["codes_display_4"],
+            results_section["codes_display_5"],
+            results_section["codes_display_6"],
+            results_section["codes_display_7"],
+            results_section["codes_display_8"],
+            results_section["lrc_display_1"],
+            results_section["lrc_display_2"],
+            results_section["lrc_display_3"],
+            results_section["lrc_display_4"],
+            results_section["lrc_display_5"],
+            results_section["lrc_display_6"],
+            results_section["lrc_display_7"],
+            results_section["lrc_display_8"],
+            results_section["details_accordion_1"],
+            results_section["details_accordion_2"],
+            results_section["details_accordion_3"],
+            results_section["details_accordion_4"],
+            results_section["details_accordion_5"],
+            results_section["details_accordion_6"],
+            results_section["details_accordion_7"],
+            results_section["details_accordion_8"],
+            results_section["restore_params_btn"],
+        ]
+    )
+    
+    results_section["next_batch_btn"].click(
+        fn=res_h.capture_current_params,
+        inputs=[
+            generation_section["captions"],
+            generation_section["lyrics"],
+            generation_section["bpm"],
+            generation_section["key_scale"],
+            generation_section["time_signature"],
+            generation_section["vocal_language"],
+            generation_section["inference_steps"],
+            generation_section["guidance_scale"],
+            generation_section["random_seed_checkbox"],
+            generation_section["seed"],
+            generation_section["reference_audio"],
+            generation_section["audio_duration"],
+            generation_section["batch_size_input"],
+            generation_section["src_audio"],
+            generation_section["text2music_audio_code_string"],
+            generation_section["repainting_start"],
+            generation_section["repainting_end"],
+            generation_section["instruction_display_gen"],
+            generation_section["audio_cover_strength"],
+            generation_section["task_type"],
+            generation_section["use_adg"],
+            generation_section["cfg_interval_start"],
+            generation_section["cfg_interval_end"],
+            generation_section["shift"],
+            generation_section["infer_method"],
+            generation_section["custom_timesteps"],
+            generation_section["audio_format"],
+            generation_section["lm_temperature"],
+            generation_section["think_checkbox"],
+            generation_section["lm_cfg_scale"],
+            generation_section["lm_top_k"],
+            generation_section["lm_top_p"],
+            generation_section["lm_negative_prompt"],
+            generation_section["use_cot_metas"],
+            generation_section["use_cot_caption"],
+            generation_section["use_cot_language"],
+            generation_section["constrained_decoding_debug"],
+            generation_section["allow_lm_batch"],
+            generation_section["auto_score"],
+            generation_section["auto_lrc"],
+            generation_section["score_scale"],
+            generation_section["lm_batch_chunk_size"],
+            generation_section["track_name"],
+            generation_section["complete_track_classes"],
+        ],
+        outputs=[results_section["generation_params_state"]]
+    ).then(
+        fn=res_h.navigate_to_next_batch,
+        inputs=[
+            generation_section["autogen_checkbox"],
+            results_section["current_batch_index"],
+            results_section["total_batches"],
+            results_section["batch_queue"],
+        ],
+        outputs=[
+            results_section["generated_audio_1"],
+            results_section["generated_audio_2"],
+            results_section["generated_audio_3"],
+            results_section["generated_audio_4"],
+            results_section["generated_audio_5"],
+            results_section["generated_audio_6"],
+            results_section["generated_audio_7"],
+            results_section["generated_audio_8"],
+            results_section["generated_audio_batch"],
+            results_section["generation_info"],
+            results_section["current_batch_index"],
+            results_section["batch_indicator"],
+            results_section["prev_batch_btn"],
+            results_section["next_batch_btn"],
+            results_section["status_output"],
+            results_section["next_batch_status"],
+            results_section["score_display_1"],
+            results_section["score_display_2"],
+            results_section["score_display_3"],
+            results_section["score_display_4"],
+            results_section["score_display_5"],
+            results_section["score_display_6"],
+            results_section["score_display_7"],
+            results_section["score_display_8"],
+            results_section["codes_display_1"],
+            results_section["codes_display_2"],
+            results_section["codes_display_3"],
+            results_section["codes_display_4"],
+            results_section["codes_display_5"],
+            results_section["codes_display_6"],
+            results_section["codes_display_7"],
+            results_section["codes_display_8"],
+            results_section["lrc_display_1"],
+            results_section["lrc_display_2"],
+            results_section["lrc_display_3"],
+            results_section["lrc_display_4"],
+            results_section["lrc_display_5"],
+            results_section["lrc_display_6"],
+            results_section["lrc_display_7"],
+            results_section["lrc_display_8"],
+            results_section["details_accordion_1"],
+            results_section["details_accordion_2"],
+            results_section["details_accordion_3"],
+            results_section["details_accordion_4"],
+            results_section["details_accordion_5"],
+            results_section["details_accordion_6"],
+            results_section["details_accordion_7"],
+            results_section["details_accordion_8"],
+            results_section["restore_params_btn"],
+        ]
+    ).then(
+        fn=lambda *args: res_h.generate_next_batch_background(dit_handler, llm_handler, *args),
+        inputs=[
+            generation_section["autogen_checkbox"],
+            results_section["generation_params_state"],
+            results_section["current_batch_index"],
+            results_section["total_batches"],
+            results_section["batch_queue"],
+            results_section["is_format_caption_state"],
+        ],
+        outputs=[
+            results_section["batch_queue"],
+            results_section["total_batches"],
+            results_section["next_batch_status"],
+            results_section["next_batch_btn"],
+        ]
+    )
+    
+    # ========== Restore Parameters Handler ==========
+    results_section["restore_params_btn"].click(
+        fn=res_h.restore_batch_parameters,
+        inputs=[
+            results_section["current_batch_index"],
+            results_section["batch_queue"]
+        ],
+        outputs=[
+            generation_section["text2music_audio_code_string"],
+            generation_section["captions"],
+            generation_section["lyrics"],
+            generation_section["bpm"],
+            generation_section["key_scale"],
+            generation_section["time_signature"],
+            generation_section["vocal_language"],
+            generation_section["audio_duration"],
+            generation_section["batch_size_input"],
+            generation_section["inference_steps"],
+            generation_section["lm_temperature"],
+            generation_section["lm_cfg_scale"],
+            generation_section["lm_top_k"],
+            generation_section["lm_top_p"],
+            generation_section["think_checkbox"],
+            generation_section["use_cot_caption"],
+            generation_section["use_cot_language"],
+            generation_section["allow_lm_batch"],
+            generation_section["track_name"],
+            generation_section["complete_track_classes"],
+        ]
+    )
+    
+    # ========== LRC Display Change Handlers ==========
+    # NEW APPROACH: Use lrc_display.change() to update audio subtitles
+    # This decouples audio value updates from subtitle updates, avoiding flickering.
+    #
+    # When lrc_display text changes (from generate, LRC button, or manual edit):
+    # 1. lrc_display.change() is triggered
+    # 2. update_audio_subtitles_from_lrc() parses LRC and updates audio subtitles
+    # 3. Audio value is NEVER updated here - only subtitles
+    for lrc_idx in range(1, 9):
+        results_section[f"lrc_display_{lrc_idx}"].change(
+            fn=res_h.update_audio_subtitles_from_lrc,
+            inputs=[
+                results_section[f"lrc_display_{lrc_idx}"],
+                # audio_duration not needed - parse_lrc_to_subtitles calculates end time from timestamps
+            ],
+            outputs=[
+                results_section[f"generated_audio_{lrc_idx}"],  # Only updates subtitles, not value
+            ]
+        )
+
+
+def setup_training_event_handlers(demo, dit_handler, llm_handler, training_section):
+    """Setup event handlers for the training tab (dataset builder and LoRA training)"""
+    
+    # ========== Load Existing Dataset (Top Section) ==========
+
+    # Load existing dataset JSON at the top of Dataset Builder
+    training_section["load_json_btn"].click(
+        fn=train_h.load_existing_dataset_for_preprocess,
+        inputs=[
+            training_section["load_json_path"],
+            training_section["dataset_builder_state"],
+        ],
+        outputs=[
+            training_section["load_json_status"],
+            training_section["audio_files_table"],
+            training_section["sample_selector"],
+            training_section["dataset_builder_state"],
+            # Also update preview fields with first sample
+            training_section["preview_audio"],
+            training_section["preview_filename"],
+            training_section["edit_caption"],
+            training_section["edit_genre"],
+            training_section["prompt_override"],
+            training_section["edit_lyrics"],
+            training_section["edit_bpm"],
+            training_section["edit_keyscale"],
+            training_section["edit_timesig"],
+            training_section["edit_duration"],
+            training_section["edit_language"],
+            training_section["edit_instrumental"],
+            training_section["raw_lyrics_display"],
+            training_section["has_raw_lyrics_state"],
+            # Update dataset-level settings
+            training_section["dataset_name"],
+            training_section["custom_tag"],
+            training_section["tag_position"],
+            training_section["all_instrumental"],
+            training_section["genre_ratio"],
+        ]
+    ).then(
+        fn=lambda has_raw: gr.update(visible=has_raw),
+        inputs=[training_section["has_raw_lyrics_state"]],
+        outputs=[training_section["raw_lyrics_display"]],
+    )
+    
+    # ========== Dataset Builder Handlers ==========
+    
+    # Scan directory for audio files
+    training_section["scan_btn"].click(
+        fn=lambda dir, name, tag, pos, instr, state: train_h.scan_directory(
+            dir, name, tag, pos, instr, state
+        ),
+        inputs=[
+            training_section["audio_directory"],
+            training_section["dataset_name"],
+            training_section["custom_tag"],
+            training_section["tag_position"],
+            training_section["all_instrumental"],
+            training_section["dataset_builder_state"],
+        ],
+        outputs=[
+            training_section["audio_files_table"],
+            training_section["scan_status"],
+            training_section["sample_selector"],
+            training_section["dataset_builder_state"],
+        ]
+    )
+    
+    # Auto-label all samples
+    training_section["auto_label_btn"].click(
+        fn=lambda state, skip, fmt_lyrics, trans_lyrics, only_unlab: train_h.auto_label_all(
+            dit_handler, llm_handler, state, skip, fmt_lyrics, trans_lyrics, only_unlab
+        ),
+        inputs=[
+            training_section["dataset_builder_state"],
+            training_section["skip_metas"],
+            training_section["format_lyrics"],
+            training_section["transcribe_lyrics"],
+            training_section["only_unlabeled"],
+        ],
+        outputs=[
+            training_section["audio_files_table"],
+            training_section["label_progress"],
+            training_section["dataset_builder_state"],
+        ]
+    ).then(
+        # Refresh preview/edit fields after labeling completes
+        fn=train_h.get_sample_preview,
+        inputs=[
+            training_section["sample_selector"],
+            training_section["dataset_builder_state"],
+        ],
+        outputs=[
+            training_section["preview_audio"],
+            training_section["preview_filename"],
+            training_section["edit_caption"],
+            training_section["edit_genre"],
+            training_section["prompt_override"],
+            training_section["edit_lyrics"],
+            training_section["edit_bpm"],
+            training_section["edit_keyscale"],
+            training_section["edit_timesig"],
+            training_section["edit_duration"],
+            training_section["edit_language"],
+            training_section["edit_instrumental"],
+            training_section["raw_lyrics_display"],
+            training_section["has_raw_lyrics_state"],
+        ]
+    ).then(
+        fn=lambda status: f"{status or '✅ Auto-label complete.'}\n✅ Preview refreshed.",
+        inputs=[training_section["label_progress"]],
+        outputs=[training_section["label_progress"]],
+    ).then(
+        fn=lambda has_raw: gr.update(visible=bool(has_raw)),
+        inputs=[training_section["has_raw_lyrics_state"]],
+        outputs=[training_section["raw_lyrics_display"]],
+    )
+
+    # Mutual exclusion: format_lyrics and transcribe_lyrics cannot both be True
+    training_section["format_lyrics"].change(
+        fn=lambda fmt: gr.update(value=False) if fmt else gr.update(),
+        inputs=[training_section["format_lyrics"]],
+        outputs=[training_section["transcribe_lyrics"]]
+    )
+
+    training_section["transcribe_lyrics"].change(
+        fn=lambda trans: gr.update(value=False) if trans else gr.update(),
+        inputs=[training_section["transcribe_lyrics"]],
+        outputs=[training_section["format_lyrics"]]
+    )
+
+    # Sample selector change - update preview
+    training_section["sample_selector"].change(
+        fn=train_h.get_sample_preview,
+        inputs=[
+            training_section["sample_selector"],
+            training_section["dataset_builder_state"],
+        ],
+        outputs=[
+            training_section["preview_audio"],
+            training_section["preview_filename"],
+            training_section["edit_caption"],
+            training_section["edit_genre"],
+            training_section["prompt_override"],
+            training_section["edit_lyrics"],
+            training_section["edit_bpm"],
+            training_section["edit_keyscale"],
+            training_section["edit_timesig"],
+            training_section["edit_duration"],
+            training_section["edit_language"],
+            training_section["edit_instrumental"],
+            training_section["raw_lyrics_display"],
+            training_section["has_raw_lyrics_state"],
+        ]
+    ).then(
+        # Show/hide raw lyrics panel based on whether raw lyrics exist
+        fn=lambda has_raw: gr.update(visible=has_raw),
+        inputs=[training_section["has_raw_lyrics_state"]],
+        outputs=[training_section["raw_lyrics_display"]],
+    )
+    
+    # Save sample edit
+    training_section["save_edit_btn"].click(
+        fn=train_h.save_sample_edit,
+        inputs=[
+            training_section["sample_selector"],
+            training_section["edit_caption"],
+            training_section["edit_genre"],
+            training_section["prompt_override"],
+            training_section["edit_lyrics"],
+            training_section["edit_bpm"],
+            training_section["edit_keyscale"],
+            training_section["edit_timesig"],
+            training_section["edit_language"],
+            training_section["edit_instrumental"],
+            training_section["dataset_builder_state"],
+        ],
+        outputs=[
+            training_section["audio_files_table"],
+            training_section["edit_status"],
+            training_section["dataset_builder_state"],
+        ]
+    )
+
+    # Update settings when changed (including genre_ratio)
+    for trigger in [training_section["custom_tag"], training_section["tag_position"], training_section["all_instrumental"], training_section["genre_ratio"]]:
+        trigger.change(
+            fn=train_h.update_settings,
+            inputs=[
+                training_section["custom_tag"],
+                training_section["tag_position"],
+                training_section["all_instrumental"],
+                training_section["genre_ratio"],
+                training_section["dataset_builder_state"],
+            ],
+            outputs=[training_section["dataset_builder_state"]]
+        )
+
+    # Save dataset
+    training_section["save_dataset_btn"].click(
+        fn=train_h.save_dataset,
+        inputs=[
+            training_section["save_path"],
+            training_section["dataset_name"],
+            training_section["dataset_builder_state"],
+        ],
+        outputs=[
+            training_section["save_status"],
+            training_section["save_path"],
+        ]
+    )
+    
+    # ========== Preprocess Handlers ==========
+    
+    # Load existing dataset JSON for preprocessing
+    # This also updates the preview section so users can view/edit samples
+    training_section["load_existing_dataset_btn"].click(
+        fn=train_h.load_existing_dataset_for_preprocess,
+        inputs=[
+            training_section["load_existing_dataset_path"],
+            training_section["dataset_builder_state"],
+        ],
+        outputs=[
+            training_section["load_existing_status"],
+            training_section["audio_files_table"],
+            training_section["sample_selector"],
+            training_section["dataset_builder_state"],
+            # Also update preview fields with first sample
+            training_section["preview_audio"],
+            training_section["preview_filename"],
+            training_section["edit_caption"],
+            training_section["edit_genre"],
+            training_section["prompt_override"],
+            training_section["edit_lyrics"],
+            training_section["edit_bpm"],
+            training_section["edit_keyscale"],
+            training_section["edit_timesig"],
+            training_section["edit_duration"],
+            training_section["edit_language"],
+            training_section["edit_instrumental"],
+            training_section["raw_lyrics_display"],
+            training_section["has_raw_lyrics_state"],
+            # Update dataset-level settings
+            training_section["dataset_name"],
+            training_section["custom_tag"],
+            training_section["tag_position"],
+            training_section["all_instrumental"],
+            training_section["genre_ratio"],
+        ]
+    ).then(
+        fn=lambda has_raw: gr.update(visible=has_raw),
+        inputs=[training_section["has_raw_lyrics_state"]],
+        outputs=[training_section["raw_lyrics_display"]],
+    )
+    
+    # Preprocess dataset to tensor files
+    training_section["preprocess_btn"].click(
+        fn=lambda output_dir, state: train_h.preprocess_dataset(
+            output_dir, dit_handler, state
+        ),
+        inputs=[
+            training_section["preprocess_output_dir"],
+            training_section["dataset_builder_state"],
+        ],
+        outputs=[training_section["preprocess_progress"]]
+    )
+    
+    # ========== Training Tab Handlers ==========
+    
+    # Load preprocessed tensor dataset
+    training_section["load_dataset_btn"].click(
+        fn=train_h.load_training_dataset,
+        inputs=[training_section["training_tensor_dir"]],
+        outputs=[training_section["training_dataset_info"]]
+    )
+    
+    # Start training from preprocessed tensors
+    def training_wrapper(tensor_dir, r, a, d, lr, ep, bs, ga, se, sh, sd, od, rc, ts):
+        from loguru import logger
+        if not isinstance(ts, dict):
+            ts = {"is_training": False, "should_stop": False}
+        try:
+            for progress, log_msg, plot, state in train_h.start_training(
+                tensor_dir, dit_handler, r, a, d, lr, ep, bs, ga, se, sh, sd, od, rc, ts
+            ):
+                yield progress, log_msg, plot, state
+        except Exception as e:
+            logger.exception("Training wrapper error")
+            yield f"❌ Error: {str(e)}", str(e), None, ts
+    
+    training_section["start_training_btn"].click(
+        fn=training_wrapper,
+        inputs=[
+            training_section["training_tensor_dir"],
+            training_section["lora_rank"],
+            training_section["lora_alpha"],
+            training_section["lora_dropout"],
+            training_section["learning_rate"],
+            training_section["train_epochs"],
+            training_section["train_batch_size"],
+            training_section["gradient_accumulation"],
+            training_section["save_every_n_epochs"],
+            training_section["training_shift"],
+            training_section["training_seed"],
+            training_section["lora_output_dir"],
+            training_section["resume_checkpoint_dir"],
+            training_section["training_state"],
+        ],
+        outputs=[
+            training_section["training_progress"],
+            training_section["training_log"],
+            training_section["training_loss_plot"],
+            training_section["training_state"],
+        ]
+    )
+    
+    # Stop training
+    training_section["stop_training_btn"].click(
+        fn=train_h.stop_training,
+        inputs=[training_section["training_state"]],
+        outputs=[
+            training_section["training_progress"],
+            training_section["training_state"],
+        ]
+    )
+    
+    # Export LoRA
+    training_section["export_lora_btn"].click(
+        fn=train_h.export_lora,
+        inputs=[
+            training_section["export_path"],
+            training_section["lora_output_dir"],
+        ],
+        outputs=[training_section["export_status"]]
+    )

acestep/gradio_ui/events/generation_handlers.py

@@ -0,0 +1,1050 @@
+"""
+Generation Input Handlers Module
+Contains event handlers and helper functions related to generation inputs
+"""
+import os
+import json
+import random
+import glob
+import gradio as gr
+from typing import Optional, List, Tuple
+from acestep.constants import (
+    TASK_TYPES_TURBO,
+    TASK_TYPES_BASE,
+)
+from acestep.gradio_ui.i18n import t
+from acestep.inference import understand_music, create_sample, format_sample
+from acestep.gpu_config import get_global_gpu_config
+
+
+def clamp_duration_to_gpu_limit(duration_value: Optional[float], llm_handler=None) -> Optional[float]:
+    """
+    Clamp duration value to GPU memory limit.
+    
+    Args:
+        duration_value: Duration in seconds (can be None or -1 for no limit)
+        llm_handler: LLM handler instance (to check if LM is initialized)
+        
+    Returns:
+        Clamped duration value, or original value if within limits
+    """
+    if duration_value is None or duration_value <= 0:
+        return duration_value
+    
+    gpu_config = get_global_gpu_config()
+    lm_initialized = llm_handler.llm_initialized if llm_handler else False
+    max_duration = gpu_config.max_duration_with_lm if lm_initialized else gpu_config.max_duration_without_lm
+    
+    if duration_value > max_duration:
+        return float(max_duration)
+    
+    return duration_value
+
+
+def parse_and_validate_timesteps(
+    timesteps_str: str,
+    inference_steps: int
+) -> Tuple[Optional[List[float]], bool, str]:
+    """
+    Parse timesteps string and validate.
+    
+    Args:
+        timesteps_str: Comma-separated timesteps string (e.g., "0.97,0.76,0.615,0.5,0.395,0.28,0.18,0.085,0")
+        inference_steps: Expected number of inference steps
+        
+    Returns:
+        Tuple of (parsed_timesteps, has_warning, warning_message)
+        - parsed_timesteps: List of float timesteps, or None if invalid/empty
+        - has_warning: Whether a warning was shown
+        - warning_message: Description of the warning
+    """
+    if not timesteps_str or not timesteps_str.strip():
+        return None, False, ""
+    
+    # Parse comma-separated values
+    values = [v.strip() for v in timesteps_str.split(",") if v.strip()]
+    
+    if not values:
+        return None, False, ""
+    
+    # Handle optional trailing 0
+    if values[-1] != "0":
+        values.append("0")
+    
+    try:
+        timesteps = [float(v) for v in values]
+    except ValueError:
+        gr.Warning(t("messages.invalid_timesteps_format"))
+        return None, True, "Invalid format"
+    
+    # Validate range [0, 1]
+    if any(ts < 0 or ts > 1 for ts in timesteps):
+        gr.Warning(t("messages.timesteps_out_of_range"))
+        return None, True, "Out of range"
+    
+    # Check if count matches inference_steps
+    actual_steps = len(timesteps) - 1
+    if actual_steps != inference_steps:
+        gr.Warning(t("messages.timesteps_count_mismatch", actual=actual_steps, expected=inference_steps))
+        return timesteps, True, f"Using {actual_steps} steps from timesteps"
+    
+    return timesteps, False, ""
+
+
+def load_metadata(file_obj, llm_handler=None):
+    """Load generation parameters from a JSON file
+    
+    Args:
+        file_obj: Uploaded file object
+        llm_handler: LLM handler instance (optional, for GPU duration limit check)
+    """
+    if file_obj is None:
+        gr.Warning(t("messages.no_file_selected"))
+        return [None] * 36 + [False]  # Return None for all fields, False for is_format_caption
+    
+    try:
+        # Read the uploaded file
+        if hasattr(file_obj, 'name'):
+            filepath = file_obj.name
+        else:
+            filepath = file_obj
+        
+        with open(filepath, 'r', encoding='utf-8') as f:
+            metadata = json.load(f)
+        
+        # Extract all fields
+        task_type = metadata.get('task_type', 'text2music')
+        captions = metadata.get('caption', '')
+        lyrics = metadata.get('lyrics', '')
+        vocal_language = metadata.get('vocal_language', 'unknown')
+        
+        # Convert bpm
+        bpm_value = metadata.get('bpm')
+        if bpm_value is not None and bpm_value != "N/A":
+            try:
+                bpm = int(bpm_value) if bpm_value else None
+            except:
+                bpm = None
+        else:
+            bpm = None
+        
+        key_scale = metadata.get('keyscale', '')
+        time_signature = metadata.get('timesignature', '')
+        
+        # Convert duration
+        duration_value = metadata.get('duration', -1)
+        if duration_value is not None and duration_value != "N/A":
+            try:
+                audio_duration = float(duration_value)
+                # Clamp duration to GPU memory limit
+                audio_duration = clamp_duration_to_gpu_limit(audio_duration, llm_handler)
+            except:
+                audio_duration = -1
+        else:
+            audio_duration = -1
+        
+        batch_size = metadata.get('batch_size', 2)
+        inference_steps = metadata.get('inference_steps', 8)
+        guidance_scale = metadata.get('guidance_scale', 7.0)
+        seed = metadata.get('seed', '-1')
+        random_seed = False  # Always set to False when loading to enable reproducibility with saved seed
+        use_adg = metadata.get('use_adg', False)
+        cfg_interval_start = metadata.get('cfg_interval_start', 0.0)
+        cfg_interval_end = metadata.get('cfg_interval_end', 1.0)
+        audio_format = metadata.get('audio_format', 'mp3')
+        lm_temperature = metadata.get('lm_temperature', 0.85)
+        lm_cfg_scale = metadata.get('lm_cfg_scale', 2.0)
+        lm_top_k = metadata.get('lm_top_k', 0)
+        lm_top_p = metadata.get('lm_top_p', 0.9)
+        lm_negative_prompt = metadata.get('lm_negative_prompt', 'NO USER INPUT')
+        use_cot_metas = metadata.get('use_cot_metas', True)  # Added: read use_cot_metas
+        use_cot_caption = metadata.get('use_cot_caption', True)
+        use_cot_language = metadata.get('use_cot_language', True)
+        audio_cover_strength = metadata.get('audio_cover_strength', 1.0)
+        think = metadata.get('thinking', True)  # Fixed: read 'thinking' not 'think'
+        audio_codes = metadata.get('audio_codes', '')
+        repainting_start = metadata.get('repainting_start', 0.0)
+        repainting_end = metadata.get('repainting_end', -1)
+        track_name = metadata.get('track_name')
+        complete_track_classes = metadata.get('complete_track_classes', [])
+        shift = metadata.get('shift', 3.0)  # Default 3.0 for base models
+        infer_method = metadata.get('infer_method', 'ode')  # Default 'ode' for diffusion inference
+        custom_timesteps = metadata.get('timesteps', '')  # Custom timesteps (stored as 'timesteps' in JSON)
+        if custom_timesteps is None:
+            custom_timesteps = ''
+        instrumental = metadata.get('instrumental', False)  # Added: read instrumental
+        
+        gr.Info(t("messages.params_loaded", filename=os.path.basename(filepath)))
+        
+        return (
+            task_type, captions, lyrics, vocal_language, bpm, key_scale, time_signature,
+            audio_duration, batch_size, inference_steps, guidance_scale, seed, random_seed,
+            use_adg, cfg_interval_start, cfg_interval_end, shift, infer_method,
+            custom_timesteps,  # Added: custom_timesteps (between infer_method and audio_format)
+            audio_format, lm_temperature, lm_cfg_scale, lm_top_k, lm_top_p, lm_negative_prompt,
+            use_cot_metas, use_cot_caption, use_cot_language, audio_cover_strength,
+            think, audio_codes, repainting_start, repainting_end,
+            track_name, complete_track_classes, instrumental,
+            True  # Set is_format_caption to True when loading from file
+        )
+        
+    except json.JSONDecodeError as e:
+        gr.Warning(t("messages.invalid_json", error=str(e)))
+        return [None] * 36 + [False]
+    except Exception as e:
+        gr.Warning(t("messages.load_error", error=str(e)))
+        return [None] * 36 + [False]
+
+
+def load_random_example(task_type: str, llm_handler=None):
+    """Load a random example from the task-specific examples directory
+    
+    Args:
+        task_type: The task type (e.g., "text2music")
+        llm_handler: LLM handler instance (optional, for GPU duration limit check)
+        
+    Returns:
+        Tuple of (caption, lyrics, think, bpm, duration, keyscale, language, timesignature) for updating UI components
+    """
+    try:
+        # Get the project root directory
+        current_file = os.path.abspath(__file__)
+        # This file is in acestep/gradio_ui/events/, need 4 levels up to reach project root
+        project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(current_file))))
+        
+        # Construct the examples directory path
+        examples_dir = os.path.join(project_root, "examples", task_type)
+        
+        # Check if directory exists
+        if not os.path.exists(examples_dir):
+            gr.Warning(f"Examples directory not found: examples/{task_type}/")
+            return "", "", True, None, None, "", "", ""
+        
+        # Find all JSON files in the directory
+        json_files = glob.glob(os.path.join(examples_dir, "*.json"))
+        
+        if not json_files:
+            gr.Warning(f"No JSON files found in examples/{task_type}/")
+            return "", "", True, None, None, "", "", ""
+        
+        # Randomly select one file
+        selected_file = random.choice(json_files)
+        
+        # Read and parse JSON
+        try:
+            with open(selected_file, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+            
+            # Extract caption (prefer 'caption', fallback to 'prompt')
+            caption_value = data.get('caption', data.get('prompt', ''))
+            if not isinstance(caption_value, str):
+                caption_value = str(caption_value) if caption_value else ''
+            
+            # Extract lyrics
+            lyrics_value = data.get('lyrics', '')
+            if not isinstance(lyrics_value, str):
+                lyrics_value = str(lyrics_value) if lyrics_value else ''
+            
+            # Extract think (default to True if not present)
+            think_value = data.get('think', True)
+            if not isinstance(think_value, bool):
+                think_value = True
+            
+            # Extract optional metadata fields
+            bpm_value = None
+            if 'bpm' in data and data['bpm'] not in [None, "N/A", ""]:
+                try:
+                    bpm_value = int(data['bpm'])
+                except (ValueError, TypeError):
+                    pass
+            
+            duration_value = None
+            if 'duration' in data and data['duration'] not in [None, "N/A", ""]:
+                try:
+                    duration_value = float(data['duration'])
+                    # Clamp duration to GPU memory limit
+                    duration_value = clamp_duration_to_gpu_limit(duration_value, llm_handler)
+                except (ValueError, TypeError):
+                    pass
+            
+            keyscale_value = data.get('keyscale', '')
+            if keyscale_value in [None, "N/A"]:
+                keyscale_value = ''
+            
+            language_value = data.get('language', '')
+            if language_value in [None, "N/A"]:
+                language_value = ''
+            
+            timesignature_value = data.get('timesignature', '')
+            if timesignature_value in [None, "N/A"]:
+                timesignature_value = ''
+            
+            gr.Info(t("messages.example_loaded", filename=os.path.basename(selected_file)))
+            return caption_value, lyrics_value, think_value, bpm_value, duration_value, keyscale_value, language_value, timesignature_value
+            
+        except json.JSONDecodeError as e:
+            gr.Warning(t("messages.example_failed", filename=os.path.basename(selected_file), error=str(e)))
+            return "", "", True, None, None, "", "", ""
+        except Exception as e:
+            gr.Warning(t("messages.example_error", error=str(e)))
+            return "", "", True, None, None, "", "", ""
+            
+    except Exception as e:
+        gr.Warning(t("messages.example_error", error=str(e)))
+        return "", "", True, None, None, "", "", ""
+
+
+def sample_example_smart(llm_handler, task_type: str, constrained_decoding_debug: bool = False):
+    """Smart sample function that uses LM if initialized, otherwise falls back to examples
+    
+    This is a Gradio wrapper that uses the understand_music API from acestep.inference
+    to generate examples when LM is available.
+    
+    Args:
+        llm_handler: LLM handler instance
+        task_type: The task type (e.g., "text2music")
+        constrained_decoding_debug: Whether to enable debug logging for constrained decoding
+        
+    Returns:
+        Tuple of (caption, lyrics, think, bpm, duration, keyscale, language, timesignature) for updating UI components
+    """
+    # Check if LM is initialized
+    if llm_handler.llm_initialized:
+        # Use LM to generate example via understand_music API
+        try:
+            result = understand_music(
+                llm_handler=llm_handler,
+                audio_codes="NO USER INPUT",  # Empty input triggers example generation
+                temperature=0.85,
+                use_constrained_decoding=True,
+                constrained_decoding_debug=constrained_decoding_debug,
+            )
+            
+            if result.success:
+                gr.Info(t("messages.lm_generated"))
+                # Clamp duration to GPU memory limit
+                clamped_duration = clamp_duration_to_gpu_limit(result.duration, llm_handler)
+                return (
+                    result.caption,
+                    result.lyrics,
+                    True,  # Always enable think when using LM-generated examples
+                    result.bpm,
+                    clamped_duration,
+                    result.keyscale,
+                    result.language,
+                    result.timesignature,
+                )
+            else:
+                gr.Warning(t("messages.lm_fallback"))
+                return load_random_example(task_type)
+                
+        except Exception as e:
+            gr.Warning(t("messages.lm_fallback"))
+            return load_random_example(task_type)
+    else:
+        # LM not initialized, use examples directory
+        return load_random_example(task_type)
+
+
+def load_random_simple_description():
+    """Load a random description from the simple_mode examples directory.
+
+    Returns:
+        Tuple of (description, instrumental, vocal_language) for updating UI components
+    """
+    try:
+        # Get the project root directory
+        current_file = os.path.abspath(__file__)
+        # This file is in acestep/gradio_ui/events/, need 4 levels up to reach project root
+        project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(current_file))))
+
+        # Construct the examples directory path
+        examples_dir = os.path.join(project_root, "examples", "simple_mode")
+
+        # Check if directory exists
+        if not os.path.exists(examples_dir):
+            gr.Warning(t("messages.simple_examples_not_found"))
+            return gr.update(), gr.update(), gr.update()
+
+        # Find all JSON files in the directory
+        json_files = glob.glob(os.path.join(examples_dir, "*.json"))
+
+        if not json_files:
+            gr.Warning(t("messages.simple_examples_empty"))
+            return gr.update(), gr.update(), gr.update()
+
+        # Randomly select one file
+        selected_file = random.choice(json_files)
+
+        # Read and parse JSON
+        try:
+            with open(selected_file, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+
+            # Extract fields
+            description = data.get('description', '')
+            instrumental = data.get('instrumental', False)
+            vocal_language = data.get('vocal_language', 'unknown')
+
+            # Ensure vocal_language is a string
+            if isinstance(vocal_language, list):
+                vocal_language = vocal_language[0] if vocal_language else 'unknown'
+
+            gr.Info(t("messages.simple_example_loaded", filename=os.path.basename(selected_file)))
+            return description, instrumental, vocal_language
+            
+        except json.JSONDecodeError as e:
+            gr.Warning(t("messages.example_failed", filename=os.path.basename(selected_file), error=str(e)))
+            return gr.update(), gr.update(), gr.update()
+        except Exception as e:
+            gr.Warning(t("messages.example_error", error=str(e)))
+            return gr.update(), gr.update(), gr.update()
+            
+    except Exception as e:
+        gr.Warning(t("messages.example_error", error=str(e)))
+        return gr.update(), gr.update(), gr.update()
+
+
+def refresh_checkpoints(dit_handler):
+    """Refresh available checkpoints"""
+    choices = dit_handler.get_available_checkpoints()
+    return gr.update(choices=choices)
+
+
+def update_model_type_settings(config_path):
+    """Update UI settings based on model type (fallback when handler not initialized yet)
+    
+    Note: This is used as a fallback when the user changes config_path dropdown 
+    before initializing the model. The actual settings are determined by the 
+    handler's is_turbo_model() method after initialization.
+    """
+    if config_path is None:
+        config_path = ""
+    config_path_lower = config_path.lower()
+    
+    # Determine is_turbo based on config_path string
+    # This is a heuristic fallback - actual model type is determined after loading
+    if "turbo" in config_path_lower:
+        is_turbo = True
+    elif "base" in config_path_lower:
+        is_turbo = False
+    else:
+        # Default to turbo settings for unknown model types
+        is_turbo = True
+    
+    return get_model_type_ui_settings(is_turbo)
+
+
+def init_service_wrapper(dit_handler, llm_handler, checkpoint, config_path, device, init_llm, lm_model_path, backend, use_flash_attention, offload_to_cpu, offload_dit_to_cpu, compile_model, quantization):
+    """Wrapper for service initialization, returns status, button state, accordion state, and model type settings"""
+    # Convert quantization checkbox to value (int8_weight_only if checked, None if not)
+    quant_value = "int8_weight_only" if quantization else None
+    
+    # Initialize DiT handler
+    status, enable = dit_handler.initialize_service(
+        checkpoint, config_path, device,
+        use_flash_attention=use_flash_attention, compile_model=compile_model, 
+        offload_to_cpu=offload_to_cpu, offload_dit_to_cpu=offload_dit_to_cpu,
+        quantization=quant_value
+    )
+    
+    # Initialize LM handler if requested
+    if init_llm:
+        # Get checkpoint directory
+        current_file = os.path.abspath(__file__)
+        # This file is in acestep/gradio_ui/events/, need 4 levels up to reach project root
+        project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(current_file))))
+        checkpoint_dir = os.path.join(project_root, "checkpoints")
+        
+        lm_status, lm_success = llm_handler.initialize(
+            checkpoint_dir=checkpoint_dir,
+            lm_model_path=lm_model_path,
+            backend=backend,
+            device=device,
+            offload_to_cpu=offload_to_cpu,
+            dtype=None,
+        )
+        
+        if lm_success:
+            status += f"\n{lm_status}"
+        else:
+            status += f"\n{lm_status}"
+            # Don't fail the entire initialization if LM fails, but log it
+            # Keep enable as is (DiT initialization result) even if LM fails
+    
+    # Check if model is initialized - if so, collapse the accordion
+    is_model_initialized = dit_handler.model is not None
+    accordion_state = gr.Accordion(open=not is_model_initialized)
+    
+    # Get model type settings based on actual loaded model
+    is_turbo = dit_handler.is_turbo_model()
+    model_type_settings = get_model_type_ui_settings(is_turbo)
+    
+    return (
+        status, 
+        gr.update(interactive=enable), 
+        accordion_state,
+        *model_type_settings
+    )
+
+
+def get_model_type_ui_settings(is_turbo: bool):
+    """Get UI settings based on whether the model is turbo or base"""
+    if is_turbo:
+        # Turbo model: max 20 steps, default 8, show shift with default 3.0, only show text2music/repaint/cover
+        return (
+            gr.update(value=8, maximum=20, minimum=1),  # inference_steps
+            gr.update(visible=False),  # guidance_scale
+            gr.update(visible=False),  # use_adg
+            gr.update(value=3.0, visible=True),  # shift (show with default 3.0)
+            gr.update(visible=False),  # cfg_interval_start
+            gr.update(visible=False),  # cfg_interval_end
+            gr.update(choices=TASK_TYPES_TURBO),  # task_type
+        )
+    else:
+        # Base model: max 200 steps, default 32, show CFG/ADG/shift, show all task types
+        return (
+            gr.update(value=32, maximum=200, minimum=1),  # inference_steps
+            gr.update(visible=True),  # guidance_scale
+            gr.update(visible=True),  # use_adg
+            gr.update(value=3.0, visible=True),  # shift (effective for base, default 3.0)
+            gr.update(visible=True),  # cfg_interval_start
+            gr.update(visible=True),  # cfg_interval_end
+            gr.update(choices=TASK_TYPES_BASE),  # task_type
+        )
+
+
+def update_negative_prompt_visibility(init_llm_checked):
+    """Update negative prompt visibility: show if Initialize 5Hz LM checkbox is checked"""
+    return gr.update(visible=init_llm_checked)
+
+
+def _has_reference_audio(reference_audio) -> bool:
+    """True if reference_audio has a usable value (Gradio Audio returns path string or (path, sr))."""
+    if reference_audio is None:
+        return False
+    if isinstance(reference_audio, str):
+        return bool(reference_audio.strip())
+    if isinstance(reference_audio, (list, tuple)) and reference_audio:
+        return bool(reference_audio[0])
+    return False
+
+
+def update_audio_cover_strength_visibility(task_type_value, init_llm_checked, reference_audio=None):
+    """Update audio_cover_strength visibility and label. Show Similarity/Denoise when reference audio is present."""
+    has_reference = _has_reference_audio(reference_audio)
+    # Show if task is cover, LM is initialized, or reference audio is present (audio-conditioned generation)
+    is_visible = (task_type_value == "cover") or init_llm_checked or has_reference
+    # Label priority: cover -> LM codes -> Similarity/Denoise (reference audio)
+    if task_type_value == "cover":
+        label = t("generation.cover_strength_label")
+        info = t("generation.cover_strength_info")
+    elif init_llm_checked:
+        label = t("generation.codes_strength_label")
+        info = t("generation.codes_strength_info")
+    elif has_reference:
+        label = t("generation.similarity_denoise_label")
+        info = t("generation.similarity_denoise_info")
+    else:
+        label = t("generation.cover_strength_label")
+        info = t("generation.cover_strength_info")
+    return gr.update(visible=is_visible, label=label, info=info)
+
+
+def convert_src_audio_to_codes_wrapper(dit_handler, src_audio):
+    """Wrapper for converting src audio to codes"""
+    codes_string = dit_handler.convert_src_audio_to_codes(src_audio)
+    return codes_string
+
+
+def update_instruction_ui(
+    dit_handler,
+    task_type_value: str, 
+    track_name_value: Optional[str], 
+    complete_track_classes_value: list, 
+    audio_codes_content: str = "",
+    init_llm_checked: bool = False,
+    reference_audio=None,
+) -> tuple:
+    """Update instruction and UI visibility based on task type."""
+    instruction = dit_handler.generate_instruction(
+        task_type=task_type_value,
+        track_name=track_name_value,
+        complete_track_classes=complete_track_classes_value
+    )
+    
+    # Show track_name for lego and extract
+    track_name_visible = task_type_value in ["lego", "extract"]
+    # Show complete_track_classes for complete
+    complete_visible = task_type_value == "complete"
+    # Show audio_cover_strength for cover, LM initialized, or reference audio present
+    has_reference = _has_reference_audio(reference_audio)
+    audio_cover_strength_visible = (task_type_value == "cover") or init_llm_checked or has_reference
+    # Label priority: cover -> LM codes -> Similarity/Denoise (reference audio)
+    if task_type_value == "cover":
+        audio_cover_strength_label = t("generation.cover_strength_label")
+        audio_cover_strength_info = t("generation.cover_strength_info")
+    elif init_llm_checked:
+        audio_cover_strength_label = t("generation.codes_strength_label")
+        audio_cover_strength_info = t("generation.codes_strength_info")
+    elif has_reference:
+        audio_cover_strength_label = t("generation.similarity_denoise_label")
+        audio_cover_strength_info = t("generation.similarity_denoise_info")
+    else:
+        audio_cover_strength_label = t("generation.cover_strength_label")
+        audio_cover_strength_info = t("generation.cover_strength_info")
+    # Show repainting controls for repaint and lego
+    repainting_visible = task_type_value in ["repaint", "lego"]
+    # Show text2music_audio_codes if task is text2music OR if it has content
+    # This allows it to stay visible even if user switches task type but has codes
+    has_audio_codes = audio_codes_content and str(audio_codes_content).strip()
+    text2music_audio_codes_visible = task_type_value == "text2music" or has_audio_codes
+    
+    return (
+        instruction,  # instruction_display_gen
+        gr.update(visible=track_name_visible),  # track_name
+        gr.update(visible=complete_visible),  # complete_track_classes
+        gr.update(visible=audio_cover_strength_visible, label=audio_cover_strength_label, info=audio_cover_strength_info),  # audio_cover_strength
+        gr.update(visible=repainting_visible),  # repainting_group
+        gr.update(visible=text2music_audio_codes_visible),  # text2music_audio_codes_group
+    )
+
+
+def transcribe_audio_codes(llm_handler, audio_code_string, constrained_decoding_debug):
+    """
+    Transcribe audio codes to metadata using LLM understanding.
+    If audio_code_string is empty, generate a sample example instead.
+    
+    This is a Gradio wrapper around the understand_music API in acestep.inference.
+    
+    Args:
+        llm_handler: LLM handler instance
+        audio_code_string: String containing audio codes (or empty for example generation)
+        constrained_decoding_debug: Whether to enable debug logging for constrained decoding
+        
+    Returns:
+        Tuple of (status_message, caption, lyrics, bpm, duration, keyscale, language, timesignature, is_format_caption)
+    """
+    # Call the inference API
+    result = understand_music(
+        llm_handler=llm_handler,
+        audio_codes=audio_code_string,
+        use_constrained_decoding=True,
+        constrained_decoding_debug=constrained_decoding_debug,
+    )
+    
+    # Handle error case with localized message
+    if not result.success:
+        # Use localized error message for LLM not initialized
+        if result.error == "LLM not initialized":
+            return t("messages.lm_not_initialized"), "", "", None, None, "", "", "", False
+        return result.status_message, "", "", None, None, "", "", "", False
+    
+    # Clamp duration to GPU memory limit
+    clamped_duration = clamp_duration_to_gpu_limit(result.duration, llm_handler)
+    
+    return (
+        result.status_message,
+        result.caption,
+        result.lyrics,
+        result.bpm,
+        clamped_duration,
+        result.keyscale,
+        result.language,
+        result.timesignature,
+        True  # Set is_format_caption to True (from Transcribe/LM understanding)
+    )
+
+
+def update_transcribe_button_text(audio_code_string):
+    """
+    Update the transcribe button text based on input content.
+    If empty: "Generate Example"
+    If has content: "Transcribe"
+    """
+    if not audio_code_string or not audio_code_string.strip():
+        return gr.update(value="Generate Example")
+    else:
+        return gr.update(value="Transcribe")
+
+
+def reset_format_caption_flag():
+    """Reset is_format_caption to False when user manually edits caption/metadata"""
+    return False
+
+
+def update_audio_uploads_accordion(reference_audio, src_audio):
+    """Update Audio Uploads accordion open state based on whether audio files are present"""
+    has_audio = (reference_audio is not None) or (src_audio is not None)
+    return gr.Accordion(open=has_audio)
+
+
+def handle_instrumental_checkbox(instrumental_checked, current_lyrics):
+    """
+    Handle instrumental checkbox changes.
+    When checked: if no lyrics, fill with [Instrumental]
+    When unchecked: if lyrics is [Instrumental], clear it
+    """
+    if instrumental_checked:
+        # If checked and no lyrics, fill with [Instrumental]
+        if not current_lyrics or not current_lyrics.strip():
+            return "[Instrumental]"
+        else:
+            # Has lyrics, don't change
+            return current_lyrics
+    else:
+        # If unchecked and lyrics is exactly [Instrumental], clear it
+        if current_lyrics and current_lyrics.strip() == "[Instrumental]":
+            return ""
+        else:
+            # Has other lyrics, don't change
+            return current_lyrics
+
+
+def handle_simple_instrumental_change(is_instrumental: bool):
+    """
+    Handle simple mode instrumental checkbox changes.
+    When checked: set vocal_language to "unknown" and disable editing.
+    When unchecked: enable vocal_language editing.
+    
+    Args:
+        is_instrumental: Whether instrumental checkbox is checked
+        
+    Returns:
+        gr.update for simple_vocal_language dropdown
+    """
+    if is_instrumental:
+        return gr.update(value="unknown", interactive=False)
+    else:
+        return gr.update(interactive=True)
+
+
+def update_audio_components_visibility(batch_size):
+    """Show/hide individual audio components based on batch size (1-8)
+    
+    Row 1: Components 1-4 (batch_size 1-4)
+    Row 2: Components 5-8 (batch_size 5-8)
+    """
+    # Clamp batch size to 1-8 range for UI
+    batch_size = min(max(int(batch_size), 1), 8)
+    
+    # Row 1 columns (1-4)
+    updates_row1 = (
+        gr.update(visible=True),  # audio_col_1: always visible
+        gr.update(visible=batch_size >= 2),  # audio_col_2
+        gr.update(visible=batch_size >= 3),  # audio_col_3
+        gr.update(visible=batch_size >= 4),  # audio_col_4
+    )
+    
+    # Row 2 container and columns (5-8)
+    show_row_5_8 = batch_size >= 5
+    updates_row2 = (
+        gr.update(visible=show_row_5_8),  # audio_row_5_8 (container)
+        gr.update(visible=batch_size >= 5),  # audio_col_5
+        gr.update(visible=batch_size >= 6),  # audio_col_6
+        gr.update(visible=batch_size >= 7),  # audio_col_7
+        gr.update(visible=batch_size >= 8),  # audio_col_8
+    )
+    
+    return updates_row1 + updates_row2
+
+
+def handle_generation_mode_change(mode: str):
+    """
+    Handle generation mode change between Simple and Custom modes.
+    
+    In Simple mode:
+    - Show simple mode group (query input, instrumental checkbox, create button)
+    - Collapse caption and lyrics accordions
+    - Hide optional parameters accordion
+    - Disable generate button until sample is created
+    
+    In Custom mode:
+    - Hide simple mode group
+    - Expand caption and lyrics accordions
+    - Show optional parameters accordion
+    - Enable generate button
+    
+    Args:
+        mode: "simple" or "custom"
+        
+    Returns:
+        Tuple of updates for:
+        - simple_mode_group (visibility)
+        - caption_accordion (open state)
+        - lyrics_accordion (open state)
+        - generate_btn (interactive state)
+        - simple_sample_created (reset state)
+        - optional_params_accordion (visibility)
+    """
+    is_simple = mode == "simple"
+    
+    return (
+        gr.update(visible=is_simple),  # simple_mode_group
+        gr.Accordion(open=not is_simple),  # caption_accordion - collapsed in simple, open in custom
+        gr.Accordion(open=not is_simple),  # lyrics_accordion - collapsed in simple, open in custom
+        gr.update(interactive=not is_simple),  # generate_btn - disabled in simple until sample created
+        False,  # simple_sample_created - reset to False on mode change
+        gr.Accordion(open=not is_simple),  # optional_params_accordion - hidden in simple mode
+    )
+
+
+def handle_create_sample(
+    llm_handler,
+    query: str,
+    instrumental: bool,
+    vocal_language: str,
+    lm_temperature: float,
+    lm_top_k: int,
+    lm_top_p: float,
+    constrained_decoding_debug: bool = False,
+):
+    """
+    Handle the Create Sample button click in Simple mode.
+    
+    Creates a sample from the user's query using the LLM, then populates
+    the caption, lyrics, and metadata fields.
+    
+    Note: cfg_scale and negative_prompt are not supported in create_sample mode.
+    
+    Args:
+        llm_handler: LLM handler instance
+        query: User's natural language music description
+        instrumental: Whether to generate instrumental music
+        vocal_language: Preferred vocal language for constrained decoding
+        lm_temperature: LLM temperature for generation
+        lm_top_k: LLM top-k sampling
+        lm_top_p: LLM top-p sampling
+        constrained_decoding_debug: Whether to enable debug logging
+        
+    Returns:
+        Tuple of updates for:
+        - captions
+        - lyrics
+        - bpm
+        - audio_duration
+        - key_scale
+        - vocal_language
+        - time_signature
+        - instrumental_checkbox
+        - caption_accordion (open)
+        - lyrics_accordion (open)
+        - generate_btn (interactive)
+        - simple_sample_created (True)
+        - think_checkbox (True)
+        - is_format_caption_state (True)
+        - status_output
+    """
+    # Check if LLM is initialized
+    if not llm_handler.llm_initialized:
+        gr.Warning(t("messages.lm_not_initialized"))
+        return (
+            gr.update(),  # captions - no change
+            gr.update(),  # lyrics - no change
+            gr.update(),  # bpm - no change
+            gr.update(),  # audio_duration - no change
+            gr.update(),  # key_scale - no change
+            gr.update(),  # vocal_language - no change
+            gr.update(),  # time_signature - no change
+            gr.update(),  # instrumental_checkbox - no change
+            gr.update(),  # caption_accordion - no change
+            gr.update(),  # lyrics_accordion - no change
+            gr.update(interactive=False),  # generate_btn - keep disabled
+            False,  # simple_sample_created - still False
+            gr.update(),  # think_checkbox - no change
+            gr.update(),  # is_format_caption_state - no change
+            t("messages.lm_not_initialized"),  # status_output
+        )
+    
+    # Convert LM parameters
+    top_k_value = None if not lm_top_k or lm_top_k == 0 else int(lm_top_k)
+    top_p_value = None if not lm_top_p or lm_top_p >= 1.0 else lm_top_p
+    
+    # Call create_sample API
+    # Note: cfg_scale and negative_prompt are not supported in create_sample mode
+    result = create_sample(
+        llm_handler=llm_handler,
+        query=query,
+        instrumental=instrumental,
+        vocal_language=vocal_language,
+        temperature=lm_temperature,
+        top_k=top_k_value,
+        top_p=top_p_value,
+        use_constrained_decoding=True,
+        constrained_decoding_debug=constrained_decoding_debug,
+    )
+    
+    # Handle error
+    if not result.success:
+        gr.Warning(result.status_message or t("messages.sample_creation_failed"))
+        return (
+            gr.update(),  # captions - no change
+            gr.update(),  # lyrics - no change
+            gr.update(),  # bpm - no change
+            gr.update(),  # audio_duration - no change
+            gr.update(),  # key_scale - no change
+            gr.update(),  # vocal_language - no change
+            gr.update(),  # simple vocal_language - no change
+            gr.update(),  # time_signature - no change
+            gr.update(),  # instrumental_checkbox - no change
+            gr.update(),  # caption_accordion - no change
+            gr.update(),  # lyrics_accordion - no change
+            gr.update(interactive=False),  # generate_btn - keep disabled
+            False,  # simple_sample_created - still False
+            gr.update(),  # think_checkbox - no change
+            gr.update(),  # is_format_caption_state - no change
+            result.status_message or t("messages.sample_creation_failed"),  # status_output
+        )
+    
+    # Success - populate fields
+    gr.Info(t("messages.sample_created"))
+    
+    # Clamp duration to GPU memory limit
+    clamped_duration = clamp_duration_to_gpu_limit(result.duration, llm_handler)
+    audio_duration_value = clamped_duration if clamped_duration and clamped_duration > 0 else -1
+    
+    return (
+        result.caption,  # captions
+        result.lyrics,  # lyrics
+        result.bpm,  # bpm
+        audio_duration_value,  # audio_duration
+        result.keyscale,  # key_scale
+        result.language,  # vocal_language
+        result.language,  # simple vocal_language
+        result.timesignature,  # time_signature
+        result.instrumental,  # instrumental_checkbox
+        gr.Accordion(open=True),  # caption_accordion - expand
+        gr.Accordion(open=True),  # lyrics_accordion - expand
+        gr.update(interactive=True),  # generate_btn - enable
+        True,  # simple_sample_created - True
+        True,  # think_checkbox - enable thinking
+        True,  # is_format_caption_state - True (LM-generated)
+        result.status_message,  # status_output
+    )
+
+
+def handle_format_sample(
+    llm_handler,
+    caption: str,
+    lyrics: str,
+    bpm,
+    audio_duration,
+    key_scale: str,
+    time_signature: str,
+    lm_temperature: float,
+    lm_top_k: int,
+    lm_top_p: float,
+    constrained_decoding_debug: bool = False,
+):
+    """
+    Handle the Format button click to format caption and lyrics.
+    
+    Takes user-provided caption and lyrics, and uses the LLM to generate
+    structured music metadata and an enhanced description.
+    
+    Note: cfg_scale and negative_prompt are not supported in format mode.
+    
+    Args:
+        llm_handler: LLM handler instance
+        caption: User's caption/description
+        lyrics: User's lyrics
+        bpm: User-provided BPM (optional, for constrained decoding)
+        audio_duration: User-provided duration (optional, for constrained decoding)
+        key_scale: User-provided key scale (optional, for constrained decoding)
+        time_signature: User-provided time signature (optional, for constrained decoding)
+        lm_temperature: LLM temperature for generation
+        lm_top_k: LLM top-k sampling
+        lm_top_p: LLM top-p sampling
+        constrained_decoding_debug: Whether to enable debug logging
+        
+    Returns:
+        Tuple of updates for:
+        - captions
+        - lyrics
+        - bpm
+        - audio_duration
+        - key_scale
+        - vocal_language
+        - time_signature
+        - is_format_caption_state
+        - status_output
+    """
+    # Check if LLM is initialized
+    if not llm_handler.llm_initialized:
+        gr.Warning(t("messages.lm_not_initialized"))
+        return (
+            gr.update(),  # captions - no change
+            gr.update(),  # lyrics - no change
+            gr.update(),  # bpm - no change
+            gr.update(),  # audio_duration - no change
+            gr.update(),  # key_scale - no change
+            gr.update(),  # vocal_language - no change
+            gr.update(),  # time_signature - no change
+            gr.update(),  # is_format_caption_state - no change
+            t("messages.lm_not_initialized"),  # status_output
+        )
+    
+    # Build user_metadata from provided values for constrained decoding
+    user_metadata = {}
+    if bpm is not None and bpm > 0:
+        user_metadata['bpm'] = int(bpm)
+    if audio_duration is not None and float(audio_duration) > 0:
+        user_metadata['duration'] = int(audio_duration)
+    if key_scale and key_scale.strip():
+        user_metadata['keyscale'] = key_scale.strip()
+    if time_signature and time_signature.strip():
+        user_metadata['timesignature'] = time_signature.strip()
+    
+    # Only pass user_metadata if we have at least one field
+    user_metadata_to_pass = user_metadata if user_metadata else None
+    
+    # Convert LM parameters
+    top_k_value = None if not lm_top_k or lm_top_k == 0 else int(lm_top_k)
+    top_p_value = None if not lm_top_p or lm_top_p >= 1.0 else lm_top_p
+    
+    # Call format_sample API
+    result = format_sample(
+        llm_handler=llm_handler,
+        caption=caption,
+        lyrics=lyrics,
+        user_metadata=user_metadata_to_pass,
+        temperature=lm_temperature,
+        top_k=top_k_value,
+        top_p=top_p_value,
+        use_constrained_decoding=True,
+        constrained_decoding_debug=constrained_decoding_debug,
+    )
+    
+    # Handle error
+    if not result.success:
+        gr.Warning(result.status_message or t("messages.format_failed"))
+        return (
+            gr.update(),  # captions - no change
+            gr.update(),  # lyrics - no change
+            gr.update(),  # bpm - no change
+            gr.update(),  # audio_duration - no change
+            gr.update(),  # key_scale - no change
+            gr.update(),  # vocal_language - no change
+            gr.update(),  # time_signature - no change
+            gr.update(),  # is_format_caption_state - no change
+            result.status_message or t("messages.format_failed"),  # status_output
+        )
+    
+    # Success - populate fields
+    gr.Info(t("messages.format_success"))
+    
+    # Clamp duration to GPU memory limit
+    clamped_duration = clamp_duration_to_gpu_limit(result.duration, llm_handler)
+    audio_duration_value = clamped_duration if clamped_duration and clamped_duration > 0 else -1
+    
+    return (
+        result.caption,  # captions
+        result.lyrics,  # lyrics
+        result.bpm,  # bpm
+        audio_duration_value,  # audio_duration
+        result.keyscale,  # key_scale
+        result.language,  # vocal_language
+        result.timesignature,  # time_signature
+        True,  # is_format_caption_state - True (LM-formatted)
+        result.status_message,  # status_output
+    )

acestep/gradio_ui/events/training_handlers.py

@@ -0,0 +1,829 @@
+"""
+Event Handlers for Training Tab
+
+Contains all event handler functions for the dataset builder and training UI.
+"""
+
+import os
+import json
+from typing import Any, Dict, List, Tuple, Optional
+from loguru import logger
+import gradio as gr
+
+from acestep.training.dataset_builder import DatasetBuilder, AudioSample
+from acestep.debug_utils import debug_log_for, debug_start_for, debug_end_for
+from acestep.gpu_config import get_global_gpu_config
+
+
+def create_dataset_builder() -> DatasetBuilder:
+    """Create a new DatasetBuilder instance."""
+    return DatasetBuilder()
+
+
+def _safe_slider(max_value: int, value: int = 0, visible: Optional[bool] = None) -> gr.Slider:
+    """Create a slider with a non-zero range to avoid Gradio math errors."""
+    max_value = max(1, int(max_value))
+    kwargs = {"maximum": max_value, "value": min(int(value), max_value)}
+    if visible is not None:
+        kwargs["visible"] = visible
+    return gr.Slider(**kwargs)
+
+def scan_directory(
+    audio_dir: str,
+    dataset_name: str,
+    custom_tag: str,
+    tag_position: str,
+    all_instrumental: bool,
+    builder_state: Optional[DatasetBuilder],
+) -> Tuple[Any, str, Any, DatasetBuilder]:
+    """Scan a directory for audio files.
+    
+    Returns:
+        Tuple of (table_data, status, slider_update, builder_state)
+    """
+    if not audio_dir or not audio_dir.strip():
+        return [], "� Please enter a directory path", _safe_slider(0, value=0, visible=False), builder_state
+    
+    # Create or use existing builder
+    builder = builder_state if builder_state else DatasetBuilder()
+    
+    # Set metadata before scanning
+    builder.metadata.name = dataset_name
+    builder.metadata.custom_tag = custom_tag
+    builder.metadata.tag_position = tag_position
+    builder.metadata.all_instrumental = all_instrumental
+    
+    # Scan directory
+    samples, status = builder.scan_directory(audio_dir.strip())
+    
+    if not samples:
+        return [], status, _safe_slider(0, value=0, visible=False), builder
+    
+    # Set instrumental and tag for all samples
+    builder.set_all_instrumental(all_instrumental)
+    if custom_tag:
+        builder.set_custom_tag(custom_tag, tag_position)
+    
+    # Get table data
+    table_data = builder.get_samples_dataframe_data()
+    
+    # Calculate slider max and return as Slider update
+    slider_max = max(0, len(samples) - 1)
+    
+    return table_data, status, _safe_slider(slider_max, value=0, visible=len(samples) > 1), builder
+
+
+def auto_label_all(
+    dit_handler,
+    llm_handler,
+    builder_state: Optional[DatasetBuilder],
+    skip_metas: bool = False,
+    format_lyrics: bool = False,
+    transcribe_lyrics: bool = False,
+    only_unlabeled: bool = False,
+    progress=None,
+) -> Tuple[List[List[Any]], str, DatasetBuilder]:
+    """Auto-label all samples in the dataset.
+
+    Args:
+        dit_handler: DiT handler for audio processing
+        llm_handler: LLM handler for caption generation
+        builder_state: Dataset builder state
+        skip_metas: If True, skip generating BPM/Key/TimeSig but still generate caption/genre
+        format_lyrics: If True, use LLM to format user-provided lyrics from .txt files
+        transcribe_lyrics: If True, use LLM to transcribe lyrics from audio (ignores .txt files)
+        only_unlabeled: If True, only label samples without caption
+        progress: Progress callback
+
+    Returns:
+        Tuple of (table_data, status, builder_state)
+    """
+    if builder_state is None:
+        return [], "� Please scan a directory first", builder_state
+
+    if not builder_state.samples:
+        return [], "� No samples to label. Please scan a directory first.", builder_state
+
+    # Check if handlers are initialized
+    if dit_handler is None or dit_handler.model is None:
+        return builder_state.get_samples_dataframe_data(), "� Model not initialized. Please initialize the service first.", builder_state
+
+    if llm_handler is None or not llm_handler.llm_initialized:
+        return builder_state.get_samples_dataframe_data(), "� LLM not initialized. Please initialize the service with LLM enabled.", builder_state
+
+    def progress_callback(msg):
+        if progress:
+            try:
+                progress(msg)
+            except:
+                pass
+
+    # Label all samples (skip_metas only skips BPM/Key/TimeSig, still generates caption/genre)
+    samples, status = builder_state.label_all_samples(
+        dit_handler=dit_handler,
+        llm_handler=llm_handler,
+        format_lyrics=format_lyrics,
+        transcribe_lyrics=transcribe_lyrics,
+        skip_metas=skip_metas,
+        only_unlabeled=only_unlabeled,
+        progress_callback=progress_callback,
+    )
+
+    # Get updated table data
+    table_data = builder_state.get_samples_dataframe_data()
+
+    # Force UI refresh for table and status
+    return gr.update(value=table_data), gr.update(value=status), builder_state
+
+
+def get_sample_preview(
+    sample_idx: int,
+    builder_state: Optional[DatasetBuilder],
+):
+    """Get preview data for a specific sample.
+
+    Returns:
+        Tuple of (audio_path, filename, caption, genre, prompt_override, lyrics, bpm, keyscale, timesig,
+                  duration, language, instrumental, raw_lyrics, raw_lyrics_visible)
+    """
+    empty = (None, "", "", "", "Use Global Ratio", "", None, "", "", 0.0, "instrumental", True, "", False)
+
+    if builder_state is None or not builder_state.samples:
+        return empty
+
+    if sample_idx is None:
+        return empty
+
+    idx = int(sample_idx)
+    if idx < 0 or idx >= len(builder_state.samples):
+        return empty
+
+    sample = builder_state.samples[idx]
+
+    # Show raw lyrics panel only when raw lyrics exist
+    has_raw = sample.has_raw_lyrics()
+
+    # Convert prompt_override to dropdown choice
+    if sample.prompt_override == "genre":
+        override_choice = "Genre"
+    elif sample.prompt_override == "caption":
+        override_choice = "Caption"
+    else:
+        override_choice = "Use Global Ratio"
+
+    display_lyrics = sample.lyrics if sample.lyrics else sample.formatted_lyrics
+
+    return (
+        sample.audio_path,
+        sample.filename,
+        sample.caption,
+        sample.genre,
+        override_choice,
+        display_lyrics,
+        sample.bpm,
+        sample.keyscale,
+        sample.timesignature,
+        sample.duration,
+        sample.language,
+        sample.is_instrumental,
+        sample.raw_lyrics if has_raw else "",
+        has_raw,
+    )
+
+
+def save_sample_edit(
+    sample_idx: int,
+    caption: str,
+    genre: str,
+    prompt_override: str,
+    lyrics: str,
+    bpm: Optional[int],
+    keyscale: str,
+    timesig: str,
+    language: str,
+    is_instrumental: bool,
+    builder_state: Optional[DatasetBuilder],
+) -> Tuple[List[List[Any]], str, DatasetBuilder]:
+    """Save edits to a sample.
+
+    Returns:
+        Tuple of (table_data, status, builder_state)
+    """
+    if builder_state is None:
+        return [], "� No dataset loaded", builder_state
+
+    idx = int(sample_idx)
+
+    # Convert dropdown choice to prompt_override value
+    if prompt_override == "Genre":
+        override_value = "genre"
+    elif prompt_override == "Caption":
+        override_value = "caption"
+    else:
+        override_value = None  # Use Global Ratio
+
+    # Update sample
+    updated_lyrics = lyrics if not is_instrumental else "[Instrumental]"
+    updated_formatted = updated_lyrics if updated_lyrics and updated_lyrics != "[Instrumental]" else ""
+    sample, status = builder_state.update_sample(
+        idx,
+        caption=caption,
+        genre=genre,
+        prompt_override=override_value,
+        lyrics=updated_lyrics,
+        formatted_lyrics=updated_formatted,
+        bpm=int(bpm) if bpm else None,
+        keyscale=keyscale,
+        timesignature=timesig,
+        language="unknown" if is_instrumental else language,
+        is_instrumental=is_instrumental,
+        labeled=True,
+    )
+
+    # Get updated table data
+    table_data = builder_state.get_samples_dataframe_data()
+
+    return table_data, status, builder_state
+
+
+def update_settings(
+    custom_tag: str,
+    tag_position: str,
+    all_instrumental: bool,
+    genre_ratio: int,
+    builder_state: Optional[DatasetBuilder],
+) -> DatasetBuilder:
+    """Update dataset settings.
+
+    Returns:
+        Updated builder_state
+    """
+    if builder_state is None:
+        return builder_state
+
+    if custom_tag:
+        builder_state.set_custom_tag(custom_tag, tag_position)
+
+    builder_state.set_all_instrumental(all_instrumental)
+    builder_state.metadata.genre_ratio = int(genre_ratio)
+
+    return builder_state
+
+
+def save_dataset(
+    save_path: str,
+    dataset_name: str,
+    builder_state: Optional[DatasetBuilder],
+) -> Tuple[str, Any]:
+    """Save the dataset to a JSON file.
+    
+    Returns:
+        Status message
+    """
+    if builder_state is None:
+        return "� No dataset to save. Please scan a directory first.", gr.update()
+    
+    if not builder_state.samples:
+        return "� No samples in dataset.", gr.update()
+    
+    if not save_path or not save_path.strip():
+        return "� Please enter a save path.", gr.update()
+    
+    save_path = save_path.strip()
+    if not save_path.lower().endswith(".json"):
+        save_path = save_path + ".json"
+    
+    # Check if any samples are labeled
+    labeled_count = builder_state.get_labeled_count()
+    if labeled_count == 0:
+        return "�️ Warning: No samples have been labeled. Consider auto-labeling first.\nSaving anyway...", gr.update(value=save_path)
+
+    return builder_state.save_dataset(save_path, dataset_name), gr.update(value=save_path)
+
+
+def load_existing_dataset_for_preprocess(
+    dataset_path: str,
+    builder_state: Optional[DatasetBuilder],
+):
+    """Load an existing dataset JSON file for preprocessing.
+
+    This allows users to load a previously saved dataset and proceed to preprocessing
+    without having to re-scan and re-label.
+
+    Returns:
+        Tuple of (status, table_data, slider_update, builder_state,
+                  audio_path, filename, caption, genre, prompt_override,
+                  lyrics, bpm, keyscale, timesig, duration, language, instrumental,
+                  raw_lyrics, has_raw)
+    """
+    # Empty preview: (audio_path, filename, caption, genre, prompt_override, lyrics, bpm, keyscale, timesig, duration, language, instrumental, raw_lyrics, has_raw)
+    empty_preview = (None, "", "", "", "Use Global Ratio", "", None, "", "", 0.0, "instrumental", True, "", False)
+
+    if not dataset_path or not dataset_path.strip():
+        updates = (gr.update(), gr.update(), gr.update(), gr.update(), gr.update())
+        return ("� Please enter a dataset path", [], _safe_slider(0, value=0, visible=False), builder_state) + empty_preview + updates
+
+    dataset_path = dataset_path.strip()
+    debug_log_for("dataset", f"UI load_existing_dataset_for_preprocess: path='{dataset_path}'")
+
+    if not os.path.exists(dataset_path):
+        updates = (gr.update(), gr.update(), gr.update(), gr.update(), gr.update())
+        return (f"� Dataset not found: {dataset_path}", [], _safe_slider(0, value=0, visible=False), builder_state) + empty_preview + updates
+
+    # Create new builder (don't reuse old state when loading a file)
+    builder = DatasetBuilder()
+
+    # Load the dataset
+    t0 = debug_start_for("dataset", "load_dataset")
+    samples, status = builder.load_dataset(dataset_path)
+    debug_end_for("dataset", "load_dataset", t0)
+
+    if not samples:
+        updates = (gr.update(), gr.update(), gr.update(), gr.update(), gr.update())
+        return (status, [], _safe_slider(0, value=0, visible=False), builder) + empty_preview + updates
+
+    # Get table data
+    table_data = builder.get_samples_dataframe_data()
+
+    # Calculate slider max
+    slider_max = max(0, len(samples) - 1)
+
+    # Create info text
+    labeled_count = builder.get_labeled_count()
+    info = f"� Loaded dataset: {builder.metadata.name}\n"
+    info += f"� Samples: {len(samples)} ({labeled_count} labeled)\n"
+    info += f"���️ Custom Tag: {builder.metadata.custom_tag or '(none)'}\n"
+    info += "� Ready for preprocessing! You can also edit samples below."
+    if any((s.formatted_lyrics and not s.lyrics) for s in builder.samples):
+        info += "\n�️ Showing formatted lyrics where lyrics are empty."
+
+    # Get first sample preview
+    first_sample = builder.samples[0]
+    has_raw = first_sample.has_raw_lyrics()
+
+    # Convert prompt_override to dropdown choice
+    if first_sample.prompt_override == "genre":
+        override_choice = "Genre"
+    elif first_sample.prompt_override == "caption":
+        override_choice = "Caption"
+    else:
+        override_choice = "Use Global Ratio"
+
+    display_lyrics = first_sample.lyrics if first_sample.lyrics else first_sample.formatted_lyrics
+
+    preview = (
+        first_sample.audio_path,
+        first_sample.filename,
+        first_sample.caption,
+        first_sample.genre,
+        override_choice,
+        display_lyrics,
+        first_sample.bpm,
+        first_sample.keyscale,
+        first_sample.timesignature,
+        first_sample.duration,
+        first_sample.language,
+        first_sample.is_instrumental,
+        first_sample.raw_lyrics if has_raw else "",
+        has_raw,
+    )
+
+    updates = (
+        gr.update(value=builder.metadata.name),
+        gr.update(value=builder.metadata.custom_tag),
+        gr.update(value=builder.metadata.tag_position),
+        gr.update(value=builder.metadata.all_instrumental),
+        gr.update(value=builder.metadata.genre_ratio),
+    )
+
+    return (info, table_data, _safe_slider(slider_max, value=0, visible=len(samples) > 1), builder) + preview + updates
+
+
+def preprocess_dataset(
+    output_dir: str,
+    dit_handler,
+    builder_state: Optional[DatasetBuilder],
+    progress=None,
+) -> str:
+    """Preprocess dataset to tensor files for fast training.
+    
+    This converts audio files to VAE latents and text to embeddings.
+    
+    Returns:
+        Status message
+    """
+    if builder_state is None:
+        return "� No dataset loaded. Please scan a directory first."
+    
+    if not builder_state.samples:
+        return "� No samples in dataset."
+    
+    labeled_count = builder_state.get_labeled_count()
+    if labeled_count == 0:
+        return "� No labeled samples. Please auto-label or manually label samples first."
+    
+    if not output_dir or not output_dir.strip():
+        return "� Please enter an output directory."
+    
+    if dit_handler is None or dit_handler.model is None:
+        return "� Model not initialized. Please initialize the service first."
+    
+    def progress_callback(msg):
+        if progress:
+            try:
+                progress(msg)
+            except:
+                pass
+    
+    # Run preprocessing
+    t0 = debug_start_for("dataset", "preprocess_to_tensors")
+    output_paths, status = builder_state.preprocess_to_tensors(
+        dit_handler=dit_handler,
+        output_dir=output_dir.strip(),
+        progress_callback=progress_callback,
+    )
+    debug_end_for("dataset", "preprocess_to_tensors", t0)
+    
+    return status
+
+
+def load_training_dataset(
+    tensor_dir: str,
+) -> str:
+    """Load a preprocessed tensor dataset for training.
+    
+    Returns:
+        Info text about the dataset
+    """
+    if not tensor_dir or not tensor_dir.strip():
+        return "� Please enter a tensor directory path"
+    
+    tensor_dir = tensor_dir.strip()
+    
+    if not os.path.exists(tensor_dir):
+        return f"� Directory not found: {tensor_dir}"
+    
+    if not os.path.isdir(tensor_dir):
+        return f"� Not a directory: {tensor_dir}"
+    
+    # Check for manifest
+    manifest_path = os.path.join(tensor_dir, "manifest.json")
+    if os.path.exists(manifest_path):
+        try:
+            with open(manifest_path, 'r') as f:
+                manifest = json.load(f)
+            
+            num_samples = manifest.get("num_samples", 0)
+            metadata = manifest.get("metadata", {})
+            name = metadata.get("name", "Unknown")
+            custom_tag = metadata.get("custom_tag", "")
+            
+            info = f"� Loaded preprocessed dataset: {name}\n"
+            info += f"� Samples: {num_samples} preprocessed tensors\n"
+            info += f"���️ Custom Tag: {custom_tag or '(none)'}"
+            
+            return info
+        except Exception as e:
+            logger.warning(f"Failed to read manifest: {e}")
+    
+    # Fallback: count .pt files
+    pt_files = [f for f in os.listdir(tensor_dir) if f.endswith('.pt')]
+    
+    if not pt_files:
+        return f"� No .pt tensor files found in {tensor_dir}"
+    
+    info = f"� Found {len(pt_files)} tensor files in {tensor_dir}\n"
+    info += "�️ No manifest.json found - using all .pt files"
+    
+    return info
+
+
+# Training handlers
+
+import time
+import re
+
+
+def _format_duration(seconds):
+    """Format seconds to human readable string."""
+    seconds = int(seconds)
+    if seconds < 60:
+        return f"{seconds}s"
+    elif seconds < 3600:
+        return f"{seconds // 60}m {seconds % 60}s"
+    else:
+        return f"{seconds // 3600}h {(seconds % 3600) // 60}m"
+
+
+def start_training(
+    tensor_dir: str,
+    dit_handler,
+    lora_rank: int,
+    lora_alpha: int,
+    lora_dropout: float,
+    learning_rate: float,
+    train_epochs: int,
+    train_batch_size: int,
+    gradient_accumulation: int,
+    save_every_n_epochs: int,
+    training_shift: float,
+    training_seed: int,
+    lora_output_dir: str,
+    resume_checkpoint_dir: str,
+    training_state: Dict,
+    progress=None,
+):
+    """Start LoRA training from preprocessed tensors.
+    
+    This is a generator function that yields progress updates.
+    """
+    if not tensor_dir or not tensor_dir.strip():
+        yield "� Please enter a tensor directory path", "", None, training_state
+        return
+    
+    tensor_dir = tensor_dir.strip()
+    
+    if not os.path.exists(tensor_dir):
+        yield f"� Tensor directory not found: {tensor_dir}", "", None, training_state
+        return
+    
+    if dit_handler is None or dit_handler.model is None:
+        yield "� Model not initialized. Please initialize the service first.", "", None, training_state
+        return
+    
+    # Training preset: LoRA training must run on non-quantized DiT.
+    if getattr(dit_handler, "quantization", None) is not None:
+        gpu_config = get_global_gpu_config()
+        if gpu_config.gpu_memory_gb <= 0:
+            yield (
+                "WARNING: CPU-only training detected. Using best-effort training path "
+                "(non-quantized DiT). Performance will be sub-optimal.",
+                "",
+                None,
+                training_state,
+            )
+        elif gpu_config.tier in {"tier1", "tier2", "tier3", "tier4"}:
+            yield (
+                f"WARNING: Low VRAM tier detected ({gpu_config.gpu_memory_gb:.1f} GB, {gpu_config.tier}). "
+                "Using best-effort training path (non-quantized DiT). Performance may be sub-optimal.",
+                "",
+                None,
+                training_state,
+            )
+
+        yield "Switching model to training preset (disable quantization)...", "", None, training_state
+        if hasattr(dit_handler, "switch_to_training_preset"):
+            switch_status, switched = dit_handler.switch_to_training_preset()
+            if not switched:
+                yield f"� {switch_status}", "", None, training_state
+                return
+            yield f"� {switch_status}", "", None, training_state
+        else:
+            yield "� Training requires non-quantized DiT, and auto-switch is unavailable in this build.", "", None, training_state
+            return
+
+    # Check for required training dependencies
+    try:
+        from lightning.fabric import Fabric
+        from peft import get_peft_model, LoraConfig
+    except ImportError as e:
+        yield f"� Missing required packages: {e}\nPlease install: pip install peft lightning", "", None, training_state
+        return
+    
+    training_state["is_training"] = True
+    training_state["should_stop"] = False
+    
+    try:
+        from acestep.training.trainer import LoRATrainer
+        from acestep.training.configs import LoRAConfig as LoRAConfigClass, TrainingConfig
+        
+        # Create configs
+        lora_config = LoRAConfigClass(
+            r=lora_rank,
+            alpha=lora_alpha,
+            dropout=lora_dropout,
+        )
+        
+        device_attr = getattr(dit_handler, "device", "")
+        if hasattr(device_attr, "type"):
+            device_type = str(device_attr.type).lower()
+        else:
+            device_type = str(device_attr).split(":", 1)[0].lower()
+
+        # Use device-tuned dataloader defaults while preserving CUDA acceleration.
+        if device_type == "cuda":
+            num_workers = 4
+            pin_memory = True
+            prefetch_factor = 2
+            persistent_workers = True
+            pin_memory_device = "cuda"
+            mixed_precision = "bf16"
+        elif device_type == "xpu":
+            num_workers = 4
+            pin_memory = True
+            prefetch_factor = 2
+            persistent_workers = True
+            pin_memory_device = None
+            mixed_precision = "bf16"
+        elif device_type == "mps":
+            num_workers = 0
+            pin_memory = False
+            prefetch_factor = 2
+            persistent_workers = False
+            pin_memory_device = None
+            mixed_precision = "fp16"
+        else:
+            cpu_count = os.cpu_count() or 2
+            num_workers = min(4, max(1, cpu_count // 2))
+            pin_memory = False
+            prefetch_factor = 2
+            persistent_workers = num_workers > 0
+            pin_memory_device = None
+            mixed_precision = "fp32"
+
+        logger.info(
+            f"Training loader config: device={device_type}, workers={num_workers}, "
+            f"pin_memory={pin_memory}, pin_memory_device={pin_memory_device}, "
+            f"persistent_workers={persistent_workers}"
+        )
+        training_config = TrainingConfig(
+            shift=training_shift,
+            learning_rate=learning_rate,
+            batch_size=train_batch_size,
+            gradient_accumulation_steps=gradient_accumulation,
+            max_epochs=train_epochs,
+            save_every_n_epochs=save_every_n_epochs,
+            seed=training_seed,
+            output_dir=lora_output_dir,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            prefetch_factor=prefetch_factor,
+            persistent_workers=persistent_workers,
+            pin_memory_device=pin_memory_device,
+            mixed_precision=mixed_precision,
+        )
+        
+        import pandas as pd
+        
+        # Initialize training log and loss history
+        log_lines = []
+        loss_data = pd.DataFrame({"step": [0], "loss": [0.0]})
+        
+        # Start timer
+        start_time = time.time()
+        
+        yield f"� Starting training from {tensor_dir}...", "", loss_data, training_state
+        
+        # Create trainer
+        trainer = LoRATrainer(
+            dit_handler=dit_handler,
+            lora_config=lora_config,
+            training_config=training_config,
+        )
+        
+        # Collect loss history
+        step_list = []
+        loss_list = []
+        training_failed = False
+        failure_message = ""
+        
+        # Train with progress updates using preprocessed tensors
+        resume_from = resume_checkpoint_dir.strip() if resume_checkpoint_dir and resume_checkpoint_dir.strip() else None
+        for step, loss, status in trainer.train_from_preprocessed(tensor_dir, training_state, resume_from=resume_from):
+            status_text = str(status)
+            status_lower = status_text.lower()
+            if (
+                status_text.startswith("❌")
+                or status_text.startswith("❌")
+                or "training failed" in status_lower
+                or "error:" in status_lower
+                or "module not found" in status_lower
+            ):
+                training_failed = True
+                failure_message = status_text
+            # Calculate elapsed time and ETA
+            elapsed_seconds = time.time() - start_time
+            time_info = f"⏱️ Elapsed: {_format_duration(elapsed_seconds)}"
+            
+            # Parse "Epoch x/y" from status to calculate ETA
+            match = re.search(r"Epoch\s+(\d+)/(\d+)", str(status))
+            if match:
+                current_ep = int(match.group(1))
+                total_ep = int(match.group(2))
+                if current_ep > 0:
+                    eta_seconds = (elapsed_seconds / current_ep) * (total_ep - current_ep)
+                    time_info += f" | ETA: ~{_format_duration(eta_seconds)}"
+            
+            # Display status with time info
+            display_status = f"{status}\n{time_info}"
+            
+            # Terminal log
+            log_msg = f"[{_format_duration(elapsed_seconds)}] Step {step}: {status}"
+            logger.info(log_msg)
+            
+            # Add to UI log
+            log_lines.append(status)
+            if len(log_lines) > 15:
+                log_lines = log_lines[-15:]
+            log_text = "\n".join(log_lines)
+            
+            # Track loss for plot (only valid values)
+            if step > 0 and loss is not None and loss == loss:  # Check for NaN
+                step_list.append(step)
+                loss_list.append(float(loss))
+                loss_data = pd.DataFrame({"step": step_list, "loss": loss_list})
+            
+            yield display_status, log_text, loss_data, training_state
+            
+            if training_state.get("should_stop", False):
+                logger.info("⏹️ Training stopped by user")
+                log_lines.append("⏹️ Training stopped by user")
+                yield f"⏹️ Stopped ({time_info})", "\n".join(log_lines[-15:]), loss_data, training_state
+                break
+        
+        total_time = time.time() - start_time
+        training_state["is_training"] = False
+        if training_failed:
+            final_msg = f"{failure_message}\nElapsed: {_format_duration(total_time)}"
+            logger.warning(final_msg)
+            log_lines.append(failure_message)
+            yield final_msg, "\n".join(log_lines[-15:]), loss_data, training_state
+            return
+        completion_msg = f"� Training completed! Total time: {_format_duration(total_time)}"
+        
+        logger.info(completion_msg)
+        log_lines.append(completion_msg)
+        
+        yield completion_msg, "\n".join(log_lines[-15:]), loss_data, training_state
+        
+    except Exception as e:
+        logger.exception("Training error")
+        training_state["is_training"] = False
+        import pandas as pd
+        empty_df = pd.DataFrame({"step": [], "loss": []})
+        yield f"� Error: {str(e)}", str(e), empty_df, training_state
+
+
+def stop_training(training_state: Dict) -> Tuple[str, Dict]:
+    """Stop the current training process.
+    
+    Returns:
+        Tuple of (status, training_state)
+    """
+    if not training_state.get("is_training", False):
+        return "�️ No training in progress", training_state
+    
+    training_state["should_stop"] = True
+    return "⏹️ Stopping training...", training_state
+
+
+def export_lora(
+    export_path: str,
+    lora_output_dir: str,
+) -> str:
+    """Export the trained LoRA weights.
+    
+    Returns:
+        Status message
+    """
+    if not export_path or not export_path.strip():
+        return "� Please enter an export path"
+    
+    # Check if there's a trained model to export
+    final_dir = os.path.join(lora_output_dir, "final")
+    checkpoint_dir = os.path.join(lora_output_dir, "checkpoints")
+    
+    # Prefer final, fallback to checkpoints
+    if os.path.exists(final_dir):
+        source_path = final_dir
+    elif os.path.exists(checkpoint_dir):
+        # Find the latest checkpoint
+        checkpoints = [d for d in os.listdir(checkpoint_dir) if d.startswith("epoch_")]
+        if not checkpoints:
+            return "� No checkpoints found"
+        
+        checkpoints.sort(key=lambda x: int(x.split("_")[1]))
+        latest = checkpoints[-1]
+        source_path = os.path.join(checkpoint_dir, latest)
+    else:
+        return f"� No trained model found in {lora_output_dir}"
+    
+    try:
+        import shutil
+        
+        export_path = export_path.strip()
+        os.makedirs(os.path.dirname(export_path) if os.path.dirname(export_path) else ".", exist_ok=True)
+        
+        if os.path.exists(export_path):
+            shutil.rmtree(export_path)
+        
+        shutil.copytree(source_path, export_path)
+        
+        return f"� LoRA exported to {export_path}"
+        
+    except Exception as e:
+        logger.exception("Export error")
+        return f"� Export failed: {str(e)}"
+
+
+

acestep/gradio_ui/i18n.py

@@ -0,0 +1,152 @@
+"""
+Internationalization (i18n) module for Gradio UI
+Supports multiple languages with easy translation management
+"""
+import os
+import json
+from typing import Dict, Optional
+
+
+class I18n:
+    """Internationalization handler"""
+    
+    def __init__(self, default_language: str = "en"):
+        """
+        Initialize i18n handler
+        
+        Args:
+            default_language: Default language code (en, zh, ja, etc.)
+        """
+        self.current_language = default_language
+        self.translations: Dict[str, Dict[str, str]] = {}
+        self._load_all_translations()
+    
+    def _load_all_translations(self):
+        """Load all translation files from i18n directory"""
+        current_file = os.path.abspath(__file__)
+        module_dir = os.path.dirname(current_file)
+        i18n_dir = os.path.join(module_dir, "i18n")
+        
+        if not os.path.exists(i18n_dir):
+            # Create i18n directory if it doesn't exist
+            os.makedirs(i18n_dir)
+            return
+        
+        # Load all JSON files in i18n directory
+        for filename in os.listdir(i18n_dir):
+            if filename.endswith(".json"):
+                lang_code = filename[:-5]  # Remove .json extension
+                filepath = os.path.join(i18n_dir, filename)
+                try:
+                    with open(filepath, 'r', encoding='utf-8') as f:
+                        self.translations[lang_code] = json.load(f)
+                except Exception as e:
+                    print(f"Error loading translation file {filename}: {e}")
+    
+    def set_language(self, language: str):
+        """Set current language"""
+        if language in self.translations:
+            self.current_language = language
+        else:
+            print(f"Warning: Language '{language}' not found, using default")
+    
+    def t(self, key: str, **kwargs) -> str:
+        """
+        Translate a key to current language
+        
+        Args:
+            key: Translation key (dot-separated for nested keys)
+            **kwargs: Optional format parameters
+        
+        Returns:
+            Translated string
+        """
+        # Get translation from current language
+        translation = self._get_nested_value(
+            self.translations.get(self.current_language, {}), 
+            key
+        )
+        
+        # Fallback to English if not found
+        if translation is None:
+            translation = self._get_nested_value(
+                self.translations.get('en', {}), 
+                key
+            )
+        
+        # Final fallback to key itself
+        if translation is None:
+            translation = key
+        
+        # Apply formatting if kwargs provided
+        if kwargs:
+            try:
+                translation = translation.format(**kwargs)
+            except KeyError:
+                pass
+        
+        return translation
+    
+    def _get_nested_value(self, data: dict, key: str) -> Optional[str]:
+        """
+        Get nested dictionary value using dot notation
+        
+        Args:
+            data: Dictionary to search
+            key: Dot-separated key (e.g., "section.subsection.key")
+        
+        Returns:
+            Value if found, None otherwise
+        """
+        keys = key.split('.')
+        current = data
+        
+        for k in keys:
+            if isinstance(current, dict) and k in current:
+                current = current[k]
+            else:
+                return None
+        
+        return current if isinstance(current, str) else None
+    
+    def get_available_languages(self) -> list:
+        """Get list of available language codes"""
+        return list(self.translations.keys())
+
+
+# Global i18n instance
+_i18n_instance: Optional[I18n] = None
+
+
+def get_i18n(language: Optional[str] = None) -> I18n:
+    """
+    Get global i18n instance
+    
+    Args:
+        language: Optional language to set
+    
+    Returns:
+        I18n instance
+    """
+    global _i18n_instance
+    
+    if _i18n_instance is None:
+        _i18n_instance = I18n(default_language=language or "en")
+    elif language is not None:
+        _i18n_instance.set_language(language)
+    
+    return _i18n_instance
+
+
+def t(key: str, **kwargs) -> str:
+    """
+    Convenience function for translation
+    
+    Args:
+        key: Translation key
+        **kwargs: Optional format parameters
+    
+    Returns:
+        Translated string
+    """
+    return get_i18n().t(key, **kwargs)

acestep/gradio_ui/i18n/en.json

@@ -0,0 +1,354 @@
+{
+  "app": {
+    "title": "🎛️ ACE-Step V1.5 Playground💡",
+    "subtitle": "Pushing the Boundaries of Open-Source Music Generation"
+  },
+  "dataset": {
+    "title": "📊 Dataset Explorer",
+    "dataset_label": "Dataset",
+    "dataset_info": "Choose dataset to explore",
+    "import_btn": "📥 Import Dataset",
+    "search_type_label": "Search Type",
+    "search_type_info": "How to find items",
+    "search_value_label": "Search Value",
+    "search_value_placeholder": "Enter keys or index (leave empty for random)",
+    "search_value_info": "Keys: exact match, Index: 0 to dataset size-1",
+    "instruction_label": "📝 Instruction",
+    "instruction_placeholder": "No instruction available",
+    "metadata_title": "📋 Item Metadata (JSON)",
+    "metadata_label": "Complete Item Information",
+    "source_audio": "Source Audio",
+    "target_audio": "Target Audio",
+    "reference_audio": "Reference Audio",
+    "get_item_btn": "🔍 Get Item",
+    "use_src_checkbox": "Use Source Audio from Dataset",
+    "use_src_info": "Check to use the source audio from dataset",
+    "data_status_label": "📊 Data Status",
+    "data_status_default": "❌ No dataset imported",
+    "autofill_btn": "📋 Auto-fill Generation Form"
+  },
+  "service": {
+    "title": "🔧 Service Configuration",
+    "checkpoint_label": "Checkpoint File",
+    "checkpoint_info": "Select a trained model checkpoint file (full path or filename)",
+    "refresh_btn": "🔄 Refresh",
+    "model_path_label": "Main Model Path",
+    "model_path_info": "Select the model configuration directory (auto-scanned from checkpoints)",
+    "device_label": "Device",
+    "device_info": "Processing device (auto-detect recommended)",
+    "lm_model_path_label": "5Hz LM Model Path",
+    "lm_model_path_info": "Select the 5Hz LM model checkpoint (auto-scanned from checkpoints)",
+    "backend_label": "5Hz LM Backend",
+    "backend_info": "Select backend for 5Hz LM: vllm (faster) or pt (PyTorch, more compatible)",
+    "init_llm_label": "Initialize 5Hz LM",
+    "init_llm_info": "Check to initialize 5Hz LM during service initialization",
+    "flash_attention_label": "Use Flash Attention",
+    "flash_attention_info_enabled": "Enable flash attention for faster inference (requires flash_attn package)",
+    "flash_attention_info_disabled": "Flash attention not available (flash_attn package not installed)",
+    "offload_cpu_label": "Offload to CPU",
+    "offload_cpu_info": "Offload models to CPU when not in use to save GPU memory",
+    "offload_dit_cpu_label": "Offload DiT to CPU",
+    "offload_dit_cpu_info": "Offload DiT to CPU (needs Offload to CPU)",
+    "compile_model_label": "Compile Model",
+    "compile_model_info": "Use torch.compile to optimize model (required for quantization)",
+    "quantization_label": "INT8 Quantization",
+    "quantization_info": "Enable INT8 weight-only quantization to reduce VRAM usage (requires Compile Model)",
+    "init_btn": "Initialize Service",
+    "status_label": "Status",
+    "language_label": "UI Language",
+    "language_info": "Select interface language"
+  },
+  "generation": {
+    "required_inputs": "📝 Required Inputs",
+    "task_type_label": "Task Type",
+    "task_type_info": "Select the task type for generation",
+    "instruction_label": "Instruction",
+    "instruction_info": "Instruction is automatically generated based on task type",
+    "load_btn": "Load",
+    "track_name_label": "Track Name",
+    "track_name_info": "Select track name for lego/extract tasks",
+    "track_classes_label": "Track Names",
+    "track_classes_info": "Select multiple track classes for complete task",
+    "audio_uploads": "🎵 Audio Uploads",
+    "reference_audio": "Reference Audio (optional)",
+    "source_audio": "Source Audio (optional)",
+    "convert_codes_btn": "Convert to Codes",
+    "lm_codes_hints": "🎼 LM Codes Hints",
+    "lm_codes_label": "LM Codes Hints",
+    "lm_codes_placeholder": "<|audio_code_10695|><|audio_code_54246|>...",
+    "lm_codes_info": "Paste LM codes hints for text2music generation",
+    "lm_codes_sample": "LM Codes Hints (Sample {n})",
+    "lm_codes_sample_info": "Codes for sample {n}",
+    "transcribe_btn": "Transcribe",
+    "repainting_controls": "🎨 Repainting Controls (seconds)",
+    "repainting_start": "Repainting Start",
+    "repainting_end": "Repainting End",
+    "mode_label": "Generation Mode",
+    "mode_info": "Simple: describe music in natural language. Custom: full control over caption and lyrics.",
+    "mode_simple": "Simple",
+    "mode_custom": "Custom",
+    "simple_query_label": "Song Description",
+    "simple_query_placeholder": "Describe the music you want to create, e.g., 'a soft Bengali love song for a quiet evening'. Leave empty for a random sample.",
+    "simple_query_info": "Enter a natural language description of the music you want to generate",
+    "simple_vocal_language_label": "Vocal Language (optional)",
+    "simple_vocal_language_info": "Select preferred language(s) for lyrics. Use 'unknown' for any language.",
+    "create_sample_btn": "Create Sample",
+    "caption_title": "📝 Music Caption",
+    "caption_label": "Music Caption (optional)",
+    "caption_placeholder": "A peaceful acoustic guitar melody with soft vocals...",
+    "caption_info": "Describe the style, genre, instruments, and mood",
+    "lyrics_title": "📝 Lyrics",
+    "lyrics_label": "Lyrics (optional)",
+    "lyrics_placeholder": "[Verse 1]\\nUnder the starry night\\nI feel so alive...",
+    "lyrics_info": "Song lyrics with structure",
+    "instrumental_label": "Instrumental",
+    "format_btn": "Format",
+    "optional_params": "⚙️ Optional Parameters",
+    "vocal_language_label": "Vocal Language (optional)",
+    "vocal_language_info": "use `unknown` for inst",
+    "bpm_label": "BPM (optional)",
+    "bpm_info": "leave empty for N/A",
+    "keyscale_label": "KeyScale (optional)",
+    "keyscale_placeholder": "Leave empty for N/A",
+    "keyscale_info": "A-G, #/♭, major/minor",
+    "timesig_label": "Time Signature (optional)",
+    "timesig_info": "2/4, 3/4, 4/4...",
+    "duration_label": "Audio Duration (seconds)",
+    "duration_info": "Use -1 for random",
+    "batch_size_label": "Batch Size",
+    "batch_size_info": "Number of audio to generate (max 8)",
+    "advanced_settings": "🔧 Advanced Settings",
+    "inference_steps_label": "DiT Inference Steps",
+    "inference_steps_info": "Turbo: max 8, Base: max 200",
+    "guidance_scale_label": "DiT Guidance Scale (Only support for base model)",
+    "guidance_scale_info": "Higher values follow text more closely",
+    "seed_label": "Seed",
+    "seed_info": "Use comma-separated values for batches",
+    "random_seed_label": "Random Seed",
+    "random_seed_info": "Enable to auto-generate seeds",
+    "audio_format_label": "Audio Format",
+    "audio_format_info": "Audio format for saved files",
+    "use_adg_label": "Use ADG",
+    "use_adg_info": "Enable Angle Domain Guidance",
+    "shift_label": "Shift",
+    "shift_info": "Timestep shift factor for base models (range 1.0~5.0, default 3.0). Not effective for turbo models.",
+    "infer_method_label": "Inference Method",
+    "infer_method_info": "Diffusion inference method. ODE (Euler) is faster, SDE (stochastic) may produce different results.",
+    "custom_timesteps_label": "Custom Timesteps",
+    "custom_timesteps_info": "Optional: comma-separated values from 1.0 to 0.0 (e.g., '0.97,0.76,0.615,0.5,0.395,0.28,0.18,0.085,0'). Overrides inference steps and shift.",
+    "cfg_interval_start": "CFG Interval Start",
+    "cfg_interval_end": "CFG Interval End",
+    "lm_params_title": "🤖 LM Generation Parameters",
+    "lm_temperature_label": "LM Temperature",
+    "lm_temperature_info": "5Hz LM temperature (higher = more random)",
+    "lm_cfg_scale_label": "LM CFG Scale",
+    "lm_cfg_scale_info": "5Hz LM CFG (1.0 = no CFG)",
+    "lm_top_k_label": "LM Top-K",
+    "lm_top_k_info": "Top-K (0 = disabled)",
+    "lm_top_p_label": "LM Top-P",
+    "lm_top_p_info": "Top-P (1.0 = disabled)",
+    "lm_negative_prompt_label": "LM Negative Prompt",
+    "lm_negative_prompt_placeholder": "Enter negative prompt for CFG (default: NO USER INPUT)",
+    "lm_negative_prompt_info": "Negative prompt (use when LM CFG Scale > 1.0)",
+    "cot_metas_label": "CoT Metas",
+    "cot_metas_info": "Use LM to generate CoT metadata (uncheck to skip LM CoT generation)",
+    "cot_language_label": "CoT Language",
+    "cot_language_info": "Generate language in CoT (chain-of-thought)",
+    "constrained_debug_label": "Constrained Decoding Debug",
+    "constrained_debug_info": "Enable debug logging for constrained decoding (check to see detailed logs)",
+    "auto_score_label": "Auto Score",
+    "auto_score_info": "Automatically calculate quality scores for all generated audios",
+    "auto_lrc_label": "Auto LRC",
+    "auto_lrc_info": "Automatically generate LRC lyrics timestamps for all generated audios",
+    "lm_batch_chunk_label": "LM Batch Chunk Size",
+    "lm_batch_chunk_info": "Max items per LM batch chunk (default: 8, limited by GPU memory)",
+    "codes_strength_label": "LM Codes Strength",
+    "codes_strength_info": "Control how many denoising steps use LM-generated codes",
+    "similarity_denoise_label": "Similarity / Denoise",
+    "similarity_denoise_info": "Controls how closely the output follows the reference audio. Higher values preserve more structure.",
+    "cover_strength_label": "Audio Cover Strength",
+    "cover_strength_info": "Control how many denoising steps use cover mode",
+    "score_sensitivity_label": "Quality Score Sensitivity",
+    "score_sensitivity_info": "Lower = more sensitive (default: 1.0). Adjusts how PMI maps to [0,1]",
+    "think_label": "Think",
+    "parallel_thinking_label": "ParallelThinking",
+    "generate_btn": "🎵 Generate Music",
+    "autogen_label": "AutoGen",
+    "caption_rewrite_label": "CaptionRewrite"
+  },
+  "results": {
+    "title": "🎵 Results",
+    "generated_music": "🎵 Generated Music (Sample {n})",
+    "send_to_src_btn": "🔗 Send To Src Audio",
+    "save_btn": "💾 Save",
+    "score_btn": "📊 Score",
+    "lrc_btn": "🎵 LRC",
+    "quality_score_label": "Quality Score (Sample {n})",
+    "quality_score_placeholder": "Click 'Score' to calculate perplexity-based quality score",
+    "codes_label": "LM Codes (Sample {n})",
+    "lrc_label": "Lyrics Timestamps (Sample {n})",
+    "lrc_placeholder": "Click 'LRC' to generate timestamps",
+    "details_accordion": "📊 Score & LRC & LM Codes",
+    "generation_status": "Generation Status",
+    "current_batch": "Current Batch",
+    "batch_indicator": "Batch {current} / {total}",
+    "next_batch_status": "Next Batch Status",
+    "prev_btn": "◀ Previous",
+    "next_btn": "Next ▶",
+    "restore_params_btn": "↙️ Apply These Settings to UI (Restore Batch Parameters)",
+    "batch_results_title": "👇 Click here to view batch results & generation details",
+    "all_files_label": "📁 All Generated Files (Download)",
+    "generation_details": "Generation Details"
+  },
+  "messages": {
+    "no_audio_to_save": "❌ No audio to save",
+    "save_success": "✅ Saved audio and metadata to {filename}",
+    "save_failed": "❌ Failed to save: {error}",
+    "no_file_selected": "⚠️ No file selected",
+    "params_loaded": "✅ Parameters loaded from {filename}",
+    "invalid_json": "❌ Invalid JSON file: {error}",
+    "load_error": "❌ Error loading file: {error}",
+    "example_loaded": "📁 Loaded example from {filename}",
+    "example_failed": "Failed to parse JSON file {filename}: {error}",
+    "example_error": "Error loading example: {error}",
+    "lm_generated": "🤖 Generated example using LM",
+    "lm_fallback": "Failed to generate example using LM, falling back to examples directory",
+    "lm_not_initialized": "❌ 5Hz LM not initialized. Please initialize it first.",
+    "autogen_enabled": "🔄 AutoGen enabled - next batch will generate after this",
+    "batch_ready": "✅ Batch {n} ready! Click 'Next' to view.",
+    "batch_generating": "🔄 Starting background generation for Batch {n}...",
+    "batch_failed": "❌ Background generation failed: {error}",
+    "viewing_batch": "✅ Viewing Batch {n}",
+    "at_first_batch": "Already at first batch",
+    "at_last_batch": "No next batch available",
+    "batch_not_found": "Batch {n} not found in queue",
+    "no_batch_data": "No batch data found to restore.",
+    "params_restored": "✅ UI Parameters restored from Batch {n}",
+    "scoring_failed": "❌ Error: Batch data not found",
+    "no_codes": "❌ No audio codes available. Please generate music first.",
+    "score_failed": "❌ Scoring failed: {error}",
+    "score_error": "❌ Error calculating score: {error}",
+    "lrc_no_batch_data": "❌ No batch data found. Please generate music first.",
+    "lrc_no_extra_outputs": "❌ No extra outputs found. Condition tensors not available.",
+    "lrc_missing_tensors": "❌ Missing required tensors for LRC generation.",
+    "lrc_sample_not_exist": "❌ Sample does not exist in current batch.",
+    "lrc_empty_result": "⚠️ LRC generation produced empty result.",
+    "empty_query": "⚠️ Please enter a music description.",
+    "sample_creation_failed": "❌ Failed to create sample. Please try again.",
+    "sample_created": "✅ Sample created! Review the caption and lyrics, then click Generate Music.",
+    "simple_examples_not_found": "⚠️ Simple mode examples directory not found.",
+    "simple_examples_empty": "⚠️ No example files found in simple mode examples.",
+    "simple_example_loaded": "🎲 Loaded random example from {filename}",
+    "format_success": "✅ Caption and lyrics formatted successfully",
+    "format_failed": "❌ Format failed: {error}",
+    "skipping_metas_cot": "⚡ Skipping Phase 1 metas COT (sample already formatted)",
+    "invalid_timesteps_format": "⚠️ Invalid timesteps format. Using default schedule.",
+    "timesteps_out_of_range": "⚠️ Timesteps must be in range [0, 1]. Using default schedule.",
+    "timesteps_count_mismatch": "⚠️ Timesteps count ({actual}) differs from inference_steps ({expected}). Using timesteps count."
+  },
+  "training": {
+    "tab_title": "🎓 LoRA Training",
+    "tab_dataset_builder": "📁 Dataset Builder",
+    "tab_train_lora": "🚀 Train LoRA",
+    "quick_start_title": "🚀 Quick Start",
+    "load_dataset_label": "Dataset JSON Path",
+    "load_dataset_info": "Load a previously saved dataset",
+    "load_btn": "📂 Load",
+    "load_status": "Load Status",
+    "scan_label": "Audio Directory Path",
+    "scan_info": "Scan for audio files (wav, mp3, flac, ogg, opus)",
+    "scan_btn": "🔍 Scan",
+    "scan_status": "Scan Status",
+    "found_audio_files": "Found Audio Files",
+    "dataset_name": "Dataset Name",
+    "dataset_name_placeholder": "Enter dataset name",
+    "dataset_settings_header": "Dataset Settings",
+    "tag_prepend": "Prepend (tag, caption)",
+    "tag_append": "Append (caption, tag)",
+    "tag_replace": "Replace caption",
+    "step2_title": "Step 2: Auto-Label with AI",
+    "step3_title": "Step 3: Preview & Edit",
+    "step4_title": "Step 4: Save Dataset",
+    "step5_title": "Step 5: Preprocess to Tensors",
+    "all_instrumental": "All Instrumental",
+    "all_instrumental_info": "Check if all tracks are instrumental (no vocals)",
+    "custom_tag": "Custom Activation Tag",
+    "custom_tag_info": "Unique tag to activate this LoRA's style",
+    "tag_position": "Tag Position",
+    "tag_position_info": "Where to place the custom tag in the caption",
+    "genre_ratio": "Genre Ratio (%)",
+    "genre_ratio_info": "0%=all Caption, 100%=all Genre. Per-sample override takes priority.",
+    "skip_metas": "Skip BPM/Key/Time Signature",
+    "skip_metas_info": "Skip BPM/Key/Time Signature generation. Caption and Genre are still generated by LLM.",
+    "only_unlabeled": "Only Unlabeled",
+    "only_unlabeled_info": "Only label samples without caption (useful for resuming failed labeling)",
+    "auto_label_btn": "🏷️ Auto-Label All",
+    "label_progress": "Labeling Progress",
+    "select_sample": "Select Sample #",
+    "select_sample_info": "Choose a sample to preview and edit",
+    "audio_preview": "Audio Preview",
+    "filename": "Filename",
+    "caption": "Caption",
+    "genre": "Genre",
+    "prompt_override_label": "Prompt Override (this sample)",
+    "prompt_override_info": "Override global ratio for this sample",
+    "lyrics_editable_label": "Lyrics (editable, used for training)",
+    "raw_lyrics_label": "Raw Lyrics (from .txt file)",
+    "no_lyrics_placeholder": "(no .txt lyrics file)",
+    "bpm": "BPM",
+    "key_label": "Key",
+    "key_placeholder": "C Major",
+    "time_sig": "Time Signature",
+    "duration_s": "Duration (s)",
+    "language": "Language",
+    "instrumental": "Instrumental",
+    "save_changes_btn": "💾 Save Changes",
+    "edit_status": "Edit Status",
+    "save_path": "Save Path",
+    "save_path_info": "Path where the dataset JSON will be saved",
+    "save_dataset_btn": "💾 Save Dataset",
+    "save_status": "Save Status",
+    "load_existing_label": "Load Existing Dataset (Optional)",
+    "load_existing_info": "Path to a previously saved dataset JSON file",
+    "load_dataset_btn": "📂 Load Dataset",
+    "tensor_output_dir": "Tensor Output Directory",
+    "tensor_output_info": "Directory to save preprocessed tensor files",
+    "preprocess_btn": "⚡ Preprocess",
+    "preprocess_progress": "Preprocessing Progress",
+    "preprocessed_tensors_dir": "Preprocessed Tensors Directory",
+    "preprocessed_tensors_info": "Directory containing preprocessed .pt tensor files",
+    "train_section_tensors": "Preprocessed Dataset Selection",
+    "train_section_lora": "LoRA Settings",
+    "train_section_params": "Training Parameters",
+    "dataset_info": "Dataset Info",
+    "lora_rank": "LoRA Rank (r)",
+    "lora_rank_info": "Higher = more capacity, more memory",
+    "lora_alpha": "LoRA Alpha",
+    "lora_alpha_info": "Scaling factor (typically 2x rank)",
+    "lora_dropout": "LoRA Dropout",
+    "learning_rate": "Learning Rate",
+    "learning_rate_info": "Start with 3e-4, adjust if needed",
+    "max_epochs": "Max Epochs",
+    "batch_size": "Batch Size",
+    "batch_size_info": "Increase if you have enough VRAM",
+    "gradient_accumulation": "Gradient Accumulation",
+    "gradient_accumulation_info": "Effective batch = batch_size × accumulation",
+    "save_every_n_epochs": "Save Every N Epochs",
+    "shift": "Shift",
+    "shift_info": "Timestep shift for turbo model",
+    "seed": "Seed",
+    "output_dir": "Output Directory",
+    "output_dir_info": "Directory to save trained LoRA weights",
+    "start_training_btn": "🚀 Start Training",
+    "stop_training_btn": "⏹️ Stop Training",
+    "training_progress": "Training Progress",
+    "training_log": "Training Log",
+    "training_loss_title": "Training Loss",
+    "step": "Step",
+    "loss": "Loss",
+    "export_header": "Export LoRA",
+    "export_path": "Export Path",
+    "export_lora_btn": "📦 Export LoRA",
+    "export_status": "Export Status"
+  }
+}

acestep/gradio_ui/i18n/he.json

@@ -0,0 +1,352 @@
+{
+  "app": {
+    "title": "🎛️ סביבת העבודה ACE-Step V1.5 Playground💡",
+    "subtitle": "פורצים את גבולות יצירת המוזיקה בקוד פתוח"
+  },
+  "dataset": {
+    "title": "📊 סייר מערכי נתונים (Dataset Explorer)",
+    "dataset_label": "מערך נתונים",
+    "dataset_info": "בחר מערך נתונים לחקירה",
+    "import_btn": "📥 ייבוא מערך נתונים",
+    "search_type_label": "סוג חיפוש",
+    "search_type_info": "כיצד למצוא פריטים",
+    "search_value_label": "ערך חיפוש",
+    "search_value_placeholder": "הזן מפתחות או אינדקס (השאר ריק לבחירה אקראית)",
+    "search_value_info": "מפתחות: התאמה מדויקת, אינדקס: 0 עד גודל המערך פחות 1",
+    "instruction_label": "📝 הנחיה (Instruction)",
+    "instruction_placeholder": "אין הנחיה זמינה",
+    "metadata_title": "📋 מטא-דאטה של הפריט (JSON)",
+    "metadata_label": "מידע מלא על הפריט",
+    "source_audio": "אודיו מקור",
+    "target_audio": "אודיו יעד",
+    "reference_audio": "אודיו ייחוס",
+    "get_item_btn": "🔍 קבל פריט",
+    "use_src_checkbox": "השתמש באודיו מקור ממערך הנתונים",
+    "use_src_info": "סמן כדי להשתמש באודיו המקור מתוך מערך הנתונים",
+    "data_status_label": "📊 מצב נתונים",
+    "data_status_default": "❌ לא יובא מערך נתונים",
+    "autofill_btn": "📋 מילוי אוטומטי של טופס היצירה"
+  },
+  "service": {
+    "title": "🔧 הגדרות שירות",
+    "checkpoint_label": "קובץ נקודת ביקורת (Checkpoint)",
+    "checkpoint_info": "בחר קובץ נקודת ביקורת של מודל מאומן (נתיב מלא או שם קובץ)",
+    "refresh_btn": "🔄 רענון",
+    "model_path_label": "נתיב מודל ראשי",
+    "model_path_info": "בחר את ספריית הגדרות המודל (נסרק אוטומטית מנקודות הביקורת)",
+    "device_label": "מכשיר (Device)",
+    "device_info": "מכשיר עיבוד (מומלץ זיהוי אוטומטי)",
+    "lm_model_path_label": "נתיב מודל 5Hz LM",
+    "lm_model_path_info": "בחר את קובץ נקודת הביקורת של מודל ה-5Hz LM",
+    "backend_label": "מנוע (Backend) 5Hz LM",
+    "backend_info": "בחר מנוע עבור 5Hz LM: vllm (מהיר יותר) או pt (PyTorch, תואם יותר)",
+    "init_llm_label": "אתחול 5Hz LM",
+    "init_llm_info": "סמן כדי לאתחל את ה-5Hz LM במהלך אתחול השירות",
+    "flash_attention_label": "השתמש ב-Flash Attention",
+    "flash_attention_info_enabled": "הפעל Flash Attention להסקה מהירה יותר (דורש חבילת flash_attn)",
+    "flash_attention_info_disabled": "Flash Attention אינו זמין (חבילת flash_attn לא מותקנת)",
+    "offload_cpu_label": "העברה ל-CPU (Offload)",
+    "offload_cpu_info": "העבר מודלים ל-CPU כשאינם בשימוש כדי לחסוך בזיכרון גרפי (VRAM)",
+    "offload_dit_cpu_label": "העברת DiT ל-CPU",
+    "offload_dit_cpu_info": "העבר DiT ל-CPU (דורש 'העברה ל-CPU')",
+    "compile_model_label": "הידור מודל (Compile)",
+    "compile_model_info": "השתמש ב-torch.compile לאופטימיזציה של המודל (נדרש עבור קוונטיזציה)",
+    "quantization_label": "קוונטיזציה INT8",
+    "quantization_info": "הפעל קוונטיזציה של משקולות בלבד (INT8) להפחתת שימוש ב-VRAM (דורש הידור מודל)",
+    "init_btn": "אתחול שירות",
+    "status_label": "מצב",
+    "language_label": "שפת ממשק",
+    "language_info": "בחר את שפת הממשק"
+  },
+  "generation": {
+    "required_inputs": "📝 קלטים נדרשים",
+    "task_type_label": "סוג משימה",
+    "task_type_info": "בחר את סוג המשימה ליצירה",
+    "instruction_label": "הנחיה",
+    "instruction_info": "ההנחיה נוצרת אוטומטית בהתאם לסוג המשימה",
+    "load_btn": "טעינה",
+    "track_name_label": "שם רצועה",
+    "track_name_info": "בחר שם רצועה עבור משימות lego/extract",
+    "track_classes_label": "שמות רצועות",
+    "track_classes_info": "בחר מספר מחלקות רצועה עבור משימה מלאה",
+    "audio_uploads": "🎵 העלאות אודיו",
+    "reference_audio": "אודיו ייחוס (אופציונלי)",
+    "source_audio": "אודיו מקור (אופציונלי)",
+    "convert_codes_btn": "המר לקודים",
+    "lm_codes_hints": "🎼 רמזי קודי LM",
+    "lm_codes_label": "רמזי קודי LM",
+    "lm_codes_placeholder": "<|audio_code_10695|><|audio_code_54246|>...",
+    "lm_codes_info": "הדבק רמזי קודי LM עבור יצירת טקסט למוזיקה (text2music)",
+    "lm_codes_sample": "רמזי קודי LM (דגימה {n})",
+    "lm_codes_sample_info": "קודים עבור דגימה {n}",
+    "transcribe_btn": "תמלול",
+    "repainting_controls": "🎨 בקרת צביעה מחדש (בשניות)",
+    "repainting_start": "תחילת צביעה מחדש",
+    "repainting_end": "סיום צביעה מחדש",
+    "mode_label": "מצב יצירה",
+    "mode_info": "פשוט: תאר מוזיקה בשפה טבעית. מותאם אישית: שליטה מלאה בתיאור ומילים.",
+    "mode_simple": "פשוט",
+    "mode_custom": "מותאם אישית",
+    "simple_query_label": "תיאור השיר",
+    "simple_query_placeholder": "תאר את המוזיקה שברצונך ליצור, למשל: 'שיר אהבה אקוסטי שקט לערב רגוע'. השאר ריק לדגימה אקראית.",
+    "simple_query_info": "הזן תיאור בשפה טבעית של המוזיקה שברצונך ליצור",
+    "simple_vocal_language_label": "שפת שירה (אופציונלי)",
+    "simple_vocal_language_info": "בחר שפות מועדפות למילים. השתמש ב-'unknown' לכל שפה.",
+    "create_sample_btn": "צור דגימה",
+    "caption_title": "📝 תיאור מוזיקלי (Caption)",
+    "caption_label": "תיאור מוזיקלי (אופציונלי)",
+    "caption_placeholder": "מנגינת גיטרה אקוסטית שלווה עם שירה רכה...",
+    "caption_info": "תאר את הסגנון, הז'אנר, הכלים והאווירה",
+    "lyrics_title": "📝 מילים",
+    "lyrics_label": "מילים (אופציונלי)",
+    "lyrics_placeholder": "[בית 1]\\nתחת שמי הלילה...\\nאני מרגיש חי...",
+    "lyrics_info": "מילות השיר עם מבנה",
+    "instrumental_label": "אינסטרומנטלי (ללא שירה)",
+    "format_btn": "פרמוט",
+    "optional_params": "⚙️ פרמטרים אופציונליים",
+    "vocal_language_label": "שפת שירה (אופציונלי)",
+    "vocal_language_info": "השתמש ב-`unknown` לקטעים כליים",
+    "bpm_label": "קצב (BPM) (אופציונלי)",
+    "bpm_info": "השאר ריק אם לא ידוע",
+    "keyscale_label": "סולם (KeyScale) (אופציונלי)",
+    "keyscale_placeholder": "השאר ריק אם לא ידוע",
+    "keyscale_info": "A-G, #/♭, מז'ור/מינור",
+    "timesig_label": "משקל מוזיקלי (אופציונלי)",
+    "timesig_info": "2/4, 3/4, 4/4...",
+    "duration_label": "אורך אודיו (שניות)",
+    "duration_info": "השתמש ב-1- לאקראי",
+    "batch_size_label": "גודל מנה (Batch Size)",
+    "batch_size_info": "מספר קטעי אודיו ליצירה (מקסימום 8)",
+    "advanced_settings": "🔧 הגדרות מתקדמות",
+    "inference_steps_label": "צעדי הסקה של DiT",
+    "inference_steps_info": "Turbo: מקסימום 8, Base: מקסימום 200",
+    "guidance_scale_label": "קנה מידה להנחיה (רק למודל base)",
+    "guidance_scale_info": "ערכים גבוהים יותר נצמדים יותר לטקסט",
+    "seed_label": "גרעין (Seed)",
+    "seed_info": "השתמש בערכים מופרדים בפסיקים עבור מנות",
+    "random_seed_label": "גרעין אקראי",
+    "random_seed_info": "אפשר ליצירה אוטומטית של גרעינים",
+    "audio_format_label": "פורמט אודיו",
+    "audio_format_info": "פורמט האודיו עבור הקבצים שיישמרו",
+    "use_adg_label": "השתמש ב-ADG",
+    "use_adg_info": "הפעל Angle Domain Guidance",
+    "shift_label": "Shift",
+    "shift_info": "פקטור הסטת צעדי זמן למודלי base (טווח 1.0~5.0, ברירת מחדל 3.0). לא משפיע על מודלי turbo.",
+    "infer_method_label": "שיטת הסקה",
+    "infer_method_info": "שיטת הסקת הדיפוזיה. ODE (Euler) מהירה יותר, SDE (stochastic) עשויה להפיק תוצאות שונות.",
+    "custom_timesteps_label": "צעדי זמן מותאמים אישית",
+    "custom_timesteps_info": "אופציונלי: ערכים מופרדים בפסיקים מ-1.0 עד 0.0. דורס את צעדי ההסקה וה-shift.",
+    "cfg_interval_start": "תחילת מרווח CFG",
+    "cfg_interval_end": "סיום מרווח CFG",
+    "lm_params_title": "🤖 פרמטרי יצירת LM",
+    "lm_temperature_label": "טמפרטורת LM",
+    "lm_temperature_info": "טמפרטורת 5Hz LM (גבוה יותר = אקראי יותר)",
+    "lm_cfg_scale_label": "קנה מידה LM CFG",
+    "lm_cfg_scale_info": "5Hz LM CFG (1.0 = ללא CFG)",
+    "lm_top_k_label": "LM Top-K",
+    "lm_top_k_info": "Top-K (0 = מושבת)",
+    "lm_top_p_label": "LM Top-P",
+    "lm_top_p_info": "Top-P (1.0 = מושבת)",
+    "lm_negative_prompt_label": "הנחיה שלילית ל-LM",
+    "lm_negative_prompt_placeholder": "הזן הנחיה שלילית עבור CFG",
+    "lm_negative_prompt_info": "הנחיה שלילית (בשימוש כאשר LM CFG Scale > 1.0)",
+    "cot_metas_label": "CoT Metas",
+    "cot_metas_info": "השתמש ב-LM ליצירת מטא-דאטה CoT (בטל סימון כדי לדלג)",
+    "cot_language_label": "שפת CoT",
+    "cot_language_info": "יצירת שפה ב-CoT (שרשרת מחשבה)",
+    "constrained_debug_label": "ניקוי באגים של פענוח מוגבל",
+    "constrained_debug_info": "הפעל לוגים של ניקוי באגים עבור פענוח מוגבל",
+    "auto_score_label": "דירוג אוטומטי",
+    "auto_score_info": "חשב אוטומטית ציוני איכות לכל קטעי האודיו שנוצרו",
+    "auto_lrc_label": "LRC אוטומטי",
+    "auto_lrc_info": "צור אוטומטית חותמות זמן למילים (LRC) לכל קטעי האודיו",
+    "lm_batch_chunk_label": "גודל מקטע מנת LM",
+    "lm_batch_chunk_info": "מקסימום פריטים למקטע מנת LM (ברירת מחדל: 8, מוגבל ע\"י זיכרון ה-GPU)",
+    "codes_strength_label": "חוזק קודי LM",
+    "codes_strength_info": "שליטה בכמות צעדי הניקוי מרעש המשתמשים בקודים שנוצרו ע\"י ה-LM",
+    "cover_strength_label": "חוזק כיסוי אודיו (Audio Cover)",
+    "cover_strength_info": "שליטה בכמות צעדי הניקוי מרעש המשתמשים במצב כיסוי",
+    "score_sensitivity_label": "רגישות ציון איכות",
+    "score_sensitivity_info": "נמוך יותר = רגיש יותר (ברירת מחדל: 1.0)",
+    "think_label": "חשיבה (Think)",
+    "parallel_thinking_label": "חשיבה מקבילית",
+    "generate_btn": "🎵 צור מוזיקה",
+    "autogen_label": "יצירה אוטומטית",
+    "caption_rewrite_label": "שכתוב תיאור"
+  },
+  "results": {
+    "title": "🎵 תוצאות",
+    "generated_music": "🎵 מוזיקה שנוצרה (דגימה {n})",
+    "send_to_src_btn": "🔗 שלח לאודיו מקור",
+    "save_btn": "💾 שמירה",
+    "score_btn": "📊 דירוג",
+    "lrc_btn": "🎵 LRC",
+    "quality_score_label": "ציון איכות (דגימה {n})",
+    "quality_score_placeholder": "לחץ על 'דירוג' לחישוב ציון איכות מבוסס מורכבות (Perplexity)",
+    "codes_label": "קודי LM (דגימה {n})",
+    "lrc_label": "חותמות זמן למילים (דגימה {n})",
+    "lrc_placeholder": "לחץ על 'LRC' ליצירת חותמות זמן",
+    "details_accordion": "📊 דירוג, LRC וקודי LM",
+    "generation_status": "מצב יצירה",
+    "current_batch": "מנה נוכחית",
+    "batch_indicator": "מנה {current} / {total}",
+    "next_batch_status": "מצב המנה הבאה",
+    "prev_btn": "◀ הקודם",
+    "next_btn": "הבא ▶",
+    "restore_params_btn": "↙️ החל הגדרות אלו על הממשק (שחזור פרמטרי מנה)",
+    "batch_results_title": "📁 תוצאות המנה ופרטי יצירה",
+    "all_files_label": "📁 כל הקבצים שנוצרו (הורדה)",
+    "generation_details": "פרטי יצירה"
+  },
+  "messages": {
+    "no_audio_to_save": "❌ אין אודיו לשמירה",
+    "save_success": "✅ האודיו והמטא-דאטה נשמרו ב-{filename}",
+    "save_failed": "❌ השמירה נכשלה: {error}",
+    "no_file_selected": "⚠️ לא נבחר קובץ",
+    "params_loaded": "✅ הפרמטרים נטענו מ-{filename}",
+    "invalid_json": "❌ קובץ JSON לא תקין: {error}",
+    "load_error": "❌ שגיאה בטעינת הקובץ: {error}",
+    "example_loaded": "📁 נטען דגם מ-{filename}",
+    "example_failed": "נכשל ניתוח קובץ ה-JSON ב-{filename}: {error}",
+    "example_error": "שגיאה בטעינת הדגם: {error}",
+    "lm_generated": "🤖 נוצר דגם באמצעות ה-LM",
+    "lm_fallback": "יצירת דגם באמצעות ה-LM נכשלה, חוזר לשימוש בספריית הדגמים",
+    "lm_not_initialized": "❌ 5Hz LM לא מאותחל. נא לאתחל אותו תחילה.",
+    "autogen_enabled": "🔄 יצירה אוטומטית הופעלה - המנה הבאה תיווצר לאחר מכן",
+    "batch_ready": "✅ מנה {n} מוכנה! לחץ על 'הבא' לצפייה.",
+    "batch_generating": "🔄 מתחיל יצירת רקע עבור מנה {n}...",
+    "batch_failed": "❌ יצירת הרקע נכשלה: {error}",
+    "viewing_batch": "✅ צופה במנה {n}",
+    "at_first_batch": "נמצא כבר במנה הראשונה",
+    "at_last_batch": "אין מנה באה זמינה",
+    "batch_not_found": "מנה {n} לא נמצאה בתור",
+    "no_batch_data": "לא נמצאו נתוני מנה לשחזור.",
+    "params_restored": "✅ פרמטרי הממשק שוחזרו ממנה {n}",
+    "scoring_failed": "❌ שגיאה: נתוני המנה לא נמצאו",
+    "no_codes": "❌ אין קודי אודיו זמינים. נא ליצור מוזיקה תחילה.",
+    "score_failed": "❌ הדירוג נכשל: {error}",
+    "score_error": "❌ שגיאה בחישוב הציון: {error}",
+    "lrc_no_batch_data": "❌ לא נמצאו נתוני מנה. נא ליצור מוזיקה תחילה.",
+    "lrc_no_extra_outputs": "❌ לא נמצאו פלטים נוספים. טנזורי התניה אינם זמינים.",
+    "lrc_missing_tensors": "❌ חסרים טנזורים נדרשים ליצירת LRC.",
+    "lrc_sample_not_exist": "❌ הדגימה אינה קיימת במנה הנוכחית.",
+    "lrc_empty_result": "⚠️ יצירת ה-LRC הפיקה תוצאה ריקה.",
+    "empty_query": "⚠️ נא להזין תיאור מוזיקלי.",
+    "sample_creation_failed": "❌ יצירת הדגימה נכשלה. נא לנסות שוב.",
+    "sample_created": "✅ הדגימה נוצרה! בדוק את התיאור והמילים, ולאחר מכן לחץ על 'צור מוזיקה'.",
+    "simple_examples_not_found": "⚠️ ספריית הדגמים של המצב הפשוט לא נמצאה.",
+    "simple_examples_empty": "⚠️ לא נמצאו קבצי דוגמה במצב פשוט.",
+    "simple_example_loaded": "🎲 נטענה דוגמה אקראית מ-{filename}",
+    "format_success": "✅ התיאור והמילים פורמטו בהצלחה",
+    "format_failed": "❌ הפירמוט נכשל: {error}",
+    "skipping_metas_cot": "⚡ מדלג על שלב 1 של מטא-דאטה COT (הדגימה כבר מפורמטת)",
+    "invalid_timesteps_format": "⚠️ פורמט צעדי זמן לא תקין. משתמש בלוח זמנים כברירת מחדל.",
+    "timesteps_out_of_range": "⚠️ צעדי הזמן חייבים להיות בטווח [0, 1]. משתמש בלוח זמנים כברירת מחדל.",
+    "timesteps_count_mismatch": "⚠️ מספר צעדי הזמן ({actual}) שונה מצעדי ההסקה ({expected}). משתמש במספר צעדי הזמן."
+  },
+  "training": {
+    "tab_title": "🎓 אימון LoRA",
+    "tab_dataset_builder": "📁 בונה מערך נתונים",
+    "tab_train_lora": "🚀 אימון LoRA",
+    "quick_start_title": "🚀 התחלה מהירה",
+    "load_dataset_label": "נתיב קובץ JSON של מערך הנתונים",
+    "load_dataset_info": "טעינת מערך נתונים שנשמר בעבר",
+    "load_btn": "📂 טעינה",
+    "load_status": "מצב טעינה",
+    "scan_label": "נתיב ספריית אודיו",
+    "scan_info": "סריקה אחר קבצי אודיו (wav, mp3, flac, ogg, opus)",
+    "scan_btn": "🔍 סריקה",
+    "scan_status": "מצב סריקה",
+    "found_audio_files": "קבצי אודיו שנמצאו",
+    "dataset_name": "שם מערך הנתונים",
+    "dataset_name_placeholder": "הזן שם למערך הנתונים",
+    "dataset_settings_header": "הגדרות מערך נתונים",
+    "tag_prepend": "הוספה בהתחלה (תגית, תיאור)",
+    "tag_append": "הוספה בסוף (תיאור, תגית)",
+    "tag_replace": "החלפת התיאור",
+    "step2_title": "שלב 2: תיוג אוטומטי באמצעות AI",
+    "step3_title": "שלב 3: תצוגה מקדימה ועריכה",
+    "step4_title": "שלב 4: שמירת מערך הנתונים",
+    "step5_title": "שלב 5: עיבוד מקדים לטנזורים (Tensors)",
+    "all_instrumental": "הכל אינסטרומנטלי",
+    "all_instrumental_info": "סמן אם כל הרצועות הן כליות (ללא שירה)",
+    "custom_tag": "תגית הפעלה מותאמת אישית",
+    "custom_tag_info": "תגית ייחודית להפעלת הסגנון של LoRA זו",
+    "tag_position": "מיקום התגית",
+    "tag_position_info": "היכן למקם את התגית המותאמת אישית בתוך התיאור",
+    "genre_ratio": "יחס ז'אנר (%)",
+    "genre_ratio_info": "0% = הכל תיאור, 100% = הכל ז'אנר. הגדרה פר-דגימה קודמת להגדרת הכלל.",
+    "skip_metas": "דלג על BPM/סולם/משקל",
+    "skip_metas_info": "דלג על יצירת BPM/סולם/משקל. התיאור והז'אנר עדיין ייווצרו על ידי ה-LLM.",
+    "only_unlabeled": "רק כאלו ללא תיוג",
+    "only_unlabeled_info": "תייג רק דגימות ללא תיאור (שימושי להמשך תיוג שנכשל)",
+    "auto_label_btn": "🏷️ תיוג אוטומטי של הכל",
+    "label_progress": "התקדמות התיוג",
+    "select_sample": "בחר דגימה #",
+    "select_sample_info": "בחר דגימה לצפייה ועריכה",
+    "audio_preview": "תצוגה מקדימה של אודיו",
+    "filename": "שם קובץ",
+    "caption": "תיאור",
+    "genre": "ז'אנר",
+    "prompt_override_label": "דריסת פרומפט (לדגימה זו)",
+    "prompt_override_info": "דריסת היחס הכללי עבור דגימה ז��",
+    "lyrics_editable_label": "מילים (ניתן לעריכה, משמש לאימון)",
+    "raw_lyrics_label": "מילים גולמיות (מתוך קובץ .txt)",
+    "no_lyrics_placeholder": "(אין קובץ מילים .txt)",
+    "bpm": "BPM",
+    "key_label": "סולם (Key)",
+    "key_placeholder": "C Major",
+    "time_sig": "משקל מוזיקלי",
+    "duration_s": "משך (שניות)",
+    "language": "שפה",
+    "instrumental": "אינסטרומנטלי",
+    "save_changes_btn": "💾 שמירת שינויים",
+    "edit_status": "מצב עריכה",
+    "save_path": "נתיב שמירה",
+    "save_path_info": "הנתיב שבו יישמר קובץ ה-JSON של מערך הנתונים",
+    "save_dataset_btn": "💾 שמירת מערך נתונים",
+    "save_status": "מצב שמירה",
+    "load_existing_label": "טעינת מערך נתונים קיים (אופציונלי)",
+    "load_existing_info": "נתיב לקובץ JSON של מערך נתונים שנשמר בעבר",
+    "load_dataset_btn": "📂 טעינת מערך נתונים",
+    "tensor_output_dir": "ספריית פלט של טנזורים",
+    "tensor_output_info": "הספרייה לשמירת קבצי טנזור שעברו עיבוד מקדים",
+    "preprocess_btn": "⚡ עיבוד מקדים",
+    "preprocess_progress": "התקדמות עיבוד מקדים",
+    "preprocessed_tensors_dir": "ספריית טנזורים מעובדים",
+    "preprocessed_tensors_info": "ספרייה המכילה קבצי .pt של טנזורים מעובדים",
+    "train_section_tensors": "בחירת מערך נתונים מעובד",
+    "train_section_lora": "הגדרות LoRA",
+    "train_section_params": "פרמטרי אימון",
+    "dataset_info": "מידע על מערך הנתונים",
+    "lora_rank": "דרגת LoRA (Rank)",
+    "lora_rank_info": "גבוה יותר = יותר קיבולת, יותר זיכרון",
+    "lora_alpha": "LoRA Alpha",
+    "lora_alpha_info": "פקטור קנה מידה (בדרך כלל פי 2 מה-Rank)",
+    "lora_dropout": "LoRA Dropout",
+    "learning_rate": "קצב למידה (Learning Rate)",
+    "learning_rate_info": "התחל עם 3e-4, שנה במידת הצורך",
+    "max_epochs": "מקסימום תקופות (Epochs)",
+    "batch_size": "גודל מנה (Batch Size)",
+    "batch_size_info": "הגדל אם יש לך מספיק זיכרון גרפי (VRAM)",
+    "gradient_accumulation": "צבירת גרדיאנטים (Accumulation)",
+    "gradient_accumulation_info": "גודל מנה אפקטיבי = גודל מנה × צבירה",
+    "save_every_n_epochs": "שמור כל N תקופות (Epochs)",
+    "shift": "Shift (הסטה)",
+    "shift_info": "הסטת צעדי זמן עבור מודל turbo",
+    "seed": "גרעין (Seed)",
+    "output_dir": "ספריית פלט",
+    "output_dir_info": "ספרייה לשמירת משקולות ה-LoRA המאומנות",
+    "start_training_btn": "🚀 התחלת אימון",
+    "stop_training_btn": "⏹️ עצירת אימון",
+    "training_progress": "התקדמות האימון",
+    "training_log": "יומן אימון",
+    "training_loss_title": "הפסד אימון (Training Loss)",
+    "step": "צעד",
+    "loss": "הפסד (Loss)",
+    "export_header": "ייצוא LoRA",
+    "export_path": "נתיב ייצוא",
+    "export_lora_btn": "📦 ייצוא LoRA",
+    "export_status": "מצב ייצוא"
+  }
+}

acestep/gradio_ui/i18n/ja.json

@@ -0,0 +1,354 @@
+{
+  "app": {
+    "title": "🎛️ ACE-Step V1.5 プレイグラウンド💡",
+    "subtitle": "オープンソース音楽生成の限界を押し広げる"
+  },
+  "dataset": {
+    "title": "📊 データセットエクスプローラー",
+    "dataset_label": "データセット",
+    "dataset_info": "探索するデータセットを選択",
+    "import_btn": "📥 データセットをインポート",
+    "search_type_label": "検索タイプ",
+    "search_type_info": "アイテムの検索方法",
+    "search_value_label": "検索値",
+    "search_value_placeholder": "キーまたはインデックスを入力(空白の場合はランダム)",
+    "search_value_info": "キー: 完全一致、インデックス: 0からデータセットサイズ-1",
+    "instruction_label": "📝 指示",
+    "instruction_placeholder": "利用可能な指示がありません",
+    "metadata_title": "📋 アイテムメタデータ (JSON)",
+    "metadata_label": "完全なアイテム情報",
+    "source_audio": "ソースオーディオ",
+    "target_audio": "ターゲットオーディオ",
+    "reference_audio": "リファレンスオーディオ",
+    "get_item_btn": "🔍 アイテムを取得",
+    "use_src_checkbox": "データセットのソースオーディオを使用",
+    "use_src_info": "データセットのソースオーディオを使用する場合はチェック",
+    "data_status_label": "📊 データステータス",
+    "data_status_default": "❌ データセットがインポートされていません",
+    "autofill_btn": "📋 生成フォームを自動入力"
+  },
+  "service": {
+    "title": "🔧 サービス設定",
+    "checkpoint_label": "チェックポイントファイル",
+    "checkpoint_info": "訓練済みモデルのチェックポイントファイルを選択(フルパスまたはファイル名)",
+    "refresh_btn": "🔄 更新",
+    "model_path_label": "メインモデルパス",
+    "model_path_info": "モデル設定ディレクトリを選択(チェックポイントから自動スキャン)",
+    "device_label": "デバイス",
+    "device_info": "処理デバイス(自動検出を推奨)",
+    "lm_model_path_label": "5Hz LM モデルパス",
+    "lm_model_path_info": "5Hz LMモデルチェックポイントを選択(チェックポイントから自動スキャン)",
+    "backend_label": "5Hz LM バックエンド",
+    "backend_info": "5Hz LMのバックエンドを選択: vllm(高速)またはpt(PyTorch、より互換性あり)",
+    "init_llm_label": "5Hz LM を初期化",
+    "init_llm_info": "サービス初期化中に5Hz LMを初期化する場合はチェック",
+    "flash_attention_label": "Flash Attention を使用",
+    "flash_attention_info_enabled": "推論を高速化するためにflash attentionを有効にする(flash_attnパッケージが必要)",
+    "flash_attention_info_disabled": "Flash attentionは利用できません(flash_attnパッケージがインストールされていません)",
+    "offload_cpu_label": "CPUにオフロード",
+    "offload_cpu_info": "使用していない時にモデルをCPUにオフロードしてGPUメモリを節約",
+    "offload_dit_cpu_label": "DiTをCPUにオフロード",
+    "offload_dit_cpu_info": "DiTをCPUにオフロード(CPUへのオフロードが必要)",
+    "compile_model_label": "モデルをコンパイル",
+    "compile_model_info": "torch.compileでモデルを最適化（量子化に必要）",
+    "quantization_label": "INT8 量子化",
+    "quantization_info": "INT8重み量子化を有効にしてVRAMを節約（モデルのコンパイルが必要）",
+    "init_btn": "サービスを初期化",
+    "status_label": "ステータス",
+    "language_label": "UI言語",
+    "language_info": "インターフェース言語を選択"
+  },
+  "generation": {
+    "required_inputs": "📝 必須入力",
+    "task_type_label": "タスクタイプ",
+    "task_type_info": "生成のタスクタイプを選択",
+    "instruction_label": "指示",
+    "instruction_info": "指示はタスクタイプに基づいて自動生成されます",
+    "load_btn": "読み込む",
+    "track_name_label": "トラック名",
+    "track_name_info": "lego/extractタスクのトラック名を選択",
+    "track_classes_label": "トラック名",
+    "track_classes_info": "completeタスクの複数のトラッククラスを選択",
+    "audio_uploads": "🎵 オーディオアップロード",
+    "reference_audio": "リファレンスオーディオ(オプション)",
+    "source_audio": "ソースオーディオ(オプション)",
+    "convert_codes_btn": "コードに変換",
+    "lm_codes_hints": "🎼 LM コードヒント",
+    "lm_codes_label": "LM コードヒント",
+    "lm_codes_placeholder": "<|audio_code_10695|><|audio_code_54246|>...",
+    "lm_codes_info": "text2music生成用のLMコードヒントを貼り付け",
+    "lm_codes_sample": "LM コードヒント(サンプル {n})",
+    "lm_codes_sample_info": "サ���プル{n}のコード",
+    "transcribe_btn": "転写",
+    "repainting_controls": "🎨 再描画コントロール(秒)",
+    "repainting_start": "再描画開始",
+    "repainting_end": "再描画終了",
+    "mode_label": "生成モード",
+    "mode_info": "シンプル：自然言語で音楽を説明。カスタム：キャプションと歌詞を完全にコントロール。",
+    "mode_simple": "シンプル",
+    "mode_custom": "カスタム",
+    "simple_query_label": "曲の説明",
+    "simple_query_placeholder": "作成したい音楽を説明してください。例：'静かな夜のための優しいベンガルのラブソング'。空欄の場合はランダムなサンプルが生成されます。",
+    "simple_query_info": "生成したい音楽の自然言語の説明を入力",
+    "simple_vocal_language_label": "ボーカル言語(オプション)",
+    "simple_vocal_language_info": "歌詞の希望言語を選択。任意の言語の場合は'unknown'を使用。",
+    "create_sample_btn": "サンプル作成",
+    "caption_title": "📝 音楽キャプション",
+    "caption_label": "音楽キャプション(オプション)",
+    "caption_placeholder": "柔らかいボーカルを伴う穏やかなアコースティックギターのメロディー...",
+    "caption_info": "スタイル、ジャンル、楽器、ムードを説明",
+    "lyrics_title": "📝 歌詞",
+    "lyrics_label": "歌詞(オプション)",
+    "lyrics_placeholder": "[バース1]\\n星空の下で\\nとても生きていると感じる...",
+    "lyrics_info": "構造を持つ曲の歌詞",
+    "instrumental_label": "インストゥルメンタル",
+    "format_btn": "フォーマット",
+    "optional_params": "⚙️ オプションパラメータ",
+    "vocal_language_label": "ボーカル言語(オプション)",
+    "vocal_language_info": "インストには`unknown`を使用",
+    "bpm_label": "BPM(オプション)",
+    "bpm_info": "空白の場合はN/A",
+    "keyscale_label": "キースケール(オプション)",
+    "keyscale_placeholder": "空白の場合はN/A",
+    "keyscale_info": "A-G, #/♭, メジャー/マイナー",
+    "timesig_label": "拍子記号(オプション)",
+    "timesig_info": "2/4, 3/4, 4/4...",
+    "duration_label": "オーディオ長(秒)",
+    "duration_info": "ランダムの場合は-1を使用",
+    "batch_size_label": "バッチサイズ",
+    "batch_size_info": "生成するオーディオの数(最大8)",
+    "advanced_settings": "🔧 詳細設定",
+    "inference_steps_label": "DiT 推論ステップ",
+    "inference_steps_info": "Turbo: 最大8、Base: 最大200",
+    "guidance_scale_label": "DiT ガイダンススケール(baseモデルのみサポート)",
+    "guidance_scale_info": "値が高いほどテキストに忠実に従う",
+    "seed_label": "シード",
+    "seed_info": "バッチにはカンマ区切りの値を使用",
+    "random_seed_label": "ランダムシード",
+    "random_seed_info": "有効にすると自動的にシードを生成",
+    "audio_format_label": "オーディオフォーマット",
+    "audio_format_info": "保存ファイルのオーディオフォーマット",
+    "use_adg_label": "ADG を使用",
+    "use_adg_info": "角度ドメインガイダンスを有効化",
+    "shift_label": "シフト",
+    "shift_info": "baseモデル用タイムステップシフト係数 (範囲 1.0~5.0、デフォルト 3.0)。turboモデルには無効。",
+    "infer_method_label": "推論方法",
+    "infer_method_info": "拡散推論方法。ODE (オイラー) は高速、SDE (確率的) は異なる結果を生成する可能性があります。",
+    "custom_timesteps_label": "カスタムタイムステップ",
+    "custom_timesteps_info": "オプション：1.0から0.0へのカンマ区切り値（例：'0.97,0.76,0.615,0.5,0.395,0.28,0.18,0.085,0'）。推論ステップとシフトを上書きします。",
+    "cfg_interval_start": "CFG 間隔開始",
+    "cfg_interval_end": "CFG 間隔終了",
+    "lm_params_title": "🤖 LM 生成パラメータ",
+    "lm_temperature_label": "LM 温度",
+    "lm_temperature_info": "5Hz LM温度(高いほどランダム)",
+    "lm_cfg_scale_label": "LM CFG スケール",
+    "lm_cfg_scale_info": "5Hz LM CFG (1.0 = CFGなし)",
+    "lm_top_k_label": "LM Top-K",
+    "lm_top_k_info": "Top-K (0 = 無効)",
+    "lm_top_p_label": "LM Top-P",
+    "lm_top_p_info": "Top-P (1.0 = 無効)",
+    "lm_negative_prompt_label": "LM ネガティブプロンプト",
+    "lm_negative_prompt_placeholder": "CFGのネガティブプロンプトを入力(デフォルト: NO USER INPUT)",
+    "lm_negative_prompt_info": "ネガティブプロンプト(LM CFGスケール > 1.0の場合に使用)",
+    "cot_metas_label": "CoT メタデータ",
+    "cot_metas_info": "LMを使用してCoTメタデータを生成(チェックを外すとLM CoT生成をスキップ)",
+    "cot_language_label": "CoT 言語",
+    "cot_language_info": "CoTで言語を生成(思考の連鎖)",
+    "constrained_debug_label": "制約付きデコーディングデバッグ",
+    "constrained_debug_info": "制約付きデコーディングのデバッグログを有効化(チェックすると詳細ログを表示)",
+    "auto_score_label": "自動スコアリング",
+    "auto_score_info": "生成されたすべてのオーディオの品質スコアを自動計算",
+    "auto_lrc_label": "自動 LRC",
+    "auto_lrc_info": "生成されたすべてのオーディオのLRC歌詞タイムスタンプを自動生成",
+    "lm_batch_chunk_label": "LM バッチチャンクサイズ",
+    "lm_batch_chunk_info": "LMバッチチャンクあたりの最大アイテム数(デフォルト: 8、GPUメモリによる制限)",
+    "codes_strength_label": "LM コード強度",
+    "codes_strength_info": "LM生成コードを使用するデノイジングステップ数を制御",
+    "similarity_denoise_label": "類似度 / ノイズ除去",
+    "similarity_denoise_info": "出力が参照オーディオにどれだけ忠実かを制御します。高い値ほど構造を保持します。",
+    "cover_strength_label": "オーディオカバー強度",
+    "cover_strength_info": "カバーモードを使用するデノイジングステップ数を制御",
+    "score_sensitivity_label": "品質スコア感度",
+    "score_sensitivity_info": "低い = より敏感(デフォルト: 1.0)。PMIが[0,1]にマッピングする方法を調整",
+    "think_label": "思考",
+    "parallel_thinking_label": "並列思考",
+    "generate_btn": "🎵 音楽を生成",
+    "autogen_label": "自動生成",
+    "caption_rewrite_label": "キャプション書き換え"
+  },
+  "results": {
+    "title": "🎵 結果",
+    "generated_music": "🎵 生成された音楽(サンプル {n})",
+    "send_to_src_btn": "🔗 ソースオーディオに送信",
+    "save_btn": "💾 保存",
+    "score_btn": "📊 スコア",
+    "lrc_btn": "🎵 LRC",
+    "quality_score_label": "品質スコア(サンプル {n})",
+    "quality_score_placeholder": "'スコア'をクリックしてパープレキシティベースの品質スコアを計算",
+    "codes_label": "LM コード(サンプル {n})",
+    "lrc_label": "歌詞タイムスタンプ(サンプル {n})",
+    "lrc_placeholder": "'LRC'をクリックしてタイムスタンプを生成",
+    "details_accordion": "📊 スコア & LRC & LM コード",
+    "generation_status": "生成ステータス",
+    "current_batch": "現在のバッチ",
+    "batch_indicator": "バッチ {current} / {total}",
+    "next_batch_status": "次のバッチステータス",
+    "prev_btn": "◀ 前へ",
+    "next_btn": "次へ ▶",
+    "restore_params_btn": "↙️ これらの設定をUIに適用(バッチパラメータを復元)",
+    "batch_results_title": "👇 クリックしてバッチ結果と生成詳細を表示",
+    "all_files_label": "📁 すべての生成ファイル(ダウンロード)",
+    "generation_details": "生成詳細"
+  },
+  "messages": {
+    "no_audio_to_save": "❌ 保存するオーディオがありません",
+    "save_success": "✅ オーディオとメタデータを {filename} に保存しました",
+    "save_failed": "❌ 保存に失敗しました: {error}",
+    "no_file_selected": "⚠️ ファイルが選択されていません",
+    "params_loaded": "✅ {filename} からパラメータを読み込みました",
+    "invalid_json": "❌ 無効なJSONファイル: {error}",
+    "load_error": "❌ ファイルの読み込みエラー: {error}",
+    "example_loaded": "📁 {filename} からサンプルを読み込みました",
+    "example_failed": "JSONファイル {filename} の解析に失敗しました: {error}",
+    "example_error": "サンプル読み込みエラー: {error}",
+    "lm_generated": "🤖 LMを使用してサンプルを生成しました",
+    "lm_fallback": "LMを使用したサンプル生成に失敗、サンプルディレクトリにフォールバック",
+    "lm_not_initialized": "❌ 5Hz LMが初期化されていません。最初に初期化してください。",
+    "autogen_enabled": "🔄 自動生成が有効 - このあと次のバッチを生成します",
+    "batch_ready": "✅ バッチ {n} の準備完了！'次へ'をクリックして表示。",
+    "batch_generating": "🔄 バッチ {n} のバックグラウンド生成を開始...",
+    "batch_failed": "❌ バックグラウンド生成に失敗しました: {error}",
+    "viewing_batch": "✅ バッチ {n} を表示中",
+    "at_first_batch": "すでに最初のバッチです",
+    "at_last_batch": "次のバッチはありません",
+    "batch_not_found": "キューにバッチ {n} が見つかりません",
+    "no_batch_data": "復元するバッチデータがありません。",
+    "params_restored": "✅ バッチ {n} からUIパラメータを復元しました",
+    "scoring_failed": "❌ エラー: バッチデータが見つかりません",
+    "no_codes": "❌ 利用可能なオーディオ��ードがありません。最初に音楽を生成してください。",
+    "score_failed": "❌ スコアリングに失敗しました: {error}",
+    "score_error": "❌ スコア計算エラー: {error}",
+    "lrc_no_batch_data": "❌ バッチデータが見つかりません。最初に音楽を生成してください。",
+    "lrc_no_extra_outputs": "❌ 追加出力が見つかりません。条件テンソルが利用できません。",
+    "lrc_missing_tensors": "❌ LRC生成に必要なテンソルがありません。",
+    "lrc_sample_not_exist": "❌ 現在のバッチにサンプルが存在しません。",
+    "lrc_empty_result": "⚠️ LRC生成の結果が空です。",
+    "empty_query": "⚠️ 音楽の説明を入力してください。",
+    "sample_creation_failed": "❌ サンプルの作成に失敗しました。もう一度お試しください。",
+    "sample_created": "✅ サンプルが作成されました！キャプションと歌詞を確認して、音楽を生成をクリックしてください。",
+    "simple_examples_not_found": "⚠️ シンプルモードサンプルディレクトリが見つかりません。",
+    "simple_examples_empty": "⚠️ シンプルモードサンプルにファイルがありません。",
+    "simple_example_loaded": "🎲 {filename} からランダムサンプルを読み込みました",
+    "format_success": "✅ キャプションと歌詞のフォーマットに成功しました",
+    "format_failed": "❌ フォーマットに失敗しました: {error}",
+    "skipping_metas_cot": "⚡ Phase 1 メタデータ COT をスキップ（サンプルは既にフォーマット済み）",
+    "invalid_timesteps_format": "⚠️ タイムステップ形式が無効です。デフォルトスケジュールを使用します。",
+    "timesteps_out_of_range": "⚠️ タイムステップは [0, 1] の範囲内である必要があります。デフォルトスケジュールを使用します。",
+    "timesteps_count_mismatch": "⚠️ タイムステップ数 ({actual}) が推論ステップ数 ({expected}) と異なります。タイムステップ数を使用します。"
+  },
+  "training": {
+    "tab_title": "🎓 LoRA トレーニング",
+    "tab_dataset_builder": "📁 データセットビルダー",
+    "tab_train_lora": "🚀 LoRA をトレーニング",
+    "quick_start_title": "🚀 クイックスタート",
+    "load_dataset_label": "データセット JSON パス",
+    "load_dataset_info": "以前保存したデータセットを読み込む",
+    "load_btn": "📂 読み込み",
+    "load_status": "読み込み状態",
+    "scan_label": "オーディオディレクトリパス",
+    "scan_info": "オーディオファイルをスキャン（wav、mp3、flac、ogg、opus）",
+    "scan_btn": "🔍 スキャン",
+    "scan_status": "スキャン状態",
+    "found_audio_files": "見つかったオーディオファイル",
+    "dataset_name": "データセット名",
+    "dataset_name_placeholder": "データセット名を入力",
+    "dataset_settings_header": "データセット設定",
+    "tag_prepend": "前置（タグ、キャプション）",
+    "tag_append": "後置（キャプション、タグ）",
+    "tag_replace": "キャプションを置換",
+    "step2_title": "ステップ 2: AI で自動ラベル",
+    "step3_title": "ステップ 3: プレビューと編集",
+    "step4_title": "ステップ 4: データセットを保存",
+    "step5_title": "ステップ 5: テンソルに前処理",
+    "all_instrumental": "すべてインストゥルメンタル",
+    "all_instrumental_info": "すべてのトラックがインストゥルメンタル（ボーカルなし）の場合にチェック",
+    "custom_tag": "カスタムアクティベーションタグ",
+    "custom_tag_info": "この LoRA のスタイルを有効にする一意のタグ",
+    "tag_position": "タグの位置",
+    "tag_position_info": "キャプション内でカスタムタグを配置する位置",
+    "genre_ratio": "ジャンル比率 (%)",
+    "genre_ratio_info": "0%=すべてキャプション、100%=すべてジャンル。サンプル単位の上書きが優先。",
+    "skip_metas": "BPM/キー/拍子をスキップ",
+    "skip_metas_info": "BPM/キー/拍子の生成をスキップ。キャプションとジャンルは LM が生成。",
+    "only_unlabeled": "未ラベルのみ",
+    "only_unlabeled_info": "キャプションのないサンプルのみラベル付け（失敗したラベル付けの再開に便利）",
+    "auto_label_btn": "🏷️ 一括自動ラベル",
+    "label_progress": "ラベル付け進捗",
+    "select_sample": "サンプル # を選択",
+    "select_sample_info": "プレビューと編集するサンプルを選択",
+    "audio_preview": "オーディオプレビュー",
+    "filename": "ファイル名",
+    "caption": "キャプション",
+    "genre": "ジャンル",
+    "prompt_override_label": "プロンプト上書き（このサ���プル）",
+    "prompt_override_info": "このサンプルのグローバル比率を上書き",
+    "lyrics_editable_label": "歌詞（編集可、トレーニング用）",
+    "raw_lyrics_label": "生歌詞（.txt ファイルから）",
+    "no_lyrics_placeholder": "（.txt 歌詞ファイルなし）",
+    "bpm": "BPM",
+    "key_label": "キー",
+    "key_placeholder": "C Major",
+    "time_sig": "拍子",
+    "duration_s": "長さ (秒)",
+    "language": "言語",
+    "instrumental": "インストゥルメンタル",
+    "save_changes_btn": "💾 変更を保存",
+    "edit_status": "編集状態",
+    "save_path": "保存パス",
+    "save_path_info": "データセット JSON の保存先パス",
+    "save_dataset_btn": "💾 データセットを保存",
+    "save_status": "保存状態",
+    "load_existing_label": "既存データセットを読み込み（任意）",
+    "load_existing_info": "以前保存したデータセット JSON ファイルのパス",
+    "load_dataset_btn": "📂 データセットを読み込み",
+    "tensor_output_dir": "テンソル出力ディレクトリ",
+    "tensor_output_info": "前処理済みテンソルファイルの保存先ディレクトリ",
+    "preprocess_btn": "⚡ 前処理",
+    "preprocess_progress": "前処理進捗",
+    "preprocessed_tensors_dir": "前処理済みテンソルディレクトリ",
+    "preprocessed_tensors_info": "前処理済み .pt テンソルファイルを含むディレクトリ",
+    "train_section_tensors": "前処理済みデータセット選択",
+    "train_section_lora": "LoRA 設定",
+    "train_section_params": "トレーニングパラメータ",
+    "dataset_info": "データセット情報",
+    "lora_rank": "LoRA ランク (r)",
+    "lora_rank_info": "高いほど容量は増えるがメモリ使用量も増加",
+    "lora_alpha": "LoRA Alpha",
+    "lora_alpha_info": "スケーリング係数（通常はランクの2倍）",
+    "lora_dropout": "LoRA Dropout",
+    "learning_rate": "学習率",
+    "learning_rate_info": "3e-4 から始め、必要に応じて調整",
+    "max_epochs": "最大エポック数",
+    "batch_size": "バッチサイズ",
+    "batch_size_info": "VRAM に余裕があれば増やせます",
+    "gradient_accumulation": "勾配累積",
+    "gradient_accumulation_info": "実効バッチ = batch_size × 累積",
+    "save_every_n_epochs": "N エポックごとに保存",
+    "shift": "Shift",
+    "shift_info": "ターボモデル用タイムステップシフト",
+    "seed": "シード",
+    "output_dir": "出力ディレクトリ",
+    "output_dir_info": "トレーニング済み LoRA 重みの保存先ディレクトリ",
+    "start_training_btn": "🚀 トレーニング開始",
+    "stop_training_btn": "⏹️ トレーニング停止",
+    "training_progress": "トレーニング進捗",
+    "training_log": "トレーニングログ",
+    "training_loss_title": "トレーニング損失",
+    "step": "ステップ",
+    "loss": "損失",
+    "export_header": "LoRA をエクスポート",
+    "export_path": "エクスポートパス",
+    "export_lora_btn": "📦 LoRA をエクスポート",
+    "export_status": "エクスポート状態"
+  }
+}

acestep/gradio_ui/i18n/zh.json

@@ -0,0 +1,350 @@
+{
+  "app": {
+    "title": "🎛️ ACE-Step V1.5 演练场💡",
+    "subtitle": "推动开源音乐生成的边界"
+  },
+  "dataset": {
+    "title": "📊 数据集浏览器",
+    "dataset_label": "数据集",
+    "dataset_info": "选择要浏览的数据集",
+    "import_btn": "📥 导入数据集",
+    "search_type_label": "搜索类型",
+    "search_type_info": "如何查找项目",
+    "search_value_label": "搜索值",
+    "search_value_placeholder": "输入键或索引(留空表示随机)",
+    "search_value_info": "键: 精确匹配, 索引: 0到数据集大小-1",
+    "instruction_label": "📝 指令",
+    "instruction_placeholder": "无可用指令",
+    "metadata_title": "📋 项目元数据 (JSON)",
+    "metadata_label": "完整项目信息",
+    "source_audio": "源音频",
+    "target_audio": "目标音频",
+    "reference_audio": "参考音频",
+    "get_item_btn": "🔍 获取项目",
+    "use_src_checkbox": "使用数据集中的源音频",
+    "use_src_info": "勾选以使用数据集中的源音频",
+    "data_status_label": "📊 数据状态",
+    "data_status_default": "❌ 未导入数据集",
+    "autofill_btn": "📋 自动填充生成表单"
+  },
+  "service": {
+    "title": "🔧 服务配置",
+    "checkpoint_label": "检查点文件",
+    "checkpoint_info": "选择训练好的模型检查点文件(完整路径或文件名)",
+    "refresh_btn": "🔄 刷新",
+    "model_path_label": "主模型路径",
+    "model_path_info": "选择模型配置目录(从检查点自动扫描)",
+    "device_label": "设备",
+    "device_info": "处理设备(建议自动检测)",
+    "lm_model_path_label": "5Hz LM 模型路径",
+    "lm_model_path_info": "选择5Hz LM模型检查点(从检查点自动扫描)",
+    "backend_label": "5Hz LM 后端",
+    "backend_info": "选择5Hz LM的后端: vllm(更快)或pt(PyTorch, 更兼容)",
+    "init_llm_label": "初始化 5Hz LM",
+    "init_llm_info": "勾选以在服务初始化期间初始化5Hz LM",
+    "flash_attention_label": "使用Flash Attention",
+    "flash_attention_info_enabled": "启用flash attention以加快推理速度(需要flash_attn包)",
+    "flash_attention_info_disabled": "Flash attention不可用(未安装flash_attn包)",
+    "offload_cpu_label": "卸载到CPU",
+    "offload_cpu_info": "不使用时将模型卸载到CPU以节省GPU内存",
+    "offload_dit_cpu_label": "将DiT卸载到CPU",
+    "offload_dit_cpu_info": "将DiT卸载到CPU(需要启用卸载到CPU)",
+    "compile_model_label": "编译模型",
+    "compile_model_info": "使用 torch.compile 优化模型（量化必需）",
+    "quantization_label": "INT8 量化",
+    "quantization_info": "启用 INT8 仅权重量化以减少显存占用（需要启用编译模型）",
+    "init_btn": "初始化服务",
+    "status_label": "状态",
+    "language_label": "界面语言",
+    "language_info": "选择界面语言"
+  },
+  "generation": {
+    "required_inputs": "📝 必需输入",
+    "task_type_label": "任务类型",
+    "task_type_info": "选择生成的任务类型",
+    "instruction_label": "指令",
+    "instruction_info": "指令根据任务类型自动生成",
+    "load_btn": "加载",
+    "track_name_label": "音轨名称",
+    "track_name_info": "为lego/extract任务选择音轨名称",
+    "track_classes_label": "音轨名称",
+    "track_classes_info": "为complete任务选择多个音轨类别",
+    "audio_uploads": "🎵 音频上传",
+    "reference_audio": "参考音频(可选)",
+    "source_audio": "源音频(可选)",
+    "convert_codes_btn": "转换为代码",
+    "lm_codes_hints": "🎼 LM 代码提示",
+    "lm_codes_label": "LM 代码提示",
+    "lm_codes_placeholder": "<|audio_code_10695|><|audio_code_54246|>...",
+    "lm_codes_info": "粘贴用于text2music生成的LM代码提示",
+    "lm_codes_sample": "LM 代码提示(样本 {n})",
+    "lm_codes_sample_info": "样本{n}的代码",
+    "transcribe_btn": "转录",
+    "repainting_controls": "🎨 重绘控制(秒)",
+    "repainting_start": "重绘开始",
+    "repainting_end": "重绘结束",
+    "mode_label": "生成模式",
+    "mode_info": "简单模式：用自然语言描述音乐。自定义模式：完全控制描述和歌词。",
+    "mode_simple": "简单",
+    "mode_custom": "自定义",
+    "simple_query_label": "歌曲描述",
+    "simple_query_placeholder": "描述你想创作的音乐，例如：'给我生成一首暗黑的戏剧古风，歌词要华丽'。留空则随机生成样本。",
+    "simple_query_info": "输入你想生成的音乐的自然语言描述",
+    "simple_vocal_language_label": "人声语言(可选)",
+    "simple_vocal_language_info": "选择歌词的首选语言。使用 'unknown' 表示任意语言。",
+    "create_sample_btn": "创建样本",
+    "caption_title": "📝 音乐描述",
+    "caption_label": "音乐描述(可选)",
+    "caption_placeholder": "一段平和的原声吉他旋律,配有柔和的人声...",
+    "caption_info": "描述风格、流派、乐器���情绪",
+    "lyrics_title": "📝 歌词",
+    "lyrics_label": "歌词(可选)",
+    "lyrics_placeholder": "[第一段]\\n在星空下\\n我感到如此活跃...",
+    "lyrics_info": "带有结构的歌曲歌词",
+    "instrumental_label": "纯音乐",
+    "format_btn": "格式化",
+    "optional_params": "⚙️ 可选参数",
+    "vocal_language_label": "人声语言(可选)",
+    "vocal_language_info": "纯音乐使用 `unknown`",
+    "bpm_label": "BPM(可选)",
+    "bpm_info": "留空表示N/A",
+    "keyscale_label": "调性(可选)",
+    "keyscale_placeholder": "留空表示N/A",
+    "keyscale_info": "A-G, #/♭, 大调/小调",
+    "timesig_label": "拍号(可选)",
+    "timesig_info": "2/4, 3/4, 4/4...",
+    "duration_label": "音频时长(秒)",
+    "duration_info": "使用-1表示随机",
+    "batch_size_label": "批量大小",
+    "batch_size_info": "要生成的音频数量(最多8个)",
+    "advanced_settings": "🔧 高级设置",
+    "inference_steps_label": "DiT 推理步数",
+    "inference_steps_info": "Turbo: 最多8, Base: 最多200",
+    "guidance_scale_label": "DiT 引导比例(仅支持base模型)",
+    "guidance_scale_info": "更高的值更紧密地遵循文本",
+    "seed_label": "种子",
+    "seed_info": "批量使用逗号分隔的值",
+    "random_seed_label": "随机种子",
+    "random_seed_info": "启用以自动生成种子",
+    "audio_format_label": "音频格式",
+    "audio_format_info": "保存文件的音频格式",
+    "use_adg_label": "使用 ADG",
+    "use_adg_info": "启用角域引导",
+    "shift_label": "Shift",
+    "shift_info": "时间步偏移因子，仅对 base 模型生效 (范围 1.0~5.0，默认 3.0)。对 turbo 模型无效。",
+    "infer_method_label": "推理方法",
+    "infer_method_info": "扩散推理方法。ODE (欧拉) 更快，SDE (随机) 可能产生不同结果。",
+    "custom_timesteps_label": "自定义时间步",
+    "custom_timesteps_info": "可选：从 1.0 到 0.0 的逗号分隔值（例如 '0.97,0.76,0.615,0.5,0.395,0.28,0.18,0.085,0'）。会覆盖推理步数和 shift 设置。",
+    "cfg_interval_start": "CFG 间隔开始",
+    "cfg_interval_end": "CFG 间隔结束",
+    "lm_params_title": "🤖 LM 生成参数",
+    "lm_temperature_label": "LM 温度",
+    "lm_temperature_info": "5Hz LM温度(越高越随机)",
+    "lm_cfg_scale_label": "LM CFG 比例",
+    "lm_cfg_scale_info": "5Hz LM CFG (1.0 = 无CFG)",
+    "lm_top_k_label": "LM Top-K",
+    "lm_top_k_info": "Top-K (0 = 禁用)",
+    "lm_top_p_label": "LM Top-P",
+    "lm_top_p_info": "Top-P (1.0 = 禁用)",
+    "lm_negative_prompt_label": "LM 负面提示",
+    "lm_negative_prompt_placeholder": "输入CFG的负面提示(默认: NO USER INPUT)",
+    "lm_negative_prompt_info": "负面提示(当LM CFG比例 > 1.0时使用)",
+    "cot_metas_label": "CoT 元数据",
+    "cot_metas_info": "使用LM生成CoT元数据(取消勾选以跳过LM CoT生成)",
+    "cot_language_label": "CoT 语言",
+    "cot_language_info": "在CoT中生成语言(思维链)",
+    "constrained_debug_label": "约束解码调试",
+    "constrained_debug_info": "启用约束解码的调试日志(勾选以查看详细日志)",
+    "auto_score_label": "自动评分",
+    "auto_score_info": "自动计算所有生成音频的质量分数",
+    "auto_lrc_label": "自动 LRC",
+    "auto_lrc_info": "自动为所有生成的音频生成LRC歌词时间戳",
+    "lm_batch_chunk_label": "LM 批量块大小",
+    "lm_batch_chunk_info": "每个LM批量块的最大项目数(默认: 8, 受GPU内存限制)",
+    "codes_strength_label": "LM 代码强度",
+    "codes_strength_info": "控制使用LM生成代码的去噪步骤数量",
+    "similarity_denoise_label": "相似度 / 降噪",
+    "similarity_denoise_info": "控制输出与参考音频的贴合程度。数值越高保留越多结构。",
+    "cover_strength_label": "音频覆盖强度",
+    "cover_strength_info": "控制使用覆盖模式的去噪步骤数量",
+    "score_sensitivity_label": "质量评分敏感度",
+    "score_sensitivity_info": "更低 = 更敏感(默认: 1.0). 调整PMI如何映射到[0,1]",
+    "think_label": "思考",
+    "parallel_thinking_label": "并行思考",
+    "generate_btn": "🎵 生成音乐",
+    "autogen_label": "自动生成",
+    "caption_rewrite_label": "描述重写"
+  },
+  "results": {
+    "title": "🎵 结果",
+    "generated_music": "🎵 生成的音乐(样本 {n})",
+    "send_to_src_btn": "🔗 发送到源音频",
+    "save_btn": "💾 保存",
+    "score_btn": "📊 评分",
+    "lrc_btn": "🎵 LRC",
+    "quality_score_label": "质量分数(样本 {n})",
+    "quality_score_placeholder": "点击'评分'以计算基于困惑度的质量分数",
+    "codes_label": "LM 代码(样本 {n})",
+    "lrc_label": "歌词时间戳(样本 {n})",
+    "lrc_placeholder": "点击'LRC'生成时间戳",
+    "details_accordion": "📊 评分与LRC与LM代码",
+    "generation_status": "生成状态",
+    "current_batch": "当前批次",
+    "batch_indicator": "��次 {current} / {total}",
+    "next_batch_status": "下一批次状态",
+    "prev_btn": "◀ 上一个",
+    "next_btn": "下一个 ▶",
+    "restore_params_btn": "↙️ 将这些设置应用到UI(恢复批次参数)",
+    "batch_results_title": "👇 点击查看批量结果和生成详情",
+    "all_files_label": "📁 所有生成的文件(下载)",
+    "generation_details": "生成详情"
+  },
+  "messages": {
+    "no_audio_to_save": "❌ 没有要保存的音频",
+    "save_success": "✅ 已将音频和元数据保存到 {filename}",
+    "save_failed": "❌ 保存失败: {error}",
+    "no_file_selected": "⚠️ 未选择文件",
+    "params_loaded": "✅ 已从 {filename} 加载参数",
+    "invalid_json": "❌ 无效的JSON文件: {error}",
+    "load_error": "❌ 加载文件时出错: {error}",
+    "example_loaded": "📁 已从 {filename} 加载示例",
+    "example_failed": "解析JSON文件 {filename} 失败: {error}",
+    "example_error": "加载示例时出错: {error}",
+    "lm_generated": "🤖 使用LM生成的示例",
+    "lm_fallback": "使用LM生成示例失败,回退到示例目录",
+    "lm_not_initialized": "❌ 5Hz LM未初始化。请先初始化它。",
+    "autogen_enabled": "🔄 已启用自动生成 - 下一批次将在此之后生成",
+    "batch_ready": "✅ 批次 {n} 就绪!点击'下一个'查看。",
+    "batch_generating": "🔄 开始为批次 {n} 进行后台生成...",
+    "batch_failed": "❌ 后台生成失败: {error}",
+    "viewing_batch": "✅ 查看批次 {n}",
+    "at_first_batch": "已在第一批次",
+    "at_last_batch": "没有下一批次可用",
+    "batch_not_found": "在队列中未找到批次 {n}",
+    "no_batch_data": "没有要恢复的批次数据。",
+    "params_restored": "✅ 已从批次 {n} 恢复UI参数",
+    "scoring_failed": "❌ 错误: 未找到批次数据",
+    "no_codes": "❌ 没有可用的音频代码。请先生成音乐。",
+    "score_failed": "❌ 评分失败: {error}",
+    "score_error": "❌ 计算分数时出错: {error}",
+    "lrc_no_batch_data": "❌ 未找到批次数据。请先生成音乐。",
+    "lrc_no_extra_outputs": "❌ 未找到额外输出。条件张量不可用。",
+    "lrc_missing_tensors": "❌ 缺少LRC生成所需的张量。",
+    "lrc_sample_not_exist": "❌ 当前批次中不存在该样本。",
+    "lrc_empty_result": "⚠️ LRC生成结果为空。",
+    "empty_query": "⚠️ 请输入音乐描述。",
+    "sample_creation_failed": "❌ 创建样本失败。请重试。",
+    "sample_created": "✅ 样本已创建！检查描述和歌词，然后点击生成音乐。",
+    "simple_examples_not_found": "⚠️ 未找到简单模式示例目录。",
+    "simple_examples_empty": "⚠️ 简单模式示例中没有示例文件。",
+    "simple_example_loaded": "🎲 已从 {filename} 加载随机示例",
+    "format_success": "✅ 描述和歌词格式化成功",
+    "format_failed": "❌ 格式化失败: {error}",
+    "skipping_metas_cot": "⚡ 跳过 Phase 1 元数据 COT（样本已格式化）",
+    "invalid_timesteps_format": "⚠️ 时间步格式无效，使用默认调度。",
+    "timesteps_out_of_range": "⚠️ 时间步必须在 [0, 1] 范围内，使用默认调度。",
+    "timesteps_count_mismatch": "⚠️ 时间步数量 ({actual}) 与推理步数 ({expected}) 不匹配，将使用时间步数量。"
+  },
+  "training": {
+    "tab_title": "🎓 LoRA 训练",
+    "tab_dataset_builder": "📁 数据集构建",
+    "tab_train_lora": "🚀 训练 LoRA",
+    "quick_start_title": "🚀 快速开始",
+    "load_dataset_label": "数据集 JSON 路径",
+    "load_btn": "📂 加载",
+    "load_status": "加载状态",
+    "scan_label": "音频目录路径",
+    "scan_info": "扫描音频文件（wav、mp3、flac、ogg、opus）",
+    "scan_btn": "🔍 扫描",
+    "scan_status": "扫描状态",
+    "found_audio_files": "已找到的音频文件",
+    "dataset_name": "数据集名称",
+    "dataset_name_placeholder": "输入数据集名称",
+    "dataset_settings_header": "数据集设置",
+    "tag_prepend": "前置（标签，描述）",
+    "tag_append": "后置（描述，标签）",
+    "tag_replace": "替换描述",
+    "step2_title": "步骤 2：AI 自动标注",
+    "step3_title": "步骤 3：预览与编辑",
+    "step4_title": "步骤 4：保存数据集",
+    "step5_title": "步骤 5：预处理为张量",
+    "all_instrumental": "全部为纯音乐",
+    "all_instrumental_info": "勾选表示所有曲目均为纯音乐（无人声）",
+    "custom_tag": "自定义激活标签",
+    "custom_tag_info": "用于激活此 LoRA 风格的唯一标签",
+    "tag_position": "标签位置",
+    "tag_position_info": "在描述中放置自定义标签的位置",
+    "genre_ratio": "风格比例 (%)",
+    "genre_ratio_info": "0%=全部描述，100%=全部风格。单样本覆盖优先。",
+    "skip_metas": "跳过 BPM/调性/拍号",
+    "skip_metas_info": "跳过 BPM/调性/拍号生成。描述和风格仍由 LM 生成。",
+    "only_unlabeled": "仅未标注",
+    "only_unlabeled_info": "仅标注无描述的样本（用于继续失败的标注）",
+    "auto_label_btn": "🏷️ 自动标注全部",
+    "label_progress": "标注进度",
+    "select_sample": "选择样本 #",
+    "select_sample_info": "选择要预览和编辑的样本",
+    "audio_preview": "音频预览",
+    "filename": "文件名",
+    "caption": "描述",
+    "genre": "风格",
+    "prompt_override_label": "提示覆盖（本样本）",
+    "prompt_override_info": "覆盖本样本的全局比例",
+    "lyrics_editable_label": "歌词（可编辑，用于训练）",
+    "raw_lyrics_label": "原始歌词（来自 .txt 文件）",
+    "no_lyrics_placeholder": "（无 .txt 歌词文件）",
+    "bpm": "BPM",
+    "key_label": "调性",
+    "key_placeholder": "C 大调",
+    "time_sig": "拍号",
+    "duration_s": "时长 (秒)",
+    "language": "语言",
+    "instrumental": "纯音乐",
+    "save_changes_btn": "💾 保存更改",
+    "edit_status": "编辑状态",
+    "save_path": "保存路径",
+    "save_path_info": "数据集 JSON 的保存路径",
+    "save_dataset_btn": "💾 保存数据集",
+    "save_status": "保存状态",
+    "load_existing_label": "加载已有数据集（可选）",
+    "load_existing_info": "之前保存的数据集 JSON 文件路径",
+    "load_dataset_btn": "📂 加载数据集",
+    "tensor_output_dir": "张量输出目录",
+    "tensor_output_info": "保存预处理张量文件的目录",
+    "preprocess_btn": "⚡ 预处理",
+    "preprocess_progress": "预处理进度",
+    "preprocessed_tensors_dir": "预处理张量目录",
+    "preprocessed_tensors_info": "包含预处理 .pt 张量文件的目录",
+    "dataset_info": "数据集信息",
+    "lora_rank": "LoRA 秩 (r)",
+    "lora_rank_info": "越高容量越大，显存占用越多",
+    "lora_alpha": "LoRA Alpha",
+    "lora_alpha_info": "缩放因子（通常为 2× 秩）",
+    "lora_dropout": "LoRA Dropout",
+    "learning_rate": "学习率",
+    "learning_rate_info": "建议从 3e-4 开始，按需调整",
+    "max_epochs": "最大轮数",
+    "batch_size": "批大小",
+    "batch_size_info": "显存充足时可增大",
+    "gradient_accumulation": "梯度累积",
+    "gradient_accumulation_info": "有效批大小 = batch_size × 累积步数",
+    "save_every_n_epochs": "每 N 轮保存",
+    "shift": "Shift",
+    "shift_info": "Turbo 模型时间步偏移",
+    "seed": "随机种子",
+    "output_dir": "输出目录",
+    "output_dir_info": "保存训练后 LoRA 权重的目录",
+    "start_training_btn": "🚀 开始训练",
+    "stop_training_btn": "⏹️ 停止训练",
+    "training_progress": "训练进度",
+    "training_log": "训练日志",
+    "training_loss_title": "训练损失",
+    "step": "步数",
+    "loss": "损失",
+    "export_header": "导出 LoRA",
+    "export_path": "导出路径",
+    "export_lora_btn": "📦 导出 LoRA",
+    "export_status": "导出状态"
+  }
+}

acestep/gradio_ui/interfaces/__init__.py

@@ -0,0 +1,94 @@
+"""
+Gradio UI Components Module
+Contains all Gradio interface component definitions and layouts
+"""
+import gradio as gr
+from acestep.gradio_ui.i18n import get_i18n, t
+from acestep.gradio_ui.interfaces.dataset import create_dataset_section
+from acestep.gradio_ui.interfaces.generation import create_generation_section
+from acestep.gradio_ui.interfaces.result import create_results_section
+from acestep.gradio_ui.interfaces.training import create_training_section
+from acestep.gradio_ui.events import setup_event_handlers, setup_training_event_handlers
+
+
+def create_gradio_interface(dit_handler, llm_handler, dataset_handler, init_params=None, language='en') -> gr.Blocks:
+    """
+    Create Gradio interface
+    
+    Args:
+        dit_handler: DiT handler instance
+        llm_handler: LM handler instance
+        dataset_handler: Dataset handler instance
+        init_params: Dictionary containing initialization parameters and state.
+                    If None, service will not be pre-initialized.
+        language: UI language code ('en', 'zh', 'ja', default: 'en')
+        
+    Returns:
+        Gradio Blocks instance
+    """
+    # Initialize i18n with selected language
+    i18n = get_i18n(language)
+    
+    with gr.Blocks(
+        title=t("app.title"),
+        theme=gr.themes.Soft(),
+        css="""
+        .main-header {
+            text-align: center;
+            margin-bottom: 2rem;
+        }
+        .section-header {
+            background: linear-gradient(90deg, #4CAF50, #45a049);
+            color: white;
+            padding: 10px;
+            border-radius: 5px;
+            margin: 10px 0;
+        }
+        .lm-hints-row {
+            align-items: stretch;
+        }
+        .lm-hints-col {
+            display: flex;
+        }
+        .lm-hints-col > div {
+            flex: 1;
+            display: flex;
+        }
+        .lm-hints-btn button {
+            height: 100%;
+            width: 100%;
+        }
+        /* Position Audio time labels lower to avoid scrollbar overlap */
+        .component-wrapper > .timestamps {
+            transform: translateY(15px);
+        }
+        """,
+    ) as demo:
+        
+        gr.HTML(f"""
+        <div class="main-header">
+            <h1>{t("app.title")}</h1>
+            <p>{t("app.subtitle")}</p>
+        </div>
+        """)
+        
+        # Dataset Explorer Section
+        dataset_section = create_dataset_section(dataset_handler)
+        
+        # Generation Section (pass init_params and language to support pre-initialization)
+        generation_section = create_generation_section(dit_handler, llm_handler, init_params=init_params, language=language)
+        
+        # Results Section
+        results_section = create_results_section(dit_handler)
+        
+        # Training Section (LoRA training and dataset builder)
+        # Pass init_params to support hiding in service mode
+        training_section = create_training_section(dit_handler, llm_handler, init_params=init_params)
+        
+        # Connect event handlers
+        setup_event_handlers(demo, dit_handler, llm_handler, dataset_handler, dataset_section, generation_section, results_section)
+        
+        # Connect training event handlers
+        setup_training_event_handlers(demo, dit_handler, llm_handler, training_section)
+    
+    return demo

acestep/gradio_ui/interfaces/dataset.py

@@ -0,0 +1,101 @@
+"""
+Gradio UI Dataset Section Module
+Contains dataset explorer section component definitions
+"""
+import gradio as gr
+
+
+def create_dataset_section(dataset_handler) -> dict:
+    """Create dataset explorer section"""
+    with gr.Accordion("📊 Dataset Explorer", open=False, visible=False):
+        with gr.Row(equal_height=True):
+            dataset_type = gr.Dropdown(
+                choices=["train", "test"],
+                value="train",
+                label="Dataset",
+                info="Choose dataset to explore",
+                scale=2
+            )
+            import_dataset_btn = gr.Button("📥 Import Dataset", variant="primary", scale=1)
+            
+            search_type = gr.Dropdown(
+                choices=["keys", "idx", "random"],
+                value="random",
+                label="Search Type",
+                info="How to find items",
+                scale=1
+            )
+            search_value = gr.Textbox(
+                label="Search Value",
+                placeholder="Enter keys or index (leave empty for random)",
+                info="Keys: exact match, Index: 0 to dataset size-1",
+                scale=2
+            )
+
+        instruction_display = gr.Textbox(
+            label="📝 Instruction",
+            interactive=False,
+            placeholder="No instruction available",
+            lines=1
+        )
+        
+        repaint_viz_plot = gr.Plot()
+        
+        with gr.Accordion("📋 Item Metadata (JSON)", open=False):
+            item_info_json = gr.Code(
+                label="Complete Item Information",
+                language="json",
+                interactive=False,
+                lines=15
+            )
+        
+        with gr.Row(equal_height=True):
+            item_src_audio = gr.Audio(
+                label="Source Audio",
+                type="filepath",
+                interactive=False,
+                scale=8
+            )
+            get_item_btn = gr.Button("🔍 Get Item", variant="secondary", interactive=False, scale=2)
+        
+        with gr.Row(equal_height=True):
+            item_target_audio = gr.Audio(
+                label="Target Audio",
+                type="filepath",
+                interactive=False,
+                scale=8
+            )
+            item_refer_audio = gr.Audio(
+                label="Reference Audio",
+                type="filepath",
+                interactive=False,
+                scale=2
+            )
+        
+        with gr.Row():
+            use_src_checkbox = gr.Checkbox(
+                label="Use Source Audio from Dataset",
+                value=True,
+                info="Check to use the source audio from dataset"
+            )
+
+        data_status = gr.Textbox(label="📊 Data Status", interactive=False, value="❌ No dataset imported")
+        auto_fill_btn = gr.Button("📋 Auto-fill Generation Form", variant="primary")
+    
+    return {
+        "dataset_type": dataset_type,
+        "import_dataset_btn": import_dataset_btn,
+        "search_type": search_type,
+        "search_value": search_value,
+        "instruction_display": instruction_display,
+        "repaint_viz_plot": repaint_viz_plot,
+        "item_info_json": item_info_json,
+        "item_src_audio": item_src_audio,
+        "get_item_btn": get_item_btn,
+        "item_target_audio": item_target_audio,
+        "item_refer_audio": item_refer_audio,
+        "use_src_checkbox": use_src_checkbox,
+        "data_status": data_status,
+        "auto_fill_btn": auto_fill_btn,
+    }
+

acestep/gradio_ui/interfaces/generation.py

@@ -0,0 +1,824 @@
+"""
+Gradio UI Generation Section Module
+Contains generation section component definitions
+"""
+import sys
+import gradio as gr
+from acestep.constants import (
+    VALID_LANGUAGES,
+    TRACK_NAMES,
+    TASK_TYPES_TURBO,
+    TASK_TYPES_BASE,
+    DEFAULT_DIT_INSTRUCTION,
+)
+from acestep.gradio_ui.i18n import t
+from acestep.gpu_config import get_global_gpu_config, GPUConfig
+
+
+def create_generation_section(dit_handler, llm_handler, init_params=None, language='en') -> dict:
+    """Create generation section
+    
+    Args:
+        dit_handler: DiT handler instance
+        llm_handler: LM handler instance
+        init_params: Dictionary containing initialization parameters and state.
+                    If None, service will not be pre-initialized.
+        language: UI language code ('en', 'zh', 'ja')
+    """
+    # Check if service is pre-initialized
+    service_pre_initialized = init_params is not None and init_params.get('pre_initialized', False)
+    
+    # Check if running in service mode (restricted UI)
+    service_mode = init_params is not None and init_params.get('service_mode', False)
+    
+    # Get current language from init_params if available
+    current_language = init_params.get('language', language) if init_params else language
+    
+    # Get GPU configuration
+    gpu_config: GPUConfig = init_params.get('gpu_config') if init_params else None
+    if gpu_config is None:
+        gpu_config = get_global_gpu_config()
+    
+    # Determine if LM is initialized (for setting appropriate limits)
+    lm_initialized = init_params.get('init_llm', False) if init_params else False
+    
+    # Calculate UI limits based on GPU config and LM state
+    max_duration = gpu_config.max_duration_with_lm if lm_initialized else gpu_config.max_duration_without_lm
+    max_batch_size = gpu_config.max_batch_size_with_lm if lm_initialized else gpu_config.max_batch_size_without_lm
+    default_batch_size = min(2, max_batch_size)  # Default to 2 or max if lower
+    init_lm_default = gpu_config.init_lm_default
+    
+    # Determine default offload setting
+    # If XPU is detected, default offload to False (keep models on device)
+    # Otherwise default to True (offload to CPU to save VRAM)
+    default_offload = True
+    try:
+        import torch
+        if hasattr(torch, 'xpu') and torch.xpu.is_available():
+            default_offload = False
+    except ImportError:
+        pass
+    
+    with gr.Group():
+        # Service Configuration - collapse if pre-initialized, hide if in service mode
+        accordion_open = not service_pre_initialized
+        accordion_visible = not service_pre_initialized  # Hide when running in service mode
+        with gr.Accordion(t("service.title"), open=accordion_open, visible=accordion_visible) as service_config_accordion:
+            # Language selector at the top
+            with gr.Row():
+                language_dropdown = gr.Dropdown(
+                    choices=[
+                        ("English", "en"),
+                        ("中文", "zh"),
+                        ("日本語", "ja"),
+                    ],
+                    value=current_language,
+                    label=t("service.language_label"),
+                    info=t("service.language_info"),
+                    scale=1,
+                )
+            
+            # Dropdown options section - all dropdowns grouped together
+            with gr.Row(equal_height=True):
+                with gr.Column(scale=4):
+                    # Set checkpoint value from init_params if pre-initialized
+                    checkpoint_value = init_params.get('checkpoint') if service_pre_initialized else None
+                    checkpoint_dropdown = gr.Dropdown(
+                        label=t("service.checkpoint_label"),
+                        choices=dit_handler.get_available_checkpoints(),
+                        value=checkpoint_value,
+                        info=t("service.checkpoint_info")
+                    )
+                with gr.Column(scale=1, min_width=90):
+                    refresh_btn = gr.Button(t("service.refresh_btn"), size="sm")
+            
+            with gr.Row():
+                # Get available acestep-v15- model list
+                available_models = dit_handler.get_available_acestep_v15_models()
+                default_model = "acestep-v15-turbo" if "acestep-v15-turbo" in available_models else (available_models[0] if available_models else None)
+                
+                # Set config_path value from init_params if pre-initialized
+                config_path_value = init_params.get('config_path', default_model) if service_pre_initialized else default_model
+                config_path = gr.Dropdown(
+                    label=t("service.model_path_label"),
+                    choices=available_models,
+                    value=config_path_value,
+                    info=t("service.model_path_info")
+                )
+                # Set device value from init_params if pre-initialized
+                device_value = init_params.get('device', 'auto') if service_pre_initialized else 'auto'
+                device = gr.Dropdown(
+                    choices=["auto", "cuda", "mps", "xpu", "cpu"],
+                    value=device_value,
+                    label=t("service.device_label"),
+                    info=t("service.device_info")
+                )
+            
+            with gr.Row():
+                # Get available 5Hz LM model list
+                available_lm_models = llm_handler.get_available_5hz_lm_models()
+                default_lm_model = "acestep-5Hz-lm-0.6B" if "acestep-5Hz-lm-0.6B" in available_lm_models else (available_lm_models[0] if available_lm_models else None)
+                
+                # Set lm_model_path value from init_params if pre-initialized
+                lm_model_path_value = init_params.get('lm_model_path', default_lm_model) if service_pre_initialized else default_lm_model
+                lm_model_path = gr.Dropdown(
+                    label=t("service.lm_model_path_label"),
+                    choices=available_lm_models,
+                    value=lm_model_path_value,
+                    info=t("service.lm_model_path_info")
+                )
+                # Set backend value from init_params if pre-initialized
+                backend_value = init_params.get('backend', 'vllm') if service_pre_initialized else 'vllm'
+                backend_dropdown = gr.Dropdown(
+                    choices=["vllm", "pt", "mlx"],
+                    value=backend_value,
+                    label=t("service.backend_label"),
+                    info=t("service.backend_info")
+                )
+            
+            # Checkbox options section - all checkboxes grouped together
+            with gr.Row():
+                # Set init_llm value from init_params if pre-initialized, otherwise use GPU config default
+                init_llm_value = init_params.get('init_llm', init_lm_default) if service_pre_initialized else init_lm_default
+                init_llm_checkbox = gr.Checkbox(
+                    label=t("service.init_llm_label"),
+                    value=init_llm_value,
+                    info=t("service.init_llm_info"),
+                )
+                # Auto-detect flash attention availability
+                flash_attn_available = dit_handler.is_flash_attention_available(device_value)
+                # Set use_flash_attention value from init_params if pre-initialized
+                use_flash_attention_value = init_params.get('use_flash_attention', flash_attn_available) if service_pre_initialized else flash_attn_available
+                use_flash_attention_checkbox = gr.Checkbox(
+                    label=t("service.flash_attention_label"),
+                    value=use_flash_attention_value,
+                    interactive=flash_attn_available,
+                    info=t("service.flash_attention_info_enabled") if flash_attn_available else t("service.flash_attention_info_disabled")
+                )
+                # Set offload_to_cpu value from init_params if pre-initialized (default True)
+                offload_to_cpu_value = init_params.get('offload_to_cpu', default_offload) if service_pre_initialized else default_offload
+                offload_to_cpu_checkbox = gr.Checkbox(
+                    label=t("service.offload_cpu_label"),
+                    value=offload_to_cpu_value,
+                    info=t("service.offload_cpu_info")
+                )
+                # Set offload_dit_to_cpu value from init_params if pre-initialized (default True)
+                offload_dit_to_cpu_value = init_params.get('offload_dit_to_cpu', default_offload) if service_pre_initialized else default_offload
+                offload_dit_to_cpu_checkbox = gr.Checkbox(
+                    label=t("service.offload_dit_cpu_label"),
+                    value=offload_dit_to_cpu_value,
+                    info=t("service.offload_dit_cpu_info")
+                )
+                # Set compile_model value from init_params if pre-initialized (default True)
+                compile_model_value = init_params.get('compile_model', True) if service_pre_initialized else True
+                compile_model_checkbox = gr.Checkbox(
+                    label=t("service.compile_model_label"),
+                    value=compile_model_value,
+                    info=t("service.compile_model_info")
+                )
+                # Set quantization value from init_params if pre-initialized.
+                # Default to False on macOS to avoid torchao incompatibilities.
+                default_quantization = False if sys.platform == "darwin" else True
+                quantization_value = init_params.get('quantization', default_quantization) if service_pre_initialized else default_quantization
+                quantization_checkbox = gr.Checkbox(
+                    label=t("service.quantization_label"),
+                    value=quantization_value,
+                    info=t("service.quantization_info")
+                )
+            
+            init_btn = gr.Button(t("service.init_btn"), variant="primary", size="lg")
+            # Set init_status value from init_params if pre-initialized
+            init_status_value = init_params.get('init_status', '') if service_pre_initialized else ''
+            init_status = gr.Textbox(label=t("service.status_label"), interactive=False, lines=3, value=init_status_value)
+            
+            # LoRA Configuration Section
+            gr.HTML("<hr><h4>🔧 LoRA Adapter</h4>")
+            with gr.Row():
+                lora_path = gr.Textbox(
+                    label="LoRA Path",
+                    placeholder="./lora_output/final/adapter",
+                    info="Path to trained LoRA adapter directory",
+                    scale=3,
+                )
+                load_lora_btn = gr.Button("📥 Load LoRA", variant="secondary", scale=1)
+                unload_lora_btn = gr.Button("🗑️ Unload", variant="secondary", scale=1)
+            with gr.Row():
+                use_lora_checkbox = gr.Checkbox(
+                    label="Use LoRA",
+                    value=False,
+                    info="Enable LoRA adapter for inference",
+                    scale=1,
+                )
+                lora_scale_slider = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=1.0,
+                    step=0.05,
+                    label="LoRA Scale",
+                    info="LoRA influence strength (0=disabled, 1=full)",
+                    scale=2,
+                )
+                lora_status = gr.Textbox(
+                    label="LoRA Status",
+                    value="No LoRA loaded",
+                    interactive=False,
+                    scale=2,
+                )
+        
+        # Inputs
+        with gr.Row():
+            with gr.Column(scale=2):
+                with gr.Accordion(t("generation.required_inputs"), open=True):
+                    # Task type
+                    # Determine initial task_type choices based on actual model in use
+                    # When service is pre-initialized, use config_path from init_params
+                    actual_model = init_params.get('config_path', default_model) if service_pre_initialized else default_model
+                    actual_model_lower = (actual_model or "").lower()
+                    if "turbo" in actual_model_lower:
+                        initial_task_choices = TASK_TYPES_TURBO
+                    else:
+                        initial_task_choices = TASK_TYPES_BASE
+                    
+                    with gr.Row(equal_height=True):
+                        with gr.Column(scale=2):
+                            task_type = gr.Dropdown(
+                                choices=initial_task_choices,
+                                value="text2music",
+                                label=t("generation.task_type_label"),
+                                info=t("generation.task_type_info"),
+                            )
+                        with gr.Column(scale=7):
+                            instruction_display_gen = gr.Textbox(
+                                label=t("generation.instruction_label"),
+                                value=DEFAULT_DIT_INSTRUCTION,
+                                interactive=False,
+                                lines=1,
+                                info=t("generation.instruction_info"),
+                            )
+                        with gr.Column(scale=1, min_width=100):
+                            load_file = gr.UploadButton(
+                                t("generation.load_btn"),
+                                file_types=[".json"],
+                                file_count="single",
+                                variant="secondary",
+                                size="sm",
+                            )
+                    
+                    track_name = gr.Dropdown(
+                        choices=TRACK_NAMES,
+                        value=None,
+                        label=t("generation.track_name_label"),
+                        info=t("generation.track_name_info"),
+                        visible=False
+                    )
+                    
+                    complete_track_classes = gr.CheckboxGroup(
+                        choices=TRACK_NAMES,
+                        label=t("generation.track_classes_label"),
+                        info=t("generation.track_classes_info"),
+                        visible=False
+                    )
+                    
+                    # Audio uploads
+                    audio_uploads_accordion = gr.Accordion(t("generation.audio_uploads"), open=False)
+                    with audio_uploads_accordion:
+                        with gr.Row(equal_height=True):
+                            with gr.Column(scale=2):
+                                reference_audio = gr.Audio(
+                                    label=t("generation.reference_audio"),
+                                    type="filepath",
+                                )
+                            with gr.Column(scale=7):
+                                src_audio = gr.Audio(
+                                    label=t("generation.source_audio"),
+                                    type="filepath",
+                                )
+                            with gr.Column(scale=1, min_width=80):
+                                convert_src_to_codes_btn = gr.Button(
+                                    t("generation.convert_codes_btn"),
+                                    variant="secondary",
+                                    size="sm"
+                                )
+                        
+                    # Audio Codes for text2music - single input for transcription or cover task
+                    with gr.Accordion(t("generation.lm_codes_hints"), open=False, visible=True) as text2music_audio_codes_group:
+                        with gr.Row(equal_height=True):
+                            text2music_audio_code_string = gr.Textbox(
+                                label=t("generation.lm_codes_label"),
+                                placeholder=t("generation.lm_codes_placeholder"),
+                                lines=6,
+                                info=t("generation.lm_codes_info"),
+                                scale=9,
+                            )
+                            transcribe_btn = gr.Button(
+                                t("generation.transcribe_btn"),
+                                variant="secondary",
+                                size="sm",
+                                scale=1,
+                            )
+                    
+                    # Repainting controls
+                    with gr.Group(visible=False) as repainting_group:
+                        gr.HTML(f"<h5>{t('generation.repainting_controls')}</h5>")
+                        with gr.Row():
+                            repainting_start = gr.Number(
+                                label=t("generation.repainting_start"),
+                                value=0.0,
+                                step=0.1,
+                            )
+                            repainting_end = gr.Number(
+                                label=t("generation.repainting_end"),
+                                value=-1,
+                                minimum=-1,
+                                step=0.1,
+                            )
+                    
+                    # Simple/Custom Mode Toggle
+                    # In service mode: only Custom mode, hide the toggle
+                    with gr.Row(visible=not service_mode):
+                        generation_mode = gr.Radio(
+                            choices=[
+                                (t("generation.mode_simple"), "simple"),
+                                (t("generation.mode_custom"), "custom"),
+                            ],
+                            value="custom" if service_mode else "simple",
+                            label=t("generation.mode_label"),
+                            info=t("generation.mode_info"),
+                        )
+                    
+                    # Simple Mode Components - hidden in service mode
+                    with gr.Group(visible=not service_mode) as simple_mode_group:
+                        with gr.Row(equal_height=True):
+                            simple_query_input = gr.Textbox(
+                                label=t("generation.simple_query_label"),
+                                placeholder=t("generation.simple_query_placeholder"),
+                                lines=2,
+                                info=t("generation.simple_query_info"),
+                                scale=12,
+                            )
+
+                            with gr.Column(scale=1, min_width=100):
+                                random_desc_btn = gr.Button(
+                                    "🎲",
+                                    variant="secondary",
+                                    size="sm",
+                                    scale=2
+                                )
+                        
+                        with gr.Row(equal_height=True):
+                            with gr.Column(scale=1, variant="compact"):
+                                simple_instrumental_checkbox = gr.Checkbox(
+                                    label=t("generation.instrumental_label"),
+                                    value=False,
+                                )
+                            with gr.Column(scale=18):
+                                create_sample_btn = gr.Button(
+                                    t("generation.create_sample_btn"),
+                                    variant="primary",
+                                    size="lg",
+                                )
+                            with gr.Column(scale=1, variant="compact"):
+                                simple_vocal_language = gr.Dropdown(
+                                    choices=VALID_LANGUAGES,
+                                    value="unknown",
+                                    allow_custom_value=True,
+                                    label=t("generation.simple_vocal_language_label"),
+                                    interactive=True,
+                                )
+                    
+                    # State to track if sample has been created in Simple mode
+                    simple_sample_created = gr.State(value=False)
+                
+                # Music Caption - wrapped in accordion that can be collapsed in Simple mode
+                # Default to expanded for better UX
+                with gr.Accordion(t("generation.caption_title"), open=True) as caption_accordion:
+                    with gr.Row(equal_height=True):
+                        captions = gr.Textbox(
+                            label=t("generation.caption_label"),
+                            placeholder=t("generation.caption_placeholder"),
+                            lines=3,
+                            info=t("generation.caption_info"),
+                            scale=12,
+                        )
+                        with gr.Column(scale=1, min_width=100):
+                            sample_btn = gr.Button(
+                                "🎲",
+                                variant="secondary",
+                                size="sm",
+                                scale=2,
+                            )
+                # Lyrics - wrapped in accordion that can be collapsed in Simple mode
+                # Default to expanded for better UX
+                with gr.Accordion(t("generation.lyrics_title"), open=True) as lyrics_accordion:
+                    lyrics = gr.Textbox(
+                        label=t("generation.lyrics_label"),
+                        placeholder=t("generation.lyrics_placeholder"),
+                        lines=8,
+                        info=t("generation.lyrics_info")
+                    )
+                    
+                    with gr.Row(variant="compact", equal_height=True):
+                        instrumental_checkbox = gr.Checkbox(
+                            label=t("generation.instrumental_label"),
+                            value=False,
+                            scale=1,
+                            min_width=120,
+                            container=True,
+                        )
+                        
+                        # 中间：语言选择 (Dropdown)
+                        # 移除 gr.HTML hack，直接使用 label 参数，Gradio 会自动处理对齐
+                        vocal_language = gr.Dropdown(
+                            choices=VALID_LANGUAGES,
+                            value="unknown",
+                            label=t("generation.vocal_language_label"),
+                            show_label=False, 
+                            container=True, 
+                            allow_custom_value=True,
+                            scale=3,
+                        )
+                        
+                        # 右侧：格式化按钮 (Button)
+                        # 放在同一行最右侧，操作更顺手
+                        format_btn = gr.Button(
+                            t("generation.format_btn"),
+                            variant="secondary",
+                            scale=1,
+                            min_width=80,
+                        )
+                
+                # Optional Parameters
+                # In service mode: auto-expand
+                with gr.Accordion(t("generation.optional_params"), open=service_mode) as optional_params_accordion:
+                    with gr.Row():
+                        bpm = gr.Number(
+                            label=t("generation.bpm_label"),
+                            value=None,
+                            step=1,
+                            info=t("generation.bpm_info")
+                        )
+                        key_scale = gr.Textbox(
+                            label=t("generation.keyscale_label"),
+                            placeholder=t("generation.keyscale_placeholder"),
+                            value="",
+                            info=t("generation.keyscale_info")
+                        )
+                        time_signature = gr.Dropdown(
+                            choices=["", "2", "3", "4", "6", "N/A"],
+                            value="",
+                            label=t("generation.timesig_label"),
+                            allow_custom_value=True,
+                            info=t("generation.timesig_info")
+                        )
+                        audio_duration = gr.Number(
+                            label=t("generation.duration_label"),
+                            value=-1,
+                            minimum=-1,
+                            maximum=float(max_duration),
+                            step=0.1,
+                            info=t("generation.duration_info") + f" (Max: {max_duration}s / {max_duration // 60} min)"
+                        )
+                        batch_size_input = gr.Number(
+                            label=t("generation.batch_size_label"),
+                            value=default_batch_size,
+                            minimum=1,
+                            maximum=max_batch_size,
+                            step=1,
+                            info=t("generation.batch_size_info") + f" (Max: {max_batch_size})",
+                            interactive=not service_mode  # Fixed in service mode
+                        )
+        
+        # Advanced Settings
+        # Default UI settings use turbo mode (max 20 steps, default 8, show shift with default 3)
+        # These will be updated after model initialization based on handler.is_turbo_model()
+        with gr.Accordion(t("generation.advanced_settings"), open=False):
+            with gr.Row():
+                inference_steps = gr.Slider(
+                    minimum=1,
+                    maximum=20,
+                    value=8,
+                    step=1,
+                    label=t("generation.inference_steps_label"),
+                    info=t("generation.inference_steps_info")
+                )
+                guidance_scale = gr.Slider(
+                    minimum=1.0,
+                    maximum=15.0,
+                    value=7.0,
+                    step=0.1,
+                    label=t("generation.guidance_scale_label"),
+                    info=t("generation.guidance_scale_info"),
+                    visible=False
+                )
+                with gr.Column():
+                    seed = gr.Textbox(
+                        label=t("generation.seed_label"),
+                        value="-1",
+                        info=t("generation.seed_info")
+                    )
+                    random_seed_checkbox = gr.Checkbox(
+                        label=t("generation.random_seed_label"),
+                        value=True,
+                        info=t("generation.random_seed_info")
+                    )
+                audio_format = gr.Dropdown(
+                    choices=["mp3", "flac"],
+                    value="mp3",
+                    label=t("generation.audio_format_label"),
+                    info=t("generation.audio_format_info"),
+                    interactive=not service_mode  # Fixed in service mode
+                )
+            
+            with gr.Row():
+                use_adg = gr.Checkbox(
+                    label=t("generation.use_adg_label"),
+                    value=False,
+                    info=t("generation.use_adg_info"),
+                    visible=False
+                )
+                shift = gr.Slider(
+                    minimum=1.0,
+                    maximum=5.0,
+                    value=3.0,
+                    step=0.1,
+                    label=t("generation.shift_label"),
+                    info=t("generation.shift_info"),
+                    visible=True
+                )
+                infer_method = gr.Dropdown(
+                    choices=["ode", "sde"],
+                    value="ode",
+                    label=t("generation.infer_method_label"),
+                    info=t("generation.infer_method_info"),
+                )
+            
+            with gr.Row():
+                custom_timesteps = gr.Textbox(
+                    label=t("generation.custom_timesteps_label"),
+                    placeholder="0.97,0.76,0.615,0.5,0.395,0.28,0.18,0.085,0",
+                    value="",
+                    info=t("generation.custom_timesteps_info"),
+                )
+            
+            with gr.Row():
+                cfg_interval_start = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.0,
+                    step=0.01,
+                    label=t("generation.cfg_interval_start"),
+                    visible=False
+                )
+                cfg_interval_end = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=1.0,
+                    step=0.01,
+                    label=t("generation.cfg_interval_end"),
+                    visible=False
+                )
+
+            # LM (Language Model) Parameters
+            gr.HTML(f"<h4>{t('generation.lm_params_title')}</h4>")
+            with gr.Row():
+                lm_temperature = gr.Slider(
+                    label=t("generation.lm_temperature_label"),
+                    minimum=0.0,
+                    maximum=2.0,
+                    value=0.85,
+                    step=0.1,
+                    scale=1,
+                    info=t("generation.lm_temperature_info")
+                )
+                lm_cfg_scale = gr.Slider(
+                    label=t("generation.lm_cfg_scale_label"),
+                    minimum=1.0,
+                    maximum=3.0,
+                    value=2.0,
+                    step=0.1,
+                    scale=1,
+                    info=t("generation.lm_cfg_scale_info")
+                )
+                lm_top_k = gr.Slider(
+                    label=t("generation.lm_top_k_label"),
+                    minimum=0,
+                    maximum=100,
+                    value=0,
+                    step=1,
+                    scale=1,
+                    info=t("generation.lm_top_k_info")
+                )
+                lm_top_p = gr.Slider(
+                    label=t("generation.lm_top_p_label"),
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.9,
+                    step=0.01,
+                    scale=1,
+                    info=t("generation.lm_top_p_info")
+                )
+            
+            with gr.Row():
+                lm_negative_prompt = gr.Textbox(
+                    label=t("generation.lm_negative_prompt_label"),
+                    value="NO USER INPUT",
+                    placeholder=t("generation.lm_negative_prompt_placeholder"),
+                    info=t("generation.lm_negative_prompt_info"),
+                    lines=2,
+                    scale=2,
+                )
+            
+            with gr.Row():
+                use_cot_metas = gr.Checkbox(
+                    label=t("generation.cot_metas_label"),
+                    value=True,
+                    info=t("generation.cot_metas_info"),
+                    scale=1,
+                )
+                use_cot_language = gr.Checkbox(
+                    label=t("generation.cot_language_label"),
+                    value=True,
+                    info=t("generation.cot_language_info"),
+                    scale=1,
+                )
+                constrained_decoding_debug = gr.Checkbox(
+                    label=t("generation.constrained_debug_label"),
+                    value=False,
+                    info=t("generation.constrained_debug_info"),
+                    scale=1,
+                    interactive=not service_mode  # Fixed in service mode
+                )
+            
+            with gr.Row():
+                auto_score = gr.Checkbox(
+                    label=t("generation.auto_score_label"),
+                    value=False,
+                    info=t("generation.auto_score_info"),
+                    scale=1,
+                    interactive=not service_mode  # Fixed in service mode
+                )
+                auto_lrc = gr.Checkbox(
+                    label=t("generation.auto_lrc_label"),
+                    value=False,
+                    info=t("generation.auto_lrc_info"),
+                    scale=1,
+                    interactive=not service_mode  # Fixed in service mode
+                )
+                lm_batch_chunk_size = gr.Number(
+                    label=t("generation.lm_batch_chunk_label"),
+                    value=8,
+                    minimum=1,
+                    maximum=32,
+                    step=1,
+                    info=t("generation.lm_batch_chunk_info"),
+                    scale=1,
+                    interactive=not service_mode  # Fixed in service mode
+                )
+            
+            with gr.Row():
+                audio_cover_strength = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=1.0,
+                    step=0.01,
+                    label=t("generation.codes_strength_label"),
+                    info=t("generation.codes_strength_info"),
+                    scale=1,
+                )
+                score_scale = gr.Slider(
+                    minimum=0.01,
+                    maximum=1.0,
+                    value=0.5,
+                    step=0.01,
+                    label=t("generation.score_sensitivity_label"),
+                    info=t("generation.score_sensitivity_info"),
+                    scale=1,
+                    visible=not service_mode  # Hidden in service mode
+                )
+        
+        # Set generate_btn to interactive if service is pre-initialized
+        generate_btn_interactive = init_params.get('enable_generate', False) if service_pre_initialized else False
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=1, variant="compact"):
+                think_checkbox = gr.Checkbox(
+                    label=t("generation.think_label"),
+                    value=True,
+                    scale=1,
+                )
+                allow_lm_batch = gr.Checkbox(
+                    label=t("generation.parallel_thinking_label"),
+                    value=True,
+                    scale=1,
+                )
+            with gr.Column(scale=18):
+                generate_btn = gr.Button(t("generation.generate_btn"), variant="primary", size="lg", interactive=generate_btn_interactive)
+            with gr.Column(scale=1, variant="compact"):
+                autogen_checkbox = gr.Checkbox(
+                    label=t("generation.autogen_label"),
+                    value=False,  # Default to False for both service and local modes
+                    scale=1,
+                    interactive=not service_mode  # Not selectable in service mode
+                )
+                use_cot_caption = gr.Checkbox(
+                    label=t("generation.caption_rewrite_label"),
+                    value=True,
+                    scale=1,
+                )
+    
+    return {
+        "service_config_accordion": service_config_accordion,
+        "language_dropdown": language_dropdown,
+        "checkpoint_dropdown": checkpoint_dropdown,
+        "refresh_btn": refresh_btn,
+        "config_path": config_path,
+        "device": device,
+        "init_btn": init_btn,
+        "init_status": init_status,
+        "lm_model_path": lm_model_path,
+        "init_llm_checkbox": init_llm_checkbox,
+        "backend_dropdown": backend_dropdown,
+        "use_flash_attention_checkbox": use_flash_attention_checkbox,
+        "offload_to_cpu_checkbox": offload_to_cpu_checkbox,
+        "offload_dit_to_cpu_checkbox": offload_dit_to_cpu_checkbox,
+        "compile_model_checkbox": compile_model_checkbox,
+        "quantization_checkbox": quantization_checkbox,
+        # LoRA components
+        "lora_path": lora_path,
+        "load_lora_btn": load_lora_btn,
+        "unload_lora_btn": unload_lora_btn,
+        "use_lora_checkbox": use_lora_checkbox,
+        "lora_scale_slider": lora_scale_slider,
+        "lora_status": lora_status,
+        "task_type": task_type,
+        "instruction_display_gen": instruction_display_gen,
+        "track_name": track_name,
+        "complete_track_classes": complete_track_classes,
+        "audio_uploads_accordion": audio_uploads_accordion,
+        "reference_audio": reference_audio,
+        "src_audio": src_audio,
+        "convert_src_to_codes_btn": convert_src_to_codes_btn,
+        "text2music_audio_code_string": text2music_audio_code_string,
+        "transcribe_btn": transcribe_btn,
+        "text2music_audio_codes_group": text2music_audio_codes_group,
+        "lm_temperature": lm_temperature,
+        "lm_cfg_scale": lm_cfg_scale,
+        "lm_top_k": lm_top_k,
+        "lm_top_p": lm_top_p,
+        "lm_negative_prompt": lm_negative_prompt,
+        "use_cot_metas": use_cot_metas,
+        "use_cot_caption": use_cot_caption,
+        "use_cot_language": use_cot_language,
+        "repainting_group": repainting_group,
+        "repainting_start": repainting_start,
+        "repainting_end": repainting_end,
+        "audio_cover_strength": audio_cover_strength,
+        # Simple/Custom Mode Components
+        "generation_mode": generation_mode,
+        "simple_mode_group": simple_mode_group,
+        "simple_query_input": simple_query_input,
+        "random_desc_btn": random_desc_btn,
+        "simple_instrumental_checkbox": simple_instrumental_checkbox,
+        "simple_vocal_language": simple_vocal_language,
+        "create_sample_btn": create_sample_btn,
+        "simple_sample_created": simple_sample_created,
+        "caption_accordion": caption_accordion,
+        "lyrics_accordion": lyrics_accordion,
+        "optional_params_accordion": optional_params_accordion,
+        # Existing components
+        "captions": captions,
+        "sample_btn": sample_btn,
+        "load_file": load_file,
+        "lyrics": lyrics,
+        "vocal_language": vocal_language,
+        "bpm": bpm,
+        "key_scale": key_scale,
+        "time_signature": time_signature,
+        "audio_duration": audio_duration,
+        "batch_size_input": batch_size_input,
+        "inference_steps": inference_steps,
+        "guidance_scale": guidance_scale,
+        "seed": seed,
+        "random_seed_checkbox": random_seed_checkbox,
+        "use_adg": use_adg,
+        "cfg_interval_start": cfg_interval_start,
+        "cfg_interval_end": cfg_interval_end,
+        "shift": shift,
+        "infer_method": infer_method,
+        "custom_timesteps": custom_timesteps,
+        "audio_format": audio_format,
+        "think_checkbox": think_checkbox,
+        "autogen_checkbox": autogen_checkbox,
+        "generate_btn": generate_btn,
+        "instrumental_checkbox": instrumental_checkbox,
+        "format_btn": format_btn,
+        "constrained_decoding_debug": constrained_decoding_debug,
+        "score_scale": score_scale,
+        "allow_lm_batch": allow_lm_batch,
+        "auto_score": auto_score,
+        "auto_lrc": auto_lrc,
+        "lm_batch_chunk_size": lm_batch_chunk_size,
+        # GPU config values for validation
+        "gpu_config": gpu_config,
+        "max_duration": max_duration,
+        "max_batch_size": max_batch_size,
+    }

acestep/gradio_ui/interfaces/result.py

@@ -0,0 +1,552 @@
+"""
+Gradio UI Results Section Module
+Contains results display section component definitions
+"""
+import gradio as gr
+from acestep.gradio_ui.i18n import t
+
+
+def create_results_section(dit_handler) -> dict:
+    """Create results display section"""
+    with gr.Accordion(t("results.title"), open=True):
+        # Hidden state to store LM-generated metadata
+        lm_metadata_state = gr.State(value=None)
+        
+        # Hidden state to track if caption/metadata is from formatted source (LM/transcription)
+        is_format_caption_state = gr.State(value=False)
+        
+        # Batch management states
+        current_batch_index = gr.State(value=0)  # Currently displayed batch index
+        total_batches = gr.State(value=1)  # Total number of batches generated
+        batch_queue = gr.State(value={})  # Dictionary storing all batch data
+        generation_params_state = gr.State(value={})  # Store generation parameters for next batches
+        is_generating_background = gr.State(value=False)  # Background generation flag
+
+        # All audio components in one row with dynamic visibility
+        with gr.Row():
+            with gr.Column(visible=True) as audio_col_1:
+                generated_audio_1 = gr.Audio(
+                    label=t("results.generated_music", n=1),
+                    type="filepath",
+                    interactive=False,
+                    buttons=[]
+                )
+                with gr.Row(equal_height=True):
+                    send_to_src_btn_1 = gr.Button(
+                        t("results.send_to_src_btn"),
+                        variant="secondary",
+                        size="sm",
+                        scale=1
+                    )
+                    save_btn_1 = gr.Button(
+                        t("results.save_btn"),
+                        variant="primary",
+                        size="sm",
+                        scale=1
+                    )
+                    score_btn_1 = gr.Button(
+                        t("results.score_btn"),
+                        variant="secondary",
+                        size="sm",
+                        scale=1
+                    )
+                    lrc_btn_1 = gr.Button(
+                        t("results.lrc_btn"),
+                        variant="secondary",
+                        size="sm",
+                        scale=1
+                    )
+                with gr.Accordion(t("results.details_accordion"), open=False, visible=True) as details_accordion_1:
+                    codes_display_1 = gr.Textbox(
+                        label=t("results.codes_label", n=1),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=4,
+                        max_lines=4,
+                        visible=True
+                    )
+                    score_display_1 = gr.Textbox(
+                        label=t("results.quality_score_label", n=1),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=6,
+                        max_lines=6,
+                        visible=True
+                    )
+                    lrc_display_1 = gr.Textbox(
+                        label=t("results.lrc_label", n=1),
+                        interactive=True,
+                        buttons=["copy"],
+                        lines=8,
+                        max_lines=8,
+                        visible=True
+                    )
+            with gr.Column(visible=True) as audio_col_2:
+                generated_audio_2 = gr.Audio(
+                    label=t("results.generated_music", n=2),
+                    type="filepath",
+                    interactive=False,
+                    buttons=[]
+                )
+                with gr.Row(equal_height=True):
+                    send_to_src_btn_2 = gr.Button(
+                        t("results.send_to_src_btn"),
+                        variant="secondary",
+                        size="sm",
+                        scale=1
+                    )
+                    save_btn_2 = gr.Button(
+                        t("results.save_btn"),
+                        variant="primary",
+                        size="sm",
+                        scale=1
+                    )
+                    score_btn_2 = gr.Button(
+                        t("results.score_btn"),
+                        variant="secondary",
+                        size="sm",
+                        scale=1
+                    )
+                    lrc_btn_2 = gr.Button(
+                        t("results.lrc_btn"),
+                        variant="secondary",
+                        size="sm",
+                        scale=1
+                    )
+                with gr.Accordion(t("results.details_accordion"), open=False, visible=True) as details_accordion_2:
+                    codes_display_2 = gr.Textbox(
+                        label=t("results.codes_label", n=2),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=4,
+                        max_lines=4,
+                        visible=True
+                    )
+                    score_display_2 = gr.Textbox(
+                        label=t("results.quality_score_label", n=2),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=6,
+                        max_lines=6,
+                        visible=True
+                    )
+                    lrc_display_2 = gr.Textbox(
+                        label=t("results.lrc_label", n=2),
+                        interactive=True,
+                        buttons=["copy"],
+                        lines=8,
+                        max_lines=8,
+                        visible=True
+                    )
+            with gr.Column(visible=False) as audio_col_3:
+                generated_audio_3 = gr.Audio(
+                    label=t("results.generated_music", n=3),
+                    type="filepath",
+                    interactive=False,
+                    buttons=[]
+                )
+                with gr.Row(equal_height=True):
+                    send_to_src_btn_3 = gr.Button(
+                        t("results.send_to_src_btn"),
+                        variant="secondary",
+                        size="sm",
+                        scale=1
+                    )
+                    save_btn_3 = gr.Button(
+                        t("results.save_btn"),
+                        variant="primary",
+                        size="sm",
+                        scale=1
+                    )
+                    score_btn_3 = gr.Button(
+                        t("results.score_btn"),
+                        variant="secondary",
+                        size="sm",
+                        scale=1
+                    )
+                    lrc_btn_3 = gr.Button(
+                        t("results.lrc_btn"),
+                        variant="secondary",
+                        size="sm",
+                        scale=1
+                    )
+                with gr.Accordion(t("results.details_accordion"), open=False, visible=True) as details_accordion_3:
+                    codes_display_3 = gr.Textbox(
+                        label=t("results.codes_label", n=3),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=4,
+                        max_lines=4,
+                        visible=True
+                    )
+                    score_display_3 = gr.Textbox(
+                        label=t("results.quality_score_label", n=3),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=6,
+                        max_lines=6,
+                        visible=True
+                    )
+                    lrc_display_3 = gr.Textbox(
+                        label=t("results.lrc_label", n=3),
+                        interactive=True,
+                        buttons=["copy"],
+                        lines=8,
+                        max_lines=8,
+                        visible=True
+                    )
+            with gr.Column(visible=False) as audio_col_4:
+                generated_audio_4 = gr.Audio(
+                    label=t("results.generated_music", n=4),
+                    type="filepath",
+                    interactive=False,
+                    buttons=[]
+                )
+                with gr.Row(equal_height=True):
+                    send_to_src_btn_4 = gr.Button(
+                        t("results.send_to_src_btn"),
+                        variant="secondary",
+                        size="sm",
+                        scale=1
+                    )
+                    save_btn_4 = gr.Button(
+                        t("results.save_btn"),
+                        variant="primary",
+                        size="sm",
+                        scale=1
+                    )
+                    score_btn_4 = gr.Button(
+                        t("results.score_btn"),
+                        variant="secondary",
+                        size="sm",
+                        scale=1
+                    )
+                    lrc_btn_4 = gr.Button(
+                        t("results.lrc_btn"),
+                        variant="secondary",
+                        size="sm",
+                        scale=1
+                    )
+                with gr.Accordion(t("results.details_accordion"), open=False, visible=True) as details_accordion_4:
+                    codes_display_4 = gr.Textbox(
+                        label=t("results.codes_label", n=4),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=4,
+                        max_lines=4,
+                        visible=True
+                    )
+                    score_display_4 = gr.Textbox(
+                        label=t("results.quality_score_label", n=4),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=6,
+                        max_lines=6,
+                        visible=True
+                    )
+                    lrc_display_4 = gr.Textbox(
+                        label=t("results.lrc_label", n=4),
+                        interactive=True,
+                        buttons=["copy"],
+                        lines=8,
+                        max_lines=8,
+                        visible=True
+                    )
+        
+        # Second row for batch size 5-8 (initially hidden)
+        with gr.Row(visible=False) as audio_row_5_8:
+            with gr.Column() as audio_col_5:
+                generated_audio_5 = gr.Audio(
+                    label=t("results.generated_music", n=5),
+                    type="filepath",
+                    interactive=False,
+                    buttons=[]
+                )
+                with gr.Row(equal_height=True):
+                    send_to_src_btn_5 = gr.Button(t("results.send_to_src_btn"), variant="secondary", size="sm", scale=1)
+                    save_btn_5 = gr.Button(t("results.save_btn"), variant="primary", size="sm", scale=1)
+                    score_btn_5 = gr.Button(t("results.score_btn"), variant="secondary", size="sm", scale=1)
+                    lrc_btn_5 = gr.Button(t("results.lrc_btn"), variant="secondary", size="sm", scale=1)
+                with gr.Accordion(t("results.details_accordion"), open=False, visible=True) as details_accordion_5:
+                    codes_display_5 = gr.Textbox(
+                        label=t("results.codes_label", n=5),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=4,
+                        max_lines=4,
+                        visible=True
+                    )
+                    score_display_5 = gr.Textbox(
+                        label=t("results.quality_score_label", n=5),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=6,
+                        max_lines=6,
+                        visible=True
+                    )
+                    lrc_display_5 = gr.Textbox(
+                        label=t("results.lrc_label", n=5),
+                        interactive=True,
+                        buttons=["copy"],
+                        lines=8,
+                        max_lines=8,
+                        visible=True
+                    )
+            with gr.Column() as audio_col_6:
+                generated_audio_6 = gr.Audio(
+                    label=t("results.generated_music", n=6),
+                    type="filepath",
+                    interactive=False,
+                    buttons=[]
+                )
+                with gr.Row(equal_height=True):
+                    send_to_src_btn_6 = gr.Button(t("results.send_to_src_btn"), variant="secondary", size="sm", scale=1)
+                    save_btn_6 = gr.Button(t("results.save_btn"), variant="primary", size="sm", scale=1)
+                    score_btn_6 = gr.Button(t("results.score_btn"), variant="secondary", size="sm", scale=1)
+                    lrc_btn_6 = gr.Button(t("results.lrc_btn"), variant="secondary", size="sm", scale=1)
+                with gr.Accordion(t("results.details_accordion"), open=False, visible=True) as details_accordion_6:
+                    codes_display_6 = gr.Textbox(
+                        label=t("results.codes_label", n=6),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=4,
+                        max_lines=4,
+                        visible=True
+                    )
+                    score_display_6 = gr.Textbox(
+                        label=t("results.quality_score_label", n=6),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=6,
+                        max_lines=6,
+                        visible=True
+                    )
+                    lrc_display_6 = gr.Textbox(
+                        label=t("results.lrc_label", n=6),
+                        interactive=True,
+                        buttons=["copy"],
+                        lines=8,
+                        max_lines=8,
+                        visible=True
+                    )
+            with gr.Column() as audio_col_7:
+                generated_audio_7 = gr.Audio(
+                    label=t("results.generated_music", n=7),
+                    type="filepath",
+                    interactive=False,
+                    buttons=[]
+                )
+                with gr.Row(equal_height=True):
+                    send_to_src_btn_7 = gr.Button(t("results.send_to_src_btn"), variant="secondary", size="sm", scale=1)
+                    save_btn_7 = gr.Button(t("results.save_btn"), variant="primary", size="sm", scale=1)
+                    score_btn_7 = gr.Button(t("results.score_btn"), variant="secondary", size="sm", scale=1)
+                    lrc_btn_7 = gr.Button(t("results.lrc_btn"), variant="secondary", size="sm", scale=1)
+                with gr.Accordion(t("results.details_accordion"), open=False, visible=True) as details_accordion_7:
+                    codes_display_7 = gr.Textbox(
+                        label=t("results.codes_label", n=7),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=4,
+                        max_lines=4,
+                        visible=True
+                    )
+                    score_display_7 = gr.Textbox(
+                        label=t("results.quality_score_label", n=7),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=6,
+                        max_lines=6,
+                        visible=True
+                    )
+                    lrc_display_7 = gr.Textbox(
+                        label=t("results.lrc_label", n=7),
+                        interactive=True,
+                        buttons=["copy"],
+                        lines=8,
+                        max_lines=8,
+                        visible=True
+                    )
+            with gr.Column() as audio_col_8:
+                generated_audio_8 = gr.Audio(
+                    label=t("results.generated_music", n=8),
+                    type="filepath",
+                    interactive=False,
+                    buttons=[]
+                )
+                with gr.Row(equal_height=True):
+                    send_to_src_btn_8 = gr.Button(t("results.send_to_src_btn"), variant="secondary", size="sm", scale=1)
+                    save_btn_8 = gr.Button(t("results.save_btn"), variant="primary", size="sm", scale=1)
+                    score_btn_8 = gr.Button(t("results.score_btn"), variant="secondary", size="sm", scale=1)
+                    lrc_btn_8 = gr.Button(t("results.lrc_btn"), variant="secondary", size="sm", scale=1)
+                with gr.Accordion(t("results.details_accordion"), open=False, visible=True) as details_accordion_8:
+                    codes_display_8 = gr.Textbox(
+                        label=t("results.codes_label", n=8),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=4,
+                        max_lines=4,
+                        visible=True
+                    )
+                    score_display_8 = gr.Textbox(
+                        label=t("results.quality_score_label", n=8),
+                        interactive=False,
+                        buttons=["copy"],
+                        lines=6,
+                        max_lines=6,
+                        visible=True
+                    )
+                    lrc_display_8 = gr.Textbox(
+                        label=t("results.lrc_label", n=8),
+                        interactive=True,
+                        buttons=["copy"],
+                        lines=8,
+                        max_lines=8,
+                        visible=True
+                    )
+        
+        status_output = gr.Textbox(label=t("results.generation_status"), interactive=False)
+        
+        # Batch navigation controls
+        with gr.Row(equal_height=True):
+            prev_batch_btn = gr.Button(
+                t("results.prev_btn"),
+                variant="secondary",
+                interactive=False,
+                scale=1,
+                size="sm"
+            )
+            batch_indicator = gr.Textbox(
+                label=t("results.current_batch"),
+                value=t("results.batch_indicator", current=1, total=1),
+                interactive=False,
+                scale=3
+            )
+            next_batch_status = gr.Textbox(
+                label=t("results.next_batch_status"),
+                value="",
+                interactive=False,
+                scale=3
+            )
+            next_batch_btn = gr.Button(
+                t("results.next_btn"),
+                variant="primary",
+                interactive=False,
+                scale=1,
+                size="sm"
+            )
+        
+        # One-click restore parameters button
+        restore_params_btn = gr.Button(
+            t("results.restore_params_btn"),
+            variant="secondary",
+            interactive=False,  # Initially disabled, enabled after generation
+            size="sm"
+        )
+        
+        with gr.Accordion(t("results.batch_results_title"), open=False):
+            generated_audio_batch = gr.File(
+                label=t("results.all_files_label"),
+                file_count="multiple",
+                interactive=False
+            )
+            generation_info = gr.Markdown(label=t("results.generation_details"))
+    
+    return {
+        "lm_metadata_state": lm_metadata_state,
+        "is_format_caption_state": is_format_caption_state,
+        "current_batch_index": current_batch_index,
+        "total_batches": total_batches,
+        "batch_queue": batch_queue,
+        "generation_params_state": generation_params_state,
+        "is_generating_background": is_generating_background,
+        "status_output": status_output,
+        "prev_batch_btn": prev_batch_btn,
+        "batch_indicator": batch_indicator,
+        "next_batch_btn": next_batch_btn,
+        "next_batch_status": next_batch_status,
+        "restore_params_btn": restore_params_btn,
+        "generated_audio_1": generated_audio_1,
+        "generated_audio_2": generated_audio_2,
+        "generated_audio_3": generated_audio_3,
+        "generated_audio_4": generated_audio_4,
+        "generated_audio_5": generated_audio_5,
+        "generated_audio_6": generated_audio_6,
+        "generated_audio_7": generated_audio_7,
+        "generated_audio_8": generated_audio_8,
+        "audio_row_5_8": audio_row_5_8,
+        "audio_col_1": audio_col_1,
+        "audio_col_2": audio_col_2,
+        "audio_col_3": audio_col_3,
+        "audio_col_4": audio_col_4,
+        "audio_col_5": audio_col_5,
+        "audio_col_6": audio_col_6,
+        "audio_col_7": audio_col_7,
+        "audio_col_8": audio_col_8,
+        "send_to_src_btn_1": send_to_src_btn_1,
+        "send_to_src_btn_2": send_to_src_btn_2,
+        "send_to_src_btn_3": send_to_src_btn_3,
+        "send_to_src_btn_4": send_to_src_btn_4,
+        "send_to_src_btn_5": send_to_src_btn_5,
+        "send_to_src_btn_6": send_to_src_btn_6,
+        "send_to_src_btn_7": send_to_src_btn_7,
+        "send_to_src_btn_8": send_to_src_btn_8,
+        "save_btn_1": save_btn_1,
+        "save_btn_2": save_btn_2,
+        "save_btn_3": save_btn_3,
+        "save_btn_4": save_btn_4,
+        "save_btn_5": save_btn_5,
+        "save_btn_6": save_btn_6,
+        "save_btn_7": save_btn_7,
+        "save_btn_8": save_btn_8,
+        "score_btn_1": score_btn_1,
+        "score_btn_2": score_btn_2,
+        "score_btn_3": score_btn_3,
+        "score_btn_4": score_btn_4,
+        "score_btn_5": score_btn_5,
+        "score_btn_6": score_btn_6,
+        "score_btn_7": score_btn_7,
+        "score_btn_8": score_btn_8,
+        "score_display_1": score_display_1,
+        "score_display_2": score_display_2,
+        "score_display_3": score_display_3,
+        "score_display_4": score_display_4,
+        "score_display_5": score_display_5,
+        "score_display_6": score_display_6,
+        "score_display_7": score_display_7,
+        "score_display_8": score_display_8,
+        "codes_display_1": codes_display_1,
+        "codes_display_2": codes_display_2,
+        "codes_display_3": codes_display_3,
+        "codes_display_4": codes_display_4,
+        "codes_display_5": codes_display_5,
+        "codes_display_6": codes_display_6,
+        "codes_display_7": codes_display_7,
+        "codes_display_8": codes_display_8,
+        "lrc_btn_1": lrc_btn_1,
+        "lrc_btn_2": lrc_btn_2,
+        "lrc_btn_3": lrc_btn_3,
+        "lrc_btn_4": lrc_btn_4,
+        "lrc_btn_5": lrc_btn_5,
+        "lrc_btn_6": lrc_btn_6,
+        "lrc_btn_7": lrc_btn_7,
+        "lrc_btn_8": lrc_btn_8,
+        "lrc_display_1": lrc_display_1,
+        "lrc_display_2": lrc_display_2,
+        "lrc_display_3": lrc_display_3,
+        "lrc_display_4": lrc_display_4,
+        "lrc_display_5": lrc_display_5,
+        "lrc_display_6": lrc_display_6,
+        "lrc_display_7": lrc_display_7,
+        "lrc_display_8": lrc_display_8,
+        "details_accordion_1": details_accordion_1,
+        "details_accordion_2": details_accordion_2,
+        "details_accordion_3": details_accordion_3,
+        "details_accordion_4": details_accordion_4,
+        "details_accordion_5": details_accordion_5,
+        "details_accordion_6": details_accordion_6,
+        "details_accordion_7": details_accordion_7,
+        "details_accordion_8": details_accordion_8,
+        "generated_audio_batch": generated_audio_batch,
+        "generation_info": generation_info,
+    }
+

acestep/gradio_ui/interfaces/training.py

@@ -0,0 +1,625 @@
+"""
+Gradio UI Training Tab Module
+
+Contains the dataset builder and LoRA training interface components.
+"""
+
+import os
+import gradio as gr
+from acestep.gradio_ui.i18n import t
+from acestep.constants import DEBUG_TRAINING
+
+
+def create_training_section(dit_handler, llm_handler, init_params=None) -> dict:
+    """Create the training tab section with dataset builder and training controls.
+    
+    Args:
+        dit_handler: DiT handler instance
+        llm_handler: LLM handler instance
+        init_params: Dictionary containing initialization parameters and state.
+                    If None, service will not be pre-initialized.
+        
+    Returns:
+        Dictionary of Gradio components for event handling
+    """
+    # Check if running in service mode (hide training tab)
+    service_mode = init_params is not None and init_params.get('service_mode', False)
+    
+    debug_training_enabled = str(DEBUG_TRAINING).strip().upper() != "OFF"
+    epoch_min = 1 if debug_training_enabled else 100
+    epoch_step = 1 if debug_training_enabled else 100
+    epoch_default = 1 if debug_training_enabled else 1000
+
+    with gr.Tab(t("training.tab_title"), visible=not service_mode):
+        gr.HTML("""
+        <div style="text-align: center; padding: 10px; margin-bottom: 15px;">
+            <h2>🎵 LoRA Training for ACE-Step</h2>
+            <p>Build datasets from your audio files and train custom LoRA adapters</p>
+        </div>
+        """)
+        
+        with gr.Tabs():
+            # ==================== Dataset Builder Tab ====================
+            with gr.Tab(t("training.tab_dataset_builder")):
+                # ========== Load Existing OR Scan New ==========
+                gr.HTML(f"""
+                <div style="padding: 10px; margin-bottom: 10px; border: 1px solid #4a4a6a; border-radius: 8px; background: linear-gradient(135deg, #2a2a4a 0%, #1a1a3a 100%);">
+                    <h3 style="margin: 0 0 5px 0;">{t("training.quick_start_title")}</h3>
+                    <p style="margin: 0; color: #aaa;">Choose one: <b>Load existing dataset</b> OR <b>Scan new directory</b></p>
+                </div>
+                """)
+                
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        gr.HTML("<h4>📂 Load Existing Dataset</h4>")
+                        with gr.Row():
+                            load_json_path = gr.Textbox(
+                                label=t("training.load_dataset_label"),
+                                placeholder="./datasets/my_lora_dataset.json",
+                                info=t("training.load_dataset_info"),
+                                scale=3,
+                            )
+                            load_json_btn = gr.Button(t("training.load_btn"), variant="primary", scale=1)
+                        load_json_status = gr.Textbox(
+                            label=t("training.load_status"),
+                            interactive=False,
+                        )
+                    
+                    with gr.Column(scale=1):
+                        gr.HTML("<h4>🔍 Scan New Directory</h4>")
+                        with gr.Row():
+                            audio_directory = gr.Textbox(
+                                label=t("training.scan_label"),
+                                placeholder="/path/to/your/audio/folder",
+                                info=t("training.scan_info"),
+                                scale=3,
+                            )
+                            scan_btn = gr.Button(t("training.scan_btn"), variant="secondary", scale=1)
+                        scan_status = gr.Textbox(
+                            label=t("training.scan_status"),
+                            interactive=False,
+                        )
+                
+                gr.HTML("<hr>")
+                
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        
+                        # Audio files table
+                        audio_files_table = gr.Dataframe(
+                            headers=["#", "Filename", "Duration", "Lyrics", "Labeled", "BPM", "Key", "Caption"],
+                            datatype=["number", "str", "str", "str", "str", "str", "str", "str"],
+                            label=t("training.found_audio_files"),
+                            interactive=False,
+                            wrap=True,
+                        )
+                    
+                    with gr.Column(scale=1):
+                        gr.HTML(f"<h3>⚙️ {t('training.dataset_settings_header')}</h3>")
+                        
+                        dataset_name = gr.Textbox(
+                            label=t("training.dataset_name"),
+                            value="my_lora_dataset",
+                            placeholder=t("training.dataset_name_placeholder"),
+                        )
+                        
+                        all_instrumental = gr.Checkbox(
+                            label=t("training.all_instrumental"),
+                            value=True,
+                            info=t("training.all_instrumental_info"),
+                        )
+
+                        format_lyrics = gr.Checkbox(
+                            label="Format Lyrics (LM)",
+                            value=False,
+                            info="Use LM to format/structure user-provided lyrics from .txt files (coming soon)",
+                            interactive=False,  # Disabled for now - model update needed
+                        )
+
+                        transcribe_lyrics = gr.Checkbox(
+                            label="Transcribe Lyrics (LM)",
+                            value=False,
+                            info="Use LM to transcribe lyrics from audio (coming soon)",
+                            interactive=False,  # Disabled for now - model update needed
+                        )
+                        
+                        custom_tag = gr.Textbox(
+                            label=t("training.custom_tag"),
+                            placeholder="e.g., 8bit_retro, my_style",
+                            info=t("training.custom_tag_info"),
+                        )
+                        
+                        tag_position = gr.Radio(
+                            choices=[
+                                (t("training.tag_prepend"), "prepend"),
+                                (t("training.tag_append"), "append"),
+                                (t("training.tag_replace"), "replace"),
+                            ],
+                            value="replace",
+                            label=t("training.tag_position"),
+                            info=t("training.tag_position_info"),
+                        )
+
+                        genre_ratio = gr.Slider(
+                            minimum=0,
+                            maximum=100,
+                            step=10,
+                            value=0,
+                            label=t("training.genre_ratio"),
+                            info=t("training.genre_ratio_info"),
+                        )
+
+                gr.HTML(f"<hr><h3>🤖 {t('training.step2_title')}</h3>")
+                
+                with gr.Row():
+                    with gr.Column(scale=3):
+                        gr.Markdown("""
+                        Click the button below to automatically generate metadata for all audio files using AI:
+                        - **Caption**: Music style, genre, mood description
+                        - **BPM**: Beats per minute
+                        - **Key**: Musical key (e.g., C Major, Am)
+                        - **Time Signature**: 4/4, 3/4, etc.
+                        """)
+                        skip_metas = gr.Checkbox(
+                            label=t("training.skip_metas"),
+                            value=False,
+                            info=t("training.skip_metas_info"),
+                        )
+                        only_unlabeled = gr.Checkbox(
+                            label=t("training.only_unlabeled"),
+                            value=False,
+                            info=t("training.only_unlabeled_info"),
+                        )
+                    with gr.Column(scale=1):
+                        auto_label_btn = gr.Button(
+                            t("training.auto_label_btn"),
+                            variant="primary",
+                            size="lg",
+                        )
+                
+                label_progress = gr.Textbox(
+                    label=t("training.label_progress"),
+                    interactive=False,
+                    lines=2,
+                )
+                
+                gr.HTML(f"<hr><h3>👀 {t('training.step3_title')}</h3>")
+                
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        sample_selector = gr.Slider(
+                            minimum=0,
+                            maximum=0,
+                            step=1,
+                            value=0,
+                            label=t("training.select_sample"),
+                            info=t("training.select_sample_info"),
+                        )
+                        
+                        preview_audio = gr.Audio(
+                            label=t("training.audio_preview"),
+                            type="filepath",
+                            interactive=False,
+                        )
+                        
+                        preview_filename = gr.Textbox(
+                            label=t("training.filename"),
+                            interactive=False,
+                        )
+
+                    with gr.Column(scale=2):
+                        with gr.Row():
+                            edit_caption = gr.Textbox(
+                                label=t("training.caption"),
+                                lines=3,
+                                placeholder="Music description...",
+                            )
+
+                        with gr.Row():
+                            edit_genre = gr.Textbox(
+                                label=t("training.genre"),
+                                lines=1,
+                                placeholder="pop, electronic, dance...",
+                            )
+                            prompt_override = gr.Dropdown(
+                                choices=["Use Global Ratio", "Caption", "Genre"],
+                                value="Use Global Ratio",
+                                label=t("training.prompt_override_label"),
+                                info=t("training.prompt_override_info"),
+                            )
+
+                        with gr.Row():
+                            edit_lyrics = gr.Textbox(
+                                label=t("training.lyrics_editable_label"),
+                                lines=6,
+                                placeholder="[Verse 1]\nLyrics here...\n\n[Chorus]\n...",
+                            )
+                            raw_lyrics_display = gr.Textbox(
+                                label=t("training.raw_lyrics_label"),
+                                lines=6,
+                                placeholder=t("training.no_lyrics_placeholder"),
+                                interactive=False,  # Read-only, can copy but not edit
+                                visible=False,  # Hidden when no raw lyrics
+                            )
+                            has_raw_lyrics_state = gr.State(False)  # Track visibility
+
+                        with gr.Row():
+                            edit_bpm = gr.Number(
+                                label=t("training.bpm"),
+                                precision=0,
+                            )
+                            edit_keyscale = gr.Textbox(
+                                label=t("training.key_label"),
+                                placeholder=t("training.key_placeholder"),
+                            )
+                            edit_timesig = gr.Dropdown(
+                                choices=["", "2", "3", "4", "6", "N/A"],
+                                label=t("training.time_sig"),
+                            )
+                            edit_duration = gr.Number(
+                                label=t("training.duration_s"),
+                                precision=1,
+                                interactive=False,
+                            )
+                        
+                        with gr.Row():
+                            edit_language = gr.Dropdown(
+                                choices=["instrumental", "en", "zh", "ja", "ko", "es", "fr", "de", "pt", "ru", "unknown"],
+                                value="instrumental",
+                                label=t("training.language"),
+                            )
+                            edit_instrumental = gr.Checkbox(
+                                label=t("training.instrumental"),
+                                value=True,
+                            )
+                            save_edit_btn = gr.Button(t("training.save_changes_btn"), variant="secondary")
+                        
+                        edit_status = gr.Textbox(
+                            label=t("training.edit_status"),
+                            interactive=False,
+                        )
+                
+                gr.HTML(f"<hr><h3>💾 {t('training.step4_title')}</h3>")
+                
+                with gr.Row():
+                    with gr.Column(scale=3):
+                        save_path = gr.Textbox(
+                            label=t("training.save_path"),
+                            value="./datasets/my_lora_dataset.json",
+                            placeholder="./datasets/dataset_name.json",
+                            info=t("training.save_path_info"),
+                        )
+                    with gr.Column(scale=1):
+                        save_dataset_btn = gr.Button(
+                            t("training.save_dataset_btn"),
+                            variant="primary",
+                            size="lg",
+                        )
+                
+                save_status = gr.Textbox(
+                    label=t("training.save_status"),
+                    interactive=False,
+                    lines=2,
+                )
+                
+                gr.HTML(f"<hr><h3>⚡ {t('training.step5_title')}</h3>")
+                
+                gr.Markdown("""
+                **Preprocessing converts your dataset to pre-computed tensors for fast training.**
+                
+                You can either:
+                - Use the dataset from Steps 1-4 above, **OR**
+                - Load an existing dataset JSON file (if you've already saved one)
+                """)
+                
+                with gr.Row():
+                    with gr.Column(scale=3):
+                        load_existing_dataset_path = gr.Textbox(
+                            label=t("training.load_existing_label"),
+                            placeholder="./datasets/my_lora_dataset.json",
+                            info=t("training.load_existing_info"),
+                        )
+                    with gr.Column(scale=1):
+                        load_existing_dataset_btn = gr.Button(
+                            t("training.load_dataset_btn"),
+                            variant="secondary",
+                            size="lg",
+                        )
+                
+                load_existing_status = gr.Textbox(
+                    label=t("training.load_status"),
+                    interactive=False,
+                )
+                
+                gr.Markdown("""
+                This step:
+                - Encodes audio to VAE latents
+                - Encodes captions and lyrics to text embeddings  
+                - Runs the condition encoder
+                - Saves all tensors to `.pt` files
+                
+                ⚠️ **This requires the model to be loaded and may take a few minutes.**
+                """)
+                
+                with gr.Row():
+                    with gr.Column(scale=3):
+                        preprocess_output_dir = gr.Textbox(
+                            label=t("training.tensor_output_dir"),
+                            value="./datasets/preprocessed_tensors",
+                            placeholder="./datasets/preprocessed_tensors",
+                            info=t("training.tensor_output_info"),
+                        )
+                    with gr.Column(scale=1):
+                        preprocess_btn = gr.Button(
+                            t("training.preprocess_btn"),
+                            variant="primary",
+                            size="lg",
+                        )
+                
+                preprocess_progress = gr.Textbox(
+                    label=t("training.preprocess_progress"),
+                    interactive=False,
+                    lines=3,
+                )
+            
+            # ==================== Training Tab ====================
+            with gr.Tab(t("training.tab_train_lora")):
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        gr.HTML(f"<h3>📊 {t('training.train_section_tensors')}</h3>")
+                        
+                        gr.Markdown("""
+                        Select the directory containing preprocessed tensor files (`.pt` files).
+                        These are created in the "Dataset Builder" tab using the "Preprocess" button.
+                        """)
+                        
+                        training_tensor_dir = gr.Textbox(
+                            label=t("training.preprocessed_tensors_dir"),
+                            placeholder="./datasets/preprocessed_tensors",
+                            value="./datasets/preprocessed_tensors",
+                            info=t("training.preprocessed_tensors_info"),
+                        )
+                        
+                        load_dataset_btn = gr.Button(t("training.load_dataset_btn"), variant="secondary")
+                        
+                        training_dataset_info = gr.Textbox(
+                            label=t("training.dataset_info"),
+                            interactive=False,
+                            lines=3,
+                        )
+                    
+                    with gr.Column(scale=1):
+                        gr.HTML(f"<h3>⚙️ {t('training.train_section_lora')}</h3>")
+                        
+                        lora_rank = gr.Slider(
+                            minimum=4,
+                            maximum=256,
+                            step=4,
+                            value=64,
+                            label=t("training.lora_rank"),
+                            info=t("training.lora_rank_info"),
+                        )
+                        
+                        lora_alpha = gr.Slider(
+                            minimum=4,
+                            maximum=512,
+                            step=4,
+                            value=128,
+                            label=t("training.lora_alpha"),
+                            info=t("training.lora_alpha_info"),
+                        )
+                        
+                        lora_dropout = gr.Slider(
+                            minimum=0.0,
+                            maximum=0.5,
+                            step=0.05,
+                            value=0.1,
+                            label=t("training.lora_dropout"),
+                        )
+                
+                gr.HTML(f"<hr><h3>🎛️ {t('training.train_section_params')}</h3>")
+                
+                with gr.Row():
+                    learning_rate = gr.Number(
+                        label=t("training.learning_rate"),
+                        value=3e-4,
+                        info=t("training.learning_rate_info"),
+                    )
+                    
+                    train_epochs = gr.Slider(
+                        minimum=epoch_min,
+                        maximum=4000,
+                        step=epoch_step,
+                        value=epoch_default,
+                        label=t("training.max_epochs"),
+                    )
+                    
+                    train_batch_size = gr.Slider(
+                        minimum=1,
+                        maximum=8,
+                        step=1,
+                        value=1,
+                        label=t("training.batch_size"),
+                        info=t("training.batch_size_info"),
+                    )
+                    
+                    gradient_accumulation = gr.Slider(
+                        minimum=1,
+                        maximum=16,
+                        step=1,
+                        value=1,
+                        label=t("training.gradient_accumulation"),
+                        info=t("training.gradient_accumulation_info"),
+                    )
+                
+                with gr.Row():
+                    save_every_n_epochs = gr.Slider(
+                        minimum=50,
+                        maximum=1000,
+                        step=50,
+                        value=200,
+                        label=t("training.save_every_n_epochs"),
+                    )
+                    
+                    training_shift = gr.Slider(
+                        minimum=1.0,
+                        maximum=5.0,
+                        step=0.5,
+                        value=3.0,
+                        label=t("training.shift"),
+                        info=t("training.shift_info"),
+                    )
+                    
+                    training_seed = gr.Number(
+                        label=t("training.seed"),
+                        value=42,
+                        precision=0,
+                    )
+                
+                with gr.Row():
+                    lora_output_dir = gr.Textbox(
+                        label=t("training.output_dir"),
+                        value="./lora_output",
+                        placeholder="./lora_output",
+                        info=t("training.output_dir_info"),
+                    )
+                
+                with gr.Row():
+                    resume_checkpoint_dir = gr.Textbox(
+                        label="Resume Checkpoint (optional)",
+                        placeholder="./lora_output/checkpoints/epoch_200",
+                        info="Directory of a saved LoRA checkpoint to resume from",
+                    )
+                
+                gr.HTML("<hr>")
+                
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        start_training_btn = gr.Button(
+                            t("training.start_training_btn"),
+                            variant="primary",
+                            size="lg",
+                        )
+                    with gr.Column(scale=1):
+                        stop_training_btn = gr.Button(
+                            t("training.stop_training_btn"),
+                            variant="stop",
+                            size="lg",
+                        )
+                
+                training_progress = gr.Textbox(
+                    label=t("training.training_progress"),
+                    interactive=False,
+                    lines=2,
+                )
+                
+                with gr.Row():
+                    training_log = gr.Textbox(
+                        label=t("training.training_log"),
+                        interactive=False,
+                        lines=10,
+                        max_lines=15,
+                        scale=1,
+                    )
+                    training_loss_plot = gr.LinePlot(
+                        x="step",
+                        y="loss",
+                        title=t("training.training_loss_title"),
+                        x_title=t("training.step"),
+                        y_title=t("training.loss"),
+                        scale=1,
+                    )
+                
+                gr.HTML(f"<hr><h3>📦 {t('training.export_header')}</h3>")
+                
+                with gr.Row():
+                    export_path = gr.Textbox(
+                        label=t("training.export_path"),
+                        value="./lora_output/final_lora",
+                        placeholder="./lora_output/my_lora",
+                    )
+                    export_lora_btn = gr.Button(t("training.export_lora_btn"), variant="secondary")
+                
+                export_status = gr.Textbox(
+                    label=t("training.export_status"),
+                    interactive=False,
+                )
+    
+    # Store dataset builder state
+    dataset_builder_state = gr.State(None)
+    training_state = gr.State({"is_training": False, "should_stop": False})
+    
+    return {
+        # Dataset Builder - Load or Scan
+        "load_json_path": load_json_path,
+        "load_json_btn": load_json_btn,
+        "load_json_status": load_json_status,
+        "audio_directory": audio_directory,
+        "scan_btn": scan_btn,
+        "scan_status": scan_status,
+        "audio_files_table": audio_files_table,
+        "dataset_name": dataset_name,
+        "all_instrumental": all_instrumental,
+        "format_lyrics": format_lyrics,
+        "transcribe_lyrics": transcribe_lyrics,
+        "custom_tag": custom_tag,
+        "tag_position": tag_position,
+        "skip_metas": skip_metas,
+        "only_unlabeled": only_unlabeled,
+        "auto_label_btn": auto_label_btn,
+        "label_progress": label_progress,
+        "sample_selector": sample_selector,
+        "preview_audio": preview_audio,
+        "preview_filename": preview_filename,
+        "edit_caption": edit_caption,
+        "edit_genre": edit_genre,
+        "prompt_override": prompt_override,
+        "genre_ratio": genre_ratio,
+        "edit_lyrics": edit_lyrics,
+        "raw_lyrics_display": raw_lyrics_display,
+        "has_raw_lyrics_state": has_raw_lyrics_state,
+        "edit_bpm": edit_bpm,
+        "edit_keyscale": edit_keyscale,
+        "edit_timesig": edit_timesig,
+        "edit_duration": edit_duration,
+        "edit_language": edit_language,
+        "edit_instrumental": edit_instrumental,
+        "save_edit_btn": save_edit_btn,
+        "edit_status": edit_status,
+        "save_path": save_path,
+        "save_dataset_btn": save_dataset_btn,
+        "save_status": save_status,
+        # Preprocessing
+        "load_existing_dataset_path": load_existing_dataset_path,
+        "load_existing_dataset_btn": load_existing_dataset_btn,
+        "load_existing_status": load_existing_status,
+        "preprocess_output_dir": preprocess_output_dir,
+        "preprocess_btn": preprocess_btn,
+        "preprocess_progress": preprocess_progress,
+        "dataset_builder_state": dataset_builder_state,
+        # Training
+        "training_tensor_dir": training_tensor_dir,
+        "load_dataset_btn": load_dataset_btn,
+        "training_dataset_info": training_dataset_info,
+        "lora_rank": lora_rank,
+        "lora_alpha": lora_alpha,
+        "lora_dropout": lora_dropout,
+        "learning_rate": learning_rate,
+        "train_epochs": train_epochs,
+        "train_batch_size": train_batch_size,
+        "gradient_accumulation": gradient_accumulation,
+        "save_every_n_epochs": save_every_n_epochs,
+        "training_shift": training_shift,
+        "training_seed": training_seed,
+        "lora_output_dir": lora_output_dir,
+        "resume_checkpoint_dir": resume_checkpoint_dir,
+        "start_training_btn": start_training_btn,
+        "stop_training_btn": stop_training_btn,
+        "training_progress": training_progress,
+        "training_log": training_log,
+        "training_loss_plot": training_loss_plot,
+        "export_path": export_path,
+        "export_lora_btn": export_lora_btn,
+        "export_status": export_status,
+        "training_state": training_state,
+    }

acestep/inference.py

@@ -0,0 +1,1310 @@
+"""
+ACE-Step Inference API Module
+
+This module provides a standardized inference interface for music generation,
+designed for third-party integration. It offers both a simplified API and
+backward-compatible Gradio UI support.
+"""
+
+import math
+import os
+import tempfile
+import shutil
+import subprocess
+import sys
+from typing import Optional, Union, List, Dict, Any, Tuple
+from dataclasses import dataclass, field, asdict
+from loguru import logger
+
+from acestep.audio_utils import AudioSaver, generate_uuid_from_params, is_audio_silent
+from acestep.constants import TASK_INSTRUCTIONS
+from acestep.gpu_config import get_gpu_config
+
+# HuggingFace Space environment detection
+IS_HUGGINGFACE_SPACE = os.environ.get("SPACE_ID") is not None
+
+def _get_spaces_gpu_decorator(duration=180):
+    """
+    Get the @spaces.GPU decorator if running in HuggingFace Space environment.
+    Returns identity decorator if not in Space environment.
+    """
+    if IS_HUGGINGFACE_SPACE:
+        try:
+            import spaces
+            return spaces.GPU(duration=duration)
+        except ImportError:
+            logger.warning("spaces package not found, GPU decorator disabled")
+            return lambda func: func
+    return lambda func: func
+
+
+@dataclass
+class GenerationParams:
+    """Configuration for music generation parameters.
+    
+    Attributes:
+        # Text Inputs
+        caption: A short text prompt describing the desired music (main prompt). < 512 characters
+        lyrics: Lyrics for the music. Use "[Instrumental]" for instrumental songs. < 4096 characters
+        instrumental: If True, generate instrumental music regardless of lyrics.
+        
+        # Music Metadata
+        bpm: BPM (beats per minute), e.g., 120. Set to None for automatic estimation. 30 ~ 300
+        keyscale: Musical key (e.g., "C Major", "Am"). Leave empty for auto-detection. A-G, #/♭, major/minor
+        timesignature: Time signature (2 for '2/4', 3 for '3/4', 4 for '4/4', 6 for '6/8'). Leave empty for auto-detection.
+        vocal_language: Language code for vocals, e.g., "en", "zh", "ja", or "unknown". see acestep/constants.py:VALID_LANGUAGES
+        duration: Target audio length in seconds. If <0 or None, model chooses automatically. 10 ~ 600
+        
+        # Generation Parameters
+        inference_steps: Number of diffusion steps (e.g., 8 for turbo, 32–100 for base model).
+        guidance_scale: CFG (classifier-free guidance) strength. Higher means following the prompt more strictly. Only support for non-turbo model.
+        seed: Integer seed for reproducibility. -1 means use random seed each time.
+        
+        # Advanced DiT Parameters
+        use_adg: Whether to use Adaptive Dual Guidance (only works for base model).
+        cfg_interval_start: Start ratio (0.0–1.0) to apply CFG.
+        cfg_interval_end: End ratio (0.0–1.0) to apply CFG.
+        shift: Timestep shift factor (default 1.0). When != 1.0, applies t = shift * t / (1 + (shift - 1) * t) to timesteps.
+        
+        # Task-Specific Parameters
+        task_type: Type of generation task. One of: "text2music", "cover", "repaint", "lego", "extract", "complete".
+        reference_audio: Path to a reference audio file for style transfer or cover tasks.
+        src_audio: Path to a source audio file for audio-to-audio tasks.
+        audio_codes: Audio semantic codes as a string (advanced use, for code-control generation).
+        repainting_start: For repaint/lego tasks: start time in seconds for region to repaint.
+        repainting_end: For repaint/lego tasks: end time in seconds for region to repaint (-1 for until end).
+        audio_cover_strength: Strength of reference audio/codes influence (range 0.0–1.0). set smaller (0.2) for style transfer tasks.
+        instruction: Optional task instruction prompt. If empty, auto-generated by system.
+        
+        # 5Hz Language Model Parameters for CoT reasoning
+        thinking: If True, enable 5Hz Language Model "Chain-of-Thought" reasoning for semantic/music metadata and codes.
+        lm_temperature: Sampling temperature for the LLM (0.0–2.0). Higher = more creative/varied results.
+        lm_cfg_scale: Classifier-free guidance scale for the LLM.
+        lm_top_k: LLM top-k sampling (0 = disabled).
+        lm_top_p: LLM top-p nucleus sampling (1.0 = disabled).
+        lm_negative_prompt: Negative prompt to use for LLM (for control).
+        use_cot_metas: Whether to let LLM generate music metadata via CoT reasoning.
+        use_cot_caption: Whether to let LLM rewrite or format the input caption via CoT reasoning.
+        use_cot_language: Whether to let LLM detect vocal language via CoT.
+    """
+    # Required Inputs
+    task_type: str = "text2music"
+    instruction: str = "Fill the audio semantic mask based on the given conditions:"
+
+    # Audio Uploads
+    reference_audio: Optional[str] = None
+    src_audio: Optional[str] = None
+
+    # LM Codes Hints
+    audio_codes: str = ""
+
+    # Text Inputs
+    caption: str = ""
+    lyrics: str = ""
+    instrumental: bool = False
+
+    # Metadata
+    vocal_language: str = "unknown"
+    bpm: Optional[int] = None
+    keyscale: str = ""
+    timesignature: str = ""
+    duration: float = -1.0
+
+    # Advanced Settings
+    inference_steps: int = 8
+    seed: int = -1
+    guidance_scale: float = 7.0
+    use_adg: bool = False
+    cfg_interval_start: float = 0.0
+    cfg_interval_end: float = 1.0
+    shift: float = 1.0
+    infer_method: str = "ode"  # "ode" or "sde" - diffusion inference method
+    # Custom timesteps (parsed from string like "0.97,0.76,0.615,0.5,0.395,0.28,0.18,0.085,0")
+    # If provided, overrides inference_steps and shift
+    timesteps: Optional[List[float]] = None
+
+    repainting_start: float = 0.0
+    repainting_end: float = -1
+    audio_cover_strength: float = 1.0
+
+    # 5Hz Language Model Parameters
+    thinking: bool = True
+    lm_temperature: float = 0.85
+    lm_cfg_scale: float = 2.0
+    lm_top_k: int = 0
+    lm_top_p: float = 0.9
+    lm_negative_prompt: str = "NO USER INPUT"
+    use_cot_metas: bool = True
+    use_cot_caption: bool = True
+    use_cot_lyrics: bool = False  # TODO: not used yet
+    use_cot_language: bool = True
+    use_constrained_decoding: bool = True
+
+    cot_bpm: Optional[int] = None
+    cot_keyscale: str = ""
+    cot_timesignature: str = ""
+    cot_duration: Optional[float] = None
+    cot_vocal_language: str = "unknown"
+    cot_caption: str = ""
+    cot_lyrics: str = ""
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert config to dictionary for JSON serialization."""
+        return asdict(self)
+
+
+@dataclass
+class GenerationConfig:
+    """Configuration for music generation.
+    
+    Attributes:
+        batch_size: Number of audio samples to generate
+        allow_lm_batch: Whether to allow batch processing in LM
+        use_random_seed: Whether to use random seed
+        seeds: Seed(s) for batch generation. Can be:
+            - None: Use random seeds (when use_random_seed=True) or params.seed (when use_random_seed=False)
+            - List[int]: List of seeds, will be padded with random seeds if fewer than batch_size
+            - int: Single seed value (will be converted to list and padded)
+        lm_batch_chunk_size: Batch chunk size for LM processing
+        constrained_decoding_debug: Whether to enable constrained decoding debug
+        audio_format: Output audio format, one of "mp3", "wav", "flac". Default: "flac"
+    """
+    batch_size: int = 2
+    allow_lm_batch: bool = False
+    use_random_seed: bool = True
+    seeds: Optional[List[int]] = None
+    lm_batch_chunk_size: int = 8
+    constrained_decoding_debug: bool = False
+    audio_format: str = "flac"  # Default to FLAC for fast saving
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert config to dictionary for JSON serialization."""
+        return asdict(self)
+
+
+@dataclass
+class GenerationResult:
+    """Result of music generation.
+    
+    Attributes:
+        # Audio Outputs
+        audios: List of audio dictionaries with paths, keys, params
+        status_message: Status message from generation
+        extra_outputs: Extra outputs from generation
+        success: Whether generation completed successfully
+        error: Error message if generation failed
+    """
+
+    # Audio Outputs
+    audios: List[Dict[str, Any]] = field(default_factory=list)
+    # Generation Information
+    status_message: str = ""
+    extra_outputs: Dict[str, Any] = field(default_factory=dict)
+    # Success Status
+    success: bool = True
+    error: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert result to dictionary for JSON serialization."""
+        return asdict(self)
+
+
+@dataclass
+class UnderstandResult:
+    """Result of music understanding from audio codes.
+    
+    Attributes:
+        # Metadata Fields
+        caption: Generated caption describing the music
+        lyrics: Generated or extracted lyrics
+        bpm: Beats per minute (None if not detected)
+        duration: Duration in seconds (None if not detected)
+        keyscale: Musical key (e.g., "C Major")
+        language: Vocal language code (e.g., "en", "zh")
+        timesignature: Time signature (e.g., "4/4")
+        
+        # Status
+        status_message: Status message from understanding
+        success: Whether understanding completed successfully
+        error: Error message if understanding failed
+    """
+    # Metadata Fields
+    caption: str = ""
+    lyrics: str = ""
+    bpm: Optional[int] = None
+    duration: Optional[float] = None
+    keyscale: str = ""
+    language: str = ""
+    timesignature: str = ""
+    
+    # Status
+    status_message: str = ""
+    success: bool = True
+    error: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert result to dictionary for JSON serialization."""
+        return asdict(self)
+
+
+def _update_metadata_from_lm(
+    metadata: Dict[str, Any],
+    bpm: Optional[int],
+    key_scale: str,
+    time_signature: str,
+    audio_duration: Optional[float],
+    vocal_language: str,
+    caption: str,
+    lyrics: str,
+) -> Tuple[Optional[int], str, str, Optional[float], str, str, str]:
+    """Update metadata fields from LM output if not provided by user."""
+
+    if bpm is None and metadata.get('bpm'):
+        bpm_value = metadata.get('bpm')
+        if bpm_value not in ["N/A", ""]:
+            try:
+                bpm = int(bpm_value)
+            except (ValueError, TypeError):
+                pass
+
+    if not key_scale and metadata.get('keyscale'):
+        key_scale_value = metadata.get('keyscale', metadata.get('key_scale', ""))
+        if key_scale_value != "N/A":
+            key_scale = key_scale_value
+
+    if not time_signature and metadata.get('timesignature'):
+        time_signature_value = metadata.get('timesignature', metadata.get('time_signature', ""))
+        if time_signature_value != "N/A":
+            time_signature = time_signature_value
+
+    if audio_duration is None or audio_duration <= 0:
+        audio_duration_value = metadata.get('duration', -1)
+        if audio_duration_value not in ["N/A", ""]:
+            try:
+                audio_duration = float(audio_duration_value)
+            except (ValueError, TypeError):
+                pass
+
+    if not vocal_language and metadata.get('vocal_language'):
+        vocal_language = metadata.get('vocal_language')
+    if not caption and metadata.get('caption'):
+        caption = metadata.get('caption')
+    if not lyrics and metadata.get('lyrics'):
+        lyrics = metadata.get('lyrics')
+    return bpm, key_scale, time_signature, audio_duration, vocal_language, caption, lyrics
+
+
+@_get_spaces_gpu_decorator(duration=180)
+def generate_music(
+    dit_handler,
+    llm_handler,
+    params: GenerationParams,
+    config: GenerationConfig,
+    save_dir: Optional[str] = None,
+    progress=None,
+) -> GenerationResult:
+    """Generate music using ACE-Step model with optional LM reasoning.
+    
+    Args:
+        dit_handler: Initialized DiT model handler (AceStepHandler instance)
+        llm_handler: Initialized LLM handler (LLMHandler instance)
+        params: Generation parameters (GenerationParams instance)
+        config: Generation configuration (GenerationConfig instance)
+        
+    Returns:
+        GenerationResult with generated audio files and metadata
+    """
+    try:
+        # Phase 1: LM-based metadata and code generation (if enabled)
+        audio_code_string_to_use = params.audio_codes
+        lm_generated_metadata = None
+        lm_generated_audio_codes_list = []
+        lm_total_time_costs = {
+            "phase1_time": 0.0,
+            "phase2_time": 0.0,
+            "total_time": 0.0,
+        }
+
+        # Extract mutable copies of metadata (will be updated by LM if needed)
+        bpm = params.bpm
+        key_scale = params.keyscale
+        time_signature = params.timesignature
+        audio_duration = params.duration
+        dit_input_caption = params.caption
+        dit_input_vocal_language = params.vocal_language
+        dit_input_lyrics = params.lyrics
+        # Determine if we need to generate audio codes
+        # If user has provided audio_codes, we don't need to generate them
+        # Otherwise, check if we need audio codes (lm_dit mode) or just metas (dit mode)
+        user_provided_audio_codes = bool(params.audio_codes and str(params.audio_codes).strip())
+
+        # Safety: cover task without any source audio or codes produces silence.
+        if params.task_type == "cover":
+            no_src_audio = not (params.reference_audio or params.src_audio)
+            if no_src_audio and not user_provided_audio_codes:
+                logger.warning("Cover task requested without source audio or audio codes. Falling back to text2music.")
+                params.task_type = "text2music"
+                if params.instruction == TASK_INSTRUCTIONS.get("cover"):
+                    params.instruction = TASK_INSTRUCTIONS.get("text2music", params.instruction)
+
+        # Determine infer_type: use "llm_dit" if we need audio codes, "dit" if only metas needed
+        # For now, we use "llm_dit" if batch mode or if user hasn't provided codes
+        # Use "dit" if user has provided codes (only need metas) or if explicitly only need metas
+        # Note: This logic can be refined based on specific requirements
+        need_audio_codes = not user_provided_audio_codes
+
+        # Determine if we should use chunk-based LM generation (always use chunks for consistency)
+        # Determine actual batch size for chunk processing
+        actual_batch_size = config.batch_size if config.batch_size is not None else 1
+
+        # Prepare seeds for batch generation
+        # Use config.seed if provided, otherwise fallback to params.seed
+        # Convert config.seed (None, int, or List[int]) to format that prepare_seeds accepts
+        seed_for_generation = ""
+        # Original code (commented out because it crashes on int seeds):
+        # if config.seeds is not None and len(config.seeds) > 0:
+        #     if isinstance(config.seeds, list):
+        #         # Convert List[int] to comma-separated string
+        #         seed_for_generation = ",".join(str(s) for s in config.seeds)
+
+        if config.seeds is not None:
+            if isinstance(config.seeds, list) and len(config.seeds) > 0:
+                # Convert List[int] to comma-separated string
+                seed_for_generation = ",".join(str(s) for s in config.seeds)
+            elif isinstance(config.seeds, int):
+                # Fix: Explicitly handle single integer seeds by converting to string.
+                # Previously, this would crash because 'len()' was called on an int.
+                seed_for_generation = str(config.seeds)
+
+        # Use dit_handler.prepare_seeds to handle seed list generation and padding
+        # This will handle all the logic: padding with random seeds if needed, etc.
+        actual_seed_list, _ = dit_handler.prepare_seeds(actual_batch_size, seed_for_generation, config.use_random_seed)
+
+        # LM-based Chain-of-Thought reasoning
+        # Skip LM for cover/repaint tasks - these tasks use reference/src audio directly
+        # and don't need LM to generate audio codes
+        skip_lm_tasks = {"cover", "repaint"}
+        
+        # Determine if we should use LLM
+        # LLM is needed for:
+        # 1. thinking=True: generate audio codes via LM
+        # 2. use_cot_caption=True: enhance/generate caption via CoT
+        # 3. use_cot_language=True: detect vocal language via CoT
+        # 4. use_cot_metas=True: fill missing metadata via CoT
+        need_lm_for_cot = params.use_cot_caption or params.use_cot_language or params.use_cot_metas
+        use_lm = (params.thinking or need_lm_for_cot) and llm_handler is not None and llm_handler.llm_initialized and params.task_type not in skip_lm_tasks
+        lm_status = []
+        
+        if params.task_type in skip_lm_tasks:
+            logger.info(f"Skipping LM for task_type='{params.task_type}' - using DiT directly")
+
+        logger.info(f"[generate_music] LLM usage decision: thinking={params.thinking}, "
+                   f"use_cot_caption={params.use_cot_caption}, use_cot_language={params.use_cot_language}, "
+                   f"use_cot_metas={params.use_cot_metas}, need_lm_for_cot={need_lm_for_cot}, "
+                   f"llm_initialized={llm_handler.llm_initialized if llm_handler else False}, use_lm={use_lm}")
+
+        def _infer_audio_duration_seconds(audio_path: str) -> Optional[float]:
+            """Best-effort duration inference for common audio formats."""
+            if not audio_path:
+                return None
+            # Try torchaudio (supports more formats when ffmpeg backend is available)
+            try:
+                import torchaudio
+                info = torchaudio.info(audio_path)
+                if info and info.num_frames and info.sample_rate:
+                    return float(info.num_frames) / float(info.sample_rate)
+            except Exception:
+                pass
+            # Try soundfile (fast for wav/flac)
+            try:
+                import soundfile as sf
+                info = sf.info(audio_path)
+                if info and info.frames and info.samplerate:
+                    return float(info.frames) / float(info.samplerate)
+            except Exception:
+                pass
+            # macOS fallback: use afinfo for m4a/aac
+            if sys.platform == "darwin" and shutil.which("afinfo"):
+                try:
+                    result = subprocess.run(
+                        ["afinfo", audio_path],
+                        check=False,
+                        capture_output=True,
+                        text=True,
+                    )
+                    if result.stdout:
+                        for line in result.stdout.splitlines():
+                            if "duration:" in line:
+                                # Example: "duration:  183.165s"
+                                parts = line.strip().split()
+                                for p in parts:
+                                    if p.endswith("s"):
+                                        try:
+                                            return float(p.rstrip("s"))
+                                        except ValueError:
+                                            continue
+                except Exception:
+                    pass
+            return None
+
+        # Clamp duration and batch size to GPU limits (applies to non-Gradio callers too)
+        try:
+            # If duration not provided, try to infer from source audio to enable safe clamping.
+            if (audio_duration is None or float(audio_duration) <= 0) and (params.src_audio or params.reference_audio):
+                audio_path = params.src_audio or params.reference_audio
+                try:
+                    inferred = _infer_audio_duration_seconds(audio_path)
+                    if inferred and inferred > 0:
+                        audio_duration = inferred
+                        params.duration = inferred
+                        logger.info(f"[generate_music] Inferred duration from audio file: {inferred:.2f}s")
+                except Exception as e:
+                    logger.warning(f"[generate_music] Failed to infer duration from audio file: {e}")
+
+            gpu_config = get_gpu_config()
+            max_duration = gpu_config.max_duration_with_lm if use_lm else gpu_config.max_duration_without_lm
+            if audio_duration is not None and float(audio_duration) > 0 and float(audio_duration) > max_duration:
+                logger.warning(f"[generate_music] Duration {audio_duration}s exceeds GPU limit {max_duration}s. Clamping.")
+                audio_duration = float(max_duration)
+                params.duration = float(max_duration)
+
+            max_batch = gpu_config.max_batch_size_with_lm if use_lm else gpu_config.max_batch_size_without_lm
+            if config.batch_size is not None and config.batch_size > max_batch:
+                logger.warning(f"[generate_music] Batch size {config.batch_size} exceeds GPU limit {max_batch}. Clamping.")
+                config.batch_size = max_batch
+
+            # Extra safety for MPS: large durations can OOM with batch > 1
+            if (
+                hasattr(dit_handler, "device")
+                and dit_handler.device == "mps"
+                and audio_duration is not None
+                and float(audio_duration) > 180
+                and config.batch_size is not None
+                and config.batch_size > 1
+            ):
+                logger.warning("[generate_music] MPS with long duration detected; reducing batch size to 1 to avoid OOM.")
+                config.batch_size = 1
+        except Exception as e:
+            logger.warning(f"[generate_music] Failed to clamp duration/batch to GPU limits: {e}")
+        
+        if use_lm:
+            # Convert sampling parameters - handle None values safely
+            top_k_value = None if not params.lm_top_k or params.lm_top_k == 0 else int(params.lm_top_k)
+            top_p_value = None if not params.lm_top_p or params.lm_top_p >= 1.0 else params.lm_top_p
+
+            # Build user_metadata from user-provided values
+            user_metadata = {}
+            if bpm is not None:
+                try:
+                    bpm_value = float(bpm)
+                    if bpm_value > 0:
+                        user_metadata['bpm'] = int(bpm_value)
+                except (ValueError, TypeError):
+                    pass
+
+            if key_scale and key_scale.strip():
+                key_scale_clean = key_scale.strip()
+                if key_scale_clean.lower() not in ["n/a", ""]:
+                    user_metadata['keyscale'] = key_scale_clean
+
+            if time_signature and time_signature.strip():
+                time_sig_clean = time_signature.strip()
+                if time_sig_clean.lower() not in ["n/a", ""]:
+                    user_metadata['timesignature'] = time_sig_clean
+
+            if audio_duration is not None:
+                try:
+                    duration_value = float(audio_duration)
+                    if duration_value > 0:
+                        user_metadata['duration'] = int(duration_value)
+                except (ValueError, TypeError):
+                    pass
+
+            user_metadata_to_pass = user_metadata if user_metadata else None
+
+            # Determine infer_type based on whether we need audio codes
+            # - "llm_dit": generates both metas and audio codes (two-phase internally)
+            # - "dit": generates only metas (single phase)
+            infer_type = "llm_dit" if need_audio_codes and params.thinking else "dit"
+
+            # Use chunk size from config, or default to batch_size if not set
+            max_inference_batch_size = int(config.lm_batch_chunk_size) if config.lm_batch_chunk_size > 0 else actual_batch_size
+            num_chunks = math.ceil(actual_batch_size / max_inference_batch_size)
+
+            all_metadata_list = []
+            all_audio_codes_list = []
+
+            for chunk_idx in range(num_chunks):
+                chunk_start = chunk_idx * max_inference_batch_size
+                chunk_end = min(chunk_start + max_inference_batch_size, actual_batch_size)
+                chunk_size = chunk_end - chunk_start
+                chunk_seeds = actual_seed_list[chunk_start:chunk_end] if chunk_start < len(actual_seed_list) else None
+
+                logger.info(f"LM chunk {chunk_idx+1}/{num_chunks} (infer_type={infer_type}) "
+                            f"(size: {chunk_size}, seeds: {chunk_seeds})")
+
+                # Use the determined infer_type
+                # - "llm_dit" will internally run two phases (metas + codes)
+                # - "dit" will only run phase 1 (metas only)
+                result = llm_handler.generate_with_stop_condition(
+                    caption=params.caption or "",
+                    lyrics=params.lyrics or "",
+                    infer_type=infer_type,
+                    temperature=params.lm_temperature,
+                    cfg_scale=params.lm_cfg_scale,
+                    negative_prompt=params.lm_negative_prompt,
+                    top_k=top_k_value,
+                    top_p=top_p_value,
+                    target_duration=audio_duration,  # Pass duration to limit audio codes generation
+                    user_metadata=user_metadata_to_pass,
+                    use_cot_caption=params.use_cot_caption,
+                    use_cot_language=params.use_cot_language,
+                    use_cot_metas=params.use_cot_metas,
+                    use_constrained_decoding=params.use_constrained_decoding,
+                    constrained_decoding_debug=config.constrained_decoding_debug,
+                    batch_size=chunk_size,
+                    seeds=chunk_seeds,
+                    progress=progress,
+                )
+
+                # Check if LM generation failed
+                if not result.get("success", False):
+                    error_msg = result.get("error", "Unknown LM error")
+                    lm_status.append(f"❌ LM Error: {error_msg}")
+                    # Return early with error
+                    return GenerationResult(
+                        audios=[],
+                        status_message=f"❌ LM generation failed: {error_msg}",
+                        extra_outputs={},
+                        success=False,
+                        error=error_msg,
+                    )
+
+                # Extract metadata and audio_codes from result dict
+                if chunk_size > 1:
+                    metadata_list = result.get("metadata", [])
+                    audio_codes_list = result.get("audio_codes", [])
+                    all_metadata_list.extend(metadata_list)
+                    all_audio_codes_list.extend(audio_codes_list)
+                else:
+                    metadata = result.get("metadata", {})
+                    audio_codes = result.get("audio_codes", "")
+                    all_metadata_list.append(metadata)
+                    all_audio_codes_list.append(audio_codes)
+
+                # Collect time costs from LM extra_outputs
+                lm_extra = result.get("extra_outputs", {})
+                lm_chunk_time_costs = lm_extra.get("time_costs", {})
+                if lm_chunk_time_costs:
+                    # Accumulate time costs from all chunks
+                    for key in ["phase1_time", "phase2_time", "total_time"]:
+                        if key in lm_chunk_time_costs:
+                            lm_total_time_costs[key] += lm_chunk_time_costs[key]
+
+                    time_str = ", ".join([f"{k}: {v:.2f}s" for k, v in lm_chunk_time_costs.items()])
+                    lm_status.append(f"✅ LM chunk {chunk_idx+1}: {time_str}")
+
+            lm_generated_metadata = all_metadata_list[0] if all_metadata_list else None
+            lm_generated_audio_codes_list = all_audio_codes_list
+
+            # Set audio_code_string_to_use based on infer_type
+            if infer_type == "llm_dit":
+                # If batch mode, use list; otherwise use single string
+                if actual_batch_size > 1:
+                    audio_code_string_to_use = all_audio_codes_list
+                else:
+                    audio_code_string_to_use = all_audio_codes_list[0] if all_audio_codes_list else ""
+            else:
+                # For "dit" mode, keep user-provided codes or empty
+                audio_code_string_to_use = params.audio_codes
+
+            # Update metadata from LM if not provided by user
+            if lm_generated_metadata:
+                bpm, key_scale, time_signature, audio_duration, vocal_language, caption, lyrics = _update_metadata_from_lm(
+                    metadata=lm_generated_metadata,
+                    bpm=bpm,
+                    key_scale=key_scale,
+                    time_signature=time_signature,
+                    audio_duration=audio_duration,
+                    vocal_language=dit_input_vocal_language,
+                    caption=dit_input_caption,
+                    lyrics=dit_input_lyrics)
+                if not params.bpm:
+                    params.cot_bpm = bpm
+                if not params.keyscale:
+                    params.cot_keyscale = key_scale
+                if not params.timesignature:
+                    params.cot_timesignature = time_signature
+                if not params.duration:
+                    params.cot_duration = audio_duration
+                if not params.vocal_language:
+                    params.cot_vocal_language = vocal_language
+                if not params.caption:
+                    params.cot_caption = caption
+                if not params.lyrics:
+                    params.cot_lyrics = lyrics
+                dit_input_lyrics = lyrics
+
+            # set cot caption and language if needed
+            if params.use_cot_caption:
+                dit_input_caption = lm_generated_metadata.get("caption", dit_input_caption)
+            if params.use_cot_language:
+                dit_input_vocal_language = lm_generated_metadata.get("vocal_language", dit_input_vocal_language)
+
+        # Phase 2: DiT music generation
+        # Use seed_for_generation (from config.seed or params.seed) instead of params.seed for actual generation
+        result = dit_handler.generate_music(
+            captions=dit_input_caption,
+            lyrics=dit_input_lyrics,
+            bpm=bpm,
+            key_scale=key_scale,
+            time_signature=time_signature,
+            vocal_language=dit_input_vocal_language,
+            inference_steps=params.inference_steps,
+            guidance_scale=params.guidance_scale,
+            use_random_seed=config.use_random_seed,
+            seed=seed_for_generation,  # Use config.seed (or params.seed fallback) instead of params.seed directly
+            reference_audio=params.reference_audio,
+            audio_duration=audio_duration,
+            batch_size=config.batch_size if config.batch_size is not None else 1,
+            src_audio=params.src_audio,
+            audio_code_string=audio_code_string_to_use,
+            repainting_start=params.repainting_start,
+            repainting_end=params.repainting_end,
+            instruction=params.instruction,
+            audio_cover_strength=params.audio_cover_strength,
+            task_type=params.task_type,
+            use_adg=params.use_adg,
+            cfg_interval_start=params.cfg_interval_start,
+            cfg_interval_end=params.cfg_interval_end,
+            shift=params.shift,
+            infer_method=params.infer_method,
+            timesteps=params.timesteps,
+            progress=progress,
+        )
+
+        # Check if generation failed
+        if not result.get("success", False):
+            return GenerationResult(
+                audios=[],
+                status_message=result.get("status_message", ""),
+                extra_outputs={},
+                success=False,
+                error=result.get("error"),
+            )
+
+        # Extract results from dit_handler.generate_music dict
+        dit_audios = result.get("audios", [])
+        status_message = result.get("status_message", "")
+        dit_extra_outputs = result.get("extra_outputs", {})
+
+        # Use the seed list already prepared above (from config.seed or params.seed fallback)
+        # actual_seed_list was computed earlier using dit_handler.prepare_seeds
+        seed_list = actual_seed_list
+
+        # Get base params dictionary
+        base_params_dict = params.to_dict()
+
+        # Save audio files using AudioSaver (format from config)
+        audio_format = config.audio_format if config.audio_format else "flac"
+        audio_saver = AudioSaver(default_format=audio_format)
+
+        # Use handler's temp_dir for saving files
+        if save_dir is not None:
+            os.makedirs(save_dir, exist_ok=True)
+
+        # Build audios list for GenerationResult with params and save files
+        # Audio saving and UUID generation handled here, outside of handler
+        audios = []
+        silent_warnings = []
+        for idx, dit_audio in enumerate(dit_audios):
+            # Create a copy of params dict for this audio
+            audio_params = base_params_dict.copy()
+
+            # Update audio-specific values
+            audio_params["seed"] = seed_list[idx] if idx < len(seed_list) else None
+
+            # Add audio codes if batch mode
+            if lm_generated_audio_codes_list and idx < len(lm_generated_audio_codes_list):
+                audio_params["audio_codes"] = lm_generated_audio_codes_list[idx]
+
+            # Get audio tensor and metadata
+            audio_tensor = dit_audio.get("tensor")
+            sample_rate = dit_audio.get("sample_rate", 48000)
+
+            # Generate UUID for this audio (moved from handler)
+            batch_seed = seed_list[idx] if idx < len(seed_list) else seed_list[0] if seed_list else -1
+            audio_code_str = lm_generated_audio_codes_list[idx] if (
+                lm_generated_audio_codes_list and idx < len(lm_generated_audio_codes_list)) else audio_code_string_to_use
+            if isinstance(audio_code_str, list):
+                audio_code_str = audio_code_str[idx] if idx < len(audio_code_str) else ""
+
+            audio_key = generate_uuid_from_params(audio_params)
+
+            silent_check = False
+            if audio_tensor is not None:
+                silent_check, rms_val, peak_val = is_audio_silent(audio_tensor, channels_first=True)
+                if silent_check:
+                    logger.warning(
+                        f"[generate_music] Silent output detected (idx={idx}, RMS={rms_val:.2e}, peak={peak_val:.2e}). "
+                        "Likely cause: LLM backend returned empty conditioning, or incompatible torch/triton/flash-attn. "
+                        "Suggest running with --backend pt."
+                    )
+                    silent_warnings.append(
+                        f"Output {idx + 1}: silent or near-silent (RMS≈{rms_val:.2e}). "
+                        "Likely causes: LLM backend failure, incompatible torch/triton/flash-attn, or CPU/fallback path. "
+                        "Try running with --backend pt."
+                    )
+
+            audio_path = None
+            if audio_tensor is not None and save_dir is not None and not silent_check:
+                try:
+                    audio_file = os.path.join(save_dir, f"{audio_key}.{audio_format}")
+                    audio_path = audio_saver.save_audio(audio_tensor,
+                                                        audio_file,
+                                                        sample_rate=sample_rate,
+                                                        format=audio_format,
+                                                        channels_first=True)
+                except Exception as e:
+                    logger.error(f"[generate_music] Failed to save audio file: {e}")
+                    audio_path = ""
+
+            audio_dict = {
+                "path": audio_path or "",
+                "tensor": audio_tensor,
+                "key": audio_key,
+                "sample_rate": sample_rate,
+                "params": audio_params,
+                "silent": silent_check,
+            }
+
+            audios.append(audio_dict)
+
+        # Merge extra_outputs: include dit_extra_outputs (latents, masks) and add LM metadata
+        extra_outputs = dit_extra_outputs.copy()
+        extra_outputs["lm_metadata"] = lm_generated_metadata
+
+        # Merge time_costs from both LM and DiT into a unified dictionary
+        unified_time_costs = {}
+
+        # Add LM time costs (if LM was used)
+        if use_lm and lm_total_time_costs:
+            for key, value in lm_total_time_costs.items():
+                unified_time_costs[f"lm_{key}"] = value
+
+        # Add DiT time costs (if available)
+        dit_time_costs = dit_extra_outputs.get("time_costs", {})
+        if dit_time_costs:
+            for key, value in dit_time_costs.items():
+                unified_time_costs[f"dit_{key}"] = value
+
+        # Calculate total pipeline time
+        if unified_time_costs:
+            lm_total = unified_time_costs.get("lm_total_time", 0.0)
+            dit_total = unified_time_costs.get("dit_total_time_cost", 0.0)
+            unified_time_costs["pipeline_total_time"] = lm_total + dit_total
+
+        # Update extra_outputs with unified time_costs
+        extra_outputs["time_costs"] = unified_time_costs
+
+        if lm_status:
+            status_message = "\n".join(lm_status) + "\n" + status_message
+        else:
+            status_message = status_message
+        if silent_warnings:
+            status_message = "⚠️ Silent output detected:\n" + "\n".join(silent_warnings) + "\n\nSuggested fix: try running with --backend pt\n\n" + (status_message or "")
+        # Create and return GenerationResult
+        return GenerationResult(
+            audios=audios,
+            status_message=status_message,
+            extra_outputs=extra_outputs,
+            success=True,
+            error=None,
+        )
+
+    except Exception as e:
+        logger.exception("Music generation failed")
+        return GenerationResult(
+            audios=[],
+            status_message=f"Error: {str(e)}",
+            extra_outputs={},
+            success=False,
+            error=str(e),
+        )
+
+
+def understand_music(
+    llm_handler,
+    audio_codes: str,
+    temperature: float = 0.85,
+    top_k: Optional[int] = None,
+    top_p: Optional[float] = None,
+    repetition_penalty: float = 1.0,
+    use_constrained_decoding: bool = True,
+    constrained_decoding_debug: bool = False,
+) -> UnderstandResult:
+    """Understand music from audio codes using the 5Hz Language Model.
+    
+    This function analyzes audio semantic codes and generates metadata about the music,
+    including caption, lyrics, BPM, duration, key scale, language, and time signature.
+    
+    If audio_codes is empty or "NO USER INPUT", the LM will generate a sample example
+    instead of analyzing existing codes.
+    
+    Note: cfg_scale and negative_prompt are not supported in understand mode.
+    
+    Args:
+        llm_handler: Initialized LLM handler (LLMHandler instance)
+        audio_codes: String of audio code tokens (e.g., "<|audio_code_123|><|audio_code_456|>...")
+                     Use empty string or "NO USER INPUT" to generate a sample example.
+        temperature: Sampling temperature for generation (0.0-2.0). Higher = more creative.
+        top_k: Top-K sampling (None or 0 = disabled)
+        top_p: Top-P (nucleus) sampling (None or 1.0 = disabled)
+        repetition_penalty: Repetition penalty (1.0 = no penalty)
+        use_constrained_decoding: Whether to use FSM-based constrained decoding for metadata
+        constrained_decoding_debug: Whether to enable debug logging for constrained decoding
+        
+    Returns:
+        UnderstandResult with parsed metadata fields and status
+        
+    Example:
+        >>> result = understand_music(llm_handler, audio_codes="<|audio_code_123|>...")
+        >>> if result.success:
+        ...     print(f"Caption: {result.caption}")
+        ...     print(f"BPM: {result.bpm}")
+        ...     print(f"Lyrics: {result.lyrics}")
+    """
+    # Check if LLM is initialized
+    if not llm_handler.llm_initialized:
+        return UnderstandResult(
+            status_message="5Hz LM not initialized. Please initialize it first.",
+            success=False,
+            error="LLM not initialized",
+        )
+    
+    # If codes are empty, use "NO USER INPUT" to generate a sample example
+    if not audio_codes or not audio_codes.strip():
+        audio_codes = "NO USER INPUT"
+    
+    try:
+        # Call LLM understanding
+        metadata, status = llm_handler.understand_audio_from_codes(
+            audio_codes=audio_codes,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            use_constrained_decoding=use_constrained_decoding,
+            constrained_decoding_debug=constrained_decoding_debug,
+        )
+        
+        # Check if LLM returned empty metadata (error case)
+        if not metadata:
+            return UnderstandResult(
+                status_message=status or "Failed to understand audio codes",
+                success=False,
+                error=status or "Empty metadata returned",
+            )
+        
+        # Extract and convert fields
+        caption = metadata.get('caption', '')
+        lyrics = metadata.get('lyrics', '')
+        keyscale = metadata.get('keyscale', '')
+        language = metadata.get('language', metadata.get('vocal_language', ''))
+        timesignature = metadata.get('timesignature', '')
+        
+        # Convert BPM to int
+        bpm = None
+        bpm_value = metadata.get('bpm')
+        if bpm_value is not None and bpm_value != 'N/A' and bpm_value != '':
+            try:
+                bpm = int(bpm_value)
+            except (ValueError, TypeError):
+                pass
+        
+        # Convert duration to float
+        duration = None
+        duration_value = metadata.get('duration')
+        if duration_value is not None and duration_value != 'N/A' and duration_value != '':
+            try:
+                duration = float(duration_value)
+            except (ValueError, TypeError):
+                pass
+        
+        # Clean up N/A values
+        if keyscale == 'N/A':
+            keyscale = ''
+        if language == 'N/A':
+            language = ''
+        if timesignature == 'N/A':
+            timesignature = ''
+        
+        return UnderstandResult(
+            caption=caption,
+            lyrics=lyrics,
+            bpm=bpm,
+            duration=duration,
+            keyscale=keyscale,
+            language=language,
+            timesignature=timesignature,
+            status_message=status,
+            success=True,
+            error=None,
+        )
+        
+    except Exception as e:
+        logger.exception("Music understanding failed")
+        return UnderstandResult(
+            status_message=f"Error: {str(e)}",
+            success=False,
+            error=str(e),
+        )
+
+
+@dataclass
+class CreateSampleResult:
+    """Result of creating a music sample from a natural language query.
+    
+    This is used by the "Simple Mode" / "Inspiration Mode" feature where users
+    provide a natural language description and the LLM generates a complete
+    sample with caption, lyrics, and metadata.
+    
+    Attributes:
+        # Metadata Fields
+        caption: Generated detailed music description/caption
+        lyrics: Generated lyrics (or "[Instrumental]" for instrumental music)
+        bpm: Beats per minute (None if not generated)
+        duration: Duration in seconds (None if not generated)
+        keyscale: Musical key (e.g., "C Major")
+        language: Vocal language code (e.g., "en", "zh")
+        timesignature: Time signature (e.g., "4")
+        instrumental: Whether this is an instrumental piece
+        
+        # Status
+        status_message: Status message from sample creation
+        success: Whether sample creation completed successfully
+        error: Error message if sample creation failed
+    """
+    # Metadata Fields
+    caption: str = ""
+    lyrics: str = ""
+    bpm: Optional[int] = None
+    duration: Optional[float] = None
+    keyscale: str = ""
+    language: str = ""
+    timesignature: str = ""
+    instrumental: bool = False
+    
+    # Status
+    status_message: str = ""
+    success: bool = True
+    error: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert result to dictionary for JSON serialization."""
+        return asdict(self)
+
+
+def create_sample(
+    llm_handler,
+    query: str,
+    instrumental: bool = False,
+    vocal_language: Optional[str] = None,
+    temperature: float = 0.85,
+    top_k: Optional[int] = None,
+    top_p: Optional[float] = None,
+    repetition_penalty: float = 1.0,
+    use_constrained_decoding: bool = True,
+    constrained_decoding_debug: bool = False,
+) -> CreateSampleResult:
+    """Create a music sample from a natural language query using the 5Hz Language Model.
+    
+    This is the "Simple Mode" / "Inspiration Mode" feature that takes a user's natural
+    language description of music and generates a complete sample including:
+    - Detailed caption/description
+    - Lyrics (unless instrumental)
+    - Metadata (BPM, duration, key, language, time signature)
+    
+    Note: cfg_scale and negative_prompt are not supported in create_sample mode.
+    
+    Args:
+        llm_handler: Initialized LLM handler (LLMHandler instance)
+        query: User's natural language music description (e.g., "a soft Bengali love song")
+        instrumental: Whether to generate instrumental music (no vocals)
+        vocal_language: Allowed vocal language for constrained decoding (e.g., "en", "zh").
+                       If provided, the model will be constrained to generate lyrics in this language.
+                       If None or "unknown", no language constraint is applied.
+        temperature: Sampling temperature for generation (0.0-2.0). Higher = more creative.
+        top_k: Top-K sampling (None or 0 = disabled)
+        top_p: Top-P (nucleus) sampling (None or 1.0 = disabled)
+        repetition_penalty: Repetition penalty (1.0 = no penalty)
+        use_constrained_decoding: Whether to use FSM-based constrained decoding
+        constrained_decoding_debug: Whether to enable debug logging
+        
+    Returns:
+        CreateSampleResult with generated sample fields and status
+        
+    Example:
+        >>> result = create_sample(llm_handler, "a soft Bengali love song for a quiet evening", vocal_language="bn")
+        >>> if result.success:
+        ...     print(f"Caption: {result.caption}")
+        ...     print(f"Lyrics: {result.lyrics}")
+        ...     print(f"BPM: {result.bpm}")
+    """
+    # Check if LLM is initialized
+    if not llm_handler.llm_initialized:
+        return CreateSampleResult(
+            status_message="5Hz LM not initialized. Please initialize it first.",
+            success=False,
+            error="LLM not initialized",
+        )
+    
+    try:
+        # Call LLM to create sample
+        metadata, status = llm_handler.create_sample_from_query(
+            query=query,
+            instrumental=instrumental,
+            vocal_language=vocal_language,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            use_constrained_decoding=use_constrained_decoding,
+            constrained_decoding_debug=constrained_decoding_debug,
+        )
+        
+        # Check if LLM returned empty metadata (error case)
+        if not metadata:
+            return CreateSampleResult(
+                status_message=status or "Failed to create sample",
+                success=False,
+                error=status or "Empty metadata returned",
+            )
+        
+        # Extract and convert fields
+        caption = metadata.get('caption', '')
+        lyrics = metadata.get('lyrics', '')
+        keyscale = metadata.get('keyscale', '')
+        language = metadata.get('language', metadata.get('vocal_language', ''))
+        timesignature = metadata.get('timesignature', '')
+        is_instrumental = metadata.get('instrumental', instrumental)
+        
+        # Convert BPM to int
+        bpm = None
+        bpm_value = metadata.get('bpm')
+        if bpm_value is not None and bpm_value != 'N/A' and bpm_value != '':
+            try:
+                bpm = int(bpm_value)
+            except (ValueError, TypeError):
+                pass
+        
+        # Convert duration to float
+        duration = None
+        duration_value = metadata.get('duration')
+        if duration_value is not None and duration_value != 'N/A' and duration_value != '':
+            try:
+                duration = float(duration_value)
+            except (ValueError, TypeError):
+                pass
+        
+        # Clean up N/A values
+        if keyscale == 'N/A':
+            keyscale = ''
+        if language == 'N/A':
+            language = ''
+        if timesignature == 'N/A':
+            timesignature = ''
+        
+        return CreateSampleResult(
+            caption=caption,
+            lyrics=lyrics,
+            bpm=bpm,
+            duration=duration,
+            keyscale=keyscale,
+            language=language,
+            timesignature=timesignature,
+            instrumental=is_instrumental,
+            status_message=status,
+            success=True,
+            error=None,
+        )
+        
+    except Exception as e:
+        logger.exception("Sample creation failed")
+        return CreateSampleResult(
+            status_message=f"Error: {str(e)}",
+            success=False,
+            error=str(e),
+        )
+
+
+@dataclass
+class FormatSampleResult:
+    """Result of formatting user-provided caption and lyrics.
+    
+    This is used by the "Format" feature where users provide caption and lyrics,
+    and the LLM formats them into structured music metadata and an enhanced description.
+    
+    Attributes:
+        # Metadata Fields
+        caption: Enhanced/formatted music description/caption
+        lyrics: Formatted lyrics (may be same as input or reformatted)
+        bpm: Beats per minute (None if not detected)
+        duration: Duration in seconds (None if not detected)
+        keyscale: Musical key (e.g., "C Major")
+        language: Vocal language code (e.g., "en", "zh")
+        timesignature: Time signature (e.g., "4")
+        
+        # Status
+        status_message: Status message from formatting
+        success: Whether formatting completed successfully
+        error: Error message if formatting failed
+    """
+    # Metadata Fields
+    caption: str = ""
+    lyrics: str = ""
+    bpm: Optional[int] = None
+    duration: Optional[float] = None
+    keyscale: str = ""
+    language: str = ""
+    timesignature: str = ""
+    
+    # Status
+    status_message: str = ""
+    success: bool = True
+    error: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert result to dictionary for JSON serialization."""
+        return asdict(self)
+
+
+def format_sample(
+    llm_handler,
+    caption: str,
+    lyrics: str,
+    user_metadata: Optional[Dict[str, Any]] = None,
+    temperature: float = 0.85,
+    top_k: Optional[int] = None,
+    top_p: Optional[float] = None,
+    repetition_penalty: float = 1.0,
+    use_constrained_decoding: bool = True,
+    constrained_decoding_debug: bool = False,
+) -> FormatSampleResult:
+    """Format user-provided caption and lyrics using the 5Hz Language Model.
+    
+    This function takes user input (caption and lyrics) and generates structured
+    music metadata including an enhanced caption, BPM, duration, key, language,
+    and time signature.
+    
+    If user_metadata is provided, those values will be used to constrain the
+    decoding, ensuring the output matches user-specified values.
+    
+    Note: cfg_scale and negative_prompt are not supported in format mode.
+    
+    Args:
+        llm_handler: Initialized LLM handler (LLMHandler instance)
+        caption: User's caption/description (e.g., "Latin pop, reggaeton")
+        lyrics: User's lyrics with structure tags
+        user_metadata: Optional dict with user-provided metadata to constrain decoding.
+                      Supported keys: bpm, duration, keyscale, timesignature, language
+        temperature: Sampling temperature for generation (0.0-2.0). Higher = more creative.
+        top_k: Top-K sampling (None or 0 = disabled)
+        top_p: Top-P (nucleus) sampling (None or 1.0 = disabled)
+        repetition_penalty: Repetition penalty (1.0 = no penalty)
+        use_constrained_decoding: Whether to use FSM-based constrained decoding for metadata
+        constrained_decoding_debug: Whether to enable debug logging for constrained decoding
+        
+    Returns:
+        FormatSampleResult with formatted metadata fields and status
+        
+    Example:
+        >>> result = format_sample(llm_handler, "Latin pop, reggaeton", "[Verse 1]\\nHola mundo...")
+        >>> if result.success:
+        ...     print(f"Caption: {result.caption}")
+        ...     print(f"BPM: {result.bpm}")
+        ...     print(f"Lyrics: {result.lyrics}")
+    """
+    # Check if LLM is initialized
+    if not llm_handler.llm_initialized:
+        return FormatSampleResult(
+            status_message="5Hz LM not initialized. Please initialize it first.",
+            success=False,
+            error="LLM not initialized",
+        )
+    
+    try:
+        # Call LLM formatting
+        metadata, status = llm_handler.format_sample_from_input(
+            caption=caption,
+            lyrics=lyrics,
+            user_metadata=user_metadata,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            use_constrained_decoding=use_constrained_decoding,
+            constrained_decoding_debug=constrained_decoding_debug,
+        )
+        
+        # Check if LLM returned empty metadata (error case)
+        if not metadata:
+            return FormatSampleResult(
+                status_message=status or "Failed to format input",
+                success=False,
+                error=status or "Empty metadata returned",
+            )
+        
+        # Extract and convert fields
+        result_caption = metadata.get('caption', '')
+        result_lyrics = metadata.get('lyrics', lyrics)  # Fall back to input lyrics
+        keyscale = metadata.get('keyscale', '')
+        language = metadata.get('language', metadata.get('vocal_language', ''))
+        timesignature = metadata.get('timesignature', '')
+        
+        # Convert BPM to int
+        bpm = None
+        bpm_value = metadata.get('bpm')
+        if bpm_value is not None and bpm_value != 'N/A' and bpm_value != '':
+            try:
+                bpm = int(bpm_value)
+            except (ValueError, TypeError):
+                pass
+        
+        # Convert duration to float
+        duration = None
+        duration_value = metadata.get('duration')
+        if duration_value is not None and duration_value != 'N/A' and duration_value != '':
+            try:
+                duration = float(duration_value)
+            except (ValueError, TypeError):
+                pass
+        
+        # Clean up N/A values
+        if keyscale == 'N/A':
+            keyscale = ''
+        if language == 'N/A':
+            language = ''
+        if timesignature == 'N/A':
+            timesignature = ''
+        
+        return FormatSampleResult(
+            caption=result_caption,
+            lyrics=result_lyrics,
+            bpm=bpm,
+            duration=duration,
+            keyscale=keyscale,
+            language=language,
+            timesignature=timesignature,
+            status_message=status,
+            success=True,
+            error=None,
+        )
+        
+    except Exception as e:
+        logger.exception("Format sample failed")
+        return FormatSampleResult(
+            status_message=f"Error: {str(e)}",
+            success=False,
+            error=str(e),
+        )

acestep/local_cache.py

@@ -0,0 +1,129 @@
+"""Local cache module to replace Redis
+
+Uses diskcache as backend, provides Redis-compatible API.
+Supports persistent storage and TTL expiration.
+"""
+
+import json
+import os
+from typing import Any, Optional
+from threading import Lock
+
+try:
+    from diskcache import Cache
+    HAS_DISKCACHE = True
+except ImportError:
+    HAS_DISKCACHE = False
+
+
+class LocalCache:
+    """
+    Local cache implementation with Redis-compatible API.
+    Uses diskcache as backend, supports persistence and TTL.
+    """
+
+    _instance = None
+    _lock = Lock()
+
+    def __new__(cls, cache_dir: Optional[str] = None):
+        """Singleton pattern"""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+                    cls._instance._initialized = False
+        return cls._instance
+
+    def __init__(self, cache_dir: Optional[str] = None):
+        if getattr(self, '_initialized', False):
+            return
+
+        if not HAS_DISKCACHE:
+            raise ImportError(
+                "diskcache not installed. Run: pip install diskcache"
+            )
+
+        if cache_dir is None:
+            cache_dir = os.path.join(
+                os.path.dirname(os.path.dirname(__file__)),
+                ".cache",
+                "local_redis"
+            )
+
+        os.makedirs(cache_dir, exist_ok=True)
+        self._cache = Cache(cache_dir)
+        self._initialized = True
+
+    def set(self, name: str, value: Any, ex: Optional[int] = None) -> bool:
+        """
+        Set key-value pair
+
+        Args:
+            name: Key name
+            value: Value (auto-serialize dict/list)
+            ex: Expiration time (seconds)
+
+        Returns:
+            bool: Success status
+        """
+        if isinstance(value, (dict, list)):
+            value = json.dumps(value, ensure_ascii=False)
+        self._cache.set(name, value, expire=ex)
+        return True
+
+    def get(self, name: str) -> Optional[str]:
+        """Get value"""
+        return self._cache.get(name)
+
+    def delete(self, name: str) -> int:
+        """Delete key, returns number of deleted items"""
+        return 1 if self._cache.delete(name) else 0
+
+    def exists(self, name: str) -> bool:
+        """Check if key exists"""
+        return name in self._cache
+
+    def keys(self, pattern: str = "*") -> list:
+        """
+        Get list of matching keys
+        Note: Simplified implementation, only supports prefix and full matching
+        """
+        if pattern == "*":
+            return list(self._cache.iterkeys())
+
+        prefix = pattern.rstrip("*")
+        return [k for k in self._cache.iterkeys() if k.startswith(prefix)]
+
+    def expire(self, name: str, seconds: int) -> bool:
+        """Set key expiration time"""
+        value = self._cache.get(name)
+        if value is not None:
+            self._cache.set(name, value, expire=seconds)
+            return True
+        return False
+
+    def ttl(self, name: str) -> int:
+        """
+        Get remaining time to live (seconds)
+        Note: diskcache does not directly support TTL queries
+        """
+        if name in self._cache:
+            return -1  # Exists but TTL unknown
+        return -2  # Key does not exist
+
+    def close(self):
+        """Close cache connection"""
+        if hasattr(self, '_cache'):
+            self._cache.close()
+
+
+# Lazily initialized global instance
+_local_cache: Optional[LocalCache] = None
+
+
+def get_local_cache(cache_dir: Optional[str] = None) -> LocalCache:
+    """Get local cache instance"""
+    global _local_cache
+    if _local_cache is None:
+        _local_cache = LocalCache(cache_dir)
+    return _local_cache

acestep/model_downloader.py

@@ -0,0 +1,634 @@
+"""
+ACE-Step Model Downloader
+
+This module provides functionality to download models from HuggingFace Hub or ModelScope.
+It supports automatic downloading when models are not found locally,
+with intelligent fallback between download sources.
+"""
+
+import os
+import sys
+import argparse
+from typing import Optional, List, Dict, Tuple
+from pathlib import Path
+
+from loguru import logger
+
+
+# =============================================================================
+# Network Detection & Smart Download
+# =============================================================================
+
+def _can_access_google(timeout: float = 3.0) -> bool:
+    """
+    Check if Google is accessible (to determine HuggingFace vs ModelScope).
+
+    Args:
+        timeout: Connection timeout in seconds
+
+    Returns:
+        True if Google is accessible, False otherwise
+    """
+    import socket
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    try:
+        sock.settimeout(timeout)
+        sock.connect(("www.google.com", 443))
+        return True
+    except (socket.timeout, socket.error, OSError):
+        return False
+    finally:
+        sock.close()
+
+
+def _download_from_huggingface_internal(
+    repo_id: str,
+    local_dir: Path,
+    token: Optional[str] = None,
+) -> None:
+    """
+    Internal function to download from HuggingFace Hub.
+
+    Args:
+        repo_id: HuggingFace repository ID (e.g., "ACE-Step/Ace-Step1.5")
+        local_dir: Local directory to save the model
+        token: HuggingFace token for private repos (optional)
+
+    Raises:
+        Exception: If download fails
+    """
+    from huggingface_hub import snapshot_download
+
+    logger.info(f"[Model Download] Downloading from HuggingFace: {repo_id} -> {local_dir}")
+
+    snapshot_download(
+        repo_id=repo_id,
+        local_dir=str(local_dir),
+        local_dir_use_symlinks=False,
+        token=token,
+    )
+
+
+def _download_from_modelscope_internal(
+    repo_id: str,
+    local_dir: Path,
+) -> None:
+    """
+    Internal function to download from ModelScope.
+
+    Args:
+        repo_id: ModelScope repository ID (e.g., "ACE-Step/Ace-Step1.5")
+        local_dir: Local directory to save the model
+
+    Raises:
+        Exception: If download fails
+    """
+    from modelscope import snapshot_download
+
+    logger.info(f"[Model Download] Downloading from ModelScope: {repo_id} -> {local_dir}")
+
+    snapshot_download(
+        model_id=repo_id,
+        local_dir=str(local_dir),
+    )
+
+
+def _smart_download(
+    repo_id: str,
+    local_dir: Path,
+    token: Optional[str] = None,
+    prefer_source: Optional[str] = None,
+) -> Tuple[bool, str]:
+    """
+    Smart download with automatic fallback between HuggingFace and ModelScope.
+
+    Automatically detects network environment and chooses the best download source.
+    If the primary source fails, automatically falls back to the alternative.
+
+    Args:
+        repo_id: Repository ID (same format for both HF and ModelScope)
+        local_dir: Local directory to save the model
+        token: HuggingFace token for private repos (optional)
+        prefer_source: Preferred download source ("huggingface", "modelscope", or None for auto-detect)
+
+    Returns:
+        Tuple of (success, message)
+    """
+    # Ensure directory exists
+    local_dir.mkdir(parents=True, exist_ok=True)
+
+    # Determine primary source
+    if prefer_source == "huggingface":
+        use_huggingface_first = True
+        logger.info("[Model Download] User preference: HuggingFace Hub")
+    elif prefer_source == "modelscope":
+        use_huggingface_first = False
+        logger.info("[Model Download] User preference: ModelScope")
+    else:
+        # Auto-detect network environment
+        can_access_google = _can_access_google()
+        use_huggingface_first = can_access_google
+        logger.info(f"[Model Download] Auto-detected: {'HuggingFace Hub' if can_access_google else 'ModelScope'}")
+
+    if use_huggingface_first:
+        logger.info("[Model Download] Using HuggingFace Hub...")
+        try:
+            _download_from_huggingface_internal(repo_id, local_dir, token)
+            return True, f"Successfully downloaded from HuggingFace: {repo_id}"
+        except Exception as e:
+            logger.warning(f"[Model Download] HuggingFace download failed: {e}")
+            logger.info("[Model Download] Falling back to ModelScope...")
+            try:
+                _download_from_modelscope_internal(repo_id, local_dir)
+                return True, f"Successfully downloaded from ModelScope: {repo_id}"
+            except Exception as e2:
+                error_msg = f"Both HuggingFace and ModelScope downloads failed. HF: {e}, MS: {e2}"
+                logger.error(error_msg)
+                return False, error_msg
+    else:
+        logger.info("[Model Download] Using ModelScope...")
+        try:
+            _download_from_modelscope_internal(repo_id, local_dir)
+            return True, f"Successfully downloaded from ModelScope: {repo_id}"
+        except Exception as e:
+            logger.warning(f"[Model Download] ModelScope download failed: {e}")
+            logger.info("[Model Download] Falling back to HuggingFace Hub...")
+            try:
+                _download_from_huggingface_internal(repo_id, local_dir, token)
+                return True, f"Successfully downloaded from HuggingFace: {repo_id}"
+            except Exception as e2:
+                error_msg = f"Both ModelScope and HuggingFace downloads failed. MS: {e}, HF: {e2}"
+                logger.error(error_msg)
+                return False, error_msg
+
+
+# =============================================================================
+# Model Registry
+# =============================================================================
+# Main model contains core components (vae, text_encoder, default DiT)
+MAIN_MODEL_REPO = "ACE-Step/Ace-Step1.5"
+
+# Sub-models that can be downloaded separately into the checkpoints directory
+SUBMODEL_REGISTRY: Dict[str, str] = {
+    # LM models
+    "acestep-5Hz-lm-0.6B": "ACE-Step/acestep-5Hz-lm-0.6B",
+    "acestep-5Hz-lm-4B": "ACE-Step/acestep-5Hz-lm-4B",
+    # DiT models
+    "acestep-v15-turbo-shift3": "ACE-Step/acestep-v15-turbo-shift3",
+    "acestep-v15-sft": "ACE-Step/acestep-v15-sft",
+    "acestep-v15-base": "ACE-Step/acestep-v15-base",
+    "acestep-v15-turbo-shift1": "ACE-Step/acestep-v15-turbo-shift1",
+    "acestep-v15-turbo-continuous": "ACE-Step/acestep-v15-turbo-continuous",
+}
+
+# Components that come from the main model repo (ACE-Step/Ace-Step1.5)
+MAIN_MODEL_COMPONENTS = [
+    "acestep-v15-turbo",      # Default DiT model
+    "vae",                     # VAE for audio encoding/decoding
+    "Qwen3-Embedding-0.6B",    # Text encoder
+    "acestep-5Hz-lm-1.7B",     # Default LM model (1.7B)
+]
+
+# Default LM model (included in main model)
+DEFAULT_LM_MODEL = "acestep-5Hz-lm-1.7B"
+
+
+def get_project_root() -> Path:
+    """Get the project root directory."""
+    current_file = Path(__file__).resolve()
+    return current_file.parent.parent
+
+
+def get_checkpoints_dir(custom_dir: Optional[str] = None) -> Path:
+    """Get the checkpoints directory path."""
+    if custom_dir:
+        return Path(custom_dir)
+    return get_project_root() / "checkpoints"
+
+
+def check_main_model_exists(checkpoints_dir: Optional[Path] = None) -> bool:
+    """
+    Check if the main model components exist in the checkpoints directory.
+    
+    Returns:
+        True if all main model components exist, False otherwise.
+    """
+    if checkpoints_dir is None:
+        checkpoints_dir = get_checkpoints_dir()
+    
+    for component in MAIN_MODEL_COMPONENTS:
+        component_path = checkpoints_dir / component
+        if not component_path.exists():
+            return False
+    return True
+
+
+def check_model_exists(model_name: str, checkpoints_dir: Optional[Path] = None) -> bool:
+    """
+    Check if a specific model exists in the checkpoints directory.
+    
+    Args:
+        model_name: Name of the model to check
+        checkpoints_dir: Custom checkpoints directory (optional)
+    
+    Returns:
+        True if the model exists, False otherwise.
+    """
+    if checkpoints_dir is None:
+        checkpoints_dir = get_checkpoints_dir()
+    
+    model_path = checkpoints_dir / model_name
+    return model_path.exists()
+
+
+def list_available_models() -> Dict[str, str]:
+    """
+    List all available models for download.
+    
+    Returns:
+        Dictionary mapping local names to HuggingFace repo IDs.
+    """
+    models = {
+        "main": MAIN_MODEL_REPO,
+        **SUBMODEL_REGISTRY
+    }
+    return models
+
+
+def download_main_model(
+    checkpoints_dir: Optional[Path] = None,
+    force: bool = False,
+    token: Optional[str] = None,
+    prefer_source: Optional[str] = None,
+) -> Tuple[bool, str]:
+    """
+    Download the main ACE-Step model from HuggingFace or ModelScope.
+
+    The main model includes:
+    - acestep-v15-turbo (default DiT model)
+    - vae (audio encoder/decoder)
+    - Qwen3-Embedding-0.6B (text encoder)
+    - acestep-5Hz-lm-1.7B (default LM model)
+
+    Args:
+        checkpoints_dir: Custom checkpoints directory (optional)
+        force: Force re-download even if model exists
+        token: HuggingFace token for private repos (optional)
+        prefer_source: Preferred download source ("huggingface", "modelscope", or None for auto-detect)
+
+    Returns:
+        Tuple of (success, message)
+    """
+    if checkpoints_dir is None:
+        checkpoints_dir = get_checkpoints_dir()
+
+    # Ensure checkpoints directory exists
+    checkpoints_dir.mkdir(parents=True, exist_ok=True)
+
+    if not force and check_main_model_exists(checkpoints_dir):
+        return True, f"Main model already exists at {checkpoints_dir}"
+
+    print(f"Downloading main model from {MAIN_MODEL_REPO}...")
+    print(f"Destination: {checkpoints_dir}")
+    print("This may take a while depending on your internet connection...")
+
+    # Use smart download with automatic fallback
+    return _smart_download(MAIN_MODEL_REPO, checkpoints_dir, token, prefer_source)
+
+
+def download_submodel(
+    model_name: str,
+    checkpoints_dir: Optional[Path] = None,
+    force: bool = False,
+    token: Optional[str] = None,
+    prefer_source: Optional[str] = None,
+) -> Tuple[bool, str]:
+    """
+    Download a specific sub-model from HuggingFace or ModelScope.
+
+    Args:
+        model_name: Name of the model to download (must be in SUBMODEL_REGISTRY)
+        checkpoints_dir: Custom checkpoints directory (optional)
+        force: Force re-download even if model exists
+        token: HuggingFace token for private repos (optional)
+        prefer_source: Preferred download source ("huggingface", "modelscope", or None for auto-detect)
+
+    Returns:
+        Tuple of (success, message)
+    """
+    if model_name not in SUBMODEL_REGISTRY:
+        available = ", ".join(SUBMODEL_REGISTRY.keys())
+        return False, f"Unknown model '{model_name}'. Available models: {available}"
+
+    if checkpoints_dir is None:
+        checkpoints_dir = get_checkpoints_dir()
+
+    # Ensure checkpoints directory exists
+    checkpoints_dir.mkdir(parents=True, exist_ok=True)
+
+    model_path = checkpoints_dir / model_name
+
+    if not force and model_path.exists():
+        return True, f"Model '{model_name}' already exists at {model_path}"
+
+    repo_id = SUBMODEL_REGISTRY[model_name]
+
+    print(f"Downloading {model_name} from {repo_id}...")
+    print(f"Destination: {model_path}")
+
+    # Use smart download with automatic fallback
+    return _smart_download(repo_id, model_path, token, prefer_source)
+
+
+def download_all_models(
+    checkpoints_dir: Optional[Path] = None,
+    force: bool = False,
+    token: Optional[str] = None,
+) -> Tuple[bool, List[str]]:
+    """
+    Download all available models.
+    
+    Args:
+        checkpoints_dir: Custom checkpoints directory (optional)
+        force: Force re-download even if models exist
+        token: HuggingFace token for private repos (optional)
+    
+    Returns:
+        Tuple of (all_success, list of messages)
+    """
+    if checkpoints_dir is None:
+        checkpoints_dir = get_checkpoints_dir()
+    
+    messages = []
+    all_success = True
+    
+    # Download main model first
+    success, msg = download_main_model(checkpoints_dir, force, token)
+    messages.append(msg)
+    if not success:
+        all_success = False
+    
+    # Download all sub-models
+    for model_name in SUBMODEL_REGISTRY:
+        success, msg = download_submodel(model_name, checkpoints_dir, force, token)
+        messages.append(msg)
+        if not success:
+            all_success = False
+    
+    return all_success, messages
+
+
+def ensure_main_model(
+    checkpoints_dir: Optional[Path] = None,
+    token: Optional[str] = None,
+    prefer_source: Optional[str] = None,
+) -> Tuple[bool, str]:
+    """
+    Ensure the main model is available, downloading if necessary.
+
+    This function is designed to be called during initialization.
+    It will only download if the model doesn't exist.
+
+    Args:
+        checkpoints_dir: Custom checkpoints directory (optional)
+        token: HuggingFace token for private repos (optional)
+        prefer_source: Preferred download source ("huggingface", "modelscope", or None for auto-detect)
+
+    Returns:
+        Tuple of (success, message)
+    """
+    if checkpoints_dir is None:
+        checkpoints_dir = get_checkpoints_dir()
+
+    if check_main_model_exists(checkpoints_dir):
+        return True, "Main model is available"
+
+    print("\n" + "=" * 60)
+    print("Main model not found. Starting automatic download...")
+    print("=" * 60 + "\n")
+
+    return download_main_model(checkpoints_dir, token=token, prefer_source=prefer_source)
+
+
+def ensure_lm_model(
+    model_name: Optional[str] = None,
+    checkpoints_dir: Optional[Path] = None,
+    token: Optional[str] = None,
+    prefer_source: Optional[str] = None,
+) -> Tuple[bool, str]:
+    """
+    Ensure an LM model is available, downloading if necessary.
+
+    Args:
+        model_name: Name of the LM model (defaults to DEFAULT_LM_MODEL)
+        checkpoints_dir: Custom checkpoints directory (optional)
+        token: HuggingFace token for private repos (optional)
+        prefer_source: Preferred download source ("huggingface", "modelscope", or None for auto-detect)
+
+    Returns:
+        Tuple of (success, message)
+    """
+    if model_name is None:
+        model_name = DEFAULT_LM_MODEL
+
+    if checkpoints_dir is None:
+        checkpoints_dir = get_checkpoints_dir()
+
+    if check_model_exists(model_name, checkpoints_dir):
+        return True, f"LM model '{model_name}' is available"
+
+    # Check if this is a known LM model
+    if model_name not in SUBMODEL_REGISTRY:
+        # Check if it might be a variant name
+        for known_model in SUBMODEL_REGISTRY:
+            if "lm" in known_model.lower() and model_name.lower() in known_model.lower():
+                model_name = known_model
+                break
+        else:
+            return False, f"Unknown LM model: {model_name}"
+
+    print("\n" + "=" * 60)
+    print(f"LM model '{model_name}' not found. Starting automatic download...")
+    print("=" * 60 + "\n")
+
+    return download_submodel(model_name, checkpoints_dir, token=token, prefer_source=prefer_source)
+
+
+def ensure_dit_model(
+    model_name: str,
+    checkpoints_dir: Optional[Path] = None,
+    token: Optional[str] = None,
+    prefer_source: Optional[str] = None,
+) -> Tuple[bool, str]:
+    """
+    Ensure a DiT model is available, downloading if necessary.
+
+    Args:
+        model_name: Name of the DiT model
+        checkpoints_dir: Custom checkpoints directory (optional)
+        token: HuggingFace token for private repos (optional)
+        prefer_source: Preferred download source ("huggingface", "modelscope", or None for auto-detect)
+
+    Returns:
+        Tuple of (success, message)
+    """
+    if checkpoints_dir is None:
+        checkpoints_dir = get_checkpoints_dir()
+
+    if check_model_exists(model_name, checkpoints_dir):
+        return True, f"DiT model '{model_name}' is available"
+
+    # Check if this is the default turbo model (part of main)
+    if model_name == "acestep-v15-turbo":
+        return ensure_main_model(checkpoints_dir, token, prefer_source)
+
+    # Check if it's a known sub-model
+    if model_name in SUBMODEL_REGISTRY:
+        print("\n" + "=" * 60)
+        print(f"DiT model '{model_name}' not found. Starting automatic download...")
+        print("=" * 60 + "\n")
+        return download_submodel(model_name, checkpoints_dir, token=token, prefer_source=prefer_source)
+
+    return False, f"Unknown DiT model: {model_name}"
+
+
+def print_model_list():
+    """Print formatted list of available models."""
+    print("\nAvailable Models for Download:")
+    print("=" * 60)
+    print("\nSupported Sources: HuggingFace Hub <-> ModelScope (auto-fallback)")
+
+    print("\n[Main Model]")
+    print(f"  main -> {MAIN_MODEL_REPO}")
+    print("  Contains: vae, Qwen3-Embedding-0.6B, acestep-v15-turbo, acestep-5Hz-lm-1.7B")
+
+    print("\n[Optional LM Models]")
+    for name, repo in SUBMODEL_REGISTRY.items():
+        if "lm" in name.lower():
+            print(f"  {name} -> {repo}")
+
+    print("\n[Optional DiT Models]")
+    for name, repo in SUBMODEL_REGISTRY.items():
+        if "lm" not in name.lower():
+            print(f"  {name} -> {repo}")
+
+    print("\n" + "=" * 60)
+
+
+def main():
+    """CLI entry point for model downloading."""
+    parser = argparse.ArgumentParser(
+        description="Download ACE-Step models with automatic fallback (HuggingFace <-> ModelScope)",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  acestep-download                          # Download main model (includes LM 1.7B)
+  acestep-download --all                    # Download all available models
+  acestep-download --model acestep-v15-sft  # Download a specific model
+  acestep-download --list                   # List all available models
+
+Network Detection:
+  Automatically detects network environment and chooses the best download source:
+  - Google accessible -> HuggingFace (fallback to ModelScope)
+  - Google blocked -> ModelScope (fallback to HuggingFace)
+
+Alternative using huggingface-cli:
+  huggingface-cli download ACE-Step/Ace-Step1.5 --local-dir ./checkpoints
+  huggingface-cli download ACE-Step/acestep-5Hz-lm-0.6B --local-dir ./checkpoints/acestep-5Hz-lm-0.6B
+        """
+    )
+    
+    parser.add_argument(
+        "--model", "-m",
+        type=str,
+        help="Specific model to download (use --list to see available models)"
+    )
+    parser.add_argument(
+        "--all", "-a",
+        action="store_true",
+        help="Download all available models"
+    )
+    parser.add_argument(
+        "--list", "-l",
+        action="store_true",
+        help="List all available models"
+    )
+    parser.add_argument(
+        "--dir", "-d",
+        type=str,
+        default=None,
+        help="Custom checkpoints directory (default: ./checkpoints)"
+    )
+    parser.add_argument(
+        "--force", "-f",
+        action="store_true",
+        help="Force re-download even if model exists"
+    )
+    parser.add_argument(
+        "--token", "-t",
+        type=str,
+        default=None,
+        help="HuggingFace token for private repos"
+    )
+    parser.add_argument(
+        "--skip-main",
+        action="store_true",
+        help="Skip downloading the main model (only download specified sub-model)"
+    )
+    
+    args = parser.parse_args()
+    
+    # Handle --list
+    if args.list:
+        print_model_list()
+        return 0
+    
+    # Get checkpoints directory
+    checkpoints_dir = get_checkpoints_dir(args.dir) if args.dir else get_checkpoints_dir()
+    print(f"Checkpoints directory: {checkpoints_dir}")
+    
+    # Handle --all
+    if args.all:
+        success, messages = download_all_models(checkpoints_dir, args.force, args.token)
+        for msg in messages:
+            print(msg)
+        return 0 if success else 1
+    
+    # Handle --model
+    if args.model:
+        if args.model == "main":
+            success, msg = download_main_model(checkpoints_dir, args.force, args.token)
+        elif args.model in SUBMODEL_REGISTRY:
+            # Download main model first if needed (unless --skip-main)
+            if not args.skip_main and not check_main_model_exists(checkpoints_dir):
+                print("Main model not found. Downloading main model first...")
+                main_success, main_msg = download_main_model(checkpoints_dir, args.force, args.token)
+                print(main_msg)
+                if not main_success:
+                    return 1
+            
+            success, msg = download_submodel(args.model, checkpoints_dir, args.force, args.token)
+        else:
+            print(f"Unknown model: {args.model}")
+            print("Use --list to see available models")
+            return 1
+        
+        print(msg)
+        return 0 if success else 1
+    
+    # Default: download main model (includes default LM 1.7B)
+    print("Downloading main model (includes vae, text encoder, DiT, and LM 1.7B)...")
+    
+    # Download main model
+    success, msg = download_main_model(checkpoints_dir, args.force, args.token)
+    print(msg)
+    
+    if success:
+        print("\nDownload complete!")
+        print(f"Models are available at: {checkpoints_dir}")
+    
+    return 0 if success else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

acestep/openrouter_adapter.py

@@ -0,0 +1,773 @@
+"""OpenRouter API adapter for ACE-Step music generation.
+
+This module provides OpenRouter-compatible endpoints that wrap the ACE-Step
+music generation API, mounted as a sub-router on the main api_server.
+
+All generation requests go through the shared asyncio.Queue, ensuring unified
+GPU scheduling with release_task.
+
+Endpoints:
+- POST /v1/chat/completions  - Generate music via chat completion format
+- GET  /v1/models            - List available models (OpenRouter format)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import json
+import os
+import re
+import tempfile
+import time
+from typing import Any, Dict, List, Optional, Tuple
+from uuid import uuid4
+
+from fastapi import APIRouter, HTTPException, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from acestep.openrouter_models import (
+    AudioConfig,
+    ChatCompletionRequest,
+    ModelInfo,
+    ModelPricing,
+    ModelsResponse,
+)
+
+
+# =============================================================================
+# Constants
+# =============================================================================
+
+MODEL_PREFIX = "acestep"
+DEFAULT_AUDIO_FORMAT = "mp3"
+
+# Generation timeout for non-streaming requests (seconds)
+GENERATION_TIMEOUT = int(os.environ.get("ACESTEP_GENERATION_TIMEOUT", "600"))
+
+
+# =============================================================================
+# Helper Functions
+# =============================================================================
+
+def _generate_completion_id() -> str:
+    """Generate a unique completion ID."""
+    return f"chatcmpl-{uuid4().hex[:24]}"
+
+
+def _get_model_id(model_name: str) -> str:
+    """Convert internal model name to OpenRouter model ID."""
+    return f"{MODEL_PREFIX}/{model_name}"
+
+
+def _parse_model_name(model_id: str) -> str:
+    """Extract internal model name from OpenRouter model ID."""
+    if "/" in model_id:
+        return model_id.split("/", 1)[1]
+    return model_id
+
+
+def _audio_to_base64_url(audio_path: str, audio_format: str = "mp3") -> str:
+    """Convert audio file to base64 data URL."""
+    if not audio_path or not os.path.exists(audio_path):
+        return ""
+
+    mime_types = {
+        "mp3": "audio/mpeg",
+        "wav": "audio/wav",
+        "flac": "audio/flac",
+        "ogg": "audio/ogg",
+        "m4a": "audio/mp4",
+        "aac": "audio/aac",
+    }
+    mime_type = mime_types.get(audio_format.lower(), "audio/mpeg")
+
+    with open(audio_path, "rb") as f:
+        audio_data = f.read()
+
+    b64_data = base64.b64encode(audio_data).decode("utf-8")
+    return f"data:{mime_type};base64,{b64_data}"
+
+
+def _format_lm_content(result: Dict[str, Any]) -> str:
+    """Format generation result as content string with metadata and lyrics."""
+    metas = result.get("metas", {})
+    lyrics = result.get("lyrics", "")
+
+    parts = []
+
+    # Add metadata section
+    meta_lines = []
+    caption = metas.get("prompt") or metas.get("caption") or result.get("prompt", "")
+    if caption:
+        meta_lines.append(f"**Caption:** {caption}")
+    if metas.get("bpm") and metas["bpm"] != "N/A":
+        meta_lines.append(f"**BPM:** {metas['bpm']}")
+    if metas.get("duration") and metas["duration"] != "N/A":
+        meta_lines.append(f"**Duration:** {metas['duration']}s")
+    if metas.get("keyscale") and metas["keyscale"] != "N/A":
+        meta_lines.append(f"**Key:** {metas['keyscale']}")
+    if metas.get("timesignature") and metas["timesignature"] != "N/A":
+        meta_lines.append(f"**Time Signature:** {metas['timesignature']}")
+
+    if meta_lines:
+        parts.append("## Metadata\n" + "\n".join(meta_lines))
+
+    # Add lyrics section
+    if lyrics and lyrics.strip() and lyrics.strip().lower() not in ("[inst]", "[instrumental]"):
+        parts.append(f"## Lyrics\n{lyrics}")
+
+    if parts:
+        return "\n\n".join(parts)
+    else:
+        return "Music generated successfully."
+
+
+def _base64_to_temp_file(b64_data: str, audio_format: str = "mp3") -> str:
+    """Save base64 audio data to temporary file."""
+    if "," in b64_data:
+        b64_data = b64_data.split(",", 1)[1]
+
+    audio_bytes = base64.b64decode(b64_data)
+    suffix = f".{audio_format}" if not audio_format.startswith(".") else audio_format
+    fd, path = tempfile.mkstemp(suffix=suffix, prefix="openrouter_audio_")
+    os.close(fd)
+
+    with open(path, "wb") as f:
+        f.write(audio_bytes)
+
+    return path
+
+
+def _extract_tagged_content(text: str) -> Tuple[Optional[str], Optional[str], str]:
+    """
+    Extract content from <prompt> and <lyrics> tags.
+
+    Returns:
+        (prompt, lyrics, remaining_text)
+    """
+    prompt = None
+    lyrics = None
+    remaining = text
+
+    prompt_match = re.search(r'<prompt>(.*?)</prompt>', text, re.DOTALL | re.IGNORECASE)
+    if prompt_match:
+        prompt = prompt_match.group(1).strip()
+        remaining = remaining.replace(prompt_match.group(0), '').strip()
+
+    lyrics_match = re.search(r'<lyrics>(.*?)</lyrics>', text, re.DOTALL | re.IGNORECASE)
+    if lyrics_match:
+        lyrics = lyrics_match.group(1).strip()
+        remaining = remaining.replace(lyrics_match.group(0), '').strip()
+
+    return prompt, lyrics, remaining
+
+
+def _looks_like_lyrics(text: str) -> bool:
+    """Heuristic to detect if text looks like song lyrics."""
+    if not text:
+        return False
+
+    lyrics_markers = [
+        "[verse", "[chorus", "[bridge", "[intro", "[outro",
+        "[hook", "[pre-chorus", "[refrain", "[inst",
+    ]
+    text_lower = text.lower()
+    for marker in lyrics_markers:
+        if marker in text_lower:
+            return True
+
+    lines = [line.strip() for line in text.split("\n") if line.strip()]
+    if len(lines) >= 4:
+        avg_line_length = sum(len(line) for line in lines) / len(lines)
+        if avg_line_length < 60:
+            return True
+
+    return False
+
+
+def _is_instrumental(lyrics: str) -> bool:
+    """Check if the music should be instrumental based on lyrics."""
+    if not lyrics:
+        return True
+    lyrics_clean = lyrics.strip().lower()
+    if not lyrics_clean:
+        return True
+    return lyrics_clean in ("[inst]", "[instrumental]")
+
+
+def _parse_messages(messages: List[Any]) -> Tuple[str, str, List[str], Optional[str], Optional[str]]:
+    """
+    Parse chat messages to extract prompt, lyrics, sample_query and audio references.
+
+    Supports two modes:
+    1. Tagged mode: Use <prompt>...</prompt> and <lyrics>...</lyrics> tags
+    2. Heuristic mode: Auto-detect based on content structure
+
+    Multiple input_audio blocks are collected in order (like multiple images).
+    The caller routes them to src_audio / reference_audio based on task_type.
+
+    Returns:
+        (prompt, lyrics, audio_paths, system_instruction, sample_query)
+    """
+    prompt_parts = []
+    lyrics = ""
+    sample_query = None
+    audio_paths: List[str] = []
+    system_instruction = None
+    has_tags = False
+
+    for msg in messages:
+        role = msg.role
+        content = msg.content
+
+        if role == "system":
+            if isinstance(content, str):
+                system_instruction = content
+            continue
+
+        if role != "user":
+            continue
+
+        if isinstance(content, str):
+            text = content.strip()
+            tagged_prompt, tagged_lyrics, remaining = _extract_tagged_content(text)
+            if tagged_prompt is not None or tagged_lyrics is not None:
+                has_tags = True
+                if tagged_prompt:
+                    prompt_parts.append(tagged_prompt)
+                if tagged_lyrics:
+                    lyrics = tagged_lyrics
+                if remaining:
+                    prompt_parts.append(remaining)
+            else:
+                if _looks_like_lyrics(text):
+                    lyrics = text
+                else:
+                    prompt_parts.append(text)
+
+        elif isinstance(content, list):
+            for part in content:
+                if isinstance(part, dict):
+                    part_type = part.get("type", "")
+
+                    if part_type == "text":
+                        text = part.get("text", "").strip()
+                        tagged_prompt, tagged_lyrics, remaining = _extract_tagged_content(text)
+                        if tagged_prompt is not None or tagged_lyrics is not None:
+                            has_tags = True
+                            if tagged_prompt:
+                                prompt_parts.append(tagged_prompt)
+                            if tagged_lyrics:
+                                lyrics = tagged_lyrics
+                            if remaining:
+                                prompt_parts.append(remaining)
+                        elif _looks_like_lyrics(text):
+                            lyrics = text
+                        else:
+                            prompt_parts.append(text)
+
+                    elif part_type == "input_audio":
+                        audio_data = part.get("input_audio", {})
+                        if isinstance(audio_data, dict):
+                            b64_data = audio_data.get("data", "")
+                            audio_format = audio_data.get("format", "mp3")
+                            if b64_data:
+                                try:
+                                    path = _base64_to_temp_file(b64_data, audio_format)
+                                    audio_paths.append(path)
+                                except Exception:
+                                    pass
+
+                elif hasattr(part, "type"):
+                    if part.type == "text":
+                        text = getattr(part, "text", "").strip()
+                        tagged_prompt, tagged_lyrics, remaining = _extract_tagged_content(text)
+                        if tagged_prompt is not None or tagged_lyrics is not None:
+                            has_tags = True
+                            if tagged_prompt:
+                                prompt_parts.append(tagged_prompt)
+                            if tagged_lyrics:
+                                lyrics = tagged_lyrics
+                            if remaining:
+                                prompt_parts.append(remaining)
+                        elif _looks_like_lyrics(text):
+                            lyrics = text
+                        else:
+                            prompt_parts.append(text)
+
+                    elif part.type == "input_audio":
+                        audio_data = getattr(part, "input_audio", None)
+                        if audio_data:
+                            b64_data = getattr(audio_data, "data", "")
+                            audio_format = getattr(audio_data, "format", "mp3")
+                            if b64_data:
+                                try:
+                                    path = _base64_to_temp_file(b64_data, audio_format)
+                                    audio_paths.append(path)
+                                except Exception:
+                                    pass
+
+    prompt = " ".join(prompt_parts).strip()
+
+    # Use sample mode when: no tags, no lyrics detected, and we have text input
+    if not has_tags and not lyrics and prompt:
+        sample_query = prompt
+        prompt = ""
+
+    return prompt, lyrics, audio_paths, system_instruction, sample_query
+
+
+def _to_generate_music_request(
+    req: ChatCompletionRequest,
+    prompt: str,
+    lyrics: str,
+    sample_query: Optional[str],
+    reference_audio_path: Optional[str],
+    src_audio_path: Optional[str],
+):
+    """
+    Convert OpenRouter ChatCompletionRequest to api_server's GenerateMusicRequest.
+
+    Audio routing depends on task_type:
+      text2music:           audio[0] → reference_audio
+      cover/repaint/lego/…: audio[0] → src_audio, audio[1] → reference_audio
+
+    task_type auto-detection:
+      text2music + reference_audio → music_continuation
+
+    Uses late import to avoid circular dependency with api_server.
+    """
+    from acestep.api_server import GenerateMusicRequest
+
+    audio_config = req.audio_config or AudioConfig()
+
+    # Resolve parameters from audio_config only
+    resolved_instrumental = audio_config.instrumental if audio_config.instrumental is not None else False
+
+    # If instrumental, set lyrics to [inst]
+    resolved_lyrics = lyrics
+    if req.lyrics:
+        resolved_lyrics = req.lyrics
+    if resolved_instrumental and not resolved_lyrics:
+        resolved_lyrics = "[inst]"
+
+    # Resolve sample_mode: explicit field takes priority, then auto-detect from messages
+    resolved_sample_mode = req.sample_mode or bool(sample_query)
+    resolved_sample_query = sample_query or ""
+
+    # Resolve seed: pass through as-is (int or comma-separated string)
+    # handler.prepare_seeds() handles both formats
+    resolved_seed = req.seed if req.seed is not None else -1
+    use_random_seed = req.seed is None
+
+    # Resolve task_type
+    # Explicit task_type from request takes priority.
+    # For text2music: auto-detect based on reference_audio.
+    resolved_task_type = req.task_type
+    if resolved_task_type == "text2music" and reference_audio_path:
+        resolved_task_type = "music_continuation"
+
+    return GenerateMusicRequest(
+        # Text input
+        prompt=prompt,
+        lyrics=resolved_lyrics,
+        sample_query=resolved_sample_query,
+        sample_mode=resolved_sample_mode,
+
+        # Music metadata
+        bpm=audio_config.bpm,
+        key_scale=audio_config.key_scale or "",
+        time_signature=audio_config.time_signature or "",
+        audio_duration=audio_config.duration if audio_config.duration else None,
+        vocal_language=audio_config.vocal_language or "en",
+
+        # LM parameters
+        lm_temperature=req.temperature if req.temperature is not None else 0.85,
+        lm_top_p=req.top_p if req.top_p is not None else 0.9,
+        lm_top_k=req.top_k if req.top_k is not None else 0,
+        thinking=req.thinking if req.thinking is not None else False,
+
+        # Generation parameters
+        inference_steps=8,
+        guidance_scale=req.guidance_scale if req.guidance_scale is not None else 7.0,
+        seed=resolved_seed,
+        use_random_seed=use_random_seed,
+        batch_size=req.batch_size if req.batch_size is not None else 1,
+
+        # Task type
+        task_type=resolved_task_type,
+
+        # Audio paths
+        reference_audio_path=reference_audio_path or None,
+        src_audio_path=src_audio_path or None,
+
+        # Audio editing
+        repainting_start=req.repainting_start,
+        repainting_end=req.repainting_end,
+        audio_cover_strength=req.audio_cover_strength,
+
+        # Format / CoT control
+        use_format=req.use_format,
+        use_cot_caption=req.use_cot_caption,
+        use_cot_language=req.use_cot_language,
+
+        # Model selection
+        model=_parse_model_name(req.model),
+
+        # Audio format
+        audio_format=(audio_config.format or DEFAULT_AUDIO_FORMAT),
+    )
+
+
+def _build_openrouter_response(
+    rec: Any,
+    model_id: str,
+    audio_format: str,
+) -> JSONResponse:
+    """Build OpenRouter non-streaming response from a completed JobRecord."""
+    if rec.status != "succeeded" or not rec.result:
+        error_msg = rec.error or "Generation failed"
+        raise HTTPException(status_code=500, detail=error_msg)
+
+    result = rec.result
+    completion_id = _generate_completion_id()
+    created_timestamp = int(time.time())
+
+    text_content = _format_lm_content(result)
+
+    # Encode audio
+    audio_obj = None
+    raw_audio_paths = result.get("raw_audio_paths", [])
+    if raw_audio_paths:
+        audio_path = raw_audio_paths[0]
+        if audio_path and os.path.exists(audio_path):
+            b64_url = _audio_to_base64_url(audio_path, audio_format)
+            if b64_url:
+                audio_obj = [{
+                    "type": "audio_url",
+                    "audio_url": {"url": b64_url},
+                }]
+
+    response_data = {
+        "id": completion_id,
+        "object": "chat.completion",
+        "created": created_timestamp,
+        "model": model_id,
+        "choices": [{
+            "index": 0,
+            "message": {
+                "role": "assistant",
+                "content": text_content,
+                "audio": audio_obj,
+            },
+            "finish_reason": "stop",
+        }],
+        "usage": {
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "total_tokens": 0,
+        },
+    }
+
+    return JSONResponse(content=response_data)
+
+
+async def _openrouter_stream_generator(
+    rec: Any,
+    model_id: str,
+    audio_format: str,
+):
+    """
+    SSE stream generator that reads from rec.progress_queue.
+
+    Yields heartbeat chunks every 2 seconds while waiting for the
+    queue worker to push the generation result.
+    """
+    completion_id = _generate_completion_id()
+    created_timestamp = int(time.time())
+
+    def _make_chunk(
+        content: Optional[str] = None,
+        role: Optional[str] = None,
+        audio: Optional[Any] = None,
+        finish_reason: Optional[str] = None,
+    ) -> str:
+        delta = {}
+        if role:
+            delta["role"] = role
+        if content is not None:
+            delta["content"] = content
+        if audio is not None:
+            delta["audio"] = audio
+
+        chunk = {
+            "id": completion_id,
+            "object": "chat.completion.chunk",
+            "created": created_timestamp,
+            "model": model_id,
+            "choices": [{
+                "index": 0,
+                "delta": delta,
+                "finish_reason": finish_reason,
+            }],
+        }
+        return f"data: {json.dumps(chunk)}\n\n"
+
+    # Initial role chunk
+    yield _make_chunk(role="assistant", content="Generating music")
+    await asyncio.sleep(0)
+
+    # Wait for result with periodic heartbeats
+    while True:
+        try:
+            msg = await asyncio.wait_for(rec.progress_queue.get(), timeout=2.0)
+        except asyncio.TimeoutError:
+            yield _make_chunk(content=".")
+            await asyncio.sleep(0)
+            continue
+
+        msg_type = msg.get("type")
+
+        if msg_type == "done":
+            break
+
+        elif msg_type == "error":
+            yield _make_chunk(content=f"\n\nError: {msg.get('content', 'Unknown error')}")
+            yield _make_chunk(finish_reason="error")
+            yield "data: [DONE]\n\n"
+            return
+
+        elif msg_type == "result":
+            result = msg.get("result", {})
+
+            # Send LM content
+            lm_content = _format_lm_content(result)
+            yield _make_chunk(content=f"\n\n{lm_content}")
+            await asyncio.sleep(0)
+
+            # Send audio
+            raw_audio_paths = result.get("raw_audio_paths", [])
+            if raw_audio_paths:
+                audio_path = raw_audio_paths[0]
+                if audio_path and os.path.exists(audio_path):
+                    b64_url = _audio_to_base64_url(audio_path, audio_format)
+                    if b64_url:
+                        audio_list = [{
+                            "type": "audio_url",
+                            "audio_url": {"url": b64_url},
+                        }]
+                        yield _make_chunk(audio=audio_list)
+                        await asyncio.sleep(0)
+
+    # Finish
+    yield _make_chunk(finish_reason="stop")
+    yield "data: [DONE]\n\n"
+
+
+# =============================================================================
+# Router Factory
+# =============================================================================
+
+def create_openrouter_router(app_state_getter) -> APIRouter:
+    """
+    Create OpenRouter-compatible API router.
+
+    Args:
+        app_state_getter: Callable that returns the FastAPI app.state object
+
+    Returns:
+        APIRouter with OpenRouter-compatible endpoints
+    """
+    router = APIRouter(tags=["OpenRouter Compatible"])
+
+    def _get_model_name_from_path(config_path: str) -> str:
+        """Extract model name from config path."""
+        if not config_path:
+            return ""
+        normalized = config_path.rstrip("/\\")
+        return os.path.basename(normalized)
+
+    @router.get("/v1/models", response_model=ModelsResponse)
+    async def list_models():
+        """List available models in OpenRouter format."""
+        state = app_state_getter()
+        models = []
+        created_timestamp = int(time.time()) - 86400 * 30
+
+        # Primary model
+        if getattr(state, "_initialized", False):
+            model_name = _get_model_name_from_path(state._config_path)
+            if model_name:
+                models.append(ModelInfo(
+                    id=_get_model_id(model_name),
+                    name=f"ACE-Step {model_name}",
+                    created=created_timestamp,
+                    input_modalities=["text", "audio"],
+                    output_modalities=["audio", "text"],
+                    context_length=4096,
+                    max_output_length=300,
+                    pricing=ModelPricing(
+                        prompt="0", completion="0", request="0",
+                    ),
+                    description="AI music generation model",
+                ))
+
+        # Secondary model
+        if getattr(state, "_initialized2", False) and getattr(state, "_config_path2", ""):
+            model_name = _get_model_name_from_path(state._config_path2)
+            if model_name:
+                models.append(ModelInfo(
+                    id=_get_model_id(model_name),
+                    name=f"ACE-Step {model_name}",
+                    created=created_timestamp,
+                    input_modalities=["text", "audio"],
+                    output_modalities=["audio", "text"],
+                    context_length=4096,
+                    max_output_length=300,
+                    pricing=ModelPricing(),
+                    description="AI music generation model",
+                ))
+
+        # Third model
+        if getattr(state, "_initialized3", False) and getattr(state, "_config_path3", ""):
+            model_name = _get_model_name_from_path(state._config_path3)
+            if model_name:
+                models.append(ModelInfo(
+                    id=_get_model_id(model_name),
+                    name=f"ACE-Step {model_name}",
+                    created=created_timestamp,
+                    input_modalities=["text", "audio"],
+                    output_modalities=["audio", "text"],
+                    context_length=4096,
+                    max_output_length=300,
+                    pricing=ModelPricing(),
+                    description="AI music generation model",
+                ))
+
+        return ModelsResponse(data=models)
+
+    @router.post("/v1/chat/completions")
+    async def chat_completions(request: Request):
+        """
+        OpenRouter-compatible chat completions endpoint for music generation.
+
+        Submits the request to the shared asyncio.Queue and waits for completion.
+        Supports both streaming (SSE) and non-streaming responses.
+        """
+        state = app_state_getter()
+
+        # Check initialization
+        if not getattr(state, "_initialized", False):
+            raise HTTPException(
+                status_code=503,
+                detail=f"Model not initialized. init_error={getattr(state, '_init_error', None)}"
+            )
+
+        # Parse request
+        try:
+            body = await request.json()
+            req = ChatCompletionRequest(**body)
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=f"Invalid request format: {str(e)}")
+
+        # Parse messages for text, audio, and system instruction
+        prompt, lyrics, audio_paths, system_instruction, sample_query = _parse_messages(req.messages)
+
+        # When lyrics or sample_mode is explicitly provided, the message text role
+        # is already known — skip auto-detection results.
+        # _parse_messages may have put raw text into prompt or sample_query;
+        # recover it as raw_text for re-assignment.
+        if req.lyrics or req.sample_mode:
+            raw_text = prompt or sample_query or ""
+            if req.lyrics:
+                # lyrics provided → message text is the prompt
+                prompt = raw_text
+                lyrics = req.lyrics
+                sample_query = None
+            else:
+                # sample_mode → message text is the sample_query
+                prompt = ""
+                lyrics = ""
+                sample_query = raw_text
+
+        if not prompt and not lyrics and not sample_query and not req.sample_mode and not audio_paths:
+            raise HTTPException(
+                status_code=400,
+                detail="No valid prompt, lyrics, sample query, or input audio found in request"
+            )
+
+        # Route audio paths based on task_type.
+        # Multiple input_audio blocks are supported (like multiple images).
+        #
+        # For cover / repaint / lego / extract / complete:
+        #   audio[0] → src_audio       (primary: the audio to edit / cover)
+        #   audio[1] → reference_audio  (optional: style conditioning)
+        #
+        # For text2music (default):
+        #   audio[0] → reference_audio  (style conditioning → music_continuation)
+        reference_audio_path = None
+        src_audio_path = None
+        _SRC_AUDIO_TASK_TYPES = {"cover", "repaint", "lego", "extract", "complete"}
+        if audio_paths:
+            if req.task_type in _SRC_AUDIO_TASK_TYPES:
+                src_audio_path = audio_paths[0]
+                if len(audio_paths) > 1:
+                    reference_audio_path = audio_paths[1]
+            else:
+                reference_audio_path = audio_paths[0]
+
+        # Convert to GenerateMusicRequest
+        gen_request = _to_generate_music_request(
+            req, prompt, lyrics, sample_query, reference_audio_path, src_audio_path
+        )
+
+        # Check queue capacity
+        job_queue = state.job_queue
+        if job_queue.full():
+            raise HTTPException(status_code=429, detail="Server busy: queue is full")
+
+        # Get audio format
+        audio_config = req.audio_config or AudioConfig()
+        audio_format = audio_config.format or DEFAULT_AUDIO_FORMAT
+
+        # Create job record and submit to queue
+        job_store = state.job_store
+        rec = job_store.create()
+
+        # Track temp files from base64 audio uploads
+        if audio_paths:
+            async with state.job_temp_files_lock:
+                state.job_temp_files.setdefault(rec.job_id, []).extend(audio_paths)
+
+        if req.stream:
+            # Streaming: use progress_queue
+            rec.progress_queue = asyncio.Queue()
+
+            async with state.pending_lock:
+                state.pending_ids.append(rec.job_id)
+
+            await job_queue.put((rec.job_id, gen_request))
+
+            return StreamingResponse(
+                _openrouter_stream_generator(rec, req.model, audio_format),
+                media_type="text/event-stream",
+            )
+        else:
+            # Non-streaming: use done_event
+            rec.done_event = asyncio.Event()
+
+            async with state.pending_lock:
+                state.pending_ids.append(rec.job_id)
+
+            await job_queue.put((rec.job_id, gen_request))
+
+            # Wait for completion with timeout
+            try:
+                await asyncio.wait_for(rec.done_event.wait(), timeout=GENERATION_TIMEOUT)
+            except asyncio.TimeoutError:
+                raise HTTPException(status_code=504, detail="Generation timeout")
+
+            return _build_openrouter_response(rec, req.model, audio_format)
+
+    return router

acestep/openrouter_models.py

@@ -0,0 +1,244 @@
+"""OpenRouter API compatible Pydantic models for ACE-Step.
+
+This module defines request/response models that conform to OpenRouter's
+chat completions API specification for audio generation.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Literal, Optional, Union
+from pydantic import BaseModel, Field
+
+
+# =============================================================================
+# Request Models
+# =============================================================================
+
+class AudioInputContent(BaseModel):
+    """Audio input content in base64 format."""
+    data: str = Field(..., description="Base64-encoded audio data")
+    format: str = Field(default="mp3", description="Audio format (mp3, wav, flac, etc.)")
+
+
+class TextContent(BaseModel):
+    """Text content block."""
+    type: Literal["text"] = "text"
+    text: str = Field(..., description="Text content")
+
+
+class AudioContent(BaseModel):
+    """Audio input content block."""
+    type: Literal["input_audio"] = "input_audio"
+    input_audio: AudioInputContent
+
+
+# Union type for message content
+ContentPart = Union[TextContent, AudioContent, Dict[str, Any]]
+
+
+class ChatMessage(BaseModel):
+    """A single message in the chat conversation."""
+    role: Literal["system", "user", "assistant"] = Field(..., description="Message role")
+    content: Union[str, List[ContentPart]] = Field(..., description="Message content")
+    name: Optional[str] = Field(default=None, description="Optional name for the message author")
+
+
+class AudioConfig(BaseModel):
+    """Audio generation configuration."""
+    duration: Optional[float] = Field(default=None, description="Target audio duration in seconds")
+    format: str = Field(default="mp3", description="Output audio format")
+    # ACE-Step specific parameters
+    bpm: Optional[int] = Field(default=None, description="Beats per minute")
+    key_scale: Optional[str] = Field(default=None, description="Musical key and scale")
+    time_signature: Optional[str] = Field(default=None, description="Time signature (e.g., 4/4)")
+    vocal_language: Optional[str] = Field(default=None, description="Vocal language code")
+    instrumental: Optional[bool] = Field(default=None, description="Generate instrumental only")
+
+
+class ChatCompletionRequest(BaseModel):
+    """OpenRouter-compatible chat completion request."""
+    model: str = Field(..., description="Model ID to use")
+    messages: List[ChatMessage] = Field(..., description="List of messages")
+
+    # Modalities
+    modalities: Optional[List[str]] = Field(
+        default=None,
+        description="Output modalities (e.g., ['audio', 'text'])"
+    )
+
+    # Audio configuration
+    audio_config: Optional[AudioConfig] = Field(
+        default=None,
+        description="Audio generation configuration"
+    )
+
+    # Standard OpenAI parameters
+    temperature: Optional[float] = Field(default=None, ge=0, le=2)
+    top_p: Optional[float] = Field(default=None, ge=0, le=1)
+    top_k: Optional[int] = Field(default=None, ge=0)
+    max_tokens: Optional[int] = Field(default=None, ge=1)
+    stream: bool = Field(default=False, description="Enable streaming response")
+    stop: Optional[Union[str, List[str]]] = Field(default=None)
+    seed: Optional[Union[int, str]] = Field(default=None, description="Seed(s) for reproducibility. Comma-separated for batch (e.g. '42,123,456')")
+
+    # ACE-Step specific parameters (extended)
+    thinking: Optional[bool] = Field(default=None, description="Use LM for audio code generation")
+    guidance_scale: Optional[float] = Field(default=None, description="Classifier-free guidance scale")
+    batch_size: Optional[int] = Field(default=None, description="Number of audio samples to generate")
+
+    # ACE-Step direct fields (bypass message parsing / audio_config)
+    lyrics: str = Field(default="", description="Direct lyrics input (bypass message parsing)")
+    sample_mode: bool = Field(default=False, description="Auto-generate caption/lyrics/metas via LM; user message becomes the query")
+    use_format: bool = Field(default=False, description="Use format_sample to enhance caption/lyrics")
+    use_cot_caption: bool = Field(default=True, description="Use CoT for caption rewriting")
+    use_cot_language: bool = Field(default=True, description="Use CoT for language detection")
+
+    # Task type
+    task_type: str = Field(default="text2music", description="Task type: text2music, cover, repaint, extract, lego, complete")
+
+    # Audio editing parameters
+    repainting_start: float = Field(default=0.0, description="Repainting region start (seconds)")
+    repainting_end: Optional[float] = Field(default=None, description="Repainting region end (seconds)")
+    audio_cover_strength: float = Field(default=1.0, description="Audio cover strength (0.0~1.0)")
+
+    class Config:
+        extra = "allow"  # Allow additional fields for forward compatibility
+
+
+# =============================================================================
+# Response Models
+# =============================================================================
+
+class AudioOutputUrl(BaseModel):
+    """Audio output URL (base64 data URL)."""
+    url: str = Field(..., description="Base64 data URL of the audio")
+
+
+class AudioOutput(BaseModel):
+    """Audio output content block."""
+    type: Literal["audio_url"] = "audio_url"
+    audio_url: AudioOutputUrl
+
+
+class AssistantMessage(BaseModel):
+    """Assistant response message."""
+    role: Literal["assistant"] = "assistant"
+    content: Optional[str] = Field(default=None, description="Text content")
+    audio: Optional[List[AudioOutput]] = Field(default=None, description="Generated audio files")
+
+
+class Choice(BaseModel):
+    """A single completion choice."""
+    index: int = Field(default=0)
+    message: AssistantMessage
+    finish_reason: Literal["stop", "length", "content_filter", "error"] = "stop"
+
+
+class Usage(BaseModel):
+    """Token usage statistics."""
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+
+
+class ChatCompletionResponse(BaseModel):
+    """OpenRouter-compatible chat completion response."""
+    id: str = Field(..., description="Unique completion ID")
+    object: Literal["chat.completion"] = "chat.completion"
+    created: int = Field(..., description="Unix timestamp")
+    model: str = Field(..., description="Model ID used")
+    choices: List[Choice] = Field(..., description="Completion choices")
+    usage: Usage = Field(default_factory=Usage)
+
+    # Extended metadata
+    system_fingerprint: Optional[str] = Field(default=None)
+
+
+# =============================================================================
+# Streaming Response Models
+# =============================================================================
+
+class DeltaContent(BaseModel):
+    """Delta content for streaming."""
+    role: Optional[Literal["assistant"]] = None
+    content: Optional[str] = None
+    audio: Optional[List[AudioOutput]] = None
+
+
+class StreamChoice(BaseModel):
+    """Streaming choice."""
+    index: int = 0
+    delta: DeltaContent
+    finish_reason: Optional[Literal["stop", "length", "content_filter", "error"]] = None
+
+
+class ChatCompletionChunk(BaseModel):
+    """Streaming chunk response."""
+    id: str
+    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
+    created: int
+    model: str
+    choices: List[StreamChoice]
+
+
+# =============================================================================
+# Models Endpoint Response
+# =============================================================================
+
+class ModelPricing(BaseModel):
+    """Model pricing information."""
+    prompt: str = Field(default="0", description="Price per prompt token in USD")
+    completion: str = Field(default="0", description="Price per completion token in USD")
+    request: str = Field(default="0", description="Price per request in USD")
+    image: str = Field(default="0", description="Price per image in USD")
+
+
+class ModelInfo(BaseModel):
+    """OpenRouter-compatible model information."""
+    id: str = Field(..., description="Model identifier")
+    name: str = Field(..., description="Display name")
+    created: int = Field(..., description="Unix timestamp of creation")
+
+    # Modalities
+    input_modalities: List[str] = Field(
+        default_factory=lambda: ["text"],
+        description="Supported input modalities"
+    )
+    output_modalities: List[str] = Field(
+        default_factory=lambda: ["audio", "text"],
+        description="Supported output modalities"
+    )
+
+    # Limits
+    context_length: int = Field(default=4096, description="Maximum context length")
+    max_output_length: int = Field(default=300, description="Maximum output length in seconds")
+
+    # Pricing
+    pricing: ModelPricing = Field(default_factory=ModelPricing)
+
+    # Metadata
+    description: Optional[str] = Field(default=None)
+    architecture: Optional[Dict[str, Any]] = Field(default=None)
+
+
+class ModelsResponse(BaseModel):
+    """Response for /v1/models endpoint."""
+    object: Literal["list"] = "list"
+    data: List[ModelInfo] = Field(default_factory=list)
+
+
+# =============================================================================
+# Error Response
+# =============================================================================
+
+class ErrorDetail(BaseModel):
+    """Error detail information."""
+    message: str
+    type: str = "invalid_request_error"
+    param: Optional[str] = None
+    code: Optional[str] = None
+
+
+class ErrorResponse(BaseModel):
+    """OpenRouter-compatible error response."""
+    error: ErrorDetail

acestep/test_time_scaling.py

@@ -0,0 +1,410 @@
+"""
+Test-Time Scaling Module
+Implements perplexity-based scoring for generated audio codes
+"""
+import torch
+import torch.nn.functional as F
+from typing import Tuple, Optional, Dict, Any, List
+from loguru import logger
+import yaml
+import math
+import re
+
+
+def pmi_score(log_prob_conditional: float, log_prob_unconditional: float) -> float:
+    """
+    Calculate Pointwise Mutual Information (PMI) score.
+    
+    PMI = log P(condition|codes) - log P(condition)
+        = log [P(codes|condition) / P(codes)]
+    
+    This removes the bias from P(condition) and measures how much the codes
+    improve our ability to predict the condition.
+    
+    Args:
+        log_prob_conditional: Average log probability of condition given codes
+        log_prob_unconditional: Average log probability of condition without codes
+        
+    Returns:
+        PMI score (higher is better, can be positive or negative)
+        - Positive: codes improve prediction → good match
+        - Zero: codes don't help → no correlation
+        - Negative: codes hurt prediction → poor match
+    """
+    return log_prob_conditional - log_prob_unconditional
+
+
+def pmi_to_normalized_score(pmi: float, scale: float = 0.1) -> float:
+    """
+    Convert PMI score to normalized [0, 1] range using sigmoid function.
+    
+    score = sigmoid(PMI / scale) = 1 / (1 + exp(-PMI / scale))
+    
+    Args:
+        pmi: PMI score (can be positive or negative)
+        scale: Scale parameter to control sensitivity (default 0.1)
+               - Smaller scale: more sensitive to PMI changes
+               - Larger scale: less sensitive to PMI changes
+        
+    Returns:
+        Normalized score in [0, 1] range, where:
+        - PMI > 0 → score > 0.5 (good match)
+        - PMI = 0 → score = 0.5 (neutral)
+        - PMI < 0 → score < 0.5 (poor match)
+        
+    Examples (scale=1.0):
+        PMI=2.0  → score≈0.88  (excellent)
+        PMI=1.0  → score≈0.73  (good)
+        PMI=0.0  → score=0.50  (neutral)
+        PMI=-1.0 → score≈0.27  (poor)
+        PMI=-2.0 → score≈0.12  (bad)
+    """
+    return 1.0 / (1.0 + math.exp(-pmi / scale))
+
+
+def _get_logits_and_target_for_scoring(llm_handler, formatted_prompt: str,
+                                       target_text: str) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Args:
+        llm_handler: The handler containing the model and tokenizer.
+        formatted_prompt: The input context.
+        target_text: The text we want to calculate probability/recall for.
+        
+    Returns:
+        Tuple of (target_logits, target_ids)
+        - target_logits: Logits used to predict the target tokens.
+        - target_ids: The ground truth token IDs of the target.
+    """
+    model = llm_handler.get_hf_model_for_scoring()
+    tokenizer = llm_handler.llm_tokenizer
+    device = llm_handler.device if llm_handler.llm_backend == "pt" else next(model.parameters()).device
+
+    # 1. Tokenize prompt ONLY to get its length (used for slicing later).
+    #    We must ensure special tokens are added to count the offset correctly.
+    prompt_tokens_temp = tokenizer(formatted_prompt, return_tensors="pt", add_special_tokens=True)
+    prompt_len = prompt_tokens_temp['input_ids'].shape[1]
+
+    # 2. Tokenize the FULL text (Prompt + Target).
+    #    This ensures subword merging at boundaries is handled correctly by the tokenizer.
+    full_text = formatted_prompt + target_text
+    full_tokens = tokenizer(full_text, return_tensors="pt", padding=False, truncation=True, add_special_tokens=True).to(device)
+
+    input_ids = full_tokens['input_ids']
+
+    # Safety check: if target was empty or truncated entirely
+    if input_ids.shape[1] <= prompt_len:
+        return torch.empty(0, device=device), torch.empty(0, device=device)
+
+    # 3. Forward Pass (Teacher Forcing)
+    with torch.no_grad():
+        with llm_handler._load_model_context():
+            outputs = model(input_ids=input_ids, attention_mask=full_tokens['attention_mask'])
+            all_logits = outputs.logits  # [1, seq_len, vocab_size]
+
+    # 4. Extract Logits and Labels
+    #    We need to predict `input_ids[i]`. The logit for this is at `all_logits[i-1]`.
+    #    Target starts at index `prompt_len`.
+    #    So we need logits from `prompt_len - 1` up to the second to last position.
+
+    target_logits = all_logits[0, prompt_len - 1:-1, :]  # [target_len, vocab_size]
+    target_ids = input_ids[0, prompt_len:]  # [target_len]
+
+    return target_logits, target_ids
+
+
+# ==============================================================================
+# Scoring Logic
+# ==============================================================================
+
+
+def _calculate_topk_recall(llm_handler,
+                           formatted_prompt: str,
+                           target_text: str,
+                           topk: int = 10) -> Tuple[float, Dict[int, float]]:
+    """
+    Calculate top-k recall for target text given prompt.
+    Checks if the ground truth token is within the top-k probabilities at each step.
+    """
+    # Use the fixed helper to get aligned logits/labels
+    pred_logits, target_ids = _get_logits_and_target_for_scoring(llm_handler, formatted_prompt, target_text)
+
+    if target_ids.shape[0] == 0:
+        return 0.0, {}
+
+    target_len = target_ids.shape[0]
+
+    # Get top-k indices for all positions at once
+    # topk_indices: [target_len, topk]
+    _, topk_indices = torch.topk(pred_logits, k=min(topk, pred_logits.shape[-1]), dim=-1)
+
+    recall_per_k = {}
+    position_scores = []
+
+    # Convert to list for faster CPU iteration
+    target_ids_list = target_ids.tolist()
+    topk_indices_list = topk_indices.tolist()
+
+    for k in range(1, topk + 1):
+        hits = 0
+        for pos in range(target_len):
+            gt_token = target_ids_list[pos]
+            # Check the top-k slice
+            topk_at_pos = topk_indices_list[pos][:k]
+
+            if gt_token in topk_at_pos:
+                hits += 1
+                # Calculate position-weighted score only once (when k=topk)
+                if k == topk:
+                    rank = topk_at_pos.index(gt_token) + 1
+                    # Rank 1 = 1.0, Rank k = small positive
+                    position_weight = 1.0 - (rank - 1) / topk
+                    position_scores.append(position_weight)
+
+        recall_per_k[k] = hits / target_len if target_len > 0 else 0.0
+
+    # Fill scores for positions where GT was NOT in top-k
+    while len(position_scores) < target_len:
+        position_scores.append(0.0)
+
+    average_recall = sum(position_scores) / len(position_scores) if position_scores else 0.0
+
+    return average_recall, recall_per_k
+
+
+def _calculate_metadata_recall(llm_handler,
+                               formatted_prompt: str,
+                               fields_dict: Dict[str, Any],
+                               topk: int = 10) -> Dict[str, float]:
+    """
+    Args:
+        fields_dict: Dictionary of {field_name: field_value}
+    """
+    if not fields_dict:
+        return {}
+
+    field_scores = {}
+
+    for field_name in sorted(fields_dict.keys()):
+        # Construct target text for this specific field
+        # e.g. <think>\nbpm: 120\n</think>\n
+        field_yaml = yaml.dump({field_name: fields_dict[field_name]}, allow_unicode=True, sort_keys=True).strip()
+        field_target_text = f"<think>\n{field_yaml}\n</think>\n"
+
+        # Calculate recall using the robust logic
+        avg_score, _ = _calculate_topk_recall(llm_handler, formatted_prompt, field_target_text, topk=topk)
+
+        field_scores[field_name] = avg_score
+        logger.debug(f"Recall for {field_name}: {avg_score:.4f}")
+
+    return field_scores
+
+
+def _calculate_log_prob(
+        llm_handler,
+        formatted_prompt: str,
+        target_text: str,
+        temperature: float = 1.0  # Kept for API compatibility, but ignored for scoring
+) -> float:
+    """
+    Calculate average log probability of target text given prompt.
+    """
+    pred_logits, target_ids = _get_logits_and_target_for_scoring(llm_handler, formatted_prompt, target_text)
+
+    if target_ids.shape[0] == 0:
+        return float('-inf')
+
+    # FIX: Do not divide by temperature.
+    # Log-probability for PMI/Perplexity should be exact.
+
+    # Calculate log probabilities (log_softmax)
+    log_probs = F.log_softmax(pred_logits, dim=-1)  # [target_len, vocab_size]
+
+    # Gather log probabilities of the ground truth tokens
+    target_log_probs = log_probs[torch.arange(target_ids.shape[0]), target_ids]
+
+    # Return average log probability
+    mean_log_prob = target_log_probs.mean().item()
+
+    return mean_log_prob
+
+
+def calculate_reward_score(
+    scores: Dict[str, float],
+    weights_config: Optional[Dict[str, float]] = None
+) -> Tuple[float, str]:
+    """
+    Reward Model Calculator: Computes a final reward based on user priorities.
+    
+    Priority Logic:
+        1. Caption (Highest): The overall vibe/style must match.
+        2. Lyrics (Medium): Content accuracy is important but secondary to vibe.
+        3. Metadata (Lowest): Technical constraints (BPM, Key) allow for slight deviations.
+    
+    Strategy: Dynamic Weighted Sum
+    - Metadata fields are aggregated into a single 'metadata' score first.
+    - Weights are dynamically renormalized if any component (e.g., lyrics) is missing.
+    
+    Args:
+        scores: Dictionary of raw scores (0.0 - 1.0) from the evaluation module.
+        weights_config: Optional custom weights. Defaults to:
+                        Caption (50%), Lyrics (30%), Metadata (20%).
+        
+    Returns:
+        final_reward: The calculated reward score (0.0 - 1.0).
+        explanation: A formatted string explaining how the score was derived.
+    """
+    
+    # 1. Default Preference Configuration
+    # These weights determine the relative importance of each component.
+    if weights_config is None:
+        weights_config = {
+            'caption': 0.50,  # High priority: Style/Vibe
+            'lyrics':  0.30,  # Medium priority: Content
+            'metadata': 0.20  # Low priority: Technical details
+        }
+    
+    # 2. Extract and Group Scores
+    # Caption and Lyrics are standalone high-level features.
+    caption_score = scores.get('caption')
+    lyrics_score = scores.get('lyrics')
+    
+    # Metadata fields (bpm, key, duration, etc.) are aggregated.
+    # We treat them as a single "Technical Score" to prevent them from 
+    # diluting the weight of Caption/Lyrics simply by having many fields.
+    meta_scores_list = [
+        val for key, val in scores.items() 
+        if key not in ['caption', 'lyrics']
+    ]
+    
+    # Calculate average of all metadata fields (if any exist)
+    meta_aggregate_score = None
+    if meta_scores_list:
+        meta_aggregate_score = sum(meta_scores_list) / len(meta_scores_list)
+    
+    # 3. specific Active Components & Dynamic Weighting
+    # We only include components that actually exist in this generation.
+    active_components = {}
+    
+    if caption_score is not None:
+        active_components['caption'] = (caption_score, weights_config['caption'])
+        
+    if lyrics_score is not None:
+        active_components['lyrics'] = (lyrics_score, weights_config['lyrics'])
+        
+    if meta_aggregate_score is not None:
+        active_components['metadata'] = (meta_aggregate_score, weights_config['metadata'])
+    
+    # 4. Calculate Final Weighted Score
+    total_base_weight = sum(w for _, w in active_components.values())
+    total_score = 0.0
+    
+    breakdown_lines = []
+    
+    if total_base_weight == 0:
+        return 0.0, "❌ No valid scores available to calculate reward."
+    
+    # Sort by weight (importance) for display
+    sorted_components = sorted(active_components.items(), key=lambda x: x[1][1], reverse=True)
+    
+    for name, (score, base_weight) in sorted_components:
+        # Renormalize weight: If lyrics are missing, caption/metadata weights scale up proportionately.
+        normalized_weight = base_weight / total_base_weight
+        weighted_contribution = score * normalized_weight
+        total_score += weighted_contribution
+        
+        breakdown_lines.append(
+            f"  • {name.title():<8} | Score: {score:.4f} | Weight: {normalized_weight:.2f} "
+            f"-> Contrib: +{weighted_contribution:.4f}"
+        )
+
+    return total_score, "\n".join(breakdown_lines)
+
+# ==============================================================================
+# Main Public API
+# ==============================================================================
+
+
+def calculate_pmi_score_per_condition(
+    llm_handler,
+    audio_codes: str,
+    caption: str = "",
+    lyrics: str = "",
+    metadata: Optional[Dict[str, Any]] = None,
+    temperature: float = 1.0,
+    topk: int = 10,
+    score_scale: float = 0.1,
+) -> Tuple[Dict[str, float], float, str]:
+    """
+    Calculate quality score separately for each condition.
+    - Metadata: Uses Top-k Recall.
+    - Caption/Lyrics: Uses PMI (Normalized).
+    """
+    if not llm_handler.llm_initialized:
+        return {}, 0.0, "❌ LLM not initialized"
+
+    if not audio_codes or not audio_codes.strip():
+        return {}, 0.0, "❌ No audio codes provided"
+
+    if "caption" not in metadata:
+        metadata['caption'] = caption
+
+    formatted_prompt = llm_handler.build_formatted_prompt_for_understanding(audio_codes=audio_codes, is_negative_prompt=False)
+    prompt_uncond = llm_handler.build_formatted_prompt_for_understanding(audio_codes="NO USER INPUT", is_negative_prompt=False)
+    try:
+        # 1. Calculate Recall for Metadata Fields
+        if metadata and isinstance(metadata, dict):
+            scores = {}
+            # Define which fields use which metric
+            metadata_recall_keys = ['bpm', 'duration', 'genres', 'keyscale', 'language', 'timesignature']
+            metadata_pmi_keys = ['caption']
+            for key in metadata_recall_keys:
+                if key in metadata and metadata[key] is not None:
+                    recall_metadata = {key: metadata[key]}
+                    field_scores = _calculate_metadata_recall(llm_handler, formatted_prompt, recall_metadata, topk=topk)
+                    scores.update(field_scores)
+
+            # 2. Calculate PMI for Caption
+            for key in metadata_pmi_keys:
+                if key in metadata and metadata[key] is not None:
+                    cot_yaml = yaml.dump({key: metadata[key]}, allow_unicode=True, sort_keys=True).strip()
+                    target_text = f"<think>\n{cot_yaml}\n</think>\n"
+
+                    log_prob_cond = _calculate_log_prob(llm_handler, formatted_prompt, target_text)
+                    log_prob_uncond = _calculate_log_prob(llm_handler, prompt_uncond, target_text)
+
+                    pmi_normalized = pmi_to_normalized_score(log_prob_cond - log_prob_uncond, scale=score_scale)
+                    scores[key] = pmi_normalized
+
+        # 3. Calculate PMI for Lyrics
+        if lyrics:
+            target_text = f"<think>\n</think>\n# Lyric\n{lyrics}\n"
+
+            log_prob_cond = _calculate_log_prob(llm_handler, formatted_prompt, target_text)
+
+            prompt_uncond = llm_handler.build_formatted_prompt_for_understanding(audio_codes="NO USER INPUT", is_negative_prompt=False)
+            log_prob_uncond = _calculate_log_prob(llm_handler, prompt_uncond, target_text)
+
+            scores['lyrics'] = pmi_to_normalized_score(log_prob_cond - log_prob_uncond, scale=score_scale)
+
+        if not scores:
+            return {}, 0.0, "❌ No conditions to evaluate"
+
+        # 4. Global Score
+        global_score = sum(scores.values()) / len(scores)
+        global_score, breakdown_lines = calculate_reward_score(scores)
+
+        # Status Message
+        status_lines = [breakdown_lines, "\n✅ Per-condition scores (0-1):"]
+        for key, score in sorted(scores.items()):
+            metric = "Top-k Recall" if key in metadata_recall_keys else "PMI (Norm)"
+            status_lines.append(f"  {key}: {score:.4f} ({metric})")
+        status = "\n".join(status_lines)
+        logger.info(f"Calculated scores: {global_score:.4f}\n{status}")
+        return scores, global_score, status
+
+    except Exception as e:
+        import traceback
+        error_msg = f"❌ Error: {str(e)}"
+        logger.error(error_msg)
+        logger.error(traceback.format_exc())
+        return {}, float('-inf'), error_msg

acestep/third_parts/nano-vllm/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Xingkai Yu
+
+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.

acestep/third_parts/nano-vllm/README.md

@@ -0,0 +1,66 @@
+<p align="center">
+<img width="300" src="assets/logo.png">
+</p>
+
+<p align="center">
+<a href="https://trendshift.io/repositories/15323" target="_blank"><img src="https://trendshift.io/api/badge/repositories/15323" alt="GeeeekExplorer%2Fnano-vllm | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
+</p>
+
+# Nano-vLLM
+
+A lightweight vLLM implementation built from scratch.
+
+## Key Features
+
+* 🚀 **Fast offline inference** - Comparable inference speeds to vLLM
+* 📖 **Readable codebase** - Clean implementation in ~ 1,200 lines of Python code
+* ⚡ **Optimization Suite** - Prefix caching, Tensor Parallelism, Torch compilation, CUDA graph, etc.
+
+## Installation
+
+```bash
+pip install git+https://github.com/GeeeekExplorer/nano-vllm.git
+```
+
+## Model Download
+
+To download the model weights manually, use the following command:
+```bash
+huggingface-cli download --resume-download Qwen/Qwen3-0.6B \
+  --local-dir ~/huggingface/Qwen3-0.6B/ \
+  --local-dir-use-symlinks False
+```
+
+## Quick Start
+
+See `example.py` for usage. The API mirrors vLLM's interface with minor differences in the `LLM.generate` method:
+```python
+from nanovllm import LLM, SamplingParams
+llm = LLM("/YOUR/MODEL/PATH", enforce_eager=True, tensor_parallel_size=1)
+sampling_params = SamplingParams(temperature=0.6, max_tokens=256)
+prompts = ["Hello, Nano-vLLM."]
+outputs = llm.generate(prompts, sampling_params)
+outputs[0]["text"]
+```
+
+## Benchmark
+
+See `bench.py` for benchmark.
+
+**Test Configuration:**
+- Hardware: RTX 4070 Laptop (8GB)
+- Model: Qwen3-0.6B
+- Total Requests: 256 sequences
+- Input Length: Randomly sampled between 100–1024 tokens
+- Output Length: Randomly sampled between 100–1024 tokens
+
+**Performance Results:**
+| Inference Engine | Output Tokens | Time (s) | Throughput (tokens/s) |
+|----------------|-------------|----------|-----------------------|
+| vLLM           | 133,966     | 98.37    | 1361.84               |
+| Nano-vLLM      | 133,966     | 93.41    | 1434.13               |
+
+
+## Star History
+
+[![Star History Chart](https://api.star-history.com/svg?repos=GeeeekExplorer/nano-vllm&type=Date)](https://www.star-history.com/#GeeeekExplorer/nano-vllm&Date)

acestep/third_parts/nano-vllm/bench.py

@@ -0,0 +1,32 @@
+import os
+import time
+from random import randint, seed
+from nanovllm import LLM, SamplingParams
+# from vllm import LLM, SamplingParams
+
+
+def main():
+    seed(0)
+    num_seqs = 256
+    max_input_len = 1024
+    max_ouput_len = 1024
+
+    path = os.path.expanduser("~/huggingface/Qwen3-0.6B/")
+    llm = LLM(path, enforce_eager=False, max_model_len=4096)
+
+    prompt_token_ids = [[randint(0, 10000) for _ in range(randint(100, max_input_len))] for _ in range(num_seqs)]
+    sampling_params = [SamplingParams(temperature=0.6, ignore_eos=True, max_tokens=randint(100, max_ouput_len)) for _ in range(num_seqs)]
+    # uncomment the following line for vllm
+    # prompt_token_ids = [dict(prompt_token_ids=p) for p in prompt_token_ids]
+
+    llm.generate(["Benchmark: "], SamplingParams())
+    t = time.time()
+    llm.generate(prompt_token_ids, sampling_params, use_tqdm=False)
+    t = (time.time() - t)
+    total_tokens = sum(sp.max_tokens for sp in sampling_params)
+    throughput = total_tokens / t
+    print(f"Total: {total_tokens}tok, Time: {t:.2f}s, Throughput: {throughput:.2f}tok/s")
+
+
+if __name__ == "__main__":
+    main()