Krea2 Modular Pipeline Support, Documentation, Test Cases by Cedric-Perauer · Pull Request #14083 · huggingface/diffusers

Cedric-Perauer · 2026-06-28T23:59:23Z

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 · 2026-06-29T09:09:22Z

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 · 2026-07-08T22:37:53Z

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 · 2026-07-08T22:57:21Z

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 · 2026-07-09T07:55:45Z

+        latent_height = block_state.height // components.vae_scale_factor
+        latent_width = block_state.width // components.vae_scale_factor


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

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

Cedric-Perauer · 2026-07-09T10:29:44Z

Thanks for the suggestions @dg845 I added the changes.

dg845

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

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

yiyixuxu · 2026-07-16T19:58:00Z

+image.save("krea2.png")
+```
+
+We additionally provide an example for using Krea2 Turbo. The distilled checkpoint disables guidance by


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

Cedric-Perauer added 6 commits June 27, 2026 13:38

towards krea2 modular

6f8e01b

add test cases

2bae576

self-review fixes: test model id to hf internal, guider change

30c5138

remove pnd

7450dc4

modular autodoc

4e4da7f

modular autodoc

4da8d49

github-actions Bot added documentation Improvements or additions to documentation tests modular-pipelines utils size/L PR with diff > 200 LOC labels Jun 28, 2026

Merge branch 'main' into master

bc0c912

Cedric-Perauer mentioned this pull request Jun 29, 2026

Modular pipeline: Krea 2 #14056

Open

Cedric-Perauer changed the title ~~Master~~ Krea2 Modular Pipeline Jun 29, 2026

Cedric-Perauer changed the title ~~Krea2 Modular Pipeline~~ Krea2 Modular Pipeline Support, Documentation, Test Cases Jun 29, 2026

Cedric-Perauer added 2 commits June 29, 2026 21:11

Merge branch 'main' into master

689ee82

Merge branch 'main' into master

d235700

github-actions Bot added the fixes-issue label Jul 2, 2026

Cedric-Perauer added 3 commits July 3, 2026 09:37

Merge branch 'main' into master

5a609ca

Merge branch 'main' into master

3de62bd

Merge branch 'main' into master

988beb9

Merge branch 'main' into master

8022f8e

yiyixuxu requested a review from dg845 July 9, 2026 00:49

dg845 reviewed Jul 9, 2026

View reviewed changes

Cedric-Perauer and others added 3 commits July 9, 2026 10:29

Update src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py

93f59d9

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

Update src/diffusers/modular_pipelines/krea2/before_denoise.py

e6e288f

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

Update docs/source/en/api/pipelines/krea2.md

afd27d3

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

Cedric-Perauer added 3 commits July 9, 2026 09:57

input params, dim multiple of scale factor * patch

3de7fad

docs for krea2 turbo modular

e3b07ab

docs fix

909bcc5

dg845 and others added 3 commits July 9, 2026 17:47

Merge branch 'main' into master

07ccbb2

Merge branch 'main' into master

ec123ba

Merge branch 'main' into master

debb4dd

dg845 reviewed Jul 15, 2026

View reviewed changes

Comment thread src/diffusers/modular_pipelines/krea2/before_denoise.py Outdated

dg845 reviewed Jul 15, 2026

View reviewed changes

Comment thread tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py Outdated

dg845 approved these changes Jul 15, 2026

View reviewed changes

Cedric-Perauer and others added 4 commits July 15, 2026 11:26

Update tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py

0eeca0d

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

Update src/diffusers/modular_pipelines/krea2/before_denoise.py

a21b285

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

Merge branch 'main' into master

421d212

Merge branch 'main' into master

2750680

yiyixuxu reviewed Jul 16, 2026

View reviewed changes

yiyixuxu mentioned this pull request Jul 16, 2026

[agents doc] notes on when to create new blocksets for checkpoint variant #14208

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Krea2 Modular Pipeline Support, Documentation, Test Cases#14083

Krea2 Modular Pipeline Support, Documentation, Test Cases#14083
Cedric-Perauer wants to merge 26 commits into
huggingface:mainfrom
Cedric-Perauer:master

Cedric-Perauer commented Jun 28, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jun 29, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Jul 8, 2026

Uh oh!

Cedric-Perauer commented Jul 8, 2026

Uh oh!

dg845 left a comment

Uh oh!

Uh oh!

dg845 Jul 9, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Cedric-Perauer commented Jul 9, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

dg845 left a comment

Uh oh!

yiyixuxu Jul 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		latent_height = block_state.height // components.vae_scale_factor
		latent_width = block_state.width // components.vae_scale_factor

	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

Uh oh!

Conversation

Cedric-Perauer commented Jun 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Validation

Docs

Before submitting

Who can review?

Uh oh!

github-actions Bot commented Jun 29, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Jul 8, 2026

Uh oh!

Cedric-Perauer commented Jul 8, 2026

Uh oh!

dg845 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dg845 Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Cedric-Perauer commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dg845 left a comment

Choose a reason for hiding this comment

Uh oh!

yiyixuxu Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Cedric-Perauer commented Jun 28, 2026 •

edited

Loading

Cedric-Perauer commented Jul 9, 2026 •

edited

Loading