Skip to content

Commit be26351

Browse files
Feature: z-image + metadata node (#8733)
## Summary Add a new "Denoise - Z-Image + Metadata" node (`ZImageDenoiseMetaInvocation`) that extends the Z-Image denoise node with metadata output for image recall functionality. This follows the same pattern as existing `denoise_latents_meta` (SD1.5/SDXL) and `flux_denoise_meta` (FLUX) nodes. **Captured metadata:** - `width` / `height` - `steps` - `guidance` (guidance_scale) - `denoising_start` / `denoising_end` - `scheduler` - `model` (transformer) - `seed` - `loras` (if applied) ## Related Issues / Discussions Enables metadata recall for Z-Image generated images, similar to existing support for SD1.5, SDXL, and FLUX models. ## QA Instructions 1. Create a workflow using the new "Denoise - Z-Image + Metadata" node 2. Connect the metadata output to a "Save Image" node 3. Generate an image 4. Check that metadata is saved with the image (visible in image info panel) 5. Verify all generation parameters are captured correctly ## Merge Plan Requires `feature/zimage-scheduler-support` #8705 branch to be merged first (base branch). ## Checklist - [x] _The PR has a short but descriptive title, suitable for a changelog_ - [ ] _Tests added / updated (if applicable)_ - [ ] _❗Changes to a redux slice have a corresponding migration_ - [ ] _Documentation added / updated (if applicable)_ - [ ] _Updated `What's New` copy (if doing a release after this PR)_
2 parents 0021404 + 384a1a6 commit be26351

File tree

4 files changed

+196
-28
lines changed

4 files changed

+196
-28
lines changed

invokeai/app/invocations/metadata_linked.py

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
)
5353
from invokeai.app.invocations.scheduler import SchedulerOutput
5454
from invokeai.app.invocations.t2i_adapter import T2IAdapterField, T2IAdapterInvocation
55+
from invokeai.app.invocations.z_image_denoise import ZImageDenoiseInvocation
5556
from invokeai.app.services.shared.invocation_context import InvocationContext
5657
from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType, SubModelType
5758
from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES
@@ -729,6 +730,52 @@ def _loras_to_json(obj: Union[Any, list[Any]]):
729730
return LatentsMetaOutput(**params, metadata=MetadataField.model_validate(md))
730731

731732

733+
@invocation(
734+
"z_image_denoise_meta",
735+
title=f"{ZImageDenoiseInvocation.UIConfig.title} + Metadata",
736+
tags=["z-image", "latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"],
737+
category="latents",
738+
version="1.0.0",
739+
)
740+
class ZImageDenoiseMetaInvocation(ZImageDenoiseInvocation, WithMetadata):
741+
"""Run denoising process with a Z-Image transformer model + metadata."""
742+
743+
def invoke(self, context: InvocationContext) -> LatentsMetaOutput:
744+
def _loras_to_json(obj: Union[Any, list[Any]]):
745+
if not isinstance(obj, list):
746+
obj = [obj]
747+
748+
output: list[dict[str, Any]] = []
749+
for item in obj:
750+
output.append(
751+
LoRAMetadataField(
752+
model=item.lora,
753+
weight=item.weight,
754+
).model_dump(exclude_none=True, exclude={"id", "type", "is_intermediate", "use_cache"})
755+
)
756+
return output
757+
758+
obj = super().invoke(context)
759+
760+
md: Dict[str, Any] = {} if self.metadata is None else self.metadata.root
761+
md.update({"width": obj.width})
762+
md.update({"height": obj.height})
763+
md.update({"steps": self.steps})
764+
md.update({"guidance": self.guidance_scale})
765+
md.update({"denoising_start": self.denoising_start})
766+
md.update({"denoising_end": self.denoising_end})
767+
md.update({"scheduler": self.scheduler})
768+
md.update({"model": self.transformer.transformer})
769+
md.update({"seed": self.seed})
770+
if len(self.transformer.loras) > 0:
771+
md.update({"loras": _loras_to_json(self.transformer.loras)})
772+
773+
params = obj.__dict__.copy()
774+
del params["type"]
775+
776+
return LatentsMetaOutput(**params, metadata=MetadataField.model_validate(md))
777+
778+
732779
@invocation(
733780
"metadata_to_vae",
734781
title="Metadata To VAE",

invokeai/frontend/web/src/features/metadata/parsing.tsx

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -725,6 +725,8 @@ const Qwen3EncoderModel: SingleMetadataHandler<ModelIdentifierField> = {
725725
return Promise.resolve(parsed);
726726
},
727727
recall: (value, store) => {
728+
// Clear conflicting Qwen3Source when setting Encoder (mutually exclusive)
729+
store.dispatch(zImageQwen3SourceModelSelected(null));
728730
store.dispatch(zImageQwen3EncoderModelSelected(value));
729731
},
730732
i18nKey: 'metadata.qwen3Encoder',
@@ -749,6 +751,8 @@ const ZImageVAEModel: SingleMetadataHandler<ModelIdentifierField> = {
749751
return Promise.resolve(parsed);
750752
},
751753
recall: (value, store) => {
754+
// Clear conflicting Qwen3Source when setting VAE (mutually exclusive)
755+
store.dispatch(zImageQwen3SourceModelSelected(null));
752756
store.dispatch(zImageVaeModelSelected(value));
753757
},
754758
i18nKey: 'metadata.vae',
@@ -773,6 +777,9 @@ const ZImageQwen3SourceModel: SingleMetadataHandler<ModelIdentifierField> = {
773777
return Promise.resolve(parsed);
774778
},
775779
recall: (value, store) => {
780+
// Clear conflicting VAE and Encoder when setting Qwen3Source (mutually exclusive)
781+
store.dispatch(zImageVaeModelSelected(null));
782+
store.dispatch(zImageQwen3EncoderModelSelected(null));
776783
store.dispatch(zImageQwen3SourceModelSelected(value));
777784
},
778785
i18nKey: 'metadata.qwen3Source',
@@ -1000,7 +1007,6 @@ export const ImageMetadataHandlers = {
10001007
CFGRescaleMultiplier,
10011008
CLIPSkip,
10021009
Guidance,
1003-
Scheduler,
10041010
Width,
10051011
Height,
10061012
Seed,
@@ -1016,6 +1022,8 @@ export const ImageMetadataHandlers = {
10161022
RefinerNegativeAestheticScore,
10171023
RefinerDenoisingStart,
10181024
MainModel,
1025+
// Scheduler must be after MainModel so that base-dependent logic (z-image scheduler) works correctly
1026+
Scheduler,
10191027
VAEModel,
10201028
Qwen3EncoderModel,
10211029
ZImageVAEModel,

invokeai/frontend/web/src/features/parameters/components/Advanced/ParamZImageQwen3VaeModelSelect.tsx

Lines changed: 13 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -17,21 +17,19 @@ import type { MainModelConfig, Qwen3EncoderModelConfig, VAEModelConfig } from 's
1717

1818
/**
1919
* Z-Image VAE Model Select - uses FLUX VAE models
20-
* Disabled when Qwen3 Source is selected (mutually exclusive)
20+
* Selecting this will clear Qwen3 Source (mutually exclusive)
2121
*/
2222
const ParamZImageVaeModelSelect = memo(() => {
2323
const dispatch = useAppDispatch();
2424
const { t } = useTranslation();
2525
const zImageVaeModel = useAppSelector(selectZImageVaeModel);
26-
const zImageQwen3SourceModel = useAppSelector(selectZImageQwen3SourceModel);
2726
const [modelConfigs, { isLoading }] = useFluxVAEModels();
2827

29-
// Disable when Qwen3 Source is selected
30-
const isDisabled = zImageQwen3SourceModel !== null;
31-
3228
const _onChange = useCallback(
3329
(model: VAEModelConfig | null) => {
3430
if (model) {
31+
// Clear conflicting Qwen3Source when setting VAE
32+
dispatch(zImageQwen3SourceModelSelected(null));
3533
dispatch(zImageVaeModelSelected(zModelIdentifierField.parse(model)));
3634
} else {
3735
dispatch(zImageVaeModelSelected(null));
@@ -48,15 +46,14 @@ const ParamZImageVaeModelSelect = memo(() => {
4846
});
4947

5048
return (
51-
<FormControl minW={0} flexGrow={1} gap={2} isDisabled={isDisabled}>
49+
<FormControl minW={0} flexGrow={1} gap={2}>
5250
<FormLabel m={0}>{t('modelManager.zImageVae')}</FormLabel>
5351
<Combobox
5452
value={value}
5553
options={options}
5654
onChange={onChange}
5755
noOptionsMessage={noOptionsMessage}
5856
isClearable
59-
isDisabled={isDisabled}
6057
placeholder={t('modelManager.zImageVaePlaceholder')}
6158
/>
6259
</FormControl>
@@ -67,21 +64,19 @@ ParamZImageVaeModelSelect.displayName = 'ParamZImageVaeModelSelect';
6764

6865
/**
6966
* Z-Image Qwen3 Encoder Model Select
70-
* Disabled when Qwen3 Source is selected (mutually exclusive)
67+
* Selecting this will clear Qwen3 Source (mutually exclusive)
7168
*/
7269
const ParamZImageQwen3EncoderModelSelect = memo(() => {
7370
const dispatch = useAppDispatch();
7471
const { t } = useTranslation();
7572
const zImageQwen3EncoderModel = useAppSelector(selectZImageQwen3EncoderModel);
76-
const zImageQwen3SourceModel = useAppSelector(selectZImageQwen3SourceModel);
7773
const [modelConfigs, { isLoading }] = useQwen3EncoderModels();
7874

79-
// Disable when Qwen3 Source is selected
80-
const isDisabled = zImageQwen3SourceModel !== null;
81-
8275
const _onChange = useCallback(
8376
(model: Qwen3EncoderModelConfig | null) => {
8477
if (model) {
78+
// Clear conflicting Qwen3Source when setting Encoder
79+
dispatch(zImageQwen3SourceModelSelected(null));
8580
dispatch(zImageQwen3EncoderModelSelected(zModelIdentifierField.parse(model)));
8681
} else {
8782
dispatch(zImageQwen3EncoderModelSelected(null));
@@ -98,15 +93,14 @@ const ParamZImageQwen3EncoderModelSelect = memo(() => {
9893
});
9994

10095
return (
101-
<FormControl minW={0} flexGrow={1} gap={2} isDisabled={isDisabled}>
96+
<FormControl minW={0} flexGrow={1} gap={2}>
10297
<FormLabel m={0}>{t('modelManager.zImageQwen3Encoder')}</FormLabel>
10398
<Combobox
10499
value={value}
105100
options={options}
106101
onChange={onChange}
107102
noOptionsMessage={noOptionsMessage}
108103
isClearable
109-
isDisabled={isDisabled}
110104
placeholder={t('modelManager.zImageQwen3EncoderPlaceholder')}
111105
/>
112106
</FormControl>
@@ -117,22 +111,20 @@ ParamZImageQwen3EncoderModelSelect.displayName = 'ParamZImageQwen3EncoderModelSe
117111

118112
/**
119113
* Z-Image Qwen3 Source Model Select - Diffusers Z-Image models for fallback
120-
* Disabled when VAE or Qwen3 Encoder is selected (mutually exclusive)
114+
* Selecting this will clear VAE and Qwen3 Encoder (mutually exclusive)
121115
*/
122116
const ParamZImageQwen3SourceModelSelect = memo(() => {
123117
const dispatch = useAppDispatch();
124118
const { t } = useTranslation();
125119
const zImageQwen3SourceModel = useAppSelector(selectZImageQwen3SourceModel);
126-
const zImageVaeModel = useAppSelector(selectZImageVaeModel);
127-
const zImageQwen3EncoderModel = useAppSelector(selectZImageQwen3EncoderModel);
128120
const [modelConfigs, { isLoading }] = useZImageDiffusersModels();
129121

130-
// Disable when VAE or Qwen3 Encoder is selected
131-
const isDisabled = zImageVaeModel !== null || zImageQwen3EncoderModel !== null;
132-
133122
const _onChange = useCallback(
134123
(model: MainModelConfig | null) => {
135124
if (model) {
125+
// Clear conflicting VAE and Encoder when setting Qwen3Source
126+
dispatch(zImageVaeModelSelected(null));
127+
dispatch(zImageQwen3EncoderModelSelected(null));
136128
dispatch(zImageQwen3SourceModelSelected(zModelIdentifierField.parse(model)));
137129
} else {
138130
dispatch(zImageQwen3SourceModelSelected(null));
@@ -149,15 +141,14 @@ const ParamZImageQwen3SourceModelSelect = memo(() => {
149141
});
150142

151143
return (
152-
<FormControl minW={0} flexGrow={1} gap={2} isDisabled={isDisabled}>
144+
<FormControl minW={0} flexGrow={1} gap={2}>
153145
<FormLabel m={0}>{t('modelManager.zImageQwen3Source')}</FormLabel>
154146
<Combobox
155147
value={value}
156148
options={options}
157149
onChange={onChange}
158150
noOptionsMessage={noOptionsMessage}
159151
isClearable
160-
isDisabled={isDisabled}
161152
placeholder={t('modelManager.zImageQwen3SourcePlaceholder')}
162153
/>
163154
</FormControl>

0 commit comments

Comments
 (0)