Skip to content

[Feat] support JoyAIImagePipeline #13437

Open
RuixiangMa wants to merge 6 commits intohuggingface:mainfrom
RuixiangMa:joyaiimage
Open

[Feat] support JoyAIImagePipeline #13437
RuixiangMa wants to merge 6 commits intohuggingface:mainfrom
RuixiangMa:joyaiimage

Conversation

@RuixiangMa
Copy link
Copy Markdown

@RuixiangMa RuixiangMa commented Apr 9, 2026 •
edited
Loading

Support JoyAIImagePipeline

Prompt image
NO dog
add hat in the dog joyai_image_seqoffload_out
Rotate the dog to show the left side view left
Rotate the dog to show the right side view right

Before submitting

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@RuixiangMa
[Feat] support JoyAIImagePipeline
c1f00a2 
Signed-off-by: Lancer <maruixiang6688@gmail.com>
@github-actions github-actions bot added models pipelines size/L PR with diff > 200 LOC labels Apr 9, 2026
@RuixiangMa RuixiangMa marked this pull request as draft April 9, 2026 11:42
@RuixiangMa
upd
d6365ec 
Signed-off-by: Lancer <maruixiang6688@gmail.com>
@github-actions github-actions bot added size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels Apr 9, 2026
@RuixiangMa
upd
e774af2 
Signed-off-by: Lancer <maruixiang6688@gmail.com>
@github-actions github-actions bot added schedulers size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels Apr 9, 2026
@RuixiangMa
upd
8190e3f 
Signed-off-by: Lancer <maruixiang6688@gmail.com>
@github-actions github-actions bot added tests size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels Apr 9, 2026
@RuixiangMa RuixiangMa marked this pull request as ready for review April 9, 2026 18:20
yiyixuxu
yiyixuxu reviewed Apr 9, 2026
View reviewed changes
Copy link
Copy Markdown
Collaborator

@yiyixuxu yiyixuxu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks for the PR, i left some initial comments

src/diffusers/models/autoencoders/autoencoder_kl_joyai_image.py Outdated
from .autoencoder_kl_wan import (
WanAttentionBlock as AttentionBlock,
)
from .autoencoder_kl_wan import (
Copy link
Copy Markdown
Collaborator

@yiyixuxu yiyixuxu Apr 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you not import from wan? let's use #Copied from isntead

Copy link
Copy Markdown
Author

@RuixiangMa RuixiangMa Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will remove the direct imports from Wan and inlined the reused helper blocks in JoyAI with # Copied from comments instead.

src/diffusers/pipelines/joyai_image/pipeline_joyai_image.py Outdated
self._joyai_force_vae_fp32 = True

@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
Copy link
Copy Markdown
Collaborator

@yiyixuxu yiyixuxu Apr 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what's going on here? is there any reason we cannot make eit work with DiffusionPipeline.from_pretrained?

Copy link
Copy Markdown
Author

@RuixiangMa RuixiangMa Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The upstream JoyAI checkpoint is not a standard diffusers repo layout, so it cannot be loaded by DiffusionPipeline.from_pretrained directly.

I agree the better long-term solution is to add a one-time conversion/export to a standard diffusers checkpoint and then rely on the base from_pretrained path for the exported directory

src/diffusers/schedulers/scheduling_joyai_flow_match_discrete.py Outdated
prev_sample: torch.FloatTensor


class JoyAIFlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
Copy link
Copy Markdown
Collaborator

@yiyixuxu yiyixuxu Apr 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@claude
can you tell me what's the difference between this and FlowMatchDiscreteSchedule?

Copy link
Copy Markdown
Author

@RuixiangMa RuixiangMa Apr 10, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I double-checked this. It is close to FlowMatchEulerDiscreteScheduler, but not numerically identical as-is because JoyAI uses its own timestep/sigma shifting behavior.

I will refactorit to subclass FlowMatchEulerDiscreteScheduler and only keep the JoyAI-specific differences in the scheduler implementation.

@RuixiangMa
upd
5ebbbd7 
Signed-off-by: Lancer <maruixiang6688@gmail.com>
@RuixiangMa
upd
2ccc45f 
Signed-off-by: Lancer <maruixiang6688@gmail.com>
@github-actions github-actions bot added documentation Improvements or additions to documentation size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels Apr 10, 2026
@RuixiangMa
Copy link
Copy Markdown
Author

RuixiangMa commented Apr 10, 2026

@yiyixuxu Thanks for the comments! PTAL

@RuixiangMa RuixiangMa mentioned this pull request Apr 10, 2026
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@yiyixuxu yiyixuxu yiyixuxu left review comments

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

documentation Improvements or additions to documentation models pipelines schedulers size/L PR with diff > 200 LOC tests

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@RuixiangMa @yiyixuxu