[model] support qwen3_tts

[Jintao-Huang]

No description provided.

[gemini-code-assist]

Code Review

This pull request adds support for Qwen3-TTS SFT training with a dual-channel architecture, introducing a custom loss function, model loader, patched forward pass, and template for data collation. Key feedback includes a critical fix to retrieve the correct hidden states from the last layer for the entire batch instead of indexing the embedding layer of a single batch element. Additionally, recommendations were made to pad reference mel spectrograms of varying durations to prevent collation crashes, slice codec IDs consistently to avoid shape mismatches, optimize audio loading/resampling with librosa, and explicitly cast boolean masks to prevent implicit type promotion issues.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Concatenating tts_ref_mel tensors directly using torch.cat will fail with a RuntimeError if the reference audios in the batch have different durations (and thus different mel spectrogram lengths).

To make the data collator robust, pad the mel spectrograms along the time dimension (dimension 1) to the maximum length in the current batch before concatenation.

        ref_mels = [data['tts_ref_mel'] for data in batch]
        max_mel_len = max(mel.shape[1] for mel in ref_mels)
        padded_ref_mels = []
        for mel in ref_mels:
            padding_len = max_mel_len - mel.shape[1]
            if padding_len > 0:
                padded_mel = F.pad(mel, (0, 0, 0, padding_len))
            else:
                padded_mel = mel
            padded_ref_mels.append(padded_mel)
        ref_mels = torch.cat(padded_ref_mels, dim=0)

[gemini-code-assist]

Multiplying float embedding tensors directly by boolean mask tensors (text_embedding_mask and codec_embedding_mask) relies on implicit type promotion, which can be slow or cause warnings/errors on certain PyTorch versions or hardware accelerators.

It is safer and more robust to explicitly cast the boolean masks to the embedding's data type before multiplication.

Suggested change

	input_text_embedding = self.talker.model.text_embedding(input_text_ids) * text_embedding_mask
	input_codec_embedding = self.talker.model.codec_embedding(input_codec_ids) * codec_embedding_mask
	input_text_embedding = self.talker.model.text_embedding(input_text_ids) * text_embedding_mask.to(input_text_embedding.dtype)
	input_codec_embedding = self.talker.model.codec_embedding(input_codec_ids) * codec_embedding_mask.to(input_codec_embedding.dtype)

[gemini-code-assist]

Currently, talker_hidden_states is indexed using codec_mask[:, :-1], while talker_codec_ids is indexed using codec_mask.

If codec_mask ever contains a True value in its last column (e.g., due to different sequence lengths or template changes), this will result in a shape mismatch between talker_codec_ids and talker_hidden_states, causing the training to crash. To ensure consistency and robustness, slice codec_ids and use codec_mask[:, :-1] for both.

Suggested change

	talker_hidden_states = hidden_states[codec_mask[:, :-1]]
	talker_codec_ids = codec_ids[codec_mask]
	talker_hidden_states = hidden_states[codec_mask[:, :-1]]
	talker_codec_ids = codec_ids[:, :-1][codec_mask[:, :-1]]

[gemini-code-assist]

Loading the audio with sr=None and then resampling it with librosa.resample is inefficient because it loads the entire audio at its native sampling rate into memory first before performing resampling.

You can load and resample the audio to 24000 Hz directly in a single highly optimized step by passing sr=24000 to librosa.load. This also guarantees a 1D mono array, making the subsequent ndim check redundant.

Suggested change

	audio, sr = librosa.load(ref_audio_path, sr=None, mono=True)
	if audio.ndim > 1:
	audio = np.mean(audio, axis=-1)
	if sr != 24000:
	audio = librosa.resample(audio, orig_sr=sr, target_sr=24000)
	audio, sr = librosa.load(ref_audio_path, sr=24000, mono=True)

[gemini-code-assist]

Multiplying the float tensor codec_i_embedding by the boolean tensor codec_mask relies on implicit type promotion. It is safer and more robust to explicitly cast the boolean mask to the embedding's data type before multiplication.

Suggested change

	codec_i_embedding = codec_i_embedding * codec_mask.unsqueeze(-1)
	codec_i_embedding = codec_i_embedding * codec_mask.unsqueeze(-1).to(codec_i_embedding.dtype)

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for Qwen3-TTS models, including dataset preprocessing, custom dual-channel template collation, patched forward training logic, and specialized saving callbacks. Feedback on these changes highlights several critical improvements: correcting the indexing of hidden states and codec IDs in the patched forward pass to prevent shape mismatches, robustly handling list types for audio paths in the preprocessor, ignoring large weight and index files during model saving to avoid loading errors, casting boolean masks to the embedding's dtype for accelerator compatibility, and sanitizing extracted zip file paths to mitigate directory traversal vulnerabilities.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

There are two critical issues in this block:

1. outputs.hidden_states[0] refers to the embedding layer output (the first element of the hidden_states tuple). To compute the sub-talker loss from the final layer's contextual representations, you should use the last layer's hidden states, which is outputs.hidden_states[-1].
2. outputs.hidden_states[0][-1] indexes the batch dimension, selecting only the last sample in the batch. If batch_size > 1, this will cause a shape mismatch and crash with an IndexError when indexed by codec_mask[:, :-1].
3. For perfect alignment and safety, codec_ids should also be sliced with [:, :-1] when indexed by codec_mask[:, :-1].

Suggested change

	# Compute sub_talker_loss from hidden states at codec positions
	hidden_states = outputs.hidden_states[0][-1]
	talker_hidden_states = hidden_states[codec_mask[:, :-1]]
	talker_codec_ids = codec_ids[codec_mask]
	_, sub_talker_loss = self.talker.forward_sub_talker_finetune(talker_codec_ids, talker_hidden_states)
	# Compute sub_talker_loss from hidden states at codec positions
	hidden_states = outputs.hidden_states[-1]
	talker_hidden_states = hidden_states[codec_mask[:, :-1]]
	talker_codec_ids = codec_ids[:, :-1][codec_mask[:, :-1]]
	_, sub_talker_loss = self.talker.forward_sub_talker_finetune(talker_codec_ids, talker_hidden_states)

[FelysNeko]

I think this will break LoRA compatibility, here's minimum fix:

        base_model = unwrap_model(model)
        from peft import PeftModel as _PeftModel
        if isinstance(base_model, _PeftModel):
            base_model = base_model.merge_and_unload()

Caution: only do this if you just want to get the last checkpoint. Otherwise, according to LLM, Swift.merge and then Swift.unmerge might be the correct way to handle this?