diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index f2d5e972b..96ca753e2 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -545,6 +545,34 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + - label: "TTS · SoulX-Singer Test" + source_file_dependencies: + - tests/e2e/offline_inference/test_soulxsinger.py + - vllm_omni/diffusion/models/soulx_singer/ + - vllm_omni/deploy/soulxsinger_svc.yaml + - vllm_omni/deploy/soulxsinger_svs.yaml + timeout_in_minutes: 30 + commands: + - | + timeout 30m bash -c " + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_WORKER_MULTIPROC_METHOD=spawn + pytest -s -v tests/e2e/offline_inference/test_soulxsinger.py -k 'svc or precomputed' -m 'advanced_model and cuda' --run-level 'advanced_model' + " + agents: + queue: "gpu_1_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + - label: "Diffusion · Qwen Image Edit Test" source_file_dependencies: - tests/e2e/online_serving/test_qwen_image_edit.py diff --git a/.gitignore b/.gitignore index 40a6916b5..a175e19e3 100644 --- a/.gitignore +++ b/.gitignore @@ -241,6 +241,7 @@ datasets/ *.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 79c44d740..75b3ea1ad 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -63,6 +63,8 @@ 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` | ✅︎ | ✅︎ | | | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index 526d2fd45..f827e79b3 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -197,6 +197,7 @@ Factory function `get_offload_backend()` selects the appropriate backend based o | StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` | | Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` | | Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` | +| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | `"layers"` | | BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | `"layers"`, `"customized modules"` | **Notes:** diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index ea1d36051..180ab9954 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -521,3 +521,100 @@ 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 new file mode 100644 index 000000000..97ed7eb97 --- /dev/null +++ b/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py @@ -0,0 +1,162 @@ +"""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 = getattr(omni_out, "request_output", None) + 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 new file mode 100644 index 000000000..116b8da4b --- /dev/null +++ b/examples/offline_inference/text_to_speech/soulxsinger/end2end.py @@ -0,0 +1,285 @@ +"""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 = getattr(out, "request_output", None) + 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 1a3096fab..e8f862654 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -24,6 +24,7 @@ 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. @@ -661,3 +662,65 @@ 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 new file mode 100644 index 000000000..91035e625 --- /dev/null +++ b/examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py @@ -0,0 +1,186 @@ +#!/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 new file mode 100755 index 000000000..3e20de39b --- /dev/null +++ b/examples/online_serving/text_to_speech/soulxsinger/run_server.sh @@ -0,0 +1,39 @@ +#!/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 196b544e6..962f58a8f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -119,6 +119,12 @@ 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" diff --git a/tests/assets/soulxsinger/music.json b/tests/assets/soulxsinger/music.json new file mode 100644 index 000000000..d9e0d8882 --- /dev/null +++ b/tests/assets/soulxsinger/music.json @@ -0,0 +1,16 @@ +[ + { + "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 new file mode 100644 index 000000000..f3126e79e Binary files /dev/null and b/tests/assets/soulxsinger/music.mp3 differ diff --git a/tests/assets/soulxsinger/zh_prompt.json b/tests/assets/soulxsinger/zh_prompt.json new file mode 100644 index 000000000..86ff3614e --- /dev/null +++ b/tests/assets/soulxsinger/zh_prompt.json @@ -0,0 +1,16 @@ +[ + { + "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 new file mode 100644 index 000000000..bc306410a Binary files /dev/null and b/tests/assets/soulxsinger/zh_prompt.mp3 differ diff --git a/tests/e2e/offline_inference/test_soulxsinger.py b/tests/e2e/offline_inference/test_soulxsinger.py new file mode 100644 index 000000000..465f5547b --- /dev/null +++ b/tests/e2e/offline_inference/test_soulxsinger.py @@ -0,0 +1,241 @@ +"""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.advanced_model, 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, + stage_configs_path=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), + stage_configs_path=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/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index edabee9a9..45bd80c7a 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -706,6 +706,8 @@ def _build_engine_args( ``engine_extras`` can still carry a stage-specific ``dtype``). """ engine_args: dict[str, Any] = {"model_arch": ps.model_arch or pipeline.model_arch} + if ps.execution_type == StageExecutionType.DIFFUSION and ps.model_arch: + engine_args.setdefault("model_class_name", ps.model_arch) if ps.engine_output_type: engine_args["engine_output_type"] = ps.engine_output_type if next_stage_proc: diff --git a/vllm_omni/deploy/soulxsinger_svc.yaml b/vllm_omni/deploy/soulxsinger_svc.yaml new file mode 100644 index 000000000..a9d815806 --- /dev/null +++ b/vllm_omni/deploy/soulxsinger_svc.yaml @@ -0,0 +1,19 @@ +# 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 new file mode 100644 index 000000000..3f021cb53 --- /dev/null +++ b/vllm_omni/deploy/soulxsinger_svs.yaml @@ -0,0 +1,19 @@ +# 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 new file mode 100644 index 000000000..30235d2c1 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/__init__.py @@ -0,0 +1,44 @@ +"""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 new file mode 100644 index 000000000..282ab12ab --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/__init__.py @@ -0,0 +1,17 @@ +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 new file mode 100644 index 000000000..465898e8a --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/convnext.py @@ -0,0 +1,46 @@ +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 new file mode 100644 index 000000000..6b1fda616 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/decoder.py @@ -0,0 +1,24 @@ +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 new file mode 100644 index 000000000..19e87b65b --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/flow_matching.py @@ -0,0 +1,423 @@ +# 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 new file mode 100644 index 000000000..23a3377f2 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/llama.py @@ -0,0 +1,393 @@ +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.cache_dit_backend 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 new file mode 100644 index 000000000..cc767dcf3 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/mel_transform.py @@ -0,0 +1,151 @@ +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 new file mode 100644 index 000000000..70933dd92 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/note_transcription.py @@ -0,0 +1,958 @@ +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 new file mode 100644 index 000000000..9e020f61b --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/__init__.py @@ -0,0 +1,7 @@ +"""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 new file mode 100644 index 000000000..0752c73f8 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/asr.py @@ -0,0 +1,326 @@ +"""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 new file mode 100644 index 000000000..825c1622f --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/mel_grid_f0.py @@ -0,0 +1,96 @@ +"""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 new file mode 100644 index 000000000..aed72a995 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py @@ -0,0 +1,431 @@ +"""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, + ) + + # Non-None placeholder so DiffusionEngine forwards custom_output to stage-1 + # (output=None triggers an early return that drops handoff). + return DiffusionOutput( + output=torch.zeros(0), + custom_output={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 new file mode 100644 index 000000000..50ae4b2fc --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rmvpe.py @@ -0,0 +1,593 @@ +"""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 new file mode 100644 index 000000000..0333b521e --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rosvot.py @@ -0,0 +1,619 @@ +"""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 new file mode 100644 index 000000000..0140182ab --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/segmenter.py @@ -0,0 +1,267 @@ +"""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 new file mode 100644 index 000000000..e29b52a29 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/stack.py @@ -0,0 +1,145 @@ +"""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 new file mode 100644 index 000000000..f135fd553 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/utils.py @@ -0,0 +1,260 @@ +"""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 new file mode 100644 index 000000000..b62902d38 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/vocal_sep.py @@ -0,0 +1,207 @@ +"""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 new file mode 100644 index 000000000..ca3ef852e --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/vocoder.py @@ -0,0 +1,983 @@ +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 new file mode 100644 index 000000000..379dda3a5 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/whisper_encoder.py @@ -0,0 +1,96 @@ +"""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 new file mode 100644 index 000000000..64fa213ab --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py @@ -0,0 +1,359 @@ +import os +import time +from collections.abc import Callable +from contextlib import contextmanager +from typing import Any, ClassVar + +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", +] + + +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, + custom_output_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 or 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=audio, + custom_output={custom_output_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 new file mode 100644 index 000000000..91d4437c2 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py @@ -0,0 +1,409 @@ +"""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, +) +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 + +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.prompts[0] + 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 audio.detach().cpu().float().numpy() + + 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: OmniDiffusionRequest) -> DiffusionOutput: + return self._forward_batch_from_request( + req, + kind="svc", + custom_output_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 new file mode 100644 index 000000000..ecf6f94a5 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py @@ -0,0 +1,353 @@ +"""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, +) +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 + +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.prompts[0] + 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 audio.detach().cpu().float().numpy() + + 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: OmniDiffusionRequest) -> DiffusionOutput: + return self._forward_batch_from_request( + req, + kind="svs", + custom_output_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 new file mode 100644 index 000000000..aed87f757 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/preprocess/__init__.py @@ -0,0 +1 @@ +"""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 new file mode 100644 index 000000000..5f54d47c9 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/preprocess/g2p.py @@ -0,0 +1,86 @@ +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 new file mode 100644 index 000000000..8a8611121 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/preprocess/metadata_utils.py @@ -0,0 +1,137 @@ +"""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 new file mode 100644 index 000000000..83f81e481 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py @@ -0,0 +1,106 @@ +"""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 DiffusionOutput.custom_output. +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 new file mode 100644 index 000000000..750175fc5 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/preprocess/pre_process.py @@ -0,0 +1,220 @@ +"""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 {}), + ) + for i, prompt in enumerate(request.prompts): + 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.prompts[i] = prompt + continue + 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.prompts[i] = 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 new file mode 100644 index 000000000..4f201c407 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/utils.py @@ -0,0 +1,565 @@ +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 5ab847024..93d7f5d75 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -276,6 +276,16 @@ _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", @@ -489,6 +499,8 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "LTX23Pipeline": "get_ltx2_post_process_func", "LTX23ImageToVideoPipeline": "get_ltx2_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", @@ -559,6 +571,8 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "HunyuanImage3ForCausalMM": "get_hunyuan_image_3_pre_process_func", "MagiHumanPipeline": "get_magi_human_pre_process_func", "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func", + "SoulXSingerPipeline": "get_soulxsinger_pre_process_func", + "SoulXSingerSVCPipeline": "get_soulxsinger_svc_pre_process_func", } diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 7960bae61..c3cf11b72 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -530,23 +530,24 @@ def get_stage_tp_size(stage_cfg: Any) -> int: def get_stage_devices_per_replica(stage_cfg: Any) -> int: """Return the number of devices consumed by one replica of *stage_cfg*.""" - if getattr(stage_cfg, "stage_type", "llm") != "diffusion": - return get_stage_tp_size(stage_cfg) + engine_args = getattr(stage_cfg, "engine_args", {}) + if getattr(stage_cfg, "stage_type", "llm") == "diffusion": + parallel_config = _get_attr_or_item(engine_args, "parallel_config") + if parallel_config is None: + return 1 - parallel_config = _get_attr_or_item(getattr(stage_cfg, "engine_args", {}), "parallel_config") - if parallel_config is None: - return 1 + world_size = _get_attr_or_item(parallel_config, "world_size") + if world_size is not None: + return max(1, int(world_size)) - world_size = _get_attr_or_item(parallel_config, "world_size") - if world_size is not None: - return max(1, int(world_size)) + try: + from vllm_omni.diffusion.data import DiffusionParallelConfig - try: - from vllm_omni.diffusion.data import DiffusionParallelConfig + return max(1, int(DiffusionParallelConfig.from_dict(_to_dict(parallel_config)).world_size)) + except Exception: + return 1 - return max(1, int(DiffusionParallelConfig.from_dict(_to_dict(parallel_config)).world_size)) - except Exception: - return 1 + return get_stage_tp_size(stage_cfg) def compute_replica_layout(