InvokeAI 6.10.0rc1__py3-none-any.whl → 6.11.0__py3-none-any.whl

invokeai/app/invocations/z_image_denoise.py

@@ -1,3 +1,4 @@
+import inspect
 import math
 from contextlib import ExitStack
 from typing import Callable, Iterator, Optional, Tuple
@@ -5,6 +6,7 @@ from typing import Callable, Iterator, Optional, Tuple
 import einops
 import torch
 import torchvision.transforms as tv_transforms
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from PIL import Image
 from torchvision.transforms.functional import resize as tv_resize
 from tqdm import tqdm
@@ -24,6 +26,7 @@ from invokeai.app.invocations.primitives import LatentsOutput
 from invokeai.app.invocations.z_image_control import ZImageControlField
 from invokeai.app.invocations.z_image_image_to_latents import ZImageImageToLatentsInvocation
 from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.flux.schedulers import ZIMAGE_SCHEDULER_LABELS, ZIMAGE_SCHEDULER_MAP, ZIMAGE_SCHEDULER_NAME_VALUES
 from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat
 from invokeai.backend.patches.layer_patcher import LayerPatcher
 from invokeai.backend.patches.lora_conversions.z_image_lora_constants import Z_IMAGE_LORA_TRANSFORMER_PREFIX
@@ -47,7 +50,7 @@ from invokeai.backend.z_image.z_image_transformer_patch import patch_transformer
     title="Denoise - Z-Image",
     tags=["image", "z-image"],
     category="image",
-    version="1.2.0",
+    version="1.4.0",
     classification=Classification.Prototype,
 )
 class ZImageDenoiseInvocation(BaseInvocation):
@@ -66,6 +69,7 @@ class ZImageDenoiseInvocation(BaseInvocation):
     )
     denoising_start: float = InputField(default=0.0, ge=0, le=1, description=FieldDescriptions.denoising_start)
     denoising_end: float = InputField(default=1.0, ge=0, le=1, description=FieldDescriptions.denoising_end)
+    add_noise: bool = InputField(default=True, description="Add noise based on denoising start.")
     transformer: TransformerField = InputField(
         description=FieldDescriptions.z_image_model, input=Input.Connection, title="Transformer"
     )
@@ -100,6 +104,13 @@ class ZImageDenoiseInvocation(BaseInvocation):
         description=FieldDescriptions.vae + " Required for control conditioning.",
         input=Input.Connection,
     )
+    # Scheduler selection for the denoising process
+    scheduler: ZIMAGE_SCHEDULER_NAME_VALUES = InputField(
+        default="euler",
+        description="Scheduler (sampler) for the denoising process. Euler is the default and recommended for "
+        "Z-Image-Turbo. Heun is 2nd-order (better quality, 2x slower). LCM is optimized for few steps.",
+        ui_choice_labels=ZIMAGE_SCHEDULER_LABELS,
+    )
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> LatentsOutput:
@@ -337,8 +348,12 @@ class ZImageDenoiseInvocation(BaseInvocation):
 
         # Prepare input latent image
         if init_latents is not None:
-            s_0 = sigmas[0]
-            latents = s_0 * noise + (1.0 - s_0) * init_latents
+            if self.add_noise:
+                # Noise the init_latents by the appropriate amount for the first timestep.
+                s_0 = sigmas[0]
+                latents = s_0 * noise + (1.0 - s_0) * init_latents
+            else:
+                latents = init_latents
         else:
             if self.denoising_start > 1e-5:
                 raise ValueError("denoising_start should be 0 when initial latents are not provided.")
@@ -361,15 +376,32 @@ class ZImageDenoiseInvocation(BaseInvocation):
             )
 
         step_callback = self._build_step_callback(context)
-        step_callback(
-            PipelineIntermediateState(
-                step=0,
-                order=1,
-                total_steps=total_steps,
-                timestep=int(sigmas[0] * 1000),
-                latents=latents,
-            ),
-        )
+
+        # Initialize the diffusers scheduler if not using built-in Euler
+        scheduler: SchedulerMixin | None = None
+        use_scheduler = self.scheduler != "euler"
+
+        if use_scheduler:
+            scheduler_class = ZIMAGE_SCHEDULER_MAP[self.scheduler]
+            scheduler = scheduler_class(
+                num_train_timesteps=1000,
+                shift=1.0,
+            )
+            # Set timesteps - LCM should use num_inference_steps (it has its own sigma schedule),
+            # while other schedulers can use custom sigmas if supported
+            is_lcm = self.scheduler == "lcm"
+            set_timesteps_sig = inspect.signature(scheduler.set_timesteps)
+            if not is_lcm and "sigmas" in set_timesteps_sig.parameters:
+                # Convert sigmas list to tensor for scheduler
+                scheduler.set_timesteps(sigmas=sigmas, device=device)
+            else:
+                # LCM or scheduler doesn't support custom sigmas - use num_inference_steps
+                scheduler.set_timesteps(num_inference_steps=total_steps, device=device)
+
+            # For Heun scheduler, the number of actual steps may differ
+            num_scheduler_steps = len(scheduler.timesteps)
+        else:
+            num_scheduler_steps = total_steps
 
         with ExitStack() as exit_stack:
             # Get transformer config to determine if it's quantized
@@ -503,91 +535,219 @@ class ZImageDenoiseInvocation(BaseInvocation):
                 )
             )
 
-            # Denoising loop
-            for step_idx in tqdm(range(total_steps)):
-                sigma_curr = sigmas[step_idx]
-                sigma_prev = sigmas[step_idx + 1]
-
-                # Timestep tensor for Z-Image model
-                # The model expects t=0 at start (noise) and t=1 at end (clean)
-                # Sigma goes from 1 (noise) to 0 (clean), so model_t = 1 - sigma
-                model_t = 1.0 - sigma_curr
-                timestep = torch.tensor([model_t], device=device, dtype=inference_dtype).expand(latents.shape[0])
-
-                # Run transformer for positive prediction
-                # Z-Image transformer expects: x as list of [C, 1, H, W] tensors, t, cap_feats as list
-                # Prepare latent input: [B, C, H, W] -> [B, C, 1, H, W] -> list of [C, 1, H, W]
-                latent_model_input = latents.to(transformer.dtype)
-                latent_model_input = latent_model_input.unsqueeze(2)  # Add frame dimension
-                latent_model_input_list = list(latent_model_input.unbind(dim=0))
-
-                # Determine if control should be applied at this step
-                apply_control = control_extension is not None and control_extension.should_apply(step_idx, total_steps)
-
-                # Run forward pass - use custom forward with control if extension is active
-                if apply_control:
-                    model_out_list, _ = z_image_forward_with_control(
-                        transformer=transformer,
-                        x=latent_model_input_list,
-                        t=timestep,
-                        cap_feats=[pos_prompt_embeds],
-                        control_extension=control_extension,
-                    )
-                else:
-                    model_output = transformer(
-                        x=latent_model_input_list,
-                        t=timestep,
-                        cap_feats=[pos_prompt_embeds],
+            # Denoising loop - supports both built-in Euler and diffusers schedulers
+            # Track user-facing step for progress (accounts for Heun's double steps)
+            user_step = 0
+
+            if use_scheduler and scheduler is not None:
+                # Use diffusers scheduler for stepping
+                # Use tqdm with total_steps (user-facing steps) not num_scheduler_steps (internal steps)
+                # This ensures progress bar shows 1/8, 2/8, etc. even when scheduler uses more internal steps
+                pbar = tqdm(total=total_steps, desc="Denoising")
+                for step_index in range(num_scheduler_steps):
+                    sched_timestep = scheduler.timesteps[step_index]
+                    # Convert scheduler timestep (0-1000) to normalized sigma (0-1)
+                    sigma_curr = sched_timestep.item() / scheduler.config.num_train_timesteps
+
+                    # For Heun scheduler, track if we're in first or second order step
+                    is_heun = hasattr(scheduler, "state_in_first_order")
+                    in_first_order = scheduler.state_in_first_order if is_heun else True
+
+                    # Timestep tensor for Z-Image model
+                    # The model expects t=0 at start (noise) and t=1 at end (clean)
+                    model_t = 1.0 - sigma_curr
+                    timestep = torch.tensor([model_t], device=device, dtype=inference_dtype).expand(latents.shape[0])
+
+                    # Run transformer for positive prediction
+                    latent_model_input = latents.to(transformer.dtype)
+                    latent_model_input = latent_model_input.unsqueeze(2)  # Add frame dimension
+                    latent_model_input_list = list(latent_model_input.unbind(dim=0))
+
+                    # Determine if control should be applied at this step
+                    apply_control = control_extension is not None and control_extension.should_apply(
+                        user_step, total_steps
                     )
-                    model_out_list = model_output[0]  # Extract list of tensors from tuple
 
-                noise_pred_cond = torch.stack([t.float() for t in model_out_list], dim=0)
-                noise_pred_cond = noise_pred_cond.squeeze(2)  # Remove frame dimension
-                noise_pred_cond = -noise_pred_cond  # Z-Image uses v-prediction with negation
+                    # Run forward pass
+                    if apply_control:
+                        model_out_list, _ = z_image_forward_with_control(
+                            transformer=transformer,
+                            x=latent_model_input_list,
+                            t=timestep,
+                            cap_feats=[pos_prompt_embeds],
+                            control_extension=control_extension,
+                        )
+                    else:
+                        model_output = transformer(
+                            x=latent_model_input_list,
+                            t=timestep,
+                            cap_feats=[pos_prompt_embeds],
+                        )
+                        model_out_list = model_output[0]
+
+                    noise_pred_cond = torch.stack([t.float() for t in model_out_list], dim=0)
+                    noise_pred_cond = noise_pred_cond.squeeze(2)
+                    noise_pred_cond = -noise_pred_cond  # Z-Image uses v-prediction with negation
+
+                    # Apply CFG if enabled
+                    if do_classifier_free_guidance and neg_prompt_embeds is not None:
+                        if apply_control:
+                            model_out_list_uncond, _ = z_image_forward_with_control(
+                                transformer=transformer,
+                                x=latent_model_input_list,
+                                t=timestep,
+                                cap_feats=[neg_prompt_embeds],
+                                control_extension=control_extension,
+                            )
+                        else:
+                            model_output_uncond = transformer(
+                                x=latent_model_input_list,
+                                t=timestep,
+                                cap_feats=[neg_prompt_embeds],
+                            )
+                            model_out_list_uncond = model_output_uncond[0]
+
+                        noise_pred_uncond = torch.stack([t.float() for t in model_out_list_uncond], dim=0)
+                        noise_pred_uncond = noise_pred_uncond.squeeze(2)
+                        noise_pred_uncond = -noise_pred_uncond
+                        noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond)
+                    else:
+                        noise_pred = noise_pred_cond
+
+                    # Use scheduler.step() for the update
+                    step_output = scheduler.step(model_output=noise_pred, timestep=sched_timestep, sample=latents)
+                    latents = step_output.prev_sample
 
-                # Apply CFG if enabled
-                if do_classifier_free_guidance and neg_prompt_embeds is not None:
+                    # Get sigma_prev for inpainting (next sigma value)
+                    if step_index + 1 < len(scheduler.sigmas):
+                        sigma_prev = scheduler.sigmas[step_index + 1].item()
+                    else:
+                        sigma_prev = 0.0
+
+                    if inpaint_extension is not None:
+                        latents = inpaint_extension.merge_intermediate_latents_with_init_latents(latents, sigma_prev)
+
+                    # For Heun, only increment user step after second-order step completes
+                    if is_heun:
+                        if not in_first_order:
+                            user_step += 1
+                            # Only call step_callback if we haven't exceeded total_steps
+                            if user_step <= total_steps:
+                                pbar.update(1)
+                                step_callback(
+                                    PipelineIntermediateState(
+                                        step=user_step,
+                                        order=2,
+                                        total_steps=total_steps,
+                                        timestep=int(sigma_curr * 1000),
+                                        latents=latents,
+                                    ),
+                                )
+                    else:
+                        # For LCM and other first-order schedulers
+                        user_step += 1
+                        # Only call step_callback if we haven't exceeded total_steps
+                        # (LCM scheduler may have more internal steps than user-facing steps)
+                        if user_step <= total_steps:
+                            pbar.update(1)
+                            step_callback(
+                                PipelineIntermediateState(
+                                    step=user_step,
+                                    order=1,
+                                    total_steps=total_steps,
+                                    timestep=int(sigma_curr * 1000),
+                                    latents=latents,
+                                ),
+                            )
+                pbar.close()
+            else:
+                # Original Euler implementation (default, optimized for Z-Image)
+                for step_idx in tqdm(range(total_steps)):
+                    sigma_curr = sigmas[step_idx]
+                    sigma_prev = sigmas[step_idx + 1]
+
+                    # Timestep tensor for Z-Image model
+                    # The model expects t=0 at start (noise) and t=1 at end (clean)
+                    # Sigma goes from 1 (noise) to 0 (clean), so model_t = 1 - sigma
+                    model_t = 1.0 - sigma_curr
+                    timestep = torch.tensor([model_t], device=device, dtype=inference_dtype).expand(latents.shape[0])
+
+                    # Run transformer for positive prediction
+                    # Z-Image transformer expects: x as list of [C, 1, H, W] tensors, t, cap_feats as list
+                    # Prepare latent input: [B, C, H, W] -> [B, C, 1, H, W] -> list of [C, 1, H, W]
+                    latent_model_input = latents.to(transformer.dtype)
+                    latent_model_input = latent_model_input.unsqueeze(2)  # Add frame dimension
+                    latent_model_input_list = list(latent_model_input.unbind(dim=0))
+
+                    # Determine if control should be applied at this step
+                    apply_control = control_extension is not None and control_extension.should_apply(
+                        step_idx, total_steps
+                    )
+
+                    # Run forward pass - use custom forward with control if extension is active
                     if apply_control:
-                        model_out_list_uncond, _ = z_image_forward_with_control(
+                        model_out_list, _ = z_image_forward_with_control(
                             transformer=transformer,
                             x=latent_model_input_list,
                             t=timestep,
-                            cap_feats=[neg_prompt_embeds],
+                            cap_feats=[pos_prompt_embeds],
                             control_extension=control_extension,
                         )
                     else:
-                        model_output_uncond = transformer(
+                        model_output = transformer(
                             x=latent_model_input_list,
                             t=timestep,
-                            cap_feats=[neg_prompt_embeds],
+                            cap_feats=[pos_prompt_embeds],
                         )
-                        model_out_list_uncond = model_output_uncond[0]  # Extract list of tensors from tuple
-
-                    noise_pred_uncond = torch.stack([t.float() for t in model_out_list_uncond], dim=0)
-                    noise_pred_uncond = noise_pred_uncond.squeeze(2)
-                    noise_pred_uncond = -noise_pred_uncond
-                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond)
-                else:
-                    noise_pred = noise_pred_cond
-
-                # Euler step
-                latents_dtype = latents.dtype
-                latents = latents.to(dtype=torch.float32)
-                latents = latents + (sigma_prev - sigma_curr) * noise_pred
-                latents = latents.to(dtype=latents_dtype)
-
-                if inpaint_extension is not None:
-                    latents = inpaint_extension.merge_intermediate_latents_with_init_latents(latents, sigma_prev)
-
-                step_callback(
-                    PipelineIntermediateState(
-                        step=step_idx + 1,
-                        order=1,
-                        total_steps=total_steps,
-                        timestep=int(sigma_curr * 1000),
-                        latents=latents,
-                    ),
-                )
+                        model_out_list = model_output[0]  # Extract list of tensors from tuple
+
+                    noise_pred_cond = torch.stack([t.float() for t in model_out_list], dim=0)
+                    noise_pred_cond = noise_pred_cond.squeeze(2)  # Remove frame dimension
+                    noise_pred_cond = -noise_pred_cond  # Z-Image uses v-prediction with negation
+
+                    # Apply CFG if enabled
+                    if do_classifier_free_guidance and neg_prompt_embeds is not None:
+                        if apply_control:
+                            model_out_list_uncond, _ = z_image_forward_with_control(
+                                transformer=transformer,
+                                x=latent_model_input_list,
+                                t=timestep,
+                                cap_feats=[neg_prompt_embeds],
+                                control_extension=control_extension,
+                            )
+                        else:
+                            model_output_uncond = transformer(
+                                x=latent_model_input_list,
+                                t=timestep,
+                                cap_feats=[neg_prompt_embeds],
+                            )
+                            model_out_list_uncond = model_output_uncond[0]  # Extract list of tensors from tuple
+
+                        noise_pred_uncond = torch.stack([t.float() for t in model_out_list_uncond], dim=0)
+                        noise_pred_uncond = noise_pred_uncond.squeeze(2)
+                        noise_pred_uncond = -noise_pred_uncond
+                        noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond)
+                    else:
+                        noise_pred = noise_pred_cond
+
+                    # Euler step
+                    latents_dtype = latents.dtype
+                    latents = latents.to(dtype=torch.float32)
+                    latents = latents + (sigma_prev - sigma_curr) * noise_pred
+                    latents = latents.to(dtype=latents_dtype)
+
+                    if inpaint_extension is not None:
+                        latents = inpaint_extension.merge_intermediate_latents_with_init_latents(latents, sigma_prev)
+
+                    step_callback(
+                        PipelineIntermediateState(
+                            step=step_idx + 1,
+                            order=1,
+                            total_steps=total_steps,
+                            timestep=int(sigma_curr * 1000),
+                            latents=latents,
+                        ),
+                    )
 
         return latents
 

invokeai/app/invocations/z_image_image_to_latents.py

@@ -20,6 +20,7 @@ from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEnc
 from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux
 
 # Z-Image can use either the Diffusers AutoencoderKL or the FLUX AutoEncoder
 ZImageVAE = Union[AutoencoderKL, FluxAutoEncoder]
@@ -47,7 +48,14 @@ class ZImageImageToLatentsInvocation(BaseInvocation, WithMetadata, WithBoard):
                 "Ensure you are using a compatible VAE model."
             )
 
-        with vae_info.model_on_device() as (_, vae):
+        # Estimate working memory needed for VAE encode
+        estimated_working_memory = estimate_vae_working_memory_flux(
+            operation="encode",
+            image_tensor=image_tensor,
+            vae=vae_info.model,
+        )
+
+        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
             if not isinstance(vae, (AutoencoderKL, FluxAutoEncoder)):
                 raise TypeError(
                     f"Expected AutoencoderKL or FluxAutoEncoder, got {type(vae).__name__}. "

invokeai/app/invocations/z_image_latents_to_image.py

@@ -21,6 +21,7 @@ from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder
 from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux
 
 # Z-Image can use either the Diffusers AutoencoderKL or the FLUX AutoEncoder
 ZImageVAE = Union[AutoencoderKL, FluxAutoEncoder]
@@ -53,12 +54,19 @@ class ZImageLatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
 
         is_flux_vae = isinstance(vae_info.model, FluxAutoEncoder)
 
+        # Estimate working memory needed for VAE decode
+        estimated_working_memory = estimate_vae_working_memory_flux(
+            operation="decode",
+            image_tensor=latents,
+            vae=vae_info.model,
+        )
+
         # FLUX VAE doesn't support seamless, so only apply for AutoencoderKL
         seamless_context = (
             nullcontext() if is_flux_vae else SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes)
         )
 
-        with seamless_context, vae_info.model_on_device() as (_, vae):
+        with seamless_context, vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
             context.util.signal_progress("Running VAE")
             if not isinstance(vae, (AutoencoderKL, FluxAutoEncoder)):
                 raise TypeError(

invokeai/app/invocations/z_image_seed_variance_enhancer.py

@@ -0,0 +1,110 @@
+import torch
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.fields import (
+    FieldDescriptions,
+    Input,
+    InputField,
+    ZImageConditioningField,
+)
+from invokeai.app.invocations.primitives import ZImageConditioningOutput
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
+    ConditioningFieldData,
+    ZImageConditioningInfo,
+)
+
+
+@invocation(
+    "z_image_seed_variance_enhancer",
+    title="Seed Variance Enhancer - Z-Image",
+    tags=["conditioning", "z-image", "variance", "seed"],
+    category="conditioning",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class ZImageSeedVarianceEnhancerInvocation(BaseInvocation):
+    """Adds seed-based noise to Z-Image conditioning to increase variance between seeds.
+
+    Z-Image-Turbo can produce relatively similar images with different seeds,
+    making it harder to explore variations of a prompt. This node implements
+    reproducible, seed-based noise injection into text embeddings to increase
+    visual variation while maintaining reproducibility.
+
+    The noise strength is auto-calibrated relative to the embedding's standard
+    deviation, ensuring consistent results across different prompts.
+    """
+
+    conditioning: ZImageConditioningField = InputField(
+        description=FieldDescriptions.cond,
+        input=Input.Connection,
+        title="Conditioning",
+    )
+    seed: int = InputField(
+        default=0,
+        ge=0,
+        description="Seed for reproducible noise generation. Different seeds produce different noise patterns.",
+    )
+    strength: float = InputField(
+        default=0.1,
+        ge=0.0,
+        le=2.0,
+        description="Noise strength as multiplier of embedding std. 0=off, 0.1=subtle, 0.5=strong.",
+    )
+    randomize_percent: float = InputField(
+        default=50.0,
+        ge=1.0,
+        le=100.0,
+        description="Percentage of embedding values to add noise to (1-100). Lower values create more selective noise patterns.",
+    )
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> ZImageConditioningOutput:
+        # Load conditioning data
+        cond_data = context.conditioning.load(self.conditioning.conditioning_name)
+        assert len(cond_data.conditionings) == 1, "Expected exactly one conditioning tensor"
+        z_image_conditioning = cond_data.conditionings[0]
+        assert isinstance(z_image_conditioning, ZImageConditioningInfo), "Expected ZImageConditioningInfo"
+
+        # Early return if strength is zero (no modification needed)
+        if self.strength == 0:
+            return ZImageConditioningOutput(conditioning=self.conditioning)
+
+        # Clone embeddings to avoid modifying the original
+        prompt_embeds = z_image_conditioning.prompt_embeds.clone()
+
+        # Calculate actual noise strength based on embedding statistics
+        # This auto-calibration ensures consistent results across different prompts
+        embed_std = torch.std(prompt_embeds).item()
+        actual_strength = self.strength * embed_std
+
+        # Generate deterministic noise using the seed
+        generator = torch.Generator(device=prompt_embeds.device)
+        generator.manual_seed(self.seed)
+        noise = torch.rand(
+            prompt_embeds.shape, generator=generator, device=prompt_embeds.device, dtype=prompt_embeds.dtype
+        )
+        noise = noise * 2 - 1  # Scale to [-1, 1)
+        noise = noise * actual_strength
+
+        # Create selective mask for noise application
+        generator.manual_seed(self.seed + 1)
+        noise_mask = torch.bernoulli(
+            torch.ones_like(prompt_embeds) * (self.randomize_percent / 100.0),
+            generator=generator,
+        ).bool()
+
+        # Apply noise only to masked positions
+        prompt_embeds = prompt_embeds + (noise * noise_mask)
+
+        # Save modified conditioning
+        new_conditioning = ZImageConditioningInfo(prompt_embeds=prompt_embeds)
+        conditioning_data = ConditioningFieldData(conditionings=[new_conditioning])
+        conditioning_name = context.conditioning.save(conditioning_data)
+
+        return ZImageConditioningOutput(
+            conditioning=ZImageConditioningField(
+                conditioning_name=conditioning_name,
+                mask=self.conditioning.mask,
+            )
+        )

invokeai/app/services/config/config_default.py

@@ -85,6 +85,7 @@ class InvokeAIAppConfig(BaseSettings):
         max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.
         max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.
         log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.
+        model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.
         device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.
         enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.
         keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.
@@ -165,9 +166,10 @@ class InvokeAIAppConfig(BaseSettings):
     max_cache_ram_gb:   Optional[float] = Field(default=None, gt=0,         description="The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.")
     max_cache_vram_gb:  Optional[float] = Field(default=None, ge=0,         description="The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.")
     log_memory_usage:              bool = Field(default=False,              description="If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.")
+    model_cache_keep_alive_min:   float = Field(default=0, ge=0,            description="How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.")
     device_working_mem_gb:        float = Field(default=3,                  description="The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.")
     enable_partial_loading:        bool = Field(default=False,              description="Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.")
-    keep_ram_copy_of_weights:      bool = Field(default=True,              description="Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.")
+    keep_ram_copy_of_weights:      bool = Field(default=True,               description="Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.")
     # Deprecated CACHE configs
     ram:                Optional[float] = Field(default=None, gt=0,         description="DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.")
     vram:               Optional[float] = Field(default=None, ge=0,         description="DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.")

invokeai/app/services/invocation_stats/invocation_stats_common.py

@@ -14,7 +14,7 @@ class NodeExecutionStatsSummary:
     node_type: str
     num_calls: int
     time_used_seconds: float
-    peak_vram_gb: float
+    delta_vram_gb: float
 
 
 @dataclass
@@ -58,10 +58,10 @@ class InvocationStatsSummary:
     def __str__(self) -> str:
         _str = ""
         _str = f"Graph stats: {self.graph_stats.graph_execution_state_id}\n"
-        _str += f"{'Node':>30} {'Calls':>7} {'Seconds':>9} {'VRAM Used':>10}\n"
+        _str += f"{'Node':>30} {'Calls':>7} {'Seconds':>9} {'VRAM Change':+>10}\n"
 
         for summary in self.node_stats:
-            _str += f"{summary.node_type:>30} {summary.num_calls:>7} {summary.time_used_seconds:>8.3f}s {summary.peak_vram_gb:>9.3f}G\n"
+            _str += f"{summary.node_type:>30} {summary.num_calls:>7} {summary.time_used_seconds:>8.3f}s {summary.delta_vram_gb:+10.3f}G\n"
 
         _str += f"TOTAL GRAPH EXECUTION TIME: {self.graph_stats.execution_time_seconds:7.3f}s\n"
 
@@ -100,7 +100,7 @@ class NodeExecutionStats:
     start_ram_gb: float  # GB
     end_ram_gb: float  # GB
 
-    peak_vram_gb: float  # GB
+    delta_vram_gb: float  # GB
 
     def total_time(self) -> float:
         return self.end_time - self.start_time
@@ -174,9 +174,9 @@ class GraphExecutionStats:
         for node_type, node_type_stats_list in node_stats_by_type.items():
             num_calls = len(node_type_stats_list)
             time_used = sum([n.total_time() for n in node_type_stats_list])
-            peak_vram = max([n.peak_vram_gb for n in node_type_stats_list])
+            delta_vram = max([n.delta_vram_gb for n in node_type_stats_list])
             summary = NodeExecutionStatsSummary(
-                node_type=node_type, num_calls=num_calls, time_used_seconds=time_used, peak_vram_gb=peak_vram
+                node_type=node_type, num_calls=num_calls, time_used_seconds=time_used, delta_vram_gb=delta_vram
             )
             summaries.append(summary)