diff --git a/.gitignore b/.gitignore index 278eb50ff..542bc5035 100644 --- a/.gitignore +++ b/.gitignore @@ -245,7 +245,6 @@ datasets/ !recipes/cosmos3/negative_prompt.json !tests/dfx/perf/tests/*.json !apps/ComfyUI-vLLM-Omni/example_workflows/*.json -!tests/assets/soulxsinger/*.json *.jsonl *.parquet diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index beb70949d..672a10174 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -76,8 +76,6 @@ th { | `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | — | | `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — | | `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — | -| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — | -| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — | | `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — | | `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) | diff --git a/docs/user_guide/diffusion/cpu_offload.md b/docs/user_guide/diffusion/cpu_offload.md index a17e4a114..93ac5e48e 100644 --- a/docs/user_guide/diffusion/cpu_offload.md +++ b/docs/user_guide/diffusion/cpu_offload.md @@ -49,7 +49,6 @@ higher-priority strategy is selected. | StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | — | yes | — | `transformer_blocks` | | Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` | | Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` | -| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` | yes | yes | — | `layers` | | BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | — | yes | — | `layers`, customized modules | | Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | yes | yes | yes | `layers`, `gen_layers` | diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index dc4f01d57..92e08fab5 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -420,8 +420,7 @@ against transformers 5.x, so nothing is lost by holding it there. (The same clash is why this is documented here rather than declared in `pyproject.toml`: an extra resolves together with the base dependencies, so -`vllm-omni[gepard]` would be unsatisfiable rather than opt-in. SoulX-Singer English -SVS documents `nemo_toolkit[asr]` the same way.) +`vllm-omni[gepard]` would be unsatisfiable rather than opt-in.) On a host whose CUDA toolkit cannot build kernels for the GPU — no `nvcc`/`ninja`, or a consumer Blackwell (`sm_120`) card — also set `VLLM_USE_FLASHINFER_SAMPLER=0`. vLLM's @@ -617,100 +616,3 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T - `--num-prompts N` replicates the prompt for performance measurement. - `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`. - Run `--help` for the full argument surface. - ---- - -## SoulX-Singer - -Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Script: `soulxsinger/end2end.py`. Deploy: `vllm_omni/deploy/soulxsinger_svs.yaml` or `soulxsinger_svc.yaml`. - -### Prerequisites - -Download DiT and preprocess weights, then set up separate SVS / SVC view directories. Copy `soulxsinger/utils/phoneme/phone_set.json` from upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) into the model weights dir as `phoneme/phone_set.json` — HuggingFace does not ship it. - -```bash -# 1. DiT weights -export BASE=path/to/SoulX-Singer -export PREPROCESS=path/to/SoulX-Singer-Preprocess -export SVC_DIR=path/to/SoulX-Singer-svc - -huggingface-cli download Soul-AILab/SoulX-Singer --local-dir "$BASE" - -# 2. Preprocess weights (required) -huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess --local-dir "$PREPROCESS" -export SOULX_PREPROCESS_WEIGHTS_DIR="$PREPROCESS" - -# 3. SVS / SVC view directories -mkdir -p "$SVC_DIR" -cp $BASE/{config.yaml,README.md,assets} $SVC_DIR -mv $BASE/model-svc.pt $SVC_DIR/model-svc.pt - -cat > "$BASE/config.json" <<'EOF' -{ - "model_type": "soulxsinger", - "architectures": ["SoulXSingerPipeline"], - "max_num_seqs": 1 -} -EOF - -cat > "$SVC_DIR/config.json" <<'EOF' -{ - "model_type": "soulxsinger", - "architectures": ["SoulXSingerSVCPipeline"], - "max_num_seqs": 1 -} -EOF -``` - -`config.yaml` hyper-parameters live under `$BASE`; each view's `config.json` `architectures` field is the single source of truth for SVS vs SVC. Point `--model` at the matching directory (`$BASE` for SVS, `$SVC_DIR` for SVC). Deploy YAML is chosen automatically from `config.json`; optional `--svs` / `--svc` only assert the mode matches. - -**Online preprocess** is the default: pass `--prompt-audio` and `--target-audio`, and the worker runs vocal separation, F0, and (for SVS) lyrics/MIDI before DiT. Install only what your run needs: - -```bash -pip install "BS-RoFormer" # vocal sep + F0 on GPU — SVS and SVC -``` - -Mandarin SVS also needs FunASR and Chinese G2P; `ffmpeg` must be on `PATH`: - -```bash -# install optional dependencies: -pip install -e ".[soulx-svs]" -``` - -English SVS adds NeMo ASR and NLTK data; pass `--language English`: - -```bash -pip install "nemo_toolkit[asr]==2.6.1" lhotse==1.32.2 -python -c "import nltk; nltk.download('cmudict'); nltk.download('averaged_perceptron_tagger_eng')" -``` - -**Precomputed metadata** is the alternative: pass both `--prompt-metadata-path` and `--target-metadata-path` and skip online ASR/ROSVOT — none of the packages above are required. JSON can be produced by integrated preprocess on a prior run, or by upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) `preprocess/` scripts if you prefer to run that outside vLLM-Omni. - -### Quick start - -```bash -# SVS — default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3 -python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \ - --model "$BASE" \ - --preprocess-weights-dir "$PREPROCESS" \ - --control score \ - --num-inference-steps 32 \ - -o output.wav - -python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \ - --model "$SVC_DIR" \ - --preprocess-weights-dir "$PREPROCESS" \ - --svc \ - --num-inference-steps 32 \ - -o output_svc.wav -``` - -`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. Long SVS targets are handled in one request. See `end2end.py --help` for `--pitch-shift`, `--vocal-sep`, `--auto-shift`, and language/control options. - -### Notes - -- Output: 24 kHz mono WAV; batch only. -- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on. -- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only. - ---- diff --git a/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py b/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py deleted file mode 100644 index 14dd5e145..000000000 --- a/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py +++ /dev/null @@ -1,162 +0,0 @@ -"""Repeatable SoulX-Singer offline benchmark (preprocess + DiT, RTF + stage timings). - -Usage: - python benchmark.py --model /path/to/SoulX-Singer --svs \\ - --prompt-audio ../SoulX-Singer/example/audio/zh_prompt.mp3 \\ - --target-audio ../SoulX-Singer/example/audio/music.mp3 \\ - --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\ - -o benchmark.wav - - python benchmark.py ... --enable-diffusion-pipeline-profiler -""" - -import argparse -import statistics -import time -from pathlib import Path - -import soundfile as sf -from end2end import ( - SVC_DEPLOY_CONFIG, - SVS_DEPLOY_CONFIG, - add_inference_args, - build_sampling, - extract_audio, - resolve_preprocess_weights_dir, -) - -from vllm_omni.entrypoints.omni import Omni - - -def _audio_duration_sec(outputs) -> float: - audio_np, sr = extract_audio(outputs) - return audio_np.size / sr - - -def _stage_durations_from_output(outputs) -> dict[str, float]: - omni_out = outputs[0] - durations = getattr(omni_out, "stage_durations", None) or {} - if durations: - return dict(durations) - ro = omni_out - if ro is not None: - return dict(getattr(ro, "stage_durations", None) or {}) - return {} - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="SoulX-Singer offline benchmark") - add_inference_args(parser) - parser.add_argument("--warmup", type=int, default=1, help="Warmup runs (excluded from stats)") - parser.add_argument("--runs", type=int, default=3, help="Measured runs") - parser.add_argument( - "--enable-diffusion-pipeline-profiler", - action="store_true", - help="Enable low-overhead stage timing in DiffusionOutput.stage_durations", - ) - parser.add_argument( - "--enforce-eager", - action="store_true", - help="Disable torch.compile on diff_estimator (eager baseline for A/B)", - ) - parser.add_argument( - "--dtype", - type=str, - default=None, - choices=["bfloat16", "bf16", "float16", "fp16", "half"], - help="DiT trunk dtype (default from deploy YAML). Use float16 to match upstream --fp16.", - ) - parser.add_argument( - "--output", - "-o", - type=str, - default=None, - help="Save WAV from the last measured run (after warmup)", - ) - parser.add_argument( - "--save-each-run", - action="store_true", - help="With --output, also write measured runs as _run1.wav, _run2.wav, ...", - ) - return parser.parse_args() - - -def main() -> None: - args = parse_args() - preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir) - deploy_config = args.deploy_config or str(SVS_DEPLOY_CONFIG if args.svs else SVC_DEPLOY_CONFIG) - mode = "SVS" if args.svs else "SVC" - - omni_kwargs: dict = { - "model": args.model, - "deploy_config": deploy_config, - "async_chunk": False, # SoulX-Singer currently supports only batch mode (pseudo-streaming was stashed) - } - if args.enable_diffusion_pipeline_profiler: - omni_kwargs["enable_diffusion_pipeline_profiler"] = True - if args.enforce_eager: - omni_kwargs["enforce_eager"] = True - if args.dtype is not None: - omni_kwargs["dtype"] = args.dtype - - compile_mode = "eager" if args.enforce_eager else "torch.compile" - dtype_label = args.dtype or "from deploy YAML" - print(f"Loading SoulX-Singer {mode} from {args.model} [{compile_mode}, dtype={dtype_label}]") - - omni = Omni(**omni_kwargs) - kind = "svs" if args.svs else "svc" - sampling = build_sampling(args, preprocess_weights_dir=preprocess_dir, kind=kind) - prompt = {"prompt_token_ids": [0]} - - latencies_ms: list[float] = [] - rtfs: list[float] = [] - last_stages: dict[str, float] = {} - last_measured_outputs = None - - total_runs = args.warmup + args.runs - for run_idx in range(total_runs): - is_warmup = run_idx < args.warmup - label = "warmup" if is_warmup else f"run {run_idx - args.warmup + 1}/{args.runs}" - t0 = time.perf_counter() - outputs = list(omni.generate([prompt], sampling)) - elapsed_ms = (time.perf_counter() - t0) * 1000.0 - audio_sec = _audio_duration_sec(outputs) - rtf = (elapsed_ms / 1000.0) / audio_sec if audio_sec > 0 else float("inf") - print(f"[{label}] client={elapsed_ms:.1f} ms, audio={audio_sec:.2f}s, RTF={rtf:.3f}") - if not is_warmup: - latencies_ms.append(elapsed_ms) - rtfs.append(rtf) - last_stages = _stage_durations_from_output(outputs) - last_measured_outputs = outputs - if args.output and args.save_each_run: - measured_idx = run_idx - args.warmup + 1 - out_path = Path(args.output) - run_path = out_path.with_name(f"{out_path.stem}_run{measured_idx}{out_path.suffix}") - audio_np, sr = extract_audio(outputs) - sf.write(str(run_path), audio_np, sr) - print(f" saved {run_path} ({sr} Hz, {audio_np.size / sr:.2f}s)") - - omni.close() - - if args.output and last_measured_outputs is not None: - audio_np, sr = extract_audio(last_measured_outputs) - out_path = Path(args.output) - out_path.parent.mkdir(parents=True, exist_ok=True) - sf.write(str(out_path), audio_np, sr) - print(f"\nSaved last measured run → {out_path} ({sr} Hz, {audio_np.size / sr:.2f}s)") - - if latencies_ms: - print("\n=== Summary (measured runs) ===") - print(f"client_ms: mean={statistics.mean(latencies_ms):.1f}, stdev={statistics.pstdev(latencies_ms):.1f}") - print(f"RTF: mean={statistics.mean(rtfs):.3f}, stdev={statistics.pstdev(rtfs):.3f}") - if last_stages: - print("\nStage durations (last measured run):") - for name, value in sorted(last_stages.items()): - if name.endswith("_ms"): - print(f" {name}: {value:.1f} ms") - else: - print(f" {name}: {value * 1000.0:.1f} ms") - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/text_to_speech/soulxsinger/end2end.py b/examples/offline_inference/text_to_speech/soulxsinger/end2end.py deleted file mode 100644 index a1bb112af..000000000 --- a/examples/offline_inference/text_to_speech/soulxsinger/end2end.py +++ /dev/null @@ -1,285 +0,0 @@ -"""Offline SoulX-Singer SVS / SVC: single-stage DiT (preprocess inline). - -Pass raw prompt/target audio; vLLM-Omni runs integrated preprocess then inference. - -Usage: - # SVS (default) - python end2end.py --model /path/to/SoulX-Singer \\ - --prompt-audio path/to/prompt/audio.mp3 \\ - --target-audio path/to/target/audio.mp3 \\ - --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\ - -o output.wav - - # SVC (config.json must declare SoulXSingerSVCPipeline; --model points at SVC view) - python end2end.py --model /path/to/SoulX-Singer-svc \\ - --prompt-audio path/to/prompt/audio.mp3 \\ - --target-audio path/to/target/audio.mp3 \\ - --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\ - -o output_svc.wav -""" - -import argparse -import os -from pathlib import Path - -import numpy as np -import soundfile as sf - -from vllm_omni.diffusion.models.soulx_singer.utils import resolve_soulx_kind, validate_soulx_extra_args -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - -REPO_ROOT = Path(__file__).resolve().parents[4] -SVS_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svs.yaml" -SVC_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svc.yaml" -DEFAULT_PROMPT_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "zh_prompt.mp3" -DEFAULT_TARGET_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "music.mp3" -_SAMPLE_RATE = 24000 - - -def _require_paths(paths: dict[str, str | None]) -> None: - missing = [name for name, p in paths.items() if not p or not os.path.isfile(p)] - if missing: - raise FileNotFoundError("Missing file(s): " + ", ".join(f"{k}={paths[k]!r}" for k in missing)) - - -def resolve_preprocess_weights_dir(explicit: str | None) -> str: - if explicit: - path = Path(explicit).expanduser().resolve() - if not (path / "rmvpe" / "rmvpe.pt").is_file(): - raise FileNotFoundError(f"--preprocess-weights-dir must contain rmvpe/rmvpe.pt: {path}") - return str(path) - env = os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR") - if env: - path = Path(env).expanduser().resolve() - if (path / "rmvpe" / "rmvpe.pt").is_file(): - return str(path) - raise FileNotFoundError( - "SoulX preprocess weights not found. Pass --preprocess-weights-dir or set " - "SOULX_PREPROCESS_WEIGHTS_DIR to a directory containing rmvpe/rmvpe.pt " - "(huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess)." - ) - - -def build_sampling( - args: argparse.Namespace, - *, - preprocess_weights_dir: str, - kind: str, -) -> OmniDiffusionSamplingParams: - extra: dict = { - "prompt_audio": os.path.abspath(args.prompt_audio), - "target_audio": os.path.abspath(args.target_audio), - "preprocess_weights_dir": preprocess_weights_dir, - "vocal_sep": args.vocal_sep, - "auto_shift": args.auto_shift, - "pitch_shift": args.pitch_shift, - } - if kind == "svs": - extra["language"] = args.language - extra["control"] = args.control - if args.prompt_metadata_path: - extra["prompt_metadata_path"] = os.path.abspath(args.prompt_metadata_path) - if args.target_metadata_path: - extra["target_metadata_path"] = os.path.abspath(args.target_metadata_path) - if args.prompt_metadata_path or args.target_metadata_path: - extra["audio_path"] = os.path.abspath(args.prompt_wav_path or args.prompt_audio) - - extra = validate_soulx_extra_args(kind, extra) - - return OmniDiffusionSamplingParams( - num_inference_steps=args.num_inference_steps, - guidance_scale=args.guidance_scale, - seed=args.seed, - extra_args=extra, - ) - - -def extract_audio(outputs) -> tuple[np.ndarray, int]: - out = outputs[0] - mm = getattr(out, "multimodal_output", None) - if mm is None: - ro = out - if ro is not None: - mm = getattr(ro, "multimodal_output", None) - if mm is None and ro.outputs: - mm = getattr(ro.outputs[0], "multimodal_output", None) - if not mm or "audio" not in mm: - raise RuntimeError("No audio in multimodal_output") - - audio_val = mm["audio"] - if isinstance(audio_val, list): - chunks = [ - np.asarray(chunk.detach().cpu().float().numpy() if hasattr(chunk, "detach") else chunk).reshape(-1) - for chunk in audio_val - if chunk is not None - ] - audio_np = np.concatenate(chunks, axis=0) if chunks else np.array([], dtype=np.float32) - elif hasattr(audio_val, "cpu"): - audio_np = audio_val.detach().cpu().numpy().squeeze() - else: - audio_np = np.asarray(audio_val, dtype=np.float32).squeeze() - - sr_val = mm.get("sr", mm.get("audio_sample_rate", _SAMPLE_RATE)) - if isinstance(sr_val, list) and sr_val: - sr_val = sr_val[-1] - if hasattr(sr_val, "item"): - sr_val = sr_val.item() - return audio_np.astype(np.float32, copy=False), int(sr_val) - - -def add_inference_args(parser: argparse.ArgumentParser) -> None: - mode = parser.add_mutually_exclusive_group() - mode.add_argument( - "--svc", - action="store_true", - help="Assert SVC mode (must match config.json architectures)", - ) - parser.add_argument( - "--model", - type=str, - required=True, - help="SoulX-Singer weight directory ($BASE for SVS, $SVC_DIR for SVC)", - ) - parser.add_argument( - "--deploy-config", - type=str, - default=None, - help="Override deploy YAML (default: soulxsinger_svs.yaml or soulxsinger_svc.yaml)", - ) - parser.add_argument( - "--prompt-audio", - type=str, - default=str(DEFAULT_PROMPT_AUDIO), - help="Prompt / reference audio (mp3 or wav)", - ) - parser.add_argument( - "--target-audio", - type=str, - default=str(DEFAULT_TARGET_AUDIO), - help="Target content audio (mp3 or wav)", - ) - parser.add_argument( - "--prompt-metadata-path", - type=str, - default=None, - help="SVS precomputed prompt metadata.json (skips online ASR/ROSVOT for prompt)", - ) - parser.add_argument( - "--target-metadata-path", - type=str, - default=None, - help="SVS precomputed target metadata.json (skips online ASR/ROSVOT for target)", - ) - parser.add_argument( - "--prompt-wav-path", - type=str, - default=None, - help="SVS prompt vocal wav for metadata processor when using --prompt-metadata-path (default: --prompt-audio)", - ) - parser.add_argument( - "--preprocess-weights-dir", - type=str, - default=None, - help="SoulX-Singer-Preprocess weights (or set SOULX_PREPROCESS_WEIGHTS_DIR)", - ) - parser.add_argument( - "--language", - type=str, - default="Mandarin", - help="SVS lyrics language (Mandarin, Cantonese, English)", - ) - parser.add_argument( - "--control", - type=str, - default="melody", - choices=["score", "melody"], - help="SVS control mode", - ) - parser.add_argument( - "--vocal-sep", - action=argparse.BooleanOptionalAction, - default=False, - help="Run vocal separation during preprocess", - ) - parser.add_argument( - "--auto-shift", - action=argparse.BooleanOptionalAction, - default=True, - help="Auto pitch shift to align target pitch with prompt", - ) - parser.add_argument("--pitch-shift", type=int, default=0) - parser.add_argument("--num-inference-steps", type=int, default=32) - parser.add_argument("--guidance-scale", type=float, default=3.0) - parser.add_argument( - "--seed", - type=int, - default=42, - help="Optional CFM RNG seed. Omit for non-deterministic sampling.", - ) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="SoulX-Singer offline SVS / SVC (preprocess + DiT)") - add_inference_args(parser) - parser.add_argument("--output", "-o", type=str, default="output.wav") - args = parser.parse_args() - return args - - -def main() -> None: - args = parse_args() - kind = resolve_soulx_kind(args.model) - - if not (args.svc ^ (kind == "svs")): - raise ValueError(f"mode mismatch: {args.model}/config.json (declares {kind!r}), --svc={args.svc}") - - required = { - "prompt_audio": args.prompt_audio, - "target_audio": args.target_audio, - } - if kind == "svs" and bool(args.prompt_metadata_path) ^ bool(args.target_metadata_path): - raise ValueError( - "SVS precomputed metadata requires both --prompt-metadata-path and " - "--target-metadata-path (or neither, for online preprocess)." - ) - if args.prompt_metadata_path: - required["prompt_metadata_path"] = args.prompt_metadata_path - required["target_metadata_path"] = args.target_metadata_path - if args.prompt_wav_path: - required["prompt_wav_path"] = args.prompt_wav_path - _require_paths(required) - preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir) - deploy_config = args.deploy_config or str(SVC_DEPLOY_CONFIG if kind == "svc" else SVS_DEPLOY_CONFIG) - mode = kind.upper() - - print(f"Loading SoulX-Singer {mode} from {args.model}") - print(f" deploy: {deploy_config}") - print(f" prompt: {args.prompt_audio}") - print(f" target: {args.target_audio}") - - omni = Omni(model=args.model, deploy_config=deploy_config, async_chunk=False) # SoulX-Singer is batch-only - sampling = build_sampling( - args, - preprocess_weights_dir=preprocess_dir, - kind=kind, - ) - prompt = {"prompt_token_ids": [0], "multi_modal_data": {}} - - print("Running generate (inline preprocess + DiT)...") - outputs = list(omni.generate([prompt], sampling)) - if not outputs: - raise RuntimeError("No output from omni.generate") - - err = getattr(outputs[0], "error", None) - if err is not None: - raise RuntimeError(f"Pipeline error: {err}") - - audio_np, sr = extract_audio(outputs) - sf.write(args.output, audio_np, sr) - print(f"Saved {args.output} ({sr} Hz, {len(audio_np) / sr:.2f}s)") - omni.close() - - -if __name__ == "__main__": - main() diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 4127eb6ef..69a50540b 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -25,7 +25,6 @@ For the full list of supported architectures across all modalities, see | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) | | VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ | | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ | -| SoulX-Singer | `Soul-AILab/SoulX-Singer` | ✓ (prompt audio) | — (batch only) | — (prompt + target audio) | — (chat client) | CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](../../offline_inference/text_to_speech/README.md#cosyvoice3) instead. @@ -715,65 +714,3 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts - Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`). - Voice cloning is gated upstream and may require a recent `mistral_common`. - A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now. - ---- - -## SoulX-Singer - -Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Single-stage DiT with inline preprocess. Uses the `/v1/chat/completions` endpoint with multimodal input (`prompt_audio` + `target_audio`). - -### Prerequisites - -Download DiT and preprocess weights, then set up separate SVS / SVC view directories and install dependencies as described in the [offline README](../../offline_inference/text_to_speech/README.md#soulx-singer). `config.json` `architectures` field is the single source of truth for SVS vs SVC — point `MODEL` at the matching directory. - -### Launch - -```bash -# SVS (default) -export MODEL=/path/to/SoulX-Singer -export PREPROCESS=/path/to/SoulX-Singer-Preprocess -bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh - -# SVC -export MODE=svc -export MODEL=/path/to/SoulX-Singer-svc -bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh -``` - -Or equivalently, set `SOULX_PREPROCESS_WEIGHTS_DIR` and launch directly: -```bash -export SOULX_PREPROCESS_WEIGHTS_DIR=$PREPROCESS -vllm serve $MODEL --omni \ - --deploy-config vllm_omni/deploy/soulxsinger_${MODE}.yaml \ - --port 8192 --trust-remote-code --enforce-eager -``` - -### Sending requests - -Audio paths must be reachable from the server host (local filesystem or data URL). The client sends prompt vocal via `input_audio` and target accompaniment via `extra_args['target_audio']`. - -```bash -# Default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3 -python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \ - --prompt-audio /path/on/server/zh_prompt.mp3 \ - --target-audio /path/on/server/music.mp3 \ - --preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \ - -o output.wav -``` - -Use precomputed metadata to skip online preprocess with following command: -```bash -python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \ - --prompt-metadata-path /path/on/server/zh_prompt.json \ - --target-metadata-path /path/on/server/music.json \ - --audio-path /path/on/server/zh_prompt.mp3 \ - -o output.wav -``` - -`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. See `openai_chat_client.py --help` for `--vocal-sep`, `--language`, `--num-inference-steps`, `--guidance-scale`, and `--seed`. - -### Notes - -- Output: 24 kHz mono WAV; batch only. -- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on. -- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only. diff --git a/examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py b/examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py deleted file mode 100644 index 91035e625..000000000 --- a/examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py +++ /dev/null @@ -1,186 +0,0 @@ -#!/usr/bin/env python3 -"""SoulX-Singer OpenAI-compatible chat client (SVS / SVC). - -Sends prompt audio via ``input_audio`` and target accompaniment via -``extra_args['target_audio']`` (server-local path). For integrated preprocess, -also pass ``preprocess_weights_dir`` in ``extra_args``. - -Usage: - python openai_chat_client.py \\ - --prompt-audio /path/on/server/zh_prompt.mp3 \\ - --target-audio /path/on/server/music.mp3 \\ - --preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \\ - -o output.wav -""" - -from __future__ import annotations - -import argparse -import base64 -import io -import sys -from pathlib import Path - -import requests -import soundfile -import torch - - -def _audio_to_data_url(path: Path) -> str: - with path.open("rb") as handle: - data = base64.b64encode(handle.read()).decode("ascii") - return f"data:audio/mpeg;base64,{data}" - - -def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None: - audio = audio.to(torch.float32) - peak = audio.abs().max().clamp(min=1e-8) - audio = audio / peak - path.parent.mkdir(parents=True, exist_ok=True) - soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16") - - -def _decode_audio_from_response(body: dict) -> tuple[torch.Tensor, int]: - for choice in body.get("choices", []): - audio_obj = choice.get("message", {}).get("audio") - if isinstance(audio_obj, dict) and audio_obj.get("data"): - data, sr = soundfile.read( - io.BytesIO(base64.b64decode(audio_obj["data"])), - dtype="float32", - always_2d=True, - ) - return torch.from_numpy(data).transpose(0, 1), sr - brief = {k: v for k, v in body.items() if k != "choices"} - raise RuntimeError(f"no audio in response message.audio: {brief}") - - -def main() -> int: - repo_root = Path(__file__).resolve().parents[4] - default_assets = repo_root / "tests" / "assets" / "soulxsinger" - - parser = argparse.ArgumentParser(description="SoulX-Singer online chat client") - parser.add_argument("--port", type=int, default=8192) - parser.add_argument("--model", default="Soul-AILab/SoulX-Singer") - parser.add_argument( - "--prompt-audio", - default=str(default_assets / "zh_prompt.mp3"), - help="Prompt vocal audio (path on server if using extra_args, or local for input_audio)", - ) - parser.add_argument( - "--target-audio", - default=str(default_assets / "music.mp3"), - help="Target accompaniment path on the server (extra_args['target_audio'])", - ) - parser.add_argument( - "--prompt-metadata-path", - default=None, - help="SVS precomputed prompt metadata.json", - ) - parser.add_argument( - "--target-metadata-path", - default=None, - help="SVS precomputed target metadata.json", - ) - parser.add_argument( - "--audio-path", - default=None, - help="SVS prompt vocal wav for precomputed metadata", - ) - parser.add_argument("--preprocess-weights-dir", default=None) - parser.add_argument("--output", "-o", default="soulxsinger_out.wav") - parser.add_argument("--svc", action="store_true", help="Use SVC mode knobs") - parser.add_argument("--language", default="Mandarin") - parser.add_argument("--num-inference-steps", type=int, default=32) - parser.add_argument("--guidance-scale", type=float, default=3.0) - parser.add_argument( - "--seed", - type=int, - default=42, - help="Optional CFM seed. Omit for non-deterministic sampling.", - ) - parser.add_argument( - "--auto-shift", - action=argparse.BooleanOptionalAction, - default=True, - help="Auto pitch shift (default: on, original upstream infer.sh)", - ) - parser.add_argument( - "--control", - default="melody", - choices=["melody", "score"], - help="SVS control mode", - ) - parser.add_argument("--vocal-sep", action="store_true") - args = parser.parse_args() - - meta_paths = (args.prompt_metadata_path, args.target_metadata_path, args.audio_path) - if any(meta_paths) and not all(meta_paths): - print( - "ERROR: precomputed metadata requires --prompt-metadata-path, " - "--target-metadata-path, and --audio-path together.", - file=sys.stderr, - ) - return 2 - - extra_args: dict = { - "vocal_sep": args.vocal_sep, - "auto_shift": args.auto_shift, - "pitch_shift": 0, - } - if all(meta_paths): - extra_args.update( - { - "prompt_metadata_path": str(Path(args.prompt_metadata_path).expanduser().resolve()), - "target_metadata_path": str(Path(args.target_metadata_path).expanduser().resolve()), - "audio_path": str(Path(args.audio_path).expanduser().resolve()), - } - ) - content = [{"type": "text", "text": "soulx-singer"}] - else: - prompt_path = Path(args.prompt_audio).expanduser().resolve() - if not prompt_path.is_file(): - print(f"ERROR: prompt audio not found: {prompt_path}", file=sys.stderr) - return 2 - extra_args["prompt_audio"] = str(prompt_path) - extra_args["target_audio"] = str(Path(args.target_audio).expanduser().resolve()) - if args.preprocess_weights_dir: - extra_args["preprocess_weights_dir"] = str(Path(args.preprocess_weights_dir).expanduser().resolve()) - content = [ - {"type": "text", "text": "soulx-singer"}, - { - "type": "input_audio", - "input_audio": {"data": _audio_to_data_url(prompt_path), "format": "mp3"}, - }, - ] - if not args.svc: - extra_args["language"] = args.language - extra_args["control"] = args.control - - payload = { - "model": args.model, - "modalities": ["audio"], - "messages": [{"role": "user", "content": content}], - "num_inference_steps": args.num_inference_steps, - "guidance_scale": args.guidance_scale, - "extra_args": extra_args, - } - if args.seed is not None: - payload["seed"] = args.seed - - print(f"POST http://localhost:{args.port}/v1/chat/completions") - response = requests.post( - f"http://localhost:{args.port}/v1/chat/completions", - headers={"Content-Type": "application/json"}, - json=payload, - timeout=1800, - ) - response.raise_for_status() - audio, sample_rate = _decode_audio_from_response(response.json()) - _save_wav(audio, Path(args.output), sample_rate) - duration = audio.shape[-1] / sample_rate - print(f"saved {args.output} sr={sample_rate}Hz duration={duration:.2f}s") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/examples/online_serving/text_to_speech/soulxsinger/run_server.sh b/examples/online_serving/text_to_speech/soulxsinger/run_server.sh deleted file mode 100755 index 3e20de39b..000000000 --- a/examples/online_serving/text_to_speech/soulxsinger/run_server.sh +++ /dev/null @@ -1,39 +0,0 @@ -#!/bin/bash -# Launch vLLM-Omni server for SoulX-Singer (single-stage DiT, preprocess inline). -# -# Usage: -# MODEL=/path/to/SoulX-Singer PREPROCESS=/path/to/Preprocess \ -# -# Audio paths in client extra_args must be readable on the server host. - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)" - -MODEL="${MODEL:-Soul-AILab/SoulX-Singer}" -MODE="${MODE:-svs}" -PORT="${PORT:-8192}" -GPUS="${GPUS:-0}" - -if [[ "$MODE" == "svc" ]]; then - DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svc.yaml}" -else - DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svs.yaml}" -fi - -echo "Starting SoulX-Singer server" -echo " MODEL=$MODEL" -echo " MODE=$MODE" -echo " PORT=$PORT" -echo " DEPLOY_CONFIG=$DEPLOY_CONFIG" -echo " CUDA_VISIBLE_DEVICES=$GPUS" - -CUDA_VISIBLE_DEVICES="$GPUS" \ -vllm serve "$MODEL" \ - --omni \ - --deploy-config "$DEPLOY_CONFIG" \ - --host 0.0.0.0 \ - --port "$PORT" \ - --trust-remote-code \ - --enforce-eager diff --git a/pyproject.toml b/pyproject.toml index a76331799..84a950719 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -135,13 +135,6 @@ docs = [ "pydantic", ] -soulx-svs = [ - "funasr>=1.0.0", - "g2pM>=0.1.2.5", - "g2p-en>=2.1.0", - "ToJyutping>=3.2.0", -] - [project.urls] Homepage = "https://github.com/vllm-project/vllm-omni" Repository = "https://github.com/vllm-project/vllm-omni" diff --git a/tests/assets/soulxsinger/music.json b/tests/assets/soulxsinger/music.json deleted file mode 100644 index d9e0d8882..000000000 --- a/tests/assets/soulxsinger/music.json +++ /dev/null @@ -1,16 +0,0 @@ -[ - { - "index": "vocal_240_51240", - "language": "Mandarin", - "time": [ - 240, - 51240 - ], - "duration": "0.21 0.18 0.34 0.38 0.26 0.22 0.50 0.20 0.33 0.13 0.60 0.22 0.38 0.52 0.30 0.12 1.82 0.98 0.20 0.26 0.38 0.36 0.58 0.54 0.26 0.50 1.46 3.03 0.24 0.24 0.29 0.25 0.20 0.14 0.80 0.22 0.54 0.30 0.24 0.16 0.58 0.28 0.38 1.73 0.79 0.24 0.30 0.34 0.30 0.30 0.34 0.21 0.21 0.36 0.34 0.23 1.85 1.90 0.23 0.39 0.68 0.50 0.31 0.43 0.76 0.38 2.00 1.87 0.68 0.72 0.56 0.62 0.80 0.40 0.42 1.68 1.79 0.70 0.66 0.54 0.24 0.48 0.68 0.40 2.34 0.14", - "text": " 只 是 因 为 为 在 人 群 群 中 多 看 了 你 一 眼 再 也 没 能 忘 掉 你 容 颜 梦 想 着 着 偶 偶 然 然 有 一 一 天 再 相 见 从 此 我 开 始 始 孤 孤 单 思 念 念 想 想 你 时 你 你 在 天 边 想 你 时 你 在 眼 前 前 想 你 时 你 你 在 脑 海 ", - "phoneme": " zh_zhi3 zh_shi4 zh_yin1 zh_wei4 zh_wei2 zh_zai4 zh_ren2 zh_qun2 zh_qun2 zh_zhong1 zh_duo1 zh_kan4 zh_le5 zh_ni3 zh_yi1 zh_yan3 zh_zai4 zh_ye3 zh_mei2 zh_neng2 zh_wang4 zh_diao4 zh_ni3 zh_rong2 zh_yan2 zh_meng4 zh_xiang3 zh_zhe5 zh_zhe5 zh_ou3 zh_ou3 zh_ran2 zh_ran2 zh_you3 zh_yi1 zh_yi1 zh_tian1 zh_zai4 zh_xiang1 zh_jian4 zh_cong2 zh_ci3 zh_wo3 zh_kai1 zh_shi3 zh_shi3 zh_gu1 zh_gu1 zh_dan1 zh_si1 zh_nian4 zh_nian4 zh_xiang3 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_tian1 zh_bian1 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_zai4 zh_yan3 zh_qian2 zh_qian2 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_nao3 zh_hai3 ", - "note_pitch": "0 64 64 64 66 68 66 66 66 64 64 64 64 66 64 60 61 0 63 63 63 64 66 63 61 59 56 0 68 66 66 68 68 66 66 64 64 0 64 66 61 61 61 64 0 62 63 63 64 64 66 65 66 61 58 59 56 0 69 71 66 68 68 71 66 64 61 0 66 61 68 66 64 63 61 59 0 71 66 68 68 71 66 64 61 0", - "note_type": "1 2 2 2 2 3 2 2 2 3 2 2 2 2 2 2 2 1 2 2 2 2 2 2 2 2 2 1 2 2 2 3 2 3 2 3 2 1 2 3 2 2 2 2 1 2 2 2 2 2 3 2 3 2 2 2 3 1 2 3 2 2 2 3 2 2 2 1 2 2 2 2 2 2 2 3 1 2 2 2 2 3 2 2 2 1", - "f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 318.9 332.2 331.2 323.9 312.3 0.0 0.0 0.0 0.0 0.0 0.0 340.9 342.3 340.2 337.4 333.5 329.3 327.8 328.3 329.1 330.5 331.9 331.5 329.2 327.6 327.4 327.7 329.6 330.5 329.5 328.1 328.5 330.1 330.4 329.5 328.9 331.0 333.6 332.8 331.7 330.2 330.6 330.9 329.5 327.6 332.4 343.5 363.9 368.7 369.3 368.2 366.7 371.3 382.0 392.9 402.3 409.8 413.8 416.3 416.4 415.4 414.1 413.4 415.3 417.4 418.3 413.6 416.0 0.0 0.0 0.0 0.0 368.2 362.5 361.3 362.0 362.9 364.3 366.4 367.5 368.4 368.9 368.1 367.6 369.7 371.1 371.1 369.5 368.0 367.3 366.6 367.1 369.5 372.0 371.3 368.2 368.1 369.6 369.9 373.4 376.5 360.2 285.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 384.0 376.7 372.0 371.7 373.9 376.3 375.5 361.2 333.0 317.8 319.0 329.1 333.5 333.2 333.2 333.3 328.0 314.5 0.0 0.0 0.0 320.6 326.5 332.9 335.7 334.1 329.9 327.4 326.3 326.6 329.1 330.5 330.0 328.0 326.7 328.6 331.1 332.4 332.1 332.8 333.2 333.1 331.8 328.6 318.9 293.2 323.6 327.5 330.0 332.8 331.1 319.4 272.2 0.0 0.0 0.0 0.0 278.5 294.5 310.0 318.5 322.7 327.4 334.1 337.2 336.3 328.5 321.0 324.2 341.5 362.2 374.7 375.8 370.5 366.0 365.2 368.7 371.7 374.4 373.0 368.7 370.4 374.8 375.1 372.5 368.7 363.0 357.1 359.0 366.8 377.1 379.5 371.5 359.9 351.0 358.4 371.7 377.9 375.8 367.8 359.7 363.6 375.3 379.9 377.5 372.9 359.8 345.2 337.1 334.5 333.4 332.8 329.6 319.2 292.8 261.7 253.4 262.8 273.0 278.1 279.3 278.6 277.6 277.9 278.3 278.0 277.3 276.7 275.6 275.3 276.1 277.8 278.2 277.7 277.6 277.2 276.4 275.5 273.9 271.8 270.7 274.7 282.7 285.5 281.4 273.1 264.4 262.2 268.9 278.4 284.3 283.6 277.3 268.5 262.2 263.2 270.7 278.8 285.0 287.3 282.0 272.1 265.2 266.0 272.1 278.2 283.8 285.1 281.4 272.6 267.2 269.1 276.1 282.7 289.3 290.5 285.6 273.9 267.5 267.4 271.8 277.7 282.6 283.1 277.8 269.8 265.6 268.6 273.8 281.1 286.2 287.2 282.3 270.4 264.5 263.3 263.0 268.5 268.3 268.7 270.9 277.2 277.4 284.3 285.7 282.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 311.5 306.4 306.0 306.8 305.8 305.0 307.2 310.8 312.4 313.5 314.4 314.6 313.7 312.2 311.1 311.1 312.8 314.6 312.7 306.5 299.7 302.7 308.7 315.2 314.2 310.9 309.4 309.8 310.6 310.7 310.8 311.4 312.2 313.0 312.5 312.2 312.2 311.9 312.0 312.2 311.4 306.1 303.3 309.3 315.0 325.3 324.4 323.9 325.0 326.4 327.7 327.7 328.1 328.6 327.8 328.3 328.3 327.6 328.6 336.1 347.2 359.4 371.0 374.9 371.3 366.5 369.5 375.5 376.2 370.2 365.7 367.3 372.9 376.9 375.3 368.4 360.1 358.0 365.4 380.8 383.5 380.4 374.4 367.6 365.9 372.0 376.3 377.9 375.7 372.6 368.6 361.2 353.7 0.0 309.4 302.3 299.2 300.6 305.7 308.5 308.6 309.0 310.9 310.9 309.5 308.6 309.0 311.3 313.0 314.4 314.5 313.0 312.1 311.2 308.0 299.5 295.8 295.0 285.7 278.0 280.1 281.2 279.9 278.3 278.3 279.5 279.7 272.2 259.1 0.0 0.0 0.0 240.7 236.0 233.8 233.5 234.6 235.7 235.6 236.2 239.8 245.4 248.8 250.1 250.4 249.9 249.1 247.2 241.7 231.7 216.0 197.6 192.1 197.2 204.3 206.5 205.4 201.7 200.8 203.5 208.2 210.9 210.3 207.1 203.5 202.3 202.1 202.8 205.1 207.3 208.1 207.3 205.4 202.1 199.3 200.5 206.1 212.9 212.9 206.5 197.1 191.1 191.2 196.8 203.6 207.8 208.7 206.4 201.7 197.3 197.4 200.7 206.5 209.9 209.9 206.9 203.3 201.2 203.4 207.7 210.5 208.9 207.6 206.5 204.2 202.0 203.6 205.7 210.8 213.1 214.3 210.6 204.1 199.7 202.3 211.9 217.7 215.1 215.0 215.3 213.3 0.0 0.0 216.3 0.0 0.0 195.2 209.0 205.3 201.3 196.8 195.7 195.8 195.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 431.9 425.8 421.0 417.7 419.3 422.7 423.4 421.1 417.9 411.9 408.6 411.5 414.9 419.6 426.7 424.2 0.0 0.0 0.0 280.8 279.4 369.7 371.8 370.6 369.8 370.5 370.5 371.4 374.6 374.0 361.0 332.7 328.0 0.0 359.0 365.0 371.1 373.8 371.9 369.7 369.6 371.0 378.3 390.4 403.4 414.1 417.9 417.5 417.0 416.6 415.7 414.9 414.5 413.8 413.5 413.6 413.8 415.2 416.9 417.3 415.4 415.0 414.1 411.4 409.2 403.5 392.7 375.6 367.5 365.8 368.8 370.7 370.8 369.0 367.0 366.0 366.2 367.8 369.2 368.8 367.2 366.7 367.3 367.3 368.8 368.5 366.8 365.3 364.0 363.0 365.0 367.3 367.7 365.5 364.0 364.9 367.4 370.3 371.3 369.9 366.5 364.1 368.3 385.8 406.7 409.4 399.7 376.0 357.2 359.9 372.4 377.5 366.9 345.1 320.6 315.8 321.8 329.0 329.9 326.4 324.5 324.8 324.9 325.3 325.1 324.7 325.9 326.8 326.6 326.6 326.4 325.4 323.2 322.5 323.9 326.2 328.6 329.9 329.5 329.0 328.2 327.7 327.9 329.6 331.2 331.4 334.2 340.1 335.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 309.0 320.1 329.6 334.4 334.4 334.5 338.2 346.0 355.1 361.4 366.1 369.5 372.7 373.3 372.9 375.2 375.8 374.7 367.6 360.1 0.0 0.0 0.0 0.0 0.0 0.0 283.6 280.4 277.2 274.3 271.8 271.0 273.3 277.2 278.8 278.1 277.2 277.4 278.8 279.3 280.2 281.4 279.2 274.9 268.3 254.1 239.9 0.0 0.0 267.7 267.7 270.1 271.8 273.1 276.3 278.3 269.1 256.1 0.0 0.0 0.0 0.0 286.0 281.1 275.7 272.7 271.7 271.7 272.5 274.0 275.5 277.9 281.1 287.1 314.0 342.2 363.0 376.8 375.4 365.8 351.8 323.3 0.0 0.0 0.0 0.0 299.6 322.7 336.5 336.9 335.4 332.9 330.0 326.6 323.4 322.3 324.4 326.7 328.5 328.5 326.4 324.0 322.5 323.2 325.5 327.7 328.6 328.1 325.4 320.9 316.5 316.1 320.8 327.7 333.7 333.2 325.6 315.5 306.9 305.5 315.1 331.0 343.3 341.6 331.4 319.5 308.3 307.6 318.1 329.5 338.2 342.0 337.4 329.4 321.0 316.6 319.8 330.8 340.6 344.5 342.3 331.4 319.3 315.7 322.9 329.3 336.5 344.0 341.1 328.3 318.7 320.0 328.2 333.5 337.9 339.7 338.5 327.3 323.3 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 285.3 284.6 284.3 288.2 295.1 305.6 312.2 315.5 287.8 0.0 0.0 0.0 0.0 0.0 309.8 312.4 313.8 312.8 310.8 310.7 312.5 313.2 310.7 305.7 305.9 308.1 309.9 311.3 311.0 309.8 309.7 310.6 311.1 311.1 311.0 311.9 313.7 316.4 319.5 311.6 281.3 283.4 306.6 0.0 0.0 0.0 334.4 326.7 322.6 321.6 321.7 324.5 332.0 336.7 333.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 332.5 331.7 332.7 334.8 334.7 335.9 339.7 345.6 353.9 362.3 370.2 374.0 373.8 371.4 370.0 368.2 367.0 368.0 370.4 371.2 371.3 370.0 360.3 0.0 0.0 0.0 319.1 326.8 331.5 332.0 331.6 327.2 328.2 334.1 343.1 350.7 358.0 363.5 368.4 370.8 371.8 371.1 370.2 368.5 368.0 371.1 375.5 376.3 369.5 0.0 0.0 0.0 280.4 267.2 263.6 265.8 271.2 275.0 276.4 276.9 277.1 280.9 287.5 284.7 280.7 0.0 0.0 0.0 0.0 0.0 0.0 231.6 226.0 222.9 224.1 226.2 228.5 234.3 240.9 246.6 249.8 249.6 246.3 245.2 245.8 247.2 248.3 249.3 249.0 245.9 242.7 235.8 226.3 217.8 213.1 209.3 208.0 207.0 205.9 205.8 203.4 201.5 205.3 207.5 209.1 210.4 208.9 203.2 199.2 199.3 201.2 205.6 206.3 204.0 202.7 202.2 203.3 206.1 205.6 201.9 198.8 195.5 195.6 198.7 204.2 214.5 218.3 214.3 207.0 199.1 192.4 189.9 193.6 201.0 210.6 212.4 209.2 202.0 195.2 190.0 190.4 196.0 204.2 209.2 208.9 203.0 195.8 190.7 189.8 194.6 200.8 208.7 213.5 210.5 202.0 193.9 187.7 187.8 193.1 199.8 202.8 204.1 203.8 200.4 197.0 193.9 191.6 192.5 198.4 204.8 203.9 203.5 201.3 200.2 198.4 198.5 201.0 204.0 204.6 205.4 202.0 199.2 194.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 426.6 432.4 436.6 437.2 434.6 433.0 434.0 453.1 473.8 492.8 503.2 501.8 491.3 479.6 478.1 489.7 503.6 508.3 503.2 489.5 482.3 484.2 495.7 505.4 506.6 501.6 495.4 492.7 497.1 500.4 496.1 490.1 480.4 453.8 412.6 373.6 357.7 354.5 356.8 361.0 365.0 367.8 369.5 370.1 371.1 371.3 370.8 370.1 369.7 369.0 369.5 370.4 371.1 370.9 370.2 370.8 371.3 372.7 380.2 386.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 435.0 427.8 417.0 411.6 412.6 411.2 410.7 412.5 413.7 414.9 414.8 415.7 415.4 414.0 413.8 414.6 415.2 416.3 415.2 413.6 413.4 414.9 416.8 416.4 414.7 413.8 413.2 414.7 417.5 420.7 428.2 443.1 459.8 481.7 502.3 510.8 509.1 497.8 490.6 493.5 500.8 508.3 511.8 504.6 493.5 487.3 492.3 507.2 518.8 514.8 496.4 489.6 490.4 491.8 493.3 0.0 0.0 0.0 0.0 0.0 391.9 368.8 360.2 358.3 360.0 362.0 361.8 360.9 360.3 361.8 365.7 368.7 369.3 367.7 364.7 365.6 368.4 370.4 371.7 369.0 366.8 367.5 367.9 369.4 374.6 384.6 396.8 388.2 382.8 0.0 0.0 0.0 0.0 0.0 352.6 339.5 328.2 326.0 326.5 326.9 328.2 329.4 329.4 329.3 337.3 349.1 350.9 343.6 333.4 326.5 327.1 331.5 333.4 326.8 318.4 0.0 0.0 0.0 282.0 283.6 281.6 279.5 273.2 267.9 268.8 273.2 277.5 279.8 278.7 276.8 273.3 268.7 265.0 266.1 272.4 282.3 284.9 279.4 268.5 255.3 253.0 259.9 270.2 281.5 284.3 279.7 270.1 260.2 256.4 260.1 266.6 273.0 279.6 283.6 282.3 272.3 262.1 255.8 259.8 269.6 274.9 278.7 278.4 271.5 261.7 255.8 256.8 263.9 272.1 279.7 278.4 269.9 259.0 255.9 262.1 269.2 274.5 279.6 282.1 280.7 274.3 269.9 270.9 270.7 273.8 276.9 274.9 271.1 266.5 266.7 269.4 276.9 282.9 282.0 279.5 274.9 271.2 267.9 263.1 270.2 281.3 285.5 283.9 280.6 271.1 262.8 263.9 263.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 362.9 364.0 366.8 366.6 363.6 359.8 357.7 359.9 367.6 371.6 365.4 359.3 358.6 359.5 365.5 375.7 378.1 371.5 359.5 350.8 358.6 374.7 381.7 376.8 363.6 353.3 356.9 367.8 378.3 380.8 375.7 372.3 372.6 374.7 372.9 365.2 346.4 314.8 285.6 277.3 275.4 275.3 277.5 279.2 279.2 278.8 277.6 276.4 276.5 277.6 278.8 279.9 281.8 282.7 281.0 279.2 278.6 279.4 279.9 279.2 278.7 281.7 285.8 289.9 288.7 284.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 428.0 417.9 410.0 409.5 411.2 414.0 417.5 417.9 418.1 419.3 419.0 417.8 417.1 417.2 416.6 415.8 415.0 415.1 414.5 410.9 401.5 388.2 371.5 359.2 363.0 371.7 375.7 370.6 359.6 353.4 358.7 367.0 374.5 378.2 373.7 362.8 358.0 361.6 368.5 373.7 377.7 375.8 368.8 361.7 360.5 366.7 372.5 375.6 380.1 384.9 383.5 376.6 0.0 0.0 0.0 0.0 0.0 331.7 320.0 316.8 320.7 325.1 327.0 327.4 326.6 326.4 325.5 325.1 326.2 327.6 328.5 328.4 327.7 328.7 329.9 329.8 329.7 327.7 326.5 327.8 328.9 329.3 329.8 329.2 328.8 329.4 330.4 331.0 331.0 330.5 329.2 328.2 328.3 328.7 329.4 327.5 322.9 314.8 299.3 290.8 292.3 299.0 305.4 309.2 312.1 313.9 316.8 320.4 327.6 333.1 338.0 341.2 342.1 338.4 331.0 0.0 0.0 0.0 0.0 0.0 0.0 311.7 292.5 281.8 276.1 275.1 276.6 277.6 277.4 274.0 264.8 249.6 239.1 238.8 243.0 245.7 245.0 241.9 237.9 234.6 237.1 242.7 250.4 252.5 248.3 240.7 231.6 228.3 234.0 240.6 247.0 250.6 248.5 242.1 233.6 226.8 227.2 234.0 240.7 246.9 248.9 244.4 237.6 230.8 228.8 236.4 246.9 250.2 249.2 245.4 239.9 232.5 224.7 226.5 238.5 252.0 257.4 255.9 248.4 237.4 230.5 231.6 238.8 247.6 252.6 253.9 253.0 248.9 241.8 236.4 234.1 236.5 246.2 257.7 255.4 248.2 238.5 234.8 237.8 244.4 250.5 253.8 251.7 246.7 240.0 238.3 244.8 251.6 256.4 256.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 484.0 484.5 487.5 490.6 487.2 479.7 478.9 483.8 494.1 505.7 510.3 501.1 485.4 472.4 472.3 495.6 515.8 517.1 503.8 486.5 475.1 479.9 491.4 507.9 509.9 505.0 496.2 485.7 484.5 490.8 495.4 497.7 496.6 489.9 469.1 424.3 387.2 373.0 367.3 364.4 364.6 366.2 366.6 368.8 371.9 373.1 371.4 368.9 367.1 367.0 367.7 368.8 370.1 369.8 369.2 369.8 370.0 369.0 369.0 371.0 375.2 385.3 391.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 425.2 418.5 409.4 409.5 408.1 406.3 408.4 410.3 412.8 413.8 412.6 411.2 410.8 410.4 410.2 410.5 412.2 413.6 413.4 412.6 412.6 413.4 415.6 417.9 418.5 417.0 415.0 415.0 416.8 421.5 437.9 455.0 472.6 493.0 502.2 496.7 484.6 480.6 487.2 498.6 506.3 506.2 499.4 489.7 489.5 496.2 503.8 507.2 501.8 491.1 489.9 501.1 520.5 533.7 525.8 0.0 0.0 0.0 0.0 0.0 421.8 365.2 351.0 352.0 358.0 362.6 365.6 365.4 363.2 362.5 362.2 362.4 364.1 365.2 364.9 363.5 363.4 366.3 368.5 370.0 369.7 368.5 366.9 365.7 366.3 368.6 370.8 372.2 370.2 373.1 377.0 376.5 371.0 346.8 326.1 325.2 328.4 331.4 334.0 332.2 327.2 324.4 331.3 349.7 363.2 362.0 346.1 327.4 321.8 329.8 340.7 348.4 349.9 346.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 269.5 275.7 282.6 282.9 278.8 275.2 274.3 273.5 273.9 275.3 276.4 277.3 276.1 272.3 268.0 268.3 272.5 279.1 281.7 280.1 271.3 261.0 257.3 262.7 273.2 280.8 283.3 279.3 269.1 258.3 258.0 266.3 278.2 286.7 287.8 283.4 274.9 264.7 257.7 260.0 272.0 286.3 294.7 289.8 274.5 263.8 263.4 268.8 277.3 284.1 286.5 285.2 281.7 272.7 264.3 260.3 267.7 281.7 289.2 289.0 281.2 266.2 256.0 254.4 261.5 276.4 286.8 288.8 286.7 273.8 260.9 260.8 270.2 283.4 291.3 292.8 286.1 273.8 265.2 264.0 271.3 283.8 290.3 289.7 277.7 266.9 260.5 263.0 267.8 281.5 286.7 286.1 285.5 280.1 279.0 283.0 284.1 284.5 285.5 282.2 0.0 0.0 280.0 276.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0" - } -] diff --git a/tests/assets/soulxsinger/music.mp3 b/tests/assets/soulxsinger/music.mp3 deleted file mode 100644 index f3126e79e..000000000 Binary files a/tests/assets/soulxsinger/music.mp3 and /dev/null differ diff --git a/tests/assets/soulxsinger/zh_prompt.json b/tests/assets/soulxsinger/zh_prompt.json deleted file mode 100644 index 86ff3614e..000000000 --- a/tests/assets/soulxsinger/zh_prompt.json +++ /dev/null @@ -1,16 +0,0 @@ -[ - { - "index": "vocal_320_10687", - "language": "Mandarin", - "time": [ - 320, - 10687 - ], - "duration": "0.23 0.34 0.26 0.70 0.52 0.46 0.36 0.44 0.14 0.24 0.64 0.47 0.51 1.10 0.28 0.38 0.32 0.32 0.38 0.32 0.31 0.19 1.45", - "text": " 除 了 想 你 你 除 了 了 爱 你 你 我 什 么 什 么 都 愿 愿 意", - "phoneme": " zh_chu2 zh_le5 zh_xiang3 zh_ni3 zh_ni3 zh_chu2 zh_le5 zh_le5 zh_ai4 zh_ni3 zh_ni3 zh_wo3 zh_shen2 zh_me5 zh_shen2 zh_me5 zh_dou1 zh_yuan4 zh_yuan4 zh_yi4", - "note_pitch": "0 62 65 67 67 69 0 67 69 67 65 67 69 67 67 66 64 64 60 60 65 67 0", - "note_type": "1 2 2 2 2 3 1 2 2 3 2 2 3 1 2 2 2 2 2 2 2 3 2", - "f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 294.1 288.2 290.6 294.6 295.7 292.8 291.5 294.4 295.4 294.7 293.5 292.2 294.4 295.8 293.2 297.7 320.1 338.1 348.3 348.7 344.4 342.6 346.2 354.8 356.9 353.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 401.1 403.8 405.0 402.1 398.4 395.6 393.3 392.0 391.4 390.3 390.2 390.3 390.3 391.3 392.8 393.6 391.7 390.5 391.4 391.6 391.5 393.5 393.8 390.8 387.4 387.8 389.3 390.8 392.2 391.6 390.2 389.8 389.1 388.4 390.0 395.5 397.2 396.7 395.5 395.1 394.6 394.9 395.6 395.2 394.6 395.4 395.9 394.0 391.7 390.7 391.7 392.6 391.6 395.7 405.8 441.7 462.5 463.8 450.2 430.9 414.5 415.2 426.7 439.8 454.4 462.9 447.8 422.5 400.6 403.2 423.5 451.3 482.3 492.8 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 435.9 414.3 406.2 402.8 398.2 396.7 396.7 395.2 398.3 416.9 441.6 446.9 441.5 437.3 435.3 432.9 431.6 432.3 434.8 436.7 433.8 422.8 406.7 390.9 382.0 381.7 384.7 382.7 368.1 357.4 355.6 355.1 352.4 348.1 346.4 348.7 351.6 355.0 354.4 351.7 349.9 349.2 348.1 346.0 345.4 344.2 344.4 345.5 346.6 349.0 349.7 349.1 349.5 349.6 349.7 349.4 349.5 352.2 354.5 355.7 355.6 356.9 359.4 361.5 363.8 360.4 354.3 357.2 363.9 372.4 382.6 399.1 402.9 400.6 395.5 390.5 388.9 390.2 391.1 391.9 391.5 390.4 390.2 391.3 391.6 391.0 388.6 386.2 389.1 403.8 430.2 441.8 449.5 448.1 443.2 438.3 432.9 430.7 434.6 442.0 447.6 446.0 440.3 434.7 431.1 435.7 442.3 445.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 423.8 402.1 398.5 397.6 393.4 390.5 391.4 402.5 427.6 442.5 442.5 435.1 430.1 430.8 439.9 447.4 442.2 426.1 412.3 399.8 391.8 389.5 388.5 387.8 386.4 384.7 384.5 387.9 391.2 391.8 392.9 393.8 392.0 392.0 395.4 398.1 398.2 396.3 393.1 391.1 388.9 386.6 383.1 381.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 357.8 367.6 370.9 365.4 359.3 355.7 358.1 372.7 396.8 404.0 398.4 392.6 389.2 388.8 383.6 362.3 341.1 325.1 326.3 327.9 331.3 333.3 326.0 319.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 356.4 353.5 349.0 346.5 338.4 328.4 323.7 326.2 334.2 338.4 331.6 309.8 280.6 256.4 252.8 256.1 258.3 256.9 257.6 261.1 259.7 258.7 258.6 260.2 262.2 262.4 262.9 264.2 263.4 260.2 256.4 253.2 238.7 223.3 231.6 257.7 258.1 258.4 258.8 258.0 256.8 255.8 254.0 255.1 258.1 261.9 263.7 262.6 256.7 253.1 250.2 246.7 258.7 294.4 327.1 342.6 346.0 344.1 341.2 342.3 345.2 350.1 364.7 382.5 396.1 396.1 389.2 381.8 381.1 387.2 397.0 399.3 390.7 374.3 360.1 350.6 346.6 347.4 350.7 354.4 354.3 351.7 349.8 348.4 346.9 347.0 348.4 349.5 351.0 352.3 353.6 353.3 350.5 348.3 345.5 344.3 344.4 347.0 350.6 352.0 351.0 350.8 350.1 347.6 345.7 347.0 350.3 351.7 350.7 348.5 346.9 347.7 349.0 349.1 348.5 346.8 346.3 348.4 349.0 349.2 351.1 349.6 348.3 350.5 351.1 348.0 347.6 349.1 351.3 356.0 361.3 360.6 354.0 341.0 316.2 302.9 302.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0" - } -] diff --git a/tests/assets/soulxsinger/zh_prompt.mp3 b/tests/assets/soulxsinger/zh_prompt.mp3 deleted file mode 100644 index bc306410a..000000000 Binary files a/tests/assets/soulxsinger/zh_prompt.mp3 and /dev/null differ diff --git a/tests/e2e/offline_inference/test_soulxsinger_expansion.py b/tests/e2e/offline_inference/test_soulxsinger_expansion.py deleted file mode 100644 index a008be526..000000000 --- a/tests/e2e/offline_inference/test_soulxsinger_expansion.py +++ /dev/null @@ -1,241 +0,0 @@ -"""E2E offline inference tests for SoulX-Singer (single-stage, preprocess inline).""" - -import functools -import importlib -import json -import os -import shutil -from pathlib import Path - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - -import numpy as np -import pytest - -from tests.helpers.mark import hardware_test -from tests.helpers.media import get_asset_path -from tests.helpers.runtime import OmniRunner -from tests.helpers.stage_config import get_deploy_config_path -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - -PROMPT_AUDIO = get_asset_path("soulxsinger/zh_prompt.mp3") -TARGET_AUDIO = get_asset_path("soulxsinger/music.mp3") -PROMPT_METADATA = get_asset_path("soulxsinger/zh_prompt.json") -TARGET_METADATA = get_asset_path("soulxsinger/music.json") -SAMPLE_RATE = 24_000 - -# phone_set.json (phoneme vocab) is not shipped on HuggingFace; SVS loads it -# from the model dir. Stage it from this pinned upstream commit so the SVS DiT -# can run in CI; the SVS tests skip gracefully if it cannot be fetched. -_PHONE_SET_URL = ( - "https://raw.githubusercontent.com/Soul-AILab/SoulX-Singer/" - "81aeb3ae772c70093c3de74dc23c92d983801ae4/soulxsinger/utils/phoneme/phone_set.json" -) - -if not PROMPT_AUDIO.is_file() or not TARGET_AUDIO.is_file(): - pytest.skip( - f"Missing SoulX-Singer audio assets: {PROMPT_AUDIO.name}, {TARGET_AUDIO.name}", - allow_module_level=True, - ) - -pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.tts] - -_CASES = ( - pytest.param( - "SoulXSingerPipeline", - "soulxsinger_svs.yaml", - { - "language": "Mandarin", - "vocal_sep": False, - "control": "score", - "auto_shift": False, - "pitch_shift": 0, - }, - ("g2pM", "g2p_en"), - id="svs", - ), - pytest.param( - "SoulXSingerSVCPipeline", - "soulxsinger_svc.yaml", - {"vocal_sep": False, "auto_shift": False, "pitch_shift": 0}, - (), - id="svc", - ), -) - - -@functools.lru_cache(maxsize=1) -def _resolve_weights() -> tuple[Path, Path, Path]: - from huggingface_hub import snapshot_download - - base = Path(snapshot_download("Soul-AILab/SoulX-Singer", allow_patterns=["*"])) - - # phone_set.json is not on HF; best-effort stage it from pinned upstream so - # SVS can load. SVS tests skip (not fail) if both this and a manual copy are absent. - phone_set = base / "phoneme" / "phone_set.json" - if not phone_set.is_file() and not (base / "phone_set.json").is_file(): - try: - from urllib.request import urlretrieve - - phone_set.parent.mkdir(exist_ok=True) - urlretrieve(_PHONE_SET_URL, phone_set) - except Exception: - pass - - # make a temporary svc directory - svc_dir = base.parent / "SoulX-Singer-SVC" - svc_dir.mkdir(parents=True, exist_ok=True) - shutil.copyfile(base / "config.yaml", svc_dir / "config.yaml") - shutil.move(base / "model-svc.pt", svc_dir / "model-svc.pt") - - if raw := os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR"): - pre = Path(raw).expanduser().resolve() - if (pre / "rmvpe" / "rmvpe.pt").is_file(): - return base, svc_dir, pre - - from huggingface_hub import snapshot_download - - pre = Path(snapshot_download("Soul-AILab/SoulX-Singer-Preprocess", allow_patterns=["*"])) - return base, svc_dir, pre - - -@pytest.fixture(scope="session") -def soulx_weights() -> tuple[Path, Path, Path]: - try: - return _resolve_weights() - except Exception as exc: - pytest.skip(f"Set SOULXSINGER_MODEL_DIR / SOULX_PREPROCESS_WEIGHTS_DIR. ({exc})") - - -def _flatten_audio(audio_val) -> np.ndarray: - import torch - - if isinstance(audio_val, list): - chunks = [c.detach().cpu().float().numpy().reshape(-1) for c in audio_val if c is not None] - return np.concatenate(chunks) if chunks else np.array([], dtype=np.float32) - if isinstance(audio_val, torch.Tensor): - return audio_val.detach().cpu().float().numpy().reshape(-1) - return np.asarray(audio_val, dtype=np.float32).reshape(-1) - - -@hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.parametrize("architecture,deploy_yaml,extra_args,py_deps", _CASES) -def test_soulxsinger_multistage_from_audio( - soulx_weights: tuple[Path, Path, Path], - architecture: str, - deploy_yaml: str, - extra_args: dict, - py_deps: tuple[str, ...], -) -> None: - for mod in py_deps: - try: - importlib.import_module(mod) - except ImportError as exc: - pytest.fail(f"SoulX SVS requires {mod}: {exc}") - - base_dir, svc_dir, preprocess_dir = soulx_weights - - # SVS mode requires phone_set.json in the model directory - if architecture == "SoulXSingerPipeline": - json.dump( - {"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1}, - open(base_dir / "config.json", "w"), - indent=4, - ) - model = str(base_dir) - if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file(): - pytest.skip( - "SoulX-Singer SVS test requires phoneme/phone_set.json. " - "Copy it from github.com/Soul-AILab/SoulX-Singer into the model dir. " - "See `examples/offline_inference/text_to_speech/README.md` for details." - ) - if architecture == "SoulXSingerSVCPipeline": - json.dump( - {"model_type": "soulxsinger", "architectures": ["SoulXSingerSVCPipeline"], "max_num_seqs": 1}, - open(svc_dir / "config.json", "w"), - indent=4, - ) - model = str(svc_dir) - - with OmniRunner( - model, - deploy_config=get_deploy_config_path(deploy_yaml), - async_chunk=False, - ) as runner: - sampling = OmniDiffusionSamplingParams( - num_inference_steps=4, - guidance_scale=3.0, - seed=42, - extra_args={ - "prompt_audio": str(PROMPT_AUDIO), - "target_audio": str(TARGET_AUDIO), - "preprocess_weights_dir": str(preprocess_dir), - **extra_args, - }, - ) - prompt = {"prompt_token_ids": [0]} - outputs = runner.generate([prompt], sampling) - - assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output" - mm = outputs[0].multimodal_output - assert isinstance(mm, dict) and "audio" in mm - audio = _flatten_audio(mm["audio"]) - assert 12_000 <= audio.size - assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4 - duration_s = audio.size / SAMPLE_RATE - assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s" - - -@hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_soulxsinger_svs_precomputed(soulx_weights: tuple[Path, Path, Path]) -> None: - """SVS from precomputed note/lyric metadata. - - Exercises the SVS DiT + vocoder serving path while skipping the heavy - upstream preprocess (funasr ASR + ROSVOT note transcription), so it runs in - the merge gate without those non-pip dependencies. - """ - base_dir, _svc_dir, preprocess_dir = soulx_weights - - if not PROMPT_METADATA.is_file() or not TARGET_METADATA.is_file(): - pytest.skip(f"Missing SoulX-Singer metadata fixtures: {PROMPT_METADATA.name}, {TARGET_METADATA.name}") - if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file(): - pytest.skip( - "SoulX-Singer SVS requires phoneme/phone_set.json (upstream fetch unavailable). " - "See `examples/offline_inference/text_to_speech/README.md`." - ) - - json.dump( - {"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1}, - open(base_dir / "config.json", "w"), - indent=4, - ) - - with OmniRunner( - str(base_dir), - deploy_config=get_deploy_config_path("soulxsinger_svs.yaml"), - async_chunk=False, - ) as runner: - sampling = OmniDiffusionSamplingParams( - num_inference_steps=4, - guidance_scale=3.0, - seed=42, - extra_args={ - "language": "Mandarin", - "control": "score", - "prompt_metadata_path": str(PROMPT_METADATA), - "target_metadata_path": str(TARGET_METADATA), - "audio_path": str(PROMPT_AUDIO), - "preprocess_weights_dir": str(preprocess_dir), - }, - ) - prompt = {"prompt_token_ids": [0]} - outputs = runner.generate([prompt], sampling) - - assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output" - mm = outputs[0].multimodal_output - assert isinstance(mm, dict) and "audio" in mm - audio = _flatten_audio(mm["audio"]) - assert 12_000 <= audio.size - assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4 - duration_s = audio.size / SAMPLE_RATE - assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s" diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 89913a073..650429bf8 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -78,8 +78,6 @@ EXCLUDED_MODELS = [ "StableDiffusionXLPipeline", "Gr00tN1d7Pipeline", "Pi0Pipeline", - "SoulXSingerPipeline", - "SoulXSingerSVCPipeline", "SanaWmPipeline", ] diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index fd6d700a9..cecb952eb 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -103,10 +103,6 @@ from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import ( ) from vllm_omni.model_executor.models.qwen3_omni.pipeline import resolve_qwen3_omni_pipeline from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE -from vllm_omni.model_executor.models.soulx_singer.pipeline import ( - SOULXSINGER_SVC_PIPELINE, - SOULXSINGER_SVS_PIPELINE, -) from vllm_omni.model_executor.models.step_audio2.pipeline import ( STEP_AUDIO2_ASR_PIPELINE, STEP_AUDIO2_PIPELINE, @@ -184,8 +180,6 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "dynin_omni": DYNIN_OMNI_PIPELINE, "indextts2": INDEXTTS2_PIPELINE, "indextts2_5": INDEXTTS25_PIPELINE, - "soulxsinger_svc": SOULXSINGER_SVC_PIPELINE, - "soulxsinger_svs": SOULXSINGER_SVS_PIPELINE, } diff --git a/vllm_omni/deploy/soulxsinger_svc.yaml b/vllm_omni/deploy/soulxsinger_svc.yaml deleted file mode 100644 index a9d815806..000000000 --- a/vllm_omni/deploy/soulxsinger_svc.yaml +++ /dev/null @@ -1,19 +0,0 @@ -# SoulX-Singer SVC: single-stage DiT (preprocess inline in pre_process_func). -# -# Use: vllm serve --omni -pipeline: soulxsinger_svc -async_chunk: false -dtype: float16 -trust_remote_code: true -distributed_executor_backend: "mp" - -stages: - - stage_id: 0 - max_num_seqs: 1 - gpu_memory_utilization: 0.5 - enforce_eager: false - devices: "0" - default_sampling_params: - num_inference_steps: 32 - guidance_scale: 3.0 - seed: 42 diff --git a/vllm_omni/deploy/soulxsinger_svs.yaml b/vllm_omni/deploy/soulxsinger_svs.yaml deleted file mode 100644 index 3f021cb53..000000000 --- a/vllm_omni/deploy/soulxsinger_svs.yaml +++ /dev/null @@ -1,19 +0,0 @@ -# SoulX-Singer SVS: single-stage DiT (preprocess inline in pre_process_func). -# -# Use: vllm serve --omni -pipeline: soulxsinger_svs -async_chunk: false -dtype: float16 -trust_remote_code: true -distributed_executor_backend: "mp" - -stages: - - stage_id: 0 - max_num_seqs: 1 - gpu_memory_utilization: 0.5 - enforce_eager: false - devices: "0" - default_sampling_params: - num_inference_steps: 32 - guidance_scale: 3.0 - seed: 42 diff --git a/vllm_omni/diffusion/models/soulx_singer/__init__.py b/vllm_omni/diffusion/models/soulx_singer/__init__.py deleted file mode 100644 index 30235d2c1..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/__init__.py +++ /dev/null @@ -1,44 +0,0 @@ -"""SoulX-Singer SVS / SVC flow-matching singing models for vLLM-Omni.""" - -__all__ = [ - "FlowMatchingAudioPipeline", - "PipelineSoulXSingerSVS", - "PipelineSoulXSingerSVC", - "get_soulxsinger_post_process_func", - "get_soulxsinger_pre_process_func", - "get_soulxsinger_svc_pre_process_func", -] - - -def __getattr__(name: str): - if name in {"FlowMatchingAudioPipeline", "get_soulxsinger_post_process_func"}: - from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import ( - FlowMatchingAudioPipeline, - get_soulxsinger_post_process_func, - ) - - return { - "FlowMatchingAudioPipeline": FlowMatchingAudioPipeline, - "get_soulxsinger_post_process_func": get_soulxsinger_post_process_func, - }[name] - if name in {"PipelineSoulXSingerSVC", "get_soulxsinger_svc_pre_process_func"}: - from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_svc import ( - PipelineSoulXSingerSVC, - get_soulxsinger_svc_pre_process_func, - ) - - return { - "PipelineSoulXSingerSVC": PipelineSoulXSingerSVC, - "get_soulxsinger_svc_pre_process_func": get_soulxsinger_svc_pre_process_func, - }[name] - if name in {"PipelineSoulXSingerSVS", "get_soulxsinger_pre_process_func"}: - from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_svs import ( - PipelineSoulXSingerSVS, - get_soulxsinger_pre_process_func, - ) - - return { - "PipelineSoulXSingerSVS": PipelineSoulXSingerSVS, - "get_soulxsinger_pre_process_func": get_soulxsinger_pre_process_func, - }[name] - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/__init__.py b/vllm_omni/diffusion/models/soulx_singer/modules/__init__.py deleted file mode 100644 index 282ab12ab..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -from .convnext import ConvNeXtV2Block -from .decoder import CFMDecoder -from .flow_matching import FlowMatchingTransformer -from .llama import DiffLlama -from .mel_transform import MelSpectrogramEncoder -from .vocoder import Vocoder -from .whisper_encoder import WhisperEncoder - -__all__ = [ - "ConvNeXtV2Block", - "FlowMatchingTransformer", - "CFMDecoder", - "DiffLlama", - "MelSpectrogramEncoder", - "Vocoder", - "WhisperEncoder", -] diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/convnext.py b/vllm_omni/diffusion/models/soulx_singer/modules/convnext.py deleted file mode 100644 index 465898e8a..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/convnext.py +++ /dev/null @@ -1,46 +0,0 @@ -import torch -import torch.nn as nn - - -class GRN(nn.Module): - def __init__(self, dim): - super().__init__() - self.gamma = nn.Parameter(torch.zeros(1, 1, dim)) - self.beta = nn.Parameter(torch.zeros(1, 1, dim)) - - def forward(self, x): - Gx = torch.norm(x, p=2, dim=1, keepdim=True) - Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) - return self.gamma * (x * Nx) + self.beta + x - - -# ref: https://github.com/SWivid/F5-TTS/blob/main/src/f5_tts/model/modules.py#L247 -class ConvNeXtV2Block(nn.Module): - def __init__( - self, - dim: int, - intermediate_dim: int, - dilation: int = 1, - ): - super().__init__() - padding = (dilation * (7 - 1)) // 2 - self.dwconv = nn.Conv1d( - dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation - ) # depthwise conv - self.norm = nn.LayerNorm(dim, eps=1e-6) - self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers - self.act = nn.GELU() - self.grn = GRN(intermediate_dim) - self.pwconv2 = nn.Linear(intermediate_dim, dim) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - residual = x - x = x.transpose(1, 2) # b n d -> b d n - x = self.dwconv(x) - x = x.transpose(1, 2) # b d n -> b n d - x = self.norm(x) - x = self.pwconv1(x) - x = self.act(x) - x = self.grn(x) - x = self.pwconv2(x) - return residual + x diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/decoder.py b/vllm_omni/diffusion/models/soulx_singer/modules/decoder.py deleted file mode 100644 index 6b1fda616..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/decoder.py +++ /dev/null @@ -1,24 +0,0 @@ -import torch -import torch.nn as nn - -from vllm_omni.diffusion.models.soulx_singer.modules.flow_matching import FlowMatchingTransformer - - -class CFMDecoder(nn.Module): - def __init__(self, config): - super().__init__() - self.model = FlowMatchingTransformer(cfg=config, **config) - - def forward(self, mel, x_mask, decoder_inp, is_prompt): - outputs = self.model(mel, x_mask, decoder_inp, is_prompt) - - noise, x, flow_pred, final_mask, prompt_len = outputs["output"] - return noise, x, flow_pred, final_mask, prompt_len - - def reverse_diffusion(self, pt_mel, pt_decoder_inp, gt_decoder_inp, n_timesteps=32, cfg=1): - diffusion_cond = torch.cat([pt_decoder_inp, gt_decoder_inp], dim=1) - diffusion_cond_emb = self.model.cond_emb(diffusion_cond) - diffusion_prompt = pt_mel - - generated = self.model.reverse_diffusion(diffusion_cond_emb, diffusion_prompt, n_timesteps=n_timesteps, cfg=cfg) - return generated diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/flow_matching.py b/vllm_omni/diffusion/models/soulx_singer/modules/flow_matching.py deleted file mode 100644 index 19e87b65b..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/flow_matching.py +++ /dev/null @@ -1,423 +0,0 @@ -# https://github.com/open-mmlab/Amphion/blob/main/models/svc/flow_matching_transformer/fmt_model.py - -# Copyright (c) 2023 Amphion. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -import math - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from transformers import LlamaConfig - -from vllm_omni.diffusion.models.soulx_singer.modules.llama import DiffLlama - - -class FlowMatchingTransformer(nn.Module): - def __init__( - self, - mel_dim=100, - hidden_size=1024, - num_layers=12, - num_heads=16, - cfg_drop_prob=0.2, - use_embedding=True, - cond_codebook_size=1024, - cond_scale_factor=1, - sigma=1e-5, - time_scheduler="linear", - cfg=None, - ): - super().__init__() - self.cfg = cfg - - if cfg is not None: - mel_dim = getattr(cfg, "mel_dim", mel_dim) - hidden_size = getattr(cfg, "hidden_size", hidden_size) - num_layers = getattr(cfg, "num_layers", num_layers) - num_heads = getattr(cfg, "num_heads", num_heads) - cfg_drop_prob = getattr(cfg, "cfg_drop_prob", cfg_drop_prob) - cond_codebook_size = getattr(cfg, "cond_codebook_size", cond_codebook_size) - time_scheduler = getattr(cfg, "time_scheduler", time_scheduler) - sigma = getattr(cfg, "sigma", sigma) - cond_scale_factor = getattr(cfg, "cond_scale_factor", cond_scale_factor) - - self.mel_dim = mel_dim - self.hidden_size = hidden_size - self.num_layers = num_layers - self.num_heads = num_heads - self.cfg_drop_prob = cfg_drop_prob - self.cond_codebook_size = cond_codebook_size - self.time_scheduler = time_scheduler - self.sigma = sigma - self.cond_scale_factor = cond_scale_factor - - if use_embedding: - self.cond_emb = nn.Embedding(cond_codebook_size, self.hidden_size) - else: - self.cond_emb = nn.Linear(cond_codebook_size, self.hidden_size) - - if cond_scale_factor != 1: - self.do_resampling = True - assert np.log2(cond_scale_factor).is_integer() - - up_layers = [] - for _ in range(int(np.log2(cond_scale_factor))): - up_layers.extend( - [ - nn.ConvTranspose1d(hidden_size, hidden_size, kernel_size=4, stride=2, padding=1), - nn.GELU(), - ] - ) - self.resampling_layers = nn.Sequential(*up_layers) - else: - self.do_resampling = False - - ### REPA: Use the Wav2Vec2Bert features to align. ### - self.use_repa = "repa" in cfg - self.repa_layer_index = None - if self.use_repa: - self.repa_layer_index = cfg.repa.layer_index - - self.repa_mlp_layer = nn.Sequential( - nn.Linear(hidden_size, hidden_size * 4), - nn.SiLU(), - nn.Linear(hidden_size * 4, cfg.repa.output_dim), - ) - - ### CTC: Use the ASR loss ### - self.use_ctc = "ctc" in cfg - self.ctc_layer_index = None - if self.use_ctc: - self.ctc_layer_index = cfg.ctc.layer_index - - self.ctc_mlp_layer = nn.Sequential( - nn.Linear(hidden_size, hidden_size * 4), - nn.SiLU(), - nn.Linear(hidden_size * 4, cfg.ctc.output_dim), - ) - - self.reset_parameters() - - # DiffLlama config must be overridden manually to make sure the - # internal RoPE head dimension is correct - llama_config = LlamaConfig( - vocab_size=0, - bos_token_id=None, - eos_token_id=None, - hidden_size=hidden_size, - num_attention_heads=num_heads, - num_hidden_layers=num_layers, - rms_norm_eps=1e-6, - output_hidden_states=False, - output_attentions=False, - use_cache=False, - ) - self.diff_estimator = DiffLlama( - mel_dim=mel_dim, hidden_size=hidden_size, num_heads=num_heads, num_layers=num_layers, config=llama_config - ) - - self.sigma = sigma - - @torch.no_grad() - def forward_diffusion(self, x, t, is_prompt=None): - """ - x: (B, T, mel_dim) - t: (B,) - """ - new_t = t - t = t.unsqueeze(-1).unsqueeze(-1) - z = torch.randn(x.shape, dtype=x.dtype, device=x.device, requires_grad=False) # (B, T, mel_dim) - - # get prompt len - if torch.rand(1) <= self.cfg_drop_prob: - prompt_len = torch.zeros(x.shape[0]).to(x) - is_prompt = torch.zeros_like(x[:, :, 0]) - else: - if is_prompt is None: - prompt_len = torch.randint(min(x.shape[1] // 4, 5), int(x.shape[1] * 0.4), (x.shape[0],)).to( - x.device - ) # (B,) - - # get is_prompt - is_prompt = torch.zeros_like(x[:, :, 0]) # (B, T) - col_indices = torch.arange(is_prompt.shape[1]).repeat(is_prompt.shape[0], 1).to(prompt_len) # (B, T) - is_prompt[col_indices < prompt_len.unsqueeze(1)] = 1 # (B, T) 1 if prompt - else: - prompt_len = is_prompt.sum(dim=1) # (B,) - - mask = torch.ones_like(x[:, :, 0]) # mask if 1, not mask if 0 - mask[is_prompt.bool()] = 0 - mask = mask[:, :, None] - - # flow matching: xt = (1 - (1 - sigma) * t) * x0 + t * x; where x0 ~ N(0, 1), x is a sample - # flow gt: x - (1 - sigma) * x0 = x - (1 - sigma) * noise - xt = ((1 - (1 - self.sigma) * t) * z + t * x) * mask + x * (1 - mask) - - return xt, z, new_t, prompt_len, mask - - def loss_t(self, x, x_mask, t, cond=None, is_prompt=None): - xt, z, new_t, prompt_len, mask = self.forward_diffusion(x, t, is_prompt) - - noise = z - - # drop all condition for cfg, so if prompt_len is 0, we also drop cond - if cond is not None: - cond = cond * torch.where( - prompt_len > 0, - torch.ones_like(prompt_len), - torch.zeros_like(prompt_len), - ).to(cond.device).unsqueeze(-1).unsqueeze(-1) - - dit_output = self.diff_estimator(xt, new_t, cond, x_mask, return_dict=True) - flow_pred = dit_output["output"] # (B, T, mel_dim) - - # final mask used for loss calculation - final_mask = mask * x_mask[..., None] # (B, T, 1) - - results = {"output": (noise, x, flow_pred, final_mask, prompt_len)} - - if self.use_repa: - repa_hidden_states = dit_output["hidden_states"][self.repa_layer_index] # (B, T, hidden_size) - - repa_pred = self.repa_mlp_layer(repa_hidden_states) # (B, T, repa_dim) - results["repa"] = repa_pred - - if self.use_ctc: - ctc_hidden_states = dit_output["hidden_states"][self.ctc_layer_index] # (B, T, hidden_size) - ctc_pred = self.ctc_mlp_layer(ctc_hidden_states) # (B, T, ctc_dim) - results["ctc"] = ctc_pred - - return results - - def compute_loss(self, x, x_mask, cond=None, is_prompt=None): - # x0: (B, T, num_quantizer) - # x_mask: (B, T) mask is 0 for padding - t = torch.rand(x.shape[0], device=x.device, requires_grad=False) - t = torch.clamp(t, 1e-5, 1.0) - # generation is harder at the start than later; use a cosine - # scheduler for timestep t when time_scheduler == "cos". - if self.time_scheduler == "cos": - t = 1 - torch.cos(t * math.pi * 0.5) - else: - pass - return self.loss_t(x, x_mask, t, cond, is_prompt) - - def reset_parameters(self): - def _reset_parameters(m): - if isinstance(m, nn.MultiheadAttention): - if m._qkv_same_embed_dim: - nn.init.normal_(m.in_proj_weight, std=0.02) - else: - nn.init.normal_(m.q_proj_weight, std=0.02) - nn.init.normal_(m.k_proj_weight, std=0.02) - nn.init.normal_(m.v_proj_weight, std=0.02) - - if m.in_proj_bias is not None: - nn.init.constant_(m.in_proj_bias, 0.0) - nn.init.constant_(m.out_proj.bias, 0.0) - if m.bias_k is not None: - nn.init.xavier_normal_(m.bias_k) - if m.bias_v is not None: - nn.init.xavier_normal_(m.bias_v) - - elif ( - isinstance(m, nn.Conv1d) - or isinstance(m, nn.ConvTranspose1d) - or isinstance(m, nn.Conv2d) - or isinstance(m, nn.ConvTranspose2d) - ): - m.weight.data.normal_(0.0, 0.02) - - elif isinstance(m, nn.Linear): - m.weight.data.normal_(mean=0.0, std=0.02) - if m.bias is not None: - m.bias.data.zero_() - - elif isinstance(m, nn.Embedding): - m.weight.data.normal_(mean=0.0, std=0.02) - if m.padding_idx is not None: - m.weight.data[m.padding_idx].zero_() - - self.apply(_reset_parameters) - - @torch.no_grad() - def reverse_diffusion( - self, - cond, - prompt, - x_mask=None, - prompt_mask=None, - n_timesteps=10, - cfg=1.0, - rescale_cfg=0.75, - ): - h = 1.0 / n_timesteps - prompt_len = prompt.shape[1] - target_len = cond.shape[1] - prompt_len - trunk_dtype = cond.dtype - - if x_mask is None: - x_mask = torch.ones(cond.shape[0], target_len, dtype=trunk_dtype, device=cond.device) - if prompt_mask is None: - prompt_mask = torch.ones(cond.shape[0], prompt_len, dtype=trunk_dtype, device=cond.device) - xt_mask = torch.cat([prompt_mask, x_mask], dim=1) - z = torch.randn( - (cond.shape[0], target_len, self.mel_dim), - dtype=trunk_dtype, - device=cond.device, - requires_grad=False, - ) - xt = z - h_tensor = torch.tensor(h, dtype=trunk_dtype, device=cond.device) - - # t from 0 to 1: x0 = z ~ N(0, 1) - for i in range(n_timesteps): - xt_input = torch.cat([prompt, xt], dim=1) - t = torch.full( - (z.shape[0],), - fill_value=(i + 0.5) * h, - dtype=trunk_dtype, - device=z.device, - ) - flow_pred = self.diff_estimator(xt_input, t, cond, xt_mask) - flow_pred = flow_pred[:, prompt_len:, :] - - # cfg - if cfg > 0: - uncond_flow_pred = self.diff_estimator(xt, t, torch.zeros_like(cond)[:, : xt.shape[1], :], x_mask) - pos_flow_pred_std = flow_pred.std() - flow_pred_cfg = flow_pred + cfg * (flow_pred - uncond_flow_pred) - rescale_flow_pred = flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std() - flow_pred = rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg - - xt = xt + flow_pred * h_tensor - - return xt - - @torch.no_grad() - def reverse_diffusion_v2( - self, - cond, - prompt, - x_mask=None, - prompt_mask=None, - n_timesteps=10, - cfg=1.0, - rescale_cfg=0.75, - ): - h = 1.0 / n_timesteps - prompt_len = prompt.shape[1] - target_len = cond.shape[1] - prompt_len * 2 - - if x_mask is None: - x_mask = torch.ones(cond.shape[0], target_len).to(cond.device) # (B, T) - if prompt_mask is None: - prompt_mask = torch.ones(cond.shape[0], prompt_len).to(cond.device) # (B, prompt_len) - xt_mask = torch.cat([prompt_mask, x_mask, prompt_mask], dim=1) - z = torch.randn( - (cond.shape[0], target_len, self.mel_dim), - dtype=cond.dtype, - device=cond.device, - requires_grad=False, - ) - xt = z - - # t from 0 to 1: x0 = z ~ N(0, 1) - for i in range(n_timesteps): - xt_input = torch.cat([prompt, xt, prompt], dim=1) - t = (0 + (i + 0.5) * h) * torch.ones(z.shape[0], dtype=z.dtype, device=z.device) - flow_pred = self.diff_estimator(xt_input, t, cond, xt_mask) - flow_pred = flow_pred[:, prompt_len:-prompt_len, :] - - # cfg - if cfg > 0: - uncond_flow_pred = self.diff_estimator(xt, t, torch.zeros_like(cond)[:, : xt.shape[1], :], x_mask) - pos_flow_pred_std = flow_pred.std() - flow_pred_cfg = flow_pred + cfg * (flow_pred - uncond_flow_pred) - rescale_flow_pred = flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std() - flow_pred = rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg - - dxt = flow_pred * h - xt = xt + dxt - - return xt - - def forward( - self, - x: torch.Tensor, - x_mask: torch.Tensor, - cond_code: torch.Tensor, - is_prompt: torch.Tensor | None = None, - ): - """ - Args: - x: (B, T, mel_dim) - x_mask: (B, T) - cond_code: (B, T), Note that cond_code might be not at 50Hz! - """ - T = x.shape[1] - - cond = self.cond_emb(cond_code) # (B, T, hidden_size) - if self.do_resampling: - # Align to the frame rate of Mels - cond = self.resampling_layers(cond.transpose(1, 2)).transpose(1, 2) - - # print("cond_code: {}, after resampling: {}".format(cond_code.shape, cond.shape)) - - if cond.shape[1] >= T: # Check time dimension - cond = cond[:, :T, :] - else: - padding_frames = T - cond.shape[1] - last_frame = cond[:, -1:, :] - padding = last_frame.repeat(1, padding_frames, 1) - cond = torch.cat([cond, padding], dim=1) - - return self.compute_loss(x, x_mask, cond, is_prompt) - - -if __name__ == "__main__": - model_cfg = { - "mel_dim": 128, - "hidden_size": 256, - "num_layers": 8, - "num_heads": 8, - "cfg_drop_prob": 0.2, - "use_embedding": False, - "cond_codebook_size": 256, - "cond_scale_factor": 1, - "sigma": 1e-5, - "time_scheduler": "cos", - } - - device = "cuda" - x = torch.randn(2, 100, 128).to(device) - x_mask = torch.ones(2, 100).to(device) - # cond_code = torch.randint(0, 16384, (2, 25)).to(device) - cond_code = torch.randn(2, 100, 256).to(device) - - model = FlowMatchingTransformer(cfg=model_cfg, **model_cfg).to(device) - outputs = model(x, x_mask, cond_code) - print(outputs) - - noise, x, flow_pred, final_mask, prompt_len = outputs["output"] - final_mask = final_mask.squeeze(-1) - - flow_gt = x - (1 - 1e-5) * noise - - # [B, n_frames, D] - diff_loss = F.l1_loss(flow_pred, flow_gt, reduction="none").float() * final_mask.unsqueeze(-1) - diff_loss = torch.mean(diff_loss, dim=2).sum() / final_mask.sum() - - print("diff_loss:", diff_loss.item()) - - diffusion_cond = torch.randn(2, 150, 256).to(device) - diffusion_cond_emb = model.cond_emb(diffusion_cond) - diffusion_prompt = torch.randn(2, 50, 128).to(device) - n_timesteps = 32 - - generated = model.reverse_diffusion(diffusion_cond_emb, diffusion_prompt, n_timesteps=n_timesteps) - print("generated:", generated.shape) diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/llama.py b/vllm_omni/diffusion/models/soulx_singer/modules/llama.py deleted file mode 100644 index b0b6afce3..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/llama.py +++ /dev/null @@ -1,393 +0,0 @@ -import math - -import torch -import torch.nn as nn -from cache_dit import ForwardPattern -from transformers import LlamaConfig, LlamaModel -from transformers.modeling_outputs import BaseModelOutputWithPast -from transformers.models.llama.modeling_llama import LlamaDecoderLayer - -from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig -from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput - - -class SinusoidalPosEmb(nn.Module): - def __init__(self, dim): - super().__init__() - self.dim = dim - - def forward(self, x): - device = x.device - dtype = x.dtype - half_dim = self.dim // 2 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, device=device, dtype=dtype) * -emb) - emb = x[:, None] * emb[None, :] - emb = torch.cat((emb.sin(), emb.cos()), dim=-1) - return emb - - -class LlamaAdaptiveRMSNorm(nn.Module): - def __init__(self, hidden_size=1024, eps=1e-6, dim_cond=1024): - super().__init__() - self.to_weight = nn.Linear(dim_cond, hidden_size) - nn.init.zeros_(self.to_weight.weight) - nn.init.ones_(self.to_weight.bias) - self.variance_epsilon = eps - self._is_hf_initialized = True - - def forward(self, hidden_states, cond_embedding): - input_dtype = hidden_states.dtype - variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - - weight = self.to_weight(cond_embedding) - if len(weight.shape) == 2: - weight = weight.unsqueeze(1) - - return (weight * hidden_states).to(input_dtype) - - -class LlamaNARDecoderLayer(LlamaDecoderLayer): - def __init__(self, config: LlamaConfig, layer_idx: int): - super().__init__(config, layer_idx) - - self.input_layernorm = LlamaAdaptiveRMSNorm( - config.hidden_size, - eps=config.rms_norm_eps, - dim_cond=config.hidden_size, - ) - self.post_attention_layernorm = LlamaAdaptiveRMSNorm( - config.hidden_size, - eps=config.rms_norm_eps, - dim_cond=config.hidden_size, - ) - - def forward( - self, - hidden_states: torch.Tensor, - cond_embedding: torch.Tensor, - attention_mask: torch.Tensor | None = None, - position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, - output_attentions: bool | None = False, - cache_position: torch.LongTensor | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor, ...]: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states, cond_embedding=cond_embedding) - - # RoPE is computed by the parent model and passed in as (cos, sin). - hidden_states, self_attn_weights = self.self_attn( - hidden_states=hidden_states, - position_embeddings=position_embeddings, - attention_mask=attention_mask, - cache_position=cache_position, - output_attentions=output_attentions, - **kwargs, - ) - - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm( - hidden_states, - cond_embedding=cond_embedding, - ) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - - outputs = (hidden_states,) - - if output_attentions: - outputs += (self_attn_weights,) - - return outputs - - -class DiffLlama(LlamaModel): - # Regional torch.compile in DiffusionModelRunner targets these blocks. - _repeated_blocks = ["LlamaNARDecoderLayer"] - _layerwise_offload_blocks_attrs = ["layers"] - - # cache-dit: Llama-style backbone with static condition embeddings injected - # before the decoder loop (like Stable Audio Open). Pattern_3 handles the - # cross-attention-style conditioning on the ``layers`` ModuleList. - _cache_dit_adapter_config = CacheDiTAdapterConfig( - block_forward_patterns={ - "layers": ForwardPattern.Pattern_3, - }, - ) - - # Sequence parallelism: shard mel sequence before the first decoder layer and - # gather after the output projection. Tensor parallelism still requires - # migrating attention to vLLM parallel linear layers. - _sp_plan = { - "layers.0": { - "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), - }, - "mel_out_mlp": SequenceParallelOutput(gather_dim=1, expected_dims=3), - } - - @staticmethod - def _is_transformer_block(name: str, module: nn.Module) -> bool: - del module - return name.startswith("layers.") and name.split(".")[-1].isdigit() - - _hsdp_shard_conditions = [_is_transformer_block] - - def __init__( - self, - mel_dim=100, - hidden_size=1024, - num_heads=16, - num_layers=16, - dropout=0.1, - ffn_dropout=0.1, - attention_dropout=0.0, - config: LlamaConfig | None = None, - ): - if config is None: - config = LlamaConfig( - vocab_size=32000, - hidden_size=hidden_size, - intermediate_size=hidden_size * 4, - num_hidden_layers=num_layers, - num_attention_heads=num_heads, - num_key_value_heads=num_heads, - max_position_embeddings=4096, - rms_norm_eps=1e-6, - attention_dropout=attention_dropout, - use_cache=False, - attn_implementation="eager", - ) - - config.use_cache = False - - super().__init__(config) - - # Keep custom decoder layers on the same attention backend. - layer_config = LlamaConfig( - vocab_size=config.vocab_size, - hidden_size=hidden_size, - intermediate_size=hidden_size * 4, - num_hidden_layers=num_layers, - num_attention_heads=num_heads, - num_key_value_heads=num_heads, - max_position_embeddings=4096, - rms_norm_eps=config.rms_norm_eps, - attention_dropout=attention_dropout, - use_cache=False, - attn_implementation="eager", - bos_token_id=None, - eos_token_id=None, - ) - - self.layers = nn.ModuleList( - [ - LlamaNARDecoderLayer( - layer_config, - layer_idx=i, - ) - for i in range(num_layers) - ] - ) - - self.norm = LlamaAdaptiveRMSNorm(hidden_size, dim_cond=hidden_size) - - self.diff_step_embedding = SinusoidalPosEmb(hidden_size) - self.diff_step_mlp = nn.Sequential( - nn.Linear(hidden_size, hidden_size * 4), - nn.SiLU(), - nn.Linear(hidden_size * 4, hidden_size), - ) - - self.cond_mlp = nn.Sequential( - nn.Linear(hidden_size, hidden_size * 4), - nn.SiLU(), - nn.Linear(hidden_size * 4, hidden_size), - ) - - self.mel_mlp = nn.Sequential( - nn.Linear(mel_dim, hidden_size * 4), - nn.SiLU(), - nn.Linear(hidden_size * 4, hidden_size), - ) - - self.mel_out_mlp = nn.Sequential( - nn.Linear(hidden_size, hidden_size * 4), - nn.SiLU(), - nn.Linear(hidden_size * 4, mel_dim), - ) - - self.embed_tokens = None - - self.post_init() - - def _prepare_decoder_attention_mask( - self, - attention_mask, - input_shape, - inputs_embeds, - past_key_values_length=0, - ): - """ - Build a 4D additive padding mask. - - This model is non-autoregressive, so no causal mask is added here. - """ - - combined_attention_mask = None - - def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None): - bsz, src_len = mask.size() - tgt_len = tgt_len if tgt_len is not None else src_len - - expanded_mask = ( - mask[:, None, None, :] - .expand( - bsz, - 1, - tgt_len, - src_len, - ) - .to(dtype) - ) - - inverted_mask = 1.0 - expanded_mask - - return inverted_mask.masked_fill( - inverted_mask.to(torch.bool), - torch.finfo(dtype).min, - ) - - if attention_mask is not None: - expanded_attn_mask = _expand_mask( - attention_mask, - inputs_embeds.dtype, - tgt_len=input_shape[-1], - ).to(inputs_embeds.device) - - combined_attention_mask = ( - expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask - ) - - return combined_attention_mask - - def forward( - self, - x, - diffusion_step, - cond, - x_mask, - input_ids: torch.LongTensor = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: list[torch.FloatTensor] | None = None, - inputs_embeds: torch.FloatTensor | None = None, - use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = False, - ) -> BaseModelOutputWithPast | torch.Tensor | dict: - batch_size, seq_length, _ = x.shape - - cond_embedding = self.cond_mlp(cond) - x = self.mel_mlp(x) - - diffusion_step = self.diff_step_embedding(diffusion_step).to(device=x.device, dtype=x.dtype) - diffusion_step = self.diff_step_mlp(diffusion_step) - - x = x + cond_embedding - - inputs_embeds = x - attention_mask = x_mask - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - - # KV cache is disabled because this diffusion decoder processes full sequences. - - if past_key_values is not None: - raise NotImplementedError("KV cache is disabled for this non-autoregressive DiffLlama implementation.") - - if position_ids is None: - device = input_ids.device if input_ids is not None else inputs_embeds.device - position_ids = torch.arange( - 0, - seq_length, - dtype=torch.long, - device=device, - ) - position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) - else: - position_ids = position_ids.view(batch_size, seq_length).long() - - if attention_mask is None: - attention_mask = torch.ones( - (batch_size, seq_length), - dtype=torch.bool, - device=inputs_embeds.device, - ) - - attention_mask = self._prepare_decoder_attention_mask( - attention_mask, - (batch_size, seq_length), - inputs_embeds, - past_key_values_length=0, - ) - - hidden_states = inputs_embeds - - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - all_layer_hidden_states = [] - - cache_position = torch.arange( - 0, - seq_length, - device=hidden_states.device, - dtype=torch.long, - ) - - # LlamaAttention expects precomputed rotary embeddings. - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - for idx, decoder_layer in enumerate(self.layers): - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - raise NotImplementedError - - layer_outputs = decoder_layer( - hidden_states, - cond_embedding=diffusion_step, - attention_mask=attention_mask, - position_embeddings=position_embeddings, - output_attentions=output_attentions, - cache_position=cache_position, - ) - - hidden_states = layer_outputs[0] - all_layer_hidden_states.append(hidden_states.clone()) - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - hidden_states = self.norm(hidden_states, cond_embedding=diffusion_step) - - if output_hidden_states: - all_hidden_states += (hidden_states,) - - hidden_states = self.mel_out_mlp(hidden_states) - - if return_dict: - return { - "output": hidden_states, - "hidden_states": all_layer_hidden_states, - "attentions": all_self_attns, - } - - return hidden_states diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/mel_transform.py b/vllm_omni/diffusion/models/soulx_singer/modules/mel_transform.py deleted file mode 100644 index cc767dcf3..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/mel_transform.py +++ /dev/null @@ -1,151 +0,0 @@ -import math -from typing import Any - -import numpy as np -import torch -import torch.nn as nn - -from vllm_omni.utils.audio import mel_filter_bank as librosa_mel_fn - - -def dynamic_range_compression(x, c=1, clip_val=1e-5): - return np.log(np.clip(x, a_min=clip_val, a_max=None) * c) - - -def dynamic_range_decompression(x, c=1): - return np.exp(x) / c - - -def dynamic_range_compression_torch(x, c=1, clip_val=1e-5): - return torch.log(torch.clamp(x, min=clip_val) * c) - - -def dynamic_range_decompression_torch(x, c=1): - return torch.exp(x) / c - - -def spectral_normalize_torch(magnitudes): - output = dynamic_range_compression_torch(magnitudes) - return output - - -def spectral_de_normalize_torch(magnitudes): - output = dynamic_range_decompression_torch(magnitudes) - return output - - -class MelSpectrogram(nn.Module): - def __init__( - self, - n_fft, - num_mels, - sampling_rate, - hop_size, - win_size, - fmin, - fmax, - center=False, - ): - super().__init__() - self.n_fft = n_fft - self.hop_size = hop_size - self.win_size = win_size - self.sampling_rate = sampling_rate - self.num_mels = num_mels - self.fmin = fmin - self.fmax = fmax - self.center = center - - mel_basis = {} - hann_window = {} - - mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax) - mel_basis = mel.to(torch.float32) - hann_window = torch.hann_window(win_size) - - self.register_buffer("mel_basis", mel_basis) - self.register_buffer("hann_window", hann_window) - - def forward(self, y): - y = torch.nn.functional.pad( - y.unsqueeze(1), - ( - int((self.n_fft - self.hop_size) / 2), - int((self.n_fft - self.hop_size) / 2), - ), - mode="reflect", - ) - y = y.squeeze(1) - spec = torch.stft( - y, - self.n_fft, - hop_length=self.hop_size, - win_length=self.win_size, - window=self.hann_window, - center=self.center, - pad_mode="reflect", - normalized=False, - onesided=True, - return_complex=True, - ) - spec = torch.view_as_real(spec) - - spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) - - spec = torch.matmul(self.mel_basis, spec) - spec = spectral_normalize_torch(spec) - - return spec - - -def load_mel_spectrogram(): - return load_mel_spectrogram_from_cfg(None) - - -def _get_from_mapping(cfg: Any, key: str, default: Any = None) -> Any: - """Safely read a field from a dict/OmegaConf-like object.""" - if cfg is None: - return default - if isinstance(cfg, dict): - return cfg.get(key, default) - return getattr(cfg, key, default) - - -def load_mel_spectrogram_from_cfg(audio_cfg: Any | None = None) -> MelSpectrogram: - """Build MelSpectrogram from `audio_config`-like config. - - Expected keys (either in dict or Hydra/OmegaConf object): - - hop_size, sample_rate (or sampling_rate), n_fft, num_mels, win_size, fmin, fmax - """ - # Defaults keep current behavior. - mel_cfg: dict[str, Any] = { - "hop_size": _get_from_mapping(audio_cfg, "hop_size", 480), - "sampling_rate": _get_from_mapping( - audio_cfg, - "sampling_rate", - _get_from_mapping(audio_cfg, "sample_rate", 24000), - ), - "n_fft": _get_from_mapping(audio_cfg, "n_fft", 1920), - "num_mels": _get_from_mapping(audio_cfg, "num_mels", 128), - "win_size": _get_from_mapping(audio_cfg, "win_size", 1920), - "fmin": _get_from_mapping(audio_cfg, "fmin", 0), - "fmax": _get_from_mapping(audio_cfg, "fmax", 12000), - } - - mel_model = MelSpectrogram(**mel_cfg) - mel_model.eval() - return mel_model - - -class MelSpectrogramEncoder(nn.Module): - def __init__(self, audio_config: dict | None = None): - super().__init__() - self.model = load_mel_spectrogram_from_cfg(audio_config) - audio_config = audio_config or {} - self.mel_mean = audio_config.get("mel_mean", -4.92) - self.mel_var = audio_config.get("mel_var", 8.14) - - def forward(self, x): - x = self.model(x).transpose(1, 2) - x = (x - self.mel_mean) / math.sqrt(self.mel_var) - return x diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/note_transcription.py b/vllm_omni/diffusion/models/soulx_singer/modules/note_transcription.py deleted file mode 100644 index 70933dd92..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/note_transcription.py +++ /dev/null @@ -1,958 +0,0 @@ -import math - -import numpy as np -import packaging.version as version -import torch -import torch.nn as nn - -# ---------------------- note transcription common layers ---------------------- - - -class LayerNorm(nn.LayerNorm): - def __init__(self, nout: int, dim: int = -1, eps: float = 1e-5): - """Construct an LayerNorm object.""" - super().__init__(nout, eps=eps) - self.dim = dim - - def forward(self, x): - if self.dim == -1: - return super().forward(x) - return super().forward(x.transpose(1, -1)).transpose(1, -1) - - -class Reshape(nn.Module): - def __init__(self, *args): - super().__init__() - self.shape = args - - def forward(self, x): - return x.view(self.shape) - - -class Permute(nn.Module): - def __init__(self, *args): - super().__init__() - self.args = args - - def forward(self, x): - return x.permute(self.args) - - -class LambdaLayer(nn.Module): - def __init__(self, lambd): - super().__init__() - self.lambd = lambd - - def forward(self, x): - return self.lambd(x) - - -class ResidualBlock(nn.Module): - """ - Implements a residual block pattern: - [Norm] → [Conv1d] → [Scale] → [Activation] → [Conv1d] - repeated n times, each as a residual unit. - """ - - def __init__( - self, - channels: int, - kernel_size: int, - dilation: int, - n: int = 2, - c_multiple: int = 2, - ln_eps: float = 1e-12, - ): - super().__init__() - - self.blocks = nn.ModuleList( - [ - nn.Sequential( - LayerNorm(channels, dim=1, eps=ln_eps), - nn.Conv1d( - channels, - c_multiple * channels, - kernel_size, - dilation=dilation, - padding=(dilation * (kernel_size - 1)) // 2, - ), - LambdaLayer(lambda x: x * kernel_size**-0.5), # scaling for stability - nn.LeakyReLU(), - nn.Conv1d(c_multiple * channels, channels, kernel_size=1, dilation=dilation), - ) - for _ in range(n) - ] - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - nonpadding = (x.abs().sum(1) > 0).float()[:, None, :] - for block in self.blocks: - x_res = block(x) - x = (x + x_res) * nonpadding - return x - - -class ConvBlocks(nn.Module): - """ - Decodes the expanded phoneme encoding into spectrograms. - """ - - def __init__( - self, - hidden_size, - out_dims, - dilations, - kernel_size, - layers_in_block=2, - c_multiple=2, - ln_eps=1e-5, - is_btc=True, - post_net_kernel=3, - ): - super().__init__() - self.is_btc = is_btc - - res_blocks = [ - ResidualBlock( - hidden_size, - kernel_size=kernel_size, - dilation=1, - n=layers_in_block, - c_multiple=c_multiple, - ln_eps=ln_eps, - ) - for _ in range(1) - ] # default 1 - self.res_blocks = nn.Sequential(*res_blocks) - - # Output normalization and conv - self.last_norm = LayerNorm(hidden_size, dim=1, eps=ln_eps) - self.post_net1 = nn.Conv1d(hidden_size, out_dims, kernel_size=post_net_kernel, padding=post_net_kernel // 2) - - def forward(self, x: torch.Tensor, nonpadding: torch.Tensor | None = None) -> torch.Tensor: - if self.is_btc: - x = x.transpose(1, 2) - if nonpadding is None: - nonpadding = (x.abs().sum(1) > 0).float()[:, None, :] - elif self.is_btc: - nonpadding = nonpadding.transpose(1, 2) - x = self.res_blocks(x) * nonpadding - x = self.last_norm(x) * nonpadding - x = self.post_net1(x) * nonpadding - if self.is_btc: - x = x.transpose(1, 2) - return x - - -def Linear(in_features: int, out_features: int, bias: bool = True, init_type: str = "xavier"): - m = nn.Linear(in_features, out_features, bias) - if init_type == "xavier": - nn.init.xavier_uniform_(m.weight) - elif init_type == "kaiming": - nn.init.kaiming_normal_(m.weight, mode="fan_in") - if bias: - nn.init.constant_(m.bias, 0.0) - return m - - -def Embedding(num_embeddings, embedding_dim, padding_idx=None, init_type="normal"): - m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx) - if init_type == "normal": - nn.init.normal_(m.weight, mean=0, std=embedding_dim**-0.5) - elif init_type == "kaiming": - nn.init.kaiming_normal_(m.weight, mode="fan_in") - if padding_idx is not None: - nn.init.constant_(m.weight[padding_idx], 0) - return m - - -# ---------------------- note transcription espnet attention ---------------------- - - -class MultiHeadedAttention(nn.Module): - """Multi-Head Attention layer. - Args: - n_head (int): The number of heads. - n_feat (int): The number of features. - dropout_rate (float): Dropout rate. - """ - - def __init__(self, n_head, n_feat, dropout_rate, flash=False): - """Construct an MultiHeadedAttention object.""" - super().__init__() - assert n_feat % n_head == 0 - # We assume d_v always equals d_k - self.d_k = n_feat // n_head - self.h = n_head - self.linear_q = nn.Linear(n_feat, n_feat) - self.linear_k = nn.Linear(n_feat, n_feat) - self.linear_v = nn.Linear(n_feat, n_feat) - self.linear_out = nn.Linear(n_feat, n_feat) - self.attn = None - self.dropout = nn.Dropout(p=dropout_rate) - self.dropout_rate = dropout_rate - self.flash = flash - - def forward_qkv(self, query, key, value): - """Transform query, key and value. - Args: - query (torch.Tensor): Query tensor (#batch, time1, size). - key (torch.Tensor): Key tensor (#batch, time2, size). - value (torch.Tensor): Value tensor (#batch, time2, size). - Returns: - torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). - torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). - torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). - """ - n_batch = query.size(0) - q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) - k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) - v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) - q = q.transpose(1, 2) # (batch, head, time1, d_k) - k = k.transpose(1, 2) # (batch, head, time2, d_k) - v = v.transpose(1, 2) # (batch, head, time2, d_k) - - return q, k, v - - def forward_attention(self, value, scores, mask): - """Compute attention context vector. - Args: - value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). - scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). - mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). - Returns: - torch.Tensor: Transformed value (#batch, time1, d_model) - weighted by the attention score (#batch, time1, time2). - """ - n_batch = value.size(0) - if mask is not None: - mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) - min_value = float(np.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) - scores = scores.masked_fill(mask, min_value) - self.attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0) # (batch, head, time1, time2) - else: - self.attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) - - p_attn = self.dropout(self.attn) - x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) - x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model) - - return self.linear_out(x) # (batch, time1, d_model) - - def forward(self, query, key, value, mask): - """Compute scaled dot product attention. - Args: - query (torch.Tensor): Query tensor (#batch, time1, size). - key (torch.Tensor): Key tensor (#batch, time2, size). - value (torch.Tensor): Value tensor (#batch, time2, size). - mask (torch.Tensor): Mask tensor (#batch, 1, time2) or - (#batch, time1, time2). - Returns: - torch.Tensor: Output tensor (#batch, time1, d_model). - """ - q, k, v = self.forward_qkv(query, key, value) - if version.parse(torch.__version__) >= version.parse("2.0") and self.flash: - n_batch = value.size(0) - x = torch.nn.functional.scaled_dot_product_attention( - q, k, v, attn_mask=mask.unsqueeze(1) if mask is not None else None, dropout_p=self.dropout_rate - ) - x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model) - return self.linear_out(x) - else: - scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) - return self.forward_attention(v, scores, mask) - - -class RelPositionMultiHeadedAttention(MultiHeadedAttention): - """Multi-Head Attention layer with relative position encoding. - Paper: https://arxiv.org/abs/1901.02860 - Args: - n_head (int): The number of heads. - n_feat (int): The number of features. - dropout_rate (float): Dropout rate. - """ - - def __init__(self, n_head, n_feat, dropout_rate): - """Construct an RelPositionMultiHeadedAttention object.""" - super().__init__(n_head, n_feat, dropout_rate) - # linear transformation for positional encoding - self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) - # these two learnable bias are used in matrix c and matrix d - # as described in https://arxiv.org/abs/1901.02860 Section 3.3 - self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k)) - self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k)) - torch.nn.init.xavier_uniform_(self.pos_bias_u) - torch.nn.init.xavier_uniform_(self.pos_bias_v) - - def rel_shift(self, x, zero_triu=False): - """Compute relative positinal encoding. - Args: - x (torch.Tensor): Input tensor (batch, time, size). - zero_triu (bool): If true, return the lower triangular part of the matrix. - Returns: - torch.Tensor: Output tensor. - """ - zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype) - x_padded = torch.cat([zero_pad, x], dim=-1) - - x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2)) - x = x_padded[:, :, 1:].view_as(x) - - if zero_triu: - ones = torch.ones((x.size(2), x.size(3))) - x = x * torch.tril(ones, x.size(3) - x.size(2))[None, None, :, :] - - return x - - def forward(self, query, key, value, pos_emb, mask): - """Compute 'Scaled Dot Product Attention' with rel. positional encoding. - Args: - query (torch.Tensor): Query tensor (#batch, time1, size). - key (torch.Tensor): Key tensor (#batch, time2, size). - value (torch.Tensor): Value tensor (#batch, time2, size). - pos_emb (torch.Tensor): Positional embedding tensor (#batch, time2, size). - mask (torch.Tensor): Mask tensor (#batch, 1, time2) or - (#batch, time1, time2). - Returns: - torch.Tensor: Output tensor (#batch, time1, d_model). - """ - q, k, v = self.forward_qkv(query, key, value) - q = q.transpose(1, 2) # (batch, time1, head, d_k) - - n_batch_pos = pos_emb.size(0) - p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) - p = p.transpose(1, 2) # (batch, head, time1, d_k) - - # (batch, head, time1, d_k) - q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) - # (batch, head, time1, d_k) - q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) - - # compute attention score - # first compute matrix a and matrix c - # as described in https://arxiv.org/abs/1901.02860 Section 3.3 - # (batch, head, time1, time2) - matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) - - # compute matrix b and matrix d - # (batch, head, time1, time2) - matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) - matrix_bd = self.rel_shift(matrix_bd) - - scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) - - return self.forward_attention(v, scores, mask) - - -# ---------------------- note transcription positional encoding ---------------------- - - -class PositionalEncoding(torch.nn.Module): - """Positional encoding. - Args: - d_model (int): Embedding dimension. - max_len (int): Maximum input length. - reverse (bool): Whether to reverse the input position. - """ - - def __init__(self, d_model, max_len=5000, reverse=False): - """Construct an PositionalEncoding object.""" - super().__init__() - self.d_model = d_model - self.reverse = reverse - self.xscale = math.sqrt(self.d_model) - self.pe = None - - self.extend_pe(torch.tensor(0.0).expand(1, max_len)) - - def extend_pe(self, x): - """Reset the positional encodings.""" - if self.pe is not None: - if self.pe.size(1) >= x.size(1): - if self.pe.dtype != x.dtype or self.pe.device != x.device: - self.pe = self.pe.to(dtype=x.dtype, device=x.device) - return - pe = torch.zeros(x.size(1), self.d_model) - if self.reverse: - position = torch.arange(x.size(1) - 1, -1, -1.0, dtype=torch.float32).unsqueeze(1) - else: - position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1) - div_term = torch.exp( - torch.arange(0, self.d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / self.d_model) - ) - pe[:, 0::2] = torch.sin(position * div_term) - pe[:, 1::2] = torch.cos(position * div_term) - pe = pe.unsqueeze(0) - self.pe = pe.to(device=x.device, dtype=x.dtype) - - def forward(self, x: torch.Tensor): - """Add positional encoding. - Args: - x (torch.Tensor): Input tensor (batch, time, `*`). - Returns: - torch.Tensor: Encoded tensor (batch, time, `*`). - """ - self.extend_pe(x) - x = x * self.xscale + self.pe[:, : x.size(1)] - return x - - -class RelPositionalEncoding(PositionalEncoding): - """Relative positional encoding module. - See : Appendix B in https://arxiv.org/abs/1901.02860 - Args: - d_model (int): Embedding dimension. - max_len (int): Maximum input length. - """ - - def __init__(self, d_model, max_len=5000): - super().__init__(d_model, max_len, reverse=True) - - def forward(self, x): - """Compute positional encoding. - Args: - x (torch.Tensor): Input tensor (batch, time, `*`). - Returns: - torch.Tensor: Encoded tensor (batch, time, `*`). - torch.Tensor: Positional embedding tensor (1, time, `*`). - """ - self.extend_pe(x) - x = x * self.xscale - pos_emb = self.pe[:, : x.size(1)] - return x, pos_emb - - -# ---------------------- note transcription conformer layers ---------------------- - - -class ConvolutionModule(nn.Module): - def __init__(self, channels: int, kernel_size: int, activation: nn.Module = nn.ReLU(), bias: bool = True): - super().__init__() - if (kernel_size - 1) % 2 != 0: - raise ValueError("kernel_size should be a odd number for 'SAME' padding") - - self.pointwise_conv1 = nn.Conv1d(channels, 2 * channels, kernel_size=1, bias=bias) - self.depthwise_conv = nn.Conv1d( - channels, channels, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, groups=channels, bias=bias - ) - self.norm = nn.BatchNorm1d(channels) - self.pointwise_conv2 = nn.Conv1d(channels, channels, kernel_size=1, bias=bias) - self.activation = activation - - def forward(self, x): - # (batch, time, channels) -> (batch, channels, time) - x = x.transpose(1, 2) - - x = self.pointwise_conv1(x) - x = nn.functional.glu(x, dim=1) - - x = self.depthwise_conv(x) - x = self.norm(x) - x = self.activation(x) - - x = self.pointwise_conv2(x) - x = x.transpose(1, 2) - return x - - -class MultiLayeredConv1d(nn.Module): - """ - Multi-layered Conv1d that replaces feed-forward in Transformer block. - - Reference: - FastSpeech: Fast, Robust and Controllable Text to Speech - https://arxiv.org/pdf/1905.09263.pdf - """ - - def __init__(self, in_chans, hidden_chans, kernel_size): - super().__init__() - padding = (kernel_size - 1) // 2 - self.w_1 = nn.Conv1d(in_chans, hidden_chans, kernel_size, padding=padding) - self.w_2 = nn.Conv1d(hidden_chans, in_chans, kernel_size, padding=padding) - - def forward(self, x): - """ - Args: - x (torch.Tensor): (batch, time, in_chans) - Returns: - torch.Tensor: (batch, time, in_chans) - """ - # Conv1d expects (batch, channels, time) - x = x.transpose(1, 2) - x = torch.relu(self.w_1(x)) - x = self.w_2(x) - x = x.transpose(1, 2) - return x - - -class Swish(nn.Module): - """Swish activation function.""" - - def forward(self, x): - return x * torch.sigmoid(x) - - -class EncoderLayer(nn.Module): - """ - A single encoder layer: optional macaron FFN, self-attn, conv, FFN, norms. - Args: - size (int): Input dimension. - self_attn (nn.Module): Self-attention module. - feed_forward (nn.Module): Standard feed-forward network. - feed_forward_macaron (nn.Module): Optional, second FFN in macaron structure. - conv_module (nn.Module): Optional, convolution module. - normalize_before (bool): If True, use pre-layer norm. Else, post-layer norm. - concat_after (bool): If True, concat attn in/out, else residual. - """ - - def __init__( - self, - size, - self_attn, - feed_forward, - feed_forward_macaron, - conv_module, - normalize_before=True, - concat_after=False, - ): - super().__init__() - self.size = size - self.self_attn = self_attn - self.feed_forward = feed_forward - self.feed_forward_macaron = feed_forward_macaron - self.conv_module = conv_module - - self.norm_ff = LayerNorm(size) - self.norm_mha = LayerNorm(size) - - self.ff_scale = 0.5 if feed_forward_macaron is not None else 1.0 - self.norm_ff_macaron = LayerNorm(size) if feed_forward_macaron is not None else None - - if conv_module is not None: - self.norm_conv = LayerNorm(size) - self.norm_final = LayerNorm(size) - else: - self.norm_conv = None - self.norm_final = None - - self.normalize_before = normalize_before - self.concat_after = concat_after - if self.concat_after: - self.concat_linear = nn.Linear(2 * size, size) - - def forward(self, x_input, mask, cache=None) -> tuple: - """ - Forward pass for the encoder layer. - - Args: - x_input (Tensor or Tuple): Input tensor of shape (batch, time, size) - or tuple (x, pos_emb) where pos_emb is the positional embedding. - mask (Tensor): Mask tensor of shape (batch, time) or (batch, 1, time). - cache (Tensor, optional): Input cache for streaming (can be None). - - Returns: - Tuple[Tensor, Tensor]: (outputs, mask), where outputs either includes - pos_emb or not, depending on input. - """ - # Unpack possible (x, pos_emb) tuple input - if isinstance(x_input, tuple): - x, pos_emb = x_input - else: - x, pos_emb = x_input, None - - # --- Macaron Feed-Forward (first FFN) --- - if self.feed_forward_macaron is not None: - residual = x - if self.normalize_before: - x = self.norm_ff_macaron(x) - x = residual + self.ff_scale * self.feed_forward_macaron(x) - if not self.normalize_before: - x = self.norm_ff_macaron(x) - - # --- Self-Attention --- - residual = x - if self.normalize_before: - x = self.norm_mha(x) - - # Handle cache for streaming - if cache is not None: - assert cache.shape == (x.shape[0], x.shape[1] - 1, self.size) - x_q = x[:, -1:, :] - residual = residual[:, -1:, :] - mask = mask[:, -1:, :] if mask is not None else None - else: - x_q = x - - # Apply self-attention (with or without positional embeddings) - if pos_emb is not None: - attn_out = self.self_attn(x_q, x, x, pos_emb, mask) - else: - attn_out = self.self_attn(x_q, x, x, mask) - - if self.concat_after: - x = torch.cat((x, attn_out), dim=-1) - x = residual + self.concat_linear(x) - else: - x = residual + attn_out - if not self.normalize_before: - x = self.norm_mha(x) - - # --- Convolution Module --- - if self.conv_module is not None: - residual = x - if self.normalize_before: - x = self.norm_conv(x) - x = residual + self.conv_module(x) - if not self.normalize_before: - x = self.norm_conv(x) - - # --- Feed-Forward (second FFN) --- - residual = x - if self.normalize_before: - x = self.norm_ff(x) - x = residual + self.ff_scale * self.feed_forward(x) - if not self.normalize_before: - x = self.norm_ff(x) - - # --- Final normalization, if present --- - if self.norm_final is not None: - x = self.norm_final(x) - - # --- Append to cache if streaming --- - if cache is not None: - x = torch.cat([cache, x], dim=1) - - # --- Final return, with or without positional embedding --- - if pos_emb is not None: - return (x, pos_emb), mask - else: - return x, mask - - -class ConformerLayers(nn.Module): - def __init__(self, hidden_size, num_layers, kernel_size=9, num_heads=4, use_last_norm=True, save_hidden=False): - super().__init__() - self.use_last_norm = use_last_norm - self.layers = nn.ModuleList() - positionwise_layer = MultiLayeredConv1d - positionwise_layer_args = (hidden_size, hidden_size * 4, 1) - self.pos_embed = RelPositionalEncoding(hidden_size) - self.encoder_layers = nn.ModuleList( - [ - EncoderLayer( - hidden_size, - RelPositionMultiHeadedAttention(num_heads, hidden_size, 0.0), - positionwise_layer(*positionwise_layer_args), - positionwise_layer(*positionwise_layer_args), - ConvolutionModule(hidden_size, kernel_size, Swish()), - ) - for _ in range(num_layers) - ] - ) - if self.use_last_norm: - self.layer_norm = nn.LayerNorm(hidden_size) - else: - self.layer_norm = nn.Linear(hidden_size, hidden_size) - self.save_hidden = save_hidden - if save_hidden: - self.hiddens = [] - - def forward(self, x, padding_mask=None): - """ - - :param x: [B, T, H] - :param padding_mask: [B, T] - :return: [B, T, H] - """ - self.hiddens = [] - nonpadding_mask = x.abs().sum(-1) > 0 - x = self.pos_embed(x) - for layer in self.encoder_layers: - x, mask = layer(x, nonpadding_mask[:, None, :]) - if self.save_hidden: - self.hiddens.append(x[0]) - x = x[0] - x = self.layer_norm(x) * nonpadding_mask.float()[:, :, None] - return x - - -# ---------------------- note transcription backbone layers ---------------------- - - -class UnetDown(nn.Module): - def __init__( - self, - hidden_size, - n_layers, - kernel_size, - down_rates, - channel_multiples=None, - is_btc=True, - constant_channels=False, - ): - super().__init__() - assert n_layers == len(down_rates) # downs, down sample rate - down_rates = [int(i) for i in down_rates] - self.n_layers = n_layers - self.hidden_size = hidden_size - self.is_btc = is_btc - channel_multiples = channel_multiples if channel_multiples is not None else down_rates - self.layers = nn.ModuleList() - self.downs = nn.ModuleList() - in_channels = hidden_size - for i in range(self.n_layers): - out_channels = int(in_channels * channel_multiples[i]) if not constant_channels else in_channels - self.layers.append( - nn.Sequential( - ResidualBlock(in_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5), - nn.Conv1d(in_channels, out_channels, kernel_size, padding=(kernel_size - 1) // 2), - ResidualBlock(out_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5), - ) - ) - self.downs.append(nn.Sequential(nn.AvgPool1d(down_rates[i]))) - in_channels = out_channels - self.last_norm = LayerNorm(out_channels, dim=1, eps=1e-6) - self.post_net = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2) - - def forward(self, x, **kwargs): - # x [B, T, C] - if self.is_btc: - x = x.transpose(1, 2) # [B, C, T] - skip_xs = [] - for i in range(self.n_layers): - skip_x = self.layers[i](x) - x = self.downs[i](skip_x) - if self.is_btc: - skip_xs.append(skip_x.transpose(1, 2)) # [B, T, C] - else: - skip_xs.append(skip_x) - x = self.post_net(self.last_norm(x)) - if self.is_btc: - x = x.transpose(1, 2) - return x, skip_xs - - -class UnetMid(nn.Module): - def __init__(self, hidden_size, kernel_size, n_layers=None, in_dims=None, out_dims=None, is_btc=True, net=None): - super().__init__() - in_dims = in_dims if in_dims is not None else hidden_size - out_dims = out_dims if out_dims is not None else hidden_size - self.pre = nn.Conv1d(in_dims, hidden_size, kernel_size, padding=kernel_size // 2) - self.post = nn.Conv1d(hidden_size, out_dims, kernel_size, padding=kernel_size // 2) - self.is_btc = is_btc - if net is not None: - self.net = net - else: - self.net = ConvBlocks( - hidden_size, - out_dims=hidden_size, - dilations=None, - kernel_size=kernel_size, - layers_in_block=2, - c_multiple=2, - post_net_kernel=3, - is_btc=is_btc, - ) - - def forward(self, x, cond=None, **kwargs): - # x [B, T, C] - if self.is_btc: - x = self.pre(x.transpose(1, 2)).transpose(1, 2) - else: - x = self.pre(x) - if cond is None: - cond = 0 - x = self.net(x + cond) - if self.is_btc: - x = self.post(x.transpose(1, 2)).transpose(1, 2) - else: - x = self.post(x) - return x - - -class UnetUp(nn.Module): - def __init__( - self, - hidden_size, - n_layers, - kernel_size, - up_rates, - channel_multiples=None, - is_btc=True, - constant_channels=False, - use_skip_layer=False, - skip_scale=1.0, - ): - super().__init__() - assert n_layers == len(up_rates) # this is reversed in up module, from the output to the interface with middle - up_rates = [int(i) for i in up_rates] - self.n_layers = n_layers - self.hidden_size = hidden_size - self.is_btc = is_btc - self.skip_scale = skip_scale - channel_multiples = channel_multiples if channel_multiples is not None else up_rates - # in_channels = int(np.cumprod(channel_multiples)[-1] * hidden_size) if not constant_channels else hidden_size - self.in_channels_lst = ( - (np.cumprod([1] + channel_multiples) * hidden_size).astype(int) - if not constant_channels - else [hidden_size for _ in range(self.n_layers + 1)] - ) - in_channels = self.in_channels_lst[-1] - self.ups = nn.ModuleList() - self.skip_layers = nn.ModuleList() - self.layers = nn.ModuleList() - for i in range(self.n_layers - 1, -1, -1): - out_channels = self.in_channels_lst[i] if not constant_channels else in_channels - self.ups.append( - nn.Sequential( - nn.ConvTranspose1d( - in_channels, - in_channels, - kernel_size=kernel_size, - stride=up_rates[i], - padding=kernel_size // 2, - output_padding=up_rates[i] - 1, - ), - LayerNorm(in_channels, dim=1, eps=1e-6), - nn.LeakyReLU(), - ) - ) - self.layers.append( - nn.Sequential( - nn.Conv1d(in_channels * 2, out_channels, kernel_size, padding=(kernel_size - 1) // 2), - ResidualBlock(out_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5), - ) - ) - if use_skip_layer: - self.skip_layers.append( - ResidualBlock(in_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5) - ) - else: - self.skip_layers.append(nn.Identity()) - - in_channels = out_channels - self.out_channels = out_channels - self.last_norm = LayerNorm(out_channels, dim=1, eps=1e-6) - self.post_net = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2) - - def forward(self, x, skips, **kwargs): - # x [B, T, C] - if self.is_btc: - x = x.transpose(1, 2) # [B, C, T] - for i in range(self.n_layers): - x = self.ups[i](x) - skip_x = ( - skips[self.n_layers - i - 1] if not self.is_btc else skips[self.n_layers - i - 1].transpose(1, 2) - ) # [B, T, C] -> [B, C, T] - skip_x = self.skip_layers[i](skip_x) * self.skip_scale - x = torch.cat((x, skip_x), dim=1) # [B, C, T] - x = self.layers[i](x) - x = self.post_net(self.last_norm(x)) - if self.is_btc: - x = x.transpose(1, 2) - return x - - -class Unet(nn.Module): - def __init__( - self, - hidden_size, - down_layers, - up_layers, - kernel_size, - updown_rates, - mid_layers=None, - channel_multiples=None, - is_btc=True, - constant_channels=False, - mid_net=None, - use_skip_layer=False, - skip_scale=1.0, - ): - super().__init__() - assert len(updown_rates) == down_layers == up_layers, f"{len(updown_rates)}, {down_layers}, {up_layers}" - if channel_multiples is not None: - assert len(channel_multiples) == len(updown_rates) - else: - channel_multiples = updown_rates - self.down = UnetDown( - hidden_size, down_layers, kernel_size, updown_rates, channel_multiples, is_btc, constant_channels - ) - down_out_dims = int(np.cumprod(channel_multiples)[-1] * hidden_size) if not constant_channels else hidden_size - self.mid = UnetMid( - hidden_size, - kernel_size, - mid_layers, - in_dims=down_out_dims, - out_dims=down_out_dims, - is_btc=is_btc, - net=mid_net, - ) - self.up = UnetUp( - hidden_size, - up_layers, - kernel_size, - updown_rates, - channel_multiples, - is_btc, - constant_channels, - use_skip_layer, - skip_scale, - ) - - def forward(self, x, mid_cond=None, **kwargs): - x, skips = self.down(x) - x = self.mid(x, mid_cond) - x = self.up(x, skips) - return x - - -class BackboneNet(nn.Module): - def __init__(self, hparams): - super().__init__() - self.hidden_size = hidden_size = hparams["hidden_size"] - updown_rates = [2, 2, 2] - channel_multiples = [1, 1, 1] - if hparams.get("updown_rates", None) is not None: - updown_rates = [int(i) for i in hparams.get("updown_rates", None).split("-")] - if hparams.get("channel_multiples", None) is not None: - channel_multiples = [float(i) for i in hparams.get("channel_multiples", None).split("-")] - assert len(updown_rates) == len(channel_multiples) - # convs - if hparams.get("bkb_net", "conv") == "conv": - self.net = Unet( - hidden_size, - down_layers=len(updown_rates), - mid_layers=hparams.get("bkb_layers", 12), - up_layers=len(updown_rates), - kernel_size=3, - updown_rates=updown_rates, - channel_multiples=channel_multiples, - is_btc=True, - constant_channels=False, - mid_net=None, - use_skip_layer=hparams.get("unet_skip_layer", False), - ) - # conformer - elif hparams.get("bkb_net", "conv") == "conformer": - mid_net = ConformerLayers( - hidden_size, - num_layers=hparams.get("bkb_layers", 12), - kernel_size=hparams.get("conformer_kernel", 9), - num_heads=4, - ) - self.net = Unet( - hidden_size, - down_layers=len(updown_rates), - up_layers=len(updown_rates), - kernel_size=3, - updown_rates=updown_rates, - channel_multiples=channel_multiples, - is_btc=True, - constant_channels=False, - mid_net=mid_net, - use_skip_layer=hparams.get("unet_skip_layer", False), - ) - - def forward(self, x): - return self.net(x) diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/__init__.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/__init__.py deleted file mode 100644 index 9e020f61b..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -"""Preprocess neural network modules for SoulX-Singer.""" - -from .pipeline import SoulXPreprocessPipeline - -__all__ = [ - "SoulXPreprocessPipeline", -] diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/asr.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/asr.py deleted file mode 100644 index 0752c73f8..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/asr.py +++ /dev/null @@ -1,326 +0,0 @@ -"""Lyric transcription (Paraformer + Parakeet) for SoulX preprocess.""" - -import os -import re -import tempfile -from pathlib import Path -from typing import Any - -import numpy as np -import soundfile as sf -import torch -import torch.nn as nn -from omegaconf import open_dict -from vllm.logger import init_logger - -from .utils import get_audio_duration - -logger = init_logger(__name__) - - -def _import_funasr_automodel(): - """Import FunASR without shelling out to ffmpeg at module import time.""" - import subprocess - from functools import wraps - - real_check_output = subprocess.check_output - - @wraps(real_check_output) - def _safe_check_output(cmd, *args, **kwargs): - if cmd and Path(str(cmd[0])).name == "ffmpeg": - raise FileNotFoundError("ffmpeg unavailable for FunASR import probe") - return real_check_output(cmd, *args, **kwargs) - - subprocess.check_output = _safe_check_output - try: - from funasr import AutoModel - finally: - subprocess.check_output = real_check_output - return AutoModel - - -def build_words_with_gaps(raw_words, raw_timestamps, audio_duration: float): - words, word_durs = [], [] - prev = 0.0 - for w, t in zip(raw_words, raw_timestamps): - s, e = float(t[0]), float(t[1]) - if s > prev: - words.append("") - word_durs.append(s - prev) - words.append(w) - word_durs.append(e - s) - prev = e - - if audio_duration > prev: - if not words: - words.append("") - word_durs.append(audio_duration) - elif words[-1] != "": - words.append("") - word_durs.append(audio_duration - prev) - else: - word_durs[-1] += audio_duration - prev - return words, word_durs - - -class ParaformerASR(nn.Module): - """Mandarin/Cantonese ASR via FunASR Paraformer.""" - - def __init__(self, model_path: str, device: str): - super().__init__() - from vllm_omni.diffusion.models.soulx_singer.utils import _patch_torchaudio_load - - _patch_torchaudio_load() - AutoModel = _import_funasr_automodel() - self.asr = AutoModel(model=model_path, disable_update=True, device=device) - logger.info("Loaded Paraformer ASR from %s", model_path) - - def _generate(self, source: str) -> tuple[list[str], list[float], float]: - out = self.asr.generate(source, output_timestamp=True)[0] - raw_words = out["text"].replace("@", "").split(" ") - raw_timestamps = [[t[0] / 1000, t[1] / 1000] for t in out["timestamp"]] - duration = float(raw_timestamps[-1][1]) if raw_timestamps else 0.0 - return raw_words, raw_timestamps, duration - - def _maybe_refine_with_sidecar_f0( - self, - wav_path: str, - words: list[str], - word_durs: list[float], - *, - sample_rate: int, - hop_size: int, - ) -> tuple[list[str], list[float]]: - """Match upstream ``_ASRZhModel.process``: refine using ``{wav}_f0.npy`` if present.""" - f0_path = os.path.splitext(wav_path)[0] + "_f0.npy" - if not os.path.isfile(f0_path): - return words, word_durs - f0 = np.load(f0_path) - if f0.size == 0: - return words, word_durs - return refine_word_durs_with_f0( - words, - word_durs, - f0, - sample_rate=sample_rate, - hop_size=hop_size, - ) - - @torch.inference_mode() - def forward(self, audio: str | np.ndarray, sample_rate: int = 16000) -> tuple[list[str], list[float]]: - if isinstance(audio, str): - raw_words, raw_timestamps, _ = self._generate(audio) - audio_duration = get_audio_duration(audio) - words, word_durs = build_words_with_gaps(raw_words, raw_timestamps, audio_duration) - return self._maybe_refine_with_sidecar_f0( - audio, - words, - word_durs, - sample_rate=24000, - hop_size=480, - ) - - waveform = np.asarray(audio, dtype=np.float32).squeeze() - audio_duration = len(waveform) / float(sample_rate) - with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: - tmp_path = tmp.name - try: - sf.write(tmp_path, waveform, sample_rate) - raw_words, raw_timestamps, _ = self._generate(tmp_path) - words, word_durs = build_words_with_gaps(raw_words, raw_timestamps, audio_duration) - return self._maybe_refine_with_sidecar_f0( - tmp_path, - words, - word_durs, - sample_rate=24000, - hop_size=480, - ) - finally: - Path(tmp_path).unlink(missing_ok=True) - Path(os.path.splitext(tmp_path)[0] + "_f0.npy").unlink(missing_ok=True) - - -def _disable_nemo_cuda_graphs(model) -> None: - decoding_cfg = model.cfg.decoding - with open_dict(decoding_cfg): - if hasattr(decoding_cfg, "greedy") and decoding_cfg.greedy is not None: - decoding_cfg.greedy.use_cuda_graph_decoder = False - if hasattr(decoding_cfg, "beam") and decoding_cfg.beam is not None: - decoding_cfg.beam.allow_cuda_graphs = False - model.change_decoding_strategy(decoding_cfg) - - -class ParakeetASR(nn.Module): - """English ASR via NeMo Parakeet-TDT.""" - - def __init__(self, model_path: str, device: str): - super().__init__() - try: - import nemo.collections.asr as nemo_asr # type: ignore - except Exception as e: # pragma: no cover - raise ImportError( - "NeMo (nemo_toolkit) is required for ASR English but is not available in this Python env." - ) from e - - self.asr = nemo_asr.models.ASRModel.restore_from( - restore_path=model_path, - map_location=device, - ) - self.asr.eval() - _disable_nemo_cuda_graphs(self.asr) - logger.info("Loaded Parakeet ASR from %s", model_path) - - @staticmethod - def _clean_word(word: str) -> str: - return re.sub(r"[\?\.,:]", "", word).strip() - - @staticmethod - def _extract_word_segments(output: Any) -> list[dict[str, Any]]: - ts = getattr(output, "timestamp", None) - if not ts or not isinstance(ts, dict): - return [] - word_ts = ts.get("word") - return word_ts if isinstance(word_ts, list) else [] - - @torch.inference_mode() - def forward(self, audio: str | np.ndarray, sample_rate: int = 16000) -> tuple[list[str], list[float]]: - if isinstance(audio, str): - sources = [audio] - duration = 0.0 - else: - waveform = np.asarray(audio, dtype=np.float32).squeeze() - duration = len(waveform) / float(sample_rate) - with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: - tmp_path = tmp.name - sf.write(tmp_path, waveform, sample_rate) - sources = [tmp_path] - - try: - outputs = self.asr.transcribe( - sources, - timestamps=True, - batch_size=1, - num_workers=0, - ) - finally: - if not isinstance(audio, str): - Path(sources[0]).unlink(missing_ok=True) - - output = outputs[0] if outputs else None - raw_words: list[str] = [] - raw_timestamps: list[list[float]] = [] - if output is not None: - for w in self._extract_word_segments(output): - s, e = float(w.get("start", 0.0)), float(w.get("end", 0.0)) - word = self._clean_word(str(w.get("word", ""))) - if word: - raw_words.append(word) - raw_timestamps.append([s, e]) - - if raw_timestamps: - duration = max(duration, float(raw_timestamps[-1][1])) - return build_words_with_gaps(raw_words, raw_timestamps, duration) - - -def refine_word_durs_with_f0( - words: list[str], - word_durs: list[float], - f0: np.ndarray, - *, - sample_rate: int = 24000, - hop_size: int = 480, -) -> tuple[list[str], list[float]]: - boundaries = np.cumsum([0, *[int(dur * sample_rate / hop_size) for dur in word_durs]]).tolist() - sil_tolerance = 5 - ext_tolerance = 5 - new_words: list[str] = [] - new_word_durs: list[float] = [] - if words: - new_words.append(words[0]) - new_word_durs.append(word_durs[0]) - - for i in range(1, len(words)): - word = words[i] - if word == "": - start_frame = boundaries[i] - end_frame = boundaries[i + 1] - num_frames = end_frame - start_frame - frame_idx = start_frame - unvoiced_count = 0 - while frame_idx < end_frame: - if f0[frame_idx] <= 1: - unvoiced_count += 1 - if unvoiced_count >= sil_tolerance: - frame_idx -= sil_tolerance - 1 - break - else: - unvoiced_count = 0 - frame_idx += 1 - - voice_frames = frame_idx - start_frame - if voice_frames >= int(num_frames * 0.9): - new_word_durs[-1] += word_durs[i] - elif voice_frames >= ext_tolerance: - dur = voice_frames * hop_size / sample_rate - new_word_durs[-1] += dur - new_words.append("") - new_word_durs.append(word_durs[i] - dur) - else: - new_words.append(word) - new_word_durs.append(word_durs[i]) - else: - new_words.append(word) - new_word_durs.append(word_durs[i]) - return new_words, new_word_durs - - -class LyricModel(nn.Module): - def __init__( - self, - zh_model_path: str, - en_model_path: str, - device: str = "cuda", - *, - target_sr: int = 24000, - hop_size: int = 480, - ): - super().__init__() - self.device_str = device - self.en_model_path = en_model_path - self.target_sr = target_sr - self.hop_size = hop_size - self.zh_asr = ParaformerASR(zh_model_path, device=device) - self.en_asr: nn.Module | None = None - - @torch.inference_mode() - def forward( - self, - audio: str | np.ndarray, - language: str = "Mandarin", - *, - sample_rate: int = 16000, - f0: np.ndarray | None = None, - ) -> tuple[list[str], list[float]]: - if language not in {"Mandarin", "Cantonese", "English"}: - raise ValueError(f"Unsupported language: {language}") - - if language.lower() == "english": - if self.en_asr is None: - self.en_asr = ParakeetASR(self.en_model_path, device=self.device_str) - self.add_module("en_asr", self.en_asr) - words, durs = self.en_asr(audio, sample_rate=sample_rate) - else: - # Path-based ASR loads ``{wav}_f0.npy`` inside ParaformerASR (upstream parity). - if isinstance(audio, str): - words, durs = self.zh_asr(audio) - else: - words, durs = self.zh_asr(audio, sample_rate=sample_rate) - if f0 is not None and len(f0) > 0: - words, durs = refine_word_durs_with_f0( - words, - durs, - f0, - sample_rate=self.target_sr, - hop_size=self.hop_size, - ) - return words, durs diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/mel_grid_f0.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/mel_grid_f0.py deleted file mode 100644 index 825c1622f..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/mel_grid_f0.py +++ /dev/null @@ -1,96 +0,0 @@ -"""Resample RMVPE F0 (16 kHz hop) onto the SoulX mel frame grid.""" - -from __future__ import annotations - -import math -import time -from typing import TYPE_CHECKING - -import numpy as np -from scipy.interpolate import interp1d -from vllm.logger import init_logger - -from .rmvpe import RMVPE_SAMPLE_RATE -from .utils import load_mono_audio - -if TYPE_CHECKING: - from .rmvpe import RMVPE - -logger = init_logger(__name__) - -RMVPE_HOP = 160 - - -def interpolate_f0_to_mel_grid( - f0_16k: np.ndarray, - original_length: int, - original_sr: int, - *, - target_sr: int, - hop_size: int, - max_duration: float = 300.0, -) -> np.ndarray: - batch_max_length = int(max_duration * target_sr / hop_size) - duration_in_seconds = original_length / original_sr - effective_target_length = int(duration_in_seconds * target_sr) - original_frames = math.ceil(effective_target_length / hop_size) - target_frames = min(original_frames, batch_max_length) - - t_16k = np.arange(len(f0_16k)) * (RMVPE_HOP / float(original_sr)) - t_target = np.arange(target_frames) * (hop_size / float(target_sr)) - - if len(f0_16k) > 0: - f_interp = interp1d( - t_16k, - f0_16k, - kind="linear", - bounds_error=False, - fill_value=0.0, - assume_sorted=True, - ) - f0 = f_interp(t_target) - else: - f0 = np.zeros(target_frames) - - if len(f0) != target_frames: - f0 = f0[:target_frames] if len(f0) > target_frames else np.pad(f0, (0, target_frames - len(f0))) - return f0.astype(np.float32, copy=False) - - -def extract_f0_file( - rmvpe: RMVPE, - wav_path: str, - *, - target_sr: int, - hop_size: int, - f0_path: str | None = None, - thread: float = 0.03, - max_duration: float = 300.0, - verbose: bool = False, -) -> np.ndarray: - if verbose: - logger.info("[f0] extracting from %s", wav_path) - t0 = time.perf_counter() - - audio, _ = load_mono_audio(wav_path, target_sr=RMVPE_SAMPLE_RATE) - f0_16k = rmvpe.infer_from_audio(audio, sample_rate=RMVPE_SAMPLE_RATE) - f0 = interpolate_f0_to_mel_grid( - np.asarray(f0_16k, dtype=np.float32), - original_length=audio.shape[-1], - original_sr=RMVPE_SAMPLE_RATE, - target_sr=target_sr, - hop_size=hop_size, - max_duration=max_duration, - ) - - if verbose: - voiced_ratio = float(np.mean(f0 > 0)) if len(f0) else 0.0 - logger.info( - "[f0] done: frames=%s voiced_ratio=%.3f time=%.3fs", - len(f0), - voiced_ratio, - time.perf_counter() - t0, - ) - if f0_path is not None: - np.save(f0_path, f0) - return f0 diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py deleted file mode 100644 index cf5f84e7a..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py +++ /dev/null @@ -1,431 +0,0 @@ -"""SoulX preprocess pipeline.""" - -from __future__ import annotations - -import json -import tempfile -from collections.abc import Iterable -from pathlib import Path -from typing import Any, ClassVar - -import numpy as np -import soundfile as sf -import torch -import torch.nn as nn -from vllm.logger import init_logger - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery -from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.stack import SoulXPreprocessStack -from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.utils import load_mono_audio, resample_mono -from vllm_omni.diffusion.models.soulx_singer.preprocess.metadata_utils import ( - SegmentMetadata, - _merge_group, - convert_metadata, -) -from vllm_omni.diffusion.models.soulx_singer.utils import ( - load_config, - load_wav, - preprocess_weight_paths, - resolve_preprocess_weights_root, - resolve_soulx_kind, -) - -logger = init_logger(__name__) - - -def build_svs_prompt_meta( - preprocess: SoulXPreprocessPipeline, - *, - prompt_audio: str | tuple[np.ndarray, int], - prompt_meta_raw: dict[str, Any], - metadata_processor, - vocal_sep: bool | None = None, -) -> dict[str, Any]: - """Build SVS prompt_meta; in-memory prompt wav is trimmed to mel2note length.""" - if isinstance(prompt_audio, str): - return metadata_processor.process(prompt_meta_raw, prompt_audio) - prompt_meta = metadata_processor.process(prompt_meta_raw, None) - prompt_vocal, prompt_sr = preprocess._extract_vocal(prompt_audio, vocal_sep=vocal_sep) - if prompt_sr != metadata_processor.sample_rate: - prompt_vocal = resample_mono( - prompt_vocal, - orig_sr=prompt_sr, - target_sr=metadata_processor.sample_rate, - ) - max_samples = prompt_meta["mel2note"].shape[1] * metadata_processor.hop_size - segment_wav = np.asarray(prompt_vocal[:max_samples], dtype=np.float32) - prompt_meta["wav"] = torch.from_numpy(segment_wav).unsqueeze(0).float().to(metadata_processor.device) - return prompt_meta - - -class SoulXPreprocessPipeline(nn.Module, SupportAudioInput, SupportsComponentDiscovery): - """Lazy-loaded preprocess stack integrated with vLLM-Omni diffusion lifecycle.""" - - support_audio_input: ClassVar[bool] = True - weights_sources: ClassVar[tuple] = () - _dit_modules: ClassVar[list[str]] = [] - _encoder_modules: ClassVar[list[str]] = [ - "stack.vocal_sep", - "stack.lyric", - "stack.rosvot", - ] - _vae_modules: ClassVar[list[str]] = [] - _resident_modules: ClassVar[list[str]] = ["stack.rmvpe"] - - def __init__( - self, - *, - od_config: OmniDiffusionConfig, - prefix: str = "", - vocal_sep: bool = True, - midi_transcribe: bool = True, - max_merge_duration_ms: int = 60000, - verbose: bool = False, - extra_args: dict[str, Any] | None = None, - ) -> None: - super().__init__() - del prefix - self.od_config = od_config - extra_args = extra_args or {} - self.kind = resolve_soulx_kind(od_config.model) - hf_config = load_config(str(Path(od_config.model) / "config.yaml")) - audio_config = hf_config.audio - if self.kind == "svc" and midi_transcribe: - logger.warning("SVC model (config.json): midi_transcribe is ignored; lyric/MIDI transcription is SVS-only.") - midi_transcribe = False - - self.vocal_sep = vocal_sep - self.midi_transcribe = midi_transcribe - self.max_merge_duration_ms = max_merge_duration_ms - self.verbose = verbose - self.target_sr = int(audio_config.sample_rate) - self.hop_size = int(audio_config.hop_size) - self._metadata_processor = None - self._weights_root: Path | None = None - self.stack: SoulXPreprocessStack | None = None - self._ensure_stack(extra_args) - - def _ensure_stack(self, extra_args: dict[str, Any] | None = None) -> None: - extra_args = extra_args or {} - weights_root = resolve_preprocess_weights_root(self.od_config) - if self.stack is not None and self._weights_root == weights_root: - return - weights = preprocess_weight_paths(weights_root) - device = str(get_local_device()) - self.stack = SoulXPreprocessStack( - weights, - device, - target_sr=self.target_sr, - hop_size=self.hop_size, - verbose=self.verbose, - ) - self._weights_root = weights_root - - def forward(self, req) -> DiffusionOutput: - from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( - SOULX_PREPROCESSED_KEY, - ) - from vllm_omni.diffusion.models.soulx_singer.preprocess.pre_process import ( - build_metadata_processor, - build_preprocess_payload, - build_warmup_payload, - is_warmup_request, - ) - from vllm_omni.inputs.data import OmniTextPrompt - - extra_args = dict(getattr(req.sampling_params, "extra_args", None) or {}) - prompt = OmniTextPrompt(prompt=req.prompts[0]) if isinstance(req.prompts[0], str) else req.prompts[0] - info = prompt.get("additional_information") or {} - if isinstance(info, dict): - extra_args.update(dict(info.get("extra_args") or {})) - - self._ensure_stack(extra_args) - if self._metadata_processor is None: - self._metadata_processor = build_metadata_processor(self.od_config) - - device = get_local_device() - if is_warmup_request(req): - payload = build_warmup_payload( - self.kind, - metadata_processor=self._metadata_processor, - device=device, - sample_rate=self.target_sr, - ) - else: - payload = build_preprocess_payload( - self.kind, - prompt=prompt, - extra_args=extra_args, - preprocess=self, - metadata_processor=self._metadata_processor, - device=device, - sample_rate=self.target_sr, - ) - - return DiffusionOutput( - output={ - "payload": {}, - "metadata": {SOULX_PREPROCESSED_KEY: payload}, - }, - to_cpu=True, - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - del weights - self._ensure_stack({}) - return set() - - def _extract_vocal( - self, - audio: str | tuple[np.ndarray, int], - *, - vocal_sep: bool | None = None, - ) -> tuple[np.ndarray, int]: - if self.stack is None: - self._ensure_stack({}) - use_sep = self.vocal_sep if vocal_sep is None else vocal_sep - if use_sep: - return self.stack.extract_vocal(audio) - return load_mono_audio(audio) - - def extract_f0(self, vocal: np.ndarray, sample_rate: int) -> np.ndarray: - if self.stack is None: - self._ensure_stack({}) - return self.stack.extract_f0(vocal, sample_rate) - - def iter_svs_metadata( - self, - audio_source: str | tuple[np.ndarray, int], - *, - language: str = "Mandarin", - vocal_sep: bool | None = None, - max_merge_duration_ms: int | None = None, - ): - """Yield merged SVS metadata dicts one segment at a time.""" - if isinstance(audio_source, tuple): - vocal, sample_rate = audio_source - origin_wav_fn = "" - else: - vocal, sample_rate = self._extract_vocal(str(audio_source), vocal_sep=vocal_sep) - origin_wav_fn = str(audio_source) - - with tempfile.TemporaryDirectory(prefix="soulx_preprocess_") as work_dir: - work = Path(work_dir) - vocal_path = work / "vocal.wav" - sf.write(vocal_path, vocal, sample_rate) - - vocal_f0_path = str(vocal_path).replace(".wav", "_f0.npy") - vocal_f0 = self.stack.extract_f0(str(vocal_path), sample_rate=sample_rate, f0_path=vocal_f0_path) - - if not self.midi_transcribe: - end_ms = int(len(vocal) / sample_rate * 1000) - item = SegmentMetadata( - item_name="full", - wav_fn=str(vocal_path), - language=language, - start_time_ms=0, - end_time_ms=end_ms, - note_text=[""], - note_dur=[end_ms / 1000.0], - note_pitch=[0], - note_type=[1], - ) - yield convert_metadata(item) - return - - segments = self.stack.ensure_segmenter().forward( - vocal, - sample_rate, - vocal_f0, - base_name="vocal", - origin_wav_fn=origin_wav_fn, - verbose=self.verbose, - ) - cut_dir = work / "cut_wavs" - cut_dir.mkdir(parents=True, exist_ok=True) - lyric = self.stack.ensure_lyric() - rosvot = self.stack.ensure_rosvot() - - long_cut_dir = work / "long_cut_wavs" - max_dur = max_merge_duration_ms or self.max_merge_duration_ms - max_gap_ms = 10000 - current_group: list[Any] = [] - current_len = 0 - prev_end = -1 - - def _flush_group() -> dict[str, Any]: - merged_item = _merge_group(vocal, sample_rate, current_group, long_cut_dir) - merged_f0_path = merged_item.wav_fn.replace(".wav", "_f0.npy") - self.stack.extract_f0(merged_item.wav_fn, sample_rate=sample_rate, f0_path=merged_f0_path) - return convert_metadata(merged_item) - - for seg in segments: - seg_key = seg["item_name"] - seg_wav_path = cut_dir / f"{seg_key}.wav" - sf.write(seg_wav_path, seg["wav"], seg["sample_rate"]) - seg_f0_path = str(seg_wav_path).replace(".wav", "_f0.npy") - self.stack.extract_f0(str(seg_wav_path), sample_rate=seg["sample_rate"], f0_path=seg_f0_path) - - words, durs = lyric.forward( - str(seg_wav_path), - language, - sample_rate=seg["sample_rate"], - ) - seg_item = { - "item_name": seg_key, - "wav_fn": str(seg_wav_path), - "start_time_ms": seg["start_time_ms"], - "end_time_ms": seg["end_time_ms"], - "origin_wav_fn": origin_wav_fn or str(vocal_path), - "words": words, - "word_durs": durs, - "language": language, - } - meta_item = rosvot.transcribe(seg_item, segment_info=seg_item, verbose=self.verbose) - - start_time = int(meta_item.get("start_time_ms", seg["start_time_ms"])) - end_time = int(meta_item.get("end_time_ms", seg["end_time_ms"])) - if current_group and ( - start_time - prev_end > max_gap_ms or current_len + end_time - start_time > max_dur - ): - yield _flush_group() - current_group = [] - current_len = 0 - - current_group.append(meta_item) - current_len += end_time - start_time - prev_end = end_time - - if current_group: - yield _flush_group() - - @staticmethod - def build_svs_payload_from_paths( - *, - prompt_metadata_path: str, - target_metadata_path: str, - audio_path: str, - metadata_processor, - ) -> dict[str, Any]: - with open(prompt_metadata_path, encoding="utf-8") as f: - prompt_metadata = json.load(f) - with open(target_metadata_path, encoding="utf-8") as f: - target_metadata = json.load(f) - if not prompt_metadata: - raise ValueError("prompt_metadata is empty") - if not target_metadata: - raise ValueError("target_metadata is empty") - - prompt_meta = metadata_processor.process(prompt_metadata[0], audio_path) - return { - "kind": "svs", - "prompt_meta": prompt_meta, - "target_meta_list": target_metadata, - } - - def build_svs_payload_from_audio( - self, - *, - prompt_audio: str | tuple[np.ndarray, int], - target_audio: str | tuple[np.ndarray, int], - metadata_processor, - language: str = "Mandarin", - vocal_sep: bool | None = None, - prompt_vocal_sep: bool | None = None, - target_vocal_sep: bool | None = None, - prompt_max_merge_duration_ms: int | None = None, - target_max_merge_duration_ms: int | None = None, - ) -> dict[str, Any]: - p_sep = prompt_vocal_sep if prompt_vocal_sep is not None else vocal_sep - t_sep = target_vocal_sep if target_vocal_sep is not None else vocal_sep - p_merge = ( - prompt_max_merge_duration_ms if prompt_max_merge_duration_ms is not None else self.max_merge_duration_ms - ) - t_merge = ( - target_max_merge_duration_ms if target_max_merge_duration_ms is not None else self.max_merge_duration_ms - ) - prompt_list = list( - self.iter_svs_metadata( - prompt_audio, - language=language, - vocal_sep=p_sep, - max_merge_duration_ms=p_merge, - ) - ) - target_list = list( - self.iter_svs_metadata( - target_audio, - language=language, - vocal_sep=t_sep, - max_merge_duration_ms=t_merge, - ) - ) - if not prompt_list or not target_list: - raise ValueError("SVS preprocess produced empty metadata") - - # Prompt wav is trimmed to mel2note length, not the metadata time window. - prompt_meta = build_svs_prompt_meta( - self, - prompt_audio=prompt_audio, - prompt_meta_raw=prompt_list[0], - metadata_processor=metadata_processor, - vocal_sep=vocal_sep, - ) - - return { - "kind": "svs", - "prompt_meta": prompt_meta, - "target_meta_list": target_list, - } - - @staticmethod - def build_svc_payload_from_paths( - *, - prompt_wav_path: str, - target_wav_path: str, - prompt_f0_path: str, - target_f0_path: str, - sample_rate: int, - device: torch.device | str, - ) -> dict[str, Any]: - prompt_wav = load_wav(prompt_wav_path, sample_rate).to(device) - target_wav = load_wav(target_wav_path, sample_rate).to(device) - prompt_f0 = torch.from_numpy(np.load(prompt_f0_path)).unsqueeze(0).to(device) - target_f0 = torch.from_numpy(np.load(target_f0_path)).unsqueeze(0).to(device) - return { - "kind": "svc", - "prompt_wav": prompt_wav, - "target_wav": target_wav, - "prompt_f0": prompt_f0, - "target_f0": target_f0, - } - - def build_svc_payload_from_audio( - self, - *, - prompt_audio: str | tuple[np.ndarray, int], - target_audio: str | tuple[np.ndarray, int], - sample_rate: int, - device: torch.device | str, - vocal_sep: bool | None = None, - ) -> dict[str, Any]: - prompt_vocal, prompt_sr = self._extract_vocal(prompt_audio, vocal_sep=vocal_sep) - target_vocal, target_sr = self._extract_vocal(target_audio, vocal_sep=vocal_sep) - - if prompt_sr != sample_rate: - prompt_vocal = resample_mono(prompt_vocal, orig_sr=prompt_sr, target_sr=sample_rate) - if target_sr != sample_rate: - target_vocal = resample_mono(target_vocal, orig_sr=target_sr, target_sr=sample_rate) - - prompt_f0 = self.extract_f0(prompt_vocal, sample_rate) - target_f0 = self.extract_f0(target_vocal, sample_rate) - - return { - "kind": "svc", - "prompt_wav": torch.from_numpy(prompt_vocal).unsqueeze(0).to(device), - "target_wav": torch.from_numpy(target_vocal).unsqueeze(0).to(device), - "prompt_f0": torch.from_numpy(prompt_f0).unsqueeze(0).to(device), - "target_f0": torch.from_numpy(target_f0).unsqueeze(0).to(device), - } diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rmvpe.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rmvpe.py deleted file mode 100644 index 50ae4b2fc..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rmvpe.py +++ /dev/null @@ -1,593 +0,0 @@ -"""RMVPE F0 extractor. - -All model code is inlined (originally from RickyL-2000/ROSVOT -``modules/pe/rmvpe/``) with two minor substitutions: - -* ``librosa.filters.mel`` → ``torchaudio.functional.melscale_fbanks`` - (already used elsewhere in vllm-omni); -* ``librosa.sequence.viterbi`` → hand-written numpy viterbi (non-default path). -* ``pyworld`` (deprecated) → removed. -""" - -from __future__ import annotations - -import math -from typing import ClassVar - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from torchaudio.functional import melscale_fbanks -from torchaudio.transforms import Resample -from vllm.logger import init_logger - -from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery - -logger = init_logger(__name__) - -RMVPE_SAMPLE_RATE = 16000 -_N_CLASS = 360 -_N_MELS = 128 -_MEL_FMIN = 30 -_MEL_FMAX = 8000 -_WINDOW_LENGTH = 1024 -_CONST = 1997.3794084376191 - - -class ConvBlockRes(nn.Module): - def __init__(self, in_channels: int, out_channels: int, momentum: float = 0.01): - super().__init__() - self.conv = nn.Sequential( - nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), - nn.BatchNorm2d(out_channels, momentum=momentum), - nn.ReLU(), - nn.Conv2d(out_channels, out_channels, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), - nn.BatchNorm2d(out_channels, momentum=momentum), - nn.ReLU(), - ) - if in_channels != out_channels: - self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1)) - self._has_shortcut = True - else: - self._has_shortcut = False - - def forward(self, x: torch.Tensor) -> torch.Tensor: - if self._has_shortcut: - return self.conv(x) + self.shortcut(x) - return self.conv(x) + x - - -class ResEncoderBlock(nn.Module): - def __init__(self, in_channels: int, out_channels: int, kernel_size, n_blocks: int = 1, momentum: float = 0.01): - super().__init__() - self.n_blocks = n_blocks - self.conv = nn.ModuleList() - self.conv.append(ConvBlockRes(in_channels, out_channels, momentum)) - for _ in range(n_blocks - 1): - self.conv.append(ConvBlockRes(out_channels, out_channels, momentum)) - self.kernel_size = kernel_size - if self.kernel_size is not None: - self.pool = nn.AvgPool2d(kernel_size=kernel_size) - - def forward(self, x: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - for i in range(self.n_blocks): - x = self.conv[i](x) - if self.kernel_size is not None: - return x, self.pool(x) - return x - - -class ResDecoderBlock(nn.Module): - def __init__(self, in_channels: int, out_channels: int, stride, n_blocks: int = 1, momentum: float = 0.01): - super().__init__() - out_padding = (0, 1) if stride == (1, 2) else (1, 1) - self.n_blocks = n_blocks - self.conv1 = nn.Sequential( - nn.ConvTranspose2d( - in_channels, - out_channels, - kernel_size=(3, 3), - stride=stride, - padding=(1, 1), - output_padding=out_padding, - bias=False, - ), - nn.BatchNorm2d(out_channels, momentum=momentum), - nn.ReLU(), - ) - self.conv2 = nn.ModuleList() - self.conv2.append(ConvBlockRes(out_channels * 2, out_channels, momentum)) - for _ in range(n_blocks - 1): - self.conv2.append(ConvBlockRes(out_channels, out_channels, momentum)) - - def forward(self, x: torch.Tensor, concat_tensor: torch.Tensor) -> torch.Tensor: - x = self.conv1(x) - x = torch.cat((x, concat_tensor), dim=1) - for i in range(self.n_blocks): - x = self.conv2[i](x) - return x - - -class Encoder(nn.Module): - def __init__( - self, - in_channels: int, - in_size: int, - n_encoders: int, - kernel_size, - n_blocks: int, - out_channels: int = 16, - momentum: float = 0.01, - ): - super().__init__() - self.n_encoders = n_encoders - self.bn = nn.BatchNorm2d(in_channels, momentum=momentum) - self.layers = nn.ModuleList() - self.latent_channels: list[list[int]] = [] - for _ in range(self.n_encoders): - self.layers.append(ResEncoderBlock(in_channels, out_channels, kernel_size, n_blocks, momentum=momentum)) - self.latent_channels.append([out_channels, in_size]) - in_channels = out_channels - out_channels *= 2 - in_size //= 2 - self.out_size = in_size - self.out_channel = out_channels - - def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: - concat_tensors: list[torch.Tensor] = [] - x = self.bn(x) - for i in range(self.n_encoders): - _c, x = self.layers[i](x) - concat_tensors.append(_c) - return x, concat_tensors - - -class Intermediate(nn.Module): - def __init__(self, in_channels: int, out_channels: int, n_inters: int, n_blocks: int, momentum: float = 0.01): - super().__init__() - self.n_inters = n_inters - self.layers = nn.ModuleList() - self.layers.append(ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum)) - for _ in range(n_inters - 1): - self.layers.append(ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - for i in range(self.n_inters): - x = self.layers[i](x) - return x - - -class Decoder(nn.Module): - def __init__(self, in_channels: int, n_decoders: int, stride, n_blocks: int, momentum: float = 0.01): - super().__init__() - self.layers = nn.ModuleList() - self.n_decoders = n_decoders - for _ in range(self.n_decoders): - out_channels = in_channels // 2 - self.layers.append(ResDecoderBlock(in_channels, out_channels, stride, n_blocks, momentum)) - in_channels = out_channels - - def forward(self, x: torch.Tensor, concat_tensors: list[torch.Tensor]) -> torch.Tensor: - for i in range(self.n_decoders): - x = self.layers[i](x, concat_tensors[-1 - i]) - return x - - -class TimbreFilter(nn.Module): - def __init__(self, latent_rep_channels: list[list[int]]): - super().__init__() - self.layers = nn.ModuleList() - for ch in latent_rep_channels: - self.layers.append(ConvBlockRes(ch[0], ch[0])) - - def forward(self, x_tensors: list[torch.Tensor]) -> list[torch.Tensor]: - return [layer(t) for layer, t in zip(self.layers, x_tensors)] - - -class DeepUnet0(nn.Module): - def __init__( - self, - kernel_size, - n_blocks: int, - en_de_layers: int = 5, - inter_layers: int = 4, - in_channels: int = 1, - en_out_channels: int = 16, - ): - super().__init__() - self.encoder = Encoder(in_channels, _N_MELS, en_de_layers, kernel_size, n_blocks, en_out_channels) - self.intermediate = Intermediate( - self.encoder.out_channel // 2, self.encoder.out_channel, inter_layers, n_blocks - ) - self.tf = TimbreFilter(self.encoder.latent_channels) - self.decoder = Decoder(self.encoder.out_channel, en_de_layers, kernel_size, n_blocks) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x, concat_tensors = self.encoder(x) - x = self.intermediate(x) - x = self.decoder(x, concat_tensors) - return x - - -class BiGRU(nn.Module): - def __init__(self, input_features: int, hidden_features: int, num_layers: int): - super().__init__() - self.gru = nn.GRU(input_features, hidden_features, num_layers=num_layers, batch_first=True, bidirectional=True) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.gru(x)[0] - - -class E2E0(nn.Module): - """E2E0 pitch estimator network (DeepUnet0 + CNN + BiGRU + classifier).""" - - def __init__( - self, - n_blocks: int, - n_gru: int, - kernel_size, - en_de_layers: int = 5, - inter_layers: int = 4, - in_channels: int = 1, - en_out_channels: int = 16, - ): - super().__init__() - self.unet = DeepUnet0(kernel_size, n_blocks, en_de_layers, inter_layers, in_channels, en_out_channels) - self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1)) - if n_gru: - self.fc = nn.Sequential( - BiGRU(3 * _N_MELS, 256, n_gru), - nn.Linear(512, _N_CLASS), - nn.Dropout(0.25), - nn.Sigmoid(), - ) - else: - self.fc = nn.Sequential( - nn.Linear(3 * _N_MELS, _N_CLASS), - nn.Dropout(0.25), - nn.Sigmoid(), - ) - - def forward(self, mel: torch.Tensor) -> torch.Tensor: - mel = mel.transpose(-1, -2).unsqueeze(1) - x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2) - x = self.fc(x) - return x - - -class MelSpectrogram(nn.Module): - """Mel-spectrogram extractor matching ROSVOT's MelSpectrogram. - - Log-mel spectrogram via torch.stft + adaptive filterbank. - """ - - def __init__( - self, - n_mel_channels: int = _N_MELS, - sampling_rate: int = RMVPE_SAMPLE_RATE, - win_length: int = _WINDOW_LENGTH, - hop_length: int = 160, - n_fft: int | None = None, - mel_fmin: float = _MEL_FMIN, - mel_fmax: float = _MEL_FMAX, - clamp: float = 1e-5, - ): - super().__init__() - n_fft = win_length if n_fft is None else n_fft - self.hann_window: dict[str, torch.Tensor] = {} - mel_basis = melscale_fbanks( - n_freqs=n_fft // 2 + 1, - f_min=float(mel_fmin), - f_max=float(mel_fmax) if mel_fmax is not None else 8000.0, - n_mels=n_mel_channels, - sample_rate=int(sampling_rate), - norm="slaney", - mel_scale="htk", - ).T # (n_freqs, n_mels) → (n_mels, n_freqs) to match librosa/layernorm convention - self.register_buffer("mel_basis", mel_basis) # (n_mels, n_freqs) - self.n_fft = n_fft - self.hop_length = hop_length - self.win_length = win_length - self.sampling_rate = sampling_rate - self.n_mel_channels = n_mel_channels - self.clamp = clamp - - def forward(self, audio: torch.Tensor, keyshift: float = 0, speed: float = 1, center: bool = True) -> torch.Tensor: - factor = 2 ** (keyshift / 12) - n_fft_new = int(np.round(self.n_fft * factor)) - win_length_new = int(np.round(self.win_length * factor)) - hop_length_new = int(np.round(self.hop_length * speed)) - - key = str(keyshift) + "_" + str(audio.device) - if key not in self.hann_window: - self.hann_window[key] = torch.hann_window(win_length_new).to(audio.device) - - if center: - pad_left = win_length_new // 2 - pad_right = (win_length_new + 1) // 2 - audio = F.pad(audio, (pad_left, pad_right)) - - fft = torch.stft( - audio, - n_fft=n_fft_new, - hop_length=hop_length_new, - win_length=win_length_new, - window=self.hann_window[key], - center=False, - return_complex=True, - ) - magnitude = fft.abs() - - if keyshift != 0: - size = self.n_fft // 2 + 1 - if magnitude.size(1) < size: - magnitude = F.pad(magnitude, (0, 0, 0, size - magnitude.size(1))) - magnitude = magnitude[:, :size, :] * self.win_length / win_length_new - - # register_buffer confuses type checkers; cast to Tensor explicitly - mel_basis: torch.Tensor = self.mel_basis # type: ignore[assignment] - mel_output = torch.matmul(mel_basis, magnitude) - log_mel = torch.log(torch.clamp(mel_output, min=self.clamp)) - return log_mel - - -def to_local_average_f0(hidden: torch.Tensor, center: torch.Tensor | None = None, thread: float = 0.03) -> np.ndarray: - """Decode hidden representation to F0 via local weighted average.""" - idx = torch.arange(_N_CLASS, device=hidden.device)[None, None, :] - idx_cents = idx * 20 + _CONST - if center is None: - center = torch.argmax(hidden, dim=2, keepdim=True) - start = torch.clip(center - 4, min=0) - end = torch.clip(center + 5, max=_N_CLASS) - idx_mask = (idx >= start) & (idx < end) - weights = hidden * idx_mask - product_sum = torch.sum(weights * idx_cents, dim=2) - weight_sum = torch.sum(weights, dim=2) - cents = product_sum / (weight_sum + (weight_sum == 0)) - f0 = 10 * 2 ** (cents / 1200) - uv = hidden.max(dim=2)[0] < thread - f0 = f0 * ~uv - return f0.cpu().numpy() - - -def numpy_viterbi(prob: np.ndarray, transition: np.ndarray) -> np.ndarray: - """Viterbi decoding on log-probabilities (numpy). - - Args: - prob: Emission probabilities, shape (n_states, n_obs). - transition: Transition matrix, shape (n_states, n_states), row-stochastic. - - Returns: - path: Most likely state sequence, shape (n_obs,). - """ - n_states, n_obs = prob.shape - log_emit = np.log(prob + 1e-12) - log_trans = np.log(transition + 1e-12) - - delta = log_emit[:, 0] + np.zeros(n_states) # initial: log(pi) omitted (uniform) - psi = np.zeros((n_obs, n_states), dtype=np.int64) - - for t in range(1, n_obs): - scores = delta[:, None] + log_trans # (n_states, n_states) - psi[t] = np.argmax(scores, axis=0) - delta = scores[psi[t], np.arange(n_states)] + log_emit[:, t] - - # Backtrack - path = np.zeros(n_obs, dtype=np.int64) - path[-1] = np.argmax(delta) - for t in range(n_obs - 2, -1, -1): - path[t] = psi[t + 1, path[t + 1]] - return path - - -def to_viterbi_f0(hidden: torch.Tensor, thread: float = 0.03) -> np.ndarray: - """Decode hidden representation to F0 via Viterbi. - - Uses a hand-written numpy Viterbi decoder (replacing - ``librosa.sequence.viterbi``) to avoid the librosa dependency. - """ - if not hasattr(to_viterbi_f0, "_transition"): - xx, yy = np.meshgrid(range(_N_CLASS), range(_N_CLASS)) - transition = np.maximum(30 - abs(xx - yy), 0) - transition = transition / transition.sum(axis=1, keepdims=True) - to_viterbi_f0._transition = transition # type: ignore[attr-defined] - - prob = hidden.squeeze(0).cpu().numpy() - prob = prob.T - prob = prob / prob.sum(axis=0, keepdims=True) - - path = numpy_viterbi(prob, to_viterbi_f0._transition) # type: ignore[attr-defined] - center = torch.from_numpy(path).unsqueeze(0).unsqueeze(-1).to(hidden.device) - return to_local_average_f0(hidden, center=center, thread=thread) - - -def resample_align_curve( - points: np.ndarray, original_timestep: float, target_timestep: float, align_length: int = -1 -) -> np.ndarray: - """Align a curve from one time grid to another via linear interpolation.""" - t_max = (len(points) - 1) * original_timestep - curve_interp = np.interp( - np.arange(0, t_max, target_timestep), - original_timestep * np.arange(len(points)), - points, - ).astype(points.dtype) - if align_length > 0: - if len(curve_interp) > align_length: - curve_interp = curve_interp[:align_length] - elif len(curve_interp) < align_length: - curve_interp = np.concatenate( - [curve_interp, np.full(align_length - len(curve_interp), fill_value=curve_interp[-1])], - axis=0, - ) - return curve_interp - - -class RMVPE(nn.Module, SupportAudioInput, SupportsComponentDiscovery): - """RMVPE F0 extractor (SoulX-Singer). - - Wraps the E2E0 model with mel-spectrogram front-end and decoding - post-processing. - """ - - support_audio_input: ClassVar[bool] = True - _dit_modules: ClassVar[list[str]] = [] - _encoder_modules: ClassVar[list[str]] = [] - _vae_modules: ClassVar[list[str]] = [] - _resident_modules: ClassVar[list[str]] = ["."] - _layerwise_offload_blocks_attrs: ClassVar[list[str]] = [] - - def __init__(self, model_path: str | None = None, hop_length: int = 160, device: str | None = None): - super().__init__() - self.resample_kernel: dict[str, Resample] = {} - if device is None: - self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - else: - self.device = torch.device(device) - - self.model = E2E0(4, 1, (2, 2)).eval() - if model_path is not None: - ckpt = torch.load(model_path, map_location=self.device) - state_dict = ckpt if not isinstance(ckpt, dict) else ckpt.get("model", ckpt) - self.model.load_state_dict(state_dict, strict=False) - logger.info("Loaded RMVPE weights from %s", model_path) - self.model = self.model.to(self.device) - - self.mel_extractor = MelSpectrogram( - _N_MELS, - RMVPE_SAMPLE_RATE, - _WINDOW_LENGTH, - hop_length, - None, - _MEL_FMIN, - _MEL_FMAX, - ).to(self.device) - self.hop_length = hop_length - - # ---- internal helpers ---- - - @torch.no_grad() - def mel2hidden(self, mel: torch.Tensor) -> torch.Tensor: - n_frames = mel.shape[-1] - pad_len = 32 * ((n_frames - 1) // 32 + 1) - n_frames - mel = F.pad(mel, (0, pad_len), mode="constant") - hidden = self.model(mel) - return hidden[:, :n_frames] - - def decode(self, hidden: torch.Tensor, thread: float = 0.03, use_viterbi: bool = False) -> np.ndarray: - if use_viterbi: - return to_viterbi_f0(hidden, thread=thread) - return to_local_average_f0(hidden, thread=thread) - - def postprocess(self, f0: np.ndarray, fmin: float = 50, fmax: float = 1000, min_gap: int = 2) -> np.ndarray: - f0[f0 < fmin] = 0 - f0[f0 > fmax] = 0 - # Eliminate glitch: if positive f0 values span fewer than min_gap frames - # between zeros, zero them out. - for idx in range(f0.shape[0] - min_gap - 1): - if f0[idx] == 0 and f0[idx + min_gap + 1] == 0 and np.sum(f0[idx : idx + min_gap + 2]) > 0: - f0[idx : idx + min_gap + 2] = 0 - return f0 - - # ---- single audio inference ---- - - def infer_from_audio( - self, audio: np.ndarray, sample_rate: int = 16000, thread: float = 0.03, use_viterbi: bool = False - ) -> np.ndarray: - audio_t = torch.from_numpy(audio).float().unsqueeze(0).to(self.device) - if sample_rate == RMVPE_SAMPLE_RATE: - audio_res = audio_t - else: - key_str = str(sample_rate) - if key_str not in self.resample_kernel: - self.resample_kernel[key_str] = Resample(sample_rate, RMVPE_SAMPLE_RATE, lowpass_filter_width=128) - self.resample_kernel[key_str] = self.resample_kernel[key_str].to(self.device) - audio_res = self.resample_kernel[key_str](audio_t) - mel = self.mel_extractor(audio_res, center=True) - hidden = self.mel2hidden(mel) - f0 = self.decode(hidden, thread=thread, use_viterbi=use_viterbi).squeeze(0) - return f0 - - def get_pitch( - self, - waveform, - sample_rate: int, - hop_size: int, - length: int, - interp_uv: bool = False, - fmin: float = 50, - fmax: float = 1000, - ) -> tuple[np.ndarray, np.ndarray]: - f0 = self.infer_from_audio(waveform, sample_rate=sample_rate) - f0 = self.postprocess(f0, fmin, fmax) - uv = f0 == 0 - time_step = hop_size / sample_rate - f0_res = resample_align_curve(f0, 0.01, time_step, length) - uv_res = resample_align_curve(uv.astype(np.float32), 0.01, time_step, length) > 0.5 - if not interp_uv: - f0_res[uv_res] = 0 - return f0_res, uv_res - - # ---- batch audio inference ---- - - def infer_from_audio_batch( - self, audios, sample_rate: int = 16000, thread: float = 0.03, use_viterbi: bool = False - ) -> list[np.ndarray]: - sizes: list[int] | None = None - if isinstance(audios, list): - audios = [torch.from_numpy(a).float() for a in audios] - sizes = [math.ceil((a.shape[0] + 1) / self.hop_length) for a in audios] - audios = nn.utils.rnn.pad_sequence(audios, batch_first=True, padding_value=0.0).to(self.device) - elif isinstance(audios, torch.Tensor): - if audios.device != self.device: - audios = audios.to(self.device) - else: - raise NotImplementedError(f"Unsupported audio type: {type(audios)}") - - if sample_rate == RMVPE_SAMPLE_RATE: - audios_res = audios - else: - key_str = str(sample_rate) - if key_str not in self.resample_kernel: - self.resample_kernel[key_str] = Resample(sample_rate, RMVPE_SAMPLE_RATE, lowpass_filter_width=128) - self.resample_kernel[key_str] = self.resample_kernel[key_str].to(self.device) - audios_res = self.resample_kernel[key_str](audios) - - mels = self.mel_extractor(audios_res, center=True) - hiddens = self.mel2hidden(mels) - f0 = self.decode(hiddens, thread=thread, use_viterbi=use_viterbi) - - f0s: list[np.ndarray] = [] - for i in range(f0.shape[0]): - f0s.append(f0[i, : sizes[i]] if sizes is not None else f0[i, :]) - return f0s - - def get_pitch_batch( - self, - waveforms, - sample_rate: int, - hop_size: int, - lengths: list[int], - interp_uv: bool = False, - fmin: float = 50, - fmax: float = 1000, - ) -> tuple[list[np.ndarray], list[np.ndarray]]: - f0s = self.infer_from_audio_batch(waveforms, sample_rate=sample_rate) - f0s_res: list[np.ndarray] = [] - uvs_res: list[np.ndarray] = [] - for idx, f0 in enumerate(f0s): - f0 = self.postprocess(f0, fmin, fmax, min_gap=6) - uv = f0 == 0 - length = lengths[idx] - time_step = hop_size / sample_rate - f0_res = resample_align_curve(f0, 0.01, time_step, length) - uv_res = resample_align_curve(uv.astype(np.float32), 0.01, time_step, length) > 0.5 - if not interp_uv: - f0_res[uv_res] = 0 - f0s_res.append(f0_res) - uvs_res.append(uv_res) - return f0s_res, uvs_res - - def release_cuda(self): - self.model = self.model.cpu() - self.mel_extractor = self.mel_extractor.cpu() diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rosvot.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rosvot.py deleted file mode 100644 index 0333b521e..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rosvot.py +++ /dev/null @@ -1,619 +0,0 @@ -"""ROSVOT note transcription.""" - -from __future__ import annotations - -import math -from pathlib import Path -from typing import Any, ClassVar - -import numpy as np -import torch -import torch.nn as nn -from vllm.logger import init_logger - -from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery -from vllm_omni.diffusion.models.soulx_singer.modules.mel_transform import spectral_normalize_torch -from vllm_omni.diffusion.models.soulx_singer.modules.note_transcription import ( - BackboneNet, - ConvBlocks, - Embedding, -) -from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.utils import ( - boundary2Interval, - denorm_f0, - f0_to_coarse, - load_model_ckpt, - load_mono_audio, - load_rosvot_config, - norm_interp_f0, - pad_or_cut_xd, - resample_mono, -) -from vllm_omni.utils.audio import mel_filter_bank as librosa_mel_fn - -logger = init_logger(__name__) - - -def regulate_real_note_itv( - note_itv: np.ndarray, - note_bd: np.ndarray, - word_bd: np.ndarray, - word_durs: np.ndarray, - hop_size: int, - audio_sample_rate: int, -) -> tuple[np.ndarray, np.ndarray]: - assert note_itv.shape[0] == np.sum(note_bd) + 1 - assert np.sum(word_bd) <= np.sum(note_bd) - assert word_durs.shape[0] == np.sum(word_bd) + 1, f"{word_durs.shape[0]} {np.sum(word_bd) + 1}" - word_bd = np.cumsum(word_bd) * word_bd - word_itv = np.zeros((word_durs.shape[0], 2)) - word_offsets = np.cumsum(word_durs) - note2words = np.zeros(note_itv.shape[0], dtype=int) - for idx in range(len(word_offsets) - 1): - word_itv[idx, 1] = word_itv[idx + 1, 0] = word_offsets[idx] - word_itv[-1, 1] = word_offsets[-1] - note_itv_secs = note_itv * hop_size / audio_sample_rate - for idx, itv in enumerate(note_itv): - start_idx, end_idx = itv - if word_bd[start_idx] > 0: - word_dur_idx = word_bd[start_idx] - note_itv_secs[idx, 0] = word_itv[word_dur_idx, 0] - note2words[idx] = word_dur_idx - if word_bd[end_idx] > 0: - word_dur_idx = word_bd[end_idx] - 1 - note_itv_secs[idx, 1] = word_itv[word_dur_idx, 1] - note2words[idx] = word_dur_idx - note2words += 1 - return note_itv_secs, note2words - - -def regulate_ill_slur( - notes: np.ndarray, note_itv: np.ndarray, note2words: np.ndarray -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - res_note2words: list[int] = [] - res_note_itv: list[list[float]] = [] - res_notes: list[int] = [] - note_idx = 0 - note_idx_end = 0 - while note_idx <= len(notes) - 1: - while note_idx <= note_idx_end < len(notes) and note2words[note_idx] == note2words[note_idx_end]: - note_idx_end += 1 - res_note2words.append(note2words[note_idx]) - res_note_itv.append(note_itv[note_idx].tolist()) - res_notes.append(notes[note_idx]) - for idx in range(note_idx + 1, note_idx_end): - if notes[idx] == notes[idx - 1]: - res_note_itv[-1][1] = note_itv[idx][1] - else: - res_note_itv.append(note_itv[idx].tolist()) - res_note2words.append(note2words[idx]) - res_notes.append(notes[idx]) - note_idx = note_idx_end - return ( - np.array(res_notes, dtype=notes.dtype), - np.array(res_note_itv, dtype=note_itv.dtype), - np.array(res_note2words, dtype=note2words.dtype), - ) - - -def regulate_boundary(bd_logits, threshold, min_gap=18, ref_bd=None, ref_bd_min_gap=8, non_padding=None): - # this doesn't preserve gradient - device = bd_logits.device - bd_logits = torch.sigmoid(bd_logits).data.cpu() - # bd_logits[0] = bd_logits[-1] = 1e-5 # avoid itv invalid problem - bd = (bd_logits > threshold).long() - bd_res = torch.zeros_like(bd).long() - for i in range(bd.shape[0]): - bd_i = bd[i] - last_bd_idx = -1 - start = -1 - for j in range(bd_i.shape[0]): - if bd_i[j] == 1: - if 0 <= start < j: - continue - elif start < 0: - start = j - else: - if 0 <= start < j: - if j - 1 > start: - bd_idx = start + int(torch.argmax(bd_logits[i, start:j]).item()) - else: - bd_idx = start - if bd_idx - last_bd_idx < min_gap and last_bd_idx > 0: - bd_idx = round((bd_idx + last_bd_idx) / 2) - bd_res[i, last_bd_idx] = 0 - bd_res[i, bd_idx] = 1 - last_bd_idx = bd_idx - start = -1 - - # assert ref_bd_min_gap <= min_gap // 2 - if ref_bd is not None and ref_bd_min_gap > 0: - ref = ref_bd.data.cpu() - for i in range(bd_res.shape[0]): - ref_bd_i = ref[i] - ref_bd_i_js = [] - for j in range(ref_bd_i.shape[0]): - if ref_bd_i[j] == 1: - ref_bd_i_js.append(j) - seg_sum = torch.sum(bd_res[i, max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap]) - if seg_sum == 0: - bd_res[i, j] = 1 - elif seg_sum == 1 and bd_res[i, j] != 1: - bd_res[i, max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap] = ref_bd_i[ - max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap - ] - elif seg_sum > 1: - for k in range(1, ref_bd_min_gap + 1): - if bd_res[i, max(0, j - k)] == 1 and ref_bd_i[max(0, j - k)] != 1: - bd_res[i, max(0, j - k)] = 0 - break - if ( - bd_res[i, min(bd_res.shape[1] - 1, j + k)] == 1 - and ref_bd_i[min(bd_res.shape[1] - 1, j + k)] != 1 - ): - bd_res[i, min(bd_res.shape[1] - 1, j + k)] = 0 - break - bd_res[i, j] = 1 - # final check - assert torch.sum(bd_res[i, ref_bd_i_js]) == len(ref_bd_i_js), ( - f"{torch.sum(bd_res[i, ref_bd_i_js])} {len(ref_bd_i_js)}" - ) - - bd_res = bd_res.to(device) - - # force valid begin and end - bd_res[:, 0] = 0 - if non_padding is not None: - for i in range(bd_res.shape[0]): - bd_res[i, sum(non_padding[i]) - 1 :] = 0 - else: - bd_res[:, -1] = 0 - - return bd_res - - -def _align_word(word_durs: list[float], mel_len: int, hop_size: int, audio_sample_rate: int) -> np.ndarray: - mel2word = np.zeros([mel_len], int) - start_time = 0.0 - for i_word, wd in enumerate(word_durs): - start_frame = int(start_time * audio_sample_rate / hop_size + 0.5) - end_frame = int((start_time + wd) * audio_sample_rate / hop_size + 0.5) - mel2word[start_frame:end_frame] = i_word + 1 - start_time += wd - return mel2word - - -class PitchDecoder(nn.Module): - def __init__(self, hparams): - super().__init__() - self.hidden_size = hidden_size = hparams["hidden_size"] - self.dropout = hparams.get("dropout", 0.0) - self.note_bd_out = nn.Linear(hidden_size, 1) - self.note_bd_temperature = max(1e-7, hparams.get("note_bd_temperature", 1.0)) - - # note prediction - self.pitch_attn_num_head = hparams.get("pitch_attn_num_head", 1) - self.multihead_dot_attn = nn.Linear(hidden_size, self.pitch_attn_num_head) - self.post = ConvBlocks( - hidden_size, - out_dims=hidden_size, - dilations=None, - kernel_size=3, - layers_in_block=1, - c_multiple=1, - post_net_kernel=3, - ) - self.pitch_out = nn.Linear(hidden_size, hparams.get("note_num", 100) + 4) - self.note_num = hparams.get("note_num", 100) - self.note_start = hparams.get("note_start", 30) - self.pitch_temperature = max(1e-7, hparams.get("note_pitch_temperature", 1.0)) - - def forward(self, feat, note_bd, train=True): - bsz, T, _ = feat.shape - - attn = torch.sigmoid(self.multihead_dot_attn(feat)) # [B, T, C] -> [B, T, num_head] - attn_feat = feat.unsqueeze(3) * attn.unsqueeze(2) # [B, T, C, 1] x [B, T, 1, num_head] -> [B, T, C, num_head] - attn_feat = torch.mean(attn_feat, dim=-1) # [B, T, C, num_head] -> [B, T, C] - mel2note = torch.cumsum(note_bd, 1) - note_length = torch.max(torch.sum(note_bd, dim=1)).item() + 1 # max length - note_lengths = torch.sum(note_bd, dim=1) + 1 # [B] - - attn = torch.mean(attn, dim=-1, keepdim=True) # [B, T, num_head] -> [B, T, 1] - denom = mel2note.new_zeros(bsz, note_length, dtype=attn.dtype).scatter_add_( - dim=1, index=mel2note, src=attn.squeeze(-1) - ) # [B, T] -> [B, note_length] count the note frames of each note (with padding excluded) - frame2note = mel2note.unsqueeze(-1).repeat(1, 1, self.hidden_size) # [B, T] -> [B, T, C], with padding included - note_aggregate = frame2note.new_zeros(bsz, note_length, self.hidden_size, dtype=attn_feat.dtype).scatter_add_( - dim=1, index=frame2note, src=attn_feat - ) # [B, T, C] -> [B, note_length, C] - note_aggregate = note_aggregate / (denom.unsqueeze(-1) + 1e-5) - note_logits = self.post(note_aggregate) - note_logits = self.pitch_out(note_logits) / self.pitch_temperature - # note_logits = torch.clamp(note_logits, min=-16., max=16.) # don't know need it or not - - note_pred = torch.softmax(note_logits, dim=-1) # [B, note_length, note_num] - note_pred = torch.argmax(note_pred, dim=-1) # [B, note_length] - # for some reason, note idx maybe 130 (why?) - note_pred[note_pred > self.note_num] = 0 - note_pred[note_pred < self.note_start] = 0 - - return note_lengths, note_logits, note_pred - - -class MidiExtractor(nn.Module): - def __init__(self, hparams): - super().__init__() - self.hparams = hparams.copy() - self.hidden_size = hidden_size = hparams["hidden_size"] - self.note_bd_threshold = hparams.get("note_bd_threshold", 0.5) - self.note_bd_min_gap = round( - hparams.get("note_bd_min_gap", 100) * hparams["audio_sample_rate"] / 1000 / hparams["hop_size"] - ) - self.note_bd_ref_min_gap = round( - hparams.get("note_bd_ref_min_gap", 50) * hparams["audio_sample_rate"] / 1000 / hparams["hop_size"] - ) - - self.mel_proj = nn.Conv1d(hparams["use_mel_bins"], hidden_size, kernel_size=3, padding=1) - self.mel_encoder = ConvBlocks( - hidden_size, - out_dims=hidden_size, - dilations=None, - kernel_size=3, - layers_in_block=2, - c_multiple=1, - post_net_kernel=3, - ) - self.use_pitch = hparams.get("use_pitch_embed", True) - if self.use_pitch: - self.pitch_embed = Embedding(300, hidden_size, 0, "kaiming") - self.uv_embed = Embedding(3, hidden_size, 0, "kaiming") - self.use_wbd = hparams.get("use_wbd", True) - if self.use_wbd: - self.word_bd_embed = Embedding(3, hidden_size, 0, "kaiming") - self.cond_encoder = ConvBlocks( - hidden_size, - out_dims=hidden_size, - dilations=None, - kernel_size=3, - layers_in_block=1, - c_multiple=1, - post_net_kernel=3, - ) - - # backbone - self.net = BackboneNet(hparams) - - # note bd prediction - self.note_bd_out = nn.Linear(hidden_size, 1) - self.note_bd_temperature = max(1e-7, hparams.get("note_bd_temperature", 1.0)) - - # note prediction - self.pitch_decoder = PitchDecoder(hparams) - - def run_encoder(self, mel, word_bd=None, pitch=None, uv=None): - mel_embed = self.mel_proj(mel.transpose(1, 2)).transpose(1, 2) - mel_embed = self.mel_encoder(mel_embed) - pitch_embed = word_bd_embed = 0 - if self.use_pitch and pitch is not None and uv is not None: - pitch_embed = self.pitch_embed(pitch) + self.uv_embed(uv) # [B, T, C] - if self.use_wbd and word_bd is not None: - word_bd_embed = self.word_bd_embed(word_bd) - feat = self.cond_encoder(mel_embed + pitch_embed + word_bd_embed) - return feat - - def forward(self, mel, word_bd=None, note_bd=None, pitch=None, uv=None, non_padding=None, train=True): - ret = {} - bsz, T, _ = mel.shape - - feat = self.run_encoder(mel, word_bd, pitch, uv) - feat = self.net(feat) # [B, T, C] - - # note bd prediction - # dropout has been dropped (inference mode) - note_bd_logits = self.note_bd_out(feat).squeeze(-1) / self.note_bd_temperature - note_bd_logits = torch.clamp(note_bd_logits, min=-16.0, max=16.0) - ret["note_bd_logits"] = note_bd_logits # [B, T] - if note_bd is None or not train: - note_bd = regulate_boundary( - note_bd_logits, - self.note_bd_threshold, - self.note_bd_min_gap, - word_bd, - self.note_bd_ref_min_gap, - non_padding, - ) - ret["note_bd_pred"] = note_bd # [B, T] - - # note pitch prediction - note_lengths, note_logits, note_pred = self.pitch_decoder(feat, note_bd, train) - ret["note_lengths"], ret["note_logits"], ret["note_pred"] = note_lengths, note_logits, note_pred - - return ret - - -class MelNet(nn.Module): - def __init__(self, hparams, device="cpu") -> None: - super().__init__() - self.n_fft = hparams["fft_size"] - self.num_mels = hparams["audio_num_mel_bins"] - self.sampling_rate = hparams["audio_sample_rate"] - self.hop_size = hparams["hop_size"] - self.win_size = hparams["win_size"] - self.fmin = hparams["fmin"] - self.fmax = hparams["fmax"] - self.device = device - - mel = librosa_mel_fn( - sr=self.sampling_rate, - n_fft=self.n_fft, - n_mels=self.num_mels, - fmin=self.fmin, - fmax=self.fmax, - ) - self.mel_basis = mel.float().to(self.device) - self.hann_window = torch.hann_window(self.win_size).to(self.device) - - def to(self, device, **kwagrs): - super().to(device=device, **kwagrs) - self.mel_basis = self.mel_basis.to(device) - self.hann_window = self.hann_window.to(device) - self.device = device - - def forward(self, y, center=False, complex=False): - if isinstance(y, np.ndarray): - y = torch.FloatTensor(y) - if len(y.shape) == 1: - y = y.unsqueeze(0) - y = y.clamp(min=-1.0, max=1.0).to(self.device) - - pad_length = math.ceil(y.shape[1] / self.hop_size) * self.hop_size - y.shape[1] - y = torch.nn.functional.pad( - y.unsqueeze(1), - [int((self.n_fft - self.hop_size) / 2), int((self.n_fft - self.hop_size) / 2 + pad_length)], - mode="reflect", - ) - y = y.squeeze(1) - - spec = torch.stft( - y, - self.n_fft, - hop_length=self.hop_size, - win_length=self.win_size, - window=self.hann_window, - center=center, - pad_mode="reflect", - normalized=False, - onesided=True, - return_complex=True, - ) - if not complex: - spec = torch.view_as_real(spec) - spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) # [B, n_fft, T] - spec = torch.matmul(self.mel_basis, spec) - spec = spectral_normalize_torch(spec) - spec = spec.transpose(1, 2) # [B, T, n_fft] - else: - B, C, T, _ = spec.shape - spec = spec.transpose(1, 2) # [B, T, n_fft, 2] - return spec - - -class RosvotModel(nn.Module, SupportAudioInput, SupportsComponentDiscovery): - support_audio_input: ClassVar[bool] = True - _dit_modules: ClassVar[list[str]] = [] - _encoder_modules: ClassVar[list[str]] = ["."] - _vae_modules: ClassVar[list[str]] = [] - _resident_modules: ClassVar[list[str]] = [] - _layerwise_offload_blocks_attrs: ClassVar[list[str]] = ["midi.net", "midi.pitch_decoder", "midi.cond_encoder"] - - def __init__( - self, - rosvot_ckpt: str | Path, - *, - config_path: str | Path = "", - pe: nn.Module | None = None, - the: float = 0.85, - verbose: bool = False, - rosvot_source_dir: str | Path | None = None, - ): - super().__init__() - self.verbose = verbose - ckpt = Path(rosvot_ckpt) - resolved_config = Path(config_path) if config_path else ckpt.with_name("config.yaml") - self.hparams = load_rosvot_config(resolved_config, hparams_str=f"note_bd_threshold={the}") - if verbose: - logger.info("ROSVOT config: %s", resolved_config) - - self.midi = MidiExtractor(self.hparams) - self.mel_net = MelNet(self.hparams) - self.pe = pe if pe is not None and self.hparams.get("use_pitch_embed", False) else None - self._checkpoint_path = str(ckpt) - self.load_checkpoint(str(ckpt), verbose=verbose) - self.eval() - - def load_checkpoint(self, checkpoint_path: str | None = None, *, verbose: bool = False) -> None: - load_model_ckpt(self.midi, checkpoint_path or self._checkpoint_path, verbose=verbose) - - @torch.no_grad() - def forward(self, wav: torch.Tensor, word_durs: list[float]) -> dict[str, Any]: - hparams = self.hparams - if wav.ndim == 1: - wav = wav.unsqueeze(0) - - mel_len = (wav.shape[-1] + hparams["hop_size"] - 1) // hparams["hop_size"] - min_word_dur = hparams.get("min_word_dur", 20) / 1000 - wd_raw = list(word_durs) - word_durs_filtered: list[float] = [] - for i, wd in enumerate(wd_raw): - if wd < min_word_dur: - if i == 0 and len(wd_raw) > 1: - wd_raw[i + 1] += wd - elif word_durs_filtered: - word_durs_filtered[-1] += wd - else: - word_durs_filtered.append(wd) - - mel2word = _align_word(word_durs_filtered, mel_len, hparams["hop_size"], hparams["audio_sample_rate"]) - if mel2word.size > 0 and mel2word[0] == 0: - mel2word = mel2word + 1 - - real_len = min(mel_len, int(np.sum(mel2word > 0))) - T = math.ceil(min(real_len, hparams["max_frames"]) / hparams["frames_multiple"]) * hparams["frames_multiple"] - device = wav.device - self.mel_net.to(device) - wav_t = pad_or_cut_xd(wav.float(), T * hparams["hop_size"], 1) - - pitch_coarse = uv_t = None - if self.pe is not None: - f0s, uvs = self.pe.get_pitch_batch( - wav_t, - sample_rate=hparams["audio_sample_rate"], - hop_size=hparams["hop_size"], - lengths=[real_len], - fmax=hparams["f0_max"], - fmin=hparams["f0_min"], - ) - f0_1d, uv_1d = norm_interp_f0(f0s[0][:T]) - f0_t = pad_or_cut_xd(torch.as_tensor(f0_1d, device=device, dtype=torch.float32), T, 0).unsqueeze(0) - uv_t = pad_or_cut_xd(torch.as_tensor(uv_1d, device=device, dtype=torch.float32), T, 0).long().unsqueeze(0) - pitch_coarse = f0_to_coarse(denorm_f0(f0_t, uv_t)).to(device) - - mel = pad_or_cut_xd(self.mel_net(wav_t)[0], T, dim=0).unsqueeze(0) - mel_nonpadding_mask = torch.zeros(1, T, device=device) - mel_nonpadding_mask[:, :real_len] = 1.0 - mel = (mel.transpose(1, 2) * mel_nonpadding_mask.unsqueeze(1)).transpose(1, 2) - mel_nonpadding = mel.abs().sum(-1) > 0 - - mel2word_t = pad_or_cut_xd(torch.as_tensor(mel2word, device=device, dtype=torch.long), T, 0) - word_bd = torch.zeros_like(mel2word_t) - word_bd[1:] = (mel2word_t[1:] != mel2word_t[:-1]).long() - word_bd[real_len:] = 0 - word_bd = word_bd.unsqueeze(0) - - outputs = self.midi( - mel=mel[:, :, : hparams.get("use_mel_bins", 80)], - word_bd=word_bd, - pitch=pitch_coarse, - uv=uv_t, - non_padding=mel_nonpadding, - ) - outputs["word_durs_filtered"] = word_durs_filtered - outputs["real_len"] = real_len - outputs["word_bd"] = word_bd - return outputs - - @staticmethod - def _load_wav(wav_src: str | np.ndarray, sample_rate: int, *, src_sample_rate: int | None = None) -> np.ndarray: - if isinstance(wav_src, str): - wav, _ = load_mono_audio(wav_src, target_sr=sample_rate) - return wav - wav = np.asarray(wav_src, dtype=np.float32) - if src_sample_rate is not None and src_sample_rate != sample_rate: - wav = resample_mono(wav, orig_sr=src_sample_rate, target_sr=sample_rate) - return wav - - @staticmethod - def _normalize_note2words(note2words: list[int]) -> list[int]: - if not note2words: - return [] - out = [note2words[0]] - for idx in range(1, len(note2words)): - out.append(max(note2words[idx], out[-1])) - return out - - @staticmethod - def _build_ep_types(note2words: list[int], align_words: list[str]) -> list[int]: - ep_types: list[int] = [] - prev = -1 - for i, w in zip(note2words, align_words): - ep_types.append(1 if w == "" else (2 if i != prev else 3)) - prev = i - return ep_types - - @torch.no_grad() - def transcribe( - self, - item: dict[str, Any], - *, - segment_info: dict[str, Any] | None = None, - verbose: bool = False, - ) -> dict[str, Any]: - if "word_durs" not in item: - raise ValueError('item must contain "word_durs" from lyric transcription') - - if item.get("wav_fn"): - wav = self._load_wav(item["wav_fn"], self.hparams["audio_sample_rate"]) - elif item.get("wav") is not None: - wav = self._load_wav( - item["wav"], - self.hparams["audio_sample_rate"], - src_sample_rate=item.get("sample_rate"), - ) - else: - raise ValueError('item must contain "wav_fn" or "wav"') - - device = next(self.parameters()).device - outputs = self(torch.from_numpy(wav).float().to(device), list(item["word_durs"])) - real_len = int(outputs["real_len"]) - word_durs_filtered = outputs["word_durs_filtered"] - word_bd = outputs["word_bd"] - item_name = item.get("item_name", "") - - note_lengths = outputs["note_lengths"].detach().cpu().numpy() - note_bd_pred = outputs["note_bd_pred"][0].detach().cpu().numpy()[:real_len] - note_pred = outputs["note_pred"][0].detach().cpu().numpy()[: note_lengths[0]] - - if note_pred.shape == (0,): - rosvot_out = {"item_name": item_name, "pitches": [], "note_durs": [], "note2words": None} - else: - note_itv_pred = boundary2Interval(note_bd_pred) - word_bd_for_reg = word_bd[0].detach().cpu().numpy()[:real_len] - hop = self.hparams["hop_size"] - sr = self.hparams["audio_sample_rate"] - - if self.hparams.get("infer_regulate_real_note_itv", True): - try: - note_itv_pred_secs, note2words = regulate_real_note_itv( - note_itv_pred, - note_bd_pred, - word_bd_for_reg, - np.array(word_durs_filtered), - hop, - sr, - ) - note_pred, note_itv_pred_secs, note2words = regulate_ill_slur( - note_pred, note_itv_pred_secs, note2words - ) - except Exception: - if verbose: - logger.exception("ROSVOT postprocess failed") - note_itv_pred_secs = note_itv_pred * hop / sr - note2words = None - else: - note_itv_pred_secs = note_itv_pred * hop / sr - note2words = None - - rosvot_out = { - "item_name": item_name, - "pitches": note_pred.tolist(), - "note_durs": [float(itv[1] - itv[0]) for itv in note_itv_pred_secs], - "note2words": note2words.tolist() if note2words is not None else None, - } - - note2words_raw = rosvot_out.get("note2words") or [] - align_words = [item["words"][idx - 1] for idx in note2words_raw if 0 < idx <= len(item["words"])] - ep_types = self._build_ep_types(self._normalize_note2words(note2words_raw), align_words) if align_words else [] - seg = segment_info or item - - return { - "item_name": seg.get("item_name", item_name), - "wav_fn": seg.get("wav_fn", item.get("wav_fn", "")), - "origin_wav_fn": seg.get("origin_wav_fn", item.get("origin_wav_fn", "")), - "start_time_ms": seg.get("start_time_ms", ""), - "end_time_ms": seg.get("end_time_ms", ""), - "language": item.get("language", ""), - "note_text": align_words, - "note_dur": rosvot_out.get("note_durs", []), - "note_type": ep_types, - "note_pitch": rosvot_out.get("pitches", []), - } diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/segmenter.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/segmenter.py deleted file mode 100644 index 0140182ab..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/segmenter.py +++ /dev/null @@ -1,267 +0,0 @@ -"""Rule-based vocal segmentation from F0 contour (no neural net).""" - -from __future__ import annotations - -import time -from dataclasses import dataclass -from typing import ClassVar - -import numpy as np -import torch.nn as nn -from vllm.logger import init_logger - -from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery - -logger = init_logger(__name__) - - -@dataclass(frozen=True) -class VocalSegmentConfig: - hop_ms: int = 20 - smooth_ms: int = 200 - start_ms: int = 120 - end_ms: int = 200 - prepad_ms: int = 80 - postpad_ms: int = 120 - min_len_ms: int = 1000 - max_len_ms: int = 20000 - short_seg_merge_gap_ms: int = 8000 - small_gap_ms: int = 500 - lookback_ms: int = 200 - lookahead_ms: int = 100 - - -def _moving_average(x: np.ndarray, win: int) -> np.ndarray: - if win <= 1: - return x - kernel = np.ones(win, dtype=np.float32) / float(win) - return np.convolve(x, kernel, mode="same") - - -def _merge_short_segments( - segments_ms: list[list[int]], - *, - min_len_ms: int, - max_len_ms: int, - short_seg_merge_gap_ms: int, - small_gap_ms: int, -) -> list[list[int]]: - if not segments_ms: - return [] - - merged: list[list[int]] = [] - cur_start, cur_end = segments_ms[0] - for next_start, next_end in segments_ms[1:]: - cur_len = cur_end - cur_start - gap_ms = next_start - cur_end - merged_len = next_end - cur_start - should_merge = (cur_len < min_len_ms and gap_ms < short_seg_merge_gap_ms) or ( - gap_ms < small_gap_ms and merged_len < max_len_ms - ) - if should_merge: - cur_end = next_end - continue - if (cur_end - cur_start) >= min_len_ms: - merged.append([cur_start, cur_end]) - cur_start, cur_end = next_start, next_end - - if (cur_end - cur_start) >= min_len_ms: - merged.append([cur_start, cur_end]) - - if not merged: - return segments_ms - - return merged - - -def _voiced_to_segments( - voiced: np.ndarray, - *, - hop_ms: int, - smooth_ms: int, - start_ms: int, - end_ms: int, - prepad_ms: int, - postpad_ms: int, - max_len_ms: int, -) -> list[list[int]]: - smooth_frames = max(1, int(round(smooth_ms / hop_ms))) - active = _moving_average(voiced.astype(np.float32), smooth_frames) >= 0.5 - - segments: list[list[int]] = [] - start_idx = None - start_frames = max(1, int(round(start_ms / hop_ms))) - end_frames = max(1, int(round(end_ms / hop_ms))) - prepad_frames = max(0, int(round(prepad_ms / hop_ms))) - postpad_frames = max(0, int(round(postpad_ms / hop_ms))) - active_count = 0 - inactive_count = 0 - - for i, flag in enumerate(active): - if flag: - active_count += 1 - inactive_count = 0 - else: - inactive_count += 1 - active_count = 0 - - if start_idx is None: - if active_count >= start_frames: - start_idx = max(0, i - start_frames + 1 - prepad_frames) - else: - if inactive_count >= end_frames: - end_idx = min(len(active) - 1, i - end_frames + 1 + postpad_frames) - start_ms_val = start_idx * hop_ms - end_ms_val = end_idx * hop_ms + hop_ms - if end_ms_val > start_ms_val: - segments.append([int(start_ms_val), int(end_ms_val)]) - start_idx = None - - if start_idx is not None: - start_ms_val = start_idx * hop_ms - end_idx = min(len(active) - 1, len(active) - 1 + postpad_frames) - end_ms_val = end_idx * hop_ms + hop_ms - if end_ms_val > start_ms_val: - segments.append([int(start_ms_val), int(end_ms_val)]) - - def _split_segment(seg: list[int]) -> list[list[int]]: - start_ms_val, end_ms_val = seg - start_frame = int(start_ms_val // hop_ms) - end_frame = int((end_ms_val - 1) // hop_ms) - end_frame = max(start_frame, min(end_frame, len(active) - 1)) - - best_start = None - best_len = 0 - cur_start = None - cur_len = 0 - for idx in range(start_frame, end_frame + 1): - if not active[idx]: - if cur_start is None: - cur_start = idx - cur_len = 1 - else: - cur_len += 1 - else: - if cur_start is not None and cur_len > best_len: - best_start, best_len = cur_start, cur_len - cur_start = None - cur_len = 0 - if cur_start is not None and cur_len > best_len: - best_start, best_len = cur_start, cur_len - - if best_start is None: - split_frame = (start_frame + end_frame) // 2 - else: - split_frame = best_start + best_len // 2 - - split_ms = split_frame * hop_ms - if split_ms <= start_ms_val: - split_ms = start_ms_val + hop_ms - if split_ms >= end_ms_val: - split_ms = end_ms_val - hop_ms - if split_ms <= start_ms_val or split_ms >= end_ms_val: - return [seg] - return [[start_ms_val, int(split_ms)], [int(split_ms), end_ms_val]] - - queue = segments[:] - segments = [] - while queue: - seg = queue.pop(0) - if (seg[1] - seg[0]) <= max_len_ms: - segments.append(seg) - continue - parts = _split_segment(seg) - if len(parts) == 1: - segments.append(seg) - else: - queue = parts + queue - return segments - - -class VocalSegmenter(nn.Module, SupportAudioInput, SupportsComponentDiscovery): - support_audio_input: ClassVar[bool] = True - _dit_modules: ClassVar[list[str]] = [] - _encoder_modules: ClassVar[list[str]] = [] - _vae_modules: ClassVar[list[str]] = [] - _resident_modules: ClassVar[list[str]] = ["."] - _layerwise_offload_blocks_attrs: ClassVar[list[str]] = [] - - def __init__(self, config: VocalSegmentConfig | None = None, *, verbose: bool = False): - super().__init__() - self.config = config or VocalSegmentConfig() - self.verbose = verbose - - def forward( - self, - audio: np.ndarray, - sample_rate: int, - f0: np.ndarray, - *, - base_name: str = "vocal", - origin_wav_fn: str = "", - verbose: bool | None = None, - ) -> list[dict]: - verbose = self.verbose if verbose is None else verbose - if verbose: - logger.info("[vocal detection] segmenting %s", base_name) - t0 = time.time() - - cfg = self.config - y = np.asarray(audio, dtype=np.float32).squeeze() - sr = int(sample_rate) - - segments_ms = _voiced_to_segments( - f0 > 0, - hop_ms=cfg.hop_ms, - smooth_ms=cfg.smooth_ms, - start_ms=cfg.start_ms, - end_ms=cfg.end_ms, - prepad_ms=cfg.prepad_ms, - postpad_ms=cfg.postpad_ms, - max_len_ms=cfg.max_len_ms, - ) - segments_ms = _merge_short_segments( - segments_ms, - min_len_ms=cfg.min_len_ms, - max_len_ms=cfg.max_len_ms, - short_seg_merge_gap_ms=cfg.short_seg_merge_gap_ms, - small_gap_ms=cfg.small_gap_ms, - ) - - adjusted_segments: list[list[int]] = [] - prev_end = 0 - for start_ms, end_ms in segments_ms: - start_ms = max(0, start_ms - cfg.lookback_ms) - end_ms = min(end_ms + cfg.lookahead_ms, int(y.shape[0] / sr * 1000)) - if start_ms < prev_end and adjusted_segments: - adjusted_segments[-1][1] = start_ms - adjusted_segments.append([start_ms, end_ms]) - prev_end = end_ms - - segment_infos = [] - for idx, (start_ms, end_ms) in enumerate(adjusted_segments): - if end_ms - start_ms > cfg.max_len_ms: - start_ms = end_ms - cfg.max_len_ms - key = f"{base_name}_{idx}" - start_sample = int(start_ms * sr / 1000) - end_sample = int(end_ms * sr / 1000) - segment_infos.append( - { - "item_name": key, - "wav": y[start_sample:end_sample], - "sample_rate": sr, - "wav_fn": key, - "start_time_ms": int(start_sample * 1000 / sr), - "end_time_ms": int(end_sample * 1000 / sr), - "origin_wav_fn": origin_wav_fn, - "duration": int((end_sample - start_sample) * 1000 / sr), - } - ) - - if verbose: - logger.info( - "[vocal detection] done n_segments=%s time=%.3fs", - len(segment_infos), - time.time() - t0, - ) - return segment_infos diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/stack.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/stack.py deleted file mode 100644 index e29b52a29..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/stack.py +++ /dev/null @@ -1,145 +0,0 @@ -"""Lazy-loaded SoulX preprocess model tree.""" - -import os -import tempfile -from pathlib import Path -from typing import ClassVar - -import numpy as np -import soundfile as sf -import torch -import torch.nn as nn -from vllm.logger import init_logger - -from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery - -from .asr import LyricModel -from .mel_grid_f0 import extract_f0_file -from .rmvpe import RMVPE -from .rosvot import RosvotModel -from .segmenter import VocalSegmenter -from .utils import resample_mono -from .vocal_sep import MelBandRoformerSeparator - -logger = init_logger(__name__) - - -class SoulXPreprocessStack(nn.Module, SupportsComponentDiscovery): - """Unified nn.Module tree for SoulX-Singer audio preprocess.""" - - _dit_modules: ClassVar[list[str]] = [] - _encoder_modules: ClassVar[list[str]] = ["vocal_sep", "lyric", "rosvot"] - _vae_modules: ClassVar[list[str]] = [] - _resident_modules: ClassVar[list[str]] = ["rmvpe"] - - def __init__( - self, - weights: dict[str, str], - device: str, - *, - target_sr: int = 24000, - hop_size: int = 480, - verbose: bool = False, - ) -> None: - super().__init__() - self.weights = weights - self.device_str = device - self.target_sr = target_sr - self.hop_size = hop_size - self.verbose = verbose - self._rmvpe: nn.Module | None = None - self._vocal_sep: nn.Module | None = None - self._segmenter: nn.Module | None = None - self._lyric: nn.Module | None = None - self._rosvot: nn.Module | None = None - - @property - def device(self) -> torch.device: - return torch.device(self.device_str) - - def ensure_rmvpe(self): - if self._rmvpe is None: - if self.verbose: - logger.info("[rmvpe] loading shared pitch model") - self._rmvpe = RMVPE(self.weights["rmvpe"], device=self.device_str) - self.add_module("rmvpe", self._rmvpe) - return self._rmvpe - - def ensure_vocal_sep(self): - if self._vocal_sep is None: - if self.verbose: - logger.info("[vocal extraction] loading model") - self._vocal_sep = MelBandRoformerSeparator( - self.weights["sep_config"], - self.weights["sep_ckpt"], - use_half=False, - ).to(self.device) - self.add_module("vocal_sep", self._vocal_sep) - if self.verbose: - logger.info("[vocal extraction] ready on %s, sr=%s", self.device_str, self._vocal_sep.sample_rate) - return self._vocal_sep - - def ensure_segmenter(self): - if self._segmenter is None: - self._segmenter = VocalSegmenter(verbose=self.verbose) - self.add_module("segmenter", self._segmenter) - return self._segmenter - - def ensure_lyric(self): - if self._lyric is None: - if self.verbose: - logger.info("[lyric transcription] loading ASR") - self._lyric = LyricModel( - self.weights["asr_zh"], - self.weights["asr_en"], - device=self.device_str, - target_sr=self.target_sr, - hop_size=self.hop_size, - ) - self.add_module("lyric", self._lyric) - return self._lyric - - def ensure_rosvot(self): - if self._rosvot is None: - if self.verbose: - logger.info("[note transcription] loading ROSVOT") - rosvot_src = os.environ.get("ROSVOT_SOURCE_DIR") - self._rosvot = RosvotModel( - self.weights["rosvot"], - pe=self.ensure_rmvpe(), - verbose=self.verbose, - rosvot_source_dir=rosvot_src, - ).to(self.device) - self.add_module("rosvot", self._rosvot) - return self._rosvot - - def extract_f0(self, vocal, sample_rate: int, *, f0_path: str | None = None) -> np.ndarray: - """Extract F0 on the mel grid; ndarray inputs are spooled to a temp wav.""" - rmvpe = self.ensure_rmvpe() - kwargs = dict( - target_sr=self.target_sr, - hop_size=self.hop_size, - f0_path=f0_path, - verbose=self.verbose, - ) - if isinstance(vocal, str): - return extract_f0_file(rmvpe, vocal, **kwargs) - tmp_path: str | None = None - try: - with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: - tmp_path = tmp.name - sf.write(tmp_path, np.asarray(vocal, dtype=np.float32).squeeze(), int(sample_rate)) - return extract_f0_file(rmvpe, tmp_path, **kwargs) - finally: - if tmp_path is not None: - Path(tmp_path).unlink(missing_ok=True) - - def extract_vocal(self, audio: str | tuple[np.ndarray, int]) -> tuple[np.ndarray, int]: - sep = self.ensure_vocal_sep() - if isinstance(audio, tuple): - mix, sr = audio - mix = np.asarray(mix, dtype=np.float32).squeeze() - if sr != sep.sample_rate: - mix = resample_mono(mix, orig_sr=sr, target_sr=sep.sample_rate) - return sep.separate_mono(mix, self.device), sep.sample_rate - return sep.separate_path(str(audio), self.device) diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/utils.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/utils.py deleted file mode 100644 index f135fd553..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/utils.py +++ /dev/null @@ -1,260 +0,0 @@ -"""Shared helpers for SoulX preprocess (I/O, config, checkpoints, pitch).""" - -from pathlib import Path -from typing import Any - -import numpy as np -import soundfile as sf -import torch -import torch.nn as nn -import torch.nn.functional as F -from omegaconf import OmegaConf -from vllm.logger import init_logger -from vllm.multimodal.audio import AudioResampler -from vllm.multimodal.media.audio import load_audio - -logger = init_logger(__name__) - - -# --- audio I/O --- - - -def resample_mono( - audio: np.ndarray, - *, - orig_sr: int, - target_sr: int, -) -> np.ndarray: - """Resample mono float waveform to ``target_sr``.""" - audio = np.asarray(audio, dtype=np.float32).squeeze() - if orig_sr == target_sr: - return audio - resampler = AudioResampler(target_sr=float(target_sr)) - return np.asarray( - resampler.resample(audio, orig_sr=float(orig_sr)), - dtype=np.float32, - ) - - -def get_audio_duration(path: str) -> float: - """Return audio duration in seconds for a file on disk.""" - return float(sf.info(path).duration) - - -def load_mono_audio( - source: str | tuple[np.ndarray, int], - *, - target_sr: int | None = None, -) -> tuple[np.ndarray, int]: - if isinstance(source, tuple): - audio, sr = source - audio = np.asarray(audio, dtype=np.float32).squeeze() - if target_sr is not None and sr != target_sr: - audio = resample_mono(audio, orig_sr=sr, target_sr=target_sr) - sr = target_sr - return audio, sr - - if target_sr is None: - audio, sr = load_audio(str(source), sr=None, mono=True) - else: - audio, sr = load_audio(str(source), sr=target_sr, mono=True) - return np.asarray(audio, dtype=np.float32).squeeze(), int(sr) - - -# --- tensor padding --- - - -def pad_or_cut_xd(values: torch.Tensor, tgt_len: int, dim: int = -1, pad_value=0) -> torch.Tensor: - ndim = values.ndim - if ndim == 1: - if values.shape[0] < tgt_len: - return F.pad(values, [0, tgt_len - values.shape[0]], value=pad_value) - return values[:tgt_len] - if ndim == 2: - if dim in (0, -2): - src_len = values.shape[0] - if src_len < tgt_len: - return F.pad(values, [0, 0, 0, tgt_len - src_len], value=pad_value) - return values[:tgt_len] - src_len = values.shape[1] - if src_len < tgt_len: - return F.pad(values, [0, tgt_len - src_len], value=pad_value) - return values[:, :tgt_len] - raise NotImplementedError(f"pad_or_cut_xd supports 1D/2D tensors, got ndim={ndim}") - - -# --- YAML config (OmegaConf) --- - - -def _apply_hparams_str(config: dict[str, Any], hparams_str: str) -> None: - if not hparams_str: - return - for item in hparams_str.split(","): - key, raw_value = item.split("=", maxsplit=1) - raw_value = raw_value.strip().strip("'\"") - node: dict[str, Any] = config - parts = key.split(".") - for part in parts[:-1]: - node = node.setdefault(part, {}) - leaf = parts[-1] - current = node.get(leaf) - if raw_value in {"True", "False"}: - node[leaf] = raw_value == "True" - elif isinstance(current, bool): - node[leaf] = raw_value.lower() == "true" - elif isinstance(current, int): - node[leaf] = int(raw_value) - elif isinstance(current, float): - node[leaf] = float(raw_value) - else: - node[leaf] = type(current)(raw_value) if current is not None else raw_value - - -def load_yaml_config( - config: str | Path, - *, - hparams_str: str = "", -) -> dict[str, Any]: - """Load YAML with optional ``base_config`` chain; returns a plain dict.""" - config_path = Path(config).expanduser() - if not config_path.is_file(): - raise FileNotFoundError(f"Config not found: {config_path}") - - merged = OmegaConf.create({}) - pending = [config_path.resolve()] - loaded: set[str] = set() - - while pending: - path = pending.pop(0) - path_key = str(path) - if path_key in loaded: - continue - loaded.add(path_key) - - current = OmegaConf.load(path_key) - base_configs = current.get("base_config") - if base_configs is not None: - entries = base_configs if OmegaConf.is_list(base_configs) else [base_configs] - for entry in entries: - base_path = Path(entry) - if not base_path.is_absolute(): - base_path = (path.parent / entry).resolve() - if not base_path.is_file(): - logger.warning("Skipping missing base config: %s", base_path) - continue - pending.insert(0, base_path) - merged = OmegaConf.merge(current, merged) - - result = OmegaConf.to_container(merged, resolve=True) - if not isinstance(result, dict): - raise TypeError(f"Expected dict config, got {type(result)}") - _apply_hparams_str(result, hparams_str) - return result - - -def load_rosvot_config(config: str | Path, *, hparams_str: str = "") -> dict[str, Any]: - return load_yaml_config(config, hparams_str=hparams_str) - - -# --- checkpoint loading --- - - -def load_model_ckpt( - model: nn.Module, - checkpoint_path: str, - *, - model_name: str = "model", - verbose: bool = True, -) -> None: - checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) - if not isinstance(checkpoint, dict) or "state_dict" not in checkpoint: - raise ValueError(f"Unexpected ROSVOT checkpoint format: {checkpoint_path}") - - state_dict = checkpoint["state_dict"] - if state_dict and "." in next(iter(state_dict)): - prefix = f"{model_name}." - state_dict = {k[len(prefix) :]: v for k, v in state_dict.items() if k.startswith(prefix)} - elif model_name in state_dict and isinstance(state_dict[model_name], dict): - state_dict = state_dict[model_name] - - model.load_state_dict(state_dict, strict=True) - if verbose: - logger.info("Loaded ROSVOT weights from %s", checkpoint_path) - - -# --- pitch helpers (ROSVOT / RMVPE) --- - - -def f0_to_coarse(f0, f0_bin=256, f0_max=900.0, f0_min=50.0): - f0_mel_min = 1127 * np.log(1 + f0_min / 700) - f0_mel_max = 1127 * np.log(1 + f0_max / 700) - is_torch = isinstance(f0, torch.Tensor) - f0_mel = 1127 * (1 + f0 / 700).log() if is_torch else 1127 * np.log(1 + f0 / 700) - f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * (f0_bin - 2) / (f0_mel_max - f0_mel_min) + 1 - f0_mel[f0_mel <= 1] = 1 - f0_mel[f0_mel > f0_bin - 1] = f0_bin - 1 - return (f0_mel + 0.5).long() if is_torch else np.rint(f0_mel).astype(int) - - -def norm_f0(f0, uv, pitch_norm="log", f0_mean=400, f0_std=100): - is_torch = isinstance(f0, torch.Tensor) - if pitch_norm == "standard": - f0 = (f0 - f0_mean) / f0_std - if pitch_norm == "log": - f0 = torch.log2(f0 + 1e-8) if is_torch else np.log2(f0 + 1e-8) - if uv is not None: - f0[uv > 0] = 0 - return f0 - - -def norm_interp_f0(f0, pitch_norm="log", f0_mean=None, f0_std=None): - is_torch = isinstance(f0, torch.Tensor) - if is_torch: - device = f0.device - f0 = f0.detach().cpu().numpy() - uv = f0 == 0 - f0 = norm_f0(f0, uv, pitch_norm, f0_mean, f0_std) - if sum(uv) == len(f0): - f0[uv] = 0 - elif sum(uv) > 0: - f0[uv] = np.interp(np.where(uv)[0], np.where(~uv)[0], f0[~uv]) - if is_torch: - uv_t = torch.FloatTensor(uv) - f0_t = torch.FloatTensor(f0) - return f0_t.to(device), uv_t.to(device) - return f0, uv - - -def denorm_f0(f0, uv, pitch_norm="log", f0_mean=400, f0_std=100, pitch_padding=None, min=50, max=900): - is_torch = isinstance(f0, torch.Tensor) - if pitch_norm == "standard": - f0 = f0 * f0_std + f0_mean - if pitch_norm == "log": - f0 = 2**f0 - f0 = f0.clamp(min=min, max=max) if is_torch else np.clip(f0, a_min=min, a_max=max) - if uv is not None: - f0[uv > 0] = 0 - if pitch_padding is not None: - f0[pitch_padding] = 0 - return f0 - - -def boundary2Interval(bd): - is_torch = isinstance(bd, torch.Tensor) - if is_torch: - device = bd.device - bd = bd.detach().cpu().numpy() - ret = np.zeros(shape=(bd.sum() + 1, 2), dtype=int) - ret_idx = 0 - ret[0, 0] = 0 - for i, u in enumerate(bd): - if i == 0: - continue - if u == 1: - ret[ret_idx, 1] = i - ret[ret_idx + 1, 0] = i - ret_idx += 1 - ret[-1, 1] = bd.shape[0] - 1 - if is_torch: - return torch.LongTensor(ret).to(device) - return ret diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/vocal_sep.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/vocal_sep.py deleted file mode 100644 index b62902d38..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/vocal_sep.py +++ /dev/null @@ -1,207 +0,0 @@ -"""BS-RoFormer vocal separation adapter (external pip package + SoulX chunking).""" - -from __future__ import annotations - -from pathlib import Path -from typing import ClassVar - -import numpy as np -import torch -import torch.nn as nn -import yaml -from omegaconf import DictConfig, OmegaConf -from vllm.logger import init_logger -from vllm.multimodal.media.audio import load_audio - -from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery -from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend - -logger = init_logger(__name__) - - -class MelBandRoformerSeparator(nn.Module, SupportAudioInput, SupportsComponentDiscovery): - support_audio_input: ClassVar[bool] = True - _dit_modules: ClassVar[list[str]] = [] - _encoder_modules: ClassVar[list[str]] = ["roformer"] - _vae_modules: ClassVar[list[str]] = [] - _resident_modules: ClassVar[list[str]] = [] - _layerwise_offload_blocks_attrs: ClassVar[list[str]] = ["roformer.layers"] - - def __init__( - self, - config_path: str | Path, - checkpoint_path: str | Path | None = None, - *, - chunk_length_sec: float = 5, - use_half: bool = True, - ): - super().__init__() - try: - from bs_roformer import MelBandRoformer as ExtModel - except ImportError as e: - raise ImportError( - 'Install BS-RoFormer: pip install "vllm-omni[soulx-preprocess]" or pip install BS-RoFormer' - ) from e - - cfg = self._load_config(Path(config_path)) - model_cfg = self._normalize_cfg(OmegaConf.to_container(cfg.model, resolve=True)) - self.roformer = ExtModel(**model_cfg) - self.config = cfg - self.sample_rate = int(cfg.audio.sample_rate) - - LayerWiseOffloadBackend.set_blocks_attr_names(self.roformer, ["layers"]) - dtype = torch.float16 if use_half else torch.float32 - self.roformer = self.roformer.to(dtype=dtype) - - chunk = int(chunk_length_sec * self.sample_rate) - self.config.inference = getattr(self.config, "inference", {}) - inf = self.config.inference - inf.chunk_size = chunk - inf.num_overlap = inf.get("num_overlap", 4) - inf.use_amp = False - inf.normalize = inf.get("normalize", True) - - fade = chunk // 10 - w = torch.ones(chunk) - w[:fade] = torch.linspace(0, 1, fade) - w[-fade:] = torch.linspace(1, 0, fade) - self.register_buffer("_window", w.to(dtype=dtype), persistent=False) - - if checkpoint_path: - self.load_checkpoint(str(checkpoint_path), use_half=use_half) - else: - self.roformer.eval() - - @staticmethod - def _normalize_cfg(cfg: dict) -> dict: - tuple_keys = {"multi_stft_resolutions_window_sizes", "freqs_per_bands", "multi_stft_window_sizes"} - - def _fix(v): - if isinstance(v, list): - return tuple(_fix(x) for x in v) - if isinstance(v, dict): - return {k: _fix(val) for k, val in v.items()} - return v - - for k in tuple_keys: - if k in cfg: - cfg[k] = _fix(cfg[k]) - return cfg - - @staticmethod - def _load_config(path: Path) -> DictConfig: - class Loader(yaml.SafeLoader): - pass - - Loader.add_constructor( - "tag:yaml.org,2002:python/tuple", - lambda loader, node: tuple(loader.construct_sequence(node)), - ) - with open(path) as f: - return OmegaConf.create(yaml.load(f, Loader=Loader)) - - def load_checkpoint(self, path: str, *, use_half: bool = True) -> None: - self.roformer.load_state_dict(torch.load(path, map_location="cpu"), strict=False) - if use_half: - self.roformer.half() - self.roformer.eval() - - @property - def chunk_size(self) -> int: - return int(self.config.inference.chunk_size) - - @property - def num_overlap(self) -> int: - return int(self.config.inference.num_overlap) - - @property - def use_amp(self) -> bool: - return bool(self.config.inference.use_amp) - - @property - def normalize_input(self) -> bool: - return bool(self.config.inference.normalize) - - @property - def num_channels(self) -> int: - return int(self.config.audio.get("num_channels", 1)) - - def _prepare_mix(self, mix: np.ndarray) -> tuple[np.ndarray, tuple[float, float] | None]: - if mix.ndim == 1: - mix = mix[np.newaxis, :] - if self.num_channels == 2 and mix.shape[0] == 1: - mix = np.repeat(mix, 2, axis=0) - if not self.normalize_input: - return mix, None - mono = mix.mean(0) - return (mix - mono.mean()) / mono.std(), (float(mono.mean()), float(mono.std())) - - @staticmethod - def _select_vocal(recon: torch.Tensor) -> torch.Tensor: - if recon.ndim == 4: - return recon[0, 0] - if recon.ndim == 3: - return recon[0] - if recon.ndim == 2: - return recon - raise ValueError(recon.shape) - - @torch.inference_mode() - def forward(self, mix: torch.Tensor) -> torch.Tensor: - if mix.ndim == 1: - mix = mix.unsqueeze(0) - if self.num_channels == 2 and mix.shape[0] == 1: - mix = mix.repeat(2, 1) - - model_dtype = next(self.roformer.parameters()).dtype - compute_dtype = torch.float32 if model_dtype == torch.float16 else model_dtype - if mix.dtype != compute_dtype: - mix = mix.to(dtype=compute_dtype) - - chunk_size = self.chunk_size - step = chunk_size // self.num_overlap - fade = chunk_size // 10 - border = chunk_size - step - n_samples = mix.shape[-1] - window = self._window.to(mix.device) - - if n_samples > 2 * border and border > 0: - mix = nn.functional.pad(mix, (border, border), mode="reflect") - - out = torch.zeros_like(mix) - weight = torch.zeros(mix.shape[-1], device=mix.device, dtype=mix.dtype) - pos = 0 - with torch.cuda.amp.autocast(enabled=self.use_amp): - while pos < mix.shape[-1]: - part = mix[:, pos : pos + chunk_size] - seg_len = part.shape[-1] - pad_mode = "reflect" if seg_len > chunk_size // 2 else "constant" - part = nn.functional.pad(part, (0, chunk_size - seg_len), mode=pad_mode, value=0) - stem = self._select_vocal(self.roformer(part.unsqueeze(0))) - win = window.clone() - if pos == 0: - win[:fade] = 1 - elif pos + step >= n_samples: - win[-fade:] = 1 - out[:, pos : pos + seg_len] += stem[:, :seg_len] * win[:seg_len] - weight[pos : pos + seg_len] += win[:seg_len] - pos += step - - out = out / weight - if n_samples > 2 * border and border > 0: - out = out[..., border:-border] - return out - - @torch.inference_mode() - def separate_mono(self, mix: np.ndarray, device: torch.device) -> np.ndarray: - m, norm = self._prepare_mix(mix) - target_dtype = next(self.roformer.parameters()).dtype - v = self.forward(torch.as_tensor(m, dtype=target_dtype, device=device)) - if norm: - v = v * norm[1] + norm[0] - v = v.mean(0) if v.ndim == 2 else v - return v.float().cpu().numpy().astype(np.float32) - - def separate_path(self, path: str, device: torch.device) -> tuple[np.ndarray, int]: - mix, _ = load_audio(path, sr=self.sample_rate, mono=False) - return self.separate_mono(np.asarray(mix, dtype=np.float32), device), self.sample_rate diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/vocoder.py b/vllm_omni/diffusion/models/soulx_singer/modules/vocoder.py deleted file mode 100644 index ca3ef852e..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/vocoder.py +++ /dev/null @@ -1,983 +0,0 @@ -import accelerate -import numpy as np -import scipy -import torch -from omegaconf import DictConfig -from torch import nn, view_as_complex, view_as_real -from torch.nn.utils import remove_weight_norm, weight_norm -from torchaudio.functional.functional import _hz_to_mel, _mel_to_hz - - -def _deep_update_dict(base: dict, override: dict | None) -> dict: - """Recursively merge `override` into `base` (dict only).""" - if not override: - return base - for k, v in override.items(): - if isinstance(v, dict) and isinstance(base.get(k), dict): - base[k] = _deep_update_dict(base[k], v) - else: - base[k] = v - return base - - -def _get_vocoder_default_cfg() -> dict: - # Defaults keep current behavior. - return { - "preprocess": { - "hop_size": 480, - "sample_rate": 24000, - "max_length": 36000, - "n_fft": 1920, - "num_mels": 128, - "win_size": 1920, - "fmin": 0, - "fmax": 12000, - "mel_var": 8.14, - "mel_mean": -4.92, - "load_phone": False, - "load_chromagram": False, - }, - "model": { - "vocos": { - "input_channels": 128, - "dim": 1024, - "intermediate_dim": 4096, - "num_layers": 30, - "n_fft": 1920, - "hop_size": 480, - "padding": "same", - }, - "period_gan": { - "max_downsample_channels": 1024, - "channels": 64, - "channel_increasing_factor": 2, - }, - "spec_gan": { - "stft_params": { - "fft_sizes": [128, 256, 512, 1024, 2048], - "hop_sizes": [32, 64, 128, 256, 512], - "win_lengths": [128, 256, 512, 1024, 2048], - "window": "hann_window", - }, - "in_channels": 1, - "out_channels": 1, - "channels": 64, - "kernel_sizes": [5, 3], - "max_downsample_channels": 1024, - "down_scales": [2, 2, 2], - "use_weight_norm": True, - "use_complex": False, - }, - }, - "loss": { - "mel_loss": {"sample_rate": 24000}, - "disc_loss_weight": 1.0, - "mel_loss_weight": 10.0, - "adv_loss_weight": 2.0, - "fm_loss_weight": 2.0, - "spec_fm_loss_weight": 1.0, - }, - } - - -def safe_log(x: torch.Tensor, clip_val: float = 1e-7) -> torch.Tensor: - """ - Computes the element-wise logarithm of the input tensor with clipping to avoid near-zero values. - - Args: - x (Tensor): Input tensor. - clip_val (float, optional): Minimum value to clip the input tensor. Defaults to 1e-7. - - Returns: - Tensor: Element-wise logarithm of the input tensor with clipping applied. - """ - return torch.log(torch.clip(x, min=clip_val)) - - -def symlog(x: torch.Tensor) -> torch.Tensor: - return torch.sign(x) * torch.log1p(x.abs()) - - -def symexp(x: torch.Tensor) -> torch.Tensor: - return torch.sign(x) * (torch.exp(x.abs()) - 1) - - -class STFT(nn.Module): - def __init__( - self, - n_fft: int, - hop_length: int, - win_length: int, - center=True, - ): - super().__init__() - self.center = center - self.n_fft = n_fft - self.hop_length = hop_length - self.win_length = win_length - window = torch.hann_window(win_length) - self.register_buffer("window", window) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - # x: (B, T * hop_length) - - if not self.center: - pad = self.win_length - self.hop_length - x = torch.nn.functional.pad(x, (pad // 2, pad // 2), mode="reflect") - - stft_spec = torch.stft( - x, - self.n_fft, - hop_length=self.hop_length, - win_length=self.win_length, - window=self.window, - center=self.center, - return_complex=False, - ) # (B, n_fft // 2 + 1, T, 2) - - rea = stft_spec[:, :, :, 0] # (B, n_fft // 2 + 1, T, 2) - imag = stft_spec[:, :, :, 1] # (B, n_fft // 2 + 1, T, 2) - - log_mag = torch.log( - torch.abs(torch.sqrt(torch.pow(rea, 2) + torch.pow(imag, 2))) + 1e-5 - ) # (B, n_fft // 2 + 1, T) - phase = torch.atan2(imag, rea) # (B, n_fft // 2 + 1, T) - - return log_mag, phase - - -class ISTFT(nn.Module): - """ - Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with - windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges. - See issue: https://github.com/pytorch/pytorch/issues/62323 - Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs. - The NOLA constraint is met as we trim padded samples anyway. - - Args: - n_fft (int): Size of Fourier transform. - hop_length (int): The distance between neighboring sliding window frames. - win_length (int): The size of window frame and STFT filter. - padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". - """ - - def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"): - super().__init__() - if padding not in ["center", "same"]: - raise ValueError("Padding must be 'center' or 'same'.") - self.padding = padding - self.n_fft = n_fft - self.hop_length = hop_length - self.win_length = win_length - window = torch.hann_window(win_length) - self.register_buffer("window", window) - - def forward(self, spec: torch.Tensor) -> torch.Tensor: - """ - Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram. - - Args: - spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size, - N is the number of frequency bins, and T is the number of time frames. - - Returns: - Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal. - """ - if self.padding == "center": - # Fallback to pytorch native implementation - return torch.istft( - spec, - self.n_fft, - self.hop_length, - self.win_length, - self.window, - center=True, - ) - elif self.padding == "same": - pad = (self.win_length - self.hop_length) // 2 - else: - raise ValueError("Padding must be 'center' or 'same'.") - - assert spec.dim() == 3, "Expected a 3D tensor as input" - B, N, T = spec.shape - - # Inverse FFT - ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") - ifft = ifft * self.window[None, :, None] - - # Overlap and Add - output_size = (T - 1) * self.hop_length + self.win_length - y = torch.nn.functional.fold( - ifft, - output_size=(1, output_size), - kernel_size=(1, self.win_length), - stride=(1, self.hop_length), - )[:, 0, 0, pad:-pad] - - # Window envelope - window_sq = self.window.square().expand(1, T, -1).transpose(1, 2) - window_envelope = torch.nn.functional.fold( - window_sq, - output_size=(1, output_size), - kernel_size=(1, self.win_length), - stride=(1, self.hop_length), - ).squeeze()[pad:-pad] - - # Normalize - assert (window_envelope > 1e-11).all() - y = y / window_envelope - - return y - - -class MDCT(nn.Module): - """ - Modified Discrete Cosine Transform (MDCT) module. - - Args: - frame_len (int): Length of the MDCT frame. - padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". - """ - - def __init__(self, frame_len: int, padding: str = "same"): - super().__init__() - if padding not in ["center", "same"]: - raise ValueError("Padding must be 'center' or 'same'.") - self.padding = padding - self.frame_len = frame_len - N = frame_len // 2 - n0 = (N + 1) / 2 - window = torch.from_numpy(scipy.signal.cosine(frame_len)).float() - self.register_buffer("window", window) - - pre_twiddle = torch.exp(-1j * torch.pi * torch.arange(frame_len) / frame_len) - post_twiddle = torch.exp(-1j * torch.pi * n0 * (torch.arange(N) + 0.5) / N) - # view_as_real: NCCL Backend does not support ComplexFloat data type - # https://github.com/pytorch/pytorch/issues/71613 - self.register_buffer("pre_twiddle", view_as_real(pre_twiddle)) - self.register_buffer("post_twiddle", view_as_real(post_twiddle)) - - def forward(self, audio: torch.Tensor) -> torch.Tensor: - """ - Apply the Modified Discrete Cosine Transform (MDCT) to the input audio. - - Args: - audio (Tensor): Input audio waveform of shape (B, T), where B is the batch size - and T is the length of the audio. - - Returns: - Tensor: MDCT coefficients of shape (B, L, N), where L is the number of output frames - and N is the number of frequency bins. - """ - if self.padding == "center": - audio = torch.nn.functional.pad(audio, (self.frame_len // 2, self.frame_len // 2)) - elif self.padding == "same": - # hop_length is 1/2 frame_len - audio = torch.nn.functional.pad(audio, (self.frame_len // 4, self.frame_len // 4)) - else: - raise ValueError("Padding must be 'center' or 'same'.") - - x = audio.unfold(-1, self.frame_len, self.frame_len // 2) - N = self.frame_len // 2 - x = x * self.window.expand(x.shape) - x = torch.fft.fft(x * view_as_complex(self.pre_twiddle).expand(x.shape), dim=-1)[..., :N] - res = x * view_as_complex(self.post_twiddle).expand(x.shape) * np.sqrt(1 / N) - return torch.real(res) * np.sqrt(2) - - -class IMDCT(nn.Module): - """ - Inverse Modified Discrete Cosine Transform (IMDCT) module. - - Args: - frame_len (int): Length of the MDCT frame. - padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". - """ - - def __init__(self, frame_len: int, padding: str = "same"): - super().__init__() - if padding not in ["center", "same"]: - raise ValueError("Padding must be 'center' or 'same'.") - self.padding = padding - self.frame_len = frame_len - N = frame_len // 2 - n0 = (N + 1) / 2 - window = torch.from_numpy(scipy.signal.cosine(frame_len)).float() - self.register_buffer("window", window) - - pre_twiddle = torch.exp(1j * torch.pi * n0 * torch.arange(N * 2) / N) - post_twiddle = torch.exp(1j * torch.pi * (torch.arange(N * 2) + n0) / (N * 2)) - self.register_buffer("pre_twiddle", view_as_real(pre_twiddle)) - self.register_buffer("post_twiddle", view_as_real(post_twiddle)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Apply the Inverse Modified Discrete Cosine Transform (IMDCT) to the input MDCT coefficients. - - Args: - x (Tensor): Input MDCT coefficients of shape (B, L, N), where B is the batch size, - L is the number of frames, and N is the number of frequency bins. - - Returns: - Tensor: Reconstructed audio waveform of shape (B, T), where T is the length of the audio. - """ - B, L, N = x.shape - Y = torch.zeros((B, L, N * 2), dtype=x.dtype, device=x.device) - Y[..., :N] = x - Y[..., N:] = -1 * torch.conj(torch.flip(x, dims=(-1,))) - y = torch.fft.ifft(Y * view_as_complex(self.pre_twiddle).expand(Y.shape), dim=-1) - y = torch.real(y * view_as_complex(self.post_twiddle).expand(y.shape)) * np.sqrt(N) * np.sqrt(2) - result = y * self.window.expand(y.shape) - output_size = (1, (L + 1) * N) - audio = torch.nn.functional.fold( - result.transpose(1, 2), - output_size=output_size, - kernel_size=(1, self.frame_len), - stride=(1, self.frame_len // 2), - )[:, 0, 0, :] - - if self.padding == "center": - pad = self.frame_len // 2 - elif self.padding == "same": - pad = self.frame_len // 4 - else: - raise ValueError("Padding must be 'center' or 'same'.") - - audio = audio[:, pad:-pad] - return audio - - -class FourierHead(nn.Module): - """Base class for inverse fourier modules.""" - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Args: - x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, - L is the sequence length, and H denotes the model dimension. - - Returns: - Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. - """ - raise NotImplementedError("Subclasses must implement the forward method.") - - -class ISTFTHead(FourierHead): - """ - ISTFT Head module for predicting STFT complex coefficients. - - Args: - dim (int): Hidden dimension of the model. - n_fft (int): Size of Fourier transform. - hop_length (int): The distance between neighboring sliding window frames, which should align with - the resolution of the input features. - padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". - """ - - def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"): - super().__init__() - out_dim = n_fft + 2 - self.out = torch.nn.Linear(dim, out_dim) - self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Forward pass of the ISTFTHead module. - - Args: - x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, - L is the sequence length, and H denotes the model dimension. - - Returns: - Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. - """ - x = self.out(x).transpose(1, 2) - mag, p = x.chunk(2, dim=1) - mag = torch.exp(mag) - mag = torch.clip(mag, max=1e2) # safeguard to prevent excessively large magnitudes - # wrapping happens here. These two lines produce real and imaginary value - x = torch.cos(p) - y = torch.sin(p) - # recalculating phase here does not produce anything new - # only costs time - # phase = torch.atan2(y, x) - # S = mag * torch.exp(phase * 1j) - # better directly produce the complex value - - # Always compute complex values in float32 to avoid ComplexHalf warning (then cast audio back if needed) - orig_dtype = mag.dtype - mag, x, y = mag.float(), x.float(), y.float() - S = mag * (x + 1j * y) - audio = self.istft(S) - if orig_dtype != torch.float32: - audio = audio.to(orig_dtype) - return audio - - -class IMDCTSymExpHead(FourierHead): - """ - IMDCT Head module for predicting MDCT coefficients with symmetric exponential function - - Args: - dim (int): Hidden dimension of the model. - mdct_frame_len (int): Length of the MDCT frame. - padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". - sample_rate (int, optional): The sample rate of the audio. If provided, the last layer will be initialized - based on perceptual scaling. Defaults to None. - clip_audio (bool, optional): If True, clip audio to [-1.0, 1.0]. Defaults to False. - """ - - def __init__( - self, - dim: int, - mdct_frame_len: int, - padding: str = "same", - sample_rate: int | None = None, - clip_audio: bool = False, - ): - super().__init__() - out_dim = mdct_frame_len // 2 - self.out = nn.Linear(dim, out_dim) - self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding) - self.clip_audio = clip_audio - - if sample_rate is not None: - # optionally init the last layer following mel-scale - m_max = _hz_to_mel(sample_rate // 2) - m_pts = torch.linspace(0, m_max, out_dim) - f_pts = _mel_to_hz(m_pts) - scale = 1 - (f_pts / f_pts.max()) - - with torch.no_grad(): - self.out.weight.mul_(scale.view(-1, 1)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Forward pass of the IMDCTSymExpHead module. - - Args: - x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, - L is the sequence length, and H denotes the model dimension. - - Returns: - Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. - """ - x = self.out(x) - x = symexp(x) - x = torch.clip(x, min=-1e2, max=1e2) # safeguard to prevent excessively large magnitudes - audio = self.imdct(x) - if self.clip_audio: - audio = torch.clip(x, min=-1.0, max=1.0) - - return audio - - -class IMDCTCosHead(FourierHead): - """ - IMDCT Head module for predicting MDCT coefficients with parametrizing MDCT = exp(m) · cos(p) - - Args: - dim (int): Hidden dimension of the model. - mdct_frame_len (int): Length of the MDCT frame. - padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". - clip_audio (bool, optional): If True, clip audio to [-1.0, 1.0]. Defaults to False. - """ - - def __init__( - self, - dim: int, - mdct_frame_len: int, - padding: str = "same", - clip_audio: bool = False, - ): - super().__init__() - self.clip_audio = clip_audio - self.out = nn.Linear(dim, mdct_frame_len) - self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Forward pass of the IMDCTCosHead module. - - Args: - x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, - L is the sequence length, and H denotes the model dimension. - - Returns: - Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. - """ - x = self.out(x) - m, p = x.chunk(2, dim=2) - m = torch.exp(m).clip(max=1e2) # safeguard to prevent excessively large magnitudes - audio = self.imdct(m * torch.cos(p)) - if self.clip_audio: - audio = torch.clip(x, min=-1.0, max=1.0) - return audio - - -class ConvNeXtBlock(nn.Module): - """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal. - - Args: - dim (int): Number of input channels. - intermediate_dim (int): Dimensionality of the intermediate layer. - layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling. - Defaults to None. - adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm. - None means non-conditional LayerNorm. Defaults to None. - """ - - def __init__( - self, - dim: int, - intermediate_dim: int, - layer_scale_init_value: float, - adanorm_num_embeddings: int | None = None, - ): - super().__init__() - self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv - self.adanorm = adanorm_num_embeddings is not None - if adanorm_num_embeddings: - self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6) - else: - self.norm = nn.LayerNorm(dim, eps=1e-6) - self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers - self.act = nn.GELU() - self.pwconv2 = nn.Linear(intermediate_dim, dim) - self.gamma = ( - nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True) - if layer_scale_init_value > 0 - else None - ) - - def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor | None = None) -> torch.Tensor: - residual = x - x = self.dwconv(x) - x = x.transpose(1, 2) # (B, C, T) -> (B, T, C) - if self.adanorm: - assert cond_embedding_id is not None - x = self.norm(x, cond_embedding_id) - else: - x = self.norm(x) - x = self.pwconv1(x) - x = self.act(x) - x = self.pwconv2(x) - if self.gamma is not None: - x = self.gamma * x - x = x.transpose(1, 2) # (B, T, C) -> (B, C, T) - - x = residual + x - return x - - -class AdaLayerNorm(nn.Module): - """ - Adaptive Layer Normalization module with learnable embeddings per `num_embeddings` classes - - Args: - num_embeddings (int): Number of embeddings. - embedding_dim (int): Dimension of the embeddings. - """ - - def __init__(self, num_embeddings: int, embedding_dim: int, eps: float = 1e-6): - super().__init__() - self.eps = eps - self.dim = embedding_dim - self.scale = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim) - self.shift = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim) - torch.nn.init.ones_(self.scale.weight) - torch.nn.init.zeros_(self.shift.weight) - - def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor) -> torch.Tensor: - scale = self.scale(cond_embedding_id) - shift = self.shift(cond_embedding_id) - x = nn.functional.layer_norm(x, (self.dim,), eps=self.eps) - x = x * scale + shift - return x - - -class ResBlock1(nn.Module): - """ - ResBlock adapted from HiFi-GAN V1 (https://github.com/jik876/hifi-gan) with dilated 1D convolutions, - but without upsampling layers. - - Args: - dim (int): Number of input channels. - kernel_size (int, optional): Size of the convolutional kernel. Defaults to 3. - dilation (tuple[int], optional): Dilation factors for the dilated convolutions. - Defaults to (1, 3, 5). - lrelu_slope (float, optional): Negative slope of the LeakyReLU activation function. - Defaults to 0.1. - layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling. - Defaults to None. - """ - - def __init__( - self, - dim: int, - kernel_size: int = 3, - dilation: tuple[int, int, int] = (1, 3, 5), - lrelu_slope: float = 0.1, - layer_scale_init_value: float | None = None, - ): - super().__init__() - self.lrelu_slope = lrelu_slope - self.convs1 = nn.ModuleList( - [ - weight_norm( - nn.Conv1d( - dim, - dim, - kernel_size, - 1, - dilation=dilation[0], - padding=self.get_padding(kernel_size, dilation[0]), - ) - ), - weight_norm( - nn.Conv1d( - dim, - dim, - kernel_size, - 1, - dilation=dilation[1], - padding=self.get_padding(kernel_size, dilation[1]), - ) - ), - weight_norm( - nn.Conv1d( - dim, - dim, - kernel_size, - 1, - dilation=dilation[2], - padding=self.get_padding(kernel_size, dilation[2]), - ) - ), - ] - ) - - self.convs2 = nn.ModuleList( - [ - weight_norm( - nn.Conv1d( - dim, - dim, - kernel_size, - 1, - dilation=1, - padding=self.get_padding(kernel_size, 1), - ) - ), - weight_norm( - nn.Conv1d( - dim, - dim, - kernel_size, - 1, - dilation=1, - padding=self.get_padding(kernel_size, 1), - ) - ), - weight_norm( - nn.Conv1d( - dim, - dim, - kernel_size, - 1, - dilation=1, - padding=self.get_padding(kernel_size, 1), - ) - ), - ] - ) - - self.gamma = nn.ParameterList( - [ - ( - nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True) - if layer_scale_init_value is not None - else None - ), - ( - nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True) - if layer_scale_init_value is not None - else None - ), - ( - nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True) - if layer_scale_init_value is not None - else None - ), - ] - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - for c1, c2, gamma in zip(self.convs1, self.convs2, self.gamma): - xt = torch.nn.functional.leaky_relu(x, negative_slope=self.lrelu_slope) - xt = c1(xt) - xt = torch.nn.functional.leaky_relu(xt, negative_slope=self.lrelu_slope) - xt = c2(xt) - if gamma is not None: - xt = gamma * xt - x = xt + x - return x - - def remove_weight_norm(self): - for conv in self.convs1: - remove_weight_norm(conv) - for conv in self.convs2: - remove_weight_norm(conv) - - @staticmethod - def get_padding(kernel_size: int, dilation: int = 1) -> int: - return int((kernel_size * dilation - dilation) / 2) - - -class Backbone(nn.Module): - """Base class for the generator's backbone. It preserves the same temporal resolution across all layers.""" - - def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: - """ - Args: - x (Tensor): Input tensor of shape (B, C, L), where B is the batch size, - C denotes output features, and L is the sequence length. - - Returns: - Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length, - and H denotes the model dimension. - """ - raise NotImplementedError("Subclasses must implement the forward method.") - - -class VocosBackbone(Backbone): - """ - Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization - - Args: - input_channels (int): Number of input features channels. - dim (int): Hidden dimension of the model. - intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock. - num_layers (int): Number of ConvNeXtBlock layers. - layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`. - adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm. - None means non-conditional model. Defaults to None. - """ - - def __init__( - self, - input_channels: int, - dim: int, - intermediate_dim: int, - num_layers: int, - layer_scale_init_value: float | None = None, - adanorm_num_embeddings: int | None = None, - ): - super().__init__() - self.input_channels = input_channels - self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3) - self.adanorm = adanorm_num_embeddings is not None - if adanorm_num_embeddings: - self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6) - else: - self.norm = nn.LayerNorm(dim, eps=1e-6) - layer_scale_init_value = layer_scale_init_value or 1 / num_layers - self.convnext = nn.ModuleList( - [ - ConvNeXtBlock( - dim=dim, - intermediate_dim=intermediate_dim, - layer_scale_init_value=layer_scale_init_value, - adanorm_num_embeddings=adanorm_num_embeddings, - ) - for _ in range(num_layers) - ] - ) - self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6) - self.apply(self._init_weights) - - def _init_weights(self, m): - if isinstance(m, nn.Conv1d | nn.Linear): - nn.init.trunc_normal_(m.weight, std=0.02) - nn.init.constant_(m.bias, 0) - - def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: - bandwidth_id = kwargs.get("bandwidth_id", None) - x = self.embed(x) - if self.adanorm: - assert bandwidth_id is not None - x = self.norm(x.transpose(1, 2), cond_embedding_id=bandwidth_id) - else: - x = self.norm(x.transpose(1, 2)) - x = x.transpose(1, 2) - for conv_block in self.convnext: - x = conv_block(x, cond_embedding_id=bandwidth_id) - x = self.final_layer_norm(x.transpose(1, 2)) - return x - - -class VocosResNetBackbone(Backbone): - """ - Vocos backbone module built with ResBlocks. - - Args: - input_channels (int): Number of input features channels. - dim (int): Hidden dimension of the model. - num_blocks (int): Number of ResBlock1 blocks. - layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to None. - """ - - def __init__( - self, - input_channels, - dim, - num_blocks, - layer_scale_init_value=None, - ): - super().__init__() - self.input_channels = input_channels - self.embed = weight_norm(nn.Conv1d(input_channels, dim, kernel_size=3, padding=1)) - layer_scale_init_value = layer_scale_init_value or 1 / num_blocks / 3 - self.resnet = nn.Sequential( - *[ResBlock1(dim=dim, layer_scale_init_value=layer_scale_init_value) for _ in range(num_blocks)] - ) - - def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: - x = self.embed(x) - x = self.resnet(x) - x = x.transpose(1, 2) - return x - - -class Vocos(nn.Module): - def __init__( - self, - input_channels: int = 256, - dim: int = 384, - intermediate_dim: int = 1152, - num_layers: int = 8, - n_fft: int = 800, - hop_size: int = 200, - padding: str = "same", - adanorm_num_embeddings=None, - cfg=None, - ): - super().__init__() - - input_channels = cfg.input_channels if cfg is not None and hasattr(cfg, "input_channels") else input_channels - dim = cfg.dim if cfg is not None and hasattr(cfg, "dim") else dim - intermediate_dim = ( - cfg.intermediate_dim if cfg is not None and hasattr(cfg, "intermediate_dim") else intermediate_dim - ) - num_layers = cfg.num_layers if cfg is not None and hasattr(cfg, "num_layers") else num_layers - adanorm_num_embeddings = ( - cfg.adanorm_num_embeddings - if cfg is not None and hasattr(cfg, "adanorm_num_embeddings") - else adanorm_num_embeddings - ) - n_fft = cfg.n_fft if cfg is not None and hasattr(cfg, "n_fft") else n_fft - hop_size = cfg.hop_size if cfg is not None and hasattr(cfg, "hop_size") else hop_size - padding = cfg.padding if cfg is not None and hasattr(cfg, "padding") else padding - - self.backbone = VocosBackbone( - input_channels=input_channels, - dim=dim, - intermediate_dim=intermediate_dim, - num_layers=num_layers, - adanorm_num_embeddings=adanorm_num_embeddings, - ) - self.head = ISTFTHead(dim, n_fft, hop_size, padding) - - def forward(self, x): - x = self.backbone(x) - x = self.head(x) - - return x[:, None, :] - - -class JsonHParams: - def __init__(self, **kwargs): - for k, v in kwargs.items(): - if isinstance(v, dict): - v = JsonHParams(**v) - self[k] = v - - def keys(self): - return self.__dict__.keys() - - def items(self): - return self.__dict__.items() - - def values(self): - return self.__dict__.values() - - def __len__(self): - return len(self.__dict__) - - def __getitem__(self, key): - return getattr(self, key) - - def __setitem__(self, key, value): - return setattr(self, key, value) - - def __contains__(self, key): - return key in self.__dict__ - - def __repr__(self): - return self.__dict__.__repr__() - - -def build_vocoder_model(cfg): - vocoder_model = Vocos(cfg=cfg.model.vocos) - vocoder_model.eval() - return vocoder_model - - -def load_checkpoint(build_model_func, cfg, ckpt_path): - model = build_model_func(cfg) - - if ckpt_path is not None: - accelerate.load_checkpoint_and_dispatch(model, ckpt_path) - return model - - -def load_vocos_model( - ckpt_path: str | None = None, - config: DictConfig = None, - device: str = "cuda", -): - """Load Vocos vocoder. - - Args: - config: DictConfig, config for vocoder. - ckpt_path: str | None = None, path to checkpoint. - device: str, device to load model. Note: accelerate dispatch handles device placement. - """ - - merged_cfg = _deep_update_dict(_get_vocoder_default_cfg(), config) - vocoder_cfg_obj = JsonHParams(**merged_cfg) - - vocoder_model = load_checkpoint(build_vocoder_model, vocoder_cfg_obj, ckpt_path) - vocoder_model.eval() - - for param in vocoder_model.parameters(): - param.requires_grad = False - - return vocoder_model - - -class Vocoder(nn.Module): - def __init__(self, vocoder: dict | None = None, ckpt_path: str | None = None): - super().__init__() - vocoder = vocoder or {} - - model_cfg = vocoder.get("model_cfg") or vocoder.get("cfg") - - self.model = load_vocos_model( - ckpt_path=ckpt_path, - config=model_cfg, - ) - - def forward(self, x): - x = self.model(x) - return x diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/whisper_encoder.py b/vllm_omni/diffusion/models/soulx_singer/modules/whisper_encoder.py deleted file mode 100644 index 379dda3a5..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/whisper_encoder.py +++ /dev/null @@ -1,96 +0,0 @@ -"""Frozen Whisper encoder wrapper (wav -> encoder embeddings).""" - -import torch -import torchaudio -from transformers import WhisperFeatureExtractor, WhisperModel - -WHISPER_MEL_FRAMES = 3000 # 3000 frames at 16000 Hz - - -class WhisperEncoder: - """Auxiliary content encoder for SVC. - - Kept in FP32 for stable HF Whisper inference under vLLM-Omni's default - BF16 diffusion trunk (same idea as Stable Audio's FP32 VAE). Outputs are - cast to the trunk dtype before fusion with the CFM conditioner. - """ - - def __init__(self, device: str | torch.device | None = None) -> None: - self.fe = WhisperFeatureExtractor.from_pretrained("openai/whisper-base") - self.model = WhisperModel.from_pretrained("openai/whisper-base") - self.device = torch.device( - device or ("cuda" if torch.cuda.is_available() else "cpu"), - ) - self.model = self.model.to(device=self.device, dtype=torch.float32) - self.model.eval() - for param in self.model.parameters(): - param.requires_grad_(False) - - @property - def dtype(self) -> torch.dtype: - return torch.float32 - - def float(self) -> "WhisperEncoder": - self.model = self.model.float() - return self - - def encode( - self, - wav: torch.Tensor, - sr: int, - *, - output_dtype: torch.dtype | None = None, - ) -> torch.Tensor: - wav = ( - torchaudio.functional.resample(wav, orig_freq=sr, new_freq=self.fe.sampling_rate) - if sr != self.fe.sampling_rate - else wav - ) - wav_np = wav.detach().float().cpu().numpy().astype("float32", copy=False) - - inputs = self.fe( - wav_np, - sampling_rate=self.fe.sampling_rate, - return_tensors="pt", - padding=False, - truncation=False, - return_attention_mask=True, - ) - - input_features = inputs.input_features - num_frames = input_features.shape[-1] - if num_frames < WHISPER_MEL_FRAMES: - pad = WHISPER_MEL_FRAMES - num_frames - input_features = torch.nn.functional.pad(input_features, (0, pad)) - else: - input_features = input_features[..., :WHISPER_MEL_FRAMES] - - input_features = input_features.to(device=wav.device, dtype=self.dtype) - if self.model.device != wav.device: - self.model = self.model.to(wav.device) - - attention_mask = inputs.attention_mask.to(wav.device) if inputs.attention_mask is not None else None - - with torch.inference_mode(): - encoder_out = self.model.encoder(input_features).last_hidden_state - - if attention_mask is not None: - valid_mel_frames = attention_mask.sum(dim=1) - valid_enc_frames = (valid_mel_frames + 1) // 2 - max_valid_enc_frames = min(int(valid_enc_frames.max().item()), encoder_out.shape[1]) - encoder_out = encoder_out[:, :max_valid_enc_frames, :] - valid_len = min(int(valid_enc_frames[0].item()), max_valid_enc_frames) - if valid_len < max_valid_enc_frames: - encoder_out[0, valid_len:, :] = 0 - - if output_dtype is not None and encoder_out.dtype != output_dtype: - encoder_out = encoder_out.to(dtype=output_dtype) - return encoder_out - - -if __name__ == "__main__": - torch.manual_seed(0) - audio = torch.randn(1, 24000 * 25).float().to("cuda") - encoder = WhisperEncoder() - whisper_encoder_out = encoder.encode(audio, sr=24000) - print(whisper_encoder_out.shape) diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py deleted file mode 100644 index 0961e6916..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py +++ /dev/null @@ -1,377 +0,0 @@ -import os -import time -from collections.abc import Callable -from contextlib import contextmanager -from typing import Any, ClassVar - -import numpy as np -import torch -import torch.nn as nn -from vllm.logger import init_logger -from vllm.utils.torch_utils import set_default_torch_dtype - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _unwrap, _wrap -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.models.interface import ( - SupportAudioInput, - SupportAudioOutput, - SupportsComponentDiscovery, -) -from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import consume_payload -from vllm_omni.diffusion.models.soulx_singer.utils import _patch_torchaudio_load, load_config -from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.platforms import current_omni_platform - -logger = init_logger(__name__) - -_SOULX_PROFILER_TARGETS = [ - "forward", -] - - -def convert_soulx_audio_output_to_numpy(audio: torch.Tensor | dict[str, Any]) -> np.ndarray | dict[str, Any]: - if isinstance(audio, dict) and isinstance(audio.get("payload"), dict): - payload = dict(audio["payload"]) - audio_payload = payload.get("audio") - if isinstance(audio_payload, torch.Tensor): - payload["audio"] = audio_payload.detach().cpu().float().numpy() - return { - "payload": payload, - "metadata": audio.get("metadata") or {}, - } - if isinstance(audio, torch.Tensor): - return audio.detach().cpu().float().numpy() - return audio - - -class FlowMatchingAudioPipeline( - nn.Module, - CFGParallelMixin, - SupportAudioInput, - SupportAudioOutput, - SupportsComponentDiscovery, - DiffusionPipelineProfilerMixin, -): - """SoulX flow-matching pipeline with CFG / CFG-parallel support.""" - - support_audio_input: ClassVar[bool] = True - support_audio_output: ClassVar[bool] = True - audio_sample_rate: ClassVar[int] = 24000 - _DEFAULT_RESCALE_CFG: ClassVar[float] = 0.75 - _dit_modules: ClassVar[list[str]] = ["cfm_decoder.model.diff_estimator"] - _encoder_modules: ClassVar[list[str]] = ["mel", "f0_encoder"] - _vae_modules: ClassVar[list[str]] = ["vocoder"] - _resident_modules: ClassVar[list[str]] = ["cfm_decoder.model.cond_emb"] - - def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): - super().__init__() - self.od_config = od_config - self.device = get_local_device() - - # 1. parse model path and yaml config - model_dir = od_config.model - assert model_dir is not None, "model_dir is not provided" - - if not os.path.exists(model_dir): - raise FileNotFoundError( - f"""Model directory {model_dir} does not exist. User should refer to - `examples/offline_inference/text_to_speech/README.md` for usage details.""" - ) - - # Empty: skip DiffusersPipelineLoader.get_all_weights() → no _prepare_weights / file glob / - # safetensors or pt iterator from the framework. Load checkpoints in __init__ instead. - self.weights_sources = () - - self.model_path = model_dir - - config_path = os.path.join(model_dir, "config.yaml") - hf_config = load_config(config_path) - - self.audio_config = hf_config.audio - self.encoder_config = hf_config.model.encoder - self.flow_matching_config = hf_config.model.flow_matching - - self.mel_dim = self.audio_config.num_mels - - self.f0_bin = self.encoder_config.f0_bin - self.f0_dim = self.encoder_config.f0_dim - self.text_dim = self.encoder_config.text_dim - self.vocab_size = self.encoder_config.vocab_size - self.pitch_dim = self.encoder_config.pitch_dim - self.type_dim = self.encoder_config.type_dim - - _patch_torchaudio_load() - - def _mel_to_audio(self, generated_mel: torch.Tensor, *, squeeze: bool = False) -> torch.Tensor: - audio = self.vocoder(generated_mel.transpose(1, 2)[0:1, ...].float()).float() - return audio.squeeze() if squeeze else audio - - def _mel_from_wav(self, wav: torch.Tensor) -> torch.Tensor: - if wav.dtype != torch.float32: - wav = wav.float() - return self.mel(wav) - - def _load_soulx_checkpoint(self, checkpoint_name: str) -> None: - weight_path = os.path.join(self.model_path, checkpoint_name) - if not os.path.isfile(weight_path): - raise FileNotFoundError( - f"Model weights not found at {weight_path}. Please check the pretrained model path." - ) - state = torch.load(weight_path, map_location=self.device) - self.mel.float() - self.load_state_dict(state["state_dict"], strict=True) - self._finalize_loaded_dtypes() - logger.info("Loaded model weights from %s", weight_path) - - def _setup_soulx_profiler(self, *, extra_targets: list[str] | None = None) -> None: - """Enable stage timing when ``od_config.enable_diffusion_pipeline_profiler`` is set.""" - targets = list(_SOULX_PROFILER_TARGETS) - if extra_targets: - targets.extend(extra_targets) - self.setup_diffusion_pipeline_profiler( - profiler_targets=targets, - enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler, - ) - - @contextmanager - def _stage_timer(self, name: str): - if not getattr(self, "enable_diffusion_pipeline_profiler", False): - yield - return - if current_omni_platform.is_available(): - current_omni_platform.synchronize() - start = time.perf_counter() - try: - yield - finally: - if current_omni_platform.is_available(): - current_omni_platform.synchronize() - duration = time.perf_counter() - start - with self._profiler_lock: - self._stage_durations[name] = self._stage_durations.get(name, 0.0) + duration - - def _profiler_stage_durations(self) -> dict[str, float] | None: - if getattr(self, "enable_diffusion_pipeline_profiler", False): - return self.stage_durations - return None - - @staticmethod - def _build_fp32_audio_modules(audio_config) -> tuple[nn.Module, nn.Module]: - """Mel/vocoder must stay FP32 even when ``od_config.dtype`` is FP16/BF16.""" - from vllm_omni.diffusion.models.soulx_singer.modules import MelSpectrogramEncoder, Vocoder - - with set_default_torch_dtype(torch.float32): - mel = MelSpectrogramEncoder(audio_config) - vocoder = Vocoder() - return mel, vocoder - - def _finalize_loaded_dtypes(self) -> None: - """Keep mel/vocoder in FP32; cast DiT trunk to ``od_config.dtype`` after load.""" - self.mel.float() - self.vocoder.float() - trunk_dtype = self.od_config.dtype - if trunk_dtype not in (torch.float16, torch.bfloat16): - return - trunk_modules = ( - "cfm_decoder", - "f0_encoder", - "preflow", - "note_text_encoder", - "note_pitch_encoder", - "note_type_encoder", - ) - for name in trunk_modules: - module = getattr(self, name, None) - if module is not None: - module.to(dtype=trunk_dtype) - whisper = getattr(self, "whisper_encoder", None) - if whisper is not None: - whisper.float() - - @property - def transformer(self) -> nn.Module: - """Alias expected by ``CFGParallelMixin`` default hooks.""" - return self.cfm_decoder.model.diff_estimator - - @property - def diffusion_trunk_dtype(self) -> torch.dtype: - """DiT trunk dtype from ``diff_estimator`` (not ``cond_emb``, which is registered first).""" - return next(self.cfm_decoder.model.diff_estimator.parameters()).dtype - - def _to_trunk_dtype(self, *tensors: torch.Tensor) -> tuple[torch.Tensor, ...]: - """Cast tensors to ``diffusion_trunk_dtype`` at the CFM/DiT boundary.""" - trunk_dtype = self.diffusion_trunk_dtype - return tuple(t if t.dtype == trunk_dtype else t.to(dtype=trunk_dtype) for t in tensors) - - def _resolve_diffusion_generator(self, sampling_params: Any) -> torch.Generator | None: - """Build a device-local RNG for CFM noise from ``generator`` or ``seed``.""" - generator = getattr(sampling_params, "generator", None) - if generator is not None: - return generator - seed = getattr(sampling_params, "seed", None) - if seed is None: - return None - return torch.Generator(device=self.device).manual_seed(int(seed)) - - def _prepare_cfm_loop_state( - self, - prompt: torch.Tensor, - cond: torch.Tensor, - n_timesteps: int, - *, - generator: torch.Generator | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: - """Cast FP32 islands to trunk dtype and build CFM loop tensors/masks.""" - trunk_dtype = self.diffusion_trunk_dtype - prompt, cond = self._to_trunk_dtype(prompt, cond) - cond_emb = self.cfm_decoder.model.cond_emb(cond) - if cond_emb.dtype != trunk_dtype: - cond_emb = self._to_trunk_dtype(cond_emb)[0] - device = cond_emb.device - - h = 1.0 / n_timesteps - h_tensor = torch.tensor(h, device=device, dtype=trunk_dtype) - prompt_len = prompt.shape[1] - target_len = cond_emb.shape[1] - prompt_len - - x_mask = torch.ones(cond_emb.shape[0], target_len, device=device, dtype=trunk_dtype) - prompt_mask = torch.ones(cond_emb.shape[0], prompt_len, device=device, dtype=trunk_dtype) - xt_mask = torch.cat([prompt_mask, x_mask], dim=1) - xt = torch.randn( - cond_emb.shape[0], - target_len, - self.mel_dim, - dtype=trunk_dtype, - device=device, - generator=generator, - ) - return cond_emb, prompt, xt, x_mask, xt_mask, h_tensor, prompt_len - - def predict_noise( - self, - *, - x: torch.Tensor, - diffusion_step: torch.Tensor, - cond: torch.Tensor, - x_mask: torch.Tensor, - prompt_len: int | None = None, - **_: Any, - ) -> torch.Tensor: - """Forward through ``diff_estimator``, optionally slice off prompt frames.""" - x, diffusion_step, cond, x_mask = self._to_trunk_dtype(x, diffusion_step, cond, x_mask) - flow_pred = self.cfm_decoder.model.diff_estimator(x, diffusion_step, cond, x_mask) - if prompt_len is not None: - flow_pred = flow_pred[:, prompt_len:, :] - return flow_pred - - def combine_cfg_noise( - self, - positive_noise_pred: torch.Tensor | tuple[torch.Tensor, ...], - negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...], - true_cfg_scale: float, - cfg_normalize: bool = False, - rescale_cfg: float | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, ...]: - """SoulX CFG: pos + scale * (pos - neg) with std rescaling.""" - del cfg_normalize # SoulX uses rescale_cfg instead of cfg_normalize_function. - if rescale_cfg is None: - rescale_cfg = getattr(self, "rescale_cfg", self._DEFAULT_RESCALE_CFG) - - pos = _unwrap(_wrap(positive_noise_pred)) - neg = _unwrap(_wrap(negative_noise_pred)) - - pos_flow_pred_std = pos.std() - flow_pred_cfg = pos + true_cfg_scale * (pos - neg) - rescale_flow_pred = flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std() - combined = rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg - return combined.to(dtype=pos.dtype) - - def _run_flow_matching_loop( - self, - prompt: torch.Tensor, - cond: torch.Tensor, - n_timesteps: int, - cfg: float, - *, - generator: torch.Generator | None = None, - ) -> torch.Tensor: - """Euler flow-matching loop aligned with ``FlowMatchingTransformer.reverse_diffusion``.""" - cond_emb, prompt, xt, x_mask, xt_mask, h_tensor, prompt_len = self._prepare_cfm_loop_state( - prompt, cond, n_timesteps, generator=generator - ) - trunk_dtype = self.diffusion_trunk_dtype - step_scale = 1.0 / n_timesteps - - for i in range(n_timesteps): - xt_input = torch.cat([prompt, xt], dim=1) - diffusion_step = torch.full( - (xt.shape[0],), - fill_value=(i + 0.5) * step_scale, - device=xt.device, - dtype=trunk_dtype, - ) - - flow_pred = self.predict_noise_maybe_with_cfg( - do_true_cfg=(cfg > 0), - true_cfg_scale=cfg, - positive_kwargs={ - "x": xt_input, - "diffusion_step": diffusion_step, - "cond": cond_emb, - "x_mask": xt_mask, - "prompt_len": prompt_len, - }, - negative_kwargs={ - "x": xt, - "diffusion_step": diffusion_step, - # CFG null cond: zero out target frames only (prompt cond stays in pos branch). - "cond": torch.zeros_like(cond_emb)[:, : xt.shape[1], :], - "x_mask": x_mask, - }, - cfg_normalize=False, - ) - - xt = xt + flow_pred * h_tensor - - return xt - - def _forward_batch_from_request( - self, - req: OmniDiffusionRequest, - *, - kind: str, - metadata_key: str, - infer_batch_fn: Callable[..., tuple[torch.Tensor, int]], - prepare_extra_args: Callable[[dict[str, Any], Any], dict[str, Any]] | None = None, - ) -> DiffusionOutput: - sampling_params = req.sampling_params - extra_args = dict(getattr(sampling_params, "extra_args", None) or {}) - if prepare_extra_args is not None: - extra_args = prepare_extra_args(extra_args, sampling_params) - sampling_params.extra_args = extra_args - - num_inference_steps = sampling_params.num_inference_steps or 32 - guidance_scale = sampling_params.guidance_scale if sampling_params.guidance_scale_provided else 3.0 - generator = self._resolve_diffusion_generator(sampling_params) - - with self._stage_timer("consume_payload"): - payload = consume_payload(req, kind, self.device) - - with self._stage_timer("infer_segments"): - audio, pitch_shift = infer_batch_fn( - payload, - extra_args=extra_args, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - generator=generator, - ) - - return DiffusionOutput( - output={ - "payload": {"audio": audio}, - "metadata": {"audio": {metadata_key: pitch_shift}}, - }, - stage_durations=self._profiler_stage_durations() or {}, - ) diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py deleted file mode 100644 index 85fed9ee7..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py +++ /dev/null @@ -1,411 +0,0 @@ -"""SoulX-Singer SVC (voice conversion) pipeline implementation.""" - -import os -from collections.abc import Iterable -from typing import Any, ClassVar - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from vllm.logger import init_logger -from vllm.utils.torch_utils import set_default_torch_dtype - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.models.soulx_singer.modules import ( - CFMDecoder, - WhisperEncoder, -) -from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import ( - FlowMatchingAudioPipeline, - convert_soulx_audio_output_to_numpy, -) -from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( - SOULX_PREPROCESSED_KEY, - get_soulx_preprocessed_payload, - has_precomputed, -) -from vllm_omni.diffusion.models.soulx_singer.preprocess.pre_process import ( - attach_preprocess_for_diffusion_request, - is_warmup_request, - resolve_preprocess_audio, -) -from vllm_omni.diffusion.models.soulx_singer.utils import ( - f0_to_coarse, - load_config, - resolve_pitch_shift, - validate_soulx_extra_args, -) -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch - -logger = init_logger(__name__) - - -def get_soulxsinger_svc_pre_process_func(od_config: OmniDiffusionConfig): - hf_config = load_config(os.path.join(od_config.model, "config.yaml")) - sample_rate = hf_config.audio.sample_rate - device = get_local_device() - _pipeline = None - - def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: - nonlocal _pipeline - - extra_args = validate_soulx_extra_args( - "svc", - dict(getattr(request.sampling_params, "extra_args", None) or {}), - ) - - # Inline build: when no warmup/no precomputed paths/no IPC payload, - # build the preprocess payload directly from audio file paths. - if not (is_warmup_request(request) or has_precomputed(extra_args, "svc")): - prompt = request.prompt - if not isinstance(prompt, str) and not get_soulx_preprocessed_payload(prompt): # type: ignore[arg-type] - prompt_audio, target_audio = resolve_preprocess_audio(prompt, extra_args) # type: ignore[arg-type] - if prompt_audio is not None and target_audio is not None: - if _pipeline is None: - from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.pipeline import ( - SoulXPreprocessPipeline, - ) - - _pipeline = SoulXPreprocessPipeline( - od_config=od_config, - vocal_sep=bool(extra_args.get("vocal_sep", False)), - verbose=bool(extra_args.get("preprocess_verbose", False)), - extra_args=extra_args, - ) - payload = _pipeline.build_svc_payload_from_audio( - prompt_audio=prompt_audio, - target_audio=target_audio, - sample_rate=sample_rate, - device=device, - vocal_sep=extra_args.get("vocal_sep"), - ) - payload.setdefault("kind", "svc") - prompt.setdefault("additional_information", {})[SOULX_PREPROCESSED_KEY] = payload # type: ignore[union-attr, assignment] - - return attach_preprocess_for_diffusion_request( - request, - kind="svc", - sample_rate=sample_rate, - device=device, - ) - - return pre_process_func - - -def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig): - def post_process_func(audio: torch.Tensor): - return convert_soulx_audio_output_to_numpy(audio) - - return post_process_func - - -class PipelineSoulXSingerSVC(FlowMatchingAudioPipeline): - """SVC pipeline for the SoulX-Singer model.""" - - _encoder_modules: ClassVar[list[str]] = [ - "whisper_encoder", - "mel", - "f0_encoder", - ] - - EXTRA_BODY_PARAMS: ClassVar[frozenset[str]] = frozenset( - { - "prompt_wav_path", - "target_wav_path", - "prompt_f0_path", - "target_f0_path", - "prompt_audio", - "target_audio", - "vocal_sep", - "preprocess_weights_dir", - "preprocess_verbose", - "auto_shift", - "pitch_shift", - } - ) - - EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]] = frozenset({"pitch_shift"}) - - def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): - super().__init__(od_config=od_config, prefix=prefix) - - self.f0_encoder = nn.Embedding(self.f0_bin, self.f0_dim) - - self.cfm_decoder = CFMDecoder(self.flow_matching_config) - - self.mel, self.vocoder = self._build_fp32_audio_modules(self.audio_config) - with set_default_torch_dtype(torch.float32): - self.whisper_encoder = WhisperEncoder(device=self.device) - - self._setup_soulx_profiler() - - def _encode_condition(self, *, whisper_features: torch.Tensor, f0_coarse: torch.Tensor) -> torch.Tensor: - cond = whisper_features + self.f0_encoder(f0_coarse) - return self._to_trunk_dtype(cond)[0] - - def _encode_prompt_whisper_feature(self, prompt_wav: torch.Tensor) -> torch.Tensor: - trunk_dtype = self.f0_encoder.weight.dtype - with self._stage_timer("whisper"): - return self.whisper_encoder.encode( - prompt_wav, - sr=self.audio_config.sample_rate, - output_dtype=trunk_dtype, - ) - - _encode_target_whisper_feature = _encode_prompt_whisper_feature - - def _infer_segment( - self, - *, - prompt_mel: torch.Tensor, - prompt_wav: torch.Tensor, - target_wav: torch.Tensor, - prompt_f0: torch.Tensor, - target_f0: torch.Tensor, - pitch_shift: int, - num_inference_steps: int, - guidance_scale: float, - prompt_feature: torch.Tensor | None = None, - generator: torch.Generator | None = None, - ) -> torch.Tensor: - """Single-segment SVC inference aligned with ``SoulXSingerSVC.infer_segment``.""" - len_prompt_mel = prompt_mel.shape[1] - prompt_f0 = F.pad(prompt_f0, (0, 0, 0, max(0, len_prompt_mel - prompt_f0.shape[1])))[:, :len_prompt_mel] - - f0_coarse_prompt = f0_to_coarse(prompt_f0) - f0_coarse_target = f0_to_coarse(target_f0, f0_shift=int(pitch_shift * 5)) - f0_coarse = torch.cat([f0_coarse_prompt, f0_coarse_target], dim=1) - - trunk_dtype = self.f0_encoder.weight.dtype - if prompt_feature is None: - prompt_feature = self._encode_prompt_whisper_feature(prompt_wav) - with self._stage_timer("whisper"): - target_feature = self.whisper_encoder.encode( - target_wav, - sr=self.audio_config.sample_rate, - output_dtype=trunk_dtype, - ) - prompt_feature = F.pad( - prompt_feature, - (0, 0, 0, max(0, f0_coarse_prompt.shape[1] - prompt_feature.shape[1])), - )[:, : f0_coarse_prompt.shape[1], :] - target_feature = F.pad( - target_feature, - (0, 0, 0, max(0, f0_coarse_target.shape[1] - target_feature.shape[1])), - )[:, : f0_coarse_target.shape[1], :] - - whisper_features = torch.cat([prompt_feature, target_feature], dim=1) - with self._stage_timer("cond_encode"): - cond = self._encode_condition(whisper_features=whisper_features, f0_coarse=f0_coarse) - - with self._stage_timer("cfm"): - generated_mel = self._run_flow_matching_loop( - prompt=prompt_mel, - cond=cond, - n_timesteps=num_inference_steps, - cfg=guidance_scale, - generator=generator, - ) - - with self._stage_timer("vocoder"): - generated_audio = self._mel_to_audio(generated_mel, squeeze=True) - target_len = target_wav.shape[-1] - if generated_audio.shape[-1] > target_len: - generated_audio = generated_audio[:target_len] - elif generated_audio.shape[-1] < target_len: - generated_audio = F.pad(generated_audio, (0, target_len - generated_audio.shape[-1])) - return generated_audio - - def _build_vocal_segments( - self, - f0: torch.Tensor, - *, - hop_size: int, - sample_rate: int, - uv_frames_th: int = 10, - min_duration_sec: float = 15.0, - max_duration_sec: float = 30.0, - ignore_silent: bool = True, - ) -> tuple[list[tuple[float, float]], list[tuple[float, float]]]: - """Build vocal segments based on F0 contour. - - Mirrors upstream ``SoulXSingerSVC.build_vocal_segments``. - Returns (overlap_segments, segments) tuples of (start_sec, end_sec). - """ - f0_np = ( - f0.detach().float().cpu().numpy().squeeze() if isinstance(f0, torch.Tensor) else np.asarray(f0).squeeze() - ) - total_frames = int(f0_np.shape[0]) - if total_frames == 0: - return [], [] - - f0_rate = sample_rate // hop_size - min_frames = max(1, int(round(min_duration_sec * f0_rate))) - max_frames = max(1, int(round(max_duration_sec * f0_rate))) - - split_points = [0] - - def append_split_point(point: int) -> None: - point = int(max(0, min(point, total_frames))) - while point - split_points[-1] > max_frames: - split_points.append(split_points[-1] + max_frames) - if point > split_points[-1]: - split_points.append(point) - - idx = 0 - while idx < total_frames: - if f0_np[idx] == 0: - run_start = idx - while idx < total_frames and f0_np[idx] == 0: - idx += 1 - run_end = idx - if (run_end - run_start) >= uv_frames_th: - split_point = max(run_end - 5, (run_start + run_end) // 2) - append_split_point(split_point) - else: - idx += 1 - append_split_point(total_frames) - - segments: list[tuple[int, int]] = [] - overlap_segments: list[tuple[int, int]] = [] - num_overlaps = 1 - - def append_segment(start_idx: int, end_idx: int, overlaps: int = num_overlaps) -> None: - segments.append((split_points[start_idx], split_points[end_idx])) - overlap_start_idx = start_idx - if start_idx > 0 and (split_points[end_idx] - split_points[start_idx - overlaps]) <= max_frames: - overlap_start_idx = start_idx - overlaps - overlap_segments.append((split_points[overlap_start_idx], split_points[end_idx])) - - seg_start, seg_end = 0, 1 - while seg_start < len(split_points) - 1: - while seg_end < len(split_points) and (split_points[seg_end] - split_points[seg_start]) < min_frames: - seg_end += 1 - if seg_end >= len(split_points): - append_segment(seg_start, len(split_points) - 1) - break - append_segment(seg_start, seg_end) - seg_start = seg_end - seg_end = seg_start + 1 - - if ignore_silent: - filtered_idx = [] - for i, (ov_start, ov_end) in enumerate(overlap_segments): - voice_ratio = np.sum(f0_np[ov_start:ov_end] > 0) / max(1, ov_end - ov_start) - voiced_frames = np.sum(f0_np[ov_start:ov_end] > 0) - if voice_ratio > 0.05 and voiced_frames >= 10: - filtered_idx.append(i) - overlap_segments = [overlap_segments[i] for i in filtered_idx] - segments = [segments[i] for i in filtered_idx] - - # convert indices to seconds - overlap_sec = [(s / f0_rate, e / f0_rate) for s, e in overlap_segments] - seg_sec = [(s / f0_rate, e / f0_rate) for s, e in segments] - return overlap_sec, seg_sec - - def infer_svc_batch( - self, - payload: dict[str, Any], - *, - extra_args: dict[str, Any], - num_inference_steps: int, - guidance_scale: float, - generator: torch.Generator | None = None, - ) -> tuple[torch.Tensor, int]: - """Batch SVC with chunked inference for long audio (matches upstream).""" - prompt_wav = payload["prompt_wav"] - target_wav = payload["target_wav"] - prompt_f0 = payload["prompt_f0"] - target_f0 = payload["target_f0"] - - pitch_shift = resolve_pitch_shift( - auto_shift=bool(extra_args.get("auto_shift", True)), - manual_shift=int(extra_args.get("pitch_shift", 0)), - prompt_f0=prompt_f0, - target_f0=target_f0, - ) - - prompt_mel = self._mel_from_wav(prompt_wav) - - # Long audio: chunk by vocal segments (same as upstream) - max_sec = 30.0 * self.audio_config.sample_rate - if target_wav.shape[-1] < max_sec: - # Short audio: one-shot - generated_audio = self._infer_segment( - prompt_mel=prompt_mel, - prompt_wav=prompt_wav, - target_wav=target_wav, - prompt_f0=prompt_f0, - target_f0=target_f0, - pitch_shift=pitch_shift, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - generator=generator, - ) - return generated_audio.unsqueeze(0), pitch_shift - - # Build vocal segments - overlap_segments, segments = self._build_vocal_segments( - target_f0, - hop_size=self.audio_config.hop_size, - sample_rate=self.audio_config.sample_rate, - ) - if len(segments) == 0: - segments = [(0.0, target_wav.shape[-1] / self.audio_config.sample_rate)] - overlap_segments = [(0.0, target_wav.shape[-1] / self.audio_config.sample_rate)] - - f0_rate = self.audio_config.sample_rate // self.audio_config.hop_size - generated_audio = torch.zeros_like(target_wav) - - for idx in range(len(segments)): - overlap_start_sec, overlap_end_sec = overlap_segments[idx] - seg_start_sec, seg_end_sec = segments[idx] - - wav_start = int(round(overlap_start_sec * self.audio_config.sample_rate)) - wav_end = int(round(overlap_end_sec * self.audio_config.sample_rate)) - f0_start = int(round(overlap_start_sec * f0_rate)) - f0_end = int(round(overlap_end_sec * f0_rate)) - - wav_start = max(0, min(wav_start, target_wav.shape[-1])) - wav_end = max(wav_start, min(wav_end, target_wav.shape[-1])) - f0_start = max(0, min(f0_start, target_f0.shape[-1])) - f0_end = max(f0_start, min(f0_end, target_f0.shape[-1])) - - segment_gt_wav = target_wav[:, wav_start:wav_end] - segment_gt_f0 = target_f0[:, f0_start:f0_end] - - segment_generated = self._infer_segment( - prompt_mel=prompt_mel, - prompt_wav=prompt_wav, - target_wav=segment_gt_wav, - prompt_f0=prompt_f0, - target_f0=segment_gt_f0, - pitch_shift=pitch_shift, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - generator=generator, - ) - - seg_start = int(round(seg_start_sec * self.audio_config.sample_rate)) - seg_end = int(round(seg_end_sec * self.audio_config.sample_rate)) - segment_generated = segment_generated[seg_start - wav_start : seg_end - wav_start] - generated_audio[:, seg_start:seg_end] = segment_generated - - return generated_audio.unsqueeze(0), pitch_shift - - @torch.inference_mode() - def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: - return self._forward_batch_from_request( - req, - kind="svc", - metadata_key="pitch_shift", - infer_batch_fn=self.infer_svc_batch, - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - del weights - self._load_soulx_checkpoint("model-svc.pt") - return {name for name, _ in self.named_parameters()} diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py deleted file mode 100644 index fb2c086f2..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py +++ /dev/null @@ -1,355 +0,0 @@ -"""SoulX-Singer SVS (score-driven) pipeline implementation.""" - -from collections.abc import Iterable -from typing import Any, ClassVar - -import torch -import torch.nn as nn -import torch.nn.functional as F -from vllm.logger import init_logger - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.models.soulx_singer.modules import ( - CFMDecoder, - ConvNeXtV2Block, -) -from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import ( - FlowMatchingAudioPipeline, - convert_soulx_audio_output_to_numpy, -) -from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( - SOULX_PREPROCESSED_KEY, - get_soulx_preprocessed_payload, - has_precomputed, -) -from vllm_omni.diffusion.models.soulx_singer.preprocess.pre_process import ( - attach_preprocess_for_diffusion_request, - build_metadata_processor, - is_warmup_request, - normalize_svs_control_extra_args, - resolve_preprocess_audio, -) -from vllm_omni.diffusion.models.soulx_singer.utils import ( - f0_to_coarse, - resolve_pitch_shift, - validate_soulx_extra_args, -) -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch - -logger = init_logger(__name__) - - -def get_soulxsinger_pre_process_func(od_config: OmniDiffusionConfig): - """Validate/load SVS preprocess payload for single-stage or stage-1 DiT.""" - metadata_processor = build_metadata_processor(od_config) - _pipeline = None - - def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: - nonlocal _pipeline - - extra_args = validate_soulx_extra_args( - "svs", - dict(getattr(request.sampling_params, "extra_args", None) or {}), - ) - - # Inline build: when no warmup/no precomputed paths/no IPC payload, - # build the preprocess payload directly from audio file paths. - if not (is_warmup_request(request) or has_precomputed(extra_args, "svs")): - prompt = request.prompt - if not isinstance(prompt, str) and not get_soulx_preprocessed_payload(prompt): # type: ignore[arg-type] - prompt_audio, target_audio = resolve_preprocess_audio(prompt, extra_args) # type: ignore[arg-type] - if prompt_audio is not None and target_audio is not None: - if _pipeline is None: - from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.pipeline import ( - SoulXPreprocessPipeline, - ) - - _pipeline = SoulXPreprocessPipeline( - od_config=od_config, - vocal_sep=bool(extra_args.get("vocal_sep", False)), - verbose=bool(extra_args.get("preprocess_verbose", False)), - extra_args=extra_args, - ) - payload = _pipeline.build_svs_payload_from_audio( - prompt_audio=prompt_audio, - target_audio=target_audio, - metadata_processor=metadata_processor, - language=str(extra_args.get("language", "Mandarin")), - vocal_sep=extra_args.get("vocal_sep"), - prompt_vocal_sep=extra_args.get("prompt_vocal_sep"), - target_vocal_sep=extra_args.get("target_vocal_sep"), - prompt_max_merge_duration_ms=extra_args.get("prompt_max_merge_duration"), - target_max_merge_duration_ms=extra_args.get("target_max_merge_duration"), - ) - payload.setdefault("kind", "svs") - prompt.setdefault("additional_information", {})[SOULX_PREPROCESSED_KEY] = payload # type: ignore[union-attr, assignment] - - return attach_preprocess_for_diffusion_request( - request, - kind="svs", - metadata_processor=metadata_processor, - ) - - return pre_process_func - - -def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig): - """Convert pipeline audio tensor output for offline consumers.""" - - def post_process_func(audio: torch.Tensor): - return convert_soulx_audio_output_to_numpy(audio) - - return post_process_func - - -def _expand_states(h: torch.Tensor, mel2token: torch.Tensor) -> torch.Tensor: - if mel2token.max() > h.size(1) - 1: - logger.warning( - "mel2token.max() (%s) is greater than h.size(1) - 1 (%s); clamping.", - mel2token.max(), - h.size(1) - 1, - ) - mel2token = torch.clamp(mel2token, 0, h.size(1) - 1) - mel2token_ = mel2token[..., None].repeat([1, 1, h.shape[-1]]) - return torch.gather(h, 1, mel2token_) - - -class PipelineSoulXSingerSVS(FlowMatchingAudioPipeline): - """Pipeline for the SoulX-Singer model (SoulX-Singer).""" - - _encoder_modules: ClassVar[list[str]] = [ - "mel", - "f0_encoder", - "preflow", - "note_text_encoder", - "note_pitch_encoder", - "note_type_encoder", - ] - - EXTRA_BODY_PARAMS: ClassVar[frozenset[str]] = frozenset( - { - "prompt_metadata_path", - "target_metadata_path", - "audio_path", - "prompt_audio", - "target_audio", - "language", - "vocal_sep", - "max_merge_duration", - "preprocess_weights_dir", - "preprocess_verbose", - "control", - "auto_shift", - "pitch_shift", - } - ) - - EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]] = frozenset({"f0_shift"}) - - def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): - super().__init__(od_config=od_config, prefix=prefix) - self.metadata_processor = build_metadata_processor(od_config) - - self.f0_encoder = nn.Embedding(self.f0_bin, self.f0_dim) - self.preflow = nn.Sequential( - *[ - ConvNeXtV2Block( - self.text_dim, - self.text_dim * 2, - ) - for _ in range(self.encoder_config.num_layers) - ] - ) - - self.cfm_decoder = CFMDecoder(self.flow_matching_config) - - self.mel, self.vocoder = self._build_fp32_audio_modules(self.audio_config) - - self.note_text_encoder = nn.Embedding(self.vocab_size, self.text_dim) - self.note_pitch_encoder = nn.Embedding(256, self.pitch_dim) - self.note_type_encoder = nn.Embedding(256, self.type_dim) - - self._setup_soulx_profiler() - - def _encode_condition( - self, - *, - note_pitch: torch.Tensor, - note_type: torch.Tensor, - note_text: torch.Tensor, - mel2note: torch.Tensor, - f0_coarse: torch.Tensor, - ) -> torch.Tensor: - features = ( - self.note_pitch_encoder(note_pitch) + self.note_type_encoder(note_type) + self.note_text_encoder(note_text) - ) - features = self.preflow(features) - features = _expand_states(features, mel2note) - features = features + self.f0_encoder(f0_coarse) - return self._to_trunk_dtype(features)[0] - - def _infer_svs_segment( - self, - prompt_meta: dict, - target_meta: dict, - *, - pitch_shift: int, - num_inference_steps: int, - guidance_scale: float, - prompt_mel: torch.Tensor | None = None, - generator: torch.Generator | None = None, - ) -> torch.Tensor: - prompt_note_text = prompt_meta["phoneme"] - prompt_mel2note = prompt_meta["mel2note"] - prompt_note_type = prompt_meta["note_type"] - prompt_note_pitch = prompt_meta["note_pitch"] - prompt_f0 = prompt_meta["f0"] - - target_note_text = target_meta["phoneme"] - target_mel2note = target_meta["mel2note"] - target_note_type = target_meta["note_type"] - target_note_pitch = target_meta["note_pitch"] - target_f0 = target_meta["f0"] - - if target_f0 is None or prompt_f0 is None: - target_f0 = torch.zeros_like(target_mel2note).float() - prompt_f0 = torch.zeros_like(prompt_mel2note).float() - if target_note_pitch is None or prompt_note_pitch is None: - target_note_pitch = torch.zeros_like(target_note_type).int() - prompt_note_pitch = torch.zeros_like(prompt_note_type).int() - - len_prompt_note = prompt_note_pitch.shape[1] - len_prompt_mel = prompt_f0.shape[1] if prompt_f0 is not None else prompt_mel2note.shape[1] - - note_pitch = torch.cat([prompt_note_pitch, target_note_pitch], dim=1) - note_text = torch.cat([prompt_note_text, target_note_text], dim=1) - note_type = torch.cat([prompt_note_type, target_note_type], dim=1) - # Target note indices follow prompt notes in the concatenated score. - mel2note = torch.cat([prompt_mel2note, target_mel2note + len_prompt_note], dim=1) - - # pitch_shift is semitones; each coarse F0 bin is 20 cents (×5 bins per semitone). - f0_coarse_prompt = f0_to_coarse(prompt_f0) - f0_coarse_target = f0_to_coarse(target_f0, f0_shift=pitch_shift * 5) - f0_coarse = torch.cat([f0_coarse_prompt, f0_coarse_target], dim=1) - - note_pitch = note_pitch.clone() - note_pitch[note_pitch > 0] = note_pitch[note_pitch > 0] + pitch_shift - note_pitch = torch.clamp(note_pitch, 0, 255) - - if prompt_mel is None: - prompt_wav = prompt_meta["wav"] - with self._stage_timer("mel"): - prompt_mel = self._mel_from_wav(prompt_wav) - - if prompt_mel.shape[1] > len_prompt_mel: - prompt_mel = prompt_mel[:, :len_prompt_mel, :] - elif prompt_mel.shape[1] < len_prompt_mel: - logger.warning( - "prompt_mel length %s is shorter than metadata frames %s; padding mel.", - prompt_mel.shape[1], - len_prompt_mel, - ) - prompt_mel = F.pad(prompt_mel, (0, 0, 0, len_prompt_mel - prompt_mel.shape[1])) - - with self._stage_timer("cond_encode"): - cond = self._encode_condition( - note_pitch=note_pitch, - note_type=note_type, - note_text=note_text, - mel2note=mel2note, - f0_coarse=f0_coarse, - ) - - with self._stage_timer("cfm"): - generated_mel = self._run_flow_matching_loop( - prompt_mel, - cond, - num_inference_steps, - guidance_scale, - generator=generator, - ) - with self._stage_timer("vocoder"): - return self._mel_to_audio(generated_mel) - - def infer_svs_batch( - self, - payload: dict[str, Any], - *, - extra_args: dict[str, Any], - num_inference_steps: int, - guidance_scale: float, - generator: torch.Generator | None = None, - ) -> tuple[torch.Tensor, int]: - """Pure batch (non-streaming): run diffusion for each target segment in order and concat.""" - prompt_meta_raw = payload.get("prompt_meta") - target_meta_list = payload["target_meta_list"] - control = str(extra_args.get("control", "score")) - - prompt_meta = None - prompt_mel = None - if prompt_meta_raw: - prompt_meta = self._ensure_processed_meta(prompt_meta_raw) - prompt_meta = self._apply_control_mode(prompt_meta, control) - if prompt_meta and prompt_meta.get("wav") is not None: - prompt_mel = self._mel_from_wav(prompt_meta["wav"]) - - pieces: list[torch.Tensor] = [] - pitch_shift = 0 - - for idx, target_raw in enumerate(target_meta_list): - target_meta = self._ensure_processed_meta(target_raw) - target_meta = self._apply_control_mode(target_meta, control) - - if idx == 0: - pitch_shift = resolve_pitch_shift( - auto_shift=bool(extra_args.get("auto_shift", True)), - manual_shift=int(extra_args.get("pitch_shift", 0)), - prompt_f0=prompt_meta["f0"] if prompt_meta else None, - target_f0=target_meta["f0"], - prompt_note_pitch=prompt_meta["note_pitch"] if prompt_meta else None, - target_note_pitch=target_meta["note_pitch"], - ) - - seg = self._infer_svs_segment( - prompt_meta or {}, - target_meta, - pitch_shift=pitch_shift, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - prompt_mel=prompt_mel, - generator=generator, - ) - pieces.append(seg.squeeze()) - - if not pieces: - return torch.zeros(1, 0, device=self.device, dtype=torch.float32), pitch_shift - - full = torch.cat(pieces, dim=-1) - return full.unsqueeze(0), pitch_shift - - @staticmethod - def _apply_control_mode(meta: dict, control: str) -> dict: - meta["note_pitch"] = meta["note_pitch"] if control == "score" else None - meta["f0"] = meta["f0"] if control == "melody" else None - return meta - - def _ensure_processed_meta(self, meta: dict) -> dict: - if isinstance(meta.get("phoneme"), torch.Tensor): - return meta - return self.metadata_processor.process(meta, None) - - @torch.inference_mode() - def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: - return self._forward_batch_from_request( - req, - kind="svs", - metadata_key="f0_shift", - infer_batch_fn=self.infer_svs_batch, - prepare_extra_args=lambda extra_args, _sampling: normalize_svs_control_extra_args(extra_args), - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - del weights - self._load_soulx_checkpoint("model.pt") - return {name for name, _ in self.named_parameters()} diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/__init__.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/__init__.py deleted file mode 100644 index aed87f757..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/preprocess/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Native SoulX-Singer preprocess helpers.""" diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/g2p.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/g2p.py deleted file mode 100644 index 5f54d47c9..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/preprocess/g2p.py +++ /dev/null @@ -1,86 +0,0 @@ -import re - -_EN_WORD_RE = re.compile(r"^[A-Za-z]+(?:'[A-Za-z]+)*$") -_ZH_WORD_RE = re.compile(r"[\u4e00-\u9fff]") - -EN_FLAG = "en_" -YUE_FLAG = "yue_" -ZH_FLAG = "zh_" - -_g2p_zh = None -_g2p_en = None - - -def _ensure_g2p_backends() -> None: - global _g2p_zh, _g2p_en - if _g2p_zh is None: - from g2pM import G2pM - - _g2p_zh = G2pM() - if _g2p_en is None: - from g2p_en import G2p as G2pE - - _g2p_en = G2pE() - - -def is_chinese_char(word: str) -> bool: - if len(word) != 1: - return False - return bool(_ZH_WORD_RE.fullmatch(word)) - - -def is_english_word(word: str) -> bool: - if not word: - return False - return bool(_EN_WORD_RE.fullmatch(word)) - - -def g2p_cantonese(sent): - import ToJyutping - - return ToJyutping.get_jyutping_list(sent) - - -def g2p_mandarin(sent): - _ensure_g2p_backends() - return _g2p_zh(sent, tone=True, char_split=False) - - -def g2p_english(word): - _ensure_g2p_backends() - return _g2p_en(word) - - -def g2p_transform(words, lang): - zh_words = [] - transformed_words = [0] * len(words) - - for idx, w in enumerate(words): - if w == "": - transformed_words[idx] = w - continue - - w = w.replace("?", "").replace(".", "").replace("!", "").replace(",", "") - - if is_chinese_char(w): - zh_words.append([idx, w]) - else: - if is_english_word(w): - w = EN_FLAG + "-".join(g2p_english(w.lower())) - else: - w = "" - transformed_words[idx] = w - - sent = "".join([k[1] for k in zh_words]) - - if len(sent) > 0: - if lang == "Cantonese": - g2pm_rst = g2p_cantonese(sent) - g2pm_rst = [YUE_FLAG + k[1] for k in g2pm_rst] - else: - g2pm_rst = g2p_mandarin(sent) - g2pm_rst = [ZH_FLAG + k for k in g2pm_rst] - for p, w in zip([k[0] for k in zh_words], g2pm_rst): - transformed_words[p] = w - - return transformed_words diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/metadata_utils.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/metadata_utils.py deleted file mode 100644 index 8a8611121..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/preprocess/metadata_utils.py +++ /dev/null @@ -1,137 +0,0 @@ -"""Metadata merge helpers for SoulX-Singer preprocess.""" - -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import numpy as np -import soundfile as sf - - -@dataclass -class SegmentMetadata: - item_name: str - wav_fn: str - language: str - start_time_ms: int - end_time_ms: int - note_text: list[str] - note_dur: list[float] - note_pitch: list[int] - note_type: list[int] - origin_wav_fn: str | None = None - wav: np.ndarray | None = None - - def get(self, name: str, default: Any = None) -> Any: - return getattr(self, name, default) - - -def _merge_group( - audio: np.ndarray, - sample_rate: int, - segments: list[SegmentMetadata | dict[str, Any]], - output_dir: Path | None, - *, - end_extension_ms: int = 0, -) -> SegmentMetadata: - if not segments: - raise ValueError("segments must not be empty") - - words: list[str] = [] - durs: list[float] = [] - pitches: list[int] = [] - types: list[int] = [] - - for i, seg in enumerate(segments): - if i > 0: - prev_seg = segments[i - 1] - gap_ms = seg.get("start_time_ms", 0) - prev_seg.get("end_time_ms", 0) - if gap_ms > 0: - words.append("") - durs.append(gap_ms / 1000.0) - pitches.append(0) - types.append(1) - - words.extend(seg.get("note_text", [])) - durs.extend(seg.get("note_dur", [])) - pitches.extend(seg.get("note_pitch", [])) - types.extend(seg.get("note_type", [])) - - if end_extension_ms > 0: - words.append("") - durs.append(end_extension_ms / 1000.0) - pitches.append(0) - types.append(1) - - merged_words, merged_durs, merged_pitches, merged_types = [], [], [], [] - for word, dur, pitch, note_type in zip(words, durs, pitches, types): - if merged_words and word == "" and merged_words[-1] == "": - merged_durs[-1] += dur - else: - merged_words.append(word) - merged_durs.append(dur) - merged_pitches.append(pitch) - merged_types.append(note_type) - - languages = [s.get("language", "Mandarin") for s in segments if s.get("language")] - language = max(languages, key=languages.count) if languages else "Mandarin" - - start_ms = int(segments[0].get("start_time_ms", 0)) - end_ms = int(segments[-1].get("end_time_ms", 0)) + end_extension_ms - start_sample = int(start_ms * sample_rate // 1000) - end_sample = int(end_ms * sample_rate // 1000) - - first_item_name = segments[0].get("item_name", "segment") - song_prefix = "_".join(str(first_item_name).split("_")[:-1]) - item_name = f"{song_prefix}_{start_ms}_{end_ms}" - - wav_fn = "" - segment_audio = audio[start_sample:end_sample] - if output_dir is not None: - output_dir.mkdir(parents=True, exist_ok=True) - wav_path = output_dir / f"{item_name}.wav" - sf.write(wav_path, segment_audio, sample_rate) - wav_fn = str(wav_path) - - return SegmentMetadata( - item_name=item_name, - wav_fn=wav_fn, - language=language, - start_time_ms=start_ms, - end_time_ms=end_ms, - note_text=merged_words, - note_dur=merged_durs, - note_pitch=merged_pitches, - note_type=merged_types, - origin_wav_fn=segments[0].get("origin_wav_fn", ""), - wav=segment_audio if output_dir is None else None, - ) - - -def convert_metadata(item: SegmentMetadata, f0: np.ndarray | None = None) -> dict[str, Any]: - """Convert merged segment metadata into SoulX JSON metadata format. - - Matches upstream ``convert_metadata``: F0 always comes from ``*_f0.npy`` written - by ``extract_f0_file`` on the merged segment wav, not from an in-memory array. - """ - from vllm_omni.diffusion.models.soulx_singer.preprocess.g2p import g2p_transform - - if item.wav_fn: - f0_path = str(Path(item.wav_fn).with_suffix("")) + "_f0.npy" - f0_arr = np.load(f0_path) - elif f0 is not None: - f0_arr = f0 - else: - raise ValueError("convert_metadata requires item.wav_fn with sidecar f0 or an explicit f0 array") - - return { - "index": item.item_name, - "language": item.language, - "time": [item.start_time_ms, item.end_time_ms], - "duration": " ".join(f"{d:.2f}" for d in item.note_dur), - "text": " ".join(item.note_text), - "phoneme": " ".join(g2p_transform(item.note_text, item.language)), - "note_pitch": " ".join(map(str, item.note_pitch)), - "note_type": " ".join(map(str, item.note_type)), - "f0": " ".join(f"{x:.1f}" for x in f0_arr), - } diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py deleted file mode 100644 index ff4e4b7ca..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py +++ /dev/null @@ -1,106 +0,0 @@ -"""SoulX-Singer payload schema (guidance/precomputed tables, dummy builders).""" - -from __future__ import annotations - -from typing import Any - -import torch - -SOULX_PRECOMPUTED_KEYS_BY_KIND: dict[str, tuple[str, ...]] = { - "svs": ("prompt_metadata_path", "target_metadata_path", "audio_path"), - "svc": ("prompt_wav_path", "target_wav_path", "prompt_f0_path", "target_f0_path"), -} - - -# Key under which the SoulX preprocess payload is stored in -# prompt["additional_information"] or diffusion output metadata. -SOULX_PREPROCESSED_KEY = "soulx_preprocessed" - - -def get_soulx_preprocessed_payload(prompt: dict[str, Any]) -> dict[str, Any] | None: - """Extract attached payload from a diffusion request prompt dict.""" - additional = prompt.get("additional_information") or {} - payload = additional.get(SOULX_PREPROCESSED_KEY) - return payload if isinstance(payload, dict) else None - - -def relocate_tensors(payload: dict[str, Any], device: torch.device | str) -> dict[str, Any]: - """Recursively move all tensors inside the payload to the target device. - - Called after decode when the payload is about to be consumed by the - diffusion model (which may live on a different device / rank). - """ - if isinstance(payload, torch.Tensor): - return payload.to(device, non_blocking=True) - if isinstance(payload, dict): - return {key: relocate_tensors(value, device) for key, value in payload.items()} - if isinstance(payload, list): - return [relocate_tensors(item, device) for item in payload] - return payload - - -def consume_payload( - req: Any, - expected_kind: str, - device: torch.device | str, -) -> dict[str, Any]: - """Read the payload attached by pre_process_func, validate the kind, - and relocate tensors to the compute device. - """ - prompts = getattr(req, "prompts", None) or [] - if not prompts: - raise ValueError(f"SoulX-Singer {expected_kind} request has no prompts.") - prompt = prompts[0] - if isinstance(prompt, str): - raise ValueError(f"SoulX-Singer {expected_kind} forward requires pre_process_func output on the prompt.") - payload = get_soulx_preprocessed_payload(prompt) - if payload is None or payload.get("kind") != expected_kind: - raise ValueError( - f"SoulX-Singer {expected_kind} forward requires additional_information['soulx_preprocessed'] " - "produced by pre_process_func." - ) - return relocate_tensors(payload, device) - - -def has_precomputed(extra_args: dict[str, Any], kind: str) -> bool: - keys = SOULX_PRECOMPUTED_KEYS_BY_KIND[kind] - return all(extra_args.get(key) for key in keys) - - -def build_dummy_payload(kind: str, device: torch.device) -> dict[str, Any]: - mel_frames = 4 - hop_size = 480 - n_fft = 1920 - wav_samples = max(mel_frames * hop_size, n_fft) - if kind == "svc": - voiced_f0 = torch.full((1, mel_frames), 100.0, device=device, dtype=torch.float64) - return { - "kind": "svc", - "prompt_wav": torch.zeros(1, wav_samples, device=device, dtype=torch.float32), - "target_wav": torch.zeros(1, wav_samples, device=device, dtype=torch.float32), - "prompt_f0": voiced_f0.clone(), - "target_f0": voiced_f0.clone(), - } - return { - "kind": "svs", - "prompt_meta": { - "phoneme": torch.zeros(1, 2, device=device, dtype=torch.long), - "note_pitch": torch.zeros(1, mel_frames, device=device, dtype=torch.long), - "note_type": torch.zeros(1, mel_frames, device=device, dtype=torch.long), - "mel2note": torch.zeros(1, mel_frames, device=device, dtype=torch.long), - "f0": torch.zeros(1, mel_frames, device=device, dtype=torch.float32), - "wav": torch.zeros(1, wav_samples, device=device, dtype=torch.float32), - }, - "target_meta_list": [ - { - "text": "dummy", - "phoneme": " ".join([""] * mel_frames), - "f0": " ".join(["100"] * mel_frames), - "language": "Mandarin", - "note_pitch": " ".join(["0"] * mel_frames), - "note_type": " ".join(["1"] * mel_frames), - "duration": " ".join(["1"] * mel_frames), - "time": [0, max(int(wav_samples / 24000 * 1000), 100)], - } - ], - } diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/pre_process.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/pre_process.py deleted file mode 100644 index a8784fcba..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/preprocess/pre_process.py +++ /dev/null @@ -1,223 +0,0 @@ -"""Shared ``pre_process_func`` helpers for SoulX-Singer pipelines (batch / full-payload only).""" - -from pathlib import Path -from typing import Any - -import torch -import torch.nn as nn -from vllm.logger import init_logger - -from vllm_omni.diffusion.data import OmniDiffusionConfig -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.pipeline import SoulXPreprocessPipeline -from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( - SOULX_PRECOMPUTED_KEYS_BY_KIND, - SOULX_PREPROCESSED_KEY, - build_dummy_payload, - get_soulx_preprocessed_payload, - has_precomputed, -) -from vllm_omni.diffusion.models.soulx_singer.utils import ( - MetadataProcessor, - load_config, - resolve_phoneset_path, - validate_soulx_extra_args, -) -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.inputs.data import OmniTextPrompt - -logger = init_logger(__name__) - - -def is_warmup_request(request: OmniDiffusionRequest) -> bool: - return request.is_dummy_run() - - -def resolve_preprocess_audio(prompt: dict[str, Any], extra_args: dict[str, Any]) -> tuple[Any, Any]: - mm = prompt.get("multi_modal_data") or {} - prompt_audio = mm.get("audio") if isinstance(mm, dict) else None - if prompt_audio is None: - prompt_audio = extra_args.get("prompt_audio") - return prompt_audio, extra_args.get("target_audio") - - -def _preprocess_inputs_missing_error(kind: str) -> ValueError: - return ValueError( - f"SoulX-Singer {kind} requires precomputed paths " - f"{list(SOULX_PRECOMPUTED_KEYS_BY_KIND[kind])}, an attached soulx_preprocessed payload " - f"(multi-stage preprocess stage), or run with pipeline soulxsinger_{kind}." - ) - - -def build_precomputed_payload( - kind: str, - extra_args: dict[str, Any], - *, - metadata_processor, - sample_rate: int | None, - device, -) -> dict[str, Any]: - if kind == "svs": - return SoulXPreprocessPipeline.build_svs_payload_from_paths( - prompt_metadata_path=str(extra_args["prompt_metadata_path"]), - target_metadata_path=str(extra_args["target_metadata_path"]), - audio_path=str(extra_args["audio_path"]), - metadata_processor=metadata_processor, - ) - if sample_rate is None or device is None: - raise ValueError("svc precomputed payload requires sample_rate and device") - return SoulXPreprocessPipeline.build_svc_payload_from_paths( # type: ignore - prompt_wav_path=str(extra_args["prompt_wav_path"]), - target_wav_path=str(extra_args["target_wav_path"]), - prompt_f0_path=str(extra_args["prompt_f0_path"]), - target_f0_path=str(extra_args["target_f0_path"]), - sample_rate=int(sample_rate), - device=device, - ) - - -def build_warmup_payload( - kind: str, - *, - metadata_processor, - device, - sample_rate: int | None, -) -> dict[str, Any]: - if kind == "svs": - if metadata_processor is None: - raise ValueError("SVS warmup requires metadata_processor") - payload = build_dummy_payload("svs", torch.device("cpu")) - dummy_prompt = payload["prompt_meta"] - processed = metadata_processor.process(dict(payload["target_meta_list"][0])) - payload["prompt_meta"] = { - key: value.clone() if isinstance(value, torch.Tensor) else value for key, value in processed.items() - } - if isinstance(dummy_prompt.get("wav"), torch.Tensor): - payload["prompt_meta"]["wav"] = dummy_prompt["wav"].clone() - return payload - if device is None or sample_rate is None: - raise ValueError("svc warmup requires device and sample_rate") - return build_dummy_payload("svc", device) - - -def build_preprocess_payload( - kind: str, - *, - prompt: dict[str, Any], - extra_args: dict[str, Any], - preprocess: nn.Module, - metadata_processor, - device, - sample_rate: int, -) -> dict[str, Any]: - if has_precomputed(extra_args, kind): - return build_precomputed_payload( - kind, - extra_args, - metadata_processor=metadata_processor, - sample_rate=sample_rate, - device=device, - ) - - prompt_audio, target_audio = resolve_preprocess_audio(prompt, extra_args) - if prompt_audio is None or target_audio is None: - raise _preprocess_inputs_missing_error(kind) - - if kind == "svc": - return preprocess.build_svc_payload_from_audio( # type: ignore - prompt_audio=prompt_audio, - target_audio=target_audio, - sample_rate=sample_rate, - device=device, - vocal_sep=extra_args.get("vocal_sep"), - ) - - return preprocess.build_svs_payload_from_audio( # type: ignore - prompt_audio=prompt_audio, - target_audio=target_audio, - metadata_processor=metadata_processor, - language=str(extra_args.get("language", "Mandarin")), - vocal_sep=extra_args.get("vocal_sep"), - prompt_vocal_sep=extra_args.get("prompt_vocal_sep"), - target_vocal_sep=extra_args.get("target_vocal_sep"), - prompt_max_merge_duration_ms=extra_args.get("prompt_max_merge_duration"), - target_max_merge_duration_ms=extra_args.get("target_max_merge_duration"), - ) - - -def attach_preprocess_for_diffusion_request( - request: OmniDiffusionRequest, - *, - kind: str, - metadata_processor=None, - sample_rate: int | None = None, - device=None, -) -> OmniDiffusionRequest: - """Resolve preprocess payload for stage-1 diffusion (warmup / IPC / precomputed paths).""" - - extra_args = validate_soulx_extra_args( - kind, - dict(getattr(request.sampling_params, "extra_args", None) or {}), - ) - prompt = request.prompt - prompt = OmniTextPrompt(prompt=prompt) if isinstance(prompt, str) else prompt - - if is_warmup_request(request): - request.sampling_params.num_inference_steps = 1 - payload = build_warmup_payload( - kind, - metadata_processor=metadata_processor, - device=device, - sample_rate=sample_rate, - ) - elif get_soulx_preprocessed_payload(prompt): - request.prompt = prompt - if kind == "svs": - extra_args = normalize_svs_control_extra_args(extra_args) - request.sampling_params.extra_args = extra_args - return request - elif has_precomputed(extra_args, kind): - payload = build_precomputed_payload( - kind, - extra_args, - metadata_processor=metadata_processor, - sample_rate=sample_rate, - device=device, - ) - else: - raise _preprocess_inputs_missing_error(kind) - - if payload.get("kind") != kind: - raise ValueError(f"Invalid {kind} preprocess payload kind: {payload.get('kind')}") - prompt.setdefault("additional_information", {})[SOULX_PREPROCESSED_KEY] = payload - request.prompt = prompt - - if kind == "svs": - extra_args = normalize_svs_control_extra_args(extra_args) - request.sampling_params.extra_args = extra_args - return request - - -def normalize_svs_control_extra_args(extra_args: dict[str, Any]) -> dict[str, Any]: - control = extra_args.get("control") - if control is None: - control = "score" - logger.warning("control is not provided, using 'score' as default") - elif control not in ("score", "melody"): - raise ValueError(f"Invalid control: {control}. Must be one of: ['score', 'melody']") - extra_args["control"] = control - return extra_args - - -def build_metadata_processor(od_config: OmniDiffusionConfig): - model_dir = od_config.model - config_path = Path(model_dir) / "config.yaml" - hf_config = load_config(str(config_path)) - audio_config = hf_config.audio - - return MetadataProcessor( - hop_size=audio_config.hop_size, - sample_rate=audio_config.sample_rate, - phoneset_path=resolve_phoneset_path(model_dir), - device=str(get_local_device()), - ) diff --git a/vllm_omni/diffusion/models/soulx_singer/utils.py b/vllm_omni/diffusion/models/soulx_singer/utils.py deleted file mode 100644 index 4f201c407..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/utils.py +++ /dev/null @@ -1,565 +0,0 @@ -import json -from pathlib import Path -from typing import Any, Literal - -import numpy as np -import torch -import torchaudio -from huggingface_hub import snapshot_download -from omegaconf import DictConfig, OmegaConf -from vllm.logger import init_logger - -from vllm_omni.diffusion.data import OmniDiffusionConfig - -logger = init_logger(__name__) - -_MODEL_WEIGHTS_REPO = "Soul-AILab/SoulX-Singer" -_PREPROCESS_WEIGHTS_REPO = "Soul-AILab/SoulX-Singer-Preprocess" - - -def _patch_torchaudio_load() -> None: - """Borrowed from MOSS-TTS-Nano. Patch torchaudio.load to use soundfile - if torchcodec is unavailable. - """ - try: - import torchaudio - - torchaudio # noqa - import torchcodec # noqa: F401 - - return - except Exception: - pass - - import soundfile as sf - - def _soundfile_load(path, frame_offset=0, num_frames=-1, normalize=True, channels_first=True, format=None): - data, sr = sf.read(str(path), dtype="float32", always_2d=True) - if frame_offset > 0: - data = data[frame_offset:] - if num_frames > 0: - data = data[:num_frames] - waveform = torch.from_numpy(data) - if channels_first: - waveform = waveform.T - return waveform, sr - - def _soundfile_save(path, src, sample_rate, channels_first=True, **kwargs): - wav = src.detach().cpu().float().numpy() - if channels_first and wav.ndim == 2: - wav = wav.T - sf.write(str(path), wav, sample_rate) - - try: - import torchaudio - - torchaudio.load = _soundfile_load - torchaudio.save = _soundfile_save - logger.info("Patched torchaudio.load/save to use soundfile (torchcodec unavailable)") - except Exception as e: - logger.warning("Could not patch torchaudio: %s", e) - - -def load_config(config_path: str | Path) -> DictConfig: - """ - Load a configuration file and optionally merges it with a base configuration. - - Args: - config_path (Path): Path to the configuration file. - """ - # Load the initial configuration from the given path - config = OmegaConf.load(str(config_path)) - - # Check if there is a base configuration specified and merge if necessary - if config.get("base_config", None) is not None: - base_config = OmegaConf.load(str(config.get("base_config"))) - config = OmegaConf.merge(base_config, config) - - return config - - -SoulXKind = Literal["svs", "svc"] - -_SVS_ARCHITECTURE = "SoulXSingerPipeline" -_SVC_ARCHITECTURE = "SoulXSingerSVCPipeline" -_ARCHITECTURE_TO_KIND: dict[str, SoulXKind] = { - _SVS_ARCHITECTURE: "svs", - _SVC_ARCHITECTURE: "svc", -} - - -def load_model_config_json(model_dir: str | Path) -> dict[str, Any]: - """Load ``config.json`` from a SoulX-Singer model view directory.""" - cfg_path = Path(model_dir) / "config.json" - if not cfg_path.is_file(): - raise FileNotFoundError( - f"SoulX-Singer config.json not found at {cfg_path}. " - "Create one with architectures ['SoulXSingerPipeline'] (SVS) or " - "['SoulXSingerSVCPipeline'] (SVC); see examples/offline_inference/text_to_speech/README.md." - ) - with cfg_path.open(encoding="utf-8") as f: - return json.load(f) - - -def resolve_soulx_kind(model_dir: str | Path) -> SoulXKind: - """Resolve SVS vs SVC from ``config.json`` ``architectures[0]``.""" - cfg = load_model_config_json(model_dir) - archs = cfg.get("architectures") or [] - if not archs: - raise ValueError( - f"SoulX-Singer config.json at {Path(model_dir) / 'config.json'} is missing 'architectures'. " - "Set architectures to ['SoulXSingerPipeline'] for SVS or ['SoulXSingerSVCPipeline'] for SVC." - ) - arch = archs[0] - kind = _ARCHITECTURE_TO_KIND.get(arch) - if kind is None: - raise ValueError( - f"Unknown SoulX-Singer architecture {arch!r} in config.json. Expected one of {list(_ARCHITECTURE_TO_KIND)}." - ) - return kind - - -def validate_soulx_extra_args(kind: SoulXKind, extra_args: dict[str, Any] | None) -> dict[str, Any]: - """Validate request ``extra_args`` against the model mode from config.json.""" - from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( - SOULX_PRECOMPUTED_KEYS_BY_KIND, - has_precomputed, - ) - - extra_args = dict(extra_args or {}) - declared = extra_args.pop("kind", None) - if declared is not None and declared != kind: - logger.warning( - "Ignoring extra_args kind=%r; model config.json declares %r.", - declared, - kind, - ) - - other_kind: SoulXKind = "svc" if kind == "svs" else "svs" - if has_precomputed(extra_args, other_kind): - raise ValueError( - f"SoulX-Singer config.json declares {kind!r}, but request supplies " - f"{other_kind} precomputed paths {list(SOULX_PRECOMPUTED_KEYS_BY_KIND[other_kind])}." - ) - - if kind == "svc": - for key in ("language", "control"): - if extra_args.get(key) is not None: - logger.warning( - "Ignoring SVS-only extra_arg %r=%r on an SVC model (config.json).", - key, - extra_args.pop(key), - ) - for key in SOULX_PRECOMPUTED_KEYS_BY_KIND["svs"]: - if extra_args.get(key): - raise ValueError(f"SoulX-Singer SVC model (config.json) does not accept SVS precomputed key {key!r}.") - else: - for key in SOULX_PRECOMPUTED_KEYS_BY_KIND["svc"]: - if extra_args.get(key): - raise ValueError(f"SoulX-Singer SVS model (config.json) does not accept SVC precomputed key {key!r}.") - - return extra_args - - -# ---------------- utils for model weights ---------------- - - -def resolve_preprocess_weights_root(od_config: OmniDiffusionConfig) -> Path: - """Locate preprocess weights on disk or download from Hugging Face.""" - - # default local dir - if Path(od_config.model).parent.is_dir(): - local_dir = Path(od_config.model).parent / "SoulX-Singer-Preprocess" - else: - local_dir = Path(__file__).parents[5] / "pretrained" / "SoulX-Singer-Preprocess" - - local_dir = local_dir.expanduser() - if local_dir.is_dir(): - logger.info("Using SoulX preprocess weights from %s", local_dir) - return local_dir - - logger.info( - "SoulX preprocess weights not found locally; downloading %s", - _PREPROCESS_WEIGHTS_REPO, - ) - downloaded = snapshot_download( - _PREPROCESS_WEIGHTS_REPO, - allow_patterns=["*"], - local_dir=local_dir, - ) - return Path(downloaded) - - -def preprocess_weight_paths(weights_root: Path) -> dict[str, str]: - """Map upstream relative paths to absolute paths under ``weights_root``.""" - root = weights_root.resolve() - return { - "rmvpe": str(root / "rmvpe/rmvpe.pt"), - "sep_ckpt": str(root / "mel-band-roformer-karaoke/mel_band_roformer_karaoke_becruily.ckpt"), - "sep_config": str(root / "mel-band-roformer-karaoke/config_karaoke_becruily.yaml"), - "asr_zh": str(root / "speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch"), - "asr_en": str(root / "parakeet-tdt-0.6b-v2/parakeet-tdt-0.6b-v2.nemo"), - "rosvot": str(root / "rosvot/rosvot/model.pt"), - } - - -def resolve_phoneset_path(model_dir: str) -> str: - """Resolve phoneme vocabulary; must be present in model directory.""" - candidates = ( - Path(model_dir) / "phoneme" / "phone_set.json", - Path(model_dir) / "phone_set.json", - ) - for path in candidates: - if path.is_file(): - return str(path) - raise FileNotFoundError( - "SoulX-Singer phoneset not found. Expected one of: " - f"{[str(p) for p in candidates]}. " - "Copy phone_set.json into the model directory. " - "See `examples/offline_inference/text_to_speech/README.md` for instructions." - ) - - -# ---------------- utils for data processing ---------------- - - -def load_wav(wav_path: str, sample_rate: int) -> torch.Tensor: - """Load wav file and resample to target sample rate. - - Args: - wav_path (str): Path to wav file. - sample_rate (int): Target sample rate. - - Returns: - torch.Tensor: Waveform tensor with shape (1, T). - """ - _patch_torchaudio_load() - waveform, sr = torchaudio.load(wav_path) - - if sr != sample_rate: - waveform = torchaudio.functional.resample(waveform, sr, sample_rate) - - if len(waveform.shape) > 1 and waveform.shape[0] > 1: - waveform = torch.mean(waveform, dim=0, keepdim=True) - - return waveform - - -def f0_to_coarse(f0, f0_bin=361, f0_min=32.7031956625, f0_shift=0): - """ - Convert continuous F0 values to discrete F0 bins (SIL and C1 - B6, 361 bins). - args: - f0: continuous F0 values - f0_bin: number of F0 bins - f0_min: minimum F0 value - f0_shift: shift value for F0 bins - returns: - f0_coarse: discrete F0 bins - """ - is_torch = isinstance(f0, torch.Tensor) - uv_mask = f0 <= 0 - - if is_torch: - f0_min_tensor = f0.new_tensor(f0_min) - f0_safe = torch.maximum(f0, f0_min_tensor) - f0_cents = 1200 * torch.log2(f0_safe / f0_min_tensor) - else: - f0_safe = np.maximum(f0, f0_min) - f0_cents = 1200 * np.log2(f0_safe / f0_min) - - f0_coarse = (f0_cents / 20) + 1 - - if is_torch: - f0_coarse = torch.round(f0_coarse).long() - f0_coarse = torch.clamp(f0_coarse, min=1, max=f0_bin - 1) - else: - f0_coarse = np.rint(f0_coarse).astype(int) - f0_coarse = np.clip(f0_coarse, 1, f0_bin - 1) - - f0_coarse[uv_mask] = 0 - - if f0_shift != 0: - if is_torch: - voiced = f0_coarse > 0 - if voiced.any(): - shifted = f0_coarse[voiced] + f0_shift - f0_coarse[voiced] = torch.clamp(shifted, 1, f0_bin - 1) - else: - voiced = f0_coarse > 0 - if np.any(voiced): - shifted = f0_coarse[voiced] + f0_shift - f0_coarse[voiced] = np.clip(shifted, 1, f0_bin - 1) - - return f0_coarse - - -def resolve_pitch_shift( - *, - auto_shift: bool, - manual_shift: int, - prompt_f0: torch.Tensor | None = None, - target_f0: torch.Tensor | None = None, - prompt_note_pitch: torch.Tensor | None = None, - target_note_pitch: torch.Tensor | None = None, -) -> int: - """Resolve semitone shift for auto_shift; ``f0_to_coarse(..., f0_shift=shift * 5)`` consumes it.""" - if not auto_shift or manual_shift != 0: - return int(manual_shift) - - if prompt_note_pitch is not None and target_note_pitch is not None: - target_pitched = target_note_pitch[target_note_pitch >= 1] - prompt_pitched = prompt_note_pitch[prompt_note_pitch >= 1] - if target_pitched.numel() > 0 and prompt_pitched.numel() > 0: - shift = torch.round(prompt_pitched.median() - target_pitched.median()) - if torch.isfinite(shift): - return int(shift.item()) - - if prompt_f0 is not None and target_f0 is not None: - target_voiced = target_f0[target_f0 > 0] - prompt_voiced = prompt_f0[prompt_f0 > 0] - if target_voiced.numel() > 0 and prompt_voiced.numel() > 0: - shift = torch.round(torch.log2(prompt_voiced.median() / target_voiced.median()) * 1200 / 100) - if torch.isfinite(shift): - return int(shift.item()) - - return 0 - - -class MetadataProcessor: - """Data processor for SoulX-Singer""" - - def __init__( - self, - hop_size: int, - sample_rate: int, - phoneset_path: str, - device: str = "cuda", - ): - """Initialize data processor. - - Args: - hop_size (int): Hop size in samples. - sample_rate (int): Sample rate in Hz. - phoneset_path (str): Path to phoneme set JSON file. - device (str): Device to use for tensor operations. - """ - self.hop_size = hop_size - self.sample_rate = sample_rate - self.device = device - self.load_phoneme_id_map(phoneset_path) - - def load_phoneme_id_map(self, phoneset_path: str): - with open(phoneset_path, encoding="utf-8") as f: - phoneset = json.load(f) - self.phone2idx = {ph: idx for idx, ph in enumerate(phoneset)} - - def merge_phoneme(self, meta): - merged_items = [] - - duration = [float(x) for x in meta["duration"].split()] - phoneme = [str(x).replace("", "") for x in meta["phoneme"].split()] - note_pitch = [int(x) for x in meta["note_pitch"].split()] - note_type = [int(x) if phoneme[i] != "" else 1 for i, x in enumerate(meta["note_type"].split())] - - for i in range(len(phoneme)): - if ( - i > 0 - and phoneme[i] == phoneme[i - 1] == "" - and note_type[i] == note_type[i - 1] - and note_pitch[i] == note_pitch[i - 1] - ): - merged_items[-1][1] += duration[i] - else: - merged_items.append([phoneme[i], duration[i], note_pitch[i], note_type[i]]) - - meta["phoneme"] = [x[0] for x in merged_items] - meta["duration"] = [x[1] for x in merged_items] - meta["note_pitch"] = [x[2] for x in merged_items] - meta["note_type"] = [x[3] for x in merged_items] - - return meta - - def preprocess( - self, - note_duration: list[float], - phonemes: list[str], - note_pitch: list[int], - note_type: list[int], - ): - """ - Insert and for each note. - Get aligned indices for each frame. - - Args: - note_duration: Duration of each note in seconds - phonemes: Phoneme sequence for each note - note_pitch: Pitch value for each note - note_type: Type value for each note - - """ - sample_rate = self.sample_rate - hop_size = self.hop_size - duration = sum(note_duration) * sample_rate / hop_size - mel2note = torch.zeros(int(duration), dtype=torch.long, device=self.device) - - ph_locations = [] # idx at mel scale and length - new_phonemes = [] - dur_sum = 0 - - note2origin = [] - - for ph_idx in range(len(phonemes)): - dur = int(np.round(dur_sum * sample_rate / hop_size)) - dur = min(dur, len(mel2note) - 1) - new_phonemes.append("") - note2origin.append(ph_idx) - if phonemes[ph_idx][:3] == "en_": - en_phs = ["en_" + x for x in phonemes[ph_idx][3:].split("-")] + [ - "" - ] # between en words in one note - ph_locations.append([dur, max(1, len(en_phs))]) - new_phonemes.extend(en_phs) - note2origin.extend([ph_idx] * len(en_phs)) - else: - ph_locations.append([dur, 1]) - new_phonemes.append(phonemes[ph_idx]) - note2origin.append(ph_idx) - new_phonemes.append("") - note2origin.append(ph_idx) - dur_sum += note_duration[ph_idx] - - ph_idx = 1 - for idx, (i, j) in enumerate(ph_locations): - next_phoneme_start = ph_locations[idx + 1][0] if idx < len(ph_locations) - 1 else len(mel2note) - if i >= len(mel2note) or i + j > len(mel2note): - break - if i < len(mel2note) and mel2note[i] > 0: - logger.warning(f"warning: overlap of {idx}: {mel2note[i]}") - while i < len(mel2note) and mel2note[i] > 0: - i += 1 - mel2note[i] = ph_idx - k = i + 1 - while k + j < next_phoneme_start: - mel2note[k : k + j] = torch.arange(ph_idx, ph_idx + j, device=self.device) + 1 - k += j - mel2note[next_phoneme_start - 1] = ph_idx + j + 1 - ph_idx += j + 2 # + ph repeats + - - new_phonemes = [""] + new_phonemes - new_note_pitch = [0] + [note_pitch[k] for k in note2origin] - new_note_type = [1] + [note_type[k] for k in note2origin] - - return { - "phoneme": torch.tensor([self.phone2idx[x] for x in new_phonemes], device=self.device).unsqueeze(0), - "note_pitch": torch.tensor(new_note_pitch, device=self.device).unsqueeze(0), - "note_type": torch.tensor(new_note_type, device=self.device).unsqueeze(0), - "mel2note": mel2note.unsqueeze(0), - } - - def process(self, meta: dict, wav_path: str | None = None) -> dict[str, torch.Tensor | None]: - meta = self.merge_phoneme(meta) - - item = self.preprocess( - meta["duration"], - meta["phoneme"], - meta["note_pitch"], - meta["note_type"], - ) - - f0 = [float(x) for x in meta.get("f0", "").split()] - min_frame = min(item["mel2note"].shape[1], len(f0)) if len(f0) > 0 else item["mel2note"].shape[1] - item["f0"] = torch.tensor(f0, device=self.device)[:min_frame].unsqueeze(0).float() if len(f0) > 0 else None - item["mel2note"] = item["mel2note"][:, :min_frame] - - if wav_path is not None: - waveform = load_wav(wav_path, self.sample_rate) - item["wav"] = waveform.to(self.device)[:, : min_frame * self.hop_size] - - return item - - -def build_vocal_segments( - f0, - *, - f0_rate: int = 50, - ignore_silent_frames_thresh: int = 5, - min_duration_sec_per_segment: float = 5.0, - max_duration_sec_per_segment: float = 30.0, - num_overlaps: int = 1, - ignore_silent_frames: bool = True, -) -> tuple[list[tuple[float, float]], list[tuple[float, float]]]: - """Build vocal segments from an F0 contour for chunked SVC inference.""" - if isinstance(f0, torch.Tensor): - f0_np = f0.detach().float().cpu().numpy() - else: - f0_np = np.asarray(f0, dtype=np.float32) - f0_np = np.squeeze(f0_np) - - total_frames = int(f0_np.shape[0]) - if total_frames == 0: - return [], [] - - min_frames = max(1, int(round(min_duration_sec_per_segment * f0_rate))) - max_frames = max(1, int(round(max_duration_sec_per_segment * f0_rate))) - - split_points = [0] - - def append_split_point(point: int): - point = int(max(0, min(point, total_frames))) - while point - split_points[-1] > max_frames: - split_points.append(split_points[-1] + max_frames) - if point > split_points[-1]: - split_points.append(point) - - idx = 0 - while idx < total_frames: - if f0_np[idx] == 0: - run_start = idx - while idx < total_frames and f0_np[idx] == 0: - idx += 1 - run_end = idx - if (run_end - run_start) >= ignore_silent_frames_thresh: - split_point = max(run_end - 5, (run_start + run_end) // 2) - append_split_point(split_point) - else: - idx += 1 - append_split_point(total_frames) - - segments: list[tuple[float, float]] = [] - overlap_segments: list[tuple[float, float]] = [] - - def append_segment(start_idx: int, end_idx: int, overlaps: int = num_overlaps): - segments.append((split_points[start_idx] / f0_rate, split_points[end_idx] / f0_rate)) - overlap_start_idx = start_idx - if start_idx > 0 and (split_points[end_idx] - split_points[start_idx - overlaps]) <= max_frames: - overlap_start_idx = start_idx - overlaps - overlap_segments.append((split_points[overlap_start_idx] / f0_rate, split_points[end_idx] / f0_rate)) - - segment_start, segment_end = 0, 1 - while segment_start < len(split_points) - 1: - while ( - segment_end < len(split_points) and (split_points[segment_end] - split_points[segment_start]) < min_frames - ): - segment_end += 1 - - if segment_end >= len(split_points): - append_segment(segment_start, len(split_points) - 1, overlaps=num_overlaps) - break - append_segment(segment_start, segment_end, overlaps=num_overlaps) - segment_start = segment_end - segment_end = segment_start + 1 - - if ignore_silent_frames: - filtered_idx = [] - for i, seg in enumerate(overlap_segments): - start_frame = int(seg[0] * f0_rate) - end_frame = int(seg[1] * f0_rate) - seg_frames = end_frame - start_frame - voice_frames = np.sum(f0_np[start_frame:end_frame] > 0) - if voice_frames / seg_frames > 0.05 and voice_frames >= 10: - filtered_idx.append(i) - - overlap_segments = [overlap_segments[i] for i in filtered_idx] - segments = [segments[i] for i in filtered_idx] - - return overlap_segments, segments diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 527ad2f74..63f9eb366 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -316,16 +316,6 @@ _DIFFUSION_MODELS = { "pipeline_cosmos3", "Cosmos3OmniDiffusersPipeline", ), - "SoulXSingerPipeline": ( - "soulx_singer", - "pipeline_soulx_singer_svs", - "PipelineSoulXSingerSVS", - ), - "SoulXSingerSVCPipeline": ( - "soulx_singer", - "pipeline_soulx_singer_svc", - "PipelineSoulXSingerSVC", - ), "DiffusersAdapterPipeline": ( "diffusers_adapter", "pipeline_diffusers_adapter", @@ -566,8 +556,6 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "MiniMaxH3Pipeline": "get_minimax_h3_post_process_func", "MiniMaxH3ModularPipeline": "get_minimax_h3_post_process_func", "StableAudioPipeline": "get_stable_audio_post_process_func", - "SoulXSingerPipeline": "get_soulxsinger_post_process_func", - "SoulXSingerSVCPipeline": "get_soulxsinger_post_process_func", "AudioXPipeline": "get_audiox_post_process_func", "WanImageToVideoPipeline": "get_wan22_i2v_post_process_func", "WanS2VPipeline": "get_wan22_s2v_post_process_func", @@ -647,8 +635,6 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "SanaWmPipeline": "get_sana_wm_pre_process_func", "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func", "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func", - "SoulXSingerPipeline": "get_soulxsinger_pre_process_func", - "SoulXSingerSVCPipeline": "get_soulxsinger_svc_pre_process_func", } diff --git a/vllm_omni/model_executor/models/soulx_singer/__init__.py b/vllm_omni/model_executor/models/soulx_singer/__init__.py deleted file mode 100644 index 1ab85a8d3..000000000 --- a/vllm_omni/model_executor/models/soulx_singer/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Keep this package lightweight. The pipeline registry imports pipeline.py -# directly, while diffusion model classes are loaded lazily from their registry. diff --git a/vllm_omni/model_executor/models/soulx_singer/pipeline.py b/vllm_omni/model_executor/models/soulx_singer/pipeline.py deleted file mode 100644 index 4d2420528..000000000 --- a/vllm_omni/model_executor/models/soulx_singer/pipeline.py +++ /dev/null @@ -1,45 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""SoulX-Singer pipeline topologies.""" - -from vllm_omni.config.stage_config import ( - PipelineConfig, - StageExecutionType, - StagePipelineConfig, -) - -SOULXSINGER_SVS_PIPELINE = PipelineConfig( - model_type="soulxsinger_svs", - default_deploy_config_name="soulxsinger_svs.yaml", - model_arch="SoulXSingerPipeline", - hf_architectures=("SoulXSingerPipeline",), - stages=( - StagePipelineConfig( - stage_id=0, - model_stage="dit", - execution_type=StageExecutionType.DIFFUSION, - input_sources=(), - final_output=True, - final_output_type="audio", - model_arch="SoulXSingerPipeline", - ), - ), -) - -SOULXSINGER_SVC_PIPELINE = PipelineConfig( - model_type="soulxsinger_svc", - default_deploy_config_name="soulxsinger_svc.yaml", - model_arch="SoulXSingerSVCPipeline", - hf_architectures=("SoulXSingerSVCPipeline",), - stages=( - StagePipelineConfig( - stage_id=0, - model_stage="dit", - execution_type=StageExecutionType.DIFFUSION, - input_sources=(), - final_output=True, - final_output_type="audio", - model_arch="SoulXSingerSVCPipeline", - ), - ), -)