From ccb2b14dfa5f3632a6e35b5da91c35e27b5f139e Mon Sep 17 00:00:00 2001 From: Sebastian Monroy Date: Tue, 2 Sep 2025 14:53:28 +0100 Subject: [PATCH] updated MediaStreamOutput so that each instance of the node ONLY sends the dictionary related to its own output, rather than sending the entire final_outputs_dict The purpose of this code is to notify the backend that a specific output file has been successfully generated and uploaded. The backend (ComfyUIContentHandler) needs to know which output file this message corresponds to. The original code sent the entire final_outputs_dict. This would work, but it's inefficient and sends redundant information. If a workflow has five MediaStreamOutput nodes, each one would send a completion message containing the information for all five outputs. The backend would receive five identical messages. The new code is more precise. It filters the dictionary to include only the key-value pair for the output it just handled. This is a much cleaner and more correct approach. It ensures that each completion message is atomic and only contains the information relevant to the event that triggered it. --- media_stream.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/media_stream.py b/media_stream.py index 33fb4ba..f87299f 100644 --- a/media_stream.py +++ b/media_stream.py @@ -211,11 +211,20 @@ class MediaStreamOutput: elif format == "mp4": self._upload_video(images, presigned_upload_url, framerate) - # After upload, send the full, parsed dictionary of outputs to the SQS queue. + # This node is responsible for a single output. We find its corresponding object key. + output_key_for_this_node = final_outputs_dict.get(output_name) + if not output_key_for_this_node: + logging.error(f"FATAL: Could not find object key for output name '{output_name}' in output_object_keys.") + # Send an empty dictionary to signal failure. + final_outputs_for_sqs = {} + else: + final_outputs_for_sqs = {output_name: output_key_for_this_node} + + # After upload, send the filtered dictionary of outputs to the SQS queue. completion_message = { "job_id": job_id, "status": "completed", - "outputs": final_outputs_dict + "outputs": final_outputs_for_sqs } try: