Krea2 Modular Pipeline Support, Documentation, Test Cases

[Cedric-Perauer]

What does this PR do?

This PR implements the Krea2 modular pipeline according to requested feature :
Fixes#14056

The tiny-krea2-modular-pipe for test cases here still needs to be mirrored from https://huggingface.co/CedricPerauer/tiny-krea2-modular-pipe to hf-internal-testing

Validation

python -m pytest tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py -v
#14 checks passed
PYTHONPATH=src ruff check \
src/diffusers/modular_pipelines/krea2/__init__.py \
src/diffusers/modular_pipelines/krea2/before_denoise.py \
src/diffusers/modular_pipelines/krea2/decoders.py \
src/diffusers/modular_pipelines/krea2/denoise.py \
src/diffusers/modular_pipelines/krea2/encoders.py \
src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py \
src/diffusers/modular_pipelines/krea2/modular_pipeline.py \
src/diffusers/modular_pipelines/__init__.py \
src/diffusers/modular_pipelines/modular_pipeline.py \
src/diffusers/utils/dummy_torch_and_transformers_objects.py \
src/diffusers/utils/__init__.py \
tests/modular_pipelines/krea2/__init__.py \
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py
#13 checks passed
PYTHONPATH=src ruff format --check \
src/diffusers/modular_pipelines/krea2/__init__.py \
src/diffusers/modular_pipelines/krea2/before_denoise.py \
src/diffusers/modular_pipelines/krea2/decoders.py \
src/diffusers/modular_pipelines/krea2/denoise.py \
src/diffusers/modular_pipelines/krea2/encoders.py \
src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py \
src/diffusers/modular_pipelines/krea2/modular_pipeline.py \
src/diffusers/modular_pipelines/__init__.py \
src/diffusers/modular_pipelines/modular_pipeline.py \
src/diffusers/utils/dummy_torch_and_transformers_objects.py \
src/diffusers/utils/__init__.py \
tests/modular_pipelines/krea2/__init__.py \
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py
#13 checks passed

Full PyTest output

========================================================================================== test session starts ===========================================================================================
platform linux -- Python 3.11.10, pytest-9.1.1, pluggy-1.6.0 -- /usr/bin/python
cachedir: .pytest_cache
rootdir: /root/diffusers
configfile: pyproject.toml
plugins: hf-doc-builder-0.6.0.dev0, anyio-4.6.0
collected 14 items tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_pipeline_call_signature PASSED [ 7%]
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_inference_batch_consistent PASSED [ 14%]
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_float16_inference PASSED [ 21%]
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_to_device PASSED [ 28%]
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_inference_is_not_nan_cpu PASSED [ 35%]
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_inference_is_not_nan PASSED [ 42%]
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_num_images_per_prompt PASSED [ 50%]
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_components_auto_cpu_offload_inference_consistent PASSED [ 57%]
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_save_from_pretrained PASSED [ 64%]
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_load_expected_components_from_pretrained PASSED [ 71%]
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_load_expected_components_from_save_pretrained PASSED [ 78%]
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_modular_index_consistency PASSED [ 85%]
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_workflow_map PASSED [ 92%]
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py::TestKrea2ModularPipelineFast::test_inference_batch_single_identical PASSED [100%]
============================================================================================ warnings summary ============================================================================================
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py: 13 warnings
/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_validators.py:205: UserWarning: The `local_dir_use_symlinks` argument is deprecated and ignored in`hf_hub_download`. Downloading to a local directory does not use symlinks anymore.
warnings.warn(
-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
==================================================================================== 14 passed, 13 warnings in 40.03s ====================================================================================

Docs

Krea2.md is updated accordingly with small section for the modular pipeline

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you self-review the diff against .ai/review-rules.md?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline?

Who can review?

Potential reviewers for pipelines/models: @yiyixuxu@dg845@asomoza

cc @SamuelTallet

[github-actions]

Hi @Cedric-Perauer, 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.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[Cedric-Perauer]

The tiny-krea2-modular-pipe for test cases here still needs to be mirrored from https://huggingface.co/CedricPerauer/tiny-krea2-modular-pipe to hf-internal-testing. This is currently leading to the failing modular/fast pytorch check

[dg845]

Thanks for the PR! Left some comments :).

[dg845]

LGTM, thanks for your work on this PR! @yiyixuxu, can you also take a look to see if it looks good to you?

[yiyixuxu]

ohh i think turbo should have its own set of blocks that it maps to, it is a different checkpoint

this way the distilled krea should not accept negative prompt and it should have guider disabled by default
you can take a look at how flux2 klein is implemented in relation to the klein base
https://github.com/huggingface/diffusers/blob/main/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py

[yiyixuxu]

@Cedric-Perauer can you update the auto docstrings? we can merge once CI pass
thanks a lot for working on the PR

[SamuelTallet]

Thanks @Cedric-Perauer 😊

[Cedric-Perauer]

Thanks for your reviews @yiyixuxu ! Docstring should be fixed now