From e804f4bc6f47c6f65f2eb7b8e971c79b6652fa64 Mon Sep 17 00:00:00 2001 From: Michal Vavak Date: Mon, 28 Sep 2026 05:50:36 +0200 Subject: [PATCH 1/4] feat: return native and VideoHelperSuite video outputs --- .changeset/tangy-steaks-sneeze.md | 5 + handler.py | 127 +++++++++++++++++++---- tests/test_handler.py | 162 ++++++++++++++++++++++++++++++ 3 files changed, 276 insertions(+), 18 deletions(-) create mode 100644 .changeset/tangy-steaks-sneeze.md diff --git a/.changeset/tangy-steaks-sneeze.md b/.changeset/tangy-steaks-sneeze.md new file mode 100644 index 000000000..d499e20cb --- /dev/null +++ b/.changeset/tangy-steaks-sneeze.md @@ -0,0 +1,5 @@ +--- +"worker-comfyui": minor +--- + +return native and VideoHelperSuite video outputs in job results diff --git a/handler.py b/handler.py index ec80043dc..046a615d4 100644 --- a/handler.py +++ b/handler.py @@ -1,6 +1,7 @@ import runpod from runpod.serverless.utils import rp_upload import json +import mimetypes import urllib.request import urllib.parse import time @@ -770,38 +771,38 @@ def get_history(prompt_id): return response.json() -def get_image_data(filename, subfolder, image_type): +def get_output_data(filename, subfolder, output_type): """ - Fetch image bytes from the ComfyUI /view endpoint. + Fetch output bytes from the ComfyUI /view endpoint. Args: - filename (str): The filename of the image. - subfolder (str): The subfolder where the image is stored. - image_type (str): The type of the image (e.g., 'output'). + filename (str): The output filename. + subfolder (str): The subfolder where the output is stored. + output_type (str): The type of the output (e.g., 'output'). Returns: - bytes: The raw image data, or None if an error occurs. + bytes: The raw output data, or None if an error occurs. """ print( - f"worker-comfyui - Fetching image data: type={image_type}, subfolder={subfolder}, filename={filename}" + f"worker-comfyui - Fetching output data: type={output_type}, subfolder={subfolder}, filename={filename}" ) - data = {"filename": filename, "subfolder": subfolder, "type": image_type} + data = {"filename": filename, "subfolder": subfolder, "type": output_type} url_values = urllib.parse.urlencode(data) try: # Use requests for consistency and timeout response = requests.get(f"http://{COMFY_HOST}/view?{url_values}", timeout=60) response.raise_for_status() - print(f"worker-comfyui - Successfully fetched image data for {filename}") + print(f"worker-comfyui - Successfully fetched output data for {filename}") return response.content except requests.Timeout: - print(f"worker-comfyui - Timeout fetching image data for {filename}") + print(f"worker-comfyui - Timeout fetching output data for {filename}") return None except requests.RequestException as e: - print(f"worker-comfyui - Error fetching image data for {filename}: {e}") + print(f"worker-comfyui - Error fetching output data for {filename}: {e}") return None except Exception as e: print( - f"worker-comfyui - Unexpected error fetching image data for {filename}: {e}" + f"worker-comfyui - Unexpected error fetching output data for {filename}: {e}" ) return None @@ -867,6 +868,7 @@ def handler(job): client_id = str(uuid.uuid4()) prompt_id = None output_data = [] + output_data_videos = [] errors = [] try: @@ -1018,7 +1020,7 @@ def handler(job): errors.append(warn_msg) continue - image_bytes = get_image_data(filename, subfolder, img_type) + image_bytes = get_output_data(filename, subfolder, img_type) if image_bytes: file_extension = os.path.splitext(filename)[1] or ".png" @@ -1084,8 +1086,91 @@ def handler(job): error_msg = f"Failed to fetch image data for {filename} from /view endpoint." errors.append(error_msg) + for video_key in ("gifs", "videos"): + video_outputs = node_output.get(video_key, []) + if video_outputs: + print( + f"worker-comfyui - Node {node_id} contains {len(video_outputs)} video(s)" + ) + for video_info in video_outputs: + filename = video_info.get("filename") + subfolder = video_info.get("subfolder", "") + video_type = video_info.get("type", "output") + + if video_type == "temp": + print( + f"worker-comfyui - Skipping video {filename} because type is 'temp'" + ) + continue + + if not filename: + warn_msg = f"Skipping video in node {node_id} due to missing filename: {video_info}" + print(f"worker-comfyui - {warn_msg}") + errors.append(warn_msg) + continue + + video_bytes = get_output_data(filename, subfolder, video_type) + if not video_bytes: + errors.append( + f"Failed to fetch video data for {filename} from /view endpoint." + ) + continue + + if os.environ.get("BUCKET_ENDPOINT_URL"): + temp_video_path = None + try: + file_extension = os.path.splitext(filename)[1] or ".mp4" + with tempfile.NamedTemporaryFile( + suffix=file_extension, + delete=False, + ) as temp_file: + temp_file.write(video_bytes) + temp_video_path = temp_file.name + + content_type, _ = mimetypes.guess_type(filename) + extra_args = ( + {"ContentType": content_type} if content_type else None + ) + s3_url = rp_upload.upload_file_to_bucket( + file_name=f"{uuid.uuid4().hex}{file_extension}", + file_location=temp_video_path, + prefix=job_id, + extra_args=extra_args, + ) + output_data_videos.append( + { + "filename": filename, + "type": "s3_url", + "data": s3_url, + } + ) + except Exception as e: + errors.append(f"Error uploading {filename} to S3: {e}") + finally: + if temp_video_path and os.path.exists(temp_video_path): + try: + os.remove(temp_video_path) + except OSError as rm_err: + print( + f"worker-comfyui - Error removing temporary video file {temp_video_path}: {rm_err}" + ) + else: + try: + base64_video = base64.b64encode(video_bytes).decode("utf-8") + output_data_videos.append( + { + "filename": filename, + "type": "base64", + "data": base64_video, + } + ) + except Exception as e: + errors.append(f"Error encoding {filename} as base64: {e}") + # Check for other output types - other_keys = [k for k in node_output.keys() if k != "images"] + other_keys = [ + k for k in node_output.keys() if k not in ("images", "gifs", "videos") + ] if other_keys: warn_msg = ( f"Node {node_id} produced unhandled output keys: {other_keys}." @@ -1121,24 +1206,30 @@ def handler(job): if output_data: final_result["images"] = output_data + if output_data_videos: + final_result["videos"] = output_data_videos + final_result.setdefault("images", []) + if errors: final_result["errors"] = errors print(f"worker-comfyui - Job completed with errors/warnings: {errors}") - if not output_data and errors: - print(f"worker-comfyui - Job failed with no output images.") + if not output_data and not output_data_videos and errors: + print(f"worker-comfyui - Job failed with no output images or videos.") return { "error": "Job processing failed", "details": errors, } - elif not output_data and not errors: + elif not output_data and not output_data_videos and not errors: print( f"worker-comfyui - Job completed successfully, but the workflow produced no images." ) final_result["status"] = "success_no_images" final_result["images"] = [] - print(f"worker-comfyui - Job completed. Returning {len(output_data)} image(s).") + print( + f"worker-comfyui - Job completed. Returning {len(output_data)} image(s) and {len(output_data_videos)} video(s)." + ) return final_result diff --git a/tests/test_handler.py b/tests/test_handler.py index fcd0916ca..7ffd13343 100644 --- a/tests/test_handler.py +++ b/tests/test_handler.py @@ -4,6 +4,10 @@ import os import json import base64 +from pathlib import Path +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import threading +from urllib.parse import parse_qs, urlparse # handler.py lives at the repository root; it imports network_volume as a # sibling module (both are ADDed to / in the Docker image), which lives in @@ -420,9 +424,167 @@ def test_valid_workflow_reaches_queue_unchanged( result = handler.handler(job) self.assertNotIn("error", result) + self.assertEqual(result, {"status": "success_no_images", "images": []}) mock_queue.assert_called_once() self.assertEqual(mock_queue.call_args[0][0], workflow) +class TestHandlerVideoOutputs(unittest.TestCase): + def _run_handler(self, node_output, use_s3=False): + uploaded_files = [] + view_requests = [] + + class ViewRequestHandler(BaseHTTPRequestHandler): + def do_GET(self): + query = parse_qs(urlparse(self.path).query, keep_blank_values=True) + view_requests.append(query) + body = f"{query['filename'][0]} bytes".encode() + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, _format, *_args): + pass + + view_server = ThreadingHTTPServer(("127.0.0.1", 0), ViewRequestHandler) + view_thread = threading.Thread(target=view_server.serve_forever, daemon=True) + + def upload_file(file_name, file_location, prefix, extra_args): + uploaded_files.append( + ( + file_name, + file_location, + Path(file_location).read_bytes(), + prefix, + extra_args, + ) + ) + return "https://bucket.example/video" + + bucket_endpoint = "https://s3.example" if use_s3 else "" + with patch.dict(os.environ, {"BUCKET_ENDPOINT_URL": bucket_endpoint}): + with ( + patch.object( + handler, "COMFY_HOST", f"127.0.0.1:{view_server.server_port}" + ), + patch("handler.validate_workflow_models", return_value=None), + patch("handler.check_server", return_value=True), + patch("handler.queue_workflow", return_value={"prompt_id": "abc"}), + patch("handler.websocket.WebSocket") as mock_ws_class, + patch( + "handler.get_history", + return_value={"abc": {"outputs": {"9": node_output}}}, + ), + patch( + "handler.rp_upload.upload_file_to_bucket", side_effect=upload_file + ) as mock_upload, + ): + mock_ws = MagicMock() + mock_ws.recv.return_value = json.dumps( + {"type": "executing", "data": {"node": None, "prompt_id": "abc"}} + ) + mock_ws_class.return_value = mock_ws + + job = { + "id": "job-video", + "input": {"workflow": {"1": {"class_type": "KSampler", "inputs": {}}}}, + } + view_thread.start() + try: + result = handler.handler(job) + finally: + view_server.shutdown() + view_server.server_close() + view_thread.join() + + return result, mock_upload, uploaded_files, view_requests + + def test_returns_videohelper_and_native_video_outputs_as_base64(self): + node_output = { + "images": [ + {"filename": "still.png", "subfolder": "", "type": "output"} + ], + "gifs": [ + {"filename": "clip.mp4", "subfolder": "renders", "type": "output"} + ], + "videos": [ + {"filename": "clip.webm", "subfolder": "", "type": "output"} + ], + } + + result, mock_upload, _, view_requests = self._run_handler(node_output) + + self.assertEqual( + result["images"], + [ + { + "filename": "still.png", + "type": "base64", + "data": base64.b64encode(b"still.png bytes").decode("utf-8"), + } + ], + ) + self.assertEqual( + result["videos"], + [ + { + "filename": "clip.mp4", + "type": "base64", + "data": base64.b64encode(b"clip.mp4 bytes").decode("utf-8"), + }, + { + "filename": "clip.webm", + "type": "base64", + "data": base64.b64encode(b"clip.webm bytes").decode("utf-8"), + }, + ], + ) + self.assertIn( + {"filename": ["clip.mp4"], "subfolder": ["renders"], "type": ["output"]}, + view_requests, + ) + self.assertIn( + {"filename": ["clip.webm"], "subfolder": [""], "type": ["output"]}, + view_requests, + ) + mock_upload.assert_not_called() + + def test_uploads_video_to_s3_and_cleans_temporary_file(self): + node_output = { + "videos": [ + {"filename": "clip.mp4", "subfolder": "renders", "type": "output"} + ] + } + + result, mock_upload, uploaded_files, view_requests = self._run_handler( + node_output, use_s3=True + ) + + self.assertEqual( + result, + { + "images": [], + "videos": [ + { + "filename": "clip.mp4", + "type": "s3_url", + "data": "https://bucket.example/video", + } + ], + }, + ) + mock_upload.assert_called_once() + self.assertTrue(uploaded_files[0][0].endswith(".mp4")) + self.assertEqual(uploaded_files[0][2], b"clip.mp4 bytes") + self.assertEqual(uploaded_files[0][3], "job-video") + self.assertEqual(uploaded_files[0][4], {"ContentType": "video/mp4"}) + self.assertFalse(os.path.exists(uploaded_files[0][1])) + self.assertIn( + {"filename": ["clip.mp4"], "subfolder": ["renders"], "type": ["output"]}, + view_requests, + ) + + if __name__ == "__main__": unittest.main() From cbd5f612f67c2c62baa4e052302dd4a3fdbab51d Mon Sep 17 00:00:00 2001 From: Michal Vavak Date: Mon, 28 Sep 2026 05:54:35 +0200 Subject: [PATCH 2/4] fix: respect bucket prefix for video uploads --- handler.py | 8 ++++++-- tests/test_handler.py | 23 +++++++++++++++++++++-- 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/handler.py b/handler.py index 046a615d4..03daedf29 100644 --- a/handler.py +++ b/handler.py @@ -1131,10 +1131,14 @@ def handler(job): extra_args = ( {"ContentType": content_type} if content_type else None ) + bucket_prefix = os.environ.get("BUCKET_PREFIX", "").strip("/") + s3_prefix = ( + f"{bucket_prefix}/{job_id}" if bucket_prefix else job_id + ) s3_url = rp_upload.upload_file_to_bucket( - file_name=f"{uuid.uuid4().hex}{file_extension}", + file_name=os.path.basename(filename), file_location=temp_video_path, - prefix=job_id, + prefix=s3_prefix, extra_args=extra_args, ) output_data_videos.append( diff --git a/tests/test_handler.py b/tests/test_handler.py index 7ffd13343..a22cd9af9 100644 --- a/tests/test_handler.py +++ b/tests/test_handler.py @@ -430,7 +430,7 @@ def test_valid_workflow_reaches_queue_unchanged( class TestHandlerVideoOutputs(unittest.TestCase): - def _run_handler(self, node_output, use_s3=False): + def _run_handler(self, node_output, use_s3=False, bucket_prefix=""): uploaded_files = [] view_requests = [] @@ -463,7 +463,13 @@ def upload_file(file_name, file_location, prefix, extra_args): return "https://bucket.example/video" bucket_endpoint = "https://s3.example" if use_s3 else "" - with patch.dict(os.environ, {"BUCKET_ENDPOINT_URL": bucket_endpoint}): + with patch.dict( + os.environ, + { + "BUCKET_ENDPOINT_URL": bucket_endpoint, + "BUCKET_PREFIX": bucket_prefix, + }, + ): with ( patch.object( handler, "COMFY_HOST", f"127.0.0.1:{view_server.server_port}" @@ -585,6 +591,19 @@ def test_uploads_video_to_s3_and_cleans_temporary_file(self): view_requests, ) + def test_uses_bucket_prefix_for_video_uploads(self): + node_output = { + "videos": [ + {"filename": "clip.mp4", "subfolder": "renders", "type": "output"} + ] + } + + _, _, uploaded_files, _ = self._run_handler( + node_output, use_s3=True, bucket_prefix="/outputs/comfyui/" + ) + + self.assertEqual(uploaded_files[0][3], "outputs/comfyui/job-video") + if __name__ == "__main__": unittest.main() From 7648c69c95e468415c29a109c14a064fb8719cfc Mon Sep 17 00:00:00 2001 From: Michal Vavak Date: Mon, 28 Sep 2026 09:50:54 +0200 Subject: [PATCH 3/4] fix: classify native ComfyUI video outputs --- handler.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/handler.py b/handler.py index 03daedf29..331015cd9 100644 --- a/handler.py +++ b/handler.py @@ -998,6 +998,7 @@ def handler(job): print(f"worker-comfyui - Processing {len(outputs)} output nodes...") for node_id, node_output in outputs.items(): + native_video_outputs = [] if "images" in node_output: print( f"worker-comfyui - Node {node_id} contains {len(node_output['images'])} image(s)" @@ -1020,6 +1021,14 @@ def handler(job): errors.append(warn_msg) continue + if os.path.splitext(filename)[1].lower() in ( + ".mp4", + ".mkv", + ".webm", + ): + native_video_outputs.append(image_info) + continue + image_bytes = get_output_data(filename, subfolder, img_type) if image_bytes: @@ -1088,6 +1097,8 @@ def handler(job): for video_key in ("gifs", "videos"): video_outputs = node_output.get(video_key, []) + if video_key == "videos" and native_video_outputs: + video_outputs = [*video_outputs, *native_video_outputs] if video_outputs: print( f"worker-comfyui - Node {node_id} contains {len(video_outputs)} video(s)" From 45cd929b75c37fb54b66ba4362bf976f5ceddf5c Mon Sep 17 00:00:00 2001 From: Michal Vavak Date: Mon, 28 Sep 2026 09:51:02 +0200 Subject: [PATCH 4/4] test: cover native SaveVideo output classification --- tests/test_handler.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/test_handler.py b/tests/test_handler.py index a22cd9af9..1c7583810 100644 --- a/tests/test_handler.py +++ b/tests/test_handler.py @@ -509,7 +509,8 @@ def upload_file(file_name, file_location, prefix, extra_args): def test_returns_videohelper_and_native_video_outputs_as_base64(self): node_output = { "images": [ - {"filename": "still.png", "subfolder": "", "type": "output"} + {"filename": "still.png", "subfolder": "", "type": "output"}, + {"filename": "native.mp4", "subfolder": "", "type": "output"}, ], "gifs": [ {"filename": "clip.mp4", "subfolder": "renders", "type": "output"} @@ -544,6 +545,11 @@ def test_returns_videohelper_and_native_video_outputs_as_base64(self): "type": "base64", "data": base64.b64encode(b"clip.webm bytes").decode("utf-8"), }, + { + "filename": "native.mp4", + "type": "base64", + "data": base64.b64encode(b"native.mp4 bytes").decode("utf-8"), + }, ], ) self.assertIn( @@ -554,6 +560,10 @@ def test_returns_videohelper_and_native_video_outputs_as_base64(self): {"filename": ["clip.webm"], "subfolder": [""], "type": ["output"]}, view_requests, ) + self.assertIn( + {"filename": ["native.mp4"], "subfolder": [""], "type": ["output"]}, + view_requests, + ) mock_upload.assert_not_called() def test_uploads_video_to_s3_and_cleans_temporary_file(self):