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__":