Update handler.py

Adding voice enrollment feature and serialization of inputs & outputs.

handler.py

@@ -3,7 +3,6 @@ import torch
 import numpy as np
 import librosa
 import soundfile as sf
-import uuid
 import traceback
 import base64
 import io
@@ -45,12 +44,6 @@ class EndpointHandler:
     except Exception as e:
       raise RuntimeError(f"Failed to load SNAC model: {e}")
 
-    self.ENROLLMENT_DIR = "enrollments"
-  
-    # Move to devices
-    self.voice_id = None
- 
-
   # Set up functions to format and encode text/audio
   def encode_text(self, text):
     return self.tokenizer.encode(text, return_tensors="pt", add_special_tokens=False)
@@ -91,9 +84,8 @@ class EndpointHandler:
     base64-encoded audio data
 
     Returns:
-    - voice_id (str): ID you can later use to synthesize speech
+    - cloning_features (str): serialized enrollment data
     """
-    os.makedirs(self.ENROLLMENT_DIR, exist_ok=True)
     enrollment_data = []
 
     for text, base64_audio in enrollment_pairs:
@@ -103,30 +95,15 @@ class EndpointHandler:
           "text_ids": text_ids,
           "audio_codes": audio_codes
       })
+    
+    # Serialize enrollment data
+    buffer = io.BytesIO()
+    torch.save(enrollment_data, buffer)
+    buffer.seek(0)
 
-    # Generate unique voice ID
-    voice_id = f"voice_{uuid.uuid4().hex[:6]}"
-    save_path = os.path.join(self.ENROLLMENT_DIR, f"{voice_id}.pt")
-    torch.save(enrollment_data, save_path)
-
-    self.voice_id = voice_id
-
-    return voice_id
-
-  def load_enrollment_by_id(self, voice_id):
-      """
-      Load encoded text/audio token blocks using voice ID
-
-      Returns:
-      - enrollment_data: list of dicts {text_ids, audio_codes}
-      """
-      path = os.path.join(self.ENROLLMENT_DIR, f"{voice_id}.pt")
-      if not os.path.exists(path):
-          raise FileNotFoundError(f"Voice ID '{voice_id}' not found.")
-
-      enrollment_data = torch.load(path, map_location="cpu")
-
-      return enrollment_data
+    # Encode as base64 string and assign to attribute
+    cloning_features = base64.b64encode(buffer.read()).decode('utf-8')
+    return cloning_features
 
   def prepare_audio_tokens_for_decoder(self, audio_codes_list):
     """
@@ -149,7 +126,6 @@ class EndpointHandler:
 
     return modified_audio_codes_list
 
-
   # Convert audio sample to codes and reconstruct
   def tokenize_audio(self, waveform):
     waveform = torch.from_numpy(waveform).unsqueeze(0).unsqueeze(0).to(self.device)
@@ -174,39 +150,26 @@ class EndpointHandler:
     """
     Preprocess input data before inference
     """
-    self.voice_cloning = data.get("parameters", {}).get("clone", False)
-
-    if isinstance(data, dict) and "inputs" in data:
-        target_text = data["inputs"]
-        parameters = data.get("parameters", {})
-    else:
-        target_text = data
-        parameters = {}
+    self.voice_cloning = data.get("clone", False)
 
     # Extract parameters from request
+    target_text = data["inputs"]
+    parameters = data.get("parameters", {})
+    cloning_features = data.get("cloning_features", None)
+
     temperature = float(parameters.get("temperature", 0.6))
     top_p = float(parameters.get("top_p", 0.95))
     max_new_tokens = int(parameters.get("max_new_tokens", 1200))
     repetition_penalty = float(parameters.get("repetition_penalty", 1.1))
-    enrollments = parameters.get("enrollments", [])
-    voice_id = parameters.get("voice_id", None)
 
     if self.voice_cloning:
+      """Handle voice cloning using cloning features"""
 
-      # Validate voice cloning input
-      enrollment_path = lambda vid: os.path.join(self.ENROLLMENT_DIR, f"{vid}.pt")
-
-      if voice_id:
-          if not os.path.exists(enrollment_path(voice_id)):
-              raise ValueError(f"Voice ID '{voice_id}' not found in {self.ENROLLMENT_DIR}")
-          enrollment_data = self.load_enrollment_by_id(voice_id)
-
-      elif enrollments:
-          voice_id = self.enroll_user(enrollments)
-          enrollment_data = self.load_enrollment_by_id(voice_id)
-
+      if not cloning_features:
+        raise ValueError("No cloning features were provided")
       else:
-          raise ValueError("You must provide either a valid voice_id or enrollment pairs.")
+        # Decode back into tensors
+        enrollment_data = torch.load(io.BytesIO(base64.b64decode(cloning_features)))
 
       # Process pre-tokenized enrollment_data
       input_sequence = []
@@ -238,7 +201,7 @@ class EndpointHandler:
       """Handle standard text-to-speech"""
 
       # Extract parameters from request
-      voice = parameters.get("voice", "tara")
+      voice = data.get("voice", "Eniola")
       prompt = f"{voice}: {target_text}"
       input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
 
@@ -297,10 +260,20 @@ class EndpointHandler:
     Main entry point for the handler
     """
     try:
-      preprocessed_inputs = self.preprocess(data)
-      model_outputs = self.inference(preprocessed_inputs)
-      response = self.postprocess(model_outputs)
-      return response
+      enroll_user = data.get("enroll_user", False)
+
+      if enroll_user:
+        # We extract cloning features for enrollment
+        enrollment_pairs = data.get("enrollments", [])
+        cloning_features = self.enroll_user(enrollment_pairs)
+        return {"cloning_features": cloning_features}
+      else:
+        # We want to generate speech using preset cloning features
+        preprocessed_inputs = self.preprocess(data)
+        model_outputs = self.inference(preprocessed_inputs)
+        response = self.postprocess(model_outputs)
+        return response
+
     # Catch that error, baby
     except Exception as e:
       traceback.print_exc()
@@ -434,10 +407,8 @@ class EndpointHandler:
     # Encode WAV bytes as base64
     audio_b64 = base64.b64encode(buffer.read()).decode('utf-8')
 
-    return {      
-        "generated_ids": generated_ids.tolist(),
-        "audio_sample": audio_sample,
-        "audio_b64": audio_b64,
-        "sample_rate": 24000,
-        "voice_id": self.voice_id
-    }    
+    return {
+      "audio_sample": audio_sample,
+      "audio_b64": audio_b64,
+      "sample_rate": 24000,
+    }