add more AQnodes
This commit is contained in:
pk
+26
-2
@@ -5,13 +5,28 @@ from .nodes.pulid.AQ_multiface_ApplyPulidFlux import AQ_multiface_ApplyPulidFlux
|
||||
from .nodes.api.AQ_Gemini import AQ_Gemini
|
||||
from .nodes.api.AQ_SendImageToAPI import AQ_SendImageToAPI
|
||||
|
||||
from .nodes.clip.AQ_CLIPSetLastLayer import AQ_CLIPSetLastLayer
|
||||
from .nodes.numbers.AQ_Numbers import AQ_Increment, AQ_Random
|
||||
from .nodes.blend.AQ_BlendImages import AQ_BlendImages
|
||||
from .nodes.image.AQ_Image import AQ_LoadImageBase64, AQ_Image_Pad, AQ_Image_DetailTransfer, AQ_ImageMaskSwitch
|
||||
from .nodes.save.AQ_SaveImageWebpReturnBase64 import AQ_SaveImageWebpReturnBase64
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AQ_MasksAndImagesAsList": AQ_MasksAndImagesAsList,
|
||||
"AQ_multiface_ApplyPulidFlux": AQ_multiface_ApplyPulidFlux,
|
||||
"AQ_Qwen": AQ_Qwen,
|
||||
"AQ_QwenLoader": AQ_QwenLoader,
|
||||
"AQ_Gemini": AQ_Gemini,
|
||||
"AQ_SendImageToAPI": AQ_SendImageToAPI
|
||||
"AQ_SendImageToAPI": AQ_SendImageToAPI,
|
||||
"AQ_CLIPSetLastLayer": AQ_CLIPSetLastLayer,
|
||||
"AQ_Increment": AQ_Increment,
|
||||
"AQ_Random": AQ_Random,
|
||||
"AQ_BlendImages": AQ_BlendImages,
|
||||
"AQ_LoadImageBase64": AQ_LoadImageBase64,
|
||||
"AQ_Image_Pad": AQ_Image_Pad,
|
||||
"AQ_Image_DetailTransfer": AQ_Image_DetailTransfer,
|
||||
"AQ_ImageMaskSwitch": AQ_ImageMaskSwitch,
|
||||
"AQ_SaveImageWebpReturnBase64": AQ_SaveImageWebpReturnBase64
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
@@ -20,7 +35,16 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AQ_Qwen": "AQ_Qwen",
|
||||
"AQ_QwenLoader": "AQ_QwenLoader",
|
||||
"AQ_Gemini": "AQ_Gemini",
|
||||
"AQ_SendImageToAPI": "AQ_SendImageToAPI"
|
||||
"AQ_SendImageToAPI": "AQ_SendImageToAPI",
|
||||
"AQ_CLIPSetLastLayer": "AQ_CLIPSetLastLayer",
|
||||
"AQ_Increment": "AQ_Increment",
|
||||
"AQ_Random": "AQ_Random",
|
||||
"AQ_BlendImages": "AQ_BlendImages",
|
||||
"AQ_LoadImageBase64": "AQ_LoadImageBase64",
|
||||
"AQ_Image_Pad": "AQ_Image_Pad",
|
||||
"AQ_Image_DetailTransfer": "AQ_Image_DetailTransfer",
|
||||
"AQ_ImageMaskSwitch": "AQ_ImageMaskSwitch",
|
||||
"AQ_SaveImageWebpReturnBase64": "AQ_SaveImageWebpReturnBase64"
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
import torch
|
||||
|
||||
class AQ_BlendImages:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"img1": ("IMAGE",),
|
||||
"img2": ("IMAGE",),
|
||||
"device": (["dont_change", "cpu", "cuda"], {"default": "dont_change"},),
|
||||
"blend_mode": (["Normal", "Dissolve", "Darken", "Multiply", "Color Burn", "Linear Burn", "Darker Color", "Lighten", "Screen", "Color Dodge", "Linear Dodge", "Lighter Color", "Overlay", "Soft Light", "Hard Light", "Vivid Light", "Linear Light", "Pin Light", "Hard Mix", "Difference", "Exclusion", "Subtract", "Divide", "Hue", "Saturation", "Color", "Luminosity"],),
|
||||
"opacity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "blend_images"
|
||||
|
||||
CATEGORY = "AQ"
|
||||
|
||||
def blend_images(self, img1, img2, device, blend_mode, opacity):
|
||||
# if device differnet than "dont_change"
|
||||
if device == "dont_change":
|
||||
device = img1.device
|
||||
else:
|
||||
img1 = img1.to(device)
|
||||
img2 = img2.to(device)
|
||||
|
||||
|
||||
# Ensure images have the same shape
|
||||
if img1.shape != img2.shape:
|
||||
raise ValueError("Images must have the same dimensions")
|
||||
|
||||
# Separate color and alpha channels
|
||||
img1, img1_a = img1[..., :3], img1[..., 3:4]
|
||||
img2, img2_a = img2[..., :3], img2[..., 3:4]
|
||||
|
||||
# Helper functions
|
||||
def lum(img):
|
||||
return 0.3 * img[..., 0] + 0.59 * img[..., 1] + 0.11 * img[..., 2]
|
||||
|
||||
def sat(img):
|
||||
min_val, _ = torch.min(img, dim=-1)
|
||||
max_val, _ = torch.max(img, dim=-1)
|
||||
return max_val - min_val
|
||||
|
||||
def set_lum(c, l):
|
||||
d = l - lum(c)
|
||||
c = c + d.unsqueeze(-1)
|
||||
return clip_color(c)
|
||||
|
||||
def set_sat(c, s):
|
||||
c_min, _ = torch.min(c, dim=-1, keepdim=True)
|
||||
c_max, _ = torch.max(c, dim=-1, keepdim=True)
|
||||
c_mid = c - c_min
|
||||
c_mid *= s.unsqueeze(-1) / (c_max - c_min + 1e-6)
|
||||
return clip_color(c_mid + c_min)
|
||||
|
||||
def clip_color(c):
|
||||
l = lum(c).unsqueeze(-1)
|
||||
n = torch.min(c, dim=-1, keepdim=True)[0]
|
||||
x = torch.max(c, dim=-1, keepdim=True)[0]
|
||||
c_clipped = torch.clamp(c, 0, 1)
|
||||
c_clipped = torch.where(n < 0, l + ((c - l) * l) / (l - n + 1e-6), c_clipped)
|
||||
c_clipped = torch.where(x > 1, l + ((c - l) * (1 - l)) / (x - l + 1e-6), c_clipped)
|
||||
return c_clipped
|
||||
|
||||
def hue_blend(img1, img2):
|
||||
# Calculate luminance and saturation for both images
|
||||
lum1 = lum(img1)
|
||||
lum2 = lum(img2)
|
||||
sat1 = sat(img1)
|
||||
sat2 = sat(img2)
|
||||
|
||||
# Combine the hue from img1 with the saturation and luminance from img2
|
||||
result = set_lum(set_sat(img1, sat2), lum2)
|
||||
|
||||
return result
|
||||
|
||||
# Blend mode calculations (only on RGB channels)
|
||||
if blend_mode == "Normal":
|
||||
result = img1
|
||||
elif blend_mode == "Dissolve":
|
||||
mask = torch.rand_like(img1) < opacity
|
||||
result = torch.where(mask, img1, img2)
|
||||
elif blend_mode == "Darken":
|
||||
result = torch.min(img1, img2)
|
||||
elif blend_mode == "Multiply":
|
||||
result = img1 * img2
|
||||
elif blend_mode == "Color Burn":
|
||||
result = torch.where(
|
||||
img1 == 1,
|
||||
img2, # If top layer (source) is white, no change
|
||||
torch.where(
|
||||
img1 > 0,
|
||||
1 - torch.clamp((1 - img2) / img1, 0, 1),
|
||||
torch.zeros_like(img2) # If top layer is black, result is black
|
||||
)
|
||||
)
|
||||
elif blend_mode == "Linear Burn":
|
||||
result = img1 + img2 - 1
|
||||
elif blend_mode == "Darker Color":
|
||||
result = torch.where(lum(img1).unsqueeze(-1) < lum(img2).unsqueeze(-1), img1, img2)
|
||||
elif blend_mode == "Lighten":
|
||||
result = torch.max(img1, img2)
|
||||
elif blend_mode == "Screen":
|
||||
result = 1 - (1 - img1) * (1 - img2)
|
||||
elif blend_mode == "Color Dodge":
|
||||
# result = torch.where(img2 < 1, torch.min(torch.ones_like(img1), img1 / (1 - img2)), torch.ones_like(img1))
|
||||
result = torch.where(img2 < 1, torch.min(torch.ones_like(img1), img1 / (1 - img2)), torch.ones_like(img1))
|
||||
elif blend_mode == "Linear Dodge":
|
||||
result = img1 + img2
|
||||
elif blend_mode == "Lighter Color":
|
||||
result = torch.where(lum(img1).unsqueeze(-1) > lum(img2).unsqueeze(-1), img1, img2)
|
||||
elif blend_mode == "Overlay":
|
||||
result = torch.where(img2 < 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2))
|
||||
elif blend_mode == "Soft Light":
|
||||
d = torch.where(img2 <= 0.25,
|
||||
((16 * img2 - 12) * img2 + 4) * img2,
|
||||
torch.sqrt(img2))
|
||||
result = torch.where(img1 <= 0.5,
|
||||
img2 - (1 - 2 * img1) * img2 * (1 - img2),
|
||||
img2 + (2 * img1 - 1) * (d - img2))
|
||||
elif blend_mode == "Hard Light":
|
||||
result = torch.where(img1 <= 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2))
|
||||
elif blend_mode == "Vivid Light":
|
||||
result = torch.where(img1 < 0.5,
|
||||
1 - (1 - img2) / (2 * img1 + 1e-8),
|
||||
img2 / (2 * (1 - img1) + 1e-8))
|
||||
elif blend_mode == "Linear Light":
|
||||
result = img2 + 2 * img1 - 1
|
||||
elif blend_mode == "Pin Light":
|
||||
result = torch.where(img1 < 0.5,
|
||||
torch.min(img2, 2 * img1),
|
||||
torch.max(img2, 2 * img1 - 1))
|
||||
elif blend_mode == "Hard Mix":
|
||||
# result = (img1 + img2).round()
|
||||
# Calculate the sum of the base and blend images
|
||||
sum_img = img1 + img2
|
||||
# Apply the Hard Mix formula
|
||||
result = torch.where(sum_img >= 1, torch.ones_like(img1), torch.zeros_like(img1))
|
||||
elif blend_mode == "Difference":
|
||||
result = torch.abs(img1 - img2)
|
||||
elif blend_mode == "Exclusion":
|
||||
result = img1 + img2 - 2 * img1 * img2
|
||||
elif blend_mode == "Subtract":
|
||||
result = img2 - img1
|
||||
elif blend_mode == "Divide":
|
||||
result = img2 / (img1 + 1e-8)
|
||||
elif blend_mode == "Hue":
|
||||
result = hue_blend(img1, img2)
|
||||
elif blend_mode == "Saturation":
|
||||
result = set_lum(set_sat(img2, sat(img1)), lum(img2))
|
||||
elif blend_mode == "Color":
|
||||
result = set_lum(img1, lum(img2))
|
||||
elif blend_mode == "Luminosity":
|
||||
result = set_lum(img2, lum(img1))
|
||||
|
||||
# self.print_tensor_shape(result, "result before opacity")
|
||||
|
||||
# Apply opacity
|
||||
opacity_tensor = torch.tensor(opacity, device=device).view(1, 1, 1, 1)
|
||||
# self.print_tensor_shape(opacity_tensor, "opacity_tensor")
|
||||
|
||||
result_rgb = img2 * (1 - opacity_tensor) + result * opacity_tensor
|
||||
|
||||
# Blend alpha channels
|
||||
result_a = img1_a * opacity_tensor + img2_a * (1 - opacity_tensor)
|
||||
|
||||
# self.print_tensor_shape(result, "final result")
|
||||
|
||||
# Combine RGB and alpha channels
|
||||
result = torch.cat([result_rgb, result_a], dim=-1)
|
||||
|
||||
return (result,)
|
||||
|
||||
# Add this method to print tensor shapes for debugging
|
||||
def print_tensor_shape(self, tensor, name):
|
||||
print(f"{name} shape: {tensor.shape}")
|
||||
@@ -0,0 +1,19 @@
|
||||
class AQ_CLIPSetLastLayer:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "clip": ("CLIP", ),
|
||||
"stop_at_clip_layer": ("INT", {"default": 0, "min": -24, "max": 0, "step": 1}),
|
||||
}}
|
||||
RETURN_TYPES = ("CLIP",)
|
||||
FUNCTION = "set_last_layer"
|
||||
|
||||
CATEGORY = "AQ/conditioning"
|
||||
|
||||
|
||||
def set_last_layer(self, clip, stop_at_clip_layer):
|
||||
if(stop_at_clip_layer == 0):
|
||||
return (clip,)
|
||||
|
||||
clip = clip.clone()
|
||||
clip.clip_layer(stop_at_clip_layer)
|
||||
return (clip,)
|
||||
@@ -0,0 +1,189 @@
|
||||
import torch
|
||||
import comfy.utils
|
||||
import comfy.model_management
|
||||
from torchvision.transforms import Resize, CenterCrop, GaussianBlur
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import base64
|
||||
from io import BytesIO
|
||||
|
||||
class AQ_LoadImageBase64:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {"image": ("STRING", {"multiline": False})}}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK")
|
||||
CATEGORY = "AQ/Image"
|
||||
FUNCTION = "load_image"
|
||||
|
||||
def load_image(self, image):
|
||||
imgdata = base64.b64decode(image)
|
||||
img = Image.open(BytesIO(imgdata))
|
||||
|
||||
if "A" in img.getbands():
|
||||
mask = np.array(img.getchannel("A")).astype(np.float32) / 255.0
|
||||
mask = 1.0 - torch.from_numpy(mask)
|
||||
else:
|
||||
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
|
||||
|
||||
img = img.convert("RGB")
|
||||
img = np.array(img).astype(np.float32) / 255.0
|
||||
img = torch.from_numpy(img)[None,]
|
||||
|
||||
return (img, mask)
|
||||
|
||||
class AQ_Image_Pad:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"left": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
"right": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
"top": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
"bottom": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "pad_image"
|
||||
|
||||
CATEGORY = "AQ/Image"
|
||||
|
||||
def pad_image(self, image, left, right, top, bottom):
|
||||
# Get the current shape of the image
|
||||
batch, height, width, channels = image.shape
|
||||
|
||||
# Create a new tensor with the padded dimensions
|
||||
new_height = height + top + bottom
|
||||
new_width = width + left + right
|
||||
padded_image = torch.zeros((batch, new_height, new_width, channels), dtype=image.dtype, device=image.device)
|
||||
|
||||
# Set the alpha channel to 0 (fully transparent) for the entire padded image
|
||||
padded_image[:, :, :, 3] = 0
|
||||
|
||||
# Copy the original image into the padded image
|
||||
padded_image[:, top:top+height, left:left+width, :] = image
|
||||
|
||||
return (padded_image,)
|
||||
|
||||
class AQ_Image_DetailTransfer:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"target": ("IMAGE",),
|
||||
"source": ("IMAGE",),
|
||||
"mode": (["add", "multiply", "screen", "overlay", "soft_light", "hard_light", "color_dodge", "color_burn", "difference", "exclusion", "divide",],{"default": "add"}),
|
||||
"blur_sigma": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 100.0, "step": 0.01}),
|
||||
"blend_factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.001, "round": 0.001}),
|
||||
"image_output": (["Hide", "Preview", "Save", "Hide/Save"], {"default": "Preview"}),
|
||||
"save_prefix": ("STRING", {"default": "ComfyUI"}),
|
||||
},
|
||||
"optional": {
|
||||
"mask": ("MASK",),
|
||||
},
|
||||
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
OUTPUT_NODE = False
|
||||
FUNCTION = "transfer"
|
||||
CATEGORY = "AQ/Image"
|
||||
|
||||
|
||||
|
||||
def transfer(self, target, source, mode, blur_sigma, blend_factor, image_output, save_prefix, mask=None, prompt=None, extra_pnginfo=None):
|
||||
batch_size, height, width, _ = target.shape
|
||||
device = comfy.model_management.get_torch_device()
|
||||
target_tensor = target.permute(0, 3, 1, 2).clone().to(device)
|
||||
source_tensor = source.permute(0, 3, 1, 2).clone().to(device)
|
||||
|
||||
if target.shape[1:] != source.shape[1:]:
|
||||
source_tensor = comfy.utils.common_upscale(source_tensor, width, height, "bilinear", "disabled")
|
||||
|
||||
if source.shape[0] < batch_size:
|
||||
source = source[0].unsqueeze(0).repeat(batch_size, 1, 1, 1)
|
||||
|
||||
kernel_size = int(6 * int(blur_sigma) + 1)
|
||||
|
||||
gaussian_blur = GaussianBlur(kernel_size=(kernel_size, kernel_size), sigma=(blur_sigma, blur_sigma))
|
||||
|
||||
blurred_target = gaussian_blur(target_tensor)
|
||||
blurred_source = gaussian_blur(source_tensor)
|
||||
|
||||
if mode == "add":
|
||||
new_image = (source_tensor - blurred_source) + blurred_target
|
||||
elif mode == "multiply":
|
||||
new_image = source_tensor * blurred_target
|
||||
elif mode == "screen":
|
||||
new_image = 1 - (1 - source_tensor) * (1 - blurred_target)
|
||||
elif mode == "overlay":
|
||||
new_image = torch.where(blurred_target < 0.5, 2 * source_tensor * blurred_target,
|
||||
1 - 2 * (1 - source_tensor) * (1 - blurred_target))
|
||||
elif mode == "soft_light":
|
||||
new_image = (1 - 2 * blurred_target) * source_tensor ** 2 + 2 * blurred_target * source_tensor
|
||||
elif mode == "hard_light":
|
||||
new_image = torch.where(source_tensor < 0.5, 2 * source_tensor * blurred_target,
|
||||
1 - 2 * (1 - source_tensor) * (1 - blurred_target))
|
||||
elif mode == "difference":
|
||||
new_image = torch.abs(blurred_target - source_tensor)
|
||||
elif mode == "exclusion":
|
||||
new_image = 0.5 - 2 * (blurred_target - 0.5) * (source_tensor - 0.5)
|
||||
elif mode == "color_dodge":
|
||||
new_image = blurred_target / (1 - source_tensor)
|
||||
elif mode == "color_burn":
|
||||
new_image = 1 - (1 - blurred_target) / source_tensor
|
||||
elif mode == "divide":
|
||||
new_image = (source_tensor / blurred_source) * blurred_target
|
||||
else:
|
||||
new_image = source_tensor
|
||||
|
||||
new_image = torch.lerp(target_tensor, new_image, blend_factor)
|
||||
if mask is not None:
|
||||
mask = mask.to(device)
|
||||
new_image = torch.lerp(target_tensor, new_image, mask)
|
||||
new_image = torch.clamp(new_image, 0, 1)
|
||||
new_image = new_image.permute(0, 2, 3, 1).cpu().float()
|
||||
|
||||
return {"ui": {},
|
||||
"result": (new_image,)}
|
||||
|
||||
class AQ_ImageMaskSwitch:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"select": ("INT", {"default": 1, "min": 1, "max": 4, "step": 1}),
|
||||
"images1": ("IMAGE",),
|
||||
},
|
||||
|
||||
"optional": {
|
||||
"mask1_opt": ("MASK",),
|
||||
"images2_opt": ("IMAGE",),
|
||||
"mask2_opt": ("MASK",),
|
||||
"images3_opt": ("IMAGE",),
|
||||
"mask3_opt": ("MASK",),
|
||||
"images4_opt": ("IMAGE",),
|
||||
"mask4_opt": ("MASK",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK",)
|
||||
|
||||
OUTPUT_NODE = False
|
||||
|
||||
FUNCTION = "doit"
|
||||
|
||||
CATEGORY = "AQ/Image"
|
||||
|
||||
def doit(self, select, images1, mask1_opt=None, images2_opt=None, mask2_opt=None, images3_opt=None, mask3_opt=None,
|
||||
images4_opt=None, mask4_opt=None):
|
||||
if select == 1:
|
||||
return images1, mask1_opt,
|
||||
elif select == 2:
|
||||
return images2_opt, mask2_opt,
|
||||
elif select == 3:
|
||||
return images3_opt, mask3_opt,
|
||||
else:
|
||||
return images4_opt, mask4_opt,
|
||||
@@ -0,0 +1,44 @@
|
||||
import random
|
||||
|
||||
class AQ_Random:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"initial": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
}}
|
||||
RETURN_TYPES = ("INT",)
|
||||
FUNCTION = "generateRandomNumber"
|
||||
|
||||
CATEGORY = "AQ/numbers"
|
||||
|
||||
|
||||
def generateRandomNumber(self, initial):
|
||||
randNumber = random.randint(0, 0xffffffffffffffff)
|
||||
|
||||
return (randNumber,)
|
||||
|
||||
|
||||
class AQ_Increment:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"initial": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
"force_initial": ("BOOLEAN", {"default": False}),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("INT",)
|
||||
FUNCTION = "incrementNumber"
|
||||
|
||||
CATEGORY = "AQ/numbers"
|
||||
|
||||
numberFromPrevExecution = -1
|
||||
|
||||
def incrementNumber(self, initial, force_initial):
|
||||
if(AQ_Increment.numberFromPrevExecution == -1 or force_initial):
|
||||
AQ_Increment.numberFromPrevExecution = initial
|
||||
return (initial,)
|
||||
|
||||
currentState = AQ_Increment.numberFromPrevExecution
|
||||
AQ_Increment.numberFromPrevExecution += 1
|
||||
|
||||
return (currentState,)
|
||||
@@ -0,0 +1,131 @@
|
||||
import os
|
||||
import folder_paths as comfy_paths
|
||||
import torch
|
||||
from PIL import Image, ImageFilter, ImageEnhance, ImageOps, ImageDraw, ImageChops, ImageFont
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
from server import PromptServer
|
||||
import piexif
|
||||
import piexif.helper
|
||||
import json
|
||||
import requests
|
||||
import io
|
||||
import base64
|
||||
|
||||
|
||||
MODELS_DIR = comfy_paths.models_dir
|
||||
AQ_IMAGES_OUTPUT_DIR = comfy_paths.get_output_directory()
|
||||
|
||||
|
||||
class AQ_SaveImageWebpReturnBase64:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"name": ("STRING", {"default": "image"}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"prompt_text": ("STRING", {"default": ""}),
|
||||
"cfg": ("FLOAT", { "default": 0, "min": 0.0, "max": 50.0 }),
|
||||
"steps": ("INT", {"default": 0}),
|
||||
"scheduler": ("STRING", {"default": ""}),
|
||||
"sampler": ("STRING", {"default": ""}),
|
||||
},
|
||||
"hidden": {"unique_id": "UNIQUE_ID",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO",
|
||||
"prompt": "PROMPT"
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "save_images"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "AQ/images"
|
||||
|
||||
def save_images(self, images, name, seed, prompt_text, cfg, steps, scheduler, sampler, unique_id, extra_pnginfo, prompt):
|
||||
import random
|
||||
|
||||
is_dev = os.environ.get("IS_DEV", False)
|
||||
resultsDev = []
|
||||
|
||||
results = []
|
||||
index = 0
|
||||
server = PromptServer.instance
|
||||
#subfolder = os.path.normpath(os.path.join('user_', server.client_id, server.last_prompt_id))
|
||||
subfolder = os.path.normpath(os.path.join('user_', server.client_id))
|
||||
if not os.path.exists(AQ_IMAGES_OUTPUT_DIR):
|
||||
os.makedirs(AQ_IMAGES_OUTPUT_DIR)
|
||||
|
||||
out_dir = os.path.join(AQ_IMAGES_OUTPUT_DIR, subfolder)
|
||||
|
||||
for tensor in images:
|
||||
index += 1
|
||||
name_with_index = f"{name}_{index}_{seed}_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) + '.webp'
|
||||
array = 255.0 * tensor.cpu().numpy()
|
||||
image = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))
|
||||
|
||||
import io
|
||||
import base64
|
||||
buffer = io.BytesIO()
|
||||
|
||||
userdata = {
|
||||
'prompt': prompt_text,
|
||||
'cfg': cfg,
|
||||
'steps': steps,
|
||||
'scheduler': scheduler,
|
||||
'sampler': sampler,
|
||||
'seed': seed,
|
||||
}
|
||||
|
||||
exifData = piexif.helper.UserComment.dump(
|
||||
json.dumps(userdata),
|
||||
encoding="unicode"
|
||||
)
|
||||
|
||||
# Create EXIF dictionary
|
||||
exif_dict = {"Exif": {piexif.ExifIFD.UserComment: exifData}}
|
||||
|
||||
# Convert the EXIF dictionary to bytes
|
||||
exif_bytes = piexif.dump(exif_dict)
|
||||
|
||||
|
||||
image.save(buffer, format="WEBP")
|
||||
webp_image_data = buffer.getvalue()
|
||||
webp_base64 = base64.b64encode(webp_image_data).decode('utf-8')
|
||||
|
||||
full_output_path = f"{out_dir}/{name_with_index}"
|
||||
if os.path.commonpath((AQ_IMAGES_OUTPUT_DIR, os.path.abspath(full_output_path))) != AQ_IMAGES_OUTPUT_DIR:
|
||||
print("Saving image outside the output folder is not allowed.")
|
||||
return {}
|
||||
|
||||
if not os.path.exists(out_dir):
|
||||
os.makedirs(out_dir)
|
||||
|
||||
image.save(full_output_path, format="WEBP", exif=exif_bytes)
|
||||
|
||||
|
||||
full_output_path_txt = full_output_path.replace(".webp", ".txt")
|
||||
with open(full_output_path_txt, "w") as f:
|
||||
f.write(prompt_text)
|
||||
|
||||
if is_dev and extra_pnginfo is not None:
|
||||
resultsDev.append({
|
||||
"filename": name_with_index + " [output]",
|
||||
"subfolder": out_dir,
|
||||
"type": "base64Image",
|
||||
"name": name_with_index
|
||||
})
|
||||
|
||||
return { "ui": { "images": resultsDev } }
|
||||
|
||||
results.append(
|
||||
{
|
||||
"type": "base64Image",
|
||||
"name": name_with_index,
|
||||
"data": webp_base64,
|
||||
"seed": seed,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
return ({"ui": {"images": results}})
|
||||
Reference in New Issue
Block a user