Update DL link

README.md

@@ -17,7 +17,7 @@ Download the preliminary model weights:
 Piano (v1)
 
 ```
-wget https://storage.googleapis.com/aria-checkpoints/amt/piano-medium-double-1.0.safetensors
+wget "https://huggingface.co/datasets/loubb/aria-midi/resolve/main/piano-medium-double-1.0.safetensors?download=true"
 ```
 
 ## Usage

amt/audio.py

@@ -235,11 +235,13 @@ def _get_noise(self, noise_paths: list):
     def apply_reverb(self, wav: torch.Tensor):
         # wav: (bz, L)
         if self.num_impulse == 0:
-            logging.warning("Warning: No files in amt/assets/impulse - skipping apply_reverb")
+            logging.warning(
+                "Warning: No files in amt/assets/impulse - skipping apply_reverb"
+            )
             return wav
         elif self.reverb_factor == 0:
             return wav
-        
+
         batch_size, _ = wav.shape
 
         reverb_strength = (
@@ -265,7 +267,9 @@ def apply_reverb(self, wav: torch.Tensor):
 
     def apply_noise(self, wav: torch.tensor):
         if self.num_noise == 0:
-            logging.warning("Warning: No files in amt/assets/noise - skipping apply_noise")
+            logging.warning(
+                "Warning: No files in amt/assets/noise - skipping apply_noise"
+            )
             return wav
 
         batch_size, _ = wav.shape
@@ -283,7 +287,9 @@ def apply_noise(self, wav: torch.tensor):
 
     def apply_applause(self, wav: torch.tensor):
         if self.num_applause == 0:
-            logging.warning("Warning: No files in amt/assets/applause - skipping apply_applause")
+            logging.warning(
+                "Warning: No files in amt/assets/applause - skipping apply_applause"
+            )
             return wav
 
         batch_size, _ = wav.shape

amt/data.py

@@ -64,6 +64,33 @@ def get_mid_segments(
         start_ms += stride_len_ms
 
 
+def _get_single_wav_segment(
+    audio_path: str,
+    sample_rate: int,
+    orig_sample_rate: int,
+    segment_samples: tuple[int] | None = None,
+):
+    if segment_samples is not None:
+        start_sample, end_sample = segment_samples
+        orig_start_sample = int(start_sample * (orig_sample_rate / sample_rate))
+        orig_end_sample = int(end_sample * (orig_sample_rate / sample_rate))
+
+        wav, _ = torchaudio.load(
+            audio_path,
+            frame_offset=orig_start_sample,
+            num_frames=orig_end_sample - orig_start_sample,
+        )
+        wav = wav.mean(0)
+    else:
+        wav, _ = torchaudio.load(audio_path)
+        wav = wav.mean(0)
+
+    wav = torchaudio.functional.resample(
+        wav, orig_freq=orig_sample_rate, new_freq=sample_rate
+    )
+    return wav
+
+
 def get_wav_segments(
     audio_path: str,
     stride_factor: int | None = None,
@@ -95,48 +122,72 @@ def get_wav_segments(
     else:
         start_sample, end_sample = 0, None
 
-    stream.add_basic_audio_stream(
-        frames_per_chunk=stride_samples,
-        stream_index=0,
-        sample_rate=sample_rate,
+    wav_info = torchaudio.info(audio_path)
+    seg_len_s = (
+        end_time_s - start_time_s
+        if segment is not None
+        else wav_info.num_frames / wav_info.sample_rate
     )
 
-    buffer = torch.tensor([], dtype=torch.float32)
-    total_samples = start_sample
-    for stride_seg in stream.stream():
-        seg_chunk = stride_seg[0].mean(1)
+    if seg_len_s <= chunk_len:
+        # Short audio segment - yield only one segment
+        wav = _get_single_wav_segment(
+            audio_path=audio_path,
+            sample_rate=sample_rate,
+            orig_sample_rate=wav_info.sample_rate,
+            segment_samples=(
+                (start_sample, end_sample) if segment is not None else None
+            ),
+        )
+        if pad_last is True:
+            yield torch.nn.functional.pad(wav, (0, chunk_samples - len(wav)))
+        else:
+            yield wav
 
-        if end_sample and total_samples + seg_chunk.shape[0] > end_sample:
-            samples_to_use = end_sample - total_samples
-            seg_chunk = seg_chunk[:samples_to_use]
+    else:
+        # Yield segments in order
+        stream.add_basic_audio_stream(
+            frames_per_chunk=stride_samples,
+            stream_index=0,
+            sample_rate=sample_rate,
+        )
 
-        total_samples += seg_chunk.shape[0]
+        buffer = torch.tensor([], dtype=torch.float32)
+        total_samples = start_sample
+        for stride_seg in stream.stream():
+            seg_chunk = stride_seg[0].mean(1)
 
-        # Pad seg_chunk if required
-        if seg_chunk.shape[0] < stride_samples:
-            seg_chunk = F.pad(
-                seg_chunk,
-                (0, stride_samples - seg_chunk.shape[0]),
-                mode="constant",
-                value=0.0,
-            )
+            if end_sample and total_samples + seg_chunk.shape[0] > end_sample:
+                samples_to_use = end_sample - total_samples
+                seg_chunk = seg_chunk[:samples_to_use]
 
-        if buffer.shape[0] < chunk_samples:
-            buffer = torch.cat((buffer, seg_chunk), dim=0)
-        else:
-            buffer = torch.cat((buffer[stride_samples:], seg_chunk), dim=0)
+            total_samples += seg_chunk.shape[0]
+
+            # Pad seg_chunk if required
+            if seg_chunk.shape[0] < stride_samples:
+                seg_chunk = F.pad(
+                    seg_chunk,
+                    (0, stride_samples - seg_chunk.shape[0]),
+                    mode="constant",
+                    value=0.0,
+                )
 
-        if buffer.shape[0] == chunk_samples:
-            yield buffer
+            if buffer.shape[0] < chunk_samples:
+                buffer = torch.cat((buffer, seg_chunk), dim=0)
+            else:
+                buffer = torch.cat((buffer[stride_samples:], seg_chunk), dim=0)
 
-        if end_sample and total_samples >= end_sample:
-            break
+            if buffer.shape[0] == chunk_samples:
+                yield buffer
 
-    if pad_last and buffer.shape[0] > stride_samples:
-        yield torch.nn.functional.pad(
-            buffer[stride_samples:],
-            (0, chunk_samples - len(buffer[stride_samples:])),
-        )
+            if end_sample and total_samples >= end_sample:
+                break
+
+        if pad_last and buffer.shape[0] > stride_samples:
+            yield torch.nn.functional.pad(
+                buffer[stride_samples:],
+                (0, chunk_samples - len(buffer[stride_samples:])),
+            )
 
 
 def get_paired_wav_mid_segments(

amt/inference/transcribe.py

@@ -757,16 +757,17 @@ def transcribe_file(
                 logger.info(f"Seen eos_tok in audio chunk {idx}: {file_path}")
                 seq = seq[:-1]
 
+            concat_seq += _shift_onset(
+                seq[init_idx:],
+                idx * CHUNK_LEN_MS,
+            )
+
             if len(next_seq) == 1:
                 logger.info(
                     f"Skipping audio chunk {idx} (silence): {file_path}"
                 )
                 seq = [tokenizer.bos_tok]
             else:
-                concat_seq += _shift_onset(
-                    seq[init_idx:],
-                    idx * CHUNK_LEN_MS,
-                )
                 seq = next_seq
 
         idx += 1
@@ -879,7 +880,7 @@ def remove_failures_from_queue_(_queue: Queue, _pid: int):
         logger.info(
             f"Finished file: {file_path} (segment: {idx if idx is not None else 'full'})"
         )
-        if len(seq) < 500:
+        if len(seq) < 10:
             logger.info(
                 f"Skipping seq - too short (segment {idx if idx is not None else 'full'})"
             )