Tag Krea 2 text encoder outputs as denoiser_input_fields - #14925
Merged
Merged
Conversation
|
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. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Tags the outputs of the two Krea 2 text-encoder blocks (
Krea2TextEncoderStep,Krea2TurboTextEncoderStep) withkwargs_type="denoiser_input_fields", as every other modular text encoder does (explicitly inz_image,flux,flux2,wan,stable_diffusion_xl; viaOutputParam.template(...)inqwenimage,ltx,helios).Without the tag, the documented split-pipeline pattern in
docs/source/en/modular_diffusers/modular_pipeline.md("Running a modular pipeline" →text_encoder_pipeline(prompt=...).get_by_kwargs("denoiser_input_fields"), thenpipeline(**text_embeddings, ...)) returns an empty dict for Krea 2, and any consumer that relies on the tag (Mellon's Encode Prompt node does exactly this) hands the denoiser nothing, which fails withRequired input 'prompt_embeds' is missinginKrea2TurboTextInputsStep. Six added lines, no behavioural change for the plainKrea2Pipelineor for the full modular pipeline run in one piece.Verified:
ruff check/ruff format --checkclean,utils/check_copies.pyclean,tests/modular_pipelines/krea246 passed / 2 skipped on CPU; Krea 2 Turbo runs end to end through Mellon (Encode Prompt → Denoise → Decode as separate pipelines sharing aComponentsManager) on an RTX 5070 Ti with this change.Fixes # (no issue filed; found while adding Krea 2 to a Mellon fork)
Self-review notes (
.ai/skills/self-review, final round)Rubric:
.ai/references/review-rules.md,modular.md,code_style.md,testing.md.Blocking issues — none.
Non-blocking issues
modular.md("InputParam / OutputParam") prefersOutputParam.template("prompt_embeds")for canonical names, while gotcha 5 says to keep a plainOutputParamwith an accurate description when the semantics differ from the template's. Krea 2'sprompt_embedsis a 4-D stack of selected decoder layers(B, text_seq_len, num_text_layers, text_hidden_dim), not the usual(B, seq, dim), so the original author's plain declaration with the precise shape in the description looks intentional; this PR keeps those descriptions and only adds the tag. If you prefer the template form, the equivalent isOutputParam.template("prompt_embeds", description="Per-prompt stacked text features (...)")for all six entries — happy to switch. Left for the reviewer.testing.mdsays to test block behaviour by running the block as a pipeline and not to assert on declaredintermediate_outputs, so the natural test would run the text-encoder block standalone with the dummy components fromtests/modular_pipelines/krea2/and assert thatget_by_kwargs("denoiser_input_fields")on the returned state containsprompt_embedsandprompt_embeds_mask(and the negative pair for the base step). No existing modular test covers the tag for any model, so this would be the first; I did not add it to keep the PR to the fix. Left for the reviewer — say the word and I'll add it.Dead code (advisory) — not applicable (no new model; the change touches declarations only).
Documentation impact — no usage doc describes Krea 2's text-encoder outputs, so nothing goes stale; the generic split-pipeline example above becomes correct for Krea 2. Agent-docs suggestion:
modular.mdcould add to "Gotchas": declaring a canonical conditioning output with a plainOutputParam(name=...)silently drops thedenoiser_input_fieldstag — either use the template or setkwargs_typeexplicitly, and check withstate.get_by_kwargs("denoiser_input_fields")after a standalone run. Not included in this PR.Summary — READY. Fix before submitting: nothing. Leave for the actual review: items 1 and 2 above.
Before submitting
self-reviewskill on the diff?Who can review?
@yiyixuxu (modular pipelines) — the Krea 2 modular blocks were added in #14083.