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