mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-20 21:38:30 +08:00
refactor: generalize image conditioning fusion
Move fusion out of Qwen Image conditioning, preserve tokenizer-owned templates, and align visual grids without resizing source images.
This commit is contained in:
parent
6aed818eff
commit
af91a7c44a
@ -6,6 +6,23 @@ import math
|
|||||||
from comfy.ldm.modules.attention import optimized_attention_for_device
|
from comfy.ldm.modules.attention import optimized_attention_for_device
|
||||||
|
|
||||||
|
|
||||||
|
def qwen2vl_image_size(height, width, min_pixels=3136, max_pixels=12845056, patch_size=14, merge_size=2):
|
||||||
|
factor = patch_size * merge_size
|
||||||
|
resized_height = round(height / factor) * factor
|
||||||
|
resized_width = round(width / factor) * factor
|
||||||
|
|
||||||
|
if resized_height * resized_width > max_pixels:
|
||||||
|
beta = math.sqrt((height * width) / max_pixels)
|
||||||
|
resized_height = max(factor, math.floor(height / beta / factor) * factor)
|
||||||
|
resized_width = max(factor, math.floor(width / beta / factor) * factor)
|
||||||
|
elif resized_height * resized_width < min_pixels:
|
||||||
|
beta = math.sqrt(min_pixels / (height * width))
|
||||||
|
resized_height = math.ceil(height * beta / factor) * factor
|
||||||
|
resized_width = math.ceil(width * beta / factor) * factor
|
||||||
|
|
||||||
|
return resized_height, resized_width
|
||||||
|
|
||||||
|
|
||||||
def process_qwen2vl_images(
|
def process_qwen2vl_images(
|
||||||
images: torch.Tensor,
|
images: torch.Tensor,
|
||||||
min_pixels: int = 3136,
|
min_pixels: int = 3136,
|
||||||
@ -30,19 +47,7 @@ def process_qwen2vl_images(
|
|||||||
grid_thw_list = []
|
grid_thw_list = []
|
||||||
img = images[0]
|
img = images[0]
|
||||||
|
|
||||||
factor = patch_size * merge_size
|
h_bar, w_bar = qwen2vl_image_size(height, width, min_pixels, max_pixels, patch_size, merge_size)
|
||||||
|
|
||||||
h_bar = round(height / factor) * factor
|
|
||||||
w_bar = round(width / factor) * factor
|
|
||||||
|
|
||||||
if h_bar * w_bar > max_pixels:
|
|
||||||
beta = math.sqrt((height * width) / max_pixels)
|
|
||||||
h_bar = max(factor, math.floor(height / beta / factor) * factor)
|
|
||||||
w_bar = max(factor, math.floor(width / beta / factor) * factor)
|
|
||||||
elif h_bar * w_bar < min_pixels:
|
|
||||||
beta = math.sqrt(min_pixels / (height * width))
|
|
||||||
h_bar = math.ceil(height * beta / factor) * factor
|
|
||||||
w_bar = math.ceil(width * beta / factor) * factor
|
|
||||||
|
|
||||||
img_resized = F.interpolate(
|
img_resized = F.interpolate(
|
||||||
img.unsqueeze(0),
|
img.unsqueeze(0),
|
||||||
|
|||||||
@ -1,8 +1,163 @@
|
|||||||
|
import torch
|
||||||
|
import torch.nn.functional as F
|
||||||
from typing_extensions import override
|
from typing_extensions import override
|
||||||
|
|
||||||
|
import comfy.text_encoders.qwen_vl
|
||||||
from comfy_api.latest import ComfyExtension, io
|
from comfy_api.latest import ComfyExtension, io
|
||||||
|
|
||||||
|
|
||||||
|
def _spatial_fusion_mask(height, width, num_sources, method, block_size, dither_ratio, device, seed=0):
|
||||||
|
rows = torch.arange(height).unsqueeze(1)
|
||||||
|
columns = torch.arange(width).unsqueeze(0)
|
||||||
|
|
||||||
|
if method == "spatial-checkerboard":
|
||||||
|
mask = (rows + columns) % num_sources
|
||||||
|
elif method == "spatial-block-interleave":
|
||||||
|
mask = (rows // block_size + columns // block_size) % num_sources
|
||||||
|
elif method == "spatial-dither-random":
|
||||||
|
generator = torch.Generator().manual_seed(seed)
|
||||||
|
random = torch.rand((height, width), generator=generator)
|
||||||
|
other_sources = 1 + ((rows + columns) % (num_sources - 1))
|
||||||
|
mask = torch.where(random < dither_ratio, 0, other_sources)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unsupported visual fusion method: {method}")
|
||||||
|
return mask.flatten().to(device)
|
||||||
|
|
||||||
|
|
||||||
|
def _visual_token_span(tokens, cond_length, visual_tokens):
|
||||||
|
if len(tokens) != 1:
|
||||||
|
raise ValueError("Image fusion requires a compatible multimodal CLIP encoder with one token stream.")
|
||||||
|
|
||||||
|
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("Image 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("Image 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 _visual_grid(image):
|
||||||
|
height, width = image.shape[1:3]
|
||||||
|
height, width = comfy.text_encoders.qwen_vl.qwen2vl_image_size(height, width, patch_size=16, merge_size=2)
|
||||||
|
return height // 32, width // 32
|
||||||
|
|
||||||
|
|
||||||
|
def _resize_visual_tokens(visual, source_grid, target_grid):
|
||||||
|
if source_grid == target_grid:
|
||||||
|
return visual
|
||||||
|
|
||||||
|
dtype = visual.dtype
|
||||||
|
batch, _, dimensions = visual.shape
|
||||||
|
height, width = source_grid
|
||||||
|
target_height, target_width = target_grid
|
||||||
|
visual = visual.reshape(batch, height, width, dimensions).permute(0, 3, 1, 2).float()
|
||||||
|
visual = F.interpolate(visual, size=(target_height, target_width), mode="bilinear", align_corners=False)
|
||||||
|
return visual.permute(0, 2, 3, 1).reshape(batch, target_height * target_width, dimensions).to(dtype=dtype)
|
||||||
|
|
||||||
|
|
||||||
|
def _fuse_conditionings(conditionings, tokens, visual_grids, method, block_size, dither_ratio, seed=0):
|
||||||
|
schedule_count = len(conditionings[0])
|
||||||
|
if any(len(source) != schedule_count for source in conditionings):
|
||||||
|
raise ValueError("All image fusion sources must use the same CLIP schedule.")
|
||||||
|
|
||||||
|
target_grid = visual_grids[0]
|
||||||
|
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], height * width) for source_tokens, cond, (height, width) in zip(tokens, source_conds, visual_grids)]
|
||||||
|
prefix_length = spans[0][0]
|
||||||
|
suffix_length = source_conds[0].shape[1] - spans[0][1]
|
||||||
|
if any(start != prefix_length or cond.shape[1] - end != suffix_length for cond, (start, end) in zip(source_conds[1:], spans[1:])):
|
||||||
|
raise ValueError("Image fusion sources produced different text token layouts.")
|
||||||
|
|
||||||
|
visuals = []
|
||||||
|
for cond, (start, end), grid in zip(source_conds, spans, visual_grids):
|
||||||
|
visual = _resize_visual_tokens(cond[:, start:end], grid, target_grid)
|
||||||
|
visuals.append(visual.to(dtype=source_conds[0].dtype, device=source_conds[0].device))
|
||||||
|
|
||||||
|
visuals = torch.stack(visuals, dim=2)
|
||||||
|
target_height, target_width = target_grid
|
||||||
|
mask = _spatial_fusion_mask(target_height, target_width, len(source_conds), method, block_size, dither_ratio, visuals.device, seed)
|
||||||
|
blended_visual = torch.take_along_dim(visuals, mask[None, :, None, None], dim=2).squeeze(2)
|
||||||
|
|
||||||
|
start, end = spans[0]
|
||||||
|
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].clone() for i in range(image.shape[0]))
|
||||||
|
return sources
|
||||||
|
|
||||||
|
|
||||||
|
class CLIPTextEncodeImageFusion(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="CLIPTextEncodeImageFusion",
|
||||||
|
display_name="CLIP Text Encode (Image Fusion)",
|
||||||
|
category="model/conditioning",
|
||||||
|
description="Encodes images separately and spatially interleaves their visual conditioning tokens.",
|
||||||
|
inputs=[
|
||||||
|
io.Clip.Input("clip"),
|
||||||
|
io.String.Input("text", 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="Probability of selecting the first source. Remaining sources are selected with a checkerboard pattern.",
|
||||||
|
),
|
||||||
|
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True, advanced=True, tooltip="Seed for the spatial-dither-random pattern."),
|
||||||
|
],
|
||||||
|
outputs=[io.Conditioning.Output()],
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def execute(cls, clip, text, images: io.Autogrow.Type, fusion_method, block_size=2, dither_ratio=0.5, seed=0) -> io.NodeOutput:
|
||||||
|
sources = _flatten_images(images)
|
||||||
|
if len(sources) < 2:
|
||||||
|
raise ValueError("Image fusion requires at least two images.")
|
||||||
|
|
||||||
|
sources = [source[:, :, :, :3] for source in sources]
|
||||||
|
visual_grids = [_visual_grid(source) for source in sources]
|
||||||
|
tokens = [clip.tokenize(text, images=[source]) for source in sources]
|
||||||
|
conditionings = [clip.encode_from_tokens_scheduled(source_tokens) for source_tokens in tokens]
|
||||||
|
conditioning = _fuse_conditionings(conditionings, tokens, visual_grids, fusion_method, block_size, dither_ratio, seed)
|
||||||
|
return io.NodeOutput(conditioning)
|
||||||
|
|
||||||
|
|
||||||
class CLIPTextEncodeControlnet(io.ComfyNode):
|
class CLIPTextEncodeControlnet(io.ComfyNode):
|
||||||
@classmethod
|
@classmethod
|
||||||
def define_schema(cls) -> io.Schema:
|
def define_schema(cls) -> io.Schema:
|
||||||
@ -61,6 +216,7 @@ class CondExtension(ComfyExtension):
|
|||||||
@override
|
@override
|
||||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||||
return [
|
return [
|
||||||
|
CLIPTextEncodeImageFusion,
|
||||||
CLIPTextEncodeControlnet,
|
CLIPTextEncodeControlnet,
|
||||||
T5TokenizerOptions,
|
T5TokenizerOptions,
|
||||||
]
|
]
|
||||||
|
|||||||
@ -7,79 +7,6 @@ import comfy.model_management
|
|||||||
import torch
|
import torch
|
||||||
import nodes
|
import nodes
|
||||||
|
|
||||||
|
|
||||||
def _spatial_fusion_mask(height, width, num_sources, method, block_size, dither_ratio, device, seed=0):
|
|
||||||
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(seed)
|
|
||||||
random = torch.rand((height, width), generator=generator, device=device)
|
|
||||||
other_sources = 1 + ((rows + columns) % (num_sources - 1))
|
|
||||||
return torch.where(random < dither_ratio, 0, other_sources).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, seed=0):
|
|
||||||
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, seed)
|
|
||||||
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].clone() for i in range(image.shape[0]))
|
|
||||||
return sources
|
|
||||||
|
|
||||||
|
|
||||||
class TextEncodeQwenImageEdit(io.ComfyNode):
|
class TextEncodeQwenImageEdit(io.ComfyNode):
|
||||||
@classmethod
|
@classmethod
|
||||||
def define_schema(cls):
|
def define_schema(cls):
|
||||||
@ -179,93 +106,6 @@ class TextEncodeQwenImageEditPlus(io.ComfyNode):
|
|||||||
return io.NodeOutput(conditioning)
|
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="Probability of selecting the first source. Remaining sources are selected with a checkerboard pattern.",
|
|
||||||
),
|
|
||||||
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True, advanced=True, tooltip="Seed for the spatial-dither-random pattern."),
|
|
||||||
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, seed=0) -> 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, seed)
|
|
||||||
|
|
||||||
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):
|
class EmptyQwenImageLayeredLatentImage(io.ComfyNode):
|
||||||
@classmethod
|
@classmethod
|
||||||
def define_schema(cls):
|
def define_schema(cls):
|
||||||
@ -296,7 +136,6 @@ class QwenExtension(ComfyExtension):
|
|||||||
return [
|
return [
|
||||||
TextEncodeQwenImageEdit,
|
TextEncodeQwenImageEdit,
|
||||||
TextEncodeQwenImageEditPlus,
|
TextEncodeQwenImageEditPlus,
|
||||||
TextEncodeQwenImageEditFusion,
|
|
||||||
EmptyQwenImageLayeredLatentImage,
|
EmptyQwenImageLayeredLatentImage,
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|||||||
@ -8,7 +8,7 @@ if not torch.cuda.is_available():
|
|||||||
cli_args.cpu = True
|
cli_args.cpu = True
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from comfy_extras.nodes_qwen import TextEncodeQwenImageEditFusion, _flatten_images, _fuse_conditionings, _spatial_fusion_mask, _visual_token_span
|
from comfy_extras.nodes_cond import CLIPTextEncodeImageFusion, _flatten_images, _fuse_conditionings, _resize_visual_tokens, _spatial_fusion_mask, _visual_grid, _visual_token_span
|
||||||
finally:
|
finally:
|
||||||
cli_args.cpu = prior_cpu
|
cli_args.cpu = prior_cpu
|
||||||
|
|
||||||
@ -72,7 +72,7 @@ def test_fusion_replaces_only_visual_tokens_and_preserves_dtype_and_metadata():
|
|||||||
[[second, {"pooled_output": torch.tensor([2.0])}]],
|
[[second, {"pooled_output": torch.tensor([2.0])}]],
|
||||||
]
|
]
|
||||||
|
|
||||||
fused = _fuse_conditionings(conditionings, tokens, 2, 2, "spatial-checkerboard", 2, 0.5)
|
fused = _fuse_conditionings(conditionings, tokens, [(2, 2), (2, 2)], "spatial-checkerboard", 2, 0.5)
|
||||||
output, output_metadata = fused[0]
|
output, output_metadata = fused[0]
|
||||||
|
|
||||||
assert output.dtype == torch.float16
|
assert output.dtype == torch.float16
|
||||||
@ -88,8 +88,8 @@ def test_dither_seed_changes_fused_conditioning():
|
|||||||
[[torch.ones((1, 6, 1)), {}]],
|
[[torch.ones((1, 6, 1)), {}]],
|
||||||
]
|
]
|
||||||
|
|
||||||
first = _fuse_conditionings(conditionings, tokens, 2, 2, "spatial-dither-random", 2, 0.5, 7)[0][0]
|
first = _fuse_conditionings(conditionings, tokens, [(2, 2), (2, 2)], "spatial-dither-random", 2, 0.5, 7)[0][0]
|
||||||
second = _fuse_conditionings(conditionings, tokens, 2, 2, "spatial-dither-random", 2, 0.5, 8)[0][0]
|
second = _fuse_conditionings(conditionings, tokens, [(2, 2), (2, 2)], "spatial-dither-random", 2, 0.5, 8)[0][0]
|
||||||
|
|
||||||
assert not torch.equal(first, second)
|
assert not torch.equal(first, second)
|
||||||
|
|
||||||
@ -103,13 +103,47 @@ def test_flatten_images_uses_numeric_input_order_and_splits_batches():
|
|||||||
|
|
||||||
sources = _flatten_images(images)
|
sources = _flatten_images(images)
|
||||||
assert [source[0, 0, 0, 0].item() for source in sources] == [1.0, 2.0, 3.0, 10.0]
|
assert [source[0, 0, 0, 0].item() for source in sources] == [1.0, 2.0, 3.0, 10.0]
|
||||||
|
images["image_2"][0, 0, 0, 0] = 99.0
|
||||||
|
assert sources[1][0, 0, 0, 0].item() == 2.0
|
||||||
|
|
||||||
|
|
||||||
def test_node_uses_custom_krea_prompt_and_returns_fused_conditioning():
|
def test_visual_tokens_interpolate_in_two_dimensions_and_restore_dtype():
|
||||||
|
visual = torch.tensor([[[0.0], [2.0]]], dtype=torch.float16)
|
||||||
|
|
||||||
|
resized = _resize_visual_tokens(visual, (2, 1), (2, 2))
|
||||||
|
|
||||||
|
assert resized.dtype == torch.float16
|
||||||
|
assert resized.flatten().tolist() == [0.0, 0.0, 2.0, 2.0]
|
||||||
|
|
||||||
|
|
||||||
|
def test_fusion_interpolates_to_first_image_grid():
|
||||||
|
first = torch.zeros((1, 6, 1), dtype=torch.float16)
|
||||||
|
second = torch.tensor([[[0.0], [0.0], [2.0], [0.0]]], dtype=torch.float16)
|
||||||
|
conditionings = [[[first, {}]], [[second, {}]]]
|
||||||
|
|
||||||
|
fused = _fuse_conditionings(conditionings, [_tokens(), _tokens()], [(2, 2), (2, 1)], "spatial-dither-random", 2, 0.0)[0][0]
|
||||||
|
|
||||||
|
assert fused.shape == first.shape
|
||||||
|
assert fused.flatten().tolist() == [0.0, 0.0, 0.0, 2.0, 2.0, 0.0]
|
||||||
|
|
||||||
|
|
||||||
|
def test_fusion_rejects_different_text_layouts():
|
||||||
|
conditionings = [
|
||||||
|
[[torch.zeros((1, 6, 1)), {}]],
|
||||||
|
[[torch.zeros((1, 7, 1)), {}]],
|
||||||
|
]
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="different text token layouts"):
|
||||||
|
_fuse_conditionings(conditionings, [_tokens(), _tokens(suffix=2)], [(2, 2), (2, 2)], "spatial-checkerboard", 2, 0.5)
|
||||||
|
|
||||||
|
|
||||||
|
def test_node_preserves_images_uses_tokenizer_template_and_returns_fused_conditioning():
|
||||||
|
seen_shapes = []
|
||||||
|
|
||||||
class FakeClip:
|
class FakeClip:
|
||||||
def tokenize(self, text, images):
|
def tokenize(self, text, images):
|
||||||
assert text.startswith("<|im_start|>system\nDescribe the image by detailing")
|
assert text == "test prompt"
|
||||||
assert "Picture 1:" not in text
|
seen_shapes.append(images[0].shape)
|
||||||
pairs = [
|
pairs = [
|
||||||
(1, 1.0),
|
(1, 1.0),
|
||||||
({"type": "image", "data": images[0]}, 1.0),
|
({"type": "image", "data": images[0]}, 1.0),
|
||||||
@ -120,27 +154,30 @@ def test_node_uses_custom_krea_prompt_and_returns_fused_conditioning():
|
|||||||
def encode_from_tokens_scheduled(self, tokens):
|
def encode_from_tokens_scheduled(self, tokens):
|
||||||
image = next(pair[0]["data"] for pair in tokens["qwen3vl_4b"][0] if isinstance(pair[0], dict))
|
image = next(pair[0]["data"] for pair in tokens["qwen3vl_4b"][0] if isinstance(pair[0], dict))
|
||||||
value = image.mean()
|
value = image.mean()
|
||||||
return [[torch.full((1, 146, 1), value, dtype=torch.float16), {"source": float(value)}]]
|
height, width = _visual_grid(image)
|
||||||
|
return [[torch.full((1, height * width + 2, 1), value, dtype=torch.float16), {"source": float(value)}]]
|
||||||
|
|
||||||
result = TextEncodeQwenImageEditFusion.execute(
|
images = {"image_1": torch.zeros(1, 32, 64, 4), "image_2": torch.ones(1, 64, 32, 4)}
|
||||||
|
result = CLIPTextEncodeImageFusion.execute(
|
||||||
FakeClip(),
|
FakeClip(),
|
||||||
"test prompt",
|
"test prompt",
|
||||||
{"image_1": torch.zeros(1, 32, 32, 3), "image_2": torch.ones(1, 32, 32, 3)},
|
images,
|
||||||
"spatial-dither-random",
|
"spatial-dither-random",
|
||||||
seed=7,
|
seed=7,
|
||||||
)
|
)
|
||||||
changed_seed = TextEncodeQwenImageEditFusion.execute(
|
changed_seed = CLIPTextEncodeImageFusion.execute(
|
||||||
FakeClip(),
|
FakeClip(),
|
||||||
"test prompt",
|
"test prompt",
|
||||||
{"image_1": torch.zeros(1, 32, 32, 3), "image_2": torch.ones(1, 32, 32, 3)},
|
images,
|
||||||
"spatial-dither-random",
|
"spatial-dither-random",
|
||||||
seed=8,
|
seed=8,
|
||||||
)
|
)
|
||||||
conditioning = result.args[0]
|
conditioning = result.args[0]
|
||||||
output, metadata = conditioning[0]
|
output, metadata = conditioning[0]
|
||||||
|
|
||||||
|
assert seen_shapes == [torch.Size([1, 32, 64, 3]), torch.Size([1, 64, 32, 3])] * 2
|
||||||
assert output.dtype == torch.float16
|
assert output.dtype == torch.float16
|
||||||
assert output.shape == (1, 146, 1)
|
assert output.shape == (1, 8, 1)
|
||||||
assert output[:, 0].item() == 0.0
|
assert output[:, 0].item() == 0.0
|
||||||
assert output[:, -1].item() == 0.0
|
assert output[:, -1].item() == 0.0
|
||||||
assert set(output[:, 1:-1].flatten().tolist()) == {0.0, 1.0}
|
assert set(output[:, 1:-1].flatten().tolist()) == {0.0, 1.0}
|
||||||
@ -148,7 +185,12 @@ def test_node_uses_custom_krea_prompt_and_returns_fused_conditioning():
|
|||||||
assert not torch.equal(output, changed_seed.args[0][0][0])
|
assert not torch.equal(output, changed_seed.args[0][0][0])
|
||||||
|
|
||||||
|
|
||||||
def test_node_exposes_seed_control():
|
def test_node_exposes_generic_interface_without_vae():
|
||||||
inputs = {value.id: value for value in TextEncodeQwenImageEditFusion.define_schema().inputs}
|
schema = CLIPTextEncodeImageFusion.define_schema()
|
||||||
|
inputs = {value.id: value for value in schema.inputs}
|
||||||
|
assert schema.node_id == "CLIPTextEncodeImageFusion"
|
||||||
|
assert schema.category == "model/conditioning"
|
||||||
|
assert "text" in inputs
|
||||||
|
assert "vae" not in inputs
|
||||||
assert inputs["seed"].default == 0
|
assert inputs["seed"].default == 0
|
||||||
assert inputs["seed"].control_after_generate is True
|
assert inputs["seed"].control_after_generate is True
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user