Added: Stereo Music Generation Support

server.py

@@ -15,8 +15,9 @@ 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("--model", type=str, default="musicgen-stereo-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")
 
@@ -25,6 +26,12 @@ args = parser.parse_args()
 # Initialize the FastAPI app
 app = FastAPI()
 
+# Build the model name based on the provided arguments
+if args.model.startswith('facebook/'):
+    args.model_name = args.model
+else:
+    args.model_name = f"facebook/{args.model}"
+
 # Load the model with the provided arguments
 musicgen_model = musicgen.MusicGen.get_pretrained(args.model_name, device=args.device)
 
@@ -37,12 +44,13 @@ 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()
+        result = result.squeeze().cpu().numpy().T
         
         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")