loop efficiency w/ cupy

augly/audio/functional.py

@@ -4,13 +4,18 @@
 import math
 from copy import deepcopy
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+import importlib
 
 import augly.audio.utils as audutils
 import numpy as np
 import torch
 from augly.utils import DEFAULT_SAMPLE_RATE
 from augly.utils.libsndfile import install_libsndfile
 
+cupy_spec = importlib.util.find_spec("cupy")
+cupy_found = cupy_spec is not None
+if cupy_found:
+    import cupy as cp
 
 install_libsndfile()
 import librosa
@@ -645,10 +650,14 @@ def loop(
     """
     assert isinstance(n, int) and n >= 0, "Expected 'n' to be a nonnegative integer"
     audio, sample_rate = audutils.validate_and_load_audio(audio, sample_rate)
+    if cupy_found:
+        audio = cp.array(audio)
 
     aug_audio = audio
     for _ in range(n):
-        aug_audio = np.append(aug_audio, audio, axis=(0 if audio.ndim == 1 else 1))
+        aug_audio = (cp if cupy_found else np).append(
+            aug_audio, audio, axis=(0 if audio.ndim == 1 else 1)
+        )
 
     audutils.get_metadata(
         metadata=metadata,

setup.py

@@ -13,7 +13,8 @@
 extra_requirements = {
     "av": [
         r for r in Path("av_requirements.txt").read_text().splitlines() if "@" not in r
-    ]
+    ],
+    "gpu": ["cupy"],
 }
 
 with open("README.md", encoding="utf8") as f: