diff --git a/handler.py b/handler.py index c9e84f32..0778903f 100644 --- a/handler.py +++ b/handler.py @@ -520,7 +520,8 @@ def handler(job): ws = None client_id = str(uuid.uuid4()) prompt_id = None - output_data = [] + output_data_images = [] + output_data_videos = [] errors = [] try: @@ -690,7 +691,7 @@ def handler(job): f"worker-comfyui - Uploaded {filename} to S3: {s3_url}" ) # Append dictionary with filename and URL - output_data.append( + output_data_images.append( { "filename": filename, "type": "s3_url", @@ -717,7 +718,7 @@ def handler(job): "utf-8" ) # Append dictionary with filename and base64 data - output_data.append( + output_data_images.append( { "filename": filename, "type": "base64", @@ -733,6 +734,54 @@ def handler(job): error_msg = f"Failed to fetch image data for {filename} from /view endpoint." errors.append(error_msg) + if "gifs" in node_output: + print( + f"worker-comfyui - Node {node_id} contains {len(node_output['gifs'])} video(s)" + ) + for video_info in node_output["gifs"]: + print(f"worker-comfyui - video_info: {video_info}") + filename = video_info.get("filename") + if video_info.get("type") == "temp": + print( + f"worker-comfyui - Skipping video {filename} because type is 'temp'" + ) + continue + if os.environ.get("BUCKET_ENDPOINT_URL"): + file_location = video_info.get("fullpath") + try: + print(f"worker-comfyui - Uploading {filename} to S3...") + s3_url = rp_upload.upload_file_to_bucket( + file_name=filename, + file_location=file_location, + prefix=f"comfy/{job_id}", + extra_args={"ContentType": f"video/{filename.split('.')[-1]}"} + ) + os.remove(file_location) # Clean up local video file + print( + f"worker-comfyui - Uploaded {filename} to S3: {s3_url}" + ) + # Append dictionary with filename and URL + output_data_videos.append( + { + "filename": filename, + "type": "s3_url", + "data": s3_url, + } + ) + except Exception as e: + error_msg = f"Error uploading {filename} to S3: {e}" + print(f"worker-comfyui - {error_msg}") + errors.append(error_msg) + if "file_location" in locals() and os.path.exists( + file_location + ): + try: + os.remove(file_location) + except OSError as rm_err: + print( + f"worker-comfyui - Error removing local video file {file_location}: {rm_err}" + ) + # Check for other output types other_keys = [k for k in node_output.keys() if k != "images"] if other_keys: @@ -766,28 +815,33 @@ def handler(job): ws.close() final_result = {} + output_data = output_data_images + output_data_videos - if output_data: - final_result["images"] = output_data + if output_data_images: + final_result["images"] = output_data_images + + if output_data_videos: + final_result["videos"] = output_data_videos 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.") + print(f"worker-comfyui - Job failed with no output images/videos.") return { "error": "Job processing failed", "details": errors, } elif not output_data and not errors: print( - f"worker-comfyui - Job completed successfully, but the workflow produced no images." + f"worker-comfyui - Job completed successfully, but the workflow produced no images/videos." ) - final_result["status"] = "success_no_images" + final_result["status"] = "success_no_images_or_videos" final_result["images"] = [] + final_result["videos"] = [] - print(f"worker-comfyui - Job completed. Returning {len(output_data)} image(s).") + print(f"worker-comfyui - Job completed. Returning {len(output_data)} image/video(s).") return final_result