Blame for pipeline.py

"""Diffusers pipeline for training-free structured diffusion guidance."""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from enum import StrEnum, auto
from typing import Final, Literal, cast

import numpy as np
import torch
from diffusers import StableDiffusionPipeline
from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
from diffusers.image_processor import PipelineImageInput
from diffusers.pipelines.stable_diffusion.pipeline_output import (
    StableDiffusionPipelineOutput,
)
from jaxtyping import Float
from PIL import Image

from .modeling_structure_diffusion import (
    StructuredAttentionProcessor,
    StructuredConditioning,
)
from .processing_structure_diffusion import (
    NounPhrase,
    align_phrase_context,
    constituency_noun_phrases,
    normalize_prompt,
    scene_graph_noun_phrases,
    token_span_from_char_span,
)


class ParserType(StrEnum):
    """Supported linguistic parsers for extracting prompt noun phrases."""

    constituency = auto()
    scene_graph = auto()


_PARSER_TYPES: Final[frozenset[ParserType]] = frozenset(ParserType)


def normalize_parser_type(parser_type: ParserType | str) -> ParserType:
    """Normalize a public parser name and reject unsupported parser choices."""
    if isinstance(parser_type, ParserType):
        normalized = parser_type
    else:
        try:
            normalized = ParserType(parser_type.lower().replace("-", "_"))
        except ValueError as exc:
            choices = ", ".join(parser.value for parser in ParserType)
            raise ValueError(
                f"Unsupported parser_type {parser_type!r}; choose one of: {choices}."
            ) from exc

    if normalized not in _PARSER_TYPES:
        raise ValueError(f"Unsupported parser_type: {parser_type}")

    return normalized


@dataclass(frozen=True)
class _StructuredCall:
    prompts: tuple[str, ...]
    parser_type: ParserType
    conjunction: bool
    save_attention_maps: bool


OutputType = Literal["pil", "np", "latent"]


class StructureDiffusionPipeline(StableDiffusionPipeline):
    """Generate images with noun-phrase-aware cross-attention.

    The pipeline retains Stable Diffusion's component loading and sampling API.
    Its three method options control noun-phrase extraction, the paper's
    conjunction key/value arrangement, and optional attention-map collection.

    Args:
        vae: Latent image decoder.
        text_encoder: CLIP text encoder.
        tokenizer: CLIP tokenizer.
        unet: Stable Diffusion denoising network.
        scheduler: Diffusion scheduler.
        safety_checker: Optional output safety checker.
        feature_extractor: Optional safety-checker image processor.
        image_encoder: Optional image encoder.
        requires_safety_checker: Whether the safety checker is required.

    Examples:
        >>> from structure_diffusion import StructureDiffusionPipeline
        >>> pipe = StructureDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
        >>> pipe("A red car and a white sheep", num_inference_steps=10).images[0]
    """

    attention_maps: tuple[tuple[str, Float[torch.Tensor, "..."]], ...] = ()

    def encode_prompt(
        self,
        prompt: str | list[str] | None,
        device: torch.device,
        num_images_per_prompt: int,
        do_classifier_free_guidance: bool,
        negative_prompt: str | list[str] | None = None,
        prompt_embeds: Float[torch.Tensor, "batch tokens dim"] | None = None,
        negative_prompt_embeds: Float[torch.Tensor, "batch tokens dim"] | None = None,
        lora_scale: float | None = None,
        clip_skip: int | None = None,
    ) -> tuple[
        Float[torch.Tensor, "batch tokens dim"],
        Float[torch.Tensor, "batch tokens dim"] | None,
    ]:
        """Encode standard prompts and attach parsed key/value contexts per call."""
        prompt_embeds, negative_prompt_embeds = super().encode_prompt(
            prompt=prompt,
            device=device,
            num_images_per_prompt=num_images_per_prompt,
            do_classifier_free_guidance=do_classifier_free_guidance,
            negative_prompt=negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            lora_scale=lora_scale,
            clip_skip=clip_skip,
        )

        request = getattr(self, "_call_request", None)
        if request is not None:
            self._install_conditioning(
                request,
                prompt_embeds,
                device,
                num_images_per_prompt,
                do_classifier_free_guidance,
                lora_scale,
                clip_skip,
            )

        return prompt_embeds, negative_prompt_embeds

    @torch.no_grad()
    def __call__(
        self,
        prompt: str | list[str],
        parser_type: ParserType | str = ParserType.constituency,
        conjunction: bool = False,
        save_attention_maps: bool = False,
        height: int | None = None,
        width: int | None = None,
        num_inference_steps: int = 50,
        timesteps: list[int] | None = None,
        sigmas: list[float] | None = None,
        guidance_scale: float = 7.5,
        negative_prompt: str | list[str] | None = None,
        num_images_per_prompt: int = 1,
        eta: float = 0.0,
        generator: torch.Generator | list[torch.Generator] | None = None,
        latents: Float[torch.Tensor, "batch channels latent_height latent_width"]
        | None = None,
        prompt_embeds: Float[torch.Tensor, "batch tokens dim"] | None = None,
        negative_prompt_embeds: Float[torch.Tensor, "batch tokens dim"] | None = None,
        ip_adapter_image: PipelineImageInput | None = None,
        ip_adapter_image_embeds: list[Float[torch.Tensor, "..."]] | None = None,
        output_type: OutputType = "pil",
        return_dict: bool = True,
        cross_attention_kwargs: dict[str, float | bool] | None = None,
        guidance_rescale: float = 0.0,
        clip_skip: int | None = None,
        callback_on_step_end: (
            PipelineCallback
            | MultiPipelineCallbacks
            | Callable[
                [
                    StableDiffusionPipeline,
                    int,
                    int,
                    dict[str, Float[torch.Tensor, "..."]],
                ],
                dict[str, Float[torch.Tensor, "..."]],
            ]
            | None
        ) = None,
        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
    ) -> (
        StableDiffusionPipelineOutput
        | tuple[
            Float[torch.Tensor, "batch channels height width"]
            | Float[np.ndarray, "batch height width channels"]
            | list[Image.Image],
            list[bool] | None,
        ]
    ):
        """Generate one or more images using structured cross-attention.

        Args:
            prompt: Prompt text or prompt batch to parse and render.
            parser_type: Constituency or scene-graph noun-phrase parser.
            conjunction: Use separate aligned noun phrases for both keys and values.
            save_attention_maps: Keep conditional attention maps on ``attention_maps``.
            height: Requested output image height.
            width: Requested output image width.
            num_inference_steps: Number of scheduler denoising steps.
            timesteps: Optional explicit scheduler timesteps.
            sigmas: Optional explicit scheduler noise levels.
            guidance_scale: Classifier-free guidance scale.
            negative_prompt: Optional negative prompt or prompt batch.
            num_images_per_prompt: Number of outputs to generate for each prompt.
            eta: Scheduler stochasticity parameter when supported.
            generator: Optional random generator or per-output generator list.
            latents: Optional initial latent tensor.
            prompt_embeds: Optional full-prompt CLIP embeddings.
            negative_prompt_embeds: Optional negative-prompt embeddings.
            ip_adapter_image: Optional IP-Adapter image input.
            ip_adapter_image_embeds: Optional precomputed IP-Adapter embeddings.
            output_type: Diffusers output format.
            return_dict: Whether to return ``StableDiffusionPipelineOutput``.
            cross_attention_kwargs: Optional Diffusers cross-attention settings.
            guidance_rescale: Optional classifier-free guidance rescaling.
            clip_skip: Number of CLIP layers to skip.
            callback_on_step_end: Optional Diffusers step callback.
            callback_on_step_end_tensor_inputs: Tensors exposed to the callback.

        Returns:
            The standard Stable Diffusion output container or tuple.

        Raises:
            ValueError: If the parser or prompt batch is invalid.
        """
        # 1. Normalize and validate the structured prompt inputs.
        normalized_parser = normalize_parser_type(parser_type)
        prompt_items = [prompt] if isinstance(prompt, str) else prompt
        normalized_prompts = tuple(normalize_prompt(item) for item in prompt_items)
        if not normalized_prompts or any(not item for item in normalized_prompts):
            raise ValueError("prompt must contain at least one non-empty string.")

        if num_images_per_prompt < 1:
            raise ValueError("num_images_per_prompt must be positive.")

        # 2. Prepare the per-call parser and attention-map state.
        request = _StructuredCall(
            prompts=normalized_prompts,
            parser_type=normalized_parser,
            conjunction=conjunction,
            save_attention_maps=save_attention_maps,
        )
        self.__dict__["attention_maps"] = ()
        previous_processors = cast(
            dict[str, Callable[..., Float[torch.Tensor, "..."]]],
            dict(self.unet.attn_processors),
        )
        active_processors: list[StructuredAttentionProcessor] = []
        replacement_processors: dict[
            str, Callable[..., Float[torch.Tensor, "..."]]
        ] = {}
        for name, processor in previous_processors.items():
            if name.endswith("attn2.processor"):
                structured_processor = StructuredAttentionProcessor(
                    name,
                    save_attention_maps,
                )
                active_processors.append(structured_processor)
                replacement_processors[name] = structured_processor
            else:
                replacement_processors[name] = processor

        if not active_processors:
            raise RuntimeError(
                "The U-Net has no cross-attention processors to replace."
            )

        # 3. Replace only cross-attention processors; self-attention stays unchanged.
        self.unet.set_attn_processor(replacement_processors)
        self._call_request = request
        self._processors_for_call = active_processors
        try:
            # 4. Run the inherited text-to-image denoising and decoding stages.
            result = super().__call__(
                prompt=(
                    normalized_prompts[0]
                    if isinstance(prompt, str)
                    else list(normalized_prompts)
                ),
                height=height,
                width=width,
                num_inference_steps=num_inference_steps,
                timesteps=timesteps,
                sigmas=sigmas,
                guidance_scale=guidance_scale,
                negative_prompt=negative_prompt,
                num_images_per_prompt=num_images_per_prompt,
                eta=eta,
                generator=generator,
                latents=latents,
                prompt_embeds=prompt_embeds,
                negative_prompt_embeds=negative_prompt_embeds,
                ip_adapter_image=ip_adapter_image,
                ip_adapter_image_embeds=ip_adapter_image_embeds,
                output_type=output_type,
                return_dict=return_dict,
                cross_attention_kwargs=cross_attention_kwargs,
                guidance_rescale=guidance_rescale,
                clip_skip=clip_skip,
                callback_on_step_end=callback_on_step_end,
                callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
            )
            self.__dict__["attention_maps"] = tuple(
                (processor.layer_name, attention_map)
                for processor in active_processors
                for attention_map in processor.attention_maps
            )
            return result
        finally:
            # 5. Restore the U-Net processors even if generation raises an error.
            del self._call_request
            del self._processors_for_call
            self.unet.set_attn_processor(dict(previous_processors))

    def _install_conditioning(
        self,
        request: _StructuredCall,
        prompt_embeds: Float[torch.Tensor, "batch tokens dim"],
        device: torch.device,
        num_images_per_prompt: int,
        do_classifier_free_guidance: bool,
        lora_scale: float | None,
        clip_skip: int | None,
    ) -> None:
        batch_size = len(request.prompts)
        expected_batch = batch_size * num_images_per_prompt
        if prompt_embeds.shape[0] != expected_batch:
            raise ValueError(
                "The encoded prompt batch does not match prompt and num_images_per_prompt."
            )

        prompt_contexts: Float[torch.Tensor, "prompts images tokens dim"] = (
            prompt_embeds.reshape(
                batch_size,
                num_images_per_prompt,
                prompt_embeds.shape[1],
                prompt_embeds.shape[2],
            )
        )
        base_contexts: Float[torch.Tensor, "batch tokens dim"] = prompt_contexts[:, 0]
        phrase_groups = [
            self._parse_prompt(prompt_text, request.parser_type)
            for prompt_text in request.prompts
        ]
        # Encode each parsed phrase separately so its CLIP tokens can be aligned below.
        phrase_texts = [phrase.text for group in phrase_groups for phrase in group]
        phrase_contexts: Float[torch.Tensor, "phrases tokens dim"] | None = None
        if phrase_texts:
            phrase_contexts, _ = StableDiffusionPipeline.encode_prompt(
                self,
                prompt=phrase_texts,
                device=device,
                num_images_per_prompt=1,
                do_classifier_free_guidance=False,
                lora_scale=lora_scale,
                clip_skip=clip_skip,
            )

        aligned_groups: list[
            list[
                tuple[
                    Float[torch.Tensor, "batch tokens dim"],
                    tuple[int, int],
                ]
            ]
        ] = []
        phrase_cursor = 0
        for prompt_index, phrase_group in enumerate(phrase_groups):
            aligned: list[
                tuple[
                    Float[torch.Tensor, "batch tokens dim"],
                    tuple[int, int],
                ]
            ] = []
            for phrase in phrase_group:
                phrase_embedding = (
                    phrase_contexts[phrase_cursor : phrase_cursor + 1]
                    if phrase_contexts is not None
                    else base_contexts[prompt_index : prompt_index + 1]
                )
                phrase_cursor += 1
                token_span = token_span_from_char_span(
                    self.tokenizer,
                    request.prompts[prompt_index],
                    phrase.start,
                    phrase.end,
                    prompt_embeds.shape[1],
                )
                if token_span is None:
                    continue

                phrase_token_count = len(self.tokenizer.tokenize(phrase.text))
                if phrase_token_count < token_span[1] - token_span[0]:
                    continue

                aligned.append(
                    (
                        align_phrase_context(
                            base_contexts[prompt_index : prompt_index + 1],
                            phrase_embedding,
                            token_span,
                            phrase_token_count,
                        ),
                        token_span,
                    )
                )

            aligned_groups.append(aligned)

        max_phrase_count = max((len(group) for group in aligned_groups), default=0)
        aligned_context_batches: list[Float[torch.Tensor, "batch tokens dim"]] = []
        for phrase_index in range(max_phrase_count):
            phrase_contexts_for_batch: list[
                Float[torch.Tensor, "batch tokens dim"]
            ] = []
            for prompt_index, group in enumerate(aligned_groups):
                if phrase_index < len(group):
                    aligned_context, _token_span = group[phrase_index]
                    prompt_phrase_context = aligned_context
                else:
                    prompt_phrase_context = base_contexts[
                        prompt_index : prompt_index + 1
                    ]

                phrase_contexts_for_batch.append(prompt_phrase_context)

            phrase_context_batch = torch.cat(phrase_contexts_for_batch, dim=0)
            aligned_context_batches.append(phrase_context_batch)

        aligned_contexts = tuple(aligned_context_batches)
        aligned_validity = tuple(
            torch.tensor(
                [index < len(group) for group in aligned_groups],
                dtype=torch.bool,
                device=device,
            )
            for index in range(max_phrase_count)
        )
        base_validity = torch.ones(batch_size, dtype=torch.bool, device=device)
        # Conjunction mode uses phrase keys with one composed value; the default uses
        # one full-prompt key with separate phrase values.
        if request.conjunction:
            combined_contexts: list[Float[torch.Tensor, "batch tokens dim"]] = []
            for prompt_index, base_context in enumerate(base_contexts):
                combined = base_context.clone()
                for aligned_context, (start, end) in aligned_groups[prompt_index]:
                    combined[start:end] = aligned_context[0, start:end]

                combined_contexts.append(combined)

            values = (torch.stack(combined_contexts),)
            keys = (base_contexts, *aligned_contexts)
            key_validity = (base_validity, *aligned_validity)
            value_validity = (base_validity,)
        else:
            keys = (base_contexts,)
            values = (base_contexts, *aligned_contexts)
            key_validity = (base_validity,)
            value_validity = (base_validity, *aligned_validity)

        repeated_keys = tuple(
            context.repeat_interleave(num_images_per_prompt, dim=0) for context in keys
        )
        repeated_values = tuple(
            context.repeat_interleave(num_images_per_prompt, dim=0)
            for context in values
        )
        repeated_key_validity = tuple(
            validity.repeat_interleave(num_images_per_prompt, dim=0)
            for validity in key_validity
        )
        repeated_value_validity = tuple(
            validity.repeat_interleave(num_images_per_prompt, dim=0)
            for validity in value_validity
        )
        conditioning = StructuredConditioning(
            keys=repeated_keys,
            values=repeated_values,
            key_validity=repeated_key_validity,
            value_validity=repeated_value_validity,
            do_classifier_free_guidance=do_classifier_free_guidance,
        )
        for processor in getattr(self, "_processors_for_call", ()):
            processor.conditioning = conditioning

    @staticmethod
    def _parse_prompt(prompt: str, parser_type: ParserType) -> list[NounPhrase]:
        if parser_type is ParserType.constituency:
            return constituency_noun_phrases(prompt)

        if parser_type is ParserType.scene_graph:
            return scene_graph_noun_phrases(prompt)

        raise ValueError(f"Unsupported parser_type: {parser_type}")