update to 0.1.2 version

app.py

@@ -4,9 +4,16 @@ HuggingFace Space entry point for OmniVoice demo.
 
 """
 
+import logging
 import os
 from typing import Any, Dict
 
+logging.basicConfig(
+    level=logging.DEBUG,
+    format="%(asctime)s %(name)s %(levelname)s: %(message)s",
+)
+logging.getLogger("omnivoice").setLevel(logging.DEBUG)
+
 import numpy as np
 import spaces
 import torch

omnivoice/cli/demo.py

@@ -243,8 +243,8 @@ def build_demo(
     def _gen_settings():
         with gr.Accordion("Generation Settings (optional)", open=False):
             sp = gr.Slider(
-                0.7,
-                1.3,
+                0.5,
+                1.5,
                 value=1.0,
                 step=0.05,
                 label="Speed",

omnivoice/models/omnivoice.py

@@ -1056,7 +1056,7 @@ class OmniVoice(PreTrainedModel):
         # Build style tokens: <|denoise|> + <|lang_start|>...<|lang_end|>
         #                      + <|instruct_start|>...<|instruct_end|>
         style_text = ""
-        if denoise:
+        if denoise and ref_audio_tokens is not None:
             style_text += "<|denoise|>"
         lang_str = lang if lang else "None"
         instruct_str = instruct if instruct else "None"
@@ -1131,6 +1131,17 @@ class OmniVoice(PreTrainedModel):
 
         B = task.batch_size
 
+        for i in range(B):
+            logger.debug(
+                "Item %d — text: %s | ref_text: %s | instruct: %s | lang: %s | target_tokens: %d",
+                i,
+                task.texts[i],
+                task.ref_texts[i],
+                task.instructs[i],
+                task.langs[i],
+                task.target_lens[i],
+            )
+
         inputs_list = [
             self._prepare_inference_inputs(
                 task.texts[i],
@@ -1173,6 +1184,9 @@ class OmniVoice(PreTrainedModel):
             batch_input_ids[B + i, :, :u_len] = inp["input_ids"][..., -u_len:]
             batch_audio_mask[B + i, :u_len] = inp["audio_mask"][..., -u_len:]
             batch_attention_mask[B + i, :, :u_len, :u_len] = True
+            if max_c_len > u_len:
+                pad_diag = torch.arange(u_len, max_c_len, device=self.device)
+                batch_attention_mask[B + i, :, pad_diag, pad_diag] = True
 
         tokens = torch.full(
             (B, self.config.num_audio_codebook, max(task.target_lens)),
@@ -1491,6 +1505,17 @@ def _combine_text(text, ref_text: Optional[str] = None) -> str:
     chinese_range = r"[\u4e00-\u9fff]"
     pattern = rf"(?<={chinese_range})\s+|\s+(?={chinese_range})"
     full_text = re.sub(pattern, "", full_text)
+
+    # Remove whitespace immediately before special emotion tags (except
+    # [laughter]).  During training these tags have no preceding space, so
+    # the text tokenizer would mis-tokenise them if spaces were present.
+    _EMOTION_TAGS = (
+        r"sigh|confirmation-en|question-en|question-ah|question-oh|"
+        r"question-ei|question-yi|surprise-ah|surprise-oh|surprise-wa|"
+        r"surprise-yo|dissatisfaction-hnn"
+    )
+    full_text = re.sub(rf"\s+(\[({_EMOTION_TAGS})\])", r"\1", full_text)
+
     return full_text
 
 

omnivoice/training/trainer.py

@@ -45,6 +45,14 @@ from omnivoice.training.checkpoint import save_checkpoint as engine_save_checkpo
 logger = logging.getLogger(__name__)
 
 
+def _to_device(batch, device):
+    """Move all tensors in a batch dict to the target device."""
+    return {
+        k: v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v
+        for k, v in batch.items()
+    }
+
+
 class OmniTrainer:
     def __init__(
         self,
@@ -211,6 +219,7 @@ class OmniTrainer:
 
         with torch.no_grad():
             for eval_batch in self.eval_dataloader:
+                eval_batch = _to_device(eval_batch, self.accelerator.device)
                 outputs = self.model(**eval_batch)
                 local_loss_sum += outputs.loss.detach()
                 eval_count += 1
@@ -269,6 +278,8 @@ class OmniTrainer:
                 train_iterator = iter(self.train_dataloader)
                 batch = next(train_iterator)
 
+            batch = _to_device(batch, self.accelerator.device)
+
             with self.accelerator.accumulate(self.model):
                 outputs = self.model(**batch)
                 loss = outputs.loss