from typing import Dict, List, Tuple
import numpy as np
import os
import torch

class PreTrainedPipeline():
    def __init__(self, path=''):
        self.model = torch.hub.load('sigsep/open-unmix-pytorch', 'umxhq')
        self.sampling_rate = int(self.model.sample_rate.item())

    def __call__(self, inputs): 
        estimates = self.model(inputs.unsqueeze(0))
        vocals = estimates[0][0].detach().numpy()
        n = vocals.shape[0]
        return vocals, self.sampling_rate, [f"label_{i}" for i in range(n)]