InvokeAI 6.10.0rc1__py3-none-any.whl → 6.11.0__py3-none-any.whl

invokeai/backend/flux/extensions/dype_extension.py

@@ -0,0 +1,91 @@
+"""DyPE extension for FLUX denoising pipeline."""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from invokeai.backend.flux.dype.base import DyPEConfig
+from invokeai.backend.flux.dype.embed import DyPEEmbedND
+
+if TYPE_CHECKING:
+    from invokeai.backend.flux.model import Flux
+
+
+@dataclass
+class DyPEExtension:
+    """Extension for Dynamic Position Extrapolation in FLUX models.
+
+    This extension manages the patching of the FLUX model's position embedder
+    and updates the step state during denoising.
+
+    Usage:
+        1. Create extension with config and target dimensions
+        2. Call patch_model() to replace pe_embedder with DyPE version
+        3. Call update_step_state() before each denoising step
+        4. Call restore_model() after denoising to restore original embedder
+    """
+
+    config: DyPEConfig
+    target_height: int
+    target_width: int
+
+    def patch_model(self, model: "Flux") -> tuple[DyPEEmbedND, object]:
+        """Patch the model's position embedder with DyPE version.
+
+        Args:
+            model: The FLUX model to patch
+
+        Returns:
+            Tuple of (new DyPE embedder, original embedder for restoration)
+        """
+        original_embedder = model.pe_embedder
+
+        dype_embedder = DyPEEmbedND.from_embednd(
+            embed_nd=original_embedder,
+            dype_config=self.config,
+        )
+
+        # Set initial state
+        dype_embedder.set_step_state(
+            sigma=1.0,
+            height=self.target_height,
+            width=self.target_width,
+        )
+
+        # Replace the embedder
+        model.pe_embedder = dype_embedder
+
+        return dype_embedder, original_embedder
+
+    def update_step_state(
+        self,
+        embedder: DyPEEmbedND,
+        timestep: float,
+        timestep_index: int,
+        total_steps: int,
+    ) -> None:
+        """Update the step state in the DyPE embedder.
+
+        This should be called before each denoising step to update the
+        current noise level for timestep-dependent scaling.
+
+        Args:
+            embedder: The DyPE embedder to update
+            timestep: Current timestep value (sigma/noise level)
+            timestep_index: Current step index (0-based)
+            total_steps: Total number of denoising steps
+        """
+        embedder.set_step_state(
+            sigma=timestep,
+            height=self.target_height,
+            width=self.target_width,
+        )
+
+    @staticmethod
+    def restore_model(model: "Flux", original_embedder: object) -> None:
+        """Restore the original position embedder.
+
+        Args:
+            model: The FLUX model to restore
+            original_embedder: The original embedder saved from patch_model()
+        """
+        model.pe_embedder = original_embedder

invokeai/backend/flux/schedulers.py

@@ -0,0 +1,62 @@
+"""Flow Matching scheduler definitions and mapping.
+
+This module provides the scheduler types and mapping for Flow Matching models
+(Flux and Z-Image), supporting multiple schedulers from the diffusers library.
+"""
+
+from typing import Literal, Type
+
+from diffusers import (
+    FlowMatchEulerDiscreteScheduler,
+    FlowMatchHeunDiscreteScheduler,
+)
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+# Note: FlowMatchLCMScheduler may not be available in all diffusers versions
+try:
+    from diffusers import FlowMatchLCMScheduler
+
+    _HAS_LCM = True
+except ImportError:
+    _HAS_LCM = False
+
+# Scheduler name literal type for type checking
+FLUX_SCHEDULER_NAME_VALUES = Literal["euler", "heun", "lcm"]
+
+# Human-readable labels for the UI
+FLUX_SCHEDULER_LABELS: dict[str, str] = {
+    "euler": "Euler",
+    "heun": "Heun (2nd order)",
+    "lcm": "LCM",
+}
+
+# Mapping from scheduler names to scheduler classes
+FLUX_SCHEDULER_MAP: dict[str, Type[SchedulerMixin]] = {
+    "euler": FlowMatchEulerDiscreteScheduler,
+    "heun": FlowMatchHeunDiscreteScheduler,
+}
+
+if _HAS_LCM:
+    FLUX_SCHEDULER_MAP["lcm"] = FlowMatchLCMScheduler
+
+
+# Z-Image scheduler types (same schedulers as Flux, both use Flow Matching)
+# Note: Z-Image-Turbo is optimized for ~8 steps with Euler, but other schedulers
+# can be used for experimentation.
+ZIMAGE_SCHEDULER_NAME_VALUES = Literal["euler", "heun", "lcm"]
+
+# Human-readable labels for the UI
+ZIMAGE_SCHEDULER_LABELS: dict[str, str] = {
+    "euler": "Euler",
+    "heun": "Heun (2nd order)",
+    "lcm": "LCM",
+}
+
+# Mapping from scheduler names to scheduler classes (same as Flux)
+ZIMAGE_SCHEDULER_MAP: dict[str, Type[SchedulerMixin]] = {
+    "euler": FlowMatchEulerDiscreteScheduler,
+    "heun": FlowMatchHeunDiscreteScheduler,
+}
+
+if _HAS_LCM:
+    ZIMAGE_SCHEDULER_MAP["lcm"] = FlowMatchLCMScheduler

invokeai/backend/flux/util.py

@@ -5,7 +5,7 @@ from typing import Literal
 
 from invokeai.backend.flux.model import FluxParams
 from invokeai.backend.flux.modules.autoencoder import AutoEncoderParams
-from invokeai.backend.model_manager.taxonomy import AnyVariant, FluxVariantType
+from invokeai.backend.model_manager.taxonomy import AnyVariant, Flux2VariantType, FluxVariantType
 
 
 @dataclass
@@ -46,6 +46,8 @@ _flux_max_seq_lengths: dict[AnyVariant, Literal[256, 512]] = {
     FluxVariantType.Dev: 512,
     FluxVariantType.DevFill: 512,
     FluxVariantType.Schnell: 256,
+    Flux2VariantType.Klein4B: 512,
+    Flux2VariantType.Klein9B: 512,
 }
 
 
@@ -117,6 +119,38 @@ _flux_transformer_params: dict[AnyVariant, FluxParams] = {
         qkv_bias=True,
         guidance_embed=True,
     ),
+    # Flux2 Klein 4B uses Qwen3 4B text encoder with stacked embeddings from layers [9, 18, 27]
+    # The context_in_dim is 3 * hidden_size of Qwen3 (3 * 2560 = 7680)
+    Flux2VariantType.Klein4B: FluxParams(
+        in_channels=64,
+        vec_in_dim=2560,  # Qwen3-4B hidden size (used for pooled output)
+        context_in_dim=7680,  # 3 layers * 2560 = 7680 for Qwen3-4B
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=19,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=True,
+    ),
+    # Flux2 Klein 9B uses Qwen3 8B text encoder with stacked embeddings from layers [9, 18, 27]
+    # The context_in_dim is 3 * hidden_size of Qwen3 (3 * 4096 = 12288)
+    Flux2VariantType.Klein9B: FluxParams(
+        in_channels=64,
+        vec_in_dim=4096,  # Qwen3-8B hidden size (used for pooled output)
+        context_in_dim=12288,  # 3 layers * 4096 = 12288 for Qwen3-8B
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=19,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=True,
+    ),
 }
 
 

invokeai/backend/flux2/__init__.py

@@ -0,0 +1,4 @@
+"""FLUX.2 backend modules.
+
+This package contains modules specific to FLUX.2 models (e.g., Klein).
+"""

invokeai/backend/flux2/denoise.py

@@ -0,0 +1,280 @@
+"""Flux2 Klein Denoising Function.
+
+This module provides the denoising function for FLUX.2 Klein models,
+which use Qwen3 as the text encoder instead of CLIP+T5.
+"""
+
+import math
+from typing import Any, Callable
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension
+from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
+
+
+def denoise(
+    model: torch.nn.Module,
+    # model input
+    img: torch.Tensor,
+    img_ids: torch.Tensor,
+    txt: torch.Tensor,
+    txt_ids: torch.Tensor,
+    # sampling parameters
+    timesteps: list[float],
+    step_callback: Callable[[PipelineIntermediateState], None],
+    cfg_scale: list[float],
+    # Negative conditioning for CFG
+    neg_txt: torch.Tensor | None = None,
+    neg_txt_ids: torch.Tensor | None = None,
+    # Scheduler for stepping (e.g., FlowMatchEulerDiscreteScheduler, FlowMatchHeunDiscreteScheduler)
+    scheduler: Any = None,
+    # Dynamic shifting parameter for FLUX.2 Klein (computed from image resolution)
+    mu: float | None = None,
+    # Inpainting extension for merging latents during denoising
+    inpaint_extension: RectifiedFlowInpaintExtension | None = None,
+    # Reference image conditioning (multi-reference image editing)
+    img_cond_seq: torch.Tensor | None = None,
+    img_cond_seq_ids: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Denoise latents using a FLUX.2 Klein transformer model.
+
+    This is a simplified denoise function for FLUX.2 Klein models that uses
+    the diffusers Flux2Transformer2DModel interface.
+
+    Note: FLUX.2 Klein has guidance_embeds=False, so no guidance parameter is used.
+    CFG is applied externally using negative conditioning when cfg_scale != 1.0.
+
+    Args:
+        model: The Flux2Transformer2DModel from diffusers.
+        img: Packed latent image tensor of shape (B, seq_len, channels).
+        img_ids: Image position IDs tensor.
+        txt: Text encoder hidden states (Qwen3 embeddings).
+        txt_ids: Text position IDs tensor.
+        timesteps: List of timesteps for denoising schedule (linear sigmas from 1.0 to 1/n).
+        step_callback: Callback function for progress updates.
+        cfg_scale: List of CFG scale values per step.
+        neg_txt: Negative text embeddings for CFG (optional).
+        neg_txt_ids: Negative text position IDs (optional).
+        scheduler: Optional diffusers scheduler (Euler, Heun, LCM). If None, uses manual Euler.
+        mu: Dynamic shifting parameter computed from image resolution. Required when scheduler
+            has use_dynamic_shifting=True.
+
+    Returns:
+        Denoised latent tensor.
+    """
+    total_steps = len(timesteps) - 1
+
+    # Store original sequence length for extracting output later (before concatenating reference images)
+    original_seq_len = img.shape[1]
+
+    # Concatenate reference image conditioning if provided (multi-reference image editing)
+    if img_cond_seq is not None and img_cond_seq_ids is not None:
+        img = torch.cat([img, img_cond_seq], dim=1)
+        img_ids = torch.cat([img_ids, img_cond_seq_ids], dim=1)
+
+    # Klein has guidance_embeds=False, but the transformer forward() still requires a guidance tensor
+    # We pass a dummy value (1.0) since it won't affect the output when guidance_embeds=False
+    guidance = torch.full((img.shape[0],), 1.0, device=img.device, dtype=img.dtype)
+
+    # Use scheduler if provided
+    use_scheduler = scheduler is not None
+    if use_scheduler:
+        # Set up scheduler with sigmas and mu for dynamic shifting
+        # Convert timesteps (0-1 range) to sigmas for the scheduler
+        # The scheduler will apply dynamic shifting internally using mu (if enabled in scheduler config)
+        sigmas = np.array(timesteps[:-1], dtype=np.float32)  # Exclude final 0.0
+
+        # Pass mu if provided - it will only be used if scheduler has use_dynamic_shifting=True
+        if mu is not None:
+            scheduler.set_timesteps(sigmas=sigmas.tolist(), mu=mu, device=img.device)
+        else:
+            scheduler.set_timesteps(sigmas=sigmas.tolist(), device=img.device)
+        num_scheduler_steps = len(scheduler.timesteps)
+        is_heun = hasattr(scheduler, "state_in_first_order")
+        user_step = 0
+
+        pbar = tqdm(total=total_steps, desc="Denoising")
+        for step_index in range(num_scheduler_steps):
+            timestep = scheduler.timesteps[step_index]
+            # Convert scheduler timestep (0-1000) to normalized (0-1) for the model
+            t_curr = timestep.item() / scheduler.config.num_train_timesteps
+            t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
+
+            # Track if we're in first or second order step (for Heun)
+            in_first_order = scheduler.state_in_first_order if is_heun else True
+
+            # Run the transformer model (matching diffusers: guidance=guidance, return_dict=False)
+            output = model(
+                hidden_states=img,
+                encoder_hidden_states=txt,
+                timestep=t_vec,
+                img_ids=img_ids,
+                txt_ids=txt_ids,
+                guidance=guidance,
+                return_dict=False,
+            )
+
+            # Extract the sample from the output (return_dict=False returns tuple)
+            pred = output[0] if isinstance(output, tuple) else output
+
+            step_cfg_scale = cfg_scale[min(user_step, len(cfg_scale) - 1)]
+
+            # Apply CFG if scale is not 1.0
+            if not math.isclose(step_cfg_scale, 1.0):
+                if neg_txt is None:
+                    raise ValueError("Negative text conditioning is required when cfg_scale is not 1.0.")
+
+                neg_output = model(
+                    hidden_states=img,
+                    encoder_hidden_states=neg_txt,
+                    timestep=t_vec,
+                    img_ids=img_ids,
+                    txt_ids=neg_txt_ids if neg_txt_ids is not None else txt_ids,
+                    guidance=guidance,
+                    return_dict=False,
+                )
+
+                neg_pred = neg_output[0] if isinstance(neg_output, tuple) else neg_output
+                pred = neg_pred + step_cfg_scale * (pred - neg_pred)
+
+            # Use scheduler.step() for the update
+            step_output = scheduler.step(model_output=pred, timestep=timestep, sample=img)
+            img = step_output.prev_sample
+
+            # Get t_prev for inpainting (next sigma value)
+            if step_index + 1 < len(scheduler.sigmas):
+                t_prev = scheduler.sigmas[step_index + 1].item()
+            else:
+                t_prev = 0.0
+
+            # Apply inpainting merge at each step
+            if inpaint_extension is not None:
+                # Separate the generated latents from the reference conditioning
+                gen_img = img[:, :original_seq_len, :]
+                ref_img = img[:, original_seq_len:, :]
+
+                # Merge only the generated part
+                gen_img = inpaint_extension.merge_intermediate_latents_with_init_latents(gen_img, t_prev)
+
+                # Concatenate back together
+                img = torch.cat([gen_img, ref_img], dim=1)
+
+            # For Heun, only increment user step after second-order step completes
+            if is_heun:
+                if not in_first_order:
+                    user_step += 1
+                    if user_step <= total_steps:
+                        pbar.update(1)
+                        preview_img = img - t_curr * pred
+                        if inpaint_extension is not None:
+                            preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents(
+                                preview_img, 0.0
+                            )
+                        step_callback(
+                            PipelineIntermediateState(
+                                step=user_step,
+                                order=2,
+                                total_steps=total_steps,
+                                timestep=int(t_curr * 1000),
+                                latents=preview_img,
+                            ),
+                        )
+            else:
+                user_step += 1
+                if user_step <= total_steps:
+                    pbar.update(1)
+                    preview_img = img - t_curr * pred
+                    if inpaint_extension is not None:
+                        preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents(preview_img, 0.0)
+                    # Extract only the generated image portion for preview (exclude reference images)
+                    callback_latents = preview_img[:, :original_seq_len, :] if img_cond_seq is not None else preview_img
+                    step_callback(
+                        PipelineIntermediateState(
+                            step=user_step,
+                            order=1,
+                            total_steps=total_steps,
+                            timestep=int(t_curr * 1000),
+                            latents=callback_latents,
+                        ),
+                    )
+
+        pbar.close()
+    else:
+        # Manual Euler stepping (original behavior)
+        for step_index, (t_curr, t_prev) in tqdm(list(enumerate(zip(timesteps[:-1], timesteps[1:], strict=True)))):
+            t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
+
+            # Run the transformer model (matching diffusers: guidance=guidance, return_dict=False)
+            output = model(
+                hidden_states=img,
+                encoder_hidden_states=txt,
+                timestep=t_vec,
+                img_ids=img_ids,
+                txt_ids=txt_ids,
+                guidance=guidance,
+                return_dict=False,
+            )
+
+            # Extract the sample from the output (return_dict=False returns tuple)
+            pred = output[0] if isinstance(output, tuple) else output
+
+            step_cfg_scale = cfg_scale[step_index]
+
+            # Apply CFG if scale is not 1.0
+            if not math.isclose(step_cfg_scale, 1.0):
+                if neg_txt is None:
+                    raise ValueError("Negative text conditioning is required when cfg_scale is not 1.0.")
+
+                neg_output = model(
+                    hidden_states=img,
+                    encoder_hidden_states=neg_txt,
+                    timestep=t_vec,
+                    img_ids=img_ids,
+                    txt_ids=neg_txt_ids if neg_txt_ids is not None else txt_ids,
+                    guidance=guidance,
+                    return_dict=False,
+                )
+
+                neg_pred = neg_output[0] if isinstance(neg_output, tuple) else neg_output
+                pred = neg_pred + step_cfg_scale * (pred - neg_pred)
+
+            # Euler step
+            preview_img = img - t_curr * pred
+            img = img + (t_prev - t_curr) * pred
+
+            # Apply inpainting merge at each step
+            if inpaint_extension is not None:
+                # Separate the generated latents from the reference conditioning
+                gen_img = img[:, :original_seq_len, :]
+                ref_img = img[:, original_seq_len:, :]
+
+                # Merge only the generated part
+                gen_img = inpaint_extension.merge_intermediate_latents_with_init_latents(gen_img, t_prev)
+
+                # Concatenate back together
+                img = torch.cat([gen_img, ref_img], dim=1)
+
+                # Handling preview images
+                preview_gen = preview_img[:, :original_seq_len, :]
+                preview_gen = inpaint_extension.merge_intermediate_latents_with_init_latents(preview_gen, 0.0)
+
+            # Extract only the generated image portion for preview (exclude reference images)
+            callback_latents = preview_img[:, :original_seq_len, :] if img_cond_seq is not None else preview_img
+            step_callback(
+                PipelineIntermediateState(
+                    step=step_index + 1,
+                    order=1,
+                    total_steps=total_steps,
+                    timestep=int(t_curr),
+                    latents=callback_latents,
+                ),
+            )
+
+    # Extract only the generated image portion (exclude concatenated reference images)
+    if img_cond_seq is not None:
+        img = img[:, :original_seq_len, :]
+
+    return img