README.md

# MOLM-Audio: SPDMark-Style Segment-Wise Audio Watermarking

LoRA-routing audio watermarks for three generators (HiFi-GAN, VibeVoice acoustic
decoder, DiffWave). Per SPDMark (https://arxiv.org/abs/2512.12090), each audio is
split into S=8 segments; each segment carries an HMAC-derived M-bit message
embedded via a parallel LoRA "basis dictionary." The verifier matches recovered
per-segment bits against the expected HMAC sequence using Hungarian assignment +
Binomial hypothesis test.

## Repo layout

```
checkpoints/
  diffwave_spdmark_spec/final/{lora_weights.pt, extractor.pt, diffwave_full.pt}
  diffwave_v5_2step/final/{lora_weights.pt, extractor.pt, diffwave_full.pt}
  hifigan_spdmark_spec/final/{lora_weights.pt, extractor.pt, hifigan_full.pt}
  vibevoice_spdmark_spec/final/{lora_weights.pt, extractor.pt, model_full.pt}
  vibevoice_14bit_v2/final/{lora_weights.pt, extractor.pt, model_full.pt}
code/                  # training + inference + smoke-test scripts (clone or download)
eval/<run>_<regime>/results.json
README.md
```

`lora_weights.pt` + the base model from torchaudio/pretrained is sufficient for
inference. `*_full.pt` is provided for one-step loading where the base model
isn't available locally.

Training/test wavs live in **MOLM-Audio/molm-audio-data** under `data/`.

## Setup

```bash
# 1. Clone or download this repo.
hf download MOLM-Audio/molm-audio --local-dir molm-audio

# 2. Drop the code/ contents into a Python env.
cd molm-audio/code
pip install -r requirements.txt

# 3. For DiffWave: also place the LJSpeech base checkpoint at
#    pretrained/diffwave-ljspeech.pt (download from the upstream DiffWave repo).
# 4. For VibeVoice: extract decoder + diffusion head components once via:
python precompute_vibevoice_data.py extract \
    --vibevoice_model microsoft/VibeVoice-1.5B \
    --output_dir pretrained/vibevoice
# Then encode your audio into latents with `precompute_vibevoice_data.py encode`.
# 5. For HiFi-GAN: nothing extra — torchaudio downloads V3 LJSpeech weights on first run.
```

## Checkpoints

### Current runs (segments, spec-trained, with Hungarian verifier)

| Run | Backbone | Routing | Paths | Bits/seg | Train attack |
|---|---|---|---|---|---|
| `hifigan_spdmark_spec` | HiFi-GAN V3 (torchaudio LJSpeech) | `0,1,2,3,4,5,6` | 4 | 14 | `nvlceqr` |
| `vibevoice_spdmark_spec` | VibeVoice acoustic decoder | `0,1,2,3,4,5,6` (×2 slots) | 2 | 14 | `nvlceqr` |
| `diffwave_spdmark_spec` | DiffWave (LJSpeech) | `0,4,8,12,16,20,24` | 4 | 14 | `nvlceqr` |

### Legacy runs (no segments)

| Run | Backbone | Routing | Paths | Bits/seg | Train attack | Notes |
|---|---|---|---|---|---|---|
| `vibevoice_14bit_v2` | VibeVoice acoustic decoder | `0,1,2,3,4,5,6` (×2 slots) | 2 | 14 | `nvlceq` | Pre-SPDMark-temporal-attacks training; eval JSON has no Hungarian verify block. |
| `diffwave_v5_2step` | DiffWave (LJSpeech) | `0,4,8,12,16,20,24` | 2 | 7 | `nvlceq` | 2-step diffusion, `lambda_perc 0.1` → louder watermark, faster inference. |

Spectral attack codes: `n` noise, `v` gain, `l` lowpass, `c` crop, `e` erase, `q` quantize, `r` resample. SPDMark temporal codes: `d` segment-drop, `s` segment-swap, `i` segment-insert.

## Training

All three trainers share the SPDMark plumbing in `code/molm_audio_adapter.py`
(HMAC keys via `sample_training_keys`, routing-mask construction, segment-map),
`code/audio_augmentations.py` (segment-aware aug with optional temporal
attacks), and `code/audio_losses.py` (`SegmentCombinedAudioLoss` with
`valid_mask`).

The exact commands used to produce the spec-trained checkpoints in this repo:

### HiFi-GAN segments
```bash
python code/training_molm_hifigan.py \
    --dataset_path data/train_full \
    --output_dir checkpoints/hifigan_spdmark_spec \
    --exp_name MOLM_HiFiGAN_SPDMark_spec \
    --routing_blocks 0,1,2,3,4,5,6 --num_paths 4 --lora_rank 64 \
    --num_segments 8 --key_bits 128 \
    --attack nvlceqr --aug_prob 0.25 \
    --lambda_perc 0.5 --lambda_perc_warmup_steps 4500 \
    --max_train_steps 75000 --train_batch_size 8 --gradient_accumulation_steps 4 \
    --learning_rate 2e-4 --lora_init_std 0.07 \
    --audio_seconds 2.0 \
    --checkpointing_steps 500 --logging_steps 10 --use_wandb
```

### DiffWave segments
```bash
python code/training_molm_audio.py \
    --diffwave_checkpoint pretrained/diffwave-ljspeech.pt \
    --dataset_path data/train_full \
    --output_dir checkpoints/diffwave_spdmark_spec \
    --exp_name DiffWave_SPDMark_spec \
    --routing_layers 0,4,8,12,16,20,24 --num_paths 4 --lora_rank 64 \
    --lora_alpha 24 --lora_init_std 0.01 \
    --diffusion_steps 4 \
    --num_segments 8 --key_bits 128 \
    --attack nvlceqr --aug_prob 0.25 \
    --adam_weight_decay 0.05 --lambda_lora_reg 0.01 \
    --lambda_perc 0.3 --lambda_perc_warmup_steps 2000 \
    --max_train_steps 40000 --train_batch_size 8 --gradient_accumulation_steps 2 \
    --learning_rate 3e-4 --max_grad_norm 1.0 --audio_seconds 2.0 \
    --checkpointing_steps 500 --logging_steps 10 --use_wandb
```

### VibeVoice spec
```bash
python code/training_molm_vibevoice.py \
    --components_dir pretrained/vibevoice \
    --data_dir data/vibevoice_latents_train_full \
    --output_dir checkpoints/vibevoice_spdmark_spec \
    --exp_name vibevoice_spdmark_spec \
    --routing_blocks 0,1,2,3,4,5,6 \
    --num_paths 2 --lora_rank 16 --lora_alpha 8 \
    --num_frames 8 --num_segments 8 --key_bits 128 \
    --attack nvlceqr --aug_prob 0.3 \
    --lambda_perc 0.5 --lambda_perc_warmup_steps 4500 --lambda_lora_reg 0.01 \
    --max_train_steps 120000 --cosine_cycle 1 \
    --train_batch_size 16 --gradient_accumulation_steps 2 \
    --learning_rate 1e-4 \
    --checkpointing_steps 500 --logging_steps 10 --use_wandb
```

### No segments training

`diffwave_v5_2step` (2-step diffusion, very low lambda_perc):
```bash
python code/training_molm_audio.py \
    --diffwave_checkpoint pretrained/diffwave-ljspeech.pt \
    --dataset_path data/train_full \
    --routing_layers 0,4,8,12,16,20,24 --num_paths 2 --lora_rank 64 --lora_alpha 64 \
    --diffusion_steps 2 \
    --max_train_steps 120000 --cosine_cycle 1 \
    --train_batch_size 8 --gradient_accumulation_steps 4 \
    --learning_rate 1e-4 \
    --attack nvlceq --aug_prob 0.3 \
    --lambda_perc 0.1 --lambda_perc_warmup_steps 4500 --lambda_lora_reg 0.01 \
    --output_dir checkpoints/diffwave_v5_2step \
    --checkpointing_steps 100 --logging_steps 10 --use_wandb \
    --exp_name diffwave_v5_2step
```

`vibevoice_14bit_v2`:
```bash
python code/training_molm_vibevoice.py \
    --components_dir pretrained/vibevoice \
    --data_dir data/vibevoice_latents_train_full \
    --routing_blocks 0,1,2,3,4,5,6 \
    --num_paths 2 --lora_rank 16 --lora_alpha 8 \
    --max_train_steps 120000 --cosine_cycle 1 \
    --train_batch_size 16 --gradient_accumulation_steps 2 \
    --learning_rate 1e-4 \
    --attack nvlceq --aug_prob 0.3 \
    --lambda_perc 0.5 --lambda_perc_warmup_steps 4500 --lambda_lora_reg 0.01 \
    --output_dir checkpoints/vibevoice_14bit_v2 \
    --exp_name vibevoice_14bit_v2 \
    --checkpointing_steps 500 --logging_steps 10 --use_wandb
```

Key SPDMark training flags (all three trainers):
- `--num_segments 8` — S, segments per audio clip.
- `--key_bits 128` — base-key width fed into HMAC-SHA256 for per-segment message derivation.
- `--attack <codes>` — see attack-code legend above. Append `dsi` to also train against segment-level temporal attacks (drop/swap/insert).

## Inference

All three generators share a common eval interface; `--verify` enables the
Hungarian + Binomial verifier (auto-sets `--message_scheme hmac`).

### HiFi-GAN
```bash
python code/generate_molm_hifigan.py \
    --lora_weights checkpoints/hifigan_spdmark_spec/final/lora_weights.pt \
    --extractor_weights checkpoints/hifigan_spdmark_spec/final/extractor.pt \
    --routing_blocks 0,1,2,3,4,5,6 --num_paths 4 --lora_rank 64 \
    --test_dir <path-to-wavs> --num_samples 20 \
    --chunked_generation --num_chunks 8 \
    --verify --gamma_f 0.01 --gamma_v 0.01 --attacks nvlceqr \
    --output_dir eval_out/hifigan_spec_spectral --device cuda
```

### DiffWave (current `spdmark_spec`)
```bash
python code/generate_molm_audio.py \
    --diffwave_checkpoint pretrained/diffwave-ljspeech.pt \
    --lora_weights checkpoints/diffwave_spdmark_spec/final/lora_weights.pt \
    --extractor_weights checkpoints/diffwave_spdmark_spec/final/extractor.pt \
    --routing_layers 0,4,8,12,16,20,24 --num_paths 4 --lora_rank 64 \
    --diffusion_steps 4 \
    --test_dir <path-to-wavs> --num_samples 20 \
    --chunked_generation --num_chunks 8 \
    --verify --gamma_f 0.01 --gamma_v 0.01 --attacks nvlceqr \
    --output_dir eval_out/diffwave_spec_spectral --device cuda
```

### DiffWave (no segments `v5_2step`)
```bash
python code/generate_molm_audio.py \
    --diffwave_checkpoint pretrained/diffwave-ljspeech.pt \
    --lora_weights checkpoints/diffwave_v5_2step/final/lora_weights.pt \
    --extractor_weights checkpoints/diffwave_v5_2step/final/extractor.pt \
    --routing_layers 0,4,8,12,16,20,24 --num_paths 2 --lora_rank 64 \
    --diffusion_steps 2 \
    --test_dir <path-to-wavs> --num_samples 20 \
    --chunked_generation --num_chunks 8 \
    --verify --gamma_f 0.01 --gamma_v 0.01 --attacks nvlceqr \
    --output_dir eval_out/diffwave_v5_2step_spectral --device cuda
```

### VibeVoice (current `spdmark_spec`)
```bash
python code/generate_molm_vibevoice.py \
    --components_dir pretrained/vibevoice \
    --lora_weights checkpoints/vibevoice_spdmark_spec/final/lora_weights.pt \
    --extractor_weights checkpoints/vibevoice_spdmark_spec/final/extractor.pt \
    --routing_blocks 0,1,2,3,4,5,6 --num_paths 2 --lora_rank 16 --lora_alpha 8 \
    --num_frames 8 \
    --data_dir <path-to-precomputed-latents> --num_samples 20 \
    --chunked_generation --num_chunks 8 \
    --verify --gamma_f 0.01 --gamma_v 0.01 --attacks nvlceqr \
    --output_dir eval_out/vibevoice_spec_spectral --device cuda
```

### VibeVoice (no segments `14bit_v2`)
Same flags as `spdmark_spec` but point at `checkpoints/vibevoice_14bit_v2/final/`.

## Smoke tests

Quick sanity check of the SPDMark plumbing:
```bash
python code/smoke_test_hifigan.py    # 9 checks incl. HMAC keys + segment-aware aug + verify
python code/smoke_test.py            # DiffWave equivalents
python code/smoke_test_vibevoice.py  # VibeVoice equivalents
```