def calculate_pitch(
    audio_array: mx.array,
    sr: int,
    min_freq: float = 75.0,
    max_freq: float = 600.0,
    frame_length: int = 400,
    hop_length: int = 160,
    threshold: float = 0.3,
) -> mx.array:
    """
    Calculate pitch frequencies for short audio clips using autocorrelation.

    Args:
        audio_array: Input audio array (1D or 2D [channels, samples])
        sr: Sampling rate
        min_freq: Minimum detectable frequency (Hz)
        max_freq: Maximum detectable frequency (Hz)
        frame_length: Analysis frame length in samples
        hop_length: Hop size in samples
        threshold: Voicing threshold (0.0-1.0)

    Returns:
        Array of pitch values (Hz) per frame
    """
    audio_np = np.array(audio_array)

    # convert to mono and ensure 1D
    if len(audio_np.shape) > 1:
        audio_np = np.mean(audio_np, axis=0)
    audio_np = np.squeeze(audio_np)

    num_samples = audio_np.shape[-1]
    pad_len = (frame_length - (num_samples % hop_length)) % hop_length
    audio_np = np.pad(audio_np, (0, pad_len))

    num_frames = (len(audio_np) - frame_length) // hop_length + 1
    frames = np.zeros((num_frames, frame_length))
    for i in range(num_frames):
        frames[i] = audio_np[i * hop_length : i * hop_length + frame_length]

    window = np.hanning(frame_length)
    frames_windowed = frames * window

    # compute autocorrelation using FFT
    fft_frames = np.fft.rfft(frames_windowed, n=2 * frame_length, axis=1)
    power_spectrum = fft_frames.real**2 + fft_frames.imag**2
    autocorr = np.fft.irfft(power_spectrum, axis=1)[:, :frame_length]

    # find valid frequency range indices
    min_idx = max(1, int(sr / max_freq))
    max_idx = min(frame_length, int(sr / min_freq))

    # find peak indices in valid range
    relevant_autocorr = autocorr[:, min_idx:max_idx]
    peak_indices = np.argmax(relevant_autocorr, axis=1) + min_idx
    peak_values = np.array([autocorr[i, peak_indices[i]] for i in range(num_frames)])

    # parabolic interpolation for sub-sample accuracy
    indices = np.clip(peak_indices, 1, frame_length - 2)
    alpha = np.array([autocorr[i, indices[i] - 1] for i in range(num_frames)])
    beta = np.array([autocorr[i, indices[i]] for i in range(num_frames)])
    gamma = np.array([autocorr[i, indices[i] + 1] for i in range(num_frames)])

    delta = 0.5 * (alpha - gamma) / (alpha - 2 * beta + gamma + 1e-8)
    valid_mask = (peak_indices > 0) & (peak_indices < frame_length - 1)
    delta = np.where(valid_mask, delta, 0.0)

    # calculate final periods and pitches
    best_period = (peak_indices + delta) / sr
    pitch = np.where(best_period > 0, 1.0 / best_period, 0.0)

    # apply voicing threshold
    autocorr_0 = autocorr[:, 0]
    voiced = (peak_values / (autocorr_0 + 1e-8)) > threshold
    pitch = np.where(voiced, pitch, 0.0)

    # clamp valid frequencies
    pitch = np.clip(pitch, min_freq, max_freq)

    return mx.array(pitch)

    Calculate pitch frequencies for short audio clips using autocorrelation.

    Args:
        audio_array: Input audio array (1D or 2D [channels, samples])
        sr: Sampling rate
        min_freq: Minimum detectable frequency (Hz)
        max_freq: Maximum detectable frequency (Hz)
        frame_length: Analysis frame length in samples
        hop_length: Hop size in samples
        threshold: Voicing threshold (0.0-1.0)

    Returns:
        Array of pitch values (Hz) per frame
    
mlx_audio/tts/models/outetts/audio_processor.py