Repository navigation
[feat]Add JoyImageEdit native model support - #14428
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds JoyImage support across ComfyUI: a new 3D transformer, Qwen3-VL-based text encoding, JoyImage model execution and detection, and node/registry wiring for image-conditioned workflows. Tiny rhyme, no crime. 🚥 Pre-merge checks | ✅ 2✅ Passed checks (2 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
nodes.py (1)
982-982: ⚡ Quick winConsider documenting the new CLIP types in the DESCRIPTION.
The
DESCRIPTIONfield on line 982 provides recipe guidance for various CLIP types, but doesn't include the newly added"joyimage"or"pixeldit"types. While the description appears to be illustrative rather than exhaustive (note the "[Recipes]" prefix), documenting the expected text encoder configuration for JoyImage would help users.Based on the PR objectives, JoyImage uses a Qwen3-VL text encoder with text_dim=4096.
📝 Suggested documentation addition
- DESCRIPTION = "[Recipes]\n\nstable_diffusion: 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\n hidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\nlens: gpt-oss-20b\n pixeldit: gemma 2 2B elm" + DESCRIPTION = "[Recipes]\n\nstable_diffusion: 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\n hidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\nlens: gpt-oss-20b\n pixeldit: gemma 2 2B elm\njoyimage: qwen3-vl"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@nodes.py` at line 982, Update the DESCRIPTION string constant to document the new CLIP type(s): add an entry for "joyimage" specifying it uses the Qwen3-VL text encoder with text_dim=4096 (and optionally note any other relevant config), and ensure "pixeldit" remains documented (it already appears as "pixeldit: gemma 2 2B elm"); locate and modify the DESCRIPTION variable to append or insert the "joyimage: qwen3-vl (text_dim=4096)" line under the [Recipes] section so users can see the expected text-encoder configuration.comfy/text_encoders/joyimage.py (1)
94-95: 💤 Low valueConsider using
Nonedefault for mutable argument.
images=[]is a mutable default argument. While it's not modified here (only read), this pattern can cause subtle bugs if the code evolves. Usingimages=Nonewith anif images is None: images = []guard is safer.♻️ Optional refactor
def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, - images=[], **kwargs): + images=None, **kwargs): + if images is None: + images = [] if text.startswith("<|im_start|>"):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@comfy/text_encoders/joyimage.py` around lines 94 - 95, The tokenize_with_weights function uses a mutable default parameter images=[], which can cause subtle bugs; change the signature of tokenize_with_weights to use images=None and add a guard at the top of the function (if images is None: images = []) so callers still get an empty list when no images are provided; update any internal references to images accordingly (function name: tokenize_with_weights).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@comfy_extras/nodes_joyimage.py`:
- Around line 62-75: In execute, ensure we don't mismatch single-image
multimodal tokens with batch latents: validate that the input image batch size
is 1 and raise/return an error if image.shape[0] > 1 (or alternatively implement
per-image conditioning), because clip.tokenize(prompt, images=[resized_image])
creates a single image token while vae.encode(resized_image) returns a batch of
latents; update the logic around clip.tokenize, vae.encode, and
node_helpers.conditioning_set_values to either enforce batch==1 early (with a
clear error) or iterate per-image to produce matching conditioning entries for
each image.
In `@comfy/model_base.py`:
- Around line 2228-2243: The current sequence builds a temporary 6D tensor
(stacked/rotated) then reshapes to 5D (variables stacked, rotated, flat),
causing extra full-volume allocations; instead, after converting refs to
device/dtype (the ref_5d list built from ref_latents using r.to(device=device,
dtype=dtype)), directly concatenate the noise xc and the refs along the temporal
channel (the T dimension) to produce the final 5D tensor without creating
stacked/rotated intermediates—i.e., replace the torch.stack/permute/reshape
dance that produces rotated and flat with a single torch.cat of [xc, *ref_5d]
along the temporal axis so b, c, n*t_noise, h, w is produced in one allocation.
In `@comfy/model_detection.py`:
- Around line 823-838: In the JoyImage autodetect branch inside
model_detection.py, handle RMSNorm params named either "weight" or "scale"
instead of hardcoding '{}double_blocks.0.attn.img_attn_q_norm.weight'; modify
the lookup that computes head_dim to check for
'{}double_blocks.0.attn.img_attn_q_norm.weight' and fall back to
'{}double_blocks.0.attn.img_attn_q_norm.scale' (using key_prefix) so head_dim is
derived from whichever exists, and similarly mirror this tolerant lookup pattern
where '{}condition_embedder.text_embedder.linear_1.weight' or other RMSNorm-like
keys are referenced (e.g., in the head_dim and text_dim computations) to avoid
KeyError for checkpoints that store .scale.
In `@comfy/sd.py`:
- Around line 1420-1422: The current conditional that returns
TEModel.QWEN3VL_8B_JOYIMAGE when keys like
"model.language_model.layers.0.self_attn.q_norm.weight" and
"model.visual.patch_embed.proj.weight" are present is too generic and causes any
Qwen3‑VL checkpoint to be routed to JoyImage; update the check in the matching
logic (the block that returns TEModel.QWEN3VL_8B_JOYIMAGE) to either return a
generic TEModel (e.g., TEModel.QWEN3VL_8B) or add an additional guard so it only
returns the JoyImage variant when clip_type == CLIPType.JOYIMAGE (or when a
JoyImage‑specific prefix/key is present), and ensure
load_text_encoder_state_dicts() / JoyImageTokenizer / joyimage.te decoding paths
are only used when that JoyImage guard is true.
---
Nitpick comments:
In `@comfy/text_encoders/joyimage.py`:
- Around line 94-95: The tokenize_with_weights function uses a mutable default
parameter images=[], which can cause subtle bugs; change the signature of
tokenize_with_weights to use images=None and add a guard at the top of the
function (if images is None: images = []) so callers still get an empty list
when no images are provided; update any internal references to images
accordingly (function name: tokenize_with_weights).
In `@nodes.py`:
- Line 982: Update the DESCRIPTION string constant to document the new CLIP
type(s): add an entry for "joyimage" specifying it uses the Qwen3-VL text
encoder with text_dim=4096 (and optionally note any other relevant config), and
ensure "pixeldit" remains documented (it already appears as "pixeldit: gemma 2
2B elm"); locate and modify the DESCRIPTION variable to append or insert the
"joyimage: qwen3-vl (text_dim=4096)" line under the [Recipes] section so users
can see the expected text-encoder configuration.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 410466e3-b042-4904-9386-d23cf3bfb680
📥 Commits
Reviewing files that changed from the base of the PR and between 822aca1 and f1b33782d9b293f8d449334d8725c7599cbc54af.
📒 Files selected for processing (9)
comfy/ldm/joyimage/model.pycomfy/model_base.pycomfy/model_detection.pycomfy/sd.pycomfy/supported_models.pycomfy/text_encoders/joyimage.pycomfy/text_encoders/qwen3_vl.pycomfy_extras/nodes_joyimage.pynodes.py
| '{}double_blocks.0.attn.img_attn_qkv.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)] | ||
| dit_config = {} | ||
| dit_config["image_model"] = "joyimage" | ||
| dit_config["in_channels"] = img_in.shape[1] | ||
| dit_config["hidden_size"] = img_in.shape[0] | ||
| dit_config["patch_size"] = list(img_in.shape[2:]) | ||
| dit_config["num_layers"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.') | ||
| head_dim = state_dict['{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix)].shape[0] | ||
| dit_config["num_attention_heads"] = dit_config["hidden_size"] // head_dim | ||
| # text_dim from the text-embedder input projection | ||
| dit_config["text_dim"] = state_dict['{}condition_embedder.text_embedder.linear_1.weight'.format(key_prefix)].shape[1] |
There was a problem hiding this comment.
Accept RMSNorm scale keys here too.
This branch hardcodes img_attn_q_norm.weight, but the rest of model_detection.py already treats RMSNorm-style params as weight or scale. A JoyImage checkpoint saved with an ops implementation that serializes .scale will enter this branch and then fail with a KeyError on Line 835, so native auto-detection becomes loader-breaking for that variant.
As per coding guidelines comfy/**: focus on backward compatibility (breaking changes affect all custom nodes).
🛠️ Proposed fix
dit_config["hidden_size"] = img_in.shape[0]
dit_config["patch_size"] = list(img_in.shape[2:])
dit_config["num_layers"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.')
- head_dim = state_dict['{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix)].shape[0]
+ q_norm_key = next(
+ (
+ f'{key_prefix}double_blocks.0.attn.img_attn_q_norm.{suffix}'
+ for suffix in ("weight", "scale")
+ if f'{key_prefix}double_blocks.0.attn.img_attn_q_norm.{suffix}' in state_dict_keys
+ ),
+ None,
+ )
+ if q_norm_key is None:
+ return None
+ head_dim = state_dict[q_norm_key].shape[0]
dit_config["num_attention_heads"] = dit_config["hidden_size"] // head_dim
# text_dim from the text-embedder input projection
dit_config["text_dim"] = state_dict['{}condition_embedder.text_embedder.linear_1.weight'.format(key_prefix)].shape[1]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| '{}double_blocks.0.attn.img_attn_qkv.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)] | |
| dit_config = {} | |
| dit_config["image_model"] = "joyimage" | |
| dit_config["in_channels"] = img_in.shape[1] | |
| dit_config["hidden_size"] = img_in.shape[0] | |
| dit_config["patch_size"] = list(img_in.shape[2:]) | |
| dit_config["num_layers"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.') | |
| head_dim = state_dict['{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix)].shape[0] | |
| dit_config["num_attention_heads"] = dit_config["hidden_size"] // head_dim | |
| # text_dim from the text-embedder input projection | |
| dit_config["text_dim"] = state_dict['{}condition_embedder.text_embedder.linear_1.weight'.format(key_prefix)].shape[1] | |
| dit_config["image_model"] = "joyimage" | |
| dit_config["in_channels"] = img_in.shape[1] | |
| dit_config["hidden_size"] = img_in.shape[0] | |
| dit_config["patch_size"] = list(img_in.shape[2:]) | |
| dit_config["num_layers"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.') | |
| q_norm_key = next( | |
| ( | |
| f'{key_prefix}double_blocks.0.attn.img_attn_q_norm.{suffix}' | |
| for suffix in ("weight", "scale") | |
| if f'{key_prefix}double_blocks.0.attn.img_attn_q_norm.{suffix}' in state_dict_keys | |
| ), | |
| None, | |
| ) | |
| if q_norm_key is None: | |
| return None | |
| head_dim = state_dict[q_norm_key].shape[0] | |
| dit_config["num_attention_heads"] = dit_config["hidden_size"] // head_dim | |
| # text_dim from the text-embedder input projection | |
| dit_config["text_dim"] = state_dict['{}condition_embedder.text_embedder.linear_1.weight'.format(key_prefix)].shape[1] |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy/model_detection.py` around lines 823 - 838, In the JoyImage autodetect
branch inside model_detection.py, handle RMSNorm params named either "weight" or
"scale" instead of hardcoding '{}double_blocks.0.attn.img_attn_q_norm.weight';
modify the lookup that computes head_dim to check for
'{}double_blocks.0.attn.img_attn_q_norm.weight' and fall back to
'{}double_blocks.0.attn.img_attn_q_norm.scale' (using key_prefix) so head_dim is
derived from whichever exists, and similarly mirror this tolerant lookup pattern
where '{}condition_embedder.text_embedder.linear_1.weight' or other RMSNorm-like
keys are referenced (e.g., in the head_dim and text_dim computations) to avoid
KeyError for checkpoints that store .scale.
Source: Coding guidelines
JoyImageEdit is an image-edit diffusion transformer from JD (jd-opensource),
Apache 2.0. This adds native ComfyUI support so it loads and runs like other
edit models (load checkpoint -> TextEncode + ReferenceLatent -> KSampler ->
VAEDecode), with no diffusers dependency.
Architecture:
- Transformer (comfy/ldm/joyimage/model.py): dual-stream (img/txt) DiT with a
Conv3d patch embed (patch_size [1,2,2]), Wan-style learnable modulation,
and 3D RoPE (rope_dim_list [16,56,56]). All attention goes through
comfy.ldm.modules.attention.optimized_attention.
- Text encoder (comfy/text_encoders/{qwen3_vl,joyimage}.py): a reusable
Qwen3-VL multimodal stack (vision tower + LM) in qwen3_vl.py, plus a thin
JoyImage-specific layer (prompt templates, drop_idx, tokenizer, te() factory)
in joyimage.py that depends on it. text_dim 4096.
- VAE: reuses the existing Wan 2.1 latent format (AutoencoderKLWan), no new
latent format.
- Edit conditioning: reuses the reference_latents mechanism. Reference and
noise latents are stacked on a new n-slot dimension and rotated at the model
boundary (model_base.JoyImage), so the transformer stays 5D-in/5D-out.
Guidance-rescale is built into the CFG path.
Model wiring:
- model_base.JoyImage uses ModelType.FLOW with sampling_settings
multiplier=1000 (the time embedding is trained on t in [0,1000]) and
shift=1.5; FLOW's linear time_snr_shift matches the diffusers
FlowMatchEuler sigma schedule.
- model_detection sniffs the transformer state-dict (double_blocks.*,
condition_embedder.*, 5D img_in Conv3d) to route image_model="joyimage".
- supported_models.JoyImage and the CLIPLoader "joyimage" type register it.
User-facing node TextEncodeJoyImageEdit (comfy_extras/nodes_joyimage.py)
bucket-resizes the input image to the nearest 1024-base bucket, encodes the
prompt with the image, and emits both the conditioning and the bucketed image
so the same pixels feed VAEEncode and the negative encode (JoyImage requires
noise and reference latents to share spatial dims).
Upstream merged native Qwen3-VL support (Comfy-Org#14298), adding comfy/text_encoders/qwen3vl.py plus helpers in qwen_vl.py / llama.py / qwen35.py. The JoyImage port previously shipped its own duplicate Qwen3-VL implementation (comfy/text_encoders/qwen3_vl.py); that duplication is now removed and the JoyImage text encoder rides on the upstream stack. - Delete comfy/text_encoders/qwen3_vl.py. - Rewrite comfy/text_encoders/joyimage.py to subclass upstream comfy.text_encoders.qwen3vl. The JoyImage checkpoint is a stock qwen3vl_8b, so only JoyImage-specific behavior is overridden: * Qwen3VL8B_JoyImage.forward builds the 3D MRoPE position ids and injects deepstack visual features on the conditioning path. Upstream Qwen3VL only does this inside generate() via build_image_inputs; SDClipModel.forward never passes those kwargs. The JoyImage node feeds an image through the encoder (clip.tokenize(prompt, images=[..])), so the override reuses build_image_inputs to reproduce the multimodal conditioning that Llama2_.forward already accepts kwargs for. * preprocess_embed keeps JoyImage's bicubic+clamp image preprocessing (process_qwen3vl_image) instead of upstream's bilinear path, to preserve validated DiT numerics. * JoyImageTokenizer keeps the JoyImage system-prompt templates, suppresses the Qwen3 <think> block, and raises on image-placeholder count mismatch. * JoyImageTEModel keeps the drop_idx=34 system-prompt strip and the pre-final-norm layer tap (layer="hidden", layer_idx=-1). - sd.py QWEN3VL_8B_JOYIMAGE branch: apply the same state-dict prefix remap the sibling QWEN3VL branch uses (model.language_model.->model., model.visual.->visual., lm_head.->model.lm_head.) so the checkpoint loads into the upstream Qwen3VL namespace, then use the module-level llama_detect. Detection ordering is preserved: the JoyImage discriminator is checked before the generic Qwen3-VL deepstack key. No changes to llama.py / qwen3vl.py / qwen_vl.py / qwen35.py.
f1b3378 to
e96bd48
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@nodes.py`:
- Line 982: The DESCRIPTION string in the CLIPLoader class is missing a recipe
hint for joyimage, which was added as a selectable type option. Add a new line
to the DESCRIPTION variable that includes the joyimage recipe information in the
same format as the other recipes listed (e.g., "joyimage: [model name]"),
ensuring the UI guidance matches the new type option that was introduced.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 28a7cc96-3679-4b4e-866c-561dbbd44eb4
📥 Commits
Reviewing files that changed from the base of the PR and between f1b33782d9b293f8d449334d8725c7599cbc54af and e96bd48.
📒 Files selected for processing (8)
comfy/ldm/joyimage/model.pycomfy/model_base.pycomfy/model_detection.pycomfy/sd.pycomfy/supported_models.pycomfy/text_encoders/joyimage.pycomfy_extras/nodes_joyimage.pynodes.py
🚧 Files skipped from review as they are similar to previous changes (6)
- comfy/supported_models.py
- comfy/model_detection.py
- comfy_extras/nodes_joyimage.py
- comfy/sd.py
- comfy/model_base.py
- comfy/ldm/joyimage/model.py
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@nodes.py`:
- Line 982: The DESCRIPTION string in the CLIPLoader class is missing a recipe
hint for joyimage, which was added as a selectable type option. Add a new line
to the DESCRIPTION variable that includes the joyimage recipe information in the
same format as the other recipes listed (e.g., "joyimage: [model name]"),
ensuring the UI guidance matches the new type option that was introduced.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 28a7cc96-3679-4b4e-866c-561dbbd44eb4
📥 Commits
Reviewing files that changed from the base of the PR and between f1b33782d9b293f8d449334d8725c7599cbc54af and e96bd48.
📒 Files selected for processing (8)
comfy/ldm/joyimage/model.pycomfy/model_base.pycomfy/model_detection.pycomfy/sd.pycomfy/supported_models.pycomfy/text_encoders/joyimage.pycomfy_extras/nodes_joyimage.pynodes.py
🚧 Files skipped from review as they are similar to previous changes (6)
- comfy/supported_models.py
- comfy/model_detection.py
- comfy_extras/nodes_joyimage.py
- comfy/sd.py
- comfy/model_base.py
- comfy/ldm/joyimage/model.py
🛑 Comments failed to post (1)
nodes.py (1)
982-982:
⚠️ Potential issue | 🟡 Minor | ⚡ Quick winAdd the missing JoyImage recipe hint in CLIPLoader description.
"joyimage"was added as a selectabletype(Line 972), but the recipes text omits it. Please add a short JoyImage line so the UI guidance matches the new option.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@nodes.py` at line 982, The DESCRIPTION string in the CLIPLoader class is missing a recipe hint for joyimage, which was added as a selectable type option. Add a new line to the DESCRIPTION variable that includes the joyimage recipe information in the same format as the other recipes listed (e.g., "joyimage: [model name]"), ensuring the UI guidance matches the new type option that was introduced.
|
Update: rebased onto latest master, now reuses the merged Qwen3-VL stack Heads-up on a significant change in this PR. When I first opened it, I didn’t realize that @kijai was implementing Qwen3-VL, so this PR shipped its own implementation of the text encoder. Today #14298 ("Support text generation with Qwen3-VL") was merged, adding an official Qwen3-VL stack (comfy/text_encoders/qwen3vl.py, qwen_vl.py, plus the Qwen3-VL configs and deepstack support in llama.py/qwen35.py). I've rebased onto current master and reworked the text encoder to build on that merged stack instead of duplicating it:
No changes to the upstream qwen3vl.py / qwen_vl.py / llama.py / qwen35.py files. Happy to adjust if you'd prefer a different layout. @kijai @comfyanonymous Could someone please take a look at this PR when you have time? Any review or feedback would be greatly appreciated! |
|
Hey, thank you for the PR, and sorry for slow response. I'm trying to find time to look at it, this month has just been extremely busy with many model releases with already pretty big backlog. |
…forward) JoyImageEditPlus is the multi-image (1-6 reference images) variant of JoyImageEdit, trained from the same base. Its diffusers transformer shares byte-identical weight structure with the single-image variant (894 keys, zero rename) but injects references differently: instead of the single-image slot-stack (stack refs + noise into a 6D tensor and rotate on the frame dim, which forces all items to share resolution), each reference is independently patchified and concatenated on the sequence dim with per-image temporal-offset 3D RoPE, allowing references at different resolutions. Since the single-image port is not yet upstream, this unifies both variants onto the Plus-style forward rather than keeping two paths; single-image is now the ref=1 special case. Verified numerically: at ref=1 with equal resolution the new path's RoPE is bit-identical to the old slot-stack layout, and the transformer output matches the diffusers Plus reference (fp32, incl. the different-resolution case). ComfyUI runs cond/uncond in one forward with a shared reference configuration, so the diffusers Plus batched RoPE, padding attention_mask, and dedicated attention processor are unnecessary here: the unified forward reuses the existing unbatched _apply_rotary_emb and JoyImageAttention. Confirmed equivalent to the diffusers batched+mask path for a single sample. - comfy/ldm/joyimage/model.py: forward takes ref_latents and builds components=[target, ref0, ...]; per-component patchify + temporal-offset RoPE; output keeps only the target segment. Old single-grid RoPE removed. - comfy/model_base.py: JoyImage drops the slot-stack / frame-rotation / shape-equality path in _apply_model, passing ref_latents straight to the transformer. Guidance-rescale and the reference_latents requirement are kept. - comfy/text_encoders/joyimage.py: the image template emits one vision block per reference (N = image count); N=1 is byte-for-byte the old template. - comfy_extras/nodes_joyimage.py: add TextEncodeJoyImageEditPlus with optional image1..image6 inputs, each bucket-resized and VAE-encoded into the reference_latents list. Detection, supported_models, and sd.py need no changes: the identical weight structure routes both variants through image_model="joyimage".
There was a problem hiding this comment.
♻️ Duplicate comments (1)
comfy_extras/nodes_joyimage.py (1)
114-144: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winSame per-image batch mismatch risk as the single-image node, now multiplied across 6 slots.
Each supplied image is independently
vae.encode'd into oneref_latentsentry, butclip.tokenize(prompt, images=resized_images)emits exactly one vision block per list entry. If any supplied image hasshape[0] > 1, that slot's latent batch won't align with its single vision-block token, desyncingreference_latentsfrom the prompt's multimodal tokens — the same issue flagged on the single-image node'sexecute(lines 62-75), just recurring here per-image.Consider validating
image.shape[0] == 1for each supplied image before encoding (or implement per-image batch iteration) to avoid silently corrupting multi-image conditioning.🐛 Proposed guard
for image in supplied: + if image.shape[0] != 1: + raise ValueError("TextEncodeJoyImageEditPlus does not support batched images; each reference image input must have batch size 1.") samples = image.movedim(-1, 1)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@comfy_extras/nodes_joyimage.py` around lines 114 - 144, The execute method in the JoyImage node builds one reference latent per supplied image and one vision token block per resized image, so multi-batch inputs can silently misalign conditioning. Update execute to either validate that each supplied image has batch size 1 before calling vae.encode and clip.tokenize, or explicitly iterate over batched images so reference_latents and the multimodal tokens stay in sync. Use the existing execute, ref_latents, resized_images, and clip.tokenize flow to place the guard where supplied images are processed.
🧹 Nitpick comments (2)
comfy/ldm/joyimage/model.py (1)
497-499: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winSlice dropped reference tokens before final projection.
norm_outandproj_outrun on all target + reference tokens, but reference tokens are discarded immediately after. Slice totarget_tokensfirst to avoid redundant GPU work in the sampling hot path.Proposed refactor
- img = self.proj_out(self.norm_out(img)) target_tokens = tt * th * tw img = img[:, :target_tokens, :] + img = self.proj_out(self.norm_out(img)) img = self.unpatchify(img, tt, th, tw)As per path instructions,
comfy/**is core ML/diffusion code where performance implications in hot paths should be prioritized.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@comfy/ldm/joyimage/model.py` around lines 497 - 499, The sampling path in `JoyImage` is doing redundant work by applying `norm_out` and `proj_out` to both target and reference tokens before immediately slicing away the reference tokens. Update the token handling in the `img` flow so `target_tokens` is computed and the slice is applied before the final projection steps, using the existing `norm_out` and `proj_out` sequence in `model.py` to limit work to only the target tokens.Source: Path instructions
comfy_extras/nodes_joyimage.py (1)
90-90: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
MAX_IMAGESconstant is unused.Declared but never referenced —
define_schema's six image inputs andexecute's six parameters are hardcoded independently of it. Either wire it in (e.g., generate inputs/params in a loop) or drop it to avoid a stale "source of truth" that can silently diverge from the actual input count.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@comfy_extras/nodes_joyimage.py` at line 90, The MAX_IMAGES constant in the node class is currently stale because define_schema and execute still hardcode six image slots independently. Either remove MAX_IMAGES from the JoyImage node if it is not intended to drive behavior, or refactor define_schema and execute to derive their image inputs and parameters from MAX_IMAGES so there is a single source of truth. Use the JoyImage class and its define_schema and execute methods to locate the update.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Duplicate comments:
In `@comfy_extras/nodes_joyimage.py`:
- Around line 114-144: The execute method in the JoyImage node builds one
reference latent per supplied image and one vision token block per resized
image, so multi-batch inputs can silently misalign conditioning. Update execute
to either validate that each supplied image has batch size 1 before calling
vae.encode and clip.tokenize, or explicitly iterate over batched images so
reference_latents and the multimodal tokens stay in sync. Use the existing
execute, ref_latents, resized_images, and clip.tokenize flow to place the guard
where supplied images are processed.
---
Nitpick comments:
In `@comfy_extras/nodes_joyimage.py`:
- Line 90: The MAX_IMAGES constant in the node class is currently stale because
define_schema and execute still hardcode six image slots independently. Either
remove MAX_IMAGES from the JoyImage node if it is not intended to drive
behavior, or refactor define_schema and execute to derive their image inputs and
parameters from MAX_IMAGES so there is a single source of truth. Use the
JoyImage class and its define_schema and execute methods to locate the update.
In `@comfy/ldm/joyimage/model.py`:
- Around line 497-499: The sampling path in `JoyImage` is doing redundant work
by applying `norm_out` and `proj_out` to both target and reference tokens before
immediately slicing away the reference tokens. Update the token handling in the
`img` flow so `target_tokens` is computed and the slice is applied before the
final projection steps, using the existing `norm_out` and `proj_out` sequence in
`model.py` to limit work to only the target tokens.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 5caad4a0-a646-4089-94e0-ab8ff4045516
📒 Files selected for processing (8)
comfy/ldm/joyimage/model.pycomfy/model_base.pycomfy/model_detection.pycomfy/sd.pycomfy/supported_models.pycomfy/text_encoders/joyimage.pycomfy_extras/nodes_joyimage.pynodes.py
🚧 Files skipped from review as they are similar to previous changes (5)
- comfy/supported_models.py
- comfy/model_detection.py
- nodes.py
- comfy/sd.py
- comfy/text_encoders/joyimage.py
Update: Add support for JoyAI-Image-Edit-Plus.Fixed #14644. JoyImageEditPlus is the multi-image variant of JoyImageEdit (supporting 1–6 reference images), trained on the same base model. Its Diffusers Transformer shares weight structure with the single-image version, but differs in how reference images are injected. In JoyImageEditPlus, each reference image is independently patchified and then concatenated along the sequence dimension. A per-image temporal offset 3D RoPE is applied, allowing reference images to have different resolutions. Since the single-image port has not yet been merged upstream, this update unifies both modes under the Plus-style forward pass, removing the need for two separate branches. The single-image mode is now treated as a special case where the number of reference images equals 1. Acknowledgements. Thanks to the community efforts that explored bringing JoyAI-Image into ComfyUI ahead of native support (apologies if anyone is missed):
Exampleworkflow: workflow_joyimage_edit_plus.json Prompt: "The woman is lovingly holding the cute puppy in her arms"
|
|
repacked checkpoints: https://huggingface.co/jdopensource/JoyAI-Image-Edit-Plus-ComfyUI |
You can quant this with https://github.com/Comfy-Org/comfy-model-tools/blob/main/quant_int8_convrot.py Tested on the previous model, considerable speed boost: 3.88s/it to 2.25s/it on a 4090 Also took the liberty to add the standard forward wrapping to pass through transformer options for attention overrides etc, sageattention seems to work fine bringing the speed down to 1.86s/it on 4090, the model was pretty slow so this should help testing. |
I’d much prefer that you host it. What we’re doing right now is just to lower the barrier to using the model before native support is fully available in ComfyUI |
Can of course do that, I'm personally going to be away for over a week starting tomorrow so I can't promise to have time to do much with this still. There's a new AGENTS.md for some best practices for model implementations, if you could refer to that with AI agent or manually that would also help. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@comfy/ldm/joyimage/model.py`:
- Around line 441-456: The shared default transformer_options in forward and
_forward can leak mutated block metadata across runs and concurrent sampling.
Update comfy.ldm.joyimage.model.JoyImageModel.forward and _forward to accept
None instead of a mutable default, then create a fresh per-call dict before
passing it to get_all_wrappers and WrapperExecutor.execute so each invocation
gets isolated transformer_options state.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 8bfdb4e6-bd49-4d7f-abeb-c56476baeec9
📒 Files selected for processing (2)
comfy/ldm/joyimage/model.pycomfy/model_base.py
🚧 Files skipped from review as they are similar to previous changes (1)
- comfy/model_base.py
ok I will hand it to my claude code |
…-VL keys JoyImage text_encoder now detects as QWEN3VL_8B and is selected by clip_type == JOYIMAGE. Also copy transformer_options in JoyImage._apply_model like the base method, so the per-block bookkeeping written during forward no longer leaks into the sampler's shared dict.
Move JoyImage CFG guidance rescale to a JoyImageGuidanceRescale node that clones the model and calls set_model_sampler_cfg_function, following the RenormCFG (nodes_lumina2.py) precedent for model-specific guidance nodes.
Some cleanup of the JoyImage transformer: - FP32LayerNorm/JoyImageModulate: drop constructor params that were never used (dtype/device on the param-free FP32LayerNorm, operations on JoyImageModulate). - modulate_table: init with torch.empty instead of torch.zeros since it is loaded from the state dict, and cast at use with comfy.ops.cast_to_input. - Replace the FeedForward nn.Dropout(0.0) no-op with nn.Identity, keeping the ModuleList slot so state-dict keys are unchanged. - Drop the always-true vec.shape guard (time_proj_dim is always hidden_size*6). - Remove the redundant per-reference .to(device,dtype) in JoyImage._apply_model; the extra_conds cast loop already moves ref_latents to the compute dtype/device. - Drop the redundant .to(xq.device) in the RoPE apply; the freqs are built on the latent device.
Replace custom FP32LayerNorm with the standard affine-free operations.LayerNorm.
JoyImage differs from the standard process_qwen2vl_images only in using bicubic interpolation and post-resize clamping. I add interpolation and clamp parameters to the shared helper function, allowing JoyImage to directly reuse process_qwen2vl_images without a duplicated implementation. (note for reviewer: this modifies qwen_vl.py. If you feel this approach is not appropriate, we can discuss alternative implementations.)
Unlike TextEncodeQwenImageEdit, the JoyImage nodes resize to a discrete bucket, so the target size can't be recomputed downstream. Keep the image output and document that it feeds VAEEncode for a matching-size init latent.
|
@kijai Hello, I followed the philosophy in AGENTS.md, pushed a round of cleanup and review fixes. 56f9142 Detect JoyImage TE by clip_type, not shared keys 1ae7a81 Move CFG guidance rescale into a node 0eafd9c Dead code / redundant casts in the transformer 3e42225 Use operations.LayerNorm 0c18c50 Reuse process_qwen2vl_images for preprocessing b0eb165 Note on why the text-encode nodes emit the bucketed image Now the workflows changed: |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
comfy_extras/nodes_joyimage.py (1)
63-79: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the shared bucket-resize-encode logic.
TextEncodeJoyImageEdit.execute(Lines 63-75) and the per-image loop body inTextEncodeJoyImageEditPlus.execute(Lines 129-137) both perform the identical sequence:movedim(-1,1)→_find_best_bucket→common_upscale→movedim(1,-1)[..., :3]. A shared helper (e.g._bucket_resize(image) -> resized_image) would remove the duplication and keep future bucket-logic tweaks (e.g. a different upscale method) in one place.♻️ Proposed refactor
+def _bucket_resize(image): + samples = image.movedim(-1, 1) + src_h, src_w = samples.shape[2], samples.shape[3] + bucket_h, bucket_w = _find_best_bucket(src_h, src_w) + resized = comfy.utils.common_upscale(samples, bucket_w, bucket_h, "bilinear", "center") + return resized.movedim(1, -1)[:, :, :, :3] + class TextEncodeJoyImageEdit(io.ComfyNode): ... `@classmethod` def execute(cls, clip, prompt, vae, image) -> io.NodeOutput: - samples = image.movedim(-1, 1) - src_h, src_w = samples.shape[2], samples.shape[3] - bucket_h, bucket_w = _find_best_bucket(src_h, src_w) - - resized = comfy.utils.common_upscale(samples, bucket_w, bucket_h, "bilinear", "center") - resized_image = resized.movedim(1, -1)[:, :, :, :3] + resized_image = _bucket_resize(image)for image in supplied: - samples = image.movedim(-1, 1) - src_h, src_w = samples.shape[2], samples.shape[3] - bucket_h, bucket_w = _find_best_bucket(src_h, src_w) - - resized = comfy.utils.common_upscale(samples, bucket_w, bucket_h, "bilinear", "center") - resized_image = resized.movedim(1, -1)[:, :, :, :3] + resized_image = _bucket_resize(image) resized_images.append(resized_image) ref_latents.append(vae.encode(resized_image))Also applies to: 117-147
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@comfy_extras/nodes_joyimage.py` around lines 63 - 79, The bucket-resize steps are duplicated in TextEncodeJoyImageEdit.execute and the per-image loop inside TextEncodeJoyImageEditPlus.execute, so extract that repeated movedim/_find_best_bucket/common_upscale/movedim[:3] flow into a shared helper such as a private bucket resize function. Update both execute paths to call the helper and keep the existing tokenization, encoding, and latent handling unchanged.
♻️ Duplicate comments (2)
comfy/ldm/joyimage/model.py (2)
422-445: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winPreviously flagged Major issue still not resolved — mutable default lives on.
The earlier review requested swapping
transformer_options={}for aNonesentinel in bothforwardand_forwardto avoid a shared dict leaking mutated block metadata (total_blocks,block_type,block_index) across concurrent runs. That fix hasn't landed in this revision.As per path instructions, "Do not add unnecessary `try`/`except` blocks..." is unrelated here, but the prior review already flagged this exact mutable-default risk in this file.🔒 Re-proposed fix
def forward( self, hidden_states: torch.Tensor, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, ref_latents=None, - transformer_options={}, + transformer_options=None, **kwargs, ) -> torch.Tensor: + if transformer_options is None: + transformer_options = {} return comfy.patcher_extension.WrapperExecutor.new_class_executor( @@ def _forward( self, hidden_states: torch.Tensor, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, ref_latents=None, - transformer_options={}, + transformer_options=None, **kwargs, ) -> torch.Tensor: + if transformer_options is None: + transformer_options = {}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@comfy/ldm/joyimage/model.py` around lines 422 - 445, The mutable default for transformer_options is still present in both JoyImageModel.forward and JoyImageModel._forward, so replace the shared {} default with a None sentinel and initialize a fresh dict inside each method before passing it into WrapperExecutor.new_class_executor or downstream logic. Keep the fix scoped to these two methods so block metadata like total_blocks, block_type, and block_index cannot leak between calls.
194-201: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winSame shared-mutable-default gremlin haunts the block's
forwardtoo.
transformer_options={}here is the identical anti-pattern already flagged as Major forJoyImageTransformer3DModel.forward/_forward— a shared dict object reused across every call/instance when the caller omits the argument. Concurrent sampling runs could cross-contaminate block metadata.🔧 Suggested fix
def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, image_rotary_emb: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]] = None, - transformer_options={}, + transformer_options=None, ) -> Tuple[torch.Tensor, torch.Tensor]: + if transformer_options is None: + transformer_options = {}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@comfy/ldm/joyimage/model.py` around lines 194 - 201, The Block.forward method has the same shared-mutable-default issue as the transformer methods: transformer_options is initialized to a single dict object and can leak state across calls. Update the Block.forward signature and any related call sites so transformer_options defaults to None and is replaced with a fresh dict inside the method, following the same pattern used in JoyImageTransformer3DModel.forward/_forward.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@comfy_extras/nodes_joyimage.py`:
- Around line 63-79: The bucket-resize steps are duplicated in
TextEncodeJoyImageEdit.execute and the per-image loop inside
TextEncodeJoyImageEditPlus.execute, so extract that repeated
movedim/_find_best_bucket/common_upscale/movedim[:3] flow into a shared helper
such as a private bucket resize function. Update both execute paths to call the
helper and keep the existing tokenization, encoding, and latent handling
unchanged.
---
Duplicate comments:
In `@comfy/ldm/joyimage/model.py`:
- Around line 422-445: The mutable default for transformer_options is still
present in both JoyImageModel.forward and JoyImageModel._forward, so replace the
shared {} default with a None sentinel and initialize a fresh dict inside each
method before passing it into WrapperExecutor.new_class_executor or downstream
logic. Keep the fix scoped to these two methods so block metadata like
total_blocks, block_type, and block_index cannot leak between calls.
- Around line 194-201: The Block.forward method has the same
shared-mutable-default issue as the transformer methods: transformer_options is
initialized to a single dict object and can leak state across calls. Update the
Block.forward signature and any related call sites so transformer_options
defaults to None and is replaced with a fresh dict inside the method, following
the same pattern used in JoyImageTransformer3DModel.forward/_forward.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: eabf544e-eca2-4cf1-b7c7-48a0e8d75cf7
📒 Files selected for processing (6)
comfy/ldm/joyimage/model.pycomfy/model_base.pycomfy/sd.pycomfy/text_encoders/joyimage.pycomfy/text_encoders/qwen_vl.pycomfy_extras/nodes_joyimage.py
|
✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged. |
|
I have read and agree to the Contributor License Agreement |
|
Thank you for your review! We really appreciate you taking the time to support our models. This allows our users to directly use the JoyImage series in ComfyUI, and also enables them to leverage ComfyUI’s extensive optimization support, such as quantization. We will continue contributing to the open-source community, and please look forward to our upcoming open-source models! |
3 participants
We are the JoyAI team, and this update provides native model support for JoyAI-Image-Edit.
We have received many requests to adapt JoyImageEdit to ComfyUI, and we are also eager to expand the influence of our model, so we have decided to proactively adapt JoyImageEdit. Since we found no existing implementation of the Qwen3VL model that JoyImageEdit depends on in ComfyUI, we have also implemented a Qwen3VL module here.
Fixes issue #10207 #13269
Model Overview
JoyAI-Image is a unified multimodal foundation model that supports image understanding, text-to-image generation, and instruction-driven image editing. The model combines an 8-billion-parameter multimodal large language model (MLLM) with a 16-billion-parameter multimodal diffusion Transformer (MMDiT). The core design philosophy of JoyAI-Image is to achieve closed-loop collaboration among understanding, generation, and editing: stronger spatial understanding enhances controllable generation and precise editing through better scene parsing, relationship localization, and instruction decomposition; generative transformations such as viewpoint conversion can provide supplementary evidence for spatial reasoning.
GitHub homepage: https://github.com/jd-opensource/JoyAI-Image
Released JoyAI-Image-Edit model weights for ComfyUI: https://huggingface.co/jdopensource/JoyAI-Image-Edit-ComfyUI
Key Features
Testing
workflow_joyimage_edit.json
Sampling settings for the reference image: seed=42, steps=40, cfg=4.0, sampler=euler, scheduler=simple, denoise=1.0, flow shift=1.5 (do not insert a ModelSampling node, it would force multiplier=1.0).
Files
New:
comfy/ldm/joyimage/model.pycomfy/text_encoders/joyimage.pycomfy/text_encoders/qwen3_vl.pycomfy_extras/nodes_joyimage.pyModified:
comfy/model_base.py,comfy/supported_models.py,comfy/model_detection.py,comfy/sd.py,nodes.pyFor reviewers
Please feel free to point out any issues with the code!