import os import model_management import torch import comfy.sd import comfy.utils import folder_paths import comfy.samplers from nodes import common_ksampler from comfy_extras.chainner_models import model_loading from PIL import Image, ImageOps, ImageFilter, ImageDraw from PIL.PngImagePlugin import PngInfo import numpy as np from torchvision.transforms import ToPILImage import cv2 from deepface import DeepFace import re import latent_preview from datetime import datetime import json import re import piexif import piexif.helper MAX_RESOLUTION=8192 # Tensor to PIL def tensor2pil(image): return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) # Convert PIL to Tensor def pil2tensor(image): return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) # model io class WLSH_Checkpoint_Loader_Model_Name: @classmethod def INPUT_TYPES(s): return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), }} RETURN_TYPES = ("MODEL", "CLIP", "VAE","STRING",) FUNCTION = "load_checkpoint" CATEGORY = "WLSH Nodes/loaders" def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True): ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) name = self.parse_name(ckpt_name) out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) new_out = list(out) new_out.pop() new_out.append(name) out = tuple(new_out) return (out) def parse_name(self, ckpt_name): path = ckpt_name filename = path.split("/")[-1] filename = filename.split(".")[:-1] filename = ".".join(filename) return filename # sampling class WLSH_KSamplerAdvancedMod: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "add_noise": (["enable", "disable"], ), "seed": ("SEED", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}), "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ), "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), "positive": ("CONDITIONING", ), "negative": ("CONDITIONING", ), "latent_image": ("LATENT", ), "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}), "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}), "return_with_leftover_noise": (["disable", "enable"], ), "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), } } RETURN_TYPES = ("LATENT",) FUNCTION = "sample" CATEGORY = "WLSH Nodes/sampling" def sample(self, model, add_noise, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise): noise_seed = seed['seed'] force_full_denoise = False if return_with_leftover_noise == "enable": force_full_denoise = False disable_noise = False if add_noise == "disable": disable_noise = True return common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise) class WLSH_Alternating_KSamplerAdvanced: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "add_noise": (["enable", "disable"], ), "seed": ("SEED", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}), "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ), "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), "clip": ("CLIP", ), "positive_prompt": ("STRING", {"forceInput": True }), "negative_prompt": ("STRING", {"forceInput": True }), "latent_image": ("LATENT", ), "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}), "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}), "return_with_leftover_noise": (["disable", "enable"], ), "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), } } RETURN_TYPES = ("LATENT",) FUNCTION = "sample" CATEGORY = "WLSH Nodes/sampling" def sample(self, model, add_noise, seed, steps, cfg, sampler_name, scheduler, clip, positive_prompt, negative_prompt, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise): noise_seed = seed['seed'] force_full_denoise = False if return_with_leftover_noise == "enable": force_full_denoise = False disable_noise = False if add_noise == "disable": disable_noise = True # alternating prompt parser # syntax: {A|B} will sequentially alternate between A and B def parse_prompt(input_string, stepnum): def replace_match(match): options = match.group(1).split('|') return options[(stepnum - 1) % len(options)] pattern = r'<(.*?)>' parsed_string = re.sub(pattern, replace_match, input_string) return parsed_string latent_input = latent_image for step in range(0,steps): positive_txt = parse_prompt(positive_prompt,step+1) positive = [[clip.encode(positive_txt), {}]] negative_txt = parse_prompt(negative_prompt,step+1) negative = [[clip.encode(negative_txt), {}]] if(step < steps): force_full_denoise = True if(step > 0): # disable_noise=True denoise=(steps-step)/(steps) latent_image = common_ksampler(model, noise_seed, 1, cfg, sampler_name, scheduler, positive, negative, latent_input, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise) latent_input = latent_image[0] return latent_image # return alternating_ksampler(clip, model, noise_seed, steps, cfg, sampler_name, scheduler, positive_prompt, negative_prompt, latent_image, # denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise) # utilities class WLSH_Seed_to_Number: def __init__(self): pass @classmethod def INPUT_TYPES(cls): return { "required": { "seed": ("SEED",), } } RETURN_TYPES = ("INT",) FUNCTION = "number_to_seed" CATEGORY = "WLSH Nodes" def number_to_seed(self, seed): return (int(seed["seed"]), ) class WLSH_SDXL_Steps: def __init__(self): pass @classmethod def INPUT_TYPES(cls): return { "required": { "precondition": ("INT", {"default": 3, "min": 1, "max": 10000}), "base": ("INT", {"default": 12, "min": 1, "max": 10000}), "total": ("INT", {"default": 20, "min": 1, "max": 10000}), } } RETURN_TYPES = ("INT","INT","INT",) FUNCTION = "set_steps" CATEGORY="WLSH Nodes" def set_steps(self,precondition,base,total): return(precondition,base,total) class WLSH_Int_Multiply: def __init__(self): pass @classmethod def INPUT_TYPES(cls): return { "required": { "number": ("INT",{"default": 2, "min": 1, "max": 10000}), "multiplier": ("INT", {"default": 2, "min": 1, "max": 10000}), } } RETURN_TYPES = ("INT",) FUNCTION = "multiply" CATEGORY="WLSH Nodes/number" def multiply(self,number,multiplier): result = number*multiplier return (int(result),) class WLSH_Time_String: time_format = ["%Y%m%d%H%M%S","%Y%m%d%H%M","%Y%m%d","%Y-%m-%d-%H%M%S", "%Y-%m-%d-%H%M", "%Y-%m-%d"] def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "style": (s.time_format,), } } RETURN_TYPES = ("STRING",) FUNCTION = "get_time" CATEGORY = "WLSH Nodes/text" def get_time(self, style): now = datetime.now() timestamp = now.strftime(style) return (timestamp,) class WLSH_SDXL_Resolutions: resolution = ["1024x1024","1152x896","1216x832","1344x768","1536x640"] direction = ["landscape","portrait"] def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "resolution": (s.resolution,), "direction": (s.direction,), } } RETURN_TYPES = ("INT","INT",) FUNCTION = "get_resolutions" CATEGORY="WLSH Nodes" def get_resolutions(self,resolution, direction): width,height = resolution.split('x') width = int(width) height = int(height) if(direction == "portrait"): width,height = height,width return(width,height) class WLSH_Resolutions_by_Ratio: aspects = ["1:1","5:4","4:3","3:2","16:10","16:9","21:9","2:1","3:1","4:1"] direction = ["landscape","portrait"] @classmethod def INPUT_TYPES(s): return {"required": { "aspect": (s.aspects,), "direction": (s.direction,), "shortside": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64})}} RETURN_TYPES = ("INT","INT",) FUNCTION = "get_resolutions" CATEGORY="WLSH Nodes" def get_resolutions(self, aspect, direction, shortside): x,y = aspect.split(':') x = int(x) y = int(y) ratio = x/y width = int(shortside * ratio) width = (width + 63) & (-64) height = shortside if(direction == "portrait"): width,height = height,width return(width,height) # latent class WLSH_Empty_Latent_Image_By_Ratio: aspects = ["1:1","5:4","4:3","3:2","16:10","16:9","21:9","2:1","3:1","4:1"] direction = ["landscape","portrait"] def __init__(self, device="cpu"): self.device = device @classmethod def INPUT_TYPES(s): return {"required": { "aspect": (s.aspects,), "direction": (s.direction,), "shortside": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}), "batch_size": ("INT", {"default": 1, "min": 1, "max": 64})}} RETURN_TYPES = ("LATENT",) FUNCTION = "generate" CATEGORY = "WLSH Nodes/latent" def generate(self, aspect, direction, shortside, batch_size=1): x,y = aspect.split(':') x = int(x) y = int(y) ratio = x/y width = int(shortside * ratio) width = (width + 63) & (-64) height = shortside if(direction == "portrait"): width,height = height,width latent = torch.zeros([batch_size, 4, height // 8, width // 8]) return ({"samples":latent}, ) class WLSH_SDXL_Quick_Empty_Latent: resolution = ["1024x1024","1152x896","1216x832","1344x768","1536x640"] direction = ["landscape","portrait"] def __init__(self, device="cpu"): self.device = device @classmethod def INPUT_TYPES(s): return {"required": { "resolution": (s.resolution,), "direction": (s.direction,), "batch_size": ("INT", {"default": 1, "min": 1, "max": 64})}} RETURN_TYPES = ("LATENT",) FUNCTION = "generate" CATEGORY = "WLSH Nodes/latent" def generate(self, resolution, direction, batch_size=1): width,height = resolution.split('x') width = int(width) height = int(height) if(direction == "portrait"): width,height = height,width latent = torch.zeros([batch_size, 4, height // 8, width // 8]) return ({"samples":latent}, ) # conditioning class WLSH_CLIP_Text_Positive_Negative: @classmethod def INPUT_TYPES(s): return {"required": {"positive": ("STRING", {"multiline": True}), "negative": ("STRING", {"multiline": True}), "clip": ("CLIP", )}} RETURN_TYPES = ("CONDITIONING","CONDITIONING","STRING","STRING") FUNCTION = "encode" CATEGORY = "WLSH Nodes/conditioning" def encode(self, clip, positive, negative): return ([[clip.encode(positive), {}]],[[clip.encode(negative), {}]],positive,negative) class WLSH_CLIP_Positive_Negative: @classmethod def INPUT_TYPES(s): return {"required": { "clip": ("CLIP", ), "positive_text": ("STRING",{"default": f'', "multiline": True}), "negative_text": ("STRING",{"default": f'', "multiline": True}) }} RETURN_TYPES = ("CONDITIONING","CONDITIONING",) FUNCTION = "encode" CATEGORY = "WLSH Nodes/conditioning" def encode(self, clip, positive_text, negative_text): return ([[clip.encode(positive_text), {}]],[[clip.encode(negative_text), {}]] ) # upscaling class WLSH_Image_Scale_By_Factor: upscale_methods = ["nearest-exact", "bilinear", "area"] @classmethod def INPUT_TYPES(s): return {"required": { "original": ("IMAGE",), "upscaled": ("IMAGE",), "upscale_method": (s.upscale_methods,), "factor": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 8.0, "step": 0.1}) }} RETURN_TYPES = ("IMAGE",) FUNCTION = "upscale" CATEGORY = "WLSH Nodes/upscaling" def upscale(self, original, upscaled, upscale_method, factor, crop): old_width = original.shape[2] old_height = original.shape[1] new_width= int(old_width * factor) new_height = int(old_height * factor) print("Processing image with shape: ",old_width,"x",old_height,"to ",new_width,"x",new_height) samples = upscaled.movedim(-1,1) s = comfy.utils.common_upscale(samples, new_width, new_height, upscale_method, crop="disabled") s = s.movedim(1,-1) return (s,) class WLSH_SDXL_Quick_Image_Scale: upscale_methods = ["nearest-exact", "bilinear", "area"] resolution = ["1024x1024","1152x896","1216x832","1344x768","1536x640"] direction = ["landscape","portrait"] crop_methods = ["disabled", "center"] @classmethod def INPUT_TYPES(s): return {"required": { "original": ("IMAGE",), "upscale_method": (s.upscale_methods,), "resolution": (s.resolution,), "direction": (s.direction,), "crop": (s.crop_methods,), }} RETURN_TYPES = ("IMAGE",) FUNCTION = "upscale" CATEGORY = "WLSH Nodes/upscaling" def upscale(self, original, upscale_method, resolution, direction, crop): width,height = resolution.split('x') new_width = int(width) new_height = int(height) if(direction == "portrait"): new_width,new_height = new_height,new_width old_width = original.shape[2] old_height = original.shape[1] #print("Processing image with shape: ",old_width,"x",old_height,"to ",new_width,"x",new_height) samples = original.movedim(-1,1) s = comfy.utils.common_upscale(samples, new_width, new_height, upscale_method, crop) s = s.movedim(1,-1) return (s,) class WLSH_Upscale_By_Factor_With_Model: upscale_methods = ["nearest-exact", "bilinear", "area"] @classmethod def INPUT_TYPES(s): return {"required": { "upscale_model": ("UPSCALE_MODEL",), "image": ("IMAGE",), "upscale_method": (s.upscale_methods,), "factor": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 8.0, "step": 0.1}) }} RETURN_TYPES = ("IMAGE",) FUNCTION = "upscale" CATEGORY = "WLSH Nodes/upscaling" def upscale(self, image, upscale_model, upscale_method, factor): # upscale image using upscaling model device = model_management.get_torch_device() upscale_model.to(device) in_img = image.movedim(-1,-3).to(device) s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=128 + 64, tile_y=128 + 64, overlap = 8, upscale_amount=upscale_model.scale) upscale_model.cpu() upscaled = torch.clamp(s.movedim(-3,-1), min=0, max=1.0) # get dimensions of orginal image old_width = image.shape[2] old_height = image.shape[1] # scale dimensions by provided factor new_width= int(old_width * factor) new_height = int(old_height * factor) print("Processing image with shape: ",old_width,"x",old_height,"to ",new_width,"x",new_height) # apply simple scaling to image samples = upscaled.movedim(-1,1) s = comfy.utils.common_upscale(samples, new_width, new_height, upscale_method, crop="disabled") s = s.movedim(1,-1) return (s,) # outpainting class WLSH_Outpaint_To_Image: directions = ["left","right","up","down"] @classmethod def INPUT_TYPES(s): return {"required": { "image": ("IMAGE",), "direction": (s.directions,), "pixels": ("INT", {"default": 128, "min": 32, "max": 512, "step": 32}), "mask_padding": ("INT",{"default": 12, "min": 0, "max": 64, "step": 4}) }} RETURN_TYPES = ("IMAGE","MASK") FUNCTION = "outpaint" CATEGORY = "WLSH Nodes/inpainting" def convert_image(self, im, direction, mask_padding): width, height = im.size im = im.convert("RGBA") alpha = Image.new('L',(width,height),255) im.putalpha(alpha) return im def outpaint(self, image, direction, mask_padding, pixels): image = tensor2pil(image) # i = 255. * image.cpu().numpy() # image = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) image = self.convert_image(image, direction, mask_padding) if direction == "right": border = (0,0,pixels,0) new_image = ImageOps.expand(image,border=border,fill=(0,0,0,0)) elif direction == "left": border = (pixels,0,0,0) new_image = ImageOps.expand(image,border=border,fill=(0,0,0,0)) elif direction == "up": border = (0,pixels,0,0) new_image = ImageOps.expand(image,border=border,fill=(0,0,0,0)) elif direction == "down": border = (0,0,0,pixels) new_image = ImageOps.expand(image,border=border,fill=(0,0,0,0)) image = new_image.convert("RGB") image = np.array(image).astype(np.float32) / 255.0 image = torch.from_numpy(image)[None,] if 'A' in new_image.getbands(): mask = np.array(new_image.getchannel('A')).astype(np.float32) / 255.0 mask = 1. - torch.from_numpy(mask) else: mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") #print("bands: ", new_image.getbands()) # if 'A' in new_image.getbands(): # mask = np.array(new_image.getchannel('A')).astype(np.float32) / 255.0 # mask = 1. - torch.from_numpy(mask) # #print("getting mask from alpha") # else: # mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") # # print("generating mask") # new_image = new_image.convert("RGB") # new_image = pil2tensor(new_image) return (image,mask) class WLSH_VAE_Encode_For_Inpaint_Padding: def __init__(self, device="cpu"): self.device = device @classmethod def INPUT_TYPES(s): return {"required": { "pixels": ("IMAGE", ), "vae": ("VAE", ), "mask": ("MASK", ), "mask_padding": ("INT",{"default": 24, "min": 6, "max": 128, "step": 2})}} RETURN_TYPES = ("LATENT",) FUNCTION = "encode" CATEGORY = "WLSH Nodes/inpainting" def encode(self, vae, pixels, mask, mask_padding=3): x = (pixels.shape[1] // 64) * 64 y = (pixels.shape[2] // 64) * 64 mask = torch.nn.functional.interpolate(mask[None,None,], size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")[0][0] pixels = pixels.clone() if pixels.shape[1] != x or pixels.shape[2] != y: pixels = pixels[:,:x,:y,:] mask = mask[:x,:y] #grow mask by a few pixels to keep things seamless in latent space kernel_tensor = torch.ones((1, 1, mask_padding, mask_padding)) mask_erosion = torch.clamp(torch.nn.functional.conv2d((mask.round())[None], kernel_tensor, padding=3), 0, 1) m = (1.0 - mask.round()) for i in range(3): pixels[:,:,:,i] -= 0.5 pixels[:,:,:,i] *= m pixels[:,:,:,i] += 0.5 t = vae.encode(pixels) return ({"samples":t, "noise_mask": (mask_erosion[0][:x,:y].round())}, ) class WLSH_Generate_Edge_Mask: directions = ["left","right","up","down"] @classmethod def INPUT_TYPES(s): return {"required": { "image": ("IMAGE",), "direction": (s.directions,), "pixels": ("INT", {"default": 128, "min": 32, "max": 512, "step": 32}), "overlap": ("INT", {"default": 64, "min": 16, "max": 256, "step": 16}) }} RETURN_TYPES = ("IMAGE",) FUNCTION = "gen_second_mask" CATEGORY = "WLSH Nodes/inpainting" def gen_second_mask(self, direction, image, pixels, overlap): image = tensor2pil(image) new_width,new_height = image.size # generate new image fully un-masked mask2 = Image.new('RGBA',(new_width,new_height),(0,0,0,255)) mask_thickness = overlap if (direction == "up"): # horizontal mask width of new image and height of 1/4 padding new_mask = Image.new('RGBA',(new_width, mask_thickness),(0,122,0,255)) mask2.paste(new_mask,(0,(pixels-int(mask_thickness/2)))) elif (direction == "down"): # horizontal mask width of new image and height of 1/4 padding new_mask = Image.new('RGBA',(new_width, mask_thickness),(0,122,0,255)) mask2.paste(new_mask,(0,new_height-pixels - int(mask_thickness/2))) elif (direction == "left"): # vertical mask height of new image and width of 1/4 padding new_mask = Image.new('RGBA',(mask_thickness,new_height),(0,122,0,255)) mask2.paste(new_mask,(pixels - int(mask_thickness/2),0)) elif (direction == "right"): # vertical mask height of new image and width of 1/4 padding new_mask = Image.new('RGBA',(mask_thickness,new_height),(0,122,0,255)) mask2.paste(new_mask,(new_width - pixels - int(mask_thickness/2),0)) mask2 = mask2.filter(ImageFilter.GaussianBlur(radius=5)) mask2 = np.array(mask2).astype(np.float32) / 255.0 mask2 = torch.from_numpy(mask2)[None,] return (mask2,) class WLSH_Generate_Face_Mask: detectors = ["opencv", "retinaface", "ssd", "mtcnn"] channels = ["red", "blue", "green"] @classmethod def INPUT_TYPES(s): return {"required": { "image": ("IMAGE",), "detector": (s.detectors,), "channel": (s.channels,), "mask_padding": ("INT",{"default": 6, "min": 0, "max": 32, "step": 2}) }} RETURN_TYPES = ("IMAGE",) FUNCTION = "gen_face_mask" CATEGORY = "WLSH Nodes/inpainting" def gen_face_mask(self, image, mask_padding, detector, channel): image = tensor2pil(image) faces = DeepFace.extract_faces(np.array(image),detector_backend=detector) # cv_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) # # Convert to grayscale # gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY) # # Detect faces in the image # face_cascade = cv2.CascadeClassifier('custom_nodes/haarcascade_frontalface_default.xml') # faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5) mask = Image.new('RGB',image.size) # Draw a rectangle on the PIL Image object colors = {"red": "RGB(255,0,0)", "green": "RGB(0,255,0)", "blue": "RGB(0,0,255)"} draw = ImageDraw.Draw(mask) for face in faces: x,y,w,h = face['facial_area'].values() draw.rectangle((x-mask_padding,y-mask_padding,x+w+mask_padding,y+h+mask_padding), outline=colors[channel], fill=colors[channel]) mask = mask.filter(ImageFilter.GaussianBlur(radius=6)) mask = np.array(mask).astype(np.float32) / 255.0 mask = torch.from_numpy(mask)[None,] return (mask,) # image I/O class WLSH_Image_Save_With_Prompt_Info: def __init__(self): self.output_dir = os.path.join(os.getcwd()+'/ComfyUI', "output") @classmethod def INPUT_TYPES(s): return { "required": { "images": ("IMAGE", ), "filename": ("STRING", {"default": f'%time_%seed', "multiline": False}), "output_path": ("STRING", {"default": './output', "multiline": False}), "extension": (['png', 'jpeg', 'tiff', 'gif'], ), "quality": ("INT", {"default": 100, "min": 1, "max": 100, "step": 1}), }, "optional": { "positive": ("STRING",{"default": '', "multiline": True}), "negative": ("STRING",{"default": '', "multiline": True}), "seed": ("SEED",), "modelname": ("STRING",{"default": '', "multiline": False}), }, "hidden": { "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO" }, } RETURN_TYPES = () FUNCTION = "save_files" OUTPUT_NODE = True CATEGORY = "WLSH Nodes/IO" def save_files(self, images, positive="uknown", negative="uknown", seed=0, modelname="uknown", filename=f'%time', output_path="./output", extension='png', quality=100, prompt=None, extra_pnginfo=None): filename = self.make_filename(seed, modelname, filename) comment = self.make_comment(positive, negative, modelname, seed) paths = self.save_images(images, output_path,filename,comment, extension, quality, prompt, extra_pnginfo) #return return { "ui": { "images": paths } } def make_comment(self, positive, negative, modelname, seed): comment = "Positive Prompt:\n" + positive + "\nNegative Prompt:\n" + negative + "\nModel: " + modelname + "\nSeed: " + str(seed['seed']) return comment def make_filename(self, seed, modelname, filename): # generate datetime string now = datetime.now() timestamp = now.strftime("%Y-%m-%d-%H%M%S") # parse input string filename = filename.replace("%time",timestamp) filename = filename.replace("%model",modelname) filename = filename.replace("%seed",str(seed['seed'])) return (filename) def save_images(self, images, output_path='', filename_prefix="ComfyUI", comment="", extension='png', quality=100, prompt=None, extra_pnginfo=None): def map_filename(filename): prefix_len = len(filename_prefix) prefix = filename[:prefix_len + 1] try: digits = int(filename[prefix_len + 1:].split('_')[0]) except: digits = 0 return (digits, prefix) # Setup custom path or default if output_path.strip() != '': if not os.path.exists(output_path.strip()): print(f'\033[34mWAS NS\033[0m Error: The path `{output_path.strip()}` specified doesn\'t exist! Defaulting to `{self.output_dir}` directory.') else: self.output_dir = os.path.normpath(output_path.strip()) imgCount = 1 paths = list() for image in images: i = 255. * image.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) 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])) prefix = filename_prefix if(images.size()[0] > 1): prefix = filename_prefix + "_{:02d}".format(imgCount) file = f"{prefix}.{extension}" if extension == 'png': img.save(os.path.join(self.output_dir, file), comment=comment, pnginfo=metadata, optimize=True) elif extension == 'webp': img.save(os.path.join(self.output_dir, file), quality=quality) elif extension == 'jpeg': img.save(os.path.join(self.output_dir, file), quality=quality, comment=comment, optimize=True) elif extension == 'tiff': img.save(os.path.join(self.output_dir, file), quality=quality, optimize=True) else: img.save(os.path.join(self.output_dir, file)) paths.append(file) imgCount += 1 return(paths) class WLSH_Image_Save_With_Prompt_File: def __init__(self): self.output_dir = os.path.join(os.getcwd()+'/ComfyUI', "output") @classmethod def INPUT_TYPES(s): return { "required": { "images": ("IMAGE", ), "filename": ("STRING", {"default": f'%time_%seed', "multiline": False}), "output_path": ("STRING", {"default": './output', "multiline": False}), "extension": (['png', 'jpeg', 'tiff', 'gif'], ), "quality": ("INT", {"default": 100, "min": 1, "max": 100, "step": 1}), }, "optional": { "positive": ("STRING",{"default": ' ', "multiline": True}), "negative": ("STRING",{"default": ' ', "multiline": True}), "seed": ("SEED",), "modelname": ("STRING",{"default": 'sd', "multiline": False}), }, "hidden": { "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO" }, } RETURN_TYPES = () FUNCTION = "save_files" OUTPUT_NODE = True CATEGORY = "WLSH Nodes/IO" def save_files(self, images, positive="unknown", negative="unknown", seed=0, modelname="unknown", filename=f'%time', output_path="./output", extension='png', quality=100, prompt=None, extra_pnginfo=None): filename = self.make_filename(seed, modelname, filename) comment = "Positive Prompt:\n" + positive + "\nNegative Prompt:\n" + negative + "\nModel: " + modelname + "\nSeed: " + str(seed['seed']) paths = self.save_images(images, output_path,filename,comment, extension, quality, prompt, extra_pnginfo) self.save_text_file(positive, negative, seed, modelname, output_path, filename) #return return { "ui": { "images": paths } } def make_filename(self, seed, modelname, filename): # generate datetime string now = datetime.now() timestamp = now.strftime("%Y-%m-%d-%H%M%S") # parse input string filename = filename.replace("%time",timestamp) filename = filename.replace("%model",modelname) filename = filename.replace("%seed",str(seed['seed'])) return (filename) def save_images(self, images, output_path='', filename_prefix="ComfyUI", comment="", extension='png', quality=100, prompt=None, extra_pnginfo=None): def map_filename(filename): prefix_len = len(filename_prefix) prefix = filename[:prefix_len + 1] try: digits = int(filename[prefix_len + 1:].split('_')[0]) except: digits = 0 return (digits, prefix) # Setup custom path or default if output_path.strip() != '': if not os.path.exists(output_path.strip()): print(f'\033[34mWAS NS\033[0m Error: The path `{output_path.strip()}` specified doesn\'t exist! Defaulting to `{self.output_dir}` directory.') else: self.output_dir = os.path.normpath(output_path.strip()) imgCount = 1 paths = list() for image in images: i = 255. * image.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) 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])) if(images.size()[0] > 1): filename_prefix += "_{:02d}".format(imgCount) file = f"{filename_prefix}.{extension}" if extension == 'png': img.save(os.path.join(self.output_dir, file), comment=comment, pnginfo=metadata, optimize=True) elif extension == 'webp': img.save(os.path.join(self.output_dir, file), quality=quality) elif extension == 'jpeg': img.save(os.path.join(self.output_dir, file), quality=quality, comment=comment, optimize=True) elif extension == 'tiff': img.save(os.path.join(self.output_dir, file), quality=quality, optimize=True) else: img.save(os.path.join(self.output_dir, file)) paths.append(file) imgCount += 1 return(paths) def save_text_file(self, positive, negative, seed, modelname, path, filename): # Ensure path exists if not os.path.exists(path): print(f'\033[34mWAS NS\033[0m Error: The path `{path}` doesn\'t exist!') # Ensure content to save if filename.strip == '': print(f'\033[34mWAS NS\033[0m Error: There is no text specified to save! Text is empty.') text = "Positive Prompt:\n" + positive + "\nNegative Prompt:\n" + negative + "\nModel: " + modelname + "\nSeed: " + str(seed['seed']) filename = self.make_filename(seed, modelname, filename) # Write text file self.writeTextFile(os.path.join(path, filename + '.txt'), text) return( text, ) # Save Text FileNotFoundError def writeTextFile(self, file, content): try: with open(file, 'w') as f: f.write(content) except OSError: print(f'\033[34mWAS Node Suite\033[0m Error: Unable to save file `{file}`') class WLSH_Save_Prompt_File: def __init__(s): pass @classmethod def INPUT_TYPES(s): return { "required": { "filename": ("STRING",{"default": 'info', "multiline": False}), "path": ("STRING", {"default": './output', "multiline": False}), "positive": ("STRING",{"default": '', "multiline": True}), }, "optional": { "negative": ("STRING",{"default": '', "multiline": True}), "modelname": ("STRING",{"default": '', "multiline": False}), "seed": ("SEED",), } } OUTPUT_NODE = True RETURN_TYPES = () FUNCTION = "save_text_file" CATEGORY = "WLSH Nodes/IO" def save_text_file(self, positive="", negative="", seed=0, path="./output", modelname="", filename="info"): # Ensure path exists if not os.path.exists(path): print(f'\033[34mWAS NS\033[0m Error: The path `{path}` doesn\'t exist!') # Ensure content to save if filename.strip == '': print(f'\033[34mWAS NS\033[0m Error: There is no text specified to save! Text is empty.') text = "Positive Prompt:\n" + positive + "\nNegative Prompt:\n" + negative + "\nSeed: " + str(seed['seed']) filename = self.make_filename(seed, modelname, filename) # Write text file self.writeTextFile(os.path.join(path, filename + '.txt'), text) return( text, ) def make_filename(self, seed, modelname="", filename=""): # generate datetime string now = datetime.now() timestamp = now.strftime("%Y-%m-%d-%H%M%S") # parse input string filename = filename.replace("%time",timestamp) filename = filename.replace("%model",modelname) filename = filename.replace("%seed",str(seed['seed'])) return (filename) # Save Text FileNotFoundError def writeTextFile(self, file, content): try: with open(file, 'w') as f: f.write(content) except OSError: print(f'\033[34mWAS Node Suite\033[0m Error: Unable to save file `{file}`') class WLSH_Save_Positive_Prompt_File: def __init__(s): pass @classmethod def INPUT_TYPES(s): return { "required": { "filename": ("STRING",{"default": 'info', "multiline": False}), "path": ("STRING", {"default": './output', "multiline": False}), "positive": ("STRING",{"default": '', "multiline": True}), } } OUTPUT_NODE = True RETURN_TYPES = () FUNCTION = "save_text_file" CATEGORY = "WLSH Nodes/IO" def save_text_file(self, positive="", path="./output", filename="info"): # Ensure path exists if not os.path.exists(path): print(f'\033[34mWAS NS\033[0m Error: The path `{path}` doesn\'t exist!') # Ensure content to save if filename.strip == '': print(f'\033[34mWAS NS\033[0m Error: There is no text specified to save! Text is empty.') # filename = self.make_filename(seed, modelname, filename) # Write text file self.writeTextFile(os.path.join(path, filename + '.txt'), positive) return( positive, ) # Save Text FileNotFoundError def writeTextFile(self, file, content): try: with open(file, 'w') as f: f.write(content) except OSError: print(f'\033[34mWAS Node Suite\033[0m Error: Unable to save file `{file}`') NODE_CLASS_MAPPINGS = { "Checkpoint Loader w/Name (WLSH)": WLSH_Checkpoint_Loader_Model_Name, "KSamplerAdvanced (WLSH)": WLSH_KSamplerAdvancedMod, "Alternating KSampler (WLSH)": WLSH_Alternating_KSamplerAdvanced, "Seed to Number (WLSH)": WLSH_Seed_to_Number, "SDXL Steps (WLSH)": WLSH_SDXL_Steps, "SDXL Resolutions (WLSH)": WLSH_SDXL_Resolutions, "Resolutions by Ratio (WLSH)": WLSH_Resolutions_by_Ratio, "Multiply Integer (WLSH)": WLSH_Int_Multiply, "Time String (WLSH)": WLSH_Time_String, "Empty Latent by Ratio (WLSH)" : WLSH_Empty_Latent_Image_By_Ratio, "SDXL Quick Empty Latent (WLSH)" : WLSH_SDXL_Quick_Empty_Latent, "CLIP Positive-Negative (WLSH)": WLSH_CLIP_Positive_Negative, "CLIP Positive-Negative w/Text (WLSH)": WLSH_CLIP_Text_Positive_Negative, "Outpaint to Image (WLSH)": WLSH_Outpaint_To_Image, "VAE Encode for Inpaint Padding (WLSH)": WLSH_VAE_Encode_For_Inpaint_Padding, "Generate Edge Mask (WLSH)": WLSH_Generate_Edge_Mask, # "Generate Face Mask (WLSH)": WLSH_Generate_Face_Mask, "Image Scale By Factor (WLSH)": WLSH_Image_Scale_By_Factor, "Upscale by Factor with Model (WLSH)": WLSH_Upscale_By_Factor_With_Model, "SDXL Quick Image Scale (WLSH)": WLSH_SDXL_Quick_Image_Scale, "Image Save with Prompt Data (WLSH)": WLSH_Image_Save_With_Prompt_Info, "Save Prompt Info (WLSH)": WLSH_Save_Prompt_File, "Image Save with Prompt File (WLSH)": WLSH_Image_Save_With_Prompt_File, "Save Positive Prompt File (WLSH)": WLSH_Save_Positive_Prompt_File, }