From 9ada89ebe945a53be086c64bac96329e3fec470d Mon Sep 17 00:00:00 2001 From: Austin Mroz Date: Tue, 25 Jun 2024 12:29:52 -0500 Subject: [PATCH] Fix chunk_index for big files, allow dump reqs Added an environment variable to dump a prompt request and exit, for easier testing. Code for big files is executing, but I'm noticing but it produces output files that are not correct and has incorrectly inserted bytes. Will either disable or swap to different method of keeping the request but re-acquiring the connection semaphore --- workflowcheckpointing.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/workflowcheckpointing.py b/workflowcheckpointing.py index f730037..54db0b6 100644 --- a/workflowcheckpointing.py +++ b/workflowcheckpointing.py @@ -185,7 +185,7 @@ class FetchLoop: ci = url.find(',') chunk_index = int(url[:ci]) url = url[ci+1:] - headers.update({'Range': f'bytes={chunk_index*chunk_size}-{(chunk_index+1)*chunk_size}'}) + headers.update({'Range': f'bytes={chunk_index*chunk_size}-{(chunk_index+1)*chunk_size}-1'}) else: chunk_index = None filename = os.path.join('fetches', string_hash(url)) @@ -219,7 +219,7 @@ class FetchLoop: chunks = [] #TODO: Race condition where multiple model downloads initiate before priority drops for i in range(0, total_size, chunk_size): - chunks.append(self.enqueue(f'{i},{url}', priority-1)) + chunks.append(self.enqueue(f'{i//chunk_size},{url}', priority-1)) await asyncio.gather(*chunks) future.set_result(filename) fetch_loop = FetchLoop() @@ -344,6 +344,11 @@ prompt_route = next(filter(lambda x: x.path == '/prompt' and x.method == 'POST' server.PromptServer.instance.routes)) original_post_prompt = prompt_route.handler async def post_prompt_remote(request): + if 'dump_req' in os.environ: + with open('resp-dump.txt', 'wb') as f: + f.write(await request.read()) + import sys + sys.exit() json_data = await request.json() if "SALAD_ORGANIZATION" in os.environ: extra_data = json_data.get("extra_data", {})