AudioCraft CLI - MusicGen Implemented [DRAFT PR]

.github/workflows/labgraph_audiogen.yml

@@ -0,0 +1,29 @@
+name: AudioCraft Tests
+
+on: [push]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v2
+
+      - name: Setup Python
+        uses: actions/setup-python@v2
+        with: 
+          python-version: '3.8'
+
+      - name: Install dependencies
+        run: |
+          cd extensions/labgraph_audiogen
+          python -m pip install --upgrade pip
+          pip install -e .
+          pip install pytest
+          sudo apt-get install ffmpeg
+
+      - name: Run Tests
+        run: |
+          cd extensions/labgraph_audiogen
+          pytest -vvv

extensions/labgraph_audiogen/README.md

@@ -0,0 +1,54 @@
+# LabGraph AudioGen Extension
+
+A powerful Python command-line command for Windows. This tool facilitates the creation, processing, and autocompletion high-quality audio and music, using the AudioCraft models, AudioGen, and MusicGen.
+
+The command is the `lg_audiogen` command, which can be installed with the `labgraph_audiogen` extension.
+
+## Features
+
+- Ability to specify duration of the generated audio.
+- Ability to generate music based on a batch file.
+- Ability to specify the model to be used for the audio generation.
+- Ability to set the output file name.
+- Ability to generate music within 4 different models.
+
+## Setup
+
+Audiocraft needs Python 3.8 or higher to run. If you have a suitable version of Python installed, you can install Audiogen with pip:
+
+```bash
+cd extensions/labgraph_audiogen  # Only if you are not already in this directory
+pip install -e .
+```
+
+## Usage
+
+```bash
+lg_audiogen --help
+```
+
+Usage: lg_audiogen [OPTIONS] [DESCRIPTION]...
+
+A command-line command to facilitate the usage of the models of Audiocraft
+
+Options:
+
+- `--version:` Show the version and exit.
+- `-d, --duration:` INTEGER with the duration of the audio
+- `-m, --model:` Name of the model to use between *[audiogen-medium | musicgen-small | musicgen-medium | musicgen-melody | musicgen-large]*
+- `-o, --output:` TEXT with the name of the output file
+- `-b, --batch:` PATH to file for batch audio description.
+- `--help:` Show this message and exit.
+
+## Examples
+
+```bash
+lg_audiogen -m musicgen-melody -d 15 -o "mariachi_rnb" "An rnb beat with some mariachi trumpets"
+
+lg_audiogen -m musicgen-small "A happy mood in a rainy day"
+```
+
+Outputs:
+
+- [mariachi_rnb.wav](https://drive.google.com/file/d/1OKXcZtRIqfL_dvfRFGtZV7VwC9CCUDX-/view)
+- [A_happy_mood_in_a_rainy_day.wav](https://drive.google.com/file/d/1cCFhL7PP8FO0uzkEeRlhbDSGoflsBJ-G/view)

extensions/labgraph_audiogen/lg_audiogen/main.py

@@ -0,0 +1,131 @@
+import click
+
+
+# All models supported by the lg_audiogen command
+SUPPORTED_MODELS = [
+    "audiogen-medium",
+    "musicgen-small",
+    "musicgen-medium",
+    "musicgen-melody",
+    "musicgen-large"
+]
+
+# Musicgen models supported by the lg_audiogen command
+MUSICGEN_MODELS = [
+    "musicgen-small",
+    "musicgen-medium",
+    "musicgen-melody",
+    "musicgen-large"
+]
+
+
+def generate_text_music(descriptions, duration, output, musicgen_model):
+    """
+    Generate music from the given descritptions and save it on the outputs folder
+
+    @param descriptions: list of descriptions to generate music from
+    @param duration: duration of the audio to generate
+    @param output: name of the output file
+    @param musicgen_model: name of the musicgen model to use
+    """
+    click.secho("\nStarting the music generation from the given descriptions", fg="bright_blue")
+
+    # Import MusicGen, audio_write in this function to avoid time consuming imports
+    from audiocraft.models import MusicGen
+    from audiocraft.data.audio import audio_write
+
+    # Load the corresponding MusicGen model and set the generation parameters
+    model = MusicGen.get_pretrained(f"facebook/{musicgen_model}")
+    model.set_generation_params(duration=duration)
+
+    filenames = []
+    # Validate the descriptions and filenames
+    for i, description in enumerate(descriptions):
+        if not output:
+            filenames.append(f"{description.replace(' ', '_')}")
+        else:
+            filenames.append(output.replace(' ', '_') + str(i))
+
+        if len(description) == 0:
+            raise click.BadParameter(
+                click.style(
+                    f"Description too short for {description}, "
+                    "please use a non-empty description",
+                    fg="bright_red"
+                )
+            )
+
+        click.secho(
+            f"Generating music from '{description}' written on the '{filenames[i]}.wav' file",
+            fg="bright_green"
+        )
+
+    try:
+        # Generate the music from the descriptions
+        music = model.generate(descriptions, progress=True)
+        # Save the music on the outputs folder with the corresponding filename
+        for i, generation in enumerate(music):
+            audio_write(f"outputs/{filenames[i]}", generation.cpu(), model.sample_rate,
+                        strategy="loudness", loudness_compressor=True)
+
+            click.secho(
+                f"Audio generated and saved on the outputs/{filenames[i]}.wav file",
+                bg="green", fg="black"
+            )
+
+    except Exception as music_error:
+        click.secho(f"Error generating music: {music_error}", bg="red", fg="white")
+
+
+@click.command()
+@click.version_option()
+@click.argument('description', nargs=-1, required=False)
+@click.option("-d", "--duration", type=int, prompt_required=False, help="Duration of the audio")
+@click.option("-m", "--model", type=click.Choice(SUPPORTED_MODELS), help="Name of the model to use")
+@click.option("-o", "--output", prompt_required=False, help="Name of the output file")
+@click.option('--batch', '-b', type=click.Path(), help='File name for batch audio description.')
+def parse_arguments(description, duration, model, output, batch):
+    """
+    A command-line command to facilitate the usage of the models of Audiocraft
+    """
+    # Validate batch and description
+    if batch:
+        try:
+            with open(batch, mode='r', encoding='utf-8') as f:
+                descriptions = [line.strip() for line in f.readlines()]
+        except FileNotFoundError as file_error:
+            raise click.FileError(filename=batch, hint=click.style(file_error, fg="bright_red"))
+    else:
+        if not description:
+            raise click.BadParameter(
+                click.style(
+                    "Description argument is required when not using --batch.",
+                    fg="bright_red"
+                )
+            )
+        descriptions = [' '.join(description)]
+
+    # Validate model and duration
+    if not model:
+        raise click.BadParameter(
+            click.style(
+                "No model provided, select a supported model with -m / --model "
+                f"(eg. -m musicgen-medium) between {', '.join(SUPPORTED_MODELS)}",
+                fg="bright_red"
+            )
+        )
+    if not duration:
+        click.secho(
+            "No duration provided, select a duration with -d / --duration, 10 seconds will be used",
+            fg="yellow"
+        )
+        duration = 10
+    if duration <= 0:
+        raise click.BadParameter(
+            click.style(
+                "Duration must be greater than 0", fg="bright_red"
+            )
+        )
+
+    if model in MUSICGEN_MODELS:
+        generate_text_music(descriptions, duration, output, model)

extensions/labgraph_audiogen/setup.py

@@ -0,0 +1,24 @@
+from setuptools import setup
+
+
+setup(
+    name="lg_audiogen",
+    version="0.0.1",
+    description="A command-line command for the usage of AudioCraft for labgraph",
+    long_description="""
+    A command-line command to facilitate the usage of the models of Audiocraft 
+    to generate and process audio on labgraph
+    """,
+    packages=["lg_audiogen"],
+    install_requires=[
+        "Click>=8.1.7",
+        "torch>=2.1.0",
+        "torchaudio>=2.1.0",
+        "audiocraft>=1.0.0",
+    ],
+    entry_points={
+        'console_scripts': [
+            'lg_audiogen=lg_audiogen.main:parse_arguments',
+        ],
+    },
+)

extensions/labgraph_audiogen/tests/test_audiocraft_cli.py

@@ -0,0 +1,24 @@
+import os
+import subprocess
+
+
+def test_text_to_smallmusicgen():
+    """
+    Test the music generation from text with the musicgen-small model
+
+    It requires the lg_audiogen command to be installed
+    """
+    # Run the lg_audiogen -m musicgen-small -d 3 -o test_audio.wav "A happy song"
+    print(subprocess.run(
+        ["lg_audiogen", "-m", "musicgen-small", "-d", "3", "-o", "test_audio", "A happy song"],
+        capture_output=True, text=True))
+    # Assert file exists and size
+    print(subprocess.run(["ls"], capture_output=True, text=True))
+    print(subprocess.run(["pwd"], capture_output=True, text=True))
+    print(subprocess.run(["ls", ".."], capture_output=True, text=True))
+    print(subprocess.run(["ls", "../.."], capture_output=True, text=True))
+    assert os.path.exists("outputs/test_audio0.wav"), "The file test_audio0.wav does not exist"
+    assert os.path.getsize("outputs/test_audio0.wav") > 0, "The file test_audio0.wav is empty"
+    # Remove the file
+    os.remove("outputs/test_audio0.wav")
+    os.rmdir("outputs")