tools/time_strech.py

import os
from pathlib import Path
from typing import Optional, Tuple

import librosa
import soundfile as sf


def align_songs_by_bpm(
    audio1_path: str,
    audio2_path: str,
    output_path: Optional[str] = None,
    output_format: str = "wav",
) -> Tuple[str, str]:
    """
    Align two songs to the same BPM by stretching the faster one to match the slower one.

    This function analyzes the tempo of two audio files and automatically stretches the faster
    track to match the BPM of the slower track, making them suitable for mixing or mashups.

    Args:
        audio1_path: Path to first audio file (supports common formats: WAV, MP3, FLAC)
        audio2_path: Path to second audio file (supports common formats: WAV, MP3, FLAC)
        output_path: Optional output directory (default: None, uses temporary directory)
        output_format: Output format ('wav' or 'mp3', default: 'wav')

    Returns:
        Tuple of (aligned_audio1_path, aligned_audio2_path): Paths to the processed audio files
        Both files will have the same BPM (the slower of the two original tempos)

    Examples:
        - Song A: 140 BPM, Song B: 128 BPM → Both become 128 BPM
        - Song A: 120 BPM, Song B: 130 BPM → Both become 120 BPM

    Note:
        Uses high-quality time-stretching to maintain audio quality
        Preserves the original pitch of both tracks
        Processing time depends on audio length and tempo difference
    """
    try:
        # Load both audio files
        y1, sr1 = librosa.load(audio1_path)
        y2, sr2 = librosa.load(audio2_path)

        # Get BPM for both tracks
        tempo1, _ = librosa.beat.beat_track(y=y1, sr=sr1)
        tempo2, _ = librosa.beat.beat_track(y=y2, sr=sr2)

        bpm1 = float(tempo1)
        bpm2 = float(tempo2)

        # Determine which track is faster and needs stretching
        if bpm1 > bpm2:
            # Stretch first track to match second track's BPM
            aligned1_path = stretch_to_bpm(
                audio1_path, bpm2, output_path, output_format
            )
            aligned2_path = stretch_to_bpm(
                audio2_path, bpm2, output_path, output_format
            )
        else:
            # Stretch second track to match first track's BPM
            aligned1_path = stretch_to_bpm(
                audio1_path, bpm1, output_path, output_format
            )
            aligned2_path = stretch_to_bpm(
                audio2_path, bpm1, output_path, output_format
            )

        return aligned1_path, aligned2_path

    except Exception as e:
        raise RuntimeError(f"Error aligning audio files: {str(e)}")


def stretch_to_bpm(
    audio_path: str,
    target_bpm: float,
    output_path: Optional[str] = None,
    output_format: str = "wav",
) -> str:
    """
    Stretch an audio file to a specific BPM.

    Args:
        audio_path: Path to audio file or URL (supports WAV, MP3, FLAC)
        target_bpm: Target BPM to stretch to
        output_path: Path to output file
        output_format: Output format ('wav' or 'mp3', default: 'wav')

    Returns:
        Path to the stretched audio file
    """
    try:
        y, sr = librosa.load(audio_path, sr=None, mono=False)

        # Get current BPM
        y_hat, sr_hat = librosa.load(audio_path)
        tempo, _ = librosa.beat.beat_track(y=y_hat, sr=sr_hat)
        current_bpm = float(tempo)

        # Calculate stretch factor
        stretch_factor = target_bpm / current_bpm

        # Apply time stretching
        y_stretched = librosa.effects.time_stretch(y, rate=stretch_factor)

        # Save to temporary file
        if not output_path:
            output_path = "output"
        os.makedirs(output_path, exist_ok=True)

        # Get original filename without extension
        original_audio_filename = Path(audio_path).stem

        # Determine output extension
        output_ext = f".{output_format.lower()}"
        output_file_path = os.path.join(
            output_path,
            f"{original_audio_filename}_stretched_to_{int(target_bpm)}_bpm{output_ext}",
        )

        if y_stretched.ndim == 2:
            y_stretched = y_stretched.T  # Transpose for multi-channel audio

        # Write in the requested format
        if output_format.lower() == "mp3":
            # For MP3, we need to use ffmpeg through subprocess
            import tempfile
            import subprocess

            # First save as WAV to temporary file
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_wav:
                sf.write(temp_wav.name, y_stretched, sr)

                # Convert to MP3 using ffmpeg
                cmd = [
                    "ffmpeg",
                    "-y",
                    "-i",
                    temp_wav.name,
                    "-c:a",
                    "libmp3lame",
                    "-b:a",
                    "192k",
                    output_file_path,
                ]
                subprocess.run(cmd, capture_output=True, check=True)

                # Clean up temp file
                os.unlink(temp_wav.name)
        else:
            # For WAV and other formats supported by soundfile
            sf.write(output_file_path, y_stretched, sr)

        return output_file_path

    except Exception as e:
        raise RuntimeError(f"Error stretching audio: {str(e)}")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Time stretch audio files")
    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # Align two songs by BPM
    align_parser = subparsers.add_parser("align", help="Align two songs to same BPM")
    align_parser.add_argument("audio1", help="Path to first audio file")
    align_parser.add_argument("audio2", help="Path to second audio file")
    align_parser.add_argument(
        "--format", default="wav", choices=["wav", "mp3"], help="Output format"
    )

    # Stretch to specific BPM
    stretch_parser = subparsers.add_parser(
        "stretch", help="Stretch audio to specific BPM"
    )
    stretch_parser.add_argument("audio", help="Path to audio file")
    stretch_parser.add_argument("target_bpm", type=float, help="Target BPM")
    stretch_parser.add_argument(
        "--format", default="wav", choices=["wav", "mp3"], help="Output format"
    )

    args = parser.parse_args()

    try:
        if args.command == "align":
            aligned1, aligned2 = align_songs_by_bpm(
                args.audio1, args.audio2, output_format=args.format
            )
            print(f"Aligned audio 1: {aligned1}")
            print(f"Aligned audio 2: {aligned2}")
        elif args.command == "stretch":
            output = stretch_to_bpm(
                args.audio, args.target_bpm, output_format=args.format
            )
            print(f"Stretched audio saved to: {output}")
        else:
            parser.print_help()
    except Exception as e:
        print(f"Error: {e}")
        exit(1)