vllm-omni/vllm_omni/diffusion/models/soulx_singer/preprocess/pre_process.py

"""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()),
    )