laion/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||²). This always points toward x0_clean — the reward R only changes the step size, never the direction. No matter how sophisticated your reward function, the model learns the same thing as pure flow matching, just faster or slower.

Differentiable rewards require gradient chain integrity

The entire path from trainable parameters to the loss function must allow gradient flow. A single @torch.no_grad() decorator anywhere in the chain (even on a seemingly internal helper function like mel spectrogram computation) will silently kill all gradient information. The training will appear to work (loss values exist, no errors) but rewards will be flat.

VoiceCLAP's mel computation kills gradients

VoiceCLAP-small (and likely other audio CLAP models) decorates compute_log_mel with @torch.no_grad(). The torch.stft inside is fully differentiable — the decorator is unnecessary and harmful for reward training. Our encode_clap_waveform_differentiable() function replicates the mel computation without the decorator, enabling proper gradient flow.

Sigma threshold is essential

At high noise levels (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:

@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 for details.