From bfa860f1342e377447c12f1c0410b7856a9bf941 Mon Sep 17 00:00:00 2001
From: Tuan Tran <{ID}+{username}@users.noreply.github.com>
Date: Mon, 3 Jun 2024 21:57:16 -0400
Subject: [PATCH] make optional dependecy group for wm

---
 .github/actions/audiocraft_build/action.yml | 2 +-
 README.md                                   | 1 +
 setup.py                                    | 3 ++-
 3 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/.github/actions/audiocraft_build/action.yml b/.github/actions/audiocraft_build/action.yml
index a264f1c1..74ab1e38 100644
--- a/.github/actions/audiocraft_build/action.yml
+++ b/.github/actions/audiocraft_build/action.yml
@@ -23,7 +23,7 @@ runs:
       python -m pip install --upgrade pip
       pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0
       pip install xformers
-      pip install -e '.[dev]'
+      pip install -e '.[dev,wm]'
   - name: System Dependencies
     shell: bash
     run: |
diff --git a/README.md b/README.md
index 8fcc03e3..4560ea24 100644
--- a/README.md
+++ b/README.md
@@ -20,6 +20,7 @@ python -m pip install setuptools wheel
 python -m pip install -U audiocraft  # stable release
 python -m pip install -U git+https://git@github.com/facebookresearch/audiocraft#egg=audiocraft  # bleeding edge
 python -m pip install -e .  # or if you cloned the repo locally (mandatory if you want to train).
+python -m pip install -e '.[wm]'  # if you want to train a watermarking model
 ```
 
 We also recommend having `ffmpeg` installed, either through your system or Anaconda:
diff --git a/setup.py b/setup.py
index d7b39e7a..83c40d6c 100644
--- a/setup.py
+++ b/setup.py
@@ -46,7 +46,8 @@
     python_requires=REQUIRES_PYTHON,
     install_requires=REQUIRED,
     extras_require={
-        'dev': ['coverage', 'flake8', 'mypy', 'pdoc3', 'pytest', 'audioseal'],
+        'dev': ['coverage', 'flake8', 'mypy', 'pdoc3', 'pytest'],
+        'wm': ['audioseal'],
     },
     packages=[p for p in find_packages() if p.startswith('audiocraft')],
     package_data={'audiocraft': ['py.typed']},