Update handler.py

Updating the handler so it supports cloning on the fly and returns the input ids & output ids.

handler.py

@@ -32,7 +32,6 @@ class EndpointHandler:
 
     # Move to devices
     self.device = "cuda" if torch.cuda.is_available() else "cpu"
-    # self.model.to(self.device)
 
     # Load SNAC model for audio decoding
     try:
@@ -147,6 +146,7 @@ class EndpointHandler:
     # Preprocess input data before inference
 
     self.voice_cloning = data.get("clone", False)
+    clone_on_the_fly = data.get("clone_on_the_fly", False)
 
     # Extract parameters from request
     target_text = data["inputs"]
@@ -159,12 +159,27 @@ class EndpointHandler:
     repetition_penalty = float(parameters.get("repetition_penalty", 1.1))
 
     if self.voice_cloning:
-      """Handle voice cloning using cloning features"""
-
-      if not cloning_features:
+      if clone_on_the_fly:
+        # Clone using text-audio enrollment pair
+        enrollment_pairs = data.get("enrollments", [])
+        enrollment_data = []
+
+        # Raise error if no enrollment is provided
+        if not enrollment_pairs:
+          raise ValueError("No enrollment pairs provided")
+        
+        for text, base64_audio in enrollment_pairs:
+          text_ids = self.encode_text(text).cpu()
+          audio_codes = self.encode_audio(base64_audio)
+          enrollment_data.append({
+              "text_ids": text_ids,
+              "audio_codes": audio_codes
+          })
+
+      elif not cloning_features:
         raise ValueError("No cloning features were provided")
       else:
-        # Decode back into tensors
+        # Clone using enrollment features gotten earlier
         enrollment_data = torch.load(io.BytesIO(base64.b64decode(cloning_features)))
 
       # Process pre-tokenized enrollment_data
@@ -187,13 +202,11 @@ class EndpointHandler:
 
       # Final input tensor
       input_ids = torch.cat(input_sequence, dim=1)
-      
-      # Heuristic to determine max_new_tokens based on empirical relationship
-      # between the length of the prompt ids and the length of the generated ids
-      prompt_ids = self.encode_text(target_text)
-      max_new_tokens = int(prompt_ids.size()[1] * 20 + 200)
 
+      # Create attention mask and move tensors to device
+      attention_mask = torch.ones_like(input_ids)
       input_ids = input_ids.to(self.device)
+      attention_mask = attention_mask.to(self.device)
 
     else:
       # Handle standard text-to-speech
@@ -237,7 +250,11 @@ class EndpointHandler:
     # Forward pass through the model
     generated_ids = self.model.generate(prompt_string, sampling_params)
 
-    return torch.tensor(generated_ids[0].outputs[0].token_ids).unsqueeze(0)
+    # return torch.tensor(generated_ids[0].outputs[0].token_ids).unsqueeze(0)
+    return {
+        "gen_ids": torch.tensor(generated_ids[0].outputs[0].token_ids).unsqueeze(0), 
+        "input_ids": input_ids
+        }
 
   def __call__(self, data):
     
@@ -293,7 +310,10 @@ class EndpointHandler:
     audio_hat = self.snac_model.decode(codes)
     return audio_hat
 
-  def postprocess(self, generated_ids):
+  def postprocess(self, model_outputs):
+
+    generated_ids = model_outputs["gen_ids"]
+    input_ids = model_outputs["input_ids"]
 
     if self.voice_cloning:
       """
@@ -357,4 +377,7 @@ class EndpointHandler:
       "audio_sample": audio_sample,
       "audio_b64": audio_b64,
       "sample_rate": 24000,
+      "gen_ids": generated_ids,
+      "input_ids_len": input_ids.shape[1],
+      "gen_ids_len": generated_ids.shape[1]
     }