commit message

.gitattributes

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

.gitignore

@@ -0,0 +1,3 @@
+tmp
+app_fin_v1.py
+__*.py

README.md

@@ -0,0 +1,13 @@
+---
+title: Demo Temp
+emoji: 📈
+colorFrom: purple
+colorTo: indigo
+sdk: gradio
+sdk_version: 3.18.0
+app_file: app.py
+pinned: false
+license: mit
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,143 @@
+import os
+
+from PIL import Image
+import gradio as gr
+
+from utils.gradio_utils import *
+from utils.direction_utils import *
+from utils.generate_synthetic import *
+
+
+if __name__=="__main__":
+
+    # populate the list of editing directions
+    d_name2desc = get_all_directions_names()
+    d_name2desc["make your own!"] = "make your own!"
+
+    with gr.Blocks(css=CSS_main) as demo:
+        # Make the header of the demo website
+        gr.HTML(HTML_header)
+
+        with gr.Row():
+            # col A: the input image or synthetic image prompt
+            with gr.Column(scale=2) as gc_left:
+                gr.HTML(" <center> <p style='font-size:150%;'> input </p> </center>")
+                img_in_real = gr.Image(type="pil", label="Start by uploading an image", elem_id="input_image")
+                img_in_synth = gr.Image(type="pil", label="Synthesized image", elem_id="input_image_synth", visible=False)
+                gr.Examples( examples="assets/test_images/", inputs=[img_in_real])
+                prompt = gr.Textbox(value="a high resolution painting of a cat in the style of van gogh", label="Or use a synthetic image. Prompt:", interactive=True)
+                with gr.Row():
+                    seed = gr.Number(value=42, label="random seed:", interactive=True)
+                    negative_guidance = gr.Number(value=5, label="negative guidance:", interactive=True)
+                btn_generate = gr.Button("Generate", label="")
+                fpath_z_gen = gr.Textbox(value="placeholder", visible=False)
+
+            # col B: the output image
+            with gr.Column(scale=2) as gc_left:
+                gr.HTML(" <center> <p style='font-size:150%;'> output </p> </center>")
+                img_out = gr.Image(type="pil", label="Output Image", visible=True)
+                with gr.Row():
+                    with gr.Column():
+                        src = gr.Dropdown(list(d_name2desc.values()), label="source", interactive=True, value="cat")
+                        src_custom = gr.Textbox(placeholder="enter new task here!", interactive=True, visible=False, label="custom source direction:")
+                        rad_src = gr.Radio(["GPT3", "flan-t5-xl (free)!", "BLOOMZ-7B (free)!", "fixed-template", "custom sentences"], label="Sentence type:", value="GPT3", interactive=True, visible=False)
+                        custom_sentences_src = gr.Textbox(placeholder="paste list of sentences here", interactive=True, visible=False, label="custom sentences:", lines=5, max_lines=20)
+                    
+
+                    with gr.Column():
+                        dest = gr.Dropdown(list(d_name2desc.values()), label="target", interactive=True, value="dog")
+                        dest_custom = gr.Textbox(placeholder="enter new task here!", interactive=True, visible=False, label="custom target direction:")
+                        rad_dest = gr.Radio(["GPT3", "flan-t5-xl (free)!", "BLOOMZ-7B (free)!", "fixed-template", "custom sentences"], label="Sentence type:", value="GPT3", interactive=True, visible=False)
+                        custom_sentences_dest = gr.Textbox(placeholder="paste list of sentences here", interactive=True, visible=False, label="custom sentences:", lines=5, max_lines=20)
+                    
+                    
+                    with gr.Row():
+                        api_key = gr.Textbox(placeholder="enter you OpenAI API key here", interactive=True, visible=False, label="OpenAI API key:", type="password")
+                        org_key = gr.Textbox(placeholder="enter you OpenAI organization key here", interactive=True, visible=False, label="OpenAI Organization:", type="password")
+                with gr.Row():
+                    btn_edit = gr.Button("Run", label="")
+                    btn_clear = gr.Button("Clear")
+
+                with gr.Accordion("Change editing settings?", open=False):
+                    num_ddim = gr.Slider(0, 200, 100, label="Number of DDIM steps", interactive=True, elem_id="slider_ddim", step=10)
+                    xa_guidance = gr.Slider(0, 0.25, 0.1, label="Cross Attention guidance", interactive=True, elem_id="slider_xa", step=0.01)
+                    edit_mul = gr.Slider(0, 2, 1.0, label="Edit multiplier", interactive=True, elem_id="slider_edit_mul", step=0.05)
+
+                with gr.Accordion("Generating your own directions", open=False):
+                    gr.Textbox("We provide 5 different ways of computing new custom directions:", show_label=False)
+                    gr.Textbox("We use GPT3 to generate a list of sentences that describe the desired edit. For this options, the users need to make an OpenAI account and enter the API and organizations keys. This option typically results is the best directions and costs roughly $0.14 for one concept.", label="1. GPT3", show_label=True)
+                    gr.Textbox("Alternatively flan-t5-xl model can also be used to to generate a list of sentences that describe the desired edit. This option is free and does not require creating any new accounts.", label="2. flan-t5-xl (free)", show_label=True)
+                    gr.Textbox("Similarly BLOOMZ-7B model can also be used to to generate the sentences for free.", label="3. BLOOMZ-7B (free)", show_label=True)
+                    gr.Textbox("Next, we provide a fixed template based sentence generation. This option does not require any language model and is therefore free and much faster. However the edit directions with this method are often entangled.", label="4. Fixed template", show_label=True)
+                    gr.Textbox("Finally, the user can also generate their own sentences.", label="5. Custom sentences", show_label=True)
+
+
+                with gr.Accordion("Tips for getting better results", open=True):
+                    gr.Textbox("The 'Cross Attention guidance' controls the amount of structure guidance to be applied when performing the edit. If the output edited image does not retain the structure from the input, increasing the value will typically address the issue. We recommend changing the value in increments of 0.05.", label="1. Controlling the image structure", show_label=True)
+                    gr.Textbox("If the output image quality is low or has some artifacts, using more steps would be helpful. This can be controlled with the 'Number of DDIM steps' slider.", label="2. Improving Image Quality", show_label=True)
+                    gr.Textbox("There can be two reasons why the output image does not have the desired edit applied. Either the cross attention guidance is too strong, or the edit is insufficient. These can be addressed by reducing the 'Cross Attention guidance' slider or increasing the 'Edit multiplier' respectively.", label="3. Amount of edit applied", show_label=True)
+
+                
+
+
+
+                
+        # txt_image_type = gr.Textbox(visible=False)
+        btn_generate.click(launch_generate_sample, [prompt, seed, negative_guidance, num_ddim], [img_in_synth, fpath_z_gen])
+        btn_generate.click(set_visible_true, [], img_in_synth)
+        btn_generate.click(set_visible_false, [], img_in_real)
+
+        def fn_clear_all():
+            return gr.update(value=None), gr.update(value=None), gr.update(value=None)
+        btn_clear.click(fn_clear_all, [], [img_out, img_in_real, img_in_synth])
+        btn_clear.click(set_visible_true, [], img_in_real)
+        btn_clear.click(set_visible_false, [], img_in_synth)
+
+        
+        # handling custom directions
+        def on_custom_seleceted(src):
+            if src=="make your own!": return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+            else: return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+        
+        src.change(on_custom_seleceted, [src], [src_custom, rad_src, api_key, org_key])
+        dest.change(on_custom_seleceted, [dest], [dest_custom, rad_dest, api_key, org_key])
+
+
+        def fn_sentence_type_change(rad):
+            print(rad)
+            if rad=="GPT3":
+                return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)
+            elif rad=="custom sentences":
+                return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+            else:
+                print("using template sentence or flan-t5-xl or bloomz-7b")
+                return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
+        rad_dest.change(fn_sentence_type_change, [rad_dest], [api_key, org_key, custom_sentences_dest])
+        rad_src.change(fn_sentence_type_change, [rad_src], [api_key, org_key, custom_sentences_src])
+
+        btn_edit.click(launch_main, 
+                        [ 
+                            img_in_real, img_in_synth,
+                            src, src_custom, dest,
+                            dest_custom, num_ddim,
+                            xa_guidance, edit_mul,
+                            fpath_z_gen, prompt,
+                            rad_src, rad_dest,
+                            api_key, org_key,
+                            custom_sentences_src, custom_sentences_dest
+                        ],
+                [img_out])
+        
+
+
+        gr.HTML("<hr>")
+
+    
+    demo.queue(concurrency_count=8)
+    demo.launch(debug=True)
+
+    # gr.close_all()
+    # demo.launch(server_port=8089, server_name="0.0.0.0", debug=True)
+
+

environment.yml

@@ -0,0 +1,23 @@
+name: pix2pix-zero
+channels:
+  - pytorch
+  - nvidia
+  - defaults
+dependencies:
+  - pip
+  - pytorch-cuda=11.6
+  - torchvision
+  - pytorch
+  - pip:
+    - accelerate
+    - diffusers
+    - einops
+    - gradio
+    - ipython
+    - numpy
+    - opencv-python-headless
+    - pillow
+    - psutil
+    - tqdm
+    - transformers
+    - salesforce-lavis

requirements.txt

@@ -0,0 +1,7 @@
+transformers
+joblib
+accelerate
+diffusers==0.12.1
+salesforce-lavis
+openai
+#git+https://github.com/pix2pixzero/pix2pix-zero.git

submodules/pix2pix-zero/.gitignore

@@ -0,0 +1,6 @@
+output
+scripts
+src/folder_*.py
+src/ig_*.py
+assets/edit_sentences
+src/utils/edit_pipeline_spatial.py

submodules/pix2pix-zero/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 pix2pixzero
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

submodules/pix2pix-zero/README.md

@@ -0,0 +1,154 @@
+# pix2pix-zero
+
+## [**[website]**](https://pix2pixzero.github.io/)
+
+
+This is author's reimplementation of "Zero-shot Image-to-Image Translation" using the diffusers library. <br>
+The results in the paper are based on the [CompVis](https://github.com/CompVis/stable-diffusion) library, which will be released later. 
+
+**[New!]** Code for editing real and synthetic images released!
+
+
+
+<br>
+<div class="gif">
+<p align="center">
+<img src='assets/main.gif' align="center">
+</p>
+</div>
+
+
+We propose pix2pix-zero, a diffusion-based image-to-image approach that allows users to specify the edit direction on-the-fly (e.g., cat to dog). Our method can directly use pre-trained [Stable Diffusion](https://github.com/CompVis/stable-diffusion), for editing real and synthetic images while preserving the input image's structure. Our method is training-free and prompt-free, as it requires neither manual text prompting for each input image nor costly fine-tuning for each task.
+
+**TL;DR**: no finetuning required, no text input needed, input structure preserved.
+
+## Results
+All our results are based on [stable-diffusion-v1-4](https://github.com/CompVis/stable-diffusion) model. Please the website for more results.
+
+<div>
+<p align="center">
+<img src='assets/results_teaser.jpg' align="center" width=800px>
+</p>
+</div>
+<hr>
+
+The top row for each of the results below show editing of real images, and the bottom row shows synthetic image editing.
+<div>
+<p align="center">
+<img src='assets/grid_dog2cat.jpg' align="center" width=800px>
+</p>
+<p align="center">
+<img src='assets/grid_zebra2horse.jpg' align="center" width=800px>
+</p>
+<p align="center">
+<img src='assets/grid_cat2dog.jpg' align="center" width=800px>
+</p>
+<p align="center">
+<img src='assets/grid_horse2zebra.jpg' align="center" width=800px>
+</p>
+<p align="center">
+<img src='assets/grid_tree2fall.jpg' align="center" width=800px>
+</p>
+</div>
+
+## Real Image Editing
+<div>
+<p align="center">
+<img src='assets/results_real.jpg' align="center" width=800px>
+</p>
+</div>
+
+## Synthetic Image Editing
+<div>
+<p align="center">
+<img src='assets/results_syn.jpg' align="center" width=800px>
+</p>
+</div>
+
+## Method Details
+
+Given an input image, we first generate text captions using [BLIP](https://github.com/salesforce/LAVIS) and apply regularized DDIM inversion to obtain our inverted noise map.
+Then, we obtain reference cross-attention maps that correspoind to the structure of the input image by denoising, guided with the CLIP embeddings 
+of our generated text (c). Next, we denoise with edited text embeddings, while enforcing a loss to match current cross-attention maps with the 
+reference cross-attention maps.
+
+<div>
+<p align="center">
+<img src='assets/method.jpeg' align="center" width=900>
+</p>
+</div>
+
+
+## Getting Started
+
+**Environment Setup**
+- We provide a [conda env file](environment.yml) that contains all the required dependencies
+  ```
+  conda env create -f environment.yml
+  ```
+- Following this, you can activate the conda environment with the command below. 
+  ```
+  conda activate pix2pix-zero
+  ```
+
+**Real Image Translation**
+- First, run the inversion command below to obtain the input noise that reconstructs the image. 
+  The command below will save the inversion in the results folder as `output/test_cat/inversion/cat_1.pt` 
+  and the BLIP-generated prompt as `output/test_cat/prompt/cat_1.txt`
+    ```
+    python src/inversion.py  \
+            --input_image "assets/test_images/cats/cat_1.png" \
+            --results_folder "output/test_cat"
+    ```
+- Next, we can perform image editing with the editing direction as shown below.
+  The command below will save the edited image as `output/test_cat/edit/cat_1.png`
+    ```
+    python src/edit_real.py \
+        --inversion "output/test_cat/inversion/cat_1.pt" \
+        --prompt "output/test_cat/prompt/cat_1.txt" \
+        --task_name "cat2dog" \
+        --results_folder "output/test_cat/" 
+    ```
+
+**Editing Synthetic Images**
+- Similarly, we can edit the synthetic images generated by Stable Diffusion with the following command.
+    ```
+    python src/edit_synthetic.py \
+        --results_folder "output/synth_editing" \
+        --prompt_str "a high resolution painting of a cat in the style of van gough" \
+        --task "cat2dog"
+    ```
+
+### **Tips and Debugging**
+  - **Controlling the Image Structure:**<br>
+    The `--xa_guidance` flag controls the amount of cross-attention guidance to be applied when performing the edit. If the output edited image does not retain the structure from the input, increasing the value will typically address the issue. We recommend changing the value in increments of 0.05. 
+
+  - **Improving Image Quality:**<br>
+    If the output image quality is low or has some artifacts, using more steps for both the inversion and editing would be helpful. 
+    This can be controlled with the `--num_ddim_steps` flag. 
+
+  - **Reducing the VRAM Requirements:**<br>
+    We can reduce the VRAM requirements using lower precision and setting the flag `--use_float_16`. 
+
+<br>
+
+**Finding Custom Edit Directions**<br>
+ - We provide some pre-computed directions in the assets [folder](assets/embeddings_sd_1.4).
+   To generate new edit directions, users can first generate two files containing a large number of sentences (~1000) and then run the command as shown below. 
+    ```
+      python src/make_edit_direction.py \
+        --file_source_sentences sentences/apple.txt \
+        --file_target_sentences sentences/orange.txt \
+        --output_folder assets/embeddings_sd_1.4
+    ```
+- After running the above command, you can set the flag `--task apple2orange` for the new edit.
+
+
+
+## Comparison
+Comparisons with different baselines, including, SDEdit + word swap, DDIM + word swap, and prompt-to-propmt. Our method successfully applies the edit, while preserving the structure of the input image. 
+<div>
+<p align="center">
+<img src='assets/comparison.jpg' align="center" width=900>
+</p>
+</div>

submodules/pix2pix-zero/environment.yml

@@ -0,0 +1,23 @@
+name: pix2pix-zero
+channels:
+  - pytorch
+  - nvidia
+  - defaults
+dependencies:
+  - pip
+  - pytorch-cuda=11.6
+  - torchvision
+  - pytorch
+  - pip:
+    - accelerate
+    - diffusers
+    - einops
+    - gradio
+    - ipython
+    - numpy
+    - opencv-python-headless
+    - pillow
+    - psutil
+    - tqdm
+    - transformers
+    - salesforce-lavis

submodules/pix2pix-zero/src/edit_real.py

@@ -0,0 +1,65 @@
+import os, pdb
+
+import argparse
+import numpy as np
+import torch
+import requests
+from PIL import Image
+
+from diffusers import DDIMScheduler
+from utils.ddim_inv import DDIMInversion
+from utils.edit_directions import construct_direction
+from utils.edit_pipeline import EditingPipeline
+
+
+if __name__=="__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--inversion', required=True)
+    parser.add_argument('--prompt', type=str, required=True)
+    parser.add_argument('--task_name', type=str, default='cat2dog')
+    parser.add_argument('--results_folder', type=str, default='output/test_cat')
+    parser.add_argument('--num_ddim_steps', type=int, default=50)
+    parser.add_argument('--model_path', type=str, default='CompVis/stable-diffusion-v1-4')
+    parser.add_argument('--xa_guidance', default=0.1, type=float)
+    parser.add_argument('--negative_guidance_scale', default=5.0, type=float)
+    parser.add_argument('--use_float_16', action='store_true')
+
+    args = parser.parse_args()
+
+    os.makedirs(os.path.join(args.results_folder, "edit"), exist_ok=True)
+    os.makedirs(os.path.join(args.results_folder, "reconstruction"), exist_ok=True)
+
+    if args.use_float_16:
+        torch_dtype = torch.float16
+    else:
+        torch_dtype = torch.float32
+
+    # if the inversion is a folder, the prompt should also be a folder
+    assert (os.path.isdir(args.inversion)==os.path.isdir(args.prompt)), "If the inversion is a folder, the prompt should also be a folder"
+    if os.path.isdir(args.inversion):
+        l_inv_paths = sorted(glob(os.path.join(args.inversion, "*.pt")))
+        l_bnames = [os.path.basename(x) for x in l_inv_paths]
+        l_prompt_paths = [os.path.join(args.prompt, x.replace(".pt",".txt")) for x in l_bnames]
+    else:
+        l_inv_paths = [args.inversion]
+        l_prompt_paths = [args.prompt]
+
+    # Make the editing pipeline
+    pipe = EditingPipeline.from_pretrained(args.model_path, torch_dtype=torch_dtype).to("cuda")
+    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+
+
+    for inv_path, prompt_path in zip(l_inv_paths, l_prompt_paths):
+        prompt_str = open(prompt_path).read().strip()
+        rec_pil, edit_pil = pipe(prompt_str,
+                num_inference_steps=args.num_ddim_steps,
+                x_in=torch.load(inv_path).unsqueeze(0),
+                edit_dir=construct_direction(args.task_name),
+                guidance_amount=args.xa_guidance,
+                guidance_scale=args.negative_guidance_scale,
+                negative_prompt=prompt_str # use the unedited prompt for the negative prompt
+        )
+        
+        bname = os.path.basename(args.inversion).split(".")[0]
+        edit_pil[0].save(os.path.join(args.results_folder, f"edit/{bname}.png"))
+        rec_pil[0].save(os.path.join(args.results_folder, f"reconstruction/{bname}.png"))

submodules/pix2pix-zero/src/edit_synthetic.py

@@ -0,0 +1,52 @@
+import os, pdb
+
+import argparse
+import numpy as np
+import torch
+import requests
+from PIL import Image
+
+from diffusers import DDIMScheduler
+from utils.edit_directions import construct_direction
+from utils.edit_pipeline import EditingPipeline
+
+
+if __name__=="__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--prompt_str', type=str, required=True)
+    parser.add_argument('--random_seed', default=0)
+    parser.add_argument('--task_name', type=str, default='cat2dog')
+    parser.add_argument('--results_folder', type=str, default='output/test_cat')
+    parser.add_argument('--num_ddim_steps', type=int, default=50)
+    parser.add_argument('--model_path', type=str, default='CompVis/stable-diffusion-v1-4')
+    parser.add_argument('--xa_guidance', default=0.15, type=float)
+    parser.add_argument('--negative_guidance_scale', default=5.0, type=float)
+    parser.add_argument('--use_float_16', action='store_true')
+    args = parser.parse_args()
+
+    os.makedirs(args.results_folder, exist_ok=True)
+
+    if args.use_float_16:
+        torch_dtype = torch.float16
+    else:
+        torch_dtype = torch.float32
+
+    # make the input noise map
+    torch.cuda.manual_seed(args.random_seed)
+    x = torch.randn((1,4,64,64), device="cuda")
+
+    # Make the editing pipeline
+    pipe = EditingPipeline.from_pretrained(args.model_path, torch_dtype=torch_dtype).to("cuda")
+    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+
+    rec_pil, edit_pil = pipe(args.prompt_str, 
+        num_inference_steps=args.num_ddim_steps,
+        x_in=x,
+        edit_dir=construct_direction(args.task_name),
+        guidance_amount=args.xa_guidance,
+        guidance_scale=args.negative_guidance_scale,
+        negative_prompt="" # use the empty string for the negative prompt
+    )
+    
+    edit_pil[0].save(os.path.join(args.results_folder, f"edit.png"))
+    rec_pil[0].save(os.path.join(args.results_folder, f"reconstruction.png"))

submodules/pix2pix-zero/src/inversion.py

@@ -0,0 +1,64 @@
+import os, pdb
+
+import argparse
+import numpy as np
+import torch
+import requests
+from PIL import Image
+
+from lavis.models import load_model_and_preprocess
+
+from utils.ddim_inv import DDIMInversion
+from utils.scheduler import DDIMInverseScheduler
+
+if __name__=="__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input_image', type=str, default='assets/test_images/cat_a.png')
+    parser.add_argument('--results_folder', type=str, default='output/test_cat')
+    parser.add_argument('--num_ddim_steps', type=int, default=50)
+    parser.add_argument('--model_path', type=str, default='CompVis/stable-diffusion-v1-4')
+    parser.add_argument('--use_float_16', action='store_true')
+    args = parser.parse_args()
+
+    # make the output folders
+    os.makedirs(os.path.join(args.results_folder, "inversion"), exist_ok=True)
+    os.makedirs(os.path.join(args.results_folder, "prompt"), exist_ok=True)
+
+    if args.use_float_16:
+        torch_dtype = torch.float16
+    else:
+        torch_dtype = torch.float32
+
+
+    # load the BLIP model
+    model_blip, vis_processors, _ = load_model_and_preprocess(name="blip_caption", model_type="base_coco", is_eval=True, device=torch.device("cuda"))
+    # make the DDIM inversion pipeline    
+    pipe = DDIMInversion.from_pretrained(args.model_path, torch_dtype=torch_dtype).to("cuda")
+    pipe.scheduler = DDIMInverseScheduler.from_config(pipe.scheduler.config)
+
+
+    # if the input is a folder, collect all the images as a list
+    if os.path.isdir(args.input_image):
+        l_img_paths = sorted(glob(os.path.join(args.input_image, "*.png")))
+    else:
+        l_img_paths = [args.input_image]
+
+
+    for img_path in l_img_paths:
+        bname = os.path.basename(args.input_image).split(".")[0]
+        img = Image.open(args.input_image).resize((512,512), Image.Resampling.LANCZOS)
+        # generate the caption
+        _image = vis_processors["eval"](img).unsqueeze(0).cuda()
+        prompt_str = model_blip.generate({"image": _image})[0]
+        x_inv, x_inv_image, x_dec_img = pipe(
+            prompt_str, 
+            guidance_scale=1,
+            num_inversion_steps=args.num_ddim_steps,
+            img=img,
+            torch_dtype=torch_dtype
+        )
+        # save the inversion
+        torch.save(x_inv[0], os.path.join(args.results_folder, f"inversion/{bname}.pt"))
+        # save the prompt string
+        with open(os.path.join(args.results_folder, f"prompt/{bname}.txt"), "w") as f:
+            f.write(prompt_str)

submodules/pix2pix-zero/src/make_edit_direction.py

@@ -0,0 +1,61 @@
+import os, pdb
+
+import argparse
+import numpy as np
+import torch
+import requests
+from PIL import Image
+
+from diffusers import DDIMScheduler
+from utils.edit_pipeline import EditingPipeline
+
+
+## convert sentences to sentence embeddings
+def load_sentence_embeddings(l_sentences, tokenizer, text_encoder, device="cuda"):
+    with torch.no_grad():
+        l_embeddings = []
+        for sent in l_sentences:
+            text_inputs = tokenizer(
+                    sent,
+                    padding="max_length",
+                    max_length=tokenizer.model_max_length,
+                    truncation=True,
+                    return_tensors="pt",
+                )
+            text_input_ids = text_inputs.input_ids
+            prompt_embeds = text_encoder(text_input_ids.to(device), attention_mask=None)[0]
+            l_embeddings.append(prompt_embeds)
+    return torch.concatenate(l_embeddings, dim=0).mean(dim=0).unsqueeze(0)
+
+
+if __name__=="__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--file_source_sentences', required=True)
+    parser.add_argument('--file_target_sentences', required=True)
+    parser.add_argument('--output_folder', required=True)
+    parser.add_argument('--model_path', type=str, default='CompVis/stable-diffusion-v1-4')
+    args = parser.parse_args()
+
+    # load the model
+    pipe = EditingPipeline.from_pretrained(args.model_path, torch_dtype=torch.float16).to("cuda")
+    bname_src = os.path.basename(args.file_source_sentences).strip(".txt")
+    outf_src = os.path.join(args.output_folder, bname_src+".pt")
+    if os.path.exists(outf_src):
+        print(f"Skipping source file {outf_src} as it already exists")
+    else:
+        with open(args.file_source_sentences, "r") as f:
+            l_sents = [x.strip() for x in f.readlines()]
+        mean_emb = load_sentence_embeddings(l_sents, pipe.tokenizer, pipe.text_encoder, device="cuda")
+        print(mean_emb.shape)
+        torch.save(mean_emb, outf_src)
+
+    bname_tgt = os.path.basename(args.file_target_sentences).strip(".txt")
+    outf_tgt = os.path.join(args.output_folder, bname_tgt+".pt")
+    if os.path.exists(outf_tgt):
+        print(f"Skipping target file {outf_tgt} as it already exists")
+    else:
+        with open(args.file_target_sentences, "r") as f:
+            l_sents = [x.strip() for x in f.readlines()]
+        mean_emb = load_sentence_embeddings(l_sents, pipe.tokenizer, pipe.text_encoder, device="cuda")
+        print(mean_emb.shape)
+        torch.save(mean_emb, outf_tgt)

submodules/pix2pix-zero/src/utils/base_pipeline.py

@@ -0,0 +1,322 @@
+
+import torch
+import inspect
+from packaging import version
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
+from diffusers import DiffusionPipeline
+from diffusers.models import AutoencoderKL, UNet2DConditionModel
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import deprecate, is_accelerate_available, logging, randn_tensor, replace_example_docstring
+from diffusers import StableDiffusionPipeline
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+
+
+
+class BasePipeline(DiffusionPipeline):
+    _optional_components = ["safety_checker", "feature_extractor"]
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        safety_checker: StableDiffusionSafetyChecker,
+        feature_extractor: CLIPFeatureExtractor,
+        requires_safety_checker: bool = True,
+    ):
+        super().__init__()
+
+        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
+                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
+                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
+                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
+                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
+                " file"
+            )
+            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["steps_offset"] = 1
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
+                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
+                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
+                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
+                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
+            )
+            deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["clip_sample"] = False
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        if safety_checker is None and requires_safety_checker:
+            logger.warning(
+                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
+                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
+                " results in services or applications open to the public. Both the diffusers team and Hugging Face"
+                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
+                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
+                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
+            )
+
+        if safety_checker is not None and feature_extractor is None:
+            raise ValueError(
+                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
+                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
+            )
+
+        is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(
+            version.parse(unet.config._diffusers_version).base_version
+        ) < version.parse("0.9.0.dev0")
+        is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64
+        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:
+            deprecation_message = (
+                "The configuration file of the unet has set the default `sample_size` to smaller than"
+                " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"
+                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
+                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
+                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
+                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
+                " in the config might lead to incorrect results in future versions. If you have downloaded this"
+                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
+                " the `unet/config.json` file"
+            )
+            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(unet.config)
+            new_config["sample_size"] = 64
+            unet._internal_dict = FrozenDict(new_config)
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.register_to_config(requires_safety_checker=requires_safety_checker)
+
+    @property
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
+    def _execution_device(self):
+        r"""
+        Returns the device on which the pipeline's models will be executed. After calling
+        `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
+        hooks.
+        """
+        if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"):
+            return self.device
+        for module in self.unet.modules():
+            if (
+                hasattr(module, "_hf_hook")
+                and hasattr(module._hf_hook, "execution_device")
+                and module._hf_hook.execution_device is not None
+            ):
+                return torch.device(module._hf_hook.execution_device)
+        return self.device
+
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+             prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            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. 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`).
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+        """
+        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]
+
+        if prompt_embeds is None:
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=self.tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.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 = self.tokenizer.batch_decode(
+                    untruncated_ids[:, self.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" {self.tokenizer.model_max_length} tokens: {removed_text}"
+                )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = text_inputs.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            prompt_embeds = self.text_encoder(
+                text_input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            prompt_embeds = prompt_embeds[0]
+
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.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)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif isinstance(negative_prompt, str):
+                uncond_tokens = [negative_prompt]
+            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
+
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = uncond_input.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.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)
+
+            # For classifier free guidance, we need to do two forward passes.
+            # Here we concatenate the unconditional and text embeddings into a single batch
+            # to avoid doing two forward passes
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        return prompt_embeds
+
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents
+    def decode_latents(self, latents):
+        latents = 1 / 0.18215 * latents
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+        image = image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+        
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    # 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
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker
+    def run_safety_checker(self, image, device, dtype):
+        if self.safety_checker is not None:
+            safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device)
+            image, has_nsfw_concept = self.safety_checker(
+                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
+            )
+        else:
+            has_nsfw_concept = None
+        return image, has_nsfw_concept
+

submodules/pix2pix-zero/src/utils/cross_attention.py

@@ -0,0 +1,57 @@
+import torch
+from diffusers.models.attention import CrossAttention
+
+class MyCrossAttnProcessor:
+    def __call__(self, attn: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None):
+        batch_size, sequence_length, _ = hidden_states.shape
+        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length)
+
+        query = attn.to_q(hidden_states)
+
+        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        # new bookkeeping to save the attn probs
+        attn.attn_probs = attention_probs
+
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        return hidden_states
+
+
+"""
+A function that prepares a U-Net model for training by enabling gradient computation 
+for a specified set of parameters and setting the forward pass to be performed by a 
+custom cross attention processor.
+
+Parameters:
+unet: A U-Net model.
+
+Returns:
+unet: The prepared U-Net model.
+"""
+def prep_unet(unet):
+    # set the gradients for XA maps to be true
+    for name, params in unet.named_parameters():
+        if 'attn2' in name:
+            params.requires_grad = True
+        else:
+            params.requires_grad = False
+    # replace the fwd function
+    for name, module in unet.named_modules():
+        module_name = type(module).__name__
+        if module_name == "CrossAttention":
+            module.set_processor(MyCrossAttnProcessor())
+    return unet

submodules/pix2pix-zero/src/utils/ddim_inv.py

@@ -0,0 +1,140 @@
+import sys
+import numpy as np
+import torch
+import torch.nn.functional as F
+from random import randrange
+from typing import Any, Callable, Dict, List, Optional, Union, Tuple
+from diffusers import DDIMScheduler
+from diffusers.schedulers.scheduling_ddim import DDIMSchedulerOutput
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+sys.path.insert(0, "src/utils")
+from base_pipeline import BasePipeline
+from cross_attention import prep_unet
+
+
+class DDIMInversion(BasePipeline):
+
+    def auto_corr_loss(self, x, random_shift=True):
+        B,C,H,W = x.shape
+        assert B==1
+        x = x.squeeze(0)
+        # x must be shape [C,H,W] now
+        reg_loss = 0.0
+        for ch_idx in range(x.shape[0]):
+            noise = x[ch_idx][None, None,:,:]
+            while True:
+                if random_shift: roll_amount = randrange(noise.shape[2]//2)
+                else: roll_amount = 1
+                reg_loss += (noise*torch.roll(noise, shifts=roll_amount, dims=2)).mean()**2
+                reg_loss += (noise*torch.roll(noise, shifts=roll_amount, dims=3)).mean()**2
+                if noise.shape[2] <= 8:
+                    break
+                noise = F.avg_pool2d(noise, kernel_size=2)
+        return reg_loss
+    
+    def kl_divergence(self, x):
+        _mu = x.mean()
+        _var = x.var()
+        return _var + _mu**2 - 1 - torch.log(_var+1e-7)
+
+
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        num_inversion_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        img=None, # the input image as a PIL image
+        torch_dtype=torch.float32,
+
+        # inversion regularization parameters
+        lambda_ac: float = 20.0,
+        lambda_kl: float = 20.0,
+        num_reg_steps: int = 5,
+        num_ac_rolls: int = 5,
+    ):
+        
+        # 0. modify the unet to be useful :D
+        self.unet = prep_unet(self.unet)
+
+        # set the scheduler to be the Inverse DDIM scheduler
+        # self.scheduler = MyDDIMScheduler.from_config(self.scheduler.config)
+
+        device = self._execution_device
+        do_classifier_free_guidance = guidance_scale > 1.0
+        self.scheduler.set_timesteps(num_inversion_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # Encode the input image with the first stage model
+        x0 = np.array(img)/255
+        x0 = torch.from_numpy(x0).type(torch_dtype).permute(2, 0, 1).unsqueeze(dim=0).repeat(1, 1, 1, 1).cuda()
+        x0 = (x0 - 0.5) * 2.
+        with torch.no_grad():
+            x0_enc = self.vae.encode(x0).latent_dist.sample().to(device, torch_dtype)
+        latents = x0_enc = 0.18215 * x0_enc
+
+        # Decode and return the image
+        with torch.no_grad():
+            x0_dec = self.decode_latents(x0_enc.detach())
+        image_x0_dec = self.numpy_to_pil(x0_dec)
+
+        with torch.no_grad():
+            prompt_embeds = self._encode_prompt(prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt).to(device)
+        extra_step_kwargs = self.prepare_extra_step_kwargs(None, eta)
+
+        # Do the inversion
+        num_warmup_steps = len(timesteps) - num_inversion_steps * self.scheduler.order # should be 0?
+        with self.progress_bar(total=num_inversion_steps) as progress_bar:
+            for i, t in enumerate(timesteps.flip(0)[1:-1]):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                # predict the noise residual
+                with torch.no_grad():
+                    noise_pred = self.unet(latent_model_input,t,encoder_hidden_states=prompt_embeds,cross_attention_kwargs=cross_attention_kwargs,).sample
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # regularization of the noise prediction
+                e_t = noise_pred
+                for _outer in range(num_reg_steps):
+                    if lambda_ac>0:
+                        for _inner in range(num_ac_rolls):
+                            _var = torch.autograd.Variable(e_t.detach().clone(), requires_grad=True)
+                            l_ac = self.auto_corr_loss(_var)
+                            l_ac.backward()
+                            _grad = _var.grad.detach()/num_ac_rolls
+                            e_t = e_t - lambda_ac*_grad
+                    if lambda_kl>0:
+                        _var = torch.autograd.Variable(e_t.detach().clone(), requires_grad=True)
+                        l_kld = self.kl_divergence(_var)
+                        l_kld.backward()
+                        _grad = _var.grad.detach()
+                        e_t = e_t - lambda_kl*_grad
+                    e_t = e_t.detach()
+                noise_pred = e_t
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, reverse=True, **extra_step_kwargs).prev_sample
+
+                # 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()
+        
+
+        x_inv = latents.detach().clone()
+        # reconstruct the image
+
+        # 8. Post-processing
+        image = self.decode_latents(latents.detach())
+        image = self.numpy_to_pil(image)
+        return x_inv, image, image_x0_dec

submodules/pix2pix-zero/src/utils/edit_directions.py

@@ -0,0 +1,29 @@
+import os
+import torch
+
+
+"""
+This function takes in a task name and returns the direction in the embedding space that transforms class A to class B for the given task.
+
+Parameters:
+task_name (str): name of the task for which direction is to be constructed.
+
+Returns:
+torch.Tensor: A tensor representing the direction in the embedding space that transforms class A to class B.
+
+Examples:
+>>> construct_direction("cat2dog")
+"""
+def construct_direction(task_name):
+    if task_name=="cat2dog":    
+        emb_dir = f"assets/embeddings_sd_1.4"
+        embs_a = torch.load(os.path.join(emb_dir, f"cat.pt"))
+        embs_b = torch.load(os.path.join(emb_dir, f"dog.pt"))
+        return (embs_b.mean(0)-embs_a.mean(0)).unsqueeze(0)
+    elif task_name=="dog2cat":    
+        emb_dir = f"assets/embeddings_sd_1.4"
+        embs_a = torch.load(os.path.join(emb_dir, f"dog.pt"))
+        embs_b = torch.load(os.path.join(emb_dir, f"cat.pt"))
+        return (embs_b.mean(0)-embs_a.mean(0)).unsqueeze(0)
+    else:
+        raise NotImplementedError

submodules/pix2pix-zero/src/utils/edit_pipeline.py

@@ -0,0 +1,179 @@
+import pdb, sys
+
+import numpy as np
+import torch
+from typing import Any, Callable, Dict, List, Optional, Union
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+sys.path.insert(0, "src/utils")
+from base_pipeline import BasePipeline
+from cross_attention import prep_unet
+
+
+class EditingPipeline(BasePipeline):
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+
+        # pix2pix parameters
+        guidance_amount=0.1,
+        edit_dir=None,
+        x_in=None,
+        only_sample=False,
+
+    ):
+
+        x_in.to(dtype=self.unet.dtype, device=self._execution_device)
+
+        # 0. modify the unet to be useful :D
+        self.unet = prep_unet(self.unet)
+        
+        # 1. setup all caching objects
+        d_ref_t2attn = {} # reference cross attention maps
+        
+        # 2. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # TODO: add the input checker function
+        # self.check_inputs( prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds )
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        do_classifier_free_guidance = guidance_scale > 1.0
+        x_in = x_in.to(dtype=self.unet.dtype, device=self._execution_device)
+        # 3. Encode input prompt = 2x77x1024
+        prompt_embeds = self._encode_prompt( prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds,)
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.in_channels
+        
+        # randomly sample a latent code if not provided
+        latents = self.prepare_latents(batch_size * num_images_per_prompt, num_channels_latents, height, width, prompt_embeds.dtype, device, generator, x_in,)
+        
+        latents_init = latents.clone()
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. First Denoising loop for getting the reference cross attention maps
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        with torch.no_grad():
+            with self.progress_bar(total=num_inference_steps) as progress_bar:
+                for i, t in enumerate(timesteps):
+                    # expand the latents if we are doing classifier free guidance
+                    latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                    latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                    # predict the noise residual
+                    noise_pred = self.unet(latent_model_input,t,encoder_hidden_states=prompt_embeds,cross_attention_kwargs=cross_attention_kwargs,).sample
+
+                    # add the cross attention map to the dictionary
+                    d_ref_t2attn[t.item()] = {}
+                    for name, module in self.unet.named_modules():
+                        module_name = type(module).__name__
+                        if module_name == "CrossAttention" and 'attn2' in name:
+                            attn_mask = module.attn_probs # size is num_channel,s*s,77
+                            d_ref_t2attn[t.item()][name] = attn_mask.detach().cpu()
+
+                    # perform guidance
+                    if do_classifier_free_guidance:
+                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                    # compute the previous noisy sample x_t -> x_t-1
+                    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+                    # 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()
+
+        # make the reference image (reconstruction)
+        image_rec = self.numpy_to_pil(self.decode_latents(latents.detach()))
+
+        if only_sample:
+            return image_rec
+
+
+        prompt_embeds_edit = prompt_embeds.clone()
+        #add the edit only to the second prompt, idx 0 is the negative prompt
+        prompt_embeds_edit[1:2] += edit_dir
+        
+        latents = latents_init
+        # Second denoising loop for editing the text prompt
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                x_in = latent_model_input.detach().clone()
+                x_in.requires_grad = True
+                
+                opt = torch.optim.SGD([x_in], lr=guidance_amount)
+
+                # predict the noise residual
+                noise_pred = self.unet(x_in,t,encoder_hidden_states=prompt_embeds_edit.detach(),cross_attention_kwargs=cross_attention_kwargs,).sample
+
+                loss = 0.0
+                for name, module in self.unet.named_modules():
+                    module_name = type(module).__name__
+                    if module_name == "CrossAttention" and 'attn2' in name:
+                        curr = module.attn_probs # size is num_channel,s*s,77
+                        ref = d_ref_t2attn[t.item()][name].detach().cuda()
+                        loss += ((curr-ref)**2).sum((1,2)).mean(0)
+                loss.backward(retain_graph=False)
+                opt.step()
+
+                # recompute the noise
+                with torch.no_grad():
+                    noise_pred = self.unet(x_in.detach(),t,encoder_hidden_states=prompt_embeds_edit,cross_attention_kwargs=cross_attention_kwargs,).sample
+                
+                latents = x_in.detach().chunk(2)[0]
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+                # 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()
+
+
+        # 8. Post-processing
+        image = self.decode_latents(latents.detach())
+
+        # 9. Run safety checker
+        image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
+
+        # 10. Convert to PIL
+        image_edit = self.numpy_to_pil(image)
+
+
+        return image_rec, image_edit

submodules/pix2pix-zero/src/utils/scheduler.py

@@ -0,0 +1,289 @@
+# Copyright 2022 Stanford University Team and 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.
+
+# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
+# and https://github.com/hojonathanho/diffusion
+import os, sys, pdb
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput, randn_tensor
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
+
+
+@dataclass
+# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM
+class DDIMSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's step function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            The predicted denoised sample (x_{0}) based on the model output from the current timestep.
+            `pred_original_sample` can be used to preview progress or for guidance.
+    """
+
+    prev_sample: torch.FloatTensor
+    pred_original_sample: Optional[torch.FloatTensor] = None
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, max_beta=0.999) -> torch.Tensor:
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
+    (1-beta) over time from t = [0,1].
+
+    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
+    to that part of the diffusion process.
+
+
+    Args:
+        num_diffusion_timesteps (`int`): the number of betas to produce.
+        max_beta (`float`): the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+
+    Returns:
+        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
+    """
+
+    def alpha_bar(time_step):
+        return math.cos((time_step + 0.008) / 1.008 * math.pi / 2) ** 2
+
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return torch.tensor(betas)
+
+
+class DDIMInverseScheduler(SchedulerMixin, ConfigMixin):
+    """
+    Denoising diffusion implicit models is a scheduler that extends the denoising procedure introduced in denoising
+    diffusion probabilistic models (DDPMs) with non-Markovian guidance.
+
+    [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__`
+    function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`.
+    [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and
+    [`~SchedulerMixin.from_pretrained`] functions.
+
+    For more details, see the original paper: https://arxiv.org/abs/2010.02502
+
+    Args:
+        num_train_timesteps (`int`): number of diffusion steps used to train the model.
+        beta_start (`float`): the starting `beta` value of inference.
+        beta_end (`float`): the final `beta` value.
+        beta_schedule (`str`):
+            the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
+            `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
+        trained_betas (`np.ndarray`, optional):
+            option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc.
+        clip_sample (`bool`, default `True`):
+            option to clip predicted sample between -1 and 1 for numerical stability.
+        set_alpha_to_one (`bool`, default `True`):
+            each diffusion step uses the value of alphas product at that step and at the previous one. For the final
+            step there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`,
+            otherwise it uses the value of alpha at step 0.
+        steps_offset (`int`, default `0`):
+            an offset added to the inference steps. You can use a combination of `offset=1` and
+            `set_alpha_to_one=False`, to make the last step use step 0 for the previous alpha product, as done in
+            stable diffusion.
+        prediction_type (`str`, default `epsilon`, optional):
+            prediction type of the scheduler function, one of `epsilon` (predicting the noise of the diffusion
+            process), `sample` (directly predicting the noisy sample`) or `v_prediction` (see section 2.4
+            https://imagen.research.google/video/paper.pdf)
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.0001,
+        beta_end: float = 0.02,
+        beta_schedule: str = "linear",
+        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        clip_sample: bool = True,
+        set_alpha_to_one: bool = True,
+        steps_offset: int = 0,
+        prediction_type: str = "epsilon",
+    ):
+        if trained_betas is not None:
+            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
+        elif beta_schedule == "linear":
+            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+        elif beta_schedule == "scaled_linear":
+            # this schedule is very specific to the latent diffusion model.
+            self.betas = (
+                torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
+            )
+        elif beta_schedule == "squaredcos_cap_v2":
+            # Glide cosine schedule
+            self.betas = betas_for_alpha_bar(num_train_timesteps)
+        else:
+            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")
+
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+
+        # At every step in ddim, we are looking into the previous alphas_cumprod
+        # For the final step, there is no previous alphas_cumprod because we are already at 0
+        # `set_alpha_to_one` decides whether we set this parameter simply to one or
+        # whether we use the final alpha of the "non-previous" one.
+        self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
+
+        # standard deviation of the initial noise distribution
+        self.init_noise_sigma = 1.0
+
+        # setable values
+        self.num_inference_steps = None
+        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
+
+    def scale_model_input(self, sample: torch.FloatTensor, timestep: Optional[int] = None) -> torch.FloatTensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.FloatTensor`): input sample
+            timestep (`int`, optional): current timestep
+
+        Returns:
+            `torch.FloatTensor`: scaled input sample
+        """
+        return sample
+
+    def _get_variance(self, timestep, prev_timestep):
+        alpha_prod_t = self.alphas_cumprod[timestep]
+        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
+        beta_prod_t = 1 - alpha_prod_t
+        beta_prod_t_prev = 1 - alpha_prod_t_prev
+
+        variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
+
+        return variance
+
+    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
+        """
+        Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
+
+        Args:
+            num_inference_steps (`int`):
+                the number of diffusion steps used when generating samples with a pre-trained model.
+        """
+
+        if num_inference_steps > self.config.num_train_timesteps:
+            raise ValueError(
+                f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"
+                f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
+                f" maximal {self.config.num_train_timesteps} timesteps."
+            )
+
+        self.num_inference_steps = num_inference_steps
+        step_ratio = self.config.num_train_timesteps // self.num_inference_steps
+        # creates integer timesteps by multiplying by ratio
+        # casting to int to avoid issues when num_inference_step is power of 3
+        timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
+        self.timesteps = torch.from_numpy(timesteps).to(device)
+        self.timesteps += self.config.steps_offset
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        sample: torch.FloatTensor,
+        eta: float = 0.0,
+        use_clipped_model_output: bool = False,
+        generator=None,
+        variance_noise: Optional[torch.FloatTensor] = None,
+        return_dict: bool = True,
+        reverse=False
+    ) -> Union[DDIMSchedulerOutput, Tuple]:
+        
+
+        e_t = model_output
+
+        x = sample
+        prev_timestep = timestep + self.config.num_train_timesteps // self.num_inference_steps
+        # print(timestep, prev_timestep)
+        a_t = alpha_prod_t = self.alphas_cumprod[timestep-1]
+        a_prev = alpha_t_prev = self.alphas_cumprod[prev_timestep-1] if prev_timestep >= 0 else self.final_alpha_cumprod
+        beta_prod_t = 1 - alpha_prod_t
+
+        pred_x0 = (x - (1-a_t)**0.5 * e_t) / a_t.sqrt()
+        # direction pointing to x_t
+        dir_xt = (1. - a_prev).sqrt() * e_t
+        x = a_prev.sqrt()*pred_x0 + dir_xt
+        if not return_dict:
+            return (x,)
+        return DDIMSchedulerOutput(prev_sample=x, pred_original_sample=pred_x0)
+
+
+
+
+
+    def add_noise(
+        self,
+        original_samples: torch.FloatTensor,
+        noise: torch.FloatTensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.FloatTensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
+        self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
+        timesteps = timesteps.to(original_samples.device)
+
+        sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+        return noisy_samples
+
+    def get_velocity(
+        self, sample: torch.FloatTensor, noise: torch.FloatTensor, timesteps: torch.IntTensor
+    ) -> torch.FloatTensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as sample
+        self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device, dtype=sample.dtype)
+        timesteps = timesteps.to(sample.device)
+
+        sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(sample.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
+        return velocity
+
+    def __len__(self):
+        return self.config.num_train_timesteps

utils/direction_utils.py

@@ -0,0 +1,79 @@
+import os, sys, pdb
+
+import torch, torchvision
+from huggingface_hub import hf_hub_url, cached_download, hf_hub_download, HfApi
+import joblib
+from pathlib import Path
+
+import json
+import requests
+import random
+
+
+"""
+Returns the list of directions currently available in the HF library
+"""
+def get_all_directions_names():
+    hf_api = HfApi()
+    info = hf_api.list_models(author="pix2pix-zero-library")
+    l_model_ids = [m.modelId for m in info]
+    l_model_ids = [m for m in l_model_ids if "_sd14" in m]
+    l_edit_names = [m.split("/")[-1] for m in l_model_ids]
+    # l_edit_names = [m for m in l_edit_names if "_sd14" in m]
+
+    # pdb.set_trace()
+    l_desc = [hf_hub_download(repo_id=m_id, filename="short_description.txt") for m_id in l_model_ids]
+    d_name2desc = {k: open(m).read() for k,m in zip(l_edit_names, l_desc)}
+
+    return d_name2desc
+
+
+def get_emb(dir_name):
+    REPO_ID = f"pix2pix-zero-library/{dir_name.replace('.pt','')}"
+    if "_sd14" not in REPO_ID: REPO_ID += "_sd14"
+    FILENAME = dir_name
+    if "_sd14" not in FILENAME: FILENAME += "_sd14"
+    if ".pt" not in FILENAME: FILENAME += ".pt"
+    ret = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
+    return torch.load(ret)
+
+
+def generate_image_prompts_with_templates(word):
+    prompts = []
+    adjectives = ['majestic', 'cute', 'colorful', 'ferocious', 'elegant', 'graceful', 'slimy', 'adorable', 'scary', 'fuzzy', 'tiny', 'gigantic', 'brave', 'fierce', 'mysterious', 'curious', 'fascinating', 'charming', 'gleaming', 'rare']
+    verbs = ['strolling', 'jumping', 'lounging', 'flying', 'sleeping', 'eating', 'playing', 'working', 'gazing', 'standing']
+    adverbs = ['gracefully', 'playfully', 'elegantly', 'fiercely', 'curiously', 'fascinatingly', 'charmingly', 'gently', 'slowly', 'quickly', 'awkwardly', 'carelessly', 'cautiously', 'innocently', 'powerfully', 'grumpily', 'mysteriously']
+    backgrounds = ['a sunny beach', 'a bustling city', 'a quiet forest', 'a cozy living room', 'a futuristic space station', 'a medieval castle', 'an enchanted forest', 'a misty graveyard', 'a snowy mountain peak', 'a crowded market']
+
+    sentence_structures = {
+        "subject verb background": lambda word, bg, verb, adj, adv: f"A {word} {verb} {bg}.",
+        "background subject verb": lambda word, bg, verb, adj, adv: f"{bg}, a {word} is {verb}.",
+        "adjective subject verb background": lambda word, bg, verb, adj, adv: f"A {adj} {word} is {verb} {bg}.",
+        "subject verb adverb background": lambda word, bg, verb, adj, adv: f"A {word} is {verb} {adv} {bg}.",
+        "adverb subject verb background": lambda word, bg, verb, adj, adv: f"{adv.capitalize()}, a {word} is {verb} {bg}.",
+        "background adjective subject verb": lambda word, bg, verb, adj, adv: f"{bg}, there is a {adj} {word} {verb}.",
+        "subject verb adjective background": lambda word, bg, verb, adj, adv: f"A {word} {verb} {adj} {bg}.",
+        "adjective subject verb": lambda word, bg, verb, adj, adv: f"A {adj} {word} is {verb}.",
+        "subject adjective verb background": lambda word, bg, verb, adj, adv: f"A {word} is {adj} and {verb} {bg}.",
+    }
+
+    sentences = []
+    for bg in backgrounds:
+        for verb in verbs:
+            for adj in adjectives:
+                adv = random.choice(adverbs)
+                sentence = f"A {adv} {adj} {word} {verb} on {bg}."
+                sentence_structure = random.choice(list(sentence_structures.keys()))
+                sentence = sentence_structures[sentence_structure](word, bg, verb, adj, adv)
+                sentences.append(sentence)
+    return sentences
+
+
+
+
+if __name__=="__main__":
+    print(get_all_directions_names())
+    # print(get_emb("dog_sd14.pt").shape)
+    # print(get_emb("dog").shape)
+    # print(generate_image_prompts("dog")[0:5])
+

utils/generate_synthetic.py

@@ -0,0 +1,316 @@
+import os, sys, time, re, pdb
+import torch, torchvision
+import numpy
+from PIL import Image
+import hashlib
+from tqdm import tqdm
+import openai
+from utils.direction_utils import *
+
+p = "submodules/pix2pix-zero/src/utils"
+if p not in sys.path:
+    sys.path.append(p)
+from diffusers import DDIMScheduler
+from edit_directions import construct_direction
+from edit_pipeline import EditingPipeline
+from ddim_inv import DDIMInversion
+from scheduler import DDIMInverseScheduler
+from lavis.models import load_model_and_preprocess
+from transformers import T5Tokenizer, AutoTokenizer, T5ForConditionalGeneration, BloomForCausalLM
+
+
+
+def load_sentence_embeddings(l_sentences, tokenizer, text_encoder, device="cuda"):
+    with torch.no_grad():
+        l_embeddings = []
+        for sent in tqdm(l_sentences):
+            text_inputs = tokenizer(
+                    sent,
+                    padding="max_length",
+                    max_length=tokenizer.model_max_length,
+                    truncation=True,
+                    return_tensors="pt",
+                )
+            text_input_ids = text_inputs.input_ids
+            prompt_embeds = text_encoder(text_input_ids.to(device), attention_mask=None)[0]
+            l_embeddings.append(prompt_embeds)
+    return torch.concatenate(l_embeddings, dim=0).mean(dim=0).unsqueeze(0)
+
+
+
+def launch_generate_sample(prompt, seed, negative_scale, num_ddim):
+    os.makedirs("tmp", exist_ok=True)
+    # do the editing
+    edit_pipe = EditingPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32).to("cuda")
+    edit_pipe.scheduler = DDIMScheduler.from_config(edit_pipe.scheduler.config)
+
+    # set the random seed and sample the input noise map
+    torch.cuda.manual_seed(int(seed))
+    z = torch.randn((1,4,64,64), device="cuda")
+
+    z_hashname = hashlib.sha256(z.cpu().numpy().tobytes()).hexdigest()
+    z_inv_fname = f"tmp/{z_hashname}_ddim_{num_ddim}_inv.pt"
+    torch.save(z, z_inv_fname)
+
+    rec_pil = edit_pipe(prompt, 
+        num_inference_steps=num_ddim, x_in=z,
+        only_sample=True, # this flag will only generate the sampled image, not the edited image
+        guidance_scale=negative_scale,
+        negative_prompt="" # use the empty string for the negative prompt
+    )
+    # print(rec_pil)
+    del edit_pipe
+    torch.cuda.empty_cache()
+
+    return rec_pil[0], z_inv_fname
+
+
+
+def clean_l_sentences(ls):
+    s = [re.sub('\d', '', x) for x in ls]
+    s = [x.replace(".","").replace("-","").replace(")","").strip() for x in s]
+    return s
+
+
+
+def gpt3_compute_word2sentences(task_type, word, num=100):
+    l_sentences = [] 
+    if task_type=="object":
+        template_prompt = f"Provide many captions for images containing {word}."
+    elif task_type=="style":
+        template_prompt = f"Provide many captions for images that are in the {word} style."
+    while True:
+        ret = openai.Completion.create(
+            model="text-davinci-002",
+            prompt=template_prompt,
+            max_tokens=1000,
+            temperature=1.0)
+        raw_return = ret.choices[0].text
+        for line in raw_return.split("\n"):
+            line = line.strip()
+            if len(line)>10:
+                skip=False 
+                for subword in word.split(" "):
+                    if subword not in line: skip=True
+                if not skip: l_sentences.append(line)
+                else:
+                    l_sentences.append(line+f", {word}")
+        time.sleep(0.05)
+        print(len(l_sentences))
+        if len(l_sentences)>=num:
+            break
+    l_sentences = clean_l_sentences(l_sentences)
+    return l_sentences
+
+
+def flant5xl_compute_word2sentences(word, num=100):
+    text_input = f"Provide a caption for images containing a {word}. The captions should be in English and should be no longer than 150 characters."
+    
+    l_sentences = []
+    tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-xl")
+    model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-xl", device_map="auto", torch_dtype=torch.float16)
+    input_ids = tokenizer(text_input, return_tensors="pt").input_ids.to("cuda")
+    input_length = input_ids.shape[1]
+    while True:
+        outputs = model.generate(input_ids,temperature=0.9, num_return_sequences=16, do_sample=True, max_length=128)
+        output = tokenizer.batch_decode(outputs[:, input_length:], skip_special_tokens=True)
+        for line in output:
+            line = line.strip()
+            skip=False 
+            for subword in word.split(" "):
+                if subword not in line: skip=True
+            if not skip: l_sentences.append(line)
+            else: l_sentences.append(line+f", {word}")
+        print(len(l_sentences))
+        if len(l_sentences)>=num:
+            break
+    l_sentences = clean_l_sentences(l_sentences)
+
+    del model
+    del tokenizer
+    torch.cuda.empty_cache()
+
+    return l_sentences
+
+def bloomz_compute_sentences(word, num=100):
+    l_sentences = []
+    tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-7b1")
+    model = BloomForCausalLM.from_pretrained("bigscience/bloomz-7b1", device_map="auto", torch_dtype=torch.float16)
+    input_text = f"Provide a caption for images containing a {word}. The captions should be in English and should be no longer than 150 characters. Caption:"
+    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
+    input_length = input_ids.shape[1]
+    t = 0.95
+    eta = 1e-5
+    min_length = 15
+
+    while True:
+        try:
+            outputs = model.generate(input_ids,temperature=t, num_return_sequences=16, do_sample=True, max_length=128, min_length=min_length, eta_cutoff=eta)
+            output = tokenizer.batch_decode(outputs[:, input_length:], skip_special_tokens=True)
+        except:
+            continue
+        for line in output:
+            line = line.strip()
+            skip=False 
+            for subword in word.split(" "):
+                if subword not in line: skip=True
+            if not skip: l_sentences.append(line)
+            else: l_sentences.append(line+f", {word}")
+        print(len(l_sentences))
+        if len(l_sentences)>=num:
+            break
+    l_sentences = clean_l_sentences(l_sentences)
+    del model
+    del tokenizer
+    torch.cuda.empty_cache()
+
+    return l_sentences
+
+
+
+def make_custom_dir(description, sent_type, api_key, org_key, l_custom_sentences):
+    if sent_type=="fixed-template":
+        l_sentences = generate_image_prompts_with_templates(description)
+    elif "GPT3" in sent_type:
+        import openai
+        openai.organization = org_key
+        openai.api_key = api_key
+        _=openai.Model.retrieve("text-davinci-002")
+        l_sentences = gpt3_compute_word2sentences("object", description, num=1000)
+    
+    elif "flan-t5-xl" in sent_type:
+        l_sentences = flant5xl_compute_word2sentences(description, num=1000)
+        # save the sentences to file
+        with open(f"tmp/flant5xl_sentences_{description}.txt", "w") as f:
+            for line in l_sentences:
+                f.write(line+"\n")
+    elif "BLOOMZ-7B" in sent_type:
+        l_sentences = bloomz_compute_sentences(description, num=1000)
+        # save the sentences to file
+        with open(f"tmp/bloomz_sentences_{description}.txt", "w") as f:
+            for line in l_sentences:
+                f.write(line+"\n")
+    
+    elif sent_type=="custom sentences":
+        l_sentences = l_custom_sentences.split("\n")
+        print(f"length of new sentence is {len(l_sentences)}")
+
+    pipe = EditingPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32).to("cuda")
+    emb = load_sentence_embeddings(l_sentences, pipe.tokenizer, pipe.text_encoder, device="cuda")
+    del pipe
+    torch.cuda.empty_cache()
+    return emb
+
+
+def launch_main(img_in_real, img_in_synth, src, src_custom, dest, dest_custom, num_ddim, xa_guidance, edit_mul, fpath_z_gen, gen_prompt, sent_type_src, sent_type_dest, api_key, org_key, custom_sentences_src, custom_sentences_dest):
+    d_name2desc = get_all_directions_names()
+    d_desc2name = {v:k for k,v in d_name2desc.items()}
+    os.makedirs("tmp", exist_ok=True)
+
+    # generate custom direction first
+    if src=="make your own!":
+        outf_name = f"tmp/template_emb_{src_custom}_{sent_type_src}.pt"
+        if not os.path.exists(outf_name):
+            src_emb = make_custom_dir(src_custom, sent_type_src, api_key, org_key, custom_sentences_src)
+            torch.save(src_emb, outf_name)
+        else:
+            src_emb = torch.load(outf_name)
+    else:
+        src_emb = get_emb(d_desc2name[src])
+    
+    if dest=="make your own!":
+        outf_name = f"tmp/template_emb_{dest_custom}_{sent_type_dest}.pt"
+        if not os.path.exists(outf_name):
+            dest_emb = make_custom_dir(dest_custom, sent_type_dest, api_key, org_key, custom_sentences_dest)
+            torch.save(dest_emb, outf_name)
+        else:
+            dest_emb = torch.load(outf_name)
+    else:
+        dest_emb = get_emb(d_desc2name[dest])
+    text_dir = (dest_emb.cuda() - src_emb.cuda())*edit_mul
+
+
+
+    if img_in_real is not None and img_in_synth is None:
+        print("using real image")
+        # resize the image so that the longer side is 512
+        width, height = img_in_real.size
+        if width > height: scale_factor = 512 / width
+        else: scale_factor = 512 / height
+        new_size = (int(width * scale_factor), int(height * scale_factor))
+        img_in_real = img_in_real.resize(new_size, Image.Resampling.LANCZOS)
+        hash = hashlib.sha256(img_in_real.tobytes()).hexdigest()
+        # print(hash)
+        inv_fname = f"tmp/{hash}_ddim_{num_ddim}_inv.pt"
+        caption_fname = f"tmp/{hash}_caption.txt"
+
+        # make the caption if it hasn't been made before
+        if not os.path.exists(caption_fname):
+            # BLIP
+            model_blip, vis_processors, _ = load_model_and_preprocess(name="blip_caption", model_type="base_coco", is_eval=True, device=torch.device("cuda"))
+            _image = vis_processors["eval"](img_in_real).unsqueeze(0).cuda()
+            prompt_str = model_blip.generate({"image": _image})[0]
+            del model_blip
+            torch.cuda.empty_cache()
+            with open(caption_fname, "w") as f:
+                f.write(prompt_str)
+        else:
+            prompt_str = open(caption_fname, "r").read().strip()
+        print(f"CAPTION: {prompt_str}")
+        
+        # do the inversion if it hasn't been done before
+        if not os.path.exists(inv_fname):
+            # inversion pipeline
+            pipe_inv = DDIMInversion.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32).to("cuda")
+            pipe_inv.scheduler = DDIMInverseScheduler.from_config(pipe_inv.scheduler.config)
+            x_inv, x_inv_image, x_dec_img = pipe_inv( prompt_str, 
+                    guidance_scale=1, num_inversion_steps=num_ddim,
+                    img=img_in_real, torch_dtype=torch.float32 )
+            x_inv = x_inv.detach()
+            torch.save(x_inv, inv_fname)
+            del pipe_inv
+            torch.cuda.empty_cache()
+        else:
+            x_inv = torch.load(inv_fname)
+
+        # do the editing
+        edit_pipe = EditingPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32).to("cuda")
+        edit_pipe.scheduler = DDIMScheduler.from_config(edit_pipe.scheduler.config)
+
+        _, edit_pil = edit_pipe(prompt_str,
+                num_inference_steps=num_ddim,
+                x_in=x_inv,
+                edit_dir=text_dir,
+                guidance_amount=xa_guidance,
+                guidance_scale=5.0,
+                negative_prompt=prompt_str # use the unedited prompt for the negative prompt
+        )
+        del edit_pipe
+        torch.cuda.empty_cache()
+        return edit_pil[0]
+
+
+    elif img_in_real is None and img_in_synth is not None:
+        print("using synthetic image")
+        x_inv = torch.load(fpath_z_gen)
+        pipe = EditingPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32).to("cuda")
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+        rec_pil, edit_pil = pipe(gen_prompt,
+            num_inference_steps=num_ddim,
+            x_in=x_inv,
+            edit_dir=text_dir,
+            guidance_amount=xa_guidance,
+            guidance_scale=5,
+            negative_prompt="" # use the empty string for the negative prompt
+        )
+        del pipe
+        torch.cuda.empty_cache()
+        return edit_pil[0]
+
+    else:
+        raise ValueError(f"Invalid image type: {image_type}")
+
+
+
+if __name__=="__main__":
+    print(flant5xl_compute_word2sentences("cat wearing sunglasses", num=100))

utils/gradio_utils.py

@@ -0,0 +1,616 @@
+import gradio as gr
+
+def set_visible_true():
+    return gr.update(visible=True)
+
+def set_visible_false():
+    return gr.update(visible=False)
+
+
+# HTML_header = f"""
+# <style>
+# {CSS_main}
+# </style>
+# <div style="text-align: center; max-width: 700px; margin: 0 auto;">
+#     <div style=" display: inline-flex; align-items: center; gap: 0.8rem; font-size: 1.75rem;">
+#         <h1 style="font-weight: 900; margin-bottom: 7px; margin-top: 5px;">
+#             Zero-shot Image-to-Image Translation
+#         </h1>
+#     </div>
+#     <p style="margin-bottom: 10px; font-size: 94%">
+#     This is the demo for <a href="https://pix2pixzero.github.io/" target="_blank">pix2pix-zero</a>. 
+#     Please visit our <a href="https://pix2pixzero.github.io/"> website</a> and <a href="https://github.com/pix2pixzero/pix2pix-zero" target="_blank">github</a> for more details.
+#     </p>
+#     <p style="margin-bottom: 10px; font-size: 94%">
+#     pix2pix-zero is a diffusion-based image-to-image approach that allows users to specify the edit direction on-the-fly 
+#     (e.g., cat to dog). Our method can directly use pre-trained text-to-image diffusion models, such as Stable Diffusion, 
+#     for editing real and synthetic images while preserving the input image's structure. Our method is training-free and prompt-free, 
+#     as it requires neither manual text prompting for each input image nor costly fine-tuning for each task.
+#     </p>
+
+# </div>
+# """
+
+
+
+
+CSS_main = """
+    body {
+    font-family: "HelveticaNeue-Light", "Helvetica Neue Light", "Helvetica Neue", Helvetica, Arial, "Lucida Grande", sans-serif; 
+    font-weight:300;
+    font-size:18px;
+    margin-left: auto;
+    margin-right: auto;
+    padding-left: 10px;
+    padding-right: 10px;
+    width: 800px;
+    }
+
+    h1 {
+        font-size:32px;
+        font-weight:300;
+        text-align: center;
+    }
+
+    h2 {
+        font-size:32px;
+        font-weight:300;
+        text-align: center;
+    }
+
+    #lbl_gallery_input{
+        font-family: 'Helvetica', 'Arial', sans-serif;
+        text-align: center;
+        color: #fff;
+        font-size: 28px;
+        display: inline
+    }
+
+
+    #lbl_gallery_comparision{
+        font-family: 'Helvetica', 'Arial', sans-serif;
+        text-align: center;
+        color: #fff;
+        font-size: 28px;
+    }
+
+    .disclaimerbox {
+        background-color: #eee;		
+        border: 1px solid #eeeeee;
+        border-radius: 10px ;
+        -moz-border-radius: 10px ;
+        -webkit-border-radius: 10px ;
+        padding: 20px;
+    }
+
+    video.header-vid {
+        height: 140px;
+        border: 1px solid black;
+        border-radius: 10px ;
+        -moz-border-radius: 10px ;
+        -webkit-border-radius: 10px ;
+    }
+
+    img.header-img {
+        height: 140px;
+        border: 1px solid black;
+        border-radius: 10px ;
+        -moz-border-radius: 10px ;
+        -webkit-border-radius: 10px ;
+    }
+
+    img.rounded {
+        border: 1px solid #eeeeee;
+        border-radius: 10px ;
+        -moz-border-radius: 10px ;
+        -webkit-border-radius: 10px ;
+    }
+
+    a:link
+    {
+        color: #941120;
+        text-decoration: none;
+    }
+    a:visited
+    {
+        color: #941120;
+        text-decoration: none;
+    }
+    a:hover {
+        color: #941120;
+    }
+
+    td.dl-link {
+        height: 160px;
+        text-align: center;
+        font-size: 22px;
+    }
+
+    .layered-paper-big { /* modified from: http://css-tricks.com/snippets/css/layered-paper/ */
+        box-shadow:
+        0px 0px 1px 1px rgba(0,0,0,0.35), /* The top layer shadow */
+        5px 5px 0 0px #fff, /* The second layer */
+        5px 5px 1px 1px rgba(0,0,0,0.35), /* The second layer shadow */
+        10px 10px 0 0px #fff, /* The third layer */
+        10px 10px 1px 1px rgba(0,0,0,0.35), /* The third layer shadow */
+        15px 15px 0 0px #fff, /* The fourth layer */
+        15px 15px 1px 1px rgba(0,0,0,0.35), /* The fourth layer shadow */
+        20px 20px 0 0px #fff, /* The fifth layer */
+        20px 20px 1px 1px rgba(0,0,0,0.35), /* The fifth layer shadow */
+        25px 25px 0 0px #fff, /* The fifth layer */
+        25px 25px 1px 1px rgba(0,0,0,0.35); /* The fifth layer shadow */
+        margin-left: 10px;
+        margin-right: 45px;
+    }
+
+    .paper-big { /* modified from: http://css-tricks.com/snippets/css/layered-paper/ */
+        box-shadow:
+        0px 0px 1px 1px rgba(0,0,0,0.35); /* The top layer shadow */
+
+        margin-left: 10px;
+        margin-right: 45px;
+    }
+
+
+    .layered-paper { /* modified from: http://css-tricks.com/snippets/css/layered-paper/ */
+        box-shadow:
+        0px 0px 1px 1px rgba(0,0,0,0.35), /* The top layer shadow */
+        5px 5px 0 0px #fff, /* The second layer */
+        5px 5px 1px 1px rgba(0,0,0,0.35), /* The second layer shadow */
+        10px 10px 0 0px #fff, /* The third layer */
+        10px 10px 1px 1px rgba(0,0,0,0.35); /* The third layer shadow */
+        margin-top: 5px;
+        margin-left: 10px;
+        margin-right: 30px;
+        margin-bottom: 5px;
+    }
+
+    .vert-cent {
+        position: relative;
+        top: 50%;
+        transform: translateY(-50%);
+    }
+
+    hr
+    {
+        border: 0;
+        height: 1px;
+        background-image: linear-gradient(to right, rgba(0, 0, 0, 0), rgba(0, 0, 0, 0.75), rgba(0, 0, 0, 0));
+    }
+
+    .card {
+        /* width: 130px;
+        height: 195px;
+        width: 1px;
+        height: 1px; */
+        position: relative;
+        display: inline-block;
+        /* margin: 50px; */
+    }
+    .card .img-top {
+        display: none;
+        position: absolute;
+        top: 0;
+        left: 0;
+        z-index: 99;
+    }
+    .card:hover .img-top {
+        display: inline;
+    }
+    details {
+    user-select: none;
+    }
+
+    details>summary span.icon {
+    width: 24px;
+    height: 24px;
+    transition: all 0.3s;
+    margin-left: auto;
+    }
+
+    details[open] summary span.icon {
+    transform: rotate(180deg);
+    }
+
+    summary {
+    display: flex;
+    cursor: pointer;
+    }
+
+    summary::-webkit-details-marker {
+    display: none;
+    }
+
+    ul {
+    display: table;
+    margin: 0 auto;
+    text-align: left;
+    }
+
+    .dark {
+        padding: 1em 2em;
+        background-color: #333;
+        box-shadow: 3px 3px 3px #333;
+        border: 1px #333;
+    }
+    .column {
+        float: left;
+        width: 20%;
+        padding: 0.5%;
+    }
+
+    .galleryImg {
+    transition: opacity 0.3s;
+    -webkit-transition: opacity 0.3s;
+    filter: grayscale(100%);
+    /* filter: blur(2px); */
+    -webkit-transition : -webkit-filter 250ms linear;
+    /* opacity: 0.5; */
+    cursor: pointer; 
+    }
+
+
+
+    .selected {	
+    /* outline: 100px solid var(--hover-background) !important; */
+    /* outline-offset: -100px; */
+    filter: grayscale(0%);
+    -webkit-transition : -webkit-filter 250ms linear;
+    /*opacity: 1.0 !important; */
+    }
+    
+    .galleryImg:hover {
+    filter: grayscale(0%);
+    -webkit-transition : -webkit-filter 250ms linear;
+
+    }
+
+    .row {
+    margin-bottom: 1em;
+    padding: 0px 1em;
+    }
+    /* Clear floats after the columns */
+    .row:after {
+    content: "";
+    display: table;
+    clear: both;
+    }
+    
+    /* The expanding image container */
+    #gallery {
+    position: relative;
+    /*display: none;*/
+    }
+
+    #section_comparison{
+        position: relative;
+        width: 100%;
+        height: max-content;
+    }
+
+    /* SLIDER
+    -------------------------------------------------- */
+
+    .slider-container {
+        position: relative;
+        height: 384px;
+        width: 512px;
+        cursor: grab;
+        overflow: hidden;
+        margin: auto;
+    }
+    .slider-after {
+        display: block;
+        position: absolute;
+        top: 0;
+        right: 0;
+        bottom: 0;
+        left: 0;
+        width: 100%;
+        height: 100%;
+        overflow: hidden;
+    }
+    .slider-before {
+        display: block;
+        position: absolute;
+        top: 0;
+        /* right: 0; */
+        bottom: 0;
+        left: 0;
+        width: 100%;
+        height: 100%;
+        z-index: 15;
+        overflow: hidden;
+    }
+    .slider-before-inset {
+        position: absolute;
+        top: 0;
+        bottom: 0;
+        left: 0;
+    }
+    .slider-after img,
+    .slider-before img {
+        object-fit: cover;
+        position: absolute;
+        width: 100%;
+        height: 100%;
+        object-position: 50% 50%;
+        top: 0;
+        bottom: 0;
+        left: 0;
+        -webkit-user-select: none;
+        -khtml-user-select: none;
+        -moz-user-select: none;
+        -o-user-select: none;
+        user-select: none;
+    }
+
+    #lbl_inset_left{
+        text-align: center;
+        position: absolute;
+        top: 384px;
+        width: 150px;
+        left: calc(50% - 256px);
+        z-index: 11;
+        font-size: 16px;
+        color: #fff;
+        margin: 10px;
+    }
+    .inset-before {
+        position: absolute;
+        width: 150px;
+        height: 150px;
+        box-shadow: 3px 3px 3px #333;
+        border: 1px #333;
+        border-style: solid;
+        z-index: 16;
+        top: 410px;
+        left: calc(50% - 256px);
+        margin: 10px;
+        font-size: 1em;
+        background-repeat: no-repeat;
+        pointer-events: none;
+    }
+
+    #lbl_inset_right{
+        text-align: center;
+        position: absolute;
+        top: 384px;
+        width: 150px;
+        right: calc(50% - 256px);
+        z-index: 11;
+        font-size: 16px;
+        color: #fff;
+        margin: 10px;
+    }
+    .inset-after {
+        position: absolute;
+        width: 150px;
+        height: 150px;
+        box-shadow: 3px 3px 3px #333;
+        border: 1px #333;
+        border-style: solid;
+        z-index: 16;
+        top: 410px;
+        right: calc(50% - 256px);
+        margin: 10px;
+        font-size: 1em;
+        background-repeat: no-repeat;
+        pointer-events: none;
+    }
+
+    #lbl_inset_input{
+        text-align: center;
+        position: absolute;
+        top: 384px;
+        width: 150px;
+        left: calc(50% - 256px + 150px + 20px);
+        z-index: 11;
+        font-size: 16px;
+        color: #fff;
+        margin: 10px;
+    }
+    .inset-target {
+        position: absolute;
+        width: 150px;
+        height: 150px;
+        box-shadow: 3px 3px 3px #333;
+        border: 1px #333;
+        border-style: solid;
+        z-index: 16;
+        top: 410px;
+        right: calc(50% - 256px + 150px + 20px);
+        margin: 10px;
+        font-size: 1em;
+        background-repeat: no-repeat;
+        pointer-events: none;
+    }
+
+    .slider-beforePosition {
+        background: #121212;
+        color: #fff;
+        left: 0;
+        pointer-events: none;
+        border-radius: 0.2rem;
+        padding: 2px 10px;
+    }
+    .slider-afterPosition {
+        background: #121212;
+        color: #fff;
+        right: 0;
+        pointer-events: none;
+        border-radius: 0.2rem;
+        padding: 2px 10px;
+    }
+    .beforeLabel {
+        position: absolute;
+        top: 0;
+        margin: 1rem;
+        font-size: 1em;
+        -webkit-user-select: none;
+        -khtml-user-select: none;
+        -moz-user-select: none;
+        -o-user-select: none;
+        user-select: none;
+    }
+    .afterLabel {
+        position: absolute;
+        top: 0;
+        margin: 1rem;
+        font-size: 1em;
+        -webkit-user-select: none;
+        -khtml-user-select: none;
+        -moz-user-select: none;
+        -o-user-select: none;
+        user-select: none;
+    }
+
+    .slider-handle {
+        height: 41px;
+        width: 41px;
+        position: absolute;
+        left: 50%;
+        top: 50%;
+        margin-left: -20px;
+        margin-top: -21px;
+        border: 2px solid #fff;
+        border-radius: 1000px;
+        z-index: 20;
+        pointer-events: none;
+        box-shadow: 0 0 10px rgb(12, 12, 12);
+    }
+    .handle-left-arrow,
+    .handle-right-arrow {
+        width: 0;
+        height: 0;
+        border: 6px inset transparent;
+        position: absolute;
+        top: 50%;
+        margin-top: -6px;
+    }
+    .handle-left-arrow {
+        border-right: 6px solid #fff;
+        left: 50%;
+        margin-left: -17px;
+    }
+    .handle-right-arrow {
+        border-left: 6px solid #fff;
+        right: 50%;
+        margin-right: -17px;
+    }
+    .slider-handle::before {
+        bottom: 50%;
+        margin-bottom: 20px;
+        box-shadow: 0 0 10px rgb(12, 12, 12);
+    }
+    .slider-handle::after {
+        top: 50%;
+        margin-top: 20.5px;
+        box-shadow: 0 0 5px rgb(12, 12, 12);
+    }
+    .slider-handle::before,
+    .slider-handle::after {
+        content: " ";
+        display: block;
+        width: 2px;
+        background: #fff;
+        height: 9999px;
+        position: absolute;
+        left: 50%;
+        margin-left: -1.5px;
+    }  
+
+
+    /* 
+    -------------------------------------------------
+    The editing results shown below inversion results
+    -------------------------------------------------
+    */
+    .edit_labels{
+        font-weight:500;
+        font-size: 24px;
+        color: #fff;
+        height: 20px;
+        margin-left: 20px;
+        position: relative; 
+        top: 20px;
+    }
+
+
+    .open > a:hover {
+        color: #555;
+        background-color: red;
+    }    
+
+
+    #directions { padding-top:30; padding-bottom:0; margin-bottom: 0px; height: 20px; }
+    #custom_task { padding-top:0; padding-bottom:0; margin-bottom: 0px; height: 20px; }
+    #slider_ddim {accent-color: #941120;}
+    #slider_ddim::-webkit-slider-thumb {background-color: #941120;}
+    #slider_xa {accent-color: #941120;}
+    #slider_xa::-webkit-slider-thumb {background-color: #941120;}
+    #slider_edit_mul {accent-color: #941120;}
+    #slider_edit_mul::-webkit-slider-thumb {background-color: #941120;}
+
+    #input_image [data-testid="image"]{
+        height: unset;
+    }
+    #input_image_synth [data-testid="image"]{
+        height: unset;
+    }
+
+
+"""
+
+
+
+HTML_header = f"""
+    <body>
+    <center>
+    <span style="font-size:36px">Zero-shot Image-to-Image Translation</span>
+    <table align=center>
+        <tr>
+            <td align=center>
+                <center>
+                    <span style="font-size:24px; margin-left: 0px;"><a href='https://pix2pixzero.github.io/'>[Website]</a></span>
+                    <span style="font-size:24px; margin-left: 20px;"><a href='https://github.com/pix2pixzero/pix2pix-zero'>[Code]</a></span>
+                </center>
+            </td>
+        </tr>
+    </table>
+    </center>
+
+    <center>
+    <div align=center>
+        <p align=left>
+        This is a demo for <span style="font-weight: bold;">pix2pix-zero</span>, a diffusion-based image-to-image approach that allows users to 
+        specify the edit direction on-the-fly (e.g., cat to dog). Our method can directly use pre-trained text-to-image diffusion models, such as Stable Diffusion, for editing real and synthetic images while preserving the input image's structure. Our method is training-free and prompt-free, as it requires neither manual text prompting for each input image nor costly fine-tuning for each task.
+        <br>
+        <span style="font-weight: 800;">TL;DR:</span> <span style=" color: #941120;"> no finetuning</span>  required; <span style=" color: #941120;"> no text input</span> needed; input <span style=" color: #941120;"> structure preserved</span>. 
+        </p>
+    </div>
+    </center>
+
+
+    <hr>
+    </body>
+"""
+
+HTML_input_header = f"""
+    <p style="font-size:150%; padding: 0px">
+    <span font-weight: 800; style=" color: #941120;"> Step 1: </span> select a real input image.
+    </p>
+"""
+
+HTML_middle_header = f"""
+    <p style="font-size:150%;">
+    <span font-weight: 800; style=" color: #941120;"> Step 2: </span> select the editing options.
+    </p>
+"""
+
+
+HTML_output_header = f"""
+    <p style="font-size:150%;">
+    <span font-weight: 800; style=" color: #941120;"> Step 3: </span> translated image!
+    </p>
+"""