From 746a79caf419c80a074da7e7763fa57d0ad9b5d9 Mon Sep 17 00:00:00 2001 From: dschnell Date: Fri, 29 May 2026 12:14:14 +0000 Subject: [PATCH] fix: correct segment order in batched k2 forced_align forced_align returned per-segment frame labels and scores in the wrong order for batches larger than one. encode_supervisions sorts the batch by length and returns the permutation, but forced_align never applied it, so each segment received another segment's alignment and score and the resulting durations did not match the token counts. Single-segment batches were correct because there was nothing to sort. The fix splits the scores by the sorted per-segment lengths and then restores the original batch order for both the labels and the scores. It also adds the local k2 import that forced_align was missing; without it the method raised a NameError, so the batched path had never run. The training loss is unaffected because it returns an order-independent scalar, which is why alignment training looked healthy while the exported alignments were scrambled. --- src/stylish_tts/train/losses.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/src/stylish_tts/train/losses.py b/src/stylish_tts/train/losses.py index a9ce029..b3397d9 100644 --- a/src/stylish_tts/train/losses.py +++ b/src/stylish_tts/train/losses.py @@ -583,6 +583,8 @@ def forced_align( input_lengths: Tensor, target_lengths: Tensor, ): + import k2 + supervision_segments, token_ids, indices = self.encode_supervisions( targets, target_lengths, input_lengths ) @@ -602,7 +604,11 @@ def forced_align( best_paths = k2.shortest_path(lattices, use_double_scores=True) frame_scores = best_paths.scores[(best_paths.labels != -1)] - frame_scores = frame_scores.split(input_lengths.tolist()) + # encode_supervisions sorted the batch by length, so best_paths come back + # in that sorted order. Split the scores by the sorted per-segment lengths + # rather than the original input lengths. + sorted_lengths = supervision_segments[:, 2].tolist() + frame_scores = frame_scores.split(sorted_lengths) scores = torch.stack([p.mean() for p in frame_scores]) batch_arc_shape = best_paths.arcs.shape().remove_axis(1) @@ -612,6 +618,13 @@ def forced_align( # k2 makes an extra frame for some reasons for i in range(len(batch_frame_labels)): batch_frame_labels[i][-1] -= 1 + # Restore the original batch order so each segment's labels and score line + # up with the caller's segment. Without this every segment receives another + # segment's alignment. + inverse = torch.empty_like(indices) + inverse[indices] = torch.arange(indices.numel(), device=indices.device) + batch_frame_labels = [batch_frame_labels[j] for j in inverse.tolist()] + scores = scores[inverse] return batch_frame_labels, scores def on_train_epoch_end(self, train):