feat: ernie image/turbo

[Pfannkuchensack]

Summary

Adds Baidu ERNIE-Image and ERNIE-Image-Turbo (HF, HF Turbo) as a new BaseModelType, mirroring the FLUX.2 Klein and Z-Image integration patterns.

The two checkpoints share the same ErnieImageTransformer2DModel architecture (3072 hidden, 24 layers, 24 heads, Mistral3 text encoder, AutoencoderKLFlux2 VAE) and only differ in inference defaults (50 steps + CFG 4.0 vs. 8 steps + CFG 1.0 for Turbo), so they live under one BaseModelType.ErnieImage without a variant enum. The optional 3B Mistral3-based prompt enhancer that ships with the pipeline is wired through with a UI toggle.

ERNIE-Image is text-to-image only in this PR — its denoise node has no denoise_mask input, so masked modes (inpaint/outpaint) are unsupported, and image-to-image is not wired.

This PR builds on top of #8859 (transformers 5.1+) and additionally bumps diffusers 0.36.0 → 0.38.0, which is the first release containing ErnieImagePipeline and ErnieImageTransformer2DModel.

Backend changes

• BaseModelType.ErnieImage, ModelType.PromptEnhancer, plus two new SubModelTypes (pe, pe_tokenizer) for the bundled prompt enhancer.
• Main_Diffusers_ErnieImage_Config and Main_Checkpoint_ErnieImage_Config with state-dict-based detection (x_embedder + text_proj + adaLN_modulation).
• A diffusers loader that loads transformer / vae / text_encoder / tokenizer plus optional pe / pe_tokenizer from a single pipeline directory.
• invokeai/backend/ernie_image/ with sampling utilities (2×2 patchify, BN normalize/denormalize, sigma schedule, padded text packing) and a rectified-flow denoise loop supporting Euler / Heun / LCM (reusing the FlowMatch* schedulers from FLUX).
• Four new invocations: ernie_image_model_loader, ernie_image_text_encoder (with prompt-enhancer toggle), ernie_image_denoise, ernie_image_vae_decode.
• ErnieImageConditioningInfo + conditioning field/output, including pickle allowlist for the disk serializer.
• Generation mode ernie_image_txt2img. ERNIE-Image is text-to-image only — its denoise node has no denoise_mask input, so masked modes (inpaint/outpaint) are unsupported, and image-to-image is not wired.
• Starter models for baidu/ERNIE-Image and baidu/ERNIE-Image-Turbo, plus a STARTER_BUNDLES entry.

Frontend changes

• services/api/schema.ts regenerated to expose the new node types.
• Type unions extended (ImageOutput, LatentToImage, DenoiseLatents, MainModelLoaderNodes). New MaskableDenoiseNodes type excludes ernie_image_denoise from the inpaint/outpaint helpers (ERNIE has no denoise_mask).
• ParamsState gains ernieImageScheduler and ernieImageUsePromptEnhancer, with reducers, selectors, and selectIsErnieImage.
• buildErnieImageGraph (txt2img only) wired into useEnqueueCanvas.
• ERNIE entries added to MODEL_BASE_TO_{COLOR,LONG_NAME,SHORT_NAME}; prompt_enhancer added to MODEL_TYPE_TO_LONG_NAME.
• ParamErnieImageScheduler and ParamErnieImagePromptEnhancer rendered conditionally in GenerationSettingsAccordion.
• addTextToImage/addImageToImage rectified-flow type guards accept ernie_image_denoise; isMainModelWithoutUnet likewise. addInpaint/addOutpaint take MaskableDenoiseNodes, which excludes ERNIE.

Diffusers 0.38 fix-out

• hotfixes.py: import LoRACompatibleConv directly. The lazy module loader in 0.38 no longer exposes diffusers.models.lora as an attribute, so the legacy patch path crashes on import.
• pyproject.toml declares prerelease = "allow" under [tool.uv], with a comment explaining that diffusers 0.38.0 itself hard-pins safetensors>=0.8.0-rc.0. We can drop this again once a diffusers patch release ships with a stable safetensors floor.

Related Issues / Discussions

• Builds on Update to transformers 5.1.0 #8859 (transformers 5.1.0).
• Diffusers 0.38.0 release notes: https://github.com/huggingface/diffusers/releases/tag/v0.38.0
• Baidu ERNIE-Image: https://huggingface.co/baidu/ERNIE-Image / https://huggingface.co/baidu/ERNIE-Image-Turbo
• Closes [enhancement]: Implement Ernie Image #9059

QA Instructions

Automated

• pytest tests/ -m "not slow" — passes (619 / 619, the same tests/model_identification LFS-skip as on main).
• pnpm lint:tsc, pnpm lint:eslint, pnpm lint:prettier — all clean.
• pnpm test:no-watch — 563 / 563 frontend tests pass.

Manual smoketest of pre-existing models (regression for the diffusers / transformers bump)

Pick at least two of these and run a single txt2img each. Confirm no crashes and visually-plausible output:

• SDXL
• FLUX.1 Dev
• FLUX.2 Klein
• SD3
• Z-Image-Turbo

Plus the relevant items from #8859's test plan: SD 1.5 prompt-weighted generation (compel path), FLUX text-to-image (T5 tokenizer path), HF model install via repo ID, NSFW checker first-time download.

Manual smoketest of ERNIE-Image (requires a GPU — 8B parameters)

1. Install baidu/ERNIE-Image-Turbo from the new starter bundle. The Model Manager should classify it as BaseModelType.ErnieImage.
2. Pick the model on the Generate tab. Confirm the scheduler dropdown and Prompt Enhancer toggle appear in the Generation accordion.
3. txt2img at 1024×1024, 8 steps, CFG 1.0 — should produce an image.
4. Toggle the prompt enhancer on with a short prompt (e.g. "a fox"); the enhancer log line in the backend should show a rewritten longer prompt before encoding.
5. Repeat with baidu/ERNIE-Image (50 steps, CFG 4.0).

Merge Plan

• This PR depends on Update to transformers 5.1.0 #8859 being merged first (or being merged in together, since they share the transformers>=5.1.0 override). I have no preference; happy to rebase whenever Update to transformers 5.1.0 #8859 lands.
• After this PR merges, a follow-up release should call out the diffusers and transformers major bumps in the changelog. The ERNIE-Image starter bundle is gated behind those bumps.
• The prerelease = "allow" line in pyproject.toml is a temporary measure tied to diffusers 0.38.0's safetensors>=0.8.0-rc.0 upstream pin. Worth revisiting (and removing) once a diffusers patch release relaxes that requirement.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable) — backend regression suite covers the new code paths via existing model-config and loader-registry tests; no new dedicated unit tests for the ERNIE sampling helpers (the patchify roundtrip was sanity-checked manually).
• ❗Changes to a redux slice have a corresponding migration — N/A: only additive fields with defaults in paramsSlice.
• Documentation added / updated (if applicable) — none yet; the docs-old/contributing/NEW_MODEL_INTEGRATION.md checklist describes the integration shape this PR follows.
• Updated What's New copy (if doing a release after this PR)

Out of scope (planned follow-ups)

• img2img / inpaint / outpaint — ERNIE-Image is text-to-image only in this PR. inpaint/outpaint would need a denoise_mask input on the backend denoise node; img2img could be added later (the denoise node already accepts optional starting latents).
• ControlNet, IP-Adapter, and LoRA support for ERNIE-Image.
• Single-file checkpoint loading (the Main_Checkpoint_ErnieImage_Config is in place as defensive scaffolding, but the loader currently raises NotImplementedError for that format).
• Metadata recall handlers in the gallery side panel for ERNIE-specific parameters.

[lstein]

Thanks for this — the ERNIE integration follows the established Z-Image/FLUX.2 shape closely, and the sampling code is a faithful port of the upstream pipeline (I diffed it against ErnieImagePipeline/ErnieImageTransformer2DModel line by line; see "Verified clean" at the bottom for what held up).

That said, I think the feature is currently non-functional on two independent counts, and there's a third issue that affects users who never touch ERNIE. Details below, reviewed at head 7696bf7f79.

---

Blockers

1. diffusers is pinned to 0.37.0 — the ERNIE classes don't exist, so no ERNIE model can load

The PR description's central premise is the 0.36 → 0.38 bump, because 0.38.0 is the first release containing ErnieImagePipeline and ErnieImageTransformer2DModel. That bump is no longer in the branch: pyproject.toml:39 reads diffusers[torch]==0.37.0, and uv.lock resolves 0.37.0. It looks like it was lost in one of the Merge remote-tracking branch 'upstream/main' merges — main moved 0.36 → 0.37 and won the conflict.

Verified in the branch's own venv:

diffusers 0.37.0
has ErnieImageTransformer2DModel: False
has ErnieImagePipeline: False

The failure mode is unpleasant: identification only reads _class_name out of model_index.json, so install succeeds. The error only surfaces on first generate, when GenericDiffusersLoader._hf_definition_to_type (generic_diffusers.py:88) calls getattr(diffusers, "ErnieImageTransformer2DModel") and raises AttributeError. So a user downloads an 8B pipeline and then hits a hard error.

2. The prompt enhancer is unconditionally disabled — config.path is relative

invokeai/app/invocations/ernie_image_model_loader.py:83-89:

config = context.models.get_config(self.model)
pe_dir = Path(config.path) / "pe"
return pe_dir.is_dir()

context.models.get_config() returns the raw record store config. Models inside the Invoke-managed models dir are stored with relative paths (model_install_default.py:1086-1088: "Models in the Invoke-managed models dir should use relative paths."). Only ModelLoader._load_and_cache rewrites config.path to absolute (load_default.py:132), and that happens well after this check.

So pe_dir.is_dir() resolves against the server process's CWD and returns False for every model installed through the starter bundle or a normal HF install. prompt_enhancer is always None, and the UI toggle is a no-op. QA step 4 in the description can't pass as written.

Fix: resolve against models_path before the check (or reuse the loader's own path resolution).

3. paramsSlice: two new required fields with no defaults and no migration — every existing user's params slice resets on upgrade

features/controlLayers/store/types.ts:827-829:

  ernieImageScheduler: zParameterErnieImageScheduler,
  ernieImageUsePromptEnhancer: z.boolean(),
  // Defaults make these resilient to rehydration of persisted state saved before the fields existed.

The two new fields landed directly above the comment that explains why they need .default() — and don't have it. _version stays at 3, so no migration branch adds them, and migrate ends at zParamsState.parse(state) (paramsSlice.ts:795). A persisted v3 state written by current main lacks both keys, so parse throws → the catch in unserialize (app/store/store.ts:166-172) → the whole slice falls back to getInitialState(). Model selection, dimensions, prompt history, every generation parameter: reset.

I verified this rather than just reading it — a throwaway vitest probe that deletes the two keys from a v3 state and calls migrate() reports MIGRATE_THREW: true.

This is the checklist item marked N/A as "only additive fields with defaults in paramsSlice" — the fields are additive, but they don't have defaults. Adding .default('euler') and .default(true) (matching the ideogram4* fields immediately below) is sufficient; no _version bump needed.

---

High

4. Heun runs 2N−1 iterations against total_steps = N, so progress overruns to ~190%

backend/ernie_image/denoise.py:59 sets total_steps = len(timesteps) - 1, but line 78 loops for step_index in range(len(scheduler.timesteps)).

FlowMatchHeunDiscreteScheduler.set_timesteps doesn't accept sigmas, so line 75 falls back to num_inference_steps=len(sigmas). Verified against the installed diffusers:

heun:  requested steps 8 -> timesteps 15, sigmas 16
euler: sigmas in 8       -> timesteps 8,  sigmas 9

With steps=8 the loop runs 15 times and emits step=9..15 against total_steps=8. The same fallback also silently discards the custom sigma schedule (and therefore denoising_start / denoising_end) whenever Heun is selected.

5. getGridSize has no 'ernie-image' case, but the backend enforces multiple_of=16

features/parameters/util/optimalDimension.ts:73-90 falls through to default: 8, and that grid size drives bbox snapping (canvasSlice.ts:1214+, paramsSlice.ts:433+). Meanwhile ernie_image_denoise.py:58-59 declares multiple_of=16 on width/height.

Concrete trigger: select an ERNIE model, set the bbox width to 1032 (legal at grid 8), invoke → the backend rejects the graph with a validation error.

Ideogram-4 handles exactly this with both a getGridSize entry and a readiness check (readiness.ts:764-790); ERNIE has neither.

6. The loader drops _apply_fp8_layerwise_casting

model_loaders/ernie_image.py:57-62 returns load_class.from_pretrained(...) directly. Both comparable loaders apply the casting pass before returning — z_image.py:167 and qwen_image.py:248.

_should_use_fp8 gates on default_settings.fp8_storage (load_default.py:198-202), which is a user-facing Model Manager toggle and is reachable for ERNIE Main configs. As written, the toggle will show up and silently do nothing.

---

Medium

7. prerelease = "allow" and the transformers>=5.1.0 override are now stale

The stated justification for prerelease = "allow" is diffusers 0.38 pinning safetensors>=0.8.0-rc.0 — that no longer applies, since diffusers is 0.37.0 and safetensors resolves to stable 0.8.0. And main already declares transformers>=5.5,<5.6, so the >=5.1.0 override adds nothing while, as a global uv override, weakening that deliberate upper bound on the next relock.

prerelease = "allow" is repo-wide, not scoped to safetensors: any future uv lock may pick up prereleases of torch, transformers, or anything else. I'd drop both lines.

8. Schedulers are default-constructed, ignoring the checkpoint's scheduler/ config

ernie_image_denoise.py:130-131 does scheduler_cls(). Upstream uses the scheduler that ships with the pipeline. The diffusers defaults are shift=1.0, num_train_timesteps=1000; FlowMatchEulerDiscreteScheduler.set_timesteps applies shift to the sigmas you pass in, so if Baidu ships shift != 1.0, use_dynamic_shifting, or a different num_train_timesteps, output silently diverges from the reference implementation. model_timestep_scale (denoise.py:65) is derived from that same config.

9. Initial noise is generated on the compute device in bf16

ernie_image_denoise.py:116-122. Every other denoise node in the repo generates on CPU in fp16 and then casts, with an explicit comment — sd3_denoise.py:160-172, z_image_denoise.py:217-227, flux2/sampling_utils.py:37-55, latent_noise.py:

# We always generate noise on the same device and dtype then cast to ensure consistency across devices/dtypes.
rand_device = "cpu"
rand_dtype = torch.float16

As written, the same seed produces different images on CUDA vs ROCm vs MPS vs CPU, and shouldUseCpuNoise is ignored.

10. The starter bundle installs both 8B pipelines

starter_models.py:1913 — ernie_image_bundle contains Turbo and base. Every other bundle is one main model plus its shared dependencies. Turbo alone seems like the better default given the download size.

11. The checkpoint config installs models that can never load

Main_Checkpoint_ErnieImage_Config matches and installs successfully, and then ErnieImageDiffusersModel._load_model raises NotImplementedError (ernie_image.py:40-41). I'd either not register it yet, or keep detection strict enough that nothing unloadable can match — as-is it's a way to end up with a permanently broken entry in the Model Manager.

---

Low / nits

• Latent preview uses the pre-step sigma with the post-step sample. denoise.py:102 computes img - t_sigma * pred, but img was already stepped at line 93. Using t_prev gives the exact x0 estimate. Same at line 139.
• Preview RGB factors are applied to BN-normalized latents. step_callback.py uses the FLUX.2 factors, but ERNIE's denoise space is BN-normalized (denormalized only at ernie_image_vae_decode.py:59), so previews will be off in color/contrast. Probably unavoidable since the BN stats live on the VAE, but the comment currently claims the factors "apply directly" — worth correcting.
• LatentsOutput.build reports half the real size. ernie_image_denoise.py:151 — the latents are patched, so width = shape[3] * 8 yields 512 for a 1024 image. Node-editor-only, but wrong.
• _SUBDIR_OVERRIDES is a no-op. ernie_image.py:30-33 maps SubModelType.PromptEnhancer -> "pe", but SubModelType.PromptEnhancer.value is already "pe" (taxonomy.py:107-108). Same for the tokenizer. The dict and its comment can go.
• Unsupported generation modes use assert rather than UnsupportedGenerationModeError, which produces a red "Failed to build graph" error instead of the friendlier warning toast useEnqueueCanvas.ts:93-96 exists for. (Ideogram-4 does the same thing, so this is precedent-consistent — just noting the nicer path exists.)
• Hardcoded English label in ParamErnieImagePromptEnhancer.tsx (<FormLabel>Prompt Enhancer</FormLabel>); the sibling scheduler component uses t(...).
• The addImageToImage ERNIE branch is dead code — the graph builder asserts txt2img before it can be reached.
• Turbo detection by name substring (configs/main.py:87-92, "turbo" in name.lower()) means renaming the model on install loses the 8-step / CFG-1.0 defaults. Z-Image models this with a variant enum.

---

Verified clean

Things I specifically attacked that held up, so you don't have to re-check them:

• Separate positive/negative text padding. Upstream concatenates uncond+cond before _pad_text so both share one Tmax; this PR pads them separately with different Tmax and runs two forwards. That's equivalent — the transformer builds attention_mask from text_lens (transformer_ernie_image.py:415-422) and image-token RoPE ids from text_lens (lines 409-411), both independent of Tmax. Two forwards also use less memory than the batched version.
• Noise shape. VAE_SCALE_FACTOR = 16 with latent_h = height // 16 looks like a 4x spatial error, since AutoencoderKLFlux2's 32-channel latents live at H/8 (cf. flux2/sampling_utils.py:44). It's correct — upstream defines vae_scale_factor = 2 ** len(block_out_channels) = 16 and uses height // vae_scale_factor for the patched grid. Confusing name, right math.
• Patchify / unpatchify roundtrip — identical to upstream's _patchify_latents / _unpatchify_latents.
• Sigma schedule — get_schedule + timesteps[:-1] reproduces upstream's linspace(1, 0, N+1)[:-1].
• bf16 timestep vector — upstream does the same thing, so not something this PR introduces.
• ModelPicker ordering — 'ernie-image' isn't in the ordered group list, but ungrouped bases are appended afterwards, same as anima/flux2. No problem.
• Loader registration (auto-imported via load/__init__.py:15-17), get_size_fs (submodel_type.value is the real subdir name for pe/pe_tokenizer), from_base signature change (only two callers; name defaults to None), and from diffusers.models.lora import LoRACompatibleConv (imports fine on 0.37.0) — all fine.
• openapi.json / schema.ts — I regenerated openapi.json locally and it's byte-identical to what's committed.

---

Repo / CI state

• The branch currently conflicts with main in invocations/model.py, util/step_callback.py, configs/main.py, starter_models.py, and openapi.json.
• Local pytest tests/ -m "not slow" gives 9 failed / 2730 passed (reproduced twice). The 4 test_qwen_image_state_dict_utils failures are a transformers-5.5 API change (_checkpoint_conversion_mapping) and main's lock resolves the same 5.5.4, so they're pre-existing; the 5 test_node_graph failures pass when that file is run alone, so they're cross-file registration pollution. Neither group looks attributable to this PR — but the description's "619 / 619" figure is stale and worth refreshing.

---

Requesting changes primarily on 1-3; the rest can be triaged. Happy to re-review once the diffusers pin and the PE path resolution are sorted.

[lstein]

Re-reviewed the commits added since my last round (7696bf7f79..4d39efba3a). All three blockers, all three High findings, and every Medium and nit from the last round are addressed — thanks for the thorough pass, and for the regression tests. I re-verified each fix rather than taking the commit message at its word; details below, then two new items the fp8 fix opened up.

Note for future rounds: the fixes are all inside the merge commit 81a920d7ce, so git log 7696bf7f79..HEAD shows nothing but merges and git diff against either parent is dominated by main's churn. git show --cc 81a920d7ce is the only view that isolates them. A separate fix commit before the merge would be easier to review (and to bisect).

---

Verified fixed

1. diffusers pin. Resolved — main itself moved to diffusers[torch]==0.39.0 in #9304, and this branch dropped the local prerelease = "allow" / transformers>=5.1.0 overrides in favour of it. pyproject.toml and uv.lock are now byte-identical to origin/main, uv lock --locked is clean, and in the branch venv from diffusers import ErnieImagePipeline, ErnieImageTransformer2DModel imports. This also closes Medium #7.

2. Prompt-enhancer path. (models_path / config.path).resolve() / "pe" matches ModelLoader._get_model_path exactly. I checked the two cases that could still bite it and both hold: an in-place install stores an absolute config.path, and Path(models_path) / <absolute> yields the absolute path unchanged (true on Windows too, since WindowsPath handles drive-absolute joins); and _load_and_cache mutates config.path to absolute in place (load_default.py:160), so on a second generate the join is idempotent rather than double-prefixed.

3. paramsSlice defaults. .default('euler') / .default(true) plus the rehydration test. Whole frontend suite green (139 files / 1686 tests).

4. Heun step accounting. Driving total_steps off len(scheduler.timesteps) is right, and I confirmed the sigma indexing can't run off the end: Heun's set_timesteps(N) produces 2N-1 timesteps against 2N sigmas, so scheduler.sigmas[step_index + 1] tops out at the last real index. Ran the loop end-to-end against a stub transformer for euler / heun / lcm at steps ∈ {1, 2, 8} with CFG on — no overrun, no IndexError, finite output, and step == total_steps on the final callback in every case (heun 8 → 15 iterations, as the test asserts). The refusal on a partial range is the right call, and the t_prev preview lines up with Heun's duplicated sigmas: at the corrector substep sigmas[i+1] is the same sigma the corrector lands on.

5. getGridSize + readiness. Both present, and getScaledBoundingBoxDimensions picks the 16 up automatically, so the scaleMethod: 'auto' path is covered as well as the explicit 'none' / 'manual' checks.

6. fp8 casting. Applied — see the new finding below for the one submodel it now reaches that it shouldn't.

8. Scheduler from the checkpoint's config. Confirmed the config actually feeds the math: FlowMatchEulerDiscreteScheduler(shift=3).set_timesteps(sigmas=[1, .8, .6, .4, .2]) gives [1, .923, .818, .667, .429] rather than the identity, and Heun honours shift through its own set_timesteps(N) too. Loading an Euler-authored scheduler_config.json into the Heun or LCM class works (diffusers drops the unknown keys), so switching schedulers in the UI doesn't blow up.

9–11 and the nits. CPU noise (fp32 rather than the repo's usual fp16, which is fine — the point was device-independence), Turbo-only bundle, Main_Checkpoint_ErnieImage_Config removed along with its now-dead _has_ernie_image_keys helper, post-step sigma in the preview, honest step_callback comment about BN-normalized latent space, real LatentsOutput size, _SUBDIR_OVERRIDES gone, UnsupportedGenerationModeError, translated PE label. All confirmed.

---

New (introduced by the fp8 fix)

A. fp8 layerwise casting now hits the prompt-enhancer LM — blocker

model_loaders/ernie_image.py:55 routes every submodel through _apply_fp8_layerwise_casting, and _should_use_fp8's exclusion set (load_default.py:255-267) lists the text encoders and tokenizers but not SubModelType.PromptEnhancer — which this PR introduces.

Trigger: install ERNIE-Image, flip FP8 storage on in Model Manager (it's rendered for every main base except Z-Image, MainModelDefaultSettings.tsx:151 — and it's exactly what someone squeezing an 8B transformer into VRAM will do), leave the prompt-enhancer toggle at its default true, generate on CUDA/ROCm. The pe submodel is a Ministral3ForCausalLM, so every nn.Linear and nn.Embedding in it gets fp8 storage plus the pre/post-forward cast hooks. _enhance_prompt then calls lm.generate(), which is one full forward per generated token — so the whole LM is cast bf16↔fp8 on every token of the rewritten prompt, on top of the fp8 rounding of a model whose entire job is text quality.

Fix: add SubModelType.PromptEnhancer to _excluded_submodel_types, alongside the existing "Don't apply FP8 to text encoders" group. Worth adding PromptEnhancerTokenizer for symmetry — it's harmless today only by accident, because _apply_fp8_layerwise_casting bails on anything that isn't an nn.Module.

B. Turbo detection now matches on the install path — nice-to-have

configs/main.py:96-102 joins name and path into one haystack. For an Invoke-managed install config.path is relative and the risk is negligible, but an in-place install records an absolute path, so any ancestor directory containing "turbo" — /mnt/turbo-nvme/models/ERNIE-Image/ — silently gives the base model 8 steps and CFG 1.0. Matching Path(path).name instead of the whole string keeps the rename-resilience you wanted without the ancestor-directory surface.

---

Nits (all nice-to-have)

• ernie_image.py:48 omits local_files_only=True, which GenericDiffusersLoader and every sibling loader (z_image.py, krea2.py, qwen_image.py) pass. It also hands torch_dtype=/variant= to the tokenizer and pe_tokenizer loads, where the siblings special-case AutoTokenizer.from_pretrained(path, local_files_only=True). I checked and transformers 5.14 accepts and ignores both kwargs (it does warn that torch_dtype is deprecated in favour of dtype, but that's repo-wide), so this is style/offline-robustness, not a bug.
• Heun now reports 2N-1 as the total, so asking for 8 steps shows a 15-step progress bar. Correct accounting, and I'd rather have that than the overrun — just flagging that it will read as surprising.
• _enhance_prompt's max_new_tokens=tokenizer.model_max_length is a faithful port of upstream _enhance_prompt_with_pe, but if the PE tokenizer config omits model_max_length the HF sentinel makes that effectively unbounded and a non-terminating rewrite hangs the graph. A literal cap would be cheap insurance. Not a regression — noting it because it's newly reachable now that the enhancer actually runs.

---

Repo / CI state

The branch conflicted with main in addImageToImage.ts and schema.ts. I've resolved both and pushed the merge to this branch as a68a02a5c8 — shout if you'd rather redo it yourself and I'll drop it.

• addImageToImage.ts — main's PiD work retyped l2i and the return value from LatentToImageNodes to ImageOutputNodes; this branch had factored the inline i2l union into an ImageToLatentsNodes alias. Both survive: ImageOutputNodes for l2i/return, ImageToLatentsNodes for i2l. Nothing to add to the i2l union, since ERNIE has no VAE-encode node.
• schema.ts regenerated rather than merged. The auto-merged openapi.json turned out byte-identical to the regenerated file.

On the merged tree: pytest tests/ is 4040 passed / 0 failed — the 9 failures I reported last round were a stale local venv, not real. Frontend tsc / eslint / prettier / vitest all clean, ruff check + ruff format --check clean, uv lock --locked clean, and regenerating schema.ts a second time is a no-op (CI typegen parity).

---

Leaving this as a comment rather than a formal review state: (A) is the one thing I'd want fixed before merge, and it's a two-line change to the exclusion set. (B) and the nits are all nice-to-haves — take or leave them. Everything else looks good to me.

[lstein]

Re-reviewed a68a02a5c8..74ae2a1144. The blocker and both nice-to-haves from last round are fixed, and the self-review pass on top of them found four real problems I had missed — the denoising_end one in particular is a genuine sampling bug, well diagnosed. Nice work. Details below, then one new blocker (pre-existing, not from these commits) and two smaller items.

---

Verified fixed

A. fp8 no longer reaches the prompt enhancer. SubModelType.PromptEnhancer and PromptEnhancerTokenizer are in _excluded_submodel_types (load_default.py:263-264), and the test pins both directions — excluded submodel False, transformer still True on the same fp8-enabled config, so it can't pass by the config simply failing to opt in.

B. Turbo detection. Path(path).name (configs/main.py:103); Path was already imported. /mnt/turbo-nvme/models/ERNIE-Image → 50 steps / CFG 4.0, .../ERNIE-Image-Turbo → 8 / 1.0, both pinned.

Nits. local_files_only=True on every from_pretrained, with tokenizers now loaded bare like the sibling loaders. I checked the one thing that could have broken silently: the variant-fallback in _load_model keys off "no file named" in str(e), and adding local_files_only=True doesn't change that message — a local dir missing *.fp16.safetensors still raises Error no file named diffusion_pytorch_model.fp16.bin found in directory ... either way, so the fallback still fires. max_new_tokens=min(tokenizer.model_max_length, PE_MAX_NEW_TOKENS) closes the int(1e30) sentinel hole.

Verified — self-review fixes

denoising_end is now honored. This is the important one and the fix is right. Passing the whole window and truncating the scheduler's appended zero — rather than patching the last sigma by hand — is what keeps the scheduler's own shift on the terminal sigma. With shift=3, denoising_end 0.75 / 0.5 / 0.25 land on 0.500 / 0.750 / 0.900, exactly 3σ/(1+2σ) of the requested raw end sigma.

The thing I most wanted to confirm is that this doesn't perturb the default path, since upstream ErnieImagePipeline does set_timesteps(sigmas=sigmas[:-1]) (pipeline_ernie_image.py:325) and you now do something different. It doesn't: for denoising_end=1.0 the requested terminal sigma is 0, shift(0) == 0, so hand-the-whole-window-then-truncate lands on the same tensors. I checked torch.equal on both sigmas and timesteps for shift ∈ {1.0, 3.0, 6.0} × N ∈ {1, 2, 8, 50} — identical in all 12 cases.

Ran the loop end-to-end against a stub transformer over {euler, euler shift=3, lcm} × steps {1, 2, 8} × windows {(0,1), (0,.75), (0,.5), (.5,1), (.25,.75)}: no IndexError from the truncated sigmas/timesteps, len(states) == len(sigmas) - 1 in every case, last emitted step equals the reported total, output finite. FlowMatchLCMScheduler does accept sigmas=, so it takes the same branch as euler. Heun still refuses partial ranges and still reports its real 15 iterations for 8 steps.

Degenerate window rejected. Confirmed the guard can't fire on a legitimate request: denoising_end=1.0 gives end == num_steps, and denoising_start < denoising_end bounds start ≤ num_steps - 1, so a full-range window is always ≥ 2 sigmas.

PE gate. Now mirrors get_hf_load_class exactly — it does config[submodel_type.value] against model_index.json and raises ValueError on KeyError (generic_diffusers.py), so requiring both the directory and the declaration is the correct predicate. json.loads failure modes are covered: missing file and a directory in its place both raise OSError, malformed JSON raises ValueError, and both are caught.

Init latents / denoising_start. The blend is the right rectified-flow form and matches diffusers' own add_noise. See item 2 below for the one thing it gets wrong.

pe_width / pe_height. Correctly placed after addTextToImage, which is what writes the resolved dimensions onto the denoise node.

---

New — the last remaining blocker

1. The LCM scheduler ignores the seed entirely

denoise.py:109 calls scheduler.step(model_output=pred, timestep=timestep, sample=img) with no generator. FlowMatchLCMScheduler.step is stochastic — it reconstructs x0_pred and re-noises:

noise = randn_tensor(x0_pred.shape, generator=generator, device=..., dtype=...)
prev_sample = (1 - sigma_next) * x0_pred + sigma_next * noise

With generator=None that draws from the global RNG, so every step injects noise the seed field does not control. The carefully seeded CPU noise from the last round only fixes the initial latent.

Trigger: select LCM in the scheduler dropdown (it's one of three offered, flux/schedulers.py:70), generate, then generate again with the same seed in the same server process. I ran exactly that through denoise() — seeding the initial latent identically and leaving the global RNG to advance between runs, as it does in a long-lived process — and got a max elementwise difference of 2.17 between the two results, on a stub model that barely moves the sample. Euler under the same test is bit-identical. So: same seed, different image; and recalling a seed from gallery metadata will not reproduce an LCM generation.

This is pre-existing rather than something these commits introduced — I missed it in both earlier rounds — but it's the one thing I'd want fixed before merge, because it's a silent violation of the seed contract on a user-selectable option.

Fix follows two existing patterns in the repo: denoise_latents.py:759 builds torch.Generator(device).manual_seed(seed ^ 0xFFFFFFFF) into scheduler_step_kwargs for the ancestral SD schedulers, and anima/scheduler_driver.py:143 passes a seeded _step_generator straight into scheduler.step. Threading self.seed down into denoise() and passing generator= is safe for all three schedulers here — euler and heun accept the kwarg too and only consult it when s_churn > 0, which is 0 on this path. (Dropping LCM from ERNIE_IMAGE_SCHEDULER_MAP would also close it, if you'd rather not carry a stochastic sampler in a Prototype node.)

---

New — nice-to-haves

2. The img2img blend uses the pre-shift sigma while the loop starts at the post-shift sigma

ernie_image_denoise.py:131 passes float(timesteps[0]) into _prepare_initial_latents. That's the raw linspace value out of get_schedule, which doesn't apply shift — the shift is applied later, inside the scheduler's set_timesteps (denoise.py:75). So the sample is built at one sigma and the first model call is told another.

Measured with shift=3: denoising_start=0.25 blends at σ=0.750 but the loop opens at σ=0.900; denoising_start=0.5 blends at 0.500 and opens at 0.750. It's exact only at denoising_start=0, where both are 1.0 — which is why the default path is unaffected.

Contrast Z-Image, which applies its own time shift in _get_sigmas before clipping to the denoising window, so its blend and its loop agree by construction. Here the shift lives in the scheduler, so the blend has to ask the scheduler for it. Cheapest fix is to let the scheduler do the noising: after set_timesteps, scheduler.scale_noise(init_latents, scheduler.timesteps[:1], noise) is exactly σ·noise + (1-σ)·sample with the shifted σ. That needs the blend to move after set_timesteps, i.e. into denoise().

Reachable only from the workflow editor — buildErnieImageGraph throws UnsupportedGenerationModeError for anything but txt2img, so denoising_start is always 0 in the linear UI — and only when the checkpoint's scheduler_config.json has shift != 1.0. Hence nice-to-have rather than a blocker.

3. Nothing in the PR can produce the latents that input documents

latents wants "VAE-encoded, BN-normalized, and patchified" latents, but the PR ships ernie_image_vae_decode with no encode counterpart, and there's no ERNIE member in the i2l union. So the only node that can feed it is another ernie_image_denoise.

That matters for the multi-stage handoff your own test docstring cites as the motivation for the denoising_end fix. Chaining stage 1 (0 → 0.5) into stage 2 (0.5 → 1.0) hands stage 2 latents that already sit at σ=0.5, and _prepare_initial_latents will noise them again to 0.5. The blend is correct for its stated contract (init latents are clean, σ=0) — it's just that the handoff is the one reachable caller and it doesn't satisfy that contract. Worth either shipping the encode node or narrowing the docstring to say the input expects fully-denoised latents.

Nit

posCond.pe_width = denoise.width picks up the scaled size, so on canvas with scaling active the enhancer is told the intermediate render size rather than the final output size. Aspect ratio is roughly preserved and that's most of what the enhancer uses it for, so this is cosmetic — getOriginalAndScaledSizesForTextToImage already hands back originalSize if you'd prefer it.

---

CI / verification

On 74ae2a1144: the 43 ERNIE + fp8 tests pass, pytest tests/backend/model_manager tests/app is 2901 passed / 130 skipped / 6 xfailed / 0 failed, ruff check and ruff format --check clean over 1165 files, and on the frontend the 10 new buildErnieImageGraph tests pass with tsc --noEmit, eslint --max-warnings=0 and prettier --check all clean.

Comment rather than a formal state again. (1) is the last remaining blocker — everything else I've raised across all three rounds is now closed. (2), (3) and the nit are yours to take or leave.

[lstein]

Re-reviewed 74ae2a1144..1748457ed9. All four items from the last round are fixed, and the two that mattered are pinned by tests that fail if the fix is reverted. No blockers left — this is good to merge from my side.

---

Verified fixed

1. The LCM seed blocker is closed

generator=torch.Generator(device="cpu").manual_seed(seed ^ 0xFFFFFFFF) is threaded from the invocation through denoise() into every scheduler.step. Re-ran my reproduction from last round against this commit: three LCM runs at the same seed with the global RNG deliberately advanced by torch.rand(1e4) between them are now bit-identical, and a different seed still diverges (max diff 1.72). Previously runs 0 and 1 differed by 2.17.

Details I checked rather than assumed:

• CPU generator, GPU sample. FlowMatchLCMScheduler.step calls randn_tensor(..., device=x0_pred.device, ...). With a CPU generator and a CUDA/ROCm tensor, randn_tensor sets rand_device="cpu", draws there and .to(device)s the result — so the step noise is device-independent in the same way the initial latent already is, which is the right call. It logs an INFO line per call, below diffusers' default WARNING verbosity, so no log spam.
• The XOR. Matches denoise_latents.py, and keeps the step noise decorrelated from the initial noise drawn at self.seed.
• Euler is genuinely unperturbed. test_deterministic_scheduler_is_unaffected_by_the_generator pins it, and it isn't a vacuous test — test_stochastic_scheduler_respects_different_seeds is the guard that stops the seeded test from passing because the generator is ignored. Good pairing.

One thing the fix covers that the comment doesn't claim: in diffusers 0.39 FlowMatchEulerDiscreteScheduler has a stochastic_sampling config flag, and when it's set step() re-noises exactly like LCM. Since _build_scheduler loads the scheduler config from the checkpoint's scheduler/ dir, a checkpoint shipping stochastic_sampling: true would have had this identical bug on the default scheduler. I ran euler both ways — same-seed identical under both. So the fix is broader than advertised.

Minor correction to the code comment at denoise.py:126: flow-match Euler in this version has no s_churn/gamma path at all (that's Heun) — it branches on config.stochastic_sampling. Doesn't change a line of code, just the rationale.

2. The img2img blend now uses the post-shift sigma

Blending moved into denoise() at float(scheduler.sigmas[0]). Verified numerically against scheduler.scale_noise across shift ∈ {1.0, 3.0, 6.0} × denoising_start ∈ {0, 0.25, 0.5, 0.75} — 12/12 exact. The pre-shift/post-shift gap I reported is visible in the matrix (shift=3: 0.750→0.900, 0.500→0.750) and gone.

The regression test is the right shape: all-ones noise against all-zeros init collapses the blend to the sigma itself, so it reads the value directly out of the first model input, and the shift != 1.0 branch pins the specific wrong value (0.5) the old code produced.

3. add_noise — the multi-stage handoff is now correct

This is the better fix for item 3; I'd flagged the docstring as the thing to narrow and this actually closes the hole. Semantics match z_image_denoise exactly (add_noise: bool = InputField(default=True, ...), ignored without init latents), so it's the field people already know.

Checked the handoff arithmetic end to end. get_schedule slices sigmas[start:end+1] with start = int(N·denoising_start), so stage 1 (0 → 0.5) at N=8 ends on sigmas[4] = 0.5 and stage 2 (0.5 → 1.0) opens on sigmas[4] = 0.5 — same raw value, same shift applied by the same scheduler, so the sigma the stage-2 loop announces is exactly where stage 1 left the sample. With add_noise=False I confirmed the first model input is the incoming sample verbatim (torch.equal), i.e. no residual blend.

addImageToImage already lists ernie_image_denoise in its type union and leaves add_noise at its default, which is the correct setting for a clean i2l latent — so whenever the encode node does land, that path is already right.

4. pe_width nit

Now getOriginalAndScaledSizesForTextToImage(state).originalSize. Same helper addTextToImage calls a few lines earlier, so the assert(false, 'Cannot get sizes for tab ...') inside it can't newly fire here. The test fixture was changed to make originalSize (832×1216) and scaledSize (512×768) actually differ, which is what makes the new assertion meaningful — the old fixture had them equal and would have passed either way.

---

Attacks that didn't find anything

• Loop sweep over {euler, euler shift=3, lcm, heun} × steps {1, 2, 8} × windows {(0,1), (0,.5), (.5,1), (.25,.75)} × {init latents, none}, all with a generator: output finite, no progress overrun, no IndexError from the truncated sigmas/timesteps. The only errors raised were the intended ones — the degenerate-window guard at steps=1, and Heun refusing a partial range.
• Lifecycle of the new state. init_latents has exactly one clear site (img, init_latents = init_latents, None) and one read per branch; generator is created once per invoke and never reset mid-loop, which is what reproducibility requires. No other caller of backend/ernie_image/denoise.denoise exists, and both new params are optional.
• Truncated schedule vs LCM's internals. step() reads sigmas[step_index + 1]; after truncation len(sigmas) == len(timesteps) + 1, so the last step's sigma_next is the requested end sigma rather than an out-of-range index. _scale_factors/_upscale_mode are unset, so its length assertion is skipped.
• Generated files. Regenerated both with CI's own commands from a clean tree — schema.ts and openapi.json are byte-identical to what's committed.

CI / verification

On 1748457ed9: 48 ERNIE + fp8 tests pass (up from 43), pytest tests/backend/model_manager tests/app is 2901 passed / 130 skipped / 6 xfailed / 0 failed, ruff check + ruff format --check clean over 1165 files, frontend tsc --noEmit / eslint --max-warnings=0 / prettier --check clean, and the 11 buildErnieImageGraph tests pass. All 17 checks green on the PR.

---

The one thing still open is not a defect: there's no ERNIE VAE encode node, so latents + add_noise=true still has no producer in-tree and the input is only reachable from another ernie_image_denoise. That's a feature gap for a follow-up, not something to hold this on — the field is documented for both cases and the chained case is now correct.

Approving. Nice work on the self-review passes — the denoising_end and add_noise findings were both yours.