Improve speed of selecting random chunks2

remfx/datasets.py

@@ -205,14 +205,11 @@ class EffectDataset(Dataset):
         if render_files:
             # Split audio file into chunks, resample, then apply random effects
             self.proc_root.mkdir(parents=True, exist_ok=True)
-            bad_files = set()
             for num_chunk in tqdm(range(self.total_chunks)):
                 chunk = None
                 while chunk is None:
                     random_dataset_choice = random.choice(self.files)
                     random_file_choice = random.choice(random_dataset_choice)
-                    if random_file_choice in bad_files:
-                        continue
                     chunk = select_random_chunk(
                         random_file_choice, self.chunk_size, self.sample_rate
                     )

remfx/utils.py

@@ -130,7 +130,7 @@ def create_sequential_chunks(
     audio_file: str, chunk_size: int, sample_rate: int
 ) -> List[torch.Tensor]:
     """Create sequential chunks of size chunk_size from an audio file.
-    Return sample_index of start of each chunk and original sr
+    Return each chunk
     """
     chunks = []
     audio, sr = torchaudio.load(audio_file)
@@ -150,13 +150,14 @@ def create_sequential_chunks(
 def select_random_chunk(
     audio_file: str, chunk_size: int, sample_rate: int
 ) -> List[torch.Tensor]:
-    """Create sequential chunks of size chunk_size (samples) from an audio file.
-    Return sample_index of start of each chunk and original sr
-    """
+    """Select random chunk of size chunk_size (samples) from an audio file."""
     audio, sr = torchaudio.load(audio_file)
-    max_len = audio.shape[-1] - int(chunk_size * (sample_rate / sr))
+    new_chunk_size = int(chunk_size * (sr / sample_rate))
+    if new_chunk_size >= audio.shape[-1]:
+        return None
+    max_len = audio.shape[-1] - new_chunk_size
     random_start = torch.randint(0, max_len, (1,)).item()
-    chunk = audio[:, random_start : random_start + chunk_size]
+    chunk = audio[:, random_start : random_start + new_chunk_size]
     resampled_chunk = torchaudio.functional.resample(chunk, sr, sample_rate)
     return resampled_chunk