mirror of
https://github.com/justUmen/Bjornulf_custom_nodes.git
synced 2026-03-21 12:42:11 -03:00
48 lines
1.3 KiB
Python
48 lines
1.3 KiB
Python
import os
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
class SaveApiImage:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", {"forceInput": True}),
|
|
}
|
|
}
|
|
|
|
FUNCTION = "save_api_image"
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "Bjornulf"
|
|
|
|
def save_api_image(self, image):
|
|
# Ensure the output directory exists
|
|
os.makedirs("./output/", exist_ok=True)
|
|
|
|
# Convert the image from ComfyUI format to PIL Image
|
|
i = 255. * image.cpu().numpy()
|
|
if i.ndim > 3:
|
|
i = np.squeeze(i)
|
|
if i.ndim == 2:
|
|
i = i[:, :, np.newaxis]
|
|
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
|
|
# Determine the next available filename
|
|
counter = 1
|
|
while True:
|
|
filename = f"./output/api_{counter:05d}.png"
|
|
if not os.path.exists(filename):
|
|
break
|
|
counter += 1
|
|
|
|
# Save the image with the determined filename
|
|
img.save(filename, format="PNG")
|
|
|
|
# Write the number of the last image to a text file with leading zeroes
|
|
with open("./output/api_next_image.txt", "w") as f:
|
|
f.write(f"api_{counter+1:05d}.png")
|
|
|
|
return ()
|