data/stemmed_dataset.py

import random
from typing import Dict, Iterable, List, Optional, Sequence, Set, Tuple
import numpy as np
import torch.utils.data
from pathlib import Path
import torch
import itertools

from conditioning.beat_embedder import Beat
from data.stem import Stem
from utils import audio as audio_utils
import json
from torch import Tensor
import torchaudio
# from collections.abc import Sized
import librosa

# import pyrubberband as pyrb
# import pylibrb

N_VALID_SAMPLES = 24
# 014f37 is removed
# EXPECTED_N_SONGS = 240
# EXPECTED_N_SONGS = 150


class StemmedDataset(torch.utils.data.Dataset):

    def __init__(self,
                 root_dir: Path,
                 stems: Set[Stem],
                 target_stem: Stem,
                 single_stem: bool,
                 min_context_seconds: int,
                 use_style_conditioning: bool,
                 use_beat_conditioning: bool,
                 type_of_context: str,
                 bpm_in_caption: bool,
                 add_click: bool,
                 sync_chunks: bool,
                 train: bool,
                 sample_rate: int,
                 chunk_size_samples: int,
                 speed_transform_p: float,
                 pitch_transform_p: float,
                 stereo: bool = False,
                 max_genres_in_description: int = 3,
                 max_moods_in_description: int = 3,
                 n_samples_per_epoch: Optional[int] = None,
                 verbose: bool = False):
        self.root_dir: Path = root_dir
        self.stems: Set[Stem] = stems
        self.single_stem: bool = single_stem
        self.min_context_seconds: int = min_context_seconds
        self.target_stem: Stem = target_stem
        self.use_style_conditioning: bool = use_style_conditioning
        self.use_beat_conditioning: bool = use_beat_conditioning
        if self.use_style_conditioning and not self.single_stem:
            raise ValueError("You can only use style conditioning if "
                             "the target is a single stem")

        if self.target_stem != Stem.ANY:
            assert self.single_stem

        self.add_click: bool = add_click
        self.bpm_in_caption: bool = bpm_in_caption
        self.sync_chunks: bool = sync_chunks
        self.stem_names: Set[str] = {s.getname() for s in self.stems}
        self.train: bool = train
        self.sample_rate: int = sample_rate
        self.chunk_size_samples: int = chunk_size_samples
        self.speed_transform_p: float = speed_transform_p
        self.pitch_transform_p: float = pitch_transform_p
        self.stereo: bool = stereo
        self.max_genres_in_description: int = max_genres_in_description
        self.max_moods_in_description: int = max_moods_in_description
        self.verbose: bool = verbose

        self.type_of_context: str = type_of_context
        assert self.type_of_context in ["stems", "beats", "stems or beats"]

        if self.add_click or self.sync_chunks:
            if not (self.root_dir / "sync.json").exists():
                raise FileNotFoundError(
                    "If you want click or sync, I need a 'sync.json' file in "
                    "the top-level dir of the dataset")

            with open(self.root_dir / "sync.json", "r") as f:
                self.syncdata: Dict[str, List[int]] = json.load(f)

        # load all song directories
        self.song_names: List[str] = sorted(
            [p.name for p in self.root_dir.iterdir() if p.is_dir()])
        # assert len(self.song_names) == EXPECTED_N_SONGS

        # if has a single target stem filter out songs that don't have that stem
        if self.target_stem != Stem.ANY:
            toremove: List[str] = []
            for song_name in self.song_names:
                if not (self.root_dir / song_name /
                        self.target_stem.getname()).exists():
                    toremove.append(song_name)
            self.song_names = [s for s in self.song_names if s not in toremove]

        # train/valid split
        if self.train:
            self.song_names = self.song_names[:-N_VALID_SAMPLES]
        else:
            self.song_names = self.song_names[-N_VALID_SAMPLES:]

        # if self.verbose:
        print(f"Loaded {len(self.song_names)} for "
              f"{'train' if self.train else 'valid'} dataset.")

        # WTF? TODO: remove
        # if self.target_stem != Stem.ANY:
        #     # if only interested in a stem, remove songs without it
        #     for songname in self.song_names:
        #         songdir = self.root_dir / songname
        #         if not (songdir / self.target_stem.getname()).exists():
        #             self.song_names.remove(songname)

        # create iterator to run n_sample times
        self.n_samples: int = n_samples_per_epoch or len(self.song_names)

        self.song_iterator: List[str] = list(
            itertools.islice(itertools.cycle(iter(self.song_names)),
                             self.n_samples))

    def __len__(self):
        return self.n_samples

    def save_sample(self, sample: Dict, path: Path):
        audio_utils.save_audio(sample["wav"], path / "input.wav",
                               self.sample_rate)
        audio_utils.save_audio(sample["conditioning"].wav, path / "cond.wav",
                               self.sample_rate)
        audio_utils.save_audio(sample["conditioning"].wav + sample["wav"],
                               path / "mix.wav", self.sample_rate)

    def get_description(self,
                        features: Dict[str, str | int],
                        instruments: Sequence[Stem],
                        speed_factor: Optional[float] = None) -> str:

        genres: List[str] = str(features["genres"]).split(",")
        moods: List[str] = str(features["moods"]).split(",")
        description = ""

        # Genre
        if len(genres) > 0:
            # if more than max number of genres, choose first few
            # if len(genres) > self.max_genres:
            #     genres = random.sample(genres, self.max_genres)
            genres = [
                s.strip() for s in genres[:self.max_genres_in_description]
            ]
            description += (f"Genre{'s' if len(genres) > 1 else ''}: "
                            f"{', '.join(genres)}. ")

        # Mood
        if len(moods) > 0:
            # if more than max number of moods, choose first few
            # if len(moods) > self.max_moods:
            #     moods = random.sample(moods, self.max_moods)
            moods = [s.strip() for s in moods[:self.max_moods_in_description]]
            description += (f"Mood{'s' if len(moods) > 1 else ''}: "
                            f"{', '.join(moods)}. ")

        # Instruments
        if not (self.single_stem and self.target_stem != Stem.ANY):
            instrument_names: List[str] = [s.name.lower() for s in instruments]
            random.shuffle(instrument_names)
            description += f"Instruments: {', '.join(instrument_names)}."

        # BPM
        if self.bpm_in_caption:
            bpm = features["bpm"]
            if speed_factor:
                bpm = round(int(bpm) / speed_factor)
            description += f" Bpm: {bpm}."

        # Key
        # key = features["key"]
        # description += f"Key: {key}."

        return description

    def _transform_chunk(self, t: Tensor, speed_factor: float,
                         pitch_factor: int, target_size: int):
        if self.stereo:
            raise NotImplementedError(
                "No augmentations for stereo audio implemented")

        stretched = audio_utils.stretch_with_timeout(
            t,
            self.sample_rate,
            speed_factor,
            pitch_factor,
            2,
        )

        if stretched.shape[-1] < target_size:
            stretched = torch.nn.functional.pad(
                stretched, (0, target_size - stretched.shape[-1]), "constant",
                0)
        elif stretched.shape[-1] > target_size:
            stretched = stretched[..., :target_size]

        return stretched

    def load_stems(
        self,
        song_path: Path,
        song_stems: Iterable[Stem],
        start_offset: int,
        n_frames: int,
    ) -> Dict[Stem, Tensor]:

        # load audio chunks for each stem
        stem_tensors: Dict[Stem, Tensor] = {}
        for stem in song_stems:
            stemdir = song_path / stem.getname()
            stem_tensor: Tensor = torch.zeros(2 if self.stereo else 1,
                                              n_frames,
                                              dtype=torch.float32)

            # for each track of a stem
            for trackpath in stemdir.iterdir():
                # load wav
                chunk = audio_utils.load_audio_chunk(trackpath,
                                                     start_offset,
                                                     n_frames,
                                                     stereo=self.stereo)
                stem_tensor += chunk

            # if not silent, include it in dictionary of stems
            if not audio_utils.is_silent(stem_tensor, threshold=0.01):
                stem_tensors[stem] = stem_tensor

        return stem_tensors

    # def choose_conditioning(
    #         self,
    #         stems: Sequence[Stem]) -> Tuple[Sequence[Stem], Sequence[Stem]]:
    #     n_stems: int = len(stems)
    #     if n_stems == 1:
    #         # if only one stem, use it as input with no conditioning
    #         return stems[:], []
    #     n_conditioning_stems: int = random.randint(1, n_stems - 1)
    #     conditioning_stems: Sequence[Stem] = random.sample(
    #         stems, n_conditioning_stems)
    #     input_stems = list(filter(lambda x: x not in conditioning_stems, stems))
    #     return input_stems, conditioning_stems

    def choose_input_and_conditioning(
            self,
            stems: Sequence[Stem]) -> Tuple[Sequence[Stem], Sequence[Stem]]:
        n_stems: int = len(stems)

        if n_stems == 1:
            raise RuntimeError("This song has only one stem")
            return stems[:], []

        # choose a random number of context stems,
        # leaving at least 1 for the input
        if self.target_stem != Stem.ANY:
            assert self.target_stem in stems
            possible_conditioning_stems = [
                s for s in stems if s != self.target_stem
            ]
            n_conditioning_stems: int = random.randint(
                1, len(possible_conditioning_stems))
            conditioning_stems: Sequence[Stem] = random.sample(
                possible_conditioning_stems, n_conditioning_stems)
        else:
            n_conditioning_stems: int = random.randint(1, n_stems - 1)
            conditioning_stems: Sequence[Stem] = random.sample(
                stems, n_conditioning_stems)

        # choose a random number of the remaining stems as input
        n_input_stems: int = 1 if self.single_stem else random.randint(
            1, n_stems - n_conditioning_stems)

        # choose input stem
        if self.target_stem != Stem.ANY:
            # if target stem != ANY, that HAS to be the input
            assert self.target_stem in stems
            input_stems: Sequence[Stem] = [self.target_stem]
        else:
            possible_input_stems: Sequence[Stem] = list(
                filter(lambda x: x not in conditioning_stems, stems))
            input_stems: Sequence[Stem] = random.sample(possible_input_stems,
                                                        n_input_stems)

        return input_stems, conditioning_stems

    def add_click_to_track(self, wav: Tensor, wav_sr: int,
                           click_frames: Sequence[int],
                           start_offset: int) -> Tensor:

        for idx, frame in enumerate(click_frames):
            if frame >= start_offset:
                first_relevant_index = idx
                break
        else:
            return wav

        shifted_click_frames: List[int] = [
            f - start_offset for f in click_frames[first_relevant_index:]
        ]

        click_track: Tensor = audio_utils.create_click(list(wav.shape), wav_sr,
                                                       shifted_click_frames)
        assert click_track.shape == wav.shape

        mix: Tensor = wav + click_track * 0.5
        return mix

    def mix_input_and_conditioning(self, stem_tensors: Dict[Stem, Tensor],
                                   input_stems: Sequence[Stem],
                                   condition_stems: Sequence[Stem]):

        assert len(input_stems) > 0
        input_tensor = torch.stack([stem_tensors[s] for s in input_stems
                                   ]).sum(dim=-0)

        if len(condition_stems) > 0:
            condition_tensor = torch.stack(
                [stem_tensors[s] for s in condition_stems]).sum(dim=-0)
        else:
            condition_tensor = None

        return input_tensor, condition_tensor

    def find_good_chunk(
            self, song_name: str, n_frames_to_take: int, song_n_frames: int,
            song_path: Path,
            song_stems: Iterable[Stem]) -> Tuple[Dict[Stem, Tensor], int]:
        found_good_chunk: bool = False
        attempts: int = 0
        while not found_good_chunk:
            attempts += 1
            if attempts > 10 and (attempts - 1) % 10 == 0:
                print(
                    f"Tried to find some non-silent chunk of song {song_name} "
                    f"for {attempts} times but it's so hard please master "
                    "I am tired let me rest")

            # choose random chunk
            if self.sync_chunks:
                choices: List[int] = self.syncdata[song_name]
                choices = list(
                    filter(lambda x: (x + n_frames_to_take) < song_n_frames,
                           choices))
                start_offset: int = random.choice(choices)

            else:
                start_offset: int = random.randint(
                    0, song_n_frames - n_frames_to_take)

            # load song stems, filter out silent ones
            stem_tensors: Dict[Stem, Tensor] = self.load_stems(
                song_path,
                song_stems,
                start_offset,
                n_frames_to_take,
            )
            nonsilent_stems: List[Stem] = list(stem_tensors.keys())

            if self.target_stem != Stem.ANY:
                if self.target_stem in nonsilent_stems and (
                        len(nonsilent_stems)
                        > (1 if len(list(song_stems)) > 1 else 0)):
                    found_good_chunk = True
            else:
                if len(nonsilent_stems) > 0:
                    found_good_chunk = True

        return stem_tensors, start_offset  # type: ignore

    def __getitem__(self, idx: int) -> Dict[str, Tensor | str]:
        # output = {
        #     "name": "",
        #     "target": torch.rand(1, 320_000),
        #     "description": "",
        #     "context": torch.rand(1, 320_000),
        #     "style": torch.rand(1, 320_000),
        #     "beat": Beat(beats, downbeats, seq_len)
        # }
        # return output

        song_name = self.song_iterator[idx]
        song_path: Path = self.root_dir / song_name

        # toss a coin to decide whether to augment data
        apply_speed_transform: bool = random.random() < self.speed_transform_p
        apply_pitch_transform: bool = random.random() < self.pitch_transform_p

        # choose random augmentation factors if needed
        speed_factor = round(random.random() * 0.4 +
                             0.80, 2) if apply_speed_transform else 1
        pitch_factor = (random.randint(-4, 4) if apply_pitch_transform else 0)

        # get song features
        with open(song_path / "features.json", "r") as f:
            features: Dict[str, str | int] = json.load(f)
        song_n_frames: int = int(features["num_frames"])
        song_sr: int = int(features["sample_rate"])
        n_frames_to_take_orig: int = int(self.chunk_size_samples /
                                         self.sample_rate * song_sr)
        n_frames_to_take = int(n_frames_to_take_orig / speed_factor)

        # get all song stems, including possibly silent ones
        song_stems: List[Stem] = []
        for stem in self.stems:
            if (song_path / stem.getname()).exists():
                song_stems.append(stem)

        if self.target_stem != Stem.ANY and self.target_stem not in song_stems:
            raise RuntimeError(f"Target stem is {self.target_stem} but song "
                               f"{song_name} has no interesting stems. "
                               "Maybe remove it from the dataset?")

        if len(song_stems) == 0:
            raise RuntimeError(f"Song {song_name} has no interesting stems. "
                               "Maybe remove it from the dataset?")
        if len(song_stems) == 1:
            raise RuntimeError(f"Song {song_name} has only one stem, which is"
                               f"{song_stems[0]}. What to do?")

        # find a good chunk of the song
        stem_tensors: Dict[Stem, Tensor]
        start_offset: int
        stem_tensors, start_offset = self.find_good_chunk(
            song_name, n_frames_to_take, song_n_frames, song_path, song_stems)
        nonsilent_stems: List[Stem] = list(stem_tensors.keys())

        # split stems between input and conditioning
        input_stems, condition_stems = self.choose_input_and_conditioning(
            nonsilent_stems)

        # mix input and condition tensors
        input_tensor, condition_tensor = self.mix_input_and_conditioning(
            stem_tensors, input_stems, condition_stems)

        # if applying style conditioning, find a good style conditioning chunk
        style_tensor: Optional[Tensor] = None
        if self.use_style_conditioning:
            assert self.single_stem
            assert len(input_stems) == 1
            inputstem = input_stems[0]
            style_tensor = self.find_good_chunk(song_name,
                                                n_frames_to_take_orig,
                                                song_n_frames, song_path,
                                                [inputstem])[0][inputstem]

        # if using beat conditioning, compute beats data for current chunk
        beats_conditioning: Optional[Beat] = None
        if self.use_beat_conditioning:
            beatfile = song_path / "beatthis.npz"
            if not beatfile.exists():
                raise FileNotFoundError(
                    f"Couldn't find beat annotations for song {song_path}")
            loaded = np.load(beatfile)
            beats_sec = torch.from_numpy(loaded["beats"])
            downbeats_sec = torch.from_numpy(loaded["downbeats"])

            # if using a speed augmentation, reposition beats
            beats_frames: Tensor = (beats_sec * speed_factor *
                                    self.sample_rate).round().long()
            downbeats_frames: Tensor = (downbeats_sec * speed_factor *
                                        self.sample_rate).round().long()
            start_offset = round(start_offset / song_sr * self.sample_rate *
                                 speed_factor)
            min_max_beat: Tensor = torch.tensor(
                [start_offset, start_offset + self.chunk_size_samples])

            beats_start_idx, beats_end_idx = torch.searchsorted(beats_frames,
                                                                min_max_beat,
                                                                right=False)
            beats_cut = beats_frames[beats_start_idx:beats_end_idx]

            downbeats_start_idx, downbeats_end_idx = torch.searchsorted(
                downbeats_frames, min_max_beat, right=False)
            downbeats_cut = downbeats_frames[
                downbeats_start_idx:downbeats_end_idx]

            assert (downbeats_end_idx + 1 >= len(downbeats_frames) or
                    downbeats_frames[downbeats_end_idx + 1]
                    >= start_offset + self.chunk_size_samples)
            assert (beats_end_idx + 1 >= len(beats_frames) or
                    beats_frames[beats_end_idx + 1]
                    >= start_offset + self.chunk_size_samples)

            beats_cut -= start_offset
            downbeats_cut -= start_offset

            beats_conditioning = Beat(beats_cut, downbeats_cut,
                                      self.chunk_size_samples)

            match self.type_of_context:
                case "beats":
                    beats_as_context = True
                case "stems":
                    beats_as_context = False
                case "stems or beats":
                    beats_as_context = random.random() < 0.5

            beats_time = beats_cut.numpy() / self.sample_rate
            downbeats_time = downbeats_cut.numpy() / self.sample_rate
            clicks = librosa.clicks(times=beats_time,
                                    sr=self.sample_rate,
                                    click_freq=1000,
                                    length=self.chunk_size_samples)
            downbeat_clicks = librosa.clicks(times=downbeats_time,
                                             sr=self.sample_rate,
                                             click_freq=1000,
                                             length=self.chunk_size_samples)
            # Combine clicks (downbeats are stronger)
            audio_cliks = clicks + downbeat_clicks
            audio_cliks = np.clip(audio_cliks, -1.0, 1.0)
            # mono audio
            if audio_cliks.ndim > 1:
                audio_cliks = audio_cliks.mean(axis=0)
            context_beats = torch.tensor(audio_cliks)
            context_beats = torch.unsqueeze(context_beats, 0)
            if beats_as_context:
                condition_tensor = context_beats

        # add click
        if self.add_click:
            raise NotImplementedError(
                "Add click is not implemented for new dataset with style")
            click_frames: List[int] = self.syncdata[song_name]
            input_tensor = self.add_click_to_track(input_tensor, song_sr,
                                                   click_frames, start_offset)
            if condition_tensor is not None:
                condition_tensor = self.add_click_to_track(
                    condition_tensor, song_sr, click_frames, start_offset)

        # get description of song input
        description: str = self.get_description(features, input_stems,
                                                speed_factor)

        # resample input and conditioning to desired sample rate
        input_tensor = torchaudio.functional.resample(input_tensor, song_sr,
                                                      self.sample_rate)
        if condition_tensor is not None and not beats_as_context:
            condition_tensor = torchaudio.functional.resample(
                condition_tensor, song_sr, self.sample_rate)

        if style_tensor is not None:
            style_tensor = torchaudio.functional.resample(
                style_tensor, song_sr, self.sample_rate)

        # data augmentation to input and conditioning
        if apply_speed_transform or apply_pitch_transform:
            input_tensor = self._transform_chunk(input_tensor, speed_factor,
                                                 pitch_factor,
                                                 self.chunk_size_samples)
            if condition_tensor is not None and not beats_as_context:
                condition_tensor = self._transform_chunk(
                    condition_tensor, speed_factor, pitch_factor,
                    self.chunk_size_samples)

        if condition_tensor is not None:
            # cut conditioning to a random length
            min_context_samples = self.min_context_seconds * self.sample_rate
            if min_context_samples < self.chunk_size_samples:
                if torch.rand(1).item() > 0.95:
                    index = self.chunk_size_samples
                else:
                    index = torch.randint(min_context_samples,
                                          self.chunk_size_samples + 1,
                                          (1,)).item()
                condition_tensor = condition_tensor[..., :index]

        output = {
            "name": song_name,
            "target": input_tensor,
            "description": description,
            "context": condition_tensor,
        }

        if self.use_beat_conditioning:
            output["beat_seconds"] = beats_time

        if style_tensor is not None:
            output["style"] = style_tensor

        if beats_conditioning is not None:
            output["beat"] = beats_conditioning

        return output


# def generate_sync_data(
#         output_path: Path) -> Dict[str, Dict[str, float | List[int]]]:
#     db = MoisesDB(data_path=str(cfg.DATA_DIR / "moisesdb"),
#                   sample_rate=32_000)
#     if not output_path.exists():
#         raise FileNotFoundError("output path doesn't seem to exist.")
#     sync_path = output_path / "sync.json"
#     if sync_path.exists():
#         raise FileExistsError()
#     data = {}
#     errors = 0
#     for song in tqdm(db, total=len(db)): # type: ignore
#         try:
#             songid = song.id
#             sr = song.sr
#             audio = librosa.to_mono(song.audio)
#             # utils.save_audio(audio)
#             tempo, beats = librosa.beat.beat_track(y=audio,
#                                                    sr=sr,
#                                                    units="samples")
#             beats = beats.tolist()
#             data[songid] = {
#                 "tempo": tempo,
#                 "beats": beats,
#             }
#         except:
#             errors += 1
#     with open(sync_path, "w") as fp:
#         json.dump(data, fp)
#     print(f"Saved sync data of {len(db) - errors} songs, with {errors} errors.")
#     return data


def prepare_data(root_dir: Path, save_mixed_drums: bool, save_mix: bool,
                 extract_features: bool, track_bpm: bool):
    from tqdm import tqdm
    import torchaudio
    from torch import Tensor
    from lag.data.auto_labelling import get_audio_features

    subdirs: List[Path] = sorted([p for p in root_dir.iterdir() if p.is_dir()],
                                 key=lambda x: x.name)
    # assert len(subdirs) == EXPECTED_N_SONGS

    if track_bpm:
        syncdata: Dict[str, List[int]] = {}

    for song in tqdm(subdirs):

        for stemdir in (p for p in song.iterdir() if p.is_dir()):
            if len(list(stemdir.iterdir())) == 0:
                raise FileNotFoundError(f"Song {song} contains no stems. WTF")

        # mix drums
        if (song / "drums").exists():
            drums_sample_rates: Set[int] = set()
            drums_audios: List[Tensor] = []
            for drum_stem in (song / "drums").iterdir():
                audio, sr = torchaudio.load(str(drum_stem))
                drums_sample_rates.add(sr)
                drums_audios.append(audio.permute(1, 0))

            if len(drums_sample_rates) != 1:
                raise ValueError(f"song {song} contains drums stems of "
                                 "different sample rates")
            drums_sr = drums_sample_rates.pop()

            # mix drums
            drums_tensor = torch.nn.utils.rnn.pad_sequence(
                drums_audios, batch_first=True,
                padding_value=0).permute(0, 2, 1).sum(dim=0)

            assert drums_tensor.ndim == 2

            if save_mixed_drums:
                target_dir = song / "drums_mixed"
                target_dir.mkdir(exist_ok=True)
                torchaudio.save(target_dir / "drums.wav", drums_tensor,
                                drums_sr)

        # load mixed song
        stem_subdirs = [p for p in song.iterdir() if p.is_dir()]

        # for each stem
        stem_tracks: List[Tensor] = []
        sample_rates: Set[int] = set()
        for stem_subdir in stem_subdirs:
            if stem_subdir.name == "drums":
                continue

            # for each track of that stem
            for audio_path in stem_subdir.iterdir():
                audio, sr = torchaudio.load(str(audio_path))
                assert audio.ndim == 2
                stem_tracks.append(audio.permute(1, 0))

                sample_rates.add(sr)
                if len(sample_rates) > 1:
                    raise ValueError(f"song {song} contains stems of "
                                     "different sample rates")
        sr = sample_rates.pop()

        # pad shorter tracks
        stems_tensor: Tensor = torch.nn.utils.rnn.pad_sequence(
            stem_tracks,
            batch_first=True,
            padding_value=0.,
        ).permute(0, 2, 1)
        assert stems_tensor.ndim == 3

        # mix song
        mixed_tensor = stems_tensor.sum(dim=0)
        num_frames = mixed_tensor.shape[-1]
        assert mixed_tensor.ndim == 2

        if save_mix:
            mix_out_path = song / "mixed.wav"
            audio_utils.save_audio(mixed_tensor, mix_out_path, sr)

        # neural classification to get metadata
        if extract_features:
            features = get_audio_features(mixed_tensor, sr, cfg.weights_dir())

            features["sample_rate"] = sr
            features["num_frames"] = num_frames
            out_file = song / "features.json"
            with open(out_file, "w") as f:
                json.dump(features, f)

        # bpm tracking
        if track_bpm:
            song_numpy = audio_utils.to_mono(mixed_tensor).squeeze().numpy()

            # audio = librosa.to_mono(song.audio)
            # utils.save_audio(audio)
            try:
                tempo, beats = librosa.beat.beat_track(y=song_numpy,
                                                       sr=sr,
                                                       units="samples")
            except Exception as e:
                print(f"Error tracking beats of song {song.name}")
                raise e

            beats = beats.tolist()

            syncdata[song.name] = beats  # type: ignore

    if track_bpm:
        sync_path: Path = root_dir / "sync.json"
        with open(sync_path, "w") as f:
            json.dump(syncdata, f)  # type: ignore


# if __name__ == "__main__":
#     # from lag import config as cfg
#     from tqdm import tqdm

#     # root_dir = Path("/home/tkol/dev/datasets") / "moisesdb" / "moisesdb_v0.1"
#     # root_dir = Path("/home/tkol/dev/datasets") / "moisesdb" / "musdb"
#     root_dir = cfg.moises_path()

#     '''prepare_data(root_dir,
#                  save_mixed_drums=False,
#                  save_mix=False,
#                  extract_features=True,
#                  track_bpm=False)'''

#     stems = {
#         Stem.DRUMS, Stem.GUITAR, Stem.BASS, Stem.PIANO, Stem.KEYBOARD,
#         Stem.STRINGS
#     }

#     dataset = StemmedDataset(
#         root_dir,
#         stems,
#         target_stem=Stem.DRUMS,
#         single_stem=True,
#         min_context_seconds=5,
#         use_style_conditioning=True,
#         use_beat_conditioning=True,
#         add_click=False,
#         bpm_in_caption=False,
#         sync_chunks=False,
#         train=False,
#         sample_rate=32_000,
#         chunk_size_samples=32_000 * 10,
#         speed_transform_p=1,
#         pitch_transform_p=1,
#         stereo=False,
#         n_samples_per_epoch=None,
#     )

#     dataset_iterator = iter(dataset)
#     for i in tqdm(range(10)):
#         sample = next(dataset_iterator)

#         target: Tensor = sample["target"]  # type: ignore
#         context: Tensor = sample["context"] if sample["context"] is not None else sample["context"]

#         audio_utils.save_audio(target, cfg.AUDIO_DIR / "temp" / f"target{i}.wav")
#         audio_utils.save_audio(context, cfg.AUDIO_DIR / "temp" / f"context{i}.wav")
#         mix = target + torch.nn.functional.pad(
#             context, (0, target.shape[-1] - context.shape[-1]))
#         audio_utils.save_audio(mix, cfg.AUDIO_DIR / "temp" / f"mix{i}.wav")