InvokeAI 6.9.0rc3__py3-none-any.whl → 6.10.0rc1__py3-none-any.whl

invokeai/backend/quantization/gguf/ggml_tensor.py

@@ -14,11 +14,31 @@ def dequantize_and_run(func, args, kwargs):
     """A helper function for running math ops on GGMLTensor inputs.
 
     Dequantizes the inputs, and runs the function.
+    Also casts other floating point tensors to match the compute_dtype of GGMLTensors
+    to avoid dtype mismatches in matrix operations.
     """
-    dequantized_args = [a.get_dequantized_tensor() if hasattr(a, "get_dequantized_tensor") else a for a in args]
-    dequantized_kwargs = {
-        k: v.get_dequantized_tensor() if hasattr(v, "get_dequantized_tensor") else v for k, v in kwargs.items()
-    }
+    # Find the compute_dtype from any GGMLTensor in the args
+    compute_dtype = None
+    for a in args:
+        if hasattr(a, "compute_dtype"):
+            compute_dtype = a.compute_dtype
+            break
+    if compute_dtype is None:
+        for v in kwargs.values():
+            if hasattr(v, "compute_dtype"):
+                compute_dtype = v.compute_dtype
+                break
+
+    def process_tensor(t):
+        if hasattr(t, "get_dequantized_tensor"):
+            return t.get_dequantized_tensor()
+        elif isinstance(t, torch.Tensor) and compute_dtype is not None and t.is_floating_point():
+            # Cast other floating point tensors to match the GGUF compute_dtype
+            return t.to(compute_dtype)
+        return t
+
+    dequantized_args = [process_tensor(a) for a in args]
+    dequantized_kwargs = {k: process_tensor(v) for k, v in kwargs.items()}
     return func(*dequantized_args, **dequantized_kwargs)
 
 
@@ -57,6 +77,9 @@ GGML_TENSOR_OP_TABLE = {
     torch.ops.aten.sub.Tensor: dequantize_and_run,  # pyright: ignore
     torch.ops.aten.allclose.default: dequantize_and_run,  # pyright: ignore
     torch.ops.aten.slice.Tensor: dequantize_and_run,  # pyright: ignore
+    torch.ops.aten.view.default: dequantize_and_run,  # pyright: ignore
+    torch.ops.aten.expand.default: dequantize_and_run,  # pyright: ignore
+    torch.ops.aten.index_put_.default: dequantize_and_run,  # pyright: ignore
 }
 
 if torch.backends.mps.is_available():

invokeai/backend/quantization/gguf/loaders.py

@@ -1,3 +1,4 @@
+import gc
 from pathlib import Path
 
 import gguf
@@ -5,18 +6,52 @@ import torch
 
 from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
 from invokeai.backend.quantization.gguf.utils import TORCH_COMPATIBLE_QTYPES
+from invokeai.backend.util.logging import InvokeAILogger
+
+logger = InvokeAILogger.get_logger()
+
+
+class WrappedGGUFReader:
+    """Wrapper around GGUFReader that adds a close() method."""
+
+    def __init__(self, path: Path):
+        self.reader = gguf.GGUFReader(path)
+
+    def __enter__(self):
+        return self.reader
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        return False
+
+    def close(self):
+        """Explicitly close the memory-mapped file."""
+        if hasattr(self.reader, "data"):
+            try:
+                self.reader.data.flush()
+                del self.reader.data
+            except (AttributeError, OSError, ValueError) as e:
+                logger.warning(f"Wasn't able to close GGUF memory map: {e}")
+        del self.reader
+        gc.collect()
 
 
 def gguf_sd_loader(path: Path, compute_dtype: torch.dtype) -> dict[str, GGMLTensor]:
-    reader = gguf.GGUFReader(path)
-
-    sd: dict[str, GGMLTensor] = {}
-    for tensor in reader.tensors:
-        torch_tensor = torch.from_numpy(tensor.data)
-        shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
-        if tensor.tensor_type in TORCH_COMPATIBLE_QTYPES:
-            torch_tensor = torch_tensor.view(*shape)
-        sd[tensor.name] = GGMLTensor(
-            torch_tensor, ggml_quantization_type=tensor.tensor_type, tensor_shape=shape, compute_dtype=compute_dtype
-        )
-    return sd
+    with WrappedGGUFReader(path) as reader:
+        sd: dict[str, GGMLTensor] = {}
+        for tensor in reader.tensors:
+            # Use .copy() to create a true copy of the data, not a view.
+            # This is critical on Windows where the memory-mapped file cannot be deleted
+            # while tensors still hold references to the mapped memory.
+            torch_tensor = torch.from_numpy(tensor.data.copy())
+
+            shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
+            if tensor.tensor_type in TORCH_COMPATIBLE_QTYPES:
+                torch_tensor = torch_tensor.view(*shape)
+            sd[tensor.name] = GGMLTensor(
+                torch_tensor,
+                ggml_quantization_type=tensor.tensor_type,
+                tensor_shape=shape,
+                compute_dtype=compute_dtype,
+            )
+        return sd

invokeai/backend/stable_diffusion/diffusion/conditioning_data.py

@@ -76,6 +76,18 @@ class CogView4ConditioningInfo:
         return self
 
 
+@dataclass
+class ZImageConditioningInfo:
+    """Z-Image text conditioning information from Qwen3 text encoder."""
+
+    prompt_embeds: torch.Tensor
+    """Text embeddings from Qwen3 encoder. Shape: (batch_size, seq_len, hidden_size)."""
+
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        self.prompt_embeds = self.prompt_embeds.to(device=device, dtype=dtype)
+        return self
+
+
 @dataclass
 class ConditioningFieldData:
     # If you change this class, adding more types, you _must_ update the instantiation of ObjectSerializerDisk in
@@ -87,6 +99,7 @@ class ConditioningFieldData:
         | List[FLUXConditioningInfo]
         | List[SD3ConditioningInfo]
         | List[CogView4ConditioningInfo]
+        | List[ZImageConditioningInfo]
     )
 
 

invokeai/backend/util/devices.py

@@ -112,3 +112,28 @@ class TorchDevice:
     @classmethod
     def _to_dtype(cls, precision_name: TorchPrecisionNames) -> torch.dtype:
         return NAME_TO_PRECISION[precision_name]
+
+    @classmethod
+    def choose_bfloat16_safe_dtype(cls, device: Optional[torch.device] = None) -> torch.dtype:
+        """Return bfloat16 if supported on the device, else fallback to float16/float32.
+
+        This is useful for models that require bfloat16 precision (e.g., Z-Image, Flux)
+        but need to run on hardware that may not support bfloat16.
+
+        Args:
+            device: The target device. If None, uses choose_torch_device().
+
+        Returns:
+            torch.bfloat16 if supported, torch.float16 for CUDA without bfloat16 support,
+            or torch.float32 for CPU/MPS.
+        """
+        device = device or cls.choose_torch_device()
+        try:
+            # Test if bfloat16 is supported on this device
+            torch.tensor([1.0], dtype=torch.bfloat16, device=device)
+            return torch.bfloat16
+        except TypeError:
+            # bfloat16 not supported - fallback based on device type
+            if device.type == "cuda":
+                return torch.float16
+            return torch.float32

invokeai/backend/util/hotfixes.py

@@ -5,7 +5,7 @@ import torch
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.loaders.single_file_model import FromOriginalModelMixin
 from diffusers.models.attention_processor import AttentionProcessor, AttnProcessor
-from diffusers.models.controlnet import ControlNetConditioningEmbedding, ControlNetOutput, zero_module
+from diffusers.models.controlnets.controlnet import ControlNetConditioningEmbedding, ControlNetOutput, zero_module
 from diffusers.models.embeddings import (
     TextImageProjection,
     TextImageTimeEmbedding,
@@ -777,7 +777,7 @@ class ControlNetModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
 
 
 diffusers.ControlNetModel = ControlNetModel
-diffusers.models.controlnet.ControlNetModel = ControlNetModel
+diffusers.models.controlnets.controlnet.ControlNetModel = ControlNetModel
 
 
 # patch LoRACompatibleConv to use original Conv2D forward function

invokeai/backend/z_image/__init__.py

@@ -0,0 +1,16 @@
+# Z-Image backend utilities
+from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter
+from invokeai.backend.z_image.z_image_control_transformer import ZImageControlTransformer2DModel
+from invokeai.backend.z_image.z_image_controlnet_extension import (
+    ZImageControlNetExtension,
+    z_image_forward_with_control,
+)
+from invokeai.backend.z_image.z_image_patchify_utils import patchify_control_context
+
+__all__ = [
+    "ZImageControlAdapter",
+    "ZImageControlTransformer2DModel",
+    "ZImageControlNetExtension",
+    "z_image_forward_with_control",
+    "patchify_control_context",
+]

invokeai/backend/z_image/extensions/__init__.py

@@ -0,0 +1 @@
+# Z-Image extensions

invokeai/backend/z_image/extensions/regional_prompting_extension.py

@@ -0,0 +1,207 @@
+from typing import Optional
+
+import torch
+import torchvision
+
+from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.mask import to_standard_float_mask
+from invokeai.backend.z_image.text_conditioning import ZImageRegionalTextConditioning, ZImageTextConditioning
+
+
+class ZImageRegionalPromptingExtension:
+    """A class for managing regional prompting with Z-Image.
+
+    This implementation is inspired by the FLUX regional prompting extension and
+    the paper https://arxiv.org/pdf/2411.02395.
+
+    Key difference from FLUX: Z-Image uses sequence order [img_tokens, txt_tokens],
+    while FLUX uses [txt_tokens, img_tokens]. The attention mask construction
+    accounts for this difference.
+    """
+
+    def __init__(
+        self,
+        regional_text_conditioning: ZImageRegionalTextConditioning,
+        regional_attn_mask: torch.Tensor | None = None,
+    ):
+        self.regional_text_conditioning = regional_text_conditioning
+        self.regional_attn_mask = regional_attn_mask
+
+    def get_attn_mask(self, block_index: int) -> torch.Tensor | None:
+        """Get the attention mask for a given block index.
+
+        Uses alternating pattern: apply mask on even blocks, no mask on odd blocks.
+        This helps balance regional control with global coherence.
+        """
+        order = [self.regional_attn_mask, None]
+        return order[block_index % len(order)]
+
+    @classmethod
+    def from_text_conditionings(
+        cls,
+        text_conditionings: list[ZImageTextConditioning],
+        img_seq_len: int,
+    ) -> "ZImageRegionalPromptingExtension":
+        """Create a ZImageRegionalPromptingExtension from a list of text conditionings.
+
+        Args:
+            text_conditionings: List of text conditionings with optional masks.
+            img_seq_len: The image sequence length (i.e. (H // patch_size) * (W // patch_size)).
+
+        Returns:
+            A configured ZImageRegionalPromptingExtension.
+        """
+        regional_text_conditioning = ZImageRegionalTextConditioning.from_text_conditionings(text_conditionings)
+        attn_mask = cls._prepare_regional_attn_mask(regional_text_conditioning, img_seq_len)
+        return cls(
+            regional_text_conditioning=regional_text_conditioning,
+            regional_attn_mask=attn_mask,
+        )
+
+    @classmethod
+    def _prepare_regional_attn_mask(
+        cls,
+        regional_text_conditioning: ZImageRegionalTextConditioning,
+        img_seq_len: int,
+    ) -> torch.Tensor | None:
+        """Prepare a regional attention mask for Z-Image.
+
+        The mask controls which tokens can attend to each other:
+        - Image tokens within a region attend only to each other
+        - Image tokens attend only to their corresponding regional text
+        - Text tokens attend only to their corresponding regional image
+        - Text tokens attend to themselves
+
+        Z-Image sequence order: [img_tokens, txt_tokens]
+
+        Args:
+            regional_text_conditioning: The regional text conditioning data.
+            img_seq_len: Number of image tokens.
+
+        Returns:
+            Attention mask of shape (img_seq_len + txt_seq_len, img_seq_len + txt_seq_len).
+            Returns None if no regional masks are present.
+        """
+        # Check if any regional masks exist
+        has_regional_masks = any(mask is not None for mask in regional_text_conditioning.image_masks)
+        if not has_regional_masks:
+            # No regional masks, return None to use default attention
+            return None
+
+        # Identify background region (area not covered by any mask)
+        background_region_mask: torch.Tensor | None = None
+        for image_mask in regional_text_conditioning.image_masks:
+            if image_mask is not None:
+                # image_mask shape: (1, 1, img_seq_len) -> flatten to (img_seq_len,)
+                mask_flat = image_mask.view(-1)
+                if background_region_mask is None:
+                    background_region_mask = torch.ones_like(mask_flat)
+                background_region_mask = background_region_mask * (1 - mask_flat)
+
+        device = TorchDevice.choose_torch_device()
+        txt_seq_len = regional_text_conditioning.prompt_embeds.shape[0]
+        total_seq_len = img_seq_len + txt_seq_len
+
+        # Initialize empty attention mask
+        # Z-Image sequence: [img_tokens (0:img_seq_len), txt_tokens (img_seq_len:total_seq_len)]
+        regional_attention_mask = torch.zeros((total_seq_len, total_seq_len), device=device, dtype=torch.float16)
+
+        for image_mask, embedding_range in zip(
+            regional_text_conditioning.image_masks,
+            regional_text_conditioning.embedding_ranges,
+            strict=True,
+        ):
+            # Calculate text token positions in the unified sequence
+            txt_start = img_seq_len + embedding_range.start
+            txt_end = img_seq_len + embedding_range.end
+
+            # 1. txt attends to itself
+            regional_attention_mask[txt_start:txt_end, txt_start:txt_end] = 1.0
+
+            if image_mask is not None:
+                # Flatten mask: (1, 1, img_seq_len) -> (img_seq_len,)
+                mask_flat = image_mask.view(img_seq_len)
+
+                # 2. img attends to corresponding regional txt
+                # Reshape mask to (img_seq_len, 1) for broadcasting
+                regional_attention_mask[:img_seq_len, txt_start:txt_end] = mask_flat.view(img_seq_len, 1)
+
+                # 3. txt attends to corresponding regional img
+                # Reshape mask to (1, img_seq_len) for broadcasting
+                regional_attention_mask[txt_start:txt_end, :img_seq_len] = mask_flat.view(1, img_seq_len)
+
+                # 4. img self-attention within region
+                # mask @ mask.T creates pairwise attention within the masked region
+                regional_attention_mask[:img_seq_len, :img_seq_len] += mask_flat.view(img_seq_len, 1) @ mask_flat.view(
+                    1, img_seq_len
+                )
+            else:
+                # Global prompt: allow attention to/from background regions only
+                if background_region_mask is not None:
+                    # 2. background img attends to global txt
+                    regional_attention_mask[:img_seq_len, txt_start:txt_end] = background_region_mask.view(
+                        img_seq_len, 1
+                    )
+
+                    # 3. global txt attends to background img
+                    regional_attention_mask[txt_start:txt_end, :img_seq_len] = background_region_mask.view(
+                        1, img_seq_len
+                    )
+                else:
+                    # No regional masks at all, allow full attention
+                    regional_attention_mask[:img_seq_len, txt_start:txt_end] = 1.0
+                    regional_attention_mask[txt_start:txt_end, :img_seq_len] = 1.0
+
+        # Allow background regions to attend to themselves
+        if background_region_mask is not None:
+            bg_mask = background_region_mask.view(img_seq_len, 1)
+            regional_attention_mask[:img_seq_len, :img_seq_len] += bg_mask @ bg_mask.T
+
+        # Convert to boolean mask
+        regional_attention_mask = regional_attention_mask > 0.5
+
+        return regional_attention_mask
+
+    @staticmethod
+    def preprocess_regional_prompt_mask(
+        mask: Optional[torch.Tensor],
+        target_height: int,
+        target_width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Preprocess a regional prompt mask to match the target image token grid.
+
+        Args:
+            mask: Input mask tensor. If None, returns a mask of all ones.
+            target_height: Height of the image token grid (H // patch_size).
+            target_width: Width of the image token grid (W // patch_size).
+            dtype: Target dtype for the mask.
+            device: Target device for the mask.
+
+        Returns:
+            Processed mask of shape (1, 1, target_height * target_width).
+        """
+        img_seq_len = target_height * target_width
+
+        if mask is None:
+            return torch.ones((1, 1, img_seq_len), dtype=dtype, device=device)
+
+        mask = to_standard_float_mask(mask, out_dtype=dtype)
+
+        # Resize mask to target dimensions
+        tf = torchvision.transforms.Resize(
+            (target_height, target_width),
+            interpolation=torchvision.transforms.InterpolationMode.NEAREST,
+        )
+
+        # Add batch dimension if needed: (h, w) -> (1, h, w) -> (1, 1, h, w)
+        if mask.ndim == 2:
+            mask = mask.unsqueeze(0)
+        if mask.ndim == 3:
+            mask = mask.unsqueeze(0)
+
+        resized_mask = tf(mask)
+
+        # Flatten to (1, 1, img_seq_len)
+        return resized_mask.flatten(start_dim=2).to(device=device)

invokeai/backend/z_image/text_conditioning.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass
+
+import torch
+
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Range
+
+
+@dataclass
+class ZImageTextConditioning:
+    """Z-Image text conditioning with optional regional mask.
+
+    Attributes:
+        prompt_embeds: Text embeddings from Qwen3 encoder. Shape: (seq_len, hidden_size).
+        mask: Optional binary mask for regional prompting. If None, the prompt is global.
+              Shape: (1, 1, img_seq_len) where img_seq_len = (H // patch_size) * (W // patch_size).
+    """
+
+    prompt_embeds: torch.Tensor
+    mask: torch.Tensor | None = None
+
+
+@dataclass
+class ZImageRegionalTextConditioning:
+    """Container for multiple regional text conditionings concatenated together.
+
+    In Z-Image, the unified sequence is [img_tokens, txt_tokens], which is different
+    from FLUX where it's [txt_tokens, img_tokens]. The attention mask must account for this.
+
+    Attributes:
+        prompt_embeds: Concatenated text embeddings from all regional prompts.
+                       Shape: (total_seq_len, hidden_size).
+        image_masks: List of binary masks for each regional prompt.
+                     image_masks[i] corresponds to embedding_ranges[i].
+                     If None, the prompt is global (applies to entire image).
+                     Shape: (1, 1, img_seq_len).
+        embedding_ranges: List of ranges indicating which portion of prompt_embeds
+                         corresponds to each regional prompt.
+    """
+
+    prompt_embeds: torch.Tensor
+    image_masks: list[torch.Tensor | None]
+    embedding_ranges: list[Range]
+
+    @classmethod
+    def from_text_conditionings(
+        cls,
+        text_conditionings: list[ZImageTextConditioning],
+    ) -> "ZImageRegionalTextConditioning":
+        """Create a ZImageRegionalTextConditioning from a list of ZImageTextConditioning objects.
+
+        Args:
+            text_conditionings: List of text conditionings, each with optional mask.
+
+        Returns:
+            A single ZImageRegionalTextConditioning with concatenated embeddings.
+        """
+        concat_embeds: list[torch.Tensor] = []
+        concat_ranges: list[Range] = []
+        image_masks: list[torch.Tensor | None] = []
+
+        cur_embed_len = 0
+        for tc in text_conditionings:
+            concat_embeds.append(tc.prompt_embeds)
+            concat_ranges.append(Range(start=cur_embed_len, end=cur_embed_len + tc.prompt_embeds.shape[0]))
+            image_masks.append(tc.mask)
+            cur_embed_len += tc.prompt_embeds.shape[0]
+
+        prompt_embeds = torch.cat(concat_embeds, dim=0)
+
+        return cls(
+            prompt_embeds=prompt_embeds,
+            image_masks=image_masks,
+            embedding_ranges=concat_ranges,
+        )