Upload 53 files

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* 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
+assets/dog_sit.png filter=lfs diff=lfs merge=lfs -text
+assets/dog.png filter=lfs diff=lfs merge=lfs -text
+assets/teaser.gif filter=lfs diff=lfs merge=lfs -text
+assets/Teaser.png filter=lfs diff=lfs merge=lfs -text
+multi_image/assets/realdog.gif filter=lfs diff=lfs merge=lfs -text

LICENSE.txt

@@ -0,0 +1,10 @@
+S-Lab License 1.0  
+ 
+Copyright 2023 S-Lab
+Redistribution and use for non-commercial purpose in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+4. In the event that redistribution and/or use for commercial purpose in source or binary forms, with or without modification is required, please contact the contributor(s) of the work.
+ 

README.md

@@ -0,0 +1,131 @@
+<p align="center">
+  <h1 align="center">DiffMorpher: Unleashing the Capability of Diffusion Models for Image Morphing</h1>
+  <h3 align="center">CVPR 2024</h3>
+  <p align="center">
+    <a href="https://kevin-thu.github.io/homepage/"><strong>Kaiwen Zhang</strong></a>
+    &nbsp;&nbsp;
+    <a href="https://zhouyifan.net/about/"><strong>Yifan Zhou</strong></a>
+    &nbsp;&nbsp;
+    <a href="https://sheldontsui.github.io/"><strong>Xudong Xu</strong></a>
+    &nbsp;&nbsp;
+    <a href="https://xingangpan.github.io/"><strong>Xingang Pan<sep>✉</sep></strong></a>
+    &nbsp;&nbsp;
+    <a href="http://daibo.info/"><strong>Bo Dai</strong></a>
+  </p>
+  <br>
+
+  <p align="center">
+    <sep>✉</sep>Corresponding Author
+  </p>
+
+  <div align="center">
+        <img src="./assets/teaser.gif", width="500">
+  </div>
+
+  <p align="center">
+    <a href="https://arxiv.org/abs/2312.07409"><img alt='arXiv' src="https://img.shields.io/badge/arXiv-2312.07409-b31b1b.svg"></a>
+    <a href="https://kevin-thu.github.io/DiffMorpher_page/"><img alt='page' src="https://img.shields.io/badge/Project-Website-orange"></a>
+    <a href="https://twitter.com/sze68zkw"><img alt='Twitter' src="https://img.shields.io/twitter/follow/sze68zkw?label=%40KaiwenZhang"></a>
+    <a href="https://twitter.com/XingangP"><img alt='Twitter' src="https://img.shields.io/twitter/follow/XingangP?label=%40XingangPan"></a>
+  </p>
+  <br>
+</p>
+
+## Web Demos
+
+[![Open in OpenXLab](https://cdn-static.openxlab.org.cn/app-center/openxlab_app.svg)](https://openxlab.org.cn/apps/detail/KaiwenZhang/DiffMorpher)
+
+<p align="left">
+  <a href="https://huggingface.co/spaces/Kevin-thu/DiffMorpher"><img alt="Huggingface" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-DiffMorpher-orange"></a>
+</p>
+
+<!-- Great thanks to [OpenXLab](https://openxlab.org.cn/home) for the NVIDIA A100 GPU support! -->
+
+## Requirements
+To install the requirements, run the following in your environment first:
+```bash
+pip install -r requirements.txt
+```
+To run the code with CUDA properly, you can comment out `torch` and `torchvision` in `requirement.txt`, and install the appropriate version of `torch` and `torchvision` according to the instructions on [PyTorch](https://pytorch.org/get-started/locally/).
+
+You can also download the pretrained model *Stable Diffusion v2.1-base* from [Huggingface](https://huggingface.co/stabilityai/stable-diffusion-2-1-base), and specify the `model_path` to your local directory.
+
+## Run Gradio UI
+To start the Gradio UI of DiffMorpher, run the following in your environment:
+```bash
+python app.py
+```
+Then, by default, you can access the UI at [http://127.0.0.1:7860](http://127.0.0.1:7860).
+
+## Run the code
+You can also run the code with the following command:
+```bash
+python main.py \
+  --image_path_0 [image_path_0] --image_path_1 [image_path_1] \ 
+  --prompt_0 [prompt_0] --prompt_1 [prompt_1] \
+  --output_path [output_path] \
+  --use_adain --use_reschedule --save_inter
+```
+The script also supports the following options:
+
+- `--image_path_0`: Path of the first image (default: "")
+- `--prompt_0`: Prompt of the first image (default: "")
+- `--image_path_1`: Path of the second image (default: "")
+- `--prompt_1`: Prompt of the second image (default: "")
+- `--model_path`: Pretrained model path (default: "stabilityai/stable-diffusion-2-1-base")
+- `--output_path`: Path of the output image (default: "")
+- `--save_lora_dir`: Path of the output lora directory (default: "./lora")
+- `--load_lora_path_0`: Path of the lora directory of the first image (default: "")
+- `--load_lora_path_1`: Path of the lora directory of the second image (default: "")
+- `--use_adain`: Use AdaIN (default: False)
+- `--use_reschedule`: Use reschedule sampling (default: False)
+- `--lamb`: Hyperparameter $\lambda \in [0,1]$ for self-attention replacement, where a larger $\lambda$ indicates more replacements (default: 0.6)
+- `--fix_lora_value`: Fix lora value (default: LoRA Interpolation, not fixed)
+- `--save_inter`: Save intermediate results (default: False)
+- `--num_frames`: Number of frames to generate (default: 50)
+- `--duration`: Duration of each frame (default: 50)
+
+Examples:
+```bash
+python main.py \
+  --image_path_0 ./assets/Trump.jpg --image_path_1 ./assets/Biden.jpg \ 
+  --prompt_0 "A photo of an American man" --prompt_1 "A photo of an American man" \
+  --output_path "./results/Trump_Biden" \
+  --use_adain --use_reschedule --save_inter
+```
+
+```bash
+python main.py \
+  --image_path_0 ./assets/vangogh.jpg --image_path_1 ./assets/pearlgirl.jpg \ 
+  --prompt_0 "An oil painting of a man" --prompt_1 "An oil painting of a woman" \
+  --output_path "./results/vangogh_pearlgirl" \
+  --use_adain --use_reschedule --save_inter
+```
+
+```bash
+python main.py \
+  --image_path_0 ./assets/lion.png --image_path_1 ./assets/tiger.png \ 
+  --prompt_0 "A photo of a lion" --prompt_1 "A photo of a tiger" \
+  --output_path "./results/lion_tiger" \
+  --use_adain --use_reschedule --save_inter
+```
+
+## MorphBench
+To evaluate the effectiveness of our methods, we present *MorphBench*, the first benchmark dataset for assessing image morphing of general objects. You can download the dataset from [Google Drive](https://drive.google.com/file/d/1NWPzJhOgP-udP_wYbd0selRG4cu8xsu4/view?usp=sharing) or [Baidu Netdisk](https://pan.baidu.com/s/1J3xE3OJdEhKyoc1QObyYaA?pwd=putk).
+
+
+## License
+The code related to the DiffMorpher algorithm is licensed under [LICENSE](LICENSE.txt). 
+
+However, this project is mostly built on the open-sourse library [diffusers](https://github.com/huggingface/diffusers), which is under a separate license terms [Apache License 2.0](https://github.com/huggingface/diffusers/blob/main/LICENSE). (Cheers to the community as well!)
+
+## Citation
+
+```bibtex
+@article{zhang2023diffmorpher,
+    title={DiffMorpher: Unleashing the Capability of Diffusion Models for Image Morphing},
+    author={Zhang, Kaiwen and Zhou, Yifan and Xu, Xudong and Pan, Xingang and Dai, Bo},
+    journal={arXiv preprint arXiv:2312.07409},
+    year={2023}
+}
+```

app.py

@@ -0,0 +1,315 @@
+import os
+import torch
+import numpy as np
+import cv2
+import gradio as gr
+from PIL import Image
+from datetime import datetime
+from model import DiffMorpherPipeline
+from utils.lora_utils import train_lora
+
+LENGTH=450
+
+def train_lora_interface(
+    image,
+    prompt,
+    model_path,
+    output_path,
+    lora_steps,
+    lora_rank,
+    lora_lr,
+    num
+):
+    os.makedirs(output_path, exist_ok=True)
+    train_lora(image, prompt, output_path, model_path,
+               lora_steps=lora_steps, lora_lr=lora_lr, lora_rank=lora_rank, weight_name=f"lora_{num}.ckpt", progress=gr.Progress())
+    return f"Train LoRA {'A' if num == 0 else 'B'} Done!"
+
+def run_diffmorpher(
+    image_0,
+    image_1,
+    prompt_0,
+    prompt_1,
+    model_path,
+    lora_mode,
+    lamb,
+    use_adain,
+    use_reschedule,
+    num_frames,
+    fps,
+    save_inter,
+    load_lora_path_0,
+    load_lora_path_1,
+    output_path
+):
+    run_id = datetime.now().strftime("%H%M") + "_" +  datetime.now().strftime("%Y%m%d")
+    os.makedirs(output_path, exist_ok=True)
+    morpher_pipeline = DiffMorpherPipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cuda")
+    if lora_mode == "Fix LoRA A":
+        fix_lora = 0
+    elif lora_mode == "Fix LoRA B":
+        fix_lora = 1
+    else:
+        fix_lora = None
+    if not load_lora_path_0:
+        load_lora_path_0 = f"{output_path}/lora_0.ckpt"
+    if not load_lora_path_1:
+        load_lora_path_1 = f"{output_path}/lora_1.ckpt"
+    images = morpher_pipeline(
+        img_0=image_0,
+        img_1=image_1,
+        prompt_0=prompt_0,
+        prompt_1=prompt_1,
+        load_lora_path_0=load_lora_path_0,
+        load_lora_path_1=load_lora_path_1,
+        lamb=lamb,
+        use_adain=use_adain,
+        use_reschedule=use_reschedule,
+        num_frames=num_frames,
+        fix_lora=fix_lora,
+        save_intermediates=save_inter,
+        progress=gr.Progress()
+    )
+    video_path = f"{output_path}/{run_id}.mp4"
+    video = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (512, 512))
+    for i, image in enumerate(images):
+        # image.save(f"{output_path}/{i}.png")
+        video.write(cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR))
+    video.release()
+    cv2.destroyAllWindows()
+    return gr.Video(value=video_path, format="mp4", label="Output video", show_label=True, height=LENGTH, width=LENGTH, interactive=False)
+
+def run_all(
+    image_0,
+    image_1,
+    prompt_0,
+    prompt_1,
+    model_path,
+    lora_mode,
+    lamb,
+    use_adain,
+    use_reschedule,
+    num_frames,
+    fps,
+    save_inter,
+    load_lora_path_0,
+    load_lora_path_1,
+    output_path,
+    lora_steps,
+    lora_rank,
+    lora_lr
+):
+    os.makedirs(output_path, exist_ok=True)
+    train_lora(image_0, prompt_0, output_path, model_path,
+        lora_steps=lora_steps, lora_lr=lora_lr, lora_rank=lora_rank, weight_name=f"lora_0.ckpt", progress=gr.Progress())
+    train_lora(image_1, prompt_1, output_path, model_path,
+        lora_steps=lora_steps, lora_lr=lora_lr, lora_rank=lora_rank, weight_name=f"lora_1.ckpt", progress=gr.Progress())
+    return run_diffmorpher(
+        image_0,
+        image_1,
+        prompt_0,
+        prompt_1,
+        model_path,
+        lora_mode,
+        lamb,
+        use_adain,
+        use_reschedule,
+        num_frames,
+        fps,
+        save_inter,
+        load_lora_path_0,
+        load_lora_path_1,
+        output_path
+    )
+
+with gr.Blocks() as demo:
+    
+    with gr.Row():
+        gr.Markdown("""
+        # Official Implementation of [DiffMorpher](https://kevin-thu.github.io/DiffMorpher_page/)
+        """)
+
+    original_image_0, original_image_1 = gr.State(Image.open("assets/Trump.jpg").convert("RGB").resize((512,512), Image.BILINEAR)), gr.State(Image.open("assets/Biden.jpg").convert("RGB").resize((512,512), Image.BILINEAR))
+    # key_points_0, key_points_1 = gr.State([]), gr.State([])
+    # to_change_points = gr.State([])
+    
+    with gr.Row():
+        with gr.Column():
+            input_img_0 = gr.Image(type="numpy", label="Input image A", value="assets/Trump.jpg", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
+            prompt_0 = gr.Textbox(label="Prompt for image A", value="a photo of an American man", interactive=True)
+            with gr.Row():
+                train_lora_0_button = gr.Button("Train LoRA A")
+                train_lora_1_button = gr.Button("Train LoRA B")
+            # show_correspond_button = gr.Button("Show correspondence points")
+        with gr.Column():
+            input_img_1 = gr.Image(type="numpy", label="Input image B ", value="assets/Biden.jpg", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
+            prompt_1 = gr.Textbox(label="Prompt for image B", value="a photo of an American man", interactive=True)
+            with gr.Row():
+                clear_button = gr.Button("Clear All")
+                run_button = gr.Button("Run w/o LoRA training")
+        with gr.Column():
+            output_video = gr.Video(format="mp4", label="Output video", show_label=True, height=LENGTH, width=LENGTH, interactive=False)
+            lora_progress_bar = gr.Textbox(label="Display LoRA training progress", interactive=False)
+            run_all_button = gr.Button("Run!")
+        # with gr.Column():
+        #     output_video = gr.Video(label="Output video", show_label=True, height=LENGTH, width=LENGTH)
+        
+    with gr.Row():
+        gr.Markdown("""
+        ### Usage:
+        1. Upload two images (with correspondence) and fill out the prompts. 
+           (It's recommended to change `[Output path]` accordingly.)
+        2. Click **"Run!"**
+        
+        Or:
+        1. Upload two images (with correspondence) and fill out the prompts.
+        2. Click the **"Train LoRA A/B"** button to fit two LoRAs for two images respectively. <br> &nbsp;&nbsp;
+           If you have trained LoRA A or LoRA B before, you can skip the step and fill the specific LoRA path in LoRA settings. <br> &nbsp;&nbsp;
+           Trained LoRAs are saved to `[Output Path]/lora_0.ckpt` and `[Output Path]/lora_1.ckpt` by default.
+        3. You might also change the settings below.
+        4. Click **"Run w/o LoRA training"**
+        
+        ### Note: 
+        1. To speed up the generation process, you can **ruduce the number of frames** or **turn off "Use Reschedule"**.
+        2. You can try the influence of different prompts. It seems that using the same prompts or aligned prompts works better.
+        ### Have fun!
+        """)
+        
+    with gr.Accordion(label="Algorithm Parameters"):
+        with gr.Tab("Basic Settings"):
+            with gr.Row():
+                # local_models_dir = 'local_pretrained_models'
+                # local_models_choice = \
+                #     [os.path.join(local_models_dir,d) for d in os.listdir(local_models_dir) if os.path.isdir(os.path.join(local_models_dir,d))]
+                model_path = gr.Text(value="stabilityai/stable-diffusion-2-1-base",
+                    label="Diffusion Model Path", interactive=True
+                )
+                lamb = gr.Slider(value=0.6, minimum=0, maximum=1, step=0.1, label="Lambda for attention replacement", interactive=True)
+                lora_mode = gr.Dropdown(value="LoRA Interp",
+                    label="LoRA Interp. or Fix LoRA",
+                    choices=["LoRA Interp", "Fix LoRA A", "Fix LoRA B"],
+                    interactive=True
+                )
+                use_adain = gr.Checkbox(value=True, label="Use AdaIN", interactive=True)
+                use_reschedule = gr.Checkbox(value=True, label="Use Reschedule", interactive=True)
+            with gr.Row():
+                num_frames = gr.Number(value=16, minimum=0, label="Number of Frames", precision=0, interactive=True)
+                fps = gr.Number(value=8, minimum=0, label="FPS (Frame rate)", precision=0, interactive=True)
+                save_inter = gr.Checkbox(value=False, label="Save Intermediate Images", interactive=True)
+                output_path = gr.Text(value="./results", label="Output Path", interactive=True)
+                
+        with gr.Tab("LoRA Settings"):
+            with gr.Row():
+                lora_steps = gr.Number(value=200, label="LoRA training steps", precision=0, interactive=True)
+                lora_lr = gr.Number(value=0.0002, label="LoRA learning rate", interactive=True)
+                lora_rank = gr.Number(value=16, label="LoRA rank", precision=0, interactive=True)
+                # save_lora_dir = gr.Text(value="./lora", label="LoRA model save path", interactive=True)
+                load_lora_path_0 = gr.Text(value="", label="LoRA model load path for image A", interactive=True)
+                load_lora_path_1 = gr.Text(value="", label="LoRA model load path for image B", interactive=True)
+    
+    def store_img(img):
+        image = Image.fromarray(img).convert("RGB").resize((512,512), Image.BILINEAR)
+        # resize the input to 512x512
+        # image = image.resize((512,512), Image.BILINEAR)
+        # image = np.array(image)
+        # when new image is uploaded, `selected_points` should be empty
+        return image
+    input_img_0.upload(
+        store_img,
+        [input_img_0],
+        [original_image_0]
+    )
+    input_img_1.upload(
+        store_img,
+        [input_img_1],
+        [original_image_1]
+    )
+    
+    def clear(LENGTH):
+        return gr.Image.update(value=None, width=LENGTH, height=LENGTH), \
+            gr.Image.update(value=None, width=LENGTH, height=LENGTH), \
+            None, None, None, None
+    clear_button.click(
+        clear,
+        [gr.Number(value=LENGTH, visible=False, precision=0)],
+        [input_img_0, input_img_1, original_image_0, original_image_1, prompt_0, prompt_1]
+    )
+        
+    train_lora_0_button.click(
+        train_lora_interface,
+        [
+         original_image_0,
+         prompt_0,
+         model_path,
+         output_path,
+         lora_steps,
+         lora_rank,
+         lora_lr,
+         gr.Number(value=0, visible=False, precision=0)
+        ],
+        [lora_progress_bar]
+    )
+    
+    train_lora_1_button.click(
+        train_lora_interface,
+        [
+         original_image_1,
+         prompt_1,
+         model_path,
+         output_path,
+         lora_steps,
+         lora_rank,
+         lora_lr,
+         gr.Number(value=1, visible=False, precision=0)
+        ],
+        [lora_progress_bar]
+    )
+    
+    run_button.click(
+        run_diffmorpher,
+        [
+         original_image_0,
+         original_image_1,
+         prompt_0,
+         prompt_1,
+         model_path,
+         lora_mode,
+         lamb,
+         use_adain,
+         use_reschedule,
+         num_frames,
+         fps,
+         save_inter,
+         load_lora_path_0,
+         load_lora_path_1,
+         output_path
+        ],
+        [output_video]
+    )
+    
+    run_all_button.click(
+        run_all,
+        [
+         original_image_0,
+         original_image_1,
+         prompt_0,
+         prompt_1,
+         model_path,
+         lora_mode,
+         lamb,
+         use_adain,
+         use_reschedule,
+         num_frames,
+         fps,
+         save_inter,
+         load_lora_path_0,
+         load_lora_path_1,
+         output_path,
+         lora_steps,
+         lora_rank,
+         lora_lr
+        ],
+        [output_video]
+    )
+        
+demo.queue().launch(debug=True)

main.py

@@ -0,0 +1,98 @@
+import os
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+from argparse import ArgumentParser
+from model import DiffMorpherPipeline
+
+parser = ArgumentParser()
+parser.add_argument(
+    "--model_path", type=str, default="stabilityai/stable-diffusion-2-1-base",
+    help="Pretrained model to use (default: %(default)s)"
+)
+parser.add_argument(
+    "--image_path_0", type=str, default="",
+    help="Path of the first image (default: %(default)s)")
+parser.add_argument(
+    "--prompt_0", type=str, default="",
+    help="Prompt of the second image (default: %(default)s)")
+parser.add_argument(
+    "--image_path_1", type=str, default="",
+    help="Path of the first image (default: %(default)s)")
+parser.add_argument(
+    "--prompt_1", type=str, default="",
+    help="Prompt of the second image (default: %(default)s)")
+parser.add_argument(
+    "--output_path", type=str, default="./results",
+    help="Path of the output image (default: %(default)s)"
+)
+parser.add_argument(
+    "--save_lora_dir", type=str, default="./lora",
+    help="Path of the output lora directory (default: %(default)s)"
+)
+parser.add_argument(
+    "--load_lora_path_0", type=str, default="",
+    help="Path of the lora directory of the first image (default: %(default)s)"
+)
+parser.add_argument(
+    "--load_lora_path_1", type=str, default="",
+    help="Path of the lora directory of the second image (default: %(default)s)"
+)
+parser.add_argument(
+    "--use_adain", action="store_true",
+    help="Use AdaIN (default: %(default)s)"
+)
+parser.add_argument(
+    "--use_reschedule",  action="store_true",
+    help="Use reschedule sampling (default: %(default)s)"
+)
+parser.add_argument(
+    "--lamb",  type=float, default=0.6,
+    help="Lambda for self-attention replacement (default: %(default)s)"
+)
+parser.add_argument(
+    "--fix_lora_value", type=float, default=None,
+    help="Fix lora value (default: LoRA Interp., not fixed)"
+)
+parser.add_argument(
+    "--save_inter", action="store_true",
+    help="Save intermediate results (default: %(default)s)"
+)
+parser.add_argument(
+    "--num_frames", type=int, default=16,
+    help="Number of frames to generate (default: %(default)s)"
+)
+parser.add_argument(
+    "--duration", type=int, default=100,
+    help="Duration of each frame (default: %(default)s ms)"
+)
+parser.add_argument(
+    "--no_lora", action="store_true"
+)
+
+args = parser.parse_args()
+
+os.makedirs(args.output_path, exist_ok=True)
+pipeline = DiffMorpherPipeline.from_pretrained(
+    args.model_path, torch_dtype=torch.float32)
+pipeline.to("cuda")
+images = pipeline(
+    img_path_0=args.image_path_0,
+    img_path_1=args.image_path_1,
+    prompt_0=args.prompt_0,
+    prompt_1=args.prompt_1,
+    save_lora_dir=args.save_lora_dir,
+    load_lora_path_0=args.load_lora_path_0,
+    load_lora_path_1=args.load_lora_path_1,
+    use_adain=args.use_adain,
+    use_reschedule=args.use_reschedule,
+    lamd=args.lamb,
+    output_path=args.output_path,
+    num_frames=args.num_frames,
+    fix_lora=args.fix_lora_value,
+    save_intermediates=args.save_inter,
+    use_lora=not args.no_lora
+)
+images[0].save(f"{args.output_path}/output.gif", save_all=True,
+               append_images=images[1:], duration=args.duration, loop=0)

model.py

@@ -0,0 +1,639 @@
+import os
+from diffusers.models import AutoencoderKL, UNet2DConditionModel
+from diffusers.models.attention_processor import AttnProcessor
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from diffusers.schedulers import KarrasDiffusionSchedulers
+import torch
+import torch.nn.functional as F
+import tqdm
+import numpy as np
+import safetensors
+from PIL import Image
+from torchvision import transforms
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+from diffusers import StableDiffusionPipeline
+from argparse import ArgumentParser
+import inspect
+
+from utils.model_utils import get_img, slerp, do_replace_attn
+from utils.lora_utils import train_lora, load_lora
+from utils.alpha_scheduler import AlphaScheduler
+
+
+class StoreProcessor():
+    def __init__(self, original_processor, value_dict, name):
+        self.original_processor = original_processor
+        self.value_dict = value_dict
+        self.name = name
+        self.value_dict[self.name] = dict()
+        self.id = 0
+
+    def __call__(self, attn, hidden_states, *args, encoder_hidden_states=None, attention_mask=None, **kwargs):
+        # Is self attention
+        if encoder_hidden_states is None:
+            self.value_dict[self.name][self.id] = hidden_states.detach()
+            self.id += 1
+        res = self.original_processor(attn, hidden_states, *args,
+                                      encoder_hidden_states=encoder_hidden_states,
+                                      attention_mask=attention_mask,
+                                      **kwargs)
+
+        return res
+
+
+class LoadProcessor():
+    def __init__(self, original_processor, name, img0_dict, img1_dict, alpha, beta=0, lamd=0.6):
+        super().__init__()
+        self.original_processor = original_processor
+        self.name = name
+        self.img0_dict = img0_dict
+        self.img1_dict = img1_dict
+        self.alpha = alpha
+        self.beta = beta
+        self.lamd = lamd
+        self.id = 0
+
+    def __call__(self, attn, hidden_states, *args, encoder_hidden_states=None, attention_mask=None, **kwargs):
+        # Is self attention
+        if encoder_hidden_states is None:
+            if self.id < 50 * self.lamd:
+                map0 = self.img0_dict[self.name][self.id]
+                map1 = self.img1_dict[self.name][self.id]
+                cross_map = self.beta * hidden_states + \
+                    (1 - self.beta) * ((1 - self.alpha) * map0 + self.alpha * map1)
+                # cross_map = self.beta * hidden_states + \
+                #     (1 - self.beta) * slerp(map0, map1, self.alpha)
+                # cross_map = slerp(slerp(map0, map1, self.alpha),
+                #                   hidden_states, self.beta)
+                # cross_map = hidden_states
+                # cross_map = torch.cat(
+                #     ((1 - self.alpha) * map0, self.alpha * map1), dim=1)
+
+                res = self.original_processor(attn, hidden_states, *args,
+                                              encoder_hidden_states=cross_map,
+                                              attention_mask=attention_mask,
+                                              **kwargs)
+            else:
+                res = self.original_processor(attn, hidden_states, *args,
+                                              encoder_hidden_states=encoder_hidden_states,
+                                              attention_mask=attention_mask,
+                                              **kwargs)
+
+            self.id += 1
+            # if self.id == len(self.img0_dict[self.name]):
+            if self.id == len(self.img0_dict[self.name]):
+                self.id = 0
+        else:
+            res = self.original_processor(attn, hidden_states, *args,
+                                          encoder_hidden_states=encoder_hidden_states,
+                                          attention_mask=attention_mask,
+                                          **kwargs)
+
+        return res
+
+
+class DiffMorpherPipeline(StableDiffusionPipeline):
+
+    def __init__(self,
+                 vae: AutoencoderKL,
+                 text_encoder: CLIPTextModel,
+                 tokenizer: CLIPTokenizer,
+                 unet: UNet2DConditionModel,
+                 scheduler: KarrasDiffusionSchedulers,
+                 safety_checker: StableDiffusionSafetyChecker,
+                 feature_extractor: CLIPImageProcessor,
+                 image_encoder=None,
+                 requires_safety_checker: bool = True,
+                 ):
+        sig = inspect.signature(super().__init__)
+        params = sig.parameters
+        if 'image_encoder' in params:
+            super().__init__(vae, text_encoder, tokenizer, unet, scheduler,
+                             safety_checker, feature_extractor, image_encoder, requires_safety_checker)
+        else:
+            super().__init__(vae, text_encoder, tokenizer, unet, scheduler,
+                             safety_checker, feature_extractor, requires_safety_checker)
+        self.img0_dict = dict()
+        self.img1_dict = dict()
+
+    def inv_step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        x: torch.FloatTensor,
+        eta=0.,
+        verbose=False
+    ):
+        """
+        Inverse sampling for DDIM Inversion
+        """
+        if verbose:
+            print("timestep: ", timestep)
+        next_step = timestep
+        timestep = min(timestep - self.scheduler.config.num_train_timesteps //
+                       self.scheduler.num_inference_steps, 999)
+        alpha_prod_t = self.scheduler.alphas_cumprod[
+            timestep] if timestep >= 0 else self.scheduler.final_alpha_cumprod
+        alpha_prod_t_next = self.scheduler.alphas_cumprod[next_step]
+        beta_prod_t = 1 - alpha_prod_t
+        pred_x0 = (x - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5
+        pred_dir = (1 - alpha_prod_t_next)**0.5 * model_output
+        x_next = alpha_prod_t_next**0.5 * pred_x0 + pred_dir
+        return x_next, pred_x0
+
+    @torch.no_grad()
+    def invert(
+            self,
+            image: torch.Tensor,
+            prompt,
+            num_inference_steps=50,
+            num_actual_inference_steps=None,
+            guidance_scale=1.,
+            eta=0.0,
+            **kwds):
+        """
+        invert a real image into noise map with determinisc DDIM inversion
+        """
+        DEVICE = torch.device(
+            "cuda") if torch.cuda.is_available() else torch.device("cpu")
+        batch_size = image.shape[0]
+        if isinstance(prompt, list):
+            if batch_size == 1:
+                image = image.expand(len(prompt), -1, -1, -1)
+        elif isinstance(prompt, str):
+            if batch_size > 1:
+                prompt = [prompt] * batch_size
+
+        # text embeddings
+        text_input = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=77,
+            return_tensors="pt"
+        )
+        text_embeddings = self.text_encoder(text_input.input_ids.to(DEVICE))[0]
+        print("input text embeddings :", text_embeddings.shape)
+        # define initial latents
+        latents = self.image2latent(image)
+
+        # unconditional embedding for classifier free guidance
+        if guidance_scale > 1.:
+            max_length = text_input.input_ids.shape[-1]
+            unconditional_input = self.tokenizer(
+                [""] * batch_size,
+                padding="max_length",
+                max_length=77,
+                return_tensors="pt"
+            )
+            unconditional_embeddings = self.text_encoder(
+                unconditional_input.input_ids.to(DEVICE))[0]
+            text_embeddings = torch.cat(
+                [unconditional_embeddings, text_embeddings], dim=0)
+
+        print("latents shape: ", latents.shape)
+        # interative sampling
+        self.scheduler.set_timesteps(num_inference_steps)
+        print("Valid timesteps: ", reversed(self.scheduler.timesteps))
+        # print("attributes: ", self.scheduler.__dict__)
+        latents_list = [latents]
+        pred_x0_list = [latents]
+        for i, t in enumerate(tqdm.tqdm(reversed(self.scheduler.timesteps), desc="DDIM Inversion")):
+            if num_actual_inference_steps is not None and i >= num_actual_inference_steps:
+                continue
+
+            if guidance_scale > 1.:
+                model_inputs = torch.cat([latents] * 2)
+            else:
+                model_inputs = latents
+
+            # predict the noise
+            noise_pred = self.unet(
+                model_inputs, t, encoder_hidden_states=text_embeddings).sample
+            if guidance_scale > 1.:
+                noise_pred_uncon, noise_pred_con = noise_pred.chunk(2, dim=0)
+                noise_pred = noise_pred_uncon + guidance_scale * \
+                    (noise_pred_con - noise_pred_uncon)
+            # compute the previous noise sample x_t-1 -> x_t
+            latents, pred_x0 = self.inv_step(noise_pred, t, latents)
+            latents_list.append(latents)
+            pred_x0_list.append(pred_x0)
+
+        return latents
+
+    @torch.no_grad()
+    def ddim_inversion(self, latent, cond):
+        timesteps = reversed(self.scheduler.timesteps)
+        with torch.autocast(device_type='cuda', dtype=torch.float32):
+            for i, t in enumerate(tqdm.tqdm(timesteps, desc="DDIM inversion")):
+                cond_batch = cond.repeat(latent.shape[0], 1, 1)
+
+                alpha_prod_t = self.scheduler.alphas_cumprod[t]
+                alpha_prod_t_prev = (
+                    self.scheduler.alphas_cumprod[timesteps[i - 1]]
+                    if i > 0 else self.scheduler.final_alpha_cumprod
+                )
+
+                mu = alpha_prod_t ** 0.5
+                mu_prev = alpha_prod_t_prev ** 0.5
+                sigma = (1 - alpha_prod_t) ** 0.5
+                sigma_prev = (1 - alpha_prod_t_prev) ** 0.5
+
+                eps = self.unet(
+                    latent, t, encoder_hidden_states=cond_batch).sample
+
+                pred_x0 = (latent - sigma_prev * eps) / mu_prev
+                latent = mu * pred_x0 + sigma * eps
+        #         if save_latents:
+        #             torch.save(latent, os.path.join(save_path, f'noisy_latents_{t}.pt'))
+        # torch.save(latent, os.path.join(save_path, f'noisy_latents_{t}.pt'))
+        return latent
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        x: torch.FloatTensor,
+    ):
+        """
+        predict the sample of the next step in the denoise process.
+        """
+        prev_timestep = timestep - \
+            self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
+        alpha_prod_t = self.scheduler.alphas_cumprod[timestep]
+        alpha_prod_t_prev = self.scheduler.alphas_cumprod[
+            prev_timestep] if prev_timestep > 0 else self.scheduler.final_alpha_cumprod
+        beta_prod_t = 1 - alpha_prod_t
+        pred_x0 = (x - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5
+        pred_dir = (1 - alpha_prod_t_prev)**0.5 * model_output
+        x_prev = alpha_prod_t_prev**0.5 * pred_x0 + pred_dir
+        return x_prev, pred_x0
+
+    @torch.no_grad()
+    def image2latent(self, image):
+        DEVICE = torch.device(
+            "cuda") if torch.cuda.is_available() else torch.device("cpu")
+        if type(image) is Image:
+            image = np.array(image)
+            image = torch.from_numpy(image).float() / 127.5 - 1
+            image = image.permute(2, 0, 1).unsqueeze(0)
+        # input image density range [-1, 1]
+        latents = self.vae.encode(image.to(DEVICE))['latent_dist'].mean
+        latents = latents * 0.18215
+        return latents
+
+    @torch.no_grad()
+    def latent2image(self, latents, return_type='np'):
+        latents = 1 / 0.18215 * latents.detach()
+        image = self.vae.decode(latents)['sample']
+        if return_type == 'np':
+            image = (image / 2 + 0.5).clamp(0, 1)
+            image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
+            image = (image * 255).astype(np.uint8)
+        elif return_type == "pt":
+            image = (image / 2 + 0.5).clamp(0, 1)
+
+        return image
+
+    def latent2image_grad(self, latents):
+        latents = 1 / 0.18215 * latents
+        image = self.vae.decode(latents)['sample']
+
+        return image  # range [-1, 1]
+
+    @torch.no_grad()
+    def cal_latent(self, num_inference_steps, guidance_scale, unconditioning, img_noise_0, img_noise_1, text_embeddings_0, text_embeddings_1, lora_0, lora_1, alpha, use_lora, fix_lora=None):
+        # latents = torch.cos(alpha * torch.pi / 2) * img_noise_0 + \
+        #     torch.sin(alpha * torch.pi / 2) * img_noise_1
+        # latents = (1 - alpha) * img_noise_0 + alpha * img_noise_1
+        # latents = latents / ((1 - alpha) ** 2 + alpha ** 2)
+        latents = slerp(img_noise_0, img_noise_1, alpha, self.use_adain)
+        text_embeddings = (1 - alpha) * text_embeddings_0 + \
+            alpha * text_embeddings_1
+
+        self.scheduler.set_timesteps(num_inference_steps)
+        if use_lora:
+            if fix_lora is not None:
+                self.unet = load_lora(self.unet, lora_0, lora_1, fix_lora)
+            else:
+                self.unet = load_lora(self.unet, lora_0, lora_1, alpha)
+
+        for i, t in enumerate(tqdm.tqdm(self.scheduler.timesteps, desc=f"DDIM Sampler, alpha={alpha}")):
+
+            if guidance_scale > 1.:
+                model_inputs = torch.cat([latents] * 2)
+            else:
+                model_inputs = latents
+            if unconditioning is not None and isinstance(unconditioning, list):
+                _, text_embeddings = text_embeddings.chunk(2)
+                text_embeddings = torch.cat(
+                    [unconditioning[i].expand(*text_embeddings.shape), text_embeddings])
+            # predict the noise
+            noise_pred = self.unet(
+                model_inputs, t, encoder_hidden_states=text_embeddings).sample
+            if guidance_scale > 1.0:
+                noise_pred_uncon, noise_pred_con = noise_pred.chunk(
+                    2, dim=0)
+                noise_pred = noise_pred_uncon + guidance_scale * \
+                    (noise_pred_con - noise_pred_uncon)
+            # compute the previous noise sample x_t -> x_t-1
+            latents = self.scheduler.step(
+                noise_pred, t, latents, return_dict=False)[0]
+        return latents
+
+    @torch.no_grad()
+    def get_text_embeddings(self, prompt, guidance_scale, neg_prompt, batch_size):
+        DEVICE = torch.device(
+            "cuda") if torch.cuda.is_available() else torch.device("cpu")
+        # text embeddings
+        text_input = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=77,
+            return_tensors="pt"
+        )
+        text_embeddings = self.text_encoder(text_input.input_ids.cuda())[0]
+
+        if guidance_scale > 1.:
+            if neg_prompt:
+                uc_text = neg_prompt
+            else:
+                uc_text = ""
+            unconditional_input = self.tokenizer(
+                [uc_text] * batch_size,
+                padding="max_length",
+                max_length=77,
+                return_tensors="pt"
+            )
+            unconditional_embeddings = self.text_encoder(
+                unconditional_input.input_ids.to(DEVICE))[0]
+            text_embeddings = torch.cat(
+                [unconditional_embeddings, text_embeddings], dim=0)
+
+        return text_embeddings
+
+    def __call__(
+            self,
+            img_0=None,
+            img_1=None,
+            img_path_0=None,
+            img_path_1=None,
+            prompt_0="",
+            prompt_1="",
+            save_lora_dir="./lora",
+            load_lora_path_0=None,
+            load_lora_path_1=None,
+            lora_steps=200,
+            lora_lr=2e-4,
+            lora_rank=16,
+            batch_size=1,
+            height=512,
+            width=512,
+            num_inference_steps=50,
+            num_actual_inference_steps=None,
+            guidance_scale=1,
+            attn_beta=0,
+            lamd=0.6,
+            use_lora=True,
+            use_adain=True,
+            use_reschedule=True,
+            output_path="./results",
+            num_frames=50,
+            fix_lora=None,
+            progress=tqdm,
+            unconditioning=None,
+            neg_prompt=None,
+            save_intermediates=False,
+            **kwds):
+
+        # if isinstance(prompt, list):
+        #     batch_size = len(prompt)
+        # elif isinstance(prompt, str):
+        #     if batch_size > 1:
+        #         prompt = [prompt] * batch_size
+        self.scheduler.set_timesteps(num_inference_steps)
+        self.use_lora = use_lora
+        self.use_adain = use_adain
+        self.use_reschedule = use_reschedule
+        self.output_path = output_path
+
+        if img_0 is None:
+            img_0 = Image.open(img_path_0).convert("RGB")
+        # else:
+        #     img_0 = Image.fromarray(img_0).convert("RGB")
+
+        if img_1 is None:
+            img_1 = Image.open(img_path_1).convert("RGB")
+        # else:
+        #     img_1 = Image.fromarray(img_1).convert("RGB")
+
+        if self.use_lora:
+            print("Loading lora...")
+            if not load_lora_path_0:
+
+                weight_name = f"{output_path.split('/')[-1]}_lora_0.ckpt"
+                load_lora_path_0 = save_lora_dir + "/" + weight_name
+                if not os.path.exists(load_lora_path_0):
+                    train_lora(img_0, prompt_0, save_lora_dir, None, self.tokenizer, self.text_encoder,
+                               self.vae, self.unet, self.scheduler, lora_steps, lora_lr, lora_rank, weight_name=weight_name)
+            print(f"Load from {load_lora_path_0}.")
+            if load_lora_path_0.endswith(".safetensors"):
+                lora_0 = safetensors.torch.load_file(
+                    load_lora_path_0, device="cpu")
+            else:
+                lora_0 = torch.load(load_lora_path_0, map_location="cpu")
+
+            if not load_lora_path_1:
+                weight_name = f"{output_path.split('/')[-1]}_lora_1.ckpt"
+                load_lora_path_1 = save_lora_dir + "/" + weight_name
+                if not os.path.exists(load_lora_path_1):
+                    train_lora(img_1, prompt_1, save_lora_dir, None, self.tokenizer, self.text_encoder,
+                               self.vae, self.unet, self.scheduler, lora_steps, lora_lr, lora_rank, weight_name=weight_name)
+            print(f"Load from {load_lora_path_1}.")
+            if load_lora_path_1.endswith(".safetensors"):
+                lora_1 = safetensors.torch.load_file(
+                    load_lora_path_1, device="cpu")
+            else:
+                lora_1 = torch.load(load_lora_path_1, map_location="cpu")
+        else:
+            lora_0 = lora_1 = None
+
+        text_embeddings_0 = self.get_text_embeddings(
+            prompt_0, guidance_scale, neg_prompt, batch_size)
+        text_embeddings_1 = self.get_text_embeddings(
+            prompt_1, guidance_scale, neg_prompt, batch_size)
+        img_0 = get_img(img_0)
+        img_1 = get_img(img_1)
+        if self.use_lora:
+            self.unet = load_lora(self.unet, lora_0, lora_1, 0)
+        img_noise_0 = self.ddim_inversion(
+            self.image2latent(img_0), text_embeddings_0)
+        if self.use_lora:
+            self.unet = load_lora(self.unet, lora_0, lora_1, 1)
+        img_noise_1 = self.ddim_inversion(
+            self.image2latent(img_1), text_embeddings_1)
+
+        print("latents shape: ", img_noise_0.shape)
+
+        original_processor = list(self.unet.attn_processors.values())[0]
+
+        def morph(alpha_list, progress, desc):
+            images = []
+            if attn_beta is not None:
+                if self.use_lora:
+                    self.unet = load_lora(
+                        self.unet, lora_0, lora_1, 0 if fix_lora is None else fix_lora)
+
+                attn_processor_dict = {}
+                for k in self.unet.attn_processors.keys():
+                    if do_replace_attn(k):
+                        if self.use_lora:
+                            attn_processor_dict[k] = StoreProcessor(self.unet.attn_processors[k],
+                                                                    self.img0_dict, k)
+                        else:
+                            attn_processor_dict[k] = StoreProcessor(original_processor,
+                                                                    self.img0_dict, k)
+                    else:
+                        attn_processor_dict[k] = self.unet.attn_processors[k]
+                self.unet.set_attn_processor(attn_processor_dict)
+
+                latents = self.cal_latent(
+                    num_inference_steps,
+                    guidance_scale,
+                    unconditioning,
+                    img_noise_0,
+                    img_noise_1,
+                    text_embeddings_0,
+                    text_embeddings_1,
+                    lora_0,
+                    lora_1,
+                    alpha_list[0],
+                    False,
+                    fix_lora
+                )
+                first_image = self.latent2image(latents)
+                first_image = Image.fromarray(first_image)
+                if save_intermediates:
+                    first_image.save(f"{self.output_path}/{0:02d}.png")
+
+                if self.use_lora:
+                    self.unet = load_lora(
+                        self.unet, lora_0, lora_1, 1 if fix_lora is None else fix_lora)
+                attn_processor_dict = {}
+                for k in self.unet.attn_processors.keys():
+                    if do_replace_attn(k):
+                        if self.use_lora:
+                            attn_processor_dict[k] = StoreProcessor(self.unet.attn_processors[k],
+                                                                    self.img1_dict, k)
+                        else:
+                            attn_processor_dict[k] = StoreProcessor(original_processor,
+                                                                    self.img1_dict, k)
+                    else:
+                        attn_processor_dict[k] = self.unet.attn_processors[k]
+
+                self.unet.set_attn_processor(attn_processor_dict)
+
+                latents = self.cal_latent(
+                    num_inference_steps,
+                    guidance_scale,
+                    unconditioning,
+                    img_noise_0,
+                    img_noise_1,
+                    text_embeddings_0,
+                    text_embeddings_1,
+                    lora_0,
+                    lora_1,
+                    alpha_list[-1],
+                    False,
+                    fix_lora
+                )
+                last_image = self.latent2image(latents)
+                last_image = Image.fromarray(last_image)
+                if save_intermediates:
+                    last_image.save(
+                        f"{self.output_path}/{num_frames - 1:02d}.png")
+
+                for i in progress.tqdm(range(1, num_frames - 1), desc=desc):
+                    alpha = alpha_list[i]
+                    if self.use_lora:
+                        self.unet = load_lora(
+                            self.unet, lora_0, lora_1, alpha if fix_lora is None else fix_lora)
+
+                    attn_processor_dict = {}
+                    for k in self.unet.attn_processors.keys():
+                        if do_replace_attn(k):
+                            if self.use_lora:
+                                attn_processor_dict[k] = LoadProcessor(
+                                    self.unet.attn_processors[k], k, self.img0_dict, self.img1_dict, alpha, attn_beta, lamd)
+                            else:
+                                attn_processor_dict[k] = LoadProcessor(
+                                    original_processor, k, self.img0_dict, self.img1_dict, alpha, attn_beta, lamd)
+                        else:
+                            attn_processor_dict[k] = self.unet.attn_processors[k]
+
+                    self.unet.set_attn_processor(attn_processor_dict)
+
+                    latents = self.cal_latent(
+                        num_inference_steps,
+                        guidance_scale,
+                        unconditioning,
+                        img_noise_0,
+                        img_noise_1,
+                        text_embeddings_0,
+                        text_embeddings_1,
+                        lora_0,
+                        lora_1,
+                        alpha_list[i],
+                        False,
+                        fix_lora
+                    )
+                    image = self.latent2image(latents)
+                    image = Image.fromarray(image)
+                    if save_intermediates:
+                        image.save(f"{self.output_path}/{i:02d}.png")
+                    images.append(image)
+
+                images = [first_image] + images + [last_image]
+
+            else:
+                for k, alpha in enumerate(alpha_list):
+
+                    latents = self.cal_latent(
+                        num_inference_steps,
+                        guidance_scale,
+                        unconditioning,
+                        img_noise_0,
+                        img_noise_1,
+                        text_embeddings_0,
+                        text_embeddings_1,
+                        lora_0,
+                        lora_1,
+                        alpha_list[k],
+                        self.use_lora,
+                        fix_lora
+                    )
+                    image = self.latent2image(latents)
+                    image = Image.fromarray(image)
+                    if save_intermediates:
+                        image.save(f"{self.output_path}/{k:02d}.png")
+                    images.append(image)
+
+            return images
+
+        with torch.no_grad():
+            if self.use_reschedule:
+                alpha_scheduler = AlphaScheduler()
+                alpha_list = list(torch.linspace(0, 1, num_frames))
+                images_pt = morph(alpha_list, progress, "Sampling...")
+                images_pt = [transforms.ToTensor()(img).unsqueeze(0)
+                             for img in images_pt]
+                alpha_scheduler.from_imgs(images_pt)
+                alpha_list = alpha_scheduler.get_list()
+                print(alpha_list)
+                images = morph(alpha_list, progress, "Reschedule..."
+                               )
+            else:
+                alpha_list = list(torch.linspace(0, 1, num_frames))
+                print(alpha_list)
+                images = morph(alpha_list, progress, "Sampling...")
+
+        return images

multi_image/README.md

@@ -0,0 +1,37 @@
+## Update
+
+Add support for multi-image input. Now you can get the morphing output among more than 2 images.
+
+## Run the code
+
+You can run the code with the following command:
+
+```
+python main.py \
+  --image_paths [image_path_0] ... [image_path_n] \ 
+  --prompts [prompt_0] ... [prompt_n] \
+  --output_path [output_path] \
+  --use_adain --use_reschedule --save_inter
+```
+
+This modification add support for the following options:
+
+- `--image_paths`: Paths of the input images
+- `--prompts`: Prompts of the images
+- `--load_lora_paths`: Paths of the lora directory of the images
+
+## Example
+
+Run the code:
+```
+python main.py \
+--image_paths ./assets/realdog0.jpg ./assets/realdog1.jpg ./assets/realdog2.jpg \
+--prompts "A photo of a dog" "A photo of a dog" "A photo of a dog"  \
+--output_path "./results/dog" \
+--use_adain --use_reschedule --save_inter
+```
+
+Output:
+<div  align="center">
+<img src="assets/realdog.gif" width="50%" height="50%">
+</div>

multi_image/main.py

@@ -0,0 +1,107 @@
+import os
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+from argparse import ArgumentParser
+from model import DiffMorpherPipeline
+
+parser = ArgumentParser()
+parser.add_argument(
+    "--model_path", type=str, default="stabilityai/stable-diffusion-2-1-base",
+    help="Pretrained model to use (default: %(default)s)"
+)
+parser.add_argument(
+    "--image_path_0", type=str, default="",
+    help="Path of the first image (default: %(default)s)")
+parser.add_argument(
+    "--prompt_0", type=str, default="",
+    help="Prompt of the second image (default: %(default)s)")
+parser.add_argument(
+    "--image_path_1", type=str, default="",
+    help="Path of the first image (default: %(default)s)")
+parser.add_argument(
+    "--prompt_1", type=str, default="",
+    help="Prompt of the second image (default: %(default)s)")
+parser.add_argument(
+    "--load_lora_path_0", type=str, default="",
+    help="Path of the lora directory of the first image (default: %(default)s)"
+)
+parser.add_argument(
+    "--load_lora_path_1", type=str, default="",
+    help="Path of the lora directory of the second image (default: %(default)s)"
+)
+parser.add_argument(
+    "--image_paths", type=str, nargs='*', default=[],
+    help="Path of the first image (default: %(default)s)")
+parser.add_argument(
+    "--prompts", type=str, nargs='*', default=[],
+    help="Prompt of the second image (default: %(default)s)")
+parser.add_argument(
+    "--output_path", type=str, default="./results",
+    help="Path of the output image (default: %(default)s)"
+)
+parser.add_argument(
+    "--save_lora_dir", type=str, default="./lora",
+    help="Path of the output lora directory (default: %(default)s)"
+)
+parser.add_argument(
+    "--load_lora_paths", type=str, nargs='*', default=[],
+    help="Path of the lora directory of the first image (default: %(default)s)"
+)
+parser.add_argument(
+    "--use_adain", action="store_true",
+    help="Use AdaIN (default: %(default)s)"
+)
+parser.add_argument(
+    "--use_reschedule",  action="store_true",
+    help="Use reschedule sampling (default: %(default)s)"
+)
+parser.add_argument(
+    "--lamb",  type=float, default=0.6,
+    help="Lambda for self-attention replacement (default: %(default)s)"
+)
+parser.add_argument(
+    "--fix_lora_value", type=float, default=None,
+    help="Fix lora value (default: LoRA Interp., not fixed)"
+)
+parser.add_argument(
+    "--save_inter", action="store_true",
+    help="Save intermediate results (default: %(default)s)"
+)
+parser.add_argument(
+    "--num_frames", type=int, default=16,
+    help="Number of frames to generate (default: %(default)s)"
+)
+parser.add_argument(
+    "--duration", type=int, default=100,
+    help="Duration of each frame (default: %(default)s ms)"
+)
+
+args = parser.parse_args()
+
+os.makedirs(args.output_path, exist_ok=True)
+pipeline = DiffMorpherPipeline.from_pretrained(
+    args.model_path, torch_dtype=torch.float32)
+pipeline.to("cuda")
+images = pipeline(
+    img_path_0=args.image_path_0,
+    img_path_1=args.image_path_1,
+    prompt_0=args.prompt_0,
+    prompt_1=args.prompt_1,
+    load_lora_path_0=args.load_lora_path_0,
+    load_lora_path_1=args.load_lora_path_1,
+    img_paths=args.image_paths,
+    prompts=args.prompts,
+    save_lora_dir=args.save_lora_dir,
+    load_lora_paths=args.load_lora_paths,
+    use_adain=args.use_adain,
+    use_reschedule=args.use_reschedule,
+    lamb=args.lamb,
+    output_path=args.output_path,
+    num_frames=args.num_frames,
+    fix_lora=args.fix_lora_value,
+    save_intermediates=args.save_inter,
+)
+images[0].save(f"{args.output_path}/output.gif", save_all=True,
+               append_images=images[1:], duration=args.duration, loop=0)

multi_image/model.py

@@ -0,0 +1,699 @@
+import os
+from diffusers.models import AutoencoderKL, UNet2DConditionModel
+from diffusers.models.attention_processor import AttnProcessor
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from diffusers.schedulers import KarrasDiffusionSchedulers
+import torch
+import torch.nn.functional as F
+import tqdm
+import numpy as np
+import safetensors
+from PIL import Image
+from torchvision import transforms
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+from diffusers import StableDiffusionPipeline
+from argparse import ArgumentParser
+
+
+from utils.model_utils import get_img, slerp, do_replace_attn
+from utils.lora_utils import train_lora, load_lora
+from utils.alpha_scheduler import AlphaScheduler
+
+class StoreProcessor():
+    def __init__(self, original_processor, value_dict, name):
+        self.original_processor = original_processor
+        self.value_dict = value_dict
+        self.name = name
+        self.value_dict[self.name] = dict()
+        self.id = 0
+
+    def __call__(self, attn, hidden_states, *args, encoder_hidden_states=None, attention_mask=None, **kwargs):
+        # Is self attention
+        if encoder_hidden_states is None:
+            self.value_dict[self.name][self.id] = hidden_states.detach()
+            self.id += 1
+        res = self.original_processor(attn, hidden_states, *args,
+                                      encoder_hidden_states=encoder_hidden_states,
+                                      attention_mask=attention_mask,
+                                      **kwargs)
+
+        return res
+    
+
+class LoadProcessor():
+    def __init__(self, original_processor, name, img0_dict, img1_dict, alpha, beta=0, lamd=0.6):
+        super().__init__()
+        self.original_processor = original_processor
+        self.name = name
+        self.img0_dict = img0_dict
+        self.img1_dict = img1_dict
+        self.alpha = alpha
+        self.beta = beta
+        self.lamd = lamd
+        self.id = 0
+
+    def parent_call(
+        self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, scale=1.0
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+        attention_mask = attn.prepare_attention_mask(
+            attention_mask, sequence_length, batch_size)
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(
+                hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states) + scale * \
+            self.original_processor.to_q_lora(hidden_states)
+        query = attn.head_to_batch_dim(query)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(
+                encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states) + scale * \
+            self.original_processor.to_k_lora(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states) + scale * \
+            self.original_processor.to_v_lora(encoder_hidden_states)
+
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(
+            query, key, attention_mask)
+        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) + scale * self.original_processor.to_out_lora(hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(
+                -1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+    def __call__(self, attn, hidden_states, *args, encoder_hidden_states=None, attention_mask=None, **kwargs):
+        # Is self attention
+        if encoder_hidden_states is None:
+            # hardcode timestep
+            if self.id < 50 * self.lamd:
+                map0 = self.img0_dict[self.name][self.id]
+                map1 = self.img1_dict[self.name][self.id]
+                cross_map = self.beta * hidden_states + \
+                    (1 - self.beta) * ((1 - self.alpha) * map0 + self.alpha * map1)
+                # cross_map = self.beta * hidden_states + \
+                #     (1 - self.beta) * slerp(map0, map1, self.alpha)
+                # cross_map = slerp(slerp(map0, map1, self.alpha),
+                #                   hidden_states, self.beta)
+                # cross_map = hidden_states
+                # cross_map = torch.cat(
+                #     ((1 - self.alpha) * map0, self.alpha * map1), dim=1)
+
+                # res = self.original_processor(attn, hidden_states, *args,
+                #                               encoder_hidden_states=cross_map,
+                #                               attention_mask=attention_mask,
+                #                               temb=temb, **kwargs)
+                res = self.parent_call(attn, hidden_states, *args,
+                                       encoder_hidden_states=cross_map,
+                                       attention_mask=attention_mask,
+                                       **kwargs)
+            else:
+                res = self.original_processor(attn, hidden_states, *args,
+                                              encoder_hidden_states=encoder_hidden_states,
+                                              attention_mask=attention_mask,
+                                              **kwargs)
+
+            self.id += 1
+            # if self.id == len(self.img0_dict[self.name]):
+            if self.id == len(self.img0_dict[self.name]):
+                self.id = 0
+        else:
+            res = self.original_processor(attn, hidden_states, *args,
+                                          encoder_hidden_states=encoder_hidden_states,
+                                          attention_mask=attention_mask,
+                                          **kwargs)
+
+        return res
+
+
+class DiffMorpherPipeline(StableDiffusionPipeline):
+
+    def __init__(self,
+                 vae: AutoencoderKL,
+                 text_encoder: CLIPTextModel,
+                 tokenizer: CLIPTokenizer,
+                 unet: UNet2DConditionModel,
+                 scheduler: KarrasDiffusionSchedulers,
+                 safety_checker: StableDiffusionSafetyChecker,
+                 feature_extractor: CLIPImageProcessor,
+                 requires_safety_checker: bool = True,
+                ):
+        super().__init__(vae, text_encoder, tokenizer, unet, scheduler,
+                         safety_checker, feature_extractor, requires_safety_checker)
+        self.img0_dict = dict()
+        self.img1_dict = dict()
+
+    def inv_step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        x: torch.FloatTensor,
+        eta=0.,
+        verbose=False
+    ):
+        """
+        Inverse sampling for DDIM Inversion
+        """
+        if verbose:
+            print("timestep: ", timestep)
+        next_step = timestep
+        timestep = min(timestep - self.scheduler.config.num_train_timesteps //
+                       self.scheduler.num_inference_steps, 999)
+        alpha_prod_t = self.scheduler.alphas_cumprod[
+            timestep] if timestep >= 0 else self.scheduler.final_alpha_cumprod
+        alpha_prod_t_next = self.scheduler.alphas_cumprod[next_step]
+        beta_prod_t = 1 - alpha_prod_t
+        pred_x0 = (x - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5
+        pred_dir = (1 - alpha_prod_t_next)**0.5 * model_output
+        x_next = alpha_prod_t_next**0.5 * pred_x0 + pred_dir
+        return x_next, pred_x0
+
+    @torch.no_grad()
+    def invert(
+            self,
+            image: torch.Tensor,
+            prompt,
+            num_inference_steps=50,
+            num_actual_inference_steps=None,
+            guidance_scale=1.,
+            eta=0.0,
+            **kwds):
+        """
+        invert a real image into noise map with determinisc DDIM inversion
+        """
+        DEVICE = torch.device(
+            "cuda") if torch.cuda.is_available() else torch.device("cpu")
+        batch_size = image.shape[0]
+        if isinstance(prompt, list):
+            if batch_size == 1:
+                image = image.expand(len(prompt), -1, -1, -1)
+        elif isinstance(prompt, str):
+            if batch_size > 1:
+                prompt = [prompt] * batch_size
+
+        # text embeddings
+        text_input = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=77,
+            return_tensors="pt"
+        )
+        text_embeddings = self.text_encoder(text_input.input_ids.to(DEVICE))[0]
+        print("input text embeddings :", text_embeddings.shape)
+        # define initial latents
+        latents = self.image2latent(image)
+
+        # unconditional embedding for classifier free guidance
+        if guidance_scale > 1.:
+            max_length = text_input.input_ids.shape[-1]
+            unconditional_input = self.tokenizer(
+                [""] * batch_size,
+                padding="max_length",
+                max_length=77,
+                return_tensors="pt"
+            )
+            unconditional_embeddings = self.text_encoder(
+                unconditional_input.input_ids.to(DEVICE))[0]
+            text_embeddings = torch.cat(
+                [unconditional_embeddings, text_embeddings], dim=0)
+
+        print("latents shape: ", latents.shape)
+        # interative sampling
+        self.scheduler.set_timesteps(num_inference_steps)
+        print("Valid timesteps: ", reversed(self.scheduler.timesteps))
+        # print("attributes: ", self.scheduler.__dict__)
+        latents_list = [latents]
+        pred_x0_list = [latents]
+        for i, t in enumerate(tqdm.tqdm(reversed(self.scheduler.timesteps), desc="DDIM Inversion")):
+            if num_actual_inference_steps is not None and i >= num_actual_inference_steps:
+                continue
+
+            if guidance_scale > 1.:
+                model_inputs = torch.cat([latents] * 2)
+            else:
+                model_inputs = latents
+
+            # predict the noise
+            noise_pred = self.unet(
+                model_inputs, t, encoder_hidden_states=text_embeddings).sample
+            if guidance_scale > 1.:
+                noise_pred_uncon, noise_pred_con = noise_pred.chunk(2, dim=0)
+                noise_pred = noise_pred_uncon + guidance_scale * \
+                    (noise_pred_con - noise_pred_uncon)
+            # compute the previous noise sample x_t-1 -> x_t
+            latents, pred_x0 = self.inv_step(noise_pred, t, latents)
+            latents_list.append(latents)
+            pred_x0_list.append(pred_x0)
+
+        return latents
+
+    @torch.no_grad()
+    def ddim_inversion(self, latent, cond):
+        timesteps = reversed(self.scheduler.timesteps)
+        with torch.autocast(device_type='cuda', dtype=torch.float32):
+            for i, t in enumerate(tqdm.tqdm(timesteps, desc="DDIM inversion")):
+                cond_batch = cond.repeat(latent.shape[0], 1, 1)
+
+                alpha_prod_t = self.scheduler.alphas_cumprod[t]
+                alpha_prod_t_prev = (
+                    self.scheduler.alphas_cumprod[timesteps[i - 1]]
+                    if i > 0 else self.scheduler.final_alpha_cumprod
+                )
+
+                mu = alpha_prod_t ** 0.5
+                mu_prev = alpha_prod_t_prev ** 0.5
+                sigma = (1 - alpha_prod_t) ** 0.5
+                sigma_prev = (1 - alpha_prod_t_prev) ** 0.5
+
+                eps = self.unet(
+                    latent, t, encoder_hidden_states=cond_batch).sample
+
+                pred_x0 = (latent - sigma_prev * eps) / mu_prev
+                latent = mu * pred_x0 + sigma * eps
+        #         if save_latents:
+        #             torch.save(latent, os.path.join(save_path, f'noisy_latents_{t}.pt'))
+        # torch.save(latent, os.path.join(save_path, f'noisy_latents_{t}.pt'))
+        return latent
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        x: torch.FloatTensor,
+    ):
+        """
+        predict the sample of the next step in the denoise process.
+        """
+        prev_timestep = timestep - \
+            self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
+        alpha_prod_t = self.scheduler.alphas_cumprod[timestep]
+        alpha_prod_t_prev = self.scheduler.alphas_cumprod[
+            prev_timestep] if prev_timestep > 0 else self.scheduler.final_alpha_cumprod
+        beta_prod_t = 1 - alpha_prod_t
+        pred_x0 = (x - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5
+        pred_dir = (1 - alpha_prod_t_prev)**0.5 * model_output
+        x_prev = alpha_prod_t_prev**0.5 * pred_x0 + pred_dir
+        return x_prev, pred_x0
+
+    @torch.no_grad()
+    def image2latent(self, image):
+        DEVICE = torch.device(
+            "cuda") if torch.cuda.is_available() else torch.device("cpu")
+        if type(image) is Image:
+            image = np.array(image)
+            image = torch.from_numpy(image).float() / 127.5 - 1
+            image = image.permute(2, 0, 1).unsqueeze(0)
+        # input image density range [-1, 1]
+        latents = self.vae.encode(image.to(DEVICE))['latent_dist'].mean
+        latents = latents * 0.18215
+        return latents
+
+    @torch.no_grad()
+    def latent2image(self, latents, return_type='np'):
+        latents = 1 / 0.18215 * latents.detach()
+        image = self.vae.decode(latents)['sample']
+        if return_type == 'np':
+            image = (image / 2 + 0.5).clamp(0, 1)
+            image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
+            image = (image * 255).astype(np.uint8)
+        elif return_type == "pt":
+            image = (image / 2 + 0.5).clamp(0, 1)
+
+        return image
+
+    def latent2image_grad(self, latents):
+        latents = 1 / 0.18215 * latents
+        image = self.vae.decode(latents)['sample']
+
+        return image  # range [-1, 1]
+
+    @torch.no_grad()
+    def cal_latent(self, num_inference_steps, guidance_scale, unconditioning, img_noise_0, img_noise_1, text_embeddings_0, text_embeddings_1, lora_0, lora_1, alpha, use_lora, fix_lora=None):
+        # latents = torch.cos(alpha * torch.pi / 2) * img_noise_0 + \
+        #     torch.sin(alpha * torch.pi / 2) * img_noise_1
+        # latents = (1 - alpha) * img_noise_0 + alpha * img_noise_1
+        # latents = latents / ((1 - alpha) ** 2 + alpha ** 2)
+        latents = slerp(img_noise_0, img_noise_1, alpha, self.use_adain)
+        text_embeddings = (1 - alpha) * text_embeddings_0 + \
+            alpha * text_embeddings_1
+
+        self.scheduler.set_timesteps(num_inference_steps)
+        if use_lora:
+            if fix_lora is not None:
+                self.unet = load_lora(self.unet, lora_0, lora_1, fix_lora)
+            else:
+                self.unet = load_lora(self.unet, lora_0, lora_1, alpha)
+
+        for i, t in enumerate(tqdm.tqdm(self.scheduler.timesteps, desc=f"DDIM Sampler, alpha={alpha}")):
+
+            if guidance_scale > 1.:
+                model_inputs = torch.cat([latents] * 2)
+            else:
+                model_inputs = latents
+            if unconditioning is not None and isinstance(unconditioning, list):
+                _, text_embeddings = text_embeddings.chunk(2)
+                text_embeddings = torch.cat(
+                    [unconditioning[i].expand(*text_embeddings.shape), text_embeddings])
+            # predict the noise
+            noise_pred = self.unet(
+                model_inputs, t, encoder_hidden_states=text_embeddings).sample
+            if guidance_scale > 1.0:
+                noise_pred_uncon, noise_pred_con = noise_pred.chunk(
+                    2, dim=0)
+                noise_pred = noise_pred_uncon + guidance_scale * \
+                    (noise_pred_con - noise_pred_uncon)
+            # compute the previous noise sample x_t -> x_t-1
+            latents = self.scheduler.step(
+                noise_pred, t, latents, return_dict=False)[0]
+        return latents
+
+    @torch.no_grad()
+    def get_text_embeddings(self, prompt, guidance_scale, neg_prompt, batch_size):
+        DEVICE = torch.device(
+            "cuda") if torch.cuda.is_available() else torch.device("cpu")
+        # text embeddings
+        text_input = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=77,
+            return_tensors="pt"
+        )
+        text_embeddings = self.text_encoder(text_input.input_ids.cuda())[0]
+
+        if guidance_scale > 1.:
+            if neg_prompt:
+                uc_text = neg_prompt
+            else:
+                uc_text = ""
+            unconditional_input = self.tokenizer(
+                [uc_text] * batch_size,
+                padding="max_length",
+                max_length=77,
+                return_tensors="pt"
+            )
+            unconditional_embeddings = self.text_encoder(
+                unconditional_input.input_ids.to(DEVICE))[0]
+            text_embeddings = torch.cat(
+                [unconditional_embeddings, text_embeddings], dim=0)
+
+        return text_embeddings
+
+    def __call__(
+            self,
+            img_0=None,
+            img_1=None,
+            img_path_0=None,
+            img_path_1=None,
+            prompt_0="",
+            prompt_1="",
+            imgs=[],
+            img_paths=None,
+            prompts=None,
+            save_lora_dir="./lora",
+            load_lora_path_0=None,
+            load_lora_path_1=None,
+            load_lora_paths=None,
+            lora_steps=200,
+            lora_lr=2e-4,
+            lora_rank=16,
+            batch_size=1,
+            height=512,
+            width=512,
+            num_inference_steps=50,
+            num_actual_inference_steps=None,
+            guidance_scale=1,
+            attn_beta=0,
+            lamd=0.6,
+            use_lora=True,
+            use_adain=True,
+            use_reschedule=True,
+            output_path = "./results",
+            num_frames=50,
+            fix_lora=None,
+            progress=tqdm,
+            unconditioning=None,
+            neg_prompt=None,
+            save_intermediates=False,
+            **kwds):
+
+
+        self.scheduler.set_timesteps(num_inference_steps)
+        self.use_lora = use_lora
+        self.use_adain = use_adain
+        self.use_reschedule = use_reschedule
+        self.output_path = output_path
+
+
+        imgs = [Image.open(img_path).convert("RGB") for img_path in img_paths]
+        assert len(prompts) == len(imgs)
+        
+        # if img_path_0 or img_0:
+        #     img_paths = [img_path_0, img_path_1]
+        #     prompts = [prompt_0, prompt_1]
+        #     load_lora_paths = [load_lora_path_0, load_lora_path_1]
+        
+        # if img_0:
+        #     imgs.append(Image.fromarray(img_0).convert("RGB"))
+        # if img_1:
+        #     imgs.append(Image.fromarray(img_1).convert("RGB"))
+        # if imgs is None:
+        #     imgs = [Image.open(img_path).convert("RGB") for img_path in img_paths]
+        # if len(prompts) < len(imgs):
+        #     prompts += ["" for _ in range(len(imgs) - len(prompts))]
+
+        if self.use_lora:
+            loras = []
+            print("Loading lora...")
+            for i, (img, prompt) in enumerate(zip(imgs, prompts)):
+                if len(load_lora_paths) == i:
+
+                    weight_name = f"{output_path.split('/')[-1]}_lora_{i}.ckpt"
+                    load_lora_paths.append(save_lora_dir + "/" + weight_name)
+                    if not os.path.exists(load_lora_paths[i]):
+                        train_lora(img, prompt, save_lora_dir, None, self.tokenizer, self.text_encoder,
+                                self.vae, self.unet, self.scheduler, lora_steps, lora_lr, lora_rank, weight_name=weight_name)
+                print(f"Load from {load_lora_paths[i]}.")
+                if load_lora_paths[i].endswith(".safetensors"):
+                    loras.append(safetensors.torch.load_file(
+                        load_lora_paths[i], device="cpu"))
+                else:
+                    loras.append(torch.load(load_lora_paths[i], map_location="cpu"))
+
+            # if not load_lora_path_1:
+            #     weight_name = f"{output_path.split('/')[-1]}_lora_1.ckpt"
+            #     load_lora_path_1 = save_lora_dir + "/" + weight_name
+            #     if not os.path.exists(load_lora_path_1):
+            #         train_lora(img_1, prompt_1, save_lora_dir, None, self.tokenizer, self.text_encoder,
+            #                    self.vae, self.unet, self.scheduler, lora_steps, lora_lr, lora_rank, weight_name=weight_name)
+            # print(f"Load from {load_lora_path_1}.")
+            # if load_lora_path_1.endswith(".safetensors"):
+            #     lora_1 = safetensors.torch.load_file(
+            #         load_lora_path_1, device="cpu")
+            # else:
+            #     lora_1 = torch.load(load_lora_path_1, map_location="cpu")
+                    
+        def morph(alpha_list, progress, desc, img_noise_0, img_noise_1, text_embeddings_0, text_embeddings_1, lora_0, lora_1):
+            images = []
+            if attn_beta is not None:
+
+                self.unet = load_lora(self.unet, lora_0, lora_1, 0 if fix_lora is None else fix_lora)
+                attn_processor_dict = {}
+                for k in self.unet.attn_processors.keys():
+                    if do_replace_attn(k):
+                        attn_processor_dict[k] = StoreProcessor(self.unet.attn_processors[k],
+                                                                self.img0_dict, k)
+                    else:
+                        attn_processor_dict[k] = self.unet.attn_processors[k]
+                self.unet.set_attn_processor(attn_processor_dict)
+
+                latents = self.cal_latent(
+                    num_inference_steps,
+                    guidance_scale,
+                    unconditioning,
+                    img_noise_0,
+                    img_noise_1,
+                    text_embeddings_0,
+                    text_embeddings_1,
+                    lora_0,
+                    lora_1,
+                    alpha_list[0],
+                    False,
+                    fix_lora
+                )
+                first_image = self.latent2image(latents)
+                first_image = Image.fromarray(first_image)
+                # if save_intermediates:
+                #     first_image.save(f"{self.output_path}/{0:02d}.png")
+
+                self.unet = load_lora(self.unet, lora_0, lora_1, 1 if fix_lora is None else fix_lora)
+                attn_processor_dict = {}
+                for k in self.unet.attn_processors.keys():
+                    if do_replace_attn(k):
+                        attn_processor_dict[k] = StoreProcessor(self.unet.attn_processors[k],
+                                                                self.img1_dict, k)
+                    else:
+                        attn_processor_dict[k] = self.unet.attn_processors[k]
+
+                self.unet.set_attn_processor(attn_processor_dict)
+
+                latents = self.cal_latent(
+                    num_inference_steps,
+                    guidance_scale,
+                    unconditioning,
+                    img_noise_0,
+                    img_noise_1,
+                    text_embeddings_0,
+                    text_embeddings_1,
+                    lora_0,
+                    lora_1,
+                    alpha_list[-1], 
+                    False,
+                    fix_lora
+                )
+                last_image = self.latent2image(latents)
+                last_image = Image.fromarray(last_image)
+                # if save_intermediates:
+                #     last_image.save(
+                #         f"{self.output_path}/{num_frames - 1:02d}.png")
+
+                for i in progress.tqdm(range(1, num_frames - 1), desc=desc):
+                    alpha = alpha_list[i]
+                    self.unet = load_lora(self.unet, lora_0, lora_1, alpha if fix_lora is None else fix_lora)
+                    attn_processor_dict = {}
+                    for k in self.unet.attn_processors.keys():
+                        if do_replace_attn(k):
+                            attn_processor_dict[k] = LoadProcessor(
+                                self.unet.attn_processors[k], k, self.img0_dict, self.img1_dict, alpha, attn_beta, lamd)
+                        else:
+                            attn_processor_dict[k] = self.unet.attn_processors[k]
+
+                    self.unet.set_attn_processor(attn_processor_dict)
+
+                    latents = self.cal_latent(
+                        num_inference_steps,
+                        guidance_scale,
+                        unconditioning,
+                        img_noise_0,
+                        img_noise_1,
+                        text_embeddings_0,
+                        text_embeddings_1,
+                        lora_0,
+                        lora_1,
+                        alpha_list[i], 
+                        False,
+                        fix_lora
+                    )
+                    image = self.latent2image(latents)
+                    image = Image.fromarray(image)
+                    # if save_intermediates:
+                    #     image.save(f"{self.output_path}/{i:02d}.png")
+                    images.append(image)
+
+                images = [first_image] + images + [last_image]
+
+            else:
+                for k, alpha in enumerate(alpha_list):
+
+                    latents = self.cal_latent(
+                        num_inference_steps,
+                        guidance_scale,
+                        unconditioning,
+                        img_noise_0,
+                        img_noise_1,
+                        text_embeddings_0,
+                        text_embeddings_1,
+                        lora_0,
+                        lora_1,
+                        alpha_list[k], 
+                        self.use_lora,
+                        fix_lora
+                    )
+                    image = self.latent2image(latents)
+                    image = Image.fromarray(image)
+                    # if save_intermediates:
+                    #     image.save(f"{self.output_path}/{k:02d}.png")
+                    images.append(image)
+
+            return images
+
+        images = []
+        
+        for img_0, img_1, prompt_0, prompt_1, lora_0, lora_1 in zip(imgs[:-1], imgs[1:], prompts[:-1], prompts[1:], loras[:-1], loras[1:]):
+            text_embeddings_0 = self.get_text_embeddings(
+                prompt_0, guidance_scale, neg_prompt, batch_size)
+            text_embeddings_1 = self.get_text_embeddings(
+                prompt_1, guidance_scale, neg_prompt, batch_size)
+            img_0 = get_img(img_0)
+            img_1 = get_img(img_1)
+            if self.use_lora:
+                self.unet = load_lora(self.unet, lora_0, lora_1, 0)
+            img_noise_0 = self.ddim_inversion(
+                self.image2latent(img_0), text_embeddings_0)
+            if self.use_lora:
+                self.unet = load_lora(self.unet, lora_0, lora_1, 1)
+            img_noise_1 = self.ddim_inversion(
+                self.image2latent(img_1), text_embeddings_1)
+
+            print("latents shape: ", img_noise_0.shape)
+
+            with torch.no_grad():
+                if self.use_reschedule:
+                    alpha_scheduler = AlphaScheduler()
+                    alpha_list = list(torch.linspace(0, 1, num_frames))
+                    images_pt = morph(alpha_list, progress, "Sampling...", img_noise_0, img_noise_1, text_embeddings_0, text_embeddings_1, lora_0, lora_1)
+                    images_pt = [transforms.ToTensor()(img).unsqueeze(0)
+                                for img in images_pt]
+                    alpha_scheduler.from_imgs(images_pt)
+                    alpha_list = alpha_scheduler.get_list()
+                    print(alpha_list)
+                    images_ = morph(alpha_list, progress, "Reschedule...", img_noise_0, img_noise_1, text_embeddings_0, text_embeddings_1, lora_0, lora_1)
+                else:
+                    alpha_list = list(torch.linspace(0, 1, num_frames))
+                    print(alpha_list)
+                    images_ = morph(alpha_list, progress, "Sampling...",  img_noise_0, img_noise_1, text_embeddings_0, text_embeddings_1, lora_0, lora_1)
+
+                if len(images) == 0:
+                    images = images_
+                else:
+                    images += images_[1:]
+
+        if save_intermediates:
+            for i, image in enumerate(images):
+                image.save(f"{self.output_path}/{i:02d}.png")
+
+        return images

requirements.txt

@@ -0,0 +1,14 @@
+accelerate==0.23.0
+diffusers==0.17.1
+einops==0.7.0
+gradio==4.7.1
+numpy==1.26.1
+opencv_python==4.5.5.64
+packaging==23.2
+Pillow==10.1.0
+safetensors==0.4.0
+tqdm==4.65.0
+transformers==4.34.1
+torch
+torchvision
+lpips