2025-03-05 15:35:26 -05:00
from __future__ import annotations
2023-11-18 04:44:17 -05:00
import nodes
2023-11-23 13:55:29 -05:00
import folder_paths
from comfy . cli_args import args
from PIL import Image
2023-11-24 11:12:10 -05:00
from PIL . PngImagePlugin import PngInfo
2023-11-23 13:55:29 -05:00
import numpy as np
import json
import os
2025-05-09 10:46:34 -07:00
import re
from io import BytesIO
from inspect import cleandoc
2025-05-18 01:09:56 -07:00
import torch
2023-11-23 13:55:29 -05:00
2025-03-05 15:35:26 -05:00
from comfy . comfy_types import FileLocator
2023-11-18 04:44:17 -05:00
MAX_RESOLUTION = nodes . MAX_RESOLUTION
class ImageCrop :
@classmethod
def INPUT_TYPES ( s ) :
return { " required " : { " image " : ( " IMAGE " , ) ,
" width " : ( " INT " , { " default " : 512 , " min " : 1 , " max " : MAX_RESOLUTION , " step " : 1 } ) ,
" height " : ( " INT " , { " default " : 512 , " min " : 1 , " max " : MAX_RESOLUTION , " step " : 1 } ) ,
" x " : ( " INT " , { " default " : 0 , " min " : 0 , " max " : MAX_RESOLUTION , " step " : 1 } ) ,
" y " : ( " INT " , { " default " : 0 , " min " : 0 , " max " : MAX_RESOLUTION , " step " : 1 } ) ,
} }
RETURN_TYPES = ( " IMAGE " , )
FUNCTION = " crop "
CATEGORY = " image/transform "
def crop ( self , image , width , height , x , y ) :
x = min ( x , image . shape [ 2 ] - 1 )
y = min ( y , image . shape [ 1 ] - 1 )
to_x = width + x
to_y = height + y
img = image [ : , y : to_y , x : to_x , : ]
return ( img , )
2023-11-19 06:09:01 -05:00
class RepeatImageBatch :
@classmethod
def INPUT_TYPES ( s ) :
return { " required " : { " image " : ( " IMAGE " , ) ,
2024-03-17 08:57:49 -04:00
" amount " : ( " INT " , { " default " : 1 , " min " : 1 , " max " : 4096 } ) ,
2023-11-19 06:09:01 -05:00
} }
RETURN_TYPES = ( " IMAGE " , )
FUNCTION = " repeat "
CATEGORY = " image/batch "
def repeat ( self , image , amount ) :
s = image . repeat ( ( amount , 1 , 1 , 1 ) )
return ( s , )
2023-11-18 04:44:17 -05:00
2024-02-12 12:46:15 -05:00
class ImageFromBatch :
@classmethod
def INPUT_TYPES ( s ) :
return { " required " : { " image " : ( " IMAGE " , ) ,
2024-03-17 08:57:49 -04:00
" batch_index " : ( " INT " , { " default " : 0 , " min " : 0 , " max " : 4095 } ) ,
" length " : ( " INT " , { " default " : 1 , " min " : 1 , " max " : 4096 } ) ,
2024-02-12 12:46:15 -05:00
} }
RETURN_TYPES = ( " IMAGE " , )
FUNCTION = " frombatch "
CATEGORY = " image/batch "
def frombatch ( self , image , batch_index , length ) :
s_in = image
batch_index = min ( s_in . shape [ 0 ] - 1 , batch_index )
length = min ( s_in . shape [ 0 ] - batch_index , length )
s = s_in [ batch_index : batch_index + length ] . clone ( )
return ( s , )
2025-05-18 01:09:56 -07:00
class ImageAddNoise :
@classmethod
def INPUT_TYPES ( s ) :
return { " required " : { " image " : ( " IMAGE " , ) ,
" seed " : ( " INT " , { " default " : 0 , " min " : 0 , " max " : 0xffffffffffffffff , " control_after_generate " : True , " tooltip " : " The random seed used for creating the noise. " } ) ,
" strength " : ( " FLOAT " , { " default " : 0.5 , " min " : 0.0 , " max " : 1.0 , " step " : 0.01 } ) ,
} }
RETURN_TYPES = ( " IMAGE " , )
FUNCTION = " repeat "
CATEGORY = " image "
def repeat ( self , image , seed , strength ) :
generator = torch . manual_seed ( seed )
s = torch . clip ( ( image + strength * torch . randn ( image . size ( ) , generator = generator , device = " cpu " ) . to ( image ) ) , min = 0.0 , max = 1.0 )
return ( s , )
2023-11-23 13:55:29 -05:00
class SaveAnimatedWEBP :
def __init__ ( self ) :
self . output_dir = folder_paths . get_output_directory ( )
self . type = " output "
self . prefix_append = " "
methods = { " default " : 4 , " fastest " : 0 , " slowest " : 6 }
@classmethod
def INPUT_TYPES ( s ) :
return { " required " :
{ " images " : ( " IMAGE " , ) ,
" filename_prefix " : ( " STRING " , { " default " : " ComfyUI " } ) ,
" fps " : ( " FLOAT " , { " default " : 6.0 , " min " : 0.01 , " max " : 1000.0 , " step " : 0.01 } ) ,
" lossless " : ( " BOOLEAN " , { " default " : True } ) ,
" quality " : ( " INT " , { " default " : 80 , " min " : 0 , " max " : 100 } ) ,
" method " : ( list ( s . methods . keys ( ) ) , ) ,
# "num_frames": ("INT", {"default": 0, "min": 0, "max": 8192}),
} ,
" hidden " : { " prompt " : " PROMPT " , " extra_pnginfo " : " EXTRA_PNGINFO " } ,
}
RETURN_TYPES = ( )
FUNCTION = " save_images "
OUTPUT_NODE = True
2023-12-17 02:37:22 -05:00
CATEGORY = " image/animation "
2023-11-23 13:55:29 -05:00
def save_images ( self , images , fps , filename_prefix , lossless , quality , method , num_frames = 0 , prompt = None , extra_pnginfo = None ) :
2023-11-23 15:06:35 -05:00
method = self . methods . get ( method )
2023-11-23 13:55:29 -05:00
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 ] )
2025-03-05 15:35:26 -05:00
results : list [ FileLocator ] = [ ]
2023-11-23 13:55:29 -05:00
pil_images = [ ]
for image in images :
i = 255. * image . cpu ( ) . numpy ( )
img = Image . fromarray ( np . clip ( i , 0 , 255 ) . astype ( np . uint8 ) )
pil_images . append ( img )
2023-11-23 15:06:35 -05:00
metadata = pil_images [ 0 ] . getexif ( )
2023-11-23 13:55:29 -05:00
if not args . disable_metadata :
if prompt is not None :
metadata [ 0x0110 ] = " prompt: {} " . format ( json . dumps ( prompt ) )
if extra_pnginfo is not None :
inital_exif = 0x010f
for x in extra_pnginfo :
metadata [ inital_exif ] = " {} : {} " . format ( x , json . dumps ( extra_pnginfo [ x ] ) )
inital_exif - = 1
if num_frames == 0 :
num_frames = len ( pil_images )
c = len ( pil_images )
for i in range ( 0 , c , num_frames ) :
file = f " { filename } _ { counter : 05 } _.webp "
pil_images [ i ] . save ( os . path . join ( full_output_folder , file ) , save_all = True , duration = int ( 1000.0 / fps ) , append_images = pil_images [ i + 1 : i + num_frames ] , exif = metadata , lossless = lossless , quality = quality , method = method )
results . append ( {
" filename " : file ,
" subfolder " : subfolder ,
" type " : self . type
} )
counter + = 1
animated = num_frames != 1
return { " ui " : { " images " : results , " animated " : ( animated , ) } }
2023-11-24 11:12:10 -05:00
class SaveAnimatedPNG :
def __init__ ( self ) :
self . output_dir = folder_paths . get_output_directory ( )
self . type = " output "
self . prefix_append = " "
@classmethod
def INPUT_TYPES ( s ) :
return { " required " :
{ " images " : ( " IMAGE " , ) ,
" filename_prefix " : ( " STRING " , { " default " : " ComfyUI " } ) ,
2023-11-24 11:19:23 -05:00
" fps " : ( " FLOAT " , { " default " : 6.0 , " min " : 0.01 , " max " : 1000.0 , " step " : 0.01 } ) ,
2023-11-24 11:12:10 -05:00
" compress_level " : ( " INT " , { " default " : 4 , " min " : 0 , " max " : 9 } )
} ,
" hidden " : { " prompt " : " PROMPT " , " extra_pnginfo " : " EXTRA_PNGINFO " } ,
}
RETURN_TYPES = ( )
FUNCTION = " save_images "
OUTPUT_NODE = True
2023-12-17 02:37:22 -05:00
CATEGORY = " image/animation "
2023-11-24 11:12:10 -05:00
def save_images ( self , images , fps , compress_level , 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 ( )
pil_images = [ ]
for image in images :
i = 255. * image . cpu ( ) . numpy ( )
img = Image . fromarray ( np . clip ( i , 0 , 255 ) . astype ( np . uint8 ) )
pil_images . append ( img )
metadata = None
if not args . disable_metadata :
metadata = PngInfo ( )
if prompt is not None :
2023-11-24 18:24:19 -05:00
metadata . add ( b " comf " , " prompt " . encode ( " latin-1 " , " strict " ) + b " \0 " + json . dumps ( prompt ) . encode ( " latin-1 " , " strict " ) , after_idat = True )
2023-11-24 11:12:10 -05:00
if extra_pnginfo is not None :
for x in extra_pnginfo :
2023-11-24 18:24:19 -05:00
metadata . add ( b " comf " , x . encode ( " latin-1 " , " strict " ) + b " \0 " + json . dumps ( extra_pnginfo [ x ] ) . encode ( " latin-1 " , " strict " ) , after_idat = True )
2023-11-24 11:12:10 -05:00
file = f " { filename } _ { counter : 05 } _.png "
pil_images [ 0 ] . save ( os . path . join ( full_output_folder , file ) , pnginfo = metadata , compress_level = compress_level , save_all = True , duration = int ( 1000.0 / fps ) , append_images = pil_images [ 1 : ] )
results . append ( {
" filename " : file ,
" subfolder " : subfolder ,
" type " : self . type
} )
return { " ui " : { " images " : results , " animated " : ( True , ) } }
2025-05-09 10:46:34 -07:00
class SVG :
"""
Stores SVG representations via a list of BytesIO objects .
"""
def __init__ ( self , data : list [ BytesIO ] ) :
self . data = data
def combine ( self , other : ' SVG ' ) - > ' SVG ' :
return SVG ( self . data + other . data )
@staticmethod
def combine_all ( svgs : list [ ' SVG ' ] ) - > ' SVG ' :
all_svgs_list : list [ BytesIO ] = [ ]
for svg_item in svgs :
all_svgs_list . extend ( svg_item . data )
return SVG ( all_svgs_list )
class SaveSVGNode :
"""
Save SVG files on disk .
"""
def __init__ ( self ) :
self . output_dir = folder_paths . get_output_directory ( )
self . type = " output "
self . prefix_append = " "
RETURN_TYPES = ( )
DESCRIPTION = cleandoc ( __doc__ or " " ) # Handle potential None value
FUNCTION = " save_svg "
CATEGORY = " image/save " # Changed
OUTPUT_NODE = True
@classmethod
def INPUT_TYPES ( s ) :
return {
" required " : {
" svg " : ( " SVG " , ) , # Changed
" filename_prefix " : ( " STRING " , { " default " : " svg/ComfyUI " , " tooltip " : " The prefix for the file to save. This may include formatting information such as %d ate:yyyy-MM-dd % o r %E mpty Latent Image.width % to include values from nodes. " } )
} ,
" hidden " : {
" prompt " : " PROMPT " ,
" extra_pnginfo " : " EXTRA_PNGINFO "
}
}
def save_svg ( self , svg : SVG , filename_prefix = " svg/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 )
results = list ( )
# Prepare metadata JSON
metadata_dict = { }
if prompt is not None :
metadata_dict [ " prompt " ] = prompt
if extra_pnginfo is not None :
metadata_dict . update ( extra_pnginfo )
# Convert metadata to JSON string
metadata_json = json . dumps ( metadata_dict , indent = 2 ) if metadata_dict else None
for batch_number , svg_bytes in enumerate ( svg . data ) :
filename_with_batch_num = filename . replace ( " % batch_num % " , str ( batch_number ) )
file = f " { filename_with_batch_num } _ { counter : 05 } _.svg "
# Read SVG content
svg_bytes . seek ( 0 )
svg_content = svg_bytes . read ( ) . decode ( ' utf-8 ' )
# Inject metadata if available
if metadata_json :
# Create metadata element with CDATA section
metadata_element = f """ <metadata>
< ! [ CDATA [
{ metadata_json }
] ] >
< / metadata >
"""
# Insert metadata after opening svg tag using regex with a replacement function
def replacement ( match ) :
# match.group(1) contains the captured <svg> tag
return match . group ( 1 ) + ' \n ' + metadata_element
# Apply the substitution
svg_content = re . sub ( r ' (<svg[^>]*>) ' , replacement , svg_content , flags = re . UNICODE )
# Write the modified SVG to file
with open ( os . path . join ( full_output_folder , file ) , ' wb ' ) as svg_file :
svg_file . write ( svg_content . encode ( ' utf-8 ' ) )
results . append ( {
" filename " : file ,
" subfolder " : subfolder ,
" type " : self . type
} )
counter + = 1
return { " ui " : { " images " : results } }
2023-11-18 04:44:17 -05:00
NODE_CLASS_MAPPINGS = {
" ImageCrop " : ImageCrop ,
2023-11-19 06:09:01 -05:00
" RepeatImageBatch " : RepeatImageBatch ,
2024-02-12 12:46:15 -05:00
" ImageFromBatch " : ImageFromBatch ,
2025-05-18 01:09:56 -07:00
" ImageAddNoise " : ImageAddNoise ,
2023-11-23 13:55:29 -05:00
" SaveAnimatedWEBP " : SaveAnimatedWEBP ,
2023-11-24 11:12:10 -05:00
" SaveAnimatedPNG " : SaveAnimatedPNG ,
2025-05-09 10:46:34 -07:00
" SaveSVGNode " : SaveSVGNode ,
2023-11-18 04:44:17 -05:00
}