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

invokeai/app/invocations/z_image_image_to_latents.py

@@ -0,0 +1,102 @@
+from typing import Union
+
+import einops
+import torch
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.fields import (
+    FieldDescriptions,
+    ImageField,
+    Input,
+    InputField,
+    WithBoard,
+    WithMetadata,
+)
+from invokeai.app.invocations.model import VAEField
+from invokeai.app.invocations.primitives import LatentsOutput
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder
+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
+
+# Z-Image can use either the Diffusers AutoencoderKL or the FLUX AutoEncoder
+ZImageVAE = Union[AutoencoderKL, FluxAutoEncoder]
+
+
+@invocation(
+    "z_image_i2l",
+    title="Image to Latents - Z-Image",
+    tags=["image", "latents", "vae", "i2l", "z-image"],
+    category="image",
+    version="1.1.0",
+    classification=Classification.Prototype,
+)
+class ZImageImageToLatentsInvocation(BaseInvocation, WithMetadata, WithBoard):
+    """Generates latents from an image using Z-Image VAE (supports both Diffusers and FLUX VAE)."""
+
+    image: ImageField = InputField(description="The image to encode.")
+    vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection)
+
+    @staticmethod
+    def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tensor:
+        if not isinstance(vae_info.model, (AutoencoderKL, FluxAutoEncoder)):
+            raise TypeError(
+                f"Expected AutoencoderKL or FluxAutoEncoder for Z-Image VAE, got {type(vae_info.model).__name__}. "
+                "Ensure you are using a compatible VAE model."
+            )
+
+        with vae_info.model_on_device() as (_, vae):
+            if not isinstance(vae, (AutoencoderKL, FluxAutoEncoder)):
+                raise TypeError(
+                    f"Expected AutoencoderKL or FluxAutoEncoder, got {type(vae).__name__}. "
+                    "VAE model type changed unexpectedly after loading."
+                )
+
+            vae_dtype = next(iter(vae.parameters())).dtype
+            image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)
+
+            with torch.inference_mode():
+                if isinstance(vae, FluxAutoEncoder):
+                    # FLUX VAE handles scaling internally
+                    generator = torch.Generator(device=TorchDevice.choose_torch_device()).manual_seed(0)
+                    latents = vae.encode(image_tensor, sample=True, generator=generator)
+                else:
+                    # AutoencoderKL - needs manual scaling
+                    vae.disable_tiling()
+                    image_tensor_dist = vae.encode(image_tensor).latent_dist
+                    latents: torch.Tensor = image_tensor_dist.sample().to(dtype=vae.dtype)
+
+                    # Apply scaling_factor and shift_factor from VAE config
+                    # Z-Image uses: latents = (latents - shift_factor) * scaling_factor
+                    scaling_factor = vae.config.scaling_factor
+                    shift_factor = getattr(vae.config, "shift_factor", None)
+
+                    if shift_factor is not None:
+                        latents = latents - shift_factor
+                    latents = latents * scaling_factor
+
+        return latents
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> LatentsOutput:
+        image = context.images.get_pil(self.image.image_name)
+
+        image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
+        if image_tensor.dim() == 3:
+            image_tensor = einops.rearrange(image_tensor, "c h w -> 1 c h w")
+
+        vae_info = context.models.load(self.vae.vae)
+        if not isinstance(vae_info.model, (AutoencoderKL, FluxAutoEncoder)):
+            raise TypeError(
+                f"Expected AutoencoderKL or FluxAutoEncoder for Z-Image VAE, got {type(vae_info.model).__name__}. "
+                "Ensure you are using a compatible VAE model."
+            )
+
+        context.util.signal_progress("Running VAE")
+        latents = self.vae_encode(vae_info=vae_info, image_tensor=image_tensor)
+
+        latents = latents.to("cpu")
+        name = context.tensors.save(tensor=latents)
+        return LatentsOutput.build(latents_name=name, latents=latents, seed=None)

invokeai/app/invocations/z_image_latents_to_image.py

@@ -0,0 +1,103 @@
+from contextlib import nullcontext
+from typing import Union
+
+import torch
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from einops import rearrange
+from PIL import Image
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.fields import (
+    FieldDescriptions,
+    Input,
+    InputField,
+    LatentsField,
+    WithBoard,
+    WithMetadata,
+)
+from invokeai.app.invocations.model import VAEField
+from invokeai.app.invocations.primitives import ImageOutput
+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
+
+# Z-Image can use either the Diffusers AutoencoderKL or the FLUX AutoEncoder
+ZImageVAE = Union[AutoencoderKL, FluxAutoEncoder]
+
+
+@invocation(
+    "z_image_l2i",
+    title="Latents to Image - Z-Image",
+    tags=["latents", "image", "vae", "l2i", "z-image"],
+    category="latents",
+    version="1.1.0",
+    classification=Classification.Prototype,
+)
+class ZImageLatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
+    """Generates an image from latents using Z-Image VAE (supports both Diffusers and FLUX VAE)."""
+
+    latents: LatentsField = InputField(description=FieldDescriptions.latents, input=Input.Connection)
+    vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection)
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        latents = context.tensors.load(self.latents.latents_name)
+
+        vae_info = context.models.load(self.vae.vae)
+        if not isinstance(vae_info.model, (AutoencoderKL, FluxAutoEncoder)):
+            raise TypeError(
+                f"Expected AutoencoderKL or FluxAutoEncoder for Z-Image VAE, got {type(vae_info.model).__name__}. "
+                "Ensure you are using a compatible VAE model."
+            )
+
+        is_flux_vae = isinstance(vae_info.model, FluxAutoEncoder)
+
+        # 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):
+            context.util.signal_progress("Running VAE")
+            if not isinstance(vae, (AutoencoderKL, FluxAutoEncoder)):
+                raise TypeError(
+                    f"Expected AutoencoderKL or FluxAutoEncoder, got {type(vae).__name__}. "
+                    "VAE model type changed unexpectedly after loading."
+                )
+
+            vae_dtype = next(iter(vae.parameters())).dtype
+            latents = latents.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)
+
+            # Disable tiling for AutoencoderKL
+            if isinstance(vae, AutoencoderKL):
+                vae.disable_tiling()
+
+            # Clear memory as VAE decode can request a lot
+            TorchDevice.empty_cache()
+
+            with torch.inference_mode():
+                if isinstance(vae, FluxAutoEncoder):
+                    # FLUX VAE handles scaling internally
+                    img = vae.decode(latents)
+                else:
+                    # AutoencoderKL - Apply scaling_factor and shift_factor from VAE config
+                    # Z-Image uses: latents = latents / scaling_factor + shift_factor
+                    scaling_factor = vae.config.scaling_factor
+                    shift_factor = getattr(vae.config, "shift_factor", None)
+
+                    latents = latents / scaling_factor
+                    if shift_factor is not None:
+                        latents = latents + shift_factor
+
+                    img = vae.decode(latents, return_dict=False)[0]
+
+            img = img.clamp(-1, 1)
+            img = rearrange(img[0], "c h w -> h w c")
+            img_pil = Image.fromarray((127.5 * (img + 1.0)).byte().cpu().numpy())
+
+        TorchDevice.empty_cache()
+
+        image_dto = context.images.save(image=img_pil)
+
+        return ImageOutput.build(image_dto)

invokeai/app/invocations/z_image_lora_loader.py

@@ -0,0 +1,153 @@
+from typing import Optional
+
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    invocation,
+    invocation_output,
+)
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField
+from invokeai.app.invocations.model import LoRAField, ModelIdentifierField, Qwen3EncoderField, TransformerField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType
+
+
+@invocation_output("z_image_lora_loader_output")
+class ZImageLoRALoaderOutput(BaseInvocationOutput):
+    """Z-Image LoRA Loader Output"""
+
+    transformer: Optional[TransformerField] = OutputField(
+        default=None, description=FieldDescriptions.transformer, title="Z-Image Transformer"
+    )
+    qwen3_encoder: Optional[Qwen3EncoderField] = OutputField(
+        default=None, description=FieldDescriptions.qwen3_encoder, title="Qwen3 Encoder"
+    )
+
+
+@invocation(
+    "z_image_lora_loader",
+    title="Apply LoRA - Z-Image",
+    tags=["lora", "model", "z-image"],
+    category="model",
+    version="1.0.0",
+)
+class ZImageLoRALoaderInvocation(BaseInvocation):
+    """Apply a LoRA model to a Z-Image transformer and/or Qwen3 text encoder."""
+
+    lora: ModelIdentifierField = InputField(
+        description=FieldDescriptions.lora_model,
+        title="LoRA",
+        ui_model_base=BaseModelType.ZImage,
+        ui_model_type=ModelType.LoRA,
+    )
+    weight: float = InputField(default=0.75, description=FieldDescriptions.lora_weight)
+    transformer: TransformerField | None = InputField(
+        default=None,
+        description=FieldDescriptions.transformer,
+        input=Input.Connection,
+        title="Z-Image Transformer",
+    )
+    qwen3_encoder: Qwen3EncoderField | None = InputField(
+        default=None,
+        title="Qwen3 Encoder",
+        description=FieldDescriptions.qwen3_encoder,
+        input=Input.Connection,
+    )
+
+    def invoke(self, context: InvocationContext) -> ZImageLoRALoaderOutput:
+        lora_key = self.lora.key
+
+        if not context.models.exists(lora_key):
+            raise ValueError(f"Unknown lora: {lora_key}!")
+
+        # Check for existing LoRAs with the same key.
+        if self.transformer and any(lora.lora.key == lora_key for lora in self.transformer.loras):
+            raise ValueError(f'LoRA "{lora_key}" already applied to transformer.')
+        if self.qwen3_encoder and any(lora.lora.key == lora_key for lora in self.qwen3_encoder.loras):
+            raise ValueError(f'LoRA "{lora_key}" already applied to Qwen3 encoder.')
+
+        output = ZImageLoRALoaderOutput()
+
+        # Attach LoRA layers to the models.
+        if self.transformer is not None:
+            output.transformer = self.transformer.model_copy(deep=True)
+            output.transformer.loras.append(
+                LoRAField(
+                    lora=self.lora,
+                    weight=self.weight,
+                )
+            )
+        if self.qwen3_encoder is not None:
+            output.qwen3_encoder = self.qwen3_encoder.model_copy(deep=True)
+            output.qwen3_encoder.loras.append(
+                LoRAField(
+                    lora=self.lora,
+                    weight=self.weight,
+                )
+            )
+
+        return output
+
+
+@invocation(
+    "z_image_lora_collection_loader",
+    title="Apply LoRA Collection - Z-Image",
+    tags=["lora", "model", "z-image"],
+    category="model",
+    version="1.0.0",
+)
+class ZImageLoRACollectionLoader(BaseInvocation):
+    """Applies a collection of LoRAs to a Z-Image transformer."""
+
+    loras: Optional[LoRAField | list[LoRAField]] = InputField(
+        default=None, description="LoRA models and weights. May be a single LoRA or collection.", title="LoRAs"
+    )
+
+    transformer: Optional[TransformerField] = InputField(
+        default=None,
+        description=FieldDescriptions.transformer,
+        input=Input.Connection,
+        title="Transformer",
+    )
+    qwen3_encoder: Qwen3EncoderField | None = InputField(
+        default=None,
+        title="Qwen3 Encoder",
+        description=FieldDescriptions.qwen3_encoder,
+        input=Input.Connection,
+    )
+
+    def invoke(self, context: InvocationContext) -> ZImageLoRALoaderOutput:
+        output = ZImageLoRALoaderOutput()
+        loras = self.loras if isinstance(self.loras, list) else [self.loras]
+        added_loras: list[str] = []
+
+        if self.transformer is not None:
+            output.transformer = self.transformer.model_copy(deep=True)
+
+        if self.qwen3_encoder is not None:
+            output.qwen3_encoder = self.qwen3_encoder.model_copy(deep=True)
+
+        for lora in loras:
+            if lora is None:
+                continue
+            if lora.lora.key in added_loras:
+                continue
+
+            if not context.models.exists(lora.lora.key):
+                raise Exception(f"Unknown lora: {lora.lora.key}!")
+
+            if lora.lora.base is not BaseModelType.ZImage:
+                raise ValueError(
+                    f"LoRA '{lora.lora.key}' is for {lora.lora.base.value if lora.lora.base else 'unknown'} models, "
+                    "not Z-Image models. Ensure you are using a Z-Image compatible LoRA."
+                )
+
+            added_loras.append(lora.lora.key)
+
+            if self.transformer is not None and output.transformer is not None:
+                output.transformer.loras.append(lora)
+
+            if self.qwen3_encoder is not None and output.qwen3_encoder is not None:
+                output.qwen3_encoder.loras.append(lora)
+
+        return output

invokeai/app/invocations/z_image_model_loader.py

@@ -0,0 +1,135 @@
+from typing import Optional
+
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    Classification,
+    invocation,
+    invocation_output,
+)
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField
+from invokeai.app.invocations.model import (
+    ModelIdentifierField,
+    Qwen3EncoderField,
+    TransformerField,
+    VAEField,
+)
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, SubModelType
+
+
+@invocation_output("z_image_model_loader_output")
+class ZImageModelLoaderOutput(BaseInvocationOutput):
+    """Z-Image base model loader output."""
+
+    transformer: TransformerField = OutputField(description=FieldDescriptions.transformer, title="Transformer")
+    qwen3_encoder: Qwen3EncoderField = OutputField(description=FieldDescriptions.qwen3_encoder, title="Qwen3 Encoder")
+    vae: VAEField = OutputField(description=FieldDescriptions.vae, title="VAE")
+
+
+@invocation(
+    "z_image_model_loader",
+    title="Main Model - Z-Image",
+    tags=["model", "z-image"],
+    category="model",
+    version="3.0.0",
+    classification=Classification.Prototype,
+)
+class ZImageModelLoaderInvocation(BaseInvocation):
+    """Loads a Z-Image model, outputting its submodels.
+
+    Similar to FLUX, you can mix and match components:
+    - Transformer: From Z-Image main model (GGUF quantized or Diffusers format)
+    - VAE: Separate FLUX VAE (shared with FLUX models) or from a Diffusers Z-Image model
+    - Qwen3 Encoder: Separate Qwen3Encoder model or from a Diffusers Z-Image model
+    """
+
+    model: ModelIdentifierField = InputField(
+        description=FieldDescriptions.z_image_model,
+        input=Input.Direct,
+        ui_model_base=BaseModelType.ZImage,
+        ui_model_type=ModelType.Main,
+        title="Transformer",
+    )
+
+    vae_model: Optional[ModelIdentifierField] = InputField(
+        default=None,
+        description="Standalone VAE model. Z-Image uses the same VAE as FLUX (16-channel). "
+        "If not provided, VAE will be loaded from the Qwen3 Source model.",
+        input=Input.Direct,
+        ui_model_base=BaseModelType.Flux,
+        ui_model_type=ModelType.VAE,
+        title="VAE",
+    )
+
+    qwen3_encoder_model: Optional[ModelIdentifierField] = InputField(
+        default=None,
+        description="Standalone Qwen3 Encoder model. "
+        "If not provided, encoder will be loaded from the Qwen3 Source model.",
+        input=Input.Direct,
+        ui_model_type=ModelType.Qwen3Encoder,
+        title="Qwen3 Encoder",
+    )
+
+    qwen3_source_model: Optional[ModelIdentifierField] = InputField(
+        default=None,
+        description="Diffusers Z-Image model to extract VAE and/or Qwen3 encoder from. "
+        "Use this if you don't have separate VAE/Qwen3 models. "
+        "Ignored if both VAE and Qwen3 Encoder are provided separately.",
+        input=Input.Direct,
+        ui_model_base=BaseModelType.ZImage,
+        ui_model_type=ModelType.Main,
+        ui_model_format=ModelFormat.Diffusers,
+        title="Qwen3 Source (Diffusers)",
+    )
+
+    def invoke(self, context: InvocationContext) -> ZImageModelLoaderOutput:
+        # Transformer always comes from the main model
+        transformer = self.model.model_copy(update={"submodel_type": SubModelType.Transformer})
+
+        # Determine VAE source
+        if self.vae_model is not None:
+            # Use standalone FLUX VAE
+            vae = self.vae_model.model_copy(update={"submodel_type": SubModelType.VAE})
+        elif self.qwen3_source_model is not None:
+            # Extract from Diffusers Z-Image model
+            self._validate_diffusers_format(context, self.qwen3_source_model, "Qwen3 Source")
+            vae = self.qwen3_source_model.model_copy(update={"submodel_type": SubModelType.VAE})
+        else:
+            raise ValueError(
+                "No VAE source provided. Either set 'VAE' to a FLUX VAE model, "
+                "or set 'Qwen3 Source' to a Diffusers Z-Image model."
+            )
+
+        # Determine Qwen3 Encoder source
+        if self.qwen3_encoder_model is not None:
+            # Use standalone Qwen3 Encoder
+            qwen3_tokenizer = self.qwen3_encoder_model.model_copy(update={"submodel_type": SubModelType.Tokenizer})
+            qwen3_encoder = self.qwen3_encoder_model.model_copy(update={"submodel_type": SubModelType.TextEncoder})
+        elif self.qwen3_source_model is not None:
+            # Extract from Diffusers Z-Image model
+            self._validate_diffusers_format(context, self.qwen3_source_model, "Qwen3 Source")
+            qwen3_tokenizer = self.qwen3_source_model.model_copy(update={"submodel_type": SubModelType.Tokenizer})
+            qwen3_encoder = self.qwen3_source_model.model_copy(update={"submodel_type": SubModelType.TextEncoder})
+        else:
+            raise ValueError(
+                "No Qwen3 Encoder source provided. Either set 'Qwen3 Encoder' to a standalone model, "
+                "or set 'Qwen3 Source' to a Diffusers Z-Image model."
+            )
+
+        return ZImageModelLoaderOutput(
+            transformer=TransformerField(transformer=transformer, loras=[]),
+            qwen3_encoder=Qwen3EncoderField(tokenizer=qwen3_tokenizer, text_encoder=qwen3_encoder),
+            vae=VAEField(vae=vae),
+        )
+
+    def _validate_diffusers_format(
+        self, context: InvocationContext, model: ModelIdentifierField, model_name: str
+    ) -> None:
+        """Validate that a model is in Diffusers format."""
+        config = context.models.get_config(model)
+        if config.format != ModelFormat.Diffusers:
+            raise ValueError(
+                f"The {model_name} model must be a Diffusers format Z-Image model. "
+                f"The selected model '{config.name}' is in {config.format.value} format."
+            )