comfyui custom node developing v2

 add one more class 

save image to png 

.

# from .imagetransfer import ImageTransfer

# __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

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:

@classmethod

def INPUT_TYPES(cls):

return {

"required": {

"image": ("IMAGE",),

},

}

RETURN_TYPES = ("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,)

else:

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

@classmethod

def INPUT_TYPES(s):

return {"required":

{"images": ("IMAGE", ),

"filename_prefix": ("STRING", {"default": "ComfyUI"})},

"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},

}

RETURN_TYPES = ()

FUNCTION = "Save_png2svg"

OUTPUT_NODE = True

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)

results.append({

"filename": file,

"subfolder": subfolder,

"type": self.type

})

counter += 1

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

NODE_CLASS_MAPPINGS = {

"ImageTransfer": ImageTransfer,

"Save_png2svg":Save_png2svg

}

NODE_DISPLAY_NAME_MAPPINGS = {

"ImageTransfer": "Image Transfer",

"Save_png2svg": "Save png2svg"

}

..