Upload 96 files

.gitattributes

@@ -33,3 +33,14 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/intro.jpg filter=lfs diff=lfs merge=lfs -text
+examples/muse_outputs/main_0.6b_0.mp3 filter=lfs diff=lfs merge=lfs -text
+examples/muse_outputs/main_0.6b_1.mp3 filter=lfs diff=lfs merge=lfs -text
+examples/muse_outputs/main_1.7b_2.mp3 filter=lfs diff=lfs merge=lfs -text
+examples/muse_outputs/main_8b_3.mp3 filter=lfs diff=lfs merge=lfs -text
+examples/train_inputs/suno_cn_0.mp3 filter=lfs diff=lfs merge=lfs -text
+examples/train_inputs/suno_cn_1.mp3 filter=lfs diff=lfs merge=lfs -text
+examples/train_inputs/suno_en_2.mp3 filter=lfs diff=lfs merge=lfs -text
+examples/train_inputs/suno_en_3.mp3 filter=lfs diff=lfs merge=lfs -text
+train/train_demo.jsonl filter=lfs diff=lfs merge=lfs -text
+train/val.jsonl filter=lfs diff=lfs merge=lfs -text

LICENSE

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

README.md

@@ -1,12 +1,80 @@
----
-title: Muse
-emoji: 🐢
-colorFrom: red
-colorTo: pink
-sdk: gradio
-sdk_version: 6.3.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Muse: Towards Reproducible Long-Form Song Generation with Fine-Grained Style Control
+
+<p align="center">
+  📄 <a href="https://arxiv.org/abs/2601.03973">Paper</a> • 📊 <a href="https://huggingface.co/datasets/bolshyC/Muse">Dataset</a> • 🤖 <a href="https://huggingface.co/bolshyC/models">Model</a> • 📚 <a href="#citation">Citation</a>
+</p>
+
+This repository is the official repository for "Muse: Towards Reproducible Long-Form Song Generation with Fine-Grained Style Control". In this repository, we provide the Muse model, training and inference scripts, pretrained checkpoints, and evaluation pipelines.
+
+## News and Updates
+
+* **2026.01.11 🔥**: We are excited to announce that all datasets and models are now fully open-sourced! 🎶 The complete training dataset (116k songs), pretrained model weights, training and evaluation code, and data pipeline are publicly available.
+
+## Installation
+
+**Requirements**: Python 3.10 is required.
+
+To set up the environment for Muse:
+
+- **For training**: Install the training framework:
+  ```bash
+  pip install ms-swift -U
+  ```
+- **For inference**: Install vLLM:
+  ```bash
+  pip install vllm
+  ```
+- **For audio encoding/decoding**: Some dependencies (e.g., `av`) require system-level packages. On Ubuntu/Debian, install FFmpeg 4.4+ first:
+  ```bash
+  sudo apt-get update
+  sudo apt-get install -y software-properties-common
+  sudo add-apt-repository ppa:savoury1/ffmpeg4 -y
+  sudo apt-get update
+  sudo apt-get install -y pkg-config ffmpeg libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev libswscale-dev libswresample-dev libavfilter-dev
+  ```
+  We recommend creating a new conda environment with Python 3.10. **Note**: Since `omegaconf==2.0.6` is required and has compatibility issues with pip 24.1+, you need to downgrade pip first:
+  ```bash
+  pip install "pip<24.1"
+  ```
+  Then install dependencies:
+  ```bash
+  pip install --default-timeout=1000 -r requirements_mucodec.txt
+  ```
+  For more details, please refer to the [MuCodec](https://github.com/tencent-ailab/MuCodec) official repository.
+
+- **For data pipeline and evaluation**: If you need to run data processing scripts (lyrics generation, metadata processing) or evaluation scripts, install additional dependencies:
+  ```bash
+  pip install -r requirements_data_eval.txt
+  ```
+
+## Repository Structure
+
+This repository contains the following main directories:
+
+- **`train/`**: Training scripts and utilities for fine-tuning the Muse model. See [`train/README.md`](train/README.md) for details.
+- **`infer/`**: Inference scripts for generating music with the Muse model. See [`infer/README.md`](infer/README.md) for details.
+- **`eval_pipeline/`**: Evaluation scripts for assessing model performance (Mulan-T, PER, AudioBox, SongEval, etc.).
+- **`data_pipeline/`**: Scripts for building and processing training data, including lyrics generation, metadata processing, and music generation utilities.
+
+## Model Architecture
+
+<p align="center">
+  <img src="assets/intro.jpg" width="800"/>
+</p>
+
+## Acknowledgments
+
+We thank [Qwen3](https://github.com/QwenLM/Qwen3) for providing the base language model, [ms-swift](https://github.com/modelscope/ms-swift) for the training framework, and [MuCodec](https://github.com/tencent-ailab/MuCodec) for discrete audio tokenization.
+
+## Citation
+
+If you find our work useful, please cite our paper:
+
+```bibtex
+@article{jiang2026muse,
+  title={Muse: Towards Reproducible Long-Form Song Generation with Fine-Grained Style Control},
+  author={Jiang, Changhao and Chen, Jiahao and Xiang, Zhenghao and Yang, Zhixiong and Wang, Hanchen and Zhuang, Jiabao and Che, Xinmeng and Sun, Jiajun and Li, Hui and Cao, Yifei and others},
+  journal={arXiv preprint arXiv:2601.03973},
+  year={2026}
+}
+```

baseline_generate/ace_step/convert.py

@@ -0,0 +1,160 @@
+"""
+Convert data to ACE-STEP acceptable format
+"""
+
+import re
+import json
+import random
+from tqdm import tqdm
+
+random.seed(42)
+
+def load_jsonl(path:str) -> list[dict]:
+    data = []
+    with open(path, 'r') as file:
+        for line in tqdm(file, desc=f"Loading {path}"):
+            data.append(json.loads(line))
+    return data
+
+def save_jsonl(data:list, path:str):
+    with open(path, 'w', encoding='utf-8') as file:
+        for ele in tqdm(data, desc=f"Saving {path}"):
+            json.dump(ele, file, ensure_ascii=False)
+            file.write("\n")
+
+START_STR = "Please generate a song in the following style:"
+END_STR = "\nNext, I will tell you the requirements and lyrics"
+
+def process_tag(content:str) -> str:
+    """Process segment label"""
+    # Extract label
+    end = content.find("[desc:")
+    tag = content[1:end-1]
+    # Lowercase & remove numbers & remove parentheses
+    tag = tag.lower()
+    tag = re.sub(r'\d+', '', tag)
+    tag = re.sub(r'\([^)]*\)', '', tag).strip()
+    if tag == "pre-chorus":
+        tag = "chorus"
+    return f"[{tag}]"
+
+def process_lyrics(content:str) -> str:
+    """Process segment lyrics"""
+    # Extract lyrics
+    start = content.find("[lyrics:\n")
+    if start == -1:
+        return ""
+    end = content.find("][phoneme:")
+    lyric = content[start+len("[lyrics:\n"):end]
+    
+    # Punctuation conversion
+    pattern = r'[,。"，:;&—‘\'.\]\[()?\n-]'
+    lyric = re.sub(pattern, '\n', lyric)
+    while lyric.find("\n\n") != -1:
+        lyric = lyric.replace("\n\n", "\n")
+    if lyric.endswith('\n'):
+        lyric = lyric[:-1]
+    return lyric
+
+def has_chinese(text) -> bool:
+    for char in text:
+        if '\u4e00' <= char <= '\u9fff':  # Basic Chinese characters
+            return True
+    return False
+
+def process_duration(lyrics:str):
+    if has_chinese(lyrics):
+        lyrics = lyrics.replace("\n", "")
+        length = len(lyrics)
+    else:
+        lyrics = lyrics.replace("\n", " ")
+        length = len(lyrics.split())
+    duration = random.randint(int(length * 0.4), int(length * 0.7))
+    return duration
+
+def process_one(messages:list[dict]):
+    """Process a conversation messages into input format, return gt_lyric and descriptions"""
+    # Overall style
+    style:str = messages[0]['content']
+    start = style.find(START_STR)
+    end = style.find(END_STR)
+    descriptions = style[start+len(START_STR):end]
+
+    # Line-by-line lyrics
+    all_lyrics = "[intro]\n\n"
+    pure_lyrics = ""
+    for message in messages[1:]:
+        if message['role'] == "assistant":
+            continue
+        content = message['content']
+        # Segment label
+        tag = process_tag(content)
+        # Segment lyrics
+        lyric = process_lyrics(content)
+        all_lyrics += f"{tag}\n{lyric}\n\n"
+        pure_lyrics += lyric
+    all_lyrics = all_lyrics[:-2]
+
+    # Duration
+    duration = process_duration(pure_lyrics)
+
+    obj = {
+        "prompt": descriptions,
+        "lyrics": all_lyrics,
+        "audio_duration": duration,
+        "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": True,
+        "use_erg_diffusion": True,
+        "oss_steps": [],
+        "actual_seeds": [
+            3299954530
+        ]
+    }
+    return obj
+
+def main():
+    path = "xxx/ACE-Step/data/inputs/messages.jsonl"
+    dataset = load_jsonl(path)
+
+    for id, ele in tqdm(enumerate(dataset), desc="Processing"):
+        messages = ele['messages']
+        data = process_one(messages)
+        path = f"./data/inputs/test_{id}.jsonl"
+        with open(path, 'w', encoding='utf-8') as file:
+            json.dump(data, file, ensure_ascii=False, indent=4)
+
+def load_jsonl(path):
+    dataset = []
+    with open(path, 'r') as file:
+        for line in file:
+            dataset.append(json.loads(line))
+    return dataset
+
+def save_jsonl(dataset, path):
+    with open(path, 'w', encoding='utf-8') as file:
+        for ele in dataset:
+            json.dump(ele, file, ensure_ascii=False)
+            file.write("\n")
+
+if __name__ == "__main__":
+    # main()
+    dataset = load_jsonl("./data/outputs/lyrics_params.jsonl")
+    for ele in dataset:
+        path = ele['audio_path']
+        ele['extra'] = int(path[len("./data/outputs/test_"):-len(".wav")])
+    sorted_data = sorted(dataset, key=lambda x: x['extra'])
+    
+    save_path = "./data/outputs/lyrics_params_.jsonl"
+    with open(save_path, 'w', encoding='utf-8') as file:
+        for ele in sorted_data:
+            del ele['extra']
+            json.dump(ele, file, ensure_ascii=False)
+            file.write("\n")

baseline_generate/ace_step/infer.py

@@ -0,0 +1,122 @@
+import click
+import os
+import json
+from acestep.pipeline_ace_step import ACEStepPipeline
+from acestep.data_sampler import DataSampler
+
+
+def sample_data(json_data):
+    return (
+        json_data["audio_duration"],
+        json_data["prompt"],
+        json_data["lyrics"],
+        json_data["infer_step"],
+        json_data["guidance_scale"],
+        json_data["scheduler_type"],
+        json_data["cfg_type"],
+        json_data["omega_scale"],
+        ", ".join(map(str, json_data["actual_seeds"])),
+        json_data["guidance_interval"],
+        json_data["guidance_interval_decay"],
+        json_data["min_guidance_scale"],
+        json_data["use_erg_tag"],
+        json_data["use_erg_lyric"],
+        json_data["use_erg_diffusion"],
+        ", ".join(map(str, json_data["oss_steps"])),
+        json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
+        (
+            json_data["guidance_scale_lyric"]
+            if "guidance_scale_lyric" in json_data
+            else 0.0
+        ),
+    )
+
+
+@click.command()
+@click.option(
+    "--checkpoint_path", type=str, default="", help="Path to the checkpoint directory"
+)
+@click.option("--bf16", type=bool, default=True, help="Whether to use bfloat16")
+@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, cpu_offload, overlapped_decode, device_id, output_path):
+    # 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()
+
+    inputs_dir = "./data/inputs"
+    for id, name in enumerate(os.listdir(inputs_dir)):
+        if not name.startswith("test"):
+            continue
+        path = os.path.join(inputs_dir, name)
+        with open(path, 'r') as file:
+            json_data = json.load(file)
+        json_data = sample_data(json_data)
+
+        pure_name = os.path.splitext(name)[0]
+        output_path = f"./data/outputs/{pure_name}.wav"
+        if os.path.exists(output_path):
+            continue
+        (
+            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
+
+        model_demo(
+            audio_duration=audio_duration,
+            prompt=prompt,
+            lyrics=lyrics,
+            infer_step=infer_step,
+            guidance_scale=guidance_scale,
+            scheduler_type=scheduler_type,
+            cfg_type=cfg_type,
+            omega_scale=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,
+        )
+
+if __name__ == "__main__":
+    main()

baseline_generate/diffrhythm2/batch_inference.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+set -euo pipefail
+
+# Navigate to script directory to ensure relative paths are consistent
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+cd "$SCRIPT_DIR"
+
+SONG_DIR="$SCRIPT_DIR/example/zh_songs"
+
+if [ ! -d "$SONG_DIR" ]; then
+  echo "Song directory not found: $SONG_DIR"
+  exit 1
+fi
+
+# Collect all song_*.jsonl files
+shopt -s nullglob
+SONG_FILES=("$SONG_DIR"/song_*.jsonl)
+shopt -u nullglob
+
+if [ ${#SONG_FILES[@]} -eq 0 ]; then
+  echo "No song_*.jsonl files in song directory"
+  exit 0
+fi
+
+export PYTHONPATH="${PYTHONPATH:-}:${SCRIPT_DIR}"
+
+espeak-ng --version
+
+# Reproducibility settings:
+# - Fixed random seed SEED
+# - DO_SAMPLE=0 tries to follow deterministic path (including fixed style prompt cropping start)
+SEED="${SEED:-42}"
+DO_SAMPLE="${DO_SAMPLE:-0}"
+
+# Further reduce cuBLAS non-determinism (enable when needed; comment out if causes errors)
+export CUBLAS_WORKSPACE_CONFIG="${CUBLAS_WORKSPACE_CONFIG:-:4096:8}"
+
+for SONG_FILE in "${SONG_FILES[@]}"; do
+  SONG_NAME="$(basename "$SONG_FILE")"
+  INPUT_PATH="./example/zh_songs/${SONG_NAME}"
+  echo "=============================="
+  echo "Starting generation: ${SONG_NAME}"
+  CMD=(python inference.py
+    --repo-id ASLP-lab/DiffRhythm2
+    --output-dir ./results/zh
+    --input-jsonl "$INPUT_PATH"
+    --cfg-strength 3.0
+    --max-secs 285.0
+    --seed "$SEED"
+  )
+  if [ "$DO_SAMPLE" -eq 1 ]; then
+    CMD+=(--do-sample)
+  fi
+  "${CMD[@]}"
+done
+
+echo "All songs generation complete, processed ${#SONG_FILES[@]} songs."

baseline_generate/diffrhythm2/batch_inference_en.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+set -euo pipefail
+
+# Navigate to script directory to ensure relative paths are consistent
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+cd "$SCRIPT_DIR"
+
+SONG_DIR="$SCRIPT_DIR/example/en_songs"
+
+if [ ! -d "$SONG_DIR" ]; then
+  echo "Song directory not found: $SONG_DIR"
+  exit 1
+fi
+
+# Collect all song_*.jsonl files
+shopt -s nullglob
+SONG_FILES=("$SONG_DIR"/song_*.jsonl)
+shopt -u nullglob
+
+if [ ${#SONG_FILES[@]} -eq 0 ]; then
+  echo "No song_*.jsonl files in song directory"
+  exit 0
+fi
+
+export PYTHONPATH="${PYTHONPATH:-}:${SCRIPT_DIR}"
+
+espeak-ng --version
+
+# Reproducibility settings:
+# - Fixed random seed SEED
+# - DO_SAMPLE=0 tries to follow deterministic path (including fixed style prompt cropping start)
+SEED="${SEED:-42}"
+DO_SAMPLE="${DO_SAMPLE:-0}"
+
+# Further reduce cuBLAS non-determinism (enable when needed; comment out if causes errors)
+export CUBLAS_WORKSPACE_CONFIG="${CUBLAS_WORKSPACE_CONFIG:-:4096:8}"
+
+for SONG_FILE in "${SONG_FILES[@]}"; do
+  SONG_NAME="$(basename "$SONG_FILE")"
+  INPUT_PATH="./example/en_songs/${SONG_NAME}"
+  echo "=============================="
+  echo "Starting generation: ${SONG_NAME}"
+  CMD=(python inference.py
+    --repo-id ASLP-lab/DiffRhythm2
+    --output-dir ./results/en
+    --input-jsonl "$INPUT_PATH"
+    --cfg-strength 3.0
+    --max-secs 285.0
+    --seed "$SEED"
+  )
+  if [ "$DO_SAMPLE" -eq 1 ]; then
+    CMD+=(--do-sample)
+  fi
+  "${CMD[@]}"
+done
+
+echo "All songs generation complete, processed ${#SONG_FILES[@]} songs."

baseline_generate/diffrhythm2/inference.py

@@ -0,0 +1,294 @@
+# Copyright 2025 ASLP Lab and Xiaomi Inc. 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 torch
+import torchaudio
+import argparse
+import json
+import os
+from tqdm import tqdm
+import random
+import pedalboard
+import numpy as np
+
+from muq import MuQMuLan
+from diffrhythm2.cfm import CFM
+from diffrhythm2.backbones.dit import DiT
+from bigvgan.model import Generator
+from huggingface_hub import hf_hub_download
+
+
+STRUCT_INFO = {
+    "[start]": 500,
+    "[end]": 501,
+    "[intro]": 502,
+    "[verse]": 503,
+    "[chorus]": 504,
+    "[outro]": 505,
+    "[inst]": 506,
+    "[solo]": 507,
+    "[bridge]": 508,
+    "[hook]": 509,
+    "[break]": 510,
+    "[stop]": 511,
+    "[space]": 512
+}
+
+lrc_tokenizer = None
+
+
+def set_seed(seed: int, deterministic: bool = True):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+    if deterministic:
+        # best-effort deterministic behavior; some ops may still be nondeterministic on certain GPUs/kernels
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        try:
+            torch.use_deterministic_algorithms(True, warn_only=True)
+        except Exception:
+            pass
+
+class CNENTokenizer():
+    def __init__(self):
+        curr_path = os.path.abspath(__file__)
+        vocab_path = os.path.join(os.path.dirname(curr_path), "g2p/g2p/vocab.json")
+        with open(vocab_path, 'r') as file:
+            self.phone2id:dict = json.load(file)['vocab']
+        self.id2phone = {v:k for (k, v) in self.phone2id.items()}
+        from g2p.g2p_generation import chn_eng_g2p
+        self.tokenizer = chn_eng_g2p
+    def encode(self, text):
+        phone, token = self.tokenizer(text)
+        token = [x+1 for x in token]
+        return token
+    def decode(self, token):
+        return "|".join([self.id2phone[x-1] for x in token])
+
+
+def prepare_model(repo_id, device):
+    diffrhythm2_ckpt_path = hf_hub_download(
+        repo_id=repo_id,
+        filename="model.safetensors",
+        local_dir="./ckpt",
+        local_files_only=False,
+    )
+    diffrhythm2_config_path = hf_hub_download(
+        repo_id=repo_id,
+        filename="config.json",
+        local_dir="./ckpt",
+        local_files_only=False,
+    )
+    with open(diffrhythm2_config_path) as f:
+        model_config = json.load(f)
+
+    model_config['use_flex_attn'] = False
+    diffrhythm2 = CFM(
+        transformer=DiT(
+            **model_config
+        ),
+        num_channels=model_config['mel_dim'],
+        block_size=model_config['block_size'],
+    )
+
+    total_params = sum(p.numel() for p in diffrhythm2.parameters())
+
+    diffrhythm2 = diffrhythm2.to(device)
+    if diffrhythm2_ckpt_path.endswith('.safetensors'):
+        from safetensors.torch import load_file
+        ckpt = load_file(diffrhythm2_ckpt_path)
+    else:
+        ckpt = torch.load(diffrhythm2_ckpt_path, map_location='cpu')
+    diffrhythm2.load_state_dict(ckpt)
+    print(f"Total params: {total_params:,}")
+
+    # load Mulan
+    mulan = MuQMuLan.from_pretrained("OpenMuQ/MuQ-MuLan-large", cache_dir="./ckpt").to(device)
+
+    # load frontend
+    lrc_tokenizer = CNENTokenizer()
+
+    # load decoder
+    decoder_ckpt_path = hf_hub_download(
+        repo_id=repo_id,
+        filename="decoder.bin",
+        local_dir="./ckpt",
+        local_files_only=False,
+    )
+    decoder_config_path = hf_hub_download(
+        repo_id=repo_id,
+        filename="decoder.json",
+        local_dir="./ckpt",
+        local_files_only=False,
+    )
+    decoder = Generator(decoder_config_path, decoder_ckpt_path)
+    decoder = decoder.to(device)
+    return diffrhythm2, mulan, lrc_tokenizer, decoder
+
+
+def parse_lyrics(lyrics: str):
+    lyrics_with_time = []
+    lyrics = lyrics.split("\n")
+    for line in lyrics:
+        struct_idx = STRUCT_INFO.get(line, None)
+        if struct_idx is not None:
+            lyrics_with_time.append([struct_idx, STRUCT_INFO['[stop]']])
+        else:
+            tokens = lrc_tokenizer.encode(line.strip())
+            tokens = tokens + [STRUCT_INFO['[stop]']]
+            lyrics_with_time.append(tokens)
+    return lyrics_with_time
+
+
+def make_fake_stereo(audio, sampling_rate):
+    left_channel = audio
+    right_channel = audio.copy()
+    right_channel = right_channel * 0.8
+    delay_samples = int(0.01 * sampling_rate)
+    right_channel = np.roll(right_channel, delay_samples)
+    right_channel[:,:delay_samples] = 0
+    stereo_audio = np.concatenate([left_channel, right_channel], axis=0)
+    
+    return stereo_audio
+    
+
+def inference(
+        model, 
+        decoder,
+        text,
+        style_prompt,
+        duration,
+        output_dir, 
+        song_name,
+        cfg_strength,
+        sample_steps=32,
+        process_bar=True,
+        fake_stereo=True,
+    ):
+    with torch.inference_mode():
+        latent = model.sample_block_cache(
+            text=text.unsqueeze(0),
+            duration=int(duration * 5),
+            style_prompt=style_prompt.unsqueeze(0),
+            steps=sample_steps,
+            cfg_strength=cfg_strength,
+            process_bar=process_bar,
+        )
+        latent = latent.transpose(1, 2)
+        audio = decoder.decode_audio(latent, overlap=5, chunk_size=20)
+
+        basename = f"{song_name}.mp3"
+        output_path = os.path.join(output_dir, basename)
+
+        num_channels = 1
+        audio = audio.float().cpu().numpy().squeeze()[None, :]
+        if fake_stereo:
+            audio = make_fake_stereo(audio, decoder.h.sampling_rate)
+            num_channels = 2
+        
+        with pedalboard.io.AudioFile(output_path, "w", decoder.h.sampling_rate, num_channels) as f:
+            f.write(audio)
+
+
+if __name__ == "__main__":
+    
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--repo-id', type=str, default=None)
+    parser.add_argument('--output-dir', type=str, default=None)
+    parser.add_argument('--input-jsonl', type=str, default=None)
+    parser.add_argument('--cfg-strength', type=float, default=2.0)
+    parser.add_argument('--max-secs', type=float, default=210.0)
+    parser.add_argument('--steps', type=int, default=16)
+    parser.add_argument('--fake-stereo', type=bool, default=True)
+    parser.add_argument('--seed', type=int, default=42)
+    parser.add_argument('--do-sample', action='store_true', default=False)
+
+    args = parser.parse_args()
+    
+    output_dir = args.output_dir
+    input_jsonl = args.input_jsonl
+    cfg_strength = args.cfg_strength
+    max_secs = args.max_secs
+    device = torch.device('cuda:7' if torch.cuda.is_available() else 'cpu')
+    dtype = torch.float16
+
+    # reproducibility
+    set_seed(args.seed, deterministic=(not args.do_sample))
+
+    # load diffrhythm2
+    diffrhythm2, mulan, lrc_tokenizer, decoder = prepare_model(args.repo_id, device)
+
+    output_dir = args.output_dir
+    os.makedirs(output_dir, exist_ok=True)
+    
+    with open(input_jsonl, 'r') as f:
+        input_info = [json.loads(i.strip()) for i in f.readlines()]
+
+    for i in tqdm(range(len(input_info))):
+        info = input_info[i]
+        song_name = info.get('song_name', f"{i:04d}")
+        lyrics = info.get('lyrics', None)
+        style_prompt = info.get('style_prompt', None)
+        if lyrics is None or style_prompt is None:
+            print(f"lyrics or style_prompt is None, skip {song_name}")
+            continue
+
+        # preprocess lyrics
+        with open(lyrics, 'r') as f:
+            lyrics = f.read()
+        lyrics_token = parse_lyrics(lyrics)
+        lyrics_token = torch.tensor(sum(lyrics_token, []), dtype=torch.long, device=device)
+
+        # preprocess style prompt
+        if os.path.isfile(style_prompt):
+            prompt_wav, sr = torchaudio.load(style_prompt)
+            prompt_wav = torchaudio.functional.resample(prompt_wav.to(device), sr, 24000)
+            if prompt_wav.shape[1] > 24000 * 10:
+                if args.do_sample:
+                    start = random.randint(0, prompt_wav.shape[1] - 24000 * 10)
+                else:
+                    start = 0
+                prompt_wav = prompt_wav[:, start:start+24000*10]
+            prompt_wav = prompt_wav.mean(dim=0, keepdim=True)
+            with torch.no_grad():
+                style_prompt_embed = mulan(wavs = prompt_wav)
+        else:
+            with torch.no_grad():
+                style_prompt_embed = mulan(texts = [style_prompt])
+        style_prompt_embed = style_prompt_embed.to(device).squeeze(0)
+
+        if device.type != 'cpu':
+            diffrhythm2 = diffrhythm2.half()
+            decoder = decoder.half()
+            style_prompt_embed = style_prompt_embed.half()
+
+        inference(
+            model=diffrhythm2,
+            decoder=decoder,
+            text=lyrics_token,
+            style_prompt=style_prompt_embed,
+            duration=max_secs,
+            output_dir=output_dir,
+            song_name=song_name,
+            sample_steps=args.steps,
+            cfg_strength=cfg_strength,
+            fake_stereo=args.fake_stereo,
+        )
+
+

baseline_generate/diffrhythm2/inference.sh

@@ -0,0 +1,10 @@
+
+export PYTHONPATH=$PYTHONPATH:$PWD
+espeak-ng --version
+
+python inference.py \
+    --repo-id ASLP-lab/DiffRhythm2 \
+    --output-dir ./results/test \
+    --input-jsonl ./example/song_1.jsonl \
+    --cfg-strength 3.0 \
+    --max-secs 285.0 \

baseline_generate/diffrhythm2/scripts/proce_song.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+Process songs.jsonl to generate corresponding lrc files and jsonl files.
+"""
+
+import json
+import os
+import re
+from pathlib import Path
+from typing import List
+
+INPUT_JSONL = Path("xxx/diffrhythm2/example/final_zh_test.jsonl")
+OUTPUT_SONG_DIR = Path("xxx/diffrhythm2/example/zh_songs")
+OUTPUT_LRC_DIR = Path("xxx/diffrhythm2/example/zh_lrc")
+
+TIMESTAMP_PATTERN = re.compile(r"\[\d{2}:\d{2}(?:\.\d+)?\]")
+STRUCTURE_PATTERN = re.compile(r"^\[[^\]]+\]$")
+
+
+def normalize_structure(tag: str) -> str:
+    """Convert structure tag to target format."""
+    tag_lower = tag.lower()
+    if tag_lower.startswith("verse"):
+        return "[verse]"
+    if "chorus" in tag_lower:
+        return "[chorus]"
+    if "bridge" in tag_lower:
+        return "[bridge]"
+    return f"[{tag_lower}]"
+
+
+def transform_lyrics(raw_lyrics: str) -> List[str]:
+    """Convert lyrics to LRC line list according to requirements."""
+    lines = ["[start]", "[intro]"]
+    for raw_line in raw_lyrics.splitlines():
+        line = raw_line.strip()
+        if not line:
+            continue
+
+        # Process structure tags separately
+        if STRUCTURE_PATTERN.match(line) and not TIMESTAMP_PATTERN.match(line):
+            tag_content = line[1:-1].strip()
+            lines.append(normalize_structure(tag_content))
+            continue
+
+        # Remove timestamps
+        text = TIMESTAMP_PATTERN.sub("", line).strip()
+        if not text:
+            continue
+        lines.append(text)
+
+    lines.append("[end]")
+    return lines
+
+
+def ensure_dirs() -> None:
+    OUTPUT_SONG_DIR.mkdir(parents=True, exist_ok=True)
+    OUTPUT_LRC_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def process_songs() -> None:
+    ensure_dirs()
+    with INPUT_JSONL.open("r", encoding="utf-8") as infile:
+        for idx, line in enumerate(infile, start=1):
+            line = line.strip()
+            if not line:
+                continue
+            data = json.loads(line)
+            description = data.get("description", "")
+            lyrics_raw = data.get("lyrics", "")
+
+            lrc_lines = transform_lyrics(lyrics_raw)
+            lrc_filename = f"song_{idx}.lrc"
+            lrc_path = OUTPUT_LRC_DIR / lrc_filename
+            lrc_path.write_text("\n".join(lrc_lines), encoding="utf-8")
+
+            song_base = f"song_{idx}"
+            song_filename = f"{song_base}.jsonl"
+            song_json_path = OUTPUT_SONG_DIR / song_filename
+            song_entry = {
+                "song_name": song_base,
+                "style_prompt": description,
+                "lyrics": f"example/zh_lrc/{lrc_filename}",
+            }
+            song_json_path.write_text(json.dumps(song_entry, ensure_ascii=False) + "\n", encoding="utf-8")
+            print(f"Processed song {idx}: {song_filename}")
+
+
+if __name__ == "__main__":
+    process_songs()

baseline_generate/diffrhythm2/scripts/proce_song_enprompt.py

@@ -0,0 +1,145 @@
+import json
+import os
+
+def extract_user_prompt(messages):
+    """
+    Extract all content from messages where role is user, and concatenate them
+    When concatenating each content, check if the concatenated prompt length exceeds 1600,
+    if so, do not concatenate and skip subsequent segments to ensure paragraph integrity
+    
+    Args:
+        messages: List of dictionaries, each containing role and content fields
+        
+    Returns:
+        Concatenated prompt string
+    """
+    # Collect all user message content, but need to check length limit
+    user_contents = []
+    current_length = 0  # Current concatenated length
+    
+    for msg in messages:
+        if msg.get("role") == "user":
+            content = msg.get("content", "")
+            if content:
+                # Calculate total length if this content is added
+                # Need to consider newline: if content already exists, need to add a newline
+                if user_contents:
+                    # Content already exists, need to add newline and current content
+                    new_length = current_length + 1 + len(content)  # 1 is newline length
+                else:
+                    # First content, no newline needed
+                    new_length = len(content)
+                
+                # If adding this content doesn't exceed 1600, add it
+                if new_length <= 1600:
+                    user_contents.append(content)
+                    current_length = new_length
+                else:
+                    # Exceeds 1600, don't add this content and stop processing subsequent segments
+                    break
+    
+    # Concatenate all content with newlines
+    if user_contents:
+        prompt = "\n".join(user_contents)
+        return prompt
+    
+    # If no user message found, return empty string
+    return ""
+
+def update_song_file(file_path, new_prompt):
+    """
+    Update style_prompt field in song file
+    
+    Args:
+        file_path: Path to song file
+        new_prompt: New prompt content
+    """
+    # Read file content
+    with open(file_path, 'r', encoding='utf-8') as f:
+        lines = [line.strip() for line in f if line.strip()]
+    
+    if not lines:
+        print(f"  Warning: File {file_path} is empty, skipping")
+        return
+    
+    # Read first JSON data
+    try:
+        data = json.loads(lines[0])
+        # Update style_prompt field
+        data['style_prompt'] = new_prompt
+        
+        # Write back to file
+        with open(file_path, 'w', encoding='utf-8') as f:
+            f.write(json.dumps(data, ensure_ascii=False) + '\n')
+            # If there's a second empty line, keep it
+            if len(lines) > 1:
+                f.write('\n')
+        
+        print(f"  ✓ Updated {file_path}")
+    except json.JSONDecodeError as e:
+        print(f"  Error: JSON parsing failed {file_path}: {e}")
+    except Exception as e:
+        print(f"  Error: Failed to update file {file_path}: {e}")
+
+def main():
+    # File paths
+    input_file = "xxx/diffrhythm2/scripts/test_messages.jsonl"
+    zh_songs_dir = "xxx/diffrhythm2/example/zh_songs"
+    en_songs_dir = "xxx/diffrhythm2/example/en_songs"
+    
+    print(f"Reading file: {input_file}")
+    
+    # Read all data
+    with open(input_file, 'r', encoding='utf-8') as f:
+        lines = [line.strip() for line in f if line.strip()]
+    
+    print(f"Read {len(lines)} entries")
+    
+    # Process each entry
+    for idx, line in enumerate(lines, 1):
+        try:
+            data = json.loads(line)
+            messages = data.get("messages", [])
+            
+            # Extract prompt
+            prompt = extract_user_prompt(messages)
+            
+            if not prompt:
+                print(f"Processing entry {idx}: No user content found, skipping")
+                continue
+            
+            # Determine if Chinese or English
+            if idx <= 50:
+                # First 50 entries: Chinese songs
+                song_num = idx
+                target_dir = zh_songs_dir
+                lang = "Chinese"
+            else:
+                # Entries 51-100: English songs
+                song_num = idx - 50  # 51->1, 52->2, ..., 100->50
+                target_dir = en_songs_dir
+                lang = "English"
+            
+            # Build file path
+            song_file = os.path.join(target_dir, f"song_{song_num}.jsonl")
+            
+            print(f"Processing entry {idx} ({lang}, song_{song_num})...")
+            print(f"  Prompt length: {len(prompt)} characters")
+            
+            # Update file
+            update_song_file(song_file, prompt)
+            
+        except json.JSONDecodeError as e:
+            print(f"JSON parsing failed for entry {idx}: {e}")
+            continue
+        except Exception as e:
+            print(f"Error processing entry {idx}: {e}")
+            import traceback
+            traceback.print_exc()
+            continue
+    
+    print(f"\nProcessing complete! Processed {len(lines)} entries")
+
+if __name__ == "__main__":
+    main()
+

baseline_generate/diffrhythm2/scripts/proce_song_enprompt_justtag.py

@@ -0,0 +1,143 @@
+import json
+import os
+
+def extract_user_prompt(messages):
+    """
+    Extract the first user role content from messages list, extract style string
+    
+    Format is usually: "Please generate a song in the following style: (style description)\n"
+    Only keep the style string part after the colon
+    
+    Args:
+        messages: List of dictionaries, each containing role and content fields
+        
+    Returns:
+        Style string
+    """
+    # Find first message with role user
+    for msg in messages:
+        if msg.get("role") == "user":
+            content = msg.get("content", "")
+            if content:
+                # Find position of "Please generate a song in the following style:"
+                style_prefix = "Please generate a song in the following style:"
+                style_index = content.find(style_prefix)
+                
+                if style_index != -1:
+                    # Find start position of content after colon
+                    start_index = style_index + len(style_prefix)
+                    # Find position of newline
+                    newline_index = content.find("\n", start_index)
+                    
+                    if newline_index != -1:
+                        # Extract content from after colon to before newline
+                        style_text = content[start_index:newline_index].strip()
+                    else:
+                        # If no newline, extract to end of string
+                        style_text = content[start_index:].strip()
+                    
+                    return style_text
+                else:
+                    # If standard format not found, return empty string
+                    return ""
+    
+    # If no user message found, return empty string
+    return ""
+
+def update_song_file(file_path, new_prompt):
+    """
+    Update style_prompt field in song file
+    
+    Args:
+        file_path: Path to song file
+        new_prompt: New prompt content
+    """
+    # Read file content
+    with open(file_path, 'r', encoding='utf-8') as f:
+        lines = [line.strip() for line in f if line.strip()]
+    
+    if not lines:
+        print(f"  Warning: File {file_path} is empty, skipping")
+        return
+    
+    # Read first JSON data
+    try:
+        data = json.loads(lines[0])
+        # Update style_prompt field
+        data['style_prompt'] = new_prompt
+        
+        # Write back to file
+        with open(file_path, 'w', encoding='utf-8') as f:
+            f.write(json.dumps(data, ensure_ascii=False) + '\n')
+            # If there's a second empty line, keep it
+            if len(lines) > 1:
+                f.write('\n')
+        
+        print(f"  ✓ Updated {file_path}")
+    except json.JSONDecodeError as e:
+        print(f"  Error: JSON parsing failed {file_path}: {e}")
+    except Exception as e:
+        print(f"  Error: Failed to update file {file_path}: {e}")
+
+def main():
+    # File paths
+    input_file = "xxx/diffrhythm2/scripts/test_messages.jsonl"
+    zh_songs_dir = "xxx/diffrhythm2/example/zh_songs"
+    en_songs_dir = "xxx/diffrhythm2/example/en_songs"
+    
+    print(f"Reading file: {input_file}")
+    
+    # Read all data
+    with open(input_file, 'r', encoding='utf-8') as f:
+        lines = [line.strip() for line in f if line.strip()]
+    
+    print(f"Read {len(lines)} entries")
+    
+    # Process each entry
+    for idx, line in enumerate(lines, 1):
+        try:
+            data = json.loads(line)
+            messages = data.get("messages", [])
+            
+            # Extract prompt
+            prompt = extract_user_prompt(messages)
+            
+            if not prompt:
+                print(f"Processing entry {idx}: No user content found, skipping")
+                continue
+            
+            # Determine if Chinese or English
+            if idx <= 50:
+                # First 50 entries: Chinese songs
+                song_num = idx
+                target_dir = zh_songs_dir
+                lang = "Chinese"
+            else:
+                # Entries 51-100: English songs
+                song_num = idx - 50  # 51->1, 52->2, ..., 100->50
+                target_dir = en_songs_dir
+                lang = "English"
+            
+            # Build file path
+            song_file = os.path.join(target_dir, f"song_{song_num}.jsonl")
+            
+            print(f"Processing entry {idx} ({lang}, song_{song_num})...")
+            print(f"  Prompt length: {len(prompt)} characters")
+            
+            # Update file
+            update_song_file(song_file, prompt)
+            
+        except json.JSONDecodeError as e:
+            print(f"JSON parsing failed for entry {idx}: {e}")
+            continue
+        except Exception as e:
+            print(f"Error processing entry {idx}: {e}")
+            import traceback
+            traceback.print_exc()
+            continue
+    
+    print(f"\nProcessing complete! Processed {len(lines)} entries")
+
+if __name__ == "__main__":
+    main()
+

baseline_generate/levo/convert.py

@@ -0,0 +1,110 @@
+"""
+Convert data to ACE-STEP acceptable format
+"""
+
+import re
+import json
+import random
+from tqdm import tqdm
+
+random.seed(42)
+
+def load_jsonl(path:str) -> list[dict]:
+    data = []
+    with open(path, 'r') as file:
+        for line in tqdm(file, desc=f"Loading {path}"):
+            data.append(json.loads(line))
+    return data
+
+def save_jsonl(data:list, path:str):
+    with open(path, 'w', encoding='utf-8') as file:
+        for ele in tqdm(data, desc=f"Saving {path}"):
+            json.dump(ele, file, ensure_ascii=False)
+            file.write("\n")
+
+START_STR = "Please generate a song in the following style:"
+END_STR = "\nNext, I will tell you the requirements and lyrics"
+
+def process_tag(content:str) -> str:
+    """Process segment label"""
+    # Extract label
+    end = content.find("[desc:")
+    tag = content[1:end-1]
+    # Lowercase & remove numbers & remove parentheses
+    tag = tag.lower()
+    tag = re.sub(r'\d+', '', tag)
+    tag = re.sub(r'\([^)]*\)', '', tag).strip()
+    if tag == "pre-chorus":
+        tag = "chorus"
+    return f"[{tag}]"
+
+def process_lyrics(content:str) -> str:
+    """Process segment lyrics"""
+    # Extract lyrics
+    start = content.find("[lyrics:\n")
+    if start == -1:
+        return ""
+    end = content.find("][phoneme:")
+    lyric = content[start+len("[lyrics:\n"):end]
+    
+    # Punctuation conversion
+    pattern = r'[,。"，:;&—‘\'.\]\[()?\n-]'
+    lyric = re.sub(pattern, '.', lyric)
+    while lyric.find("..") != -1:
+        lyric = lyric.replace("..", ".")
+    if lyric.endswith('.'):
+        lyric = lyric[:-1]
+    return lyric
+
+def random_size() -> str:
+    # Intro/outro length
+    sizes = ['short', 'medium', 'long']
+    return random.choice(sizes)
+
+def process_one(messages:list[dict]):
+    """Process a conversation messages into input format, return gt_lyric and descriptions"""
+    # Overall style
+    style:str = messages[0]['content']
+    start = style.find(START_STR)
+    end = style.find(END_STR)
+    descriptions = style[start+len(START_STR):end]
+
+    # Line-by-line lyrics
+    start_tag = "intro-" + random_size()
+    end_tag = "outro-" + random_size()
+    gt_lyric = f"[{start_tag}] ;"
+    for message in messages[1:]:
+        if message['role'] == "assistant":
+            continue
+        content = message['content']
+        # Segment label
+        tag = process_tag(content)
+        # Segment lyrics
+        lyric = process_lyrics(content)
+        if lyric == "" or tag.startswith("[outro"):
+            gt_lyric += f" [{end_tag}]"
+            break
+        gt_lyric += f" {tag} {lyric} ;"
+    if not gt_lyric.endswith(f" [{end_tag}]"):
+        gt_lyric += f" [{end_tag}]"
+    return descriptions, gt_lyric
+
+def main():
+    path = "xxx/SongGeneration/data/inputs/test_messages.jsonl"
+    dataset = load_jsonl(path)
+    save_path = "xxx/SongGeneration/data/inputs/lyrics.jsonl"
+
+    with open(save_path, 'w', encoding='utf-8') as file:
+        for id, ele in tqdm(enumerate(dataset), desc="Processing"):
+            messages = ele['messages']
+            descriptions, gt_lyric = process_one(messages)
+            data = {
+                "idx": f"test_{id}",
+                "descriptions": descriptions,
+                "gt_lyric": gt_lyric
+            }
+            json.dump(data, file, ensure_ascii=False)
+            file.write("\n")
+
+if __name__ == "__main__":
+    main()

baseline_generate/levo/generate.py

@@ -0,0 +1,591 @@
+from hmac import new
+import sys
+import os
+import argparse
+
+import time
+import json
+import torch
+import torchaudio
+import numpy as np
+from omegaconf import OmegaConf
+from codeclm.models import builders
+import gc
+from codeclm.trainer.codec_song_pl import CodecLM_PL
+from codeclm.models import CodecLM
+from third_party.demucs.models.pretrained import get_model_from_yaml
+import re
+import subprocess
+
+auto_prompt_type = ['Pop', 'R&B', 'Dance', 'Jazz', 'Folk', 'Rock', 'Chinese Style', 'Chinese Tradition', 'Metal', 'Reggae', 'Chinese Opera', 'Auto']
+
+def get_free_gpu() -> int:
+    """Return the GPU ID with the least memory usage"""
+    cmd = "nvidia-smi --query-gpu=index,memory.free --format=csv,noheader,nounits"
+    result = subprocess.check_output(cmd.split()).decode().strip().split("\n")
+
+    free_list = []
+    for line in result:
+        idx, free_mem = line.split(",")
+        free_list.append((int(idx), int(free_mem)))  # (GPU id, free memory MiB)
+    
+    # Sort by remaining memory
+    free_list.sort(key=lambda x: x[1], reverse=True)
+    return free_list[0][0]
+
+class Separator:
+    def __init__(self, dm_model_path='third_party/demucs/ckpt/htdemucs.pth', dm_config_path='third_party/demucs/ckpt/htdemucs.yaml', gpu_id=0) -> None:
+        gpu_id = get_free_gpu()
+        self.device = f"cuda:{gpu_id}"
+        print(f"Using {self.device}")
+        
+        # if torch.cuda.is_available() and gpu_id < torch.cuda.device_count():
+        #     self.device = torch.device(f"cuda:{gpu_id}")
+        # else:
+        #     self.device = torch.device("cpu")
+
+        self.demucs_model = self.init_demucs_model(dm_model_path, dm_config_path)
+
+    def init_demucs_model(self, model_path, config_path):
+        model = get_model_from_yaml(config_path, model_path)
+        model.to(self.device)
+        model.eval()
+        return model
+    
+    def load_audio(self, f):
+        a, fs = torchaudio.load(f)
+        if (fs != 48000):
+            a = torchaudio.functional.resample(a, fs, 48000)
+        if a.shape[-1] >= 48000*10:
+            a = a[..., :48000*10]
+        return a[:, 0:48000*10]
+    
+    def run(self, audio_path, output_dir='tmp', ext=".flac"):
+        os.makedirs(output_dir, exist_ok=True)
+        name, _ = os.path.splitext(os.path.split(audio_path)[-1])
+        output_paths = []
+
+        for stem in self.demucs_model.sources:
+            output_path = os.path.join(output_dir, f"{name}_{stem}{ext}")
+            if os.path.exists(output_path):
+                output_paths.append(output_path)
+        if len(output_paths) == 1:  # 4
+            vocal_path = output_paths[0]
+        else:
+            drums_path, bass_path, other_path, vocal_path = self.demucs_model.separate(audio_path, output_dir, device=self.device)
+            for path in [drums_path, bass_path, other_path]:
+                os.remove(path)
+        full_audio = self.load_audio(audio_path)
+        vocal_audio = self.load_audio(vocal_path)
+        bgm_audio = full_audio - vocal_audio
+        return full_audio, vocal_audio, bgm_audio
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Song Generation Script')
+    
+    # Required parameters
+    parser.add_argument('--ckpt_path', type=str, required=True,
+                      help='Path to the checkpoint directory containing config.yaml and model.pt')
+    parser.add_argument('--input_jsonl', type=str, required=True,
+                      help='Path to input JSONL file containing generation tasks')
+    parser.add_argument('--save_dir', type=str, required=True,
+                      help='Directory to save generated audio files and results')
+    # Optional parameters
+    parser.add_argument('--generate_type', type=str, default='mixed',
+                      help='Type of generation: "vocal" or "bgm" or "separate" or "mixed" (default: "mixed")')
+    parser.add_argument('--use_flash_attn', action='store_true',
+                      help='Whether to use flash attention (default: False)')
+    parser.add_argument('--low_mem', action='store_true',
+                      help='Whether to use low memory mode (default: False)')
+    return parser.parse_args()
+
+def generate(args):
+    torch.set_num_threads(1)
+    ckpt_path = args.ckpt_path
+    input_jsonl = args.input_jsonl
+    save_dir = args.save_dir
+    cfg_path = os.path.join(ckpt_path, 'config.yaml')
+    ckpt_path = os.path.join(ckpt_path, 'model.pt')
+    cfg = OmegaConf.load(cfg_path)
+    cfg.lm.use_flash_attn_2 = args.use_flash_attn
+    print(f"use_flash_attn: {args.use_flash_attn}")
+    cfg.mode = 'inference'
+    max_duration = cfg.max_dur
+    gen_type = args.generate_type
+    
+
+    separator = Separator()
+    auto_prompt = torch.load('tools/new_prompt.pt')
+    audio_tokenizer = builders.get_audio_tokenizer_model(cfg.audio_tokenizer_checkpoint, cfg)
+    audio_tokenizer = audio_tokenizer.eval().cuda()
+    with open(input_jsonl, "r") as fp:
+        lines = fp.readlines()
+
+        
+    new_items = []
+    for line in lines:
+        item = json.loads(line)
+        target_wav_name = f"{save_dir}/audios/{item['idx']}.flac"
+        # get prompt audio
+        if "prompt_audio_path" in item:
+            assert os.path.exists(item['prompt_audio_path']), f"prompt_audio_path {item['prompt_audio_path']} not found"
+            assert 'auto_prompt_audio_type' not in item, f"auto_prompt_audio_type and prompt_audio_path cannot be used together"
+            with torch.no_grad():
+                pmt_wav, vocal_wav, bgm_wav = separator.run(item['prompt_audio_path'])
+            item['raw_pmt_wav'] = pmt_wav
+            item['raw_vocal_wav'] = vocal_wav
+            item['raw_bgm_wav'] = bgm_wav
+            if pmt_wav.dim() == 2:
+                pmt_wav = pmt_wav[None]
+            if pmt_wav.dim() != 3:
+                raise ValueError("Melody wavs should have a shape [B, C, T].")
+            pmt_wav = list(pmt_wav)
+            if vocal_wav.dim() == 2:
+                vocal_wav = vocal_wav[None]
+            if vocal_wav.dim() != 3:
+                raise ValueError("Vocal wavs should have a shape [B, C, T].")
+            vocal_wav = list(vocal_wav)
+            if bgm_wav.dim() == 2:
+                bgm_wav = bgm_wav[None]
+            if bgm_wav.dim() != 3:
+                raise ValueError("BGM wavs should have a shape [B, C, T].")
+            bgm_wav = list(bgm_wav)
+            if type(pmt_wav) == list:
+                pmt_wav = torch.stack(pmt_wav, dim=0)
+            if type(vocal_wav) == list:
+                vocal_wav = torch.stack(vocal_wav, dim=0)
+            if type(bgm_wav) == list:
+                bgm_wav = torch.stack(bgm_wav, dim=0)
+            pmt_wav = pmt_wav
+            vocal_wav = vocal_wav
+            bgm_wav = bgm_wav
+            with torch.no_grad():
+                pmt_wav, _ = audio_tokenizer.encode(pmt_wav.cuda())
+            melody_is_wav = False
+        elif "auto_prompt_audio_type" in item:
+            assert item["auto_prompt_audio_type"] in auto_prompt_type, f"auto_prompt_audio_type {item['auto_prompt_audio_type']} not found"
+            prompt_token = auto_prompt[item["auto_prompt_audio_type"]][np.random.randint(0, len(auto_prompt[item["auto_prompt_audio_type"]]))]
+            pmt_wav = prompt_token[:,[0],:]
+            vocal_wav = prompt_token[:,[1],:]
+            bgm_wav = prompt_token[:,[2],:]
+            melody_is_wav = False
+        else:
+            pmt_wav = None
+            vocal_wav = None
+            bgm_wav = None
+            melody_is_wav = True
+        item['pmt_wav'] = pmt_wav
+        item['vocal_wav'] = vocal_wav
+        item['bgm_wav'] = bgm_wav
+        item['melody_is_wav'] = melody_is_wav
+        item["idx"] = f"{item['idx']}"
+        item["wav_path"] = target_wav_name
+        new_items.append(item)
+
+    del audio_tokenizer
+    del separator
+
+    torch.cuda.empty_cache()
+
+    if "audio_tokenizer_checkpoint_sep" in cfg.keys():
+        seperate_tokenizer = builders.get_audio_tokenizer_model(cfg.audio_tokenizer_checkpoint_sep, cfg) 
+    else:
+        seperate_tokenizer = None
+    
+    if seperate_tokenizer is not None:
+        seperate_tokenizer = seperate_tokenizer.eval().cuda()
+
+    for item in new_items:
+        if "prompt_audio_path" in item:
+            with torch.no_grad():
+                vocal_wav, bgm_wav = seperate_tokenizer.encode(item['vocal_wav'].cuda(), item['bgm_wav'].cuda())
+            item['vocal_wav'] = vocal_wav
+            item['bgm_wav'] = bgm_wav
+
+    torch.cuda.empty_cache()
+    audiolm = builders.get_lm_model(cfg)
+    checkpoint = torch.load(ckpt_path, map_location='cpu')
+    audiolm_state_dict = {k.replace('audiolm.', ''): v for k, v in checkpoint.items() if k.startswith('audiolm')}
+    audiolm.load_state_dict(audiolm_state_dict, strict=False)
+    audiolm = audiolm.eval()
+    audiolm = audiolm.cuda().to(torch.float16)
+
+    model = CodecLM(name = "tmp",
+        lm = audiolm,
+        audiotokenizer = None,
+        max_duration = max_duration,
+        seperate_tokenizer = seperate_tokenizer,
+    )
+
+    cfg_coef = 1.5 #25
+    temp = 0.9
+    top_k = 50
+    top_p = 0.0
+    record_tokens = True
+    record_window = 50
+
+    model.set_generation_params(duration=max_duration, extend_stride=5, temperature=temp, cfg_coef=cfg_coef,
+                                top_k=top_k, top_p=top_p, record_tokens=record_tokens, record_window=record_window)
+    os.makedirs(save_dir, exist_ok=True)
+    os.makedirs(save_dir + "/audios", exist_ok=True)
+    os.makedirs(save_dir + "/jsonl", exist_ok=True)
+
+    for item in new_items:
+        lyric = item["gt_lyric"]
+        descriptions = item["descriptions"] if "descriptions" in item else None
+        pmt_wav = item['pmt_wav']
+        vocal_wav = item['vocal_wav']
+        bgm_wav = item['bgm_wav']
+        melody_is_wav = item['melody_is_wav']
+        target_wav_name = f"{save_dir}/audios/{item['idx']}.flac"
+
+
+        generate_inp = {
+            'lyrics': [lyric.replace("  ", " ")],
+            'descriptions': [descriptions],
+            'melody_wavs': pmt_wav,
+            'vocal_wavs': vocal_wav,
+            'bgm_wavs': bgm_wav,
+            'melody_is_wav': melody_is_wav,
+        }
+        start_time = time.time()
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            with torch.no_grad():
+                tokens = model.generate(**generate_inp, return_tokens=True)
+        mid_time = time.time()
+
+        with torch.no_grad():
+            if 'raw_pmt_wav' in item:
+                if gen_type == 'separate':
+                    wav_seperate = model.generate_audio(tokens, item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'], chunked=True, gen_type='mixed')
+                    wav_vocal = model.generate_audio(tokens, item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'], chunked=True, gen_type='vocal')
+                    wav_bgm = model.generate_audio(tokens, item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'], chunked=True, gen_type='bgm')
+                elif gen_type == 'mixed':
+                    wav_seperate = model.generate_audio(tokens, item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'],chunked=True, gen_type=gen_type)
+                else:
+                    wav_seperate = model.generate_audio(tokens,chunked=True, gen_type=gen_type)
+                del item['raw_pmt_wav']
+                del item['raw_vocal_wav']
+                del item['raw_bgm_wav']
+            else:
+                if gen_type == 'separate':
+                    wav_vocal = model.generate_audio(tokens, chunked=True, gen_type='vocal')
+                    wav_bgm = model.generate_audio(tokens, chunked=True, gen_type='bgm')
+                    wav_seperate = model.generate_audio(tokens, chunked=True, gen_type='mixed')
+                else:
+                    wav_seperate = model.generate_audio(tokens, chunked=True, gen_type=gen_type)
+        del item['pmt_wav']
+        del item['vocal_wav']
+        del item['bgm_wav']
+        del item['melody_is_wav']
+        end_time = time.time()
+        if gen_type == 'separate':
+            torchaudio.save(target_wav_name.replace('.flac', '_vocal.flac'), wav_vocal[0].cpu().float(), cfg.sample_rate)
+            torchaudio.save(target_wav_name.replace('.flac', '_bgm.flac'), wav_bgm[0].cpu().float(), cfg.sample_rate)
+            torchaudio.save(target_wav_name, wav_seperate[0].cpu().float(), cfg.sample_rate)
+        else:
+            torchaudio.save(target_wav_name, wav_seperate[0].cpu().float(), cfg.sample_rate)
+
+        print(f"process{item['idx']}, lm cost {mid_time - start_time}s, diffusion cost {end_time - mid_time}")
+        item["idx"] = f"{item['idx']}"
+        item["wav_path"] = target_wav_name
+    
+    src_jsonl_name = os.path.split(input_jsonl)[-1]
+    with open(f"{save_dir}/jsonl/{src_jsonl_name}.jsonl", "w", encoding='utf-8') as fw:
+        for item in new_items:
+            fw.writelines(json.dumps(item, ensure_ascii=False)+"\n")
+
+def generate_lowmem(args):
+    torch.set_num_threads(1)
+    ckpt_path = args.ckpt_path
+    input_jsonl = args.input_jsonl
+    save_dir = args.save_dir
+    cfg_path = os.path.join(ckpt_path, 'config.yaml')
+    ckpt_path = os.path.join(ckpt_path, 'model.pt')
+    cfg = OmegaConf.load(cfg_path)
+    cfg.lm.use_flash_attn_2 = args.use_flash_attn
+    print(f"use_flash_attn: {args.use_flash_attn}")
+    cfg.mode = 'inference'
+    max_duration = cfg.max_dur
+    gen_type = args.generate_type
+    chunk_size = 128
+    use_audio_tokenizer = False
+    with open(input_jsonl, "r") as fp:
+        lines = fp.readlines()
+    for line in lines:
+        item = json.loads(line)
+        if "prompt_audio_path" in item:
+            use_audio_tokenizer = True
+            break
+    if use_audio_tokenizer:
+        separator = Separator()
+        audio_tokenizer = builders.get_audio_tokenizer_model(cfg.audio_tokenizer_checkpoint, cfg)
+        audio_tokenizer = audio_tokenizer.eval().cuda()
+    auto_prompt = torch.load('tools/new_prompt.pt')
+    new_items = []
+    for line in lines:
+        item = json.loads(line)
+        target_wav_name = f"{save_dir}/audios/{item['idx']}.flac"
+        # get prompt audio
+        if "prompt_audio_path" in item:
+            assert os.path.exists(item['prompt_audio_path']), f"prompt_audio_path {item['prompt_audio_path']} not found"
+            assert 'auto_prompt_audio_type' not in item, f"auto_prompt_audio_type and prompt_audio_path cannot be used together"
+            with torch.no_grad():
+                pmt_wav, vocal_wav, bgm_wav = separator.run(item['prompt_audio_path'])
+            item['raw_pmt_wav'] = pmt_wav
+            item['raw_vocal_wav'] = vocal_wav
+            item['raw_bgm_wav'] = bgm_wav
+            if pmt_wav.dim() == 2:
+                pmt_wav = pmt_wav[None]
+            if pmt_wav.dim() != 3:
+                raise ValueError("Melody wavs should have a shape [B, C, T].")
+            pmt_wav = list(pmt_wav)
+            if vocal_wav.dim() == 2:
+                vocal_wav = vocal_wav[None]
+            if vocal_wav.dim() != 3:
+                raise ValueError("Vocal wavs should have a shape [B, C, T].")
+            vocal_wav = list(vocal_wav)
+            if bgm_wav.dim() == 2:
+                bgm_wav = bgm_wav[None]
+            if bgm_wav.dim() != 3:
+                raise ValueError("BGM wavs should have a shape [B, C, T].")
+            bgm_wav = list(bgm_wav)
+            if type(pmt_wav) == list:
+                pmt_wav = torch.stack(pmt_wav, dim=0)
+            if type(vocal_wav) == list:
+                vocal_wav = torch.stack(vocal_wav, dim=0)
+            if type(bgm_wav) == list:
+                bgm_wav = torch.stack(bgm_wav, dim=0)
+            with torch.no_grad():
+                pmt_wav, _ = audio_tokenizer.encode(pmt_wav.cuda())
+            melody_is_wav = False
+        elif "auto_prompt_audio_type" in item:
+            assert item["auto_prompt_audio_type"] in auto_prompt_type, f"auto_prompt_audio_type {item['auto_prompt_audio_type']} not found"
+            prompt_token = auto_prompt[item["auto_prompt_audio_type"]][np.random.randint(0, len(auto_prompt[item["auto_prompt_audio_type"]]))]
+            pmt_wav = prompt_token[:,[0],:]
+            vocal_wav = prompt_token[:,[1],:]
+            bgm_wav = prompt_token[:,[2],:]
+            melody_is_wav = False
+        else:
+            pmt_wav = None
+            vocal_wav = None
+            bgm_wav = None
+            melody_is_wav = True
+        item['pmt_wav'] = pmt_wav
+        item['vocal_wav'] = vocal_wav
+        item['bgm_wav'] = bgm_wav
+        item['melody_is_wav'] = melody_is_wav
+        item["idx"] = f"{item['idx']}"
+        item["wav_path"] = target_wav_name
+        new_items.append(item)
+
+    if use_audio_tokenizer:
+        del audio_tokenizer
+        del separator
+
+    torch.cuda.empty_cache()
+    
+    if "audio_tokenizer_checkpoint_sep" in cfg.keys() and use_audio_tokenizer:
+        seperate_tokenizer = builders.get_audio_tokenizer_model(cfg.audio_tokenizer_checkpoint_sep, cfg) 
+    else:
+        seperate_tokenizer = None
+    
+    if seperate_tokenizer is not None:
+        seperate_tokenizer = seperate_tokenizer.eval().cuda()
+
+    for item in new_items:
+        if "prompt_audio_path" in item:
+            with torch.no_grad():
+                vocal_wav, bgm_wav = seperate_tokenizer.encode(item['vocal_wav'].cuda(), item['bgm_wav'].cuda())
+            item['vocal_wav'] = vocal_wav
+            item['bgm_wav'] = bgm_wav
+
+    if use_audio_tokenizer:
+        del seperate_tokenizer
+
+    torch.cuda.empty_cache()
+
+    # Define model or load pretrained model
+    audiolm = builders.get_lm_model(cfg)
+    checkpoint = torch.load(ckpt_path, map_location='cpu')
+    audiolm_state_dict = {k.replace('audiolm.', ''): v for k, v in checkpoint.items() if k.startswith('audiolm')}
+    audiolm.load_state_dict(audiolm_state_dict, strict=False)
+    audiolm = audiolm.eval()
+
+    offload_audiolm = True if 'offload' in cfg.keys() and 'audiolm' in cfg.offload else False
+    if offload_audiolm:
+        audiolm_offload_param = OffloadParamParse.parse_config(audiolm, cfg.offload.audiolm)
+        audiolm_offload_param.show()
+        offload_profiler = OffloadProfiler(device_index=0, **(audiolm_offload_param.init_param_dict()))
+        offload_profiler.offload_layer(**(audiolm_offload_param.offload_layer_param_dict()))
+        offload_profiler.clean_cache_wrapper(**(audiolm_offload_param.clean_cache_param_dict()))
+    else:
+        audiolm = audiolm.cuda().to(torch.float16)
+
+    model = CodecLM(name = "tmp",
+        lm = audiolm,
+        audiotokenizer = None,
+        max_duration = max_duration,
+        seperate_tokenizer = None,
+    )
+    
+    cfg_coef = 1.5 #25
+    temp = 0.9
+    top_k = 50
+    top_p = 0.0
+    record_tokens = True
+    record_window = 50
+    
+
+    model.set_generation_params(duration=max_duration, extend_stride=5, temperature=temp, cfg_coef=cfg_coef,
+                                top_k=top_k, top_p=top_p, record_tokens=record_tokens, record_window=record_window)
+    os.makedirs(save_dir, exist_ok=True)
+    os.makedirs(save_dir + "/audios", exist_ok=True)
+    os.makedirs(save_dir + "/jsonl", exist_ok=True)
+
+    
+    for item in new_items:
+        lyric = item["gt_lyric"]
+        descriptions = item["descriptions"] if "descriptions" in item else None
+        pmt_wav = item['pmt_wav']
+        vocal_wav = item['vocal_wav']
+        bgm_wav = item['bgm_wav']
+        melody_is_wav = item['melody_is_wav']
+            
+        generate_inp = {
+            'lyrics': [lyric.replace("  ", " ")],
+            'descriptions': [descriptions],
+            'melody_wavs': pmt_wav,
+            'vocal_wavs': vocal_wav,
+            'bgm_wavs': bgm_wav,
+            'melody_is_wav': melody_is_wav,
+        }
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            with torch.no_grad():
+                tokens = model.generate(**generate_inp, return_tokens=True)
+                if offload_audiolm:
+                    offload_profiler.reset_empty_cache_mem_line()
+        item['tokens'] = tokens
+    if offload_audiolm:
+        offload_profiler.stop()
+        del offload_profiler
+        del audiolm_offload_param
+    del model
+    audiolm = audiolm.cpu()
+    del audiolm
+    del checkpoint
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    seperate_tokenizer = builders.get_audio_tokenizer_model_cpu(cfg.audio_tokenizer_checkpoint_sep, cfg)
+    device = "cuda:0"
+    seperate_tokenizer.model.device = device
+    seperate_tokenizer.model.vae = seperate_tokenizer.model.vae.to(device)
+    seperate_tokenizer.model.model.device = torch.device(device)
+    seperate_tokenizer = seperate_tokenizer.eval()
+
+    # offload_wav_tokenizer_diffusion =  True if 'offload' in cfg.keys() and 'wav_tokenizer_diffusion' in cfg.offload else False
+    offload_wav_tokenizer_diffusion =  False
+    if offload_wav_tokenizer_diffusion:
+        sep_offload_param = OffloadParamParse.parse_config(seperate_tokenizer, cfg.offload.wav_tokenizer_diffusion)
+        sep_offload_param.show()
+        sep_offload_profiler = OffloadProfiler(device_index=0, **(sep_offload_param.init_param_dict()))
+        sep_offload_profiler.offload_layer(**(sep_offload_param.offload_layer_param_dict()))
+        sep_offload_profiler.clean_cache_wrapper(**(sep_offload_param.clean_cache_param_dict()))
+    else:
+        seperate_tokenizer.model.model = seperate_tokenizer.model.model.to(device)
+
+    model = CodecLM(name = "tmp",
+        lm = None,
+        audiotokenizer = None,
+        max_duration = max_duration,
+        seperate_tokenizer = seperate_tokenizer,
+    )
+
+    for item in new_items:
+        with torch.no_grad():
+            if 'raw_pmt_wav' in item:
+                if gen_type == 'separate':
+                    wav_seperate = model.generate_audio(item['tokens'], item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'],chunked=True, gen_type='mixed')
+                    wav_vocal = model.generate_audio(item['tokens'],chunked=True, gen_type='vocal')
+                    wav_bgm = model.generate_audio(item['tokens'], chunked=True, gen_type='bgm')
+                elif gen_type == 'mixed':
+                    wav_seperate = model.generate_audio(item['tokens'], item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'],chunked=True, gen_type=gen_type)
+                else:
+                    wav_seperate = model.generate_audio(item['tokens'], chunked=True, gen_type=gen_type)
+                del item['raw_pmt_wav']
+                del item['raw_vocal_wav']
+                del item['raw_bgm_wav']
+            else:
+                if gen_type == 'separate':
+                    wav_vocal = model.generate_audio(item['tokens'], chunked=True, gen_type='vocal')
+                    wav_bgm = model.generate_audio(item['tokens'], chunked=True, gen_type='bgm')
+                    wav_seperate = model.generate_audio(item['tokens'], chunked=True, gen_type='mixed')
+                else:
+                    wav_seperate = model.generate_audio(item['tokens'], chunked=True, gen_type=gen_type)
+        if gen_type == 'separate':
+            torchaudio.save(item['wav_path'].replace('.flac', '_vocal.flac'), wav_vocal[0].cpu().float(), cfg.sample_rate)
+            torchaudio.save(item['wav_path'].replace('.flac', '_bgm.flac'), wav_bgm[0].cpu().float(), cfg.sample_rate)
+            torchaudio.save(item['wav_path'], wav_seperate[0].cpu().float(), cfg.sample_rate)
+        else:
+            torchaudio.save(item['wav_path'], wav_seperate[0].cpu().float(), cfg.sample_rate)
+        del item['tokens']
+        del item['pmt_wav']
+        del item['vocal_wav']
+        del item['bgm_wav']
+        del item['melody_is_wav']
+        if offload_wav_tokenizer_diffusion:
+            sep_offload_profiler.reset_empty_cache_mem_line()
+    
+    if offload_wav_tokenizer_diffusion:
+        sep_offload_profiler.stop()
+    torch.cuda.empty_cache()
+    src_jsonl_name = os.path.split(input_jsonl)[-1]
+    with open(f"{save_dir}/jsonl/{src_jsonl_name}.jsonl", "w", encoding='utf-8') as fw:
+        for item in new_items:
+            fw.writelines(json.dumps(item, ensure_ascii=False)+"\n")
+
+
+if __name__ == "__main__":
+    torch.backends.cudnn.enabled = False
+    OmegaConf.register_new_resolver("eval", lambda x: eval(x))
+    OmegaConf.register_new_resolver("concat", lambda *x: [xxx for xx in x for xxx in xx])
+    OmegaConf.register_new_resolver("get_fname", lambda: os.path.splitext(os.path.basename(sys.argv[1]))[0])
+    OmegaConf.register_new_resolver("load_yaml", lambda x: list(OmegaConf.load(x)))
+    np.random.seed(int(time.time()))
+    # Parse command line arguments
+    args = parse_args()
+    if torch.cuda.is_available():
+        device = torch.cuda.current_device()
+        reserved = torch.cuda.memory_reserved(device)
+        total = torch.cuda.get_device_properties(device).total_memory
+        res_mem = (total - reserved) / 1024 / 1024 / 1024
+        print(f"reserved memory: {res_mem}GB")
+
+        model_name = args.ckpt_path.split("/")[-1].lower().replace('-', '_')
+        assert model_name in ['songgeneration_base', 'songgeneration_base_new', 'songgeneration_base_full', 'songgeneration_large'], f'{model_name} is not supported, currently only songgeneration_base, songgeneration_base_new, songgeneration_base_full, songgeneration_large are supported. Please download correct files and rename the folder to the corresponding version name.'
+        if model_name == 'songgeneration_base' or model_name == 'songgeneration_base_new' or model_name == 'songgeneration_base_full':
+            if res_mem > 24 and not args.low_mem:
+                print("use generate")
+                generate(args)
+            else:
+                from codeclm.utils.offload_profiler import OffloadProfiler, OffloadParamParse
+                print("use generate_lowmem")
+                generate_lowmem(args)
+        elif model_name == 'songgeneration_large':
+            if res_mem > 36 and not args.low_mem:
+                print("use generate")
+                generate(args)
+            else:                
+                print("use generate_lowmem")   
+                from codeclm.utils.offload_profiler import OffloadProfiler, OffloadParamParse
+                generate_lowmem(args)
+            
+
+        # elif model_name == 'songgeneration_base_full':
+
+    else:
+        print("CUDA is not available")
+        exit()
+

baseline_generate/mureka_o2/generate.py

@@ -0,0 +1,390 @@
+#!/usr/bin/env python3
+"""
+Script for batch generating songs using Mureka API
+Processes the first 100 songs in cleaned_data_no_desc.json file
+"""
+
+import json
+import os
+import time
+import requests
+from typing import Dict, List, Optional
+from pathlib import Path
+
+# API Configuration
+API_URL = "https://api.mureka.cn/v1/song/generate"
+QUERY_API_URL = "https://api.mureka.cn/v1/song/query"
+API_KEY_ENV = "MUREKA_API_KEY"
+MODEL = "mureka-o2"
+
+# Configuration Parameters
+MAX_SONGS = 100
+RETRY_TIMES = 3
+RETRY_DELAY = 2  # seconds
+REQUEST_DELAY = 60  # Delay between requests (seconds) - set to 60 seconds (1 minute) to avoid rate limiting
+RATE_LIMIT_DELAY = 60  # Wait time when encountering 429 error (seconds)
+QUERY_INTERVAL = 10  # Interval for querying task status (seconds)
+MAX_QUERY_TIME = 3600  # Maximum query time (seconds), 1 hour
+
+def load_songs(json_file: str, max_count: int = MAX_SONGS) -> List[Dict]:
+    """Load song data from JSON file"""
+    with open(json_file, 'r', encoding='utf-8') as f:
+        data = json.load(f)
+    
+    # Only take first max_count songs
+    return data[:max_count]
+
+def is_song_processed(output_file: Path) -> bool:
+    """
+    Check if song has been processed (including completed tasks)
+    
+    Args:
+        output_file: Output file path
+    
+    Returns:
+        True if file exists and contains valid API response
+    """
+    if not output_file.exists():
+        return False
+    
+    try:
+        with open(output_file, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+            # Check if contains api_response field
+            if 'api_response' in data and data['api_response']:
+                status = data['api_response'].get('status', '')
+                # If status is succeeded, failed, timeouted or cancelled, consider as processed
+                if status in ['succeeded', 'failed', 'timeouted', 'cancelled']:
+                    return True
+                # If status is preparing, queued, running, streaming, reviewing, also consider as processed (task created)
+                if status in ['preparing', 'queued', 'running', 'streaming', 'reviewing']:
+                    return True
+    except (json.JSONDecodeError, KeyError, IOError):
+        # File corrupted or format incorrect, consider as not processed
+        return False
+    
+    return False
+
+def load_processed_song(output_file: Path) -> Optional[Dict]:
+    """
+    Load processed song results from existing file
+    
+    Args:
+        output_file: Output file path
+    
+    Returns:
+        Processed song data, returns None if loading fails
+    """
+    try:
+        with open(output_file, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+            return data.get('api_response')
+    except (json.JSONDecodeError, KeyError, IOError):
+        return None
+
+def query_task_status(task_id: str, api_key: str) -> Optional[Dict]:
+    """
+    Query task status
+    
+    Args:
+        task_id: Task ID
+        api_key: API key
+    
+    Returns:
+        Task status data, returns None on failure
+    """
+    headers = {
+        "Authorization": f"Bearer {api_key}"
+    }
+    
+    url = f"{QUERY_API_URL}/{task_id}"
+    
+    try:
+        response = requests.get(url, headers=headers, timeout=30)
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        print(f"    Failed to query task status: {str(e)}")
+        return None
+
+def wait_for_task_completion(task_id: str, api_key: str) -> Optional[Dict]:
+    """
+    Wait for task completion and return final result
+    
+    Args:
+        task_id: Task ID
+        api_key: API key
+    
+    Returns:
+        Complete data after task completion, returns None on failure
+    """
+    start_time = time.time()
+    last_status = None
+    
+    print(f"  Waiting for task completion (Task ID: {task_id})...")
+    
+    while time.time() - start_time < MAX_QUERY_TIME:
+        result = query_task_status(task_id, api_key)
+        
+        if not result:
+            time.sleep(QUERY_INTERVAL)
+            continue
+        
+        status = result.get('status', '')
+        
+        # If status changed, print new status
+        if status != last_status:
+            print(f"    Status: {status}")
+            last_status = status
+        
+        # Task completed (success or failure)
+        if status in ['succeeded', 'failed', 'timeouted', 'cancelled']:
+            if status == 'succeeded':
+                print(f"  ✓ Task completed!")
+                if 'choices' in result and result['choices']:
+                    print(f"    Found {len(result['choices'])} generated songs")
+            else:
+                print(f"  ✗ Task failed: {status}")
+                if 'failed_reason' in result:
+                    print(f"    Failure reason: {result['failed_reason']}")
+            return result
+        
+        # Task still processing, continue waiting
+        time.sleep(QUERY_INTERVAL)
+    
+    print(f"  ⚠ Query timeout (exceeded {MAX_QUERY_TIME} seconds)")
+    return None
+
+def generate_song(lyrics: str, prompt: str, api_key: str) -> Optional[Dict]:
+    """
+    Call API to generate a single song (serial processing, ensuring concurrency = 1)
+    
+    Args:
+        lyrics: Lyrics content
+        prompt: Prompt (corresponds to description)
+        api_key: API key
+    
+    Returns:
+        API response data, returns None on failure
+    """
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    
+    payload = {
+        "lyrics": lyrics,
+        "model": MODEL,
+        "prompt": prompt
+    }
+    
+    for attempt in range(RETRY_TIMES):
+        try:
+            response = requests.post(API_URL, headers=headers, json=payload, timeout=300)
+            
+            # Check if it's a 429 error (rate limit)
+            if response.status_code == 429:
+                # Try to get Retry-After time from response header
+                retry_after = response.headers.get('Retry-After')
+                if retry_after:
+                    wait_time = int(retry_after)
+                else:
+                    wait_time = RATE_LIMIT_DELAY
+                
+                print(f"  Attempt {attempt + 1}/{RETRY_TIMES} failed: 429 Too Many Requests")
+                print(f"  Waiting {wait_time} seconds before retry...")
+                if attempt < RETRY_TIMES - 1:
+                    time.sleep(wait_time)
+                    continue
+                else:
+                    print(f"  All retries failed")
+                    return None
+            
+            response.raise_for_status()
+            return response.json()
+            
+        except requests.exceptions.HTTPError as e:
+            if e.response and e.response.status_code == 429:
+                # 429 error already handled above
+                continue
+            print(f"  Attempt {attempt + 1}/{RETRY_TIMES} failed: {str(e)}")
+            if attempt < RETRY_TIMES - 1:
+                time.sleep(RETRY_DELAY)
+            else:
+                print(f"  All retries failed")
+                return None
+        except requests.exceptions.RequestException as e:
+            print(f"  Attempt {attempt + 1}/{RETRY_TIMES} failed: {str(e)}")
+            if attempt < RETRY_TIMES - 1:
+                time.sleep(RETRY_DELAY)
+            else:
+                print(f"  All retries failed")
+                return None
+    
+    return None
+
+def main():
+    """Main function"""
+    # Check API key
+    api_key = os.getenv(API_KEY_ENV)
+    if not api_key:
+        print(f"Error: Please set environment variable {API_KEY_ENV}")
+        print(f"Example: export {API_KEY_ENV}=your_api_key")
+        return
+    
+    # Load song data
+    json_file = "cleaned_data_no_desc.json"
+    if not os.path.exists(json_file):
+        print(f"Error: File not found {json_file}")
+        return
+    
+    print(f"Loading song data from {json_file}...")
+    songs = load_songs(json_file, MAX_SONGS)
+    print(f"Loaded {len(songs)} songs")
+    
+    # Create output directory
+    output_dir = Path("generated_songs")
+    output_dir.mkdir(exist_ok=True)
+    
+    # List to save results
+    results = []
+    
+    # Process each song (serial processing, ensuring concurrency = 1)
+    for idx, song in enumerate(songs, 1):
+        print(f"\n[{idx}/{len(songs)}] Processing song...")
+        print(f"  Description: {song.get('description', 'N/A')[:50]}...")
+        
+        # Check output file path
+        output_file = output_dir / f"song_{idx:03d}.json"
+        
+        # Check if already processed
+        if is_song_processed(output_file):
+            print(f"  ⊙ Already processed, skipping")
+            # Load processed results from file
+            existing_result = load_processed_song(output_file)
+            results.append({
+                "index": idx,
+                "status": "already_processed",
+                "output_file": str(output_file)
+            })
+            # Already processed songs don't need delay
+            continue
+        
+        lyrics = song.get('lyrics', '')
+        description = song.get('description', '')
+        
+        if not lyrics or not description:
+            print(f"  Skipping: missing lyrics or description")
+            results.append({
+                "index": idx,
+                "status": "skipped",
+                "reason": "missing data"
+            })
+            continue
+        
+        # Call API (serial execution, ensuring concurrency = 1)
+        result = generate_song(lyrics, description, api_key)
+        
+        if result:
+            task_id = result.get('id')
+            initial_status = result.get('status', '')
+            print(f"  ✓ Task created (ID: {task_id}, Status: {initial_status})")
+            
+            # If task status is not final, wait for task completion
+            if initial_status not in ['succeeded', 'failed', 'timeouted', 'cancelled']:
+                final_result = wait_for_task_completion(task_id, api_key)
+                if final_result:
+                    result = final_result
+                else:
+                    # Query timeout, use initial result
+                    print(f"  ⚠ Using initial result (query timeout)")
+            
+            # Save single result (including final status and choices)
+            with open(output_file, 'w', encoding='utf-8') as f:
+                json.dump({
+                    "index": idx,
+                    "original_data": song,
+                    "api_response": result,
+                    "task_id": task_id
+                }, f, ensure_ascii=False, indent=2)
+            
+            # Check if successfully completed
+            final_status = result.get('status', '')
+            if final_status == 'succeeded':
+                results.append({
+                    "index": idx,
+                    "status": "success",
+                    "output_file": str(output_file),
+                    "task_id": task_id,
+                    "has_audio": 'choices' in result and len(result.get('choices', [])) > 0
+                })
+            elif final_status in ['failed', 'timeouted', 'cancelled']:
+                results.append({
+                    "index": idx,
+                    "status": "failed",
+                    "output_file": str(output_file),
+                    "task_id": task_id,
+                    "failed_reason": result.get('failed_reason', final_status)
+                })
+            else:
+                # Task still processing
+                results.append({
+                    "index": idx,
+                    "status": "processing",
+                    "output_file": str(output_file),
+                    "task_id": task_id,
+                    "current_status": final_status
+                })
+        else:
+            print(f"  ✗ Generation failed")
+            # Save failure information, including original data
+            error_file = output_dir / f"song_{idx:03d}_error.json"
+            with open(error_file, 'w', encoding='utf-8') as f:
+                json.dump({
+                    "index": idx,
+                    "original_data": song,
+                    "error": "API call failed, generate_song returned None",
+                    "timestamp": time.time()
+                }, f, ensure_ascii=False, indent=2)
+            
+            results.append({
+                "index": idx,
+                "status": "failed",
+                "error_file": str(error_file),
+                "reason": "API call failed"
+            })
+        
+        # Delay between requests to avoid rate limiting (ensuring concurrency = 1)
+        if idx < len(songs):
+            print(f"  Waiting {REQUEST_DELAY} seconds before processing next song...")
+            time.sleep(REQUEST_DELAY)
+    
+    # Save summary results
+    summary_file = output_dir / "summary.json"
+    with open(summary_file, 'w', encoding='utf-8') as f:
+        json.dump({
+            "total": len(songs),
+            "success": sum(1 for r in results if r.get("status") == "success"),
+            "processing": sum(1 for r in results if r.get("status") == "processing"),
+            "already_processed": sum(1 for r in results if r.get("status") == "already_processed"),
+            "failed": sum(1 for r in results if r.get("status") == "failed"),
+            "skipped": sum(1 for r in results if r.get("status") == "skipped"),
+            "results": results
+        }, f, ensure_ascii=False, indent=2)
+    
+    # Print statistics
+    print(f"\n{'='*50}")
+    print(f"Processing complete!")
+    print(f"Total: {len(songs)} songs")
+    print(f"Successfully completed: {sum(1 for r in results if r.get('status') == 'success')} songs")
+    print(f"Processing: {sum(1 for r in results if r.get('status') == 'processing')} songs")
+    print(f"Already processed: {sum(1 for r in results if r.get('status') == 'already_processed')} songs")
+    print(f"Failed: {sum(1 for r in results if r.get('status') == 'failed')} songs")
+    print(f"Skipped: {sum(1 for r in results if r.get('status') == 'skipped')} songs")
+    print(f"Results saved in: {output_dir}/")
+    print(f"Summary file: {summary_file}")
+    print(f"\nTip: If tasks are still processing, you can rerun the script later to check status")
+
+if __name__ == "__main__":
+    main()
+

baseline_generate/suno/config.py

@@ -0,0 +1,70 @@
+"""
+Configuration file example
+Copy this file as config.py and fill in your actual configuration
+"""
+
+# ============== API Configuration ==============
+
+# Suno API key (obtain from https://sunoapi.org)
+SUNO_API_KEY = ""
+
+# API base URL (usually no need to modify)
+SUNO_API_BASE_URL = "https://api.sunoapi.org"
+
+# ============== Generation Configuration ==============
+
+# Default model version
+DEFAULT_MODEL_VERSION = "V5"  # Options: V3_5, V4, V4_5, V4_5PLUS, V5
+
+# Whether to enable custom mode
+DEFAULT_CUSTOM_MODE = True
+
+# Whether to generate instrumental by default
+DEFAULT_INSTRUMENTAL = False
+
+# ============== Task Configuration ==============
+
+# Maximum wait time (seconds)
+MAX_WAIT_TIME = 300
+
+# Check interval (seconds)
+CHECK_INTERVAL = 10
+
+# Retry count
+MAX_RETRIES = 3
+
+# ============== File Configuration ==============
+
+# Music file save directory
+OUTPUT_DIRECTORY = "./generated_music"
+
+# Audio format
+AUDIO_FORMAT = "mp3"  # Options: mp3, wav
+
+# ============== Batch Generation Configuration ==============
+
+# Concurrency for batch generation
+BATCH_CONCURRENCY = 5
+
+# Batch generation delay (seconds, to avoid rate limiting)
+BATCH_DELAY = 2
+
+# ============== Logging Configuration ==============
+
+# Log level
+LOG_LEVEL = "INFO"  # Options: DEBUG, INFO, WARNING, ERROR
+
+# Log file path
+LOG_FILE = "./suno_api.log"
+
+# Whether to output to console
+LOG_TO_CONSOLE = True
+
+# ============== Webhook Configuration ==============
+
+# Webhook callback URL (optional)
+WEBHOOK_URL = None  # Example: "https://your-domain.com/webhook"
+
+# Webhook secret (for verifying callback requests)
+WEBHOOK_SECRET = None
+

baseline_generate/suno/suno_4_5.py

@@ -0,0 +1,766 @@
+# -*- coding: utf-8 -*-
+"""
+Suno API Batch Generation - V4.5 Special Edition
+Supported models: V4_5 (default), V4_5PLUS, V4_5ALL
+"""
+import json
+import time
+import requests
+import os
+import logging
+import csv
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from collections import deque
+from threading import Lock, Semaphore
+from tqdm import tqdm
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import SUNO_API_KEY
+
+
+# Configure logging
+def setup_logging(output_dir):
+    log_file = os.path.join(output_dir, f"run_log_v4_5_{time.strftime('%Y%m%d_%H%M%S')}.txt")
+    
+    # Create logger
+    logger = logging.getLogger('SunoBatchV4_5')
+    logger.setLevel(logging.INFO)
+    
+    # Clear old handlers
+    if logger.hasHandlers():
+        logger.handlers.clear()
+        
+    # File Handler
+    file_handler = logging.FileHandler(log_file, encoding='utf-8')
+    file_handler.setFormatter(logging.Formatter('%(message)s'))
+    logger.addHandler(file_handler)
+    
+    # Console Handler
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(logging.Formatter('%(message)s'))
+    logger.addHandler(console_handler)
+    
+    return logger, log_file
+
+# Global logger
+logger = logging.getLogger('SunoBatchV4_5')
+
+# Replace print with logger.info
+def print_log(msg):
+    logger.info(msg)
+
+
+class SunoAPI:
+    """Simplified Suno API client"""
+    
+    def __init__(self, api_key):
+        self.api_key = api_key
+        self.base_url = 'https://api.sunoapi.org/api/v1'
+        self.headers = {
+            'Authorization': f'Bearer {api_key}',
+            'Content-Type': 'application/json'
+        }
+        
+        # Configure retry strategy
+        self.session = requests.Session()
+        retry_strategy = Retry(
+            total=5,  # Maximum retry count
+            backoff_factor=1,  # Retry interval (1s, 2s, 4s, 8s...)
+            status_forcelist=[500, 502, 503, 504],  # Status codes that need retry
+            allowed_methods=["HEAD", "GET", "POST", "OPTIONS"]  # Allowed retry methods
+        )
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        self.session.mount("https://", adapter)
+        self.session.mount("http://", adapter)
+    
+    def generate_music(self, prompt, model='V4_5', vocalGender=None, **options):
+        """Generate music"""
+        payload = {
+            'prompt': prompt,
+            'model': model,
+            'callBackUrl': 'https://example.com/callback',
+            **options
+        }
+        
+        if vocalGender:
+            payload['vocalGender'] = vocalGender
+        
+        try:
+            response = self.session.post(
+                f'{self.base_url}/generate',
+                headers=self.headers,
+                json=payload,
+                timeout=30
+            )
+            
+            # Check HTTP errors
+            response.raise_for_status()
+            
+            # Try to parse JSON
+            try:
+                result = response.json()
+            except json.JSONDecodeError:
+                raise Exception(f"API returned non-JSON response: {response.text[:200]}")
+                
+            if result.get('code') != 200:
+                raise Exception(f"Generation failed: {result.get('msg', result)}")
+            
+            return result['data']['taskId']
+            
+        except requests.exceptions.RequestException as e:
+            raise Exception(f"Request exception: {str(e)}")
+    
+    def get_task_status(self, task_id):
+        """Get task status"""
+        try:
+            response = self.session.get(
+                f'{self.base_url}/generate/record-info?taskId={task_id}',
+                headers={'Authorization': f'Bearer {self.api_key}'},
+                timeout=30
+            )
+            response.raise_for_status()
+            return response.json().get('data', {})
+        except Exception as e:
+            # Status query failure should not crash the program, return empty dict or throw specific exception
+            # print_log(f"Failed to get status: {e}")
+            raise e
+    
+    def get_timestamped_lyrics(self, task_id, audio_id):
+        """Get timestamped lyrics"""
+        payload = {
+            'taskId': task_id,
+            'audioId': audio_id
+        }
+        
+        try:
+            response = self.session.post(
+                f'{self.base_url}/generate/get-timestamped-lyrics',
+                headers=self.headers,
+                json=payload,
+                timeout=30
+            )
+            response.raise_for_status()
+            return response.json()
+        except Exception:
+            return {}  # Lyrics retrieval failure is non-fatal error
+    
+    def wait_for_completion(self, task_id, max_wait_time=600, check_interval=5):
+        """Wait for task completion, return result and polling statistics"""
+        start_time = time.time()
+        poll_count = 0
+        total_poll_time = 0
+        
+        while time.time() - start_time < max_wait_time:
+            try:
+                poll_start = time.time()
+                status = self.get_task_status(task_id)
+                poll_time = time.time() - poll_start
+                poll_count += 1
+                total_poll_time += poll_time
+                
+                current_status = status.get('status')
+                
+                if current_status == 'SUCCESS':
+                    return {
+                        'result': status.get('response'),
+                        'wait_time': time.time() - start_time,
+                        'poll_count': poll_count,
+                        'avg_poll_time': total_poll_time / poll_count if poll_count > 0 else 0
+                    }
+                elif current_status == 'FAILED':
+                    raise Exception(f"Task failed: {status.get('errorMessage')}")
+                
+                time.sleep(check_interval)
+            except Exception as e:
+                if time.time() - start_time >= max_wait_time:
+                    raise
+                time.sleep(check_interval)
+        
+        raise Exception('Task timeout')
+    
+    def download_file(self, url, save_path):
+        """Download file to local, return download statistics"""
+        try:
+            start_time = time.time()
+            downloaded_bytes = 0
+            
+            # Use session to download
+            with self.session.get(url, stream=True, timeout=60) as r:
+                r.raise_for_status()
+                with open(save_path, 'wb') as f:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        downloaded_bytes += len(chunk)
+            
+            download_time = time.time() - start_time
+            return {
+                'success': True,
+                'bytes': downloaded_bytes,
+                'time': download_time,
+                'speed': downloaded_bytes / download_time if download_time > 0 else 0
+            }
+        except Exception as e:
+            print_log(f"Download failed {url}: {e}")
+            return {'success': False, 'error': str(e)}
+
+
+# Result record lock
+result_lock = Lock()
+
+def save_result_record(output_dir, record):
+    """Save single result to CSV in real-time"""
+    file_path = os.path.join(output_dir, "generation_results.csv")
+    file_exists = os.path.isfile(file_path)
+    
+    # Only record key information
+    row = {
+        'song_id': record.get('song_id'),
+        'task_id': record.get('task_id'),
+        'status': 'SUCCESS' if record.get('success') else 'FAILED',
+        'error': record.get('error', ''),
+        'submit_time': time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(record.get('submit_time', 0))),
+        'total_time': f"{record.get('total_time', 0):.1f}",
+        'tracks_count': record.get('tracks_count', 0)
+    }
+    
+    with result_lock:
+        with open(file_path, 'a', newline='', encoding='utf-8') as f:
+            writer = csv.DictWriter(f, fieldnames=['song_id', 'task_id', 'status', 'error', 'submit_time', 'total_time', 'tracks_count'])
+            if not file_exists:
+                writer.writeheader()
+            writer.writerow(row)
+
+
+class ImprovedRateLimiter:
+    """Improved rate limiter (with statistics)
+    
+    Precise control: maximum 8 requests per 10 seconds
+    Uses sliding window algorithm to ensure no more than 8 requests in any 10-second time window
+    """
+    
+    def __init__(self, max_requests=5, time_window=10):
+        self.max_requests = max_requests
+        self.time_window = time_window
+        self.request_times = deque()
+        self.lock = Lock()
+        self.semaphore = Semaphore(max_requests)
+        
+        # Statistics
+        self.total_wait_time = 0
+        self.wait_count = 0
+        self.total_requests = 0
+    
+    def acquire(self):
+        """Acquire request permission"""
+        with self.lock:
+            now = time.time()
+            
+            # Clean expired request records
+            while self.request_times and now - self.request_times[0] >= self.time_window:
+                self.request_times.popleft()
+            
+            # If limit reached, calculate wait time needed
+            wait_time = 0
+            if len(self.request_times) >= self.max_requests:
+                oldest_request = self.request_times[0]
+                wait_time = self.time_window - (now - oldest_request) + 0.05  # Add buffer
+                
+                if wait_time > 0:
+                    print_log(f"  [Rate Limit] Waiting {wait_time:.2f} seconds...")
+                    time.sleep(wait_time)
+                    
+                    # Record wait time
+                    self.total_wait_time += wait_time
+                    self.wait_count += 1
+                    
+                    # Re-clean
+                    now = time.time()
+                    while self.request_times and now - self.request_times[0] >= self.time_window:
+                        self.request_times.popleft()
+            
+            # Record this request time
+            self.request_times.append(time.time())
+            self.total_requests += 1
+            
+    def get_current_rate(self):
+        """Get current rate (number of requests in last 10 seconds)"""
+        with self.lock:
+            now = time.time()
+            while self.request_times and now - self.request_times[0] >= self.time_window:
+                self.request_times.popleft()
+            return len(self.request_times)
+    
+    def get_stats(self):
+        """Get statistics"""
+        with self.lock:
+            return {
+                'total_requests': self.total_requests,
+                'total_wait_time': self.total_wait_time,
+                'wait_count': self.wait_count,
+                'avg_wait_time': self.total_wait_time / self.wait_count if self.wait_count > 0 else 0
+            }
+
+
+# Global rate limiter (5 requests per 10 seconds)
+rate_limiter = ImprovedRateLimiter(max_requests=5, time_window=10)
+
+
+def submit_generation_task(api, song_index, data):
+    """Phase 1: Submit generation task (rate limited)"""
+    # Use sunov4_5_000001 format
+    song_id = data.get("id", f"sunov4_5_{song_index:06d}")
+    
+    try:
+        description = data.get("description", "")
+        lyrics = data.get("lyrics", "")
+        vocal_gender = data.get("vocalGender")
+        
+        print_log(f"[Song {song_id}] Submitting task... (current rate: {rate_limiter.get_current_rate()}/5)")
+        
+        # Record request start time
+        request_start = time.time()
+        
+        # Rate limiting
+        rate_limiter.acquire()
+        
+        # Submit task
+        submit_start = time.time()
+        task_id = api.generate_music(
+            prompt=lyrics,
+            style=description,
+            title=f"Song_{song_id}",
+            model='V4_5', # Explicitly specify V4.5 model
+            customMode=True,
+            instrumental=False,
+            vocalGender=vocal_gender
+        )
+        request_time = time.time() - submit_start
+        
+        print_log(f"[Song {song_id}] ✓ Task submitted, ID: {task_id}")
+        
+        return {
+            'song_id': song_id,
+            'song_index': song_index,
+            'task_id': task_id,
+            'data': data,
+            'submit_time': time.time(),
+            'request_time': request_time,
+            'success': True
+        }
+        
+    except Exception as e:
+        print_log(f"[Song {song_id}] ✗ Submission failed: {e}")
+        # If submission fails, also record it (even though not at download stage yet)
+        return {
+            'song_id': song_id,
+            'song_index': song_index,
+            'success': False,
+            'error': str(e)
+        }
+
+
+def wait_and_download_result(api, task_info, output_dir):
+    """Phase 2: Wait for result and download (not rate limited)"""
+    if not task_info['success']:
+        return task_info
+    
+    song_id = task_info['song_id']
+    song_index = task_info['song_index']
+    task_id = task_info['task_id']
+    data = task_info['data']
+    start_time = task_info['submit_time']
+    
+    try:
+        original_lyrics = data.get("original_lyrics", data.get("lyrics", ""))
+        lyrics = data.get("lyrics", "")
+        description = data.get("description", "")
+        
+        print_log(f"[Song {song_id}] Waiting for generation to complete...")
+        
+        # Wait for completion (returns detailed statistics)
+        wait_result = api.wait_for_completion(task_id, max_wait_time=600, check_interval=8)
+        result = wait_result['result']
+        
+        # Process returned result
+        tracks = []
+        if isinstance(result, dict):
+            if 'data' in result:
+                tracks = result['data']
+            elif 'sunoData' in result:
+                tracks = result['sunoData']
+            else:
+                for key, value in result.items():
+                    if isinstance(value, list) and len(value) > 0 and 'audioUrl' in value[0]:
+                        tracks = value
+                        break
+        
+        if not tracks:
+            raise Exception("Audio track data not found")
+        
+        # Download phase statistics
+        download_start = time.time()
+        downloaded_files = []
+        total_download_bytes = 0
+        download_count = 0
+        
+        # Process each track
+        for track_idx, track in enumerate(tracks):
+            audio_url = track.get('audioUrl') or track.get('audio_url')
+            audio_id = track.get('id')
+            
+            base_filename = f"{song_id}_{track_idx}"
+            audio_path = os.path.join(output_dir, f"{base_filename}.mp3")
+            lyrics_path = os.path.join(output_dir, f"{base_filename}_lyrics.json")
+            
+            # Download audio
+            if audio_url:
+                download_result = api.download_file(audio_url, audio_path)
+                if download_result['success']:
+                    downloaded_files.append(audio_path)
+                    total_download_bytes += download_result['bytes']
+                    download_count += 1
+            
+            # Get timestamped lyrics
+            timestamped_lyrics_data = None
+            if audio_id:
+                try:
+                    lyrics_response = api.get_timestamped_lyrics(task_id, audio_id)
+                    if lyrics_response.get('code') == 200:
+                        timestamped_lyrics_data = lyrics_response.get('data')
+                except Exception as e:
+                    print_log(f"[Song {song_id}] Track {track_idx+1}: Failed to get lyrics: {e}")
+            
+            # Save lyrics and metadata
+            lyrics_content = {
+                "song_id": song_id,
+                "song_index": song_index,
+                "track_index": track_idx,
+                "original_lyrics": original_lyrics,
+                "cleaned_lyrics": lyrics,
+                "timestamped_lyrics": timestamped_lyrics_data,
+                "style": description,
+                "full_track_data": track
+            }
+            
+            with open(lyrics_path, 'w', encoding='utf-8') as f:
+                json.dump(lyrics_content, f, ensure_ascii=False, indent=2)
+            downloaded_files.append(lyrics_path)
+        
+        download_time = time.time() - download_start
+        total_time = time.time() - start_time
+        
+        print_log(f"[Song {song_id}] ✓ Complete! {len(tracks)} tracks, took {total_time:.1f} seconds")
+        
+        final_result = {
+            'song_id': song_id,
+            'song_index': song_index,
+            'task_id': task_id,
+            'success': True,
+            'tracks_count': len(tracks),
+            'files': downloaded_files,
+            'total_time': total_time,
+            'submit_time': start_time,
+            'wait_time': wait_result['wait_time'],
+            'poll_count': wait_result['poll_count'],
+            'avg_poll_time': wait_result['avg_poll_time'],
+            'download_time': download_time,
+            'download_bytes': total_download_bytes,
+            'download_count': download_count,
+            'avg_download_speed': total_download_bytes / download_time if download_time > 0 else 0
+        }
+        
+        # Save result in real-time
+        save_result_record(output_dir, final_result)
+        return final_result
+        
+    except Exception as e:
+        total_time = time.time() - start_time
+        print_log(f"[Song {song_id}] ✗ Processing failed: {e} (took {total_time:.1f} seconds)")
+        
+        error_result = {
+            'song_id': song_id,
+            'song_index': song_index,
+            'task_id': task_id,
+            'success': False,
+            'error': str(e),
+            'total_time': total_time,
+            'submit_time': start_time
+        }
+        
+        # Save result in real-time
+        save_result_record(output_dir, error_result)
+        return error_result
+
+
+def format_bytes(bytes_size):
+    """Format byte size"""
+    for unit in ['B', 'KB', 'MB', 'GB']:
+        if bytes_size < 1024.0:
+            return f"{bytes_size:.2f} {unit}"
+        bytes_size /= 1024.0
+    return f"{bytes_size:.2f} TB"
+
+
+def format_speed(bytes_per_sec):
+    """Format speed"""
+    return f"{format_bytes(bytes_per_sec)}/s"
+
+
+def main():
+    """Main program - two-phase concurrent processing"""
+    input_file = "cleaned_data_truncated.json"
+    output_dir = "sunov4_5_truncated"
+    # Create output directory
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Initialize logging
+    global logger
+    logger, log_file = setup_logging(output_dir)
+    
+    print_log("=" * 70)
+    print_log("Suno API Batch Generation - V4.5 Special Edition")
+    print_log("Strategy: Fast submission (5 requests/10s) + Parallel waiting + Detailed performance analysis")
+    print_log(f"Log file: {log_file}")
+    print_log("=" * 70)
+    
+    # Read input file
+    try:
+        all_data = []
+        if input_file.endswith('.jsonl'):
+            try:
+                with open(input_file, 'r', encoding='utf-8') as f:
+                    # Try reading first line to determine format
+                    first_line = f.readline().strip()
+                    if first_line.startswith('['):
+                        # Looks like regular JSON array
+                        f.seek(0)
+                        all_data = json.load(f)
+                    else:
+                        # Try reading line by line
+                        f.seek(0)
+                        for i, line in enumerate(f):
+                            line = line.strip()
+                            if line:
+                                all_data.append(json.loads(line))
+            except json.JSONDecodeError:
+                # If above parsing fails, try one final read as regular JSON
+                print_log(f"Note: Failed to parse {input_file} as JSONL format, trying as regular JSON...")
+                with open(input_file, 'r', encoding='utf-8') as f:
+                    all_data = json.load(f)
+        else:
+            with open(input_file, 'r', encoding='utf-8') as f:
+                all_data = json.load(f)
+        
+    except FileNotFoundError:
+        print_log(f"File {input_file} not found.")
+        return
+    except json.JSONDecodeError as e:
+        print_log(f"JSON parsing error: {e}")
+        return
+    
+    # Initialize API
+    api = SunoAPI(SUNO_API_KEY)
+    
+    print_log(f"\nPreparing to generate {len(all_data)} songs...")
+    print_log(f"Start time: {time.strftime('%H:%M:%S')}\n")
+    
+    overall_start_time = time.time()
+    
+    # ===== Phase 1: Batch Submission =====
+    print_log("\n" + "=" * 70)
+    print_log("Phase 1: Batch Submission")
+    print_log("=" * 70 + "\n")
+    
+    submit_start_time = time.time()
+    submitted_tasks = []
+    total_request_time = 0
+    
+    # Adjust rate limit: maximum 5 requests per 10 seconds
+    rate_limiter.max_requests = 5
+    rate_limiter.time_window = 10
+    rate_limiter.request_times.clear()
+    print_log(f"Rate limit: {rate_limiter.max_requests} requests / {rate_limiter.time_window} seconds")
+    
+    # Only submit tasks that need to run
+    tasks_to_run = []
+    for i, data in enumerate(all_data, 1):
+        tasks_to_run.append((i, data))
+        
+    print_log(f"Number of tasks to submit: {len(tasks_to_run)}")
+
+    # Use thread pool for submission
+    # Submission concurrency is controlled by rate_limiter, can be set to 5
+    with ThreadPoolExecutor(max_workers=5) as executor:
+        submit_futures = {
+            executor.submit(submit_generation_task, api, idx, data): idx
+            for idx, data in tasks_to_run
+        }
+        
+        with tqdm(total=len(tasks_to_run), desc="Submitting tasks", unit="song") as pbar:
+            for future in as_completed(submit_futures):
+                result = future.result()
+                submitted_tasks.append(result)
+                if result.get('success') and 'request_time' in result:
+                    total_request_time += result['request_time']
+                pbar.update(1)
+    
+    submit_phase_time = time.time() - submit_start_time
+    success_submits = sum(1 for t in submitted_tasks if t['success'])
+    
+    # Get rate limit statistics
+    rate_limit_stats = rate_limiter.get_stats()
+    
+    print_log(f"\nSubmission phase complete: {success_submits}/{len(tasks_to_run)} successful")
+    print_log(f"  Total time: {submit_phase_time:.1f} seconds")
+    print_log(f"  Actual request time: {total_request_time:.2f} seconds")
+    print_log(f"  Rate limit waiting: {rate_limit_stats['total_wait_time']:.2f} seconds ({rate_limit_stats['wait_count']} times)")
+    if rate_limit_stats['wait_count'] > 0:
+        print_log(f"  Average wait time: {rate_limit_stats['avg_wait_time']:.2f} seconds/time")
+    
+    # ===== Phase 2: Parallel Waiting and Download =====
+    print_log("\n" + "=" * 70)
+    print_log("Phase 2: Wait for Generation and Download")
+    print_log("=" * 70 + "\n")
+    
+    wait_start_time = time.time()
+    final_results = []
+    
+    # Use more threads for parallel waiting (not rate limited)
+    with ThreadPoolExecutor(max_workers=20) as executor:
+        download_futures = {
+            executor.submit(wait_and_download_result, api, task, output_dir): task
+            for task in submitted_tasks if task['success']
+        }
+        
+        # Add failed submission tasks to results
+        for task in submitted_tasks:
+            if not task['success']:
+                final_results.append(task)
+        
+        with tqdm(total=len(download_futures), desc="Downloading results", unit="song") as pbar:
+            for future in as_completed(download_futures):
+                result = future.result()
+                final_results.append(result)
+                pbar.update(1)
+    
+    wait_phase_time = time.time() - wait_start_time
+    
+    # ===== Detailed Statistics and Report =====
+    overall_time = time.time() - overall_start_time
+    
+    print_log("\n" + "=" * 70)
+    print_log("Batch Generation Complete - Detailed Performance Report")
+    print_log("=" * 70)
+    
+    success_count = sum(1 for r in final_results if r.get('success'))
+    fail_count = len(final_results) - success_count
+    total_tracks = sum(r.get('tracks_count', 0) for r in final_results if r.get('success'))
+    
+    successful_results = [r for r in final_results if r.get('success')]
+    
+    # Basic Statistics
+    print_log(f"\n[Basic Statistics]")
+    print_log(f"  Total songs: {len(all_data)}")
+    print_log(f"  Successful: {success_count}")
+    print_log(f"  Failed: {fail_count}")
+    print_log(f"  Total tracks: {total_tracks}")
+    if success_count > 0:
+        avg_tracks = total_tracks / success_count
+        print_log(f"  Average tracks per song: {avg_tracks:.2f}")
+    
+    # Time Statistics
+    print_log(f"\n[Time Statistics]")
+    print_log(f"  ├── Submission phase: {submit_phase_time:.1f} seconds")
+    print_log(f"  │   ├── Actual request time: {total_request_time:.2f} seconds")
+    print_log(f"  │   └── Rate limit waiting: {rate_limit_stats['total_wait_time']:.2f} seconds")
+    print_log(f"  ├── Generation waiting phase: {wait_phase_time:.1f} seconds")
+    
+    if successful_results:
+        wait_times = [r.get('wait_time', 0) for r in successful_results if 'wait_time' in r]
+        download_times = [r.get('download_time', 0) for r in successful_results if 'download_time' in r]
+        
+        if wait_times:
+            avg_wait = sum(wait_times) / len(wait_times)
+            min_wait = min(wait_times)
+            max_wait = max(wait_times)
+            print_log(f"  │   ├── Average wait time: {avg_wait:.1f} seconds/song")
+            print_log(f"  │   ├── Fastest: {min_wait:.1f} seconds")
+            print_log(f"  │   └── Slowest: {max_wait:.1f} seconds")
+        
+        if download_times:
+            total_download_time = sum(download_times)
+            avg_download = total_download_time / len(download_times)
+            print_log(f"  ├── Download phase: {total_download_time:.1f} seconds")
+            print_log(f"  │   └── Average download time: {avg_download:.2f} seconds/song")
+    
+    print_log(f"  └── Total time: {overall_time:.1f} seconds ({overall_time/60:.1f} minutes)")
+    
+    # Single Song Generation Statistics
+    if successful_results:
+        total_times = [r.get('total_time', 0) for r in successful_results if 'total_time' in r]
+        if total_times:
+            print_log(f"\n[Single Song Generation Statistics]")
+            avg_time = sum(total_times) / len(total_times)
+            min_time = min(total_times)
+            max_time = max(total_times)
+            print_log(f"  Average total time per song: {avg_time:.1f} seconds")
+            print_log(f"  Fastest generation: {min_time:.1f} seconds")
+            print_log(f"  Slowest generation: {max_time:.1f} seconds")
+    
+    # Download Statistics
+    total_download_bytes = sum(r.get('download_bytes', 0) for r in successful_results)
+    total_download_count = sum(r.get('download_count', 0) for r in successful_results)
+    
+    if total_download_bytes > 0:
+        print_log(f"\n[Download Statistics]")
+        print_log(f"  Total download: {format_bytes(total_download_bytes)}")
+        print_log(f"  Number of files: {total_download_count}")
+        print_log(f"  Average file size: {format_bytes(total_download_bytes / total_download_count)}")
+        
+        download_speeds = [r.get('avg_download_speed', 0) for r in successful_results if r.get('avg_download_speed', 0) > 0]
+        if download_speeds:
+            avg_speed = sum(download_speeds) / len(download_speeds)
+            print_log(f"  Average download speed: {format_speed(avg_speed)}")
+    
+    # Polling Statistics
+    poll_counts = [r.get('poll_count', 0) for r in successful_results if 'poll_count' in r]
+    if poll_counts:
+        total_polls = sum(poll_counts)
+        avg_polls = total_polls / len(poll_counts)
+        print_log(f"\n[Polling Statistics]")
+        print_log(f"  Total polling count: {total_polls}")
+        print_log(f"  Average polling per song: {avg_polls:.1f}")
+    
+    # Efficiency Analysis
+    print_log(f"\n[Efficiency Analysis]")
+    if success_count > 0:
+        throughput = success_count / (overall_time / 60)
+        print_log(f"  Actual throughput: {throughput:.2f} songs/minute")
+        
+        # Theoretical fastest time (assuming no rate limit)
+        if wait_times:
+            ideal_time = submit_phase_time - rate_limit_stats['total_wait_time'] + max(wait_times)
+            efficiency = (ideal_time / overall_time) * 100
+            print_log(f"  Theoretical fastest time: {ideal_time:.1f} seconds")
+            print_log(f"  Concurrency efficiency: {efficiency:.1f}%")
+    
+    # Show failed songs
+    if fail_count > 0:
+        print_log("\n" + "=" * 70)
+        print_log("Failed Songs List")
+        print_log("=" * 70)
+        for r in sorted(final_results, key=lambda x: x.get('song_index', 0)):
+            if not r.get('success'):
+                song_id = r.get('song_id', r.get('song_index', 'Unknown'))
+                print_log(f"  [{song_id}] {r.get('error', 'Unknown error')}")
+    
+    print_log("\n" + "=" * 70)
+    print_log(f"All files saved to: {os.path.abspath(output_dir)}")
+    print_log("=" * 70)
+
+
+if __name__ == '__main__':
+    main()
+

baseline_generate/suno/suno_5.py

@@ -0,0 +1,768 @@
+# -*- coding: utf-8 -*-
+"""
+Suno API Batch Generation - V5 Version (5 requests per 10 seconds)
+Changes:
+1. Rate control: 5 requests within 10 seconds
+2. ID format: sunov5_000001
+"""
+import json
+import time
+import requests
+import os
+import logging
+import csv
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from collections import deque
+from threading import Lock, Semaphore
+from tqdm import tqdm
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import SUNO_API_KEY
+
+
+# Configure logging
+def setup_logging(output_dir):
+    log_file = os.path.join(output_dir, f"run_log_{time.strftime('%Y%m%d_%H%M%S')}.txt")
+    
+    # Create logger
+    logger = logging.getLogger('SunoBatchV5')
+    logger.setLevel(logging.INFO)
+    
+    # Clear old handlers
+    if logger.hasHandlers():
+        logger.handlers.clear()
+        
+    # File Handler
+    file_handler = logging.FileHandler(log_file, encoding='utf-8')
+    file_handler.setFormatter(logging.Formatter('%(message)s'))
+    logger.addHandler(file_handler)
+    
+    # Console Handler
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(logging.Formatter('%(message)s'))
+    logger.addHandler(console_handler)
+    
+    return logger, log_file
+
+# Global logger
+logger = logging.getLogger('SunoBatchV5')
+
+# Replace print with logger.info
+def print_log(msg):
+    logger.info(msg)
+
+
+class SunoAPI:
+    """Simplified Suno API client"""
+    
+    def __init__(self, api_key):
+        self.api_key = api_key
+        self.base_url = 'https://api.sunoapi.org/api/v1'
+        self.headers = {
+            'Authorization': f'Bearer {api_key}',
+            'Content-Type': 'application/json'
+        }
+        
+        # Configure retry strategy
+        self.session = requests.Session()
+        retry_strategy = Retry(
+            total=5,  # Maximum retry count
+            backoff_factor=1,  # Retry interval (1s, 2s, 4s, 8s...)
+            status_forcelist=[500, 502, 503, 504],  # Status codes that need retry
+            allowed_methods=["HEAD", "GET", "POST", "OPTIONS"]  # Allowed retry methods
+        )
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        self.session.mount("https://", adapter)
+        self.session.mount("http://", adapter)
+    
+    def generate_music(self, prompt, model='V5', vocalGender=None, **options):
+        """Generate music"""
+        payload = {
+            'prompt': prompt,
+            'model': model,
+            'callBackUrl': 'https://example.com/callback',
+            **options
+        }
+        
+        if vocalGender:
+            payload['vocalGender'] = vocalGender
+        
+        try:
+            response = self.session.post(
+                f'{self.base_url}/generate',
+                headers=self.headers,
+                json=payload,
+                timeout=30
+            )
+            
+            # Check HTTP errors
+            response.raise_for_status()
+            
+            # Try to parse JSON
+            try:
+                result = response.json()
+            except json.JSONDecodeError:
+                raise Exception(f"API returned non-JSON response: {response.text[:200]}")
+                
+            if result.get('code') != 200:
+                raise Exception(f"Generation failed: {result.get('msg', result)}")
+            
+            return result['data']['taskId']
+            
+        except requests.exceptions.RequestException as e:
+            raise Exception(f"Request exception: {str(e)}")
+    
+    def get_task_status(self, task_id):
+        """Get task status"""
+        try:
+            response = self.session.get(
+                f'{self.base_url}/generate/record-info?taskId={task_id}',
+                headers={'Authorization': f'Bearer {self.api_key}'},
+                timeout=30
+            )
+            response.raise_for_status()
+            return response.json().get('data', {})
+        except Exception as e:
+            # Status query failure should not crash the program, return empty dict or throw specific exception
+            # print_log(f"Failed to get status: {e}")
+            raise e
+    
+    def get_timestamped_lyrics(self, task_id, audio_id):
+        """Get timestamped lyrics"""
+        payload = {
+            'taskId': task_id,
+            'audioId': audio_id
+        }
+        
+        try:
+            response = self.session.post(
+                f'{self.base_url}/generate/get-timestamped-lyrics',
+                headers=self.headers,
+                json=payload,
+                timeout=30
+            )
+            response.raise_for_status()
+            return response.json()
+        except Exception:
+            return {}  # Lyrics retrieval failure is non-fatal error
+    
+    def wait_for_completion(self, task_id, max_wait_time=600, check_interval=5):
+        """Wait for task completion, return result and polling statistics"""
+        start_time = time.time()
+        poll_count = 0
+        total_poll_time = 0
+        
+        while time.time() - start_time < max_wait_time:
+            try:
+                poll_start = time.time()
+                status = self.get_task_status(task_id)
+                poll_time = time.time() - poll_start
+                poll_count += 1
+                total_poll_time += poll_time
+                
+                current_status = status.get('status')
+                
+                if current_status == 'SUCCESS':
+                    return {
+                        'result': status.get('response'),
+                        'wait_time': time.time() - start_time,
+                        'poll_count': poll_count,
+                        'avg_poll_time': total_poll_time / poll_count if poll_count > 0 else 0
+                    }
+                elif current_status == 'FAILED':
+                    raise Exception(f"Task failed: {status.get('errorMessage')}")
+                
+                time.sleep(check_interval)
+            except Exception as e:
+                if time.time() - start_time >= max_wait_time:
+                    raise
+                time.sleep(check_interval)
+        
+        raise Exception('Task timeout')
+    
+    def download_file(self, url, save_path):
+        """Download file to local, return download statistics"""
+        try:
+            start_time = time.time()
+            downloaded_bytes = 0
+            
+            # Use session to download
+            with self.session.get(url, stream=True, timeout=60) as r:
+                r.raise_for_status()
+                with open(save_path, 'wb') as f:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        downloaded_bytes += len(chunk)
+            
+            download_time = time.time() - start_time
+            return {
+                'success': True,
+                'bytes': downloaded_bytes,
+                'time': download_time,
+                'speed': downloaded_bytes / download_time if download_time > 0 else 0
+            }
+        except Exception as e:
+            print_log(f"Download failed {url}: {e}")
+            return {'success': False, 'error': str(e)}
+
+
+# Result record lock
+result_lock = Lock()
+
+def save_result_record(output_dir, record):
+    """Save single result to CSV in real-time"""
+    file_path = os.path.join(output_dir, "generation_results.csv")
+    file_exists = os.path.isfile(file_path)
+    
+    # Only record key information
+    row = {
+        'song_id': record.get('song_id'),
+        'task_id': record.get('task_id'),
+        'status': 'SUCCESS' if record.get('success') else 'FAILED',
+        'error': record.get('error', ''),
+        'submit_time': time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(record.get('submit_time', 0))),
+        'total_time': f"{record.get('total_time', 0):.1f}",
+        'tracks_count': record.get('tracks_count', 0)
+    }
+    
+    with result_lock:
+        with open(file_path, 'a', newline='', encoding='utf-8') as f:
+            writer = csv.DictWriter(f, fieldnames=['song_id', 'task_id', 'status', 'error', 'submit_time', 'total_time', 'tracks_count'])
+            if not file_exists:
+                writer.writeheader()
+            writer.writerow(row)
+
+
+class ImprovedRateLimiter:
+    """Improved rate limiter (with statistics)
+    
+    Precise control: maximum 8 requests per 10 seconds
+    Uses sliding window algorithm to ensure no more than 8 requests in any 10-second time window
+    """
+    
+    def __init__(self, max_requests=5, time_window=10):
+        self.max_requests = max_requests
+        self.time_window = time_window
+        self.request_times = deque()
+        self.lock = Lock()
+        self.semaphore = Semaphore(max_requests)
+        
+        # Statistics
+        self.total_wait_time = 0
+        self.wait_count = 0
+        self.total_requests = 0
+    
+    def acquire(self):
+        """Acquire request permission"""
+        with self.lock:
+            now = time.time()
+            
+            # Clean expired request records
+            while self.request_times and now - self.request_times[0] >= self.time_window:
+                self.request_times.popleft()
+            
+            # If limit reached, calculate wait time needed
+            wait_time = 0
+            if len(self.request_times) >= self.max_requests:
+                oldest_request = self.request_times[0]
+                wait_time = self.time_window - (now - oldest_request) + 0.05  # Add buffer
+                
+                if wait_time > 0:
+                    print_log(f"  [Rate Limit] Waiting {wait_time:.2f} seconds...")
+                    time.sleep(wait_time)
+                    
+                    # Record wait time
+                    self.total_wait_time += wait_time
+                    self.wait_count += 1
+                    
+                    # Re-clean
+                    now = time.time()
+                    while self.request_times and now - self.request_times[0] >= self.time_window:
+                        self.request_times.popleft()
+            
+            # Record this request time
+            self.request_times.append(time.time())
+            self.total_requests += 1
+            
+    def get_current_rate(self):
+        """Get current rate (number of requests in last 10 seconds)"""
+        with self.lock:
+            now = time.time()
+            while self.request_times and now - self.request_times[0] >= self.time_window:
+                self.request_times.popleft()
+            return len(self.request_times)
+    
+    def get_stats(self):
+        """Get statistics"""
+        with self.lock:
+            return {
+                'total_requests': self.total_requests,
+                'total_wait_time': self.total_wait_time,
+                'wait_count': self.wait_count,
+                'avg_wait_time': self.total_wait_time / self.wait_count if self.wait_count > 0 else 0
+            }
+
+
+# Global rate limiter (5 requests per 10 seconds)
+rate_limiter = ImprovedRateLimiter(max_requests=5, time_window=10)
+
+
+def submit_generation_task(api, song_index, data):
+    """Phase 1: Submit generation task (rate limited)"""
+    # Use sunov5_000001 format
+    song_id = data.get("id", f"sunov5_{song_index:06d}")
+    
+    try:
+        description = data.get("description", "")
+        lyrics = data.get("lyrics", "")
+        vocal_gender = data.get("vocalGender")
+        
+        print_log(f"[Song {song_id}] Submitting task... (current rate: {rate_limiter.get_current_rate()}/5)")
+        
+        # Record request start time
+        request_start = time.time()
+        
+        # Rate limiting
+        rate_limiter.acquire()
+        
+        # Submit task
+        submit_start = time.time()
+        task_id = api.generate_music(
+            prompt=lyrics,
+            style=description,
+            title=f"Song_{song_id}",
+            model='V5',
+            customMode=True,
+            instrumental=False,
+            vocalGender=vocal_gender
+        )
+        request_time = time.time() - submit_start
+        
+        print_log(f"[Song {song_id}] ✓ Task submitted, ID: {task_id}")
+        
+        return {
+            'song_id': song_id,
+            'song_index': song_index,
+            'task_id': task_id,
+            'data': data,
+            'submit_time': time.time(),
+            'request_time': request_time,
+            'success': True
+        }
+        
+    except Exception as e:
+        print_log(f"[Song {song_id}] ✗ Submission failed: {e}")
+        # If submission fails, also record it (even though not at download stage yet)
+        return {
+            'song_id': song_id,
+            'song_index': song_index,
+            'success': False,
+            'error': str(e)
+        }
+
+
+def wait_and_download_result(api, task_info, output_dir):
+    """Phase 2: Wait for result and download (not rate limited)"""
+    if not task_info['success']:
+        return task_info
+    
+    song_id = task_info['song_id']
+    song_index = task_info['song_index']
+    task_id = task_info['task_id']
+    data = task_info['data']
+    start_time = task_info['submit_time']
+    
+    try:
+        original_lyrics = data.get("original_lyrics", data.get("lyrics", ""))
+        lyrics = data.get("lyrics", "")
+        description = data.get("description", "")
+        
+        print_log(f"[Song {song_id}] Waiting for generation to complete...")
+        
+        # Wait for completion (returns detailed statistics)
+        wait_result = api.wait_for_completion(task_id, max_wait_time=600, check_interval=8)
+        result = wait_result['result']
+        
+        # Process returned result
+        tracks = []
+        if isinstance(result, dict):
+            if 'data' in result:
+                tracks = result['data']
+            elif 'sunoData' in result:
+                tracks = result['sunoData']
+            else:
+                for key, value in result.items():
+                    if isinstance(value, list) and len(value) > 0 and 'audioUrl' in value[0]:
+                        tracks = value
+                        break
+        
+        if not tracks:
+            raise Exception("Audio track data not found")
+        
+        # Download phase statistics
+        download_start = time.time()
+        downloaded_files = []
+        total_download_bytes = 0
+        download_count = 0
+        
+        # Process each track
+        for track_idx, track in enumerate(tracks):
+            audio_url = track.get('audioUrl') or track.get('audio_url')
+            audio_id = track.get('id')
+            
+            base_filename = f"{song_id}_{track_idx}"
+            audio_path = os.path.join(output_dir, f"{base_filename}.mp3")
+            lyrics_path = os.path.join(output_dir, f"{base_filename}_lyrics.json")
+            
+            # Download audio
+            if audio_url:
+                download_result = api.download_file(audio_url, audio_path)
+                if download_result['success']:
+                    downloaded_files.append(audio_path)
+                    total_download_bytes += download_result['bytes']
+                    download_count += 1
+            
+            # Get timestamped lyrics
+            timestamped_lyrics_data = None
+            if audio_id:
+                try:
+                    lyrics_response = api.get_timestamped_lyrics(task_id, audio_id)
+                    if lyrics_response.get('code') == 200:
+                        timestamped_lyrics_data = lyrics_response.get('data')
+                except Exception as e:
+                    print_log(f"[Song {song_id}] Track {track_idx+1}: Failed to get lyrics: {e}")
+            
+            # Save lyrics and metadata
+            lyrics_content = {
+                "song_id": song_id,
+                "song_index": song_index,
+                "track_index": track_idx,
+                "original_lyrics": original_lyrics,
+                "cleaned_lyrics": lyrics,
+                "timestamped_lyrics": timestamped_lyrics_data,
+                "style": description,
+                "full_track_data": track
+            }
+            
+            with open(lyrics_path, 'w', encoding='utf-8') as f:
+                json.dump(lyrics_content, f, ensure_ascii=False, indent=2)
+            downloaded_files.append(lyrics_path)
+        
+        download_time = time.time() - download_start
+        total_time = time.time() - start_time
+        
+        print_log(f"[Song {song_id}] ✓ Complete! {len(tracks)} tracks, took {total_time:.1f} seconds")
+        
+        final_result = {
+            'song_id': song_id,
+            'song_index': song_index,
+            'task_id': task_id,
+            'success': True,
+            'tracks_count': len(tracks),
+            'files': downloaded_files,
+            'total_time': total_time,
+            'submit_time': start_time,
+            'wait_time': wait_result['wait_time'],
+            'poll_count': wait_result['poll_count'],
+            'avg_poll_time': wait_result['avg_poll_time'],
+            'download_time': download_time,
+            'download_bytes': total_download_bytes,
+            'download_count': download_count,
+            'avg_download_speed': total_download_bytes / download_time if download_time > 0 else 0
+        }
+        
+        # Save result in real-time
+        save_result_record(output_dir, final_result)
+        return final_result
+        
+    except Exception as e:
+        total_time = time.time() - start_time
+        print_log(f"[Song {song_id}] ✗ Processing failed: {e} (took {total_time:.1f} seconds)")
+        
+        error_result = {
+            'song_id': song_id,
+            'song_index': song_index,
+            'task_id': task_id,
+            'success': False,
+            'error': str(e),
+            'total_time': total_time,
+            'submit_time': start_time
+        }
+        
+        # Save result in real-time
+        save_result_record(output_dir, error_result)
+        return error_result
+
+
+def format_bytes(bytes_size):
+    """Format byte size"""
+    for unit in ['B', 'KB', 'MB', 'GB']:
+        if bytes_size < 1024.0:
+            return f"{bytes_size:.2f} {unit}"
+        bytes_size /= 1024.0
+    return f"{bytes_size:.2f} TB"
+
+
+def format_speed(bytes_per_sec):
+    """Format speed"""
+    return f"{format_bytes(bytes_per_sec)}/s"
+
+
+def main():
+    """Main program - two-phase concurrent processing"""
+    input_file = "cleaned_data_truncated.json"
+    output_dir = "sunov5_truncated"
+    # Create output directory
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Initialize logging
+    global logger
+    logger, log_file = setup_logging(output_dir)
+    
+    print_log("=" * 70)
+    print_log("Suno API Batch Generation - V5 Special Edition")
+    print_log("Strategy: Fast submission (5 requests/10s) + Parallel waiting + Detailed performance analysis")
+    print_log(f"Log file: {log_file}")
+    print_log("=" * 70)
+    
+    # Read input file
+    try:
+        all_data = []
+        if input_file.endswith('.jsonl'):
+            try:
+                with open(input_file, 'r', encoding='utf-8') as f:
+                    # Try reading first line to determine format
+                    first_line = f.readline().strip()
+                    if first_line.startswith('['):
+                        # Looks like regular JSON array
+                        f.seek(0)
+                        all_data = json.load(f)
+                    else:
+                        # Try reading line by line
+                        f.seek(0)
+                        for i, line in enumerate(f):
+                            line = line.strip()
+                            if line:
+                                all_data.append(json.loads(line))
+            except json.JSONDecodeError:
+                # If above parsing fails, try one final read as regular JSON
+                print_log(f"Note: Failed to parse {input_file} as JSONL format, trying as regular JSON...")
+                with open(input_file, 'r', encoding='utf-8') as f:
+                    all_data = json.load(f)
+        else:
+            with open(input_file, 'r', encoding='utf-8') as f:
+                all_data = json.load(f)
+        
+    except FileNotFoundError:
+        print_log(f"File {input_file} not found.")
+        return
+    except json.JSONDecodeError as e:
+        print_log(f"JSON parsing error: {e}")
+        return
+    
+    # Initialize API
+    api = SunoAPI(SUNO_API_KEY)
+    
+    print_log(f"\nPreparing to generate {len(all_data)} songs...")
+    print_log(f"Start time: {time.strftime('%H:%M:%S')}\n")
+    
+    overall_start_time = time.time()
+    
+    # ===== Phase 1: Batch Submission =====
+    print_log("\n" + "=" * 70)
+    print_log("Phase 1: Batch Submission")
+    print_log("=" * 70 + "\n")
+    
+    submit_start_time = time.time()
+    submitted_tasks = []
+    total_request_time = 0
+    
+    # Adjust rate limit: maximum 5 requests per 10 seconds
+    rate_limiter.max_requests = 5
+    rate_limiter.time_window = 10
+    rate_limiter.request_times.clear()
+    print_log(f"Rate limit: {rate_limiter.max_requests} requests / {rate_limiter.time_window} seconds")
+    
+    # Only submit tasks that need to run
+    tasks_to_run = []
+    for i, data in enumerate(all_data, 1):
+        tasks_to_run.append((i, data))
+        
+    print_log(f"Number of tasks to submit: {len(tasks_to_run)}")
+
+    # Use thread pool for submission
+    # Submission concurrency is controlled by rate_limiter, can be set to 5
+    with ThreadPoolExecutor(max_workers=5) as executor:
+        submit_futures = {
+            executor.submit(submit_generation_task, api, idx, data): idx
+            for idx, data in tasks_to_run
+        }
+        
+        with tqdm(total=len(tasks_to_run), desc="Submitting tasks", unit="song") as pbar:
+            for future in as_completed(submit_futures):
+                result = future.result()
+                submitted_tasks.append(result)
+                if result.get('success') and 'request_time' in result:
+                    total_request_time += result['request_time']
+                pbar.update(1)
+    
+    submit_phase_time = time.time() - submit_start_time
+    success_submits = sum(1 for t in submitted_tasks if t['success'])
+    
+    # Get rate limit statistics
+    rate_limit_stats = rate_limiter.get_stats()
+    
+    print_log(f"\nSubmission phase complete: {success_submits}/{len(tasks_to_run)} successful")
+    print_log(f"  Total time: {submit_phase_time:.1f} seconds")
+    print_log(f"  Actual request time: {total_request_time:.2f} seconds")
+    print_log(f"  Rate limit waiting: {rate_limit_stats['total_wait_time']:.2f} seconds ({rate_limit_stats['wait_count']} times)")
+    if rate_limit_stats['wait_count'] > 0:
+        print_log(f"  Average wait time: {rate_limit_stats['avg_wait_time']:.2f} seconds/time")
+    
+    # ===== Phase 2: Parallel Waiting and Download =====
+    print_log("\n" + "=" * 70)
+    print_log("Phase 2: Wait for Generation and Download")
+    print_log("=" * 70 + "\n")
+    
+    wait_start_time = time.time()
+    final_results = []
+    
+    # Use more threads for parallel waiting (not rate limited)
+    with ThreadPoolExecutor(max_workers=20) as executor:
+        download_futures = {
+            executor.submit(wait_and_download_result, api, task, output_dir): task
+            for task in submitted_tasks if task['success']
+        }
+        
+        # Add failed submission tasks to results
+        for task in submitted_tasks:
+            if not task['success']:
+                final_results.append(task)
+        
+        with tqdm(total=len(download_futures), desc="Downloading results", unit="song") as pbar:
+            for future in as_completed(download_futures):
+                result = future.result()
+                final_results.append(result)
+                pbar.update(1)
+    
+    wait_phase_time = time.time() - wait_start_time
+    
+    # ===== Detailed Statistics and Report =====
+    overall_time = time.time() - overall_start_time
+    
+    print_log("\n" + "=" * 70)
+    print_log("Batch Generation Complete - Detailed Performance Report")
+    print_log("=" * 70)
+    
+    success_count = sum(1 for r in final_results if r.get('success'))
+    fail_count = len(final_results) - success_count
+    total_tracks = sum(r.get('tracks_count', 0) for r in final_results if r.get('success'))
+    
+    successful_results = [r for r in final_results if r.get('success')]
+    
+    # Basic Statistics
+    print_log(f"\n[Basic Statistics]")
+    print_log(f"  Total songs: {len(all_data)}")
+    print_log(f"  Successful: {success_count}")
+    print_log(f"  Failed: {fail_count}")
+    print_log(f"  Total tracks: {total_tracks}")
+    if success_count > 0:
+        avg_tracks = total_tracks / success_count
+        print_log(f"  Average tracks per song: {avg_tracks:.2f}")
+    
+    # Time Statistics
+    print_log(f"\n[Time Statistics]")
+    print_log(f"  ├── Submission phase: {submit_phase_time:.1f} seconds")
+    print_log(f"  │   ├── Actual request time: {total_request_time:.2f} seconds")
+    print_log(f"  │   └── Rate limit waiting: {rate_limit_stats['total_wait_time']:.2f} seconds")
+    print_log(f"  ├── Generation waiting phase: {wait_phase_time:.1f} seconds")
+    
+    if successful_results:
+        wait_times = [r.get('wait_time', 0) for r in successful_results if 'wait_time' in r]
+        download_times = [r.get('download_time', 0) for r in successful_results if 'download_time' in r]
+        
+        if wait_times:
+            avg_wait = sum(wait_times) / len(wait_times)
+            min_wait = min(wait_times)
+            max_wait = max(wait_times)
+            print_log(f"  │   ├── Average wait time: {avg_wait:.1f} seconds/song")
+            print_log(f"  │   ├── Fastest: {min_wait:.1f} seconds")
+            print_log(f"  │   └── Slowest: {max_wait:.1f} seconds")
+        
+        if download_times:
+            total_download_time = sum(download_times)
+            avg_download = total_download_time / len(download_times)
+            print_log(f"  ├── Download phase: {total_download_time:.1f} seconds")
+            print_log(f"  │   └── Average download time: {avg_download:.2f} seconds/song")
+    
+    print_log(f"  └── Total time: {overall_time:.1f} seconds ({overall_time/60:.1f} minutes)")
+    
+    # Single Song Generation Statistics
+    if successful_results:
+        total_times = [r.get('total_time', 0) for r in successful_results if 'total_time' in r]
+        if total_times:
+            print_log(f"\n[Single Song Generation Statistics]")
+            avg_time = sum(total_times) / len(total_times)
+            min_time = min(total_times)
+            max_time = max(total_times)
+            print_log(f"  Average total time per song: {avg_time:.1f} seconds")
+            print_log(f"  Fastest generation: {min_time:.1f} seconds")
+            print_log(f"  Slowest generation: {max_time:.1f} seconds")
+    
+    # Download Statistics
+    total_download_bytes = sum(r.get('download_bytes', 0) for r in successful_results)
+    total_download_count = sum(r.get('download_count', 0) for r in successful_results)
+    
+    if total_download_bytes > 0:
+        print_log(f"\n[Download Statistics]")
+        print_log(f"  Total download: {format_bytes(total_download_bytes)}")
+        print_log(f"  Number of files: {total_download_count}")
+        print_log(f"  Average file size: {format_bytes(total_download_bytes / total_download_count)}")
+        
+        download_speeds = [r.get('avg_download_speed', 0) for r in successful_results if r.get('avg_download_speed', 0) > 0]
+        if download_speeds:
+            avg_speed = sum(download_speeds) / len(download_speeds)
+            print_log(f"  Average download speed: {format_speed(avg_speed)}")
+    
+    # Polling Statistics
+    poll_counts = [r.get('poll_count', 0) for r in successful_results if 'poll_count' in r]
+    if poll_counts:
+        total_polls = sum(poll_counts)
+        avg_polls = total_polls / len(poll_counts)
+        print_log(f"\n[Polling Statistics]")
+        print_log(f"  Total polling count: {total_polls}")
+        print_log(f"  Average polling per song: {avg_polls:.1f}")
+    
+    # Efficiency Analysis
+    print_log(f"\n[Efficiency Analysis]")
+    if success_count > 0:
+        throughput = success_count / (overall_time / 60)
+        print_log(f"  Actual throughput: {throughput:.2f} songs/minute")
+        
+        # Theoretical fastest time (assuming no rate limit)
+        if wait_times:
+            ideal_time = submit_phase_time - rate_limit_stats['total_wait_time'] + max(wait_times)
+            efficiency = (ideal_time / overall_time) * 100
+            print_log(f"  Theoretical fastest time: {ideal_time:.1f} seconds")
+            print_log(f"  Concurrency efficiency: {efficiency:.1f}%")
+    
+    # Show failed songs
+    if fail_count > 0:
+        print_log("\n" + "=" * 70)
+        print_log("Failed Songs List")
+        print_log("=" * 70)
+        for r in sorted(final_results, key=lambda x: x.get('song_index', 0)):
+            if not r.get('success'):
+                song_id = r.get('song_id', r.get('song_index', 'Unknown'))
+                print_log(f"  [{song_id}] {r.get('error', 'Unknown error')}")
+    
+    print_log("\n" + "=" * 70)
+    print_log(f"All files saved to: {os.path.abspath(output_dir)}")
+    print_log("=" * 70)
+
+
+if __name__ == '__main__':
+    main()
+

baseline_generate/yue/batch.sh

@@ -0,0 +1,55 @@
+#!/bin/bash
+
+# Batch music generation example script
+# Requires yue source repository, main code change is infer_batch.py replacing infer.py
+
+# Get absolute path of script directory (to avoid filesystem mount issues)
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd 2>/dev/null || echo "xxx/YuE/inference")"
+
+# Change to script directory (if possible, otherwise use absolute path)
+cd "$SCRIPT_DIR" 2>/dev/null || true
+
+
+# Set HuggingFace mirror
+export HF_ENDPOINT=https://hf-mirror.com
+
+# Set PyTorch CUDA memory management optimization
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+
+export CUDA_VISIBLE_DEVICES=1
+
+# Set JSONL file path
+JSONL_PATH=""
+
+# Set output directory
+OUTPUT_DIR=""
+
+# Set processing range (optional)
+# Example: only process first 5 songs
+START_IDX=0
+END_IDX=-1
+
+# Set generation parameters
+MAX_NEW_TOKENS=3500
+REPETITION_PENALTY=1.1
+RUN_N_SEGMENTS=24
+STAGE2_BATCH_SIZE=16
+CUDA_IDX=0
+SEED=42
+NO_SAMPLE=0
+
+# Run batch generation (using absolute path)
+python "$SCRIPT_DIR/infer_batch.py" \
+    --jsonl_path "$JSONL_PATH" \
+    --output_dir "$OUTPUT_DIR" \
+    --start_idx $START_IDX \
+    --end_idx $END_IDX \
+    --max_new_tokens $MAX_NEW_TOKENS \
+    --repetition_penalty $REPETITION_PENALTY \
+    --run_n_segments $RUN_N_SEGMENTS \
+    --stage2_batch_size $STAGE2_BATCH_SIZE \
+    --cuda_idx $CUDA_IDX \
+    --seed $SEED \
+    --rescale \
+    $( [ "$NO_SAMPLE" -eq 1 ] && echo "--no_sample" )
+

baseline_generate/yue/codecmanipulator.py

@@ -0,0 +1,204 @@
+import json
+import numpy as np
+import einops
+
+
+class CodecManipulator(object):
+    r"""
+    **mm tokenizer v0.1**
+    see codeclm/hf/mm_tokenizer_v0.1_hf/id2vocab.json
+
+    text tokens: 
+        llama tokenizer 0~31999
+    
+    special tokens: "32000": "<EOD>", "32001": "<SOA>", "32002": "<EOA>", "32003": "<SOI>", "32004": "<EOI>", "32005": "<SOV>", "32006": "<EOV>", "32007": "<s_local>", "32008": "<e_local>", "32009": "<s_global>", "32010": "<e_global>", "32011": "<semantic>", "32012": "<acoustic>", "32013": "<low_level>", "32014": "<dac_16k>", "32015": "<dac_44k>", "32016": "<xcodec>", "32017": "<placeholder>", "32018": "<semantic_mert>", "32019": "<semantic_hubert>", "32020": "<visual>", "32021": "<semanticodec>"
+
+    mm tokens:
+        dac_16k: 4 codebook, 1024 vocab, 32022 - 36117
+        dac_44k: 9 codebook, 1024 vocab, 36118 - 45333
+        xcodec: 12 codebook, 1024 vocab, 45334 - 57621
+        semantic mert: 1024, 57622 - 58645
+        semantic hubert: 512, 58646 - 59157
+        visual: 64000, not included in v0.1
+        semanticodec 100tps 16384: semantic=16384, 59158 - 75541, acoustic=8192, 75542 - 83733
+    """
+    def __init__(self, codec_type, quantizer_begin=None, n_quantizer=None, teacher_forcing=False, data_feature="codec"):
+        self.codec_type = codec_type
+        self.mm_v0_2_cfg = {
+            "dac16k": {"codebook_size": 1024, "num_codebooks": 4, "global_offset": 32022, "sep": ["<dac_16k>"], "fps": 50},
+            "dac44k": {"codebook_size": 1024, "num_codebooks": 9, "global_offset": 36118, "sep": ["<dac_44k>"]},
+            "xcodec": {"codebook_size": 1024, "num_codebooks": 12, "global_offset": 45334, "sep": ["<xcodec>"], "fps": 50},
+            "mert": {"codebook_size": 1024, "global_offset": 57622, "sep": ["<semantic_mert>"]},
+            "hubert": {"codebook_size": 512, "global_offset": 58646, "sep": ["<semantic_hubert>"]},
+            "semantic/s": {"codebook_size": 16384, "num_codebooks": 1, "global_offset": 59158, "sep": ["<semanticodec>", "<semantic>"]},
+            "semantic/a": {"codebook_size": 8192, "num_codebooks": 1, "global_offset": 75542, "sep": ["<semanticodec>", "<acoustic>"]},
+            "semanticodec": {"codebook_size": [16384, 8192], "num_codebooks": 2, "global_offset": 59158, "sep": ["<semanticodec>"], "fps": 50},
+            "special_tokens": {
+                '<EOD>': 32000, '<SOA>': 32001, '<EOA>': 32002, '<SOI>': 32003, '<EOI>': 32004, '<SOV>': 32005, '<EOV>': 32006, '<s_local>': 32007, '<e_local>': 32008, '<s_global>': 32009, '<e_global>': 32010, '<semantic>': 32011, '<acoustic>': 32012, '<stage_1>': 32013, '<dac_16k>': 32014, '<dac_44k>': 32015, '<xcodec>': 32016, '<stage_2>': 32017, '<semantic_mert>': 32018, '<semantic_hubert>': 32019, '<visual>': 32020, '<semanticodec>': 32021
+            },
+            "metadata": {
+                "len": 83734,
+                "text_range": [0, 31999],
+                "special_range": [32000, 32021],
+                "mm_range": [32022, 83733]
+            },
+            "codec_range": {
+                "dac16k": [32022, 36117],
+                "dac44k": [36118, 45333],
+                "xcodec": [45334, 57621],
+                # "hifi16k": [53526, 57621],
+                "mert": [57622, 58645],
+                "hubert": [58646, 59157],
+                "semantic/s": [59158, 75541],
+                "semantic/a": [75542, 83733],
+                "semanticodec": [59158, 83733]
+            }
+        }
+        self.sep = self.mm_v0_2_cfg[self.codec_type]["sep"]
+        self.sep_ids = [self.mm_v0_2_cfg["special_tokens"][s] for s in self.sep]
+        self.codebook_size = self.mm_v0_2_cfg[self.codec_type]["codebook_size"]
+        self.num_codebooks = self.mm_v0_2_cfg[self.codec_type]["num_codebooks"]
+        self.global_offset = self.mm_v0_2_cfg[self.codec_type]["global_offset"]
+        self.fps = self.mm_v0_2_cfg[self.codec_type]["fps"] if "fps" in self.mm_v0_2_cfg[self.codec_type] else None
+
+        self.quantizer_begin = quantizer_begin if quantizer_begin is not None else 0
+        self.n_quantizer = n_quantizer if n_quantizer is not None else self.num_codebooks  
+        self.teacher_forcing = teacher_forcing 
+        self.data_feature = data_feature
+
+
+    def offset_tok_ids(self, x, global_offset=0, codebook_size=2048, num_codebooks=4):
+        """
+        x: (K, T)
+        """
+        if isinstance(codebook_size, int):
+            assert x.max() < codebook_size, f"max(x)={x.max()}, codebook_size={codebook_size}"
+        elif isinstance(codebook_size, list):
+            for i, cs in enumerate(codebook_size):
+                assert x[i].max() < cs, f"max(x)={x[i].max()}, codebook_size={cs}, layer_id={i}"
+        else:
+            raise ValueError(f"codebook_size={codebook_size}")
+        assert x.min() >= 0, f"min(x)={x.min()}"
+        assert x.shape[0] == num_codebooks or x.shape[0] == self.n_quantizer, \
+            f"x.shape[0]={x.shape[0]}, num_codebooks={num_codebooks}, n_quantizer={self.n_quantizer}"
+
+        _x = x.copy()
+        _x = _x.astype(np.uint32)
+        cum_offset = 0
+        quantizer_begin = self.quantizer_begin
+        quantizer_end = quantizer_begin+self.n_quantizer
+        for k in range(self.quantizer_begin, quantizer_end): # k: quantizer_begin to quantizer_end - 1
+            if isinstance(codebook_size, int):
+                _x[k] += global_offset + k * codebook_size
+            elif isinstance(codebook_size, list):
+                _x[k] += global_offset + cum_offset
+                cum_offset += codebook_size[k]
+            else:
+                raise ValueError(f"codebook_size={codebook_size}")
+        return _x[quantizer_begin:quantizer_end]
+
+    def unoffset_tok_ids(self, x, global_offset=0, codebook_size=2048, num_codebooks=4):
+        """
+        x: (K, T)
+        """
+        if isinstance(codebook_size, int):
+            assert x.max() < global_offset + codebook_size * num_codebooks, f"max(x)={x.max()}, codebook_size={codebook_size}"
+        elif isinstance(codebook_size, list):
+            assert x.max() < global_offset + sum(codebook_size), f"max(x)={x.max()}, codebook_size={codebook_size}"
+        assert x.min() >= global_offset, f"min(x)={x.min()}, global_offset={global_offset}"
+        assert x.shape[0] == num_codebooks or x.shape[0] == self.n_quantizer, \
+            f"x.shape[0]={x.shape[0]}, num_codebooks={num_codebooks}, n_quantizer={self.n_quantizer}"
+        
+        _x = x.copy()
+        _x = _x.astype(np.uint32)
+        cum_offset = 0
+        quantizer_begin = self.quantizer_begin
+        quantizer_end = quantizer_begin+self.n_quantizer
+        for k in range(quantizer_begin, quantizer_end):
+            if isinstance(codebook_size, int):
+                _x[k-quantizer_begin] -= global_offset + k * codebook_size
+            elif isinstance(codebook_size, list):
+                _x[k-quantizer_begin] -= global_offset + cum_offset
+                cum_offset += codebook_size[k]
+            else:
+                raise ValueError(f"codebook_size={codebook_size}")
+        return _x
+
+    def flatten(self, x):
+        if len(x.shape) > 2:
+            x = x.squeeze()
+        assert x.shape[0] == self.num_codebooks or x.shape[0] == self.n_quantizer, \
+            f"x.shape[0]={x.shape[0]}, num_codebooks={self.num_codebooks}, n_quantizer={self.n_quantizer}"
+        return einops.rearrange(x, 'K T -> (T K)')
+
+    def unflatten(self, x, n_quantizer=None):
+        if x.ndim > 1 and x.shape[0] == 1:
+            x = x.squeeze(0)
+        assert len(x.shape) == 1
+        assert x.shape[0] % self.num_codebooks == 0 or x.shape[0] % self.n_quantizer == 0, \
+            f"x.shape[0]={x.shape[0]}, num_codebooks={self.num_codebooks}, n_quantizer={self.n_quantizer}"
+        if n_quantizer!=self.num_codebooks:
+            return einops.rearrange(x, '(T K) -> K T', K=n_quantizer)
+        return einops.rearrange(x, '(T K) -> K T', K=self.num_codebooks)
+    
+    # def check_codec_type_from_path(self, path):
+    #     if self.codec_type == "hifi16k":
+    #         assert "academicodec_hifi_16k_320d_large_uni" in path
+    
+    def get_codec_type_from_range(self, ids):
+        ids_range = [ids.min(), ids.max()]
+        codec_range = self.mm_v0_2_cfg["codec_range"]
+        for codec_type, r in codec_range.items():
+            if ids_range[0] >= r[0] and ids_range[1] <= r[1]:
+                return codec_type
+        raise ValueError(f"ids_range={ids_range}, codec_range={codec_range}")
+
+    def npy2ids(self, npy):
+        if isinstance(npy, str):
+            data = np.load(npy)
+        elif isinstance(npy, np.ndarray):
+            data = npy
+        else:
+            raise ValueError(f"not supported type: {type(npy)}")
+        # data = data.squeeze()
+
+        assert len(data.shape)==2,  f'data shape: {data.shape} is not (n_codebook, seq_len)'
+        data = self.offset_tok_ids(
+            data, 
+            global_offset=self.global_offset, 
+            codebook_size=self.codebook_size, 
+            num_codebooks=self.num_codebooks, 
+        )
+        data = self.flatten(data)
+        codec_range = self.get_codec_type_from_range(data)
+        assert codec_range == self.codec_type, f"get_codec_type_from_range(data)={codec_range}, self.codec_type={self.codec_type}"
+        data = data.tolist()
+        return data
+    
+    def ids2npy(self, token_ids):
+        # make sure token_ids starts with codebook 0
+        if isinstance(self.codebook_size, int):
+            codebook_0_range = (self.global_offset + self.quantizer_begin*self.codebook_size, self.global_offset + (self.quantizer_begin+1)*self.codebook_size)
+        elif isinstance(self.codebook_size, list):
+            codebook_0_range = (self.global_offset, self.global_offset + self.codebook_size[0])
+        assert token_ids[0] >= codebook_0_range[0] \
+            and token_ids[0] < codebook_0_range[1], f"token_ids[0]={token_ids[self.quantizer_begin]}, codebook_0_range={codebook_0_range}"
+        data = np.array(token_ids)
+        data = self.unflatten(data, n_quantizer=self.n_quantizer)
+        data = self.unoffset_tok_ids(
+            data, 
+            global_offset=self.global_offset, 
+            codebook_size=self.codebook_size, 
+            num_codebooks=self.num_codebooks, 
+        )
+        return data
+
+    def npy_to_json_str(self, npy_path):
+        data = self.npy2ids(npy_path)
+        return json.dumps({"text": data, "src": npy_path, "codec": self.codec_type})
+    
+    def sep(self):
+        return ''.join(self.sep)
+    
+    def sep_ids(self):
+        return self.sep_ids

baseline_generate/yue/infer_batch.py

@@ -0,0 +1,904 @@
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer'))
+sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer', 'descriptaudiocodec'))
+import re
+import random
+import uuid
+import copy
+import json
+from tqdm import tqdm
+from collections import Counter
+import argparse
+import numpy as np
+import torch
+import torchaudio
+from torchaudio.transforms import Resample
+import soundfile as sf
+from einops import rearrange
+from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
+from omegaconf import OmegaConf
+from codecmanipulator import CodecManipulator
+from mmtokenizer import _MMSentencePieceTokenizer
+from models.soundstream_hubert_new import SoundStream
+from vocoder import build_codec_model, process_audio
+from post_process_audio import replace_low_freq_with_energy_matched
+
+os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
+
+parser = argparse.ArgumentParser()
+# Model Configuration:
+parser.add_argument("--stage1_model", type=str, default="m-a-p/YuE-s1-7B-anneal-en-cot", help="The model checkpoint path or identifier for the Stage 1 model.")
+parser.add_argument("--stage2_model", type=str, default="m-a-p/YuE-s2-1B-general", help="The model checkpoint path or identifier for the Stage 2 model.")
+parser.add_argument("--max_new_tokens", type=int, default=3000, help="The maximum number of new tokens to generate in one pass during text generation.")
+parser.add_argument("--repetition_penalty", type=float, default=1.1, help="repetition_penalty ranges from 1.0 to 2.0 (or higher in some cases). It controls the diversity and coherence of the audio tokens generated. The higher the value, the greater the discouragement of repetition. Setting value to 1.0 means no penalty.")
+parser.add_argument("--run_n_segments", type=int, default=2, help="The number of segments to process during generation. Each segment is ~30s (with default max_new_tokens=3000). For example: 2=~1min, 6=~3min, 8=~4min.")
+parser.add_argument("--stage2_batch_size", type=int, default=4, help="The batch size used in Stage 2 inference.")
+parser.add_argument(
+    "--no_sample",
+    action="store_true",
+    help="If set, disable sampling in Stage 1 generation (i.e., use deterministic decoding). When enabled, top_p/temperature will be ignored.",
+)
+# Prompt - Batch processing parameters
+parser.add_argument("--jsonl_path", type=str, required=True, help="The file path to a JSONL file containing genre and lyrics for batch processing.")
+parser.add_argument("--start_idx", type=int, default=0, help="Start index in the JSONL file for batch processing.")
+parser.add_argument("--end_idx", type=int, default=-1, help="End index in the JSONL file for batch processing. -1 means process all.")
+parser.add_argument("--use_audio_prompt", action="store_true", help="If set, the model will use an audio file as a prompt during generation. The audio file should be specified using --audio_prompt_path.")
+parser.add_argument("--audio_prompt_path", type=str, default="", help="The file path to an audio file to use as a reference prompt when --use_audio_prompt is enabled.")
+parser.add_argument("--prompt_start_time", type=float, default=0.0, help="The start time in seconds to extract the audio prompt from the given audio file.")
+parser.add_argument("--prompt_end_time", type=float, default=30.0, help="The end time in seconds to extract the audio prompt from the given audio file.")
+parser.add_argument("--use_dual_tracks_prompt", action="store_true", help="If set, the model will use dual tracks as a prompt during generation. The vocal and instrumental files should be specified using --vocal_track_prompt_path and --instrumental_track_prompt_path.")
+parser.add_argument("--vocal_track_prompt_path", type=str, default="", help="The file path to a vocal track file to use as a reference prompt when --use_dual_tracks_prompt is enabled.")
+parser.add_argument("--instrumental_track_prompt_path", type=str, default="", help="The file path to an instrumental track file to use as a reference prompt when --use_dual_tracks_prompt is enabled.")
+# Output 
+parser.add_argument("--output_dir", type=str, default="./output", help="The directory where generated outputs will be saved.")
+parser.add_argument("--keep_intermediate", action="store_true", help="If set, intermediate outputs will be saved during processing.")
+parser.add_argument("--disable_offload_model", action="store_true", help="If set, the model will not be offloaded from the GPU to CPU after Stage 1 inference.")
+parser.add_argument("--cuda_idx", type=int, default=0)
+parser.add_argument("--seed", type=int, default=42, help="An integer value to reproduce generation.")
+# Config for xcodec and upsampler
+parser.add_argument('--basic_model_config', default='./xcodec_mini_infer/final_ckpt/config.yaml', help='YAML files for xcodec configurations.')
+parser.add_argument('--resume_path', default='./xcodec_mini_infer/final_ckpt/ckpt_00360000.pth', help='Path to the xcodec checkpoint.')
+parser.add_argument('--config_path', type=str, default='./xcodec_mini_infer/decoders/config.yaml', help='Path to Vocos config file.')
+parser.add_argument('--vocal_decoder_path', type=str, default='./xcodec_mini_infer/decoders/decoder_131000.pth', help='Path to Vocos decoder weights.')
+parser.add_argument('--inst_decoder_path', type=str, default='./xcodec_mini_infer/decoders/decoder_151000.pth', help='Path to Vocos decoder weights.')
+parser.add_argument('-r', '--rescale', action='store_true', help='Rescale output to avoid clipping.')
+
+
+args = parser.parse_args()
+if args.use_audio_prompt and not args.audio_prompt_path:
+    raise FileNotFoundError("Please offer audio prompt filepath using '--audio_prompt_path', when you enable 'use_audio_prompt'!")
+if args.use_dual_tracks_prompt and not args.vocal_track_prompt_path and not args.instrumental_track_prompt_path:
+    raise FileNotFoundError("Please offer dual tracks prompt filepath using '--vocal_track_prompt_path' and '--inst_decoder_path', when you enable '--use_dual_tracks_prompt'!")
+
+stage1_model = args.stage1_model
+stage2_model = args.stage2_model
+cuda_idx = args.cuda_idx
+max_new_tokens = args.max_new_tokens
+do_sample_stage1 = (not args.no_sample)
+
+def seed_everything(seed=42): 
+    random.seed(seed) 
+    np.random.seed(seed) 
+    torch.manual_seed(seed) 
+    torch.cuda.manual_seed_all(seed) 
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+seed_everything(args.seed)
+
+# Read JSONL file
+print(f"Reading JSONL file: {args.jsonl_path}")
+music_data_list = []
+with open(args.jsonl_path, 'r', encoding='utf-8') as f:
+    for line in f:
+        if line.strip():
+            music_data_list.append(json.loads(line))
+
+# Determine processing range
+start_idx = args.start_idx
+end_idx = len(music_data_list) if args.end_idx == -1 else min(args.end_idx, len(music_data_list))
+music_data_list = music_data_list[start_idx:end_idx]
+print(f"Total {len(music_data_list)} songs to generate (indices {start_idx} to {end_idx-1})")
+
+# Detect processed songs - check completion status of each stage
+def check_song_status(song_idx, output_dir):
+    """
+    Check song processing status
+    Returns: (stage1_done, stage2_done, stage3_done, song_dir, stage1_output_set, stage2_output_dir)
+    """
+    if not os.path.exists(output_dir):
+        return False, False, False, None, None, None
+    
+    # Find song directory (may have multiple, take the latest or first)
+    song_dirs = []
+    for item in os.listdir(output_dir):
+        if item.startswith('song_') and os.path.isdir(os.path.join(output_dir, item)):
+            try:
+                idx = int(item.split('_')[1])
+                if idx == song_idx:
+                    song_dirs.append(os.path.join(output_dir, item))
+            except (ValueError, IndexError):
+                continue
+    
+    if not song_dirs:
+        return False, False, False, None, None, None
+    
+    # Use the latest directory (sorted by modification time)
+    song_dir = max(song_dirs, key=lambda x: os.path.getmtime(x))
+    
+    # Check Stage 1: whether stage1 directory has vtrack and itrack .npy files
+    stage1_dir = os.path.join(song_dir, "stage1")
+    stage1_done = False
+    stage1_output_set = []
+    if os.path.exists(stage1_dir):
+        stage1_files = [f for f in os.listdir(stage1_dir) if f.endswith('.npy')]
+        vtrack_files = [f for f in stage1_files if '_vtrack' in f]
+        itrack_files = [f for f in stage1_files if '_itrack' in f]
+        if vtrack_files and itrack_files:
+            stage1_done = True
+            # Build stage1_output_set
+            for f in vtrack_files + itrack_files:
+                stage1_output_set.append(os.path.join(stage1_dir, f))
+    
+    # Check Stage 2: whether stage2 directory has corresponding .npy files
+    stage2_dir = os.path.join(song_dir, "stage2")
+    stage2_done = False
+    if stage1_done and os.path.exists(stage2_dir):
+        stage2_files = [f for f in os.listdir(stage2_dir) if f.endswith('.npy')]
+        # Check if all stage1 files have corresponding stage2 files
+        if stage1_output_set:
+            stage1_basenames = {os.path.basename(f) for f in stage1_output_set}
+            stage2_basenames = set(stage2_files)
+            if stage1_basenames.issubset(stage2_basenames):
+                stage2_done = True
+    
+    # Check Stage 3: whether there is a final mixed file (in song_dir root directory)
+    stage3_done = False
+    for root, dirs, files in os.walk(song_dir):
+        if any(f.endswith('_mixed.mp3') for f in files):
+            stage3_done = True
+            break
+    
+    return stage1_done, stage2_done, stage3_done, song_dir, stage1_output_set, stage2_dir
+
+# Detect processing status of all songs
+song_status_map = {}  # {song_idx: (stage1_done, stage2_done, stage3_done, song_dir, stage1_output_set, stage2_output_dir)}
+if os.path.exists(args.output_dir):
+    print(f"\nDetecting processed songs...")
+    for list_idx in range(len(music_data_list)):
+        song_idx = start_idx + list_idx
+        stage1_done, stage2_done, stage3_done, song_dir, stage1_output_set, stage2_output_dir = check_song_status(song_idx, args.output_dir)
+        if stage1_done or stage2_done or stage3_done:
+            song_status_map[song_idx] = (stage1_done, stage2_done, stage3_done, song_dir, stage1_output_set, stage2_output_dir)
+    
+    if song_status_map:
+        fully_completed = [idx for idx, (s1, s2, s3, _, _, _) in song_status_map.items() if s3]
+        partial_completed = [idx for idx, (s1, s2, s3, _, _, _) in song_status_map.items() if not s3]
+        print(f"✓ Found {len(fully_completed)} fully completed songs: {sorted(fully_completed)}")
+        if partial_completed:
+            print(f"✓ Found {len(partial_completed)} partially completed songs: {sorted(partial_completed)}")
+            for idx in sorted(partial_completed):
+                s1, s2, s3, _, _, _ = song_status_map[idx]
+                status_parts = []
+                if s1: status_parts.append("Stage1")
+                if s2: status_parts.append("Stage2")
+                if s3: status_parts.append("Stage3")
+                print(f"  Index {idx}: Completed {', '.join(status_parts)}")
+        remaining_count = len(music_data_list) - len(fully_completed)
+        print(f"✓ Will skip fully completed songs, {remaining_count} songs remaining to process")
+    else:
+        print(f"✓ No processed songs found, will start from the beginning")
+else:
+    print(f"✓ Output directory does not exist, will start from the beginning")
+
+# Load tokenizer and model
+device = torch.device(f"cuda:{cuda_idx}" if torch.cuda.is_available() else "cpu")
+mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
+print("Loading Stage 1 model...")
+model = AutoModelForCausalLM.from_pretrained(
+    stage1_model, 
+    torch_dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2", # Using flash_attention_2 for better performance
+    # device_map="auto",
+    )
+# to device, if gpu is available
+model.to(device)
+model.eval()
+
+if torch.__version__ >= "2.0.0":
+    try:
+        model = torch.compile(model)
+    except Exception as e:
+        print(f"Warning: torch.compile not available: {e}")
+
+codectool = CodecManipulator("xcodec", 0, 1)
+codectool_stage2 = CodecManipulator("xcodec", 0, 8)
+model_config = OmegaConf.load(args.basic_model_config)
+codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
+# Load checkpoint with weights_only=False to allow OmegaConf types
+# Note: Only use this if you trust the checkpoint source
+parameter_dict = torch.load(args.resume_path, map_location='cpu', weights_only=False)
+codec_model.load_state_dict(parameter_dict['codec_model'])
+codec_model.to(device)
+codec_model.eval()
+
+class BlockTokenRangeProcessor(LogitsProcessor):
+    def __init__(self, start_id, end_id):
+        self.blocked_token_ids = list(range(start_id, end_id))
+
+    def __call__(self, input_ids, scores):
+        scores[:, self.blocked_token_ids] = -float("inf")
+        return scores
+
+def load_audio_mono(filepath, sampling_rate=16000):
+    audio, sr = torchaudio.load(filepath)
+    # Convert to mono
+    audio = torch.mean(audio, dim=0, keepdim=True)
+    # Resample if needed
+    if sr != sampling_rate:
+        resampler = Resample(orig_freq=sr, new_freq=sampling_rate)
+        audio = resampler(audio)
+    return audio
+
+def encode_audio(codec_model, audio_prompt, device, target_bw=0.5):
+    if len(audio_prompt.shape) < 3:
+        audio_prompt.unsqueeze_(0)
+    with torch.no_grad():
+        raw_codes = codec_model.encode(audio_prompt.to(device), target_bw=target_bw)
+    raw_codes = raw_codes.transpose(0, 1)
+    raw_codes = raw_codes.cpu().numpy().astype(np.int16)
+    return raw_codes
+
+def split_lyrics(lyrics):
+    """
+    Split lyrics by segments, following YuE official best practices:
+    
+    Official requirements:
+    1. Lyrics should be segmented using structure tags: [verse], [chorus], [bridge], [outro], etc.
+    2. Each segment is separated by two newlines "\n\n"
+    3. Each segment is about 30 seconds (when --max_new_tokens 3000), don't put too many words
+    4. Avoid using [intro] tag (not very stable), recommend starting with [verse] or [chorus]
+    5. Supports multiple languages: English, Chinese, Cantonese, Japanese, Korean, etc.
+    
+    Args:
+        lyrics: Raw lyrics string
+    
+    Returns:
+        Structured lyrics segment list, each segment in [tag]\ncontent\n\n format
+    """
+    # Regular expression: match [any tag] and its following content
+    # Supports: [Verse 1], [Pre-Chorus], [Chorus (Outro)] and other complex tags
+    pattern = r"\[([^\]]+)\](.*?)(?=\[|\Z)"
+    segments = re.findall(pattern, lyrics, re.DOTALL)
+    structured_lyrics = [f"[{seg[0]}]\n{seg[1].strip()}\n\n" for seg in segments]
+    return structured_lyrics
+
+def save_audio(wav: torch.Tensor, path, sample_rate: int, rescale: bool = False):
+    folder_path = os.path.dirname(path)
+    if not os.path.exists(folder_path):
+        os.makedirs(folder_path)
+    limit = 0.99
+    max_val = wav.abs().max()
+    wav = wav * min(limit / max_val, 1) if rescale else wav.clamp(-limit, limit)
+    torchaudio.save(str(path), wav, sample_rate=sample_rate, encoding='PCM_S', bits_per_sample=16)
+
+def stage2_generate(model, prompt, batch_size=16):
+    codec_ids = codectool.unflatten(prompt, n_quantizer=1)
+    codec_ids = codectool.offset_tok_ids(
+                    codec_ids, 
+                    global_offset=codectool.global_offset, 
+                    codebook_size=codectool.codebook_size, 
+                    num_codebooks=codectool.num_codebooks, 
+                ).astype(np.int32)
+    
+    # Prepare prompt_ids based on batch size or single input
+    if batch_size > 1:
+        codec_list = []
+        for i in range(batch_size):
+            idx_begin = i * 300
+            idx_end = (i + 1) * 300
+            codec_list.append(codec_ids[:, idx_begin:idx_end])
+
+        codec_ids = np.concatenate(codec_list, axis=0)
+        prompt_ids = np.concatenate(
+            [
+                np.tile([mmtokenizer.soa, mmtokenizer.stage_1], (batch_size, 1)),
+                codec_ids,
+                np.tile([mmtokenizer.stage_2], (batch_size, 1)),
+            ],
+            axis=1
+        )
+    else:
+        prompt_ids = np.concatenate([
+            np.array([mmtokenizer.soa, mmtokenizer.stage_1]),
+            codec_ids.flatten(),  # Flatten the 2D array to 1D
+            np.array([mmtokenizer.stage_2])
+        ]).astype(np.int32)
+        prompt_ids = prompt_ids[np.newaxis, ...]
+
+    codec_ids = torch.as_tensor(codec_ids).to(device)
+    prompt_ids = torch.as_tensor(prompt_ids).to(device)
+    len_prompt = prompt_ids.shape[-1]
+    
+    block_list = LogitsProcessorList([BlockTokenRangeProcessor(0, 46358), BlockTokenRangeProcessor(53526, mmtokenizer.vocab_size)])
+
+    # Teacher forcing generate loop
+    for frames_idx in range(codec_ids.shape[1]):
+        cb0 = codec_ids[:, frames_idx:frames_idx+1]
+        prompt_ids = torch.cat([prompt_ids, cb0], dim=1)
+        input_ids = prompt_ids
+
+        with torch.no_grad():
+            stage2_output = model.generate(input_ids=input_ids, 
+                min_new_tokens=7,
+                max_new_tokens=7,
+                eos_token_id=mmtokenizer.eoa,
+                pad_token_id=mmtokenizer.eoa,
+                logits_processor=block_list,
+            )
+        
+        assert stage2_output.shape[1] - prompt_ids.shape[1] == 7, f"output new tokens={stage2_output.shape[1]-prompt_ids.shape[1]}"
+        prompt_ids = stage2_output
+
+    # Return output based on batch size
+    if batch_size > 1:
+        output = prompt_ids.cpu().numpy()[:, len_prompt:]
+        output_list = [output[i] for i in range(batch_size)]
+        output = np.concatenate(output_list, axis=0)
+    else:
+        output = prompt_ids[0].cpu().numpy()[len_prompt:]
+
+    return output
+
+def sanitize_genres_for_filename(genres, max_length=80):
+    """
+    Clean and truncate genres string for filename generation
+    Ensure filename is not too long (Linux filename limit is 255 bytes)
+    
+    Args:
+        genres: Raw genres string
+        max_length: Maximum length of genres part (default 80, leaving space for other parameters)
+    
+    Returns:
+        Cleaned genres string
+    """
+    if not genres:
+        return "Unknown"
+    
+    # Clean unsafe characters
+    genres_clean = re.sub(r'[<>:"/\\|?*\x00-\x1f]', '_', genres)
+    genres_clean = genres_clean.strip('_').strip()
+    
+    # If contains comma-separated tags, try to keep first few tags
+    if ',' in genres_clean:
+        tags = [tag.strip() for tag in genres_clean.split(',')]
+        # Try to keep first few tags until reaching length limit
+        result_tags = []
+        current_length = 0
+        for tag in tags:
+            if current_length + len(tag) + 1 <= max_length:  # +1 for comma
+                result_tags.append(tag)
+                current_length += len(tag) + 1
+            else:
+                break
+        if result_tags:
+            genres_clean = ','.join(result_tags)
+        else:
+            # If first tag is too long, directly truncate
+            genres_clean = tags[0][:max_length] if tags else genres_clean[:max_length]
+    
+    # If still too long, directly truncate
+    if len(genres_clean) > max_length:
+        genres_clean = genres_clean[:max_length]
+    
+    # Replace spaces with hyphens (for consistency)
+    genres_clean = genres_clean.replace(' ', '-')
+    
+    return genres_clean
+
+def stage2_inference(model, stage1_output_set, stage2_output_dir, batch_size=4):
+    stage2_result = []
+    for i in tqdm(range(len(stage1_output_set)), desc="Stage 2 inference"):
+        output_filename = os.path.join(stage2_output_dir, os.path.basename(stage1_output_set[i]))
+        
+        if os.path.exists(output_filename):
+            print(f'{output_filename} stage2 has done.')
+            stage2_result.append(output_filename)
+            continue
+        
+        # Load the prompt
+        prompt = np.load(stage1_output_set[i]).astype(np.int32)
+        
+        # Only accept 6s segments
+        output_duration = prompt.shape[-1] // 50 // 6 * 6
+        num_batch = output_duration // 6
+        
+        if num_batch <= batch_size:
+            # If num_batch is less than or equal to batch_size, we can infer the entire prompt at once
+            output = stage2_generate(model, prompt[:, :output_duration*50], batch_size=num_batch)
+        else:
+            # If num_batch is greater than batch_size, process in chunks of batch_size
+            segments = []
+            num_segments = (num_batch // batch_size) + (1 if num_batch % batch_size != 0 else 0)
+
+            for seg in range(num_segments):
+                start_idx = seg * batch_size * 300
+                # Ensure the end_idx does not exceed the available length
+                end_idx = min((seg + 1) * batch_size * 300, output_duration*50)  # Adjust the last segment
+                current_batch_size = batch_size if seg != num_segments-1 or num_batch % batch_size == 0 else num_batch % batch_size
+                segment = stage2_generate(
+                    model,
+                    prompt[:, start_idx:end_idx],
+                    batch_size=current_batch_size
+                )
+                segments.append(segment)
+
+            # Concatenate all the segments
+            output = np.concatenate(segments, axis=0)
+        
+        # Process the ending part of the prompt
+        if output_duration*50 != prompt.shape[-1]:
+            ending = stage2_generate(model, prompt[:, output_duration*50:], batch_size=1)
+            output = np.concatenate([output, ending], axis=0)
+        output = codectool_stage2.ids2npy(output)
+
+        # Fix invalid codes (a dirty solution, which may harm the quality of audio)
+        # We are trying to find better one
+        fixed_output = copy.deepcopy(output)
+        for i, line in enumerate(output):
+            for j, element in enumerate(line):
+                if element < 0 or element > 1023:
+                    counter = Counter(line)
+                    most_frequant = sorted(counter.items(), key=lambda x: x[1], reverse=True)[0][0]
+                    fixed_output[i, j] = most_frequant
+        # save output
+        np.save(output_filename, fixed_output)
+        stage2_result.append(output_filename)
+    return stage2_result
+
+def process_one_song(music_data, song_idx, total_songs):
+    """Process Stage 1 for a single song"""
+    
+    # Compatible with genre and description fields
+    genres = music_data.get('genre') or music_data.get('description', '')
+    lyrics_raw = music_data['lyrics']
+    description = music_data.get('description', '')
+    
+    print(f"Description: {description[:100]}...")
+    print(f"Genre tags: {genres}")
+    
+    # ===== Print original lyrics =====
+    print("\n" + "="*60)
+    print("【Original Lyrics (lyrics_raw)】")
+    print("="*60)
+    print(lyrics_raw)
+    print("="*60 + "\n")
+    
+    lyrics = split_lyrics(lyrics_raw)
+    
+    # Validate lyrics format and give warnings (following official best practices)
+    print(f"Lyrics analysis: Identified {len(lyrics)} segments")
+    
+    # ===== Print segmented lyrics =====
+    print("\n" + "="*60)
+    print("【Segmented Lyrics (lyrics)】")
+    print("="*60)
+    for i, seg in enumerate(lyrics):
+        tag = seg.split('\n')[0].strip()
+        # Check if unstable [intro] tag is used
+        if 'intro' in tag.lower():
+            print(f"  ⚠️  Warning: Segment {i+1} uses {tag} tag, official recommendation is to avoid [intro], use [verse] or [chorus] instead")
+        else:
+            print(f"  Segment {i+1}. {tag}")
+        # Print each segment's content (limit length)
+        content = seg.strip()
+        if len(content) > 150:
+            print(f"    Content preview: {content[:150]}...")
+        else:
+            print(f"    Content: {content}")
+        print()
+    print("="*60 + "\n")
+    
+    # Create output directory for this song
+    random_id = uuid.uuid4()
+    song_output_dir = os.path.join(args.output_dir, f"song_{song_idx:04d}_{random_id}")
+    stage1_output_dir = os.path.join(song_output_dir, "stage1")
+    stage2_output_dir = os.path.join(song_output_dir, "stage2")
+    os.makedirs(stage1_output_dir, exist_ok=True)
+    os.makedirs(stage2_output_dir, exist_ok=True)
+    
+    # Stage 1: Generate audio tokens
+    print("--- Stage 1: Generate audio tokens ---")
+    stage1_output_set = []
+    full_lyrics = "\n".join(lyrics)
+    prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genres}\n{full_lyrics}"]
+    prompt_texts += lyrics
+    
+    # ===== Print prompt texts passed to model =====
+    print("\n" + "="*60)
+    print("【Prompt Texts Passed to Model (prompt_texts)】")
+    print("="*60)
+    print(f"Total {len(prompt_texts)} prompts (first is full prompt, subsequent are segments)\n")
+    for i, pt in enumerate(prompt_texts):
+        if i == 0:
+            print(f"Prompt {i} [Full prompt header]:")
+            if len(pt) > 300:
+                print(f"{pt[:300]}...")
+            else:
+                print(pt)
+        else:
+            print(f"\nPrompt {i} [Segment {i}]:")
+            if len(pt) > 200:
+                print(f"{pt[:200]}...")
+            else:
+                print(pt)
+    print("="*60 + "\n")
+    
+    output_seq = None
+    # Here is suggested decoding config
+    top_p = 0.93
+    temperature = 1.0
+    repetition_penalty = args.repetition_penalty
+    if not do_sample_stage1:
+        print("Note: --no_sample is enabled, Stage 1 will use deterministic decoding; top_p/temperature will be ignored.")
+    # special tokens
+    start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
+    end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
+    # Format text prompt
+    # +1 because prompt_texts[0] is the full prompt which will be skipped, so need len(lyrics)+1 to process all segments
+    run_n_segments = min(args.run_n_segments+1, len(lyrics)+1)
+    
+    for i, p in enumerate(tqdm(prompt_texts[:run_n_segments], desc="Stage1 inference")):
+        section_text = p.replace('[start_of_segment]', '').replace('[end_of_segment]', '')
+        guidance_scale = 1.5 if i <=1 else 1.2
+        
+        # ===== Print currently processing segment =====
+        if i == 0:
+            print(f"\n[Segment {i}] Skipped (full prompt header)")
+        else:
+            print(f"\n" + "-"*60)
+            print(f"[Processing segment {i}/{len(prompt_texts[:run_n_segments])-1}]")
+            print("-"*60)
+            tag_line = section_text.split('\n')[0] if '\n' in section_text else section_text[:50]
+            print(f"Segment tag: {tag_line}")
+            print(f"Segment content length: {len(section_text)} characters")
+            if len(section_text) > 200:
+                print(f"Segment content preview: {section_text[:200]}...")
+            else:
+                print(f"Segment content: {section_text}")
+            print("-"*60)
+        
+        if i==0:
+            continue
+        if i==1:
+            if args.use_dual_tracks_prompt or args.use_audio_prompt:
+                if args.use_dual_tracks_prompt:
+                    vocals_ids = load_audio_mono(args.vocal_track_prompt_path)
+                    instrumental_ids = load_audio_mono(args.instrumental_track_prompt_path)
+                    vocals_ids = encode_audio(codec_model, vocals_ids, device, target_bw=0.5)
+                    instrumental_ids = encode_audio(codec_model, instrumental_ids, device, target_bw=0.5)
+                    vocals_ids = codectool.npy2ids(vocals_ids[0])
+                    instrumental_ids = codectool.npy2ids(instrumental_ids[0])
+                    ids_segment_interleaved = rearrange([np.array(vocals_ids), np.array(instrumental_ids)], 'b n -> (n b)')
+                    audio_prompt_codec = ids_segment_interleaved[int(args.prompt_start_time*50*2): int(args.prompt_end_time*50*2)]
+                    audio_prompt_codec = audio_prompt_codec.tolist()
+                elif args.use_audio_prompt:
+                    audio_prompt = load_audio_mono(args.audio_prompt_path)
+                    raw_codes = encode_audio(codec_model, audio_prompt, device, target_bw=0.5)
+                    # Format audio prompt
+                    code_ids = codectool.npy2ids(raw_codes[0])
+                    audio_prompt_codec = code_ids[int(args.prompt_start_time *50): int(args.prompt_end_time *50)] # 50 is tps of xcodec
+                audio_prompt_codec_ids = [mmtokenizer.soa] + codectool.sep_ids + audio_prompt_codec + [mmtokenizer.eoa]
+                sentence_ids = mmtokenizer.tokenize("[start_of_reference]") +  audio_prompt_codec_ids + mmtokenizer.tokenize("[end_of_reference]")
+                head_id = mmtokenizer.tokenize(prompt_texts[0]) + sentence_ids
+            else:
+                head_id = mmtokenizer.tokenize(prompt_texts[0])
+            prompt_ids = head_id + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
+        else:
+            prompt_ids = end_of_segment + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
+
+        prompt_ids = torch.as_tensor(prompt_ids).unsqueeze(0).to(device) 
+        input_ids = torch.cat([raw_output, prompt_ids], dim=1) if i > 1 else prompt_ids
+        # Use window slicing in case output sequence exceeds the context of model
+        max_context = 16384-max_new_tokens-1
+        if input_ids.shape[-1] > max_context:
+            print(f'Section {i}: output length {input_ids.shape[-1]} exceeding context length {max_context}, now using the last {max_context} tokens.')
+            input_ids = input_ids[:, -(max_context):]
+        with torch.no_grad():
+            output_seq = model.generate(
+                input_ids=input_ids, 
+                max_new_tokens=max_new_tokens, 
+                min_new_tokens=100, 
+                do_sample=do_sample_stage1,
+                top_p=top_p,
+                temperature=temperature, 
+                repetition_penalty=repetition_penalty, 
+                eos_token_id=mmtokenizer.eoa,
+                pad_token_id=mmtokenizer.eoa,
+                logits_processor=LogitsProcessorList([BlockTokenRangeProcessor(0, 32002), BlockTokenRangeProcessor(32016, 32016)]),
+                guidance_scale=guidance_scale,
+                )
+            if output_seq[0][-1].item() != mmtokenizer.eoa:
+                tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)
+                output_seq = torch.cat((output_seq, tensor_eoa), dim=1)
+        if i > 1:
+            raw_output = torch.cat([raw_output, prompt_ids, output_seq[:, input_ids.shape[-1]:]], dim=1)
+        else:
+            raw_output = output_seq
+
+    # save raw output and check sanity
+    ids = raw_output[0].cpu().numpy()
+    soa_idx = np.where(ids == mmtokenizer.soa)[0].tolist()
+    eoa_idx = np.where(ids == mmtokenizer.eoa)[0].tolist()
+    if len(soa_idx)!=len(eoa_idx):
+        raise ValueError(f'invalid pairs of soa and eoa, Num of soa: {len(soa_idx)}, Num of eoa: {len(eoa_idx)}')
+
+    vocals = []
+    instrumentals = []
+    range_begin = 1 if args.use_audio_prompt or args.use_dual_tracks_prompt else 0
+    for i in range(range_begin, len(soa_idx)):
+        codec_ids = ids[soa_idx[i]+1:eoa_idx[i]]
+        if codec_ids[0] == 32016:
+            codec_ids = codec_ids[1:]
+        codec_ids = codec_ids[:2 * (codec_ids.shape[0] // 2)]
+        vocals_ids = codectool.ids2npy(rearrange(codec_ids,"(n b) -> b n", b=2)[0])
+        vocals.append(vocals_ids)
+        instrumentals_ids = codectool.ids2npy(rearrange(codec_ids,"(n b) -> b n", b=2)[1])
+        instrumentals.append(instrumentals_ids)
+    vocals = np.concatenate(vocals, axis=1)
+    instrumentals = np.concatenate(instrumentals, axis=1)
+    # Clean genres string to avoid filename being too long
+    genres_clean = sanitize_genres_for_filename(genres, max_length=80)
+    vocal_save_path = os.path.join(stage1_output_dir, f"{genres_clean}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}_{random_id}_vtrack".replace('.', '@')+'.npy')
+    inst_save_path = os.path.join(stage1_output_dir, f"{genres_clean}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}_{random_id}_itrack".replace('.', '@')+'.npy')
+    np.save(vocal_save_path, vocals)
+    np.save(inst_save_path, instrumentals)
+    stage1_output_set.append(vocal_save_path)
+    stage1_output_set.append(inst_save_path)
+    
+    return stage1_output_set, stage2_output_dir, song_output_dir
+
+# Load Stage 2 model and vocoder (load only once)
+print("\n" + "="*60)
+print("Loading Stage 2 model...")
+print("="*60)
+model_stage2 = AutoModelForCausalLM.from_pretrained(
+    stage2_model, 
+    torch_dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2", # Using flash_attention_2 for better performance
+    # device_map="auto",
+    )
+model_stage2.to(device)
+model_stage2.eval()
+
+if torch.__version__ >= "2.0.0":
+    try:
+        model_stage2 = torch.compile(model_stage2)
+    except Exception as e:
+        print(f"Warning: torch.compile not available: {e}")
+
+print("Loading vocoder...")
+vocal_decoder, inst_decoder = build_codec_model(args.config_path, args.vocal_decoder_path, args.inst_decoder_path)
+
+# Batch process all songs - process each song completely before continuing to next
+all_results = []
+skipped_count = 0
+for list_idx, music_data in enumerate(music_data_list):
+    # Calculate actual song index (considering start_idx offset)
+    song_idx = start_idx + list_idx
+    
+    try:
+        # Compatible with genre and description fields
+        genres = music_data.get('genre') or music_data.get('description', '')
+        
+        # Check processing status
+        stage1_done = False
+        stage2_done = False
+        stage3_done = False
+        song_output_dir = None
+        stage1_output_set = None
+        stage2_output_dir = None
+        
+        if song_idx in song_status_map:
+            stage1_done, stage2_done, stage3_done, song_output_dir, stage1_output_set, stage2_output_dir = song_status_map[song_idx]
+        
+        # If all completed, skip
+        if stage3_done:
+            print(f"\n{'='*60}")
+            print(f"⏭️  Skipping song {list_idx+1}/{len(music_data_list)} (index {song_idx}, fully completed)")
+            print(f"{'='*60}")
+            skipped_count += 1
+            continue
+        
+        # Decide which stage to start from based on completion status
+        print(f"\n{'='*60}")
+        print(f"Starting to process song {list_idx+1}/{len(music_data_list)} (index {song_idx})")
+        if stage1_done:
+            print(f"  ✓ Stage 1 completed, will start from Stage 2")
+        if stage2_done:
+            print(f"  ✓ Stage 2 completed, will start from Stage 3")
+        print(f"{'='*60}")
+        
+        # Stage 1: Generate audio tokens (if not completed)
+        if not stage1_done:
+            stage1_output_set, stage2_output_dir, song_output_dir = process_one_song(music_data, song_idx, len(music_data_list))
+            print(f"��� Stage 1 completed, generated {len(stage1_output_set)} files")
+            for f in stage1_output_set:
+                print(f"  - {os.path.basename(f)}")
+        else:
+            print(f"⏭️  Skipping Stage 1 (completed)")
+            print(f"  Using existing Stage 1 outputs:")
+            for f in stage1_output_set:
+                print(f"    - {os.path.basename(f)}")
+        
+        # Note: Do not unload Stage 1 model here, as subsequent songs still need it
+        # Stage 1 model will be unloaded uniformly after all songs are processed
+        
+        # Stage 2: Process audio tokens (if not completed)
+        if not stage2_done:
+            print(f"\n--- Stage 2: Processing song {list_idx+1} (index {song_idx}) ---")
+            stage2_result = stage2_inference(model_stage2, stage1_output_set, stage2_output_dir, batch_size=args.stage2_batch_size)
+            print(f"✓ Stage 2 completed, generated {len(stage2_result)} files")
+            for f in stage2_result:
+                print(f"  - {os.path.basename(f)}")
+        else:
+            print(f"\n⏭️  Skipping Stage 2 (completed)")
+            # Get existing stage2 results
+            stage2_result = []
+            if os.path.exists(stage2_output_dir):
+                for f in stage1_output_set:
+                    basename = os.path.basename(f)
+                    stage2_file = os.path.join(stage2_output_dir, basename)
+                    if os.path.exists(stage2_file):
+                        stage2_result.append(stage2_file)
+            print(f"  Using existing Stage 2 outputs:")
+            for f in stage2_result:
+                print(f"    - {os.path.basename(f)}")
+        
+        # Stage 3: Reconstruct audio and mix (if not completed)
+        final_output = None
+        if not stage3_done:
+            print(f"\n--- Stage 3: Reconstructing audio for song {list_idx+1} (index {song_idx}) ---")
+            
+            # reconstruct tracks
+            recons_output_dir = os.path.join(song_output_dir, "recons")
+            recons_mix_dir = os.path.join(recons_output_dir, 'mix')
+            os.makedirs(recons_mix_dir, exist_ok=True)
+            tracks = []
+            for npy in stage2_result:
+                codec_result = np.load(npy)
+                decodec_rlt=[]
+                with torch.no_grad():
+                    decoded_waveform = codec_model.decode(torch.as_tensor(codec_result.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(device))
+                decoded_waveform = decoded_waveform.cpu().squeeze(0)
+                decodec_rlt.append(torch.as_tensor(decoded_waveform))
+                decodec_rlt = torch.cat(decodec_rlt, dim=-1)
+                save_path = os.path.join(recons_output_dir, os.path.splitext(os.path.basename(npy))[0] + ".mp3")
+                tracks.append(save_path)
+                save_audio(decodec_rlt, save_path, 16000)
+            
+            # mix tracks
+            recons_mix = None
+            for inst_path in tracks:
+                try:
+                    if (inst_path.endswith('.wav') or inst_path.endswith('.mp3')) \
+                        and '_itrack' in inst_path:
+                        # find pair
+                        vocal_path = inst_path.replace('_itrack', '_vtrack')
+                        if not os.path.exists(vocal_path):
+                            continue
+                        # mix
+                        recons_mix = os.path.join(recons_mix_dir, os.path.basename(inst_path).replace('_itrack', '_mixed'))
+                        vocal_stem, sr = sf.read(inst_path)
+                        instrumental_stem, _ = sf.read(vocal_path)
+                        mix_stem = (vocal_stem + instrumental_stem) / 1
+                        sf.write(recons_mix, mix_stem, sr)
+                except Exception as e:
+                    print(e)
+
+            # vocoder to upsample audios
+            vocoder_output_dir = os.path.join(song_output_dir, 'vocoder')
+            vocoder_stems_dir = os.path.join(vocoder_output_dir, 'stems')
+            vocoder_mix_dir = os.path.join(vocoder_output_dir, 'mix')
+            os.makedirs(vocoder_mix_dir, exist_ok=True)
+            os.makedirs(vocoder_stems_dir, exist_ok=True)
+            
+            for npy in stage2_result:
+                if '_itrack' in npy:
+                    # Process instrumental
+                    instrumental_output = process_audio(
+                        npy,
+                        os.path.join(vocoder_stems_dir, 'itrack.mp3'),
+                        args.rescale,
+                        args,
+                        inst_decoder,
+                        codec_model
+                    )
+                else:
+                    # Process vocal
+                    vocal_output = process_audio(
+                        npy,
+                        os.path.join(vocoder_stems_dir, 'vtrack.mp3'),
+                        args.rescale,
+                        args,
+                        vocal_decoder,
+                        codec_model
+                    )
+            
+            # mix tracks
+            vocoder_mix = None
+            try:
+                mix_output = instrumental_output + vocal_output
+                vocoder_mix = os.path.join(vocoder_mix_dir, os.path.basename(recons_mix))
+                save_audio(mix_output, vocoder_mix, 44100, args.rescale)
+                print(f"Created mix: {vocoder_mix}")
+            except RuntimeError as e:
+                print(e)
+                print(f"Mix failed! inst: {instrumental_output.shape}, vocal: {vocal_output.shape}")
+
+            # Post process
+            if recons_mix and vocoder_mix:
+                final_output = os.path.join(song_output_dir, os.path.basename(recons_mix))
+                replace_low_freq_with_energy_matched(
+                    a_file=recons_mix,     # 16kHz
+                    b_file=vocoder_mix,     # 48kHz
+                    c_file=final_output,
+                    cutoff_freq=5500.0
+                )
+                print(f"✓ Song {list_idx+1} (index {song_idx}) completed! Output: {final_output}")
+        else:
+            print(f"\n⏭️  Skipping Stage 3 (completed)")
+            # Find final output file (usually in song_dir root directory)
+            # First check root directory
+            root_files = [f for f in os.listdir(song_output_dir) if f.endswith('_mixed.mp3')]
+            if root_files:
+                final_output = os.path.join(song_output_dir, root_files[0])
+            else:
+                # If root directory doesn't have it, traverse subdirectories to find
+                for root, dirs, files in os.walk(song_output_dir):
+                    for f in files:
+                        if f.endswith('_mixed.mp3'):
+                            final_output = os.path.join(root, f)
+                            break
+                    if final_output:
+                        break
+            if final_output:
+                print(f"  Final output: {final_output}")
+        
+        all_results.append({
+            'song_idx': song_idx,
+            'genres': genres,
+            'output_path': final_output if recons_mix and vocoder_mix else None
+        })
+        
+    except Exception as e:
+        print(f"✗ Error processing song {list_idx+1} (index {song_idx}): {e}")
+        import traceback
+        traceback.print_exc()
+        continue
+
+# After all songs are processed, unload models to free memory
+if not args.disable_offload_model:
+    print("\nCleaning up models to free memory...")
+    if 'model' in locals():
+        model.cpu()
+        del model
+    if 'model_stage2' in locals():
+        model_stage2.cpu()
+        del model_stage2
+    torch.cuda.empty_cache()
+    print("Models unloaded")
+
+print("\n" + "="*60)
+print("Batch generation complete!")
+newly_processed = len([r for r in all_results if r.get('output_path')])
+print(f"✓ Newly processed: {newly_processed} songs")
+if skipped_count > 0:
+    print(f"⏭️  Skipped (already completed): {skipped_count} songs")
+print(f"📊 Total completed: {newly_processed + skipped_count} songs")
+print("="*60)
+for result in all_results:
+    if result.get('output_path'):
+        print(f"Song {result['song_idx']+1}: {result['output_path']}")
+

baseline_generate/yue/mmtokenizer.py

@@ -0,0 +1,367 @@
+from abc import ABC
+from abc import abstractmethod
+
+
+class AbstractTokenizer(ABC):
+    """Abstract class for tokenizer."""
+
+    def __init__(self, name):
+        self.name = name
+        super().__init__()
+
+    @property
+    @abstractmethod
+    def vocab_size(self):
+        pass
+
+    @property
+    @abstractmethod
+    def vocab(self):
+        """Dictionary from vocab text token to id token."""
+        pass
+
+    @property
+    @abstractmethod
+    def inv_vocab(self):
+        """Dictionary from vocab id token to text token."""
+        pass
+
+    @abstractmethod
+    def tokenize(self, text):
+        pass
+
+    def detokenize(self, token_ids):
+        raise NotImplementedError('detokenizer is not implemented for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def cls(self):
+        raise NotImplementedError('CLS is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def sep(self):
+        raise NotImplementedError('SEP is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def pad(self):
+        raise NotImplementedError('PAD is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def eod(self):
+        raise NotImplementedError('EOD is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def mask(self):
+        raise NotImplementedError('MASK is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+
+class _SentencePieceTokenizer(AbstractTokenizer):
+    """SentencePieceTokenizer-Megatron wrapper"""
+
+    def __init__(self, model_file, vocab_extra_ids=0):
+        name = 'SentencePieceTokenizer'
+        super().__init__(name)
+
+        import sentencepiece
+        self.tokenizer = sentencepiece.SentencePieceProcessor(model_file=model_file)
+        self._initalize(vocab_extra_ids)
+
+    def _populate_vocab(self):
+        self._vocab = {}
+        self._inv_vocab = {}
+
+        for i in range(len(self.tokenizer)):
+            t = self.tokenizer.id_to_piece(i)
+            self._inv_vocab[i] = t
+            self._vocab[t] = i
+
+    def _initalize(self, vocab_extra_ids):
+        self._populate_vocab()
+        self._special_tokens = {}
+        self._inv_special_tokens = {}
+
+        self._t5_tokens = []
+
+        def _add_special_token(t):
+            if t not in self._vocab:
+                next_id = len(self._vocab)
+                self._vocab[t] = next_id
+                self._inv_vocab[next_id] = t
+            self._special_tokens[t] = self._vocab[t]
+            self._inv_special_tokens[self._vocab[t]] = t
+
+        _add_special_token('<CLS>')
+        self._cls_id = self._vocab['<CLS>']
+        _add_special_token('<SEP>')
+        self._sep_id = self._vocab['<SEP>']
+        _add_special_token('<EOD>')
+        self._eod_id = self._vocab['<EOD>']
+        _add_special_token('<MASK>')
+        self._mask_id = self._vocab['<MASK>']
+
+        pad_id = self.tokenizer.pad_id()
+        try:
+            pad_token = self.tokenizer.id_to_piece(pad_id)
+        except IndexError:
+            pad_token = '<PAD>'
+        _add_special_token(pad_token)
+        self._pad_id = self._vocab[pad_token]
+
+        bos_id = self.tokenizer.bos_id()
+        try:
+            bos_token = self.tokenizer.id_to_piece(bos_id)
+        except IndexError:
+            bos_token = '<BOS>'
+        _add_special_token(bos_token)
+        self._bos_id = self._vocab[bos_token]
+
+        eos_id = self.tokenizer.eos_id()
+        try:
+            eos_token = self.tokenizer.id_to_piece(eos_id)
+        except IndexError:
+            eos_token = '<EOS>'
+        _add_special_token(eos_token)
+        self._eos_id = self._vocab[eos_token]
+
+        for i in range(vocab_extra_ids):
+            t = "<extra_id_{}>".format(i)
+            _add_special_token(t)
+            self._t5_tokens += [t]
+
+    @property
+    def vocab_size(self):
+        return len(self._vocab)
+
+    @property
+    def vocab(self):
+        return self._vocab
+
+    @property
+    def inv_vocab(self):
+        return self._inv_vocab
+
+    @property
+    def decoder(self):
+        return self._inv_vocab
+
+    @property
+    def encoder(self):
+        return self._vocab
+
+    # From:
+    # https://github.com/NVIDIA/NeMo/blob/c8fa217e811d60d11d014827c7f3845ff6c99ae7/nemo/collections/common/tokenizers/sentencepiece_tokenizer.py#L89
+    def tokenize(self, text):
+        ids = []
+        idx = 0
+
+        while 1:
+            indices = {}
+            for token in self._special_tokens:
+                try:
+                    indices[token] = text[idx:].index(token)
+                except ValueError:
+                    continue
+            if len(indices) == 0:
+                break
+
+            next_token = min(indices, key=indices.get)
+            next_idx = idx + indices[next_token]
+
+            ids.extend(self.tokenizer.encode_as_ids(text[idx:next_idx]))
+            ids.append(self._special_tokens[next_token])
+            idx = next_idx + len(next_token)
+
+        ids.extend(self.tokenizer.encode_as_ids(text[idx:]))
+        return ids
+
+    # From:
+    # https://github.com/NVIDIA/NeMo/blob/c8fa217e811d60d11d014827c7f3845ff6c99ae7/nemo/collections/common/tokenizers/sentencepiece_tokenizer.py#L125
+    def detokenize(self, ids):
+        text = ""
+        last_i = 0
+
+        for i, id in enumerate(ids):
+            if id in self._inv_special_tokens:
+                text += self.tokenizer.decode_ids(ids[last_i:i]) + " "
+                text += self._inv_special_tokens[id] + " "
+                last_i = i + 1
+
+        text += self.tokenizer.decode_ids(ids[last_i:])
+        return text
+
+    @property
+    def cls(self):
+        return self._cls_id
+
+    @property
+    def sep(self):
+        return self._sep_id
+
+    @property
+    def pad(self):
+        return self._pad_id
+
+    @property
+    def bos_token_id(self):
+        return self._bos_id
+
+    @property
+    def bos(self):
+        return self._bos_id
+
+    @property
+    def eod(self):
+        return self._eod_id
+
+    @property
+    def eos_token_id(self):
+        return self._eos_id
+
+    @property
+    def eos(self):
+        return self._eos_id
+
+    @property
+    def mask(self):
+        return self._mask_id
+
+    @property
+    def additional_special_tokens_ids(self):
+        return [self.vocab[k] for k in self._t5_tokens]
+
+class _MMSentencePieceTokenizer(_SentencePieceTokenizer):
+    """SentencePieceTokenizer-Megatron wrapper"""
+
+    def __init__(self, model_file, vocab_extra_ids=0):
+        super().__init__(model_file, vocab_extra_ids)
+
+
+    def _initalize(self, vocab_extra_ids):
+        self._populate_vocab()
+        self._special_tokens = {}
+        self._inv_special_tokens = {}
+
+        self._t5_tokens = []
+
+        def _add_special_token(t):
+            if t not in self._vocab:
+                next_id = len(self._vocab)
+                self._vocab[t] = next_id
+                self._inv_vocab[next_id] = t
+            self._special_tokens[t] = self._vocab[t]
+            self._inv_special_tokens[self._vocab[t]] = t
+
+        _add_special_token('<CLS>')
+        self._cls_id = self._vocab['<CLS>']
+        _add_special_token('<SEP>')
+        self._sep_id = self._vocab['<SEP>']
+        _add_special_token('<EOD>')
+        self._eod_id = self._vocab['<EOD>']
+        _add_special_token('<MASK>')
+        self._mask_id = self._vocab['<MASK>']
+
+        _add_special_token('<SOA>')
+        self._soa_id = self._vocab['<SOA>']
+        _add_special_token('<EOA>')
+        self._eoa_id = self._vocab['<EOA>']
+        _add_special_token('<SOV>')
+        self._sov_id = self._vocab['<SOV>']
+        _add_special_token('<EOV>')
+        self._eov_id = self._vocab['<EOV>']
+        _add_special_token('<SOI>')
+        self._soi_id = self._vocab['<SOI>']
+        _add_special_token('<EOI>')
+        self._eoi_id = self._vocab['<EOI>']
+        _add_special_token('<s_local>')
+        self._s_local_id = self._vocab['<s_local>']
+        _add_special_token('<e_local>')
+        self._e_local_id = self._vocab['<e_local>']
+        _add_special_token('<s_global>')
+        self._s_global_id = self._vocab['<s_global>']
+        _add_special_token('<e_global>')
+        self._e_global_id = self._vocab['<e_global>']
+        _add_special_token('<stage_1>')
+        self._stage_1_id = self._vocab['<stage_1>']
+        _add_special_token('<stage_2>')
+        self._stage_2_id = self._vocab['<stage_2>']
+        pad_id = self.tokenizer.pad_id()
+        try:
+            pad_token = self.tokenizer.id_to_piece(pad_id)
+        except IndexError:
+            pad_token = '<PAD>'
+        _add_special_token(pad_token)
+        self._pad_id = self._vocab[pad_token]
+
+        bos_id = self.tokenizer.bos_id()
+        try:
+            bos_token = self.tokenizer.id_to_piece(bos_id)
+        except IndexError:
+            bos_token = '<BOS>'
+        _add_special_token(bos_token)
+        self._bos_id = self._vocab[bos_token]
+
+        eos_id = self.tokenizer.eos_id()
+        try:
+            eos_token = self.tokenizer.id_to_piece(eos_id)
+        except IndexError:
+            eos_token = '<EOS>'
+        _add_special_token(eos_token)
+        self._eos_id = self._vocab[eos_token]
+
+        for i in range(vocab_extra_ids):
+            t = "<extra_id_{}>".format(i)
+            _add_special_token(t)
+            self._t5_tokens += [t]
+    
+    @property
+    def soa(self):
+        return self._soa_id
+    
+    @property
+    def eoa(self):
+        return self._eoa_id
+    
+    @property
+    def sov(self):
+        return self._sov_id
+    
+    @property
+    def eov(self):
+        return self._eov_id
+    
+    @property
+    def soi(self):
+        return self._soi_id
+    
+    @property
+    def eoi(self):
+        return self._eoi_id
+    
+    @property
+    def s_local(self):
+        return self._s_local_id
+    
+    @property
+    def e_local(self):
+        return self._e_local_id
+    
+    @property
+    def s_global(self):
+        return self._s_global_id
+
+    @property
+    def e_global(self):
+        return self._e_global_id
+    
+    @property
+    def stage_1(self):
+        return self._stage_1_id
+
+    @property
+    def stage_2(self):
+        return self._stage_2_id

data_pipeline/lyrics_gene/filter_all_cn.py

@@ -0,0 +1,272 @@
+import json
+import os
+import re
+import torch
+from sentence_transformers import SentenceTransformer, util
+from tqdm import tqdm
+
+# os.environ["HTTP_PROXY"] = "http://127.0.0.1:7890"
+# os.environ["HTTPS_PROXY"] = "http://127.0.0.1:7890"
+
+# Use HuggingFace mirror site
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+# Set HuggingFace cache directory so SentenceTransformer can recognize downloaded models
+os.environ["HF_HOME"] = os.path.expanduser("~/.cache/huggingface")
+
+
+def clean_lyrics(lyrics):
+    """
+    Clean lyrics by removing segment tags, timestamp tags, and newlines, keeping only pure lyric text
+    
+    Args:
+        lyrics: Raw lyrics text (contains segment tags like [Verse 1], timestamps like [00:07.00], and newlines)
+        
+    Returns:
+        Cleaned lyrics text (plain text, no tags and newlines)
+    """
+    # Use regex to remove all [tag] format content (including segment tags and timestamps)
+    # Pattern matching [any content]
+    cleaned = re.sub(r'\[.*?\]', '', lyrics)
+    
+    # Remove all newlines, replace with spaces
+    cleaned = cleaned.replace('\n', ' ')
+    
+    # Remove extra spaces (replace multiple consecutive spaces with single space)
+    cleaned = re.sub(r'\s+', ' ', cleaned)
+    
+    # Remove leading and trailing spaces
+    cleaned = cleaned.strip()
+    
+    return cleaned
+
+
+def load_music_data(input_file, max_count=None):
+    """
+    Load music data from jsonl file
+    
+    Args:
+        input_file: Path to input jsonl file
+        max_count: Maximum number to read, None means read all
+        
+    Returns:
+        List of music data
+    """
+    music_list = []
+    print(f"Loading music data: {input_file}")
+    if max_count:
+        print(f"Limiting to first {max_count} songs")
+    
+    with open(input_file, 'r', encoding='utf-8') as f:
+        for line in tqdm(f, desc="Loading data"):
+            try:
+                data = json.loads(line.strip())
+                # Ensure required fields are present
+                if 'description' in data and 'lyrics' in data:
+                    music_list.append(data)
+                    # If reached maximum count, stop reading
+                    if max_count and len(music_list) >= max_count:
+                        break
+            except json.JSONDecodeError:
+                continue
+    print(f"Successfully loaded {len(music_list)} songs")
+    return music_list
+
+
+def deduplicate_music(music_list, texts, model, threshold=0.90, output_file=None, save_interval=10000, matrix_save_dir=None):
+    """
+    Deduplicate music data based on text similarity
+    
+    Args:
+        music_list: List of music data
+        texts: List of texts for comparison
+        model: SentenceTransformer model
+        threshold: Similarity threshold
+        output_file: Output file path, if provided supports incremental saving
+        save_interval: Save every N valid songs processed
+        matrix_save_dir: Directory to save matrices
+        
+    Returns:
+        Deduplicated music data list
+    """
+    print(f"Computing embeddings for {len(texts)} texts...")
+    embeddings = model.encode(texts, convert_to_tensor=True, show_progress_bar=True)
+    
+    print("Computing similarity matrix...")
+    cos_scores = util.pytorch_cos_sim(embeddings, embeddings)
+    
+    # Save similarity matrix and embeddings
+    if matrix_save_dir:
+        os.makedirs(matrix_save_dir, exist_ok=True)
+        embeddings_path = os.path.join(matrix_save_dir, 'embeddings.pt')
+        cos_scores_path = os.path.join(matrix_save_dir, 'cos_scores.pt')
+        print(f"Saving embeddings to: {embeddings_path}")
+        torch.save(embeddings.cpu(), embeddings_path)
+        print(f"Saving similarity matrix to: {cos_scores_path}")
+        torch.save(cos_scores.cpu(), cos_scores_path)
+        print("Matrix saving complete!")
+    
+    print(f"Deduplicating (threshold: {threshold})...")
+    keep_idx = []
+    removed = set()
+    
+    # If output file provided, open in write mode
+    f = None
+    if output_file:
+        os.makedirs(os.path.dirname(output_file), exist_ok=True)
+        f = open(output_file, 'w', encoding='utf-8')
+    
+    saved_count = 0
+    
+    for i in tqdm(range(len(music_list)), desc="Deduplication progress"):
+        if i in removed:
+            continue
+        keep_idx.append(i)
+        
+        # If incremental saving enabled, save every save_interval songs
+        if f and len(keep_idx) - saved_count >= save_interval:
+            # Save all valid songs from saved_count to current
+            for idx in range(saved_count, len(keep_idx)):
+                music = music_list[keep_idx[idx]]
+                f.write(json.dumps(music, ensure_ascii=False) + '\n')
+            f.flush()  # Ensure write to disk
+            saved_count = len(keep_idx)
+            print(f"Saved {saved_count} valid songs to file")
+        
+        for j in range(i+1, len(music_list)):
+            if cos_scores[i][j] > threshold:
+                removed.add(j)
+    
+    # Save remaining valid songs
+    if f:
+        for idx in range(saved_count, len(keep_idx)):
+            music = music_list[keep_idx[idx]]
+            f.write(json.dumps(music, ensure_ascii=False) + '\n')
+        f.close()
+        print(f"Saved all {len(keep_idx)} valid songs to file")
+    
+    deduped_music_list = [music_list[i] for i in keep_idx]
+    print(f"Deduplication complete: {len(music_list)} -> {len(deduped_music_list)} (removed {len(removed)} songs)")
+    
+    return deduped_music_list
+
+
+def dedup_by_description_and_lyrics(input_file, output_file, threshold=0.95, max_count=None, device='cuda:1', save_interval=10000, matrix_save_dir=None):
+    """
+    Method 2: Deduplicate based on description + lyrics
+    
+    Args:
+        input_file: Path to input jsonl file
+        output_file: Path to output jsonl file
+        threshold: Similarity threshold
+        max_count: Maximum number to read, None means read all
+        device: Device to use, default cuda:1 (GPU1)
+        save_interval: Save every N valid songs processed, default 10000
+        matrix_save_dir: Directory to save matrices, if provided saves embeddings and similarity matrix
+    """
+    print("\n========== Method 2: Deduplicate based on description + lyrics ==========")
+    print(f"Using device: {device}")
+    
+    # Load data
+    music_list = load_music_data(input_file, max_count=max_count)
+    
+    # Extract combined text from description + lyrics
+    combined_texts = []
+    for music in music_list:
+        description = music.get('description', '')
+        lyrics = music.get('lyrics', '')
+        # Clean lyrics, remove structure tags
+        cleaned_lyrics = clean_lyrics(lyrics)
+        # Concatenate description and cleaned lyrics (separated by delimiter)
+        combined_text = f"{description} [SEP] {cleaned_lyrics}"
+        combined_texts.append(combined_text)
+    
+    # Load Chinese model and specify device
+    # Check if local model exists, if so use local path directly to avoid re-downloading
+    hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
+    model_cache_dir = os.path.join(hf_home, "hub", "models--shibing624--text2vec-bge-large-chinese", "snapshots")
+    
+    # Find model snapshot directory
+    local_model_path = None
+    if os.path.exists(model_cache_dir):
+        snapshots = [d for d in os.listdir(model_cache_dir) if os.path.isdir(os.path.join(model_cache_dir, d))]
+        if snapshots:
+            # Use latest snapshot (usually unique)
+            local_model_path = os.path.join(model_cache_dir, snapshots[0])
+            if os.path.exists(os.path.join(local_model_path, "config.json")):
+                print(f"Detected local model, using path: {local_model_path}")
+                model = SentenceTransformer(local_model_path, device=device)
+            else:
+                local_model_path = None
+    
+    if local_model_path is None:
+        print(f"Loading model to {device}...")
+        model = SentenceTransformer('shibing624/text2vec-bge-large-chinese', device=device)
+    
+    # Deduplicate (supports incremental saving)
+    deduped_music_list = deduplicate_music(
+        music_list, 
+        combined_texts, 
+        model, 
+        threshold, 
+        output_file=output_file,
+        save_interval=save_interval,
+        matrix_save_dir=matrix_save_dir
+    )
+    
+    # Deduplication function already handled saving, just print info here
+    if output_file:
+        print(f"✓ Save complete! Remaining {len(deduped_music_list)} songs after deduplication\n")
+    else:
+        # If no output file provided, save once here (compatibility with old code)
+        print(f"Saving results to: {output_file}")
+        os.makedirs(os.path.dirname(output_file), exist_ok=True)
+        with open(output_file, 'w', encoding='utf-8') as f:
+            for music in deduped_music_list:
+                f.write(json.dumps(music, ensure_ascii=False) + '\n')
+        print(f"✓ Save complete! Remaining {len(deduped_music_list)} songs after deduplication\n")
+    
+    return deduped_music_list
+
+
+if __name__ == '__main__':
+    # Input file path
+    input_file = 'lrc_4w_single_pro_des.jsonl'
+    
+    # Output file path
+    output_file = 'filter_all_4w.jsonl'
+    
+    # Matrix save directory
+    matrix_save_dir = 'generate_lrc'
+    
+    # Set maximum read count (for testing, None means read all)
+    max_count = None  # Test first 5 songs
+        
+    # Deduplicate based on description + lyrics
+    print("\nDeduplicating based on description + lyrics")
+    dedup_by_description_and_lyrics(
+        input_file, 
+        output_file, 
+        threshold=0.90, 
+        max_count=max_count, 
+        device='cuda:7',
+        save_interval=10000,  # Save every 10000 valid songs
+        matrix_save_dir=matrix_save_dir  # Save similarity matrix
+    )
+    print(f"\nComplete! Results saved to: {output_file}")
+    print(f"Similarity matrix saved to: {matrix_save_dir}")
+    # Test lyrics cleaning effect
+    # print("\n========== Test Lyrics Cleaning Effect ==========")
+    # music_list = load_music_data(input_file, max_count=max_count)
+    
+    # print("\n" + "="*80)
+    # for i, music in enumerate(music_list, 1):
+    #     print(f"\n[Song {i}]")
+    #     print(f"Description: {music.get('description', '')}")
+    #     print("\n--- Original Lyrics ---")
+    #     original_lyrics = music.get('lyrics', '')
+    #     print(original_lyrics[:500] + "..." if len(original_lyrics) > 500 else original_lyrics)
+    #     print("\n--- Cleaned Lyrics ---")
+    #     cleaned_lyrics = clean_lyrics(original_lyrics)
+    #     print(cleaned_lyrics)
+    #     print("\n" + "-"*80)
+

data_pipeline/lyrics_gene/filter_all_en.py

@@ -0,0 +1,256 @@
+import json
+import os
+import re
+import torch
+from sentence_transformers import SentenceTransformer, util
+from tqdm import tqdm
+
+# os.environ["HTTP_PROXY"] = "http://127.0.0.1:7890"
+# os.environ["HTTPS_PROXY"] = "http://127.0.0.1:7890"
+
+# Use HuggingFace mirror site
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+# Set HuggingFace cache directory so SentenceTransformer can recognize downloaded models
+os.environ["HF_HOME"] = os.path.expanduser("~/.cache/huggingface")
+
+MODEL_EN = "sentence-transformers/all-MiniLM-L6-v2"
+
+
+def clean_lyrics(lyrics):
+    """
+    Clean lyrics by removing segment tags, timestamp tags, and newlines, keeping only pure lyric text
+    
+    Args:
+        lyrics: Raw lyrics text (contains segment tags like [Verse 1], timestamps like [00:07.00], and newlines)
+        
+    Returns:
+        Cleaned lyrics text (plain text, no tags and newlines)
+    """
+    # Use regex to remove all [tag] format content (including segment tags and timestamps)
+    # Pattern matching [any content]
+    cleaned = re.sub(r'\[.*?\]', '', lyrics)
+    
+    # Remove all newlines, replace with spaces
+    cleaned = cleaned.replace('\n', ' ')
+    
+    # Remove extra spaces (replace multiple consecutive spaces with single space)
+    cleaned = re.sub(r'\s+', ' ', cleaned)
+    
+    # Remove leading and trailing spaces
+    cleaned = cleaned.strip()
+    
+    return cleaned
+
+
+def load_music_data(input_file, max_count=None):
+    """
+    Load music data from jsonl file
+    
+    Args:
+        input_file: Path to input jsonl file
+        max_count: Maximum number to read, None means read all
+        
+    Returns:
+        List of music data
+    """
+    music_list = []
+    print(f"Loading music data: {input_file}")
+    if max_count:
+        print(f"Limiting to first {max_count} songs")
+    
+    with open(input_file, 'r', encoding='utf-8') as f:
+        for line in tqdm(f, desc="Loading data"):
+            try:
+                data = json.loads(line.strip())
+                # Ensure required fields are present
+                if 'description' in data and 'lyrics' in data:
+                    music_list.append(data)
+                    # If reached maximum count, stop reading
+                    if max_count and len(music_list) >= max_count:
+                        break
+            except json.JSONDecodeError:
+                continue
+    print(f"Successfully loaded {len(music_list)} songs")
+    return music_list
+
+
+def deduplicate_music(music_list, texts, model, threshold=0.90, output_file=None, save_interval=10000, matrix_save_dir=None):
+    """
+    Deduplicate music data based on text similarity
+    
+    Args:
+        music_list: List of music data
+        texts: List of texts for comparison
+        model: SentenceTransformer model
+        threshold: Similarity threshold
+        output_file: Output file path, if provided supports incremental saving
+        save_interval: Save every N valid songs processed
+        matrix_save_dir: Directory to save matrices
+        
+    Returns:
+        Deduplicated music data list
+    """
+    print(f"Computing embeddings for {len(texts)} texts...")
+    embeddings = model.encode(texts, convert_to_tensor=True, show_progress_bar=True)
+    
+    print("Computing similarity matrix...")
+    cos_scores = util.pytorch_cos_sim(embeddings, embeddings)
+    
+    # Save similarity matrix and embeddings
+    if matrix_save_dir:
+        os.makedirs(matrix_save_dir, exist_ok=True)
+        embeddings_path = os.path.join(matrix_save_dir, 'embeddings.pt')
+        cos_scores_path = os.path.join(matrix_save_dir, 'cos_scores.pt')
+        print(f"Saving embeddings to: {embeddings_path}")
+        torch.save(embeddings.cpu(), embeddings_path)
+        print(f"Saving similarity matrix to: {cos_scores_path}")
+        torch.save(cos_scores.cpu(), cos_scores_path)
+        print("Matrix saving complete!")
+    
+    print(f"Deduplicating (threshold: {threshold})...")
+    keep_idx = []
+    removed = set()
+    
+    # If output file provided, open in write mode
+    f = None
+    if output_file:
+        os.makedirs(os.path.dirname(output_file), exist_ok=True)
+        f = open(output_file, 'w', encoding='utf-8')
+    
+    saved_count = 0
+    
+    for i in tqdm(range(len(music_list)), desc="Deduplication progress"):
+        if i in removed:
+            continue
+        keep_idx.append(i)
+        
+        # If incremental saving enabled, save every save_interval songs
+        if f and len(keep_idx) - saved_count >= save_interval:
+            # Save all valid songs from saved_count to current
+            for idx in range(saved_count, len(keep_idx)):
+                music = music_list[keep_idx[idx]]
+                f.write(json.dumps(music, ensure_ascii=False) + '\n')
+            f.flush()  # Ensure write to disk
+            saved_count = len(keep_idx)
+            print(f"Saved {saved_count} valid songs to file")
+        
+        for j in range(i+1, len(music_list)):
+            if cos_scores[i][j] > threshold:
+                removed.add(j)
+    
+    # Save remaining valid songs
+    if f:
+        for idx in range(saved_count, len(keep_idx)):
+            music = music_list[keep_idx[idx]]
+            f.write(json.dumps(music, ensure_ascii=False) + '\n')
+        f.close()
+        print(f"Saved all {len(keep_idx)} valid songs to file")
+    
+    deduped_music_list = [music_list[i] for i in keep_idx]
+    print(f"Deduplication complete: {len(music_list)} -> {len(deduped_music_list)} (removed {len(removed)} songs)")
+    
+    return deduped_music_list
+
+
+def dedup_by_description_and_lyrics(input_file, output_file, threshold=0.95, max_count=None, device='cuda:1', save_interval=10000, matrix_save_dir=None):
+    """
+    Method 2: Deduplicate based on description + lyrics
+    
+    Args:
+        input_file: Path to input jsonl file
+        output_file: Path to output jsonl file
+        threshold: Similarity threshold
+        max_count: Maximum number to read, None means read all
+        device: Device to use, default cuda:1 (GPU1)
+        save_interval: Save every N valid songs processed, default 10000
+        matrix_save_dir: Directory to save matrices, if provided saves embeddings and similarity matrix
+    """
+    print("\n========== Method 2: Deduplicate based on description + lyrics ==========")
+    print(f"Using device: {device}")
+    
+    # Load data
+    music_list = load_music_data(input_file, max_count=max_count)
+    
+    # Extract combined text from description + lyrics
+    combined_texts = []
+    for music in music_list:
+        description = music.get('description', '')
+        lyrics = music.get('lyrics', '')
+        # Clean lyrics, remove structure tags
+        cleaned_lyrics = clean_lyrics(lyrics)
+        # Concatenate description and cleaned lyrics (separated by delimiter)
+        combined_text = f"{description} [SEP] {cleaned_lyrics}"
+        combined_texts.append(combined_text)
+    
+    # Load English model
+    print(f"Loading English model {MODEL_EN} to {device}...")
+    model = SentenceTransformer(MODEL_EN, device=device)
+    
+    # Deduplicate (supports incremental saving)
+    deduped_music_list = deduplicate_music(
+        music_list, 
+        combined_texts, 
+        model, 
+        threshold, 
+        output_file=output_file,
+        save_interval=save_interval,
+        matrix_save_dir=matrix_save_dir
+    )
+    
+    # Deduplication function already handled saving, just print info here
+    if output_file:
+        print(f"✓ Save complete! Remaining {len(deduped_music_list)} songs after deduplication\n")
+    else:
+        # If no output file provided, save once here (compatibility with old code)
+        print(f"Saving results to: {output_file}")
+        os.makedirs(os.path.dirname(output_file), exist_ok=True)
+        with open(output_file, 'w', encoding='utf-8') as f:
+            for music in deduped_music_list:
+                f.write(json.dumps(music, ensure_ascii=False) + '\n')
+        print(f"✓ Save complete! Remaining {len(deduped_music_list)} songs after deduplication\n")
+    
+    return deduped_music_list
+
+
+if __name__ == '__main__':
+    # Input file path
+    input_file = 'en_lrc_4w_single_pro_des.jsonl'
+    
+    # Output file path
+    output_file = 'filter_en_single_4w(0.9).jsonl'
+    
+    # Matrix save directory
+    matrix_save_dir = 'en_matrix'
+    
+    # Set maximum read count (for testing, None means read all)
+    max_count = None  # Test first 5 songs
+        
+    # Deduplicate based on description + lyrics
+    print("\nDeduplicating based on description + lyrics")
+    dedup_by_description_and_lyrics(
+        input_file, 
+        output_file, 
+        threshold=0.90, 
+        max_count=max_count, 
+        device='cuda:7',
+        save_interval=10000,  # Save every 10000 valid songs
+        matrix_save_dir=matrix_save_dir  # Save similarity matrix
+    )
+    print(f"\nComplete! Results saved to: {output_file}")
+    print(f"Similarity matrix saved to: {matrix_save_dir}")
+    # Test lyrics cleaning effect
+    # print("\n========== Test Lyrics Cleaning Effect ==========")
+    # music_list = load_music_data(input_file, max_count=max_count)
+    
+    # print("\n" + "="*80)
+    # for i, music in enumerate(music_list, 1):
+    #     print(f"\n[Song {i}]")
+    #     print(f"Description: {music.get('description', '')}")
+    #     print("\n--- Original Lyrics ---")
+    #     original_lyrics = music.get('lyrics', '')
+    #     print(original_lyrics[:500] + "..." if len(original_lyrics) > 500 else original_lyrics)
+    #     print("\n--- Cleaned Lyrics ---")
+    #     cleaned_lyrics = clean_lyrics(original_lyrics)
+    #     print(cleaned_lyrics)
+    #     print("\n" + "-"*80)
+

data_pipeline/lyrics_gene/gen_lyrics_cn.py

@@ -0,0 +1,568 @@
+import os
+import json
+import time
+import random
+import re
+from openai import OpenAI
+from tqdm import tqdm
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+# Set environment variables
+# Note: Set these environment variables before running the script
+# export OPENAI_API_KEY="your-api-key"
+# export OPENAI_BASE_URL="https://api.openai.com/v1"  # or your custom API URL
+if not os.environ.get("OPENAI_API_KEY"):
+    os.environ["OPENAI_API_KEY"] = ""  # Replace with your API key or set via environment variable
+if not os.environ.get("OPENAI_BASE_URL"):
+    os.environ["OPENAI_BASE_URL"] = "https://api.openai.com/v1"  # Replace with API URL or set via environment variable
+
+# Initialize client
+client = OpenAI()
+
+
+def _extract_lyrics_timestamps(lyrics_text):
+    """
+    Extract timestamps from lyrics and convert to seconds
+    Args:
+        lyrics_text: Lyrics string
+    Returns:
+        List[float]: Timestamps in order (seconds)
+    """
+    if not isinstance(lyrics_text, str):
+        return []
+    pattern = re.compile(r'\[(\d{2}):(\d{2})(?:\.(\d{2}))?\]')
+    timestamps = []
+    for match in pattern.finditer(lyrics_text):
+        minutes = int(match.group(1))
+        seconds = int(match.group(2))
+        fraction = match.group(3)
+        total_seconds = minutes * 60 + seconds
+        if fraction is not None:
+            divisor = 100 if len(fraction) == 2 else 10 ** len(fraction)
+            total_seconds += int(fraction) / divisor
+        timestamps.append(total_seconds)
+    return timestamps
+
+
+def _validate_timestamps(lyrics_text, min_last_timestamp=170, max_interval=35):
+    """
+    Validate if lyrics timestamps meet requirements
+    Args:
+        lyrics_text: Lyrics string
+        min_last_timestamp: Minimum value of last timestamp (seconds)
+        max_interval: Maximum interval between last two timestamps (seconds)
+    Returns:
+        bool: Whether validation passes
+    """
+    timestamps = _extract_lyrics_timestamps(lyrics_text)
+    if len(timestamps) < 2:
+        print("Validation failed: Timestamp count less than 2")
+        return False
+    last = timestamps[-1]
+    second_last = timestamps[-2]
+    if last < min_last_timestamp:
+        print(f"Validation failed: Last timestamp {last:.2f}s is less than {min_last_timestamp}s")
+        return False
+    if last - second_last > max_interval:
+        print(f"Validation failed: Interval between last two timestamps {last - second_last:.2f}s is greater than {max_interval}s")
+        return False
+    return True
+
+
+def chat_gpt(text, model='gpt-4o-mini'):
+    while True:
+        try:
+            # Call OpenAI chat completions API
+            completion = client.chat.completions.create(
+                model=model,  # Use GPT-4o-mini model
+                messages=[
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": text}
+                ]
+            )
+            # Get reply content
+            if getattr(completion.choices[0].message, 'content', None):
+                content = completion.choices[0].message.content.strip()
+                return content
+            else:
+                print('error_wait_2s')
+        except Exception as e:
+            print(f"Error: {e}")
+        time.sleep(2)
+
+
+def chat_gpt_call(text, model='gpt-4o-mini'):
+    # Call OpenAI chat completions API
+    completion = client.chat.completions.create(
+        model=model,  # Use GPT-4o-mini model
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": text}
+        ]
+    )
+    # Get reply content
+    if getattr(completion.choices[0].message, 'content', None):
+        content = completion.choices[0].message.content.strip()
+        return content
+    else:
+        print('error_wait_2s')
+
+
+def generate_music_descriptions(all_music_data, index_pool, output_file, file_lock, sample_size=20, model='gpt-4o-mini', max_retries=0):
+    """
+    Read music data file, randomly sample, call GPT to generate new music descriptions and lyrics
+    
+    Args:
+        all_music_data: All music data list
+        index_pool: Index pool object (thread-safe)
+        output_file: Output jsonl file path
+        file_lock: File write lock
+        sample_size: Number of random samples
+        model: Model name to use
+        max_retries: Maximum retry count
+    
+    Returns:
+        (used_indices, success_count): List of used indices and count of successful generations
+    """
+    # duration_ranges = [
+    #     ("3.00", "3.15", 50), ("3.15", "3.30", 50), ("3.30", "3.45", 60),
+    #  ("3.45", "4.00", 60),
+    #     ("4.00", "4.15", 70), ("4.15", "4.30", 70), ("4.30", "4.45", 70)
+    # ]
+    duration_ranges = [
+        ("3.00", "3.15", 60), ("3.15", "3.30", 70), ("3.30", "3.45", 80),("3.45", "4.00", 90),
+        ("4.00", "4.15", 100), ("4.15", "4.30", 100), ("4.30", "4.45", 100)
+    ]
+    selected_range = random.choice(duration_ranges)
+    require_length = selected_range[2]
+    
+    # Directly convert to timestamp format (strictly corresponds to left and right ends of tuple)
+    start_timestamp = f"[{selected_range[0].replace('.', ':')}.00]"
+    end_timestamp = f"[{selected_range[1].replace('.', ':')}.00]"
+    
+    # Generate duration description
+    start_duration = f"{selected_range[0].replace('.', '分')}秒"
+    end_duration = f"{selected_range[1].replace('.', '分')}秒"
+    
+    # Generate random timestamps for examples (randomly generated within time range)
+    # Parse time string to minutes and seconds
+    start_parts = selected_range[0].split('.')
+    end_parts = selected_range[1].split('.')
+    
+    start_minutes = int(start_parts[0])
+    start_seconds = int(start_parts[1])
+    end_minutes = int(end_parts[0])
+    end_seconds = int(end_parts[1])
+    
+    # Convert to total seconds
+    start_total_seconds = start_minutes * 60 + start_seconds
+    end_total_seconds = end_minutes * 60 + end_seconds
+    
+    # Randomly generate within range
+    example1_seconds = random.randint(start_total_seconds, end_total_seconds)
+    example2_seconds = random.randint(start_total_seconds, end_total_seconds)
+    
+    example1_minutes = example1_seconds // 60
+    example1_secs = example1_seconds % 60
+    example2_minutes = example2_seconds // 60
+    example2_secs = example2_seconds % 60
+    
+    example1_timestamp = f"[{example1_minutes:02d}:{example1_secs:02d}.00]"
+    example2_timestamp = f"[{example2_minutes:02d}:{example2_secs:02d}.00]"
+    
+    # Get sample indices from index pool (thread-safe)
+    selected_indices = index_pool.get_indices(sample_size)
+    if not selected_indices:
+        return [], 0
+    
+    sample_data = [all_music_data[i] for i in selected_indices]
+    
+    # Extract all unique styles
+    styles = []
+    for data in sample_data:
+        style = data.get('style', '')
+        if style and style not in styles:
+            styles.append(style)
+    
+    styles_text = "、".join(styles)
+    
+    # Build example text - include all sampled data (excluding style)
+    examples = []
+    for i, data in enumerate(sample_data, 1):  
+        lyrics_text = " ".join(data.get('lyrics', [])) if isinstance(data.get('lyrics'), list) else data.get('lyrics', '')
+        description = data.get('description', '')
+        examples.append(f"示例{i}:\ndescription: {description}\nlyrics: {lyrics_text}")
+    
+    examples_text = "\n\n".join(examples)
+
+    prompt = f"""生成2首完整的歌曲，每首歌必须满足以下硬性指标：
+- 严禁生成小于{require_length}行的歌词！
+- 每首歌的歌词行数必须严格大于{require_length}行，这是硬性要求！
+- 最后一句时间戳必须在{start_timestamp}到{end_timestamp}之间
+- 两首歌的时长、行数必须有差异，严禁最后的时间戳均相同
+- 相邻歌词行的时间戳间隔不得超过10秒！必须保证时间戳连续自然递进
+- 严禁出现如"[03:25.00]在心中[04:25.00]最后一行歌词"的生硬间隔,严禁超过10s的间隔
+如果生成的歌曲不满足以上任意一项，则视为不合格，请重新生成。
+请生成2首新的、具有多样性的音乐描述和LRC格式歌词，语言为中文。
+
+
+创作要求：
+1.风格与流派需要确保多样性。
+2.Description标签化要求（必须严格遵守）：
+   description字段必须使用结构化的标签格式，包括以下标签，用逗号分隔：
+   - 音乐风格标签
+   - 音乐流派标签
+   - 乐器标签
+   - 情感基调标签
+   - 氛围标签
+   - 演唱方式和人声标签，仅限男声或女生二选一，单人独唱
+   注意：每个标签简洁明了，多个同类标签可用斜杠分隔（如"钢琴/小提琴"）
+3.歌词创造力：lyrics应该具有深度和艺术性：
+   - 主题可以涉及爱情、人生、社会、自然、哲思、梦想、回忆等各个方面
+   - 运用丰富的文学手法：比喻、意象、对比、排比等
+   - 情感真挚，注重韵律和节奏感
+   - 可以是叙事性、抒情性或意识流风格
+4.歌词结构和长度要求（必须严格遵守）：
+   - lyrics必须按照以下结构组织，并使用段落标签标注每个部分
+   - 结构顺序必须严格遵循该顺序,共8个段落标签：[Verse 1]主歌1 → [Pre-Chorus]预副歌 → [Chorus]副歌 → [Verse 2]主歌2 → [Pre-Chorus]预副歌 → [Chorus]副歌 → [Bridge]桥段 → [Chorus (Outro)]副歌（结尾）
+   - 一首歌词的段落标签只有8个，即[Verse 1]和[Verse 2]只出现一次,[Pre-Chorus]和[Chorus]各出现两次，[Bridge]和[Chorus (Outro)]各出现一次，禁止额外添加或重复更多的段落标签
+   - 每个段落标签（如[Verse 1]、[Chorus]等）必须独占一行，后面紧跟该段落的LRC格式歌词
+   - 段落之间用空行分隔
+   - **总行数要求**：整首歌必须包含至少{require_length}行带时间戳的歌词（不包括段落标签行和空行）
+5.LRC格式强制规则（必须��格遵守）：
+   - 每行歌词格式必须为 `[mm:ss.xx]歌词内容`，时间戳与歌词间无空格，歌词内容需完整连贯
+   - **每一行只能包含一小句歌词**，遇到逗号、句号等标点符号时必须换行。
+   - **严禁将多句歌词合并在同一行**
+   - 时间戳需自然分布，**第一句歌词起始时间不得为 [00:00.00]**，需考虑前奏空白（建议从[00:05.00]到[00:15.00]之间开始）
+   - 时间戳间隔要求多样性：每首歌内部的时间戳间隔必须多样化，多采用小数点数间隔，严禁使用固定间隔：
+     * 同一首歌内必须包含多种不同的间隔，不要所有句子都使用相同间隔（如不要全部都是4秒间隔）
+     * 根据歌词内容的情感强度和音乐节拍来动态调整间隔
+     * 相邻歌词行的间隔应该有所变化，体现音乐的节奏起伏
+   - 时间戳分配应根据歌曲的风格、情感、节奏来合理推测，而非机械地按照歌词长度分配
+   - 每行歌词长度应自然变化，切勿长度一致
+   - **歌曲总时长必须达到{start_duration}到{end_duration}（即最后一句时间戳必须在{start_timestamp}到{end_timestamp}之间）这是硬性要求！**
+6.歌词长度要求：lyrics字段的歌词行数必须大于{require_length}行，若生成长度过短请重新生成。
+7.独特性和原创性：每首作品都应该是独一无二的，避免简单重复示例的内容。
+8.格式要求：
+   - 直接返回JSON数组格式，包含2个歌曲对象，每个对象只有description和lyrics两个字段
+   - description字段：必须是标签格式，不是叙述性文本
+   - lyrics字段：带段落标签的LRC格式字符串
+   - 严禁在JSON中插入任何额外的符号、标记、注释或说明文字
+
+LRC格式示例（带段落标签）：
+[Verse 1]
+[00:08.00]第一句歌词
+[00:12.50]第二句歌词
+[00:17.20]第三句歌词
+
+[Pre-Chorus]
+[00:22.00]预副歌歌词
+[00:26.50]预副歌歌词
+
+[Chorus]
+[00:31.00]副歌歌词
+[00:35.50]副歌歌词
+
+负面示例（禁止出现）：
+- 错误：[01:30.00](钢琴间奏) - 禁止在时间戳后加括号注释
+- 错误：[00:00.00]开始的歌词 - 第一句不能从00:00.00开始
+- 错误: [00:05.00]在那片熟悉的田野，阳光洒满金色的麦穗 - 严禁多句歌词放在同一行
+
+现在，请充分发挥你的创造力，生成2首全新的、完整的音乐描述和LRC格式歌词作品。
+特别提醒：每首歌必须是完整歌曲，不要缩写或省略！必须包含完整的8个段落（Verse 1, Pre-Chorus, Chorus, Verse 2, Pre-Chorus, Chorus, Bridge, Chorus Outro），严格确保大于{require_length}行歌词。
+
+直接返回JSON数组格式：
+[
+  {{"description": "...", "lyrics": "..."}},
+  {{"description": "...", "lyrics": "..."}}
+]"""
+    # Try to generate with retry mechanism
+    for attempt in range(max_retries + 1):
+        try:
+            # Call OpenAI API
+            completion = client.chat.completions.create(
+                model=model,
+                messages=[
+                    {"role": "system", "content": f"You are a creative music lyricist and composer. Please generate diverse and creative music tag-based descriptions and LRC format lyrics with song structure tags. CRITICAL REQUIREMENTS: 1) Description must be structured tags separated by commas, NOT narrative text. 2) Return ONLY pure, valid JSON format without any extra symbols, markers, or comments. 3) Each song must include structure tags like [Verse 1], [Chorus], [Bridge], etc., followed by LRC format lyrics [mm:ss.xx]lyric_content. 4) MANDATORY: Each song must have MORE than {require_length} lines of lyrics with timestamps. "},
+                    {"role": "user", "content": prompt}
+                ],
+                n=1,  
+                temperature=1.0,  
+            )
+            #print(prompt)
+            # Extract all responses
+            results = []
+            filtered_count = 0
+            last_content = None
+            
+            for i, choice in enumerate(completion.choices, 1):
+                try:
+                    content = choice.message.content.strip()
+                    last_content = content
+                    print(f"\n=== GPT Response {i} ===")
+                    print(content)
+                    print("=" * 50)
+                    # Try to extract JSON content
+                    if "```json" in content:
+                        content = content.split("```json")[1].split("```")[0].strip()
+                    elif "```" in content:
+                        content = content.split("```")[1].split("```")[0].strip()
+                    
+                    # Clean trailing commas in JSON (extra commas)
+                    # Remove commas after last element of object/array
+                    content = re.sub(r',(\s*[}\]])', r'\1', content)
+                    
+                    # Parse JSON array
+                    result_array = json.loads(content)
+                    
+                    # Ensure it's a list
+                    if isinstance(result_array, list):
+                        # Validate each object in array
+                        for song in result_array:
+                            if isinstance(song, dict) and 'description' in song and 'lyrics' in song:
+                                if _validate_timestamps(song.get('lyrics', '')):
+                                    results.append(song)
+                                else:
+                                    filtered_count += 1
+                    # If returned a single object (compatibility with old format)
+                    elif isinstance(result_array, dict) and 'description' in result_array and 'lyrics' in result_array:
+                        if _validate_timestamps(result_array.get('lyrics', '')):
+                            results.append(result_array)
+                        else:
+                            filtered_count += 1
+                        
+                except json.JSONDecodeError:
+                    continue
+            
+            if filtered_count:
+                print(f"Total {filtered_count} songs filtered due to timestamp validation failure")
+            
+            # Print parsing results
+            print(f"\nParsing complete, results length: {len(results)}")
+            print(f"Results content: {results}")
+            print(start_duration, end_duration,example1_timestamp,example2_timestamp,require_length)
+            
+            # If parsed result length is not 5, write model response content to test.txt
+            if len(results) != 2:
+                print(f"Warning: Parsed result length is not 2, actual is {len(results)}, will write to test.txt")
+                with open('test.txt', 'w', encoding='utf-8') as f:
+                    if last_content is not None:
+                        f.write(last_content)
+                print("Written to test.txt file")
+            
+            # Check if successfully generated 50 songs (10 responses * 5 each)
+            if len(results) >= 50:
+                # Append save results to file (use lock to ensure thread safety)
+                with file_lock:
+                    with open(output_file, 'a', encoding='utf-8') as f:
+                        for result in results[:50]:  # Only save first 50 songs
+                            f.write(json.dumps(result, ensure_ascii=False) + '\n')
+                
+                return selected_indices, min(len(results), 50)
+            elif attempt < max_retries:
+                print(f"Only successfully parsed {len(results)}/50 songs, retrying...")
+                time.sleep(2)
+            else:
+                # Last attempt, save even if not 50 songs
+                if len(results) > 0:
+                    with file_lock:
+                        with open(output_file, 'a', encoding='utf-8') as f:
+                            for result in results:
+                                f.write(json.dumps(result, ensure_ascii=False) + '\n')
+                return selected_indices, len(results)
+                
+        except Exception as e:
+            if attempt < max_retries:
+                print(f"Error occurred during generation: {e}, retrying...")
+                time.sleep(2)
+            else:
+                print(f"Generation failed: {e}")
+                return selected_indices, 0
+    
+    return selected_indices, 0
+
+
+class IndexPool:
+    """Thread-safe index pool with automatic reset support"""
+    
+    def __init__(self, total_size, selected_file):
+        self.total_size = total_size
+        self.selected_file = selected_file
+        self.lock = threading.Lock()
+        self.available_indices = []
+        self.selected_indices = set()
+        self.reset_count = 0  # Record reset count
+        
+        # Load selected indices from file
+        self._load_selected_indices()
+        # Initialize available indices
+        self._reset_pool()
+    
+    def _load_selected_indices(self):
+        """Load selected indices from file"""
+        if os.path.exists(self.selected_file):
+            with open(self.selected_file, 'r', encoding='utf-8') as f:
+                for line in f:
+                    self.selected_indices.add(int(line.strip()))
+    
+    def _reset_pool(self):
+        """Reset index pool"""
+        # Calculate available indices
+        self.available_indices = [i for i in range(self.total_size) if i not in self.selected_indices]
+        random.shuffle(self.available_indices)  # Shuffle order
+        
+        if len(self.available_indices) == 0:
+            # If no available indices, all have been used, reset selected_indices
+            self.reset_count += 1
+            print(f"\nIndex pool exhausted, resetting pool for the {self.reset_count}th time, re-selecting from {self.total_size} songs")
+            self.selected_indices.clear()
+            self.available_indices = list(range(self.total_size))
+            random.shuffle(self.available_indices)
+    
+    def get_indices(self, count):
+        """
+        Thread-safe get specified number of indices
+        
+        Args:
+            count: Number of indices needed
+            
+        Returns:
+            List of selected indices
+        """
+        with self.lock:
+            # Check if pool needs to be reset
+            if len(self.available_indices) < count:
+                self._reset_pool()
+            
+            # Get indices
+            selected = self.available_indices[:count]
+            self.available_indices = self.available_indices[count:]
+            
+            # Add to selected set
+            for idx in selected:
+                self.selected_indices.add(idx)
+            
+            # Write to file
+            with open(self.selected_file, 'a', encoding='utf-8') as f:
+                for idx in selected:
+                    f.write(f"{idx}\n")
+            
+            return selected
+    
+    def get_stats(self):
+        """Get statistics"""
+        with self.lock:
+            return {
+                'available': len(self.available_indices),
+                'selected': len(self.selected_indices),
+                'reset_count': self.reset_count
+            }
+
+
+def batch_generate_music(input_file, output_file, selected_file, total_songs=1000, sample_size=20, model='gpt-4o-mini', num_threads=10):
+    """
+    Batch generate music descriptions and lyrics (multi-threaded version)
+    
+    Args:
+        input_file: Path to input jsonl file
+        output_file: Path to output jsonl file
+        selected_file: Path to file recording selected indices
+        total_songs: Total number of songs to generate
+        sample_size: Number of samples to extract each time
+        model: Model name to use
+        num_threads: Number of threads
+    """
+    # Load all music data
+    print("Loading music data...")
+    all_music_data = []
+    with open(input_file, 'r', encoding='utf-8') as f:
+        for line in f:
+            data = json.loads(line.strip())
+            all_music_data.append(data)
+    print(f"Loaded {len(all_music_data)} songs")
+    
+    # Create thread-safe index pool
+    index_pool = IndexPool(len(all_music_data), selected_file)
+    stats = index_pool.get_stats()
+    print(f"Currently selected indices: {stats['selected']}")
+    print(f"Currently available indices: {stats['available']}")
+    
+    # Calculate number of calls needed (5 songs per call)
+    num_iterations = (total_songs + 1) // 2  # Round up
+    print(f"Need to call {num_iterations} times to generate approximately {total_songs} songs (5 per call)")
+    print(f"Using {num_threads} threads for parallel processing\n")
+    
+    # Create file write lock
+    file_lock = threading.Lock()
+    
+    # Statistics
+    total_generated = 0
+    generated_lock = threading.Lock()
+    
+    def worker_task(task_id):
+        """Worker thread task"""
+        try:
+            used_indices, success_count = generate_music_descriptions(
+                all_music_data=all_music_data,
+                index_pool=index_pool,
+                output_file=output_file,
+                file_lock=file_lock,
+                sample_size=sample_size,
+                model=model,
+                max_retries=0  # Retry
+            )
+            return success_count
+        except Exception as e:
+            print(f"Task {task_id} failed: {e}")
+            return 0
+    
+    # Use thread pool and progress bar
+    with ThreadPoolExecutor(max_workers=num_threads) as executor:
+        # Submit all tasks
+        futures = {executor.submit(worker_task, i): i for i in range(num_iterations)}
+        
+        # Use tqdm to show progress
+        with tqdm(total=num_iterations, desc="Generation progress", unit="batch") as pbar:
+            for future in as_completed(futures):
+                success_count = future.result()
+                
+                with generated_lock:
+                    total_generated += success_count
+                
+                # Get current statistics
+                stats = index_pool.get_stats()
+                
+                # Update progress bar
+                pbar.set_postfix({
+                    'Batch': f'{success_count}/5',
+                    'Total': total_generated,
+                    'Remaining': stats['available'],
+                    'Resets': stats['reset_count']
+                })
+                pbar.update(1)
+    
+    # Final statistics
+    stats = index_pool.get_stats()
+    print(f"\nGeneration complete!")
+    print(f"Total generated: {total_generated} songs")
+    print(f"Used {stats['selected']} indices")
+    print(f"Remaining available indices: {stats['available']}")
+    print(f"Pool reset count: {stats['reset_count']}")
+
+
+if __name__ == '__main__':
+    input_file = 'tagged_musics.jsonl'
+    output_file = 'generate_lrc_5mini.jsonl'
+    selected_file = 'selected.txt'
+    # n=1, max_retries=0, sample 10 songs each time, generate 5 new songs
+    batch_generate_music(
+        input_file=input_file,
+        output_file=output_file,
+        selected_file=selected_file,
+        total_songs=10,
+        sample_size=4,
+        model='gpt-5-mini',
+        num_threads=5  # Test with 1 thread first
+    )
+    # Append to txt file

data_pipeline/lyrics_gene/gen_lyrics_en.py

@@ -0,0 +1,577 @@
+import os
+import json
+import time
+import random
+import re
+from openai import OpenAI
+from tqdm import tqdm
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+# Set environment variables
+# Note: Set these environment variables before running the script
+# export OPENAI_API_KEY="your-api-key"
+# export OPENAI_BASE_URL="https://api.openai.com/v1"  # or your custom API URL
+if not os.environ.get("OPENAI_API_KEY"):
+    os.environ["OPENAI_API_KEY"] = ""  # Replace with your API key or set via environment variable
+if not os.environ.get("OPENAI_BASE_URL"):
+    os.environ["OPENAI_BASE_URL"] = "https://api.openai.com/v1"  # Replace with API URL or set via environment variable
+
+# Initialize client
+client = OpenAI()
+
+
+def _extract_lyrics_timestamps(lyrics_text):
+    """
+    Extract timestamps from lyrics and convert to seconds
+    Args:
+        lyrics_text: Lyrics string
+    Returns:
+        List[float]: Timestamps in order (seconds)
+    """
+    if not isinstance(lyrics_text, str):
+        return []
+    pattern = re.compile(r'\[(\d{2}):(\d{2})(?:\.(\d{2}))?\]')
+    timestamps = []
+    for match in pattern.finditer(lyrics_text):
+        minutes = int(match.group(1))
+        seconds = int(match.group(2))
+        fraction = match.group(3)
+        total_seconds = minutes * 60 + seconds
+        if fraction is not None:
+            divisor = 100 if len(fraction) == 2 else 10 ** len(fraction)
+            total_seconds += int(fraction) / divisor
+        timestamps.append(total_seconds)
+    return timestamps
+
+
+def _validate_timestamps(lyrics_text, min_last_timestamp=170, max_interval=35):
+    """
+    Validate if timestamps in lyrics meet requirements
+    Args:
+        lyrics_text: Lyrics string
+        min_last_timestamp: Minimum value of last timestamp (seconds)
+        max_interval: Maximum interval between last two timestamps (seconds)
+    Returns:
+        bool: Whether validation passed
+    """
+    timestamps = _extract_lyrics_timestamps(lyrics_text)
+    if len(timestamps) < 2:
+        print("Validation failed: Timestamp count less than 2")
+        return False
+    last = timestamps[-1]
+    second_last = timestamps[-2]
+    if last < min_last_timestamp:
+        print(f"Validation failed: Last timestamp {last:.2f}s is less than {min_last_timestamp}s")
+        return False
+    if last - second_last > max_interval:
+        print(f"Validation failed: Interval between last two timestamps {last - second_last:.2f}s is greater than {max_interval}s")
+        return False
+    return True
+
+
+def chat_gpt(text, model='gpt-4o-mini'):
+    while True:
+        try:
+            # Call OpenAI chat completions API
+            completion = client.chat.completions.create(
+                model=model,  # Use GPT-4o-mini model
+                messages=[
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": text}
+                ]
+            )
+            # Get response content
+            if getattr(completion.choices[0].message, 'content', None):
+                content = completion.choices[0].message.content.strip()
+                return content
+            else:
+                print('error_wait_2s')
+        except Exception as e:
+            print(f"Error: {e}")
+        time.sleep(2)
+
+
+def chat_gpt_call(text, model='gpt-4o-mini'):
+    # Call OpenAI chat completions API
+    completion = client.chat.completions.create(
+        model=model,  # Use GPT-4o-mini model
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": text}
+        ]
+    )
+    # Get response content
+    if getattr(completion.choices[0].message, 'content', None):
+        content = completion.choices[0].message.content.strip()
+        return content
+    else:
+        print('error_wait_2s')
+
+
+def generate_music_descriptions(all_music_data, index_pool, output_file, file_lock, sample_size=20, model='gpt-4o-mini', max_retries=0):
+    """
+    Read music data file, randomly sample, call GPT to generate new music descriptions and lyrics
+    
+    Args:
+        all_music_data: List of all music data
+        index_pool: Index pool object (thread-safe)
+        output_file: Path to output jsonl file
+        file_lock: File write lock
+        sample_size: Number of samples to randomly extract
+        model: Model name to use
+        max_retries: Maximum retry count
+    
+    Returns:
+        (used_indices, success_count): List of used indices and number of successful generations
+    """
+    # duration_ranges = [
+    #     ("3.00", "3.15", 50), ("3.15", "3.30", 50), ("3.30", "3.45", 60),
+    #  ("3.45", "4.00", 60),
+    #     ("4.00", "4.15", 70), ("4.15", "4.30", 70), ("4.30", "4.45", 70)
+    # ]
+    duration_ranges = [
+        ("3.00", "3.15", 60), ("3.15", "3.30", 70), ("3.30", "3.45", 80),("3.45", "4.00", 90),
+        ("4.00", "4.15", 100), ("4.15", "4.30", 100), ("4.30", "4.45", 100)
+    ]
+    selected_range = random.choice(duration_ranges)
+    require_length = selected_range[2]
+    
+    # Directly convert to timestamp format (strictly corresponding to left and right ends of tuple)
+    start_timestamp = f"[{selected_range[0].replace('.', ':')}.00]"
+    end_timestamp = f"[{selected_range[1].replace('.', ':')}.00]"
+    
+    # Generate duration description
+    # Convert seconds to minutes and seconds format
+    start_seconds = float(selected_range[0])
+    start_minutes = int(start_seconds // 60)
+    start_secs = int(start_seconds % 60)
+    start_duration = f"{start_minutes}min {start_secs}sec"
+    
+    end_seconds = float(selected_range[1])
+    end_minutes = int(end_seconds // 60)
+    end_secs = int(end_seconds % 60)
+    end_duration = f"{end_minutes}min {end_secs}sec"
+    
+    # Generate random timestamps in examples (randomly generated within time range)
+    # Parse time string to minutes and seconds
+    start_parts = selected_range[0].split('.')
+    end_parts = selected_range[1].split('.')
+    
+    start_minutes = int(start_parts[0])
+    start_seconds = int(start_parts[1])
+    end_minutes = int(end_parts[0])
+    end_seconds = int(end_parts[1])
+    
+    # Convert to total seconds
+    start_total_seconds = start_minutes * 60 + start_seconds
+    end_total_seconds = end_minutes * 60 + end_seconds
+    
+    # Randomly generate within range
+    example1_seconds = random.randint(start_total_seconds, end_total_seconds)
+    example2_seconds = random.randint(start_total_seconds, end_total_seconds)
+    
+    example1_minutes = example1_seconds // 60
+    example1_secs = example1_seconds % 60
+    example2_minutes = example2_seconds // 60
+    example2_secs = example2_seconds % 60
+    
+    example1_timestamp = f"[{example1_minutes:02d}:{example1_secs:02d}.00]"
+    example2_timestamp = f"[{example2_minutes:02d}:{example2_secs:02d}.00]"
+    
+    # Get sample indices from index pool (thread-safe)
+    selected_indices = index_pool.get_indices(sample_size)
+    if not selected_indices:
+        return [], 0
+    
+    sample_data = [all_music_data[i] for i in selected_indices]
+    
+    # Extract all unique styles
+    styles = []
+    for data in sample_data:
+        style = data.get('style', '')
+        if style and style not in styles:
+            styles.append(style)
+    
+    styles_text = "、".join(styles)
+    
+    # Build example text - include all sampled data (excluding style)
+    examples = []
+    for i, data in enumerate(sample_data, 1):  
+        lyrics_text = " ".join(data.get('lyrics', [])) if isinstance(data.get('lyrics'), list) else data.get('lyrics', '')
+        description = data.get('description', '')
+        examples.append(f"Example {i}:\ndescription: {description}\nlyrics: {lyrics_text}")
+    
+    examples_text = "\n\n".join(examples)
+
+    prompt = f"""Generate 2 complete songs. Each song must meet the following hard requirements:
+- Strictly forbidden to generate lyrics with fewer than {require_length} lines!
+- The number of lyric lines for each song must be strictly greater than {require_length}. This is a hard requirement!
+- The timestamp of the final line must be between {start_timestamp} and {end_timestamp}.
+- The two songs must differ in duration and line count; their final timestamps must not be identical.
+- The timestamp interval between adjacent lyric lines must not exceed 10 seconds! Timestamps must be continuous and progress naturally.
+- Awkward gaps like "[03:25.00]in the heart[04:25.00]the last lyric" are strictly forbidden. Do not exceed a 10-second interval.
+- It is strictly forbidden to repeat the entire structure or its sections after one iteration is complete. It is also strictly forbidden to repeat the same lyric line multiple times.
+If any of the above requirements are not met, the generation is considered a failure. Please regenerate.
+Please generate 2 new, diverse music descriptions and LRC format lyrics. The language should be English.
+
+
+Creative Requirements:
+1. Style and Genre must be diverse.
+2. Description Tagging Requirements (Must be strictly followed):
+   The description field must use a structured tag format, including the following tags, separated by commas:
+   - Music Style tag
+   - Music Genre tag
+   - Instruments tag
+   - Emotional Tone tag
+   - Mood/Atmosphere tag
+   - Vocal Style and Voice tag, limited to either "male voice" or "female voice", solo performance only.
+   Note: Each tag should be concise. Multiple tags of the same category can be separated by a slash (e.g., "Piano/Violin").
+3. Lyric Creativity: The lyrics should have depth and artistry:
+   - Themes can cover various aspects such as love, life, society, nature, philosophy, dreams, memories, etc.
+   - Use rich literary devices: metaphors, imagery, contrast, parallelism, etc.
+   - Express sincere emotions with a focus on rhyme and rhythm.
+   - The style can be narrative, lyrical, or stream-of-consciousness.
+4. Lyric Structure and Length Requirements (Must be strictly followed):
+   - The lyrics must be organized using the following structure, with section tags annotating each part.
+   - The structure must strictly follow this order, for a total of 8 section tags: [Verse 1] → [Pre-Chorus] → [Chorus] → [Verse 2] → [Pre-Chorus] → [Chorus] → [Bridge] → [Chorus (Outro)].
+   - A single song can only have these 8 section tags. [Verse 1] and [Verse 2] appear once; [Pre-Chorus] and [Chorus] appear twice; [Bridge] and [Chorus (Outro)] appear once. Do not add or repeat extra section tags.
+   - Each section tag (e.g., [Verse 1], [Chorus]) must be on its own line, immediately followed by the LRC format lyrics for that section.
+   - Separate sections with a blank line.
+   - **Total Line Count Requirement**: The entire song must contain at least {require_length} lines of timestamped lyrics (not including section tags or blank lines).
+5. LRC Format Mandatory Rules (Must be strictly followed):
+   - Each line of lyrics must be in the format `[mm:ss.xx]Lyric content`, with no space between the timestamp and the lyrics. The lyric content should be coherent.
+   - **Each line must contain only one short phrase of lyrics.** Start a new line when encountering punctuation like commas or periods.
+   - **Strictly forbidden to merge multiple sentences or clauses onto the same line.**
+   - Timestamps must be distributed naturally. **The first line's timestamp must not be [00:00.00]**. Allow for an instrumental intro (suggestion: start between [00:05.00] and [00:15.00]).
+   - Timestamp intervals must be varied: The intervals within each song must be diverse, often using decimal values. Do not use a fixed interval:
+     * A single song must contain a variety of different intervals; do not use the same interval for all lines (e.g., not all 4-second gaps).
+     * Dynamically adjust intervals based on the emotional intensity and rhythm of the lyrics.
+     * The gap between adjacent lines should vary to reflect the musical rhythm.
+   - Timestamp allocation should be reasonably inferred based on the song's style, emotion, and rhythm, not mechanically assigned based on lyric length.
+   - The length of each lyric line should vary naturally; do not make them all uniform.
+   - **The total song duration must be between {start_duration} and {end_duration} (meaning the final line's timestamp must be between {start_timestamp} and {end_timestamp}). This is a hard requirement!**
+6. Lyric Length Requirement: The number of lyric lines in the lyrics field must be greater than {require_length}. If the generated length is too short, please regenerate.
+7. Uniqueness and Originality: Each piece should be unique. Avoid simply repeating the content from examples.
+8. Format Requirements:
+   - Directly return a JSON array containing 2 song objects. Each object must have only "description" and "lyrics" fields.
+   - `description` field: Must be in tag format, not narrative text.
+   - `lyrics` field: A string in LRC format with section tags.
+   - Strictly forbidden to insert any extra symbols, markers, comments, or explanatory text within the JSON.
+
+LRC Format Example (with section tags):
+[Verse 1]
+[00:08.00]First line of lyrics
+[00:12.50]Second line of lyrics
+[00:17.20]Third line of lyrics
+
+[Pre-Chorus]
+[00:22.00]Pre-chorus lyrics
+[00:26.50]Pre-chorus lyrics
+
+[Chorus]
+[00:31.00]Chorus lyrics
+[00:35.50]Chorus lyrics
+
+Negative Examples (to avoid):
+- Incorrect: [01:30.00](Piano Interlude) - Do not add parenthetical comments after the timestamp.
+- Incorrect: [00:00.00]Starting lyric - The first line cannot start at 00:00.00.
+- Incorrect: [00:05.00]In the familiar field, the sun casts golden rays upon the wheat - Strictly forbidden to place multiple clauses on the same line.
+- Incorrect: [03:00.00] In the light of hope[03:05.50] In the light of hope[03:10.20] In the light of hope -Excessive repetition of the exact same lyric line is strictly forbidden. Lyrical content must show variation.
+Now, please fully unleash your creativity and generate 2 new, complete works of music descriptions and LRC format lyrics.
+Special Reminder: Each song must be complete, not abbreviated or omitted! It must contain the full 8 sections (Verse 1, Pre-Chorus, Chorus, Verse 2, Pre-Chorus, Chorus, Bridge, Chorus Outro) and strictly ensure more than {require_length} lines of lyrics.
+
+Directly return in JSON array format:
+[
+  {{"description": "...", "lyrics": "..."}},
+  {{"description": "...", "lyrics": "..."}}
+]"""
+    # Try to generate with retry mechanism
+    for attempt in range(max_retries + 1):
+        try:
+            # Call OpenAI API
+            completion = client.chat.completions.create(
+                model=model,
+                messages=[
+                    {"role": "system", "content": f"You are a creative music lyricist and composer. Please generate diverse and creative music tag-based descriptions and LRC format lyrics with song structure tags. CRITICAL REQUIREMENTS: 1) Description must be structured tags separated by commas, NOT narrative text. 2) Return ONLY pure, valid JSON format without any extra symbols, markers, or comments. 3) Each song must include structure tags like [Verse 1], [Chorus], [Bridge], etc., followed by LRC format lyrics [mm:ss.xx]lyric_content. 4) MANDATORY: Each song must have MORE than {require_length} lines of lyrics with timestamps. "},
+                    {"role": "user", "content": prompt}
+                ],
+                n=1,  
+                temperature=1.0,  
+            )
+            #print(prompt)
+            # Extract all responses
+            results = []
+            filtered_count = 0
+            last_content = None
+            
+            for i, choice in enumerate(completion.choices, 1):
+                try:
+                    content = choice.message.content.strip()
+                    last_content = content
+                    print(f"\n=== GPT Response {i} ===")
+                    print(content)
+                    print("=" * 50)
+                    # Try to extract JSON content
+                    if "```json" in content:
+                        content = content.split("```json")[1].split("```")[0].strip()
+                    elif "```" in content:
+                        content = content.split("```")[1].split("```")[0].strip()
+                    
+                    # Clean trailing commas in JSON (extra commas)
+                    # Remove commas after last element of object/array
+                    content = re.sub(r',(\s*[}\]])', r'\1', content)
+                    
+                    # Parse JSON array
+                    result_array = json.loads(content)
+                    
+                    # Ensure it's a list
+                    if isinstance(result_array, list):
+                        # Validate each object in array
+                        for song in result_array:
+                            if isinstance(song, dict) and 'description' in song and 'lyrics' in song:
+                                if _validate_timestamps(song.get('lyrics', '')):
+                                    results.append(song)
+                                else:
+                                    filtered_count += 1
+                    # If returned a single object (compatibility with old format)
+                    elif isinstance(result_array, dict) and 'description' in result_array and 'lyrics' in result_array:
+                        if _validate_timestamps(result_array.get('lyrics', '')):
+                            results.append(result_array)
+                        else:
+                            filtered_count += 1
+                        
+                except json.JSONDecodeError:
+                    continue
+            
+            if filtered_count:
+                print(f"Total {filtered_count} songs filtered due to timestamp validation failure")
+            
+            # Print parsing results
+            print(f"\nParsing complete, results length: {len(results)}")
+            print(f"Results content: {results}")
+            print(start_duration, end_duration,example1_timestamp,example2_timestamp,require_length)
+            
+            # If parsed result length is not 2, write model response content to test.txt
+            if len(results) != 2:
+                print(f"Warning: Parsed result length is not 2, actual is {len(results)}, will write to test.txt")
+                with open('test.txt', 'w', encoding='utf-8') as f:
+                    if last_content is not None:
+                        f.write(last_content)
+                print("Written to test.txt file")
+            
+            # Check if successfully generated 50 songs (10 responses * 5 each)
+            if len(results) >= 50:
+                # Append save results to file (use lock to ensure thread safety)
+                with file_lock:
+                    with open(output_file, 'a', encoding='utf-8') as f:
+                        for result in results[:50]:  # Only save first 50 songs
+                            f.write(json.dumps(result, ensure_ascii=False) + '\n')
+                
+                return selected_indices, min(len(results), 50)
+            elif attempt < max_retries:
+                print(f"Only successfully parsed {len(results)}/50 songs, retrying...")
+                time.sleep(2)
+            else:
+                # Last attempt, save even if not 50 songs
+                if len(results) > 0:
+                    with file_lock:
+                        with open(output_file, 'a', encoding='utf-8') as f:
+                            for result in results:
+                                f.write(json.dumps(result, ensure_ascii=False) + '\n')
+                return selected_indices, len(results)
+                
+        except Exception as e:
+            if attempt < max_retries:
+                print(f"Error occurred during generation: {e}, retrying...")
+                time.sleep(2)
+            else:
+                print(f"Generation failed: {e}")
+                return selected_indices, 0
+    
+    return selected_indices, 0
+
+
+class IndexPool:
+    """Thread-safe index pool with automatic reset support"""
+    
+    def __init__(self, total_size, selected_file):
+        self.total_size = total_size
+        self.selected_file = selected_file
+        self.lock = threading.Lock()
+        self.available_indices = []
+        self.selected_indices = set()
+        self.reset_count = 0  # Record reset count
+        
+        # Load selected indices from file
+        self._load_selected_indices()
+        # Initialize available indices
+        self._reset_pool()
+    
+    def _load_selected_indices(self):
+        """Load selected indices from file"""
+        if os.path.exists(self.selected_file):
+            with open(self.selected_file, 'r', encoding='utf-8') as f:
+                for line in f:
+                    self.selected_indices.add(int(line.strip()))
+    
+    def _reset_pool(self):
+        """Reset index pool"""
+        # Calculate available indices
+        self.available_indices = [i for i in range(self.total_size) if i not in self.selected_indices]
+        random.shuffle(self.available_indices)  # Shuffle order
+        
+        if len(self.available_indices) == 0:
+            # If no available indices, all have been used, reset selected_indices
+            self.reset_count += 1
+            print(f"\nIndex pool exhausted, resetting pool for the {self.reset_count}th time, re-selecting from {self.total_size} songs")
+            self.selected_indices.clear()
+            self.available_indices = list(range(self.total_size))
+            random.shuffle(self.available_indices)
+    
+    def get_indices(self, count):
+        """
+        Thread-safe get specified number of indices
+        
+        Args:
+            count: Number of indices needed
+            
+        Returns:
+            List of selected indices
+        """
+        with self.lock:
+            # Check if pool needs to be reset
+            if len(self.available_indices) < count:
+                self._reset_pool()
+            
+            # Get indices
+            selected = self.available_indices[:count]
+            self.available_indices = self.available_indices[count:]
+            
+            # Add to selected set
+            for idx in selected:
+                self.selected_indices.add(idx)
+            
+            # Write to file
+            with open(self.selected_file, 'a', encoding='utf-8') as f:
+                for idx in selected:
+                    f.write(f"{idx}\n")
+            
+            return selected
+    
+    def get_stats(self):
+        """Get statistics"""
+        with self.lock:
+            return {
+                'available': len(self.available_indices),
+                'selected': len(self.selected_indices),
+                'reset_count': self.reset_count
+            }
+
+
+def batch_generate_music(input_file, output_file, selected_file, total_songs=1000, sample_size=20, model='gpt-4o-mini', num_threads=10):
+    """
+    Batch generate music descriptions and lyrics (multi-threaded version)
+    
+    Args:
+        input_file: Path to input jsonl file
+        output_file: Path to output jsonl file
+        selected_file: Path to file recording selected indices
+        total_songs: Total number of songs to generate
+        sample_size: Number of samples to extract each time
+        model: Model name to use
+        num_threads: Number of threads
+    """
+    # Load all music data
+    print("Loading music data...")
+    all_music_data = []
+    with open(input_file, 'r', encoding='utf-8') as f:
+        for line in f:
+            data = json.loads(line.strip())
+            all_music_data.append(data)
+    print(f"Loaded {len(all_music_data)} songs")
+    
+    # Create thread-safe index pool
+    index_pool = IndexPool(len(all_music_data), selected_file)
+    stats = index_pool.get_stats()
+    print(f"Currently selected indices: {stats['selected']}")
+    print(f"Currently available indices: {stats['available']}")
+    
+    # Calculate number of calls needed (5 songs per call)
+    num_iterations = (total_songs + 1) // 2  # Round up
+    print(f"Need to call {num_iterations} times to generate approximately {total_songs} songs (5 per call)")
+    print(f"Using {num_threads} threads for parallel processing\n")
+    
+    # Create file write lock
+    file_lock = threading.Lock()
+    
+    # Statistics
+    total_generated = 0
+    generated_lock = threading.Lock()
+    
+    def worker_task(task_id):
+        """Worker thread task"""
+        try:
+            used_indices, success_count = generate_music_descriptions(
+                all_music_data=all_music_data,
+                index_pool=index_pool,
+                output_file=output_file,
+                file_lock=file_lock,
+                sample_size=sample_size,
+                model=model,
+                max_retries=0  # Retry
+            )
+            return success_count
+        except Exception as e:
+            print(f"Task {task_id} failed: {e}")
+            return 0
+    
+    # Use thread pool and progress bar
+    with ThreadPoolExecutor(max_workers=num_threads) as executor:
+        # Submit all tasks
+        futures = {executor.submit(worker_task, i): i for i in range(num_iterations)}
+        
+        # Use tqdm to show progress
+        with tqdm(total=num_iterations, desc="Generation progress", unit="batch") as pbar:
+            for future in as_completed(futures):
+                success_count = future.result()
+                
+                with generated_lock:
+                    total_generated += success_count
+                
+                # Get current statistics
+                stats = index_pool.get_stats()
+                
+                # Update progress bar
+                pbar.set_postfix({
+                    'Batch': f'{success_count}/5',
+                    'Total': total_generated,
+                    'Remaining': stats['available'],
+                    'Resets': stats['reset_count']
+                })
+                pbar.update(1)
+    
+    # Final statistics
+    stats = index_pool.get_stats()
+    print(f"\nGeneration complete!")
+    print(f"Total generated: {total_generated} songs")
+    print(f"Used {stats['selected']} indices")
+    print(f"Remaining available indices: {stats['available']}")
+    print(f"Pool reset count: {stats['reset_count']}")
+
+
+if __name__ == '__main__':
+    input_file = 'tagged_musics.jsonl'
+    output_file = 'generate_en_lrc.jsonl'
+    selected_file = 'selected.txt'
+    # n=1, max_retries=0, sample 10 songs each time, generate 5 new songs
+    batch_generate_music(
+        input_file=input_file,
+        output_file=output_file,
+        selected_file=selected_file,
+        total_songs=100,
+        sample_size=4,
+        model='gpt-4o-mini',
+        num_threads=20  # Test with 1 thread first
+    )
+    # Append to txt file

data_pipeline/meta_process/convert_convs.py

@@ -0,0 +1,98 @@
+"""
+Generate multi-turn dialogue data, each turn contains lyric text and corresponding audio token slices
+"""
+
+import json
+import re
+import os
+import torch
+from tqdm import tqdm
+from my_tool import load_jsonl
+
+TOKEN_PER_SECOND = 25  # Number of tokens per second of audio
+NUM_ITEMS = 100000  # Process N items first
+
+timestamp_pattern = re.compile(r"\[([0-9]{1,2}):([0-9]{1,2})(?:[.:]([0-9]{1,3}))?\]")
+
+def _parse_lyric_with_timestamps(lyric: str):
+    """
+    Return [(start_time_s, text), ...] sorted by timestamp
+    """
+    result = []
+    for match in timestamp_pattern.finditer(lyric):
+        start_idx = match.end()
+        end_idx = lyric.find("[", start_idx)
+        text = lyric[start_idx:end_idx].strip() if end_idx != -1 else lyric[start_idx:].strip()
+        if not text:
+            continue
+        minute = int(match.group(1))
+        second = int(match.group(2))
+        ms = int(match.group(3)) if match.group(3) else 0
+        total_seconds = minute * 60 + second + ms / 1000
+        result.append((total_seconds, text))
+    return result
+
+def _load_audio_tokens(pt_file):
+    """
+    Load MuCodec encoding of audio
+    """
+    audio_ids = torch.load(pt_file, map_location="cpu").squeeze().long()
+    return audio_ids
+
+def _get_token_slice(audio_tokens, start_s, end_s):
+    """Split encoding by time segment"""
+    start_idx = int(start_s * TOKEN_PER_SECOND)
+    end_idx = int(end_s * TOKEN_PER_SECOND)
+    sliced = audio_tokens[start_idx:end_idx]
+    return "[SOA]" + "".join([f"<AUDIO_{i.item()}>" for i in sliced]) + "[EOA]"
+
+def _process_item(item, pt_dir:str):
+    song_name = item.get("song") or item.get("name")
+    song_name = song_name.split('.mp3')[0]              # For mucodec, remove extension
+    pt_file = os.path.join(pt_dir, f"{song_name}.pt")
+    if not os.path.exists(pt_file):
+        return None
+
+    audio_tokens = _load_audio_tokens(pt_file)
+    tlyric_ = item.get('tlyric', "")
+    lyric_ = item.get('lyric', "")
+    lyric = tlyric_ if len(tlyric_) > len(lyric_) else lyric_
+    lyrics_ts = _parse_lyric_with_timestamps(lyric)
+
+    if not lyrics_ts:
+        # Skip if no lyrics
+        return None
+
+    rounds = []
+
+    # First generate a system message containing song information
+    intro_text = (
+        f"请生成一首歌曲，歌名为《{item.get('name', '')}》，风格是{item.get('style','')}"
+        f"，情绪为{item.get('emotion','')}，节奏：{item.get('rhythm','')}，"
+        f"{item.get('description','')}，由{item.get('singer','')}演唱，语言：{item.get('lang','')}。"
+        f"歌词如下：" + " ".join([text for _, text in lyrics_ts]) + "接下来我会逐句告诉你需要生成歌曲片段的歌词，\n请先生成前奏"
+    )
+    rounds.append({"role": "user", "content": intro_text})
+    rounds.append({"role": "assistant", "content": _get_token_slice(audio_tokens, 0, lyrics_ts[0][0])})  # Intro tokens
+
+    # Each lyric line corresponds to one round
+    for idx, (start_s, text) in enumerate(lyrics_ts[:-1]):  ## Last line handled separately
+        end_s = lyrics_ts[idx + 1][0] if idx + 1 < len(lyrics_ts) else len(audio_tokens)/TOKEN_PER_SECOND  # Last line to end of audio
+        rounds.append({"role": "user", "content": text})
+        rounds.append({"role": "assistant", "content": _get_token_slice(audio_tokens, start_s, end_s)})
+
+    # Tail processing logic
+    rounds.append({"role": "user", "content": f"请生成歌词{lyrics_ts[-1][1]}以及歌曲结尾"})
+    rounds.append({"role": "assistant", "content": _get_token_slice(audio_tokens, lyrics_ts[-1][0], len(audio_tokens)/TOKEN_PER_SECOND)})
+
+    return rounds
+
+# ===== External Interface =====
+
+def get_convert_convs(dataset:list[dict], pt_dir:str, save_path:str):
+    with open(save_path, "w", encoding="utf-8") as fout:
+        for item in tqdm(dataset, desc="Converting convs"):
+            rounds = _process_item(item, pt_dir)
+            if not rounds:
+                continue
+            fout.write(json.dumps({"messages": rounds}, ensure_ascii=False) + "\n")

data_pipeline/meta_process/convert_lyrics.py

@@ -0,0 +1,180 @@
+import os
+import re
+import json
+import copy
+from tqdm import tqdm
+from my_tool import dict_sort_print
+from collections import defaultdict
+from convert_convs import _parse_lyric_with_timestamps
+
+# ===== Lyric Parsing =====
+
+def _parse_lyrics(text:str) -> dict:
+    """Parse metadata, lyrics and timestamps from lyric information"""
+    segs = text.split("\n")
+    metadata = {
+        "lyrics_meta": {},
+        "lyrics": [],
+        "lyrics_time": [],
+    }
+    for seg in segs:
+        # Format: [time] metadata / lyrics
+        results = _parse_lyric_with_timestamps(seg)
+        for time, content in results:
+            if ":" in content or "：" in content:
+                # Metadata
+                pos1 = content.find(":")
+                pos2 = content.find("：")
+                pos = pos1 if pos1 != -1 else pos2
+                key = content[:pos].strip()
+                value = content[pos+1:].strip()
+                metadata["lyrics_meta"][key] = value
+            elif time == "00:00.00":
+                # Unstructured metadata at the beginning
+                continue
+            elif len(metadata['lyrics']) == 0 and "/" in content:
+                # Unstructured metadata at the beginning
+                continue
+            else:
+                # Only keep English and space punctuation
+                if len(content) == 0:
+                    # Middle gap/end
+                    if len(metadata['lyrics']) != 0 and metadata['lyrics'][-1] != "<nop>":
+                        # If there's no previous segment (beginning), or previous segment is empty, don't record (merge)
+                        metadata['lyrics'].append("<nop>")
+                        metadata['lyrics_time'].append(time)
+                else:
+                    if len(metadata['lyrics_time']) != 0 and metadata['lyrics_time'][-1] == time and time != "<nop>":
+                        # Same timestamp means it's a translation (don't record)
+                        continue
+                    # Actual lyrics
+                    metadata['lyrics'].append(content)
+                    metadata['lyrics_time'].append(time)
+    return metadata
+
+# ===== Language Detection =====
+
+def _count_ch_nan(text:str):
+    """Count the number of Chinese and other non-English characters in a string"""
+    ch_num = 0
+    nan_num = 0
+    nan = ""
+    for c in text:
+        if '\u4e00' <= c <= '\u9fff':
+            ch_num += 1
+        elif ('a' <= c <= 'z') or ('A' <= c <= 'Z') or len(c.strip()) == 0:
+            continue
+        else:
+            nan_num += 1
+            nan += c
+    # if len(nan) > 0:
+    #     print(nan)
+    return ch_num, nan_num
+
+def _lang_decide(lyrics:list[str], val_limit:int=5, word_limit=3) -> str:
+    """
+    Determine the language type of lyrics (en/zh/ez/instrument/nan)
+    - val_limit: Only count if there are at least this many sentences
+    - word_limit: Only count if a sentence has at least this many words
+    """
+    ch_lyrics = 0
+    en_lyrics = 0
+    nan_lyrics = 0
+    for lyric in lyrics:
+        lyric = copy.deepcopy(lyric)
+        if lyric.strip() == "<nop>":
+            continue
+        lyric = re.sub(r"[''￥·′´（），。？""!@#$%^&*()?.'/,=+_—— ！…《》<>0-9～※~;－・\"、☆｜△【】＃「」‖{}\[\]-]", " ", lyric)
+        ch_num, nan_num = _count_ch_nan(lyric)
+        
+        if nan_num > word_limit:
+            nan_lyrics += 1
+            continue
+        elif ch_num > word_limit:
+            ch_lyrics += 1
+        
+        lyric = re.sub(r'[\u4e00-\u9fff]+', '', lyric)
+        # Count English words by space separation
+        en_num = len(lyric.split(" "))
+        if en_num > word_limit:
+            en_lyrics += 1
+
+    if nan_lyrics > val_limit:
+        return "nan"
+    if ch_lyrics > val_limit and en_lyrics > val_limit:
+        return "ez"
+    if ch_lyrics > val_limit:
+        return "zh"
+    if en_lyrics > val_limit:
+        return "en"
+    return "instrument"
+
+# ===== External Interface =====
+
+def get_convert_lyrics(dataset:list[dict], save_path:str, dir:str, src_subfix:str=""):
+    """Convert lyrics and annotate language type (need to locate corresponding song)"""
+    new_dataset = []
+    lang_count = defaultdict(int)
+    unmatch = []
+    with open(save_path, 'w', encoding='utf-8') as file:
+        for ele in tqdm(dataset, desc="Converting Lyrics"):
+            ele = copy.deepcopy(ele)
+            # Skip if no lyrics
+            if not ele['has_lyric']:
+                # Don't add to final result
+                continue
+            # Get lyrics
+            lyric = ele['lyric']
+            if lyric == "":
+                lyric = ele['tlyric']
+
+            # Parse lyrics
+            new_data = _parse_lyrics(lyric)
+
+            # Language detection
+            lang = _lang_decide(new_data['lyrics'])
+            lang_count[lang] += 1
+
+            # Remove redundant fields
+            del ele['artists']
+            del ele['lyric']
+            del ele['tlyric']
+            del ele['has_lyric']
+
+            # Add new fields
+            ele['lyric_lang'] = lang
+            ele['source'] += src_subfix
+            for key, value in new_data.items():
+                ele[key] = value
+            
+            new_dataset.append(ele)
+            json.dump(ele, file, ensure_ascii=False)
+            file.write("\n")
+
+    dict_sort_print(lang_count)
+    return new_dataset, unmatch
+
+def get_match_music(music_data:list[dict], lyric_data:list[dict]):
+    """Get songs that match or don't match with lyrics"""
+    # 1. Build lookup set from songs
+    name_map = {}
+    for ele in tqdm(lyric_data, desc="Existing Lyrics"):
+        name = ele['name']
+        name = re.sub(" ", "", name)
+        artist = ele['artist']
+        complete_name = f"{name} - {artist}.mp3"
+        name_map[complete_name] = ele
+    
+    # 2. Iterate through songs to find remaining ones
+    matches = []
+    unmatches = []
+    for ele in tqdm(music_data, desc="Check Matching"):
+        path = ele['path']
+        name = os.path.basename(path)
+        if name not in name_map:
+            unmatches.append(ele)
+        else:
+            meta = name_map[name]
+            meta['path'] = path
+            matches.append(meta)
+    return matches, unmatches

data_pipeline/meta_process/convert_messages.py

@@ -0,0 +1,593 @@
+"""
+Generate multi-turn dialogue data based on segment descriptions (segment-by-segment generation).
+
+Input:
+- MuseData/data/meta_suno_cn.jsonl      # Base meta (contains src_path, lyrics, lang, tag, etc.)
+- 3_block_80000_cn_desc.jsonl           # Segment descriptions, contains startS/endS/text/desc for each section
+- mucodec pt directory: same as PT_DIR_CN in multi_data_suno.py
+
+Output:
+- MuseData/sft_dataset_suno_cn.jsonl    # Multi-turn dialogues generated segment by segment
+
+Dialogue format example (refer to temp1.jsonl):
+- First user: Summary prompt (Chinese), explains "segment-by-segment" + provides [Intro dsec...] description
+- First assistant: Intro tokens (0 ~ first section's startS)
+- Subsequent segments:
+  user.content = "[{Section} dsec]{desc}\\n{text}"
+  assistant.content = corresponding time segment tokens
+"""
+
+import json
+import os
+import re
+from typing import Dict, List, Tuple, Optional
+
+import torch
+from tqdm import tqdm
+from my_tool import load_jsonl, clean_newlines
+
+# Language configuration
+LANG = "en"
+# Path configuration
+META_FILE = f"meta_suno_{LANG}.jsonl"
+# desc folder, read all *.jsonl files in it
+DESC_DIR = "desc"
+PT_DIR = f"suno_mucodec_{LANG}"
+# Output directory (different from meta directory), each desc file generates a set of three files
+OUTPUT_DIR = "outputs"
+OUTPUT_BASENAME = "minus_phonemes"
+
+TOKEN_PER_SECOND = 25
+
+
+LOG_FILE = os.path.join(OUTPUT_DIR, "section_mismatch_cn.log")   # Place in output directory
+
+def _log_warning(msg: str):
+    """Print and save to file"""
+    print(msg, end='\n')
+    with open(LOG_FILE, 'a', encoding='utf-8') as f:
+        f.write(msg + '\n')
+
+
+# Timestamp parsing regex
+timestamp_pattern = re.compile(
+    r"\[([0-9]{1,2}):([0-9]{1,2})(?:[.:]([0-9]{1,3}))?\]"
+)
+
+
+def load_pt(pt_file: str) -> torch.Tensor:
+    return torch.load(pt_file, map_location="cpu").squeeze().long()
+
+
+def get_token_slice(audio: torch.Tensor, start_s: float, end_s: float) -> str:
+    if start_s < 0:
+        start_s = 0
+    if end_s < 0:
+        end_s = 0
+    s_idx = int(start_s * TOKEN_PER_SECOND)
+    e_idx = int(end_s * TOKEN_PER_SECOND)
+    s_idx = max(0, min(s_idx, audio.shape[0]))
+    e_idx = max(0, min(e_idx, audio.shape[0]))
+    if e_idx <= s_idx:
+        sliced = []
+    else:
+        sliced = audio[s_idx:e_idx]
+    return "[SOA]" + "".join(f"<AUDIO_{int(i)}>" for i in sliced) + "[EOA]"
+
+
+def infer_pt_path(src_path: str) -> Optional[str]:
+    if not src_path:
+        return None
+    stem = os.path.splitext(os.path.basename(src_path))[0]
+    return os.path.join(PT_DIR, f"{stem}.pt")
+
+
+def parse_lyric_with_timestamps(lyric: str) -> List[Tuple[float, str]]:
+    """
+    Parse [(start_time_s, text), ...] from lyrics with timestamps, sorted by time ascending.
+    The returned timestamps come from the lyrics field in the meta file.
+    """
+    result: List[Tuple[float, str]] = []
+    matches = list(timestamp_pattern.finditer(lyric))
+    
+    for i, match in enumerate(matches):
+        start_idx = match.end()
+        if i + 1 < len(matches):
+            end_idx = matches[i + 1].start()
+        else:
+            end_idx = len(lyric)
+        
+        text = lyric[start_idx:end_idx].strip()
+        minutes = int(match.group(1))
+        seconds = int(match.group(2))
+        ms_str = match.group(3) if match.group(3) else "0"
+        
+        if len(ms_str) == 2:
+            fractional_seconds = int(ms_str) / 100.0
+        elif len(ms_str) == 3:
+            fractional_seconds = int(ms_str) / 1000.0
+        else:
+            fractional_seconds = int(ms_str) / 1000.0 if ms_str else 0.0
+        
+        total_seconds = minutes * 60 + seconds + fractional_seconds
+        result.append((total_seconds, text))
+    return result
+
+
+def extract_section_from_text(text: str) -> Optional[str]:
+    """
+    Relaxed rule: As long as [] contains English words (≥2 letters), return the entire bracket content as-is.
+    """
+    # Match if English words appear in the first []
+    m = re.search(r'\[([A-Za-z][A-Za-z0-9\s\-\(\)]*)\]', text)
+    if m:
+        return m.group(1).strip()   # Remove leading and trailing spaces
+    return None
+
+def format_section_label(sec_name: str) -> str:
+    """Keep original spaces, only trim leading and trailing whitespace."""
+    return sec_name.strip()
+
+
+def normalize_section_name(sec_name: str) -> str:
+    """
+    Normalize section name for matching:
+    - Remove all spaces
+    - Convert to lowercase
+    - Remove trailing digits (if any)
+    """
+    # Remove all spaces
+    normalized = sec_name.replace(" ", "").lower()
+    # Remove trailing digits (e.g., "verse1" -> "verse", "chorus1" -> "chorus")
+    normalized = re.sub(r"\d+$", "", normalized)
+    return normalized
+
+
+def clean_desc(desc: str) -> str:
+    """
+    Clean desc field:
+    1. If starts with [desc], remove it
+    2. If both ends are brackets, remove brackets
+    """
+    if not desc:
+        return desc
+    
+    desc = desc.strip()
+    
+    # If starts with [desc], remove it
+    if desc.startswith("[desc]"):
+        desc = desc[6:].strip()
+    
+    # If both ends are brackets, remove brackets
+    if desc.startswith("[") and desc.endswith("]"):
+        desc = desc[1:-1].strip()
+    
+    return desc
+
+
+def build_desc_map(desc_path_or_dir: str) -> Dict[Tuple[str, int], List[dict]]:
+    """
+    Support passing a single jsonl file or a directory containing multiple jsonl files.
+    In directory scenario, files are sorted by name and read sequentially, later records with the same key will overwrite earlier ones.
+    """
+    mapping: Dict[Tuple[str, int], List[dict]] = {}
+
+    paths: List[str] = []
+    if os.path.isdir(desc_path_or_dir):
+        for name in sorted(os.listdir(desc_path_or_dir)):
+            if name.endswith(".jsonl"):
+                paths.append(os.path.join(desc_path_or_dir, name))
+    else:
+        paths.append(desc_path_or_dir)
+
+    for path in paths:
+        with open(path, "r", encoding="utf-8") as f:
+            for line in f:
+                try:
+                    obj = json.loads(line)
+                except Exception:
+                    with open("error.txt", 'a', encoding='utf-8') as error_file:
+                        error_file.write(line + "\n")
+                    continue
+
+                song_id = obj.get("song_id")
+                track_idx = obj.get("track_index", 0)
+
+                # Change: Put the entire object
+                # sections = obj.get("sections", [])
+                # mapping[(song_id, track_idx)] = sections
+                mapping[(song_id, track_idx)] = obj
+    return mapping
+
+
+def extract_suffix_from_desc_path(desc_path: str, fallback_idx: int) -> str:
+    """
+    Extract suffix from desc filename for output file naming.
+    Rule: Try to extract <suffix> from filename pattern "3_block_<suffix>_(cn|en)_desc.jsonl".
+    If no match, use fallback_idx (starting from 0) converted to string.
+    """
+    fname = os.path.basename(desc_path)
+    m = re.search(r"3_block_([^_]+)_(?:cn|en)_desc\.jsonl", fname, re.IGNORECASE)
+    if m:
+        return m.group(1)
+    return str(fallback_idx)
+
+
+def extract_suffix_num(desc_path: str, fallback_idx: int) -> int:
+    """
+    Extract sortable numeric suffix for processing desc files in numeric order.
+    Use fallback_idx if parsing fails.
+    """
+    suffix = extract_suffix_from_desc_path(desc_path, fallback_idx)
+    try:
+        return int(suffix)
+    except ValueError:
+        return fallback_idx
+
+
+def build_messages(item: dict, obj: dict, audio: torch.Tensor) -> Optional[dict]:
+    """
+    Use timestamps from meta to split audio, desc_sections provide desc and text.
+    """
+    if not obj:
+        return None
+    desc_sections = obj.get("sections", [])
+
+    # Parse timestamps from meta's lyrics
+    lyrics_raw = item.get("lyrics", "") or ""
+    meta_ts_list = parse_lyric_with_timestamps(lyrics_raw)
+    if not meta_ts_list:
+        return None
+
+    total_seconds = audio.shape[0] / float(TOKEN_PER_SECOND)
+    
+    # Sort desc_sections (by startS, for matching)
+    desc_sections = sorted(desc_sections, key=lambda x: x.get("startS", 0.0))
+    
+    # Get middle sections from desc_sections (skip first Intro and last Outro)
+    # Intro and Outro are not in meta's lyrics, need separate handling
+    # Middle sections are matched in order
+    middle_desc_sections = desc_sections[1:-1] if len(desc_sections) > 2 else desc_sections[1:] if len(desc_sections) > 1 else []
+    
+    # Identify sections from meta timestamps and build mapping
+    # One section may contain multiple lyric lines, need to merge
+    section_timestamps: List[Tuple[str, float, float, str, str]] = []  # (section_name, start_s, end_s, text, desc)
+    current_section: Optional[Tuple[str, float, str]] = None  # (section_name, start_s, accumulated_text)
+    desc_idx = 0  # For matching desc in order (from middle_desc_sections)
+    
+    for idx, (start_s, text) in enumerate(meta_ts_list):
+        # Extract section name
+        section_name = extract_section_from_text(text)
+        
+        if section_name:
+            # Encountered new section label
+            # First save previous section (if any)
+            if current_section:
+                # Determine end time of previous section (current timestamp)
+                prev_sec_name, prev_start_s, prev_text = current_section
+                # Only remove timestamp, keep all other content (including line breaks, section labels, etc.)
+                clean_prev_text = re.sub(r"\[([0-9]{1,2}):([0-9]{1,2})(?:[.:]([0-9]{1,3}))?\]", "", prev_text)
+                
+                # Get desc in order (match from middle sections in order)
+                prev_desc = ""
+                if desc_idx < len(middle_desc_sections):
+                    prev_desc = clean_desc(middle_desc_sections[desc_idx].get("desc", ""))
+                    desc_idx += 1
+                
+                section_timestamps.append((prev_sec_name, prev_start_s, start_s, clean_prev_text, prev_desc))
+            
+            # Start new section
+            current_section = (section_name, start_s, text)
+        else:
+            # No section label, belongs to subsequent lines of current section
+            if current_section:
+                sec_name, sec_start, sec_text = current_section
+                # Preserve line breaks, connect with line breaks
+                current_section = (sec_name, sec_start, sec_text + "\n" + text)
+            # If no current_section, skip (might be empty line before Intro)
+    
+    # Process last section
+    # Check if last timestamp is empty text (indicates end marker)
+    outro_start_s: Optional[float] = None
+    if meta_ts_list and not meta_ts_list[-1][1].strip():
+        # Last timestamp is empty text, indicates end marker
+        outro_start_s = meta_ts_list[-1][0]
+    
+    if current_section:
+        sec_name, sec_start, sec_text = current_section
+        # If last timestamp is empty text, last section's end time should be this timestamp
+        # Otherwise use total duration
+        if outro_start_s is not None:
+            end_s = outro_start_s
+        else:
+            end_s = total_seconds
+        
+        # Only remove timestamp, keep all other content
+        clean_text = re.sub(r"\[([0-9]{1,2}):([0-9]{1,2})(?:[.:]([0-9]{1,3}))?\]", "", sec_text)
+        
+        # Get desc in order (match from middle sections in order)
+        desc = ""
+        if desc_idx < len(middle_desc_sections):
+            desc = clean_desc(middle_desc_sections[desc_idx].get("desc", ""))
+            desc_idx += 1
+        
+        section_timestamps.append((sec_name, sec_start, end_s, clean_text, desc))
+    
+    # Check if counts match (only check middle sections, excluding Intro and Outro)
+    if desc_idx != len(middle_desc_sections):
+        _log_warning(f"⚠️ Warning: Section count mismatch! meta has {len(section_timestamps)} sections, desc has {len(middle_desc_sections)} middle sections (excluding Intro and Outro) (song_id: {item.get('song_id')}, track_index: {item.get('track_index')})")
+    
+    if not section_timestamps:
+        return None
+
+    # Intro segment: from 0 to first section's start time
+    # Intro's desc should have been obtained in sequential matching, but Intro itself is not in meta's lyrics
+    # So need to get from desc_sections' first section (usually Intro)
+    first_section_start = section_timestamps[0][1] if section_timestamps else total_seconds
+    intro_desc = ""
+    if desc_sections and desc_sections[0].get("section", "").lower() == "intro":
+        intro_desc = clean_desc(desc_sections[0].get("desc", ""))
+
+    # Change: Use desc tag
+    # tag = item.get("tag", "")
+    song_id:str = obj.get("song_id", "")
+    omni_tag = obj.get("omni", "")
+    style_tag = obj.get("style", "")
+
+    if song_id.find("cn") != -1:
+        # Chinese songs use omni directly
+        tag = omni_tag
+    else:
+        # English songs compare omni / style
+        style_sim = obj.get("style_sim", 0)
+        omni_sim = obj.get("omni_sim", 0)
+        try:
+            tag = omni_tag if omni_sim > style_sim else style_tag
+        except Exception as e:
+            # If sim score is invalid, default to omni
+            tag = omni_tag
+            print(f"Error: {song_id}, {e}")
+
+    # Change: English
+    intro_prompt = (
+        f"Please generate a song in the following style:{tag}.\n"
+        "Next, I will tell you the requirements and lyrics for the song fragment to be generated, section by section.\n"
+        f"[Intro][desc:{intro_desc}]"
+    )
+
+    messages: List[dict] = []
+    messages.append({"role": "user", "content": intro_prompt})
+    messages.append(
+        {
+            "role": "assistant",
+            "content": get_token_slice(audio, 0.0, first_section_start),
+        }
+    )
+
+    # Process segment by segment (using timestamps from meta)
+    for idx, (sec_name, start_s, end_s, text, desc) in enumerate(section_timestamps):
+        # user content: [Section dsec : desc][Section lyrics : ...] (preserve original spaces)
+        label = format_section_label(sec_name)
+        content = f"[{label}]"
+        content += f"[desc:{desc}]"
+        lyrics_text = re.sub(r'^\[.*?\]\s*\n?', '', text.strip())
+        if lyrics_text:
+            content += f"[lyrics:\n{lyrics_text}]"
+        messages.append({"role": "user", "content": content})
+        messages.append(
+            {
+                "role": "assistant",
+                "content": get_token_slice(audio, start_s, end_s),
+            }
+        )
+    
+    # If last timestamp is empty text, add Outro segment
+    # Outro's desc should be obtained from desc_sections' last section (usually Outro)
+    if outro_start_s is not None and outro_start_s < total_seconds:
+        outro_desc = ""
+        if desc_sections and desc_sections[-1].get("section", "").lower() == "outro":
+            outro_desc = clean_desc(desc_sections[-1].get("desc", ""))
+        
+        messages.append({"role": "user", "content": f"[Outro][desc:{outro_desc}]"})
+        messages.append(
+            {
+                "role": "assistant",
+                "content": get_token_slice(audio, outro_start_s, total_seconds),
+            }
+        )
+
+    sample = {
+        "song_id": item.get("song_id"),
+        "track_index": item.get("track_index"),
+        "src_path": item.get("src_path"),
+        "tag": item.get("tag"),
+        "lang": item.get("lang"),
+        "duration": item.get("duration"),
+        "messages": messages,
+    }
+    return sample
+
+
+def process_with_desc(desc_path: str, suffix: str) -> None:
+    """
+    Generate output files using a single desc file (messages-only / meta-only).
+    Does not generate main file.
+    """
+    desc_map = build_desc_map(desc_path)
+
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    # messages-only naming: remove suffix, only keep block number (e.g., ..._8000.jsonl)
+    out_msg = os.path.join(OUTPUT_DIR, f"{OUTPUT_BASENAME}_{suffix}.jsonl")
+    out_meta = os.path.join(OUTPUT_DIR, f"{OUTPUT_BASENAME}_{suffix}_meta_only.jsonl")
+
+    dataset = load_jsonl(META_FILE)
+    total = len(dataset)
+    kept = 0
+    skipped = 0
+
+    for path in [out_msg, out_meta]:
+        if os.path.exists(path):
+            # Change: use past
+            already_msg = load_jsonl(out_msg)
+            already_meta = load_jsonl(out_meta)
+            assert len(already_meta) == len(already_msg)
+            dataset = dataset[len(already_msg):]
+    
+    with open(out_msg, "a", encoding="utf-8") as fout_msg, \
+         open(out_meta, "a", encoding="utf-8") as fout_meta:
+        
+        for item in tqdm(dataset, desc=f"Processing meta_suno_{LANG}.jsonl (desc: {suffix})"):
+            key = (item.get("song_id"), item.get("track_index", 0))
+
+            obj = desc_map.get(key)
+
+            if not obj:
+                skipped += 1
+                continue
+
+            pt_path = infer_pt_path(item.get("src_path", ""))
+            if not pt_path or not os.path.exists(pt_path):
+                skipped += 1
+                continue
+
+            audio = load_pt(pt_path)
+            sample = build_messages(item, obj, audio)
+            if not sample:
+                skipped += 1
+                continue
+
+            # Write messages-only (remove _messages_only suffix)
+            messages_only = {"messages": sample.get("messages", [])}
+            fout_msg.write(json.dumps(messages_only, ensure_ascii=False) + "\n")
+            # Write meta-only
+            meta_only = {k: v for k, v in sample.items() if k != "messages"}
+            fout_meta.write(json.dumps(meta_only, ensure_ascii=False) + "\n")
+            kept += 1
+
+    print(f"✅ messages-only: {out_msg}")
+    print(f"✅ meta-only: {out_meta}")
+    print(f"Total {total}, kept {kept}, skipped {skipped}")
+
+
+def convert_train_valid():
+    # Collect desc files
+    if not os.path.isdir(DESC_DIR):
+        print(f"⚠️ DESC_DIR does not exist: {DESC_DIR}")
+        return
+
+    unsorted_files = [
+        os.path.join(DESC_DIR, name)
+        for name in os.listdir(DESC_DIR)
+        if name.endswith(".jsonl")
+    ]
+
+    # Sort by extracted numeric suffix, ensure 8000 comes before 16000
+    desc_files = sorted(
+        unsorted_files,
+        key=lambda p: extract_suffix_num(p, 0),
+    )
+    if not desc_files:
+        print(f"⚠️ No desc files found: {DESC_DIR}")
+        return
+
+    # Change
+    # for idx, desc_path in enumerate(desc_files):
+    #     suffix = extract_suffix_from_desc_path(desc_path, idx)
+    #     process_with_desc(desc_path, suffix)
+
+    for desc_path in desc_files:
+        name = os.path.splitext(os.path.basename(desc_path))[0]
+        if name.endswith(LANG):
+            process_with_desc(desc_path, name)
+
+# assistant needs, but content is empty
+# Then each section inside has 3 descs, you take the one corresponding to the maximum value
+# omni also has two, take the larger one (style is not used)
+
+from meta_phonemes import _get_lyrics, _trans_sentences
+
+def _form_section(section:dict, en:bool) -> str:
+    """Process inside a section, determine which desc to select"""
+    # Segment label
+    section_tag = f"[{section['section']}]"
+    # Segment description
+    descs = [section['desc1'], section['desc2'], section['desc3']]
+    sims = [section['desc1_sim'], section['desc2_sim'], section['desc3_sim']]
+    max_sim = max(sims)
+    max_index = sims.index(max_sim)
+    desc:str = descs[max_index]
+    
+    if desc == "音频过短":  # "Audio too short" - keep original Chinese as it's part of data processing logic
+        desc = "[desc:]"
+    else:
+        DESC_START = "[desc] "
+        if desc.startswith(DESC_START):
+            desc = desc[len(DESC_START):] + "]"
+        desc = "[desc:" + desc[1:]
+    
+    # Lyrics & phonemes
+    text:str = section['text']
+    if text.find(']') != -1:
+        # Remove preceding segment label
+        start = text.rfind(']')
+        text = text[start+1:]
+    if len(text.strip()) == 0:
+        # Opening segment has no lyrics/phonemes
+        lyrics = ""
+        phonemes = ""
+    else:
+        if en:
+            lyrics = "[lyrics:\n" + clean_newlines(text) + "]"
+        else:
+            lyrics = "[lyrics:" + text + "]"
+
+        sentences, lyrics = _get_lyrics(lyrics)
+        phonemes = _trans_sentences(sentences)
+    return section_tag + desc + lyrics + phonemes
+
+def _form_intro(ele:dict) -> str:
+    """Process to get multi-turn dialogue opening"""
+    omni1 = ele['omni1']
+    omni2 = ele['omni2']
+    omni_sim1 = ele['omni1_sim']
+    omni_sim2 = ele['omni2_sim']
+    tag = omni1 if omni_sim1 > omni_sim2 else omni2
+    return (
+        f"Please generate a song in the following style:{tag}.\n"
+        "Next, I will tell you the requirements and lyrics for the song fragment to be generated, section by section.\n"
+    )
+
+def convert_test():
+    path = "filter.jsonl"
+    dataset = load_jsonl(path)
+    save_path = "messages.jsonl"
+    with open(save_path, 'w', encoding='utf-8') as file:
+        for ele in tqdm(dataset, desc=f"Converting {path}"):
+            messages = []
+            # Segment processing
+            sections = ele['sections']
+            id:str = ele['song_id']
+            english = id.startswith("suno_test_en")
+            for section in sections:
+                content = _form_section(section, english)
+                messages += [
+                    {
+                        "role": "user",
+                        "content": content
+                    },
+                    {
+                        "role": "assistant",
+                        "content": ""
+                    }
+                ]
+            # Initial addition
+            first_content = messages[0]['content']
+            intro = _form_intro(ele)
+            messages[0]['content'] = intro + first_content
+
+            data = {"messages": messages}
+            json.dump(data, file, ensure_ascii=False)
+            file.write("\n")
+
+if __name__ == "__main__":
+    convert_test()

data_pipeline/meta_process/convert_segments.py

@@ -0,0 +1,93 @@
+import os
+import json
+from tqdm import tqdm
+from my_tool import path_join, load_json
+from concurrent.futures import ProcessPoolExecutor, as_completed
+
+def _check_label(label:str, max_length:int=30) -> bool:
+    """Check if label is valid (non-empty, not timestamp, not long lyrics)"""
+    length = len(label.strip())
+    if length == 0:
+        # print("Error Label: Empty")
+        return False
+    if length > max_length:
+        # print(f"Error Label: Words - {label}")
+        return False
+    if label.find(":") != -1 and label.find(".") != -1:
+        # Considered as timestamp
+        # print(f"Error Label: Timestamp - {label}")
+        return False
+    return True
+
+def _convert_one(path:str):
+    """Segment a song's metadata, remove redundant content"""
+    data = load_json(path)
+    dir = os.path.dirname(path)
+    name = f"{data['song_id']}_{data['track_index']}.mp3"
+    path = path_join(dir, name)
+    new_data = {
+        "path": path,
+        "song_id": data['song_id'],
+        "segments": []
+    }
+    words_info = data['timestamped_lyrics']['alignedWords']  # Sentence-by-sentence information
+    seg_info = None
+
+    empty_head = False
+    for id, word_info in enumerate(words_info):
+        if not word_info['success']:
+            continue
+        word:str = word_info['word']
+        
+        label = ""
+        if word.startswith('['):
+            if seg_info is not None:
+                new_data['segments'].append(seg_info)
+            label_end = word.find(']')
+            label = word[1:label_end]
+            if not _check_label(label):
+                label = ""
+
+        if label != "":
+            seg_info = {
+                "start": word_info['startS'],
+                "end": 0,
+                "label": label,
+                "word": word[label_end+2:]
+            }
+        elif seg_info is not None:
+            seg_info['end'] = word_info['endS']
+            seg_info['word'] += word
+        else:
+            empty_head = True
+    if seg_info is not None:
+        seg_info['end'] = word_info['endS']
+        seg_info['word'] += word
+    else:
+        empty_head = True
+    if empty_head:
+        # print(f"Empty Head, segment: {len(new_data['segments'])}, path: {path}")
+        pass
+    return new_data
+
+# ===== External Interface =====
+
+def get_convert_segments(data_dir:str, save_path:str, max_workers:int=10):
+    paths = []
+    for name in tqdm(os.listdir(data_dir), desc="Getting the JSON Paths"):
+        if name.endswith(".json"):
+            path = path_join(data_dir, name)
+            paths.append(path)
+
+    dataset = []
+    with open(save_path, 'w', encoding='utf-8') as file:
+        with ProcessPoolExecutor(max_workers=max_workers) as executor:
+            futures = [executor.submit(_convert_one, path) for path in paths]
+            with tqdm(total=len(futures), desc="Converting Segments") as pbar:
+                for future in as_completed(futures):
+                    result = future.result()
+                    dataset.append(result)
+                    json.dump(result, file, ensure_ascii=False)
+                    file.write("\n")
+                    pbar.update(1)
+    return dataset

data_pipeline/meta_process/evaluate_polyphones.py

@@ -0,0 +1,62 @@
+import re
+import jieba
+from tqdm import tqdm
+from my_tool import load_jsonl, save_json, load_json
+from pypinyin import pinyin, Style, load_phrases_dict
+from pypinyin_dict.phrase_pinyin_data import cc_cedict
+
+cc_cedict.load()
+re_special_pinyin = re.compile(r'^(n|ng|m)$')
+
+# Add
+reference = load_json("poly_correct.json")
+load_phrases_dict(reference)
+
+def _filter(dataset:list[dict]):
+    """Filter non-polyphone characters in test set"""
+    new_dataset = []
+    for ele in tqdm(dataset, desc="Filtering"):
+        pos = ele['pos']
+        sentence = ele['sentence']
+        word = sentence[pos]
+        phones = pinyin(word, style=Style.NORMAL, heteronym=True)[0]
+        if len(phones) > 1:
+            new_dataset.append(ele)
+    print(f"Filter non polyphone, {len(dataset)} -> {len(new_dataset)}")
+    return new_dataset
+
+def evaluate_polyphones(dataset:list[dict], save_fail:str):
+    """Check pinyin processing accuracy for polyphones"""
+    dataset = _filter(dataset)
+    total = len(dataset)
+    right = 0
+    correct_dic = {}
+    for ele in tqdm(dataset):
+        pos = ele['pos']
+        phone = ele['phone']
+        sentence = ele['sentence']
+        seg_list = jieba.cut(sentence)
+        length = 0
+        for seg in seg_list:
+            if length <= pos and length + len(seg) > pos:
+                delta = pos - length    # Position in segment
+                break
+            length += len(seg)
+        pred_phones = pinyin(seg, style=Style.NORMAL)
+        pred_phone = pred_phones[delta][0]
+        if pred_phone == phone or pred_phone.endswith("v"):
+            right += 1
+        elif len(pred_phones) > 1:
+            # Corrected pronunciation (only meaningful for phrases)
+            pred_phones[delta] = [phone]
+            correct_dic[seg] = pred_phones
+    print(f"Acc: {(right / total):.2f}")
+    
+    origin_dic = load_json(save_fail)
+    merge_dic = origin_dic | correct_dic
+    save_json(merge_dic, save_fail)
+        
+if __name__ == "__main__":
+    path = "polyphones.jsonl"
+    dataset = load_jsonl(path)
+    evaluate_polyphones(dataset, "poly_correct.json")

data_pipeline/meta_process/filter.py

@@ -0,0 +1,46 @@
+import librosa
+from tqdm import tqdm
+from concurrent.futures import ProcessPoolExecutor, as_completed
+
+def filter_lang(dataset:list[dict], langs:list[str]) -> list[dict]:
+    """Filter dataset, only keep items with lang tag matching specified languages"""
+    new_dataset = []
+    for ele in tqdm(dataset, desc="Filtering Lang"):
+        if 'lang' not in ele or ele['lang'] not in langs :
+            continue
+        new_dataset.append(ele)
+    print(f"filter: {len(dataset)} -> {len(new_dataset)}")
+    return new_dataset
+
+def _check_duration(ele, lower_bound, upper_bound):
+    """Subprocess task: Check if audio duration is within range"""
+    duration = librosa.get_duration(filename=ele['path'])
+    if lower_bound != -1 and duration < lower_bound:
+        return None
+    if upper_bound != -1 and duration > upper_bound:
+        return None
+    return ele
+
+def filter_length(dataset:list[dict], lower_bound:int=-1, upper_bound:int=-1, max_worker:int=4) -> list[dict]:
+    """Filter dataset, only keep items with length in [lower_bound, upper_bound], if set to -1 then no limit on that side"""
+    new_dataset = []
+    with ProcessPoolExecutor(max_workers=max_worker) as executor:
+        futures = [
+            executor.submit(_check_duration, ele, lower_bound, upper_bound)
+            for ele in dataset
+        ]
+        with tqdm(total=len(futures), desc="Filtering Length") as pbar:
+            for future in as_completed(futures):
+                result = future.result()
+                if result is not None:
+                    new_dataset.append(result)
+                pbar.update(1)
+        # for ele in tqdm(dataset, desc="Filtering Length"):
+        #     duration = librosa.get_duration(filename=ele['path'])
+        #     if lower_bound != -1 and duration < lower_bound:
+        #         continue
+        #     if upper_bound != -1 and duration > upper_bound:
+        #         continue
+        #     new_dataset.append(ele)
+        print(f"filter: {len(dataset)} -> {len(new_dataset)}")
+    return new_dataset

data_pipeline/meta_process/main.py

@@ -0,0 +1,77 @@
+from my_tool import (
+    load_json,
+    load_jsonl, 
+    load_txt, 
+    save_jsonl, 
+    format_meta, 
+    pure_name, 
+    BASE_DIR,
+    compose_analyze,
+    get_sample,
+    get_field_suno,
+    tags_analyze,
+    find_json,
+    show_dir,
+    convert_mp3,
+    tar_dir,
+    tar_size_check,
+    clean_newlines,
+    dict_sort_print,
+)
+from meta_lang import load_asr_model, get_lang_meta
+from meta_tags import load_tag_model, get_tags_meta
+from meta_endpoints import get_endpoints_meta
+from meta_phonemes import get_phonemes_meta
+from filter import filter_lang, filter_length
+from convert_convs import get_convert_convs
+from convert_segments import get_convert_segments
+from convert_lyrics import get_convert_lyrics, get_match_music
+
+def pipeline():
+    import os
+    dir = "suno_batch"
+    name = pure_name(dir)
+    save_dir = BASE_DIR / f"data/{name}"
+
+    # Initialize paths (only once)
+    os.makedirs(save_dir, exist_ok=True)
+    raw_path = os.path.join(save_dir, "raw.jsonl")
+    if os.path.exists(raw_path):
+        dataset = load_jsonl(raw_path)
+    else:
+        dataset = format_meta(dir)
+    save_jsonl(dataset, raw_path)
+
+    # Length filtering
+    dataset = dataset[:1000]
+    max_workers = 10
+    dataset = filter_length(dataset, 120, 360, max_workers)
+
+    # Language tagging
+    # lang_bs = 8
+    # model = load_asr_model(lang_bs)
+    # lang_path = os.path.join(save_dir, "meta_lang.jsonl")
+    # dataset = get_lang_meta(model, dataset, lang_bs, lang_path)
+
+    # Language filtering
+    # dataset = filter_lang(dataset, ['zh', 'en'])
+
+    # Style tagging
+    tag_bs = 4
+    tag_path = os.path.join(save_dir, "meta_tags.jsonl")
+    model, processor = load_tag_model()
+    prompt_path = BASE_DIR / "prompts/new_tags.md"
+    prompt = load_txt(prompt_path)
+    get_tags_meta(model, processor, dataset, prompt, tag_bs, tag_path)
+
+def repeat(func):
+    while True:
+        try:
+            func()
+            break
+        except Exception as e:
+            print(f"Error: {e}")
+            continue
+
+if __name__ == "__main__":
+    repeat(pipeline)

data_pipeline/meta_process/meta_endpoints.py

@@ -0,0 +1,118 @@
+import json
+import webrtcvad
+import collections
+from tqdm import tqdm
+from my_tool import dup_remove
+from pydub import AudioSegment
+from concurrent.futures import ProcessPoolExecutor, as_completed
+
+def _frame_generator(frame_duration_ms, audio, sample_rate):
+    """Split audio into frames"""
+    bytes_per_sample = 2
+    frame_size = int(sample_rate * frame_duration_ms / 1000.0) * bytes_per_sample
+    offset = 0
+    timestamp = 0.0
+    frame_duration = frame_duration_ms / 1000.0
+    while offset + frame_size < len(audio):
+        yield audio[offset:offset + frame_size], timestamp
+        timestamp += frame_duration
+        offset += frame_size
+
+def _vad_collector(sample_rate, frame_duration_ms, padding_duration_ms, vad, frames):
+    """Merge continuous vocal segments based on webrtcvad"""
+    num_padding_frames = int(padding_duration_ms / frame_duration_ms)
+    ring_buffer = collections.deque(maxlen=num_padding_frames)
+
+    triggered = False
+    speech_segments = []
+
+    for frame_bytes, timestamp in frames:
+        is_speech = vad.is_speech(frame_bytes, sample_rate)
+
+        if not triggered:
+            ring_buffer.append((frame_bytes, timestamp, is_speech))
+            num_voiced = len([f for f in ring_buffer if f[2]])
+            if num_voiced > 0.9 * ring_buffer.maxlen:
+                triggered = True
+                start_time = ring_buffer[0][1]
+                ring_buffer.clear()
+        else:
+            ring_buffer.append((frame_bytes, timestamp, is_speech))
+            num_unvoiced = len([f for f in ring_buffer if not f[2]])
+            if num_unvoiced > 0.9 * ring_buffer.maxlen:
+                end_time = timestamp + (frame_duration_ms / 1000.0)
+                speech_segments.append((start_time, end_time))
+                triggered = False
+                ring_buffer.clear()
+
+    # If still in speech state at the end, close the last segment
+    if triggered:
+        end_time = timestamp + (frame_duration_ms / 1000.0)
+        speech_segments.append((start_time, end_time))
+
+    return speech_segments
+
+def _one_process(path):
+    """Detect vocal segments in an audio"""
+    # 1. Compress audio
+    audio = AudioSegment.from_file(path)
+    audio = audio.set_frame_rate(16000).set_channels(1).set_sample_width(2)
+    sample_rate = audio.frame_rate
+    audio_data = audio.raw_data
+
+    # 2. Initialize VAD (0-3, higher value means more likely to be considered as speech)
+    vad = webrtcvad.Vad(0)
+
+    # 3. Generate frames
+    frames = list(_frame_generator(30, audio_data, sample_rate))
+
+    # 4. Detect vocal intervals
+    segments = _vad_collector(sample_rate, 30, 300, vad, frames)
+
+    # If no vocals, set both start and end to -1
+    if len(segments) == 0:
+        return {
+            "start": -1,
+            "end": -1,
+            "segments": [],
+        }
+
+    return {
+        "start": segments[0][0],
+        "end": segments[-1][1],
+        "segments": segments,
+    }
+
+# ===== External Interface =====
+
+def get_endpoints_meta(dataset:list[dict], save_path:str, max_workers:int=4, save_middle:bool=True):
+    """
+    Add endpoint labels to each audio in dataset (mainly for separated vocal audio)
+    - Requires 'path' field in each data entry in dataset
+    - Write fields: endpoints.start/end
+    - Write to save_path in real-time
+    - save_middle determines whether to record each sentence's endpoints to save.segments field
+    """
+    dataset = dup_remove(dataset, save_path, 'path', 'endpoints')
+    new_dataset = []
+    with open(save_path, 'a', encoding='utf-8') as file:
+        with ProcessPoolExecutor(max_workers=max_workers) as executor:
+            futures = {executor.submit(_one_process, ele['path']): ele for ele in dataset}
+            for future in tqdm(as_completed(futures), desc="Detecting endpoints"):
+                ele = futures[future]   # Get original element
+                try:
+                    result = future.result()
+                    ele['endpoints'] = {
+                        "start": result['start'],
+                        "end": result['end']
+                    }
+                    if save_middle:
+                        if "save" not in ele:
+                            ele['save'] = {}
+                        ele['save']['segments'] = result['segments']
+                    new_dataset.append(ele)
+                    json.dump(ele, file, ensure_ascii=False)
+                    file.write("\n")
+                except Exception:
+                    pass
+    return new_dataset

data_pipeline/meta_process/meta_lang.py

@@ -0,0 +1,125 @@
+import os
+import json
+from tqdm import tqdm
+from funasr import AutoModel
+from typing import List, Tuple
+from collections import defaultdict
+from my_tool import get_free_gpu, dup_remove
+
+# ===== ASR Model (External) =====
+
+def load_asr_model(bs:int):
+    """Load lyric recognition model"""
+    device = f"cuda:{get_free_gpu()}"
+    model = AutoModel(
+        model="iic/SenseVoiceSmall",
+        trust_remote_code=True,
+        vad_model="fsmn-vad",
+        vad_kwargs={"max_single_segment_time": 30000},
+        device=device,
+        batch_size=bs,
+        max_batch_size=bs * 2,
+    )
+    print(f"Using {device}")
+    return model
+
+# ===== ASR Parsing =====
+
+def _struct2lang(text: str) -> str:
+    """Extract language identifier from structured representation"""
+    start = text.find("|")
+    text = text[start+1:]
+    end = text.find("|")
+    return text[:end]
+
+def _struct2lyrics(text:str) -> str:
+    start = text.rfind(">")
+    lyric = text[start+1:]
+    return lyric
+
+def _struct_parse(text: str) -> Tuple[List[str], List[str]]:
+    """Split structured information sentence by sentence and then parse"""
+    texts = text.split(" <")
+    langs, lyrics = [], []
+    for ele in texts:
+        langs.append(_struct2lang(ele))
+        lyrics.append(_struct2lyrics(ele))
+    return langs, lyrics
+
+# ===== ASR Processing =====
+
+def _batch_asr(model, paths:List[str]) -> List[Tuple[List[str], List[str]]]:
+    """Batch speech recognition"""
+    outputs = model.generate(
+        input=paths,
+        cache=None,
+        language="auto",
+        use_itn=True,
+        batch_size_s=240,
+        merge_vad=True,
+        merge_length_s=15,
+    )
+    return [_struct_parse(output['text']) for output in outputs]
+
+# ===== Overall Language Detection =====
+
+def _lang_decide(lang_lyrics:list[tuple[str, str]], val_limit:int=5, word_limit=5) -> str:
+    """
+    Determine language based on sentence recognition information
+    - val_limit: Only count if there are at least this many sentences
+    - word_limit: Only count if a sentence has at least this many words
+    """
+    lang_count = defaultdict(int)
+    seg_langs, seg_lyrics = lang_lyrics
+    for lang, lyric in zip(seg_langs, seg_lyrics):
+        lyric = lyric.strip()
+        if lang == "en":
+            words_num = len(lyric.split())
+        else:
+            words_num = len(lyric)
+        if words_num >= word_limit:
+            lang_count[lang] += 1
+    langs = []
+    for lang, count in lang_count.items():
+        if count >= val_limit:
+            langs.append(lang)
+    if len(langs) == 0:
+        return "pure"
+    elif len(langs) == 1:
+        return langs[0]
+    else:
+        return "multi: " + " ".join(langs)
+
+# ===== External Interface =====
+
+def get_lang_meta(model, dataset:list[dict], bs:int, save_path:str, save_middle:bool=True) -> list[dict]:
+    """
+    Perform language recognition on a JSONL dataset
+    - Final language tag is saved to lang field, types include zh, en, ja, ko, yue, pure, multi, etc.
+    - save_middle determines whether to save intermediate recognition results (sentence languages and lyrics) to save.langs, save.lyrics
+    """
+    data_num = len(dataset)
+    dataset = dup_remove(dataset, save_path, 'path', 'lang')
+    new_dataset = []
+    with open(save_path, 'a', encoding='utf-8') as file:
+        for i in tqdm(range(0, data_num, bs), desc="Lang detecting"):
+            batch = []
+            paths = []
+            for ele in dataset[i:i+bs]:
+                path = ele['path']
+                if os.path.exists(path):
+                    batch.append(ele)
+                    paths.append(path)
+            lang_lyrics_lis = _batch_asr(model, paths)
+            langs = [_lang_decide(lang_lyrics) for lang_lyrics in lang_lyrics_lis]
+            for ele, (seg_langs, seg_lyrics), lang in zip(batch, lang_lyrics_lis, langs):
+                ele['lang'] = lang
+                if save_middle:
+                    if 'save' not in ele:
+                        ele['save'] = {}
+                    ele['save']['langs'] = seg_langs
+                    ele['save']['lyrics'] = seg_lyrics
+                new_dataset.append(ele)
+                json.dump(ele, file, ensure_ascii=False)
+                file.write("\n")
+    return new_dataset

data_pipeline/meta_process/meta_phonemes.py

@@ -0,0 +1,283 @@
+import os
+import re
+import json
+import copy
+import jieba
+import string
+from tqdm import tqdm
+from g2p_en import G2p
+from my_tool import BASE_DIR, load_json
+from pypinyin import pinyin, Style, load_phrases_dict
+from pypinyin_dict.phrase_pinyin_data import cc_cedict
+
+cc_cedict.load()
+re_special_pinyin = re.compile(r'^(n|ng|m)$')
+reference = load_json("poly_correct.json")
+load_phrases_dict(reference)
+
+# ===== Chinese Conversion =====
+
+def _split_py(py):
+    """Split pinyin with tone number into initial (sm) and final (ym) parts"""
+    tone = py[-1]
+    py = py[:-1]
+    sm = ""
+    ym = ""
+    suf_r = ""
+    if re_special_pinyin.match(py):
+        py = 'e' + py
+    if py[-1] == 'r':
+        suf_r = 'r'
+        py = py[:-1]
+
+    if len(py) == 0:
+        # rx
+        return "", suf_r + tone
+
+    if py == 'zi' or py == 'ci' or py == 'si' or py == 'ri':
+        sm = py[:1]
+        ym = "ii"
+    elif py == 'zhi' or py == 'chi' or py == 'shi':
+        sm = py[:2]
+        ym = "iii"
+    elif py == 'ya' or py == 'yan' or py == 'yang' or py == 'yao' or py == 'ye' or py == 'yong' or py == 'you':
+        sm = ""
+        ym = 'i' + py[1:]
+    elif py == 'yi' or py == 'yin' or py == 'ying':
+        sm = ""
+        ym = py[1:]
+    elif py == 'yu' or py == 'yv' or py == 'yuan' or py == 'yvan' or py == 'yue ' or py == 'yve' or py == 'yun' or py == 'yvn':
+        sm = ""
+        ym = 'v' + py[2:]
+    elif py == 'wu':
+        sm = ""
+        ym = "u"
+    elif py[0] == 'w':
+        sm = ""
+        ym = "u" + py[1:]
+    elif len(py) >= 2 and (py[0] == 'j' or py[0] == 'q' or py[0] == 'x') and py[1] == 'u':
+        sm = py[0]
+        ym = 'v' + py[2:]
+    else:
+        seg_pos = re.search('a|e|i|o|u|v', py)
+        try:
+            sm = py[:seg_pos.start()]
+            ym = py[seg_pos.start():]
+            if ym == 'ui':
+                ym = 'uei'
+            elif ym == 'iu':
+                ym = 'iou'
+            elif ym == 'un':
+                ym = 'uen'
+            elif ym == 'ue':
+                ym = 've'
+        except Exception:
+            sm = ym = ""
+            return sm, ym
+    ym += suf_r + tone
+    return sm, ym
+
+# All Chinese punctuation
+chinese_punctuation_pattern = r'[\u3002\uff0c\uff1f\uff01\uff1b\uff1a\u201c\u201d\u2018\u2019\u300a\u300b\u3008\u3009\u3010\u3011\u300e\u300f\u2014\u2026\u3001\uff08\uff09]'
+
+def _has_ch_punc(text):
+    match = re.search(chinese_punctuation_pattern, text)
+    return match is not None
+
+def _has_en_punc(text):
+    return text in string.punctuation
+
+def _trans_cn(text:str, with_sp=True):
+    """Convert Chinese to phonemes"""
+    phonemes = []
+    # Word segmentation
+    seg_list = jieba.cut(text)
+    # Process word by word
+    for seg in seg_list:
+        # String validity
+        if seg.strip() == "": continue
+        # seg_tn = tn_chinese(seg)
+        # Convert to pinyin (without tone)
+        py =[_py[0] for _py in pinyin(seg, style=Style.TONE3, neutral_tone_with_five=True)]
+        # Punctuation detection (skip if present)
+        if any([_has_ch_punc(_py) for _py in py])  or any([_has_en_punc(_py) for _py in py]):
+            continue
+        # Split pinyin
+        # phonemes += _split_py(_py)
+        for _py in py:
+            sm, ym = _split_py(_py)
+            if sm != "":
+                phonemes.append(sm)
+            if ym != "":
+                phonemes.append(ym)
+        if with_sp:
+            phonemes += ["sp"]
+    return phonemes
+
+# ===== English Conversion =====
+
+def _read_lexicon(lex_path):
+    """Read English lexicon"""
+    lexicon = {}
+    with open(lex_path) as f:
+        for line in f:
+            temp = re.split(r"\s+", line.strip("\n"))
+            word = temp[0]
+            phones = temp[1:]
+            if word.lower() not in lexicon:
+                lexicon[word.lower()] = phones
+    return lexicon
+
+LEX_PATH = BASE_DIR / f"data/ref/lexion.txt"
+lexicon = _read_lexicon(LEX_PATH)
+
+g2p = G2p()
+
+def _trans_en(word:str, with_sp=True):
+    """Convert English (word) to phonemes"""
+    w_lower = word.lower()
+    phonemes = []
+    if w_lower in lexicon:
+        # Use lexicon if available (cannot directly get reference)
+        phonemes += lexicon[w_lower]
+    else:
+        # Use G2P if not in lexicon
+        phonemes = g2p(w_lower)
+        if not phonemes:
+            phonemes = []
+        # Add to lexicon
+        lexicon[w_lower] = phonemes
+    if len(phonemes) > 0 and with_sp:
+        phonemes.append("sp")
+    return phonemes
+
+# ===== Single Sentence Processing =====
+
+def _char_lang(c:str) -> int:
+    """
+    Check if a character is Chinese, English, or other
+    0 - Chinese
+    1 - English
+    2 - Number
+    3 - Other
+    """
+    if '\u4e00' <= c <= '\u9fff':
+        return 0
+    elif ('a' <= c <= 'z') or ('A' <= c <= 'Z'):
+        return 1
+    elif c.isdigit():
+        return 2
+    else:
+        return 3
+
+NUMBER_MAP = {
+    "0": "zero",
+    "1": "one",
+    "2": "two",
+    "3": "three",
+    "4": "four",
+    "5": "five",
+    "6": "six",
+    "7": "seven",
+    "8": "eight",
+    "9": "nine",
+}
+
+def _lang_seperate(text:str) -> list[str]:
+    """Split string by language"""
+    lang_segs = []      # Set of split strings
+    lang_tags = []      # Tags for each string segment
+    lang_seg = ""       # Previous continuous language string
+    lang_tag = -1       # Language type of previous character
+    en_count = 0
+    for c in text:
+        lang = _char_lang(c)
+        if lang_tag != lang:
+            # Different from previous character type
+            if lang_seg != "":
+                lang_segs.append(lang_seg)
+                lang_tags.append(lang_tag)
+                if lang_tag == 1:
+                    en_count += 1
+                lang_seg = ""
+            if lang == 2 and en_count >= 4:
+                # Number conversion in English
+                lang_segs.append(NUMBER_MAP[c])
+                lang_tags.append(1)
+            lang_tag = lang
+        if lang < 2:
+            lang_seg += c
+    if lang_seg != "":
+        # Last valid segment
+        lang_segs.append(lang_seg)
+        lang_tags.append(lang_tag)
+    return lang_segs, lang_tags
+
+def _phoneme_trans(text:str, with_sp=True):
+    """Convert a lyric segment to phonemes"""
+    # Split by language
+    lang_segs, lang_tags = _lang_seperate(text)
+    # Convert segment by segment
+    phonemes = []
+    for lang_seg, lang_tag in zip(lang_segs, lang_tags):
+        if lang_tag == 0:
+            # Chinese
+            phonemes += _trans_cn(lang_seg, with_sp)
+        else:
+            # English
+            phonemes += _trans_en(lang_seg, with_sp)
+    return phonemes
+
+# ===== Dynamic Adaptation =====
+
+def _get_lyrics(raw_content:str) -> list[str]:
+    """Extract lyric content from dialogue, format like '[stage][dsec:xxx][lyrics:xxx\nxxx]'"""
+    START_FORMAT = "[lyrics:"
+    start = raw_content.find(START_FORMAT)
+    if start == -1:
+        return None, None
+    content = raw_content[start+len(START_FORMAT):-1]
+    # Filter brackets
+    content = re.sub(r'\[.*?\]', '', content)   # Complete brackets
+    content = re.sub(r'[\[\]]', '', content)    # Unclosed brackets
+    # Split sentences
+    sentences = content.split("\n")
+    # Reconstruct
+    new_content = raw_content[:start] + START_FORMAT + content + "]"
+    return sentences, new_content
+
+def _trans_sentences(sentences:list[str], with_sp:bool=True) -> str:
+    """Convert sentence list to wrapped phoneme string"""
+    phonemes_lis = []
+    for sentence in sentences:
+        phonemes = _phoneme_trans(sentence, with_sp)
+        phonemes_lis.append(" ".join(phonemes))
+    # Wrap
+    phonemes_str = '\n'.join(phonemes_lis)
+    envelope = f"[phoneme:{phonemes_str}]"
+    envelope = re.sub(r'\d+', '', envelope)     # Remove tones
+    return envelope
+
+# ===== External Interface =====
+
+def get_phonemes_meta(dataset:list[dict], save_path:str, with_sp:bool=True):
+    """Add phonemes to lyrics in dataset"""
+    new_dataset = []
+    with open(save_path, 'w', encoding='utf-8') as file:
+        for ele in tqdm(dataset, desc="Phoneme trans"):
+            ele = copy.deepcopy(ele)
+            messages = ele['messages']
+            # Skip first message, process subsequent ones sentence by sentence
+            for message in messages[1:]:
+                if message['role'] == "assistant":
+                    continue
+                content = message['content']
+                sentences, new_content = _get_lyrics(content)
+                if sentences is None:
+                    continue
+                phonemes = _trans_sentences(sentences, with_sp)
+                message['content'] = new_content + phonemes
+            new_dataset.append(ele)
+            json.dump(ele, file, ensure_ascii=False)
+            file.write("\n")
+    return new_dataset

data_pipeline/meta_process/meta_tags.py

@@ -0,0 +1,124 @@
+import os
+import json
+import torch
+from tqdm import tqdm
+from transformers import (
+    Qwen3OmniMoeProcessor,
+    Qwen3OmniMoeForConditionalGeneration
+)
+from qwen_omni_utils import process_mm_info
+from my_tool import get_free_gpu, audio_cut, extract_json, dup_remove, BASE_DIR
+
+# ===== Tag Model and Processor (External) =====
+
+def load_tag_model():
+    """Load tag model"""
+    device = f"cuda:{get_free_gpu()}"
+    print(f"Using {device}")
+    model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+    model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(
+        model_name,
+        dtype=torch.bfloat16,
+        # local_files_only=True
+    ).to(device)
+    model.disable_talker()
+    model.eval()
+
+    processor = Qwen3OmniMoeProcessor.from_pretrained(
+        model_name,
+        # local_files_only=True
+    )
+    return model, processor
+
+# ===== Tag Annotation =====
+
+def _format_messages(prompt:str, path:str) -> list[dict]:
+    """Construct messages to pass to omni"""
+    messages = [
+        {
+            "role": "system",
+            "content": [
+                {"type": "text", "text": prompt}
+            ]
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "audio", "audio": path},
+            ]
+        }
+    ]
+    return messages
+
+def _batch_tagging(model, processor, paths:list[str], prompt:str, mode="random"):
+    """Annotate a batch of songs"""
+    convs = []
+    middle_paths = []
+    output_dir = BASE_DIR / "data/temp"
+    for path in paths:
+        seg_path = audio_cut(path, mode, output_dir)
+        middle_paths.append(seg_path)
+        messages = _format_messages(prompt, seg_path)
+        convs.append(messages)
+
+    USE_AUDIO_IN_VIDEO = False
+
+    with torch.no_grad():
+        text = processor.apply_chat_template(convs, add_generation_prompt=True, tokenize=False)
+        audios, images, videos = process_mm_info(convs, use_audio_in_video=USE_AUDIO_IN_VIDEO)
+        inputs = processor(
+            text=text,
+            audio=audios,
+            padding=True, 
+            images=images,
+            videos=videos,
+            return_tensors="pt",
+            use_audio_in_video=USE_AUDIO_IN_VIDEO
+        )
+        inputs = inputs.to(model.device).to(model.dtype)
+
+        text_ids = model.generate(
+            **inputs,
+            max_new_tokens=2048,
+            return_audio=False,
+            thinker_return_dict_in_generate=True,
+            use_audio_in_video=USE_AUDIO_IN_VIDEO
+        )
+        gene_texts = processor.batch_decode(
+            text_ids[0].sequences[:, inputs["input_ids"].shape[1] :],
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False
+        )
+
+    torch.cuda.empty_cache()
+    # Delete audio segments
+    for path in middle_paths:
+        if os.path.exists(path):
+            os.remove(path)
+    return gene_texts
+
+# ===== External Interface =====
+
+def get_tags_meta(model, processor, dataset:list[dict], prompt:str, bs:int, save_path:str):
+    data_num = len(dataset)
+    dataset = dup_remove(dataset, save_path, 'path', 'tags')
+    new_dataset = []
+    with open(save_path, 'a', encoding="utf-8") as file:
+        for i in tqdm(range(0, data_num, bs)):
+            batch = []
+            paths = []
+            for ele in dataset[i:i+bs]:
+                path = ele['path']
+                if os.path.exists(path):
+                    batch.append(ele)
+                    paths.append(path)
+            contents = _batch_tagging(model, processor, paths, prompt)
+            for ele, content in zip(batch, contents):
+                ckeck, json_data = extract_json(content)
+                if not ckeck:
+                    continue
+                ele['tags'] = json_data['tags']
+                new_dataset.append(ele)
+                json.dump(ele, file, ensure_ascii=False)
+                file.write('\n')
+    return new_dataset

data_pipeline/meta_process/meta_vocal.py

@@ -0,0 +1,141 @@
+import os
+import math
+from typing import List, Dict
+
+import torch
+import librosa
+import numpy as np
+
+from demucs.pretrained import get_model
+from demucs.apply import apply_model
+
+
+# ======================
+# Basic Configuration
+# ======================
+
+SAMPLE_RATE = 44100
+MAX_DURATION = 5          # Only take first 30 seconds
+BATCH_SIZE = 8               # Can adjust to 16~32 for 80G GPU memory
+VOCAL_DB_THRESHOLD = -35.0   # Vocal presence threshold (empirical value)
+# DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE = "cuda:1"
+
+
+# ======================
+# Audio Loading (first 30s only)
+# ======================
+
+def load_audio_30s(path: str, sr: int = SAMPLE_RATE) -> torch.Tensor:
+    """
+    Returns shape: [channels=2, samples]
+    """
+    y, _ = librosa.load(
+        path,
+        sr=sr,
+        mono=False,
+        duration=MAX_DURATION
+    )
+
+    if y.ndim == 1:
+        y = np.stack([y, y], axis=0)
+
+    return torch.from_numpy(y)
+
+
+# ======================
+# dB Calculation
+# ======================
+
+def rms_db(wav: torch.Tensor) -> float:
+    """
+    wav: [channels, samples]
+    """
+    rms = torch.sqrt(torch.mean(wav ** 2))
+    db = 20 * torch.log10(rms + 1e-8)
+    return db.item()
+
+
+# ======================
+# Demucs Vocal Detection
+# ======================
+
+class DemucsVocalDetector:
+    def __init__(self):
+        self.model = (
+            get_model("htdemucs")
+            .to(DEVICE)
+            .eval()
+        )
+        self.vocal_idx = self.model.sources.index("vocals")
+
+    @torch.no_grad()
+    def batch_has_vocal(self, audio_paths: List[str]) -> Dict[str, bool]:
+        """
+        Input: List of audio paths
+        Output: {path: whether has vocals}
+        """
+        results = {}
+
+        batch_wavs = []
+        batch_paths = []
+
+        print("start load")
+        for path in audio_paths:
+            try:
+                wav = load_audio_30s(path)
+                batch_wavs.append(wav)
+                batch_paths.append(path)
+            except Exception as e:
+                results[path] = False
+                continue
+
+        print("finish load")
+        self._process_batch(batch_wavs, batch_paths, results)
+
+        return results
+
+    def _process_batch(self, wavs, paths, results):
+        max_len = max(w.shape[1] for w in wavs)
+        padded = []
+
+        for w in wavs:
+            if w.shape[1] < max_len:
+                w = torch.nn.functional.pad(w, (0, max_len - w.shape[1]))
+            padded.append(w)
+
+        batch = torch.stack(padded, dim=0).to(DEVICE)
+
+        print("start demucs")
+        torch.cuda.synchronize()
+
+        sources = apply_model(
+            self.model,
+            batch,
+            SAMPLE_RATE,
+            device=DEVICE,
+            split=False,       # 🔥 Core
+            progress=False
+        )
+
+        torch.cuda.synchronize()
+        print("demucs done")
+
+        vocals = sources[:, self.vocal_idx]
+
+        for i, path in enumerate(paths):
+            db = rms_db(vocals[i])
+            results[path] = db > VOCAL_DB_THRESHOLD
+
+# ======================
+# Usage Example
+# ======================
+
+if __name__ == "__main__":
+    audio_list = []
+
+    detector = DemucsVocalDetector()
+    result = detector.batch_has_vocal(audio_list)
+
+    for k, v in result.items():
+        print(f"{k}: {'Has' if v else 'No'} vocals")

data_pipeline/meta_process/my_tool.py

@@ -0,0 +1,551 @@
+import os
+import re
+import json
+import random
+import tarfile
+import subprocess
+import json_repair
+from tqdm import tqdm
+from pathlib import Path
+from pydub import AudioSegment
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed
+
+# ===== Macros =====
+
+BASE_DIR = Path(__file__).parent
+
+# ===== Helper Functions =====
+
+def pure_name(path:str):
+    """Get the original name of a file path (without extension)"""
+    basename = os.path.basename(path)
+    dot_pos = basename.rfind('.')
+    if dot_pos == -1:
+        return basename
+    return basename[:dot_pos]
+
+def extract_json(text: str) -> tuple[bool, dict]:
+    """Extract and repair JSON data from text (enhanced error-tolerant version)
+    
+    Features:
+    1. Automatically identify code block markers (```json``` or ```)
+    2. Fix common JSON errors (mismatched quotes, trailing commas, etc.)
+    3. Support lenient parsing mode
+    
+    Returns: (success, parsed dictionary)
+    """
+    # Preprocessing: extract possible JSON content area
+    content = text
+    
+    # Case 1: Check ```json``` code block
+    if '```json' in text:
+        start = text.find('```json')
+        end = text.find('```', start + 6)
+        content = text[start + 6:end].strip()
+    # Case 2: Check regular ``` code block
+    elif '```' in text:
+        start = text.find('```')
+        end = text.find('```', start + 3)
+        content = text[start + 3:end].strip()
+    
+    # Clean common interference items in content
+    content = re.sub(r'^[^{[]*', '', content)  # Remove unstructured content before JSON
+    content = re.sub(r'[^}\]]*$', '', content)  # Remove unstructured content after JSON
+    
+    # Try standard parsing
+    try:
+        json_data = json.loads(content)
+        return True, json_data
+    except json.JSONDecodeError as e:
+        standard_error = e
+    
+    # Try to repair with json_repair
+    try:
+        repaired = json_repair.repair_json(content)
+        json_data = json.loads(repaired)
+        return True, json_data
+    except Exception as e:
+        repair_error = e
+        return False, {
+            "standard_error": standard_error,
+            "repair_error": repair_error
+        }
+
+def path_join(dir, name):
+    return os.path.join(dir, name)
+
+def dict_sort_print(dic:dict, value:bool=True, reverse=True):
+    """Sort a dictionary by value size and output"""
+    idx = 1 if value else 0
+    sorted_lis = sorted(dic.items(), key=lambda x: x[idx], reverse=reverse)
+    sorted_dic = {}
+    for key, value in sorted_lis:
+        sorted_dic[key] = value
+    print(json.dumps(sorted_dic, indent=4, ensure_ascii=False))
+
+def clean_newlines(text: str) -> str:
+    """
+    Clean lyric line breaks:
+    1. Keep line breaks after punctuation
+    2. Convert line breaks after non-punctuation → space
+    3. Fix extra spaces after English apostrophes
+    4. Merge redundant spaces
+    5. Preserve paragraph structure, ensure line breaks after punctuation
+    """
+    if not text:
+        return ""
+
+    text = text.strip()
+
+    # First unify line breaks to \n
+    text = text.replace('\r\n', '\n').replace('\r', '\n')
+
+    # Merge non-empty lines into one sentence (remove original line breaks first)
+    lines = [line.strip() for line in text.split('\n')]
+    text = ' '.join(line for line in lines if line)
+
+    # Add line break after sentence-ending punctuation (Chinese and English punctuation)
+    text = re.sub(r'([.,!?:;，。！？；])\s*', r'\1\n', text)
+
+    # Fix spaces after English apostrophes
+    text = re.sub(r"'\s+", "'", text)
+
+    # Merge redundant spaces
+    text = re.sub(r'[ \t]+', ' ', text)
+
+    # Remove leading and trailing spaces from lines
+    text = '\n'.join(line.strip() for line in text.split('\n'))
+
+    return text.strip()
+
+# ===== Detection Functions =====
+def is_ch_char(char:str):
+    """Determine if a single character is a Chinese character"""
+    if len(char) != 1:
+        return False
+    
+    # Unicode ranges for Chinese characters
+    # 1. Basic Chinese: 0x4E00-0x9FFF
+    # 2. Extension A: 0x3400-0x4DBF
+    # 3. Extension B: 0x20000-0x2A6DF
+    # 4. Extension C: 0x2A700-0x2B73F
+    # 5. Extension D: 0x2B740-0x2B81F
+    # 6. Extension E: 0x2B820-0x2CEAF
+    
+    code = ord(char)
+    
+    # Common check (covers most cases)
+    if 0x4E00 <= code <= 0x9FFF:
+        return True
+    # Extension A
+    if 0x3400 <= code <= 0x4DBF:
+        return True
+    # Other extensions not considered for now
+    
+    return False
+
+# ===== File Operations =====
+
+def load_txt(path:str) -> str:
+    """Load a file as plain text"""
+    with open(path, 'r') as file:
+        content = file.read()
+    return content
+
+def load_json(path:str):
+    """Load a JSON file"""
+    if not os.path.exists(path):
+        return {}
+    with open(path, 'r') as file:
+        data = json.load(file)
+    return data
+
+def load_jsonl(path:str, limit=-1) -> list[dict]:
+    """Load a JSONL file"""
+    data = []
+    with open(path, 'r') as file:
+        for id, line in tqdm(enumerate(file), desc=f"Loading {path}"):
+            if limit != -1 and id == limit:
+                break
+            data.append(json.loads(line))
+    return data
+
+def save_json(data, path:str):
+    """Save a JSON file"""
+    with open(path, 'w', encoding='utf-8') as file:
+        json.dump(data, file, ensure_ascii=False, indent=4)
+
+def save_jsonl(data:list[dict], path:str, mode='w'):
+    """Save a JSONL file"""
+    with open(path, mode, encoding='utf-8') as file:
+        for ele in tqdm(data, desc=f"Saving to {path}"):
+            json.dump(ele, file, ensure_ascii=False)
+            file.write("\n")
+
+def audio_cut(input_path, mode:str, output_dir:str, segment_length:int=30000):
+    """
+    Extract a segment of specified length from an audio file
+    - mode: Cut type (random / middle)
+    - output_dir: Output folder
+    - segment_length: Segment length (milliseconds)
+    """
+    assert mode in ['random', 'middle']
+
+    # Check if file exists
+    if not os.path.exists(input_path):
+        raise FileNotFoundError(f"Audio file not found: {input_path}")
+    
+    # Load audio file
+    audio = AudioSegment.from_file(input_path)
+    audio = audio.set_frame_rate(44100).set_channels(1) # Set sample rate and channels
+    audio_duration = len(audio) # Duration control
+    
+    # If audio length is less than target segment length, use entire audio
+    if audio_duration <= segment_length:
+        print(f"Warning: Audio too short ({audio_duration}ms), using full audio: {input_path}")
+        segment = audio
+    else:
+        # Calculate slice position based on mode
+        if mode == "random":
+            # Random cut
+            max_start = max(0, audio_duration - segment_length)
+            start = random.randint(0, max_start)
+            end = start + segment_length
+        else:
+            # Cut from middle
+            middle_point = audio_duration // 2
+            start = max(0, middle_point - (segment_length // 2))
+            end = min(audio_duration, start + segment_length)
+            
+            # If cutting from middle would exceed boundaries, adjust start position
+            if end > audio_duration:
+                end = audio_duration
+                start = end - segment_length
+            elif start < 0:
+                start = 0
+                end = segment_length
+        
+        # Ensure slice range is valid
+        start = max(0, min(start, audio_duration))
+        end = max(0, min(end, audio_duration))
+        
+        if start >= end:
+            raise ValueError(f"Invalid slice range: start={start}, end={end}, duration={audio_duration}")
+        
+        # Execute slice
+        segment = audio[start:end]
+    
+    # Generate output path
+    basename = pure_name(input_path)
+    output_path = os.path.join(output_dir, f"seg_{basename}.wav")
+    
+    # Save segment
+    segment.export(
+        output_path, 
+        format="wav",
+        codec="pcm_s16le",  # 16-bit little-endian encoding
+        parameters=["-acodec", "pcm_s16le"]  # ffmpeg parameters
+    )
+    return output_path
+
+def format_meta(dir:str, show:bool=True) -> list[dict]:
+    """Recursively get all audio paths (wav / mp3) in a folder and build JSONL"""
+    if not os.path.isdir(dir):
+        return []
+    dataset = []
+    if show:
+        for name in tqdm(os.listdir(dir), desc=f"Formating {dir}"):
+            path = os.path.join(dir, name)
+            if os.path.isdir(path):
+                dataset += format_meta(path, False)
+            elif name.endswith('.mp3') or name.endswith('.wav'):
+                dataset.append({"path": path})
+    else:
+        for name in os.listdir(dir):
+            path = os.path.join(dir, name)
+            if os.path.isdir(path):
+                dataset += format_meta(path, False)
+            elif name.endswith('.mp3') or name.endswith('.wav'):
+                dataset.append({"path": path})
+    return dataset
+
+def dup_remove(raw_data:list[dict], save_path:str, key:str, seg:str):
+    """
+    Remove already generated items from dataset
+    - key is the primary key in raw dataset, foreign key in save
+    - seg is the target field
+    """
+    if not os.path.exists(save_path):
+        print(f"Dup num: 0")
+        return raw_data
+    save_data = load_jsonl(save_path)
+    keys = set()
+    for ele in tqdm(save_data, desc="Constructing Dup Set"):
+        if seg in ele:
+            keys.add(ele[key])
+    rest_data = []
+    dup_count = 0
+    for ele in tqdm(raw_data, desc="Checking Dup"):
+        if ele[key] not in keys:
+            rest_data.append(ele)
+        else:
+            dup_count += 1
+    print(f"Dup num: {dup_count}")
+    return rest_data
+
+def tar_size_check(data_dir:str, subfixes:list[str], per:int, max_size:int):
+    """
+    Determine the number of files that can fit in a block before compression (assuming uniform file sizes)
+    - data_dir: Folder to compress
+    - subfixes: File suffixes to compress (e.g., .mp3)
+    - per: Check every N files on average
+    - max_size: Maximum limit in GB
+    """
+    names = sorted(list(os.listdir(data_dir)))
+    count = 0
+    size_sum = 0
+    for name in tqdm(names, desc="Size Checking"):
+        path = os.path.join(data_dir, name)
+        subfix = os.path.splitext(name)[1]
+        if subfix not in subfixes:
+            continue
+        count += 1
+        size_sum += os.path.getsize(path)
+        if count % per == 0:
+            gb_size = size_sum / 1024 / 1024 / 1024
+            if gb_size > max_size:
+                break
+            print(f"Count: {count}, Size: {gb_size:.2f}GB")
+
+def tar_dir(
+        data_dir:str,
+        subfixes:list[str],
+        save_dir:str,
+        group_size:int,
+        tmp_dir:str,
+        mark:str,
+        max_workers:int=10,
+        ):
+    """Compress files in a directory in chunks (non-recursive)"""
+    names = sorted(list(os.listdir(data_dir)))
+    file_num = len(names)
+    for i in range(0, file_num, group_size):
+        names_subset = names[i:i+group_size]
+        size_sum = 0
+        name_path = os.path.join(tmp_dir, f"name_{i}_{mark}")
+        with open(name_path, 'w', encoding='utf-8') as file:
+            for name in tqdm(names_subset, desc=f"Counting Block {i}"):
+                path = os.path.join(data_dir, name)
+                subfix = os.path.splitext(path)[1]
+                if subfix not in subfixes:
+                    continue
+                file.write("./" + name + "\n")
+                size_sum += os.path.getsize(path)
+        gb_size = size_sum / 1024 / 1024 / 1024
+        print(f"Zipping block {i+1}, size: {gb_size:.2f}GB")
+        
+        tar_cmd = [
+            'tar', 
+            '--no-recursion',
+            '--files-from', str(name_path),
+            '-cf', '-'
+        ]
+        pigz_cmd = ['pigz', '-p', str(max_workers), '-c']
+
+        tar_process = subprocess.Popen(tar_cmd, stdout=subprocess.PIPE, cwd=data_dir)
+        pigz_process = subprocess.Popen(pigz_cmd, stdin=tar_process.stdout, stdout=subprocess.PIPE, cwd=data_dir)
+
+        save_path = os.path.join(save_dir, f"block_{i}_{mark}.tar.gz")
+        with open(save_path, 'wb') as out_file:
+            while True:
+                data = pigz_process.stdout.read(4096)
+                if not data:
+                    break
+                out_file.write(data)
+        
+        tar_process.wait()
+        pigz_process.wait()
+
+        if tar_process.returncode == 0 and pigz_process.returncode == 0:
+            print(f"Compression completed: {save_path}")
+        else:
+            print(f"Compression failed: tar return code={tar_process.returncode}, pigz return code={pigz_process.returncode}")
+
+def music_avg_size(dir:str):
+    """Average music size (MB), length (s)"""
+    dataset = format_meta(dir)
+    dataset = dataset[:50]
+    size_sum = 0
+    length_sum = 0
+    for ele in tqdm(dataset, desc=f"Counting Music Size in {dir}"):
+        path = ele['path']
+        audio = AudioSegment.from_file(path)
+        length_sum += len(audio) / 1000.0
+        size_sum += os.path.getsize(path)
+    size_avg = size_sum / len(dataset) / 1024 / 1024
+    length_avg = length_sum / len(dataset)
+    return size_avg, length_avg
+
+def get_sample(path:str, save_path:str="tmp.jsonl", num:int=100):
+    """Get N records from a JSONL file"""
+    if not os.path.exists(path):
+        return
+    if path.endswith(".jsonl"):
+        dataset = load_jsonl(path)
+    elif path.endswith(".json"):
+        dataset = load_json(path)
+    else:
+        print(f"Unsupport file: {path}")
+        return
+    sub_dataset = random.sample(dataset, num)
+    save_jsonl(sub_dataset, save_path)
+
+def _get_field_one(path:str, field:str):
+    """Process data from one path"""
+    with open(path, 'r') as file:
+        data = json.load(file)
+    new_data = {
+        "id": f"{data['song_id']}_{data['track_index']}",
+        field: data[field]
+    }
+    return new_data
+
+def get_field_suno(dir:str, save_path:str, field:str, max_workers:int=8):
+    """Extract a specific field from scattered JSON files in suno"""
+    paths = []
+    for name in tqdm(os.listdir(dir), desc="Getting names"):
+        if not name.endswith(".json"):
+            continue
+        paths.append(os.path.join(dir, name))
+
+    with ProcessPoolExecutor(max_workers=max_workers) as executor:
+        futures = [executor.submit(_get_field_one, path, field) for path in paths]
+        with open(save_path, 'w', encoding='utf-8') as file:
+            with tqdm(total=len(paths), desc="Processing the JSONs") as pbar:
+                for future in as_completed(futures):
+                    result = future.result()
+                    json.dump(result, file, ensure_ascii=False)
+                    file.write("\n")
+                    pbar.update(1)
+
+def find_json(dir:str) -> list[str]:
+    """Find JSONL / JSON files in a folder"""
+    names = []
+    for name in tqdm(os.listdir(dir), desc="Finding JSON/JSONL"):
+        if name.endswith(".json") or name.endswith(".jsonl"):
+            names.append(name)
+    return names
+
+def show_dir(dir:str):
+    """Display all contents in a directory"""
+    if not os.path.isdir(dir):
+        return
+    for name in os.listdir(dir):
+        print(name)
+
+def _convert_mp3(path:str, dir:str):
+    """Process a single audio file"""
+    purename = pure_name(path)
+    output_path = os.path.join(dir, purename + ".mp3")
+    if os.path.exists(output_path):
+        # Already completed
+        return "pass"
+    try:
+        audio = AudioSegment.from_file(path)
+    except Exception:
+        # Failed to read file
+        print(f"fail to load {path}")
+        return "fail"
+    audio.export(output_path, format='mp3')
+    return "finish"
+
+def convert_mp3(meta_path:str, dir:str, max_workers:int=10):
+    """Convert all specified audio files to mp3 and save in specified directory"""
+    os.makedirs(dir, exist_ok=True)
+    dataset = load_jsonl(meta_path)
+    pass_num = 0
+    finish_num = 0
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = [executor.submit(_convert_mp3, ele['path'], dir) for ele in dataset]
+        with tqdm(total=len(dataset), desc=f"Converting {meta_path}") as pbar:
+            for future in as_completed(futures):
+                res = future.result()
+                if res == "pass":
+                    pass_num += 1
+                else:
+                    finish_num += 1
+                pbar.update(1)
+    print(f"Finish {finish_num}, Pass {pass_num}")
+
+# ===== GPU and Models =====
+
+def get_free_gpu() -> int:
+    """Return the GPU ID with the least memory usage"""
+    cmd = "nvidia-smi --query-gpu=index,memory.free --format=csv,noheader,nounits"
+    result = subprocess.check_output(cmd.split()).decode().strip().split("\n")
+
+    free_list = []
+    for line in result:
+        idx, free_mem = line.split(",")
+        free_list.append((int(idx), int(free_mem)))  # (GPU id, free memory MiB)
+    
+    # Sort by remaining memory
+    free_list.sort(key=lambda x: x[1], reverse=True)
+    return free_list[0][0]
+
+# ===== Data Analysis =====
+
+def compose_analyze(dataset:list[dict]):
+    """Statistical analysis of music structure composition"""
+    # Label count
+    labels = defaultdict(int)
+    for ele in tqdm(dataset):
+        segments = ele['segments']
+        for segment in segments:
+            label = segment['label']
+            labels[label] += 1
+    print(f"Number of labels: {len(labels)}")
+    print(dict_sort_print(labels))
+
+    # Different combinations
+    label_combs = defaultdict(int)
+    for ele in tqdm(dataset):
+        segments = ele['segments']
+        labels = []
+        for segment in segments:
+            label = segment['label']
+            labels.append(label)
+        if len(labels) == 0:
+            continue
+        label_comb = " | ".join(labels)
+        label_combs[label_comb] += 1
+    print(f"Number of combinations: {len(label_combs)}")
+    print(dict_sort_print(label_combs))
+
+def _filter_tag(content:str) -> list[str]:
+    """Split and format tag fields"""
+    tags = []
+    raws = re.split(r'[,，.]', content)
+    for raw in raws:
+        raw = raw.strip().lower()   # Remove spaces and convert to lowercase
+        if raw == "":
+            continue
+        seg_pos = raw.find(":")
+        if seg_pos != -1:
+            # If colon exists, only take the part after it
+            tag = raw[seg_pos+1:].strip()
+        else:
+            tag = raw
+        tags.append(tag)
+    return tags
+
+def tags_analyze(dataset:list[dict]):
+    """Song tag analysis"""
+    tag_count = defaultdict(int)
+    for ele in tqdm(dataset, desc="Tag analyzing"):
+        tags = _filter_tag(ele['style'])
+        for tag in tags:
+            tag_count[tag] += 1
+    print(f"Number of tags: {len(tag_count.keys())}")
+    print(dict_sort_print(tag_count))