Update handler.py

handler.py

@@ -17,7 +17,11 @@ class EndpointHandler:
         # process input
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", None)
-
+        
+        duration = data.pop("duration", 5.0)
+        duration = max(min(duration, 30.0), 0.0)
+        max_new_tokens = int(duration * self.model.config.audio_encoder.frame_rate)
+        
         # preprocess
         inputs = self.processor(
             text=[inputs],
@@ -26,9 +30,9 @@ class EndpointHandler:
 
         # pass inputs with all kwargs in data
         if parameters is not None:
-            outputs = self.model.generate(**inputs, max_new_tokens=256, **parameters)
+            outputs = self.model.generate(**inputs, max_new_tokens=max_new_tokens, **parameters)
         else:
-            outputs = self.model.generate(**inputs, max_new_tokens=256)
+            outputs = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
 
         # postprocess the prediction
         prediction = outputs[0].cpu().numpy()