From 22be8138d97061e918927c9f291fb504ab6db990 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Fri, 16 Oct 2026 17:20:42 +0200 Subject: [PATCH] fix(krea2): hold Raw's schedule shift at the end of its fit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Krea-2 Raw's mu kept climbing past 6400 image tokens (2.21 at 2048², 7.41 at 4096²); calculate_shift now holds it at the fit's end, as the FLUX.1, Z-Image and LTX-2 fits do. krea2_denoise 1.3.1; Turbo and an explicit shift are unchanged. --- .../03.Models/Local Models/krea-2.mdx | 9 +++- .../app/invocations/krea2/krea2_denoise.py | 11 +++-- invokeai/backend/krea2/sampling_utils.py | 10 +++- invokeai/frontend/api/openapi.json | 4 +- invokeai/frontend/api/schema.ts | 2 +- tests/app/invocations/test_krea2_denoise.py | 47 ++++++++++++++++++- 6 files changed, 70 insertions(+), 13 deletions(-) diff --git a/docs/src/content/docs/Users Guide/03.Models/Local Models/krea-2.mdx b/docs/src/content/docs/Users Guide/03.Models/Local Models/krea-2.mdx index d87a444e5cd..81fc1496276 100644 --- a/docs/src/content/docs/Users Guide/03.Models/Local Models/krea-2.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/Local Models/krea-2.mdx @@ -1,7 +1,7 @@ --- title: Krea-2 description: Generate images with the Krea-2 text-to-image models (Turbo and Raw), including the GGUF / single-file workflow, the conditioning enhancers and training-free style reference. -lastUpdated: 2026-09-30 +lastUpdated: 2026-10-16 sidebar: order: 7 --- @@ -18,6 +18,13 @@ published checkpoints: The variant is detected automatically on install, and selecting a Krea-2 model sets sensible defaults (steps, CFG, 1024×1024) for that variant. +:::note[Image size and the noise schedule] +Raw puts more of its steps at high noise the larger the image, up to the area of 1280×1280 (about 1.6 +megapixels). Larger images keep that schedule, which is also the one Turbo always uses. Without the limit, a +2048×2048 Raw image would spend about half of its steps above 90% noise, and a 4096×4096 one all of them. +To sample with a different schedule, set **Shift** on the **Denoise - Krea-2** node in a workflow. +::: + ## License Krea-2 Turbo and Raw are released under the **Krea 2 Community License**. Commercial use of the model and its images is allowed only for companies with less than $1M in annual revenue; above that, an Enterprise license is needed. Redistributed copies must add "Krea" to the model name. The full Diffusers repositories are gated: accept the license on the model's HuggingFace page and add your HuggingFace token under **API Keys** in the Model Manager. The quantized GGUF and NVFP4 copies are not gated, but the same license applies. diff --git a/invokeai/app/invocations/krea2/krea2_denoise.py b/invokeai/app/invocations/krea2/krea2_denoise.py index 294fd018397..04aa13a51d5 100644 --- a/invokeai/app/invocations/krea2/krea2_denoise.py +++ b/invokeai/app/invocations/krea2/krea2_denoise.py @@ -136,7 +136,7 @@ def log_summary(self) -> None: title="Denoise - Krea-2", tags=["image", "krea2", "krea-2"], category="image", - version="1.3.0", + version="1.3.1", classification=Classification.Prototype, ) class Krea2DenoiseInvocation(BaseInvocation, WithMetadata, WithBoard): @@ -172,7 +172,7 @@ class Krea2DenoiseInvocation(BaseInvocation, WithMetadata, WithBoard): default=None, description="Override the resolution-aware timestep shift (mu). Leave unset -- or 0, which is what " "the node editor fills in for an untouched optional float -- to use the model default: mu=1.15 for " - "the distilled Turbo checkpoint, resolution-dependent for Raw.", + "the distilled Turbo checkpoint, resolution-dependent for Raw up to the area of 1280x1280 and held above it.", ) style_reference: Optional[Krea2StyleReferenceField] = InputField( default=None, @@ -402,9 +402,10 @@ def _run_diffusion(self, context: InvocationContext): elif self._is_distilled(context): mu = KREA2_DISTILLED_MU else: - # Resolution-aware shift. Honor the loaded scheduler's config so a Raw checkpoint that ships a - # customized scheduler_config.json is sampled with its own shift params, falling back to the - # Krea-2 defaults for the stock scheduler (and for any field the config omits). + # Resolution-aware shift, held at the end of its fit (see `calculate_shift`). Honor the loaded + # scheduler's config so a Raw checkpoint that ships a customized scheduler_config.json is sampled + # with its own shift params. A key that file omits takes the diffusers scheduler's own default, + # not Krea-2's; the Krea-2 defaults below apply only to the stock scheduler built above. mu = calculate_shift( image_seq_len, base_image_seq_len=getattr(scheduler.config, "base_image_seq_len", KREA2_BASE_IMAGE_SEQ_LEN), diff --git a/invokeai/backend/krea2/sampling_utils.py b/invokeai/backend/krea2/sampling_utils.py index 875c7cefd46..4c137609fe3 100644 --- a/invokeai/backend/krea2/sampling_utils.py +++ b/invokeai/backend/krea2/sampling_utils.py @@ -82,14 +82,20 @@ def calculate_shift( base_shift: float = KREA2_BASE_SHIFT, max_shift: float = KREA2_MAX_SHIFT, ) -> float: - """Resolution-aware mu (linear interpolation by sequence length). + """Resolution-aware mu: a line through the two config points, held at the last one. + + Past ``max_image_seq_len`` (6400 tokens, 1280x1280 of area) the line would keep climbing: mu 2.2 at + 2048x2048 puts about half of 28 steps above sigma 0.9, and mu 7.4 at 4096x4096 puts all of them + there, so an img2img start has no step near the sigma it asks for. Held with the stock config, Raw + ends at Turbo's fixed mu, which ComfyUI also uses for Krea-2 at every size; diffusers' + ``Krea2Pipeline`` extrapolates. NOTE: mu is fed straight into ``FlowMatchEulerDiscreteScheduler.set_timesteps(..., mu=mu)``; the exponential time-shift happens inside the scheduler. Do NOT ``exp()`` this value. """ m = (max_shift - base_shift) / (max_image_seq_len - base_image_seq_len) b = base_shift - m * base_image_seq_len - return image_seq_len * m + b + return min(image_seq_len, max_image_seq_len) * m + b def build_sigmas(steps: int) -> List[float]: diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json index 018206e89dc..dc2c1c8c7af 100644 --- a/invokeai/frontend/api/openapi.json +++ b/invokeai/frontend/api/openapi.json @@ -61946,7 +61946,7 @@ } ], "default": null, - "description": "Override the resolution-aware timestep shift (mu). Leave unset -- or 0, which is what the node editor fills in for an untouched optional float -- to use the model default: mu=1.15 for the distilled Turbo checkpoint, resolution-dependent for Raw.", + "description": "Override the resolution-aware timestep shift (mu). Leave unset -- or 0, which is what the node editor fills in for an untouched optional float -- to use the model default: mu=1.15 for the distilled Turbo checkpoint, resolution-dependent for Raw up to the area of 1280x1280 and held above it.", "field_kind": "input", "input": "any", "orig_default": null, @@ -61999,7 +61999,7 @@ "tags": ["image", "krea2", "krea-2"], "title": "Denoise - Krea-2", "type": "object", - "version": "1.3.0", + "version": "1.3.1", "output": { "$ref": "#/components/schemas/LatentsOutput" } diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts index 733984fe92e..af053c11b34 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -24310,7 +24310,7 @@ export type components = { seed?: number; /** * Shift - * @description Override the resolution-aware timestep shift (mu). Leave unset -- or 0, which is what the node editor fills in for an untouched optional float -- to use the model default: mu=1.15 for the distilled Turbo checkpoint, resolution-dependent for Raw. + * @description Override the resolution-aware timestep shift (mu). Leave unset -- or 0, which is what the node editor fills in for an untouched optional float -- to use the model default: mu=1.15 for the distilled Turbo checkpoint, resolution-dependent for Raw up to the area of 1280x1280 and held above it. * @default null */ shift?: number | null; diff --git a/tests/app/invocations/test_krea2_denoise.py b/tests/app/invocations/test_krea2_denoise.py index a2ca338534c..d1cc6bbfa36 100644 --- a/tests/app/invocations/test_krea2_denoise.py +++ b/tests/app/invocations/test_krea2_denoise.py @@ -775,16 +775,22 @@ def _mu_branch_context(tmp_path, transformer: _Transformer, variant=Krea2Variant ) -def _install_mu_capturing_scheduler(monkeypatch, captured: list[float], config_kwargs: dict) -> None: +def _install_mu_capturing_scheduler( + monkeypatch, captured: list[float], config_kwargs: dict, stop: bool = False +) -> None: class _MuCapturingScheduler: def __init__(self, **_kwargs) -> None: # The real diffusers scheduler exposes shift params on .config (a FrozenDict, attribute-access); # the denoise reads them via getattr(scheduler.config, , ). Omitting a key - # here (SimpleNamespace without the attr) exercises the fallback-to-default path. + # here (SimpleNamespace without the attr) exercises that fallback, which only the stock scheduler + # reaches: a real scheduler_config.json missing a key gets the diffusers scheduler's default. self.config = SimpleNamespace(num_train_timesteps=1000, **config_kwargs) def set_timesteps(self, *, sigmas, mu, device) -> None: captured.append(mu) + # A frame of real size need not be denoised to know the mu it was handed. + if stop: + raise _StopAtTimesteps self.sigmas = torch.tensor([*sigmas, 0.0], device=device) self.timesteps = self.sigmas[:-1] * self.config.num_train_timesteps @@ -813,6 +819,43 @@ def test_run_diffusion_mu_honors_scheduler_config_shift_params(monkeypatch, tmp_ assert captured[0] != pytest.approx(calculate_shift(image_seq_len)) +class _StopAtTimesteps(Exception): + """Ends a run once the scheduler has been handed its mu.""" + + +@pytest.mark.parametrize( + ("width", "height", "mu"), + [ + (512, 512, 0.58125), # 1024 tokens: on the line + (1024, 1024, 0.90625), # 4096 tokens + (2048, 768, 1.1229166666666667), # 6144 tokens: area decides, not a side + (1280, 1280, 1.15), # 6400 tokens: the line's end + (2048, 2048, 1.15), # 16384 tokens: would be 2.206 extrapolated + (4096, 4096, 1.15), # 65536 tokens: would be 7.41 + ], +) +def test_raw_mu_follows_the_fit_up_to_its_end_and_holds_there(width: int, height: int, mu: float) -> None: + from invokeai.backend.krea2.sampling_utils import calculate_shift + + # One token per 16x16 pixels, as `_run_diffusion` counts them. + assert calculate_shift((width // 16) * (height // 16)) == pytest.approx(mu) + + +def test_raw_mu_holds_at_the_end_of_a_custom_scheduler_fit(monkeypatch, tmp_path) -> None: + # A checkpoint's own scheduler_config.json moves the end of the fit, and the hold moves with it. + _patch_runtime(monkeypatch) + captured: list[float] = [] + custom = {"base_image_seq_len": 256, "max_image_seq_len": 1024, "base_shift": 0.3, "max_shift": 1.5} + _install_mu_capturing_scheduler(monkeypatch, captured, custom, stop=True) + invocation = _mu_branch_invocation() + invocation.width = invocation.height = 2048 + + with pytest.raises(_StopAtTimesteps): + invocation._run_diffusion(_mu_branch_context(tmp_path, _Transformer())) + + assert captured == [pytest.approx(1.5)] + + def test_run_diffusion_mu_falls_back_to_krea2_defaults_for_absent_config_keys(monkeypatch, tmp_path) -> None: # A scheduler config that omits some shift params must fall back to the Krea-2 defaults for exactly the # missing keys, honoring the ones it does provide.