add app_haircolor

.gitignore

@@ -0,0 +1,3 @@
+.vscode
+.DS_Store
+__pycache__

README.md

@@ -10,4 +10,5 @@ pinned: false
 license: mit
 ---
 
+Modified from: https://huggingface.co/spaces/turboedit/turbo_edit
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,13 @@
+import gradio as gr
+
+# from app_base import create_demo as create_demo_face
+from app_haircolor import create_demo as create_demo_haircolor
+
+with gr.Blocks(css="style.css") as demo:
+    with gr.Tabs():
+        # with gr.Tab(label="Face"):
+        #     create_demo_face()
+        with gr.Tab(label="Hair Color"):
+            create_demo_haircolor()
+
+demo.launch()

app_haircolor.py

@@ -0,0 +1,127 @@
+import spaces
+import gradio as gr
+import time
+import torch
+
+from PIL import Image
+from segment_utils import(
+    segment_image,
+    restore_result,
+)
+from enhance_utils import enhance_image
+from inversion_run_adapter import run as adapter_run
+
+
+DEFAULT_SRC_PROMPT = "a woman, with hair"
+DEFAULT_EDIT_PROMPT = "a woman, with baby blue hair, 8k, high quality"
+
+DEFAULT_CATEGORY = "hair"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+@spaces.GPU(duration=15)
+def image_to_image(
+    input_image: Image,
+    input_image_prompt: str,
+    edit_prompt: str,
+    seed: int,
+    w1: float,
+    num_steps: int,
+    start_step: int,
+    guidance_scale: float,
+    generate_size: int,
+    lineart_scale: float,
+    canny_scale: float,
+    lineart_detect: float,
+    canny_detect: float,
+):
+    w2 = 1.0
+    run_task_time = 0
+    time_cost_str = ''
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+    run_model = adapter_run
+    generated_image = run_model(
+        input_image,
+        input_image_prompt,
+        edit_prompt,
+        generate_size,
+        seed,
+        w1,
+        w2,
+        num_steps,
+        start_step,
+        guidance_scale,
+        lineart_scale,
+        canny_scale,
+        lineart_detect,
+        canny_detect,
+    )
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+    enhanced_image = enhance_image(generated_image, False)
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+
+    return enhanced_image, generated_image, time_cost_str
+
+def get_time_cost(run_task_time, time_cost_str):
+    now_time = int(time.time()*1000)
+    if run_task_time == 0:
+        time_cost_str = 'start'
+    else:
+        if time_cost_str != '': 
+            time_cost_str += f'-->'
+        time_cost_str += f'{now_time - run_task_time}'
+    run_task_time = now_time
+    return run_task_time, time_cost_str
+
+def create_demo() -> gr.Blocks:
+    with gr.Blocks() as demo:
+        croper = gr.State()
+        with gr.Row():
+            with gr.Column():
+                input_image_prompt = gr.Textbox(lines=1, label="Input Image Prompt", value=DEFAULT_SRC_PROMPT)
+                edit_prompt = gr.Textbox(lines=1, label="Edit Prompt", value=DEFAULT_EDIT_PROMPT)
+                category = gr.Textbox(label="Category", value=DEFAULT_CATEGORY, visible=False)
+            with gr.Column():
+                num_steps = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Num Steps")
+                start_step = gr.Slider(minimum=1, maximum=100, value=2, step=1, label="Start Step")
+                with gr.Accordion("Advanced Options", open=False):
+                    guidance_scale = gr.Slider(minimum=0, maximum=20, value=0, step=0.5, label="Guidance Scale", visible=True)
+                    generate_size = gr.Number(label="Generate Size", value=768)
+                    mask_expansion = gr.Number(label="Mask Expansion", value=10, visible=True)
+                    mask_dilation = gr.Slider(minimum=0, maximum=10, value=2, step=1, label="Mask Dilation")
+                    lineart_scale = gr.Slider(minimum=0, maximum=5, value=1.6, step=0.1, label="Lineart Weights", visible=True)
+                    canny_scale = gr.Slider(minimum=0, maximum=5, value=0.8, step=0.1, label="Canny Weights", visible=True)
+                    lineart_detect = gr.Number(label="Lineart Detect", value=0.375, visible=True)
+                    canny_detect = gr.Number(label="Canny Detect", value=0.375, visible=True)
+            with gr.Column():
+                seed = gr.Number(label="Seed", value=8)
+                w1 = gr.Number(label="W1", value=2.5)
+                g_btn = gr.Button("Edit Image")
+                
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Input Image", type="pil")
+            with gr.Column():
+                restored_image = gr.Image(label="Restored Image", type="pil", interactive=False)
+                download_path = gr.File(label="Download the output image", interactive=False)
+            with gr.Column():
+                origin_area_image = gr.Image(label="Origin Area Image", type="pil", interactive=False)
+                enhanced_image = gr.Image(label="Enhanced Image", type="pil", interactive=False)
+                generated_cost = gr.Textbox(label="Time cost by step (ms):", visible=True, interactive=False)
+                generated_image = gr.Image(label="Generated Image", type="pil", interactive=False)
+        
+        g_btn.click(
+            fn=segment_image,
+            inputs=[input_image, category, generate_size, mask_expansion, mask_dilation],
+            outputs=[origin_area_image, croper],
+        ).success(
+            fn=image_to_image,
+            inputs=[origin_area_image, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step, guidance_scale, generate_size, lineart_scale, canny_scale, lineart_detect, canny_detect],
+            outputs=[enhanced_image, generated_image, generated_cost],
+        ).success(
+            fn=restore_result,
+            inputs=[croper, category, enhanced_image],
+            outputs=[restored_image, download_path],
+        )
+
+    return demo

checkpoints/selfie_multiclass_256x256.tflite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6748b1253a99067ef71f7e26ca71096cd449baefa8f101900ea23016507e0e0
+size 16371837

config.py

@@ -0,0 +1,141 @@
+from ml_collections import config_dict
+import yaml
+from diffusers.schedulers import (
+    DDIMScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+    DDPMScheduler,
+)
+from inversion_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"
+
+    cfg.max_norm_zs = [-1] * (num_steps_actual - 1) + [15.5]
+    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}\
+"""

croper.py

@@ -0,0 +1,108 @@
+import PIL
+import numpy as np
+
+from PIL import Image
+
+class Croper:
+    def __init__(
+        self,
+        input_image: PIL.Image,
+        target_mask: np.ndarray,
+        mask_size: int = 256,
+        mask_expansion: int = 20,
+    ):
+        self.input_image = input_image
+        self.target_mask = target_mask
+        self.mask_size = mask_size
+        self.mask_expansion = mask_expansion
+    
+    def corp_mask_image(self):
+        target_mask = self.target_mask
+        input_image = self.input_image
+        mask_expansion = self.mask_expansion
+        original_width, original_height = input_image.size
+        mask_indices = np.where(target_mask)
+        start_y = np.min(mask_indices[0])
+        end_y = np.max(mask_indices[0])
+        start_x = np.min(mask_indices[1])
+        end_x = np.max(mask_indices[1])
+        mask_height = end_y - start_y
+        mask_width = end_x - start_x
+        # choose the max side length
+        max_side_length = max(mask_height, mask_width)
+        # expand the mask area
+        height_diff = (max_side_length - mask_height) // 2
+        width_diff = (max_side_length - mask_width) // 2
+        start_y = start_y - mask_expansion - height_diff
+        if start_y < 0:
+            start_y = 0
+        end_y = end_y + mask_expansion + height_diff
+        if end_y > original_height:
+            end_y = original_height
+        start_x = start_x - mask_expansion - width_diff
+        if start_x < 0:
+            start_x = 0
+        end_x = end_x + mask_expansion + width_diff
+        if end_x > original_width:
+            end_x = original_width
+        expanded_height = end_y - start_y
+        expanded_width = end_x - start_x
+        expanded_max_side_length = max(expanded_height, expanded_width)
+        # calculate the crop area
+        crop_mask = target_mask[start_y:end_y, start_x:end_x]
+        crop_mask_start_y = (expanded_max_side_length - expanded_height) // 2
+        crop_mask_end_y = crop_mask_start_y + expanded_height
+        crop_mask_start_x = (expanded_max_side_length - expanded_width) // 2
+        crop_mask_end_x = crop_mask_start_x + expanded_width
+        # create a square mask
+        square_mask = np.zeros((expanded_max_side_length, expanded_max_side_length), dtype=target_mask.dtype)
+        square_mask[crop_mask_start_y:crop_mask_end_y, crop_mask_start_x:crop_mask_end_x] = crop_mask
+        square_mask_image = Image.fromarray((square_mask * 255).astype(np.uint8))
+
+        crop_image = input_image.crop((start_x, start_y, end_x, end_y))
+        square_image = Image.new("RGB", (expanded_max_side_length, expanded_max_side_length))
+        square_image.paste(crop_image, (crop_mask_start_x, crop_mask_start_y))
+
+        self.origin_start_x = start_x
+        self.origin_start_y = start_y
+        self.origin_end_x = end_x
+        self.origin_end_y = end_y
+
+        self.square_start_x = crop_mask_start_x
+        self.square_start_y = crop_mask_start_y
+        self.square_end_x = crop_mask_end_x
+        self.square_end_y = crop_mask_end_y
+
+        self.square_length = expanded_max_side_length
+        self.square_mask_image = square_mask_image
+        self.square_image = square_image
+        self.corp_mask = crop_mask
+
+        mask_size = self.mask_size
+        self.resized_square_mask_image = square_mask_image.resize((mask_size, mask_size))
+        self.resized_square_image = square_image.resize((mask_size, mask_size))
+
+        return self.resized_square_mask_image
+    
+    def restore_result(self, generated_image):
+        square_length = self.square_length
+        generated_image = generated_image.resize((square_length, square_length))
+        square_mask_image = self.square_mask_image
+        cropped_generated_image = generated_image.crop((self.square_start_x, self.square_start_y, self.square_end_x, self.square_end_y))
+        cropped_square_mask_image = square_mask_image.crop((self.square_start_x, self.square_start_y, self.square_end_x, self.square_end_y))
+
+        restored_image = self.input_image.copy()
+        restored_image.paste(cropped_generated_image, (self.origin_start_x, self.origin_start_y), cropped_square_mask_image)
+        
+        return restored_image
+    
+    def restore_result_v2(self, generated_image):
+        square_length = self.square_length
+        generated_image = generated_image.resize((square_length, square_length))
+        cropped_generated_image = generated_image.crop((self.square_start_x, self.square_start_y, self.square_end_x, self.square_end_y))
+
+        restored_image = self.input_image.copy()
+        restored_image.paste(cropped_generated_image, (self.origin_start_x, self.origin_start_y))
+        
+        return restored_image
+        

enhance_utils.py

@@ -0,0 +1,41 @@
+import os
+import torch
+import cv2
+import numpy as np
+
+from PIL import Image
+from gfpgan.utils import GFPGANer
+from basicsr.archs.srvgg_arch import SRVGGNetCompact
+from realesrgan.utils import RealESRGANer
+
+os.system("pip freeze")
+if not os.path.exists('GFPGANv1.4.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
+if not os.path.exists('realesr-general-x4v3.pth'):
+    os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
+
+os.makedirs('output', exist_ok=True)
+
+model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+model_path = 'realesr-general-x4v3.pth'
+half = True if torch.cuda.is_available() else False
+upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
+
+face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=1, arch='clean', channel_multiplier=2)
+
+def enhance_image(
+    pil_image: Image,
+    enhance_face: bool = True,
+):
+    img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
+
+    h, w = img.shape[0:2]
+    if h < 300:
+        img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
+    if enhance_face:
+        _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=True, paste_back=True)
+    else:
+        output, _ = upsampler.enhance(img, outscale=2)
+    pil_output = Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))
+
+    return pil_output

inversion_run_adapter.py

@@ -0,0 +1,282 @@
+import torch
+
+from diffusers import (
+    DDPMScheduler,
+    DiffusionPipeline,
+    T2IAdapter,
+    MultiAdapter,
+)
+from controlnet_aux import (
+    LineartDetector,
+    CannyDetector,
+    MidasDetector,
+    PidiNetDetector,
+)
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img import retrieve_timesteps, retrieve_latents
+from PIL import Image
+from inversion_utils import get_ddpm_inversion_scheduler, create_xts
+from config import get_config, get_num_steps_actual
+from functools import partial
+from compel import Compel, ReturnedEmbeddingsType
+
+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"
+
+generator = None
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
+BASE_MODEL = "stabilityai/sdxl-turbo"
+# BASE_MODEL = "SG161222/RealVisXL_V5.0_Lightning"
+# BASE_MODEL = "Lykon/dreamshaper-xl-v2-turbo"
+# BASE_MODEL = "RunDiffusion/Juggernaut-XL-Lightning"
+
+lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators")
+lineart_detector = lineart_detector.to(device)
+
+pidinet_detector = PidiNetDetector.from_pretrained("lllyasviel/Annotators")
+pidinet_detector = pidinet_detector.to(device)
+
+canndy_detector = CannyDetector()
+
+midas_detector = MidasDetector.from_pretrained(
+    "valhalla/t2iadapter-aux-models", filename="dpt_large_384.pt", model_type="dpt_large"
+)
+midas_detector = midas_detector.to(device)
+
+adapters = MultiAdapter(
+    [
+        T2IAdapter.from_pretrained(
+            "TencentARC/t2i-adapter-lineart-sdxl-1.0",
+            torch_dtype=torch.float16,
+            varient="fp16",
+        ),
+        T2IAdapter.from_pretrained(
+            "TencentARC/t2i-adapter-canny-sdxl-1.0",
+            torch_dtype=torch.float16,
+            varient="fp16",
+        ),
+        # T2IAdapter.from_pretrained(
+        #     "TencentARC/t2i-adapter-sketch-sdxl-1.0",
+        #     torch_dtype=torch.float16,
+        #     varient="fp16",
+        # ),
+        # T2IAdapter.from_pretrained(
+        #     "TencentARC/t2i-adapter-depth-midas-sdxl-1.0",
+        #     torch_dtype=torch.float16,
+        #     varient="fp16",
+        # ),
+    ]
+)
+adapters = adapters.to(torch.float16)
+
+pipeline = DiffusionPipeline.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.float16, 
+    variant="fp16", 
+    use_safetensors=True,
+    adapter=adapters,
+    custom_pipeline="./pipelines/pipeline_sdxl_adapter_img2img.py",
+)
+pipeline = pipeline.to(device)
+
+pipeline.scheduler = DDPMScheduler.from_pretrained(
+    BASE_MODEL,
+    subfolder="scheduler",
+)
+
+config = get_config(args)
+
+compel_proc = Compel(
+  tokenizer=[pipeline.tokenizer, pipeline.tokenizer_2] ,
+  text_encoder=[pipeline.text_encoder, pipeline.text_encoder_2],
+  returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+  requires_pooled=[False, True]
+)
+
+def run(
+    input_image:Image,
+    src_prompt:str,
+    tgt_prompt:str,
+    generate_size:int,
+    seed:int,
+    w1:float,
+    w2:float,
+    num_steps:int,
+    start_step:int,
+    guidance_scale:float,
+    lineart_scale:float = 0.5,
+    canny_scale:float = 0.5,
+    lineart_detect:float = 0.375,
+    canny_detect:float = 0.375,
+):
+    generator = torch.Generator().manual_seed(seed)
+
+    config.num_steps_inversion = num_steps
+    config.step_start = start_step
+    num_steps_actual = get_num_steps_actual(config)
+    
+
+    num_steps_inversion = config.num_steps_inversion
+    denoising_start = (num_steps_inversion - num_steps_actual) / num_steps_inversion
+    print(f"-------->num_steps_inversion: {num_steps_inversion} num_steps_actual: {num_steps_actual} denoising_start: {denoising_start}")
+    
+    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,
+        denoising_start=denoising_start,
+        strength=0,
+        device=device,
+    )
+    timesteps = timesteps.type(torch.int64)
+
+    timesteps = [torch.tensor(t) for t in timesteps.tolist()]
+    timesteps_len = len(timesteps)
+    config.step_start = start_step + num_steps_actual - timesteps_len
+    num_steps_actual = timesteps_len
+    config.max_norm_zs = [-1] * (num_steps_actual - 1) + [15.5]
+    print(f"-------->num_steps_inversion: {num_steps_inversion} num_steps_actual: {num_steps_actual} step_start: {config.step_start}")
+    print(f"-------->timesteps len: {len(timesteps)} max_norm_zs len: {len(config.max_norm_zs)}")
+    lineart_image = lineart_detector(input_image, detect_resolution=int(generate_size * lineart_detect), image_resolution=generate_size)
+    canny_image = canndy_detector(input_image, detect_resolution=int(generate_size * canny_detect), image_resolution=generate_size)
+    # pidinet_image = pidinet_detector(input_image, detect_resolution=512, image_resolution=generate_size, apply_filter=True)
+    # depth_image = midas_detector(input_image, detect_resolution=512, image_resolution=generate_size)
+    cond_image = [lineart_image, canny_image]
+    conditioning_scale = [lineart_scale, canny_scale]
+    pipeline.__call__ = partial(
+        pipeline.__call__,
+        num_inference_steps=num_steps_inversion,
+        guidance_scale=guidance_scale,
+        generator=generator,
+        denoising_start=denoising_start,
+        strength=0,
+        adapter_image=cond_image,
+        adapter_conditioning_scale=conditioning_scale,
+    )
+
+    x_0_image = input_image
+    x_0 = encode_image(x_0_image, pipeline)
+    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] * num_steps_actual
+    config.ws2 = [w2] * num_steps_actual
+    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]
+    conditioning, pooled = compel_proc(prompt)
+
+    image = pipeline.__call__(
+        image=latent,
+        prompt_embeds=conditioning,
+        pooled_prompt_embeds=pooled,
+        eta=1,
+    ).images
+    return image[2]
+
+def encode_image(image, pipe):
+    image = pipe.image_processor.preprocess(image)
+    originDtype = pipe.dtype
+    image = image.to(device=device, dtype=originDtype)
+
+    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(originDtype)
+
+    init_latents = init_latents.to(originDtype)
+    init_latents = pipe.vae.config.scaling_factor * init_latents
+
+    return init_latents.to(dtype=torch.float16)
+
+def get_timesteps(pipe, 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 = pipe.scheduler.timesteps[t_start * pipe.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(
+                pipe.scheduler.config.num_train_timesteps
+                - (denoising_start * pipe.scheduler.config.num_train_timesteps)
+            )
+        )
+
+        num_inference_steps = (timesteps < discrete_timestep_cutoff).sum().item()
+        if pipe.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 derivative) 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

inversion_utils.py

@@ -0,0 +1,794 @@
+import torch
+import os
+import PIL
+
+from typing import List, Optional, Union
+from diffusers.schedulers.scheduling_ddim import DDIMSchedulerOutput
+from PIL import Image
+from diffusers.utils import logging
+
+VECTOR_DATA_FOLDER = "vector_data"
+VECTOR_DATA_DICT = "vector_data"
+
+logger = logging.get_logger(__name__)
+
+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 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 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
+
+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 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 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

pipelines/pipeline_sdxl_adapter_img2img.py

@@ -0,0 +1,1672 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import PIL.Image
+import torch
+from transformers import (
+    CLIPImageProcessor,
+    CLIPTextModel,
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    CLIPVisionModelWithProjection,
+)
+
+from diffusers.callbacks import (
+    MultiPipelineCallbacks,
+    PipelineCallback,
+)
+
+from diffusers.image_processor import (
+    PipelineImageInput,
+    VaeImageProcessor,
+)
+
+from diffusers.loaders import (
+    FromSingleFileMixin,
+    IPAdapterMixin,
+    StableDiffusionXLLoraLoaderMixin,
+    TextualInversionLoaderMixin,
+)
+
+from diffusers.models import (
+    AutoencoderKL,
+    ImageProjection,
+    MultiAdapter,
+    T2IAdapter,
+    UNet2DConditionModel,
+)
+
+from diffusers.models.attention_processor import (
+    AttnProcessor2_0,
+    XFormersAttnProcessor,
+)
+
+from diffusers.models.lora import (
+    adjust_lora_scale_text_encoder,
+)
+
+from diffusers.schedulers import (
+    KarrasDiffusionSchedulers,
+)
+
+from diffusers.utils import (
+    PIL_INTERPOLATION,
+    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,
+)
+
+
+if is_invisible_watermark_available():
+    from diffusers.pipelines.stable_diffusion_xl.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]
+        ```
+"""
+
+
+def _preprocess_adapter_image(image, height, width):
+    if isinstance(image, torch.Tensor):
+        return image
+    elif isinstance(image, PIL.Image.Image):
+        image = [image]
+
+    if isinstance(image[0], PIL.Image.Image):
+        image = [np.array(i.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])) for i in image]
+        image = [
+            i[None, ..., None] if i.ndim == 2 else i[None, ...] for i in image
+        ]  # expand [h, w] or [h, w, c] to [b, h, w, c]
+        image = np.concatenate(image, axis=0)
+        image = np.array(image).astype(np.float32) / 255.0
+        image = image.transpose(0, 3, 1, 2)
+        image = torch.from_numpy(image)
+    elif isinstance(image[0], torch.Tensor):
+        if image[0].ndim == 3:
+            image = torch.stack(image, dim=0)
+        elif image[0].ndim == 4:
+            image = torch.cat(image, dim=0)
+        else:
+            raise ValueError(
+                f"Invalid image tensor! Expecting image tensor with 3 or 4 dimension, but recive: {image[0].ndim}"
+            )
+    return image
+
+
+# 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,
+    sigmas: Optional[List[float]] = 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 override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` 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 and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    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)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, 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,
+        adapter: Union[T2IAdapter, MultiAdapter, List[T2IAdapter]],
+        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,
+            adapter=adapter,
+            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.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.Tensor] = 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.Tensor`, *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.Tensor`, *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.Tensor`, *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.Tensor`, *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_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
+    ):
+        image_embeds = []
+        if do_classifier_free_guidance:
+            negative_image_embeds = []
+        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."
+                )
+
+            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
+                )
+
+                image_embeds.append(single_image_embeds[None, :])
+                if do_classifier_free_guidance:
+                    negative_image_embeds.append(single_negative_image_embeds[None, :])
+        else:
+            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)
+                    negative_image_embeds.append(single_negative_image_embeds)
+                image_embeds.append(single_image_embeds)
+
+        ip_adapter_image_embeds = []
+        for i, single_image_embeds in enumerate(image_embeds):
+            single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0)
+            if do_classifier_free_guidance:
+                single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0)
+                single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], dim=0)
+
+            single_image_embeds = single_image_embeds.to(device=device)
+            ip_adapter_image_embeds.append(single_image_embeds)
+
+        return ip_adapter_image_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 derivative) 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)}"
+            )
+
+        latents_mean = latents_std = None
+        if hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None:
+            latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1)
+        if hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None:
+            latents_std = torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1)
+
+        # 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):
+                if image.shape[0] < batch_size and batch_size % image.shape[0] == 0:
+                    image = torch.cat([image] * (batch_size // image.shape[0]), dim=0)
+                elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0:
+                    raise ValueError(
+                        f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} "
+                    )
+
+                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)
+            if latents_mean is not None and latents_std is not None:
+                latents_mean = latents_mean.to(device=device, dtype=dtype)
+                latents_std = latents_std.to(device=device, dtype=dtype)
+                init_latents = (init_latents - latents_mean) * self.vae.config.scaling_factor / latents_std
+            else:
+                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
+    ):
+        image_embeds = []
+        if do_classifier_free_guidance:
+            negative_image_embeds = []
+        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."
+                )
+
+            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
+                )
+
+                image_embeds.append(single_image_embeds[None, :])
+                if do_classifier_free_guidance:
+                    negative_image_embeds.append(single_negative_image_embeds[None, :])
+        else:
+            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)
+                    negative_image_embeds.append(single_negative_image_embeds)
+                image_embeds.append(single_image_embeds)
+
+        ip_adapter_image_embeds = []
+        for i, single_image_embeds in enumerate(image_embeds):
+            single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0)
+            if do_classifier_free_guidance:
+                single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0)
+                single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], dim=0)
+
+            single_image_embeds = single_image_embeds.to(device=device)
+            ip_adapter_image_embeds.append(single_image_embeds)
+
+        return ip_adapter_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,
+            ),
+        )
+        # 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.t2i_adapter.pipeline_stable_diffusion_adapter.StableDiffusionAdapterPipeline._default_height_width
+    def _default_height_width(self, height, width, image):
+        # NOTE: It is possible that a list of images have different
+        # dimensions for each image, so just checking the first image
+        # is not _exactly_ correct, but it is simple.
+        while isinstance(image, list):
+            image = image[0]
+
+        if height is None:
+            if isinstance(image, PIL.Image.Image):
+                height = image.height
+            elif isinstance(image, torch.Tensor):
+                height = image.shape[-2]
+
+            # round down to nearest multiple of `self.adapter.downscale_factor`
+            height = (height // self.adapter.downscale_factor) * self.adapter.downscale_factor
+
+        if width is None:
+            if isinstance(image, PIL.Image.Image):
+                width = image.width
+            elif isinstance(image, torch.Tensor):
+                width = image.shape[-1]
+
+            # round down to nearest multiple of `self.adapter.downscale_factor`
+            width = (width // self.adapter.downscale_factor) * self.adapter.downscale_factor
+
+        return height, width
+
+    # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
+    def get_guidance_scale_embedding(
+        self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32
+    ) -> torch.Tensor:
+        """
+        See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
+
+        Args:
+            w (`torch.Tensor`):
+                Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings.
+            embedding_dim (`int`, *optional*, defaults to 512):
+                Dimension of the embeddings to generate.
+            dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
+                Data type of the generated embeddings.
+
+        Returns:
+            `torch.Tensor`: Embedding vectors with shape `(len(w), 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,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        adapter_image: PipelineImageInput = None,
+        strength: float = 0.3,
+        num_inference_steps: int = 50,
+        timesteps: List[int] = None,
+        sigmas: List[float] = 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.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = 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,
+        adapter_conditioning_scale: Union[float, List[float]] = 1.0,
+        adapter_conditioning_factor: float = 1.0,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[
+            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+        ] = 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.Tensor` or `PIL.Image.Image` or `np.ndarray` or `List[torch.Tensor]` 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.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            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.Tensor`, *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.Tensor`, *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.Tensor`, *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.Tensor`, *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.Tensor`, *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.Tensor]`, *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`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
+                A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
+                each denoising step during the inference. 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.
+        """
+        height, width = self._default_height_width(height, width, adapter_image)
+        device = self._execution_device
+
+        if isinstance(self.adapter, MultiAdapter):
+            adapter_input = []
+
+            for one_image in adapter_image:
+                one_image = _preprocess_adapter_image(one_image, height, width)
+                one_image = one_image.to(device=device, dtype=self.adapter.dtype)
+                adapter_input.append(one_image)
+        else:
+            adapter_input = _preprocess_adapter_image(adapter_image, height, width)
+            adapter_input = adapter_input.to(device=device, dtype=self.adapter.dtype)
+
+        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`",
+            )
+
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        # 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, sigmas
+        )
+        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
+        if latents is None:
+            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
+        # adapter_input = adapter_input.type(latents.dtype)
+        if isinstance(self.adapter, MultiAdapter):
+            adapter_state = self.adapter(adapter_input, adapter_conditioning_scale)
+            for k, v in enumerate(adapter_state):
+                adapter_state[k] = v
+        else:
+            adapter_state = self.adapter(adapter_input)
+            for k, v in enumerate(adapter_state):
+                adapter_state[k] = v * adapter_conditioning_scale
+        if num_images_per_prompt > 1:
+            for k, v in enumerate(adapter_state):
+                adapter_state[k] = v.repeat(num_images_per_prompt, 1, 1, 1)
+        if self.do_classifier_free_guidance:
+            for k, v in enumerate(adapter_state):
+                adapter_state[k] = torch.cat([v] * 2, dim=0)
+
+        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
+                
+                if i < int(num_inference_steps * adapter_conditioning_factor):
+                    down_intrablock_additional_residuals = [state.clone() for state in adapter_state]
+                else:
+                    down_intrablock_additional_residuals = None
+                
+                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,
+                    down_intrablock_additional_residuals=down_intrablock_additional_residuals,
+                )[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
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                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)
+            elif latents.dtype != self.vae.dtype:
+                if torch.backends.mps.is_available():
+                    # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                    self.vae = self.vae.to(latents.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)

requirements.txt

@@ -0,0 +1,17 @@
+ml-collections
+gradio
+torch
+torchvision
+diffusers
+transformers
+accelerate
+mediapipe
+spaces
+sentencepiece
+compel
+gfpgan
+git+https://github.com/XPixelGroup/BasicSR@master
+facexlib
+realesrgan
+controlnet_aux
+peft

run_configs/noise_shift_3_steps.yaml

@@ -0,0 +1,19 @@
+breakdown: "x_t_hat_c"
+cross_r: 0.9
+eta_reconstruct: 1
+eta_retrieve: 1
+max_norm_zs: [-1, -1, 15.5]
+model: "stabilityai/sdxl-turbo"
+noise_shift_delta: 1
+noise_timesteps: [599, 299, 0]
+timesteps: [799, 499, 199]
+num_steps_inversion: 5
+step_start: 1
+real_cfg_scale: 0
+real_cfg_scale_save: 0
+scheduler_type: "ddpm"
+seed: 2
+self_r: 0.5
+ws1: 1.5
+ws2: 1
+clean_step_timestep: 0

run_configs/noise_shift_guidance_1_5.yaml

@@ -0,0 +1,18 @@
+breakdown: "x_t_hat_c"
+cross_r: 0.9
+eta: 1
+max_norm_zs: [-1, -1, -1, 15.5]
+model: ""
+noise_shift_delta: 1
+noise_timesteps: null
+num_steps_inversion: 20
+step_start: 5
+real_cfg_scale: 0
+real_cfg_scale_save: 0
+scheduler_type: "ddpm"
+seed: 2
+self_r: 0.5
+timesteps: null
+ws1: 1.5
+ws2: 1
+clean_step_timestep: 0

segment_utils.py

@@ -0,0 +1,98 @@
+import numpy as np
+import mediapipe as mp
+import uuid
+
+from PIL import Image
+from mediapipe.tasks import python
+from mediapipe.tasks.python import vision
+from scipy.ndimage import binary_dilation
+from croper import Croper
+
+segment_model = "checkpoints/selfie_multiclass_256x256.tflite"
+base_options = python.BaseOptions(model_asset_path=segment_model)
+options = vision.ImageSegmenterOptions(base_options=base_options,output_category_mask=True)
+segmenter = vision.ImageSegmenter.create_from_options(options)
+
+def restore_result(croper, category, generated_image):
+    square_length = croper.square_length
+    generated_image = generated_image.resize((square_length, square_length))
+
+    cropped_generated_image = generated_image.crop((croper.square_start_x, croper.square_start_y, croper.square_end_x, croper.square_end_y))
+    cropped_square_mask_image = get_restore_mask_image(croper, category, cropped_generated_image)
+
+    restored_image = croper.input_image.copy()
+    restored_image.paste(cropped_generated_image, (croper.origin_start_x, croper.origin_start_y), cropped_square_mask_image)
+
+    extension = 'png'
+    if restored_image.mode == 'RGBA':
+        extension = 'png'
+    else:
+        extension = 'jpg'
+
+    path = f"output/{uuid.uuid4()}.{extension}"
+    restored_image.save(path)
+
+    return restored_image, path
+
+def segment_image(input_image, category, input_size, mask_expansion, mask_dilation):
+    mask_size = int(input_size)
+    mask_expansion = int(mask_expansion)
+
+    image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(input_image))
+    segmentation_result = segmenter.segment(image)
+    category_mask = segmentation_result.category_mask
+    category_mask_np = category_mask.numpy_view()
+
+    if category == "hair":
+        target_mask = get_hair_mask(category_mask_np, mask_dilation)
+    elif category == "clothes":
+        target_mask = get_clothes_mask(category_mask_np, mask_dilation)
+    elif category == "face":
+        target_mask = get_face_mask(category_mask_np, mask_dilation)
+    else:
+        target_mask = get_face_mask(category_mask_np, mask_dilation)
+    
+    croper = Croper(input_image, target_mask, mask_size, mask_expansion)
+    croper.corp_mask_image()
+    origin_area_image = croper.resized_square_image
+
+    return origin_area_image, croper
+
+def get_face_mask(category_mask_np, dilation=1):
+    face_skin_mask = category_mask_np == 3
+    if dilation > 0:
+        face_skin_mask = binary_dilation(face_skin_mask, iterations=dilation)
+
+    return face_skin_mask
+
+def get_clothes_mask(category_mask_np, dilation=1):
+    body_skin_mask = category_mask_np == 2
+    clothes_mask = category_mask_np == 4
+    combined_mask = np.logical_or(body_skin_mask, clothes_mask)
+    combined_mask = binary_dilation(combined_mask, iterations=4)
+    if dilation > 0:
+        combined_mask = binary_dilation(combined_mask, iterations=dilation)
+    return combined_mask
+
+def get_hair_mask(category_mask_np, dilation=1):
+    hair_mask = category_mask_np == 1
+    if dilation > 0:
+        hair_mask = binary_dilation(hair_mask, iterations=dilation)
+    return hair_mask
+
+def get_restore_mask_image(croper, category, generated_image):
+    image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(generated_image))
+    segmentation_result = segmenter.segment(image)
+    category_mask = segmentation_result.category_mask
+    category_mask_np = category_mask.numpy_view()
+
+    if category == "hair":
+        target_mask = get_hair_mask(category_mask_np, 0)
+    elif category == "clothes":
+        target_mask = get_clothes_mask(category_mask_np, 0)
+    elif category == "face":
+        target_mask = get_face_mask(category_mask_np, 0)
+    
+    combined_mask = np.logical_or(target_mask, croper.corp_mask)
+    mask_image = Image.fromarray((combined_mask * 255).astype(np.uint8))
+    return mask_image