diff --git a/comfy_extras/nodes_qwen.py b/comfy_extras/nodes_qwen.py index ab21acf32..b5088fae2 100644 --- a/comfy_extras/nodes_qwen.py +++ b/comfy_extras/nodes_qwen.py @@ -1,3 +1,25 @@ +import node_helpers +import comfy.utils + +PREFERRED_QWENIMAGE_RESOLUTIONS = [ + (672, 1568), + (688, 1504), + (720, 1456), + (752, 1392), + (800, 1328), + (832, 1248), + (880, 1184), + (944, 1104), + (1024, 1024), + (1104, 944), + (1184, 880), + (1248, 832), + (1328, 800), + (1392, 752), + (1456, 720), + (1504, 688), + (1568, 672), +] class TextEncodeQwenImageEdit: @@ -7,19 +29,33 @@ class TextEncodeQwenImageEdit: "clip": ("CLIP", ), "prompt": ("STRING", {"multiline": True, "dynamicPrompts": True}), }, - "optional": {"image": ("IMAGE", ),}} + "optional": {"vae": ("VAE", ), + "image": ("IMAGE", ),}} + RETURN_TYPES = ("CONDITIONING",) FUNCTION = "encode" CATEGORY = "advanced/conditioning" - def encode(self, clip, prompt, image=None): + def encode(self, clip, prompt, vae=None, image=None): + ref_latent = None if image is None: images = [] else: images = [image] + if vae is not None: + width = image.shape[2] + height = image.shape[1] + aspect_ratio = width / height + _, width, height = min((abs(aspect_ratio - w / h), w, h) for w, h in PREFERRED_QWENIMAGE_RESOLUTIONS) + image = comfy.utils.common_upscale(image.movedim(-1, 1), width, height, "lanczos", "center").movedim(1, -1) + ref_latent = vae.encode(image[:, :, :, :3]) + tokens = clip.tokenize(prompt, images=images) - return (clip.encode_from_tokens_scheduled(tokens), ) + 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, ) NODE_CLASS_MAPPINGS = {