Added enhancement_VAE/README.md

enhancement_VAE/README.md

@@ -0,0 +1,85 @@
+# 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.