# LaionBox v0.1 — Differentiable-Reward LoRA for DramaBox TTS A LoRA adapter for the [DramaBox](https://github.com/ResembleAI/DramaBox) LTX-2.3 text-to-speech model, trained with **differentiable auxiliary rewards** (CLAP naturalness, centroid real/fake scoring, WavLM speaker similarity) to improve voice naturalness while preserving speaker identity. ## What This Is This is a **rank-128 LoRA** (226M parameters, 906MB) that modifies the audio self-attention and feed-forward layers of the LTX-2.3 22B transformer. It was trained on 3,845 samples (3,247 DramaBox synthetic + 598 Emilia real speech) for 19 epochs total: - **Epochs 1–12**: Base IC-LoRA training (voice cloning with flow matching loss only) - **Epochs 13–14**: + CLAP auxiliary loss (positive/negative text similarity) - **Epochs 15–19**: + Differentiable reward training with 3 auxiliary losses through frozen reward models The LoRA targets 6 module types across 10 transformer blocks: ``` audio_attn1.to_q, audio_attn1.to_k, audio_attn1.to_v, audio_attn1.to_out.0 audio_ff.net.0.proj, audio_ff.net.2 ``` ## What It Does Better Than Baseline | Metric | Baseline (epoch 14) | This LoRA (epoch 19) | Change | |--------|---------------------|----------------------|--------| | Naturalness reward | 0.265 | 0.354 | **+33% relative** | | Quality probability | 0.873 | 0.939 | +7.6% | | Speaker similarity | 0.904 | 0.898 | -0.7% (stable) | | Flow matching loss | 0.509 | 0.521 | +2.4% (expected tradeoff) | The model produces audio that scores significantly higher on CLAP-based naturalness metrics — sounding less robotic and more like genuine human speech — while maintaining speaker identity fidelity above 0.89 cosine similarity. ## How to Use ### Prerequisites - [DramaBox](https://github.com/ResembleAI/DramaBox) repository cloned with models downloaded - Python 3.10+, PyTorch 2.6+, PEFT, safetensors - An A100 80GB or equivalent GPU (the base model is 22B parameters) ### Inference ```python # From the DramaBox repository root: python src/inference.py \ --voice-sample reference.wav \ --prompt "(softly, with warmth) \"I've been thinking about what you said.\"" \ --checkpoint models/ltx-2.3-22b-dev-audio-only-v13-merged.safetensors \ --full-checkpoint models/ltx-2.3-22b-dev.safetensors \ --lora path/to/lora_epoch5.safetensors \ --lora-rank 128 \ --output output.wav ``` The `--lora` flag loads the LoRA weights, applies them via PEFT, and merges into the base model before inference. The inference script auto-detects the PEFT weight format and handles key mapping. ### Unconditional (no voice reference) ```python python src/inference.py \ --prompt "(excitedly) \"This is amazing!\"" \ --checkpoint models/ltx-2.3-22b-dev-audio-only-v13-merged.safetensors \ --full-checkpoint models/ltx-2.3-22b-dev.safetensors \ --lora path/to/lora_epoch5.safetensors \ --lora-rank 128 \ --output output.wav ``` ### Loading Manually ```python from peft import LoraConfig, get_peft_model from safetensors.torch import load_file # After building the LTX velocity model: lora_config = LoraConfig( r=128, lora_alpha=128, lora_dropout=0.0, bias="none", target_modules=[ "audio_attn1.to_k", "audio_attn1.to_q", "audio_attn1.to_v", "audio_attn1.to_out.0", "audio_ff.net.0.proj", "audio_ff.net.2", ], ) velocity_model = get_peft_model(velocity_model, lora_config) lora_sd = load_file("lora_epoch5.safetensors") # Map key format if needed (lora_A.weight -> lora_A.default.weight) mapped = {} for k, v in lora_sd.items(): k = k.replace(".lora_A.weight", ".lora_A.default.weight") k = k.replace(".lora_B.weight", ".lora_B.default.weight") mapped[k] = v velocity_model.load_state_dict(mapped, strict=False) velocity_model = velocity_model.merge_and_unload() ``` ## Prompt Format DramaBox uses a specific prompt format with stage directions in parentheses and dialogue in double quotes: ``` A young woman with a warm, slightly breathy voice delivers this studio-quality recording. (gently, almost whispering) "I never thought I'd see you again." (pause, then with growing emotion) "After all these years..." ``` For multi-scene prompts with emotional transitions: ``` A confident male speaker with a deep baritone delivers this high-quality studio recording. (calmly, measured) "The reports all check out. Everything's in order." CUT TO: (urgently, voice cracking) "We need to evacuate. Now. Right now." ``` ## Training Data The training dataset consists of 3,845 samples: - **3,247 DramaBox samples**: Synthetic TTS audio generated by the base DramaBox model, with MOSS-Audio-refined prompts that match the actual performance (not the original generation prompt) - **598 Emilia samples**: Real human speech from the [Emilia](https://huggingface.co/datasets/amphion/Emilia-Dataset) dataset, filtered to the top 10% by DNS MOS score (≥ 3.478) Each sample includes: - Pre-encoded audio latents (DramaBox VAE, 8×T×16 at ~25fps) - Pre-encoded text conditions (Gemma-3-12B-IT 4-bit hidden states) - Training mode labels: `voice_clone_fwd`, `voice_clone_rev`, `unconditional` Training dataset available at: [laion/laionbox-voice-acting-training-data](https://huggingface.co/datasets/laion/laionbox-voice-acting-training-data) (upload in progress) ## Training Process ### Base Training (Epochs 1–12) Standard IC-LoRA training with flow matching loss on the 3-mode setup: - **voice_clone_fwd**: Reference audio prepended, target follows - **voice_clone_rev**: Target first, reference appended - **unconditional**: No reference audio, text conditioning only Best flow loss: 0.4645 at step 920. ### CLAP Auxiliary Loss (Epochs 13–14) Added a CLAP-based auxiliary loss using [VoiceCLAP-small](https://huggingface.co/laion/voiceclap-small): - Decode x0 prediction → waveform → CLAP embedding - Score against positive text ("realistic, genuine, spontaneous...") and negative text ("distorted, unnatural, robotic...") - Reward = cos(pred, positive) - cos(pred, negative) - Reward modulates x0 reconstruction loss weight This was computed inside `torch.no_grad()` — the reward didn't backpropagate through the decoder or CLAP model. It provided a weak signal (rewards stayed flat). Best flow loss: 0.4474. ### Differentiable Reward Training (Epochs 15–19) — The Key Innovation **The critical insight**: reward-weighted reconstruction loss always produces gradients in the same direction as flow matching (toward x0_clean), regardless of the reward value. The reward only scales the magnitude, never the direction. To actually steer the model toward higher-reward outputs, you need **gradients through the reward function itself**. We removed `torch.no_grad()` from the reward computation path, allowing gradients to flow: ``` LoRA params → velocity prediction → x0 recovery → VAE decoder → waveform → VoiceCLAP → CLAP embedding → naturalness loss → WavLM-SV → speaker embedding → speaker similarity loss ``` This required fixing two bugs: 1. **VoiceCLAP's `compute_log_mel` has `@torch.no_grad()`**: Despite `torch.stft` being fully differentiable, the decorator kills all gradients. We wrote `encode_clap_waveform_differentiable()` that replicates the mel computation without the decorator. 2. **`Wav2Vec2ForXVector` loads wrong weights for WavLM**: The checkpoint stores keys as `wavlm.encoder.layers.*` but `Wav2Vec2ForXVector` expects `wav2vec2.encoder.layers.*`. Fix: use `WavLMForXVector`. ### Three Auxiliary Losses 1. **Naturalness** (CLAP text similarity + quality MLP): - `0.5 * (cos(emb, pos_text) - cos(emb, neg_text)) + 0.5 * (2 * quality_prob - 1)` - Quality MLP: 768→128→32→1 binary classifier trained on clean vs distorted audio 2. **Centroid real/fake** (zero-shot linear classifier): - `cos(emb, real_centroid) - cos(emb, synth_centroid)` - Centroids precomputed from ~2,600 real (Emilia) and ~2,600 synthetic (DramaBox) CLAP embeddings - 94.3% accuracy as zero-shot classifier, with continuous scores suitable for gradient optimization 3. **Speaker similarity** (WavLM-SV, only for voice clone modes): - `cos(WavLM(pred_wav), WavLM(ref_wav))` - Reference embedding is detached (only prediction carries gradient) - Skipped for unconditional mode (no reference) ### Adaptive Coefficient Normalization Each aux loss is independently normalized using EMA-based adaptive coefficients: ```python ema_flow = 0.95 * ema_flow + 0.05 * flow_loss ema_aux_i = 0.95 * ema_aux_i + 0.05 * abs(aux_i_loss) coeff_i = min(target_ratio * ema_flow / max(ema_aux_i, 1e-8), coeff_cap) total_loss = flow + coeff1 * aux1 + coeff2 * aux2 + coeff3 * aux3 ``` With `target_ratio=5.0` and `coeff_cap=10.0`, each aux loss targets 5× the flow matching magnitude. ### Sigma Threshold Auxiliary losses are only computed when the diffusion noise level `sigma 0.4), the x0 prediction `noisy - sigma * velocity_pred` is dominated by noise. Decoding this through the VAE produces meaningless audio, and reward signals computed on it are random noise. Only compute auxiliary losses in the low-sigma regime where predictions are close to clean audio. ### Centroid loss may conflict with naturalness The centroid real/fake score always hit its coefficient cap (10.0) and never improved, while naturalness steadily increased. These two objectives may be in tension — the "real" direction in CLAP space (toward Emilia training data distribution) is not identical to the "natural-sounding" direction (toward positive text prompts). ### Speaker similarity is already near ceiling WavLM-SV speaker similarity started at 0.904 and stayed within 0.89-0.91 throughout training. The IC-LoRA architecture already preserves speaker identity well. Differentiable speaker loss provides minimal additional benefit but doesn't hurt. ### Quality probability trends upward reliably The quality MLP probability (clean vs distorted detector) showed the clearest monotonic improvement: 0.873 → 0.939 over 5 epochs. This suggests the differentiable CLAP path successfully steers audio away from distortion patterns. ## Limitations 1. **Centroid score degraded slightly** (-0.012 over 5 epochs). The model moved toward CLAP naturalness at the expense of embedding proximity to the real speech centroid. These objectives partially conflict. 2. **Flow matching loss increased slightly** (+0.012). This is the expected tradeoff when auxiliary losses steer the model — it sacrifices some distributional matching for perceptual quality. 3. **Validation inference failed** during training due to a missing `bitsandbytes` dependency (needed for Gemma 4-bit inference in the validation subprocess). The LoRA checkpoints themselves are fine. 4. **Training data is small** (3,845 samples). Larger datasets may yield stronger improvements. 5. **Single naturalness objective**. The positive/negative text prompts are hardcoded. Different prompt choices would steer quality in different directions. 6. **DramaBox-specific**. This LoRA only works with the LTX-2.3 22B audio-only transformer architecture. ## Files in This Repository | File | Description | |------|-------------| | `lora_epoch5.safetensors` | LoRA weights (226M params, 906MB, float32) | | `adapter_config.json` | PEFT adapter configuration | | `training_args.yaml` | Exact training configuration used | | `metrics.jsonl` | Step-level training metrics (102 logged steps) | | `training_code/dramabox_finetune_train_multi_aux.py` | Full training script (1878 lines) | | `configs/finetune_diff_reward_5ep.yaml` | Training config YAML | ## Citation If you use this work, please cite: ```bibtex @misc{laionbox2026, title={LaionBox v0.1: Differentiable-Reward LoRA for DramaBox TTS}, author={LAION}, year={2026}, url={https://huggingface.co/laion/laionbox-v0.1-wip} } ``` ## License This LoRA adapter follows the same license as the base DramaBox model. See [DramaBox](https://github.com/ResembleAI/DramaBox) for details.
