Add Image-Morpher as folder, without assets.

Image-Morpher/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.
+ 

Image-Morpher/README.md

@@ -0,0 +1,133 @@
+# to be edited. 
+
+<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}
+}
+```

Image-Morpher/lcm_lora/lcm_lora_test.py

@@ -0,0 +1,67 @@
+import torch
+# from diffusers import AutoPipelineForText2Image
+from lcm_schedule import LCMScheduler
+from diffusers import DiffusionPipeline
+
+# Load the SD v1.5 base model
+pipe = DiffusionPipeline.from_pretrained(
+    # "runwayml/stable-diffusion-v1-5",
+    "Lykon/dreamshaper-7", # seems to produce better results, its a fine-tuned sd1-5
+    torch_dtype=torch.float16,
+    # variant="fp16", # Removing variant since torch_dtype is set
+).to("cuda")
+
+# Set the LCMScheduler
+pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+
+# Load LCM LoRA for SD v1.5
+pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
+# pipe.fuse_lora() # my ver of Stable Diffusion Pipeline above does not have this...
+
+# Ensure the pipeline is on CUDA with the proper dtype
+# pipe.to(device="cuda", dtype=torch.float16) # Removed device argument
+
+prompt = "A hyperrealistic portrait of a cat, highly detailed, cute pet"
+negative_prompt = "blurry, low quality"
+generator = torch.manual_seed(0)  # for reproducibility
+
+# Run inference (using very few steps as typical for LCM accelerated inference)
+image = pipe(prompt, 
+             negative_prompt=negative_prompt, # this works well for dreamshaper ! best output is with this, num_inference = 8, and guidance =0 on dreamshaper
+             num_inference_steps=8, # 8 seems to work best for both... but finetune it. 
+             guidance_scale=0, # the dreamshaper model works better with guidance_scale = 0!! disabled essentially. 
+             generator=generator).images[0]
+
+# Save the resulting image to disk
+image.save("lcm_accelerated_cat3.png")
+
+# print(pipe.scheduler)
+
+
+
+
+
+
+
+
+
+
+
+
+
+# model_id = "Lykon/dreamshaper-7"
+# adapter_id = "latent-consistency/lcm-lora-sdv1-5"
+
+# pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16")
+# pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+# pipe.to("cuda")
+
+# # load and fuse lcm lora
+# pipe.load_lora_weights(adapter_id)
+# pipe.fuse_lora()
+
+
+# prompt = "Self-portrait oil painting, a beautiful cyborg with golden hair, 8k"
+
+# # disable guidance_scale by passing 0
+# image = pipe(prompt=prompt, num_inference_steps=4, guidance_scale=0).images[0]

Image-Morpher/lcm_lora/lcm_schedule.py

@@ -0,0 +1,654 @@
+# Copyright 2024 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 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, logging
+# from diffusers.utils.torch_utils import randn_tensor
+# from diffusers.scheduling_utils import SchedulerMixin
+from diffusers.schedulers.scheduling_ddim import *
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class LCMSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.Tensor` 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.Tensor` 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.Tensor
+    denoised: Optional[torch.Tensor] = None
+
+
+# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
+def betas_for_alpha_bar(
+    num_diffusion_timesteps,
+    max_beta=0.999,
+    alpha_transform_type="cosine",
+):
+    """
+    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.
+        alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
+                     Choose from `cosine` or `exp`
+
+    Returns:
+        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
+    """
+    if alpha_transform_type == "cosine":
+
+        def alpha_bar_fn(t):
+            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+
+    elif alpha_transform_type == "exp":
+
+        def alpha_bar_fn(t):
+            return math.exp(t * -12.0)
+
+    else:
+        raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
+
+    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_fn(t2) / alpha_bar_fn(t1), max_beta))
+    return torch.tensor(betas, dtype=torch.float32)
+
+
+# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr
+def rescale_zero_terminal_snr(betas: torch.Tensor) -> torch.Tensor:
+    """
+    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
+
+
+    Args:
+        betas (`torch.Tensor`):
+            the betas that the scheduler is being initialized with.
+
+    Returns:
+        `torch.Tensor`: rescaled betas with zero terminal SNR
+    """
+    # Convert betas to alphas_bar_sqrt
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+    alphas_bar_sqrt = alphas_cumprod.sqrt()
+
+    # Store old values.
+    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
+    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
+
+    # Shift so the last timestep is zero.
+    alphas_bar_sqrt -= alphas_bar_sqrt_T
+
+    # Scale so the first timestep is back to the old value.
+    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+
+    # Convert alphas_bar_sqrt to betas
+    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
+    alphas = alphas_bar[1:] / alphas_bar[:-1]  # Revert cumprod
+    alphas = torch.cat([alphas_bar[0:1], alphas])
+    betas = 1 - alphas
+
+    return betas
+
+
+class LCMScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `LCMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with
+    non-Markovian guidance.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. [`~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.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        beta_start (`float`, defaults to 0.0001):
+            The starting `beta` value of inference.
+        beta_end (`float`, defaults to 0.02):
+            The final `beta` value.
+        beta_schedule (`str`, defaults to `"linear"`):
+            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*):
+            Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
+        original_inference_steps (`int`, *optional*, defaults to 50):
+            The default number of inference steps used to generate a linearly-spaced timestep schedule, from which we
+            will ultimately take `num_inference_steps` evenly spaced timesteps to form the final timestep schedule.
+        clip_sample (`bool`, defaults to `True`):
+            Clip the predicted sample for numerical stability.
+        clip_sample_range (`float`, defaults to 1.0):
+            The maximum magnitude for sample clipping. Valid only when `clip_sample=True`.
+        set_alpha_to_one (`bool`, defaults to `True`):
+            Each diffusion step uses the alphas product value 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 alpha value at step 0.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps, as required by some model families.
+        prediction_type (`str`, defaults to `epsilon`, *optional*):
+            Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
+            `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
+            Video](https://imagen.research.google/video/paper.pdf) paper).
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+            as Stable Diffusion.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True`.
+        timestep_spacing (`str`, defaults to `"leading"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        timestep_scaling (`float`, defaults to 10.0):
+            The factor the timesteps will be multiplied by when calculating the consistency model boundary conditions
+            `c_skip` and `c_out`. Increasing this will decrease the approximation error (although the approximation
+            error at the default of `10.0` is already pretty small).
+        rescale_betas_zero_snr (`bool`, defaults to `False`):
+            Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
+            dark samples instead of limiting it to samples with medium brightness. Loosely related to
+            [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
+    """
+
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.00085,
+        beta_end: float = 0.012,
+        beta_schedule: str = "scaled_linear",
+        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        original_inference_steps: int = 50,
+        clip_sample: bool = False,
+        clip_sample_range: float = 1.0,
+        set_alpha_to_one: bool = True,
+        steps_offset: int = 0,
+        prediction_type: str = "epsilon",
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        sample_max_value: float = 1.0,
+        timestep_spacing: str = "leading",
+        timestep_scaling: float = 10.0,
+        rescale_betas_zero_snr: bool = False,
+    ):
+        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} is not implemented for {self.__class__}")
+
+        # Rescale for zero SNR
+        if rescale_betas_zero_snr:
+            self.betas = rescale_zero_terminal_snr(self.betas)
+
+        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))
+        self.custom_timesteps = False
+
+        self._step_index = None
+        self._begin_index = None
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
+    def _init_step_index(self, timestep):
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    @property
+    def step_index(self):
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.Tensor`):
+                The input sample.
+            timestep (`int`, *optional*):
+                The current timestep in the diffusion chain.
+        Returns:
+            `torch.Tensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+
+        return sample
+
+    def set_timesteps(
+        self,
+        num_inference_steps: Optional[int] = None,
+        device: Union[str, torch.device] = None,
+        original_inference_steps: Optional[int] = None,
+        timesteps: Optional[List[int]] = None,
+        strength: int = 1.0,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`, *optional*):
+                The number of diffusion steps used when generating samples with a pre-trained model. If used,
+                `timesteps` must be `None`.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            original_inference_steps (`int`, *optional*):
+                The original number of inference steps, which will be used to generate a linearly-spaced timestep
+                schedule (which is different from the standard `diffusers` implementation). We will then take
+                `num_inference_steps` timesteps from this schedule, evenly spaced in terms of indices, and use that as
+                our final timestep schedule. If not set, this will default to the `original_inference_steps` attribute.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
+                timestep spacing strategy of equal spacing between timesteps on the training/distillation timestep
+                schedule is used. If `timesteps` is passed, `num_inference_steps` must be `None`.
+        """
+        # 0. Check inputs
+        if num_inference_steps is None and timesteps is None:
+            raise ValueError("Must pass exactly one of `num_inference_steps` or `custom_timesteps`.")
+
+        if num_inference_steps is not None and timesteps is not None:
+            raise ValueError("Can only pass one of `num_inference_steps` or `custom_timesteps`.")
+
+        # 1. Calculate the LCM original training/distillation timestep schedule.
+        original_steps = (
+            original_inference_steps if original_inference_steps is not None else self.config.original_inference_steps
+        )
+
+        if original_steps > self.config.num_train_timesteps:
+            raise ValueError(
+                f"`original_steps`: {original_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."
+            )
+
+        # LCM Timesteps Setting
+        # The skipping step parameter k from the paper.
+        k = self.config.num_train_timesteps // original_steps
+        # LCM Training/Distillation Steps Schedule
+        # Currently, only a linearly-spaced schedule is supported (same as in the LCM distillation scripts).
+        lcm_origin_timesteps = np.asarray(list(range(1, int(original_steps * strength) + 1))) * k - 1
+
+        # 2. Calculate the LCM inference timestep schedule.
+        if timesteps is not None:
+            # 2.1 Handle custom timestep schedules.
+            train_timesteps = set(lcm_origin_timesteps)
+            non_train_timesteps = []
+            for i in range(1, len(timesteps)):
+                if timesteps[i] >= timesteps[i - 1]:
+                    raise ValueError("`custom_timesteps` must be in descending order.")
+
+                if timesteps[i] not in train_timesteps:
+                    non_train_timesteps.append(timesteps[i])
+
+            if timesteps[0] >= self.config.num_train_timesteps:
+                raise ValueError(
+                    f"`timesteps` must start before `self.config.train_timesteps`:"
+                    f" {self.config.num_train_timesteps}."
+                )
+
+            # Raise warning if timestep schedule does not start with self.config.num_train_timesteps - 1
+            if strength == 1.0 and timesteps[0] != self.config.num_train_timesteps - 1:
+                logger.warning(
+                    f"The first timestep on the custom timestep schedule is {timesteps[0]}, not"
+                    f" `self.config.num_train_timesteps - 1`: {self.config.num_train_timesteps - 1}. You may get"
+                    f" unexpected results when using this timestep schedule."
+                )
+
+            # Raise warning if custom timestep schedule contains timesteps not on original timestep schedule
+            if non_train_timesteps:
+                logger.warning(
+                    f"The custom timestep schedule contains the following timesteps which are not on the original"
+                    f" training/distillation timestep schedule: {non_train_timesteps}. You may get unexpected results"
+                    f" when using this timestep schedule."
+                )
+
+            # Raise warning if custom timestep schedule is longer than original_steps
+            if len(timesteps) > original_steps:
+                logger.warning(
+                    f"The number of timesteps in the custom timestep schedule is {len(timesteps)}, which exceeds the"
+                    f" the length of the timestep schedule used for training: {original_steps}. You may get some"
+                    f" unexpected results when using this timestep schedule."
+                )
+
+            timesteps = np.array(timesteps, dtype=np.int64)
+            self.num_inference_steps = len(timesteps)
+            self.custom_timesteps = True
+
+            # Apply strength (e.g. for img2img pipelines) (see StableDiffusionImg2ImgPipeline.get_timesteps)
+            init_timestep = min(int(self.num_inference_steps * strength), self.num_inference_steps)
+            t_start = max(self.num_inference_steps - init_timestep, 0)
+            timesteps = timesteps[t_start * self.order :]
+            # TODO: also reset self.num_inference_steps?
+        else:
+            # 2.2 Create the "standard" LCM inference timestep schedule.
+            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."
+                )
+
+            skipping_step = len(lcm_origin_timesteps) // num_inference_steps
+
+            if skipping_step < 1:
+                raise ValueError(
+                    f"The combination of `original_steps x strength`: {original_steps} x {strength} is smaller than `num_inference_steps`: {num_inference_steps}. Make sure to either reduce `num_inference_steps` to a value smaller than {int(original_steps * strength)} or increase `strength` to a value higher than {float(num_inference_steps / original_steps)}."
+                )
+
+            self.num_inference_steps = num_inference_steps
+
+            if num_inference_steps > original_steps:
+                raise ValueError(
+                    f"`num_inference_steps`: {num_inference_steps} cannot be larger than `original_inference_steps`:"
+                    f" {original_steps} because the final timestep schedule will be a subset of the"
+                    f" `original_inference_steps`-sized initial timestep schedule."
+                )
+
+            # LCM Inference Steps Schedule
+            lcm_origin_timesteps = lcm_origin_timesteps[::-1].copy()
+            # Select (approximately) evenly spaced indices from lcm_origin_timesteps.
+            inference_indices = np.linspace(0, len(lcm_origin_timesteps), num=num_inference_steps, endpoint=False)
+            inference_indices = np.floor(inference_indices).astype(np.int64)
+            timesteps = lcm_origin_timesteps[inference_indices]
+
+        self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.long)
+
+        self._step_index = None
+        self._begin_index = None
+
+    def get_scalings_for_boundary_condition_discrete(self, timestep):
+        self.sigma_data = 0.5  # Default: 0.5
+        scaled_timestep = timestep * self.config.timestep_scaling
+
+        c_skip = self.sigma_data**2 / (scaled_timestep**2 + self.sigma_data**2)
+        c_out = scaled_timestep / (scaled_timestep**2 + self.sigma_data**2) ** 0.5
+        return c_skip, c_out
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: int,
+        sample: torch.Tensor,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+    ) -> Union[LCMSchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`.
+        Returns:
+            [`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # 1. get previous step value
+        prev_step_index = self.step_index + 1
+        if prev_step_index < len(self.timesteps):
+            prev_timestep = self.timesteps[prev_step_index]
+        else:
+            prev_timestep = timestep
+
+        # 2. compute alphas, betas
+        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
+
+        # 3. Get scalings for boundary conditions
+        c_skip, c_out = self.get_scalings_for_boundary_condition_discrete(timestep)
+
+        # 4. Compute the predicted original sample x_0 based on the model parameterization
+        if self.config.prediction_type == "epsilon":  # noise-prediction
+            predicted_original_sample = (sample - beta_prod_t.sqrt() * model_output) / alpha_prod_t.sqrt()
+        elif self.config.prediction_type == "sample":  # x-prediction
+            predicted_original_sample = model_output
+        elif self.config.prediction_type == "v_prediction":  # v-prediction
+            predicted_original_sample = alpha_prod_t.sqrt() * sample - beta_prod_t.sqrt() * model_output
+        else:
+            raise ValueError(
+                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or"
+                " `v_prediction` for `LCMScheduler`."
+            )
+
+        # 5. Clip or threshold "predicted x_0"
+        if self.config.thresholding:
+            predicted_original_sample = self._threshold_sample(predicted_original_sample)
+        elif self.config.clip_sample:
+            predicted_original_sample = predicted_original_sample.clamp(
+                -self.config.clip_sample_range, self.config.clip_sample_range
+            )
+
+        # 6. Denoise model output using boundary conditions
+        denoised = c_out * predicted_original_sample + c_skip * sample
+
+        # 7. Sample and inject noise z ~ N(0, I) for MultiStep Inference
+        # Noise is not used on the final timestep of the timestep schedule.
+        # This also means that noise is not used for one-step sampling.
+        if self.step_index != self.num_inference_steps - 1:
+            noise = randn_tensor(
+                model_output.shape, generator=generator, device=model_output.device, dtype=denoised.dtype
+            )
+            prev_sample = alpha_prod_t_prev.sqrt() * denoised + beta_prod_t_prev.sqrt() * noise
+        else:
+            prev_sample = denoised
+
+        # upon completion increase step index by one
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample, denoised)
+
+        return LCMSchedulerOutput(prev_sample=prev_sample, denoised=denoised)
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.Tensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
+        # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement
+        # for the subsequent add_noise calls
+        self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device)
+        alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype)
+        timesteps = timesteps.to(original_samples.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[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
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity
+    def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as sample
+        self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device)
+        alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype)
+        timesteps = timesteps.to(sample.device)
+
+        sqrt_alpha_prod = 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 - 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
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.previous_timestep
+    def previous_timestep(self, timestep):
+        if self.custom_timesteps or self.num_inference_steps:
+            index = (self.timesteps == timestep).nonzero(as_tuple=True)[0][0]
+            if index == self.timesteps.shape[0] - 1:
+                prev_t = torch.tensor(-1)
+            else:
+                prev_t = self.timesteps[index + 1]
+        else:
+            prev_t = timestep - 1
+        return prev_t

Image-Morpher/logs/71seconds.log

@@ -0,0 +1,14 @@
+2025-03-09 17:59:17,957 - ERROR - Model loading failed: [Errno 2] No such file or directory: '/mnt/slurm_home/nalin/.cache/huggingface/hub/models--DarkFlameUniverse--Stable-Diffusion-2-1-Base-8bit/refs/main', switching to standard 2-1 base.
+2025-03-09 18:00:29,214 - INFO - Total execution time: 71.52 seconds
+2025-03-09 18:00:29,215 - INFO - Model Path: DarkFlameUniverse/Stable-Diffusion-2-1-Base-8bit
+2025-03-09 18:00:29,215 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 18:00:29,215 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 18:00:29,215 - INFO - Use LCM: True
+2025-03-09 18:00:29,215 - INFO - Number of inference steps: 8
+2025-03-09 18:00:29,215 - INFO - Guidance scale: 1
+
+71 seconds with standard 2-1.. and the vae attention slicing ?! and cudnn benchmark.
+
+lcm-lora works with the stabilityai/sd2-1 model card? 
+
+i think for this run the lora was already partially trained so it picked it up and trained the rest, so it was quicker

Image-Morpher/logs/97 seconds legitimate.log

@@ -0,0 +1,23 @@
+2025-03-09 23:23:30,516 - INFO - Total execution time: 97.10 seconds
+2025-03-09 23:23:30,516 - INFO - Model Path: stabilityai/stable-diffusion-2-1-base
+2025-03-09 23:23:30,516 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 23:23:30,516 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 23:23:30,516 - INFO - Use LCM: True
+2025-03-09 23:23:30,516 - INFO - Number of inference steps: 8
+2025-03-09 23:23:30,516 - INFO - Guidance scale: 1
+
+
+so first time under 100 seconds.
+
+
+this was with
+# Configure compiler settings (safe defaults)
+inductor.config.force_fuse_int_mm_with_mul = True  # Better for diffusion models
+
+and adding 
+# Adding these AFTER device movement, instead of before, so it can optimise convolutions for the specific GPU / cuda kernel
+torch.backends.cudnn.benchmark = True
+torch.set_float32_matmul_precision("high")  # Better for modern GPUs
+
+
+In all i was able to shave off about 20-30 seconds from the optimisations, from the lcm-lora it was about a minute (?) so 1.5min saved

Image-Morpher/logs/SD2-1 compatible, confirm on main.py.log

@@ -0,0 +1,10 @@
+2025-03-09 18:44:19,754 - INFO - Total execution time: 104.91 seconds
+2025-03-09 18:44:19,754 - INFO - Model Path: stabilityai/stable-diffusion-2-1-base
+2025-03-09 18:44:19,754 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 18:44:19,755 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 18:44:19,755 - INFO - Use LCM: True
+2025-03-09 18:44:19,755 - INFO - Number of inference steps: 8
+2025-03-09 18:44:19,755 - INFO - Guidance scale: 1
+
+
+confirm that sd2-1 is faster than sdv1-5 in training, saved about 15-20s here. 

Image-Morpher/logs/SD2-1,vae&atnnslicing,benchmark,lcm.log

@@ -0,0 +1,11 @@
+2025-03-09 18:05:43,787 - ERROR - Model loading failed: [Errno 2] No such file or directory: '/mnt/slurm_home/nalin/.cache/huggingface/hub/models--DarkFlameUniverse--Stable-Diffusion-2-1-Base-8bit/refs/main', switching to standard 2-1 base.
+2025-03-09 18:07:28,182 - INFO - Total execution time: 104.66 seconds
+2025-03-09 18:07:28,182 - INFO - Model Path: DarkFlameUniverse/Stable-Diffusion-2-1-Base-8bit
+2025-03-09 18:07:28,182 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 18:07:28,182 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 18:07:28,182 - INFO - Use LCM: True
+2025-03-09 18:07:28,182 - INFO - Number of inference steps: 8
+2025-03-09 18:07:28,182 - INFO - Guidance scale: 1
+
+yup it does seem that it switched to 2-1, its compatible with lcm-lora 1-5, and that the vae and attention slicing were ENABLED
+note that cudnn benchmark was also ENABLED

Image-Morpher/logs/best config, matmul precision /"high/" ++.log

@@ -0,0 +1,17 @@
+2025-03-09 23:57:53,227 - INFO - Total execution time: 97.25 seconds
+2025-03-09 23:57:53,227 - INFO - Model Path: stabilityai/stable-diffusion-2-1-base
+2025-03-09 23:57:53,227 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 23:57:53,227 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 23:57:53,227 - INFO - Use LCM: True
+2025-03-09 23:57:53,228 - INFO - Number of inference steps: 8
+2025-03-09 23:57:53,228 - INFO - Guidance scale: 1
+
+
+
+best result, this was with 
+torch.set_float32_matmul_precision("high")  # Better for modern GPUs
+
+it instructs cuda to not be so precise in intermediate computations, does some approximations, rounding that would not be crucial in a 
+deep learning workload..... on "medium" or "low" settings they have additional steps to ensure this precision apparently. 
+
+and this was particularly useful after moving pipeline.to("cuda"). 

Image-Morpher/logs/execution_20250302_224812.log

@@ -0,0 +1,9 @@
+2025-03-02 22:50:16,961 - INFO - Total execution time: 124.71 seconds
+2025-03-02 22:50:16,961 - INFO - Model Path: sd-legacy/stable-diffusion-v1-5
+2025-03-02 22:50:16,961 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-02 22:50:16,961 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-02 22:50:16,962 - INFO - Number of inference steps: 8
+2025-03-02 22:50:16,962 - INFO - Guidance scale: 1
+2025-03-02 22:50:16,962 - INFO - Use LCM: True
+
+base

Image-Morpher/logs/execution_20250302_225132.log

@@ -0,0 +1,9 @@
+2025-03-02 22:53:49,437 - INFO - Total execution time: 136.99 seconds
+2025-03-02 22:53:49,438 - INFO - Model Path: sd-legacy/stable-diffusion-v1-5
+2025-03-02 22:53:49,438 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-02 22:53:49,438 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-02 22:53:49,438 - INFO - Number of inference steps: 8
+2025-03-02 22:53:49,438 - INFO - Guidance scale: 1
+2025-03-02 22:53:49,438 - INFO - Use LCM: True
+
+cudnn benchamrk and precison 32

Image-Morpher/logs/execution_20250302_230601.log

@@ -0,0 +1,9 @@
+2025-03-02 23:08:10,022 - INFO - Total execution time: 128.30 seconds
+2025-03-02 23:08:10,022 - INFO - Model Path: sd-legacy/stable-diffusion-v1-5
+2025-03-02 23:08:10,022 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-02 23:08:10,022 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-02 23:08:10,022 - INFO - Number of inference steps: 8
+2025-03-02 23:08:10,022 - INFO - Guidance scale: 1
+2025-03-02 23:08:10,022 - INFO - Use LCM: True
+
+pipeline memory format changes, bad results

Image-Morpher/logs/execution_20250302_231103.log

@@ -0,0 +1,9 @@
+2025-03-02 23:13:22,730 - INFO - Total execution time: 139.31 seconds
+2025-03-02 23:13:22,758 - INFO - Model Path: sd-legacy/stable-diffusion-v1-5
+2025-03-02 23:13:22,758 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-02 23:13:22,758 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-02 23:13:22,758 - INFO - Number of inference steps: 8
+2025-03-02 23:13:22,758 - INFO - Guidance scale: 1
+2025-03-02 23:13:22,758 - INFO - Use LCM: True
+
+cudnn benchmark... doesnt help either

Image-Morpher/logs/execution_20250303_000722.log

@@ -0,0 +1,11 @@
+2025-03-03 00:08:10,065 - INFO - Total execution time: 47.30 seconds
+2025-03-03 00:08:10,065 - INFO - Model Path: sd-legacy/stable-diffusion-v1-5
+2025-03-03 00:08:10,065 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-03 00:08:10,065 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-03 00:08:10,065 - INFO - Use LCM: True
+2025-03-03 00:08:10,065 - INFO - Number of inference steps: 8
+2025-03-03 00:08:10,065 - INFO - Guidance scale: 1
+
+was with the other torch flags, but i think loras were loaded in alr
+
+

Image-Morpher/logs/execution_20250303_150757.log

@@ -0,0 +1,9 @@
+2025-03-03 15:10:09,361 - INFO - Total execution time: 131.64 seconds
+2025-03-03 15:10:09,362 - INFO - Model Path: sd-legacy/stable-diffusion-v1-5
+2025-03-03 15:10:09,362 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-03 15:10:09,362 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-03 15:10:09,362 - INFO - Use LCM: True
+2025-03-03 15:10:09,362 - INFO - Number of inference steps: 8
+2025-03-03 15:10:09,362 - INFO - Guidance scale: 1
+
+this was with the flags, no diff.

Image-Morpher/logs/execution_20250309_171159.log

@@ -0,0 +1,12 @@
+2025-03-09 17:14:01,299 - INFO - Total execution time: 121.91 seconds
+2025-03-09 17:14:01,299 - INFO - Model Path: sd-legacy/stable-diffusion-v1-5
+2025-03-09 17:14:01,299 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 17:14:01,299 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 17:14:01,299 - INFO - Use LCM: True
+2025-03-09 17:14:01,299 - INFO - Number of inference steps: 8
+2025-03-09 17:14:01,299 - INFO - Guidance scale: 1
+
+
+pipeline.enable_attention_slicing(1)
+
+extreme memory saving, not required tbh.

Image-Morpher/logs/execution_20250309_172206.log

@@ -0,0 +1,12 @@
+2025-03-09 17:28:03,402 - INFO - Total execution time: 357.21 seconds
+2025-03-09 17:28:03,402 - INFO - Model Path: sd-legacy/stable-diffusion-v1-5
+2025-03-09 17:28:03,402 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 17:28:03,402 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 17:28:03,403 - INFO - Use LCM: True
+2025-03-09 17:28:03,403 - INFO - Number of inference steps: 8
+2025-03-09 17:28:03,403 - INFO - Guidance scale: 1
+
+pipeline.unet = torch.compile(pipeline.unet)  # Optimize further
+torch.cuda.empty_cache()
+
+on args.use_lcm

Image-Morpher/logs/execution_20250309_173119.log

@@ -0,0 +1,12 @@
+2025-03-09 17:35:58,495 - INFO - Total execution time: 279.30 seconds
+2025-03-09 17:35:58,495 - INFO - Model Path: sd-legacy/stable-diffusion-v1-5
+2025-03-09 17:35:58,495 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 17:35:58,495 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 17:35:58,495 - INFO - Use LCM: True
+2025-03-09 17:35:58,495 - INFO - Number of inference steps: 8
+2025-03-09 17:35:58,496 - INFO - Guidance scale: 1
+
+pipeline.unet = torch.compile(pipeline.unet)  # Optimize further
+torch.cuda.empty_cache()
+
+after the initial pipeline loaded in. DiffMorpherPipeline

Image-Morpher/logs/execution_20250309_173628.log

@@ -0,0 +1,13 @@
+2025-03-09 17:41:02,525 - INFO - Total execution time: 273.81 seconds
+2025-03-09 17:41:02,526 - INFO - Model Path: sd-legacy/stable-diffusion-v1-5
+2025-03-09 17:41:02,526 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 17:41:02,526 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 17:41:02,526 - INFO - Use LCM: True
+2025-03-09 17:41:02,526 - INFO - Number of inference steps: 8
+2025-03-09 17:41:02,526 - INFO - Guidance scale: 1
+
+
+pipeline.unet = torch.compile(pipeline.unet)  # Optimize further
+# torch.cuda.empty_cache()
+
+without empty_cache, no help.

Image-Morpher/logs/execution_20250309_175711.log

@@ -0,0 +1 @@
+2025-03-09 17:57:11,357 - ERROR - Model loading failed: [Errno 2] No such file or directory: '/mnt/slurm_home/nalin/.cache/huggingface/hub/models--DarkFlameUniverse--Stable-Diffusion-2-1-Base-8bit/refs/main', switching to standard 2-1 base.

Image-Morpher/logs/execution_20250309_184725.log

@@ -0,0 +1,10 @@
+2025-03-09 18:49:26,905 - INFO - Total execution time: 121.72 seconds
+2025-03-09 18:49:26,906 - INFO - Model Path: sd-legacy/stable-diffusion-v1-5
+2025-03-09 18:49:26,906 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 18:49:26,906 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 18:49:26,906 - INFO - Use LCM: True
+2025-03-09 18:49:26,906 - INFO - Number of inference steps: 8
+2025-03-09 18:49:26,906 - INFO - Guidance scale: 1
+
+
+yup sdv1-5 slower and worse results (biden dun have teeth)

Image-Morpher/logs/execution_20250309_215757.log

@@ -0,0 +1,15 @@
+2025-03-09 21:57:58,103 - ERROR - Model loading failed: [Errno 2] No such file or directory: '/mnt/slurm_home/nalin/.cache/huggingface/hub/models--DarkFlameUniverse--Stable-Diffusion-2-1-Base-8bit/refs/main', switching to standard 2-1 base.
+2025-03-09 22:01:16,108 - INFO - Total execution time: 198.35 seconds
+2025-03-09 22:01:16,108 - INFO - Model Path: DarkFlameUniverse/Stable-Diffusion-2-1-Base-8bit
+2025-03-09 22:01:16,108 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 22:01:16,108 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 22:01:16,108 - INFO - Use LCM: False
+2025-03-09 22:01:16,108 - INFO - Number of inference steps: 50
+2025-03-09 22:01:16,108 - INFO - Guidance scale: 1
+
+
+So this was with LCM false and the optimisations for Flags enabled. 
+inductor.config.conv_1x1_as_mm = True
+inductor.config.coordinate_descent_tuning = True
+inductor.config.epilogue_fusion = False
+inductor.config.coordinate_descent_check_all_directions = True

Image-Morpher/logs/execution_20250309_220805.log

@@ -0,0 +1,16 @@
+2025-03-09 22:08:06,169 - ERROR - Model loading failed: [Errno 2] No such file or directory: '/mnt/slurm_home/nalin/.cache/huggingface/hub/models--DarkFlameUniverse--Stable-Diffusion-2-1-Base-8bit/refs/main', switching to standard 2-1 base.
+2025-03-09 22:11:24,031 - INFO - Total execution time: 198.23 seconds
+2025-03-09 22:11:24,031 - INFO - Model Path: DarkFlameUniverse/Stable-Diffusion-2-1-Base-8bit
+2025-03-09 22:11:24,031 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 22:11:24,031 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 22:11:24,031 - INFO - Use LCM: False
+2025-03-09 22:11:24,031 - INFO - Number of inference steps: 50
+2025-03-09 22:11:24,031 - INFO - Guidance scale: 1
+
+
+inductor.config.conv_1x1_as_mm = True
+inductor.config.coordinate_descent_tuning = True
+inductor.config.epilogue_fusion = False
+inductor.config.coordinate_descent_check_all_directions = True
+
+yeah this makes no difference at all. 

Image-Morpher/logs/execution_20250309_221759.log

@@ -0,0 +1,26 @@
+2025-03-09 22:17:59,860 - ERROR - Model loading failed: [Errno 2] No such file or directory: '/mnt/slurm_home/nalin/.cache/huggingface/hub/models--DarkFlameUniverse--Stable-Diffusion-2-1-Base-8bit/refs/main', switching to standard 2-1 base.
+
+torch.compile() doesnt really work, the lora training is very slow.
+
+change: 
+
+def compile_components(pipe):
+    # Compile with maximum optimization
+    pipe.unet = torch.compile(
+        pipe.unet,
+        mode="max-autotune",
+        fullgraph=True,
+        dynamic=False
+    )
+    
+    # Special compilation for VAE decoder
+    pipe.vae.decode = torch.compile(
+        pipe.vae.decode,
+        mode="reduce-overhead",
+        fullgraph=True
+    )
+    
+    return pipe
+
+if not args.use_lcm:  # LCM may conflict with full compilation
+    pipeline = compile_components(pipeline)

Image-Morpher/logs/execution_20250309_224623.log

@@ -0,0 +1,17 @@
+2025-03-09 22:46:23,911 - ERROR - Model loading failed: [Errno 2] No such file or directory: '/mnt/slurm_home/nalin/.cache/huggingface/hub/models--DarkFlameUniverse--Stable-Diffusion-2-1-Base-8bit/refs/main', switching to standard 2-1 base.
+2025-03-09 22:48:08,545 - INFO - Total execution time: 104.98 seconds
+2025-03-09 22:48:08,545 - INFO - Model Path: DarkFlameUniverse/Stable-Diffusion-2-1-Base-8bit
+2025-03-09 22:48:08,545 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 22:48:08,545 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 22:48:08,545 - INFO - Use LCM: True
+2025-03-09 22:48:08,545 - INFO - Number of inference steps: 8
+2025-03-09 22:48:08,545 - INFO - Guidance scale: 1
+
+
+trying some lcm lora optimisations, this does not seem to help too much ? its with 
+
+    pipeline.scheduler = LCMScheduler.from_config(
+        pipeline.scheduler.config,
+        timestep_spacing="trailing",  # Optimized for LCM
+        prediction_type = "epsilion" # this is default anyway. 
+    )

Image-Morpher/logs/execution_20250309_224915.log

@@ -0,0 +1,10 @@
+2025-03-09 22:49:15,959 - ERROR - Model loading failed: [Errno 2] No such file or directory: '/mnt/slurm_home/nalin/.cache/huggingface/hub/models--DarkFlameUniverse--Stable-Diffusion-2-1-Base-8bit/refs/main', switching to standard 2-1 base.
+2025-03-09 22:49:50,993 - INFO - Total execution time: 35.29 seconds
+2025-03-09 22:49:50,993 - INFO - Model Path: DarkFlameUniverse/Stable-Diffusion-2-1-Base-8bit
+2025-03-09 22:49:50,993 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 22:49:50,993 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 22:49:50,993 - INFO - Use LCM: True
+2025-03-09 22:49:50,993 - INFO - Number of inference steps: 8
+2025-03-09 22:49:50,993 - INFO - Guidance scale: 1
+
+fail with v_prediction prediction_type

Image-Morpher/logs/execution_20250309_232833.log

@@ -0,0 +1,17 @@
+2025-03-09 23:30:11,054 - INFO - Total execution time: 97.80 seconds
+2025-03-09 23:30:11,054 - INFO - Model Path: stabilityai/stable-diffusion-2-1-base
+2025-03-09 23:30:11,054 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 23:30:11,054 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 23:30:11,055 - INFO - Use LCM: True
+2025-03-09 23:30:11,055 - INFO - Number of inference steps: 8
+2025-03-09 23:30:11,055 - INFO - Guidance scale: 1
+
+
+so this was without the 
+
+# Configure compiler settings (safe defaults)
+# inductor.config.force_fuse_int_mm_with_mul = True  # Better for diffusion models
+
+it was 0.7 seconds off, so the real benefit (of the real cudnn benchmark OR the set_float32_matmul_precision actually!!, u didnt realise bef) was 
+shifting the cudnn benchmark test AFTER pipeline.to("cuda")
+

Image-Morpher/logs/execution_20250309_235119.log

@@ -0,0 +1,16 @@
+2025-03-09 23:53:03,895 - INFO - Total execution time: 104.19 seconds
+2025-03-09 23:53:03,895 - INFO - Model Path: stabilityai/stable-diffusion-2-1-base
+2025-03-09 23:53:03,896 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 23:53:03,896 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 23:53:03,896 - INFO - Use LCM: True
+2025-03-09 23:53:03,896 - INFO - Number of inference steps: 8
+2025-03-09 23:53:03,896 - INFO - Guidance scale: 1
+
+
+this was base, with the 
+
+torch.backends.cudnn.benchmark = True
+torch.set_float32_matmul_precision("high")  # Better for modern GPUs
+
+commented OUT. Hence the previous does make a difference. 
+

Image-Morpher/logs/execution_20250309_235339.log

@@ -0,0 +1,15 @@
+2025-03-09 23:55:24,160 - INFO - Total execution time: 104.42 seconds
+2025-03-09 23:55:24,160 - INFO - Model Path: stabilityai/stable-diffusion-2-1-base
+2025-03-09 23:55:24,160 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 23:55:24,160 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 23:55:24,160 - INFO - Use LCM: True
+2025-03-09 23:55:24,160 - INFO - Number of inference steps: 8
+2025-03-09 23:55:24,160 - INFO - Guidance scale: 1
+
+
+this was with 
+
+torch.backends.cudnn.benchmark = True
+
+not much speed up! infact its a bit slower (sometimes its slightly tiny bit faster)
+

Image-Morpher/logs/slight_saving_cudnn_benchmark.log

@@ -0,0 +1,9 @@
+2025-03-02 22:57:31,314 - INFO - Total execution time: 123.84 seconds
+2025-03-02 22:57:31,314 - INFO - Model Path: sd-legacy/stable-diffusion-v1-5
+2025-03-02 22:57:31,314 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-02 22:57:31,314 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-02 22:57:31,315 - INFO - Number of inference steps: 8
+2025-03-02 22:57:31,315 - INFO - Guidance scale: 1
+2025-03-02 22:57:31,315 - INFO - Use LCM: True
+
+cudnn benchmark

Image-Morpher/logs/strange_test.log

@@ -0,0 +1,12 @@
+2025-03-09 17:45:32,554 - ERROR - Model loading failed: [Errno 2] No such file or directory: '/mnt/slurm_home/nalin/.cache/huggingface/hub/models--DarkFlameUniverse--Stable-Diffusion-2-1-Base-8bit/refs/main'
+2025-03-09 17:47:22,395 - INFO - Total execution time: 110.10 seconds
+2025-03-09 17:47:22,396 - INFO - Model Path: DarkFlameUniverse/Stable-Diffusion-2-1-Base-8bit # i think this was NOT the model path that was actually used.
+2025-03-09 17:47:22,396 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 17:47:22,396 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 17:47:22,396 - INFO - Use LCM: True
+2025-03-09 17:47:22,396 - INFO - Number of inference steps: 8
+2025-03-09 17:47:22,396 - INFO - Guidance scale: 1
+
+
+so it worked with the quantized 8-bit model (no it didnt) ! but... it was not able to load it in at first, then had to use my local cache of the failed run i tried befo i think (nope)
+

Image-Morpher/logs/successful_memory_optimisation1.log

@@ -0,0 +1,17 @@
+2025-03-09 17:01:38,747 - INFO - Total execution time: 121.97 seconds
+2025-03-09 17:01:38,747 - INFO - Model Path: sd-legacy/stable-diffusion-v1-5
+2025-03-09 17:01:38,747 - INFO - Image Path 0: ./assets/Trump.jpg
+2025-03-09 17:01:38,747 - INFO - Image Path 1: ./assets/Biden.jpg
+2025-03-09 17:01:38,747 - INFO - Use LCM: True
+2025-03-09 17:01:38,747 - INFO - Number of inference steps: 8
+2025-03-09 17:01:38,747 - INFO - Guidance scale: 1
+
+
+fastest one yet, 
+with optimisation
+pipeline.enable_vae_slicing()
+pipeline.enable_attention_slicing()
+
+about 5-8 seconds saved on average.
+
+and the cudnn benchmark MAYBE helps but not really.

Image-Morpher/main.py

@@ -0,0 +1,162 @@
+"""
+Modified main.py on DiffMorpher, LCM-LoRA support + param additions + logging + optimizations for speed-up
+"""
+import os
+import torch
+import numpy as np
+from PIL import Image
+from argparse import ArgumentParser
+from model import DiffMorpherPipeline
+import time
+import logging
+
+logs_folder = "logs"
+os.makedirs(logs_folder, exist_ok=True)
+
+# Create a unique log filename using the current time 
+log_filename = os.path.join(logs_folder, f"execution_{time.strftime('%Y%m%d_%H%M%S')}.log")
+logging.basicConfig(
+    filename=log_filename,
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+
+start_time = time.time()
+
+parser = ArgumentParser()
+parser.add_argument(
+    "--model_path", type=str, default="stabilityai/stable-diffusion-2-1-base",
+    help="Pretrained model to use (default: %(default)s)"
+)
+# Available SDV1-5 versions: 
+# sd-legacy/stable-diffusion-v1-5
+# lykon/dreamshaper-7
+
+# Original DiffMorpher SD: 
+# stabilityai/stable-diffusion-2-1-base
+
+# Quantized models to try (non-functional, possible extension for future)
+# DarkFlameUniverse/Stable-Diffusion-2-1-Base-8bit
+# Xerox32/SD2.1-base-Int8
+
+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 first image (default: %(default)s)"
+)
+parser.add_argument(
+    "--image_path_1", type=str, default="",
+    help="Path of the second 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 for saving LoRA weights (default: %(default)s)"
+)
+parser.add_argument(
+    "--load_lora_path_0", type=str, default="",
+    help="Path of the LoRA weights for the first image (default: %(default)s)"
+)
+parser.add_argument(
+    "--load_lora_path_1", type=str, default="",
+    help="Path of the LoRA weights for the second image (default: %(default)s)"
+)
+parser.add_argument(
+    "--num_inference_steps", type=int, default=50, 
+    help="Number of inference steps (default: %(default)s)")
+parser.add_argument(
+    "--guidance_scale", type=float, default=1,  # To match current diffmorpher
+    help="Guidance scale for classifier-free guidance (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", help="Disable style LoRA (default: %(default)s)")
+
+# New argument for LCM LoRA acceleration
+parser.add_argument("--use_lcm", action="store_true", help="Enable LCM-LoRA acceleration for faster sampling")
+
+args = parser.parse_args()
+os.makedirs(args.output_path, exist_ok=True)
+
+# Create the pipeline from the given model path
+pipeline = DiffMorpherPipeline.from_pretrained(args.model_path, torch_dtype=torch.float32)
+
+# memory optimisations for vae and attention slicing - breaks computations into smaller chunks to fit better in mem
+# can lead to more efficient caching and memory access. better memory locality
+# found that its helpful with GPUs with limited VRAM memory in particular.
+pipeline.enable_vae_slicing()
+pipeline.enable_attention_slicing()
+
+pipeline.to("cuda")
+
+# Add these AFTER device movement
+torch.backends.cudnn.benchmark = True # finds efficient convolution algo by running short benchmark, minimal speed-up.
+torch.set_float32_matmul_precision("high")  # Better for modern GPUs, reduces about 7 seconds of inference time.
+
+# Integrate LCM-LoRA if flagged, OUTSIDE any of the style LoRA loading / training steps.
+if args.use_lcm:
+    from lcm_lora.lcm_schedule import LCMScheduler
+    # Replace scheduler using LCM's configuration
+    pipeline.scheduler = LCMScheduler.from_config(pipeline.scheduler.config)
+    # Load the LCM LoRA weights (LCM provides an add-on network)
+    pipeline.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
+    # Set the lcm_inference_steps
+    args.num_inference_steps = 8  # Override with LCM-recommended steps
+    # set CFG (range allowed by legacy code: 0 to 1, 1 performs best)
+    args.guidance_scale = 1
+
+# Run the pipeline inference using existing parameters
+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,
+    num_inference_steps = args.num_inference_steps, # enforce when LCM enabled
+    fix_lora=args.fix_lora_value,
+    save_intermediates=args.save_inter,
+    use_lora=not args.no_lora,
+    use_lcm = args.use_lcm,
+    guidance_scale=args.guidance_scale, # enforce when LCM enabled
+)
+
+# Save the resulting GIF output from the sequence of images
+images[0].save(f"{args.output_path}/output.gif", save_all=True,
+               append_images=images[1:], duration=args.duration, loop=0)
+
+end_time = time.time()
+elapsed_time = end_time - start_time
+
+# Log the execution details and parameters
+logging.info(f"Total execution time: {elapsed_time:.2f} seconds")
+logging.info(f"Model Path: {args.model_path}")
+logging.info(f"Image Path 0: {args.image_path_0}")
+logging.info(f"Image Path 1: {args.image_path_1}")
+logging.info(f"Use LCM: {args.use_lcm}")
+logging.info(f"Number of inference steps: {args.num_inference_steps}")
+logging.info(f"Guidance scale: {args.guidance_scale}")
+
+print(f"Total execution time: {elapsed_time:.2f} seconds, log file saved as {log_filename}")

Image-Morpher/model.py

@@ -0,0 +1,647 @@
+# model.py
+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.alphas_cumprod[0]
+        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.alphas_cumprod[0]
+                )
+
+                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.alphas_cumprod[0]
+        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, use_lcm, 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)
+
+        if use_lcm:
+            sampler_desc = "LCM multi-step sampler"
+        else: 
+            sampler_desc = "DDIM Sampler" # currently defaults to this   
+
+        for i, t in enumerate(tqdm.tqdm(self.scheduler.timesteps, desc=f"{sampler_desc}, 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_lcm = False,
+            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
+        self.use_lcm = use_lcm
+
+        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

Image-Morpher/requirements_diffmorpher.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

Image-Morpher/run.sh

@@ -0,0 +1,21 @@
+# srun -p rtx3090_slab -w slabgpu08 --gres=gpu:1 \
+#     --job-name=test --kill-on-bad-exit=1 python3 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 --use_lcm 
+
+# srun -p rtx3090_slab -w slabgpu08 --gres=gpu:1 \
+#     --job-name=test --kill-on-bad-exit=1 python3 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 --use_lcm
+
+srun -p rtx3090_slab -w slabgpu05 --gres=gpu:1 \
+    --job-name=test --kill-on-bad-exit=1 python3 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 --use_lcm
+

Image-Morpher/utils/alpha_scheduler.py

@@ -0,0 +1,54 @@
+import bisect
+import torch
+import torch.nn.functional as F
+import lpips
+
+perceptual_loss = lpips.LPIPS()
+
+
+def distance(img_a, img_b):
+    return perceptual_loss(img_a, img_b).item()
+    # return F.mse_loss(img_a, img_b).item()
+
+
+class AlphaScheduler:
+    def __init__(self):
+        ...
+
+    def from_imgs(self, imgs):
+        self.__num_values = len(imgs)
+        self.__values = [0]
+        for i in range(self.__num_values - 1):
+            dis = distance(imgs[i], imgs[i + 1])
+            self.__values.append(dis)
+            self.__values[i + 1] += self.__values[i]
+        for i in range(self.__num_values):
+            self.__values[i] /= self.__values[-1]
+
+    def save(self, filename):
+        torch.save(torch.tensor(self.__values), filename)
+
+    def load(self, filename):
+        self.__values = torch.load(filename).tolist()
+        self.__num_values = len(self.__values)
+
+    def get_x(self, y):
+        assert y >= 0 and y <= 1
+        id = bisect.bisect_left(self.__values, y)
+        id -= 1
+        if id < 0:
+            id = 0
+        yl = self.__values[id]
+        yr = self.__values[id + 1]
+        xl = id * (1 / (self.__num_values - 1))
+        xr = (id + 1) * (1 / (self.__num_values - 1))
+        x = (y - yl) / (yr - yl) * (xr - xl) + xl
+        return x
+
+    def get_list(self, len=None):
+        if len is None:
+            len = self.__num_values
+
+        ys = torch.linspace(0, 1, len)
+        res = [self.get_x(y) for y in ys]
+        return res

Image-Morpher/utils/lora_utils.py

@@ -0,0 +1,284 @@
+# utils / lora_utils.py
+from timeit import default_timer as timer
+from datetime import timedelta
+from PIL import Image
+import os
+import numpy as np
+from einops import rearrange
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+import transformers
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from packaging import version
+from PIL import Image
+import tqdm
+
+from transformers import AutoTokenizer, PretrainedConfig
+
+import diffusers
+from diffusers import (
+    AutoencoderKL,
+    DDPMScheduler,
+    DiffusionPipeline,
+    DPMSolverMultistepScheduler,
+    StableDiffusionPipeline,
+    UNet2DConditionModel,
+)
+from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
+from diffusers.models.attention_processor import (
+    AttnAddedKVProcessor,
+    AttnAddedKVProcessor2_0,
+    LoRAAttnAddedKVProcessor,
+    LoRAAttnProcessor,
+    LoRAAttnProcessor2_0,
+    SlicedAttnAddedKVProcessor,
+)
+from diffusers.optimization import get_scheduler
+from diffusers.utils import check_min_version
+from diffusers.utils.import_utils import is_xformers_available
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.17.0")
+
+
+def import_model_class_from_model_name_or_path(pretrained_model_name_or_path: str, revision: str):
+    text_encoder_config = PretrainedConfig.from_pretrained(
+        pretrained_model_name_or_path,
+        subfolder="text_encoder",
+        revision=revision,
+    )
+    model_class = text_encoder_config.architectures[0]
+
+    if model_class == "CLIPTextModel":
+        from transformers import CLIPTextModel
+
+        return CLIPTextModel
+    elif model_class == "RobertaSeriesModelWithTransformation":
+        from diffusers.pipelines.alt_diffusion.modeling_roberta_series import RobertaSeriesModelWithTransformation
+
+        return RobertaSeriesModelWithTransformation
+    elif model_class == "T5EncoderModel":
+        from transformers import T5EncoderModel
+
+        return T5EncoderModel
+    else:
+        raise ValueError(f"{model_class} is not supported.")
+
+def tokenize_prompt(tokenizer, prompt, tokenizer_max_length=None):
+    if tokenizer_max_length is not None:
+        max_length = tokenizer_max_length
+    else:
+        max_length = tokenizer.model_max_length
+
+    text_inputs = tokenizer(
+        prompt,
+        truncation=True,
+        padding="max_length",
+        max_length=max_length,
+        return_tensors="pt",
+    )
+
+    return text_inputs
+
+def encode_prompt(text_encoder, input_ids, attention_mask, text_encoder_use_attention_mask=False):
+    text_input_ids = input_ids.to(text_encoder.device)
+
+    if text_encoder_use_attention_mask:
+        attention_mask = attention_mask.to(text_encoder.device)
+    else:
+        attention_mask = None
+
+    prompt_embeds = text_encoder(
+        text_input_ids,
+        attention_mask=attention_mask,
+    )
+    prompt_embeds = prompt_embeds[0]
+
+    return prompt_embeds
+
+# model_path: path of the model
+# image: input image, have not been pre-processed
+# save_lora_dir: the path to save the lora
+# prompt: the user input prompt
+# lora_steps: number of lora training step
+# lora_lr: learning rate of lora training
+# lora_rank: the rank of lora
+def train_lora(image, prompt, save_lora_dir, model_path=None, tokenizer=None, text_encoder=None, vae=None, unet=None, noise_scheduler=None, lora_steps=200, lora_lr=2e-4, lora_rank=16, weight_name=None, safe_serialization=False, progress=tqdm):
+    # initialize accelerator
+    accelerator = Accelerator(
+        gradient_accumulation_steps=1,
+        # mixed_precision='fp16'
+    )
+    set_seed(0)
+
+    # Load the tokenizer
+    if tokenizer is None:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            subfolder="tokenizer",
+            revision=None,
+            use_fast=False,
+        )
+    # initialize the model
+    if noise_scheduler is None:
+        noise_scheduler = DDPMScheduler.from_pretrained(model_path, subfolder="scheduler")
+    if text_encoder is None:
+        text_encoder_cls = import_model_class_from_model_name_or_path(model_path, revision=None)
+        text_encoder = text_encoder_cls.from_pretrained(
+            model_path, subfolder="text_encoder", revision=None
+        )
+    if vae is None:
+        vae = AutoencoderKL.from_pretrained(
+            model_path, subfolder="vae", revision=None
+        )
+    if unet is None:
+        unet = UNet2DConditionModel.from_pretrained(
+            model_path, subfolder="unet", revision=None
+        )
+
+    # set device and dtype
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+    vae.requires_grad_(False)
+    text_encoder.requires_grad_(False)
+    unet.requires_grad_(False)
+
+    unet.to(device)
+    vae.to(device)
+    text_encoder.to(device)
+
+    # initialize UNet LoRA
+    unet_lora_attn_procs = {}
+    for name, attn_processor in unet.attn_processors.items():
+        cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
+        if name.startswith("mid_block"):
+            hidden_size = unet.config.block_out_channels[-1]
+        elif name.startswith("up_blocks"):
+            block_id = int(name[len("up_blocks.")])
+            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+        elif name.startswith("down_blocks"):
+            block_id = int(name[len("down_blocks.")])
+            hidden_size = unet.config.block_out_channels[block_id]
+        else:
+            raise NotImplementedError("name must start with up_blocks, mid_blocks, or down_blocks")
+
+        if isinstance(attn_processor, (AttnAddedKVProcessor, SlicedAttnAddedKVProcessor, AttnAddedKVProcessor2_0)):
+            lora_attn_processor_class = LoRAAttnAddedKVProcessor
+        else:
+            lora_attn_processor_class = (
+                LoRAAttnProcessor2_0 if hasattr(F, "scaled_dot_product_attention") else LoRAAttnProcessor
+            )
+        unet_lora_attn_procs[name] = lora_attn_processor_class(
+            hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, rank=lora_rank
+        )
+    unet.set_attn_processor(unet_lora_attn_procs)
+    unet_lora_layers = AttnProcsLayers(unet.attn_processors)
+
+    # Optimizer creation
+    params_to_optimize = (unet_lora_layers.parameters())
+    optimizer = torch.optim.AdamW(
+        params_to_optimize,
+        lr=lora_lr,
+        betas=(0.9, 0.999),
+        weight_decay=1e-2,
+        eps=1e-08,
+    )
+
+    lr_scheduler = get_scheduler(
+        "constant",
+        optimizer=optimizer,
+        num_warmup_steps=0,
+        num_training_steps=lora_steps,
+        num_cycles=1,
+        power=1.0,
+    )
+
+    # prepare accelerator
+    unet_lora_layers = accelerator.prepare_model(unet_lora_layers)
+    optimizer = accelerator.prepare_optimizer(optimizer)
+    lr_scheduler = accelerator.prepare_scheduler(lr_scheduler)
+
+    # initialize text embeddings
+    with torch.no_grad():
+        text_inputs = tokenize_prompt(tokenizer, prompt, tokenizer_max_length=None)
+        text_embedding = encode_prompt(
+            text_encoder,
+            text_inputs.input_ids,
+            text_inputs.attention_mask,
+            text_encoder_use_attention_mask=False
+        )
+
+    if type(image) == np.ndarray:
+        image = Image.fromarray(image)
+        
+    # initialize latent distribution
+    image_transforms = transforms.Compose(
+        [
+            transforms.Resize(512, interpolation=transforms.InterpolationMode.BILINEAR),
+            # transforms.RandomCrop(512),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5]),
+        ]
+    )
+
+    image = image_transforms(image).to(device)
+    image = image.unsqueeze(dim=0)
+    
+    latents_dist = vae.encode(image).latent_dist
+    for _ in progress.tqdm(range(lora_steps), desc="Training LoRA..."):
+        unet.train()
+        model_input = latents_dist.sample() * vae.config.scaling_factor
+        # Sample noise that we'll add to the latents
+        noise = torch.randn_like(model_input)
+        bsz, channels, height, width = model_input.shape
+        # Sample a random timestep for each image
+        timesteps = torch.randint(
+            0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device
+        )
+        timesteps = timesteps.long()
+
+        # Add noise to the model input according to the noise magnitude at each timestep
+        # (this is the forward diffusion process)
+        noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps)
+
+        # Predict the noise residual
+        model_pred = unet(noisy_model_input, timesteps, text_embedding).sample
+
+        # Get the target for loss depending on the prediction type
+        if noise_scheduler.config.prediction_type == "epsilon":
+            target = noise
+        elif noise_scheduler.config.prediction_type == "v_prediction":
+            target = noise_scheduler.get_velocity(model_input, noise, timesteps)
+        else:
+            raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
+
+        loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+        accelerator.backward(loss)
+        optimizer.step()
+        lr_scheduler.step()
+        optimizer.zero_grad()
+
+    # save the trained lora
+    # unet = unet.to(torch.float32)
+    # vae = vae.to(torch.float32)
+    # text_encoder = text_encoder.to(torch.float32)
+
+    # unwrap_model is used to remove all special modules added when doing distributed training
+    # so here, there is no need to call unwrap_model
+    # unet_lora_layers = accelerator.unwrap_model(unet_lora_layers)
+    LoraLoaderMixin.save_lora_weights(
+        save_directory=save_lora_dir,
+        unet_lora_layers=unet_lora_layers,
+        text_encoder_lora_layers=None,
+        weight_name=weight_name,
+        safe_serialization=safe_serialization
+    )
+    
+def load_lora(unet, lora_0, lora_1, alpha):
+    lora = {}
+    for key in lora_0:
+        lora[key] = (1 - alpha) * lora_0[key] + alpha * lora_1[key]
+    unet.load_attn_procs(lora)
+    return unet

Image-Morpher/utils/model_utils.py

@@ -0,0 +1,87 @@
+# utils/model_utils.py
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+
+def calc_mean_std(feat, eps=1e-5):
+    # eps is a small value added to the variance to avoid divide-by-zero.
+    size = feat.size()
+
+    N, C = size[:2]
+    feat_var = feat.view(N, C, -1).var(dim=2) + eps
+    if len(size) == 3:
+        feat_std = feat_var.sqrt().view(N, C, 1)
+        feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1)
+    else:
+        feat_std = feat_var.sqrt().view(N, C, 1, 1)
+        feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
+    return feat_mean, feat_std
+
+
+def get_img(img, resolution=512):
+    norm_mean = [0.5, 0.5, 0.5]
+    norm_std = [0.5, 0.5, 0.5]
+    transform = transforms.Compose([
+        transforms.Resize((resolution, resolution)),
+        transforms.ToTensor(),
+        transforms.Normalize(norm_mean, norm_std)
+    ])
+    img = transform(img)
+    return img.unsqueeze(0)
+
+@torch.no_grad()
+def slerp(p0, p1, fract_mixing: float, adain=True):
+    r""" Copied from lunarring/latentblending
+    Helper function to correctly mix two random variables using spherical interpolation.
+    The function will always cast up to float64 for sake of extra 4.
+    Args:
+        p0: 
+            First tensor for interpolation
+        p1: 
+            Second tensor for interpolation
+        fract_mixing: float 
+            Mixing coefficient of interval [0, 1]. 
+            0 will return in p0
+            1 will return in p1
+            0.x will return a mix between both preserving angular velocity.
+    """
+    if p0.dtype == torch.float16:
+        recast_to = 'fp16'
+    else:
+        recast_to = 'fp32'
+
+    p0 = p0.double()
+    p1 = p1.double()
+
+    if adain:
+        mean1, std1 = calc_mean_std(p0)
+        mean2, std2 = calc_mean_std(p1)
+        mean = mean1 * (1 - fract_mixing) + mean2 * fract_mixing
+        std = std1 * (1 - fract_mixing) + std2 * fract_mixing
+        
+    norm = torch.linalg.norm(p0) * torch.linalg.norm(p1)
+    epsilon = 1e-7
+    dot = torch.sum(p0 * p1) / norm
+    dot = dot.clamp(-1+epsilon, 1-epsilon)
+
+    theta_0 = torch.arccos(dot)
+    sin_theta_0 = torch.sin(theta_0)
+    theta_t = theta_0 * fract_mixing
+    s0 = torch.sin(theta_0 - theta_t) / sin_theta_0
+    s1 = torch.sin(theta_t) / sin_theta_0
+    interp = p0*s0 + p1*s1
+
+    if adain:
+        interp = F.instance_norm(interp) * std + mean
+
+    if recast_to == 'fp16':
+        interp = interp.half()
+    elif recast_to == 'fp32':
+        interp = interp.float()
+
+    return interp
+
+
+def do_replace_attn(key: str):
+    # return key.startswith('up_blocks.2') or key.startswith('up_blocks.3')
+    return key.startswith('up')