Spaces:
Sleeping
Sleeping
| import asyncio | |
| import time | |
| from clip_transform import CLIPTransform | |
| from chat_service import ChatService | |
| from dotenv import load_dotenv | |
| from speech_service import SpeechService | |
| from concurrent.futures import ThreadPoolExecutor | |
| from audio_stream_processor import AudioStreamProcessor | |
| from streaming_chat_service import StreamingChatService | |
| from pipeline import Pipeline, Node, Job | |
| from typing import List | |
| class ChatJob(Job): | |
| def __init__(self, data, chat_service: ChatService): | |
| super().__init__(data) | |
| self.chat_service = chat_service | |
| class Node1(Node): | |
| next_id = 0 | |
| async def process_job(self, job: ChatJob): | |
| # input job.data is the input string | |
| # output job.data is the next sentance | |
| async for sentence in job.chat_service.get_responses_as_sentances_async(job.data): | |
| if job.chat_service.ignore_sentence(sentence): | |
| continue | |
| print(f"{sentence}") | |
| new_job = ChatJob(sentence, job.chat_service) | |
| new_job.id = self.next_id | |
| self.next_id += 1 | |
| yield new_job | |
| class Node2(Node): | |
| next_id = 0 | |
| async def process_job(self, job: ChatJob): | |
| # input job.data is the sentance | |
| # output job.data is the streamed speech bytes | |
| async for chunk in job.chat_service.get_speech_chunks_async(job.data): | |
| new_job = ChatJob(chunk, job.chat_service) | |
| new_job.id = self.next_id | |
| self.next_id += 1 | |
| yield new_job | |
| class Node3(Node): | |
| # sync_size = 64 | |
| # sync = [] | |
| async def process_job(self, job: ChatJob): | |
| # input job.data is the streamed speech bytes | |
| # Node3.sync.append(job.data) | |
| job.chat_service.enqueue_speech_bytes_to_play([job.data]) | |
| yield job | |
| # if len(Node3.sync) >= Node3.sync_size: | |
| # audio_chunks = Node3.sync[:Node3.sync_size] | |
| # Node3.sync = Node3.sync[Node3.sync_size:] | |
| # job.chat_service.enqueue_speech_bytes_to_play(audio_chunks) | |
| # yield job | |
| class ChatPipeline(): | |
| def __init__(self): | |
| load_dotenv() | |
| self.pipeline = Pipeline() | |
| self.audio_processor = AudioStreamProcessor() | |
| self.chat_service = StreamingChatService(self.audio_processor, voice_id="2OviOUQc1JsQRQgNkVBj") # Chales003 | |
| def __enter__(self): | |
| return self | |
| def __exit__(self, exc_type, exc_value, traceback): | |
| self.audio_processor.close() | |
| self.audio_processor = None | |
| def __del__(self): | |
| if self.audio_processor: | |
| self.audio_processor.close() | |
| self.audio_processor = None | |
| async def start(self): | |
| self.node1_queue = asyncio.Queue() | |
| self.node2_queue = asyncio.Queue() | |
| self.node3_queue = asyncio.Queue() | |
| self.sync = [] | |
| await self.pipeline.add_node(Node1, 1, self.node1_queue, self.node2_queue, sequential_node=True) | |
| await self.pipeline.add_node(Node2, 1, self.node2_queue, self.node3_queue, sequential_node=True) | |
| await self.pipeline.add_node(Node3, 1, self.node3_queue, None, sequential_node=True) | |
| async def enqueue(self, prompt): | |
| job = ChatJob(prompt, self.chat_service) | |
| await self.pipeline.enqueue_job(job) | |
| async def wait_until_all_jobs_idle(self): | |
| # TODO - implement this | |
| while True: | |
| await asyncio.sleep(0.1) | |