Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 63 additions & 9 deletions handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -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:
Expand Down Expand Up @@ -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


Expand Down