towhee
/
audio-embedding
copied
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Readme
Files and versions
36 lines
1.0 KiB
36 lines
1.0 KiB
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), conf=op_config)
|
|
.output('vec')
|
|
)
|