Upload app

app.py

@@ -0,0 +1,354 @@
+from diffusers import DDIMScheduler
+from pipeline_stable_diffusion_xl_opt import StableDiffusionXLPipeline
+from injection_utils import regiter_attention_editor_diffusers
+from bounded_attention import BoundedAttention
+from pytorch_lightning import seed_everything
+
+import gradio as gr
+import torch
+import numpy as np
+from PIL import Image, ImageDraw
+
+from functools import partial
+
+RESOLUTION = 512
+MIN_SIZE = 0.01
+COLORS = ["red", "blue", "green", "orange", "purple", "turquoise", "olive"]
+
+
+def inference(
+    device,
+    model,
+    boxes,
+    prompts,
+    subject_token_indices,
+    filter_token_indices,
+    num_tokens,
+    init_step_size,
+    final_step_size,
+    num_clusters_per_subject,
+    cross_loss_scale,
+    self_loss_scale,
+    classifier_free_guidance_scale,
+    num_iterations,
+    loss_threshold,
+    num_guidance_steps,
+    seed,
+):
+    seed_everything(seed)
+    start_code = torch.randn([len(prompts), 4, 128, 128], device=device)
+    editor = BoundedAttention(
+        boxes,
+        prompts,
+        subject_token_indices,
+        list(range(70, 82)),
+        list(range(70, 82)),
+        eos_token_index=num_tokens + 1,
+        cross_loss_coef=cross_loss_scale,
+        self_loss_coef=self_loss_scale,
+        filter_token_indices=filter_token_indices,
+        max_guidance_iter=num_guidance_steps,
+        max_guidance_iter_per_step=num_iterations,
+        start_step_size=init_step_size,
+        end_step_size=final_step_size,
+        loss_stopping_value=loss_threshold,
+        num_clusters_per_box=num_clusters_per_subject,
+        debug=False,
+    )
+
+    regiter_attention_editor_diffusers(model, editor)
+    return model(prompts, latents=start_code, guidance_scale=classifier_free_guidance_scale).images
+
+
+def generate(
+    device,
+    model,
+    prompt,
+    subject_token_indices,
+    filter_token_indices,
+    num_tokens,
+    init_step_size,
+    final_step_size,
+    num_clusters_per_subject,
+    cross_loss_scale,
+    self_loss_scale,
+    classifier_free_guidance_scale,
+    batch_size,
+    num_iterations,
+    loss_threshold,
+    num_guidance_steps,
+    seed,
+    boxes
+):
+    if 'boxes' not in boxes:
+        boxes['boxes'] = []
+
+    boxes = boxes['boxes']
+    subject_token_indices = convert_token_indices(subject_token_indices, nested=True)
+    if len(boxes) != len(subject_token_indices):
+        raise ValueError("""
+            The number of boxes should be equal to the number of subject token indices.
+            Number of boxes drawn: {}, number of grounding tokens: {}.
+        """.format(len(boxes), len(subject_token_indices)))
+
+    filter_token_indices = convert_token_indices(filter_token_indices) if len(filter_token_indices.strip()) > 0 else None
+    num_tokens = int(num_tokens) if len(num_tokens.strip()) > 0 else None
+    prompts = [prompt.strip('.').strip(',').strip()] * batch_size
+
+    images = inference(
+        device, model, boxes, prompts, subject_token_indices, filter_token_indices, num_tokens, init_step_size,
+        final_step_size, num_clusters_per_subject, cross_loss_scale, self_loss_scale, classifier_free_guidance_scale,
+        num_iterations, loss_threshold, num_guidance_steps, seed)
+
+    blank_samples = batch_size % 2 if batch_size > 1 else 0
+    images = [gr.Image.update(value=x, visible=True) for i, x in enumerate(images)] \
+                 + [gr.Image.update(value=None, visible=True) for _ in range(blank_samples)] \
+                 + [gr.Image.update(value=None, visible=False) for _ in range(4 - batch_size - blank_samples)]
+
+    return images
+
+
+def convert_token_indices(token_indices, nested=False):
+    if nested:
+        return [convert_token_indices(indices, nested=False) for indices in token_indices.split(';')]
+
+    return [int(index.strip()) for index in token_indices.split(',') if len(index.strip()) > 0]
+
+
+def draw(boxes, mask):
+    print('Called draw')
+    print('before boxes', boxes)
+    if mask.ndim == 3:
+        mask = 255 - mask[..., 0]
+
+    mask = (mask != 0).astype('uint8') * 255
+    if mask.sum() > 0:
+        x1x2 = np.where(mask.max(0) != 0)[0] / RESOLUTION
+        y1y2 = np.where(mask.max(1) != 0)[0] / RESOLUTION
+        y1, y2 = y1y2.min(), y1y2.max()
+        x1, x2 = x1x2.min(), x1x2.max()
+
+        if (x2 - x1 > MIN_SIZE) and (y2 - y1 > MIN_SIZE):
+            boxes.append((x1, y1, x2, y2))
+            layout_image = draw_boxes(np.array(boxes) * RESOLUTION)
+
+    print('after boxes', boxes)
+    return [boxes, None, layout_image]
+
+
+def draw_boxes(boxes):
+    if len(boxes) == 0:
+        return None
+
+    image = Image.new('RGB', (RESOLUTION, RESOLUTION), (255, 255, 255))
+    drawing = ImageDraw.Draw(image)
+    print(boxes)
+    for i, box in enumerate(boxes):
+        drawing.rectangle(box, outline=COLORS[i % len(COLORS)], width=4)
+
+    return image
+
+
+def clear(batch_size):
+    blank_samples = batch_size % 2 if batch_size > 1 else 0
+    out_images = [gr.Image.update(value=None, visible=True) for i in range(batch_size)]
+    return [[], None, None] + out_images
+
+
+def main():
+    css = """
+    #paper-info a {
+        color:#008AD7;
+        text-decoration: none;
+    }
+    #paper-info a:hover {
+        cursor: pointer;
+        text-decoration: none;
+    }
+
+    .tooltip {
+        color: #555;
+        position: relative;
+        display: inline-block;
+        cursor: pointer;
+    }
+
+    .tooltip .tooltiptext {
+        visibility: hidden;
+        width: 400px;
+        background-color: #555;
+        color: #fff;
+        text-align: center;
+        padding: 5px;
+        border-radius: 5px;
+        position: absolute;
+        z-index: 1; /* Set z-index to 1 */
+        left: 10px;
+        top: 100%;
+        opacity: 0;
+        transition: opacity 0.3s;
+    }
+
+    .tooltip:hover .tooltiptext {
+        visibility: visible;
+        opacity: 1;
+        z-index: 9999; /* Set a high z-index value when hovering */
+    }
+    """
+
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    model_path = "stabilityai/stable-diffusion-xl-base-1.0"
+    scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
+    model = StableDiffusionXLPipeline.from_pretrained(model_path, scheduler=scheduler, torch_dtype=torch.float16).to(device)
+    model.unet.set_default_attn_processor()
+    model.enable_xformers_memory_efficient_attention()
+    model.enable_sequential_cpu_offload()
+
+    with gr.Blocks(
+            css=css,
+            title="Bounded Attention demo",
+    ) as demo:
+        description = """<p style="text-align: center; font-weight: bold;">
+            <span style="font-size: 28px">Bounded Attention</span>
+            <br>
+            <span style="font-size: 18px" id="paper-info">
+                [<a href="https://omer11a.github.io/bounded-attention/" target="_blank">Project Page</a>]
+                [<a href=" " target="_blank">Paper</a>]
+                [<a href="https://github.com/omer11a/bounded-attention" target="_blank">GitHub</a>]
+            </span>
+        </p>
+        """
+        gr.HTML(description)
+        with gr.Column():
+            prompt = gr.Textbox(
+                label="Text prompt",
+            )
+
+            subject_token_indices = gr.Textbox(
+                label="The token indices of each subject (separate indices for the same subject with commas, and between different subjects with semicolons)",
+            )
+
+            filter_token_indices = gr.Textbox(
+                label="The token indices to filter, i.e. conjunctions, number, postional relations, etc. (if left empty, this will be automatically inferred)",
+            )
+
+            num_tokens = gr.Textbox(
+                label="The number of tokens in the prompt (can be left empty, but we recommend filling this, so we can verify your input, as sometimes rare words are split into more than one token)",
+            )
+
+            with gr.Row():
+                sketch_pad = gr.Sketchpad(label="Sketch Pad", shape=(RESOLUTION, RESOLUTION))
+                layout_image = gr.Image(type="pil", label="Bounding Boxes")
+                out_images = gr.Image(type="pil", visible=True, label="Generated Image")
+
+            with gr.Row():
+                clear_button = gr.Button(value='Clear')
+                generate_button = gr.Button(value='Generate')
+
+            with gr.Accordion("Advanced Options", open=False):
+                with gr.Column():
+                    description = """
+                        <div class="tooltip">Batch size &#9432
+                        <span class="tooltiptext">The number of images to generate.</span>
+                        </div>
+                        <div class="tooltip">Initial step size &#9432
+                        <span class="tooltiptext">The initial step size of the linear step size scheduler when performing guidance.</span>
+                        </div>
+                        <div class="tooltip">Final step size &#9432
+                        <span class="tooltiptext">The final step size of the linear step size scheduler when performing guidance.</span>
+                        </div>
+                        <div class="tooltip">Number of self-attention clusters per subject &#9432
+                        <span class="tooltiptext">Determines the number of clusters when clustering the self-attention maps (#clusters = #subject x #clusters_per_subject). Changing this value might improve semantics (adherence to the prompt), especially when the subjects exceed their bounding boxes.</span>
+                        </div>
+                        <div class="tooltip">Cross-attention loss scale factor &#9432
+                        <span class="tooltiptext">The scale factor of the cross-attention loss term. Increasing it will improve semantic control (adherence to the prompt), but may reduce image quality.</span>
+                        </div>
+                        <div class="tooltip">Self-attention loss scale factor &#9432
+                        <span class="tooltiptext">The scale factor of the self-attention loss term. Increasing it will improve layout control (adherence to the bounding boxes), but may reduce image quality.</span>
+                        </div>
+                        <div class="tooltip">Classifier-free guidance scale &#9432
+                        <span class="tooltiptext">The scale factor of classifier-free guidance.</span>
+                        </div>
+                        <div class="tooltip" >Number of Gradient Descent iterations per timestep &#9432
+                        <span class="tooltiptext">The number of Gradient Descent iterations for each timestep when performing guidance.</span>
+                        </div>
+                        <div class="tooltip" >Loss Threshold &#9432
+                        <span class="tooltiptext">If the loss is below the threshold, Gradient Descent stops for that timestep. </span>
+                        </div>
+                        <div class="tooltip" >Number of guidance steps &#9432
+                        <span class="tooltiptext">The number of timesteps in which to perform guidance.</span>
+                        </div>
+                    """
+                    gr.HTML(description)
+                    batch_size = gr.Slider(minimum=1, maximum=5, step=1, value=1, label="Number of samples")
+                    init_step_size = gr.Slider(minimum=0, maximum=50, step=0.5, value=25, label="Initial step size")
+                    final_step_size = gr.Slider(minimum=0, maximum=20, step=0.5, value=10, label="Final step size")
+                    num_clusters_per_subject = gr.Slider(minimum=0, maximum=5, step=0.5, value=3, label="Number of clusters per subject")
+                    cross_loss_scale = gr.Slider(minimum=0, maximum=2, step=0.1, value=1, label="Cross-attention loss scale factor")
+                    self_loss_scale = gr.Slider(minimum=0, maximum=2, step=0.1, value=1, label="Self-attention loss scale factor")
+                    classifier_free_guidance_scale = gr.Slider(minimum=0, maximum=50, step=0.5, value=7.5, label="Classifier-free guidance Scale")
+                    num_iterations = gr.Slider(minimum=0, maximum=10, step=1, value=5, label="Number of Gradient Descent iterations")
+                    loss_threshold = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.2, label="Loss threshold")
+                    num_guidance_steps = gr.Slider(minimum=10, maximum=20, step=1, value=15, label="Number of timesteps to perform guidance")
+                    seed = gr.Slider(minimum=0, maximum=1000, step=1, value=445, label="Random Seed")
+
+            boxes = gr.State([])
+
+            demo.load(
+                clear,
+                inputs=[batch_size],
+                outputs=[boxes, sketch_pad, layout_image, out_images],
+                queue=False
+            )
+
+            sketch_pad.edit(
+                draw,
+                inputs=[boxes, sketch_pad],
+                outputs=[boxes, sketch_pad, layout_image],
+                queue=False,
+            )
+
+            clear_button.click(
+                clear,
+                inputs=[batch_size],
+                outputs=[boxes, sketch_pad, layout_image, out_images],
+                queue=False,
+            )
+
+            generate_button.click(
+                fn=partial(generate, device, model),
+                inputs=[
+                    prompt, subject_token_indices, filter_token_indices, num_tokens,
+                    init_step_size, final_step_size, num_clusters_per_subject, cross_loss_scale, self_loss_scale,
+                    classifier_free_guidance_scale, batch_size, num_iterations, loss_threshold, num_guidance_steps,
+                    seed,
+                    boxes,
+                ],
+                outputs=[out_images],
+                queue=True,
+            )
+
+        #with gr.Column():
+        #    gr.Examples(
+        #        examples=[
+        #            [
+        #                [[0.35, 0.4, 0.65, 0.9], [0, 0.6, 0.3, 0.9], [0.7, 0.55, 1, 0.85]],
+        #                "3D Pixar animation of a cute unicorn and a pink hedgehog and a nerdy owl traveling in a magical forest",
+        #                "7,8,17;11,12,17;15,16,17",
+        #                "5,6,9,10,13,14,18,19",
+        #                286,
+        #            ],
+        #        ],
+        #        inputs=[boxes, prompt, subject_token_indices, filter_token_indices, seed],
+        #        outputs=None,
+        #        fn=None,
+        #        cache_examples=False,
+        #    )
+        description = """<p> The source code of this demo is based on the <a href="https://huggingface.co/spaces/gligen/demo/tree/main">GLIGEN demo</a>.</p>"""
+        gr.HTML(description)
+
+    demo.queue(concurrency_count=1, api_open=False)
+    demo.launch(share=False, show_api=False, show_error=True)
+
+if __name__ == '__main__':
+    main()

bounded_attention.py

@@ -0,0 +1,522 @@
+import nltk
+import einops
+import torch
+import torch.nn.functional as F
+import torchvision.utils
+from torch_kmeans import KMeans
+
+import os
+
+import injection_utils
+import utils
+
+
+class BoundedAttention(injection_utils.AttentionBase):
+    EPSILON = 1e-5
+    FILTER_TAGS = {
+        'CC', 'CD', 'DT', 'EX', 'IN', 'LS', 'MD', 'PDT', 'POS', 'PRP', 'PRP$', 'RP', 'TO', 'UH', 'WDT', 'WP', 'WRB'}
+    TAG_RULES = {'left': 'IN', 'right': 'IN', 'top': 'IN', 'bottom': 'IN'}
+
+    def __init__(
+        self,
+        boxes,
+        prompts,
+        subject_token_indices,
+        cross_loss_layers,
+        self_loss_layers,
+        cross_mask_layers=None,
+        self_mask_layers=None,
+        eos_token_index=None,
+        filter_token_indices=None,
+        leading_token_indices=None,
+        mask_cross_during_guidance=True,
+        mask_eos=True,
+        cross_loss_coef=1,
+        self_loss_coef=1,
+        max_guidance_iter=15,
+        max_guidance_iter_per_step=5,
+        start_step_size=30,
+        end_step_size=10,
+        loss_stopping_value=0.2,
+        cross_mask_threshold=0.2,
+        self_mask_threshold=0.2,
+        delta_refine_mask_steps=5,
+        pca_rank=None,
+        num_clusters=None,
+        num_clusters_per_box=3,
+        map_dir=None,
+        debug=False,
+        delta_debug_attention_steps=20,
+        delta_debug_mask_steps=5,
+        debug_layers=None,
+        saved_resolution=64,
+    ):
+        super().__init__()
+        self.boxes = boxes
+        self.prompts = prompts
+        self.subject_token_indices = subject_token_indices
+        self.cross_loss_layers = set(cross_loss_layers)
+        self.self_loss_layers = set(self_loss_layers)
+        self.cross_mask_layers = self.cross_loss_layers if cross_mask_layers is None else set(cross_mask_layers)
+        self.self_mask_layers = self.self_loss_layers if self_mask_layers is None else set(self_mask_layers)
+
+        self.eos_token_index = eos_token_index
+        self.filter_token_indices = filter_token_indices
+        self.leading_token_indices = leading_token_indices
+        self.mask_cross_during_guidance = mask_cross_during_guidance
+        self.mask_eos = mask_eos
+        self.cross_loss_coef = cross_loss_coef
+        self.self_loss_coef = self_loss_coef
+        self.max_guidance_iter = max_guidance_iter
+        self.max_guidance_iter_per_step = max_guidance_iter_per_step
+        self.start_step_size = start_step_size
+        self.step_size_coef = (end_step_size - start_step_size) / max_guidance_iter
+        self.loss_stopping_value = loss_stopping_value
+        self.cross_mask_threshold = cross_mask_threshold
+        self.self_mask_threshold = self_mask_threshold
+
+        self.delta_refine_mask_steps = delta_refine_mask_steps
+        self.pca_rank = pca_rank
+        num_clusters = len(boxes) * num_clusters_per_box if num_clusters is None else num_clusters
+        self.clustering = KMeans(n_clusters=num_clusters, num_init=100)
+        self.centers = None
+
+        self.map_dir = map_dir
+        self.debug = debug
+        self.delta_debug_attention_steps = delta_debug_attention_steps
+        self.delta_debug_mask_steps = delta_debug_mask_steps
+        self.debug_layers = self.cross_loss_layers | self.self_loss_layers if debug_layers is None else debug_layers
+        self.saved_resolution = saved_resolution
+
+        self.optimized = False
+        self.cross_foreground_values = []
+        self.self_foreground_values = []
+        self.cross_background_values = []
+        self.self_background_values = []
+        self.cross_maps = []
+        self.self_maps = []
+        self.self_masks = None
+
+    def clear_values(self, include_maps=False):
+        lists = (
+            self.cross_foreground_values,
+            self.self_foreground_values,
+            self.cross_background_values,
+            self.self_background_values,
+        )
+
+        if include_maps:
+            lists = (
+                *all_values,
+                self.cross_maps,
+                self.self_maps,
+            )
+
+        for values in lists:
+            values.clear()
+
+    def before_step(self):
+        self.clear_values()
+        if self.cur_step == 0:
+            self._determine_tokens()
+
+    def reset(self):
+        self.clear_values(include_maps=True)
+        super().reset()
+
+    def forward(self, q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs):
+        self._display_attention_maps(attn, is_cross, num_heads)
+
+        _, n, d = sim.shape
+        sim_u, sim_c = sim.reshape(-1, num_heads, n, d).chunk(2)  # b h n d
+        if is_cross:
+            sim_c = self._hide_other_subjects_from_tokens(sim_c)
+        else:
+            sim_u = self._hide_other_subjects_from_subjects(sim_u)
+            sim_c = self._hide_other_subjects_from_subjects(sim_c)
+
+        sim = torch.cat((sim_u, sim_c)).reshape(-1, n, d)
+        attn = sim.softmax(-1)
+        self._save(attn, is_cross, num_heads)
+        self._display_attention_maps(attn, is_cross, num_heads, prefix='masked')
+        self._debug_hook(q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs)
+
+        out = torch.bmm(attn, v)
+        out = einops.rearrange(out, '(b h) n d -> b n (h d)', h=num_heads)
+        return out
+
+    def update_loss(self, forward_pass, latents, i):
+        if i >= self.max_guidance_iter:
+            return latents
+
+        step_size = self.start_step_size + self.step_size_coef * i
+        updated_latents = latents
+
+        self.optimized = True
+        normalized_loss = torch.tensor(10000)
+        with torch.enable_grad():
+            latents = updated_latents = updated_latents.clone().detach().requires_grad_(True)
+            for guidance_iter in range(self.max_guidance_iter_per_step):
+                if normalized_loss < self.loss_stopping_value:
+                    break
+
+                latent_model_input = torch.cat([latents] * 2)
+                cur_step = self.cur_step
+                forward_pass(latent_model_input)
+                self.cur_step = cur_step
+
+                loss, normalized_loss = self._compute_loss()
+                grad_cond = torch.autograd.grad(loss, [updated_latents])[0]
+                latents = updated_latents = updated_latents - step_size * grad_cond
+                if self.debug:
+                    print(f'Loss at step={i}, iter={guidance_iter}: {normalized_loss}')
+                    grad_norms = grad_cond.flatten(start_dim=2).norm(dim=1)
+                    grad_norms = grad_norms / grad_norms.max(dim=1, keepdim=True)[0]
+                    self._save_maps(grad_norms, 'grad_norms')
+
+        self.optimized = False
+        return latents
+
+    def _tokenize(self):
+        ids = self.model.tokenizer.encode(self.prompts[0])
+        tokens = self.model.tokenizer.convert_ids_to_tokens(ids, skip_special_tokens=True)
+        return [token[:-4] for token in tokens]  # remove ending </w>
+
+    def _tag_tokens(self):
+        tagged_tokens = nltk.pos_tag(self._tokenize())
+        return [type(self).TAG_RULES.get(token, tag) for token, tag in tagged_tokens]
+
+    def _determine_eos_token(self):
+        tokens = self._tokenize()
+        eos_token_index = len(tokens) + 1
+        if self.eos_token_index is None:
+            self.eos_token_index = eos_token_index
+        elif eos_token_index != self.eos_token_index:
+            raise ValueError(f'Wrong EOS token index. Tokens are: {tokens}.')
+
+    def _determine_filter_tokens(self):
+        if self.filter_token_indices is not None:
+            return
+
+        tags = self._tag_tokens()
+        self.filter_token_indices = [i + 1 for i, tag in enumerate(tags) if tag in type(self).FILTER_TAGS]
+
+    def _determine_leading_tokens(self):
+        if self.leading_token_indices is not None:
+            return
+
+        tags = self._tag_tokens()
+        leading_token_indices = []
+        for indices in self.subject_token_indices:
+            subject_noun_indices = [i for i in indices if tags[i - 1].startswith('NN')]
+            leading_token_candidates = subject_noun_indices if len(subject_noun_indices) > 0 else indices
+            leading_token_indices.append(leading_token_candidates[-1])
+
+        self.leading_token_indices = leading_token_indices
+
+    def _determine_tokens(self):
+        self._determine_eos_token()
+        self._determine_filter_tokens()
+        self._determine_leading_tokens()
+
+    def _split_references(self, tensor, num_heads):
+        tensor = tensor.reshape(-1, num_heads, *tensor.shape[1:])
+        unconditional, conditional = tensor.chunk(2)
+
+        num_subjects = len(self.boxes)
+        batch_unconditional = unconditional[:-num_subjects]
+        references_unconditional = unconditional[-num_subjects:]
+        batch_conditional = conditional[:-num_subjects]
+        references_conditional = conditional[-num_subjects:]
+
+        batch = torch.cat((batch_unconditional, batch_conditional))
+        references = torch.cat((references_unconditional, references_conditional))
+        batch = batch.reshape(-1, *batch_unconditional.shape[2:])
+        references = references.reshape(-1, *references_unconditional.shape[2:])
+        return batch, references
+
+    def _hide_other_subjects_from_tokens(self, sim):  # b h i j
+        dtype = sim.dtype
+        device = sim.device
+        batch_size = sim.size(0)
+        resolution = int(sim.size(2) ** 0.5)
+        subject_masks, background_masks = self._obtain_masks(resolution, batch_size=batch_size, device=device)  # b s n
+        include_background = self.optimized or (not self.mask_cross_during_guidance and self.cur_step < self.max_guidance_iter_per_step)
+        subject_masks = torch.logical_or(subject_masks, background_masks.unsqueeze(1)) if include_background else subject_masks
+        min_value = torch.finfo(sim.dtype).min
+        sim_masks = torch.zeros_like(sim[:, 0, :, :])  # b i j
+        for token_indices in (*self.subject_token_indices, self.filter_token_indices):
+            sim_masks[:, :, token_indices] = min_value
+
+        for batch_index in range(batch_size):
+            for subject_mask, token_indices in zip(subject_masks[batch_index], self.subject_token_indices):
+                for token_index in token_indices:
+                    sim_masks[batch_index, subject_mask, token_index] = 0
+
+        if self.mask_eos and not include_background:
+            for batch_index, background_mask in zip(range(batch_size), background_masks):
+                sim_masks[batch_index, background_mask, self.eos_token_index] = min_value
+
+        return sim + sim_masks.unsqueeze(1)
+
+    def _hide_other_subjects_from_subjects(self, sim):  # b h i j
+        dtype = sim.dtype
+        device = sim.device
+        batch_size = sim.size(0)
+        resolution = int(sim.size(2) ** 0.5)
+        subject_masks, background_masks = self._obtain_masks(resolution, batch_size=batch_size, device=device)  # b s n
+        min_value = torch.finfo(dtype).min
+        sim_masks = torch.zeros_like(sim[:, 0, :, :])  # b i j
+        for batch_index, background_mask in zip(range(batch_size), background_masks):
+            sim_masks[batch_index, ~background_mask, ~background_mask] = min_value
+
+        for batch_index in range(batch_size):
+            for subject_mask in subject_masks[batch_index]:
+                subject_sim_mask = sim_masks[batch_index, subject_mask]
+                condition = torch.logical_or(subject_sim_mask == 0, subject_mask.unsqueeze(0))
+                sim_masks[batch_index, subject_mask] = torch.where(condition, 0, min_value).to(dtype=dtype)
+
+        return sim + sim_masks.unsqueeze(1)
+
+    def _save(self, attn, is_cross, num_heads):
+        _, attn = attn.chunk(2)
+        attn = attn.reshape(-1, num_heads, *attn.shape[-2:])  # b h n k
+
+        self._save_mask_maps(attn, is_cross)
+        self._save_loss_values(attn, is_cross)
+
+    def _save_mask_maps(self, attn, is_cross):
+        if (
+            (self.optimized) or
+            (is_cross and self.cur_att_layer not in self.cross_mask_layers) or
+            ((not is_cross) and (self.cur_att_layer not in self.self_mask_layers))
+        ):
+            return
+
+        if is_cross:
+            attn = attn[..., self.leading_token_indices]
+            mask_maps = self.cross_maps
+        else:
+            mask_maps = self.self_maps
+
+        mask_maps.append(attn.mean(dim=1))  # mean over heads
+        if self.cur_step > 0:
+            mask_maps = mask_maps[1:]  # throw away old maps
+
+    def _save_loss_values(self, attn, is_cross):
+        if (
+            (not self.optimized) or
+            (is_cross and (self.cur_att_layer not in self.cross_loss_layers)) or
+            ((not is_cross) and (self.cur_att_layer not in self.self_loss_layers))
+        ):
+            return
+
+        resolution = int(attn.size(2) ** 0.5)
+        boxes = self._convert_boxes_to_masks(resolution, device=attn.device)  # s n
+        background_mask = boxes.sum(dim=0) == 0
+
+        if is_cross:
+            saved_foreground_values = self.cross_foreground_values
+            saved_background_values = self.cross_background_values
+            contexts = [indices + [self.eos_token_index] for indices in self.subject_token_indices]  # TODO: fix EOS loss term
+        else:
+            saved_foreground_values = self.self_foreground_values
+            saved_background_values = self.self_background_values
+            contexts = boxes
+
+        foreground_values = []
+        background_values = []
+        for i, (box, context) in enumerate(zip(boxes, contexts)):
+            context_attn = attn[:, :, :, context]
+            
+            # sum over heads, pixels and contexts
+            foreground_values.append(context_attn[:, :, box].sum(dim=(1, 2, 3)))
+            background_values.append(context_attn[:, :, background_mask].sum(dim=(1, 2, 3)))
+
+        saved_foreground_values.append(torch.stack(foreground_values, dim=1))
+        saved_background_values.append(torch.stack(background_values, dim=1))
+
+    def _compute_loss(self):
+        cross_losses = self._compute_loss_term(self.cross_foreground_values, self.cross_background_values)
+        self_losses = self._compute_loss_term(self.self_foreground_values, self.self_background_values)
+        b, s = cross_losses.shape
+
+        # sum over samples and subjects
+        total_cross_loss = cross_losses.sum()
+        total_self_loss = self_losses.sum()
+
+        loss = self.cross_loss_coef * total_cross_loss + self.self_loss_coef * total_self_loss
+        normalized_loss = loss / b / s
+        return loss, normalized_loss
+
+    def _compute_loss_term(self, foreground_values, background_values):
+        # mean over layers
+        mean_foreground = torch.stack(foreground_values).mean(dim=0)
+        mean_background = torch.stack(background_values).mean(dim=0)
+        iou = mean_foreground / (mean_foreground + len(self.boxes) * mean_background)
+        return (1 - iou) ** 2
+
+    def _obtain_masks(self, resolution, return_boxes=False, return_existing=False, batch_size=None, device=None):
+        return_boxes = return_boxes or (return_existing and self.self_masks is None)
+        if return_boxes or self.cur_step < self.max_guidance_iter:
+            masks = self._convert_boxes_to_masks(resolution, device=device).unsqueeze(0)
+            if batch_size is not None:
+                masks = masks.expand(batch_size, *masks.shape[1:])
+        else:
+            masks = self._obtain_self_masks(resolution, return_existing=return_existing)
+            if device is not None:
+                masks = masks.to(device=device)
+
+        background_mask = masks.sum(dim=1) == 0
+        return masks, background_mask
+
+    def _convert_boxes_to_masks(self, resolution, device=None):  # s n
+        boxes = torch.zeros(len(self.boxes), resolution, resolution, dtype=bool, device=device)
+        for i, box in enumerate(self.boxes):
+            x0, x1 = box[0] * resolution, box[2] * resolution
+            y0, y1 = box[1] * resolution, box[3] * resolution
+
+            boxes[i, round(y0) : round(y1), round(x0) : round(x1)] = True
+
+        return boxes.flatten(start_dim=1)
+
+    def _obtain_self_masks(self, resolution, return_existing=False):
+        if (
+            (self.self_masks is None) or
+            (
+                (self.cur_step % self.delta_refine_mask_steps == 0) and
+                (self.cur_att_layer == 0) and
+                (not return_existing)
+            )
+        ):
+            self.self_masks = self._fix_zero_masks(self._build_self_masks())
+
+        b, s, n = self.self_masks.shape
+        mask_resolution = int(n ** 0.5)
+        self_masks = self.self_masks.reshape(b, s, mask_resolution, mask_resolution).float()
+        self_masks = F.interpolate(self_masks, resolution, mode='nearest-exact')
+        return self_masks.flatten(start_dim=2).bool()
+
+    def _cluster_self_maps(self):  # b s n
+        self_maps = self._aggregate_maps(self.self_maps)  # b n m
+        if self.pca_rank is not None:
+            dtype = self_maps.dtype
+            _, _, eigen_vectors = torch.pca_lowrank(self_maps.float(), self.pca_rank)
+            self_maps = torch.matmul(self_maps, eigen_vectors.to(dtype=dtype))
+
+        clustering_results = self.clustering(self_maps, centers=self.centers)
+        self.clustering.num_init = 1  # clustering is deterministic after the first time
+        self.centers = clustering_results.centers
+        clusters = clustering_results.labels
+        num_clusters = self.clustering.n_clusters
+        self._save_maps(clusters / num_clusters, f'clusters')
+        return num_clusters, clusters
+
+    def _build_self_masks(self):
+        c, clusters = self._cluster_self_maps()  # b n
+        cluster_masks = torch.stack([(clusters == cluster_index) for cluster_index in range(c)], dim=2)  # b n c
+        cluster_area = cluster_masks.sum(dim=1, keepdim=True)  # b 1 c
+
+        n = clusters.size(1)
+        resolution = int(n ** 0.5)
+        cross_masks = self._obtain_cross_masks(resolution)  # b s n
+        cross_mask_area = cross_masks.sum(dim=2, keepdim=True)  # b s 1
+
+        intersection = torch.bmm(cross_masks.float(), cluster_masks.float())  # b s c
+        min_area = torch.minimum(cross_mask_area, cluster_area)  # b s c
+        score_per_cluster, subject_per_cluster = torch.max(intersection / min_area, dim=1)  # b c
+        subjects = torch.gather(subject_per_cluster, 1, clusters)  # b n
+        scores = torch.gather(score_per_cluster, 1, clusters)  # b n
+
+        s = cross_masks.size(1)
+        self_masks = torch.stack([(subjects == subject_index) for subject_index in range(s)], dim=1)  # b s n
+        scores = scores.unsqueeze(1).expand(-1 ,s, n)  # b s n
+        self_masks[scores < self.self_mask_threshold] = False
+        self._save_maps(self_masks, 'self_masks')
+        return self_masks
+
+    def _obtain_cross_masks(self, resolution, scale=10):
+        maps = self._aggregate_maps(self.cross_maps, resolution=resolution)  # b n k
+        maps = F.sigmoid(scale * (maps - self.cross_mask_threshold))
+        maps = self._normalize_maps(maps, reduce_min=True)
+        maps = maps.transpose(1, 2)  # b k n
+        existing_masks, _ = self._obtain_masks(
+            resolution, return_existing=True, batch_size=maps.size(0), device=maps.device)
+        maps = maps * existing_masks.to(dtype=maps.dtype)
+        self._save_maps(maps, 'cross_masks')
+        return maps
+
+    def _fix_zero_masks(self, masks):
+        b, s, n = masks.shape
+        resolution = int(n ** 0.5)
+        boxes = self._convert_boxes_to_masks(resolution, device=masks.device)  # s n
+
+        for i in range(b):
+            for j in range(s):
+                if masks[i, j].sum() == 0:
+                    print('******Found a zero mask!******')
+                    for k in range(s):
+                        masks[i, k] = boxes[j] if (k == j) else masks[i, k].logical_and(~boxes[j])
+
+        return masks
+
+    def _aggregate_maps(self, maps, resolution=None):  # b n k
+        maps = torch.stack(maps).mean(0)  # mean over layers
+        if resolution is not None:
+            b, n, k = maps.shape
+            original_resolution = int(n ** 0.5)
+            maps = maps.transpose(1, 2).reshape(b, k, original_resolution, original_resolution)
+            maps = F.interpolate(maps, resolution, mode='bilinear', antialias=True)
+            maps = maps.reshape(b, k, -1).transpose(1, 2)
+
+        maps = self._normalize_maps(maps)
+        return maps
+
+    @classmethod
+    def _normalize_maps(cls, maps, reduce_min=False):  # b n k
+        max_values = maps.max(dim=1, keepdim=True)[0]
+        min_values = maps.min(dim=1, keepdim=True)[0] if reduce_min else 0
+        numerator = maps - min_values
+        denominator = max_values - min_values + cls.EPSILON
+        return numerator / denominator
+
+    def _save_maps(self, maps, prefix):
+        if self.map_dir is None or self.cur_step % self.delta_debug_mask_steps != 0:
+            return
+
+        resolution = int(maps.size(-1) ** 0.5)
+        maps = maps.reshape(-1, 1, resolution, resolution).float()
+        maps = F.interpolate(maps, self.saved_resolution, mode='bilinear', antialias=True)
+        path = os.path.join(self.map_dir, f'map_{prefix}_{self.cur_step}_{self.cur_att_layer}.png')
+        torchvision.utils.save_image(maps, path)
+
+    def _display_attention_maps(self, attention_maps, is_cross, num_heads, prefix=None):
+        if (not self.debug) or (self.cur_step == 0) or (self.cur_step % self.delta_debug_attention_steps > 0) or (self.cur_att_layer not in self.debug_layers):
+            return
+
+        dir_name = self.map_dir
+        if prefix is not None:
+            splits = list(os.path.split(dir_name))
+            splits[-1] = '_'.join((prefix, splits[-1]))
+            dir_name = os.path.join(*splits)
+
+        resolution = int(attention_maps.size(-2) ** 0.5)
+        if is_cross:
+            attention_maps = einops.rearrange(attention_maps, 'b (r1 r2) k -> b k r1 r2', r1=resolution)
+            attention_maps = F.interpolate(attention_maps, self.saved_resolution, mode='bilinear', antialias=True)
+            attention_maps = einops.rearrange(attention_maps, 'b k r1 r2 -> b (r1 r2) k')
+
+        utils.display_attention_maps(
+            attention_maps,
+            is_cross,
+            num_heads,
+            self.model.tokenizer,
+            self.prompts,
+            dir_name,
+            self.cur_step,
+            self.cur_att_layer,
+            resolution,
+        )
+
+    def _debug_hook(self, q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs):
+        pass

injection_utils.py

@@ -0,0 +1,234 @@
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from typing import Optional, Union, Tuple, List, Callable, Dict
+
+from torchvision.utils import save_image
+from einops import rearrange, repeat
+
+
+class AttentionBase:
+    def __init__(self):
+        self.cur_step = 0
+        self.num_att_layers = -1
+        self.cur_att_layer = 0
+
+    def before_step(self):
+        pass
+
+    def after_step(self):
+        pass
+
+    def __call__(self, q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs):
+        if self.cur_att_layer == 0:
+            self.before_step()
+
+        out = self.forward(q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs)
+        self.cur_att_layer += 1
+        if self.cur_att_layer == self.num_att_layers:
+            self.cur_att_layer = 0
+            self.cur_step += 1
+            # after step
+            self.after_step()
+        return out
+
+    def forward(self, q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs):
+        out = torch.einsum('b i j, b j d -> b i d', attn, v)
+        out = rearrange(out, '(b h) n d -> b n (h d)', h=num_heads)
+        return out
+
+    def reset(self):
+        self.cur_step = 0
+        self.cur_att_layer = 0
+
+
+class AttentionStore(AttentionBase):
+    def __init__(self, res=[32], min_step=0, max_step=1000):
+        super().__init__()
+        self.res = res
+        self.min_step = min_step
+        self.max_step = max_step
+        self.valid_steps = 0
+
+        self.self_attns = []  # store the all attns
+        self.cross_attns = []
+
+        self.self_attns_step = []  # store the attns in each step
+        self.cross_attns_step = []
+
+    def after_step(self):
+        if self.cur_step > self.min_step and self.cur_step < self.max_step:
+            self.valid_steps += 1
+            if len(self.self_attns) == 0:
+                self.self_attns = self.self_attns_step
+                self.cross_attns = self.cross_attns_step
+            else:
+                for i in range(len(self.self_attns)):
+                    self.self_attns[i] += self.self_attns_step[i]
+                    self.cross_attns[i] += self.cross_attns_step[i]
+        self.self_attns_step.clear()
+        self.cross_attns_step.clear()
+
+    def forward(self, q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs):
+        if attn.shape[1] <= 64 ** 2:  # avoid OOM
+            if is_cross:
+                self.cross_attns_step.append(attn)
+            else:
+                self.self_attns_step.append(attn)
+        return super().forward(q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs)
+
+
+def regiter_attention_editor_diffusers(model, editor: AttentionBase):
+    """
+    Register a attention editor to Diffuser Pipeline, refer from [Prompt-to-Prompt]
+    """
+    def ca_forward(self, place_in_unet):
+        def forward(x, encoder_hidden_states=None, attention_mask=None, context=None, mask=None):
+            """
+            The attention is similar to the original implementation of LDM CrossAttention class
+            except adding some modifications on the attention
+            """
+            if encoder_hidden_states is not None:
+                context = encoder_hidden_states
+            if attention_mask is not None:
+                mask = attention_mask
+
+            to_out = self.to_out
+            if isinstance(to_out, nn.modules.container.ModuleList):
+                to_out = self.to_out[0]
+            else:
+                to_out = self.to_out
+
+            h = self.heads
+            q = self.to_q(x)
+            is_cross = context is not None
+            context = context if is_cross else x
+            k = self.to_k(context)
+            v = self.to_v(context)
+            q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+
+            sim = torch.einsum('b i d, b j d -> b i j', q, k) * self.scale
+
+            if mask is not None:
+                mask = rearrange(mask, 'b ... -> b (...)')
+                max_neg_value = -torch.finfo(sim.dtype).max
+                mask = repeat(mask, 'b j -> (b h) () j', h=h)
+                mask = mask[:, None, :].repeat(h, 1, 1)
+                sim.masked_fill_(~mask, max_neg_value)
+
+            attn = sim.softmax(dim=-1)
+            # the only difference
+            out = editor(
+                q, k, v, sim, attn, is_cross, place_in_unet,
+                self.heads, scale=self.scale)
+
+            return to_out(out)
+
+        return forward
+
+    #def conv_forward(self):
+    #    def forward(x):
+    #        return editor.conv(self, x)
+
+    #    return forward
+
+    def register_editor(net, count, place_in_unet, prefix=''):
+        #print(prefix + f'Found net: {net.__class__.__name__}')
+        #if 'Conv' in net.__class__.__name__:
+        #    print(prefix + f'Found conv with kernel size: {net.kernel_size}')
+        #    net.old_forward = net.forward
+        #    net.forward = conv_forward(net)
+        for name, subnet in net.named_children():
+            if net.__class__.__name__ == 'Attention':  # spatial Transformer layer
+                net.forward = ca_forward(net, place_in_unet)
+                return count + 1
+            elif hasattr(net, 'children'):
+                count = register_editor(subnet, count, place_in_unet, prefix=prefix + '\t')
+        return count
+
+    cross_att_count = 0
+    for net_name, net in model.unet.named_children():
+        if "down" in net_name:
+            cross_att_count += register_editor(net, 0, "down")
+            #print(f'Down number of attention layers {cross_att_count}!')
+        elif "mid" in net_name:
+            cross_att_count += register_editor(net, 0, "mid")
+            #print(f'Mid number of attention layers {cross_att_count}!')
+        elif "up" in net_name:
+            cross_att_count += register_editor(net, 0, "up")
+            #print(f'Up number of attention layers {cross_att_count}!')
+    print(f'Number of attention layers {cross_att_count}!')
+    editor.num_att_layers = cross_att_count
+    editor.model = model
+    model.editor = editor
+
+
+def regiter_attention_editor_ldm(model, editor: AttentionBase):
+    """
+    Register a attention editor to Stable Diffusion model, refer from [Prompt-to-Prompt]
+    """
+    def ca_forward(self, place_in_unet):
+        def forward(x, encoder_hidden_states=None, attention_mask=None, context=None, mask=None):
+            """
+            The attention is similar to the original implementation of LDM CrossAttention class
+            except adding some modifications on the attention
+            """
+            if encoder_hidden_states is not None:
+                context = encoder_hidden_states
+            if attention_mask is not None:
+                mask = attention_mask
+
+            to_out = self.to_out
+            if isinstance(to_out, nn.modules.container.ModuleList):
+                to_out = self.to_out[0]
+            else:
+                to_out = self.to_out
+
+            h = self.heads
+            q = self.to_q(x)
+            is_cross = context is not None
+            context = context if is_cross else x
+            k = self.to_k(context)
+            v = self.to_v(context)
+            q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+
+            sim = torch.einsum('b i d, b j d -> b i j', q, k) * self.scale
+
+            if mask is not None:
+                mask = rearrange(mask, 'b ... -> b (...)')
+                max_neg_value = -torch.finfo(sim.dtype).max
+                mask = repeat(mask, 'b j -> (b h) () j', h=h)
+                mask = mask[:, None, :].repeat(h, 1, 1)
+                sim.masked_fill_(~mask, max_neg_value)
+
+            attn = sim.softmax(dim=-1)
+            # the only difference
+            out = editor(
+                q, k, v, sim, attn, is_cross, place_in_unet,
+                self.heads, scale=self.scale)
+
+            return to_out(out)
+
+        return forward
+
+    def register_editor(net, count, place_in_unet):
+        for name, subnet in net.named_children():
+            if net.__class__.__name__ == 'CrossAttention':  # spatial Transformer layer
+                net.forward = ca_forward(net, place_in_unet)
+                return count + 1
+            elif hasattr(net, 'children'):
+                count = register_editor(subnet, count, place_in_unet)
+        return count
+
+    cross_att_count = 0
+    for net_name, net in model.model.diffusion_model.named_children():
+        if "input" in net_name:
+            cross_att_count += register_editor(net, 0, "input")
+        elif "middle" in net_name:
+            cross_att_count += register_editor(net, 0, "middle")
+        elif "output" in net_name:
+            cross_att_count += register_editor(net, 0, "output")
+    editor.num_att_layers = cross_att_count

pipeline_stable_diffusion_xl_opt.py

@@ -0,0 +1,968 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import os
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL, UNet2DConditionModel
+from diffusers.models.attention_processor import (
+    AttnProcessor2_0,
+    LoRAAttnProcessor2_0,
+    LoRAXFormersAttnProcessor,
+    XFormersAttnProcessor,
+)
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import (
+    is_accelerate_available,
+    is_accelerate_version,
+    is_invisible_watermark_available,
+    logging,
+    randn_tensor,
+    replace_example_docstring,
+)
+from diffusers.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.stable_diffusion_xl import StableDiffusionXLPipelineOutput
+
+
+if is_invisible_watermark_available():
+    from .watermark import StableDiffusionXLWatermarker
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import StableDiffusionXLPipeline
+
+        >>> pipe = StableDiffusionXLPipeline.from_pretrained(
+        ...     "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
+        ... )
+        >>> pipe = pipe.to("cuda")
+
+        >>> prompt = "a photo of an astronaut riding a horse on mars"
+        >>> image = pipe(prompt).images[0]
+        ```
+"""
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+    """
+    Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
+    Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
+    """
+    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    # rescale the results from guidance (fixes overexposure)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+    return noise_cfg
+
+
+class StableDiffusionXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoaderMixin):
+    r"""
+    Pipeline for text-to-image generation using Stable Diffusion XL.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+
+    In addition the pipeline inherits the following loading methods:
+        - *LoRA*: [`StableDiffusionXLPipeline.load_lora_weights`]
+        - *Ckpt*: [`loaders.FromSingleFileMixin.from_single_file`]
+
+    as well as the following saving methods:
+        - *LoRA*: [`loaders.StableDiffusionXLPipeline.save_lora_weights`]
+
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            Frozen text-encoder. Stable Diffusion XL uses the text portion of
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        text_encoder_2 ([` CLIPTextModelWithProjection`]):
+            Second frozen text-encoder. Stable Diffusion XL uses the text and pool portion of
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
+            specifically the
+            [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)
+            variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_2 (`CLIPTokenizer`):
+            Second Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
+        scheduler ([`SchedulerMixin`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
+            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
+    """
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        text_encoder_2: CLIPTextModelWithProjection,
+        tokenizer: CLIPTokenizer,
+        tokenizer_2: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        force_zeros_for_empty_prompt: bool = True,
+        add_watermarker: Optional[bool] = None,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            unet=unet,
+            scheduler=scheduler,
+        )
+        self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt)
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.default_sample_size = self.unet.config.sample_size
+
+        add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available()
+
+        if add_watermarker:
+            self.watermark = StableDiffusionXLWatermarker()
+        else:
+            self.watermark = None
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_slicing
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_slicing
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_tiling
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.vae.enable_tiling()
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_tiling
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    def enable_model_cpu_offload(self, gpu_id=0):
+        r"""
+        Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
+        to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
+        method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
+        `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
+        """
+        if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
+            from accelerate import cpu_offload_with_hook
+        else:
+            raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.")
+
+        device = torch.device(f"cuda:{gpu_id}")
+
+        if self.device.type != "cpu":
+            self.to("cpu", silence_dtype_warnings=True)
+            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
+
+        model_sequence = (
+            [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2]
+        )
+        model_sequence.extend([self.unet, self.vae])
+
+        hook = None
+        for cpu_offloaded_model in model_sequence:
+            _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook)
+
+        # We'll offload the last model manually.
+        self.final_offload_hook = hook
+
+    def encode_prompt(
+        self,
+        prompt: str,
+        prompt_2: Optional[str] = None,
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: Optional[str] = None,
+        negative_prompt_2: Optional[str] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        lora_scale: Optional[float] = None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in both text-encoders
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            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. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
+            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.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            lora_scale (`float`, *optional*):
+                A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+        """
+        device = device or self._execution_device
+
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, LoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+        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]
+
+        # Define tokenizers and text encoders
+        tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2]
+        text_encoders = (
+            [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2]
+        )
+
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            # textual inversion: procecss multi-vector tokens if necessary
+            prompt_embeds_list = []
+            prompts = [prompt, prompt_2]
+            for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
+                if isinstance(self, TextualInversionLoaderMixin):
+                    prompt = self.maybe_convert_prompt(prompt, tokenizer)
+
+                text_inputs = tokenizer(
+                    prompt,
+                    padding="max_length",
+                    max_length=tokenizer.model_max_length,
+                    truncation=True,
+                    return_tensors="pt",
+                )
+
+                text_input_ids = text_inputs.input_ids
+                untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+                if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                    text_input_ids, untruncated_ids
+                ):
+                    removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1])
+                    logger.warning(
+                        "The following part of your input was truncated because CLIP can only handle sequences up to"
+                        f" {tokenizer.model_max_length} tokens: {removed_text}"
+                    )
+
+                prompt_embeds = text_encoder(
+                    text_input_ids.to(device),
+                    output_hidden_states=True,
+                )
+
+                # We are only ALWAYS interested in the pooled output of the final text encoder
+                pooled_prompt_embeds = prompt_embeds[0]
+                ### TODO: remove
+                null_text_inputs = tokenizer(
+                    ['a realistic photo of an empty background'] * batch_size,
+                    padding="max_length",
+                    max_length=tokenizer.model_max_length,
+                    truncation=True,
+                    return_tensors="pt",
+                )
+                null_input_ids = null_text_inputs.input_ids
+                null_prompt_embeds = text_encoder(
+                    null_input_ids.to(device),
+                    output_hidden_states=True,
+                )
+                pooled_prompt_embeds = null_prompt_embeds[0]
+                ### TODO: remove
+                prompt_embeds = prompt_embeds.hidden_states[-2]
+
+                prompt_embeds_list.append(prompt_embeds)
+
+            prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
+
+        # get unconditional embeddings for classifier free guidance
+        zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt
+        if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:
+            negative_prompt_embeds = torch.zeros_like(prompt_embeds)
+            negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
+        elif do_classifier_free_guidance and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt_2 = negative_prompt_2 or negative_prompt
+
+            uncond_tokens: List[str]
+            if prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif isinstance(negative_prompt, str):
+                uncond_tokens = [negative_prompt, negative_prompt_2]
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+            else:
+                uncond_tokens = [negative_prompt, negative_prompt_2]
+
+            negative_prompt_embeds_list = []
+            for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders):
+                if isinstance(self, TextualInversionLoaderMixin):
+                    negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer)
+
+                max_length = prompt_embeds.shape[1]
+                uncond_input = tokenizer(
+                    negative_prompt,
+                    padding="max_length",
+                    max_length=max_length,
+                    truncation=True,
+                    return_tensors="pt",
+                )
+
+                negative_prompt_embeds = text_encoder(
+                    uncond_input.input_ids.to(device),
+                    output_hidden_states=True,
+                )
+                # We are only ALWAYS interested in the pooled output of the final text encoder
+                negative_pooled_prompt_embeds = negative_prompt_embeds[0]
+                negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
+
+                negative_prompt_embeds_list.append(negative_prompt_embeds)
+
+            negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
+
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
+            bs_embed * num_images_per_prompt, -1
+        )
+        if do_classifier_free_guidance:
+            negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
+                bs_embed * num_images_per_prompt, -1
+            )
+
+        return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        height,
+        width,
+        callback_steps,
+        negative_prompt=None,
+        negative_prompt_2=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        pooled_prompt_embeds=None,
+        negative_pooled_prompt_embeds=None,
+    ):
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_2 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+        if prompt_embeds is not None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
+            )
+
+        if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
+            )
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, dtype):
+        add_time_ids = list(original_size + crops_coords_top_left + target_size)
+
+        passed_add_embed_dim = (
+            self.unet.config.addition_time_embed_dim * len(add_time_ids) + self.text_encoder_2.config.projection_dim
+        )
+        expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
+
+        if expected_add_embed_dim != passed_add_embed_dim:
+            raise ValueError(
+                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
+            )
+
+        add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
+        return add_time_ids
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae
+    def upcast_vae(self):
+        dtype = self.vae.dtype
+        self.vae.to(dtype=torch.float32)
+        use_torch_2_0_or_xformers = isinstance(
+            self.vae.decoder.mid_block.attentions[0].processor,
+            (
+                AttnProcessor2_0,
+                XFormersAttnProcessor,
+                LoRAXFormersAttnProcessor,
+                LoRAAttnProcessor2_0,
+            ),
+        )
+        # if xformers or torch_2_0 is used attention block does not need
+        # to be in float32 which can save lots of memory
+        if use_torch_2_0_or_xformers:
+            self.vae.post_quant_conv.to(dtype)
+            self.vae.decoder.conv_in.to(dtype)
+            self.vae.decoder.mid_block.to(dtype)
+
+    def update_loss(self, latents, i, t, prompt_embeds, cross_attention_kwargs, add_text_embeds, add_time_ids):
+        def forward_pass(latent_model_input):
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+            added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+            self.unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=prompt_embeds,
+                cross_attention_kwargs=cross_attention_kwargs,
+                added_cond_kwargs=added_cond_kwargs,
+                return_dict=False,
+            )
+            self.unet.zero_grad()
+
+        return self.editor.update_loss(forward_pass, latents, i)
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        denoising_end: Optional[float] = None,
+        guidance_scale: float = 5.0,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: 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,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_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,
+        guidance_rescale: float = 0.0,
+        original_size: Optional[Tuple[int, int]] = None,
+        crops_coords_top_left: Tuple[int, int] = (0, 0),
+        target_size: Optional[Tuple[int, int]] = None,
+    ):
+        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.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in both text-encoders
+            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.
+            denoising_end (`float`, *optional*):
+                When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be
+                completed before it is intentionally prematurely terminated. As a result, the returned sample will
+                still retain a substantial amount of noise as determined by the discrete timesteps selected by the
+                scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a
+                "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image
+                Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output)
+            guidance_scale (`float`, *optional*, defaults to 5.0):
+                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. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
+            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.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled 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_xl.StableDiffusionXLPipelineOutput`] 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 `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            guidance_rescale (`float`, *optional*, defaults to 0.7):
+                Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
+                Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
+                [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
+                Guidance rescale factor should fix overexposure when using zero terminal SNR.
+            original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
+                `original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as
+                explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+            crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+                `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
+                `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
+                `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+            target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                For most cases, `target_size` should be set to the desired height and width of the generated image. If
+                not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in
+                section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
+            `tuple`. When returning a tuple, the first element is a list with the generated images.
+        """
+        # 0. Default height and width to unet
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        original_size = original_size or (height, width)
+        target_size = target_size or (height, width)
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            callback_steps,
+            negative_prompt,
+            negative_prompt_2,
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_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
+        text_encoder_lora_scale = (
+            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
+        )
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            lora_scale=text_encoder_lora_scale,
+        )
+
+        # 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.config.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)
+
+        # 7. Prepare added time ids & embeddings
+        add_text_embeds = pooled_prompt_embeds
+        add_time_ids = self._get_add_time_ids(
+            original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype
+        )
+
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
+            add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0)
+
+        prompt_embeds = prompt_embeds.to(device)
+        add_text_embeds = add_text_embeds.to(device)
+        add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
+
+        # 8. Denoising loop
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+        # 7.1 Apply denoising_end
+        if denoising_end is not None and type(denoising_end) == float and denoising_end > 0 and denoising_end < 1:
+            discrete_timestep_cutoff = int(
+                round(
+                    self.scheduler.config.num_train_timesteps
+                    - (denoising_end * self.scheduler.config.num_train_timesteps)
+                )
+            )
+            num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
+            timesteps = timesteps[:num_inference_steps]
+
+        latents = latents.half()
+        prompt_embeds = prompt_embeds.half()
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                latents = self.update_loss(latents, i, t, prompt_embeds, cross_attention_kwargs, add_text_embeds, add_time_ids)
+
+                # 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
+                added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # 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)
+
+                if do_classifier_free_guidance and guidance_rescale > 0.0:
+                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+                # 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)
+
+        # make sure the VAE is in float32 mode, as it overflows in float16
+        if self.vae.dtype == torch.float16 and self.vae.config.force_upcast:
+            self.upcast_vae()
+            latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
+
+        if not output_type == "latent":
+            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+        else:
+            image = latents
+            return StableDiffusionXLPipelineOutput(images=image)
+
+        # apply watermark if available
+        if self.watermark is not None:
+            image = self.watermark.apply_watermark(image)
+
+        image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # 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,)
+
+        return StableDiffusionXLPipelineOutput(images=image)
+
+    # Overrride to properly handle the loading and unloading of the additional text encoder.
+    def load_lora_weights(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs):
+        # We could have accessed the unet config from `lora_state_dict()` too. We pass
+        # it here explicitly to be able to tell that it's coming from an SDXL
+        # pipeline.
+        state_dict, network_alphas = self.lora_state_dict(
+            pretrained_model_name_or_path_or_dict,
+            unet_config=self.unet.config,
+            **kwargs,
+        )
+        self.load_lora_into_unet(state_dict, network_alphas=network_alphas, unet=self.unet)
+
+        text_encoder_state_dict = {k: v for k, v in state_dict.items() if "text_encoder." in k}
+        if len(text_encoder_state_dict) > 0:
+            self.load_lora_into_text_encoder(
+                text_encoder_state_dict,
+                network_alphas=network_alphas,
+                text_encoder=self.text_encoder,
+                prefix="text_encoder",
+                lora_scale=self.lora_scale,
+            )
+
+        text_encoder_2_state_dict = {k: v for k, v in state_dict.items() if "text_encoder_2." in k}
+        if len(text_encoder_2_state_dict) > 0:
+            self.load_lora_into_text_encoder(
+                text_encoder_2_state_dict,
+                network_alphas=network_alphas,
+                text_encoder=self.text_encoder_2,
+                prefix="text_encoder_2",
+                lora_scale=self.lora_scale,
+            )
+
+    @classmethod
+    def save_lora_weights(
+        self,
+        save_directory: Union[str, os.PathLike],
+        unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+        text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+        text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+        is_main_process: bool = True,
+        weight_name: str = None,
+        save_function: Callable = None,
+        safe_serialization: bool = True,
+    ):
+        state_dict = {}
+
+        def pack_weights(layers, prefix):
+            layers_weights = layers.state_dict() if isinstance(layers, torch.nn.Module) else layers
+            layers_state_dict = {f"{prefix}.{module_name}": param for module_name, param in layers_weights.items()}
+            return layers_state_dict
+
+        state_dict.update(pack_weights(unet_lora_layers, "unet"))
+
+        if text_encoder_lora_layers and text_encoder_2_lora_layers:
+            state_dict.update(pack_weights(text_encoder_lora_layers, "text_encoder"))
+            state_dict.update(pack_weights(text_encoder_2_lora_layers, "text_encoder_2"))
+
+        self.write_lora_layers(
+            state_dict=state_dict,
+            save_directory=save_directory,
+            is_main_process=is_main_process,
+            weight_name=weight_name,
+            save_function=save_function,
+            safe_serialization=safe_serialization,
+        )
+
+    def _remove_text_encoder_monkey_patch(self):
+        self._remove_text_encoder_monkey_patch_classmethod(self.text_encoder)
+        self._remove_text_encoder_monkey_patch_classmethod(self.text_encoder_2)

requirements.txt

@@ -0,0 +1,22 @@
+accelerate==0.21.0
+diffusers==0.20.0
+einops==0.6.1
+lightning-utilities==0.9.0
+matplotlib==3.7.3
+nltk==3.8.1
+numpy
+opencv-python==4.8.1.78
+Pillow==9.4.0
+pytorch-lightning==2.0.7
+scikit-image==0.22.0
+scikit-learn==1.3.1
+scipy==1.11.2
+tokenizers==0.13.3
+torch==2.0.1
+torch-kmeans==0.2.0
+torchaudio==2.0.2
+torchmetrics==1.0.3
+torchvision==0.15.2
+tqdm==4.66.1
+transformers==4.32.0
+xformers==0.0.21

utils.py

@@ -0,0 +1,231 @@
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+import cv2
+from sklearn.decomposition import PCA
+from torchvision import transforms
+import matplotlib.pyplot as plt
+import torch
+
+import os
+
+
+def display_attention_maps(
+    attention_maps,
+    is_cross,
+    num_heads,
+    tokenizer,
+    prompts,
+    dir_name,
+    step,
+    layer,
+    resolution,
+    is_query=False,
+    is_key=False,
+    points=None,
+    image_path=None,
+):
+    attention_maps = attention_maps.reshape(-1, num_heads, attention_maps.size(-2), attention_maps.size(-1))
+    num_samples = len(attention_maps) // 2
+    attention_type = 'cross' if is_cross else 'self'
+    for i, attention_map in enumerate(attention_maps):
+        if is_query:
+            attention_type = f'{attention_type}_queries'
+        elif is_key:
+            attention_type = f'{attention_type}_keys'
+
+        cond = 'uncond' if i < num_samples else 'cond'
+        i = i % num_samples
+        cur_dir_name = f'{dir_name}/{resolution}/{attention_type}/{layer}/{cond}/{i}'
+        os.makedirs(cur_dir_name, exist_ok=True)
+
+        if is_cross and not is_query:
+            fig = show_cross_attention(attention_map, tokenizer, prompts[i % num_samples])
+        else:
+            fig = show_self_attention(attention_map)
+            if points is not None:
+                point_dir_name = f'{cur_dir_name}/points'
+                os.makedirs(point_dir_name, exist_ok=True)
+                for j, point in enumerate(points):
+                    specific_point_dir_name = f'{point_dir_name}/{j}'
+                    os.makedirs(specific_point_dir_name, exist_ok=True)
+                    point_path = f'{specific_point_dir_name}/{step}.png'
+                    point_fig = show_individual_self_attention(attention_map, point, image_path=image_path)
+                    point_fig.save(point_path)
+                    point_fig.close()
+
+        fig.save(f'{cur_dir_name}/{step}.png')
+        fig.close()
+
+
+def text_under_image(image: np.ndarray, text: str, text_color: tuple[int, int, int] = (0, 0, 0)):
+    h, w, c = image.shape
+    offset = int(h * .2)
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    # font = ImageFont.truetype("/usr/share/fonts/truetype/noto/NotoMono-Regular.ttf", font_size)
+    text_size = cv2.getTextSize(text, font, 1, 2)[0]
+    lines = text.splitlines()
+    img = np.ones((h + offset + (text_size[1] + 2) * len(lines) - 2, w, c), dtype=np.uint8) * 255
+    img[:h, :w] = image
+
+    for i, line in enumerate(lines):
+        text_size = cv2.getTextSize(line, font, 1, 2)[0]
+        text_x, text_y = ((w - text_size[0]) // 2, h + offset + i * (text_size[1] + 2))
+        cv2.putText(img, line, (text_x, text_y), font, 1, text_color, 2)
+
+    return img
+
+def view_images(images, num_rows=1, offset_ratio=0.02):
+    if type(images) is list:
+        num_empty = len(images) % num_rows
+    elif images.ndim == 4:
+        num_empty = images.shape[0] % num_rows
+    else:
+        images = [images]
+        num_empty = 0
+
+    empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
+    images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
+    num_items = len(images)
+
+    h, w, c = images[0].shape
+    offset = int(h * offset_ratio)
+    num_cols = num_items // num_rows
+    image_ = np.ones((h * num_rows + offset * (num_rows - 1),
+                      w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
+    for i in range(num_rows):
+        for j in range(num_cols):
+            image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
+                i * num_cols + j]
+
+    return Image.fromarray(image_)
+
+
+def show_cross_attention(attention_maps, tokenizer, prompt, k_norms=None, v_norms=None):
+    attention_maps = attention_maps.mean(dim=0)
+    res = int(attention_maps.size(-2) ** 0.5)
+    attention_maps = attention_maps.reshape(res, res, -1)
+    tokens = tokenizer.encode(prompt)
+    decoder = tokenizer.decode
+    if k_norms is not None:
+        k_norms = k_norms.round(decimals=1)
+    if v_norms is not None:
+        v_norms = v_norms.round(decimals=1)
+    images = []
+    for i in range(len(tokens) + 5):
+        image = attention_maps[:, :, i]
+        image = 255 * image / image.max()
+        image = image.unsqueeze(-1).expand(*image.shape, 3)
+        image = image.detach().cpu().numpy().astype(np.uint8)
+        image = np.array(Image.fromarray(image).resize((256, 256)))
+        token = tokens[i] if i < len(tokens) else tokens[-1]
+        text = decoder(int(token))
+        if k_norms is not None and v_norms is not None:
+            text += f'\n{k_norms[i]}\n{v_norms[i]})'
+        image = text_under_image(image, text)
+        images.append(image)
+    return view_images(np.stack(images, axis=0))
+
+
+def show_queries_keys(queries, keys, colors, labels):  # [h ni d]
+    num_queries = [query.size(1) for query in queries]
+    num_keys = [key.size(1) for key in keys]
+    h, _, d = queries[0].shape
+
+    data = torch.cat((*queries, *keys), dim=1)  # h n d
+    data = data.permute(1, 0, 2)  # n h d
+    data = data.reshape(-1, h * d).detach().cpu().numpy()
+    pca = PCA(n_components=2)
+    data = pca.fit_transform(data)  # n 2
+
+    query_indices = np.array(num_queries).cumsum()
+    total_num_queries = query_indices[-1]
+    queries = np.split(data[:total_num_queries], query_indices[:-1])
+    if len(num_keys) == 0:
+        keys = [None, ] * len(labels)
+    else:
+        key_indices = np.array(num_keys).cumsum()
+        keys = np.split(data[total_num_queries:], key_indices[:-1])
+
+    fig, ax = plt.subplots()
+    marker_size = plt.rcParams['lines.markersize'] ** 2
+    query_size = int(1.25 * marker_size)
+    key_size = int(2 * marker_size)
+    for query, key, color, label in zip(queries, keys, colors, labels):
+        print(f'# queries of {label}', query.shape[0])
+        ax.scatter(query[:, 0], query[:, 1], s=query_size, color=color, marker='o', label=f'"{label}" queries')
+
+        if key is None:
+            continue
+
+        print(f'# keys of {label}', key.shape[0])
+        keys_label = f'"{label}" key'
+        if key.shape[0] > 1:
+            keys_label += 's'
+        ax.scatter(key[:, 0], key[:, 1], s=key_size, color=color, marker='x', label=keys_label)
+
+    ax.set_axis_off()
+    #ax.set_xlabel('X-axis')
+    #ax.set_ylabel('Y-axis')
+    #ax.set_title('Scatter Plot with Circles and Crosses')
+
+    #ax.legend()
+    return fig
+
+
+def show_self_attention(attention_maps):  # h n m
+    attention_maps = attention_maps.transpose(0, 1).flatten(start_dim=1).detach().cpu().numpy()
+    pca = PCA(n_components=3)
+    pca_img = pca.fit_transform(attention_maps)  # N X 3
+    h = w = int(pca_img.shape[0] ** 0.5)
+    pca_img = pca_img.reshape(h, w, 3)
+    pca_img_min = pca_img.min(axis=(0, 1))
+    pca_img_max = pca_img.max(axis=(0, 1))
+    pca_img = (pca_img - pca_img_min) / (pca_img_max - pca_img_min)
+    pca_img = Image.fromarray((pca_img * 255).astype(np.uint8))
+    pca_img = transforms.Resize(256, interpolation=transforms.InterpolationMode.NEAREST)(pca_img)
+    return pca_img
+
+
+def draw_box(pil_img, bboxes, colors=None, width=5):
+    draw = ImageDraw.Draw(pil_img)
+    #font = ImageFont.truetype('./FreeMono.ttf', 25)
+    w, h = pil_img.size
+    colors = ['red'] * len(bboxes) if colors is None else colors
+    for obj_bbox, color in zip(bboxes, colors):
+        x_0, y_0, x_1, y_1 = obj_bbox[0], obj_bbox[1], obj_bbox[2], obj_bbox[3]
+        draw.rectangle([int(x_0 * w), int(y_0 * h), int(x_1 * w), int(y_1 * h)], outline=color, width=width)
+    return pil_img
+
+
+def show_individual_self_attention(attn, point, image_path=None):
+    resolution = int(attn.size(-1) ** 0.5)
+    attn = attn.mean(dim=0).reshape(resolution, resolution, resolution, resolution)
+    attn = attn[round(point[1] * resolution), round(point[0] * resolution)]
+    attn = (attn - attn.min()) / (attn.max() - attn.min())
+    image = None if image_path is None else Image.open(image_path).convert('RGB')
+    image = show_image_relevance(attn, image=image)
+    return Image.fromarray(image)
+
+
+def show_image_relevance(image_relevance, image: Image.Image = None, relevnace_res=16):
+    # create heatmap from mask on image
+    def show_cam_on_image(img, mask):
+        img = img.resize((relevnace_res ** 2, relevnace_res ** 2))
+        img = np.array(img)
+        img = (img - img.min()) / (img.max() - img.min())
+        heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
+        heatmap = np.float32(heatmap) / 255
+        cam = heatmap + np.float32(img)
+        cam = cam / np.max(cam)
+        return cam
+
+    image_relevance = image_relevance.reshape(1, 1, image_relevance.shape[-1], image_relevance.shape[-1])
+    image_relevance = image_relevance.cuda() # because float16 precision interpolation is not supported on cpu
+    image_relevance = torch.nn.functional.interpolate(image_relevance, size=relevnace_res ** 2, mode='bilinear')
+    image_relevance = image_relevance.cpu() # send it back to cpu
+    image_relevance = (image_relevance - image_relevance.min()) / (image_relevance.max() - image_relevance.min())
+    image_relevance = image_relevance.reshape(relevnace_res ** 2, relevnace_res ** 2)
+    vis = image_relevance if image is None else show_cam_on_image(image, image_relevance)
+    vis = np.uint8(255 * vis)
+    vis = cv2.cvtColor(np.array(vis), cv2.COLOR_RGB2BGR)
+    return vis