feat(transformers): add support to Dia

Signed-off-by: Ettore Di Giacinto <[email protected]>

Author: mudler

backend/python/transformers/backend.py

@@ -22,7 +22,7 @@
 
 XPU=os.environ.get("XPU", "0") == "1"
 from transformers import AutoTokenizer, AutoModel, set_seed, TextIteratorStreamer, StoppingCriteriaList, StopStringCriteria, MambaConfig, MambaForCausalLM
-from transformers import AutoProcessor, MusicgenForConditionalGeneration
+from transformers import AutoProcessor, MusicgenForConditionalGeneration, DiaForConditionalGeneration
 from scipy.io import wavfile
 import outetts
 from sentence_transformers import SentenceTransformer
@@ -90,6 +90,7 @@ def LoadModel(self, request, context):
         self.CUDA = torch.cuda.is_available()
         self.OV=False
         self.OuteTTS=False
+        self.DiaTTS=False
         self.SentenceTransformer = False
 
         device_map="cpu"
@@ -202,6 +203,11 @@ def LoadModel(self, request, context):
                 autoTokenizer = False
                 self.processor = AutoProcessor.from_pretrained(model_name)
                 self.model = MusicgenForConditionalGeneration.from_pretrained(model_name)
+            elif request.Type == "DiaForConditionalGeneration":
+                autoTokenizer = False
+                self.processor = AutoProcessor.from_pretrained(model_name)
+                self.model = DiaForConditionalGeneration.from_pretrained(model_name)
+                self.DiaTTS = True
             elif request.Type == "OuteTTS":
                 autoTokenizer = False
                 options = request.Options
@@ -506,6 +512,50 @@ def SoundGeneration(self, request, context):
             return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
         return backend_pb2.Result(success=True)
 
+
+    def DiaTTS(self, request, context):
+        """
+        Generates dialogue audio using the Dia model.
+        
+        Args:
+            request: A TTSRequest containing text dialogue and generation parameters
+            context: The gRPC context
+            
+        Returns:
+            A Result object indicating success or failure
+        """
+        try:
+            print("[DiaTTS] generating dialogue audio", file=sys.stderr)
+            
+            # Prepare text input - expect dialogue format like [S1] ... [S2] ...
+            text = [request.text]
+            
+            # Process the input
+            inputs = self.processor(text=text, padding=True, return_tensors="pt")
+            
+            # Generate audio with default Dia parameters
+            outputs = self.model.generate(
+                **inputs, 
+                max_new_tokens=3072, 
+                guidance_scale=3.0, 
+                temperature=1.8, 
+                top_p=0.90, 
+                top_k=45
+            )
+            
+            # Decode and save audio
+            outputs = self.processor.batch_decode(outputs)
+            self.processor.save_audio(outputs, request.dst)
+            
+            print("[DiaTTS] Generated dialogue audio", file=sys.stderr)
+            print("[DiaTTS] Audio saved to", request.dst, file=sys.stderr)
+            print("[DiaTTS] Dialogue generation done", file=sys.stderr)
+            
+        except Exception as err:
+            return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
+        return backend_pb2.Result(success=True)
+
+
     def OuteTTS(self, request, context):
         try:
             print("[OuteTTS] generating TTS", file=sys.stderr)
@@ -529,6 +579,9 @@ def OuteTTS(self, request, context):
     def TTS(self, request, context):
         if self.OuteTTS:
             return self.OuteTTS(request, context)
+        
+        if self.DiaTTS:
+            return self.DiaTTS(request, context)
 
         model_name = request.model
         try: