dude_tts/frontend.py

"""XML conditioning for the two synchronized inference lanes."""
import torch
from .interleave import interleave_xml
from .layout import native_prefix

def duplex_prefix(record, config, encode_text, encode_event=None):
    body = interleave_xml(record['xml'], encode_text, -1, encode_event=encode_event)
    headers, voices, speaker_offsets = [], [], []
    for channel in ['A', 'B']:
        row = record['channels'][channel]
        ids = row['text_ids']
        if len(ids) < 8:
            raise ValueError('Native assistant template must include its header and trailer')
        # Preserve native prompts/controls, while the actual content comes only
        # from XML -> DAG. The speaker-ID embedding is replaced by the frozen
        # source speaker's reference vector below.
        pt, pc = native_prefix(dict(row, text_ids=ids[:3] + ids[-5:]), config)
        speaker = config.talker_config.spk_id[row['speaker'].lower()]
        sites = (pc[:, 0] == speaker).nonzero().flatten()
        if len(sites) != 1:
            raise ValueError('Expected exactly one native speaker conditioning site')
        offset = int(sites[0])
        pc[offset] = -1
        vector = torch.as_tensor(row['voice_embedding'], dtype=torch.float32)
        if vector.shape != (config.talker_config.hidden_size,) or not torch.isfinite(vector).all():
            raise ValueError('Invalid source-speaker reference embedding')
        voices.append(vector)
        speaker_offsets.append(offset)
        headers.append((pt[:-2], pc[:-2]))  # Reserve text EOS/audio BOS for after the body.
    header_len = max(len(pt) for pt, _ in headers)
    body_len = len(body.lanes['A'])
    length = 2 * (header_len + body_len + 2)
    text = torch.full((length,), -1, dtype=torch.long)
    codec = torch.full((length, 16), -1, dtype=torch.long)
    mask = torch.zeros(length, dtype=torch.bool)
    voice_positions = []
    for ch, channel in enumerate(['A', 'B']):
        pt, pc = headers[ch]
        left = header_len - len(pt)
        sites = 2 * torch.arange(left, header_len) + ch
        text[sites], codec[sites], mask[sites] = pt, pc, True
        voice_positions.append(2 * (left + speaker_offsets[ch]) + ch)
        ids = torch.tensor(body.lanes[channel], dtype=torch.long)
        occupied = torch.tensor(body.occupied[channel])
        sites = 2 * torch.arange(header_len, header_len + body_len) + ch
        text[sites] = ids
        codec[sites[occupied], 0] = config.talker_config.codec_pad_id
        mask[sites] = occupied
        last = occupied.nonzero().flatten()
        lane_end = int(last[-1]) + 1 if len(last) else 0
        end = 2 * (header_len + lane_end) + ch
        bos = 2 * (header_len + body_len + 1) + ch
        # Text EOS follows this lane's final text/event token. Padding keeps
        # every anchor position and the common audio BOS unchanged.
        text[end:bos:2] = config.tts_pad_token_id
        codec[end:bos:2, 0] = config.talker_config.codec_pad_id
        text[end], codec[end, 0], mask[end] = config.tts_eos_token_id, config.talker_config.codec_pad_id, True
        text[bos], codec[bos, 0], mask[bos] = config.tts_pad_token_id, config.talker_config.codec_bos_id, True
    return dict(text_ids=text, codec_ids=codec, attention_mask=mask,
                voice_positions=torch.tensor(voice_positions), voice_embeddings=torch.stack(voices))