134 lines
5.1 KiB
Python
134 lines
5.1 KiB
Python
"""
|
|
XB-ToolBox Qwen 图像编辑编码节点
|
|
=================================
|
|
复刻官方 TextEncodeQwenImageEdit / TextEncodeQwenImageEditPlus。
|
|
转换为传统 ComfyUI 节点格式，图片输入动态增加（最多9张）。
|
|
"""
|
|
|
|
import node_helpers
|
|
import comfy.utils
|
|
import math
|
|
import torch
|
|
import nodes
|
|
from .nodes_reference_any import FlexibleOptionalInputType
|
|
|
|
LLAMA_TEMPLATE = (
|
|
"<|im_start|>system\n"
|
|
"Describe 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"
|
|
)
|
|
|
|
|
|
class XB_TextEncodeQwenImageEdit:
|
|
"""Qwen 图像编辑编码（基础版）— 图片输入动态扩展，最多9张。"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"prompt": ("STRING", {"multiline": True, "dynamicPrompts": True}),
|
|
},
|
|
"optional": FlexibleOptionalInputType("IMAGE", {
|
|
"vae": ("VAE", {"optional": True}),
|
|
"image1": ("IMAGE", {"optional": True}),
|
|
}),
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
RETURN_NAMES = ("条件",)
|
|
FUNCTION = "encode"
|
|
CATEGORY = "XB_ToolBox/Qwen"
|
|
|
|
def encode(self, clip, prompt, vae=None, image1=None, **kwargs):
|
|
ref_latents = []
|
|
images_vl = []
|
|
|
|
# 收集图片
|
|
imgs = [image1] + [v for k, v in kwargs.items() if v is not None and isinstance(v, torch.Tensor)]
|
|
imgs = [i for i in imgs if i is not None]
|
|
|
|
for img in imgs:
|
|
samples = img.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")
|
|
img_resized = s.movedim(1, -1)
|
|
images_vl.append(img_resized[:, :, :, :3])
|
|
if vae is not None:
|
|
ref_latents.append(vae.encode(img_resized[:, :, :, :3]))
|
|
|
|
tokens = clip.tokenize(prompt, images=images_vl)
|
|
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,)
|
|
|
|
|
|
class XB_TextEncodeQwenImageEditPlus:
|
|
"""Qwen 图像编辑编码（增强版）— 多图VL标记 + 动态图片槽位，最多9张。"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"prompt": ("STRING", {"multiline": True, "dynamicPrompts": True}),
|
|
},
|
|
"optional": FlexibleOptionalInputType("IMAGE", {
|
|
"vae": ("VAE", {"optional": True}),
|
|
"image1": ("IMAGE", {"optional": True}),
|
|
}),
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
RETURN_NAMES = ("条件",)
|
|
FUNCTION = "encode"
|
|
CATEGORY = "XB_ToolBox/Qwen"
|
|
|
|
def encode(self, clip, prompt, vae=None, image1=None, **kwargs):
|
|
ref_latents = []
|
|
images_vl = []
|
|
|
|
# 收集图片
|
|
imgs = [image1] + [v for k, v in kwargs.items() if v is not None and isinstance(v, torch.Tensor)]
|
|
imgs = [i for i in imgs if i is not None]
|
|
|
|
image_prompt = ""
|
|
for i, image in enumerate(imgs):
|
|
samples = image.movedim(-1, 1)
|
|
|
|
# VL 尺寸 (384x384 区域)
|
|
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))
|
|
|
|
# VAE 编码 (1024x1024 区域)
|
|
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,)
|