Browse Source
Update pipelines with pydantic
Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
1 changed files with
15 additions and
13 deletions
-
audio_embedding.py
|
|
@ -1,20 +1,22 @@ |
|
|
|
from posixpath import basename |
|
|
|
from typing import Optional |
|
|
|
from pydantic import BaseModel |
|
|
|
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 |
|
|
|
class AudioEmbeddingConfig(BaseModel): |
|
|
|
# config for audio_decode.ffmpeg |
|
|
|
batch_size: Optional[int] = -1 |
|
|
|
sample_rate: Optional[float] = None |
|
|
|
layout: Optional[str] = None |
|
|
|
|
|
|
|
# config for audio_embedding.vggish |
|
|
|
weights_path: Optional[str] = None |
|
|
|
framework: Optional[str] = 'pytorch' |
|
|
|
|
|
|
|
# config for triton |
|
|
|
device: Optional[int] = -1 |
|
|
|
|
|
|
|
|
|
|
|
@AutoPipes.register |
|
|
|