scripts/utils/gpu_probe.py

import os

import torch

# --- MONKEYPATCHES START ---
try:
    import torchaudio

    # Patch 1: AudioMetaData
    if not hasattr(torchaudio, "AudioMetaData"):

        class AudioMetaData:
            def __init__(self, sample_rate, num_frames, num_channels, bits_per_sample, encoding):
                self.sample_rate = sample_rate
                self.num_frames = num_frames
                self.num_channels = num_channels
                self.bits_per_sample = bits_per_sample
                self.encoding = encoding

        torchaudio.AudioMetaData = AudioMetaData

    # Patch 2: list_audio_backends
    if not hasattr(torchaudio, "list_audio_backends"):
        torchaudio.list_audio_backends = lambda: ["soundfile"]

    # Patch 3: load/info (Mock)
    import soundfile as sf

    def robust_load(filepath, **kwargs):
        data, sr = sf.read(filepath, dtype="float32")  # Alway float32
        return torch.tensor(data).float().unsqueeze(0), sr

    torchaudio.load = robust_load

    class MockAudioInfo:
        def __init__(self, frames, samplerate, channels):
            self.num_frames = frames
            self.sample_rate = samplerate
            self.num_channels = channels

    def robust_info(filepath, **kwargs):
        s = sf.info(filepath)
        return MockAudioInfo(s.frames, s.samplerate, s.channels)

    torchaudio.info = robust_info

except ImportError:
    pass

import semver

# Patch 4: Semver
original_parse = semver.VersionInfo.parse


def lenient_parse(version_str):
    try:
        return original_parse(version_str)
    except ValueError:
        return semver.VersionInfo(0, 0, 0)  # Dummy


semver.VersionInfo.parse = lenient_parse

# --- MONKEYPATCHES END ---

from transformers import pipeline
from pyannote.audio import Pipeline, Model, Inference


def test_whisper():
    print("\n--- Testing Whisper on GPU ---")
    try:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Device: {device}")
        if device == "cpu":
            print("FAIL: CUDA not available")
            return False

        pipe = pipeline(
            "automatic-speech-recognition", model="openai/whisper-large-v3", torch_dtype=torch.float16, device=device
        )
        # Dummy inference
        # Generate dummy 1s audio
        import numpy as np

        dummy_audio = np.zeros(16000, dtype=np.float32)

        print("Running Inference...")
        pipe(dummy_audio)
        print("SUCCESS: Whisper runs on GPU")
        return True
    except Exception as e:
        print(f"FAIL: Whisper crashed on GPU: {e}")
        import traceback

        traceback.print_exc()
        return False


def test_pyannote_embedding():
    print("\n--- Testing Pyannote Embedding on GPU ---")
    try:
        device = torch.device("cuda")
        print("Loading Model...")
        model = Model.from_pretrained(
            "pyannote/wespeaker-voxceleb-resnet34-LM", use_auth_token=os.environ.get("HF_TOKEN")
        )
        model.to(device)
        inference = Inference(model, window="whole", device=device)

        # Dummy file (needs a real path or Mock)
        # We'll creating a dummy wav
        import soundfile as sf
        import numpy as np

        dummy_wav = "temp_gpu_test.wav"
        sf.write(dummy_wav, np.random.uniform(-1, 1, 16000), 16000)

        print("Running Embedding...")
        inference(dummy_wav)
        print("SUCCESS: Pyannote Embedding runs on GPU")
        os.remove(dummy_wav)
        return True
    except Exception as e:
        print(f"FAIL: Pyannote Embedding crashed on GPU: {e}")
        import traceback

        traceback.print_exc()
        if os.path.exists("temp_gpu_test.wav"):
            os.remove("temp_gpu_test.wav")
        return False


def test_pyannote_diarization():
    print("\n--- Testing Pyannote Diarization on GPU ---")
    try:
        device = torch.device("cuda")
        print("Loading Pipeline...")
        pipeline = Pipeline.from_pretrained(
            "pyannote/speaker-diarization-3.1", use_auth_token=os.environ.get("HF_TOKEN")
        )
        pipeline.to(device)

        dummy_wav = "temp_gpu_test_dia.wav"
        import soundfile as sf
        import numpy as np

        sf.write(dummy_wav, np.random.uniform(-1, 1, 16000 * 5), 16000)  # 5s

        print("Running Diarization...")
        pipeline(dummy_wav)
        print("SUCCESS: Pyannote Diarization runs on GPU")
        os.remove(dummy_wav)
        return True
    except Exception as e:
        print(f"FAIL: Pyannote Diarization crashed on GPU: {e}")
        import traceback

        traceback.print_exc()
        if os.path.exists(dummy_wav):
            os.remove(dummy_wav)
        return False


if __name__ == "__main__":
    w_ok = test_whisper()
    e_ok = test_pyannote_embedding()
    d_ok = test_pyannote_diarization()

    print("\nSUMMARY:")
    print(f"Whisper: {'OK' if w_ok else 'FAIL'}")
    print(f"Embedding: {'OK' if e_ok else 'FAIL'}")
    print(f"Diarization: {'OK' if d_ok else 'FAIL'}")