@yiyixuxu @dg845
Thank you very much for your valuable feedback. I've made some modifications. See my latest commits.

Specifically, I refactored the attention module. However, since the weight key names in the Diffusers model are already fixed, I didn't change the actual keys in the attention part. Additionally, I will consider refactoring the image pre-processing logic, since the logic is quite complex, I directly copied it over from the training code.

If you have any further suggestions, please feel free to share. Thank you so much!

@claude can you do a final review too?

Claude finished @yiyixuxu's task in 3m 24s —— View job

Final Review of JoyAI-JoyImage-Edit PR

• Read repository guidelines (CLAUDE.md, models.md, pipelines.md)
• Review transformer model (transformer_joyimage.py)
• Review image processor (image_processor.py)
• Review pipeline (pipeline_joyimage_edit.py)
• Review pipeline output and init/export files
• Review tests
• Check outstanding reviewer comments
• Post final review

The PR has improved significantly since the last round. Many of the earlier issues (einops removal, attention refactor, gradient checkpointing, assert→ValueError, etc.) have been addressed. Below are the remaining issues I found, organized by severity. I've also noted which items align with outstanding reviewer comments from @yiyixuxu and @dg845 that still need resolution.

Must Fix

1. _keep_in_fp32_modules references non-existent modules — transformer_joyimage.py:381-387

_keep_in_fp32_modules = [
    "time_embedder",
    "scale_shift_table",  # ← does not exist in this model
    "norm1",
    "norm2",
    "norm3",              # ← does not exist in this model
]

scale_shift_table and norm3 are WAN concepts copied over but not present in JoyImageEditTransformer3DModel. These entries are harmless (they just match nothing) but are misleading. Update to match the actual module names in this model.

2. __init__.py export ordering — src/diffusers/__init__.py:584-586

"JoyImageEditPipeline",
"Kandinsky3Img2ImgPipeline",
"JoyImageEditPipelineOutput",  # ← out of alphabetical order

JoyImageEditPipelineOutput should come immediately after JoyImageEditPipeline, before Kandinsky3Img2ImgPipeline. This will likely be caught by make quality.

3. enable_tiling still in docstring — pipeline_joyimage_edit.py:698-699

            enable_tiling (`bool`, *optional*, defaults to `False`):
                Enable tiled VAE decoding to reduce peak memory usage.

This parameter does not exist in __call__. Remove from docstring.

4. **kwargs in __call__ silently swallows arguments — pipeline_joyimage_edit.py:649

yiyixuxu flagged this. **kwargs is accepted but never used — it silently discards any mistyped keyword arguments. Remove it.

Should Fix (Outstanding reviewer comments)

5. encode_prompt_multiple_images should be called from __call__ directly — yiyixuxu's comment

Currently encode_prompt internally routes to encode_prompt_multiple_images when images is not None. yiyixuxu asked to move the routing to __call__:

# In __call__:
if processed_image is None:
    ... = self.encode_prompt(...)
else:
    ... = self.encode_prompt_multiple_images(...)

Additionally, encode_prompt_multiple_images should accept pre-computed prompt_embeds/prompt_embeds_mask and handle num_images_per_prompt expansion, similar to encode_prompt.

6. Timestep dtype inference via weight introspection — transformer_joyimage.py:356-358, yiyixuxu's comment

time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype
if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8:
    timestep = timestep.to(time_embedder_dtype)

Per models.md gotcha #4: "don't cast activations by reading a weight's dtype — the stored weight dtype isn't the compute dtype under gguf / quantized loading." yiyixuxu suggested simplifying to:

timestep = timestep.to(encoder_hidden_states.dtype)
temb = self.time_embedder(timestep)

7. nn.RMSNorm vs diffusers RMSNorm — transformer_joyimage.py:207-208, 212-213

The attention sublayer QK norms use nn.RMSNorm which does not upcast to FP32 during computation. The diffusers RMSNorm (from normalization.py) does upcast. dg845 flagged this earlier. If the original model's training used FP32 upcast for QK norms, this could cause numerical drift. Please verify whether the original training code upcasts, and use the appropriate norm class.

Nice to Fix

8. Two forward passes instead of batched CFG — pipeline_joyimage_edit.py:823-836

The denoising loop makes two separate transformer forward passes (conditional + unconditional) instead of batching them together. This is ~2x slower than the standard diffusers pattern:

latent_model_input = torch.cat([latents] * 2)
prompt_concat = torch.cat([negative_prompt_embeds, prompt_embeds])
noise_pred = self.transformer(latent_model_input, ..., prompt_concat, ...)
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)

I understand this may be intentional due to the multi-item latent structure, but it's worth considering for performance.

9. Always-on guidance rescaling — pipeline_joyimage_edit.py:839-842

cond_norm = torch.norm(noise_pred, dim=2, keepdim=True)
noise_norm = torch.norm(comb_pred, dim=2, keepdim=True)
noise_pred = comb_pred * (cond_norm / noise_norm.clamp_min(1e-6))

This norm-based guidance rescaling is always applied during CFG. Most diffusers pipelines make this optional via a guidance_rescale parameter. If this is an intentional design choice for JoyImage, a brief comment explaining why would help readers.

10. Mixed typing style — pipeline_joyimage_edit.py:2, 620-649

The pipeline mixes old-style typing (Dict, List, Optional, Union) and modern style (str | list[str], int | None). Similarly, the transformer uses from typing import Tuple instead of tuple. Should be consistent — prefer Python 3.9+ builtins per dg845's earlier suggestion.

11. Tuple in transformer — transformer_joyimage.py:17

from typing import Tuple — replace all Tuple[...] with tuple[...] for Python 3.9+ consistency.

12. image_processor.py bucket validation — image_processor.py:57-63

find_best_bucket will KeyError if basesize != 1024. Add a check:

if basesize not in BUCKETS:
    raise ValueError(f"Unsupported basesize {basesize}. Supported: {list(BUCKETS.keys())}")

13. @register_to_config on VaeImageProcessor subclass — image_processor.py:91

VaeImageProcessor.__init__ already uses register_to_config. Double-decorating in the subclass may cause issues with config serialization. Verify this works correctly with from_config.

Style / CI

14. make style and make quality needed — dg845 asked for this. The style bot doesn't seem to be working. The authors should run make style && make quality && make fix-copies locally and push the result.

Tests

15. test_from_save_pretrained_dtype_inference unconditionally skipped — test_models_transformer_joyimage.py:90-91

pytest.skip("Tolerance requirements too high for meaningful test")

This unconditionally skips a core test. If there are tolerance issues, investigate the root cause or use @pytest.mark.xfail with a reason so it's tracked. pytest.skip hides the problem entirely.

16. Offloading test failures — dg845 identified that test_group_offloading_inference, test_pipeline_level_group_offloading_inference, test_sequential_cpu_offload_forward_pass, and test_sequential_offload_forward_pass_twice fail due to Qwen3VL incompatibility with leaf-level offloading. These need to be either overridden (as dg845 suggested) or skipped with clear documentation.

Summary

The architecture is clean — JoyImageAttention + JoyImageAttnProcessor with dispatch_attention_fn follows the recommended pattern well. Gradient checkpointing, _repeated_blocks, and _no_split_modules are all properly set. The _get_last_decoder_hidden_states hook for transformers 5.x compat is a pragmatic solution.

The main outstanding items are yiyixuxu's latest review comments (#5, #6, #4), the stale _keep_in_fp32_modules entries (#1), and running make style && make quality (#14). Once those are addressed, this should be ready to merge.

ohh calude did a good job here, can you look into fixing the #1 and #3 listed in the claude review?

@claude are you able to do a dead code analysis since it's a new model here?
(also, can you let me know why you did not include this in the review already? we listed in the review-fules.md)

Claude finished @yiyixuxu's task in 2m 9s —— View job

Dead Code Analysis for JoyImage-Edit Model

• Read review rules to understand the dead code analysis requirements
• Read the pipeline __call__ and trace all calls into the transformer
• Read the transformer model __init__ and forward
• Read the published model config (jdopensource/JoyAI-Image-Edit-Diffusers)
• Cross-reference: identify unused parameters, methods, and unreachable code paths
• Explain why previous reviews missed this analysis
• Post final dead code analysis report

Why this wasn't in the previous reviews

The review rules in .ai/review-rules.md explicitly require a dead code analysis section for PRs that add new models ("Trace the call path... Check the default model config... Flag unused parameters and methods"). My two earlier reviews focused on correctness, API design, and convention issues but failed to follow the dead-code-analysis checklist. That was an oversight on my part — the rule was there, I just didn't execute it. Sorry about that.

Call path trace

Pipeline __call__ (lines 619–880) calls into the transformer at line 823:

noise_pred = self.transformer(
    hidden_states=latent_model_input,   # (B, num_items, C, T, H, W) — 6D when editing
    timestep=t_expand,
    encoder_hidden_states=prompt_embeds,
    return_dict=False,
)[0]

Only three arguments are ever passed: hidden_states, timestep, encoder_hidden_states, and return_dict. The transformer's forward signature is:

def forward(self, hidden_states, timestep, encoder_hidden_states=None, return_dict=True):

All four parameters are used. Good — no unused forward parameters.

Published config (jdopensource/JoyAI-Image-Edit-Diffusers/transformer/config.json):

{
  "hidden_size": 4096,
  "num_attention_heads": 32,
  "rope_type": "rope",
  "rope_dim_list": [16, 56, 56],
  "patch_size": [1, 2, 2],
  "theta": 10000
}

Findings (advisory — not blocking)

Under the published config and the pipeline's call path, the following code appears unreachable or unused:

1. Text RoPE branch is dead code — transformer_joyimage.py:561

txt_rope_size=txt_seq_len if self.rope_type == "mrope" else None,

The published config has rope_type = "rope", so txt_rope_size is always None. This means:

• get_rotary_pos_embed always takes the early return at line 490–491 (if txt_rope_size is None: return vis_freqs, None)
• The entire text-RoPE computation (lines 494–501) is never executed
• In JoyImageAttnProcessor.__call__, the text RoPE branch at line 143–144 (if txt_freqs is not None) is never taken — text tokens get no RoPE

This is likely intentional (the model was trained with rope_type="rope"), but the "mrope" code path is untested dead code.

2. rope_dim_list is None fallback is unreachable — transformer_joyimage.py:466-467

if rope_dim_list is None:
    rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)]

The published config sets rope_dim_list = [16, 56, 56], and the __init__ default is also [16, 56, 56]. There is no config where rope_dim_list would be None. This fallback is defensive but unreachable.

3. _keep_in_fp32_modules references non-existent modules — transformer_joyimage.py:381-387

_keep_in_fp32_modules = [
    "time_embedder",       # ✓ exists (inside condition_embedder.time_embedder)
    "scale_shift_table",   # ✗ does NOT exist anywhere in the model
    "norm1",               # ✓ exists (img_norm1, txt_norm1 in blocks)
    "norm2",               # ✓ exists (img_norm2, txt_norm2 in blocks)
    "norm3",               # ✗ does NOT exist anywhere in the model
]

scale_shift_table and norm3 are WAN model concepts that were copied but don't exist in JoyImageEditTransformer3DModel. These entries silently match nothing — harmless but misleading.

4. num_items = 0 initialization — semantically misleading — transformer_joyimage.py:532

num_items = 0
if is_multi_item:
    num_items = hidden_states.shape[1]

When is_multi_item is False, num_items remains 0. The un-multi-item block at line 581 (if is_multi_item) won't execute in that case, so num_items = 0 is never used meaningfully. However, if is_multi_item were accidentally set while num_items = 0, it would cause a silent reshape bug. Consider leaving num_items unset until needed.

5. Pipeline: template_type="image" passed to encode_prompt is overridden when images are present — pipeline_joyimage_edit.py:757

prompt_embeds, prompt_embeds_mask = self.encode_prompt(
    ...
    images=processed_image,
    ...
    template_type="image",   # <-- this value is ignored when images is not None
)

When images is not None, encode_prompt immediately delegates to encode_prompt_multiple_images (line 386-392), which hard-codes template_type="multiple_images". The template_type="image" argument passed from __call__ is silently discarded in the editing path.

6. _skip_layerwise_casting_patterns includes "condition_embedder" indirectly via "norm" — transformer_joyimage.py:378

_skip_layerwise_casting_patterns = ["img_in", "condition_embedder", "norm"]

The "norm" pattern matches norm_out, img_norm1, img_norm2, txt_norm1, txt_norm2, and also the RMSNorm QK norms inside attention (img_attn_q_norm, etc.). This is a very broad match. It also matches condition_embedder which contains the time embedder — this is likely intentional but means most of the model is excluded from layerwise casting.

7. vec.shape[-1] > self.hidden_size check in forward — transformer_joyimage.py:553-554

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

JoyImageTimeTextImageEmbedding.forward returns temb from time_proj which has output dim hidden_size * 6 (set at init: time_proj_dim=hidden_size * 6). So vec.shape[-1] is always hidden_size * 6, which is always > self.hidden_size. This branch is always taken — the else path (where vec stays 1D) is dead code. The check could be removed and replaced with an unconditional unflatten.

Methods and layers that ARE used (confirming no false positives)

All initialized layers and helper methods are reachable from the pipeline call path. No dead layers or methods found.

Summary

The model is clean — all initialized layers and methods are used. The dead code is limited to:

• An unreachable rope_type == "mrope" text-RoPE path (finding Add glide modeling files #1)
• A stale rope_dim_list is None fallback (finding + cosine schedule and unet config #2)
• _keep_in_fp32_modules entries (scale_shift_table, norm3) that match nothing (finding make from hub import work #3 — also flagged in my earlier review as "must fix" Add glide modeling files #1)
• An always-true condition on vec.shape (finding add vqvae #7)

None of these are bugs — they're defensive/copied code that happens to be unreachable under the published config. The authors may have configs or future models that exercise these paths.

@dg845 @yiyixuxu
Thanks for the suggestions! Updates:

Refactor

• Separate encode_prompt_multiple_images from encode_prompt. Dispatching is now handled in __call__ based on whether processed_image is provided, instead of branching inside encode_prompt.
• Add prompt_embeds, prompt_embeds_mask, and num_images_per_prompt support to encode_prompt_multiple_images so it matches the encode_prompt interface.

Test updates

• Skipp test_group_offloading_inference, test_pipeline_level_group_offloading_inference, test_sequential_cpu_offload_forward_pass, and test_sequential_offload_forward_pass_twice.

Not changed

• We keep the current time_embedder dtype inference pattern (next(iter(...)).dtype). We tested simplifying it, but that causes dtype mismatch errors.

@dg845 fix these in e8c4db7. hoping this resolves the CI failures.

@yiyixuxu
We directly used the testing script from #13444 (comment). If you apply the changes below, you will hit the error mentioned in #13444 (comment).

Let us know if you need more details!

    def forward(
        self,
        timestep: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
    ):
        timestep = self.timesteps_proj(timestep)

        if timestep.dtype != encoder_hidden_states.dtype:
            timestep = timestep.to(encoder_hidden_states.dtype)
        temb = self.time_embedder(timestep)
        timestep_proj = self.time_proj(self.act_fn(temb))

        encoder_hidden_states = self.text_embedder(encoder_hidden_states)

        return temb, timestep_proj, encoder_hidden_states

thanks a lot for the PR!

thank you so much for your suggestions!

Hey folks,

@Moran232 @feice-huang docs is missing for JoyImage. Could you please add docs in a separate PR?