Spaces:
Running
Running
| # 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) | |