Skip to content

Commit 1030249

Browse files
Moran232zhangmaoquan.1feice-huangdg845github-actions[bot]
authored
[feat] JoyAI-JoyImage-Edit support (huggingface#13444)
* [feat] JoyAI-JoyImage-Edit support * [fix] remove rearrange * [refactor] two pass when do cfg * [refactor] remove repa, use wantimetextembeding, refactor modulate code * [refactor] Joyimage Attention refactor * remove vae tiling and autocast * [fix] remove einops from setup.py * [refactor] Refactor JoyImageEditPipeline to use explicit arguments instead of namespace and remove _build_arg * [fix] remove deprecated method decode_latents * [refactor] refactor the image pre-processing logic into a separate VaeImageProcessor subclass * [refactor] add JoyImageAttention to align with Attention + AttnProcessor design and update conversion script for new weight key mapping (e.g. img_attn_qkv -> attn.img_attn_qkv) * [refactor] simplify bucket logic in JoyImageEditImageProcessor by replacing runtime generation with precomputed lookup tables * [fix] remove leftover training-only parameters * [fix] add layerwise casting and fp32 module patterns to JoyImageTransformer3DModel. Reference WanTransformer3DModel to fix layer casting errors during inference. * [test] add JoyImageEditPipeline fast tests and JoyImageEditTransformer3DModel model tests * [fix] fix some pipeline args to support batch inference * [fix] duplicate images to match batch size when fewer images than prompts in JoyImageEditPipeline * [fix] remove no longer used config parameters * Apply style fixes * [fix] remove unused dataclass and rewrite helpers as inline functions * [fix] make dummy objects for JoyImageEdit * [fix] allow test_torch_compile_repeated_blocks to pass * [fix] add examples on JoyImageEditPipeline * fix code style issues with ruff and black * Apply style fixes * [fix] change default num_inference_steps to 40 * [fix] use forward hook to extract pre-norm hidden states for transformers 5.x compatibility * [fix] change the assert to ValueError in pipeline * [fix] rename JoyImageTransformer3DModel to JoyImageEditTransformer3DModel, clean up anything about the alias * [fix] support gradient checkpointing * [refactor] simplify RoPE utilities, inline helpers, copy WanTimeTextImageEmbedding locally and remove unused parameters * [fix] remove _get_text_encoder_ckpt and qwen_processor * [fix] change nn.RMSNorm to FP32LayerNorm * [fix] small fixes for suggestions given by Claude * [refactor] build model using from _pretained instead of config * [refactor] auto-wrap prompt and support text-to-image in JoyImage Edit pipeline * make style, make quality and make fix-copies * [test] small fix to use vocab_size=1024 * [refactor] separate encode_prompt_multiple_images from encode_prompt, support prompt_embeds/prompt_embesd_mask/num_images_per_prompt in edit mode * [test] fix CI: use strict=False for xfail and add @require_torch_accelerator to group offloading test * [refactor] separate image_latents from latents in prepare_latents to align with flux2 * make style --------- Co-authored-by: zhangmaoquan.1 <zhangmaoquan.1@jd.com> Co-authored-by: huangfeice <huangfeice@gmail.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: YiYi Xu <yixu310@gmail.com>
1 parent 3577280 commit 1030249

15 files changed

Lines changed: 2453 additions & 0 deletions

‎scripts/convert_joyimage_edit_to_diffusers.py‎

Lines changed: 366 additions & 0 deletions
Large diffs are not rendered by default.

‎src/diffusers/__init__.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,7 @@
252252
"HunyuanVideoFramepackTransformer3DModel",
253253
"HunyuanVideoTransformer3DModel",
254254
"I2VGenXLUNet",
255+
"JoyImageEditTransformer3DModel",
255256
"Kandinsky3UNet",
256257
"Kandinsky5Transformer3DModel",
257258
"LatteTransformer3DModel",
@@ -582,6 +583,8 @@
582583
"IFPipeline",
583584
"IFSuperResolutionPipeline",
584585
"ImageTextPipelineOutput",
586+
"JoyImageEditPipeline",
587+
"JoyImageEditPipelineOutput",
585588
"Kandinsky3Img2ImgPipeline",
586589
"Kandinsky3Pipeline",
587590
"Kandinsky5I2IPipeline",
@@ -1071,6 +1074,7 @@
10711074
HunyuanVideoFramepackTransformer3DModel,
10721075
HunyuanVideoTransformer3DModel,
10731076
I2VGenXLUNet,
1077+
JoyImageEditTransformer3DModel,
10741078
Kandinsky3UNet,
10751079
Kandinsky5Transformer3DModel,
10761080
LatteTransformer3DModel,
@@ -1376,6 +1380,8 @@
13761380
IFPipeline,
13771381
IFSuperResolutionPipeline,
13781382
ImageTextPipelineOutput,
1383+
JoyImageEditPipeline,
1384+
JoyImageEditPipelineOutput,
13791385
Kandinsky3Img2ImgPipeline,
13801386
Kandinsky3Pipeline,
13811387
Kandinsky5I2IPipeline,

‎src/diffusers/models/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,9 @@
113113
_import_structure["transformers.transformer_hunyuan_video15"] = ["HunyuanVideo15Transformer3DModel"]
114114
_import_structure["transformers.transformer_hunyuan_video_framepack"] = ["HunyuanVideoFramepackTransformer3DModel"]
115115
_import_structure["transformers.transformer_hunyuanimage"] = ["HunyuanImageTransformer2DModel"]
116+
_import_structure["transformers.transformer_joyimage"] = [
117+
"JoyImageEditTransformer3DModel",
118+
]
116119
_import_structure["transformers.transformer_kandinsky"] = ["Kandinsky5Transformer3DModel"]
117120
_import_structure["transformers.transformer_longcat_audio_dit"] = ["LongCatAudioDiTTransformer"]
118121
_import_structure["transformers.transformer_longcat_image"] = ["LongCatImageTransformer2DModel"]
@@ -236,6 +239,7 @@
236239
HunyuanVideo15Transformer3DModel,
237240
HunyuanVideoFramepackTransformer3DModel,
238241
HunyuanVideoTransformer3DModel,
242+
JoyImageEditTransformer3DModel,
239243
Kandinsky5Transformer3DModel,
240244
LatteTransformer3DModel,
241245
LongCatAudioDiTTransformer,

‎src/diffusers/models/transformers/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
from .transformer_hunyuan_video15 import HunyuanVideo15Transformer3DModel
3737
from .transformer_hunyuan_video_framepack import HunyuanVideoFramepackTransformer3DModel
3838
from .transformer_hunyuanimage import HunyuanImageTransformer2DModel
39+
from .transformer_joyimage import JoyImageEditTransformer3DModel
3940
from .transformer_kandinsky import Kandinsky5Transformer3DModel
4041
from .transformer_longcat_audio_dit import LongCatAudioDiTTransformer
4142
from .transformer_longcat_image import LongCatImageTransformer2DModel

0 commit comments

Comments
 (0)