InvokeAI 6.11.0__py3-none-any.whl → 6.11.1__py3-none-any.whl

invokeai/app/invocations/flux2_denoise.py

@@ -329,15 +329,13 @@ class Flux2DenoiseInvocation(BaseInvocation):
         noise_packed = pack_flux2(noise)
         x = pack_flux2(x)
 
-        # Apply BN normalization BEFORE denoising (as per diffusers Flux2KleinPipeline)
-        # BN normalization: y = (x - mean) / std
-        # This transforms latents to normalized space for the transformer
-        # IMPORTANT: Also normalize init_latents and noise for inpainting to maintain consistency
-        if bn_mean is not None and bn_std is not None:
-            x = self._bn_normalize(x, bn_mean, bn_std)
-            if init_latents_packed is not None:
-                init_latents_packed = self._bn_normalize(init_latents_packed, bn_mean, bn_std)
-            noise_packed = self._bn_normalize(noise_packed, bn_mean, bn_std)
+        # BN normalization for txt2img:
+        # - DO NOT normalize random noise (it's already N(0,1) distributed)
+        # - Diffusers only normalizes image latents from VAE (for img2img/kontext)
+        # - Output MUST be denormalized after denoising before VAE decode
+        #
+        # For img2img with init_latents, we should normalize init_latents on unpacked
+        # shape (B, 128, H/16, W/16) - this is handled by _bn_normalize_unpacked below
 
         # Verify packed dimensions
         assert packed_h * packed_w == x.shape[1]
@@ -366,16 +364,24 @@ class Flux2DenoiseInvocation(BaseInvocation):
         if self.scheduler in FLUX_SCHEDULER_MAP and not is_inpainting:
             # Only use scheduler for txt2img - use manual Euler for inpainting to preserve exact timesteps
             scheduler_class = FLUX_SCHEDULER_MAP[self.scheduler]
-            scheduler = scheduler_class(
-                num_train_timesteps=1000,
-                shift=3.0,
-                use_dynamic_shifting=True,
-                base_shift=0.5,
-                max_shift=1.15,
-                base_image_seq_len=256,
-                max_image_seq_len=4096,
-                time_shift_type="exponential",
-            )
+            # FlowMatchHeunDiscreteScheduler only supports num_train_timesteps and shift parameters
+            # FlowMatchEulerDiscreteScheduler and FlowMatchLCMScheduler support dynamic shifting
+            if self.scheduler == "heun":
+                scheduler = scheduler_class(
+                    num_train_timesteps=1000,
+                    shift=3.0,
+                )
+            else:
+                scheduler = scheduler_class(
+                    num_train_timesteps=1000,
+                    shift=3.0,
+                    use_dynamic_shifting=True,
+                    base_shift=0.5,
+                    max_shift=1.15,
+                    base_image_seq_len=256,
+                    max_image_seq_len=4096,
+                    time_shift_type="exponential",
+                )
 
         # Prepare reference image extension for FLUX.2 Klein built-in editing
         ref_image_extension = None

invokeai/app/invocations/flux2_vae_decode.py

@@ -57,20 +57,6 @@ class Flux2VaeDecodeInvocation(BaseInvocation, WithMetadata, WithBoard):
             # Decode using diffusers API
             decoded = vae.decode(latents, return_dict=False)[0]
 
-        # Debug: Log decoded output statistics
-        print(
-            f"[FLUX.2 VAE] Decoded output: shape={decoded.shape}, "
-            f"min={decoded.min().item():.4f}, max={decoded.max().item():.4f}, "
-            f"mean={decoded.mean().item():.4f}"
-        )
-        # Check per-channel statistics to diagnose color issues
-        for c in range(min(3, decoded.shape[1])):
-            ch = decoded[0, c]
-            print(
-                f"[FLUX.2 VAE] Channel {c}: min={ch.min().item():.4f}, "
-                f"max={ch.max().item():.4f}, mean={ch.mean().item():.4f}"
-            )
-
         # Convert from [-1, 1] to [0, 1] then to [0, 255] PIL image
         img = (decoded / 2 + 0.5).clamp(0, 1)
         img = rearrange(img[0], "c h w -> h w c")

invokeai/app/invocations/flux_denoise.py

@@ -71,7 +71,7 @@ from invokeai.backend.util.devices import TorchDevice
     title="FLUX Denoise",
     tags=["image", "flux"],
     category="image",
-    version="4.5.0",
+    version="4.5.1",
 )
 class FluxDenoiseInvocation(BaseInvocation):
     """Run denoising process with a FLUX transformer model."""
@@ -176,7 +176,10 @@ class FluxDenoiseInvocation(BaseInvocation):
     # DyPE (Dynamic Position Extrapolation) for high-resolution generation
     dype_preset: DyPEPreset = InputField(
         default=DYPE_PRESET_OFF,
-        description="DyPE preset for high-resolution generation. 'auto' enables automatically for resolutions > 1536px. '4k' uses optimized settings for 4K output.",
+        description=(
+            "DyPE preset for high-resolution generation. 'auto' enables automatically for resolutions > 1536px. "
+            "'area' enables automatically based on image area. '4k' uses optimized settings for 4K output."
+        ),
         ui_order=100,
         ui_choice_labels=DYPE_PRESET_LABELS,
     )

invokeai/backend/flux/dype/__init__.py

@@ -10,11 +10,13 @@ from invokeai.backend.flux.dype.base import DyPEConfig
 from invokeai.backend.flux.dype.embed import DyPEEmbedND
 from invokeai.backend.flux.dype.presets import (
     DYPE_PRESET_4K,
+    DYPE_PRESET_AREA,
     DYPE_PRESET_AUTO,
     DYPE_PRESET_LABELS,
     DYPE_PRESET_MANUAL,
     DYPE_PRESET_OFF,
     DyPEPreset,
+    get_dype_config_for_area,
     get_dype_config_for_resolution,
 )
 
@@ -25,7 +27,9 @@ __all__ = [
     "DYPE_PRESET_OFF",
     "DYPE_PRESET_MANUAL",
     "DYPE_PRESET_AUTO",
+    "DYPE_PRESET_AREA",
     "DYPE_PRESET_4K",
     "DYPE_PRESET_LABELS",
+    "get_dype_config_for_area",
     "get_dype_config_for_resolution",
 ]

invokeai/backend/flux/dype/presets.py

@@ -1,17 +1,19 @@
 """DyPE presets and automatic configuration."""
 
+import math
 from dataclasses import dataclass
 from typing import Literal
 
 from invokeai.backend.flux.dype.base import DyPEConfig
 
 # DyPE preset type - using Literal for proper frontend dropdown support
-DyPEPreset = Literal["off", "manual", "auto", "4k"]
+DyPEPreset = Literal["off", "manual", "auto", "area", "4k"]
 
 # Constants for preset values
 DYPE_PRESET_OFF: DyPEPreset = "off"
 DYPE_PRESET_MANUAL: DyPEPreset = "manual"
 DYPE_PRESET_AUTO: DyPEPreset = "auto"
+DYPE_PRESET_AREA: DyPEPreset = "area"
 DYPE_PRESET_4K: DyPEPreset = "4k"
 
 # Human-readable labels for the UI
@@ -19,6 +21,7 @@ DYPE_PRESET_LABELS: dict[str, str] = {
     "off": "Off",
     "manual": "Manual",
     "auto": "Auto (>1536px)",
+    "area": "Area (auto)",
     "4k": "4K Optimized",
 }
 
@@ -88,6 +91,50 @@ def get_dype_config_for_resolution(
     )
 
 
+def get_dype_config_for_area(
+    width: int,
+    height: int,
+    base_resolution: int = 1024,
+) -> DyPEConfig | None:
+    """Automatically determine DyPE config based on target area.
+
+    Uses sqrt(area/base_area) as an effective side-length ratio.
+    DyPE is enabled only when target area exceeds base area.
+
+    Returns:
+        DyPEConfig if DyPE should be enabled, None otherwise
+    """
+    area = width * height
+    base_area = base_resolution**2
+
+    if area <= base_area:
+        return None
+
+    area_ratio = area / base_area
+    effective_side_ratio = math.sqrt(area_ratio)  # 1.0 at base, 2.0 at 2K (if base is 1K)
+
+    # Strength: 0 at base area, 8 at sat_area, clamped thereafter.
+    sat_area = 2027520  # Determined by experimentation where a vertical line appears
+    sat_side_ratio = math.sqrt(sat_area / base_area)
+    dynamic_dype_scale = 8.0 * (effective_side_ratio - 1.0) / (sat_side_ratio - 1.0)
+    dynamic_dype_scale = max(0.0, min(dynamic_dype_scale, 8.0))
+
+    # Continuous exponent schedule:
+    # r=1 -> 0.5, r=2 -> 1.0, r=4 -> 2.0 (exact), smoothly varying in between.
+    x = math.log2(effective_side_ratio)
+    dype_exponent = 0.25 * (x**2) + 0.25 * x + 0.5
+    dype_exponent = max(0.5, min(dype_exponent, 2.0))
+
+    return DyPEConfig(
+        enable_dype=True,
+        base_resolution=base_resolution,
+        method="vision_yarn",
+        dype_scale=dynamic_dype_scale,
+        dype_exponent=dype_exponent,
+        dype_start_sigma=1.0,
+    )
+
+
 def get_dype_config_from_preset(
     preset: DyPEPreset,
     width: int,
@@ -133,6 +180,14 @@ def get_dype_config_from_preset(
             activation_threshold=1536,
         )
 
+    if preset == DYPE_PRESET_AREA:
+        # Area-based preset - custom values are ignored
+        return get_dype_config_for_area(
+            width=width,
+            height=height,
+            base_resolution=1024,
+        )
+
     # Use preset configuration (4K etc.) - custom values are ignored
     preset_config = DYPE_PRESETS.get(preset)
     if preset_config is None:

invokeai/backend/flux2/denoise.py

@@ -4,6 +4,7 @@ This module provides the denoising function for FLUX.2 Klein models,
 which use Qwen3 as the text encoder instead of CLIP+T5.
 """
 
+import inspect
 import math
 from typing import Any, Callable
 
@@ -87,11 +88,18 @@ def denoise(
         # 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:
+        # Check if scheduler supports sigmas parameter using inspect.signature
+        # FlowMatchHeunDiscreteScheduler and FlowMatchLCMScheduler don't support sigmas
+        set_timesteps_sig = inspect.signature(scheduler.set_timesteps)
+        supports_sigmas = "sigmas" in set_timesteps_sig.parameters
+        if supports_sigmas and mu is not None:
+            # Pass mu if provided - it will only be used if scheduler has use_dynamic_shifting=True
             scheduler.set_timesteps(sigmas=sigmas.tolist(), mu=mu, device=img.device)
-        else:
+        elif supports_sigmas:
             scheduler.set_timesteps(sigmas=sigmas.tolist(), device=img.device)
+        else:
+            # Scheduler doesn't support sigmas (e.g., Heun, LCM) - use num_inference_steps
+            scheduler.set_timesteps(num_inference_steps=len(sigmas), device=img.device)
         num_scheduler_steps = len(scheduler.timesteps)
         is_heun = hasattr(scheduler, "state_in_first_order")
         user_step = 0

invokeai/backend/flux2/sampling_utils.py

@@ -108,33 +108,27 @@ def unpack_flux2(x: torch.Tensor, height: int, width: int) -> torch.Tensor:
 
 
 def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float:
-    """Compute empirical mu for FLUX.2 schedule shifting.
+    """Compute mu for FLUX.2 schedule shifting.
 
-    This matches the diffusers Flux2Pipeline implementation.
-    The mu value controls how much the schedule is shifted towards higher timesteps.
+    Uses a fixed mu value of 2.02, matching ComfyUI's proven FLUX.2 configuration.
+
+    The previous implementation (from diffusers' FLUX.1 pipeline) computed mu as a
+    linear function of image_seq_len, which produced excessively high values at
+    high resolutions (e.g., mu=3.23 at 2048x2048). This over-shifted the sigma
+    schedule, compressing almost all values above 0.9 and forcing the model to
+    denoise everything in the final 1-2 steps, causing severe grid/diamond artifacts.
+
+    ComfyUI uses a fixed shift=2.02 for FLUX.2 Klein at all resolutions and produces
+    artifact-free images even at 2048x2048.
 
     Args:
-        image_seq_len: Number of image tokens (packed_h * packed_w).
-        num_steps: Number of denoising steps.
+        image_seq_len: Number of image tokens (packed_h * packed_w). Currently unused.
+        num_steps: Number of denoising steps. Currently unused.
 
     Returns:
-        The empirical mu value.
+        The mu value (fixed at 2.02).
     """
-    a1, b1 = 8.73809524e-05, 1.89833333
-    a2, b2 = 0.00016927, 0.45666666
-
-    if image_seq_len > 4300:
-        mu = a2 * image_seq_len + b2
-        return float(mu)
-
-    m_200 = a2 * image_seq_len + b2
-    m_10 = a1 * image_seq_len + b1
-
-    a = (m_200 - m_10) / 190.0
-    b = m_200 - 200.0 * a
-    mu = a * num_steps + b
-
-    return float(mu)
+    return 2.02
 
 
 def get_schedule_flux2(
@@ -169,11 +163,14 @@ def get_schedule_flux2(
 
 
 def generate_img_ids_flux2(h: int, w: int, batch_size: int, device: torch.device) -> torch.Tensor:
-    """Generate tensor of image position ids for FLUX.2.
+    """Generate tensor of image position ids for FLUX.2 with RoPE scaling.
 
     FLUX.2 uses 4D position coordinates (T, H, W, L) for its rotary position embeddings.
     This is different from FLUX.1 which uses 3D coordinates.
 
+    RoPE Scaling: For resolutions >1536x1536, position IDs are scaled down using
+    Position Interpolation to prevent RoPE degradation and diamond/grid artifacts.
+
     IMPORTANT: Position IDs must use int64 (long) dtype like diffusers, not bfloat16.
     Using floating point dtype for position IDs can cause NaN in rotary embeddings.