diff --git a/src/rp_handler.py b/src/rp_handler.py index 64779a8b..9b0b291b 100644 --- a/src/rp_handler.py +++ b/src/rp_handler.py @@ -250,7 +250,7 @@ def process_output_images(outputs, job_id): for image in node_output["images"]: output_images = os.path.join(image["subfolder"], image["filename"]) - print(f"runpod-worker-comfy - image generation is done") + print(f"runpod-worker-comfy - image generation is done - single image method") # expected image output folder local_image_path = f"{COMFY_OUTPUT_PATH}/{output_images}" @@ -309,7 +309,7 @@ def calculate_size_of_images(images): import base64 size_of_images = 0 for image in images: - if(images is not None): + if(image is not None): #img = image @@ -346,7 +346,11 @@ def process_output_images_list(outputs, job_id, save_to_storage=None): errors = [] images = [] + print("List of Images method - outputs:") + + for node_id, node_output in outputs.items(): + print(node_output) if "images" in node_output: images.append(node_output["images"]) @@ -363,7 +367,7 @@ def process_output_images_list(outputs, job_id, save_to_storage=None): COMFY_OUTPUT_PATH, image["subfolder"], image["filename"] ) - print(f"runpod-worker-comfy - image generation is done") + print(f"runpod-worker-comfy - image generation is done - list of images method") if os.path.exists(local_image_path): # Determine the return type: S3 URL or Base64 string @@ -385,11 +389,11 @@ def process_output_images_list(outputs, job_id, save_to_storage=None): image_list.append(image_json) # Save to RunPod storage if requested - if save_to_storage: - storage_path = os.path.join(save_to_storage, image["filename"]) - Path(save_to_storage).mkdir(parents=True, exist_ok=True) - shutil.copy2(local_image_path, storage_path) - print(f"Saved image to RunPod storage: {storage_path}") + #if save_to_storage: + #storage_path = os.path.join(save_to_storage, image["filename"]) + #Path(save_to_storage).mkdir(parents=True, exist_ok=True) + #shutil.copy2(local_image_path, storage_path) + #print(f"Saved image to RunPod storage: {storage_path}") else: # Log the error if the file doesn't exist @@ -476,7 +480,7 @@ def handler(job): try: # Get the generated image and return it as URL in an AWS bucket or as base64 - images_result = process_output_images_list(history[prompt_id].get("outputs"), job["id"], save_to_storage=True) + images_result = process_output_images_list(history[prompt_id].get("outputs"), job["id"]) except Exception as e: print(f"Error while saving list of images: {str(e)}") # Get the generated image and return it as URL in an AWS bucket or as base64