Added base Server and Client for MusicGen

.gitignore

@@ -0,0 +1,167 @@
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+
+*.wav
+*.mp3

client.py

@@ -0,0 +1,39 @@
+import requests
+import argparse
+
+# Parse command line arguments
+parser = argparse.ArgumentParser(description="Music Generation Client")
+parser.add_argument(
+    "--server_url", type=str, default="http://localhost:8000", help="URL of the server"
+)
+parser.add_argument(
+    "--prompts",
+    nargs="+",
+    type=str,
+    default=["Lofi Music for Coding"],
+    help="Prompts for music generation",
+)
+parser.add_argument(
+    "--output_file", type=str, default="output.wav", help="Output file name"
+)
+
+args = parser.parse_args()
+
+
+def generate_music(server_url, prompts, output_file):
+    url = f"{server_url}/generate_music"
+    headers = {"Content-Type": "application/json"}
+    data = {"prompts": prompts}
+
+    response = requests.post(url, json=data, headers=headers)
+
+    if response.status_code == 200:
+        with open(output_file, "wb") as f:
+            f.write(response.content)
+        print(f"Music saved to {output_file}")
+    else:
+        print(f"Failed to generate music: {response.status_code}, {response.text}")
+
+
+if __name__ == "__main__":
+    generate_music(args.server_url, args.prompts, args.output_file)

run_test.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+
+echo "Script started."
+
+# Run server
+echo "Starting server..."
+python server.py --duration 10 &
+echo "Server started."
+
+# Sleep
+echo "Waiting for the server to startup..."
+sleep 10
+
+# Run client
+echo "Starting client..."
+python client.py --server_url http://localhost:8000 --prompts "Lofi Music for Coding" --output_file output.wav 
+echo "Client finished." 
+
+
+# Kill server
+echo "Killing server..."
+kill $(ps aux | grep 'server.py' | awk '{print $2}') 
+
+
+# Done
+sleep 5
+echo "Script finished."
+

server.py

@@ -0,0 +1,53 @@
+import warnings
+import argparse
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List
+import torch
+from audiocraft.models import musicgen
+import numpy as np
+import io
+from fastapi.responses import StreamingResponse
+from scipy.io.wavfile import write as wav_write
+import uvicorn
+
+warnings.simplefilter('ignore')
+
+# Parse command line arguments
+parser = argparse.ArgumentParser(description="Music Generation Server")
+parser.add_argument("--model_name", type=str, default="small", help="Pretrained model name")
+parser.add_argument("--device", type=str, default="cuda", help="Device to load the model on")
+parser.add_argument("--duration", type=int, default=10, help="Duration of generated music in seconds")
+parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to run the server on")
+parser.add_argument("--port", type=int, default=8000, help="Port to run the server on")
+
+args = parser.parse_args()
+
+# Initialize the FastAPI app
+app = FastAPI()
+
+# Load the model with the provided arguments
+musicgen_model = musicgen.MusicGen.get_pretrained(args.model_name, device=args.device)
+musicgen_model.set_generation_params(duration=args.duration)
+
+class MusicRequest(BaseModel):
+    prompts: List[str]
+
+@app.post("/generate_music")
+def generate_music(request: MusicRequest):
+    try:
+        result = musicgen_model.generate(request.prompts, progress=False)
+        result = result.squeeze().cpu().numpy()
+        
+        sample_rate = musicgen_model.sample_rate
+        
+        buffer = io.BytesIO()
+        wav_write(buffer, sample_rate, result)
+        buffer.seek(0)
+        
+        return StreamingResponse(buffer, media_type="audio/wav")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    uvicorn.run(app, host=args.host, port=args.port)