code added

.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,13 @@
+ViT-L-14.pt
+*.png
+*.jpg
+*.jpeg
+./output/
+./output_old/
+./metrics/
+./dataset/
+__pycache__/
+./met/*
+*.csv
+./.idea/
+./.gradio/

README.md

@@ -0,0 +1,14 @@
+---
+title: Cora
+emoji: 🖼️
+colorFrom: yellow
+colorTo: yellow
+sdk: gradio
+sdk_version: 5.26.0
+app_file: app.py
+pinned: false
+license: mit
+short_description: Demo for Cora. Our few-step image editing method
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,207 @@
+import gradio as gr
+import spaces
+
+from PIL import Image
+
+import numpy as np
+import torch.nn.functional as F
+
+from utils.pipeline_utils import *
+from utils import get_args
+
+from main import run
+
+pipeline = None
+
+
+def get_pipeline():
+    global pipeline
+    if pipeline is None:
+        pipeline = load_pipeline(fp16=False, cache_dir=None)
+    return pipeline
+
+
+def process_masks(masks):
+    # masks: list of file paths
+    processed_masks = []
+    mask_composit = torch.zeros((512, 512), dtype=torch.float32, device='cuda')
+    for mask_path in masks:
+        mask = Image.open(mask_path).convert("L").resize((512, 512))
+        mask = torch.tensor(np.array(mask), dtype=torch.float32, device='cuda')
+        mask[mask > 0] = 255.0
+        mask = mask / 255.0
+        processed_masks.append(mask)
+        mask_composit += mask
+    mask_composit = torch.clamp(mask_composit, 0, 1)
+    mask_composit = F.interpolate(mask_composit[None, None, :, :], size=(64, 64), mode="nearest")[0, 0]
+    if mask_composit.sum() == 0:
+        mask_composit = None
+
+    return mask_composit
+
+
+@spaces.GPU
+def main_pipeline(
+        input_image: str,
+        src_prompt: str,
+        tgt_prompt: str,
+        alpha: float,
+        beta: float,
+        w1: float,
+        seed: int,
+        dift_correction: bool = True,
+):
+    args = get_args()
+    pipeline = get_pipeline()
+
+    args.alpha = alpha
+    args.beta = beta
+    args.w1 = w1
+    args.seed = seed
+    args.structural_alignment = True
+    args.support_new_object = True
+    args.apply_dift_correction = dift_correction
+    torch.cuda.empty_cache()
+    res_image = run(input_image['background'], src_prompt, tgt_prompt, masks=process_masks(input_image['layers']),
+                    pipeline=pipeline, args=args)[2]
+
+    return res_image
+
+
+DESCRIPTION = """# Cora 🖼️🐱🦅
+        ## Fast & Controllable Image Editing
+
+        ### 🛠️ Quick start  
+        1. **Upload** or drag-and-drop the image you’d like to edit.  
+        2. **Source prompt** – describe what’s in the original image.  
+        3. **Target prompt** – describe the result you want.  
+        4. Adjust the parameters as needed.  
+        5. *(Optional)* Paint a mask to specify the area to edit.  
+        6. Click **Edit** and wait a few seconds for the output.
+
+        ### ⚙️ Parameter cheat-sheet  
+
+        | Parameter | What it does | `0` (minimum) | `1` (maximum) |
+        |-----------|--------------|---------------|---------------|
+        | **alpha** | Appearance transfer control | preserve source appearance | target prompt affects appearance |
+        | **beta**  | Structural change control | preserve original structure | full layout change |
+        | **w**     | Prompt strength    | subtle tweaks | strong changes |
+        | **Seed**  | Fixes randomness for reproducibility | – | – |
+        | **Apply correspondence correction** | Uses correspondence-aware latent fix | – | – |
+
+        ### 📜 Tips  
+        - To replicate **TurboEdit**, set **alpha = 1**, **beta = 1**, and turn **off**  *Apply correspondence correction*.  
+        - To test reconstruction quality of the inversion, use identical source & target prompts with **alpha = 1** and **beta = 1**.
+        #### 🙏 Acknowledgements  
+        The demo template is largely adapted from **[TurboEdit on Hugging Face Spaces](https://huggingface.co/spaces/turboedit/turbo_edit)**.  
+    """
+
+with gr.Blocks(css="app/style.css") as demo:
+    gr.HTML(
+        """<a href="https://huggingface.co/spaces/armikaeili/cora?duplicate=true">
+        <img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>Duplicate the Space to run privately without waiting in queue"""
+    )
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.ImageMask(
+                label="Input image", type="filepath", height=512, width=512, brush=gr.Brush(color_mode='defaults')
+            )
+            result = gr.Image(label="Result", type="pil", height=512, width=512)
+        with gr.Column():
+            src_prompt = gr.Text(
+                label="Source Prompt",
+                max_lines=1,
+                placeholder="Source Prompt",
+            )
+            tgt_prompt = gr.Text(
+                label="Target Prompt",
+                max_lines=1,
+                placeholder="Target Prompt",
+            )
+            with gr.Accordion("Advanced Options", open=False):
+                seed = gr.Slider(
+                    label="seed", minimum=0, maximum=16 * 1024, value=200, step=1
+                )
+                w1 = gr.Slider(
+                    label="w", minimum=1.0, maximum=3.0, value=1.9, step=0.05
+                )
+                alpha = gr.Slider(
+                    label="alpha", minimum=0, maximum=1, value=0, step=0.01
+                )
+                beta = gr.Slider(
+                    label="beta", minimum=0, maximum=1, value=0.04, step=0.01
+                )
+                with gr.Row():
+                    dift_correction = gr.Checkbox(
+                        label="Apply correspondence correction",
+                        value=True)
+            run_button = gr.Button("Edit")
+
+            examples = [
+                [
+                    "assets/white_cat.png",  # input_image
+                    "a cat",  # src_prompt
+                    "a cat wearing a suit",  # tgt_prompt
+                    0.1,  # alpha
+                    0.1,  # beta
+                    1.9,  # w1
+                    7,  # seed
+                    True  # dift_correction
+                ],
+                [
+                    "assets/bear.png",  # input_image
+                    "a sitting brown bear",  # src_prompt
+                    "a roaring blue bear",  # tgt_prompt
+                    0.7,  # alpha
+                    0.1,  # beta
+                    1.9,  # w1
+                    7,  # seed
+                    True  # dift_correction
+                ],
+                [
+                    "assets/cat.jpg",  # input_image
+                    "a cat",  # src_prompt
+                    "an eagle",  # tgt_prompt
+                    0.7,  # alpha
+                    0.3,  # beta
+                    1.9,  # w1
+                    7,  # seed
+                    True  # dift_correction
+                ],
+                [
+                    "assets/dog.png",  # input_image
+                    "a photo of a dog",  # src_prompt
+                    "a photo of a dog lying",  # tgt_prompt
+                    0.0,  # alpha
+                    1,  # beta
+                    1.9,  # w1
+                    7,  # seed
+                    True  # dift_correction
+                ],
+
+            ]
+
+            inputs = [
+                input_image,
+                src_prompt,
+                tgt_prompt,
+                alpha,
+                beta,
+                w1,
+                seed,
+                dift_correction
+            ]
+            outputs = [result]
+            #
+            gr.Examples(
+                examples=examples,
+                inputs=inputs,
+                outputs=outputs,
+                fn=main_pipeline,
+                cache_examples=False,
+            )
+
+    run_button.click(fn=main_pipeline, inputs=inputs, outputs=outputs)
+demo.queue(max_size=50).launch(share=False)

main.py

@@ -0,0 +1,126 @@
+import os
+import torch
+from PIL import Image
+import jsonc as json
+
+from model import (
+    DirectionalAttentionControl,
+    StableDiffusionXLImg2ImgPipeline,
+    register_attention_editor_diffusers,
+)
+
+from utils.pipeline_utils import *
+from utils import get_args, extract_mask
+from src import get_ddpm_inversion_scheduler
+from visualization import save_results
+
+
+def run(
+    image_path,
+    src_prompt,
+    tgt_prompt,
+    masks,
+    pipeline: StableDiffusionXLImg2ImgPipeline,
+    args,
+):
+    seed = args.seed
+    num_timesteps = args.timesteps
+    torch.manual_seed(seed)
+    generator = torch.Generator(device=SAMPLING_DEVICE).manual_seed(seed)
+
+    timesteps, config = set_pipeline(pipeline, num_timesteps, generator, args)
+
+    x_0_image = Image.open(image_path).convert("RGB").resize((512, 512), RESIZE_TYPE)
+    x_0 = encode_image(x_0_image, pipeline, generator)
+    x_ts = create_xts(
+        config.noise_shift_delta,
+        config.noise_timesteps,
+        generator,
+        pipeline.scheduler,
+        timesteps,
+        x_0,
+    )
+    x_ts = [xt.to(dtype=x_0.dtype) for xt in x_ts]
+    latents = [x_ts[0]]
+
+    if not isinstance(masks, torch.Tensor):
+        mask = extract_mask(masks, 512, 512)
+    else:
+        mask = masks
+
+    pipeline.scheduler = get_ddpm_inversion_scheduler(
+        pipeline.scheduler,
+        config,
+        timesteps,
+        latents,
+        x_ts,
+        w1=args.w1,
+        dift_timestep=args.dift_timestep,
+        movement_intensifier=args.movement_intensifier,
+        apply_dift_correction=args.apply_dift_correction,
+        mask=mask,
+    )
+
+    
+    step, layer = 0, 44
+    editor = DirectionalAttentionControl(
+                                        step, layer, total_steps=11, 
+                                        model_type="SDXL", 
+                                        alpha=args.alpha, mode=args.mode, beta=1-args.beta,
+                                        structural_alignment=args.structural_alignment,
+                                        support_new_object=args.support_new_object
+                                        )
+    register_attention_editor_diffusers(pipeline, editor)
+
+    latent = latents[0].expand(3, -1, -1, -1)
+    prompt = [src_prompt, src_prompt, tgt_prompt]
+    pipeline.unet.latent_store.reset()
+    image = pipeline.__call__(image=latent, prompt=prompt).images
+    return [x_0_image, image[0], image[2]]
+
+
+if __name__ == "__main__":
+    args = get_args()
+
+    img_paths_to_prompts = json.load(open(args.prompts_file, "r"))
+    eval_dataset_folder = args.eval_dataset_folder
+
+    img_paths = [
+        f"{eval_dataset_folder}/{img_name}" for img_name in img_paths_to_prompts.keys()
+    ]
+    pipeline = load_pipeline(args.fp16, args.cache_dir)
+
+    sim_scores_total = 0
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    images_to_plot = []
+    output_dir = args.output_dir
+
+    for i, img_path in enumerate(img_paths):
+        args.img_path = img_path
+        img_name = img_path.split("/")[-1]
+        prompt = img_paths_to_prompts[img_name]["src_prompt"]
+        edit_prompts = img_paths_to_prompts[img_name]["tgt_prompt"]
+        args.alpha = img_paths_to_prompts[img_name].get("alpha", 0.7)
+        args.beta = img_paths_to_prompts[img_name].get("beta", 1)
+        masks = img_paths_to_prompts[img_name].get("masks", None)
+        args.mask = masks
+        args.source_prompt = prompt
+        args.target_prompt = edit_prompts[0]
+
+        res = run(
+            img_path,
+            prompt,
+            edit_prompts[0],
+            masks,
+            pipeline=pipeline,
+            args=args,
+        )
+        
+        torch.cuda.empty_cache()
+        save_results(
+            args=args,
+            source_prompt=prompt,
+            target_prompt=edit_prompts[0],
+            images=res
+            )

model/__init__.py

@@ -0,0 +1,5 @@
+from .directional_attentions import DirectionalAttentionControl
+from .modules.register_attention import register_attention_editor_diffusers
+from .pipeline_sdxl import StableDiffusionXLImg2ImgPipeline
+from .modules.dift_utils import gen_nn_map
+from .modules.freq_filters import freq_exp

model/directional_attentions.py

@@ -0,0 +1,478 @@
+import torch
+import torch.nn.functional as F
+import ctypes
+import numpy as np
+
+from einops import rearrange, repeat
+from scipy.optimize import linear_sum_assignment
+from typing import Optional, Union, Tuple, List, Callable, Dict
+
+from model.modules.dift_utils import gen_nn_map
+
+
+class AttentionBase:
+    def __init__(self):
+        self.cur_step = 0
+        self.num_att_layers = -1
+        self.cur_att_layer = 0
+
+    def after_step(self):
+        pass
+
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        sim: torch.Tensor,
+        attn: torch.Tensor,
+        is_cross: bool,
+        place_in_unet: str,
+        num_heads: int,
+        **kwargs
+    ) -> torch.Tensor:
+        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
+            self.after_step()
+
+        return out
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        sim: torch.Tensor,
+        attn: torch.Tensor,
+        is_cross: bool,
+        place_in_unet: str,
+        num_heads: int,
+        **kwargs
+    ) -> torch.Tensor:
+        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 DirectionalAttentionControl(AttentionBase):
+    MODEL_TYPE = {"SD": 16, "SDXL": 70}
+
+    def __init__(
+        self,
+        start_step: int = 4,
+        start_layer: int = 10,
+        layer_idx: Optional[List[int]] = None,
+        step_idx: Optional[List[int]] = None,
+        total_steps: int = 50,
+        model_type: str = "SD",
+        **kwargs
+    ):
+        super().__init__()
+        self.total_steps = total_steps
+        self.total_layers = self.MODEL_TYPE.get(model_type, 16)
+        self.start_step = start_step
+        self.start_layer = start_layer
+        self.layer_idx = layer_idx if layer_idx is not None else list(range(start_layer, self.total_layers))
+        self.step_idx = step_idx if step_idx is not None else list(range(start_step, total_steps))
+
+        self.w = 1.0
+        self.structural_alignment = kwargs.get("structural_alignment", False)
+        self.style_transfer_only = kwargs.get("style_transfer_only", False)
+        self.alpha = kwargs.get("alpha", 0.5)
+        self.beta = kwargs.get("beta", 0.5)
+        self.newness = kwargs.get("support_new_object", True)
+        self.mode = kwargs.get("mode", "normal")
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        sim: torch.Tensor,
+        attn: torch.Tensor,
+        is_cross: bool,
+        place_in_unet: str,
+        num_heads: int,
+        **kwargs
+    ) -> torch.Tensor:
+        if is_cross or self.cur_step not in self.step_idx or self.cur_att_layer // 2 not in self.layer_idx:
+            return super().forward(q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs)
+        
+        q_s, q_middle, q_t = q.chunk(3)
+        k_s, k_middle, k_t = k.chunk(3)
+        v_s, v_middle, v_t = v.chunk(3)
+        attn_s, attn_middle, attn_t = attn.chunk(3)
+
+        out_s = self.attn_batch(q_s, k_s, v_s, sim, attn_s, is_cross, place_in_unet, num_heads, **kwargs)
+        out_middle = self.attn_batch(q_middle, k_middle, v_middle, sim, attn_middle, is_cross, place_in_unet, num_heads, **kwargs)
+        
+        if self.cur_step <= 0 and self.beta > 0 and \
+                self.structural_alignment:
+            q_t = self.align_queries_via_matching(q_s, q_t, beta=self.beta)
+        
+        out_t = self.apply_mode(q_t, k_s, k_t, v_s, v_t, attn_t, sim, is_cross, place_in_unet, num_heads, **kwargs)
+
+        out = torch.cat([out_s, out_middle, out_t], dim=0)
+        return out
+
+    def apply_mode(
+        self,
+        q_t: torch.Tensor,
+        k_s: torch.Tensor,
+        k_t: torch.Tensor,
+        v_s: torch.Tensor,
+        v_t: torch.Tensor,
+        attn_t: torch.Tensor,
+        sim: torch.Tensor,
+        is_cross: bool,
+        place_in_unet: str,
+        num_heads: int,
+        **kwargs
+    ) -> torch.Tensor:
+        mode = self.mode
+
+        if 'dift' in mode and self.cur_step <= 0:
+            mode = 'normal'
+
+        if mode == "concat":
+            out_t = self.attn_batch(
+                q_t, torch.cat([k_s, 0.85 * k_t]), torch.cat([v_s, v_t]),
+                sim, attn_t, is_cross, place_in_unet, num_heads, **kwargs
+            )
+
+        elif mode == "concat_dift":
+            updated_k_s, updated_v_s, _ = self.process_dift_features(kwargs.get("dift_features"), k_s, k_t, v_s, v_t)
+            out_t = self.attn_batch(
+                q_t, torch.cat([updated_k_s, k_t]), torch.cat([updated_v_s, v_t]),
+                sim, attn_t, is_cross, place_in_unet, num_heads, **kwargs
+            )
+
+        elif mode == "masa":
+            out_t = self.attn_batch(q_t, k_s, v_s, sim, attn_t, is_cross, place_in_unet, num_heads, **kwargs)
+
+        elif mode == "normal":
+            out_t = self.attn_batch(q_t, k_t, v_t, sim, attn_t, is_cross, place_in_unet, num_heads, **kwargs)
+
+        elif mode == "lerp":
+            time = self.alpha
+            k_lerp = k_s + time * (k_t - k_s)
+            v_lerp = v_s + time * (v_t - v_s)
+            out_t = self.attn_batch(q_t, k_lerp, v_lerp, sim, attn_t, is_cross, place_in_unet, num_heads, **kwargs)
+
+        elif mode == "lerp_dift":
+            updated_k_s, updated_v_s, newness = self.process_dift_features(
+                kwargs.get("dift_features"), k_s, k_t, v_s, v_t, return_newness=self.newness
+            )
+            out_t = self.apply_lerp_dift(q_t, k_s, k_t, v_s, v_t, updated_k_s, updated_v_s, newness, sim, attn_t, is_cross, place_in_unet, num_heads, **kwargs)
+
+        elif mode in ("slerp", "log_slerp"):
+            time = self.alpha
+            k_slerp = self.slerp_fixed_length_batch(k_s, k_t, t=time)
+            v_slerp = self.slerp_batch(v_s, v_t, t=time, log_slerp="log" in mode)
+            out_t = self.attn_batch(q_t, k_slerp, v_slerp, sim, attn_t, is_cross, place_in_unet, num_heads, **kwargs)
+
+        elif mode in ("slerp_dift", "log_slerp_dift"):
+            out_t = self.apply_slerp_dift(q_t, k_s, k_t, v_s, v_t, sim, attn_t, is_cross, place_in_unet, num_heads, **kwargs)
+
+        else:
+            out_t = self.attn_batch(q_t, k_t, v_t, sim, attn_t, is_cross, place_in_unet, num_heads, **kwargs)
+
+        return out_t
+
+    def attn_batch(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        sim: torch.Tensor,
+        attn: torch.Tensor,
+        is_cross: bool,
+        place_in_unet: str,
+        num_heads: int,
+        **kwargs
+    ) -> torch.Tensor:
+        b = q.shape[0] // num_heads
+        q = rearrange(q, "(b h) n d -> h (b n) d", h=num_heads)
+        k = rearrange(k, "(b h) n d -> h (b n) d", h=num_heads)
+        v = rearrange(v, "(b h) n d -> h (b n) d", h=num_heads)
+
+        scale = kwargs.get("scale", 1.0)
+        sim_batched = torch.einsum("h i d, h j d -> h i j", q, k) * scale
+        attn_batched = sim_batched.softmax(-1)
+
+        out = torch.einsum("h i j, h j d -> h i d", attn_batched, v)
+        out = rearrange(out, "h (b n) d -> b n (h d)", b=b)
+        return out
+
+    def slerp(self, x: torch.Tensor, y: torch.Tensor, t: float = 0.5) -> torch.Tensor:
+        x_norm = x.norm(p=2)
+        y_norm = y.norm(p=2)
+        if y_norm < 1e-12:
+            return x
+
+        y_normalized = y / y_norm
+        y_same_length = y_normalized * x_norm
+        dot_xy = (x * y_same_length).sum()
+        cos_theta = torch.clamp(dot_xy / (x_norm * x_norm), -1.0, 1.0)
+        theta = torch.acos(cos_theta)
+        if torch.isclose(theta, torch.tensor(0.0)):
+            return x
+
+        sin_theta = torch.sin(theta)
+        s1 = torch.sin((1.0 - t) * theta) / sin_theta
+        s2 = torch.sin(t * theta) / sin_theta
+        return s1 * x + s2 * y_same_length
+
+    def slerp_batch(
+        self, 
+        x: torch.Tensor,
+        y: torch.Tensor,
+        t: float = 0.5,
+        eps: float = 1e-12,
+        log_slerp: bool = False
+    ) -> torch.Tensor:
+        """
+        Variation of SLERP for batches that allows for linear or logarithmic interpolation of magnitudes.
+        """
+        x_norm = x.norm(p=2, dim=-1, keepdim=True)
+        y_norm = y.norm(p=2, dim=-1, keepdim=True)
+        y_zero_mask = (y_norm < eps)
+
+        x_unit = x / (x_norm + eps)
+        y_unit = y / (y_norm + eps)
+        dot_xy = (x_unit * y_unit).sum(dim=-1, keepdim=True)
+        cos_theta = torch.clamp(dot_xy, -1.0, 1.0)
+
+        theta = torch.acos(cos_theta)
+        sin_theta = torch.sin(theta)
+        theta_zero_mask = (theta.abs() < 1e-7)
+
+        sin_theta_safe = torch.where(sin_theta.abs() < eps, torch.ones_like(sin_theta), sin_theta)
+        s1 = torch.sin((1.0 - t) * theta) / sin_theta_safe
+        s2 = torch.sin(t * theta) / sin_theta_safe
+        dir_interp = s1 * x_unit + s2 * y_unit
+
+        if not log_slerp:
+            mag_interp = (1.0 - t) * x_norm + t * y_norm
+        else:
+            mag_interp = (x_norm ** (1.0 - t)) * (y_norm ** t)
+
+        out = mag_interp * dir_interp
+        out = torch.where(y_zero_mask | theta_zero_mask, x, out)
+        return out
+
+
+    def slerp_fixed_length_batch(
+        self, 
+        x: torch.Tensor, 
+        y: torch.Tensor, 
+        t: float = 0.5, 
+        eps: float = 1e-12
+    ) -> torch.Tensor:
+        """
+        performing SLERP while preserving the norm of the source tensor x
+        """
+        x_norm = x.norm(p=2, dim=-1, keepdim=True)
+        y_norm = y.norm(p=2, dim=-1, keepdim=True)
+        y_zero_mask = (y_norm < eps)
+        y_normalized = y / (y_norm + eps)
+        y_same_length = y_normalized * x_norm
+        dot_xy = (x * y_same_length).sum(dim=-1, keepdim=True)
+        cos_theta = torch.clamp(dot_xy / (x_norm * x_norm + eps), -1.0, 1.0)
+        theta = torch.acos(cos_theta)
+        sin_theta = torch.sin(theta)
+
+        sin_theta_safe = torch.where(sin_theta.abs() < eps, torch.ones_like(sin_theta), sin_theta)
+        s1 = torch.sin((1.0 - t) * theta) / sin_theta_safe
+        s2 = torch.sin(t * theta) / sin_theta_safe
+        out = s1 * x + s2 * y_same_length
+        theta_zero_mask = (theta.abs() < 1e-7)
+
+        out = torch.where(y_zero_mask | theta_zero_mask, x, out)
+        return out
+
+    def apply_lerp_dift(
+        self,
+        q_t: torch.Tensor,
+        k_s: torch.Tensor,
+        k_t: torch.Tensor,
+        v_s: torch.Tensor,
+        v_t: torch.Tensor,
+        updated_k_s: torch.Tensor,
+        updated_v_s: torch.Tensor,
+        newness: torch.Tensor,
+        sim: torch.Tensor,
+        attn_t: torch.Tensor,
+        is_cross: bool,
+        place_in_unet: str,
+        num_heads: int,
+        **kwargs
+    ) -> torch.Tensor:
+        alpha = self.alpha
+        k_lerp = k_s + alpha * (k_t - k_s)
+        v_lerp = v_s + alpha * (v_t - v_s)
+        if alpha > 0:
+            k_t_new = newness * k_t + (1 - newness) * k_lerp
+            v_t_new = newness * v_t + (1 - newness) * v_lerp
+        else:
+            k_t_new = k_s
+            v_t_new = v_s
+
+        out_t = self.attn_batch(q_t, k_t_new, v_t_new, sim, attn_t, is_cross, place_in_unet, num_heads, **kwargs)
+        return out_t
+
+    def apply_slerp_dift(
+        self,
+        q_t: torch.Tensor,
+        k_s: torch.Tensor,
+        k_t: torch.Tensor,
+        v_s: torch.Tensor,
+        v_t: torch.Tensor,
+        sim: torch.Tensor,
+        attn_t: torch.Tensor,
+        is_cross: bool,
+        place_in_unet: str,
+        num_heads: int,
+        **kwargs
+    ) -> torch.Tensor:
+        updated_k_s, updated_v_s, newness = self.process_dift_features(
+            kwargs.get("dift_features"), k_s, k_t, v_s, v_t, return_newness=self.newness
+        )
+        alpha = self.alpha
+        log_slerp = "log" in self.mode
+
+        # Interpolate from k_t->updated_k_s so that if alpha=0, we get k_t
+        k_slerp = self.slerp_fixed_length_batch(k_t, updated_k_s, t=1-alpha)
+        v_slerp = self.slerp_batch(v_t, updated_v_s, t=1-alpha, log_slerp=log_slerp)
+
+        if alpha > 0:
+            k_t_new = newness * k_t + (1 - newness) * k_slerp
+            v_t_new = newness * v_t + (1 - newness) * v_slerp
+        else:
+            k_t_new = k_s
+            v_t_new = v_s
+
+        out_t = self.attn_batch(q_t, k_t_new, v_t_new, sim, attn_t, is_cross, place_in_unet, num_heads, **kwargs)
+        return out_t
+
+    def process_dift_features(
+        self,
+        dift_features: torch.Tensor,
+        k_s: torch.Tensor,
+        k_t: torch.Tensor,
+        v_s: torch.Tensor,
+        v_t: torch.Tensor,
+        return_newness: bool = True
+    ):
+        dift_s, _, dift_t = dift_features.chunk(3)
+        k_s1 = k_s.permute(0, 2, 1).reshape(k_s.shape[0], k_s.shape[2], int(k_s.shape[1]**0.5), -1)
+        v_s1 = v_s.permute(0, 2, 1).reshape(v_s.shape[0], v_s.shape[2], int(v_s.shape[1]**0.5), -1)
+
+        k_s1 = k_s1.reshape(-1, k_s1.shape[-2], k_s1.shape[-1])
+        v_s1 = v_s1.reshape(-1, v_s1.shape[-2], v_s1.shape[-1])
+
+        ################# uncomment only for visualization #################
+        # result = gen_nn_map(
+        #     [dift_s[0], dift_s[0]],
+        #     dift_s[0],
+        #     dift_t[0],
+        #     kernel_size=1,
+        #     stride=1,
+        #     device=k_s.device,
+        #     timestep=self.cur_step,
+        #     visualize=True,
+        #     return_newness=return_newness
+        # )
+        #####################################################################
+        
+        resized_src = F.interpolate(dift_s[0].unsqueeze(0), size=k_s1.shape[-1], mode='bilinear', align_corners=False).squeeze(0)
+        resized_tgt = F.interpolate(dift_t[0].unsqueeze(0), size=k_s1.shape[-1], mode='bilinear', align_corners=False).squeeze(0)
+
+        result = gen_nn_map(
+            [k_s1, v_s1],
+            resized_src,
+            resized_tgt,
+            kernel_size=1,
+            stride=1,
+            device=k_s.device,
+            timestep=self.cur_step,
+            return_newness=return_newness
+        )
+
+        if return_newness:
+            updated_k_s, updated_v_s, newness = result
+        else:
+            updated_k_s, updated_v_s = result
+            newness = torch.zeros_like(updated_k_s[:1]).to(k_s.device)
+
+        newness = newness.view(-1).unsqueeze(0).unsqueeze(-1)
+        updated_k_s = updated_k_s.reshape(k_s.shape[0], k_s.shape[2], -1).permute(0, 2, 1)
+        updated_v_s = updated_v_s.reshape(v_s.shape[0], v_s.shape[2], -1).permute(0, 2, 1)
+
+        return updated_k_s, updated_v_s, newness
+    
+    def sinkhorn(self, cost_matrix, max_iter=50, epsilon=1e-8):
+        n, m = cost_matrix.shape
+        K = torch.exp(-cost_matrix / cost_matrix.std())  # Kernelized cost matrix
+        u = torch.ones(n, device=cost_matrix.device) / n
+        v = torch.ones(m, device=cost_matrix.device) / m
+        
+        for _ in range(max_iter):
+            u_prev = u.clone()
+            u = 1.0 / (K @ v)
+            v = 1.0 / (K.T @ u)
+            if torch.max(torch.abs(u - u_prev)) < epsilon:
+                break
+        
+        P = torch.diag(u) @ K @ torch.diag(v)
+        return P
+
+    def align_queries_via_matching(self, q_s: torch.Tensor, q_t: torch.Tensor, beta: float = 0.5, device: str = "cuda"):
+        q_s = q_s.to(device)
+        q_t = q_t.to(device)
+
+        B, _, _ = q_s.shape
+        q_t_updated = torch.zeros_like(q_t, device=device)
+
+        for b in range(B):
+            ########################### L2 ##############################
+            # cost_matrix1 = (q_s[b].unsqueeze(1) - q_t[b].unsqueeze(0)).pow(2).sum(dim=-1)
+            ######################### cosine ############################
+            cost_matrix1 = - F.cosine_similarity(
+                            q_s[b].unsqueeze(1), q_t[b].unsqueeze(0), dim=-1)
+            #############################################################
+            # cost_matrix2 = (q_t[b].unsqueeze(1) - q_t[b].unsqueeze(0)).pow(2).sum(dim=-1)
+            cost_matrix2 = torch.abs(torch.arange(q_t[b].shape[0], device=device).unsqueeze(0) - 
+                         torch.arange(q_t[b].shape[0], device=device).unsqueeze(1)).float()
+            cost_matrix2 = cost_matrix2 ** 0.5
+            # cost_matrix2 = torch.where(cost_matrix2 > 0, 1.0, 0.0)
+            
+            mean1 = cost_matrix1.mean()
+            std1  = cost_matrix1.std()
+            mean2 = cost_matrix2.mean()
+            std2  = cost_matrix2.std()
+            cost_func_1_std = (cost_matrix1 - mean1) / (std1 + 1e-8)
+            cost_func_2_std = (cost_matrix2 - mean2) / (std2 + 1e-8)
+
+            cost_matrix = beta * cost_func_1_std + (1.0 - beta) * cost_func_2_std
+            cost_np = cost_matrix.detach().cpu().numpy()
+            row_ind, col_ind = linear_sum_assignment(cost_np)
+            q_t_updated[b] = q_t[b][col_ind]
+
+            # P = self.sinkhorn(cost_matrix)
+            # col_ind = P.argmax(dim=1)
+            # idea 1
+            # q_t_updated[b] = q_t[b][col_ind]
+            # idea 2
+            # q_t_updated[b] = P @ q_t[b]
+
+        return q_t_updated
+
+    

model/modules/dift_utils.py

@@ -0,0 +1,262 @@
+import os
+import math
+import torch
+import torch.nn.functional as F
+import numpy as np
+from typing import List
+from sklearn.decomposition import PCA
+from typing import Optional, Tuple
+from PIL import Image
+from model.modules.new_object_detection import *
+
+
+class DIFTLatentStore:
+    def __init__(self, steps: List[int], up_ft_indices: List[int]):
+        self.steps = steps
+        self.up_ft_indices = up_ft_indices
+        self.dift_features = {}
+        self.smoothed_dift_features = {}
+
+    def __call__(self, features: torch.Tensor, t: int, layer_index: int):
+        if t in self.steps and layer_index in self.up_ft_indices:
+            self.dift_features[f'{int(t)}_{layer_index}'] = features
+    
+    def smooth(self, kernel_size=3, sigma=1):
+        for key, value in self.dift_features.items():
+            if key not in self.smoothed_dift_features:
+                self.smoothed_dift_features[key] = torch.stack([gaussian_smooth(x, kernel_size=kernel_size, sigma=sigma) for x in value], dim=0)
+                
+    def copy(self):
+        copy_dift = DIFTLatentStore(self.steps, self.up_ft_indices)
+
+        for key, value in self.dift_features.items():
+            copy_dift.dift_features[key] = value.clone()
+
+        return copy_dift
+
+    def reset(self):
+        self.dift_features = {}
+        self.smoothed_dift_features = {}
+
+def gaussian_smooth(input_tensor, kernel_size=3, sigma=1):
+    kernel = np.fromfunction(
+        lambda x, y: (1/ (2 * np.pi * sigma ** 2)) * 
+                      np.exp(-((x - (kernel_size - 1) / 2) ** 2 + (y - (kernel_size - 1) / 2) ** 2) / (2 * sigma ** 2)),
+        (kernel_size, kernel_size)
+    )
+    kernel = torch.Tensor(kernel / kernel.sum()).to(input_tensor.dtype).to(input_tensor.device)
+    
+    kernel = kernel.unsqueeze(0).unsqueeze(0)
+    
+    smoothed_slices = []
+    for i in range(input_tensor.size(0)):
+        slice_tensor = input_tensor[i, :, :]
+        slice_tensor = F.conv2d(slice_tensor.unsqueeze(0).unsqueeze(0), kernel, padding=kernel_size // 2)[0, 0]
+        smoothed_slices.append(slice_tensor)
+    
+    smoothed_tensor = torch.stack(smoothed_slices, dim=0)
+
+    return smoothed_tensor
+
+def cos_dist(a, b):
+    a_norm = F.normalize(a, dim=-1)
+    b_norm = F.normalize(b, dim=-1)
+    res = a_norm @ b_norm.T
+    return 1 - res
+
+def extract_patches(feature_map: torch.Tensor, patch_size: int, stride: int) -> torch.Tensor:
+    # feature_map is (C, H, W). Unfold requires (B, C, H, W).
+    feature_map = feature_map.unsqueeze(0)  # (1, C, H, W)
+
+    # Unfold: output shape will be (B, C * patch_size^2, num_patches)
+    patches = F.unfold(
+        feature_map,
+        kernel_size=patch_size,
+        stride=stride
+    )
+    # Now patches is (1, C*patch_size^2, num_patches)
+
+    # Transpose to get shape (num_patches, C*patch_size^2)
+    patches = patches.squeeze(0).transpose(0, 1)  # (num_patches, C*patch_size^2)
+    return patches
+
+def reassemble_patches(
+    patches: torch.Tensor,
+    out_shape: Tuple[int, int, int],
+    patch_size: int,
+    stride: int
+) -> torch.Tensor:
+    C, H, W = out_shape
+    
+    # 1) Convert from (num_patches, C*patch_size^2) to (B=1, C*patch_size^2, num_patches)
+    patches_4d = patches.transpose(0, 1).unsqueeze(0)  # (1, C*patch_size^2, num_patches)
+    
+    # 2) fold: reassemble patches to (1, C, H, W)
+    reassembled = F.fold(
+        patches_4d,
+        output_size=(H, W),
+        kernel_size=patch_size,
+        stride=stride
+    )
+    
+    # 3) Create a divisor mask to account for overlapping regions.
+    #    We do this by folding a "ones" tensor of the same shape as patches_4d.
+    ones_input = torch.ones_like(patches_4d)
+    overlap_count = F.fold(
+        ones_input,
+        output_size=(H, W),
+        kernel_size=patch_size,
+        stride=stride
+    )
+    
+    # 4) Divide to normalize overlapping areas
+    reassembled = reassembled / overlap_count.clamp_min(1e-8)
+    
+    # 5) Remove the batch dimension -> (C, H, W)
+    reassembled = reassembled.squeeze(0)
+    
+    return reassembled
+
+def calculate_patch_distance(index1: int, index2: int, grid_size: int, stride: int, patch_size: int) -> float:
+    row1, col1 = index1 // grid_size, index1 % grid_size
+    row2, col2 = index2 // grid_size, index2 % grid_size
+    # print('row1, col1:', row1, col1)
+    x_center1, y_center1 = (row1 * stride) + (patch_size / 2), (col1 * stride) + (patch_size / 2)
+    x_center2, y_center2 = (row2 * stride) + (patch_size / 2), (col2 * stride) + (patch_size / 2)
+    return math.sqrt((x_center2 - x_center1)**2 + (y_center2 - y_center1)**2)
+
+def gen_nn_map(
+    latent, 
+    src_features, 
+    tgt_features, 
+    device, 
+    kernel_size=3, 
+    stride=1, 
+    return_newness=False,
+    **kwargs
+    ):
+    batch_size = kwargs.get("batch_size", None)
+    timestep = kwargs.get("timestep", None)
+    
+    if kwargs.get("visualize", False):
+        dift_visualization(src_features, tgt_features, filename_out=f"output/feat_colors_{timestep}.png")
+        
+    src_patches = extract_patches(src_features, kernel_size, stride)
+    tgt_patches = extract_patches(tgt_features, kernel_size, stride)
+    
+    if isinstance(latent, list):
+        latent_patches = [extract_patches(l, kernel_size, stride) for l in latent]
+    else:
+        latent_patches = extract_patches(latent, kernel_size, stride)
+
+    num_tgt = src_patches.size(0)
+    batch = batch_size or num_tgt
+    nearest_neighbor_indices = torch.empty(num_tgt, dtype=torch.long, device=device)
+    nearest_neighbor_distances = torch.empty(num_tgt, dtype=torch.long, device=device)
+    dist_chunks = []
+
+    for start in range(0, num_tgt, batch):
+        sims = cos_dist(src_patches, tgt_patches[start : start + batch])
+        dist_chunks.append(sims)
+        min_distances, best_idx = sims.min(0)
+        nearest_neighbor_indices[start : start + batch] = best_idx
+        nearest_neighbor_distances[start : start + batch] = min_distances
+
+    if not isinstance(latent, list):
+        aligned_latent = latent_patches[nearest_neighbor_indices]
+        aligned_latent = reassemble_patches(aligned_latent, latent.shape, kernel_size, stride)
+    else:
+        aligned_latent = [latent_patches[i][nearest_neighbor_indices] for i in range(len(latent_patches))]
+        aligned_latent = [reassemble_patches(l, latent[0].shape, kernel_size, stride) for l in aligned_latent]
+    
+    if return_newness:
+        dist_matrix = torch.cat(dist_chunks, dim=0)
+        newness_method = 'two_sided'
+        # newness_method = 'distance'
+        if newness_method.lower() == "distance":
+            newness = detect_newness_distance(nearest_neighbor_distances, quantile=0.97)
+        
+        elif newness_method.lower() == "two_sided":
+            newness = detect_newness_two_sided(dist_matrix, k=4)
+
+        out_shape = latent[0].shape if isinstance(latent, list) else latent.shape
+        out_shape = (1, out_shape[1], out_shape[2])
+
+        newness = reassemble_patches(newness.unsqueeze(-1), out_shape, kernel_size, stride)
+    
+
+    del src_patches, tgt_patches, latent_patches, nearest_neighbor_indices, nearest_neighbor_distances
+    
+    ################## visualization of changing source features to match target   ##################
+    if False:
+        updated_src_patches = src_patches[nearest_neighbor_indices]
+        updated_src_patches = reassemble_patches(updated_src_patches, src_features.shape, kernel_size, stride)
+        dift_visualization(
+            updated_src_patches, tgt_features,
+            filename_out=f"output/updated_feat_colors_{timestep}.png",
+        )
+    
+    if return_newness:
+        if isinstance(aligned_latent, list):
+            aligned_latent.append(newness)
+        else:
+            return aligned_latent, newness
+    return aligned_latent
+
+def dift_visualization(
+    src_feature: torch.Tensor,
+    tgt_feature: torch.Tensor,
+    filename_out: str,
+    resize_to: Optional[Tuple[int, int]] = (512, 512)
+):
+    """
+    Flatten features, apply PCA for 3D embedding, normalize for RGB, then reshape and save as image
+    """
+
+    C, H_s, W_s = src_feature.shape
+    _, H_t, W_t = tgt_feature.shape
+
+    src_flat = src_feature.permute(1, 2, 0).reshape(-1, C)  # (H_s*W_s, C)
+    tgt_flat = tgt_feature.permute(1, 2, 0).reshape(-1, C)  # (H_t*W_t, C)
+
+    all_features = torch.cat([src_flat, tgt_flat], dim=0)  # shape: (N_total, C)
+
+    all_features_np = all_features.detach().cpu().numpy()
+
+    num_components = 3
+    pca = PCA(n_components=num_components)
+    all_features_3d = pca.fit_transform(all_features_np)  # shape: (N_total, 3)
+
+    # 6) Normalize each dimension to [0,1]
+    def normalize_to_01(array_2d):
+        min_vals = array_2d.min(axis=0)
+        max_vals = array_2d.max(axis=0)
+        denom = (max_vals - min_vals) + 1e-8
+        return (array_2d - min_vals) / denom
+
+    all_features_rgb = normalize_to_01(all_features_3d)
+
+    N_src = H_s * W_s
+    src_rgb_flat = all_features_rgb[:N_src]      # (N_src, 3)
+    tgt_rgb_flat = all_features_rgb[N_src:]      # (N_tgt, 3)
+
+    src_color_map = src_rgb_flat.reshape(H_s, W_s, 3)
+    tgt_color_map = tgt_rgb_flat.reshape(H_t, W_t, 3)
+
+    src_img = Image.fromarray((src_color_map * 255).astype(np.uint8))
+    tgt_img = Image.fromarray((tgt_color_map * 255).astype(np.uint8))
+
+    src_img_resized = src_img.resize(resize_to, Image.Resampling.LANCZOS)
+    tgt_img_resized = tgt_img.resize(resize_to, Image.Resampling.LANCZOS)
+
+    combined_width = resize_to[0] * 2
+    combined_height = resize_to[1]
+    combined_img = Image.new("RGB", (combined_width, combined_height))
+    combined_img.paste(src_img_resized, (0, 0))
+    combined_img.paste(tgt_img_resized, (resize_to[0], 0))
+
+    os.makedirs(os.path.dirname(filename_out), exist_ok=True)
+    combined_img.save(filename_out)
+
+    print(f"Saved visualization to {filename_out}")
+

model/modules/freq_filters.py

@@ -0,0 +1,173 @@
+# Adapted from: https://github.com/kookie12/FlexiEdit/blob/main/flexiedit/frequency_utils.py
+
+import torch
+import torch.fft as fft
+import math
+import torch.nn.functional as F
+
+''' define hyperparameters '''
+# low-pass filter settings
+filter_type= "gaussian" #"butterworth" 
+n= 4 # gaussian parameter
+# Sampling process settings
+global alpha, reinversion_step, d_s, d_t, refined_step, masa_step_original, masa_step_target_branch, masa_step_retarget_branch
+alpha = 0.7
+d_t= 0.3
+d_s= 0.3
+refined_step = 0 
+masa_step_original = 4
+masa_step_target_branch = 51
+masa_step_retarget_branch = 0
+
+def gaussian_low_pass_filter(shape, d_s=0.25, d_t=0.25):
+    T, H, W = shape[-3], shape[-2], shape[-1]
+    mask = torch.zeros(shape)
+    if d_s==0 or d_t==0:
+        return mask
+    for t in range(T):
+        for h in range(H):
+            for w in range(W):
+                d_square = (((d_s/d_t)*(2*t/T-1))**2 + (2*h/H-1)**2 + (2*w/W-1)**2)
+                mask[..., t,h,w] = math.exp(-1/(2*d_s**2) * d_square)
+    return mask
+
+
+def butterworth_low_pass_filter(shape, n=4, d_s=0.25, d_t=0.25):
+    T, H, W = shape[-3], shape[-2], shape[-1]
+    mask = torch.zeros(shape)
+    if d_s==0 or d_t==0:
+        return mask
+    for t in range(T):
+        for h in range(H):
+            for w in range(W):
+                d_square = (((d_s/d_t)*(2*t/T-1))**2 + (2*h/H-1)**2 + (2*w/W-1)**2)
+                mask[..., t,h,w] = 1 / (1 + (d_square / d_s**2)**n)
+    return mask
+
+
+def ideal_low_pass_filter(shape, d_s=0.25, d_t=0.25):
+    T, H, W = shape[-3], shape[-2], shape[-1]
+    mask = torch.zeros(shape)
+    if d_s==0 or d_t==0:
+        return mask
+    for t in range(T):
+        for h in range(H):
+            for w in range(W):
+                d_square = (((d_s/d_t)*(2*t/T-1))**2 + (2*h/H-1)**2 + (2*w/W-1)**2)
+                mask[..., t,h,w] =  1 if d_square <= d_s*2 else 0
+    return mask
+
+
+def box_low_pass_filter(shape, d_s=0.25, d_t=0.25):
+    T, H, W = shape[-3], shape[-2], shape[-1]
+    mask = torch.zeros(shape)
+    if d_s==0 or d_t==0:
+        return mask
+
+    threshold_s = round(int(H // 2) * d_s)
+    threshold_t = round(T // 2 * d_t)
+
+    cframe, crow, ccol = T // 2, H // 2, W //2
+    #mask[..., cframe - threshold_t:cframe + threshold_t, crow - threshold_s:crow + threshold_s, ccol - threshold_s:ccol + threshold_s] = 1.0
+    mask[..., crow - threshold_s:crow + threshold_s, ccol - threshold_s:ccol + threshold_s] = 1.0
+
+    return mask
+
+def get_freq_filter(shape, device, filter_type, n, d_s, d_t):
+    if filter_type == "gaussian":
+        return gaussian_low_pass_filter(shape=shape, d_s=d_s, d_t=d_t).to(device)
+    elif filter_type == "ideal":
+        return ideal_low_pass_filter(shape=shape, d_s=d_s, d_t=d_t).to(device)
+    elif filter_type == "box":
+        return box_low_pass_filter(shape=shape, d_s=d_s, d_t=d_t).to(device)
+    elif filter_type == "butterworth":
+        return butterworth_low_pass_filter(shape=shape, n=n, d_s=d_s, d_t=d_t).to(device)
+    else:
+        raise NotImplementedError
+
+def freq_2d(x, LPF, alpha):
+    # FFT
+    x_freq = fft.fftn(x, dim=(-3, -2, -1))
+    x_freq = fft.fftshift(x_freq, dim=(-3, -2, -1))
+
+    # frequency mix
+    HPF = 1 - LPF
+
+    x_freq_low = x_freq * LPF
+    x_freq_high = x_freq * HPF
+    
+    x_freq_sum = x_freq
+
+    # IFFT
+    _x_freq_low = fft.ifftshift(x_freq_low, dim=(-3, -2, -1))
+    x_low = fft.ifftn(_x_freq_low, dim=(-3, -2, -1)).real
+    x_low_alpha = fft.ifftn(_x_freq_low*alpha, dim=(-3, -2, -1)).real
+    
+    _x_freq_high = fft.ifftshift(x_freq_high, dim=(-3, -2, -1))
+    x_high = fft.ifftn(_x_freq_high, dim=(-3, -2, -1)).real
+    x_high_alpha = fft.ifftn(_x_freq_high*alpha, dim=(-3, -2, -1)).real
+    
+    _x_freq_sum = fft.ifftshift(x_freq_sum, dim=(-3, -2, -1))
+    x_sum = fft.ifftn(_x_freq_sum, dim=(-3, -2, -1)).real
+    
+    _x_freq_low_alpha_high = fft.ifftshift(x_freq_low + x_freq_high*alpha, dim=(-3, -2, -1))
+    x_low_alpha_high = fft.ifftn(_x_freq_low_alpha_high, dim=(-3, -2, -1)).real
+    
+    _x_freq_high_alpha_low = fft.ifftshift(x_freq_low*alpha + x_freq_high, dim=(-3, -2, -1))
+    x_high_alpha_low = fft.ifftn(_x_freq_high_alpha_low, dim=(-3, -2, -1)).real
+
+    _x_freq_alpha_high_alpha_low = fft.ifftshift(x_freq_low*alpha + x_freq_high*alpha, dim=(-3, -2, -1))
+    x_alpha_high_alpha_low = fft.ifftn(_x_freq_alpha_high_alpha_low, dim=(-3, -2, -1)).real
+
+    return x_low, x_high, x_sum, x_low_alpha, x_high_alpha, x_low_alpha_high, x_high_alpha_low, x_alpha_high_alpha_low
+
+
+def freq_exp(feat, mode, user_mask, auto_mask, movement_intensifier):
+    movement_intensifier = 1 - movement_intensifier
+    """ Frequency manipulation for latent space. """
+    feat = feat.view(4,1,64,64)
+    f_shape = feat.shape # 1, 4, 64, 64
+    LPF = get_freq_filter(f_shape, feat.device, filter_type, n, d_s, d_t) # d_s, d_t
+    f_dtype = feat.dtype
+    feat_low, feat_high, feat_sum, feat_low_alpha, feat_high_alpha, feat_low_alpha_high, feat_high_alpha_low, x_alpha_high_alpha_low = freq_2d(feat.to(torch.float64), LPF, movement_intensifier)
+    feat_low = feat_low.to(f_dtype)
+    feat_high = feat_high.to(f_dtype)
+    feat_sum = feat_sum.to(f_dtype)
+    feat_low_alpha = feat_low_alpha.to(f_dtype)
+    feat_high_alpha = feat_high_alpha.to(f_dtype)
+    feat_low_alpha_high = feat_low_alpha_high.to(f_dtype)
+    feat_high_alpha_low = feat_high_alpha_low.to(f_dtype)
+
+    latent_low = feat_low.view(1,4,64,64)
+    
+    latent_high = feat_high.view(1,4,64,64)
+    
+    latent_sum = feat_sum.view(1,4,64,64)
+    
+    latent_low_alpha_high = feat_low_alpha_high.view(1,4,64,64)
+    latent_high_alpha_low = feat_high_alpha_low.view(1,4,64,64)
+    
+    mask = torch.zeros_like(latent_sum)
+    if mode == "auto_mask":
+        auto_mask = auto_mask.unsqueeze(1) # [1,64,64] => [1,1,64,64]
+        mask = auto_mask.expand_as(latent_sum) # [1,1,64,64] => [1,4,64,64]
+        
+    elif mode == "user_mask":
+        bbx_start_point, bbx_end_point = user_mask
+        mask[:, :, bbx_start_point[1]//8:bbx_end_point[1]//8, bbx_start_point[0]//8:bbx_end_point[0]//8] = 1
+        
+    latents_shape = latent_sum.shape
+    random_gaussian = torch.randn(latents_shape, device=latent_sum.device)
+    
+    # Apply gaussian scaling
+    g_range = random_gaussian.max() - random_gaussian.min()
+    l_range = latent_low_alpha_high.max() - latent_low_alpha_high.min()
+    random_gaussian = random_gaussian * (l_range/g_range)
+
+    # No scaling applied. If you wish to apply scaling to the mask, replace the following lines accordingly.
+    s_range, r_range, s_range2, r_range2 = 1, 1, 1, 1
+        
+    latent_mask_h = latent_sum * (1 - mask) + (latent_low_alpha_high + (1-movement_intensifier)*random_gaussian) * (s_range/r_range) *mask # edit 할 부분에 high frequency가 줄어들고 가우시안 더하기
+    latent_mask_l = latent_sum * (1 - mask) + (latent_high_alpha_low + (1-movement_intensifier)*random_gaussian) * (s_range2/r_range2) *mask # edit 할 부분에 low frequency가 줄어들고 가우시안 더하기
+    
+    return latent_mask_h, latent_mask_l, latent_sum # latent_low, latent_high, latent_sum

model/modules/new_object_detection.py

@@ -0,0 +1,187 @@
+import torch
+
+def distance_to_similarity(distances, temperature=1.0):
+    """
+    Turns a distance matrix into a similarity matrix so it works with distribution-based metrics.
+    """
+    similarities = torch.exp(-distances / temperature)
+    similarities = torch.clamp(similarities, min=1e-8)
+    return similarities
+
+#################################
+##   "New Object" Detection    ##
+#################################
+
+def detect_newness_two_sided(distances, k=3, quantile=0.97):
+    device = distances.device
+    N_src, N_tgt = distances.shape
+
+    topk_src_idx_t = torch.topk(distances, k, dim=0, largest=False).indices  # [k, N_tgt]
+    topk_tgt_idx_s = torch.topk(distances, k, dim=1, largest=False).indices  # [N_src, k]
+
+    src_to_tgt_mask = torch.zeros((N_src, N_tgt), device=device)
+    tgt_to_src_mask = torch.zeros((N_src, N_tgt), device=device)
+
+    row_indices = topk_src_idx_t  # [k, N_tgt]
+    col_indices = torch.arange(N_tgt, device=device).unsqueeze(0).repeat(k, 1)  # [k, N_tgt]
+    src_to_tgt_mask[row_indices, col_indices] = 1.0  # Assign 1.0 at the top-k positions
+
+    row_indices = torch.arange(N_src, device=device).unsqueeze(1).repeat(1, k)  # [N_src, k]
+    col_indices = topk_tgt_idx_s  # [N_src, k]
+    tgt_to_src_mask[row_indices, col_indices] = 1.0  # Assign 1.0 at the top-k positions
+
+    overlap_mask = (src_to_tgt_mask * tgt_to_src_mask).sum(dim=0) > 0  # [N_tgt]
+
+    distances[:, overlap_mask] = 0.0
+
+    two_sided_mask = (~overlap_mask).float()
+
+    min_distances, _ = distances.min(dim=0)
+    threshold = torch.quantile(min_distances, quantile)
+    threshold_mask = (min_distances > threshold).float()
+
+    combined_mask = two_sided_mask * threshold_mask
+    return combined_mask
+
+def detect_newness_distance(min_distances, quantile=0.97):
+    """
+    Old approach: threshold on min distance at a chosen percentile.
+    """
+    threshold = torch.quantile(min_distances, quantile)
+    newness_mask = (min_distances > threshold).float()
+    return newness_mask
+
+def detect_newness_topk_margin(distances, top_k=2, quantile=0.03):
+    """
+    Top-k margin approach in distance space.
+    distances: [N_src, N_tgt]
+    Sort each column ascending => best match is index 0, second best is index 1, etc.
+    A smaller margin => ambiguous => likely new.
+    We threshold the margin at some percentile.
+    """
+    sorted_dists, _ = torch.sort(distances, dim=0)  
+    best = sorted_dists[0]                        # [N_tgt]
+    second_best = sorted_dists[1] if top_k >= 2 else sorted_dists[0]  # [N_tgt]
+    margin = second_best - best  # [N_tgt]
+
+    # If margin < threshold => ambiguous => "new"
+    # We'll pick threshold as a quantile of margin
+    threshold = torch.quantile(margin, quantile)
+    newness_mask = (margin < threshold).float()
+    return newness_mask
+
+def detect_newness_entropy(distances, temperature=1.0, quantile=0.97):
+    """
+    Entropy-based approach. First convert distance->similarity with an exponential.
+    Then normalize to get a distribution for each target patch, compute Shannon entropy.
+    High entropy => new object (no strong match).
+    """
+    similarities = distance_to_similarity(distances, temperature=temperature)
+    probs = similarities / similarities.sum(dim=0, keepdim=True)  # [N_src, N_tgt]
+    # Shannon Entropy: -sum(p log p)
+    entropy = -torch.sum(probs * torch.log(probs), dim=0)  # [N_tgt]
+
+    # threshold
+    threshold = torch.quantile(entropy, quantile)
+    newness_mask = (entropy > threshold).float()
+    return newness_mask
+
+def detect_newness_gini(distances, temperature=1.0, quantile=0.97):
+    """
+    Gini impurity-based approach. Convert distances to similarities,
+    get a distribution, compute Gini.
+    High Gini => wide distribution => new object.
+    """
+    similarities = distance_to_similarity(distances, temperature=temperature)
+    probs = similarities / similarities.sum(dim=0, keepdim=True)
+    # Gini: sum(p_i*(1-p_i)) => high if spread out
+    gini = torch.sum(probs * (1.0 - probs), dim=0)  # [N_tgt]
+
+    threshold = torch.quantile(gini, quantile)
+    newness_mask = (gini > threshold).float()
+    return newness_mask
+
+def detect_newness_kl(distances, temperature=1.0, quantile=0.97):
+    """
+    KL-based approach. Compare distribution to uniform => if close to uniform => new object.
+    1) Convert distances -> similarities
+    2) p(x) = similarities / sum(similarities)
+    3) KL(p || uniform) => sum p(x) log (p(x)/(1/N_src))
+    4) If p is near uniform => KL small => new object.
+       We'll invert it => newness ~ 1/KL.
+    """
+    similarities = distance_to_similarity(distances, temperature=temperature)
+    N_src = distances.shape[0]
+    probs = similarities / similarities.sum(dim=0, keepdim=True)
+
+    uniform_val = 1.0 / float(N_src)
+    kl_vals = torch.sum(probs * torch.log(probs / uniform_val), dim=0)  # [N_tgt]
+    inv_kl = 1.0 / (kl_vals + 1e-8)  # big => distribution is near uniform => new
+
+    threshold = torch.quantile(inv_kl, quantile)
+    newness_mask = (inv_kl > threshold).float()
+    return newness_mask
+
+def detect_newness_variation_ratio(distances, temperature=1.0, quantile=0.97):
+    """
+    Variation Ratio: 1 - max(prob).
+    1) Convert distance->similarity
+    2) p(x) = sim(x) / sum_x'(sim(x'))
+    3) var_ratio = 1 - max(p)
+    High var_ratio => new object.
+    """
+    similarities = distance_to_similarity(distances, temperature=temperature)
+    probs = similarities / similarities.sum(dim=0, keepdim=True)
+    max_prob, _ = torch.max(probs, dim=0)  # [N_tgt]
+    var_ratio = 1.0 - max_prob
+
+    threshold = torch.quantile(var_ratio, quantile)
+    newness_mask = (var_ratio > threshold).float()
+    return newness_mask
+
+
+def detect_newness_two_sided_ratio(
+    distances,
+    top_k_ratio_quantile=0.03,
+    two_sided=True
+):
+    """
+    Two-sided matching + ratio test in distance space.
+
+    Ratio test: For each t, let d0 = best distance, d1 = second best.
+        ratio = d0 / (d1 + 1e-8).
+        If ratio < ratio_threshold => ambiguous => new.
+        (Typically a smaller ratio means a better match, but we invert logic:
+        a patch can be "new" if the ratio is extremely small or ambiguous.)  
+    """
+
+    N_src, N_tgt = distances.shape
+
+    # Target → Source: best match
+    min_vals_t, best_s_for_t = torch.min(distances, dim=0)
+
+    # Source → Target: best match
+    min_vals_s, best_t_for_s = torch.min(distances, dim=1)
+
+    # Two-sided consistency check
+    twosided_mask = torch.zeros(N_tgt, device=distances.device)
+    if two_sided:
+        for t in range(N_tgt):
+            s = best_s_for_t[t]
+            if best_t_for_s[s] != t:
+                twosided_mask[t] = 1.0
+
+    # Ratio test: ambiguous if best match is not clearly better than second best
+    sorted_dists, _ = torch.sort(distances, dim=0)
+    d0 = sorted_dists[0]
+    d1 = sorted_dists[1]
+    ratio = d0 / (d1 + 1e-8)
+    ratio_threshold = torch.quantile(ratio, top_k_ratio_quantile)
+    ratio_mask = (ratio < ratio_threshold).float()
+
+    # Combine checks (currently using only two-sided result)
+    newness_mask = twosided_mask
+
+    return newness_mask
+
+

model/modules/register_attention.py

@@ -0,0 +1,77 @@
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+
+from typing import Any
+from model.directional_attentions import DirectionalAttentionControl, AttentionBase
+from utils.utils import find_smallest_key_with_suffix
+
+
+def register_attention_editor_diffusers(model: Any, editor: AttentionBase):
+    def ca_forward(self, place_in_unet):
+        def forward(
+            x: torch.Tensor, 
+            encoder_hidden_states: torch.Tensor = None,
+            attention_mask: torch.Tensor = None,
+            context: torch.Tensor = None,
+            mask: torch.Tensor = None
+        ):
+            if encoder_hidden_states is not None:
+                context = encoder_hidden_states
+            if attention_mask is not None:
+                mask = attention_mask
+
+            h = self.heads
+            is_cross = context is not None
+            context = context if is_cross else x
+
+            q = self.to_q(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)
+                sim.masked_fill_(~mask, max_neg_value)
+
+            dift_features_dict = getattr(model.unet.latent_store, 'dift_features', {})
+            dift_features_key = find_smallest_key_with_suffix(dift_features_dict, suffix='_1')
+            dift_features = dift_features_dict.get(dift_features_key, None)
+
+            attn = sim.softmax(dim=-1)
+            out = editor(
+                q, k, v, sim, attn, is_cross, place_in_unet,
+                self.heads,
+                scale=self.scale,
+                dift_features=dift_features
+            )
+
+            to_out = self.to_out
+            if isinstance(to_out, nn.modules.container.ModuleList):
+                to_out = self.to_out[0]
+
+            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__ == '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)
+        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")
+        elif "mid" in net_name:
+            cross_att_count += register_editor(net, 0, "mid")
+        elif "up" in net_name:
+            cross_att_count += register_editor(net, 0, "up")
+    editor.num_att_layers = cross_att_count

model/pipeline_sdxl.py

@@ -0,0 +1,1440 @@
+# Modified only lines 1360–1370; search for 'changed section' to locate the update easily.
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import PIL.Image
+import torch
+from transformers import (
+    CLIPImageProcessor,
+    CLIPTextModel,
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    CLIPVisionModelWithProjection,
+)
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import (
+    FromSingleFileMixin,
+    IPAdapterMixin,
+    StableDiffusionXLLoraLoaderMixin,
+    TextualInversionLoaderMixin,
+)
+from diffusers.models import AutoencoderKL, ImageProjection, UNet2DConditionModel
+from diffusers.models.attention_processor import (
+    AttnProcessor2_0,
+    LoRAAttnProcessor2_0,
+    LoRAXFormersAttnProcessor,
+    XFormersAttnProcessor,
+)
+from diffusers.models.lora import adjust_lora_scale_text_encoder
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    deprecate,
+    is_invisible_watermark_available,
+    is_torch_xla_available,
+    logging,
+    replace_example_docstring,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
+from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
+
+from model.modules.dift_utils import gaussian_smooth, gen_nn_map
+
+
+if is_invisible_watermark_available():
+    from .watermark import StableDiffusionXLWatermarker
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import StableDiffusionXLImg2ImgPipeline
+        >>> from diffusers.utils import load_image
+
+        >>> pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+        ...     "stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16
+        ... )
+        >>> pipe = pipe.to("cuda")
+        >>> url = "https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/aa_xl/000000009.png"
+
+        >>> init_image = load_image(url).convert("RGB")
+        >>> prompt = "a photo of an astronaut riding a horse on mars"
+        >>> image = pipe(prompt, image=init_image).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
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used,
+            `timesteps` must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+                Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
+                timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
+                must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class StableDiffusionXLImg2ImgPipeline(
+    DiffusionPipeline,
+    StableDiffusionMixin,
+    TextualInversionLoaderMixin,
+    FromSingleFileMixin,
+    StableDiffusionXLLoraLoaderMixin,
+    IPAdapterMixin,
+):
+    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.)
+
+    The pipeline also inherits the following loading methods:
+        - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
+        - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
+        - [`~loaders.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] for loading LoRA weights
+        - [`~loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] for saving LoRA weights
+        - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
+
+    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`].
+        requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`):
+            Whether the `unet` requires an `aesthetic_score` condition to be passed during inference. Also see the
+            config of `stabilityai/stable-diffusion-xl-refiner-1-0`.
+        force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"True"`):
+            Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of
+            `stabilityai/stable-diffusion-xl-base-1-0`.
+        add_watermarker (`bool`, *optional*):
+            Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to
+            watermark output images. If not defined, it will default to True if the package is installed, otherwise no
+            watermarker will be used.
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae"
+    _optional_components = [
+        "tokenizer",
+        "tokenizer_2",
+        "text_encoder",
+        "text_encoder_2",
+        "image_encoder",
+        "feature_extractor",
+    ]
+    _callback_tensor_inputs = [
+        "latents",
+        "prompt_embeds",
+        "negative_prompt_embeds",
+        "add_text_embeds",
+        "add_time_ids",
+        "negative_pooled_prompt_embeds",
+        "add_neg_time_ids",
+    ]
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        text_encoder_2: CLIPTextModelWithProjection,
+        tokenizer: CLIPTokenizer,
+        tokenizer_2: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        image_encoder: CLIPVisionModelWithProjection = None,
+        feature_extractor: CLIPImageProcessor = None,
+        requires_aesthetics_score: bool = False,
+        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,
+            image_encoder=image_encoder,
+            feature_extractor=feature_extractor,
+            scheduler=scheduler,
+        )
+        self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt)
+        self.register_to_config(requires_aesthetics_score=requires_aesthetics_score)
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+        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_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.encode_prompt
+    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,
+        clip_skip: Optional[int] = 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.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+        """
+        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, StableDiffusionXLLoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if self.text_encoder is not None:
+                if not USE_PEFT_BACKEND:
+                    adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
+                else:
+                    scale_lora_layers(self.text_encoder, lora_scale)
+
+            if self.text_encoder_2 is not None:
+                if not USE_PEFT_BACKEND:
+                    adjust_lora_scale_text_encoder(self.text_encoder_2, lora_scale)
+                else:
+                    scale_lora_layers(self.text_encoder_2, lora_scale)
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        if prompt is not None:
+            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
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            # textual inversion: process 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]
+                if clip_skip is None:
+                    prompt_embeds = prompt_embeds.hidden_states[-2]
+                else:
+                    # "2" because SDXL always indexes from the penultimate layer.
+                    prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 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
+
+            # normalize str to list
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            negative_prompt_2 = (
+                batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+            )
+
+            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 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)
+
+        if self.text_encoder_2 is not None:
+            prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
+        else:
+            prompt_embeds = prompt_embeds.to(dtype=self.unet.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]
+
+            if self.text_encoder_2 is not None:
+                negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
+            else:
+                negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.unet.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
+            )
+
+        if self.text_encoder is not None:
+            if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder, lora_scale)
+
+        if self.text_encoder_2 is not None:
+            if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder_2, lora_scale)
+
+        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,
+        strength,
+        num_inference_steps,
+        callback_steps,
+        negative_prompt=None,
+        negative_prompt_2=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        ip_adapter_image=None,
+        ip_adapter_image_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+    ):
+        if strength < 0 or strength > 1:
+            raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
+        if num_inference_steps is None:
+            raise ValueError("`num_inference_steps` cannot be None.")
+        elif not isinstance(num_inference_steps, int) or num_inference_steps <= 0:
+            raise ValueError(
+                f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type"
+                f" {type(num_inference_steps)}."
+            )
+        if 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 callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        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 ip_adapter_image is not None and ip_adapter_image_embeds is not None:
+            raise ValueError(
+                "Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined."
+            )
+
+        if ip_adapter_image_embeds is not None:
+            if not isinstance(ip_adapter_image_embeds, list):
+                raise ValueError(
+                    f"`ip_adapter_image_embeds` has to be of type `list` but is {type(ip_adapter_image_embeds)}"
+                )
+            elif ip_adapter_image_embeds[0].ndim not in [3, 4]:
+                raise ValueError(
+                    f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
+                )
+
+    def get_timesteps(self, num_inference_steps, strength, device, denoising_start=None):
+        # get the original timestep using init_timestep
+        if denoising_start is None:
+            init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
+            t_start = max(num_inference_steps - init_timestep, 0)
+        else:
+            t_start = 0
+
+        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
+
+        # Strength is irrelevant if we directly request a timestep to start at;
+        # that is, strength is determined by the denoising_start instead.
+        if denoising_start is not None:
+            discrete_timestep_cutoff = int(
+                round(
+                    self.scheduler.config.num_train_timesteps
+                    - (denoising_start * self.scheduler.config.num_train_timesteps)
+                )
+            )
+
+            num_inference_steps = (timesteps < discrete_timestep_cutoff).sum().item()
+            if self.scheduler.order == 2 and num_inference_steps % 2 == 0:
+                # if the scheduler is a 2nd order scheduler we might have to do +1
+                # because `num_inference_steps` might be even given that every timestep
+                # (except the highest one) is duplicated. If `num_inference_steps` is even it would
+                # mean that we cut the timesteps in the middle of the denoising step
+                # (between 1st and 2nd devirative) which leads to incorrect results. By adding 1
+                # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler
+                num_inference_steps = num_inference_steps + 1
+
+            # because t_n+1 >= t_n, we slice the timesteps starting from the end
+            timesteps = timesteps[-num_inference_steps:]
+            return timesteps, num_inference_steps
+
+        return timesteps, num_inference_steps - t_start
+
+    def prepare_latents(
+        self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None, add_noise=True
+    ):
+        if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
+            raise ValueError(
+                f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
+            )
+
+        # Offload text encoder if `enable_model_cpu_offload` was enabled
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.text_encoder_2.to("cpu")
+            torch.cuda.empty_cache()
+
+        image = image.to(device=device, dtype=dtype)
+
+        batch_size = batch_size * num_images_per_prompt
+
+        if image.shape[1] == 4:
+            init_latents = image
+
+        else:
+            # make sure the VAE is in float32 mode, as it overflows in float16
+            if self.vae.config.force_upcast:
+                image = image.float()
+                self.vae.to(dtype=torch.float32)
+
+            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."
+                )
+
+            elif isinstance(generator, list):
+                init_latents = [
+                    retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i])
+                    for i in range(batch_size)
+                ]
+                init_latents = torch.cat(init_latents, dim=0)
+            else:
+                init_latents = retrieve_latents(self.vae.encode(image), generator=generator)
+
+            if self.vae.config.force_upcast:
+                self.vae.to(dtype)
+
+            init_latents = init_latents.to(dtype)
+            init_latents = self.vae.config.scaling_factor * init_latents
+
+        if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:
+            # expand init_latents for batch_size
+            additional_image_per_prompt = batch_size // init_latents.shape[0]
+            init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0)
+        elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0:
+            raise ValueError(
+                f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
+            )
+        else:
+            init_latents = torch.cat([init_latents], dim=0)
+
+        if add_noise:
+            shape = init_latents.shape
+            noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+            # get latents
+            init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
+
+        latents = init_latents
+
+        return latents
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image
+    def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):
+        dtype = next(self.image_encoder.parameters()).dtype
+
+        if not isinstance(image, torch.Tensor):
+            image = self.feature_extractor(image, return_tensors="pt").pixel_values
+
+        image = image.to(device=device, dtype=dtype)
+        if output_hidden_states:
+            image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
+            image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
+            uncond_image_enc_hidden_states = self.image_encoder(
+                torch.zeros_like(image), output_hidden_states=True
+            ).hidden_states[-2]
+            uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
+                num_images_per_prompt, dim=0
+            )
+            return image_enc_hidden_states, uncond_image_enc_hidden_states
+        else:
+            image_embeds = self.image_encoder(image).image_embeds
+            image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+            uncond_image_embeds = torch.zeros_like(image_embeds)
+
+            return image_embeds, uncond_image_embeds
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds
+    def prepare_ip_adapter_image_embeds(
+        self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance
+    ):
+        if ip_adapter_image_embeds is None:
+            if not isinstance(ip_adapter_image, list):
+                ip_adapter_image = [ip_adapter_image]
+
+            if len(ip_adapter_image) != len(self.unet.encoder_hid_proj.image_projection_layers):
+                raise ValueError(
+                    f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters."
+                )
+
+            image_embeds = []
+            for single_ip_adapter_image, image_proj_layer in zip(
+                ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers
+            ):
+                output_hidden_state = not isinstance(image_proj_layer, ImageProjection)
+                single_image_embeds, single_negative_image_embeds = self.encode_image(
+                    single_ip_adapter_image, device, 1, output_hidden_state
+                )
+                single_image_embeds = torch.stack([single_image_embeds] * num_images_per_prompt, dim=0)
+                single_negative_image_embeds = torch.stack(
+                    [single_negative_image_embeds] * num_images_per_prompt, dim=0
+                )
+
+                if do_classifier_free_guidance:
+                    single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
+                    single_image_embeds = single_image_embeds.to(device)
+
+                image_embeds.append(single_image_embeds)
+        else:
+            repeat_dims = [1]
+            image_embeds = []
+            for single_image_embeds in ip_adapter_image_embeds:
+                if do_classifier_free_guidance:
+                    single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2)
+                    single_image_embeds = single_image_embeds.repeat(
+                        num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
+                    )
+                    single_negative_image_embeds = single_negative_image_embeds.repeat(
+                        num_images_per_prompt, *(repeat_dims * len(single_negative_image_embeds.shape[1:]))
+                    )
+                    single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
+                else:
+                    single_image_embeds = single_image_embeds.repeat(
+                        num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
+                    )
+                image_embeds.append(single_image_embeds)
+
+        return image_embeds
+
+    def _get_add_time_ids(
+        self,
+        original_size,
+        crops_coords_top_left,
+        target_size,
+        aesthetic_score,
+        negative_aesthetic_score,
+        negative_original_size,
+        negative_crops_coords_top_left,
+        negative_target_size,
+        dtype,
+        text_encoder_projection_dim=None,
+    ):
+        if self.config.requires_aesthetics_score:
+            add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,))
+            add_neg_time_ids = list(
+                negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,)
+            )
+        else:
+            add_time_ids = list(original_size + crops_coords_top_left + target_size)
+            add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_target_size)
+
+        passed_add_embed_dim = (
+            self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim
+        )
+        expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
+
+        if (
+            expected_add_embed_dim > passed_add_embed_dim
+            and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model."
+            )
+        elif (
+            expected_add_embed_dim < passed_add_embed_dim
+            and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model."
+            )
+        elif 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)
+        add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype)
+
+        return add_time_ids, add_neg_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)
+
+    # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
+    def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=torch.float32):
+        """
+        See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
+
+        Args:
+            timesteps (`torch.Tensor`):
+                generate embedding vectors at these timesteps
+            embedding_dim (`int`, *optional*, defaults to 512):
+                dimension of the embeddings to generate
+            dtype:
+                data type of the generated embeddings
+
+        Returns:
+            `torch.FloatTensor`: Embedding vectors with shape `(len(timesteps), embedding_dim)`
+        """
+        assert len(w.shape) == 1
+        w = w * 1000.0
+
+        half_dim = embedding_dim // 2
+        emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
+        emb = w.to(dtype)[:, None] * emb[None, :]
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+        if embedding_dim % 2 == 1:  # zero pad
+            emb = torch.nn.functional.pad(emb, (0, 1))
+        assert emb.shape == (w.shape[0], embedding_dim)
+        return emb
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def guidance_rescale(self):
+        return self._guidance_rescale
+
+    @property
+    def clip_skip(self):
+        return self._clip_skip
+
+    # 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.
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
+
+    @property
+    def cross_attention_kwargs(self):
+        return self._cross_attention_kwargs
+
+    @property
+    def denoising_end(self):
+        return self._denoising_end
+
+    @property
+    def denoising_start(self):
+        return self._denoising_start
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @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,
+        image: PipelineImageInput = None,
+        strength: float = 0.3,
+        num_inference_steps: int = 50,
+        timesteps: List[int] = None,
+        denoising_start: Optional[float] = None,
+        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,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        original_size: Tuple[int, int] = None,
+        crops_coords_top_left: Tuple[int, int] = (0, 0),
+        target_size: Tuple[int, int] = None,
+        negative_original_size: Optional[Tuple[int, int]] = None,
+        negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
+        negative_target_size: Optional[Tuple[int, int]] = None,
+        aesthetic_score: float = 6.0,
+        negative_aesthetic_score: float = 2.5,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        **kwargs,
+    ):
+        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
+            image (`torch.FloatTensor` or `PIL.Image.Image` or `np.ndarray` or `List[torch.FloatTensor]` or `List[PIL.Image.Image]` or `List[np.ndarray]`):
+                The image(s) to modify with the pipeline.
+            strength (`float`, *optional*, defaults to 0.3):
+                Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image`
+                will be used as a starting point, adding more noise to it the larger the `strength`. The number of
+                denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will
+                be maximum and the denoising process will run for the full number of iterations specified in
+                `num_inference_steps`. A value of 1, therefore, essentially ignores `image`. Note that in the case of
+                `denoising_start` being declared as an integer, the value of `strength` will be ignored.
+            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.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            denoising_start (`float`, *optional*):
+                When specified, indicates the fraction (between 0.0 and 1.0) of the total denoising process to be
+                bypassed before it is initiated. Consequently, the initial part of the denoising process is skipped and
+                it is assumed that the passed `image` is a partly denoised image. Note that when this is specified,
+                strength will be ignored. The `denoising_start` parameter is particularly beneficial when this pipeline
+                is integrated into a "Mixture of Denoisers" multi-pipeline setup, as detailed in [**Refine Image
+                Quality**](https://huggingface.co/docs/diffusers/using-diffusers/sdxl#refine-image-quality).
+            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 (ca. final 20% of timesteps still needed) and should be
+                denoised by a successor pipeline that has `denoising_start` set to 0.8 so that it only denoises the
+                final 20% of 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 [**Refine Image
+                Quality**](https://huggingface.co/docs/diffusers/using-diffusers/sdxl#refine-image-quality).
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. 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.
+            ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+            ip_adapter_image_embeds (`List[torch.FloatTensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters.
+                Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding
+                if `do_classifier_free_guidance` is set to `True`.
+                If not provided, embeddings are computed from the `ip_adapter_image` 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.StableDiffusionXLPipelineOutput`] instead of a
+                plain tuple.
+            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.0):
+                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 `(height, width)` 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 `(height, width)`. 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).
+            negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                To negatively condition the generation process based on a specific image resolution. 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). For more
+                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
+            negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+                To negatively condition the generation process based on a specific crop coordinates. 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). For more
+                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
+            negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                To negatively condition the generation process based on a target image resolution. It should be as same
+                as the `target_size` for most cases. 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). For more
+                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
+            aesthetic_score (`float`, *optional*, defaults to 6.0):
+                Used to simulate an aesthetic score of the generated image by influencing the positive text condition.
+                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).
+            negative_aesthetic_score (`float`, *optional*, defaults to 2.5):
+                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). Can be used to
+                simulate an aesthetic score of the generated image by influencing the negative text condition.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
+            `tuple. When returning a tuple, the first element is a list with the generated images.
+        """
+
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+
+        if callback is not None:
+            deprecate(
+                "callback",
+                "1.0.0",
+                "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+            )
+        if callback_steps is not None:
+            deprecate(
+                "callback_steps",
+                "1.0.0",
+                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+            )
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            strength,
+            num_inference_steps,
+            callback_steps,
+            negative_prompt,
+            negative_prompt_2,
+            prompt_embeds,
+            negative_prompt_embeds,
+            ip_adapter_image,
+            ip_adapter_image_embeds,
+            callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._clip_skip = clip_skip
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._denoising_end = denoising_end
+        self._denoising_start = denoising_start
+        self._interrupt = False
+        # 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
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            self.cross_attention_kwargs.get("scale", None) if self.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=self.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,
+            clip_skip=self.clip_skip,
+        )
+        # 4. Preprocess image
+        image = self.image_processor.preprocess(image)
+
+        # 5. Prepare timesteps
+        def denoising_value_valid(dnv):
+            return isinstance(dnv, float) and 0 < dnv < 1
+
+        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+        timesteps, num_inference_steps = self.get_timesteps(
+            num_inference_steps,
+            strength,
+            device,
+            denoising_start=self.denoising_start if denoising_value_valid(self.denoising_start) else None,
+        )
+        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
+
+        add_noise = True if self.denoising_start is None else False
+        # 6. Prepare latent variables
+        latents = self.prepare_latents(
+            image,
+            latent_timestep,
+            batch_size,
+            num_images_per_prompt,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            add_noise,
+        )
+        # 7. Prepare extra step kwargs.
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        height, width = latents.shape[-2:]
+        height = height * self.vae_scale_factor
+        width = width * self.vae_scale_factor
+
+        original_size = original_size or (height, width)
+        target_size = target_size or (height, width)
+
+        # 8. Prepare added time ids & embeddings
+        if negative_original_size is None:
+            negative_original_size = original_size
+        if negative_target_size is None:
+            negative_target_size = target_size
+
+        add_text_embeds = pooled_prompt_embeds
+        if self.text_encoder_2 is None:
+            text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
+        else:
+            text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
+
+        add_time_ids, add_neg_time_ids = self._get_add_time_ids(
+            original_size,
+            crops_coords_top_left,
+            target_size,
+            aesthetic_score,
+            negative_aesthetic_score,
+            negative_original_size,
+            negative_crops_coords_top_left,
+            negative_target_size,
+            dtype=prompt_embeds.dtype,
+            text_encoder_projection_dim=text_encoder_projection_dim,
+        )
+        add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1)
+
+        if self.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_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1)
+            add_time_ids = torch.cat([add_neg_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)
+
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+                self.do_classifier_free_guidance,
+            )
+
+        # 9. Denoising loop
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+        # 9.1 Apply denoising_end
+        if (
+            self.denoising_end is not None
+            and self.denoising_start is not None
+            and denoising_value_valid(self.denoising_end)
+            and denoising_value_valid(self.denoising_start)
+            and self.denoising_start >= self.denoising_end
+        ):
+            raise ValueError(
+                f"`denoising_start`: {self.denoising_start} cannot be larger than or equal to `denoising_end`: "
+                + f" {self.denoising_end} when using type float."
+            )
+        elif self.denoising_end is not None and denoising_value_valid(self.denoising_end):
+            discrete_timestep_cutoff = int(
+                round(
+                    self.scheduler.config.num_train_timesteps
+                    - (self.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]
+
+        # 9.2 Optionally get Guidance Scale Embedding
+        timestep_cond = None
+        if self.unet.config.time_cond_proj_dim is not None:
+            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)
+            timestep_cond = self.get_guidance_scale_embedding(
+                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+            ).to(device=device, dtype=latents.dtype)
+
+        self._num_timesteps = len(timesteps)
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if self.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}
+                if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+                    added_cond_kwargs["image_embeds"] = image_embeds
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    timestep_cond=timestep_cond,
+                    cross_attention_kwargs=self.cross_attention_kwargs,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                if self.do_classifier_free_guidance and self.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=self.guidance_rescale)
+
+                # compute the previous noisy sample x_t -> x_t-1
+
+                ########################## changed section ##########################
+                # dift_features = self.unet.latent_store.dift_features[f'{t.item()}_0']
+                dift_features = self.unet.latent_store.smoothed_dift_features[f'{t.item()}_1']
+
+
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 
+                                              noise_pred_uncond= noise_pred_uncond if self.do_classifier_free_guidance else noise_pred,
+                                              dift_features=dift_features, 
+                                              return_dict=False)[0]
+                #####################################################################
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+                    add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds)
+                    negative_pooled_prompt_embeds = callback_outputs.pop(
+                        "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
+                    )
+                    add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
+                    add_neg_time_ids = callback_outputs.pop("add_neg_time_ids", add_neg_time_ids)
+
+                # 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:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if not output_type == "latent":
+            # make sure the VAE is in float32 mode, as it overflows in float16
+            needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
+
+            if needs_upcasting:
+                self.upcast_vae()
+                latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
+
+            # unscale/denormalize the latents
+            # denormalize with the mean and std if available and not None
+            has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None
+            has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None
+            if has_latents_mean and has_latents_std:
+                latents_mean = (
+                    torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype)
+                )
+                latents_std = (
+                    torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype)
+                )
+                latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean
+            else:
+                latents = latents / self.vae.config.scaling_factor
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+
+            # cast back to fp16 if needed
+            if needs_upcasting:
+                self.vae.to(dtype=torch.float16)
+        else:
+            image = latents
+
+        # 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 all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return StableDiffusionXLPipelineOutput(images=image)

model/unet_sdxl.py

@@ -0,0 +1,381 @@
+# Modified only lines 360–365; search for 'changed section' to locate the update easily.
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin
+from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.models.activations import get_activation
+from diffusers.models.attention_processor import (
+    ADDED_KV_ATTENTION_PROCESSORS,
+    CROSS_ATTENTION_PROCESSORS,
+    Attention,
+    AttentionProcessor,
+    AttnAddedKVProcessor,
+    AttnProcessor,
+)
+from diffusers.models.embeddings import (
+    GaussianFourierProjection,
+    GLIGENTextBoundingboxProjection,
+    ImageHintTimeEmbedding,
+    ImageProjection,
+    ImageTimeEmbedding,
+    TextImageProjection,
+    TextImageTimeEmbedding,
+    TextTimeEmbedding,
+    TimestepEmbedding,
+    Timesteps,
+)
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.unets.unet_2d_blocks import (
+    get_down_block,
+    get_mid_block,
+    get_up_block,
+)
+
+from diffusers.models.unets import UNet2DConditionModel
+from model.modules.dift_utils import DIFTLatentStore
+
+
+@dataclass
+class UNet2DConditionOutput(BaseOutput):
+    sample: torch.FloatTensor = None
+
+class OursUNet2DConditionModel(UNet2DConditionModel):
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: Optional[int] = None,
+        in_channels: int = 4,
+        out_channels: int = 4,
+        center_input_sample: bool = False,
+        flip_sin_to_cos: bool = True,
+        freq_shift: int = 0,
+        down_block_types: Tuple[str] = (
+            "CrossAttnDownBlock2D",
+            "CrossAttnDownBlock2D",
+            "CrossAttnDownBlock2D",
+            "DownBlock2D",
+        ),
+        mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
+        up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
+        only_cross_attention: Union[bool, Tuple[bool]] = False,
+        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+        layers_per_block: Union[int, Tuple[int]] = 2,
+        downsample_padding: int = 1,
+        mid_block_scale_factor: float = 1,
+        dropout: float = 0.0,
+        act_fn: str = "silu",
+        norm_num_groups: Optional[int] = 32,
+        norm_eps: float = 1e-5,
+        cross_attention_dim: Union[int, Tuple[int]] = 1280,
+        transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
+        reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None,
+        encoder_hid_dim: Optional[int] = None,
+        encoder_hid_dim_type: Optional[str] = None,
+        attention_head_dim: Union[int, Tuple[int]] = 8,
+        num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
+        dual_cross_attention: bool = False,
+        use_linear_projection: bool = False,
+        class_embed_type: Optional[str] = None,
+        addition_embed_type: Optional[str] = None,
+        addition_time_embed_dim: Optional[int] = None,
+        num_class_embeds: Optional[int] = None,
+        upcast_attention: bool = False,
+        resnet_time_scale_shift: str = "default",
+        resnet_skip_time_act: bool = False,
+        resnet_out_scale_factor: float = 1.0,
+        time_embedding_type: str = "positional",
+        time_embedding_dim: Optional[int] = None,
+        time_embedding_act_fn: Optional[str] = None,
+        timestep_post_act: Optional[str] = None,
+        time_cond_proj_dim: Optional[int] = None,
+        conv_in_kernel: int = 3,
+        conv_out_kernel: int = 3,
+        projection_class_embeddings_input_dim: Optional[int] = None,
+        attention_type: str = "default",
+        class_embeddings_concat: bool = False,
+        mid_block_only_cross_attention: Optional[bool] = None,
+        cross_attention_norm: Optional[str] = None,
+        addition_embed_type_num_heads: int = 64,
+        steps: List[int] = list(range(1, 1000)),
+    ):
+        super().__init__(
+            sample_size=sample_size,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            center_input_sample=center_input_sample,
+            flip_sin_to_cos=flip_sin_to_cos,
+            freq_shift=freq_shift,
+            down_block_types=down_block_types,
+            mid_block_type=mid_block_type,
+            up_block_types=up_block_types,
+            only_cross_attention=only_cross_attention,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            downsample_padding=downsample_padding,
+            mid_block_scale_factor=mid_block_scale_factor,
+            dropout=dropout,
+            act_fn=act_fn,
+            norm_num_groups=norm_num_groups,
+            norm_eps=norm_eps,
+            cross_attention_dim=cross_attention_dim,
+            transformer_layers_per_block=transformer_layers_per_block,
+            reverse_transformer_layers_per_block=reverse_transformer_layers_per_block,
+            encoder_hid_dim=encoder_hid_dim,
+            encoder_hid_dim_type=encoder_hid_dim_type,
+            attention_head_dim=attention_head_dim,
+            num_attention_heads=num_attention_heads,
+            dual_cross_attention=dual_cross_attention,
+            use_linear_projection=use_linear_projection,
+            class_embed_type=class_embed_type,
+            addition_embed_type=addition_embed_type,
+            addition_time_embed_dim=addition_time_embed_dim,
+            num_class_embeds=num_class_embeds,
+            upcast_attention=upcast_attention,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            resnet_skip_time_act=resnet_skip_time_act,
+            resnet_out_scale_factor=resnet_out_scale_factor,
+            time_embedding_type=time_embedding_type,
+            time_embedding_dim=time_embedding_dim,
+            time_embedding_act_fn=time_embedding_act_fn,
+            timestep_post_act=timestep_post_act,
+            time_cond_proj_dim=time_cond_proj_dim,
+            conv_in_kernel=conv_in_kernel,
+            conv_out_kernel=conv_out_kernel,
+            projection_class_embeddings_input_dim=projection_class_embeddings_input_dim,
+            attention_type=attention_type,
+            class_embeddings_concat=class_embeddings_concat,
+            mid_block_only_cross_attention=mid_block_only_cross_attention,
+            cross_attention_norm=cross_attention_norm,
+            addition_embed_type_num_heads=addition_embed_type_num_heads,
+        )
+
+        self.latent_store = DIFTLatentStore(steps=steps, up_ft_indices=[0, 1, 2])
+    
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        class_labels: Optional[torch.Tensor] = None,
+        timestep_cond: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        mid_block_additional_residual: Optional[torch.Tensor] = None,
+        down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[UNet2DConditionOutput, Tuple]:
+        default_overall_up_factor = 2**self.num_upsamplers
+
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
+
+        for dim in sample.shape[-2:]:
+            if dim % default_overall_up_factor != 0:
+                # Forward upsample size to force interpolation output size.
+                forward_upsample_size = True
+                break
+
+        if attention_mask is not None:
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        if encoder_attention_mask is not None:
+            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        # 0. center input if necessary
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
+
+        # 1. time
+        t_emb = self.get_time_embed(sample=sample, timestep=timestep)
+        emb = self.time_embedding(t_emb, timestep_cond)
+        aug_emb = None
+
+        class_emb = self.get_class_embed(sample=sample, class_labels=class_labels)
+        if class_emb is not None:
+            if self.config.class_embeddings_concat:
+                emb = torch.cat([emb, class_emb], dim=-1)
+            else:
+                emb = emb + class_emb
+
+        aug_emb = self.get_aug_embed(
+            emb=emb, encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
+        )
+        if self.config.addition_embed_type == "image_hint":
+            aug_emb, hint = aug_emb
+            sample = torch.cat([sample, hint], dim=1)
+
+        emb = emb + aug_emb if aug_emb is not None else emb
+
+        if self.time_embed_act is not None:
+            emb = self.time_embed_act(emb)
+
+        encoder_hidden_states = self.process_encoder_hidden_states(
+            encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
+        )
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        # 2.5 GLIGEN position net
+        if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None:
+            cross_attention_kwargs = cross_attention_kwargs.copy()
+            gligen_args = cross_attention_kwargs.pop("gligen")
+            cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)}
+
+        # 3. down
+        # we're popping the `scale` instead of getting it because otherwise `scale` will be propagated
+        # to the internal blocks and will raise deprecation warnings. this will be confusing for our users.
+        if cross_attention_kwargs is not None:
+            cross_attention_kwargs = cross_attention_kwargs.copy()
+            lora_scale = cross_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+
+        is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
+        # using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets
+        is_adapter = down_intrablock_additional_residuals is not None
+
+        if not is_adapter and mid_block_additional_residual is None and down_block_additional_residuals is not None:
+            deprecate(
+                "T2I should not use down_block_additional_residuals",
+                "1.3.0",
+                "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \
+                       and will be removed in diffusers 1.3.0.  `down_block_additional_residuals` should only be used \
+                       for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ",
+                standard_warn=False,
+            )
+            down_intrablock_additional_residuals = down_block_additional_residuals
+            is_adapter = True
+
+        down_block_res_samples = (sample,)
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                # For t2i-adapter CrossAttnDownBlock2D
+                additional_residuals = {}
+                if is_adapter and len(down_intrablock_additional_residuals) > 0:
+                    additional_residuals["additional_residuals"] = down_intrablock_additional_residuals.pop(0)
+
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                    **additional_residuals,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+                if is_adapter and len(down_intrablock_additional_residuals) > 0:
+                    sample += down_intrablock_additional_residuals.pop(0)
+
+            down_block_res_samples += res_samples
+
+        if is_controlnet:
+            new_down_block_res_samples = ()
+
+            for down_block_res_sample, down_block_additional_residual in zip(
+                down_block_res_samples, down_block_additional_residuals
+            ):
+                down_block_res_sample = down_block_res_sample + down_block_additional_residual
+                new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
+
+            down_block_res_samples = new_down_block_res_samples
+
+        # 4. mid
+        if self.mid_block is not None:
+            if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention:
+                sample = self.mid_block(
+                    sample,
+                    emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+            else:
+                sample = self.mid_block(sample, emb)
+
+            # To support T2I-Adapter-XL
+            if (
+                is_adapter
+                and len(down_intrablock_additional_residuals) > 0
+                and sample.shape == down_intrablock_additional_residuals[0].shape
+            ):
+                sample += down_intrablock_additional_residuals.pop(0)
+
+        if is_controlnet:
+            sample = sample + mid_block_additional_residual
+
+        # 5. up
+        for i, upsample_block in enumerate(self.up_blocks):
+            is_final_block = i == len(self.up_blocks) - 1
+
+            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+            # if we have not reached the final block and need to forward the
+            # upsample size, we do it here
+            if not is_final_block and forward_upsample_size:
+                upsample_size = down_block_res_samples[-1].shape[2:]
+            
+            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    upsample_size=upsample_size,
+                    attention_mask=attention_mask,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+            else:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    upsample_size=upsample_size,
+                )
+            
+            ########################## changed section ##########################
+            self.latent_store(sample.detach(), t=timestep, layer_index=i)
+            self.latent_store.smooth(kernel_size=3, sigma=1)
+            #####################################################################
+
+        # 6. post-process
+        if self.conv_norm_out:
+            sample = self.conv_norm_out(sample)
+            sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (sample,)
+
+        return UNet2DConditionOutput(sample=sample)
+
+

requirements.txt

@@ -0,0 +1,15 @@
+torch==2.5.1
+torchvision==0.20.1
+ml-collections==0.1.1
+diffusers>=0.29.0
+transformers>=4.46
+accelerate==0.33.0
+json-with-comments==1.2.7
+gradio
+spaces
+scipy
+einops
+scikit-learn
+opencv-python
+rembg
+onnxruntime

src/__init__.py

@@ -0,0 +1 @@
+from .ddpm_inversion import get_ddpm_inversion_scheduler

src/ddpm_inversion.py

@@ -0,0 +1,222 @@
+import torch
+from torch import Tensor
+import torch.nn.functional as F
+from typing import Optional, Union, Tuple
+from utils import normalize
+from model import freq_exp, gen_nn_map
+from src.ddpm_step import deterministic_ddpm_step
+from diffusers.schedulers.scheduling_ddim import DDIMSchedulerOutput
+
+
+# Kernel sizes for the DIFT correction at successive time-ranges
+DIFT_KERNELS: Tuple[int, int, int, int] = (12, 7, 5, 3)
+
+def _get_kernel_for_timestep(timestep: int) -> Tuple[int, int]:
+    if timestep >= 799:
+        return DIFT_KERNELS[0], 1
+    if timestep >= 599:
+        return DIFT_KERNELS[1], 1
+    if timestep >= 299:
+        return DIFT_KERNELS[2], 1
+    return DIFT_KERNELS[3], 1
+
+def step_save_latents(
+    self,
+    model_output: torch.FloatTensor,
+    timestep: int,
+    sample: torch.FloatTensor,
+    return_dict: bool = True,
+    noise_pred_uncond: Optional[torch.FloatTensor] = None,
+    **kwargs,
+):
+    timestep_index = self._timesteps.index(timestep)
+    next_timestep_index = timestep_index + 1
+    
+    u_hat_t, beta_coef = deterministic_ddpm_step(
+        model_output=model_output,
+        timestep=timestep,
+        sample=sample,
+        scheduler=self,
+    )
+
+    x_t_minus_1 = self.x_ts[next_timestep_index]
+    self.x_ts_c_predicted.append(u_hat_t)
+
+    z_t = x_t_minus_1 - u_hat_t
+    self.latents.append(z_t)
+
+    z_t, _ = normalize(z_t, timestep_index, self._config.max_norm_zs)
+
+    x_t_minus_1_predicted = u_hat_t + z_t
+
+    if not return_dict:
+        return (x_t_minus_1_predicted,)
+
+    return DDIMSchedulerOutput(prev_sample=x_t_minus_1, pred_original_sample=None)
+
+
+def step_use_latents(
+    self,
+    model_output: torch.FloatTensor,
+    timestep: int,
+    sample: torch.FloatTensor,
+    return_dict: bool = True,
+    noise_pred_uncond: Optional[torch.FloatTensor] = None,
+    **kwargs,
+):
+    timestep_index = self._timesteps.index(timestep)
+    next_timestep_index = timestep_index + 1
+
+    z_t = self.latents[next_timestep_index]
+    _, normalize_coefficient = normalize(
+        z_t,
+        timestep_index,
+        self._config.max_norm_zs,
+    )
+
+    x_t_hat_c_hat, beta_coef = deterministic_ddpm_step(
+        model_output=model_output,
+        timestep=timestep,
+        sample=sample,
+        scheduler=self,
+    )
+
+    x_t_minus_1_exact = self.x_ts[next_timestep_index]
+    x_t_minus_1_exact = x_t_minus_1_exact.expand_as(x_t_hat_c_hat)
+
+    x_t_c_predicted: torch.Tensor = self.x_ts_c_predicted[next_timestep_index]
+
+    x_t_c = x_t_c_predicted[0].expand_as(x_t_hat_c_hat)
+    
+    mask: Optional[Tensor] = kwargs.get("mask", None)
+    if mask is not None and timestep > 300:
+        mask = mask.to(x_t_hat_c_hat.device)
+        movement_intensifier = kwargs.get("movement_intensifier", 0.0)
+
+        if timestep > 900 and movement_intensifier > 0.0:
+            latent_mask_h, *_  = freq_exp(
+                x_t_hat_c_hat[1:], 
+                "auto_mask", 
+                None, 
+                mask.unsqueeze(0),
+                movement_intensifier
+                )
+            x_t_hat_c_hat[1:] = latent_mask_h
+
+        x_t_hat_c_hat[-1] = x_t_hat_c_hat[-1] * mask + (1-mask) * x_t_c[-1]
+
+    edit_prompts_num = model_output.size(0) // 2
+    x_t_hat_c_indices = (
+        0,
+        edit_prompts_num,
+    )
+    edit_images_indices = (
+        edit_prompts_num,
+        (model_output.size(0)),
+    )
+
+    x_t_hat_c = torch.zeros_like(x_t_hat_c_hat)
+    x_t_hat_c[edit_images_indices[0] : edit_images_indices[1]] = x_t_hat_c_hat[
+        x_t_hat_c_indices[0] : x_t_hat_c_indices[1]
+    ]
+
+    w1 = kwargs.get("w1", 1.9)
+    cross_prompt_term = x_t_hat_c_hat - x_t_hat_c
+    cross_trajectory_term = x_t_hat_c - normalize_coefficient * x_t_c
+
+    x_t_minus_1_hat_ = (
+        normalize_coefficient * x_t_minus_1_exact
+        + cross_trajectory_term
+        + w1 * cross_prompt_term
+    )
+
+    x_t_minus_1_hat_[x_t_hat_c_indices[0] : x_t_hat_c_indices[1]] = x_t_minus_1_hat_[
+        edit_images_indices[0] : edit_images_indices[1]
+    ]
+    
+    dift_timestep = kwargs.get("dift_timestep", 700)
+    
+    if timestep < dift_timestep and kwargs.get("apply_dift_correction", False):
+        z_t = torch.cat([z_t]*x_t_hat_c_hat.shape[0], dim=0)
+
+        dift_features: Optional[Tensor] = kwargs.get("dift_features", None)
+        dift_s, _, dift_t = dift_features.chunk(3)
+
+        resized_src_features = F.interpolate(dift_s[0].unsqueeze(0), size=z_t.shape[-1], mode='bilinear', align_corners=False).squeeze(0)
+        resized_tgt_features = F.interpolate(dift_t[0].unsqueeze(0), size=z_t.shape[-1], mode='bilinear', align_corners=False).squeeze(0)
+        
+        kernel_size, stride = _get_kernel_for_timestep(timestep)
+        torch.cuda.empty_cache()
+        
+        updated_z_t = gen_nn_map(z_t[1], resized_src_features, resized_tgt_features, 
+                                 kernel_size=kernel_size, stride=stride,
+                                 device=z_t.device, timestep=timestep)
+
+        alpha = 1.0
+        z_t[1] = alpha * updated_z_t + (1 - alpha) * z_t[1]
+
+        x_t_minus_1_hat = x_t_hat_c_hat + z_t * normalize_coefficient
+    else:
+        x_t_minus_1_hat = x_t_minus_1_hat_
+
+    if not return_dict:
+        return (x_t_minus_1_hat,)
+
+    return DDIMSchedulerOutput(
+        prev_sample=x_t_minus_1_hat,
+        pred_original_sample=None,
+    )
+
+
+def get_ddpm_inversion_scheduler(
+    scheduler,
+    config,
+    timesteps,
+    latents,
+    x_ts,
+    **kwargs,
+):
+    def step(
+        model_output: torch.FloatTensor,
+        timestep: int,
+        sample: torch.FloatTensor,
+        eta: float = 0.0,
+        use_clipped_model_output: bool = False,
+        generator=None,
+        variance_noise: Optional[torch.FloatTensor] = None,
+        noise_pred_uncond: Optional[torch.FloatTensor] = None,
+        dift_features: Optional[torch.FloatTensor] = None,
+        return_dict: bool = True,
+    ):
+        # predict and save x_t_c
+        res_inv = step_save_latents(
+            scheduler,
+            model_output[:1, :, :, :],
+            timestep,
+            sample[:1, :, :, :],
+            return_dict,
+            noise_pred_uncond[:1, :, :, :],
+            **kwargs,
+        )
+
+        res_inf = step_use_latents(
+            scheduler,
+            model_output[1:, :, :, :],
+            timestep,
+            sample[1:, :, :, :],
+            return_dict,
+            noise_pred_uncond[1:, :, :, :],
+            dift_features=dift_features,
+            **kwargs,
+        )
+        res = (torch.cat((res_inv[0], res_inf[0]), dim=0),)
+        return res
+
+    scheduler._timesteps = timesteps
+    scheduler._config = config
+    scheduler.latents = latents
+    scheduler.x_ts = x_ts
+    scheduler.x_ts_c_predicted = [None]
+    scheduler.step = step
+    return scheduler
+

src/ddpm_step.py

@@ -0,0 +1,70 @@
+import torch
+from typing import Union
+
+def deterministic_ddpm_step(
+    model_output: torch.FloatTensor,
+    timestep: Union[float, torch.FloatTensor],
+    sample: torch.FloatTensor,
+    scheduler,
+):
+    """
+    Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+    process from the learned model outputs (most often the predicted noise).
+    """
+    t = timestep
+
+    prev_t = scheduler.previous_timestep(t)
+
+    if model_output.shape[1] == sample.shape[1] * 2 and scheduler.variance_type in [
+        "learned",
+        "learned_range",
+    ]:
+        model_output, predicted_variance = torch.split(
+            model_output, sample.shape[1], dim=1
+        )
+    else:
+        predicted_variance = None
+
+    # 1. compute alphas, betas
+    alpha_prod_t = scheduler.alphas_cumprod[t]
+    alpha_prod_t_prev = (
+        scheduler.alphas_cumprod[prev_t] if prev_t >= 0 else scheduler.one
+    )
+    beta_prod_t = 1 - alpha_prod_t
+    beta_prod_t_prev = 1 - alpha_prod_t_prev
+    current_alpha_t = alpha_prod_t / alpha_prod_t_prev
+    current_beta_t = 1 - current_alpha_t
+
+    # 2. compute predicted original sample from predicted noise also called
+    # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
+    if scheduler.config.prediction_type == "epsilon":
+        pred_original_sample = (
+            sample - beta_prod_t ** (0.5) * model_output
+        ) / alpha_prod_t ** (0.5)
+    elif scheduler.config.prediction_type == "sample":
+        pred_original_sample = model_output
+    elif scheduler.config.prediction_type == "v_prediction":
+        pred_original_sample = (alpha_prod_t**0.5) * sample - (
+            beta_prod_t**0.5
+        ) * model_output
+    else:
+        raise ValueError(
+            f"prediction_type given as {scheduler.config.prediction_type} must be one of `epsilon`, `sample` or"
+            " `v_prediction`  for the DDPMScheduler."
+        )
+
+    # 3. Clip or threshold "predicted x_0"
+    if scheduler.config.thresholding:
+        pred_original_sample = scheduler._threshold_sample(pred_original_sample)
+    elif scheduler.config.clip_sample:
+        pred_original_sample = pred_original_sample.clamp(
+            -scheduler.config.clip_sample_range, scheduler.config.clip_sample_range
+        )
+
+    current_sample_coeff = current_alpha_t ** (0.5) * beta_prod_t_prev / beta_prod_t
+    coef_D = current_sample_coeff * (beta_prod_t ** (0.5))   ## it is equal to coef_D
+    pred_prev_sample = (alpha_prod_t_prev ** (0.5) * pred_original_sample) + (
+        coef_D * model_output
+    )
+
+    return pred_prev_sample, coef_D

utils/__init__.py

@@ -0,0 +1,3 @@
+from .args import get_args
+from .pipeline_utils import load_pipeline, set_pipeline, encode_image, create_xts
+from .utils import extract_mask, find_smallest_key_with_suffix, normalize

utils/args.py

@@ -0,0 +1,65 @@
+import argparse
+
+def add_general_arguments(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--cache_dir", type=str, default=None)
+    parser.add_argument("--prompts_file", type=str, default="")
+    parser.set_defaults(fp16=False)
+    parser.add_argument("--fp16", action="store_true")
+    # parser.add_argument("--seeds", type=int, nargs='+', default=[7], help="List of seed values (e.g., --seed 22 42)")
+    parser.add_argument("--seed", type=int, default=7, help="Seed value for random number generation.")
+    parser.add_argument("--output_dir", type=str, default="output")
+    parser.add_argument("--eval_dataset_folder", type=str, default="dataset")
+    parser.add_argument("--num_of_timesteps", type=int, default=5)  # 3 or 4
+
+def add_extra_arguments(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--guidance_scale", type=float, default=0.0, help="Guidance scale value.")
+    parser.add_argument("--apply_dift_correction", action="store_true", help="Apply DIFT correction.")
+    parser.set_defaults(apply_dift_correction=False)
+    parser.add_argument("--w1", type=float, default=1.9, help="Weight for CTRL-X mode.")
+    parser.add_argument("--support_new_object", action="store_true", help="Enable support for new object detection.")
+    parser.add_argument("--mode", type=str, default="slerp_dift", help="Attention Type (e.g., normal, slerp, lerp, ...).")
+    parser.add_argument("--dift_timestep", type=int, default=400, help="DIFT timestep.")
+    parser.add_argument("--movement_intensifier", type=float, default=0.2, help="Movement intensifier factor.")
+    parser.add_argument("--structural_alignment", action="store_true", help="Enable structural alignment.")
+
+def add_editing_arguments(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--max_norm_zs", type=float, nargs="+", default=[-1, -1, -1, 15.5],
+                        help="A list of floats for max_norm_zs.")
+    parser.add_argument("--noise_shift_delta", type=float, default=1)
+    parser.add_argument("--noise_timesteps", type=int, nargs="+", default=[799, 499, 199, 0],
+                        help="A list of ints for noise_timesteps.")
+    parser.add_argument("--timesteps", type=int, nargs="+", default=[999, 799, 499, 199],
+                        help="A list of ints for timesteps.")
+    parser.add_argument("--num_steps_inversion", type=int, default=5)
+    parser.add_argument("--step_start", type=int, default=1)
+
+def check_args(args):
+    if args.num_of_timesteps not in [3, 4, 5, 10]:
+        raise ValueError("num_timesteps must be 3, 4, or 5 or 10")
+
+    if args.timesteps is not None:
+        num_steps_actual = len(args.timesteps)
+    else:
+        num_steps_actual = args.num_steps_inversion - args.step_start
+    
+    if isinstance(args.max_norm_zs, (int, float)):
+        args.max_norm_zs = [args.max_norm_zs] * num_steps_actual
+    
+    assert (
+        len(args.max_norm_zs) == num_steps_actual
+    ), f"len(args.max_norm_zs) ({len(args.max_norm_zs)}) != num_steps_actual ({num_steps_actual})"
+
+    assert args.noise_timesteps is None or len(args.noise_timesteps) == (
+        num_steps_actual
+    ), f"len(args.noise_timesteps) ({len(args.noise_timesteps)}) != num_steps_actual ({num_steps_actual})"
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    add_general_arguments(parser)
+    add_editing_arguments(parser)
+    add_extra_arguments(parser)
+    args = parser.parse_args()
+    check_args(args)
+    return args
+

utils/dino_utils.py

@@ -0,0 +1,38 @@
+import torch
+from torchvision import transforms
+
+class DINOv2Processor:
+    def __init__(self, model_name="dinov2_vitb14", device="cpu", image_size=518):
+        self.model_name = model_name
+        self.device = device
+        self.image_size = image_size
+        self.model = self._load_model()
+
+    def _load_model(self):
+        model = torch.hub.load('facebookresearch/dinov2', self.model_name)
+        model.eval()
+        model.to(self.device)
+        return model
+
+    def _preprocess_image(self, image):
+        preprocess = transforms.Compose([
+            transforms.Resize(self.image_size, interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.CenterCrop(self.image_size),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225]
+            ),
+        ])
+        return preprocess(image)
+
+    def compute_similarity(self, pil_image1, pil_image2):
+        img1_t = self._preprocess_image(pil_image1).unsqueeze(0).to(self.device)
+        img2_t = self._preprocess_image(pil_image2).unsqueeze(0).to(self.device)
+        with torch.no_grad():
+            feat1 = self.model(img1_t)
+            feat2 = self.model(img2_t)
+        feat1 = feat1 / feat1.norm(dim=1, keepdim=True)
+        feat2 = feat2 / feat2.norm(dim=1, keepdim=True)
+        similarity = (feat1 * feat2).sum(dim=1)
+        return similarity.item()

utils/general_utils.py

@@ -0,0 +1,155 @@
+import torch 
+import torch.nn.functional as F
+import io
+import cv2
+import numpy as np
+from PIL import Image
+from rembg import remove
+
+
+def normalize(
+    z_t,
+    i,
+    max_norm_zs,
+):
+    max_norm = max_norm_zs[i]
+    if max_norm < 0:
+        return z_t, 1
+
+    norm = torch.norm(z_t)
+    if norm < max_norm:
+        return z_t, 1
+
+    coeff = max_norm / norm
+    z_t = z_t * coeff
+    return z_t, coeff
+
+def normalize2(x, dim):
+    x_mean = x.mean(dim=dim, keepdim=True)
+    x_std = x.std(dim=dim, keepdim=True)
+    x_normalized = (x - x_mean) / x_std
+    return x_normalized
+
+def find_lambda_via_newton_batched(Qp, K_source, K_target, max_iter=50, tol=1e-7):
+    dot_QpK_source = torch.einsum("bcd,bmd->bcm", Qp, K_source) # shape [B]
+    dot_QpK_target = torch.einsum("bcd,bmd->bcm", Qp, K_target) # shape [B]
+    X = torch.exp(dot_QpK_source)
+
+    lmbd = torch.zeros([1], device=Qp.device, dtype=Qp.dtype) + 0.7
+    for it in range(max_iter):
+        y = torch.exp(lmbd * dot_QpK_target)
+        Z = (X + y).sum(dim=(2), keepdim=True)
+        x = X / Z
+        y = y / Z
+        val = (x.sum(dim=(1,2)) - y.sum(dim=(1,2))).sum()
+
+        grad = - (dot_QpK_target * y).sum()
+
+        if not (val.abs() > tol and grad.abs() > 1e-12):
+            break
+
+        lmbd = lmbd - val / grad
+        if lmbd.item() < 0.4:
+            return 0.1
+        elif lmbd.item() > 0.9:
+            return 0.65
+        
+    return lmbd.item()
+
+def find_lambda_via_super_halley(Qp, K_source, K_target, max_iter=50, tol=1e-7):
+    dot_QpK_source = torch.einsum("bcd,bmd->bcm", Qp, K_source)
+    dot_QpK_target = torch.einsum("bcd,bmd->bcm", Qp, K_target)
+    X = torch.exp(dot_QpK_source)
+
+    lmbd = torch.zeros([], device=Qp.device, dtype=Qp.dtype) + 0.8
+
+    for it in range(max_iter):
+        y = torch.exp(lmbd * dot_QpK_target)
+
+        Z = (X + y).sum(dim=2, keepdim=True)
+        x = X / Z
+        y = y / Z
+        
+        val = (x.sum(dim=(1,2)) - y.sum(dim=(1,2))).sum()
+
+        grad = - (dot_QpK_target * y).sum()
+
+        f2 = - (dot_QpK_target**2 * y).sum()
+
+        if not (val.abs() > tol and grad.abs() > 1e-12):
+            break
+
+        denom = grad**2 - val * f2
+        if denom.abs() < 1e-20:
+            break
+
+        update = (val * grad) / denom
+        lmbd = lmbd - update
+
+        print(f"iter={it}, λ={lmbd.item():.6f}, val={val.item():.6e}, grad={grad.item():.6e}")
+
+    return lmbd
+
+def find_smallest_key_with_suffix(features_dict: dict, suffix: str = "_1") -> str:
+        smallest_key = None
+        smallest_number = float('inf')
+        for key in features_dict.keys():
+            if key.endswith(suffix):
+                try:
+                    number = int(key.split('_')[0])
+                    if number < smallest_number:
+                        smallest_number = number
+                        smallest_key = key
+                except ValueError:
+                    continue
+        return smallest_key
+
+def extract_mask(masks, original_width, original_height):
+    if not masks:
+        return None
+
+    combined_mask = torch.zeros(512, 512)
+    scale_x = 512 / original_width
+    scale_y = 512 / original_height
+
+    for mask in masks:
+        start_x, start_y = mask["start_point"]
+        end_x, end_y = mask["end_point"]
+
+        start_x, end_x = min(start_x, end_x), max(start_x, end_x)
+        start_y, end_y = min(start_y, end_y), max(start_y, end_y)
+
+        scaled_start_x, scaled_start_y = int(start_x * scale_x), int(start_y * scale_y)
+        scaled_end_x, scaled_end_y = int(end_x * scale_x), int(end_y * scale_y)
+        combined_mask[scaled_start_y:scaled_end_y, scaled_start_x:scaled_end_x] += 1
+
+    binary_mask = (combined_mask > 0).float()
+    resized_mask = F.interpolate(binary_mask[None, None, :, :], size=(64, 64), mode="nearest")[0, 0]
+
+    return resized_mask
+
+def remove_foreground(pil_image, threshold=128):
+    try:
+        with io.BytesIO() as input_buffer:
+            pil_image.save(input_buffer, format="PNG")
+            input_image_bytes = input_buffer.getvalue()
+        
+        output_image_bytes = remove(input_image_bytes, alpha_matting=True)
+        
+        output_image = Image.open(io.BytesIO(output_image_bytes))
+        
+        mask = output_image.split()[-1]
+        
+        mask_array = np.array(mask)
+        binary_mask = (mask_array >= threshold).astype(np.float32)
+
+        kernel = np.ones((15, 15), np.uint8)
+        binary_mask = cv2.erode(binary_mask, kernel, iterations=1)
+        
+        mask_tensor = torch.from_numpy(binary_mask)
+        resized_mask = F.interpolate(mask_tensor[None, None, :, :], size=(64, 64), mode="nearest")[0, 0]
+        
+        return resized_mask
+    except Exception as e:
+        print(f"Error while removing foreground: {e}")
+        return None

utils/pipeline_utils.py

@@ -0,0 +1,148 @@
+
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img import (
+    retrieve_timesteps,
+    retrieve_latents,
+)
+
+import torch
+from functools import partial
+from diffusers import DDPMScheduler
+from model.pipeline_sdxl import StableDiffusionXLImg2ImgPipeline
+
+SAMPLING_DEVICE = "cpu"  # "cuda"
+VAE_SAMPLE = "argmax"  # "argmax" or "sample"
+RESIZE_TYPE = None  # Image.LANCZOS
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def encode_image(image, pipe, generator):
+    pipe_dtype = pipe.dtype
+    image = pipe.image_processor.preprocess(image)
+    image = image.to(device=device, dtype=pipe.dtype)
+
+    if pipe.vae.config.force_upcast:
+        image = image.float()
+        pipe.vae.to(dtype=torch.float32)
+
+    init_latents = retrieve_latents(
+        pipe.vae.encode(image), generator=generator, sample_mode=VAE_SAMPLE
+    )
+
+    if pipe.vae.config.force_upcast:
+        pipe.vae.to(pipe_dtype)
+
+    init_latents = init_latents.to(pipe_dtype)
+    init_latents = pipe.vae.config.scaling_factor * init_latents
+
+    return init_latents
+
+def create_xts(
+    noise_shift_delta,
+    noise_timesteps,
+    generator,
+    scheduler,
+    timesteps,
+    x_0,
+):
+    if noise_timesteps is None:
+        noising_delta = noise_shift_delta * (timesteps[0] - timesteps[1])
+        noise_timesteps = [timestep - int(noising_delta) for timestep in timesteps]
+        # noise_timesteps = [timestep for timestep in timesteps]
+    
+    # print(noise_timesteps, timesteps)
+    first_x_0_idx = len(noise_timesteps)
+    for i in range(len(noise_timesteps)):
+        if noise_timesteps[i] <= 0:
+            first_x_0_idx = i
+            break
+
+    noise_timesteps = noise_timesteps[:first_x_0_idx]
+
+    x_0_expanded = x_0.expand(len(noise_timesteps), -1, -1, -1)
+    noise = torch.randn(
+        x_0_expanded.size(), generator=generator, device=SAMPLING_DEVICE
+    ).to(x_0.device)
+
+    x_ts = scheduler.add_noise(
+        x_0_expanded,
+        noise,
+        torch.IntTensor(noise_timesteps),
+    )
+    x_ts = [t.unsqueeze(dim=0) for t in list(x_ts)]
+    x_ts += [x_0] * (len(timesteps) - first_x_0_idx)
+    x_ts += [x_0]
+    return x_ts
+
+def load_pipeline(fp16, cache_dir):
+    kwargs = (
+        {
+            "torch_dtype": torch.float16,
+            "variant": "fp16",
+        }
+        if fp16
+        else {}
+    )
+    from model.unet_sdxl import OursUNet2DConditionModel
+    unet = OursUNet2DConditionModel.from_pretrained(
+        "stabilityai/sdxl-turbo", 
+        subfolder="unet", 
+        cache_dir=cache_dir,
+        safety_checker=None,
+        **kwargs,
+    )
+    
+    pipeline = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+        "stabilityai/sdxl-turbo",
+        unet=unet,
+        cache_dir=cache_dir,
+        safety_checker=None,
+        **kwargs,
+    )
+    
+    pipeline = pipeline.to(device)
+    pipeline.scheduler = DDPMScheduler.from_pretrained(  # type: ignore
+        "stabilityai/sdxl-turbo",
+        subfolder="scheduler",
+    )
+
+    return pipeline
+
+def set_pipeline(pipeline: StableDiffusionXLImg2ImgPipeline, num_timesteps, generator, config):
+    if config.timesteps is None:
+        denoising_start = config.step_start / config.num_steps_inversion
+        timesteps, num_inference_steps = retrieve_timesteps(
+            pipeline.scheduler, config.num_steps_inversion, device, None
+        )
+        timesteps, num_inference_steps = pipeline.get_timesteps(
+            num_inference_steps=num_inference_steps,
+            device=device,
+            denoising_start=denoising_start,
+            strength=0,
+        )
+        timesteps = timesteps.type(torch.int64)
+        pipeline.__call__ = partial(
+            pipeline.__call__,
+            num_inference_steps=config.num_steps_inversion,
+            guidance_scale=config.guidance_scale,
+            generator=generator,
+            denoising_start=denoising_start,
+            strength=0,
+        )
+        pipeline.scheduler.set_timesteps(
+            timesteps=timesteps.cpu(),
+        )
+    else:
+        timesteps = torch.tensor(config.timesteps, dtype=torch.int64)
+        pipeline.__call__ = partial(
+            pipeline.__call__,
+            timesteps=timesteps,
+            guidance_scale=config.guidance_scale,
+            denoising_start=0,
+            strength=1,
+        )
+        pipeline.scheduler.set_timesteps(
+            timesteps=config.timesteps,  # device=pipeline.device
+        )
+    timesteps = [torch.tensor(t) for t in timesteps.tolist()]
+    return timesteps, config
+

utils/utils.py

@@ -0,0 +1,129 @@
+import torch 
+import torch.nn.functional as F
+import io
+import cv2
+import numpy as np
+from PIL import Image
+
+
+def normalize(
+    z_t,
+    i,
+    max_norm_zs,
+):
+    max_norm = max_norm_zs[i]
+    if max_norm < 0:
+        return z_t, 1
+
+    norm = torch.norm(z_t)
+    if norm < max_norm:
+        return z_t, 1
+
+    coeff = max_norm / norm
+    z_t = z_t * coeff
+    return z_t, coeff
+
+def normalize2(x, dim):
+    x_mean = x.mean(dim=dim, keepdim=True)
+    x_std = x.std(dim=dim, keepdim=True)
+    x_normalized = (x - x_mean) / x_std
+    return x_normalized
+
+def find_lambda_via_newton_batched(Qp, K_source, K_target, max_iter=50, tol=1e-7):
+    dot_QpK_source = torch.einsum("bcd,bmd->bcm", Qp, K_source) # shape [B]
+    dot_QpK_target = torch.einsum("bcd,bmd->bcm", Qp, K_target) # shape [B]
+    X = torch.exp(dot_QpK_source)
+
+    lmbd = torch.zeros([1], device=Qp.device, dtype=Qp.dtype) + 0.7
+    for it in range(max_iter):
+        y = torch.exp(lmbd * dot_QpK_target)
+        Z = (X + y).sum(dim=(2), keepdim=True)
+        x = X / Z
+        y = y / Z
+        val = (x.sum(dim=(1,2)) - y.sum(dim=(1,2))).sum()
+
+        grad = - (dot_QpK_target * y).sum()
+
+        if not (val.abs() > tol and grad.abs() > 1e-12):
+            break
+
+        lmbd = lmbd - val / grad
+        if lmbd.item() < 0.4:
+            return 0.1
+        elif lmbd.item() > 0.9:
+            return 0.65
+        
+    return lmbd.item()
+
+def find_lambda_via_super_halley(Qp, K_source, K_target, max_iter=50, tol=1e-7):
+    dot_QpK_source = torch.einsum("bcd,bmd->bcm", Qp, K_source)
+    dot_QpK_target = torch.einsum("bcd,bmd->bcm", Qp, K_target)
+    X = torch.exp(dot_QpK_source)
+
+    lmbd = torch.zeros([], device=Qp.device, dtype=Qp.dtype) + 0.8
+
+    for it in range(max_iter):
+        y = torch.exp(lmbd * dot_QpK_target)
+
+        Z = (X + y).sum(dim=2, keepdim=True)
+        x = X / Z
+        y = y / Z
+        
+        val = (x.sum(dim=(1,2)) - y.sum(dim=(1,2))).sum()
+
+        grad = - (dot_QpK_target * y).sum()
+
+        f2 = - (dot_QpK_target**2 * y).sum()
+
+        if not (val.abs() > tol and grad.abs() > 1e-12):
+            break
+
+        denom = grad**2 - val * f2
+        if denom.abs() < 1e-20:
+            break
+
+        update = (val * grad) / denom
+        lmbd = lmbd - update
+
+        print(f"iter={it}, λ={lmbd.item():.6f}, val={val.item():.6e}, grad={grad.item():.6e}")
+
+    return lmbd
+
+def find_smallest_key_with_suffix(features_dict: dict, suffix: str = "_1") -> str:
+        smallest_key = None
+        smallest_number = float('inf')
+        for key in features_dict.keys():
+            if key.endswith(suffix):
+                try:
+                    number = int(key.split('_')[0])
+                    if number < smallest_number:
+                        smallest_number = number
+                        smallest_key = key
+                except ValueError:
+                    continue
+        return smallest_key
+
+def extract_mask(masks, original_width, original_height):
+    if not masks:
+        return None
+
+    combined_mask = torch.zeros(512, 512)
+    scale_x = 512 / original_width
+    scale_y = 512 / original_height
+
+    for mask in masks:
+        start_x, start_y = mask["start_point"]
+        end_x, end_y = mask["end_point"]
+
+        start_x, end_x = min(start_x, end_x), max(start_x, end_x)
+        start_y, end_y = min(start_y, end_y), max(start_y, end_y)
+
+        scaled_start_x, scaled_start_y = int(start_x * scale_x), int(start_y * scale_y)
+        scaled_end_x, scaled_end_y = int(end_x * scale_x), int(end_y * scale_y)
+        combined_mask[scaled_start_y:scaled_end_y, scaled_start_x:scaled_end_x] += 1
+
+    binary_mask = (combined_mask > 0).float()
+    resized_mask = F.interpolate(binary_mask[None, None, :, :], size=(64, 64), mode="nearest")[0, 0]
+
+    return resized_mask
+

visualization/__init__.py

@@ -0,0 +1 @@
+from .image_utils import save_results

visualization/draw_box.py

@@ -0,0 +1,44 @@
+import cv2
+
+bbox_start = (-1, -1)
+bbox_end = (-1, -1)
+drawing = False  
+
+image_path = 'images/fatty-corgi.jpg'
+
+image = cv2.resize(cv2.imread(image_path), (512, 512))
+image_copy = image.copy()
+
+def draw_bbox(event, x, y, flags, param):
+    global bbox_start, bbox_end, drawing, image
+
+    if event == cv2.EVENT_LBUTTONDOWN:
+        drawing = True
+        bbox_start = (x, y)
+        bbox_end = bbox_start
+
+    elif event == cv2.EVENT_MOUSEMOVE:
+        if drawing:
+            image = image_copy.copy()
+            cv2.rectangle(image, bbox_start, (x, y), (0, 0, 255), 2)
+
+    elif event == cv2.EVENT_LBUTTONUP:
+        drawing = False
+        bbox_end = (x, y)
+        cv2.rectangle(image, bbox_start, bbox_end, (0, 0, 255), 2)
+        # print(f"BBox Coordinates: Start: {bbox_start}, End: {bbox_end}\n")
+        print(f"bbx_start_point= ({bbox_start[0]}, {bbox_start[1]}), ")
+        print(f"bbx_end_point= ({bbox_end[0]}, {bbox_end[1]})")
+        
+        x1, y1 = bbox_start
+        x2, y2 = bbox_end
+
+cv2.namedWindow("Image")
+cv2.setMouseCallback("Image", draw_bbox)
+
+while True:
+    cv2.imshow("Image", image)
+    if cv2.waitKey(1) & 0xFF == ord('q'):
+        break
+
+cv2.destroyAllWindows()

visualization/image_utils.py

@@ -0,0 +1,106 @@
+import torch
+import numpy as np
+
+import os
+import re
+import jsonc as json
+from PIL import Image
+
+
+def img_list_to_pil(img_list, cond_image = None, seperation = 10):
+    if cond_image is not None:
+        img_list.append(cond_image)
+
+    widths, heights = zip(*(i.size for i in img_list))
+    total_width = sum(widths) + seperation * len(img_list)
+    max_height = max(heights)
+
+    new_im = Image.new('RGB', (total_width, max_height))
+
+    x_offset = 0
+    for im in img_list:
+        new_im.paste(im, (x_offset, 0))
+        x_offset += im.size[0] + seperation
+
+    return new_im
+
+
+def grid_image_visualize(images, row_size):
+    widths, heights = zip(*(i.size for i in images))
+    total_width = max(widths) * row_size + 10 * (row_size - 1)
+    max_height = max(heights) * ((len(images) + row_size - 1) // row_size)
+    new_im = Image.new('RGB', (total_width, max_height))
+
+    x_offset = 0
+    y_offset = 0
+    for i, im in enumerate(images):
+        new_im.paste(im, (x_offset, y_offset))
+        x_offset += im.size[0] + 10
+        if (i + 1) % row_size == 0:
+            x_offset = 0
+            y_offset += im.size[1]
+
+    return new_im
+
+def process_images(images, res=512):
+    res_images = []
+    for image in images:
+        crop_size = min(image.size)
+        
+        left = (image.size[0] - crop_size) // 2
+        top = (image.size[1] - crop_size) // 2
+        right = (image.size[0] + crop_size) // 2
+        bottom = (image.size[1] + crop_size) // 2
+
+        image = image.crop((left, top, right, bottom))
+        image = image.resize((res, res), Image.BILINEAR)
+        res_images.append(image)
+    return res_images
+
+def sanitize_prompt(prompt: str, max_len: int = 50) -> str:
+    sanitized = re.sub(r'[^a-zA-Z0-9_\-]+', '_', prompt)
+    return sanitized[:max_len].strip("_")
+
+def get_next_index(folder_path: str) -> int:
+    if not os.path.exists(folder_path):
+        return 0
+
+    pattern = re.compile(r'.*_(\d+)\.(?:png|json)$')
+    max_index = -1
+
+    for filename in os.listdir(folder_path):
+        match = pattern.match(filename)
+        if match:
+            idx = int(match.group(1))
+            if idx > max_index:
+                max_index = idx
+
+    return max_index + 1
+
+def save_results(
+    args,
+    source_prompt: str,
+    target_prompt: str,
+    images: Image.Image,
+):
+    src_name = sanitize_prompt(source_prompt)
+    tgt_name = sanitize_prompt(target_prompt)
+    folder_name = f"{src_name}#{tgt_name}"
+
+    output_dir = os.path.join(args.output_dir, folder_name)
+    os.makedirs(output_dir, exist_ok=True)
+
+    next_idx = get_next_index(output_dir)
+
+    concated_image = img_list_to_pil([images[0], images[-1]], cond_image=None, seperation=10)
+    concated_image.save(os.path.join(output_dir, f"concat_{next_idx}.png"))
+    images[0].save(os.path.join(output_dir, f"input_{next_idx}.png"))
+    images[-1].save(os.path.join(output_dir, f"output_{next_idx}.png"))
+
+    args_filename = f"args_{next_idx}.json"
+    args_path = os.path.join(output_dir, args_filename)
+
+    with open(args_path, "w") as f:
+        json.dump(vars(args), f, indent=4)
+
+    print(f"Saved image to {output_dir} and args to {args_path}")