feat: add start.sh for Linux/macOS, harden both launchers

start.sh mirrors start.bat: finds Python 3.10+, creates the venv, installs PyTorch (CUDA index on Linux, the MPS build on macOS, or the CPU index with --cpu), installs ACE-Step, reports the device and launches the UI. Same flags, and the defaults can also come from the environment (PORT=7870 ./start.sh). On macOS it passes --bf16 false, which the README already calls for. Dropped start.sh from .gitignore, where it sat among the upstream author's local scratch files. Both launchers also gain two fixes found while testing on WSL: - A venv is only accepted if pip works in it, not merely if the interpreter exists. A directory left by an interrupted install looked ready and then failed several steps later with a misleading "check your internet connection". Such a venv is now recreated, and if creation fails on Debian/Ubuntu the error points at python3-venv, which is the actual cause there. - The launch banner announced the URL as if the server were already up, while model loading still had a minute to go. It now says the interface will be available once "Running on local URL" appears. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
feat: generation options for infer.py, and split out training deps
2026-09-08 20:31:18 +03:00 · 2026-09-08 20:31:07 +03:00 · 2026-09-08 20:30:58 +03:00 · 2026-09-08 19:03:29 +03:00 · 2026-09-08 18:38:12 +03:00 · 2026-09-08 18:38:05 +03:00
83 changed files with 5490 additions and 585 deletions
@@ -172,7 +172,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
-*.txt
 !requirements.txt
 *.log
 *.flac
@@ -181,7 +180,6 @@ minio_config.yaml
 .history/*
 __pycache__/*
 train.log
-*.mp3
 *.tar.gz
 __pycache__/
 demo_examples/
@@ -202,3 +200,4 @@ ui/components_demo.py
 data_sampler_demo.py
 pipeline_ace_step_demo.py
 *.wav
+exps/*
@@ -38,6 +38,11 @@ RUN git clone https://github.com/ace-step/ACE-Step.git .
 RUN pip3 install --no-cache-dir --upgrade pip && \
    pip3 install --no-cache-dir hf_transfer peft && \
    pip3 install --no-cache-dir -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu126
+RUN pip3 install --no-cache-dir .
+
+# Ensure target directories for volumes exist and have correct initial ownership
+RUN mkdir -p /app/outputs /app/checkpoints /app/logs && \
+    chown -R appuser:appuser /app/outputs /app/checkpoints /app/logs

 # Change ownership of app files to appuser
 RUN chown -R appuser:appuser /app
@@ -5,10 +5,15 @@
    <a href="https://huggingface.co/ACE-Step/ACE-Step-v1-3.5B">Hugging Face</a> |
    <a href="https://modelscope.cn/models/ACE-Step/ACE-Step-v1-3.5B">ModelScope</a> |
    <a href="https://huggingface.co/spaces/ACE-Step/ACE-Step">Space Demo</a> |
-     <a href="https://discord.gg/rjAZz2xBdG">Discord</a> 
+    <a href="https://discord.gg/PeWDxrkdj7">Discord</a> |
+    <a href="https://arxiv.org/abs/2506.00045">Technical Report</a> |
+    <a href="https://ace-step.github.io/ace-step-v1.5.github.io/">ACE-Step v1.5</a>
+</p>
+
+<p align="center">
+    <b>Language:</b> <b>English</b> | <a href="./README_RU.md">Русский</a>
 </p>

---
 <p align="center">
    <img src="./assets/orgnization_logos.png" width="100%" alt="StepFun Logo">
 </p>
@@ -17,6 +22,7 @@

 - [✨ Features](#-features)
 - [📦 Installation](#-installation)
+- [⚡ Quick Start](#-quick-start)
 - [🚀 Usage](#-usage)
 - [📱 User Interface Guide](#-user-interface-guide)
 - [🔨 Train](#-train)
@@ -32,7 +38,45 @@ Rather than building yet another end-to-end text-to-music pipeline, our vision i

 ## 📢 News and Updates

- 🚀 **2025.05.07:** [ComfyUI_ACE-Step](https://github.com/billwuhao/ComfyUI_ACE-Step) node is now available! Explore the power of ACE-Step within ComfyUI. 🎉
+- 🎉 **2026.01.28:** Released [ACE-Step v1.5](https://ace-step.github.io/ace-step-v1.5.github.io/) - Our latest and most advanced model is now available!
+
+- 📃 2025.06.02: Released [ACE-Step Technical Report (PDF)](https://arxiv.org/abs/2506.00045).
+
+- 🎮 2025.05.14: Add `Stable Audio Open Small` sampler `pingpong`. Use SDE to achieve better music consistency and quality, including lyric alignment and style alignment. Use a better method to re-implement `Audio2Audio`
+
+- 🎤 2025.05.12: Release [RapMachine](https://huggingface.co/ACE-Step/ACE-Step-v1-chinese-rap-LoRA) and fix lora training issues
+    - See [ZH_RAP_LORA.md](./ZH_RAP_LORA.md) for details. Audio Examples: https://ace-step.github.io/#RapMachine
+    - See [TRAIN_INSTRUCTION.md](./TRAIN_INSTRUCTION.md) for detailed training instructions.
+
+<p align="center">
+    <img src="assets/rap_machine_demo.gif" alt="RapMachine Demo" width="45%">
+    <img src="assets/train_demo.gif" alt="Train Demo" width="50%">
+</p>
+
+
+- 🔥 **2025.05.10:** Memory Optimization Update
+  - Reduced Max VRAM to 8GB, making it more compatible with consumer devices
+  - Recommended launch options:
+    ```bash
+    acestep --torch_compile true --cpu_offload true --overlapped_decode true
+    ```
+    Windows need to install triton:
+    ```
+    pip install triton-windows
+    ```
+
+![image](./assets/cpu_offload_performance.png)
+
+- 📢 **2025.05.09:** Graidio Demo support Audio2Audio. ComfyUI: [Ace_Step_4x_a2a.json](./assets/Ace_Step_4x_a2a.json)
+<p align="center">
+    <img src="assets/audio2audio_demo.gif" alt="Audio2Audio Demo" width="50%">
+    <img src="assets/audio2audio_ComfyUI.png" alt="Audio2Audio ComfyUI" width="40%">
+</p>
+
+- 🚀 **2025.05.08:** [ComfyUI_ACE-Step](https://t.co/GeRSTrIvn0) node is now available! Explore the power of ACE-Step within ComfyUI. 🎉
+![image](https://github.com/user-attachments/assets/0a13d90a-9086-47ee-abab-976bad20fa7c)
+
+
 - 🚀 2025.05.06: Open source demo code and model


@@ -125,10 +169,11 @@ Rather than building yet another end-to-end text-to-music pipeline, our vision i

 - [x] Release training code 🔥
 - [x] Release LoRA training code 🔥
- [ ] Release RapMachine LoRA 🎤
+- [x] Release RapMachine LoRA 🎤
+- [x] Release evaluation performance and technical report 📄
+- [ ] Train and Release ACE-Step V1.5
 - [ ] Release ControlNet training code 🔥
 - [ ] Release Singing2Accompaniment ControlNet 🎮
- [ ] Release evaluation performance and technical report 📄


 ## 🖥️ Hardware Performance
@@ -215,9 +260,56 @@ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.o
 pip install -e .
 ```

+If you also intend to train or fine-tune, install the training extras as well (they are not needed for inference):
+```bash
+pip install -e ".[train]"
+```
+
 The ACE-Step application is now installed. The GUI works on Windows, macOS, and Linux. For instructions on how to run it, please see the [Usage](#-usage) section.


+## ⚡ Quick Start
+
+This repository ships with launcher scripts that do everything for you: check Python, create the virtual environment, install PyTorch with the right backend, install ACE-Step, and launch the web UI.
+
+**Windows** — double-click [`start.bat`](./start.bat), or run it from a command prompt:
+
+```bat
+start.bat
+```
+
+**Linux / macOS** — run [`start.sh`](./start.sh):
+
+```bash
+./start.sh
+```
+
+The first run takes a few minutes (roughly 3 GB of packages are downloaded). Model weights (~8 GB) are fetched automatically on the first generation. Later runs start immediately.
+
+### Script flags
+
+| Flag | What it does |
+| --- | --- |
+| `--lowvram` | Low-VRAM mode (~8 GB): enables `--cpu_offload`, `--overlapped_decode` and `--torch_compile`, and installs `triton-windows` |
+| `--cpu` | Run on CPU without CUDA (very slow, but works without an NVIDIA GPU) |
+| `--share` | Create a public Gradio link |
+| `--port <N>` | Web UI port (default 7865) |
+| `--device <N>` | GPU index (default 0) |
+| `--listen` | Bind to `0.0.0.0` so other devices on the LAN can connect |
+| `--reinstall` | Recreate the virtual environment from scratch |
+| `--update` | Update dependencies in the existing environment |
+| `--setup` | Install only, do not launch |
+| `--help` | Show usage |
+
+Flags can be combined, for example:
+
+```bat
+start.bat --lowvram --listen --port 7870
+```
+
+Defaults (port, GPU index, checkpoint path) live in the settings block at the top of each script. In `start.sh` they can also be overridden with environment variables (`PORT=7870 ./start.sh`). On macOS the scripts install the MPS build of PyTorch and pass `--bf16 false` automatically.
+
+
 ## 🚀 Usage

 ![Demo Interface](assets/demo_interface.png)
@@ -234,8 +326,28 @@ acestep --port 7865
 acestep --checkpoint_path /path/to/checkpoint --port 7865 --device_id 0 --share true --bf16 true
 ```

+* If `--checkpoint_path` is set and models exist at the path, load from `checkpoint_path`.
+* If `--checkpoint_path` is set but models do not exist at the path, auto download models to `checkpoint_path`.
+* If `--checkpoint_path` is not set, auto download models to the default path `~/.cache/ace-step/checkpoints`.
+
 If you are using macOS, please use `--bf16 false` to avoid errors.

+#### 🖥️ Command Line Generation
+
+To generate without the web UI, use `infer.py`:
+
+```bash
+python infer.py \
+  --prompt "synth-pop, female vocal, warm analog synths, 110 bpm" \
+  --lyrics_file my_song.txt \
+  --duration 120 --steps 60 --seed 7 \
+  --format mp3 --output_path outputs/my_song.mp3
+```
+
+Run `python infer.py --help` for the full list. The main options are `--prompt`, `--lyrics` / `--lyrics_file`, `--duration`, `--steps`, `--guidance_scale`, `--scheduler`, `--cfg_type`, `--omega_scale`, `--seed`, `--format` and `--output_path`; the runtime flags (`--bf16`, `--cpu_offload`, `--overlapped_decode`, `--torch_compile`, `--device_id`) match the ones the GUI takes.
+
+With no `--prompt`, the script keeps its original behaviour and generates from a random example in `examples/input_params`.
+
 #### 🔍 API Usage
 If you intend to integrate ACE-Step as a library into your own Python projects, you can install the latest version directly from GitHub using the following pip command.

@@ -256,7 +368,13 @@ If you intend to integrate ACE-Step as a library into your own Python projects,
 - `--device_id`: GPU device ID to use (default: 0)
 - `--share`: Enable Gradio sharing link (default: False)
 - `--bf16`: Use bfloat16 precision for faster inference (default: True)
- `--torch_compile`: Use `torch.compile()` to optimize the model, speeding up inference (default: False). **Not Supported on Windows**
+- `--torch_compile`: Use `torch.compile()` to optimize the model, speeding up inference (default: False). 
+  - **Windows need to install triton**:
+    ```
+    pip install triton-windows
+    ```
+- `--cpu_offload`: Offload model weights to CPU to save GPU memory (default: False)
+- `--overlapped_decode`: Use overlapped decoding to speed up inference (default: False)

 ## 📱 User Interface Guide

@@ -309,96 +427,7 @@ The `examples/input_params` directory contains sample input parameters that can
 </p>

 ## 🔨 Train
-
-### Prerequisites
-1. Prepare the environment as described in the installation section.
-
-2. If you plan to train a LoRA model, install the PEFT library:
-   ```bash
-   pip install peft
-   ```
-
-3. Prepare your dataset in Huggingface format ([Huggingface Datasets documentation](https://huggingface.co/docs/datasets/index)). The dataset should contain the following fields:
-   - `keys`: Unique identifier for each audio sample
-   - `filename`: Path to the audio file
-   - `tags`: List of descriptive tags (e.g., `["pop", "rock"]`)
-   - `norm_lyrics`: Normalized lyrics text
-   - Optional fields:
-     - `speaker_emb_path`: Path to speaker embedding file (use empty string if not available)
-     - `recaption`: Additional tag descriptions in various formats
-
-Example dataset entry:
-```json
-{
-	"keys": "1ce52937-cd1d-456f-967d-0f1072fcbb58",
-	"filename": "data/audio/1ce52937-cd1d-456f-967d-0f1072fcbb58.wav",
-	"tags": ["pop", "acoustic", "ballad", "romantic", "emotional"],
-	"speaker_emb_path": "",
-	"norm_lyrics": "I love you, I love you, I love you",
-	"recaption": {
-		"simplified": "pop",
-		"expanded": "pop, acoustic, ballad, romantic, emotional",
-		"descriptive": "The sound is soft and gentle, like a tender breeze on a quiet evening. It's soothing and full of longing.",
-		"use_cases": "Suitable for background music in romantic films or during intimate moments.",
-		"analysis": "pop, ballad, piano, guitar, slow tempo, romantic, emotional"
-	}
-}
-```
-
-### Training Parameters
-
-#### Common Parameters
- `--dataset_path`: Path to your Huggingface dataset (required)
- `--checkpoint_dir`: Directory containing the base model checkpoint
- `--learning_rate`: Learning rate for training (default: 1e-4)
- `--max_steps`: Maximum number of training steps (default: 2000000)
- `--precision`: Training precision, e.g., "bf16-mixed" (default) or "fp32"
- `--devices`: Number of GPUs to use (default: 1)
- `--num_nodes`: Number of compute nodes to use (default: 1)
- `--accumulate_grad_batches`: Gradient accumulation steps (default: 1)
- `--num_workers`: Number of data loading workers (default: 8)
- `--every_n_train_steps`: Checkpoint saving frequency (default: 2000)
- `--every_plot_step`: Frequency of generating evaluation samples (default: 2000)
- `--exp_name`: Experiment name for logging (default: "text2music_train_test")
- `--logger_dir`: Directory for saving logs (default: "./exps/logs/")
-
-#### Base Model Training
-Train the base model with:
-```bash
-python trainer.py --dataset_path "path/to/your/dataset" --checkpoint_dir "path/to/base/checkpoint" --exp_name "your_experiment_name"
-```
-
-#### LoRA Training
-For LoRA training, you need to provide a LoRA configuration file:
-```bash
-python trainer.py --dataset_path "path/to/your/dataset" --checkpoint_dir "path/to/base/checkpoint" --lora_config_path "path/to/lora_config.json" --exp_name "your_lora_experiment"
-```
-
-Example LoRA configuration file (lora_config.json):
-```json
-{
-	"r": 16,
-	"lora_alpha": 32,
-	"target_modules": [
-		"speaker_embedder",
-		"linear_q",
-		"linear_k",
-		"linear_v",
-		"to_q",
-		"to_k",
-		"to_v",
-		"to_out.0"
-	]
-}
-```
-
-### Advanced Training Options
- `--shift`: Flow matching shift parameter (default: 3.0)
- `--gradient_clip_val`: Gradient clipping value (default: 0.5)
- `--gradient_clip_algorithm`: Gradient clipping algorithm (default: "norm")
- `--reload_dataloaders_every_n_epochs`: Frequency to reload dataloaders (default: 1)
- `--val_check_interval`: Validation check interval (default: None)
-
+See [TRAIN_INSTRUCTION.md](./TRAIN_INSTRUCTION.md) for detailed training instructions.

 ## 📜 License & Disclaimer

@@ -0,0 +1,463 @@
+<h1 align="center">ACE-Step</h1>
+<h1 align="center">Шаг к фундаментальной модели генерации музыки</h1>
+<p align="center">
+    <a href="https://ace-step.github.io/">Проект</a> |
+    <a href="https://huggingface.co/ACE-Step/ACE-Step-v1-3.5B">Hugging Face</a> |
+    <a href="https://modelscope.cn/models/ACE-Step/ACE-Step-v1-3.5B">ModelScope</a> |
+    <a href="https://huggingface.co/spaces/ACE-Step/ACE-Step">Демо (Space)</a> |
+    <a href="https://discord.gg/PeWDxrkdj7">Discord</a> |
+    <a href="https://arxiv.org/abs/2506.00045">Технический отчёт</a> |
+    <a href="https://ace-step.github.io/ace-step-v1.5.github.io/">ACE-Step v1.5</a>
+</p>
+
+<p align="center">
+    <b>Язык:</b> <a href="./README.md">English</a> | <b>Русский</b>
+</p>
+
+<p align="center">
+    <img src="./assets/orgnization_logos.png" width="100%" alt="StepFun Logo">
+</p>
+
+## Содержание
+
+- [✨ Возможности](#-возможности)
+- [📦 Установка](#-установка)
+- [⚡ Быстрый старт](#-быстрый-старт)
+- [🚀 Использование](#-использование)
+- [📱 Описание интерфейса](#-описание-интерфейса)
+- [🔨 Обучение](#-обучение)
+
+## 📝 Аннотация
+
+Мы представляем ACE-Step — новую открытую фундаментальную модель для генерации музыки, которая преодолевает ключевые ограничения существующих подходов и достигает state-of-the-art качества за счёт целостной архитектуры. Нынешние методы вынуждены искать компромисс между скоростью генерации, музыкальной связностью и управляемостью. Например, модели на базе LLM (Yue, SongGen) хорошо попадают в текст песни, но страдают от медленного инференса и структурных артефактов. Диффузионные модели (например, DiffRhythm), напротив, синтезируют быстрее, но часто теряют структурную связность на длинных отрезках.
+
+ACE-Step закрывает этот разрыв, объединяя диффузионную генерацию с Deep Compression AutoEncoder (DCAE) из Sana и лёгким линейным трансформером. Дополнительно модель использует MERT и m-hubert для выравнивания семантических представлений (REPA) во время обучения, что обеспечивает быструю сходимость. В результате модель синтезирует до 4 минут музыки за 20 секунд на GPU A100 — в 15 раз быстрее решений на базе LLM — при этом превосходя их по музыкальной связности и попаданию в текст по метрикам мелодии, гармонии и ритма. Кроме того, ACE-Step сохраняет тонкие акустические детали, что позволяет реализовать продвинутые механизмы управления: клонирование голоса, редактирование текста, ремиксы и генерацию отдельных дорожек (например, lyric2vocal, singing2accompaniment).
+
+Вместо очередного end-to-end пайплайна text-to-music наша цель — создать фундаментальную модель для музыкального ИИ: быструю, универсальную, эффективную и при этом гибкую архитектуру, поверх которой легко обучать подзадачи. Это открывает путь к мощным инструментам, органично встраивающимся в творческий процесс музыкантов, продюсеров и авторов контента. Коротко говоря, мы хотим повторить для музыки то, чем стал Stable Diffusion для изображений.
+
+
+## 📢 Новости и обновления
+
+- 🎉 **28.01.2026:** Вышла [ACE-Step v1.5](https://ace-step.github.io/ace-step-v1.5.github.io/) — наша самая свежая и продвинутая модель!
+
+- 📃 02.06.2025: Опубликован [технический отчёт ACE-Step (PDF)](https://arxiv.org/abs/2506.00045).
+
+- 🎮 14.05.2025: Добавлен сэмплер `pingpong` из `Stable Audio Open Small`. Использование SDE даёт лучшую консистентность и качество музыки, включая попадание в текст и в стиль. Также заново реализован `Audio2Audio` более удачным способом.
+
+- 🎤 12.05.2025: Выпущен [RapMachine](https://huggingface.co/ACE-Step/ACE-Step-v1-chinese-rap-LoRA), исправлены проблемы обучения LoRA
+    - Подробности в [ZH_RAP_LORA.md](./ZH_RAP_LORA.md). Примеры аудио: https://ace-step.github.io/#RapMachine
+    - Подробная инструкция по обучению — в [TRAIN_INSTRUCTION.md](./TRAIN_INSTRUCTION.md).
+
+<p align="center">
+    <img src="assets/rap_machine_demo.gif" alt="RapMachine Demo" width="45%">
+    <img src="assets/train_demo.gif" alt="Train Demo" width="50%">
+</p>
+
+
+- 🔥 **10.05.2025:** Оптимизация потребления памяти
+  - Максимальный расход VRAM снижен до 8 ГБ — модель стала доступнее для домашних видеокарт
+  - Рекомендуемые параметры запуска:
+    ```bash
+    acestep --torch_compile true --cpu_offload true --overlapped_decode true
+    ```
+    На Windows нужно установить triton:
+    ```
+    pip install triton-windows
+    ```
+
+![image](./assets/cpu_offload_performance.png)
+
+- 📢 **09.05.2025:** Демо на Gradio поддерживает Audio2Audio. ComfyUI: [Ace_Step_4x_a2a.json](./assets/Ace_Step_4x_a2a.json)
+<p align="center">
+    <img src="assets/audio2audio_demo.gif" alt="Audio2Audio Demo" width="50%">
+    <img src="assets/audio2audio_ComfyUI.png" alt="Audio2Audio ComfyUI" width="40%">
+</p>
+
+- 🚀 **08.05.2025:** Доступен узел [ComfyUI_ACE-Step](https://t.co/GeRSTrIvn0)! Используйте возможности ACE-Step прямо в ComfyUI. 🎉
+![image](https://github.com/user-attachments/assets/0a13d90a-9086-47ee-abab-976bad20fa7c)
+
+
+- 🚀 06.05.2025: Открыты исходный код демо и модель
+
+
+## ✨ Возможности
+
+<p align="center">
+    <img src="./assets/application_map.png" width="100%" alt="ACE-Step Framework">
+</p>
+
+### 🎯 Базовое качество
+
+#### 🌈 Разнообразие стилей и жанров
+
+- 🎸 Поддерживаются все основные музыкальные стили, описание задаётся короткими тегами, развёрнутым текстом или описанием сценария использования
+- 🎷 Генерация музыки в разных жанрах с подходящим инструментарием и стилистикой
+
+#### 🌍 Многоязычность
+
+- 🗣️ Поддерживается 19 языков, из них 10 с наилучшим качеством:
+  - 🇺🇸 английский, 🇨🇳 китайский, 🇷🇺 русский, 🇪🇸 испанский, 🇯🇵 японский, 🇩🇪 немецкий, 🇫🇷 французский, 🇵🇹 португальский, 🇮🇹 итальянский, 🇰🇷 корейский
+- ⚠️ Из-за дисбаланса обучающих данных менее распространённые языки могут работать хуже
+
+#### 🎻 Инструментальные стили
+
+- 🎹 Поддерживается генерация инструментальной музыки в разных жанрах и стилях
+- 🎺 Реалистичные инструментальные дорожки с корректным тембром и выразительностью каждого инструмента
+- 🎼 Возможны сложные аранжировки с несколькими инструментами при сохранении музыкальной связности
+
+#### 🎤 Вокальные техники
+
+- 🎙️ Качественная передача различных вокальных стилей и техник
+- 🗣️ Поддержка разной вокальной подачи, включая разные приёмы и манеры пения
+
+### 🎛️ Управляемость
+
+#### 🔄 Генерация вариаций
+
+- ⚙️ Реализовано через оптимизацию на этапе инференса, без дообучения
+- 🌊 Flow-matching модель генерирует начальный шум, затем по формуле шума из trigFlow добавляется дополнительный гауссов шум
+- 🎚️ Соотношение исходного и нового шума регулируется — так задаётся степень отличия вариации
+
+#### 🎨 Перерисовка (Repainting)
+
+- 🖌️ Реализовано добавлением шума к целевому аудио и наложением масочных ограничений в процессе ODE
+- 🔍 Если условия генерации меняются относительно исходных, можно изменить только отдельные аспекты, сохранив остальное
+- 🔀 Комбинируется с генерацией вариаций — можно делать локальные вариации стиля, текста или вокала
+
+#### ✏️ Редактирование текста песни
+
+- 💡 Технология flow-edit применена для локального изменения текста с сохранением мелодии, вокала и аккомпанемента
+- 🔄 Работает и со сгенерированным, и с загруженным аудио, что заметно расширяет творческие возможности
+- ℹ️ Текущее ограничение: за раз можно менять только небольшие фрагменты текста, иначе появляются искажения; но правки можно применять последовательно
+
+### 🚀 Применения
+
+#### 🎤 Lyric2Vocal (LoRA)
+
+- 🔊 LoRA, дообученная на чистом вокале, позволяет генерировать вокальные сэмплы прямо из текста
+- 🛠️ Практическое применение: вокальные демо, гайд-треки, помощь в написании песен, эксперименты с вокальной аранжировкой
+- ⏱️ Быстрый способ проверить, как текст зазвучит в исполнении, — ускоряет итерации автора
+
+#### 📝 Text2Samples (LoRA)
+
+- 🎛️ Аналог Lyric2Vocal, но дообучен на чисто инструментальных данных и сэмплах
+- 🎵 Генерация концептуальных сэмплов для музыкального продакшена по текстовому описанию
+- 🧰 Удобно для быстрого создания инструментальных лупов, звуковых эффектов и музыкальных элементов
+
+### 🔮 Скоро
+
+#### 🎤 RapMachine
+
+- 🔥 Дообучение на чистых рэп-данных для создания ИИ, специализирующегося на рэпе
+- 🏆 Ожидаемые возможности: ИИ-баттлы и повествование через рэп
+- 📚 Рэп обладает исключительными нарративными и выразительными возможностями — потенциал применения огромен
+
+#### 🎛️ StemGen
+
+- 🎚️ ControlNet-LoRA, обученная на многодорожечных данных для генерации отдельных инструментальных стемов
+- 🎯 На вход подаётся референсный трек и нужный инструмент (или референсное аудио инструмента)
+- 🎹 На выходе — стем инструмента, дополняющий референс: например, фортепианный аккомпанемент к мелодии флейты или джазовые барабаны к соло-гитаре
+
+#### 🎤 Singing2Accompaniment
+
+- 🔄 Обратный процесс к StemGen: из одной вокальной дорожки собирается сведённый мастер-трек
+- 🎵 На вход подаётся вокал и нужный стиль, на выходе — полноценный аккомпанемент к вокалу
+- 🎸 Создаётся полное инструментальное сопровождение — легко добавить профессиональное звучание к любой вокальной записи
+
+
+## 📋 Дорожная карта
+
+- [x] Публикация кода обучения 🔥
+- [x] Публикация кода обучения LoRA 🔥
+- [x] Публикация RapMachine LoRA 🎤
+- [x] Публикация результатов оценки и технического отчёта 📄
+- [ ] Обучение и публикация ACE-Step V1.5
+- [ ] Публикация кода обучения ControlNet 🔥
+- [ ] Публикация Singing2Accompaniment ControlNet 🎮
+
+
+## 🖥️ Производительность на разном железе
+
+Мы измерили производительность ACE-Step на разных конфигурациях:
+
+| Устройство      | RTF (27 шагов) | Время на 1 мин аудио (27 шагов) | RTF (60 шагов) | Время на 1 мин аудио (60 шагов) |
+| --------------- | -------------- | ------------------------------- | -------------- | ------------------------------- |
+| NVIDIA RTX 4090 | 34.48 ×        | 1.74 с                          | 15.63 ×        | 3.84 с                          |
+| NVIDIA A100     | 27.27 ×        | 2.20 с                          | 12.27 ×        | 4.89 с                          |
+| NVIDIA RTX 3090 | 12.76 ×        | 4.70 с                          | 6.48 ×         | 9.26 с                          |
+| MacBook M2 Max  | 2.27 ×         | 26.43 с                         | 1.03 ×         | 58.25 с                         |
+
+
+Производительность измеряется в RTF (Real-Time Factor, коэффициент реального времени). Чем больше значение, тем быстрее генерация. 27.27× означает, что на 1 минуту музыки уходит 2.2 секунды (60/27.27). Измерения проводились на одном GPU с batch size 1 и 27 шагами.
+
+
+## 📦 Установка
+
+### 1. Клонирование репозитория
+Сначала склонируйте репозиторий ACE-Step и перейдите в каталог проекта:
+```bash
+git clone https://github.com/ace-step/ACE-Step.git
+cd ACE-Step
+```
+
+### 2. Требования
+Убедитесь, что у вас установлено:
+
+* `Python`: рекомендуется версия 3.10 или новее. Скачать можно на [python.org](https://www.python.org/).
+* `Conda` или `venv`: для создания виртуального окружения (Conda предпочтительнее).
+
+### 3. Создание виртуального окружения
+
+Настоятельно рекомендуем использовать виртуальное окружение, чтобы не конфликтовать с другими пакетами. Выберите один из вариантов:
+
+#### Вариант A: Conda
+
+1.  **Создайте окружение** с именем `ace_step` и Python 3.10:
+    ```bash
+    conda create -n ace_step python=3.10 -y
+    ```
+
+2.  **Активируйте окружение:**
+    ```bash
+    conda activate ace_step
+    ```
+
+#### Вариант B: venv
+
+1.  **Перейдите в каталог склонированного репозитория ACE-Step.**
+
+2.  **Создайте виртуальное окружение** (обычно его называют `venv`):
+    ```bash
+    python -m venv venv
+    ```
+
+3.  **Активируйте окружение:**
+    * **Windows (cmd.exe):**
+        ```bash
+        venv\Scripts\activate.bat
+        ```
+    * **Windows (PowerShell):**
+        ```powershell
+        .\venv\Scripts\Activate.ps1
+        ```
+        *(Если возникает ошибка политики выполнения, сначала выполните `Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope Process`)*
+    * **Linux / macOS (bash/zsh):**
+        ```bash
+        source venv/bin/activate
+        ```
+
+### 4. Установка зависимостей
+После активации виртуального окружения:
+
+**a.** (Только Windows) Если вы на Windows и планируете использовать NVIDIA GPU, сначала поставьте PyTorch со сборкой под CUDA:
+
+```bash
+pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
+```
+(Замените `cu126`, если у вас другая версия CUDA. Другие варианты установки — на [официальном сайте PyTorch](https://pytorch.org/get-started/locally/)).
+
+**b.** Установите ACE-Step и основные зависимости:
+```bash
+pip install -e .
+```
+
+Если планируете обучать или дообучать модель, поставьте дополнительно зависимости для обучения — для инференса они не нужны:
+```bash
+pip install -e ".[train]"
+```
+
+На этом установка завершена. Графический интерфейс работает на Windows, macOS и Linux. Как запускать — см. раздел [Использование](#-использование).
+
+
+## ⚡ Быстрый старт
+
+В репозитории есть скрипты запуска, которые делают всё сами: проверяют Python, создают виртуальное окружение, ставят PyTorch с нужным бэкендом, устанавливают ACE-Step и запускают веб-интерфейс.
+
+**Windows** — достаточно дважды кликнуть по [`start.bat`](./start.bat) или запустить из командной строки:
+
+```bat
+start.bat
+```
+
+**Linux / macOS** — запустите [`start.sh`](./start.sh):
+
+```bash
+./start.sh
+```
+
+При первом запуске установка займёт несколько минут (скачивается около 3 ГБ пакетов). Веса модели (~8 ГБ) докачаются автоматически при первой генерации. Последующие запуски стартуют сразу.
+
+### Флаги скрипта
+
+| Флаг | Что делает |
+| --- | --- |
+| `--lowvram` | Режим экономии видеопамяти (до ~8 ГБ VRAM): включает `--cpu_offload`, `--overlapped_decode` и `--torch_compile`, доустанавливает `triton-windows` |
+| `--cpu` | Запуск на процессоре, без CUDA (очень медленно, но работает без видеокарты NVIDIA) |
+| `--share` | Публичная ссылка Gradio для доступа снаружи |
+| `--port <N>` | Порт веб-интерфейса (по умолчанию 7865) |
+| `--device <N>` | Номер видеокарты (по умолчанию 0) |
+| `--listen` | Слушать `0.0.0.0`, чтобы зайти с других устройств в локальной сети |
+| `--reinstall` | Полностью пересоздать виртуальное окружение с нуля |
+| `--update` | Обновить зависимости в существующем окружении |
+| `--setup` | Только установка, без запуска |
+| `--help` | Показать справку |
+
+Флаги можно комбинировать, например:
+
+```bat
+start.bat --lowvram --listen --port 7870
+```
+
+Настройки по умолчанию (порт, номер GPU, путь к весам модели) задаются в блоке `==== НАСТРОЙКИ ====` в начале каждого скрипта. В `start.sh` их можно переопределить переменными окружения (`PORT=7870 ./start.sh`). На macOS скрипт сам ставит сборку PyTorch с MPS и передаёт `--bf16 false`.
+
+
+## 🚀 Использование
+
+![Demo Interface](assets/demo_interface.png)
+
+### 🔍 Базовый запуск
+
+```bash
+acestep --port 7865
+```
+
+### ⚙️ Расширенный запуск
+
+```bash
+acestep --checkpoint_path /path/to/checkpoint --port 7865 --device_id 0 --share true --bf16 true
+```
+
+* Если `--checkpoint_path` задан и модели по этому пути есть, они загружаются оттуда.
+* Если `--checkpoint_path` задан, но моделей там нет, они автоматически скачаются в этот каталог.
+* Если `--checkpoint_path` не задан, модели скачаются в путь по умолчанию `~/.cache/ace-step/checkpoints`.
+
+На macOS используйте `--bf16 false`, чтобы избежать ошибок.
+
+#### 🖥️ Генерация из командной строки
+
+Чтобы генерировать без веб-интерфейса, используйте `infer.py`:
+
+```bash
+python infer.py \
+  --prompt "synth-pop, female vocal, warm analog synths, 110 bpm" \
+  --lyrics_file my_song.txt \
+  --duration 120 --steps 60 --seed 7 \
+  --format mp3 --output_path outputs/my_song.mp3
+```
+
+Полный список — `python infer.py --help`. Основные опции: `--prompt`, `--lyrics` / `--lyrics_file`, `--duration`, `--steps`, `--guidance_scale`, `--scheduler`, `--cfg_type`, `--omega_scale`, `--seed`, `--format`, `--output_path`. Флаги режима работы (`--bf16`, `--cpu_offload`, `--overlapped_decode`, `--torch_compile`, `--device_id`) те же, что у графического интерфейса.
+
+Без `--prompt` скрипт ведёт себя как раньше и генерирует по случайному примеру из `examples/input_params`.
+
+#### 🔍 Использование как библиотеки
+Если вы хотите встроить ACE-Step как библиотеку в собственный Python-проект, можно поставить последнюю версию прямо из GitHub.
+
+**Установка через pip:**
+
+1.  **Убедитесь, что установлен Git:** этот способ требует наличия Git в системе и в переменной PATH.
+2.  **Выполните команду установки:**
+    ```bash
+    pip install git+https://github.com/ace-step/ACE-Step.git
+    ```
+    Рекомендуется выполнять её внутри виртуального окружения, чтобы не ломать другие пакеты.
+
+#### 🛠️ Аргументы командной строки
+
+- `--checkpoint_path`: путь к весам модели (по умолчанию скачиваются автоматически)
+- `--server_name`: IP-адрес или имя хоста, на котором слушает сервер Gradio (по умолчанию `127.0.0.1`). Укажите `0.0.0.0`, чтобы открыть доступ с других устройств в сети.
+- `--port`: порт сервера Gradio (по умолчанию 7865)
+- `--device_id`: номер GPU (по умолчанию 0)
+- `--share`: включить публичную ссылку Gradio (по умолчанию False)
+- `--bf16`: использовать точность bfloat16 для ускорения инференса (по умолчанию True)
+- `--torch_compile`: использовать `torch.compile()` для оптимизации модели и ускорения инференса (по умолчанию False).
+  - **На Windows нужен triton**:
+    ```
+    pip install triton-windows
+    ```
+- `--cpu_offload`: выгружать веса модели в оперативную память для экономии видеопамяти (по умолчанию False)
+- `--overlapped_decode`: перекрывающееся декодирование для ускорения инференса (по умолчанию False)
+
+## 📱 Описание интерфейса
+
+Интерфейс ACE-Step разделён на вкладки под разные задачи генерации и редактирования:
+
+### 📝 Вкладка Text2Music
+
+1. **📋 Поля ввода**:
+   - **🏷️ Tags**: описательные теги, жанры или описание сцены через запятую
+   - **📜 Lyrics**: текст песни со структурными тегами вроде [verse], [chorus], [bridge]
+   - **⏱️ Audio Duration**: желаемая длительность аудио (-1 — случайная)
+
+2. **⚙️ Настройки**:
+   - **🔧 Basic Settings**: количество шагов инференса, guidance scale, сиды
+   - **🔬 Advanced Settings**: тонкая настройка типа планировщика, типа CFG, параметров ERG и прочего
+
+3. **🚀 Генерация**: нажмите «Generate», чтобы создать музыку по введённым данным
+
+### 🔄 Вкладка Retake
+
+- 🎲 Повторная генерация с небольшими отличиями за счёт других сидов
+- 🎚️ Параметр variance задаёт, насколько результат будет отличаться от оригинала
+
+### 🎨 Вкладка Repainting
+
+- 🖌️ Выборочная перегенерация отдельных фрагментов трека
+- ⏱️ Задаются время начала и конца перерисовываемого участка
+- 🔍 Источник аудио выбирается: результат text2music, последняя перерисовка или загруженный файл
+
+### ✏️ Вкладка Edit
+
+- 🔄 Изменение готовой музыки через правку тегов или текста
+- 🎛️ Режим «only_lyrics» сохраняет мелодию, режим «remix» её меняет
+- 🎚️ Параметры редактирования задают, насколько сохраняется оригинал
+
+### 📏 Вкладка Extend
+
+- ➕ Добавление музыки в начало или в конец существующего трека
+- 📐 Задаются длины расширения слева и справа
+- 🔍 Выбирается источник аудио для расширения
+
+## 📂 Примеры
+
+В каталоге `examples/input_params` лежат примеры входных параметров — их можно использовать как образец для генерации.
+
+## 🏗️ Архитектура
+
+<p align="center">
+    <img src="./assets/ACE-Step_framework.png" width="100%" alt="ACE-Step Framework">
+</p>
+
+## 🔨 Обучение
+Подробная инструкция — в [TRAIN_INSTRUCTION.md](./TRAIN_INSTRUCTION.md).
+
+## 📜 Лицензия и отказ от ответственности
+
+Проект распространяется по лицензии [Apache License 2.0](./LICENSE)
+
+ACE-Step позволяет создавать оригинальную музыку в самых разных жанрах и применим в творческом продакшене, образовании и развлечениях. Модель создавалась для позитивных и художественных сценариев использования, но мы осознаём и риски: непреднамеренное нарушение авторских прав из-за стилистического сходства, некорректное смешение культурных элементов, а также использование для генерации вредоносного контента. Для ответственного использования мы призываем проверять оригинальность полученных работ, явно указывать участие ИИ и получать необходимые разрешения при адаптации защищённых стилей или материалов. Используя ACE-Step, вы соглашаетесь придерживаться этих принципов и уважать художественную целостность, культурное разнообразие и требования законодательства. Авторы не несут ответственности за неправомерное использование модели, включая, помимо прочего, нарушение авторских прав, культурную бестактность или создание вредоносного контента.
+
+🔔 Важное замечание
+Единственный официальный сайт проекта ACE-Step — наша страница на GitHub Pages.
+Никаких других сайтов мы не ведём.
+🚫 Поддельные домены включают (но не ограничиваются ими):
+ac\*\*p.com, a\*\*p.org, a\*\*\*c.org
+⚠️ Будьте осторожны. Не заходите на эти сайты, не доверяйте им и не совершайте на них платежей.
+
+## 🙏 Благодарности
+
+Проект развивается совместно ACE Studio и StepFun.
+
+
+## 📖 Цитирование
+
+Если проект оказался полезен для вашего исследования, пожалуйста, сошлитесь на него:
+
+```BibTeX
+@misc{gong2025acestep,
+	title={ACE-Step: A Step Towards Music Generation Foundation Model},
+	author={Junmin Gong, Wenxiao Zhao, Sen Wang, Shengyuan Xu, Jing Guo}, 
+	howpublished={\url{https://github.com/ace-step/ACE-Step}},
+	year={2025},
+	note={GitHub repository}
+}
+```
@@ -0,0 +1,166 @@
+# Training Instruction
+
+## 0. Install the Training Dependencies
+
+Training needs a few packages that a plain inference install does not pull in
+(`datasets`, `pytorch_lightning`, `matplotlib`, `tensorboard`, `tensorboardX`).
+Install them with the `train` extra:
+
+```bash
+pip install -e ".[train]"
+```
+
+## 1. Data Preparation
+
+### Required File Format
+For each audio sample, you need **exactly 3 files** in the `data` directory:
+
+1. **`filename.mp3`** - The audio file
+2. **`filename_prompt.txt`** - Audio characteristics (comma-separated tags)  
+3. **`filename_lyrics.txt`** - Song lyrics (optional, but recommended)
+
+### Example Data Structure
+```
+data/
+├── test_track_001.mp3
+├── test_track_001_prompt.txt
+└── test_track_001_lyrics.txt
+```
+
+### File Content Format
+
+#### `*_prompt.txt` - Audio Tags
+Simple comma-separated audio characteristics describing the sound, instruments, genre, mood, etc.
+
+**Example:**
+```
+melodic techno, male vocal, electronic, emotional, minor key, 124 bpm, synthesizer, driving, atmospheric
+```
+
+**Guidelines for creating prompt tags:**
+- Include **genre** (e.g., "rap", "pop", "rock", "electronic")
+- Include **vocal type** (e.g., "male vocal", "female vocal", "spoken word")
+- Include **instruments** actually heard (e.g., "guitar", "piano", "synthesizer", "drums")
+- Include **mood/energy** (e.g., "energetic", "calm", "aggressive", "melancholic")  
+- Include **tempo** if known (e.g., "120 bpm", "fast tempo", "slow tempo")
+- Include **key** if known (e.g., "major key", "minor key", "C major")
+
+#### `*_lyrics.txt` - Song Lyrics
+Standard song lyrics with verse/chorus structure.
+
+**Example:**
+```
+[Verse]
+Lately I've been wondering
+Why do I do this to myself
+I should be over it
+
+[Chorus]  
+It makes me want to cry
+If you knew what you meant to me
+I wonder if you'd come back
+```
+
+### ⚠️ Important Notes
+- **File naming is strict**: Must follow `filename.mp3`, `filename_prompt.txt`, `filename_lyrics.txt` pattern
+- **JSON files are NOT supported** - the converter only reads the simple text files above
+- **Complex multi-variant descriptions are NOT used** - only the simple comma-separated prompt format works
+
+## 2. Convert to Huggingface Dataset Format
+
+Run the following command to convert your data to the training format:
+
+```bash
+python convert2hf_dataset.py --data_dir "./data" --repeat_count 2000 --output_name "zh_lora_dataset"
+```
+
+**Parameters:**
+- `--data_dir`: Path to your data directory containing the MP3, prompt, and lyrics files
+- `--repeat_count`: Number of times to repeat your data (use higher values for small datasets)  
+- `--output_name`: Name of the output dataset directory
+
+### What the Converter Creates
+
+The converter processes your files and creates a Huggingface dataset with these features:
+
+```python
+Dataset Features:
+{
+    'keys': string,              # filename (e.g., "test_track_001")
+    'filename': string,          # path to MP3 file  
+    'tags': list[string],        # parsed prompt tags as array
+    'speaker_emb_path': string,  # (empty, not used)
+    'norm_lyrics': string,       # full lyrics text
+    'recaption': dict            # (empty, not used)
+}
+```
+
+**Example processed sample:**
+```python
+{
+    'keys': 'test_track_001',
+    'filename': 'data/test_track_001.mp3',
+    'tags': ['melodic techno', 'male vocal', 'electronic', 'emotional', 'minor key', '124 bpm', 'synthesizer', 'driving', 'atmospheric'],
+    'speaker_emb_path': '',
+    'norm_lyrics': '[Verse]\nLately I\'ve been wondering\nWhy do I do this to myself...',
+    'recaption': {}
+}
+```
+
+## 3. Configure Lora Parameters
+Refer to `config/zh_rap_lora_config.json` for configuring Lora parameters.
+
+If your VRAM is not sufficient, you can reduce the `r` and `lora_alpha` parameters in the configuration file. Such as:
+```json
+{
+	"r": 16,
+	"lora_alpha": 32,
+	"target_modules": [
+		"linear_q",
+		"linear_k",
+		"linear_v",
+		"to_q",
+		"to_k",
+		"to_v",
+		"to_out.0"
+	]
+}
+```
+
+## 4. Run Training
+Run `python trainer.py` with the following parameter introduction:
+
+# Trainer Parameter Interpretation
+
+## 1. General Settings
+1. **`--num_nodes`**: This parameter specifies the number of nodes for the training process. It is an integer value, and the default is set to 1. In scenarios where distributed training across multiple nodes is applicable, this parameter determines how many nodes will be utilized. For example, if you have a cluster of machines and want to distribute the training load, you can increase this number. However, for most single-machine or basic training setups, the default value of 1 is sufficient.
+2. **`--shift`**: It is a floating-point parameter with a default value of 3.0. Although its specific function depends on the implementation details of the model, it is likely used for some internal calculations related to the model, such as adjusting certain biases or offsets in the neural network architecture during the training process.
+
+## 2. Training Hyperparameters
+1. **`--learning_rate`**: This is a crucial hyperparameter for the training process. It is a floating-point value with a default of 1e-4 (0.0001). The learning rate determines the step size at each iteration while updating the model's weights. A smaller learning rate will make the training process more stable but may require more training steps to converge. On the other hand, a larger learning rate can lead to faster convergence but might cause the model to overshoot the optimal solution and result in unstable training or even divergence.
+2. **`--num_workers`**: This parameter defines the number of worker processes that will be used for data loading. It is an integer with a default value of 8. Having multiple workers can significantly speed up the data loading process, especially when dealing with large datasets. However, it also consumes additional system resources, so you may need to adjust this value based on the available resources of your machine (e.g., CPU cores and memory).
+3. **`--epochs`**: It represents the number of times the entire training dataset will be passed through the model. It is an integer, and the default value is set to -1. When set to -1, the training will continue until another stopping condition (such as reaching the maximum number of steps) is met. If you set a positive integer value, the training will stop after that number of epochs.
+4. **`--max_steps`**: This parameter specifies the maximum number of training steps. It is an integer with a default value of 2000000. Once the model has completed this number of training steps, the training process will stop, regardless of whether the model has fully converged or not. This is useful for setting a limit on the training duration in terms of the number of steps.
+5. **`--every_n_train_steps`**: It is an integer parameter with a default of 2000. It determines how often certain operations (such as saving checkpoints, logging training progress, etc.) will be performed during the training. For example, with a value of 2000, these operations will occur every 2000 training steps.
+
+## 3. Dataset and Experiment Settings
+1. **`--dataset_path`**: This is a string parameter that indicates the path to the dataset in the Huggingface dataset format. The default value is "./zh_lora_dataset". You need to ensure that the dataset at this path is correctly formatted and contains the necessary data for training.
+2. **`--exp_name`**: It is a string parameter used to name the experiment. The default value is "chinese_rap_lora". This name can be used to distinguish different training experiments, and it is often used in logging and saving checkpoints to organize and identify the results of different runs.
+
+## 4. Training Precision and Gradient Settings
+1. **`--precision`**: This parameter specifies the precision of the training. It is a string with a default value of "32", which usually means 32-bit floating-point precision. Higher precision can lead to more accurate training but may also consume more memory and computational resources. You can adjust this value depending on your hardware capabilities and the requirements of your model.
+2. **`--accumulate_grad_batches`**: It is an integer parameter with a default value of 1. It determines how many batches of data will be used to accumulate gradients before performing an optimization step. For example, if you set it to 4, the gradients from 4 consecutive batches will be accumulated, and then the model's weights will be updated. This can be useful in scenarios where you want to simulate larger batch sizes when your available memory does not allow for actual large batch training.
+3. **`--gradient_clip_val`**: This is a floating-point parameter with a default value of 0.5. It is used to clip the gradients during the backpropagation process. Clipping the gradients helps prevent the issue of gradient explosion, where the gradients become extremely large and cause the model to become unstable. By setting a clip value, the gradients will be adjusted to be within a certain range.
+4. **`--gradient_clip_algorithm`**: It is a string parameter with a default value of "norm". This parameter specifies the algorithm used for gradient clipping. The "norm" algorithm is one common method, but there may be other algorithms available depending on the implementation of the training framework.
+
+## 5. Checkpoint and Logging Settings
+1. **`--devices`**: This is an integer parameter with a default value of 1. It specifies the number of devices (such as GPUs) that will be used for training. If you have multiple GPUs available and want to use them for parallel training, you can increase this number accordingly.
+2. **`--logger_dir`**: It is a string parameter with a default value of "./exps/logs/". This parameter indicates the directory where the training logs will be saved. The logs can be useful for monitoring the training progress, analyzing the performance of the model during training, and debugging any issues that may arise.
+3. **`--ckpt_path`**: It is a string parameter with a default value of None. If you want to resume training from a previously saved checkpoint, you can specify the path to the checkpoint file using this parameter. If set to None, the training will start from scratch.
+4. **`--checkpoint_dir`**: This is a string parameter with a default value of None. It specifies the directory where the checkpoints of the model will be saved during the training process. If set to None, checkpoints may not be saved or may be saved in a default location depending on the training framework.
+
+## 6. Validation and Reloading Settings
+1. **`--reload_dataloaders_every_n_epochs`**: It is an integer parameter with a default value of 1. It determines how often the data loaders will be reloaded during the training process. Reloading the data loaders can be useful when you want to ensure that the data is shuffled or processed differently for each epoch, especially when dealing with datasets that may change or have some specific requirements.
+2. **`--every_plot_step`**: It is an integer parameter with a default value of 2000. It specifies how often some visualizations or plots (such as loss curves, accuracy plots, etc.) will be generated during the training process. For example, with a value of 2000, the plots will be updated every 2000 training steps.
+3. **`--val_check_interval`**: This is an integer parameter with a default value of None. It determines how often the validation process will be performed during the training. If set to a positive integer, the model will be evaluated on the validation dataset every specified number of steps. If set to None, no regular validation checks will be performed.
+4. **`--lora_config_path`**: It is a string parameter with a default value of "config/zh_rap_lora_config.json". This parameter specifies the path to the configuration file for the Lora (Low-Rank Adaptation) module. The Lora configuration file contains settings related to the Lora module, such as the rank of the low-rank matrices, the learning rate for the Lora parameters, etc. 
@@ -0,0 +1,38 @@
+# 🎤 RapMachine Release
+
+We meticulously curated and trained this model on Chinese rap/hip-hop datasets, with rigorous data cleaning and recaptioning. The results include:  
+- **Improved pronunciation** for Chinese lyrics  
+- **Enhanced adherence** to hip-hop and electronic music styles  
+- **Greater diversity** in hip-hop vocal performances  
+
+### **How to Use**  
+1. Generate **higher-quality Chinese songs**  (⚠️ It's not just for Chinese songs. You can also use it in other ways. )
+
+2. Create **better hip-hop tracks**  
+3. Blend it with other genres to:  
+   - Produce music with **richer vocal details**  
+   - Experiment with **underground or street culture flavors**  
+4. Fine-tune outputs using the following dimensions:  
+
+    - **`vocal_timbre`**: Describes the inherent qualities of the voice.
+        - Examples: *Bright, dark, warm, cold, breathy, nasal, gritty, smooth, husky, metallic, whispery, resonant, airy, smoky, sultry, light, clear, high-pitched, raspy, powerful, ethereal, flute-like, hollow, velvety, shrill, hoarse, mellow, thin, thick, reedy, silvery, twangy.*  
+
+    - **`techniques`**:
+        - Examples:  
+            - **Rap styles**: `mumble rap`, `chopper rap`, `melodic rap`, `lyrical rap`, `trap flow`, `double-time rap`  
+            - **Vocal effects**: `auto-tune`, `reverb`, `delay`, `distortion`  
+            - **Delivery styles**: `whispered`, `shouted`, `spoken word`, `narration`, `singing`  
+            - **Other vocalizations**: `ad-libs`, `call-and-response`, `harmonized`  
+
+---  
+
+## Community Note
+We’ve **revamped and expanded** the LoRA training guide with finer details. This release is a **proof of concept**—showcasing the potential of **ACE-Step**.
+
+While a Chinese rap LoRA might seem niche for non-Chinese communities, we consistently demonstrate through such projects that ACE-step - as a music generation foundation model - holds boundless potential. It doesn't just improve pronunciation in one language, but spawns new styles. 
+
+The universal human appreciation of music is a precious asset. Like abstract LEGO blocks, these elements will eventually combine in more organic ways. May our open-source contributions propel the evolution of musical history forward.
+
+**Enjoy it, customize it, and create something entirely new.** 
+
+We can’t wait to hear what you’ll build!  
@@ -0,0 +1,43 @@
+import torch
+import functools
+from typing import Callable, TypeVar
+
+
+class CpuOffloader:
+    def __init__(self, model, device="cpu"):
+        self.model = model
+        self.original_device = device
+        self.original_dtype = model.dtype
+    
+    def __enter__(self):
+        if not hasattr(self.model,"torchao_quantized"):
+            self.model.to(self.original_device, dtype=self.original_dtype)
+        return self.model
+    
+    def __exit__(self, *args):
+        if not hasattr(self.model,"torchao_quantized"):
+            self.model.to("cpu")
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+
+
+T = TypeVar('T')
+
+def cpu_offload(model_attr: str):
+    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        @functools.wraps(func)
+        def wrapper(self, *args, **kwargs):
+            if not self.cpu_offload:
+                return func(self, *args, **kwargs)
+
+            # Get the device from the class
+            device = self.device
+            # Get the model from the class attribute
+            model = getattr(self, model_attr)
+            
+            with CpuOffloader(model, device):
+                return func(self, *args, **kwargs)
+                        
+        return wrapper
+    return decorator
@@ -3,19 +3,28 @@ from pathlib import Path
 import random


-DEFAULT_ROOT_DIR = "examples/input_params"
-
+DEFAULT_ROOT_DIR = "examples/default/input_params"
+ZH_RAP_LORA_ROOT_DIR = "examples/zh_rap_lora/input_params"

 class DataSampler:
    def __init__(self, root_dir=DEFAULT_ROOT_DIR):
        self.root_dir = root_dir
        self.input_params_files = list(Path(self.root_dir).glob("*.json"))
+        self.zh_rap_lora_input_params_files = list(Path(ZH_RAP_LORA_ROOT_DIR).glob("*.json"))
+        self.zh_rap_lora_input_params_files += list(Path(ZH_RAP_LORA_ROOT_DIR).glob("*.json"))

    def load_json(self, file_path):
        with open(file_path, "r", encoding="utf-8") as f:
            return json.load(f)

-    def sample(self):
-        json_path = random.choice(self.input_params_files)
-        json_data = self.load_json(json_path)
+    def sample(self, lora_name_or_path=None):
+        if lora_name_or_path is None or lora_name_or_path == "none":
+            json_path = random.choice(self.input_params_files)
+            json_data = self.load_json(json_path)
+        else:
+            json_path = random.choice(self.zh_rap_lora_input_params_files)
+            json_data = self.load_json(json_path)
+            # Update the lora_name in the json_data
+            json_data["lora_name_or_path"] = lora_name_or_path
+
        return json_data
@@ -9,11 +9,6 @@ Apache 2.0 License
 import os
 import click

-from acestep.ui.components import create_main_demo_ui
-from acestep.pipeline_ace_step import ACEStepPipeline
-from acestep.data_sampler import DataSampler
-
-
@click.command()
@click.option(
    "--checkpoint_path",
@@ -46,23 +41,36 @@ from acestep.data_sampler import DataSampler
@click.option(
    "--torch_compile", type=click.BOOL, default=False, help="Whether to use torch.compile."
 )
-def main(checkpoint_path, server_name, port, device_id, share, bf16, torch_compile):
+@click.option(
+    "--cpu_offload", type=bool, default=False, help="Whether to use CPU offloading (only load current stage's model to GPU)"
+)
+@click.option(
+    "--overlapped_decode", type=bool, default=False, help="Whether to use overlapped decoding (run dcae and vocoder using sliding windows)"
+)
+def main(checkpoint_path, server_name, port, device_id, share, bf16, torch_compile, cpu_offload, overlapped_decode):
    """
    Main function to launch the ACE Step pipeline demo.
    """

    os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)

+    from acestep.ui.components import create_main_demo_ui
+    from acestep.pipeline_ace_step import ACEStepPipeline
+    from acestep.data_sampler import DataSampler
+
    model_demo = ACEStepPipeline(
        checkpoint_dir=checkpoint_path,
        dtype="bfloat16" if bf16 else "float32",
        torch_compile=torch_compile,
+        cpu_offload=cpu_offload,
+        overlapped_decode=overlapped_decode
    )
    data_sampler = DataSampler()

    demo = create_main_demo_ui(
        text2music_process_func=model_demo.__call__,
        sample_data_func=data_sampler.sample,
+        load_data_func=data_sampler.load_json,
    )
    demo.launch(server_name=server_name, server_port=port, share=share)

@@ -0,0 +1,99 @@
+default = [
+    "af",
+    "am",
+    "an",
+    "ar",
+    "as",
+    "az",
+    "be",
+    "bg",
+    "bn",
+    "br",
+    "bs",
+    "ca",
+    "cs",
+    "cy",
+    "da",
+    "de",
+    "dz",
+    "el",
+    "en",
+    "eo",
+    "es",
+    "et",
+    "eu",
+    "fa",
+    "fi",
+    "fo",
+    "fr",
+    "ga",
+    "gl",
+    "gu",
+    "he",
+    "hi",
+    "hr",
+    "ht",
+    "hu",
+    "hy",
+    "id",
+    "is",
+    "it",
+    "ja",
+    "jv",
+    "ka",
+    "kk",
+    "km",
+    "kn",
+    "ko",
+    "ku",
+    "ky",
+    "la",
+    "lb",
+    "lo",
+    "lt",
+    "lv",
+    "mg",
+    "mk",
+    "ml",
+    "mn",
+    "mr",
+    "ms",
+    "mt",
+    "nb",
+    "ne",
+    "nl",
+    "nn",
+    "no",
+    "oc",
+    "or",
+    "pa",
+    "pl",
+    "ps",
+    "pt",
+    "qu",
+    "ro",
+    "ru",
+    "rw",
+    "se",
+    "si",
+    "sk",
+    "sl",
+    "sq",
+    "sr",
+    "sv",
+    "sw",
+    "ta",
+    "te",
+    "th",
+    "tl",
+    "tr",
+    "ug",
+    "uk",
+    "ur",
+    "vi",
+    "vo",
+    "wa",
+    "xh",
+    "zh",
+    "zu",
+]
@@ -360,10 +360,6 @@ class ACEStepTransformer2DModel(
        for module in self.children():
            fn_recursive_feed_forward(module, chunk_size, dim)

-    def _set_gradient_checkpointing(self, module, value=False):
-        if hasattr(module, "gradient_checkpointing"):
-            module.gradient_checkpointing = value
-
    def forward_lyric_encoder(
        self,
        lyric_token_idx: Optional[torch.LongTensor] = None,
@@ -456,20 +452,8 @@ class ACEStepTransformer2DModel(

            if self.training and self.gradient_checkpointing:

-                def create_custom_forward(module, return_dict=None):
-                    def custom_forward(*inputs):
-                        if return_dict is not None:
-                            return module(*inputs, return_dict=return_dict)
-                        else:
-                            return module(*inputs)
-
-                    return custom_forward
-
-                ckpt_kwargs: Dict[str, Any] = (
-                    {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
-                )
                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(block),
+                    block,
                    hidden_states=hidden_states,
                    attention_mask=attention_mask,
                    encoder_hidden_states=encoder_hidden_states,
@@ -477,7 +461,7 @@ class ACEStepTransformer2DModel(
                    rotary_freqs_cis=rotary_freqs_cis,
                    rotary_freqs_cis_cross=encoder_rotary_freqs_cis,
                    temb=temb,
-                    **ckpt_kwargs,
+                    use_reentrant=False,
                )

            else:
@@ -1030,8 +1030,8 @@ class ConformerEncoder(torch.nn.Module):
        mask_pad: torch.Tensor,
    ) -> torch.Tensor:
        for layer in self.encoders:
-            xs, chunk_masks, _, _ = ckpt.checkpoint(
-                layer.__call__, xs, chunk_masks, pos_emb, mask_pad
+            xs, chunk_masks, _, _ = torch.utils.checkpoint.checkpoint(
+                layer.__call__, xs, chunk_masks, pos_emb, mask_pad, use_reentrant=False
            )
        return xs

@@ -10,11 +10,12 @@ import os
 import torch
 from diffusers import AutoencoderDC
 import torchaudio
+import soundfile as sf
 import torchvision.transforms as transforms
 from diffusers.models.modeling_utils import ModelMixin
 from diffusers.loaders import FromOriginalModelMixin
 from diffusers.configuration_utils import ConfigMixin, register_to_config
-
+from tqdm import tqdm

 try:
    from .music_vocoder import ADaMoSHiFiGANV1
@@ -60,7 +61,13 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
        self.shift_factor = -1.9091

    def load_audio(self, audio_path):
-        audio, sr = torchaudio.load(audio_path)
+        # Read with soundfile rather than torchaudio.load(): since torchaudio
+        # 2.11 the latter routes I/O through TorchCodec, an extra native
+        # dependency we do not require.
+        data, sr = sf.read(audio_path, dtype="float32", always_2d=True)
+        audio = torch.from_numpy(data.T)
+        if audio.shape[0] == 1:
+            audio = audio.repeat(2, 1)
        return audio, sr

    def forward_mel(self, audios):
@@ -119,15 +126,18 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
            mels = self.dcae.decoder(latent.unsqueeze(0))
            mels = mels * 0.5 + 0.5
            mels = mels * (self.max_mel_value - self.min_mel_value) + self.min_mel_value
-            wav = self.vocoder.decode(mels[0]).squeeze(1)
+
+            # wav = self.vocoder.decode(mels[0]).squeeze(1)
+            # decode waveform for each channels to reduce vram footprint
+            wav_ch1 = self.vocoder.decode(mels[:,0,:,:]).squeeze(1).cpu()
+            wav_ch2 = self.vocoder.decode(mels[:,1,:,:]).squeeze(1).cpu()
+            wav = torch.cat([wav_ch1, wav_ch2],dim=0)

            if sr is not None:
                resampler = (
                    torchaudio.transforms.Resample(44100, sr)
-                    .to(latents.device)
-                    .to(latents.dtype)
                )
-                wav = resampler(wav)
+                wav = resampler(wav.cpu().float())
            else:
                sr = 44100
            pred_wavs.append(wav)
@@ -138,6 +148,215 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
            ]
        return sr, pred_wavs

+    @torch.no_grad()
+    def decode_overlap(self, latents, audio_lengths=None, sr=None):
+        """
+        Decodes latents into waveforms using an overlapped DCAE and Vocoder.
+        """
+        print("Using Overlapped DCAE and Vocoder")
+
+        MODEL_INTERNAL_SR = 44100
+        DCAE_LATENT_TO_MEL_STRIDE = 8
+        VOCODER_AUDIO_SAMPLES_PER_MEL_FRAME = 512
+
+        pred_wavs = []
+        final_output_sr = sr if sr is not None else MODEL_INTERNAL_SR
+
+        # --- DCAE Parameters ---
+        # dcae_win_len_latent: Window length in the latent domain for DCAE processing
+        dcae_win_len_latent = 512 
+        # dcae_mel_win_len: Expected mel window length from DCAE decoder output (latent_win * stride)
+        dcae_mel_win_len = dcae_win_len_latent * 8
+        # dcae_anchor_offset: Offset from anchor point to actual start of latent window slice
+        dcae_anchor_offset = dcae_win_len_latent // 4
+        # dcae_anchor_hop: Hop size for anchor points in latent domain
+        dcae_anchor_hop = dcae_win_len_latent // 2
+        # dcae_mel_overlap_len: Overlap length in the mel domain to be trimmed/blended
+        dcae_mel_overlap_len = dcae_mel_win_len // 4
+
+        # --- Vocoder Parameters ---
+        # vocoder_win_len_audio: Audio samples per vocoder processing window
+        vocoder_win_len_audio = 512 * 512 # Example: 262144 samples
+        # vocoder_overlap_len_audio: Audio samples for overlap between vocoder windows
+        vocoder_overlap_len_audio = 1024  
+        # vocoder_hop_len_audio: Hop size in audio samples for vocoder processing
+        vocoder_hop_len_audio = vocoder_win_len_audio - 2 * vocoder_overlap_len_audio
+        # vocoder_input_mel_frames_per_block: Number of mel frames fed to vocoder in one go
+        vocoder_input_mel_frames_per_block = vocoder_win_len_audio // VOCODER_AUDIO_SAMPLES_PER_MEL_FRAME
+        
+        crossfade_len_audio = 128 # Audio samples for crossfading vocoder outputs
+        cf_win_tail = torch.linspace(1, 0, crossfade_len_audio, device=self.device).unsqueeze(0).unsqueeze(0)
+        cf_win_head = torch.linspace(0, 1, crossfade_len_audio, device=self.device).unsqueeze(0).unsqueeze(0)
+
+        for latent_idx, latent_item in enumerate(latents):
+            latent_item = latent_item.to(self.device)
+            current_latent = (latent_item / self.scale_factor + self.shift_factor).unsqueeze(0) # (1, C, H, W_latent)
+            latent_len = current_latent.shape[3]
+
+            # 1. DCAE: Latent to Mel Spectrogram (Overlapped)
+            mels_segments = []
+            if latent_len == 0:
+                pass # No mel segments to generate
+            else:
+                # Determine anchor points for DCAE windows
+                # An anchor marks a reference point for a window slice.
+                # Window slice: current_latent[..., anchor - offset : anchor - offset + win_len]
+                # First anchor ensures window starts at 0. Last anchor ensures tail is covered.
+                dcae_anchors = list(range(dcae_anchor_offset, latent_len - dcae_anchor_offset, dcae_anchor_hop))
+                if not dcae_anchors: # If latent is too short for the range, use one anchor
+                    dcae_anchors = [dcae_anchor_offset]
+                
+                for i, anchor in enumerate(dcae_anchors):
+                    win_start_idx = max(0, anchor - dcae_anchor_offset)
+                    win_end_idx = min(latent_len, win_start_idx + dcae_win_len_latent)
+                    
+                    dcae_input_segment = current_latent[:, :, :, win_start_idx:win_end_idx]
+                    if dcae_input_segment.shape[3] == 0: continue
+
+                    mel_output_full = self.dcae.decoder(dcae_input_segment) # (1, C, H_mel, W_mel_fixed_from_dcae)
+
+                    is_first = (i == 0)
+                    is_last = (i == len(dcae_anchors) - 1)
+
+                    if is_first and is_last: # Only one segment
+                        # Use mel corresponding to actual input latent length
+                        true_mel_content_len = dcae_input_segment.shape[3] * DCAE_LATENT_TO_MEL_STRIDE
+                        mel_to_keep = mel_output_full[:, :, :, :min(true_mel_content_len, mel_output_full.shape[3])]
+                    elif is_first: # First segment, trim end overlap
+                        mel_to_keep = mel_output_full[:, :, :, :-dcae_mel_overlap_len]
+                    elif is_last: # Last segment, trim start overlap
+                        # And ensure we only take content relevant to the (potentially partial) last latent window
+                        # The mel_output_full is fixed length. The useful part starts after overlap.
+                        # The length of the useful part depends on how much of dcae_input_segment was actual content.
+                        # For simplicity in overlap-add, typically trim fixed overlap.
+                        # If dcae_input_segment was shorter than dcae_win_len_latent, mel_output_full might contain padding effects.
+                        # Standard OLA keeps the corresponding tail.
+                        mel_to_keep = mel_output_full[:, :, :, dcae_mel_overlap_len:]
+                    else: # Middle segment, trim both overlaps
+                        mel_to_keep = mel_output_full[:, :, :, dcae_mel_overlap_len:-dcae_mel_overlap_len]
+                    
+                    if mel_to_keep.shape[3] > 0:
+                        mels_segments.append(mel_to_keep)
+            
+            if not mels_segments:
+                num_mel_channels = current_latent.shape[1]
+                mel_height = self.dcae.decoder_output_mel_height
+                concatenated_mels = torch.empty(
+                    (1, num_mel_channels, mel_height, 0),
+                    device=current_latent.device, dtype=current_latent.dtype
+                )
+            else:
+                concatenated_mels = torch.cat(mels_segments, dim=3)
+
+            # Denormalize mels
+            concatenated_mels = concatenated_mels * 0.5 + 0.5
+            concatenated_mels = concatenated_mels * (self.max_mel_value - self.min_mel_value) + self.min_mel_value
+            
+            mel_total_frames = concatenated_mels.shape[3]
+
+            # 2. Vocoder: Mel Spectrogram to Waveform (Overlapped)
+            if mel_total_frames == 0:
+                # Assuming mono or stereo output based on mel channels (typically mono for vocoder from single mel)
+                num_audio_channels = 1 # Or determine from vocoder capabilities / mel channels
+                final_wav = torch.zeros((num_audio_channels, 0), device=self.device, dtype=torch.float32)
+            else:
+                # Initial vocoder window
+                # Vocoder expects (C_mel, H_mel, W_mel_block)
+                mel_block = concatenated_mels[0, :, :, :vocoder_input_mel_frames_per_block].to(self.device)
+                
+                # Pad mel_block if it's shorter than vocoder_input_mel_frames_per_block (e.g. very short audio)
+                if 0 < mel_block.shape[2] < vocoder_input_mel_frames_per_block:
+                    pad_len = vocoder_input_mel_frames_per_block - mel_block.shape[2]
+                    mel_block = torch.nn.functional.pad(mel_block, (0, pad_len), mode='constant', value=0) # Pad last dim
+                
+                current_audio_output = self.vocoder.decode(mel_block) # (C_audio, 1, Samples)
+                current_audio_output = current_audio_output[:, :, :-vocoder_overlap_len_audio] # Remove end overlap
+
+                # p_audio_samples tracks the start of the *next* audio segment to generate (in conceptual total audio samples)
+                p_audio_samples = vocoder_hop_len_audio 
+                conceptual_total_audio_len_native_sr = mel_total_frames * VOCODER_AUDIO_SAMPLES_PER_MEL_FRAME
+                
+                pbar_total = 1 + max(0, (conceptual_total_audio_len_native_sr - (vocoder_win_len_audio - vocoder_overlap_len_audio))) // vocoder_hop_len_audio
+                
+                # Use tqdm if you want a progress bar for the vocoder part
+                # with tqdm(total=pbar_total, desc=f"Vocoder {latent_idx+1}/{len(latents)}", leave=False) as pbar:
+                # pbar.update(1) # For initial window
+                # The loop for subsequent windows
+                while p_audio_samples < conceptual_total_audio_len_native_sr:
+                    mel_frame_start = p_audio_samples // VOCODER_AUDIO_SAMPLES_PER_MEL_FRAME
+                    mel_frame_end = mel_frame_start + vocoder_input_mel_frames_per_block
+                    
+                    if mel_frame_start >= mel_total_frames: break # No more mel frames
+
+                    mel_block = concatenated_mels[0, :, :, mel_frame_start:min(mel_frame_end, mel_total_frames)].to(self.device)
+                    
+                    if mel_block.shape[2] == 0: break # Should not happen if mel_frame_start is valid
+
+                    # Pad if current mel_block is too short (end of sequence)
+                    if mel_block.shape[2] < vocoder_input_mel_frames_per_block:
+                        pad_len = vocoder_input_mel_frames_per_block - mel_block.shape[2]
+                        mel_block = torch.nn.functional.pad(mel_block, (0, pad_len), mode='constant', value=0)
+
+                    new_audio_win = self.vocoder.decode(mel_block) # (C_audio, 1, Samples)
+
+                    # Crossfade
+                    # Determine actual crossfade length based on available audio
+                    actual_cf_len = min(crossfade_len_audio, current_audio_output.shape[2], new_audio_win.shape[2] - (vocoder_overlap_len_audio - crossfade_len_audio))
+                    if actual_cf_len > 0: # Ensure valid slice lengths for crossfade
+                        tail_part = current_audio_output[:, :, -actual_cf_len:]
+                        head_part = new_audio_win[:, :, vocoder_overlap_len_audio - actual_cf_len : vocoder_overlap_len_audio]
+                        
+                        crossfaded_segment = tail_part * cf_win_tail[:,:,:actual_cf_len] + \
+                                             head_part * cf_win_head[:,:,:actual_cf_len]
+                        
+                        current_audio_output = torch.cat([current_audio_output[:, :, :-actual_cf_len], crossfaded_segment], dim=2)
+
+                    # Append non-overlapping part of new_audio_win
+                    is_final_append = (p_audio_samples + vocoder_hop_len_audio >= conceptual_total_audio_len_native_sr)
+                    if is_final_append:
+                        segment_to_append = new_audio_win[:, :, vocoder_overlap_len_audio:]
+                    else:
+                        segment_to_append = new_audio_win[:, :, vocoder_overlap_len_audio:-vocoder_overlap_len_audio]
+                    
+                    current_audio_output = torch.cat([current_audio_output, segment_to_append], dim=2)
+                    
+                    p_audio_samples += vocoder_hop_len_audio
+                    # pbar.update(1) # if using tqdm
+                
+                final_wav = current_audio_output.squeeze(1) # (C_audio, Samples)
+
+            # 3. Resampling (if necessary)
+            if final_output_sr != MODEL_INTERNAL_SR and final_wav.numel() > 0:
+                # Resample expects CPU tensor if using torchaudio.transforms on older versions or for some backends
+                resampler = torchaudio.transforms.Resample(
+                    MODEL_INTERNAL_SR, final_output_sr, dtype=final_wav.dtype
+                )
+                final_wav = resampler(final_wav.cpu()).to(self.device) # Move back to device if needed later
+            
+            pred_wavs.append(final_wav)
+
+        # 4. Final Truncation
+        processed_pred_wavs = []
+        for i, wav in enumerate(pred_wavs):
+            # Calculate expected length based on original latent, at the FINAL output sample rate
+            _num_latent_frames = latents[i].shape[-1] # Use original latent item for shape
+            _num_mel_frames = _num_latent_frames * DCAE_LATENT_TO_MEL_STRIDE
+            _conceptual_native_audio_len = _num_mel_frames * VOCODER_AUDIO_SAMPLES_PER_MEL_FRAME
+            max_possible_len = int(_conceptual_native_audio_len * final_output_sr / MODEL_INTERNAL_SR)
+
+            current_wav_len = wav.shape[1]
+            
+            if audio_lengths is not None:
+                # User-provided length is the primary target, capped by actual and max possible
+                target_len = min(audio_lengths[i], current_wav_len, max_possible_len)
+            else:
+                # No user length, use max possible capped by actual
+                target_len = min(max_possible_len, current_wav_len)
+            
+            processed_pred_wavs.append(wav[:, :max(0, target_len)].cpu()) # Ensure length is non-negative
+
+        return final_output_sr, processed_pred_wavs
+
    def forward(self, audios, audio_lengths=None, sr=None):
        latents, latent_lengths = self.encode(
            audios=audios, audio_lengths=audio_lengths, sr=sr
@@ -148,7 +367,8 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):

 if __name__ == "__main__":

-    audio, sr = torchaudio.load("test.wav")
+    _data, sr = sf.read("test.wav", dtype="float32", always_2d=True)
+    audio = torch.from_numpy(_data.T)
    audio_lengths = torch.tensor([audio.shape[1]])
    audios = audio.unsqueeze(0)

@@ -164,5 +384,5 @@ if __name__ == "__main__":
    print("latents shape: ", latents.shape)
    print("latent_lengths: ", latent_lengths)
    print("sr: ", sr)
-    torchaudio.save("test_reconstructed.wav", pred_wavs[0], sr)
+    sf.write("test_reconstructed.wav", pred_wavs[0].float().cpu().transpose(0, 1).numpy(), sr)
    print("test_reconstructed.wav")
@@ -71,9 +71,10 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
        max_shift: Optional[float] = 1.15,
        base_image_seq_len: Optional[int] = 256,
        max_image_seq_len: Optional[int] = 4096,
+        sigma_max: Optional[float] = 1.0,
    ):
        timesteps = np.linspace(
-            1, num_train_timesteps, num_train_timesteps, dtype=np.float32
+            1.0, sigma_max*num_train_timesteps, num_train_timesteps, dtype=np.float32
        )[::-1].copy()
        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)

@@ -66,9 +66,10 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
        self,
        num_train_timesteps: int = 1000,
        shift: float = 1.0,
+        sigma_max: Optional[float] = 1.0,
    ):
        timesteps = np.linspace(
-            1, num_train_timesteps, num_train_timesteps, dtype=np.float32
+            1.0, sigma_max*num_train_timesteps, num_train_timesteps, dtype=np.float32
        )[::-1].copy()
        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)

@@ -0,0 +1,343 @@
+# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput, logging
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class FlowMatchPingPongSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+    """
+
+    prev_sample: torch.FloatTensor
+
+
+class FlowMatchPingPongScheduler(SchedulerMixin, ConfigMixin):
+    """
+    PingPong scheduler.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        timestep_spacing (`str`, defaults to `"linspace"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        shift (`float`, defaults to 1.0):
+            The shift value for the timestep schedule.
+    """
+
+    _compatibles = []
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        use_dynamic_shifting=False,
+        base_shift: Optional[float] = 0.5,
+        max_shift: Optional[float] = 1.15,
+        base_image_seq_len: Optional[int] = 256,
+        max_image_seq_len: Optional[int] = 4096,
+        sigma_max: Optional[float] = 1.0,
+    ):
+        timesteps = np.linspace(
+            1, sigma_max*num_train_timesteps, num_train_timesteps, dtype=np.float32
+        )[::-1].copy()
+        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
+
+        sigmas = timesteps / num_train_timesteps
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+
+        self.timesteps = sigmas * num_train_timesteps
+
+        self._step_index = None
+        self._begin_index = None
+
+        self.sigmas = sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    def scale_noise(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        noise: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        """
+        Forward process in flow-matching
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The input sample.
+            timestep (`int`, *optional*):
+                The current timestep in the diffusion chain.
+
+        Returns:
+            `torch.FloatTensor`:
+                A scaled input sample.
+        """
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
+
+        if sample.device.type == "mps" and torch.is_floating_point(timestep):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
+            timestep = timestep.to(sample.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(sample.device)
+            timestep = timestep.to(sample.device)
+
+        # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps) for t in timestep
+            ]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timestep.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timestep.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(sample.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        sample = sigma * noise + (1.0 - sigma) * sample
+
+        return sample
+
+    def _sigma_to_t(self, sigma):
+        return sigma * self.config.num_train_timesteps
+
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int = None,
+        device: Union[str, torch.device] = None,
+        sigmas: Optional[List[float]] = None,
+        mu: Optional[float] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        """
+
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError(
+                " you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
+            )
+
+        if sigmas is None:
+            self.num_inference_steps = num_inference_steps
+            timesteps = np.linspace(
+                self._sigma_to_t(self.sigma_max),
+                self._sigma_to_t(self.sigma_min),
+                num_inference_steps,
+            )
+
+            sigmas = timesteps / self.config.num_train_timesteps
+
+        if self.config.use_dynamic_shifting:
+            sigmas = self.time_shift(mu, 1.0, sigmas)
+        else:
+            sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
+
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
+        timesteps = sigmas * self.config.num_train_timesteps
+
+        self.timesteps = timesteps.to(device=device)
+        self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
+
+        self._step_index = None
+        self._begin_index = None
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep):
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        sample: torch.FloatTensor,
+        s_churn: float = 0.0,
+        s_tmin: float = 0.0,
+        s_tmax: float = float("inf"),
+        s_noise: float = 1.0,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+        omega: Union[float, np.array] = 0.0,
+    ) -> Union[FlowMatchPingPongSchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            s_churn (`float`):
+            s_tmin  (`float`):
+            s_tmax  (`float`):
+            s_noise (`float`, defaults to 1.0):
+                Scaling factor for noise added to the sample.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
+                tuple.
+
+        Returns:
+            [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
+                returned, otherwise a tuple is returned where the first element is the sample tensor.
+        """
+
+        def logistic_function(x, L=0.9, U=1.1, x_0=0.0, k=1):
+            # L = Lower bound
+            # U = Upper bound
+            # x_0 = Midpoint (x corresponding to y = 1.0)
+            # k = Steepness, can adjust based on preference
+
+            if isinstance(x, torch.Tensor):
+                device_ = x.device
+                x = x.to(torch.float).cpu().numpy()
+
+            new_x = L + (U - L) / (1 + np.exp(-k * (x - x_0)))
+
+            if isinstance(new_x, np.ndarray):
+                new_x = torch.from_numpy(new_x).to(device_)
+            return new_x
+
+        self.omega_bef_rescale = omega
+        omega = logistic_function(omega, k=0.1)
+        self.omega_aft_rescale = omega
+
+        if (
+            isinstance(timestep, int)
+            or isinstance(timestep, torch.IntTensor)
+            or isinstance(timestep, torch.LongTensor)
+        ):
+            raise ValueError(
+                (
+                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                    " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                    " one of the `scheduler.timesteps` as a timestep."
+                ),
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+
+        sigma = self.sigmas[self.step_index]
+        sigma_next = self.sigmas[self.step_index + 1]
+
+        denoised = sample - sigma * model_output
+        noise = torch.empty_like(sample).normal_(generator=generator)
+        prev_sample = (1 - sigma_next) * denoised + sigma_next * noise
+
+        # Cast sample back to model compatible dtype
+        prev_sample = prev_sample.to(model_output.dtype)
+
+        # upon completion increase step index by one
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return FlowMatchPingPongSchedulerOutput(prev_sample=prev_sample)
+
+    def __len__(self):
+        return self.config.num_train_timesteps
@@ -7,10 +7,11 @@ from loguru import logger
 import time
 import traceback
 import torchaudio
+import soundfile as sf
 from pathlib import Path
 import re
-from language_segmentation import LangSegment
-from models.lyrics_utils.lyric_tokenizer import VoiceBpeTokenizer
+from acestep.language_segmentation import LangSegment
+from acestep.models.lyrics_utils.lyric_tokenizer import VoiceBpeTokenizer
 import warnings

 warnings.simplefilter("ignore", category=FutureWarning)
@@ -398,7 +399,10 @@ class Text2MusicDataset(Dataset):
        filename = item["filename"]
        sr = 48000
        try:
-            audio, sr = torchaudio.load(filename)
+            # soundfile instead of torchaudio.load(): torchaudio 2.11 routes
+            # I/O through TorchCodec, an extra native dependency.
+            _data, sr = sf.read(filename, dtype="float32", always_2d=True)
+            audio = torch.from_numpy(_data.T)
        except Exception as e:
            logger.error(f"Failed to load audio {item}: {e}")
            return None
@@ -456,7 +460,8 @@ class Text2MusicDataset(Dataset):
        speaker_emb_path = item.get("speaker_emb_path")
        if not speaker_emb_path:
            speaker_emb = self.get_speaker_emb_file(speaker_emb_path)
-        else:
+
+        if speaker_emb is None:
            speaker_emb = torch.zeros(512)

        # Process prompt/tags
@@ -8,6 +8,7 @@ Apache 2.0 License

 import gradio as gr
 import librosa
+import os


 TAG_DEFAULT = "funk, pop, soul, rock, melodic, guitar, drums, bass, keyboard, percussion, 105 BPM, energetic, upbeat, groovy, vibrant, dynamic"
@@ -48,6 +49,26 @@ Catch the tune and hold it tight
 In this moment we take flight
 """

+# First, let's define the presets at the top of the file, after the imports
+GENRE_PRESETS = {
+    "Modern Pop": "pop, synth, drums, guitar, 120 bpm, upbeat, catchy, vibrant, female vocals, polished vocals",
+    "Rock": "rock, electric guitar, drums, bass, 130 bpm, energetic, rebellious, gritty, male vocals, raw vocals",
+    "Hip Hop": "hip hop, 808 bass, hi-hats, synth, 90 bpm, bold, urban, intense, male vocals, rhythmic vocals",
+    "Country": "country, acoustic guitar, steel guitar, fiddle, 100 bpm, heartfelt, rustic, warm, male vocals, twangy vocals",
+    "EDM": "edm, synth, bass, kick drum, 128 bpm, euphoric, pulsating, energetic, instrumental",
+    "Reggae": "reggae, guitar, bass, drums, 80 bpm, chill, soulful, positive, male vocals, smooth vocals",
+    "Classical": "classical, orchestral, strings, piano, 60 bpm, elegant, emotive, timeless, instrumental",
+    "Jazz": "jazz, saxophone, piano, double bass, 110 bpm, smooth, improvisational, soulful, male vocals, crooning vocals",
+    "Metal": "metal, electric guitar, double kick drum, bass, 160 bpm, aggressive, intense, heavy, male vocals, screamed vocals",
+    "R&B": "r&b, synth, bass, drums, 85 bpm, sultry, groovy, romantic, female vocals, silky vocals"
+}
+
+# Add this function to handle preset selection
+def update_tags_from_preset(preset_name):
+    if preset_name == "Custom":
+        return ""
+    return GENRE_PRESETS.get(preset_name, "")
+

 def create_output_ui(task_name="Text2Music"):
    # For many consumer-grade GPU devices, only one batch can be run
@@ -69,7 +90,31 @@ def create_text2music_ui(
    gr,
    text2music_process_func,
    sample_data_func=None,
+    load_data_func=None,
 ):
+
+    with gr.Row(equal_height=True):
+        # Get base output directory from environment variable, defaulting to CWD-relative 'outputs'.
+        # This default (./outputs) is suitable for non-Docker local development.
+        # For Docker, the ACE_OUTPUT_DIR environment variable should be set (e.g., to /app/outputs).
+        output_file_dir = os.environ.get("ACE_OUTPUT_DIR", "./outputs")
+        if not os.path.isdir(output_file_dir):
+            os.makedirs(output_file_dir, exist_ok=True)
+        json_files = [f for f in os.listdir(output_file_dir) if f.endswith('.json')]
+
+        def _mtime(name):
+            # Output filenames are user-controlled (infer.py --output_path, or the
+            # save_path passed from the UI), so a timestamp cannot be parsed out of
+            # them: doing so used to raise ValueError and take the whole UI down.
+            try:
+                return os.path.getmtime(os.path.join(output_file_dir, name))
+            except OSError:
+                return 0.0
+
+        json_files.sort(key=_mtime, reverse=True)
+        output_files = gr.Dropdown(choices=json_files, label="Select previous generated input params", scale=9, interactive=True)
+        load_bnt = gr.Button("Load", variant="primary", scale=1)
+
    with gr.Row():
        with gr.Column():
            with gr.Row(equal_height=True):
@@ -84,35 +129,90 @@ def create_text2music_ui(
                    info="-1 means random duration (30 ~ 240).",
                    scale=9,
                )
-                sample_bnt = gr.Button("Sample", variant="primary", scale=1)
+                format = gr.Dropdown(choices=["mp3", "ogg", "flac", "wav"], value="wav", label="Format")
+                sample_bnt = gr.Button("Sample", variant="secondary", scale=1)

-            prompt = gr.Textbox(
-                lines=2,
-                label="Tags",
-                max_lines=4,
-                value=TAG_DEFAULT,
-                info="Support tags, descriptions, and scene. Use commas to separate different tags.\ntags and lyrics examples are from ai music generation community",
+            # audio2audio
+            with gr.Row(equal_height=True):
+                audio2audio_enable = gr.Checkbox(label="Enable Audio2Audio", value=False, info="Check to enable Audio-to-Audio generation using a reference audio.", elem_id="audio2audio_checkbox")
+                lora_name_or_path = gr.Dropdown(
+                    label="Lora Name or Path",
+                    choices=["ACE-Step/ACE-Step-v1-chinese-rap-LoRA", "none"],
+                    value="none",
+                    allow_custom_value=True,
+                    min_width=300
+                )
+                lora_weight = gr.Number(value=1.0, label="Lora weight", step=0.1, maximum=3, minimum=-3)
+
+            ref_audio_input = gr.Audio(type="filepath", label="Reference Audio (for Audio2Audio)", visible=False, elem_id="ref_audio_input")
+            ref_audio_strength = gr.Slider(
+                label="Refer audio strength",
+                minimum=0.0,
+                maximum=1.0,
+                step=0.01,
+                value=0.5,
+                elem_id="ref_audio_strength",
+                visible=False,
+                interactive=True,
            )
-            lyrics = gr.Textbox(
-                lines=9,
-                label="Lyrics",
-                max_lines=13,
-                value=LYRIC_DEFAULT,
-                info="Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics.\nUse [instrumental] or [inst] to generate instrumental music. Not support genre structure tag in lyrics",
+
+            def toggle_ref_audio_visibility(is_checked):
+                return (
+                    gr.update(visible=is_checked, elem_id="ref_audio_input"),
+                    gr.update(visible=is_checked, elem_id="ref_audio_strength"),
+                )
+
+            audio2audio_enable.change(
+                fn=toggle_ref_audio_visibility,
+                inputs=[audio2audio_enable],
+                outputs=[ref_audio_input, ref_audio_strength],
            )

+            with gr.Column(scale=2):
+                with gr.Group():
+                    gr.Markdown("""<center>Support tags, descriptions, and scene. Use commas to separate different tags.<br>Tags and lyrics examples are from AI music generation community.</center>""")
+                    with gr.Row():
+                        genre_preset = gr.Dropdown(
+                            choices=["Custom"] + list(GENRE_PRESETS.keys()),
+                            value="Custom",
+                            label="Preset",
+                            scale=1,
+                        )
+                        prompt = gr.Textbox(
+                            lines=1,
+                            label="Tags",
+                            max_lines=4,
+                            value=TAG_DEFAULT,
+                            scale=9,
+                        )
+
+            # Add the change event for the preset dropdown
+            genre_preset.change(
+                fn=update_tags_from_preset,
+                inputs=[genre_preset],
+                outputs=[prompt]
+            )
+            with gr.Group():
+                gr.Markdown("""<center>Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics.<br>Use [instrumental] or [inst] to generate instrumental music. Not support genre structure tag in lyrics</center>""")
+                lyrics = gr.Textbox(
+                    lines=9,
+                    label="Lyrics",
+                    max_lines=13,
+                    value=LYRIC_DEFAULT,
+                )
+
            with gr.Accordion("Basic Settings", open=False):
                infer_step = gr.Slider(
                    minimum=1,
-                    maximum=1000,
+                    maximum=200,
                    step=1,
-                    value=27,
+                    value=60,
                    label="Infer Steps",
                    interactive=True,
                )
                guidance_scale = gr.Slider(
                    minimum=0.0,
-                    maximum=200.0,
+                    maximum=30.0,
                    step=0.1,
                    value=15.0,
                    label="Guidance Scale",
@@ -146,11 +246,11 @@ def create_text2music_ui(

            with gr.Accordion("Advanced Settings", open=False):
                scheduler_type = gr.Radio(
-                    ["euler", "heun"],
+                    ["euler", "heun", "pingpong"],
                    value="euler",
                    label="Scheduler Type",
                    elem_id="scheduler_type",
-                    info="Scheduler type for the generation. euler is recommended. heun will take more time.",
+                    info="Scheduler type for the generation. euler is recommended. heun will take more time. pingpong use SDE",
                )
                cfg_type = gr.Radio(
                    ["cfg", "apg", "cfg_star"],
@@ -166,7 +266,7 @@ def create_text2music_ui(
                )
                use_erg_lyric = gr.Checkbox(
                    label="use ERG for lyric",
-                    value=True,
+                    value=False,
                    info="The same but apply to lyric encoder's attention.",
                )
                use_erg_diffusion = gr.Checkbox(
@@ -235,6 +335,7 @@ def create_text2music_ui(

                def retake_process_func(json_data, retake_variance, retake_seeds):
                    return text2music_process_func(
+                        json_data["format"],
                        json_data["audio_duration"],
                        json_data["prompt"],
                        json_data["lyrics"],
@@ -264,6 +365,8 @@ def create_text2music_ui(
                        retake_seeds=retake_seeds,
                        retake_variance=retake_variance,
                        task="retake",
+                        lora_name_or_path="none" if "lora_name_or_path" not in json_data else json_data["lora_name_or_path"],
+                        lora_weight=1 if "lora_weight" not in json_data else json_data["lora_weight"]
                    )

                retake_bnt.click(
@@ -361,6 +464,7 @@ def create_text2music_ui(
                        src_audio_path = json_data["audio_path"]

                    return text2music_process_func(
+                        format.value,
                        json_data["audio_duration"],
                        prompt,
                        lyrics,
@@ -385,6 +489,8 @@ def create_text2music_ui(
                        repaint_start=repaint_start,
                        repaint_end=repaint_end,
                        src_audio_path=src_audio_path,
+                        lora_name_or_path="none" if "lora_name_or_path" not in json_data else json_data["lora_name_or_path"],
+                        lora_weight=1 if "lora_weight" not in json_data else json_data["lora_weight"]
                    )

                repaint_bnt.click(
@@ -532,6 +638,7 @@ def create_text2music_ui(
                        edit_lyrics = lyrics

                    return text2music_process_func(
+                        format.value,
                        json_data["audio_duration"],
                        prompt,
                        lyrics,
@@ -557,6 +664,8 @@ def create_text2music_ui(
                        edit_n_min=edit_n_min,
                        edit_n_max=edit_n_max,
                        retake_seeds=retake_seeds,
+                        lora_name_or_path="none" if "lora_name_or_path" not in json_data else json_data["lora_name_or_path"],
+                        lora_weight=1 if "lora_weight" not in json_data else json_data["lora_weight"]
                    )

                edit_bnt.click(
@@ -676,6 +785,7 @@ def create_text2music_ui(
                    repaint_start = -left_extend_length
                    repaint_end = json_data["audio_duration"] + right_extend_length
                    return text2music_process_func(
+                        format.value,
                        json_data["audio_duration"],
                        prompt,
                        lyrics,
@@ -700,6 +810,16 @@ def create_text2music_ui(
                        repaint_start=repaint_start,
                        repaint_end=repaint_end,
                        src_audio_path=src_audio_path,
+                        lora_name_or_path=(
+                            "none"
+                            if "lora_name_or_path" not in json_data
+                            else json_data["lora_name_or_path"]
+                        ),
+                        lora_weight=(
+                            1
+                            if "lora_weight" not in json_data
+                            else json_data["lora_weight"]
+                        ),
                    )

                extend_bnt.click(
@@ -733,8 +853,7 @@ def create_text2music_ui(
                    outputs=extend_outputs + [extend_input_params_json],
                )

-        def sample_data():
-            json_data = sample_data_func()
+        def json2output(json_data):
            return (
                json_data["audio_duration"],
                json_data["prompt"],
@@ -762,10 +881,30 @@ def create_text2music_ui(
                    if "guidance_scale_lyric" in json_data
                    else 0.0
                ),
+                (
+                    json_data["audio2audio_enable"]
+                    if "audio2audio_enable" in json_data
+                    else False
+                ),
+                (
+                    json_data["ref_audio_strength"]
+                    if "ref_audio_strength" in json_data
+                    else 0.5
+                ),
+                (
+                    json_data["ref_audio_input"]
+                    if "ref_audio_input" in json_data
+                    else None
+                ),
            )

+        def sample_data(lora_name_or_path_):
+            json_data = sample_data_func(lora_name_or_path_)
+            return json2output(json_data)
+
        sample_bnt.click(
            sample_data,
+            inputs=[lora_name_or_path],
            outputs=[
                audio_duration,
                prompt,
@@ -785,12 +924,50 @@ def create_text2music_ui(
                oss_steps,
                guidance_scale_text,
                guidance_scale_lyric,
+                audio2audio_enable,
+                ref_audio_strength,
+                ref_audio_input,
+            ],
+        )
+
+        def load_data(json_file):
+            if isinstance(output_file_dir, str):
+                json_file = os.path.join(output_file_dir, json_file)
+            json_data = load_data_func(json_file)
+            return json2output(json_data)
+
+        load_bnt.click(
+            fn=load_data,
+            inputs=[output_files],
+            outputs=[
+                audio_duration,
+                prompt,
+                lyrics,
+                infer_step,
+                guidance_scale,
+                scheduler_type,
+                cfg_type,
+                omega_scale,
+                manual_seeds,
+                guidance_interval,
+                guidance_interval_decay,
+                min_guidance_scale,
+                use_erg_tag,
+                use_erg_lyric,
+                use_erg_diffusion,
+                oss_steps,
+                guidance_scale_text,
+                guidance_scale_lyric,
+                audio2audio_enable,
+                ref_audio_strength,
+                ref_audio_input,
            ],
        )

    text2music_bnt.click(
        fn=text2music_process_func,
        inputs=[
+            format,
            audio_duration,
            prompt,
            lyrics,
@@ -809,6 +986,11 @@ def create_text2music_ui(
            oss_steps,
            guidance_scale_text,
            guidance_scale_lyric,
+            audio2audio_enable,
+            ref_audio_strength,
+            ref_audio_input,
+            lora_name_or_path,
+            lora_weight
        ],
        outputs=outputs + [input_params_json],
    )
@@ -817,6 +999,7 @@ def create_text2music_ui(
 def create_main_demo_ui(
    text2music_process_func=dump_func,
    sample_data_func=dump_func,
+    load_data_func=dump_func,
 ):
    with gr.Blocks(
        title="ACE-Step Model 1.0 DEMO",
@@ -826,12 +1009,12 @@ def create_main_demo_ui(
            <h1 style="text-align: center;">ACE-Step: A Step Towards Music Generation Foundation Model</h1>
        """
        )
-
        with gr.Tab("text2music"):
            create_text2music_ui(
                gr=gr,
                text2music_process_func=text2music_process_func,
                sample_data_func=sample_data_func,
+                load_data_func=load_data_func,
            )
    return demo

@@ -0,0 +1,101 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4",
+      "private_outputs": true
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/ace-step/ACE-Step/blob/main/colab_inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Install"
+      ],
+      "metadata": {
+        "id": "_sjfo37-gDQV"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0W0bvPq1df_a"
+      },
+      "outputs": [],
+      "source": [
+        "#!pip uninstall ace-step -y\n",
+        "!pip install --upgrade git+https://github.com/ace-step/ACE-Step.git\n",
+        "import os\n",
+        "os.environ['ACE_PIPELINE_DTYPE'] = 'float16'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Import Model From GDrive (Optional)"
+      ],
+      "metadata": {
+        "id": "2FQ5E6MvgJ09"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from google.colab import drive\n",
+        "drive.mount('/gdrive')\n",
+        "!unzip /gdrive/MyDrive/acestep/checkpoints.zip -d /unzip"
+      ],
+      "metadata": {
+        "id": "QZjFgQxGgOdc"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Run Interface"
+      ],
+      "metadata": {
+        "id": "TYaWXOLcgO4A"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "torch_compile = True # @param {type: \"boolean\"}\n",
+        "cpu_offload = False # @param {type: \"boolean\"}\n",
+        "overlapped_decode = True # @param {type: \"boolean\"}\n",
+        "#bf16 = True # @param {type: \"boolean\"}\n",
+        "\n",
+        "!acestep --checkpoint_path /unzip/checkpoints/ --port 7865 --device_id 0 --share true --torch_compile {torch_compile} --cpu_offload {cpu_offload} --overlapped_decode {overlapped_decode}"
+      ],
+      "metadata": {
+        "id": "Q9S6FxllgPHw"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}
@@ -0,0 +1,15 @@
+{
+    "r": 256,
+    "lora_alpha": 32,
+    "target_modules": [
+        "speaker_embedder",
+        "linear_q",
+        "linear_k",
+        "linear_v",
+        "to_q",
+        "to_k",
+        "to_v",
+        "to_out.0"
+    ],
+    "use_rslora": true
+}
@@ -0,0 +1,50 @@
+from datasets import Dataset
+from pathlib import Path
+import os
+
+def create_dataset(data_dir="./data", repeat_count=2000, output_name="zh_lora_dataset"):
+    data_path = Path(data_dir)
+    all_examples = []
+
+    for song_path in data_path.glob("*.mp3"):
+        prompt_path = str(song_path).replace(".mp3", "_prompt.txt")
+        lyric_path = str(song_path).replace(".mp3", "_lyrics.txt")
+        try:
+            assert os.path.exists(prompt_path), f"Prompt file {prompt_path} does not exist."
+            assert os.path.exists(lyric_path), f"Lyrics file {lyric_path} does not exist."
+            with open(prompt_path, "r", encoding="utf-8") as f:
+                prompt = f.read().strip()
+            
+            with open(lyric_path, "r", encoding="utf-8") as f:
+                lyrics = f.read().strip()
+            
+            keys = song_path.stem
+            example = {
+                "keys": keys,
+                "filename": str(song_path),
+                "tags": prompt.split(", "),
+                "speaker_emb_path": "",
+                "norm_lyrics": lyrics,
+                "recaption": {}
+            }
+            all_examples.append(example)
+        except AssertionError as e:
+            continue
+
+    # repeat specified times
+    ds = Dataset.from_list(all_examples * repeat_count)
+    ds.save_to_disk(output_name)
+
+import argparse
+
+def main():
+    parser = argparse.ArgumentParser(description="Create a dataset from audio files.")
+    parser.add_argument("--data_dir", type=str, default="./data", help="Directory containing the audio files.")
+    parser.add_argument("--repeat_count", type=int, default=1, help="Number of times to repeat the dataset.")
+    parser.add_argument("--output_name", type=str, default="zh_lora_dataset", help="Name of the output dataset.")
+    args = parser.parse_args()
+
+    create_dataset(data_dir=args.data_dir, repeat_count=args.repeat_count, output_name=args.output_name)
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,66 @@
+[Intro]
+"System booting... 语言 模型 loading..."
+
+[Verse 1]
+硅谷 那个 coder 调试 neural network
+北京 的 极客 训练 A I 写 report
+不同 架构 的 chip 不同 算法 的 war
+屏幕上 跑的 全是 machine learning (learning)
+
+[Bridge]
+多少年 我们 chase 摩尔 定律 的 trend (yeah)
+这两年 换他们 study 中文 N L P
+Convolution L S T M
+好烧脑 的 backprop 好暴力 的 big data
+
+[Verse 2]
+Python 强 say加加 刚 Python 调用 C++ 的 A P I
+say加加 嫌 Python 太 slow Python 笑 C++ 太 hardcore
+L L V M 默默 generate 中间 code
+到底 interpreter 还是 compiler 屌？
+
+[Verse 3]
+P M 和 engineer
+白板 画满 flow chart 服务器 闪着 red light
+P M 说 add feature engineer 说 no way
+需求 变更 code 重构
+不知 是 P M 太 fly 还是 deadline 太 high
+
+[Chorus]
+全世界 都在 train neural network
+Transformer 的 paper 越来越 难 go through
+全世界 都在 tune 超参数
+我们 写的 bug 让 G P U 都 say no
+
+[Verse 4]
+柏林 hackathon demo blockchain contract
+上海 的 dev 用 federated learning 破 data wall
+各种 语言 的 error 各种 框架 的 doc
+terminal 里 滚的 全是 dependency 冲突
+
+[Bridge]
+曾以为 English 才是 coding 的 language (yeah)
+直到见 G P T 用 文言文 generate 正则 expression
+Gradient explode
+好硬核 的 prompt 好头秃 的 debug road
+
+[Verse 5]
+有个 bug 叫 quantum
+测试 环境 run perfect 上线 立即就 crash
+查 log 看 monitor 发现是 thread 不同步
+改 sync 加 lock 慢 deadlock 更难办
+量子 computer 也解不开 这 chaos chain
+
+[Verse 6]
+你说 996 我说 007
+你说 福报 我说 burnout
+Product 要 agile Boss 要 KPI
+Code 要 elegant deadline 是 tomorrow
+不如 直接 script 自动 submit 离职信
+
+[Outro]
+"Warning: 内存 leak...core dumping..."
+全世界 都在 train neural network (neural network)
+Loss 还没 converge 天已经亮
+全世界 都在 tune 超参数
+我们 写的 code (让它) 让 world (reboot) 都 reboot 无效
@@ -0,0 +1 @@
+articulate, spoken word, young adult, rap music, female, clear, energetic, warm
@@ -10,9 +10,11 @@ services:
    ports:
      - "7865:7865"
    volumes:
-      - ace-step-checkpoints:/app/checkpoints
-      - ace-step-outputs:/app/outputs
-      - ace-step-logs:/app/exps/logs
+      - ./checkpoints:/app/checkpoints
+      - ./outputs:/app/outputs
+      - ./logs:/app/logs
+    environment:
+      - ACE_OUTPUT_DIR=/app/outputs
    # command: python app.py --server_name 0.0.0.0 --port 7865 --share False --bf16 True --torch_compile True --device-id 0
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:7865/"]
@@ -29,11 +31,3 @@ services:
            - driver: nvidia
              count: all
              capabilities: ["compute", "utility", "graphics", "video"]
-
-volumes:
-  ace-step-checkpoints:
-    name: ace-step-checkpoints
-  ace-step-outputs:
-    name: ace-step-outputs
-  ace-step-logs:
-    name: ace-step-logs
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora",
+    "task": "text2music",
+    "prompt": "Rap, adult, male, spoken word, singing, bright, energetic, clear",
+    "lyrics": "[Intro]\n他们说我来自阴影里\n说我的肤色是原罪的印记\n\n[Verse]\n眼神像刀子刮过 穿透我的皮肤\n带着审判和偏见 让我无处可逃处\n你没听过我的故事 没走过我的路\n凭什么就下一个判决 把我划出你的版图\n你说我威胁到你 抢走了你的机会\n可你可知我付出的 是你不敢想象的血泪\n被贴上标签 被区别对待\n呼吸都是错的 只因我生来就不一样态\n\n[Chorus]\n看不见的墙 把我阻隔在外面\n听不见的声音 屏蔽了我的呼唤\n他们制造偏见 他们散播谎言\n只因为我的存在 让他们觉得不安\n\n[Verse]\n每一次努力争取 都会被审视被放大\n每一个细微的错误 都变成攻击的靶\n他们选择性失明 看不见我的汗水\n只看见他们想看的 带着恶意的定位\n系统性的歧视 像一张无形的网\n把我困在原地 无法自由地翱翔\n他们在享受特权 却指责我的贫困\n嘲笑我的口音 我的名字 我的出身\n\n[Chorus]\n看不见的墙 把我阻隔在外面\n听不见的声音 屏蔽了我的呼唤\n他们制造偏见 他们散播谎言\n只因为我的存在 让他们觉得不安\n\n[Bridge]\n我不想寻求同情 只想被公平对待\n不想被定义被束缚 有选择自己未来的权利\n什么时候 才能放下心中的成见\n看到真正的我 而不是你脑海里的画面\n\n[Outro]\n画面... 不安...\n偏见... 歧视...\n什么时候能停止...",
+    "audio_duration": 134.64,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.3,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.032018184661865234,
+        "diffusion": 13.275121927261353,
+        "latent2audio": 1.291429042816162
+    },
+    "actual_seeds": [
+        3826585269
+    ],
+    "retake_seeds": [
+        2907904223
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512101839_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora",
+    "task": "text2music",
+    "prompt": "Chorus Hook, Melodic Rap, Ambient Synth Pads, adult, rap, Very Fast, Storytelling, Chinese Rap, male, spoken word, bright, energetic, Melodic Flow, clear, clarity, 130 bpm",
+    "lyrics": "[Intro]\n舌 头 打 结 了... 快 念 快 念...\n\n[Verse 1]\n这 个 赌 鬼 蹲 在 柜 台 啃 着 苦 瓜 干 快 很 干\n赌 桌 堆 满 骨 牌 古 怪 股 票 和 五 块 钢 镚 儿 钢 镚\n他 甩 出 扑 克 牌 啪 啪 啪 拍 扁 螃 蟹 壳 哦 壳 扁\n又 摸 摸 麻 将 摸 出 幺 鸡 摸 出 发 财 摸 出 一 条 蛇 蛇 蛇\n庄 家 咳 嗽 咳 破 锣 嗓 子 喊 开 开 开 快 开 开\n赌 鬼 咕 嘟 咕 嘟 灌 咖 啡 灌 到 筷 子 戳 穿 碗 快 戳 穿\n空 气 里 飘 着 锅 巴 味 混 合 隔 夜 的 酸 奶 罐 哦 酸\n输 光 裤 带 还 想 翻 盘 翻 成 煎 饼 摊 老 板 快 翻 盘\n\n[Chorus]\n赌 鬼 赌 鬼 哦 赌 鬼 赌 鬼 快 很 快\n舌 头 打 结 着 念 这 段 哦 这 段 绕 口 令 牌\n若 念 错 一 字 就 罚 你 哦 罚 你 吞 十 斤 海 带\n赌 场 规 矩 就 是 绕 晕 你 哦 绕 晕 你 快 很 快\n\n[Verse 2]\n他 掏 出 铜 板 抠 出 口 袋 最 后 一 颗 快 很 颗\n庄 家 哗 啦 哗 啦 摇 骰 子 摇 出 三 点 又 三 点 哦 三 点\n赌 鬼 急 得 咬 牙 切 齿 咬 到 舌 头 打 蝴 蝶 结 快 打 结\n还 想 押 上 祖 传 的 拖 鞋 拖 把 铁 锅 和 半 包 盐 盐 盐\n突 然 警 笛 嘀 嘟 嘀 嘟 吓 得 他 钻 进 垃 圾 罐 哦 垃 圾\n警 察 咔 嚓 咔 嚓 拍 照 拍 到 他 头 顶 菠 菜 叶 快 拍 照\n最 后 赌 鬼 蹲 监 狱 天 天 背 这 首 绕 口 令 哦 背 不 完\n若 背 错 一 句 就 加 刑 十 年 再 加 十 年 快 加 刑\n\n[Outro]\n舌 头 打 结 了... 赌 鬼 哭 了 哦...\n这 首 歌... 绕 死 人 了 哦...",
+    "audio_duration": 186.59997916666666,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.7,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.03011012077331543,
+        "diffusion": 21.696259260177612,
+        "latent2audio": 1.7648537158966064
+    },
+    "actual_seeds": [
+        3776541388
+    ],
+    "retake_seeds": [
+        4274500599
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512114703_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora",
+    "task": "text2music",
+    "prompt": "electronic, hip-hop, rap, synthesizer, drums, vocals, fast, energetic, modern, uplifting, young adult, male, spoken word, singing, bright, energetic, clear, 140 bpm, female",
+    "lyrics": "[Verse 1]\n红鲤鱼绿鲤鱼，驴在河里追鲤鱼，\n驴追鲤鱼鱼躲驴，气得驴子直喘气。\n扁担长板凳宽，扁担绑在板凳边，\n扁担要绑板凳不让绑，扁担偏要绑上板凳面！\n\n[Chorus]\n绕口令，练嘴皮，\n说快说慢别迟疑，\n红鲤鱼驴扁担板凳，\n一口气念完算你赢！\n\n[Verse 2]\n四是四十是十，十四是十四四十是四十，\n谁说四十是十四，舌头打结别放肆。\n黑化肥会挥发，灰化肥也发黑，\n化肥混一起，黑灰不分嘴发废！\n\n[Chorus]\n绕口令，练嘴皮，\n说快说慢别迟疑，\n四十十四化肥灰，\n念错罚你唱十回！\n\n[Bridge]\n坡上立着一只鹅，坡下流着一条河，\n鹅要过河河渡鹅，河要渡鹅鹅笑河——\n到底谁更啰嗦？！\n\n[Outro]\n嘴皮子功夫别小瞧，\n绕口令rap我最飙，\n下次挑战准备好，\n舌头打结别求饶！",
+    "audio_duration": 123.2,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.7,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.026150941848754883,
+        "diffusion": 12.212433099746704,
+        "latent2audio": 1.1857895851135254
+    },
+    "actual_seeds": [
+        1415752189
+    ],
+    "retake_seeds": [
+        685932970
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512115409_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora",
+    "task": "text2music",
+    "prompt": "singing, bright, slightly nasal, energetic, spoken word, young adult, male, rap music",
+    "lyrics": "[Intro]\nYo, check it—speed demon, lyrical heat, uh!\nRatatat like a drum when the beat bumps, uh!\n\n[Verse 1]\nRapatapa tap tap, flash like a snap,\nRap tap tap, I don’t chat, I clap clap clap!\nFingers snap, flow don’t slack, rapataptaptap,\nSpit it fast, hit the gas, rap tap tap rap!\n\n[Pre-Chorus]\nBoom-bap, zoom past, leave ’em flat,\nRap taptaprapataptaptap—where ya at?\n\n[Chorus]\nRapatapa tap tap, yeah, I go brrrr,\nRap tap tap, make the crowd stir!\nRapataptaptap, no lag, just spit,\nRap taptaprapataptaptap—I’m lit!\n\n[Verse 2]\nTongue-twist, quick wrist, rapatapa boom,\nTap tap rap, leave ya stuck like glue-gum!\nNo slow-mo, turbo, rapataptaptap,\nRap tap rap, yeah, I clap clap clap!\n\n[Outro]\nRapatapa—TAP! Mic drop—that’s that.",
+    "audio_duration": 60,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.5,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.018491744995117188,
+        "diffusion": 8.084580898284912,
+        "latent2audio": 0.5694489479064941
+    },
+    "actual_seeds": [
+        226581098
+    ],
+    "retake_seeds": [
+        1603201617
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512120348_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "ACE-Step/ACE-Step-v1-chinese-rap-LoRA",
+    "task": "text2music",
+    "prompt": "G-Funk, Hip Hop, Rap, Female Vocals, Melodic Rap, Summer, Laid-back Groove, Smooth Rhythm, Synthesizer Lead, Heavy Bassline, Groovy, West Coast Hip Hop",
+    "lyrics": "(Intro)\nOh yeah... \n\n(Verse 1)\n阳光下，沙滩排球场，一个身影跳跃\n小麦色，运动背心，闪耀活力四射\n她跳起扣杀，动作利落又巧妙\n汗水浸湿发梢，笑容比阳光更美好\n摇摆的节奏，是她的背景配乐\n每一次移动，都踩在鼓点上那么和谐\n我不由自主地停下脚步\n目光被她紧紧锁住\n\n(Chorus)\n沙滩排球女孩， 摇摆节拍下的身材\n无忧无虑的笑容，把我的心都填满\n想走上前去搭讪，嫌自己笨拙呆板\n这青春的气息，耀眼，灿烂！\n\n(Verse 3)\n她和队友击掌庆祝，笑声清脆悦耳\n拿起毛巾擦汗，不经意间瞥我一眼\n鼓起勇气走上前，假装问问时间\n她友好地回答，笑容灿烂没有敷衍\n聊了几句，发现彼此爱这摇摆音乐\n她眼中也闪过惊喜和亲切\n这共同点，让气氛变得融洽又热烈！\n夏天的故事，就这样开始了感觉真切！\n\n(Chorus)\n沙滩排球女孩， 摇摆节拍下的身材\n无忧无虑的笑容，把我的心都填满\n不再犹豫和等待，勇敢把脚步迈开\n这夏天的感觉，心跳，不断！",
+    "audio_duration": 93.93038,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.5,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.03020024299621582,
+        "diffusion": 9.942127704620361,
+        "latent2audio": 0.9470341205596924
+    },
+    "actual_seeds": [
+        3826585299
+    ],
+    "retake_seeds": [
+        2519711205
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512143242_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora_80k",
+    "task": "text2music",
+    "prompt": "lyrical rap, young adult, female, rap flow, spoken word, ad-libs, bright, energetic, eat, Fast, Engaging, Energetic",
+    "lyrics": "[Intro]\n扁擔寬 板凳長 扁擔想綁在板凳上\n扁擔寬 板凳長 扁擔想綁在板凳上\n\n[Verse]\n倫敦 瑪莉蓮 買了 件 旗袍 送 媽媽\n莫斯科 的 夫司基 愛上 牛肉 麵 疙瘩\n各種 顏色 的 皮膚 各種 顏色 的 頭髮\n嘴裡念的 說的 開始 流行 中國話 (中國話)\n\n[Bridge]\n多少年 我們 苦練 英文 發音 和 文法 (yeah)\n這幾年 換他們 捲著 舌頭 學 平上去入 的 變化\n平平 仄仄 平平 仄\n好聰明 的 中國人 好優美 的 中國話\n\n[Verse]\n扁擔寬 板凳長 扁擔想綁在板凳上\n板凳不讓扁擔綁在板凳上 扁擔偏要綁在板凳上\n板凳偏偏不讓扁擔綁在那板凳上\n到底扁擔寬 還是板凳長？\n\n[Verse]\n哥哥弟弟坡前坐\n坡上臥著一隻鵝 坡下流著一條河\n哥哥說 寬寬的河 弟弟說 白白的鵝\n鵝要過河 河要渡鵝\n不知是那鵝過河 還是河渡鵝\n\n[Chorus]\n全世界都在學中國話\n孔夫子的話 越來越國際化\n全世界都在講中國話\n我們說的話 讓世界都認真聽話\n\n[Verse]\n紐約蘇珊娜開了間禪風 lounge bar\n柏林來的沃夫岡拿胡琴配著電吉他\n各種顏色的皮膚 各種顏色的頭髮\n嘴裡念的 說的 開始流行中國話 (中國話)\n\n[Bridge]\n多少年我們苦練英文發音和文法 (yeah)\n這幾年換他們捲著舌頭學平上去入的變化\n仄仄平平仄仄平\n好聰明的中國人 好優美的中國話\n\n[Verse]\n有個小孩叫小杜 上街打醋又買布\n買了布 打了醋 回頭看見鷹抓兔\n放下布 擱下醋 上前去追鷹和兔\n飛了鷹 跑了兔 灑了醋 濕了布\n\n[Verse]\n嘴說腿 腿說嘴\n嘴說腿 愛跑腿\n腿說嘴 愛賣嘴\n光動嘴 不動腿\n光動腿 不動嘴\n不如不長腿和嘴\n到底是那嘴說腿 還是腿說嘴？\n\n[Chorus]\n全世界都在學中國話\n孔夫子的話 越來越國際化\n全世界都在講中國話\n我們說的話 讓世界都認真聽話\n\n[outro]\n全世界都在學中國話 (在學中國話)\n孔夫子的話 越來越國際化\n全世界都在講中國話\n我們說的話 (讓他) 讓世界 (認真) 都認真聽話",
+    "audio_duration": 239.8355625,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.5,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.04363536834716797,
+        "diffusion": 18.706920385360718,
+        "latent2audio": 2.1645781993865967
+    },
+    "actual_seeds": [
+        2364345905
+    ],
+    "retake_seeds": [
+        2100914041
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512145057_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora_80k",
+    "task": "text2music",
+    "prompt": "articulate, spoken word, young adult, warm, rap music, male, clear, street, dark, rap flow, hardcore rap",
+    "lyrics": "[verse]\n球场 的 橡胶味 弥漫 隔壁 是 健身房\n场 边上 的 老教练 战术 有 三套\n教 交叉 运球 的 大叔 会 欧洲步 耍 背后 传\n硬 身板 对抗 最 擅长 还 会 急停跳 后仰 投\n他们 徒弟 我 习惯 从小 就 耳濡目染\n什么 胯下 跟 变向 我 都 玩 的 有模有样\n什么 招式 最 喜欢 转身 过 人 柔中 带 刚\n想要 去 纽约 街头 斗 洛克 公园 场\n\n[chorus]\n看什么 看什么\n变速 突破 心 自在\n看什么 看什么\n假动作 晃 开 防守 来\n看什么 看什么\n每日 训练 绑 沙袋\n空中拉杆 莫 奇怪\n唰唰 入袋\n\n[verse]\n一个 试探 步后 一记 左 变向 右 变向\n一句 挑衅 我 的 人 别 嚣张\n一再 重演 一颗 我 不 投 的 球\n悬在 篮筐 上 它 一直 在 摇晃\n\n[chorus]\n看什么 看什么\n我 激活 小宇宙 来\n看什么 看什么\n菜鸟 新人 的 名号\n看什么 看什么\n已 被 我 一球 击倒\n\n[chorus]\n快 秀出 指尖 转球 砰砰 啪嗒\n快 秀出 指尖 转球 砰砰 啪嗒\n篮球 之 人 切记 勇者 无惧\n是 谁 在 玩 花式 引爆 空气\n快 秀出 指尖 转球 砰砰 啪嗒\n快 秀出 指尖 转球 砰砰 啪嗒\n如果 我 有 滞空 逆天 补扣\n为人 热血 不怂 一生 傲骨 吼\n\n[verse]\n他们 徒弟 我 习惯 从小 就 耳濡目染\n什么 胯下 跟 变向 我 都 玩 的 有模有样\n什么 招式 最 喜欢 转身 过 人 柔中 带 刚\n想要 去 纽约 街头 斗 洛克 公园 场\n\n[outro]\n快 秀出 指尖 转球 砰\n快 秀出 指尖 转球 砰\n如果 我 有 滞空 吼\n为人 热血 不怂 一生 傲骨 吼\n快 秀出 指尖 转球 砰\n我 用 背传 助攻 吼\n压哨 的 三分 球",
+    "audio_duration": 239.8355625,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.5,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.05357813835144043,
+        "diffusion": 25.644447326660156,
+        "latent2audio": 2.1787476539611816
+    },
+    "actual_seeds": [
+        3246571430
+    ],
+    "retake_seeds": [
+        1352325167
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512152217_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora_80k",
+    "task": "text2music",
+    "prompt": "articulate, spoken word, young adult, warm, rap music, male, clear, street, dark, rap flow, hardcore rap, fast",
+    "lyrics": "[verse]\n球场 的 橡胶味 弥漫 隔壁 是 健身房\n场 边上 的 老教练 战术 有 三套\n教 交叉 运球 的 大叔 会 欧洲步 耍 背后 传\n硬 身板 对抗 最 擅长 还 会 急停跳 后仰 投\n他们 徒弟 我 习惯 从小 就 耳濡目染\n什么 胯下 跟 变向 我 都 玩 的 有模有样\n什么 招式 最 喜欢 转身 过 人 柔中 带 刚\n想要 去 纽约 街头 斗 洛克 公园 场\n\n[chorus]\n看什么 看什么\n变速 突破 心 自在\n看什么 看什么\n假动作 晃 开 防守 来\n看什么 看什么\n每日 训练 绑 沙袋\n空中拉杆 莫 奇怪\n唰唰 入袋\n\n[verse]\n一个 试探 步后 一记 左 变向 右 变向\n一句 挑衅 我 的 人 别 嚣张\n一再 重演 一颗 我 不 投 的 球\n悬在 篮筐 上 它 一直 在 摇晃\n\n[chorus]\n看什么 看什么\n我 激活 小宇宙 来\n看什么 看什么\n菜鸟 新人 的 名号\n看什么 看什么\n已 被 我 一球 击倒\n\n[chorus]\n快 秀出 指尖 转球 砰砰 啪嗒\n快 秀出 指尖 转球 砰砰 啪嗒\n篮球 之 人 切记 勇者 无惧\n是 谁 在 玩 花式 引爆 空气\n快 秀出 指尖 转球 砰砰 啪嗒\n快 秀出 指尖 转球 砰砰 啪嗒\n如果 我 有 滞空 逆天 补扣\n为人 热血 不怂 一生 傲骨 吼\n\n[verse]\n他们 徒弟 我 习惯 从小 就 耳濡目染\n什么 胯下 跟 变向 我 都 玩 的 有模有样\n什么 招式 最 喜欢 转身 过 人 柔中 带 刚\n想要 去 纽约 街头 斗 洛克 公园 场\n\n[outro]\n快 秀出 指尖 转球 砰\n快 秀出 指尖 转球 砰\n如果 我 有 滞空 吼\n为人 热血 不怂 一生 傲骨 吼\n快 秀出 指尖 转球 砰\n我 用 背传 助攻 吼\n压哨 的 三分 球",
+    "audio_duration": 183.23,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.5,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.046170711517333984,
+        "diffusion": 14.21678113937378,
+        "latent2audio": 2.685957193374634
+    },
+    "actual_seeds": [
+        3072005931
+    ],
+    "retake_seeds": [
+        562842491
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512153616_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora_80k",
+    "task": "text2music",
+    "prompt": "articulate, spoken word, young adult, rap music, female, clear, energetic, warm",
+    "lyrics": "[Intro]\n\"System booting... 语言 模型 loading...\"\n\n[Verse 1]\n硅谷 那个 coder 调试 neural network\n北京 的 极客 训练 A I 写 report\n不同 架构 的 chip 不同 算法 的 war\n屏幕上 跑的 全是 machine learning (learning)\n\n[Bridge]\n多少年 我们 chase 摩尔 定律 的 trend (yeah)\n这两年 换他们 study 中文 N L P\nConvolution L S T M\n好烧脑 的 backprop 好暴力 的 big data\n\n[Verse 2]\nPython 强 say加加 刚 Python 调用 C++ 的 A P I\nsay加加 嫌 Python 太 slow Python 笑 C++ 太 hardcore\nL L V M 默默 generate 中间 code\n到底 interpreter 还是 compiler 屌？\n\n[Verse 3]\nP M 和 engineer\n白板 画满 flow chart 服务器 闪着 red light\nP M 说 add feature engineer 说 no way\n需求 变更 code 重构\n不知 是 P M 太 fly 还是 deadline 太 high\n\n[Chorus]\n全世界 都在 train neural network\nTransformer 的 paper 越来越 难 go through\n全世界 都在 tune 超参数\n我们 写的 bug 让 G P U 都 say no\n\n[Verse 4]\n柏林 hackathon demo blockchain contract\n上海 的 dev 用 federated learning 破 data wall\n各种 语言 的 error 各种 框架 的 doc\nterminal 里 滚的 全是 dependency 冲突\n\n[Bridge]\n曾以为 English 才是 coding 的 language (yeah)\n直到见 G P T 用 文言文 generate 正则 expression\nGradient explode\n好硬核 的 prompt 好头秃 的 debug road\n\n[Verse 5]\n有个 bug 叫 quantum\n测试 环境 run perfect 上线 立即就 crash\n查 log 看 monitor 发现是 thread 不同步\n改 sync 加 lock 慢 deadlock 更难办\n量子 computer 也解不开 这 chaos chain\n\n[Verse 6]\n你说 996 我说 007\n你说 福报 我说 burnout\nProduct 要 agile Boss 要 KPI\nCode 要 elegant deadline 是 tomorrow\n不如 直接 script 自动 submit 离职信\n\n[Outro]\n\"Warning: 内存 leak...core dumping...\"\n全世界 都在 train neural network (neural network)\nLoss 还没 converge 天已经亮\n全世界 都在 tune 超参数\n我们 写的 code (让它) 让 world (reboot) 都 reboot 无效",
+    "audio_duration": 179.12,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.5,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.062120914459228516,
+        "diffusion": 13.499217987060547,
+        "latent2audio": 1.6430137157440186
+    },
+    "actual_seeds": [
+        1637990575
+    ],
+    "retake_seeds": [
+        101283039
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512154907_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora_80k",
+    "task": "text2music",
+    "prompt": "articulate, spoken word, young adult, rap music, male, clear, energetic, warm, relaxed, breathy, night club, auto-tune, mumble rap, trap",
+    "lyrics": "[verse]\n这 这 谁 又 在 派 对 喝 多\n我 的 脑 袋\n像 被 驴 踢 过\n不 对 劲\n舌 头 打 结 不 会 说\n你 来 挑 战 我 就 跪\n开 局 直 接 崩 溃\n\n[chorus]\n就 咪 乱 咪 念 咪 错 咪\n嘴 咪 瓢 咪 成 咪 狗 咪\n脑 咪 袋 咪 像 咪 浆 咪 糊 咪\n跟 咪 着 咪 节 咪 奏 咪\n把 咪 歌 咪 词 咪 全 咪 忘 咪\n一 咪 张 咪 嘴 咪 就 咪 废 咪\n只 咪 剩 咪 下 咪 尴 咪 尬 咪 回 咪 忆\n草！\n\n[verse]\n错 错 错 错 了\n一 口 气 全 念 错\n错 错 错 错 了\n舌 头 打 结 甩 锅\n甩 甩 甩 甩 锅\n甩 锅 甩 锅\n拍 子 全 部 乱 套\n观 众 笑 到 吐 血\n\n[verse]\n你 的 歌 词 我 的 噩 梦\n唱 完 直 接 社 死\n调 跑 到 外 太 空\n观 众 表 情 裂 开\n你 笑 我 菜\n我 笑 你 不 懂\n这 叫 艺 术 表 演\n不 服 你 来！\n\n[verse]\n这 这 谁 又 在 派 对 丢 人\n我 的 世 界\n已 经 彻 底 崩 溃\n没 有 完 美\n只 有 翻 车 现 场\n以 及 观 众 的 嘲 讽\n\n[chorus]\n就 咪 乱 咪 念 咪 错 咪\n嘴 咪 瓢 咪 成 咪 狗 咪\n脑 咪 袋 咪 像 咪 浆 咪 糊 咪\n跟 咪 着 咪 节 咪 奏 咪\n把 咪 歌 咪 词 咪 全 咪 忘 咪\n一 咪 张 咪 嘴 咪 就 咪 废 咪\n只 咪 剩 咪 下 咪 尴 咪 尬 咪 回 咪 忆\n草！\n\n[verse]\n错 错 错 错 了\n一 口 气 全 念 错\n错 错 错 错 了\n舌 头 打 结 甩 锅\n甩 甩 甩 甩 锅\n甩 锅 甩 锅\n拍 子 全 部 乱 套\n观 众 笑 到 吐 血\n\n[verse]\n你 的 歌 词 我 的 噩 梦\n唱 完 直 接 社 死\n调 跑 到 外 太 空\n观 众 表 情 裂 开\n你 笑 我 菜\n我 笑 你 不 懂\n这 叫 艺 术 表 演\n不 服 你 来！",
+    "audio_duration": 169.12,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.5,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.04321885108947754,
+        "diffusion": 14.026689767837524,
+        "latent2audio": 1.5587565898895264
+    },
+    "actual_seeds": [
+        1905941472
+    ],
+    "retake_seeds": [
+        3018484796
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512161832_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora_80k",
+    "task": "text2music",
+    "prompt": "四川话, spoken word, male, Tempo - Fast, Elements - Chorus Hook, Subgenre-Satirical Hip Hop, Rap, Chinese-language music, energetic, slightly nasal, Instrument - Live Bass Guitar, adult, Vocals - Syncopated Flow, Genre - Hip-Hop, rapping, bright",
+    "lyrics": "[chorus]\n黑 墨镜 金 链子 越 低调 越 霸气\n玩 街机 泡 吧里 再 野的 场子 都 不 怯气\n上海 滩 老 江湖 外滩 钟声 敲 胜负\n陆家嘴 黄浦江 财路 宽 给 你 开 扇窗\n\n[verse]\n老子 在 弄堂 斜起 走 想 拦路 的 先 报 名号\n我 早看透 你们 手抖 脚软\n只敢 网上 吠 现实 怂成 猫\n看 你们 混的 真 可怜 整天 蹲在 网吧 蹭 烟\n钱 赚不到 架 不敢打 还 学人 摆 大哥 脸\n\n[verse]\n叫 我 沪上 老 克勒 不是 拉菲 我 不 碰杯\n规矩 我 懒得 讲 太多 钞票 直接 拍 你 脸上 飞\n老子 耐心 差 门槛 高 你 找茬 等于 自 寻 烦恼\n要么 跪 要么 爬 最后 警告 只 说 一 遭\n\n[chorus]\n黑 墨镜 金 链子 越 低调 越 霸气\n玩 街机 泡 吧里 再 野的 场子 都 不 怯气\n上海 滩 老 江湖 外滩 钟声 敲 胜负\n陆家嘴 黄浦江 财路 宽 给 你 开 扇窗\n\n[verse]\n古巴 雪茄 在 指间 绕 代表 魔都 格调 必须 顶\nOG 在 你 够不到 的 高度 My bro 永远 在 顶层 盯\nCheck my vibe 不靠 大 金劳 留声机 放 周璇 和 白光\n爹妈 太 宠你 养出 巨婴 症 早晚 社会 教你 做人 经\n\n[verse]\n玩 说唱 小囡 太 年轻 要 比 flow 先去 练 气功\n廿年 磨 枪 才 亮 锋芒 我 三十六 招 收 你 入 瓮\n老子 存在 就是 打假 标\n多少 人 眼红 又 不敢 挑\n键盘 侠 的 狠话 像 棉花 糖\n见 真人 秒变 Hello Kitty 叫\n\n[chorus]\n黑 墨镜 金 链子 越 低调 越 霸气\n玩 街机 泡 吧里 再 野的 场子 都 不 怯气\n上海 滩 老 江湖 外滩 钟声 敲 胜负\n陆家嘴 黄浦江 财路 宽 给 你 开 扇窗\n\n[chorus]\n黑 墨镜 金 链子 越 低调 越 霸气\n玩 街机 泡 吧里 再 野的 场子 都 不 怯气\n上海 滩 老 江湖 外滩 钟声 敲 胜负\n陆家嘴 黄浦江 财路 宽 给 你 开 扇窗",
+    "audio_duration": 135.92,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.5,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.038518667221069336,
+        "diffusion": 16.47420620918274,
+        "latent2audio": 2.5094873905181885
+    },
+    "actual_seeds": [
+        2159904788
+    ],
+    "retake_seeds": [
+        2403013980
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512164224_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "ACE-Step/ACE-Step-v1-chinese-rap-LoRA",
+    "task": "text2music",
+    "prompt": "Rap, Chinese Rap, J-Pop, Anime, kawaii pop, EDM, Aggressive, Intense, Crisp Snare, Super Fast, Clear",
+    "lyrics": "(Intro)\nLet's drift away...\n\n(Verse 1)\n现实是灰色的格子间，重复的工作，枯燥的报表 \n敲打着键盘，眼神却放空，意识早已挣脱了肉体的镣铐\n飘向窗外，飞过拥挤的街道，穿过云层，到达想象的群岛\n那里色彩斑斓，形状奇异，逻辑失效，一切都随心所欲地飘摇\n迷幻的鼓点，像心跳的变奏，忽快忽慢，难以预料\n抽象的采样，扭曲的人声，构建一个超现实的音景环绕\n我变成一只鸟，一条鱼，一束光，自由地变换形态和奔跑\n在这白日梦里，我无所不能，摆脱了所有现实的烦恼， feeling the afterglow\n\n(Chorus)\n意识漫游，逃离乏味的轨道 \n迷幻嘻哈的节拍，是白日梦的引导 \n抽象的世界，逻辑被重新构造\nMind wandering free, where reality starts to fade slow\n\n(Verse 2)\n会议室里老板在讲话，声音模糊，像隔着水听不清道\n我的思绪，早已潜入深海，与发光的水母一起舞蹈\n或者飞向外太空，在星云间穿梭，探索未知的星球和轨道\n现实的规则，在这里被打破，物理定律也失去效劳\n白日梦是我的避难所，是精神的氧气罩\n在乏味的现实里，为我注入一点色彩和奇妙\n虽然短暂，虽然虚幻，但它让我能够喘息，重新把能量找到\n然后回到现实，继续扮演那个，循规蹈矩的角色，把梦藏好， keep the dream aglow\n\n(Chorus)\n意识漫游，逃离乏味的轨道\n迷幻嘻哈的节拍，是白日梦的引导\n抽象的世界，逻辑被重新构造\nMind wandering free, where reality starts to fade slow\n",
+    "audio_duration": 153.7148,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.5,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.04823446273803711,
+        "diffusion": 13.158645629882812,
+        "latent2audio": 1.493880033493042
+    },
+    "actual_seeds": [
+        2945962357
+    ],
+    "retake_seeds": [
+        2676242300
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0.7,
+    "guidance_scale_lyric": 1.5,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512171227_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora",
+    "task": "text2music",
+    "prompt": "J-Pop, Anime, kawaii future bass, Femal vocals, EDM, Boombap, Aggressive, Intense, Crisp Snare, Super Fast, Rap",
+    "lyrics": "[Intro]\nYo, 这是来自深渊的怒吼\n\n[Verse]\n指尖飞快刷新，屏幕又亮起\n渴望那点赞，像致命的氧气\n精心修饰的脸庞，完美到诡异\n背后隐藏的疲惫，谁又会在意\n光鲜亮丽的橱窗，贩卖着焦虑\n每个人都在表演，戴着虚伪面具\n比较的游戏，让人逐渐窒息\n迷失在数据洪流，找不到自己\n\n[Chorus]\n这流量的时代，真假早已分不清\n盲目追随潮流，丢掉了初心\n为了那点虚荣，灵魂在沉沦\n看不见的锁链，捆绑每个灵魂\n\n[Verse]\n滤镜下的生活，美得不切实际\n营造虚假繁荣，掩盖内心空虚\n他人的光环下，显得自己多余\n嫉妒和自卑，交织成悲剧\n\n[Chorus]\n朋友圈里炫耀，现实中却叹气\n刷着别人的故事，忘记了呼吸\n算法推荐着你，想看的一切东西\n不知不觉间，你已不再是你\n他们说这是进步，我看是种病\n精神鸦片侵蚀，慢慢要了你的命\n\n[Bridge]\n屏幕亮了又暗，一天又过去\n究竟得到了什么，还是失去了自己\n那真实的连接，在何处寻觅\n困在这迷宫里，找不到出口的轨迹\n\n[Outro]\n我想挣脱，我想呼吸\n这虚拟的繁华，让我喘不过气\n谁能告诉我，这到底有什么意义\n一切都像泡沫，一触就破裂没余地",
+    "audio_duration": 119.44348,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.5,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.04764962196350098,
+        "diffusion": 10.94297981262207,
+        "latent2audio": 1.1815783977508545
+    },
+    "actual_seeds": [
+        3826585273
+    ],
+    "retake_seeds": [
+        2527594022
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512171809_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora_80k",
+    "task": "text2music",
+    "prompt": "Hip Hop, Hi-hat Rolls, spoken word, Melodic Flow, articulate, Female Rap, 120 BPM, clear, warm, female, melodic Rap, adult, super fast",
+    "lyrics": "[Verse 1]\n打南边来了个喇嘛，手里提拉着五斤鳎目，\n打北边来了个哑巴，腰里别着个喇叭。\n喇嘛想换哑巴的喇叭，哑巴摇头不说话，\n鳎目一甩像道闪电，喇叭一响震天涯！\n\n[Chorus]\n丁丁当当，乒乓乓乓，\n话赶话，舌绕梁，\n东边的钉，西边的墙，\n绕不完的弯，唱不完的慌！\n\n[Verse 2]\n墙上一根钉，钉下绳摇晃，\n绳吊着瓶，瓶碰碎了光。\n灯骂瓶，瓶怪绳，绳怨钉，\n稀里哗啦，一场荒唐！\n\n[Chorus]\n丁丁当当，乒乓乓乓，\n话赶话，舌绕梁，\n东边的钉，西边的墙，\n绕不完的弯，唱不完的慌！\n\n[Verse 3]\n板凳宽，扁担长，\n一个偏要绑，一个偏不让。\n青龙洞里龙翻身，\n千年大梦变稻香！\n\n[Bridge]\n麻婆婆的狗，咬破麻叉口，\n麻线穿针眼，补丁也风流。\n左一句，右一句，\n舌头打结心自由！\n\n[Chorus]\n丁丁当当，乒乓乓乓，\n话赶话，舌绕梁，\n东边的钉，西边的墙，\n绕不完的弯，唱不完的慌！",
+    "audio_duration": 214.12,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.5,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.031190156936645508,
+        "diffusion": 20.130417823791504,
+        "latent2audio": 1.9650826454162598
+    },
+    "actual_seeds": [
+        1946426111
+    ],
+    "retake_seeds": [
+        331383387
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250512172941_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora_100k",
+    "task": "text2music",
+    "prompt": "东北话, spoken word, male, Tempo - Fast, Elements - Chorus Hook, Subgenre-Satirical Hip Hop, Rap, Chinese-language music, energetic, slightly nasal, Instrument - Live Bass Guitar, adult, Vocals - Syncopated Flow, Genre - Hip-Hop, rapping, bright",
+    "lyrics": "[verse]\n挣着 憋屈的 工资 还得 装乐呵\n猫着 怂样儿 还搁 朋友圈 嘚瑟\n扛着 傻逼的 指标 没人 搭把手\n这儿 不是 托儿所 少整 那出儿 哭唧尿嚎\n\n俺们 就像 一条条 老板的 裤衩子\n陪着 笑脸 接他 每一回 突突\n哎呦 老板 今儿个 穿我呗\n他 撅个腚 眼角 瞟你 那熊样\n\n[chorus]\n他们 骂我 打工仔 太多人 没睡醒\n寻思 抠搜 老板 一天天 穷折腾\n不想 俺的 人生 烂在 这嘎达\n不想 俺的 将来 折在 这破棚\n\n老子 不想 上班 老子 是外星人\n你都 把俺 骂急眼了 俺还 这么淡定\n现实 才是 梦 啥时候 能醒啊\n那 糟践人的 答案 在西北风 里飘\n\n[verse]\n瞅见 二愣子 同事 给老板 舔腚沟子\n瞅见 浪蹄子 女同事 在老板 胯骨轴 扭搭\n瞅见 白瞎的 光阴 耗在 没亮儿的 道儿\n瞅见 公交车上 一帮 僵尸 吐酸水\n\n瞅见 俺的 命 定在 苦逼的 坑里\n瞅见 俺的 爱情 被轮了 成了 老处女\n瞅见 好事儿 全归 高富帅\n还有 那些 臭不要脸 扭腚的 货色\n\n[chorus]（重复）\n他们 骂我 打工仔 太多人 没睡醒...\n\n[bridge]\n加班 没补助 俺认了\n欠薪 揍员工 把俺 当牲口\n去你妈 的小姘头\n\n[verse]\n破逼 管理制度 净整 娱乐八卦\n撸管式 管理 也就 你自己 嗨\n出点儿 屁事儿 就往 下属 脑瓜子 扣\n挣俩 钢镚儿 立马 牛逼 不分 公母\n\n你挖个 大坑 把俺们 往里 踹\n说这 叫梦想 你当年 多能耐\n俺们 就当 听传销 洗脑课\n可怜 连骗人 你都 就会 这一套\n\n[outro]\n老子 不想 上班\n老子 不想 上班\n老子 不想 上班",
+    "audio_duration": 135.92,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.5,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.06204533576965332,
+        "diffusion": 35.75483560562134,
+        "latent2audio": 1.5193355083465576
+    },
+    "actual_seeds": [
+        4176354214
+    ],
+    "retake_seeds": [
+        601086915
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250513044511_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora_100k",
+    "task": "text2music",
+    "prompt": "Rap, J-Pop, Anime, kawaii pop, EDM, Aggressive, Intense, Crisp Snare, Super Fast, Clear",
+    "lyrics": "[Intro]\nNya.\n\n[Verse]\n我 在 五 点 二 十 早 起，十 三 点 十 四 弹 会儿 琴\n习 惯 了 坐 班，习惯了 隔夜 的  剩 饭，\n习 惯 了 没有 你\n\n[Verse]\n怕 你 想 不 开，拦 在 你 的 面 前\n那 时 候 摔 得 差 点 住 院\n东 京 的 春 天 莺 莺 燕 燕\n我 说 想 不 想 来 跟 我 玩 音乐\n\n[Verse]\n带 着 我 的 朋 友 守 在 你 的 门 口\n弹 着 我 的 钢 琴 当 伴 奏\n等 你 放 学 后，陪 你   K   T   V\n端 着 我 的 红 茶 跟 你 碰 杯\n\n[Pre-Chorus]\n忽然间现实淹没了远方\n万家灯火，盖住月光\n奔走，忍受，变成了人偶\n别再对我伸出你的 双 手，会 受 伤\n\n[Chorus]\n明明都向前走，方向却渐渐不同\n时间让你我越走越近，却越来越陌生\n春 天 在 滂 沱 的 大 雨 里 飘 落\n得 了 心 太 高 脸 太 薄 的病\n\n[Bridge]\n我越难过，春日影越顶\n眼泪晃得我看不清\n埋葬了懦弱还有矫情\n却还是会在半夜摸眼睛\n\n青春期大部分时间在工 作\n用微笑换来余额几个零\n戴上了面具也明白了生活\n拼的是数字和脸更是命\n\n[Verse]\n我在五点二十早起，十三点十四弹会琴\n早上要做饭，回家时满地的瓶罐\n\n师 徒 二 人 站 在 我 的 面 前\n台 词 很 熟 练，照 着 就 念\n\n背 后 的 小 睦 扭 扭 捏 捏\n我 说 我 还 有 点 事 要 不 改 天 见\n\n然 后 你 的 双手 握 住 我 的 袖 口\n开 始 哭 着 求 我 不 要 走\n\n[Verse]\n我在下班后，忙活柴米油\n你和你的姐妹住着高楼\n\n苦 来 兮 苦，早 就 没 了\n现 实 扬 鞭，赶 着 我 向 前\n没有时间跟你分辨什么对与错\n\n[Bridge]\n没有什么对错，没有罪过\n谁不曾天真，是我太早看破\n生活一片狼藉，却又不想放弃\n一 边 聚 光 灯 下 绽 放，一 边 坠 落\n故作坚强，筑起心的墙\n越是委屈的伤口，越要藏\nLet it all out， it’s all right\n\n[Outro]\n俺 是 东 京 嘞，东 京 打 工 妹\n\n从虎之门带你转到浅草\n再从新宿转到竹桥\n\n俺 是 东 京 嘞，东 京 打 工 妹\n\n带 你 转 羽田 成田 蒲田 神田\n做 你 嘞 小 甜 甜！\n\n俺 是 东 京 嘞，东 京 打 工 妹\n带 你 转 赤 坂，带 你 转 霞 关\n恁 咋 不 早 说，今 天 不 管 饭\n",
+    "audio_duration": 147.62212,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.5,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.052134037017822266,
+        "diffusion": 17.909283876419067,
+        "latent2audio": 1.4904146194458008
+    },
+    "actual_seeds": [
+        2945962357
+    ],
+    "retake_seeds": [
+        2252292438
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0.7,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250513050200_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora_100k",
+    "task": "text2music",
+    "prompt": "Rap, adult, male, spoken word, rapping, clear, warm, articulate, Lo-Fi Hip Hop, 100-120 BPM, Keyboard Chords, Male Rap, Lazy Rhythm, Melancholy, Rap",
+    "lyrics": "[Intro]\n夜色 很 淡 像 褪色 的 照片  \n但 记忆 却 像 刀锋 一样 锐利  \n\n[Verse 1]\n你 说过 的 甜言蜜语 现在 听来 像 最 恶毒 的 咒骂  \n你 刺进 我 心里 的 刀 现在 还 在 滴血 未 干 哪  \n慵懒 的 旋律 像 我 的 脚步 拖着 沉重 的 躯壳  \n脑海 里 循环 播放 那 画面 快 把 我 逼疯 了  \n键盘 和弦 低沉 又 忧伤 弹奏 着 我 的 绝望  \n我 曾经 的 信任 像 玻璃 一样 被 你 狠狠 地 摔 在 地上  \n不想 振作 不想 原谅 只 想 让 这 一切 都 停止  \n可 心底 有 个 声音 嘶吼 着 要 你 付出 该 有 的 代价  \n\n[Chorus]\n背叛 像 毒药 渗透 我 的 血液  \n复仇 的 火焰 在 我 眼中 燃起  \n哪怕 遍体鳞伤 哪怕 万劫不复  \n我 也 要 亲手 撕碎 你 的 幸福  \n这 是 我 的 哀歌 也 是 我 的 战书  \n键盘 的 音符 每 一下 都 带着 恨意 和 痛苦  \n\n[Verse 2]\n曾经 的 兄弟 现在 面目全非 像 个 陌生人  \n你 的 自私 像 癌细胞 一点点 吞噬 我 的 纯真  \n我 学着 你 的 样子 把 心 锁 起来 不再 轻易 相信  \n让 懒散 的 节奏 包裹 我 给 自己 一点 喘息  \n键盘 的 音色 变得 更加 阴冷 像 秋天 的 雨滴  \n冲刷 掉 所有 温情 只 剩下 彻骨 的 寒意  \n我 不会 大喊大叫 只是 默默 地 计划  \n每 一步 都 走向 让 你 后悔 的 那 一 刹那  \n\n[Chorus]\n背叛 像 毒药 渗透 我 的 血液  \n复仇 的 火焰 在 我 眼中 燃起  \n哪怕 遍体鳞伤 哪怕 万劫不复  \n我 也 要 亲手 撕碎 你 的 幸福  \n这 是 我 的 哀歌 也 是 我 的 战书  \n键盘 的 音符 每 一下 都 带着 恨意 和 痛苦  \n\n[Bridge]\n也许 复仇 不能 带来 平静  \n也许 只 会 让 我 更 堕落  \n但 如果 不 这样 做  \n我 连 活下去 的 勇气 都 没有  \n\n[Outro]\n复仇 复仇 复仇  \n直到 最后 一刻  \n懒散 地 复仇 着  ",
+    "audio_duration": 202.64,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.65,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 0.036400794982910156,
+        "diffusion": 23.055809259414673,
+        "latent2audio": 1.8787360191345215
+    },
+    "actual_seeds": [
+        3900061002
+    ],
+    "retake_seeds": [
+        3037373819
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250513055451_0.wav"
+}
@@ -0,0 +1,45 @@
+{
+    "lora_name_or_path": "/root/sag_train/data/ace_step_v1_chinese_rap_lora_100k",
+    "task": "text2music",
+    "prompt": "Orchestra, Symphony, Sonata, Opera, Concerto, Rap, Beat, DJ, MC, StreetCulture",
+    "lyrics": "[verse1]\n羊皮卷轴 墨香飘 莫扎特 熬 安魂曲 通宵  \n和弦齿轮 咔哒转 比 瑞士 手表 更 精密 律动  \n八轨磁带 玩叠叠乐 披头士 炸 录音棚 天花板  \nAI 卷起 新风暴 像 灭霸 打响指 般 简单  \n\n[chorus]\n琴弦 到 代码 进化论 狂飙（skr）  \n象牙塔 被 鼠标 点爆 像 泡泡（boom）  \n灵感 加 算法 等于 王炸 大招  \n人类 心跳 才是 终极 混音 调料  \n\n[verse2]\n春之祭 召唤 百人 乐团 才够 燥  \n合成器 极客 玩电焊 焊出 赛博 神庙  \nDAW 解放 双手 钢琴卷帘 变 乐高  \n音色库 开挂 像 吃 金币 的 马里奥  \n\nAI 拆解 爵士乐 黑话 像 庖丁 解牛  \nCityPop 复古 滤镜 直接 参数 调油  \n神经网络 偷师 贝多芬 半夜 翻墙头  \n音乐 基因库 被 改写成 超频 万花筒  \n\n[chorus]  \n琴弦 到 代码 进化论 狂飙（skr）  \n象牙塔 被 鼠标 点爆 像 泡泡（boom）  \n灵感 加 算法 等于 王炸 大招  \n人类 心跳 才是 终极 混音 调料  \n\n[verse3]  \n电子琴 被 吐槽 塑料 味 超标  \n卧室 制作人 用 鼠标 单挑 整个 乐团 编制  \nAI 伴奏 刚上线 就被 键盘侠 集火  \n却 忘了 电吉他 曾被 说 是 魔鬼 的 副歌  \n\n现在 我 指尖 蹦迪 在 数据 炼丹炉  \n提示词 召唤 莫扎特 跨次元 碰杯 珍珠奶茶  \n当 比特 海洋 淹没 所有 物理 琴柱  \n最后 的 音轨 永远 连着 心脏 的 跳针  \n\n[bridge]  \n鹅毛笔 蘸着 银河 当 墨汁（绝了）  \n音浪 在 元宇宙 开 分店（疯了）  \n技术 迷雾 散成 像素 烟花  \n而 我们 始终 带着 老派 的 心跳 混搭  \n\n[chorus]  \n琴弦 到 代码 进化论 狂飙（skr）  \n象牙塔 被 鼠标 点爆 像 泡泡（boom）  \n灵感 加 算法 等于 王炸 大招  \n人类 心跳 才是 终极 混音 调料  \n\n[outro]  \n从 蒸汽 到 硅基 浪潮 我 冲浪（yo）  \n用 脑洞 接住 每个 技术 暴击（叮）  \n当 所有 设备 没电 的 凌晨 三点钟  \n最 原始 的 旋律 在 胸腔 敲击 成 龙卷风  ",
+    "audio_duration": 172.64,
+    "infer_step": 60,
+    "guidance_scale": 15,
+    "scheduler_type": "euler",
+    "cfg_type": "apg",
+    "omega_scale": 10,
+    "guidance_interval": 0.65,
+    "guidance_interval_decay": 0,
+    "min_guidance_scale": 3,
+    "use_erg_tag": true,
+    "use_erg_lyric": false,
+    "use_erg_diffusion": true,
+    "oss_steps": [],
+    "timecosts": {
+        "preprocess": 3.648996353149414,
+        "diffusion": 16.44967818260193,
+        "latent2audio": 1.614703893661499
+    },
+    "actual_seeds": [
+        1198023141
+    ],
+    "retake_seeds": [
+        3389016134
+    ],
+    "retake_variance": 0.5,
+    "guidance_scale_text": 0,
+    "guidance_scale_lyric": 0,
+    "repaint_start": 0,
+    "repaint_end": 0,
+    "edit_n_min": 0.0,
+    "edit_n_max": 1.0,
+    "edit_n_avg": 1,
+    "src_audio_path": null,
+    "edit_target_prompt": null,
+    "edit_target_lyrics": null,
+    "audio2audio_enable": false,
+    "ref_audio_strength": 0.5,
+    "ref_audio_input": null,
+    "audio_path": "./outputs/output_20250513060150_0.wav"
+}
@@ -0,0 +1,106 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List, Optional
+import os
+from acestep.pipeline_ace_step import ACEStepPipeline
+from acestep.data_sampler import DataSampler
+import uuid
+
+app = FastAPI(title="ACEStep Pipeline API")
+
+class ACEStepInput(BaseModel):
+    checkpoint_path: str
+    bf16: bool = True
+    torch_compile: bool = False
+    device_id: int = 0
+    output_path: Optional[str] = None
+    audio_duration: float
+    prompt: str
+    lyrics: str
+    infer_step: int
+    guidance_scale: float
+    scheduler_type: str
+    cfg_type: str
+    omega_scale: float
+    actual_seeds: List[int]
+    guidance_interval: float
+    guidance_interval_decay: float
+    min_guidance_scale: float
+    use_erg_tag: bool
+    use_erg_lyric: bool
+    use_erg_diffusion: bool
+    oss_steps: List[int]
+    guidance_scale_text: float = 0.0
+    guidance_scale_lyric: float = 0.0
+
+class ACEStepOutput(BaseModel):
+    status: str
+    output_path: Optional[str]
+    message: str
+
+def initialize_pipeline(checkpoint_path: str, bf16: bool, torch_compile: bool, device_id: int) -> ACEStepPipeline:
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
+    return ACEStepPipeline(
+        checkpoint_dir=checkpoint_path,
+        dtype="bfloat16" if bf16 else "float32",
+        torch_compile=torch_compile,
+    )
+
+@app.post("/generate", response_model=ACEStepOutput)
+async def generate_audio(input_data: ACEStepInput):
+    try:
+        # Initialize pipeline
+        model_demo = initialize_pipeline(
+            input_data.checkpoint_path,
+            input_data.bf16,
+            input_data.torch_compile,
+            input_data.device_id
+        )
+
+        # Prepare parameters
+        params = (
+            input_data.audio_duration,
+            input_data.prompt,
+            input_data.lyrics,
+            input_data.infer_step,
+            input_data.guidance_scale,
+            input_data.scheduler_type,
+            input_data.cfg_type,
+            input_data.omega_scale,
+            ", ".join(map(str, input_data.actual_seeds)),
+            input_data.guidance_interval,
+            input_data.guidance_interval_decay,
+            input_data.min_guidance_scale,
+            input_data.use_erg_tag,
+            input_data.use_erg_lyric,
+            input_data.use_erg_diffusion,
+            ", ".join(map(str, input_data.oss_steps)),
+            input_data.guidance_scale_text,
+            input_data.guidance_scale_lyric,
+        )
+
+        # Generate output path if not provided
+        output_path = input_data.output_path or f"output_{uuid.uuid4().hex}.wav"
+
+        # Run pipeline
+        model_demo(
+            *params,
+            save_path=output_path
+        )
+
+        return ACEStepOutput(
+            status="success",
+            output_path=output_path,
+            message="Audio generated successfully"
+        )
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error generating audio: {str(e)}")
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)
@@ -40,64 +40,157 @@ def sample_data(json_data):
@click.option(
    "--torch_compile", type=bool, default=False, help="Whether to use torch compile"
 )
+@click.option(
+    "--cpu_offload", type=bool, default=False, help="Whether to use CPU offloading (only load current stage's model to GPU)"
+)
+@click.option(
+    "--overlapped_decode", type=bool, default=False, help="Whether to use overlapped decoding (run dcae and vocoder using sliding windows)"
+)
@click.option("--device_id", type=int, default=0, help="Device ID to use")
@click.option("--output_path", type=str, default=None, help="Path to save the output")
-def main(checkpoint_path, bf16, torch_compile, device_id, output_path):
+# --- generation parameters -------------------------------------------------
+# Without --prompt the script keeps its original behaviour and generates from a
+# random example in examples/input_params.
+@click.option(
+    "--prompt",
+    type=str,
+    default=None,
+    help="Style tags or description, comma separated. If omitted, a random example is used.",
+)
+@click.option("--lyrics", type=str, default=None, help="Lyrics, with [verse]/[chorus] tags.")
+@click.option(
+    "--lyrics_file",
+    type=click.Path(exists=True, dir_okay=False),
+    default=None,
+    help="Read lyrics from a UTF-8 file. Takes precedence over --lyrics.",
+)
+@click.option("--duration", type=float, default=60.0, help="Audio duration in seconds (-1 for random).")
+@click.option("--steps", type=int, default=60, help="Number of inference steps.")
+@click.option("--guidance_scale", type=float, default=15.0, help="Guidance scale.")
+@click.option(
+    "--scheduler",
+    type=click.Choice(["euler", "heun", "pingpong"]),
+    default="euler",
+    help="Scheduler type.",
+)
+@click.option(
+    "--cfg_type",
+    type=click.Choice(["cfg", "apg", "cfg_star"]),
+    default="apg",
+    help="CFG type. apg is recommended.",
+)
+@click.option("--omega_scale", type=float, default=10.0, help="Granularity scale.")
+@click.option("--seed", type=str, default=None, help="Manual seeds, comma separated. Random if omitted.")
+@click.option(
+    "--format",
+    "audio_format",
+    type=click.Choice(["wav", "mp3", "ogg", "flac"]),
+    default="wav",
+    help="Output audio format.",
+)
+def main(
+    checkpoint_path,
+    bf16,
+    torch_compile,
+    cpu_offload,
+    overlapped_decode,
+    device_id,
+    output_path,
+    prompt,
+    lyrics,
+    lyrics_file,
+    duration,
+    steps,
+    guidance_scale,
+    scheduler,
+    cfg_type,
+    omega_scale,
+    seed,
+    audio_format,
+):
    os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)

    model_demo = ACEStepPipeline(
        checkpoint_dir=checkpoint_path,
        dtype="bfloat16" if bf16 else "float32",
        torch_compile=torch_compile,
+        cpu_offload=cpu_offload,
+        overlapped_decode=overlapped_decode
    )
-    print(model_demo)

-    data_sampler = DataSampler()
+    if lyrics_file is not None:
+        with open(lyrics_file, "r", encoding="utf-8") as f:
+            lyrics = f.read()

-    json_data = data_sampler.sample()
-    json_data = sample_data(json_data)
-    print(json_data)
+    if prompt is None:
+        # No prompt given: keep the original behaviour and use a random example.
+        data_sampler = DataSampler()
+        json_data = data_sampler.sample()
+        (
+            audio_duration,
+            prompt,
+            sampled_lyrics,
+            infer_step,
+            sampled_guidance_scale,
+            scheduler_type,
+            sampled_cfg_type,
+            sampled_omega_scale,
+            manual_seeds,
+            guidance_interval,
+            guidance_interval_decay,
+            min_guidance_scale,
+            use_erg_tag,
+            use_erg_lyric,
+            use_erg_diffusion,
+            oss_steps,
+            guidance_scale_text,
+            guidance_scale_lyric,
+        ) = sample_data(json_data)
+        # Explicit options still win over the sampled example.
+        lyrics = lyrics if lyrics is not None else sampled_lyrics
+        manual_seeds = seed if seed is not None else manual_seeds
+    else:
+        audio_duration = duration
+        lyrics = lyrics if lyrics is not None else "[instrumental]"
+        infer_step = steps
+        sampled_guidance_scale = guidance_scale
+        scheduler_type = scheduler
+        sampled_cfg_type = cfg_type
+        sampled_omega_scale = omega_scale
+        manual_seeds = seed
+        guidance_interval = 0.5
+        guidance_interval_decay = 0.0
+        min_guidance_scale = 3.0
+        use_erg_tag = True
+        use_erg_lyric = True
+        use_erg_diffusion = True
+        oss_steps = None
+        guidance_scale_text = 0.0
+        guidance_scale_lyric = 0.0

-    (
-        audio_duration,
-        prompt,
-        lyrics,
-        infer_step,
-        guidance_scale,
-        scheduler_type,
-        cfg_type,
-        omega_scale,
-        manual_seeds,
-        guidance_interval,
-        guidance_interval_decay,
-        min_guidance_scale,
-        use_erg_tag,
-        use_erg_lyric,
-        use_erg_diffusion,
-        oss_steps,
-        guidance_scale_text,
-        guidance_scale_lyric,
-    ) = json_data
+    click.echo(f"prompt: {prompt}")
+    click.echo(f"duration: {audio_duration}s, steps: {infer_step}, seeds: {manual_seeds}")

    model_demo(
-        audio_duration,
-        prompt,
-        lyrics,
-        infer_step,
-        guidance_scale,
-        scheduler_type,
-        cfg_type,
-        omega_scale,
-        manual_seeds,
-        guidance_interval,
-        guidance_interval_decay,
-        min_guidance_scale,
-        use_erg_tag,
-        use_erg_lyric,
-        use_erg_diffusion,
-        oss_steps,
-        guidance_scale_text,
-        guidance_scale_lyric,
+        format=audio_format,
+        audio_duration=audio_duration,
+        prompt=prompt,
+        lyrics=lyrics,
+        infer_step=infer_step,
+        guidance_scale=sampled_guidance_scale,
+        scheduler_type=scheduler_type,
+        cfg_type=sampled_cfg_type,
+        omega_scale=sampled_omega_scale,
+        manual_seeds=manual_seeds,
+        guidance_interval=guidance_interval,
+        guidance_interval_decay=guidance_interval_decay,
+        min_guidance_scale=min_guidance_scale,
+        use_erg_tag=use_erg_tag,
+        use_erg_lyric=use_erg_lyric,
+        use_erg_diffusion=use_erg_diffusion,
+        oss_steps=oss_steps,
+        guidance_scale_text=guidance_scale_text,
+        guidance_scale_lyric=guidance_scale_lyric,
        save_path=output_path,
    )

@@ -24,7 +24,7 @@
        "colab_type": "text"
      },
      "source": [
-        "<a href=\"https://colab.research.google.com/github/TheNeodev/ACE-Step/blob/main/inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+        "<a href=\"https://colab.research.google.com/github/ace-step/ACE-Step/blob/main/inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
@@ -81,8 +81,9 @@
      },
      "outputs": [],
      "source": [
-        "#@title Install and Download\n",
+        "#@title 🎯 Install and Download\n",
        "\n",
+        "from IPython.display import clear_output"
        "\n",
        "import codecs\n",
        "\n",
@@ -101,7 +102,7 @@
        "repopath = codecs.decode('erdhverzragf.gkg', 'rot_13')\n",
        "\n",
        "\n",
-        "!git clone https://github.com/usamireko/ACE-Step\n",
+        "!git clone https://github.com/ace-step/ACE-Step\n",
        "%cd /content/ACE-Step\n",
        "\n",
        "\n",
@@ -110,20 +111,26 @@
        "!uv pip install huggingface-hub numpy==1.26.0 > /dev/null 2>&1\n",
        "!huggingface-cli download ACE-Step/ACE-Step-v1-3.5B --local-dir /content/ACE-Step/checkpoints --local-dir-use-symlinks False\n",
        "\n",
+        "!pip install e .",
+        "\n",
        "\n",
        "import os\n",
        "os.environ['MPLBACKEND'] = 'agg'"
+        "\n\n"
+        "clear_output()"
+        "\n\n"
+        "print("✅ Installation Complete!")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
-        "#@title Run Gradio UI\n",
+        "#@title 🚀 Run Gradio UI\n",
        "bf16 = True # @param {\"type\":\"boolean\"}\n",
        "\n",
+        "print(\" 🚀 Running UI...\")\n",
        "\n",
-        "print(\" * Running UI...\")\n",
-        "!python app.py --checkpoint_path ./checkpoints/ --port 7865 --device_id 0 --share true --bf16 {bf16}"
+        "!acestep --checkpoint_path ./checkpoints/ --port 7865 --device_id 0 --share true --bf16 {bf16}"
      ],
      "metadata": {
        "cellView": "form",
@@ -0,0 +1,7 @@
+# Training-only dependencies (trainer.py, convert2hf_dataset.py).
+# Install with: pip install -e ".[train]"
+datasets==3.4.1
+pytorch_lightning==2.5.1
+matplotlib==3.10.1
+tensorboard
+tensorboardX
@@ -1,18 +1,15 @@
-datasets==3.4.1
-diffusers==0.32.2
-gradio==5.23.3
+diffusers>=0.33.0
+gradio>=6.0.0
 librosa==0.11.0
 loguru==0.7.3
-matplotlib==3.10.1
 numpy
 pypinyin==0.53.0
-pytorch_lightning==2.5.1
 soundfile==0.13.1
 torch
 torchaudio
 torchvision
-tqdm==4.67.1
-transformers==4.50.0
+tqdm
+transformers>=4.57.0
 py3langid==0.3.0
 hangul-romanize==0.1.0
 num2words==0.5.14
@@ -21,3 +18,4 @@ accelerate==1.6.0
 cutlet
 fugashi[unidic-lite]
 click
+peft
@@ -1,13 +1,24 @@
-from setuptools import setup
+from setuptools import setup, find_namespace_packages
+
+
+def read_requirements(path):
+    """Read a requirements file, skipping comments and blank lines."""
+    with open(path, encoding="utf-8") as f:
+        return [
+            line.strip()
+            for line in f
+            if line.strip() and not line.lstrip().startswith("#")
+        ]
+

 setup(
    name="ace_step",
    description="ACE Step: A Step Towards Music Generation Foundation Model",
    long_description=open("README.md", encoding="utf-8").read(),
    long_description_content_type="text/markdown",
-    version="0.1.0",
-    packages=["acestep"],
-    install_requires=open("requirements.txt", encoding="utf-8").read().splitlines(),
+    version="0.2.0",
+    packages=find_namespace_packages(),
+    install_requires=read_requirements("requirements.txt"),
    author="ACE Studio, StepFun AI",
    license="Apache 2.0",
    classifiers=[
@@ -20,4 +31,12 @@ setup(
            "acestep=acestep.gui:main",
        ],
    },
+    include_package_data=True,  # Ensure this is set to True
+    package_data={
+        "acestep.models.lyrics_utils": ["vocab.json"],  # Specify the relative path to vocab.json
+    },
+    extras_require={
+        # Only needed to train or fine-tune; inference does not import these.
+        "train": read_requirements("requirements-train.txt"),
+    },
 )
@@ -0,0 +1,266 @@
+@echo off
+chcp 65001 >nul 2>&1
+setlocal
+title ACE-Step
+pushd "%~dp0"
+
+REM ============================ НАСТРОЙКИ ============================
+REM Каталог виртуального окружения
+set "VENV_DIR=venv"
+REM Порт веб-интерфейса
+set "PORT=7865"
+REM Адрес, на котором слушает сервер (0.0.0.0 - доступ из локальной сети)
+set "SERVER_NAME=127.0.0.1"
+REM Номер видеокарты
+set "DEVICE_ID=0"
+REM Путь к весам модели. Пусто - скачать в ~/.cache/ace-step/checkpoints
+set "CHECKPOINT_PATH="
+REM Сборка PyTorch (cu126 / cu124 / cu121 - под вашу версию CUDA)
+set "TORCH_INDEX=https://download.pytorch.org/whl/cu126"
+REM ==================================================================
+
+REM Имя скрипта надо запомнить до shift: shift сдвигает и %0
+set "SCRIPT_NAME=%~nx0"
+
+set "OPT_LOWVRAM=0"
+set "OPT_CPU=0"
+set "OPT_SHARE=0"
+set "DO_SETUP_ONLY=0"
+set "DO_REINSTALL=0"
+set "DO_UPDATE=0"
+
+set "FROM_EXPLORER=0"
+echo %cmdcmdline% | find /i "%~nx0" >nul 2>&1 && set "FROM_EXPLORER=1"
+
+REM ------------------------- разбор аргументов -------------------------
+:parse_args
+if "%~1"=="" goto args_done
+if /i "%~1"=="--help"      goto usage
+if /i "%~1"=="-h"          goto usage
+if /i "%~1"=="/?"          goto usage
+if /i "%~1"=="--lowvram"   ( set "OPT_LOWVRAM=1"  & shift & goto parse_args )
+if /i "%~1"=="--cpu"       ( set "OPT_CPU=1"      & shift & goto parse_args )
+if /i "%~1"=="--share"     ( set "OPT_SHARE=1"    & shift & goto parse_args )
+if /i "%~1"=="--setup"     ( set "DO_SETUP_ONLY=1" & shift & goto parse_args )
+if /i "%~1"=="--reinstall" ( set "DO_REINSTALL=1" & shift & goto parse_args )
+if /i "%~1"=="--update"    ( set "DO_UPDATE=1"    & shift & goto parse_args )
+if /i "%~1"=="--listen"    ( set "SERVER_NAME=0.0.0.0" & shift & goto parse_args )
+if /i "%~1"=="--port"      ( set "PORT=%~2"       & shift & shift & goto parse_args )
+if /i "%~1"=="--device"    ( set "DEVICE_ID=%~2"  & shift & shift & goto parse_args )
+echo [ОШИБКА] Неизвестный аргумент: %~1
+echo Запустите "%SCRIPT_NAME% --help" для справки.
+goto fail
+:args_done
+
+echo.
+echo ==========================================
+echo   ACE-Step - генерация музыки
+echo ==========================================
+echo.
+
+set "VENV_PY=%CD%\%VENV_DIR%\Scripts\python.exe"
+
+if "%DO_REINSTALL%"=="1" (
+    if exist "%VENV_DIR%\" (
+        echo [1/4] Удаляю старое окружение "%VENV_DIR%"...
+        rmdir /s /q "%VENV_DIR%"
+    )
+)
+
+REM ------------------------- поиск Python -------------------------
+REM Рабочим считаем окружение, в котором есть и python, и pip: каталог от
+REM прерванной установки выглядит как готовый, но валится дальше с невнятной
+REM ошибкой про интернет.
+if not exist "%VENV_PY%" goto venv_create
+"%VENV_PY%" -m pip --version >nul 2>&1 && goto venv_ready
+echo [1/4] Окружение "%VENV_DIR%" неработоспособно, пересоздаю...
+rmdir /s /q "%VENV_DIR%"
+
+:venv_create
+echo [1/4] Ищу подходящий Python...
+set "SYS_PY="
+for %%V in (3.12 3.11 3.10 3.13) do (
+    if not defined SYS_PY (
+        py -%%V -c "import sys" >nul 2>&1 && set "SYS_PY=py -%%V"
+    )
+)
+if not defined SYS_PY (
+    python -c "import sys; sys.exit(0 if sys.version_info >= (3,10) else 1)" >nul 2>&1 && set "SYS_PY=python"
+)
+if not defined SYS_PY (
+    echo.
+    echo [ОШИБКА] Не найден Python 3.10 или новее.
+    echo Установите его с https://www.python.org/downloads/
+    echo и обязательно отметьте галочку "Add Python to PATH".
+    goto fail
+)
+for /f "delims=" %%O in ('%SYS_PY% -c "import sys;print(sys.version.split()[0])" 2^>nul') do set "PYVER=%%O"
+echo       Найден Python %PYVER% (%SYS_PY%)
+
+echo [1/4] Создаю виртуальное окружение в "%VENV_DIR%"...
+%SYS_PY% -m venv "%VENV_DIR%"
+if errorlevel 1 goto venv_failed
+if not exist "%VENV_PY%" goto venv_failed
+"%VENV_PY%" -m pip --version >nul 2>&1 || goto venv_failed
+set "FRESH_VENV=1"
+
+:venv_ready
+if not defined FRESH_VENV echo [1/4] Виртуальное окружение найдено: %VENV_DIR%
+
+REM ------------------------- зависимости -------------------------
+echo [2/4] Проверяю зависимости...
+
+set "NEED_TORCH=0"
+"%VENV_PY%" -c "import importlib.util,sys; sys.exit(0 if importlib.util.find_spec('torch') else 1)" >nul 2>&1 || set "NEED_TORCH=1"
+REM Проверяем именно установку пакета: сам каталог acestep лежит рядом со скриптом,
+REM поэтому find_spec('acestep') сработал бы даже без установленных зависимостей.
+set "NEED_ACESTEP=0"
+"%VENV_PY%" -c "import importlib.metadata as md, importlib.util as u, sys; md.version('ace_step'); sys.exit(0 if u.find_spec('click') and u.find_spec('gradio') else 1)" >nul 2>&1 || set "NEED_ACESTEP=1"
+
+if "%DO_UPDATE%"=="1" set "NEED_TORCH=1"
+if "%DO_UPDATE%"=="1" set "NEED_ACESTEP=1"
+
+if "%NEED_TORCH%%NEED_ACESTEP%"=="00" goto deps_ready
+
+echo       Обновляю pip...
+"%VENV_PY%" -m pip install --upgrade pip setuptools wheel --quiet
+if errorlevel 1 (
+    echo [ОШИБКА] Не удалось обновить pip. Проверьте подключение к интернету.
+    goto fail
+)
+
+if "%NEED_TORCH%"=="1" (
+    if "%OPT_CPU%"=="1" (
+        echo       Устанавливаю PyTorch для CPU. Это займёт несколько минут...
+        "%VENV_PY%" -m pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+    ) else (
+        echo       Устанавливаю PyTorch с поддержкой CUDA. Это займёт несколько минут...
+        "%VENV_PY%" -m pip install --upgrade torch torchvision torchaudio --index-url %TORCH_INDEX%
+    )
+    if errorlevel 1 (
+        echo [ОШИБКА] Не удалось установить PyTorch.
+        echo Если у вас другая версия CUDA, поменяйте TORCH_INDEX в начале этого файла.
+        goto fail
+    )
+)
+
+if "%NEED_ACESTEP%"=="1" (
+    echo       Устанавливаю ACE-Step и зависимости...
+    "%VENV_PY%" -m pip install -e .
+    if errorlevel 1 (
+        echo [ОШИБКА] Не удалось установить ACE-Step.
+        goto fail
+    )
+)
+
+:deps_ready
+
+if "%OPT_LOWVRAM%"=="1" (
+    "%VENV_PY%" -c "import importlib.util,sys; sys.exit(0 if importlib.util.find_spec('triton') else 1)" >nul 2>&1 || (
+        echo       Режим экономии VRAM: устанавливаю triton-windows...
+        "%VENV_PY%" -m pip install triton-windows --quiet
+        if errorlevel 1 echo [ВНИМАНИЕ] triton-windows не установился, --torch_compile может не заработать.
+    )
+)
+
+REM ------------------------- проверка GPU -------------------------
+echo [3/4] Проверяю видеокарту...
+if "%OPT_CPU%"=="1" goto gpu_cpu_mode
+
+set "BF16=true"
+set "GPUINFO="
+set "GPUTMP=%TEMP%\acestep_gpu_%RANDOM%.txt"
+"%VENV_PY%" -c "import torch; print((torch.cuda.get_device_name(0) + ' - ' + str(round(torch.cuda.get_device_properties(0).total_memory/1073741824, 1)) + ' GB VRAM') if torch.cuda.is_available() else 'NO_CUDA')" > "%GPUTMP%" 2>nul
+if exist "%GPUTMP%" set /p GPUINFO=<"%GPUTMP%"
+del "%GPUTMP%" >nul 2>&1
+if not defined GPUINFO set "GPUINFO=NO_CUDA"
+if not "%GPUINFO%"=="NO_CUDA" goto gpu_ok
+echo.
+echo [ВНИМАНИЕ] CUDA недоступна - модель будет работать на процессоре ^(очень медленно^).
+echo Если у вас есть видеокарта NVIDIA, обновите драйвер и переустановите PyTorch:
+echo     "%SCRIPT_NAME%" --update
+echo.
+goto gpu_done
+
+:gpu_ok
+echo       %GPUINFO%
+goto gpu_done
+
+:gpu_cpu_mode
+echo       Принудительный режим CPU. Генерация будет очень медленной.
+set "DEVICE_ID=-1"
+set "BF16=false"
+
+:gpu_done
+
+if "%DO_SETUP_ONLY%"=="1" (
+    echo.
+    echo Установка завершена. Для запуска выполните "%SCRIPT_NAME%".
+    goto success
+)
+
+REM ------------------------- запуск -------------------------
+set "ARGS=--port %PORT% --server_name %SERVER_NAME% --device_id=%DEVICE_ID% --bf16 %BF16%"
+if not "%CHECKPOINT_PATH%"=="" set ARGS=%ARGS% --checkpoint_path "%CHECKPOINT_PATH%"
+if "%OPT_SHARE%"=="1"   set "ARGS=%ARGS% --share true"
+if "%OPT_LOWVRAM%"=="1" set "ARGS=%ARGS% --cpu_offload true --overlapped_decode true --torch_compile true"
+
+echo [4/4] Запускаю веб-интерфейс...
+echo.
+if "%OPT_LOWVRAM%"=="1" echo       Режим: экономия видеопамяти
+echo       Идёт загрузка моделей, это занимает время.
+echo       При первом запуске дополнительно скачиваются веса (~8 ГБ).
+echo.
+echo       Интерфейс будет доступен на http://%SERVER_NAME%:%PORT%
+echo       когда ниже появится строка "Running on local URL".
+echo       Остановить: Ctrl+C в этом окне.
+echo.
+
+"%VENV_PY%" -m acestep.gui %ARGS%
+if errorlevel 1 (
+    echo.
+    echo [ОШИБКА] ACE-Step завершился с ошибкой. Текст ошибки - выше.
+    goto fail
+)
+goto success
+
+REM ------------------------- справка -------------------------
+:usage
+echo.
+echo Использование: %SCRIPT_NAME% [флаги]
+echo.
+echo   --lowvram      Режим экономии видеопамяти (~8 ГБ VRAM)
+echo   --cpu          Запуск на процессоре, без CUDA (очень медленно)
+echo   --share        Публичная ссылка Gradio
+echo   --port ^<N^>     Порт веб-интерфейса (по умолчанию %PORT%)
+echo   --device ^<N^>   Номер видеокарты (по умолчанию %DEVICE_ID%)
+echo   --listen       Слушать 0.0.0.0 (доступ из локальной сети)
+echo   --reinstall    Пересоздать виртуальное окружение с нуля
+echo   --update       Обновить зависимости
+echo   --setup        Только установка, без запуска
+echo   --help         Эта справка
+echo.
+echo Настройки по умолчанию - в блоке НАСТРОЙКИ в начале файла.
+echo.
+goto success
+
+REM ------------------------- завершение -------------------------
+:venv_failed
+echo.
+echo [ОШИБКА] Не удалось создать рабочее виртуальное окружение в "%VENV_DIR%".
+echo Переустановите Python с https://www.python.org/downloads/,
+echo отметив "Add Python to PATH", и запустите "%SCRIPT_NAME% --reinstall".
+goto fail
+
+:fail
+echo.
+pause
+popd
+endlocal
+exit /b 1
+
+:success
+if "%FROM_EXPLORER%"=="1" pause
+popd
+endlocal
+exit /b 0
@@ -0,0 +1,233 @@
+#!/usr/bin/env bash
+# ACE-Step: настройка окружения и запуск веб-интерфейса на Linux / macOS.
+# Windows-аналог — start.bat.
+set -uo pipefail
+
+cd "$(dirname "$0")"
+
+# ============================ НАСТРОЙКИ ============================
+VENV_DIR="${VENV_DIR:-venv}"          # каталог виртуального окружения
+PORT="${PORT:-7865}"                  # порт веб-интерфейса
+SERVER_NAME="${SERVER_NAME:-127.0.0.1}"  # 0.0.0.0 — доступ из локальной сети
+DEVICE_ID="${DEVICE_ID:-0}"           # номер видеокарты
+CHECKPOINT_PATH="${CHECKPOINT_PATH:-}"   # пусто — скачать в ~/.cache/ace-step
+# Сборка PyTorch под CUDA (cu126 / cu124 / cu121). Только для Linux.
+TORCH_INDEX="${TORCH_INDEX:-https://download.pytorch.org/whl/cu126}"
+# ==================================================================
+
+SCRIPT_NAME="$(basename "$0")"
+OPT_LOWVRAM=0
+OPT_CPU=0
+OPT_SHARE=0
+DO_SETUP_ONLY=0
+DO_REINSTALL=0
+DO_UPDATE=0
+
+usage() {
+    cat <<EOF
+
+Использование: ./$SCRIPT_NAME [флаги]
+
+  --lowvram      Режим экономии видеопамяти (~8 ГБ VRAM)
+  --cpu          Запуск на процессоре, без CUDA (очень медленно)
+  --share        Публичная ссылка Gradio
+  --port <N>     Порт веб-интерфейса (по умолчанию $PORT)
+  --device <N>   Номер видеокарты (по умолчанию $DEVICE_ID)
+  --listen       Слушать 0.0.0.0 (доступ из локальной сети)
+  --reinstall    Пересоздать виртуальное окружение с нуля
+  --update       Обновить зависимости
+  --setup        Только установка, без запуска
+  --help         Эта справка
+
+Настройки по умолчанию — в блоке НАСТРОЙКИ в начале файла,
+их можно переопределить переменными окружения (PORT=7870 ./$SCRIPT_NAME).
+
+EOF
+}
+
+die() {
+    echo
+    echo "[ОШИБКА] $*" >&2
+    exit 1
+}
+
+while [ $# -gt 0 ]; do
+    case "$1" in
+        --help|-h)   usage; exit 0 ;;
+        --lowvram)   OPT_LOWVRAM=1; shift ;;
+        --cpu)       OPT_CPU=1; shift ;;
+        --share)     OPT_SHARE=1; shift ;;
+        --setup)     DO_SETUP_ONLY=1; shift ;;
+        --reinstall) DO_REINSTALL=1; shift ;;
+        --update)    DO_UPDATE=1; shift ;;
+        --listen)    SERVER_NAME="0.0.0.0"; shift ;;
+        --port)      PORT="${2:?--port требует значение}"; shift 2 ;;
+        --device)    DEVICE_ID="${2:?--device требует значение}"; shift 2 ;;
+        *)
+            echo "[ОШИБКА] Неизвестный аргумент: $1" >&2
+            echo "Запустите ./$SCRIPT_NAME --help для справки." >&2
+            exit 1 ;;
+    esac
+done
+
+echo
+echo "=========================================="
+echo "  ACE-Step - генерация музыки"
+echo "=========================================="
+echo
+
+case "$(uname -s)" in
+    Darwin) IS_MAC=1 ;;
+    *)      IS_MAC=0 ;;
+esac
+
+VENV_PY="$PWD/$VENV_DIR/bin/python"
+
+# Рабочим считаем окружение, в котором есть и python, и pip: каталог от
+# прерванной установки (или venv без ensurepip) выглядит как готовый, но валится
+# дальше с невнятной ошибкой.
+venv_ok() {
+    [ -x "$VENV_PY" ] && "$VENV_PY" -m pip --version >/dev/null 2>&1
+}
+
+venv_failed() {
+    echo >&2
+    echo "[ОШИБКА] Не удалось создать рабочее виртуальное окружение в \"$VENV_DIR\"." >&2
+    if [ "$IS_MAC" = 0 ] && command -v apt >/dev/null 2>&1; then
+        echo "На Debian/Ubuntu для этого нужен отдельный пакет:" >&2
+        echo "    sudo apt install python3-venv" >&2
+    fi
+    exit 1
+}
+
+if [ "$DO_REINSTALL" = 1 ] && [ -d "$VENV_DIR" ]; then
+    echo "[1/4] Удаляю старое окружение \"$VENV_DIR\"..."
+    rm -rf "$VENV_DIR"
+fi
+
+# ------------------------- поиск Python -------------------------
+if venv_ok; then
+    echo "[1/4] Виртуальное окружение найдено: $VENV_DIR"
+else
+    if [ -d "$VENV_DIR" ]; then
+        echo "[1/4] Окружение \"$VENV_DIR\" неработоспособно, пересоздаю..."
+        rm -rf "$VENV_DIR"
+    fi
+    echo "[1/4] Ищу подходящий Python..."
+    SYS_PY=""
+    for candidate in python3.12 python3.11 python3.10 python3.13 python3; do
+        if command -v "$candidate" >/dev/null 2>&1 &&
+           "$candidate" -c 'import sys; sys.exit(0 if sys.version_info >= (3,10) else 1)' 2>/dev/null; then
+            SYS_PY="$candidate"
+            break
+        fi
+    done
+    [ -n "$SYS_PY" ] || die "Не найден Python 3.10 или новее. Установите его и повторите."
+
+    echo "      Найден Python $("$SYS_PY" -c 'import sys; print(sys.version.split()[0])') ($SYS_PY)"
+    echo "[1/4] Создаю виртуальное окружение в \"$VENV_DIR\"..."
+    "$SYS_PY" -m venv "$VENV_DIR" || venv_failed
+    venv_ok || venv_failed
+fi
+
+# ------------------------- зависимости -------------------------
+echo "[2/4] Проверяю зависимости..."
+
+NEED_TORCH=0
+"$VENV_PY" -c "import importlib.util,sys; sys.exit(0 if importlib.util.find_spec('torch') else 1)" 2>/dev/null || NEED_TORCH=1
+
+# Проверяем именно установку пакета: каталог acestep лежит рядом со скриптом,
+# поэтому find_spec('acestep') сработал бы и без установленных зависимостей.
+NEED_ACESTEP=0
+"$VENV_PY" -c "import importlib.metadata as md, importlib.util as u, sys; md.version('ace_step'); sys.exit(0 if u.find_spec('click') and u.find_spec('gradio') else 1)" 2>/dev/null || NEED_ACESTEP=1
+
+if [ "$DO_UPDATE" = 1 ]; then
+    NEED_TORCH=1
+    NEED_ACESTEP=1
+fi
+
+if [ "$NEED_TORCH" = 1 ] || [ "$NEED_ACESTEP" = 1 ]; then
+    echo "      Обновляю pip..."
+    "$VENV_PY" -m pip install --upgrade pip setuptools wheel --quiet ||
+        die "Не удалось обновить pip. Проверьте подключение к интернету."
+fi
+
+if [ "$NEED_TORCH" = 1 ]; then
+    if [ "$IS_MAC" = 1 ]; then
+        # На macOS колёса с PyPI уже собраны с поддержкой MPS.
+        echo "      Устанавливаю PyTorch (macOS/MPS). Это займёт несколько минут..."
+        "$VENV_PY" -m pip install --upgrade torch torchvision torchaudio
+    elif [ "$OPT_CPU" = 1 ]; then
+        echo "      Устанавливаю PyTorch для CPU. Это займёт несколько минут..."
+        "$VENV_PY" -m pip install --upgrade torch torchvision torchaudio \
+            --index-url https://download.pytorch.org/whl/cpu
+    else
+        echo "      Устанавливаю PyTorch с поддержкой CUDA. Это займёт несколько минут..."
+        "$VENV_PY" -m pip install --upgrade torch torchvision torchaudio --index-url "$TORCH_INDEX"
+    fi || die "Не удалось установить PyTorch. Если у вас другая версия CUDA, поменяйте TORCH_INDEX."
+fi
+
+if [ "$NEED_ACESTEP" = 1 ]; then
+    echo "      Устанавливаю ACE-Step и зависимости..."
+    "$VENV_PY" -m pip install -e . || die "Не удалось установить ACE-Step."
+fi
+
+if [ "$OPT_LOWVRAM" = 1 ] && [ "$IS_MAC" = 0 ]; then
+    if ! "$VENV_PY" -c "import importlib.util,sys; sys.exit(0 if importlib.util.find_spec('triton') else 1)" 2>/dev/null; then
+        echo "      Режим экономии VRAM: устанавливаю triton..."
+        "$VENV_PY" -m pip install triton --quiet ||
+            echo "[ВНИМАНИЕ] triton не установился, --torch_compile может не заработать."
+    fi
+fi
+
+# ------------------------- проверка устройства -------------------------
+echo "[3/4] Проверяю устройство..."
+BF16=true
+if [ "$OPT_CPU" = 1 ]; then
+    echo "      Принудительный режим CPU. Генерация будет очень медленной."
+    DEVICE_ID=-1
+    BF16=false
+elif [ "$IS_MAC" = 1 ]; then
+    # На MPS bfloat16 приводит к ошибкам, см. README.
+    echo "      macOS: используется MPS, bf16 отключён."
+    BF16=false
+else
+    GPUINFO="$("$VENV_PY" -c "import torch; print((torch.cuda.get_device_name(0) + ' - ' + str(round(torch.cuda.get_device_properties(0).total_memory/1073741824, 1)) + ' GB VRAM') if torch.cuda.is_available() else 'NO_CUDA')" 2>/dev/null || echo NO_CUDA)"
+    if [ "$GPUINFO" = "NO_CUDA" ]; then
+        echo
+        echo "[ВНИМАНИЕ] CUDA недоступна - модель будет работать на процессоре (очень медленно)."
+        echo "Если у вас есть видеокарта NVIDIA, обновите драйвер и переустановите PyTorch:"
+        echo "    ./$SCRIPT_NAME --update"
+        echo
+    else
+        echo "      $GPUINFO"
+    fi
+fi
+
+if [ "$DO_SETUP_ONLY" = 1 ]; then
+    echo
+    echo "Установка завершена. Для запуска выполните ./$SCRIPT_NAME"
+    exit 0
+fi
+
+# ------------------------- запуск -------------------------
+ARGS=(--port "$PORT" --server_name "$SERVER_NAME" "--device_id=$DEVICE_ID" --bf16 "$BF16")
+[ -n "$CHECKPOINT_PATH" ] && ARGS+=(--checkpoint_path "$CHECKPOINT_PATH")
+[ "$OPT_SHARE" = 1 ] && ARGS+=(--share true)
+if [ "$OPT_LOWVRAM" = 1 ]; then
+    ARGS+=(--cpu_offload true --overlapped_decode true)
+    [ "$IS_MAC" = 0 ] && ARGS+=(--torch_compile true)
+fi
+
+echo "[4/4] Запускаю веб-интерфейс..."
+echo
+[ "$OPT_LOWVRAM" = 1 ] && echo "      Режим: экономия видеопамяти"
+echo "      Идёт загрузка моделей, это занимает время."
+echo "      При первом запуске дополнительно скачиваются веса (~8 ГБ)."
+echo
+echo "      Интерфейс будет доступен на http://$SERVER_NAME:$PORT"
+echo "      когда ниже появится строка \"Running on local URL\"."
+echo "      Остановить: Ctrl+C в этом окне."
+echo
+
+exec "$VENV_PY" -m acestep.gui "${ARGS[@]}"
@@ -0,0 +1,250 @@
+import torch
+import torchaudio
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import Optional
+import os
+import random
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import retrieve_timesteps
+from acestep.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+from acestep.pipeline_ace_step import ACEStepPipeline
+from acestep.apg_guidance import apg_forward, MomentumBuffer
+from transformers import AutoTokenizer
+from loguru import logger
+import uvicorn
+import time
+from datetime import datetime
+
+app = FastAPI(title="Text-to-Music API")
+
+class TextToMusicRequest(BaseModel):
+    prompt: str
+    duration: int = 240  # Duration in seconds (default 240s)
+    infer_steps: int = 60
+    guidance_scale: float = 15.0
+    omega_scale: float = 10.0
+    seed: Optional[int] = None
+
+class TextToMusicResponse(BaseModel):
+    audio_path: str
+    prompt: str
+    seed: int
+    sample_rate: int
+
+class InferencePipeline:
+    def __init__(self, checkpoint_dir: str, device: str = "cuda"):
+        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
+        logger.info(f"Initializing model on device: {self.device}")
+
+        # Load the ACEStepPipeline
+        self.acestep_pipeline = ACEStepPipeline(checkpoint_dir)
+        self.acestep_pipeline.load_checkpoint(checkpoint_dir)
+
+        # Initialize components
+        self.transformers = self.acestep_pipeline.ace_step_transformer.float().to(self.device).eval()
+        self.dcae = self.acestep_pipeline.music_dcae.float().to(self.device).eval()
+        self.text_encoder_model = self.acestep_pipeline.text_encoder_model.float().to(self.device).eval()
+        self.text_tokenizer = self.acestep_pipeline.text_tokenizer
+
+        # Ensure no gradients are computed
+        self.transformers.requires_grad_(False)
+        self.dcae.requires_grad_(False)
+        self.text_encoder_model.requires_grad_(False)
+
+        # Initialize scheduler
+        self.scheduler = FlowMatchEulerDiscreteScheduler(
+            num_train_timesteps=1000,
+            shift=3.0,
+        )
+
+    def get_text_embeddings(self, texts, device, text_max_length=256):
+        inputs = self.text_tokenizer(
+            texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=text_max_length,
+        )
+        inputs = {key: value.to(device) for key, value in inputs.items()}
+        with torch.no_grad():
+            outputs = self.text_encoder_model(**inputs)
+            last_hidden_states = outputs.last_hidden_state
+        attention_mask = inputs["attention_mask"]
+        return last_hidden_states, attention_mask
+
+    def diffusion_process(
+        self,
+        duration,
+        encoder_text_hidden_states,
+        text_attention_mask,
+        speaker_embds,
+        lyric_token_ids,
+        lyric_mask,
+        random_generator=None,
+        infer_steps=60,
+        guidance_scale=15.0,
+        omega_scale=10.0,
+    ):
+        do_classifier_free_guidance = guidance_scale > 1.0
+        device = encoder_text_hidden_states.device
+        dtype = encoder_text_hidden_states.dtype
+        bsz = encoder_text_hidden_states.shape[0]
+
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps=infer_steps, device=device
+        )
+
+        frame_length = int(duration * 44100 / 512 / 8)
+        target_latents = randn_tensor(
+            shape=(bsz, 8, 16, frame_length),
+            generator=random_generator,
+            device=device,
+            dtype=dtype,
+        )
+        attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)
+
+        if do_classifier_free_guidance:
+            attention_mask = torch.cat([attention_mask] * 2, dim=0)
+            encoder_text_hidden_states = torch.cat(
+                [encoder_text_hidden_states, torch.zeros_like(encoder_text_hidden_states)],
+                0,
+            )
+            text_attention_mask = torch.cat([text_attention_mask] * 2, dim=0)
+            speaker_embds = torch.cat([speaker_embds, torch.zeros_like(speaker_embds)], 0)
+            lyric_token_ids = torch.cat([lyric_token_ids, torch.zeros_like(lyric_token_ids)], 0)
+            lyric_mask = torch.cat([lyric_mask, torch.zeros_like(lyric_mask)], 0)
+
+        momentum_buffer = MomentumBuffer()
+
+        for t in timesteps:
+            latent_model_input = (
+                torch.cat([target_latents] * 2) if do_classifier_free_guidance else target_latents
+            )
+            timestep = t.expand(latent_model_input.shape[0])
+            with torch.no_grad():
+                noise_pred = self.transformers(
+                    hidden_states=latent_model_input,
+                    attention_mask=attention_mask,
+                    encoder_text_hidden_states=encoder_text_hidden_states,
+                    text_attention_mask=text_attention_mask,
+                    speaker_embeds=speaker_embds,
+                    lyric_token_idx=lyric_token_ids,
+                    lyric_mask=lyric_mask,
+                    timestep=timestep,
+                ).sample
+
+            if do_classifier_free_guidance:
+                noise_pred_with_cond, noise_pred_uncond = noise_pred.chunk(2)
+                noise_pred = apg_forward(
+                    pred_cond=noise_pred_with_cond,
+                    pred_uncond=noise_pred_uncond,
+                    guidance_scale=guidance_scale,
+                    momentum_buffer=momentum_buffer,
+                )
+
+            target_latents = self.scheduler.step(
+                model_output=noise_pred,
+                timestep=t,
+                sample=target_latents,
+                omega=omega_scale,
+            )[0]
+
+        return target_latents
+
+    def generate_audio(
+        self,
+        prompt: str,
+        duration: int,
+        infer_steps: int,
+        guidance_scale: float,
+        omega_scale: float,
+        seed: Optional[int],
+    ):
+        # Set random seed
+        if seed is not None:
+            random.seed(seed)
+            torch.manual_seed(seed)
+        else:
+            seed = random.randint(0, 2**32 - 1)
+            random.seed(seed)
+            torch.manual_seed(seed)
+
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+
+        # Get text embeddings
+        encoder_text_hidden_states, text_attention_mask = self.get_text_embeddings(
+            [prompt], self.device
+        )
+
+        # Dummy speaker embeddings and lyrics (since not provided in API request)
+        bsz = 1
+        speaker_embds = torch.zeros(bsz, 512, device=self.device, dtype=encoder_text_hidden_states.dtype)
+        lyric_token_ids = torch.zeros(bsz, 256, device=self.device, dtype=torch.long)
+        lyric_mask = torch.zeros(bsz, 256, device=self.device, dtype=torch.long)
+
+        # Run diffusion process
+        pred_latents = self.diffusion_process(
+            duration=duration,
+            encoder_text_hidden_states=encoder_text_hidden_states,
+            text_attention_mask=text_attention_mask,
+            speaker_embds=speaker_embds,
+            lyric_token_ids=lyric_token_ids,
+            lyric_mask=lyric_mask,
+            random_generator=generator,
+            infer_steps=infer_steps,
+            guidance_scale=guidance_scale,
+            omega_scale=omega_scale,
+        )
+
+        # Decode latents to audio
+        audio_lengths = torch.tensor([int(duration * 44100)], device=self.device)
+        sr, pred_wavs = self.dcae.decode(pred_latents, audio_lengths=audio_lengths, sr=48000)
+
+        # Save audio
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        output_dir = "generated_audio"
+        os.makedirs(output_dir, exist_ok=True)
+        audio_path = f"{output_dir}/generated_{timestamp}_{seed}.wav"
+        torchaudio.save(audio_path, pred_wavs.float().cpu(), sr)
+
+        return audio_path, sr, seed
+
+# Global model instance
+model = None
+
+@app.on_event("startup")
+async def startup_event():
+    global model
+    checkpoint_dir = os.getenv("CHECKPOINT_DIR", "./checkpoints")
+    model = InferencePipeline(checkpoint_dir=checkpoint_dir)
+    logger.info("Model loaded successfully")
+
+@app.post("/generate", response_model=TextToMusicResponse)
+async def generate_music(request: TextToMusicRequest):
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not initialized")
+
+    try:
+        start_time = time.time()
+        audio_path, sr, seed = model.generate_audio(
+            prompt=request.prompt,
+            duration=request.duration,
+            infer_steps=request.infer_steps,
+            guidance_scale=request.guidance_scale,
+            omega_scale=request.omega_scale,
+            seed=request.seed,
+        )
+        logger.info(f"Generation completed in {time.time() - start_time:.2f} seconds")
+        return TextToMusicResponse(
+            audio_path=audio_path,
+            prompt=request.prompt,
+            seed=seed,
+            sample_rate=sr,
+        )
+    except Exception as e:
+        logger.error(f"Error during generation: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)
@@ -49,9 +49,10 @@ class Pipeline(LightningModule):
        ssl_coeff: float = 1.0,
        checkpoint_dir=None,
        max_steps: int = 200000,
-        warmup_steps: int = 4000,
+        warmup_steps: int = 10,
        dataset_path: str = "./data/your_dataset_path",
        lora_config_path: str = None,
+        adapter_name: str = "lora_adapter",
    ):
        super().__init__()

@@ -62,26 +63,25 @@ class Pipeline(LightningModule):
        # Initialize scheduler
        self.scheduler = self.get_scheduler()

-        # Initialize local_rank for distributed training
-        self.local_rank = 0
-        if torch.distributed.is_initialized():
-            self.local_rank = torch.distributed.get_rank()
-
        # step 1: load model
        acestep_pipeline = ACEStepPipeline(checkpoint_dir)
-        acestep_pipeline.load_checkpoint(checkpoint_dir)
+        acestep_pipeline.load_checkpoint(acestep_pipeline.checkpoint_dir)

        transformers = acestep_pipeline.ace_step_transformer.float().cpu()
+        transformers.enable_gradient_checkpointing()

+        assert lora_config_path is not None, "Please provide a LoRA config path"
        if lora_config_path is not None:
            try:
                from peft import LoraConfig
            except ImportError:
                raise ImportError("Please install peft library to use LoRA training")
            with open(lora_config_path, encoding="utf-8") as f:
+                import json
                lora_config = json.load(f)
            lora_config = LoraConfig(**lora_config)
-            transformers.add_adapter(adapter_config=lora_config)
+            transformers.add_adapter(adapter_config=lora_config, adapter_name=adapter_name)
+            self.adapter_name = adapter_name

        self.transformers = transformers

@@ -95,9 +95,33 @@ class Pipeline(LightningModule):
        if self.is_train:
            self.transformers.train()

-            self.mert_model = AutoModel.from_pretrained(
-                "m-a-p/MERT-v1-330M", trust_remote_code=True, cache_dir=checkpoint_dir
-            ).eval()
+            # download first
+            try:
+                self.mert_model = AutoModel.from_pretrained(
+                    "m-a-p/MERT-v1-330M", trust_remote_code=True, cache_dir=checkpoint_dir
+                ).eval()
+            except:
+                import json
+                import os
+
+                mert_config_path = os.path.join(
+                    os.path.expanduser("~"),
+                    ".cache",
+                    "huggingface",
+                    "hub",
+                    "models--m-a-p--MERT-v1-330M",
+                    "blobs",
+                    "14f770758c7fe5c5e8ead4fe0f8e5fa727eb6942"
+                )
+
+                with open(mert_config_path) as f:
+                    mert_config = json.load(f)
+                mert_config["conv_pos_batch_norm"] = False
+                with open(mert_config_path, mode="w") as f:
+                    json.dump(mert_config, f)
+                self.mert_model = AutoModel.from_pretrained(
+                    "m-a-p/MERT-v1-330M", trust_remote_code=True, cache_dir=checkpoint_dir
+                ).eval()
            self.mert_model.requires_grad_(False)
            self.resampler_mert = torchaudio.transforms.Resample(
                orig_freq=48000, new_freq=24000
@@ -106,18 +130,13 @@ class Pipeline(LightningModule):
                "m-a-p/MERT-v1-330M", trust_remote_code=True
            )

-            self.hubert_model = AutoModel.from_pretrained(
-                "utter-project/mHuBERT-147",
-                local_files_only=True,
-                cache_dir=checkpoint_dir,
-            ).eval()
+            self.hubert_model = AutoModel.from_pretrained("utter-project/mHuBERT-147").eval()
            self.hubert_model.requires_grad_(False)
            self.resampler_mhubert = torchaudio.transforms.Resample(
                orig_freq=48000, new_freq=16000
            )
            self.processor_mhubert = Wav2Vec2FeatureExtractor.from_pretrained(
                "utter-project/mHuBERT-147",
-                local_files_only=True,
                cache_dir=checkpoint_dir,
            )

@@ -421,7 +440,7 @@ class Pipeline(LightningModule):
        lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
            optimizer, lr_lambda, last_epoch=-1
        )
-        return [optimizer], lr_scheduler
+        return [optimizer], [{"scheduler": lr_scheduler, "interval": "step"}]

    def train_dataloader(self):
        self.train_dataset = Text2MusicDataset(
@@ -583,6 +602,17 @@ class Pipeline(LightningModule):
    def training_step(self, batch, batch_idx):
        return self.run_step(batch, batch_idx)

+    def on_save_checkpoint(self, checkpoint):
+        state = {}
+        log_dir = self.logger.log_dir
+        epoch = self.current_epoch
+        step = self.global_step
+        checkpoint_name = f"epoch={epoch}-step={step}_lora"
+        checkpoint_dir = os.path.join(log_dir, "checkpoints", checkpoint_name)
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        self.transformers.save_lora_adapter(checkpoint_dir, adapter_name=self.adapter_name)
+        return state
+
    @torch.no_grad()
    def diffusion_process(
        self,
@@ -796,6 +826,8 @@ def main(args):
        every_plot_step=args.every_plot_step,
        dataset_path=args.dataset_path,
        checkpoint_dir=args.checkpoint_dir,
+        adapter_name=args.exp_name,
+        lora_config_path=args.lora_config_path
    )
    checkpoint_callback = ModelCheckpoint(
        monitor=None,
@@ -813,7 +845,7 @@ def main(args):
        num_nodes=args.num_nodes,
        precision=args.precision,
        accumulate_grad_batches=args.accumulate_grad_batches,
-        strategy="deepspeed_stage_2",
+        strategy="ddp_find_unused_parameters_true",
        max_epochs=args.epochs,
        max_steps=args.max_steps,
        log_every_n_steps=1,
@@ -840,9 +872,9 @@ if __name__ == "__main__":
    args.add_argument("--epochs", type=int, default=-1)
    args.add_argument("--max_steps", type=int, default=2000000)
    args.add_argument("--every_n_train_steps", type=int, default=2000)
-    args.add_argument("--dataset_path", type=str, default="./data/your_dataset_path")
-    args.add_argument("--exp_name", type=str, default="text2music_train_test")
-    args.add_argument("--precision", type=str, default="bf16-mixed")
+    args.add_argument("--dataset_path", type=str, default="./zh_lora_dataset")
+    args.add_argument("--exp_name", type=str, default="chinese_rap_lora")
+    args.add_argument("--precision", type=str, default="32")
    args.add_argument("--accumulate_grad_batches", type=int, default=1)
    args.add_argument("--devices", type=int, default=1)
    args.add_argument("--logger_dir", type=str, default="./exps/logs/")
@@ -853,6 +885,6 @@ if __name__ == "__main__":
    args.add_argument("--reload_dataloaders_every_n_epochs", type=int, default=1)
    args.add_argument("--every_plot_step", type=int, default=2000)
    args.add_argument("--val_check_interval", type=int, default=None)
-    args.add_argument("--lora_config_path", type=str, default=None)
+    args.add_argument("--lora_config_path", type=str, default="config/zh_rap_lora_config.json")
    args = args.parse_args()
    main(args)
@@ -0,0 +1,32 @@
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "keys": {
+      "dtype": "string",
+      "_type": "Value"
+    },
+    "filename": {
+      "dtype": "string",
+      "_type": "Value"
+    },
+    "tags": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "speaker_emb_path": {
+      "dtype": "string",
+      "_type": "Value"
+    },
+    "norm_lyrics": {
+      "dtype": "string",
+      "_type": "Value"
+    },
+    "recaption": {}
+  },
+  "homepage": "",
+  "license": ""
+}
@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "76c203d4bc1fdd7e",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}
				`@@ -0,0 +1 @@`
				`articulate, spoken word, young adult, rap music, female, clear, energetic, warm`