from towhee import pipe, ops, AutoPipes, AutoConfig


@AutoConfig.register
class AudioEmbeddingConfig:
    def __init__(self):
        # config for audio_decode.ffmpeg
        self.batch_size = -1
        self.sample_rate = None
        self.layout = None
        
        # config for audio_embedding.vggish
        self.weights_path: str = None
        self.framework: str = 'pytorch'
        
        # config for triton
        self.device = -1


@AutoPipes.register
def AudioEmbedding(config=None):
    if not config:
        config = AudioEmbeddingConfig()

    if config.device >= 0:
        op_config = AutoConfig.TritonGPUConfig(device_ids=[config.device], max_batch_size=128)
    else:
        op_config = AutoConfig.TritonCPUConfig()

    
    return (
        pipe.input('path')
            .map('path', 'frame', ops.audio_decode.ffmpeg(config.batch_size, config.sample_rate, config.layout))
            .map('frame', 'vec', ops.audio_embedding.vggish(config.weights_path, config.framework), config=op_config)
            .output('vec')
    )