feat: add JoyImage edit plus by tangyanf · Pull Request #14032 · huggingface/diffusers

tangyanf · 2026-06-22T06:54:14Z

Description

We are the JoyAI Team, and this is the Diffusers implementation for the JoyAI-Image-Edit-Plus model.

GitHub Repository: [https://github.com/jd-opensource/JoyAI-Image]
Hugging Face Model: [https://huggingface.co/jdopensource/JoyAI-Image-Edit-Plus-Diffusers]
Original opensource weights: [https://huggingface.co/jdopensource/JoyAI-Image-Edit-Plus]
Fixes#14049

Model Overview

JoyAI-Image-Edit-Plus extends JoyAI-Image-Edit with multi-image editing capabilities. While JoyAI-Image-Edit operates on a single reference image, Edit-Plus accepts multiple reference
images as input and performs instruction-guided editing across them — enabling tasks such as subject composition, style transfer from multiple sources, and multi-view consistent editing.

It combines an 8B Multimodal Large Language Model (MLLM) with a 16B Multimodal Diffusion Transformer (MMDiT), supporting variable-resolution reference images that are independently
encoded and jointly denoised.

Key Features

Multi-Image Input: Accepts multiple reference images with different resolutions, enabling complex editing scenarios that require information from multiple visual sources.
Subject Composition: Combine elements from separate images into a coherent output guided by text instructions (e.g., "Let the person lovingly play with the dog" given separate person
and dog images).
Cross-Image Style Transfer: Apply style or attributes from one reference image to subjects in another.
Variable Resolution Support: Each reference image is independently resized and encoded at its optimal resolution, preserving fine-grained details regardless of input size.
Instruction-Guided Generation: Natural language prompts control how multiple reference images are composed and edited in the final output.

github-actions · 2026-06-22T09:30:53Z

Hi @tangyanf, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

sergereview

🤗 Serge says:

This PR adds the JoyImage Edit Plus model and pipeline. There are several blocking issues that need to be addressed before merging.

Blocking — Debug artifacts left in production code

Multiple torch.save() calls, a print() statement, and a commented-out exit(0) are left in pipeline_joyimage_edit_plus.py. These will write files to the user's working directory and print to stdout during every inference call.

Blocking — `einops` dependency

Per .ai/models.md: "No new mandatory dependency without discussion (e.g. einops). Optional deps guarded with is_X_available() and a dummy in utils/dummy_*.py." The pipeline directly imports from einops import rearrange — this is the only non-comment usage of einops in src/diffusers/. The rearrange calls should be rewritten with native PyTorch (reshape, permute, unflatten).

Blocking — `sglang` integration code in model forward

The transformer's forward method contains sglang-specific code: list-unwrapping for "SglangXvideo CFG branches" (lines 272-276) and a try: from sglang... fallback (lines 279-287). Per .ai/AGENTS.md: "No defensive code, unused code paths, or legacy stubs — do not add fallback paths, safety checks, or configuration options 'just in case'." This code doesn't belong in the diffusers model — the pipeline always passes the required arguments.

Blocking — Missing dummy objects

JoyImageEditPlusTransformer3DModel, JoyImageEditPlusPipeline, and JoyImageEditPlusPipelineOutput are not registered in dummy_pt_objects.py / dummy_torch_and_transformers_objects.py. This will cause ImportError when torch/transformers are not installed.

Blocking — Missing tests

No test files were added for the new model or pipeline.

Blocking — Hardcoded `device_type="cuda"` in `torch.autocast`

torch.autocast(device_type="cuda", ...) is hardcoded in two places in the pipeline. This will fail on MPS, XPU, and other non-CUDA devices.

Non-blocking — Inlined scheduler sigma math

Per .ai/pipelines.md gotcha #3, the pipeline manually computes shifted sigmas and temporarily overrides self.scheduler.shift — this is exactly what FlowMatchEulerDiscreteScheduler does with its shift config. The scheduler should own this logic.

Non-blocking — Unused imports and parameters

import inspect in transformer_joyimage_edit_plus.py is unused.
enable_denormalization parameter is declared in prepare_latents and __call__ but never read.
retrieve_timesteps is duplicated from the existing pipeline without a # Copied from annotation.

serge v0.1.0 · model: claude-opus-4-6 · 29 LLM turns · 50 tool calls · 190.2s · 1602502 in / 7369 out tokens

- Remove einops dependency: replace rearrange with reshape/permute - Remove sglang-specific code from transformer forward - Remove unused import inspect from transformer - Fix hardcoded device_type="cuda" to use device.type - Simplify scheduler sigma math: delegate to retrieve_timesteps - Remove unused enable_denormalization parameter - Fix callback latents variable binding - Fix output_type="pt" to return stacked tensor - Set return_dict default to True in transformer forward - Add dummy objects for JoyImageEditPlus classes - Add transformer and pipeline test files

yiyixuxu · 2026-06-24T18:36:07Z

@claude can you do a review here?

github-actions · 2026-06-24T18:36:29Z

Claude encountered an error —— View job

I'll analyze this and get back to you.

yiyixuxu

thanks, i left some feedbacks

tangyanf · 2026-06-25T02:48:23Z

thanks, i left some feedbacks

@yiyixuxu Thank you for taking the time to review this PR! I've addressed all the feedback — here's a summary of the changes:

Removed cross-imports from transformer_joyimage.py; copied and renamed classes with # Copied from annotations.
Removed **kwargs from JoyImageEditPlusAttnProcessor.call.
Removed cos.ndim == 2 branch in _apply_rotary_emb_batched (only batched path kept).
Made shape_list a required argument, removed None default and ValueError check.
Removed conditional on vec.unflatten.
Replaced _resize_center_crop with self.vae_image_processor.resize_center_crop().
Replaced _get_bucket_size with self.vae_image_processor.get_default_height_width().
Added # Copied from comment to retrieve_timesteps.

Please let me know if there's anything else that needs to be updated!

tangyanf · 2026-06-25T03:01:44Z

🤗 Serge says:
This PR adds the JoyImage Edit Plus model and pipeline. There are several blocking issues that need to be addressed before merging.
Blocking — Debug artifacts left in production code
Multiple torch.save() calls, a print() statement, and a commented-out exit(0) are left in pipeline_joyimage_edit_plus.py. These will write files to the user's working directory and print to stdout during every inference call.
Blocking — einops dependency
Per .ai/models.md: "No new mandatory dependency without discussion (e.g. einops). Optional deps guarded with is_X_available() and a dummy in utils/dummy_*.py." The pipeline directly imports from einops import rearrange — this is the only non-comment usage of einops in src/diffusers/. The rearrange calls should be rewritten with native PyTorch (reshape, permute, unflatten).
Blocking — sglang integration code in model forward
The transformer's forward method contains sglang-specific code: list-unwrapping for "SglangXvideo CFG branches" (lines 272-276) and a try: from sglang... fallback (lines 279-287). Per .ai/AGENTS.md: "No defensive code, unused code paths, or legacy stubs — do not add fallback paths, safety checks, or configuration options 'just in case'." This code doesn't belong in the diffusers model — the pipeline always passes the required arguments.
Blocking — Missing dummy objects
JoyImageEditPlusTransformer3DModel, JoyImageEditPlusPipeline, and JoyImageEditPlusPipelineOutput are not registered in dummy_pt_objects.py / dummy_torch_and_transformers_objects.py. This will cause ImportError when torch/transformers are not installed.
Blocking — Missing tests
No test files were added for the new model or pipeline.
Blocking — Hardcoded device_type="cuda" in torch.autocast
torch.autocast(device_type="cuda", ...) is hardcoded in two places in the pipeline. This will fail on MPS, XPU, and other non-CUDA devices.
Non-blocking — Inlined scheduler sigma math
Per .ai/pipelines.md gotcha #3, the pipeline manually computes shifted sigmas and temporarily overrides self.scheduler.shift — this is exactly what FlowMatchEulerDiscreteScheduler does with its shift config. The scheduler should own this logic.
Non-blocking — Unused imports and parameters
import inspect in transformer_joyimage_edit_plus.py is unused.
enable_denormalization parameter is declared in prepare_latents and __call__ but never read.
retrieve_timesteps is duplicated from the existing pipeline without a # Copied from annotation.
serge v0.1.0 · model: claude-opus-4-6 · 29 LLM turns · 50 tool calls · 190.2s · 1602502 in / 7369 out tokens

re check

tangyanf · 2026-06-25T09:40:10Z

@tarekziade Hi! I've addressed all the feedback from sergereview[bot] in my latest commits. However, the bot's "changes requested" review is still blocking the merge. Could you help me with one of the following:

Dismiss the bot's review now that the changes have been addressed, or
Let me know how to trigger a re-review from the bot so it can re-evaluate the updated code?

Thanks!

…t_to_diffusers.py JoyImage Edit and Edit Plus share identical VAE and transformer weight layouts — only the target model class differs. Consolidate both into a single script with a --model_type flag (edit | edit_plus) instead of maintaining two nearly-duplicate files.

resolved

Replace `torch.autocast(dtype=torch.float32)` + `.float()` with `.to(self.vae.dtype)` for both VAE encode and decode calls. The previous approach caused dtype mismatch (float32 input vs bfloat16 bias) on CPU where autocast does not automatically cast conv weights, breaking the CI `test_layerwise_casting_inference` test.

tangyanf · 2026-07-03T07:25:56Z

Hi @yiyixuxu @dg845 , all review comments have been addressed and CI is green. Happy to make any further changes if needed. Let me know if there is anything on my end needed for the merge.
Thanks!

dg845

Thanks for the PR! Left some design comments :).

tarekziade · 2026-07-03T11:43:54Z

@tarekziade Hi! I've addressed all the feedback from sergereview[bot] in my latest commits. However, the bot's "changes requested" review is still blocking the merge. Could you help me with one of the following:
1. Dismiss the bot's review now that the changes have been addressed, or
2. Let me know how to trigger a re-review from the bot so it can re-evaluate the updated code?
Thanks!

Right, a maintainer need to trigger a new review with @askserge and it would review again

sergereview

🤗 Serge says:

The PR adds JoyImage Edit Plus support (multi-image editing variant). The model, pipeline, conversion script, docs, and tests are structurally sound and follow existing patterns from the JoyImage Edit family.

Correctness

Division-by-zero in CFG norm rescaling: The existing pipeline_joyimage_edit.py guards against zero-norm with .clamp_min(1e-6), but the new pipeline omits this guard. When comb_pred is all-zeros for a patch (e.g. padded patches), noise_norm will be zero, producing NaN.

Style

Redundant double-resize of reference images: Images are bucket-resized in __call__ (lines 528-530) and then the same resize is applied again inside prepare_latents (lines 337-338). The second resize is a no-op on already-bucketed images, but it's unnecessary work and confusing to readers. Consider removing the resize from prepare_latents since the caller already handles it.

serge v0.1.0 · model: claude-opus-4-6 · 29 LLM turns · 38 tool calls · 159.5s · 1907972 in / 5737 out tokens

tarekziade · 2026-07-03T11:53:01Z

haha my message triggered it

Transformer: - Simplify RoPE by removing dead txt_freqs path - Fix patch_size type hint to list[int] - Use self.config.xxx instead of explicit attributes Pipeline: - Merge VaeImageProcessor into JoyImageEditImageProcessor - Inline normalize/denormalize latents at call sites - Use image_processor.preprocess() for reference images - Remove duplicate resize_center_crop in prepare_latents - Rename padded_latents to latents in __call__ - Support pre-computed latents parameter - Remove redundant padding token zeroing (attention mask suffices) - Use image_processor.postprocess() for output conversion - Return (image,) tuple when return_dict=False

tangyanf · 2026-07-06T03:43:40Z

@dg845 Thanks for the thorough review! All 12 comments have been addressed:

Transformer (transformer_joyimage_edit_plus.py):

Simplified RoPE in AttnProcessor — removed dead txt_freqs path, now passes vis_freqs directly instead of (vis_freqs, None).
Fixed type hint: patch_size: list → patch_size: list[int].
Removed explicit self.xxx attributes, using self.config.xxx instead (kept self.out_channels since it has computed logic).

Pipeline (pipeline_joyimage_edit_plus.py):
4. Merged image processors — removed VaeImageProcessor, kept only JoyImageEditImageProcessor as self.image_processor.
5. Inlined normalize_latents() / denormalize_latents() at call sites, removed the methods.
6. Replaced manual numpy→tensor preprocessing with self.image_processor.preprocess().
7. Removed duplicate resize_center_crop in prepare_latents (already done in call).
8. Renamed padded_latents → latents throughout call.
9. Added pre-computed latents parameter support in prepare_latents.
10. Removed padding token zeroing — the transformer already handles masking via encoder_hidden_states_mask.
11. Replaced manual output conversion with self.image_processor.postprocess().
12. Changed return_dict=False to return (image,) tuple.

Also cleaned up the unused numpy import.

Add .clamp_min(1e-6) to noise_norm to prevent NaN when comb_pred is all-zeros (e.g. padded patches), matching pipeline_joyimage_edit.py.

tangyanf · 2026-07-07T04:41:26Z

Hi @dg845 @yiyixuxu , all CI failures in this PR are pre-existing issues on main and unrelated to our changes:

build_pr_documentation — ModuleNotFoundError: No module named 'sentencepiece' in the Kolors tokenizer import chain (modular_pipelines → auto_pipeline → pag → kolors → tokenizer.py).
test_local_files_only_with_sharded_checkpoint — Assertion mismatch due to a huggingface_hub error message format change (test expects old wording).
test_lcm_multistep — Numerical drift in LCM img2img output (max diff 0.0067 vs threshold 0.001), likely caused by an upstream dependency update.

None of these are related to the JoyImage Edit Plus code. The PR should be ready for review/merge. Let me know if there's anything else needed on my end!

Requested changes have been addressed.

dg845

Thanks for the changes! I think this PR is close to merge.

- Fix vae_image_processor -> image_processor in docs - Remove empty __init__ from AttnProcessor - Restore # Copied from on retrieve_timesteps and sync with source - Add check_inputs method for basic input validation - Remove unnecessary test skip for test_from_save_pretrained_dtype_inference

dg845 · 2026-07-07T23:51:46Z

@bot /style

github-actions · 2026-07-07T23:52:58Z

Style bot fixed some files and pushed the changes.

dg845

Thanks for your work on this PR!

dg845 · 2026-07-08T00:44:59Z

Merging as the CI failures are unrelated.

github-actionsBot added models pipelines size/L PR with diff > 200 LOC labels Jun 22, 2026

yiyixuxu added the no-issue-needed for PRs that do not require link to an issue label Jun 22, 2026

sergereviewBot previously requested changes Jun 22, 2026
View reviewed changes

github-actionsBot added tests utils labels Jun 23, 2026

tangyanfei.8 added 5 commits June 24, 2026 09:35

feat: add image edit plus
a082315

refactor: remove debug code
d6375a8

fix: add missing newline at end of pipeline_output.py
aa2f563

fix: add missing newline at end of pipeline_output.py
8a911e5

tangyanf force-pushed the add-joyimage-edit-plus branch from 6f2763a to 8a911e5Compare June 24, 2026 09:48

github-actionsBot added the fixes-issue label Jun 24, 2026

doc: add joyimage-edit-plus doc
1344fd0

github-actionsBot added the documentation Improvements or additions to documentation label Jun 24, 2026

refactor: update code format
53d0b33

yiyixuxu reviewed Jun 24, 2026
View reviewed changes

refactor: update code formate
4f6cb76

tangyanf requested a review from yiyixuxu June 25, 2026 02:37

yiyixuxu reviewed Jul 2, 2026
View reviewed changes

Comment threadsrc/diffusers/models/transformers/transformer_joyimage.py Outdated

yiyixuxu requested a review from dg845 July 2, 2026 04:40

dg845 reviewed Jul 3, 2026
View reviewed changes

sergereviewBot previously requested changes Jul 3, 2026
View reviewed changes

Comment threadsrc/diffusers/pipelines/joyimage/pipeline_joyimage_edit_plus.py Outdated
Comment threadsrc/diffusers/pipelines/joyimage/pipeline_joyimage_edit_plus.py Outdated

tarekziade mentioned this pull request Jul 3, 2026
do not trigger if @askserge is not the first word huggingface/serge#44
Closed

tangyanfei.8and others added 2 commits July 6, 2026 03:53

fix: guard against division-by-zero in CFG norm rescaling
6cbdd17
Add .clamp_min(1e-6) to noise_norm to prevent NaN when comb_pred is all-zeros (e.g. padded patches), matching pipeline_joyimage_edit.py.

Merge branch 'main' into add-joyimage-edit-plus
cd1328f

tangyanf requested a review from dg845 July 6, 2026 03:57

dg845and others added 3 commits July 6, 2026 15:41

Merge branch 'main' into add-joyimage-edit-plus
da5c16c

style: run ruff format on edit-plus pipeline
87920e8

Merge branch 'main' into add-joyimage-edit-plus
4e93efc

Merge branch 'main' into add-joyimage-edit-plus
8822e45

dg845 reviewed Jul 7, 2026
View reviewed changes

tangyanfei.8and others added 2 commits July 7, 2026 08:50

Merge branch 'main' into add-joyimage-edit-plus
5a831a3

tangyanf requested a review from dg845 July 7, 2026 09:27

Merge branch 'main' into add-joyimage-edit-plus
1be91a0

Apply style fixes
c9a6f17

dg845 approved these changes Jul 8, 2026
View reviewed changes

dg845 merged commit 408ce8d into huggingface:mainJul 8, 2026
13 of 15 checks passed

danielxmed mentioned this pull request Jul 10, 2026
Flux-family attention: flash-attn and other backends fail in an autocast context #14104
Open

Uh oh!

Conversation

tangyanf commented Jun 22, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Model Overview

Key Features

Uh oh!

github-actionsBot commented Jun 22, 2026

Uh oh!

sergereviewBot left a comment

Choose a reason for hiding this comment

Blocking — Debug artifacts left in production code

Blocking — einops dependency

Blocking — sglang integration code in model forward

Blocking — Missing dummy objects

Blocking — Missing tests

Blocking — Hardcoded device_type="cuda" in torch.autocast

Non-blocking — Inlined scheduler sigma math

Non-blocking — Unused imports and parameters

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yiyixuxu commented Jun 24, 2026

Uh oh!

github-actionsBot commented Jun 24, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yiyixuxu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tangyanf commented Jun 25, 2026

Uh oh!

tangyanf commented Jun 25, 2026

Blocking — Debug artifacts left in production code

Blocking — einops dependency

Blocking — sglang integration code in model forward

Blocking — Missing dummy objects

Blocking — Missing tests

Blocking — Hardcoded device_type="cuda" in torch.autocast

Non-blocking — Inlined scheduler sigma math

Non-blocking — Unused imports and parameters

Uh oh!

tangyanf commented Jun 25, 2026

Uh oh!

Uh oh!

tangyanf commented Jul 3, 2026

Uh oh!

dg845 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tarekziade commented Jul 3, 2026

Uh oh!

sergereviewBot left a comment

tangyanf commented Jun 22, 2026•
edited
Loading

Blocking — `einops` dependency

Blocking — `sglang` integration code in model forward

Blocking — Hardcoded `device_type="cuda"` in `torch.autocast`

github-actionsBot commented Jun 24, 2026•
edited
Loading

Blocking — `einops` dependency

Blocking — `sglang` integration code in model forward

Blocking — Hardcoded `device_type="cuda"` in `torch.autocast`

github-actionsBot commented Jul 7, 2026•
edited
Loading