inference.py

"""
Guitar Domain Expert HTDemucs 6s Standalone Inference Script

Extracts a clean, isolated guitar stem from any mixed audio file using
the fine-tuned htdemucs_6s model (guitar stem index 4).

Usage:
    python inference.py --input song.mp3 --output guitar_stem.wav
    python inference.py --input song.wav --output guitar.wav --device cpu
    python inference.py --input song.mp3  # outputs to song_guitar.wav alongside input

Requirements:
    pip install torch torchaudio demucs==4.0.1 soundfile tqdm
"""

import argparse
import logging
import sys
from pathlib import Path

import soundfile as sf
import torch
import torchaudio
from tqdm import tqdm

logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
GUITAR_STEM_INDEX = 4          # Index 4 in htdemucs_6s: [drums, bass, other, vocals, guitar, piano]
MODEL_SAMPLE_RATE = 44100      # Model was fine-tuned at 44.1 kHz
SEGMENT_SECONDS = 7.8          # Receptive field used during fine-tuning
OVERLAP = 0.25                 # Chunk overlap for smooth boundary stitching


# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------
def load_model(checkpoint_path: str, device: str = "auto") -> torch.nn.Module:
    """
    Load the fine-tuned guitar extraction model from a local checkpoint.

    Args:
        checkpoint_path: Path to guitar_htdemucs_6s.pt
        device: "cuda", "cpu", or "auto" (auto-selects GPU if available)

    Returns:
        HTDemucs model in eval mode, moved to the target device.
    """
    from demucs.pretrained import get_model

    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"

    logger.info(f"Loading model on {device} ...")
    model = get_model("htdemucs_6s")

    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)

    # Checkpoint may be a full training snapshot (epoch, model_state_dict, optimizer_state_dict, ...)
    # or a bare state dict handle both gracefully.
    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
        state_dict = checkpoint["model_state_dict"]
        logger.info(f"  Checkpoint: epoch {checkpoint.get('epoch', '?')} | "
                    f"val_sdr={checkpoint.get('val_sdr_db', '?')} dB")
    else:
        state_dict = checkpoint

    # demucs get_model() returns a BagOfModels wrapper whose state_dict prefixes
    # all keys with "models.0.". The training notebook saved the inner HTDemucs
    # model's state_dict directly (no prefix). Remap to match the wrapper schema.
    wrapper_keys = set(model.state_dict().keys())
    if not any(k in wrapper_keys for k in list(state_dict.keys())[:3]):
        state_dict = {f"models.0.{k}": v for k, v in state_dict.items()}

    model.load_state_dict(state_dict)

    model.to(device)
    model.eval()
    logger.info("Model loaded.")
    return model, device


# ---------------------------------------------------------------------------
# Audio I/O
# ---------------------------------------------------------------------------
def load_audio(path: str) -> tuple[torch.Tensor, int]:
    """
    Load any audio file and return a stereo float32 tensor.

    Uses soundfile instead of torchaudio.load to avoid the Windows torchcodec
    DLL crash (libtorchcodec*.dll not found). torchaudio is only used for
    resampling, not file I/O.

    Args:
        path: Path to audio file (.wav, .flac; use soundfile-supported formats)

    Returns:
        waveform: Tensor of shape (2, T) stereo, float32
        original_sr: Original sample rate of the file
    """
    audio_np, sr = sf.read(str(path), always_2d=True, dtype="float32")
    # soundfile returns (frames, channels) transpose to (channels, frames)
    waveform = torch.from_numpy(audio_np.T)  # (C, T)

    # Ensure stereo
    if waveform.shape[0] == 1:
        waveform = waveform.repeat(2, 1)
    elif waveform.shape[0] > 2:
        waveform = waveform[:2]

    return waveform, sr


def resample_if_needed(waveform: torch.Tensor, src_sr: int, target_sr: int) -> torch.Tensor:
    """Resample to model sample rate if the source differs."""
    if src_sr == target_sr:
        return waveform
    logger.info(f"Resampling from {src_sr} Hz → {target_sr} Hz ...")
    return torchaudio.functional.resample(waveform, src_sr, target_sr)


# ---------------------------------------------------------------------------
# Chunked inference
# ---------------------------------------------------------------------------
def separate_guitar(
    model: torch.nn.Module,
    waveform: torch.Tensor,
    device: str,
    shifts: int = 1,
    segment_seconds: float = SEGMENT_SECONDS,
    overlap: float = OVERLAP,
) -> torch.Tensor:
    """
    Extract the guitar stem using demucs' native apply_model API.

    BagOfModels (returned by get_model) does not implement forward() directly
    demucs mandates using apply_model, which handles chunking, overlap-add,
    and cross-fade internally.

    Args:
        model: Loaded fine-tuned HTDemucs model in eval mode.
        waveform: Stereo input tensor of shape (2, T) at 44100 Hz.
        device: "cuda" or "cpu"
        shifts: Number of overlapping shifts (default 1). Higher = cleaner, slower.
        segment_seconds: Chunk length in seconds (default 7.8 s).
        overlap: Fraction of chunk to overlap between adjacent segments.

    Returns:
        guitar_stem: Stereo tensor of shape (2, T) isolated guitar waveform.
    """
    from demucs.apply import apply_model

    # apply_model expects (batch, channels, time) on CPU. 
    # It will move chunks to the device internally to save VRAM.
    mix = waveform.unsqueeze(0)  # (1, 2, T)

    with torch.no_grad():
        sources = apply_model(
            model,
            mix,
            shifts=shifts,
            segment=segment_seconds,
            overlap=overlap,
            device=device,
            progress=True,       # shows tqdm bar
        )
    # sources shape: (batch, stems, channels, time)
    guitar_stem = sources[0, GUITAR_STEM_INDEX].cpu().clamp(-1.0, 1.0)
    return guitar_stem


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Extract a clean guitar stem from a mixed audio track.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        "--input", "-i", required=True, type=str,
        help="Path to the input audio file (wav, mp3, flac, etc.)",
    )
    parser.add_argument(
        "--output", "-o", type=str, default=None,
        help="Path for the output guitar stem WAV. Defaults to <input>_guitar.wav",
    )
    parser.add_argument(
        "--checkpoint", "-c", type=str,
        default=str(Path(__file__).parent / "guitar_htdemucs_6s.pt"),
        help="Path to the model checkpoint (default: guitar_htdemucs_6s.pt in script dir)",
    )
    parser.add_argument(
        "--device", "-d", type=str, default="auto",
        choices=["auto", "cuda", "cpu"],
        help="Compute device (default: auto uses GPU if available)",
    )
    parser.add_argument(
        "--shifts", "-s", type=int, default=1,
        help="Number of shifts for overlap-add. Higher = better quality, slower (default: 1)",
    )
    parser.add_argument(
        "--overlap", type=float, default=OVERLAP,
        help=f"Chunk overlap fraction (default: {OVERLAP})",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    input_path = Path(args.input)
    if not input_path.exists():
        logger.error(f"Error: input file not found: {input_path}")
        sys.exit(1)

    output_path = Path(args.output) if args.output else input_path.with_name(
        f"{input_path.stem}_guitar.wav"
    )

    # Load model
    model, device = load_model(args.checkpoint, args.device)

    # Load and prepare audio
    waveform, original_sr = load_audio(str(input_path))
    waveform = resample_if_needed(waveform, original_sr, MODEL_SAMPLE_RATE)

    duration_s = waveform.shape[1] / MODEL_SAMPLE_RATE
    logger.info(f"Input:    {input_path.name}  ({duration_s:.1f}s, {waveform.shape[0]}ch @ {MODEL_SAMPLE_RATE} Hz)")

    # Run separation
    guitar_stem = separate_guitar(
        model, 
        waveform, 
        device, 
        shifts=args.shifts, 
        overlap=args.overlap
    )

    # Save output
    sf.write(
        str(output_path),
        guitar_stem.T.numpy(),   # soundfile expects (T, channels)
        samplerate=MODEL_SAMPLE_RATE,
        subtype="PCM_24",
    )
    logger.info(f"Output:   {output_path}")


if __name__ == "__main__":
    main()