fireredasr2s/fireredvad/__init__.py

# Copyright 2026 Xiaohongshu. (Author: Kaituo Xu, Kai Huang)

import os
import sys

__version__ = "0.0.1"

_CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))

try:
    from fireredasr2s.fireredvad.aed import FireRedAed, FireRedAedConfig
    from fireredasr2s.fireredvad.stream_vad import FireRedStreamVad, FireRedStreamVadConfig
    from fireredasr2s.fireredvad.vad import FireRedVad, FireRedVadConfig
except ImportError:
    if _CURRENT_DIR not in sys.path:
        sys.path.insert(0, _CURRENT_DIR)
    from .aed import FireRedAed, FireRedAedConfig
    from .stream_vad import FireRedStreamVad, FireRedStreamVadConfig
    from .vad import FireRedVad, FireRedVadConfig


def non_stream_vad(wav_path, model_dir="pretrained_models/FireRedVAD/VAD", **kwargs):
    """Quick VAD inference"""
    config = FireRedVadConfig(**kwargs)
    vad = FireRedVad.from_pretrained(model_dir, config)
    result, probs = vad.detect(wav_path)
    return result


def stream_vad_full(wav_path, model_dir="pretrained_models/FireRedVAD/Stream-VAD", **kwargs):
    """Quick Stream VAD inference"""
    config = FireRedStreamVadConfig(**kwargs)
    svad = FireRedStreamVad.from_pretrained(model_dir, config)
    frame_results, result = svad.detect_full(wav_path)
    return frame_results, result


def non_stream_aed(wav_path, model_dir="pretrained_models/FireRedVAD/AED", **kwargs):
    """Quick AED inference"""
    config = FireRedAedConfig(**kwargs)
    aed = FireRedAed.from_pretrained(model_dir, config)
    result, probs = aed.detect(wav_path)
    return result


__all__ = [
    '__version__',
    'FireRedVad',
    'FireRedVadConfig',
    'FireRedAed',
    'FireRedAedConfig', 
    'FireRedStreamVad',
    'FireRedStreamVadConfig',
    'non_stream_vad',
    'stream_vad_full',
    'non_stream_aed'
]