added params parser to handler

handler.py

@@ -2,6 +2,32 @@ from typing import Dict, List, Any
 from transformers import AutoProcessor, MusicgenForConditionalGeneration
 import torch
 
+def create_params(params, fr):
+    # default
+    out = { "do_sample": True,
+          "guidance_scale": 3, 
+          "max_new_tokens": 256
+          }
+
+    has_tokens = False
+
+    if params is None:
+       return out
+
+    if 'duration' in params:
+        out['max_new_tokens'] =  params['duration'] * fr
+        has_tokens = True
+
+    for k, p in params.items():
+        if k in out: 
+          if has_tokens and k == 'max_new_tokens':
+            continue
+
+          out[k] = p
+
+    return out
+
+
 class EndpointHandler:
     def __init__(self, path="pbotsaris/musicgen-small"):
         # load model and processor
@@ -25,17 +51,21 @@ class EndpointHandler:
             return_tensors="pt"
         ).to('cuda')
 
-        if params is not None:
-            with torch.cuda.amp.autocast():
-                outputs = self.model.generate(**inputs, **params)
-        else:
-            with torch.cuda.amp.autocast():
-                outputs = self.model.generate(**inputs)
+        params = create_params(params, self.model.config.audio_encoder.frame_rate)
+
+        with torch.cuda.amp.autocast():
+            outputs = self.model.generate(**inputs, **params)
 
         pred = outputs[0].cpu().numpy().tolist()
+        sr = 32000
+
+        try:
+             sr = self.model.config.audio_encoder.sampling_rate
 
-        return [{"audio": pred, "sr": self.model.config.sampling_rate}]
+        except:
+             sr = 32000
 
+        return [{"audio": pred, "sr":sr}]
 
 
 if __name__ == "__main__":