up

handler.py

@@ -22,15 +22,15 @@ class EndpointHandler:
         inputs = self.processor(
             text=[inputs],
             padding=True,
-            return_tensors="pt",)
+            return_tensors="pt",).to("cuda")
 
         # 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=256, **parameters)
         else:
-            outputs = self.model.generate(inputs, max_new_tokens=256)
+            outputs = self.model.generate(**inputs, max_new_tokens=256)
 
         # postprocess the prediction
-        prediction = outputs[0].numpy()
+        prediction = outputs[0].cpu().numpy()
 
         return [{"generated_audio": prediction}]