| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
Feat: Add PiD (Pixel Diffusion Decoder) 4× super-resolution decode for FLUX / FLUX.2 / SD3 / SDXL / Z-Image / Qwen-Image (#9281) * 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. * 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 * docs: regenerate invocation-context data for offload_from_vram * 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 --------- Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com> Co-authored-by: JPPhoto <jpollack@jpollackphoto.com> | 1 个月前 | |
Add opt-in PiD memory optimizations (#9460) * perf(pid): chunk pixel activations * perf(pid): gate memory optimizations * chore: openapi schema * refactor(pid): rename memory optimization setting * Fix generated JSON settings file * fix(pid): make the memory optimization actually free VRAM, and say what it costs Addresses every point from the review of #9460. The setting freed activation memory that the cache then withheld anyway. `estimate_pid_decode_working_memory` was not flag-aware, so a decode that peaks at 1.5 GiB reserved the unoptimized 3.9 GiB; the cache takes max(working_mem_bytes, device_working_mem_gb) and subtracts that from the weight budget, so the saving never became weight residency - it only avoided a hard OOM, on precisely the low-VRAM machines this feature exists for. The estimate now takes the flag, and each node reads the setting once and feeds both the estimate and the decode from it, so the two cannot drift apart. Calibrated against measured peaks (RTX 4090, fp32 PidNet, bf16 autocast, 4 steps, B=1): 1024px 509 MiB 1536px 934 MiB 2048px 1533 MiB which is 85.3 * U + 167 MiB, not a pure multiple of the output area - chunking bounds the per-block activations to a fixed working set, so a single scaling constant would under-reserve at small sizes or over-reserve at large ones. The constants carry ~15% headroom. Below the chunk size the pixel blocks run unchunked, so the estimate clamps to the unoptimized one rather than charging for a working set that is never allocated. The documented cost was the wrong cost. "at the cost of slower decoding" is not observable - 2.78 s either way at 2048px, median of 3 with warmup - while the cost that is real went unmentioned: the option changes the decoded image. Both parts are non-bit-exact and the few-step sampler amplifies them (~43 dB PSNR end-to-end, visually indistinguishable, not reproducible against an unoptimized decode). The setting description and both docs pages now state that, with the measured VRAM numbers. The equivalence test could not fail. It asserted `assert_close` at pixel_hidden_size=4 / BL=8 / chunk=3 on the CPU, while the shipped path is CUDA under bf16 autocast with BL in the thousands. Measured at production dimensions: CPU fp32 is bit-identical (including B=2, where chunks straddle images), CUDA fp32 differs by 9.5e-07, CUDA bf16 by 1.57e-02 - systematic, not noise, both paths being internally deterministic. So chunking is exact as *mathematics*; the divergence is cuBLAS picking different kernels for 1024-row slices. `test_pid_chunked_equivalence.py` now pins both halves: exactness on the CPU, and an absolute tolerance contract on CUDA with 3x headroom (relative tolerances are meaningless here - activations cross zero, so max|rel| reaches 1e3 on elements whose absolute error is one bf16 ULP). Two review points did not survive measurement, and are documented rather than "fixed": - The fp32 `_velocity_to_x0` branch is kept. Fusing the multiply-subtract in fp64 is bit-identical to the default expression but frees nothing (288 MiB either way at 2048px), so the 192 MiB is bought entirely with precision - max|diff| 4.8e-07 per call, 8.6% of the flag's total saving. That is a fair trade for someone who opted into trading quality for VRAM, but it is now written down, in the function, in the setting description and in the docs. - The doubled `adaLN_modulation` per chunk stays. Global attention sits between the two halves, so reusing the slices means holding them for every chunk - the full-resolution tensor the path exists to avoid (536 MiB in bf16 at 2048px). Projecting only the needed slices is no better: they are interleaved per pixel position, so selecting them gathers rows of a 1536x24576 weight, ~50 MiB per call. The recompute is ~9.9 TFLOP per 2048px decode, about 4% wall clock - and the honest source of the "slower decoding" the setting advertises. Observability: a decode with the flag on now logs the resolution, the patch-token count and whether chunking actually engaged. The setting is server-wide and never enters image metadata, so this is the only record that a given decode ran optimized, and the only feedback that a yaml-only, restart-required knob took effect at all. Tests: AST sweeps assert that every module building a `PiDDecodeConfig` forwards the flag, estimates working memory for the same mode, and reads the setting exactly once - so an eighth PiD node is covered the day it lands. Plus estimate behaviour (shrinks when enabled, keeps the fixed term, never exceeds the unoptimized estimate, still returns 0 for unsupported backbones), batch > 1 across chunk boundaries, and the two paths the flag must not reach: the discriminator feature extraction returns before the pixel blocks, and context parallelism is unreachable in this codebase (its only caller lives in a vendored class InvokeAI never instantiates) - both pinned so a refactor cannot quietly put them under the flag. All new tests mutation-verified: breaking the chunked assembly, making the estimate ignore the flag, dropping the fixed term, and omitting the flag from a node's decode or estimate each fail at least one test. tests/app + tests/backend/pid: 2202 passed. The 9 failures are the pre-existing network-dependent ones in test_model_install / test_load_api / test_download_queue. * test(pid): fix a chunking test that was both unportable and partly vacuous CI caught this on macos-default py3.11; every other job in the matrix was cancelled by fail-fast. Test-only change, no production code touched. Two separate mistakes, both mine: 1. The CPU comparison asserted `torch.equal`. That held on x86-64 with MKL and failed on macOS/Accelerate. Splitting a GEMM along its row dimension can select a micro-kernel with different K-blocking, so bit-exactness there is a property of the BLAS, not of the chunking. Only reassociation-closeness is portable. 2. Worse, and only found while investigating the first: the `batch_size=1` parametrization never entered the chunked path at all. The dispatch guard is `BL > chunk_size`, and 512px with B=1 puts BL at exactly 1024 - so it compared the unchunked path against itself and passed for the wrong reason. Verified by spying on `_forward_chunked`: zero calls. Both cases now demonstrably chunk - 768px/B=1 (BL 2304, boundaries inside one image) and 512px/B=2 (BL 2048, boundaries straddling images) - and a context manager fails the test if `_forward_chunked` is not entered, so the comparison cannot silently empty out again. Bit-equality is replaced by a signal-relative bound, calibrated rather than guessed. At these dimensions the signal is ~5.7, so one fp32 ULP is ~6.8e-07: correct code, x86-64/MKL max|diff| = 0 attention contribution off by 1e-6 max|diff| = 7.2e-07 (1.3e-07 relative, sub-ULP) attention contribution off by 1e-4 max|diff| = 1.0e-05 (1.9e-06 relative, ~15 ULP) 1e-5 relative is ~84 ULP: above any BLAS reassociation, four orders of magnitude below a structural break. The docstring states what that gives up - a uniform scaling error below ~2e-06 relative is indistinguishable from legitimate reassociation and no portable test can claim it - and what it still guards, which is the bug class that matters. Mutation-verified against realistic breakage: an off-by-one on the last chunk, a wrong `s_cond` slice, and skipping the chunked path each fail 5 of the 7 tests. 98 passed locally, ruff clean. * fix(pid): account for batch size in memory estimate * test(pid): avoid unsupported Windows CPU bf16 * fix(pid): account for autocast cache in memory estimate * docs(pid): explain working memory floor * fix(pid): correct working memory calibration * test(pid): guard against cache-term regression --------- Co-authored-by: Alexander Eichhorn <alex@eichhorn.dev> | 1 个月前 | |
Feat: Add PiD (Pixel Diffusion Decoder) 4× super-resolution decode for FLUX / FLUX.2 / SD3 / SDXL / Z-Image / Qwen-Image (#9281) * 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. * 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 * docs: regenerate invocation-context data for offload_from_vram * 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 --------- Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com> Co-authored-by: JPPhoto <jpollack@jpollackphoto.com> | 1 个月前 | |
fix(model loaders): ignore unexpected checkpoint keys, report them at debug only (#9581) * fix(model loaders): ignore unexpected checkpoint keys, report them at debug only Single-file loaders disagreed, loader by loader, about what an unexpected key from `load_state_dict(strict=False)` means: some raised, one warned, most ignored it silently, and the rest used `strict=True` and let torch raise. The ones that hard-failed turned any harmless extra tensor an exporter happened to serialize into a user-facing crash that needed a code change and a release — Anima went through this twice (#9201, #9402) for tensors the model does not need and the official checkpoint does not contain. Per the team decision on #9437, extra keys are now reported at DEBUG and otherwise ignored everywhere. New `invokeai/backend/util/state_dict_loading.py` holds the single policy: - `log_unexpected_keys()` — DEBUG only, never raises. - `load_state_dict_ignoring_extras()` — a drop-in for `strict=True` that keeps the strictness that matters (every required parameter must be filled, shape mismatches still raise) and drops the strictness that only produces whack-a-mole. - `reject_incomplete_load()` — the meta-device completeness sweep, generalized out of krea2. Stronger than `missing_keys` for models built under `init_empty_weights()`: immune to non-persistent buffers and tied weights. Every previously existing missing-key guard is preserved exactly; only the unexpected-key policy changed. `flux.py`'s bare `assert len(unexpected_keys) == 0` — which carried no message and was stripped entirely under `python -O` — is gone with it. Two consequences worth calling out: - `configs/pid_decoder.py` rejected unexpected keys at *identification* time, deliberately mirroring the loader ("both are fatal there"). Left alone, the PiD relaxation would have been unreachable and the installer would refuse a file that now loads fine. It keeps refusing non-string keys, which `load_state_dict` genuinely cannot survive. - Anima's unexpected-key `RuntimeError` was its only hard load-time guard, so it is replaced with the meta-device sweep rather than dropped — otherwise an incomplete checkpoint would fail mid-inference with "Cannot copy out of meta tensor" instead of at load time. `wan.py::_raise_for_incompatible_keys` deliberately keeps raising: Wan derivatives (Animate, S2V, Fun-Camera) are supersets whose extra branches are the feature the checkpoint exists for, not exporter noise, and it strips the benign extras before that check. Closes #9437 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01V4nLTdWBQaaDUFVLx9TyZH * docs(model loaders): fix two stale unexpected-key comments Review follow-up for #9581. The ideogram4 fp8 encoder branch and the z_image SDNQ Qwen3 loader still described unexpected checkpoint keys as fatal, but both now go through helpers that log them at DEBUG and ignore them (`load_fp8_state_dict`, `raise_on_incomplete_sdnq_load`). Update the comments to match the code; missing-key strictness is unchanged. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NxYxw5RUtuRTBRWTn79M9d --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com> | 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> | 1 个月前 |
| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
| 1 个月前 | ||
| 1 个月前 | ||
| 1 个月前 | ||
| 9 天前 | ||
| 1 个月前 |