Spaces:
Sleeping
Sleeping
File size: 1,015 Bytes
89040ed |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
# Prediction interface for Cog ⚙️
# https://github.com/replicate/cog/blob/main/docs/python.md
import torch
from cog import BasePredictor, Input, Path
from pipeline import build_audiosep, inference
class Predictor(BasePredictor):
def setup(self) -> None:
"""Load the model into memory to make running multiple predictions efficient"""
self.model = build_audiosep(
config_yaml="config/audiosep_base.yaml",
checkpoint_path="checkpoint/audiosep_base_4M_steps.ckpt",
device="cuda",
)
def predict(
self,
audio_file: Path = Input(description="Input audio file."),
text: str = Input(description="Input text.", default="water drops"),
) -> Path:
"""Run a single prediction on the model"""
output_file = "/tmp/separated_audio.wav"
# AudioSep processes the audio at 32 kHz sampling rate
inference(self.model, str(audio_file), text, output_file, "cuda")
return Path(output_file)
|