Publish Structure Diffusion pipeline at 647bfc6e796a

README.md

@@ -0,0 +1,318 @@
+---
+language:
+  - en
+license: apache-2.0
+library_name: diffusers
+pipeline_tag: text-to-image
+tags:
+  - structured-diffusion-guidance
+  - text-to-image
+  - image-generation
+datasets: []
+base_model: CompVis/stable-diffusion-v1-4
+model-index:
+  - name: "StructureDiffusion"
+    results:
+      - task:
+          type: "text-to-image"
+          name: "Text-to-image generation"
+        dataset:
+          type: "synthetic"
+          name: "CPU parity fixture"
+          split: "one prompt at 64x64 and 512x512"
+        metrics:
+          - type: "vendor-parity"
+            value: 0
+            name: "Bit-exact CPU image comparison against reference images"
+---
+
+<!-- --8<-- [start:card] -->
+
+# Model Card for StructureDiffusion
+
+[![arXiv](https://img.shields.io/static/v1?label=arXiv&message=2212.05032&color=b31b1b&style=flat-square&logo=arxiv&logoColor=white)](https://arxiv.org/abs/2212.05032)
+![venue](https://img.shields.io/static/v1?label=venue&message=ICLR+2023&color=purple&style=flat-square)
+![license](https://img.shields.io/static/v1?label=license&message=Apache-2.0&color=green&style=flat-square&logo=apache&logoColor=white)
+![base](https://img.shields.io/static/v1?label=base&message=diffusers&color=blue&style=flat-square&logo=huggingface&logoColor=white)
+![task](https://img.shields.io/static/v1?label=task&message=text-to-image&color=purple&style=flat-square)
+![vendor-parity](https://img.shields.io/static/v1?label=vendor-parity&message=bit-exact&color=success&style=flat-square)
+![hub](https://img.shields.io/static/v1?label=hub&message=published-code-only&color=brightgreen&style=flat-square&logo=huggingface&logoColor=white)
+
+This package implements [StructureDiffusion](https://arxiv.org/abs/2212.05032), a training-free text-to-image method from ICLR 2023. It parses a prompt into noun phrases, encodes the full prompt and each phrase, then combines cross-attention keys and values to improve object and attribute binding. The [`🧨diffusers`](https://huggingface.co/docs/diffusers/index) pipeline uses the [Stable Diffusion 1.4 checkpoint](https://huggingface.co/CompVis/stable-diffusion-v1-4).
+
+## Model Details
+
+### Model Description
+
+`StructureDiffusionPipeline` subclasses [`diffusers.StableDiffusionPipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img), so it uses the framework's standard loading, device placement, callbacks, scheduler, and image output behavior. `parser_type` selects constituency parsing or optional scene-graph parsing. `conjunction=True` uses phrase-specific keys with one combined value; the default uses the full prompt as the key and the full prompt plus phrase encodings as values.
+
+- **Developed by:** Weixi Feng et al.
+- **Shared by:** creative-graphic-design.
+- **Model type:** image generation; task: text-to-image; conditioning: text.
+- **Language(s):** English prompts.
+- **Code license:** Apache-2.0.
+- **Reference implementation license:** MIT.
+- **Base checkpoint:** CompVis Stable Diffusion v1.4; its weights have separate CreativeML OpenRAIL-M terms.
+
+### Model Sources
+
+- **Repository:** [official Structured Diffusion Guidance implementation](https://github.com/weixi-feng/Structured-Diffusion-Guidance/tree/773b2ba068c6cc60443d5d747a725413c8c9d46f)
+- **Paper:** [arXiv:2212.05032](https://arxiv.org/abs/2212.05032)
+- **Released base weights:** [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4)
+
+## Supported Checkpoints
+
+| Checkpoint                           | Hub ID                                        | Status                                                                         |
+| ------------------------------------ | --------------------------------------------- | ------------------------------------------------------------------------------ |
+| Stable Diffusion v1.4                | `CompVis/stable-diffusion-v1-4`               | Gated base checkpoint; load with `from_pretrained` after access is approved.   |
+| Structure Diffusion package artifact | `creative-graphic-design/structure-diffusion` | Published code only; this package does not upload Structure Diffusion weights. |
+
+## Uses
+
+### Direct Use
+
+Use the pipeline for research inference on English prompts that describe multiple objects, attributes, or relations. The method changes cross-attention conditioning and does not train or fine-tune Stable Diffusion weights.
+
+### Downstream Use
+
+Use generated images as draft visual material for creative-design workflows after inspecting them for prompt adherence and content suitability.
+
+### Out-of-Scope Use
+
+Do not treat generated images as verified depictions of real people or events, or use them as evidence that prompt attributes were rendered correctly. The pipeline does not provide an image safety guarantee.
+
+## Bias, Risks, and Limitations
+
+Generation behavior inherits the training data, biases, and restrictions of the selected Stable Diffusion checkpoint. The default constituency parser is English-only and requires the [Stanza](https://stanfordnlp.github.io/stanza/) English tokenizer, POS, and constituency models to be installed separately. Scene-graph parsing is an optional extra with its own [spaCy](https://spacy.io/) English model prerequisite. Outputs vary with the checkpoint, parser, scheduler, device, precision, and library versions.
+
+### Recommendations
+
+Inspect generated images before downstream use, and record the checkpoint revision, parser, scheduler, dimensions, seed, and inference steps when comparing results.
+
+## How to Get Started with the Model
+
+Install the package from this repository archive and download the English parser models before the first inference:
+
+```bash
+pip install "structure-diffusion @ https://github.com/creative-graphic-design/design-generators/archive/647bfc6e796a36794bb7141047dfc6c1d930679a.zip#subdirectory=models/structure-diffusion"
+python -c "import stanza; stanza.download('en', processors='tokenize,pos,constituency')"
+```
+
+[`CompVis/stable-diffusion-v1-4`](https://huggingface.co/CompVis/stable-diffusion-v1-4) is a gated model. Request access and accept its CreativeML OpenRAIL-M terms on the model page before downloading its weights. Set `HF_TOKEN_FILE` to a private file containing a token that can access the checkpoint.
+
+```python
+import os
+from pathlib import Path
+
+import torch
+from diffusers import DiffusionPipeline
+
+token = Path(os.environ["HF_TOKEN_FILE"]).read_text().strip()
+pipe = DiffusionPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-4",
+    revision="133a221b8aa7292a167afc5127cb63fb5005638b",
+    custom_pipeline="creative-graphic-design/structure-diffusion",
+    torch_dtype=torch.float32,
+    token=token,
+    trust_remote_code=True,
+)
+pipe.to("cpu")
+result = pipe(
+    "A red car and a white sheep",
+    parser_type="constituency",
+    height=512,
+    width=512,
+    num_inference_steps=10,
+    generator=torch.Generator(device="cpu").manual_seed(42),
+)
+image = result.images[0]
+print("image size: {}x{}".format(image.size[0], image.size[1]))
+```
+
+```text
+image size: 512x512
+```
+
+The default constituency parser needs the English Stanza model assets downloaded separately with the command above. The optional scene-graph parser is available through the scene-graph extra and requires its spaCy English model. See [REPRODUCING.md](https://github.com/creative-graphic-design/design-generators/blob/main/models/structure-diffusion/REPRODUCING.md) for local testing and CPU comparison instructions.
+
+## Examples
+
+The gallery compares `diffusers.StableDiffusionPipeline` with `StructureDiffusionPipeline` using the same prompt and seed. It includes every output for seeds `42`, `0`, `1`, and `2` for each prompt. The results vary by seed and do not show a consistent compositional improvement from structured guidance. For the car and sheep prompt, the sheep is absent from both seed-42 images; at seed 0 the structured result shows only part of a sheep, at seed 1 the baseline shows a sheep while the structured result does not, and both show a sheep at seed 2. For the backpack and elephant prompt, the baseline shows both objects while the structured result omits the elephant at seed 42; both omit the backpack at seed 0 and the elephant at seed 1; at seed 2 both show the objects, with the elephant more fully visible in the structured result. For the fruit prompt, neither pipeline clearly shows a banana at seeds 42, 0, or 1; at seed 2 both show several apples and banana-like shapes. The prompts are taken from Figure 1 (car and sheep; backpack and elephant) and Figure 13 (fruit examples) of the [paper](https://openreview.net/pdf?id=PUIqjT4rzq7). All outputs are shown without selecting among seeds.
+
+### A red car and a white sheep
+
+Source: [Figure 1, attribute leakage](https://openreview.net/pdf?id=PUIqjT4rzq7).
+
+| Seed | Stable Diffusion v1.4                                                                                                                                                                                                              | StructureDiffusion                                                                                                                                                                                                            |
+| ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 42   | ![Stable Diffusion v1.4, seed 42, A red car and a white sheep](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/stable-diffusion-v1-4-red-car-white-sheep-seed42.png) | ![StructureDiffusion, seed 42, A red car and a white sheep](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/structure-diffusion-red-car-white-sheep-seed42.png) |
+| 0    | ![Stable Diffusion v1.4, seed 0, A red car and a white sheep](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/stable-diffusion-v1-4-red-car-white-sheep-seed0.png)   | ![StructureDiffusion, seed 0, A red car and a white sheep](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/structure-diffusion-red-car-white-sheep-seed0.png)   |
+| 1    | ![Stable Diffusion v1.4, seed 1, A red car and a white sheep](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/stable-diffusion-v1-4-red-car-white-sheep-seed1.png)   | ![StructureDiffusion, seed 1, A red car and a white sheep](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/structure-diffusion-red-car-white-sheep-seed1.png)   |
+| 2    | ![Stable Diffusion v1.4, seed 2, A red car and a white sheep](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/stable-diffusion-v1-4-red-car-white-sheep-seed2.png)   | ![StructureDiffusion, seed 2, A red car and a white sheep](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/structure-diffusion-red-car-white-sheep-seed2.png)   |
+
+### A blue backpack and a brown elephant
+
+Source: [Figure 1, missing objects](https://openreview.net/pdf?id=PUIqjT4rzq7).
+
+| Seed | Stable Diffusion v1.4                                                                                                                                                                                                                                | StructureDiffusion                                                                                                                                                                                                                              |
+| ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 42   | ![Stable Diffusion v1.4, seed 42, A blue backpack and a brown elephant](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/stable-diffusion-v1-4-blue-backpack-brown-elephant-seed42.png) | ![StructureDiffusion, seed 42, A blue backpack and a brown elephant](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/structure-diffusion-blue-backpack-brown-elephant-seed42.png) |
+| 0    | ![Stable Diffusion v1.4, seed 0, A blue backpack and a brown elephant](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/stable-diffusion-v1-4-blue-backpack-brown-elephant-seed0.png)   | ![StructureDiffusion, seed 0, A blue backpack and a brown elephant](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/structure-diffusion-blue-backpack-brown-elephant-seed0.png)   |
+| 1    | ![Stable Diffusion v1.4, seed 1, A blue backpack and a brown elephant](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/stable-diffusion-v1-4-blue-backpack-brown-elephant-seed1.png)   | ![StructureDiffusion, seed 1, A blue backpack and a brown elephant](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/structure-diffusion-blue-backpack-brown-elephant-seed1.png)   |
+| 2    | ![Stable Diffusion v1.4, seed 2, A blue backpack and a brown elephant](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/stable-diffusion-v1-4-blue-backpack-brown-elephant-seed2.png)   | ![StructureDiffusion, seed 2, A blue backpack and a brown elephant](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/structure-diffusion-blue-backpack-brown-elephant-seed2.png)   |
+
+### A green apple and a red banana
+
+Source: [Figure 13, CC-500 qualitative results](https://openreview.net/pdf?id=PUIqjT4rzq7).
+
+| Seed | Stable Diffusion v1.4                                                                                                                                                                                                                    | StructureDiffusion                                                                                                                                                                                                                  |
+| ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 42   | ![Stable Diffusion v1.4, seed 42, A green apple and a red banana](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/stable-diffusion-v1-4-green-apple-red-banana-seed42.png) | ![StructureDiffusion, seed 42, A green apple and a red banana](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/structure-diffusion-green-apple-red-banana-seed42.png) |
+| 0    | ![Stable Diffusion v1.4, seed 0, A green apple and a red banana](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/stable-diffusion-v1-4-green-apple-red-banana-seed0.png)   | ![StructureDiffusion, seed 0, A green apple and a red banana](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/structure-diffusion-green-apple-red-banana-seed0.png)   |
+| 1    | ![Stable Diffusion v1.4, seed 1, A green apple and a red banana](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/stable-diffusion-v1-4-green-apple-red-banana-seed1.png)   | ![StructureDiffusion, seed 1, A green apple and a red banana](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/structure-diffusion-green-apple-red-banana-seed1.png)   |
+| 2    | ![Stable Diffusion v1.4, seed 2, A green apple and a red banana](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/stable-diffusion-v1-4-green-apple-red-banana-seed2.png)   | ![StructureDiffusion, seed 2, A green apple and a red banana](https://huggingface.co/datasets/creative-graphic-design/documentation-images/resolve/main/structure-diffusion/structure-diffusion-green-apple-red-banana-seed2.png)   |
+
+Every image uses `CompVis/stable-diffusion-v1-4` at revision `133a221b8aa7292a167afc5127cb63fb5005638b`, the checkpoint's PNDM scheduler configuration (`skip_prk_steps=True`), 30 inference steps, 512×512 resolution, guidance scale `7.5`, and CPU float32. The structured pipeline uses the constituency parser with `conjunction=True`; the baseline uses `diffusers.StableDiffusionPipeline`. The base checkpoint weights are licensed under CreativeML OpenRAIL-M. The image files are hosted in the public documentation-images dataset and are not committed to this repository.
+
+This code recreates all displayed pairs with the same prompts and seeds. Set `HF_TOKEN_FILE` to a private file containing a Hugging Face token after accepting the checkpoint's access terms. Install the package and download the English Stanza parser models as shown above before running it.
+
+```python
+import gc
+import os
+from pathlib import Path
+
+import torch
+from diffusers import StableDiffusionPipeline
+from structure_diffusion import StructureDiffusionPipeline
+
+MODEL_ID = "CompVis/stable-diffusion-v1-4"
+MODEL_REVISION = "133a221b8aa7292a167afc5127cb63fb5005638b"
+OUTPUT_DIR = Path("structure-diffusion-examples")
+OUTPUT_DIR.mkdir(exist_ok=True)
+EXAMPLES = (
+    ("red-car-white-sheep", "A red car and a white sheep"),
+    ("blue-backpack-brown-elephant", "A blue backpack and a brown elephant"),
+    ("green-apple-red-banana", "A green apple and a red banana"),
+)
+SEEDS = (42, 0, 1, 2)
+
+torch.set_num_threads(2)
+token = Path(os.environ["HF_TOKEN_FILE"]).read_text().strip()
+
+
+def make_generator(seed):
+    return torch.Generator(device="cpu").manual_seed(seed)
+
+
+baseline = StableDiffusionPipeline.from_pretrained(
+    MODEL_ID,
+    revision=MODEL_REVISION,
+    torch_dtype=torch.float32,
+    token=token,
+)
+baseline.to("cpu")
+for slug, prompt in EXAMPLES:
+    for seed in SEEDS:
+        image = baseline(
+            prompt,
+            height=512,
+            width=512,
+            num_inference_steps=30,
+            guidance_scale=7.5,
+            generator=make_generator(seed),
+        ).images[0]
+        image.save(OUTPUT_DIR / f"stable-diffusion-v1-4-{slug}-seed{seed}.png")
+del baseline
+
+gc.collect()
+structured = StructureDiffusionPipeline.from_pretrained(
+    MODEL_ID,
+    revision=MODEL_REVISION,
+    torch_dtype=torch.float32,
+    token=token,
+)
+structured.to("cpu")
+for slug, prompt in EXAMPLES:
+    for seed in SEEDS:
+        image = structured(
+            prompt,
+            parser_type="constituency",
+            conjunction=True,
+            height=512,
+            width=512,
+            num_inference_steps=30,
+            guidance_scale=7.5,
+            generator=make_generator(seed),
+        ).images[0]
+        image.save(OUTPUT_DIR / f"structure-diffusion-{slug}-seed{seed}.png")
+```
+
+## Training Details
+
+### Training Data
+
+This method adds no training data or learned weights. The selected Stable Diffusion checkpoint retains its original data provenance and terms.
+
+### Training Procedure
+
+No training or checkpoint conversion is performed by this package. It loads the released pipeline-format base checkpoint and changes the denoising pipeline's text-conditioning path.
+
+## Evaluation
+
+### Testing Data, Factors & Metrics
+
+The independent CPU comparison covers 64×64 and 512×512 outputs for `A red car and a white sheep`, seed 42, 10 PNDM steps through the PLMS path in `diffusers`, guidance scale 7.5, and CPU float32 execution. Each reference/package pair receives the same precomputed starting latent. The vendor demo calls `preprocess_prompts` before parsing and token alignment, which lowercases the prompt; its `get_token_alignment_map` helper is case-sensitive only when called directly with mismatched prompt and tokenizer casing. The package also accepts capitalized input.
+
+The test directly compares vendor-derived noun phrases and token spans, and compares independently assembled key/value conditioning with the package conditioning. For the image comparison, a separate base `diffusers` pipeline uses the pinned vendor attention processors to generate the reference image; the test compares that final image with the `StructureDiffusionPipeline` output. It does not capture or compare per-layer attention outputs. Both model loading and sampling use `diffusers`, so this checks structured conditioning and its end-to-end image result without reproducing the vendor checkpoint loader or sampler implementation.
+
+### Parity Results
+
+| Check                                                | Cases | Criterion                                                                          | Result                                                                                                                                                                                         |
+| ---------------------------------------------------- | ----: | ---------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 64×64 CPU image output against the reference image   |     1 | Exact phrase and token-span match; K/V max difference ≤ `1e-6`; image MAE < `1e-3` | `a red car` / `a white sheep`; spans `[1, 4)` / `[5, 8)`; 1 key / 3 values; K/V max differences `0.0` / `0.0`; pixel max/mean/RMS differences `0.0` / `0.0` / `0.0`; bit-exact, SHA-256 match. |
+| 512×512 CPU image output against the reference image |     1 | Exact phrase and token-span match; K/V max difference ≤ `1e-6`; image MAE < `1e-3` | `a red car` / `a white sheep`; spans `[1, 4)` / `[5, 8)`; 1 key / 3 values; K/V max differences `0.0` / `0.0`; pixel max/mean/RMS differences `0.0` / `0.0` / `0.0`; bit-exact, SHA-256 match. |
+
+The independent comparison ran on commit `786fcf6cfbcfb68b9c9565d8ee275c32cd5a28ac`; the reviewer recorded the full inputs, initial latents, measurements, and runtimes in the [formal parity result](https://github.com/creative-graphic-design/design-generators/issues/485#issuecomment-6608854871). Both resolutions were bit-exact: the package and reference image SHA-256 values matched.
+
+## Reproducibility
+
+Run the CPU comparison test in [REPRODUCING.md](https://github.com/creative-graphic-design/design-generators/blob/main/models/structure-diffusion/REPRODUCING.md) to compare noun phrases, token spans, and K/V conditioning directly, then compare final generated images using the same scheduler, seed, and initial latent. A separate base `diffusers` pipeline runs the pinned reference attention processors to generate the reference image; the test does not compare per-layer attention outputs. `diffusers` supplies model loading and sampling in both paths, so the result does not claim parity with the reference checkpoint loader or sampler. The independently reviewed comparison on commit `786fcf6` matched phrase and span outputs exactly, had zero measured K/V differences, and produced bit-exact 64×64 and 512×512 images with identical SHA-256 values.
+
+## Environmental Impact
+
+No model training is performed. The formal standard-CPU comparison took 126.777 seconds at 64×64 and 948.902 seconds at 512×512; no GPU comparison is included.
+
+## Technical Specifications
+
+### Model Architecture and Objective
+
+The pipeline retains Stable Diffusion's VAE, CLIP text encoder, U-Net, and scheduler. It extracts phrase spans, encodes each phrase independently, aligns phrase tokens to the full prompt, and applies the official method's attention-map and value-composition rules in U-Net cross-attention layers.
+
+### Compute Infrastructure
+
+The package supports CPU and accelerator execution with `diffusers`. The recorded comparison timings are test-process runtimes on a standard Colab CPU; setup time and GPU performance are not included.
+
+#### Software
+
+Runtime dependencies are declared in `models/structure-diffusion/pyproject.toml`. The default parser requires Stanza models; scene-graph parsing also requires the optional `scene-graph` extra and its spaCy model.
+
+## Citation
+
+```bibtex
+@inproceedings{feng2023trainingfree,
+  title={Training-Free Structured Diffusion Guidance for Compositional Text-to-Image Synthesis},
+  author={Feng, Weixi and He, Xuehai and Fu, Tsu-Jui and Jampani, Varun and Akula, Arjun and Narayana, Pradyumna and Basu, Sugato and Wang, Xin Eric and Wang, William Yang},
+  booktitle={International Conference on Learning Representations},
+  year={2023},
+  eprint={2212.05032},
+  archivePrefix={arXiv},
+  primaryClass={cs.CV},
+  url = "https://arxiv.org/abs/2212.05032"
+}
+```
+
+## Model Card Contact
+
+Open an issue or pull request in the [design-generators repository](https://github.com/creative-graphic-design/design-generators).
+
+<!-- --8<-- [end:card] -->

modeling_structure_diffusion.py

@@ -0,0 +1,369 @@
+"""Cross-attention processor for structured diffusion guidance."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from diffusers.models.attention_processor import Attention
+from jaxtyping import Bool, Float
+
+
+@dataclass(frozen=True)
+class StructuredConditioning:
+    """Full-prompt keys, phrase-composed values, and per-example validity."""
+
+    keys: tuple[Float[torch.Tensor, "batch tokens dim"], ...]
+    values: tuple[Float[torch.Tensor, "batch tokens dim"], ...]
+    key_validity: tuple[Bool[torch.Tensor, "batch"], ...]
+    value_validity: tuple[Bool[torch.Tensor, "batch"], ...]
+    do_classifier_free_guidance: bool
+
+
+class StructuredAttentionProcessor:
+    """Apply noun-phrase key/value attention in a Diffusers cross-attention layer."""
+
+    def __init__(self, layer_name: str, save_attention_maps: bool) -> None:
+        """Create per-layer state for the requested attention-map behavior.
+
+        Args:
+            layer_name: U-Net attention processor name used to label saved maps.
+            save_attention_maps: Whether to retain conditional attention maps.
+        """
+        self.layer_name = layer_name
+        self.save_attention_maps = save_attention_maps
+        self.conditioning: StructuredConditioning | None = None
+        self.attention_maps: list[Float[torch.Tensor, "..."]] = []
+        self._last_probabilities: list[
+            Float[torch.Tensor, "batch_heads query tokens"]
+        ] = []
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: Float[torch.Tensor, "..."],
+        encoder_hidden_states: Float[torch.Tensor, "batch tokens dim"] | None = None,
+        attention_mask: Bool[torch.Tensor, "..."]
+        | Float[torch.Tensor, "..."]
+        | None = None,
+        temb: Float[torch.Tensor, "batch channels"] | None = None,
+        scale: float | None = None,
+    ) -> Float[torch.Tensor, "..."]:
+        """Run ordinary attention or noun-phrase-aware key/value attention."""
+        del scale
+        if (
+            attn.to_q is None
+            or attn.to_k is None
+            or attn.to_v is None
+            or attn.to_out is None
+        ):
+            raise ValueError(
+                "Attention layers must define query, key, value, and output projections."
+            )
+
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+        spatial_shape: tuple[int, int] | None = None
+        spatial_dimensions: tuple[int, int, int] | None = None
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            spatial_shape = (height, width)
+            spatial_dimensions = (channel, height, width)
+            spatial_tokens: Float[torch.Tensor, "batch channels tokens"] = (
+                hidden_states.view(batch_size, channel, height * width)
+            )
+            hidden_states = spatial_tokens.transpose(1, 2)
+
+        batch_size = hidden_states.shape[0]
+        if attn.group_norm is not None:
+            channels_first: Float[torch.Tensor, "batch channels tokens"] = (
+                hidden_states.transpose(1, 2)
+            )
+            normalized_channels = attn.group_norm(channels_first)
+            hidden_states = normalized_channels.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
+            )
+
+        conditioning = self.conditioning
+        if conditioning is None:
+            output = self._normal_attention(
+                attn,
+                query,
+                encoder_hidden_states,
+                attention_mask,
+            )
+        # CFG concatenates unconditional and conditional embeddings before this processor.
+        # Keep the unconditional half on ordinary attention and structure the conditional half.
+        elif conditioning.do_classifier_free_guidance:
+            base_batch = conditioning.keys[0].shape[0]
+            if batch_size != 2 * base_batch:
+                raise ValueError(
+                    "Classifier-free guidance queries must contain matching unconditional and conditional batches."
+                )
+
+            unconditional_query, conditional_query = query.split(base_batch, dim=0)
+            unconditional_context = encoder_hidden_states[:base_batch]
+            unconditional_mask, conditional_mask = self._split_mask(
+                attention_mask,
+                base_batch,
+                batch_size,
+            )
+            unconditional_output = self._normal_attention(
+                attn,
+                unconditional_query,
+                unconditional_context,
+                unconditional_mask,
+            )
+            conditional_output = self._structured_attention(
+                attn,
+                conditional_query,
+                conditional_mask,
+            )
+            output = torch.cat((unconditional_output, conditional_output), dim=0)
+        else:
+            if batch_size != conditioning.keys[0].shape[0]:
+                raise ValueError(
+                    "Structured conditioning and query batches must match."
+                )
+
+            output = self._structured_attention(attn, query, attention_mask)
+
+        if self.save_attention_maps:
+            if spatial_shape is None:
+                query_length = self._last_probabilities[0].shape[-2]
+                side_length = int(query_length**0.5)
+                if side_length * side_length == query_length:
+                    spatial_shape = (side_length, side_length)
+
+        if self.save_attention_maps and spatial_shape is not None:
+            height, width = spatial_shape
+            for probabilities in self._last_probabilities:
+                map_batch = probabilities.shape[0] // attn.heads
+                spatial_attention: Float[
+                    torch.Tensor, "batch heads height width tokens"
+                ] = probabilities.reshape(
+                    map_batch,
+                    attn.heads,
+                    height,
+                    width,
+                    probabilities.shape[-1],
+                )
+                saved_tokens = spatial_attention[..., :20]
+                saved_attention = saved_tokens.detach().cpu()
+                self.attention_maps.append(saved_attention)
+
+        output = attn.to_out[0](output)
+        output = attn.to_out[1](output)
+        if spatial_dimensions is not None:
+            channel, height, width = spatial_dimensions
+            spatial_tokens = output.transpose(-1, -2)
+            output = spatial_tokens.reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            output = output + residual
+
+        return output / attn.rescale_output_factor
+
+    def _normal_attention(
+        self,
+        attn: Attention,
+        query: Float[torch.Tensor, "batch sequence channels"],
+        context: Float[torch.Tensor, "batch tokens dim"],
+        attention_mask: Bool[torch.Tensor, "..."] | Float[torch.Tensor, "..."] | None,
+    ) -> Float[torch.Tensor, "batch sequence channels"]:
+        if attn.to_k is None or attn.to_v is None:
+            raise ValueError("Attention layers must define key and value projections.")
+
+        query_heads = self._to_heads(query, attn.heads)
+        key_projection = attn.to_k(context)
+        key_heads = self._to_heads(key_projection, attn.heads)
+        value_projection = attn.to_v(context)
+        value_heads = self._to_heads(value_projection, attn.heads)
+        if attn.norm_q is not None:
+            query_heads = attn.norm_q(query_heads)
+
+        if attn.norm_k is not None:
+            key_heads = attn.norm_k(key_heads)
+
+        probabilities = self._attention_probabilities(
+            attn,
+            query_heads,
+            key_heads,
+            attention_mask,
+        )
+        self._last_probabilities = [probabilities]
+        return self._from_heads(torch.bmm(probabilities, value_heads), attn.heads)
+
+    def _structured_attention(
+        self,
+        attn: Attention,
+        query: Float[torch.Tensor, "batch sequence channels"],
+        attention_mask: Bool[torch.Tensor, "..."] | Float[torch.Tensor, "..."] | None,
+    ) -> Float[torch.Tensor, "batch sequence channels"]:
+        conditioning = self.conditioning
+        if conditioning is None:
+            raise RuntimeError("Structured conditioning has not been installed.")
+
+        if attn.to_k is None or attn.to_v is None:
+            raise ValueError("Attention layers must define key and value projections.")
+
+        query_heads = self._to_heads(query, attn.heads)
+        if attn.norm_q is not None:
+            query_heads = attn.norm_q(query_heads)
+
+        # Project keys and values from their separate contexts before pairing them below.
+        key_heads: list[Float[torch.Tensor, "batch_heads tokens head_dim"]] = []
+        for context in conditioning.keys:
+            normalized = (
+                attn.norm_encoder_hidden_states(context) if attn.norm_cross else context
+            )
+            key_projection = attn.to_k(normalized)
+            projected = self._to_heads(key_projection, attn.heads)
+            if attn.norm_k is not None:
+                projected = attn.norm_k(projected)
+
+            key_heads.append(projected)
+
+        value_heads: list[Float[torch.Tensor, "batch_heads tokens head_dim"]] = []
+        for context in conditioning.values:
+            value_projection = attn.to_v(context)
+            value_heads.append(self._to_heads(value_projection, attn.heads))
+
+        probabilities = [
+            self._attention_probabilities(attn, query_heads, key, attention_mask)
+            for key in key_heads
+        ]
+
+        # Pair each key's attention scores with the corresponding value contexts.
+        # A singleton side is reused when only one key or value context is present.
+        if len(probabilities) == len(value_heads):
+            pairs = list(
+                zip(range(len(probabilities)), range(len(value_heads)), strict=True)
+            )
+        elif len(probabilities) == 1:
+            pairs = [(0, value_index) for value_index in range(len(value_heads))]
+        elif len(value_heads) == 1:
+            pairs = [(key_index, 0) for key_index in range(len(probabilities))]
+        else:
+            raise ValueError(
+                "Structured attention requires matching lists or a singleton side."
+            )
+
+        outputs = torch.stack(
+            [
+                torch.bmm(probabilities[key_index], value_heads[value_index])
+                for key_index, value_index in pairs
+            ]
+        )
+        validity = torch.stack(
+            [
+                conditioning.key_validity[key_index]
+                & conditioning.value_validity[value_index]
+                for key_index, value_index in pairs
+            ]
+        ).repeat_interleave(attn.heads, dim=1)
+        # Placeholder phrase contexts are excluded when averaging the valid pairs.
+        output_weights: Float[torch.Tensor, "pairs batch_heads 1 1"] = (
+            validity.to(outputs.dtype).unsqueeze(-1).unsqueeze(-1)
+        )
+        weighted_outputs = outputs * output_weights
+        summed_weights = output_weights.sum(dim=0).clamp_min(1)
+        outputs = weighted_outputs.sum(dim=0) / summed_weights
+
+        saved_probabilities: list[Float[torch.Tensor, "batch_heads query tokens"]] = []
+        for key_index, attention in enumerate(probabilities):
+            key_validity = conditioning.key_validity[key_index]
+            head_validity = key_validity.repeat_interleave(attn.heads)
+            attention_weights: Float[torch.Tensor, "batch_heads 1 1"] = (
+                head_validity.to(attention.dtype).unsqueeze(-1).unsqueeze(-1)
+            )
+            saved_probabilities.append(attention * attention_weights)
+
+        self._last_probabilities = saved_probabilities
+        return self._from_heads(outputs, attn.heads)
+
+    @staticmethod
+    def _to_heads(
+        projected: Float[torch.Tensor, "batch sequence inner_dim"], heads: int
+    ) -> Float[torch.Tensor, "batch_heads sequence head_dim"]:
+        batch_size, sequence_length, inner_dim = projected.shape
+        head_dim = inner_dim // heads
+        head_layout: Float[torch.Tensor, "batch sequence heads head_dim"] = (
+            projected.view(batch_size, sequence_length, heads, head_dim)
+        )
+        heads_first: Float[torch.Tensor, "batch heads sequence head_dim"] = (
+            head_layout.transpose(1, 2)
+        )
+        merged_batch = heads_first.reshape(
+            batch_size * heads, sequence_length, head_dim
+        )
+        return merged_batch
+
+    @staticmethod
+    def _from_heads(
+        values: Float[torch.Tensor, "batch_heads sequence head_dim"], heads: int
+    ) -> Float[torch.Tensor, "batch sequence inner_dim"]:
+        batch_heads, sequence_length, head_dim = values.shape
+        batch_size = batch_heads // heads
+        head_layout: Float[torch.Tensor, "batch heads sequence head_dim"] = values.view(
+            batch_size, heads, sequence_length, head_dim
+        )
+        batch_first: Float[torch.Tensor, "batch sequence heads head_dim"] = (
+            head_layout.transpose(1, 2)
+        )
+        merged_heads = batch_first.reshape(
+            batch_size, sequence_length, heads * head_dim
+        )
+        return merged_heads
+
+    @staticmethod
+    def _attention_probabilities(
+        attn: Attention,
+        query: Float[torch.Tensor, "batch_heads query head_dim"],
+        key: Float[torch.Tensor, "batch_heads tokens head_dim"],
+        attention_mask: Bool[torch.Tensor, "..."] | Float[torch.Tensor, "..."] | None,
+    ) -> Float[torch.Tensor, "batch_heads query tokens"]:
+        scores = torch.bmm(query, key.transpose(1, 2)) * attn.scale
+        if attention_mask is not None:
+            if attention_mask.ndim == 2:
+                mask_per_head = attention_mask.repeat_interleave(attn.heads, dim=0)
+                mask = mask_per_head.unsqueeze(1)
+            elif attention_mask.ndim == 4:
+                mask = attention_mask.reshape(
+                    -1, attention_mask.shape[-2], attention_mask.shape[-1]
+                )
+            elif attention_mask.shape[0] * attn.heads == scores.shape[0]:
+                mask = attention_mask.repeat_interleave(attn.heads, dim=0)
+            else:
+                mask = attention_mask
+
+            if mask.dtype == torch.bool:
+                scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
+            else:
+                scores = scores + mask
+
+        return scores.softmax(dim=-1)
+
+    @staticmethod
+    def _split_mask(
+        attention_mask: Bool[torch.Tensor, "..."] | Float[torch.Tensor, "..."] | None,
+        base_batch: int,
+        total_batch: int,
+    ) -> tuple[
+        Bool[torch.Tensor, "..."] | Float[torch.Tensor, "..."] | None,
+        Bool[torch.Tensor, "..."] | Float[torch.Tensor, "..."] | None,
+    ]:
+        if attention_mask is None or attention_mask.shape[0] != total_batch:
+            return attention_mask, attention_mask
+
+        unconditional_mask, conditional_mask = attention_mask.split(base_batch, dim=0)
+        return unconditional_mask, conditional_mask

pipeline.py

@@ -0,0 +1,499 @@
+"""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}")

processing_structure_diffusion.py

@@ -0,0 +1,204 @@
+"""Prompt parsing and CLIP token alignment for structured diffusion guidance."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Protocol, cast
+
+import torch
+from jaxtyping import Float
+from nltk.tree import Tree
+
+
+class _StanzaWord(Protocol):
+    start_char: int | None
+    end_char: int | None
+
+
+class _StanzaSentence(Protocol):
+    constituency: str
+    words: list[_StanzaWord]
+
+
+class _StanzaDocument(Protocol):
+    sentences: list[_StanzaSentence]
+
+
+class _StanzaPipeline(Protocol):
+    def __call__(self, text: str) -> _StanzaDocument: ...
+
+
+class _Tokenizer(Protocol):
+    def tokenize(self, text: str) -> list[str]: ...
+
+
+@dataclass(frozen=True)
+class NounPhrase:
+    """A parsed noun phrase with character offsets into normalized prompt text."""
+
+    text: str
+    start: int
+    end: int
+
+
+def normalize_prompt(prompt: str) -> str:
+    """Normalize casing and trailing punctuation before parsing and token alignment."""
+    return prompt.lower().strip().strip(".").strip()
+
+
+def constituency_noun_phrases(prompt: str) -> list[NounPhrase]:
+    """Parse a prompt and return its multi-token noun phrases.
+
+    The full prompt is represented separately by the pipeline. Returned offsets
+    refer to normalized prompt text so they can be aligned to CLIP BPE tokens.
+    Stanza models must be installed before this function runs.
+    """
+    document = _constituency_pipeline()(prompt)
+    if not document.sentences:
+        return []
+
+    sentence = document.sentences[0]
+    tree = Tree.fromstring(str(sentence.constituency))
+    word_offsets: list[tuple[int, int]] = []
+    for word in sentence.words:
+        if word.start_char is None or word.end_char is None:
+            raise ValueError("The constituency parser did not return word offsets.")
+
+        word_offsets.append((word.start_char, word.end_char))
+
+    if len(tree.leaves()) != len(word_offsets):
+        raise ValueError(
+            "The constituency tree and parser word offsets have different lengths."
+        )
+
+    phrases: list[NounPhrase] = []
+
+    def visit(node: Tree | str, word_start: int) -> int:
+        if isinstance(node, str):
+            return 1
+
+        cursor = word_start
+        for child in node:
+            cursor += visit(child, cursor)
+
+        if node.label() == "NP" and cursor - word_start > 1:
+            char_start = word_offsets[word_start][0]
+            char_end = word_offsets[cursor - 1][1]
+            text = prompt[char_start:char_end].strip()
+            if text and text != prompt:
+                phrases.append(NounPhrase(text=text, start=char_start, end=char_end))
+
+        return cursor - word_start
+
+    visit(tree, 0)
+    return phrases
+
+
+def scene_graph_noun_phrases(prompt: str) -> list[NounPhrase]:
+    """Return entity and relation spans from the optional scene-graph parser."""
+    try:
+        import sng_parser
+    except ImportError as exc:
+        raise ImportError(
+            "Install structure-diffusion[scene-graph] to use parser_type='scene_graph'."
+        ) from exc
+
+    word_offsets = [
+        (match.start(), match.end()) for match in re.finditer(r"\S+", prompt)
+    ]
+    graph = sng_parser.parse(prompt)
+    spans: list[tuple[int, int]] = []
+
+    for entity in graph["entities"]:
+        start, end = entity["span_bounds"]
+        if end - start > 1 and 0 <= start < end <= len(word_offsets):
+            spans.append((start, end))
+
+    for relation in graph["relations"]:
+        subject_start, subject_end = graph["entities"][relation["subject"]][
+            "span_bounds"
+        ]
+        object_start, object_end = graph["entities"][relation["object"]]["span_bounds"]
+        start = min(subject_start, object_start)
+        end = max(subject_end, object_end)
+        if end - start > 1 and 0 <= start < end <= len(word_offsets):
+            spans.append((start, end))
+
+    phrases: list[NounPhrase] = []
+    seen: set[tuple[int, int]] = set()
+    for start, end in spans:
+        if (start, end) in seen:
+            continue
+
+        seen.add((start, end))
+        char_start = word_offsets[start][0]
+        char_end = word_offsets[end - 1][1]
+        text = prompt[char_start:char_end].strip()
+        if text and text != prompt:
+            phrases.append(NounPhrase(text=text, start=char_start, end=char_end))
+
+    return phrases
+
+
+def token_span_from_char_span(
+    tokenizer: _Tokenizer,
+    prompt: str,
+    start: int,
+    end: int,
+    sequence_length: int,
+) -> tuple[int, int] | None:
+    """Map prompt character offsets to a BOS-shifted, truncated CLIP token span."""
+    if start < 0 or end <= start or end > len(prompt):
+        raise ValueError("Phrase character offsets must be a non-empty prompt span.")
+
+    # CLIP reserves token index 0 for BOS, so shift text-token offsets by one.
+    start_token = len(tokenizer.tokenize(prompt[:start])) + 1
+    end_token = len(tokenizer.tokenize(prompt[:end])) + 1
+    end_token = min(end_token, sequence_length - 1)
+    if start_token >= end_token:
+        return None
+
+    return start_token, end_token
+
+
+def align_phrase_context(
+    prompt_context: Float[torch.Tensor, "batch tokens dim"],
+    phrase_context: Float[torch.Tensor, "batch tokens dim"],
+    token_span: tuple[int, int],
+    phrase_token_count: int,
+) -> Float[torch.Tensor, "batch tokens dim"]:
+    """Copy only actual phrase-token embeddings into the matching prompt span."""
+    start, end = token_span
+    span_length = end - start
+    if start < 1 or end > prompt_context.shape[1] or span_length < 1:
+        raise ValueError("The aligned token span is outside the prompt context.")
+
+    if phrase_token_count < span_length:
+        raise ValueError("The phrase encoding has fewer real tokens than its span.")
+
+    if phrase_context.shape[0] != prompt_context.shape[0]:
+        raise ValueError("Prompt and phrase context batches must match.")
+
+    aligned = prompt_context.clone()
+    # Skip the phrase encoding's BOS token and copy only the aligned phrase tokens.
+    aligned[:, start:end] = phrase_context[:, 1 : 1 + span_length]
+    return aligned
+
+
+@lru_cache(maxsize=1)
+def _constituency_pipeline() -> _StanzaPipeline:
+    """Build a reusable English constituency parser without downloading models."""
+    import stanza
+
+    return cast(
+        _StanzaPipeline,
+        stanza.Pipeline(
+            lang="en",
+            processors="tokenize,pos,constituency",
+            download_method=None,
+            use_gpu=False,
+            verbose=False,
+        ),
+    )