Spaces:
Sleeping
Sleeping
| import asyncio | |
| import ray | |
| from ray.util.queue import Queue, Full | |
| class FFMpegConverter: | |
| def __init__(self, output_queue: Queue, buffer_size: int = 1920, output_format: str='s16le'): | |
| self.output_queue = output_queue | |
| self.buffer_size = buffer_size | |
| self.output_format = output_format | |
| self.input_pipe = None | |
| self.output_pipe = None | |
| self.process = None | |
| self.running = True | |
| async def run(self): | |
| while self.running: | |
| try: | |
| chunk = await self.output_pipe.readexactly(self.buffer_size) | |
| except asyncio.IncompleteReadError: | |
| # exit if we have finsihsed the process | |
| if self.running == False: | |
| return | |
| # If the pipe is broken, restart the process. | |
| await self.start_process() | |
| continue | |
| if self.running == False: | |
| return | |
| # print(f"FFMpegConverter: read {len(chunk)} bytes") | |
| chunk_ref = ray.put(chunk) | |
| keep_trying = True | |
| while keep_trying: | |
| try: | |
| await self.output_queue.put_async(chunk_ref, timeout=0.01) | |
| keep_trying = False | |
| except Full: | |
| if self.running == False: | |
| return | |
| await asyncio.sleep(0.01) | |
| async def start_process(self): | |
| cmd = [ | |
| 'ffmpeg', | |
| '-i', 'pipe:0', # read from stdin | |
| '-f', self.output_format, | |
| '-ar', '48000', | |
| '-ac', '1', | |
| 'pipe:1' # write to stdout | |
| ] | |
| self.process = await asyncio.create_subprocess_exec( | |
| *cmd, | |
| stdin=asyncio.subprocess.PIPE, | |
| stdout=asyncio.subprocess.PIPE, | |
| stderr=asyncio.subprocess.PIPE | |
| ) | |
| self.input_pipe = self.process.stdin | |
| self.output_pipe = self.process.stdout | |
| assert self.input_pipe is not None, "input_pipe was not initialized" | |
| # print (f"input_pipe: {self.input_pipe}") | |
| async def push_chunk(self, chunk): | |
| if self.running == False: | |
| return | |
| try: | |
| self.input_pipe.write(chunk) | |
| await self.input_pipe.drain() | |
| except BrokenPipeError: | |
| # If the pipe is broken, restart the process. | |
| await self.start_process() | |
| self.input_pipe.write(chunk) | |
| await self.input_pipe.drain() | |
| async def close(self): | |
| self.running = False # Stop the loop inside run() | |
| if self.process: | |
| self.process.stdin.transport.close() | |
| self.process.kill() | |
| self.process.terminate() | |