Upload 7 files

README.md

@@ -1,13 +1,25 @@
----
-title: SynGen
-emoji: 🦀
-colorFrom: gray
-colorTo: purple
-sdk: gradio
-sdk_version: 3.34.0
-app_file: app.py
-pinned: false
-license: cc-by-sa-4.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Linguistic Binding in Diffusion Models: Enhancing Attribute Correspondence through Attention Map Alignment
+
+## Setup
+Clone this repository and create a conda environment:
+```
+conda env create -f environment.yaml
+conda activate syngen
+```
+
+If you rather use an existing environment, just run:
+```
+pip install -r requirements.txt
+```
+
+Finally, run:
+```
+python -m spacy download en_core_web_trf
+```
+
+## Inference
+```
+python run.py --prompt "a horned lion and a spotted monkey" --seed 1269
+```
+
+Note that this will download the stable diffusion model `CompVis/stable-diffusion-v1-4`. If you rather use an existing copy of the model, provide the absolute path using `--model_path`.

app.py

@@ -0,0 +1,37 @@
+import gradio as gr
+import torch
+
+from syngen_diffusion_pipeline import SynGenDiffusionPipeline
+
+model_path = 'CompVis/stable-diffusion-v1-4'
+device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+pipe = SynGenDiffusionPipeline.from_pretrained(model_path).to(device)
+
+
+def generate_fn(prompt, seed):
+    generator = torch.Generator(device.type).manual_seed(int(seed))
+    result = pipe(prompt=prompt, generator=generator, num_inference_steps=50)
+    return result['images'][0]
+
+title = "SynGen"
+description = """
+This is the demo for [SynGen](https://github.com/RoyiRa/Syntax-Guided-Generation), an image synthesis approach which first syntactically analyses the prompt to identify entities and their modifiers, and then uses a novel loss function that encourages the cross-attention maps to agree with the linguistic binding reflected by the syntax. Preprint: \"Linguistic Binding in Diffusion Models: Enhancing Attribute Correspondence through Attention Map Alignment\" (arxiv link coming soon).
+"""
+
+examples = [
+    ["a yellow flamingo and a pink sunflower", "16"],
+    ["a yellow flamingo and a pink sunflower", "60"],
+    ["a checkered bowl in a cluttered room", "69"],
+    ["a checkered bowl in a cluttered room", "77"],
+    ["a horned lion and a spotted monkey", "1269"],
+    ["a horned lion and a spotted monkey", "9146"]
+]
+
+prompt_textbox = gr.Textbox(label="Prompt", placeholder="A yellow flamingo and a pink sunflower", lines=1)
+seed_textbox = gr.Textbox(label="Seed", placeholder="42", lines=1)
+
+output = gr.Image(label="generation")
+demo = gr.Interface(fn=generate_fn, inputs=[prompt_textbox, seed_textbox], outputs=output, examples=examples,
+                    title=title, description=description, allow_flagging=False)
+
+demo.launch()

compute_loss.py

@@ -0,0 +1,211 @@
+import torch.distributions as dist
+from typing import List, Dict
+import itertools
+
+start_token = "<|startoftext|>"
+end_token = "<|endoftext|>"
+
+
+def _get_outside_indices(subtree_indices, attn_map_idx_to_wp):
+    flattened_subtree_indices = _flatten_indices(subtree_indices)
+    outside_indices = [
+        map_idx
+        for map_idx in attn_map_idx_to_wp.keys() if (map_idx not in flattened_subtree_indices)
+    ]
+    return outside_indices
+
+
+def _flatten_indices(related_indices):
+    flattened_related_indices = []
+    for item in related_indices:
+        if isinstance(item, list):
+            flattened_related_indices.extend(item)
+        else:
+            flattened_related_indices.append(item)
+    return flattened_related_indices
+
+
+def split_indices(related_indices: List[int]):
+    noun = [related_indices[-1]]  # assumes noun is always last in the list
+    modifier = related_indices[:-1]
+    if isinstance(modifier, int):
+        modifier = [modifier]
+    return noun, modifier
+
+
+def _symmetric_kl(attention_map1, attention_map2):
+    # Convert map into a single distribution: 16x16 -> 256
+    if len(attention_map1.shape) > 1:
+        attention_map1 = attention_map1.reshape(-1)
+    if len(attention_map2.shape) > 1:
+        attention_map2 = attention_map2.reshape(-1)
+
+    p = dist.Categorical(probs=attention_map1)
+    q = dist.Categorical(probs=attention_map2)
+
+    kl_divergence_pq = dist.kl_divergence(p, q)
+    kl_divergence_qp = dist.kl_divergence(q, p)
+
+    avg_kl_divergence = (kl_divergence_pq + kl_divergence_qp) / 2
+    return avg_kl_divergence
+
+
+def calculate_positive_loss(attention_maps, modifier, noun):
+    src_indices = modifier
+    dest_indices = noun
+
+    if isinstance(src_indices, list) and isinstance(dest_indices, list):
+        wp_pos_loss = [
+            _symmetric_kl(attention_maps[s], attention_maps[d])
+            for (s, d) in itertools.product(src_indices, dest_indices)
+        ]
+        positive_loss = max(wp_pos_loss)
+    elif isinstance(dest_indices, list):
+        wp_pos_loss = [
+            _symmetric_kl(attention_maps[src_indices], attention_maps[d])
+            for d in dest_indices
+        ]
+        positive_loss = max(wp_pos_loss)
+    elif isinstance(src_indices, list):
+        wp_pos_loss = [
+            _symmetric_kl(attention_maps[s], attention_maps[dest_indices])
+            for s in src_indices
+        ]
+        positive_loss = max(wp_pos_loss)
+    else:
+        positive_loss = _symmetric_kl(
+            attention_maps[src_indices], attention_maps[dest_indices]
+        )
+
+    return positive_loss
+
+
+def _calculate_outside_loss(attention_maps, src_indices, outside_loss):
+    negative_loss = []
+    computed_pairs = set()
+    pair_counter = 0
+
+    for outside_idx in outside_loss:
+        if isinstance(src_indices, list):
+            wp_neg_loss = []
+            for t in src_indices:
+                pair_key = (t, outside_idx)
+                if pair_key not in computed_pairs:
+                    wp_neg_loss.append(
+                        _symmetric_kl(
+                            attention_maps[t], attention_maps[outside_idx]
+                        )
+                    )
+                    computed_pairs.add(pair_key)
+            negative_loss.append(max(wp_neg_loss) if wp_neg_loss else 0)
+            pair_counter += 1
+
+        else:
+            pair_key = (src_indices, outside_idx)
+            if pair_key not in computed_pairs:
+                negative_loss.append(
+                    _symmetric_kl(
+                        attention_maps[src_indices], attention_maps[outside_idx]
+                    )
+                )
+                computed_pairs.add(pair_key)
+                pair_counter += 1
+
+    return negative_loss, pair_counter
+
+
+def align_wordpieces_indices(
+        wordpieces2indices, start_idx, target_word
+):
+    """
+    Aligns a `target_word` that contains more than one wordpiece (the first wordpiece is `start_idx`)
+    """
+
+    wp_indices = [start_idx]
+    wp = wordpieces2indices[start_idx].replace("</w>", "")
+
+    # Run over the next wordpieces in the sequence (which is why we use +1)
+    for wp_idx in range(start_idx + 1, len(wordpieces2indices)):
+        if wp == target_word:
+            break
+
+        wp2 = wordpieces2indices[wp_idx].replace("</w>", "")
+        if target_word.startswith(wp + wp2) and wp2 != target_word:
+            wp += wordpieces2indices[wp_idx].replace("</w>", "")
+            wp_indices.append(wp_idx)
+        else:
+            wp_indices = (
+                []
+            )  # if there's no match, you want to clear the list and finish
+            break
+
+    return wp_indices
+
+
+def extract_attribution_indices(prompt, parser):
+    doc = parser(prompt)
+    subtrees = []
+    modifiers = ["amod", "nmod", "compound", "npadvmod", "advmod", "acomp"]
+
+    for w in doc:
+        if w.pos_ not in ["NOUN", "PROPN"] or w.dep_ in modifiers:
+            continue
+        subtree = []
+        stack = []
+        for child in w.children:
+            if child.dep_ in modifiers:
+                subtree.append(child)
+                stack.extend(child.children)
+
+        while stack:
+            node = stack.pop()
+            if node.dep_ in modifiers or node.dep_ == "conj":
+                subtree.append(node)
+                stack.extend(node.children)
+        if subtree:
+            subtree.append(w)
+            subtrees.append(subtree)
+    return subtrees
+
+
+def calculate_negative_loss(
+        attention_maps, modifier, noun, subtree_indices, attn_map_idx_to_wp
+):
+    outside_indices = _get_outside_indices(subtree_indices, attn_map_idx_to_wp)
+    negative_modifier_loss, num_modifier_pairs = _calculate_outside_loss(
+        attention_maps, modifier, outside_indices
+    )
+    negative_noun_loss, num_noun_pairs = _calculate_outside_loss(
+        attention_maps, noun, outside_indices
+    )
+
+    negative_modifier_loss = -sum(negative_modifier_loss) / len(outside_indices)
+    negative_noun_loss = -sum(negative_noun_loss) / len(outside_indices)
+
+    negative_loss = (negative_modifier_loss + negative_noun_loss) / 2
+
+    return negative_loss
+
+def get_indices(tokenizer, prompt: str) -> Dict[str, int]:
+    """Utility function to list the indices of the tokens you wish to alte"""
+    ids = tokenizer(prompt).input_ids
+    indices = {
+        i: tok
+        for tok, i in zip(
+            tokenizer.convert_ids_to_tokens(ids), range(len(ids))
+        )
+    }
+    return indices
+
+def get_attention_map_index_to_wordpiece(tokenizer, prompt):
+    attn_map_idx_to_wp = {}
+
+    wordpieces2indices = get_indices(tokenizer, prompt)
+
+    # Ignore `start_token` and `end_token`
+    for i in list(wordpieces2indices.keys())[1:-1]:
+        wordpiece = wordpieces2indices[i]
+        wordpiece = wordpiece.replace("</w>", "")
+        attn_map_idx_to_wp[i] = wordpiece
+
+    return attn_map_idx_to_wp

environment.yaml

@@ -0,0 +1,15 @@
+name: syngen
+channels:
+  - pytorch
+  - defaults
+dependencies:
+  - python=3.11.2
+  - pip=23.0.1
+  - pip:
+    - diffusers==0.14.0
+    - numpy==1.23.3
+    - spacy==3.5.2
+    - tqdm==4.65.0
+    - transformers @ git+https://github.com/huggingface/transformers.git@dbc12269ed5546b2da9236b9f1078b95b6a4d3d5
+    - torch==2.0.0
+    - accelerate==0.18.0

requirements.txt

@@ -0,0 +1,7 @@
+diffusers==0.14.0
+numpy==1.23.3
+spacy==3.5.2
+tqdm==4.65.0
+transformers @ git+https://github.com/huggingface/transformers.git@dbc12269ed5546b2da9236b9f1078b95b6a4d3d5
+torch==2.0.0
+accelerate==0.18.0

run.py

@@ -0,0 +1,66 @@
+import argparse
+import os
+
+import torch
+from syngen_diffusion_pipeline import SynGenDiffusionPipeline
+
+
+def main(prompt, seed, output_directory, model_path):
+    pipe = load_model(model_path)
+    image = generate(pipe, prompt, seed)
+    save_image(image, prompt, seed, output_directory)
+
+
+def load_model(model_path):
+    device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+    pipe = SynGenDiffusionPipeline.from_pretrained(model_path).to(device)
+
+    return pipe
+
+
+def generate(pipe, prompt, seed):
+    device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+    generator = torch.Generator(device.type).manual_seed(seed)
+    result = pipe(prompt=prompt, generator=generator)
+    return result['images'][0]
+
+
+def save_image(image, prompt, seed, output_directory):
+    if not os.path.exists(output_directory):
+        os.makedirs(output_directory)
+
+    file_name = f"{output_directory}/{prompt}_{seed}.png"
+    image.save(file_name)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default="a checkered bowl on a red and blue table"
+    )
+
+    parser.add_argument(
+        '--seed',
+        type=int,
+        default=1924
+    )
+
+    parser.add_argument(
+        '--output_directory',
+        type=str,
+        default='./output'
+    )
+
+    parser.add_argument(
+        '--model_path',
+        type=str,
+        default='CompVis/stable-diffusion-v1-4',
+        help='The path to the model (this will download the model if the path doesn\'t exist)'
+    )
+
+    args = parser.parse_args()
+
+    main(args.prompt, args.seed, args.output_directory, args.model_path)

syngen_diffusion_pipeline.py

@@ -0,0 +1,495 @@
+import itertools
+from typing import Any, Callable, Dict, Optional, Union, List
+
+import spacy
+import torch
+from diffusers import StableDiffusionPipeline, AutoencoderKL, UNet2DConditionModel
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput, StableDiffusionSafetyChecker
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import (
+    EXAMPLE_DOC_STRING,
+)
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_attend_and_excite import (
+    AttentionStore,
+    AttendExciteCrossAttnProcessor,
+)
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import (
+    logging,
+    replace_example_docstring,
+)
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPFeatureExtractor
+
+from compute_loss import get_attention_map_index_to_wordpiece, split_indices, calculate_positive_loss, calculate_negative_loss, get_indices, start_token, end_token, \
+    align_wordpieces_indices, extract_attribution_indices
+
+logger = logging.get_logger(__name__)
+
+
+class SynGenDiffusionPipeline(StableDiffusionPipeline):
+    def __init__(self,
+                 vae: AutoencoderKL,
+                 text_encoder: CLIPTextModel,
+                 tokenizer: CLIPTokenizer,
+                 unet: UNet2DConditionModel,
+                 scheduler: KarrasDiffusionSchedulers,
+                 safety_checker: StableDiffusionSafetyChecker,
+                 feature_extractor: CLIPFeatureExtractor,
+                 requires_safety_checker: bool = True,
+                 ):
+        super().__init__(vae, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor,
+                         requires_safety_checker)
+
+        self.parser = spacy.load("en_core_web_trf")
+
+    def _aggregate_and_get_attention_maps_per_token(self):
+        attention_maps = self.attention_store.aggregate_attention(
+            from_where=("up", "down", "mid"),
+        )
+        attention_maps_list = _get_attention_maps_list(
+            attention_maps=attention_maps
+        )
+        return attention_maps_list
+
+    @staticmethod
+    def _update_latent(
+            latents: torch.Tensor, loss: torch.Tensor, step_size: float
+    ) -> torch.Tensor:
+        """Update the latent according to the computed loss."""
+        grad_cond = torch.autograd.grad(
+            loss.requires_grad_(True), [latents], retain_graph=True
+        )[0]
+        latents = latents - step_size * grad_cond
+        return latents
+
+    def register_attention_control(self):
+        attn_procs = {}
+        cross_att_count = 0
+        for name in self.unet.attn_processors.keys():
+            if name.startswith("mid_block"):
+                place_in_unet = "mid"
+            elif name.startswith("up_blocks"):
+                place_in_unet = "up"
+            elif name.startswith("down_blocks"):
+                place_in_unet = "down"
+            else:
+                continue
+
+            cross_att_count += 1
+            attn_procs[name] = AttendExciteCrossAttnProcessor(
+                attnstore=self.attention_store, place_in_unet=place_in_unet
+            )
+
+        self.unet.set_attn_processor(attn_procs)
+        self.attention_store.num_att_layers = cross_att_count
+
+    # Based on StableDiffusionPipeline.__call__ . New code is annotated with NEW.
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+            self,
+            prompt: Union[str, List[str]] = None,
+            height: Optional[int] = None,
+            width: Optional[int] = None,
+            num_inference_steps: int = 50,
+            guidance_scale: float = 7.5,
+            negative_prompt: Optional[Union[str, List[str]]] = None,
+            num_images_per_prompt: Optional[int] = 1,
+            eta: float = 0.0,
+            generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+            latents: Optional[torch.FloatTensor] = None,
+            prompt_embeds: Optional[torch.FloatTensor] = None,
+            negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+            output_type: Optional[str] = "pil",
+            return_dict: bool = True,
+            callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+            callback_steps: int = 1,
+            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+            syngen_step_size: int = 20,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
+                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
+                plain tuple.
+            callback (`Callable`, *optional*):
+                A function that will be called every `callback_steps` steps during inference. The function will be
+                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function will be called. If not specified, the callback will be
+                called at every step.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttnProcessor` as defined under
+                `self.processor` in
+                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
+            syngen_step_size (`int`, *optional*, default to 20):
+                Controls the step size of each SynGen update.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
+            When returning a tuple, the first element is a list with the generated images, and the second element is a
+            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
+            (nsfw) content, according to the `safety_checker`.
+        """
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            callback_steps,
+            negative_prompt,
+            prompt_embeds,
+            negative_prompt_embeds,
+        )
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        prompt_embeds = self._encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+        )
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # NEW - stores the attention calculated in the unet
+        self.attention_store = AttentionStore()
+        self.register_attention_control()
+
+        # NEW
+        text_embeddings = (
+            prompt_embeds[batch_size * num_images_per_prompt:] if do_classifier_free_guidance else prompt_embeds
+        )
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # NEW
+                latents = self._syngen_step(
+                    latents,
+                    text_embeddings,
+                    t,
+                    i,
+                    syngen_step_size,
+                    cross_attention_kwargs,
+                    prompt,
+                    max_iter_to_alter=25,
+                )
+
+                # expand the latents if we are doing classifier free guidance
+                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, t
+                )
+
+                # predict the noise residual
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                ).sample
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (
+                            noise_pred_text - noise_pred_uncond
+                    )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(
+                    noise_pred, t, latents, **extra_step_kwargs
+                ).prev_sample
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                        (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+
+        if output_type == "latent":
+            image = latents
+            has_nsfw_concept = None
+        elif output_type == "pil":
+            # 8. Post-processing
+            image = self.decode_latents(latents)
+
+            # 9. Run safety checker
+            image, has_nsfw_concept = self.run_safety_checker(
+                image, device, prompt_embeds.dtype
+            )
+
+            # 10. Convert to PIL
+            image = self.numpy_to_pil(image)
+        else:
+            # 8. Post-processing
+            image = self.decode_latents(latents)
+
+            # 9. Run safety checker
+            image, has_nsfw_concept = self.run_safety_checker(
+                image, device, prompt_embeds.dtype
+            )
+
+        # Offload last model to CPU
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.final_offload_hook.offload()
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return StableDiffusionPipelineOutput(
+            images=image, nsfw_content_detected=has_nsfw_concept
+        )
+
+    def _syngen_step(
+            self,
+            latents,
+            text_embeddings,
+            t,
+            i,
+            step_size,
+            cross_attention_kwargs,
+            prompt,
+            max_iter_to_alter=25,
+    ):
+        with torch.enable_grad():
+            latents = latents.clone().detach().requires_grad_(True)
+            updated_latents = []
+            for latent, text_embedding in zip(latents, text_embeddings):
+                # Forward pass of denoising with text conditioning
+                latent = latent.unsqueeze(0)
+                text_embedding = text_embedding.unsqueeze(0)
+
+                self.unet(
+                    latent,
+                    t,
+                    encoder_hidden_states=text_embedding,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                ).sample
+                self.unet.zero_grad()
+
+                # Get attention maps
+                attention_maps = self._aggregate_and_get_attention_maps_per_token()
+
+                loss = self._compute_loss(attention_maps=attention_maps, prompt=prompt)
+
+                # Perform gradient update
+                if i < max_iter_to_alter:
+                    if loss != 0:
+                        latent = self._update_latent(
+                            latents=latent, loss=loss, step_size=step_size
+                        )
+                    logger.info(f"Iteration {i} | Loss: {loss:0.4f}")
+
+            updated_latents.append(latent)
+
+        latents = torch.cat(updated_latents, dim=0)
+
+        return latents
+
+    def _compute_loss(
+            self, attention_maps: List[torch.Tensor], prompt: Union[str, List[str]]
+    ) -> torch.Tensor:
+        attn_map_idx_to_wp = get_attention_map_index_to_wordpiece(self.tokenizer, prompt)
+        loss = self._attribution_loss(attention_maps, prompt, attn_map_idx_to_wp)
+
+        return loss
+
+
+    def _attribution_loss(
+            self,
+            attention_maps: List[torch.Tensor],
+            prompt: Union[str, List[str]],
+            attn_map_idx_to_wp,
+    ) -> torch.Tensor:
+        subtrees_indices = self._extract_attribution_indices(prompt)
+        loss = 0
+
+        for subtree_indices in subtrees_indices:
+            noun, modifier = split_indices(subtree_indices)
+            all_subtree_pairs = list(itertools.product(noun, modifier))
+            positive_loss, negative_loss = self._calculate_losses(
+                attention_maps,
+                all_subtree_pairs,
+                subtree_indices,
+                attn_map_idx_to_wp,
+            )
+            loss += positive_loss
+            loss += negative_loss
+
+        return loss
+
+    def _calculate_losses(
+            self,
+            attention_maps,
+            all_subtree_pairs,
+            subtree_indices,
+            attn_map_idx_to_wp,
+    ):
+        positive_loss = []
+        negative_loss = []
+        for pair in all_subtree_pairs:
+            noun, modifier = pair
+            positive_loss.append(
+                calculate_positive_loss(attention_maps, modifier, noun)
+            )
+            negative_loss.append(
+                calculate_negative_loss(
+                    attention_maps, modifier, noun, subtree_indices, attn_map_idx_to_wp
+                )
+            )
+
+        positive_loss = sum(positive_loss)
+        negative_loss = sum(negative_loss)
+
+        return positive_loss, negative_loss
+
+    def _align_indices(self, prompt, spacy_pairs):
+        wordpieces2indices = get_indices(self.tokenizer, prompt)
+        paired_indices = []
+        collected_spacy_indices = (
+            set()
+        )  # helps track recurring nouns across different relations (i.e., cases where there is more than one instance of the same word)
+
+        for pair in spacy_pairs:
+            curr_collected_wp_indices = (
+                []
+            )  # helps track which nouns and amods were added to the current pair (this is useful in sentences with repeating amod on the same relation (e.g., "a red red red bear"))
+            for member in pair:
+                for idx, wp in wordpieces2indices.items():
+                    if wp in [start_token, end_token]:
+                        continue
+
+                    wp = wp.replace("</w>", "")
+                    if member.text == wp:
+                        if idx not in curr_collected_wp_indices and idx not in collected_spacy_indices:
+                            curr_collected_wp_indices.append(idx)
+                            break
+                    # take care of wordpieces that are split up
+                    elif member.text.startswith(wp) and wp != member.text:  # can maybe be while loop
+                        wp_indices = align_wordpieces_indices(
+                            wordpieces2indices, idx, member.text
+                        )
+                        # check if all wp_indices are not already in collected_spacy_indices
+                        if wp_indices and (wp_indices not in curr_collected_wp_indices) and all([wp_idx not in collected_spacy_indices for wp_idx in wp_indices]):
+                            curr_collected_wp_indices.append(wp_indices)
+                            break
+
+            for collected_idx in curr_collected_wp_indices:
+                if isinstance(collected_idx, list):
+                    for idx in collected_idx:
+                        collected_spacy_indices.add(idx)
+                else:
+                    collected_spacy_indices.add(collected_idx)
+
+            paired_indices.append(curr_collected_wp_indices)
+            
+        return paired_indices
+
+    def _extract_attribution_indices(self, prompt):
+        pairs = extract_attribution_indices(prompt, self.parser)
+        paired_indices = self._align_indices(prompt, pairs)
+        return paired_indices
+
+
+
+def _get_attention_maps_list(
+        attention_maps: torch.Tensor
+) -> List[torch.Tensor]:
+    attention_maps *= 100
+    attention_maps_list = [
+        attention_maps[:, :, i] for i in range(attention_maps.shape[2])
+    ]
+
+    return attention_maps_list