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

invokeai/app/invocations/z_image_text_encoder.py

@@ -0,0 +1,197 @@
+from contextlib import ExitStack
+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,
+    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
+from invokeai.backend.patches.layer_patcher import LayerPatcher
+from invokeai.backend.patches.lora_conversions.z_image_lora_constants import Z_IMAGE_LORA_QWEN3_PREFIX
+from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
+    ConditioningFieldData,
+    ZImageConditioningInfo,
+)
+from invokeai.backend.util.devices import TorchDevice
+
+# Z-Image max sequence length based on diffusers default
+Z_IMAGE_MAX_SEQ_LEN = 512
+
+
+@invocation(
+    "z_image_text_encoder",
+    title="Prompt - Z-Image",
+    tags=["prompt", "conditioning", "z-image"],
+    category="conditioning",
+    version="1.1.0",
+    classification=Classification.Prototype,
+)
+class ZImageTextEncoderInvocation(BaseInvocation):
+    """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(
+            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.
+
+        Based on the ZImagePipeline._encode_prompt method from diffusers.
+        """
+        prompt = self.prompt
+        device = TorchDevice.choose_torch_device()
+
+        text_encoder_info = context.models.load(self.qwen3_encoder.text_encoder)
+        tokenizer_info = context.models.load(self.qwen3_encoder.tokenizer)
+
+        with ExitStack() as exit_stack:
+            (_, text_encoder) = exit_stack.enter_context(text_encoder_info.model_on_device())
+            (_, tokenizer) = exit_stack.enter_context(tokenizer_info.model_on_device())
+
+            # Apply LoRA models to the text encoder
+            lora_dtype = TorchDevice.choose_bfloat16_safe_dtype(device)
+            exit_stack.enter_context(
+                LayerPatcher.apply_smart_model_patches(
+                    model=text_encoder,
+                    patches=self._lora_iterator(context),
+                    prefix=Z_IMAGE_LORA_QWEN3_PREFIX,
+                    dtype=lora_dtype,
+                )
+            )
+
+            context.util.signal_progress("Running Qwen3 text encoder")
+            if not isinstance(text_encoder, PreTrainedModel):
+                raise TypeError(
+                    f"Expected PreTrainedModel for text encoder, got {type(text_encoder).__name__}. "
+                    "The Qwen3 encoder model may be corrupted or incompatible."
+                )
+            if not isinstance(tokenizer, PreTrainedTokenizerBase):
+                raise TypeError(
+                    f"Expected PreTrainedTokenizerBase for tokenizer, got {type(tokenizer).__name__}. "
+                    "The Qwen3 tokenizer may be corrupted or incompatible."
+                )
+
+            # Apply chat template similar to diffusers ZImagePipeline
+            # The chat template formats the prompt for the Qwen3 model
+            try:
+                prompt_formatted = tokenizer.apply_chat_template(
+                    [{"role": "user", "content": prompt}],
+                    tokenize=False,
+                    add_generation_prompt=True,
+                    enable_thinking=True,
+                )
+            except (AttributeError, TypeError) as e:
+                # Fallback if tokenizer doesn't support apply_chat_template or enable_thinking
+                context.logger.warning(f"Chat template failed ({e}), using raw prompt.")
+                prompt_formatted = prompt
+
+            # Tokenize the formatted prompt
+            text_inputs = tokenizer(
+                prompt_formatted,
+                padding="max_length",
+                max_length=max_seq_len,
+                truncation=True,
+                return_attention_mask=True,
+                return_tensors="pt",
+            )
+
+            text_input_ids = text_inputs.input_ids
+            attention_mask = text_inputs.attention_mask
+            if not isinstance(text_input_ids, torch.Tensor):
+                raise TypeError(
+                    f"Expected torch.Tensor for input_ids, got {type(text_input_ids).__name__}. "
+                    "Tokenizer returned unexpected type."
+                )
+            if not isinstance(attention_mask, torch.Tensor):
+                raise TypeError(
+                    f"Expected torch.Tensor for attention_mask, got {type(attention_mask).__name__}. "
+                    "Tokenizer returned unexpected type."
+                )
+
+            # Check for truncation
+            untruncated_ids = tokenizer(prompt_formatted, padding="longest", return_tensors="pt").input_ids
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = tokenizer.batch_decode(untruncated_ids[:, max_seq_len - 1 : -1])
+                context.logger.warning(
+                    f"The following part of your input was truncated because `max_sequence_length` is set to "
+                    f"{max_seq_len} tokens: {removed_text}"
+                )
+
+            # Get hidden states from the text encoder
+            # Use the second-to-last hidden state like diffusers does
+            prompt_mask = attention_mask.to(device).bool()
+            outputs = text_encoder(
+                text_input_ids.to(device),
+                attention_mask=prompt_mask,
+                output_hidden_states=True,
+            )
+
+            # Validate hidden_states output
+            if not hasattr(outputs, "hidden_states") or outputs.hidden_states is None:
+                raise RuntimeError(
+                    "Text encoder did not return hidden_states. "
+                    "Ensure output_hidden_states=True is supported by this model."
+                )
+            if len(outputs.hidden_states) < 2:
+                raise RuntimeError(
+                    f"Expected at least 2 hidden states from text encoder, got {len(outputs.hidden_states)}. "
+                    "This may indicate an incompatible model or configuration."
+                )
+            prompt_embeds = outputs.hidden_states[-2]
+
+            # Z-Image expects a 2D tensor [seq_len, hidden_dim] with only valid tokens
+            # Based on diffusers ZImagePipeline implementation:
+            # embeddings_list.append(prompt_embeds[i][prompt_masks[i]])
+            # Since batch_size=1, we take the first item and filter by mask
+            prompt_embeds = prompt_embeds[0][prompt_mask[0]]
+
+        if not isinstance(prompt_embeds, torch.Tensor):
+            raise TypeError(
+                f"Expected torch.Tensor for prompt embeddings, got {type(prompt_embeds).__name__}. "
+                "Text encoder returned unexpected type."
+            )
+        return prompt_embeds
+
+    def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[ModelPatchRaw, float]]:
+        """Iterate over LoRA models to apply to the Qwen3 text encoder."""
+        for lora in self.qwen3_encoder.loras:
+            lora_info = context.models.load(lora.lora)
+            if not isinstance(lora_info.model, ModelPatchRaw):
+                raise TypeError(
+                    f"Expected ModelPatchRaw for LoRA '{lora.lora.key}', got {type(lora_info.model).__name__}. "
+                    "The LoRA model may be corrupted or incompatible."
+                )
+            yield (lora_info.model, lora.weight)
+            del lora_info

invokeai/app/services/model_install/model_install_common.py

@@ -85,9 +85,12 @@ class LocalModelSource(StringLikeSource):
 
 class HFModelSource(StringLikeSource):
     """
-    A HuggingFace repo_id with optional variant, sub-folder and access token.
+    A HuggingFace repo_id with optional variant, sub-folder(s) and access token.
     Note that the variant option, if not provided to the constructor, will default to fp16, which is
     what people (almost) always want.
+
+    The subfolder can be a single path or multiple paths joined by '+' (e.g., "text_encoder+tokenizer").
+    When multiple subfolders are specified, all of them will be downloaded and combined into the model directory.
     """
 
     repo_id: str
@@ -103,6 +106,16 @@ class HFModelSource(StringLikeSource):
             raise ValueError(f"{v}: invalid repo_id format")
         return v
 
+    @property
+    def subfolders(self) -> list[Path]:
+        """Return list of subfolders (supports '+' separated multiple subfolders)."""
+        if self.subfolder is None:
+            return []
+        subfolder_str = self.subfolder.as_posix()
+        if "+" in subfolder_str:
+            return [Path(s.strip()) for s in subfolder_str.split("+")]
+        return [self.subfolder]
+
     def __str__(self) -> str:
         """Return string version of repoid when string rep needed."""
         base: str = self.repo_id

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 @@ class ModelInstallService(ModelInstallServiceBase):
         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 @@ class ModelInstallService(ModelInstallServiceBase):
             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():
@@ -417,10 +435,15 @@ class ModelInstallService(ModelInstallServiceBase):
         model_path.mkdir(parents=True, exist_ok=True)
         model_source = self._guess_source(str(source))
         remote_files, _ = self._remote_files_from_source(model_source)
+        # Handle multiple subfolders for HFModelSource
+        subfolders = model_source.subfolders if isinstance(model_source, HFModelSource) else []
         job = self._multifile_download(
             dest=model_path,
             remote_files=remote_files,
-            subfolder=model_source.subfolder if isinstance(model_source, HFModelSource) else None,
+            subfolder=model_source.subfolder
+            if isinstance(model_source, HFModelSource) and len(subfolders) <= 1
+            else None,
+            subfolders=subfolders if len(subfolders) > 1 else None,
         )
         files_string = "file" if len(remote_files) == 1 else "files"
         self._logger.info(f"Queuing model download: {source} ({len(remote_files)} {files_string})")
@@ -438,10 +461,13 @@ class ModelInstallService(ModelInstallServiceBase):
         if isinstance(source, HFModelSource):
             metadata = HuggingFaceMetadataFetch(self._session).from_id(source.repo_id, source.variant)
             assert isinstance(metadata, ModelMetadataWithFiles)
+            # Use subfolders property which handles '+' separated multiple subfolders
+            subfolders = source.subfolders
             return (
                 metadata.download_urls(
                     variant=source.variant or self._guess_variant(),
-                    subfolder=source.subfolder,
+                    subfolder=source.subfolder if len(subfolders) <= 1 else None,
+                    subfolders=subfolders if len(subfolders) > 1 else None,
                     session=self._session,
                 ),
                 metadata,
@@ -492,6 +518,39 @@ class ModelInstallService(ModelInstallServiceBase):
         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:
@@ -521,7 +580,7 @@ class ModelInstallService(ModelInstallServiceBase):
             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")
@@ -566,7 +625,7 @@ class ModelInstallService(ModelInstallServiceBase):
         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."""
@@ -741,10 +800,13 @@ class ModelInstallService(ModelInstallServiceBase):
         install_job._install_tmpdir = destdir
         install_job.total_bytes = sum((x.size or 0) for x in remote_files)
 
+        # Handle multiple subfolders for HFModelSource
+        subfolders = source.subfolders if isinstance(source, HFModelSource) else []
         multifile_job = self._multifile_download(
             remote_files=remote_files,
             dest=destdir,
-            subfolder=source.subfolder if isinstance(source, HFModelSource) else None,
+            subfolder=source.subfolder if isinstance(source, HFModelSource) and len(subfolders) <= 1 else None,
+            subfolders=subfolders if len(subfolders) > 1 else None,
             access_token=source.access_token,
             submit_job=False,  # Important! Don't submit the job until we have set our _download_cache dict
         )
@@ -771,6 +833,7 @@ class ModelInstallService(ModelInstallServiceBase):
         remote_files: List[RemoteModelFile],
         dest: Path,
         subfolder: Optional[Path] = None,
+        subfolders: Optional[List[Path]] = None,
         access_token: Optional[str] = None,
         submit_job: bool = True,
     ) -> MultiFileDownloadJob:
@@ -778,24 +841,61 @@ class ModelInstallService(ModelInstallServiceBase):
         # we are installing the "vae" subfolder, we do not want to create an additional folder level, such
         # as "sdxl-turbo/vae", nor do we want to put the contents of the vae folder directly into "sdxl-turbo".
         # So what we do is to synthesize a folder named "sdxl-turbo_vae" here.
-        if subfolder:
+        #
+        # For multiple subfolders (e.g., text_encoder+tokenizer), we create a combined folder name
+        # (e.g., sdxl-turbo_text_encoder_tokenizer) and keep each subfolder's contents in its own
+        # subdirectory within the model folder.
+
+        if subfolders and len(subfolders) > 1:
+            # Multiple subfolders: create combined name and keep subfolder structure
+            top = Path(remote_files[0].path.parts[0])  # e.g. "Z-Image-Turbo/"
+            subfolder_names = [sf.name.replace("/", "_").replace("\\", "_") for sf in subfolders]
+            combined_name = "_".join(subfolder_names)
+            path_to_add = Path(f"{top}_{combined_name}")
+
+            parts: List[RemoteModelFile] = []
+            for model_file in remote_files:
+                assert model_file.size is not None
+                # Determine which subfolder this file belongs to
+                file_path = model_file.path
+                new_path: Optional[Path] = None
+                for sf in subfolders:
+                    try:
+                        # Try to get relative path from this subfolder
+                        relative = file_path.relative_to(top / sf)
+                        # Keep the subfolder name as a subdirectory
+                        new_path = path_to_add / sf.name / relative
+                        break
+                    except ValueError:
+                        continue
+
+                if new_path is None:
+                    # File doesn't match any subfolder, keep original path structure
+                    new_path = path_to_add / file_path.relative_to(top)
+
+                parts.append(RemoteModelFile(url=model_file.url, path=new_path))
+        elif subfolder:
+            # Single subfolder: flatten into renamed folder
             top = Path(remote_files[0].path.parts[0])  # e.g. "sdxl-turbo/"
             path_to_remove = top / subfolder  # sdxl-turbo/vae/
             subfolder_rename = subfolder.name.replace("/", "_").replace("\\", "_")
             path_to_add = Path(f"{top}_{subfolder_rename}")
-        else:
-            path_to_remove = Path(".")
-            path_to_add = Path(".")
-
-        parts: List[RemoteModelFile] = []
-        for model_file in remote_files:
-            assert model_file.size is not None
-            parts.append(
-                RemoteModelFile(
-                    url=model_file.url,  # if a subfolder, then sdxl-turbo_vae/config.json
-                    path=path_to_add / model_file.path.relative_to(path_to_remove),
+
+            parts = []
+            for model_file in remote_files:
+                assert model_file.size is not None
+                parts.append(
+                    RemoteModelFile(
+                        url=model_file.url,
+                        path=path_to_add / model_file.path.relative_to(path_to_remove),
+                    )
                 )
-            )
+        else:
+            # No subfolder specified - pass through unchanged
+            parts = []
+            for model_file in remote_files:
+                assert model_file.size is not None
+                parts.append(RemoteModelFile(url=model_file.url, path=model_file.path))
 
         return self._download_queue.multifile_download(
             parts=parts,

invokeai/app/services/model_records/model_records_base.py

@@ -138,6 +138,18 @@ class ModelRecordServiceBase(ABC):
         """
         pass
 
+    @abstractmethod
+    def replace_model(self, key: str, new_config: AnyModelConfig) -> AnyModelConfig:
+        """
+        Replace the model record entirely, returning the new record.
+
+        This is used when we re-identify a model and have a new config object.
+
+        :param key: Unique key for the model to be updated.
+        :param new_config: The new model config to write.
+        """
+        pass
+
     @abstractmethod
     def get_model(self, key: str) -> AnyModelConfig:
         """

invokeai/app/services/model_records/model_records_sql.py

@@ -179,6 +179,23 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
 
         return self.get_model(key)
 
+    def replace_model(self, key: str, new_config: AnyModelConfig) -> AnyModelConfig:
+        if key != new_config.key:
+            raise ValueError("key does not match new_config.key")
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                UPDATE models
+                SET
+                    config=?
+                WHERE id=?;
+                """,
+                (new_config.model_dump_json(), key),
+            )
+            if cursor.rowcount == 0:
+                raise UnknownModelException("model not found")
+        return self.get_model(key)
+
     def get_model(self, key: str) -> AnyModelConfig:
         """
         Retrieve the ModelConfigBase instance for the indicated model.