[feat]Add JoyImageEdit native model support

comfy/model_base.py

@@ -58,6 +58,7 @@
 import comfy.ldm.seedvr.model
 import comfy.ldm.boogu.model
 import comfy.ldm.qwen_image.model
+import comfy.ldm.joyimage.model
 import comfy.ldm.ideogram4.model
 import comfy.ldm.krea2.model
 import comfy.ldm.kandinsky5.model
@@ -2276,6 +2277,28 @@ def extra_conds_shapes(self, **kwargs):
             out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
         return out
 
+class JoyImage(BaseModel):
+    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.joyimage.model.JoyImageTransformer3DModel)
+        self.memory_usage_factor_conds = ("ref_latents",)
+
+    def extra_conds(self, **kwargs):
+        out = super().extra_conds(**kwargs)
+        cross_attn = kwargs.get("cross_attn", None)
+        if cross_attn is not None:
+            out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
+        ref_latents = kwargs.get("reference_latents", None)
+        if ref_latents is not None:
+            out['ref_latents'] = comfy.conds.CONDList([self.process_latent_in(lat) for lat in ref_latents])
+        return out
+
+    def extra_conds_shapes(self, **kwargs):
+        out = {}
+        ref_latents = kwargs.get("reference_latents", None)
+        if ref_latents is not None:
+            out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
+        return out
+
 class Ideogram4(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.ideogram4.model.Ideogram4Transformer2DModel)

comfy/model_detection.py

@@ -1058,6 +1058,25 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
                 dit_config["image_model"] = "SAM31"
             return dit_config
 
+    if (
+        '{}double_blocks.0.attn.img_attn_qkv.weight'.format(key_prefix) in state_dict_keys
+        and '{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix) in state_dict_keys
+        and '{}condition_embedder.time_embedder.linear_1.weight'.format(key_prefix) in state_dict_keys
+        and '{}img_in.weight'.format(key_prefix) in state_dict_keys
+        and len(state_dict['{}img_in.weight'.format(key_prefix)].shape) == 5
+    ):
+        img_in = state_dict['{}img_in.weight'.format(key_prefix)]
+        head_dim = state_dict['{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix)].shape[0]
+        return {
+            "image_model": "joyimage",
+            "in_channels": img_in.shape[1],
+            "hidden_size": img_in.shape[0],
+            "patch_size": list(img_in.shape[2:]),
+            "num_layers": count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.'),
+            "num_attention_heads": img_in.shape[0] // head_dim,
+            "text_dim": 4096,
+        }
+
     if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
         return None
 

comfy/sd.py

@@ -76,6 +76,7 @@
 import comfy.text_encoders.cogvideo
 import comfy.text_encoders.sa3
 import comfy.text_encoders.gpt_oss
+import comfy.text_encoders.joyimage
 
 import comfy.model_patcher
 import comfy.lora
@@ -1377,6 +1378,7 @@ class CLIPType(Enum):
     IDEOGRAM4 = 30
     BOOGU = 31
     KREA2 = 32
+    JOYIMAGE = 33
 
 
 
@@ -1706,6 +1708,10 @@ class EmptyClass:
                 clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
                 clip_target.clip = comfy.text_encoders.krea2.te(**llama_detect(clip_data))
                 clip_target.tokenizer = comfy.text_encoders.krea2.Krea2Tokenizer
+            elif clip_type == CLIPType.JOYIMAGE and te_model == TEModel.QWEN3VL_8B:  # JoyImageEdit: full Qwen3-VL-8B, edit-conditioning template + drop_idx.
+                clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
+                clip_target.clip = comfy.text_encoders.joyimage.te(**llama_detect(clip_data))
+                clip_target.tokenizer = comfy.text_encoders.joyimage.JoyImageTokenizer
             elif clip_type in (CLIPType.FLUX, CLIPType.FLUX2):  # Flux2 Klein reuses the Qwen3-VL LM (3-layer tap -> 12288); visual unused.
                 klein_model_type = "qwen3_8b" if te_model == TEModel.QWEN3VL_8B else "qwen3_4b"
                 clip_target.clip = comfy.text_encoders.flux.klein_te(**llama_detect(clip_data), model_type=klein_model_type)

comfy/supported_models.py

@@ -27,6 +27,7 @@
 import comfy.text_encoders.ideogram4
 import comfy.text_encoders.boogu
 import comfy.text_encoders.krea2
+import comfy.text_encoders.joyimage
 import comfy.text_encoders.anima
 import comfy.text_encoders.ace15
 import comfy.text_encoders.longcat_image
@@ -1911,6 +1912,38 @@ def clip_target(self, state_dict={}):
         hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
         return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image.QwenImageTokenizer, comfy.text_encoders.qwen_image.te(**hunyuan_detect))
 
+class JoyImage(supported_models_base.BASE):
+    unet_config = {
+        "image_model": "joyimage",
+    }
+
+    sampling_settings = {
+        "multiplier": 1000,
+        "shift": 1.5,
+    }
+
+    memory_usage_factor = 1.8
+
+    unet_extra_config = {
+        "theta": 10000,
+        "rope_dim_list": [16, 56, 56],
+    }
+
+    latent_format = latent_formats.Wan21
+
+    supported_inference_dtypes = [torch.bfloat16, torch.float32]
+
+    vae_key_prefix = ["vae."]
+    text_encoder_key_prefix = ["text_encoders."]
+
+    def get_model(self, state_dict, prefix="", device=None):
+        return model_base.JoyImage(self, device=device)
+
+    def clip_target(self, state_dict={}):
+        pref = self.text_encoder_key_prefix[0]
+        qwen3vl_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl.transformer.".format(pref))
+        return supported_models_base.ClipTarget(comfy.text_encoders.joyimage.JoyImageTokenizer, comfy.text_encoders.joyimage.te(**qwen3vl_detect))
+
 class HunyuanImage21(HunyuanVideo):
     unet_config = {
         "image_model": "hunyuan_video",
@@ -2389,6 +2422,7 @@ def get_model(self, state_dict, prefix="", device=None):
     Omnigen2,
     Boogu,
     QwenImage,
+    JoyImage,
     Ideogram4,
     Krea2,
     Flux2,

comfy/text_encoders/joyimage.py

@@ -0,0 +1,97 @@
+import torch
+
+from comfy import sd1_clip
+import comfy.text_encoders.qwen_vl
+from comfy.text_encoders.qwen3vl import Qwen3VL, Qwen3VLTokenizer
+
+JOYIMAGE_VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>"
+JOYIMAGE_TEMPLATE_TEXT = (
+    "<|im_start|>system\n \\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{}<|im_end|>\n<|im_start|>assistant\n"
+)
+JOYIMAGE_TEMPLATE_IMAGE = (
+    "<|im_start|>system\n \\nDescribe the image by detailing the color, shape, size, texture, "
+    "quantity, text, spatial relationships of the objects and background:<|im_end|>\n"
+    f"<|im_start|>user\n{JOYIMAGE_VISION_BLOCK}{{}}<|im_end|>\n<|im_start|>assistant\n"
+)
+# The DiT was trained without the leading system-prompt tokens.
+JOYIMAGE_DROP_IDX = 34
+PAD_TOKEN = 151643
+
+
+class Qwen3VL8B_JoyImage(Qwen3VL):
+    model_type = "qwen3vl_8b"
+
+    def preprocess_embed(self, embed, device):
+        if embed["type"] == "image":
+            image, grid = comfy.text_encoders.qwen_vl.process_qwen2vl_images(
+                embed["data"], min_pixels=65536, max_pixels=16777216, patch_size=16,
+                image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5],
+                interpolation="bicubic",
+            )
+            merged, deepstack = self.visual(image.to(device, dtype=torch.float32), grid)
+            return merged, {"grid": grid, "deepstack": deepstack}
+        return None, None
+
+
+class JoyImageTokenizer(Qwen3VLTokenizer):
+    def __init__(self, embedding_directory=None, tokenizer_data={}):
+        super().__init__(
+            embedding_directory=embedding_directory, tokenizer_data=tokenizer_data,
+            model_type="qwen3vl_8b",
+        )
+        self.llama_template = JOYIMAGE_TEMPLATE_TEXT
+        self.llama_template_images = JOYIMAGE_TEMPLATE_IMAGE
+
+    def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=None, **kwargs):
+        kwargs.pop("thinking", None)
+        return super().tokenize_with_weights(
+            text, return_word_ids=return_word_ids, llama_template=llama_template,
+            images=images or [], thinking=True, **kwargs,
+        )
+
+
+class _JoyImageClipModel(sd1_clip.SDClipModel):
+    def __init__(self, device="cpu", layer="hidden", layer_idx=-1, dtype=None,
+                 attention_mask=True, model_options={}):
+        super().__init__(
+            device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={},
+            # JoyImage conditions on the pre-final-norm output of the last decoder layer.
+            dtype=dtype, special_tokens={"pad": PAD_TOKEN}, layer_norm_hidden_state=False,
+            model_class=Qwen3VL8B_JoyImage, enable_attention_masks=attention_mask,
+            return_attention_masks=attention_mask, model_options=model_options,
+        )
+
+
+class JoyImageTEModel(sd1_clip.SD1ClipModel):
+    def __init__(self, device="cpu", dtype=None, model_options={}):
+        super().__init__(
+            device=device, dtype=dtype, name="qwen3vl_8b",
+            clip_model=_JoyImageClipModel, model_options=model_options,
+        )
+
+    def encode_token_weights(self, token_weight_pairs):
+        out, pooled, extra = super().encode_token_weights(token_weight_pairs)
+        if out.shape[1] <= JOYIMAGE_DROP_IDX:
+            raise ValueError(
+                f"JoyImageTEModel: encoded sequence length {out.shape[1]} is shorter "
+                f"than drop_idx={JOYIMAGE_DROP_IDX}; the prompt did not include the "
+                f"template prefix."
+            )
+        out = out[:, JOYIMAGE_DROP_IDX:]
+        if "attention_mask" in extra:
+            extra["attention_mask"] = extra["attention_mask"][:, JOYIMAGE_DROP_IDX:]
+        return out, pooled, extra
+
+
+def te(dtype_llama=None, llama_quantization_metadata=None):
+    class JoyImageTEModel_(JoyImageTEModel):
+        def __init__(self, device="cpu", dtype=None, model_options={}):
+            if llama_quantization_metadata is not None:
+                model_options = model_options.copy()
+                model_options["quantization_metadata"] = llama_quantization_metadata
+            if dtype_llama is not None:
+                dtype = dtype_llama
+            super().__init__(device=device, dtype=dtype, model_options=model_options)
+    return JoyImageTEModel_

comfy/text_encoders/qwen_vl.py

@@ -15,6 +15,7 @@ def process_qwen2vl_images(
     merge_size: int = 2,
     image_mean: list = None,
     image_std: list = None,
+    interpolation: str = "bilinear",
 ):
     if image_mean is None:
         image_mean = [0.48145466, 0.4578275, 0.40821073]
@@ -47,10 +48,9 @@ def process_qwen2vl_images(
     img_resized = F.interpolate(
         img.unsqueeze(0),
         size=(h_bar, w_bar),
-        mode='bilinear',
+        mode=interpolation,
         align_corners=False
     ).squeeze(0)
-
     normalized = img_resized.clone()
     for c in range(3):
         normalized[c] = (img_resized[c] - image_mean[c]) / image_std[c]

comfy_extras/nodes_joyimage.py

@@ -0,0 +1,102 @@
+from typing_extensions import override
+
+import comfy.utils
+import node_helpers
+from comfy_api.latest import ComfyExtension, io
+
+
+# fmt: off
+BUCKETS_1024 = [
+    (512, 1792), (512, 1856), (512, 1920), (512, 1984), (512, 2048),
+    (576, 1600), (576, 1664), (576, 1728), (576, 1792),
+    (640, 1472), (640, 1536), (640, 1600),
+    (704, 1344), (704, 1408), (704, 1472),
+    (768, 1216), (768, 1280), (768, 1344),
+    (832, 1152), (832, 1216),
+    (896, 1088), (896, 1152),
+    (960, 1024), (960, 1088),
+    (1024, 960), (1024, 1024),
+    (1088, 896), (1088, 960),
+    (1152, 832), (1152, 896),
+    (1216, 768), (1216, 832),
+    (1280, 768),
+    (1344, 704), (1344, 768),
+    (1408, 704),
+    (1472, 640), (1472, 704),
+    (1536, 640),
+    (1600, 576), (1600, 640),
+    (1664, 576),
+    (1728, 576),
+    (1792, 512), (1792, 576),
+    (1856, 512),
+    (1920, 512),
+    (1984, 512),
+    (2048, 512),
+]
+# fmt: on
+
+
+def _find_best_bucket(height: int, width: int) -> tuple[int, int]:
+    target_ratio = height / width
+    return min(BUCKETS_1024, key=lambda hw: abs(hw[0] / hw[1] - target_ratio))
+
+
+def _resize_reference(image):
+    if image.shape[0] != 1:
+        raise ValueError("JoyImage reference inputs must contain one image each")
+    samples = image.movedim(-1, 1)
+    bucket_h, bucket_w = _find_best_bucket(samples.shape[2], samples.shape[3])
+    resized = comfy.utils.common_upscale(samples, bucket_w, bucket_h, "bilinear", "center")
+    return resized.movedim(1, -1)[:, :, :, :3]
+
+
+def _encode(clip, prompt, vae, images):
+    resized_images = [_resize_reference(image) for image in images]
+    conditioning = clip.encode_from_tokens_scheduled(clip.tokenize(prompt, images=resized_images))
+    if vae is not None and resized_images:
+        ref_latents = [vae.encode(image) for image in resized_images]
+        conditioning = node_helpers.conditioning_set_values(
+            conditioning, {"reference_latents": ref_latents}, append=True,
+        )
+    return conditioning
+
+
+class TextEncodeJoyImageEdit(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        image_template = io.Autogrow.TemplatePrefix(
+            io.Image.Input("image"),
+            prefix="image",
+            min=0,
+            max=6,
+        )
+        return io.Schema(
+            node_id="TextEncodeJoyImageEdit",
+            category="model/conditioning/joyimage",
+            inputs=[
+                io.Clip.Input("clip"),
+                io.String.Input("prompt", multiline=True, dynamic_prompts=True),
+                io.Vae.Input("vae", optional=True),
+                io.Autogrow.Input("images", template=image_template, optional=True),
+            ],
+            outputs=[
+                io.Conditioning.Output(),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, clip, prompt, vae=None, images: io.Autogrow.Type = None) -> io.NodeOutput:
+        images = images or {}
+        return io.NodeOutput(_encode(clip, prompt, vae, list(images.values())))
+
+
+class JoyImageExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            TextEncodeJoyImageEdit,
+        ]
+
+
+async def comfy_entrypoint() -> JoyImageExtension:
+    return JoyImageExtension()

nodes.py

@@ -992,7 +992,7 @@ class CLIPLoader:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
-                              "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2"], ),
+                              "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2", "joyimage"], ),
                               },
                 "optional": {
                               "device": (["default", "cpu"], {"advanced": True}),
@@ -1002,7 +1002,7 @@ def INPUT_TYPES(s):
 
     CATEGORY = "model/loaders"
 
-    DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm"
+    DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\njoyimage: qwen3-vl 8B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm"
 
     def load_clip(self, clip_name, type="stable_diffusion", device="default"):
         clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
@@ -2462,6 +2462,7 @@ async def init_builtin_extra_nodes():
         "nodes_seedvr.py",
         "nodes_context_windows.py",
         "nodes_qwen.py",
+        "nodes_joyimage.py",
         "nodes_boogu.py",
         "nodes_chroma_radiance.py",
         "nodes_pid.py",