Added: Music Duration Control from the client end

client.py

@@ -16,14 +16,16 @@ parser.add_argument(
 parser.add_argument(
     "--output_file", type=str, default="output.wav", help="Output file name"
 )
+parser.add_argument(
+    "--duration", type=int, default=10, help="Duration of generated music in seconds"
+)
 
 args = parser.parse_args()
 
-
-def generate_music(server_url, prompts, output_file):
+def generate_music(server_url, prompts, duration, output_file):
     url = f"{server_url}/generate_music"
     headers = {"Content-Type": "application/json"}
-    data = {"prompts": prompts}
+    data = {"prompts": prompts, "duration": duration}
 
     response = requests.post(url, json=data, headers=headers)
 
@@ -34,6 +36,5 @@ def generate_music(server_url, prompts, 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)
+    generate_music(args.server_url, args.prompts, args.duration, args.output_file)

server.py

@@ -2,7 +2,7 @@ import warnings
 import argparse
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from typing import List
+from typing import List, Optional
 import torch
 from audiocraft.models import musicgen
 import numpy as np
@@ -17,7 +17,6 @@ warnings.simplefilter('ignore')
 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")
 
@@ -28,14 +27,15 @@ 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]
+    duration: Optional[int] = 10  # Default duration is 10 seconds if not provided
 
 @app.post("/generate_music")
 def generate_music(request: MusicRequest):
     try:
+        musicgen_model.set_generation_params(duration=request.duration)
         result = musicgen_model.generate(request.prompts, progress=False)
         result = result.squeeze().cpu().numpy()