mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-20 05:18:15 +08:00
306 lines
14 KiB
Python
306 lines
14 KiB
Python
import node_helpers
|
|
import comfy.utils
|
|
import math
|
|
from typing_extensions import override
|
|
from comfy_api.latest import ComfyExtension, io
|
|
import comfy.model_management
|
|
import torch
|
|
import nodes
|
|
|
|
|
|
def _spatial_fusion_mask(height, width, num_sources, method, block_size, dither_ratio, device):
|
|
rows = torch.arange(height, device=device).unsqueeze(1)
|
|
columns = torch.arange(width, device=device).unsqueeze(0)
|
|
|
|
if method == "spatial-checkerboard":
|
|
return ((rows + columns) % num_sources).flatten()
|
|
if method == "spatial-block-interleave":
|
|
return ((rows // block_size + columns // block_size) % num_sources).flatten()
|
|
if method == "spatial-dither-random":
|
|
generator = torch.Generator(device=device).manual_seed(42)
|
|
random = torch.rand((height, width), generator=generator, device=device)
|
|
if num_sources == 2:
|
|
return torch.where(random < dither_ratio, 0, 1).flatten()
|
|
return (random * num_sources).long().flatten()
|
|
raise ValueError(f"Unsupported visual fusion method: {method}")
|
|
|
|
|
|
def _visual_token_span(tokens, cond_length, visual_tokens):
|
|
if len(tokens) != 1:
|
|
raise ValueError("Visual fusion requires a Qwen3-VL or Krea2 text encoder.")
|
|
|
|
token_pairs = next(iter(tokens.values()))[0]
|
|
image_positions = [i for i, pair in enumerate(token_pairs) if isinstance(pair[0], dict) and pair[0].get("type") == "image"]
|
|
if len(image_positions) != 1:
|
|
raise ValueError("Visual fusion requires exactly one visual token block per encoding pass.")
|
|
|
|
image_position = image_positions[0]
|
|
if any(not isinstance(pair[0], (int, float)) for pair in token_pairs[image_position + 1:]):
|
|
raise ValueError("Visual fusion does not support embeddings after the image token block.")
|
|
|
|
end = cond_length - (len(token_pairs) - image_position - 1)
|
|
start = end - visual_tokens
|
|
if start < 0 or end > cond_length:
|
|
raise ValueError("Could not locate the visual token block in the encoded conditioning.")
|
|
return start, end
|
|
|
|
|
|
def _fuse_conditionings(conditionings, tokens, visual_height, visual_width, method, block_size, dither_ratio):
|
|
schedule_count = len(conditionings[0])
|
|
if any(len(source) != schedule_count for source in conditionings):
|
|
raise ValueError("All visual fusion sources must use the same CLIP schedule.")
|
|
|
|
visual_tokens = visual_height * visual_width
|
|
fused = []
|
|
for schedule in range(schedule_count):
|
|
source_conds = [source[schedule][0] for source in conditionings]
|
|
spans = [_visual_token_span(source_tokens, cond.shape[1], visual_tokens) for source_tokens, cond in zip(tokens, source_conds)]
|
|
if any(span != spans[0] for span in spans[1:]):
|
|
raise ValueError("Visual fusion sources produced different token layouts.")
|
|
|
|
start, end = spans[0]
|
|
visuals = torch.stack([cond[:, start:end] for cond in source_conds], dim=2)
|
|
mask = _spatial_fusion_mask(visual_height, visual_width, len(source_conds), method, block_size, dither_ratio, visuals.device)
|
|
blended_visual = torch.take_along_dim(visuals, mask[None, :, None, None], dim=2).squeeze(2)
|
|
|
|
blended = source_conds[0].clone()
|
|
blended[:, start:end] = blended_visual
|
|
fused.append([blended, conditionings[0][schedule][1].copy()])
|
|
return fused
|
|
|
|
|
|
def _flatten_images(images):
|
|
sources = []
|
|
for name in sorted(images, key=lambda value: int(value.rsplit("_", 1)[-1])):
|
|
image = images[name]
|
|
if image is None:
|
|
continue
|
|
if image.ndim == 3:
|
|
image = image.unsqueeze(0)
|
|
sources.extend(image[i:i + 1] for i in range(image.shape[0]))
|
|
return sources
|
|
|
|
|
|
class TextEncodeQwenImageEdit(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="TextEncodeQwenImageEdit",
|
|
category="model/conditioning/qwen image",
|
|
inputs=[
|
|
io.Clip.Input("clip"),
|
|
io.String.Input("prompt", multiline=True, dynamic_prompts=True),
|
|
io.Vae.Input("vae", optional=True),
|
|
io.Image.Input("image", optional=True),
|
|
],
|
|
outputs=[
|
|
io.Conditioning.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, clip, prompt, vae=None, image=None) -> io.NodeOutput:
|
|
ref_latent = None
|
|
if image is None:
|
|
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 io.NodeOutput(conditioning)
|
|
|
|
|
|
class TextEncodeQwenImageEditPlus(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="TextEncodeQwenImageEditPlus",
|
|
category="model/conditioning/qwen image",
|
|
inputs=[
|
|
io.Clip.Input("clip"),
|
|
io.String.Input("prompt", multiline=True, dynamic_prompts=True),
|
|
io.Vae.Input("vae", optional=True),
|
|
io.Image.Input("image1", optional=True),
|
|
io.Image.Input("image2", optional=True),
|
|
io.Image.Input("image3", optional=True),
|
|
],
|
|
outputs=[
|
|
io.Conditioning.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, clip, prompt, vae=None, image1=None, image2=None, image3=None) -> io.NodeOutput:
|
|
ref_latents = []
|
|
images = [image1, image2, 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 in enumerate(images):
|
|
if image is not None:
|
|
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 io.NodeOutput(conditioning)
|
|
|
|
|
|
class TextEncodeQwenImageEditFusion(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
images = io.Autogrow.TemplateNames(
|
|
io.Image.Input("image"),
|
|
names=[f"image_{i}" for i in range(1, 17)],
|
|
min=2,
|
|
)
|
|
return io.Schema(
|
|
node_id="TextEncodeQwenImageEditFusion",
|
|
display_name="Text Encode Qwen Image Edit (Visual Fusion)",
|
|
category="model/conditioning/qwen image",
|
|
description="Encodes images separately and spatially interleaves their Qwen3-VL visual conditioning tokens.",
|
|
inputs=[
|
|
io.Clip.Input("clip"),
|
|
io.String.Input("prompt", multiline=True, dynamic_prompts=True),
|
|
io.Autogrow.Input("images", template=images),
|
|
io.Combo.Input(
|
|
"fusion_method",
|
|
options=["spatial-checkerboard", "spatial-block-interleave", "spatial-dither-random"],
|
|
default="spatial-checkerboard",
|
|
),
|
|
io.Int.Input("block_size", default=2, min=1, max=8, step=1, advanced=True),
|
|
io.Float.Input(
|
|
"dither_ratio",
|
|
default=0.5,
|
|
min=0.0,
|
|
max=1.0,
|
|
step=0.01,
|
|
advanced=True,
|
|
tooltip="For two sources, the probability of selecting the first source. Three or more sources are selected uniformly.",
|
|
),
|
|
io.Vae.Input("vae", optional=True),
|
|
],
|
|
outputs=[io.Conditioning.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, clip, prompt, images: io.Autogrow.Type, fusion_method, block_size=2, dither_ratio=0.5, vae=None) -> io.NodeOutput:
|
|
sources = _flatten_images(images)
|
|
if len(sources) < 2:
|
|
raise ValueError("Visual fusion requires at least two images.")
|
|
|
|
first = sources[0].movedim(-1, 1)
|
|
total = 384 * 384
|
|
scale_by = math.sqrt(total / (first.shape[3] * first.shape[2]))
|
|
width = max(32, round(first.shape[3] * scale_by))
|
|
height = max(32, round(first.shape[2] * scale_by))
|
|
|
|
processed = []
|
|
for source in sources:
|
|
samples = source[:, :, :, :3].movedim(-1, 1)
|
|
resized = comfy.utils.common_upscale(samples, width, height, "area", "center")
|
|
processed.append(resized.movedim(1, -1))
|
|
|
|
factor = 32
|
|
visual_height = max(factor, round(height / factor) * factor) // factor
|
|
visual_width = max(factor, round(width / factor) * factor) // factor
|
|
|
|
full_prompt = (
|
|
"<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n"
|
|
"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" + prompt + "<|im_end|>\n"
|
|
"<|im_start|>assistant\n"
|
|
)
|
|
tokens = [clip.tokenize(full_prompt, images=[image]) for image in processed]
|
|
token_key = next(iter(tokens[0]), None)
|
|
if token_key not in ("qwen3vl_4b", "qwen3vl_8b") or any(next(iter(source_tokens), None) != token_key for source_tokens in tokens):
|
|
raise ValueError("Visual fusion requires a Qwen3-VL or Krea2 text encoder.")
|
|
|
|
conditionings = [clip.encode_from_tokens_scheduled(source_tokens) for source_tokens in tokens]
|
|
conditioning = _fuse_conditionings(conditionings, tokens, visual_height, visual_width, fusion_method, block_size, dither_ratio)
|
|
|
|
if vae is not None:
|
|
ref_latents = []
|
|
for source in sources:
|
|
samples = source[:, :, :, :3].movedim(-1, 1)
|
|
scale_by = math.sqrt((1024 * 1024) / (samples.shape[3] * samples.shape[2]))
|
|
latent_width = max(8, round(samples.shape[3] * scale_by / 8.0) * 8)
|
|
latent_height = max(8, round(samples.shape[2] * scale_by / 8.0) * 8)
|
|
resized = comfy.utils.common_upscale(samples, latent_width, latent_height, "area", "disabled")
|
|
ref_latents.append(vae.encode(resized.movedim(1, -1)))
|
|
conditioning = node_helpers.conditioning_set_values(conditioning, {"reference_latents": ref_latents}, append=True)
|
|
|
|
return io.NodeOutput(conditioning)
|
|
|
|
|
|
class EmptyQwenImageLayeredLatentImage(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="EmptyQwenImageLayeredLatentImage",
|
|
display_name="Empty Qwen Image Layered Latent",
|
|
category="model/latent/qwen",
|
|
inputs=[
|
|
io.Int.Input("width", default=640, min=16, max=nodes.MAX_RESOLUTION, step=16),
|
|
io.Int.Input("height", default=640, min=16, max=nodes.MAX_RESOLUTION, step=16),
|
|
io.Int.Input("layers", default=3, min=0, max=nodes.MAX_RESOLUTION, step=1),
|
|
io.Int.Input("batch_size", default=1, min=1, max=4096),
|
|
],
|
|
outputs=[
|
|
io.Latent.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, width, height, layers, batch_size=1) -> io.NodeOutput:
|
|
latent = torch.zeros([batch_size, 16, layers + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device())
|
|
return io.NodeOutput({"samples": latent})
|
|
|
|
|
|
class QwenExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
TextEncodeQwenImageEdit,
|
|
TextEncodeQwenImageEditPlus,
|
|
TextEncodeQwenImageEditFusion,
|
|
EmptyQwenImageLayeredLatentImage,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> QwenExtension:
|
|
return QwenExtension()
|