labhamlet/wavjepa-base

1import torch2 3def normalize(audio):4    mean = audio.mean(dim=(-2, -1), keepdim=True)5    std = audio.std(dim=(-2, -1), keepdim=True)6    audio = (audio - mean) / (std + 1e-5) # Add epsilon for stability7    return audio8 9def calculate_padding_mask(pad_frames, total_frames, sr, output_steps, process_seconds, device, B):10    # How many 2 seconds chunks does this audio have?11    # Find it and then multiply by the output_steps.12    total_frames = int((total_frames / sr) / process_seconds)13    total_output_steps = output_steps * total_frames14    mask = torch.zeros((B, total_output_steps), dtype = torch.bool, device = device)15 16    # Check the number of padding tokens that we have in the audio.17    output_sr = int(output_steps / process_seconds)18    pad_seconds = pad_frames / sr19    pad_steps = int(pad_seconds * output_sr)20    # Create the mask21 22    mask[..., total_output_steps - pad_steps:] = True 23    return mask, total_output_steps - pad_steps24 25 26def get_timestamps(sample_rate, B, input_audio_len, x):27    audio_len = input_audio_len28    sec = audio_len / sample_rate29    x_len = x.shape[1]30    step = sec / x_len * 1000  # sec -> ms31    ts = torch.tensor([step * i for i in range(x_len)]).unsqueeze(0)32    ts = ts.repeat(B, 1)33    return ts