comfyui custom node developing v2

 add one more class 

save image to png 


# from .imagetransfer import ImageTransfer

from PIL import Image
import torch
from torchvision.transforms.functional import to_pil_image, to_tensor
import os
import folder_paths
from comfy.cli_args import args
from PIL.PngImagePlugin import PngInfo
import json
import numpy as np

class ImageTransfer:
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),

FUNCTION = "transfer"
CATEGORY = "Utility"

def transfer(self, image):
# Ensure we're not tracking gradients for this operation
print(f"Type of the image tensor: {image.dtype}")
print(f"Dimensions of the image tensor: {image.shape}")

with torch.no_grad():
# Check if image tensor is in the expected format [1, H, W, C]
if image.dim() == 4 and image.shape[-1] in [1, 3]:
# Permute the tensor to match [C, H, W] format expected by to_pil_image
image_permuted = image.squeeze(0).permute(2, 0, 1)
# Convert the PyTorch tensor to a PIL Image
image_pil = to_pil_image(image_permuted)
# Resize the PIL Image
resized_image_pil = image_pil.resize((256, 256), Image.LANCZOS)
# Convert the PIL Image back to a PyTorch tensor
resized_image_tensor = to_tensor(resized_image_pil)
# Permute dimensions back to [1, H, W, C] and add the batch dimension
resized_image_tensor = resized_image_tensor.permute(1, 2, 0).unsqueeze(0)
print(f"Type of the resized image tensor: {resized_image_tensor.dtype}")
print(f"Dimensions of the resized image tensor: {resized_image_tensor.shape}")

return (resized_image_tensor,)
raise ValueError("Image tensor format not recognized. Expected format: [1, H, W, C].")

class Save_png2svg:
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
self.type = "output"
self.prefix_append = ""
self.compress_level = 4

return {"required":
{"images": ("IMAGE", ),
"filename_prefix": ("STRING", {"default": "ComfyUI"})},
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},

FUNCTION = "Save_png2svg"
CATEGORY = "image"

def Save_png2svg(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
filename_prefix += self.prefix_append
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
results = list()
for (batch_number, image) in enumerate(images):
i = 255. * image.cpu().numpy()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
metadata = None
if not args.disable_metadata:
metadata = PngInfo()
if prompt is not None:
metadata.add_text("prompt", json.dumps(prompt))
if extra_pnginfo is not None:
for x in extra_pnginfo:
metadata.add_text(x, json.dumps(extra_pnginfo[x]))

filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
file = f"{filename_with_batch_num}_{counter:05}_.png"
img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level)
"filename": file,
"subfolder": subfolder,
"type": self.type
counter += 1

return { "ui": { "images": results } }
"ImageTransfer": ImageTransfer,

"ImageTransfer": "Image Transfer",
"Save_png2svg": "Save png2svg"


No comments:

Post a Comment