1movmedia/pyworker
2
0
Fork 0
Code Pull Requests Activity

Handle errors from model for comfyui-json

Browse Source
This commit is contained in:
Lucas Armand
2025-10-08 12:00:45 -07:00
parent 0397af719d
commit 7d3be849d9
1 changed files with 32 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
workers/comfyui-json/server.py
+32 -26
View File
@@ -33,33 +33,39 @@ log = logging.getLogger(__file__)
async def generate_client_response( async def generate_client_response(
client_request: web.Request, model_response: ClientResponse client_request: web.Request, model_response: ClientResponse
) -> Union[web.Response, web.StreamResponse]: ) -> Union[web.Response, web.StreamResponse]:
# Check if the response is actually streaming based on response headers/content-type match model_response.status:
is_streaming_response = ( case 200:
model_response.content_type == "text/event-stream" log.debug("SUCCESS")
or model_response.content_type == "application/x-ndjson" # Check if the response is actually streaming based on response headers/content-type
or model_response.headers.get("Transfer-Encoding") == "chunked" is_streaming_response = (
or "stream" in model_response.content_type.lower() model_response.content_type == "text/event-stream"
) or model_response.content_type == "application/x-ndjson"
or model_response.headers.get("Transfer-Encoding") == "chunked"
or "stream" in model_response.content_type.lower()
)
if is_streaming_response: if is_streaming_response:
log.debug("Detected streaming response...") log.debug("Detected streaming response...")
res = web.StreamResponse() res = web.StreamResponse()
res.content_type = model_response.content_type res.content_type = model_response.content_type
await res.prepare(client_request) await res.prepare(client_request)
async for chunk in model_response.content: async for chunk in model_response.content:
await res.write(chunk) await res.write(chunk)
await res.write_eof() await res.write_eof()
log.debug("Done streaming response") log.debug("Done streaming response")
return res return res
else: else:
log.debug("Detected non-streaming response...") log.debug("Detected non-streaming response...")
content = await model_response.read() content = await model_response.read()
return web.Response( return web.Response(
body=content, body=content,
status=model_response.status, status=model_response.status,
content_type=model_response.content_type content_type=model_response.content_type
) )
case code:
log.debug(f"Model responded with error {code}")
return web.Response(status=code)
@dataclasses.dataclass @dataclasses.dataclass
class ComfyWorkflowHandler(EndpointHandler[ComfyWorkflowData]): class ComfyWorkflowHandler(EndpointHandler[ComfyWorkflowData]):