anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[CI]Remove SoulX-Singer support (#6362)

Browse Source 
Signed-off-by: wangyu <410167048@qq.com>
This commit is contained in:
wangyu authored and GitHub committed 2026-08-23 07:16:32 +08:00
1 parent b2c9376539
commit 2e096788a6
51 files changed
+1 -9575
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.gitignore
-1
View File
@@ -245,7 +245,6 @@ datasets/
!recipes/cosmos3/negative_prompt.json
!tests/dfx/perf/tests/*.json
!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
!tests/assets/soulxsinger/*.json
*.jsonl
*.parquet
docs/models/supported_models.md
-2
View File
@@ -76,8 +76,6 @@ th {
| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | — |
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — |
| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — |
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — |
| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
docs/user_guide/diffusion/cpu_offload.md
-1
View File
@@ -49,7 +49,6 @@ higher-priority strategy is selected.
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | — | yes | — | `transformer_blocks` |
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` | yes | yes | — | `layers` |
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | — | yes | — | `layers`, customized modules |
| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | yes | yes | yes | `layers`, `gen_layers` |
examples/offline_inference/text_to_speech/README.md
+1 -99
View File
@@ -420,8 +420,7 @@ against transformers 5.x, so nothing is lost by holding it there.
(The same clash is why this is documented here rather than declared in
`pyproject.toml`: an extra resolves together with the base dependencies, so
`vllm-omni[gepard]` would be unsatisfiable rather than opt-in. SoulX-Singer English
SVS documents `nemo_toolkit[asr]` the same way.)
`vllm-omni[gepard]` would be unsatisfiable rather than opt-in.)
On a host whose CUDA toolkit cannot build kernels for the GPU — no `nvcc`/`ninja`, or a
consumer Blackwell (`sm_120`) card — also set `VLLM_USE_FLASHINFER_SAMPLER=0`. vLLM's
@@ -617,100 +616,3 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
- `--num-prompts N` replicates the prompt for performance measurement.
- `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
- Run `--help` for the full argument surface.
---
## SoulX-Singer
Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Script: `soulxsinger/end2end.py`. Deploy: `vllm_omni/deploy/soulxsinger_svs.yaml` or `soulxsinger_svc.yaml`.
### Prerequisites
Download DiT and preprocess weights, then set up separate SVS / SVC view directories. Copy `soulxsinger/utils/phoneme/phone_set.json` from upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) into the model weights dir as `phoneme/phone_set.json` — HuggingFace does not ship it.
```bash
# 1. DiT weights
export BASE=path/to/SoulX-Singer
export PREPROCESS=path/to/SoulX-Singer-Preprocess
export SVC_DIR=path/to/SoulX-Singer-svc
huggingface-cli download Soul-AILab/SoulX-Singer --local-dir "$BASE"
# 2. Preprocess weights (required)
huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess --local-dir "$PREPROCESS"
export SOULX_PREPROCESS_WEIGHTS_DIR="$PREPROCESS"
# 3. SVS / SVC view directories
mkdir -p "$SVC_DIR"
cp $BASE/{config.yaml,README.md,assets} $SVC_DIR
mv $BASE/model-svc.pt $SVC_DIR/model-svc.pt
cat > "$BASE/config.json" <<'EOF'
{
"model_type": "soulxsinger",
"architectures": ["SoulXSingerPipeline"],
"max_num_seqs": 1
}
EOF
cat > "$SVC_DIR/config.json" <<'EOF'
{
"model_type": "soulxsinger",
"architectures": ["SoulXSingerSVCPipeline"],
"max_num_seqs": 1
}
EOF
```
`config.yaml` hyper-parameters live under `$BASE`; each view's `config.json` `architectures` field is the single source of truth for SVS vs SVC. Point `--model` at the matching directory (`$BASE` for SVS, `$SVC_DIR` for SVC). Deploy YAML is chosen automatically from `config.json`; optional `--svs` / `--svc` only assert the mode matches.
**Online preprocess** is the default: pass `--prompt-audio` and `--target-audio`, and the worker runs vocal separation, F0, and (for SVS) lyrics/MIDI before DiT. Install only what your run needs:
```bash
pip install "BS-RoFormer" # vocal sep + F0 on GPU — SVS and SVC
```
Mandarin SVS also needs FunASR and Chinese G2P; `ffmpeg` must be on `PATH`:
```bash
# install optional dependencies:
pip install -e ".[soulx-svs]"
```
English SVS adds NeMo ASR and NLTK data; pass `--language English`:
```bash
pip install "nemo_toolkit[asr]==2.6.1" lhotse==1.32.2
python -c "import nltk; nltk.download('cmudict'); nltk.download('averaged_perceptron_tagger_eng')"
```
**Precomputed metadata** is the alternative: pass both `--prompt-metadata-path` and `--target-metadata-path` and skip online ASR/ROSVOT — none of the packages above are required. JSON can be produced by integrated preprocess on a prior run, or by upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) `preprocess/` scripts if you prefer to run that outside vLLM-Omni.
### Quick start
```bash
# SVS — default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3
python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
--model "$BASE" \
--preprocess-weights-dir "$PREPROCESS" \
--control score \
--num-inference-steps 32 \
-o output.wav
python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
--model "$SVC_DIR" \
--preprocess-weights-dir "$PREPROCESS" \
--svc \
--num-inference-steps 32 \
-o output_svc.wav
```
`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. Long SVS targets are handled in one request. See `end2end.py --help` for `--pitch-shift`, `--vocal-sep`, `--auto-shift`, and language/control options.
### Notes
- Output: 24 kHz mono WAV; batch only.
- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on.
- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only.
---
examples/offline_inference/text_to_speech/soulxsinger/benchmark.py
-162
View File
@@ -1,162 +0,0 @@
"""Repeatable SoulX-Singer offline benchmark (preprocess + DiT, RTF + stage timings).
Usage:
python benchmark.py --model /path/to/SoulX-Singer --svs \\
--prompt-audio ../SoulX-Singer/example/audio/zh_prompt.mp3 \\
--target-audio ../SoulX-Singer/example/audio/music.mp3 \\
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
-o benchmark.wav
python benchmark.py ... --enable-diffusion-pipeline-profiler
"""
import argparse
import statistics
import time
from pathlib import Path
import soundfile as sf
from end2end import (
SVC_DEPLOY_CONFIG,
SVS_DEPLOY_CONFIG,
add_inference_args,
build_sampling,
extract_audio,
resolve_preprocess_weights_dir,
)
from vllm_omni.entrypoints.omni import Omni
def _audio_duration_sec(outputs) -> float:
audio_np, sr = extract_audio(outputs)
return audio_np.size / sr
def _stage_durations_from_output(outputs) -> dict[str, float]:
omni_out = outputs[0]
durations = getattr(omni_out, "stage_durations", None) or {}
if durations:
return dict(durations)
ro = omni_out
if ro is not None:
return dict(getattr(ro, "stage_durations", None) or {})
return {}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="SoulX-Singer offline benchmark")
add_inference_args(parser)
parser.add_argument("--warmup", type=int, default=1, help="Warmup runs (excluded from stats)")
parser.add_argument("--runs", type=int, default=3, help="Measured runs")
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable low-overhead stage timing in DiffusionOutput.stage_durations",
)
parser.add_argument(
"--enforce-eager",
action="store_true",
help="Disable torch.compile on diff_estimator (eager baseline for A/B)",
)
parser.add_argument(
"--dtype",
type=str,
default=None,
choices=["bfloat16", "bf16", "float16", "fp16", "half"],
help="DiT trunk dtype (default from deploy YAML). Use float16 to match upstream --fp16.",
)
parser.add_argument(
"--output",
"-o",
type=str,
default=None,
help="Save WAV from the last measured run (after warmup)",
)
parser.add_argument(
"--save-each-run",
action="store_true",
help="With --output, also write measured runs as <stem>_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()
examples/offline_inference/text_to_speech/soulxsinger/end2end.py
-285
View File
@@ -1,285 +0,0 @@
"""Offline SoulX-Singer SVS / SVC: single-stage DiT (preprocess inline).
Pass raw prompt/target audio; vLLM-Omni runs integrated preprocess then inference.
Usage:
# SVS (default)
python end2end.py --model /path/to/SoulX-Singer \\
--prompt-audio path/to/prompt/audio.mp3 \\
--target-audio path/to/target/audio.mp3 \\
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
-o output.wav
# SVC (config.json must declare SoulXSingerSVCPipeline; --model points at SVC view)
python end2end.py --model /path/to/SoulX-Singer-svc \\
--prompt-audio path/to/prompt/audio.mp3 \\
--target-audio path/to/target/audio.mp3 \\
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
-o output_svc.wav
"""
import argparse
import os
from pathlib import Path
import numpy as np
import soundfile as sf
from vllm_omni.diffusion.models.soulx_singer.utils import resolve_soulx_kind, validate_soulx_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
REPO_ROOT = Path(__file__).resolve().parents[4]
SVS_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svs.yaml"
SVC_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svc.yaml"
DEFAULT_PROMPT_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "zh_prompt.mp3"
DEFAULT_TARGET_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "music.mp3"
_SAMPLE_RATE = 24000
def _require_paths(paths: dict[str, str | None]) -> None:
missing = [name for name, p in paths.items() if not p or not os.path.isfile(p)]
if missing:
raise FileNotFoundError("Missing file(s): " + ", ".join(f"{k}={paths[k]!r}" for k in missing))
def resolve_preprocess_weights_dir(explicit: str | None) -> str:
if explicit:
path = Path(explicit).expanduser().resolve()
if not (path / "rmvpe" / "rmvpe.pt").is_file():
raise FileNotFoundError(f"--preprocess-weights-dir must contain rmvpe/rmvpe.pt: {path}")
return str(path)
env = os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR")
if env:
path = Path(env).expanduser().resolve()
if (path / "rmvpe" / "rmvpe.pt").is_file():
return str(path)
raise FileNotFoundError(
"SoulX preprocess weights not found. Pass --preprocess-weights-dir or set "
"SOULX_PREPROCESS_WEIGHTS_DIR to a directory containing rmvpe/rmvpe.pt "
"(huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess)."
)
def build_sampling(
args: argparse.Namespace,
*,
preprocess_weights_dir: str,
kind: str,
) -> OmniDiffusionSamplingParams:
extra: dict = {
"prompt_audio": os.path.abspath(args.prompt_audio),
"target_audio": os.path.abspath(args.target_audio),
"preprocess_weights_dir": preprocess_weights_dir,
"vocal_sep": args.vocal_sep,
"auto_shift": args.auto_shift,
"pitch_shift": args.pitch_shift,
}
if kind == "svs":
extra["language"] = args.language
extra["control"] = args.control
if args.prompt_metadata_path:
extra["prompt_metadata_path"] = os.path.abspath(args.prompt_metadata_path)
if args.target_metadata_path:
extra["target_metadata_path"] = os.path.abspath(args.target_metadata_path)
if args.prompt_metadata_path or args.target_metadata_path:
extra["audio_path"] = os.path.abspath(args.prompt_wav_path or args.prompt_audio)
extra = validate_soulx_extra_args(kind, extra)
return OmniDiffusionSamplingParams(
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
extra_args=extra,
)
def extract_audio(outputs) -> tuple[np.ndarray, int]:
out = outputs[0]
mm = getattr(out, "multimodal_output", None)
if mm is None:
ro = out
if ro is not None:
mm = getattr(ro, "multimodal_output", None)
if mm is None and ro.outputs:
mm = getattr(ro.outputs[0], "multimodal_output", None)
if not mm or "audio" not in mm:
raise RuntimeError("No audio in multimodal_output")
audio_val = mm["audio"]
if isinstance(audio_val, list):
chunks = [
np.asarray(chunk.detach().cpu().float().numpy() if hasattr(chunk, "detach") else chunk).reshape(-1)
for chunk in audio_val
if chunk is not None
]
audio_np = np.concatenate(chunks, axis=0) if chunks else np.array([], dtype=np.float32)
elif hasattr(audio_val, "cpu"):
audio_np = audio_val.detach().cpu().numpy().squeeze()
else:
audio_np = np.asarray(audio_val, dtype=np.float32).squeeze()
sr_val = mm.get("sr", mm.get("audio_sample_rate", _SAMPLE_RATE))
if isinstance(sr_val, list) and sr_val:
sr_val = sr_val[-1]
if hasattr(sr_val, "item"):
sr_val = sr_val.item()
return audio_np.astype(np.float32, copy=False), int(sr_val)
def add_inference_args(parser: argparse.ArgumentParser) -> None:
mode = parser.add_mutually_exclusive_group()
mode.add_argument(
"--svc",
action="store_true",
help="Assert SVC mode (must match config.json architectures)",
)
parser.add_argument(
"--model",
type=str,
required=True,
help="SoulX-Singer weight directory ($BASE for SVS, $SVC_DIR for SVC)",
)
parser.add_argument(
"--deploy-config",
type=str,
default=None,
help="Override deploy YAML (default: soulxsinger_svs.yaml or soulxsinger_svc.yaml)",
)
parser.add_argument(
"--prompt-audio",
type=str,
default=str(DEFAULT_PROMPT_AUDIO),
help="Prompt / reference audio (mp3 or wav)",
)
parser.add_argument(
"--target-audio",
type=str,
default=str(DEFAULT_TARGET_AUDIO),
help="Target content audio (mp3 or wav)",
)
parser.add_argument(
"--prompt-metadata-path",
type=str,
default=None,
help="SVS precomputed prompt metadata.json (skips online ASR/ROSVOT for prompt)",
)
parser.add_argument(
"--target-metadata-path",
type=str,
default=None,
help="SVS precomputed target metadata.json (skips online ASR/ROSVOT for target)",
)
parser.add_argument(
"--prompt-wav-path",
type=str,
default=None,
help="SVS prompt vocal wav for metadata processor when using --prompt-metadata-path (default: --prompt-audio)",
)
parser.add_argument(
"--preprocess-weights-dir",
type=str,
default=None,
help="SoulX-Singer-Preprocess weights (or set SOULX_PREPROCESS_WEIGHTS_DIR)",
)
parser.add_argument(
"--language",
type=str,
default="Mandarin",
help="SVS lyrics language (Mandarin, Cantonese, English)",
)
parser.add_argument(
"--control",
type=str,
default="melody",
choices=["score", "melody"],
help="SVS control mode",
)
parser.add_argument(
"--vocal-sep",
action=argparse.BooleanOptionalAction,
default=False,
help="Run vocal separation during preprocess",
)
parser.add_argument(
"--auto-shift",
action=argparse.BooleanOptionalAction,
default=True,
help="Auto pitch shift to align target pitch with prompt",
)
parser.add_argument("--pitch-shift", type=int, default=0)
parser.add_argument("--num-inference-steps", type=int, default=32)
parser.add_argument("--guidance-scale", type=float, default=3.0)
parser.add_argument(
"--seed",
type=int,
default=42,
help="Optional CFM RNG seed. Omit for non-deterministic sampling.",
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="SoulX-Singer offline SVS / SVC (preprocess + DiT)")
add_inference_args(parser)
parser.add_argument("--output", "-o", type=str, default="output.wav")
args = parser.parse_args()
return args
def main() -> None:
args = parse_args()
kind = resolve_soulx_kind(args.model)
if not (args.svc ^ (kind == "svs")):
raise ValueError(f"mode mismatch: {args.model}/config.json (declares {kind!r}), --svc={args.svc}")
required = {
"prompt_audio": args.prompt_audio,
"target_audio": args.target_audio,
}
if kind == "svs" and bool(args.prompt_metadata_path) ^ bool(args.target_metadata_path):
raise ValueError(
"SVS precomputed metadata requires both --prompt-metadata-path and "
"--target-metadata-path (or neither, for online preprocess)."
)
if args.prompt_metadata_path:
required["prompt_metadata_path"] = args.prompt_metadata_path
required["target_metadata_path"] = args.target_metadata_path
if args.prompt_wav_path:
required["prompt_wav_path"] = args.prompt_wav_path
_require_paths(required)
preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir)
deploy_config = args.deploy_config or str(SVC_DEPLOY_CONFIG if kind == "svc" else SVS_DEPLOY_CONFIG)
mode = kind.upper()
print(f"Loading SoulX-Singer {mode} from {args.model}")
print(f" deploy: {deploy_config}")
print(f" prompt: {args.prompt_audio}")
print(f" target: {args.target_audio}")
omni = Omni(model=args.model, deploy_config=deploy_config, async_chunk=False) # SoulX-Singer is batch-only
sampling = build_sampling(
args,
preprocess_weights_dir=preprocess_dir,
kind=kind,
)
prompt = {"prompt_token_ids": [0], "multi_modal_data": {}}
print("Running generate (inline preprocess + DiT)...")
outputs = list(omni.generate([prompt], sampling))
if not outputs:
raise RuntimeError("No output from omni.generate")
err = getattr(outputs[0], "error", None)
if err is not None:
raise RuntimeError(f"Pipeline error: {err}")
audio_np, sr = extract_audio(outputs)
sf.write(args.output, audio_np, sr)
print(f"Saved {args.output} ({sr} Hz, {len(audio_np) / sr:.2f}s)")
omni.close()
if __name__ == "__main__":
main()
examples/online_serving/text_to_speech/README.md
-63
View File
@@ -25,7 +25,6 @@ For the full list of supported architectures across all modalities, see
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
| VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
| SoulX-Singer | `Soul-AILab/SoulX-Singer` | ✓ (prompt audio) | — (batch only) | — (prompt + target audio) | — (chat client) |
CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](../../offline_inference/text_to_speech/README.md#cosyvoice3) instead.
@@ -715,65 +714,3 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
- Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
- Voice cloning is gated upstream and may require a recent `mistral_common`.
- A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now.
---
## SoulX-Singer
Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Single-stage DiT with inline preprocess. Uses the `/v1/chat/completions` endpoint with multimodal input (`prompt_audio` + `target_audio`).
### Prerequisites
Download DiT and preprocess weights, then set up separate SVS / SVC view directories and install dependencies as described in the [offline README](../../offline_inference/text_to_speech/README.md#soulx-singer). `config.json` `architectures` field is the single source of truth for SVS vs SVC — point `MODEL` at the matching directory.
### Launch
```bash
# SVS (default)
export MODEL=/path/to/SoulX-Singer
export PREPROCESS=/path/to/SoulX-Singer-Preprocess
bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh
# SVC
export MODE=svc
export MODEL=/path/to/SoulX-Singer-svc
bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh
```
Or equivalently, set `SOULX_PREPROCESS_WEIGHTS_DIR` and launch directly:
```bash
export SOULX_PREPROCESS_WEIGHTS_DIR=$PREPROCESS
vllm serve $MODEL --omni \
--deploy-config vllm_omni/deploy/soulxsinger_${MODE}.yaml \
--port 8192 --trust-remote-code --enforce-eager
```
### Sending requests
Audio paths must be reachable from the server host (local filesystem or data URL). The client sends prompt vocal via `input_audio` and target accompaniment via `extra_args['target_audio']`.
```bash
# Default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3
python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \
--prompt-audio /path/on/server/zh_prompt.mp3 \
--target-audio /path/on/server/music.mp3 \
--preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \
-o output.wav
```
Use precomputed metadata to skip online preprocess with following command:
```bash
python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \
--prompt-metadata-path /path/on/server/zh_prompt.json \
--target-metadata-path /path/on/server/music.json \
--audio-path /path/on/server/zh_prompt.mp3 \
-o output.wav
```
`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. See `openai_chat_client.py --help` for `--vocal-sep`, `--language`, `--num-inference-steps`, `--guidance-scale`, and `--seed`.
### Notes
- Output: 24 kHz mono WAV; batch only.
- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on.
- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only.
examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py
-186
View File
@@ -1,186 +0,0 @@
#!/usr/bin/env python3
"""SoulX-Singer OpenAI-compatible chat client (SVS / SVC).
Sends prompt audio via ``input_audio`` and target accompaniment via
``extra_args['target_audio']`` (server-local path). For integrated preprocess,
also pass ``preprocess_weights_dir`` in ``extra_args``.
Usage:
python openai_chat_client.py \\
--prompt-audio /path/on/server/zh_prompt.mp3 \\
--target-audio /path/on/server/music.mp3 \\
--preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \\
-o output.wav
"""
from __future__ import annotations
import argparse
import base64
import io
import sys
from pathlib import Path
import requests
import soundfile
import torch
def _audio_to_data_url(path: Path) -> str:
with path.open("rb") as handle:
data = base64.b64encode(handle.read()).decode("ascii")
return f"data:audio/mpeg;base64,{data}"
def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None:
audio = audio.to(torch.float32)
peak = audio.abs().max().clamp(min=1e-8)
audio = audio / peak
path.parent.mkdir(parents=True, exist_ok=True)
soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16")
def _decode_audio_from_response(body: dict) -> tuple[torch.Tensor, int]:
for choice in body.get("choices", []):
audio_obj = choice.get("message", {}).get("audio")
if isinstance(audio_obj, dict) and audio_obj.get("data"):
data, sr = soundfile.read(
io.BytesIO(base64.b64decode(audio_obj["data"])),
dtype="float32",
always_2d=True,
)
return torch.from_numpy(data).transpose(0, 1), sr
brief = {k: v for k, v in body.items() if k != "choices"}
raise RuntimeError(f"no audio in response message.audio: {brief}")
def main() -> int:
repo_root = Path(__file__).resolve().parents[4]
default_assets = repo_root / "tests" / "assets" / "soulxsinger"
parser = argparse.ArgumentParser(description="SoulX-Singer online chat client")
parser.add_argument("--port", type=int, default=8192)
parser.add_argument("--model", default="Soul-AILab/SoulX-Singer")
parser.add_argument(
"--prompt-audio",
default=str(default_assets / "zh_prompt.mp3"),
help="Prompt vocal audio (path on server if using extra_args, or local for input_audio)",
)
parser.add_argument(
"--target-audio",
default=str(default_assets / "music.mp3"),
help="Target accompaniment path on the server (extra_args['target_audio'])",
)
parser.add_argument(
"--prompt-metadata-path",
default=None,
help="SVS precomputed prompt metadata.json",
)
parser.add_argument(
"--target-metadata-path",
default=None,
help="SVS precomputed target metadata.json",
)
parser.add_argument(
"--audio-path",
default=None,
help="SVS prompt vocal wav for precomputed metadata",
)
parser.add_argument("--preprocess-weights-dir", default=None)
parser.add_argument("--output", "-o", default="soulxsinger_out.wav")
parser.add_argument("--svc", action="store_true", help="Use SVC mode knobs")
parser.add_argument("--language", default="Mandarin")
parser.add_argument("--num-inference-steps", type=int, default=32)
parser.add_argument("--guidance-scale", type=float, default=3.0)
parser.add_argument(
"--seed",
type=int,
default=42,
help="Optional CFM seed. Omit for non-deterministic sampling.",
)
parser.add_argument(
"--auto-shift",
action=argparse.BooleanOptionalAction,
default=True,
help="Auto pitch shift (default: on, original upstream infer.sh)",
)
parser.add_argument(
"--control",
default="melody",
choices=["melody", "score"],
help="SVS control mode",
)
parser.add_argument("--vocal-sep", action="store_true")
args = parser.parse_args()
meta_paths = (args.prompt_metadata_path, args.target_metadata_path, args.audio_path)
if any(meta_paths) and not all(meta_paths):
print(
"ERROR: precomputed metadata requires --prompt-metadata-path, "
"--target-metadata-path, and --audio-path together.",
file=sys.stderr,
)
return 2
extra_args: dict = {
"vocal_sep": args.vocal_sep,
"auto_shift": args.auto_shift,
"pitch_shift": 0,
}
if all(meta_paths):
extra_args.update(
{
"prompt_metadata_path": str(Path(args.prompt_metadata_path).expanduser().resolve()),
"target_metadata_path": str(Path(args.target_metadata_path).expanduser().resolve()),
"audio_path": str(Path(args.audio_path).expanduser().resolve()),
}
)
content = [{"type": "text", "text": "soulx-singer"}]
else:
prompt_path = Path(args.prompt_audio).expanduser().resolve()
if not prompt_path.is_file():
print(f"ERROR: prompt audio not found: {prompt_path}", file=sys.stderr)
return 2
extra_args["prompt_audio"] = str(prompt_path)
extra_args["target_audio"] = str(Path(args.target_audio).expanduser().resolve())
if args.preprocess_weights_dir:
extra_args["preprocess_weights_dir"] = str(Path(args.preprocess_weights_dir).expanduser().resolve())
content = [
{"type": "text", "text": "soulx-singer"},
{
"type": "input_audio",
"input_audio": {"data": _audio_to_data_url(prompt_path), "format": "mp3"},
},
]
if not args.svc:
extra_args["language"] = args.language
extra_args["control"] = args.control
payload = {
"model": args.model,
"modalities": ["audio"],
"messages": [{"role": "user", "content": content}],
"num_inference_steps": args.num_inference_steps,
"guidance_scale": args.guidance_scale,
"extra_args": extra_args,
}
if args.seed is not None:
payload["seed"] = args.seed
print(f"POST http://localhost:{args.port}/v1/chat/completions")
response = requests.post(
f"http://localhost:{args.port}/v1/chat/completions",
headers={"Content-Type": "application/json"},
json=payload,
timeout=1800,
)
response.raise_for_status()
audio, sample_rate = _decode_audio_from_response(response.json())
_save_wav(audio, Path(args.output), sample_rate)
duration = audio.shape[-1] / sample_rate
print(f"saved {args.output} sr={sample_rate}Hz duration={duration:.2f}s")
return 0
if __name__ == "__main__":
raise SystemExit(main())
examples/online_serving/text_to_speech/soulxsinger/run_server.sh
-39
View File
@@ -1,39 +0,0 @@
#!/bin/bash
# Launch vLLM-Omni server for SoulX-Singer (single-stage DiT, preprocess inline).
#
# Usage:
# MODEL=/path/to/SoulX-Singer PREPROCESS=/path/to/Preprocess \
#
# Audio paths in client extra_args must be readable on the server host.
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)"
MODEL="${MODEL:-Soul-AILab/SoulX-Singer}"
MODE="${MODE:-svs}"
PORT="${PORT:-8192}"
GPUS="${GPUS:-0}"
if [[ "$MODE" == "svc" ]]; then
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svc.yaml}"
else
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svs.yaml}"
fi
echo "Starting SoulX-Singer server"
echo " MODEL=$MODEL"
echo " MODE=$MODE"
echo " PORT=$PORT"
echo " DEPLOY_CONFIG=$DEPLOY_CONFIG"
echo " CUDA_VISIBLE_DEVICES=$GPUS"
CUDA_VISIBLE_DEVICES="$GPUS" \
vllm serve "$MODEL" \
--omni \
--deploy-config "$DEPLOY_CONFIG" \
--host 0.0.0.0 \
--port "$PORT" \
--trust-remote-code \
--enforce-eager
pyproject.toml
-7
View File
@@ -135,13 +135,6 @@ docs = [
"pydantic",
]
soulx-svs = [
"funasr>=1.0.0",
"g2pM>=0.1.2.5",
"g2p-en>=2.1.0",
"ToJyutping>=3.2.0",
]
[project.urls]
Homepage = "https://github.com/vllm-project/vllm-omni"
Repository = "https://github.com/vllm-project/vllm-omni"
tests/assets/soulxsinger/music.json
-16
View File
@@ -1,16 +0,0 @@
[
{
"index": "vocal_240_51240",
"language": "Mandarin",
"time": [
240,
51240
],
"duration": "0.21 0.18 0.34 0.38 0.26 0.22 0.50 0.20 0.33 0.13 0.60 0.22 0.38 0.52 0.30 0.12 1.82 0.98 0.20 0.26 0.38 0.36 0.58 0.54 0.26 0.50 1.46 3.03 0.24 0.24 0.29 0.25 0.20 0.14 0.80 0.22 0.54 0.30 0.24 0.16 0.58 0.28 0.38 1.73 0.79 0.24 0.30 0.34 0.30 0.30 0.34 0.21 0.21 0.36 0.34 0.23 1.85 1.90 0.23 0.39 0.68 0.50 0.31 0.43 0.76 0.38 2.00 1.87 0.68 0.72 0.56 0.62 0.80 0.40 0.42 1.68 1.79 0.70 0.66 0.54 0.24 0.48 0.68 0.40 2.34 0.14",
"text": "<SP> 只 是 因 为 为 在 人 群 群 中 多 看 了 你 一 眼 <SP> 再 也 没 能 忘 掉 你 容 颜 <SP> 梦 想 着 着 偶 偶 然 然 有 <SP> 一 一 天 再 相 见 <SP> 从 此 我 开 始 始 孤 孤 单 思 念 念 <SP> 想 想 你 时 你 你 在 天 边 <SP> 想 你 时 你 在 眼 前 前 <SP> 想 你 时 你 你 在 脑 海 <SP>",
"phoneme": "<SP> 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 <SP> zh_zai4 zh_ye3 zh_mei2 zh_neng2 zh_wang4 zh_diao4 zh_ni3 zh_rong2 zh_yan2 <SP> zh_meng4 zh_xiang3 zh_zhe5 zh_zhe5 zh_ou3 zh_ou3 zh_ran2 zh_ran2 zh_you3 <SP> zh_yi1 zh_yi1 zh_tian1 zh_zai4 zh_xiang1 zh_jian4 <SP> zh_cong2 zh_ci3 zh_wo3 zh_kai1 zh_shi3 zh_shi3 zh_gu1 zh_gu1 zh_dan1 zh_si1 zh_nian4 zh_nian4 <SP> zh_xiang3 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_tian1 zh_bian1 <SP> zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_zai4 zh_yan3 zh_qian2 zh_qian2 <SP> zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_nao3 zh_hai3 <SP>",
"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.Line truncated
}
]
tests/assets/soulxsinger/music.mp3
BIN
View File
Binary file not shown.
tests/assets/soulxsinger/zh_prompt.json
-16
View File
@@ -1,16 +0,0 @@
[
{
"index": "vocal_320_10687",
"language": "Mandarin",
"time": [
320,
10687
],
"duration": "0.23 0.34 0.26 0.70 0.52 0.46 0.36 0.44 0.14 0.24 0.64 0.47 0.51 1.10 0.28 0.38 0.32 0.32 0.38 0.32 0.31 0.19 1.45",
"text": "<SP> 除 了 想 你 你 <SP> 除 了 了 爱 你 你 <SP> 我 什 么 什 么 都 愿 愿 意",
"phoneme": "<SP> zh_chu2 zh_le5 zh_xiang3 zh_ni3 zh_ni3 <SP> zh_chu2 zh_le5 zh_le5 zh_ai4 zh_ni3 zh_ni3 <SP> 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"
}
]
tests/assets/soulxsinger/zh_prompt.mp3
BIN
View File
Binary file not shown.
tests/e2e/offline_inference/test_soulxsinger_expansion.py
-241
View File
@@ -1,241 +0,0 @@
"""E2E offline inference tests for SoulX-Singer (single-stage, preprocess inline)."""
import functools
import importlib
import json
import os
import shutil
from pathlib import Path
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
import numpy as np
import pytest
from tests.helpers.mark import hardware_test
from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
PROMPT_AUDIO = get_asset_path("soulxsinger/zh_prompt.mp3")
TARGET_AUDIO = get_asset_path("soulxsinger/music.mp3")
PROMPT_METADATA = get_asset_path("soulxsinger/zh_prompt.json")
TARGET_METADATA = get_asset_path("soulxsinger/music.json")
SAMPLE_RATE = 24_000
# phone_set.json (phoneme vocab) is not shipped on HuggingFace; SVS loads it
# from the model dir. Stage it from this pinned upstream commit so the SVS DiT
# can run in CI; the SVS tests skip gracefully if it cannot be fetched.
_PHONE_SET_URL = (
"https://raw.githubusercontent.com/Soul-AILab/SoulX-Singer/"
"81aeb3ae772c70093c3de74dc23c92d983801ae4/soulxsinger/utils/phoneme/phone_set.json"
)
if not PROMPT_AUDIO.is_file() or not TARGET_AUDIO.is_file():
pytest.skip(
f"Missing SoulX-Singer audio assets: {PROMPT_AUDIO.name}, {TARGET_AUDIO.name}",
allow_module_level=True,
)
pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.tts]
_CASES = (
pytest.param(
"SoulXSingerPipeline",
"soulxsinger_svs.yaml",
{
"language": "Mandarin",
"vocal_sep": False,
"control": "score",
"auto_shift": False,
"pitch_shift": 0,
},
("g2pM", "g2p_en"),
id="svs",
),
pytest.param(
"SoulXSingerSVCPipeline",
"soulxsinger_svc.yaml",
{"vocal_sep": False, "auto_shift": False, "pitch_shift": 0},
(),
id="svc",
),
)
@functools.lru_cache(maxsize=1)
def _resolve_weights() -> tuple[Path, Path, Path]:
from huggingface_hub import snapshot_download
base = Path(snapshot_download("Soul-AILab/SoulX-Singer", allow_patterns=["*"]))
# phone_set.json is not on HF; best-effort stage it from pinned upstream so
# SVS can load. SVS tests skip (not fail) if both this and a manual copy are absent.
phone_set = base / "phoneme" / "phone_set.json"
if not phone_set.is_file() and not (base / "phone_set.json").is_file():
try:
from urllib.request import urlretrieve
phone_set.parent.mkdir(exist_ok=True)
urlretrieve(_PHONE_SET_URL, phone_set)
except Exception:
pass
# make a temporary svc directory
svc_dir = base.parent / "SoulX-Singer-SVC"
svc_dir.mkdir(parents=True, exist_ok=True)
shutil.copyfile(base / "config.yaml", svc_dir / "config.yaml")
shutil.move(base / "model-svc.pt", svc_dir / "model-svc.pt")
if raw := os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR"):
pre = Path(raw).expanduser().resolve()
if (pre / "rmvpe" / "rmvpe.pt").is_file():
return base, svc_dir, pre
from huggingface_hub import snapshot_download
pre = Path(snapshot_download("Soul-AILab/SoulX-Singer-Preprocess", allow_patterns=["*"]))
return base, svc_dir, pre
@pytest.fixture(scope="session")
def soulx_weights() -> tuple[Path, Path, Path]:
try:
return _resolve_weights()
except Exception as exc:
pytest.skip(f"Set SOULXSINGER_MODEL_DIR / SOULX_PREPROCESS_WEIGHTS_DIR. ({exc})")
def _flatten_audio(audio_val) -> np.ndarray:
import torch
if isinstance(audio_val, list):
chunks = [c.detach().cpu().float().numpy().reshape(-1) for c in audio_val if c is not None]
return np.concatenate(chunks) if chunks else np.array([], dtype=np.float32)
if isinstance(audio_val, torch.Tensor):
return audio_val.detach().cpu().float().numpy().reshape(-1)
return np.asarray(audio_val, dtype=np.float32).reshape(-1)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("architecture,deploy_yaml,extra_args,py_deps", _CASES)
def test_soulxsinger_multistage_from_audio(
soulx_weights: tuple[Path, Path, Path],
architecture: str,
deploy_yaml: str,
extra_args: dict,
py_deps: tuple[str, ...],
) -> None:
for mod in py_deps:
try:
importlib.import_module(mod)
except ImportError as exc:
pytest.fail(f"SoulX SVS requires {mod}: {exc}")
base_dir, svc_dir, preprocess_dir = soulx_weights
# SVS mode requires phone_set.json in the model directory
if architecture == "SoulXSingerPipeline":
json.dump(
{"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1},
open(base_dir / "config.json", "w"),
indent=4,
)
model = str(base_dir)
if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file():
pytest.skip(
"SoulX-Singer SVS test requires phoneme/phone_set.json. "
"Copy it from github.com/Soul-AILab/SoulX-Singer into the model dir. "
"See `examples/offline_inference/text_to_speech/README.md` for details."
)
if architecture == "SoulXSingerSVCPipeline":
json.dump(
{"model_type": "soulxsinger", "architectures": ["SoulXSingerSVCPipeline"], "max_num_seqs": 1},
open(svc_dir / "config.json", "w"),
indent=4,
)
model = str(svc_dir)
with OmniRunner(
model,
deploy_config=get_deploy_config_path(deploy_yaml),
async_chunk=False,
) as runner:
sampling = OmniDiffusionSamplingParams(
num_inference_steps=4,
guidance_scale=3.0,
seed=42,
extra_args={
"prompt_audio": str(PROMPT_AUDIO),
"target_audio": str(TARGET_AUDIO),
"preprocess_weights_dir": str(preprocess_dir),
**extra_args,
},
)
prompt = {"prompt_token_ids": [0]}
outputs = runner.generate([prompt], sampling)
assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output"
mm = outputs[0].multimodal_output
assert isinstance(mm, dict) and "audio" in mm
audio = _flatten_audio(mm["audio"])
assert 12_000 <= audio.size
assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4
duration_s = audio.size / SAMPLE_RATE
assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s"
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_soulxsinger_svs_precomputed(soulx_weights: tuple[Path, Path, Path]) -> None:
"""SVS from precomputed note/lyric metadata.
Exercises the SVS DiT + vocoder serving path while skipping the heavy
upstream preprocess (funasr ASR + ROSVOT note transcription), so it runs in
the merge gate without those non-pip dependencies.
"""
base_dir, _svc_dir, preprocess_dir = soulx_weights
if not PROMPT_METADATA.is_file() or not TARGET_METADATA.is_file():
pytest.skip(f"Missing SoulX-Singer metadata fixtures: {PROMPT_METADATA.name}, {TARGET_METADATA.name}")
if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file():
pytest.skip(
"SoulX-Singer SVS requires phoneme/phone_set.json (upstream fetch unavailable). "
"See `examples/offline_inference/text_to_speech/README.md`."
)
json.dump(
{"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1},
open(base_dir / "config.json", "w"),
indent=4,
)
with OmniRunner(
str(base_dir),
deploy_config=get_deploy_config_path("soulxsinger_svs.yaml"),
async_chunk=False,
) as runner:
sampling = OmniDiffusionSamplingParams(
num_inference_steps=4,
guidance_scale=3.0,
seed=42,
extra_args={
"language": "Mandarin",
"control": "score",
"prompt_metadata_path": str(PROMPT_METADATA),
"target_metadata_path": str(TARGET_METADATA),
"audio_path": str(PROMPT_AUDIO),
"preprocess_weights_dir": str(preprocess_dir),
},
)
prompt = {"prompt_token_ids": [0]}
outputs = runner.generate([prompt], sampling)
assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output"
mm = outputs[0].multimodal_output
assert isinstance(mm, dict) and "audio" in mm
audio = _flatten_audio(mm["audio"])
assert 12_000 <= audio.size
assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4
duration_s = audio.size / SAMPLE_RATE
assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s"
tests/model_tests/diffusion/test_alignment.py
-2
View File
@@ -78,8 +78,6 @@ EXCLUDED_MODELS = [
"StableDiffusionXLPipeline",
"Gr00tN1d7Pipeline",
"Pi0Pipeline",
"SoulXSingerPipeline",
"SoulXSingerSVCPipeline",
"SanaWmPipeline",
]
vllm_omni/config/pipeline_registry.py
-6
View File
@@ -103,10 +103,6 @@ from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import (
)
from vllm_omni.model_executor.models.qwen3_omni.pipeline import resolve_qwen3_omni_pipeline
from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE
from vllm_omni.model_executor.models.soulx_singer.pipeline import (
SOULXSINGER_SVC_PIPELINE,
SOULXSINGER_SVS_PIPELINE,
)
from vllm_omni.model_executor.models.step_audio2.pipeline import (
STEP_AUDIO2_ASR_PIPELINE,
STEP_AUDIO2_PIPELINE,
@@ -184,8 +180,6 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
"dynin_omni": DYNIN_OMNI_PIPELINE,
"indextts2": INDEXTTS2_PIPELINE,
"indextts2_5": INDEXTTS25_PIPELINE,
"soulxsinger_svc": SOULXSINGER_SVC_PIPELINE,
"soulxsinger_svs": SOULXSINGER_SVS_PIPELINE,
}
vllm_omni/deploy/soulxsinger_svc.yaml
-19
View File
@@ -1,19 +0,0 @@
# SoulX-Singer SVC: single-stage DiT (preprocess inline in pre_process_func).
#
# Use: vllm serve <SoulX-Singer-svc> --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
vllm_omni/deploy/soulxsinger_svs.yaml
-19
View File
@@ -1,19 +0,0 @@
# SoulX-Singer SVS: single-stage DiT (preprocess inline in pre_process_func).
#
# Use: vllm serve <SoulX-Singer> --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
vllm_omni/diffusion/models/soulx_singer/__init__.py
-44
View File
@@ -1,44 +0,0 @@
"""SoulX-Singer SVS / SVC flow-matching singing models for vLLM-Omni."""
__all__ = [
"FlowMatchingAudioPipeline",
"PipelineSoulXSingerSVS",
"PipelineSoulXSingerSVC",
"get_soulxsinger_post_process_func",
"get_soulxsinger_pre_process_func",
"get_soulxsinger_svc_pre_process_func",
]
def __getattr__(name: str):
if name in {"FlowMatchingAudioPipeline", "get_soulxsinger_post_process_func"}:
from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import (
FlowMatchingAudioPipeline,
get_soulxsinger_post_process_func,
)
return {
"FlowMatchingAudioPipeline": FlowMatchingAudioPipeline,
"get_soulxsinger_post_process_func": get_soulxsinger_post_process_func,
}[name]
if name in {"PipelineSoulXSingerSVC", "get_soulxsinger_svc_pre_process_func"}:
from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_svc import (
PipelineSoulXSingerSVC,
get_soulxsinger_svc_pre_process_func,
)
return {
"PipelineSoulXSingerSVC": PipelineSoulXSingerSVC,
"get_soulxsinger_svc_pre_process_func": get_soulxsinger_svc_pre_process_func,
}[name]
if name in {"PipelineSoulXSingerSVS", "get_soulxsinger_pre_process_func"}:
from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_svs import (
PipelineSoulXSingerSVS,
get_soulxsinger_pre_process_func,
)
return {
"PipelineSoulXSingerSVS": PipelineSoulXSingerSVS,
"get_soulxsinger_pre_process_func": get_soulxsinger_pre_process_func,
}[name]
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
vllm_omni/diffusion/models/soulx_singer/modules/__init__.py
-17
View File
@@ -1,17 +0,0 @@
from .convnext import ConvNeXtV2Block
from .decoder import CFMDecoder
from .flow_matching import FlowMatchingTransformer
from .llama import DiffLlama
from .mel_transform import MelSpectrogramEncoder
from .vocoder import Vocoder
from .whisper_encoder import WhisperEncoder
__all__ = [
"ConvNeXtV2Block",
"FlowMatchingTransformer",
"CFMDecoder",
"DiffLlama",
"MelSpectrogramEncoder",
"Vocoder",
"WhisperEncoder",
]
vllm_omni/diffusion/models/soulx_singer/modules/convnext.py
-46
View File
@@ -1,46 +0,0 @@
import torch
import torch.nn as nn
class GRN(nn.Module):
def __init__(self, dim):
super().__init__()
self.gamma = nn.Parameter(torch.zeros(1, 1, dim))
self.beta = nn.Parameter(torch.zeros(1, 1, dim))
def forward(self, x):
Gx = torch.norm(x, p=2, dim=1, keepdim=True)
Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
return self.gamma * (x * Nx) + self.beta + x
# ref: https://github.com/SWivid/F5-TTS/blob/main/src/f5_tts/model/modules.py#L247
class ConvNeXtV2Block(nn.Module):
def __init__(
self,
dim: int,
intermediate_dim: int,
dilation: int = 1,
):
super().__init__()
padding = (dilation * (7 - 1)) // 2
self.dwconv = nn.Conv1d(
dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation
) # depthwise conv
self.norm = nn.LayerNorm(dim, eps=1e-6)
self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers
self.act = nn.GELU()
self.grn = GRN(intermediate_dim)
self.pwconv2 = nn.Linear(intermediate_dim, dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
residual = x
x = x.transpose(1, 2) # b n d -> b d n
x = self.dwconv(x)
x = x.transpose(1, 2) # b d n -> b n d
x = self.norm(x)
x = self.pwconv1(x)
x = self.act(x)
x = self.grn(x)
x = self.pwconv2(x)
return residual + x
vllm_omni/diffusion/models/soulx_singer/modules/decoder.py
-24
View File
@@ -1,24 +0,0 @@
import torch
import torch.nn as nn
from vllm_omni.diffusion.models.soulx_singer.modules.flow_matching import FlowMatchingTransformer
class CFMDecoder(nn.Module):
def __init__(self, config):
super().__init__()
self.model = FlowMatchingTransformer(cfg=config, **config)
def forward(self, mel, x_mask, decoder_inp, is_prompt):
outputs = self.model(mel, x_mask, decoder_inp, is_prompt)
noise, x, flow_pred, final_mask, prompt_len = outputs["output"]
return noise, x, flow_pred, final_mask, prompt_len
def reverse_diffusion(self, pt_mel, pt_decoder_inp, gt_decoder_inp, n_timesteps=32, cfg=1):
diffusion_cond = torch.cat([pt_decoder_inp, gt_decoder_inp], dim=1)
diffusion_cond_emb = self.model.cond_emb(diffusion_cond)
diffusion_prompt = pt_mel
generated = self.model.reverse_diffusion(diffusion_cond_emb, diffusion_prompt, n_timesteps=n_timesteps, cfg=cfg)
return generated
vllm_omni/diffusion/models/soulx_singer/modules/flow_matching.py
-423
View File
@@ -1,423 +0,0 @@
# https://github.com/open-mmlab/Amphion/blob/main/models/svc/flow_matching_transformer/fmt_model.py
# Copyright (c) 2023 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import LlamaConfig
from vllm_omni.diffusion.models.soulx_singer.modules.llama import DiffLlama
class FlowMatchingTransformer(nn.Module):
def __init__(
self,
mel_dim=100,
hidden_size=1024,
num_layers=12,
num_heads=16,
cfg_drop_prob=0.2,
use_embedding=True,
cond_codebook_size=1024,
cond_scale_factor=1,
sigma=1e-5,
time_scheduler="linear",
cfg=None,
):
super().__init__()
self.cfg = cfg
if cfg is not None:
mel_dim = getattr(cfg, "mel_dim", mel_dim)
hidden_size = getattr(cfg, "hidden_size", hidden_size)
num_layers = getattr(cfg, "num_layers", num_layers)
num_heads = getattr(cfg, "num_heads", num_heads)
cfg_drop_prob = getattr(cfg, "cfg_drop_prob", cfg_drop_prob)
cond_codebook_size = getattr(cfg, "cond_codebook_size", cond_codebook_size)
time_scheduler = getattr(cfg, "time_scheduler", time_scheduler)
sigma = getattr(cfg, "sigma", sigma)
cond_scale_factor = getattr(cfg, "cond_scale_factor", cond_scale_factor)
self.mel_dim = mel_dim
self.hidden_size = hidden_size
self.num_layers = num_layers
self.num_heads = num_heads
self.cfg_drop_prob = cfg_drop_prob
self.cond_codebook_size = cond_codebook_size
self.time_scheduler = time_scheduler
self.sigma = sigma
self.cond_scale_factor = cond_scale_factor
if use_embedding:
self.cond_emb = nn.Embedding(cond_codebook_size, self.hidden_size)
else:
self.cond_emb = nn.Linear(cond_codebook_size, self.hidden_size)
if cond_scale_factor != 1:
self.do_resampling = True
assert np.log2(cond_scale_factor).is_integer()
up_layers = []
for _ in range(int(np.log2(cond_scale_factor))):
up_layers.extend(
[
nn.ConvTranspose1d(hidden_size, hidden_size, kernel_size=4, stride=2, padding=1),
nn.GELU(),
]
)
self.resampling_layers = nn.Sequential(*up_layers)
else:
self.do_resampling = False
### REPA: Use the Wav2Vec2Bert features to align. ###
self.use_repa = "repa" in cfg
self.repa_layer_index = None
if self.use_repa:
self.repa_layer_index = cfg.repa.layer_index
self.repa_mlp_layer = nn.Sequential(
nn.Linear(hidden_size, hidden_size * 4),
nn.SiLU(),
nn.Linear(hidden_size * 4, cfg.repa.output_dim),
)
### CTC: Use the ASR loss ###
self.use_ctc = "ctc" in cfg
self.ctc_layer_index = None
if self.use_ctc:
self.ctc_layer_index = cfg.ctc.layer_index
self.ctc_mlp_layer = nn.Sequential(
nn.Linear(hidden_size, hidden_size * 4),
nn.SiLU(),
nn.Linear(hidden_size * 4, cfg.ctc.output_dim),
)
self.reset_parameters()
# DiffLlama config must be overridden manually to make sure the
# internal RoPE head dimension is correct
llama_config = LlamaConfig(
vocab_size=0,
bos_token_id=None,
eos_token_id=None,
hidden_size=hidden_size,
num_attention_heads=num_heads,
num_hidden_layers=num_layers,
rms_norm_eps=1e-6,
output_hidden_states=False,
output_attentions=False,
use_cache=False,
)
self.diff_estimator = DiffLlama(
mel_dim=mel_dim, hidden_size=hidden_size, num_heads=num_heads, num_layers=num_layers, config=llama_config
)
self.sigma = sigma
@torch.no_grad()
def forward_diffusion(self, x, t, is_prompt=None):
"""
x: (B, T, mel_dim)
t: (B,)
"""
new_t = t
t = t.unsqueeze(-1).unsqueeze(-1)
z = torch.randn(x.shape, dtype=x.dtype, device=x.device, requires_grad=False) # (B, T, mel_dim)
# get prompt len
if torch.rand(1) <= self.cfg_drop_prob:
prompt_len = torch.zeros(x.shape[0]).to(x)
is_prompt = torch.zeros_like(x[:, :, 0])
else:
if is_prompt is None:
prompt_len = torch.randint(min(x.shape[1] // 4, 5), int(x.shape[1] * 0.4), (x.shape[0],)).to(
x.device
) # (B,)
# get is_prompt
is_prompt = torch.zeros_like(x[:, :, 0]) # (B, T)
col_indices = torch.arange(is_prompt.shape[1]).repeat(is_prompt.shape[0], 1).to(prompt_len) # (B, T)
is_prompt[col_indices < prompt_len.unsqueeze(1)] = 1 # (B, T) 1 if prompt
else:
prompt_len = is_prompt.sum(dim=1) # (B,)
mask = torch.ones_like(x[:, :, 0]) # mask if 1, not mask if 0
mask[is_prompt.bool()] = 0
mask = mask[:, :, None]
# flow matching: xt = (1 - (1 - sigma) * t) * x0 + t * x; where x0 ~ N(0, 1), x is a sample
# flow gt: x - (1 - sigma) * x0 = x - (1 - sigma) * noise
xt = ((1 - (1 - self.sigma) * t) * z + t * x) * mask + x * (1 - mask)
return xt, z, new_t, prompt_len, mask
def loss_t(self, x, x_mask, t, cond=None, is_prompt=None):
xt, z, new_t, prompt_len, mask = self.forward_diffusion(x, t, is_prompt)
noise = z
# drop all condition for cfg, so if prompt_len is 0, we also drop cond
if cond is not None:
cond = cond * torch.where(
prompt_len > 0,
torch.ones_like(prompt_len),
torch.zeros_like(prompt_len),
).to(cond.device).unsqueeze(-1).unsqueeze(-1)
dit_output = self.diff_estimator(xt, new_t, cond, x_mask, return_dict=True)
flow_pred = dit_output["output"] # (B, T, mel_dim)
# final mask used for loss calculation
final_mask = mask * x_mask[..., None] # (B, T, 1)
results = {"output": (noise, x, flow_pred, final_mask, prompt_len)}
if self.use_repa:
repa_hidden_states = dit_output["hidden_states"][self.repa_layer_index] # (B, T, hidden_size)
repa_pred = self.repa_mlp_layer(repa_hidden_states) # (B, T, repa_dim)
results["repa"] = repa_pred
if self.use_ctc:
ctc_hidden_states = dit_output["hidden_states"][self.ctc_layer_index] # (B, T, hidden_size)
ctc_pred = self.ctc_mlp_layer(ctc_hidden_states) # (B, T, ctc_dim)
results["ctc"] = ctc_pred
return results
def compute_loss(self, x, x_mask, cond=None, is_prompt=None):
# x0: (B, T, num_quantizer)
# x_mask: (B, T) mask is 0 for padding
t = torch.rand(x.shape[0], device=x.device, requires_grad=False)
t = torch.clamp(t, 1e-5, 1.0)
# generation is harder at the start than later; use a cosine
# scheduler for timestep t when time_scheduler == "cos".
if self.time_scheduler == "cos":
t = 1 - torch.cos(t * math.pi * 0.5)
else:
pass
return self.loss_t(x, x_mask, t, cond, is_prompt)
def reset_parameters(self):
def _reset_parameters(m):
if isinstance(m, nn.MultiheadAttention):
if m._qkv_same_embed_dim:
nn.init.normal_(m.in_proj_weight, std=0.02)
else:
nn.init.normal_(m.q_proj_weight, std=0.02)
nn.init.normal_(m.k_proj_weight, std=0.02)
nn.init.normal_(m.v_proj_weight, std=0.02)
if m.in_proj_bias is not None:
nn.init.constant_(m.in_proj_bias, 0.0)
nn.init.constant_(m.out_proj.bias, 0.0)
if m.bias_k is not None:
nn.init.xavier_normal_(m.bias_k)
if m.bias_v is not None:
nn.init.xavier_normal_(m.bias_v)
elif (
isinstance(m, nn.Conv1d)
or isinstance(m, nn.ConvTranspose1d)
or isinstance(m, nn.Conv2d)
or isinstance(m, nn.ConvTranspose2d)
):
m.weight.data.normal_(0.0, 0.02)
elif isinstance(m, nn.Linear):
m.weight.data.normal_(mean=0.0, std=0.02)
if m.bias is not None:
m.bias.data.zero_()
elif isinstance(m, nn.Embedding):
m.weight.data.normal_(mean=0.0, std=0.02)
if m.padding_idx is not None:
m.weight.data[m.padding_idx].zero_()
self.apply(_reset_parameters)
@torch.no_grad()
def reverse_diffusion(
self,
cond,
prompt,
x_mask=None,
prompt_mask=None,
n_timesteps=10,
cfg=1.0,
rescale_cfg=0.75,
):
h = 1.0 / n_timesteps
prompt_len = prompt.shape[1]
target_len = cond.shape[1] - prompt_len
trunk_dtype = cond.dtype
if x_mask is None:
x_mask = torch.ones(cond.shape[0], target_len, dtype=trunk_dtype, device=cond.device)
if prompt_mask is None:
prompt_mask = torch.ones(cond.shape[0], prompt_len, dtype=trunk_dtype, device=cond.device)
xt_mask = torch.cat([prompt_mask, x_mask], dim=1)
z = torch.randn(
(cond.shape[0], target_len, self.mel_dim),
dtype=trunk_dtype,
device=cond.device,
requires_grad=False,
)
xt = z
h_tensor = torch.tensor(h, dtype=trunk_dtype, device=cond.device)
# t from 0 to 1: x0 = z ~ N(0, 1)
for i in range(n_timesteps):
xt_input = torch.cat([prompt, xt], dim=1)
t = torch.full(
(z.shape[0],),
fill_value=(i + 0.5) * h,
dtype=trunk_dtype,
device=z.device,
)
flow_pred = self.diff_estimator(xt_input, t, cond, xt_mask)
flow_pred = flow_pred[:, prompt_len:, :]
# cfg
if cfg > 0:
uncond_flow_pred = self.diff_estimator(xt, t, torch.zeros_like(cond)[:, : xt.shape[1], :], x_mask)
pos_flow_pred_std = flow_pred.std()
flow_pred_cfg = flow_pred + cfg * (flow_pred - uncond_flow_pred)
rescale_flow_pred = flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std()
flow_pred = rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg
xt = xt + flow_pred * h_tensor
return xt
@torch.no_grad()
def reverse_diffusion_v2(
self,
cond,
prompt,
x_mask=None,
prompt_mask=None,
n_timesteps=10,
cfg=1.0,
rescale_cfg=0.75,
):
h = 1.0 / n_timesteps
prompt_len = prompt.shape[1]
target_len = cond.shape[1] - prompt_len * 2
if x_mask is None:
x_mask = torch.ones(cond.shape[0], target_len).to(cond.device) # (B, T)
if prompt_mask is None:
prompt_mask = torch.ones(cond.shape[0], prompt_len).to(cond.device) # (B, prompt_len)
xt_mask = torch.cat([prompt_mask, x_mask, prompt_mask], dim=1)
z = torch.randn(
(cond.shape[0], target_len, self.mel_dim),
dtype=cond.dtype,
device=cond.device,
requires_grad=False,
)
xt = z
# t from 0 to 1: x0 = z ~ N(0, 1)
for i in range(n_timesteps):
xt_input = torch.cat([prompt, xt, prompt], dim=1)
t = (0 + (i + 0.5) * h) * torch.ones(z.shape[0], dtype=z.dtype, device=z.device)
flow_pred = self.diff_estimator(xt_input, t, cond, xt_mask)
flow_pred = flow_pred[:, prompt_len:-prompt_len, :]
# cfg
if cfg > 0:
uncond_flow_pred = self.diff_estimator(xt, t, torch.zeros_like(cond)[:, : xt.shape[1], :], x_mask)
pos_flow_pred_std = flow_pred.std()
flow_pred_cfg = flow_pred + cfg * (flow_pred - uncond_flow_pred)
rescale_flow_pred = flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std()
flow_pred = rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg
dxt = flow_pred * h
xt = xt + dxt
return xt
def forward(
self,
x: torch.Tensor,
x_mask: torch.Tensor,
cond_code: torch.Tensor,
is_prompt: torch.Tensor | None = None,
):
"""
Args:
x: (B, T, mel_dim)
x_mask: (B, T)
cond_code: (B, T), Note that cond_code might be not at 50Hz!
"""
T = x.shape[1]
cond = self.cond_emb(cond_code) # (B, T, hidden_size)
if self.do_resampling:
# Align to the frame rate of Mels
cond = self.resampling_layers(cond.transpose(1, 2)).transpose(1, 2)
# print("cond_code: {}, after resampling: {}".format(cond_code.shape, cond.shape))
if cond.shape[1] >= T: # Check time dimension
cond = cond[:, :T, :]
else:
padding_frames = T - cond.shape[1]
last_frame = cond[:, -1:, :]
padding = last_frame.repeat(1, padding_frames, 1)
cond = torch.cat([cond, padding], dim=1)
return self.compute_loss(x, x_mask, cond, is_prompt)
if __name__ == "__main__":
model_cfg = {
"mel_dim": 128,
"hidden_size": 256,
"num_layers": 8,
"num_heads": 8,
"cfg_drop_prob": 0.2,
"use_embedding": False,
"cond_codebook_size": 256,
"cond_scale_factor": 1,
"sigma": 1e-5,
"time_scheduler": "cos",
}
device = "cuda"
x = torch.randn(2, 100, 128).to(device)
x_mask = torch.ones(2, 100).to(device)
# cond_code = torch.randint(0, 16384, (2, 25)).to(device)
cond_code = torch.randn(2, 100, 256).to(device)
model = FlowMatchingTransformer(cfg=model_cfg, **model_cfg).to(device)
outputs = model(x, x_mask, cond_code)
print(outputs)
noise, x, flow_pred, final_mask, prompt_len = outputs["output"]
final_mask = final_mask.squeeze(-1)
flow_gt = x - (1 - 1e-5) * noise
# [B, n_frames, D]
diff_loss = F.l1_loss(flow_pred, flow_gt, reduction="none").float() * final_mask.unsqueeze(-1)
diff_loss = torch.mean(diff_loss, dim=2).sum() / final_mask.sum()
print("diff_loss:", diff_loss.item())
diffusion_cond = torch.randn(2, 150, 256).to(device)
diffusion_cond_emb = model.cond_emb(diffusion_cond)
diffusion_prompt = torch.randn(2, 50, 128).to(device)
n_timesteps = 32
generated = model.reverse_diffusion(diffusion_cond_emb, diffusion_prompt, n_timesteps=n_timesteps)
print("generated:", generated.shape)
vllm_omni/diffusion/models/soulx_singer/modules/llama.py
-393
View File
@@ -1,393 +0,0 @@
import math
import torch
import torch.nn as nn
from cache_dit import ForwardPattern
from transformers import LlamaConfig, LlamaModel
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.models.llama.modeling_llama import LlamaDecoderLayer
from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
class SinusoidalPosEmb(nn.Module):
def __init__(self, dim):
super().__init__()
self.dim = dim
def forward(self, x):
device = x.device
dtype = x.dtype
half_dim = self.dim // 2
emb = math.log(10000) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim, device=device, dtype=dtype) * -emb)
emb = x[:, None] * emb[None, :]
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
return emb
class LlamaAdaptiveRMSNorm(nn.Module):
def __init__(self, hidden_size=1024, eps=1e-6, dim_cond=1024):
super().__init__()
self.to_weight = nn.Linear(dim_cond, hidden_size)
nn.init.zeros_(self.to_weight.weight)
nn.init.ones_(self.to_weight.bias)
self.variance_epsilon = eps
self._is_hf_initialized = True
def forward(self, hidden_states, cond_embedding):
input_dtype = hidden_states.dtype
variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
weight = self.to_weight(cond_embedding)
if len(weight.shape) == 2:
weight = weight.unsqueeze(1)
return (weight * hidden_states).to(input_dtype)
class LlamaNARDecoderLayer(LlamaDecoderLayer):
def __init__(self, config: LlamaConfig, layer_idx: int):
super().__init__(config, layer_idx)
self.input_layernorm = LlamaAdaptiveRMSNorm(
config.hidden_size,
eps=config.rms_norm_eps,
dim_cond=config.hidden_size,
)
self.post_attention_layernorm = LlamaAdaptiveRMSNorm(
config.hidden_size,
eps=config.rms_norm_eps,
dim_cond=config.hidden_size,
)
def forward(
self,
hidden_states: torch.Tensor,
cond_embedding: torch.Tensor,
attention_mask: torch.Tensor | None = None,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
output_attentions: bool | None = False,
cache_position: torch.LongTensor | None = None,
**kwargs,
) -> tuple[torch.FloatTensor, ...]:
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states, cond_embedding=cond_embedding)
# RoPE is computed by the parent model and passed in as (cos, sin).
hidden_states, self_attn_weights = self.self_attn(
hidden_states=hidden_states,
position_embeddings=position_embeddings,
attention_mask=attention_mask,
cache_position=cache_position,
output_attentions=output_attentions,
**kwargs,
)
hidden_states = residual + hidden_states
residual = hidden_states
hidden_states = self.post_attention_layernorm(
hidden_states,
cond_embedding=cond_embedding,
)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states
outputs = (hidden_states,)
if output_attentions:
outputs += (self_attn_weights,)
return outputs
class DiffLlama(LlamaModel):
# Regional torch.compile in DiffusionModelRunner targets these blocks.
_repeated_blocks = ["LlamaNARDecoderLayer"]
_layerwise_offload_blocks_attrs = ["layers"]
# cache-dit: Llama-style backbone with static condition embeddings injected
# before the decoder loop (like Stable Audio Open). Pattern_3 handles the
# cross-attention-style conditioning on the ``layers`` ModuleList.
_cache_dit_adapter_config = CacheDiTAdapterConfig(
block_forward_patterns={
"layers": ForwardPattern.Pattern_3,
},
)
# Sequence parallelism: shard mel sequence before the first decoder layer and
# gather after the output projection. Tensor parallelism still requires
# migrating attention to vLLM parallel linear layers.
_sp_plan = {
"layers.0": {
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
},
"mel_out_mlp": SequenceParallelOutput(gather_dim=1, expected_dims=3),
}
@staticmethod
def _is_transformer_block(name: str, module: nn.Module) -> bool:
del module
return name.startswith("layers.") and name.split(".")[-1].isdigit()
_hsdp_shard_conditions = [_is_transformer_block]
def __init__(
self,
mel_dim=100,
hidden_size=1024,
num_heads=16,
num_layers=16,
dropout=0.1,
ffn_dropout=0.1,
attention_dropout=0.0,
config: LlamaConfig | None = None,
):
if config is None:
config = LlamaConfig(
vocab_size=32000,
hidden_size=hidden_size,
intermediate_size=hidden_size * 4,
num_hidden_layers=num_layers,
num_attention_heads=num_heads,
num_key_value_heads=num_heads,
max_position_embeddings=4096,
rms_norm_eps=1e-6,
attention_dropout=attention_dropout,
use_cache=False,
attn_implementation="eager",
)
config.use_cache = False
super().__init__(config)
# Keep custom decoder layers on the same attention backend.
layer_config = LlamaConfig(
vocab_size=config.vocab_size,
hidden_size=hidden_size,
intermediate_size=hidden_size * 4,
num_hidden_layers=num_layers,
num_attention_heads=num_heads,
num_key_value_heads=num_heads,
max_position_embeddings=4096,
rms_norm_eps=config.rms_norm_eps,
attention_dropout=attention_dropout,
use_cache=False,
attn_implementation="eager",
bos_token_id=None,
eos_token_id=None,
)
self.layers = nn.ModuleList(
[
LlamaNARDecoderLayer(
layer_config,
layer_idx=i,
)
for i in range(num_layers)
]
)
self.norm = LlamaAdaptiveRMSNorm(hidden_size, dim_cond=hidden_size)
self.diff_step_embedding = SinusoidalPosEmb(hidden_size)
self.diff_step_mlp = nn.Sequential(
nn.Linear(hidden_size, hidden_size * 4),
nn.SiLU(),
nn.Linear(hidden_size * 4, hidden_size),
)
self.cond_mlp = nn.Sequential(
nn.Linear(hidden_size, hidden_size * 4),
nn.SiLU(),
nn.Linear(hidden_size * 4, hidden_size),
)
self.mel_mlp = nn.Sequential(
nn.Linear(mel_dim, hidden_size * 4),
nn.SiLU(),
nn.Linear(hidden_size * 4, hidden_size),
)
self.mel_out_mlp = nn.Sequential(
nn.Linear(hidden_size, hidden_size * 4),
nn.SiLU(),
nn.Linear(hidden_size * 4, mel_dim),
)
self.embed_tokens = None
self.post_init()
def _prepare_decoder_attention_mask(
self,
attention_mask,
input_shape,
inputs_embeds,
past_key_values_length=0,
):
"""
Build a 4D additive padding mask.
This model is non-autoregressive, so no causal mask is added here.
"""
combined_attention_mask = None
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None):
bsz, src_len = mask.size()
tgt_len = tgt_len if tgt_len is not None else src_len
expanded_mask = (
mask[:, None, None, :]
.expand(
bsz,
1,
tgt_len,
src_len,
)
.to(dtype)
)
inverted_mask = 1.0 - expanded_mask
return inverted_mask.masked_fill(
inverted_mask.to(torch.bool),
torch.finfo(dtype).min,
)
if attention_mask is not None:
expanded_attn_mask = _expand_mask(
attention_mask,
inputs_embeds.dtype,
tgt_len=input_shape[-1],
).to(inputs_embeds.device)
combined_attention_mask = (
expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
)
return combined_attention_mask
def forward(
self,
x,
diffusion_step,
cond,
x_mask,
input_ids: torch.LongTensor = None,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: list[torch.FloatTensor] | None = None,
inputs_embeds: torch.FloatTensor | None = None,
use_cache: bool | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = False,
) -> BaseModelOutputWithPast | torch.Tensor | dict:
batch_size, seq_length, _ = x.shape
cond_embedding = self.cond_mlp(cond)
x = self.mel_mlp(x)
diffusion_step = self.diff_step_embedding(diffusion_step).to(device=x.device, dtype=x.dtype)
diffusion_step = self.diff_step_mlp(diffusion_step)
x = x + cond_embedding
inputs_embeds = x
attention_mask = x_mask
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
# KV cache is disabled because this diffusion decoder processes full sequences.
if past_key_values is not None:
raise NotImplementedError("KV cache is disabled for this non-autoregressive DiffLlama implementation.")
if position_ids is None:
device = input_ids.device if input_ids is not None else inputs_embeds.device
position_ids = torch.arange(
0,
seq_length,
dtype=torch.long,
device=device,
)
position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
else:
position_ids = position_ids.view(batch_size, seq_length).long()
if attention_mask is None:
attention_mask = torch.ones(
(batch_size, seq_length),
dtype=torch.bool,
device=inputs_embeds.device,
)
attention_mask = self._prepare_decoder_attention_mask(
attention_mask,
(batch_size, seq_length),
inputs_embeds,
past_key_values_length=0,
)
hidden_states = inputs_embeds
all_hidden_states = () if output_hidden_states else None
all_self_attns = () if output_attentions else None
all_layer_hidden_states = []
cache_position = torch.arange(
0,
seq_length,
device=hidden_states.device,
dtype=torch.long,
)
# LlamaAttention expects precomputed rotary embeddings.
position_embeddings = self.rotary_emb(hidden_states, position_ids)
for idx, decoder_layer in enumerate(self.layers):
if output_hidden_states:
all_hidden_states += (hidden_states,)
if self.gradient_checkpointing and self.training:
raise NotImplementedError
layer_outputs = decoder_layer(
hidden_states,
cond_embedding=diffusion_step,
attention_mask=attention_mask,
position_embeddings=position_embeddings,
output_attentions=output_attentions,
cache_position=cache_position,
)
hidden_states = layer_outputs[0]
all_layer_hidden_states.append(hidden_states.clone())
if output_attentions:
all_self_attns += (layer_outputs[1],)
hidden_states = self.norm(hidden_states, cond_embedding=diffusion_step)
if output_hidden_states:
all_hidden_states += (hidden_states,)
hidden_states = self.mel_out_mlp(hidden_states)
if return_dict:
return {
"output": hidden_states,
"hidden_states": all_layer_hidden_states,
"attentions": all_self_attns,
}
return hidden_states
vllm_omni/diffusion/models/soulx_singer/modules/mel_transform.py
-151
View File
@@ -1,151 +0,0 @@
import math
from typing import Any
import numpy as np
import torch
import torch.nn as nn
from vllm_omni.utils.audio import mel_filter_bank as librosa_mel_fn
def dynamic_range_compression(x, c=1, clip_val=1e-5):
return np.log(np.clip(x, a_min=clip_val, a_max=None) * c)
def dynamic_range_decompression(x, c=1):
return np.exp(x) / c
def dynamic_range_compression_torch(x, c=1, clip_val=1e-5):
return torch.log(torch.clamp(x, min=clip_val) * c)
def dynamic_range_decompression_torch(x, c=1):
return torch.exp(x) / c
def spectral_normalize_torch(magnitudes):
output = dynamic_range_compression_torch(magnitudes)
return output
def spectral_de_normalize_torch(magnitudes):
output = dynamic_range_decompression_torch(magnitudes)
return output
class MelSpectrogram(nn.Module):
def __init__(
self,
n_fft,
num_mels,
sampling_rate,
hop_size,
win_size,
fmin,
fmax,
center=False,
):
super().__init__()
self.n_fft = n_fft
self.hop_size = hop_size
self.win_size = win_size
self.sampling_rate = sampling_rate
self.num_mels = num_mels
self.fmin = fmin
self.fmax = fmax
self.center = center
mel_basis = {}
hann_window = {}
mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
mel_basis = mel.to(torch.float32)
hann_window = torch.hann_window(win_size)
self.register_buffer("mel_basis", mel_basis)
self.register_buffer("hann_window", hann_window)
def forward(self, y):
y = torch.nn.functional.pad(
y.unsqueeze(1),
(
int((self.n_fft - self.hop_size) / 2),
int((self.n_fft - self.hop_size) / 2),
),
mode="reflect",
)
y = y.squeeze(1)
spec = torch.stft(
y,
self.n_fft,
hop_length=self.hop_size,
win_length=self.win_size,
window=self.hann_window,
center=self.center,
pad_mode="reflect",
normalized=False,
onesided=True,
return_complex=True,
)
spec = torch.view_as_real(spec)
spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))
spec = torch.matmul(self.mel_basis, spec)
spec = spectral_normalize_torch(spec)
return spec
def load_mel_spectrogram():
return load_mel_spectrogram_from_cfg(None)
def _get_from_mapping(cfg: Any, key: str, default: Any = None) -> Any:
"""Safely read a field from a dict/OmegaConf-like object."""
if cfg is None:
return default
if isinstance(cfg, dict):
return cfg.get(key, default)
return getattr(cfg, key, default)
def load_mel_spectrogram_from_cfg(audio_cfg: Any | None = None) -> MelSpectrogram:
"""Build MelSpectrogram from `audio_config`-like config.
Expected keys (either in dict or Hydra/OmegaConf object):
- hop_size, sample_rate (or sampling_rate), n_fft, num_mels, win_size, fmin, fmax
"""
# Defaults keep current behavior.
mel_cfg: dict[str, Any] = {
"hop_size": _get_from_mapping(audio_cfg, "hop_size", 480),
"sampling_rate": _get_from_mapping(
audio_cfg,
"sampling_rate",
_get_from_mapping(audio_cfg, "sample_rate", 24000),
),
"n_fft": _get_from_mapping(audio_cfg, "n_fft", 1920),
"num_mels": _get_from_mapping(audio_cfg, "num_mels", 128),
"win_size": _get_from_mapping(audio_cfg, "win_size", 1920),
"fmin": _get_from_mapping(audio_cfg, "fmin", 0),
"fmax": _get_from_mapping(audio_cfg, "fmax", 12000),
}
mel_model = MelSpectrogram(**mel_cfg)
mel_model.eval()
return mel_model
class MelSpectrogramEncoder(nn.Module):
def __init__(self, audio_config: dict | None = None):
super().__init__()
self.model = load_mel_spectrogram_from_cfg(audio_config)
audio_config = audio_config or {}
self.mel_mean = audio_config.get("mel_mean", -4.92)
self.mel_var = audio_config.get("mel_var", 8.14)
def forward(self, x):
x = self.model(x).transpose(1, 2)
x = (x - self.mel_mean) / math.sqrt(self.mel_var)
return x
vllm_omni/diffusion/models/soulx_singer/modules/note_transcription.py
-958
View File
@@ -1,958 +0,0 @@
import math
import numpy as np
import packaging.version as version
import torch
import torch.nn as nn
# ---------------------- note transcription common layers ----------------------
class LayerNorm(nn.LayerNorm):
def __init__(self, nout: int, dim: int = -1, eps: float = 1e-5):
"""Construct an LayerNorm object."""
super().__init__(nout, eps=eps)
self.dim = dim
def forward(self, x):
if self.dim == -1:
return super().forward(x)
return super().forward(x.transpose(1, -1)).transpose(1, -1)
class Reshape(nn.Module):
def __init__(self, *args):
super().__init__()
self.shape = args
def forward(self, x):
return x.view(self.shape)
class Permute(nn.Module):
def __init__(self, *args):
super().__init__()
self.args = args
def forward(self, x):
return x.permute(self.args)
class LambdaLayer(nn.Module):
def __init__(self, lambd):
super().__init__()
self.lambd = lambd
def forward(self, x):
return self.lambd(x)
class ResidualBlock(nn.Module):
"""
Implements a residual block pattern:
[Norm] → [Conv1d] → [Scale] → [Activation] → [Conv1d]
repeated n times, each as a residual unit.
"""
def __init__(
self,
channels: int,
kernel_size: int,
dilation: int,
n: int = 2,
c_multiple: int = 2,
ln_eps: float = 1e-12,
):
super().__init__()
self.blocks = nn.ModuleList(
[
nn.Sequential(
LayerNorm(channels, dim=1, eps=ln_eps),
nn.Conv1d(
channels,
c_multiple * channels,
kernel_size,
dilation=dilation,
padding=(dilation * (kernel_size - 1)) // 2,
),
LambdaLayer(lambda x: x * kernel_size**-0.5), # scaling for stability
nn.LeakyReLU(),
nn.Conv1d(c_multiple * channels, channels, kernel_size=1, dilation=dilation),
)
for _ in range(n)
]
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
nonpadding = (x.abs().sum(1) > 0).float()[:, None, :]
for block in self.blocks:
x_res = block(x)
x = (x + x_res) * nonpadding
return x
class ConvBlocks(nn.Module):
"""
Decodes the expanded phoneme encoding into spectrograms.
"""
def __init__(
self,
hidden_size,
out_dims,
dilations,
kernel_size,
layers_in_block=2,
c_multiple=2,
ln_eps=1e-5,
is_btc=True,
post_net_kernel=3,
):
super().__init__()
self.is_btc = is_btc
res_blocks = [
ResidualBlock(
hidden_size,
kernel_size=kernel_size,
dilation=1,
n=layers_in_block,
c_multiple=c_multiple,
ln_eps=ln_eps,
)
for _ in range(1)
] # default 1
self.res_blocks = nn.Sequential(*res_blocks)
# Output normalization and conv
self.last_norm = LayerNorm(hidden_size, dim=1, eps=ln_eps)
self.post_net1 = nn.Conv1d(hidden_size, out_dims, kernel_size=post_net_kernel, padding=post_net_kernel // 2)
def forward(self, x: torch.Tensor, nonpadding: torch.Tensor | None = None) -> torch.Tensor:
if self.is_btc:
x = x.transpose(1, 2)
if nonpadding is None:
nonpadding = (x.abs().sum(1) > 0).float()[:, None, :]
elif self.is_btc:
nonpadding = nonpadding.transpose(1, 2)
x = self.res_blocks(x) * nonpadding
x = self.last_norm(x) * nonpadding
x = self.post_net1(x) * nonpadding
if self.is_btc:
x = x.transpose(1, 2)
return x
def Linear(in_features: int, out_features: int, bias: bool = True, init_type: str = "xavier"):
m = nn.Linear(in_features, out_features, bias)
if init_type == "xavier":
nn.init.xavier_uniform_(m.weight)
elif init_type == "kaiming":
nn.init.kaiming_normal_(m.weight, mode="fan_in")
if bias:
nn.init.constant_(m.bias, 0.0)
return m
def Embedding(num_embeddings, embedding_dim, padding_idx=None, init_type="normal"):
m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)
if init_type == "normal":
nn.init.normal_(m.weight, mean=0, std=embedding_dim**-0.5)
elif init_type == "kaiming":
nn.init.kaiming_normal_(m.weight, mode="fan_in")
if padding_idx is not None:
nn.init.constant_(m.weight[padding_idx], 0)
return m
# ---------------------- note transcription espnet attention ----------------------
class MultiHeadedAttention(nn.Module):
"""Multi-Head Attention layer.
Args:
n_head (int): The number of heads.
n_feat (int): The number of features.
dropout_rate (float): Dropout rate.
"""
def __init__(self, n_head, n_feat, dropout_rate, flash=False):
"""Construct an MultiHeadedAttention object."""
super().__init__()
assert n_feat % n_head == 0
# We assume d_v always equals d_k
self.d_k = n_feat // n_head
self.h = n_head
self.linear_q = nn.Linear(n_feat, n_feat)
self.linear_k = nn.Linear(n_feat, n_feat)
self.linear_v = nn.Linear(n_feat, n_feat)
self.linear_out = nn.Linear(n_feat, n_feat)
self.attn = None
self.dropout = nn.Dropout(p=dropout_rate)
self.dropout_rate = dropout_rate
self.flash = flash
def forward_qkv(self, query, key, value):
"""Transform query, key and value.
Args:
query (torch.Tensor): Query tensor (#batch, time1, size).
key (torch.Tensor): Key tensor (#batch, time2, size).
value (torch.Tensor): Value tensor (#batch, time2, size).
Returns:
torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k).
torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k).
torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k).
"""
n_batch = query.size(0)
q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
q = q.transpose(1, 2) # (batch, head, time1, d_k)
k = k.transpose(1, 2) # (batch, head, time2, d_k)
v = v.transpose(1, 2) # (batch, head, time2, d_k)
return q, k, v
def forward_attention(self, value, scores, mask):
"""Compute attention context vector.
Args:
value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k).
scores (torch.Tensor): Attention score (#batch, n_head, time1, time2).
mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2).
Returns:
torch.Tensor: Transformed value (#batch, time1, d_model)
weighted by the attention score (#batch, time1, time2).
"""
n_batch = value.size(0)
if mask is not None:
mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2)
min_value = float(np.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min)
scores = scores.masked_fill(mask, min_value)
self.attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0) # (batch, head, time1, time2)
else:
self.attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2)
p_attn = self.dropout(self.attn)
x = torch.matmul(p_attn, value) # (batch, head, time1, d_k)
x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model)
return self.linear_out(x) # (batch, time1, d_model)
def forward(self, query, key, value, mask):
"""Compute scaled dot product attention.
Args:
query (torch.Tensor): Query tensor (#batch, time1, size).
key (torch.Tensor): Key tensor (#batch, time2, size).
value (torch.Tensor): Value tensor (#batch, time2, size).
mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
(#batch, time1, time2).
Returns:
torch.Tensor: Output tensor (#batch, time1, d_model).
"""
q, k, v = self.forward_qkv(query, key, value)
if version.parse(torch.__version__) >= version.parse("2.0") and self.flash:
n_batch = value.size(0)
x = torch.nn.functional.scaled_dot_product_attention(
q, k, v, attn_mask=mask.unsqueeze(1) if mask is not None else None, dropout_p=self.dropout_rate
)
x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model)
return self.linear_out(x)
else:
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
return self.forward_attention(v, scores, mask)
class RelPositionMultiHeadedAttention(MultiHeadedAttention):
"""Multi-Head Attention layer with relative position encoding.
Paper: https://arxiv.org/abs/1901.02860
Args:
n_head (int): The number of heads.
n_feat (int): The number of features.
dropout_rate (float): Dropout rate.
"""
def __init__(self, n_head, n_feat, dropout_rate):
"""Construct an RelPositionMultiHeadedAttention object."""
super().__init__(n_head, n_feat, dropout_rate)
# linear transformation for positional encoding
self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
# these two learnable bias are used in matrix c and matrix d
# as described in https://arxiv.org/abs/1901.02860 Section 3.3
self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k))
self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k))
torch.nn.init.xavier_uniform_(self.pos_bias_u)
torch.nn.init.xavier_uniform_(self.pos_bias_v)
def rel_shift(self, x, zero_triu=False):
"""Compute relative positinal encoding.
Args:
x (torch.Tensor): Input tensor (batch, time, size).
zero_triu (bool): If true, return the lower triangular part of the matrix.
Returns:
torch.Tensor: Output tensor.
"""
zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype)
x_padded = torch.cat([zero_pad, x], dim=-1)
x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2))
x = x_padded[:, :, 1:].view_as(x)
if zero_triu:
ones = torch.ones((x.size(2), x.size(3)))
x = x * torch.tril(ones, x.size(3) - x.size(2))[None, None, :, :]
return x
def forward(self, query, key, value, pos_emb, mask):
"""Compute 'Scaled Dot Product Attention' with rel. positional encoding.
Args:
query (torch.Tensor): Query tensor (#batch, time1, size).
key (torch.Tensor): Key tensor (#batch, time2, size).
value (torch.Tensor): Value tensor (#batch, time2, size).
pos_emb (torch.Tensor): Positional embedding tensor (#batch, time2, size).
mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
(#batch, time1, time2).
Returns:
torch.Tensor: Output tensor (#batch, time1, d_model).
"""
q, k, v = self.forward_qkv(query, key, value)
q = q.transpose(1, 2) # (batch, time1, head, d_k)
n_batch_pos = pos_emb.size(0)
p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)
p = p.transpose(1, 2) # (batch, head, time1, d_k)
# (batch, head, time1, d_k)
q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
# (batch, head, time1, d_k)
q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
# compute attention score
# first compute matrix a and matrix c
# as described in https://arxiv.org/abs/1901.02860 Section 3.3
# (batch, head, time1, time2)
matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
# compute matrix b and matrix d
# (batch, head, time1, time2)
matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
matrix_bd = self.rel_shift(matrix_bd)
scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2)
return self.forward_attention(v, scores, mask)
# ---------------------- note transcription positional encoding ----------------------
class PositionalEncoding(torch.nn.Module):
"""Positional encoding.
Args:
d_model (int): Embedding dimension.
max_len (int): Maximum input length.
reverse (bool): Whether to reverse the input position.
"""
def __init__(self, d_model, max_len=5000, reverse=False):
"""Construct an PositionalEncoding object."""
super().__init__()
self.d_model = d_model
self.reverse = reverse
self.xscale = math.sqrt(self.d_model)
self.pe = None
self.extend_pe(torch.tensor(0.0).expand(1, max_len))
def extend_pe(self, x):
"""Reset the positional encodings."""
if self.pe is not None:
if self.pe.size(1) >= x.size(1):
if self.pe.dtype != x.dtype or self.pe.device != x.device:
self.pe = self.pe.to(dtype=x.dtype, device=x.device)
return
pe = torch.zeros(x.size(1), self.d_model)
if self.reverse:
position = torch.arange(x.size(1) - 1, -1, -1.0, dtype=torch.float32).unsqueeze(1)
else:
position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
div_term = torch.exp(
torch.arange(0, self.d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / self.d_model)
)
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
self.pe = pe.to(device=x.device, dtype=x.dtype)
def forward(self, x: torch.Tensor):
"""Add positional encoding.
Args:
x (torch.Tensor): Input tensor (batch, time, `*`).
Returns:
torch.Tensor: Encoded tensor (batch, time, `*`).
"""
self.extend_pe(x)
x = x * self.xscale + self.pe[:, : x.size(1)]
return x
class RelPositionalEncoding(PositionalEncoding):
"""Relative positional encoding module.
See : Appendix B in https://arxiv.org/abs/1901.02860
Args:
d_model (int): Embedding dimension.
max_len (int): Maximum input length.
"""
def __init__(self, d_model, max_len=5000):
super().__init__(d_model, max_len, reverse=True)
def forward(self, x):
"""Compute positional encoding.
Args:
x (torch.Tensor): Input tensor (batch, time, `*`).
Returns:
torch.Tensor: Encoded tensor (batch, time, `*`).
torch.Tensor: Positional embedding tensor (1, time, `*`).
"""
self.extend_pe(x)
x = x * self.xscale
pos_emb = self.pe[:, : x.size(1)]
return x, pos_emb
# ---------------------- note transcription conformer layers ----------------------
class ConvolutionModule(nn.Module):
def __init__(self, channels: int, kernel_size: int, activation: nn.Module = nn.ReLU(), bias: bool = True):
super().__init__()
if (kernel_size - 1) % 2 != 0:
raise ValueError("kernel_size should be a odd number for 'SAME' padding")
self.pointwise_conv1 = nn.Conv1d(channels, 2 * channels, kernel_size=1, bias=bias)
self.depthwise_conv = nn.Conv1d(
channels, channels, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, groups=channels, bias=bias
)
self.norm = nn.BatchNorm1d(channels)
self.pointwise_conv2 = nn.Conv1d(channels, channels, kernel_size=1, bias=bias)
self.activation = activation
def forward(self, x):
# (batch, time, channels) -> (batch, channels, time)
x = x.transpose(1, 2)
x = self.pointwise_conv1(x)
x = nn.functional.glu(x, dim=1)
x = self.depthwise_conv(x)
x = self.norm(x)
x = self.activation(x)
x = self.pointwise_conv2(x)
x = x.transpose(1, 2)
return x
class MultiLayeredConv1d(nn.Module):
"""
Multi-layered Conv1d that replaces feed-forward in Transformer block.
Reference:
FastSpeech: Fast, Robust and Controllable Text to Speech
https://arxiv.org/pdf/1905.09263.pdf
"""
def __init__(self, in_chans, hidden_chans, kernel_size):
super().__init__()
padding = (kernel_size - 1) // 2
self.w_1 = nn.Conv1d(in_chans, hidden_chans, kernel_size, padding=padding)
self.w_2 = nn.Conv1d(hidden_chans, in_chans, kernel_size, padding=padding)
def forward(self, x):
"""
Args:
x (torch.Tensor): (batch, time, in_chans)
Returns:
torch.Tensor: (batch, time, in_chans)
"""
# Conv1d expects (batch, channels, time)
x = x.transpose(1, 2)
x = torch.relu(self.w_1(x))
x = self.w_2(x)
x = x.transpose(1, 2)
return x
class Swish(nn.Module):
"""Swish activation function."""
def forward(self, x):
return x * torch.sigmoid(x)
class EncoderLayer(nn.Module):
"""
A single encoder layer: optional macaron FFN, self-attn, conv, FFN, norms.
Args:
size (int): Input dimension.
self_attn (nn.Module): Self-attention module.
feed_forward (nn.Module): Standard feed-forward network.
feed_forward_macaron (nn.Module): Optional, second FFN in macaron structure.
conv_module (nn.Module): Optional, convolution module.
normalize_before (bool): If True, use pre-layer norm. Else, post-layer norm.
concat_after (bool): If True, concat attn in/out, else residual.
"""
def __init__(
self,
size,
self_attn,
feed_forward,
feed_forward_macaron,
conv_module,
normalize_before=True,
concat_after=False,
):
super().__init__()
self.size = size
self.self_attn = self_attn
self.feed_forward = feed_forward
self.feed_forward_macaron = feed_forward_macaron
self.conv_module = conv_module
self.norm_ff = LayerNorm(size)
self.norm_mha = LayerNorm(size)
self.ff_scale = 0.5 if feed_forward_macaron is not None else 1.0
self.norm_ff_macaron = LayerNorm(size) if feed_forward_macaron is not None else None
if conv_module is not None:
self.norm_conv = LayerNorm(size)
self.norm_final = LayerNorm(size)
else:
self.norm_conv = None
self.norm_final = None
self.normalize_before = normalize_before
self.concat_after = concat_after
if self.concat_after:
self.concat_linear = nn.Linear(2 * size, size)
def forward(self, x_input, mask, cache=None) -> tuple:
"""
Forward pass for the encoder layer.
Args:
x_input (Tensor or Tuple): Input tensor of shape (batch, time, size)
or tuple (x, pos_emb) where pos_emb is the positional embedding.
mask (Tensor): Mask tensor of shape (batch, time) or (batch, 1, time).
cache (Tensor, optional): Input cache for streaming (can be None).
Returns:
Tuple[Tensor, Tensor]: (outputs, mask), where outputs either includes
pos_emb or not, depending on input.
"""
# Unpack possible (x, pos_emb) tuple input
if isinstance(x_input, tuple):
x, pos_emb = x_input
else:
x, pos_emb = x_input, None
# --- Macaron Feed-Forward (first FFN) ---
if self.feed_forward_macaron is not None:
residual = x
if self.normalize_before:
x = self.norm_ff_macaron(x)
x = residual + self.ff_scale * self.feed_forward_macaron(x)
if not self.normalize_before:
x = self.norm_ff_macaron(x)
# --- Self-Attention ---
residual = x
if self.normalize_before:
x = self.norm_mha(x)
# Handle cache for streaming
if cache is not None:
assert cache.shape == (x.shape[0], x.shape[1] - 1, self.size)
x_q = x[:, -1:, :]
residual = residual[:, -1:, :]
mask = mask[:, -1:, :] if mask is not None else None
else:
x_q = x
# Apply self-attention (with or without positional embeddings)
if pos_emb is not None:
attn_out = self.self_attn(x_q, x, x, pos_emb, mask)
else:
attn_out = self.self_attn(x_q, x, x, mask)
if self.concat_after:
x = torch.cat((x, attn_out), dim=-1)
x = residual + self.concat_linear(x)
else:
x = residual + attn_out
if not self.normalize_before:
x = self.norm_mha(x)
# --- Convolution Module ---
if self.conv_module is not None:
residual = x
if self.normalize_before:
x = self.norm_conv(x)
x = residual + self.conv_module(x)
if not self.normalize_before:
x = self.norm_conv(x)
# --- Feed-Forward (second FFN) ---
residual = x
if self.normalize_before:
x = self.norm_ff(x)
x = residual + self.ff_scale * self.feed_forward(x)
if not self.normalize_before:
x = self.norm_ff(x)
# --- Final normalization, if present ---
if self.norm_final is not None:
x = self.norm_final(x)
# --- Append to cache if streaming ---
if cache is not None:
x = torch.cat([cache, x], dim=1)
# --- Final return, with or without positional embedding ---
if pos_emb is not None:
return (x, pos_emb), mask
else:
return x, mask
class ConformerLayers(nn.Module):
def __init__(self, hidden_size, num_layers, kernel_size=9, num_heads=4, use_last_norm=True, save_hidden=False):
super().__init__()
self.use_last_norm = use_last_norm
self.layers = nn.ModuleList()
positionwise_layer = MultiLayeredConv1d
positionwise_layer_args = (hidden_size, hidden_size * 4, 1)
self.pos_embed = RelPositionalEncoding(hidden_size)
self.encoder_layers = nn.ModuleList(
[
EncoderLayer(
hidden_size,
RelPositionMultiHeadedAttention(num_heads, hidden_size, 0.0),
positionwise_layer(*positionwise_layer_args),
positionwise_layer(*positionwise_layer_args),
ConvolutionModule(hidden_size, kernel_size, Swish()),
)
for _ in range(num_layers)
]
)
if self.use_last_norm:
self.layer_norm = nn.LayerNorm(hidden_size)
else:
self.layer_norm = nn.Linear(hidden_size, hidden_size)
self.save_hidden = save_hidden
if save_hidden:
self.hiddens = []
def forward(self, x, padding_mask=None):
"""
:param x: [B, T, H]
:param padding_mask: [B, T]
:return: [B, T, H]
"""
self.hiddens = []
nonpadding_mask = x.abs().sum(-1) > 0
x = self.pos_embed(x)
for layer in self.encoder_layers:
x, mask = layer(x, nonpadding_mask[:, None, :])
if self.save_hidden:
self.hiddens.append(x[0])
x = x[0]
x = self.layer_norm(x) * nonpadding_mask.float()[:, :, None]
return x
# ---------------------- note transcription backbone layers ----------------------
class UnetDown(nn.Module):
def __init__(
self,
hidden_size,
n_layers,
kernel_size,
down_rates,
channel_multiples=None,
is_btc=True,
constant_channels=False,
):
super().__init__()
assert n_layers == len(down_rates) # downs, down sample rate
down_rates = [int(i) for i in down_rates]
self.n_layers = n_layers
self.hidden_size = hidden_size
self.is_btc = is_btc
channel_multiples = channel_multiples if channel_multiples is not None else down_rates
self.layers = nn.ModuleList()
self.downs = nn.ModuleList()
in_channels = hidden_size
for i in range(self.n_layers):
out_channels = int(in_channels * channel_multiples[i]) if not constant_channels else in_channels
self.layers.append(
nn.Sequential(
ResidualBlock(in_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5),
nn.Conv1d(in_channels, out_channels, kernel_size, padding=(kernel_size - 1) // 2),
ResidualBlock(out_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5),
)
)
self.downs.append(nn.Sequential(nn.AvgPool1d(down_rates[i])))
in_channels = out_channels
self.last_norm = LayerNorm(out_channels, dim=1, eps=1e-6)
self.post_net = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2)
def forward(self, x, **kwargs):
# x [B, T, C]
if self.is_btc:
x = x.transpose(1, 2) # [B, C, T]
skip_xs = []
for i in range(self.n_layers):
skip_x = self.layers[i](x)
x = self.downs[i](skip_x)
if self.is_btc:
skip_xs.append(skip_x.transpose(1, 2)) # [B, T, C]
else:
skip_xs.append(skip_x)
x = self.post_net(self.last_norm(x))
if self.is_btc:
x = x.transpose(1, 2)
return x, skip_xs
class UnetMid(nn.Module):
def __init__(self, hidden_size, kernel_size, n_layers=None, in_dims=None, out_dims=None, is_btc=True, net=None):
super().__init__()
in_dims = in_dims if in_dims is not None else hidden_size
out_dims = out_dims if out_dims is not None else hidden_size
self.pre = nn.Conv1d(in_dims, hidden_size, kernel_size, padding=kernel_size // 2)
self.post = nn.Conv1d(hidden_size, out_dims, kernel_size, padding=kernel_size // 2)
self.is_btc = is_btc
if net is not None:
self.net = net
else:
self.net = ConvBlocks(
hidden_size,
out_dims=hidden_size,
dilations=None,
kernel_size=kernel_size,
layers_in_block=2,
c_multiple=2,
post_net_kernel=3,
is_btc=is_btc,
)
def forward(self, x, cond=None, **kwargs):
# x [B, T, C]
if self.is_btc:
x = self.pre(x.transpose(1, 2)).transpose(1, 2)
else:
x = self.pre(x)
if cond is None:
cond = 0
x = self.net(x + cond)
if self.is_btc:
x = self.post(x.transpose(1, 2)).transpose(1, 2)
else:
x = self.post(x)
return x
class UnetUp(nn.Module):
def __init__(
self,
hidden_size,
n_layers,
kernel_size,
up_rates,
channel_multiples=None,
is_btc=True,
constant_channels=False,
use_skip_layer=False,
skip_scale=1.0,
):
super().__init__()
assert n_layers == len(up_rates) # this is reversed in up module, from the output to the interface with middle
up_rates = [int(i) for i in up_rates]
self.n_layers = n_layers
self.hidden_size = hidden_size
self.is_btc = is_btc
self.skip_scale = skip_scale
channel_multiples = channel_multiples if channel_multiples is not None else up_rates
# in_channels = int(np.cumprod(channel_multiples)[-1] * hidden_size) if not constant_channels else hidden_size
self.in_channels_lst = (
(np.cumprod([1] + channel_multiples) * hidden_size).astype(int)
if not constant_channels
else [hidden_size for _ in range(self.n_layers + 1)]
)
in_channels = self.in_channels_lst[-1]
self.ups = nn.ModuleList()
self.skip_layers = nn.ModuleList()
self.layers = nn.ModuleList()
for i in range(self.n_layers - 1, -1, -1):
out_channels = self.in_channels_lst[i] if not constant_channels else in_channels
self.ups.append(
nn.Sequential(
nn.ConvTranspose1d(
in_channels,
in_channels,
kernel_size=kernel_size,
stride=up_rates[i],
padding=kernel_size // 2,
output_padding=up_rates[i] - 1,
),
LayerNorm(in_channels, dim=1, eps=1e-6),
nn.LeakyReLU(),
)
)
self.layers.append(
nn.Sequential(
nn.Conv1d(in_channels * 2, out_channels, kernel_size, padding=(kernel_size - 1) // 2),
ResidualBlock(out_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5),
)
)
if use_skip_layer:
self.skip_layers.append(
ResidualBlock(in_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5)
)
else:
self.skip_layers.append(nn.Identity())
in_channels = out_channels
self.out_channels = out_channels
self.last_norm = LayerNorm(out_channels, dim=1, eps=1e-6)
self.post_net = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2)
def forward(self, x, skips, **kwargs):
# x [B, T, C]
if self.is_btc:
x = x.transpose(1, 2) # [B, C, T]
for i in range(self.n_layers):
x = self.ups[i](x)
skip_x = (
skips[self.n_layers - i - 1] if not self.is_btc else skips[self.n_layers - i - 1].transpose(1, 2)
) # [B, T, C] -> [B, C, T]
skip_x = self.skip_layers[i](skip_x) * self.skip_scale
x = torch.cat((x, skip_x), dim=1) # [B, C, T]
x = self.layers[i](x)
x = self.post_net(self.last_norm(x))
if self.is_btc:
x = x.transpose(1, 2)
return x
class Unet(nn.Module):
def __init__(
self,
hidden_size,
down_layers,
up_layers,
kernel_size,
updown_rates,
mid_layers=None,
channel_multiples=None,
is_btc=True,
constant_channels=False,
mid_net=None,
use_skip_layer=False,
skip_scale=1.0,
):
super().__init__()
assert len(updown_rates) == down_layers == up_layers, f"{len(updown_rates)}, {down_layers}, {up_layers}"
if channel_multiples is not None:
assert len(channel_multiples) == len(updown_rates)
else:
channel_multiples = updown_rates
self.down = UnetDown(
hidden_size, down_layers, kernel_size, updown_rates, channel_multiples, is_btc, constant_channels
)
down_out_dims = int(np.cumprod(channel_multiples)[-1] * hidden_size) if not constant_channels else hidden_size
self.mid = UnetMid(
hidden_size,
kernel_size,
mid_layers,
in_dims=down_out_dims,
out_dims=down_out_dims,
is_btc=is_btc,
net=mid_net,
)
self.up = UnetUp(
hidden_size,
up_layers,
kernel_size,
updown_rates,
channel_multiples,
is_btc,
constant_channels,
use_skip_layer,
skip_scale,
)
def forward(self, x, mid_cond=None, **kwargs):
x, skips = self.down(x)
x = self.mid(x, mid_cond)
x = self.up(x, skips)
return x
class BackboneNet(nn.Module):
def __init__(self, hparams):
super().__init__()
self.hidden_size = hidden_size = hparams["hidden_size"]
updown_rates = [2, 2, 2]
channel_multiples = [1, 1, 1]
if hparams.get("updown_rates", None) is not None:
updown_rates = [int(i) for i in hparams.get("updown_rates", None).split("-")]
if hparams.get("channel_multiples", None) is not None:
channel_multiples = [float(i) for i in hparams.get("channel_multiples", None).split("-")]
assert len(updown_rates) == len(channel_multiples)
# convs
if hparams.get("bkb_net", "conv") == "conv":
self.net = Unet(
hidden_size,
down_layers=len(updown_rates),
mid_layers=hparams.get("bkb_layers", 12),
up_layers=len(updown_rates),
kernel_size=3,
updown_rates=updown_rates,
channel_multiples=channel_multiples,
is_btc=True,
constant_channels=False,
mid_net=None,
use_skip_layer=hparams.get("unet_skip_layer", False),
)
# conformer
elif hparams.get("bkb_net", "conv") == "conformer":
mid_net = ConformerLayers(
hidden_size,
num_layers=hparams.get("bkb_layers", 12),
kernel_size=hparams.get("conformer_kernel", 9),
num_heads=4,
)
self.net = Unet(
hidden_size,
down_layers=len(updown_rates),
up_layers=len(updown_rates),
kernel_size=3,
updown_rates=updown_rates,
channel_multiples=channel_multiples,
is_btc=True,
constant_channels=False,
mid_net=mid_net,
use_skip_layer=hparams.get("unet_skip_layer", False),
)
def forward(self, x):
return self.net(x)
vllm_omni/diffusion/models/soulx_singer/modules/preprocess/__init__.py
-7
View File
@@ -1,7 +0,0 @@
"""Preprocess neural network modules for SoulX-Singer."""
from .pipeline import SoulXPreprocessPipeline
__all__ = [
"SoulXPreprocessPipeline",
]
vllm_omni/diffusion/models/soulx_singer/modules/preprocess/asr.py
-326
View File
@@ -1,326 +0,0 @@
"""Lyric transcription (Paraformer + Parakeet) for SoulX preprocess."""
import os
import re
import tempfile
from pathlib import Path
from typing import Any
import numpy as np
import soundfile as sf
import torch
import torch.nn as nn
from omegaconf import open_dict
from vllm.logger import init_logger
from .utils import get_audio_duration
logger = init_logger(__name__)
def _import_funasr_automodel():
"""Import FunASR without shelling out to ffmpeg at module import time."""
import subprocess
from functools import wraps
real_check_output = subprocess.check_output
@wraps(real_check_output)
def _safe_check_output(cmd, *args, **kwargs):
if cmd and Path(str(cmd[0])).name == "ffmpeg":
raise FileNotFoundError("ffmpeg unavailable for FunASR import probe")
return real_check_output(cmd, *args, **kwargs)
subprocess.check_output = _safe_check_output
try:
from funasr import AutoModel
finally:
subprocess.check_output = real_check_output
return AutoModel
def build_words_with_gaps(raw_words, raw_timestamps, audio_duration: float):
words, word_durs = [], []
prev = 0.0
for w, t in zip(raw_words, raw_timestamps):
s, e = float(t[0]), float(t[1])
if s > prev:
words.append("<SP>")
word_durs.append(s - prev)
words.append(w)
word_durs.append(e - s)
prev = e
if audio_duration > prev:
if not words:
words.append("<SP>")
word_durs.append(audio_duration)
elif words[-1] != "<SP>":
words.append("<SP>")
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 == "<SP>":
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("<SP>")
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
vllm_omni/diffusion/models/soulx_singer/modules/preprocess/mel_grid_f0.py
-96
View File
@@ -1,96 +0,0 @@
"""Resample RMVPE F0 (16 kHz hop) onto the SoulX mel frame grid."""
from __future__ import annotations
import math
import time
from typing import TYPE_CHECKING
import numpy as np
from scipy.interpolate import interp1d
from vllm.logger import init_logger
from .rmvpe import RMVPE_SAMPLE_RATE
from .utils import load_mono_audio
if TYPE_CHECKING:
from .rmvpe import RMVPE
logger = init_logger(__name__)
RMVPE_HOP = 160
def interpolate_f0_to_mel_grid(
f0_16k: np.ndarray,
original_length: int,
original_sr: int,
*,
target_sr: int,
hop_size: int,
max_duration: float = 300.0,
) -> np.ndarray:
batch_max_length = int(max_duration * target_sr / hop_size)
duration_in_seconds = original_length / original_sr
effective_target_length = int(duration_in_seconds * target_sr)
original_frames = math.ceil(effective_target_length / hop_size)
target_frames = min(original_frames, batch_max_length)
t_16k = np.arange(len(f0_16k)) * (RMVPE_HOP / float(original_sr))
t_target = np.arange(target_frames) * (hop_size / float(target_sr))
if len(f0_16k) > 0:
f_interp = interp1d(
t_16k,
f0_16k,
kind="linear",
bounds_error=False,
fill_value=0.0,
assume_sorted=True,
)
f0 = f_interp(t_target)
else:
f0 = np.zeros(target_frames)
if len(f0) != target_frames:
f0 = f0[:target_frames] if len(f0) > target_frames else np.pad(f0, (0, target_frames - len(f0)))
return f0.astype(np.float32, copy=False)
def extract_f0_file(
rmvpe: RMVPE,
wav_path: str,
*,
target_sr: int,
hop_size: int,
f0_path: str | None = None,
thread: float = 0.03,
max_duration: float = 300.0,
verbose: bool = False,
) -> np.ndarray:
if verbose:
logger.info("[f0] extracting from %s", wav_path)
t0 = time.perf_counter()
audio, _ = load_mono_audio(wav_path, target_sr=RMVPE_SAMPLE_RATE)
f0_16k = rmvpe.infer_from_audio(audio, sample_rate=RMVPE_SAMPLE_RATE)
f0 = interpolate_f0_to_mel_grid(
np.asarray(f0_16k, dtype=np.float32),
original_length=audio.shape[-1],
original_sr=RMVPE_SAMPLE_RATE,
target_sr=target_sr,
hop_size=hop_size,
max_duration=max_duration,
)
if verbose:
voiced_ratio = float(np.mean(f0 > 0)) if len(f0) else 0.0
logger.info(
"[f0] done: frames=%s voiced_ratio=%.3f time=%.3fs",
len(f0),
voiced_ratio,
time.perf_counter() - t0,
)
if f0_path is not None:
np.save(f0_path, f0)
return f0
vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py
-431
View File
@@ -1,431 +0,0 @@
"""SoulX preprocess pipeline."""
from __future__ import annotations
import json
import tempfile
from collections.abc import Iterable
from pathlib import Path
from typing import Any, ClassVar
import numpy as np
import soundfile as sf
import torch
import torch.nn as nn
from vllm.logger import init_logger
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery
from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.stack import SoulXPreprocessStack
from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.utils import load_mono_audio, resample_mono
from vllm_omni.diffusion.models.soulx_singer.preprocess.metadata_utils import (
SegmentMetadata,
_merge_group,
convert_metadata,
)
from vllm_omni.diffusion.models.soulx_singer.utils import (
load_config,
load_wav,
preprocess_weight_paths,
resolve_preprocess_weights_root,
resolve_soulx_kind,
)
logger = init_logger(__name__)
def build_svs_prompt_meta(
preprocess: SoulXPreprocessPipeline,
*,
prompt_audio: str | tuple[np.ndarray, int],
prompt_meta_raw: dict[str, Any],
metadata_processor,
vocal_sep: bool | None = None,
) -> dict[str, Any]:
"""Build SVS prompt_meta; in-memory prompt wav is trimmed to mel2note length."""
if isinstance(prompt_audio, str):
return metadata_processor.process(prompt_meta_raw, prompt_audio)
prompt_meta = metadata_processor.process(prompt_meta_raw, None)
prompt_vocal, prompt_sr = preprocess._extract_vocal(prompt_audio, vocal_sep=vocal_sep)
if prompt_sr != metadata_processor.sample_rate:
prompt_vocal = resample_mono(
prompt_vocal,
orig_sr=prompt_sr,
target_sr=metadata_processor.sample_rate,
)
max_samples = prompt_meta["mel2note"].shape[1] * metadata_processor.hop_size
segment_wav = np.asarray(prompt_vocal[:max_samples], dtype=np.float32)
prompt_meta["wav"] = torch.from_numpy(segment_wav).unsqueeze(0).float().to(metadata_processor.device)
return prompt_meta
class SoulXPreprocessPipeline(nn.Module, SupportAudioInput, SupportsComponentDiscovery):
"""Lazy-loaded preprocess stack integrated with vLLM-Omni diffusion lifecycle."""
support_audio_input: ClassVar[bool] = True
weights_sources: ClassVar[tuple] = ()
_dit_modules: ClassVar[list[str]] = []
_encoder_modules: ClassVar[list[str]] = [
"stack.vocal_sep",
"stack.lyric",
"stack.rosvot",
]
_vae_modules: ClassVar[list[str]] = []
_resident_modules: ClassVar[list[str]] = ["stack.rmvpe"]
def __init__(
self,
*,
od_config: OmniDiffusionConfig,
prefix: str = "",
vocal_sep: bool = True,
midi_transcribe: bool = True,
max_merge_duration_ms: int = 60000,
verbose: bool = False,
extra_args: dict[str, Any] | None = None,
) -> None:
super().__init__()
del prefix
self.od_config = od_config
extra_args = extra_args or {}
self.kind = resolve_soulx_kind(od_config.model)
hf_config = load_config(str(Path(od_config.model) / "config.yaml"))
audio_config = hf_config.audio
if self.kind == "svc" and midi_transcribe:
logger.warning("SVC model (config.json): midi_transcribe is ignored; lyric/MIDI transcription is SVS-only.")
midi_transcribe = False
self.vocal_sep = vocal_sep
self.midi_transcribe = midi_transcribe
self.max_merge_duration_ms = max_merge_duration_ms
self.verbose = verbose
self.target_sr = int(audio_config.sample_rate)
self.hop_size = int(audio_config.hop_size)
self._metadata_processor = None
self._weights_root: Path | None = None
self.stack: SoulXPreprocessStack | None = None
self._ensure_stack(extra_args)
def _ensure_stack(self, extra_args: dict[str, Any] | None = None) -> None:
extra_args = extra_args or {}
weights_root = resolve_preprocess_weights_root(self.od_config)
if self.stack is not None and self._weights_root == weights_root:
return
weights = preprocess_weight_paths(weights_root)
device = str(get_local_device())
self.stack = SoulXPreprocessStack(
weights,
device,
target_sr=self.target_sr,
hop_size=self.hop_size,
verbose=self.verbose,
)
self._weights_root = weights_root
def forward(self, req) -> DiffusionOutput:
from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import (
SOULX_PREPROCESSED_KEY,
)
from vllm_omni.diffusion.models.soulx_singer.preprocess.pre_process import (
build_metadata_processor,
build_preprocess_payload,
build_warmup_payload,
is_warmup_request,
)
from vllm_omni.inputs.data import OmniTextPrompt
extra_args = dict(getattr(req.sampling_params, "extra_args", None) or {})
prompt = OmniTextPrompt(prompt=req.prompts[0]) if isinstance(req.prompts[0], str) else req.prompts[0]
info = prompt.get("additional_information") or {}
if isinstance(info, dict):
extra_args.update(dict(info.get("extra_args") or {}))
self._ensure_stack(extra_args)
if self._metadata_processor is None:
self._metadata_processor = build_metadata_processor(self.od_config)
device = get_local_device()
if is_warmup_request(req):
payload = build_warmup_payload(
self.kind,
metadata_processor=self._metadata_processor,
device=device,
sample_rate=self.target_sr,
)
else:
payload = build_preprocess_payload(
self.kind,
prompt=prompt,
extra_args=extra_args,
preprocess=self,
metadata_processor=self._metadata_processor,
device=device,
sample_rate=self.target_sr,
)
return DiffusionOutput(
output={
"payload": {},
"metadata": {SOULX_PREPROCESSED_KEY: payload},
},
to_cpu=True,
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
del weights
self._ensure_stack({})
return set()
def _extract_vocal(
self,
audio: str | tuple[np.ndarray, int],
*,
vocal_sep: bool | None = None,
) -> tuple[np.ndarray, int]:
if self.stack is None:
self._ensure_stack({})
use_sep = self.vocal_sep if vocal_sep is None else vocal_sep
if use_sep:
return self.stack.extract_vocal(audio)
return load_mono_audio(audio)
def extract_f0(self, vocal: np.ndarray, sample_rate: int) -> np.ndarray:
if self.stack is None:
self._ensure_stack({})
return self.stack.extract_f0(vocal, sample_rate)
def iter_svs_metadata(
self,
audio_source: str | tuple[np.ndarray, int],
*,
language: str = "Mandarin",
vocal_sep: bool | None = None,
max_merge_duration_ms: int | None = None,
):
"""Yield merged SVS metadata dicts one segment at a time."""
if isinstance(audio_source, tuple):
vocal, sample_rate = audio_source
origin_wav_fn = ""
else:
vocal, sample_rate = self._extract_vocal(str(audio_source), vocal_sep=vocal_sep)
origin_wav_fn = str(audio_source)
with tempfile.TemporaryDirectory(prefix="soulx_preprocess_") as work_dir:
work = Path(work_dir)
vocal_path = work / "vocal.wav"
sf.write(vocal_path, vocal, sample_rate)
vocal_f0_path = str(vocal_path).replace(".wav", "_f0.npy")
vocal_f0 = self.stack.extract_f0(str(vocal_path), sample_rate=sample_rate, f0_path=vocal_f0_path)
if not self.midi_transcribe:
end_ms = int(len(vocal) / sample_rate * 1000)
item = SegmentMetadata(
item_name="full",
wav_fn=str(vocal_path),
language=language,
start_time_ms=0,
end_time_ms=end_ms,
note_text=["<SP>"],
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),
}
vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rmvpe.py
-593
View File
@@ -1,593 +0,0 @@
"""RMVPE F0 extractor.
All model code is inlined (originally from RickyL-2000/ROSVOT
``modules/pe/rmvpe/``) with two minor substitutions:
* ``librosa.filters.mel`` → ``torchaudio.functional.melscale_fbanks``
(already used elsewhere in vllm-omni);
* ``librosa.sequence.viterbi`` → hand-written numpy viterbi (non-default path).
* ``pyworld`` (deprecated) → removed.
"""
from __future__ import annotations
import math
from typing import ClassVar
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchaudio.functional import melscale_fbanks
from torchaudio.transforms import Resample
from vllm.logger import init_logger
from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery
logger = init_logger(__name__)
RMVPE_SAMPLE_RATE = 16000
_N_CLASS = 360
_N_MELS = 128
_MEL_FMIN = 30
_MEL_FMAX = 8000
_WINDOW_LENGTH = 1024
_CONST = 1997.3794084376191
class ConvBlockRes(nn.Module):
def __init__(self, in_channels: int, out_channels: int, momentum: float = 0.01):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
nn.BatchNorm2d(out_channels, momentum=momentum),
nn.ReLU(),
nn.Conv2d(out_channels, out_channels, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
nn.BatchNorm2d(out_channels, momentum=momentum),
nn.ReLU(),
)
if in_channels != out_channels:
self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1))
self._has_shortcut = True
else:
self._has_shortcut = False
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self._has_shortcut:
return self.conv(x) + self.shortcut(x)
return self.conv(x) + x
class ResEncoderBlock(nn.Module):
def __init__(self, in_channels: int, out_channels: int, kernel_size, n_blocks: int = 1, momentum: float = 0.01):
super().__init__()
self.n_blocks = n_blocks
self.conv = nn.ModuleList()
self.conv.append(ConvBlockRes(in_channels, out_channels, momentum))
for _ in range(n_blocks - 1):
self.conv.append(ConvBlockRes(out_channels, out_channels, momentum))
self.kernel_size = kernel_size
if self.kernel_size is not None:
self.pool = nn.AvgPool2d(kernel_size=kernel_size)
def forward(self, x: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
for i in range(self.n_blocks):
x = self.conv[i](x)
if self.kernel_size is not None:
return x, self.pool(x)
return x
class ResDecoderBlock(nn.Module):
def __init__(self, in_channels: int, out_channels: int, stride, n_blocks: int = 1, momentum: float = 0.01):
super().__init__()
out_padding = (0, 1) if stride == (1, 2) else (1, 1)
self.n_blocks = n_blocks
self.conv1 = nn.Sequential(
nn.ConvTranspose2d(
in_channels,
out_channels,
kernel_size=(3, 3),
stride=stride,
padding=(1, 1),
output_padding=out_padding,
bias=False,
),
nn.BatchNorm2d(out_channels, momentum=momentum),
nn.ReLU(),
)
self.conv2 = nn.ModuleList()
self.conv2.append(ConvBlockRes(out_channels * 2, out_channels, momentum))
for _ in range(n_blocks - 1):
self.conv2.append(ConvBlockRes(out_channels, out_channels, momentum))
def forward(self, x: torch.Tensor, concat_tensor: torch.Tensor) -> torch.Tensor:
x = self.conv1(x)
x = torch.cat((x, concat_tensor), dim=1)
for i in range(self.n_blocks):
x = self.conv2[i](x)
return x
class Encoder(nn.Module):
def __init__(
self,
in_channels: int,
in_size: int,
n_encoders: int,
kernel_size,
n_blocks: int,
out_channels: int = 16,
momentum: float = 0.01,
):
super().__init__()
self.n_encoders = n_encoders
self.bn = nn.BatchNorm2d(in_channels, momentum=momentum)
self.layers = nn.ModuleList()
self.latent_channels: list[list[int]] = []
for _ in range(self.n_encoders):
self.layers.append(ResEncoderBlock(in_channels, out_channels, kernel_size, n_blocks, momentum=momentum))
self.latent_channels.append([out_channels, in_size])
in_channels = out_channels
out_channels *= 2
in_size //= 2
self.out_size = in_size
self.out_channel = out_channels
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
concat_tensors: list[torch.Tensor] = []
x = self.bn(x)
for i in range(self.n_encoders):
_c, x = self.layers[i](x)
concat_tensors.append(_c)
return x, concat_tensors
class Intermediate(nn.Module):
def __init__(self, in_channels: int, out_channels: int, n_inters: int, n_blocks: int, momentum: float = 0.01):
super().__init__()
self.n_inters = n_inters
self.layers = nn.ModuleList()
self.layers.append(ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum))
for _ in range(n_inters - 1):
self.layers.append(ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum))
def forward(self, x: torch.Tensor) -> torch.Tensor:
for i in range(self.n_inters):
x = self.layers[i](x)
return x
class Decoder(nn.Module):
def __init__(self, in_channels: int, n_decoders: int, stride, n_blocks: int, momentum: float = 0.01):
super().__init__()
self.layers = nn.ModuleList()
self.n_decoders = n_decoders
for _ in range(self.n_decoders):
out_channels = in_channels // 2
self.layers.append(ResDecoderBlock(in_channels, out_channels, stride, n_blocks, momentum))
in_channels = out_channels
def forward(self, x: torch.Tensor, concat_tensors: list[torch.Tensor]) -> torch.Tensor:
for i in range(self.n_decoders):
x = self.layers[i](x, concat_tensors[-1 - i])
return x
class TimbreFilter(nn.Module):
def __init__(self, latent_rep_channels: list[list[int]]):
super().__init__()
self.layers = nn.ModuleList()
for ch in latent_rep_channels:
self.layers.append(ConvBlockRes(ch[0], ch[0]))
def forward(self, x_tensors: list[torch.Tensor]) -> list[torch.Tensor]:
return [layer(t) for layer, t in zip(self.layers, x_tensors)]
class DeepUnet0(nn.Module):
def __init__(
self,
kernel_size,
n_blocks: int,
en_de_layers: int = 5,
inter_layers: int = 4,
in_channels: int = 1,
en_out_channels: int = 16,
):
super().__init__()
self.encoder = Encoder(in_channels, _N_MELS, en_de_layers, kernel_size, n_blocks, en_out_channels)
self.intermediate = Intermediate(
self.encoder.out_channel // 2, self.encoder.out_channel, inter_layers, n_blocks
)
self.tf = TimbreFilter(self.encoder.latent_channels)
self.decoder = Decoder(self.encoder.out_channel, en_de_layers, kernel_size, n_blocks)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x, concat_tensors = self.encoder(x)
x = self.intermediate(x)
x = self.decoder(x, concat_tensors)
return x
class BiGRU(nn.Module):
def __init__(self, input_features: int, hidden_features: int, num_layers: int):
super().__init__()
self.gru = nn.GRU(input_features, hidden_features, num_layers=num_layers, batch_first=True, bidirectional=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.gru(x)[0]
class E2E0(nn.Module):
"""E2E0 pitch estimator network (DeepUnet0 + CNN + BiGRU + classifier)."""
def __init__(
self,
n_blocks: int,
n_gru: int,
kernel_size,
en_de_layers: int = 5,
inter_layers: int = 4,
in_channels: int = 1,
en_out_channels: int = 16,
):
super().__init__()
self.unet = DeepUnet0(kernel_size, n_blocks, en_de_layers, inter_layers, in_channels, en_out_channels)
self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1))
if n_gru:
self.fc = nn.Sequential(
BiGRU(3 * _N_MELS, 256, n_gru),
nn.Linear(512, _N_CLASS),
nn.Dropout(0.25),
nn.Sigmoid(),
)
else:
self.fc = nn.Sequential(
nn.Linear(3 * _N_MELS, _N_CLASS),
nn.Dropout(0.25),
nn.Sigmoid(),
)
def forward(self, mel: torch.Tensor) -> torch.Tensor:
mel = mel.transpose(-1, -2).unsqueeze(1)
x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2)
x = self.fc(x)
return x
class MelSpectrogram(nn.Module):
"""Mel-spectrogram extractor matching ROSVOT's MelSpectrogram.
Log-mel spectrogram via torch.stft + adaptive filterbank.
"""
def __init__(
self,
n_mel_channels: int = _N_MELS,
sampling_rate: int = RMVPE_SAMPLE_RATE,
win_length: int = _WINDOW_LENGTH,
hop_length: int = 160,
n_fft: int | None = None,
mel_fmin: float = _MEL_FMIN,
mel_fmax: float = _MEL_FMAX,
clamp: float = 1e-5,
):
super().__init__()
n_fft = win_length if n_fft is None else n_fft
self.hann_window: dict[str, torch.Tensor] = {}
mel_basis = melscale_fbanks(
n_freqs=n_fft // 2 + 1,
f_min=float(mel_fmin),
f_max=float(mel_fmax) if mel_fmax is not None else 8000.0,
n_mels=n_mel_channels,
sample_rate=int(sampling_rate),
norm="slaney",
mel_scale="htk",
).T # (n_freqs, n_mels) → (n_mels, n_freqs) to match librosa/layernorm convention
self.register_buffer("mel_basis", mel_basis) # (n_mels, n_freqs)
self.n_fft = n_fft
self.hop_length = hop_length
self.win_length = win_length
self.sampling_rate = sampling_rate
self.n_mel_channels = n_mel_channels
self.clamp = clamp
def forward(self, audio: torch.Tensor, keyshift: float = 0, speed: float = 1, center: bool = True) -> torch.Tensor:
factor = 2 ** (keyshift / 12)
n_fft_new = int(np.round(self.n_fft * factor))
win_length_new = int(np.round(self.win_length * factor))
hop_length_new = int(np.round(self.hop_length * speed))
key = str(keyshift) + "_" + str(audio.device)
if key not in self.hann_window:
self.hann_window[key] = torch.hann_window(win_length_new).to(audio.device)
if center:
pad_left = win_length_new // 2
pad_right = (win_length_new + 1) // 2
audio = F.pad(audio, (pad_left, pad_right))
fft = torch.stft(
audio,
n_fft=n_fft_new,
hop_length=hop_length_new,
win_length=win_length_new,
window=self.hann_window[key],
center=False,
return_complex=True,
)
magnitude = fft.abs()
if keyshift != 0:
size = self.n_fft // 2 + 1
if magnitude.size(1) < size:
magnitude = F.pad(magnitude, (0, 0, 0, size - magnitude.size(1)))
magnitude = magnitude[:, :size, :] * self.win_length / win_length_new
# register_buffer confuses type checkers; cast to Tensor explicitly
mel_basis: torch.Tensor = self.mel_basis # type: ignore[assignment]
mel_output = torch.matmul(mel_basis, magnitude)
log_mel = torch.log(torch.clamp(mel_output, min=self.clamp))
return log_mel
def to_local_average_f0(hidden: torch.Tensor, center: torch.Tensor | None = None, thread: float = 0.03) -> np.ndarray:
"""Decode hidden representation to F0 via local weighted average."""
idx = torch.arange(_N_CLASS, device=hidden.device)[None, None, :]
idx_cents = idx * 20 + _CONST
if center is None:
center = torch.argmax(hidden, dim=2, keepdim=True)
start = torch.clip(center - 4, min=0)
end = torch.clip(center + 5, max=_N_CLASS)
idx_mask = (idx >= start) & (idx < end)
weights = hidden * idx_mask
product_sum = torch.sum(weights * idx_cents, dim=2)
weight_sum = torch.sum(weights, dim=2)
cents = product_sum / (weight_sum + (weight_sum == 0))
f0 = 10 * 2 ** (cents / 1200)
uv = hidden.max(dim=2)[0] < thread
f0 = f0 * ~uv
return f0.cpu().numpy()
def numpy_viterbi(prob: np.ndarray, transition: np.ndarray) -> np.ndarray:
"""Viterbi decoding on log-probabilities (numpy).
Args:
prob: Emission probabilities, shape (n_states, n_obs).
transition: Transition matrix, shape (n_states, n_states), row-stochastic.
Returns:
path: Most likely state sequence, shape (n_obs,).
"""
n_states, n_obs = prob.shape
log_emit = np.log(prob + 1e-12)
log_trans = np.log(transition + 1e-12)
delta = log_emit[:, 0] + np.zeros(n_states) # initial: log(pi) omitted (uniform)
psi = np.zeros((n_obs, n_states), dtype=np.int64)
for t in range(1, n_obs):
scores = delta[:, None] + log_trans # (n_states, n_states)
psi[t] = np.argmax(scores, axis=0)
delta = scores[psi[t], np.arange(n_states)] + log_emit[:, t]
# Backtrack
path = np.zeros(n_obs, dtype=np.int64)
path[-1] = np.argmax(delta)
for t in range(n_obs - 2, -1, -1):
path[t] = psi[t + 1, path[t + 1]]
return path
def to_viterbi_f0(hidden: torch.Tensor, thread: float = 0.03) -> np.ndarray:
"""Decode hidden representation to F0 via Viterbi.
Uses a hand-written numpy Viterbi decoder (replacing
``librosa.sequence.viterbi``) to avoid the librosa dependency.
"""
if not hasattr(to_viterbi_f0, "_transition"):
xx, yy = np.meshgrid(range(_N_CLASS), range(_N_CLASS))
transition = np.maximum(30 - abs(xx - yy), 0)
transition = transition / transition.sum(axis=1, keepdims=True)
to_viterbi_f0._transition = transition # type: ignore[attr-defined]
prob = hidden.squeeze(0).cpu().numpy()
prob = prob.T
prob = prob / prob.sum(axis=0, keepdims=True)
path = numpy_viterbi(prob, to_viterbi_f0._transition) # type: ignore[attr-defined]
center = torch.from_numpy(path).unsqueeze(0).unsqueeze(-1).to(hidden.device)
return to_local_average_f0(hidden, center=center, thread=thread)
def resample_align_curve(
points: np.ndarray, original_timestep: float, target_timestep: float, align_length: int = -1
) -> np.ndarray:
"""Align a curve from one time grid to another via linear interpolation."""
t_max = (len(points) - 1) * original_timestep
curve_interp = np.interp(
np.arange(0, t_max, target_timestep),
original_timestep * np.arange(len(points)),
points,
).astype(points.dtype)
if align_length > 0:
if len(curve_interp) > align_length:
curve_interp = curve_interp[:align_length]
elif len(curve_interp) < align_length:
curve_interp = np.concatenate(
[curve_interp, np.full(align_length - len(curve_interp), fill_value=curve_interp[-1])],
axis=0,
)
return curve_interp
class RMVPE(nn.Module, SupportAudioInput, SupportsComponentDiscovery):
"""RMVPE F0 extractor (SoulX-Singer).
Wraps the E2E0 model with mel-spectrogram front-end and decoding
post-processing.
"""
support_audio_input: ClassVar[bool] = True
_dit_modules: ClassVar[list[str]] = []
_encoder_modules: ClassVar[list[str]] = []
_vae_modules: ClassVar[list[str]] = []
_resident_modules: ClassVar[list[str]] = ["."]
_layerwise_offload_blocks_attrs: ClassVar[list[str]] = []
def __init__(self, model_path: str | None = None, hop_length: int = 160, device: str | None = None):
super().__init__()
self.resample_kernel: dict[str, Resample] = {}
if device is None:
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
else:
self.device = torch.device(device)
self.model = E2E0(4, 1, (2, 2)).eval()
if model_path is not None:
ckpt = torch.load(model_path, map_location=self.device)
state_dict = ckpt if not isinstance(ckpt, dict) else ckpt.get("model", ckpt)
self.model.load_state_dict(state_dict, strict=False)
logger.info("Loaded RMVPE weights from %s", model_path)
self.model = self.model.to(self.device)
self.mel_extractor = MelSpectrogram(
_N_MELS,
RMVPE_SAMPLE_RATE,
_WINDOW_LENGTH,
hop_length,
None,
_MEL_FMIN,
_MEL_FMAX,
).to(self.device)
self.hop_length = hop_length
# ---- internal helpers ----
@torch.no_grad()
def mel2hidden(self, mel: torch.Tensor) -> torch.Tensor:
n_frames = mel.shape[-1]
pad_len = 32 * ((n_frames - 1) // 32 + 1) - n_frames
mel = F.pad(mel, (0, pad_len), mode="constant")
hidden = self.model(mel)
return hidden[:, :n_frames]
def decode(self, hidden: torch.Tensor, thread: float = 0.03, use_viterbi: bool = False) -> np.ndarray:
if use_viterbi:
return to_viterbi_f0(hidden, thread=thread)
return to_local_average_f0(hidden, thread=thread)
def postprocess(self, f0: np.ndarray, fmin: float = 50, fmax: float = 1000, min_gap: int = 2) -> np.ndarray:
f0[f0 < fmin] = 0
f0[f0 > fmax] = 0
# Eliminate glitch: if positive f0 values span fewer than min_gap frames
# between zeros, zero them out.
for idx in range(f0.shape[0] - min_gap - 1):
if f0[idx] == 0 and f0[idx + min_gap + 1] == 0 and np.sum(f0[idx : idx + min_gap + 2]) > 0:
f0[idx : idx + min_gap + 2] = 0
return f0
# ---- single audio inference ----
def infer_from_audio(
self, audio: np.ndarray, sample_rate: int = 16000, thread: float = 0.03, use_viterbi: bool = False
) -> np.ndarray:
audio_t = torch.from_numpy(audio).float().unsqueeze(0).to(self.device)
if sample_rate == RMVPE_SAMPLE_RATE:
audio_res = audio_t
else:
key_str = str(sample_rate)
if key_str not in self.resample_kernel:
self.resample_kernel[key_str] = Resample(sample_rate, RMVPE_SAMPLE_RATE, lowpass_filter_width=128)
self.resample_kernel[key_str] = self.resample_kernel[key_str].to(self.device)
audio_res = self.resample_kernel[key_str](audio_t)
mel = self.mel_extractor(audio_res, center=True)
hidden = self.mel2hidden(mel)
f0 = self.decode(hidden, thread=thread, use_viterbi=use_viterbi).squeeze(0)
return f0
def get_pitch(
self,
waveform,
sample_rate: int,
hop_size: int,
length: int,
interp_uv: bool = False,
fmin: float = 50,
fmax: float = 1000,
) -> tuple[np.ndarray, np.ndarray]:
f0 = self.infer_from_audio(waveform, sample_rate=sample_rate)
f0 = self.postprocess(f0, fmin, fmax)
uv = f0 == 0
time_step = hop_size / sample_rate
f0_res = resample_align_curve(f0, 0.01, time_step, length)
uv_res = resample_align_curve(uv.astype(np.float32), 0.01, time_step, length) > 0.5
if not interp_uv:
f0_res[uv_res] = 0
return f0_res, uv_res
# ---- batch audio inference ----
def infer_from_audio_batch(
self, audios, sample_rate: int = 16000, thread: float = 0.03, use_viterbi: bool = False
) -> list[np.ndarray]:
sizes: list[int] | None = None
if isinstance(audios, list):
audios = [torch.from_numpy(a).float() for a in audios]
sizes = [math.ceil((a.shape[0] + 1) / self.hop_length) for a in audios]
audios = nn.utils.rnn.pad_sequence(audios, batch_first=True, padding_value=0.0).to(self.device)
elif isinstance(audios, torch.Tensor):
if audios.device != self.device:
audios = audios.to(self.device)
else:
raise NotImplementedError(f"Unsupported audio type: {type(audios)}")
if sample_rate == RMVPE_SAMPLE_RATE:
audios_res = audios
else:
key_str = str(sample_rate)
if key_str not in self.resample_kernel:
self.resample_kernel[key_str] = Resample(sample_rate, RMVPE_SAMPLE_RATE, lowpass_filter_width=128)
self.resample_kernel[key_str] = self.resample_kernel[key_str].to(self.device)
audios_res = self.resample_kernel[key_str](audios)
mels = self.mel_extractor(audios_res, center=True)
hiddens = self.mel2hidden(mels)
f0 = self.decode(hiddens, thread=thread, use_viterbi=use_viterbi)
f0s: list[np.ndarray] = []
for i in range(f0.shape[0]):
f0s.append(f0[i, : sizes[i]] if sizes is not None else f0[i, :])
return f0s
def get_pitch_batch(
self,
waveforms,
sample_rate: int,
hop_size: int,
lengths: list[int],
interp_uv: bool = False,
fmin: float = 50,
fmax: float = 1000,
) -> tuple[list[np.ndarray], list[np.ndarray]]:
f0s = self.infer_from_audio_batch(waveforms, sample_rate=sample_rate)
f0s_res: list[np.ndarray] = []
uvs_res: list[np.ndarray] = []
for idx, f0 in enumerate(f0s):
f0 = self.postprocess(f0, fmin, fmax, min_gap=6)
uv = f0 == 0
length = lengths[idx]
time_step = hop_size / sample_rate
f0_res = resample_align_curve(f0, 0.01, time_step, length)
uv_res = resample_align_curve(uv.astype(np.float32), 0.01, time_step, length) > 0.5
if not interp_uv:
f0_res[uv_res] = 0
f0s_res.append(f0_res)
uvs_res.append(uv_res)
return f0s_res, uvs_res
def release_cuda(self):
self.model = self.model.cpu()
self.mel_extractor = self.mel_extractor.cpu()
vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rosvot.py
-619
View File
@@ -1,619 +0,0 @@
"""ROSVOT note transcription."""
from __future__ import annotations
import math
from pathlib import Path
from typing import Any, ClassVar
import numpy as np
import torch
import torch.nn as nn
from vllm.logger import init_logger
from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery
from vllm_omni.diffusion.models.soulx_singer.modules.mel_transform import spectral_normalize_torch
from vllm_omni.diffusion.models.soulx_singer.modules.note_transcription import (
BackboneNet,
ConvBlocks,
Embedding,
)
from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.utils import (
boundary2Interval,
denorm_f0,
f0_to_coarse,
load_model_ckpt,
load_mono_audio,
load_rosvot_config,
norm_interp_f0,
pad_or_cut_xd,
resample_mono,
)
from vllm_omni.utils.audio import mel_filter_bank as librosa_mel_fn
logger = init_logger(__name__)
def regulate_real_note_itv(
note_itv: np.ndarray,
note_bd: np.ndarray,
word_bd: np.ndarray,
word_durs: np.ndarray,
hop_size: int,
audio_sample_rate: int,
) -> tuple[np.ndarray, np.ndarray]:
assert note_itv.shape[0] == np.sum(note_bd) + 1
assert np.sum(word_bd) <= np.sum(note_bd)
assert word_durs.shape[0] == np.sum(word_bd) + 1, f"{word_durs.shape[0]} {np.sum(word_bd) + 1}"
word_bd = np.cumsum(word_bd) * word_bd
word_itv = np.zeros((word_durs.shape[0], 2))
word_offsets = np.cumsum(word_durs)
note2words = np.zeros(note_itv.shape[0], dtype=int)
for idx in range(len(word_offsets) - 1):
word_itv[idx, 1] = word_itv[idx + 1, 0] = word_offsets[idx]
word_itv[-1, 1] = word_offsets[-1]
note_itv_secs = note_itv * hop_size / audio_sample_rate
for idx, itv in enumerate(note_itv):
start_idx, end_idx = itv
if word_bd[start_idx] > 0:
word_dur_idx = word_bd[start_idx]
note_itv_secs[idx, 0] = word_itv[word_dur_idx, 0]
note2words[idx] = word_dur_idx
if word_bd[end_idx] > 0:
word_dur_idx = word_bd[end_idx] - 1
note_itv_secs[idx, 1] = word_itv[word_dur_idx, 1]
note2words[idx] = word_dur_idx
note2words += 1
return note_itv_secs, note2words
def regulate_ill_slur(
notes: np.ndarray, note_itv: np.ndarray, note2words: np.ndarray
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
res_note2words: list[int] = []
res_note_itv: list[list[float]] = []
res_notes: list[int] = []
note_idx = 0
note_idx_end = 0
while note_idx <= len(notes) - 1:
while note_idx <= note_idx_end < len(notes) and note2words[note_idx] == note2words[note_idx_end]:
note_idx_end += 1
res_note2words.append(note2words[note_idx])
res_note_itv.append(note_itv[note_idx].tolist())
res_notes.append(notes[note_idx])
for idx in range(note_idx + 1, note_idx_end):
if notes[idx] == notes[idx - 1]:
res_note_itv[-1][1] = note_itv[idx][1]
else:
res_note_itv.append(note_itv[idx].tolist())
res_note2words.append(note2words[idx])
res_notes.append(notes[idx])
note_idx = note_idx_end
return (
np.array(res_notes, dtype=notes.dtype),
np.array(res_note_itv, dtype=note_itv.dtype),
np.array(res_note2words, dtype=note2words.dtype),
)
def regulate_boundary(bd_logits, threshold, min_gap=18, ref_bd=None, ref_bd_min_gap=8, non_padding=None):
# this doesn't preserve gradient
device = bd_logits.device
bd_logits = torch.sigmoid(bd_logits).data.cpu()
# bd_logits[0] = bd_logits[-1] = 1e-5 # avoid itv invalid problem
bd = (bd_logits > threshold).long()
bd_res = torch.zeros_like(bd).long()
for i in range(bd.shape[0]):
bd_i = bd[i]
last_bd_idx = -1
start = -1
for j in range(bd_i.shape[0]):
if bd_i[j] == 1:
if 0 <= start < j:
continue
elif start < 0:
start = j
else:
if 0 <= start < j:
if j - 1 > start:
bd_idx = start + int(torch.argmax(bd_logits[i, start:j]).item())
else:
bd_idx = start
if bd_idx - last_bd_idx < min_gap and last_bd_idx > 0:
bd_idx = round((bd_idx + last_bd_idx) / 2)
bd_res[i, last_bd_idx] = 0
bd_res[i, bd_idx] = 1
last_bd_idx = bd_idx
start = -1
# assert ref_bd_min_gap <= min_gap // 2
if ref_bd is not None and ref_bd_min_gap > 0:
ref = ref_bd.data.cpu()
for i in range(bd_res.shape[0]):
ref_bd_i = ref[i]
ref_bd_i_js = []
for j in range(ref_bd_i.shape[0]):
if ref_bd_i[j] == 1:
ref_bd_i_js.append(j)
seg_sum = torch.sum(bd_res[i, max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap])
if seg_sum == 0:
bd_res[i, j] = 1
elif seg_sum == 1 and bd_res[i, j] != 1:
bd_res[i, max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap] = ref_bd_i[
max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap
]
elif seg_sum > 1:
for k in range(1, ref_bd_min_gap + 1):
if bd_res[i, max(0, j - k)] == 1 and ref_bd_i[max(0, j - k)] != 1:
bd_res[i, max(0, j - k)] = 0
break
if (
bd_res[i, min(bd_res.shape[1] - 1, j + k)] == 1
and ref_bd_i[min(bd_res.shape[1] - 1, j + k)] != 1
):
bd_res[i, min(bd_res.shape[1] - 1, j + k)] = 0
break
bd_res[i, j] = 1
# final check
assert torch.sum(bd_res[i, ref_bd_i_js]) == len(ref_bd_i_js), (
f"{torch.sum(bd_res[i, ref_bd_i_js])} {len(ref_bd_i_js)}"
)
bd_res = bd_res.to(device)
# force valid begin and end
bd_res[:, 0] = 0
if non_padding is not None:
for i in range(bd_res.shape[0]):
bd_res[i, sum(non_padding[i]) - 1 :] = 0
else:
bd_res[:, -1] = 0
return bd_res
def _align_word(word_durs: list[float], mel_len: int, hop_size: int, audio_sample_rate: int) -> np.ndarray:
mel2word = np.zeros([mel_len], int)
start_time = 0.0
for i_word, wd in enumerate(word_durs):
start_frame = int(start_time * audio_sample_rate / hop_size + 0.5)
end_frame = int((start_time + wd) * audio_sample_rate / hop_size + 0.5)
mel2word[start_frame:end_frame] = i_word + 1
start_time += wd
return mel2word
class PitchDecoder(nn.Module):
def __init__(self, hparams):
super().__init__()
self.hidden_size = hidden_size = hparams["hidden_size"]
self.dropout = hparams.get("dropout", 0.0)
self.note_bd_out = nn.Linear(hidden_size, 1)
self.note_bd_temperature = max(1e-7, hparams.get("note_bd_temperature", 1.0))
# note prediction
self.pitch_attn_num_head = hparams.get("pitch_attn_num_head", 1)
self.multihead_dot_attn = nn.Linear(hidden_size, self.pitch_attn_num_head)
self.post = ConvBlocks(
hidden_size,
out_dims=hidden_size,
dilations=None,
kernel_size=3,
layers_in_block=1,
c_multiple=1,
post_net_kernel=3,
)
self.pitch_out = nn.Linear(hidden_size, hparams.get("note_num", 100) + 4)
self.note_num = hparams.get("note_num", 100)
self.note_start = hparams.get("note_start", 30)
self.pitch_temperature = max(1e-7, hparams.get("note_pitch_temperature", 1.0))
def forward(self, feat, note_bd, train=True):
bsz, T, _ = feat.shape
attn = torch.sigmoid(self.multihead_dot_attn(feat)) # [B, T, C] -> [B, T, num_head]
attn_feat = feat.unsqueeze(3) * attn.unsqueeze(2) # [B, T, C, 1] x [B, T, 1, num_head] -> [B, T, C, num_head]
attn_feat = torch.mean(attn_feat, dim=-1) # [B, T, C, num_head] -> [B, T, C]
mel2note = torch.cumsum(note_bd, 1)
note_length = torch.max(torch.sum(note_bd, dim=1)).item() + 1 # max length
note_lengths = torch.sum(note_bd, dim=1) + 1 # [B]
attn = torch.mean(attn, dim=-1, keepdim=True) # [B, T, num_head] -> [B, T, 1]
denom = mel2note.new_zeros(bsz, note_length, dtype=attn.dtype).scatter_add_(
dim=1, index=mel2note, src=attn.squeeze(-1)
) # [B, T] -> [B, note_length] count the note frames of each note (with padding excluded)
frame2note = mel2note.unsqueeze(-1).repeat(1, 1, self.hidden_size) # [B, T] -> [B, T, C], with padding included
note_aggregate = frame2note.new_zeros(bsz, note_length, self.hidden_size, dtype=attn_feat.dtype).scatter_add_(
dim=1, index=frame2note, src=attn_feat
) # [B, T, C] -> [B, note_length, C]
note_aggregate = note_aggregate / (denom.unsqueeze(-1) + 1e-5)
note_logits = self.post(note_aggregate)
note_logits = self.pitch_out(note_logits) / self.pitch_temperature
# note_logits = torch.clamp(note_logits, min=-16., max=16.) # don't know need it or not
note_pred = torch.softmax(note_logits, dim=-1) # [B, note_length, note_num]
note_pred = torch.argmax(note_pred, dim=-1) # [B, note_length]
# for some reason, note idx maybe 130 (why?)
note_pred[note_pred > self.note_num] = 0
note_pred[note_pred < self.note_start] = 0
return note_lengths, note_logits, note_pred
class MidiExtractor(nn.Module):
def __init__(self, hparams):
super().__init__()
self.hparams = hparams.copy()
self.hidden_size = hidden_size = hparams["hidden_size"]
self.note_bd_threshold = hparams.get("note_bd_threshold", 0.5)
self.note_bd_min_gap = round(
hparams.get("note_bd_min_gap", 100) * hparams["audio_sample_rate"] / 1000 / hparams["hop_size"]
)
self.note_bd_ref_min_gap = round(
hparams.get("note_bd_ref_min_gap", 50) * hparams["audio_sample_rate"] / 1000 / hparams["hop_size"]
)
self.mel_proj = nn.Conv1d(hparams["use_mel_bins"], hidden_size, kernel_size=3, padding=1)
self.mel_encoder = ConvBlocks(
hidden_size,
out_dims=hidden_size,
dilations=None,
kernel_size=3,
layers_in_block=2,
c_multiple=1,
post_net_kernel=3,
)
self.use_pitch = hparams.get("use_pitch_embed", True)
if self.use_pitch:
self.pitch_embed = Embedding(300, hidden_size, 0, "kaiming")
self.uv_embed = Embedding(3, hidden_size, 0, "kaiming")
self.use_wbd = hparams.get("use_wbd", True)
if self.use_wbd:
self.word_bd_embed = Embedding(3, hidden_size, 0, "kaiming")
self.cond_encoder = ConvBlocks(
hidden_size,
out_dims=hidden_size,
dilations=None,
kernel_size=3,
layers_in_block=1,
c_multiple=1,
post_net_kernel=3,
)
# backbone
self.net = BackboneNet(hparams)
# note bd prediction
self.note_bd_out = nn.Linear(hidden_size, 1)
self.note_bd_temperature = max(1e-7, hparams.get("note_bd_temperature", 1.0))
# note prediction
self.pitch_decoder = PitchDecoder(hparams)
def run_encoder(self, mel, word_bd=None, pitch=None, uv=None):
mel_embed = self.mel_proj(mel.transpose(1, 2)).transpose(1, 2)
mel_embed = self.mel_encoder(mel_embed)
pitch_embed = word_bd_embed = 0
if self.use_pitch and pitch is not None and uv is not None:
pitch_embed = self.pitch_embed(pitch) + self.uv_embed(uv) # [B, T, C]
if self.use_wbd and word_bd is not None:
word_bd_embed = self.word_bd_embed(word_bd)
feat = self.cond_encoder(mel_embed + pitch_embed + word_bd_embed)
return feat
def forward(self, mel, word_bd=None, note_bd=None, pitch=None, uv=None, non_padding=None, train=True):
ret = {}
bsz, T, _ = mel.shape
feat = self.run_encoder(mel, word_bd, pitch, uv)
feat = self.net(feat) # [B, T, C]
# note bd prediction
# dropout has been dropped (inference mode)
note_bd_logits = self.note_bd_out(feat).squeeze(-1) / self.note_bd_temperature
note_bd_logits = torch.clamp(note_bd_logits, min=-16.0, max=16.0)
ret["note_bd_logits"] = note_bd_logits # [B, T]
if note_bd is None or not train:
note_bd = regulate_boundary(
note_bd_logits,
self.note_bd_threshold,
self.note_bd_min_gap,
word_bd,
self.note_bd_ref_min_gap,
non_padding,
)
ret["note_bd_pred"] = note_bd # [B, T]
# note pitch prediction
note_lengths, note_logits, note_pred = self.pitch_decoder(feat, note_bd, train)
ret["note_lengths"], ret["note_logits"], ret["note_pred"] = note_lengths, note_logits, note_pred
return ret
class MelNet(nn.Module):
def __init__(self, hparams, device="cpu") -> None:
super().__init__()
self.n_fft = hparams["fft_size"]
self.num_mels = hparams["audio_num_mel_bins"]
self.sampling_rate = hparams["audio_sample_rate"]
self.hop_size = hparams["hop_size"]
self.win_size = hparams["win_size"]
self.fmin = hparams["fmin"]
self.fmax = hparams["fmax"]
self.device = device
mel = librosa_mel_fn(
sr=self.sampling_rate,
n_fft=self.n_fft,
n_mels=self.num_mels,
fmin=self.fmin,
fmax=self.fmax,
)
self.mel_basis = mel.float().to(self.device)
self.hann_window = torch.hann_window(self.win_size).to(self.device)
def to(self, device, **kwagrs):
super().to(device=device, **kwagrs)
self.mel_basis = self.mel_basis.to(device)
self.hann_window = self.hann_window.to(device)
self.device = device
def forward(self, y, center=False, complex=False):
if isinstance(y, np.ndarray):
y = torch.FloatTensor(y)
if len(y.shape) == 1:
y = y.unsqueeze(0)
y = y.clamp(min=-1.0, max=1.0).to(self.device)
pad_length = math.ceil(y.shape[1] / self.hop_size) * self.hop_size - y.shape[1]
y = torch.nn.functional.pad(
y.unsqueeze(1),
[int((self.n_fft - self.hop_size) / 2), int((self.n_fft - self.hop_size) / 2 + pad_length)],
mode="reflect",
)
y = y.squeeze(1)
spec = torch.stft(
y,
self.n_fft,
hop_length=self.hop_size,
win_length=self.win_size,
window=self.hann_window,
center=center,
pad_mode="reflect",
normalized=False,
onesided=True,
return_complex=True,
)
if not complex:
spec = torch.view_as_real(spec)
spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) # [B, n_fft, T]
spec = torch.matmul(self.mel_basis, spec)
spec = spectral_normalize_torch(spec)
spec = spec.transpose(1, 2) # [B, T, n_fft]
else:
B, C, T, _ = spec.shape
spec = spec.transpose(1, 2) # [B, T, n_fft, 2]
return spec
class RosvotModel(nn.Module, SupportAudioInput, SupportsComponentDiscovery):
support_audio_input: ClassVar[bool] = True
_dit_modules: ClassVar[list[str]] = []
_encoder_modules: ClassVar[list[str]] = ["."]
_vae_modules: ClassVar[list[str]] = []
_resident_modules: ClassVar[list[str]] = []
_layerwise_offload_blocks_attrs: ClassVar[list[str]] = ["midi.net", "midi.pitch_decoder", "midi.cond_encoder"]
def __init__(
self,
rosvot_ckpt: str | Path,
*,
config_path: str | Path = "",
pe: nn.Module | None = None,
the: float = 0.85,
verbose: bool = False,
rosvot_source_dir: str | Path | None = None,
):
super().__init__()
self.verbose = verbose
ckpt = Path(rosvot_ckpt)
resolved_config = Path(config_path) if config_path else ckpt.with_name("config.yaml")
self.hparams = load_rosvot_config(resolved_config, hparams_str=f"note_bd_threshold={the}")
if verbose:
logger.info("ROSVOT config: %s", resolved_config)
self.midi = MidiExtractor(self.hparams)
self.mel_net = MelNet(self.hparams)
self.pe = pe if pe is not None and self.hparams.get("use_pitch_embed", False) else None
self._checkpoint_path = str(ckpt)
self.load_checkpoint(str(ckpt), verbose=verbose)
self.eval()
def load_checkpoint(self, checkpoint_path: str | None = None, *, verbose: bool = False) -> None:
load_model_ckpt(self.midi, checkpoint_path or self._checkpoint_path, verbose=verbose)
@torch.no_grad()
def forward(self, wav: torch.Tensor, word_durs: list[float]) -> dict[str, Any]:
hparams = self.hparams
if wav.ndim == 1:
wav = wav.unsqueeze(0)
mel_len = (wav.shape[-1] + hparams["hop_size"] - 1) // hparams["hop_size"]
min_word_dur = hparams.get("min_word_dur", 20) / 1000
wd_raw = list(word_durs)
word_durs_filtered: list[float] = []
for i, wd in enumerate(wd_raw):
if wd < min_word_dur:
if i == 0 and len(wd_raw) > 1:
wd_raw[i + 1] += wd
elif word_durs_filtered:
word_durs_filtered[-1] += wd
else:
word_durs_filtered.append(wd)
mel2word = _align_word(word_durs_filtered, mel_len, hparams["hop_size"], hparams["audio_sample_rate"])
if mel2word.size > 0 and mel2word[0] == 0:
mel2word = mel2word + 1
real_len = min(mel_len, int(np.sum(mel2word > 0)))
T = math.ceil(min(real_len, hparams["max_frames"]) / hparams["frames_multiple"]) * hparams["frames_multiple"]
device = wav.device
self.mel_net.to(device)
wav_t = pad_or_cut_xd(wav.float(), T * hparams["hop_size"], 1)
pitch_coarse = uv_t = None
if self.pe is not None:
f0s, uvs = self.pe.get_pitch_batch(
wav_t,
sample_rate=hparams["audio_sample_rate"],
hop_size=hparams["hop_size"],
lengths=[real_len],
fmax=hparams["f0_max"],
fmin=hparams["f0_min"],
)
f0_1d, uv_1d = norm_interp_f0(f0s[0][:T])
f0_t = pad_or_cut_xd(torch.as_tensor(f0_1d, device=device, dtype=torch.float32), T, 0).unsqueeze(0)
uv_t = pad_or_cut_xd(torch.as_tensor(uv_1d, device=device, dtype=torch.float32), T, 0).long().unsqueeze(0)
pitch_coarse = f0_to_coarse(denorm_f0(f0_t, uv_t)).to(device)
mel = pad_or_cut_xd(self.mel_net(wav_t)[0], T, dim=0).unsqueeze(0)
mel_nonpadding_mask = torch.zeros(1, T, device=device)
mel_nonpadding_mask[:, :real_len] = 1.0
mel = (mel.transpose(1, 2) * mel_nonpadding_mask.unsqueeze(1)).transpose(1, 2)
mel_nonpadding = mel.abs().sum(-1) > 0
mel2word_t = pad_or_cut_xd(torch.as_tensor(mel2word, device=device, dtype=torch.long), T, 0)
word_bd = torch.zeros_like(mel2word_t)
word_bd[1:] = (mel2word_t[1:] != mel2word_t[:-1]).long()
word_bd[real_len:] = 0
word_bd = word_bd.unsqueeze(0)
outputs = self.midi(
mel=mel[:, :, : hparams.get("use_mel_bins", 80)],
word_bd=word_bd,
pitch=pitch_coarse,
uv=uv_t,
non_padding=mel_nonpadding,
)
outputs["word_durs_filtered"] = word_durs_filtered
outputs["real_len"] = real_len
outputs["word_bd"] = word_bd
return outputs
@staticmethod
def _load_wav(wav_src: str | np.ndarray, sample_rate: int, *, src_sample_rate: int | None = None) -> np.ndarray:
if isinstance(wav_src, str):
wav, _ = load_mono_audio(wav_src, target_sr=sample_rate)
return wav
wav = np.asarray(wav_src, dtype=np.float32)
if src_sample_rate is not None and src_sample_rate != sample_rate:
wav = resample_mono(wav, orig_sr=src_sample_rate, target_sr=sample_rate)
return wav
@staticmethod
def _normalize_note2words(note2words: list[int]) -> list[int]:
if not note2words:
return []
out = [note2words[0]]
for idx in range(1, len(note2words)):
out.append(max(note2words[idx], out[-1]))
return out
@staticmethod
def _build_ep_types(note2words: list[int], align_words: list[str]) -> list[int]:
ep_types: list[int] = []
prev = -1
for i, w in zip(note2words, align_words):
ep_types.append(1 if w == "<SP>" 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", []),
}
vllm_omni/diffusion/models/soulx_singer/modules/preprocess/segmenter.py
-267
View File
@@ -1,267 +0,0 @@
"""Rule-based vocal segmentation from F0 contour (no neural net)."""
from __future__ import annotations
import time
from dataclasses import dataclass
from typing import ClassVar
import numpy as np
import torch.nn as nn
from vllm.logger import init_logger
from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery
logger = init_logger(__name__)
@dataclass(frozen=True)
class VocalSegmentConfig:
hop_ms: int = 20
smooth_ms: int = 200
start_ms: int = 120
end_ms: int = 200
prepad_ms: int = 80
postpad_ms: int = 120
min_len_ms: int = 1000
max_len_ms: int = 20000
short_seg_merge_gap_ms: int = 8000
small_gap_ms: int = 500
lookback_ms: int = 200
lookahead_ms: int = 100
def _moving_average(x: np.ndarray, win: int) -> np.ndarray:
if win <= 1:
return x
kernel = np.ones(win, dtype=np.float32) / float(win)
return np.convolve(x, kernel, mode="same")
def _merge_short_segments(
segments_ms: list[list[int]],
*,
min_len_ms: int,
max_len_ms: int,
short_seg_merge_gap_ms: int,
small_gap_ms: int,
) -> list[list[int]]:
if not segments_ms:
return []
merged: list[list[int]] = []
cur_start, cur_end = segments_ms[0]
for next_start, next_end in segments_ms[1:]:
cur_len = cur_end - cur_start
gap_ms = next_start - cur_end
merged_len = next_end - cur_start
should_merge = (cur_len < min_len_ms and gap_ms < short_seg_merge_gap_ms) or (
gap_ms < small_gap_ms and merged_len < max_len_ms
)
if should_merge:
cur_end = next_end
continue
if (cur_end - cur_start) >= min_len_ms:
merged.append([cur_start, cur_end])
cur_start, cur_end = next_start, next_end
if (cur_end - cur_start) >= min_len_ms:
merged.append([cur_start, cur_end])
if not merged:
return segments_ms
return merged
def _voiced_to_segments(
voiced: np.ndarray,
*,
hop_ms: int,
smooth_ms: int,
start_ms: int,
end_ms: int,
prepad_ms: int,
postpad_ms: int,
max_len_ms: int,
) -> list[list[int]]:
smooth_frames = max(1, int(round(smooth_ms / hop_ms)))
active = _moving_average(voiced.astype(np.float32), smooth_frames) >= 0.5
segments: list[list[int]] = []
start_idx = None
start_frames = max(1, int(round(start_ms / hop_ms)))
end_frames = max(1, int(round(end_ms / hop_ms)))
prepad_frames = max(0, int(round(prepad_ms / hop_ms)))
postpad_frames = max(0, int(round(postpad_ms / hop_ms)))
active_count = 0
inactive_count = 0
for i, flag in enumerate(active):
if flag:
active_count += 1
inactive_count = 0
else:
inactive_count += 1
active_count = 0
if start_idx is None:
if active_count >= start_frames:
start_idx = max(0, i - start_frames + 1 - prepad_frames)
else:
if inactive_count >= end_frames:
end_idx = min(len(active) - 1, i - end_frames + 1 + postpad_frames)
start_ms_val = start_idx * hop_ms
end_ms_val = end_idx * hop_ms + hop_ms
if end_ms_val > start_ms_val:
segments.append([int(start_ms_val), int(end_ms_val)])
start_idx = None
if start_idx is not None:
start_ms_val = start_idx * hop_ms
end_idx = min(len(active) - 1, len(active) - 1 + postpad_frames)
end_ms_val = end_idx * hop_ms + hop_ms
if end_ms_val > start_ms_val:
segments.append([int(start_ms_val), int(end_ms_val)])
def _split_segment(seg: list[int]) -> list[list[int]]:
start_ms_val, end_ms_val = seg
start_frame = int(start_ms_val // hop_ms)
end_frame = int((end_ms_val - 1) // hop_ms)
end_frame = max(start_frame, min(end_frame, len(active) - 1))
best_start = None
best_len = 0
cur_start = None
cur_len = 0
for idx in range(start_frame, end_frame + 1):
if not active[idx]:
if cur_start is None:
cur_start = idx
cur_len = 1
else:
cur_len += 1
else:
if cur_start is not None and cur_len > best_len:
best_start, best_len = cur_start, cur_len
cur_start = None
cur_len = 0
if cur_start is not None and cur_len > best_len:
best_start, best_len = cur_start, cur_len
if best_start is None:
split_frame = (start_frame + end_frame) // 2
else:
split_frame = best_start + best_len // 2
split_ms = split_frame * hop_ms
if split_ms <= start_ms_val:
split_ms = start_ms_val + hop_ms
if split_ms >= end_ms_val:
split_ms = end_ms_val - hop_ms
if split_ms <= start_ms_val or split_ms >= end_ms_val:
return [seg]
return [[start_ms_val, int(split_ms)], [int(split_ms), end_ms_val]]
queue = segments[:]
segments = []
while queue:
seg = queue.pop(0)
if (seg[1] - seg[0]) <= max_len_ms:
segments.append(seg)
continue
parts = _split_segment(seg)
if len(parts) == 1:
segments.append(seg)
else:
queue = parts + queue
return segments
class VocalSegmenter(nn.Module, SupportAudioInput, SupportsComponentDiscovery):
support_audio_input: ClassVar[bool] = True
_dit_modules: ClassVar[list[str]] = []
_encoder_modules: ClassVar[list[str]] = []
_vae_modules: ClassVar[list[str]] = []
_resident_modules: ClassVar[list[str]] = ["."]
_layerwise_offload_blocks_attrs: ClassVar[list[str]] = []
def __init__(self, config: VocalSegmentConfig | None = None, *, verbose: bool = False):
super().__init__()
self.config = config or VocalSegmentConfig()
self.verbose = verbose
def forward(
self,
audio: np.ndarray,
sample_rate: int,
f0: np.ndarray,
*,
base_name: str = "vocal",
origin_wav_fn: str = "",
verbose: bool | None = None,
) -> list[dict]:
verbose = self.verbose if verbose is None else verbose
if verbose:
logger.info("[vocal detection] segmenting %s", base_name)
t0 = time.time()
cfg = self.config
y = np.asarray(audio, dtype=np.float32).squeeze()
sr = int(sample_rate)
segments_ms = _voiced_to_segments(
f0 > 0,
hop_ms=cfg.hop_ms,
smooth_ms=cfg.smooth_ms,
start_ms=cfg.start_ms,
end_ms=cfg.end_ms,
prepad_ms=cfg.prepad_ms,
postpad_ms=cfg.postpad_ms,
max_len_ms=cfg.max_len_ms,
)
segments_ms = _merge_short_segments(
segments_ms,
min_len_ms=cfg.min_len_ms,
max_len_ms=cfg.max_len_ms,
short_seg_merge_gap_ms=cfg.short_seg_merge_gap_ms,
small_gap_ms=cfg.small_gap_ms,
)
adjusted_segments: list[list[int]] = []
prev_end = 0
for start_ms, end_ms in segments_ms:
start_ms = max(0, start_ms - cfg.lookback_ms)
end_ms = min(end_ms + cfg.lookahead_ms, int(y.shape[0] / sr * 1000))
if start_ms < prev_end and adjusted_segments:
adjusted_segments[-1][1] = start_ms
adjusted_segments.append([start_ms, end_ms])
prev_end = end_ms
segment_infos = []
for idx, (start_ms, end_ms) in enumerate(adjusted_segments):
if end_ms - start_ms > cfg.max_len_ms:
start_ms = end_ms - cfg.max_len_ms
key = f"{base_name}_{idx}"
start_sample = int(start_ms * sr / 1000)
end_sample = int(end_ms * sr / 1000)
segment_infos.append(
{
"item_name": key,
"wav": y[start_sample:end_sample],
"sample_rate": sr,
"wav_fn": key,
"start_time_ms": int(start_sample * 1000 / sr),
"end_time_ms": int(end_sample * 1000 / sr),
"origin_wav_fn": origin_wav_fn,
"duration": int((end_sample - start_sample) * 1000 / sr),
}
)
if verbose:
logger.info(
"[vocal detection] done n_segments=%s time=%.3fs",
len(segment_infos),
time.time() - t0,
)
return segment_infos
vllm_omni/diffusion/models/soulx_singer/modules/preprocess/stack.py
-145
View File
@@ -1,145 +0,0 @@
"""Lazy-loaded SoulX preprocess model tree."""
import os
import tempfile
from pathlib import Path
from typing import ClassVar
import numpy as np
import soundfile as sf
import torch
import torch.nn as nn
from vllm.logger import init_logger
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from .asr import LyricModel
from .mel_grid_f0 import extract_f0_file
from .rmvpe import RMVPE
from .rosvot import RosvotModel
from .segmenter import VocalSegmenter
from .utils import resample_mono
from .vocal_sep import MelBandRoformerSeparator
logger = init_logger(__name__)
class SoulXPreprocessStack(nn.Module, SupportsComponentDiscovery):
"""Unified nn.Module tree for SoulX-Singer audio preprocess."""
_dit_modules: ClassVar[list[str]] = []
_encoder_modules: ClassVar[list[str]] = ["vocal_sep", "lyric", "rosvot"]
_vae_modules: ClassVar[list[str]] = []
_resident_modules: ClassVar[list[str]] = ["rmvpe"]
def __init__(
self,
weights: dict[str, str],
device: str,
*,
target_sr: int = 24000,
hop_size: int = 480,
verbose: bool = False,
) -> None:
super().__init__()
self.weights = weights
self.device_str = device
self.target_sr = target_sr
self.hop_size = hop_size
self.verbose = verbose
self._rmvpe: nn.Module | None = None
self._vocal_sep: nn.Module | None = None
self._segmenter: nn.Module | None = None
self._lyric: nn.Module | None = None
self._rosvot: nn.Module | None = None
@property
def device(self) -> torch.device:
return torch.device(self.device_str)
def ensure_rmvpe(self):
if self._rmvpe is None:
if self.verbose:
logger.info("[rmvpe] loading shared pitch model")
self._rmvpe = RMVPE(self.weights["rmvpe"], device=self.device_str)
self.add_module("rmvpe", self._rmvpe)
return self._rmvpe
def ensure_vocal_sep(self):
if self._vocal_sep is None:
if self.verbose:
logger.info("[vocal extraction] loading model")
self._vocal_sep = MelBandRoformerSeparator(
self.weights["sep_config"],
self.weights["sep_ckpt"],
use_half=False,
).to(self.device)
self.add_module("vocal_sep", self._vocal_sep)
if self.verbose:
logger.info("[vocal extraction] ready on %s, sr=%s", self.device_str, self._vocal_sep.sample_rate)
return self._vocal_sep
def ensure_segmenter(self):
if self._segmenter is None:
self._segmenter = VocalSegmenter(verbose=self.verbose)
self.add_module("segmenter", self._segmenter)
return self._segmenter
def ensure_lyric(self):
if self._lyric is None:
if self.verbose:
logger.info("[lyric transcription] loading ASR")
self._lyric = LyricModel(
self.weights["asr_zh"],
self.weights["asr_en"],
device=self.device_str,
target_sr=self.target_sr,
hop_size=self.hop_size,
)
self.add_module("lyric", self._lyric)
return self._lyric
def ensure_rosvot(self):
if self._rosvot is None:
if self.verbose:
logger.info("[note transcription] loading ROSVOT")
rosvot_src = os.environ.get("ROSVOT_SOURCE_DIR")
self._rosvot = RosvotModel(
self.weights["rosvot"],
pe=self.ensure_rmvpe(),
verbose=self.verbose,
rosvot_source_dir=rosvot_src,
).to(self.device)
self.add_module("rosvot", self._rosvot)
return self._rosvot
def extract_f0(self, vocal, sample_rate: int, *, f0_path: str | None = None) -> np.ndarray:
"""Extract F0 on the mel grid; ndarray inputs are spooled to a temp wav."""
rmvpe = self.ensure_rmvpe()
kwargs = dict(
target_sr=self.target_sr,
hop_size=self.hop_size,
f0_path=f0_path,
verbose=self.verbose,
)
if isinstance(vocal, str):
return extract_f0_file(rmvpe, vocal, **kwargs)
tmp_path: str | None = None
try:
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
tmp_path = tmp.name
sf.write(tmp_path, np.asarray(vocal, dtype=np.float32).squeeze(), int(sample_rate))
return extract_f0_file(rmvpe, tmp_path, **kwargs)
finally:
if tmp_path is not None:
Path(tmp_path).unlink(missing_ok=True)
def extract_vocal(self, audio: str | tuple[np.ndarray, int]) -> tuple[np.ndarray, int]:
sep = self.ensure_vocal_sep()
if isinstance(audio, tuple):
mix, sr = audio
mix = np.asarray(mix, dtype=np.float32).squeeze()
if sr != sep.sample_rate:
mix = resample_mono(mix, orig_sr=sr, target_sr=sep.sample_rate)
return sep.separate_mono(mix, self.device), sep.sample_rate
return sep.separate_path(str(audio), self.device)
vllm_omni/diffusion/models/soulx_singer/modules/preprocess/utils.py
-260
View File
@@ -1,260 +0,0 @@
"""Shared helpers for SoulX preprocess (I/O, config, checkpoints, pitch)."""
from pathlib import Path
from typing import Any
import numpy as np
import soundfile as sf
import torch
import torch.nn as nn
import torch.nn.functional as F
from omegaconf import OmegaConf
from vllm.logger import init_logger
from vllm.multimodal.audio import AudioResampler
from vllm.multimodal.media.audio import load_audio
logger = init_logger(__name__)
# --- audio I/O ---
def resample_mono(
audio: np.ndarray,
*,
orig_sr: int,
target_sr: int,
) -> np.ndarray:
"""Resample mono float waveform to ``target_sr``."""
audio = np.asarray(audio, dtype=np.float32).squeeze()
if orig_sr == target_sr:
return audio
resampler = AudioResampler(target_sr=float(target_sr))
return np.asarray(
resampler.resample(audio, orig_sr=float(orig_sr)),
dtype=np.float32,
)
def get_audio_duration(path: str) -> float:
"""Return audio duration in seconds for a file on disk."""
return float(sf.info(path).duration)
def load_mono_audio(
source: str | tuple[np.ndarray, int],
*,
target_sr: int | None = None,
) -> tuple[np.ndarray, int]:
if isinstance(source, tuple):
audio, sr = source
audio = np.asarray(audio, dtype=np.float32).squeeze()
if target_sr is not None and sr != target_sr:
audio = resample_mono(audio, orig_sr=sr, target_sr=target_sr)
sr = target_sr
return audio, sr
if target_sr is None:
audio, sr = load_audio(str(source), sr=None, mono=True)
else:
audio, sr = load_audio(str(source), sr=target_sr, mono=True)
return np.asarray(audio, dtype=np.float32).squeeze(), int(sr)
# --- tensor padding ---
def pad_or_cut_xd(values: torch.Tensor, tgt_len: int, dim: int = -1, pad_value=0) -> torch.Tensor:
ndim = values.ndim
if ndim == 1:
if values.shape[0] < tgt_len:
return F.pad(values, [0, tgt_len - values.shape[0]], value=pad_value)
return values[:tgt_len]
if ndim == 2:
if dim in (0, -2):
src_len = values.shape[0]
if src_len < tgt_len:
return F.pad(values, [0, 0, 0, tgt_len - src_len], value=pad_value)
return values[:tgt_len]
src_len = values.shape[1]
if src_len < tgt_len:
return F.pad(values, [0, tgt_len - src_len], value=pad_value)
return values[:, :tgt_len]
raise NotImplementedError(f"pad_or_cut_xd supports 1D/2D tensors, got ndim={ndim}")
# --- YAML config (OmegaConf) ---
def _apply_hparams_str(config: dict[str, Any], hparams_str: str) -> None:
if not hparams_str:
return
for item in hparams_str.split(","):
key, raw_value = item.split("=", maxsplit=1)
raw_value = raw_value.strip().strip("'\"")
node: dict[str, Any] = config
parts = key.split(".")
for part in parts[:-1]:
node = node.setdefault(part, {})
leaf = parts[-1]
current = node.get(leaf)
if raw_value in {"True", "False"}:
node[leaf] = raw_value == "True"
elif isinstance(current, bool):
node[leaf] = raw_value.lower() == "true"
elif isinstance(current, int):
node[leaf] = int(raw_value)
elif isinstance(current, float):
node[leaf] = float(raw_value)
else:
node[leaf] = type(current)(raw_value) if current is not None else raw_value
def load_yaml_config(
config: str | Path,
*,
hparams_str: str = "",
) -> dict[str, Any]:
"""Load YAML with optional ``base_config`` chain; returns a plain dict."""
config_path = Path(config).expanduser()
if not config_path.is_file():
raise FileNotFoundError(f"Config not found: {config_path}")
merged = OmegaConf.create({})
pending = [config_path.resolve()]
loaded: set[str] = set()
while pending:
path = pending.pop(0)
path_key = str(path)
if path_key in loaded:
continue
loaded.add(path_key)
current = OmegaConf.load(path_key)
base_configs = current.get("base_config")
if base_configs is not None:
entries = base_configs if OmegaConf.is_list(base_configs) else [base_configs]
for entry in entries:
base_path = Path(entry)
if not base_path.is_absolute():
base_path = (path.parent / entry).resolve()
if not base_path.is_file():
logger.warning("Skipping missing base config: %s", base_path)
continue
pending.insert(0, base_path)
merged = OmegaConf.merge(current, merged)
result = OmegaConf.to_container(merged, resolve=True)
if not isinstance(result, dict):
raise TypeError(f"Expected dict config, got {type(result)}")
_apply_hparams_str(result, hparams_str)
return result
def load_rosvot_config(config: str | Path, *, hparams_str: str = "") -> dict[str, Any]:
return load_yaml_config(config, hparams_str=hparams_str)
# --- checkpoint loading ---
def load_model_ckpt(
model: nn.Module,
checkpoint_path: str,
*,
model_name: str = "model",
verbose: bool = True,
) -> None:
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
if not isinstance(checkpoint, dict) or "state_dict" not in checkpoint:
raise ValueError(f"Unexpected ROSVOT checkpoint format: {checkpoint_path}")
state_dict = checkpoint["state_dict"]
if state_dict and "." in next(iter(state_dict)):
prefix = f"{model_name}."
state_dict = {k[len(prefix) :]: v for k, v in state_dict.items() if k.startswith(prefix)}
elif model_name in state_dict and isinstance(state_dict[model_name], dict):
state_dict = state_dict[model_name]
model.load_state_dict(state_dict, strict=True)
if verbose:
logger.info("Loaded ROSVOT weights from %s", checkpoint_path)
# --- pitch helpers (ROSVOT / RMVPE) ---
def f0_to_coarse(f0, f0_bin=256, f0_max=900.0, f0_min=50.0):
f0_mel_min = 1127 * np.log(1 + f0_min / 700)
f0_mel_max = 1127 * np.log(1 + f0_max / 700)
is_torch = isinstance(f0, torch.Tensor)
f0_mel = 1127 * (1 + f0 / 700).log() if is_torch else 1127 * np.log(1 + f0 / 700)
f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * (f0_bin - 2) / (f0_mel_max - f0_mel_min) + 1
f0_mel[f0_mel <= 1] = 1
f0_mel[f0_mel > f0_bin - 1] = f0_bin - 1
return (f0_mel + 0.5).long() if is_torch else np.rint(f0_mel).astype(int)
def norm_f0(f0, uv, pitch_norm="log", f0_mean=400, f0_std=100):
is_torch = isinstance(f0, torch.Tensor)
if pitch_norm == "standard":
f0 = (f0 - f0_mean) / f0_std
if pitch_norm == "log":
f0 = torch.log2(f0 + 1e-8) if is_torch else np.log2(f0 + 1e-8)
if uv is not None:
f0[uv > 0] = 0
return f0
def norm_interp_f0(f0, pitch_norm="log", f0_mean=None, f0_std=None):
is_torch = isinstance(f0, torch.Tensor)
if is_torch:
device = f0.device
f0 = f0.detach().cpu().numpy()
uv = f0 == 0
f0 = norm_f0(f0, uv, pitch_norm, f0_mean, f0_std)
if sum(uv) == len(f0):
f0[uv] = 0
elif sum(uv) > 0:
f0[uv] = np.interp(np.where(uv)[0], np.where(~uv)[0], f0[~uv])
if is_torch:
uv_t = torch.FloatTensor(uv)
f0_t = torch.FloatTensor(f0)
return f0_t.to(device), uv_t.to(device)
return f0, uv
def denorm_f0(f0, uv, pitch_norm="log", f0_mean=400, f0_std=100, pitch_padding=None, min=50, max=900):
is_torch = isinstance(f0, torch.Tensor)
if pitch_norm == "standard":
f0 = f0 * f0_std + f0_mean
if pitch_norm == "log":
f0 = 2**f0
f0 = f0.clamp(min=min, max=max) if is_torch else np.clip(f0, a_min=min, a_max=max)
if uv is not None:
f0[uv > 0] = 0
if pitch_padding is not None:
f0[pitch_padding] = 0
return f0
def boundary2Interval(bd):
is_torch = isinstance(bd, torch.Tensor)
if is_torch:
device = bd.device
bd = bd.detach().cpu().numpy()
ret = np.zeros(shape=(bd.sum() + 1, 2), dtype=int)
ret_idx = 0
ret[0, 0] = 0
for i, u in enumerate(bd):
if i == 0:
continue
if u == 1:
ret[ret_idx, 1] = i
ret[ret_idx + 1, 0] = i
ret_idx += 1
ret[-1, 1] = bd.shape[0] - 1
if is_torch:
return torch.LongTensor(ret).to(device)
return ret
vllm_omni/diffusion/models/soulx_singer/modules/preprocess/vocal_sep.py
-207
View File
@@ -1,207 +0,0 @@
"""BS-RoFormer vocal separation adapter (external pip package + SoulX chunking)."""
from __future__ import annotations
from pathlib import Path
from typing import ClassVar
import numpy as np
import torch
import torch.nn as nn
import yaml
from omegaconf import DictConfig, OmegaConf
from vllm.logger import init_logger
from vllm.multimodal.media.audio import load_audio
from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery
from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend
logger = init_logger(__name__)
class MelBandRoformerSeparator(nn.Module, SupportAudioInput, SupportsComponentDiscovery):
support_audio_input: ClassVar[bool] = True
_dit_modules: ClassVar[list[str]] = []
_encoder_modules: ClassVar[list[str]] = ["roformer"]
_vae_modules: ClassVar[list[str]] = []
_resident_modules: ClassVar[list[str]] = []
_layerwise_offload_blocks_attrs: ClassVar[list[str]] = ["roformer.layers"]
def __init__(
self,
config_path: str | Path,
checkpoint_path: str | Path | None = None,
*,
chunk_length_sec: float = 5,
use_half: bool = True,
):
super().__init__()
try:
from bs_roformer import MelBandRoformer as ExtModel
except ImportError as e:
raise ImportError(
'Install BS-RoFormer: pip install "vllm-omni[soulx-preprocess]" or pip install BS-RoFormer'
) from e
cfg = self._load_config(Path(config_path))
model_cfg = self._normalize_cfg(OmegaConf.to_container(cfg.model, resolve=True))
self.roformer = ExtModel(**model_cfg)
self.config = cfg
self.sample_rate = int(cfg.audio.sample_rate)
LayerWiseOffloadBackend.set_blocks_attr_names(self.roformer, ["layers"])
dtype = torch.float16 if use_half else torch.float32
self.roformer = self.roformer.to(dtype=dtype)
chunk = int(chunk_length_sec * self.sample_rate)
self.config.inference = getattr(self.config, "inference", {})
inf = self.config.inference
inf.chunk_size = chunk
inf.num_overlap = inf.get("num_overlap", 4)
inf.use_amp = False
inf.normalize = inf.get("normalize", True)
fade = chunk // 10
w = torch.ones(chunk)
w[:fade] = torch.linspace(0, 1, fade)
w[-fade:] = torch.linspace(1, 0, fade)
self.register_buffer("_window", w.to(dtype=dtype), persistent=False)
if checkpoint_path:
self.load_checkpoint(str(checkpoint_path), use_half=use_half)
else:
self.roformer.eval()
@staticmethod
def _normalize_cfg(cfg: dict) -> dict:
tuple_keys = {"multi_stft_resolutions_window_sizes", "freqs_per_bands", "multi_stft_window_sizes"}
def _fix(v):
if isinstance(v, list):
return tuple(_fix(x) for x in v)
if isinstance(v, dict):
return {k: _fix(val) for k, val in v.items()}
return v
for k in tuple_keys:
if k in cfg:
cfg[k] = _fix(cfg[k])
return cfg
@staticmethod
def _load_config(path: Path) -> DictConfig:
class Loader(yaml.SafeLoader):
pass
Loader.add_constructor(
"tag:yaml.org,2002:python/tuple",
lambda loader, node: tuple(loader.construct_sequence(node)),
)
with open(path) as f:
return OmegaConf.create(yaml.load(f, Loader=Loader))
def load_checkpoint(self, path: str, *, use_half: bool = True) -> None:
self.roformer.load_state_dict(torch.load(path, map_location="cpu"), strict=False)
if use_half:
self.roformer.half()
self.roformer.eval()
@property
def chunk_size(self) -> int:
return int(self.config.inference.chunk_size)
@property
def num_overlap(self) -> int:
return int(self.config.inference.num_overlap)
@property
def use_amp(self) -> bool:
return bool(self.config.inference.use_amp)
@property
def normalize_input(self) -> bool:
return bool(self.config.inference.normalize)
@property
def num_channels(self) -> int:
return int(self.config.audio.get("num_channels", 1))
def _prepare_mix(self, mix: np.ndarray) -> tuple[np.ndarray, tuple[float, float] | None]:
if mix.ndim == 1:
mix = mix[np.newaxis, :]
if self.num_channels == 2 and mix.shape[0] == 1:
mix = np.repeat(mix, 2, axis=0)
if not self.normalize_input:
return mix, None
mono = mix.mean(0)
return (mix - mono.mean()) / mono.std(), (float(mono.mean()), float(mono.std()))
@staticmethod
def _select_vocal(recon: torch.Tensor) -> torch.Tensor:
if recon.ndim == 4:
return recon[0, 0]
if recon.ndim == 3:
return recon[0]
if recon.ndim == 2:
return recon
raise ValueError(recon.shape)
@torch.inference_mode()
def forward(self, mix: torch.Tensor) -> torch.Tensor:
if mix.ndim == 1:
mix = mix.unsqueeze(0)
if self.num_channels == 2 and mix.shape[0] == 1:
mix = mix.repeat(2, 1)
model_dtype = next(self.roformer.parameters()).dtype
compute_dtype = torch.float32 if model_dtype == torch.float16 else model_dtype
if mix.dtype != compute_dtype:
mix = mix.to(dtype=compute_dtype)
chunk_size = self.chunk_size
step = chunk_size // self.num_overlap
fade = chunk_size // 10
border = chunk_size - step
n_samples = mix.shape[-1]
window = self._window.to(mix.device)
if n_samples > 2 * border and border > 0:
mix = nn.functional.pad(mix, (border, border), mode="reflect")
out = torch.zeros_like(mix)
weight = torch.zeros(mix.shape[-1], device=mix.device, dtype=mix.dtype)
pos = 0
with torch.cuda.amp.autocast(enabled=self.use_amp):
while pos < mix.shape[-1]:
part = mix[:, pos : pos + chunk_size]
seg_len = part.shape[-1]
pad_mode = "reflect" if seg_len > chunk_size // 2 else "constant"
part = nn.functional.pad(part, (0, chunk_size - seg_len), mode=pad_mode, value=0)
stem = self._select_vocal(self.roformer(part.unsqueeze(0)))
win = window.clone()
if pos == 0:
win[:fade] = 1
elif pos + step >= n_samples:
win[-fade:] = 1
out[:, pos : pos + seg_len] += stem[:, :seg_len] * win[:seg_len]
weight[pos : pos + seg_len] += win[:seg_len]
pos += step
out = out / weight
if n_samples > 2 * border and border > 0:
out = out[..., border:-border]
return out
@torch.inference_mode()
def separate_mono(self, mix: np.ndarray, device: torch.device) -> np.ndarray:
m, norm = self._prepare_mix(mix)
target_dtype = next(self.roformer.parameters()).dtype
v = self.forward(torch.as_tensor(m, dtype=target_dtype, device=device))
if norm:
v = v * norm[1] + norm[0]
v = v.mean(0) if v.ndim == 2 else v
return v.float().cpu().numpy().astype(np.float32)
def separate_path(self, path: str, device: torch.device) -> tuple[np.ndarray, int]:
mix, _ = load_audio(path, sr=self.sample_rate, mono=False)
return self.separate_mono(np.asarray(mix, dtype=np.float32), device), self.sample_rate
vllm_omni/diffusion/models/soulx_singer/modules/vocoder.py
-983
View File
@@ -1,983 +0,0 @@
import accelerate
import numpy as np
import scipy
import torch
from omegaconf import DictConfig
from torch import nn, view_as_complex, view_as_real
from torch.nn.utils import remove_weight_norm, weight_norm
from torchaudio.functional.functional import _hz_to_mel, _mel_to_hz
def _deep_update_dict(base: dict, override: dict | None) -> dict:
"""Recursively merge `override` into `base` (dict only)."""
if not override:
return base
for k, v in override.items():
if isinstance(v, dict) and isinstance(base.get(k), dict):
base[k] = _deep_update_dict(base[k], v)
else:
base[k] = v
return base
def _get_vocoder_default_cfg() -> dict:
# Defaults keep current behavior.
return {
"preprocess": {
"hop_size": 480,
"sample_rate": 24000,
"max_length": 36000,
"n_fft": 1920,
"num_mels": 128,
"win_size": 1920,
"fmin": 0,
"fmax": 12000,
"mel_var": 8.14,
"mel_mean": -4.92,
"load_phone": False,
"load_chromagram": False,
},
"model": {
"vocos": {
"input_channels": 128,
"dim": 1024,
"intermediate_dim": 4096,
"num_layers": 30,
"n_fft": 1920,
"hop_size": 480,
"padding": "same",
},
"period_gan": {
"max_downsample_channels": 1024,
"channels": 64,
"channel_increasing_factor": 2,
},
"spec_gan": {
"stft_params": {
"fft_sizes": [128, 256, 512, 1024, 2048],
"hop_sizes": [32, 64, 128, 256, 512],
"win_lengths": [128, 256, 512, 1024, 2048],
"window": "hann_window",
},
"in_channels": 1,
"out_channels": 1,
"channels": 64,
"kernel_sizes": [5, 3],
"max_downsample_channels": 1024,
"down_scales": [2, 2, 2],
"use_weight_norm": True,
"use_complex": False,
},
},
"loss": {
"mel_loss": {"sample_rate": 24000},
"disc_loss_weight": 1.0,
"mel_loss_weight": 10.0,
"adv_loss_weight": 2.0,
"fm_loss_weight": 2.0,
"spec_fm_loss_weight": 1.0,
},
}
def safe_log(x: torch.Tensor, clip_val: float = 1e-7) -> torch.Tensor:
"""
Computes the element-wise logarithm of the input tensor with clipping to avoid near-zero values.
Args:
x (Tensor): Input tensor.
clip_val (float, optional): Minimum value to clip the input tensor. Defaults to 1e-7.
Returns:
Tensor: Element-wise logarithm of the input tensor with clipping applied.
"""
return torch.log(torch.clip(x, min=clip_val))
def symlog(x: torch.Tensor) -> torch.Tensor:
return torch.sign(x) * torch.log1p(x.abs())
def symexp(x: torch.Tensor) -> torch.Tensor:
return torch.sign(x) * (torch.exp(x.abs()) - 1)
class STFT(nn.Module):
def __init__(
self,
n_fft: int,
hop_length: int,
win_length: int,
center=True,
):
super().__init__()
self.center = center
self.n_fft = n_fft
self.hop_length = hop_length
self.win_length = win_length
window = torch.hann_window(win_length)
self.register_buffer("window", window)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: (B, T * hop_length)
if not self.center:
pad = self.win_length - self.hop_length
x = torch.nn.functional.pad(x, (pad // 2, pad // 2), mode="reflect")
stft_spec = torch.stft(
x,
self.n_fft,
hop_length=self.hop_length,
win_length=self.win_length,
window=self.window,
center=self.center,
return_complex=False,
) # (B, n_fft // 2 + 1, T, 2)
rea = stft_spec[:, :, :, 0] # (B, n_fft // 2 + 1, T, 2)
imag = stft_spec[:, :, :, 1] # (B, n_fft // 2 + 1, T, 2)
log_mag = torch.log(
torch.abs(torch.sqrt(torch.pow(rea, 2) + torch.pow(imag, 2))) + 1e-5
) # (B, n_fft // 2 + 1, T)
phase = torch.atan2(imag, rea) # (B, n_fft // 2 + 1, T)
return log_mag, phase
class ISTFT(nn.Module):
"""
Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
See issue: https://github.com/pytorch/pytorch/issues/62323
Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
The NOLA constraint is met as we trim padded samples anyway.
Args:
n_fft (int): Size of Fourier transform.
hop_length (int): The distance between neighboring sliding window frames.
win_length (int): The size of window frame and STFT filter.
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
"""
def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"):
super().__init__()
if padding not in ["center", "same"]:
raise ValueError("Padding must be 'center' or 'same'.")
self.padding = padding
self.n_fft = n_fft
self.hop_length = hop_length
self.win_length = win_length
window = torch.hann_window(win_length)
self.register_buffer("window", window)
def forward(self, spec: torch.Tensor) -> torch.Tensor:
"""
Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
Args:
spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
N is the number of frequency bins, and T is the number of time frames.
Returns:
Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal.
"""
if self.padding == "center":
# Fallback to pytorch native implementation
return torch.istft(
spec,
self.n_fft,
self.hop_length,
self.win_length,
self.window,
center=True,
)
elif self.padding == "same":
pad = (self.win_length - self.hop_length) // 2
else:
raise ValueError("Padding must be 'center' or 'same'.")
assert spec.dim() == 3, "Expected a 3D tensor as input"
B, N, T = spec.shape
# Inverse FFT
ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
ifft = ifft * self.window[None, :, None]
# Overlap and Add
output_size = (T - 1) * self.hop_length + self.win_length
y = torch.nn.functional.fold(
ifft,
output_size=(1, output_size),
kernel_size=(1, self.win_length),
stride=(1, self.hop_length),
)[:, 0, 0, pad:-pad]
# Window envelope
window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
window_envelope = torch.nn.functional.fold(
window_sq,
output_size=(1, output_size),
kernel_size=(1, self.win_length),
stride=(1, self.hop_length),
).squeeze()[pad:-pad]
# Normalize
assert (window_envelope > 1e-11).all()
y = y / window_envelope
return y
class MDCT(nn.Module):
"""
Modified Discrete Cosine Transform (MDCT) module.
Args:
frame_len (int): Length of the MDCT frame.
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
"""
def __init__(self, frame_len: int, padding: str = "same"):
super().__init__()
if padding not in ["center", "same"]:
raise ValueError("Padding must be 'center' or 'same'.")
self.padding = padding
self.frame_len = frame_len
N = frame_len // 2
n0 = (N + 1) / 2
window = torch.from_numpy(scipy.signal.cosine(frame_len)).float()
self.register_buffer("window", window)
pre_twiddle = torch.exp(-1j * torch.pi * torch.arange(frame_len) / frame_len)
post_twiddle = torch.exp(-1j * torch.pi * n0 * (torch.arange(N) + 0.5) / N)
# view_as_real: NCCL Backend does not support ComplexFloat data type
# https://github.com/pytorch/pytorch/issues/71613
self.register_buffer("pre_twiddle", view_as_real(pre_twiddle))
self.register_buffer("post_twiddle", view_as_real(post_twiddle))
def forward(self, audio: torch.Tensor) -> torch.Tensor:
"""
Apply the Modified Discrete Cosine Transform (MDCT) to the input audio.
Args:
audio (Tensor): Input audio waveform of shape (B, T), where B is the batch size
and T is the length of the audio.
Returns:
Tensor: MDCT coefficients of shape (B, L, N), where L is the number of output frames
and N is the number of frequency bins.
"""
if self.padding == "center":
audio = torch.nn.functional.pad(audio, (self.frame_len // 2, self.frame_len // 2))
elif self.padding == "same":
# hop_length is 1/2 frame_len
audio = torch.nn.functional.pad(audio, (self.frame_len // 4, self.frame_len // 4))
else:
raise ValueError("Padding must be 'center' or 'same'.")
x = audio.unfold(-1, self.frame_len, self.frame_len // 2)
N = self.frame_len // 2
x = x * self.window.expand(x.shape)
x = torch.fft.fft(x * view_as_complex(self.pre_twiddle).expand(x.shape), dim=-1)[..., :N]
res = x * view_as_complex(self.post_twiddle).expand(x.shape) * np.sqrt(1 / N)
return torch.real(res) * np.sqrt(2)
class IMDCT(nn.Module):
"""
Inverse Modified Discrete Cosine Transform (IMDCT) module.
Args:
frame_len (int): Length of the MDCT frame.
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
"""
def __init__(self, frame_len: int, padding: str = "same"):
super().__init__()
if padding not in ["center", "same"]:
raise ValueError("Padding must be 'center' or 'same'.")
self.padding = padding
self.frame_len = frame_len
N = frame_len // 2
n0 = (N + 1) / 2
window = torch.from_numpy(scipy.signal.cosine(frame_len)).float()
self.register_buffer("window", window)
pre_twiddle = torch.exp(1j * torch.pi * n0 * torch.arange(N * 2) / N)
post_twiddle = torch.exp(1j * torch.pi * (torch.arange(N * 2) + n0) / (N * 2))
self.register_buffer("pre_twiddle", view_as_real(pre_twiddle))
self.register_buffer("post_twiddle", view_as_real(post_twiddle))
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Apply the Inverse Modified Discrete Cosine Transform (IMDCT) to the input MDCT coefficients.
Args:
x (Tensor): Input MDCT coefficients of shape (B, L, N), where B is the batch size,
L is the number of frames, and N is the number of frequency bins.
Returns:
Tensor: Reconstructed audio waveform of shape (B, T), where T is the length of the audio.
"""
B, L, N = x.shape
Y = torch.zeros((B, L, N * 2), dtype=x.dtype, device=x.device)
Y[..., :N] = x
Y[..., N:] = -1 * torch.conj(torch.flip(x, dims=(-1,)))
y = torch.fft.ifft(Y * view_as_complex(self.pre_twiddle).expand(Y.shape), dim=-1)
y = torch.real(y * view_as_complex(self.post_twiddle).expand(y.shape)) * np.sqrt(N) * np.sqrt(2)
result = y * self.window.expand(y.shape)
output_size = (1, (L + 1) * N)
audio = torch.nn.functional.fold(
result.transpose(1, 2),
output_size=output_size,
kernel_size=(1, self.frame_len),
stride=(1, self.frame_len // 2),
)[:, 0, 0, :]
if self.padding == "center":
pad = self.frame_len // 2
elif self.padding == "same":
pad = self.frame_len // 4
else:
raise ValueError("Padding must be 'center' or 'same'.")
audio = audio[:, pad:-pad]
return audio
class FourierHead(nn.Module):
"""Base class for inverse fourier modules."""
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
L is the sequence length, and H denotes the model dimension.
Returns:
Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
"""
raise NotImplementedError("Subclasses must implement the forward method.")
class ISTFTHead(FourierHead):
"""
ISTFT Head module for predicting STFT complex coefficients.
Args:
dim (int): Hidden dimension of the model.
n_fft (int): Size of Fourier transform.
hop_length (int): The distance between neighboring sliding window frames, which should align with
the resolution of the input features.
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
"""
def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
super().__init__()
out_dim = n_fft + 2
self.out = torch.nn.Linear(dim, out_dim)
self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Forward pass of the ISTFTHead module.
Args:
x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
L is the sequence length, and H denotes the model dimension.
Returns:
Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
"""
x = self.out(x).transpose(1, 2)
mag, p = x.chunk(2, dim=1)
mag = torch.exp(mag)
mag = torch.clip(mag, max=1e2) # safeguard to prevent excessively large magnitudes
# wrapping happens here. These two lines produce real and imaginary value
x = torch.cos(p)
y = torch.sin(p)
# recalculating phase here does not produce anything new
# only costs time
# phase = torch.atan2(y, x)
# S = mag * torch.exp(phase * 1j)
# better directly produce the complex value
# Always compute complex values in float32 to avoid ComplexHalf warning (then cast audio back if needed)
orig_dtype = mag.dtype
mag, x, y = mag.float(), x.float(), y.float()
S = mag * (x + 1j * y)
audio = self.istft(S)
if orig_dtype != torch.float32:
audio = audio.to(orig_dtype)
return audio
class IMDCTSymExpHead(FourierHead):
"""
IMDCT Head module for predicting MDCT coefficients with symmetric exponential function
Args:
dim (int): Hidden dimension of the model.
mdct_frame_len (int): Length of the MDCT frame.
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
sample_rate (int, optional): The sample rate of the audio. If provided, the last layer will be initialized
based on perceptual scaling. Defaults to None.
clip_audio (bool, optional): If True, clip audio to [-1.0, 1.0]. Defaults to False.
"""
def __init__(
self,
dim: int,
mdct_frame_len: int,
padding: str = "same",
sample_rate: int | None = None,
clip_audio: bool = False,
):
super().__init__()
out_dim = mdct_frame_len // 2
self.out = nn.Linear(dim, out_dim)
self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding)
self.clip_audio = clip_audio
if sample_rate is not None:
# optionally init the last layer following mel-scale
m_max = _hz_to_mel(sample_rate // 2)
m_pts = torch.linspace(0, m_max, out_dim)
f_pts = _mel_to_hz(m_pts)
scale = 1 - (f_pts / f_pts.max())
with torch.no_grad():
self.out.weight.mul_(scale.view(-1, 1))
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Forward pass of the IMDCTSymExpHead module.
Args:
x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
L is the sequence length, and H denotes the model dimension.
Returns:
Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
"""
x = self.out(x)
x = symexp(x)
x = torch.clip(x, min=-1e2, max=1e2) # safeguard to prevent excessively large magnitudes
audio = self.imdct(x)
if self.clip_audio:
audio = torch.clip(x, min=-1.0, max=1.0)
return audio
class IMDCTCosHead(FourierHead):
"""
IMDCT Head module for predicting MDCT coefficients with parametrizing MDCT = exp(m) · cos(p)
Args:
dim (int): Hidden dimension of the model.
mdct_frame_len (int): Length of the MDCT frame.
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
clip_audio (bool, optional): If True, clip audio to [-1.0, 1.0]. Defaults to False.
"""
def __init__(
self,
dim: int,
mdct_frame_len: int,
padding: str = "same",
clip_audio: bool = False,
):
super().__init__()
self.clip_audio = clip_audio
self.out = nn.Linear(dim, mdct_frame_len)
self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Forward pass of the IMDCTCosHead module.
Args:
x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
L is the sequence length, and H denotes the model dimension.
Returns:
Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
"""
x = self.out(x)
m, p = x.chunk(2, dim=2)
m = torch.exp(m).clip(max=1e2) # safeguard to prevent excessively large magnitudes
audio = self.imdct(m * torch.cos(p))
if self.clip_audio:
audio = torch.clip(x, min=-1.0, max=1.0)
return audio
class ConvNeXtBlock(nn.Module):
"""ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.
Args:
dim (int): Number of input channels.
intermediate_dim (int): Dimensionality of the intermediate layer.
layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
Defaults to None.
adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
None means non-conditional LayerNorm. Defaults to None.
"""
def __init__(
self,
dim: int,
intermediate_dim: int,
layer_scale_init_value: float,
adanorm_num_embeddings: int | None = None,
):
super().__init__()
self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
self.adanorm = adanorm_num_embeddings is not None
if adanorm_num_embeddings:
self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6)
else:
self.norm = nn.LayerNorm(dim, eps=1e-6)
self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers
self.act = nn.GELU()
self.pwconv2 = nn.Linear(intermediate_dim, dim)
self.gamma = (
nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
if layer_scale_init_value > 0
else None
)
def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor | None = None) -> torch.Tensor:
residual = x
x = self.dwconv(x)
x = x.transpose(1, 2) # (B, C, T) -> (B, T, C)
if self.adanorm:
assert cond_embedding_id is not None
x = self.norm(x, cond_embedding_id)
else:
x = self.norm(x)
x = self.pwconv1(x)
x = self.act(x)
x = self.pwconv2(x)
if self.gamma is not None:
x = self.gamma * x
x = x.transpose(1, 2) # (B, T, C) -> (B, C, T)
x = residual + x
return x
class AdaLayerNorm(nn.Module):
"""
Adaptive Layer Normalization module with learnable embeddings per `num_embeddings` classes
Args:
num_embeddings (int): Number of embeddings.
embedding_dim (int): Dimension of the embeddings.
"""
def __init__(self, num_embeddings: int, embedding_dim: int, eps: float = 1e-6):
super().__init__()
self.eps = eps
self.dim = embedding_dim
self.scale = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim)
self.shift = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim)
torch.nn.init.ones_(self.scale.weight)
torch.nn.init.zeros_(self.shift.weight)
def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor) -> torch.Tensor:
scale = self.scale(cond_embedding_id)
shift = self.shift(cond_embedding_id)
x = nn.functional.layer_norm(x, (self.dim,), eps=self.eps)
x = x * scale + shift
return x
class ResBlock1(nn.Module):
"""
ResBlock adapted from HiFi-GAN V1 (https://github.com/jik876/hifi-gan) with dilated 1D convolutions,
but without upsampling layers.
Args:
dim (int): Number of input channels.
kernel_size (int, optional): Size of the convolutional kernel. Defaults to 3.
dilation (tuple[int], optional): Dilation factors for the dilated convolutions.
Defaults to (1, 3, 5).
lrelu_slope (float, optional): Negative slope of the LeakyReLU activation function.
Defaults to 0.1.
layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
Defaults to None.
"""
def __init__(
self,
dim: int,
kernel_size: int = 3,
dilation: tuple[int, int, int] = (1, 3, 5),
lrelu_slope: float = 0.1,
layer_scale_init_value: float | None = None,
):
super().__init__()
self.lrelu_slope = lrelu_slope
self.convs1 = nn.ModuleList(
[
weight_norm(
nn.Conv1d(
dim,
dim,
kernel_size,
1,
dilation=dilation[0],
padding=self.get_padding(kernel_size, dilation[0]),
)
),
weight_norm(
nn.Conv1d(
dim,
dim,
kernel_size,
1,
dilation=dilation[1],
padding=self.get_padding(kernel_size, dilation[1]),
)
),
weight_norm(
nn.Conv1d(
dim,
dim,
kernel_size,
1,
dilation=dilation[2],
padding=self.get_padding(kernel_size, dilation[2]),
)
),
]
)
self.convs2 = nn.ModuleList(
[
weight_norm(
nn.Conv1d(
dim,
dim,
kernel_size,
1,
dilation=1,
padding=self.get_padding(kernel_size, 1),
)
),
weight_norm(
nn.Conv1d(
dim,
dim,
kernel_size,
1,
dilation=1,
padding=self.get_padding(kernel_size, 1),
)
),
weight_norm(
nn.Conv1d(
dim,
dim,
kernel_size,
1,
dilation=1,
padding=self.get_padding(kernel_size, 1),
)
),
]
)
self.gamma = nn.ParameterList(
[
(
nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True)
if layer_scale_init_value is not None
else None
),
(
nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True)
if layer_scale_init_value is not None
else None
),
(
nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True)
if layer_scale_init_value is not None
else None
),
]
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
for c1, c2, gamma in zip(self.convs1, self.convs2, self.gamma):
xt = torch.nn.functional.leaky_relu(x, negative_slope=self.lrelu_slope)
xt = c1(xt)
xt = torch.nn.functional.leaky_relu(xt, negative_slope=self.lrelu_slope)
xt = c2(xt)
if gamma is not None:
xt = gamma * xt
x = xt + x
return x
def remove_weight_norm(self):
for conv in self.convs1:
remove_weight_norm(conv)
for conv in self.convs2:
remove_weight_norm(conv)
@staticmethod
def get_padding(kernel_size: int, dilation: int = 1) -> int:
return int((kernel_size * dilation - dilation) / 2)
class Backbone(nn.Module):
"""Base class for the generator's backbone. It preserves the same temporal resolution across all layers."""
def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
"""
Args:
x (Tensor): Input tensor of shape (B, C, L), where B is the batch size,
C denotes output features, and L is the sequence length.
Returns:
Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length,
and H denotes the model dimension.
"""
raise NotImplementedError("Subclasses must implement the forward method.")
class VocosBackbone(Backbone):
"""
Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization
Args:
input_channels (int): Number of input features channels.
dim (int): Hidden dimension of the model.
intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
num_layers (int): Number of ConvNeXtBlock layers.
layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`.
adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
None means non-conditional model. Defaults to None.
"""
def __init__(
self,
input_channels: int,
dim: int,
intermediate_dim: int,
num_layers: int,
layer_scale_init_value: float | None = None,
adanorm_num_embeddings: int | None = None,
):
super().__init__()
self.input_channels = input_channels
self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3)
self.adanorm = adanorm_num_embeddings is not None
if adanorm_num_embeddings:
self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6)
else:
self.norm = nn.LayerNorm(dim, eps=1e-6)
layer_scale_init_value = layer_scale_init_value or 1 / num_layers
self.convnext = nn.ModuleList(
[
ConvNeXtBlock(
dim=dim,
intermediate_dim=intermediate_dim,
layer_scale_init_value=layer_scale_init_value,
adanorm_num_embeddings=adanorm_num_embeddings,
)
for _ in range(num_layers)
]
)
self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Conv1d | nn.Linear):
nn.init.trunc_normal_(m.weight, std=0.02)
nn.init.constant_(m.bias, 0)
def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
bandwidth_id = kwargs.get("bandwidth_id", None)
x = self.embed(x)
if self.adanorm:
assert bandwidth_id is not None
x = self.norm(x.transpose(1, 2), cond_embedding_id=bandwidth_id)
else:
x = self.norm(x.transpose(1, 2))
x = x.transpose(1, 2)
for conv_block in self.convnext:
x = conv_block(x, cond_embedding_id=bandwidth_id)
x = self.final_layer_norm(x.transpose(1, 2))
return x
class VocosResNetBackbone(Backbone):
"""
Vocos backbone module built with ResBlocks.
Args:
input_channels (int): Number of input features channels.
dim (int): Hidden dimension of the model.
num_blocks (int): Number of ResBlock1 blocks.
layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to None.
"""
def __init__(
self,
input_channels,
dim,
num_blocks,
layer_scale_init_value=None,
):
super().__init__()
self.input_channels = input_channels
self.embed = weight_norm(nn.Conv1d(input_channels, dim, kernel_size=3, padding=1))
layer_scale_init_value = layer_scale_init_value or 1 / num_blocks / 3
self.resnet = nn.Sequential(
*[ResBlock1(dim=dim, layer_scale_init_value=layer_scale_init_value) for _ in range(num_blocks)]
)
def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
x = self.embed(x)
x = self.resnet(x)
x = x.transpose(1, 2)
return x
class Vocos(nn.Module):
def __init__(
self,
input_channels: int = 256,
dim: int = 384,
intermediate_dim: int = 1152,
num_layers: int = 8,
n_fft: int = 800,
hop_size: int = 200,
padding: str = "same",
adanorm_num_embeddings=None,
cfg=None,
):
super().__init__()
input_channels = cfg.input_channels if cfg is not None and hasattr(cfg, "input_channels") else input_channels
dim = cfg.dim if cfg is not None and hasattr(cfg, "dim") else dim
intermediate_dim = (
cfg.intermediate_dim if cfg is not None and hasattr(cfg, "intermediate_dim") else intermediate_dim
)
num_layers = cfg.num_layers if cfg is not None and hasattr(cfg, "num_layers") else num_layers
adanorm_num_embeddings = (
cfg.adanorm_num_embeddings
if cfg is not None and hasattr(cfg, "adanorm_num_embeddings")
else adanorm_num_embeddings
)
n_fft = cfg.n_fft if cfg is not None and hasattr(cfg, "n_fft") else n_fft
hop_size = cfg.hop_size if cfg is not None and hasattr(cfg, "hop_size") else hop_size
padding = cfg.padding if cfg is not None and hasattr(cfg, "padding") else padding
self.backbone = VocosBackbone(
input_channels=input_channels,
dim=dim,
intermediate_dim=intermediate_dim,
num_layers=num_layers,
adanorm_num_embeddings=adanorm_num_embeddings,
)
self.head = ISTFTHead(dim, n_fft, hop_size, padding)
def forward(self, x):
x = self.backbone(x)
x = self.head(x)
return x[:, None, :]
class JsonHParams:
def __init__(self, **kwargs):
for k, v in kwargs.items():
if isinstance(v, dict):
v = JsonHParams(**v)
self[k] = v
def keys(self):
return self.__dict__.keys()
def items(self):
return self.__dict__.items()
def values(self):
return self.__dict__.values()
def __len__(self):
return len(self.__dict__)
def __getitem__(self, key):
return getattr(self, key)
def __setitem__(self, key, value):
return setattr(self, key, value)
def __contains__(self, key):
return key in self.__dict__
def __repr__(self):
return self.__dict__.__repr__()
def build_vocoder_model(cfg):
vocoder_model = Vocos(cfg=cfg.model.vocos)
vocoder_model.eval()
return vocoder_model
def load_checkpoint(build_model_func, cfg, ckpt_path):
model = build_model_func(cfg)
if ckpt_path is not None:
accelerate.load_checkpoint_and_dispatch(model, ckpt_path)
return model
def load_vocos_model(
ckpt_path: str | None = None,
config: DictConfig = None,
device: str = "cuda",
):
"""Load Vocos vocoder.
Args:
config: DictConfig, config for vocoder.
ckpt_path: str | None = None, path to checkpoint.
device: str, device to load model. Note: accelerate dispatch handles device placement.
"""
merged_cfg = _deep_update_dict(_get_vocoder_default_cfg(), config)
vocoder_cfg_obj = JsonHParams(**merged_cfg)
vocoder_model = load_checkpoint(build_vocoder_model, vocoder_cfg_obj, ckpt_path)
vocoder_model.eval()
for param in vocoder_model.parameters():
param.requires_grad = False
return vocoder_model
class Vocoder(nn.Module):
def __init__(self, vocoder: dict | None = None, ckpt_path: str | None = None):
super().__init__()
vocoder = vocoder or {}
model_cfg = vocoder.get("model_cfg") or vocoder.get("cfg")
self.model = load_vocos_model(
ckpt_path=ckpt_path,
config=model_cfg,
)
def forward(self, x):
x = self.model(x)
return x
vllm_omni/diffusion/models/soulx_singer/modules/whisper_encoder.py
-96
View File
@@ -1,96 +0,0 @@
"""Frozen Whisper encoder wrapper (wav -> encoder embeddings)."""
import torch
import torchaudio
from transformers import WhisperFeatureExtractor, WhisperModel
WHISPER_MEL_FRAMES = 3000 # 3000 frames at 16000 Hz
class WhisperEncoder:
"""Auxiliary content encoder for SVC.
Kept in FP32 for stable HF Whisper inference under vLLM-Omni's default
BF16 diffusion trunk (same idea as Stable Audio's FP32 VAE). Outputs are
cast to the trunk dtype before fusion with the CFM conditioner.
"""
def __init__(self, device: str | torch.device | None = None) -> None:
self.fe = WhisperFeatureExtractor.from_pretrained("openai/whisper-base")
self.model = WhisperModel.from_pretrained("openai/whisper-base")
self.device = torch.device(
device or ("cuda" if torch.cuda.is_available() else "cpu"),
)
self.model = self.model.to(device=self.device, dtype=torch.float32)
self.model.eval()
for param in self.model.parameters():
param.requires_grad_(False)
@property
def dtype(self) -> torch.dtype:
return torch.float32
def float(self) -> "WhisperEncoder":
self.model = self.model.float()
return self
def encode(
self,
wav: torch.Tensor,
sr: int,
*,
output_dtype: torch.dtype | None = None,
) -> torch.Tensor:
wav = (
torchaudio.functional.resample(wav, orig_freq=sr, new_freq=self.fe.sampling_rate)
if sr != self.fe.sampling_rate
else wav
)
wav_np = wav.detach().float().cpu().numpy().astype("float32", copy=False)
inputs = self.fe(
wav_np,
sampling_rate=self.fe.sampling_rate,
return_tensors="pt",
padding=False,
truncation=False,
return_attention_mask=True,
)
input_features = inputs.input_features
num_frames = input_features.shape[-1]
if num_frames < WHISPER_MEL_FRAMES:
pad = WHISPER_MEL_FRAMES - num_frames
input_features = torch.nn.functional.pad(input_features, (0, pad))
else:
input_features = input_features[..., :WHISPER_MEL_FRAMES]
input_features = input_features.to(device=wav.device, dtype=self.dtype)
if self.model.device != wav.device:
self.model = self.model.to(wav.device)
attention_mask = inputs.attention_mask.to(wav.device) if inputs.attention_mask is not None else None
with torch.inference_mode():
encoder_out = self.model.encoder(input_features).last_hidden_state
if attention_mask is not None:
valid_mel_frames = attention_mask.sum(dim=1)
valid_enc_frames = (valid_mel_frames + 1) // 2
max_valid_enc_frames = min(int(valid_enc_frames.max().item()), encoder_out.shape[1])
encoder_out = encoder_out[:, :max_valid_enc_frames, :]
valid_len = min(int(valid_enc_frames[0].item()), max_valid_enc_frames)
if valid_len < max_valid_enc_frames:
encoder_out[0, valid_len:, :] = 0
if output_dtype is not None and encoder_out.dtype != output_dtype:
encoder_out = encoder_out.to(dtype=output_dtype)
return encoder_out
if __name__ == "__main__":
torch.manual_seed(0)
audio = torch.randn(1, 24000 * 25).float().to("cuda")
encoder = WhisperEncoder()
whisper_encoder_out = encoder.encode(audio, sr=24000)
print(whisper_encoder_out.shape)
vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py
-377
View File
@@ -1,377 +0,0 @@
import os
import time
from collections.abc import Callable
from contextlib import contextmanager
from typing import Any, ClassVar
import numpy as np
import torch
import torch.nn as nn
from vllm.logger import init_logger
from vllm.utils.torch_utils import set_default_torch_dtype
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _unwrap, _wrap
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.models.interface import (
SupportAudioInput,
SupportAudioOutput,
SupportsComponentDiscovery,
)
from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import consume_payload
from vllm_omni.diffusion.models.soulx_singer.utils import _patch_torchaudio_load, load_config
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.platforms import current_omni_platform
logger = init_logger(__name__)
_SOULX_PROFILER_TARGETS = [
"forward",
]
def convert_soulx_audio_output_to_numpy(audio: torch.Tensor | dict[str, Any]) -> np.ndarray | dict[str, Any]:
if isinstance(audio, dict) and isinstance(audio.get("payload"), dict):
payload = dict(audio["payload"])
audio_payload = payload.get("audio")
if isinstance(audio_payload, torch.Tensor):
payload["audio"] = audio_payload.detach().cpu().float().numpy()
return {
"payload": payload,
"metadata": audio.get("metadata") or {},
}
if isinstance(audio, torch.Tensor):
return audio.detach().cpu().float().numpy()
return audio
class FlowMatchingAudioPipeline(
nn.Module,
CFGParallelMixin,
SupportAudioInput,
SupportAudioOutput,
SupportsComponentDiscovery,
DiffusionPipelineProfilerMixin,
):
"""SoulX flow-matching pipeline with CFG / CFG-parallel support."""
support_audio_input: ClassVar[bool] = True
support_audio_output: ClassVar[bool] = True
audio_sample_rate: ClassVar[int] = 24000
_DEFAULT_RESCALE_CFG: ClassVar[float] = 0.75
_dit_modules: ClassVar[list[str]] = ["cfm_decoder.model.diff_estimator"]
_encoder_modules: ClassVar[list[str]] = ["mel", "f0_encoder"]
_vae_modules: ClassVar[list[str]] = ["vocoder"]
_resident_modules: ClassVar[list[str]] = ["cfm_decoder.model.cond_emb"]
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__()
self.od_config = od_config
self.device = get_local_device()
# 1. parse model path and yaml config
model_dir = od_config.model
assert model_dir is not None, "model_dir is not provided"
if not os.path.exists(model_dir):
raise FileNotFoundError(
f"""Model directory {model_dir} does not exist. User should refer to
`examples/offline_inference/text_to_speech/README.md` for usage details."""
)
# Empty: skip DiffusersPipelineLoader.get_all_weights() → no _prepare_weights / file glob /
# safetensors or pt iterator from the framework. Load checkpoints in __init__ instead.
self.weights_sources = ()
self.model_path = model_dir
config_path = os.path.join(model_dir, "config.yaml")
hf_config = load_config(config_path)
self.audio_config = hf_config.audio
self.encoder_config = hf_config.model.encoder
self.flow_matching_config = hf_config.model.flow_matching
self.mel_dim = self.audio_config.num_mels
self.f0_bin = self.encoder_config.f0_bin
self.f0_dim = self.encoder_config.f0_dim
self.text_dim = self.encoder_config.text_dim
self.vocab_size = self.encoder_config.vocab_size
self.pitch_dim = self.encoder_config.pitch_dim
self.type_dim = self.encoder_config.type_dim
_patch_torchaudio_load()
def _mel_to_audio(self, generated_mel: torch.Tensor, *, squeeze: bool = False) -> torch.Tensor:
audio = self.vocoder(generated_mel.transpose(1, 2)[0:1, ...].float()).float()
return audio.squeeze() if squeeze else audio
def _mel_from_wav(self, wav: torch.Tensor) -> torch.Tensor:
if wav.dtype != torch.float32:
wav = wav.float()
return self.mel(wav)
def _load_soulx_checkpoint(self, checkpoint_name: str) -> None:
weight_path = os.path.join(self.model_path, checkpoint_name)
if not os.path.isfile(weight_path):
raise FileNotFoundError(
f"Model weights not found at {weight_path}. Please check the pretrained model path."
)
state = torch.load(weight_path, map_location=self.device)
self.mel.float()
self.load_state_dict(state["state_dict"], strict=True)
self._finalize_loaded_dtypes()
logger.info("Loaded model weights from %s", weight_path)
def _setup_soulx_profiler(self, *, extra_targets: list[str] | None = None) -> None:
"""Enable stage timing when ``od_config.enable_diffusion_pipeline_profiler`` is set."""
targets = list(_SOULX_PROFILER_TARGETS)
if extra_targets:
targets.extend(extra_targets)
self.setup_diffusion_pipeline_profiler(
profiler_targets=targets,
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler,
)
@contextmanager
def _stage_timer(self, name: str):
if not getattr(self, "enable_diffusion_pipeline_profiler", False):
yield
return
if current_omni_platform.is_available():
current_omni_platform.synchronize()
start = time.perf_counter()
try:
yield
finally:
if current_omni_platform.is_available():
current_omni_platform.synchronize()
duration = time.perf_counter() - start
with self._profiler_lock:
self._stage_durations[name] = self._stage_durations.get(name, 0.0) + duration
def _profiler_stage_durations(self) -> dict[str, float] | None:
if getattr(self, "enable_diffusion_pipeline_profiler", False):
return self.stage_durations
return None
@staticmethod
def _build_fp32_audio_modules(audio_config) -> tuple[nn.Module, nn.Module]:
"""Mel/vocoder must stay FP32 even when ``od_config.dtype`` is FP16/BF16."""
from vllm_omni.diffusion.models.soulx_singer.modules import MelSpectrogramEncoder, Vocoder
with set_default_torch_dtype(torch.float32):
mel = MelSpectrogramEncoder(audio_config)
vocoder = Vocoder()
return mel, vocoder
def _finalize_loaded_dtypes(self) -> None:
"""Keep mel/vocoder in FP32; cast DiT trunk to ``od_config.dtype`` after load."""
self.mel.float()
self.vocoder.float()
trunk_dtype = self.od_config.dtype
if trunk_dtype not in (torch.float16, torch.bfloat16):
return
trunk_modules = (
"cfm_decoder",
"f0_encoder",
"preflow",
"note_text_encoder",
"note_pitch_encoder",
"note_type_encoder",
)
for name in trunk_modules:
module = getattr(self, name, None)
if module is not None:
module.to(dtype=trunk_dtype)
whisper = getattr(self, "whisper_encoder", None)
if whisper is not None:
whisper.float()
@property
def transformer(self) -> nn.Module:
"""Alias expected by ``CFGParallelMixin`` default hooks."""
return self.cfm_decoder.model.diff_estimator
@property
def diffusion_trunk_dtype(self) -> torch.dtype:
"""DiT trunk dtype from ``diff_estimator`` (not ``cond_emb``, which is registered first)."""
return next(self.cfm_decoder.model.diff_estimator.parameters()).dtype
def _to_trunk_dtype(self, *tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:
"""Cast tensors to ``diffusion_trunk_dtype`` at the CFM/DiT boundary."""
trunk_dtype = self.diffusion_trunk_dtype
return tuple(t if t.dtype == trunk_dtype else t.to(dtype=trunk_dtype) for t in tensors)
def _resolve_diffusion_generator(self, sampling_params: Any) -> torch.Generator | None:
"""Build a device-local RNG for CFM noise from ``generator`` or ``seed``."""
generator = getattr(sampling_params, "generator", None)
if generator is not None:
return generator
seed = getattr(sampling_params, "seed", None)
if seed is None:
return None
return torch.Generator(device=self.device).manual_seed(int(seed))
def _prepare_cfm_loop_state(
self,
prompt: torch.Tensor,
cond: torch.Tensor,
n_timesteps: int,
*,
generator: torch.Generator | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
"""Cast FP32 islands to trunk dtype and build CFM loop tensors/masks."""
trunk_dtype = self.diffusion_trunk_dtype
prompt, cond = self._to_trunk_dtype(prompt, cond)
cond_emb = self.cfm_decoder.model.cond_emb(cond)
if cond_emb.dtype != trunk_dtype:
cond_emb = self._to_trunk_dtype(cond_emb)[0]
device = cond_emb.device
h = 1.0 / n_timesteps
h_tensor = torch.tensor(h, device=device, dtype=trunk_dtype)
prompt_len = prompt.shape[1]
target_len = cond_emb.shape[1] - prompt_len
x_mask = torch.ones(cond_emb.shape[0], target_len, device=device, dtype=trunk_dtype)
prompt_mask = torch.ones(cond_emb.shape[0], prompt_len, device=device, dtype=trunk_dtype)
xt_mask = torch.cat([prompt_mask, x_mask], dim=1)
xt = torch.randn(
cond_emb.shape[0],
target_len,
self.mel_dim,
dtype=trunk_dtype,
device=device,
generator=generator,
)
return cond_emb, prompt, xt, x_mask, xt_mask, h_tensor, prompt_len
def predict_noise(
self,
*,
x: torch.Tensor,
diffusion_step: torch.Tensor,
cond: torch.Tensor,
x_mask: torch.Tensor,
prompt_len: int | None = None,
**_: Any,
) -> torch.Tensor:
"""Forward through ``diff_estimator``, optionally slice off prompt frames."""
x, diffusion_step, cond, x_mask = self._to_trunk_dtype(x, diffusion_step, cond, x_mask)
flow_pred = self.cfm_decoder.model.diff_estimator(x, diffusion_step, cond, x_mask)
if prompt_len is not None:
flow_pred = flow_pred[:, prompt_len:, :]
return flow_pred
def combine_cfg_noise(
self,
positive_noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
true_cfg_scale: float,
cfg_normalize: bool = False,
rescale_cfg: float | None = None,
) -> torch.Tensor | tuple[torch.Tensor, ...]:
"""SoulX CFG: pos + scale * (pos - neg) with std rescaling."""
del cfg_normalize # SoulX uses rescale_cfg instead of cfg_normalize_function.
if rescale_cfg is None:
rescale_cfg = getattr(self, "rescale_cfg", self._DEFAULT_RESCALE_CFG)
pos = _unwrap(_wrap(positive_noise_pred))
neg = _unwrap(_wrap(negative_noise_pred))
pos_flow_pred_std = pos.std()
flow_pred_cfg = pos + true_cfg_scale * (pos - neg)
rescale_flow_pred = flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std()
combined = rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg
return combined.to(dtype=pos.dtype)
def _run_flow_matching_loop(
self,
prompt: torch.Tensor,
cond: torch.Tensor,
n_timesteps: int,
cfg: float,
*,
generator: torch.Generator | None = None,
) -> torch.Tensor:
"""Euler flow-matching loop aligned with ``FlowMatchingTransformer.reverse_diffusion``."""
cond_emb, prompt, xt, x_mask, xt_mask, h_tensor, prompt_len = self._prepare_cfm_loop_state(
prompt, cond, n_timesteps, generator=generator
)
trunk_dtype = self.diffusion_trunk_dtype
step_scale = 1.0 / n_timesteps
for i in range(n_timesteps):
xt_input = torch.cat([prompt, xt], dim=1)
diffusion_step = torch.full(
(xt.shape[0],),
fill_value=(i + 0.5) * step_scale,
device=xt.device,
dtype=trunk_dtype,
)
flow_pred = self.predict_noise_maybe_with_cfg(
do_true_cfg=(cfg > 0),
true_cfg_scale=cfg,
positive_kwargs={
"x": xt_input,
"diffusion_step": diffusion_step,
"cond": cond_emb,
"x_mask": xt_mask,
"prompt_len": prompt_len,
},
negative_kwargs={
"x": xt,
"diffusion_step": diffusion_step,
# CFG null cond: zero out target frames only (prompt cond stays in pos branch).
"cond": torch.zeros_like(cond_emb)[:, : xt.shape[1], :],
"x_mask": x_mask,
},
cfg_normalize=False,
)
xt = xt + flow_pred * h_tensor
return xt
def _forward_batch_from_request(
self,
req: OmniDiffusionRequest,
*,
kind: str,
metadata_key: str,
infer_batch_fn: Callable[..., tuple[torch.Tensor, int]],
prepare_extra_args: Callable[[dict[str, Any], Any], dict[str, Any]] | None = None,
) -> DiffusionOutput:
sampling_params = req.sampling_params
extra_args = dict(getattr(sampling_params, "extra_args", None) or {})
if prepare_extra_args is not None:
extra_args = prepare_extra_args(extra_args, sampling_params)
sampling_params.extra_args = extra_args
num_inference_steps = sampling_params.num_inference_steps or 32
guidance_scale = sampling_params.guidance_scale if sampling_params.guidance_scale_provided else 3.0
generator = self._resolve_diffusion_generator(sampling_params)
with self._stage_timer("consume_payload"):
payload = consume_payload(req, kind, self.device)
with self._stage_timer("infer_segments"):
audio, pitch_shift = infer_batch_fn(
payload,
extra_args=extra_args,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
generator=generator,
)
return DiffusionOutput(
output={
"payload": {"audio": audio},
"metadata": {"audio": {metadata_key: pitch_shift}},
},
stage_durations=self._profiler_stage_durations() or {},
)
vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py
-411
View File
@@ -1,411 +0,0 @@
"""SoulX-Singer SVC (voice conversion) pipeline implementation."""
import os
from collections.abc import Iterable
from typing import Any, ClassVar
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from vllm.logger import init_logger
from vllm.utils.torch_utils import set_default_torch_dtype
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.models.soulx_singer.modules import (
CFMDecoder,
WhisperEncoder,
)
from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import (
FlowMatchingAudioPipeline,
convert_soulx_audio_output_to_numpy,
)
from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import (
SOULX_PREPROCESSED_KEY,
get_soulx_preprocessed_payload,
has_precomputed,
)
from vllm_omni.diffusion.models.soulx_singer.preprocess.pre_process import (
attach_preprocess_for_diffusion_request,
is_warmup_request,
resolve_preprocess_audio,
)
from vllm_omni.diffusion.models.soulx_singer.utils import (
f0_to_coarse,
load_config,
resolve_pitch_shift,
validate_soulx_extra_args,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
logger = init_logger(__name__)
def get_soulxsinger_svc_pre_process_func(od_config: OmniDiffusionConfig):
hf_config = load_config(os.path.join(od_config.model, "config.yaml"))
sample_rate = hf_config.audio.sample_rate
device = get_local_device()
_pipeline = None
def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
nonlocal _pipeline
extra_args = validate_soulx_extra_args(
"svc",
dict(getattr(request.sampling_params, "extra_args", None) or {}),
)
# Inline build: when no warmup/no precomputed paths/no IPC payload,
# build the preprocess payload directly from audio file paths.
if not (is_warmup_request(request) or has_precomputed(extra_args, "svc")):
prompt = request.prompt
if not isinstance(prompt, str) and not get_soulx_preprocessed_payload(prompt): # type: ignore[arg-type]
prompt_audio, target_audio = resolve_preprocess_audio(prompt, extra_args) # type: ignore[arg-type]
if prompt_audio is not None and target_audio is not None:
if _pipeline is None:
from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.pipeline import (
SoulXPreprocessPipeline,
)
_pipeline = SoulXPreprocessPipeline(
od_config=od_config,
vocal_sep=bool(extra_args.get("vocal_sep", False)),
verbose=bool(extra_args.get("preprocess_verbose", False)),
extra_args=extra_args,
)
payload = _pipeline.build_svc_payload_from_audio(
prompt_audio=prompt_audio,
target_audio=target_audio,
sample_rate=sample_rate,
device=device,
vocal_sep=extra_args.get("vocal_sep"),
)
payload.setdefault("kind", "svc")
prompt.setdefault("additional_information", {})[SOULX_PREPROCESSED_KEY] = payload # type: ignore[union-attr, assignment]
return attach_preprocess_for_diffusion_request(
request,
kind="svc",
sample_rate=sample_rate,
device=device,
)
return pre_process_func
def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig):
def post_process_func(audio: torch.Tensor):
return convert_soulx_audio_output_to_numpy(audio)
return post_process_func
class PipelineSoulXSingerSVC(FlowMatchingAudioPipeline):
"""SVC pipeline for the SoulX-Singer model."""
_encoder_modules: ClassVar[list[str]] = [
"whisper_encoder",
"mel",
"f0_encoder",
]
EXTRA_BODY_PARAMS: ClassVar[frozenset[str]] = frozenset(
{
"prompt_wav_path",
"target_wav_path",
"prompt_f0_path",
"target_f0_path",
"prompt_audio",
"target_audio",
"vocal_sep",
"preprocess_weights_dir",
"preprocess_verbose",
"auto_shift",
"pitch_shift",
}
)
EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]] = frozenset({"pitch_shift"})
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__(od_config=od_config, prefix=prefix)
self.f0_encoder = nn.Embedding(self.f0_bin, self.f0_dim)
self.cfm_decoder = CFMDecoder(self.flow_matching_config)
self.mel, self.vocoder = self._build_fp32_audio_modules(self.audio_config)
with set_default_torch_dtype(torch.float32):
self.whisper_encoder = WhisperEncoder(device=self.device)
self._setup_soulx_profiler()
def _encode_condition(self, *, whisper_features: torch.Tensor, f0_coarse: torch.Tensor) -> torch.Tensor:
cond = whisper_features + self.f0_encoder(f0_coarse)
return self._to_trunk_dtype(cond)[0]
def _encode_prompt_whisper_feature(self, prompt_wav: torch.Tensor) -> torch.Tensor:
trunk_dtype = self.f0_encoder.weight.dtype
with self._stage_timer("whisper"):
return self.whisper_encoder.encode(
prompt_wav,
sr=self.audio_config.sample_rate,
output_dtype=trunk_dtype,
)
_encode_target_whisper_feature = _encode_prompt_whisper_feature
def _infer_segment(
self,
*,
prompt_mel: torch.Tensor,
prompt_wav: torch.Tensor,
target_wav: torch.Tensor,
prompt_f0: torch.Tensor,
target_f0: torch.Tensor,
pitch_shift: int,
num_inference_steps: int,
guidance_scale: float,
prompt_feature: torch.Tensor | None = None,
generator: torch.Generator | None = None,
) -> torch.Tensor:
"""Single-segment SVC inference aligned with ``SoulXSingerSVC.infer_segment``."""
len_prompt_mel = prompt_mel.shape[1]
prompt_f0 = F.pad(prompt_f0, (0, 0, 0, max(0, len_prompt_mel - prompt_f0.shape[1])))[:, :len_prompt_mel]
f0_coarse_prompt = f0_to_coarse(prompt_f0)
f0_coarse_target = f0_to_coarse(target_f0, f0_shift=int(pitch_shift * 5))
f0_coarse = torch.cat([f0_coarse_prompt, f0_coarse_target], dim=1)
trunk_dtype = self.f0_encoder.weight.dtype
if prompt_feature is None:
prompt_feature = self._encode_prompt_whisper_feature(prompt_wav)
with self._stage_timer("whisper"):
target_feature = self.whisper_encoder.encode(
target_wav,
sr=self.audio_config.sample_rate,
output_dtype=trunk_dtype,
)
prompt_feature = F.pad(
prompt_feature,
(0, 0, 0, max(0, f0_coarse_prompt.shape[1] - prompt_feature.shape[1])),
)[:, : f0_coarse_prompt.shape[1], :]
target_feature = F.pad(
target_feature,
(0, 0, 0, max(0, f0_coarse_target.shape[1] - target_feature.shape[1])),
)[:, : f0_coarse_target.shape[1], :]
whisper_features = torch.cat([prompt_feature, target_feature], dim=1)
with self._stage_timer("cond_encode"):
cond = self._encode_condition(whisper_features=whisper_features, f0_coarse=f0_coarse)
with self._stage_timer("cfm"):
generated_mel = self._run_flow_matching_loop(
prompt=prompt_mel,
cond=cond,
n_timesteps=num_inference_steps,
cfg=guidance_scale,
generator=generator,
)
with self._stage_timer("vocoder"):
generated_audio = self._mel_to_audio(generated_mel, squeeze=True)
target_len = target_wav.shape[-1]
if generated_audio.shape[-1] > target_len:
generated_audio = generated_audio[:target_len]
elif generated_audio.shape[-1] < target_len:
generated_audio = F.pad(generated_audio, (0, target_len - generated_audio.shape[-1]))
return generated_audio
def _build_vocal_segments(
self,
f0: torch.Tensor,
*,
hop_size: int,
sample_rate: int,
uv_frames_th: int = 10,
min_duration_sec: float = 15.0,
max_duration_sec: float = 30.0,
ignore_silent: bool = True,
) -> tuple[list[tuple[float, float]], list[tuple[float, float]]]:
"""Build vocal segments based on F0 contour.
Mirrors upstream ``SoulXSingerSVC.build_vocal_segments``.
Returns (overlap_segments, segments) tuples of (start_sec, end_sec).
"""
f0_np = (
f0.detach().float().cpu().numpy().squeeze() if isinstance(f0, torch.Tensor) else np.asarray(f0).squeeze()
)
total_frames = int(f0_np.shape[0])
if total_frames == 0:
return [], []
f0_rate = sample_rate // hop_size
min_frames = max(1, int(round(min_duration_sec * f0_rate)))
max_frames = max(1, int(round(max_duration_sec * f0_rate)))
split_points = [0]
def append_split_point(point: int) -> None:
point = int(max(0, min(point, total_frames)))
while point - split_points[-1] > max_frames:
split_points.append(split_points[-1] + max_frames)
if point > split_points[-1]:
split_points.append(point)
idx = 0
while idx < total_frames:
if f0_np[idx] == 0:
run_start = idx
while idx < total_frames and f0_np[idx] == 0:
idx += 1
run_end = idx
if (run_end - run_start) >= uv_frames_th:
split_point = max(run_end - 5, (run_start + run_end) // 2)
append_split_point(split_point)
else:
idx += 1
append_split_point(total_frames)
segments: list[tuple[int, int]] = []
overlap_segments: list[tuple[int, int]] = []
num_overlaps = 1
def append_segment(start_idx: int, end_idx: int, overlaps: int = num_overlaps) -> None:
segments.append((split_points[start_idx], split_points[end_idx]))
overlap_start_idx = start_idx
if start_idx > 0 and (split_points[end_idx] - split_points[start_idx - overlaps]) <= max_frames:
overlap_start_idx = start_idx - overlaps
overlap_segments.append((split_points[overlap_start_idx], split_points[end_idx]))
seg_start, seg_end = 0, 1
while seg_start < len(split_points) - 1:
while seg_end < len(split_points) and (split_points[seg_end] - split_points[seg_start]) < min_frames:
seg_end += 1
if seg_end >= len(split_points):
append_segment(seg_start, len(split_points) - 1)
break
append_segment(seg_start, seg_end)
seg_start = seg_end
seg_end = seg_start + 1
if ignore_silent:
filtered_idx = []
for i, (ov_start, ov_end) in enumerate(overlap_segments):
voice_ratio = np.sum(f0_np[ov_start:ov_end] > 0) / max(1, ov_end - ov_start)
voiced_frames = np.sum(f0_np[ov_start:ov_end] > 0)
if voice_ratio > 0.05 and voiced_frames >= 10:
filtered_idx.append(i)
overlap_segments = [overlap_segments[i] for i in filtered_idx]
segments = [segments[i] for i in filtered_idx]
# convert indices to seconds
overlap_sec = [(s / f0_rate, e / f0_rate) for s, e in overlap_segments]
seg_sec = [(s / f0_rate, e / f0_rate) for s, e in segments]
return overlap_sec, seg_sec
def infer_svc_batch(
self,
payload: dict[str, Any],
*,
extra_args: dict[str, Any],
num_inference_steps: int,
guidance_scale: float,
generator: torch.Generator | None = None,
) -> tuple[torch.Tensor, int]:
"""Batch SVC with chunked inference for long audio (matches upstream)."""
prompt_wav = payload["prompt_wav"]
target_wav = payload["target_wav"]
prompt_f0 = payload["prompt_f0"]
target_f0 = payload["target_f0"]
pitch_shift = resolve_pitch_shift(
auto_shift=bool(extra_args.get("auto_shift", True)),
manual_shift=int(extra_args.get("pitch_shift", 0)),
prompt_f0=prompt_f0,
target_f0=target_f0,
)
prompt_mel = self._mel_from_wav(prompt_wav)
# Long audio: chunk by vocal segments (same as upstream)
max_sec = 30.0 * self.audio_config.sample_rate
if target_wav.shape[-1] < max_sec:
# Short audio: one-shot
generated_audio = self._infer_segment(
prompt_mel=prompt_mel,
prompt_wav=prompt_wav,
target_wav=target_wav,
prompt_f0=prompt_f0,
target_f0=target_f0,
pitch_shift=pitch_shift,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
generator=generator,
)
return generated_audio.unsqueeze(0), pitch_shift
# Build vocal segments
overlap_segments, segments = self._build_vocal_segments(
target_f0,
hop_size=self.audio_config.hop_size,
sample_rate=self.audio_config.sample_rate,
)
if len(segments) == 0:
segments = [(0.0, target_wav.shape[-1] / self.audio_config.sample_rate)]
overlap_segments = [(0.0, target_wav.shape[-1] / self.audio_config.sample_rate)]
f0_rate = self.audio_config.sample_rate // self.audio_config.hop_size
generated_audio = torch.zeros_like(target_wav)
for idx in range(len(segments)):
overlap_start_sec, overlap_end_sec = overlap_segments[idx]
seg_start_sec, seg_end_sec = segments[idx]
wav_start = int(round(overlap_start_sec * self.audio_config.sample_rate))
wav_end = int(round(overlap_end_sec * self.audio_config.sample_rate))
f0_start = int(round(overlap_start_sec * f0_rate))
f0_end = int(round(overlap_end_sec * f0_rate))
wav_start = max(0, min(wav_start, target_wav.shape[-1]))
wav_end = max(wav_start, min(wav_end, target_wav.shape[-1]))
f0_start = max(0, min(f0_start, target_f0.shape[-1]))
f0_end = max(f0_start, min(f0_end, target_f0.shape[-1]))
segment_gt_wav = target_wav[:, wav_start:wav_end]
segment_gt_f0 = target_f0[:, f0_start:f0_end]
segment_generated = self._infer_segment(
prompt_mel=prompt_mel,
prompt_wav=prompt_wav,
target_wav=segment_gt_wav,
prompt_f0=prompt_f0,
target_f0=segment_gt_f0,
pitch_shift=pitch_shift,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
generator=generator,
)
seg_start = int(round(seg_start_sec * self.audio_config.sample_rate))
seg_end = int(round(seg_end_sec * self.audio_config.sample_rate))
segment_generated = segment_generated[seg_start - wav_start : seg_end - wav_start]
generated_audio[:, seg_start:seg_end] = segment_generated
return generated_audio.unsqueeze(0), pitch_shift
@torch.inference_mode()
def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
return self._forward_batch_from_request(
req,
kind="svc",
metadata_key="pitch_shift",
infer_batch_fn=self.infer_svc_batch,
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
del weights
self._load_soulx_checkpoint("model-svc.pt")
return {name for name, _ in self.named_parameters()}
vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py
-355
View File
@@ -1,355 +0,0 @@
"""SoulX-Singer SVS (score-driven) pipeline implementation."""
from collections.abc import Iterable
from typing import Any, ClassVar
import torch
import torch.nn as nn
import torch.nn.functional as F
from vllm.logger import init_logger
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.models.soulx_singer.modules import (
CFMDecoder,
ConvNeXtV2Block,
)
from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import (
FlowMatchingAudioPipeline,
convert_soulx_audio_output_to_numpy,
)
from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import (
SOULX_PREPROCESSED_KEY,
get_soulx_preprocessed_payload,
has_precomputed,
)
from vllm_omni.diffusion.models.soulx_singer.preprocess.pre_process import (
attach_preprocess_for_diffusion_request,
build_metadata_processor,
is_warmup_request,
normalize_svs_control_extra_args,
resolve_preprocess_audio,
)
from vllm_omni.diffusion.models.soulx_singer.utils import (
f0_to_coarse,
resolve_pitch_shift,
validate_soulx_extra_args,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
logger = init_logger(__name__)
def get_soulxsinger_pre_process_func(od_config: OmniDiffusionConfig):
"""Validate/load SVS preprocess payload for single-stage or stage-1 DiT."""
metadata_processor = build_metadata_processor(od_config)
_pipeline = None
def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
nonlocal _pipeline
extra_args = validate_soulx_extra_args(
"svs",
dict(getattr(request.sampling_params, "extra_args", None) or {}),
)
# Inline build: when no warmup/no precomputed paths/no IPC payload,
# build the preprocess payload directly from audio file paths.
if not (is_warmup_request(request) or has_precomputed(extra_args, "svs")):
prompt = request.prompt
if not isinstance(prompt, str) and not get_soulx_preprocessed_payload(prompt): # type: ignore[arg-type]
prompt_audio, target_audio = resolve_preprocess_audio(prompt, extra_args) # type: ignore[arg-type]
if prompt_audio is not None and target_audio is not None:
if _pipeline is None:
from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.pipeline import (
SoulXPreprocessPipeline,
)
_pipeline = SoulXPreprocessPipeline(
od_config=od_config,
vocal_sep=bool(extra_args.get("vocal_sep", False)),
verbose=bool(extra_args.get("preprocess_verbose", False)),
extra_args=extra_args,
)
payload = _pipeline.build_svs_payload_from_audio(
prompt_audio=prompt_audio,
target_audio=target_audio,
metadata_processor=metadata_processor,
language=str(extra_args.get("language", "Mandarin")),
vocal_sep=extra_args.get("vocal_sep"),
prompt_vocal_sep=extra_args.get("prompt_vocal_sep"),
target_vocal_sep=extra_args.get("target_vocal_sep"),
prompt_max_merge_duration_ms=extra_args.get("prompt_max_merge_duration"),
target_max_merge_duration_ms=extra_args.get("target_max_merge_duration"),
)
payload.setdefault("kind", "svs")
prompt.setdefault("additional_information", {})[SOULX_PREPROCESSED_KEY] = payload # type: ignore[union-attr, assignment]
return attach_preprocess_for_diffusion_request(
request,
kind="svs",
metadata_processor=metadata_processor,
)
return pre_process_func
def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig):
"""Convert pipeline audio tensor output for offline consumers."""
def post_process_func(audio: torch.Tensor):
return convert_soulx_audio_output_to_numpy(audio)
return post_process_func
def _expand_states(h: torch.Tensor, mel2token: torch.Tensor) -> torch.Tensor:
if mel2token.max() > h.size(1) - 1:
logger.warning(
"mel2token.max() (%s) is greater than h.size(1) - 1 (%s); clamping.",
mel2token.max(),
h.size(1) - 1,
)
mel2token = torch.clamp(mel2token, 0, h.size(1) - 1)
mel2token_ = mel2token[..., None].repeat([1, 1, h.shape[-1]])
return torch.gather(h, 1, mel2token_)
class PipelineSoulXSingerSVS(FlowMatchingAudioPipeline):
"""Pipeline for the SoulX-Singer model (SoulX-Singer)."""
_encoder_modules: ClassVar[list[str]] = [
"mel",
"f0_encoder",
"preflow",
"note_text_encoder",
"note_pitch_encoder",
"note_type_encoder",
]
EXTRA_BODY_PARAMS: ClassVar[frozenset[str]] = frozenset(
{
"prompt_metadata_path",
"target_metadata_path",
"audio_path",
"prompt_audio",
"target_audio",
"language",
"vocal_sep",
"max_merge_duration",
"preprocess_weights_dir",
"preprocess_verbose",
"control",
"auto_shift",
"pitch_shift",
}
)
EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]] = frozenset({"f0_shift"})
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__(od_config=od_config, prefix=prefix)
self.metadata_processor = build_metadata_processor(od_config)
self.f0_encoder = nn.Embedding(self.f0_bin, self.f0_dim)
self.preflow = nn.Sequential(
*[
ConvNeXtV2Block(
self.text_dim,
self.text_dim * 2,
)
for _ in range(self.encoder_config.num_layers)
]
)
self.cfm_decoder = CFMDecoder(self.flow_matching_config)
self.mel, self.vocoder = self._build_fp32_audio_modules(self.audio_config)
self.note_text_encoder = nn.Embedding(self.vocab_size, self.text_dim)
self.note_pitch_encoder = nn.Embedding(256, self.pitch_dim)
self.note_type_encoder = nn.Embedding(256, self.type_dim)
self._setup_soulx_profiler()
def _encode_condition(
self,
*,
note_pitch: torch.Tensor,
note_type: torch.Tensor,
note_text: torch.Tensor,
mel2note: torch.Tensor,
f0_coarse: torch.Tensor,
) -> torch.Tensor:
features = (
self.note_pitch_encoder(note_pitch) + self.note_type_encoder(note_type) + self.note_text_encoder(note_text)
)
features = self.preflow(features)
features = _expand_states(features, mel2note)
features = features + self.f0_encoder(f0_coarse)
return self._to_trunk_dtype(features)[0]
def _infer_svs_segment(
self,
prompt_meta: dict,
target_meta: dict,
*,
pitch_shift: int,
num_inference_steps: int,
guidance_scale: float,
prompt_mel: torch.Tensor | None = None,
generator: torch.Generator | None = None,
) -> torch.Tensor:
prompt_note_text = prompt_meta["phoneme"]
prompt_mel2note = prompt_meta["mel2note"]
prompt_note_type = prompt_meta["note_type"]
prompt_note_pitch = prompt_meta["note_pitch"]
prompt_f0 = prompt_meta["f0"]
target_note_text = target_meta["phoneme"]
target_mel2note = target_meta["mel2note"]
target_note_type = target_meta["note_type"]
target_note_pitch = target_meta["note_pitch"]
target_f0 = target_meta["f0"]
if target_f0 is None or prompt_f0 is None:
target_f0 = torch.zeros_like(target_mel2note).float()
prompt_f0 = torch.zeros_like(prompt_mel2note).float()
if target_note_pitch is None or prompt_note_pitch is None:
target_note_pitch = torch.zeros_like(target_note_type).int()
prompt_note_pitch = torch.zeros_like(prompt_note_type).int()
len_prompt_note = prompt_note_pitch.shape[1]
len_prompt_mel = prompt_f0.shape[1] if prompt_f0 is not None else prompt_mel2note.shape[1]
note_pitch = torch.cat([prompt_note_pitch, target_note_pitch], dim=1)
note_text = torch.cat([prompt_note_text, target_note_text], dim=1)
note_type = torch.cat([prompt_note_type, target_note_type], dim=1)
# Target note indices follow prompt notes in the concatenated score.
mel2note = torch.cat([prompt_mel2note, target_mel2note + len_prompt_note], dim=1)
# pitch_shift is semitones; each coarse F0 bin is 20 cents (×5 bins per semitone).
f0_coarse_prompt = f0_to_coarse(prompt_f0)
f0_coarse_target = f0_to_coarse(target_f0, f0_shift=pitch_shift * 5)
f0_coarse = torch.cat([f0_coarse_prompt, f0_coarse_target], dim=1)
note_pitch = note_pitch.clone()
note_pitch[note_pitch > 0] = note_pitch[note_pitch > 0] + pitch_shift
note_pitch = torch.clamp(note_pitch, 0, 255)
if prompt_mel is None:
prompt_wav = prompt_meta["wav"]
with self._stage_timer("mel"):
prompt_mel = self._mel_from_wav(prompt_wav)
if prompt_mel.shape[1] > len_prompt_mel:
prompt_mel = prompt_mel[:, :len_prompt_mel, :]
elif prompt_mel.shape[1] < len_prompt_mel:
logger.warning(
"prompt_mel length %s is shorter than metadata frames %s; padding mel.",
prompt_mel.shape[1],
len_prompt_mel,
)
prompt_mel = F.pad(prompt_mel, (0, 0, 0, len_prompt_mel - prompt_mel.shape[1]))
with self._stage_timer("cond_encode"):
cond = self._encode_condition(
note_pitch=note_pitch,
note_type=note_type,
note_text=note_text,
mel2note=mel2note,
f0_coarse=f0_coarse,
)
with self._stage_timer("cfm"):
generated_mel = self._run_flow_matching_loop(
prompt_mel,
cond,
num_inference_steps,
guidance_scale,
generator=generator,
)
with self._stage_timer("vocoder"):
return self._mel_to_audio(generated_mel)
def infer_svs_batch(
self,
payload: dict[str, Any],
*,
extra_args: dict[str, Any],
num_inference_steps: int,
guidance_scale: float,
generator: torch.Generator | None = None,
) -> tuple[torch.Tensor, int]:
"""Pure batch (non-streaming): run diffusion for each target segment in order and concat."""
prompt_meta_raw = payload.get("prompt_meta")
target_meta_list = payload["target_meta_list"]
control = str(extra_args.get("control", "score"))
prompt_meta = None
prompt_mel = None
if prompt_meta_raw:
prompt_meta = self._ensure_processed_meta(prompt_meta_raw)
prompt_meta = self._apply_control_mode(prompt_meta, control)
if prompt_meta and prompt_meta.get("wav") is not None:
prompt_mel = self._mel_from_wav(prompt_meta["wav"])
pieces: list[torch.Tensor] = []
pitch_shift = 0
for idx, target_raw in enumerate(target_meta_list):
target_meta = self._ensure_processed_meta(target_raw)
target_meta = self._apply_control_mode(target_meta, control)
if idx == 0:
pitch_shift = resolve_pitch_shift(
auto_shift=bool(extra_args.get("auto_shift", True)),
manual_shift=int(extra_args.get("pitch_shift", 0)),
prompt_f0=prompt_meta["f0"] if prompt_meta else None,
target_f0=target_meta["f0"],
prompt_note_pitch=prompt_meta["note_pitch"] if prompt_meta else None,
target_note_pitch=target_meta["note_pitch"],
)
seg = self._infer_svs_segment(
prompt_meta or {},
target_meta,
pitch_shift=pitch_shift,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
prompt_mel=prompt_mel,
generator=generator,
)
pieces.append(seg.squeeze())
if not pieces:
return torch.zeros(1, 0, device=self.device, dtype=torch.float32), pitch_shift
full = torch.cat(pieces, dim=-1)
return full.unsqueeze(0), pitch_shift
@staticmethod
def _apply_control_mode(meta: dict, control: str) -> dict:
meta["note_pitch"] = meta["note_pitch"] if control == "score" else None
meta["f0"] = meta["f0"] if control == "melody" else None
return meta
def _ensure_processed_meta(self, meta: dict) -> dict:
if isinstance(meta.get("phoneme"), torch.Tensor):
return meta
return self.metadata_processor.process(meta, None)
@torch.inference_mode()
def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
return self._forward_batch_from_request(
req,
kind="svs",
metadata_key="f0_shift",
infer_batch_fn=self.infer_svs_batch,
prepare_extra_args=lambda extra_args, _sampling: normalize_svs_control_extra_args(extra_args),
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
del weights
self._load_soulx_checkpoint("model.pt")
return {name for name, _ in self.named_parameters()}
vllm_omni/diffusion/models/soulx_singer/preprocess/__init__.py
-1
View File
@@ -1 +0,0 @@
"""Native SoulX-Singer preprocess helpers."""
vllm_omni/diffusion/models/soulx_singer/preprocess/g2p.py
-86
View File
@@ -1,86 +0,0 @@
import re
_EN_WORD_RE = re.compile(r"^[A-Za-z]+(?:'[A-Za-z]+)*$")
_ZH_WORD_RE = re.compile(r"[\u4e00-\u9fff]")
EN_FLAG = "en_"
YUE_FLAG = "yue_"
ZH_FLAG = "zh_"
_g2p_zh = None
_g2p_en = None
def _ensure_g2p_backends() -> None:
global _g2p_zh, _g2p_en
if _g2p_zh is None:
from g2pM import G2pM
_g2p_zh = G2pM()
if _g2p_en is None:
from g2p_en import G2p as G2pE
_g2p_en = G2pE()
def is_chinese_char(word: str) -> bool:
if len(word) != 1:
return False
return bool(_ZH_WORD_RE.fullmatch(word))
def is_english_word(word: str) -> bool:
if not word:
return False
return bool(_EN_WORD_RE.fullmatch(word))
def g2p_cantonese(sent):
import ToJyutping
return ToJyutping.get_jyutping_list(sent)
def g2p_mandarin(sent):
_ensure_g2p_backends()
return _g2p_zh(sent, tone=True, char_split=False)
def g2p_english(word):
_ensure_g2p_backends()
return _g2p_en(word)
def g2p_transform(words, lang):
zh_words = []
transformed_words = [0] * len(words)
for idx, w in enumerate(words):
if w == "<SP>":
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 = "<SP>"
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
vllm_omni/diffusion/models/soulx_singer/preprocess/metadata_utils.py
-137
View File
@@ -1,137 +0,0 @@
"""Metadata merge helpers for SoulX-Singer preprocess."""
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import numpy as np
import soundfile as sf
@dataclass
class SegmentMetadata:
item_name: str
wav_fn: str
language: str
start_time_ms: int
end_time_ms: int
note_text: list[str]
note_dur: list[float]
note_pitch: list[int]
note_type: list[int]
origin_wav_fn: str | None = None
wav: np.ndarray | None = None
def get(self, name: str, default: Any = None) -> Any:
return getattr(self, name, default)
def _merge_group(
audio: np.ndarray,
sample_rate: int,
segments: list[SegmentMetadata | dict[str, Any]],
output_dir: Path | None,
*,
end_extension_ms: int = 0,
) -> SegmentMetadata:
if not segments:
raise ValueError("segments must not be empty")
words: list[str] = []
durs: list[float] = []
pitches: list[int] = []
types: list[int] = []
for i, seg in enumerate(segments):
if i > 0:
prev_seg = segments[i - 1]
gap_ms = seg.get("start_time_ms", 0) - prev_seg.get("end_time_ms", 0)
if gap_ms > 0:
words.append("<SP>")
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("<SP>")
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 == "<SP>" and merged_words[-1] == "<SP>":
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),
}
vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py
-106
View File
@@ -1,106 +0,0 @@
"""SoulX-Singer payload schema (guidance/precomputed tables, dummy builders)."""
from __future__ import annotations
from typing import Any
import torch
SOULX_PRECOMPUTED_KEYS_BY_KIND: dict[str, tuple[str, ...]] = {
"svs": ("prompt_metadata_path", "target_metadata_path", "audio_path"),
"svc": ("prompt_wav_path", "target_wav_path", "prompt_f0_path", "target_f0_path"),
}
# Key under which the SoulX preprocess payload is stored in
# prompt["additional_information"] or diffusion output metadata.
SOULX_PREPROCESSED_KEY = "soulx_preprocessed"
def get_soulx_preprocessed_payload(prompt: dict[str, Any]) -> dict[str, Any] | None:
"""Extract attached payload from a diffusion request prompt dict."""
additional = prompt.get("additional_information") or {}
payload = additional.get(SOULX_PREPROCESSED_KEY)
return payload if isinstance(payload, dict) else None
def relocate_tensors(payload: dict[str, Any], device: torch.device | str) -> dict[str, Any]:
"""Recursively move all tensors inside the payload to the target device.
Called after decode when the payload is about to be consumed by the
diffusion model (which may live on a different device / rank).
"""
if isinstance(payload, torch.Tensor):
return payload.to(device, non_blocking=True)
if isinstance(payload, dict):
return {key: relocate_tensors(value, device) for key, value in payload.items()}
if isinstance(payload, list):
return [relocate_tensors(item, device) for item in payload]
return payload
def consume_payload(
req: Any,
expected_kind: str,
device: torch.device | str,
) -> dict[str, Any]:
"""Read the payload attached by pre_process_func, validate the kind,
and relocate tensors to the compute device.
"""
prompts = getattr(req, "prompts", None) or []
if not prompts:
raise ValueError(f"SoulX-Singer {expected_kind} request has no prompts.")
prompt = prompts[0]
if isinstance(prompt, str):
raise ValueError(f"SoulX-Singer {expected_kind} forward requires pre_process_func output on the prompt.")
payload = get_soulx_preprocessed_payload(prompt)
if payload is None or payload.get("kind") != expected_kind:
raise ValueError(
f"SoulX-Singer {expected_kind} forward requires additional_information['soulx_preprocessed'] "
"produced by pre_process_func."
)
return relocate_tensors(payload, device)
def has_precomputed(extra_args: dict[str, Any], kind: str) -> bool:
keys = SOULX_PRECOMPUTED_KEYS_BY_KIND[kind]
return all(extra_args.get(key) for key in keys)
def build_dummy_payload(kind: str, device: torch.device) -> dict[str, Any]:
mel_frames = 4
hop_size = 480
n_fft = 1920
wav_samples = max(mel_frames * hop_size, n_fft)
if kind == "svc":
voiced_f0 = torch.full((1, mel_frames), 100.0, device=device, dtype=torch.float64)
return {
"kind": "svc",
"prompt_wav": torch.zeros(1, wav_samples, device=device, dtype=torch.float32),
"target_wav": torch.zeros(1, wav_samples, device=device, dtype=torch.float32),
"prompt_f0": voiced_f0.clone(),
"target_f0": voiced_f0.clone(),
}
return {
"kind": "svs",
"prompt_meta": {
"phoneme": torch.zeros(1, 2, device=device, dtype=torch.long),
"note_pitch": torch.zeros(1, mel_frames, device=device, dtype=torch.long),
"note_type": torch.zeros(1, mel_frames, device=device, dtype=torch.long),
"mel2note": torch.zeros(1, mel_frames, device=device, dtype=torch.long),
"f0": torch.zeros(1, mel_frames, device=device, dtype=torch.float32),
"wav": torch.zeros(1, wav_samples, device=device, dtype=torch.float32),
},
"target_meta_list": [
{
"text": "dummy",
"phoneme": " ".join(["<SP>"] * 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)],
}
],
}
vllm_omni/diffusion/models/soulx_singer/preprocess/pre_process.py
-223
View File
@@ -1,223 +0,0 @@
"""Shared ``pre_process_func`` helpers for SoulX-Singer pipelines (batch / full-payload only)."""
from pathlib import Path
from typing import Any
import torch
import torch.nn as nn
from vllm.logger import init_logger
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.pipeline import SoulXPreprocessPipeline
from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import (
SOULX_PRECOMPUTED_KEYS_BY_KIND,
SOULX_PREPROCESSED_KEY,
build_dummy_payload,
get_soulx_preprocessed_payload,
has_precomputed,
)
from vllm_omni.diffusion.models.soulx_singer.utils import (
MetadataProcessor,
load_config,
resolve_phoneset_path,
validate_soulx_extra_args,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.inputs.data import OmniTextPrompt
logger = init_logger(__name__)
def is_warmup_request(request: OmniDiffusionRequest) -> bool:
return request.is_dummy_run()
def resolve_preprocess_audio(prompt: dict[str, Any], extra_args: dict[str, Any]) -> tuple[Any, Any]:
mm = prompt.get("multi_modal_data") or {}
prompt_audio = mm.get("audio") if isinstance(mm, dict) else None
if prompt_audio is None:
prompt_audio = extra_args.get("prompt_audio")
return prompt_audio, extra_args.get("target_audio")
def _preprocess_inputs_missing_error(kind: str) -> ValueError:
return ValueError(
f"SoulX-Singer {kind} requires precomputed paths "
f"{list(SOULX_PRECOMPUTED_KEYS_BY_KIND[kind])}, an attached soulx_preprocessed payload "
f"(multi-stage preprocess stage), or run with pipeline soulxsinger_{kind}."
)
def build_precomputed_payload(
kind: str,
extra_args: dict[str, Any],
*,
metadata_processor,
sample_rate: int | None,
device,
) -> dict[str, Any]:
if kind == "svs":
return SoulXPreprocessPipeline.build_svs_payload_from_paths(
prompt_metadata_path=str(extra_args["prompt_metadata_path"]),
target_metadata_path=str(extra_args["target_metadata_path"]),
audio_path=str(extra_args["audio_path"]),
metadata_processor=metadata_processor,
)
if sample_rate is None or device is None:
raise ValueError("svc precomputed payload requires sample_rate and device")
return SoulXPreprocessPipeline.build_svc_payload_from_paths( # type: ignore
prompt_wav_path=str(extra_args["prompt_wav_path"]),
target_wav_path=str(extra_args["target_wav_path"]),
prompt_f0_path=str(extra_args["prompt_f0_path"]),
target_f0_path=str(extra_args["target_f0_path"]),
sample_rate=int(sample_rate),
device=device,
)
def build_warmup_payload(
kind: str,
*,
metadata_processor,
device,
sample_rate: int | None,
) -> dict[str, Any]:
if kind == "svs":
if metadata_processor is None:
raise ValueError("SVS warmup requires metadata_processor")
payload = build_dummy_payload("svs", torch.device("cpu"))
dummy_prompt = payload["prompt_meta"]
processed = metadata_processor.process(dict(payload["target_meta_list"][0]))
payload["prompt_meta"] = {
key: value.clone() if isinstance(value, torch.Tensor) else value for key, value in processed.items()
}
if isinstance(dummy_prompt.get("wav"), torch.Tensor):
payload["prompt_meta"]["wav"] = dummy_prompt["wav"].clone()
return payload
if device is None or sample_rate is None:
raise ValueError("svc warmup requires device and sample_rate")
return build_dummy_payload("svc", device)
def build_preprocess_payload(
kind: str,
*,
prompt: dict[str, Any],
extra_args: dict[str, Any],
preprocess: nn.Module,
metadata_processor,
device,
sample_rate: int,
) -> dict[str, Any]:
if has_precomputed(extra_args, kind):
return build_precomputed_payload(
kind,
extra_args,
metadata_processor=metadata_processor,
sample_rate=sample_rate,
device=device,
)
prompt_audio, target_audio = resolve_preprocess_audio(prompt, extra_args)
if prompt_audio is None or target_audio is None:
raise _preprocess_inputs_missing_error(kind)
if kind == "svc":
return preprocess.build_svc_payload_from_audio( # type: ignore
prompt_audio=prompt_audio,
target_audio=target_audio,
sample_rate=sample_rate,
device=device,
vocal_sep=extra_args.get("vocal_sep"),
)
return preprocess.build_svs_payload_from_audio( # type: ignore
prompt_audio=prompt_audio,
target_audio=target_audio,
metadata_processor=metadata_processor,
language=str(extra_args.get("language", "Mandarin")),
vocal_sep=extra_args.get("vocal_sep"),
prompt_vocal_sep=extra_args.get("prompt_vocal_sep"),
target_vocal_sep=extra_args.get("target_vocal_sep"),
prompt_max_merge_duration_ms=extra_args.get("prompt_max_merge_duration"),
target_max_merge_duration_ms=extra_args.get("target_max_merge_duration"),
)
def attach_preprocess_for_diffusion_request(
request: OmniDiffusionRequest,
*,
kind: str,
metadata_processor=None,
sample_rate: int | None = None,
device=None,
) -> OmniDiffusionRequest:
"""Resolve preprocess payload for stage-1 diffusion (warmup / IPC / precomputed paths)."""
extra_args = validate_soulx_extra_args(
kind,
dict(getattr(request.sampling_params, "extra_args", None) or {}),
)
prompt = request.prompt
prompt = OmniTextPrompt(prompt=prompt) if isinstance(prompt, str) else prompt
if is_warmup_request(request):
request.sampling_params.num_inference_steps = 1
payload = build_warmup_payload(
kind,
metadata_processor=metadata_processor,
device=device,
sample_rate=sample_rate,
)
elif get_soulx_preprocessed_payload(prompt):
request.prompt = prompt
if kind == "svs":
extra_args = normalize_svs_control_extra_args(extra_args)
request.sampling_params.extra_args = extra_args
return request
elif has_precomputed(extra_args, kind):
payload = build_precomputed_payload(
kind,
extra_args,
metadata_processor=metadata_processor,
sample_rate=sample_rate,
device=device,
)
else:
raise _preprocess_inputs_missing_error(kind)
if payload.get("kind") != kind:
raise ValueError(f"Invalid {kind} preprocess payload kind: {payload.get('kind')}")
prompt.setdefault("additional_information", {})[SOULX_PREPROCESSED_KEY] = payload
request.prompt = prompt
if kind == "svs":
extra_args = normalize_svs_control_extra_args(extra_args)
request.sampling_params.extra_args = extra_args
return request
def normalize_svs_control_extra_args(extra_args: dict[str, Any]) -> dict[str, Any]:
control = extra_args.get("control")
if control is None:
control = "score"
logger.warning("control is not provided, using 'score' as default")
elif control not in ("score", "melody"):
raise ValueError(f"Invalid control: {control}. Must be one of: ['score', 'melody']")
extra_args["control"] = control
return extra_args
def build_metadata_processor(od_config: OmniDiffusionConfig):
model_dir = od_config.model
config_path = Path(model_dir) / "config.yaml"
hf_config = load_config(str(config_path))
audio_config = hf_config.audio
return MetadataProcessor(
hop_size=audio_config.hop_size,
sample_rate=audio_config.sample_rate,
phoneset_path=resolve_phoneset_path(model_dir),
device=str(get_local_device()),
)
vllm_omni/diffusion/models/soulx_singer/utils.py
-565
View File
@@ -1,565 +0,0 @@
import json
from pathlib import Path
from typing import Any, Literal
import numpy as np
import torch
import torchaudio
from huggingface_hub import snapshot_download
from omegaconf import DictConfig, OmegaConf
from vllm.logger import init_logger
from vllm_omni.diffusion.data import OmniDiffusionConfig
logger = init_logger(__name__)
_MODEL_WEIGHTS_REPO = "Soul-AILab/SoulX-Singer"
_PREPROCESS_WEIGHTS_REPO = "Soul-AILab/SoulX-Singer-Preprocess"
def _patch_torchaudio_load() -> None:
"""Borrowed from MOSS-TTS-Nano. Patch torchaudio.load to use soundfile
if torchcodec is unavailable.
"""
try:
import torchaudio
torchaudio # noqa
import torchcodec # noqa: F401
return
except Exception:
pass
import soundfile as sf
def _soundfile_load(path, frame_offset=0, num_frames=-1, normalize=True, channels_first=True, format=None):
data, sr = sf.read(str(path), dtype="float32", always_2d=True)
if frame_offset > 0:
data = data[frame_offset:]
if num_frames > 0:
data = data[:num_frames]
waveform = torch.from_numpy(data)
if channels_first:
waveform = waveform.T
return waveform, sr
def _soundfile_save(path, src, sample_rate, channels_first=True, **kwargs):
wav = src.detach().cpu().float().numpy()
if channels_first and wav.ndim == 2:
wav = wav.T
sf.write(str(path), wav, sample_rate)
try:
import torchaudio
torchaudio.load = _soundfile_load
torchaudio.save = _soundfile_save
logger.info("Patched torchaudio.load/save to use soundfile (torchcodec unavailable)")
except Exception as e:
logger.warning("Could not patch torchaudio: %s", e)
def load_config(config_path: str | Path) -> DictConfig:
"""
Load a configuration file and optionally merges it with a base configuration.
Args:
config_path (Path): Path to the configuration file.
"""
# Load the initial configuration from the given path
config = OmegaConf.load(str(config_path))
# Check if there is a base configuration specified and merge if necessary
if config.get("base_config", None) is not None:
base_config = OmegaConf.load(str(config.get("base_config")))
config = OmegaConf.merge(base_config, config)
return config
SoulXKind = Literal["svs", "svc"]
_SVS_ARCHITECTURE = "SoulXSingerPipeline"
_SVC_ARCHITECTURE = "SoulXSingerSVCPipeline"
_ARCHITECTURE_TO_KIND: dict[str, SoulXKind] = {
_SVS_ARCHITECTURE: "svs",
_SVC_ARCHITECTURE: "svc",
}
def load_model_config_json(model_dir: str | Path) -> dict[str, Any]:
"""Load ``config.json`` from a SoulX-Singer model view directory."""
cfg_path = Path(model_dir) / "config.json"
if not cfg_path.is_file():
raise FileNotFoundError(
f"SoulX-Singer config.json not found at {cfg_path}. "
"Create one with architectures ['SoulXSingerPipeline'] (SVS) or "
"['SoulXSingerSVCPipeline'] (SVC); see examples/offline_inference/text_to_speech/README.md."
)
with cfg_path.open(encoding="utf-8") as f:
return json.load(f)
def resolve_soulx_kind(model_dir: str | Path) -> SoulXKind:
"""Resolve SVS vs SVC from ``config.json`` ``architectures[0]``."""
cfg = load_model_config_json(model_dir)
archs = cfg.get("architectures") or []
if not archs:
raise ValueError(
f"SoulX-Singer config.json at {Path(model_dir) / 'config.json'} is missing 'architectures'. "
"Set architectures to ['SoulXSingerPipeline'] for SVS or ['SoulXSingerSVCPipeline'] for SVC."
)
arch = archs[0]
kind = _ARCHITECTURE_TO_KIND.get(arch)
if kind is None:
raise ValueError(
f"Unknown SoulX-Singer architecture {arch!r} in config.json. Expected one of {list(_ARCHITECTURE_TO_KIND)}."
)
return kind
def validate_soulx_extra_args(kind: SoulXKind, extra_args: dict[str, Any] | None) -> dict[str, Any]:
"""Validate request ``extra_args`` against the model mode from config.json."""
from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import (
SOULX_PRECOMPUTED_KEYS_BY_KIND,
has_precomputed,
)
extra_args = dict(extra_args or {})
declared = extra_args.pop("kind", None)
if declared is not None and declared != kind:
logger.warning(
"Ignoring extra_args kind=%r; model config.json declares %r.",
declared,
kind,
)
other_kind: SoulXKind = "svc" if kind == "svs" else "svs"
if has_precomputed(extra_args, other_kind):
raise ValueError(
f"SoulX-Singer config.json declares {kind!r}, but request supplies "
f"{other_kind} precomputed paths {list(SOULX_PRECOMPUTED_KEYS_BY_KIND[other_kind])}."
)
if kind == "svc":
for key in ("language", "control"):
if extra_args.get(key) is not None:
logger.warning(
"Ignoring SVS-only extra_arg %r=%r on an SVC model (config.json).",
key,
extra_args.pop(key),
)
for key in SOULX_PRECOMPUTED_KEYS_BY_KIND["svs"]:
if extra_args.get(key):
raise ValueError(f"SoulX-Singer SVC model (config.json) does not accept SVS precomputed key {key!r}.")
else:
for key in SOULX_PRECOMPUTED_KEYS_BY_KIND["svc"]:
if extra_args.get(key):
raise ValueError(f"SoulX-Singer SVS model (config.json) does not accept SVC precomputed key {key!r}.")
return extra_args
# ---------------- utils for model weights ----------------
def resolve_preprocess_weights_root(od_config: OmniDiffusionConfig) -> Path:
"""Locate preprocess weights on disk or download from Hugging Face."""
# default local dir
if Path(od_config.model).parent.is_dir():
local_dir = Path(od_config.model).parent / "SoulX-Singer-Preprocess"
else:
local_dir = Path(__file__).parents[5] / "pretrained" / "SoulX-Singer-Preprocess"
local_dir = local_dir.expanduser()
if local_dir.is_dir():
logger.info("Using SoulX preprocess weights from %s", local_dir)
return local_dir
logger.info(
"SoulX preprocess weights not found locally; downloading %s",
_PREPROCESS_WEIGHTS_REPO,
)
downloaded = snapshot_download(
_PREPROCESS_WEIGHTS_REPO,
allow_patterns=["*"],
local_dir=local_dir,
)
return Path(downloaded)
def preprocess_weight_paths(weights_root: Path) -> dict[str, str]:
"""Map upstream relative paths to absolute paths under ``weights_root``."""
root = weights_root.resolve()
return {
"rmvpe": str(root / "rmvpe/rmvpe.pt"),
"sep_ckpt": str(root / "mel-band-roformer-karaoke/mel_band_roformer_karaoke_becruily.ckpt"),
"sep_config": str(root / "mel-band-roformer-karaoke/config_karaoke_becruily.yaml"),
"asr_zh": str(root / "speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch"),
"asr_en": str(root / "parakeet-tdt-0.6b-v2/parakeet-tdt-0.6b-v2.nemo"),
"rosvot": str(root / "rosvot/rosvot/model.pt"),
}
def resolve_phoneset_path(model_dir: str) -> str:
"""Resolve phoneme vocabulary; must be present in model directory."""
candidates = (
Path(model_dir) / "phoneme" / "phone_set.json",
Path(model_dir) / "phone_set.json",
)
for path in candidates:
if path.is_file():
return str(path)
raise FileNotFoundError(
"SoulX-Singer phoneset not found. Expected one of: "
f"{[str(p) for p in candidates]}. "
"Copy phone_set.json into the model directory. "
"See `examples/offline_inference/text_to_speech/README.md` for instructions."
)
# ---------------- utils for data processing ----------------
def load_wav(wav_path: str, sample_rate: int) -> torch.Tensor:
"""Load wav file and resample to target sample rate.
Args:
wav_path (str): Path to wav file.
sample_rate (int): Target sample rate.
Returns:
torch.Tensor: Waveform tensor with shape (1, T).
"""
_patch_torchaudio_load()
waveform, sr = torchaudio.load(wav_path)
if sr != sample_rate:
waveform = torchaudio.functional.resample(waveform, sr, sample_rate)
if len(waveform.shape) > 1 and waveform.shape[0] > 1:
waveform = torch.mean(waveform, dim=0, keepdim=True)
return waveform
def f0_to_coarse(f0, f0_bin=361, f0_min=32.7031956625, f0_shift=0):
"""
Convert continuous F0 values to discrete F0 bins (SIL and C1 - B6, 361 bins).
args:
f0: continuous F0 values
f0_bin: number of F0 bins
f0_min: minimum F0 value
f0_shift: shift value for F0 bins
returns:
f0_coarse: discrete F0 bins
"""
is_torch = isinstance(f0, torch.Tensor)
uv_mask = f0 <= 0
if is_torch:
f0_min_tensor = f0.new_tensor(f0_min)
f0_safe = torch.maximum(f0, f0_min_tensor)
f0_cents = 1200 * torch.log2(f0_safe / f0_min_tensor)
else:
f0_safe = np.maximum(f0, f0_min)
f0_cents = 1200 * np.log2(f0_safe / f0_min)
f0_coarse = (f0_cents / 20) + 1
if is_torch:
f0_coarse = torch.round(f0_coarse).long()
f0_coarse = torch.clamp(f0_coarse, min=1, max=f0_bin - 1)
else:
f0_coarse = np.rint(f0_coarse).astype(int)
f0_coarse = np.clip(f0_coarse, 1, f0_bin - 1)
f0_coarse[uv_mask] = 0
if f0_shift != 0:
if is_torch:
voiced = f0_coarse > 0
if voiced.any():
shifted = f0_coarse[voiced] + f0_shift
f0_coarse[voiced] = torch.clamp(shifted, 1, f0_bin - 1)
else:
voiced = f0_coarse > 0
if np.any(voiced):
shifted = f0_coarse[voiced] + f0_shift
f0_coarse[voiced] = np.clip(shifted, 1, f0_bin - 1)
return f0_coarse
def resolve_pitch_shift(
*,
auto_shift: bool,
manual_shift: int,
prompt_f0: torch.Tensor | None = None,
target_f0: torch.Tensor | None = None,
prompt_note_pitch: torch.Tensor | None = None,
target_note_pitch: torch.Tensor | None = None,
) -> int:
"""Resolve semitone shift for auto_shift; ``f0_to_coarse(..., f0_shift=shift * 5)`` consumes it."""
if not auto_shift or manual_shift != 0:
return int(manual_shift)
if prompt_note_pitch is not None and target_note_pitch is not None:
target_pitched = target_note_pitch[target_note_pitch >= 1]
prompt_pitched = prompt_note_pitch[prompt_note_pitch >= 1]
if target_pitched.numel() > 0 and prompt_pitched.numel() > 0:
shift = torch.round(prompt_pitched.median() - target_pitched.median())
if torch.isfinite(shift):
return int(shift.item())
if prompt_f0 is not None and target_f0 is not None:
target_voiced = target_f0[target_f0 > 0]
prompt_voiced = prompt_f0[prompt_f0 > 0]
if target_voiced.numel() > 0 and prompt_voiced.numel() > 0:
shift = torch.round(torch.log2(prompt_voiced.median() / target_voiced.median()) * 1200 / 100)
if torch.isfinite(shift):
return int(shift.item())
return 0
class MetadataProcessor:
"""Data processor for SoulX-Singer"""
def __init__(
self,
hop_size: int,
sample_rate: int,
phoneset_path: str,
device: str = "cuda",
):
"""Initialize data processor.
Args:
hop_size (int): Hop size in samples.
sample_rate (int): Sample rate in Hz.
phoneset_path (str): Path to phoneme set JSON file.
device (str): Device to use for tensor operations.
"""
self.hop_size = hop_size
self.sample_rate = sample_rate
self.device = device
self.load_phoneme_id_map(phoneset_path)
def load_phoneme_id_map(self, phoneset_path: str):
with open(phoneset_path, encoding="utf-8") as f:
phoneset = json.load(f)
self.phone2idx = {ph: idx for idx, ph in enumerate(phoneset)}
def merge_phoneme(self, meta):
merged_items = []
duration = [float(x) for x in meta["duration"].split()]
phoneme = [str(x).replace("<AP>", "<SP>") for x in meta["phoneme"].split()]
note_pitch = [int(x) for x in meta["note_pitch"].split()]
note_type = [int(x) if phoneme[i] != "<SP>" 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] == "<SP>"
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 <BOW> and <EOW> 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("<BOW>")
note2origin.append(ph_idx)
if phonemes[ph_idx][:3] == "en_":
en_phs = ["en_" + x for x in phonemes[ph_idx][3:].split("-")] + [
"<SEP>"
] # <sep> 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("<EOW>")
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 # <BOW> + ph repeats + <EOW>
new_phonemes = ["<PAD>"] + 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
vllm_omni/diffusion/registry.py
-14
View File
@@ -316,16 +316,6 @@ _DIFFUSION_MODELS = {
"pipeline_cosmos3",
"Cosmos3OmniDiffusersPipeline",
),
"SoulXSingerPipeline": (
"soulx_singer",
"pipeline_soulx_singer_svs",
"PipelineSoulXSingerSVS",
),
"SoulXSingerSVCPipeline": (
"soulx_singer",
"pipeline_soulx_singer_svc",
"PipelineSoulXSingerSVC",
),
"DiffusersAdapterPipeline": (
"diffusers_adapter",
"pipeline_diffusers_adapter",
@@ -566,8 +556,6 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
"MiniMaxH3Pipeline": "get_minimax_h3_post_process_func",
"MiniMaxH3ModularPipeline": "get_minimax_h3_post_process_func",
"StableAudioPipeline": "get_stable_audio_post_process_func",
"SoulXSingerPipeline": "get_soulxsinger_post_process_func",
"SoulXSingerSVCPipeline": "get_soulxsinger_post_process_func",
"AudioXPipeline": "get_audiox_post_process_func",
"WanImageToVideoPipeline": "get_wan22_i2v_post_process_func",
"WanS2VPipeline": "get_wan22_s2v_post_process_func",
@@ -647,8 +635,6 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
"SanaWmPipeline": "get_sana_wm_pre_process_func",
"Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func",
"Cosmos3OmniPipeline": "get_cosmos3_pre_process_func",
"SoulXSingerPipeline": "get_soulxsinger_pre_process_func",
"SoulXSingerSVCPipeline": "get_soulxsinger_svc_pre_process_func",
}
vllm_omni/model_executor/models/soulx_singer/__init__.py
-5
View File
@@ -1,5 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Keep this package lightweight. The pipeline registry imports pipeline.py
# directly, while diffusion model classes are loaded lazily from their registry.
vllm_omni/model_executor/models/soulx_singer/pipeline.py
-45
View File
@@ -1,45 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""SoulX-Singer pipeline topologies."""
from vllm_omni.config.stage_config import (
PipelineConfig,
StageExecutionType,
StagePipelineConfig,
)
SOULXSINGER_SVS_PIPELINE = PipelineConfig(
model_type="soulxsinger_svs",
default_deploy_config_name="soulxsinger_svs.yaml",
model_arch="SoulXSingerPipeline",
hf_architectures=("SoulXSingerPipeline",),
stages=(
StagePipelineConfig(
stage_id=0,
model_stage="dit",
execution_type=StageExecutionType.DIFFUSION,
input_sources=(),
final_output=True,
final_output_type="audio",
model_arch="SoulXSingerPipeline",
),
),
)
SOULXSINGER_SVC_PIPELINE = PipelineConfig(
model_type="soulxsinger_svc",
default_deploy_config_name="soulxsinger_svc.yaml",
model_arch="SoulXSingerSVCPipeline",
hf_architectures=("SoulXSingerSVCPipeline",),
stages=(
StagePipelineConfig(
stage_id=0,
model_stage="dit",
execution_type=StageExecutionType.DIFFUSION,
input_sources=(),
final_output=True,
final_output_type="audio",
model_arch="SoulXSingerSVCPipeline",
),
),
)