# Audex Enhancement VAE

This directory contains the Audex enhancement VAE used to convert
XCodec1-decoded 16 kHz mono WAV audio into enhanced 48 kHz mono WAV audio.

The enhancement model is a postprocessor for Audex text-to-audio generation:

```text
Audex generation -> XCodec1 decode -> 16 kHz WAV -> Enhancement VAE -> 48 kHz WAV
```

The model is intended for XCodec1-decoded Audex audio as inputs. It is not a general
purpose enhancer for arbitrary audio.


## Requirements
- `torch, numpy, scipy`

## Command Line Usage

For a folder of outputs:

```bash
python enhancement_VAE/enhance_audio_48k.py \
  --input tta_outputs \
  --output-dir tta_outputs_enhanced_48k
```

For a single WAV file:

```bash
python enhancement_VAE/enhance_audio_48k.py \
  --input tta_outputs/example.wav \
  --output-dir tta_outputs_enhanced_48k
```

Outputs are written as:

```text
<input_stem>_enhanced_48k.wav
```

Supported options:

- `--device`: inference device. Defaults to `cuda` when available, otherwise `cpu`.
- `--seed`: torch seed for stochastic VAE sampling. Defaults to `0`.
- `--deterministic`: use the posterior mean instead of VAE sampling.

## Python API

```python
from pathlib import Path
import torch

from enhancement_VAE.enhancement_vae import enhance_file, load_model

root = Path("enhancement_VAE")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = load_model(
    checkpoint_path=root / "XCodec_RVQ4_mono_causal_fp32.safetensors",
    config_path=root / "config.json",
    device=device,
)

enhance_file(
    model=model,
    input_path=Path("input_16k.wav"),
    output_path=Path("input_16k_enhanced_48k.wav"),
    deterministic=False,
)
```

## Input and Output

Input:

- XCodec1-decoded Mono 16 kHz WAV file

Output:

- Mono 48 kHz audio WAV file

If a directory is passed to `--input`, all `.wav` files directly inside that
directory are processed. Directory traversal is not recursive.