diarize.py

def diarize_audio(audio_file, num_speakers=None, device=None):
    # Set device
    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    # Load pipeline
    print("Loading diarization pipeline...")
    pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-community-1")
    pipeline.to(device)

    # Perform diarization
    print(f"Performing diarization on {audio_file}...")
    with ProgressHook() as hook:
        if num_speakers:
            diarization = pipeline(audio_file, hook=hook, num_speakers=num_speakers)
        else:
            diarization = pipeline(audio_file, hook=hook)

    print("Diarization complete!")
    return diarization