Create handler.py

handler.py

@@ -0,0 +1,104 @@
+import torch
+import torchaudio
+import numpy as np
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from livekit import rtc
+import asyncio
+import os
+
+class EndpointHandler:
+    def __init__(self, path: str = ""):
+        # Load the Orpheus TTS model and tokenizer from the given path (Hub repository).
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        path = "atharva27/orpheus"
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+        self.model.to(self.device)
+        self.model.eval()
+
+    def __call__(self, data: dict) -> list:
+        # Extract input text and optional voice and LiveKit parameters.
+        text_input = data.get("inputs") or data.get("text") or ""
+        if not isinstance(text_input, str) or text_input.strip() == "":
+            raise ValueError("No text input provided for TTS")
+        voice = data.get("voice", "tara")  # default voice (e.g., "tara")
+        
+        # Format prompt with voice name (Orpheus expects prompts like "voice: text").
+        prompt = f"{voice}: {text_input}"
+        
+        # Encode prompt and generate output tokens with the TTS model.
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
+        generate_kwargs = {
+            "max_new_tokens": 1024,               # allow sufficient tokens for audio output
+            "do_sample": True,
+            "temperature": 0.8,
+            "top_p": 0.95,
+            "repetition_penalty": 1.1,           # >=1.1 for stable speech generation
+            "pad_token_id": self.tokenizer.eos_token_id,
+        }
+        output_ids = self.model.generate(input_ids, **generate_kwargs)
+        # The generated sequence includes the prompt; isolate newly generated tokens:
+        generated_tokens = output_ids[0, input_ids.size(1):]
+        output_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=False)
+        
+        # Extract audio token IDs (assume tokens are in the output_text)
+        # This is a placeholder for token extraction, replace with actual logic.
+        audio_token_ids = [int(m) for m in output_text.split()]
+
+        # Example: convert the audio token IDs to waveform data
+        waveform = self.generate_waveform_from_tokens(audio_token_ids)
+
+        # Save or stream waveform
+        torchaudio.save("output_audio.wav", waveform, 24000)  # Save as a 24 kHz audio file
+
+        # For real-time streaming, we will use LiveKit to stream the audio
+        lk_url = data.get("livekit_url")
+        lk_token = data.get("livekit_token")
+        room_name = data.get("livekit_room", "default-room")
+
+        # Streaming logic
+        asyncio.run(self.stream_audio(lk_url, lk_token, room_name, waveform))
+
+        return [{"status": "success"}]
+
+    def generate_waveform_from_tokens(self, audio_token_ids):
+        """
+        Convert audio tokens into a waveform (this part is for demonstration).
+        You should implement a proper method to decode tokens to actual audio.
+        """
+        # Here we're simulating the waveform by generating random data based on the tokens
+        # Replace this logic with actual audio generation
+        num_samples = len(audio_token_ids) * 100  # Estimate number of samples based on tokens
+        waveform = torch.randn(1, num_samples)  # Simulate random audio waveform
+        return waveform
+
+    async def stream_audio(self, lk_url, lk_token, room_name, waveform):
+        room = rtc.Room()
+        try:
+            await room.connect(lk_url, lk_token, options=rtc.RoomOptions(auto_subscribe=True))
+        except Exception as e:
+            return f"Failed to connect to LiveKit: {e}"
+
+        # Create an audio track for streaming the TTS output
+        source = rtc.AudioSource(sample_rate=24000, num_channels=1)
+        track = rtc.LocalAudioTrack.create_audio_track("tts-audio", source)
+        await room.local_participant.publish_track(track, rtc.TrackPublishOptions(name="TTS Audio"))
+
+        # Stream the waveform data in chunks for real-time playback
+        frame_duration = 0.05  # 50 ms per frame
+        frame_samples = int(24000 * frame_duration)  # 50 ms of audio at 24 kHz sample rate
+        total_samples = waveform.size(1)
+        for start in range(0, total_samples, frame_samples):
+            end = min(start + frame_samples, total_samples)
+            chunk = waveform[:, start:end].numpy().astype(np.int16)  # Convert chunk to 16-bit PCM
+
+            # Create an AudioFrame and send to LiveKit
+            audio_frame = rtc.AudioFrame.create(24000, 1, len(chunk))
+            np.copyto(audio_frame.data, chunk)
+            await source.capture_frame(audio_frame)
+
+            # Sleep to maintain real-time pace (synchronize with frame duration)
+            await asyncio.sleep(frame_duration)
+
+        # Disconnect from the room after streaming is finished
+        await room.disconnect()