Publish Collage Diffusion pipeline at 9fff847022ee
30cbfdc verified about 19 hours ago
33 kB
"""Diffusers pipeline and layer helpers for collage harmonization."""

from __future__ import annotations

import math
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Literal, Protocol

import numpy as np
import torch
import torch.nn.functional as F
from diffusers import StableDiffusionImg2ImgPipeline
from diffusers.models.attention_processor import Attention
from diffusers.pipelines.stable_diffusion.pipeline_output import (
    StableDiffusionPipelineOutput,
)
from jaxtyping import Float
from laygen.common.randomness import randn, resolve_torch_generator
from PIL import Image, ImageFilter


class TokenizerProtocol(Protocol):
    """The tokenizer operation needed to locate layer phrases in a prompt."""

    def encode(self, text: str, *, add_special_tokens: bool) -> list[int]:
        """Encode text into token ids."""


@dataclass(frozen=True, slots=True)
class CollageLayer:
    """One positioned RGBA layer and the controls associated with its prompt.
    Args:
        image: RGBA image for the layer.
        prompt: Layer description that also appears in the scene prompt.
        position: Pixel ``(x, y)`` position on the output canvas.
        noise_level: Layer-specific amount of source preservation, in ``[0, 1]``.
        positive_attention_weight: Attention-logit boost inside the visible mask.
        negative_attention_weight: Attention-logit penalty outside the visible mask.
        textual_inversion_token: Optional token inserted before ``prompt``.
    Raises:
        TypeError: If the image or position has the wrong type.
        ValueError: If the image mode, prompt, or numeric control is invalid.
    """

    image: Image.Image
    prompt: str
    position: tuple[int, int] = (0, 0)
    noise_level: float = 0.7
    positive_attention_weight: float = 0.0
    negative_attention_weight: float = 0.0
    textual_inversion_token: str | None = None

    def __post_init__(self) -> None:
        """Validate layer input modes and scalar controls."""
        if not isinstance(self.image, Image.Image):
            raise TypeError("image must be a PIL Image")

        if self.image.mode != "RGBA":
            raise ValueError("image must use RGBA mode")

        if not self.prompt.strip():
            raise ValueError("prompt must be a non-empty layer description")

        if len(self.position) != 2 or any(
            not isinstance(value, int) or isinstance(value, bool)
            for value in self.position
        ):
            raise TypeError("position must contain two integer pixel coordinates")

        if not math.isfinite(self.noise_level) or not 0.0 <= self.noise_level <= 1.0:
            raise ValueError("noise_level must be finite and between 0 and 1")

        for name, weight in (
            ("positive_attention_weight", self.positive_attention_weight),
            ("negative_attention_weight", self.negative_attention_weight),
        ):
            if not math.isfinite(weight) or weight < 0:
                raise ValueError(f"{name} must be finite and non-negative")

        if (
            self.textual_inversion_token is not None
            and not self.textual_inversion_token.strip()
        ):
            raise ValueError("textual_inversion_token must be non-empty when provided")


def _placed_alpha(
    layer: CollageLayer,
    *,
    width: int,
    height: int,
) -> Float[np.ndarray, "height width"]:
    canvas = Image.new("RGBA", (width, height), (0, 0, 0, 0))
    # Using the RGBA image as its own paste mask multiplies partial alpha by itself.
    canvas.paste(layer.image, layer.position, layer.image)

    # Keep alpha arithmetic at float64 until conversion to the 8-bit noise map.
    return np.asarray(canvas.getchannel("A"), dtype=np.float64) / 255.0


def compose_collage_layers(
    layers: Sequence[CollageLayer],
    *,
    width: int = 512,
    height: int = 512,
    background_color: tuple[int, int, int] = (150, 150, 150),
) -> tuple[Image.Image, list[Float[np.ndarray, "height width"]]]:
    """Composite layers back-to-front and return their visible alpha masks.
    Args:
        layers: Layers ordered from back to front.
        width: Output canvas width in pixels.
        height: Output canvas height in pixels.
        background_color: RGB color used outside all layers.
    Returns:
        The RGB collage and one visible alpha mask per input layer.
    Raises:
        ValueError: If there are no layers or the canvas is invalid.
    Examples:
        >>> layer = CollageLayer(Image.new("RGBA", (2, 2), (255, 0, 0, 255)), "red square")
        >>> image, masks = compose_collage_layers([layer], width=2, height=2)
        >>> image.getpixel((0, 0)), float(masks[0][0, 0])
        ((255, 0, 0), 1.0)
    """
    if not layers:
        raise ValueError("layers must contain at least one RGBA image")

    if width < 1 or height < 1:
        raise ValueError("canvas width and height must be positive")

    if len(background_color) != 3 or any(
        not 0 <= channel <= 255 for channel in background_color
    ):
        raise ValueError("background_color must contain three byte values")

    composite = Image.new("RGB", (width, height), background_color)
    for layer in layers:
        # Use each RGBA layer's alpha as its paste mask while compositing.
        composite.paste(layer.image, layer.position, layer.image)

    visible_masks: list[Float[np.ndarray, "height width"]] = []
    for index, layer in enumerate(layers):
        layer_alpha = _placed_alpha(layer, width=width, height=height)

        top_canvas = Image.new("RGBA", (width, height), (0, 0, 0, 0))
        for top_layer in layers[index + 1 :]:
            top_canvas.paste(top_layer.image, top_layer.position, top_layer.image)

        top_alpha = np.asarray(top_canvas.getchannel("A"), dtype=np.float64) / 255.0
        visible_masks.append(layer_alpha * (1.0 - top_alpha))

    return composite, visible_masks


def make_layer_noise_map(
    visible_masks: Sequence[Float[np.ndarray, "height width"]],
    layers: Sequence[CollageLayer],
    *,
    strength: float,
    blur_radius: float = 10.0,
) -> Float[np.ndarray, "height width"]:
    """Blend per-layer noise controls over the global image-to-image strength.
    Args:
        visible_masks: Visible alpha masks in back-to-front layer order.
        layers: Layers whose ``noise_level`` values drive the map.
        strength: Global image-to-image strength used outside layer masks.
        blur_radius: Gaussian blur radius for smooth layer boundaries.
    Returns:
        A grayscale float map with values in ``[0, 1]``.
    Raises:
        ValueError: If masks, layers, or controls are inconsistent.
    """
    if len(visible_masks) != len(layers) or not layers:
        raise ValueError("visible_masks and layers must have the same non-zero length")

    if not math.isfinite(strength) or not 0.0 < strength <= 1.0:
        raise ValueError("strength must be finite and between 0 and 1")

    if not math.isfinite(blur_radius) or blur_radius < 0:
        raise ValueError("blur_radius must be finite and non-negative")

    noise_map = np.full_like(visible_masks[0], int(strength * 255), dtype=np.uint8)
    for mask, layer in zip(visible_masks, layers, strict=True):
        if mask.shape != noise_map.shape:
            raise ValueError("all visible masks must have the same shape")

        if not np.isfinite(mask).all() or np.any(mask < 0.0) or np.any(mask > 1.0):
            raise ValueError("visible masks must contain finite values between 0 and 1")

        layer_values = Image.fromarray(
            (mask * layer.noise_level * 255).astype(np.uint8)
        )
        layer_mask = Image.fromarray((mask * 255).astype(np.uint8))
        collage_map = Image.fromarray(noise_map)
        collage_map.paste(layer_values, (0, 0), layer_mask)
        noise_map = np.asarray(collage_map, dtype=np.uint8)

    grayscale = Image.fromarray(noise_map)
    if blur_radius > 0:
        grayscale = grayscale.filter(ImageFilter.GaussianBlur(blur_radius))

    return np.asarray(grayscale, dtype=np.float32) / 255.0


def _find_token_span(
    prompt_token_ids: Sequence[int],
    layer_token_ids: Sequence[int],
) -> tuple[int, ...]:
    if not layer_token_ids:
        raise ValueError("layer descriptions must produce at least one tokenizer token")

    token_count = len(layer_token_ids)
    for start in range(1, len(prompt_token_ids) - token_count + 1):
        if tuple(prompt_token_ids[start : start + token_count]) == tuple(
            layer_token_ids
        ):
            return tuple(range(start, start + token_count))

    raise ValueError("each layer description must appear as a token span in the prompt")


def _prompt_token_spans(
    tokenizer: TokenizerProtocol,
    prompt: str,
    layers: Sequence[CollageLayer],
) -> tuple[tuple[int, ...], ...]:
    prompt_ids = tokenizer.encode(prompt, add_special_tokens=True)
    spans: list[tuple[int, ...]] = []
    textual_inversion_enabled = any(
        layer.textual_inversion_token is not None for layer in layers
    )

    for layer in layers:
        try:
            layer_ids = tokenizer.encode(f" {layer.prompt}", add_special_tokens=False)
            span = _find_token_span(prompt_ids, layer_ids)
        except ValueError:
            layer_ids = tokenizer.encode(layer.prompt, add_special_tokens=False)
            span = _find_token_span(prompt_ids, layer_ids)

        if textual_inversion_enabled:
            # Extend each span's start one token left to include a prepended TI token.
            span = tuple(range(span[0] - 1, span[-1] + 1))

        spans.append(span)

    return tuple(spans)


def _inject_textual_inversion_tokens(
    prompt: str,
    layers: Sequence[CollageLayer],
    embedding_sources: Mapping[str, str | Path] | None,
) -> str:
    requested_tokens = {
        layer.textual_inversion_token
        for layer in layers
        if layer.textual_inversion_token is not None
    }
    supplied_tokens = set() if embedding_sources is None else set(embedding_sources)
    unused_tokens = supplied_tokens - requested_tokens
    if unused_tokens:
        raise ValueError(
            f"textual inversion embeddings have no matching layer: {sorted(unused_tokens)}"
        )

    for layer in layers:
        token = layer.textual_inversion_token
        if token is None:
            continue

        if prompt.count(layer.prompt) != 1:
            raise ValueError(
                "each layer description must occur exactly once in the scene prompt"
            )

        prompt = prompt.replace(layer.prompt, f"{token} {layer.prompt}", 1)

    return prompt


def _resize_mask_to_query(
    mask: Float[np.ndarray, "height width"],
    *,
    size: tuple[int, int],
    device: torch.device,
) -> Float[torch.Tensor, "query"]:
    """Resize a visible mask with the reference tensor-sampling semantics."""
    mask_tensor = torch.as_tensor(mask, device=device, dtype=torch.float64)
    batched_mask = mask_tensor.unsqueeze(0).unsqueeze(0)
    resized_mask = F.interpolate(
        batched_mask,
        size=size,
        mode="bilinear",
        align_corners=False,
        antialias=False,
    )
    resized_mask = resized_mask.squeeze(0).squeeze(0)
    query_length = size[0] * size[1]

    return resized_mask.reshape(query_length)


class _CollageCrossAttentionProcessor:
    """Apply visible-layer logit guidance through Diffusers' processor API."""

    def __init__(
        self,
        *,
        visible_masks: Sequence[Float[np.ndarray, "height width"]],
        token_spans: Sequence[Sequence[int]],
        layers: Sequence[CollageLayer],
        output_size: tuple[int, int],
        batch_size: int,
        classifier_free_guidance: bool,
    ) -> None:
        self.visible_masks = tuple(visible_masks)
        self.token_spans = tuple(tuple(span) for span in token_spans)
        self.layers = tuple(layers)
        self.output_size = output_size
        self.batch_size = batch_size
        self.classifier_free_guidance = classifier_free_guidance
        self.sigma = 1.0

    def __call__(
        self,
        attn: Attention,
        hidden_states: Float[torch.Tensor, "batch sequence channels"],
        encoder_hidden_states: Float[torch.Tensor, "batch tokens channels"]
        | None = None,
        attention_mask: Float[torch.Tensor, "..."] | None = None,
        temb: Float[torch.Tensor, "..."] | None = None,
        scale: float | None = None,
    ) -> Float[torch.Tensor, "batch sequence channels"]:
        # Add per-layer score bias before softmax; CFG leaves unconditional rows unchanged.
        del scale
        residual = hidden_states
        spatial_dimensions: tuple[int, int, int] | None = None
        if attn.spatial_norm is not None:
            hidden_states = attn.spatial_norm(hidden_states, temb)

        input_ndim = hidden_states.ndim
        if input_ndim == 4:
            batch_size, channel, height, width = hidden_states.shape
            spatial_dimensions = (channel, height, width)
            hidden_states = hidden_states.view(batch_size, channel, height * width)
            hidden_states = hidden_states.transpose(1, 2)

        batch_size, sequence_length, _ = (
            hidden_states.shape
            if encoder_hidden_states is None
            else encoder_hidden_states.shape
        )
        if attention_mask is not None:
            attention_mask = attn.prepare_attention_mask(
                attention_mask, sequence_length, batch_size
            )
            attention_mask = attention_mask.view(
                batch_size, attn.heads, -1, attention_mask.shape[-1]
            )

        if attn.group_norm is not None:
            hidden_states = hidden_states.transpose(1, 2)
            hidden_states = attn.group_norm(hidden_states)
            hidden_states = hidden_states.transpose(1, 2)

        query = attn.to_q(hidden_states)
        if encoder_hidden_states is None:
            encoder_hidden_states = hidden_states
        elif attn.norm_cross:
            encoder_hidden_states = attn.norm_encoder_hidden_states(
                encoder_hidden_states
            )

        to_k = attn.to_k
        to_v = attn.to_v
        to_out = attn.to_out
        if to_k is None or to_v is None or to_out is None:
            raise ValueError(
                "cross-attention requires key, value, and output projections"
            )

        key = to_k(encoder_hidden_states)
        value = to_v(encoder_hidden_states)
        head_dim = key.shape[-1] // attn.heads
        query = query.view(batch_size, -1, attn.heads, head_dim)
        query = query.transpose(1, 2)
        key = key.view(batch_size, -1, attn.heads, head_dim)
        key = key.transpose(1, 2)
        value = value.view(batch_size, -1, attn.heads, head_dim)
        value = value.transpose(1, 2)

        if attn.norm_q is not None:
            query = attn.norm_q(query)

        if attn.norm_k is not None:
            key = attn.norm_k(key)

        score_query = query.float() if attn.upcast_attention else query
        score_key = key.float() if attn.upcast_attention else key
        scores = torch.matmul(score_query, score_key.transpose(-1, -2)) * attn.scale
        if attention_mask is not None:
            scores = scores + attention_mask

        self._apply_layer_guidance(scores)
        softmax_scores = scores.float() if attn.upcast_softmax else scores
        attention_probs = softmax_scores.softmax(dim=-1).to(value.dtype)
        output = torch.matmul(attention_probs, value)
        output = output.transpose(1, 2)
        output = output.reshape(batch_size, -1, attn.heads * head_dim)
        output = output.to(query.dtype)
        output = to_out[0](output)
        output = to_out[1](output)

        if spatial_dimensions is not None:
            channel, height, width = spatial_dimensions
            output = output.transpose(-1, -2)
            output = output.reshape(batch_size, channel, height, width)

        if attn.residual_connection:
            output = output + residual

        return output / attn.rescale_output_factor

    def _apply_layer_guidance(
        self, scores: Float[torch.Tensor, "batch heads query tokens"]
    ) -> None:
        query_length = scores.shape[-2]
        output_height, output_width = self.output_size
        target_height = max(
            1, int(round(math.sqrt(query_length * output_height / output_width)))
        )
        while query_length % target_height:
            target_height -= 1

        target_width = query_length // target_height
        latent_masks = [
            _resize_mask_to_query(
                mask,
                size=(target_height, target_width),
                device=scores.device,
            )
            for mask in self.visible_masks
        ]
        # CFG batches unconditional rows before conditional rows; guide only the latter.
        cfg_offset = self.batch_size if self.classifier_free_guidance else 0
        sigma_scale = math.log(math.log(max(self.sigma, 0.0) + 1.01) + 1.0)

        for batch_index in range(self.batch_size):
            conditional_index = batch_index + cfg_offset
            for mask, span, layer in zip(
                latent_masks, self.token_spans, self.layers, strict=True
            ):
                for token_index in span:
                    if token_index >= scores.shape[-1]:
                        raise ValueError(
                            "layer token span exceeds the text encoder sequence length"
                        )

                    token_scores = scores[conditional_index, :, :, token_index]
                    max_score = token_scores.max()
                    visible_mask = mask.unsqueeze(0)
                    non_visible_mask = 1.0 - visible_mask
                    positive_bias = visible_mask * layer.positive_attention_weight
                    negative_bias = non_visible_mask * layer.negative_attention_weight
                    bias = (positive_bias - negative_bias) * sigma_scale * max_score
                    scores[conditional_index, :, :, token_index] += bias


class CollageDiffusionPipeline(StableDiffusionImg2ImgPipeline):
    """Harmonize positioned image layers with layer-aware diffusion controls.
    The pipeline loads the standard Stable Diffusion image-to-image components and
    replaces only the sampling orchestration needed for visible-layer attention
    and noise guidance.
    Examples:
        >>> from PIL import Image
        >>> from collage_diffusion import CollageLayer
        >>> layer = CollageLayer(Image.new("RGBA", (16, 16), (220, 80, 30, 255)), "orange toy")
        >>> layer.prompt
        'orange toy'
    """

    @torch.no_grad()
    def __call__(
        self,
        prompt: str,
        layers: Sequence[CollageLayer],
        *,
        negative_prompt: str | None = None,
        strength: float = 0.7,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        num_images_per_prompt: int = 1,
        generator: torch.Generator | None = None,
        seed: int | None = None,
        textual_inversion_embeddings: Mapping[str, str | Path] | None = None,
        height: int = 512,
        width: int = 512,
        noise_blur: float = 10.0,
        eta: float = 0.0,
        output_type: Literal["pil", "np"] = "pil",
        return_dict: bool = True,
        step_callback: Callable[
            [int, float, Float[torch.Tensor, "batch channels height width"]], None
        ]
        | None = None,
    ) -> (
        StableDiffusionPipelineOutput
        | tuple[
            list[Image.Image] | Float[np.ndarray, "batch height width channels"],
            list[bool] | None,
        ]
    ):
        """Generate a harmonized collage from an ordered set of RGBA layers.
        Args:
            prompt: Global scene description containing every layer description.
            layers: Back-to-front RGBA layers with descriptions and controls.
            negative_prompt: Optional text to suppress during classifier-free guidance.
            strength: Global image-to-image strength in ``(0, 1]``.
            num_inference_steps: Number of scheduler steps before strength adjustment.
            guidance_scale: Classifier-free guidance scale.
            num_images_per_prompt: Number of images to generate for the prompt.
            generator: Explicit random generator; takes precedence over ``seed``.
            seed: Optional local seed used when ``generator`` is absent.
            textual_inversion_embeddings: Map from layer token to caller-provided weights.
            height: Output canvas height; must be divisible by the VAE scale factor.
            width: Output canvas width; must be divisible by the VAE scale factor.
            noise_blur: Gaussian blur radius for the blended per-layer noise map.
            eta: Optional scheduler noise parameter used by DDIM-compatible schedulers.
            output_type: ``"pil"`` or ``"np"`` image output format.
            return_dict: Whether to return ``StableDiffusionPipelineOutput``.
            step_callback: Optional callback receiving step index, timestep, and latents.
        Returns:
            A Diffusers image output or its standard tuple representation.
        Raises:
            TypeError: If a layer sequence contains an unsupported value.
            ValueError: If the scheduler does not expose sigmas or generation inputs are invalid.
        """
        # Add layer-aware attention guidance and restore source regions during denoising.
        # 1. Validate scheduler, prompts, layers, and generation settings.
        if getattr(self.scheduler, "sigmas", None) is None:
            raise ValueError(
                "scheduler must expose `sigmas` for cross-attention guidance; "
                "use EulerAncestralDiscreteScheduler"
            )

        if not prompt.strip():
            raise ValueError("prompt must be non-empty")

        if not layers:
            raise ValueError("layers must contain at least one layer")

        if any(not isinstance(layer, CollageLayer) for layer in layers):
            raise TypeError("layers must contain CollageLayer values")

        if not math.isfinite(strength) or not 0.0 < strength <= 1.0:
            raise ValueError("strength must be finite and between 0 and 1")

        if num_inference_steps < 1 or num_images_per_prompt < 1:
            raise ValueError(
                "num_inference_steps and num_images_per_prompt must be positive"
            )

        if not math.isfinite(guidance_scale) or guidance_scale < 0.0:
            raise ValueError("guidance_scale must be finite and non-negative")

        if height < 1 or width < 1:
            raise ValueError("height and width must be positive")

        if output_type not in {"pil", "np"}:
            raise ValueError("output_type must be 'pil' or 'np'")

        device = self._execution_device
        vae_scale_factor = self.vae_scale_factor
        if height % vae_scale_factor or width % vae_scale_factor:
            raise ValueError(
                "height and width must be divisible by the VAE scale factor"
            )

        # 2. Prepare layer prompts, collage image, visible masks, and noise map.
        generator = resolve_torch_generator(generator=generator, seed=seed)
        prompt = _inject_textual_inversion_tokens(
            prompt, layers, textual_inversion_embeddings
        )
        if textual_inversion_embeddings is not None:
            for token, embedding_source in textual_inversion_embeddings.items():
                self.load_textual_inversion(embedding_source, token=token)

        token_spans = _prompt_token_spans(self.tokenizer, prompt, layers)
        collage, visible_masks = compose_collage_layers(
            layers, width=width, height=height
        )
        noise_map = make_layer_noise_map(
            visible_masks, layers, strength=strength, blur_radius=noise_blur
        )
        image = self.image_processor.preprocess(collage, height=height, width=width)
        image = image.to(device=device, dtype=self.unet.dtype)

        # 3. Encode the prompt and batch both guidance branches when guidance_scale > 1.
        do_classifier_free_guidance = guidance_scale > 1.0
        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            do_classifier_free_guidance,
            negative_prompt=negative_prompt,
        )
        if do_classifier_free_guidance:
            if negative_prompt_embeds is None:
                raise ValueError(
                    "classifier-free guidance requires negative prompt embeddings"
                )

            # Concatenate both branches so the U-Net handles them in one call.
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])

        # 4. Encode the collage and prepare its initial latent and noise.
        encoded = self.vae.encode(image)
        init_latents = encoded.latent_dist.sample(generator=generator)
        init_latents = init_latents * self.vae.config.scaling_factor
        init_latents = init_latents.repeat(num_images_per_prompt, 1, 1, 1)
        initial_noise = randn(
            init_latents.shape,
            generator=generator,
            device=device,
            dtype=init_latents.dtype,
        )

        # 5. Set img2img timesteps and build the latent-resolution preservation mask.
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        scheduler_offset = int(self.scheduler.config.get("steps_offset", 0))
        init_timestep = min(
            int(num_inference_steps * strength) + scheduler_offset,
            num_inference_steps,
        )
        timestep_start = max(
            num_inference_steps - init_timestep + scheduler_offset,
            0,
        )
        timesteps = self.scheduler.timesteps[timestep_start:]
        effective_steps = num_inference_steps - timestep_start
        if len(timesteps) == 0:
            raise ValueError(
                "strength and num_inference_steps produced no denoising steps"
            )

        latent_timestep = timesteps[:1].repeat(num_images_per_prompt)
        latents = self.scheduler.add_noise(init_latents, initial_noise, latent_timestep)
        noise_mask_pixels = np.rint(noise_map * 255.0).astype(np.uint8)
        # Quantize the source strength map before resizing it to the latent grid.
        noise_mask_image = Image.fromarray(noise_mask_pixels)
        noise_mask_image = noise_mask_image.resize(
            (latents.shape[-1], latents.shape[-2]),
            resample=Image.Resampling.NEAREST,
        )
        noise_mask_values = np.asarray(noise_mask_image, dtype=np.float32)
        noise_mask_values = noise_mask_values.copy()
        normalized_noise_mask = noise_mask_values / 255.0
        preservation_mask = 1.0 - normalized_noise_mask
        preservation_mask_tensor = torch.as_tensor(
            preservation_mask, device=device, dtype=latents.dtype
        )
        batched_preservation_mask = preservation_mask_tensor.unsqueeze(0).unsqueeze(0)
        noise_mask = batched_preservation_mask.expand(
            num_images_per_prompt, 1, *latents.shape[-2:]
        )

        # 6. Install layer-aware attention processors and prepare scheduler arguments.
        previous_processors = dict(self.unet.attn_processors)
        controlled_processors = {
            name: _CollageCrossAttentionProcessor(
                visible_masks=visible_masks,
                token_spans=token_spans,
                layers=layers,
                output_size=(height, width),
                batch_size=num_images_per_prompt,
                classifier_free_guidance=do_classifier_free_guidance,
            )
            if name.endswith(".attn2.processor")
            else processor
            for name, processor in previous_processors.items()
        }
        controlled_attention_processors = tuple(
            processor
            for processor in controlled_processors.values()
            if isinstance(processor, _CollageCrossAttentionProcessor)
        )
        attention_control_enabled = any(
            layer.positive_attention_weight or layer.negative_attention_weight
            for layer in layers
        )
        if attention_control_enabled and not controlled_attention_processors:
            raise ValueError("the U-Net exposes no cross-attention processors")

        if attention_control_enabled:
            self.unet.set_attn_processor(controlled_processors)

        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
        self._num_timesteps = len(timesteps)
        try:
            # 7. Denoise, applying attention guidance and layer-specific preservation.
            with self.progress_bar(total=effective_steps) as progress_bar:
                for step_index, timestep in enumerate(timesteps):
                    if attention_control_enabled:
                        sigma = float(self.scheduler.sigmas[step_index].item())
                        for processor in controlled_attention_processors:
                            processor.sigma = sigma

                    latent_model_input = (
                        torch.cat([latents] * 2)
                        if do_classifier_free_guidance
                        else latents
                    )
                    latent_model_input = self.scheduler.scale_model_input(
                        latent_model_input, timestep
                    )
                    noise_prediction = self.unet(
                        latent_model_input,
                        timestep,
                        encoder_hidden_states=prompt_embeds,
                    ).sample
                    if do_classifier_free_guidance:
                        # Standard CFG: unconditional + scale * (conditional - unconditional).
                        unconditional, conditional = noise_prediction.chunk(2)
                        noise_prediction = unconditional + guidance_scale * (
                            conditional - unconditional
                        )

                    # Advance x_t to x_(t-1); the mask below restores source regions.
                    latents = self.scheduler.step(
                        noise_prediction,
                        timestep,
                        latents,
                        **extra_step_kwargs,
                    ).prev_sample
                    step_noise = randn(
                        latents.shape,
                        generator=generator,
                        device=device,
                        dtype=latents.dtype,
                    )
                    timestep_batch = timestep.expand(latents.shape[0])
                    preserved_latents = self.scheduler.add_noise(
                        init_latents, step_noise, timestep_batch
                    )
                    timestep_value = float(timestep.item())
                    # Apply the timestep-dependent preservation threshold and spread it to neighbors.
                    mask_t = (noise_mask > 1.0 - timestep_value / 1000.0).to(
                        latents.dtype
                    )
                    dilation = int(4.0 * timestep_value / 1000.0)
                    if dilation > 0:
                        mask_t = F.max_pool2d(
                            mask_t,
                            kernel_size=dilation * 2 + 1,
                            stride=1,
                            padding=dilation,
                        )

                    latents = preserved_latents * mask_t + latents * (1.0 - mask_t)
                    if step_callback is not None:
                        step_callback(step_index, timestep_value, latents)

                    progress_bar.update()
        finally:
            if attention_control_enabled:
                self.unet.set_attn_processor(previous_processors)

        # 8. Decode the final latent and return the standard Diffusers output.
        decoded = self.vae.decode(
            latents / self.vae.config.scaling_factor, return_dict=False
        )[0]
        decoded, has_nsfw_concept = self.run_safety_checker(
            decoded, device, prompt_embeds.dtype
        )
        do_denormalize = (
            [not flag for flag in has_nsfw_concept]
            if has_nsfw_concept is not None
            else None
        )
        images = self.image_processor.postprocess(
            decoded, output_type=output_type, do_denormalize=do_denormalize
        )
        self.maybe_free_model_hooks()

        if not return_dict:
            return images, has_nsfw_concept

        return StableDiffusionPipelineOutput(
            images=images,
            nsfw_content_detected=has_nsfw_concept,
        )