Update handler.py

handler.py

@@ -1,43 +1,29 @@
-from typing import Dict, Any
+from typing import Dict, List, Any
 from transformers import AutoProcessor, MusicgenForConditionalGeneration
-import scipy
-import io
+import torch
 
 class EndpointHandler:
-    def __init__(self, path=""):
-        # Explicitly load processor with local files
-        self.processor = AutoProcessor.from_pretrained(
-            path,
-            local_files_only=True,
-            trust_remote_code=True
-        )
-        self.model = MusicgenForConditionalGeneration.from_pretrained(
-            path,
-            local_files_only=True,
-            trust_remote_code=True
-        )
+    def __init__(self, model_path):
+        self.processor = AutoProcessor.from_pretrained(model_path)
+        self.model = MusicgenForConditionalGeneration.from_pretrained(model_path)
+        if torch.cuda.is_available():
+            self.model = self.model.to("cuda")
 
-    def __call__(self, data: Dict[str, Any]) -> bytes:
-        text = data.get("inputs", "")
-        duration = data.get("parameters", {}).get("duration", 5)
-        
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         inputs = self.processor(
-            text=[text],
-            return_tensors="pt",
+            text=data["text"],
+            audio=data.get("audio", None),
             padding=True,
-            truncation=True
+            sampling_rate=data.get("sampling_rate", None),
+            return_tensors="pt",
         )
-        
+        if torch.cuda.is_available():
+            inputs = {k: v.to("cuda") for k, v in inputs.items()}
+
         audio_values = self.model.generate(
             **inputs,
-            max_new_tokens=int(duration * 50)
+            do_sample=data.get("do_sample", True),
+            guidance_scale=data.get("guidance_scale", 3),
+            max_new_tokens=data.get("max_new_tokens", 256),
         )
-        
-        sampling_rate = self.model.config.audio_encoder.sampling_rate
-        with io.BytesIO() as wav_io:
-            scipy.io.wavfile.write(
-                wav_io,
-                rate=sampling_rate,
-                data=audio_values[0, 0].numpy()
-            )
-            return wav_io.getvalue()
+        return {"audio_values": audio_values.cpu().numpy()}