add gpu implementation (#2)

* add wav2aug gpu

* make freq drop safer

* fix pol inv

* add no grads

* fix speed perturb

* add torchaudio

* removed cpu implementation

* remove cpu tests

* add general utils and tests

* add linter

* clean docstrings

* resamp test aud, cpu when no gpu avail

* fix dev install

* touch up readme, formatting

* remove cuda hard requirement on gpu ops

* remove cuda requirement from tests

Author: gfdb

.github/workflows/black.yml

@@ -0,0 +1,25 @@
+name: black
+
+on:
+  push:
+    branches: [ main, master, add-gpu ]
+  pull_request:
+    branches: [ main, master ]
+
+jobs:
+  format-check:
+    name: Check Python formatting with Black
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+      - name: Install dev dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install .[dev]
+      - name: Run Black (check only)
+        run: |
+          black --check .

.gitignore

@@ -16,3 +16,4 @@ htmlcov/
 # OS
 .DS_Store
 
+test.py

pyproject.toml

@@ -27,7 +27,8 @@ classifiers = [
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 dependencies = [
-    "torch>=2.8.0",
+    "torch>=2.0.0",
+    "torchaudio>=2.8.0",
     "torchcodec>=0.7.0",
 ]
 
@@ -36,6 +37,9 @@ test = [
     "pytest>=8",
     "pytest-cov",
 ]
+dev = [
+  "black>=23.1.0,<25",
+]
 
 [project.urls]
 Homepage = "https://github.com/gfdb/wav2aug"
@@ -49,4 +53,15 @@ include-package-data = true
 packages = { find = { where = ["."], include = ["wav2aug*"], exclude = ["tests*"], namespaces = true } }
 
 [tool.setuptools.package-data]
-wav2aug = ["py.typed", "**/*.yaml", "assets/**"]
+wav2aug = ["py.typed", "**/*.yaml", "assets/**"]
+
+
+[tool.black]
+line-length = 88
+target-version = ["py310"]
+# Include only python source files
+include = "\\.pyi?$"
+# Exclude common build, virtualenv and cache directories
+exclude = '''
+/(\.|build|dist|__pycache__|\.venv|venv|\.eggs|\.git|\.hg|\.mypy_cache|\.pytest_cache)/
+'''

readme.md

@@ -1,64 +1,117 @@
-# Wav2Aug: Toward Universal Time-Domain Speech Augmentation 
+# 🎛️ Wav2Aug: Toward Universal Time-Domain Speech Augmentation
 
-A minimalistic PyTorch-based audio augmentation library for speech and audio processing. The goal of this library is to provide a general purpose speech augmentation policy that can be used on any task and perform well without having to tune augmentation hyperparameters. Just install, and start augmenting. Applies two random augmentations per call.
+A minimalistic PyTorch-based audio augmentation library for speech and audio augmentation. The goal of this library is to provide a general purpose speech augmentation policy that can be used on any task and perform well without having to tune augmentation hyperparameters. Just install, and start augmenting. Applies two random augmentations per call.
 
 ![Diagram](https://raw.githubusercontent.com/gfdb/wav2aug/main/wav2aug.png)
 
-## Features
+## ⚙️ Features
 
-- **Minimal dependencies**: We only rely on PyTorch and torchcodec.
-- **9 core augmentations**: amplitude scaling/clipping, noise addition, frequency dropout, polarity inversion, chunk swapping, speed perturbation, time dropout, and babble noise.
-- **In-place operations**: All cpu augmentations are done in place.
+* **Minimal dependencies**: we only rely on PyTorch, torchcodec, and torchaudio.
+* **9 core augmentations**: amplitude scaling/clipping, noise addition, frequency dropout, polarity inversion, chunk swapping, speed perturbation, time dropout, and babble noise.
+* **Simplicity**: just install and start augmenting!
+* **Randomness**: all stochastic ops use PyTorch RNGs. Set a single seed and be done, e.g. torch.manual_seed(0); torch.cuda.manual_seed_all(0)
 
-## Installation
+## 📦 Installation
 
 ### pip
+
 ```bash
 pip install wav2aug
 ```
 
 ### uv
+
 ```bash
 uv add wav2aug
 ```
 
-## Quick Start
+## 🚀 Quick Start
 
 ```python
 import torch
-from wav2aug import Wav2Aug
+from wav2aug.gpu import Wav2Aug
 
-# Initialize augmenter
-aug = Wav2Aug(sample_rate=16000)
+# Initialize the augmenter once
+augmenter = Wav2Aug(sample_rate=16000)
 
-# Process audio (supports [T] mono or [C, T] multi-channel)
-waveform = torch.randn(8000)  # 0.5s at 16kHz
-augmented = aug(waveform)
+# in the forward pass
+wavs = torch.randn(3, 50000)
+lens = torch.ones((wavs.size(0)))
+
+aug_wavs, aug_lens = augmenter(wavs, lens)
 ```
 
-## Augmentation Types
+That's it!
+
+## 🧪 Augmentation Types
 
-- **Amplitude Scaling/Clipping**: Random gain and peak limiting
-- **Noise Addition**: Environmental noise with SNR control
-- **Frequency Dropout**: Spectral masking with random notch filters
-- **Polarity Inversion**: Random phase flip
-- **Chunk Swapping**: Temporal segment reordering
-- **Speed Perturbation**: Time-scale modification
-- **Time Dropout**: Random silence insertion
-- **Babble Noise**: Multi-speaker background (auto-enabled with sufficient buffer)
+* 🔊 **Amplitude Scaling/Clipping**: Random gain and peak limiting
+* 🌫️ **Noise Addition**: Environmental noise with SNR control
+* 📶 **Frequency Dropout**: Spectral masking with random notch filters
+* 🔄 **Polarity Inversion**: Random phase flip
+* 🧩 **Chunk Swapping**: Temporal segment reordering
+* ⏱️ **Speed Perturbation**: Time-scale modification
+* 🕳️ **Time Dropout**: Random silence insertion
+* 👥 **Babble Noise**: Multi-speaker background (auto-enabled with sufficient buffer)
 
-## Development Installation
+## 🛠️ Development Installation
+
+### uv
 
 ```bash
 git clone https://github.com/gfdb/wav2aug
 cd wav2aug
-uv python pin 3.10 # or greater
+
+# create venv and pin Python
+uv venv
+source .venv/bin/activate
+uv python pin 3.10  # or 3.11/3.12
+
+# runtime only
 uv sync
-uv sync --extra test # for test deps
+
+# extras
+uv sync --extra dev
+uv sync --extra test
 ```
 
-## Tests
+### pip
 
 ```bash
-uv run pytest tests/
+git clone https://github.com/gfdb/wav2aug
+cd wav2aug
+
+# create venv
+python -m venv .venv
+source .venv/bin/activate
+
+# runtime only
+python -m pip install .
+
+# editable + extras for development
+python -m pip install -e '.[dev,test]'
 ```
+
+## ✅ Tests
+
+### uv
+
+```bash
+uv run pytest -q tests/
+```
+
+### pip
+
+```bash
+pytest -q tests/
+```
+
+## 🤝 Contributing
+
+* Issues and PRs are welcome and encouraged!
+
+* Bug reports: please open an issue with a minimal repro (env, torch/torchaudio/torchcodec versions, code snippet, expected vs. actual, traceback).
+
+* Feature requests: please open an issue with use-case and proposed feature.
+
+* PRs: keep them focused. Add tests when behavior changes. Don't forget to run formatters and tests before submitting!