BedVibe-TTS: the architecture, the conditioning, and the pipeline that feeds it.

{
  "id": 437, "text": "Δεν ήξερα", "language": "Greek", "emotion": "Angry",
  "dataset": "B_BLEND",
  "emotion_vector": [6, 0, 0, 0, 0, 0],
  "trait_vector":  [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
  "speaker_a": "Speaker_01", "speaker_b": "Speaker_02",
  "speaker_a_weight": 0.67,  "speaker_b_weight": 0.33,
  "blend_vector": [0.67, 0.33],
  "tokens_a": [
    "tokens/.../Speaker_01/Greek/Angry/0437.pt",    "tokens/.../Speaker_01/Greek/Happy/0437.pt",
    "tokens/.../Speaker_01/Greek/Neutral/0437.pt",  "tokens/.../Speaker_01/Greek/Scared/0437.pt",
    "tokens/.../Speaker_01/Greek/Shouting/0437.pt", "tokens/.../Speaker_01/Greek/Whisper/0437.pt"
  ],
  "tokens_b": [
    "tokens/.../Speaker_02/Greek/Angry/0437.pt",    "tokens/.../Speaker_02/Greek/Happy/0437.pt",
    "tokens/.../Speaker_02/Greek/Neutral/0437.pt",  "tokens/.../Speaker_02/Greek/Scared/0437.pt",
    "tokens/.../Speaker_02/Greek/Shouting/0437.pt", "tokens/.../Speaker_02/Greek/Whisper/0437.pt"
  ],
  "text_ids": [207, 149, 207, 182, 207, 190, 33, 207, 175, ...]
}

uint16-safe by design (all IDs validated < 65,536 at pack time); conditioning vectors stored as float16 and upcast to float32 on read — roughly half the storage of a naive float32/uint32 layout.

BF16 autocast (GradScaler disabled by design), env-selectable SDPA kernel (flash / mem-efficient / math — the same switch that surfaced the Blackwell failure documented in the report), gradient accumulation, and a deterministic per-epoch sampler (seed = base + epoch) for reproducible shuffles.

# per-epoch skip counters (returned in each epoch's summary)
skip_oom = skip_nonfinite_loss = skip_nonfinite_grad = 0
scaler = torch.cuda.amp.GradScaler(enabled=False)    # BF16, not FP16

# --- TEXT-USED test (batch 0): prove the model actually reads the text ---
out_norm = model(x=text,     x_lens=text_lens, y=audio, ...)
text_abl = torch.zeros_like(text)                    # ablate the text
out_abl  = model(x=text_abl, x_lens=torch.ones_like(text_lens), y=audio, ...)
diff = (out_norm["logits"] - out_abl["logits"]).abs().max().item()
if diff <= 0:
    raise RuntimeError("[FATAL] TEXT-USED test FAILED: logits did not "
                       "change when text was ablated. STOP TRAINING.")

try:
    with torch.autocast("cuda", dtype=torch.bfloat16):
        loss = model(...)["loss"]

    if not torch.isfinite(loss):                     # NON-FINITE LOSS GUARD
        print(f"[NONFINITE-LOSS-SKIP] batch_idx={i} loss={loss}")
        skip_nonfinite_loss += 1
        optimizer.zero_grad(set_to_none=True); continue

    (loss / grad_accum_steps).backward()
    if accum >= grad_accum_steps:                    # NON-FINITE GRAD GUARD
        bad_grad = False
        for name, p in model.named_parameters():
            if p.grad is not None and not torch.isfinite(p.grad).all():
                bad_grad = True; break
        if bad_grad:
            print(f"[NONFINITE-GRAD-SKIP] :: bad={name} :: skipping step")
            skip_nonfinite_grad += 1
            optimizer.zero_grad(set_to_none=True); continue
        optimizer.step(); optimizer.zero_grad(set_to_none=True)

except torch.cuda.OutOfMemoryError as e:             # OOM RECOVERY
    print(f"[OOM-SKIP] batch_idx={i} :: {e}")
    skip_oom += 1
    torch.cuda.empty_cache()
    optimizer.zero_grad(set_to_none=True); continue

def _make_segments(T, max_t, overlap=0, min_tail=0):
    """Slice [0, T) into VRAM-safe windows of <= max_t frames so a long
    utterance still fits in memory; overlap keeps continuity; a too-short
    tail is shifted left to stay full-length. max_t <= 0 = no segmenting."""
    if max_t <= 0 or T <= 0:
        return [(0, T)]
    step = max(1, max_t - overlap)
    segs, s = [], 0
    while s < T:
        e = min(T, s + max_t)
        segs.append((s, e))
        if e >= T: break
        s += step
    if min_tail > 0 and segs[-1][1] - segs[-1][0] < min_tail:
        segs[-1] = (max(0, T - max_t), T)            # shift the last window left
    return segs