add universal noise and optional denoiser noise inputs

Author: JPPhoto

docs/contributing/NEW_MODEL_INTEGRATION.md

@@ -1209,6 +1209,25 @@ export const NewModelSchedulerSelect = () => {
 - [ ] Frontend UI component
 - [ ] State management
 
+**External Noise:**
+- [ ] Add optional `noise: LatentsField` input to the denoise invocation
+- [ ] Validate external noise shape against the architecture's expected
+      latent shape
+- [ ] Preserve existing behavior when `noise` is not connected
+- [ ] Extend `Universal Noise` when the architecture's latent noise contract
+      can be represented there
+- [ ] Add a dedicated architecture-compatible noise invocation only when
+      `Universal Noise` cannot support the architecture cleanly
+
+If your model supports external noise, the denoise invocation should accept
+it as an optional input rather than replacing the existing seed-driven path.
+When possible, wire the architecture into `Universal Noise` instead of
+creating a separate noise node. Only create a dedicated noise invocation if
+the architecture has a noise tensor contract that `Universal Noise` cannot
+express cleanly. When external noise is connected, validate rank, channel
+count, and spatial shape before blending it with init latents or using it as
+the initial latent state.
+
 ---
 
 ## Summary: Minimal Integration

invokeai/app/invocations/anima_denoise.py

@@ -38,6 +38,7 @@
 )
 from invokeai.app.invocations.model import TransformerField
 from invokeai.app.invocations.primitives import LatentsOutput
+from invokeai.app.invocations.universal_noise import validate_noise_tensor_shape
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.anima.anima_transformer_patch import patch_anima_for_regional_prompting
 from invokeai.backend.anima.conditioning_data import AnimaRegionalTextConditioning, AnimaTextConditioning
@@ -165,7 +166,7 @@ def merge_intermediate_latents_with_init_latents(
     title="Denoise - Anima",
     tags=["image", "anima"],
     category="image",
-    version="1.2.0",
+    version="1.3.0",
     classification=Classification.Prototype,
 )
 class AnimaDenoiseInvocation(BaseInvocation):
@@ -181,6 +182,9 @@ class AnimaDenoiseInvocation(BaseInvocation):
     latents: Optional[LatentsField] = InputField(
         default=None, description=FieldDescriptions.latents, input=Input.Connection
     )
+    noise: Optional[LatentsField] = InputField(
+        default=None, description=FieldDescriptions.noise, input=Input.Connection
+    )
     # denoise_mask is used for inpainting. Only the masked region is modified.
     denoise_mask: Optional[DenoiseMaskField] = InputField(
         default=None, description=FieldDescriptions.denoise_mask, input=Input.Connection
@@ -459,7 +463,7 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
                 init_latents = init_latents.unsqueeze(2)  # [B, C, H, W] -> [B, C, 1, H, W]
 
         # Generate initial noise (3D latent: [B, C, T, H, W])
-        noise = self._get_noise(self.height, self.width, inference_dtype, device, self.seed)
+        noise = self._prepare_noise_tensor(context, inference_dtype, device)
 
         # Prepare input latents
         if init_latents is not None:
@@ -696,6 +700,16 @@ def _run_transformer(ctx: torch.Tensor, x: torch.Tensor, t: torch.Tensor) -> tor
         # Remove temporal dimension for output: [B, C, 1, H, W] -> [B, C, H, W]
         return latents.squeeze(2)
 
+    def _prepare_noise_tensor(
+        self, context: InvocationContext, inference_dtype: torch.dtype, device: torch.device
+    ) -> torch.Tensor:
+        if self.noise is not None:
+            noise = context.tensors.load(self.noise.latents_name).to(device=device, dtype=inference_dtype)
+            validate_noise_tensor_shape(noise, "Anima", self.width, self.height)
+            return noise
+
+        return self._get_noise(self.height, self.width, inference_dtype, device, self.seed)
+
     def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]:
         def step_callback(state: PipelineIntermediateState) -> None:
             context.util.sd_step_callback(state, BaseModelType.Anima)

invokeai/app/invocations/cogview4_denoise.py

@@ -20,6 +20,7 @@
 )
 from invokeai.app.invocations.model import TransformerField
 from invokeai.app.invocations.primitives import LatentsOutput
+from invokeai.app.invocations.universal_noise import validate_noise_tensor_shape
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.flux.sampling_utils import clip_timestep_schedule_fractional
 from invokeai.backend.model_manager.taxonomy import BaseModelType
@@ -34,7 +35,7 @@
     title="Denoise - CogView4",
     tags=["image", "cogview4"],
     category="image",
-    version="1.0.0",
+    version="1.1.0",
     classification=Classification.Prototype,
 )
 class CogView4DenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
@@ -44,6 +45,9 @@ class CogView4DenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
     latents: Optional[LatentsField] = InputField(
         default=None, description=FieldDescriptions.latents, input=Input.Connection
     )
+    noise: Optional[LatentsField] = InputField(
+        default=None, description=FieldDescriptions.noise, input=Input.Connection
+    )
     # denoise_mask is used for image-to-image inpainting. Only the masked region is modified.
     denoise_mask: Optional[DenoiseMaskField] = InputField(
         default=None, description=FieldDescriptions.denoise_mask, input=Input.Connection
@@ -245,15 +249,7 @@ def _run_diffusion(
         # Generate initial latent noise.
         num_channels_latents = transformer_info.model.config.in_channels  # type: ignore
         assert isinstance(num_channels_latents, int)
-        noise = self._get_noise(
-            batch_size=1,
-            num_channels_latents=num_channels_latents,
-            height=self.height,
-            width=self.width,
-            dtype=inference_dtype,
-            device=device,
-            seed=self.seed,
-        )
+        noise = self._prepare_noise_tensor(context, num_channels_latents, inference_dtype, device)
 
         # Prepare input latent image.
         if init_latents is not None:
@@ -356,6 +352,24 @@ def _run_diffusion(
 
         return latents
 
+    def _prepare_noise_tensor(
+        self, context: InvocationContext, num_channels_latents: int, inference_dtype: torch.dtype, device: torch.device
+    ) -> torch.Tensor:
+        if self.noise is not None:
+            noise = context.tensors.load(self.noise.latents_name).to(device=device, dtype=inference_dtype)
+            validate_noise_tensor_shape(noise, "CogView4", self.width, self.height, num_channels=num_channels_latents)
+            return noise
+
+        return self._get_noise(
+            batch_size=1,
+            num_channels_latents=num_channels_latents,
+            height=self.height,
+            width=self.width,
+            dtype=inference_dtype,
+            device=device,
+            seed=self.seed,
+        )
+
     def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]:
         def step_callback(state: PipelineIntermediateState) -> None:
             context.util.sd_step_callback(state, BaseModelType.CogView4)

invokeai/app/invocations/flux2_denoise.py

@@ -23,6 +23,7 @@
 )
 from invokeai.app.invocations.model import TransformerField, VAEField
 from invokeai.app.invocations.primitives import LatentsOutput
+from invokeai.app.invocations.universal_noise import validate_noise_tensor_shape
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.flux.sampling_utils import clip_timestep_schedule_fractional
 from invokeai.backend.flux.schedulers import FLUX_SCHEDULER_LABELS, FLUX_SCHEDULER_MAP, FLUX_SCHEDULER_NAME_VALUES
@@ -54,7 +55,7 @@
     title="FLUX2 Denoise",
     tags=["image", "flux", "flux2", "klein", "denoise"],
     category="image",
-    version="1.4.0",
+    version="1.5.0",
     classification=Classification.Prototype,
 )
 class Flux2DenoiseInvocation(BaseInvocation):
@@ -69,6 +70,11 @@ class Flux2DenoiseInvocation(BaseInvocation):
         description=FieldDescriptions.latents,
         input=Input.Connection,
     )
+    noise: Optional[LatentsField] = InputField(
+        default=None,
+        description=FieldDescriptions.noise,
+        input=Input.Connection,
+    )
     denoise_mask: Optional[DenoiseMaskField] = InputField(
         default=None,
         description=FieldDescriptions.denoise_mask,
@@ -240,14 +246,7 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
             init_latents = init_latents.to(device=device, dtype=inference_dtype)
 
         # Prepare input noise (FLUX.2 uses 32 channels)
-        noise = get_noise_flux2(
-            num_samples=1,
-            height=self.height,
-            width=self.width,
-            device=device,
-            dtype=inference_dtype,
-            seed=self.seed,
-        )
+        noise = self._prepare_noise_tensor(context, inference_dtype, device)
         b, _c, latent_h, latent_w = noise.shape
         packed_h = latent_h // 2
         packed_w = latent_w // 2
@@ -486,6 +485,23 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
         x = unpack_flux2(x.float(), self.height, self.width)
         return x
 
+    def _prepare_noise_tensor(
+        self, context: InvocationContext, inference_dtype: torch.dtype, device: torch.device
+    ) -> torch.Tensor:
+        if self.noise is not None:
+            noise = context.tensors.load(self.noise.latents_name).to(device=device, dtype=inference_dtype)
+            validate_noise_tensor_shape(noise, "FLUX.2", self.width, self.height)
+            return noise
+
+        return get_noise_flux2(
+            num_samples=1,
+            height=self.height,
+            width=self.width,
+            device=device,
+            dtype=inference_dtype,
+            seed=self.seed,
+        )
+
     def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor) -> Optional[torch.Tensor]:
         """Prepare the inpaint mask."""
         if self.denoise_mask is None:

invokeai/app/invocations/flux_denoise.py

@@ -28,6 +28,7 @@
 from invokeai.app.invocations.ip_adapter import IPAdapterField
 from invokeai.app.invocations.model import ControlLoRAField, LoRAField, TransformerField, VAEField
 from invokeai.app.invocations.primitives import LatentsOutput
+from invokeai.app.invocations.universal_noise import validate_noise_tensor_shape
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.flux.controlnet.instantx_controlnet_flux import InstantXControlNetFlux
 from invokeai.backend.flux.controlnet.xlabs_controlnet_flux import XLabsControlNetFlux
@@ -71,7 +72,7 @@
     title="FLUX Denoise",
     tags=["image", "flux"],
     category="image",
-    version="4.5.1",
+    version="4.6.0",
 )
 class FluxDenoiseInvocation(BaseInvocation):
     """Run denoising process with a FLUX transformer model."""
@@ -82,6 +83,11 @@ class FluxDenoiseInvocation(BaseInvocation):
         description=FieldDescriptions.latents,
         input=Input.Connection,
     )
+    noise: Optional[LatentsField] = InputField(
+        default=None,
+        description=FieldDescriptions.noise,
+        input=Input.Connection,
+    )
     # denoise_mask is used for image-to-image inpainting. Only the masked region is modified.
     denoise_mask: Optional[DenoiseMaskField] = InputField(
         default=None,
@@ -211,21 +217,15 @@ def _run_diffusion(
         context: InvocationContext,
     ):
         inference_dtype = torch.bfloat16
+        device = TorchDevice.choose_torch_device()
 
         # Load the input latents, if provided.
         init_latents = context.tensors.load(self.latents.latents_name) if self.latents else None
         if init_latents is not None:
-            init_latents = init_latents.to(device=TorchDevice.choose_torch_device(), dtype=inference_dtype)
+            init_latents = init_latents.to(device=device, dtype=inference_dtype)
 
         # Prepare input noise.
-        noise = get_noise(
-            num_samples=1,
-            height=self.height,
-            width=self.width,
-            device=TorchDevice.choose_torch_device(),
-            dtype=inference_dtype,
-            seed=self.seed,
-        )
+        noise = self._prepare_noise_tensor(context, inference_dtype, device)
         b, _c, latent_h, latent_w = noise.shape
         packed_h = latent_h // 2
         packed_w = latent_w // 2
@@ -237,7 +237,7 @@ def _run_diffusion(
             packed_height=packed_h,
             packed_width=packed_w,
             dtype=inference_dtype,
-            device=TorchDevice.choose_torch_device(),
+            device=device,
         )
         neg_text_conditionings: list[FluxTextConditioning] | None = None
         if self.negative_text_conditioning is not None:
@@ -247,14 +247,14 @@ def _run_diffusion(
                 packed_height=packed_h,
                 packed_width=packed_w,
                 dtype=inference_dtype,
-                device=TorchDevice.choose_torch_device(),
+                device=device,
             )
         redux_conditionings: list[FluxReduxConditioning] = self._load_redux_conditioning(
             context=context,
             redux_cond_field=self.redux_conditioning,
             packed_height=packed_h,
             packed_width=packed_w,
-            device=TorchDevice.choose_torch_device(),
+            device=device,
             dtype=inference_dtype,
         )
         pos_regional_prompting_extension = RegionalPromptingExtension.from_text_conditioning(
@@ -331,9 +331,7 @@ def _run_diffusion(
         img_cond: torch.Tensor | None = None
         is_flux_fill = transformer_config.variant is FluxVariantType.DevFill
         if is_flux_fill:
-            img_cond = self._prep_flux_fill_img_cond(
-                context, device=TorchDevice.choose_torch_device(), dtype=inference_dtype
-            )
+            img_cond = self._prep_flux_fill_img_cond(context, device=device, dtype=inference_dtype)
         else:
             if self.fill_conditioning is not None:
                 raise ValueError("fill_conditioning was provided, but the model is not a FLUX Fill model.")
@@ -391,7 +389,7 @@ def _run_diffusion(
                 if isinstance(self.kontext_conditioning, list)
                 else [self.kontext_conditioning],
                 vae_field=self.controlnet_vae,
-                device=TorchDevice.choose_torch_device(),
+                device=device,
                 dtype=inference_dtype,
             )
 
@@ -508,6 +506,23 @@ def _run_diffusion(
         x = unpack(x.float(), self.height, self.width)
         return x
 
+    def _prepare_noise_tensor(
+        self, context: InvocationContext, inference_dtype: torch.dtype, device: torch.device
+    ) -> torch.Tensor:
+        if self.noise is not None:
+            noise = context.tensors.load(self.noise.latents_name).to(device=device, dtype=inference_dtype)
+            validate_noise_tensor_shape(noise, "FLUX", self.width, self.height)
+            return noise
+
+        return get_noise(
+            num_samples=1,
+            height=self.height,
+            width=self.width,
+            device=device,
+            dtype=inference_dtype,
+            seed=self.seed,
+        )
+
     def _load_text_conditioning(
         self,
         context: InvocationContext,

invokeai/app/invocations/metadata_linked.py

@@ -717,7 +717,7 @@ def _loras_to_json(obj: Union[Any, list[Any]]):
         md.update({"denoising_start": self.denoising_start})
         md.update({"denoising_end": self.denoising_end})
         md.update({"model": self.transformer.transformer})
-        md.update({"seed": self.seed})
+        md.update({"seed": self.noise.seed if self.noise is not None and self.noise.seed is not None else self.seed})
         md.update({"cfg_scale": self.cfg_scale})
         md.update({"cfg_scale_start_step": self.cfg_scale_start_step})
         md.update({"cfg_scale_end_step": self.cfg_scale_end_step})
@@ -735,7 +735,7 @@ def _loras_to_json(obj: Union[Any, list[Any]]):
     title=f"{ZImageDenoiseInvocation.UIConfig.title} + Metadata",
     tags=["z-image", "latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"],
     category="latents",
-    version="1.0.0",
+    version="1.1.0",
 )
 class ZImageDenoiseMetaInvocation(ZImageDenoiseInvocation, WithMetadata):
     """Run denoising process with a Z-Image transformer model + metadata."""
@@ -766,7 +766,7 @@ def _loras_to_json(obj: Union[Any, list[Any]]):
         md.update({"denoising_end": self.denoising_end})
         md.update({"scheduler": self.scheduler})
         md.update({"model": self.transformer.transformer})
-        md.update({"seed": self.seed})
+        md.update({"seed": self.noise.seed if self.noise is not None and self.noise.seed is not None else self.seed})
         if len(self.transformer.loras) > 0:
             md.update({"loras": _loras_to_json(self.transformer.loras)})