audio2midi 0.2.0__py2.py3-none-any.whl → 0.4.0__py2.py3-none-any.whl

audio2midi/crepe_pitch_detector.py

@@ -1,54 +1,872 @@
-from math import ceil as math_ceil
-from typing import Callable
-from numpy.lib.stride_tricks import as_strided
-from keras.layers import Input, Reshape, Conv2D, BatchNormalization
-from keras.layers import MaxPool2D, Dropout, Permute, Flatten, Dense
-from keras.models import Model
-from keras.callbacks import Callback
-from hmmlearn.hmm import CategoricalHMM
-from librosa import load as librosa_load
-from pretty_midi_fix import PrettyMIDI , PitchBend , Note ,Instrument
+import warnings
 import numpy as np
+import torch
+import librosa
+from torch.nn import functional as F
+from tqdm import tqdm
+from functools import partial
 from huggingface_hub import hf_hub_download
+from scipy.stats import triang
+
+###############################################################################
+# Constants
+###############################################################################
+
+CENTS_PER_BIN = 20  # cents
+MAX_FMAX = 2006.  # hz
+PITCH_BINS = 360
+SAMPLE_RATE = 16000  # hz
+WINDOW_SIZE = 1024  # samples
+UNVOICED = np.nan
+# Minimum decibel level
+MIN_DB = -100.
+
+# Reference decibel level
+REF_DB = 20.
+
+
+
+
+###############################################################################
+# Probability sequence decoding methods
+###############################################################################
+
+
+def argmax(logits):
+    """Sample observations by taking the argmax"""
+    bins = logits.argmax(dim=1)
+
+    # Convert to frequency in Hz
+    return bins, bins_to_frequency(bins)
+
+
+
+
+
+
+
+
+
+
+
+###############################################################################
+# Pitch unit conversions
+###############################################################################
+
+
+def bins_to_cents(bins):
+    """Converts pitch bins to cents"""
+    cents = CENTS_PER_BIN * bins + 1997.3794084376191
+
+    # Trade quantization error for noise
+    return dither(cents)
+
+
+def bins_to_frequency(bins):
+    """Converts pitch bins to frequency in Hz"""
+    return cents_to_frequency(bins_to_cents(bins))
+
+
+def cents_to_bins(cents, quantize_fn=torch.floor):
+    """Converts cents to pitch bins"""
+    bins = (cents - 1997.3794084376191) / CENTS_PER_BIN
+    return quantize_fn(bins).int()
+
+
+def cents_to_frequency(cents):
+    """Converts cents to frequency in Hz"""
+    return 10 * 2 ** (cents / 1200)
+
+
+def frequency_to_bins(frequency, quantize_fn=torch.floor):
+    """Convert frequency in Hz to pitch bins"""
+    return cents_to_bins(frequency_to_cents(frequency), quantize_fn)
+
+
+def frequency_to_cents(frequency):
+    """Convert frequency in Hz to cents"""
+    return 1200 * torch.log2(frequency / 10.)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+###############################################################################
+# Pitch thresholding methods
+###############################################################################
+
+
+class At:
+    """Simple thresholding at a specified probability value"""
+
+    def __init__(self, value):
+        self.value = value
+
+    def __call__(self, pitch, periodicity):
+        # Make a copy to prevent in-place modification
+        pitch = torch.clone(pitch)
+
+        # Threshold
+        pitch[periodicity < self.value] = UNVOICED
+        return pitch
+
+
+class Hysteresis:
+    """Hysteresis thresholding"""
+
+    def __init__(self,
+                 lower_bound=.19,
+                 upper_bound=.31,
+                 width=.2,
+                 stds=1.7,
+                 return_threshold=False):
+        self.lower_bound = lower_bound
+        self.upper_bound = upper_bound
+        self.width = width
+        self.stds = stds
+        self.return_threshold = return_threshold
+
+    def __call__(self, pitch, periodicity):
+        # Save output device
+        device = pitch.device
+
+        # Perform hysteresis in log-2 space
+        pitch = torch.log2(pitch).detach().flatten().cpu().numpy()
+
+        # Flatten periodicity
+        periodicity = periodicity.flatten().cpu().numpy()
+
+        # Ignore confidently unvoiced pitch
+        pitch[periodicity < self.lower_bound] = UNVOICED
+
+        # Whiten pitch
+        mean, std = np.nanmean(pitch), np.nanstd(pitch)
+        pitch = (pitch - mean) / std
+
+        # Require high confidence to make predictions far from the mean
+        parabola = self.width * pitch ** 2 - self.width * self.stds ** 2
+        threshold = \
+            self.lower_bound + np.clip(parabola, 0, 1 - self.lower_bound)
+        threshold[np.isnan(threshold)] = self.lower_bound
+
+        # Apply hysteresis to prevent short, unconfident voiced regions
+        i = 0
+        while i < len(periodicity) - 1:
+
+            # Detect unvoiced to voiced transition
+            if periodicity[i] < threshold[i] and \
+               periodicity[i + 1] > threshold[i + 1]:
+
+                # Grow region until next unvoiced or end of array
+                start, end, keep = i + 1, i + 1, False
+                while end < len(periodicity) and \
+                      periodicity[end] > threshold[end]:
+                    if periodicity[end] > self.upper_bound:
+                        keep = True
+                    end += 1
+
+                # Force unvoiced if we didn't pass the confidence required by
+                # the hysteresis
+                if not keep:
+                    threshold[start:end] = 1
+
+                i = end
+
+            else:
+                i += 1
+
+        # Remove pitch with low periodicity
+        pitch[periodicity < threshold] = UNVOICED
+
+        # Unwhiten
+        pitch = pitch * std + mean
+
+        # Convert to Hz
+        pitch = torch.tensor(2 ** pitch, device=device)[None, :]
+
+        # Optionally return threshold
+        if self.return_threshold:
+            return pitch, torch.tensor(threshold, device=device)
+
+        return pitch
+
+
+###############################################################################
+# Periodicity thresholding methods
+###############################################################################
+
+
+class Silence:
+    """Set periodicity to zero in silent regions"""
+
+    def __init__(self, value=-60):
+        self.value = value
+        self.a_weighted_weights = self.perceptual_weights()
+    def perceptual_weights(self):
+        """A-weighted frequency-dependent perceptual loudness weights"""
+        frequencies = librosa.fft_frequencies(sr=SAMPLE_RATE,n_fft=WINDOW_SIZE)
+
+        # A warning is raised for nearly inaudible frequencies, but it ends up
+        # defaulting to -100 db. That default is fine for our purposes.
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', RuntimeWarning)
+            return librosa.A_weighting(frequencies)[:, None] - REF_DB
+
+    def a_weighted(self,audio, sample_rate, hop_length=None, pad=True):
+        """Retrieve the per-frame loudness"""
+        # Save device
+        device = audio.device
+
+        # Default hop length of 10 ms
+        hop_length = sample_rate // 100 if hop_length is None else hop_length
+
+        # Convert to numpy
+        audio = audio.detach().cpu().numpy().squeeze(0)
+
+        # Take stft
+        stft = librosa.stft(audio,
+                            n_fft=WINDOW_SIZE,
+                            hop_length=hop_length,
+                            win_length=WINDOW_SIZE,
+                            center=pad,
+                            pad_mode='constant')
+
+        # Compute magnitude on db scale
+        db = librosa.amplitude_to_db(np.abs(stft))
+
+        # Apply A-weighting
+        weighted = db + self.a_weighted_weights
+
+        # Threshold
+        weighted[weighted < MIN_DB] = MIN_DB
+
+        # Average over weighted frequencies
+        return torch.from_numpy(weighted.mean(axis=0)).float().to(device)[None]
+
+    def __call__(self,
+                 periodicity,
+                 audio,
+                 sample_rate=SAMPLE_RATE,
+                 hop_length=None,
+                 pad=True):
+        # Don't modify in-place
+        periodicity = torch.clone(periodicity)
+
+        # Compute loudness
+        loudness = self.a_weighted(
+            audio, sample_rate, hop_length, pad)
+
+        # Threshold silence
+        periodicity[loudness < self.value] = 0.
+
+        return periodicity
+
+
+
+
+
+
+
+###############################################################################
+# Sequence filters
+###############################################################################
+
+
+def mean(signals, win_length=9):
+    """Averave filtering for signals containing nan values
+
+    Arguments
+        signals (torch.tensor (shape=(batch, time)))
+            The signals to filter
+        win_length
+            The size of the analysis window
+
+    Returns
+        filtered (torch.tensor (shape=(batch, time)))
+    """
+
+    assert signals.dim() == 2, "Input tensor must have 2 dimensions (batch_size, width)"
+    signals = signals.unsqueeze(1)
+
+    # Apply the mask by setting masked elements to zero, or make NaNs zero
+    mask = ~torch.isnan(signals)
+    masked_x = torch.where(mask, signals, torch.zeros_like(signals))
+
+    # Create a ones kernel with the same number of channels as the input tensor
+    ones_kernel = torch.ones(signals.size(1), 1, win_length, device=signals.device)
+
+    # Perform sum pooling
+    sum_pooled = F.conv1d(
+        masked_x,
+        ones_kernel,
+        stride=1,
+        padding=win_length // 2,
+    )
+
+    # Count the non-masked (valid) elements in each pooling window
+    valid_count = F.conv1d(
+        mask.float(),
+        ones_kernel,
+        stride=1,
+        padding=win_length // 2,
+    )
+    valid_count = valid_count.clamp(min=1)  # Avoid division by zero
+
+    # Perform masked average pooling
+    avg_pooled = sum_pooled / valid_count
+
+    # Fill zero values with NaNs
+    avg_pooled[avg_pooled == 0] = float("nan")
+
+    return avg_pooled.squeeze(1)
+
+
+def median(signals, win_length):
+    """Median filtering for signals containing nan values
+
+    Arguments
+        signals (torch.tensor (shape=(batch, time)))
+            The signals to filter
+        win_length
+            The size of the analysis window
+
+    Returns
+        filtered (torch.tensor (shape=(batch, time)))
+    """
+
+    assert signals.dim() == 2, "Input tensor must have 2 dimensions (batch_size, width)"
+    signals = signals.unsqueeze(1)
+
+    mask = ~torch.isnan(signals)
+    masked_x = torch.where(mask, signals, torch.zeros_like(signals))
+    padding = win_length // 2
+
+    x = F.pad(masked_x, (padding, padding), mode="reflect")
+    mask = F.pad(mask.float(), (padding, padding), mode="constant", value=0)
+
+    x = x.unfold(2, win_length, 1)
+    mask = mask.unfold(2, win_length, 1)
+
+    x = x.contiguous().view(x.size()[:3] + (-1,))
+    mask = mask.contiguous().view(mask.size()[:3] + (-1,))
+
+    # Combine the mask with the input tensor
+    x_masked = torch.where(mask.bool(), x.float(), float("inf")).to(x)
+
+    # Sort the masked tensor along the last dimension
+    x_sorted, _ = torch.sort(x_masked, dim=-1)
+
+    # Compute the count of non-masked (valid) values
+    valid_count = mask.sum(dim=-1)
+
+    # Calculate the index of the median value for each pooling window
+    median_idx = ((valid_count - 1) // 2).clamp(min=0)
+
+    # Gather the median values using the calculated indices
+    median_pooled = x_sorted.gather(-1, median_idx.unsqueeze(-1).long()).squeeze(-1)
+
+    # Fill infinite values with NaNs
+    median_pooled[torch.isinf(median_pooled)] = float("nan")
 
-class PredictProgressCallback(Callback):
-    def __init__(self, total_batches,progress_callback: Callable[[int, int], None] = None):
+    return median_pooled.squeeze(1)
+
+
+###############################################################################
+# Utilities
+###############################################################################
+
+
+def nanfilter(signals, win_length, filter_fn):
+    """Filters a sequence, ignoring nan values
+
+    Arguments
+        signals (torch.tensor (shape=(batch, time)))
+            The signals to filter
+        win_length
+            The size of the analysis window
+        filter_fn (function)
+            The function to use for filtering
+
+    Returns
+        filtered (torch.tensor (shape=(batch, time)))
+    """
+    # Output buffer
+    filtered = torch.empty_like(signals)
+
+    # Loop over frames
+    for i in range(signals.size(1)):
+
+        # Get analysis window bounds
+        start = max(0, i - win_length // 2)
+        end = min(signals.size(1), i + win_length // 2 + 1)
+
+        # Apply filter to window
+        filtered[:, i] = filter_fn(signals[:, start:end])
+
+    return filtered
+
+
+def nanmean(signals):
+    """Computes the mean, ignoring nans
+
+    Arguments
+        signals (torch.tensor [shape=(batch, time)])
+            The signals to filter
+
+    Returns
+        filtered (torch.tensor [shape=(batch, time)])
+    """
+    signals = signals.clone()
+
+    # Find nans
+    nans = torch.isnan(signals)
+
+    # Set nans to 0.
+    signals[nans] = 0.
+
+    # Compute average
+    return signals.sum(dim=1) / (~nans).float().sum(dim=1)
+
+
+def nanmedian(signals):
+    """Computes the median, ignoring nans
+
+    Arguments
+        signals (torch.tensor [shape=(batch, time)])
+            The signals to filter
+
+    Returns
+        filtered (torch.tensor [shape=(batch, time)])
+    """
+    # Find nans
+    nans = torch.isnan(signals)
+
+    # Compute median for each slice
+    medians = [nanmedian1d(signal[~nan]) for signal, nan in zip(signals, nans)]
+
+    # Stack results
+    return torch.tensor(medians, dtype=signals.dtype, device=signals.device)
+
+
+def nanmedian1d(signal):
+    """Computes the median. If signal is empty, returns torch.nan
+
+    Arguments
+        signal (torch.tensor [shape=(time,)])
+
+    Returns
+        median (torch.tensor [shape=(1,)])
+    """
+    return torch.median(signal) if signal.numel() else np.nan
+
+
+def dither(cents):
+    """Dither the predicted pitch in cents to remove quantization error"""
+    noise = triang.rvs(c=0.5,loc=-CENTS_PER_BIN,scale=2 * CENTS_PER_BIN,size=cents.size())
+    return cents + cents.new_tensor(noise)
+
+def periodicity(probabilities, bins):
+    """Computes the periodicity from the network output and pitch bins"""
+    # shape=(batch * time / hop_length, 360)
+    probs_stacked = probabilities.transpose(1, 2).reshape(-1, PITCH_BINS)
+
+    # shape=(batch * time / hop_length, 1)
+    bins_stacked = bins.reshape(-1, 1).to(torch.int64)
+
+    # Use maximum logit over pitch bins as periodicity
+    periodicity = probs_stacked.gather(1, bins_stacked)
+
+    # shape=(batch, time / hop_length)
+    return periodicity.reshape(probabilities.size(0), probabilities.size(2))
+
+
+
+
+
+
+
+
+
+
+
+
+
+class CrepeTorch(torch.nn.Module):
+
+    def __init__(self, model_type='full',model_path=None):
         super().__init__()
-        self.total_batches = total_batches
-        self.progress_callback = progress_callback
-    def on_predict_begin(self, logs=None):
-        self.progress_callback(0,self.total_batches)
-    def on_predict_batch_end(self, batch, logs=None):
-        self.progress_callback(batch,self.total_batches)
-    def on_predict_end(self, logs=None):
-        self.progress_callback(self.total_batches,self.total_batches)
+        xx, yy = np.meshgrid(range(360), range(360))
+        transition = np.maximum(12 - abs(xx - yy), 0)
+        self.viterbi_transition = transition / transition.sum(axis=1, keepdims=True)
+        model_type_importance = {'tiny': 4, 'small': 8, 'medium': 16, 'large': 24, 'full': 32}[model_type]
+        out_channels = [n * model_type_importance for n in [32, 4, 4, 4, 8, 16]]
+        in_channels = [n * model_type_importance for n in [32, 4, 4, 4, 8]]
+        in_channels.insert(0,1)
+        self.in_features = 64*model_type_importance
+        # Shared layer parameters
+        kernel_sizes = [(512, 1)] + 5 * [(64, 1)]
+        strides = [(4, 1)] + 5 * [(1, 1)]
+
+        # Overload with eps and momentum conversion given by MMdnn
+        batch_norm_fn = partial(torch.nn.BatchNorm2d,eps=0.0010000000474974513,momentum=0.0)
+
+        # Layer definitions
+        self.conv1 = torch.nn.Conv2d(
+            in_channels=in_channels[0],
+            out_channels=out_channels[0],
+            kernel_size=kernel_sizes[0],
+            stride=strides[0])
+        self.conv1_BN = batch_norm_fn(
+            num_features=out_channels[0])
+
+        self.conv2 = torch.nn.Conv2d(
+            in_channels=in_channels[1],
+            out_channels=out_channels[1],
+            kernel_size=kernel_sizes[1],
+            stride=strides[1])
+        self.conv2_BN = batch_norm_fn(
+            num_features=out_channels[1])
+
+        self.conv3 = torch.nn.Conv2d(
+            in_channels=in_channels[2],
+            out_channels=out_channels[2],
+            kernel_size=kernel_sizes[2],
+            stride=strides[2])
+        self.conv3_BN = batch_norm_fn(
+            num_features=out_channels[2])
+
+        self.conv4 = torch.nn.Conv2d(
+            in_channels=in_channels[3],
+            out_channels=out_channels[3],
+            kernel_size=kernel_sizes[3],
+            stride=strides[3])
+        self.conv4_BN = batch_norm_fn(
+            num_features=out_channels[3])
+
+        self.conv5 = torch.nn.Conv2d(
+            in_channels=in_channels[4],
+            out_channels=out_channels[4],
+            kernel_size=kernel_sizes[4],
+            stride=strides[4])
+        self.conv5_BN = batch_norm_fn(
+            num_features=out_channels[4])
+
+        self.conv6 = torch.nn.Conv2d(
+            in_channels=in_channels[5],
+            out_channels=out_channels[5],
+            kernel_size=kernel_sizes[5],
+            stride=strides[5])
+        self.conv6_BN = batch_norm_fn(
+            num_features=out_channels[5])
+
+        self.classifier = torch.nn.Linear(
+            in_features=self.in_features,
+            out_features=PITCH_BINS)
+        if not model_path:
+            model_path = hf_hub_download("shethjenil/Audio2Midi_Models",f"crepe_{model_type}.pt")
+        self.load_state_dict(torch.load(model_path))
+        self.eval()
+
+    def forward(self, x, embed=False):
+        # Forward pass through first five layers
+        x = self.embed(x)
+
+        if embed:
+            return x
+
+        # Forward pass through layer six
+        x = self.layer(x, self.conv6, self.conv6_BN)
+
+        # shape=(batch, self.in_features)
+        x = x.permute(0, 2, 1, 3).reshape(-1, self.in_features)
+
+        # Compute logits
+        return torch.sigmoid(self.classifier(x))
+
+    def embed(self, x):
+        """Map input audio to pitch embedding"""
+        # shape=(batch, 1, 1024, 1)
+        x = x[:, None, :, None]
 
+        # Forward pass through first five layers
+        x = self.layer(x, self.conv1, self.conv1_BN, (0, 0, 254, 254))
+        x = self.layer(x, self.conv2, self.conv2_BN)
+        x = self.layer(x, self.conv3, self.conv3_BN)
+        x = self.layer(x, self.conv4, self.conv4_BN)
+        x = self.layer(x, self.conv5, self.conv5_BN)
+        return x
+
+    def layer(self, x, conv, batch_norm, padding=(0, 0, 31, 32)):
+        """Forward pass through one layer"""
+        x = F.pad(x, padding)
+        x = conv(x)
+        x = F.relu(x)
+        x = batch_norm(x)
+        return F.max_pool2d(x, (2, 1), (2, 1))
+
+    def viterbi(self,logits):
+        """Sample observations using viterbi decoding"""
+        # Normalize logits
+        with torch.no_grad():
+            probs = torch.nn.functional.softmax(logits, dim=1)
+
+        # Convert to numpy
+        sequences = probs.cpu().numpy()
+
+        # Perform viterbi decoding
+        bins = np.array([
+            librosa.sequence.viterbi(sequence, self.viterbi_transition).astype(np.int64)
+            for sequence in sequences])
+
+        # Convert to pytorch
+        bins = torch.tensor(bins, device=probs.device)
+
+        # Convert to frequency in Hz
+        return bins, bins_to_frequency(bins)
+
+    def get_device(self):
+        return next(self.parameters()).device
+
+    def preprocess(self,audio,sample_rate,hop_length=None,batch_size=None,pad=True):
+        """Convert audio to model input
+
+        Arguments
+            audio (torch.tensor [shape=(1, time)])
+                The audio signals
+            sample_rate (int)
+                The sampling rate in Hz
+            hop_length (int)
+                The hop_length in samples
+            batch_size (int)
+                The number of frames per batch
+            pad (bool)
+                Whether to zero-pad the audio
+
+        Returns
+            frames (torch.tensor [shape=(1 + int(time // hop_length), 1024)])
+        """
+        # Default hop length of 10 ms
+        hop_length = sample_rate // 100 if hop_length is None else hop_length
+
+        # Get total number of frames
+
+        # Maybe pad
+        if pad:
+            total_frames = 1 + int(audio.size(1) // hop_length)
+            audio = torch.nn.functional.pad(
+                audio,
+                (WINDOW_SIZE // 2, WINDOW_SIZE // 2))
+        else:
+            total_frames = 1 + int((audio.size(1) - WINDOW_SIZE) // hop_length)
+
+        # Default to running all frames in a single batch
+        batch_size = total_frames if batch_size is None else batch_size
+
+        # Generate batches
+        for i in range(0, total_frames, batch_size):
+
+            # Batch indices
+            start = max(0, i * hop_length)
+            end = min(audio.size(1),
+                    (i + batch_size - 1) * hop_length + WINDOW_SIZE)
+
+            # Chunk
+            frames = torch.nn.functional.unfold(
+                audio[:, None, None, start:end],
+                kernel_size=(1, WINDOW_SIZE),
+                stride=(1, hop_length))
+
+            # shape=(1 + int(time / hop_length, 1024)
+            frames = frames.transpose(1, 2).reshape(-1, WINDOW_SIZE)
+
+            # Place on device
+            frames = frames.to(self.get_device())
+
+            # Mean-center
+            frames -= frames.mean(dim=1, keepdim=True)
+
+            # Scale
+            # Note: during silent frames, this produces very large values. But
+            # this seems to be what the network expects.
+            frames /= torch.max(torch.tensor(1e-10, device=frames.device),frames.std(dim=1, keepdim=True))
+
+            yield frames
+
+    def postprocess(self,probabilities,fmin=0.,fmax=MAX_FMAX,return_periodicity=False):
+        """Convert model output to F0 and periodicity
+
+        Arguments
+            probabilities (torch.tensor [shape=(1, 360, time / hop_length)])
+                The probabilities for each pitch bin inferred by the network
+            fmin (float)
+                The minimum allowable frequency in Hz
+            fmax (float)
+                The maximum allowable frequency in Hz
+            viterbi (bool)
+                Whether to use viterbi decoding
+            return_periodicity (bool)
+                Whether to also return the network confidence
+
+        Returns
+            pitch (torch.tensor [shape=(1, 1 + int(time // hop_length))])
+            periodicity (torch.tensor [shape=(1, 1 + int(time // hop_length))])
+        """
+        # Sampling is non-differentiable, so remove from graph
+        probabilities = probabilities.detach()
+
+        # Convert frequency range to pitch bin range
+        minidx = frequency_to_bins(torch.tensor(fmin))
+        maxidx = frequency_to_bins(torch.tensor(fmax),
+                                                    torch.ceil)
+
+        # Remove frequencies outside of allowable range
+        probabilities[:, :minidx] = -float('inf')
+        probabilities[:, maxidx:] = -float('inf')
+
+        # Perform argmax or viterbi sampling
+        bins, pitch = self.viterbi(probabilities)
+
+        if not return_periodicity:
+            return pitch
+
+        # Compute periodicity from probabilities and decoded pitch bins
+        return pitch, periodicity(probabilities, bins)
+
+    def predict(self,audio,sample_rate,hop_length=None,fmin=50.,fmax=MAX_FMAX,return_periodicity=False,batch_size=None,pad=True):
+        """Performs pitch estimation
+
+        Arguments
+            audio (torch.tensor [shape=(1, time)])
+                The audio signal
+            sample_rate (int)
+                The sampling rate in Hz
+            hop_length (int)
+                The hop_length in samples
+            fmin (float)
+                The minimum allowable frequency in Hz
+            fmax (float)
+                The maximum allowable frequency in Hz
+            return_periodicity (bool)
+                Whether to also return the network confidence
+            batch_size (int)
+                The number of frames per batch
+            pad (bool)
+                Whether to zero-pad the audio
+
+        Returns
+            pitch (torch.tensor [shape=(1, 1 + int(time // hop_length))])
+            (Optional) periodicity (torch.tensor
+                                    [shape=(1, 1 + int(time // hop_length))])
+        """
+
+        results = []
+        with torch.no_grad():
+            print("prediction started")
+            for frames in self.preprocess(audio,sample_rate,hop_length,batch_size,pad):
+                # shape=(batch, 360, time / hop_length)
+                result = self.postprocess(self.forward(frames, embed=False).reshape(audio.size(0), -1, PITCH_BINS).transpose(1, 2),fmin,fmax,return_periodicity)
+                if isinstance(result, tuple):
+                    result = (result[0].to(audio.device),result[1].to(audio.device))
+                else:
+                    result = result.to(audio.device)
+                results.append(result)
+            print("prediction finished")
+        if return_periodicity:
+            pitch, periodicity = zip(*results)
+            return torch.cat(pitch, 1), torch.cat(periodicity, 1)
+        return torch.cat(results, 1)
+
+    def predict_from_file(self,audio_file,hop_length=None,fmin=50.,fmax=MAX_FMAX,return_periodicity=False,batch_size=None,pad=True):
+        audio, sample_rate = librosa.load(audio_file, sr=SAMPLE_RATE)
+        return self.predict(torch.from_numpy(audio).unsqueeze(0),sample_rate,hop_length,fmin,fmax,return_periodicity,batch_size,pad)
+
+    def predict_from_file_to_file(self,audio_file,output_pitch_file,output_periodicity_file=None,hop_length=None,fmin=50.,fmax=MAX_FMAX,batch_size=None,pad=True):
+        prediction = self.predict_from_file(audio_file,hop_length,fmin,fmax,False,output_periodicity_file is not None,batch_size,pad)
+        if output_periodicity_file is not None:
+            torch.save(prediction[0].detach(), output_pitch_file)
+            torch.save(prediction[1].detach(), output_periodicity_file)
+        else:
+            torch.save(prediction.detach(), output_pitch_file)
+
+    def predict_from_files_to_files(self,audio_files,output_pitch_files,output_periodicity_files=None,hop_length=None,fmin=50.,fmax=MAX_FMAX,batch_size=None,pad=True):
+        if output_periodicity_files is None:
+            output_periodicity_files = len(audio_files) * [None]
+        for audio_file, output_pitch_file, output_periodicity_file in tqdm(zip(audio_files, output_pitch_files, output_periodicity_files), desc='torchcrepe', dynamic_ncols=True):
+            self.predict_from_file_to_file(audio_file,output_pitch_file,None,output_periodicity_file,hop_length,fmin,fmax,batch_size,pad)
+
+    def embedding(self,audio,sample_rate,hop_length=None,batch_size=None,pad=True):
+        """Embeds audio to the output of CREPE's fifth maxpool layer
+
+        Arguments
+            audio (torch.tensor [shape=(1, time)])
+                The audio signals
+            sample_rate (int)
+                The sampling rate in Hz
+            hop_length (int)
+                The hop_length in samples
+            batch_size (int)
+                The number of frames per batch
+            pad (bool)
+                Whether to zero-pad the audio
+
+        Returns
+            embedding (torch.tensor [shape=(1,
+                                            1 + int(time // hop_length), 32, -1)])
+        """
+        # shape=(batch, time / hop_length, 32, embedding_size)
+        with torch.no_grad():
+            return torch.cat([self.forward(frames, embed=True).reshape(audio.size(0), frames.size(0), 32, -1).to(audio.device) for frames in self.preprocess(audio,sample_rate,hop_length,batch_size,pad)], 1)
+
+    def embedding_from_file(self,audio_file,hop_length=None,batch_size=None,pad=True):
+        audio, sample_rate = librosa.load(audio_file, sr=SAMPLE_RATE)
+        return self.embed(torch.from_numpy(audio).unsqueeze(0),sample_rate,hop_length,batch_size,pad)
+
+    def embedding_from_file_to_file(self,audio_file,output_file,hop_length=None,batch_size=None,pad=True):
+        with torch.no_grad():
+            torch.save(self.embed_from_file(audio_file,hop_length,batch_size,pad).detach(), output_file)
+
+    def embedding_from_files_to_files(self,audio_files,output_files,hop_length=None,batch_size=None,pad=True):
+        for audio_file, output_file in tqdm(zip(audio_files, output_files), desc='torchcrepe', dynamic_ncols=True):
+            self.embed_from_file_to_file(audio_file,output_file,hop_length,batch_size,pad)
+
+
+
+
+
+
+
+
+
+
+
+
+from hmmlearn.hmm import CategoricalHMM
+from typing import Callable
+from numpy.lib.stride_tricks import as_strided
+from pretty_midi_fix import PrettyMIDI , PitchBend , Note ,Instrument
+from huggingface_hub import hf_hub_download
+from torch.utils.data import DataLoader, TensorDataset
 
 class Crepe():
+
     def __init__(self,model_type="full",model_path=None):
-        if not model_path:
-            model_path = hf_hub_download("shethjenil/Audio2Midi_Models",f"crepe_{model_type}.h5")
-        model_type_importance = {'tiny': 4, 'small': 8, 'medium': 16, 'large': 24, 'full': 32}[model_type]
-        filters = [n * model_type_importance for n in [32, 4, 4, 4, 8, 16]]
-        widths = [512, 64, 64, 64, 64, 64]
-        strides = [(4, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1)]
-        x = Input(shape=(1024,), name='input', dtype='float32')
-        y = Reshape(target_shape=(1024, 1, 1), name='input-reshape')(x)
-        layers = [1, 2, 3, 4, 5, 6]
-        for l, f, w, s in zip(layers, filters, widths, strides):
-            y = Conv2D(f, (w, 1), strides=s, padding='same', activation='relu', name="conv%d" % l)(y)
-            y = BatchNormalization(name="conv%d-BN" % l)(y)
-            y = MaxPool2D(pool_size=(2, 1), strides=None, padding='valid', name="conv%d-maxpool" % l)(y)
-            y = Dropout(0.25, name="conv%d-dropout" % l)(y)
-        y = Permute((2, 1, 3), name="transpose")(y)
-        y = Flatten(name="flatten")(y)
-        y = Dense(360, activation='sigmoid', name="classifier")(y)
-        self.model = Model(inputs=x, outputs=y)
-        self.model.load_weights(model_path)
-        self.model.compile('adam', 'binary_crossentropy')
+        self.model = CrepeTorch(model_type,model_path)
         self.cents_mapping=(np.linspace(0, 7180, 360) + 1997.3794084376191)
 
     def to_local_average_cents(self, salience, center=None):
+        if isinstance(salience, torch.Tensor):
+            salience = salience.numpy()
+
         if salience.ndim == 1:
             if center is None:
                 center = int(np.argmax(salience))
@@ -63,6 +881,8 @@ class Crepe():
         raise Exception("label should be either 1d or 2d ndarray")
 
     def to_viterbi_cents(self,salience):
+        if isinstance(salience, torch.Tensor):
+            salience = salience.numpy()
         starting = np.ones(360) / 360
         xx, yy = np.meshgrid(range(360), range(360))
         transition = np.maximum(12 - abs(xx - yy), 0)
@@ -84,11 +904,21 @@ class Crepe():
         frames = frames.transpose().copy()
         frames -= np.mean(frames, axis=1)[:, np.newaxis]
         frames /= np.clip(np.std(frames, axis=1)[:, np.newaxis], 1e-8, None)
-        return self.model.predict(frames,batch_size,0,callbacks=[PredictProgressCallback(math_ceil(len(frames) / batch_size),progress_callback)])
-    
+        device = self.model.get_device()
+        all_outputs = []
+        all_batch = list(DataLoader(TensorDataset(torch.from_numpy(frames)), batch_size=batch_size, shuffle=False))
+        total_batch = len(all_batch)
+        with torch.no_grad():
+            for i , batch in enumerate(all_batch):
+                inputs = batch[0].to(device)
+                outputs = self.model(inputs)
+                all_outputs.append(outputs.cpu())
+                progress_callback(i,total_batch)
+        return torch.cat(all_outputs, dim=0)
+
     def model_predict(self,audio:np.ndarray,viterbi, center, step_size,progress_callback,batch_size):
         activation = self.get_activation(audio.astype(np.float32), center, step_size,progress_callback,batch_size)
-        confidence = activation.max(axis=1)
+        confidence = activation.max(axis=1).values # Access the values from the named tuple
         cents = self.to_viterbi_cents(activation) if viterbi else self.to_local_average_cents(activation)
         frequency = 10 * 2 ** (cents / 1200)
         frequency[np.isnan(frequency)] = 0
@@ -96,7 +926,7 @@ class Crepe():
         return time, frequency, confidence
 
     def predict(self,audio_path,viterbi=False, center=True, step_size=10,min_confidence=0.8,batch_size=32,progress_callback: Callable[[int, int], None] = None,output_file= "output.mid"):
-        time, frequency, confidence = self.model_predict(librosa_load(audio_path, sr=16000, mono=True)[0],viterbi,center,step_size,progress_callback,batch_size)
+        time, frequency, confidence = self.model_predict(librosa.load(audio_path, sr=16000, mono=True)[0],viterbi,center,step_size,progress_callback,batch_size)
         mask = confidence > min_confidence
         times = time[mask]
         frequencies = frequency[mask]