Skip to content

Tag Krea 2 text encoder outputs as denoiser_input_fields - #14925

Merged
yiyixuxu merged 1 commit into
huggingface:mainfrom
tbsbngrtz:krea2-denoiser-input-fields
Oct 2, 2026
Merged

yiyixuxu merged 1 commit into
huggingface:mainfrom
tbsbngrtz:krea2-denoiser-input-fields

Conversation

@tbsbngrtz

Copy link
Copy Markdown
Contributor

What does this PR do?

Tags the outputs of the two Krea 2 text-encoder blocks (Krea2TextEncoderStep, Krea2TurboTextEncoderStep) with kwargs_type="denoiser_input_fields", as every other modular text encoder does (explicitly in z_image, flux, flux2, wan, stable_diffusion_xl; via OutputParam.template(...) in qwenimage, 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"), then pipeline(**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 with Required input 'prompt_embeds' is missing in Krea2TurboTextInputsStep. Six added lines, no behavioural change for the plain Krea2Pipeline or for the full modular pipeline run in one piece.

Verified: ruff check / ruff format --check clean, utils/check_copies.py clean, tests/modular_pipelines/krea2 46 passed / 2 skipped on CPU; Krea 2 Turbo runs end to end through Mellon (Encode Prompt → Denoise → Decode as separate pipelines sharing a ComponentsManager) 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

  1. Explicit tag vs. template. modular.md ("InputParam / OutputParam") prefers OutputParam.template("prompt_embeds") for canonical names, while gotcha 5 says to keep a plain OutputParam with an accurate description when the semantics differ from the template's. Krea 2's prompt_embeds is 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 is OutputParam.template("prompt_embeds", description="Per-prompt stacked text features (...)") for all six entries — happy to switch. Left for the reviewer.
  2. No regression test. testing.md says to test block behaviour by running the block as a pipeline and not to assert on declared intermediate_outputs, so the natural test would run the text-encoder block standalone with the dummy components from tests/modular_pipelines/krea2/ and assert that get_by_kwargs("denoiser_input_fields") on the returned state contains prompt_embeds and prompt_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.md could add to "Gotchas": declaring a canonical conditioning output with a plain OutputParam(name=...) silently drops the denoiser_input_fields tag — either use the template or set kwargs_type explicitly, and check with state.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

  • 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 run the self-review skill on the diff?
    • Did you share the final self-review notes in the PR description or a comment?
  • 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? (none needed — see "Documentation impact")
  • Did you write any new necessary tests? (see non-blocking item 2)
  • Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

@yiyixuxu (modular pipelines) — the Krea 2 modular blocks were added in #14083.

@github-actions github-actions Bot added modular-pipelines size/S PR with diff < 50 LOC labels Oct 2, 2026

@yiyixuxu yiyixuxu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

thanks

@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

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.

@yiyixuxu
yiyixuxu merged commit 4156630 into huggingface:main Oct 2, 2026
18 of 19 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

modular-pipelines size/S PR with diff < 50 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants