Merge branch 'main' into copilot/add-unload-model-option

Author: lstein

docs/RELEASE.md

@@ -89,7 +89,7 @@ The publish jobs will not run if any of the previous jobs fail.
 
 They use [GitHub environments], which are configured as [trusted publishers] on PyPI.
 
-Both jobs require a @hipsterusername or @psychedelicious to approve them from the workflow's **Summary** tab.
+Both jobs require a @lstein or @blessedcoolant to approve them from the workflow's **Summary** tab.
 
 - Click the **Review deployments** button
 - Select the environment (either `testpypi` or `pypi` - typically you select both)
@@ -101,7 +101,7 @@ Both jobs require a @hipsterusername or @psychedelicious to approve them from th
 
 Check the [python infrastructure status page] for incidents.
 
-If there are no incidents, contact @hipsterusername or @lstein, who have owner access to GH and PyPI, to see if access has expired or something like that.
+If there are no incidents, contact @lstein or @blessedcoolant, who have owner access to GH and PyPI, to see if access has expired or something like that.
 
 #### `publish-testpypi` Job
 

invokeai/app/invocations/fields.py

@@ -333,6 +333,11 @@ class ZImageConditioningField(BaseModel):
     """A Z-Image conditioning tensor primitive value"""
 
     conditioning_name: str = Field(description="The name of conditioning tensor")
+    mask: Optional[TensorField] = Field(
+        default=None,
+        description="The mask associated with this conditioning tensor for regional prompting. "
+        "Excluded regions should be set to False, included regions should be set to True.",
+    )
 
 
 class ConditioningField(BaseModel):

invokeai/app/invocations/z_image_denoise.py

@@ -32,23 +32,29 @@
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ZImageConditioningInfo
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.z_image.extensions.regional_prompting_extension import ZImageRegionalPromptingExtension
+from invokeai.backend.z_image.text_conditioning import ZImageTextConditioning
 from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter
 from invokeai.backend.z_image.z_image_controlnet_extension import (
     ZImageControlNetExtension,
     z_image_forward_with_control,
 )
+from invokeai.backend.z_image.z_image_transformer_patch import patch_transformer_for_regional_prompting
 
 
 @invocation(
     "z_image_denoise",
     title="Denoise - Z-Image",
     tags=["image", "z-image"],
     category="image",
-    version="1.1.0",
+    version="1.2.0",
     classification=Classification.Prototype,
 )
 class ZImageDenoiseInvocation(BaseInvocation):
-    """Run the denoising process with a Z-Image model."""
+    """Run the denoising process with a Z-Image model.
+
+    Supports regional prompting by connecting multiple conditioning inputs with masks.
+    """
 
     # If latents is provided, this means we are doing image-to-image.
     latents: Optional[LatentsField] = InputField(
@@ -63,10 +69,10 @@ class ZImageDenoiseInvocation(BaseInvocation):
     transformer: TransformerField = InputField(
         description=FieldDescriptions.z_image_model, input=Input.Connection, title="Transformer"
     )
-    positive_conditioning: ZImageConditioningField = InputField(
+    positive_conditioning: ZImageConditioningField | list[ZImageConditioningField] = InputField(
         description=FieldDescriptions.positive_cond, input=Input.Connection
     )
-    negative_conditioning: Optional[ZImageConditioningField] = InputField(
+    negative_conditioning: ZImageConditioningField | list[ZImageConditioningField] | None = InputField(
         default=None, description=FieldDescriptions.negative_cond, input=Input.Connection
     )
     # Z-Image-Turbo works best without CFG (guidance_scale=1.0)
@@ -126,25 +132,50 @@ def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor)
     def _load_text_conditioning(
         self,
         context: InvocationContext,
-        conditioning_name: str,
+        cond_field: ZImageConditioningField | list[ZImageConditioningField],
+        img_height: int,
+        img_width: int,
         dtype: torch.dtype,
         device: torch.device,
-    ) -> torch.Tensor:
-        """Load Z-Image text conditioning."""
-        cond_data = context.conditioning.load(conditioning_name)
-        if len(cond_data.conditionings) != 1:
-            raise ValueError(
-                f"Expected exactly 1 conditioning entry for Z-Image, got {len(cond_data.conditionings)}. "
-                "Ensure you are using the Z-Image text encoder."
-            )
-        z_image_conditioning = cond_data.conditionings[0]
-        if not isinstance(z_image_conditioning, ZImageConditioningInfo):
-            raise TypeError(
-                f"Expected ZImageConditioningInfo, got {type(z_image_conditioning).__name__}. "
-                "Ensure you are using the Z-Image text encoder."
-            )
-        z_image_conditioning = z_image_conditioning.to(dtype=dtype, device=device)
-        return z_image_conditioning.prompt_embeds
+    ) -> list[ZImageTextConditioning]:
+        """Load Z-Image text conditioning with optional regional masks.
+
+        Args:
+            context: The invocation context.
+            cond_field: Single conditioning field or list of fields.
+            img_height: Height of the image token grid (H // patch_size).
+            img_width: Width of the image token grid (W // patch_size).
+            dtype: Target dtype.
+            device: Target device.
+
+        Returns:
+            List of ZImageTextConditioning objects with embeddings and masks.
+        """
+        # Normalize to a list
+        cond_list = [cond_field] if isinstance(cond_field, ZImageConditioningField) else cond_field
+
+        text_conditionings: list[ZImageTextConditioning] = []
+        for cond in cond_list:
+            # Load the text embeddings
+            cond_data = context.conditioning.load(cond.conditioning_name)
+            assert len(cond_data.conditionings) == 1
+            z_image_conditioning = cond_data.conditionings[0]
+            assert isinstance(z_image_conditioning, ZImageConditioningInfo)
+            z_image_conditioning = z_image_conditioning.to(dtype=dtype, device=device)
+            prompt_embeds = z_image_conditioning.prompt_embeds
+
+            # Load the mask, if provided
+            mask: torch.Tensor | None = None
+            if cond.mask is not None:
+                mask = context.tensors.load(cond.mask.tensor_name)
+                mask = mask.to(device=device)
+                mask = ZImageRegionalPromptingExtension.preprocess_regional_prompt_mask(
+                    mask, img_height, img_width, dtype, device
+                )
+
+            text_conditionings.append(ZImageTextConditioning(prompt_embeds=prompt_embeds, mask=mask))
+
+        return text_conditionings
 
     def _get_noise(
         self,
@@ -221,14 +252,33 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
 
         transformer_info = context.models.load(self.transformer.transformer)
 
-        # Load positive conditioning
-        pos_prompt_embeds = self._load_text_conditioning(
+        # Calculate image token grid dimensions
+        patch_size = 2  # Z-Image uses patch_size=2
+        latent_height = self.height // LATENT_SCALE_FACTOR
+        latent_width = self.width // LATENT_SCALE_FACTOR
+        img_token_height = latent_height // patch_size
+        img_token_width = latent_width // patch_size
+        img_seq_len = img_token_height * img_token_width
+
+        # Load positive conditioning with regional masks
+        pos_text_conditionings = self._load_text_conditioning(
             context=context,
-            conditioning_name=self.positive_conditioning.conditioning_name,
+            cond_field=self.positive_conditioning,
+            img_height=img_token_height,
+            img_width=img_token_width,
             dtype=inference_dtype,
             device=device,
         )
 
+        # Create regional prompting extension
+        regional_extension = ZImageRegionalPromptingExtension.from_text_conditionings(
+            text_conditionings=pos_text_conditionings,
+            img_seq_len=img_seq_len,
+        )
+
+        # Get the concatenated prompt embeddings for the transformer
+        pos_prompt_embeds = regional_extension.regional_text_conditioning.prompt_embeds
+
         # Load negative conditioning if provided and guidance_scale != 1.0
         # CFG formula: pred = pred_uncond + cfg_scale * (pred_cond - pred_uncond)
         # At cfg_scale=1.0: pred = pred_cond (no effect, skip uncond computation)
@@ -238,21 +288,22 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
             not math.isclose(self.guidance_scale, 1.0) and self.negative_conditioning is not None
         )
         if do_classifier_free_guidance:
-            if self.negative_conditioning is None:
-                raise ValueError("Negative conditioning is required when guidance_scale != 1.0")
-            neg_prompt_embeds = self._load_text_conditioning(
+            assert self.negative_conditioning is not None
+            # Load all negative conditionings and concatenate embeddings
+            # Note: We ignore masks for negative conditioning as regional negative prompting is not fully supported
+            neg_text_conditionings = self._load_text_conditioning(
                 context=context,
-                conditioning_name=self.negative_conditioning.conditioning_name,
+                cond_field=self.negative_conditioning,
+                img_height=img_token_height,
+                img_width=img_token_width,
                 dtype=inference_dtype,
                 device=device,
             )
-
-        # Calculate image sequence length for timestep shifting
-        patch_size = 2  # Z-Image uses patch_size=2
-        image_seq_len = ((self.height // LATENT_SCALE_FACTOR) * (self.width // LATENT_SCALE_FACTOR)) // (patch_size**2)
+            # Concatenate all negative embeddings
+            neg_prompt_embeds = torch.cat([tc.prompt_embeds for tc in neg_text_conditionings], dim=0)
 
         # Calculate shift based on image sequence length
-        mu = self._calculate_shift(image_seq_len)
+        mu = self._calculate_shift(img_seq_len)
 
         # Generate sigma schedule with time shift
         sigmas = self._get_sigmas(mu, self.steps)
@@ -443,6 +494,15 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
                 )
             )
 
+            # Apply regional prompting patch if we have regional masks
+            exit_stack.enter_context(
+                patch_transformer_for_regional_prompting(
+                    transformer=transformer,
+                    regional_attn_mask=regional_extension.regional_attn_mask,
+                    img_seq_len=img_seq_len,
+                )
+            )
+
             # Denoising loop
             for step_idx in tqdm(range(total_steps)):
                 sigma_curr = sigmas[step_idx]

invokeai/app/invocations/z_image_text_encoder.py

@@ -1,11 +1,18 @@
 from contextlib import ExitStack
-from typing import Iterator, Tuple
+from typing import Iterator, Optional, Tuple
 
 import torch
 from transformers import PreTrainedModel, PreTrainedTokenizerBase
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
-from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, UIComponent
+from invokeai.app.invocations.fields import (
+    FieldDescriptions,
+    Input,
+    InputField,
+    TensorField,
+    UIComponent,
+    ZImageConditioningField,
+)
 from invokeai.app.invocations.model import Qwen3EncoderField
 from invokeai.app.invocations.primitives import ZImageConditioningOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
@@ -27,25 +34,34 @@
     title="Prompt - Z-Image",
     tags=["prompt", "conditioning", "z-image"],
     category="conditioning",
-    version="1.0.0",
+    version="1.1.0",
     classification=Classification.Prototype,
 )
 class ZImageTextEncoderInvocation(BaseInvocation):
-    """Encodes and preps a prompt for a Z-Image image."""
+    """Encodes and preps a prompt for a Z-Image image.
+
+    Supports regional prompting by connecting a mask input.
+    """
 
     prompt: str = InputField(description="Text prompt to encode.", ui_component=UIComponent.Textarea)
     qwen3_encoder: Qwen3EncoderField = InputField(
         title="Qwen3 Encoder",
         description=FieldDescriptions.qwen3_encoder,
         input=Input.Connection,
     )
+    mask: Optional[TensorField] = InputField(
+        default=None,
+        description="A mask defining the region that this conditioning prompt applies to.",
+    )
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ZImageConditioningOutput:
         prompt_embeds = self._encode_prompt(context, max_seq_len=Z_IMAGE_MAX_SEQ_LEN)
         conditioning_data = ConditioningFieldData(conditionings=[ZImageConditioningInfo(prompt_embeds=prompt_embeds)])
         conditioning_name = context.conditioning.save(conditioning_data)
-        return ZImageConditioningOutput.build(conditioning_name)
+        return ZImageConditioningOutput(
+            conditioning=ZImageConditioningField(conditioning_name=conditioning_name, mask=self.mask)
+        )
 
     def _encode_prompt(self, context: InvocationContext, max_seq_len: int) -> torch.Tensor:
         """Encode prompt using Qwen3 text encoder.

invokeai/app/services/model_install/model_install_default.py

@@ -1,8 +1,10 @@
 """Model installation class."""
 
+import gc
 import locale
 import os
 import re
+import sys
 import threading
 import time
 from copy import deepcopy
@@ -187,6 +189,22 @@ def register_path(
         config.source_type = ModelSourceType.Path
         return self._register(model_path, config)
 
+    # TODO: Replace this with a proper fix for underlying problem of Windows holding open
+    # the file when it needs to be moved.
+    @staticmethod
+    def _move_with_retries(src: Path, dst: Path, attempts: int = 5, delay: float = 0.5) -> None:
+        """Workaround for Windows file-handle issues when moving files."""
+        for tries_left in range(attempts, 0, -1):
+            try:
+                move(src, dst)
+                return
+            except PermissionError:
+                gc.collect()
+                if tries_left == 1:
+                    raise
+                time.sleep(delay)
+                delay *= 2  # Exponential backoff
+
     def install_path(
         self,
         model_path: Union[Path, str],
@@ -205,7 +223,7 @@ def install_path(
             dest_dir.mkdir(parents=True)
             dest_path = dest_dir / model_path.name if model_path.is_file() else dest_dir
             if model_path.is_file():
-                move(model_path, dest_path)
+                self._move_with_retries(model_path, dest_path)  # Windows workaround TODO: fix root cause
             elif model_path.is_dir():
                 # Move the contents of the directory, not the directory itself
                 for item in model_path.iterdir():
@@ -500,6 +518,39 @@ def _start_installer_thread(self) -> None:
         self._install_thread.start()
         self._running = True
 
+    @staticmethod
+    def _safe_rmtree(path: Path, logger: Any) -> None:
+        """Remove a directory tree with retry logic for Windows file locking issues.
+
+        On Windows, memory-mapped files may not be immediately released even after
+        the file handle is closed. This function retries the removal with garbage
+        collection to help release any lingering references.
+        """
+        max_retries = 3
+        retry_delay = 0.5  # seconds
+
+        for attempt in range(max_retries):
+            try:
+                # Force garbage collection to release any lingering file references
+                gc.collect()
+                rmtree(path)
+                return
+            except PermissionError as e:
+                if attempt < max_retries - 1 and sys.platform == "win32":
+                    logger.warning(
+                        f"Failed to remove {path} (attempt {attempt + 1}/{max_retries}): {e}. "
+                        f"Retrying in {retry_delay}s..."
+                    )
+                    time.sleep(retry_delay)
+                    retry_delay *= 2  # Exponential backoff
+                else:
+                    logger.error(f"Failed to remove temporary directory {path}: {e}")
+                    # On final failure, don't raise - the temp dir will be cleaned up on next startup
+                    return
+            except Exception as e:
+                logger.error(f"Unexpected error removing {path}: {e}")
+                return
+
     def _install_next_item(self) -> None:
         self._logger.debug(f"Installer thread {threading.get_ident()} starting")
         while True:
@@ -529,7 +580,7 @@ def _install_next_item(self) -> None:
             finally:
                 # if this is an install of a remote file, then clean up the temporary directory
                 if job._install_tmpdir is not None:
-                    rmtree(job._install_tmpdir)
+                    self._safe_rmtree(job._install_tmpdir, self._logger)
                 self._install_completed_event.set()
                 self._install_queue.task_done()
         self._logger.info(f"Installer thread {threading.get_ident()} exiting")
@@ -574,7 +625,7 @@ def _remove_dangling_install_dirs(self) -> None:
         path = self._app_config.models_path
         for tmpdir in path.glob(f"{TMPDIR_PREFIX}*"):
             self._logger.info(f"Removing dangling temporary directory {tmpdir}")
-            rmtree(tmpdir)
+            self._safe_rmtree(tmpdir, self._logger)
 
     def _scan_for_missing_models(self) -> list[AnyModelConfig]:
         """Scan the models directory for missing models and return a list of them."""