src/diffsynth_fov/pipeline.py

"""
FLUX2 Foveated Image Generation Pipeline for DiffSynth.

This pipeline integrates foveation logic from pipeline_flux2_klein_foveation.py
into the DiffSynth FLUX2 pipeline architecture.

Key features:
- Mixed-resolution latent processing for efficient foveated rendering
- X0 (clean image) prediction for better upsampling quality
- Phase-aligned RoPE for consistent position encoding across resolutions
- CRPA attention for efficient mixed-resolution attention
"""

import os
import torch
import matplotlib
matplotlib.use("Agg")
import math
import torchvision
import numpy as np
import torch.nn.functional as F
from PIL import Image
from typing import Union, List, Optional, Tuple
from tqdm import tqdm
from einops import rearrange
from torchvision.utils import save_image

from diffsynth.core.device.npu_compatible_device import get_device_type
from diffsynth.diffusion import FlowMatchScheduler
from diffsynth.core import ModelConfig, gradient_checkpoint_forward
from diffsynth.diffusion.base_pipeline import BasePipeline, PipelineUnit, ControlNetInput

from transformers import AutoProcessor, AutoTokenizer
from diffsynth.models.flux2_text_encoder import Flux2TextEncoder
from .dit import Flux2DiTFoveated
from diffsynth.models.flux2_vae import Flux2VAE
from diffsynth.models.z_image_text_encoder import ZImageTextEncoder
from scipy.ndimage import zoom, gaussian_filter
import time

# Optional saliency-detection dependencies: only required when
# foveated_training_mode == "saliency". Import lazily so a default
# install does not need them.
try:
    from scipy.datasets import face as scipy_face
except Exception:
    scipy_face = None
try:
    import deepgaze_pytorch
except Exception:
    deepgaze_pytorch = None


def _gaussian_blur_mask_2d(
    mask: torch.Tensor,
    sigma_pixels: float,
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    """Apply separable Gaussian blur to a [1, 1, H, W] mask. sigma_pixels controls falloff width."""
    # Kernel size odd, ~3 sigma on each side
    k = max(3, int(math.ceil(3 * sigma_pixels)) * 2 + 1)
    if k % 2 == 0:
        k += 1
    x = torch.arange(k, device=device, dtype=dtype) - (k - 1) / 2.0
    g = torch.exp(-(x ** 2) / (2 * sigma_pixels ** 2 + 1e-6))
    g = g / g.sum()
    # [1, 1, k, 1] and [1, 1, 1, k] for separable conv
    kernel_h = g.view(1, 1, k, 1)
    kernel_w = g.view(1, 1, 1, k)
    pad = k // 2
    out = F.conv2d(mask, kernel_h, padding=(pad, 0))
    out = F.conv2d(out, kernel_w, padding=(0, pad))
    return out


class Flux2FoveatedImagePipeline(BasePipeline):
    """
    FLUX2 Foveated Image Generation Pipeline.
    
    This pipeline extends the standard FLUX2 pipeline with support for:
    - Foveated rendering with mixed-resolution latents
    - X0 prediction upsampling for better quality
    - CRPA attention for efficient attention computation
    """

    def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16):
        super().__init__(
            device=device, torch_dtype=torch_dtype,
            height_division_factor=16, width_division_factor=16,
        )
        self.scheduler = FlowMatchScheduler("FLUX.2")
        self.text_encoder: Flux2TextEncoder = None
        self.text_encoder_qwen3: ZImageTextEncoder = None
        self.dit: Flux2DiTFoveated = None
        self.vae: Flux2VAE = None
        self.tokenizer: AutoProcessor = None
        self.in_iteration_models = ("dit",)
        self.units = [
            Flux2FoveatedUnit_ShapeChecker(),
            Flux2FoveatedUnit_PromptEmbedder(),
            Flux2FoveatedUnit_Qwen3PromptEmbedder(),
            Flux2FoveatedUnit_NoiseInitializer(),
            Flux2FoveatedUnit_InputImageEmbedder(),
            Flux2FoveatedUnit_ImageIDs(),
        ]
        
        # fixed latent for tracking validation
        self.fixed_clean_latent = None
        self.detach_z_mix = False

        self.model_fn = model_fn_flux2_foveated
        
    
    
    @staticmethod
    def from_pretrained(
        torch_dtype: torch.dtype = torch.bfloat16,
        device: Union[str, torch.device] = get_device_type(),
        model_configs: list[ModelConfig] = [],
        tokenizer_config: ModelConfig = ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="tokenizer/"),
        vram_limit: float = None,
    ):
        """
        Load a FLUX2 Foveated Image Pipeline from pretrained weights.
        
        The pipeline will first try to fetch a registered "flux2_dit_foveated" model.
        If not found, it will fetch a standard "flux2_dit" model and convert it to
        the foveated version (same weights, different attention processing).
        """
        # Initialize pipeline
        pipe = Flux2FoveatedImagePipeline(device=device, torch_dtype=torch_dtype)
        model_pool = pipe.download_and_load_models(model_configs, vram_limit)
        
        # Fetch models
        pipe.text_encoder = model_pool.fetch_model("flux2_text_encoder")
        pipe.text_encoder_qwen3 = model_pool.fetch_model("z_image_text_encoder")
        
        # Try to fetch foveated model first, fall back to base model
        pipe.dit = model_pool.fetch_model("flux2_dit_foveated")
        if pipe.dit is None:
            # Load base FLUX2 model and convert to foveated version
            base_dit = model_pool.fetch_model("flux2_dit")
            if base_dit is not None:
                pipe.dit = Flux2FoveatedImagePipeline._convert_to_foveated_dit(base_dit)
                print("Converted base FLUX2 DiT to foveated version.")
        
        pipe.vae = model_pool.fetch_model("flux2_vae")
        if tokenizer_config is not None:
            tokenizer_config.download_if_necessary()
            pipe.tokenizer = AutoTokenizer.from_pretrained(tokenizer_config.path)
        
        # VRAM Management
        pipe.vram_management_enabled = pipe.check_vram_management_state()
        return pipe

    @staticmethod
    def _convert_to_foveated_dit(base_dit):
        """
        Convert a base FLUX2 DiT to a foveated version.
        
        The foveated model has the exact same architecture as the base model,
        just with CRPA attention processors. We infer config from state dict.
        
        Args:
            base_dit: A loaded Flux2DiT model instance.
            
        Returns:
            A Flux2DiTFoveated model with the same weights.
        """
        # Get the state dict from the base model
        state_dict = base_dit.state_dict()
        
        # Infer configuration from state dict shapes
        # inner_dim from x_embedder.weight: [inner_dim, in_channels]
        inner_dim = state_dict['x_embedder.weight'].shape[0]
        in_channels = state_dict['x_embedder.weight'].shape[1]
        
        # joint_attention_dim from context_embedder.weight: [inner_dim, joint_attention_dim]
        joint_attention_dim = state_dict['context_embedder.weight'].shape[1]
        
        # Count transformer blocks
        num_layers = 0
        num_single_layers = 0
        for key in state_dict.keys():
            if key.startswith('transformer_blocks.') and key.endswith('.attn.to_q.weight'):
                idx = int(key.split('.')[1])
                num_layers = max(num_layers, idx + 1)
            if key.startswith('single_transformer_blocks.') and key.endswith('.attn.to_qkv_mlp_proj.weight'):
                idx = int(key.split('.')[1])
                num_single_layers = max(num_single_layers, idx + 1)
        
        # Infer attention config from to_q weight: [inner_dim, inner_dim]
        # inner_dim = num_heads * head_dim
        # Default head_dim is 128, so num_heads = inner_dim / 128
        attention_head_dim = 128
        num_attention_heads = inner_dim // attention_head_dim
        
        # Check for guidance embeddings
        guidance_embeds = 'time_guidance_embed.guidance_embedder.linear_1.weight' in state_dict
        
        print(f"Inferred config: inner_dim={inner_dim}, in_channels={in_channels}, "
              f"joint_attention_dim={joint_attention_dim}, num_layers={num_layers}, "
              f"num_single_layers={num_single_layers}, num_attention_heads={num_attention_heads}, "
              f"attention_head_dim={attention_head_dim}, guidance_embeds={guidance_embeds}")
        
        # Create foveated model with inferred parameters
        foveated_dit = Flux2DiTFoveated(
            in_channels=in_channels,
            num_layers=num_layers,
            num_single_layers=num_single_layers,
            attention_head_dim=attention_head_dim,
            num_attention_heads=num_attention_heads,
            joint_attention_dim=joint_attention_dim,
            guidance_embeds=guidance_embeds,
        )
        
        # Load state dict directly - architectures are identical
        foveated_dit.load_state_dict(state_dict)
        
        # Move to same device and dtype as base model
        device = next(base_dit.parameters()).device
        dtype = next(base_dit.parameters()).dtype
        foveated_dit = foveated_dit.to(device=device, dtype=dtype)
        
        return foveated_dit

    # =====================================================================
    # Patchify / Unpatchify Operations (FLUX2 2x2 patches)
    # =====================================================================

    @staticmethod
    def _patchify_latents(latents):
        """Convert [B, C, H, W] to 2x2 patched format [B, C*4, H/2, W/2]."""
        batch_size, num_channels, height, width = latents.shape
        latents = latents.view(batch_size, num_channels, height // 2, 2, width // 2, 2)
        latents = latents.permute(0, 1, 3, 5, 2, 4)
        latents = latents.reshape(batch_size, num_channels * 4, height // 2, width // 2)
        return latents

    @staticmethod
    def _unpatchify_latents(latents):
        """Convert [B, C*4, H/2, W/2] back to [B, C, H, W]."""
        batch_size, num_channels, height, width = latents.shape
        latents = latents.reshape(batch_size, num_channels // 4, 2, 2, height, width)
        latents = latents.permute(0, 1, 4, 2, 5, 3)
        latents = latents.reshape(batch_size, num_channels // 4, height * 2, width * 2)
        return latents

    # =====================================================================
    # Foveation-Specific Latent Operations
    # =====================================================================

    @staticmethod
    def _prepare_latent_image_ids_foveation(
        batch_size: int, 
        height: int, 
        width: int, 
        device: torch.device, 
        dtype: torch.dtype, 
        foveation_mask: Optional[torch.Tensor] = None,
        lr_factor: int = 2,
    ):
        """
        Prepare 4D position IDs for image latents with foveation support.
        
        For FLUX2, position IDs have format (T, H, W, L) where:
        - T: Time coordinate (always 0 for image latents)
        - H: Height coordinate
        - W: Width coordinate  
        - L: Layer coordinate (always 0 for image latents)
        
        With foveation:
        - High-res blocks keep all 4 tokens
        - Low-res blocks keep only top-left token
        
        Returns:
            latent_image_ids: [Seq_Len, 4] position IDs
            resolution_mask: [Seq_Len] mask (1=HR, 0=LR)
            resolution_mask_top_left: [Seq_Len] mask for key subsampling
        """
        # Create base coordinate grid (all must have same dtype for cartesian_prod)
        t = torch.zeros(1, dtype=dtype)
        h = torch.arange(height, dtype=dtype)
        w = torch.arange(width, dtype=dtype)
        l = torch.zeros(1, dtype=dtype)
        
        latent_image_ids = torch.cartesian_prod(t, h, w, l).to(device=device)
        
        resolution_mask_grid = torch.ones(height, width, device=device, dtype=dtype)

        if foveation_mask is not None:
            if height % lr_factor != 0 or width % lr_factor != 0:
                raise ValueError(f"Height and width must be divisible by lr_factor ({lr_factor}) for mixed-resolution.")

            n_per_block = lr_factor * lr_factor
            h_d, w_d = height // lr_factor, width // lr_factor

            grid_2d = latent_image_ids.view(height, width, 4)
            grid_blocks = grid_2d.view(h_d, lr_factor, w_d, lr_factor, 4).permute(0, 2, 1, 3, 4).reshape(h_d, w_d, n_per_block, 4)

            mask_blocks = foveation_mask.view(h_d, lr_factor, w_d, lr_factor).permute(0, 2, 1, 3).reshape(h_d, w_d, n_per_block)
            is_high_res_block = (mask_blocks.sum(dim=-1) > 0)
            is_low_res_block = ~is_high_res_block

            low_res_coords = grid_blocks[..., 0, :]
            output_grid = grid_blocks.clone()
            # Flatten before mask indexing — torch 2.5.1 mishandles 2D bool mask + int/slice.
            output_grid.view(h_d * w_d, n_per_block, 4)[is_low_res_block.view(-1), 0, :] = low_res_coords[is_low_res_block]

            res_mask_blocks = resolution_mask_grid.view(h_d, lr_factor, w_d, lr_factor).permute(0, 2, 1, 3).reshape(h_d, w_d, n_per_block)
            output_res_mask = res_mask_blocks.clone()
            output_res_mask[is_low_res_block, :] = 0.0

            valid_tokens = torch.ones(h_d, w_d, n_per_block, device=device, dtype=torch.bool)
            valid_tokens.view(h_d * w_d, n_per_block)[is_low_res_block.view(-1), 1:] = False

            output_grid = output_grid.reshape(-1, 4)
            output_res_mask = output_res_mask.reshape(-1)
            valid_tokens = valid_tokens.reshape(-1)

            tl_mask_blocks = torch.zeros_like(res_mask_blocks)
            tl_mask_blocks[..., 0] = 1.0
            tl_mask_blocks[is_low_res_block, :] = 1.0
            resolution_mask_top_left = tl_mask_blocks.reshape(-1)[valid_tokens]

            latent_image_ids = output_grid[valid_tokens]
            resolution_mask = output_res_mask[valid_tokens]

            return (
                latent_image_ids.to(device=device, dtype=dtype),
                resolution_mask.to(device=device, dtype=dtype),
                resolution_mask_top_left.to(device=device, dtype=dtype)
            )
        else:
            return (
                latent_image_ids.to(device=device, dtype=dtype),
                None,
                None
            )

    @staticmethod
    def _downsample_latents(
        latents: torch.Tensor,
        height: int,
        width: int,
        foveation_mask: Optional[torch.Tensor] = None,
        sample_mode: str = "nearest",
        latents_downsampled: Optional[torch.Tensor] = None,
        lr_factor: int = 2,
    ):
        """
        Downsample latents spatially based on foveation mask.

        For low-res blocks, averages/samples lr_factor x lr_factor regions to single tokens.
        High-res blocks keep all lr_factor*lr_factor tokens.
        """
        if foveation_mask is None:
            return latents

        mask = foveation_mask.to(latents.device)
        batch_size, seq_len, channels = latents.shape

        if height * width != seq_len:
            raise ValueError(f"Height ({height}) * Width ({width}) != seq_len ({seq_len})")

        output_latents = latents.clone()

        # Unpack to spatial format for downsampling
        latents_spatial = latents.view(batch_size, height, width, channels // 4, 2, 2)
        latents_spatial = latents_spatial.permute(0, 3, 1, 4, 2, 5)
        latents_spatial = latents_spatial.reshape(batch_size, channels // 4, height * 2, width * 2)

        # Downsample (naive interpolation when latents_downsampled isn't precomputed).
        if latents_downsampled is None:
            if sample_mode == "average":
                latents_down = F.interpolate(latents_spatial, scale_factor=1.0 / lr_factor, mode="bilinear") * lr_factor
            elif sample_mode == "top_left":
                latents_down = F.interpolate(latents_spatial, scale_factor=1.0 / lr_factor, mode="nearest")
            else:
                raise ValueError(f"Unknown sample_mode: {sample_mode}")

            latents_down = latents_down.reshape(batch_size, channels // 4, height // lr_factor, 2, width // lr_factor, 2)
            latents_down = latents_down.permute(0, 2, 4, 1, 3, 5)
            low_res_packed = latents_down.reshape(batch_size, height // lr_factor, width // lr_factor, channels)
        else:
            low_res_packed = latents_downsampled.reshape(batch_size, height // lr_factor, width // lr_factor, channels)

        n_per_block = lr_factor * lr_factor
        h_d, w_d = height // lr_factor, width // lr_factor
        output_latents = output_latents.view(batch_size, height, width, channels)
        output_latents = output_latents.view(batch_size, h_d, lr_factor, w_d, lr_factor, channels)
        output_latents = output_latents.permute(0, 1, 3, 2, 4, 5).reshape(batch_size, h_d, w_d, n_per_block, channels)

        mask_blocks = mask.view(h_d, lr_factor, w_d, lr_factor).permute(0, 2, 1, 3).reshape(h_d, w_d, n_per_block)
        is_high_res_block = (mask_blocks.sum(dim=-1) > 0)
        is_low_res_block = ~is_high_res_block

        # Flatten the (h_d, w_d) dims before mask indexing — torch 2.5.1 mishandles
        # mixing a 2D bool mask with an integer index in adjacent slots.
        output_latents_flat = output_latents.view(batch_size, h_d * w_d, n_per_block, channels)
        low_res_packed_flat = low_res_packed.view(batch_size, h_d * w_d, channels)
        is_low_res_block_flat = is_low_res_block.view(-1)
        output_latents_flat[:, is_low_res_block_flat, 0, :] = low_res_packed_flat[:, is_low_res_block_flat, :].to(output_latents.dtype)

        valid_tokens = torch.ones(h_d, w_d, n_per_block, device=latents.device, dtype=torch.bool)
        # Flatten before mask indexing — torch 2.5.1 mishandles 2D bool mask + slice.
        valid_tokens_flat = valid_tokens.view(h_d * w_d, n_per_block)
        valid_tokens_flat[is_low_res_block.view(-1), 1:] = False

        output_latents = output_latents.view(batch_size, -1, channels)
        valid_tokens = valid_tokens.view(-1)
        
        return output_latents[:, valid_tokens, :]

    @staticmethod
    def _downsample_remaining_high_res(
        latents: torch.Tensor,
        height: int,
        width: int,
        foveation_mask: torch.Tensor,
        downsample_mode: str = "nearest",
        lr_factor: int = 2,
    ):
        """
        Downsample all blocks to uniform low-res (single token per block).
        """
        batch_size, reduced_len, channels = latents.shape
        device = latents.device
        dtype = latents.dtype

        n_per_block = lr_factor * lr_factor
        h_d, w_d = height // lr_factor, width // lr_factor

        mask = foveation_mask.to(device)
        mask_blocks = mask.view(h_d, lr_factor, w_d, lr_factor).permute(0, 2, 1, 3).reshape(h_d, w_d, n_per_block)

        is_high_res_block = (mask_blocks.sum(dim=-1) > 0)
        is_low_res_block = ~is_high_res_block

        layout_mask = torch.ones(h_d, w_d, n_per_block, device=device, dtype=torch.bool)
        # Flatten before mask indexing — torch 2.5.1 mishandles 2D bool mask + slice.
        layout_mask.view(h_d * w_d, n_per_block)[is_low_res_block.view(-1), 1:] = False

        if layout_mask.sum() != reduced_len:
            raise ValueError(f"Latent length {reduced_len} != mask topology {layout_mask.sum()}")

        reconstructed_blocks = torch.zeros(batch_size, h_d * w_d * n_per_block, channels, device=device, dtype=dtype)
        reconstructed_blocks[:, layout_mask.view(-1), :] = latents
        reconstructed_blocks = reconstructed_blocks.view(batch_size, h_d, w_d, n_per_block, channels)

        if is_high_res_block.any():
            high_res_tokens = reconstructed_blocks[:, is_high_res_block, :, :]

            if downsample_mode == "average":
                pooled_val = high_res_tokens.mean(dim=2)
                # Flatten spatial dims before mask indexing — torch 2.5.1 mishandles
                # mixing a 2D bool mask with an integer index in adjacent slots.
                reconstructed_blocks_flat = reconstructed_blocks.view(batch_size, h_d * w_d, n_per_block, channels)
                reconstructed_blocks_flat[:, is_high_res_block.view(-1), 0, :] = pooled_val

        final_grid = reconstructed_blocks[:, :, :, 0, :]
        return final_grid.reshape(batch_size, h_d * w_d, channels)

    @staticmethod
    def _upsample_latents(
        latents: torch.Tensor,
        height: int,
        width: int,
        foveation_mask: Optional[torch.Tensor] = None,
        upsample_mode: str = "nearest",
        lr_factor: int = 2,
    ):
        """
        Upsamples reduced latents back to the original spatial grid size.
        """
        if foveation_mask is None:
            return latents

        batch_size, reduced_len, channels = latents.shape
        device = latents.device
        dtype = latents.dtype

        n_per_block = lr_factor * lr_factor
        h_d, w_d = height // lr_factor, width // lr_factor

        mask = foveation_mask.to(device)
        mask_blocks = mask.view(h_d, lr_factor, w_d, lr_factor).permute(0, 2, 1, 3).reshape(h_d, w_d, n_per_block)

        is_high_res_block = (mask_blocks.sum(dim=-1) > 0)
        is_low_res_block = ~is_high_res_block

        layout_mask = torch.ones(h_d, w_d, n_per_block, device=device, dtype=torch.bool)
        # Flatten before mask indexing — torch 2.5.1 mishandles 2D bool mask + slice.
        layout_mask.view(h_d * w_d, n_per_block)[is_low_res_block.view(-1), 1:] = False

        if layout_mask.sum() != reduced_len:
            raise ValueError(f"Latent length {reduced_len} does not match mask topology {layout_mask.sum()}.")

        reconstructed_blocks = torch.zeros(batch_size, h_d * w_d * n_per_block, channels, device=device, dtype=dtype)
        reconstructed_blocks[:, layout_mask.view(-1), :] = latents
        reconstructed_blocks = reconstructed_blocks.view(batch_size, h_d, w_d, n_per_block, channels)

        if upsample_mode == "nearest":
            if is_low_res_block.any():
                # Flatten before mask indexing — torch 2.5.1 mishandles 2D bool mask + slice.
                rb_flat = reconstructed_blocks.view(batch_size, h_d * w_d, n_per_block, channels)
                low_res_block_flat = is_low_res_block.view(-1)
                low_res_vals = rb_flat[:, low_res_block_flat, 0:1, :]
                rb_flat[:, low_res_block_flat, :, :] = low_res_vals.expand(-1, -1, n_per_block, -1)

            final_output = reconstructed_blocks.view(batch_size, h_d, w_d, lr_factor, lr_factor, channels)
            final_output = final_output.permute(0, 1, 3, 2, 4, 5).reshape(batch_size, height, width, channels)
        else:
            coarse_grid = reconstructed_blocks[..., 0, :].clone()

            if is_high_res_block.any():
                high_res_means = reconstructed_blocks[:, is_high_res_block, :, :].mean(dim=2)
                coarse_grid[:, is_high_res_block, :] = high_res_means

            coarse_grid = coarse_grid.permute(0, 3, 1, 2)

            fine_grid = F.interpolate(
                coarse_grid,
                size=(height, width),
                mode=upsample_mode,
                align_corners=False,
                antialias=True if upsample_mode != 'nearest' else False
            )

            final_output = fine_grid.permute(0, 2, 3, 1)

            high_res_raster = reconstructed_blocks.view(batch_size, h_d, w_d, lr_factor, lr_factor, channels)
            high_res_raster = high_res_raster.permute(0, 1, 3, 2, 4, 5).reshape(batch_size, height, width, channels)

            pixel_mask = is_high_res_block.repeat_interleave(lr_factor, dim=0).repeat_interleave(lr_factor, dim=1)
            final_output[:, pixel_mask, :] = high_res_raster[:, pixel_mask, :]

        return final_output.reshape(batch_size, height * width, channels)

    def foveated_training_forward(
        self,
        inputs: dict,
        timestep: torch.Tensor,
        timestep_id: int,
        prediction_type: str,
        lr_downsample_factor: int = 2,
    ):
        """Single-step foveated forward for training. Used by FoveatedFlowMatchSFTLoss.

        Downsamples latents with foveation, runs the DiT, and returns the foveated
        noise prediction [B, L, C] matching the mixed-resolution training target.
        """
        foveation_mask = inputs.get("foveation_mask")
        if foveation_mask is None:
            models = {name: getattr(self, name) for name in self.in_iteration_models}
            return self.model_fn(**models, **inputs, timestep=timestep)

        # fixed clean latent for tracking validation
        if self.fixed_clean_latent is None:
            self.fixed_clean_latent = inputs["input_latents"]

        latents = inputs["latents"]
        height = inputs["height"]
        width = inputs["width"]
        latent_height, latent_width = height // 16, width // 16
        batch_size = latents.shape[0]
        foveation_mask = foveation_mask.to(latents.device)

        latents_input = self._downsample_latents(
            latents, latent_height, latent_width,
            foveation_mask=foveation_mask,
            latents_downsampled=inputs["latents_downsampled"],
            lr_factor=lr_downsample_factor,
        )

        latent_ids, resolution_mask, resolution_mask_top_left = self._prepare_latent_image_ids_foveation(
            batch_size, latent_height, latent_width, self.device, latents.dtype,
            foveation_mask, lr_factor=lr_downsample_factor,
        )
        if latent_ids.ndim == 2:
            latent_ids = latent_ids.unsqueeze(0).expand(batch_size, -1, -1)

        inputs["latents"] = latents_input
        inputs["image_ids"] = latent_ids
        inputs["resolution_mask"] = resolution_mask
        inputs["resolution_mask_top_left"] = resolution_mask_top_left

        models = {name: getattr(self, name) for name in self.in_iteration_models}
        return self.model_fn(**models, **inputs, timestep=timestep)

    @torch.no_grad()
    def __call__(
        self,
        # Prompt
        prompt: str,
        negative_prompt: str = "",
        cfg_scale: float = 1.0,
        embedded_guidance: float = 4.0,
        # Image
        input_image: Image.Image = None,
        denoising_strength: float = 1.0,
        # Shape
        height: int = 1024,
        width: int = 1024,
        # Randomness
        seed: int = None,
        rand_device: str = "cpu",
        # Steps
        num_inference_steps: int = 30,
        progress_bar_cmd=tqdm,
        # Foveation
        foveation_mask: Optional[torch.Tensor] = None,
        decode_mode: str = "direct",
        prediction_type: str = "clean",
        soft_foveation_blend: bool = False,
        full_res_foveation_mask: Optional[torch.Tensor] = None,
        lr_downsample_factor: int = 2,
    ):
        """Generate images with optional foveated rendering.

        Args:
            prompt: Text prompt for generation.
            foveation_mask: [H, W] binary mask (1 = high-res, 0 = low-res).
                Shape should match the token grid (height//16, width//16).
            decode_mode: "direct" (HR-only decode) or "merge" (blend HR + LR decodes).
            soft_foveation_blend: If True, Gaussian-blur the mask boundary before
                blending in merge decode.
        """
        self.scheduler.set_timesteps(
            num_inference_steps,
            denoising_strength=denoising_strength,
            dynamic_shift_len=height // 16 * width // 16,
        )

        latent_height = height // 16
        latent_width = width // 16

        inputs_posi = {"prompt": prompt}
        inputs_nega = {"negative_prompt": negative_prompt}
        inputs_shared = {
            "cfg_scale": cfg_scale, "embedded_guidance": embedded_guidance,
            "input_image": input_image, "denoising_strength": denoising_strength,
            "height": height, "width": width,
            "seed": seed, "rand_device": rand_device,
            "num_inference_steps": num_inference_steps,
            "foveation_mask": foveation_mask,
            "soft_foveation_blend": soft_foveation_blend,
            "lr_downsample_factor": lr_downsample_factor,
        }
        for unit in self.units:
            inputs_shared, inputs_posi, inputs_nega = self.unit_runner(
                unit, self, inputs_shared, inputs_posi, inputs_nega,
            )

        # Denoise with foveation support
        self.load_models_to_device(self.in_iteration_models)
        models = {name: getattr(self, name) for name in self.in_iteration_models}

        latents = inputs_shared["latents"]
        # All downsampling happens once up front; the DiT operates on the mixed-resolution
        # token sequence for every denoising step.
        latents = self._downsample_latents(
            latents, latent_height, latent_width,
            foveation_mask=foveation_mask,
            latents_downsampled=inputs_shared["latents_downsampled"],
            lr_factor=lr_downsample_factor,
        )
        print('decode mode', decode_mode)

        start_t = time.time()
        for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)):
            timestep_val = timestep.unsqueeze(0).to(dtype=self.torch_dtype, device=self.device)
            batch_size = latents.shape[0]

            if foveation_mask is not None:
                latents_input = latents
                latent_ids, resolution_mask, resolution_mask_top_left = self._prepare_latent_image_ids_foveation(
                    batch_size, latent_height, latent_width,
                    self.device, latents.dtype, foveation_mask,
                    lr_factor=lr_downsample_factor,
                )
                if latent_ids.ndim == 2:
                    latent_ids = latent_ids.unsqueeze(0).expand(batch_size, -1, -1)
            else:
                latents_input = latents
                resolution_mask = None
                resolution_mask_top_left = None
                latent_ids = inputs_shared["image_ids"]

            inputs_shared["latents"] = latents_input
            inputs_shared["image_ids"] = latent_ids
            inputs_shared["resolution_mask"] = resolution_mask
            inputs_shared["resolution_mask_top_left"] = resolution_mask_top_left

            noise_pred = self.cfg_guided_model_fn(
                self.model_fn, cfg_scale,
                inputs_shared, inputs_posi, inputs_nega,
                **models,
                timestep=timestep_val,
                progress_id=progress_id,
            )

            if foveation_mask is None:
                noise_pred = noise_pred[:, :latents.size(1), :]

            latents = self.step(
                self.scheduler, progress_id=progress_id,
                noise_pred=noise_pred, latents=latents,
            )
        vit_end_t = time.time()

        inputs_shared["latents"] = latents

        # Decode
        self.load_models_to_device(['vae'])

        if decode_mode == "direct":
            latents_spatial = rearrange(latents, "B (H W) C -> B C H W", H=latent_height, W=latent_width)
            image = self.vae.decode(latents_spatial)
            image = self.vae_output_to_image(image)
        elif decode_mode == "merge" and foveation_mask is not None:
            # Merge decoding: decode low-res and high-res separately, then blend.
            # Note: DiffSynth VAE expects patched latents [B, 128, H/2, W/2].
            mixed_res_latents = latents  # already mixed-resolution

            # Step 2: Low-res path - downsample mixed-res to uniform low-res grid
            low_res_latents = self._downsample_remaining_high_res(
                mixed_res_latents, latent_height, latent_width, foveation_mask=foveation_mask,
                lr_factor=lr_downsample_factor
            )
            h_d, w_d = latent_height // lr_downsample_factor, latent_width // lr_downsample_factor
            # VAE expects patched format [B, 128, H/2, W/2] - just reshape, don't unpatchify
            low_res_spatial = low_res_latents.view(batch_size, h_d, w_d, -1).permute(0, 3, 1, 2)
            low_res_image = self.vae.decode(low_res_spatial)

            # Step 3: High-res path - upsample mixed-res to uniform high-res grid
            high_res_latents = self._upsample_latents(
                mixed_res_latents, latent_height, latent_width, foveation_mask=foveation_mask,
                lr_factor=lr_downsample_factor
            )
            # VAE expects patched format [B, 128, H/2, W/2] - just reshape, don't unpatchify
            high_res_spatial = high_res_latents.view(
                batch_size, latent_height, latent_width, -1
            ).permute(0, 3, 1, 2)
            high_res_image = self.vae.decode(high_res_spatial)

            # Step 4: Upsample low-res image and merge
            low_res_image = F.interpolate(low_res_image, size=(height, width), mode='bicubic', align_corners=False)
            
            if full_res_foveation_mask is not None:
                high_res_foveation_mask = full_res_foveation_mask[None, None, :, :].float()
            else:
                high_res_foveation_mask = F.interpolate(
                    foveation_mask[None, None, :, :].float(), 
                    size=(height, width), 
                    mode='bicubic'
                )
            high_res_foveation_mask = high_res_foveation_mask.to(latents.dtype).to(latents.device)
                
            if soft_foveation_blend:
                # Soft edge: Gaussian falloff over ~upscale_factor pixels at high-res
                upscale = height / latent_height
                high_res_foveation_mask = _gaussian_blur_mask_2d(
                    high_res_foveation_mask, sigma_pixels=upscale, device=latents.device, dtype=latents.dtype
                )
                high_res_foveation_mask = high_res_foveation_mask.clamp(0.0, 1.0)
            
            merged_image = low_res_image * (1 - high_res_foveation_mask) + high_res_image * high_res_foveation_mask
            image = self.vae_output_to_image(merged_image)
        else:
            # Fallback to direct
            latents_spatial = rearrange(latents, "B (H W) C -> B C H W", H=latent_height, W=latent_width)
            image = self.vae.decode(latents_spatial)
            image = self.vae_output_to_image(image)
        vae_end_t = time.time()

        # individually track time
        self.vit_timing = vit_end_t - start_t
        self.vae_timing = vae_end_t - vit_end_t
            
        self.load_models_to_device([])

        return image


# =====================================================================
# Pipeline Units (adapted from flux2_image.py)
# =====================================================================

class Flux2FoveatedUnit_ShapeChecker(PipelineUnit):
    def __init__(self):
        super().__init__(
            input_params=("height", "width"),
            output_params=("height", "width"),
        )

    def process(self, pipe: Flux2FoveatedImagePipeline, height, width):
        height, width = pipe.check_resize_height_width(height, width)
        return {"height": height, "width": width}


class Flux2FoveatedUnit_PromptEmbedder(PipelineUnit):
    def __init__(self):
        super().__init__(
            seperate_cfg=True,
            input_params_posi={"prompt": "prompt"},
            input_params_nega={"prompt": "negative_prompt"},
            output_params=("prompt_emb", "prompt_emb_mask"),
            onload_model_names=("text_encoder",)
        )
        self.system_message = "You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object attribution and actions without speculation."

    def format_text_input(self, prompts: List[str], system_message: str = None):
        cleaned_txt = [prompt.replace("[IMG]", "") for prompt in prompts]
        return [
            [
                {"role": "system", "content": [{"type": "text", "text": system_message}]},
                {"role": "user", "content": [{"type": "text", "text": prompt}]},
            ]
            for prompt in cleaned_txt
        ]

    def get_mistral_3_small_prompt_embeds(
        self,
        text_encoder,
        tokenizer,
        prompt: Union[str, List[str]],
        dtype: Optional[torch.dtype] = None,
        device: Optional[torch.device] = None,
        max_sequence_length: int = 512,
        system_message: str = None,
        hidden_states_layers: List[int] = (10, 20, 30),
    ):
        dtype = text_encoder.dtype if dtype is None else dtype
        device = text_encoder.device if device is None else device
        prompt = [prompt] if isinstance(prompt, str) else prompt

        messages_batch = self.format_text_input(prompts=prompt, system_message=system_message)
        inputs = tokenizer.apply_chat_template(
            messages_batch,
            add_generation_prompt=False,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
            padding="max_length",
            truncation=True,
            max_length=max_sequence_length,
        )

        input_ids = inputs["input_ids"].to(device)
        attention_mask = inputs["attention_mask"].to(device)

        output = text_encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True,
            use_cache=False,
        )

        out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1)
        out = out.to(dtype=dtype, device=device)

        batch_size, num_channels, seq_len, hidden_dim = out.shape
        prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim)

        return prompt_embeds
    
    def prepare_text_ids(self, x: torch.Tensor, t_coord: Optional[torch.Tensor] = None):
        B, L, _ = x.shape
        out_ids = []
        for i in range(B):
            t = torch.arange(1) if t_coord is None else t_coord[i]
            h = torch.arange(1)
            w = torch.arange(1)
            l = torch.arange(L)
            coords = torch.cartesian_prod(t, h, w, l)
            out_ids.append(coords)
        return torch.stack(out_ids)

    def encode_prompt(
        self,
        text_encoder,
        tokenizer,
        prompt: Union[str, List[str]],
        dtype = None,
        device: Optional[torch.device] = None,
        num_images_per_prompt: int = 1,
        prompt_embeds: Optional[torch.Tensor] = None,
        max_sequence_length: int = 512,
        text_encoder_out_layers: Tuple[int] = (10, 20, 30),
    ):
        prompt = [prompt] if isinstance(prompt, str) else prompt
        if prompt_embeds is None:
            prompt_embeds = self.get_mistral_3_small_prompt_embeds(
                text_encoder=text_encoder,
                tokenizer=tokenizer,
                prompt=prompt,
                dtype=dtype,
                device=device,
                max_sequence_length=max_sequence_length,
                system_message=self.system_message,
                hidden_states_layers=text_encoder_out_layers,
            )

        batch_size, seq_len, _ = prompt_embeds.shape
        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)

        text_ids = self.prepare_text_ids(prompt_embeds)
        text_ids = text_ids.to(device)
        return prompt_embeds, text_ids

    def process(self, pipe: Flux2FoveatedImagePipeline, prompt):
        if pipe.text_encoder_qwen3 is not None:
            return {}
        
        pipe.load_models_to_device(self.onload_model_names)
        prompt_embeds, text_ids = self.encode_prompt(
            pipe.text_encoder, pipe.tokenizer, prompt,
            dtype=pipe.torch_dtype, device=pipe.device,
        )
        return {"prompt_embeds": prompt_embeds, "text_ids": text_ids}


class Flux2FoveatedUnit_Qwen3PromptEmbedder(PipelineUnit):
    def __init__(self):
        super().__init__(
            seperate_cfg=True,
            input_params_posi={"prompt": "prompt"},
            input_params_nega={"prompt": "negative_prompt"},
            output_params=("prompt_emb", "prompt_emb_mask"),
            onload_model_names=("text_encoder_qwen3",)
        )
        self.hidden_states_layers = (9, 18, 27)

    def get_qwen3_prompt_embeds(
        self,
        text_encoder,
        tokenizer,
        prompt: Union[str, List[str]],
        dtype: Optional[torch.dtype] = None,
        device: Optional[torch.device] = None,
        max_sequence_length: int = 512,
    ):
        dtype = text_encoder.dtype if dtype is None else dtype
        device = text_encoder.device if device is None else device
        prompt = [prompt] if isinstance(prompt, str) else prompt

        all_input_ids = []
        all_attention_masks = []

        for single_prompt in prompt:
            messages = [{"role": "user", "content": single_prompt}]
            text = tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
                enable_thinking=False,
            )
            inputs = tokenizer(
                text,
                return_tensors="pt",
                padding="max_length",
                truncation=True,
                max_length=max_sequence_length,
            )
            all_input_ids.append(inputs["input_ids"])
            all_attention_masks.append(inputs["attention_mask"])

        input_ids = torch.cat(all_input_ids, dim=0).to(device)
        attention_mask = torch.cat(all_attention_masks, dim=0).to(device)

        with torch.inference_mode():
            output = text_encoder(
                input_ids=input_ids,
                attention_mask=attention_mask,
                output_hidden_states=True,
                use_cache=False,
            )

        out = torch.stack([output.hidden_states[k] for k in self.hidden_states_layers], dim=1)
        out = out.to(dtype=dtype, device=device)

        batch_size, num_channels, seq_len, hidden_dim = out.shape
        prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim)
        return prompt_embeds

    def prepare_text_ids(self, x: torch.Tensor, t_coord: Optional[torch.Tensor] = None):
        B, L, _ = x.shape
        out_ids = []
        for i in range(B):
            t = torch.arange(1) if t_coord is None else t_coord[i]
            h = torch.arange(1)
            w = torch.arange(1)
            l = torch.arange(L)
            coords = torch.cartesian_prod(t, h, w, l)
            out_ids.append(coords)
        return torch.stack(out_ids)

    def encode_prompt(
        self,
        text_encoder,
        tokenizer,
        prompt: Union[str, List[str]],
        dtype = None,
        device: Optional[torch.device] = None,
        num_images_per_prompt: int = 1,
        prompt_embeds: Optional[torch.Tensor] = None,
        max_sequence_length: int = 512,
    ):
        prompt = [prompt] if isinstance(prompt, str) else prompt
        if prompt_embeds is None:
            prompt_embeds = self.get_qwen3_prompt_embeds(
                text_encoder=text_encoder,
                tokenizer=tokenizer,
                prompt=prompt,
                dtype=dtype,
                device=device,
                max_sequence_length=max_sequence_length,
            )

        batch_size, seq_len, _ = prompt_embeds.shape
        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)

        text_ids = self.prepare_text_ids(prompt_embeds)
        text_ids = text_ids.to(device)
        return prompt_embeds, text_ids

    def process(self, pipe: Flux2FoveatedImagePipeline, prompt):
        if pipe.text_encoder_qwen3 is None:
            return {}
        
        pipe.load_models_to_device(self.onload_model_names)
        prompt_embeds, text_ids = self.encode_prompt(
            pipe.text_encoder_qwen3, pipe.tokenizer, prompt,
            dtype=pipe.torch_dtype, device=pipe.device,
        )
        return {"prompt_embeds": prompt_embeds, "text_ids": text_ids}


class Flux2FoveatedUnit_NoiseInitializer(PipelineUnit):
    def __init__(self):
        super().__init__(
            input_params=("height", "width", "seed", "rand_device", "lr_downsample_factor"),
            output_params=("noise", "downsampled_noise"),
        )

    def process(self, pipe: Flux2FoveatedImagePipeline, height, width, seed, rand_device, lr_downsample_factor=2):
        noise = pipe.generate_noise((1, 128, height//16, width//16), seed=seed, rand_device=rand_device, rand_torch_dtype=pipe.torch_dtype)
        noise_downsampled = pipe.generate_noise((1, 128, height//16//lr_downsample_factor, width//16//lr_downsample_factor), seed=seed, rand_device=rand_device, rand_torch_dtype=pipe.torch_dtype)
        noise = noise.reshape(1, 128, height//16 * width//16).permute(0, 2, 1)
        noise_downsampled = noise_downsampled.reshape(1, 128, height//16//lr_downsample_factor * width//16//lr_downsample_factor).permute(0, 2, 1)
        return {"noise": noise, "noise_downsampled": noise_downsampled}


# Saliency-to-mask constants (same as run_saliency_detection.py / fovetaion_masks_from_video.py)
_INFERENCE_SCALE = 0.5
_SPATIAL_SIGMA = 20.0
_R_MIN, _R_MAX = 0.2, 0.5
_COVERAGE = 0.9
_DENSE_THRESHOLD = 0.3


def bbox_to_foveation_mask(
    height: int,
    width: int,
    bboxes_xyxy,
    sigma: float = 10.0,
):
    """
    Create a full-resolution foveation mask (H, W) in [0, 1] from a list of
    bounding boxes in absolute pixel coordinates (x1, y1, x2, y2).

    This function mirrors bbox_to_foveation_mask in misc/run_bbox_detection.py:
    the mask is 1 inside any bbox, 0 outside, followed by Gaussian smoothing
    and normalization. This yields a smooth "high-acuity" region covering all
    detected objects.
    """
    mask = np.zeros((height, width), dtype=np.float32)
    for box in bboxes_xyxy:
        x1, y1, x2, y2 = box
        x1_i = max(0, min(width - 1, int(np.floor(x1))))
        y1_i = max(0, min(height - 1, int(np.floor(y1))))
        x2_i = max(0, min(width, int(np.ceil(x2))))
        y2_i = max(0, min(height, int(np.ceil(y2))))
        if x2_i <= x1_i or y2_i <= y1_i:
            continue
        mask[y1_i:y2_i, x1_i:x2_i] = 1.0

    if sigma > 0:
        mask = gaussian_filter(mask, sigma=sigma)
    max_val = mask.max()
    if max_val > 0:
        mask = mask / max_val
    return mask


class Flux2FoveatedUnit_InputImageEmbedder(PipelineUnit):
    def __init__(self):
        super().__init__(
            input_params=("input_image", "noise", "noise_downsampled", "lr_downsample_factor"),
            output_params=("latents", "latents_downsampled", "input_latents", "input_latents_downsampled", "foveation_mask"),
            onload_model_names=("vae",)
        )
        try:
            self.saliency_detector = deepgaze_pytorch.DeepGazeIIE(pretrained=True)
            self.saliency_detector.eval()
            print("Saliency detector loaded")
        except Exception as e:
            print(e)
            self.saliency_detector = None
            print("Saliency detector not loaded")
        self.bbox_detector = None

    def get_foveation_mask_from_image(
        self,
        pipe: "Flux2FoveatedImagePipeline",
        image: torch.Tensor,
    ) -> Optional[torch.Tensor]:
        """
        Create a foveation mask from the input image, depending on the
        foveated_training_mode on the pipeline:
        - "saliency": current DeepGaze-based saliency-to-mask logic.
        - "bbox": use bbox_to_foveation_mask (same logic as run_bbox_detection.py)
          on provided bounding boxes, then downsample to latent grid size.
        image: (1, 3, H, W) in [-1, 1]. Returns (H//16, W//16) mask.
        """
        mode = getattr(pipe, "foveated_training_mode", "saliency")

        if mode == "bbox":
            # Run YOLO detector directly on the input image to obtain bboxes,
            # then convert them to a foveation mask using the exact same logic
            # as misc/run_bbox_detection.py.
            if image.dim() == 3:
                image = image.unsqueeze(0)
            device = image.device
            h_orig, w_orig = int(image.shape[2]), int(image.shape[3])

            # Lazy-load YOLO detector with same defaults as run_bbox_detection.py
            if self.bbox_detector is None:
                try:
                    from ultralytics import YOLO
                except ImportError as exc:
                    print(
                        "ultralytics not installed; bbox-based foveation requires "
                        "`pip install ultralytics`. Falling back to no foveation mask."
                    )
                    return None
                # Use same default model and keep device selection at predict-time
                self.bbox_detector = YOLO("yolov8n.pt")
                self.bbox_detector.model.float()

            # Convert input tensor ([-1, 1]) to uint8 RGB numpy array, same scaling
            # convention as saliency path (0-255).
            img_01 = (image[0].float() + 1.0) / 2.0
            img_255 = (img_01.clamp(0, 1) * 255.0).permute(1, 2, 0).cpu().numpy().astype(np.uint8)

            # Device string for YOLO; follow run_bbox_detection default conf params.
            device_str = str(pipe.device) if hasattr(pipe, "device") else "cpu"
            results = self.bbox_detector.predict(
                img_255,
                conf=0.25,
                device=device_str,
                verbose=False,
            )
            if not results:
                # No detections result: use full high-res mask
                mask_np = np.ones((h_orig, w_orig), dtype=np.float32)
            else:
                res = results[0]
                if res.boxes is None or res.boxes.xyxy is None or len(res.boxes) == 0:
                    # No boxes: use full high-res mask
                    mask_np = np.ones((h_orig, w_orig), dtype=np.float32)
                else:
                    boxes_np = res.boxes.xyxy.cpu().numpy()
                    # Full-resolution [H, W] mask using the exact bbox->mask logic.
                    mask_np = bbox_to_foveation_mask(h_orig, w_orig, boxes_np)
            mask_hr = torch.from_numpy(mask_np).to(device=device, dtype=torch.float32)
            # Downsample by 32, binarize, then upsample 2x so effective factor is 16
            mask_hr_4d = mask_hr.unsqueeze(0).unsqueeze(0)
            h_32, w_32 = h_orig // 32, w_orig // 32
            mask_32 = F.interpolate(mask_hr_4d, size=(h_32, w_32), mode="bilinear", align_corners=False)
            mask_binary_32 = (mask_32.squeeze(0).squeeze(0) > 0.5).to(torch.float32).unsqueeze(0).unsqueeze(0)
            latent_h, latent_w = h_orig // 16, w_orig // 16
            mask_binary = F.interpolate(mask_binary_32, size=(latent_h, latent_w), mode="nearest")
            return mask_binary

        # Default: saliency-based foveation mask
        if self.saliency_detector is None:
            return None
        if image.dim() == 3:
            image = image.unsqueeze(0)
        device = image.device
        self.saliency_detector = self.saliency_detector.to(device)
        h_orig, w_orig = int(image.shape[2]), int(image.shape[3])
        # DeepGaze input: [0, 255] float, same as run_saliency_detection _frame_to_tensor
        img_01 = (image.float() + 1.0) / 2.0
        img_255 = img_01.clamp(0, 1) * 255.0
        h_inf = int(h_orig * _INFERENCE_SCALE)
        w_inf = int(w_orig * _INFERENCE_SCALE)
        img_inf = F.interpolate(img_255, size=(h_inf, w_inf), mode="bilinear", align_corners=False)
        centerbias = torch.zeros(1, h_inf, w_inf, device=device, dtype=torch.float32)
        with torch.no_grad():
            log_density = self.saliency_detector(img_inf, centerbias)
        smap = torch.exp(log_density[0, 0]).cpu().numpy()
        smax = float(smap.max())
        if smax > 0:
            smap = smap / smax
        smoothed = gaussian_filter(smap, sigma=_SPATIAL_SIGMA)
        smax = float(smoothed.max())
        if smax > 0:
            smoothed = smoothed / smax
        # _saliency_to_params (same as run_saliency_detection)
        h, w = smoothed.shape
        eps = 1e-8
        peak = float(smoothed.max())
        if peak < eps:
            cx_norm, cy_norm, r = 0.0, 0.0, _R_MIN
        else:
            thresh = _DENSE_THRESHOLD * peak
            dense_mask = smoothed >= thresh
            dense_smap = smoothed * dense_mask
            dense_total = dense_smap.sum() + eps
            y_coords = np.arange(h, dtype=np.float64)
            x_coords = np.arange(w, dtype=np.float64)
            x_grid, y_grid = np.meshgrid(x_coords, y_coords)
            cx_px = (dense_smap * x_grid).sum() / dense_total
            cy_px = (dense_smap * y_grid).sum() / dense_total
            cx_norm = (cx_px / w) - 0.5
            cy_norm = (cy_px / h) - 0.5
            dist = np.sqrt((x_grid - cx_px) ** 2 + (y_grid - cy_px) ** 2)
            dense_pixels = dense_mask.ravel()
            flat_dist = dist.ravel()[dense_pixels]
            flat_sal = dense_smap.ravel()[dense_pixels]
            sort_idx = np.argsort(flat_dist)
            cumsum = np.cumsum(flat_sal[sort_idx])
            target = _COVERAGE * dense_total
            hit = min(np.searchsorted(cumsum, target), len(flat_dist) - 1)
            radius_px = flat_dist[sort_idx[hit]]
            half_diag = 0.5 * np.sqrt(h ** 2 + w ** 2)
            r = float(np.clip(radius_px / half_diag, _R_MIN, _R_MAX))
        # Binarized circular mask at full image res
        cx = (cx_norm + 0.5) * w_orig
        cy = (cy_norm + 0.5) * h_orig
        diagonal = (h_orig ** 2 + w_orig ** 2) ** 0.5
        radius_px = r * (diagonal / 2.0)
        y = np.arange(h_orig, dtype=np.float64)
        x = np.arange(w_orig, dtype=np.float64)
        xx, yy = np.meshgrid(x, y)
        dist_sq = (xx - cx) ** 2 + (yy - cy) ** 2
        mask_hr = (dist_sq <= radius_px ** 2).astype(np.float32)
        mask_hr = torch.from_numpy(mask_hr).to(device=device, dtype=torch.float32)
        # Downsample by 32, binarize, then upsample 2x so effective factor is 16
        mask_hr_4d = mask_hr.unsqueeze(0).unsqueeze(0)
        h_32, w_32 = h_orig // 32, w_orig // 32
        mask_32 = F.interpolate(mask_hr_4d, size=(h_32, w_32), mode="bilinear", align_corners=False)
        mask_binary_32 = (mask_32.squeeze(0).squeeze(0) > 0.5).to(torch.float32).unsqueeze(0).unsqueeze(0)
        latent_h, latent_w = h_orig // 16, w_orig // 16
        mask_binary = F.interpolate(mask_binary_32, size=(latent_h, latent_w), mode="nearest")
        return mask_binary

    def process(self, pipe: Flux2FoveatedImagePipeline, input_image, noise, noise_downsampled, lr_downsample_factor=2):
        if input_image is None:
            return {
                "latents": noise,
                "latents_downsampled": noise_downsampled,
                "input_latents": None,
                "input_latents_downsampled": None,
                "foveation_mask": None,
            }
        pipe.load_models_to_device(["vae"])
        image = pipe.preprocess_image(input_image)
        input_latents = pipe.vae.encode(image)
        # Decide between saliency-based and bbox-based mask creation.
        foveation_mask = self.get_foveation_mask_from_image(pipe, image)
        #save_image((image + 1) / 2.0, "/home/brianchc/foveated_diffusion/paper_figures/input_image.png")
        #save_image(foveation_mask, "/home/brianchc/foveated_diffusion/paper_figures/foveation_mask.png")
        #exit()
        _, _, latent_height, latent_width = input_latents.shape
        input_latents = rearrange(input_latents, "B C H W -> B (H W) C")

        # create downsample
        image_downsampled = F.interpolate(image, size=(image.shape[-2]//lr_downsample_factor, image.shape[-1]//lr_downsample_factor), mode='bicubic', align_corners=False)
        input_latents_downsampled = pipe.vae.encode(image_downsampled)
        input_latents_downsampled = rearrange(input_latents_downsampled, "B C H W -> B (H W) C") # 4X less tokens


        if pipe.scheduler.training:
            return {
                "latents": noise,
                "latents_downsampled": noise_downsampled,
                "input_latents": input_latents,
                "input_latents_downsampled": input_latents_downsampled,
                "foveation_mask": foveation_mask,
            }
        else:
            latents = pipe.scheduler.add_noise(input_latents, noise, timestep=pipe.scheduler.timesteps[0])
            latents_downsampled = pipe.scheduler.add_noise(input_latents_downsampled, noise_downsampled, timestep=pipe.scheduler.timesteps[0])
            return {
                "latents": latents,
                "latents_downsampled": latents_downsampled,
                "input_latents": input_latents,
                "input_latents_downsampled": input_latents_downsampled,
                "foveation_mask": foveation_mask,
            }


class Flux2FoveatedUnit_ImageIDs(PipelineUnit):
    def __init__(self):
        super().__init__(
            input_params=("height", "width", "foveation_mask", "lr_downsample_factor"),
            output_params=("image_ids",),
        )

    def prepare_latent_ids(self, height, width, device, dtype, foveation_mask=None, lr_factor=2):
        return Flux2FoveatedImagePipeline._prepare_latent_image_ids_foveation(
            batch_size=1,
            height=height,
            width=width,
            device=device,
            dtype=dtype,
            foveation_mask=foveation_mask,
            lr_factor=lr_factor,
        )

    def process(self, pipe: Flux2FoveatedImagePipeline, height, width, foveation_mask, lr_downsample_factor=2):
        latent_ids, resolution_mask, resolution_mask_top_left = self.prepare_latent_ids(
            height // 16, width // 16, pipe.device, pipe.torch_dtype, foveation_mask, lr_factor=lr_downsample_factor
        )
        latent_ids = latent_ids.unsqueeze(0).expand(1, -1, -1)
        return {
            "image_ids": latent_ids,
            "resolution_mask": resolution_mask,
            "resolution_mask_top_left": resolution_mask_top_left,
        }


# =====================================================================
# Model Function for Foveated FLUX2
# =====================================================================

def model_fn_flux2_foveated(
    dit: Flux2DiTFoveated,
    latents=None,
    timestep=None,
    embedded_guidance=None,
    prompt_embeds=None,
    text_ids=None,
    image_ids=None,
    resolution_mask=None,
    resolution_mask_top_left=None,
    lr_downsample_factor=2,
    use_gradient_checkpointing=False,
    use_gradient_checkpointing_offload=False,
    **kwargs,
):
    """Model function for foveated FLUX2 generation."""
    image_seq_len = latents.shape[1]
    embedded_guidance = torch.tensor([embedded_guidance], device=latents.device)

    model_output = dit(
        hidden_states=latents,
        timestep=timestep / 1000,
        guidance=embedded_guidance,
        encoder_hidden_states=prompt_embeds,
        txt_ids=text_ids,
        img_ids=image_ids,
        resolution_mask=resolution_mask,
        resolution_mask_top_left=resolution_mask_top_left,
        lr_factor=lr_downsample_factor,
        use_gradient_checkpointing=use_gradient_checkpointing,
        use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
    )
    
    model_output = model_output[:, :image_seq_len]
    return model_output