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

invokeai/backend/patches/lora_conversions/z_image_lora_conversion_utils.py

@@ -0,0 +1,189 @@
+"""Z-Image LoRA conversion utilities.
+
+Z-Image uses S3-DiT transformer architecture with Qwen3 text encoder.
+LoRAs for Z-Image typically follow the diffusers PEFT format.
+"""
+
+from typing import Dict
+
+import torch
+
+from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
+from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict
+from invokeai.backend.patches.lora_conversions.z_image_lora_constants import (
+    Z_IMAGE_LORA_QWEN3_PREFIX,
+    Z_IMAGE_LORA_TRANSFORMER_PREFIX,
+)
+from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+
+
+def is_state_dict_likely_z_image_lora(state_dict: dict[str | int, torch.Tensor]) -> bool:
+    """Checks if the provided state dict is likely a Z-Image LoRA.
+
+    Z-Image LoRAs can have keys for transformer and/or Qwen3 text encoder.
+    They may use various prefixes depending on the training framework.
+    """
+    str_keys = [k for k in state_dict.keys() if isinstance(k, str)]
+
+    # Check for Z-Image transformer keys (S3-DiT architecture)
+    # Various training frameworks use different prefixes
+    has_transformer_keys = any(
+        k.startswith(
+            (
+                "transformer.",
+                "base_model.model.transformer.",
+                "diffusion_model.",
+            )
+        )
+        for k in str_keys
+    )
+
+    # Check for Qwen3 text encoder keys
+    has_qwen3_keys = any(k.startswith(("text_encoder.", "base_model.model.text_encoder.")) for k in str_keys)
+
+    return has_transformer_keys or has_qwen3_keys
+
+
+def lora_model_from_z_image_state_dict(
+    state_dict: Dict[str, torch.Tensor], alpha: float | None = None
+) -> ModelPatchRaw:
+    """Convert a Z-Image LoRA state dict to a ModelPatchRaw.
+
+    Z-Image LoRAs can contain layers for:
+    - Transformer (S3-DiT architecture)
+    - Qwen3 text encoder
+
+    Z-Image LoRAs may use various key prefixes depending on how they were trained:
+    - "transformer." or "base_model.model.transformer." for diffusers PEFT format
+    - "diffusion_model." for some training frameworks
+    - "text_encoder." or "base_model.model.text_encoder." for Qwen3 encoder
+
+    Args:
+        state_dict: The LoRA state dict
+        alpha: The alpha value for LoRA scaling. If None, uses rank as alpha.
+
+    Returns:
+        A ModelPatchRaw containing the LoRA layers
+    """
+    layers: dict[str, BaseLayerPatch] = {}
+
+    # Group keys by layer
+    grouped_state_dict = _group_by_layer(state_dict)
+
+    for layer_key, layer_dict in grouped_state_dict.items():
+        # Convert PEFT format keys to internal format
+        values = _get_lora_layer_values(layer_dict, alpha)
+
+        # Determine the appropriate prefix based on the layer type and clean up the key
+        clean_key = layer_key
+
+        # Handle various transformer prefixes
+        transformer_prefixes = [
+            "base_model.model.transformer.diffusion_model.",
+            "base_model.model.transformer.",
+            "transformer.diffusion_model.",
+            "transformer.",
+            "diffusion_model.",
+        ]
+
+        # Handle text encoder prefixes
+        text_encoder_prefixes = [
+            "base_model.model.text_encoder.",
+            "text_encoder.",
+        ]
+
+        is_text_encoder = False
+
+        # Check and strip text encoder prefixes first
+        for prefix in text_encoder_prefixes:
+            if layer_key.startswith(prefix):
+                clean_key = layer_key[len(prefix) :]
+                is_text_encoder = True
+                break
+
+        # If not text encoder, check transformer prefixes
+        if not is_text_encoder:
+            for prefix in transformer_prefixes:
+                if layer_key.startswith(prefix):
+                    clean_key = layer_key[len(prefix) :]
+                    break
+
+        # Apply the appropriate internal prefix
+        if is_text_encoder:
+            final_key = f"{Z_IMAGE_LORA_QWEN3_PREFIX}{clean_key}"
+        else:
+            final_key = f"{Z_IMAGE_LORA_TRANSFORMER_PREFIX}{clean_key}"
+
+        layer = any_lora_layer_from_state_dict(values)
+        layers[final_key] = layer
+
+    return ModelPatchRaw(layers=layers)
+
+
+def _get_lora_layer_values(layer_dict: dict[str, torch.Tensor], alpha: float | None) -> dict[str, torch.Tensor]:
+    """Convert layer dict keys from PEFT format to internal format."""
+    if "lora_A.weight" in layer_dict:
+        # PEFT format: lora_A.weight, lora_B.weight
+        values = {
+            "lora_down.weight": layer_dict["lora_A.weight"],
+            "lora_up.weight": layer_dict["lora_B.weight"],
+        }
+        if alpha is not None:
+            values["alpha"] = torch.tensor(alpha)
+        return values
+    elif "lora_down.weight" in layer_dict:
+        # Already in internal format
+        return layer_dict
+    else:
+        # Unknown format, return as-is
+        return layer_dict
+
+
+def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]:
+    """Groups the keys in the state dict by layer.
+
+    Z-Image LoRAs have keys like:
+    - diffusion_model.layers.17.attention.to_k.alpha
+    - diffusion_model.layers.17.attention.to_k.dora_scale
+    - diffusion_model.layers.17.attention.to_k.lora_down.weight
+    - diffusion_model.layers.17.attention.to_k.lora_up.weight
+
+    We need to group these by the full layer path (e.g., diffusion_model.layers.17.attention.to_k)
+    and extract the suffix (alpha, dora_scale, lora_down.weight, lora_up.weight).
+    """
+    layer_dict: dict[str, dict[str, torch.Tensor]] = {}
+
+    # Known suffixes that indicate the end of a layer name
+    known_suffixes = [
+        ".lora_A.weight",
+        ".lora_B.weight",
+        ".lora_down.weight",
+        ".lora_up.weight",
+        ".dora_scale",
+        ".alpha",
+    ]
+
+    for key in state_dict:
+        if not isinstance(key, str):
+            continue
+
+        # Try to find a known suffix
+        layer_name = None
+        key_name = None
+        for suffix in known_suffixes:
+            if key.endswith(suffix):
+                layer_name = key[: -len(suffix)]
+                key_name = suffix[1:]  # Remove leading dot
+                break
+
+        if layer_name is None:
+            # Fallback to original logic for unknown formats
+            parts = key.rsplit(".", maxsplit=2)
+            layer_name = parts[0]
+            key_name = ".".join(parts[1:])
+
+        if layer_name not in layer_dict:
+            layer_dict[layer_name] = {}
+        layer_dict[layer_name][key_name] = state_dict[key]
+
+    return layer_dict

invokeai/backend/quantization/gguf/ggml_tensor.py

@@ -14,11 +14,42 @@ 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 and target_device from any GGMLTensor in the args
+    compute_dtype = None
+    target_device = None
+    for a in args:
+        if hasattr(a, "compute_dtype"):
+            compute_dtype = a.compute_dtype
+        if isinstance(a, torch.Tensor) and target_device is None:
+            target_device = a.device
+        if compute_dtype is not None and target_device is not None:
+            break
+    if compute_dtype is None or target_device is None:
+        for v in kwargs.values():
+            if hasattr(v, "compute_dtype") and compute_dtype is None:
+                compute_dtype = v.compute_dtype
+            if isinstance(v, torch.Tensor) and target_device is None:
+                target_device = v.device
+            if compute_dtype is not None and target_device is not None:
+                break
+
+    def process_tensor(t):
+        if hasattr(t, "get_dequantized_tensor"):
+            result = t.get_dequantized_tensor()
+            # Ensure the dequantized tensor is on the target device
+            if target_device is not None and result.device != target_device:
+                result = result.to(target_device)
+            return result
+        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 +88,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,205 @@
+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.
+
+        This uses an 'unrestricted' image self-attention approach (similar to FLUX):
+        - Image tokens can attend to ALL other image tokens (unrestricted self-attention)
+        - Image tokens attend only to their corresponding regional text
+        - Text tokens attend only to their corresponding regional image
+        - Text tokens attend to themselves
+
+        The unrestricted image self-attention allows the model to maintain global
+        coherence across regions, preventing the generation of separate/disconnected
+        images for each region.
+
+        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)
+            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
+
+        # 4. Allow unrestricted image self-attention
+        # This is the key difference from the restricted approach - all image tokens
+        # can attend to each other, which helps maintain global coherence across regions
+        regional_attention_mask[:img_seq_len, :img_seq_len] = 1.0
+
+        # 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)