Add Qwen3-TTS data and forward-loss contracts (1/4)

cookbook/qwen3_tts/.gitignore

@@ -0,0 +1,3 @@
+outputs/
+__pycache__/
+*.py[cod]

cookbook/qwen3_tts/README.md

@@ -0,0 +1,29 @@
+# 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.
+
+
+## Causal conditioning correction
+
+Residual-codebook training uses the Talker state **before** consuming the target
+audio frame, matching generation's `past_hidden`. Selecting the state at the
+target frame leaks its codebook embeddings into the conditioning. A regression
+changes the target residual codes and checks that the current conditioning stays
+unchanged while the next frame's conditioning changes. Earlier update/reload
+tests did not establish this causal equivalence.
+
+The corrected forward and targeted contracts were validated on a single Ascend
+NPU workflow, followed by real-data multi-step training and paired listening.
+These results do not establish quality for arbitrary datasets or checkpoints.

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

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)

cookbook/qwen3_tts/tests/batch1/isolated_support.py

@@ -0,0 +1,46 @@
+"""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.common', 'tts/common'), ('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)
+    from tts.common.multi_codebook_loss import MultiCodebookLoss
+    sys.modules['tts.common'].MultiCodebookLoss = MultiCodebookLoss

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)

cookbook/qwen3_tts/tts/__init__.py

@@ -0,0 +1 @@
+"""Local Qwen3-TTS cookbook support; no global Twinkle registration."""

cookbook/qwen3_tts/tts/common/__init__.py

@@ -0,0 +1 @@
+from .multi_codebook_loss import MultiCodebookLoss