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

yiyixuxu · 2026-06-24T19:28:59Z

+from ..modeling_utils import ModelMixin
+from ..normalization import FP32LayerNorm
+from .transformer_joyimage import (
+    JoyImageAttention,


we should not cross-import between transformer models
for exact same blocks you want to reuse, can you copy paste over with #Copied from .... with JoyImage-> JoyImagePlus statements? , or create new ones for joy image edit

yiyixuxu · 2026-06-24T19:55:37Z

+        encoder_hidden_states: torch.Tensor = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,


Suggested change

**kwargs,

yiyixuxu · 2026-06-24T20:31:39Z

+    if cos.ndim == 2:
+        # unbatched: [S, D] -> [1, S, 1, D]
+        cos = cos.unsqueeze(0).unsqueeze(2)
+        sin = sin.unsqueeze(0).unsqueeze(2)


Suggested change

if cos.ndim == 2:

# unbatched: [S, D] -> [1, S, 1, D]

cos = cos.unsqueeze(0).unsqueeze(2)

sin = sin.unsqueeze(0).unsqueeze(2)

yiyixuxu · 2026-06-24T20:42:15Z

+        if shape_list is None:
+            raise ValueError(
+                "shape_list must be provided either as an argument or via forward_batch.vae_image_sizes"
+            )


Suggested change

if shape_list is None:

raise ValueError(

"shape_list must be provided either as an argument or via forward_batch.vae_image_sizes"

)

let's just make sure shape_list is a required argument

yiyixuxu · 2026-06-24T20:45:23Z

+        if vec.shape[-1] > self.hidden_size:
+            vec = vec.unflatten(1, (6, -1))


Suggested change

if vec.shape[-1] > self.hidden_size:

vec = vec.unflatten(1, (6, -1))

vec = vec.unflatten(1, (6, -1))

yiyixuxu · 2026-06-24T20:50:07Z

+"""
+
+
+def retrieve_timesteps(


can we add a #Copied from ... here?

yiyixuxu · 2026-06-24T20:52:18Z

+        img = img.crop((left, top, left + bw, top + bh))
+        return img
+
+    def _get_bucket_size(self, img: Image.Image) -> tuple[int, int]:


can you move this to image_processor.py?

yiyixuxu · 2026-06-24T20:58:11Z

+            latent = latent / self.vae.config.scaling_factor
+        return latent
+
+    def _resize_center_crop(self, img: Image.Image, target_size: tuple[int, int]) -> Image.Image:


you already have a _resize_center_crop method on the image processor? can you just make it work with self.image_processor.preprocess(...)?

github-actions Bot 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

sergereview Bot requested changes Jun 22, 2026

View reviewed changes

github-actions Bot 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 8a911e5 Compare June 24, 2026 09:48

github-actions Bot added the fixes-issue label Jun 24, 2026

doc: add joyimage-edit-plus doc

1344fd0

github-actions Bot 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

		if vec.shape[-1] > self.hidden_size:
		vec = vec.unflatten(1, (6, -1))

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-actions Bot commented Jun 22, 2026

Uh oh!

sergereview Bot 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-actions Bot 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

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-actions Bot commented Jun 24, 2026 •

edited

Loading