guidellm.utils.audio

def encode_audio(
    audio: libs.AudioDecoder
    | bytes
    | str
    | Path
    | np.ndarray
    | torch.Tensor
    | dict[str, Any],
    sample_rate: int | None = None,
    file_name: str | None = None,
    encode_sample_rate: int = 16000,
    max_duration: float | None = None,
    mono: bool = True,
    audio_format: str | None = None,
    bitrate: str = "64k",
) -> dict[
    Literal[
        "type",
        "audio",
        "format",
        "mimetype",
        "audio_samples",
        "audio_seconds",
        "audio_bytes",
        "file_name",
    ],
    str | int | float | bytes | None,
]:
    """
    Decode audio (if necessary) and re-encode to specified format.

    If ``audio_format`` is not provided, the format is detected from the
    source codec (via torchcodec metadata). When detection is not possible
    (e.g. raw float tensors), falls back to WAV with a one-time warning.
    The returned ``audio_samples`` value is the number of frames represented
    at the encoded output sample rate.

    :param audio: Audio input in any supported form.
    :param sample_rate: Sample rate hint for raw decoded audio.
    :param file_name: Override file name in output metadata. Defaults to a
        name derived from the resolved format.
    :param encode_sample_rate: Target sample rate for the encoded output.
    :param max_duration: Truncate audio to this duration in seconds.
    :param mono: Convert to mono if True.
    :param audio_format: Target encoding format. If None, detected from source.
    :param bitrate: Bitrate for lossy formats like mp3.
    :return: Dict containing encoded audio bytes and metadata.
    """
    samples, source_codec = _decode_audio(
        audio, sample_rate=sample_rate, max_duration=max_duration
    )

    # Resolve format: explicit > codec-detected > WAV fallback
    if audio_format is None:
        if source_codec:
            audio_format = _codec_to_format(source_codec)
        if audio_format is None:
            audio_format = "wav"
            logger.debug("Falling back to WAV audio formatting")

    bitrate_val = (
        int(bitrate.rstrip("k")) * 1000 if bitrate.endswith("k") else int(bitrate)
    )
    format_val = audio_format.lower()

    encoded_audio = _encode_audio(
        samples=samples,
        resample_rate=encode_sample_rate,
        bitrate=bitrate_val,
        audio_format=format_val,
        mono=mono,
    )
    encoded_sample_frames = round(samples.duration_seconds * encode_sample_rate)

    resolved_file_name = file_name
    if not resolved_file_name and isinstance(audio, str | Path):
        resolved_file_name = get_file_name(audio)
    if not resolved_file_name:
        resolved_file_name = f"audio.{format_val}"

    return {
        "type": "audio_file",
        "audio": encoded_audio,
        "file_name": resolved_file_name,
        "format": audio_format,
        "mimetype": f"audio/{format_val}",
        "audio_samples": encoded_sample_frames,
        "audio_seconds": samples.duration_seconds,
        "audio_bytes": len(encoded_audio),
    }

def get_file_name(path: Path | str) -> str:
    """Get file name from path."""
    return Path(path).name

def pcm16_append_b64_chunks(
    audio_item: dict[str, Any] | bytes,
    *,
    target_sample_rate: int = 16000,
    chunk_samples: int = 3200,
) -> list[str]:
    """
    Decode audio to base64-encoded PCM16 mono chunks for realtime ``append`` events.

    Matches vLLM ``input_audio_buffer.append`` (PCM16 mono at ``target_sample_rate``
    Hz), split into ``chunk_samples``-frame segments.
    Equivalent conversion flow to vLLM's realtime microphone client example, but
    generalized for dataset/file inputs used by GuideLLM benchmarks.
    """
    # Accept common audio column shapes used in GuideLLM datasets.
    if isinstance(audio_item, dict):
        if "audio" in audio_item:
            decode_sr = _sample_rate_hint_from_audio_column_dict(audio_item)
            samples, _ = _decode_audio(
                audio_item["audio"],
                sample_rate=decode_sr,
            )
        elif "data" in audio_item or "url" in audio_item:
            samples, _ = _decode_audio(audio_item)
        else:
            raise ValueError(
                "audio_column dict must include 'audio', 'data', or 'url' "
                "(same shapes as encode_audio / _decode_audio); "
                f"got keys {list(audio_item)!r}"
            )
    else:
        samples, _ = _decode_audio(audio_item)

    # Ensure channel-first shape, then downmix to mono for realtime PCM input.
    data = samples.data
    if data.dim() == 1:
        data = data.unsqueeze(0)
    if data.shape[0] > 1:
        data = data.mean(dim=0, keepdim=True)

    # Realtime endpoint expects 16 kHz PCM16 mono.
    sr = _require_positive_sample_rate(samples.sample_rate)
    if sr != target_sample_rate:
        t_in = data.shape[1]
        t_out = max(1, int(round(t_in * target_sample_rate / sr)))
        data = torch.nn.functional.interpolate(
            data.unsqueeze(0),
            size=t_out,
            mode="linear",
            align_corners=False,
        ).squeeze(0)

    # Convert float waveform to signed little-endian PCM16 bytes.
    wave = data.squeeze(0)
    pcm_i16 = (
        (
            wave.clamp(_PCM16_WAVE_CLIP_MIN, _PCM16_WAVE_CLIP_MAX)
            * _PCM16_FLOAT_TO_INT16_SCALE
        )
        .round()
        .to(torch.int16)
    )
    buf = pcm_i16.cpu().numpy().tobytes()

    # Split PCM bytes into chunk-sized base64 payloads for append events.
    chunk_bytes = max(1, chunk_samples) * _BYTES_PER_PCM16_SAMPLE
    out: list[str] = []
    for i in range(0, len(buf), chunk_bytes):
        pcm_chunk = buf[i : i + chunk_bytes]
        if pcm_chunk:
            out.append(base64.b64encode(pcm_chunk).decode("ascii"))
    if not out:
        raise ValueError("Decoded audio produced no PCM data")
    return out