# model_handler.py
from speechbrain.inference.separation import SepformerSeparation


class MySeparationModel:
    def __init__(self):
        self.model = SepformerSeparation.from_hparams(
            source="speechbrain/sepformer-wham",
            savedir="pretrained_models/sepformer-wham"
        )

    def predict(self, audio_path):
        return self.model.separate_file(audio_path)