diff --git a/__init__.py b/__init__.py index ee43d88..4975ac5 100755 --- a/__init__.py +++ b/__init__.py @@ -11,6 +11,8 @@ 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 from .nodes.filters.AQ_filters import AQ_BatchAverageImage, AQ_ColorMatchImage +from .nodes.text_encoders.AQ_TextEncodeQwenImage import AQ_TextEncodeQwenImageEdit +from .nodes.text_encoders.AQ_TextEncodeQwenImage import AQ_TextEncodeQwenImageEditPlus NODE_CLASS_MAPPINGS = { "AQ_MasksAndImagesAsList": AQ_MasksAndImagesAsList, @@ -29,7 +31,9 @@ NODE_CLASS_MAPPINGS = { "AQ_ImageMaskSwitch": AQ_ImageMaskSwitch, "AQ_SaveImageWebpReturnBase64": AQ_SaveImageWebpReturnBase64, "AQ_BatchAverageImage": AQ_BatchAverageImage, - "AQ_ColorMatchImage": AQ_ColorMatchImage + "AQ_ColorMatchImage": AQ_ColorMatchImage, + "AQ_TextEncodeQwenImageEdit": AQ_TextEncodeQwenImageEdit, + "AQ_TextEncodeQwenImageEditPlus": AQ_TextEncodeQwenImageEditPlus } NODE_DISPLAY_NAME_MAPPINGS = { @@ -49,7 +53,10 @@ NODE_DISPLAY_NAME_MAPPINGS = { "AQ_ImageMaskSwitch": "AQ_ImageMaskSwitch", "AQ_SaveImageWebpReturnBase64": "AQ_SaveImageWebpReturnBase64", "AQ_BatchAverageImage": "AQ_BatchAverageImage", - "AQ_ColorMatchImage": "AQ_ColorMatchImage" + "AQ_ColorMatchImage": "AQ_ColorMatchImage", + "AQ_TextEncodeQwenImageEdit": "AQ_TextEncodeQwenImageEdit", + "AQ_TextEncodeQwenImageEditPlus": "AQ_TextEncodeQwenImageEditPlus" + } diff --git a/nodes/text_encoders/AQ_TextEncodeQwenImage.py b/nodes/text_encoders/AQ_TextEncodeQwenImage.py new file mode 100755 index 0000000..17151a3 --- /dev/null +++ b/nodes/text_encoders/AQ_TextEncodeQwenImage.py @@ -0,0 +1,123 @@ +import node_helpers +import comfy.utils +import math + + +class AQ_TextEncodeQwenImageEdit: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "clip": ("CLIP",), + "prompt": ("STRING", {"multiline": True, "dynamicPrompts": True}), + "use_image": ("BOOLEAN", {"default": True}), + }, + "optional": { + "vae": ("VAE",), + "image": ("IMAGE",), + }, + } + + RETURN_TYPES = ("CONDITIONING",) + FUNCTION = "encode" + + CATEGORY = "advanced/conditioning" + + def encode(self, clip, prompt, use_image=True, vae=None, image=None): + ref_latent = None + if image is None or not use_image: + images = [] + else: + samples = image.movedim(-1, 1) + total = int(1024 * 1024) + + scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2])) + width = round(samples.shape[3] * scale_by) + height = round(samples.shape[2] * scale_by) + + s = comfy.utils.common_upscale(samples, width, height, "area", "disabled") + image = s.movedim(1, -1) + images = [image[:, :, :, :3]] + if vae is not None: + ref_latent = vae.encode(image[:, :, :, :3]) + + tokens = clip.tokenize(prompt, images=images) + conditioning = clip.encode_from_tokens_scheduled(tokens) + if ref_latent is not None: + conditioning = node_helpers.conditioning_set_values( + conditioning, {"reference_latents": [ref_latent]}, append=True + ) + return (conditioning,) + + +class AQ_TextEncodeQwenImageEditPlus: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "clip": ("CLIP",), + "prompt": ("STRING", {"multiline": True, "dynamicPrompts": True}), + "use_image1": ("BOOLEAN", {"default": True}), + "use_image2": ("BOOLEAN", {"default": True}), + "use_image3": ("BOOLEAN", {"default": True}), + }, + "optional": { + "vae": ("VAE",), + "image1": ("IMAGE",), + "image2": ("IMAGE",), + "image3": ("IMAGE",), + }, + } + + RETURN_TYPES = ("CONDITIONING",) + FUNCTION = "encode" + + CATEGORY = "advanced/conditioning" + + def encode(self, clip, prompt, use_image1=True, use_image2=True, use_image3=True, vae=None, image1=None, image2=None, image3=None): + ref_latents = [] + images = [image1, image2, image3] + use_images = [use_image1, use_image2, use_image3] + images_vl = [] + llama_template = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + image_prompt = "" + + for i, (image, use_image) in enumerate(zip(images, use_images)): + if image is not None and use_image: + samples = image.movedim(-1, 1) + total = int(384 * 384) + + scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2])) + width = round(samples.shape[3] * scale_by) + height = round(samples.shape[2] * scale_by) + + s = comfy.utils.common_upscale( + samples, width, height, "area", "disabled" + ) + images_vl.append(s.movedim(1, -1)) + if vae is not None: + total = int(1024 * 1024) + scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2])) + width = round(samples.shape[3] * scale_by / 8.0) * 8 + height = round(samples.shape[2] * scale_by / 8.0) * 8 + + s = comfy.utils.common_upscale( + samples, width, height, "area", "disabled" + ) + ref_latents.append(vae.encode(s.movedim(1, -1)[:, :, :, :3])) + + image_prompt += ( + "Picture {}: <|vision_start|><|image_pad|><|vision_end|>".format( + i + 1 + ) + ) + + tokens = clip.tokenize( + image_prompt + prompt, images=images_vl, llama_template=llama_template + ) + conditioning = clip.encode_from_tokens_scheduled(tokens) + if len(ref_latents) > 0: + conditioning = node_helpers.conditioning_set_values( + conditioning, {"reference_latents": ref_latents}, append=True + ) + return (conditioning,)