diff --git a/invokeai/backend/model_manager/configs/mistral_encoder.py b/invokeai/backend/model_manager/configs/mistral_encoder.py index 77dcf205820..b5b075965db 100644 --- a/invokeai/backend/model_manager/configs/mistral_encoder.py +++ b/invokeai/backend/model_manager/configs/mistral_encoder.py @@ -263,10 +263,14 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - return cls(variant=variant, **override_fields) +# Comfy-Org ``*_fp4_mixed`` files probe as checkpoints here, but the loader cannot dequantize them: +# it scales the packed uint8 instead of unpacking the two nibbles a byte holds, which is where the +# shape mismatch in issue #9565 comes from. Rejecting them at probe is a follow-up. Kept out of the +# class docstring below because pydantic publishes that as the schema description. class MistralEncoder_Checkpoint_Config(Checkpoint_Config_Base, Config_Base): """Configuration for a single-file Mistral text encoder (safetensors). - Accepts both 30-layer cow (Comfy-Org bf16/fp8/fp4) and 40-layer Mistral Small 3 + Accepts both 30-layer cow (Comfy-Org bf16/fp8) and 40-layer Mistral Small 3 (BFL canonical / upstream Mistral 3.x single-files). The loader uses the detected variant to decide whether to keep or strip the final RMSNorm. """ diff --git a/invokeai/backend/model_manager/load/model_loaders/mistral_encoder.py b/invokeai/backend/model_manager/load/model_loaders/mistral_encoder.py index b933aebd6ad..639585f719c 100644 --- a/invokeai/backend/model_manager/load/model_loaders/mistral_encoder.py +++ b/invokeai/backend/model_manager/load/model_loaders/mistral_encoder.py @@ -4,9 +4,14 @@ FLUX.2 [dev] uses BFL's 30-layer "cow-mistral3-small" distillation as its sole text encoder. The diffusers release wraps it in the multimodal ``Mistral3ForConditionalGeneration``; standalone single-file safetensors -(Comfy-Org bf16/fp8/fp4) and GGUF redistributions (gguf-org cow variants) ship +(Comfy-Org bf16/fp8) and GGUF redistributions (gguf-org cow variants) ship only the text tower, which we load as an encoder-only ``MistralModel``. +Comfy-Org's ``*_fp4_mixed`` files are deliberately not listed: FP4 packs two +values per byte, and the dequantization below multiplies the packed uint8 by the +scale instead of unpacking the nibbles, so it produces the wrong tensor (and a +shape mismatch at load - issue #9565). Rejecting them outright is a follow-up. + Both single-file packagings embed the canonical Tekken tokenizer as a U8 tensor named ``tekken_model`` (~19 MB). When ``mistral_common`` is installed we use that embedded tokenizer directly; otherwise we fall back to fetching the @@ -340,18 +345,24 @@ def _warn_if_40_layer_mistral(variant: MistralVariantType, logger: Any) -> None: "If this is NOT BFL's canonical FLUX.2-dev/text_encoder, expect degraded " "prompt adherence — upstream Mistral 3.1 / 3.2 weights (GGUFs from " "unsloth, gguf-org, etc.) are not what FLUX.2's joint attention was " - "trained against. Recommended encoders: Comfy-Org bf16/fp8/fp4 or " + "trained against. Recommended encoders: Comfy-Org bf16/fp8 or " "gguf-org cow-mistral3-small quants (all 30-layer cow distillation)." ) def _drop_quantization_metadata(sd: dict[str, Any], logger, target_dtype: torch.dtype | None = None) -> dict[str, Any]: - """Dequantize Comfy-Org-style FP8/FP4 weights and drop their metadata keys. + """Dequantize Comfy-Org-style FP8 weights and drop their metadata keys. Comfy-Org's Mistral FLUX.2 redistributions store quantized weights alongside ``*.weight_scale`` (and occasionally ``*.input_scale``) tensors. We apply the scale in-place and remove the metadata so transformers can load the result. + FP8 only. This multiplies the stored weight by its scale, which is correct for + one value per byte and wrong for FP4, where a byte holds two packed nibbles + that have to be unpacked first - the source of the shape mismatch in #9565. + ``*_fp4_mixed`` files are not supported; the starter entry for one was removed + rather than left pointing at a 12 GB download that cannot load. + Dequantization runs in fp32 for numerical accuracy, but each result is cast back down to ``target_dtype`` immediately (when provided) so the transient peak is a single fp32 weight at a time rather than the whole dict held at fp32. For a diff --git a/invokeai/backend/model_manager/starter_models.py b/invokeai/backend/model_manager/starter_models.py index eb8812f02f5..8d1c221b102 100644 --- a/invokeai/backend/model_manager/starter_models.py +++ b/invokeai/backend/model_manager/starter_models.py @@ -1221,14 +1221,6 @@ class StarterModelBundle(BaseModel): type=ModelType.MistralEncoder, ) -flux2_dev_comfy_mistral_fp4 = StarterModel( - name="FLUX.2 [dev] Mistral Encoder (Comfy FP4 mixed)", - base=BaseModelType.Any, - source="https://huggingface.co/Comfy-Org/flux2-dev/resolve/main/split_files/text_encoders/mistral_3_small_flux2_fp4_mixed.safetensors", - description="Comfy-Org FP4-mixed of BFL's 30-layer cow-mistral3-small. Smallest safetensors variant; embeds Tekken tokenizer. ~12.3GB", - type=ModelType.MistralEncoder, -) - # gguf-org cow GGUF variants (30-layer cow, llama.cpp packaging, also embed Tekken). # Lower memory footprint than the Comfy safetensors but slightly lower fidelity. flux2_dev_cow_mistral_q4 = StarterModel( @@ -2375,7 +2367,6 @@ def _gemini_3_resolution_presets( flux2_klein_qwen3_4b_encoder, flux2_klein_qwen3_8b_encoder, flux2_dev_comfy_mistral_bf16, - flux2_dev_comfy_mistral_fp4, flux2_dev_comfy_mistral_fp8, flux2_dev_cow_mistral_iq4_xs, flux2_dev_cow_mistral_q4, diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 7e9315af642..cce48362e2b 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -67212,7 +67212,7 @@ "variant" ], "title": "MistralEncoder_Checkpoint_Config", - "description": "Configuration for a single-file Mistral text encoder (safetensors).\n\nAccepts both 30-layer cow (Comfy-Org bf16/fp8/fp4) and 40-layer Mistral Small 3\n(BFL canonical / upstream Mistral 3.x single-files). The loader uses the\ndetected variant to decide whether to keep or strip the final RMSNorm." + "description": "Configuration for a single-file Mistral text encoder (safetensors).\n\nAccepts both 30-layer cow (Comfy-Org bf16/fp8) and 40-layer Mistral Small 3\n(BFL canonical / upstream Mistral 3.x single-files). The loader uses the\ndetected variant to decide whether to keep or strip the final RMSNorm." }, "MistralEncoder_Diffusers_Config": { "properties": { diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index 96741b4f075..62ff0f70dc8 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -1802,6 +1802,7 @@ "duplicateWanTransformer": "The same model is selected as both Transformer and Transformer (Low Noise). An A14B expert pair needs two different models.", "noZImageVaeSourceSelected": "No VAE source: Select VAE (FLUX) or Qwen3 Source model", "noZImageQwen3EncoderSourceSelected": "No Qwen3 Encoder source: Select Qwen3 Encoder or Qwen3 Source model", + "zImageQwen3EncoderIncompatible": "Selected Qwen3 Encoder is incompatible with Z-Image: select a Qwen3 4B encoder", "noKrea2VaeModelSelected": "Non-diffusers Krea-2: select a VAE in Advanced settings", "noKrea2Qwen3VlEncoderModelSelected": "Non-diffusers Krea-2: select a Qwen3-VL Encoder in Advanced settings", "krea2RebalanceWeightsInvalid": "Krea-2 Conditioning Rebalance weights must be exactly 12 finite comma-separated numbers", diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.test.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.test.ts index 7afecbb4c3b..3f6b43d7753 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.test.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.test.ts @@ -32,6 +32,15 @@ const mockFluxVAE = { format: 'checkpoint' as const, }; +const mockZImageQwen3Encoder = { + key: 'zimage-qwen3-4b-key', + hash: 'zimage-qwen3-4b-hash', + name: 'Z-Image Qwen3 4B Encoder', + base: 'any' as const, + type: 'qwen3_encoder' as const, + variant: 'qwen3_4b' as const, +}; + const mockAnimaMainModel = { key: 'anima-main-key', hash: 'anima-main-hash', @@ -153,7 +162,7 @@ const mockSelectQwen3VLEncoderModels = vi.fn((_state: unknown) => [mockKrea2Qwen const mockSelectZImageDiffusersModels = vi.fn((_state: unknown) => [] as unknown[]); // Z-Image borrows the FLUX.1 VAE pool - flux2 VAEs are deliberately not part of it. const mockSelectFlux1VAEModels = vi.fn((_state: unknown) => [] as unknown[]); -const mockSelectQwen3EncoderModels = vi.fn((_state: unknown) => [] as unknown[]); +const mockSelectZImageQwen3EncoderModels = vi.fn((_state: unknown) => [] as unknown[]); vi.mock('services/api/hooks/modelsByType', () => ({ selectAnimaQwen3EncoderModels: (state: unknown) => mockSelectAnimaQwen3EncoderModels(state), @@ -161,7 +170,7 @@ vi.mock('services/api/hooks/modelsByType', () => ({ selectAnimaCompatibleVAEModels: (state: unknown) => mockSelectAnimaCompatibleVAEModels(state), selectQwenImageVAEModels: (state: unknown) => mockSelectQwenImageVAEModels(state), selectQwen3VLEncoderModels: (state: unknown) => mockSelectQwen3VLEncoderModels(state), - selectQwen3EncoderModels: (state: unknown) => mockSelectQwen3EncoderModels(state), + selectZImageQwen3EncoderModels: (state: unknown) => mockSelectZImageQwen3EncoderModels(state), selectZImageDiffusersModels: (state: unknown) => mockSelectZImageDiffusersModels(state), selectFlux1VAEModels: (state: unknown) => mockSelectFlux1VAEModels(state), selectGlobalRefImageModels: vi.fn(() => []), @@ -247,6 +256,7 @@ const paramsSliceActual = (await vi.importActual('features/controlLayers/store/p animaVaeModelSelected: { type: string }; krea2VaeModelSelected: { type: string }; krea2Qwen3VlEncoderModelSelected: { type: string }; + zImageQwen3EncoderModelSelected: { type: string }; zImageQwen3SourceModelSelected: { type: string }; zImageVaeModelSelected: { type: string }; }; @@ -255,6 +265,7 @@ const { animaVaeModelSelected, krea2VaeModelSelected, krea2Qwen3VlEncoderModelSelected, + zImageQwen3EncoderModelSelected, zImageQwen3SourceModelSelected, zImageVaeModelSelected, } = paramsSliceActual; @@ -764,7 +775,7 @@ describe('modelSelected listener - Z-Image VAE defaulting', () => { mockDispatch.mockClear(); // No diffusers model installed, so the listener falls through to the encoder + VAE branch. mockSelectZImageDiffusersModels.mockReturnValue([]); - mockSelectQwen3EncoderModels.mockReturnValue([mockAnimaQwen3Encoder]); + mockSelectZImageQwen3EncoderModels.mockReturnValue([mockZImageQwen3Encoder]); mockSelectFlux1VAEModels.mockReturnValue([mockFluxVAE]); }); @@ -778,6 +789,42 @@ describe('modelSelected listener - Z-Image VAE defaulting', () => { expect(vaeDispatch!.payload).toMatchObject({ key: mockFluxVAE.key, base: 'flux' }); }); + // The encoder slot draws from the 4B pool: the general Qwen3 pool also lists Klein 9B's 8B encoder, + // and defaulting to it made the first denoise step fail with a 4096 vs 2560 shape mismatch (#9526). + it('should default the Z-Image encoder slot from the 4B pool', () => { + const state = buildMockState({ model: mockFluxMainModel }); + const action = modelSelected(zParameterModel.parse(mockZImageTurboMain)); + + capturedEffect!(action, { getState: () => state, dispatch: mockDispatch }); + + const encoderDispatch = dispatched.find( + (a) => a.type === zImageQwen3EncoderModelSelected.type && a.payload !== null + ); + // The full identifier: the slot's reducer parses with zModelIdentifierField, which silently drops a + // payload without `hash` and `type` - the slot then stayed empty despite the dispatch. + expect(encoderDispatch!.payload).toEqual({ + key: mockZImageQwen3Encoder.key, + hash: mockZImageQwen3Encoder.hash, + name: mockZImageQwen3Encoder.name, + base: mockZImageQwen3Encoder.base, + type: mockZImageQwen3Encoder.type, + }); + }); + + it('should not default the Z-Image encoder slot when the 4B pool is empty', () => { + mockSelectZImageQwen3EncoderModels.mockReturnValue([]); + + const state = buildMockState({ model: mockFluxMainModel }); + const action = modelSelected(zParameterModel.parse(mockZImageTurboMain)); + + capturedEffect!(action, { getState: () => state, dispatch: mockDispatch }); + + const encoderDispatch = dispatched.find( + (a) => a.type === zImageQwen3EncoderModelSelected.type && a.payload !== null + ); + expect(encoderDispatch).toBeUndefined(); + }); + it('should not default the Z-Image VAE slot when the FLUX.1 pool is empty', () => { mockSelectFlux1VAEModels.mockReturnValue([]); diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts index 2b9b67740af..4da890bcd55 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts @@ -67,7 +67,6 @@ import { selectAnimaVAEModels, selectFlux1VAEModels, selectGlobalRefImageModels, - selectQwen3EncoderModels, selectQwen3VLEncoderModels, selectQwenImageDiffusersModels, selectQwenImageVAEModels, @@ -77,6 +76,7 @@ import { selectWanT5EncoderModels, selectWanVAEModels, selectZImageDiffusersModels, + selectZImageQwen3EncoderModels, } from 'services/api/hooks/modelsByType'; import type { FLUXKontextModelConfig, FLUXReduxModelConfig, IPAdapterModelConfig } from 'services/api/types'; import { @@ -176,7 +176,8 @@ export const addModelSelectedListener = (startAppListening: AppStartListening) = } } else { // Fallback: try to set Qwen3 Encoder + VAE - const availableQwen3Encoders = selectQwen3EncoderModels(state); + // 4B encoders only - the 8B one Klein 9B uses is listed too, but Z-Image cannot consume it (#9526). + const availableQwen3Encoders = selectZImageQwen3EncoderModels(state); // FLUX.1 VAEs only - the Z-Image VAE picker is built from `isFlux1VAEModelConfig` and // Z-Image cannot use a FLUX.2 VAE, so a wider flux+flux2 pool would default the slot to // a model the user can neither see in the picker nor generate with. @@ -190,8 +191,10 @@ export const addModelSelectedListener = (startAppListening: AppStartListening) = dispatch( zImageQwen3EncoderModelSelected({ key: qwen3Encoder.key, + hash: qwen3Encoder.hash, name: qwen3Encoder.name, base: qwen3Encoder.base, + type: qwen3Encoder.type, }) ); } diff --git a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts index f7c8420dade..6f95038f36f 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts @@ -300,10 +300,7 @@ const slice = createSlice({ } state.zImageVaeModel = result.data; }, - zImageQwen3EncoderModelSelected: ( - state, - action: PayloadAction<{ key: string; name: string; base: string } | null> - ) => { + zImageQwen3EncoderModelSelected: (state, action: PayloadAction) => { const result = zParamsState.shape.zImageQwen3EncoderModel.safeParse(action.payload); if (!result.success) { return; diff --git a/invokeai/frontend/web/src/features/metadata/parsing.test.tsx b/invokeai/frontend/web/src/features/metadata/parsing.test.tsx index 895944c38ca..4ee658e66d6 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.test.tsx +++ b/invokeai/frontend/web/src/features/metadata/parsing.test.tsx @@ -872,6 +872,21 @@ describe('ImageMetadataHandlers — Anima / Z-Image / FLUX.1 recall gating', () ).rejects.toThrow(); }); + // Klein 9B's 8B encoder is 4096 wide; Z-Image's caption embedder takes 2560 and fails at the first + // denoise step (#9526). A workflow-built Z-Image image can still record one. + it('rejects Klein 9Bs 8B encoder', async () => { + currentBase = 'z-image'; + nextResolved = largeEncoder('qwen3_8b'); + const store = makeStore(); + + await expect( + ImageMetadataHandlers.ZImageQwen3EncoderModel.parse( + { model: fakeMain('z-image'), qwen3_encoder: nextResolved }, + store + ) + ).rejects.toThrow(); + }); + // This handler recalls into the Z-Image slots (and nulls zImageQwen3SourceModel). Anima and FLUX.2 // Klein write the same metadata field, so without the base gate they would clobber those slots. it.each(['anima', 'flux2'])('rejects when the current base is %s', async (base) => { diff --git a/invokeai/frontend/web/src/features/metadata/parsing.tsx b/invokeai/frontend/web/src/features/metadata/parsing.tsx index a653afc10b8..24402b94ec6 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.tsx +++ b/invokeai/frontend/web/src/features/metadata/parsing.tsx @@ -160,6 +160,7 @@ import { isFlux1VAEModelConfig, isFlux2VAEModelConfig, isQwen3EncoderModelConfig, + isZImageQwen3EncoderModelConfig, } from 'services/api/types'; import { assert } from 'tsafe'; import z from 'zod'; @@ -1598,18 +1599,18 @@ const ZImageQwen3EncoderModel: SingleMetadataHandler = { // Check provenance: `qwen3_encoder` is also written by Anima and FLUX.2 Klein, and this handler // clears `zImageQwen3SourceModel` on recall (review 4966712044). assertMetadataModelBase(metadata, 'z-image', 'ZImageQwen3EncoderModel'); - // The picker's domain (`useQwen3EncoderModels`): the 4B/8B encoders, i.e. everything except Anima's - // 0.6B, whose 1024-wide embeddings Z-Image cannot consume. That split lives in `variant`, so the - // full config is needed - the identifier alone cannot tell the two apart. + // The picker's domain (`useZImageQwen3EncoderModels`): the 4B encoder only. Anima's 0.6B (1024 wide) + // and Klein 9B's 8B (4096 wide) produce embeddings Z-Image cannot consume (#9526). That split lives in + // `variant`, so the full config is needed - the identifier alone cannot tell them apart. const parsed = await parseModelIdentifierMatching({ raw: getProperty(metadata, 'qwen3_encoder'), store, type: 'qwen3_encoder', - isCompatible: isQwen3EncoderModelConfig, + isCompatible: isZImageQwen3EncoderModelConfig, handlerType: 'ZImageQwen3EncoderModel', }); - // Klein and Z-Image encoders both satisfy isQwen3EncoderModelConfig, so the variant cannot separate - // those two - the currently selected base does. + // Klein 4B and Z-Image share the 4B encoder, so the variant cannot separate those two - the currently + // selected base does. const base = selectBase(store.getState()); assert(base === 'z-image', 'ZImageQwen3EncoderModel handler only works with Z-Image models'); return Promise.resolve(parsed); diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelEdit.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelEdit.tsx index 0d46fc3605a..b1554279eb7 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelEdit.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelEdit.tsx @@ -47,7 +47,7 @@ type ModelEditFormValues = UpdateModelBody & { }; const stringFieldOptions = { - validate: (value?: string | null) => (value && value.trim().length > 3) || 'Must be at least 3 characters', + validate: (value?: string | null) => (value && value.trim().length >= 3) || 'Must be at least 3 characters', }; export const ModelEdit = memo(({ modelConfig }: Props) => { @@ -142,7 +142,7 @@ export const ModelEdit = memo(({ modelConfig }: Props) => { leftIcon={} onClick={form.handleSubmit(onSubmit)} isLoading={isSubmitting} - isDisabled={Boolean(Object.keys(form.formState.errors).length)} + isDisabled={!form.formState.isValid} > {t('common.save')} diff --git a/invokeai/frontend/web/src/features/parameters/components/Advanced/ParamZImageQwen3VaeModelSelect.tsx b/invokeai/frontend/web/src/features/parameters/components/Advanced/ParamZImageQwen3VaeModelSelect.tsx index 9159a1ede02..75d9ac4bb93 100644 --- a/invokeai/frontend/web/src/features/parameters/components/Advanced/ParamZImageQwen3VaeModelSelect.tsx +++ b/invokeai/frontend/web/src/features/parameters/components/Advanced/ParamZImageQwen3VaeModelSelect.tsx @@ -12,7 +12,11 @@ import { import { type ModelIdentifierField, zModelIdentifierField } from 'features/nodes/types/common'; import { memo, useCallback } from 'react'; import { useTranslation } from 'react-i18next'; -import { useFlux1VAEModels, useQwen3EncoderModels, useZImageDiffusersModels } from 'services/api/hooks/modelsByType'; +import { + useFlux1VAEModels, + useZImageDiffusersModels, + useZImageQwen3EncoderModels, +} from 'services/api/hooks/modelsByType'; import type { MainModelConfig, Qwen3EncoderModelConfig, VAEModelConfig } from 'services/api/types'; /** @@ -70,7 +74,7 @@ const ParamZImageQwen3EncoderModelSelect = memo(() => { const dispatch = useAppDispatch(); const { t } = useTranslation(); const zImageQwen3EncoderModel = useAppSelector(selectZImageQwen3EncoderModel); - const [modelConfigs, { isLoading }] = useQwen3EncoderModels(); + const [modelConfigs, { isLoading }] = useZImageQwen3EncoderModels(); const _onChange = useCallback( (model: Qwen3EncoderModelConfig | null) => { diff --git a/invokeai/frontend/web/src/features/queue/store/readiness.test.ts b/invokeai/frontend/web/src/features/queue/store/readiness.test.ts index ca81039bf99..ea936f8b37d 100644 --- a/invokeai/frontend/web/src/features/queue/store/readiness.test.ts +++ b/invokeai/frontend/web/src/features/queue/store/readiness.test.ts @@ -294,6 +294,7 @@ const buildZImageTabArg = (overrides: { zImageVaeModel?: unknown; zImageQwen3EncoderModel?: unknown; zImageQwen3SourceModel?: unknown; + zImageQwen3EncoderConfig?: AnyModelConfig | null; }) => ({ isConnected: true, model: overrides.model ?? zImageGgufModel, @@ -306,6 +307,7 @@ const buildZImageTabArg = (overrides: { refImages: baseRefImages, loras: [], dynamicPrompts: baseDynamicPrompts, + zImageQwen3EncoderConfig: overrides.zImageQwen3EncoderConfig ?? null, hasFlux2DiffusersVaeSource: false, hasFlux2DiffusersQwen3Source: false, hasFlux2DevDiffusersSource: false, @@ -318,6 +320,9 @@ const hasZImageVaeReason = (reasons: { content: string }[]) => const hasZImageQwen3Reason = (reasons: { content: string }[]) => reasons.some((r) => r.content.includes('noZImageQwen3EncoderSourceSelected')); +const hasZImageQwen3IncompatibleReason = (reasons: { content: string }[]) => + reasons.some((r) => r.content.includes('zImageQwen3EncoderIncompatible')); + describe('Z-Image readiness checks – generate tab', () => { it('no errors when main model is a self-contained SDNQ pipeline (no component source selected)', () => { const reasons = getReasonsWhyCannotEnqueueGenerateTab(buildZImageTabArg({ model: zImageSdnqPipelineModel })); @@ -347,6 +352,32 @@ describe('Z-Image readiness checks – generate tab', () => { expect(hasZImageQwen3Reason(reasons)).toBe(false); }); + // A slot persisted before the picker was narrowed to 4B encoders can still hold Klein 9B's 8B encoder. + // The picker no longer shows it, so without this check the slot looks empty yet generation fails with a + // 4096 vs 2560 shape mismatch (#9526). + it('errors when the standalone encoder is an 8B encoder', () => { + const reasons = getReasonsWhyCannotEnqueueGenerateTab( + buildZImageTabArg({ + model: zImageGgufModel, + zImageQwen3EncoderModel: { key: 'enc' }, + zImageQwen3EncoderConfig: { key: 'enc', type: 'qwen3_encoder', variant: 'qwen3_8b' } as AnyModelConfig, + }) + ); + expect(hasZImageQwen3IncompatibleReason(reasons)).toBe(true); + expect(hasZImageQwen3Reason(reasons)).toBe(false); + }); + + it('does not error when the standalone encoder is a 4B encoder', () => { + const reasons = getReasonsWhyCannotEnqueueGenerateTab( + buildZImageTabArg({ + model: zImageGgufModel, + zImageQwen3EncoderModel: { key: 'enc' }, + zImageQwen3EncoderConfig: { key: 'enc', type: 'qwen3_encoder', variant: 'qwen3_4b' } as AnyModelConfig, + }) + ); + expect(hasZImageQwen3IncompatibleReason(reasons)).toBe(false); + }); + it('does not treat a non-pipeline SDNQ model (no submodels) as self-contained', () => { const zImageSdnqSingleFile = { ...zImageSdnqPipelineModel, @@ -368,6 +399,88 @@ describe('Z-Image readiness checks – generate tab', () => { }); }); +const buildZImageCanvasArg = (overrides: { + model?: MainModelConfig | null; + zImageQwen3EncoderModel?: unknown; + zImageQwen3EncoderConfig?: AnyModelConfig | null; +}) => ({ + ...buildCanvasTabArg({}), + model: overrides.model ?? zImageGgufModel, + params: { + ...baseParams, + zImageVaeModel: null, + zImageQwen3EncoderModel: overrides.zImageQwen3EncoderModel ?? null, + zImageQwen3SourceModel: null, + } as unknown as ParamsState, + zImageQwen3EncoderConfig: overrides.zImageQwen3EncoderConfig ?? null, +}); + +describe('Z-Image encoder-slot compatibility is independent of the main model format', () => { + // ZImageModelLoaderInvocation resolves the encoder as standalone slot -> Qwen3 Source -> + // self-contained main, so a populated slot wins even when the main ships its own encoder. Gating + // the compatibility check on "not a self-contained pipeline" let a leftover 8B slot reach an SDNQ + // pipeline main and fail at the first denoise step with 4096 vs 2560 - the #9526 crash, from the + // one class of main the check is supposed to cover. + const encoder8B = { key: 'enc', type: 'qwen3_encoder', variant: 'qwen3_8b' } as AnyModelConfig; + const encoder4B = { key: 'enc', type: 'qwen3_encoder', variant: 'qwen3_4b' } as AnyModelConfig; + + it('generate: errors when an 8B encoder slot is set alongside a self-contained SDNQ main', () => { + const reasons = getReasonsWhyCannotEnqueueGenerateTab( + buildZImageTabArg({ + model: zImageSdnqPipelineModel, + zImageQwen3EncoderModel: { key: 'enc' }, + zImageQwen3EncoderConfig: encoder8B, + }) + ); + expect(hasZImageQwen3IncompatibleReason(reasons)).toBe(true); + }); + + it('generate: no error when the slot holds a 4B encoder alongside a self-contained SDNQ main', () => { + const reasons = getReasonsWhyCannotEnqueueGenerateTab( + buildZImageTabArg({ + model: zImageSdnqPipelineModel, + zImageQwen3EncoderModel: { key: 'enc' }, + zImageQwen3EncoderConfig: encoder4B, + }) + ); + expect(hasZImageQwen3IncompatibleReason(reasons)).toBe(false); + expect(hasZImageQwen3Reason(reasons)).toBe(false); + }); + + it('canvas: errors when an 8B encoder slot is set alongside a self-contained SDNQ main', () => { + const reasons = getReasonsWhyCannotEnqueueCanvasTab( + buildZImageCanvasArg({ + model: zImageSdnqPipelineModel, + zImageQwen3EncoderModel: { key: 'enc' }, + zImageQwen3EncoderConfig: encoder8B, + }) as never + ); + expect(hasZImageQwen3IncompatibleReason(reasons)).toBe(true); + }); + + it('canvas: errors when an 8B encoder slot is set alongside a GGUF main', () => { + const reasons = getReasonsWhyCannotEnqueueCanvasTab( + buildZImageCanvasArg({ + model: zImageGgufModel, + zImageQwen3EncoderModel: { key: 'enc' }, + zImageQwen3EncoderConfig: encoder8B, + }) as never + ); + expect(hasZImageQwen3IncompatibleReason(reasons)).toBe(true); + }); + + it('canvas: no error when the slot holds a 4B encoder alongside a self-contained SDNQ main', () => { + const reasons = getReasonsWhyCannotEnqueueCanvasTab( + buildZImageCanvasArg({ + model: zImageSdnqPipelineModel, + zImageQwen3EncoderModel: { key: 'enc' }, + zImageQwen3EncoderConfig: encoder4B, + }) as never + ); + expect(hasZImageQwen3IncompatibleReason(reasons)).toBe(false); + }); +}); + describe('FLUX.2 Klein readiness checks – canvas tab', () => { it('no errors when main model is diffusers', () => { const reasons = getReasonsWhyCannotEnqueueCanvasTab(buildCanvasTabArg({ model: flux2DiffusersModel }) as never); diff --git a/invokeai/frontend/web/src/features/queue/store/readiness.ts b/invokeai/frontend/web/src/features/queue/store/readiness.ts index e1497d78e96..eaf8cde3c9e 100644 --- a/invokeai/frontend/web/src/features/queue/store/readiness.ts +++ b/invokeai/frontend/web/src/features/queue/store/readiness.ts @@ -58,6 +58,7 @@ import { isSelfContainedSDNQFlux1Pipeline, isSelfContainedSDNQPipeline, isWanSingleFileMainModelConfig, + isZImageQwen3EncoderModelConfig, } from 'services/api/types'; import { $isConnected } from 'services/events/stores'; @@ -120,6 +121,17 @@ const selectWanWiredConfigs = (state: RootState) => { }; }; +/** Resolve the standalone Z-Image encoder slot to its installed config. Null when the slot is empty or + * points at a deleted model. A slot persisted before the picker was narrowed to 4B encoders can still + * hold Klein 9B's 8B encoder, which Z-Image cannot consume (#9526). */ +const selectZImageQwen3EncoderConfig = (state: RootState) => { + const { zImageQwen3EncoderModel } = selectParamsSlice(state); + const query = selectModelConfigsQuery(state); + return zImageQwen3EncoderModel && query.data + ? (modelConfigsAdapterSelectors.selectById(query.data, zImageQwen3EncoderModel.key) ?? null) + : null; +}; + const debouncedUpdateReasons = debounce(async (arg: UpdateReasonsArg) => { const { tab, @@ -160,6 +172,7 @@ const debouncedUpdateReasons = debounce(async (arg: UpdateReasonsArg) => { hasFlux2DiffusersQwen3Source, hasFlux2DevDiffusersSource, wanWiredConfigs: selectWanWiredConfigs(store.getState()), + zImageQwen3EncoderConfig: selectZImageQwen3EncoderConfig(store.getState()), }); $reasonsWhyCannotEnqueue.set(reasons); } else if (tab === 'canvas') { @@ -188,6 +201,7 @@ const debouncedUpdateReasons = debounce(async (arg: UpdateReasonsArg) => { hasFlux2DiffusersQwen3Source, hasFlux2DevDiffusersSource, wanWiredConfigs: selectWanWiredConfigs(store.getState()), + zImageQwen3EncoderConfig: selectZImageQwen3EncoderConfig(store.getState()), }); $reasonsWhyCannotEnqueue.set(reasons); } else if (tab === 'workflows') { @@ -347,6 +361,8 @@ export const getReasonsWhyCannotEnqueueGenerateTab = (arg: { refImages: RefImagesState; loras: LoRA[]; dynamicPrompts: DynamicPromptsState; + /** Resolved config of the standalone Z-Image encoder slot - see `selectZImageQwen3EncoderConfig`. */ + zImageQwen3EncoderConfig?: AnyModelConfig | null; hasFlux2DiffusersVaeSource: boolean; hasFlux2DiffusersQwen3Source: boolean; hasFlux2DevDiffusersSource: boolean; @@ -366,6 +382,7 @@ export const getReasonsWhyCannotEnqueueGenerateTab = (arg: { refImages, loras, dynamicPrompts, + zImageQwen3EncoderConfig, hasFlux2DiffusersVaeSource, hasFlux2DiffusersQwen3Source, hasFlux2DevDiffusersSource, @@ -524,6 +541,14 @@ export const getReasonsWhyCannotEnqueueGenerateTab = (arg: { reasons.push({ content: i18n.t('parameters.invoke.noZImageQwen3EncoderSourceSelected') }); } } + // Deliberately outside the branch above: ZImageModelLoaderInvocation resolves the encoder as + // standalone slot -> Qwen3 Source -> self-contained main, so a populated slot wins even when the + // main ships its own encoder. Gating this on !mainIsSelfContainedPipeline let a leftover 8B slot + // reach an SDNQ pipeline main and fail at the first denoise step with 4096 vs 2560 - issue #9526, + // from the one class of main this check is supposed to cover. + if (zImageQwen3EncoderConfig && !isZImageQwen3EncoderModelConfig(zImageQwen3EncoderConfig)) { + reasons.push({ content: i18n.t('parameters.invoke.zImageQwen3EncoderIncompatible') }); + } // PiD decode (Z-Image reuses the FLUX decoder) needs both a PiD decoder and the Gemma-2 caption encoder. if (params.pidMode !== 'off') { if (!params.pidDecoderModel) { @@ -746,6 +771,8 @@ export const getReasonsWhyCannotEnqueueCanvasTab = (arg: { canvasIsRasterizing: boolean; canvasIsCompositing: boolean; canvasIsSelectingObject: boolean; + /** Resolved config of the standalone Z-Image encoder slot - see `selectZImageQwen3EncoderConfig`. */ + zImageQwen3EncoderConfig?: AnyModelConfig | null; hasFlux2DiffusersVaeSource: boolean; hasFlux2DiffusersQwen3Source: boolean; hasFlux2DevDiffusersSource: boolean; @@ -771,6 +798,7 @@ export const getReasonsWhyCannotEnqueueCanvasTab = (arg: { canvasIsRasterizing, canvasIsCompositing, canvasIsSelectingObject, + zImageQwen3EncoderConfig, hasFlux2DiffusersVaeSource, hasFlux2DiffusersQwen3Source, hasFlux2DevDiffusersSource, @@ -1279,6 +1307,14 @@ export const getReasonsWhyCannotEnqueueCanvasTab = (arg: { reasons.push({ content: i18n.t('parameters.invoke.noZImageQwen3EncoderSourceSelected') }); } } + // Deliberately outside the branch above: ZImageModelLoaderInvocation resolves the encoder as + // standalone slot -> Qwen3 Source -> self-contained main, so a populated slot wins even when the + // main ships its own encoder. Gating this on !mainIsSelfContainedPipeline let a leftover 8B slot + // reach an SDNQ pipeline main and fail at the first denoise step with 4096 vs 2560 - issue #9526, + // from the one class of main this check is supposed to cover. + if (zImageQwen3EncoderConfig && !isZImageQwen3EncoderModelConfig(zImageQwen3EncoderConfig)) { + reasons.push({ content: i18n.t('parameters.invoke.zImageQwen3EncoderIncompatible') }); + } // PiD decode on the Canvas: decoder + Gemma-2 encoder required, and "Scale Before Processing" must be off. if (params.pidMode !== 'off') { if (!params.pidDecoderModel) { diff --git a/invokeai/frontend/web/src/services/api/hooks/modelsByType.ts b/invokeai/frontend/web/src/services/api/hooks/modelsByType.ts index 1cbccef9290..f84fa32d385 100644 --- a/invokeai/frontend/web/src/services/api/hooks/modelsByType.ts +++ b/invokeai/frontend/web/src/services/api/hooks/modelsByType.ts @@ -48,6 +48,7 @@ import { isWanT5EncoderModelConfig, isWanVAEModelConfig, isZImageDiffusersMainModelConfig, + isZImageQwen3EncoderModelConfig, } from 'services/api/types'; const buildModelsHook = @@ -126,6 +127,7 @@ export const useQwenImageDiffusersModels = () => buildModelsHook(isQwenImageDiff export const useQwenImageVAEModels = () => buildModelsHook(isQwenImageVAEModelConfig)(); export const useQwenVLEncoderModels = () => buildModelsHook(isQwenVLEncoderModelConfig)(); export const useQwen3EncoderModels = () => buildModelsHook(isQwen3EncoderModelConfig)(); +export const useZImageQwen3EncoderModels = () => buildModelsHook(isZImageQwen3EncoderModelConfig)(); export const useQwen3VLEncoderModels = () => buildModelsHook(isQwen3VLEncoderModelConfig)(); export const useWanDiffusersModels = () => buildModelsHook(isWanDiffusersMainModelConfig)(); export const useWanSingleFileLowNoiseModels = () => buildModelsHook(isWanLowNoisePartnerOption)(); @@ -172,7 +174,7 @@ export const selectRegionalRefImageModels = buildModelsSelector( (config) => isIPAdapterModelConfig(config) || isFluxReduxModelConfig(config) ); export const selectAnimaQwen3EncoderModels = buildModelsSelector(isAnimaQwen3EncoderModelConfig); -export const selectQwen3EncoderModels = buildModelsSelector(isQwen3EncoderModelConfig); +export const selectZImageQwen3EncoderModels = buildModelsSelector(isZImageQwen3EncoderModelConfig); export const selectQwenImageDiffusersModels = buildModelsSelector(isQwenImageDiffusersMainModelConfig); export const selectQwenImageVAEModels = buildModelsSelector(isQwenImageVAEModelConfig); export const selectQwenVLEncoderModels = buildModelsSelector(isQwenVLEncoderModelConfig); diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 600a97bff25..deeb5d10dd9 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -28446,7 +28446,7 @@ export type components = { * MistralEncoder_Checkpoint_Config * @description Configuration for a single-file Mistral text encoder (safetensors). * - * Accepts both 30-layer cow (Comfy-Org bf16/fp8/fp4) and 40-layer Mistral Small 3 + * Accepts both 30-layer cow (Comfy-Org bf16/fp8) and 40-layer Mistral Small 3 * (BFL canonical / upstream Mistral 3.x single-files). The loader uses the * detected variant to decide whether to keep or strip the final RMSNorm. */ diff --git a/invokeai/frontend/web/src/services/api/types.ts b/invokeai/frontend/web/src/services/api/types.ts index 93589541742..282ffef64a4 100644 --- a/invokeai/frontend/web/src/services/api/types.ts +++ b/invokeai/frontend/web/src/services/api/types.ts @@ -432,6 +432,11 @@ export const isAnimaQwen3EncoderModelConfig = (config: AnyModelConfig): config i return config.type === 'qwen3_encoder' && config.variant === 'qwen3_06b'; }; +/** Z-Image consumes 2560-wide embeddings, so only the 4B encoder fits - Klein 9B's 8B one (4096) does not. */ +export const isZImageQwen3EncoderModelConfig = (config: AnyModelConfig): config is Qwen3EncoderModelConfig => { + return config.type === 'qwen3_encoder' && config.variant === 'qwen3_4b'; +}; + export const isMistralEncoderModelConfig = (config: AnyModelConfig): config is MistralEncoderModelConfig => { return config.type === 'mistral_encoder'; };