initial commit

README.md

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

app.py

@@ -0,0 +1,10 @@
+import gradio as gr
+
+from app_base import create_demo as create_demo_face
+
+with gr.Blocks(css="style.css") as demo:
+    with gr.Tabs():
+        with gr.Tab(label="Face"):
+            create_demo_face()
+
+demo.launch(server_name="0.0.0.0")

app_base.py

@@ -0,0 +1,209 @@
+import spaces
+import gradio as gr
+import time
+import torch
+import tempfile
+import os
+import gc
+
+from loading_utils import load_image
+
+from segment_utils import(
+    segment_image,
+    restore_result,
+)
+from enhance_utils import enhance_sd_image
+from inversion_run_base import run as base_run
+
+DEFAULT_SRC_PROMPT = "a person"
+DEFAULT_EDIT_PROMPT = "a person with perfect face"
+
+DEFAULT_CATEGORY = "face"
+
+def image_to_image(
+    input_image_path: str,
+    input_image_prompt: str,
+    edit_prompt: str,
+    seed: int,
+    w1: float,
+    num_steps: int,
+    start_step: int,
+    guidance_scale: float,
+    generate_size: int,
+    mask_expansion: int = 50,
+    mask_dilation: int = 2,
+    save_quality: int = 95,
+    enable_segment: bool = True,
+):
+    segment_category = "face"
+    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)
+    input_image = load_image(input_image_path)
+    icc_profile = input_image.info.get('icc_profile')
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'load_image done')
+
+    if enable_segment:
+        target_area_image, croper = segment_image(
+            input_image,
+            segment_category,
+            generate_size,
+            mask_expansion,
+            mask_dilation,
+        )
+    else:
+        target_area_image = resize_image(input_image, generate_size)
+        croper = None
+    
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'segment_image done')
+    
+    run_model = base_run
+    try:
+        res_image = run_model(
+            target_area_image,
+            input_image_prompt,
+            edit_prompt ,
+            seed,
+            w1,
+            w2,
+            num_steps,
+            start_step,
+            guidance_scale,
+        )
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'run_sd_model done')
+        
+    finally:
+        torch.cuda.empty_cache()
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'cuda_empty_cache done')
+
+    enhanced_image = res_image
+    enhanced_image = enhance_sd_image(res_image)
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'enhance_image done')
+
+    if enable_segment:
+        restored_image = restore_result(croper, segment_category, enhanced_image)
+    else:
+        restored_image = enhanced_image.resize(input_image.size)
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'restore_result done')
+
+    torch.cuda.empty_cache()
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'cuda_empty_cache done')
+    if os.getenv('ENABLE_GC', False):
+        gc.collect()
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'gc_collect done')
+
+    extension = 'png'
+    if restored_image.mode == 'RGBA':
+        extension = 'png'
+    else:
+        extension = 'webp'
+
+    output_path = tempfile.mktemp(suffix=f".{extension}")
+    restored_image.save(output_path, format=extension, quality=save_quality, icc_profile=icc_profile)
+
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'save_image done')
+
+    return output_path, restored_image, time_cost_str
+
+def get_time_cost(
+    run_task_time, 
+    time_cost_str,
+    step: 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}'
+        if step != '':
+            time_cost_str += f'-->{step}'
+    run_task_time = now_time
+    return run_task_time, time_cost_str
+
+def resize_image(image, target_size = 1024):
+    h, w = image.size
+    if h >= w:
+        w = int(w * target_size / h)
+        h = target_size
+    else:
+        h = int(h * target_size / w)
+        w = target_size
+    return image.resize((w, h))
+
+
+def infer(
+    input_image_path: str,
+    input_image_prompt: str,
+    edit_prompt: str,
+    seed: int,
+    w1: float,
+    num_steps: int,
+    start_step: int,
+    guidance_scale: float,
+    generate_size: int,
+    mask_expansion: int = 50,
+    mask_dilation: int = 2,
+    save_quality: int = 95,
+    enable_segment: bool = True,
+):
+    return image_to_image(
+        input_image_path,
+        input_image_prompt,
+        edit_prompt,
+        seed,
+        w1,
+        num_steps,
+        start_step,
+        guidance_scale,
+        generate_size,
+        mask_expansion,
+        mask_dilation,
+        save_quality,
+        enable_segment
+    )
+
+infer = spaces.GPU(infer)
+
+def create_demo() -> gr.Blocks:
+
+    with gr.Blocks() as demo:
+        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)
+                with gr.Accordion("Advanced Options", open=False):
+                    enable_segment = gr.Checkbox(label="Enable Segment", value=True)
+                    mask_expansion = gr.Number(label="Mask Expansion", value=50, visible=True)
+                    mask_dilation = gr.Slider(minimum=0, maximum=10, value=2, step=1, label="Mask Dilation")
+                    save_quality = gr.Slider(minimum=1, maximum=100, value=95, step=1, label="Save Quality")
+            with gr.Column():
+                num_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Num Steps")
+                start_step = gr.Slider(minimum=1, maximum=100, value=15, step=1, label="Start Step")
+                g_btn = gr.Button("Edit Image")
+                with gr.Accordion("Advanced Options", open=False):
+                    guidance_scale = gr.Slider(minimum=0, maximum=20, value=0, step=0.5, label="Guidance Scale")
+                    seed = gr.Number(label="Seed", value=8)
+                    w1 = gr.Number(label="W1", value=1.5)
+                    generate_size = gr.Number(label="Generate Size", value=1024)
+        
+        with gr.Row():
+            with gr.Column():
+                input_image_path = gr.Image(label="Input Image", type="filepath", interactive=True)
+            with gr.Column():
+                restored_image = gr.Image(label="Restored Image", format="png", type="pil", interactive=False)
+                download_path = gr.File(label="Download the output image", interactive=False)
+                generated_cost = gr.Textbox(label="Time cost by step (ms):", visible=True, interactive=False)
+        
+        g_btn.click(
+            fn=infer,
+            inputs=[input_image_path, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step, guidance_scale, generate_size, mask_expansion, mask_dilation, save_quality, enable_segment],
+            outputs=[download_path, restored_image, generated_cost],
+        )
+    
+    
+
+    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,44 @@
+import os
+
+from PIL import Image
+from model_handler import MODELS
+
+upscaler = MODELS.upscaler
+upscaler4SD = MODELS.upscaler4SD
+
+def enhance_image(
+    input_image: Image,
+):
+    
+    h, w = input_image.size
+    max_scale_size = 1024
+    if h > max_scale_size:
+        w = int(w * max_scale_size / h)
+        h = max_scale_size
+    if w > max_scale_size:
+        h = int(h * max_scale_size / w)
+        w = max_scale_size
+
+    if h != input_image.size[1] or w != input_image.size[0]:
+        input_image = input_image.resize((w, h))
+
+    if os.environ.get("TILING", False):
+        tileSizeStr = os.environ.get("TILE_SIZE", 1024)
+        tileSize = int(tileSizeStr)
+        enhanced_image = upscaler(input_image, tiling=True, tile_width=tileSize, tile_height=tileSize)
+    else:
+        enhanced_image = upscaler(input_image)
+
+    return enhanced_image
+
+def enhance_sd_image(
+    input_image: Image,
+):
+    if os.environ.get("TILING", False):
+        tileSizeStr = os.environ.get("TILE_SIZE", 1024)
+        tileSize = int(tileSizeStr)
+        enhanced_image = upscaler4SD(input_image, tiling=True, tile_width=tileSize, tile_height=tileSize)
+    else:
+        enhanced_image = upscaler4SD(input_image)
+
+    return enhanced_image

inversion_run_base.py

@@ -0,0 +1,198 @@
+import torch
+
+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
+
+from model_handler import MODELS
+
+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"
+
+pipeline = MODELS.base_pipe
+
+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,
+    seed:int,
+    w1:float,
+    w2:float,
+    num_steps:int,
+    start_step:int,
+    guidance_scale:float,
+):
+    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
+    
+    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]
+
+    pipeline.__call__ = partial(
+        pipeline.__call__,
+        num_inference_steps=num_steps_inversion,
+        guidance_scale=guidance_scale,
+        generator=generator,
+        denoising_start=denoising_start,
+        strength=0,
+    )
+
+    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,793 @@
+import torch
+import os
+import PIL
+
+from typing import List, Optional, Union
+from diffusers.schedulers.scheduling_ddim import DDIMSchedulerOutput
+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

loading_utils.py

@@ -0,0 +1,60 @@
+import os
+from typing import  Callable, Optional, Union
+
+import PIL.Image
+import PIL.ImageOps
+import requests
+
+
+session = requests.Session()
+adapter = requests.adapters.HTTPAdapter(
+    pool_connections=10,  
+    pool_maxsize=10,
+    max_retries=3
+)
+session.mount('http://', adapter)
+session.mount('https://', adapter)
+
+
+def load_image(
+    image: Union[str, PIL.Image.Image], convert_method: Optional[Callable[[PIL.Image.Image], PIL.Image.Image]] = None
+) -> PIL.Image.Image:
+    """
+    Loads `image` to a PIL Image.
+
+    Args:
+        image (`str` or `PIL.Image.Image`):
+            The image to convert to the PIL Image format.
+        convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], *optional*):
+            A conversion method to apply to the image after loading it. When set to `None` the image will be converted
+            "RGB".
+
+    Returns:
+        `PIL.Image.Image`:
+            A PIL Image.
+    """
+    if isinstance(image, str):
+        if image.startswith("http://") or image.startswith("https://"):
+            image = PIL.Image.open(session.get(image, stream=True).raw)
+        elif os.path.isfile(image):
+            image = PIL.Image.open(image)
+        else:
+            raise ValueError(
+                f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {image} is not a valid path."
+            )
+    elif isinstance(image, PIL.Image.Image):
+        image = image
+    else:
+        raise ValueError(
+            "Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image."
+        )
+
+    image = PIL.ImageOps.exif_transpose(image)
+
+    if convert_method is not None:
+        image = convert_method(image)
+    else:
+        image = image.convert("RGB")
+
+    return image
+

model_handler.py

@@ -0,0 +1,87 @@
+import torch
+import os
+
+from diffusers import (
+    DDPMScheduler,
+    StableDiffusionXLImg2ImgPipeline,
+    LTXPipeline,
+    AutoencoderKL,
+)
+
+from hidiffusion import apply_hidiffusion
+
+from mediapipe.tasks import python
+from mediapipe.tasks.python import vision
+
+from image_gen_aux import UpscaleWithModel
+
+BASE_MODEL = "stabilityai/sdxl-turbo"
+VIDEO_MODEL = "Lightricks/LTX-Video"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+class ModelHandler:
+    def __init__(self):
+        self.base_pipe = None
+        self.video_pipe = None
+        self.compiled_model = None
+        self.segmenter = None
+        self.upscaler = None
+        self.upscaler4SD = None
+        self.load_models()
+
+    def load_base(self):
+        vae = AutoencoderKL.from_pretrained(
+            "madebyollin/sdxl-vae-fp16-fix", 
+            torch_dtype=torch.float16,
+        )
+
+        base_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+            BASE_MODEL, 
+            vae=vae,
+            torch_dtype=torch.float16, 
+            variant="fp16", 
+            use_safetensors=True,
+        )
+        base_pipe = base_pipe.to(device, silence_dtype_warnings=True)
+        base_pipe.scheduler = DDPMScheduler.from_pretrained(
+            BASE_MODEL,
+            subfolder="scheduler",
+        )
+        apply_hidiffusion(base_pipe)
+        
+        return base_pipe
+    
+    def load_video_pipe(self):
+        pipe = LTXPipeline.from_pretrained(VIDEO_MODEL, torch_dtype=torch.bfloat16)
+        pipe.to(device)
+        return pipe
+    
+    def load_segmenter(self):
+        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)
+        return segmenter
+    
+    def load_upscaler(self):
+        model_name = os.environ.get("UPSCALE_MODEL", "Phips/4xNomosWebPhoto_RealPLKSR")
+        upscaler = UpscaleWithModel.from_pretrained(model_name).to(device)
+        return upscaler
+    
+    def load_upscaler4SD(self):
+        model_name = os.environ.get("UPSCALE_FOR_SD_MODEL", "Phips/1xDeJPG_realplksr_otf")
+        upscaler = UpscaleWithModel.from_pretrained(model_name).to(device)
+        return upscaler
+
+    def load_models(self):
+        base_pipe = self.load_base()
+        segmenter = self.load_segmenter()
+        upscaler = self.load_upscaler()
+        upscaler4SD = self.load_upscaler4SD()
+
+        self.base_pipe = base_pipe
+        self.segmenter = segmenter
+        self.upscaler = upscaler
+        self.upscaler4SD = upscaler4SD
+
+MODELS = ModelHandler()

requirements.txt

@@ -0,0 +1,12 @@
+torch
+torchvision
+ml-collections
+mediapipe
+gradio
+diffusers
+transformers
+accelerate
+sentencepiece
+compel
+hidiffusion
+git+https://github.com/asomoza/image_gen_aux.git

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,86 @@
+import numpy as np
+import mediapipe as mp
+import uuid
+
+from PIL import Image
+from scipy.ndimage import binary_dilation
+from croper import Croper
+
+from model_handler import MODELS
+
+segmenter = MODELS.segmenter
+
+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)
+
+    return restored_image
+
+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