diff --git a/cookbook/qwen3_tts/.gitignore b/cookbook/qwen3_tts/.gitignore new file mode 100644 index 00000000..76782217 --- /dev/null +++ b/cookbook/qwen3_tts/.gitignore @@ -0,0 +1,3 @@ +outputs/ +__pycache__/ +*.py[cod] diff --git a/cookbook/qwen3_tts/README.md b/cookbook/qwen3_tts/README.md new file mode 100644 index 00000000..404cff89 --- /dev/null +++ b/cookbook/qwen3_tts/README.md @@ -0,0 +1,15 @@ +# Qwen3-TTS cookbook — layer 1 + +This contribution is part 1 of a four-part series. Scope is cumulative: +1. Audio features, padded input processing, two-stage forward and codebook loss. +2. Twinkle model bridge and pretrained LoRA update/save/reload. +3. Accumulation windows and update-boundary checkpoint restoration. +4. Offline manifest preparation and the user-facing training CLI. + +Only the features through part 1 are present here. No generation, quality evaluation or changes to Twinkle core are included. + +Install Twinkle from the repository root (`pip install -e .`) and optional dependencies with `pip install -r cookbook/qwen3_tts/requirements.txt`. The real-model tests require an existing local Qwen3-TTS 12Hz 1.7B Base and a single Ascend NPU. Development baseline: Twinkle 5337cc3, Qwen 022e286, transformers 4.57.3, torch 2.10 / torch_npu 2.10, bf16. No model download is needed. + +From this cookbook directory, run `python tests/batch1/test_contracts.py`. CPU development without Twinkle/Qwen can use `--isolated-twinkle`; this substitutes infrastructure and skips two official package checks. It is not native or NPU validation. + +The final series tree is byte-identical to the integrated training baseline 0466bd2 that passed 40 tests on one Ascend NPU. Individual intermediate layers receive CPU checks; they have not independently been rerun on NPU. No convergence or voice-similarity result is claimed. The model bridge depends on Twinkle internal optimizer-group and wrapping APIs; these remain confined to this cookbook. diff --git a/cookbook/qwen3_tts/requirements.txt b/cookbook/qwen3_tts/requirements.txt new file mode 100644 index 00000000..50874f35 --- /dev/null +++ b/cookbook/qwen3_tts/requirements.txt @@ -0,0 +1,8 @@ +accelerate==1.12.0 +numpy>=2.0,<2.3 +peft==0.17.1 +qwen-tts @ git+https://github.com/QwenLM/Qwen3-TTS.git@022e286b98fbec7e1e916cb940cdf532cd9f488e +safetensors +scipy +soundfile +transformers==4.57.3 diff --git a/cookbook/qwen3_tts/scripts/audio_features.py b/cookbook/qwen3_tts/scripts/audio_features.py new file mode 100644 index 00000000..cf955bd9 --- /dev/null +++ b/cookbook/qwen3_tts/scripts/audio_features.py @@ -0,0 +1,42 @@ +"""Official 24 kHz mel extraction and 16-codebook audio encoding.""" + +import numpy as np + +MEL_DIM = 128 +MEL_HOP = 256 +MEL_WIN = 1024 +MEL_FMAX = 12000.0 +TARGET_SR = 24000 + +def extract_mel(wav_24k: np.ndarray) -> np.ndarray: + """使用官方 mel 算法，返回单条 [F,128]，不使用替代 fbank。""" + import torch + from qwen_tts.core.models.modeling_qwen3_tts import mel_spectrogram + + wav = np.asarray(wav_24k, dtype=np.float32) + if wav.ndim != 1: + raise ValueError('Expected mono waveform [samples]') + with torch.no_grad(): + mel = mel_spectrogram( + torch.from_numpy(wav).unsqueeze(0), n_fft=MEL_WIN, + num_mels=MEL_DIM, sampling_rate=TARGET_SR, hop_size=MEL_HOP, + win_size=MEL_WIN, fmin=0, fmax=MEL_FMAX, + ).transpose(1, 2) + return mel[0].cpu().numpy() + +def encode_audio_codes(wav_24k: np.ndarray, tokenizer) -> np.ndarray: + """单条输入必须返回一个 [T,16]；兼容列表、二维张量和单批张量。""" + result = tokenizer.encode(wav_24k, sr=TARGET_SR) # wav_24k 已升采样到 24kHz + codes = result.audio_codes + if isinstance(codes, (list, tuple)): + if len(codes) != 1: + raise ValueError('Single waveform must return exactly one audio_codes entry') + codes = codes[0] + if hasattr(codes, 'cpu'): + codes = codes.cpu().numpy() + codes = np.asarray(codes) + if codes.ndim == 3 and codes.shape[0] == 1: + codes = codes[0] + if codes.ndim != 2 or codes.shape[0] == 0 or codes.shape[1] != 16: + raise ValueError(f'Expected nonempty audio_codes [T,16], got {codes.shape}') + return codes.astype(np.int64) diff --git a/cookbook/qwen3_tts/tests/batch1/isolated_support.py b/cookbook/qwen3_tts/tests/batch1/isolated_support.py new file mode 100644 index 00000000..0f30fc48 --- /dev/null +++ b/cookbook/qwen3_tts/tests/batch1/isolated_support.py @@ -0,0 +1,44 @@ +"""Explicit CPU test doubles for Twinkle infrastructure only. + +Production tensor operations and transformers loss remain real. This is NOT +Twinkle/NPU integration evidence. Only loaded with --isolated-twinkle. +""" +import sys +import types +from pathlib import Path +import torch + + +def install(root): + def module(name, **attrs): + obj = types.ModuleType(name) + obj.__dict__.update(attrs) + sys.modules[name] = obj + return obj + + for name, path in [('tts', 'tts'), ('tts.qwen3', 'tts/qwen3')]: + module(name, __path__=[str(root / path)]) + + class Processor: + padding_map = {} + VLM_CONCAT_FIELDS = set() + def __init__(self, **kwargs): + self.padding_free = kwargs.get('padding_free', False) + + def pad_stack(values, pad_value, concat): + assert concat is False + return torch.nn.utils.rnn.pad_sequence(values, batch_first=True, padding_value=pad_value) + + module('twinkle', remote_function=lambda **kw: lambda fn: fn, + Platform=types.SimpleNamespace(get_local_device=lambda: 'cpu')) + module('twinkle.data_format', InputFeature=dict, Trajectory=dict, LossOutput=dict) + module('twinkle.data_format.output', LossOutput=dict) + module('twinkle.processor') + module('twinkle.processor.base', InputProcessor=Processor) + module('twinkle.utils') + module('twinkle.utils.torch_utils', pad_and_stack_tensors=pad_stack) + module('twinkle.loss') + module('twinkle.loss.base', Loss=object) + module('twinkle.infra', collect_tensor_dict=lambda x: x) + module('twinkle.model') + module('twinkle.model.transformers', TransformersModel=object) diff --git a/cookbook/qwen3_tts/tests/batch1/test_contracts.py b/cookbook/qwen3_tts/tests/batch1/test_contracts.py new file mode 100644 index 00000000..107dc25c --- /dev/null +++ b/cookbook/qwen3_tts/tests/batch1/test_contracts.py @@ -0,0 +1,234 @@ +"""Run: python tests/batch1/test_contracts.py [--isolated-twinkle]. + +Default requires installed Twinkle + Qwen3-TTS + transformers. Isolated mode +substitutes only Twinkle infrastructure, and skips official mel parity. +All tensors, autograd and ForCausalLMLoss use real torch/transformers. +""" +import ast +import importlib.util +import inspect +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +import numpy as np +import torch +from torch import nn +from transformers.loss.loss_utils import ForCausalLMLoss + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT)) +ISOLATED = '--isolated-twinkle' in sys.argv +if ISOLATED: + sys.argv.remove('--isolated-twinkle') + from isolated_support import install + install(ROOT) + +from tts.qwen3.data import Qwen3TTSDataConfig, build_qwen3_sample, collate_qwen3_batch +from tts.qwen3.input_processor import Qwen3TTSInputProcessor +from tts.qwen3.forward import qwen3_tts_forward_step +from tts.qwen3.loss import Qwen3TTSLoss + +spec = importlib.util.spec_from_file_location('prepare', ROOT / 'scripts/audio_features.py') +prepare = importlib.util.module_from_spec(spec) +spec.loader.exec_module(prepare) + + +def samples(n=2): + cfg = Qwen3TTSDataConfig( + tts_pad_token_id=0, tts_bos_token_id=1, tts_eos_token_id=2, + codec_pad_id=3, codec_bos_id=4, codec_eos_token_id=5, + codec_nothink_id=6, codec_think_bos_id=7, codec_think_eos_id=8) + rows = [build_qwen3_sample(np.arange(10, 15-i), + np.arange((3-i)*16).reshape(3-i, 16) % 20 + 9, cfg) for i in range(n)] + mels = [np.full((10, 128), i+1, dtype=np.float32) for i in range(n)] + return cfg, rows, mels + + +class Talker(nn.Module): + """Small differentiable double of Qwen's two-stage API, with real HF loss.""" + def __init__(self): + super().__init__() + self.config = types.SimpleNamespace(num_code_groups=16) + self.model = nn.Module() + self.model.text_embedding = nn.Embedding(64, 8) + self.text_projection = nn.Identity() + self.model.codec_embedding = nn.Embedding(64, 8) + self.code_predictor = nn.Module() + self.code_predictor.embeddings = nn.ModuleList([nn.Embedding(64, 8) for _ in range(15)]) + self.code_predictor.get_input_embeddings = lambda: self.code_predictor.embeddings + self.head = nn.Linear(8, 64) + self.heads = nn.ModuleList([nn.Linear(8, 64) for _ in range(15)]) + + def forward(self, inputs_embeds, attention_mask, labels, **kwargs): + self.labels = labels + self.hidden = inputs_embeds + self.logits = self.head(inputs_embeds) + self.logits.retain_grad() + return types.SimpleNamespace( + loss=ForCausalLMLoss(self.logits, labels, 64), + hidden_states=((inputs_embeds,), None)) + + def forward_sub_talker_finetune(self, codes, hidden): + self.selected_hidden = hidden + self.selected_codes = codes + logits = torch.stack([head(hidden) for head in self.heads], dim=1) + logits.retain_grad() + self.sub_logits = logits + return logits, ForCausalLMLoss(logits, codes[:, 1:], 64) + + +class Model(nn.Module): + dtype = torch.float32 + device = torch.device('cpu') + def __init__(self): + super().__init__() + self.talker = Talker() + self.speaker_encoder = lambda mels: mels.mean(1)[:, :8] + + +class Contracts(unittest.TestCase): + def test_current_frame_residual_codes_do_not_leak_into_condition(self): + cfg, rows, mels = samples(1) + raw = collate_qwen3_batch(rows, cfg, mels) + model = Model() + processor = Qwen3TTSInputProcessor() + batch = processor(raw, model=model) + qwen3_tts_forward_step(model, batch) + before = model.talker.selected_hidden.detach().clone() + changed = {k: v.clone() for k,v in batch.items()} + first = int(torch.where(changed['codec_mask'][0])[0][0]) + changed['codec_ids'][0, first, 1:] = (changed['codec_ids'][0, first, 1:] + 1) % 64 + qwen3_tts_forward_step(model, changed) + torch.testing.assert_close(model.talker.selected_hidden[0], before[0], rtol=0, atol=0) + self.assertFalse(torch.equal(model.talker.selected_hidden[1], before[1])) + + + + def test_codec_single_list_tensor_and_batch_tensor(self): + expected = np.arange(48).reshape(3, 16) + for value in ([torch.tensor(expected)], [expected], torch.tensor(expected), + torch.tensor(expected)[None], expected): + tokenizer = types.SimpleNamespace(encode=lambda *a, **kw: types.SimpleNamespace(audio_codes=value)) + np.testing.assert_array_equal(prepare.encode_audio_codes(np.ones(4096), tokenizer), expected) + + def test_codec_rejects_empty_multiple_and_wrong_shapes(self): + for value in ([], [np.zeros((2,16))]*2, np.zeros((2,3,16)), np.zeros((2,15)), np.zeros((0,16))): + with self.subTest(shape=str(type(value))), self.assertRaises(ValueError): + tokenizer = types.SimpleNamespace(encode=lambda *a, **kw: types.SimpleNamespace(audio_codes=value)) + prepare.encode_audio_codes(np.ones(4096), tokenizer) + + def test_numpy_mel_batch_identity_and_singleton(self): + for n in (1,2): + cfg, rows, mels = samples(n) + for refs in (mels, [m[None] for m in mels]): + batch = collate_qwen3_batch(rows, cfg, refs) + self.assertEqual(batch['ref_mels'].shape, (n,10,128)) + for i in range(n): + np.testing.assert_array_equal(batch['ref_mels'][i], mels[i]) + + def test_mel_length_and_count_fail_explicitly(self): + cfg, rows, mels = samples() + for refs in ([mels[0], mels[1][:7]], mels[:1], [], [mels[0], np.zeros((10,80))]): + with self.assertRaises(ValueError): + collate_qwen3_batch(rows, cfg, refs) + feats = [dict(row, ref_mels=mel) for row, mel in zip(rows, refs)] + if len(feats) == 2: + with self.assertRaises(ValueError): + Qwen3TTSInputProcessor().collate_fn(feats) + + def test_processor_entire_call_samples_and_collated(self): + for n in (1,2): + cfg, rows, mels = samples(n) + expected = collate_qwen3_batch(rows, cfg, mels) + inputs = [[dict(row, ref_mels=mel) for row,mel in zip(rows,mels)], expected, + {k:torch.as_tensor(v) for k,v in expected.items()}] + for value in inputs: + actual = Qwen3TTSInputProcessor()(value, model=Model()) + self.assertEqual(set(actual), set(expected)) + for key in expected: + np.testing.assert_array_equal(actual[key].numpy(), expected[key]) + self.assertEqual(actual['codec_mask'].dtype, torch.bool) + + def test_processor_rejects_missing_mel_bad_rank_and_packing(self): + cfg, rows, mels = samples() + batch = collate_qwen3_batch(rows, cfg, mels) + missing = dict(batch); del missing['ref_mels'] + for value in (missing, dict(batch, ref_mels=np.zeros((20,128)))): + with self.assertRaises(ValueError): + Qwen3TTSInputProcessor()(value, model=Model()) + with self.assertRaises(ValueError): + Qwen3TTSInputProcessor(padding_free=True)(batch, model=Model()) + + + def test_talker_time_shift_independent_of_processor(self): + cfg, rows, mels = samples() + batch = {key:torch.as_tensor(value) for key,value in collate_qwen3_batch(rows,cfg).items()} + batch['ref_mels'] = torch.as_tensor(np.stack(mels)) + model = Model() + outputs = qwen3_tts_forward_step(model,batch) + # Independent oracle: position t predicts original label[t+1], including + # the first frame and EOS. This fails the old external+internal shift. + expected = nn.functional.cross_entropy(model.talker.logits[:,:-1].reshape(-1,64), + batch['codec_0_labels'][:,1:].reshape(-1)) + torch.testing.assert_close(outputs['talker_loss'], expected) + self.assertEqual(model.talker.logits.shape[:2],batch['codec_0_labels'].shape) + + def test_official_mel_call_contract(self): + fake = types.ModuleType('qwen_tts.core.models.modeling_qwen3_tts') + def mel(wave, **kw): + self.assertEqual(wave.shape, (1,4096)) + self.assertEqual(kw, dict(n_fft=1024,num_mels=128,sampling_rate=24000, + hop_size=256,win_size=1024,fmin=0,fmax=12000.0)) + return torch.ones(1,128,16) + fake.mel_spectrogram = mel + with patch.dict(sys.modules, {'qwen_tts.core.models.modeling_qwen3_tts':fake}): + self.assertEqual(prepare.extract_mel(np.ones(4096)).shape, (16,128)) + + @unittest.skipIf(ISOLATED, 'Official qwen_tts not installed in isolated CPU run') + def test_official_sub_talker_methods_and_last_head_gradient(self): + from qwen_tts.core.models.modeling_qwen3_tts import ( + Qwen3TTSTalkerForConditionalGeneration, + Qwen3TTSTalkerCodePredictorModelForConditionalGeneration, + ) + # Run the installed official assembly and head-selection methods. Only + # the expensive transformer decoder is replaced by a tiny causal sum. + embeddings = nn.ModuleList([nn.Embedding(64,8) for _ in range(15)]) + heads = nn.ModuleList([nn.Linear(8,64) for _ in range(15)]) + cp = types.SimpleNamespace( + config=types.SimpleNamespace(num_code_groups=16, vocab_size=64, + output_attentions=False, output_hidden_states=False), + small_to_mtp_projection=nn.Identity(), lm_head=heads, + get_input_embeddings=lambda:embeddings, loss_function=ForCausalLMLoss, + model=lambda inputs_embeds, **kw: types.SimpleNamespace(last_hidden_state=inputs_embeds.cumsum(1))) + cp.forward_finetune = types.MethodType( + Qwen3TTSTalkerCodePredictorModelForConditionalGeneration.forward_finetune, cp) + talker = types.SimpleNamespace( + config=types.SimpleNamespace(num_code_groups=16,hidden_size=8), + code_predictor=cp, get_input_embeddings=lambda:embeddings[0]) + codes = torch.arange(32).reshape(2,16) + hidden = torch.randn(2,8,requires_grad=True) + logits, _ = Qwen3TTSTalkerForConditionalGeneration.forward_sub_talker_finetune(talker,codes,hidden) + logits.retain_grad() + outputs = dict(talker_loss=hidden.sum()*0,sub_talker_logits=logits,sub_talker_labels=codes[:,1:]) + Qwen3TTSLoss()({},outputs)['loss'].backward() + expected = logits.detach().softmax(-1) + expected.scatter_add_(-1,codes[:,1:,None],-torch.ones(2,15,1)) + torch.testing.assert_close(logits.grad, expected*(0.3/30)) + for head in heads: + self.assertGreater(head.weight.grad.abs().sum().item(),0) + + @unittest.skipIf(ISOLATED, 'Official qwen_tts not installed in isolated CPU run') + def test_official_mel_numerical_parity(self): + from qwen_tts.core.models.modeling_qwen3_tts import mel_spectrogram + wave = np.sin(np.arange(8192, dtype=np.float32)*0.03) + expected = mel_spectrogram(torch.from_numpy(wave)[None], n_fft=1024, + num_mels=128,sampling_rate=24000,hop_size=256,win_size=1024,fmin=0,fmax=12000).transpose(1,2)[0] + np.testing.assert_allclose(prepare.extract_mel(wave), expected.numpy(), rtol=1e-6, atol=1e-6) + + +if __name__ == '__main__': + print('CPU contracts / Twinkle infrastructure: ' + ('ISOLATED DOUBLES' if ISOLATED else 'INSTALLED LIBRARY')) + unittest.main(verbosity=2) diff --git a/cookbook/qwen3_tts/tts/__init__.py b/cookbook/qwen3_tts/tts/__init__.py new file mode 100644 index 00000000..28cc1f3b --- /dev/null +++ b/cookbook/qwen3_tts/tts/__init__.py @@ -0,0 +1 @@ +"""Local Qwen3-TTS cookbook support; no global Twinkle registration.""" diff --git a/cookbook/qwen3_tts/tts/qwen3/__init__.py b/cookbook/qwen3_tts/tts/qwen3/__init__.py new file mode 100644 index 00000000..c2e53cdf --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/__init__.py @@ -0,0 +1 @@ +"""Qwen3-TTS single-card training components.""" diff --git a/cookbook/qwen3_tts/tts/qwen3/data.py b/cookbook/qwen3_tts/tts/qwen3/data.py new file mode 100644 index 00000000..c5023483 --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/data.py @@ -0,0 +1,190 @@ +"""Qwen3-TTS 双通道样本构建与 batch 组织。特殊 token 应从目标模型 config 读取。""" + +from dataclasses import dataclass, field +from typing import Dict, List, Optional + +import numpy as np + +NUM_CODE_GROUPS = 16 # Qwen3-TTS 12Hz codec. + + +@dataclass +class Qwen3TTSDataConfig: + """Token layout for the 12Hz model. + + The loader supplies actual IDs from the local model configuration; defaults + document the 1.7B Base layout and should not be assumed for other variants. + """ + + # 文本通道（Qwen3 文本词表段，151xxx，config 顶层 tts_*_token_id） + tts_pad_token_id: int = 151671 + tts_bos_token_id: int = 151672 + tts_eos_token_id: int = 151673 + # codec 通道（codec 词表段，2xxx，config.talker_config.codec_*_id） + codec_pad_id: int = 2148 + codec_bos_id: int = 2149 + codec_eos_token_id: int = 2150 + codec_nothink_id: int = 2155 + codec_think_bos_id: int = 2156 + codec_think_eos_id: int = 2157 + num_code_groups: int = NUM_CODE_GROUPS + + +def build_assistant_text(text: str) -> str: + """官方 `_build_assistant_text`：把目标文本包成 assistant 轮（ChatML）。 + + 真实 text_ids 由 Qwen3 文本 tokenizer 对该串编码、再 `[:, :-5]` 得到。 + 本模块只负责布局，text_ids 由上游（Preprocessor/Template）传入。 + """ + return f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n" + + +def build_qwen3_sample( + text_ids: np.ndarray, + audio_codes: np.ndarray, + cfg: Qwen3TTSDataConfig, +) -> Dict[str, np.ndarray]: + """把单条 (text_ids, audio_codes) 组成双通道 InputFeature（未 pad）。 + + Args: + text_ids: 1D int array [L]，L>=4（官方 `[:, :-5]` 后的文本 token，含 assistant 标签）。 + audio_codes: 2D int array [C, num_code_groups]，目标音频多码本。 + cfg: 特殊 token id。 + + Returns: + dict，每项是长度 T=8+L+C 的 numpy 数组（input_ids/codec_ids 是 2D）。 + """ + text_ids = np.asarray(text_ids).reshape(-1) + audio_codes = np.asarray(audio_codes) + assert audio_codes.ndim == 2 and audio_codes.shape[1] == cfg.num_code_groups, ( + f"audio_codes shape 应为 [C, {cfg.num_code_groups}]，实得 {audio_codes.shape}" + ) + L = int(text_ids.shape[0]) + C = int(audio_codes.shape[0]) + assert L >= 4, f"text_ids_len L={L} 过短（需 >=4 以满足 [3:] 非空）" + + T = 8 + L + C + audio_codec_0 = audio_codes[:, 0] # [C] + ncb = cfg.num_code_groups + + input_ids = np.zeros((T, 2), dtype=np.int64) + codec_ids = np.zeros((T, ncb), dtype=np.int64) + codec_0_labels = np.full((T,), -100, dtype=np.int64) + text_embedding_mask = np.zeros((T,), dtype=bool) + codec_embedding_mask = np.zeros((T,), dtype=bool) + codec_mask = np.zeros((T,), dtype=bool) + attention_mask = np.zeros((T,), dtype=np.int64) + + # ---- 文本通道 input_ids[:,0] ---- + input_ids[0:3, 0] = text_ids[0:3] + input_ids[3:7, 0] = cfg.tts_pad_token_id + input_ids[7, 0] = cfg.tts_bos_token_id + input_ids[8 : 8 + L - 3, 0] = text_ids[3:] # L-3 个 + input_ids[8 + L - 3, 0] = cfg.tts_eos_token_id + input_ids[8 + L - 2 : 8 + L + C, 0] = cfg.tts_pad_token_id # [8+L-2 : T] + + # ---- codec 通道 input_ids[:,1] ---- + input_ids[3:8, 1] = [ + cfg.codec_nothink_id, + cfg.codec_think_bos_id, + cfg.codec_think_eos_id, + 0, # speaker emb 槽（id=0 占位，训练步注入 speaker_emb） + cfg.codec_pad_id, + ] + input_ids[8 : 8 + L - 3, 1] = cfg.codec_pad_id + input_ids[8 + L - 3, 1] = cfg.codec_pad_id + input_ids[8 + L - 2, 1] = cfg.codec_bos_id + input_ids[8 + L - 1 : 8 + L - 1 + C, 1] = audio_codec_0 # [8+L-1 : 8+L-1+C] + input_ids[8 + L - 1 + C, 1] = cfg.codec_eos_token_id # pos T-1 + + # ---- codec_0_labels（码本0 label，-100 mask） ---- + codec_0_labels[8 + L - 1 : 8 + L - 1 + C] = audio_codec_0 + codec_0_labels[8 + L - 1 + C] = cfg.codec_eos_token_id # pos T-1 + + # ---- codec_ids（全码本，喂 Code predictor） ---- + codec_ids[8 + L - 1 : 8 + L - 1 + C, :] = audio_codes + + # ---- masks ---- + text_embedding_mask[: 8 + L + C] = True # [:T] + codec_embedding_mask[3 : 8 + L + C] = True # [3:T] + codec_embedding_mask[6] = False # speaker emb 槽 + codec_mask[8 + L - 1 : 8 + L - 1 + C] = True # [8+L-1 : 8+L+C-1) = C 个目标帧 + attention_mask[: 8 + L + C] = True # [:T] + + return { + 'input_ids': input_ids, + 'codec_ids': codec_ids, + 'codec_0_labels': codec_0_labels, + 'text_embedding_mask': text_embedding_mask, + 'codec_embedding_mask': codec_embedding_mask, + 'codec_mask': codec_mask, + 'attention_mask': attention_mask, + 'length': np.int64(T), + } + + +def collate_qwen3_batch( + samples: List[Dict[str, np.ndarray]], + cfg: Qwen3TTSDataConfig, + ref_mels: Optional[List[np.ndarray]] = None, + max_length: Optional[int] = None, +) -> Dict[str, np.ndarray]: + """把多个 build_qwen3_sample 输出 pad 成一个 batch（对齐官方 collate_fn 的 padding）。 + + padding 规则（与官方一致）： + input_ids/codec_ids -> 0；masks -> False；attention_mask -> 0；codec_0_labels -> -100。 + ref_mels 单条为 [F,128] 或 [1,F,128]，同长度 stack 为 [B,F,128]。 + 不同参考长度拒绝合批；使用相同参考音频或按参考长度分批，避免无掩码补零影响池化。 + """ + b = len(samples) + if max_length is None: + max_length = int(max(int(s['length']) for s in samples)) + t = max_length + ncb = cfg.num_code_groups + + input_ids = np.zeros((b, t, 2), dtype=np.int64) + codec_ids = np.zeros((b, t, ncb), dtype=np.int64) + codec_0_labels = np.full((b, t), -100, dtype=np.int64) + text_embedding_mask = np.zeros((b, t), dtype=bool) + codec_embedding_mask = np.zeros((b, t), dtype=bool) + codec_mask = np.zeros((b, t), dtype=bool) + attention_mask = np.zeros((b, t), dtype=np.int64) + + for i, s in enumerate(samples): + Ti = int(s['length']) + input_ids[i, :Ti, :] = s['input_ids'] + codec_ids[i, :Ti, :] = s['codec_ids'] + codec_0_labels[i, :Ti] = s['codec_0_labels'] + text_embedding_mask[i, :Ti] = s['text_embedding_mask'] + codec_embedding_mask[i, :Ti] = s['codec_embedding_mask'] + codec_mask[i, :Ti] = s['codec_mask'] + attention_mask[i, :Ti] = s['attention_mask'] + + batch = { + 'input_ids': input_ids, + 'codec_ids': codec_ids, + 'codec_0_labels': codec_0_labels, + 'text_embedding_mask': text_embedding_mask[:, :, None], # [b,t,1] 对齐官方 + 'codec_embedding_mask': codec_embedding_mask[:, :, None], # [b,t,1] + 'codec_mask': codec_mask, + 'attention_mask': attention_mask, + } + if ref_mels is not None: + if len(ref_mels) != b: + raise ValueError('Each sample requires one ref_mel') + batch['ref_mels'] = np.stack(normalize_ref_mels(ref_mels), axis=0) + return batch + + +def normalize_ref_mels(ref_mels): + """Validate per-sample mel shapes without padding or merging speakers.""" + normalized = [] + for mel in ref_mels: + if mel.ndim == 3 and mel.shape[0] == 1: + mel = mel[0] + if mel.ndim != 2 or mel.shape[0] == 0 or mel.shape[1] != 128: + raise ValueError('Each ref_mel must have shape [F,128] or [1,F,128], F > 0') + normalized.append(mel) + if not normalized or any(m.shape != normalized[0].shape for m in normalized): + raise ValueError('ref_mels must have equal lengths; use a shared reference or length buckets') + return normalized diff --git a/cookbook/qwen3_tts/tts/qwen3/forward.py b/cookbook/qwen3_tts/tts/qwen3/forward.py new file mode 100644 index 00000000..2b14c688 --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/forward.py @@ -0,0 +1,88 @@ +"""Qwen3-TTS two-stage forward. Talker 使用内部时间 shift；残差码本 logits 交自定义同位 CE。 +小型替身可测试调用契约；真实模型/NPU/LoRA 仍需单独验收。""" + +from typing import Dict, Optional +import torch + + +def qwen3_tts_forward_step( + model, + batch: Dict[str, torch.Tensor], + target_speaker_embedding: Optional[torch.Tensor] = None, +): + """Qwen3-TTS 训练一步前向（产出中间态 dict，供 Qwen3TTSLoss 算总 loss）。 + + Args: + model: `Qwen3TTSForConditionalGeneration`（底层，= `Qwen3TTSModel.model`），需含顶层 + `.speaker_encoder` + `.talker`（含 `.model.codec_embedding`、`.model.text_embedding`、 + `.text_projection`、`.codec_head`、`.code_predictor`、`.forward_sub_talker_finetune`）。 + batch: collate_qwen3_batch 输出转 torch 后的张量字典（input_ids[B,T,2]、codec_ids[B,T,16]、 + text_embedding_mask[B,T,1]、codec_embedding_mask[B,T,1]、codec_mask[B,T]、 + attention_mask[B,T]、codec_0_labels[B,T]、ref_mels[B,F,128]）。 + target_speaker_embedding: 首步算出的目标说话人 emb（官方写进 codec_embedding.weight[3000]； + 本项目走 clone 推理，可传 None，仅用于一致性记录）。 + + Returns: + dict: + - 'talker_loss': 码本0 的 mean CE（model.talker.forward 的 outputs.loss，标准时间轴 shift 正确）。 + - 'sub_talker_logits': code_predictor 输出 [N, 15, V]（N=批次目标帧总数；15=码本1-15 头数）。 + - 'sub_talker_labels': [N, 15] = talker_codec_ids[:, 1:]（码本1-15 真值，已是 head 预测目标序列）。 + - 'target_speaker_embedding': 首步缓存的 speaker emb（供外部记录）。 + 总 loss = talker_loss + 0.3 × code_predictor CE，由 Qwen3TTSLoss 从本 dict 算（不在编排里算）。 + """ + talker = model.talker + # ---- ① speaker_encoder 前置（独立前向，detach 冻结）---- + with torch.no_grad(): + speaker_embedding = model.speaker_encoder(batch['ref_mels'].to(model.dtype)) + if target_speaker_embedding is None: + target_speaker_embedding = speaker_embedding + + # ---- ② 手动组装 inputs_embeds（不走 model(input_ids=...)）---- + input_text_ids = batch['input_ids'][..., 0] # [B,T] + input_codec_ids = batch['input_ids'][..., 1] # [B,T] + text_emb_mask = batch['text_embedding_mask'].to(model.dtype) + codec_emb_mask = batch['codec_embedding_mask'].to(model.dtype) + codec_ids = batch['codec_ids'] # [B,T,16] + codec_mask = batch['codec_mask'] # [B,T] + + # 0.6B: text_hidden_size=2048, hidden_size=1024. Even equal-sized + # variants must use the pretrained projection, as official inference does. + input_text_emb = talker.text_projection(talker.model.text_embedding(input_text_ids)) * text_emb_mask + input_codec_emb = talker.model.codec_embedding(input_codec_ids) * codec_emb_mask + input_codec_emb[:, 6, :] = speaker_embedding # 注入 pos6 + input_embeddings = input_text_emb + input_codec_emb + + # 叠加码本 1-15 的输入嵌入（code_predictor.get_input_embeddings() 返回 ModuleList） + cp_input_embs = talker.code_predictor.get_input_embeddings() # ModuleList[15] + for i in range(1, talker.config.num_code_groups): + e = cp_input_embs[i - 1](codec_ids[..., i]) * codec_mask.unsqueeze(-1).to(model.dtype) + input_embeddings = input_embeddings + e + + # ---- ③ 完整时间序列交 Talker；内部 ForCausalLMLoss 仅 shift 一次。---- + outputs = talker( + inputs_embeds=input_embeddings, + attention_mask=batch['attention_mask'], + labels=batch['codec_0_labels'], + output_hidden_states=True, + use_cache=False, + ) + talker_loss = outputs.loss # 码本0 mean shift+CE（HF loss_function，时间轴正确） + + # ---- ④ 第二段 forward：Code predictor 预测码本 1-15（向量化，一次喂所有目标帧）---- + # 注：forward_sub_talker_finetune 已向量化（输入 [N,16]+[N,H] → 内部 [N,16,H] 一次前向）， + hidden = outputs.hidden_states[0][-1] # [B, T, H] + talker_hidden = hidden[:, :-1][codec_mask[:, 1:]] # Use the state that predicted cb0, before consuming this frame; no target leakage + talker_codec_ids = codec_ids[codec_mask] # [N, 16] + sub_talker_logits, _sub_loss_ignored = talker.forward_sub_talker_finetune( + talker_codec_ids, talker_hidden + ) + # 忽略 _sub_loss_ignored（继承的 ForCausalLMLoss 对码本轴有 off-by-one，见 loss.py）； + # Qwen3TTSLoss 用 sub_talker_logits + sub_talker_labels 自己算正确的 per-head CE。 + sub_talker_labels = talker_codec_ids[:, 1:] # [N, 15] = 码本1-15 真值 + + return { + 'talker_loss': talker_loss, + 'sub_talker_logits': sub_talker_logits, # [N, 15, V] + 'sub_talker_labels': sub_talker_labels, # [N, 15] + 'target_speaker_embedding': target_speaker_embedding, + } diff --git a/cookbook/qwen3_tts/tts/qwen3/input_processor.py b/cookbook/qwen3_tts/tts/qwen3/input_processor.py new file mode 100644 index 00000000..a1399498 --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/input_processor.py @@ -0,0 +1,171 @@ +"""Qwen3-TTS 单卡 padded batch processor。 + +编码样本按时间 pad 并 stack；已 collate 的 dict 直接校验并转换设备/类型。 +参考 mel 保留 [B,F,128]，不同参考长度明确拒绝；不走通用 LM 字段过滤管线。 +""" + +from typing import Dict, List + +import numpy as np +import torch + +from twinkle.data_format import InputFeature +from twinkle import Platform, remote_function +from twinkle.processor.base import InputProcessor +from twinkle.utils.torch_utils import pad_and_stack_tensors +from .data import normalize_ref_mels + +# 序列字段：强制 stack（concat=False）→ [B, maxT, *] +_SEQ_STACK_FIELDS = ( + 'input_ids', 'codec_ids', 'codec_0_labels', + 'text_embedding_mask', 'codec_embedding_mask', 'codec_mask', + 'attention_mask', 'position_ids', +) +# mask 字段 stack 后补 channel 维 → [B, maxT, 1]（对齐官方 collate_fn + forward 广播） +_MASK_CHANNEL_FIELDS = ('text_embedding_mask', 'codec_embedding_mask') + +# to_transformers_dict 白名单追加（基类 `_keys` 之外的自定义字段） +_EXTRA_WHITELIST = ( + 'codec_ids', 'codec_0_labels', + 'text_embedding_mask', 'codec_embedding_mask', 'codec_mask', + 'ref_mels', +) + + +def _to_tensor(v) -> torch.Tensor: + """numpy/list → torch；bool → uint8（规避 F.pad 对 bool 的版本兼容问题；值仍是 0/1）。""" + if isinstance(v, np.ndarray): + v = torch.from_numpy(v) + elif isinstance(v, list): + v = torch.tensor(v) + if v.dtype == torch.bool: + v = v.to(torch.uint8) + return v + + +class Qwen3TTSInputProcessor(InputProcessor): + """Qwen3-TTS 批组织：2D 强制 stack + ref_mels stack + 自定义字段白名单。""" + + padding_map = { + **InputProcessor.padding_map, + 'codec_ids': 0, + 'codec_0_labels': -100, + 'text_embedding_mask': False, + 'codec_embedding_mask': False, + 'codec_mask': False, + 'ref_mels': 0.0, + } + + @remote_function() + def __call__(self, inputs, **kwargs): + """Single-card, padded TTS: encoded samples or an already collated dict. + + Return one tensor dict. Do not run the generic LM preparation pipeline: + it iterates dict keys, changes ranks and filters out TTS fields. + """ + if self.padding_free or kwargs.get('enable_sp', False): + raise ValueError('Qwen3 TTS supports padded single-card batches only') + if isinstance(inputs, dict) and getattr(inputs['input_ids'], 'ndim', None) == 3: + batch = self.to_transformers_dict([inputs])[0] + else: + samples = [inputs] if isinstance(inputs, dict) else inputs + if not samples: + raise ValueError('Empty TTS batch') + batch = self.to_transformers_dict(self.collate_fn(samples))[0] + required = set(_EXTRA_WHITELIST) | {'input_ids', 'attention_mask'} + missing = required - batch.keys() + if missing: + raise ValueError(f'Missing TTS fields: {sorted(missing)}') + batch = {key: torch.as_tensor(value) for key, value in batch.items()} + b, t, channels = batch['input_ids'].shape + if channels != 2 or b == 0 or t == 0: + raise ValueError('input_ids must have shape [B,T,2]') + shapes = {'codec_ids': (b, t, 16), 'codec_0_labels': (b, t), + 'attention_mask': (b, t), 'codec_mask': (b, t), + 'text_embedding_mask': (b, t, 1), 'codec_embedding_mask': (b, t, 1)} + for key, shape in shapes.items(): + if tuple(batch[key].shape) != shape: + raise ValueError(f'{key} must have shape {shape}') + mel = batch['ref_mels'] + if mel.ndim != 3 or mel.shape[0] != b or mel.shape[1] == 0 or mel.shape[2] != 128: + raise ValueError('ref_mels must have shape [B,F,128], F > 0') + if kwargs.get('micro_batch_size') not in (None, b): + raise ValueError('Split TTS micro-batches in the DataLoader') + model = kwargs.get('model') + device = getattr(model, 'device', None) + if device is None: + device = Platform.get_local_device() + for key, value in batch.items(): + dtype = (torch.float32 if key == 'ref_mels' else torch.bool + if key in (*_MASK_CHANNEL_FIELDS, 'codec_mask') else torch.long) + batch[key] = value.to(device=device, dtype=dtype) + return batch + + def collate_fn(self, inputs, micro_batch_size=None, variable_seq_lengths=False, **kwargs): + seq_out: Dict[str, torch.Tensor] = {} + + for key in _SEQ_STACK_FIELDS: + vals: List[torch.Tensor] = [] + for inp in inputs: + v = inp.get(key) + if v is None: + continue + v = _to_tensor(v) + # 撤销 prepare_inputs 对 1D 字段的 unsqueeze（[1,T]→[T]）； + # 2D 字段（[T,2]/[T,16]）shape[0]=T≠1 不受影响。 + if key not in ('input_ids', 'codec_ids') and v.dim() == 2 and v.shape[0] == 1: + v = v.squeeze(0) + vals.append(v) + if not vals: + continue + pad = int(self.padding_map.get(key, 0)) + if len(vals) == 1: + stacked = vals[0].unsqueeze(0) # 保证有 batch 维 + else: + stacked = pad_and_stack_tensors(vals, pad_value=pad, concat=False) + if key in _MASK_CHANNEL_FIELDS: + stacked = stacked.unsqueeze(-1) # [B, maxT, 1] + seq_out[key] = stacked + + # 保留每条参考的 batch 归属；无长度掩码时不自动补零。 + ref_mels: List[torch.Tensor] = [] + for inp in inputs: + v = inp.get('ref_mels') + if v is None: + continue + v = _to_tensor(v) + ref_mels.append(v) + if len(ref_mels) != len(inputs): + raise ValueError('Each sample requires one ref_mel') + seq_out['ref_mels'] = torch.stack(normalize_ref_mels(ref_mels), dim=0) + + # length：标量元数据，取 batch 内最大 + lengths = [int(inp.get('length')) for inp in inputs if inp.get('length') is not None] + if lengths: + seq_out['length'] = max(lengths) + + return [InputFeature(**seq_out)] + + @staticmethod + def to_transformers_dict(inputs, **kwargs): + results = [] + # 基类白名单 + 本模型自定义字段（基类 _keys 见 processor/base.py:663-680） + _keys = { + 'input_ids', 'inputs_embeds', 'attention_mask', 'position_ids', 'labels', + 'completion_mask', 'cu_seq_lens_q', 'cu_seq_lens_k', 'cu_seqlens_q', + 'cu_seqlens_kv', 'max_length_q', 'max_length_k', 'packed_seq_params', + 'routed_experts', 'mm_token_type_ids', 'second_per_grid_ts', + } | set(InputProcessor.VLM_CONCAT_FIELDS) | set(_EXTRA_WHITELIST) + + for _input in inputs: + output = {} + for key in list(_input.keys()): + if key not in _keys: + continue + value = _input[key] + if isinstance(value, torch.Tensor) or not isinstance(value, (list, np.ndarray)): + output[key] = value + else: + output[key] = np.array(value) + results.append(InputFeature(**output)) + return results diff --git a/cookbook/qwen3_tts/tts/qwen3/loss.py b/cookbook/qwen3_tts/tts/qwen3/loss.py new file mode 100644 index 00000000..b501422f --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/loss.py @@ -0,0 +1,48 @@ +"""Qwen3-TTS 两段 loss：Talker 时间轴 CE + 0.3 倍残差码本同位 CE。 +源码依据：Qwen3-TTS 022e286。其 residual logits 与 cb1..cb15 已对齐，再 shift 会丢 cb1 的直接监督， +并让最后 head 无直接 CE；cb15 标签会被错配给前一 head。此结论不代表官方权重的训练方式或音质收益。""" + +import torch +import torch.nn.functional as F + +from twinkle.data_format import LossOutput + +from twinkle.loss.base import Loss + + +class Qwen3TTSLoss(Loss): + """Qwen3-TTS 2 段式 loss。 + + Args: + sub_loss_weight: code_predictor 段权重（官方 0.3）。 + ignore_index: 忽略的 label 值（非音频帧 / pad）。 + """ + + require_logps = False + require_logits = True + + def __init__( + self, + sub_loss_weight: float = 0.3, + ignore_index: int = -100, + ): + self.sub_loss_weight = sub_loss_weight + self.ignore_index = ignore_index + + def __call__(self, inputs, outputs, **kwargs) -> LossOutput: + # talker 段（码本0）：透传编排算好的 mean CE（标准时间轴 shift 正确） + talker_loss = outputs['talker_loss'] + # code_predictor 段（码本1-15） + sub_logits = outputs['sub_talker_logits'] # [N, 15, V] + sub_labels = outputs['sub_talker_labels'] # [N, 15] + + vocab = sub_logits.shape[-1] + sub_loss = F.cross_entropy( + sub_logits.reshape(-1, vocab), sub_labels.reshape(-1), + ignore_index=self.ignore_index, reduction='mean') + + total = talker_loss + self.sub_loss_weight * sub_loss + # 两段均 per-token mean（镜像官方 sft_12hz.py：outputs.loss 与 sub_loss 都是 mean）。 + # num_tokens=0 走 Twinkle PER-TOKEN-MEAN 归一（grad_accum 路径）， + # 与 accelerate.backward(loss)+gradient_accumulation_steps 一致。 + return LossOutput(loss=total, num_tokens=0)