Laionbox v0.1 Wip

LaionBox v0.1 — Differentiable-Reward LoRA for DramaBox TTS

A LoRA adapter for the 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 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

# 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 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

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 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 (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:

• 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:

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. At high sigma, the x0 prediction is dominated by noise, so decoded audio is meaningless — reward signals would be random. About 23% of micro-batches pass this threshold.

Training Runs Summary

We ran 7 iterations to arrive at the final approach:

Run	Approach	Naturalness Trend	Verdict
1	Reward-weighted recon, cap=2	Flat	Same gradient direction as flow matching
2	Higher coefficients (ratio=5)	Flat	Same problem, stronger magnitude
3	+ Sigma filter + large batch	Flat	Still same gradient direction
4	Rejection sampling (top 50%)	Flat	Filters bad examples but no direction
5	Differentiable reward (broken CLAP)	Flat	@torch.no_grad() killed gradient chain
6	Diff reward + gradient fix (1ep)	+0.013	Gradients flow, rewards move
7	Diff reward + gradient fix (5ep)	+0.089	Sustained improvement confirmed

How to Further Fine-Tune

Resume Training

cd Voice-Acting-Pipeline

# Edit configs/finetune_diff_reward_5ep.yaml:
#   resume_lora: ./lora_epoch5.safetensors
#   epochs: 3  # or however many more
#   output_dir: ./finetune_output/your_run

accelerate launch --num_processes=8 \
    scripts/dramabox_finetune_train_multi_aux.py \
    --config configs/finetune_diff_reward_5ep.yaml

Prepare Your Own Data

Each training sample needs:

• Audio latent: Encoded through DramaBox's VAE (8×T×16 at ~25fps)
• Text condition: Gemma-3-12B-IT hidden states for the prompt
• Mode: voice_clone_fwd, voice_clone_rev, or unconditional
• Reference latent (optional): For voice cloning modes

See training_code/dramabox_finetune_train_multi_aux.py for the full data loading pipeline.

Key Hyperparameters

Parameter	Our value	Notes
lora_rank	128	Higher = more capacity, more VRAM
lora_alpha	128	Equal to rank (standard)
lr	4e-5	Peak learning rate
grad_accum	32	Global batch = GPUs × grad_accum
aux_target_ratio	5.0	Each aux loss targets 5× flow magnitude
coeff_cap	10.0	Max coefficient multiplier
aux_sigma_max	0.4	Only compute aux when sigma < this
differentiable_reward	true	Backprop through frozen reward models
diff_reward_checkpoint	true	Gradient checkpointing on aux models

VRAM Requirements

Component	VRAM
LTX-2.3 transformer (bf16, grad ckpt)	~44 GB
LoRA (rank 128) + AdamW optimizer	~3.7 GB
AudioDecoder + BigVGAN vocoder (frozen)	~0.5 GB
VoiceCLAP-small (frozen)	~0.3 GB
WavLM-SV (frozen, float32)	~0.2 GB
Quality MLP + centroids	~0.01 GB
Activations (checkpointed)	~18-25 GB
Total	~67-72 GB

Requires A100 80GB or equivalent. H100 80GB also works.

Insights: What We Learned About Auxiliary Losses for Diffusion Models

Reward-weighted reconstruction does not work

If you compute a scalar reward R and multiply it by the x0 reconstruction loss, the gradient is R * ∇(||x0_pred - x0_clean||2). This...