| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
feat(fp8): enable FP8 storage for Anima (#9415) * feat(fp8): enable FP8 storage for Z-Image Z-Image was excluded from FP8 storage in #8945 because diffusers' enable_layerwise_casting() was called with the global torch dtype (fp16) while Z-Image loads in bf16: skipped modules stayed bf16, hooked ones produced fp16, and attention crashed. That root cause was fixed later in the same PR — the compute dtype now comes from the model's own parameters — so the exclusion is obsolete. Removing it alone is not enough. Our hook-based cast (#9231) dropped one thing diffusers' enable_layerwise_casting() did: honoring the model's declared _skip_layerwise_casting_patterns. Z-Image needs it, and not for quality — TimestepEmbedder.forward reads self.mlp[0].weight.dtype and casts its *input* to it. With an fp8 weight the input becomes float8 before our pre-hook restores the weight, and F.linear dies with: RuntimeError: "addmm_cuda" not implemented for 'Float8_e4m3fn' which is why ZImageTransformer2DModel declares ['t_embedder', 'cap_embedder']. _apply_fp8_to_nn_module now takes extra_skip_patterns and the caller passes the model's list. For other models this is a strict superset of our defaults (FLUX/SD3 pos_embed+norm, UNet norm, CogView4 also proj_out), so it only ever skips more. Also wire the cast into ZImageCheckpointModel: only the diffusers loader called it, so the toggle was a silent no-op for single-file Z-Image models even though both paths build the same ZImageTransformer2DModel. Tested end to end on CUDA: transformer resident VRAM drops from ~11.5GB to 5880MB for both Z-Image-Turbo (diffusers) and Z-Image-Turbo (checkpoint, 14.37GB file), with clean output images in both cases. * Chore openapi * feat(fp8): enable FP8 storage for Anima The fp8_storage toggle was shown for Anima main models but did nothing: AnimaCheckpointModel never called _apply_fp8_layerwise_casting. Wire it in — the state dict is cast to a single model_dtype before load_state_dict, so the layerwise cast has one unambiguous compute dtype to restore to. Wiring alone renders a heavily dithered image with no fine detail. The cause is t_embedder: it produces the adaln_lora conditioning consumed by every block, so casting it to FP8 corrupts every token everywhere. None of the generic skip patterns match it — they target diffusers' module names (norm, pos_embed, patch_embed, proj_in/out) and this architecture names things differently. AnimaTransformer now declares _skip_layerwise_casting_patterns, the same attribute diffusers models use, so the loader needs no special-casing. Measured on CUDA, same seed/steps/CFG each run: casting nothing = broken at 1994MB; t_embedder alone = clean at 2010MB; adding x_embedder and final_layer changes nothing further (2012MB) and is kept as margin on the I/O layers; adaln_modulation was tested too and is deliberately not listed — it costs 168MB and made no difference. Against a bf16 reference (3988MB) the FP8 result keeps the same composition and loses only a little micro-detail. * test(fp8): drop the Z-Image entry from the exclusion parametrize main added a device-probe parametrize listing Z-Image as an excluded model. This branch removes that exclusion, so the entry contradicts `test_should_use_fp8_allows_z_image` and the case now returns the probe's value instead of False. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(fp8): pin Anima's skip patterns to real modules and guard the wiring Review follow-ups for #9415. Add `tests/.../test_anima_fp8_wiring.py`. Deleting the `_apply_fp8_layerwise_casting` call from the Anima single-file loader previously left the whole model_manager and anima suites green, so the dead `fp8_storage` toggle this PR fixes could come straight back with CI passing. The new boundary test fails on that mutation. The pattern test now instantiates the real `AnimaTransformer` under `accelerate.init_empty_weights()` and pins all three declared patterns to actual dotted module paths, instead of asserting a string is in a list against a hand-built stand-in. A second test records that `_FP8_DEFAULT_SKIP_PATTERNS` covers zero modules in this architecture, so the declared list is demonstrably not redundant. Lift the transformer kwargs to `ANIMA_TRANSFORMER_CONFIG` so tests build the real graph without duplicating them, mirroring `KREA2_TRANSFORMER_CONFIG`. Correct the skip-list comment. `adaln_modulation` "made no difference" was not supported by measurement: relative L2 against bf16 on a single forward goes 0.134 -> 0.091 when it is skipped, making it the largest remaining error source. The 168MB call still stands, but it rests on a 35-step A/B showing no visible difference, and the comment now says so. Also note that most of what the `final_layer` entry shields is `final_layer.adaln_modulation.*` (1.57 of 1.70M params). Stop offering FP8 storage for Anima LLLite ControlNets in the model manager. `AnimaControlNetLLLiteModel` never calls the layerwise cast, so the toggle was rendered and inert; at 16-63MB per adapter, hiding it beats wiring it. * fix(fp8): stop the hidden Anima ControlNet fp8 toggle from re-persisting Two fixes from an adversarial review of the merge: - `ControlAdapterModelDefaultSettings` hid the FP8 storage control for Anima LLLite adapters but kept sending its value. react-hook-form keeps unrendered fields in `defaultValues` (`shouldUnregister` defaults to false), so a value persisted before the control was hidden was re-sent verbatim on every save, with no UI left to clear it. Null it out wherever the control is hidden. - `test_single_file_loader_applies_fp8_layerwise_casting` passed `fp8_storage` as a top-level kwarg to `model_construct`. It is not a field of `Main_Checkpoint_Anima_Config` and the model has no `extra="allow"`, so pydantic silently discarded it and `default_settings` stayed `None` -- the toggle was off in the test that exists to prove the toggle is wired up. Build a real `MainModelDefaultSettings(fp8_storage=True)` instead. --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> | 13 天前 | |
feat: ernie image/turbo (#9115) * Update to Transformers 5.1.0 * remove extra stuff * feat: ERNIE-Image integration (incl. diffusers 0.38 + transformers 5) Adds Baidu ERNIE-Image and ERNIE-Image-Turbo as a new BaseModelType, mirroring the FLUX.2 / Z-Image integration pattern. Both models share the ErnieImageTransformer2DModel architecture (3072 hidden, 24 layers, 24 heads) and an AutoencoderKLFlux2 VAE; they differ only in default inference settings (50 steps + CFG 4.0 vs 8 steps + CFG 1.0 for Turbo). Built on top of PR #8859 (transformers 5.1+) and additionally bumps diffusers 0.36.0 -> 0.38.0, which is the first release containing the ErnieImagePipeline and ErnieImageTransformer2DModel. Backend - BaseModelType.ErnieImage, ModelType.PromptEnhancer, two new SubModelTypes (pe, pe_tokenizer) for the bundled prompt enhancer - Main_Diffusers_ErnieImage_Config + Main_Checkpoint_ErnieImage_Config with state-dict-based detection (x_embedder + text_proj + adaLN_modulation) - Diffusers loader registered for ERNIE-Image; uses upstream subdir conventions, loads transformer / vae / text_encoder / tokenizer plus optional pe / pe_tokenizer - New invokeai/backend/ernie_image/ with sampling utilities (2x2 patchify, BN normalize/denormalize, sigma schedule, padded text packing) and a rectified-flow denoise loop supporting Euler/Heun/LCM - Five invocations: model_loader, text_encoder (with prompt-enhancer toggle), denoise, vae_encode, vae_decode - ErnieImageConditioningInfo + ConditioningField/Output + pickle allowlist - ERNIE_IMAGE_SCHEDULER_MAP reusing the FlowMatch* scheduler classes - New generation modes ernie_image_{txt2img,img2img,inpaint,outpaint} - Starter models for baidu/ERNIE-Image and baidu/ERNIE-Image-Turbo + STARTER_BUNDLES entry Frontend - Regenerated services/api/schema.ts to expose the new node types - Type unions extended (ImageOutput / LatentToImage / ImageToLatents / DenoiseLatents / MainModelLoaderNodes) - ParamsState gains ernieImageScheduler + ernieImageUsePromptEnhancer, with reducers, selectors, and selectIsErnieImage - buildErnieImageGraph (txt2img/img2img/inpaint/outpaint) wired into useEnqueueCanvas - ERNIE entries added to MODEL_BASE_TO_{COLOR,LONG_NAME,SHORT_NAME} and prompt_enhancer to MODEL_TYPE_TO_LONG_NAME - ParamErnieImageScheduler and ParamErnieImagePromptEnhancer rendered conditionally in GenerationSettingsAccordion - All add{TextTo,ImageTo,Inpaint,Outpaint}Image type guards extended to accept ernie_image_denoise; isMainModelWithoutUnet ditto Diffusers 0.38 fix-out - hotfixes.py: import LoRACompatibleConv directly; the lazy-module __getattr__ no longer exposes diffusers.models.lora as an attribute Verification - pytest tests/ -m "not slow": 619 passed, 0 failed - pnpm lint:tsc / lint:eslint / lint:prettier: clean - pnpm test:no-watch: 563 passed, 0 failed - Manual smoketest pending: requires baidu/ERNIE-Image weights and a GPU (8B parameters; CPU not practical) Out of scope (follow-up phases) - ControlNet, IP-Adapter, LoRA support for ERNIE-Image - Single-file checkpoint loading (defensive scaffolding only) - Metadata recall handlers in the gallery side panel * Chore Ruff + Typegen * Chore ruff * fix(ernie-image): timestep scale, live preview, missing graph case, and UI cleanup - Pass timesteps in [0, num_train_timesteps] to the transformer instead of [0, 1]; the diffusers Timesteps embedding expects the unnormalised range, which produced mosaic-pattern garbage instead of an image. - Unpatchify predicted-x0 before the denoise step callback and route through sd_step_callback so the canvas shows a live preview during sampling (uses FLUX.2's RGB factors -- same AutoencoderKLFlux2 / 32 latent channels). - Add the ernie-image case to useEnqueueGenerate (Generate tab); was only wired up in useEnqueueCanvas, so plain text-to-image failed with "No graph builders for base ernie-image". - Move the Prompt Enhancer toggle from the Generation accordion into a dedicated ERNIE-Image block in the Advanced accordion; hide the rest of the SD-style advanced controls (CLIP skip, CFG rescale, seamless, color comp., separate VAE) since none apply to ERNIE-Image. - Detect ERNIE-Image-Turbo by name in MainModelDefaultSettings.from_base so installs (starter or manual) get steps=8, cfg_scale=1.0 instead of the standard 50/4.0. - Pin compel to Cstannahill/compel5@chore/transformers5-diffusers-smoke for transformers>=5 compatibility (PR damian0815/compel#129). * Merged missed * Chore Ruff * Chore Fix UV lock * fix(ernie): restrict ERNIE-Image to text-to-image only ERNIE's denoise node has no denoise_mask input, so masked modes are unsupported. Drop img2img/inpaint/outpaint from the ERNIE graph builder, exclude ernie_image_denoise from a new MaskableDenoiseNodes type used by addInpaint/addOutpaint, and align addImageToImage unions with the node type aliases. Fixes tsc failures on the ernie-image branch. * fix(ernie): restrict ERNIE-Image to text-to-image only ERNIE-Image's denoise node has no denoise_mask input and no mask logic, so masked modes (inpaint/outpaint) are impossible and image-to-image is dropped as well. Frontend: - buildErnieImageGraph now builds txt2img only; asserts on other modes - add MaskableDenoiseNodes (DenoiseLatentsNodes minus ernie_image_denoise) and use it in addInpaint/addOutpaint - drop ernie_image_vae_encode from ImageToLatentsNodes; align addImageToImage unions with the LatentToImage/ImageToLatents aliases - regenerate schema.ts Backend: - delete the ernie_image_vae_encode invocation (i2l, only used by the removed image-input modes) - drop ernie_image_{img2img,inpaint,outpaint} from GENERATION_MODES Fixes the tsc failures on the ernie-image branch. * Chore OpenApi * fix(model_manager): remove duplicate _has_anima_keys shadowing ComfyUI-bundled prefix support A duplicate _has_anima_keys definition (older, net.-only) shadowed the complete version that also recognizes the `model.diffusion_model.` ComfyUI-bundled prefix, causing Anima identification to reject bundled checkpoints. * Chore Openapi * fix(ernie-image): exclude prompt enhancer from fp8, tighten turbo detection Address review findings on #9115: - Exclude SubModelType.PromptEnhancer and PromptEnhancerTokenizer from fp8 layerwise casting. The prompt enhancer is a causal LM driven by generate() — one full forward per generated token — so casting made the whole LM round-trip bf16<->fp8 per token, on top of fp8 rounding a model whose entire job is text quality. - Match turbo detection on the install directory's leaf name instead of the whole path string. An in-place install records an absolute path, so an ancestor directory like /mnt/turbo-nvme/ silently gave the base model Turbo's 8 steps and CFG 1.0. - Pass local_files_only=True on all ERNIE-Image from_pretrained calls and load tokenizers bare, matching the sibling loaders. - Cap the prompt enhancer's max_new_tokens at 1024. Driving it off model_max_length hangs the graph if the tokenizer config omits it and transformers substitutes its int(1e30) sentinel. Adds regression tests for the fp8 exclusion and the turbo path matching. * fix(ernie-image): honor denoising window, noise init latents, harden PE gate Self-review follow-ups on #9115: - Honor denoising_end. Every FlowMatch scheduler appends its own terminal 0 sigma, and passing the window minus its last entry let that zero stand in for the requested end sigma - so denoising_end < 1.0 ran a full denoise in fewer, coarser steps. Hand the scheduler the whole window and truncate its appended zero instead, which also keeps the scheduler's own `shift` applied to the terminal sigma. - Reject a denoising window that rounds down to a single sigma. It yielded zero steps, so the loop returned its input untouched and the graph decoded raw noise with no error at all. - Blend image-to-image init latents with noise at the first sigma, and reject denoising_start > 0 when no latents are provided. Both cases previously lied to the model about where the sample sits on the rectified-flow path. The shape check now covers batch and spatial dims, and its message no longer points at a VAE-encode node that does not exist. - Require both `pe` and `pe_tokenizer` to be present AND declared in model_index.json before offering the prompt enhancer. get_hf_load_class resolves submodels from that file, so a directory-only check let a partial install pass the gate and then hard-fail the generation - by default, since the toggle is on. - Pass the real generation dimensions to the prompt enhancer instead of leaving it on its 1024x1024 defaults. Adds regression tests for each, including a graph-builder test suite for ERNIE (which had none). Added ERNIE-Image + Krea 2 Raw into README.md under Supported Model. * Readme.me update with ERNIE-Image-Turbo * fix(ernie-image): seed the stochastic scheduler, blend at the shifted sigma Round-4 review follow-ups on #9115: - Pass a seeded generator into scheduler.step. FlowMatchLCMScheduler is stochastic - it re-noises the sample every step - so with generator=None it drew from the global RNG. The seed field only controlled the initial latent, so the same seed produced a different image on every run and seed recall from gallery metadata could not reproduce an LCM generation. Seeded from `seed ^ 0xFFFFFFFF` (as denoise_latents.py does) so the step noise stays decorrelated from the initial noise, and from a CPU generator so it is device-independent like the initial noise. - Blend image-to-image init latents at the scheduler's post-shift first sigma instead of the raw schedule value. get_schedule emits raw linspace values and the scheduler applies `shift` in set_timesteps, so the blend built the sample at one sigma and then told the first model call it was at another. The blend moves into denoise(), which is the only layer that knows the shifted sigma. - Add an `add_noise` field mirroring z_image_denoise. The only node that can currently feed `latents` is another ernie_image_denoise, whose output already sits at the handoff sigma - re-noising it broke the very multi-stage handoff the denoising_end fix enables. - Give the prompt enhancer the original size rather than the intermediate scaled render size. Regenerates openapi.json / schema.ts for the new field. Adds regression tests for the seed contract (same seed identical, different seeds different, euler unaffected) and for the post-shift blend. --------- Co-authored-by: Your Name <you@example.com> Co-authored-by: 4pointoh <97913726+4pointoh@users.noreply.github.com> Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> | 1 个月前 | |
feat: ernie image/turbo (#9115) * Update to Transformers 5.1.0 * remove extra stuff * feat: ERNIE-Image integration (incl. diffusers 0.38 + transformers 5) Adds Baidu ERNIE-Image and ERNIE-Image-Turbo as a new BaseModelType, mirroring the FLUX.2 / Z-Image integration pattern. Both models share the ErnieImageTransformer2DModel architecture (3072 hidden, 24 layers, 24 heads) and an AutoencoderKLFlux2 VAE; they differ only in default inference settings (50 steps + CFG 4.0 vs 8 steps + CFG 1.0 for Turbo). Built on top of PR #8859 (transformers 5.1+) and additionally bumps diffusers 0.36.0 -> 0.38.0, which is the first release containing the ErnieImagePipeline and ErnieImageTransformer2DModel. Backend - BaseModelType.ErnieImage, ModelType.PromptEnhancer, two new SubModelTypes (pe, pe_tokenizer) for the bundled prompt enhancer - Main_Diffusers_ErnieImage_Config + Main_Checkpoint_ErnieImage_Config with state-dict-based detection (x_embedder + text_proj + adaLN_modulation) - Diffusers loader registered for ERNIE-Image; uses upstream subdir conventions, loads transformer / vae / text_encoder / tokenizer plus optional pe / pe_tokenizer - New invokeai/backend/ernie_image/ with sampling utilities (2x2 patchify, BN normalize/denormalize, sigma schedule, padded text packing) and a rectified-flow denoise loop supporting Euler/Heun/LCM - Five invocations: model_loader, text_encoder (with prompt-enhancer toggle), denoise, vae_encode, vae_decode - ErnieImageConditioningInfo + ConditioningField/Output + pickle allowlist - ERNIE_IMAGE_SCHEDULER_MAP reusing the FlowMatch* scheduler classes - New generation modes ernie_image_{txt2img,img2img,inpaint,outpaint} - Starter models for baidu/ERNIE-Image and baidu/ERNIE-Image-Turbo + STARTER_BUNDLES entry Frontend - Regenerated services/api/schema.ts to expose the new node types - Type unions extended (ImageOutput / LatentToImage / ImageToLatents / DenoiseLatents / MainModelLoaderNodes) - ParamsState gains ernieImageScheduler + ernieImageUsePromptEnhancer, with reducers, selectors, and selectIsErnieImage - buildErnieImageGraph (txt2img/img2img/inpaint/outpaint) wired into useEnqueueCanvas - ERNIE entries added to MODEL_BASE_TO_{COLOR,LONG_NAME,SHORT_NAME} and prompt_enhancer to MODEL_TYPE_TO_LONG_NAME - ParamErnieImageScheduler and ParamErnieImagePromptEnhancer rendered conditionally in GenerationSettingsAccordion - All add{TextTo,ImageTo,Inpaint,Outpaint}Image type guards extended to accept ernie_image_denoise; isMainModelWithoutUnet ditto Diffusers 0.38 fix-out - hotfixes.py: import LoRACompatibleConv directly; the lazy-module __getattr__ no longer exposes diffusers.models.lora as an attribute Verification - pytest tests/ -m "not slow": 619 passed, 0 failed - pnpm lint:tsc / lint:eslint / lint:prettier: clean - pnpm test:no-watch: 563 passed, 0 failed - Manual smoketest pending: requires baidu/ERNIE-Image weights and a GPU (8B parameters; CPU not practical) Out of scope (follow-up phases) - ControlNet, IP-Adapter, LoRA support for ERNIE-Image - Single-file checkpoint loading (defensive scaffolding only) - Metadata recall handlers in the gallery side panel * Chore Ruff + Typegen * Chore ruff * fix(ernie-image): timestep scale, live preview, missing graph case, and UI cleanup - Pass timesteps in [0, num_train_timesteps] to the transformer instead of [0, 1]; the diffusers Timesteps embedding expects the unnormalised range, which produced mosaic-pattern garbage instead of an image. - Unpatchify predicted-x0 before the denoise step callback and route through sd_step_callback so the canvas shows a live preview during sampling (uses FLUX.2's RGB factors -- same AutoencoderKLFlux2 / 32 latent channels). - Add the ernie-image case to useEnqueueGenerate (Generate tab); was only wired up in useEnqueueCanvas, so plain text-to-image failed with "No graph builders for base ernie-image". - Move the Prompt Enhancer toggle from the Generation accordion into a dedicated ERNIE-Image block in the Advanced accordion; hide the rest of the SD-style advanced controls (CLIP skip, CFG rescale, seamless, color comp., separate VAE) since none apply to ERNIE-Image. - Detect ERNIE-Image-Turbo by name in MainModelDefaultSettings.from_base so installs (starter or manual) get steps=8, cfg_scale=1.0 instead of the standard 50/4.0. - Pin compel to Cstannahill/compel5@chore/transformers5-diffusers-smoke for transformers>=5 compatibility (PR damian0815/compel#129). * Merged missed * Chore Ruff * Chore Fix UV lock * fix(ernie): restrict ERNIE-Image to text-to-image only ERNIE's denoise node has no denoise_mask input, so masked modes are unsupported. Drop img2img/inpaint/outpaint from the ERNIE graph builder, exclude ernie_image_denoise from a new MaskableDenoiseNodes type used by addInpaint/addOutpaint, and align addImageToImage unions with the node type aliases. Fixes tsc failures on the ernie-image branch. * fix(ernie): restrict ERNIE-Image to text-to-image only ERNIE-Image's denoise node has no denoise_mask input and no mask logic, so masked modes (inpaint/outpaint) are impossible and image-to-image is dropped as well. Frontend: - buildErnieImageGraph now builds txt2img only; asserts on other modes - add MaskableDenoiseNodes (DenoiseLatentsNodes minus ernie_image_denoise) and use it in addInpaint/addOutpaint - drop ernie_image_vae_encode from ImageToLatentsNodes; align addImageToImage unions with the LatentToImage/ImageToLatents aliases - regenerate schema.ts Backend: - delete the ernie_image_vae_encode invocation (i2l, only used by the removed image-input modes) - drop ernie_image_{img2img,inpaint,outpaint} from GENERATION_MODES Fixes the tsc failures on the ernie-image branch. * Chore OpenApi * fix(model_manager): remove duplicate _has_anima_keys shadowing ComfyUI-bundled prefix support A duplicate _has_anima_keys definition (older, net.-only) shadowed the complete version that also recognizes the `model.diffusion_model.` ComfyUI-bundled prefix, causing Anima identification to reject bundled checkpoints. * Chore Openapi * fix(ernie-image): exclude prompt enhancer from fp8, tighten turbo detection Address review findings on #9115: - Exclude SubModelType.PromptEnhancer and PromptEnhancerTokenizer from fp8 layerwise casting. The prompt enhancer is a causal LM driven by generate() — one full forward per generated token — so casting made the whole LM round-trip bf16<->fp8 per token, on top of fp8 rounding a model whose entire job is text quality. - Match turbo detection on the install directory's leaf name instead of the whole path string. An in-place install records an absolute path, so an ancestor directory like /mnt/turbo-nvme/ silently gave the base model Turbo's 8 steps and CFG 1.0. - Pass local_files_only=True on all ERNIE-Image from_pretrained calls and load tokenizers bare, matching the sibling loaders. - Cap the prompt enhancer's max_new_tokens at 1024. Driving it off model_max_length hangs the graph if the tokenizer config omits it and transformers substitutes its int(1e30) sentinel. Adds regression tests for the fp8 exclusion and the turbo path matching. * fix(ernie-image): honor denoising window, noise init latents, harden PE gate Self-review follow-ups on #9115: - Honor denoising_end. Every FlowMatch scheduler appends its own terminal 0 sigma, and passing the window minus its last entry let that zero stand in for the requested end sigma - so denoising_end < 1.0 ran a full denoise in fewer, coarser steps. Hand the scheduler the whole window and truncate its appended zero instead, which also keeps the scheduler's own `shift` applied to the terminal sigma. - Reject a denoising window that rounds down to a single sigma. It yielded zero steps, so the loop returned its input untouched and the graph decoded raw noise with no error at all. - Blend image-to-image init latents with noise at the first sigma, and reject denoising_start > 0 when no latents are provided. Both cases previously lied to the model about where the sample sits on the rectified-flow path. The shape check now covers batch and spatial dims, and its message no longer points at a VAE-encode node that does not exist. - Require both `pe` and `pe_tokenizer` to be present AND declared in model_index.json before offering the prompt enhancer. get_hf_load_class resolves submodels from that file, so a directory-only check let a partial install pass the gate and then hard-fail the generation - by default, since the toggle is on. - Pass the real generation dimensions to the prompt enhancer instead of leaving it on its 1024x1024 defaults. Adds regression tests for each, including a graph-builder test suite for ERNIE (which had none). Added ERNIE-Image + Krea 2 Raw into README.md under Supported Model. * Readme.me update with ERNIE-Image-Turbo * fix(ernie-image): seed the stochastic scheduler, blend at the shifted sigma Round-4 review follow-ups on #9115: - Pass a seeded generator into scheduler.step. FlowMatchLCMScheduler is stochastic - it re-noises the sample every step - so with generator=None it drew from the global RNG. The seed field only controlled the initial latent, so the same seed produced a different image on every run and seed recall from gallery metadata could not reproduce an LCM generation. Seeded from `seed ^ 0xFFFFFFFF` (as denoise_latents.py does) so the step noise stays decorrelated from the initial noise, and from a CPU generator so it is device-independent like the initial noise. - Blend image-to-image init latents at the scheduler's post-shift first sigma instead of the raw schedule value. get_schedule emits raw linspace values and the scheduler applies `shift` in set_timesteps, so the blend built the sample at one sigma and then told the first model call it was at another. The blend moves into denoise(), which is the only layer that knows the shifted sigma. - Add an `add_noise` field mirroring z_image_denoise. The only node that can currently feed `latents` is another ernie_image_denoise, whose output already sits at the handoff sigma - re-noising it broke the very multi-stage handoff the denoising_end fix enables. - Give the prompt enhancer the original size rather than the intermediate scaled render size. Regenerates openapi.json / schema.ts for the new field. Adds regression tests for the seed contract (same seed identical, different seeds different, euler unaffected) and for the post-shift blend. --------- Co-authored-by: Your Name <you@example.com> Co-authored-by: 4pointoh <97913726+4pointoh@users.noreply.github.com> Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> | 1 个月前 | |
fix(flux2): estimate working memory for denoise and both VAE directions (#9519) * fix(flux2): estimate working memory for denoise and both VAE directions The FLUX.2 path called model_on_device() with no working_mem_bytes anywhere, so the model cache reserved only the default device_working_mem_gb and filled the rest of the card with the model. Reference images make that fatal rather than merely tight: their latents are concatenated onto the image stream, so three 1024x1024 references quadruple the attended sequence of a 1024x1024 generation -- 6.5GB of activations against a 3GB reservation. Measured on CUDA in bf16 as peak reserved memory: transformer activations scale linearly at ~0.39 MB/token (no O(seq^2) term, SDPA) and are independent of block count; the FLUX.2 VAE costs ~2170 (decode) / ~1070 (encode) bytes per pixel per element byte, so a 1024x1024 decode peaks at ~4.3GB. Add Flux2DenoiseInvocation._estimate_working_memory() and estimate_vae_working_memory_flux2(), and pass them at every load site so the cache evicts enough to make room instead of hitting the shortfall as an OOM. Closes #9500 * fix(flux2): budget SDPA's materialized score matrix where it is real The FLUX.2 working-memory estimates were linear in the sequence length, which holds only while SDPA picks a fused kernel. That is a property of the torch build, not of FLUX.2: ROCm's fused kernels cap the head dim at 128 and reject arbitrary additive masks, so both the VAE's 512-wide mid-block head and the dense S x S bias regional prompting attaches fall through to the math fallback and materialize the score matrix -- ~17GB for a 1536px decode, and heads x S^2 for a masked forward. Rather than assume either way, ask torch: sdpa_score_matrix_bytes() queries can_use_flash/efficient/cudnn_attention for the real head dim, dtype and mask, and adds 13 bytes per score element only when no fused kernel is eligible. Measured on CUDA with SDPBackend.MATH forced: 12.9 bytes/element at 4k tokens, 10.3 at 8k, 9.7 at 16k, identical for bf16, fp16 and fp32 because the fallback's softmax intermediates are always fp32. On CUDA every shape reports fused, so the term is zero and the existing calibration is untouched. Non-CUDA devices keep the fused assumption -- torch exposes no equivalent query there, and guessing would reserve double-digit GB on no evidence. * fix(flux2): ask the real dispatcher which SDPA path a build takes The score-matrix term probed torch's CUDA eligibility helpers and read everything else as fused. That was wrong twice over: MPS has no fused SDPA kernel at all and runs the MPSGraph math transcription, so a 1024px VAE decode was admitted ~3.5GB short; and a failed probe returned "fused" too, turning "we don't know" into the one answer that can OOM. Ask `_fused_sdp_choice` instead -- the same dispatch query `scaled_dot_product_attention` runs to pick its kernel. Torch registers it for CPU, CUDA/ROCm and XPU only, so the call raises on exactly the devices that fall through to `math`, and every other failure lands on the conservative side by the same branch. Diffusers models do not reach torch's SDPA directly, so also consult `dispatch_attention_fn`'s active backend: a user on `_native_math` materializes the score matrix on hardware whose probe reports fused. Only the transformer needs this -- the FLUX.2 VAE's mid-block attention still calls SDPA itself through `AttnProcessor2_0` -- and a test pins that asymmetry. On CUDA with the stock backend every one of these terms remains zero. * fix(flux2): read the attention backend live instead of caching it once `_diffusers_attention_dispatch()` was `lru_cache`d, so the first estimate in a process pinned the answer forever. A switch to `_native_math` after that kept reserving zero for the S x S score matrix -- the exact case the lookup was added to catch. Read it live; it is a dict lookup against an already-imported module, priced once per invocation. The torch probe had the same defect one level down: its answer depends on the global SDPA kernel toggles, which `sdpa_kernel()` and `enable_flash_sdp()` flip at runtime. That probe allocates and dispatches, so it stays cached -- but keyed on the toggles, so a switch invalidates it. Per-model overrides need no plumbing: `set_attention_backend()` stamps its choice onto the process-wide registry as well as onto the model's processors, deliberately, so the estimate sees it without holding the model it is priced ahead of. A test pins that propagation. * fix(flux2): stop caching the SDPA probe and scale the VAE estimate by batch The probe's cache key held the four per-backend enable flags, but torch takes the *first eligible* backend in a priority order that `sdpa_kernel(..., set_priority=True)` reorders while leaving every flag untouched -- measured: same flags, EFFICIENT outside and MATH inside. A fused answer cached before the switch would suppress the score-matrix reservation after it. Rather than adding the priority order to the key -- the next thing to forget is always one more -- drop the cache. The probe costs ~6us against a multi-second forward, so there is nothing to protect. `vae.decode` is also handed whatever batch the latents carry, and a LatentsField is not pinned to one, so an estimate built from H and W alone gave a two-sample decode a single sample's reservation. Measured at 1024px: 4.23GB at batch 1, 7.96GB at 2, 11.89GB at 3 -- linear, slightly sub-linear per sample, so the scaled single-sample estimate stays an upper bound. The score matrix is (batch, heads, S, S) and scales with it. * fix(flux2): scale the denoise reservation by the latent batch The node had `b` in hand from preparing the latents and never passed it, so a two-sample run reserved one sample's activations and the cache admitted it to a card that could not run it. Batched latents do not come from the stock UI, but the API and custom graphs reach this node. Batch multiplies the token count and nothing else. Measured on the Klein geometry with a reduced block count: 4608 tokens at B=1 peaks at 2570MB, the same 4608 at B=2 at 5126MB, and 9728 tokens at B=1 at 5584MB -- per total token that is 0.554-0.578MB across every combination, so batch and sequence are interchangeable. Reference latents are repeated per sample by `ensure_batch_size`, so they scale too, and the score matrix is (batch, heads, S, S). The fixed base does not scale -- it covers weight casts and allocator slack -- and neither does the regional bias, built as (1, 1, S, S) and broadcast. * fix(flux2): take the reservation's batch from the blended latents `b` was read from the noise tensor, which this node builds at batch 1 from width/height/seed whenever `add_noise` is set. Batched init latents then broadcast against it in the img2img preblend, producing a two-sample `x` against a one-sample reservation. Read `x.shape[0]` instead. It is already in scope at the estimate -- the blend, the pack and the BN normalize all run above it, and none of them change the batch -- and it is the only thing that knows how many samples reach the transformer. Expanding the noise to match would have worked too, but that changes the noise, and with it the output. Note that the same `b` still feeds `generate_img_ids_flux2`, which is a correctness question rather than a memory one and is left alone here. * fix(flux2): scale the estimate by transformer width, not just token count Per-token activation cost is linear in the transformer's hidden width, and the constant was calibrated on Klein 9B (4096) but applied to every variant. FLUX.2 [dev] is 6144 and reaches this node as a first-class path, so 1024x1024 with three references reserved 7.6GB against ~10GB needed -- the same shortfall #9500 describes, on the model where partial loading makes the estimate decide residency. Measured slope between 4608 and 9216 tokens, everything else held fixed: 0.291 MB/tok at 3072, 0.386 at 4096, 0.555 at 6144 -- 0.755 / 1.00 / 1.438 against width ratios of 0.75 / 1.00 / 1.50. Scale by width, and take the head count from the same number instead of always charging the widest. Also: raise SDPA_MATH_BYTES_PER_SCORE_ELEMENT to 14, which ROCm's 13.62 at the smallest measured shape needs; drop the claim that ROCm rejects additive masks, which gfx1100 disproves; log at info when the score-matrix term fires, since it decides residency and nothing else said so; guard the warning-filter swap with a lock now that the probe runs on every estimate; and give the VAE encode node the same compute device as the decode node. * fix(flux2): raise both calibrated constants to bound the AMD measurements Two ROCm runs came in through the new calibration script and both shipped constants were under their worst point. SDPA_MATH_BYTES_PER_SCORE_ELEMENT goes 14 -> 17 (gfx1201 costs 16.38 for the shape where CUDA costs 12.88), and the per-token activation constant 0.40 -> 0.42 MB (gfx1201 measures 0.4067 at the reference width). Both are now pinned against all three platforms' measured points rather than only being self-consistent. The runs also disprove the ROCm framing this feature carried: gfx1100 reports MATH for the VAE's 512-wide head, gfx1201 reports FLASH. Two cards, same vendor, same torch, opposite answers -- which is the case for asking torch rather than hard-coding a rule, but it means the docstrings could not keep saying "ROCm caps the head dim at 128". The VAE linear constants are left alone despite measuring short on gfx1201. That run had MIOPEN_FIND_MODE=2 and a HIP allocator garbage_collection threshold set, and its series is non-monotonic above 1024px -- peak reserved falls as resolution rises, which is what a GC threshold does to this measurement. The two AMD cards are also 1.8x apart. The gap is documented where the constants are defined; the script now reports the environment and flags a non-monotonic series so the next run cannot be ambiguous about it. * fix(flux2): fit the VAE constants per convolution backend A clean ROCm run — the earlier one had MIOPEN_FIND_MODE=2, worth a uniform 1.28x, and a HIP allocator GC threshold that clipped the high-resolution points — puts the gfx1201 numbers at 3453/2688 bytes per pixel per element byte against cuDNN's 2185/1072. Flat across 512-1024px on both, so the linear model holds; only the coefficient moves. It is MIOpen's convolution workspaces, not the attention term: identical on the fused path. Shipping the MIOpen numbers everywhere would add ~60% to every cuDNN decode for nothing, so the constant follows the backend, keyed on torch.version.hip rather than the device string (a HIP build reports device.type == "cuda"). The two operations also stop sharing a ratio. "Encoding costs half of decoding" holds on cuDNN (0.49) and not on MIOpen (0.78), so it was a backend property masquerading as an architectural one. MIOPEN_FIND_MODE=2 is deliberately not budgeted for: it is not the default and would tax everyone else. Noted where the constants are defined. * fix(flux2): take the larger VAE term, not the sum, and refit MIOpen The W7900 run shows forced math and the fused path measuring identically to three decimals at every resolution, with the total flat-linear in area. The score matrix does not add to the convolution peak: the mid-block sits alone at the 8x-downsampled bottleneck, so the full-resolution feature maps are not live while it runs, and peak reserved is a high-water mark rather than a running total. Measured on cuDNN, forcing math stays *below* the fused path until 1536px and then exceeds it by 2.6GB against the 21.5GB the term prices standalone. Summing reserved 11.1GB for a 1024px gfx1100 decode that measures 6.7. Take the max: 7.0GB. A max model is weakest at the crossover, and one measured point sits there -- a 768px encode with cuDNN's constant and a materializing kernel wants 1.80GB against a 1.35GB max, reproducibly. It is not reachable as a shortfall: the cache floors every reservation at device_working_mem_gb and the whole crossover region is below it. Pinned rather than rounded away. The MIOpen decode constant also goes 3500 -> 3600; the W7900 asks for 3525 at 512px where gfx1201 asks for 3453. Its encode column agrees with gfx1201's to the byte, so this is MIOpen rather than a per-card quirk. --------- Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com> | 3 天前 | |
fix(fp8): resolve compute dtype instead of reading model.dtype (#9412) * fix(fp8): resolve compute dtype instead of reading model.dtype SDXL with fp8_storage crashed before the UNet was ever called: NotImplementedError: "pow_cuda" not implemented for 'Float8_e4m3fn' After layerwise casting the UNet's weights are float8_e4m3fn, and diffusers derives `model.dtype` from the first parameter — so `unet.dtype` reports a storage-only dtype. The legacy SD/SDXL denoise path used it for every tensor it built, so the latents were created in float8 and the first bit of scheduler math (`sigma ** 2` in `add_noise`) blew up. torch has no arithmetic kernels for float8; it is only valid for weights that the forward hooks cast up per layer. Add `get_model_compute_dtype()`: returns `model.dtype` for normal models and the compute dtype for fp8 ones. The loader records the compute dtype on the model when it applies the cast; if the marker is missing (older cache entry, Krea2 encoder path) the resolver scans for the first non-fp8 float param, which works because the cast skips norm layers. Converted every site that derived a tensor dtype from a possibly-fp8 model: latents, noise, mask, masked_latents, conditioning, IP-Adapter and LoRA patch weights in denoise_latents and tiled_multi_diffusion_denoise_latents, plus the LoRA and T2I-Adapter extensions on the modular path. ControlNet and T2I-Adapter control images had the same latent bug — those configs expose an fp8_storage toggle too, so their control image would have been built in float8. Also point LayerPatcher at the shared FP8_STORAGE_DTYPES constant. Regression test covers the real loader path: `model.dtype` is float8 while the resolver returns fp16, and the resolved dtype survives the scheduler arithmetic that crashed. * fix(fp8): harden compute-dtype marker against double-cast poisoning Review follow-up on the compute-dtype resolver. `_apply_fp8_layerwise_casting` derives the compute dtype from the first parameter and is not idempotent. Called on an already-cast model, the first param is float8, so it would record float8 as the *compute* dtype — and `get_model_compute_dtype` trusts the marker, silently reintroducing the "pow_cuda" not implemented for 'Float8_e4m3fn' crash. Two guards close the class: `set_fp8_compute_dtype` rejects any storage-only dtype, and the cast early-returns when the marker is already present. Move the marker-setting into `_apply_fp8_to_nn_module` itself. It was duplicated at both call sites (load_default and krea2's text encoder), so a third caller would have to remember it — the exact failure the fallback scan exists to paper over. Log a warning when the last-resort fallback fires (fp8 storage, no marker, no non-fp8 float param): it returns the global torch dtype, which is wrong for a bf16-compute model and would otherwise surface as an unexplained mismatch deep in the forward pass. Note the same bug in the vendored HiDiffusion pipeline, which builds control images from `controlnet.dtype` at four sites. Dead code today — only apply_hidiffusion/remove_hidiffusion are imported — but it would reproduce the crash if ever wired up with an fp8 ControlNet. Tests: the float8-marker guard for both fp8 dtypes, the marker is set by the cast itself, and a double cast is a no-op (skipped norm layer stays in compute dtype, hooks registered once). The marker-missing fallback test now simulates a legacy model with delattr instead of locking in the old split. --------- Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com> | 30 天前 | |
Feat(model support): ideogram4 support (#9303) * Update to Transformers 5.1.0 * remove extra stuff * chore(deps): compel fork + transformers>=5.9.0 + remove override Switches compel from PyPI 2.1.1 to invoke-ai/compel@main fork which supports transformers 5.x. Bumps transformers floor to 5.9.0. Removes the transformers>=5.1.0 uv override that was only needed to bypass compel 2.1.1's <5.0 constraint. NOTE: compel fork pulls notebook dep (full Jupyter stack); flag to maintainer for cleanup. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * fix(z_image): resolve rope_theta from rope_parameters for transformers 5.x transformers 5.x no longer exposes rope_theta as a top-level attribute on Qwen3Config; the value is stored in the rope_parameters (and rope_scaling) dict instead. Read it from there with a getattr fallback so the inv_freq buffer is computed from the configured base (1e6 / 256) instead of raising AttributeError. Applies to both the safetensors and GGUF Qwen3 encoder paths. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * fix(model_manager): replace removed hf_hub get_token_permission with whoami huggingface_hub 1.x removed get_token_permission(). HFTokenHelper.get_status() now validates the token via whoami(), which returns user info for a valid token and raises HfHubHTTPError for an invalid one. Preserves the original three-way status: VALID on success, INVALID on HfHubHTTPError (e.g. 401), UNKNOWN on any other error (e.g. network failure). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * chore(deps): regenerate uv.lock after upstream merge Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * fix(sd3): resolve merge conflict marker, drop T5TokenizerFast The upstream merge left an unresolved conflict marker in _t5_encode and reintroduced T5TokenizerFast. Keep our v5 assertion (T5Tokenizer only) plus upstream's new t5_device logic, and drop the now-dead T5TokenizerFast monkeypatch in the test (the name no longer exists in the module). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * style: ruff fixes on merge-resolved files - flux_text_encoder.py: drop unused typing.Union (F401) left by v5 import merge - huggingface.py: ruff format (wrap append(SimpleNamespace(...))) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * chore(deps): pin transformers <5.6 (diffusers single-file CLIP incompat) transformers 5.6 flattened CLIPTextModel (removed the self.text_model wrapper, hoisted embeddings/encoder/final_layer_norm to the top level). diffusers' single-file checkpoint loader (create_diffusers_clip_model_from_ldm) still assumes the nested layout, so loading SD1.5 .safetensors checkpoints fails on 5.6+ with 'CLIPTextModel object has no attribute text_model' and, once that read is shimmed, 'Cannot copy out of meta tensor' (weights never populate the flattened model). Pin to >=5.5,<5.6 (last pre-flattening release) which keeps both the single-file and from_pretrained paths working. The invoke-ai/compel fork accepts any 5.x. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * @ chore(deps): replace compel fork with official compel 2.4.0 compel 2.4.0 (released 2026-05-30) merges the transformers-5 support that the invoke-ai fork carried (both descend from upstream PR #129), plus the maintainer-reviewed padding rework and added diffusers/T5 smoke coverage. Switch from the git fork to the PyPI release. - pyproject: compel git+main -> compel>=2.4.0,<3 - uv.lock: compel 2.3.1 (git 8f404b45) -> 2.4.0 (pypi) - transformers stays 5.5.4 (satisfies compel >=5,<6 and our <5.6 pin) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> @ * feat(ideogram4): backend + model-manager registration for Ideogram 4 Vendor the Apache-2.0 Ideogram 4 reference model (DiT, FLUX2-style VAE, logit-normal flow-match scheduler, nf4/fp8 quant loading) into invokeai/backend/ideogram4/, plus InvokeAI glue (Qwen3-VL text encoding, packed-input build, dual-branch Euler denoise loop). Register the model: BaseModelType.Ideogram4, Main_Diffusers_Ideogram4_Config (detected via the Ideogram4Pipeline class name in model_index.json), and the Ideogram4DiffusersModel loader that loads both transformers as one Ideogram4TransformerPair submodel plus the Qwen3-VL encoder and VAE. Text-to-image only. * feat(ideogram4): Ideogram 4 backend — model manager, invocations, nf4 loading End-to-end text-to-image backend for Ideogram 4, validated through the real session runner. Vendors the Apache-2.0 reference model (DiT, FLUX2-style VAE, logit-normal flow-match scheduler) into invokeai/backend/ideogram4/ with InvokeAI glue. Registers BaseModelType.Ideogram4, Main_Diffusers_Ideogram4_Config, and the Ideogram4DiffusersModel loader (two transformers as one Ideogram4TransformerPair; Qwen3-VL encoder + VAE). Both transformers and the encoder load via InvokeLinearNF4 so they work with the partial-load cache. Adds Ideogram4ConditioningInfo/Field/Output and the model_loader/text_encoder/denoise/l2i invocations. Text-to-image only. * feat(ideogram4): frontend — Regions→JSON prompt, graph builder, UI Wires Ideogram 4 into the canvas/generate UI. buildIdeogram4Prompt assembles the structured JSON caption from the global prompt + Canvas Regional Guidance layers (each region → an obj element with a 0–1000 bbox + desc), with raw-JSON passthrough and a plain-text fallback when there are no regions. Adds buildIdeogram4Graph (text-to-image only, no negative prompt) and the enqueue switch. Structured captions use a static string node + a decoy positive-prompt node so the linear batch can't clobber the assembled JSON; plain text uses the real node so dynamic prompts/batching still work. Registers the 'ideogram-4' base (enums, color, names, model picker, grid size 16), a sampler-preset param (V4_QUALITY_48/V4_DEFAULT_20/V4_TURBO_12) replacing the steps/CFG controls, ParamIdeogram4SamplerPreset, and metadata recall. Regenerates schema.ts. * feat(ideogram4): advanced sampler overrides + color palette Advanced accordion now shows only Ideogram 4-relevant controls. Adds optional overrides of the sampler preset — steps, guidance scale (overrides the main gw, preserves the preset's polish tail), and schedule shift (mu) — plus a color palette editor that injects style_description.color_palette into the auto-built JSON caption (uppercase #RRGGBB, max 16, ignored for raw-JSON prompts). All are nullable (null = use preset), recallable from metadata, and the irrelevant controls (VAE, CLIP skip, CFG rescale, seamless, color compensation) are hidden for Ideogram 4. Backend denoise gains steps/guidance_scale/mu fields; schema.ts regenerated. * Use existing keys + fix select size * Update Readme * feat(ideogram4): add Ideogram 4 to starter models with non-commercial license hint - Implement the weight-only fp8 text-encoder load path (was NotImplementedError); validated against the real fp8 build + add CPU unit tests for the fp8 mechanism - Show Ideogram 4 handlers in the Recall Parameters tab - Recall the assembled JSON caption back into the positive prompt - Translate the metadata "Auto" values - Document Ideogram 4 (install/license, regional-guidance JSON prompting, presets) - Add Ideogram 4 nf4 (CUDA) + fp8 (any device) starter models and bundle - Surface the FLUX-style Non-Commercial License popover for Ideogram 4 models - Note the gated HuggingFace license requirement in the model descriptions * Chore Ruff * Chore Ruff * Chore OpenApi * Chore Knit * fix(deps): regenerate uv.lock to remove duplicate packages from bad merge * fix(ideogram4): make bitsandbytes import lazy in quantized_loading bitsandbytes has no macOS wheels and is excluded on darwin, but the module-level import broke test collection on macOS CI. Move the import into the two bnb-only functions and a TYPE_CHECKING block so the fp8 path imports without bitsandbytes installed. * fix(ideogram4): make bitsandbytes import lazy in quantized_loading bitsandbytes has no macOS wheels and is excluded on darwin, but the module-level import broke test collection on macOS CI. Move the import into the two bnb-only functions and a TYPE_CHECKING block so the fp8 path imports without bitsandbytes installed. * Fix: address ideogram4 review (strict load, 1-step guidance, i18n, runtime caption) - Non-fp8 Ideogram 4 text-encoder load now validates the state dict: unexpected keys raise, missing keys warn (mirrors the fp8 helper) instead of silently accepting a partial load. - Guidance schedule: cap the polish tail at num_steps-1 so at least one main step always remains (the guidance_scale override was silently dropped at num_steps=1), and require steps >= 2 (backend field + frontend slider/marks). - Localize the Ideogram sampler-preset option labels via t() with the step count interpolated; add the three preset i18n keys. - Assemble the structured JSON caption at generation time in a new ideogram4_caption_builder node (Python port of buildIdeogram4Caption) instead of at graph-build time. The graph now wires the real prompt node -> caption builder -> text encoder and returns it as positivePrompt, so dynamic prompts / prompt batching vary the encoded caption (the decoy that dropped them is removed). The builder's output is wired to a new declared ideogram4_caption metadata field via an edge, so each batched image records its actual caption. Regenerates schema.ts for the new node + metadata field. Adds tests for caption assembly, the guidance schedule, and the graph wiring. * feat(ideogram4): step previews + document the model's built-in safety filter - Emit a low-res progress preview each denoise step so the forming image is visible during generation, like the other denoise nodes. Ideogram uses a FLUX.2-style 32-channel VAE, so the packed latent is unpatchified/denormalized (get_latent_norm) and run through the FLUX.2 latent->RGB factors — no full VAE decode per step. The denoise loop now hands the callback the packed grid latent. - Document Ideogram 4's built-in content safety filter in models.mdx: it lives in the model weights (not Invoke's NSFW checker, can't be disabled from Invoke) and false-positives on benign prompts; structured JSON prompts trip it less. * feat(ideogram4): avoid safety-filter false-positives + step previews + caption visibility The main fix: Ideogram 4's built-in safety filter (baked into the model weights) false-positives and returns an "Image blocked by safety filter" placeholder for "degenerate" captions — empirically, an empty `compositional_deconstruction.elements` list, or a single full-frame [0,0,1000,1000] element whose desc just repeats the high_level_description. Our assembly produced empty elements whenever the user drew no regions, so plain prompts were blocked. - Caption assembly (build_ideogram4_caption): - Always emit a structured JSON caption; never bare plain text (the filter false-positives far more on plain text). Raw-JSON pastes still pass through. - When there are no regions, synthesize one default element describing the whole scene from the prompt with a *partial* (non-full-frame) bbox [100,100,900,900]. This never yields an empty/degenerate elements list. Verified end-to-end against the model: the previously-blocked "golden retriever on a skateboard" now renders. - Metadata: always wire the caption builder's output to the ideogram4_caption metadata field, so the viewer's "Structured Caption" row shows the exact JSON that was encoded (not just the raw prompt) for every generation. - Denoise: emit a low-res progress preview each step (unpatchify + FLUX.2 latent->RGB factors, since Ideogram uses a FLUX.2-style 32-channel VAE) so the forming image is visible during generation, like the other denoise nodes. - Docs: document the model's built-in safety filter in models.mdx (it's not Invoke's NSFW checker, can't be disabled from Invoke, and false-positives). Updates the caption/graph tests accordingly (also fixes latent tsc errors in the graph-builder test's core_metadata / ideogram4_caption comparisons). * Chore typegen + openapi + Ruff * Fix Knit * fix(ideogram4): enforce steps>=2 client-side and validate region bbox Address review on the Ideogram 4 PR: - The backend denoise node requires steps >= 2, but the client still accepted ideogram4_steps = 1 in three places, letting a recalled or rehydrated value build a graph that violates the backend schema. Tighten the zod schema to min(2) with `.catch(null)` (a stale/out-of-range value normalizes to null = use the preset instead of failing the whole persisted slice), normalize dispatched values through the schema in setIdeogram4Steps, and refuse an out-of-range value in the ideogram4_steps metadata recall parser. The slider was already min=2. - Ideogram4Region.bbox was an unconstrained Optional[list[int]], so a workflow/API caller could pass a wrong-length or out-of-range box that the caption builder serialized verbatim into the structured prompt. Add a field validator requiring exactly four coordinates, each in 0..1000. Add tests for both: the region bbox contract (valid/None accepted; short, long, negative, and >1000 rejected) and the ideogram4Steps normalization (valid kept, null kept, stale 1 normalized to null on both dispatch and rehydrate). * fix(ideogram4): block unsupported canvas modes/bbox, warn dropped region inputs, constrain bbox Address the latest review on the Ideogram 4 PR: - Canvas readiness allowed unsupported generation modes: Ideogram 4 is txt2img-only (buildIdeogram4Graph asserts it), but a raster layer or inpaint mask makes the compositor pick img2img/outpaint/inpaint, failing only at graph build. Warn in getRasterLayerWarnings/getInpaintMaskWarnings for Ideogram 4 (these already flow into canvas readiness reasons), blocking enqueue up front. - Canvas readiness had no Ideogram 4 bbox check; the backend requires multiples of 16. Add the 16-grid check mirroring the other 16-grid models so an off-grid bbox (e.g. 1025x1024) is blocked instead of failing backend validation. - getRegionalGuidanceWarnings had no Ideogram 4 branch, so a region whose only input is a negative prompt, auto-negative, or reference image looked effective while the graph silently drops it. Warn those inputs are unsupported. - Ideogram4Region.bbox now also rejects inverted boxes (y_min <= y_max, x_min <= x_max). - The advanced-settings badge selector lumped Ideogram 4 into the generic branch, showing stale VAE/clip-skip/CFG-rescale/seamless badges for controls that are hidden for Ideogram. Exclude Ideogram 4 from that branch. - Model bbox as a constrained type (exactly 4 ints, each 0..1000) so the OpenAPI schema advertises minItems/maxItems and item minimum/maximum. Add readiness tests (bbox grid, raster/inpaint blocking, empty-layer allowance, regional-guidance negative/reference-image warnings) and bbox ordering-rejection tests. * fix(ideogram4): reject text encoders with weights left on the meta device _load_text_encoder() builds the encoder under accelerate.init_empty_weights() and previously downgraded missing keys to a warning (both the fp8 path via load_fp8_state_dict(strict=False) and the non-fp8 path). A missing non-tied weight therefore stayed on the meta device, so a bad or mismatched encoder appeared to load and only failed later during device movement or encoding. Add _verify_encoder_fully_materialized(): call tie_weights() to materialize tied weights from their source, then hard-fail if any parameter or buffer remains on the meta device. Wire it into both the fp8 and non-fp8 (incl. bnb-nf4) paths and drop the missing-key warning — genuinely missing non-tied weights are now caught as leftover meta tensors, while tied weights are tolerated. This is a state-based, path-agnostic check. Add tests: passes when fully materialized, raises on a leftover meta tensor from a missing non-tied weight, and tolerates a tied weight resolved by tie_weights(). * chore(ui): prettier formatting for AdvancedSettingsAccordion The Ideogram 4 badge-suppression branch left the wrapped block at its old indentation, failing `pnpm lint:prettier` in frontend-checks. Formatting only. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Your Name <you@example.com> Co-authored-by: 4pointoh <97913726+4pointoh@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com> Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> | 1 个月前 | |
feat: add native Intel XPU (torch.xpu) device support (#9401) * feat(backend): add Intel XPU (torch.xpu) device support Additive xpu branches only: device selection and normalization, float16 default, VRAM queries with a passthrough-VM fallback (missing SYCL free-memory aspect), fp8 layerwise casting via a runtime probe, VAE auto-tiling, partial loading, stats/OOM handling, multi-GPU parallel session execution (device enumeration, config/API validation, worker pinning, and the generation-device options endpoint), and the auxiliary image utilities (depth/SAM/DINO pipelines accept xpu instead of falling back to CPU; cache clearing is device-agnostic). CUDA (incl. ROCm), MPS, and CPU behavior unchanged. Verified end to end on Arc Pro B70 hardware, including dual-GPU worker startup. * test(backend): add XPU coverage for TorchDevice Mock-based, mirroring the CUDA/MPS suites: device choice, dtype, normalize, the xpu_mem_get_info fallback branches, and multi-GPU generation_devices resolution/validation/labeling on XPU. Also makes the auto-without-CUDA generation-devices test hermetic on XPU machines. * build: add [xpu] extra torch 2.7.1+xpu / torchvision 0.22.1+xpu / pytorch-triton-xpu 3.3.1 from the torch-xpu index, gated to linux-x86_64 and win_amd64; uv.lock regenerated. * feat(backend): extend idle-GPU text encoder offload to XPU The idle-device arbiter and the session processor's borrow path both gated on `device.type == "cuda"`, so on a multi-XPU system no device ever registered and `offload_text_encoders_to_idle_gpus` (enabled by default) silently did nothing: encoders kept churning the denoise model in and out of VRAM. Register and lend XPU devices alongside CUDA. MPS is deliberately excluded -- it is always a single shared device, so there is never another GPU to borrow. Verified on a dual Intel Arc Pro B70 host: a text encoder node now runs on the idle GPU while the session denoises on the other ("Running compel on idle device xpu:0 (session device xpu:1)"). * feat(ui): show the executing GPU for XPU sessions Queue items already persist the executing device generically (e.g. "xpu:1"), but both readers dropped it: the session event only forwarded devices starting with "cuda", and the frontend index parser only matched /^cuda:(\d+)$/. On a multi-XPU system the progress circle and queue-item badges were therefore always blank. Accept indexed XPU devices in both places, and correct the queue-item field description, which claimed the device is set only on CUDA. * fix(mm): gate Krea 2 fp8 encoder casting on fp8 storage support The Qwen3-VL encoder kept its fp8 storage only on CUDA, so elsewhere an fp8 checkpoint was loaded as full bf16 (~8.9GB instead of ~4.4GB) and thrashed partial loading when sharing a GPU with a large transformer. Reuse the existing cached `_device_supports_fp8_storage` probe, which already backs the layerwise-casting path. It returns True unconditionally on CUDA, so CUDA behaviour is unchanged. * chore: label XPU devices by index in load logs and fp8 help text Model load lines printed the device index only for CUDA, so every model on a multi-XPU host logged as a bare "xpu device", making it impossible to tell the GPUs apart. The FP8 Storage tooltip likewise claimed CUDA-only support. * test: cover XPU config validation, progress device and fp8 probe Three paths changed by this branch had no coverage: - The `device` field pattern was untested. `test_device_choice_xpu` looks like it covers it, but the config model does not enable `validate_assignment`, so assigning `config.device` skips validation entirely; only constructing the model exercises the pattern. Added constructor-based valid/invalid cases. - `generation_devices` validation was parametrized for cuda/cpu/mps only. - The progress event's device field, which now reports XPU sessions. Also cover `_device_supports_fp8_storage`, which gates FP8 storage in both the generic layerwise-casting path and the Krea 2 encoder: CUDA answers True without probing, CPU is rejected, and a failing XPU probe returns False instead of raising. Each new test was verified to fail when the corresponding fix is reverted. * fix(nodes): recognise XPU out-of-memory errors in the Anima VAE retry The Anima VAE decode catches OOM and retries once with tiling, which caps peak allocation. Detection matched `torch.cuda.OutOfMemoryError` or the words "out of memory" in the message, so it missed XPU entirely: torch's XPU backend does not raise a recoverable `torch.OutOfMemoryError` on exhaustion, it surfaces the Level Zero/UR result code as a plain RuntimeError -- and `UR_RESULT_ERROR_OUT_OF_DEVICE_MEMORY` contains no spaces, so the existing substring never matched. The decode therefore failed outright instead of retrying tiled. Match the `*_OUT_OF_DEVICE_MEMORY` / `*_OUT_OF_HOST_MEMORY` spellings (both UR and ZE prefixes) alongside the existing conditions, and fold the cuDNN/cuBLAS checks into the same case-insensitive comparison. Extends the existing parametrized retry test with the three XPU spellings; each was verified to fail before this change. Note the driver behaviour itself is not reproducible on the hardware used here -- this stack overcommits into host RAM and hangs rather than raising -- so the tests pin the classifier, not the driver. * style: wrap long vram_usage_gb ternary for ruff * fix: drop CUDA-only wording from progress device description Matches the committed openapi/schema artifacts, which already say "on a GPU". * docs: regenerate settings data for xpu device values * fix: stop xpu VRAM probe from reporting an unknown total as zero (0, 0) made the cache's available-VRAM arithmetic collapse to a constant -working_mem budget for the life of the process. Also widen the except: the failure type moves between torch releases (RuntimeError for the missing SYCL aspect, AssertionError from _lazy_init), and warn once when the blind estimate is in use. * fix: probe fp8 support on the target device, per device, without caching failures The probe allocated via an index-less "xpu", which resolves through the thread's current XPU device rather than the device being loaded onto -- so during idle-GPU encoder offload it measured the busy denoise GPU. It was also keyed on device type, letting one device decide for another, and memoised transient failures (it runs during a load, when the device may be momentarily full) with no way back but a restart. Also probe the bf16 upcast, which is the runtime path for Krea-2/FLUX. * fix: pin torch current device when borrowing an idle GPU Worker startup set both the session device and torch's per-thread current device; the offload borrow set only the former, leaving index-less allocations on the worker's own GPU. Extracted the shared helper and guarded it on backend availability. * fix: keep idle-GPU borrows within one device type generation_devices accepts a mixed list, so a cuda session could be handed an xpu device for its text encoder. * feat: detect Intel integrated GPUs via Level Zero torch exposes no is-integrated flag, but Level Zero does (ZE_DEVICE_PROPERTY_FLAG_INTEGRATED), and its loader already ships with the torch+xpu runtime -- so no new dependency and no compiled extension. Use it to keep iGPUs out of `generation_devices: auto` when a discrete GPU exists, and to stop budgeting them as dedicated VRAM (they share system RAM, like MPS). An unknown answer keeps the previous behaviour, an iGPU-only machine keeps its device, and an explicit device list can still opt one in. * feat: add xpu torch index to pins.json Gives the launcher an Intel install option instead of requiring a manual pip install of the extra. * fix: report VRAM diagnostics for the device in use All three sites dispatched on torch.cuda.is_available() first, so a mixed NVIDIA + Arc box running on xpu reported a constant 0.0 GB and logged "CUDA Memory Allocated" -- which would make XPU bug reports unactionable. * docs: record why xpu takes the CUDA VAE constants and keeps the broad OOM needle XPU SDPA was measured on Arc Pro B70 / torch 2.13+xpu: peak memory doubles when sequence length doubles (2.00x across 2048-16384; 2.0 MB at seq=16384 vs 512 MB for a materialised score matrix). So XPU is in CUDA's O(area) regime, not ROCm's math-attention regime, and the existing constants are correct rather than accidental. * fix: derive rand_device metadata from the backend's devices Was hardcoded to 'cuda' for any non-CPU noise, which is wrong on Arc. Falls back to 'cuda' when the device query has not resolved, so Nvidia metadata is unchanged. * docs: add Intel Arc install, driver and VRAM-reporting notes * fix: probe fp8 device support only when a model requests it The probe was the first statement in _should_use_fp8, so it allocated on the GPU during the first load of any model at all -- tokenizer, VAE, scheduler -- and on API/install threads it forced XPU lazy SYCL init on a thread that never generates. Moved below the exclusions. * fix: query Level Zero Sysman for driver-global free VRAM on xpu The blind estimate (total minus this process's reserved bytes) is what made _get_vram_available over-commit on a shared GPU: it feeds a formula that assumes a driver-global figure. Sysman's zesMemoryGetState reports that figure and is often available when the SYCL ext_intel_free_memory aspect is not, so try it before estimating. Measured on Arc Pro B70 with 16 GiB held by another process: Sysman reported 15.553 GiB free, the estimate 31.725 -- a 16.172 GiB error, exactly the foreign allocation. Sysman is not a guaranteed substitute (torch's query bottoms out in the same layer), so the estimate remains as a last resort. * fix: make the fp8 probe mirror the runtime cast path The storage cast happens on CPU while params are still CPU-resident, then the fp8 tensor is copied to the device and the pre-hook upcasts there. Probing all three steps on the device would pass on a build where the host->device fp8 copy or one upcast target fails, and break at forward time instead. Verified on Arc Pro B70 / torch 2.13+xpu: the full sequence works on both cards. * fix: degrade gracefully when a backend cannot name a device torch.xpu.get_device_name goes through _lazy_init, which raises AssertionError on a build without XPU. Naming is used only for labels and logs, so fall back to the device string rather than propagating. * fix: resolve an index-less device in the Sysman VRAM query Returning None for a device with no index would skip the driver-global query and fall through to the blind estimate with no visible symptom. Callers currently always pass a concrete device, so this is a latent hazard rather than a live bug. * fix: declare ctypes prototypes for the Level Zero calls Handles come back from (c_void_p * n)() as plain Python ints, and ctypes converts an undeclared int argument to a C int -- 32 bits. Any handle above 2**31 was being silently truncated; a direct test of that path segfaults. It happened to work on the B70 because the handles fit. Also: release the idle-GPU borrow if re-pinning raises (the setup was outside the try, so a failure there stranded the lock for the life of the process), and report a failing fp8 probe once per device instead of on every model load. * fix: drop the ZES_ENABLE_SYSMAN mutation from the Sysman probe Setting a process-wide environment variable from a read-only query leaks into child processes. It also bought nothing: the variable only gates Sysman on runtimes predating zesInit and must be set before Level Zero initialises, which torch has already done by then. Verified on Arc Pro B70 that zesInit succeeds with the variable unset. * refactor: tidy up the xpu additions after a cleanup review level_zero: cache the loader so it is opened and its prototypes configured once rather than twice, share the driver/device enumeration and its ordering guard between the two probes, and collapse the Sysman pair of globals into one nullable tuple. Also: fp8 support cache is a set (it only ever stored True), the pbr_maps empty_cache is routed through TorchDevice like the PR's other conversions, the shared-memory VRAM branch stops re-testing the device type it matched on, `_auto_generation_devices` partitions in one pass, and rand_device only answers when every generation device is the same accelerator. Merges three duplicate mem_get_info tests into one parametrized case and drops two fp8 probe tests fully subsumed by the cast-sequence test. * build: pin the xpu extra to torch 2.13.0 Intel's XPU backend matured considerably after 2.7.1: torch.xpu.mem_get_info() works on driver/kernel combinations where it previously raised, and the oneAPI user-space runtime ships with the wheel, so upgrading torch upgrades it too. Follows the rocm extra, which already pins ahead of cpu/cuda. pytorch-triton-xpu was renamed triton-xpu upstream. The darwin/aarch64 fallbacks stay on 2.7.1 to match the other extras and the project's torch<2.8.0 constraint on darwin. cpu/cuda/rocm exports are unchanged package-for-package (196/211/197); the only delta is a dropped "via pytorch-triton-xpu" comment annotation from the rename. * test: stub Sysman in the unknown-total xpu probe test Without it the test only passes where Level Zero cannot answer -- never on the Intel hardware the probe exists for, where Sysman returns before the tier under test is reached. * build: teach the pins check about the xpu index Its per-platform allowlist rejects anything unlisted, so pins.json's xpu entry fails it. PyTorch publishes XPU wheels for win32 and linux x86_64, matching the extra's markers. * fix: defer the xpu device pin like cuda's torch.xpu.set_device() brings up a SYCL context that holds VRAM in an otherwise idle process, the same reason the CUDA pin waits for the first claimed queue item. * docs: regenerate settings data on linux Regenerating on Windows flips two path defaults to backslashes, which the docs check rejects. * fix(mm): handle shared memory on integrated GPUs Their VRAM is system RAM, so a RAM copy doubles each model's footprint against the same pool. Drop it, letting a full load move weights rather than copy them. Keep partial loading on -- it is the only path that respects vram_available -- and raise a clean error when a full-load-only model cannot fit, instead of walking into an uncatchable OOM-kill. Warn once when a setting is overridden. Scoped to integrated XPU; CPU and MPS are unchanged. * docs: note intel device selection and integrated-GPU memory `auto` prefers CUDA on a mixed Nvidia/Arc box, and keep_ram_copy_of_weights is ignored on an integrated GPU. * chore(ui): typegen for the xpu device values * fix(mm): compare the integrated-GPU full-load guard against bytes still to move A resident model's weights occupy the same DRAM that vram_available is read from, so its total can exceed "available" precisely because it is loaded. lock() runs on every use and full_load_to_vram() is a no-op when resident; comparing the total refused the re-lock and evicted a healthy model on every other generation. Compare what full_load_to_vram() will actually move instead. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test: run the integrated-GPU cache tests on CPU-only torch ModelCache.__init__ sizes the RAM cache from the device's total VRAM, which on an xpu execution device reads torch.xpu.get_device_properties() -- an AssertionError on the CPU-only builds CI runs, failing 9 of these tests before they reached their subject. Stub a fixed total during construction, and add a regression test for the resident-model re-lock guard. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: LexiconCode <aaronwalker@protonmail.com> Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> | 22 天前 | |
ruff format | 11 个月前 | |
perf(qwen-image): add a tiling option to the Qwen-Image VAE nodes (#9427) * feat(qwen-image): add a tiling option to the image-to-latents node The Qwen-Image i2l node hardcoded vae.disable_tiling(), so a full-frame encode was the only option. At 2560x1440 that peaks at 9.26 GiB — on top of a resident multi-GB transformer, which is what makes an upscale round-trip run out of headroom exactly at this node while every other node fits. Adds `tiled` / `tile_size` input fields following the SD/SDXL i2l node, OR'd with the global force_tiled_decode setting. Off by default, so behaviour is unchanged unless enabled. estimate_vae_working_memory_qwen_image gains a matching tile_size parameter. Without it the change would be inert: the cache would keep reserving the full-frame figure (10.99 GiB at 2560x1440) and evict models to honour it, no matter what the VAE actually does. Tiled, it budgets one tile plus 25% overlap plus the resident RGB image, mirroring estimate_vae_working_memory_wan. Measured through the node at 2560x1440: 10.99 -> 0.26 GiB reserved, 9.26 -> 0.17 GiB actual peak, identical latent shape. Tiled latents differ by ~1.4% relative L2 on noise input (worst case for tile blending; real images blend far better), which is why this stays opt-in. * Add test * feat(qwen-image): make VAE tiling usable on both Qwen-Image VAE nodes Both nodes reserve working memory for a full-frame operation, which at high resolutions exceeds a 24 GB card, so the model cache evicts everything else to honour it. On CUDA at 2560x1440: 19.91 GiB for the decode and 10.99 GiB for the encode. Tiling is the intended escape hatch, but it did not work on either node: - qwen_image_i2l hardcoded vae.disable_tiling(), so it could not be enabled. - qwen_image_l2i honoured the global force_tiled_decode, but computed its working-memory estimate before and independently of that flag. Tiling bounded the VAE while the cache still reserved the full-frame figure, so the memory was never freed for anything else — effectively inert. Adds `tiled` / `tile_size` input fields to both nodes following the SD/SDXL i2l/l2i nodes, OR'd with force_tiled_decode. Off by default; behaviour is unchanged unless enabled. estimate_vae_working_memory_qwen_image gains a matching tile_size parameter, and both nodes resolve tile_size=0 to the VAE default (256px) before estimating. Tiled it budgets one tile plus 25% overlap plus the resident RGB image, mirroring estimate_vae_working_memory_wan. Without this the change would be cosmetic on i2l and remain inert on l2i. Measured through the i2l node at 2560x1440: 10.99 -> 0.26 GiB reserved, 9.26 -> 0.17 GiB actual peak, identical latent shape. Verified across eight resolutions that tiled and untiled encodes produce the same latent dimensions. Tiled latents differ by ~1.4% relative L2 on noise input (worst case for tile blending), which is why this stays opt-in. Also fixes a crash in qwen_image_i2l: `width`/`height` are `int | None`, but the workflow UI sends 0 for an unset number input, and `0 is not None` reached `image.resize((0, 0))` -> "height and width must be > 0". Non-positive values are now treated as unset, matching how tile_size uses 0. * fix(qwen-image): pass a matched tile stride and scope the tiling state enable_tiling() was called with tile_sample_min_* only, leaving the stride at the module's 192px default. The tile loops step by stride but slice each accumulated tile to min, so any tile_size below 192 silently dropped whole bands of the image -- a 128px tile turned a 512x512 decode into 384x384 with no error -- while sizes above 256 grew every tile without removing any, making compute scale with tile_size^2 (8x a full frame at 512px). Pass all four parameters with the stock 4:3 ratio, rounding the stride down to a multiple of the 8x spatial compression so the pixel and latent steps agree. Tile sizes below 64px are clamped; the field carries the 0 "use default" sentinel and so cannot take a pydantic lower bound. enable_tiling() also writes straight onto the module, and disable_tiling() only clears use_tiling. That module is the model cache's own instance, so a tile size set once persisted for the lifetime of the cache entry and leaked across invocations and into anima_latents_to_image, which shares the instance. Apply the geometry through a context manager that restores it, and resolve the 0 sentinel against a constant instead of the module's current value. Also budget the pixel-space buffers tiled_decode holds simultaneously (~5 frames, not 1) -- the term that grows with output area, so it degraded in exactly the regime tiling exists for. * test(qwen-image): pin the multiple-of-8 tile-stride rounding The rounding in `_tile_stride_for` was load-bearing but uncovered: dropping it left the whole suite green, yet a raw 3/4 stride silently truncates the encode for any tile size whose 3/4 is not a multiple of 8. `tile_size` is `multiple_of=8`, so 72 and 80 are both reachable from the workflow UI and both land there (54 and 60). Cover it at the argument level (72 -> 48, 80 -> 56) and end-to-end against a real tiny VAE, plus a pin on the failure mode itself: min=72 with the un-rounded stride 54 turns a 512x512 encode into a 57x57 latent instead of 64x64, with no exception. Also correct the rationale on QWEN_IMAGE_VAE_MIN_TILE_SIZE. `_tile_stride_for` already floors the stride at 8, so the derived latent step never collapses to 0, and 8/24/32 all round to clean multiples of 8 and produce correctly sized output -- 64 is not the smallest valid tile. It is a cost floor: the tile count grows with the inverse square of the stride (1620 tiles at 64px versus 57,600 at 8px on a 2560x1440 frame). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> --------- Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> | 19 天前 | |
Run ruff | 2 年前 | |
fix(model cache): evict records at shutdown() instead of only releasing shared weights (#9494) * fix(model cache): release shared weights when a cache goes away Nothing released a cache's SharedCpuWeightsStore references except _delete_cache_entry(): shutdown() left every resident record's refcount held, and a cache dropped without shutdown() (test teardown; any future wiring that rebuilds caches at runtime) stranded the canonical tensors and their accounting forever. Today's production wiring tears the store down together with its caches, so the live exposure is cross-test pollution of the process-global store and RAM pinned past ModelManagerService.stop() — but the refcount invariant ('every acquire is paired with exactly one release') was simply not upheld, and this makes it self-healing before any wiring change turns it into a real peer-accounting bug. Two mechanisms, for the two ways a cache goes away: - shutdown() now releases its resident records' shared references synchronously — it runs in a normal thread context, so the direct (locking) release is safe there, and teardown does not depend on a later store operation happening. - Each wrapper registers a weakref.finalize fallback for the dropped-without-shutdown case. The finalizer runs in GC context, where taking the store's non-reentrant lock could self-deadlock (a collection can fire inside acquire()'s critical section on the same thread — the rule ModelCache.release_first_use_grace documents), so it only ENQUEUES into a SimpleQueue; every public store method drains the queue under the lock. The finalizer is registered inside the acquire's try (a registration failure must release too), its args carry the key and canonical dict rather than the wrapper (finalize holds args strongly — referencing self would make the wrapper immortal), and release_shared_weights() detaches it before releasing synchronously so eviction-then-collection releases exactly once. The state-dict identity keeps releases correct across invalidate()'s retired entries. RamBudget.total_in_use() now documents why its store read must stay outside the budget lock: the drain allocates under the store lock, so GC can run _on_cache_collected (store→budget) there, and a budget→store order anywhere would complete the deadlock cycle. Six regression tests, verified to fail before the fix, covering: shutdown releases synchronously with an empty queue; collection returns refcount/bytes/budget to zero; the collection-time release is enqueue-only (never applied inline by GC); eviction + collection release exactly once across two caches; a retired (invalidated) entry is freed by a collected holder; and the partial-load wrapper behaves like the full-load one. One existing test relied on an abandoned wrapper leaking its reference and now binds it. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(model cache): evict records at shutdown() instead of only releasing shared weights shutdown() released the resident records' shared-store references while retaining the records themselves, so the accounting stopped describing reality: - The store (and RamBudget) reported zero for bytes whose tensors the retained wrappers still held. - A post-shutdown load of the same key on a peer cache registered a duplicate canonical alongside the still-resident released copy. - A post-shutdown eviction of a released record (put() after shutdown() is reachable: Invoker.stop() stops the model manager before the session processor) read uses_shared_weights as already-False and debited the non-shared budget for bytes that were admitted as shared. shutdown() now routes idle records through _delete_cache_entry(), which releases shared ownership and budget accounting together, exactly once. Records still in use — locked by an in-flight generation or inside the put()->lock() admission window — keep their references and are marked stale; unlock() evicts them through the existing stale path when the generation lets go, so the accounting stays truthful at every point. All five regression tests verified to fail against the previous shutdown() behavior. Follow-on to #9403, addressing JPPhoto's review comment there. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(model cache): match records by identity in stale eviction and _delete_cache_entry Surfaced by adversarial review of the shutdown() change: a stale-marked record can be detached while still locked (the VRAM-move error paths call _delete_cache_entry on a locked record) and its key re-admitted before the record's last unlock(). The stale-eviction path matched by key only, so it popped the NEW record — detaching it from the cache and all accounting — and, the old record's shared release having already happened, read uses_shared_weights as False and debited the non-shared budget for bytes that were admitted as shared. The hazard predates the shutdown() change (drop_model() sets the same flag), but shutdown() now arms stale marks at every server stop that overlaps in-flight work, so close it here: _delete_cache_entry() and unlock()'s stale eviction act only when the record passed in IS the record currently held under its key; a delete of a detached record is a full no-op. Regression test verified to fail against the key-only matching. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(model cache): track get()->lock() holders through shutdown and abandonment Two defects found in review of the shutdown eviction change (JPPhoto, 2026-08-13): 1. shutdown() racing the gap between get() and the LoadedModel's first lock evicted the warm record out from under its holder: the holder locked a detached record whose shared-store ownership had just been released, so a peer's reload of the same key minted a duplicate canonical copy while the budget counted one. 2. A record retained by the shutdown sweep for a never-locked holder could never be evicted if that holder was simply dropped: the abandonment finalizer's deferred work was discarded post-shutdown (and the worker was stopped), pinning the record, its shared-store refcount and its budget bytes for the life of the process. The fix tracks every wrapper's get()->lock() window with a per-record hold count (CacheRecord.first_use_holds), armed in LoadedModelWithoutConfig's constructor and released exactly once per wrapper — on its first lock, or by its weakref finalizer if it is dropped un-entered. Held records are treated like locked ones by every eviction path (shutdown, budget reconcile, peer-requested eviction, make_room, drop_model, unlock's stale eviction); stale-marked records whose last holder is abandoned are evicted by the deferred worker, which now outlives shutdown() for exactly that purpose (it already exits via the cache-collection finalizer). Holds are only granted while a worker is alive to carry the finalizer's release, and a worker death zeroes surviving holds at the next start so no record can stay shielded with nothing left to unshield it. Admissions landing after shutdown() are marked stale at birth so their final release evicts them too. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(model cache): epoch-guard hold releases and recover stranded holds at shutdown Hardening from adversarial review of the first-use-hold mechanism: - Hold releases (the wrapper's first-lock release and the abandonment finalizer's deferred release) now quote the epoch the hold was armed under, and dead-worker recovery bumps the record's epoch when it zeroes stranded holds. Without this, a surviving wrapper's late release — or a release enqueued before the worker died and drained after the restart — would decrement a fresh hold armed by a different wrapper under the healthy replacement worker, silently unshielding that wrapper's window. - shutdown() now runs the dead-worker hold recovery itself (and clears the put()-grace flags in the same situation): a hold whose abandonment release was dropped by the dead-thread dispatch check has no other releaser, and after shutdown no put() is guaranteed to run the usual next-start recovery — the sweep would stale-retain the record, its shared-store refcount and its budget bytes for the life of the process. - register_first_use_hold() declines to arm on a record that is no longer the occupant under its key: an eviction already won the race against the wrapper's construction and a hold on a detached record shields nothing. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(model cache): withhold the post-shutdown grace and recover from the worker's own death Two follow-ups from review. put() after shutdown() no longer arms the post-admission grace. That flag's backstop releaser is the sweep at the top of the next put(), and after shutdown no further put() is guaranteed: a load cancelled between put() and the LoadedModel's construction leaves no wrapper (hence no finalizer either), so an armed flag would stand for the life of the process, hiding the record from every asynchronous eviction path while its bytes stayed charged to the shared budget. Withholding it costs only the shield -- the record stays stale at birth, so its eventual release still evicts it, and a loader that does come back gets the ordinary first_use_holds shield. The deferred worker now runs stranded-shield recovery from inside its own dying frame. Previously recovery depended on something else happening first -- the next admission, or shutdown() -- and neither is guaranteed when the worker dies *after* shutdown()'s liveness check: the records the shutdown sweep retained for a live holder were left shielded by holds nothing could release. The recovery is scoped by thread identity (a replacement worker's shields are its own) and retires the worker slot before sweeping, so a concurrent admission cannot arm a shield the recovery is about to zero. It also drains the queue the dead worker left behind, whose _AbandonedHolderRelease items pin their models' CPU weights. _ensure_deferred_worker and shutdown() now share the same recovery, which also lifts orphaned admission graces and evicts whatever that leaves unshielded. Three tests, each reverted-and-confirmed-failing against the code it guards. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GjAw9hpJe8d1GyvpdKJFGw * fix(model cache): narrow the dead-worker recovery and stop it pinning records Adversarial review of the previous commit found three problems with the shared recovery it introduced. The recovery cleared the put()-set admission grace unconditionally. On a live cache that is a new failure mode, not a fix: the dying worker is is_alive() for as long as it unwinds, so a cold load landing in that window starts no replacement worker and is admitted with the ordinary grace, which the recovery then zeroed while the loader was still between put() and get() -- a reconcile could evict the record and the loader's get() would raise IndexError. The grace only actually loses a releaser once the cache is shut down (its backstop is the next put()'s sweep, not the worker), so it is now lifted only then. The recovery also evicted stale-unshielded records from _ensure_deferred_worker, which register_first_use_hold calls before arming -- so a second wrapper's construction could detach the very record it was about to shield, releasing shared-store ownership while live wrappers still held the tensors. That is the accounting lie shutdown() itself refuses to make. The eviction moved to _evict_stale_unshielded_entries, called only from the dying worker and only on a shut-down cache, where nothing else can ever run it; it now also collects and empties the device cache the way the other abandonment path does. Keeping shutdown()'s call to pure field assignments restores its old property that the branch cannot raise before the resident-record sweep. The queue drain the previous commit added did not close the pin it targeted: _dispatch_deferred's liveness gate is unsynchronized, so a finalizer that read the worker slot just before it was retired still enqueues after the drain. The drain is gone; _AbandonedHolderRelease now holds its record weakly instead, so a stranded item pins nothing, and the worker clears the strong reference it resolves before parking on the next get(). Also moves _reconcile_budget_if_pending's lock acquisition adjacent to its try: a BaseException between the two leaked the cache RLock to an unwinding thread, blocking every other thread for the life of the process. Five tests, each reverted-and-confirmed-failing against the code it guards. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GjAw9hpJe8d1GyvpdKJFGw * fix(model cache): key the dead-worker backstops on worker liveness, not the slot A second adversarial pass found that retiring the worker slot from inside the dying worker silently disabled both remaining recovery sites, which gated on "a dead thread still occupying the slot". The dying recovery deliberately leaves a live cache's admission grace standing -- the next put()'s sweep is still its backstop -- and hands the lift to shutdown(); with the slot already empty, shutdown() skipped it and stale-retained the record, its shared-store reference and its budget bytes for the life of the process. Both gates now key on "no live worker": shutdown() lifts when the slot is empty or dead, and the worker start recovers unconditionally (it has already returned if a worker is alive). That also makes a failed recovery retryable, which matters because the recovery was not exception-safe and had already retired the slot by the time it could raise. _clear_stranded_first_use_holds now unshields every record before reporting any of them -- a logging handler that raises is one of the ways the worker dies in the first place, and logging inline let that same handler abort the sweep partway -- and the post-eviction gc/empty_cache housekeeping, which the codebase already documents can raise from a sick CUDA context, no longer takes the eviction down with it. Also corrects two overstated claims in the weakref rationale: a stranded queue item can be drained later by a replacement worker (the queue is per-cache, not per-worker), and the hold decrement in _release_abandoned_holder runs before the identity check -- it is inert on a detached record for a different reason, which the docstring now gives. Moving _reconcile_budget_if_pending's acquire adjacent to its try narrows the RLock-leak window rather than closing it; said so. Three tests, each reverted-and-confirmed-failing against the code it guards. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GjAw9hpJe8d1GyvpdKJFGw * fix(model cache): claim the first-use window at the lookup, and let a cancelled admission release itself Two findings from review. The first-use shield was armed by LoadedModel's constructor, leaving the whole stretch between the cache lookup and that constructor unshielded -- and that stretch is not a couple of instructions: the configured loader retrieves its record inside _load_and_cache and then does the shared-store shell registration and two returns before load_model wraps it. A shutdown sweep or a peer's reconcile landing there detached the record its holder was about to lock, releasing shared-store ownership while the tensors lived on, so a peer's reload minted a duplicate canonical the budget counted once. ModelCache.get_with_first_use_claim() now arms the hold in the same lock acquisition as the lookup and hands back a FirstUseClaim that owns it: the wrapper adopts the claim and releases it at its first lock, and a claim dropped without ever being adopted -- the load raised before a wrapper existed -- releases the hold by dying. shutdown() stale-retained a record carrying only the put()-set admission grace. That grace's three releasers are the loader's own get()->lock(), the abandonment finalizer of a wrapper built from the record, and the sweep at the top of the next put(); a load cancelled between its put() and its retrieval has neither of the first two, and after shutdown no further put() is guaranteed to run the third, so the record, its shared-store reference and its budget charge stood until the cache object was collected. put(claim_admission=True) now hands the loader a claim over that window too, so such a load releases its admission by dying and the shutdown sweep finds an ordinary idle record. Retiring the grace at shutdown instead -- the obvious shortcut, and what the first two drafts of this commit did -- is not safe. The flag is unowned, so a standing grace does not mean nobody is working on the record: it is equally the state of a load still between its put() and its retrieval, and of a live un-entered wrapper whose hold a worker death zeroed. Both were evicted out from under their holder, with the duplicate-canonical accounting lie and an IndexError from a retrieval that no longer found its own model. For the same reason the admission window is not shielded by either flag but by a weak reference to the claim (CacheRecord.admission_claim_ref): nothing has to release it, so neither a worker death (which zeroes holds) nor another holder's abandonment (which clears the grace) can make a running load look finished. _recover_stranded_shields retires it once the cache is shut down, where the eviction its expiry should trigger would otherwise travel through a dead worker -- the same trade that method already makes for holds. The claim is armed only after put() has committed its accounting, so a failure to allocate it cannot leave a resident, store-owning record the budget never counted, and both hand-back guards release the hold when the object that was to carry its release cannot be built. Nine tests, each reverted-and-confirmed-failing against the code it guards. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FCitU6EFcp76AaauNfaWTz * fix(model cache): refuse post-shutdown prefetch admissions, and keep abandonment releases Two records could survive a shut-down cache with nothing left to retire them. A prefetch=True admission after shutdown() was marked stale and inserted anyway. prefetch is the promise that no loader will come back for the record, so there is no get() -> lock() -> unlock() to run the stale eviction, no wrapper whose finalizer could carry an abandonment release, and no claim whose expiry could stand in for either. What is left are the paths that may or may not run -- another admission's make_room, a budget reconcile, a peer's eviction request -- and after shutdown none of them is guaranteed to come. put() now refuses such an admission, the same standard the post-admission grace is already withheld under. It costs only a reload: the sole caller takes the submodel it asked for from the pipeline object, not from the cache. The refusal sits after _ensure_deferred_worker() (a post-shutdown prefetch must still revive the worker that carries the retained records' abandonment releases), after the stale-grace sweep (on a shut-down cache with a LIVE worker that sweep is the only backstop a stale grace has left, since shutdown() runs the recovery only when no worker is alive), and before _make_room_internal (nothing resident should be evicted to house a model this call is about to refuse). _dispatch_deferred dropped every item while no worker was running, which included _AbandonedHolderRelease. Its holder is already gone -- finalizers fire once -- so no lock, no unlock and no second finalizer is coming, and that item is the only thing left that can retire the record. Dropping it stranded a record the shutdown sweep had retained, resident and charged, for the life of the process; no later sweep can repair that, because once dead-worker recovery zeroes the hold, a record whose holder is gone is indistinguishable from one a live wrapper is still holding, where retention is required. Abandonment releases are now kept and drained by whichever worker runs next; reconciles are still dropped, since cached_model_keys() can request one on every call and the next cache operation's release hook re-runs it anyway. Evicting from _ensure_deferred_worker() is deliberately NOT the fix: it runs _recover_stranded_shields() immediately before, so at that instant a record a live wrapper still holds looks unshielded, and evicting it would release shared-store ownership while the tensors live on. Four tests, each checked for sensitivity by reverting the line it guards. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015n81CKSi7bUA217L2E9SwN * fix(model cache): bound the abandonment queue, own worker-less admissions, keep live claims through recovery JPPhoto's round-6 findings on the shutdown accounting, all three confirmed: - A grace-only abandonment enqueued one kept item per dropped wrapper while no worker could be started, and nothing short of a lock() or another put() ever cleared the grace that kept them coming, so a warm get/drop loop grew the queue without bound. release_first_use_grace now clears the grace itself, lock-free (the flag is monotonic), queues an eviction only for a stale record and only once per record (CacheRecord.abandonment_release_pending, re-opened by the worker the moment it dequeues the item), and still wakes the worker for a pending budget reconcile the drained item used to run. - A post-shutdown put(claim_admission=True) with no startable worker got no claim, so nothing owned the record once its loader died and no later admission swept it. _claim_first_use now mints a hold-less FirstUseClaim (hold_epoch=None) whenever the record is still the occupant, so the admission stays owned through CacheRecord.admission_claim_ref and its finalizer still queues the eviction a stale record owes; and put() runs _evict_stale_unshielded_entries() when the cache is shut down and no worker is alive after its revival attempt, the same terminal sweep the dying worker runs, placed after the grace sweep and before the prefetch refusal. - _recover_stranded_shields retired a live admission claim on a shut-down cache, so a worker death followed by shutdown() evicted a record whose loader was between put() and get(), and that loader's retrieval raised IndexError. A live claim now survives every recovery; the eviction owed once it dies travels through the kept queue item or the next admission's sweep. Two further defects surfaced by the adversarial passes over the fix: shutdown() now marks a record stale BEFORE consulting its shield, so a hold-less abandonment racing the sweep either sees the mark and queues the eviction or has already cleared its shield when the sweep looks; and the coalescing gate is opened at the dequeue site rather than in the handler, so a raise in the handler, a worker death inside it, or a fork cannot leave it closed with nothing queued behind it. Ten new tests, two rewritten; each production change reverted individually and confirmed to fail only the tests that guard it. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011JbML8Y8mHdXH72bcUry9o * fix(model cache): refuse worker-less post-shutdown admissions JPPhoto's round-7 blocker, the residual I disclosed rather than closed last round: a claimed admission into a shut-down cache with no startable worker has no releaser (no worker to drain the eviction its dropped claim's finalizer queues) and no guaranteed future cache operation to stand in for one, so a load that dies before retrieving it pins the record, its shared-store reference and its budget bytes for the life of the cache object. put() now refuses every admission -- claimed, plain, and prefetch -- once the cache is shut down and no worker is alive after its revival attempt, the same standard the post-shutdown prefetch was already refused under, generalized. Nothing that could be stranded is admitted. The refusal is narrowly scoped: a live cache still admits worker-less (a future op cleans up), and a normal shutdown keeps its worker alive so the graceful retain-and-reclaim path is unchanged -- only thread exhaustion reaches the refusal, where a load racing shutdown is failing regardless (its retrieval raises IndexError, as a refused prefetch's does; both loaders already handle a None put()). The terminal sweep (_evict_stale_unshielded_entries) still runs before the refusal returns, so a record already orphaned by a worker death is reclaimed even though this admission -- possibly the last cache operation -- is refused. Three tests whose orphan came from the now-refused admission are removed; their mechanism coverage survives elsewhere. Two added: the refusal with its clean loader-style IndexError, and the terminal-sweep-still-runs case. Both new production lines mutated and confirmed to fail only their guarding tests. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011JbML8Y8mHdXH72bcUry9o * fix(model cache): clean shutdown-retained records --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com> Co-authored-by: JPPhoto <jpollack@jpollackphoto.com> | 1 天前 | |
fix(backend): revert non-blocking device transfer In #6490 we enabled non-blocking torch device transfers throughout the model manager's memory management code. When using this torch feature, torch attempts to wait until the tensor transfer has completed before allowing any access to the tensor. Theoretically, that should make this a safe feature to use. This provides a small performance improvement but causes race conditions in some situations. Specific platforms/systems are affected, and complicated data dependencies can make this unsafe. - Intermittent black images on MPS devices - reported on discord and #6545, fixed with special handling in #6549. - Intermittent OOMs and black images on a P4000 GPU on Windows - reported in #6613, fixed in this commit. On my system, I haven't experience any issues with generation, but targeted testing of non-blocking ops did expose a race condition when moving tensors from CUDA to CPU. One workaround is to use torch streams with manual sync points. Our application logic is complicated enough that this would be a lot of work and feels ripe for edge cases and missed spots. Much safer is to fully revert non-locking - which is what this change does. | 2 年前 | |
fix(krea2-lora): validate kohya keys by un-flattening, not by prefix (#9518) `lora_unet_blocks_<idx>_` is not Krea-2's spelling alone — Wan writes it verbatim and Anima writes `lora_unet_[llm_adapter_]blocks_<idx>_`. Matching it as a prefix accepted a mislabeled Wan or Anima LoRA under an explicit Krea-2 override, where it installed and then silently no-op'd at generation time: the un-flattener rejects `self_attn`/`cross_attn`/`mlp_layer0`, so every layer warn-skipped. Install-time validation exists to prevent exactly that failure mode. Ask the converter's own un-flattener instead: a kohya key is a Krea-2 key only if its flattened path reconstructs to a leaf of the native module vocabulary. That also fixes the converse — the doubled-separator spelling (`lora_unet__blocks_...`) that the converter deliberately tolerates but no prefix spelled out, so a transformer-only adapter written that way could not be installed at all. Both come from one change: `split_kohya_krea2_key()` is now the only place that splits a kohya key, so the converter's per-module gate, the rewrite it guards, and identification cannot disagree about which module a key is in. The reconstruction helpers move to `krea2_lora_constants` because `configs/lora.py` cannot import the converter — it pulls in the patch layers, which import `model_manager.load`, closing a cycle back into `model_manager.configs`. Same reason `anima_lora_constants` exists. Side effect, tested: an adapter on an `nn.Sequential` position that holds no Linear (`lora_unet_tmlp_1`, `tproj_0`, `txtmlp_2`) no longer installs. It previously matched by prefix and then warn-skipped at apply time. Follow-up to #9449, addressing the two non-blocking notes from review 4888833569. | 9 天前 | |
Feat: pid followup (#9474) * fix(pid): harden PiD decoder identification, GGUF loader tests and docs Follow-up to #9281, addressing the review items: - Require the complete LQ projection. Identification accepted a checkpoint with a single `lq_proj.*` key and `load_pid_decoder` tolerated every missing `lq_proj.*`. Models are built under `skip_torch_weight_init()`, so those weights stayed uninitialised and would decode to garbage/NaNs. `required_lq_proj_keys()` derives the expected key set from the vendored network, and both identification and load now reject any missing key. - Match the install source when identifying backbone and variant. A direct single-file install is stored as `<uuid>/model_ema_bf16.pth`, so the name carried no `res2k…` marker and no backbone hint: SDXL/Qwen-Image decoders were labelled `res2k_sr4x` although only the 2K-to-4K preset exists, and SD3/Qwen-Image decoders were registered as `flux`, which their decode node then rejects. The source (HF path/URL) survives the download and is now matched alongside the on-disk name, with a per-backbone variant fallback. - Replace the hardcoded-path Gemma-2 GGUF loader test (always skipped in CI) with a synthetic tiny Gemma-2 built from mocked GGUF tensors: asserts quantized retention, norm materialisation, meta-buffer repair, absence of meta parameters and a finite forward. The real-file comparison is now opt-in via INVOKEAI_TEST_GEMMA2_GGUF. - Drop stale doc claims: PiD as a decode is now documented separately from the prototype `pid_upscale` node, the starter checkpoints are spread over `checkpoints/` and `checkpoints_deprecated/`, and the GGUF encoder is loaded natively instead of being dequantized by transformers. * fix(pid): check LQ completeness before the backbone, stop probing the RNG Review follow-ups for #9474. 1. `_raise_if_lq_projection_incomplete` ran after `_validate_base`, but the backbone is read from `lq_proj.latent_proj.0.weight` — one of the weights a truncated file may be missing. Such a file therefore failed with "cannot determine PiD decoder backbone" instead of the "missing … LQ projection weights" message the install flow promises. Completeness is now checked first, against `common_required_lq_proj_keys()`, the key set every backbone requires. The per-backbone check stays after the backbone is known: the two sets are the same 71 keys today, so it is a no-op that keeps the check from silently weakening to the intersection if a backbone ever adds LQ parameters of its own. Note this does not change the `Unknown_Config` fallback: `ModelConfigFactory` applies that to any file no config matches, and `allow_unknown_models` defaults to true. With it disabled the truncated file is rejected outright. The PR's QA step is worded as if rejection were unconditional; it is not, and that is a model-manager-wide behaviour rather than anything PiD-specific. 2. `required_lq_proj_keys()` built a real `LQProjection2D` just to read parameter names, running every `reset_parameters()` and so drawing from the global CPU RNG during model identification — leaving later unseeded randomness dependent on how many candidate files were probed. It is now built on the meta device inside `torch.random.fork_rng`. 3. The `net.` normalisation existed twice, and the copies had already diverged: only the loader dropped the distill-only submodules (`net_ema.`, `fake_score.`, `discriminator.`). Since `net_ema.*` shadows PidNet's own parameter names, the drift ran in the direction where identification accepts what the loader then refuses. Both now share `backend/pid/state_dict_utils.py`. * fix(pid): reject a recognised-but-broken PiD checkpoint instead of registering it Identification rejected a truncated PiD checkpoint with NotAMatchError, which only means "not my kind of model": ModelConfigFactory collects those, finds no match, and with allow_unknown_models (default: true) falls back to Unknown_Config. A file that had already identified itself as a PiD decoder and was then found to be missing LQ projection weights was therefore still installed, as an unknown model with a database record, and only failed once something tried to load it. Add InvalidMatchError for "recognised, and unusable". It is deliberately not a NotAMatchError subclass, since the factory catches that one per candidate class and would swallow it. When no config class matched and at least one raised it, classification returns no config regardless of allow_unknown, and ModelInstallService._probe reports the specific reason rather than the misleading "could not identify model". Order the architecture check ahead of the completeness check. A v1.5 checkpoint is intact, just built to a shape InvokeAI cannot construct; judged against the legacy key set it would be misreported as truncated and now hard-rejected on top of that. It stays a plain no-match, so it remains registrable as an unknown model - only a broken file is fatal. This costs the completeness check nothing: the hidden dim is read from lq_proj.latent_proj.0.weight, so a file truncated past that weight falls straight through to it. Collapse the LQ key contract to one entry point. common_required_lq_proj_keys() and the per-backbone re-check are gone; required_lq_proj_keys() takes no backbone, the probe lives in the private _probe_lq_proj_keys(), and test_pid_decode.py pins that every backbone agrees, so key drift fails in CI instead of silently weakening the install-time check. * fix(pid): make every backbone-independent PiD rejection final The previous commit made a truncated checkpoint fatal but left the architecture check a plain no-match, on the reasoning that an intact v1.5 file is merely unsupported and should stay registrable. That opened a hole: a file that is both 1024-dim and truncated is rejected by the architecture check first, never reaches the completeness check, and lands back in Unknown_Config - the exact outcome the previous commit set out to prevent. The distinction does not survive contact with the failure mode. Once a file has identified itself as a PiD decoder, any rejection that does not depend on which backbone it is will be raised identically by all five config classes, so the file ends up with no match and is registered through the Unknown_Config fallback. Those rejections are now all InvalidMatchError: unsupported lq_hidden_dim, an incomplete LQ projection, a latent channel count no backbone uses, and a checkpoint whose backbone cannot be determined at all. Splitting them out of _validate_base is what makes that legible. _validate_base now only ever answers "not *this* backbone", which four of the five classes are supposed to say about every valid checkpoint, and every rejection in it stays a NotAMatchError. The backbone-independent checks run ahead of it in from_model_on_disk, architecture first so an intact v1.5 file is diagnosed as unsupported rather than judged against the legacy key set and misreported as truncated. Also handle InvalidModelConfigException in the startup orphan scan. ModelSearch._walk_directory already contains anything the on_model_found callback raises, so startup was never actually at risk; catching it in the callback makes skipping a bad file a property of the scan rather than of its caller, and names the file and the reason in the log. * fix(pid): hold a checkpoint to PidNet's whole contract, and stop guessing from paths Identification checked less than the loader demands and inferred the rest from file paths. Three consequences, all reported in review: A checkpoint with every lq_proj weight and none of the 385 backbone weights was registered, then refused by load_pid_decoder. Only the 71-key LQ projection was ever checked. A subset check is not a milder version of the same guarantee: loaders run under skip_torch_weight_init(), so a weight the checkpoint does not supply is uninitialised memory rather than a default. required_pid_net_shapes() now derives the whole contract - 456 keys and their shapes - from a meta-device PidNet, the same trick the LQ probe already used but applied to the real network instead of one submodule. Missing keys, unexpected keys and wrong shapes are all fatal, because all three are fatal in load_pid_decoder; a stricter installer cannot reject a file that would have loaded. Probing the real net also removes the reason _LQ_PROBE_DIM, _LQ_NUM_RES_BLOCKS_DEFAULT and the hand-copied num_outputs derivation existed, along with the test that kept them in sync. Wrong-shaped tensors were accepted when a filename supplied a backbone. The architecture, the backbone and the kernel are all read off lq_proj.latent_proj.0 .weight, and each read answered None when it was not a 4D conv - so one malformed tensor made all three abstain at once and the file fell through to name-only matching. That weight is now validated first, and the three reads only run when it is there; its absence is a truncation, which the contract check diagnoses better than a guess about the architecture. Backbone detection concatenated the install source, the parent directory and the filename into one string and substring-matched it with a fixed precedence, so /flux/model_sd3.pth matched flux first and was registered as FLUX although the file says sd3. Name components are now matched most-specific-first, a component naming two different backbones decides nothing rather than being resolved by precedence, and a local-path source is not evidence at all - the model manager sets source to the file's own path when there is no remote one, so trusting it means matching arbitrary ancestor directories of the user's model library. Nothing is lost: install_path identifies a local file before it moves it. Requiring the full contract also makes the latent channel count always readable, which retires the name-only backbone path entirely. The name can now only break the FLUX.1 / SD3 / Qwen-Image tie, never pick a backbone outright, and an explicit base override - already validated against one class's Literal - beats it. The checks that would rule out all five configs move out of _validate_base, leaving it to answer only "not this backbone". Also fixes, orthogonally: from_model_on_disk popped `variant` out of the override dict the factory builds once and shares across every candidate class, so the first PiD class to run consumed it and a later one that actually matched fell back to name inference. Verified against all 11 NVIDIA checkpoints: every one matches the contract exactly (missing=0, unexpected=0, no shape mismatch), the 9 supported decoders identify with the right base and variant both in place and as a direct single-file install, and the dinov2 / siglip decoders are rejected by latent channel count rather than registered as unknown models. * fix(pid): tolerate a checkpoint whose keys are not all strings A bare (un-prefixed) PidNet checkpoint is passed through strip_net_prefix untouched, on purpose: without the net. prefix there is no evidence the file is a distill serialisation, so a stray key must reach the unexpected-key checks rather than be dropped. A .pth unpickles to whatever it contains, so those keys need not all be strings - and reporting the unexpected ones sorts them. Sorting {1, "not_a_pid_key"} raises TypeError. That failure does not surface as a failure. ModelConfigFactory catches an unexpected exception from a candidate class as a generic no-match, so all five PiD configs drop out and allow_unknown_models registers the file as Unknown_Config - the exact fallback these checks exist to close. A complete bare contract plus one non-string key and one unexpected string key was therefore installed as an unknown model. Sort both key sets with key=str, and stop the type annotations claiming otherwise: strip_net_prefix and pid_net_shapes return dict[Any, ...], not dict[str, ...], and _Shapes follows. The old signature was not merely imprecise - it carried a type: ignore for the pass-through return, which is what let a str-only assumption look checked. Verified against the eleven real NVIDIA checkpoints: unchanged, all five supported backbones identify in place and as a direct single-file install, and the dinov2 / siglip decoders are still rejected by latent channel count. * fix(pid): reject non-string keys before torch trips over them The previous commit taught identification to tolerate a bare checkpoint whose keys are not all strings, and justified keeping those keys with a claim about the loader that is simply wrong: nn.Module.load_state_dict calls .startswith() on every key, so a non-string one raises AttributeError from inside torch before any unexpected key is reported. Passing a complete state dict plus {1: tensor} to load_pid_decoder raised that AttributeError rather than the RuntimeError the function reports every other unusable checkpoint with. load_pid_decoder now checks for non-string keys before it hands anything to torch, and says what is actually wrong with the file. Identification already rejects such a checkpoint, so this is the second line rather than the first - but load_pid_decoder is public, the model cache reaches it for records written before this PR, and a file can be swapped on disk after install. The reasoning in strip_net_prefix and its test is corrected to match what torch does. Keeping non-string keys is still right - dropping them would hide a malformed file from the checks meant to catch it - but the burden it puts on consumers is the opposite of what was written there: neither may assume the key type, so identification sorts its key reports with key=str and the loader rejects non-strings up front. Verified: the reviewer's repro now raises "PiD checkpoint has 1 keys that are not strings and so cannot name a PidNet parameter: [1]". The eleven real NVIDIA checkpoints are unaffected, 22/22 as before. --------- Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com> | 24 天前 | |
feat: Adding Support for SD.Next Quantization Engine (SDNQ) (Flux1&Flux2klein4B/9B&Z-Image) (#9228) * feat(quantization): add SDNQ (SD.Next Quantization) support Add support for loading SDNQ-quantized models with on-the-fly CPU dequantization, similar to existing GGUF support. New features: - SDNQTensor class with __torch_dispatch__ for automatic dequantization - Support for symmetric/asymmetric int8/uint8/fp8 quantization - Optional SVD correction (low-rank approximation) - Model loaders for Flux and Z-Image SDNQ models - Automatic format detection via weight+scale key pairs New files: - invokeai/backend/quantization/sdnq/ (core module) - tests/backend/quantization/sdnq/ (unit tests) Modified files: - taxonomy.py: Add ModelFormat.SDNQQuantized - configs/main.py: Add Main_SDNQ_FLUX_Config, Main_SDNQ_ZImage_Config - configs/factory.py: Register SDNQ configs - model_loaders/flux.py: Add FluxSDNQCheckpointModel - model_loaders/z_image.py: Add ZImageSDNQCheckpointModel * fix(sdnq): improve uint4 dequantization and add diffusers format support - Add uint4 per-group quantization with packed weight unpacking - Handle 1D flattened weights (reshape to 2D before unpacking) - Support SDNQ diffusers format for FLUX transformer and T5 - Add SDNQ VAE loading with AutoencoderKL - Add diagnostic logging for debugging dequantization - Fix bit order in uint4 unpacking (lower, upper) * test(sdnq): align asymmetric dequant expectation with upstream convention The test was checking `(weight - zero_point) * scale`, but SDNQ (Disty0/sdnq) defines asymmetric dequantization as `zero_point + weight * scale` (via torch.addcmul), where zero_point is a post-scale bias rather than a pre-scale integer offset. The implementation already follows this convention; only the test expectation was wrong. * test(sdnq): align asymmetric dequant expectation with upstream convention The test was checking `(weight - zero_point) * scale`, but SDNQ (Disty0/sdnq) defines asymmetric dequantization as `zero_point + weight * scale` (via torch.addcmul), where zero_point is a post-scale bias rather than a pre-scale integer offset. The implementation already follows this convention; only the test expectation was wrong. feat(sdnq): support sidecar LoRA application on SDNQ-quantized layers Bring SDNQ to feature parity with GGUF in the sidecar patching path so LoRA, LoKr, DoRA, FullLayer, and FluxControl patches apply correctly to SDNQ-quantized Linear and Conv2d modules. Without this, the sidecar aggregate replaced the SDNQTensor weight with a meta tensor and patches silently produced wrong results. - Add SDNQTensor branch in CustomModuleMixin._aggregate_patch_parameters mirroring the GGMLTensor branch. - Extend the (GGMLTensor) dtype-cast exclusion to also cover SDNQTensor in CustomLinear, CustomConv2d, CustomInvokeLinearNF4, and CustomInvokeLinear8bitLt. - Add `linear_with_sdnq_quantized_tensor` and `linear_sdnq_quantized` fixtures so the existing custom-module test matrix exercises SDNQ alongside GGUF, BnB-8bit, and NF4. * feat(sdnq): support SDNQ-quantized T5 encoders Add T5Encoder_SDNQ_Config for diffusers-style T5 bundles whose text_encoder_2/ folder holds SDNQ-quantized safetensors (detected via quantization_config.json's quant_method or via the SDNQ-style weight+scale key pairs). Add T5EncoderSDNQLoader that materializes the T5EncoderModel on meta, then loads the SDNQ state dict, and re-shares the embed_tokens/shared weight per HuggingFace's tied- weight convention. * feat(sdnq): support SDNQ-quantized FLUX.2 transformers Add Main_SDNQ_Flux2_Config covering Klein 4B/9B and their Base variants (detected via _get_flux2_variant on the dequantized SDNQTensor shapes plus the existing filename heuristic), and Flux2SDNQCheckpointModel that loads diffusers-layout SDNQ FLUX.2 checkpoints straight into Flux2Transformer2DModel. Architecture (num_layers, hidden_size, attention head count, guidance presence) is detected from state-dict shapes the same way the fp16 loader does, since SDNQTensor.shape reports the dequantized shape. BFL-layout SDNQ FLUX.2 checkpoints are not supported here — that would require an SDNQTensor-aware port of the _convert_flux2_bfl_to_diffusers fuse logic. * feat(sdnq): support full ZImagePipeline diffusers folders Add Main_SDNQ_Diffusers_ZImage_Config so a complete SDNQ ZImagePipeline folder (model_index.json + transformer/ + text_encoder/ + tokenizer/ + vae/) is recognised on install and its submodels are wired up. Extend ZImageSDNQCheckpointModel to load the transformer from the subfolder using ZImageTransformer2DModel.from_config() so non-default architecture parameters (e.g. axes_lens [1536,512,512] in newer Z-Image Turbo SDNQ exports) are honoured instead of the single-file path's hardcoded [1024,512,512]. Verified end-to-end against Tongyi-MAI/Z-Image-Turbo-SDNQ-uint4-svd-r32: 269 quantized + 252 regular tensors load into a 6.15B-param model with 0 missing / 0 unexpected keys. * fix(sdnq): match T5 SDNQ submodel layout in FluxPipeline bundles T5Encoder_SDNQ_Config originally only looked for text_encoder_2/ as a subfolder of mod.path, which works for standalone T5 bundles but misses the case where a parent FluxPipeline / similar config registers its T5 submodel with path_or_prefix pointing straight at the text_encoder_2 folder. Allow both layouts in both the config's detection logic and T5EncoderSDNQLoader's te_dir resolution. Verified end-to-end with Disty0/FLUX.1-schnell-SDNQ-uint4-svd-r32. * fix(sdnq): swap scale/shift halves in FLUX BFL converter's norm_out The diffusers→BFL state-dict converter renamed norm_out.linear.{weight,bias} to final_layer.adaLN_modulation.1.{weight,bias} but did not swap the two halves along dim 0. diffusers' AdaLayerNormContinuous packs the linear output as (scale, shift); BFL's LastLayer packs as (shift, scale). Without the swap, the final adaLN modulation runs with scale and shift permuted, which produces structured-but-very-noisy output for every pixel. Reuse the same pattern the FLUX.2 converter applies for the analogous adaLN_modulation key. * feat(sdnq): dispatch all ZImagePipeline submodels via SDNQ loader ZImageSDNQCheckpointModel only handled the Transformer submodel, so attempts to use an SDNQ ZImagePipeline as the "Qwen3 & VAE source model" (which triggers loads for TextEncoder / Tokenizer / VAE) crashed with "Only Transformer submodels are currently supported". Add per-submodel handlers that load text_encoder/ via sdnq_sd_loader into an empty Qwen3ForCausalLM (re-sharing lm_head with embed_tokens when tied), tokenizer/ via AutoTokenizer, and vae/ via AutoencoderKL.from_pretrained. The single-file SDNQ checkpoint path keeps its transformer-only behaviour but now raises a clearer error when asked for a different submodel. * feat(sdnq): support FLUX.2 Klein dynamic mixed-precision pipelines Add support for SDNQ-quantized Flux2KleinPipeline folders, which mix uint4 and int5 dtypes across layers (chosen dynamically by SDNQ during quantization to stay under a per-group loss budget). Core changes: - Add INT5_ASYM quantization type + unpack_uint5 + dequantize_int5_per_group. Sign-extension matches Disty0/sdnq's unpack_int convention (raw 0..31 - 16). zero_point is optional (dynamic-mixed sometimes emits scale-only int5 tensors). - _infer_quantization_type now takes a per_tensor_dtype override; the loader builds an inverted map from quantization_config.json's modules_dtype_dict. - _get_original_shape uses the packed weight size as the authoritative source for in_features, fixing a bug where Klein 4B's group_size=64 layers were misread as group_size=128 (the previous fallback). Pipeline integration: - Add Main_SDNQ_Diffusers_Flux2_Config matching Flux2Pipeline / Flux2KleinPipeline folders with quantized transformer. - Flux2SDNQCheckpointModel now dispatches all pipeline submodels: transformer (Flux2Transformer2DModel.from_config + sdnq state dict), text_encoder (Qwen3ForCausalLM SDNQ + lm_head/embed_tokens tie), tokenizer (AutoTokenizer), vae (AutoencoderKLFlux2 / AutoencoderKL). - Extend flux2_klein_model_loader._validate_diffusers_format and the isFlux2DiffusersMainModelConfig FE filter to also accept SDNQ pipeline configs (when submodels is populated). Verified against Disty0/FLUX.2-klein-4B-SDNQ-4bit-dynamic: 98 uint4 + 2 int5 tensors load into a 3.88B-param Flux2Transformer2DModel with 0 missing / 0 unexpected keys; both dequant paths produce reasonable zero-centred weight distributions. * - Reject SDNQ-quantized folders in Main_Diffusers_FLUX_Config and Main_Diffusers_Flux2_Config so identification routes them to the SDNQ configs instead. Without this both configs accept the folder and the plain diffusers loader wins, then crashes when reading packed uint8 weights as bf16. * - Merge multi-shard safetensors in sdnq_sd_loader so Klein 9B's diffusion_pytorch_model-{00001,00002}-of-00002.safetensors and FLUX.2 dev's sharded transformer both load. Detect cross-shard key collisions as a corruption signal. * - Treat SDNQ ZImagePipeline / Flux2KleinPipeline folders as "main_is_diffusers" in z_image_model_loader and flux2_klein_model_loader so the auto-extract-submodels branch handles them. Without this the loader demanded a separate VAE/Qwen3 source even though the SDNQ pipeline carries those submodels itself. - Drop the ui_model_format=Diffusers hint on Klein's qwen3_source_model field so the FE combobox can also show SDNQ pipeline configs (the FE filter already accepts them). * fix(sdnq): unblock FLUX.2 Klein SDNQ pipelines in the UI Loading the Klein 4B SDNQ pipeline as the main model errored with "No Qwen3 Encoder selected" in the UI even though the pipeline carries its own Qwen3 + VAE submodels, and the Model Manager showed no format badge at all on SDNQ models. - flux2_klein_model_loader now treats SDNQ-with-submodels as main_is_diffusers, so the auto-extract-submodels branch handles SDNQ pipelines exactly like plain diffusers. Drop the ui_model_format=Diffusers hint on qwen3_source_model so the combobox can also show SDNQ pipeline configs. - readiness.ts no longer demands a standalone VAE/Qwen3 for FLUX.2 Klein when the main model is itself a pipeline (diffusers or SDNQ-with-submodels). Without this the Invoke button stayed disabled with "Non-diffusers FLUX.2 Klein models require a standalone Qwen3 Encoder" even when the SDNQ pipeline could self-source everything. - Register sdnq_quantized in zModelFormat, the manually-edited OpenAPI schema, ModelFormatBadge, and MODEL_FORMAT_TO_LONG_NAME so SDNQ models render an "sdnq" badge instead of an empty placeholder. * feat(sdnq): add starter models and user-facing docs - 4 new starter models covering all SDNQ pipelines verified end-to-end in this branch: FLUX.1 schnell, Z-Image Turbo, FLUX.2 Klein 4B (dynamic mixed), FLUX.2 Klein 9B (dynamic mixed + SVD). Each entry is self-contained (no separate encoder/VAE dependencies because the SDNQ pipeline folder bundles them). - New /configuration/sdnq-quantization/ page: support matrix, VRAM footprints, install steps (Starter Models + HF + Folder), LoRA compatibility notes, SDNQ-vs-SVDQuant/Nunchaku disambiguation, comparison with GGUF/NF4/FP8, troubleshooting. - Cross-link from fp8-storage.mdx's "no-op on quantized" caution. * Chore Fix Path * fix(sdnq): add missing variant/cpu_only fields to SDNQ configs Z-Image and Qwen3 SDNQ configs were missing `variant` (and `cpu_only` on Qwen3) fields that exist on the other variants of the same union, breaking TypeScript narrowing on the FE. - Main_SDNQ_ZImage_Config: add variant (default Turbo) - Main_SDNQ_Diffusers_ZImage_Config: add variant, detect from scheduler_config.json shift value - Qwen3Encoder_SDNQ_Config: add cpu_only + variant, detect from embed_tokens shape - Qwen3Encoder_SDNQ_Folder_Config: add cpu_only + variant, detect from config.json hidden_size - Regenerate FE schema.ts Discriminator tags are unchanged since variant has no default. * Fix openapi schema.ts * Fix Path * feat(pid): vendor PiD decoder backend (phase A of integration) Adds a vendored subset of NVIDIA's PiD (Pixel Diffusion Decoder) at invokeai/backend/pid/ as the foundation for upcoming FLUX / FLUX.2 / SD3 / Z-Image PiD decode nodes plus a future PiD-based 4x upscale node. Upstream: https://github.com/nv-tlabs/PiD (Apache 2.0). Vendor scope: * _src/{networks,models,modules}: PidNet, PixDiT_T2I, LQProjection2D, PidModel, PidDistillModel, PixelDiTModel, GeneralConditioner. * _ext/imaginaire: minimal Imaginaire framework subset (lazy_config, model, utils/{log,misc,distributed,device,count_params}). * configs/, tokenizers/, checkpointer/, trainer.py, visualize/, _demo_*, from_*, easy_io/, S3/wandb training helpers were intentionally excluded. Dependency stripping (no new hard deps introduced): * loguru, termcolor -> stdlib logging shim * iopath PathManager -> stdlib pathlib stub * fvcore Registry -> minimal stdlib Registry * lazy_config/lazy.py: yaml/dill/cloudpickle/detectron2 save/load paths replaced with a minimal LazyCall stub * lazy_config/instantiate.py: omegaconf DictConfig/ListConfig branches removed; configs are plain dict / LazyCall mappings * megatron, pynvml, boto3/wandb imports are try/except-guarded or local to functions and stay inert in our inference path All pid.* imports rewritten to invokeai.backend.pid.*; SPDX-Apache-2.0 headers retained on vendored files; attribution and detailed list of local modifications added in LICENSE-PiD.txt. The pre-trained PiD checkpoints distributed by NVIDIA remain under NSCLv1 (non-commercial); this commit only vendors code. Smoke test: PidNet, PidModel, PidDistillModel, GeneralConditioner import cleanly; LazyCall -> instantiate round-trip resolves to the expected nn.Module. ruff check passes. * feat(pid): wire PiD + Gemma-2 into model-manager and add decode nodes Adds the model-manager plumbing and workflow nodes needed to use the vendored PiD decoder (phase A) end-to-end with FLUX, SD3 and Z-Image. Model manager (Phase B + B.5): * taxonomy: ModelType.PiDDecoder, PiDDecoderVariantType (Res2k_Sr4x / Res2kTo4k_Sr4x), ModelType.Gemma2Encoder + ModelFormat.Gemma2Encoder, both added to AnyVariant + variant_type_adapter. * configs/pid_decoder.py: per-backbone PiD configs (FLUX / FLUX.2 / SD3) with state-dict probing on 'lq_proj' substring and backbone/variant detection from the official NVIDIA filenames. * configs/gemma2_encoder.py: Gemma-2 directory probing on Gemma2ForCausalLM architecture + tokenizer files. * AnyModelConfig union updated. * model_loaders/pid_decoder.py: loads .pth / .safetensors, strips the upstream 'net.' prefix, supports torch.load(weights_only=True). * model_loaders/gemma2_encoder.py: SubModelType.{Tokenizer, TextEncoder} dispatch; returns the causal LM's inner Gemma2Model (transformers 4.56's get_decoder() returns None for Gemma2). Decode pipeline (Phase C): * backend/pid/decode.py: build_pid_net + load_pid_decoder (per-backbone PixDiT_T2I hyperparams derived from PiD's pid_sr4x base + per-experiment overrides), encode_caption_for_pid (chi-prompt + Gemma encoding, mirrors PixelDiTModel._encode_text_raw), and a PiDDecoder wrapper with a reimplemented few-step distill sampler (no autocast / no distributed / no PixelDiTModel init paths from upstream). Invocations (Phase 6.x): * Gemma2EncoderField + PiDDecoderField in invocations/model.py. * gemma2_encoder_loader / pid_decoder_loader: thin ModelIdentifierField pickers that emit the corresponding fields. * z_image_pid_decode (pilot), flux_pid_decode, sd3_pid_decode: caption encode -> Gemma offload -> PiD state dict load -> PidNet construct -> decode. Per-backbone latent denormalisation (FLUX1 ae_params, SD3 hardcoded 1.5305/0.0609, Z-Image piggybacks on FLUX VAE). End-to-end validated with the released PiD_res2k_sr4x_official_flux_distill_4step.pth checkpoint and gemma-2-2b-it: PidNet rebuilds at exactly 456 keys / 1.36B params, sampler runs at ~5 GB VRAM peak (Gemma dominates), output shape and range match. FLUX.2 PiD decode is deliberately deferred: it needs BN-based latent denormalisation and 32->128 channel packing, and we have no FLUX.2 checkpoint to validate against yet. * feat(pid): end-to-end PiD pixel-diffusion decoder integration Adds the NVIDIA PiD decoder as a 4x super-resolution alternative to the regular VAE/RAE decode path. Includes model-manager configs and loaders for both the PiD checkpoints and the Gemma-2 caption encoder they require, plus four invocations: latent-in decode for FLUX / SD3 / Z-Image and an image-in pid_upscale node. - Decode pipeline keeps PidNet params in fp32 and uses bf16 autocast only for matmuls; caption embeddings have outliers that overflow bf16 RMSNorm. - encode_caption_for_pid forces tokenizer padding_side="right" (Gemma defaults to left, PiD trained with right) and returns the attention mask as bool so it stays compatible with SDPA. - Z-Image reuses the FLUX-trained checkpoint and reads scale/shift from the VAE config at runtime (PiD upstream notes they are checkpoint-specific). - TextLLM config now excludes Gemma2ForCausalLM so it falls through to the dedicated Gemma2 encoder config instead of being misclassified. - Frontend: new model_type / model_format / variant enums, type guards and category metadata; schema.ts regenerated via pnpm typegen. * Fix closing Step Tag * Chore Ruff * Chore Typegen * Chore Knip * fix(pid): remove unused vendored models/utils.py (broken easy_io import) * feat(pid): identify decoder backbone from weight shapes, not filename Read latent channel count from lq_proj.latent_proj.0.weight (FLUX.2=128, FLUX.1/SD3=16) as the primary discriminator; fall back to filename/dir name only to disambiguate the architecturally identical FLUX.1/SD3 pair. Fixes FLUX.2 checkpoints (model_ema_bf16.pth) not being recognised, and correctly rejects unsupported backbones (RAE/dinov2, 768ch). Fix Flux2 docstring 32->128. * feat(ui): PiD decode (Fit mode) for FLUX text-to-image Add a "PiD Decode" mode select (Off / Fit / Native) to the FLUX advanced settings with PiD decoder + Gemma-2 encoder pickers. In Fit mode the FLUX graph swaps the VAE decode for a PiD 4x super-resolution decode and downscales back to the requested size. Adds params state (pidMode, decoder, encoder, steps) with a v3->v4 migration, model hooks, readiness checks, and graph guards for the not-yet-wired Native and non-txt2img paths. * feat(ui): PiD Native 4x mode for FLUX text-to-image Make the generation dimension helpers PiD-aware via an optional pidScale: in Native mode the user-facing dimensions are the 4x target (grid 64, optimal 2048), generation runs at target/4, and PiD's 4x output is used directly with no downscale. Thread pidScale through the params dimension reducers and the optimal-dimension/grid-size selectors, resync dimensions when toggling Native, and wire the Native path in the FLUX graph builder. Add working_mem_bytes for PiD Decode * feat(ui): PiD Fit decode for FLUX image-to-image Extract the PiD decode chain into buildPidDecodeChain (loaders + decode + fit-downscale, no denoise setup) so it can substitute for the VAE decode across generation modes. Widen addImageToImage's l2i param to ImageOutputNodes (it only consumes .image) and wire the PiD chain into the img2img branch in Fit mode. Native stays txt2img-only (a 4x result can't composite onto the bbox); inpaint/outpaint remain gated off for now. * feat(ui): PiD Native 4x decode for FLUX image-to-image Add addPidImageToImageNative: the canvas bbox is the 4x target, so the init image is downscaled to bbox/4, denoised at that resolution, and PiD decodes straight back up to the full bbox with no post-decode downscale - preserving all PiD detail while still compositing cleanly onto the region. Wire it into the img2img branch of buildFLUXGraph (native vs fit vs off) and drop the native-txt2img-only guard. Make the canvas FLUX grid check PiD-aware so a native bbox must be a multiple of 64 (16 * 4) for bbox/4 to land on the grid. * feat(ui): add informational popover to PiD Decode setting Explain PiD usage on hover, mirroring the DyPE popover: what the decoder is (NVIDIA Pixel Diffusion Decoder, 4x SR, needs a PiD decoder + Gemma-2 encoder), Fit vs Native modes, the 2K / 2K-to-4K target resolutions, that Steps can be lowered, and that Scale Before Processing must be off. Links to nv-tlabs/PiD. * feat(models): add PiD decoder + Gemma-2 encoder to starter models Register NVIDIA's PiD FLUX decoders (2K and 2K-to-4K presets, from nvidia/PiD) and the Efficient-Large-Model/gemma-2-2b-it caption encoder as starter models so they can be installed from the Model Manager. The Gemma-2 encoder is wired as a dependency of each decoder (and offered standalone). * feat(pid): add FLUX.2 Klein PiD 4x-SR decode support Add a flux2_pid_decode node that packs the stored FLUX.2 latent (32ch @ H/8) into PiD's 128ch @ H/16 layout before decoding; FLUX.2's BatchNorm denormalization is already applied in flux2_denoise, so no scalar denorm is needed (optional vae input reads identity constants). Generalize the frontend PiD decode chain (decodeNodeType, optional vaeSource) and wire the isFlux2 graph path for txt2img/img2img (Fit & Native). Base-aware PiD gating/decoder-filter, FLUX.2 readiness checks, and two nvidia/PiD FLUX.2 starter decoders (2K, 2Kto4K). Standard FLUX PiD path unchanged. * feat(pid): add SD3 PiD 4x-SR decode support Wire the existing sd3_pid_decode node into the SD3 graph builder (txt2img and img2img, Fit & Native) with a PiD guard, base-aware gating/decoder-filter (sd-3), and SD3 readiness checks. Add two nvidia/PiD SD3 starter decoders (2K, 2Kto4K). Harden the PiD config probe against the 16-channel FLUX.1/SD3 ambiguity: when the checkpoint's directory name is silent (the HF single-file download renames it), trust an explicit base override so SD3 checkpoints are not misidentified as FLUX.1. Also benefits Qwen. FLUX / FLUX.2 identification is unchanged. * feat(pid): add SDXL PiD 4x-SR decode support Build the full SDXL PiD backend stack: _PER_BACKBONE[SDXL] (4ch/down8), PiDDecoder_Checkpoint_SDXL_Config with a 4-channel latent-map entry, factory union + loader registration, and a new sdxl_pid_decode node (reads the VAE's scaling_factor/shift at runtime; SDXL fallbacks 0.13025/0.0). 4-channel latents are unambiguous, so no directory-name disambiguation is needed. Generalize the shared PiD decode chain to support SD-family denoise: denoise_latents has no width/height, so thread an optional noise node for sizing and round to the model's native grid (8 for SDXL, 16 for FLUX). Wire buildSDXLGraph (txt2img + img2img, Fit & Native) with the VAE as the decode's scaling source, base-aware gating/readiness, and a starter decoder (SDXL 2Kto4K only). PiD + SDXL refiner is blocked for now via a graph guard and a readiness reason. FLUX/FLUX.2/SD3 paths are unchanged. * feat(pid): add Z-Image PiD 4x-SR decode support Wire the existing z_image_pid_decode node into the Z-Image graph builder (txt2img and img2img, Fit & Native) with a PiD guard and readiness checks. Z-Image shares FLUX.1's 16-channel VAE and has no PiD checkpoints of its own, so it reuses the FLUX decoder: the decoder filter maps z-image -> flux, showing FLUX PiD decoders when a Z-Image model is active. The Z-Image VAE is passed to the decode node so it reads the real scaling_factor / shift instead of the fallback constants. No backend, schema, or starter-model changes. FLUX/FLUX.2/SD3/SDXL paths are unchanged. * feat(pid): add Qwen-Image PiD 4x-SR decode support Build the full Qwen-Image PiD backend stack: _PER_BACKBONE[QwenImage] (16ch/down8), PiDDecoder_Checkpoint_QwenImage_Config (added to the 16-channel latent map + filename heuristic), factory union + loader registration, and a new qwen_image_pid_decode node. Unlike the scalar-scaling bases, the Qwen-Image VAE normalizes per channel (latents_mean / latents_std) and stores a 5D video-style latent, so the node denormalizes per-channel (z * std + mean, read from the VAE config) and drops the singleton temporal frame before decoding - matching qwen_image_l2i. Wire buildQwenImageGraph (txt2img + img2img, Fit & Native) with the Qwen-Image VAE as the decode's normalization source, base-aware gating/readiness, and a starter decoder (Qwen-Image 2Kto4K only). The 16-channel FLUX/SD3/Qwen ambiguity is handled by the existing trusted- base-override probe hardening. FLUX/FLUX.2/SD3/SDXL/Z-Image paths are unchanged. * Chore Ruff * Add Docs * fix(pid): green up frontend tests and knip for the PiD branch - graph-builder tests: set pidMode 'off' in the FLUX / Qwen-Image / SDXL+SD3 param fixtures so the PiD guard doesn't fire on an undefined pidMode and call the (unmocked) size helpers - paramsSlice migration test: expect _version 4 (v3→v4 adds the PiD fields) - remove the unused setPidSteps action and selectPidSteps selector flagged by knip; the pidSteps state field stays at its default of 4 * Chore openapi + typegen * Merge remote-tracking branch 'origin/main' into feature/svd-quantization Resolve conflicts in qwen3_encoder.py (keep both SDNQ tensor/key detection and the Qwen-VL visual-tower guard) and regenerate openapi.json + schema.ts from the merged backend. Fix T5TokenizerFast under transformers 5.x: it is now an alias of T5Tokenizer. Use T5Tokenizer in flux.py (was used unimported -> F821) and also match the new "T5Tokenizer" class name in the SDNQ FLUX model_index.json submodel probe. * Fix SDNQ mixed-group dequant, Z-Image self-contained pipelines, and FLUX.2 single-file identification Address three SDNQ review findings (PR #9228): - loaders.py: derive per-tensor group_size from the scale tensor's group dimension instead of trusting quantization_config.json. Dynamic-mixed models (e.g. FLUX.2 Klein 4B) quantize some layers with a group size that differs from the nominal config value; trusting the config produced a num_groups that disagreed with scale.shape[1] and broke the reshape/ broadcast in dequantize_uint4/int5_per_group. - z_image_model_loader: fall back to the main model for the VAE/Qwen3 submodels when it is a self-contained SDNQ pipeline (format=sdnq_quantized with submodels). Mirror the same self-contained handling in the frontend readiness checks and buildZImageGraph, so a freshly installed SDNQ Z-Image pipeline generates without manually selecting a component source. - FLUX.2 single-file identification: list SDNQ FLUX.2 before FLUX.1 in the config union, reject FLUX.2 state dicts in Main_SDNQ_FLUX_Config so the two stay mutually exclusive regardless of iteration order, and accept the bare diffusers transformer layout (transformer_blocks./context_embedder.) in Main_SDNQ_Flux2_Config so the real single-file SDNQ FLUX.2 checkpoint is identified instead of falling through to FLUX.1 or unknown. Add regression tests for all three: mixed-group uint4 dequant, the Z-Image loader/readiness self-contained fallback, and SDNQ FLUX.2 identification. * Chore typegen + openapi * Fix SDNQ mixed-group dequant, Z-Image self-contained pipelines, and FLUX.2/Z-Image identification Address five SDNQ review findings (PR #9228): - loaders.py: derive per-tensor group_size from the scale tensor's group dimension instead of trusting quantization_config.json. Dynamic-mixed models (e.g. FLUX.2 Klein 4B) quantize some layers with a group size that differs from the nominal config value; trusting the config produced a num_groups that disagreed with scale.shape[1] and broke the reshape/ broadcast in dequantize_uint4/int5_per_group. - z_image_model_loader: fall back to the main model for the VAE/Qwen3 submodels when it is a self-contained SDNQ pipeline (format=sdnq_quantized with submodels). Mirror the same self-contained handling in the frontend readiness checks and buildZImageGraph, so a freshly installed SDNQ Z-Image pipeline generates without manually selecting a component source. - FLUX.2 single-file identification: list SDNQ FLUX.2 before FLUX.1 in the config union, reject FLUX.2 state dicts in Main_SDNQ_FLUX_Config so the two stay mutually exclusive regardless of iteration order, and accept the bare diffusers transformer layout (transformer_blocks./context_embedder.) in Main_SDNQ_Flux2_Config so the real single-file SDNQ FLUX.2 checkpoint is identified instead of falling through to FLUX.1 or unknown. - Main_Diffusers_ZImage_Config: reject SDNQ-quantized ZImagePipeline folders so Main_SDNQ_Diffusers_ZImage_Config matches them instead (mirrors the FLUX and FLUX.2 guards). Without this a full SDNQ Z-Image pipeline could be classified as plain diffusers, mis-reading packed weights and breaking the self-contained path. - z_image loader: check missing/unexpected keys after load_state_dict in the SDNQ transformer folder path and fail fast with the offending key list, instead of silently returning a model with required params on meta tensors that fails later during device movement or inference. Add regression tests for all five: mixed-group uint4 dequant, the Z-Image loader/readiness self-contained fallback, SDNQ FLUX.2 identification, and SDNQ ZImagePipeline-folder identification. * Harden SDNQ model identification, loading, and Z-Image pipeline support Address the SDNQ review findings on the feature branch: Dequantization / loading correctness - sdnq/loaders.py: derive per-tensor group_size from the scale tensor's group dimension instead of trusting quantization_config.json. Dynamic-mixed models (e.g. FLUX.2 Klein 4B) quantize some layers with a group size that differs from the nominal config value; trusting the config produced a num_groups that disagreed with scale.shape[1] and broke the reshape/broadcast in dequantize_uint4/int5_per_group. - Add shared raise_on_incomplete_sdnq_load() and apply it to every SDNQ folder loader that used load_state_dict(strict=False) without checking the result: FLUX.2 pipeline transformer, SDNQ CLIP/T5/VAE (flux.py), standalone SDNQ VAE (vae.py), and both Z-Image Qwen3 text-encoder paths (z_image.py). A partial or mismatched export now fails fast with the offending key list instead of returning a model with parameters left on the meta device. Tied/re-shared weights (T5 encoder.embed_tokens, Qwen3 lm_head, CLIP position_ids) are allow-listed. Model identification - Reject FLUX.2 in Main_SDNQ_FLUX_Config and accept the bare diffusers transformer layout (transformer_blocks./context_embedder.) in Main_SDNQ_Flux2_Config, and list FLUX.2 before FLUX.1 in the config union, so a single-file SDNQ FLUX.2 checkpoint identifies as FLUX.2 instead of falling through to FLUX.1 or unknown. - Main_Diffusers_ZImage_Config: reject SDNQ-quantized ZImagePipeline folders so Main_SDNQ_Diffusers_ZImage_Config matches them instead (mirrors the FLUX and FLUX.2 guards). - Qwen3Encoder_SDNQ_Folder_Config: verify the folder is actually a Qwen3 encoder (Qwen3 config class name or Qwen3 state-dict keys) before claiming it, so an SDNQ transformer/VAE folder is not stored as type=qwen3_encoder. - Qwen3Encoder_Qwen3Encoder_Config: reject SDNQ folders so it no longer competes with the SDNQ folder config (both share the Qwen3Encoder type, making the factory tiebreak non-deterministic). Z-Image self-contained SDNQ pipelines - z_image_model_loader: fall back to the main model for the VAE/Qwen3 submodels when it is a self-contained SDNQ pipeline (format=sdnq_quantized with submodels), and drop the ui_model_format=diffusers hint from qwen3_source_model so SDNQ pipelines are not filtered out of the node/workflow model pickers. - Mirror the self-contained handling in the frontend readiness checks and buildZImageGraph so a freshly installed SDNQ Z-Image pipeline generates without a manually selected component source. Add regression tests for every fix: mixed-group uint4 dequant, the load-guard helper, SDNQ FLUX.2 identification, SDNQ ZImagePipeline-folder identification, Qwen3 encoder folder identification, the Z-Image self-contained loader/readiness path, and the qwen3_source_model field template. * Chore Typegen + OpenApi * Harden SDNQ model identification, loading, and Z-Image pipeline support Address the SDNQ review findings on the feature branch: Dequantization / loading correctness - sdnq/loaders.py: derive per-tensor group_size from the scale tensor's group dimension instead of trusting quantization_config.json. Dynamic-mixed models (e.g. FLUX.2 Klein 4B) quantize some layers with a group size that differs from the nominal config value; trusting the config produced a num_groups that disagreed with scale.shape[1] and broke the reshape/broadcast in dequantize_uint4/int5_per_group. - Add shared raise_on_incomplete_sdnq_load() and use it in every SDNQ folder loader that called load_state_dict(strict=False) without checking the result: FLUX.2 pipeline transformer, SDNQ CLIP/T5/VAE (flux.py), standalone SDNQ VAE (vae.py), and both Z-Image Qwen3 text-encoder paths (z_image.py). A partial or mismatched export now fails fast with the offending key list instead of returning a model with parameters left on the meta device. Tied/re-shared weights (T5 encoder.embed_tokens, Qwen3 lm_head, CLIP position_ids) are allow-listed. Replace the T5 SDNQ assert checks (stripped by python -O) and add the missing unexpected-key check to the standalone Qwen3 SDNQ loader. Model identification - Reject FLUX.2 in Main_SDNQ_FLUX_Config and accept the bare diffusers transformer layout (transformer_blocks./context_embedder.) in Main_SDNQ_Flux2_Config, and list FLUX.2 before FLUX.1 in the config union, so a single-file SDNQ FLUX.2 checkpoint identifies as FLUX.2 instead of falling through to FLUX.1 or unknown. - Main_Diffusers_ZImage_Config: reject SDNQ-quantized ZImagePipeline folders so Main_SDNQ_Diffusers_ZImage_Config matches them instead (mirrors the FLUX/FLUX.2 guards). - Qwen3Encoder_SDNQ_Folder_Config: verify the folder is actually a Qwen3 encoder (Qwen3 config class or Qwen3 state-dict keys) and reject complete causal LMs (root config.json + tokenizer files) so an SDNQ transformer/VAE/TextLLM folder is not stored as type=qwen3_encoder. - Qwen3Encoder_Qwen3Encoder_Config: reject SDNQ folders so it no longer competes with the SDNQ folder config (both share the Qwen3Encoder type, making the factory tiebreak non-deterministic). Z-Image self-contained SDNQ pipelines - z_image_model_loader: fall back to the main model for the VAE/Qwen3 submodels when it is a self-contained SDNQ pipeline, but only when the pipeline actually exposes the required vae + text_encoder + tokenizer submodels (a truthy submodels dict is not enough — a partial pipeline would fail later on missing folders). Drop the ui_model_format=diffusers hint from qwen3_source_model so SDNQ pipelines are not filtered out of the node/workflow model pickers. - Mirror both the self-contained handling and the specific-submodels requirement in the frontend readiness checks and buildZImageGraph, so a freshly installed SDNQ Z-Image pipeline generates without a manually selected component source. Add regression tests for every fix: mixed-group uint4 dequant, the load-guard helper, SDNQ FLUX.2 identification, SDNQ ZImagePipeline-folder identification, Qwen3 encoder folder identification (SDNQ rejection, causal-LM rejection), the Z-Image self-contained loader/readiness path incl. partial-pipeline handling, and the qwen3_source_model field template. * Chore Ruff * Harden SDNQ model identification, loading, and Z-Image / FLUX.2 pipeline support Address the SDNQ review findings on the feature branch: Dequantization / loading correctness - sdnq/loaders.py: derive per-tensor group_size from the scale tensor's group dimension instead of trusting quantization_config.json. Dynamic-mixed models (e.g. FLUX.2 Klein 4B) quantize some layers with a group size that differs from the nominal config value; trusting the config produced a num_groups that disagreed with scale.shape[1] and broke the reshape/broadcast in dequantize_uint4/int5_per_group. - Add shared raise_on_incomplete_sdnq_load() and use it in every SDNQ folder loader that called load_state_dict(strict=False) without checking the result: FLUX.2 pipeline transformer, SDNQ CLIP/T5/VAE (flux.py), standalone SDNQ VAE (vae.py), and both Z-Image Qwen3 text-encoder paths (z_image.py). A partial or mismatched export now fails fast with the offending key list instead of returning a model with parameters left on the meta device. Tied/re-shared weights (T5 encoder.embed_tokens, Qwen3 lm_head, CLIP position_ids) are allow-listed. Replace the T5 SDNQ assert checks (stripped by python -O) and add the missing unexpected-key check to the standalone Qwen3 SDNQ loader. Model identification - Reject FLUX.2 in Main_SDNQ_FLUX_Config and accept the bare diffusers transformer layout (transformer_blocks./context_embedder.) in Main_SDNQ_Flux2_Config, and list FLUX.2 before FLUX.1 in the config union, so a single-file SDNQ FLUX.2 checkpoint identifies as FLUX.2 instead of falling through to FLUX.1 or unknown. - Main_Diffusers_ZImage_Config: reject SDNQ-quantized ZImagePipeline folders so Main_SDNQ_Diffusers_ZImage_Config matches them instead (mirrors the FLUX/FLUX.2 guards). - Qwen3Encoder_SDNQ_Folder_Config: verify the folder is actually a Qwen3 encoder, reject complete causal LMs (root config.json + tokenizer files), accept the same compatible Qwen architectures as the unquantized config (Qwen2VLForConditionalGeneration / Qwen2ForCausalLM / Qwen3ForCausalLM via a shared constant), and make the state-dict fallback resilient to sharded folders, so an SDNQ transformer/VAE/TextLLM folder is not stored as type=qwen3_encoder and a compatible sharded encoder is not wrongly rejected. - Qwen3Encoder_Qwen3Encoder_Config: reject SDNQ folders so it no longer competes with the SDNQ folder config (both share the Qwen3Encoder type, making the factory tiebreak non-deterministic). Z-Image / FLUX.2 self-contained SDNQ pipelines - Add a shared is_self_contained_sdnq_pipeline() helper requiring the specific VAE + Qwen3 (text_encoder + tokenizer) submodels a pipeline install must ship. A truthy submodels dict is not enough: Main_SDNQ_Diffusers_* configs record whichever submodels they recognize, so a partial pipeline can expose only the transformer and would fail at runtime on missing fixed subfolders. - z_image_model_loader and flux2_klein_model_loader: use the helper for both the main-model self-contained fallback and the explicit qwen3_source validator, so a partial pipeline (as main model or selected source) requires an explicit VAE/Qwen3 source instead of emitting requests against missing folders. - Drop the ui_model_format=diffusers hint from z_image qwen3_source_model so SDNQ pipelines are not filtered out of the node/workflow model pickers. - Mirror the same specific-submodels requirement in the frontend readiness checks (Z-Image + FLUX.2, generate and canvas tabs) and buildZImageGraph, and bring the FLUX.2 canvas readiness block to parity with the generate tab (it previously used format !== 'diffusers' and never recognized SDNQ pipelines). Add regression tests for every fix: mixed-group uint4 dequant, the load-guard helper, SDNQ FLUX.2 identification, SDNQ ZImagePipeline-folder identification, Qwen3 encoder folder identification (SDNQ rejection, causal-LM rejection, compatible Qwen2 class names, sharded folders), the Z-Image and FLUX.2 self-contained loader/readiness paths incl. partial-pipeline handling, and the qwen3_source_model field template. * Harden SDNQ model identification, loading, and Z-Image / FLUX.2 pipeline support Address the SDNQ review findings on the feature branch: Dequantization / loading correctness - sdnq/loaders.py: derive per-tensor group_size from the scale tensor's group dimension instead of trusting quantization_config.json. Dynamic-mixed models (e.g. FLUX.2 Klein 4B) quantize some layers with a group size that differs from the nominal config value; trusting the config produced a num_groups that disagreed with scale.shape[1] and broke the reshape/broadcast in dequantize_uint4/int5_per_group. - Add shared raise_on_incomplete_sdnq_load() and use it in every SDNQ folder loader that called load_state_dict(strict=False) without checking the result: FLUX.2 pipeline transformer, SDNQ CLIP/T5/VAE (flux.py), standalone SDNQ VAE (vae.py), and both Z-Image Qwen3 text-encoder paths (z_image.py). A partial or mismatched export now fails fast with the offending key list instead of returning a model with parameters left on the meta device. Tied/re-shared weights (T5 encoder.embed_tokens, Qwen3 lm_head, CLIP position_ids) are allow-listed. Replace the T5 SDNQ assert checks (stripped by python -O) and add the missing unexpected-key check to the standalone Qwen3 SDNQ loader. Model identification - Reject FLUX.2 in Main_SDNQ_FLUX_Config and accept the bare diffusers transformer layout (transformer_blocks./context_embedder.) in Main_SDNQ_Flux2_Config, and list FLUX.2 before FLUX.1 in the config union, so a single-file SDNQ FLUX.2 checkpoint identifies as FLUX.2 instead of falling through to FLUX.1 or unknown. - Main_Diffusers_ZImage_Config: reject SDNQ-quantized ZImagePipeline folders so Main_SDNQ_Diffusers_ZImage_Config matches them instead (mirrors the FLUX/FLUX.2 guards). - Qwen3Encoder_SDNQ_Folder_Config: verify the folder is actually a Qwen3 encoder, reject complete causal LMs (root config.json + tokenizer files), accept the same compatible Qwen architectures as the unquantized config (Qwen2VLForConditionalGeneration / Qwen2ForCausalLM / Qwen3ForCausalLM via a shared _QWEN3_ENCODER_ARCHITECTURES constant), and make the state-dict fallback resilient to sharded folders — so an SDNQ transformer/VAE/TextLLM folder is not stored as type=qwen3_encoder and a compatible sharded encoder is not wrongly rejected. - Qwen3Encoder_Qwen3Encoder_Config: reject SDNQ folders so it no longer competes with the SDNQ folder config (both share the Qwen3Encoder type, making the factory tiebreak non-deterministic). SDNQ pipeline submodel discovery - Main_SDNQ_Diffusers_Flux2_Config / Main_SDNQ_Diffusers_ZImage_Config _get_submodels(): record the TextEncoder for any compatible encoder class (_QWEN3_ENCODER_ARCHITECTURES) and the Tokenizer for the slow and fast Qwen2 tokenizer classes, not just Qwen3ForCausalLM / Qwen2Tokenizer. Otherwise a valid pipeline advertising e.g. Qwen2ForCausalLM or Qwen2TokenizerFast was recorded as partial and forced to use separate VAE/Qwen3 sources. Z-Image / FLUX.2 self-contained SDNQ pipelines - Add a shared is_self_contained_sdnq_pipeline() helper requiring the specific VAE + Qwen3 (text_encoder + tokenizer) submodels a pipeline install must ship. A truthy submodels dict is not enough: a partial (or partially recognized) pipeline can expose only the transformer and would fail at runtime on missing fixed subfolders. - z_image_model_loader and flux2_klein_model_loader: use the helper for both the main-model self-contained fallback and the explicit qwen3_source validator, so a partial pipeline (as main model or selected source) requires an explicit VAE/Qwen3 source instead of emitting requests against missing folders. - Drop the ui_model_format=diffusers hint from z_image qwen3_source_model so SDNQ pipelines are not filtered out of the node/workflow model pickers. - Mirror the same specific-submodels requirement in the frontend readiness checks (Z-Image + FLUX.2, generate and canvas tabs) and buildZImageGraph, and bring the FLUX.2 canvas readiness block to parity with the generate tab (it previously used format !== 'diffusers' and never recognized SDNQ pipelines). Add regression tests for every fix: mixed-group uint4 dequant, the load-guard helper, SDNQ FLUX.2 identification, SDNQ ZImagePipeline-folder identification, Qwen3 encoder folder identification (SDNQ rejection, causal-LM rejection, compatible Qwen2 class names, sharded folders), SDNQ pipeline submodel discovery across compatible encoder/tokenizer classes, the Z-Image and FLUX.2 self-contained loader/readiness paths incl. partial-pipeline handling, and the qwen3_source_model field template. * Exclude Qwen2-VL encoder from SDNQ pipeline self-contained discovery The SDNQ FLUX.2 / Z-Image pipeline loaders instantiate a text-only Qwen3ForCausalLM for the discovered text_encoder/ folder, so they cannot load a Qwen2VLForConditionalGeneration model (multimodal, visual tower). Recording it as a self-contained TextEncoder marked the pipeline complete even though the loader would fail on the visual-tower weights, causing readiness/invocation to select the main model as the Qwen source. Narrow _get_submodels()'s accepted TextEncoder classes to the text-only Qwen causal-LM classes, excluding the Qwen-VL class. * Harden SDNQ model identification, loading, dequantization, and cache accounting Address the SDNQ review findings on the feature branch: Dequantization / tensor correctness - sdnq/loaders.py: derive per-tensor group_size from the scale tensor's group dimension instead of trusting quantization_config.json, so dynamic-mixed models (e.g. FLUX.2 Klein 4B) whose per-layer group size differs from the config value no longer break the reshape/broadcast in dequantize_uint4/int5_per_group. - sdnq/utils.py: normalize 2D [out_features, num_groups] scale/zero_point with a trailing singleton before per-group arithmetic; a 2D param previously right-aligned against the 3D grouped weight and failed broadcasting. - sdnq/sdnq_tensor.py: when an SDNQTensor is moved with .to(device) (aten _to_copy), move its scale/zero_point/svd payloads too. Previously only quantized_data moved, leaving auxiliary tensors in system RAM and forcing a host->device copy of all of them on every dequantization of every layer, every step. - Add SDNQTensor.sdnq_storage_nbytes() and make calc_tensor_size() SDNQ-aware (duck-typed) so cache eviction / partial-load budgets use real storage (packed data + scale + zero_point + svd) instead of the wrapper's dequantized shape × uint8, which over-counted packed weights and omitted the auxiliary tensors. Loader robustness - Add shared raise_on_incomplete_sdnq_load() and use it in every SDNQ folder loader that called load_state_dict(strict=False) without checking the result (FLUX.2 pipeline transformer, SDNQ CLIP/T5/VAE, standalone SDNQ VAE, both Z-Image Qwen3 text-encoder paths). Partial/mismatched exports now fail fast with the offending keys. Tied/re-shared weights are allow-listed; replace the T5 assert checks (stripped by python -O) and add the missing unexpected-key check to the standalone Qwen3 SDNQ loader. - vae.py: SDNQ VAE detection now inspects all safetensors shards (union of keys) and the quantization_config.json marker, and _load_sdnq_vae hands sharded directories to sdnq_sd_loader — a VAE split across standard shard files is detected and loaded instead of falling through to the generic diffusers loader. Model identification - Single-file SDNQ FLUX.2: accept only the bare diffusers layout (transformer_blocks./context_embedder.) that the loader supports and reject BFL / ComfyUI (model.diffusion_model.double_blocks.*) checkpoints, which the loader cannot convert. Reject FLUX.2 in Main_SDNQ_FLUX_Config and list FLUX.2 before FLUX.1 in the config union so an ambiguous checkpoint never reaches the FLUX.1 loader. - Main_Diffusers_ZImage_Config: reject SDNQ-quantized ZImagePipeline folders so the SDNQ config matches them instead (mirrors the FLUX/FLUX.2 guards). - Qwen3Encoder_Qwen3Encoder_Config: reject SDNQ folders so it no longer competes (non-deterministically) with the SDNQ folder config. - Qwen3Encoder_SDNQ_Folder_Config: verify the folder is a Qwen3 encoder, reject complete causal LMs (root config.json + tokenizer files), and accept only the architecture the SDNQ loader can build (Qwen3ForCausalLM) — rejecting a declared Qwen2ForCausalLM / Qwen2VLForConditionalGeneration folder rather than accepting one the loader's strict guard would fail on. SDNQ pipeline submodel discovery - Main_SDNQ_Diffusers_Flux2_Config / Main_SDNQ_Diffusers_ZImage_Config _get_submodels(): record the Tokenizer for the slow/fast Qwen2 tokenizer classes, and the TextEncoder only for the loadable Qwen3ForCausalLM class. A pipeline advertising Qwen2ForCausalLM or the multimodal Qwen2VLForConditionalGeneration is therefore not marked self-contained, since the pipeline loaders build a text-only Qwen3ForCausalLM. Self-contained SDNQ pipelines (Z-Image / FLUX.2) - Add shared is_self_contained_sdnq_pipeline() requiring the specific VAE + Qwen3 (text_encoder + tokenizer) submodels; a truthy submodels dict is not enough. - z_image_model_loader and flux2_klein_model_loader use it for both the main-model fallback and the explicit qwen3_source validator; drop the ui_model_format=diffusers hint from z_image qwen3_source_model so SDNQ pipelines aren't filtered out of the pickers. - Mirror the specific-submodels requirement in the frontend readiness checks (Z-Image + FLUX.2, generate and canvas tabs) and buildZImageGraph, and bring the FLUX.2 canvas readiness block to parity with the generate tab. Docs - sdnq-quantization.mdx: stop recommending the unsupported FLUX.2 dev SDNQ export; recommend a supported Klein 9B export and note FLUX.2 dev is not implemented. Add regression tests for every fix: mixed-group and 2D-scale uint4/int5 dequant, device movement of all payloads, SDNQ-aware size accounting, the load-guard helper, sharded SDNQ VAE detection, SDNQ FLUX.2 identification (bare accepted / BFL rejected), SDNQ ZImagePipeline-folder identification, Qwen3 encoder folder identification (SDNQ/causal-LM rejection, Qwen3-only architecture), SDNQ pipeline submodel discovery (loadable vs unloadable encoder classes, fast tokenizer), the self-contained loader/readiness paths incl. partial-pipeline handling, and the qwen3_source_model field template. * Chore Ruff * docs: regenerate invocation-context data for offload_from_vram * Fix: address SDNQ review merge blockers (partial pipelines, Qwen/T5 identification, cat/stack) - SDNQ FLUX.2 & Z-Image submodel discovery now records a component only when its subfolder exists on disk, not merely because model_index.json advertises it. A partial download that keeps the full index no longer classifies as a self-contained pipeline and then fails in the loader on missing fixed vae/, text_encoder/, tokenizer/ subfolders. - Single-file SDNQ Qwen encoder identification now requires Qwen3-only QK-norm (q_norm/k_norm) weights and rejects a bundled Qwen-VL visual tower, so a Qwen2 causal LM or a Qwen-VL checkpoint no longer resolves to Qwen3Encoder_SDNQ_Config (the loader builds a text-only Qwen3ForCausalLM and would fail on either). - SDNQ T5 encoder: share resolve_text_encoder_dir/resolve_tokenizer_dir so the tokenizer is located per layout (child in the standalone bundle, sibling in the inline text_encoder_2 layout) instead of always <path>/tokenizer_2, and reject an install with no resolvable tokenizer_2 at identification time. - SDNQTensor.dequantize_and_run recurses into list/tuple args so torch.cat / torch.stack dequantize their operands instead of redispatching onto the SDNQTensors until RecursionError. Adds identification tests for partial pipelines, single-file Qwen2/Qwen-VL rejection, inline/standalone T5 tokenizer resolution, and torch.cat/torch.stack. * fix(pid): resolve PiD decoder review findings (state, readiness, steps, race) Address all reviewer findings on the PiD decoder feature: - Guard offload_model_from_vram with @synchronized, matching its sibling drop_model, to prevent a race when the VRAM working set is mutated concurrently during model swaps. - On base change (modelChanged), clear a PiD decoder that is incompatible with the new main model's decoder base at the root, respecting the Z-Image -> FLUX decoder reuse so a still-valid decoder is kept. - On switching to a base without PiD support, reset pidMode to 'off' and refit the dimensions so no hidden 4x native grid survives. - Extend the scaled-grid bbox readiness validation to SD3, SDXL and Z-Image, mirroring the existing FLUX.2 native-grid check. - Add a PiD Steps control (slider + number input, 1-8, default 4) with a pidStepsChanged action and selectPidSteps selector, so the documented step count is actually configurable and flows into the graph. Add readiness and paramsSlice tests covering the scaled-grid blocking and the base-change decoder/pidMode behavior * fix(pid): address review — cap steps at 4, honor encoder device, decoder/base guards Merge blockers: - PiD steps 5-8 produced duplicate timesteps: the student schedule has only 4 transitions (a 5-point list), so sub-sampling to >4 steps rounded distinct indices onto the same point and wasted network forwards on repeated timesteps. Cap the public range at 4 across the backend fields (le=4), the UI slider/input (max=4), and the pidSteps zod schema (int, 1-4), and harden _get_t_list with a strictly-decreasing assertion as a safety net. - CPU-only Gemma encoders crashed on CUDA hosts: each PiD invocation passed the global compute device to caption encoding, pushing the tokenizer output to CUDA while a cpu_only encoder stayed on the CPU. Encode on the encoder's actual device instead (next(encoder.parameters()).device), honoring model_on_device(). - Gemma 2 could no longer be configured as a generic TextLLM: the specialised- architecture exclusion rejected Gemma2ForCausalLM unconditionally. Only defer to the encoder config during automatic classification; keep an explicit type=text_llm request valid (the generic causal-LM loader supports it). Follow-ups folded in: - Reject incompatible Gemma 2 sizes before execution: PiD's caption projection is fixed at Gemma-2-2b's 2304-dim hidden state, so the encoder config now rejects 9B (3584) / 27B (4608) up front instead of failing deep in inference. - Validate the PiD decoder's base against each base-specific decode node: the base-agnostic loader let the Nodes editor wire any decoder into any node. Add assert_pid_decoder_matches_base and call it in all seven decode nodes, preserving the Z-Image-reuses-FLUX-decoder case (its node backbone is FLUX). Add regression tests: the distill schedule (strictly decreasing 1-4, safety net trips at 5) and decoder/base validation; the Gemma2 hidden-size gate; and TextLLM classification (auto-defers, explicit type still matches, plain causal LMs match). The expand-prompt pipeline always sent a dedicated "system" role message, which some chat templates (notably Gemma) reject with "System role not supported", 500-ing prompt expansion for those models. When applying the chat template fails with a system-role error, fold the system prompt into the first user turn and retry instead of failing. Adds regression tests for both the fallback and the normal (system-supported) path. * Chore fix * fix(pid): keep large Gemma2 as TextLLM, raise (not assert) schedule guard, compute_device, narrow pid_upscale VAE Address the latest review on the PiD PR: - Merge blocker: automatic classification sent Gemma 2 9B/27B to Unknown. The PiD Gemma2 encoder config rejects their non-2304 hidden size, and TextLLM deferred *every* Gemma2ForCausalLM, so neither matched. TextLLM now defers only the size the encoder config accepts (2304 = Gemma-2-2b); larger variants stay TextLLM. - Schedule safety net used assert, which is stripped under `python -O`, leaving _get_t_list(num_steps=5) returning a duplicate schedule. Raise ValueError instead so the guard holds in optimized runtimes; the regression test now asserts ValueError and passes under `python -O`. - All seven PiD caption paths derived the Gemma device from the first parameter, which is wrong under partial loading (first param on CPU, later modules on CUDA). Use the cache contract's LoadedModel.compute_device instead. - pid_upscale advertised Z-Image / 16-channel VAEs but delegates to the FLUX-only vae_encode. Narrow the field description and validate the VAE is a FLUX AutoEncoder up front (a diffusers AutoencoderKL now fails with a clear error instead of a stripped-assert failure inside vae_encode). Update the TextLLM/Gemma2 tests (per-size config-level + a factory-level check that 2304 -> Gemma2Encoder and 3584/4608 -> TextLLM) and the schedule test (ValueError, green under python -O). * chore: regenerate OpenAPI schema and frontend types * feat(pid): accept single-file GGUF Gemma-2-2b as the PiD caption encoder The PiD Gemma encoder was directory + HuggingFace only, so a llama.cpp GGUF (e.g. gemma-2-2b-it-Q4_K_M.gguf) could not be used. Add GGUF support: - Gemma2Encoder_GGUF_Config: identifies a single .gguf file, reads the GGUF metadata and requires general.architecture == "gemma2" and <arch>.embedding_length == 2304 (Gemma-2-2b), rejecting 9B/27B as the directory config does. - Gemma2EncoderGGUFLoader (format gguf_quantized): loads via transformers from_pretrained(<dir>, gguf_file=<name>) — transformers dequantizes gemma2 GGUFs and reads the tokenizer from the GGUF metadata — then exposes the Gemma2Model decoder, matching the directory loader. PiD encodes the caption once and offloads the encoder, so dequantizing at load is acceptable. - Register the config in the AnyModelConfig union. No frontend change: the PiD encoder picker filters by type=gemma2_encoder, so the GGUF variant appears automatically. Verified end-to-end against a real q4_k_m file: it classifies as Gemma2Encoder_GGUF_Config and loads to a Gemma2Model producing 2304-dim hidden states. Adds config identification tests (match, 9B/27B rejected, non-gemma2 rejected, non-.gguf rejected). * Chore Ruff * fix(models): stop Qwen3 GGUF config from matching Gemma-2 GGUF encoders A Gemma-2 GGUF satisfies the generic Qwen3 GGUF key heuristic (token_embd.weight + blk.* keys), so it matched both Qwen3Encoder_GGUF_Config and the intended Gemma2Encoder_GGUF_Config. On a fresh install the Gemma config happened to win, but re-identification could pick Qwen3, mis-classifying the model. Add _has_gemma2_keys (Gemma uses blk.*.post_attention_norm / post_ffw_norm, which a Qwen3 encoder never has — Qwen3 has attn_q_norm/attn_k_norm instead) and reject such state dicts in both Qwen3 encoder configs' _validate_looks_like_qwen3_model (GGUF and checkpoint), mirroring the existing T5 / Qwen-VL exclusions. The Gemma config already rejects Qwen3 GGUFs via the general.architecture metadata, so the two are now mutually exclusive and identification is deterministic. Add regression tests: _has_gemma2_keys detection and that the Qwen3 GGUF config rejects a Gemma-keyed state dict. * test(models): assert re-identified Gemma GGUF drops the stale Qwen3 variant The Gemma2 GGUF encoder config has no `variant` field, so re-identifying a model previously mis-detected as a Qwen3 GGUF (which carries a variant) drops it — the serialized record has no variant key and replace_model overwrites it away. Assert this explicitly in the Gemma GGUF identification test. * fix(pid): point 2K-to-4K starter decoders at NVIDIA's v1.5 replacements NVIDIA deprecated the FLUX / FLUX.2 / Qwen-Image `res2kto4k_sr4x` PiD decoders and moved them to `checkpoints_deprecated/`, replacing them with the recommended `v1pt5_res2kto4k_sr4x` checkpoints. Our starter models still pointed at the old `checkpoints/` paths, which now 404 on install. Repoint the three affected 2K-to-4K starters (FLUX, FLUX.2, Qwen-Image) to the v1.5 successors and note the upgrade in their descriptions. The 2K (`res2k_sr4x`) decoders and the SD3 / SDXL 2K-to-4K decoders are not deprecated and are unchanged. Base and variant are still sent as explicit overrides, so config identification is unaffected by the new directory name (res2kto4k -> Res2kTo4k_Sr4x). * feat(pid): load Gemma-2 GGUF encoders natively (keep weights quantized) The GGUF Gemma encoder used transformers' from_pretrained(gguf_file=...), which dequantizes every weight at load — so a quantized Gemma cost the same VRAM as the unquantized model. Load it via InvokeAI's GGMLTensor path instead: read the config from GGUF metadata, map llama.cpp tensor names to Gemma2Model, and keep the 2D projection weights as GGMLTensor (dequantized on demand by the model cache). Materialize only the embedding and the RMSNorm weights, subtracting 1 from the norms (llama.cpp folds +1 in; Gemma2RMSNorm re-adds it), and assert nothing is left on meta. Verified: hidden states match the fully-dequantized loader within quantization tolerance. Adds key-mapping tests and a local load/compare test. NVIDIA's v1.5 decoders use a different network (lq_hidden_dim=1024, PiT injection) that build_pid_net (512-dim legacy) cannot load, causing a size-mismatch crash. - Point the FLUX/FLUX.2/Qwen 2K-to-4K starters back at the legacy checkpoints (moved to checkpoints_deprecated/) that the current network loads. - Reject a checkpoint whose lq_proj hidden dim is not the supported 512 at identification time, instead of accepting it and failing inside the decode. - Enumerate all supported backbones (add SDXL, Qwen-Image) in the loader title and correct the variant enum docs (not every backbone ships both presets). Full v1.5 architecture support is planned as a follow-up. Adds PiD decoder identification tests (legacy accepted, 1024-dim v1.5 rejected). * Chore openapi + typegen * Docs update * Chore Ruff format * Chore openapi * fix(sdnq): harden self-contained pipeline detection (review 9228) Addresses the four merge blockers from the PR #9228 review and checks in the reviewer's failing tests. - _get_submodels (FLUX.2 + Z-Image) only checked is_dir(), so the empty component folders an interrupted download leaves behind were recorded as present. Require the files the component's loader actually needs: config + weights for transformer/text_encoder/vae, a vocab/config file for the weightless tokenizer folder. - is_self_contained_sdnq_pipeline() did not require the Transformer submodel, so a malformed model_index.json advertising only VAE/encoder/ tokenizer passed the check even though every loader requests the transformer. - Qwen3Encoder_SDNQ_Folder_Config's config-less fallback matched the generic Qwen keys (model.layers.* / model.embed_tokens.weight), which Qwen2 and Qwen2-VL folders carry too. Mirror the single-file path: reject a bundled visual tower and require the Qwen3-only q_norm/k_norm weights. - isZImage/isFlux2DiffusersMainModelConfig returned true for any nonempty SDNQ submodels map, offering a transformer-only pipeline in the source pickers. Added isSelfContainedSDNQPipeline() mirroring the backend's required-submodel set and reused it in readiness.ts and buildZImageGraph.ts so both sides agree on what "complete" means. Test fixtures now write real component content, since discovery no longer accepts bare directories; the empty-dir case is requested explicitly. * fix(sdnq): address the two merge blockers from the 9228 review flux_vae_decode raised TypeError on a VAE without a shift factor. `shift_factor` is optional on diffusers' AutoencoderKL: the FLUX VAE sets one, a plain SD-style config leaves it None, and `tensor + None` raises. Guard it the way every sibling decode path already does (getattr + explicit None check) — a sweep found this was the only unguarded one. Self-contained pipeline detection trusted the advertised class name. model_index.json records what a pipeline claims each component is, and that claim selects the loader, but nothing forces it to be true. A folder can advertise Qwen3ForCausalLM for text_encoder/ and ship a Qwen2 or multimodal Qwen2-VL model: the class check passes, the file-presence check passes, the pipeline is recorded as self-contained, and the mismatch only surfaces when the loader builds a Qwen3ForCausalLM against a state dict that cannot satisfy it. Components must now also be what the index says they are, read from the component's own config (architectures / _class_name / tokenizer_class). Strict on contradiction, lenient on silence: a component declaring no class at all is still accepted, so repos omitting those keys keep working — the same trade-off as Qwen3Encoder_SDNQ_Config._validate_is_qwen3_encoder. The FLUX.1 _get_submodels is deliberately untouched; the review lists it under follow-on candidates, not blockers. Also carries two fixups from the main merge: the Klein loader's error message now names the accepted shape ("Diffusers-style FLUX.2 pipeline") since it also accepts SDNQ pipelines, and the SDNQ-added readiness test builders gained main's new required hasFlux2DevDiffusersSource field. * fix(sdnq): reject pipeline components that do not identify themselves Follow-up to the previous blocker fix, which only went half way: it rejected a component whose config *contradicted* the class model_index advertises, but still accepted one that declared nothing at all. "Some config.json plus some weight file" is exactly the shape an unrelated model has, so that left the hole open for every mismatched component whose config happens to omit the key — which is what the review means by "any component directory containing a config file and weight file". Reproduced before fixing: a text_encoder/ holding a Qwen2-VL state dict (visual.* keys) with a config declaring no architecture, under an index advertising Qwen3ForCausalLM, was recorded as a complete self-contained pipeline. Weight-bearing components (transformer / text_encoder / vae) must now name their own class. save_pretrained writes _class_name (diffusers) or architectures (transformers) for anything it saves, so one that names nothing is not a normally-produced component. The tokenizer stays lenient: it carries no weights, so nothing can be mis-instantiated against it, and tokenizer_class is written less consistently. Rejecting is not fatal to the install — the component is left unrecorded, the pipeline is simply not self-contained, and the user wires it up explicitly, which is the right outcome for a folder we cannot confirm. The fixtures wrote empty component configs and would have failed under the stricter rule; they now carry the class a real save_pretrained writes. The old leniency test is replaced by two: rejection for text_encoder/vae, and the remaining leniency for the tokenizer. * fix(sdnq): work through the follow-on findings from the 9228 review Discovery (main.py) - FLUX.1 recorded components on the index's word alone. It now runs the same two guards as the FLUX.2 / Z-Image pipelines: the folder must hold the files its loader needs, and must declare a class consistent with what the index advertises. Generalising this surfaced a latent bug — _sdnq_component_dir_is_populated only exempted Tokenizer, so FLUX.1's second tokenizer slot would have been rejected for carrying no weights. Hence _WEIGHTLESS_SUBMODEL_TYPES. - Submodels are no longer taken from a serialized config. A persisted map records what existed at install time, so a component deleted since stayed "present" and the pipeline stayed self-contained until a loader opened a folder that was gone. Loading (flux.py, vae.py) - The SDNQ loader reconstructed model_path / submodel_type.value, which assumes each index key equals its slot name. A pipeline that names its CLIP encoder something else was discovered fine and then loaded from a path that does not exist. It now follows the discovered path. - The SDNQ loaders build modules with init_empty_weights + load_state_dict rather than from_pretrained, so nothing called .eval() and anything dropout- or norm-sensitive stayed in training mode. Identification (t5_encoder.py) - The SDNQ weight/scale pair is matched across the whole directory, not within each shard. Sharding splits by tensor order, so the pair routinely straddles files; per-file matching reported such a checkpoint as unquantized whenever it carried no quantization marker. Diagnostics (sdnq/utils.py, sdnq/loaders.py, sdnq/sdnq_tensor.py) - The uint4 diagnostic ran full-tensor reductions and a unique() sort on the first dequantization of every model, before anything checked whether the output would be read. Now gated on the log level and bounded to a fixed sample. Three stdout prints removed; two were verbatim duplicates of the logger call beneath them. flux_model_loader - A complete SDNQ pipeline ships its own T5, CLIP and VAE, which is what the SDNQ docs promise, but the node required separate identifiers anyway and users had to install duplicates. The three inputs are now optional and fall back to the pipeline. FLUX.1 needs six submodels (the T5 pair on top), so this uses its own completeness check rather than the single-encoder one. Models that cannot supply the parts are still rejected, naming every missing part at once. Node version 1.1.0. * fix(sdnq): close the review-4891383742 blockers and generalise the two loader fixes Merge blockers - The frontend still demanded standalone T5 / CLIP / VAE for every FLUX.1 model, in readiness (both tabs) and as an assert in buildFLUXGraph, so the backend's self-contained path was unreachable through the UI. A complete SDNQ pipeline now satisfies readiness on its own and the graph omits the inputs, letting the node resolve them from the main model. FLUX.1 needs six submodels (the T5 pair on top), so the frontend check mirrors is_self_contained_sdnq_flux1_pipeline rather than the single-encoder one. Documented in sdnq-quantization.mdx. - The component population check accepted .bin/.pt/.pth/.ckpt/.gguf, but sdnq_sd_loader globs *.safetensors and raises when it finds none, so a correctly declared component holding an unsupported format was recorded as self-contained and failed at load time. Made loader-specific rather than narrowing everything: a *quantized* component (quantization_config .json present) must ship safetensors, an unquantized one — typically the VAE, which SDNQ exports leave in bfloat16 — still goes through from_pretrained and may ship what diffusers reads. Generalising the two follow-ons, per review and beyond it A survey of every loader that hand-builds a module found 43 of 52 returning without .eval(), across FLUX, FLUX.2, Z-Image, Krea-2, Qwen-Image, Anima, Wan, Mistral and Ideogram — not the handful named. Rather than 43 edits, put_in_eval_mode now runs in load_default._load_and_cache, the single construction choke point every loader passes through, so it also covers loaders that do not exist yet. It is idempotent for the nine that already did it, and passes tokenizers, schedulers and pipelines through untouched. The local _in_eval_mode helper added to flux.py earlier is gone; one mechanism, not two. Likewise the path fix: resolve_submodel_path replaces the per-site Path(config.path) / "<slot>" reconstruction in the FLUX.1 dispatch, the FLUX.2 transformer/encoder/tokenizer/VAE loaders and the Z-Image equivalents — the loaders whose configs carry an index-key-derived submodels map. The plain-diffusers loaders (Stable Diffusion, ONNX, CogView4, Wan, Qwen-Image, Krea-2) keep the conventional layout on purpose: they have no such map. Verified against real installs: Z-Image-Turbo-SDNQ-uint4-svd-r32 and FLUX.2-klein-9B-SDNQ-4bit-dynamic-svd-r32 identify unchanged, with all four submodels and correct variants. Both declare a class on every weight-bearing component and ship only safetensors, which is what makes the stricter discovery safe. * fix(sdnq): decide component weight formats by loader, and size what we load Blocker — the quantization marker was the wrong criterion Whether `sdnq_sd_loader` reads a component is a property of the loader, not of the folder, but the population check asked `_is_sdnq_folder`, which only looks for quantization_config.json. A component carrying SDNQ weight/scale keys without that marker read as unquantized, so a .bin or .gguf passed discovery and then failed in a loader that was always going to want safetensors. Modelled explicitly instead, because it genuinely differs per pipeline: FLUX.2 and Z-Image build their transformer *and* Qwen3 text encoder with sdnq_sd_loader unconditionally, while FLUX.1 does so only for the transformer — its CLIP, T5 and VAE branch on the folder marker and fall back to from_pretrained. Requiring safetensors across the board would have rejected a legitimate unquantized CLIP encoder shipping .bin. Finding — cache sizing had drifted from the load path Loading follows the discovered component path now, but get_size_fs still appended submodel_type.value. For an index key like `clip_encoder` that measures a `text_encoder/` that does not exist, returns 0, and make_room reserves nothing before a multi-GB component is read. It now goes through resolve_submodel_path, on a branch that only triggers when the discovered path differs from the conventional one — the normal case is byte-identical, which matters because calc_model_size_by_fs has a standalone-install fallback keyed on the subfolder name. Docs The reviewer asked for documentation of the weight-format rule. It is a sentence on the existing "no safetensors files found" troubleshooting entry rather than a section in the how-to flow: published SDNQ repos ship safetensors throughout, so a user installing one can never hit it, and the symptom already had a home. Both mutation-checked: ignoring the per-loader set fails four tests, reverting get_size_fs fails one. Z-Image-Turbo-SDNQ-uint4-svd-r32 and FLUX.2-klein-9B-SDNQ-4bit-dynamic-svd-r32 identify unchanged. * fix(ui): keep the FLUX.1 readiness test off MainModelConfig.submodels The fixture is cast to MainModelConfig, which has no `submodels` field — the generated schema doesn't carry the SDNQ variants — so reading it back through the cast failed tsc in CI. Keep the map as its own const. vitest's typecheck only covers *.test-d.ts, so a full `tsc --noEmit` is what catches this; running the suite is not enough after touching a test file. Also formats readiness.ts, which prettier would have failed on next. * fix(sdnq): one SDNQ detector, and bump the two nodes whose contract changed Detection had four near-identical implementations — configs/main.py, the FLUX and Z-Image loaders, and vae.py — and only vae.py looked past the quantization_config.json marker. That divergence is a defect, not untidiness: identification and loading inspect the same folder and must reach the same verdict. On a markerless export they didn't, so identification handed the folder to a plain-diffusers config and the loader then ran from_pretrained() over packed SDNQ weights. They now all delegate to quantization/sdnq/detection.py, which checks the marker first and falls back to a weight/scale key pair resolved across the union of every shard (sharding routinely separates a weight from its scale). t5_encoder's key-only check is the same code. The tests drive every detector over the same cases so they cannot drift apart again. Node versions: flux2_klein_model_loader 1.0.0 -> 1.1.0 and z_image_model_loader 3.0.0 -> 3.1.0. Both dropped their ui_model_format=Diffusers pin and gained a self-contained-pipeline fallback in this PR, so a node serialized against the old version still carries the old field contract. Also reads the FLUX.2 variant through flux2_variant.py instead of comparing the transformer geometry by hand. The hand-rolled comparison had no [dev] branch, so [dev]'s 15360 / 6144 fell into the else and a [dev] SDNQ pipeline was silently identified as Klein4B. This is not [dev] support — that is separate work — only a correct label instead of a wrong one; the Klein path is covered by its own test. Verified against real installs: Z-Image-Turbo-SDNQ-uint4-svd-r32 and FLUX.2-klein-9B-SDNQ-4bit-dynamic-svd-r32 identify unchanged. * fix(sdnq): reject a transformer whose weights the loader cannot read `_validate_has_sdnq_transformer` only checked that the folder looked SDNQ. Discovery, which uses a stricter predicate, then left Transformer out of `submodels` — but the config was created anyway, so the model installed cleanly and only failed when a loader opened the transformer path at generation time. Identification now requires `_sdnq_component_dir_is_populated` for the transformer: the same predicate discovery uses, so the two cannot disagree about the slot every SDNQ loader needs. Applied in all three pipeline configs. Docs: the detection section still named `quant_method: "sdnq"` as the criterion. Markerless exports are recognized by their weight/scale key pairs now, across shards, and the transformer's weights have to be readable safetensors — both documented. --------- Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com> Co-authored-by: JPPhoto <jpollack@jpollackphoto.com> | 26 天前 | |
feat(qwen3): bundle Qwen3 tokenizer for offline single-file/GGUF encoders (#9338) * feat(qwen3): bundle Qwen3 tokenizer for offline single-file/GGUF encoders" -m "Single-file (safetensors) and GGUF Qwen3 encoder checkpoints used by Anima (0.6B) and Z-Image (4B/8B) ship weights only — no tokenizer files. The loader pulled the tokenizer from Qwen/Qwen3-4B on HuggingFace, which fails offline / airgapped and whenever the HF cache is not persisted (e.g. Docker without a cache volume). Vendor the self-contained Qwen3 fast tokenizer (Apache-2.0, from Qwen/Qwen3-4B) in the package and load it locally, mirroring the bundled T5-XXL tokenizer (#9244). The Qwen3 BPE tokenizer is identical across the 0.6B/4B/8B variants, so a single copy serves every Qwen3 encoder. Removes the HuggingFace download path from both the checkpoint and GGUF loaders. * fix(qwen3): gzip bundled tokenizer to pass LFS check The vendored Qwen3 tokenizer.json is ~11MB, over the repo's 10MB lfs-warning threshold, failing the "lfs checks" CI job. Git LFS is unsuitable here since the file must ship inside the wheel for offline use. Vendor it gzip-compressed (~2MB) instead and decompress into a temp dir at load time. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * refactor(qwen3): fix stale tokenizer-loader comments and method name The single-file/GGUF Qwen3 loaders now use the vendored tokenizer, but the call-site comments still described the removed HuggingFace download path and the method was still named _load_tokenizer_with_offline_fallback despite having no fallback. Rename to _load_bundled_tokenizer and update the comments to match. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(qwen3): restore chat_template in bundled tokenizer config The vendored tokenizer_config.json was missing the chat_template that Qwen/Qwen3-4B ships. The Z-Image text encoder formats prompts via tokenizer.apply_chat_template(), which raises ValueError: Cannot use chat template functions because tokenizer.chat_template is not set ... so GGUF/single-file Qwen3 encoders failed at encode time. The old HF-download path pulled the full config (template included), so this was a regression introduced by bundling. Restore the exact upstream Qwen3-4B chat_template and add a regression test. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com> Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> | 1 个月前 | |
feat(model): Add ER SDE / DPM++ 2M Scheduler Support For Anima (#9125) * refactor(anima): reshape ANIMA_SCHEDULER_MAP to (class, kwargs) tuples Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * refactor(anima): address Task 1 review feedback Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * feat(anima): add dpmpp_2m and dpmpp_2m_sde schedulers * refactor(anima): unify ANIMA_SHIFT in schedulers.py and add Literal-coverage test * fix(anima): seed generator into scheduler.step for SDE reproducibility * feat(anima): add pure ancestral-Euler step helper Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * docs(anima): fix _anima_euler_ancestral_step docstring formula to match code Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor(anima): address Task 4 review feedback * feat(anima): add euler_a (rectified-flow ancestral Euler) scheduler * fix(anima): sample euler_a noise in float32 to avoid bfloat16 quantization * chore(anima): bump anima_denoise to v1.3.0 and regen schema * fix(frontend): revert Windows path-separator drift in schema regen * fix(anima): gate step_generator construction to schedulers that need it Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * chore(anima): apply ruff lint and format fixes Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * feat(frontend): expose new Anima schedulers in dropdown and metadata recall * chore(frontend): apply prettier wrap to setAnimaScheduler PayloadAction * fix(anima): correct euler_a math — variance-preserving noise mix, not biased Euler * revert(anima): remove euler_a scheduler — quality not worth the complexity * chore(anima): apply ruff format (trim trailing blank lines) * feat(rectified-flow): add order-1 ER-SDE stepper for rectified flow Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * docs(rectified-flow): document lambda_next>0 precondition; tighten terminal-step test Address code review feedback on order-1 ER-SDE stepper: - Add docstring preconditions to integral helpers noting the logarithmic singularity at lam=0 (callers must guard sigma_next>0). - Tighten terminal-step test from atol=1e-5 to torch.equal — the algebra is exact when sigma_next=0, not approximate. * feat(rectified-flow): add 2nd-order Taylor extension to ER-SDE stepper * test(rectified-flow): tighten Task 2 state-mutation test with value assertion Address code review feedback on the 2nd-order Taylor extension tests: - Assert state.old_d_x0 equals the analytically-computed d_x0 = 0.2v / (1.5 - 4.0) rather than just checking it's non-None. - Document that x_t is intentionally re-used across calls (state threading test, not trajectory correctness). - Document the order-2 correction coefficient magnitude that justifies the atol=1e-3 threshold in the engagement test. * feat(rectified-flow): add 3rd-order Taylor extension to ER-SDE stepper * docs(rectified-flow): document have_two_back invariant + order-3 test margin Address code review feedback on Task 3: - Comment why have_two_back checks both old_d_x0 and sigma_prev_prev (the sigma~=1 boundary path can break the joint invariant). - Document the analytically verified ~0.0004 per-element correction magnitude that justifies the atol=1e-3 threshold in the order-3 test. * feat(anima): register er_sde scheduler choice * docs(anima): document custom-code-path scheduler convention; tighten test Address code review feedback on Task 4: - Add an in-file comment above ANIMA_SCHEDULER_MAP explaining the convention: schedulers with custom code paths (er_sde) live in the Literal+labels only, not the map. - Hoist `import typing` to module-level in test_anima_schedulers.py (was inline-imported in two test functions). - Pin the er_sde label value (== "ER-SDE"), not just key existence. * feat(anima): wire er_sde scheduler into denoise loop * docs(anima): document float32 noise dtype and sigma_next/sigma_prev naming Address code review feedback on Task 5: - Comment explaining why fresh_noise is float32 (matches er_sde_rf_step's dtype contract with x_t.to(float32)). - Bridging comment at the inpaint extension call clarifying that sigma_next here means the same thing as sigma_prev in the Euler branch and the AnimaInpaintExtension API. * chore(anima): bump anima_denoise to v1.4.0 and regen schema * feat(frontend): expose er_sde scheduler in dropdown and metadata recall Address code review on Task 6 — er_sde was registered in the OpenAPI schema but missing from the frontend's own Zod enums and Redux PayloadAction types, so: - The combobox dropdown didn't list it. - setAnimaScheduler('er_sde') would fail TypeScript at the call site. - Metadata recall for er_sde-generated images would silently no-op (the scheduler value couldn't pass zParameterScheduler validation). Changes: - Add er_sde to zAnimaSchedulerField (the per-Anima Zod enum). - Add er_sde to the animaScheduler state-shape Zod enum. - Widen setAnimaScheduler's PayloadAction union. - Add ER-SDE option to the ParamAnimaScheduler combobox. - Make the metadata Scheduler handler accept ParameterAnimaScheduler too, with a fallback parse and a narrowing guard before the SD/SDXL dispatch. * fix(rectified-flow): guard 2nd-order branch against sigma_prev_curr=1.0 When step 0 goes through the sigma_curr=1 closed-form limit branch it writes state.sigma_prev_curr=1.0. On step 1, have_one_back was True and the 2nd-order path called _lambda(1.0) = 1.0/(1.0-1.0), crashing with ZeroDivisionError in every real denoise run. Fix: extend the have_one_back guard to require that sigma_prev_curr is more than _SIGMA_ONE_TOLERANCE below 1.0. The finite-difference derivative across the limit step is not meaningful, so skipping the 2nd-order term on that transition is correct. Order-3 is already gated behind old_d_x0 being set, which this path never sets, so no additional guard is needed there. Adds a regression test that runs the full sigma=1.0->0.9->0.7->0.0 sequence and asserts no ZeroDivisionError and all-finite output. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * feat(scheduler): add ERSDEScheduler available to SD/SDXL ER-SDE solver (Cui et al., arXiv:2309.06169) usable across SD/SDXL (VP-SDE) and rectified-flow models. Anima migration follows in subsequent commits. - ERSDEScheduler(SchedulerMixin, ConfigMixin) with prediction_type (epsilon | v_prediction | flow_prediction), use_flow_sigmas, solver_order (1/2/3 with auto-warmup), and stochastic toggle - set_timesteps(sigmas=) for pre-shifted Anima/FLUX/Z-Image schedules - Closed-form limit at sigma=1 in flow mode - Unit tests + VP smoke + 5/5 Anima parity vs er_sde_rf_step (worst delta 5.137e-07) - Frontend wiring: zSchedulerField, SCHEDULER_OPTIONS, OpenAPI regen - parsing.tsx cleanup: removes the AnyScheduler widening since er_sde is now a first-class general scheduler Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * feat(anima): wire ERSDEScheduler into ANIMA_SCHEDULER_MAP Adds er_sde to the standard scheduler dispatch map with rectified-flow kwargs (flow_prediction, use_flow_sigmas=True, flow_shift=3.0, solver_order=3, stochastic=True). Anima still routes through the legacy elif is_er_sde: branch — that's removed in a follow-up commit. This is the additive prerequisite that lets the cutover happen without a window where Anima can't use ER-SDE. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * test(anima): add ER-SDE dispatch integration tests Verifies the ANIMA_SCHEDULER_MAP['er_sde'] entry instantiates correctly, accepts pre-shifted sigmas via set_timesteps(sigmas=...), and resets multistep state. Catches wiring regressions that the algorithm-level parity test cannot. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor(anima): remove elif is_er_sde branch, dispatch through ANIMA_SCHEDULER_MAP Anima ER-SDE now flows through the same standard scheduler path as dpmpp_2m_sde — pre-shifted sigmas via scheduler.set_timesteps(sigmas=...), inpaint extension via inpaint_extension.merge_intermediate_latents_with_init_latents, step_callback per-step. The custom code path was the only thing keeping ER-SDE off the universal pipeline. Bumps invocation version to 1.5.0. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * docs(er_sde): mark module as internal reference and parity oracle ERSDEScheduler is now the production code path. er_sde_rf_step is kept as the comparison oracle for the scheduler's parity test, and as a self-contained mathematical reference for the rectified-flow algebra. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * chore(rectified-flow): remove er_sde.py reference helper ERSDEScheduler is the production code path. The pure-function helper was retained as a parity oracle but YAGNI — keeping ~200 lines of code purely as a regression net for hypothetical future drift isn't worth the maintenance signal it generates. Removes: - invokeai/backend/rectified_flow/er_sde.py - tests/backend/rectified_flow/test_er_sde.py - tests/backend/rectified_flow/test_er_sde_scheduler_anima_parity.py ERSDEScheduler's own tests (test_er_sde_scheduler.py) remain — they exercise both VP-SDE and rectified-flow paths directly. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * fix(schema): restore forward slashes in @default cache dir paths Windows-side typegen run flipped these to backslashes. Restore the canonical forward-slash form to match upstream. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * style: apply ruff lint and format to ER-SDE files Sort imports + format per project ruff config (line-length 120). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor(state): use zParameterAnimaScheduler in state shape The animaScheduler field inlined its enum (originally to add er_sde when the shared schema didn't have it yet). Now that zAnimaSchedulerField already includes er_sde, reference the shared zParameterAnimaScheduler to match the pattern used by scheduler/fluxScheduler/zImageScheduler. Drops the redundant .default('euler') — initial value comes from getInitialParamsState. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * style(types): sort imports per simple-import-sort Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * fix(anima): honour clipped sigma schedule for DPM++ img2img/inpaint DPMSolverMultistepScheduler doesn't accept sigmas= in diffusers 0.35.1, so the fallback previously called set_timesteps(num_inference_steps=total_steps) which regenerated a full schedule from sigma_max, ignoring denoising_start/end. When the scheduler supports set_begin_index, call set_timesteps with the full step count and offset into it, so the internal flow_shift applies correctly and denoising starts at the right sigma. Also fixes the inpaint sigma_prev lookup and the timestep loop to use the same offset, and corrects the false parity-test reference in the ER-SDE dispatch test docstring. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * style: apply ruff format to anima_denoise dispatch block Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * refactor(anima): extract scheduler driver and fix Heun progress/inpaint bugs Encapsulate per-scheduler dispatch quirks (sigmas= vs num_inference_steps=, Heun's doubled-array index, set_begin_index path) in AnimaSchedulerDriver and tighten two latent bugs in the Heun path of anima_denoise: * Heun's terminal first-order step never reported a user-step completion, so progress capped at N-1 of N. The driver now flags it via sigma_prev==0, and the <= total_steps clamp that papered over the off-by-one is gone. * The inpaint mix ran after every Heun half-step, corrupting the second-order corrector's input (RectifiedFlowInpaintExtension's docstring says it should be called after each denoising step — i.e. once per user step). Mix is now gated on completes_user_step, which is unconditionally True for non-Heun. Also: Heun shift kwarg switched to ANIMA_SHIFT (its set_timesteps doesn't accept sigmas=, so it builds its own internal schedule); narrative comments in scheduler_driver and er_sde_scheduler trimmed; new tests covering driver iteration counts, terminal sigma_prev, seed determinism, and the begin_index fallback for clipped DPM++ schedules. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * style: apply ruff lint and format to anima scheduler driver Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * ci: re-run after flaky token-expiration test The 1-second JWT token-expiration test in test_token_service.py is timing sensitive — passes locally on retry. Empty commit to retrigger CI. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Co-authored-by: Your Name <you@example.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com> | 3 个月前 | |
feat(mm): siglip model loading supports partial loading In the previous commit, the LLaVA model was updated to support partial loading. In this commit, the SigLIP model is updated in the same way. This model is used for FLUX Redux. It's <4GB and only ever run in isolation, so it won't benefit from partial loading for the vast majority of users. Regardless, I think it is best if we make _all_ models work with partial loading. PS: I also fixed the initial load dtype issue, described in the prev commit. It's probably a non-issue for this model, but we may as well fix it. | 1 年前 | |
feat: add native Intel XPU (torch.xpu) device support (#9401) * feat(backend): add Intel XPU (torch.xpu) device support Additive xpu branches only: device selection and normalization, float16 default, VRAM queries with a passthrough-VM fallback (missing SYCL free-memory aspect), fp8 layerwise casting via a runtime probe, VAE auto-tiling, partial loading, stats/OOM handling, multi-GPU parallel session execution (device enumeration, config/API validation, worker pinning, and the generation-device options endpoint), and the auxiliary image utilities (depth/SAM/DINO pipelines accept xpu instead of falling back to CPU; cache clearing is device-agnostic). CUDA (incl. ROCm), MPS, and CPU behavior unchanged. Verified end to end on Arc Pro B70 hardware, including dual-GPU worker startup. * test(backend): add XPU coverage for TorchDevice Mock-based, mirroring the CUDA/MPS suites: device choice, dtype, normalize, the xpu_mem_get_info fallback branches, and multi-GPU generation_devices resolution/validation/labeling on XPU. Also makes the auto-without-CUDA generation-devices test hermetic on XPU machines. * build: add [xpu] extra torch 2.7.1+xpu / torchvision 0.22.1+xpu / pytorch-triton-xpu 3.3.1 from the torch-xpu index, gated to linux-x86_64 and win_amd64; uv.lock regenerated. * feat(backend): extend idle-GPU text encoder offload to XPU The idle-device arbiter and the session processor's borrow path both gated on `device.type == "cuda"`, so on a multi-XPU system no device ever registered and `offload_text_encoders_to_idle_gpus` (enabled by default) silently did nothing: encoders kept churning the denoise model in and out of VRAM. Register and lend XPU devices alongside CUDA. MPS is deliberately excluded -- it is always a single shared device, so there is never another GPU to borrow. Verified on a dual Intel Arc Pro B70 host: a text encoder node now runs on the idle GPU while the session denoises on the other ("Running compel on idle device xpu:0 (session device xpu:1)"). * feat(ui): show the executing GPU for XPU sessions Queue items already persist the executing device generically (e.g. "xpu:1"), but both readers dropped it: the session event only forwarded devices starting with "cuda", and the frontend index parser only matched /^cuda:(\d+)$/. On a multi-XPU system the progress circle and queue-item badges were therefore always blank. Accept indexed XPU devices in both places, and correct the queue-item field description, which claimed the device is set only on CUDA. * fix(mm): gate Krea 2 fp8 encoder casting on fp8 storage support The Qwen3-VL encoder kept its fp8 storage only on CUDA, so elsewhere an fp8 checkpoint was loaded as full bf16 (~8.9GB instead of ~4.4GB) and thrashed partial loading when sharing a GPU with a large transformer. Reuse the existing cached `_device_supports_fp8_storage` probe, which already backs the layerwise-casting path. It returns True unconditionally on CUDA, so CUDA behaviour is unchanged. * chore: label XPU devices by index in load logs and fp8 help text Model load lines printed the device index only for CUDA, so every model on a multi-XPU host logged as a bare "xpu device", making it impossible to tell the GPUs apart. The FP8 Storage tooltip likewise claimed CUDA-only support. * test: cover XPU config validation, progress device and fp8 probe Three paths changed by this branch had no coverage: - The `device` field pattern was untested. `test_device_choice_xpu` looks like it covers it, but the config model does not enable `validate_assignment`, so assigning `config.device` skips validation entirely; only constructing the model exercises the pattern. Added constructor-based valid/invalid cases. - `generation_devices` validation was parametrized for cuda/cpu/mps only. - The progress event's device field, which now reports XPU sessions. Also cover `_device_supports_fp8_storage`, which gates FP8 storage in both the generic layerwise-casting path and the Krea 2 encoder: CUDA answers True without probing, CPU is rejected, and a failing XPU probe returns False instead of raising. Each new test was verified to fail when the corresponding fix is reverted. * fix(nodes): recognise XPU out-of-memory errors in the Anima VAE retry The Anima VAE decode catches OOM and retries once with tiling, which caps peak allocation. Detection matched `torch.cuda.OutOfMemoryError` or the words "out of memory" in the message, so it missed XPU entirely: torch's XPU backend does not raise a recoverable `torch.OutOfMemoryError` on exhaustion, it surfaces the Level Zero/UR result code as a plain RuntimeError -- and `UR_RESULT_ERROR_OUT_OF_DEVICE_MEMORY` contains no spaces, so the existing substring never matched. The decode therefore failed outright instead of retrying tiled. Match the `*_OUT_OF_DEVICE_MEMORY` / `*_OUT_OF_HOST_MEMORY` spellings (both UR and ZE prefixes) alongside the existing conditions, and fold the cuDNN/cuBLAS checks into the same case-insensitive comparison. Extends the existing parametrized retry test with the three XPU spellings; each was verified to fail before this change. Note the driver behaviour itself is not reproducible on the hardware used here -- this stack overcommits into host RAM and hangs rather than raising -- so the tests pin the classifier, not the driver. * style: wrap long vram_usage_gb ternary for ruff * fix: drop CUDA-only wording from progress device description Matches the committed openapi/schema artifacts, which already say "on a GPU". * docs: regenerate settings data for xpu device values * fix: stop xpu VRAM probe from reporting an unknown total as zero (0, 0) made the cache's available-VRAM arithmetic collapse to a constant -working_mem budget for the life of the process. Also widen the except: the failure type moves between torch releases (RuntimeError for the missing SYCL aspect, AssertionError from _lazy_init), and warn once when the blind estimate is in use. * fix: probe fp8 support on the target device, per device, without caching failures The probe allocated via an index-less "xpu", which resolves through the thread's current XPU device rather than the device being loaded onto -- so during idle-GPU encoder offload it measured the busy denoise GPU. It was also keyed on device type, letting one device decide for another, and memoised transient failures (it runs during a load, when the device may be momentarily full) with no way back but a restart. Also probe the bf16 upcast, which is the runtime path for Krea-2/FLUX. * fix: pin torch current device when borrowing an idle GPU Worker startup set both the session device and torch's per-thread current device; the offload borrow set only the former, leaving index-less allocations on the worker's own GPU. Extracted the shared helper and guarded it on backend availability. * fix: keep idle-GPU borrows within one device type generation_devices accepts a mixed list, so a cuda session could be handed an xpu device for its text encoder. * feat: detect Intel integrated GPUs via Level Zero torch exposes no is-integrated flag, but Level Zero does (ZE_DEVICE_PROPERTY_FLAG_INTEGRATED), and its loader already ships with the torch+xpu runtime -- so no new dependency and no compiled extension. Use it to keep iGPUs out of `generation_devices: auto` when a discrete GPU exists, and to stop budgeting them as dedicated VRAM (they share system RAM, like MPS). An unknown answer keeps the previous behaviour, an iGPU-only machine keeps its device, and an explicit device list can still opt one in. * feat: add xpu torch index to pins.json Gives the launcher an Intel install option instead of requiring a manual pip install of the extra. * fix: report VRAM diagnostics for the device in use All three sites dispatched on torch.cuda.is_available() first, so a mixed NVIDIA + Arc box running on xpu reported a constant 0.0 GB and logged "CUDA Memory Allocated" -- which would make XPU bug reports unactionable. * docs: record why xpu takes the CUDA VAE constants and keeps the broad OOM needle XPU SDPA was measured on Arc Pro B70 / torch 2.13+xpu: peak memory doubles when sequence length doubles (2.00x across 2048-16384; 2.0 MB at seq=16384 vs 512 MB for a materialised score matrix). So XPU is in CUDA's O(area) regime, not ROCm's math-attention regime, and the existing constants are correct rather than accidental. * fix: derive rand_device metadata from the backend's devices Was hardcoded to 'cuda' for any non-CPU noise, which is wrong on Arc. Falls back to 'cuda' when the device query has not resolved, so Nvidia metadata is unchanged. * docs: add Intel Arc install, driver and VRAM-reporting notes * fix: probe fp8 device support only when a model requests it The probe was the first statement in _should_use_fp8, so it allocated on the GPU during the first load of any model at all -- tokenizer, VAE, scheduler -- and on API/install threads it forced XPU lazy SYCL init on a thread that never generates. Moved below the exclusions. * fix: query Level Zero Sysman for driver-global free VRAM on xpu The blind estimate (total minus this process's reserved bytes) is what made _get_vram_available over-commit on a shared GPU: it feeds a formula that assumes a driver-global figure. Sysman's zesMemoryGetState reports that figure and is often available when the SYCL ext_intel_free_memory aspect is not, so try it before estimating. Measured on Arc Pro B70 with 16 GiB held by another process: Sysman reported 15.553 GiB free, the estimate 31.725 -- a 16.172 GiB error, exactly the foreign allocation. Sysman is not a guaranteed substitute (torch's query bottoms out in the same layer), so the estimate remains as a last resort. * fix: make the fp8 probe mirror the runtime cast path The storage cast happens on CPU while params are still CPU-resident, then the fp8 tensor is copied to the device and the pre-hook upcasts there. Probing all three steps on the device would pass on a build where the host->device fp8 copy or one upcast target fails, and break at forward time instead. Verified on Arc Pro B70 / torch 2.13+xpu: the full sequence works on both cards. * fix: degrade gracefully when a backend cannot name a device torch.xpu.get_device_name goes through _lazy_init, which raises AssertionError on a build without XPU. Naming is used only for labels and logs, so fall back to the device string rather than propagating. * fix: resolve an index-less device in the Sysman VRAM query Returning None for a device with no index would skip the driver-global query and fall through to the blind estimate with no visible symptom. Callers currently always pass a concrete device, so this is a latent hazard rather than a live bug. * fix: declare ctypes prototypes for the Level Zero calls Handles come back from (c_void_p * n)() as plain Python ints, and ctypes converts an undeclared int argument to a C int -- 32 bits. Any handle above 2**31 was being silently truncated; a direct test of that path segfaults. It happened to work on the B70 because the handles fit. Also: release the idle-GPU borrow if re-pinning raises (the setup was outside the try, so a failure there stranded the lock for the life of the process), and report a failing fp8 probe once per device instead of on every model load. * fix: drop the ZES_ENABLE_SYSMAN mutation from the Sysman probe Setting a process-wide environment variable from a read-only query leaks into child processes. It also bought nothing: the variable only gates Sysman on runtimes predating zesInit and must be set before Level Zero initialises, which torch has already done by then. Verified on Arc Pro B70 that zesInit succeeds with the variable unset. * refactor: tidy up the xpu additions after a cleanup review level_zero: cache the loader so it is opened and its prototypes configured once rather than twice, share the driver/device enumeration and its ordering guard between the two probes, and collapse the Sysman pair of globals into one nullable tuple. Also: fp8 support cache is a set (it only ever stored True), the pbr_maps empty_cache is routed through TorchDevice like the PR's other conversions, the shared-memory VRAM branch stops re-testing the device type it matched on, `_auto_generation_devices` partitions in one pass, and rand_device only answers when every generation device is the same accelerator. Merges three duplicate mem_get_info tests into one parametrized case and drops two fp8 probe tests fully subsumed by the cast-sequence test. * build: pin the xpu extra to torch 2.13.0 Intel's XPU backend matured considerably after 2.7.1: torch.xpu.mem_get_info() works on driver/kernel combinations where it previously raised, and the oneAPI user-space runtime ships with the wheel, so upgrading torch upgrades it too. Follows the rocm extra, which already pins ahead of cpu/cuda. pytorch-triton-xpu was renamed triton-xpu upstream. The darwin/aarch64 fallbacks stay on 2.7.1 to match the other extras and the project's torch<2.8.0 constraint on darwin. cpu/cuda/rocm exports are unchanged package-for-package (196/211/197); the only delta is a dropped "via pytorch-triton-xpu" comment annotation from the rename. * test: stub Sysman in the unknown-total xpu probe test Without it the test only passes where Level Zero cannot answer -- never on the Intel hardware the probe exists for, where Sysman returns before the tier under test is reached. * build: teach the pins check about the xpu index Its per-platform allowlist rejects anything unlisted, so pins.json's xpu entry fails it. PyTorch publishes XPU wheels for win32 and linux x86_64, matching the extra's markers. * fix: defer the xpu device pin like cuda's torch.xpu.set_device() brings up a SYCL context that holds VRAM in an otherwise idle process, the same reason the CUDA pin waits for the first claimed queue item. * docs: regenerate settings data on linux Regenerating on Windows flips two path defaults to backslashes, which the docs check rejects. * fix(mm): handle shared memory on integrated GPUs Their VRAM is system RAM, so a RAM copy doubles each model's footprint against the same pool. Drop it, letting a full load move weights rather than copy them. Keep partial loading on -- it is the only path that respects vram_available -- and raise a clean error when a full-load-only model cannot fit, instead of walking into an uncatchable OOM-kill. Warn once when a setting is overridden. Scoped to integrated XPU; CPU and MPS are unchanged. * docs: note intel device selection and integrated-GPU memory `auto` prefers CUDA on a mixed Nvidia/Arc box, and keep_ram_copy_of_weights is ignored on an integrated GPU. * chore(ui): typegen for the xpu device values * fix(mm): compare the integrated-GPU full-load guard against bytes still to move A resident model's weights occupy the same DRAM that vram_available is read from, so its total can exceed "available" precisely because it is loaded. lock() runs on every use and full_load_to_vram() is a no-op when resident; comparing the total refused the re-lock and evicted a healthy model on every other generation. Compare what full_load_to_vram() will actually move instead. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test: run the integrated-GPU cache tests on CPU-only torch ModelCache.__init__ sizes the RAM cache from the device's total VRAM, which on an xpu execution device reads torch.xpu.get_device_properties() -- an AssertionError on the CPU-only builds CI runs, failing 9 of these tests before they reached their subject. Stub a fixed total during construction, and add a regression test for the resident-model re-lock guard. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: LexiconCode <aaronwalker@protonmail.com> Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> | 22 天前 | |
feat: t5 encoder gguf support (#9324) * feat(model-manager): support GGUF-quantized T5 text encoders Add loading support for single-file GGUF T5 encoders (e.g. city96/t5-v1_1-xxl-encoder-gguf, llama.cpp naming), mirroring the existing Qwen3 GGUF encoder path. - Add T5Encoder_GGUF_Config (single-file, detects enc.blk.* keys + GGML tensors) and register it in the AnyModelConfig union - Add T5EncoderGGUFModel loader: remaps llama.cpp T5 keys to transformers naming, infers T5Config from tensor shapes, dequantizes token/relative-attention-bias embeddings, ties embed_tokens to shared - Work around transformers T5DenseGatedActDense casting activations to the uint8 GGML weight dtype (int8 guard doesn't cover uint8), which would corrupt the feed-forward output - Reject T5 encoders in the Qwen3 GGUF/checkpoint configs so the two stay mutually exclusive (both carry token_embd.weight; the factory resolves multi-matches from a set, so this is not order-safe) Reuse the vendored T5-XXL tokenizer instead of downloading it: move it out of Anima into a neutral invokeai/backend/t5 module shared by Anima and the GGUF loader, and update the package-data path accordingly. * Chore Typegen + Openapi * Add T5 Recalling * Add 2 gguf T5 to the Starter Models * Chore Ruff * Chore Typegen * test(t5-gguf): add unit tests for GGUF T5 loader helpers + fail-loud FFN patch guard - Add unit coverage for the pure, high-risk parts of T5EncoderGGUFModel: key remapping (_convert_t5_gguf_to_transformers), config inference (_infer_t5_config_from_state_dict), and the wo-dtype workaround. - Make _make_feed_forward_gguf_safe raise if it patches no feed-forward modules, so a future transformers class rename fails loudly at load time instead of silently corrupting encoder output. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * refactor(t5-gguf): drop dead _shape_of branch, document T5 v1.1 XXL assumptions - Inline tensor.shape in _infer_t5_config_from_state_dict: GGMLTensor.shape already returns the dequantized (logical) shape, so the _shape_of helper's fallback branch was unreachable. Remove the helper. - Document that config inference targets the T5 v1.1 XXL family and that the hardcoded architectural constants (rel-attention max distance, layer-norm epsilon, gated-gelu) are that family's defaults. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> | 1 个月前 | |
feat(nodes): extract LATENT_SCALE_FACTOR to constants.py | 2 年前 | |
fix(flux2): estimate working memory for denoise and both VAE directions (#9519) * fix(flux2): estimate working memory for denoise and both VAE directions The FLUX.2 path called model_on_device() with no working_mem_bytes anywhere, so the model cache reserved only the default device_working_mem_gb and filled the rest of the card with the model. Reference images make that fatal rather than merely tight: their latents are concatenated onto the image stream, so three 1024x1024 references quadruple the attended sequence of a 1024x1024 generation -- 6.5GB of activations against a 3GB reservation. Measured on CUDA in bf16 as peak reserved memory: transformer activations scale linearly at ~0.39 MB/token (no O(seq^2) term, SDPA) and are independent of block count; the FLUX.2 VAE costs ~2170 (decode) / ~1070 (encode) bytes per pixel per element byte, so a 1024x1024 decode peaks at ~4.3GB. Add Flux2DenoiseInvocation._estimate_working_memory() and estimate_vae_working_memory_flux2(), and pass them at every load site so the cache evicts enough to make room instead of hitting the shortfall as an OOM. Closes #9500 * fix(flux2): budget SDPA's materialized score matrix where it is real The FLUX.2 working-memory estimates were linear in the sequence length, which holds only while SDPA picks a fused kernel. That is a property of the torch build, not of FLUX.2: ROCm's fused kernels cap the head dim at 128 and reject arbitrary additive masks, so both the VAE's 512-wide mid-block head and the dense S x S bias regional prompting attaches fall through to the math fallback and materialize the score matrix -- ~17GB for a 1536px decode, and heads x S^2 for a masked forward. Rather than assume either way, ask torch: sdpa_score_matrix_bytes() queries can_use_flash/efficient/cudnn_attention for the real head dim, dtype and mask, and adds 13 bytes per score element only when no fused kernel is eligible. Measured on CUDA with SDPBackend.MATH forced: 12.9 bytes/element at 4k tokens, 10.3 at 8k, 9.7 at 16k, identical for bf16, fp16 and fp32 because the fallback's softmax intermediates are always fp32. On CUDA every shape reports fused, so the term is zero and the existing calibration is untouched. Non-CUDA devices keep the fused assumption -- torch exposes no equivalent query there, and guessing would reserve double-digit GB on no evidence. * fix(flux2): ask the real dispatcher which SDPA path a build takes The score-matrix term probed torch's CUDA eligibility helpers and read everything else as fused. That was wrong twice over: MPS has no fused SDPA kernel at all and runs the MPSGraph math transcription, so a 1024px VAE decode was admitted ~3.5GB short; and a failed probe returned "fused" too, turning "we don't know" into the one answer that can OOM. Ask `_fused_sdp_choice` instead -- the same dispatch query `scaled_dot_product_attention` runs to pick its kernel. Torch registers it for CPU, CUDA/ROCm and XPU only, so the call raises on exactly the devices that fall through to `math`, and every other failure lands on the conservative side by the same branch. Diffusers models do not reach torch's SDPA directly, so also consult `dispatch_attention_fn`'s active backend: a user on `_native_math` materializes the score matrix on hardware whose probe reports fused. Only the transformer needs this -- the FLUX.2 VAE's mid-block attention still calls SDPA itself through `AttnProcessor2_0` -- and a test pins that asymmetry. On CUDA with the stock backend every one of these terms remains zero. * fix(flux2): read the attention backend live instead of caching it once `_diffusers_attention_dispatch()` was `lru_cache`d, so the first estimate in a process pinned the answer forever. A switch to `_native_math` after that kept reserving zero for the S x S score matrix -- the exact case the lookup was added to catch. Read it live; it is a dict lookup against an already-imported module, priced once per invocation. The torch probe had the same defect one level down: its answer depends on the global SDPA kernel toggles, which `sdpa_kernel()` and `enable_flash_sdp()` flip at runtime. That probe allocates and dispatches, so it stays cached -- but keyed on the toggles, so a switch invalidates it. Per-model overrides need no plumbing: `set_attention_backend()` stamps its choice onto the process-wide registry as well as onto the model's processors, deliberately, so the estimate sees it without holding the model it is priced ahead of. A test pins that propagation. * fix(flux2): stop caching the SDPA probe and scale the VAE estimate by batch The probe's cache key held the four per-backend enable flags, but torch takes the *first eligible* backend in a priority order that `sdpa_kernel(..., set_priority=True)` reorders while leaving every flag untouched -- measured: same flags, EFFICIENT outside and MATH inside. A fused answer cached before the switch would suppress the score-matrix reservation after it. Rather than adding the priority order to the key -- the next thing to forget is always one more -- drop the cache. The probe costs ~6us against a multi-second forward, so there is nothing to protect. `vae.decode` is also handed whatever batch the latents carry, and a LatentsField is not pinned to one, so an estimate built from H and W alone gave a two-sample decode a single sample's reservation. Measured at 1024px: 4.23GB at batch 1, 7.96GB at 2, 11.89GB at 3 -- linear, slightly sub-linear per sample, so the scaled single-sample estimate stays an upper bound. The score matrix is (batch, heads, S, S) and scales with it. * fix(flux2): scale the denoise reservation by the latent batch The node had `b` in hand from preparing the latents and never passed it, so a two-sample run reserved one sample's activations and the cache admitted it to a card that could not run it. Batched latents do not come from the stock UI, but the API and custom graphs reach this node. Batch multiplies the token count and nothing else. Measured on the Klein geometry with a reduced block count: 4608 tokens at B=1 peaks at 2570MB, the same 4608 at B=2 at 5126MB, and 9728 tokens at B=1 at 5584MB -- per total token that is 0.554-0.578MB across every combination, so batch and sequence are interchangeable. Reference latents are repeated per sample by `ensure_batch_size`, so they scale too, and the score matrix is (batch, heads, S, S). The fixed base does not scale -- it covers weight casts and allocator slack -- and neither does the regional bias, built as (1, 1, S, S) and broadcast. * fix(flux2): take the reservation's batch from the blended latents `b` was read from the noise tensor, which this node builds at batch 1 from width/height/seed whenever `add_noise` is set. Batched init latents then broadcast against it in the img2img preblend, producing a two-sample `x` against a one-sample reservation. Read `x.shape[0]` instead. It is already in scope at the estimate -- the blend, the pack and the BN normalize all run above it, and none of them change the batch -- and it is the only thing that knows how many samples reach the transformer. Expanding the noise to match would have worked too, but that changes the noise, and with it the output. Note that the same `b` still feeds `generate_img_ids_flux2`, which is a correctness question rather than a memory one and is left alone here. * fix(flux2): scale the estimate by transformer width, not just token count Per-token activation cost is linear in the transformer's hidden width, and the constant was calibrated on Klein 9B (4096) but applied to every variant. FLUX.2 [dev] is 6144 and reaches this node as a first-class path, so 1024x1024 with three references reserved 7.6GB against ~10GB needed -- the same shortfall #9500 describes, on the model where partial loading makes the estimate decide residency. Measured slope between 4608 and 9216 tokens, everything else held fixed: 0.291 MB/tok at 3072, 0.386 at 4096, 0.555 at 6144 -- 0.755 / 1.00 / 1.438 against width ratios of 0.75 / 1.00 / 1.50. Scale by width, and take the head count from the same number instead of always charging the widest. Also: raise SDPA_MATH_BYTES_PER_SCORE_ELEMENT to 14, which ROCm's 13.62 at the smallest measured shape needs; drop the claim that ROCm rejects additive masks, which gfx1100 disproves; log at info when the score-matrix term fires, since it decides residency and nothing else said so; guard the warning-filter swap with a lock now that the probe runs on every estimate; and give the VAE encode node the same compute device as the decode node. * fix(flux2): raise both calibrated constants to bound the AMD measurements Two ROCm runs came in through the new calibration script and both shipped constants were under their worst point. SDPA_MATH_BYTES_PER_SCORE_ELEMENT goes 14 -> 17 (gfx1201 costs 16.38 for the shape where CUDA costs 12.88), and the per-token activation constant 0.40 -> 0.42 MB (gfx1201 measures 0.4067 at the reference width). Both are now pinned against all three platforms' measured points rather than only being self-consistent. The runs also disprove the ROCm framing this feature carried: gfx1100 reports MATH for the VAE's 512-wide head, gfx1201 reports FLASH. Two cards, same vendor, same torch, opposite answers -- which is the case for asking torch rather than hard-coding a rule, but it means the docstrings could not keep saying "ROCm caps the head dim at 128". The VAE linear constants are left alone despite measuring short on gfx1201. That run had MIOPEN_FIND_MODE=2 and a HIP allocator garbage_collection threshold set, and its series is non-monotonic above 1024px -- peak reserved falls as resolution rises, which is what a GC threshold does to this measurement. The two AMD cards are also 1.8x apart. The gap is documented where the constants are defined; the script now reports the environment and flags a non-monotonic series so the next run cannot be ambiguous about it. * fix(flux2): fit the VAE constants per convolution backend A clean ROCm run — the earlier one had MIOPEN_FIND_MODE=2, worth a uniform 1.28x, and a HIP allocator GC threshold that clipped the high-resolution points — puts the gfx1201 numbers at 3453/2688 bytes per pixel per element byte against cuDNN's 2185/1072. Flat across 512-1024px on both, so the linear model holds; only the coefficient moves. It is MIOpen's convolution workspaces, not the attention term: identical on the fused path. Shipping the MIOpen numbers everywhere would add ~60% to every cuDNN decode for nothing, so the constant follows the backend, keyed on torch.version.hip rather than the device string (a HIP build reports device.type == "cuda"). The two operations also stop sharing a ratio. "Encoding costs half of decoding" holds on cuDNN (0.49) and not on MIOpen (0.78), so it was a backend property masquerading as an architectural one. MIOPEN_FIND_MODE=2 is deliberately not budgeted for: it is not the default and would tax everyone else. Noted where the constants are defined. * fix(flux2): take the larger VAE term, not the sum, and refit MIOpen The W7900 run shows forced math and the fused path measuring identically to three decimals at every resolution, with the total flat-linear in area. The score matrix does not add to the convolution peak: the mid-block sits alone at the 8x-downsampled bottleneck, so the full-resolution feature maps are not live while it runs, and peak reserved is a high-water mark rather than a running total. Measured on cuDNN, forcing math stays *below* the fused path until 1536px and then exceeds it by 2.6GB against the 21.5GB the term prices standalone. Summing reserved 11.1GB for a 1024px gfx1100 decode that measures 6.7. Take the max: 7.0GB. A max model is weakest at the crossover, and one measured point sits there -- a 768px encode with cuDNN's constant and a materializing kernel wants 1.80GB against a 1.35GB max, reproducibly. It is not reachable as a shortfall: the cache floors every reservation at device_working_mem_gb and the whole crossover region is below it. Pinned rather than rounded away. The MIOpen decode constant also goes 3500 -> 3600; the W7900 asks for 3525 at 512px where gfx1201 asks for 3453. Its encode column agrees with gfx1201's to the byte, so this is MIOpen rather than a per-card quirk. --------- Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com> | 3 天前 | |
Add opt-in low-VRAM mode for Wan generation (#9462) * feat(video): optimize Wan memory usage * chore: openapi schema * Added additional optimizations * fix(wan): address video memory optimization review * test(cache): keep VRAM budget test CPU-safe * fix(wan): address memory optimization review * fix(wan): calibrate VAE from checkpoint files * feat(wan): add tiled VAE calibration mode * fix(wan): align VAE calibration with estimator | 21 天前 | |
Fix: z-image regional guidance split mismatch (#9273) * fix(z-image): repair regional guidance forward after diffusers refactor Z-Image Regional Guidance crashed with "split_with_sizes expects split_sizes to sum exactly to 162 ... but got split_sizes=[160]". The regional-prompting patch was a hand-copied snapshot of an outdated ZImageTransformer2DModel.forward. The installed diffusers version changed _pad_with_ids so caption pos_ids are now longer than the caption feature tensor, while the stale patch split RoPE embeddings by feature lengths instead of pos_ids lengths. Rewrite create_regional_forward to delegate to the model's own helpers (patchify_and_embed, _prepare_sequence, _build_unified_sequence) and only override the main-layer attention mask to inject the regional mask. This keeps the patch in sync with upstream diffusers and stops re-implementing the drift-prone patchify/RoPE/padding logic. * fix(z-image): repair & realign regional guidance after diffusers refactor Z-Image Regional Guidance crashed with "split_with_sizes expects split_sizes to sum exactly to 162 ... but got split_sizes=[160]". The regional-prompting patch was a hand-copied snapshot of an outdated ZImageTransformer2DModel.forward; the installed diffusers version changed _pad_with_ids so caption pos_ids are longer than the caption feature tensor, while the stale patch split RoPE embeddings by feature lengths instead of pos_ids lengths. Rewrite create_regional_forward to delegate to the model's own helpers (patchify_and_embed, _prepare_sequence, _build_unified_sequence) so it stays in sync with upstream diffusers, and only override the main-layer attention mask. Also fix two reasons regional guidance had no visible effect: - Mask alignment: the unified sequence pads the image and caption blocks individually to a multiple of 32, so the real layout is [img_real | img_pad | txt_real | txt_pad]. Scatter the four regional sub-blocks into their padding-aware positions instead of assuming a contiguous top-left block (which only matched square 1024x1024). - CFG pass: the patched forward also runs for the negative prompt; only apply the regional mask to passes whose caption length matches the positive prompt, otherwise fall back to the plain padding mask. * Chore Ruff + Typegen * fix(z-image): use identity to gate regional mask onto the positive pass The regional attention patch ran for both the conditioned and negative/CFG forward passes and distinguished them by comparing the padded caption length against the positive prompt's expected length. Two short prompts that round up to the same multiple of 32 collided, so the positive regional mask could be injected into the unconditional prediction and silently corrupt CFG. Discriminate the conditioned pass by tensor identity (cap_feats is the exact positive_cap_feats the mask was built for) instead of a length heuristic, so the positive and negative passes can never be confused. The context manager now requires positive_cap_feats whenever a regional mask is provided, turning the previously inferred invariant into an enforced one rather than a silent no-op. Also build the (bsz, 1, S, S) float mask lazily: compute applied_regional from cheap scalar checks first and skip materializing/cloning the full mask on passes that never match (every negative pass), avoiding a ~33 MB bf16 clone per call. --------- Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> | 2 个月前 | |
consolidate model manager parts into a single class | 2 年前 | |
feat(prompt): show phase + token progress for LLM prompt expansion (#9204) * feat(prompt): show phase + token progress for LLM prompt expansion Both /utilities/expand-prompt and /utilities/image-to-prompt were blocking calls with only a spinner — users had no signal whether the model was still loading or already generating, which is rough on larger LLMs. The pipelines now stream tokens via HuggingFace's TextIteratorStreamer, and the routes emit new llm_task_progress / complete / error socket events correlated to a client-supplied task_id and routed privately to the requesting user. The Expand Prompt and Image-to-Prompt popovers render a phase label ("Loading model…" / "Generating…") plus a token progress bar (current / max_tokens) while the request is in flight. * Chore Ruff * Chore openapi * fix(utilities): restore missing events binding in image_to_prompt The merge with upstream/main dropped `events = ApiDependencies.invoker.services.events` from `image_to_prompt` when inserting the new image-read access check, leaving five `events.emit_llm_task_*` calls dangling as NameErrors whenever a task_id is sent. Restore the binding to match `expand_prompt`. * fix(prompt): address review findings for LLM prompt progress Critical: prevent a deadlock when model.generate() raises in the producer thread. transformers only calls streamer.end() on the normal exit of the generation loop, so on failure (e.g. CUDA OOM) the consuming loop blocked forever, leaking the worker and pinning the model on device. Both pipelines now call streamer.end() in the worker's except block and pass a generous STREAM_TIMEOUT to the streamer as a backstop for hangs that don't raise; queue.Empty is surfaced as an error instead of blocking on thread.join(). Adds regression tests that drive generate().side_effect through the real streamer for both the text and LLaVA pipelines. Medium: fix a per-request store leak. llm_task_complete/llm_task_error now delete the store key instead of writing a terminal state, so a socket event arriving after the mutation's finally-clear can't orphan an entry. The now dead error branch of LLMTaskProgressDisplay is removed (errors surface via the RTK Query toast); the LLMTaskState union collapses to the progress case. Minor: throttle progress emissions (>=100ms) with a guaranteed final emit to bound the O(n^2) re-encode + socket cost at large max_tokens; use the generated S['LLMTask*Event'] schema types instead of hand-typed payloads; extract a shared _make_progress_callback helper to keep the two endpoints in sync; add LLaVA pipeline test coverage. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com> Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> | 1 个月前 | |
refactor: model manager v3 (#8607) * feat(mm): add UnknownModelConfig * refactor(ui): move model categorisation-ish logic to central location, simplify model manager models list * refactor(ui)refactor(ui): more cleanup of model categories * refactor(ui): remove unused excludeSubmodels I can't remember what this was for and don't see any reference to it. Maybe it's just remnants from a previous implementation? * feat(nodes): add unknown as model base * chore(ui): typegen * feat(ui): add unknown model base support in ui * feat(ui): allow changing model type in MM, fix up base and variant selects * feat(mm): omit model description instead of making it "base type filename model" * feat(app): add setting to allow unknown models * feat(ui): allow changing model format in MM * feat(app): add the installed model config to install complete events * chore(ui): typegen * feat(ui): toast warning when installed model is unidentified * docs: update config docstrings * chore(ui): typegen * tests(mm): fix test for MM, leave the UnknownModelConfig class in the list of configs * tidy(ui): prefer types from zod schemas for model attrs * chore(ui): lint * fix(ui): wrong translation string * feat(mm): normalized model storage Store models in a flat directory structure. Each model is in a dir named its unique key (a UUID). Inside that dir is either the model file or the model dir. * feat(mm): add migration to flat model storage * fix(mm): normalized multi-file/diffusers model installation no worky now worky * refactor: port MM probes to new api - Add concept of match certainty to new probe - Port CLIP Embed models to new API - Fiddle with stuff * feat(mm): port TIs to new API * tidy(mm): remove unused probes * feat(mm): port spandrel to new API * fix(mm): parsing for spandrel * fix(mm): loader for clip embed * fix(mm): tis use existing weight_files method * feat(mm): port vae to new API * fix(mm): vae class inheritance and config_path * tidy(mm): patcher types and import paths * feat(mm): better errors when invalid model config found in db * feat(mm): port t5 to new API * feat(mm): make config_path optional * refactor(mm): simplify model classification process Previously, we had a multi-phase strategy to identify models from their files on disk: 1. Run each model config classes' `matches()` method on the files. It checks if the model could possibly be an identified as the candidate model type. This was intended to be a quick check. Break on the first match. 2. If we have a match, run the config class's `parse()` method. It derive some additional model config attrs from the model files. This was intended to encapsulate heavier operations that may require loading the model into memory. 3. Derive the common model config attrs, like name, description, calculate the hash, etc. Some of these are also heavier operations. This strategy has some issues: - It is not clear how the pieces fit together. There is some back-and-forth between different methods and the config base class. It is hard to trace the flow of logic until you fully wrap your head around the system and therefore difficult to add a model architecture to the probe. - The assumption that we could do quick, lightweight checks before heavier checks is incorrect. We often _must_ load the model state dict in the `matches()` method. So there is no practical perf benefit to splitting up the responsibility of `matches()` and `parse()`. - Sometimes we need to do the same checks in `matches()` and `parse()`. In these cases, splitting the logic is has a negative perf impact because we are doing the same work twice. - As we introduce the concept of an "unknown" model config (i.e. a model that we cannot identify, but still record in the db; see #8582), we will _always_ run _all_ the checks for every model. Therefore we need not try to defer heavier checks or resource-intensive ops like hashing. We are going to do them anyways. - There are situations where a model may match multiple configs. One known case are SD pipeline models with merged LoRAs. In the old probe API, we relied on the implicit order of checks to know that if a model matched for pipeline _and_ LoRA, we prefer the pipeline match. But, in the new API, we do not have this implicit ordering of checks. To resolve this in a resilient way, we need to get all matches up front, then use tie-breaker logic to figure out which should win (or add "differential diagnosis" logic to the matchers). - Field overrides weren't handled well by this strategy. They were only applied at the very end, if a model matched successfully. This means we cannot tell the system "Hey, this model is type X with base Y. Trust me bro.". We cannot override the match logic. As we move towards letting users correct mis-identified models (see #8582), this is a requirement. We can simplify the process significantly and better support "unknown" models. Firstly, model config classes now have a single `from_model_on_disk()` method that attempts to construct an instance of the class from the model files. This replaces the `matches()` and `parse()` methods. If we fail to create the config instance, a special exception is raised that indicates why we think the files cannot be identified as the given model config class. Next, the flow for model identification is a bit simpler: - Derive all the common fields up-front (name, desc, hash, etc). - Merge in overrides. - Call `from_model_on_disk()` for every config class, passing in the fields. Overrides are handled in this method. - Record the results for each config class and choose the best one. The identification logic is a bit more verbose, with the special exceptions and handling of overrides, but it is very clear what is happening. The one downside I can think of for this strategy is we do need to check every model type, instead of stopping at the first match. It's a bit less efficient. In practice, however, this isn't a hot code path, and the improved clarity is worth far more than perf optimizations that the end user will likely never notice. * refactor(mm): remove unused methods in config.py * refactor(mm): add model config parsing utils * fix(mm): abstractmethod bork * tidy(mm): clarify that model id utils are private * fix(mm): fall back to UnknownModelConfig correctly * feat(mm): port CLIPVisionDiffusersConfig to new api * feat(mm): port SigLIPDiffusersConfig to new api * feat(mm): make match helpers more succint * feat(mm): port flux redux to new api * feat(mm): port ip adapter to new api * tidy(mm): skip optimistic override handling for now * refactor(mm): continue iterating on config * feat(mm): port flux "control lora" and t2i adapter to new api * tidy(ui): use Extract to get model config types * fix(mm): t2i base determination * feat(mm): port cnet to new api * refactor(mm): add config validation utils, make it all consistent and clean * feat(mm): wip port of main models to new api * feat(mm): wip port of main models to new api * feat(mm): wip port of main models to new api * docs(mm): add todos * tidy(mm): removed unused model merge class * feat(mm): wip port main models to new api * tidy(mm): clean up model heuristic utils * tidy(mm): clean up ModelOnDisk caching * tidy(mm): flux lora format util * refactor(mm): make config classes narrow Simpler logic to identify, less complexity to add new model, fewer useless attrs that do not relate to the model arch, etc * refactor(mm): diffusers loras w * feat(mm): consistent naming for all model config classes * fix(mm): tag generation & scattered probe fixes * tidy(mm): consistent class names * refactor(mm): split configs into separate files * docs(mm): add comments for identification utils * chore(ui): typegen * refactor(mm): remove legacy probe, new configs dir structure, update imports * fix(mm): inverted condition * docs(mm): update docsstrings in factory.py * docs(mm): document flux variant attr * feat(mm): add helper method for legacy configs * feat(mm): satisfy type checker in flux denoise * docs(mm): remove extraneous comment * fix(mm): ensure unknown model configs get unknown attrs * fix(mm): t5 identification * fix(mm): sdxl ip adapter identification * feat(mm): more flexible config matching utils * fix(mm): clip vision identification * feat(mm): add sanity checks before probing paths * docs(mm): add reminder for self for field migrations * feat(mm): clearer naming for main config class hierarchy * feat(mm): fix clip vision starter model bases, add ref to actual models * feat(mm): add model config schema migration logic * fix(mm): duplicate import * refactor(mm): split big migration into 3 Split the big migration that did all of these things into 3: - Migration 22: Remove unique contraint on base/name/type in models table - Migration 23: Migrate configs to v6.8.0 schemas - Migration 24: Normalize file storage * fix(mm): pop base/type/format when creating unknown model config * fix(db): migration 22 insert only real cols * fix(db): migration 23 fall back to unknown model when config change fails * feat(db): run migrations 23 and 24 * fix(mm): false negative on flux lora * fix(mm): vae checkpoint probe checking for dir instead of file * fix(mm): ModelOnDisk skips dirs when looking for weights Previously a path w/ any of the known weights suffixes would be seen as a weights file, even if it was a directory. We now check to ensure the candidate path is actually a file before adding it to the list of weights. * feat(mm): add method to get main model defaults from a base * feat(mm): do not log when multiple non-unknown model matches * refactor(mm): continued iteration on model identifcation * tests(mm): refactor model identification tests Overhaul of model identification (probing) tests. Previously we didn't test the correctness of probing except in a few narrow cases - now we do. See tests/model_identification/README.md for a detailed overview of the new test setup. It includes instructions for adding a new test case. In brief: - Download the model you want to add as a test case - Run a script against it to generate the test model files - Fill in the expected model type/format/base/etc in the generated test metadata JSON file Included test cases: - All starter models - A handful of other models that I had installed - Models present in the previous test cases as smoke tests, now also tested for correctness * fix(mm): omit type/format/base when creating unknown config instance * feat(mm): use ValueError for model id sanity checks * feat(mm): add flag for updating models to allow class changes * tests(mm): fix remaining MM tests * feat: allow users to edit models freely * feat(ui): add warning for model settings edit * tests(mm): flux state dict tests * tidy: remove unused file * fix(mm): lora state dict loading in model id * feat(ui): use translation string for model edit warning * docs(db): update version numbers in migration comments * chore: bump version to v6.9.0a1 * docs: update model id readme * tests(mm): attempt to fix windows model id tests * fix(mm): issue with deleting single file models * feat(mm): just delete the dir w/ rmtree when deleting model * tests(mm): windows CI issue * fix(ui): typegen schema sync * fix(mm): fixes for migration 23 - Handle CLIP Embed and Main SD models missing variant field - Handle errors when calling the discriminator function, previously only handled ValidationError but it could be a ValueError or something else - Better logging for config migration * chore: bump version to v6.9.0a2 * chore: bump version to v6.9.0a3 | 10 个月前 | |
Merge branch 'main' into ryan/spandrel-upscale | 2 年前 | |
Add tiling support to the SpoandrelImageToImage node. | 2 年前 | |
Disable thinking in LLMs so prompt expansion will work properly with thinking models (#9380) * Update text_llm_pipeline.py Fixes issue #9379 * test(text-llm): cover thinking-disabled prompt rendering; fix duplicate BOS Add hermetic coverage for enable_thinking=False using the bundled Qwen3 tokenizer and its real chat template - no network, weights or GPU needed - including the system-role retry branch, which had no coverage at all. Also stop re-adding special tokens when tokenizing a rendered chat template. The template already emits its own control tokens, so the default add_special_tokens=True duplicated BOS for Gemma- and Llama-family models. Qwen has no BOS, which is why manual testing could not surface it. Document how prompt expansion behaves with reasoning models. * test(text-llm): accept add_special_tokens in the _TinyTokenizer fake The pipeline now passes add_special_tokens explicitly, which the fake tokenizer's narrow signature rejected. It landed on main after this branch was cut, so the breakage only showed up in CI, which tests the merge with main. --------- Co-authored-by: Alexander Eichhorn <alex@eichhorn.dev> | 23 天前 | |
feat: add native Intel XPU (torch.xpu) device support (#9401) * feat(backend): add Intel XPU (torch.xpu) device support Additive xpu branches only: device selection and normalization, float16 default, VRAM queries with a passthrough-VM fallback (missing SYCL free-memory aspect), fp8 layerwise casting via a runtime probe, VAE auto-tiling, partial loading, stats/OOM handling, multi-GPU parallel session execution (device enumeration, config/API validation, worker pinning, and the generation-device options endpoint), and the auxiliary image utilities (depth/SAM/DINO pipelines accept xpu instead of falling back to CPU; cache clearing is device-agnostic). CUDA (incl. ROCm), MPS, and CPU behavior unchanged. Verified end to end on Arc Pro B70 hardware, including dual-GPU worker startup. * test(backend): add XPU coverage for TorchDevice Mock-based, mirroring the CUDA/MPS suites: device choice, dtype, normalize, the xpu_mem_get_info fallback branches, and multi-GPU generation_devices resolution/validation/labeling on XPU. Also makes the auto-without-CUDA generation-devices test hermetic on XPU machines. * build: add [xpu] extra torch 2.7.1+xpu / torchvision 0.22.1+xpu / pytorch-triton-xpu 3.3.1 from the torch-xpu index, gated to linux-x86_64 and win_amd64; uv.lock regenerated. * feat(backend): extend idle-GPU text encoder offload to XPU The idle-device arbiter and the session processor's borrow path both gated on `device.type == "cuda"`, so on a multi-XPU system no device ever registered and `offload_text_encoders_to_idle_gpus` (enabled by default) silently did nothing: encoders kept churning the denoise model in and out of VRAM. Register and lend XPU devices alongside CUDA. MPS is deliberately excluded -- it is always a single shared device, so there is never another GPU to borrow. Verified on a dual Intel Arc Pro B70 host: a text encoder node now runs on the idle GPU while the session denoises on the other ("Running compel on idle device xpu:0 (session device xpu:1)"). * feat(ui): show the executing GPU for XPU sessions Queue items already persist the executing device generically (e.g. "xpu:1"), but both readers dropped it: the session event only forwarded devices starting with "cuda", and the frontend index parser only matched /^cuda:(\d+)$/. On a multi-XPU system the progress circle and queue-item badges were therefore always blank. Accept indexed XPU devices in both places, and correct the queue-item field description, which claimed the device is set only on CUDA. * fix(mm): gate Krea 2 fp8 encoder casting on fp8 storage support The Qwen3-VL encoder kept its fp8 storage only on CUDA, so elsewhere an fp8 checkpoint was loaded as full bf16 (~8.9GB instead of ~4.4GB) and thrashed partial loading when sharing a GPU with a large transformer. Reuse the existing cached `_device_supports_fp8_storage` probe, which already backs the layerwise-casting path. It returns True unconditionally on CUDA, so CUDA behaviour is unchanged. * chore: label XPU devices by index in load logs and fp8 help text Model load lines printed the device index only for CUDA, so every model on a multi-XPU host logged as a bare "xpu device", making it impossible to tell the GPUs apart. The FP8 Storage tooltip likewise claimed CUDA-only support. * test: cover XPU config validation, progress device and fp8 probe Three paths changed by this branch had no coverage: - The `device` field pattern was untested. `test_device_choice_xpu` looks like it covers it, but the config model does not enable `validate_assignment`, so assigning `config.device` skips validation entirely; only constructing the model exercises the pattern. Added constructor-based valid/invalid cases. - `generation_devices` validation was parametrized for cuda/cpu/mps only. - The progress event's device field, which now reports XPU sessions. Also cover `_device_supports_fp8_storage`, which gates FP8 storage in both the generic layerwise-casting path and the Krea 2 encoder: CUDA answers True without probing, CPU is rejected, and a failing XPU probe returns False instead of raising. Each new test was verified to fail when the corresponding fix is reverted. * fix(nodes): recognise XPU out-of-memory errors in the Anima VAE retry The Anima VAE decode catches OOM and retries once with tiling, which caps peak allocation. Detection matched `torch.cuda.OutOfMemoryError` or the words "out of memory" in the message, so it missed XPU entirely: torch's XPU backend does not raise a recoverable `torch.OutOfMemoryError` on exhaustion, it surfaces the Level Zero/UR result code as a plain RuntimeError -- and `UR_RESULT_ERROR_OUT_OF_DEVICE_MEMORY` contains no spaces, so the existing substring never matched. The decode therefore failed outright instead of retrying tiled. Match the `*_OUT_OF_DEVICE_MEMORY` / `*_OUT_OF_HOST_MEMORY` spellings (both UR and ZE prefixes) alongside the existing conditions, and fold the cuDNN/cuBLAS checks into the same case-insensitive comparison. Extends the existing parametrized retry test with the three XPU spellings; each was verified to fail before this change. Note the driver behaviour itself is not reproducible on the hardware used here -- this stack overcommits into host RAM and hangs rather than raising -- so the tests pin the classifier, not the driver. * style: wrap long vram_usage_gb ternary for ruff * fix: drop CUDA-only wording from progress device description Matches the committed openapi/schema artifacts, which already say "on a GPU". * docs: regenerate settings data for xpu device values * fix: stop xpu VRAM probe from reporting an unknown total as zero (0, 0) made the cache's available-VRAM arithmetic collapse to a constant -working_mem budget for the life of the process. Also widen the except: the failure type moves between torch releases (RuntimeError for the missing SYCL aspect, AssertionError from _lazy_init), and warn once when the blind estimate is in use. * fix: probe fp8 support on the target device, per device, without caching failures The probe allocated via an index-less "xpu", which resolves through the thread's current XPU device rather than the device being loaded onto -- so during idle-GPU encoder offload it measured the busy denoise GPU. It was also keyed on device type, letting one device decide for another, and memoised transient failures (it runs during a load, when the device may be momentarily full) with no way back but a restart. Also probe the bf16 upcast, which is the runtime path for Krea-2/FLUX. * fix: pin torch current device when borrowing an idle GPU Worker startup set both the session device and torch's per-thread current device; the offload borrow set only the former, leaving index-less allocations on the worker's own GPU. Extracted the shared helper and guarded it on backend availability. * fix: keep idle-GPU borrows within one device type generation_devices accepts a mixed list, so a cuda session could be handed an xpu device for its text encoder. * feat: detect Intel integrated GPUs via Level Zero torch exposes no is-integrated flag, but Level Zero does (ZE_DEVICE_PROPERTY_FLAG_INTEGRATED), and its loader already ships with the torch+xpu runtime -- so no new dependency and no compiled extension. Use it to keep iGPUs out of `generation_devices: auto` when a discrete GPU exists, and to stop budgeting them as dedicated VRAM (they share system RAM, like MPS). An unknown answer keeps the previous behaviour, an iGPU-only machine keeps its device, and an explicit device list can still opt one in. * feat: add xpu torch index to pins.json Gives the launcher an Intel install option instead of requiring a manual pip install of the extra. * fix: report VRAM diagnostics for the device in use All three sites dispatched on torch.cuda.is_available() first, so a mixed NVIDIA + Arc box running on xpu reported a constant 0.0 GB and logged "CUDA Memory Allocated" -- which would make XPU bug reports unactionable. * docs: record why xpu takes the CUDA VAE constants and keeps the broad OOM needle XPU SDPA was measured on Arc Pro B70 / torch 2.13+xpu: peak memory doubles when sequence length doubles (2.00x across 2048-16384; 2.0 MB at seq=16384 vs 512 MB for a materialised score matrix). So XPU is in CUDA's O(area) regime, not ROCm's math-attention regime, and the existing constants are correct rather than accidental. * fix: derive rand_device metadata from the backend's devices Was hardcoded to 'cuda' for any non-CPU noise, which is wrong on Arc. Falls back to 'cuda' when the device query has not resolved, so Nvidia metadata is unchanged. * docs: add Intel Arc install, driver and VRAM-reporting notes * fix: probe fp8 device support only when a model requests it The probe was the first statement in _should_use_fp8, so it allocated on the GPU during the first load of any model at all -- tokenizer, VAE, scheduler -- and on API/install threads it forced XPU lazy SYCL init on a thread that never generates. Moved below the exclusions. * fix: query Level Zero Sysman for driver-global free VRAM on xpu The blind estimate (total minus this process's reserved bytes) is what made _get_vram_available over-commit on a shared GPU: it feeds a formula that assumes a driver-global figure. Sysman's zesMemoryGetState reports that figure and is often available when the SYCL ext_intel_free_memory aspect is not, so try it before estimating. Measured on Arc Pro B70 with 16 GiB held by another process: Sysman reported 15.553 GiB free, the estimate 31.725 -- a 16.172 GiB error, exactly the foreign allocation. Sysman is not a guaranteed substitute (torch's query bottoms out in the same layer), so the estimate remains as a last resort. * fix: make the fp8 probe mirror the runtime cast path The storage cast happens on CPU while params are still CPU-resident, then the fp8 tensor is copied to the device and the pre-hook upcasts there. Probing all three steps on the device would pass on a build where the host->device fp8 copy or one upcast target fails, and break at forward time instead. Verified on Arc Pro B70 / torch 2.13+xpu: the full sequence works on both cards. * fix: degrade gracefully when a backend cannot name a device torch.xpu.get_device_name goes through _lazy_init, which raises AssertionError on a build without XPU. Naming is used only for labels and logs, so fall back to the device string rather than propagating. * fix: resolve an index-less device in the Sysman VRAM query Returning None for a device with no index would skip the driver-global query and fall through to the blind estimate with no visible symptom. Callers currently always pass a concrete device, so this is a latent hazard rather than a live bug. * fix: declare ctypes prototypes for the Level Zero calls Handles come back from (c_void_p * n)() as plain Python ints, and ctypes converts an undeclared int argument to a C int -- 32 bits. Any handle above 2**31 was being silently truncated; a direct test of that path segfaults. It happened to work on the B70 because the handles fit. Also: release the idle-GPU borrow if re-pinning raises (the setup was outside the try, so a failure there stranded the lock for the life of the process), and report a failing fp8 probe once per device instead of on every model load. * fix: drop the ZES_ENABLE_SYSMAN mutation from the Sysman probe Setting a process-wide environment variable from a read-only query leaks into child processes. It also bought nothing: the variable only gates Sysman on runtimes predating zesInit and must be set before Level Zero initialises, which torch has already done by then. Verified on Arc Pro B70 that zesInit succeeds with the variable unset. * refactor: tidy up the xpu additions after a cleanup review level_zero: cache the loader so it is opened and its prototypes configured once rather than twice, share the driver/device enumeration and its ordering guard between the two probes, and collapse the Sysman pair of globals into one nullable tuple. Also: fp8 support cache is a set (it only ever stored True), the pbr_maps empty_cache is routed through TorchDevice like the PR's other conversions, the shared-memory VRAM branch stops re-testing the device type it matched on, `_auto_generation_devices` partitions in one pass, and rand_device only answers when every generation device is the same accelerator. Merges three duplicate mem_get_info tests into one parametrized case and drops two fp8 probe tests fully subsumed by the cast-sequence test. * build: pin the xpu extra to torch 2.13.0 Intel's XPU backend matured considerably after 2.7.1: torch.xpu.mem_get_info() works on driver/kernel combinations where it previously raised, and the oneAPI user-space runtime ships with the wheel, so upgrading torch upgrades it too. Follows the rocm extra, which already pins ahead of cpu/cuda. pytorch-triton-xpu was renamed triton-xpu upstream. The darwin/aarch64 fallbacks stay on 2.7.1 to match the other extras and the project's torch<2.8.0 constraint on darwin. cpu/cuda/rocm exports are unchanged package-for-package (196/211/197); the only delta is a dropped "via pytorch-triton-xpu" comment annotation from the rename. * test: stub Sysman in the unknown-total xpu probe test Without it the test only passes where Level Zero cannot answer -- never on the Intel hardware the probe exists for, where Sysman returns before the tier under test is reached. * build: teach the pins check about the xpu index Its per-platform allowlist rejects anything unlisted, so pins.json's xpu entry fails it. PyTorch publishes XPU wheels for win32 and linux x86_64, matching the extra's markers. * fix: defer the xpu device pin like cuda's torch.xpu.set_device() brings up a SYCL context that holds VRAM in an otherwise idle process, the same reason the CUDA pin waits for the first claimed queue item. * docs: regenerate settings data on linux Regenerating on Windows flips two path defaults to backslashes, which the docs check rejects. * fix(mm): handle shared memory on integrated GPUs Their VRAM is system RAM, so a RAM copy doubles each model's footprint against the same pool. Drop it, letting a full load move weights rather than copy them. Keep partial loading on -- it is the only path that respects vram_available -- and raise a clean error when a full-load-only model cannot fit, instead of walking into an uncatchable OOM-kill. Warn once when a setting is overridden. Scoped to integrated XPU; CPU and MPS are unchanged. * docs: note intel device selection and integrated-GPU memory `auto` prefers CUDA on a mixed Nvidia/Arc box, and keep_ram_copy_of_weights is ignored on an integrated GPU. * chore(ui): typegen for the xpu device values * fix(mm): compare the integrated-GPU full-load guard against bytes still to move A resident model's weights occupy the same DRAM that vram_available is read from, so its total can exceed "available" precisely because it is loaded. lock() runs on every use and full_load_to_vram() is a no-op when resident; comparing the total refused the re-lock and evicted a healthy model on every other generation. Compare what full_load_to_vram() will actually move instead. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test: run the integrated-GPU cache tests on CPU-only torch ModelCache.__init__ sizes the RAM cache from the device's total VRAM, which on an xpu execution device reads torch.xpu.get_device_properties() -- an AssertionError on the CPU-only builds CI runs, failing 9 of these tests before they reached their subject. Stub a fixed total during construction, and add a regression test for the resident-model re-lock guard. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: LexiconCode <aaronwalker@protonmail.com> Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> | 22 天前 |
| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
| 13 天前 | ||
| 1 个月前 | ||
| 1 个月前 | ||
| 3 天前 | ||
| 30 天前 | ||
| 1 个月前 | ||
| 22 天前 | ||
| 11 个月前 | ||
| 19 天前 | ||
| 2 年前 | ||
| 1 天前 | ||
| 2 年前 | ||
| 9 天前 | ||
| 24 天前 | ||
| 26 天前 | ||
| 1 个月前 | ||
| 3 个月前 | ||
| 1 年前 | ||
| 22 天前 | ||
| 1 个月前 | ||
| 2 年前 | ||
| 3 天前 | ||
| 21 天前 | ||
| 2 个月前 | ||
| 2 年前 | ||
| 1 个月前 | ||
| 10 个月前 | ||
| 2 年前 | ||
| 2 年前 | ||
| 23 天前 | ||
| 22 天前 |