Merge branch 'main' into feature/inference-speed-up

examples/test_streaming.py

@@ -23,7 +23,7 @@ clone_model = Qwen3TTSModel.from_pretrained(
 start = log_time(start, "Load Base model")
 
 # for real speedup, use vLLM for LM inference (or SGlang probably)
-# torch.compile doesn't help much for autoregressive generation due to dynamic shapes ( I think but idk )
+# torch.compile doesn't help much for autoregressive generation due to dynamic shapes
 
 ref_audio_path = "kuklina-1.wav"
 ref_text = (
@@ -48,7 +48,7 @@ print("\n--- Standard generation ---")
 start = time.time()
 wavs, sr = clone_model.generate_voice_clone(
     text=test_text,
-    language="Russian",
+    language="English",
     voice_clone_prompt=voice_clone_prompt,
 )
 standard_time = time.time() - start
@@ -64,7 +64,7 @@ chunk_count = 0
 
 for chunk, chunk_sr in clone_model.stream_generate_voice_clone(
     text=test_text,
-    language="Russian",
+    language="English",
     voice_clone_prompt=voice_clone_prompt,
     emit_every_frames=8,
     decode_window_frames=80,
@@ -89,4 +89,4 @@ print(f"Streaming first chunk: {first_chunk_time:.2f}s")
 print(f"Streaming total: {streaming_time:.2f}s")
 print(f"Latency improvement: {standard_time - first_chunk_time:.2f}s faster to first audio")
 
-print(f"\n[{time.time() - total_start:.2f}s] TOTAL")
+print(f"\n[{time.time() - total_start:.2f}s] TOTAL")

finetuning/prepare_data.py

@@ -48,14 +48,18 @@ def main():
         if len(batch_lines) >= BATCH_INFER_NUM:
             enc_res = tokenizer_12hz.encode(batch_audios)
             for code, line in zip(enc_res.audio_codes, batch_lines):
-                line['audio_codes'] = code
+                line['audio_codes'] = code.cpu().tolist()
                 final_lines.append(line)
+            batch_lines.clear()
+            batch_audios.clear()
 
     if len(batch_audios) > 0:
         enc_res = tokenizer_12hz.encode(batch_audios)
         for code, line in zip(enc_res.audio_codes, batch_lines):
             line['audio_codes'] = code.cpu().tolist()
             final_lines.append(line)
+        batch_lines.clear()
+        batch_audios.clear()
 
     final_lines = [json.dumps(line, ensure_ascii=False) for line in final_lines]
 

qwen_tts/core/models/modeling_qwen3_tts.py

@@ -519,7 +519,7 @@ class Qwen3TTSPreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = ["Qwen3TTSDecoderLayer"]
     _skip_keys_device_placement = "past_key_values"
-    _supports_flash_attn_2 = True
+    _supports_flash_attn = True
     _supports_sdpa = True
     _supports_cache_class = True
     _supports_static_cache = False
@@ -550,8 +550,7 @@ class Qwen3TTSTalkerTextPreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = []
     _skip_keys_device_placement = ["past_key_values"]
-    _supports_flash_attn_3 = True
-    _supports_flash_attn_2 = True
+    _supports_flash_attn = True
     _supports_sdpa = True
     _supports_flex_attn = True
     _supports_cache_class = True
@@ -2119,6 +2118,11 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
         weights_only=True,
         **kwargs,
     ):
+        # Hotfix to enable passing the correct attn implementation which is stored in the config but not in kwargs
+        requested_attn_implementation = kwargs.pop("attn_implementation", None)
+        if requested_attn_implementation is None and config and config._attn_implementation:
+            requested_attn_implementation = config._attn_implementation
+
         model = super().from_pretrained(
             pretrained_model_name_or_path,
             *model_args,
@@ -2131,6 +2135,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
             revision=revision,
             use_safetensors=use_safetensors,
             weights_only=weights_only,
+            attn_implementation=requested_attn_implementation,
             **kwargs,
         )
         if not local_files_only and not os.path.isdir(pretrained_model_name_or_path):

qwen_tts/inference/qwen3_tts_model.py

@@ -809,7 +809,7 @@ class Qwen3TTSModel:
         text: Union[str, List[str]],
         instruct: Union[str, List[str]],
         language: Union[str, List[str]] = None,
-        non_streaming_mode: bool = False,
+        non_streaming_mode: bool = True,
         **kwargs,
     ) -> Tuple[List[np.ndarray], int]:
         """
@@ -905,7 +905,7 @@ class Qwen3TTSModel:
         speaker: Union[str, List[str]],
         language: Union[str, List[str]] = None,
         instruct: Optional[Union[str, List[str]]] = None,
-        non_streaming_mode: bool = False,
+        non_streaming_mode: bool = True,
         **kwargs,
     ) -> Tuple[List[np.ndarray], int]:
         """