invoke-ai--invokeai/invokeai/app/invocations/flux2_denoise.py

"""Flux2 Klein Denoise Invocation.

Run denoising process with a FLUX.2 Klein transformer model.
Uses Qwen3 conditioning instead of CLIP+T5.
"""

from contextlib import ExitStack
from typing import Callable, Iterator, Optional, Tuple

import torch
import torchvision.transforms as tv_transforms
from torchvision.transforms.functional import resize as tv_resize

from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
from invokeai.app.invocations.fields import (
    DenoiseMaskField,
    FieldDescriptions,
    FluxConditioningField,
    FluxKontextConditioningField,
    Input,
    InputField,
    LatentsField,
)
from invokeai.app.invocations.latent_noise import validate_noise_tensor_shape
from invokeai.app.invocations.model import TransformerField, VAEField
from invokeai.app.invocations.primitives import LatentsOutput
from invokeai.app.services.shared.invocation_context import InvocationContext
from invokeai.backend.flux.sampling_utils import clip_timestep_schedule_fractional
from invokeai.backend.flux.schedulers import FLUX_SCHEDULER_LABELS, FLUX_SCHEDULER_MAP, FLUX_SCHEDULER_NAME_VALUES
from invokeai.backend.flux2.denoise import denoise
from invokeai.backend.flux2.ref_image_extension import Flux2RefImageExtension
from invokeai.backend.flux2.sampling_utils import (
    compute_empirical_mu,
    generate_img_ids_flux2,
    get_noise_flux2,
    get_schedule_flux2,
    pack_flux2,
    unpack_flux2,
)
from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType
from invokeai.backend.patches.layer_patcher import LayerPatcher
from invokeai.backend.patches.lora_conversions.flux_bfl_peft_lora_conversion_utils import (
    convert_bfl_lora_patch_to_diffusers,
)
from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension
from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import FLUXConditioningInfo
from invokeai.backend.util.devices import TorchDevice


@invocation(
    "flux2_denoise",
    title="FLUX2 Denoise",
    tags=["image", "flux", "flux2", "klein", "denoise"],
    category="latents",
    version="1.5.0",
    classification=Classification.Prototype,
)
class Flux2DenoiseInvocation(BaseInvocation):
    """Run denoising process with a FLUX.2 Klein transformer model.

    This node is designed for FLUX.2 Klein models which use Qwen3 as the text encoder.
    It does not support ControlNet, IP-Adapters, or regional prompting.
    """

    latents: Optional[LatentsField] = InputField(
        default=None,
        description=FieldDescriptions.latents,
        input=Input.Connection,
    )
    noise: Optional[LatentsField] = InputField(
        default=None,
        description=FieldDescriptions.noise,
        input=Input.Connection,
    )
    denoise_mask: Optional[DenoiseMaskField] = InputField(
        default=None,
        description=FieldDescriptions.denoise_mask,
        input=Input.Connection,
    )
    denoising_start: float = InputField(
        default=0.0,
        ge=0,
        le=1,
        description=FieldDescriptions.denoising_start,
    )
    denoising_end: float = InputField(
        default=1.0,
        ge=0,
        le=1,
        description=FieldDescriptions.denoising_end,
    )
    add_noise: bool = InputField(default=True, description="Add noise based on denoising start.")
    transformer: TransformerField = InputField(
        description=FieldDescriptions.flux_model,
        input=Input.Connection,
        title="Transformer",
    )
    positive_text_conditioning: FluxConditioningField = InputField(
        description=FieldDescriptions.positive_cond,
        input=Input.Connection,
    )
    negative_text_conditioning: Optional[FluxConditioningField] = InputField(
        default=None,
        description="Negative conditioning tensor. Can be None if cfg_scale is 1.0.",
        input=Input.Connection,
    )
    guidance: float = InputField(
        default=4.0,
        ge=0,
        le=20,
        description="Guidance strength for distilled guidance-embedding models. "
        "Inert for all current FLUX.2 Klein variants (their guidance_embeds weights are absent/zero); "
        "kept for node-graph compatibility and future guidance-embedded models.",
    )
    cfg_scale: float = InputField(
        default=1.0,
        description=FieldDescriptions.cfg_scale,
        title="CFG Scale",
    )
    width: int = InputField(default=1024, multiple_of=16, description="Width of the generated image.")
    height: int = InputField(default=1024, multiple_of=16, description="Height of the generated image.")
    num_steps: int = InputField(
        default=4,
        description="Number of diffusion steps. Use 4 for distilled models, 28+ for base models.",
    )
    scheduler: FLUX_SCHEDULER_NAME_VALUES = InputField(
        default="euler",
        description="Scheduler (sampler) for the denoising process. 'euler' is fast and standard. "
        "'heun' is 2nd-order (better quality, 2x slower). 'lcm' is optimized for few steps.",
        ui_choice_labels=FLUX_SCHEDULER_LABELS,
    )
    seed: int = InputField(default=0, description="Randomness seed for reproducibility.")
    vae: VAEField = InputField(
        description="FLUX.2 VAE model (required for BN statistics).",
        input=Input.Connection,
    )
    kontext_conditioning: FluxKontextConditioningField | list[FluxKontextConditioningField] | None = InputField(
        default=None,
        description="FLUX Kontext conditioning (reference images for multi-reference image editing).",
        input=Input.Connection,
        title="Reference Images",
    )

    def _get_bn_stats(self, context: InvocationContext) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
        """Extract BN statistics from the FLUX.2 VAE.

        The FLUX.2 VAE uses batch normalization on the patchified 128-channel representation.
        IMPORTANT: BFL FLUX.2 VAE uses affine=False, so there are NO learnable weight/bias.

        BN formula (affine=False): y = (x - mean) / std
        Inverse: x = y * std + mean

        Returns:
            Tuple of (bn_mean, bn_std) tensors of shape (128,), or None if BN layer not found.
        """
        with context.models.load(self.vae.vae).model_on_device() as (_, vae):
            # Ensure VAE is in eval mode to prevent BN stats from being updated
            vae.eval()

            # Try to find the BN layer - it may be at different locations depending on model format
            bn_layer = None
            if hasattr(vae, "bn"):
                bn_layer = vae.bn
            elif hasattr(vae, "batch_norm"):
                bn_layer = vae.batch_norm
            elif hasattr(vae, "encoder") and hasattr(vae.encoder, "bn"):
                bn_layer = vae.encoder.bn

            if bn_layer is None:
                return None

            # Verify running statistics are initialized
            if bn_layer.running_mean is None or bn_layer.running_var is None:
                return None

            # Get BN running statistics from VAE
            bn_mean = bn_layer.running_mean.clone()  # Shape: (128,)
            bn_var = bn_layer.running_var.clone()  # Shape: (128,)
            bn_eps = bn_layer.eps if hasattr(bn_layer, "eps") else 1e-4  # BFL uses 1e-4
            bn_std = torch.sqrt(bn_var + bn_eps)

        return bn_mean, bn_std

    def _bn_normalize(
        self,
        x: torch.Tensor,
        bn_mean: torch.Tensor,
        bn_std: torch.Tensor,
    ) -> torch.Tensor:
        """Apply BN normalization to packed latents.

        BN formula (affine=False): y = (x - mean) / std

        Args:
            x: Packed latents of shape (B, seq, 128).
            bn_mean: BN running mean of shape (128,).
            bn_std: BN running std of shape (128,).

        Returns:
            Normalized latents of same shape.
        """
        # x: (B, seq, 128), params: (128,) -> broadcast over batch and sequence dims
        bn_mean = bn_mean.to(x.device, x.dtype)
        bn_std = bn_std.to(x.device, x.dtype)
        return (x - bn_mean) / bn_std

    def _bn_denormalize(
        self,
        x: torch.Tensor,
        bn_mean: torch.Tensor,
        bn_std: torch.Tensor,
    ) -> torch.Tensor:
        """Apply BN denormalization to packed latents (inverse of normalization).

        Inverse BN (affine=False): x = y * std + mean

        Args:
            x: Packed latents of shape (B, seq, 128).
            bn_mean: BN running mean of shape (128,).
            bn_std: BN running std of shape (128,).

        Returns:
            Denormalized latents of same shape.
        """
        # x: (B, seq, 128), params: (128,) -> broadcast over batch and sequence dims
        bn_mean = bn_mean.to(x.device, x.dtype)
        bn_std = bn_std.to(x.device, x.dtype)
        return x * bn_std + bn_mean

    @torch.no_grad()
    def invoke(self, context: InvocationContext) -> LatentsOutput:
        latents = self._run_diffusion(context)
        latents = latents.detach().to("cpu")

        name = context.tensors.save(tensor=latents)
        return LatentsOutput.build(latents_name=name, latents=latents, seed=None)

    def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
        inference_dtype = torch.bfloat16
        device = TorchDevice.choose_torch_device()

        # Get BN statistics from VAE for latent denormalization (optional)
        # BFL FLUX.2 VAE uses affine=False, so only mean/std are needed
        # Some VAE formats (e.g. diffusers) may not expose BN stats directly
        bn_stats = self._get_bn_stats(context)
        bn_mean, bn_std = bn_stats if bn_stats is not None else (None, None)

        # Load the input latents, if provided
        init_latents = context.tensors.load(self.latents.latents_name) if self.latents else None
        if init_latents is not None:
            init_latents = init_latents.to(device=device, dtype=inference_dtype)

        # Prepare input noise (FLUX.2 uses 32 channels).
        # If noise will never be consumed, avoid validating/loading it.
        should_ignore_noise = init_latents is not None and not self.add_noise and self.denoise_mask is None
        noise: Optional[torch.Tensor]
        if should_ignore_noise:
            noise = None
            b, _c, latent_h, latent_w = init_latents.shape
        else:
            noise = self._prepare_noise_tensor(context, inference_dtype, device)
            b, _c, latent_h, latent_w = noise.shape
        packed_h = latent_h // 2
        packed_w = latent_w // 2

        # Load the conditioning data
        pos_cond_data = context.conditioning.load(self.positive_text_conditioning.conditioning_name)
        assert len(pos_cond_data.conditionings) == 1
        pos_flux_conditioning = pos_cond_data.conditionings[0]
        assert isinstance(pos_flux_conditioning, FLUXConditioningInfo)
        pos_flux_conditioning = pos_flux_conditioning.to(dtype=inference_dtype, device=device)

        # Qwen3 stacked embeddings (stored in t5_embeds field for compatibility)
        txt = pos_flux_conditioning.t5_embeds

        # Generate text position IDs (4D format for FLUX.2: T, H, W, L)
        # FLUX.2 uses 4D position coordinates for its rotary position embeddings
        # IMPORTANT: Position IDs must be int64 (long) dtype
        # Diffusers uses: T=0, H=0, W=0, L=0..seq_len-1
        seq_len = txt.shape[1]
        txt_ids = torch.zeros(1, seq_len, 4, device=device, dtype=torch.long)
        txt_ids[..., 3] = torch.arange(seq_len, device=device, dtype=torch.long)  # L coordinate varies

        # Load negative conditioning if provided
        neg_txt = None
        neg_txt_ids = None
        if self.negative_text_conditioning is not None:
            neg_cond_data = context.conditioning.load(self.negative_text_conditioning.conditioning_name)
            assert len(neg_cond_data.conditionings) == 1
            neg_flux_conditioning = neg_cond_data.conditionings[0]
            assert isinstance(neg_flux_conditioning, FLUXConditioningInfo)
            neg_flux_conditioning = neg_flux_conditioning.to(dtype=inference_dtype, device=device)
            neg_txt = neg_flux_conditioning.t5_embeds
            # For text tokens: T=0, H=0, W=0, L=0..seq_len-1 (only L varies per token)
            neg_seq_len = neg_txt.shape[1]
            neg_txt_ids = torch.zeros(1, neg_seq_len, 4, device=device, dtype=torch.long)
            neg_txt_ids[..., 3] = torch.arange(neg_seq_len, device=device, dtype=torch.long)

        # Validate transformer config
        transformer_config = context.models.get_config(self.transformer.transformer)
        assert transformer_config.base == BaseModelType.Flux2 and transformer_config.type == ModelType.Main

        # Calculate the timestep schedule using FLUX.2 specific schedule
        # This matches diffusers' Flux2Pipeline implementation
        # Note: Schedule shifting is handled by the scheduler via mu parameter
        image_seq_len = packed_h * packed_w
        timesteps = get_schedule_flux2(
            num_steps=self.num_steps,
            image_seq_len=image_seq_len,
        )
        # Compute mu for dynamic schedule shifting (used by FlowMatchEulerDiscreteScheduler)
        mu = compute_empirical_mu(image_seq_len=image_seq_len, num_steps=self.num_steps)

        # Clip the timesteps schedule based on denoising_start and denoising_end
        timesteps = clip_timestep_schedule_fractional(timesteps, self.denoising_start, self.denoising_end)

        # Prepare input latent image
        if init_latents is not None:
            if self.add_noise:
                assert noise is not None
                # Noise the init latents using the first timestep from the clipped
                # InvokeAI schedule.
                #
                # Known limitation: if a scheduler later uses a different first
                # effective timestep/sigma than this precomputed schedule, the
                # img2img preblend below may not match that scheduler exactly.
                # This is an existing pipeline limitation and applies to both
                # seed-generated noise and externally supplied noise.
                t_0 = timesteps[0]
                x = t_0 * noise + (1.0 - t_0) * init_latents
            else:
                x = init_latents
        else:
            if self.denoising_start > 1e-5:
                raise ValueError("denoising_start should be 0 when initial latents are not provided.")
            assert noise is not None
            x = noise

        # If len(timesteps) == 1, then short-circuit
        if len(timesteps) <= 1:
            return x

        # Generate image position IDs (FLUX.2 uses 4D coordinates)
        # Position IDs use int64 dtype like diffusers
        img_ids = generate_img_ids_flux2(h=latent_h, w=latent_w, batch_size=b, device=device)

        # Prepare inpaint mask
        inpaint_mask = self._prep_inpaint_mask(context, x)

        # Pack all latent tensors
        init_latents_packed = pack_flux2(init_latents) if init_latents is not None else None
        inpaint_mask_packed = pack_flux2(inpaint_mask) if inpaint_mask is not None else None
        noise_packed = pack_flux2(noise) if noise is not None else None
        x = pack_flux2(x)

        # BN normalization for img2img/inpainting:
        # - The init_latents from VAE encode are NOT BN-normalized
        # - The transformer operates in BN-normalized space
        # - We must normalize x, init_latents, AND noise for InpaintExtension
        # - Output MUST be denormalized after denoising before VAE decode
        #
        # This ensures that:
        # 1. x starts in the correct normalized space for the transformer
        # 2. When InpaintExtension merges intermediate_latents with noised_init_latents,
        #    both are in the same scale/space (noise and init_latents must be in same space
        #    for the linear interpolation: noised = noise * t + init * (1-t))
        if bn_mean is not None and bn_std is not None:
            if init_latents_packed is not None:
                init_latents_packed = self._bn_normalize(init_latents_packed, bn_mean, bn_std)
                # Also normalize noise for InpaintExtension - it's used to compute
                # noised_init_latents = noise * t + init_latents * (1-t)
                # Both operands must be in the same normalized space
                if noise_packed is not None:
                    noise_packed = self._bn_normalize(noise_packed, bn_mean, bn_std)
            # For img2img/inpainting, x is computed from init_latents and must also be normalized
            # For txt2img, x is pure noise (already N(0,1)) - normalizing it would be incorrect
            # We detect img2img by checking if init_latents was provided
            if init_latents is not None:
                x = self._bn_normalize(x, bn_mean, bn_std)

        # Verify packed dimensions
        assert packed_h * packed_w == x.shape[1]

        # Prepare inpaint extension
        inpaint_extension: Optional[RectifiedFlowInpaintExtension] = None
        if inpaint_mask_packed is not None:
            assert init_latents_packed is not None
            assert noise_packed is not None
            inpaint_extension = RectifiedFlowInpaintExtension(
                init_latents=init_latents_packed,
                inpaint_mask=inpaint_mask_packed,
                noise=noise_packed,
            )

        # Prepare CFG scale list
        num_steps = len(timesteps) - 1
        cfg_scale_list = [self.cfg_scale] * num_steps

        # Check if we're doing inpainting (have a mask or a clipped schedule)
        is_inpainting = self.denoise_mask is not None or self.denoising_start > 1e-5

        # Create scheduler with FLUX.2 Klein configuration
        # For inpainting/img2img, use manual Euler stepping to preserve the exact
        # clipped timestep schedule used for the initial latent/noise preblend.
        # For txt2img, use the scheduler with dynamic shifting for optimal results.
        #
        # This split is intentional. Reusing a scheduler for img2img here can
        # change the first effective timestep/sigma and break parity with the
        # preblend computed above.
        scheduler = None
        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]
            # 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
        if self.kontext_conditioning:
            ref_image_extension = Flux2RefImageExtension(
                context=context,
                ref_image_conditioning=self.kontext_conditioning
                if isinstance(self.kontext_conditioning, list)
                else [self.kontext_conditioning],
                vae_field=self.vae,
                device=device,
                dtype=inference_dtype,
                bn_mean=bn_mean,
                bn_std=bn_std,
            )

        with ExitStack() as exit_stack:
            # Load the transformer model
            (cached_weights, transformer) = exit_stack.enter_context(
                context.models.load(self.transformer.transformer).model_on_device()
            )
            config = transformer_config

            # Determine if the model is quantized
            if config.format in [ModelFormat.Diffusers]:
                model_is_quantized = False
            elif config.format in [
                ModelFormat.BnbQuantizedLlmInt8b,
                ModelFormat.BnbQuantizednf4b,
                ModelFormat.GGUFQuantized,
            ]:
                model_is_quantized = True
            else:
                model_is_quantized = False

            # Apply LoRA models to the transformer
            exit_stack.enter_context(
                LayerPatcher.apply_smart_model_patches(
                    model=transformer,
                    patches=self._lora_iterator(context),
                    prefix=FLUX_LORA_TRANSFORMER_PREFIX,
                    dtype=inference_dtype,
                    cached_weights=cached_weights,
                    force_sidecar_patching=model_is_quantized,
                )
            )

            # Prepare reference image conditioning if provided
            img_cond_seq = None
            img_cond_seq_ids = None
            if ref_image_extension is not None:
                # Ensure batch sizes match
                ref_image_extension.ensure_batch_size(x.shape[0])
                img_cond_seq, img_cond_seq_ids = (
                    ref_image_extension.ref_image_latents,
                    ref_image_extension.ref_image_ids,
                )

            x = denoise(
                model=transformer,
                img=x,
                img_ids=img_ids,
                txt=txt,
                txt_ids=txt_ids,
                timesteps=timesteps,
                step_callback=self._build_step_callback(context),
                guidance=self.guidance,
                cfg_scale=cfg_scale_list,
                neg_txt=neg_txt,
                neg_txt_ids=neg_txt_ids,
                scheduler=scheduler,
                mu=mu,
                inpaint_extension=inpaint_extension,
                img_cond_seq=img_cond_seq,
                img_cond_seq_ids=img_cond_seq_ids,
            )

        # Apply BN denormalization if BN stats are available
        # The diffusers Flux2KleinPipeline applies: latents = latents * bn_std + bn_mean
        # This transforms latents from normalized space to VAE's expected input space
        if bn_mean is not None and bn_std is not None:
            x = self._bn_denormalize(x, bn_mean, bn_std)

        x = unpack_flux2(x.float(), self.height, self.width)
        return x

    def _prepare_noise_tensor(
        self, context: InvocationContext, inference_dtype: torch.dtype, device: torch.device
    ) -> torch.Tensor:
        if self.noise is not None:
            noise = context.tensors.load(self.noise.latents_name).to(device=device, dtype=inference_dtype)
            validate_noise_tensor_shape(noise, "FLUX.2", self.width, self.height)
            return noise

        return get_noise_flux2(
            num_samples=1,
            height=self.height,
            width=self.width,
            device=device,
            dtype=inference_dtype,
            seed=self.seed,
        )

    def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor) -> Optional[torch.Tensor]:
        """Prepare the inpaint mask."""
        if self.denoise_mask is None:
            return None

        mask = context.tensors.load(self.denoise_mask.mask_name)
        mask = 1.0 - mask

        _, _, latent_height, latent_width = latents.shape
        mask = tv_resize(
            img=mask,
            size=[latent_height, latent_width],
            interpolation=tv_transforms.InterpolationMode.BILINEAR,
            antialias=False,
        )

        mask = mask.to(device=latents.device, dtype=latents.dtype)
        return mask.expand_as(latents)

    def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[ModelPatchRaw, float]]:
        """Iterate over LoRA models to apply.

        Converts BFL-format LoRA keys to diffusers format if needed, since FLUX.2 Klein
        uses Flux2Transformer2DModel (diffusers naming) but LoRAs may have been loaded
        with BFL naming (e.g. when a Klein 4B LoRA is misidentified as FLUX.1).
        """
        for lora in self.transformer.loras:
            lora_info = context.models.load(lora.lora)
            assert isinstance(lora_info.model, ModelPatchRaw)
            converted = convert_bfl_lora_patch_to_diffusers(lora_info.model)
            yield (converted, lora.weight)
            del lora_info

    def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]:
        """Build a callback for step progress updates."""

        def step_callback(state: PipelineIntermediateState) -> None:
            latents = state.latents.float()
            state.latents = unpack_flux2(latents, self.height, self.width).squeeze()
            context.util.flux2_step_callback(state)

        return step_callback