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]

Should we round block_state.height and block_state.width to a multiple of self.vae_scale_factor * self.patch_size before calculating latent_height and latent_width like the standard Krea2 pipeline does?

diffusers/src/diffusers/pipelines/krea2/pipeline_krea2.py

Lines 534 to 542 in 208704a

	multiple = self.vae_scale_factor * self.patch_size
	if height % multiple != 0 or width % multiple != 0:
	rounded_height = ((height + multiple - 1) // multiple) * multiple
	rounded_width = ((width + multiple - 1) // multiple) * multiple
	logger.warning(
	f"`height` and `width` must be multiples of {multiple}; rounding up from {height}x{width} to"
	f" {rounded_height}x{rounded_width}."
	)
	height, width = rounded_height, rounded_width

[SamuelTallet]

Hello :)

I'm not sure if this is the right time to bring this up or if it would be better to address it in a follow-up PR: it would be great if Krea's modular pipeline could also handle reference images. I’m thinking of these new cool LoRAs:

• https://huggingface.co/ostris/krea2_turbo_style_reference
• https://huggingface.co/conradlocke/krea2-identity-edit

[Cedric-Perauer]

Thanks for the suggestions @dg845 I added the changes.