Upload folder using huggingface_hub

app.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+import gradio as gr
+import spaces
+from PIL import Image
+import torch
+
+from my_run import run as run_model
+
+
+DESCRIPTION = '''# Turbo Edit
+'''
+
+@spaces.GPU
+def main_pipeline(
+        input_image: str,
+        src_prompt: str,
+        tgt_prompt: str,
+        seed: int,
+        w1: float,
+        # w2: float,
+        ):
+
+        w2 = 1.0
+        res_image = run_model(input_image, src_prompt, tgt_prompt, seed, w1, w2)
+
+        return res_image
+
+
+with gr.Blocks(css='app/style.css') as demo:
+    gr.Markdown(DESCRIPTION)
+
+    gr.HTML(
+        '''<a href="https://huggingface.co/spaces/garibida/ReNoise-Inversion?duplicate=true">
+        <img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>Duplicate the Space to run privately without waiting in queue''')
+
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(
+                label="Input image",
+                type="filepath",
+                height=512,
+                width=512
+            )
+            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=7865,
+                    step=1
+                )
+                w1 = gr.Slider(
+                    label='w',
+                    minimum=1.0,
+                    maximum=3.0,
+                    value=1.5,
+                    step=0.05
+                )
+                # w2 = gr.Slider(
+                #     label='w2',
+                #     minimum=1.0,
+                #     maximum=3.0,
+                #     value=1.0,
+                #     step=0.05
+                # )
+
+            run_button = gr.Button('Edit')
+        with gr.Column():
+            # result = gr.Gallery(label='Result')
+            result = gr.Image(
+                label="Result",
+                type="pil",
+                height=512,
+                width=512
+            )
+
+            examples = [
+                [
+                    "demo_im/WhatsApp Image 2024-05-17 at 17.32.53.jpeg", #input_image
+                    "a painting of a white cat sleeping on a lotus flower", #src_prompt
+                    "a painting of a white cat sleeping on a lotus flower", #tgt_prompt
+                    4759, #seed
+                    1.0, #w1
+                    # 1.1, #w2
+                ],
+                [
+                    "demo_im/pexels-pixabay-458976.less.png", #input_image
+                    "a squirrel standing in the grass", #src_prompt
+                    "a squirrel standing in the grass", #tgt_prompt
+                    6128, #seed
+                    1.25, #w1
+                    # 1.1, #w2
+                ],
+            ]
+
+            gr.Examples(examples=examples,
+                        inputs=[
+                            input_image,
+                            src_prompt,
+                            tgt_prompt,
+                            seed,
+                            w1,
+                            # w2,
+                        ],
+                        outputs=[
+                            result
+                        ],
+                        fn=main_pipeline,
+                        cache_examples=True)
+
+
+    inputs = [
+        input_image,
+        src_prompt,
+        tgt_prompt,
+        seed,
+        w1,
+        # w2,
+    ]
+    outputs = [
+        result
+    ]
+    run_button.click(fn=main_pipeline, inputs=inputs, outputs=outputs)
+
+demo.queue(max_size=50).launch(share=True, max_threads=100)

config.py

@@ -0,0 +1,140 @@
+from ml_collections import config_dict
+import yaml
+from diffusers.schedulers import (
+    DDIMScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+    DDPMScheduler,
+)
+from utils import (
+    deterministic_ddim_step,
+    deterministic_ddpm_step,
+    deterministic_euler_step,
+    deterministic_non_ancestral_euler_step,
+)
+
+BREAKDOWNS = ["x_t_c_hat", "x_t_hat_c", "no_breakdown", "x_t_hat_c_with_zeros"]
+SCHEDULERS = ["ddpm", "ddim", "euler", "euler_non_ancestral"]
+MODELS = [
+    "stabilityai/sdxl-turbo",
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    "CompVis/stable-diffusion-v1-4",
+]
+
+def get_num_steps_actual(cfg):
+    return (
+        cfg.num_steps_inversion
+        - cfg.step_start
+        + (1 if cfg.clean_step_timestep > 0 else 0)
+        if cfg.timesteps is None
+        else len(cfg.timesteps) + (1 if cfg.clean_step_timestep > 0 else 0)
+    )
+
+
+def get_config(args):
+    if args.config_from_file and args.config_from_file != "":
+        with open(args.config_from_file, "r") as f:
+            cfg = config_dict.ConfigDict(yaml.safe_load(f))
+
+        num_steps_actual = get_num_steps_actual(cfg)
+
+    else:
+        cfg = config_dict.ConfigDict()
+
+        cfg.seed = 2
+        cfg.self_r = 0.5
+        cfg.cross_r = 0.9
+        cfg.eta = 1
+        cfg.scheduler_type = SCHEDULERS[0]
+
+        cfg.num_steps_inversion = 50  # timesteps: 999, 799, 599, 399, 199
+        cfg.step_start = 20
+        cfg.timesteps = None
+        cfg.noise_timesteps = None
+        num_steps_actual = get_num_steps_actual(cfg)
+        cfg.ws1 = [2] * num_steps_actual
+        cfg.ws2 = [1] * num_steps_actual
+        cfg.real_cfg_scale = 0
+        cfg.real_cfg_scale_save = 0
+        cfg.breakdown = BREAKDOWNS[1]
+        cfg.noise_shift_delta = 1
+        cfg.max_norm_zs = [-1] * (num_steps_actual - 1) + [15.5]
+
+        cfg.clean_step_timestep = 0
+
+        cfg.model = MODELS[1]
+
+    if cfg.scheduler_type == "ddim":
+        cfg.scheduler_class = DDIMScheduler
+        cfg.step_function = deterministic_ddim_step
+    elif cfg.scheduler_type == "ddpm":
+        cfg.scheduler_class = DDPMScheduler
+        cfg.step_function = deterministic_ddpm_step
+    elif cfg.scheduler_type == "euler":
+        cfg.scheduler_class = EulerAncestralDiscreteScheduler
+        cfg.step_function = deterministic_euler_step
+    elif cfg.scheduler_type == "euler_non_ancestral":
+        cfg.scheduler_class = EulerDiscreteScheduler
+        cfg.step_function = deterministic_non_ancestral_euler_step
+    else:
+        raise ValueError(f"Unknown scheduler type: {cfg.scheduler_type}")
+
+    with cfg.ignore_type():
+        if isinstance(cfg.max_norm_zs, (int, float)):
+            cfg.max_norm_zs = [cfg.max_norm_zs] * num_steps_actual
+
+        if isinstance(cfg.ws1, (int, float)):
+            cfg.ws1 = [cfg.ws1] * num_steps_actual
+
+        if isinstance(cfg.ws2, (int, float)):
+            cfg.ws2 = [cfg.ws2] * num_steps_actual
+
+    if not hasattr(cfg, "update_eta"):
+        cfg.update_eta = False
+
+    if not hasattr(cfg, "save_timesteps"):
+        cfg.save_timesteps = None
+
+    if not hasattr(cfg, "scheduler_timesteps"):
+        cfg.scheduler_timesteps = None
+
+    assert (
+        cfg.scheduler_type == "ddpm" or cfg.timesteps is None
+    ), "timesteps must be None for ddim/euler"
+
+    assert (
+        len(cfg.max_norm_zs) == num_steps_actual
+    ), f"len(cfg.max_norm_zs) ({len(cfg.max_norm_zs)}) != num_steps_actual ({num_steps_actual})"
+
+    assert (
+        len(cfg.ws1) == num_steps_actual
+    ), f"len(cfg.ws1) ({len(cfg.ws1)}) != num_steps_actual ({num_steps_actual})"
+
+    assert (
+        len(cfg.ws2) == num_steps_actual
+    ), f"len(cfg.ws2) ({len(cfg.ws2)}) != num_steps_actual ({num_steps_actual})"
+
+    assert cfg.noise_timesteps is None or len(cfg.noise_timesteps) == (
+        num_steps_actual - (1 if cfg.clean_step_timestep > 0 else 0)
+    ), f"len(cfg.noise_timesteps) ({len(cfg.noise_timesteps)}) != num_steps_actual ({num_steps_actual})"
+
+    assert cfg.save_timesteps is None or len(cfg.save_timesteps) == (
+        num_steps_actual - (1 if cfg.clean_step_timestep > 0 else 0)
+    ), f"len(cfg.save_timesteps) ({len(cfg.save_timesteps)}) != num_steps_actual ({num_steps_actual})"
+
+    return cfg
+
+
+def get_config_name(config, args):
+    if args.folder_name is not None and args.folder_name != "":
+        return args.folder_name
+    timesteps_str = (
+        f"step_start {config.step_start}"
+        if config.timesteps is None
+        else f"timesteps {config.timesteps}"
+    )
+    return f"""\
+ws1 {config.ws1[0]} ws2 {config.ws2[0]} real_cfg_scale {config.real_cfg_scale} {timesteps_str} \
+real_cfg_scale_save {config.real_cfg_scale_save} seed {config.seed} max_norm_zs {config.max_norm_zs[-1]} noise_shift_delta {config.noise_shift_delta} \
+scheduler_type {config.scheduler_type} fp16 {args.fp16}\
+"""

my_run.py

@@ -0,0 +1,476 @@
+from diffusers import AutoPipelineForImage2Image
+from diffusers import DDPMScheduler
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img import retrieve_timesteps, retrieve_latents
+from diffusers.schedulers.scheduling_ddim import DDIMSchedulerOutput
+import torch
+from PIL import Image
+
+num_steps_inversion = 5
+strngth = 0.8
+generator = None
+device = "cuda"
+image_path = "edit_dataset/01.jpg"
+src_prompt = "butterfly perched on purple flower"
+tgt_prompt = "dragonfly perched on purple flower"
+ws1 = [1.5, 1.5, 1.5, 1.5]
+ws2 = [1, 1, 1, 1]
+
+def encode_image(image, pipe):
+    image = pipe.image_processor.preprocess(image)
+    image = image.to(device=device, dtype=pipeline.dtype)
+
+    if pipe.vae.config.force_upcast:
+        image = image.float()
+        pipe.vae.to(dtype=torch.float32)
+
+    if isinstance(generator, list):
+        init_latents = [
+            retrieve_latents(pipe.vae.encode(image[i : i + 1]), generator=generator[i])
+            for i in range(1)
+        ]
+        init_latents = torch.cat(init_latents, dim=0)
+    else:
+        init_latents = retrieve_latents(pipe.vae.encode(image), generator=generator)
+
+    if pipe.vae.config.force_upcast:
+        pipe.vae.to(pipeline.dtype)
+
+    init_latents = init_latents.to(pipeline.dtype)
+    init_latents = pipe.vae.config.scaling_factor * init_latents
+
+    return init_latents.to(dtype=torch.float16)
+
+# def create_xts(scheduler, timesteps, x_0, noise_shift_delta=1, generator=None):
+#     noising_delta = noise_shift_delta * (timesteps[0] - timesteps[1])
+#     noise_timesteps = [timestep - int(noising_delta) for timestep in timesteps]
+#     noise_timesteps = noise_timesteps[:3]
+
+#     x_0_expanded = x_0.expand(len(noise_timesteps), -1, -1, -1)
+#     noise = torch.randn(x_0_expanded.size(), generator=generator, device="cpu", dtype=x_0.dtype).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]
+#     return x_ts
+
+def deterministic_ddpm_step(
+    model_output: torch.FloatTensor,
+    timestep,
+    sample: torch.FloatTensor,
+    eta,
+    use_clipped_model_output,
+    generator,
+    variance_noise,
+    return_dict,
+    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).
+
+    Args:
+        model_output (`torch.FloatTensor`):
+            The direct output from learned diffusion model.
+        timestep (`float`):
+            The current discrete timestep in the diffusion chain.
+        sample (`torch.FloatTensor`):
+            A current instance of a sample created by the diffusion process.
+        generator (`torch.Generator`, *optional*):
+            A random number generator.
+        return_dict (`bool`, *optional*, defaults to `True`):
+            Whether or not to return a [`~schedulers.scheduling_ddpm.DDPMSchedulerOutput`] or `tuple`.
+
+    Returns:
+        [`~schedulers.scheduling_ddpm.DDPMSchedulerOutput`] or `tuple`:
+            If return_dict is `True`, [`~schedulers.scheduling_ddpm.DDPMSchedulerOutput`] is returned, otherwise a
+            tuple is returned where the first element is the sample tensor.
+
+    """
+    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
+        )
+
+    # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
+    # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
+    pred_original_sample_coeff = (
+        alpha_prod_t_prev ** (0.5) * current_beta_t
+    ) / beta_prod_t
+    current_sample_coeff = current_alpha_t ** (0.5) * beta_prod_t_prev / beta_prod_t
+
+    # 5. Compute predicted previous sample µ_t
+    # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
+    pred_prev_sample = (
+        pred_original_sample_coeff * pred_original_sample
+        + current_sample_coeff * sample
+    )
+
+    return pred_prev_sample
+
+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 step_save_latents(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        sample: torch.FloatTensor,
+        eta: float = 0.0,
+        use_clipped_model_output: bool = False,
+        generator=None,
+        variance_noise= None,
+        return_dict: bool = True,
+    ):
+
+    timestep_index = self._inner_index
+    next_timestep_index = timestep_index + 1
+    u_hat_t = deterministic_ddpm_step(
+        model_output=model_output,
+        timestep=timestep,
+        sample=sample,
+        eta=eta,
+        use_clipped_model_output=use_clipped_model_output,
+        generator=generator,
+        variance_noise=variance_noise,
+        return_dict=False,
+        scheduler=self,
+    )
+    x_t_minus_1 = self.x_ts[timestep_index]
+    self.x_ts_c_hat.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, [-1, -1, -1, 15.5])
+    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,
+    eta: float = 0.0,
+    use_clipped_model_output: bool = False,
+    generator=None,
+    variance_noise= None,
+    return_dict: bool = True,
+):
+    print(f'_inner_index: {self._inner_index}')
+    timestep_index = self._inner_index
+    next_timestep_index = timestep_index + 1
+    z_t = self.latents[timestep_index]  # + 1 because latents[0] is X_T
+
+    _, normalize_coefficient = normalize(
+        z_t,
+        timestep_index,
+        [-1, -1, -1, 15.5],
+    )
+
+    if normalize_coefficient == 0:
+        eta = 0
+
+    # eta = normalize_coefficient
+
+    x_t_hat_c_hat = deterministic_ddpm_step(
+        model_output=model_output,
+        timestep=timestep,
+        sample=sample,
+        eta=eta,
+        use_clipped_model_output=use_clipped_model_output,
+        generator=generator,
+        variance_noise=variance_noise,
+        return_dict=False,
+        scheduler=self,
+    )
+
+    w1 = ws1[timestep_index]
+    w2 = ws2[timestep_index]
+
+    x_t_minus_1_exact = self.x_ts[timestep_index]
+    x_t_minus_1_exact = x_t_minus_1_exact.expand_as(x_t_hat_c_hat)
+
+    x_t_c_hat: torch.Tensor = self.x_ts_c_hat[timestep_index]
+
+    x_t_c = x_t_c_hat[0].expand_as(x_t_hat_c_hat)
+    
+    zero_index_reconstruction = 0
+    edit_prompts_num = (model_output.size(0) - zero_index_reconstruction) // 2
+    x_t_hat_c_indices = (zero_index_reconstruction, edit_prompts_num + zero_index_reconstruction)
+    edit_images_indices = (
+        edit_prompts_num + zero_index_reconstruction,
+        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]
+    ]
+    v1 = x_t_hat_c_hat - x_t_hat_c
+    v2 = x_t_hat_c - normalize_coefficient * x_t_c
+
+    x_t_minus_1 = normalize_coefficient * x_t_minus_1_exact + w1 * v1 + w2 * v2
+
+    x_t_minus_1[x_t_hat_c_indices[0] : x_t_hat_c_indices[1]] = x_t_minus_1[
+        edit_images_indices[0] : edit_images_indices[1]
+    ] # update x_t_hat_c to be x_t_hat_c_hat
+    
+
+    if not return_dict:
+        return (x_t_minus_1,)
+
+    return DDIMSchedulerOutput(
+        prev_sample=x_t_minus_1,
+        pred_original_sample=None,
+    )
+
+
+class myDDPMScheduler(DDPMScheduler):
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        sample: torch.FloatTensor,
+        eta: float = 0.0,
+        use_clipped_model_output: bool = False,
+        generator=None,
+        variance_noise= None,
+        return_dict: bool = True,
+        ):
+        print(f"timestep: {timestep}")
+
+        res_inv =  step_save_latents(
+                self,
+                model_output[:1, :, :, :],
+                timestep,
+                sample[:1, :, :, :],
+                eta,
+                use_clipped_model_output,
+                generator,
+                variance_noise,
+                return_dict,
+            )
+
+        res_inf = step_use_latents(
+                self,
+                model_output[1:, :, :, :],
+                timestep,
+                sample[1:, :, :, :],
+                eta,
+                use_clipped_model_output,
+                generator,
+                variance_noise,
+                return_dict,
+            )
+        
+        self._inner_index+=1
+
+        res = (torch.cat((res_inv[0], res_inf[0]), dim=0),)
+        return res
+
+
+pipeline = AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", safety_checker = None)
+pipeline = pipeline.to("cuda")
+pipeline.scheduler = DDPMScheduler.from_pretrained(  # type: ignore
+                'stabilityai/sdxl-turbo',
+                subfolder="scheduler",
+                # cache_dir="/home/joberant/NLP_2223/giladd/test_dir/sdxl-turbo/models_cache",
+            )
+# pipeline.scheduler = DDPMScheduler.from_config(pipeline.scheduler.config)
+
+denoising_start = 0.2
+timesteps, num_inference_steps = retrieve_timesteps(
+                        pipeline.scheduler, 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)
+from functools import partial
+
+timesteps = [torch.tensor(t) for t in timesteps.tolist()]
+pipeline.__call__ = partial(
+    pipeline.__call__,
+    num_inference_steps=num_steps_inversion,
+    guidance_scale=0,
+    generator=generator,
+    denoising_start=denoising_start,
+    strength=0,
+)
+
+# timesteps, num_inference_steps = retrieve_timesteps(pipeline.scheduler, num_steps_inversion, device, None)
+# timesteps, num_inference_steps = pipeline.get_timesteps(num_inference_steps=num_inference_steps, device=device, strength=strngth)
+
+
+from utils import get_ddpm_inversion_scheduler, create_xts
+
+
+
+from config import get_config, get_config_name
+import argparse
+
+# parser = argparse.ArgumentParser()
+# parser.add_argument("--images_paths", type=str, default=None)
+# parser.add_argument("--images_folder", type=str, default=None)
+# parser.set_defaults(force_use_cpu=False)
+# parser.add_argument("--force_use_cpu", action="store_true")
+# parser.add_argument("--folder_name", type=str, default='test_measure_time')
+# parser.add_argument("--config_from_file", type=str, default='run_configs/noise_shift_guidance_1_5.yaml')
+# parser.set_defaults(save_intermediate_results=False)
+# parser.add_argument("--save_intermediate_results", action="store_true")
+# parser.add_argument("--batch_size", type=int, default=None)
+# parser.set_defaults(skip_p_to_p=False)
+# parser.add_argument("--skip_p_to_p", action="store_true", default=True)
+# parser.set_defaults(only_p_to_p=False)
+# parser.add_argument("--only_p_to_p", action="store_true")
+# parser.set_defaults(fp16=False)
+# parser.add_argument("--fp16", action="store_true", default=False)
+# parser.add_argument("--prompts_file", type=str, default='dataset_measure_time/dataset.json')
+# parser.add_argument("--images_in_prompts_file", type=str, default=None)
+# parser.add_argument("--seed", type=int, default=2)
+# parser.add_argument("--time_measure_n", type=int, default=1)
+
+# args = parser.parse_args()
+class Object(object):
+    pass
+
+args = Object()
+args.images_paths = None
+args.images_folder = None
+args.force_use_cpu = False
+args.folder_name = 'test_measure_time'
+args.config_from_file = 'run_configs/noise_shift_guidance_1_5.yaml'
+args.save_intermediate_results = False
+args.batch_size = None
+args.skip_p_to_p = True
+args.only_p_to_p = False
+args.fp16 = False
+args.prompts_file = 'dataset_measure_time/dataset.json'
+args.images_in_prompts_file = None
+args.seed = 986
+args.time_measure_n = 1
+
+
+assert (
+    args.batch_size is None or args.save_intermediate_results is False
+), "save_intermediate_results is not implemented for batch_size > 1"
+
+config = get_config(args)
+
+
+
+
+
+# latent = latents[0].expand(3, -1, -1, -1)
+# prompt = [src_prompt, src_prompt, tgt_prompt]
+
+# image = pipeline.__call__(image=latent, prompt=prompt, eta=1).images
+
+# for i, im in enumerate(image):
+#     im.save(f"output_{i}.png")
+
+def run(image_path, src_prompt, tgt_prompt, seed, w1, w2):
+    generator = torch.Generator().manual_seed(seed)
+    x_0_image = Image.open(image_path).convert("RGB").resize((512, 512), Image.LANCZOS)
+    x_0 = encode_image(x_0_image, pipeline)
+    # x_ts = create_xts(pipeline.scheduler, timesteps, x_0, noise_shift_delta=1, generator=generator)
+    x_ts = create_xts(1, None, 0, generator, pipeline.scheduler, timesteps, x_0, no_add_noise=False)
+    x_ts = [xt.to(dtype=torch.float16) for xt in x_ts]
+    latents = [x_ts[0]]
+    x_ts_c_hat = [None]
+    config.ws1 = [w1] * 4
+    config.ws2 = [w2] * 4
+    pipeline.scheduler = get_ddpm_inversion_scheduler(
+                    pipeline.scheduler,
+                    config.step_function,
+                    config,
+                    timesteps,
+                    config.save_timesteps,
+                    latents,
+                    x_ts,
+                    x_ts_c_hat,
+                    args.save_intermediate_results,
+                    pipeline,
+                    x_0,
+                    v1s_images := [],
+                    v2s_images := [],
+                    deltas_images := [],
+                    v1_x0s := [],
+                    v2_x0s := [],
+                    deltas_x0s := [],
+                    "res12",
+                    image_name="im_name",
+                    time_measure_n=args.time_measure_n,
+                )
+    latent = latents[0].expand(3, -1, -1, -1)
+    prompt = [src_prompt, src_prompt, tgt_prompt]
+    image = pipeline.__call__(image=latent, prompt=prompt, eta=1).images
+    return image[2]
+
+if __name__ == "__main__":
+    res = run(image_path, src_prompt, tgt_prompt, args.seed, 1.5, 1.0)
+    res.save("output.png")

resize.py

@@ -0,0 +1,18 @@
+from PIL import Image
+
+def resize_image(input_path, output_path, new_size):
+    # Open the image
+    image = Image.open(input_path)
+
+    # Resize the image
+    resized_image = image.resize(new_size)
+
+    # Save the resized image
+    resized_image.save(output_path)
+
+# Example usage
+input_path = "demo_im/pexels-pixabay-458976.png"
+output_path = "demo_im/pexels-pixabay-458976.less.png"
+new_size = (512, 512)
+
+resize_image(input_path, output_path, new_size)

utils.py

@@ -0,0 +1,1356 @@
+import itertools
+from typing import List, Optional, Union
+import PIL
+import PIL.Image
+import torch
+from diffusers.schedulers.scheduling_ddim import DDIMSchedulerOutput
+from diffusers.utils import make_image_grid
+from PIL import Image, ImageDraw, ImageFont
+import os
+from diffusers.utils import (
+    logging,
+    USE_PEFT_BACKEND,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.loaders import (
+    StableDiffusionXLLoraLoaderMixin,
+)
+from diffusers.image_processor import VaeImageProcessor
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+from diffusers.models.lora import adjust_lora_scale_text_encoder
+from diffusers import DiffusionPipeline
+
+
+VECTOR_DATA_FOLDER = "vector_data"
+VECTOR_DATA_DICT = "vector_data"
+
+
+def encode_image(image: PIL.Image, pipe: DiffusionPipeline):
+    pipe.image_processor: VaeImageProcessor = pipe.image_processor  # type: ignore
+    image = pipe.image_processor.pil_to_numpy(image)
+    image = pipe.image_processor.numpy_to_pt(image)
+    image = image.to(pipe.device)
+    return (
+        pipe.vae.encode(
+            pipe.image_processor.preprocess(image),
+        ).latent_dist.mode()
+        * pipe.vae.config.scaling_factor
+    )
+
+
+def decode_latents(latent, pipe):
+    latent_img = pipe.vae.decode(
+        latent / pipe.vae.config.scaling_factor, return_dict=False
+    )[0]
+    return pipe.image_processor.postprocess(latent_img, output_type="pil")
+
+
+def get_device(argv, args=None):
+    import sys
+
+    def debugger_is_active():
+        return hasattr(sys, "gettrace") and sys.gettrace() is not None
+
+    if args:
+        return (
+            torch.device("cuda")
+            if (torch.cuda.is_available() and not debugger_is_active())
+            and not args.force_use_cpu
+            else torch.device("cpu")
+        )
+
+    return (
+        torch.device("cuda")
+        if (torch.cuda.is_available() and not debugger_is_active())
+        and not "cpu" in set(argv[1:])
+        else torch.device("cpu")
+    )
+
+
+def deterministic_ddim_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,
+    return_dict: bool = True,
+    scheduler=None,
+):
+
+    if scheduler.num_inference_steps is None:
+        raise ValueError(
+            "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+        )
+
+    prev_timestep = (
+        timestep - scheduler.config.num_train_timesteps // scheduler.num_inference_steps
+    )
+
+    # 2. compute alphas, betas
+    alpha_prod_t = scheduler.alphas_cumprod[timestep]
+    alpha_prod_t_prev = (
+        scheduler.alphas_cumprod[prev_timestep]
+        if prev_timestep >= 0
+        else scheduler.final_alpha_cumprod
+    )
+
+    beta_prod_t = 1 - alpha_prod_t
+
+    if scheduler.config.prediction_type == "epsilon":
+        pred_original_sample = (
+            sample - beta_prod_t ** (0.5) * model_output
+        ) / alpha_prod_t ** (0.5)
+        pred_epsilon = model_output
+    elif scheduler.config.prediction_type == "sample":
+        pred_original_sample = model_output
+        pred_epsilon = (
+            sample - alpha_prod_t ** (0.5) * pred_original_sample
+        ) / beta_prod_t ** (0.5)
+    elif scheduler.config.prediction_type == "v_prediction":
+        pred_original_sample = (alpha_prod_t**0.5) * sample - (
+            beta_prod_t**0.5
+        ) * model_output
+        pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
+    else:
+        raise ValueError(
+            f"prediction_type given as {scheduler.config.prediction_type} must be one of `epsilon`, `sample`, or"
+            " `v_prediction`"
+        )
+
+    # 4. 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,
+        )
+
+    # 5. compute variance: "sigma_t(η)" -> see formula (16)
+    # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
+    variance = scheduler._get_variance(timestep, prev_timestep)
+    std_dev_t = eta * variance ** (0.5)
+
+    if use_clipped_model_output:
+        # the pred_epsilon is always re-derived from the clipped x_0 in Glide
+        pred_epsilon = (
+            sample - alpha_prod_t ** (0.5) * pred_original_sample
+        ) / beta_prod_t ** (0.5)
+
+    # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
+    pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (
+        0.5
+    ) * pred_epsilon
+
+    # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
+    prev_sample = (
+        alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
+    )
+    return prev_sample
+
+
+def deterministic_euler_step(
+    model_output: torch.FloatTensor,
+    timestep: Union[float, torch.FloatTensor],
+    sample: torch.FloatTensor,
+    eta,
+    use_clipped_model_output,
+    generator,
+    variance_noise,
+    return_dict,
+    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).
+
+    Args:
+        model_output (`torch.FloatTensor`):
+            The direct output from learned diffusion model.
+        timestep (`float`):
+            The current discrete timestep in the diffusion chain.
+        sample (`torch.FloatTensor`):
+            A current instance of a sample created by the diffusion process.
+        generator (`torch.Generator`, *optional*):
+            A random number generator.
+        return_dict (`bool`):
+            Whether or not to return a
+            [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] or tuple.
+
+    Returns:
+        [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] or `tuple`:
+            If return_dict is `True`,
+            [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] is returned,
+            otherwise a tuple is returned where the first element is the sample tensor.
+
+    """
+
+    if (
+        isinstance(timestep, int)
+        or isinstance(timestep, torch.IntTensor)
+        or isinstance(timestep, torch.LongTensor)
+    ):
+        raise ValueError(
+            (
+                "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                " one of the `scheduler.timesteps` as a timestep."
+            ),
+        )
+
+    if scheduler.step_index is None:
+        scheduler._init_step_index(timestep)
+
+    sigma = scheduler.sigmas[scheduler.step_index]
+
+    # Upcast to avoid precision issues when computing prev_sample
+    sample = sample.to(torch.float32)
+
+    # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
+    if scheduler.config.prediction_type == "epsilon":
+        pred_original_sample = sample - sigma * model_output
+    elif scheduler.config.prediction_type == "v_prediction":
+        # * c_out + input * c_skip
+        pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (
+            sample / (sigma**2 + 1)
+        )
+    elif scheduler.config.prediction_type == "sample":
+        raise NotImplementedError("prediction_type not implemented yet: sample")
+    else:
+        raise ValueError(
+            f"prediction_type given as {scheduler.config.prediction_type} must be one of `epsilon`, or `v_prediction`"
+        )
+
+    sigma_from = scheduler.sigmas[scheduler.step_index]
+    sigma_to = scheduler.sigmas[scheduler.step_index + 1]
+    sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5
+    sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
+
+    # 2. Convert to an ODE derivative
+    derivative = (sample - pred_original_sample) / sigma
+
+    dt = sigma_down - sigma
+
+    prev_sample = sample + derivative * dt
+
+    # Cast sample back to model compatible dtype
+    prev_sample = prev_sample.to(model_output.dtype)
+
+    # upon completion increase step index by one
+    scheduler._step_index += 1
+
+    return prev_sample
+
+
+def deterministic_non_ancestral_euler_step(
+    model_output: torch.FloatTensor,
+    timestep: Union[float, torch.FloatTensor],
+    sample: torch.FloatTensor,
+    eta: float = 0.0,
+    use_clipped_model_output: bool = False,
+    s_churn: float = 0.0,
+    s_tmin: float = 0.0,
+    s_tmax: float = float("inf"),
+    s_noise: float = 1.0,
+    generator: Optional[torch.Generator] = None,
+    variance_noise: Optional[torch.FloatTensor] = None,
+    return_dict: bool = True,
+    scheduler=None,
+):
+    """
+    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).
+
+    Args:
+        model_output (`torch.FloatTensor`):
+            The direct output from learned diffusion model.
+        timestep (`float`):
+            The current discrete timestep in the diffusion chain.
+        sample (`torch.FloatTensor`):
+            A current instance of a sample created by the diffusion process.
+        s_churn (`float`):
+        s_tmin  (`float`):
+        s_tmax  (`float`):
+        s_noise (`float`, defaults to 1.0):
+            Scaling factor for noise added to the sample.
+        generator (`torch.Generator`, *optional*):
+            A random number generator.
+        return_dict (`bool`):
+            Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
+            tuple.
+
+    Returns:
+        [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
+            If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
+            returned, otherwise a tuple is returned where the first element is the sample tensor.
+    """
+
+    if (
+        isinstance(timestep, int)
+        or isinstance(timestep, torch.IntTensor)
+        or isinstance(timestep, torch.LongTensor)
+    ):
+        raise ValueError(
+            (
+                "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                " one of the `scheduler.timesteps` as a timestep."
+            ),
+        )
+
+    if not scheduler.is_scale_input_called:
+        logger.warning(
+            "The `scale_model_input` function should be called before `step` to ensure correct denoising. "
+            "See `StableDiffusionPipeline` for a usage example."
+        )
+
+    if scheduler.step_index is None:
+        scheduler._init_step_index(timestep)
+
+    # Upcast to avoid precision issues when computing prev_sample
+    sample = sample.to(torch.float32)
+
+    sigma = scheduler.sigmas[scheduler.step_index]
+
+    gamma = (
+        min(s_churn / (len(scheduler.sigmas) - 1), 2**0.5 - 1)
+        if s_tmin <= sigma <= s_tmax
+        else 0.0
+    )
+
+    sigma_hat = sigma * (gamma + 1)
+
+    # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
+    # NOTE: "original_sample" should not be an expected prediction_type but is left in for
+    # backwards compatibility
+    if (
+        scheduler.config.prediction_type == "original_sample"
+        or scheduler.config.prediction_type == "sample"
+    ):
+        pred_original_sample = model_output
+    elif scheduler.config.prediction_type == "epsilon":
+        pred_original_sample = sample - sigma_hat * model_output
+    elif scheduler.config.prediction_type == "v_prediction":
+        # denoised = model_output * c_out + input * c_skip
+        pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (
+            sample / (sigma**2 + 1)
+        )
+    else:
+        raise ValueError(
+            f"prediction_type given as {scheduler.config.prediction_type} must be one of `epsilon`, or `v_prediction`"
+        )
+
+    # 2. Convert to an ODE derivative
+    derivative = (sample - pred_original_sample) / sigma_hat
+
+    dt = scheduler.sigmas[scheduler.step_index + 1] - sigma_hat
+
+    prev_sample = sample + derivative * dt
+
+    # Cast sample back to model compatible dtype
+    prev_sample = prev_sample.to(model_output.dtype)
+
+    # upon completion increase step index by one
+    scheduler._step_index += 1
+
+    return prev_sample
+
+
+def deterministic_ddpm_step(
+    model_output: torch.FloatTensor,
+    timestep: Union[float, torch.FloatTensor],
+    sample: torch.FloatTensor,
+    eta,
+    use_clipped_model_output,
+    generator,
+    variance_noise,
+    return_dict,
+    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).
+
+    Args:
+        model_output (`torch.FloatTensor`):
+            The direct output from learned diffusion model.
+        timestep (`float`):
+            The current discrete timestep in the diffusion chain.
+        sample (`torch.FloatTensor`):
+            A current instance of a sample created by the diffusion process.
+        generator (`torch.Generator`, *optional*):
+            A random number generator.
+        return_dict (`bool`, *optional*, defaults to `True`):
+            Whether or not to return a [`~schedulers.scheduling_ddpm.DDPMSchedulerOutput`] or `tuple`.
+
+    Returns:
+        [`~schedulers.scheduling_ddpm.DDPMSchedulerOutput`] or `tuple`:
+            If return_dict is `True`, [`~schedulers.scheduling_ddpm.DDPMSchedulerOutput`] is returned, otherwise a
+            tuple is returned where the first element is the sample tensor.
+
+    """
+    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
+        )
+
+    # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
+    # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
+    pred_original_sample_coeff = (
+        alpha_prod_t_prev ** (0.5) * current_beta_t
+    ) / beta_prod_t
+    current_sample_coeff = current_alpha_t ** (0.5) * beta_prod_t_prev / beta_prod_t
+
+    # 5. Compute predicted previous sample µ_t
+    # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
+    pred_prev_sample = (
+        pred_original_sample_coeff * pred_original_sample
+        + current_sample_coeff * sample
+    )
+
+    return pred_prev_sample
+
+
+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 find_index(timesteps, timestep):
+    for i, t in enumerate(timesteps):
+        if t == timestep:
+            return i
+    return -1
+
+map_timpstep_to_index = {
+    torch.tensor(799): 0,
+    torch.tensor(599): 1,
+    torch.tensor(399): 2,
+    torch.tensor(199): 3,
+    torch.tensor(799, device='cuda:0'): 0,
+    torch.tensor(599, device='cuda:0'): 1,
+    torch.tensor(399, device='cuda:0'): 2,
+    torch.tensor(199, device='cuda:0'): 3,
+}
+
+def step_save_latents(
+    self,
+    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,
+    return_dict: bool = True,
+):
+    # print(self._save_timesteps)
+    # timestep_index = map_timpstep_to_index[timestep]
+    # timestep_index = ((self._save_timesteps == timestep).nonzero(as_tuple=True)[0]).item()
+    timestep_index = self._save_timesteps.index(timestep) if not self.clean_step_run else -1
+    next_timestep_index = timestep_index + 1 if not self.clean_step_run else -1
+    u_hat_t = self.step_function(
+        model_output=model_output,
+        timestep=timestep,
+        sample=sample,
+        eta=eta,
+        use_clipped_model_output=use_clipped_model_output,
+        generator=generator,
+        variance_noise=variance_noise,
+        return_dict=False,
+        scheduler=self,
+    )
+
+    x_t_minus_1 = self.x_ts[next_timestep_index]
+    self.x_ts_c_hat.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,
+    eta: float = 0.0,
+    use_clipped_model_output: bool = False,
+    generator=None,
+    variance_noise: Optional[torch.FloatTensor] = None,
+    return_dict: bool = True,
+):
+    # timestep_index = ((self._save_timesteps == timestep).nonzero(as_tuple=True)[0]).item()
+    timestep_index = self._timesteps.index(timestep) if not self.clean_step_run else -1
+    next_timestep_index = (
+        timestep_index + 1 if not self.clean_step_run else -1
+    )
+    z_t = self.latents[next_timestep_index]  # + 1 because latents[0] is X_T
+
+    _, normalize_coefficient = normalize(
+        z_t[0] if self._config.breakdown == "x_t_hat_c_with_zeros" else z_t,
+        timestep_index,
+        self._config.max_norm_zs,
+    )
+
+    if normalize_coefficient == 0:
+        eta = 0
+
+    # eta = normalize_coefficient
+
+    x_t_hat_c_hat = self.step_function(
+        model_output=model_output,
+        timestep=timestep,
+        sample=sample,
+        eta=eta,
+        use_clipped_model_output=use_clipped_model_output,
+        generator=generator,
+        variance_noise=variance_noise,
+        return_dict=False,
+        scheduler=self,
+    )
+
+    w1 = self._config.ws1[timestep_index]
+    w2 = self._config.ws2[timestep_index]
+
+    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_hat: torch.Tensor = self.x_ts_c_hat[next_timestep_index]
+    if self._config.breakdown == "x_t_c_hat":
+        raise NotImplementedError("breakdown x_t_c_hat not implemented yet")
+
+    # x_t_c_hat = x_t_c_hat.expand_as(x_t_hat_c_hat)
+    x_t_c = x_t_c_hat[0].expand_as(x_t_hat_c_hat)
+
+    # if self._config.breakdown == "x_t_c_hat":
+    #     v1 = x_t_hat_c_hat - x_t_c_hat
+    #     v2 = x_t_c_hat - x_t_c
+    if (
+        self._config.breakdown == "x_t_hat_c"
+        or self._config.breakdown == "x_t_hat_c_with_zeros"
+    ):
+        zero_index_reconstruction = 1 if not self.time_measure_n else 0
+        edit_prompts_num = (
+            (model_output.size(0) - zero_index_reconstruction) // 3
+            if self._config.breakdown == "x_t_hat_c_with_zeros" and not self.p_to_p
+            else (model_output.size(0) - zero_index_reconstruction) // 2
+        )
+        x_t_hat_c_indices = (zero_index_reconstruction, edit_prompts_num + zero_index_reconstruction)
+        edit_images_indices = (
+            edit_prompts_num + zero_index_reconstruction,
+            (
+                model_output.size(0)
+                if self._config.breakdown == "x_t_hat_c"
+                else zero_index_reconstruction + 2 * edit_prompts_num
+            ),
+        )
+        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]
+        ]
+        v1 = x_t_hat_c_hat - x_t_hat_c
+        v2 = x_t_hat_c - normalize_coefficient * x_t_c
+        if self._config.breakdown == "x_t_hat_c_with_zeros" and not self.p_to_p:
+            path = os.path.join(
+                self.folder_name,
+                VECTOR_DATA_FOLDER,
+                self.image_name,
+            )
+            if not hasattr(self, VECTOR_DATA_DICT):
+                os.makedirs(path, exist_ok=True)
+                self.vector_data = dict()
+
+            x_t_0 = x_t_c_hat[1]
+            empty_prompt_indices = (1 + 2 * edit_prompts_num, 1 + 3 * edit_prompts_num)
+            x_t_hat_0 = x_t_hat_c_hat[empty_prompt_indices[0] : empty_prompt_indices[1]]
+
+            self.vector_data[timestep.item()] = dict()
+            self.vector_data[timestep.item()]["x_t_hat_c"] = x_t_hat_c[
+                edit_images_indices[0] : edit_images_indices[1]
+            ]
+            self.vector_data[timestep.item()]["x_t_hat_0"] = x_t_hat_0
+            self.vector_data[timestep.item()]["x_t_c"] = x_t_c[0].expand_as(x_t_hat_0)
+            self.vector_data[timestep.item()]["x_t_0"] = x_t_0.expand_as(x_t_hat_0)
+            self.vector_data[timestep.item()]["x_t_hat_c_hat"] = x_t_hat_c_hat[
+                edit_images_indices[0] : edit_images_indices[1]
+            ]
+            self.vector_data[timestep.item()]["x_t_minus_1_noisy"] = x_t_minus_1_exact[
+                0
+            ].expand_as(x_t_hat_0)
+            self.vector_data[timestep.item()]["x_t_minus_1_clean"] = self.x_0s[
+                next_timestep_index
+            ].expand_as(x_t_hat_0)
+
+    else:  # no breakdown
+        v1 = x_t_hat_c_hat - normalize_coefficient * x_t_c
+        v2 = 0
+
+    if self.save_intermediate_results and not self.p_to_p:
+        delta = v1 + v2
+        v1_plus_x0 = self.x_0s[next_timestep_index] + v1
+        v2_plus_x0 = self.x_0s[next_timestep_index] + v2
+        delta_plus_x0 = self.x_0s[next_timestep_index] + delta
+
+        v1_images = decode_latents(v1, self.pipe)
+        self.v1s_images.append(v1_images)
+        v2_images = (
+            decode_latents(v2, self.pipe)
+            if self._config.breakdown != "no_breakdown"
+            else [PIL.Image.new("RGB", (1, 1))]
+        )
+        self.v2s_images.append(v2_images)
+        delta_images = decode_latents(delta, self.pipe)
+        self.deltas_images.append(delta_images)
+        v1_plus_x0_images = decode_latents(v1_plus_x0, self.pipe)
+        self.v1_x0s.append(v1_plus_x0_images)
+        v2_plus_x0_images = (
+            decode_latents(v2_plus_x0, self.pipe)
+            if self._config.breakdown != "no_breakdown"
+            else [PIL.Image.new("RGB", (1, 1))]
+        )
+        self.v2_x0s.append(v2_plus_x0_images)
+        delta_plus_x0_images = decode_latents(delta_plus_x0, self.pipe)
+        self.deltas_x0s.append(delta_plus_x0_images)
+
+    # print(f"v1 norm: {torch.norm(v1, dim=0).mean()}")
+    # if self._config.breakdown != "no_breakdown":
+    #     print(f"v2 norm: {torch.norm(v2, dim=0).mean()}")
+    #     print(f"v sum norm: {torch.norm(v1 + v2, dim=0).mean()}")
+
+    x_t_minus_1 = normalize_coefficient * x_t_minus_1_exact + w1 * v1 + w2 * v2
+
+    if (
+        self._config.breakdown == "x_t_hat_c"
+        or self._config.breakdown == "x_t_hat_c_with_zeros"
+    ):
+        x_t_minus_1[x_t_hat_c_indices[0] : x_t_hat_c_indices[1]] = x_t_minus_1[
+            edit_images_indices[0] : edit_images_indices[1]
+        ] # update x_t_hat_c to be x_t_hat_c_hat
+        if self._config.breakdown == "x_t_hat_c_with_zeros" and not self.p_to_p:
+            x_t_minus_1[empty_prompt_indices[0] : empty_prompt_indices[1]] = (
+                x_t_minus_1[edit_images_indices[0] : edit_images_indices[1]]
+            )
+            self.vector_data[timestep.item()]["x_t_minus_1_edited"] = x_t_minus_1[
+                edit_images_indices[0] : edit_images_indices[1]
+            ]
+            if timestep == self._timesteps[-1]:
+                torch.save(
+                    self.vector_data,
+                    os.path.join(
+                        path,
+                        f"{VECTOR_DATA_DICT}.pt",
+                    ),
+                )
+    # p_to_p_force_perfect_reconstruction
+    if not self.time_measure_n:
+        x_t_minus_1[0] = x_t_minus_1_exact[0]
+
+    if not return_dict:
+        return (x_t_minus_1,)
+
+    return DDIMSchedulerOutput(
+        prev_sample=x_t_minus_1,
+        pred_original_sample=None,
+    )
+
+
+
+def get_ddpm_inversion_scheduler(
+    scheduler,
+    step_function,
+    config,
+    timesteps,
+    save_timesteps,
+    latents,
+    x_ts,
+    x_ts_c_hat,
+    save_intermediate_results,
+    pipe,
+    x_0,
+    v1s_images,
+    v2s_images,
+    deltas_images,
+    v1_x0s,
+    v2_x0s,
+    deltas_x0s,
+    folder_name,
+    image_name,
+    time_measure_n,
+):
+    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,
+        return_dict: bool = True,
+    ):
+        # if scheduler.is_save:
+        # start = timer()
+        res_inv =  step_save_latents(
+            scheduler,
+            model_output[:1, :, :, :],
+            timestep,
+            sample[:1, :, :, :],
+            eta,
+            use_clipped_model_output,
+            generator,
+            variance_noise,
+            return_dict,
+        )
+        # end = timer()
+        # print(f"Run Time Inv: {end - start}")
+
+        res_inf = step_use_latents(
+            scheduler,
+            model_output[1:, :, :, :],
+            timestep,
+            sample[1:, :, :, :],
+            eta,
+            use_clipped_model_output,
+            generator,
+            variance_noise,
+            return_dict,
+        )
+        # res = res_inv
+        res = (torch.cat((res_inv[0], res_inf[0]), dim=0),)
+        return res
+        # return res
+
+    scheduler.step_function = step_function
+    scheduler.is_save = True
+    scheduler._timesteps = timesteps
+    scheduler._save_timesteps = save_timesteps if save_timesteps else timesteps
+    scheduler._config = config
+    scheduler.latents = latents
+    scheduler.x_ts = x_ts
+    scheduler.x_ts_c_hat = x_ts_c_hat
+    scheduler.step = step
+    scheduler.save_intermediate_results = save_intermediate_results
+    scheduler.pipe = pipe
+    scheduler.v1s_images = v1s_images
+    scheduler.v2s_images = v2s_images
+    scheduler.deltas_images = deltas_images
+    scheduler.v1_x0s = v1_x0s
+    scheduler.v2_x0s = v2_x0s
+    scheduler.deltas_x0s = deltas_x0s
+    scheduler.clean_step_run = False
+    scheduler.x_0s = create_xts(
+        config.noise_shift_delta,
+        config.noise_timesteps,
+        config.clean_step_timestep,
+        None,
+        pipe.scheduler,
+        timesteps,
+        x_0,
+        no_add_noise=True,
+    )
+    scheduler.folder_name = folder_name
+    scheduler.image_name = image_name
+    scheduler.p_to_p = False
+    scheduler.p_to_p_replace = False
+    scheduler.time_measure_n = time_measure_n
+    return scheduler
+
+
+def create_grid(
+    images,
+    p_to_p_images,
+    prompts,
+    original_image_path,
+):
+    images_len = len(images) if len(images) > 0 else len(p_to_p_images)
+    images_size = images[0].size if len(images) > 0 else p_to_p_images[0].size
+    x_0 = Image.open(original_image_path).resize(images_size)
+
+    images_ = [x_0] + images + ([x_0] + p_to_p_images if p_to_p_images else [])
+
+    l1 = 1 if len(images) > 0 else 0
+    l2 = 1 if len(p_to_p_images) else 0
+    grid = make_image_grid(images_, rows=l1 + l2, cols=images_len + 1, resize=None)
+
+    width = images_size[0]
+    height = width // 5
+    font = ImageFont.truetype("font.ttf", width // 14)
+
+    grid1 = Image.new("RGB", size=(grid.size[0], grid.size[1] + height))
+    grid1.paste(grid, (0, 0))
+
+    draw = ImageDraw.Draw(grid1)
+
+    c_width = 0
+    for prompt in prompts:
+        if len(prompt) > 30:
+            prompt = prompt[:30] + "\n" + prompt[30:]
+        draw.text((c_width, width * 2), prompt, font=font, fill=(255, 255, 255))
+        c_width += width
+
+    return grid1
+
+
+def save_intermediate_results(
+    v1s_images,
+    v2s_images,
+    deltas_images,
+    v1_x0s,
+    v2_x0s,
+    deltas_x0s,
+    folder_name,
+    original_prompt,
+):
+    from diffusers.utils import make_image_grid
+
+    path = f"{folder_name}/{original_prompt}_intermediate_results/"
+    os.makedirs(path, exist_ok=True)
+    make_image_grid(
+        list(itertools.chain(*v1s_images)),
+        rows=len(v1s_images),
+        cols=len(v1s_images[0]),
+    ).save(f"{path}v1s_images.png")
+    make_image_grid(
+        list(itertools.chain(*v2s_images)),
+        rows=len(v2s_images),
+        cols=len(v2s_images[0]),
+    ).save(f"{path}v2s_images.png")
+    make_image_grid(
+        list(itertools.chain(*deltas_images)),
+        rows=len(deltas_images),
+        cols=len(deltas_images[0]),
+    ).save(f"{path}deltas_images.png")
+    make_image_grid(
+        list(itertools.chain(*v1_x0s)),
+        rows=len(v1_x0s),
+        cols=len(v1_x0s[0]),
+    ).save(f"{path}v1_x0s.png")
+    make_image_grid(
+        list(itertools.chain(*v2_x0s)),
+        rows=len(v2_x0s),
+        cols=len(v2_x0s[0]),
+    ).save(f"{path}v2_x0s.png")
+    make_image_grid(
+        list(itertools.chain(*deltas_x0s)),
+        rows=len(deltas_x0s[0]),
+        cols=len(deltas_x0s),
+    ).save(f"{path}deltas_x0s.png")
+    for i, image in enumerate(list(itertools.chain(*deltas_x0s))):
+        image.save(f"{path}deltas_x0s_{i}.png")
+
+
+# copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.py and removed the add_noise line
+def prepare_latents_no_add_noise(
+    self,
+    image,
+    timestep,
+    batch_size,
+    num_images_per_prompt,
+    dtype,
+    device,
+    generator=None,
+):
+    from diffusers.utils import deprecate
+
+    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)}"
+        )
+
+    image = image.to(device=device, dtype=dtype)
+
+    batch_size = batch_size * num_images_per_prompt
+
+    if image.shape[1] == 4:
+        init_latents = image
+
+    else:
+        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 = [
+                self.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 = self.retrieve_latents(
+                self.vae.encode(image), generator=generator
+            )
+
+        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
+        deprecation_message = (
+            f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial"
+            " images (`image`). Initial images are now duplicating to match the number of text prompts. Note"
+            " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update"
+            " your script to pass as many initial images as text prompts to suppress this warning."
+        )
+        deprecate(
+            "len(prompt) != len(image)",
+            "1.0.0",
+            deprecation_message,
+            standard_warn=False,
+        )
+        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)
+
+    # get latents
+    latents = init_latents
+
+    return latents
+
+
+# Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.encode_prompt
+def encode_prompt_empty_prompt_zeros_sdxl(
+    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: procecss multi-vector tokens if necessary
+        prompt_embeds_list = []
+        prompts = [prompt, prompt_2]
+        for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
+
+            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)]
+
+            if self.config.force_zeros_for_empty_prompt:
+                prompt_embeds[[i for i in range(len(prompt)) if prompt[i] == ""]] = 0
+                pooled_prompt_embeds[
+                    [i for i in range(len(prompt)) if prompt[i] == ""]
+                ] = 0
+
+            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
+        ):
+
+            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,
+    )
+
+
+def create_xts(
+    noise_shift_delta,
+    noise_timesteps,
+    clean_step_timestep,
+    generator,
+    scheduler,
+    timesteps,
+    x_0,
+    no_add_noise=False,
+):
+    if noise_timesteps is None:
+        noising_delta = noise_shift_delta * (timesteps[0] - timesteps[1])
+        noise_timesteps = [timestep - int(noising_delta) for timestep in 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="cpu").to(
+            x_0.device
+        )
+        if not no_add_noise
+        else torch.zeros_like(x_0_expanded)
+    )
+    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]
+    if clean_step_timestep > 0:
+        x_ts += [x_0]
+    return x_ts
+
+
+# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
+def add_noise(
+    self,
+    original_samples: torch.FloatTensor,
+    noise: torch.FloatTensor,
+    image_timesteps: torch.IntTensor,
+    noise_timesteps: torch.IntTensor,
+) -> torch.FloatTensor:
+    # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
+    # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement
+    # for the subsequent add_noise calls
+    self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device)
+    alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype)
+    timesteps = timesteps.to(original_samples.device)
+
+    sqrt_alpha_prod = alphas_cumprod[image_timesteps] ** 0.5
+    sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+    while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
+        sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+    sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[noise_timesteps]) ** 0.5
+    sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+    while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+    noisy_samples = (
+        sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+    )
+    return noisy_samples
+
+
+def make_image_grid(
+    images: List[PIL.Image.Image], rows: int, cols: int, resize: int = None, size=None
+) -> PIL.Image.Image:
+    """
+    Prepares a single grid of images. Useful for visualization purposes.
+    """
+    assert len(images) == rows * cols
+
+    if resize is not None:
+        images = [img.resize((resize, resize)) for img in images]
+
+    w, h = size
+    grid = Image.new("RGB", size=(cols * w, rows * h))
+
+    for i, img in enumerate(images):
+        grid.paste(img, box=(i % cols * w, i // cols * h))
+    return grid