Spaces:
Running
Running
Download media_file_response.py from ZHIWEI666/ComfyUI-Ranking-API: direct link, hf CLI and curl.
- Browser
- Download file 5.32 kB
-
https://huggingface.co/spaces/ZHIWEI666/ComfyUI-Ranking-API/resolve/main/media_file_response.py
- Command line
-
hf download hf://spaces/ZHIWEI666/ComfyUI-Ranking-API/media_file_response.py
-
curl -L -o media_file_response.py https://huggingface.co/spaces/ZHIWEI666/ComfyUI-Ranking-API/resolve/main/media_file_response.py
5.32 kB
| """媒体缓存文件的有界读取、Range交付与真实线程关闭。""" | |
| import asyncio | |
| import os | |
| from secrets import token_hex | |
| import stat | |
| from fastapi.responses import FileResponse | |
| from starlette.datastructures import MutableHeaders | |
| from owned_thread import run_owned_thread | |
| from workflow_stream_response import OwnedDelivery | |
| class OwnedMediaFileResponse(OwnedDelivery, FileResponse): | |
| total_timeout = 300 | |
| operation_timeout = 15 | |
| eof_timeout = 5 | |
| def __init__(self, path, **kwargs): | |
| snapshot = os.stat(path) | |
| if not stat.S_ISREG(snapshot.st_mode): | |
| raise ValueError('invalid media cache file') | |
| super().__init__(path, stat_result=snapshot, **kwargs) | |
| self._file = None | |
| self._init_delivery(self._close_file) | |
| async def send_owned(self, scope, receive, send): | |
| async def bounded_send(message): | |
| final = message['type'] == 'http.response.body' and not message.get('more_body', False) | |
| async with asyncio.timeout(self.eof_timeout if final else self.operation_timeout): | |
| await send(message) | |
| async with asyncio.timeout(self.total_timeout): | |
| return await super().send_owned(scope, receive, bounded_send) | |
| async def _file_work(self, function, *args): | |
| async with asyncio.timeout(self.operation_timeout): | |
| return await run_owned_thread(function, *args) | |
| async def _open_file(self): | |
| def open_file(): | |
| self._file = open(self.path, 'rb') | |
| self._check_snapshot(os.fstat(self._file.fileno())) | |
| await self._file_work(open_file) | |
| def _check_snapshot(self, current): | |
| fields = ('st_dev', 'st_ino', 'st_size', 'st_mtime_ns') | |
| if os.name != 'nt': | |
| fields += ('st_ctime_ns',) | |
| if (not stat.S_ISREG(current.st_mode) or not current.st_ino or not self.stat_result.st_ino | |
| or any(getattr(current, key) != getattr(self.stat_result, key) for key in fields)): | |
| raise RuntimeError('media cache file changed') | |
| async def _prepare_file(self, header_only): | |
| if header_only: | |
| current = await self._file_work(os.stat, self.path) | |
| self._check_snapshot(current) | |
| else: | |
| await self._open_file() | |
| async def _close_file(self): | |
| if self._file is None: | |
| return | |
| def close_file(): | |
| try: | |
| self._file.close() | |
| finally: | |
| self._file = None | |
| await run_owned_thread(close_file) | |
| async def _send_part(self, send, start, end, *, final): | |
| await self._file_work(self._file.seek, start) | |
| remaining = end - start | |
| if remaining == 0: | |
| await send({'type': 'http.response.body', 'body': b'', 'more_body': not final}) | |
| while remaining: | |
| chunk = await self._file_work(self._file.read, min(self.chunk_size, remaining)) | |
| if not chunk: | |
| raise RuntimeError('media cache file truncated') | |
| remaining -= len(chunk) | |
| await send({'type': 'http.response.body', 'body': chunk, 'more_body': remaining > 0 or not final}) | |
| async def _handle_simple(self, send, send_header_only, send_pathsend): | |
| await self._prepare_file(send_header_only) | |
| await send({'type': 'http.response.start', 'status': self.status_code, 'headers': self.raw_headers}) | |
| if send_header_only: | |
| await send({'type': 'http.response.body', 'body': b'', 'more_body': False}) | |
| return | |
| await self._send_part(send, 0, self.stat_result.st_size, final=True) | |
| async def _handle_single_range(self, send, start, end, file_size, send_header_only): | |
| await self._prepare_file(send_header_only) | |
| headers = MutableHeaders(raw=list(self.raw_headers)) | |
| headers['content-range'] = f'bytes {start}-{end - 1}/{file_size}' | |
| headers['content-length'] = str(end - start) | |
| await send({'type': 'http.response.start', 'status': 206, 'headers': headers.raw}) | |
| if send_header_only: | |
| await send({'type': 'http.response.body', 'body': b'', 'more_body': False}) | |
| return | |
| await self._send_part(send, start, end, final=True) | |
| async def _handle_multiple_ranges(self, send, ranges, file_size, send_header_only): | |
| await self._prepare_file(send_header_only) | |
| boundary = token_hex(13) | |
| length, header = self.generate_multipart(ranges, boundary, file_size, self.headers['content-type']) | |
| headers = MutableHeaders(raw=list(self.raw_headers)) | |
| headers['content-type'] = f'multipart/byteranges; boundary={boundary}' | |
| headers['content-length'] = str(length) | |
| await send({'type': 'http.response.start', 'status': 206, 'headers': headers.raw}) | |
| if send_header_only: | |
| await send({'type': 'http.response.body', 'body': b'', 'more_body': False}) | |
| return | |
| for start, end in ranges: | |
| await send({'type': 'http.response.body', 'body': header(start, end), 'more_body': True}) | |
| await self._send_part(send, start, end, final=False) | |
| await send({'type': 'http.response.body', 'body': b'\r\n', 'more_body': True}) | |
| await send({'type': 'http.response.body', 'body': f'--{boundary}--'.encode('latin-1'), 'more_body': False}) | |