Upload 17 files

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

MotionDirector_inference.py

@@ -0,0 +1,284 @@
+import argparse
+import os
+import platform
+import re
+import warnings
+from typing import Optional
+
+import torch
+from diffusers import DDIMScheduler, TextToVideoSDPipeline
+from einops import rearrange
+from torch import Tensor
+from torch.nn.functional import interpolate
+from tqdm import trange
+import random
+
+from MotionDirector_train import export_to_video, handle_memory_attention, load_primary_models, unet_and_text_g_c, freeze_models
+from utils.lora_handler import LoraHandler
+from utils.ddim_utils import ddim_inversion
+import imageio
+
+
+def initialize_pipeline(
+    model: str,
+    device: str = "cuda",
+    xformers: bool = False,
+    sdp: bool = False,
+    lora_path: str = "",
+    lora_rank: int = 64,
+    lora_scale: float = 1.0,
+):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+
+        scheduler, tokenizer, text_encoder, vae, unet = load_primary_models(model)
+
+    # Freeze any necessary models
+    freeze_models([vae, text_encoder, unet])
+
+    # Enable xformers if available
+    handle_memory_attention(xformers, sdp, unet)
+
+    lora_manager_temporal = LoraHandler(
+        version="cloneofsimo",
+        use_unet_lora=True,
+        use_text_lora=False,
+        save_for_webui=False,
+        only_for_webui=False,
+        unet_replace_modules=["TransformerTemporalModel"],
+        text_encoder_replace_modules=None,
+        lora_bias=None
+    )
+
+    unet_lora_params, unet_negation = lora_manager_temporal.add_lora_to_model(
+        True, unet, lora_manager_temporal.unet_replace_modules, 0, lora_path, r=lora_rank, scale=lora_scale)
+
+    unet.eval()
+    text_encoder.eval()
+    unet_and_text_g_c(unet, text_encoder, False, False)
+
+    pipe = TextToVideoSDPipeline.from_pretrained(
+        pretrained_model_name_or_path=model,
+        scheduler=scheduler,
+        tokenizer=tokenizer,
+        text_encoder=text_encoder.to(device=device, dtype=torch.half),
+        vae=vae.to(device=device, dtype=torch.half),
+        unet=unet.to(device=device, dtype=torch.half),
+    )
+    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+
+    return pipe
+
+
+def inverse_video(pipe, latents, num_steps):
+    ddim_inv_scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+    ddim_inv_scheduler.set_timesteps(num_steps)
+
+    ddim_inv_latent = ddim_inversion(
+        pipe, ddim_inv_scheduler, video_latent=latents.to(pipe.device),
+        num_inv_steps=num_steps, prompt="")[-1]
+    return ddim_inv_latent
+
+
+def prepare_input_latents(
+    pipe: TextToVideoSDPipeline,
+    batch_size: int,
+    num_frames: int,
+    height: int,
+    width: int,
+    latents_path:str,
+    noise_prior: float
+):
+    # initialize with random gaussian noise
+    scale = pipe.vae_scale_factor
+    shape = (batch_size, pipe.unet.config.in_channels, num_frames, height // scale, width // scale)
+    if noise_prior > 0.:
+        cached_latents = torch.load(latents_path)
+        if 'inversion_noise' not in cached_latents:
+            latents = inverse_video(pipe, cached_latents['latents'].unsqueeze(0), 50).squeeze(0)
+        else:
+            latents = torch.load(latents_path)['inversion_noise'].unsqueeze(0)
+        if latents.shape[0] != batch_size:
+            latents = latents.repeat(batch_size, 1, 1, 1, 1)
+        if latents.shape != shape:
+            latents = interpolate(rearrange(latents, "b c f h w -> (b f) c h w", b=batch_size), (height // scale, width // scale), mode='bilinear')
+            latents = rearrange(latents, "(b f) c h w -> b c f h w", b=batch_size)
+        noise = torch.randn_like(latents, dtype=torch.half)
+        latents = (noise_prior) ** 0.5 * latents + (1 - noise_prior) ** 0.5 * noise
+    else:
+        latents = torch.randn(shape, dtype=torch.half)
+
+    return latents
+
+
+def encode(pipe: TextToVideoSDPipeline, pixels: Tensor, batch_size: int = 8):
+    nf = pixels.shape[2]
+    pixels = rearrange(pixels, "b c f h w -> (b f) c h w")
+
+    latents = []
+    for idx in trange(
+        0, pixels.shape[0], batch_size, desc="Encoding to latents...", unit_scale=batch_size, unit="frame"
+    ):
+        pixels_batch = pixels[idx : idx + batch_size].to(pipe.device, dtype=torch.half)
+        latents_batch = pipe.vae.encode(pixels_batch).latent_dist.sample()
+        latents_batch = latents_batch.mul(pipe.vae.config.scaling_factor).cpu()
+        latents.append(latents_batch)
+    latents = torch.cat(latents)
+
+    latents = rearrange(latents, "(b f) c h w -> b c f h w", f=nf)
+
+    return latents
+
+
+@torch.inference_mode()
+def inference(
+    model: str,
+    prompt: str,
+    negative_prompt: Optional[str] = None,
+    width: int = 256,
+    height: int = 256,
+    num_frames: int = 24,
+    num_steps: int = 50,
+    guidance_scale: float = 15,
+    device: str = "cuda",
+    xformers: bool = False,
+    sdp: bool = False,
+    lora_path: str = "",
+    lora_rank: int = 64,
+    lora_scale: float = 1.0,
+    seed: Optional[int] = None,
+    latents_path: str="",
+    noise_prior: float = 0.,
+    repeat_num: int = 1,
+):
+    if seed is not None:
+        random_seed = seed
+        torch.manual_seed(seed)
+
+    with torch.autocast(device, dtype=torch.half):
+        # prepare models
+        pipe = initialize_pipeline(model, device, xformers, sdp, lora_path, lora_rank, lora_scale)
+
+        for i in range(repeat_num):
+            if seed is None:
+                random_seed = random.randint(100, 10000000)
+                torch.manual_seed(random_seed)
+
+            # prepare input latents
+            init_latents = prepare_input_latents(
+                pipe=pipe,
+                batch_size=len(prompt),
+                num_frames=num_frames,
+                height=height,
+                width=width,
+                latents_path=latents_path,
+                noise_prior=noise_prior
+            )
+
+            with torch.no_grad():
+                video_frames = pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    width=width,
+                    height=height,
+                    num_frames=num_frames,
+                    num_inference_steps=num_steps,
+                    guidance_scale=guidance_scale,
+                    latents=init_latents
+                ).frames
+
+            # =========================================
+            # ========= write outputs to file =========
+            # =========================================
+            os.makedirs(args.output_dir, exist_ok=True)
+
+            # save to mp4
+            export_to_video(video_frames, f"{out_name}_{random_seed}.mp4", args.fps)
+
+            # # save to gif
+            file_name = f"{out_name}_{random_seed}.gif"
+            imageio.mimsave(file_name, video_frames, 'GIF', duration=1000 * 1 / args.fps, loop=0)
+
+    return video_frames
+
+
+if __name__ == "__main__":
+    import decord
+
+    decord.bridge.set_bridge("torch")
+
+    # fmt: off
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-m", "--model", type=str, required=True,
+                        help="HuggingFace repository or path to model checkpoint directory")
+    parser.add_argument("-p", "--prompt", type=str, required=True, help="Text prompt to condition on")
+    parser.add_argument("-n", "--negative-prompt", type=str, default=None, help="Text prompt to condition against")
+    parser.add_argument("-o", "--output_dir", type=str, default="./outputs/inference", help="Directory to save output video to")
+    parser.add_argument("-B", "--batch-size", type=int, default=1, help="Batch size for inference")
+    parser.add_argument("-W", "--width", type=int, default=384, help="Width of output video")
+    parser.add_argument("-H", "--height", type=int, default=384, help="Height of output video")
+    parser.add_argument("-T", "--num-frames", type=int, default=16, help="Total number of frames to generate")
+    parser.add_argument("-s", "--num-steps", type=int, default=30, help="Number of diffusion steps to run per frame.")
+    parser.add_argument("-g", "--guidance-scale", type=float, default=12, help="Scale for guidance loss (higher values = more guidance, but possibly more artifacts).")
+    parser.add_argument("-f", "--fps", type=int, default=8, help="FPS of output video")
+    parser.add_argument("-d", "--device", type=str, default="cuda", help="Device to run inference on (defaults to cuda).")
+    parser.add_argument("-x", "--xformers", action="store_true", help="Use XFormers attnetion, a memory-efficient attention implementation (requires `pip install xformers`).")
+    parser.add_argument("-S", "--sdp", action="store_true", help="Use SDP attention, PyTorch's built-in memory-efficient attention implementation.")
+    parser.add_argument("-cf", "--checkpoint_folder", type=str, default=None, help="Path to Low Rank Adaptation checkpoint file (defaults to empty string, which uses no LoRA).")
+    parser.add_argument("-lr", "--lora_rank", type=int, default=32, help="Size of the LoRA checkpoint's projection matrix (defaults to 32).")
+    parser.add_argument("-ls", "--lora_scale", type=float, default=1.0, help="Scale of LoRAs.")
+    parser.add_argument("-r", "--seed", type=int, default=None, help="Random seed to make generations reproducible.")
+    parser.add_argument("-np", "--noise_prior", type=float, default=0., help="Scale of the influence of inversion noise.")
+    parser.add_argument("-ci", "--checkpoint_index", type=int, required=True,
+                        help="The index of checkpoint, such as 300.")
+    parser.add_argument("-rn", "--repeat_num", type=int, default=1,
+                        help="How many results to generate with the same prompt.")
+
+    args = parser.parse_args()
+    # fmt: on
+
+    # =========================================
+    # ====== validate and prepare inputs ======
+    # =========================================
+
+    out_name = f"{args.output_dir}/"
+    prompt = re.sub(r'[<>:"/\\|?*\x00-\x1F]', "_", args.prompt) if platform.system() == "Windows" else args.prompt
+    out_name += f"{prompt}".replace(' ','_').replace(',', '').replace('.', '')
+
+    args.prompt = [prompt] * args.batch_size
+    if args.negative_prompt is not None:
+        args.negative_prompt = [args.negative_prompt] * args.batch_size
+
+    # =========================================
+    # ============= sample videos =============
+    # =========================================
+    if args.checkpoint_index is not None:
+        lora_path = f"{args.checkpoint_folder}/checkpoint-{args.checkpoint_index}/temporal/lora"
+    else:
+        lora_path = f"{args.checkpoint_folder}/checkpoint-default/temporal/lora"
+    latents_folder = f"{args.checkpoint_folder}/cached_latents"
+    latents_path = f"{latents_folder}/{random.choice(os.listdir(latents_folder))}"
+    assert os.path.exists(lora_path)
+    video_frames = inference(
+        model=args.model,
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        width=args.width,
+        height=args.height,
+        num_frames=args.num_frames,
+        num_steps=args.num_steps,
+        guidance_scale=args.guidance_scale,
+        device=args.device,
+        xformers=args.xformers,
+        sdp=args.sdp,
+        lora_path=lora_path,
+        lora_rank=args.lora_rank,
+        lora_scale = args.lora_scale,
+        seed=args.seed,
+        latents_path=latents_path,
+        noise_prior=args.noise_prior,
+        repeat_num=args.repeat_num
+    )
+
+
+

MotionDirector_inference_batch.py

@@ -0,0 +1,290 @@
+import argparse
+import os
+import platform
+import re
+import warnings
+from typing import Optional
+
+import torch
+from diffusers import DDIMScheduler, TextToVideoSDPipeline
+from einops import rearrange
+from torch import Tensor
+from torch.nn.functional import interpolate
+from tqdm import trange
+import random
+
+from MotionDirector_train import export_to_video, handle_memory_attention, load_primary_models, unet_and_text_g_c, freeze_models
+from utils.lora_handler import LoraHandler
+from utils.ddim_utils import ddim_inversion
+import imageio
+
+
+def initialize_pipeline(
+    model: str,
+    device: str = "cuda",
+    xformers: bool = False,
+    sdp: bool = False,
+    lora_path: str = "",
+    lora_rank: int = 64,
+    lora_scale: float = 1.0,
+):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+
+        scheduler, tokenizer, text_encoder, vae, unet = load_primary_models(model)
+
+    # Freeze any necessary models
+    freeze_models([vae, text_encoder, unet])
+
+    # Enable xformers if available
+    handle_memory_attention(xformers, sdp, unet)
+
+    lora_manager_temporal = LoraHandler(
+        version="cloneofsimo",
+        use_unet_lora=True,
+        use_text_lora=False,
+        save_for_webui=False,
+        only_for_webui=False,
+        unet_replace_modules=["TransformerTemporalModel"],
+        text_encoder_replace_modules=None,
+        lora_bias=None
+    )
+
+    unet_lora_params, unet_negation = lora_manager_temporal.add_lora_to_model(
+        True, unet, lora_manager_temporal.unet_replace_modules, 0, lora_path, r=lora_rank, scale=lora_scale)
+
+    unet.eval()
+    text_encoder.eval()
+    unet_and_text_g_c(unet, text_encoder, False, False)
+
+    pipe = TextToVideoSDPipeline.from_pretrained(
+        pretrained_model_name_or_path=model,
+        scheduler=scheduler,
+        tokenizer=tokenizer,
+        text_encoder=text_encoder.to(device=device, dtype=torch.half),
+        vae=vae.to(device=device, dtype=torch.half),
+        unet=unet.to(device=device, dtype=torch.half),
+    )
+    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+
+    return pipe
+
+
+def inverse_video(pipe, latents, num_steps):
+    ddim_inv_scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+    ddim_inv_scheduler.set_timesteps(num_steps)
+
+    ddim_inv_latent = ddim_inversion(
+        pipe, ddim_inv_scheduler, video_latent=latents.to(pipe.device),
+        num_inv_steps=num_steps, prompt="")[-1]
+    return ddim_inv_latent
+
+
+def prepare_input_latents(
+    pipe: TextToVideoSDPipeline,
+    batch_size: int,
+    num_frames: int,
+    height: int,
+    width: int,
+    latents_path:str,
+    noise_prior: float
+):
+    # initialize with random gaussian noise
+    scale = pipe.vae_scale_factor
+    shape = (batch_size, pipe.unet.config.in_channels, num_frames, height // scale, width // scale)
+    if noise_prior > 0.:
+        cached_latents = torch.load(latents_path)
+        if 'inversion_noise' not in cached_latents:
+            latents = inverse_video(pipe, cached_latents['latents'].unsqueeze(0), 50).squeeze(0)
+        else:
+            latents = torch.load(latents_path)['inversion_noise'].unsqueeze(0)
+        if latents.shape[0] != batch_size:
+            latents = latents.repeat(batch_size, 1, 1, 1, 1)
+        if latents.shape != shape:
+            latents = interpolate(rearrange(latents, "b c f h w -> (b f) c h w", b=batch_size), (height // scale, width // scale), mode='bilinear')
+            latents = rearrange(latents, "(b f) c h w -> b c f h w", b=batch_size)
+        noise = torch.randn_like(latents, dtype=torch.half)
+        latents = (noise_prior) ** 0.5 * latents + (1 - noise_prior) ** 0.5 * noise
+    else:
+        latents = torch.randn(shape, dtype=torch.half)
+
+    return latents
+
+
+def encode(pipe: TextToVideoSDPipeline, pixels: Tensor, batch_size: int = 8):
+    nf = pixels.shape[2]
+    pixels = rearrange(pixels, "b c f h w -> (b f) c h w")
+
+    latents = []
+    for idx in trange(
+        0, pixels.shape[0], batch_size, desc="Encoding to latents...", unit_scale=batch_size, unit="frame"
+    ):
+        pixels_batch = pixels[idx : idx + batch_size].to(pipe.device, dtype=torch.half)
+        latents_batch = pipe.vae.encode(pixels_batch).latent_dist.sample()
+        latents_batch = latents_batch.mul(pipe.vae.config.scaling_factor).cpu()
+        latents.append(latents_batch)
+    latents = torch.cat(latents)
+
+    latents = rearrange(latents, "(b f) c h w -> b c f h w", f=nf)
+
+    return latents
+
+
+
+
+@torch.inference_mode()
+def inference(
+    model: str,
+    prompt: str,
+    negative_prompt: Optional[str] = None,
+    width: int = 256,
+    height: int = 256,
+    num_frames: int = 24,
+    num_steps: int = 50,
+    guidance_scale: float = 15,
+    device: str = "cuda",
+    xformers: bool = False,
+    sdp: bool = False,
+    lora_path: str = "",
+    lora_rank: int = 64,
+    lora_scale: float = 1.0,
+    seed: Optional[int] = None,
+    latents_path: str="",
+    noise_prior: float = 0.,
+    repeat_num: int = 1,
+):
+
+    with torch.autocast(device, dtype=torch.half):
+        # prepare models
+        pipe = initialize_pipeline(model, device, xformers, sdp, lora_path, lora_rank, lora_scale)
+
+        for i in range(repeat_num):
+            if seed is not None:
+                random_seed = seed
+                torch.manual_seed(seed)
+            else:
+                random_seed = random.randint(100, 10000000)
+                torch.manual_seed(random_seed)
+
+            # prepare input latents
+            init_latents = prepare_input_latents(
+                pipe=pipe,
+                batch_size=len(prompt),
+                num_frames=num_frames,
+                height=height,
+                width=width,
+                latents_path=latents_path,
+                noise_prior=noise_prior
+            )
+
+            video_frames = pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                width=width,
+                height=height,
+                num_frames=num_frames,
+                num_inference_steps=num_steps,
+                guidance_scale=guidance_scale,
+                latents=init_latents
+            ).frames
+            # =========================================
+            # ========= write outputs to file =========
+            # =========================================
+            os.makedirs(args.output_dir, exist_ok=True)
+
+            # save to mp4
+            export_to_video(video_frames, f"{out_name}_{random_seed}.mp4", args.fps)
+
+            # # save to gif
+            file_name = f"{out_name}_{random_seed}.gif"
+            imageio.mimsave(file_name, video_frames, 'GIF', duration=1000 * 1 / args.fps, loop=0)
+
+    return video_frames
+
+
+if __name__ == "__main__":
+    import decord
+
+    decord.bridge.set_bridge("torch")
+
+    # fmt: off
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-m", "--model", type=str, default='/Users/rui/data/models/zeroscope_v2_576w/',
+                        help="HuggingFace repository or path to model checkpoint directory")
+    parser.add_argument("-p", "--prompt", type=str, default=None, help="Text prompt to condition on")
+    parser.add_argument("-n", "--negative-prompt", type=str, default=None, help="Text prompt to condition against")
+    parser.add_argument("-o", "--output_dir", type=str, default="./outputs/inference", help="Directory to save output video to")
+    parser.add_argument("-B", "--batch-size", type=int, default=1, help="Batch size for inference")
+    parser.add_argument("-W", "--width", type=int, default=384, help="Width of output video")
+    parser.add_argument("-H", "--height", type=int, default=384, help="Height of output video")
+    parser.add_argument("-T", "--num-frames", type=int, default=16, help="Total number of frames to generate")
+    parser.add_argument("-s", "--num-steps", type=int, default=30, help="Number of diffusion steps to run per frame.")
+    parser.add_argument("-g", "--guidance-scale", type=float, default=12, help="Scale for guidance loss (higher values = more guidance, but possibly more artifacts).")
+    parser.add_argument("-f", "--fps", type=int, default=8, help="FPS of output video")
+    parser.add_argument("-d", "--device", type=str, default="cuda", help="Device to run inference on (defaults to cuda).")
+    parser.add_argument("-x", "--xformers", action="store_true", help="Use XFormers attnetion, a memory-efficient attention implementation (requires `pip install xformers`).")
+    parser.add_argument("-S", "--sdp", action="store_true", help="Use SDP attention, PyTorch's built-in memory-efficient attention implementation.")
+    parser.add_argument("-cf", "--checkpoint_folder", type=str, default=None, help="Path to Low Rank Adaptation checkpoint file (defaults to empty string, which uses no LoRA).")
+    parser.add_argument("-lr", "--lora_rank", type=int, default=32, help="Size of the LoRA checkpoint's projection matrix (defaults to 32).")
+    parser.add_argument("-ls", "--lora_scale", type=float, default=1.0, help="Scale of LoRAs.")
+    parser.add_argument("-r", "--seed", type=int, default=None, help="Random seed to make generations reproducible.")
+    parser.add_argument("-np", "--noise_prior", type=float, default=0., help="Random seed to make generations reproducible.")
+    parser.add_argument("-ci", "--checkpoint_index", type=int, default=None,
+                        help="Random seed to make generations reproducible.")
+    parser.add_argument("-rn", "--repeat_num", type=int, default=None,
+                        help="Random seed to make generations reproducible.")
+
+    args = parser.parse_args()
+    # fmt: on
+
+    # =========================================
+    # ====== validate and prepare inputs ======
+    # =========================================
+
+    # args.prompt = ["A firefighter standing in front of a burning forest captured with a dolly zoom.",
+    #                "A spaceman standing on the moon with earth behind him captured with a dolly zoom."]
+    args.prompt = "A person is riding a bicycle past the Eiffel Tower."
+    args.checkpoint_folder = './outputs/train/train_2023-12-02T11-45-22/'
+    args.checkpoint_index = 500
+    args.noise_prior = 0.
+    args.repeat_num = 10
+
+    out_name = f"{args.output_dir}/"
+    prompt = re.sub(r'[<>:"/\\|?*\x00-\x1F]', "_", args.prompt) if platform.system() == "Windows" else args.prompt
+    out_name += f"{prompt}".replace(' ','_').replace(',', '').replace('.', '')
+
+    args.prompt = [prompt] * args.batch_size
+    if args.negative_prompt is not None:
+        args.negative_prompt = [args.negative_prompt] * args.batch_size
+
+    # =========================================
+    # ============= sample videos =============
+    # =========================================
+
+    lora_path = f"{args.checkpoint_folder}/checkpoint-{args.checkpoint_index}/temporal/lora"
+    latents_folder = f"{args.checkpoint_folder}/cached_latents"
+    latents_path = f"{latents_folder}/{random.choice(os.listdir(latents_folder))}"
+    # if args.seed is None:
+    #     args.seed = random.randint(100, 10000000)
+    assert os.path.exists(lora_path)
+    video_frames = inference(
+        model=args.model,
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        width=args.width,
+        height=args.height,
+        num_frames=args.num_frames,
+        num_steps=args.num_steps,
+        guidance_scale=args.guidance_scale,
+        device=args.device,
+        xformers=args.xformers,
+        sdp=args.sdp,
+        lora_path=lora_path,
+        lora_rank=args.lora_rank,
+        lora_scale = args.lora_scale,
+        seed=args.seed,
+        latents_path=latents_path,
+        noise_prior=args.noise_prior,
+        repeat_num=args.repeat_num
+    )
+

MotionDirector_train.py

@@ -0,0 +1,1021 @@
+import argparse
+import datetime
+import logging
+import inspect
+import math
+import os
+import random
+import gc
+import copy
+
+from typing import Dict, Optional, Tuple
+from omegaconf import OmegaConf
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+import diffusers
+import transformers
+
+from torchvision import transforms
+from tqdm.auto import tqdm
+
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import set_seed
+
+from models.unet_3d_condition import UNet3DConditionModel
+from diffusers.models import AutoencoderKL
+from diffusers import DDIMScheduler, TextToVideoSDPipeline
+from diffusers.optimization import get_scheduler
+from diffusers.utils.import_utils import is_xformers_available
+from diffusers.models.attention_processor import AttnProcessor2_0, Attention
+from diffusers.models.attention import BasicTransformerBlock
+
+from transformers import CLIPTextModel, CLIPTokenizer
+from transformers.models.clip.modeling_clip import CLIPEncoder
+from utils.dataset import VideoJsonDataset, SingleVideoDataset, \
+    ImageDataset, VideoFolderDataset, CachedDataset
+from einops import rearrange, repeat
+from utils.lora_handler import LoraHandler
+from utils.lora import extract_lora_child_module
+from utils.ddim_utils import ddim_inversion
+import imageio
+import numpy as np
+
+
+already_printed_trainables = False
+
+logger = get_logger(__name__, log_level="INFO")
+
+
+def create_logging(logging, logger, accelerator):
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+
+
+def accelerate_set_verbose(accelerator):
+    if accelerator.is_local_main_process:
+        transformers.utils.logging.set_verbosity_warning()
+        diffusers.utils.logging.set_verbosity_info()
+    else:
+        transformers.utils.logging.set_verbosity_error()
+        diffusers.utils.logging.set_verbosity_error()
+
+
+def get_train_dataset(dataset_types, train_data, tokenizer):
+    train_datasets = []
+
+    # Loop through all available datasets, get the name, then add to list of data to process.
+    for DataSet in [VideoJsonDataset, SingleVideoDataset, ImageDataset, VideoFolderDataset]:
+        for dataset in dataset_types:
+            if dataset == DataSet.__getname__():
+                train_datasets.append(DataSet(**train_data, tokenizer=tokenizer))
+
+    if len(train_datasets) > 0:
+        return train_datasets
+    else:
+        raise ValueError("Dataset type not found: 'json', 'single_video', 'folder', 'image'")
+
+
+def extend_datasets(datasets, dataset_items, extend=False):
+    biggest_data_len = max(x.__len__() for x in datasets)
+    extended = []
+    for dataset in datasets:
+        if dataset.__len__() == 0:
+            del dataset
+            continue
+        if dataset.__len__() < biggest_data_len:
+            for item in dataset_items:
+                if extend and item not in extended and hasattr(dataset, item):
+                    print(f"Extending {item}")
+
+                    value = getattr(dataset, item)
+                    value *= biggest_data_len
+                    value = value[:biggest_data_len]
+
+                    setattr(dataset, item, value)
+
+                    print(f"New {item} dataset length: {dataset.__len__()}")
+                    extended.append(item)
+
+
+def export_to_video(video_frames, output_video_path, fps):
+    video_writer = imageio.get_writer(output_video_path, fps=fps)
+    for img in video_frames:
+        video_writer.append_data(np.array(img))
+    video_writer.close()
+
+
+def create_output_folders(output_dir, config):
+    now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
+    out_dir = os.path.join(output_dir, f"train_{now}")
+
+    os.makedirs(out_dir, exist_ok=True)
+    os.makedirs(f"{out_dir}/samples", exist_ok=True)
+    # OmegaConf.save(config, os.path.join(out_dir, 'config.yaml'))
+
+    return out_dir
+
+
+def load_primary_models(pretrained_model_path):
+    noise_scheduler = DDIMScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler")
+    tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
+    text_encoder = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder")
+    vae = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae")
+    unet = UNet3DConditionModel.from_pretrained(pretrained_model_path, subfolder="unet")
+
+    return noise_scheduler, tokenizer, text_encoder, vae, unet
+
+
+def unet_and_text_g_c(unet, text_encoder, unet_enable, text_enable):
+    unet._set_gradient_checkpointing(value=unet_enable)
+    text_encoder._set_gradient_checkpointing(CLIPEncoder, value=text_enable)
+
+
+def freeze_models(models_to_freeze):
+    for model in models_to_freeze:
+        if model is not None: model.requires_grad_(False)
+
+
+def is_attn(name):
+    return ('attn1' or 'attn2' == name.split('.')[-1])
+
+
+def set_processors(attentions):
+    for attn in attentions: attn.set_processor(AttnProcessor2_0())
+
+
+def set_torch_2_attn(unet):
+    optim_count = 0
+
+    for name, module in unet.named_modules():
+        if is_attn(name):
+            if isinstance(module, torch.nn.ModuleList):
+                for m in module:
+                    if isinstance(m, BasicTransformerBlock):
+                        set_processors([m.attn1, m.attn2])
+                        optim_count += 1
+    if optim_count > 0:
+        print(f"{optim_count} Attention layers using Scaled Dot Product Attention.")
+
+
+def handle_memory_attention(enable_xformers_memory_efficient_attention, enable_torch_2_attn, unet):
+    try:
+        is_torch_2 = hasattr(F, 'scaled_dot_product_attention')
+        enable_torch_2 = is_torch_2 and enable_torch_2_attn
+
+        if enable_xformers_memory_efficient_attention and not enable_torch_2:
+            if is_xformers_available():
+                from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
+                unet.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)
+            else:
+                raise ValueError("xformers is not available. Make sure it is installed correctly")
+
+        if enable_torch_2:
+            set_torch_2_attn(unet)
+
+    except:
+        print("Could not enable memory efficient attention for xformers or Torch 2.0.")
+
+
+def param_optim(model, condition, extra_params=None, is_lora=False, negation=None):
+    extra_params = extra_params if len(extra_params.keys()) > 0 else None
+    return {
+        "model": model,
+        "condition": condition,
+        'extra_params': extra_params,
+        'is_lora': is_lora,
+        "negation": negation
+    }
+
+
+def create_optim_params(name='param', params=None, lr=5e-6, extra_params=None):
+    params = {
+        "name": name,
+        "params": params,
+        "lr": lr
+    }
+    if extra_params is not None:
+        for k, v in extra_params.items():
+            params[k] = v
+
+    return params
+
+
+def negate_params(name, negation):
+    # We have to do this if we are co-training with LoRA.
+    # This ensures that parameter groups aren't duplicated.
+    if negation is None: return False
+    for n in negation:
+        if n in name and 'temp' not in name:
+            return True
+    return False
+
+
+def create_optimizer_params(model_list, lr):
+    import itertools
+    optimizer_params = []
+
+    for optim in model_list:
+        model, condition, extra_params, is_lora, negation = optim.values()
+        # Check if we are doing LoRA training.
+        if is_lora and condition and isinstance(model, list):
+            params = create_optim_params(
+                params=itertools.chain(*model),
+                extra_params=extra_params
+            )
+            optimizer_params.append(params)
+            continue
+
+        if is_lora and condition and not isinstance(model, list):
+            for n, p in model.named_parameters():
+                if 'lora' in n:
+                    params = create_optim_params(n, p, lr, extra_params)
+                    optimizer_params.append(params)
+            continue
+
+        # If this is true, we can train it.
+        if condition:
+            for n, p in model.named_parameters():
+                should_negate = 'lora' in n and not is_lora
+                if should_negate: continue
+
+                params = create_optim_params(n, p, lr, extra_params)
+                optimizer_params.append(params)
+
+    return optimizer_params
+
+
+def get_optimizer(use_8bit_adam):
+    if use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
+            )
+
+        return bnb.optim.AdamW8bit
+    else:
+        return torch.optim.AdamW
+
+
+def is_mixed_precision(accelerator):
+    weight_dtype = torch.float32
+
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    return weight_dtype
+
+
+def cast_to_gpu_and_type(model_list, accelerator, weight_dtype):
+    for model in model_list:
+        if model is not None: model.to(accelerator.device, dtype=weight_dtype)
+
+
+def inverse_video(pipe, latents, num_steps):
+    ddim_inv_scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+    ddim_inv_scheduler.set_timesteps(num_steps)
+
+    ddim_inv_latent = ddim_inversion(
+        pipe, ddim_inv_scheduler, video_latent=latents.to(pipe.device),
+        num_inv_steps=num_steps, prompt="")[-1]
+    return ddim_inv_latent
+
+
+def handle_cache_latents(
+        should_cache,
+        output_dir,
+        train_dataloader,
+        train_batch_size,
+        vae,
+        unet,
+        pretrained_model_path,
+        noise_prior,
+        cached_latent_dir=None,
+):
+    # Cache latents by storing them in VRAM.
+    # Speeds up training and saves memory by not encoding during the train loop.
+    if not should_cache: return None
+    vae.to('cuda', dtype=torch.float16)
+    vae.enable_slicing()
+
+    pipe = TextToVideoSDPipeline.from_pretrained(
+        pretrained_model_path,
+        vae=vae,
+        unet=copy.deepcopy(unet).to('cuda', dtype=torch.float16)
+    )
+    pipe.text_encoder.to('cuda', dtype=torch.float16)
+
+    cached_latent_dir = (
+        os.path.abspath(cached_latent_dir) if cached_latent_dir is not None else None
+    )
+
+    if cached_latent_dir is None:
+        cache_save_dir = f"{output_dir}/cached_latents"
+        os.makedirs(cache_save_dir, exist_ok=True)
+
+        for i, batch in enumerate(tqdm(train_dataloader, desc="Caching Latents.")):
+
+            save_name = f"cached_{i}"
+            full_out_path = f"{cache_save_dir}/{save_name}.pt"
+
+            pixel_values = batch['pixel_values'].to('cuda', dtype=torch.float16)
+            batch['latents'] = tensor_to_vae_latent(pixel_values, vae)
+            if noise_prior > 0.:
+                batch['inversion_noise'] = inverse_video(pipe, batch['latents'], 50)
+            for k, v in batch.items(): batch[k] = v[0]
+
+            torch.save(batch, full_out_path)
+            del pixel_values
+            del batch
+
+            # We do this to avoid fragmentation from casting latents between devices.
+            torch.cuda.empty_cache()
+    else:
+        cache_save_dir = cached_latent_dir
+
+    return torch.utils.data.DataLoader(
+        CachedDataset(cache_dir=cache_save_dir),
+        batch_size=train_batch_size,
+        shuffle=True,
+        num_workers=0
+    )
+
+
+def handle_trainable_modules(model, trainable_modules=None, is_enabled=True, negation=None):
+    global already_printed_trainables
+
+    # This can most definitely be refactored :-)
+    unfrozen_params = 0
+    if trainable_modules is not None:
+        for name, module in model.named_modules():
+            for tm in tuple(trainable_modules):
+                if tm == 'all':
+                    model.requires_grad_(is_enabled)
+                    unfrozen_params = len(list(model.parameters()))
+                    break
+
+                if tm in name and 'lora' not in name:
+                    for m in module.parameters():
+                        m.requires_grad_(is_enabled)
+                        if is_enabled: unfrozen_params += 1
+
+    if unfrozen_params > 0 and not already_printed_trainables:
+        already_printed_trainables = True
+        print(f"{unfrozen_params} params have been unfrozen for training.")
+
+
+def tensor_to_vae_latent(t, vae):
+    video_length = t.shape[1]
+
+    t = rearrange(t, "b f c h w -> (b f) c h w")
+    latents = vae.encode(t).latent_dist.sample()
+    latents = rearrange(latents, "(b f) c h w -> b c f h w", f=video_length)
+    latents = latents * 0.18215
+
+    return latents
+
+
+def sample_noise(latents, noise_strength, use_offset_noise=False):
+    b, c, f, *_ = latents.shape
+    noise_latents = torch.randn_like(latents, device=latents.device)
+
+    if use_offset_noise:
+        offset_noise = torch.randn(b, c, f, 1, 1, device=latents.device)
+        noise_latents = noise_latents + noise_strength * offset_noise
+
+    return noise_latents
+
+
+def enforce_zero_terminal_snr(betas):
+    """
+    Corrects noise in diffusion schedulers.
+    From: Common Diffusion Noise Schedules and Sample Steps are Flawed
+    https://arxiv.org/pdf/2305.08891.pdf
+    """
+    # Convert betas to alphas_bar_sqrt
+    alphas = 1 - betas
+    alphas_bar = alphas.cumprod(0)
+    alphas_bar_sqrt = alphas_bar.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
+    alphas = alphas_bar[1:] / alphas_bar[:-1]
+    alphas = torch.cat([alphas_bar[0:1], alphas])
+    betas = 1 - alphas
+
+    return betas
+
+
+def should_sample(global_step, validation_steps, validation_data):
+    return global_step % validation_steps == 0 and validation_data.sample_preview
+
+
+def save_pipe(
+        path,
+        global_step,
+        accelerator,
+        unet,
+        text_encoder,
+        vae,
+        output_dir,
+        lora_manager_spatial: LoraHandler,
+        lora_manager_temporal: LoraHandler,
+        unet_target_replace_module=None,
+        text_target_replace_module=None,
+        is_checkpoint=False,
+        save_pretrained_model=True
+):
+    if is_checkpoint:
+        save_path = os.path.join(output_dir, f"checkpoint-{global_step}")
+        os.makedirs(save_path, exist_ok=True)
+    else:
+        save_path = output_dir
+
+    # Save the dtypes so we can continue training at the same precision.
+    u_dtype, t_dtype, v_dtype = unet.dtype, text_encoder.dtype, vae.dtype
+
+    # Copy the model without creating a reference to it. This allows keeping the state of our lora training if enabled.
+    unet_out = copy.deepcopy(accelerator.unwrap_model(unet.cpu(), keep_fp32_wrapper=False))
+    text_encoder_out = copy.deepcopy(accelerator.unwrap_model(text_encoder.cpu(), keep_fp32_wrapper=False))
+
+    pipeline = TextToVideoSDPipeline.from_pretrained(
+        path,
+        unet=unet_out,
+        text_encoder=text_encoder_out,
+        vae=vae,
+    ).to(torch_dtype=torch.float32)
+
+    lora_manager_spatial.save_lora_weights(model=copy.deepcopy(pipeline), save_path=save_path+'/spatial', step=global_step)
+    lora_manager_temporal.save_lora_weights(model=copy.deepcopy(pipeline), save_path=save_path+'/temporal', step=global_step)
+
+    if save_pretrained_model:
+        pipeline.save_pretrained(save_path)
+
+    if is_checkpoint:
+        unet, text_encoder = accelerator.prepare(unet, text_encoder)
+        models_to_cast_back = [(unet, u_dtype), (text_encoder, t_dtype), (vae, v_dtype)]
+        [x[0].to(accelerator.device, dtype=x[1]) for x in models_to_cast_back]
+
+    logger.info(f"Saved model at {save_path} on step {global_step}")
+
+    del pipeline
+    del unet_out
+    del text_encoder_out
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+def main(
+        pretrained_model_path: str,
+        output_dir: str,
+        train_data: Dict,
+        validation_data: Dict,
+        extra_train_data: list = [],
+        dataset_types: Tuple[str] = ('json'),
+        validation_steps: int = 100,
+        trainable_modules: Tuple[str] = None,  # Eg: ("attn1", "attn2")
+        extra_unet_params=None,
+        train_batch_size: int = 1,
+        max_train_steps: int = 500,
+        learning_rate: float = 5e-5,
+        lr_scheduler: str = "constant",
+        lr_warmup_steps: int = 0,
+        adam_beta1: float = 0.9,
+        adam_beta2: float = 0.999,
+        adam_weight_decay: float = 1e-2,
+        adam_epsilon: float = 1e-08,
+        gradient_accumulation_steps: int = 1,
+        gradient_checkpointing: bool = False,
+        text_encoder_gradient_checkpointing: bool = False,
+        checkpointing_steps: int = 500,
+        resume_from_checkpoint: Optional[str] = None,
+        resume_step: Optional[int] = None,
+        mixed_precision: Optional[str] = "fp16",
+        use_8bit_adam: bool = False,
+        enable_xformers_memory_efficient_attention: bool = True,
+        enable_torch_2_attn: bool = False,
+        seed: Optional[int] = None,
+        use_offset_noise: bool = False,
+        rescale_schedule: bool = False,
+        offset_noise_strength: float = 0.1,
+        extend_dataset: bool = False,
+        cache_latents: bool = False,
+        cached_latent_dir=None,
+        use_unet_lora: bool = False,
+        unet_lora_modules: Tuple[str] = [],
+        text_encoder_lora_modules: Tuple[str] = [],
+        save_pretrained_model: bool = True,
+        lora_rank: int = 16,
+        lora_path: str = '',
+        lora_unet_dropout: float = 0.1,
+        logger_type: str = 'tensorboard',
+        **kwargs
+):
+    *_, config = inspect.getargvalues(inspect.currentframe())
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=gradient_accumulation_steps,
+        mixed_precision=mixed_precision,
+        log_with=logger_type,
+        project_dir=output_dir
+    )
+
+    # Make one log on every process with the configuration for debugging.
+    create_logging(logging, logger, accelerator)
+
+    # Initialize accelerate, transformers, and diffusers warnings
+    accelerate_set_verbose(accelerator)
+
+    # Handle the output folder creation
+    if accelerator.is_main_process:
+        output_dir = create_output_folders(output_dir, config)
+
+    # Load scheduler, tokenizer and models.
+    noise_scheduler, tokenizer, text_encoder, vae, unet = load_primary_models(pretrained_model_path)
+
+    # Freeze any necessary models
+    freeze_models([vae, text_encoder, unet])
+
+    # Enable xformers if available
+    handle_memory_attention(enable_xformers_memory_efficient_attention, enable_torch_2_attn, unet)
+
+    # Initialize the optimizer
+    optimizer_cls = get_optimizer(use_8bit_adam)
+
+    # Get the training dataset based on types (json, single_video, image)
+    train_datasets = get_train_dataset(dataset_types, train_data, tokenizer)
+
+    # If you have extra train data, you can add a list of however many you would like.
+    # Eg: extra_train_data: [{: {dataset_types, train_data: {etc...}}}]
+    try:
+        if extra_train_data is not None and len(extra_train_data) > 0:
+            for dataset in extra_train_data:
+                d_t, t_d = dataset['dataset_types'], dataset['train_data']
+                train_datasets += get_train_dataset(d_t, t_d, tokenizer)
+
+    except Exception as e:
+        print(f"Could not process extra train datasets due to an error : {e}")
+
+    # Extend datasets that are less than the greatest one. This allows for more balanced training.
+    attrs = ['train_data', 'frames', 'image_dir', 'video_files']
+    extend_datasets(train_datasets, attrs, extend=extend_dataset)
+
+    # Process one dataset
+    if len(train_datasets) == 1:
+        train_dataset = train_datasets[0]
+
+    # Process many datasets
+    else:
+        train_dataset = torch.utils.data.ConcatDataset(train_datasets)
+
+    # Create parameters to optimize over with a condition (if "condition" is true, optimize it)
+    extra_unet_params = extra_unet_params if extra_unet_params is not None else {}
+    extra_text_encoder_params = extra_unet_params if extra_unet_params is not None else {}
+
+    # Use LoRA if enabled.
+    # one temporal lora
+    lora_manager_temporal = LoraHandler(use_unet_lora=use_unet_lora, unet_replace_modules=["TransformerTemporalModel"])
+
+    unet_lora_params_temporal, unet_negation_temporal = lora_manager_temporal.add_lora_to_model(
+        use_unet_lora, unet, lora_manager_temporal.unet_replace_modules, lora_unet_dropout,
+        lora_path + '/temporal/lora/', r=lora_rank)
+
+    optimizer_temporal = optimizer_cls(
+        create_optimizer_params([param_optim(unet_lora_params_temporal, use_unet_lora, is_lora=True,
+                                             extra_params={**{"lr": learning_rate}, **extra_text_encoder_params}
+                                             )], learning_rate),
+        lr=learning_rate,
+        betas=(adam_beta1, adam_beta2),
+        weight_decay=adam_weight_decay,
+        eps=adam_epsilon,
+    )
+
+    lr_scheduler_temporal = get_scheduler(
+        lr_scheduler,
+        optimizer=optimizer_temporal,
+        num_warmup_steps=lr_warmup_steps * gradient_accumulation_steps,
+        num_training_steps=max_train_steps * gradient_accumulation_steps,
+    )
+
+    # one spatial lora for each video
+    if 'folder' in dataset_types:
+        spatial_lora_num = train_dataset.__len__()
+    else:
+        spatial_lora_num = 1
+
+    lora_manager_spatials = []
+    unet_lora_params_spatial_list = []
+    optimizer_spatial_list = []
+    lr_scheduler_spatial_list = []
+    for i in range(spatial_lora_num):
+        lora_manager_spatial = LoraHandler(use_unet_lora=use_unet_lora, unet_replace_modules=["Transformer2DModel"])
+        lora_manager_spatials.append(lora_manager_spatial)
+        unet_lora_params_spatial, unet_negation_spatial = lora_manager_spatial.add_lora_to_model(
+            use_unet_lora, unet, lora_manager_spatial.unet_replace_modules, lora_unet_dropout,
+            lora_path + '/spatial/lora/', r=lora_rank)
+
+        unet_lora_params_spatial_list.append(unet_lora_params_spatial)
+
+        optimizer_spatial = optimizer_cls(
+            create_optimizer_params([param_optim(unet_lora_params_spatial, use_unet_lora, is_lora=True,
+                                                 extra_params={**{"lr": learning_rate}, **extra_text_encoder_params}
+                                                 )], learning_rate),
+            lr=learning_rate,
+            betas=(adam_beta1, adam_beta2),
+            weight_decay=adam_weight_decay,
+            eps=adam_epsilon,
+        )
+
+        optimizer_spatial_list.append(optimizer_spatial)
+
+        # Scheduler
+        lr_scheduler_spatial = get_scheduler(
+            lr_scheduler,
+            optimizer=optimizer_spatial,
+            num_warmup_steps=lr_warmup_steps * gradient_accumulation_steps,
+            num_training_steps=max_train_steps * gradient_accumulation_steps,
+        )
+        lr_scheduler_spatial_list.append(lr_scheduler_spatial)
+
+        unet_negation_all = unet_negation_spatial + unet_negation_temporal
+
+    # DataLoaders creation:
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset,
+        batch_size=train_batch_size,
+        shuffle=True
+    )
+
+    # Latents caching
+    cached_data_loader = handle_cache_latents(
+        cache_latents,
+        output_dir,
+        train_dataloader,
+        train_batch_size,
+        vae,
+        unet,
+        pretrained_model_path,
+        validation_data.noise_prior,
+        cached_latent_dir,
+    )
+
+    if cached_data_loader is not None:
+        train_dataloader = cached_data_loader
+
+    # Prepare everything with our `accelerator`.
+    unet, optimizer_spatial_list, optimizer_temporal, train_dataloader, lr_scheduler_spatial_list, lr_scheduler_temporal, text_encoder = accelerator.prepare(
+        unet,
+        optimizer_spatial_list, optimizer_temporal,
+        train_dataloader,
+        lr_scheduler_spatial_list, lr_scheduler_temporal,
+        text_encoder
+    )
+
+    # Use Gradient Checkpointing if enabled.
+    unet_and_text_g_c(
+        unet,
+        text_encoder,
+        gradient_checkpointing,
+        text_encoder_gradient_checkpointing
+    )
+
+    # Enable VAE slicing to save memory.
+    vae.enable_slicing()
+
+    # For mixed precision training we cast the text_encoder and vae weights to half-precision
+    # as these models are only used for inference, keeping weights in full precision is not required.
+    weight_dtype = is_mixed_precision(accelerator)
+
+    # Move text encoders, and VAE to GPU
+    models_to_cast = [text_encoder, vae]
+    cast_to_gpu_and_type(models_to_cast, accelerator, weight_dtype)
+
+    # Fix noise schedules to predcit light and dark areas if available.
+    if not use_offset_noise and rescale_schedule:
+        noise_scheduler.betas = enforce_zero_terminal_snr(noise_scheduler.betas)
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
+
+    # Afterwards we recalculate our number of training epochs
+    num_train_epochs = math.ceil(max_train_steps / num_update_steps_per_epoch)
+
+    # We need to initialize the trackers we use, and also store our configuration.
+    # The trackers initializes automatically on the main process.
+    if accelerator.is_main_process:
+        accelerator.init_trackers("text2video-fine-tune")
+
+    # Train!
+    total_batch_size = train_batch_size * accelerator.num_processes * gradient_accumulation_steps
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num Epochs = {num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {max_train_steps}")
+    global_step = 0
+    first_epoch = 0
+
+    # Only show the progress bar once on each machine.
+    progress_bar = tqdm(range(global_step, max_train_steps), disable=not accelerator.is_local_main_process)
+    progress_bar.set_description("Steps")
+
+    def finetune_unet(batch, step, mask_spatial_lora=False, mask_temporal_lora=False):
+        nonlocal use_offset_noise
+        nonlocal rescale_schedule
+
+        # Unfreeze UNET Layers
+        if global_step == 0:
+            already_printed_trainables = False
+            unet.train()
+            handle_trainable_modules(
+                unet,
+                trainable_modules,
+                is_enabled=True,
+                negation=unet_negation_all
+            )
+
+        # Convert videos to latent space
+        if not cache_latents:
+            latents = tensor_to_vae_latent(batch["pixel_values"], vae)
+        else:
+            latents = batch["latents"]
+
+        # Sample noise that we'll add to the latents
+        use_offset_noise = use_offset_noise and not rescale_schedule
+        noise = sample_noise(latents, offset_noise_strength, use_offset_noise)
+        bsz = latents.shape[0]
+
+        # Sample a random timestep for each video
+        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
+        timesteps = timesteps.long()
+
+        # Add noise to the latents according to the noise magnitude at each timestep
+        # (this is the forward diffusion process)
+        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+        # *Potentially* Fixes gradient checkpointing training.
+        # See: https://github.com/prigoyal/pytorch_memonger/blob/master/tutorial/Checkpointing_for_PyTorch_models.ipynb
+        if kwargs.get('eval_train', False):
+            unet.eval()
+            text_encoder.eval()
+
+        # Encode text embeddings
+        token_ids = batch['prompt_ids']
+        encoder_hidden_states = text_encoder(token_ids)[0]
+        detached_encoder_state = encoder_hidden_states.clone().detach()
+
+        # 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(latents, noise, timesteps)
+
+        else:
+            raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
+
+        encoder_hidden_states = detached_encoder_state
+
+        if mask_spatial_lora:
+            loras = extract_lora_child_module(unet, target_replace_module=["Transformer2DModel"])
+            for lora_i in loras:
+                lora_i.scale = 0.
+            loss_spatial = None
+        else:
+            loras = extract_lora_child_module(unet, target_replace_module=["Transformer2DModel"])
+            for lora_i in loras:
+                lora_i.scale = 1.
+
+            for lora_idx in range(0, len(loras), spatial_lora_num):
+                loras[lora_idx + step].scale = 1.
+
+            loras = extract_lora_child_module(unet, target_replace_module=["TransformerTemporalModel"])
+            for lora_i in loras:
+                lora_i.scale = 0.
+
+            ran_idx = torch.randint(0, noisy_latents.shape[2], (1,)).item()
+
+            if random.uniform(0, 1) < -0.5:
+                pixel_values_spatial = transforms.functional.hflip(
+                    batch["pixel_values"][:, ran_idx, :, :, :]).unsqueeze(1)
+                latents_spatial = tensor_to_vae_latent(pixel_values_spatial, vae)
+                noise_spatial = sample_noise(latents_spatial, offset_noise_strength, use_offset_noise)
+                noisy_latents_input = noise_scheduler.add_noise(latents_spatial, noise_spatial, timesteps)
+                target_spatial = noise_spatial
+                model_pred_spatial = unet(noisy_latents_input, timesteps,
+                                          encoder_hidden_states=encoder_hidden_states).sample
+                loss_spatial = F.mse_loss(model_pred_spatial[:, :, 0, :, :].float(),
+                                          target_spatial[:, :, 0, :, :].float(), reduction="mean")
+            else:
+                noisy_latents_input = noisy_latents[:, :, ran_idx, :, :]
+                target_spatial = target[:, :, ran_idx, :, :]
+                model_pred_spatial = unet(noisy_latents_input.unsqueeze(2), timesteps,
+                                          encoder_hidden_states=encoder_hidden_states).sample
+                loss_spatial = F.mse_loss(model_pred_spatial[:, :, 0, :, :].float(),
+                                          target_spatial.float(), reduction="mean")
+
+        if mask_temporal_lora:
+            loras = extract_lora_child_module(unet, target_replace_module=["TransformerTemporalModel"])
+            for lora_i in loras:
+                lora_i.scale = 0.
+            loss_temporal = None
+        else:
+            loras = extract_lora_child_module(unet, target_replace_module=["TransformerTemporalModel"])
+            for lora_i in loras:
+                lora_i.scale = 1.
+            model_pred = unet(noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states).sample
+            loss_temporal = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+
+            beta = 1
+            alpha = (beta ** 2 + 1) ** 0.5
+            ran_idx = torch.randint(0, model_pred.shape[2], (1,)).item()
+            model_pred_decent = alpha * model_pred - beta * model_pred[:, :, ran_idx, :, :].unsqueeze(2)
+            target_decent = alpha * target - beta * target[:, :, ran_idx, :, :].unsqueeze(2)
+            loss_ad_temporal = F.mse_loss(model_pred_decent.float(), target_decent.float(), reduction="mean")
+            loss_temporal = loss_temporal + loss_ad_temporal
+
+        return loss_spatial, loss_temporal, latents, noise
+
+    for epoch in range(first_epoch, num_train_epochs):
+        train_loss_spatial = 0.0
+        train_loss_temporal = 0.0
+
+        for step, batch in enumerate(train_dataloader):
+            # Skip steps until we reach the resumed step
+            if resume_from_checkpoint and epoch == first_epoch and step < resume_step:
+                if step % gradient_accumulation_steps == 0:
+                    progress_bar.update(1)
+                continue
+
+            with accelerator.accumulate(unet), accelerator.accumulate(text_encoder):
+
+                text_prompt = batch['text_prompt'][0]
+
+                for optimizer_spatial in optimizer_spatial_list:
+                    optimizer_spatial.zero_grad(set_to_none=True)
+
+                optimizer_temporal.zero_grad(set_to_none=True)
+
+                mask_temporal_lora = False
+                # mask_spatial_lora = False
+                mask_spatial_lora = random.uniform(0, 1) < 0.1 and not mask_temporal_lora
+
+                with accelerator.autocast():
+                    loss_spatial, loss_temporal, latents, init_noise = finetune_unet(batch, step, mask_spatial_lora=mask_spatial_lora, mask_temporal_lora=mask_temporal_lora)
+
+                # Gather the losses across all processes for logging (if we use distributed training).
+                if not mask_spatial_lora:
+                    avg_loss_spatial = accelerator.gather(loss_spatial.repeat(train_batch_size)).mean()
+                    train_loss_spatial += avg_loss_spatial.item() / gradient_accumulation_steps
+
+                if not mask_temporal_lora:
+                    avg_loss_temporal = accelerator.gather(loss_temporal.repeat(train_batch_size)).mean()
+                    train_loss_temporal += avg_loss_temporal.item() / gradient_accumulation_steps
+
+                # Backpropagate
+                if not mask_spatial_lora:
+                    accelerator.backward(loss_spatial, retain_graph = True)
+                    optimizer_spatial_list[step].step()
+
+                if not mask_temporal_lora:
+                    accelerator.backward(loss_temporal)
+                    optimizer_temporal.step()
+
+                lr_scheduler_spatial_list[step].step()
+                lr_scheduler_temporal.step()
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+                accelerator.log({"train_loss": train_loss_temporal}, step=global_step)
+                train_loss_temporal = 0.0
+                if global_step % checkpointing_steps == 0 and global_step > 0:
+                    save_pipe(
+                        pretrained_model_path,
+                        global_step,
+                        accelerator,
+                        unet,
+                        text_encoder,
+                        vae,
+                        output_dir,
+                        lora_manager_spatial,
+                        lora_manager_temporal,
+                        unet_lora_modules,
+                        text_encoder_lora_modules,
+                        is_checkpoint=True,
+                        save_pretrained_model=save_pretrained_model
+                    )
+
+                if should_sample(global_step, validation_steps, validation_data):
+                    if accelerator.is_main_process:
+                        with accelerator.autocast():
+                            unet.eval()
+                            text_encoder.eval()
+                            unet_and_text_g_c(unet, text_encoder, False, False)
+                            loras = extract_lora_child_module(unet, target_replace_module=["Transformer2DModel"])
+                            for lora_i in loras:
+                                lora_i.scale = validation_data.spatial_scale
+
+                            if validation_data.noise_prior > 0:
+                                preset_noise = (validation_data.noise_prior) ** 0.5 * batch['inversion_noise'] + (
+                                    1-validation_data.noise_prior) ** 0.5 * torch.randn_like(batch['inversion_noise'])
+                            else:
+                                preset_noise = None
+
+                            pipeline = TextToVideoSDPipeline.from_pretrained(
+                                pretrained_model_path,
+                                text_encoder=text_encoder,
+                                vae=vae,
+                                unet=unet
+                            )
+
+                            diffusion_scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
+                            pipeline.scheduler = diffusion_scheduler
+
+                            prompt_list = text_prompt if len(validation_data.prompt) <= 0 else validation_data.prompt
+                            for prompt in prompt_list:
+                                save_filename = f"{global_step}_{prompt.replace('.', '')}"
+
+                                out_file = f"{output_dir}/samples/{save_filename}.mp4"
+
+                                with torch.no_grad():
+                                    video_frames = pipeline(
+                                        prompt,
+                                        width=validation_data.width,
+                                        height=validation_data.height,
+                                        num_frames=validation_data.num_frames,
+                                        num_inference_steps=validation_data.num_inference_steps,
+                                        guidance_scale=validation_data.guidance_scale,
+                                        latents=preset_noise
+                                    ).frames
+                                export_to_video(video_frames, out_file, train_data.get('fps', 8))
+                                logger.info(f"Saved a new sample to {out_file}")
+                            del pipeline
+                            torch.cuda.empty_cache()
+
+                    unet_and_text_g_c(
+                        unet,
+                        text_encoder,
+                        gradient_checkpointing,
+                        text_encoder_gradient_checkpointing
+                    )
+
+            accelerator.log({"loss_temporal": loss_temporal.detach().item()}, step=step)
+
+            if global_step >= max_train_steps:
+                break
+
+    # Create the pipeline using the trained modules and save it.
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+        save_pipe(
+            pretrained_model_path,
+            global_step,
+            accelerator,
+            unet,
+            text_encoder,
+            vae,
+            output_dir,
+            lora_manager_spatial,
+            lora_manager_temporal,
+            unet_lora_modules,
+            text_encoder_lora_modules,
+            is_checkpoint=False,
+            save_pretrained_model=save_pretrained_model
+        )
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str, default='./configs/config_multi_videos.yaml')
+    args = parser.parse_args()
+    main(**OmegaConf.load(args.config))
+

README.md

@@ -1,13 +1,364 @@
----
-title: MotionDirector
-emoji: 🏆
-colorFrom: green
-colorTo: yellow
-sdk: gradio
-sdk_version: 4.8.0
-app_file: app.py
-pinned: false
-license: apache-2.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# MotionDirector
+
+This is the official repository of [MotionDirector](https://showlab.github.io/MotionDirector).
+
+**MotionDirector: Motion Customization of Text-to-Video Diffusion Models.**
+<br/>
+[Rui Zhao](https://ruizhaocv.github.io/),
+[Yuchao Gu](https://ycgu.site/), 
+[Jay Zhangjie Wu](https://zhangjiewu.github.io/), 
+[David Junhao Zhang](https://junhaozhang98.github.io/),
+[Jiawei Liu](https://jia-wei-liu.github.io/),
+[Weijia Wu](https://weijiawu.github.io/),
+[Jussi Keppo](https://www.jussikeppo.com/),
+[Mike Zheng Shou](https://sites.google.com/view/showlab)
+<br/>
+
+[![Project Page](https://img.shields.io/badge/Project-Website-orange)](https://showlab.github.io/MotionDirector)
+[![arXiv](https://img.shields.io/badge/arXiv-MotionDirector-b31b1b.svg)](https://arxiv.org/abs/2310.08465)
+
+<p align="center">
+<img src="https://github.com/showlab/MotionDirector/blob/page/assets/teaser.gif" width="1080px"/>  
+<br>
+<em>MotionDirector can customize text-to-video diffusion models to generate videos with desired motions.</em>
+</p>
+
+<table class="center">
+<tr>
+  <td style="text-align:center;" colspan="4"><b>Astronaut's daily life on Mars (Motion concepts learned by MotionDirector)</b></td>
+</tr>
+<tr>
+<td style="text-align:center;"><b>Lifting Weights</b></td>
+<td style="text-align:center;"><b>Playing Golf</b></td>
+<td style="text-align:center;"><b>Riding Horse</b></td>
+<td style="text-align:center;"><b>Riding Bicycle</b></td>
+</tr>
+<tr>
+  <td><img src=assets/astronaut_mars/An_astronaut_is_lifting_weights_on_Mars_4K_high_quailty_highly_detailed_4008521.gif></td>
+  <td><img src=assets/astronaut_mars/Astronaut_playing_golf_on_Mars_659514.gif></td>
+  <td><img src=assets/astronaut_mars/An_astronaut_is_riding_a_horse_on_Mars_4K_high_quailty_highly_detailed_1913261.gif></td>              
+  <td><img src=assets/astronaut_mars/An_astronaut_is_riding_a_bicycle_past_the_pyramids_Mars_4K_high_quailty_highly_detailed_5532778.gif></td>
+</tr>
+<tr>
+  <td width=25% style="text-align:center;">"An astronaut is lifting weights on Mars, 4K, high quailty, highly detailed.” </br> seed: 4008521</td>
+  <td width=25% style="text-align:center;">"Astronaut playing golf on Mars” </br> seed: 659514</td>
+  <td width=25% style="text-align:center;">"An astronaut is riding a horse on Mars, 4K, high quailty, highly detailed."  </br> seed: 1913261</td>
+  <td width=25% style="text-align:center;">"An astronaut is riding a bicycle past the pyramids Mars, 4K, high quailty, highly detailed."  </br> seed: 5532778</td>
+<tr>
+</table>
+
+## News
+- [2023.12.06] [MotionDirector for Sports](#MotionDirector_for_Sports) released! Lifting weights, riding horse, palying golf, etc.
+- [2023.12.05] [Colab demo](https://github.com/camenduru/MotionDirector-colab) is available. Thanks to [Camenduru](https://twitter.com/camenduru).
+- [2023.12.04] [MotionDirector for Cinematic Shots](#MotionDirector_for_Cinematic_Shots) released. Now, you can make AI films with professional cinematic shots!
+- [2023.12.02] Code and model weights released!
+
+## ToDo
+- [ ] Gradio Demo
+- [ ] More trained weights of MotionDirector
+
+## Setup
+### Requirements
+
+```shell
+# create virtual environment
+conda create -n motiondirector python=3.8
+conda activate motiondirector
+# install packages
+pip install -r requirements.txt
+```
+
+### Weights of Foundation Models
+```shell
+git lfs install
+## You can choose the ModelScopeT2V or ZeroScope, etc., as the foundation model.
+## ZeroScope
+git clone https://huggingface.co/cerspense/zeroscope_v2_576w ./models/zeroscope_v2_576w/
+## ModelScopeT2V
+git clone https://huggingface.co/damo-vilab/text-to-video-ms-1.7b ./models/model_scope/
+```
+### Weights of trained MotionDirector <a name="download_weights"></a>
+```shell
+# Make sure you have git-lfs installed (https://git-lfs.com)
+git lfs install
+git clone https://huggingface.co/ruizhaocv/MotionDirector_weights ./outputs
+```
+
+## Usage
+### Training
+
+#### Train MotionDirector on multiple videos:
+```bash
+python MotionDirector_train.py --config ./configs/config_multi_videos.yaml
+```
+#### Train MotionDirector on a single video:
+```bash
+python MotionDirector_train.py --config ./configs/config_single_video.yaml
+```
+
+Note:  
+- Before running the above command, 
+make sure you replace the path to foundational model weights and training data with your own in the config files `config_multi_videos.yaml` or `config_single_video.yaml`.
+- Generally, training on multiple 16-frame videos usually takes `300~500` steps, about `9~16` minutes using one A5000 GPU. Training on a single video takes `50~150` steps, about `1.5~4.5` minutes using one A5000 GPU. The required VRAM for training is around `14GB`.
+- Reduce `n_sample_frames` if your GPU memory is limited.
+- Reduce the learning rate and increase the training steps for better performance.
+
+
+### Inference
+```bash
+python MotionDirector_inference.py --model /path/to/the/foundation/model  --prompt "Your prompt" --checkpoint_folder /path/to/the/trained/MotionDirector --checkpoint_index 300 --noise_prior 0.
+```
+Note: 
+- Replace `/path/to/the/foundation/model` with your own path to the foundation model, like ZeroScope.
+- The value of `checkpoint_index` means the checkpoint saved at which the training step is selected.
+- The value of `noise_prior` indicates how much the inversion noise of the reference video affects the generation. 
+We recommend setting it to `0` for MotionDirector trained on multiple videos to achieve the highest diverse generation, while setting it to `0.1~0.5` for MotionDirector trained on a single video for faster convergence and better alignment with the reference video.
+
+
+## Inference with pre-trained MotionDirector
+All available weights are at official [Huggingface Repo](https://huggingface.co/ruizhaocv/MotionDirector_weights).
+Run the [download command](#download_weights), the weights will be downloaded to the folder `outputs`, then run the following inference command to generate videos.
+
+### MotionDirector trained on multiple videos:
+```bash
+python MotionDirector_inference.py --model /path/to/the/ZeroScope  --prompt "A person is riding a bicycle past the Eiffel Tower." --checkpoint_folder ./outputs/train/riding_bicycle/ --checkpoint_index 300 --noise_prior 0. --seed 7192280
+```
+Note:  
+- Replace `/path/to/the/ZeroScope` with your own path to the foundation model, i.e. the ZeroScope.
+- Change the `prompt` to generate different videos. 
+- The `seed` is set to a random value by default. Set it to a specific value will obtain certain results, as provided in the table below.
+
+Results:
+
+<table class="center">
+<tr>
+  <td style="text-align:center;"><b>Reference Videos</b></td>
+  <td style="text-align:center;" colspan="3"><b>Videos Generated by MotionDirector</b></td>
+</tr>
+<tr>
+  <td><img src=assets/multi_videos_results/reference_videos.gif></td>
+  <td><img src=assets/multi_videos_results/A_person_is_riding_a_bicycle_past_the_Eiffel_Tower_7192280.gif></td>
+  <td><img src=assets/multi_videos_results/A_panda_is_riding_a_bicycle_in_a_garden_2178639.gif></td>              
+  <td><img src=assets/multi_videos_results/An_alien_is_riding_a_bicycle_on_Mars_2390886.gif></td>
+</tr>
+<tr>
+  <td width=25% style="text-align:center;color:gray;">"A person is riding a bicycle."</td>
+  <td width=25% style="text-align:center;">"A person is riding a bicycle past the Eiffel Tower.” </br> seed: 7192280</td>
+  <td width=25% style="text-align:center;">"A panda is riding a bicycle in a garden."  </br> seed: 2178639</td>
+  <td width=25% style="text-align:center;">"An alien is riding a bicycle on Mars."  </br> seed: 2390886</td>
+</table>
+
+### MotionDirector trained on a single video:
+16 frames:
+```bash
+python MotionDirector_inference.py --model /path/to/the/ZeroScope  --prompt "A tank is running on the moon." --checkpoint_folder ./outputs/train/car_16/ --checkpoint_index 150 --noise_prior 0.5 --seed 8551187
+```
+<table class="center">
+<tr>
+  <td style="text-align:center;"><b>Reference Video</b></td>
+  <td style="text-align:center;" colspan="3"><b>Videos Generated by MotionDirector</b></td>
+</tr>
+<tr>
+  <td><img src=assets/single_video_results/reference_video.gif></td>
+  <td><img src=assets/single_video_results/A_tank_is_running_on_the_moon_8551187.gif></td>
+  <td><img src=assets/single_video_results/A_lion_is_running_past_the_pyramids_431554.gif></td>              
+  <td><img src=assets/single_video_results/A_spaceship_is_flying_past_Mars_8808231.gif></td>
+</tr>
+<tr>
+  <td width=25% style="text-align:center;color:gray;">"A car is running on the road."</td>
+  <td width=25% style="text-align:center;">"A tank is running on the moon.” </br> seed: 8551187</td>
+  <td width=25% style="text-align:center;">"A lion is running past the pyramids." </br> seed: 431554</td>
+  <td width=25% style="text-align:center;">"A spaceship is flying past Mars."  </br> seed: 8808231</td>
+</tr>
+</table>
+
+24 frames:
+```bash
+python MotionDirector_inference.py --model /path/to/the/ZeroScope  --prompt "A truck is running past the Arc de Triomphe." --checkpoint_folder ./outputs/train/car_24/ --checkpoint_index 150 --noise_prior 0.5 --width 576 --height 320 --num-frames 24 --seed 34543
+```
+<table class="center">
+<tr>
+  <td style="text-align:center;"><b>Reference Video</b></td>
+  <td style="text-align:center;" colspan="3"><b>Videos Generated by MotionDirector</b></td>
+</tr>
+<tr>
+  <td><img src=assets/single_video_results/24_frames/reference_video.gif></td>
+  <td><img src=assets/single_video_results/24_frames/A_truck_is_running_past_the_Arc_de_Triomphe_34543.gif></td>
+  <td><img src=assets/single_video_results/24_frames/An_elephant_is_running_in_a_forest_2171736.gif></td>              
+</tr>
+<tr>
+  <td width=25% style="text-align:center;color:gray;">"A car is running on the road."</td>
+  <td width=25% style="text-align:center;">"A truck is running past the Arc de Triomphe.” </br> seed: 34543</td>
+  <td width=25% style="text-align:center;">"An elephant is running in a forest." </br> seed: 2171736</td>
+ </tr>
+<tr>
+  <td><img src=assets/single_video_results/24_frames/reference_video.gif></td>
+  <td><img src=assets/single_video_results/24_frames/A_person_on_a_camel_is_running_past_the_pyramids_4904126.gif></td>              
+  <td><img src=assets/single_video_results/24_frames/A_spacecraft_is_flying_past_the_Milky_Way_galaxy_3235677.gif></td>
+</tr>
+<tr>
+  <td width=25% style="text-align:center;color:gray;">"A car is running on the road."</td>
+  <td width=25% style="text-align:center;">"A person on a camel is running past the pyramids." </br> seed: 4904126</td>
+  <td width=25% style="text-align:center;">"A spacecraft is flying past the Milky Way galaxy."  </br> seed: 3235677</td>
+</tr>
+</table>
+
+## MotionDirector for Sports <a name="MotionDirector_for_Sports"></a>
+
+```bash
+python MotionDirector_inference.py --model /path/to/the/ZeroScope  --prompt "A panda is lifting weights in a garden." --checkpoint_folder ./outputs/train/lifting_weights/ --checkpoint_index 300 --noise_prior 0. --seed 9365597
+```
+<table class="center">
+<tr>
+  <td style="text-align:center;" colspan="4"><b>Videos Generated by MotionDirector</b></td>
+</tr>
+<tr>
+<td style="text-align:center;" colspan="2"><b>Lifting Weights</b></td>
+<td style="text-align:center;" colspan="2"><b>Riding Bicycle</b></td>
+</tr>
+<tr>
+  <td><img src=assets/sports_results/lifting_weights/A_panda_is_lifting_weights_in_a_garden_1699276.gif></td>
+  <td><img src=assets/sports_results/lifting_weights/A_police_officer_is_lifting_weights_in_front_of_the_police_station_6804745.gif></td>
+  <td><img src=assets/multi_videos_results/A_panda_is_riding_a_bicycle_in_a_garden_2178639.gif></td>              
+  <td><img src=assets/multi_videos_results/An_alien_is_riding_a_bicycle_on_Mars_2390886.gif></td>
+</tr>
+<tr>
+  <td width=25% style="text-align:center;">"A panda is lifting weights in a garden.” </br> seed: 1699276</td>
+  <td width=25% style="text-align:center;">"A police officer is lifting weights in front of the police station.” </br> seed: 6804745</td>
+  <td width=25% style="text-align:center;">"A panda is riding a bicycle in a garden."  </br> seed: 2178639</td>
+  <td width=25% style="text-align:center;">"An alien is riding a bicycle on Mars."  </br> seed: 2390886</td>
+<tr>
+<td style="text-align:center;" colspan="2"><b>Riding Horse</b></td>
+<td style="text-align:center;" colspan="2"><b>Playing Golf</b></td>
+</tr>
+<tr>
+  <td><img src=assets/sports_results/riding_horse/A_Royal_Guard_riding_a_horse_in_front_of_Buckingham_Palace_4490970.gif></td>
+  <td><img src=assets/sports_results/riding_horse/A_man_riding_an_elephant_through_the_jungle_6230765.gif></td>
+  <td><img src=assets/sports_results/playing_golf/A_man_is_playing_golf_in_front_of_the_White_House_8870450.gif></td>              
+  <td><img src=assets/sports_results/playing_golf/A_monkey_is_playing_golf_on_a_field_full_of_flowers_2989633.gif></td>
+</tr>
+<tr>
+  <td width=25% style="text-align:center;">"A Royal Guard riding a horse in front of Buckingham Palace.” </br> seed: 4490970</td>
+  <td width=25% style="text-align:center;">"A man riding an elephant through the jungle.” </br> seed: 6230765</td>
+  <td width=25% style="text-align:center;">"A man is playing golf in front of the White House."  </br> seed: 8870450</td>
+  <td width=25% style="text-align:center;">"A monkey is playing golf on a field full of flowers."  </br> seed: 2989633</td>
+<tr>
+</table>
+
+More sports, to be continued ...
+
+## MotionDirector for Cinematic Shots <a name="MotionDirector_for_Cinematic_Shots"></a>
+
+### 1. Zoom
+#### 1.1 Dolly Zoom (Hitchcockian Zoom)
+```bash
+python MotionDirector_inference.py --model /path/to/the/ZeroScope  --prompt "A firefighter standing in front of a burning forest captured with a dolly zoom." --checkpoint_folder ./outputs/train/dolly_zoom/ --checkpoint_index 150 --noise_prior 0.5 --seed 9365597
+```
+<table class="center">
+<tr>
+  <td style="text-align:center;"><b>Reference Video</b></td>
+  <td style="text-align:center;" colspan="3"><b>Videos Generated by MotionDirector</b></td>
+</tr>
+<tr>
+  <td><img src=assets/cinematic_shots_results/dolly_zoom_16.gif></td>
+  <td><img src=assets/cinematic_shots_results/A_firefighter_standing_in_front_of_a_burning_forest_captured_with_a_dolly_zoom_9365597.gif></td>
+  <td><img src=assets/cinematic_shots_results/A_lion_sitting_on_top_of_a_cliff_captured_with_a_dolly_zoom_1675932.gif></td>              
+  <td><img src=assets/cinematic_shots_results/A_Roman_soldier_standing_in_front_of_the_Colosseum_captured_with_a_dolly_zoom_2310805.gif></td>
+</tr>
+<tr>
+  <td width=25% style="text-align:center;color:gray;">"A man standing in room captured with a dolly zoom."</td>
+  <td width=25% style="text-align:center;">"A firefighter standing in front of a burning forest captured with a dolly zoom." </br> seed: 9365597 </br> noise_prior: 0.5</td>
+  <td width=25% style="text-align:center;">"A lion sitting on top of a cliff captured with a dolly zoom." </br> seed: 1675932 </br> noise_prior: 0.5</td>
+  <td width=25% style="text-align:center;">"A Roman soldier standing in front of the Colosseum captured with a dolly zoom."  </br> seed: 2310805 </br> noise_prior: 0.5 </td>
+</tr>
+<tr>
+  <td><img src=assets/cinematic_shots_results/dolly_zoom_16.gif></td>
+  <td><img src=assets/cinematic_shots_results/A_firefighter_standing_in_front_of_a_burning_forest_captured_with_a_dolly_zoom_4615820.gif></td>
+  <td><img src=assets/cinematic_shots_results/A_lion_sitting_on_top_of_a_cliff_captured_with_a_dolly_zoom_4114896.gif></td>              
+  <td><img src=assets/cinematic_shots_results/A_Roman_soldier_standing_in_front_of_the_Colosseum_captured_with_a_dolly_zoom_7492004.gif></td>
+</tr>
+<tr>
+  <td width=25% style="text-align:center;color:gray;">"A man standing in room captured with a dolly zoom."</td>
+  <td width=25% style="text-align:center;">"A firefighter standing in front of a burning forest captured with a dolly zoom." </br> seed: 4615820 </br> noise_prior: 0.3</td>
+  <td width=25% style="text-align:center;">"A lion sitting on top of a cliff captured with a dolly zoom." </br> seed: 4114896 </br> noise_prior: 0.3</td>
+  <td width=25% style="text-align:center;">"A Roman soldier standing in front of the Colosseum captured with a dolly zoom."  </br> seed: 7492004</td>
+</tr>
+</table>
+
+#### 1.2 Zoom In
+The reference video is shot with my own water cup. You can also pick up your cup or any other object to practice camera movements and turn it into imaginative videos. Create your AI films with customized camera movements!
+
+```bash
+python MotionDirector_inference.py --model /path/to/the/ZeroScope  --prompt "A firefighter standing in front of a burning forest captured with a zoom in." --checkpoint_folder ./outputs/train/zoom_in/ --checkpoint_index 150 --noise_prior 0.3 --seed 1429227
+```
+<table class="center">
+<tr>
+  <td style="text-align:center;"><b>Reference Video</b></td>
+  <td style="text-align:center;" colspan="3"><b>Videos Generated by MotionDirector</b></td>
+</tr>
+<tr>
+  <td><img src=assets/cinematic_shots_results/zoom_in_16.gif></td>
+  <td><img src=assets/cinematic_shots_results/A_firefighter_standing_in_front_of_a_burning_forest_captured_with_a_zoom_in_1429227.gif></td>
+  <td><img src=assets/cinematic_shots_results/A_lion_sitting_on_top_of_a_cliff_captured_with_a_zoom_in_487239.gif></td>              
+  <td><img src=assets/cinematic_shots_results/A_Roman_soldier_standing_in_front_of_the_Colosseum_captured_with_a_zoom_in_1393184.gif></td>
+</tr>
+<tr>
+  <td width=25% style="text-align:center;color:gray;">"A cup in a lab captured with a zoom in."</td>
+  <td width=25% style="text-align:center;">"A firefighter standing in front of a burning forest captured with a zoom in." </br> seed: 1429227</td>
+  <td width=25% style="text-align:center;">"A lion sitting on top of a cliff captured with a zoom in." </br> seed: 487239 </td>
+  <td width=25% style="text-align:center;">"A Roman soldier standing in front of the Colosseum captured with a zoom in."  </br> seed: 1393184</td>
+</tr>
+</table>
+
+#### 1.3 Zoom Out
+```bash
+python MotionDirector_inference.py --model /path/to/the/ZeroScope  --prompt "A firefighter standing in front of a burning forest captured with a zoom out." --checkpoint_folder ./outputs/train/zoom_out/ --checkpoint_index 150 --noise_prior 0.3 --seed 4971910
+```
+<table class="center">
+<tr>
+  <td style="text-align:center;"><b>Reference Video</b></td>
+  <td style="text-align:center;" colspan="3"><b>Videos Generated by MotionDirector</b></td>
+</tr>
+<tr>
+  <td><img src=assets/cinematic_shots_results/zoom_out_16.gif></td>
+  <td><img src=assets/cinematic_shots_results/A_firefighter_standing_in_front_of_a_burning_forest_captured_with_a_zoom_out_4971910.gif></td>
+  <td><img src=assets/cinematic_shots_results/A_lion_sitting_on_top_of_a_cliff_captured_with_a_zoom_out_1767994.gif></td>              
+  <td><img src=assets/cinematic_shots_results/A_Roman_soldier_standing_in_front_of_the_Colosseum_captured_with_a_zoom_out_8203639.gif></td>
+</tr>
+<tr>
+  <td width=25% style="text-align:center;color:gray;">"A cup in a lab captured with a zoom out."</td>
+  <td width=25% style="text-align:center;">"A firefighter standing in front of a burning forest captured with a zoom out." </br> seed: 4971910</td>
+  <td width=25% style="text-align:center;">"A lion sitting on top of a cliff captured with a zoom out." </br> seed: 1767994 </td>
+  <td width=25% style="text-align:center;">"A Roman soldier standing in front of the Colosseum captured with a zoom out."  </br> seed: 8203639</td>
+</tr>
+</table>
+
+More Cinematic Shots, to be continued ....
+
+## More results
+
+If you have a more impressive MotionDirector or generated videos, please feel free to open an issue and share them with us. We would greatly appreciate it.
+Improvements to the code are also highly welcome.
+
+Please refer to [Project Page](https://showlab.github.io/MotionDirector) for more results.
+
+
+## Citation
+
+
+```bibtex
+
+@article{zhao2023motiondirector,
+  title={MotionDirector: Motion Customization of Text-to-Video Diffusion Models},
+  author={Zhao, Rui and Gu, Yuchao and Wu, Jay Zhangjie and Zhang, David Junhao and Liu, Jiawei and Wu, Weijia and Keppo, Jussi and Shou, Mike Zheng},
+  journal={arXiv preprint arXiv:2310.08465},
+  year={2023}
+}
+
+```
+
+## Shoutouts
+
+- This code builds on [diffusers](https://github.com/huggingface/diffusers) and [Text-To-Video-Finetuning](https://github.com/ExponentialML/Text-To-Video-Finetuning). Thanks for open-sourcing!
+- Thanks to [camenduru](https://twitter.com/camenduru) for the [colab demo](https://github.com/camenduru/MotionDirector-colab).
+- Thanks to [yhyu13](https://github.com/yhyu13) for the [Huggingface Repo](https://huggingface.co/Yhyu13/MotionDirector_LoRA).

app.py

@@ -0,0 +1,100 @@
+import gradio as gr
+import os
+from demo.motiondirector import MotionDirector
+
+from huggingface_hub import snapshot_download
+
+snapshot_download(repo_id="cerspense/zeroscope_v2_576w", local_dir="./zeroscope_v2_576w/")
+snapshot_download(repo_id="ruizhaocv/MotionDirector", local_dir="./MotionDirector_pretrained")
+
+is_spaces = True if "SPACE_ID" in os.environ else False
+true_for_shared_ui = False  # This will be true only if you are in a shared UI
+if (is_spaces):
+    true_for_shared_ui = True if "ruizhaocv/MotionDirector" in os.environ['SPACE_ID'] else False
+
+
+
+runner = MotionDirector()
+
+
+def motiondirector(model_select, text_pormpt, neg_text_pormpt, random_seed=1, steps=25, guidance_scale=7.5, baseline_select=False):
+    return runner(model_select, text_pormpt, neg_text_pormpt, int(random_seed) if random_seed != "" else 1, int(steps), float(guidance_scale), baseline_select)
+
+
+with gr.Blocks() as demo:
+    gr.HTML(
+        """
+        <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
+        <a href="https://github.com/showlab/MotionDirector" style="margin-right: 20px; text-decoration: none; display: flex; align-items: center;">
+        </a>
+        <div>
+            <h1 >MotionDirector: Motion Customization of Text-to-Video Diffusion Models</h1>
+            <h5 style="margin: 0;">More MotionDirectors are on the way. Stay tuned 🔥! Give us a star ✨ on Github for the latest update.</h5>
+            </br>
+            <div style="display: flex; justify-content: center; align-items: center; text-align: center;>
+                <a href="https://arxiv.org/abs/2310.08465"><img src="https://img.shields.io/badge/arXiv-MotionDirector-b31b1b.svg"></a>&nbsp;&nbsp;
+                <a href='https://showlab.github.io/MotionDirector'><img src='https://img.shields.io/badge/Project_Page-MotionDirector-green'></a>&nbsp;&nbsp;
+                <a href='https://github.com/showlab/MotionDirector'><img src='https://img.shields.io/badge/Github-MotionDirector-blue'></a>&nbsp;&nbsp;
+            </div>
+        </div>
+        </div>
+        """)
+    with gr.Row():
+        generated_video_baseline = gr.Video(format="mp4", label="Video Generated by base model (ZeroScope with same seed)")
+        generated_video = gr.Video(format="mp4", label="Video Generated by MotionDirector")
+
+        with gr.Column():
+            baseline_select = gr.Checkbox(label="Compare with baseline (ZeroScope with same seed)", info="Run baseline? Note: Inference time will be doubled.")
+            random_seed = gr.Textbox(label="Random seed", value=1, info="default: 1")
+            sampling_steps = gr.Textbox(label="Sampling steps", value=30, info="default: 30")
+            guidance_scale = gr.Textbox(label="Guidance scale", value=12, info="default: 12")
+
+    with gr.Row():
+        model_select = gr.Dropdown(
+            ["1-1: [Cinematic Shots] -- Dolly Zoom (Hitchcockian Zoom)",
+             "1-2: [Cinematic Shots] -- Zoom In",
+             "1-3: [Cinematic Shots] -- Zoom Out",
+             "2-1: [Object Trajectory] -- Right to Left",
+             "2-2: [Object Trajectory] -- Left to Right",
+             "3-1: [Sports Concepts] -- Riding Bicycle",
+             "3-2: [Sports Concepts] -- Riding Horse",
+             "3-3: [Sports Concepts] -- Lifting Weights",
+             "3-4: [Sports Concepts] -- Playing Golf"
+             ],
+            label="MotionDirector",
+            info="Which MotionDirector would you like to use!"
+        )
+
+        text_pormpt = gr.Textbox(label="Text Prompt", value='', placeholder="Input your text prompt here!")
+        neg_text_pormpt = gr.Textbox(label="Negative Text Prompt", value='', placeholder="default: None")
+
+    submit = gr.Button("Generate")
+
+    # when the `submit` button is clicked
+    submit.click(
+        motiondirector,
+        [model_select, text_pormpt, neg_text_pormpt, random_seed, sampling_steps, guidance_scale, baseline_select],
+        [generated_video, generated_video_baseline]
+    )
+
+    # Examples
+    gr.Markdown("## Examples")
+    gr.Examples(
+        fn=motiondirector,
+        examples=[
+            ["1-1: [Cinematic Shots] -- Dolly Zoom (Hitchcockian Zoom)", "A lion sitting on top of a cliff captured with a dolly zoom.", 1675932],
+            ["1-2: [Cinematic Shots] -- Zoom In", "A firefighter standing in front of a burning forest captured with a zoom in.", 1429227],
+            ["1-3: [Cinematic Shots] -- Zoom Out", "A lion sitting on top of a cliff captured with a zoom out.", 1767994],
+            ["2-1: [Object Trajectory] -- Right to Left", "A tank is running on the moon.", 8551187],
+            ["2-2: [Object Trajectory] -- Left to Right", "A tiger is running in the forest.", 3463673],
+            ["3-1: [Sports Concepts] -- Riding Bicycle", "An astronaut is riding a bicycle past the pyramids Mars 4K high quailty highly detailed.", 4422954],
+            ["3-2: [Sports Concepts] -- Riding Horse", "A man riding an elephant through the jungle.", 6230765],
+            ["3-3: [Sports Concepts] -- Lifting Weights", "A panda is lifting weights in a garden.", 1699276],
+            ["3-4: [Sports Concepts] -- Playing Golf", "A man is playing golf in front of the White House.", 8870450],
+        ],
+        inputs=[model_select, text_pormpt, random_seed],
+        outputs=generated_video,
+    )
+
+demo.queue(max_size=15)
+demo.launch(share=True)

demo/MotionDirector_gradio.py

@@ -0,0 +1,92 @@
+import argparse
+import imageio
+import numpy as np
+import gradio as gr
+from PIL import Image
+
+from demo.motiondirector import MotionDirector
+
+runner = MotionDirector()
+
+
+def motiondirector(model_select, text_pormpt, neg_text_pormpt, random_seed=1, steps=25, guidance_scale=7.5, baseline_select=False):
+    return runner(model_select, text_pormpt, neg_text_pormpt, int(random_seed) if random_seed != "" else 1, int(steps), float(guidance_scale), baseline_select)
+
+
+with gr.Blocks() as demo:
+    gr.HTML(
+        """
+        <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
+        <a href="https://github.com/showlab/MotionDirector" style="margin-right: 20px; text-decoration: none; display: flex; align-items: center;">
+        </a>
+        <div>
+            <h1 >MotionDirector: Motion Customization of Text-to-Video Diffusion Models</h1>
+            <h5 style="margin: 0;">More MotionDirectors are on the way. Stay tuned 🔥! Give us a star ✨ on Github for the latest update.</h5>
+            </br>
+            <div style="display: flex; justify-content: center; align-items: center; text-align: center;>
+                <a href="https://arxiv.org/abs/2310.08465"><img src="https://img.shields.io/badge/arXiv-MotionDirector-b31b1b.svg"></a>&nbsp;&nbsp;
+                <a href='https://showlab.github.io/MotionDirector'><img src='https://img.shields.io/badge/Project_Page-MotionDirector-green'></a>&nbsp;&nbsp;
+                <a href='https://github.com/showlab/MotionDirector'><img src='https://img.shields.io/badge/Github-MotionDirector-blue'></a>&nbsp;&nbsp;
+            </div>
+        </div>
+        </div>
+        """)
+    with gr.Row():
+        generated_video_baseline = gr.Video(format="mp4", label="Video Generated by base model (ZeroScope with same seed)")
+        generated_video = gr.Video(format="mp4", label="Video Generated by MotionDirector")
+
+        with gr.Column():
+            baseline_select = gr.Checkbox(label="Compare with baseline (ZeroScope with same seed)", info="Run baseline? Note: Inference time will be doubled.")
+            random_seed = gr.Textbox(label="Random seed", value=1, info="default: 1")
+            sampling_steps = gr.Textbox(label="Sampling steps", value=30, info="default: 30")
+            guidance_scale = gr.Textbox(label="Guidance scale", value=12, info="default: 12")
+
+    with gr.Row():
+        model_select = gr.Dropdown(
+            ["1-1: [Cinematic Shots] -- Dolly Zoom (Hitchcockian Zoom)",
+             "1-2: [Cinematic Shots] -- Zoom In",
+             "1-3: [Cinematic Shots] -- Zoom Out",
+             "2-1: [Object Trajectory] -- Right to Left",
+             "2-2: [Object Trajectory] -- Left to Right",
+             "3-1: [Sports Concepts] -- Riding Bicycle",
+             "3-2: [Sports Concepts] -- Riding Horse",
+             "3-3: [Sports Concepts] -- Lifting Weights",
+             "3-4: [Sports Concepts] -- Playing Golf"
+             ],
+            label="MotionDirector",
+            info="Which MotionDirector would you like to use!"
+        )
+
+        text_pormpt = gr.Textbox(label="Text Prompt", value='', placeholder="Input your text prompt here!")
+        neg_text_pormpt = gr.Textbox(label="Negative Text Prompt", value='', placeholder="default: None")
+
+    submit = gr.Button("Generate")
+
+    # when the `submit` button is clicked
+    submit.click(
+        motiondirector,
+        [model_select, text_pormpt, neg_text_pormpt, random_seed, sampling_steps, guidance_scale, baseline_select],
+        [generated_video, generated_video_baseline]
+    )
+
+    # Examples
+    gr.Markdown("## Examples")
+    gr.Examples(
+        fn=motiondirector,
+        examples=[
+            ["1-1: [Cinematic Shots] -- Dolly Zoom (Hitchcockian Zoom)", "A lion sitting on top of a cliff captured with a dolly zoom.", 1675932],
+            ["1-2: [Cinematic Shots] -- Zoom In", "A firefighter standing in front of a burning forest captured with a zoom in.", 1429227],
+            ["1-3: [Cinematic Shots] -- Zoom Out", "A lion sitting on top of a cliff captured with a zoom out.", 1767994],
+            ["2-1: [Object Trajectory] -- Right to Left", "A tank is running on the moon.", 8551187],
+            ["2-2: [Object Trajectory] -- Left to Right", "A tiger is running in the forest.", 3463673],
+            ["3-1: [Sports Concepts] -- Riding Bicycle", "An astronaut is riding a bicycle past the pyramids Mars 4K high quailty highly detailed.", 4422954],
+            ["3-2: [Sports Concepts] -- Riding Horse", "A man riding an elephant through the jungle.", 6230765],
+            ["3-3: [Sports Concepts] -- Lifting Weights", "A panda is lifting weights in a garden.", 1699276],
+            ["3-4: [Sports Concepts] -- Playing Golf", "A man is playing golf in front of the White House.", 8870450],
+        ],
+        inputs=[model_select, text_pormpt, random_seed],
+        outputs=generated_video,
+    )
+
+demo.queue(max_size=15)
+demo.launch(share=True)

demo/motiondirector.py

@@ -0,0 +1,218 @@
+import os
+import warnings
+from typing import Optional
+
+import torch
+from diffusers import DDIMScheduler, TextToVideoSDPipeline
+from einops import rearrange
+from torch import Tensor
+from torch.nn.functional import interpolate
+from tqdm import trange
+import random
+
+from MotionDirector_train import export_to_video, handle_memory_attention, load_primary_models, unet_and_text_g_c, freeze_models
+from utils.lora_handler import LoraHandler
+from utils.ddim_utils import ddim_inversion
+from utils.lora import extract_lora_child_module
+import imageio
+
+
+def initialize_pipeline(
+    model: str,
+    device: str = "cuda",
+    xformers: bool = True,
+    sdp: bool = True,
+    lora_path: str = "",
+    lora_rank: int = 32,
+    lora_scale: float = 1.0,
+):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+
+        scheduler, tokenizer, text_encoder, vae, unet = load_primary_models(model)
+
+    # Freeze any necessary models
+    freeze_models([vae, text_encoder, unet])
+
+    # Enable xformers if available
+    handle_memory_attention(xformers, sdp, unet)
+
+    lora_manager_temporal = LoraHandler(
+        version="cloneofsimo",
+        use_unet_lora=True,
+        use_text_lora=False,
+        save_for_webui=False,
+        only_for_webui=False,
+        unet_replace_modules=["TransformerTemporalModel"],
+        text_encoder_replace_modules=None,
+        lora_bias=None
+    )
+
+    unet_lora_params, unet_negation = lora_manager_temporal.add_lora_to_model(
+        True, unet, lora_manager_temporal.unet_replace_modules, 0, lora_path, r=lora_rank, scale=lora_scale)
+
+    unet.eval()
+    text_encoder.eval()
+    unet_and_text_g_c(unet, text_encoder, False, False)
+
+    pipe = TextToVideoSDPipeline.from_pretrained(
+        pretrained_model_name_or_path=model,
+        scheduler=scheduler,
+        tokenizer=tokenizer,
+        text_encoder=text_encoder.to(device=device, dtype=torch.half),
+        vae=vae.to(device=device, dtype=torch.half),
+        unet=unet.to(device=device, dtype=torch.half),
+    )
+    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+
+    return pipe
+
+
+def inverse_video(pipe, latents, num_steps):
+    ddim_inv_scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+    ddim_inv_scheduler.set_timesteps(num_steps)
+
+    ddim_inv_latent = ddim_inversion(
+        pipe, ddim_inv_scheduler, video_latent=latents.to(pipe.device),
+        num_inv_steps=num_steps, prompt="")[-1]
+    return ddim_inv_latent
+
+
+def prepare_input_latents(
+    pipe: TextToVideoSDPipeline,
+    batch_size: int,
+    num_frames: int,
+    height: int,
+    width: int,
+    latents_path:str,
+    noise_prior: float
+):
+    # initialize with random gaussian noise
+    scale = pipe.vae_scale_factor
+    shape = (batch_size, pipe.unet.config.in_channels, num_frames, height // scale, width // scale)
+    if noise_prior > 0.:
+        cached_latents = torch.load(latents_path)
+        if 'inversion_noise' not in cached_latents:
+            latents = inverse_video(pipe, cached_latents['latents'].unsqueeze(0), 50).squeeze(0)
+        else:
+            latents = torch.load(latents_path)['inversion_noise'].unsqueeze(0)
+        if latents.shape[0] != batch_size:
+            latents = latents.repeat(batch_size, 1, 1, 1, 1)
+        if latents.shape != shape:
+            latents = interpolate(rearrange(latents, "b c f h w -> (b f) c h w", b=batch_size), (height // scale, width // scale), mode='bilinear')
+            latents = rearrange(latents, "(b f) c h w -> b c f h w", b=batch_size)
+        noise = torch.randn_like(latents, dtype=torch.half)
+        latents_base = noise
+        latents = (noise_prior) ** 0.5 * latents + (1 - noise_prior) ** 0.5 * noise
+    else:
+        latents = torch.randn(shape, dtype=torch.half)
+        latents_base = latents
+
+    return latents, latents_base
+
+
+class MotionDirector():
+    def __init__(self):
+        self.version = "0.0.0"
+        self.foundation_model_path = "./zeroscope_v2_576w/"
+        self.lora_path = "./MotionDirector_pretrained/dolly_zoom_(hitchcockian_zoom)/checkpoint-default/temporal/lora"
+        with torch.autocast("cuda", dtype=torch.half):
+            self.pipe = initialize_pipeline(model=self.foundation_model_path, lora_path=self.lora_path, lora_scale=1)
+
+    def reload_lora(self, lora_path):
+        if lora_path != self.lora_path:
+            self.lora_path = lora_path
+            with torch.autocast("cuda", dtype=torch.half):
+                self.pipe = initialize_pipeline(model=self.foundation_model_path, lora_path=self.lora_path)
+
+    def __call__(self, model_select, text_pormpt, neg_text_pormpt, random_seed, steps, guidance_scale, baseline_select):
+        model_select = str(model_select)
+        out_name = f"./outputs/inference"
+        out_name += f"{text_pormpt}".replace(' ', '_').replace(',', '').replace('.', '')
+
+        model_select_type = model_select.split('--')[1].strip()
+        model_select_type = model_select_type.lower().replace(' ', '_')
+
+        lora_path = f"./MotionDirector_pretrained/{model_select_type}/checkpoint-default/temporal/lora"
+        self.reload_lora(lora_path)
+        latents_folder = f"./MotionDirector_pretrained/{model_select_type}/cached_latents"
+        latents_path = f"{latents_folder}/{random.choice(os.listdir(latents_folder))}"
+        assert os.path.exists(lora_path)
+
+        if '3-' in model_select:
+            noise_prior = 0.
+        elif '2-' in model_select:
+            noise_prior = 0.5
+        else:
+            noise_prior = 0.3
+
+        if random_seed > 1000:
+            torch.manual_seed(random_seed)
+        else:
+            random_seed = random.randint(100, 10000000)
+            torch.manual_seed(random_seed)
+        device = "cuda"
+        with torch.autocast(device, dtype=torch.half):
+            # prepare input latents
+            with torch.no_grad():
+                init_latents,init_latents_base = prepare_input_latents(
+                    pipe=self.pipe,
+                    batch_size=1,
+                    num_frames=16,
+                    height=384,
+                    width=384,
+                    latents_path=latents_path,
+                    noise_prior=noise_prior
+                )
+                video_frames = self.pipe(
+                    prompt=text_pormpt,
+                    negative_prompt=neg_text_pormpt,
+                    width=384,
+                    height=384,
+                    num_frames=16,
+                    num_inference_steps=steps,
+                    guidance_scale=guidance_scale,
+                    latents=init_latents
+                ).frames
+
+
+                out_file = f"{out_name}_{random_seed}.mp4"
+                os.makedirs(os.path.dirname(out_file), exist_ok=True)
+                export_to_video(video_frames, out_file, 8)
+
+                if baseline_select:
+                    with torch.autocast("cuda", dtype=torch.half):
+
+                        loras = extract_lora_child_module(self.pipe.unet, target_replace_module=["TransformerTemporalModel"])
+                        for lora_i in loras:
+                            lora_i.scale = 0.
+
+                        # self.pipe = initialize_pipeline(model=self.foundation_model_path, lora_path=self.lora_path,
+                        #                                 lora_scale=0.)
+                        with torch.no_grad():
+                            video_frames = self.pipe(
+                                prompt=text_pormpt,
+                                negative_prompt=neg_text_pormpt,
+                                width=384,
+                                height=384,
+                                num_frames=16,
+                                num_inference_steps=steps,
+                                guidance_scale=guidance_scale,
+                                latents=init_latents_base,
+                            ).frames
+
+                            out_file_baseline = f"{out_name}_{random_seed}_baseline.mp4"
+                            os.makedirs(os.path.dirname(out_file_baseline), exist_ok=True)
+                            export_to_video(video_frames, out_file_baseline, 8)
+                    # with torch.autocast("cuda", dtype=torch.half):
+                    #     self.pipe = initialize_pipeline(model=self.foundation_model_path, lora_path=self.lora_path,
+                    #                                     lora_scale=1.)
+                    loras = extract_lora_child_module(self.pipe.unet,
+                                                      target_replace_module=["TransformerTemporalModel"])
+                    for lora_i in loras:
+                        lora_i.scale = 1.
+
+                else:
+                    out_file_baseline = None
+
+        return [out_file, out_file_baseline]

models/unet_3d_blocks.py

@@ -0,0 +1,842 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.utils.checkpoint as checkpoint
+from torch import nn
+from diffusers.models.resnet import Downsample2D, ResnetBlock2D, TemporalConvLayer, Upsample2D
+from diffusers.models.transformer_2d import Transformer2DModel
+from diffusers.models.transformer_temporal import TransformerTemporalModel
+
+# Assign gradient checkpoint function to simple variable for readability.
+g_c = checkpoint.checkpoint
+
+def use_temporal(module, num_frames, x):
+    if num_frames == 1:
+        if isinstance(module, TransformerTemporalModel):
+            return {"sample": x}
+        else:
+            return x
+
+def custom_checkpoint(module, mode=None):
+    if mode == None: raise ValueError('Mode for gradient checkpointing cannot be none.')
+    custom_forward = None
+
+    if mode == 'resnet':
+        def custom_forward(hidden_states, temb):
+            inputs = module(hidden_states, temb)
+            return inputs
+
+    if mode == 'attn':
+        def custom_forward(
+            hidden_states, 
+            encoder_hidden_states=None, 
+            cross_attention_kwargs=None
+        ):
+            inputs = module(
+                hidden_states,
+                encoder_hidden_states,
+                cross_attention_kwargs
+            )
+            return inputs
+
+    if mode == 'temp':
+         def custom_forward(hidden_states, num_frames=None):
+            inputs = use_temporal(module, num_frames, hidden_states)
+            if inputs is None: inputs = module(
+                hidden_states, 
+                num_frames=num_frames
+            )
+            return inputs
+
+    return custom_forward
+
+def transformer_g_c(transformer, sample, num_frames):
+    sample = g_c(custom_checkpoint(transformer, mode='temp'), 
+        sample, num_frames, use_reentrant=False
+    )['sample']
+
+    return sample
+
+def cross_attn_g_c(
+        attn, 
+        temp_attn, 
+        resnet, 
+        temp_conv, 
+        hidden_states, 
+        encoder_hidden_states, 
+        cross_attention_kwargs, 
+        temb, 
+        num_frames,
+        inverse_temp=False
+    ):
+    
+    def ordered_g_c(idx):
+
+        # Self and CrossAttention
+        if idx == 0: return g_c(custom_checkpoint(attn, mode='attn'),
+            hidden_states, encoder_hidden_states,cross_attention_kwargs, use_reentrant=False
+        )['sample']
+
+        # Temporal Self and CrossAttention
+        if idx == 1: return g_c(custom_checkpoint(temp_attn, mode='temp'), 
+            hidden_states, num_frames, use_reentrant=False)['sample']
+
+        # Resnets
+        if idx == 2: return g_c(custom_checkpoint(resnet, mode='resnet'), 
+            hidden_states, temb, use_reentrant=False)
+        
+        # Temporal Convolutions
+        if idx == 3: return g_c(custom_checkpoint(temp_conv, mode='temp'), 
+            hidden_states, num_frames, use_reentrant=False
+    )
+
+    # Here we call the function depending on the order in which they are called. 
+    # For some layers, the orders are different, so we access the appropriate one by index.
+    
+    if not inverse_temp:
+        for idx in [0,1,2,3]: hidden_states = ordered_g_c(idx) 
+    else:
+        for idx in [2,3,0,1]: hidden_states = ordered_g_c(idx)
+
+    return hidden_states
+
+def up_down_g_c(resnet, temp_conv, hidden_states, temb, num_frames):
+    hidden_states = g_c(custom_checkpoint(resnet, mode='resnet'), hidden_states, temb, use_reentrant=False)
+    hidden_states = g_c(custom_checkpoint(temp_conv, mode='temp'), 
+        hidden_states, num_frames,  use_reentrant=False
+    )
+    return hidden_states
+
+def get_down_block(
+    down_block_type,
+    num_layers,
+    in_channels,
+    out_channels,
+    temb_channels,
+    add_downsample,
+    resnet_eps,
+    resnet_act_fn,
+    attn_num_head_channels,
+    resnet_groups=None,
+    cross_attention_dim=None,
+    downsample_padding=None,
+    dual_cross_attention=False,
+    use_linear_projection=True,
+    only_cross_attention=False,
+    upcast_attention=False,
+    resnet_time_scale_shift="default",
+):
+    if down_block_type == "DownBlock3D":
+        return DownBlock3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            temb_channels=temb_channels,
+            add_downsample=add_downsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            downsample_padding=downsample_padding,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    elif down_block_type == "CrossAttnDownBlock3D":
+        if cross_attention_dim is None:
+            raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock3D")
+        return CrossAttnDownBlock3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            temb_channels=temb_channels,
+            add_downsample=add_downsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            downsample_padding=downsample_padding,
+            cross_attention_dim=cross_attention_dim,
+            attn_num_head_channels=attn_num_head_channels,
+            dual_cross_attention=dual_cross_attention,
+            use_linear_projection=use_linear_projection,
+            only_cross_attention=only_cross_attention,
+            upcast_attention=upcast_attention,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    raise ValueError(f"{down_block_type} does not exist.")
+
+
+def get_up_block(
+    up_block_type,
+    num_layers,
+    in_channels,
+    out_channels,
+    prev_output_channel,
+    temb_channels,
+    add_upsample,
+    resnet_eps,
+    resnet_act_fn,
+    attn_num_head_channels,
+    resnet_groups=None,
+    cross_attention_dim=None,
+    dual_cross_attention=False,
+    use_linear_projection=True,
+    only_cross_attention=False,
+    upcast_attention=False,
+    resnet_time_scale_shift="default",
+):
+    if up_block_type == "UpBlock3D":
+        return UpBlock3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    elif up_block_type == "CrossAttnUpBlock3D":
+        if cross_attention_dim is None:
+            raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock3D")
+        return CrossAttnUpBlock3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            cross_attention_dim=cross_attention_dim,
+            attn_num_head_channels=attn_num_head_channels,
+            dual_cross_attention=dual_cross_attention,
+            use_linear_projection=use_linear_projection,
+            only_cross_attention=only_cross_attention,
+            upcast_attention=upcast_attention,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    raise ValueError(f"{up_block_type} does not exist.")
+
+
+class UNetMidBlock3DCrossAttn(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        attn_num_head_channels=1,
+        output_scale_factor=1.0,
+        cross_attention_dim=1280,
+        dual_cross_attention=False,
+        use_linear_projection=True,
+        upcast_attention=False,
+    ):
+        super().__init__()
+
+        self.gradient_checkpointing = False
+        self.has_cross_attention = True
+        self.attn_num_head_channels = attn_num_head_channels
+        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+
+        # there is always at least one resnet
+        resnets = [
+            ResnetBlock2D(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                temb_channels=temb_channels,
+                eps=resnet_eps,
+                groups=resnet_groups,
+                dropout=dropout,
+                time_embedding_norm=resnet_time_scale_shift,
+                non_linearity=resnet_act_fn,
+                output_scale_factor=output_scale_factor,
+                pre_norm=resnet_pre_norm,
+            )
+        ]
+        temp_convs = [
+            TemporalConvLayer(
+                in_channels,
+                in_channels,
+                dropout=0.1
+            )
+        ]
+        attentions = []
+        temp_attentions = []
+
+        for _ in range(num_layers):
+            attentions.append(
+                Transformer2DModel(
+                    in_channels // attn_num_head_channels,
+                    attn_num_head_channels,
+                    in_channels=in_channels,
+                    num_layers=1,
+                    cross_attention_dim=cross_attention_dim,
+                    norm_num_groups=resnet_groups,
+                    use_linear_projection=use_linear_projection,
+                    upcast_attention=upcast_attention,
+                )
+            )
+            temp_attentions.append(
+                TransformerTemporalModel(
+                    in_channels // attn_num_head_channels,
+                    attn_num_head_channels,
+                    in_channels=in_channels,
+                    num_layers=1,
+                    cross_attention_dim=cross_attention_dim,
+                    norm_num_groups=resnet_groups,
+                )
+            )
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+            temp_convs.append(
+                TemporalConvLayer(
+                    in_channels,
+                    in_channels,
+                    dropout=0.1
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+        self.temp_convs = nn.ModuleList(temp_convs)
+        self.attentions = nn.ModuleList(attentions)
+        self.temp_attentions = nn.ModuleList(temp_attentions)
+
+    def forward(
+        self,
+        hidden_states,
+        temb=None,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        num_frames=1,
+        cross_attention_kwargs=None,
+    ):
+        if self.gradient_checkpointing:
+            hidden_states = up_down_g_c(
+                    self.resnets[0], 
+                    self.temp_convs[0], 
+                    hidden_states, 
+                    temb, 
+                    num_frames
+                )
+        else:
+            hidden_states = self.resnets[0](hidden_states, temb)
+            hidden_states = self.temp_convs[0](hidden_states, num_frames=num_frames)
+            
+        for attn, temp_attn, resnet, temp_conv in zip(
+            self.attentions, self.temp_attentions, self.resnets[1:], self.temp_convs[1:]
+        ):
+            if self.gradient_checkpointing:
+                hidden_states = cross_attn_g_c(
+                        attn, 
+                        temp_attn, 
+                        resnet, 
+                        temp_conv, 
+                        hidden_states, 
+                        encoder_hidden_states, 
+                        cross_attention_kwargs, 
+                        temb, 
+                        num_frames
+                    )
+            else:
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                ).sample
+                
+                if num_frames > 1:
+                    hidden_states = temp_attn(hidden_states, num_frames=num_frames).sample
+
+                hidden_states = resnet(hidden_states, temb)
+
+                if num_frames > 1:
+                    hidden_states = temp_conv(hidden_states, num_frames=num_frames)
+
+        return hidden_states
+
+
+class CrossAttnDownBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        attn_num_head_channels=1,
+        cross_attention_dim=1280,
+        output_scale_factor=1.0,
+        downsample_padding=1,
+        add_downsample=True,
+        dual_cross_attention=False,
+        use_linear_projection=False,
+        only_cross_attention=False,
+        upcast_attention=False,
+    ):
+        super().__init__()
+        resnets = []
+        attentions = []
+        temp_attentions = []
+        temp_convs = []
+
+        self.gradient_checkpointing = False
+        self.has_cross_attention = True
+        self.attn_num_head_channels = attn_num_head_channels
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+            temp_convs.append(
+                TemporalConvLayer(
+                    out_channels,
+                    out_channels,
+                    dropout=0.1
+                )
+            )
+            attentions.append(
+                Transformer2DModel(
+                    out_channels // attn_num_head_channels,
+                    attn_num_head_channels,
+                    in_channels=out_channels,
+                    num_layers=1,
+                    cross_attention_dim=cross_attention_dim,
+                    norm_num_groups=resnet_groups,
+                    use_linear_projection=use_linear_projection,
+                    only_cross_attention=only_cross_attention,
+                    upcast_attention=upcast_attention,
+                )
+            )
+            temp_attentions.append(
+                TransformerTemporalModel(
+                    out_channels // attn_num_head_channels,
+                    attn_num_head_channels,
+                    in_channels=out_channels,
+                    num_layers=1,
+                    cross_attention_dim=cross_attention_dim,
+                    norm_num_groups=resnet_groups,
+                )
+            )
+        self.resnets = nn.ModuleList(resnets)
+        self.temp_convs = nn.ModuleList(temp_convs)
+        self.attentions = nn.ModuleList(attentions)
+        self.temp_attentions = nn.ModuleList(temp_attentions)
+
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    Downsample2D(
+                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+    def forward(
+        self,
+        hidden_states,
+        temb=None,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        num_frames=1,
+        cross_attention_kwargs=None,
+    ):
+        # TODO(Patrick, William) - attention mask is not used
+        output_states = ()
+
+        for resnet, temp_conv, attn, temp_attn in zip(
+            self.resnets, self.temp_convs, self.attentions, self.temp_attentions
+        ):
+        
+            if self.gradient_checkpointing:
+                hidden_states = cross_attn_g_c(
+                        attn, 
+                        temp_attn, 
+                        resnet, 
+                        temp_conv, 
+                        hidden_states, 
+                        encoder_hidden_states, 
+                        cross_attention_kwargs, 
+                        temb, 
+                        num_frames,
+                        inverse_temp=True
+                    )
+            else:
+                hidden_states = resnet(hidden_states, temb)
+
+                if num_frames > 1:
+                    hidden_states = temp_conv(hidden_states, num_frames=num_frames)
+
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                ).sample
+
+                if num_frames > 1:
+                    hidden_states = temp_attn(hidden_states, num_frames=num_frames).sample
+
+            output_states += (hidden_states,)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+            output_states += (hidden_states,)
+
+        return hidden_states, output_states
+
+
+class DownBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        output_scale_factor=1.0,
+        add_downsample=True,
+        downsample_padding=1,
+    ):
+        super().__init__()
+        resnets = []
+        temp_convs = []
+
+        self.gradient_checkpointing = False
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+            temp_convs.append(
+                TemporalConvLayer(
+                    out_channels,
+                    out_channels,
+                    dropout=0.1
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+        self.temp_convs = nn.ModuleList(temp_convs)
+
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    Downsample2D(
+                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+    def forward(self, hidden_states, temb=None, num_frames=1):
+        output_states = ()
+
+        for resnet, temp_conv in zip(self.resnets, self.temp_convs):
+            if self.gradient_checkpointing:
+                hidden_states = up_down_g_c(resnet, temp_conv, hidden_states, temb, num_frames)
+            else:
+                hidden_states = resnet(hidden_states, temb)
+
+                if num_frames > 1:
+                    hidden_states = temp_conv(hidden_states, num_frames=num_frames)
+
+            output_states += (hidden_states,)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+            output_states += (hidden_states,)
+
+        return hidden_states, output_states
+
+
+class CrossAttnUpBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        prev_output_channel: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        attn_num_head_channels=1,
+        cross_attention_dim=1280,
+        output_scale_factor=1.0,
+        add_upsample=True,
+        dual_cross_attention=False,
+        use_linear_projection=False,
+        only_cross_attention=False,
+        upcast_attention=False,
+    ):
+        super().__init__()
+        resnets = []
+        temp_convs = []
+        attentions = []
+        temp_attentions = []
+
+        self.gradient_checkpointing = False
+        self.has_cross_attention = True
+        self.attn_num_head_channels = attn_num_head_channels
+
+        for i in range(num_layers):
+            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
+            resnet_in_channels = prev_output_channel if i == 0 else out_channels
+
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=resnet_in_channels + res_skip_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+            temp_convs.append(
+                TemporalConvLayer(
+                    out_channels,
+                    out_channels,
+                    dropout=0.1
+                )
+            )
+            attentions.append(
+                Transformer2DModel(
+                    out_channels // attn_num_head_channels,
+                    attn_num_head_channels,
+                    in_channels=out_channels,
+                    num_layers=1,
+                    cross_attention_dim=cross_attention_dim,
+                    norm_num_groups=resnet_groups,
+                    use_linear_projection=use_linear_projection,
+                    only_cross_attention=only_cross_attention,
+                    upcast_attention=upcast_attention,
+                )
+            )
+            temp_attentions.append(
+                TransformerTemporalModel(
+                    out_channels // attn_num_head_channels,
+                    attn_num_head_channels,
+                    in_channels=out_channels,
+                    num_layers=1,
+                    cross_attention_dim=cross_attention_dim,
+                    norm_num_groups=resnet_groups,
+                )
+            )
+        self.resnets = nn.ModuleList(resnets)
+        self.temp_convs = nn.ModuleList(temp_convs)
+        self.attentions = nn.ModuleList(attentions)
+        self.temp_attentions = nn.ModuleList(temp_attentions)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
+        else:
+            self.upsamplers = None
+
+    def forward(
+        self,
+        hidden_states,
+        res_hidden_states_tuple,
+        temb=None,
+        encoder_hidden_states=None,
+        upsample_size=None,
+        attention_mask=None,
+        num_frames=1,
+        cross_attention_kwargs=None,
+    ):
+        # TODO(Patrick, William) - attention mask is not used
+        for resnet, temp_conv, attn, temp_attn in zip(
+            self.resnets, self.temp_convs, self.attentions, self.temp_attentions
+        ):
+            # pop res hidden states
+            res_hidden_states = res_hidden_states_tuple[-1]
+            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+            if self.gradient_checkpointing:
+                hidden_states = cross_attn_g_c(
+                        attn, 
+                        temp_attn, 
+                        resnet, 
+                        temp_conv, 
+                        hidden_states, 
+                        encoder_hidden_states, 
+                        cross_attention_kwargs, 
+                        temb, 
+                        num_frames,
+                        inverse_temp=True
+                    )
+            else:
+                hidden_states = resnet(hidden_states, temb)
+
+                if num_frames > 1:
+                    hidden_states = temp_conv(hidden_states, num_frames=num_frames)
+
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                ).sample
+
+                if num_frames > 1:
+                    hidden_states = temp_attn(hidden_states, num_frames=num_frames).sample
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states, upsample_size)
+
+        return hidden_states
+
+
+class UpBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        prev_output_channel: int,
+        out_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        output_scale_factor=1.0,
+        add_upsample=True,
+    ):
+        super().__init__()
+        resnets = []
+        temp_convs = []
+        self.gradient_checkpointing = False
+        for i in range(num_layers):
+            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
+            resnet_in_channels = prev_output_channel if i == 0 else out_channels
+
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=resnet_in_channels + res_skip_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+            temp_convs.append(
+                TemporalConvLayer(
+                    out_channels,
+                    out_channels,
+                    dropout=0.1
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+        self.temp_convs = nn.ModuleList(temp_convs)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
+        else:
+            self.upsamplers = None
+
+    def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, num_frames=1):
+        for resnet, temp_conv in zip(self.resnets, self.temp_convs):
+            # pop res hidden states
+            res_hidden_states = res_hidden_states_tuple[-1]
+            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+            if self.gradient_checkpointing:
+                hidden_states = up_down_g_c(resnet, temp_conv, hidden_states, temb, num_frames)
+            else:
+                hidden_states = resnet(hidden_states, temb)
+
+                if num_frames > 1:
+                    hidden_states = temp_conv(hidden_states, num_frames=num_frames)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states, upsample_size)
+
+        return hidden_states

models/unet_3d_condition.py

@@ -0,0 +1,500 @@
+# Copyright 2023 Alibaba DAMO-VILAB and The HuggingFace Team. All rights reserved.
+# Copyright 2023 The ModelScope Team.
+#
+# 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.
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput, logging
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformer_temporal import TransformerTemporalModel
+from .unet_3d_blocks import (
+    CrossAttnDownBlock3D,
+    CrossAttnUpBlock3D,
+    DownBlock3D,
+    UNetMidBlock3DCrossAttn,
+    UpBlock3D,
+    get_down_block,
+    get_up_block,
+    transformer_g_c
+)
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class UNet3DConditionOutput(BaseOutput):
+    """
+    Args:
+        sample (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`):
+            Hidden states conditioned on `encoder_hidden_states` input. Output of last layer of model.
+    """
+
+    sample: torch.FloatTensor
+
+
+class UNet3DConditionModel(ModelMixin, ConfigMixin):
+    r"""
+    UNet3DConditionModel is a conditional 2D UNet model that takes in a noisy sample, conditional state, and a timestep
+    and returns sample shaped output.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library
+    implements for all the models (such as downloading or saving, etc.)
+
+    Parameters:
+        sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+            Height and width of input/output sample.
+        in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample.
+        out_channels (`int`, *optional*, defaults to 4): The number of channels in the output.
+        down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
+            The tuple of downsample blocks to use.
+        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D",)`):
+            The tuple of upsample blocks to use.
+        block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+            The tuple of output channels for each block.
+        layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
+        downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
+        mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.
+        act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
+        norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
+            If `None`, it will skip the normalization and activation layers in post-processing
+        norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
+        cross_attention_dim (`int`, *optional*, defaults to 1280): The dimension of the cross attention features.
+        attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: Optional[int] = None,
+        in_channels: int = 4,
+        out_channels: int = 4,
+        down_block_types: Tuple[str] = (
+            "CrossAttnDownBlock3D",
+            "CrossAttnDownBlock3D",
+            "CrossAttnDownBlock3D",
+            "DownBlock3D",
+        ),
+        up_block_types: Tuple[str] = ("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D"),
+        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+        layers_per_block: int = 2,
+        downsample_padding: int = 1,
+        mid_block_scale_factor: float = 1,
+        act_fn: str = "silu",
+        norm_num_groups: Optional[int] = 32,
+        norm_eps: float = 1e-5,
+        cross_attention_dim: int = 1024,
+        attention_head_dim: Union[int, Tuple[int]] = 64,
+    ):
+        super().__init__()
+
+        self.sample_size = sample_size
+        self.gradient_checkpointing = False
+        # Check inputs
+        if len(down_block_types) != len(up_block_types):
+            raise ValueError(
+                f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
+            )
+
+        if len(block_out_channels) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}."
+            )
+
+        # input
+        conv_in_kernel = 3
+        conv_out_kernel = 3
+        conv_in_padding = (conv_in_kernel - 1) // 2
+        self.conv_in = nn.Conv2d(
+            in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
+        )
+
+        # time
+        time_embed_dim = block_out_channels[0] * 4
+        self.time_proj = Timesteps(block_out_channels[0], True, 0)
+        timestep_input_dim = block_out_channels[0]
+
+        self.time_embedding = TimestepEmbedding(
+            timestep_input_dim,
+            time_embed_dim,
+            act_fn=act_fn,
+        )
+
+        self.transformer_in = TransformerTemporalModel(
+            num_attention_heads=8,
+            attention_head_dim=attention_head_dim,
+            in_channels=block_out_channels[0],
+            num_layers=1,
+        )
+
+        # class embedding
+        self.down_blocks = nn.ModuleList([])
+        self.up_blocks = nn.ModuleList([])
+
+        if isinstance(attention_head_dim, int):
+            attention_head_dim = (attention_head_dim,) * len(down_block_types)
+
+        # down
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            down_block = get_down_block(
+                down_block_type,
+                num_layers=layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=time_embed_dim,
+                add_downsample=not is_final_block,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=cross_attention_dim,
+                attn_num_head_channels=attention_head_dim[i],
+                downsample_padding=downsample_padding,
+                dual_cross_attention=False,
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        self.mid_block = UNetMidBlock3DCrossAttn(
+            in_channels=block_out_channels[-1],
+            temb_channels=time_embed_dim,
+            resnet_eps=norm_eps,
+            resnet_act_fn=act_fn,
+            output_scale_factor=mid_block_scale_factor,
+            cross_attention_dim=cross_attention_dim,
+            attn_num_head_channels=attention_head_dim[-1],
+            resnet_groups=norm_num_groups,
+            dual_cross_attention=False,
+        )
+
+        # count how many layers upsample the images
+        self.num_upsamplers = 0
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        reversed_attention_head_dim = list(reversed(attention_head_dim))
+
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            is_final_block = i == len(block_out_channels) - 1
+
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
+
+            # add upsample block for all BUT final layer
+            if not is_final_block:
+                add_upsample = True
+                self.num_upsamplers += 1
+            else:
+                add_upsample = False
+
+            up_block = get_up_block(
+                up_block_type,
+                num_layers=layers_per_block + 1,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                prev_output_channel=prev_output_channel,
+                temb_channels=time_embed_dim,
+                add_upsample=add_upsample,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=cross_attention_dim,
+                attn_num_head_channels=reversed_attention_head_dim[i],
+                dual_cross_attention=False,
+            )
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        if norm_num_groups is not None:
+            self.conv_norm_out = nn.GroupNorm(
+                num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps
+            )
+            self.conv_act = nn.SiLU()
+        else:
+            self.conv_norm_out = None
+            self.conv_act = None
+
+        conv_out_padding = (conv_out_kernel - 1) // 2
+        self.conv_out = nn.Conv2d(
+            block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding
+        )
+
+    def set_attention_slice(self, slice_size):
+        r"""
+        Enable sliced attention computation.
+
+        When this option is enabled, the attention module will split the input tensor in slices, to compute attention
+        in several steps. This is useful to save some memory in exchange for a small speed decrease.
+
+        Args:
+            slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
+                When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If
+                `"max"`, maxium amount of memory will be saved by running only one slice at a time. If a number is
+                provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
+                must be a multiple of `slice_size`.
+        """
+        sliceable_head_dims = []
+
+        def fn_recursive_retrieve_slicable_dims(module: torch.nn.Module):
+            if hasattr(module, "set_attention_slice"):
+                sliceable_head_dims.append(module.sliceable_head_dim)
+
+            for child in module.children():
+                fn_recursive_retrieve_slicable_dims(child)
+
+        # retrieve number of attention layers
+        for module in self.children():
+            fn_recursive_retrieve_slicable_dims(module)
+
+        num_slicable_layers = len(sliceable_head_dims)
+
+        if slice_size == "auto":
+            # half the attention head size is usually a good trade-off between
+            # speed and memory
+            slice_size = [dim // 2 for dim in sliceable_head_dims]
+        elif slice_size == "max":
+            # make smallest slice possible
+            slice_size = num_slicable_layers * [1]
+
+        slice_size = num_slicable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
+
+        if len(slice_size) != len(sliceable_head_dims):
+            raise ValueError(
+                f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
+                f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
+            )
+
+        for i in range(len(slice_size)):
+            size = slice_size[i]
+            dim = sliceable_head_dims[i]
+            if size is not None and size > dim:
+                raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
+
+        # Recursively walk through all the children.
+        # Any children which exposes the set_attention_slice method
+        # gets the message
+        def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+            if hasattr(module, "set_attention_slice"):
+                module.set_attention_slice(slice_size.pop())
+
+            for child in module.children():
+                fn_recursive_set_attention_slice(child, slice_size)
+
+        reversed_slice_size = list(reversed(slice_size))
+        for module in self.children():
+            fn_recursive_set_attention_slice(module, reversed_slice_size)
+
+    def _set_gradient_checkpointing(self, value=False):
+        self.gradient_checkpointing = value
+        self.mid_block.gradient_checkpointing = value
+        for module in self.down_blocks + self.up_blocks:
+            if isinstance(module, (CrossAttnDownBlock3D, DownBlock3D, CrossAttnUpBlock3D, UpBlock3D)):
+                module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        class_labels: Optional[torch.Tensor] = None,
+        timestep_cond: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        mid_block_additional_residual: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[UNet3DConditionOutput, Tuple]:
+        r"""
+        Args:
+            sample (`torch.FloatTensor`): (batch, num_frames, channel, height, width) noisy inputs tensor
+            timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
+            encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`models.unet_2d_condition.UNet3DConditionOutput`] instead of a plain tuple.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
+
+        Returns:
+            [`~models.unet_2d_condition.UNet3DConditionOutput`] or `tuple`:
+            [`~models.unet_2d_condition.UNet3DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When
+            returning a tuple, the first element is the sample tensor.
+        """
+        # By default samples have to be AT least a multiple of the overall upsampling factor.
+        # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
+        # However, the upsampling interpolation output size can be forced to fit any upsampling size
+        # on the fly if necessary.
+        default_overall_up_factor = 2**self.num_upsamplers
+
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
+
+        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
+            logger.info("Forward upsample size to force interpolation output size.")
+            forward_upsample_size = True
+
+        # prepare attention_mask
+        if attention_mask is not None:
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # 1. time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        num_frames = sample.shape[2]
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+
+        # timesteps does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=self.dtype)
+
+        emb = self.time_embedding(t_emb, timestep_cond)
+        emb = emb.repeat_interleave(repeats=num_frames, dim=0)
+        encoder_hidden_states = encoder_hidden_states.repeat_interleave(repeats=num_frames, dim=0)
+
+        # 2. pre-process
+        sample = sample.permute(0, 2, 1, 3, 4).reshape((sample.shape[0] * num_frames, -1) + sample.shape[3:])
+        sample = self.conv_in(sample)
+        
+        if num_frames > 1:
+            if self.gradient_checkpointing:
+                sample = transformer_g_c(self.transformer_in, sample, num_frames)
+            else:
+                sample = self.transformer_in(sample, num_frames=num_frames).sample
+
+        # 3. down
+        down_block_res_samples = (sample,)
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    num_frames=num_frames,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb, num_frames=num_frames)
+
+            down_block_res_samples += res_samples
+
+        if down_block_additional_residuals is not None:
+            new_down_block_res_samples = ()
+
+            for down_block_res_sample, down_block_additional_residual in zip(
+                down_block_res_samples, down_block_additional_residuals
+            ):
+                down_block_res_sample = down_block_res_sample + down_block_additional_residual
+                new_down_block_res_samples += (down_block_res_sample,)
+
+            down_block_res_samples = new_down_block_res_samples
+
+        # 4. mid
+        if self.mid_block is not None:
+            sample = self.mid_block(
+                sample,
+                emb,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=attention_mask,
+                num_frames=num_frames,
+                cross_attention_kwargs=cross_attention_kwargs,
+            )
+
+        if mid_block_additional_residual is not None:
+            sample = sample + mid_block_additional_residual
+
+        # 5. up
+        for i, upsample_block in enumerate(self.up_blocks):
+            is_final_block = i == len(self.up_blocks) - 1
+
+            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+            # if we have not reached the final block and need to forward the
+            # upsample size, we do it here
+            if not is_final_block and forward_upsample_size:
+                upsample_size = down_block_res_samples[-1].shape[2:]
+
+            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states,
+                    upsample_size=upsample_size,
+                    attention_mask=attention_mask,
+                    num_frames=num_frames,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                )
+            else:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    upsample_size=upsample_size,
+                    num_frames=num_frames,
+                )
+
+        # 6. post-process
+        if self.conv_norm_out:
+            sample = self.conv_norm_out(sample)
+            sample = self.conv_act(sample)
+
+        sample = self.conv_out(sample)
+
+        # reshape to (batch, channel, framerate, width, height)
+        sample = sample[None, :].reshape((-1, num_frames) + sample.shape[1:]).permute(0, 2, 1, 3, 4)
+
+        if not return_dict:
+            return (sample,)
+
+        return UNet3DConditionOutput(sample=sample)

requirements.txt

@@ -0,0 +1,19 @@
+accelerate==0.18.0
+decord==0.6.0
+deepspeed==0.10.0
+diffusers==0.18.0
+huggingface-hub==0.16.4
+lora-diffusion @ git+https://github.com/cloneofsimo/lora.git@bdd51b04c49fa90a88919a19850ec3b4cf3c5ecd
+loralib==0.1.0
+numpy==1.24.0
+omegaconf==2.3.0
+opencv-python==4.8.0.74
+torch==2.0.1
+torchaudio==2.0.2
+torchvision==0.15.2
+tqdm==4.65.0
+transformers==4.27.4
+einops==0.7.0
+imageio==2.33.0
+imageio-ffmpeg==0.4.9
+gradio==3.26.0

utils/bucketing.py

@@ -0,0 +1,32 @@
+from PIL import Image
+
+def min_res(size, min_size): return 192 if size < 192 else size
+
+def up_down_bucket(m_size, in_size, direction):
+    if direction == 'down': return abs(int(m_size - in_size))
+    if direction == 'up': return abs(int(m_size + in_size))
+
+def get_bucket_sizes(size, direction: 'down', min_size):
+    multipliers = [64, 128]
+    for i, m in enumerate(multipliers):
+        res =  up_down_bucket(m, size, direction)
+        multipliers[i] = min_res(res, min_size=min_size)
+    return multipliers
+
+def closest_bucket(m_size, size, direction, min_size):
+    lst = get_bucket_sizes(m_size, direction, min_size)
+    return lst[min(range(len(lst)), key=lambda i: abs(lst[i]-size))]
+
+def resolve_bucket(i,h,w): return  (i / (h / w))
+
+def sensible_buckets(m_width, m_height, w, h, min_size=192):
+    if h > w:
+        w = resolve_bucket(m_width, h, w)
+        w = closest_bucket(m_width, w, 'down', min_size=min_size)
+        return w, m_height
+    if h < w:
+        h = resolve_bucket(m_height, w, h)
+        h = closest_bucket(m_height, h, 'down', min_size=min_size)
+        return m_width, h
+
+    return m_width, m_height

utils/convert_diffusers_to_original_ms_text_to_video.py

@@ -0,0 +1,465 @@
+# Script for converting a HF Diffusers saved pipeline to a Stable Diffusion checkpoint.
+# *Only* converts the UNet, and Text Encoder.
+# Does not convert optimizer state or any other thing.
+
+import argparse
+import os.path as osp
+import re
+
+import torch
+from safetensors.torch import load_file, save_file
+
+# =================#
+# UNet Conversion #
+# =================#
+
+print ('Initializing the conversion map')
+
+unet_conversion_map = [
+    # (ModelScope, HF Diffusers)
+
+    # from Vanilla ModelScope/StableDiffusion
+    ("time_embed.0.weight", "time_embedding.linear_1.weight"),
+    ("time_embed.0.bias", "time_embedding.linear_1.bias"),
+    ("time_embed.2.weight", "time_embedding.linear_2.weight"),
+    ("time_embed.2.bias", "time_embedding.linear_2.bias"),
+
+
+    # from Vanilla ModelScope/StableDiffusion
+    ("input_blocks.0.0.weight", "conv_in.weight"),
+    ("input_blocks.0.0.bias", "conv_in.bias"),
+
+
+    # from Vanilla ModelScope/StableDiffusion
+    ("out.0.weight", "conv_norm_out.weight"),
+    ("out.0.bias", "conv_norm_out.bias"),
+    ("out.2.weight", "conv_out.weight"),
+    ("out.2.bias", "conv_out.bias"),
+]
+
+unet_conversion_map_resnet = [
+    # (ModelScope, HF Diffusers)
+
+    # SD
+    ("in_layers.0", "norm1"),
+    ("in_layers.2", "conv1"),
+    ("out_layers.0", "norm2"),
+    ("out_layers.3", "conv2"),
+    ("emb_layers.1", "time_emb_proj"),
+    ("skip_connection", "conv_shortcut"),
+
+    # MS
+    #("temopral_conv", "temp_convs"), # ROFL, they have a typo here --kabachuha
+]
+
+unet_conversion_map_layer = []
+
+# Convert input TemporalTransformer
+unet_conversion_map_layer.append(('input_blocks.0.1', 'transformer_in'))
+
+# Reference for the default settings
+
+# "model_cfg": {
+#     "unet_in_dim": 4,
+#     "unet_dim": 320,
+#     "unet_y_dim": 768,
+#     "unet_context_dim": 1024,
+#     "unet_out_dim": 4,
+#     "unet_dim_mult": [1, 2, 4, 4],
+#     "unet_num_heads": 8,
+#     "unet_head_dim": 64,
+#     "unet_res_blocks": 2,
+#     "unet_attn_scales": [1, 0.5, 0.25],
+#     "unet_dropout": 0.1,
+#     "temporal_attention": "True",
+#     "num_timesteps": 1000,
+#     "mean_type": "eps",
+#     "var_type": "fixed_small",
+#     "loss_type": "mse"
+# }
+
+# hardcoded number of downblocks and resnets/attentions...
+# would need smarter logic for other networks.
+for i in range(4):
+    # loop over downblocks/upblocks
+
+    for j in range(2):
+        # loop over resnets/attentions for downblocks
+
+        # Spacial SD stuff
+        hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
+        sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
+        unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
+
+        if i < 3:
+            # no attention layers in down_blocks.3
+            hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
+            sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
+            unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+        
+        # Temporal MS stuff
+        hf_down_res_prefix = f"down_blocks.{i}.temp_convs.{j}."
+        sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0.temopral_conv."
+        unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
+
+        if i < 3:
+            # no attention layers in down_blocks.3
+            hf_down_atn_prefix = f"down_blocks.{i}.temp_attentions.{j}."
+            sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.2."
+            unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+
+    for j in range(3):
+        # loop over resnets/attentions for upblocks
+
+        # Spacial SD stuff
+        hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
+        sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
+        unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
+
+        if i > 0:
+            # no attention layers in up_blocks.0
+            hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
+            sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
+            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
+        
+        # loop over resnets/attentions for upblocks
+        hf_up_res_prefix = f"up_blocks.{i}.temp_convs.{j}."
+        sd_up_res_prefix = f"output_blocks.{3*i + j}.0.temopral_conv."
+        unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
+
+        if i > 0:
+            # no attention layers in up_blocks.0
+            hf_up_atn_prefix = f"up_blocks.{i}.temp_attentions.{j}."
+            sd_up_atn_prefix = f"output_blocks.{3*i + j}.2."
+            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
+
+    # Up/Downsamplers are 2D, so don't need to touch them
+    if i < 3:
+        # no downsample in down_blocks.3
+        hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
+        sd_downsample_prefix = f"input_blocks.{3*(i+1)}.op."
+        unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+
+        # no upsample in up_blocks.3
+        hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+        sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 3}."
+        unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
+
+
+# Handle the middle block
+
+# Spacial
+hf_mid_atn_prefix = "mid_block.attentions.0."
+sd_mid_atn_prefix = "middle_block.1."
+unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
+
+for j in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{j}."
+    sd_mid_res_prefix = f"middle_block.{3*j}."
+    unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+# Temporal
+hf_mid_atn_prefix = "mid_block.temp_attentions.0."
+sd_mid_atn_prefix = "middle_block.2."
+unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
+
+for j in range(2):
+    hf_mid_res_prefix = f"mid_block.temp_convs.{j}."
+    sd_mid_res_prefix = f"middle_block.{3*j}.temopral_conv."
+    unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+# The pipeline
+def convert_unet_state_dict(unet_state_dict, strict_mapping=False):
+    print ('Converting the UNET')
+    # buyer beware: this is a *brittle* function,
+    # and correct output requires that all of these pieces interact in
+    # the exact order in which I have arranged them.
+    mapping = {k: k for k in unet_state_dict.keys()}
+
+    for sd_name, hf_name in unet_conversion_map:
+        if strict_mapping:
+            if hf_name in mapping:
+                mapping[hf_name] = sd_name
+        else:
+            mapping[hf_name] = sd_name
+    for k, v in mapping.items():
+        if "resnets" in k:
+            for sd_part, hf_part in unet_conversion_map_resnet:
+                v = v.replace(hf_part, sd_part)
+            mapping[k] = v
+        # elif "temp_convs" in k:
+        #     for sd_part, hf_part in unet_conversion_map_resnet:
+        #         v = v.replace(hf_part, sd_part)
+        #     mapping[k] = v
+    for k, v in mapping.items():
+        for sd_part, hf_part in unet_conversion_map_layer:
+            v = v.replace(hf_part, sd_part)
+        mapping[k] = v
+    
+
+    # there must be a pattern, but I don't want to bother atm
+    do_not_unsqueeze = [f'output_blocks.{i}.1.proj_out.weight' for i in range(3, 12)] + [f'output_blocks.{i}.1.proj_in.weight' for i in range(3, 12)] + ['middle_block.1.proj_in.weight', 'middle_block.1.proj_out.weight'] + [f'input_blocks.{i}.1.proj_out.weight' for i in [1, 2, 4, 5, 7, 8]] + [f'input_blocks.{i}.1.proj_in.weight' for i in [1, 2, 4, 5, 7, 8]]
+    print (do_not_unsqueeze)
+
+    new_state_dict = {v: (unet_state_dict[k].unsqueeze(-1) if ('proj_' in k and ('bias' not in k) and (k not in do_not_unsqueeze)) else unet_state_dict[k]) for k, v in mapping.items()}
+    # HACK: idk why the hell it does not work with list comprehension
+    for k, v in new_state_dict.items():
+        has_k = False
+        for n in do_not_unsqueeze:
+            if k == n:
+                has_k = True
+
+        if has_k:
+            v = v.squeeze(-1)
+        new_state_dict[k] = v
+
+    return new_state_dict
+
+# TODO: VAE conversion. We doesn't train it in the most cases, but may be handy for the future --kabachuha
+
+# =========================#
+# Text Encoder Conversion #
+# =========================#
+
+# IT IS THE SAME CLIP ENCODER, SO JUST COPYPASTING IT --kabachuha
+
+# =========================#
+# Text Encoder Conversion #
+# =========================#
+
+
+textenc_conversion_lst = [
+    # (stable-diffusion, HF Diffusers)
+    ("resblocks.", "text_model.encoder.layers."),
+    ("ln_1", "layer_norm1"),
+    ("ln_2", "layer_norm2"),
+    (".c_fc.", ".fc1."),
+    (".c_proj.", ".fc2."),
+    (".attn", ".self_attn"),
+    ("ln_final.", "transformer.text_model.final_layer_norm."),
+    ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"),
+    ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"),
+]
+protected = {re.escape(x[1]): x[0] for x in textenc_conversion_lst}
+textenc_pattern = re.compile("|".join(protected.keys()))
+
+# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
+code2idx = {"q": 0, "k": 1, "v": 2}
+
+
+def convert_text_enc_state_dict_v20(text_enc_dict):
+    #print ('Converting the text encoder')
+    new_state_dict = {}
+    capture_qkv_weight = {}
+    capture_qkv_bias = {}
+    for k, v in text_enc_dict.items():
+        if (
+            k.endswith(".self_attn.q_proj.weight")
+            or k.endswith(".self_attn.k_proj.weight")
+            or k.endswith(".self_attn.v_proj.weight")
+        ):
+            k_pre = k[: -len(".q_proj.weight")]
+            k_code = k[-len("q_proj.weight")]
+            if k_pre not in capture_qkv_weight:
+                capture_qkv_weight[k_pre] = [None, None, None]
+            capture_qkv_weight[k_pre][code2idx[k_code]] = v
+            continue
+
+        if (
+            k.endswith(".self_attn.q_proj.bias")
+            or k.endswith(".self_attn.k_proj.bias")
+            or k.endswith(".self_attn.v_proj.bias")
+        ):
+            k_pre = k[: -len(".q_proj.bias")]
+            k_code = k[-len("q_proj.bias")]
+            if k_pre not in capture_qkv_bias:
+                capture_qkv_bias[k_pre] = [None, None, None]
+            capture_qkv_bias[k_pre][code2idx[k_code]] = v
+            continue
+
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k)
+        new_state_dict[relabelled_key] = v
+
+    for k_pre, tensors in capture_qkv_weight.items():
+        if None in tensors:
+            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
+        new_state_dict[relabelled_key + ".in_proj_weight"] = torch.cat(tensors)
+
+    for k_pre, tensors in capture_qkv_bias.items():
+        if None in tensors:
+            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
+        new_state_dict[relabelled_key + ".in_proj_bias"] = torch.cat(tensors)
+
+    return new_state_dict
+
+
+def convert_text_enc_state_dict(text_enc_dict):
+    return text_enc_dict
+
+textenc_conversion_lst = [
+    # (stable-diffusion, HF Diffusers)
+    ("resblocks.", "text_model.encoder.layers."),
+    ("ln_1", "layer_norm1"),
+    ("ln_2", "layer_norm2"),
+    (".c_fc.", ".fc1."),
+    (".c_proj.", ".fc2."),
+    (".attn", ".self_attn"),
+    ("ln_final.", "transformer.text_model.final_layer_norm."),
+    ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"),
+    ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"),
+]
+protected = {re.escape(x[1]): x[0] for x in textenc_conversion_lst}
+textenc_pattern = re.compile("|".join(protected.keys()))
+
+# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
+code2idx = {"q": 0, "k": 1, "v": 2}
+
+
+def convert_text_enc_state_dict_v20(text_enc_dict):
+    new_state_dict = {}
+    capture_qkv_weight = {}
+    capture_qkv_bias = {}
+    for k, v in text_enc_dict.items():
+        if (
+            k.endswith(".self_attn.q_proj.weight")
+            or k.endswith(".self_attn.k_proj.weight")
+            or k.endswith(".self_attn.v_proj.weight")
+        ):
+            k_pre = k[: -len(".q_proj.weight")]
+            k_code = k[-len("q_proj.weight")]
+            if k_pre not in capture_qkv_weight:
+                capture_qkv_weight[k_pre] = [None, None, None]
+            capture_qkv_weight[k_pre][code2idx[k_code]] = v
+            continue
+
+        if (
+            k.endswith(".self_attn.q_proj.bias")
+            or k.endswith(".self_attn.k_proj.bias")
+            or k.endswith(".self_attn.v_proj.bias")
+        ):
+            k_pre = k[: -len(".q_proj.bias")]
+            k_code = k[-len("q_proj.bias")]
+            if k_pre not in capture_qkv_bias:
+                capture_qkv_bias[k_pre] = [None, None, None]
+            capture_qkv_bias[k_pre][code2idx[k_code]] = v
+            continue
+
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k)
+        new_state_dict[relabelled_key] = v
+
+    for k_pre, tensors in capture_qkv_weight.items():
+        if None in tensors:
+            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
+        new_state_dict[relabelled_key + ".in_proj_weight"] = torch.cat(tensors)
+
+    for k_pre, tensors in capture_qkv_bias.items():
+        if None in tensors:
+            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
+        new_state_dict[relabelled_key + ".in_proj_bias"] = torch.cat(tensors)
+
+    return new_state_dict
+
+
+def convert_text_enc_state_dict(text_enc_dict):
+    return text_enc_dict
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--model_path", default=None, type=str, required=True, help="Path to the model to convert.")
+    parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.")
+    parser.add_argument("--clip_checkpoint_path", default=None, type=str, help="Path to the output CLIP model.")
+    parser.add_argument("--half", action="store_true", help="Save weights in half precision.")
+    parser.add_argument(
+        "--use_safetensors", action="store_true", help="Save weights use safetensors, default is ckpt."
+    )
+
+    args = parser.parse_args()
+
+    assert args.model_path is not None, "Must provide a model path!"
+
+    assert args.checkpoint_path is not None, "Must provide a checkpoint path!"
+
+    assert args.clip_checkpoint_path is not None, "Must provide a CLIP checkpoint path!"
+
+    # Path for safetensors
+    unet_path = osp.join(args.model_path, "unet", "diffusion_pytorch_model.safetensors")
+    #vae_path = osp.join(args.model_path, "vae", "diffusion_pytorch_model.safetensors")
+    text_enc_path = osp.join(args.model_path, "text_encoder", "model.safetensors")
+
+    # Load models from safetensors if it exists, if it doesn't pytorch
+    if osp.exists(unet_path):
+        unet_state_dict = load_file(unet_path, device="cpu")
+    else:
+        unet_path = osp.join(args.model_path, "unet", "diffusion_pytorch_model.bin")
+        unet_state_dict = torch.load(unet_path, map_location="cpu")
+
+    # if osp.exists(vae_path):
+    #     vae_state_dict = load_file(vae_path, device="cpu")
+    # else:
+    #     vae_path = osp.join(args.model_path, "vae", "diffusion_pytorch_model.bin")
+    #     vae_state_dict = torch.load(vae_path, map_location="cpu")
+
+    if osp.exists(text_enc_path):
+        text_enc_dict = load_file(text_enc_path, device="cpu")
+    else:
+        text_enc_path = osp.join(args.model_path, "text_encoder", "pytorch_model.bin")
+        text_enc_dict = torch.load(text_enc_path, map_location="cpu")
+
+    # Convert the UNet model
+    unet_state_dict = convert_unet_state_dict(unet_state_dict)
+    #unet_state_dict = {"model.diffusion_model." + k: v for k, v in unet_state_dict.items()}
+
+    # Convert the VAE model
+    # vae_state_dict = convert_vae_state_dict(vae_state_dict)
+    # vae_state_dict = {"first_stage_model." + k: v for k, v in vae_state_dict.items()}
+
+    # Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
+    is_v20_model = "text_model.encoder.layers.22.layer_norm2.bias" in text_enc_dict
+
+    if is_v20_model:
+
+        # MODELSCOPE always uses the 2.X encoder, btw --kabachuha
+
+        # Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
+        text_enc_dict = {"transformer." + k: v for k, v in text_enc_dict.items()}
+        text_enc_dict = convert_text_enc_state_dict_v20(text_enc_dict)
+        #text_enc_dict = {"cond_stage_model.model." + k: v for k, v in text_enc_dict.items()}
+    else:
+        text_enc_dict = convert_text_enc_state_dict(text_enc_dict)
+        #text_enc_dict = {"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}
+
+    # DON'T PUT TOGETHER FOR THE NEW CHECKPOINT AS MODELSCOPE USES THEM IN THE SPLITTED FORM --kabachuha
+    # Save CLIP and the Diffusion model to their own files
+
+    #state_dict = {**unet_state_dict, **vae_state_dict, **text_enc_dict}
+    print ('Saving UNET')
+    state_dict = {**unet_state_dict}
+
+    if args.half:
+        state_dict = {k: v.half() for k, v in state_dict.items()}
+
+    if args.use_safetensors:
+        save_file(state_dict, args.checkpoint_path)
+    else:
+        #state_dict = {"state_dict": state_dict}
+        torch.save(state_dict, args.checkpoint_path)
+
+    # TODO: CLIP conversion doesn't work atm
+    # print ('Saving CLIP')
+    # state_dict = {**text_enc_dict}
+
+    # if args.half:
+    #     state_dict = {k: v.half() for k, v in state_dict.items()}
+
+    # if args.use_safetensors:
+    #     save_file(state_dict, args.checkpoint_path)
+    # else:
+    #     #state_dict = {"state_dict": state_dict}
+    #     torch.save(state_dict, args.clip_checkpoint_path)
+    
+    print('Operation successfull')

utils/dataset.py

@@ -0,0 +1,578 @@
+import os
+import decord
+import numpy as np
+import random
+import json
+import torchvision
+import torchvision.transforms as T
+import torch
+
+from glob import glob
+from PIL import Image
+from itertools import islice
+from pathlib import Path
+from .bucketing import sensible_buckets
+
+decord.bridge.set_bridge('torch')
+
+from torch.utils.data import Dataset
+from einops import rearrange, repeat
+
+
+def get_prompt_ids(prompt, tokenizer):
+    prompt_ids = tokenizer(
+            prompt,
+            truncation=True,
+            padding="max_length",
+            max_length=tokenizer.model_max_length,
+            return_tensors="pt",
+    ).input_ids
+
+    return prompt_ids
+
+
+def read_caption_file(caption_file):
+        with open(caption_file, 'r', encoding="utf8") as t:
+            return t.read()
+
+
+def get_text_prompt(
+        text_prompt: str = '', 
+        fallback_prompt: str= '',
+        file_path:str = '', 
+        ext_types=['.mp4'],
+        use_caption=False
+    ):
+    try:
+        if use_caption:
+            if len(text_prompt) > 1: return text_prompt
+            caption_file = ''
+            # Use caption on per-video basis (One caption PER video)
+            for ext in ext_types:
+                maybe_file = file_path.replace(ext, '.txt')
+                if maybe_file.endswith(ext_types): continue
+                if os.path.exists(maybe_file): 
+                    caption_file = maybe_file
+                    break
+
+            if os.path.exists(caption_file):
+                return read_caption_file(caption_file)
+            
+            # Return fallback prompt if no conditions are met.
+            return fallback_prompt
+
+        return text_prompt
+    except:
+        print(f"Couldn't read prompt caption for {file_path}. Using fallback.")
+        return fallback_prompt
+
+    
+def get_video_frames(vr, start_idx, sample_rate=1, max_frames=24):
+    max_range = len(vr)
+    frame_number = sorted((0, start_idx, max_range))[1]
+
+    frame_range = range(frame_number, max_range, sample_rate)
+    frame_range_indices = list(frame_range)[:max_frames]
+
+    return frame_range_indices
+
+
+def process_video(vid_path, use_bucketing, w, h, get_frame_buckets, get_frame_batch):
+    if use_bucketing:
+        vr = decord.VideoReader(vid_path)
+        resize = get_frame_buckets(vr)
+        video = get_frame_batch(vr, resize=resize)
+
+    else:
+        vr = decord.VideoReader(vid_path, width=w, height=h)
+        video = get_frame_batch(vr)
+
+    return video, vr
+
+
+# https://github.com/ExponentialML/Video-BLIP2-Preprocessor
+class VideoJsonDataset(Dataset):
+    def __init__(
+            self,
+            tokenizer = None,
+            width: int = 256,
+            height: int = 256,
+            n_sample_frames: int = 4,
+            sample_start_idx: int = 1,
+            frame_step: int = 1,
+            json_path: str ="",
+            json_data = None,
+            vid_data_key: str = "video_path",
+            preprocessed: bool = False,
+            use_bucketing: bool = False,
+            **kwargs
+    ):
+        self.vid_types = (".mp4", ".avi", ".mov", ".webm", ".flv", ".mjpeg")
+        self.use_bucketing = use_bucketing
+        self.tokenizer = tokenizer
+        self.preprocessed = preprocessed
+        
+        self.vid_data_key = vid_data_key
+        self.train_data = self.load_from_json(json_path, json_data)
+
+        self.width = width
+        self.height = height
+
+        self.n_sample_frames = n_sample_frames
+        self.sample_start_idx = sample_start_idx
+        self.frame_step = frame_step
+
+    def build_json(self, json_data):
+        extended_data = []
+        for data in json_data['data']:
+            for nested_data in data['data']:
+                self.build_json_dict(
+                    data, 
+                    nested_data, 
+                    extended_data
+                )
+        json_data = extended_data
+        return json_data
+
+    def build_json_dict(self, data, nested_data, extended_data):
+        clip_path = nested_data['clip_path'] if 'clip_path' in nested_data else None
+        
+        extended_data.append({
+            self.vid_data_key: data[self.vid_data_key],
+            'frame_index': nested_data['frame_index'],
+            'prompt': nested_data['prompt'],
+            'clip_path': clip_path
+        })
+        
+    def load_from_json(self, path, json_data):
+        try:
+            with open(path) as jpath:
+                print(f"Loading JSON from {path}")
+                json_data = json.load(jpath)
+
+                return self.build_json(json_data)
+
+        except:
+            self.train_data = []
+            print("Non-existant JSON path. Skipping.")
+            
+    def validate_json(self, base_path, path):
+        return os.path.exists(f"{base_path}/{path}")
+
+    def get_frame_range(self, vr):
+        return get_video_frames(
+            vr, 
+            self.sample_start_idx, 
+            self.frame_step, 
+            self.n_sample_frames
+        )
+    
+    def get_vid_idx(self, vr, vid_data=None):
+        frames = self.n_sample_frames
+
+        if vid_data is not None:
+            idx = vid_data['frame_index']
+        else:
+            idx = self.sample_start_idx
+
+        return idx
+
+    def get_frame_buckets(self, vr):
+        _, h, w = vr[0].shape        
+        width, height = sensible_buckets(self.width, self.height, h, w)
+        # width, height = self.width, self.height
+        resize = T.transforms.Resize((height, width), antialias=True)
+
+        return resize
+
+    def get_frame_batch(self, vr, resize=None):
+        frame_range = self.get_frame_range(vr)
+        frames = vr.get_batch(frame_range)
+        video = rearrange(frames, "f h w c -> f c h w")
+
+        if resize is not None: video = resize(video)
+        return video
+
+    def process_video_wrapper(self, vid_path):
+        video, vr = process_video(
+                vid_path,
+                self.use_bucketing,
+                self.width, 
+                self.height, 
+                self.get_frame_buckets, 
+                self.get_frame_batch
+            )
+        
+        return video, vr 
+
+    def train_data_batch(self, index):
+
+        # If we are training on individual clips.
+        if 'clip_path' in self.train_data[index] and \
+            self.train_data[index]['clip_path'] is not None:
+
+            vid_data = self.train_data[index]
+
+            clip_path = vid_data['clip_path']
+            
+            # Get video prompt
+            prompt = vid_data['prompt']
+
+            video, _ = self.process_video_wrapper(clip_path)
+
+            prompt_ids = get_prompt_ids(prompt, self.tokenizer)
+
+            return video, prompt, prompt_ids
+
+         # Assign train data
+        train_data = self.train_data[index]
+        
+        # Get the frame of the current index.
+        self.sample_start_idx = train_data['frame_index']
+        
+        # Initialize resize
+        resize = None
+
+        video, vr = self.process_video_wrapper(train_data[self.vid_data_key])
+
+        # Get video prompt
+        prompt = train_data['prompt']
+        vr.seek(0)
+
+        prompt_ids = get_prompt_ids(prompt, self.tokenizer)
+
+        return video, prompt, prompt_ids
+
+    @staticmethod
+    def __getname__(): return 'json'
+
+    def __len__(self):
+        if self.train_data is not None:
+            return len(self.train_data)
+        else: 
+            return 0
+
+    def __getitem__(self, index):
+        
+        # Initialize variables
+        video = None
+        prompt = None
+        prompt_ids = None
+
+        # Use default JSON training
+        if self.train_data is not None:
+            video, prompt, prompt_ids = self.train_data_batch(index)
+
+        example = {
+            "pixel_values": (video / 127.5 - 1.0),
+            "prompt_ids": prompt_ids[0],
+            "text_prompt": prompt,
+            'dataset': self.__getname__()
+        }
+
+        return example
+
+
+class SingleVideoDataset(Dataset):
+    def __init__(
+        self,
+            tokenizer = None,
+            width: int = 256,
+            height: int = 256,
+            n_sample_frames: int = 4,
+            frame_step: int = 1,
+            single_video_path: str = "",
+            single_video_prompt: str = "",
+            use_caption: bool = False,
+            use_bucketing: bool = False,
+            **kwargs
+    ):
+        self.tokenizer = tokenizer
+        self.use_bucketing = use_bucketing
+        self.frames = []
+        self.index = 1
+
+        self.vid_types = (".mp4", ".avi", ".mov", ".webm", ".flv", ".mjpeg")
+        self.n_sample_frames = n_sample_frames
+        self.frame_step = frame_step
+
+        self.single_video_path = single_video_path
+        self.single_video_prompt = single_video_prompt
+
+        self.width = width
+        self.height = height
+    def create_video_chunks(self):
+        vr = decord.VideoReader(self.single_video_path)
+        vr_range = range(0, len(vr), self.frame_step)
+
+        self.frames = list(self.chunk(vr_range, self.n_sample_frames))
+        return self.frames
+
+    def chunk(self, it, size):
+        it = iter(it)
+        return iter(lambda: tuple(islice(it, size)), ())
+
+    def get_frame_batch(self, vr, resize=None):
+        index = self.index
+        frames = vr.get_batch(self.frames[self.index])
+        video = rearrange(frames, "f h w c -> f c h w")
+
+        if resize is not None: video = resize(video)
+        return video
+
+    def get_frame_buckets(self, vr):
+        _, h, w = vr[0].shape        
+        # width, height = sensible_buckets(self.width, self.height, h, w)
+        width, height = self.width, self.height
+        resize = T.transforms.Resize((height, width), antialias=True)
+
+        return resize
+    
+    def process_video_wrapper(self, vid_path):
+        video, vr = process_video(
+                vid_path,
+                self.use_bucketing,
+                self.width, 
+                self.height, 
+                self.get_frame_buckets, 
+                self.get_frame_batch
+            )
+        
+        return video, vr 
+
+    def single_video_batch(self, index):
+        train_data = self.single_video_path
+        self.index = index
+
+        if train_data.endswith(self.vid_types):
+            video, _ = self.process_video_wrapper(train_data)
+
+            prompt = self.single_video_prompt
+            prompt_ids = get_prompt_ids(prompt, self.tokenizer)
+
+            return video, prompt, prompt_ids
+        else:
+            raise ValueError(f"Single video is not a video type. Types: {self.vid_types}")
+    
+    @staticmethod
+    def __getname__(): return 'single_video'
+
+    def __len__(self):
+        
+        return len(self.create_video_chunks())
+
+    def __getitem__(self, index):
+
+        video, prompt, prompt_ids = self.single_video_batch(index)
+
+        example = {
+            "pixel_values": (video / 127.5 - 1.0),
+            "prompt_ids": prompt_ids[0],
+            "text_prompt": prompt,
+            'dataset': self.__getname__()
+        }
+
+        return example
+
+
+class ImageDataset(Dataset):
+    
+    def __init__(
+        self,
+        tokenizer = None,
+        width: int = 256,
+        height: int = 256,
+        base_width: int = 256,
+        base_height: int = 256,
+        use_caption:     bool = False,
+        image_dir: str = '',
+        single_img_prompt: str = '',
+        use_bucketing: bool = False,
+        fallback_prompt: str = '',
+        **kwargs
+    ):
+        self.tokenizer = tokenizer
+        self.img_types = (".png", ".jpg", ".jpeg", '.bmp')
+        self.use_bucketing = use_bucketing
+
+        self.image_dir = self.get_images_list(image_dir)
+        self.fallback_prompt = fallback_prompt
+
+        self.use_caption = use_caption
+        self.single_img_prompt = single_img_prompt
+
+        self.width = width
+        self.height = height
+
+    def get_images_list(self, image_dir):
+        if os.path.exists(image_dir):
+            imgs = [x for x in os.listdir(image_dir) if x.endswith(self.img_types)]
+            full_img_dir = []
+
+            for img in imgs: 
+                full_img_dir.append(f"{image_dir}/{img}")
+
+            return sorted(full_img_dir)
+
+        return ['']
+
+    def image_batch(self, index):
+        train_data = self.image_dir[index]
+        img = train_data
+
+        try:
+            img = torchvision.io.read_image(img, mode=torchvision.io.ImageReadMode.RGB)
+        except:
+            img = T.transforms.PILToTensor()(Image.open(img).convert("RGB"))
+
+        width = self.width
+        height = self.height
+
+        if self.use_bucketing:
+            _, h, w = img.shape
+            width, height = sensible_buckets(width, height, w, h)
+              
+        resize = T.transforms.Resize((height, width), antialias=True)
+
+        img = resize(img) 
+        img = repeat(img, 'c h w -> f c h w', f=1)
+
+        prompt = get_text_prompt(
+            file_path=train_data,
+            text_prompt=self.single_img_prompt,
+            fallback_prompt=self.fallback_prompt,
+            ext_types=self.img_types,  
+            use_caption=True
+        )
+        prompt_ids = get_prompt_ids(prompt, self.tokenizer)
+
+        return img, prompt, prompt_ids
+
+    @staticmethod
+    def __getname__(): return 'image'
+    
+    def __len__(self):
+        # Image directory
+        if os.path.exists(self.image_dir[0]):
+            return len(self.image_dir)
+        else:
+            return 0
+
+    def __getitem__(self, index):
+        img, prompt, prompt_ids = self.image_batch(index)
+        example = {
+            "pixel_values": (img / 127.5 - 1.0),
+            "prompt_ids": prompt_ids[0],
+            "text_prompt": prompt, 
+            'dataset': self.__getname__()
+        }
+
+        return example
+
+
+class VideoFolderDataset(Dataset):
+    def __init__(
+        self,
+        tokenizer=None,
+        width: int = 256,
+        height: int = 256,
+        n_sample_frames: int = 16,
+        fps: int = 8,
+        path: str = "./data",
+        fallback_prompt: str = "",
+        use_bucketing: bool = False,
+        **kwargs
+    ):
+        self.tokenizer = tokenizer
+        self.use_bucketing = use_bucketing
+
+        self.fallback_prompt = fallback_prompt
+
+        self.video_files = glob(f"{path}/*.mp4")
+
+        self.width = width
+        self.height = height
+
+        self.n_sample_frames = n_sample_frames
+        self.fps = fps
+
+    def get_frame_buckets(self, vr):
+        _, h, w = vr[0].shape        
+        width, height = sensible_buckets(self.width, self.height, h, w)
+        # width, height = self.width, self.height
+        resize = T.transforms.Resize((height, width), antialias=True)
+
+        return resize
+
+    def get_frame_batch(self, vr, resize=None):
+        n_sample_frames = self.n_sample_frames
+        native_fps = vr.get_avg_fps()
+        
+        every_nth_frame = max(1, round(native_fps / self.fps))
+        every_nth_frame = min(len(vr), every_nth_frame)
+        
+        effective_length = len(vr) // every_nth_frame
+        if effective_length < n_sample_frames:
+            n_sample_frames = effective_length
+
+        effective_idx = random.randint(0, (effective_length - n_sample_frames))
+        idxs = every_nth_frame * np.arange(effective_idx, effective_idx + n_sample_frames)
+
+        video = vr.get_batch(idxs)
+        video = rearrange(video, "f h w c -> f c h w")
+
+        if resize is not None: video = resize(video)
+        return video, vr
+        
+    def process_video_wrapper(self, vid_path):
+        video, vr = process_video(
+                vid_path,
+                self.use_bucketing,
+                self.width, 
+                self.height, 
+                self.get_frame_buckets, 
+                self.get_frame_batch
+            )
+        return video, vr
+    
+    def get_prompt_ids(self, prompt):
+        return self.tokenizer(
+            prompt,
+            truncation=True,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            return_tensors="pt",
+        ).input_ids
+
+    @staticmethod
+    def __getname__(): return 'folder'
+
+    def __len__(self):
+        return len(self.video_files)
+
+    def __getitem__(self, index):
+
+        video, _ = self.process_video_wrapper(self.video_files[index])
+
+        prompt = self.fallback_prompt
+
+        prompt_ids = self.get_prompt_ids(prompt)
+
+        return {"pixel_values": (video[0] / 127.5 - 1.0), "prompt_ids": prompt_ids[0], "text_prompt": prompt, 'dataset': self.__getname__()}
+
+
+class CachedDataset(Dataset):
+    def __init__(self,cache_dir: str = ''):
+        self.cache_dir = cache_dir
+        self.cached_data_list = self.get_files_list()
+
+    def get_files_list(self):
+        tensors_list = [f"{self.cache_dir}/{x}" for x in os.listdir(self.cache_dir) if x.endswith('.pt')]
+        return sorted(tensors_list)
+
+    def __len__(self):
+        return len(self.cached_data_list)
+
+    def __getitem__(self, index):
+        cached_latent = torch.load(self.cached_data_list[index], map_location='cuda:0')
+        return cached_latent

utils/ddim_utils.py

@@ -0,0 +1,65 @@
+import numpy as np
+from typing import Union
+
+import torch
+
+from tqdm import tqdm
+
+
+# DDIM Inversion
+@torch.no_grad()
+def init_prompt(prompt, pipeline):
+    uncond_input = pipeline.tokenizer(
+        [""], padding="max_length", max_length=pipeline.tokenizer.model_max_length,
+        return_tensors="pt"
+    )
+    uncond_embeddings = pipeline.text_encoder(uncond_input.input_ids.to(pipeline.device))[0]
+    text_input = pipeline.tokenizer(
+        [prompt],
+        padding="max_length",
+        max_length=pipeline.tokenizer.model_max_length,
+        truncation=True,
+        return_tensors="pt",
+    )
+    text_embeddings = pipeline.text_encoder(text_input.input_ids.to(pipeline.device))[0]
+    context = torch.cat([uncond_embeddings, text_embeddings])
+
+    return context
+
+
+def next_step(model_output: Union[torch.FloatTensor, np.ndarray], timestep: int,
+              sample: Union[torch.FloatTensor, np.ndarray], ddim_scheduler):
+    timestep, next_timestep = min(
+        timestep - ddim_scheduler.config.num_train_timesteps // ddim_scheduler.num_inference_steps, 999), timestep
+    alpha_prod_t = ddim_scheduler.alphas_cumprod[timestep] if timestep >= 0 else ddim_scheduler.final_alpha_cumprod
+    alpha_prod_t_next = ddim_scheduler.alphas_cumprod[next_timestep]
+    beta_prod_t = 1 - alpha_prod_t
+    next_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
+    next_sample_direction = (1 - alpha_prod_t_next) ** 0.5 * model_output
+    next_sample = alpha_prod_t_next ** 0.5 * next_original_sample + next_sample_direction
+    return next_sample
+
+
+def get_noise_pred_single(latents, t, context, unet):
+    noise_pred = unet(latents, t, encoder_hidden_states=context)["sample"]
+    return noise_pred
+
+
+@torch.no_grad()
+def ddim_loop(pipeline, ddim_scheduler, latent, num_inv_steps, prompt):
+    context = init_prompt(prompt, pipeline)
+    uncond_embeddings, cond_embeddings = context.chunk(2)
+    all_latent = [latent]
+    latent = latent.clone().detach()
+    for i in tqdm(range(num_inv_steps)):
+        t = ddim_scheduler.timesteps[len(ddim_scheduler.timesteps) - i - 1]
+        noise_pred = get_noise_pred_single(latent, t, cond_embeddings, pipeline.unet)
+        latent = next_step(noise_pred, t, latent, ddim_scheduler)
+        all_latent.append(latent)
+    return all_latent
+
+
+@torch.no_grad()
+def ddim_inversion(pipeline, ddim_scheduler, video_latent, num_inv_steps, prompt=""):
+    ddim_latents = ddim_loop(pipeline, ddim_scheduler, video_latent, num_inv_steps, prompt)
+    return ddim_latents

utils/lora.py

@@ -0,0 +1,1483 @@
+import json
+import math
+from itertools import groupby
+import os
+from typing import Callable, Dict, List, Optional, Set, Tuple, Type, Union
+
+import numpy as np
+import PIL
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from safetensors.torch import safe_open
+    from safetensors.torch import save_file as safe_save
+
+    safetensors_available = True
+except ImportError:
+    from .safe_open import safe_open
+
+    def safe_save(
+        tensors: Dict[str, torch.Tensor],
+        filename: str,
+        metadata: Optional[Dict[str, str]] = None,
+    ) -> None:
+        raise EnvironmentError(
+            "Saving safetensors requires the safetensors library. Please install with pip or similar."
+        )
+
+    safetensors_available = False
+
+
+class LoraInjectedLinear(nn.Module):
+    def __init__(
+        self, in_features, out_features, bias=False, r=4, dropout_p=0.1, scale=1.0
+    ):
+        super().__init__()
+
+        if r > min(in_features, out_features):
+            #raise ValueError(
+            #    f"LoRA rank {r} must be less or equal than {min(in_features, out_features)}"
+            #)
+            print(f"LoRA rank {r} is too large. setting to: {min(in_features, out_features)}")
+            r = min(in_features, out_features)
+
+        self.r = r
+        self.linear = nn.Linear(in_features, out_features, bias)
+        self.lora_down = nn.Linear(in_features, r, bias=False)
+        self.dropout = nn.Dropout(dropout_p)
+        self.lora_up = nn.Linear(r, out_features, bias=False)
+        self.scale = scale
+        self.selector = nn.Identity()
+
+        nn.init.normal_(self.lora_down.weight, std=1 / r)
+        nn.init.zeros_(self.lora_up.weight)
+
+    def forward(self, input):
+        return (
+            self.linear(input)
+            + self.dropout(self.lora_up(self.selector(self.lora_down(input))))
+            * self.scale
+        )
+
+    def realize_as_lora(self):
+        return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
+
+    def set_selector_from_diag(self, diag: torch.Tensor):
+        # diag is a 1D tensor of size (r,)
+        assert diag.shape == (self.r,)
+        self.selector = nn.Linear(self.r, self.r, bias=False)
+        self.selector.weight.data = torch.diag(diag)
+        self.selector.weight.data = self.selector.weight.data.to(
+            self.lora_up.weight.device
+        ).to(self.lora_up.weight.dtype)
+
+
+class MultiLoraInjectedLinear(nn.Module):
+    def __init__(
+        self, in_features, out_features, bias=False, r=4, dropout_p=0.1, lora_num=1, scales=[1.0]
+    ):
+        super().__init__()
+
+        if r > min(in_features, out_features):
+            #raise ValueError(
+            #    f"LoRA rank {r} must be less or equal than {min(in_features, out_features)}"
+            #)
+            print(f"LoRA rank {r} is too large. setting to: {min(in_features, out_features)}")
+            r = min(in_features, out_features)
+
+        self.r = r
+        self.linear = nn.Linear(in_features, out_features, bias)
+
+        for i in range(lora_num):
+            if i==0:
+                self.lora_down =[nn.Linear(in_features, r, bias=False)]
+                self.dropout = [nn.Dropout(dropout_p)]
+                self.lora_up = [nn.Linear(r, out_features, bias=False)]
+                self.scale = scales[i]
+                self.selector = [nn.Identity()]
+            else:
+                self.lora_down.append(nn.Linear(in_features, r, bias=False))
+                self.dropout.append( nn.Dropout(dropout_p))
+                self.lora_up.append( nn.Linear(r, out_features, bias=False))
+                self.scale.append(scales[i])
+
+        nn.init.normal_(self.lora_down.weight, std=1 / r)
+        nn.init.zeros_(self.lora_up.weight)
+
+    def forward(self, input):
+        return (
+            self.linear(input)
+            + self.dropout(self.lora_up(self.selector(self.lora_down(input))))
+            * self.scale
+        )
+
+    def realize_as_lora(self):
+        return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
+
+    def set_selector_from_diag(self, diag: torch.Tensor):
+        # diag is a 1D tensor of size (r,)
+        assert diag.shape == (self.r,)
+        self.selector = nn.Linear(self.r, self.r, bias=False)
+        self.selector.weight.data = torch.diag(diag)
+        self.selector.weight.data = self.selector.weight.data.to(
+            self.lora_up.weight.device
+        ).to(self.lora_up.weight.dtype)
+
+
+class LoraInjectedConv2d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups: int = 1,
+        bias: bool = True,
+        r: int = 4,
+        dropout_p: float = 0.1,
+        scale: float = 1.0,
+    ):
+        super().__init__()
+        if r > min(in_channels, out_channels):
+            print(f"LoRA rank {r} is too large. setting to: {min(in_channels, out_channels)}")
+            r = min(in_channels, out_channels)
+
+        self.r = r
+        self.conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+        )
+
+        self.lora_down = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=r,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=False,
+        )
+        self.dropout = nn.Dropout(dropout_p)
+        self.lora_up = nn.Conv2d(
+            in_channels=r,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+        self.selector = nn.Identity()
+        self.scale = scale
+
+        nn.init.normal_(self.lora_down.weight, std=1 / r)
+        nn.init.zeros_(self.lora_up.weight)
+
+    def forward(self, input):
+        return (
+            self.conv(input)
+            + self.dropout(self.lora_up(self.selector(self.lora_down(input))))
+            * self.scale
+        )
+
+    def realize_as_lora(self):
+        return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
+
+    def set_selector_from_diag(self, diag: torch.Tensor):
+        # diag is a 1D tensor of size (r,)
+        assert diag.shape == (self.r,)
+        self.selector = nn.Conv2d(
+            in_channels=self.r,
+            out_channels=self.r,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+        self.selector.weight.data = torch.diag(diag)
+
+        # same device + dtype as lora_up
+        self.selector.weight.data = self.selector.weight.data.to(
+            self.lora_up.weight.device
+        ).to(self.lora_up.weight.dtype)
+
+class LoraInjectedConv3d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: (3, 1, 1),
+        padding: (1, 0, 0),
+        bias: bool = False,
+        r: int = 4,
+        dropout_p: float = 0,
+        scale: float = 1.0,
+    ):
+        super().__init__()
+        if r > min(in_channels, out_channels):
+            print(f"LoRA rank {r} is too large. setting to: {min(in_channels, out_channels)}")
+            r = min(in_channels, out_channels)
+
+        self.r = r
+        self.kernel_size = kernel_size
+        self.padding = padding
+        self.conv = nn.Conv3d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            padding=padding,
+        )
+
+        self.lora_down = nn.Conv3d(
+            in_channels=in_channels,
+            out_channels=r,
+            kernel_size=kernel_size,
+            bias=False,
+            padding=padding
+        )
+        self.dropout = nn.Dropout(dropout_p)
+        self.lora_up = nn.Conv3d(
+            in_channels=r,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+        self.selector = nn.Identity()
+        self.scale = scale
+
+        nn.init.normal_(self.lora_down.weight, std=1 / r)
+        nn.init.zeros_(self.lora_up.weight)
+
+    def forward(self, input):
+        return (
+            self.conv(input)
+            + self.dropout(self.lora_up(self.selector(self.lora_down(input))))
+            * self.scale
+        )
+
+    def realize_as_lora(self):
+        return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
+
+    def set_selector_from_diag(self, diag: torch.Tensor):
+        # diag is a 1D tensor of size (r,)
+        assert diag.shape == (self.r,)
+        self.selector = nn.Conv3d(
+            in_channels=self.r,
+            out_channels=self.r,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+        self.selector.weight.data = torch.diag(diag)
+
+        # same device + dtype as lora_up
+        self.selector.weight.data = self.selector.weight.data.to(
+            self.lora_up.weight.device
+        ).to(self.lora_up.weight.dtype)
+
+UNET_DEFAULT_TARGET_REPLACE = {"CrossAttention", "Attention", "GEGLU"}
+
+UNET_EXTENDED_TARGET_REPLACE = {"ResnetBlock2D", "CrossAttention", "Attention", "GEGLU"}
+
+TEXT_ENCODER_DEFAULT_TARGET_REPLACE = {"CLIPAttention"}
+
+TEXT_ENCODER_EXTENDED_TARGET_REPLACE = {"CLIPAttention"}
+
+DEFAULT_TARGET_REPLACE = UNET_DEFAULT_TARGET_REPLACE
+
+EMBED_FLAG = "<embed>"
+
+
+def _find_children(
+    model,
+    search_class: List[Type[nn.Module]] = [nn.Linear],
+):
+    """
+    Find all modules of a certain class (or union of classes).
+
+    Returns all matching modules, along with the parent of those moduless and the
+    names they are referenced by.
+    """
+    # For each target find every linear_class module that isn't a child of a LoraInjectedLinear
+    for parent in model.modules():
+        for name, module in parent.named_children():
+            if any([isinstance(module, _class) for _class in search_class]):
+                yield parent, name, module
+
+
+def _find_modules_v2(
+    model,
+    ancestor_class: Optional[Set[str]] = None,
+    search_class: List[Type[nn.Module]] = [nn.Linear],
+    exclude_children_of: Optional[List[Type[nn.Module]]] = None,
+    #     [
+    #     LoraInjectedLinear,
+    #     LoraInjectedConv2d,
+    #     LoraInjectedConv3d
+    # ],
+):
+    """
+    Find all modules of a certain class (or union of classes) that are direct or
+    indirect descendants of other modules of a certain class (or union of classes).
+
+    Returns all matching modules, along with the parent of those moduless and the
+    names they are referenced by.
+    """
+
+    # Get the targets we should replace all linears under
+    if ancestor_class is not None:
+        ancestors = (
+            module
+            for name, module in model.named_modules()
+            if module.__class__.__name__ in ancestor_class # and ('transformer_in' not in name)
+        )
+    else:
+        # this, incase you want to naively iterate over all modules.
+        ancestors = [module for module in model.modules()]
+
+    # For each target find every linear_class module that isn't a child of a LoraInjectedLinear
+    for ancestor in ancestors:
+        for fullname, module in ancestor.named_modules():
+            if any([isinstance(module, _class) for _class in search_class]):
+                continue_flag = True
+                if 'Transformer2DModel' in ancestor_class and ('attn1' in fullname or 'ff' in fullname):
+                    continue_flag = False
+                if 'TransformerTemporalModel' in ancestor_class and ('attn1' in fullname or 'attn2' in fullname or 'ff' in fullname):
+                    continue_flag = False
+                if continue_flag:
+                    continue
+                # Find the direct parent if this is a descendant, not a child, of target
+                *path, name = fullname.split(".")
+                parent = ancestor
+                while path:
+                    parent = parent.get_submodule(path.pop(0))
+                # Skip this linear if it's a child of a LoraInjectedLinear
+                if exclude_children_of and any(
+                    [isinstance(parent, _class) for _class in exclude_children_of]
+                ):
+                    continue
+                if name in ['lora_up', 'dropout', 'lora_down']:
+                    continue
+                # Otherwise, yield it
+                yield parent, name, module
+
+
+def _find_modules_old(
+    model,
+    ancestor_class: Set[str] = DEFAULT_TARGET_REPLACE,
+    search_class: List[Type[nn.Module]] = [nn.Linear],
+    exclude_children_of: Optional[List[Type[nn.Module]]] = [LoraInjectedLinear],
+):
+    ret = []
+    for _module in model.modules():
+        if _module.__class__.__name__ in ancestor_class:
+
+            for name, _child_module in _module.named_modules():
+                if _child_module.__class__ in search_class:
+                    ret.append((_module, name, _child_module))
+    print(ret)
+    return ret
+
+
+_find_modules = _find_modules_v2
+
+
+def inject_trainable_lora(
+    model: nn.Module,
+    target_replace_module: Set[str] = DEFAULT_TARGET_REPLACE,
+    r: int = 4,
+    loras=None,  # path to lora .pt
+    verbose: bool = False,
+    dropout_p: float = 0.0,
+    scale: float = 1.0,
+):
+    """
+    inject lora into model, and returns lora parameter groups.
+    """
+
+    require_grad_params = []
+    names = []
+
+    if loras != None:
+        loras = torch.load(loras)
+
+    for _module, name, _child_module in _find_modules(
+        model, target_replace_module, search_class=[nn.Linear]
+    ):
+        weight = _child_module.weight
+        bias = _child_module.bias
+        if verbose:
+            print("LoRA Injection : injecting lora into ", name)
+            print("LoRA Injection : weight shape", weight.shape)
+        _tmp = LoraInjectedLinear(
+            _child_module.in_features,
+            _child_module.out_features,
+            _child_module.bias is not None,
+            r=r,
+            dropout_p=dropout_p,
+            scale=scale,
+        )
+        _tmp.linear.weight = weight
+        if bias is not None:
+            _tmp.linear.bias = bias
+
+        # switch the module
+        _tmp.to(_child_module.weight.device).to(_child_module.weight.dtype)
+        _module._modules[name] = _tmp
+
+        require_grad_params.append(_module._modules[name].lora_up.parameters())
+        require_grad_params.append(_module._modules[name].lora_down.parameters())
+
+        if loras != None:
+            _module._modules[name].lora_up.weight = loras.pop(0)
+            _module._modules[name].lora_down.weight = loras.pop(0)
+
+        _module._modules[name].lora_up.weight.requires_grad = True
+        _module._modules[name].lora_down.weight.requires_grad = True
+        names.append(name)
+
+    return require_grad_params, names
+
+
+def inject_trainable_lora_extended(
+    model: nn.Module,
+    target_replace_module: Set[str] = UNET_EXTENDED_TARGET_REPLACE,
+    r: int = 4,
+    loras=None,  # path to lora .pt
+    dropout_p: float = 0.0,
+    scale: float = 1.0,
+):
+    """
+    inject lora into model, and returns lora parameter groups.
+    """
+
+    require_grad_params = []
+    names = []
+
+    if loras != None:
+        loras = torch.load(loras)
+    if True:
+        for target_replace_module_i in target_replace_module:
+            for _module, name, _child_module in _find_modules(
+                model, [target_replace_module_i], search_class=[nn.Linear, nn.Conv2d, nn.Conv3d]
+            ):
+                # if name == 'to_q':
+                #     continue
+                if _child_module.__class__ == nn.Linear:
+                    weight = _child_module.weight
+                    bias = _child_module.bias
+                    _tmp = LoraInjectedLinear(
+                        _child_module.in_features,
+                        _child_module.out_features,
+                        _child_module.bias is not None,
+                        r=r,
+                        dropout_p=dropout_p,
+                        scale=scale,
+                    )
+                    _tmp.linear.weight = weight
+                    if bias is not None:
+                        _tmp.linear.bias = bias
+                elif _child_module.__class__ == nn.Conv2d:
+                    weight = _child_module.weight
+                    bias = _child_module.bias
+                    _tmp = LoraInjectedConv2d(
+                        _child_module.in_channels,
+                        _child_module.out_channels,
+                        _child_module.kernel_size,
+                        _child_module.stride,
+                        _child_module.padding,
+                        _child_module.dilation,
+                        _child_module.groups,
+                        _child_module.bias is not None,
+                        r=r,
+                        dropout_p=dropout_p,
+                        scale=scale,
+                    )
+
+                    _tmp.conv.weight = weight
+                    if bias is not None:
+                        _tmp.conv.bias = bias
+
+                elif _child_module.__class__ == nn.Conv3d:
+                    weight = _child_module.weight
+                    bias = _child_module.bias
+                    _tmp = LoraInjectedConv3d(
+                        _child_module.in_channels,
+                        _child_module.out_channels,
+                        bias=_child_module.bias is not None,
+                        kernel_size=_child_module.kernel_size,
+                        padding=_child_module.padding,
+                        r=r,
+                        dropout_p=dropout_p,
+                        scale=scale,
+                    )
+
+                    _tmp.conv.weight = weight
+                    if bias is not None:
+                        _tmp.conv.bias = bias
+                # switch the module
+                _tmp.to(_child_module.weight.device).to(_child_module.weight.dtype)
+                if bias is not None:
+                    _tmp.to(_child_module.bias.device).to(_child_module.bias.dtype)
+
+                _module._modules[name] = _tmp
+                require_grad_params.append(_module._modules[name].lora_up.parameters())
+                require_grad_params.append(_module._modules[name].lora_down.parameters())
+
+                if loras != None:
+                    _module._modules[name].lora_up.weight = loras.pop(0)
+                    _module._modules[name].lora_down.weight = loras.pop(0)
+
+                _module._modules[name].lora_up.weight.requires_grad = True
+                _module._modules[name].lora_down.weight.requires_grad = True
+                names.append(name)
+    else:
+        for _module, name, _child_module in _find_modules(
+                model, target_replace_module, search_class=[nn.Linear, nn.Conv2d, nn.Conv3d]
+        ):
+            if _child_module.__class__ == nn.Linear:
+                weight = _child_module.weight
+                bias = _child_module.bias
+                _tmp = LoraInjectedLinear(
+                    _child_module.in_features,
+                    _child_module.out_features,
+                    _child_module.bias is not None,
+                    r=r,
+                    dropout_p=dropout_p,
+                    scale=scale,
+                )
+                _tmp.linear.weight = weight
+                if bias is not None:
+                    _tmp.linear.bias = bias
+            elif _child_module.__class__ == nn.Conv2d:
+                weight = _child_module.weight
+                bias = _child_module.bias
+                _tmp = LoraInjectedConv2d(
+                    _child_module.in_channels,
+                    _child_module.out_channels,
+                    _child_module.kernel_size,
+                    _child_module.stride,
+                    _child_module.padding,
+                    _child_module.dilation,
+                    _child_module.groups,
+                    _child_module.bias is not None,
+                    r=r,
+                    dropout_p=dropout_p,
+                    scale=scale,
+                )
+
+                _tmp.conv.weight = weight
+                if bias is not None:
+                    _tmp.conv.bias = bias
+
+            elif _child_module.__class__ == nn.Conv3d:
+                weight = _child_module.weight
+                bias = _child_module.bias
+                _tmp = LoraInjectedConv3d(
+                    _child_module.in_channels,
+                    _child_module.out_channels,
+                    bias=_child_module.bias is not None,
+                    kernel_size=_child_module.kernel_size,
+                    padding=_child_module.padding,
+                    r=r,
+                    dropout_p=dropout_p,
+                    scale=scale,
+                )
+
+                _tmp.conv.weight = weight
+                if bias is not None:
+                    _tmp.conv.bias = bias
+            # switch the module
+            _tmp.to(_child_module.weight.device).to(_child_module.weight.dtype)
+            if bias is not None:
+                _tmp.to(_child_module.bias.device).to(_child_module.bias.dtype)
+
+            _module._modules[name] = _tmp
+            require_grad_params.append(_module._modules[name].lora_up.parameters())
+            require_grad_params.append(_module._modules[name].lora_down.parameters())
+
+            if loras != None:
+                _module._modules[name].lora_up.weight = loras.pop(0)
+                _module._modules[name].lora_down.weight = loras.pop(0)
+
+            _module._modules[name].lora_up.weight.requires_grad = True
+            _module._modules[name].lora_down.weight.requires_grad = True
+            names.append(name)
+
+    return require_grad_params, names
+
+
+def inject_inferable_lora(
+        model, 
+        lora_path='', 
+        unet_replace_modules=["UNet3DConditionModel"], 
+        text_encoder_replace_modules=["CLIPEncoderLayer"],
+        is_extended=False, 
+        r=16
+    ):    
+    from transformers.models.clip import CLIPTextModel
+    from diffusers import UNet3DConditionModel
+
+    def is_text_model(f): return 'text_encoder' in f and isinstance(model.text_encoder, CLIPTextModel)
+    def is_unet(f): return 'unet' in f and model.unet.__class__.__name__ == "UNet3DConditionModel"
+
+    if os.path.exists(lora_path):
+        try:
+            for f in os.listdir(lora_path):
+                if f.endswith('.pt'):
+                    lora_file = os.path.join(lora_path, f)
+
+                    if is_text_model(f):
+                        monkeypatch_or_replace_lora(
+                            model.text_encoder,
+                            torch.load(lora_file),
+                            target_replace_module=text_encoder_replace_modules,
+                            r=r
+                        )
+                        print("Successfully loaded Text Encoder LoRa.")
+                        continue
+
+                    if is_unet(f):
+                        monkeypatch_or_replace_lora_extended(
+                            model.unet,
+                            torch.load(lora_file),
+                            target_replace_module=unet_replace_modules,
+                            r=r
+                        )
+                        print("Successfully loaded UNET LoRa.")
+                        continue
+
+                    print("Found a .pt file, but doesn't have the correct name format. (unet.pt, text_encoder.pt)")
+
+        except Exception as e:
+            print(e)
+            print("Couldn't inject LoRA's due to an error.")
+
+def extract_lora_ups_down(model, target_replace_module=DEFAULT_TARGET_REPLACE):
+
+    loras = []
+
+    for target_replace_module_i in target_replace_module:
+
+        for _m, _n, _child_module in _find_modules(
+            model,
+            [target_replace_module_i],
+            search_class=[LoraInjectedLinear, LoraInjectedConv2d, LoraInjectedConv3d],
+        ):
+            loras.append((_child_module.lora_up, _child_module.lora_down))
+
+    if len(loras) == 0:
+        raise ValueError("No lora injected.")
+
+    return loras
+
+
+def extract_lora_child_module(model, target_replace_module=DEFAULT_TARGET_REPLACE):
+
+    loras = []
+
+    for target_replace_module_i in target_replace_module:
+
+        for _m, _n, _child_module in _find_modules(
+            model,
+            [target_replace_module_i],
+            search_class=[LoraInjectedLinear, LoraInjectedConv2d, LoraInjectedConv3d],
+        ):
+            loras.append(_child_module)
+
+    if len(loras) == 0:
+        raise ValueError("No lora injected.")
+
+    return loras
+
+def extract_lora_as_tensor(
+    model, target_replace_module=DEFAULT_TARGET_REPLACE, as_fp16=True
+):
+
+    loras = []
+
+    for _m, _n, _child_module in _find_modules(
+        model,
+        target_replace_module,
+        search_class=[LoraInjectedLinear, LoraInjectedConv2d, LoraInjectedConv3d],
+    ):
+        up, down = _child_module.realize_as_lora()
+        if as_fp16:
+            up = up.to(torch.float16)
+            down = down.to(torch.float16)
+
+        loras.append((up, down))
+
+    if len(loras) == 0:
+        raise ValueError("No lora injected.")
+
+    return loras
+
+
+def save_lora_weight(
+    model,
+    path="./lora.pt",
+    target_replace_module=DEFAULT_TARGET_REPLACE,
+    flag=None
+):  
+    weights = []
+    for _up, _down in extract_lora_ups_down(
+        model, target_replace_module=target_replace_module
+    ):
+        weights.append(_up.weight.to("cpu").to(torch.float32))
+        weights.append(_down.weight.to("cpu").to(torch.float32))
+    if not flag:
+        torch.save(weights, path)
+    else:
+        weights_new=[]
+        for i in range(0, len(weights), 4):
+            subset = weights[i+(flag-1)*2:i+(flag-1)*2+2]
+            weights_new.extend(subset)
+        torch.save(weights_new, path)
+
+def save_lora_as_json(model, path="./lora.json"):
+    weights = []
+    for _up, _down in extract_lora_ups_down(model):
+        weights.append(_up.weight.detach().cpu().numpy().tolist())
+        weights.append(_down.weight.detach().cpu().numpy().tolist())
+
+    import json
+
+    with open(path, "w") as f:
+        json.dump(weights, f)
+
+
+def save_safeloras_with_embeds(
+    modelmap: Dict[str, Tuple[nn.Module, Set[str]]] = {},
+    embeds: Dict[str, torch.Tensor] = {},
+    outpath="./lora.safetensors",
+):
+    """
+    Saves the Lora from multiple modules in a single safetensor file.
+
+    modelmap is a dictionary of {
+        "module name": (module, target_replace_module)
+    }
+    """
+    weights = {}
+    metadata = {}
+
+    for name, (model, target_replace_module) in modelmap.items():
+        metadata[name] = json.dumps(list(target_replace_module))
+
+        for i, (_up, _down) in enumerate(
+            extract_lora_as_tensor(model, target_replace_module)
+        ):
+            rank = _down.shape[0]
+
+            metadata[f"{name}:{i}:rank"] = str(rank)
+            weights[f"{name}:{i}:up"] = _up
+            weights[f"{name}:{i}:down"] = _down
+
+    for token, tensor in embeds.items():
+        metadata[token] = EMBED_FLAG
+        weights[token] = tensor
+
+    print(f"Saving weights to {outpath}")
+    safe_save(weights, outpath, metadata)
+
+
+def save_safeloras(
+    modelmap: Dict[str, Tuple[nn.Module, Set[str]]] = {},
+    outpath="./lora.safetensors",
+):
+    return save_safeloras_with_embeds(modelmap=modelmap, outpath=outpath)
+
+
+def convert_loras_to_safeloras_with_embeds(
+    modelmap: Dict[str, Tuple[str, Set[str], int]] = {},
+    embeds: Dict[str, torch.Tensor] = {},
+    outpath="./lora.safetensors",
+):
+    """
+    Converts the Lora from multiple pytorch .pt files into a single safetensor file.
+
+    modelmap is a dictionary of {
+        "module name": (pytorch_model_path, target_replace_module, rank)
+    }
+    """
+
+    weights = {}
+    metadata = {}
+
+    for name, (path, target_replace_module, r) in modelmap.items():
+        metadata[name] = json.dumps(list(target_replace_module))
+
+        lora = torch.load(path)
+        for i, weight in enumerate(lora):
+            is_up = i % 2 == 0
+            i = i // 2
+
+            if is_up:
+                metadata[f"{name}:{i}:rank"] = str(r)
+                weights[f"{name}:{i}:up"] = weight
+            else:
+                weights[f"{name}:{i}:down"] = weight
+
+    for token, tensor in embeds.items():
+        metadata[token] = EMBED_FLAG
+        weights[token] = tensor
+
+    print(f"Saving weights to {outpath}")
+    safe_save(weights, outpath, metadata)
+
+
+def convert_loras_to_safeloras(
+    modelmap: Dict[str, Tuple[str, Set[str], int]] = {},
+    outpath="./lora.safetensors",
+):
+    convert_loras_to_safeloras_with_embeds(modelmap=modelmap, outpath=outpath)
+
+
+def parse_safeloras(
+    safeloras,
+) -> Dict[str, Tuple[List[nn.parameter.Parameter], List[int], List[str]]]:
+    """
+    Converts a loaded safetensor file that contains a set of module Loras
+    into Parameters and other information
+
+    Output is a dictionary of {
+        "module name": (
+            [list of weights],
+            [list of ranks],
+            target_replacement_modules
+        )
+    }
+    """
+    loras = {}
+    metadata = safeloras.metadata()
+
+    get_name = lambda k: k.split(":")[0]
+
+    keys = list(safeloras.keys())
+    keys.sort(key=get_name)
+
+    for name, module_keys in groupby(keys, get_name):
+        info = metadata.get(name)
+
+        if not info:
+            raise ValueError(
+                f"Tensor {name} has no metadata - is this a Lora safetensor?"
+            )
+
+        # Skip Textual Inversion embeds
+        if info == EMBED_FLAG:
+            continue
+
+        # Handle Loras
+        # Extract the targets
+        target = json.loads(info)
+
+        # Build the result lists - Python needs us to preallocate lists to insert into them
+        module_keys = list(module_keys)
+        ranks = [4] * (len(module_keys) // 2)
+        weights = [None] * len(module_keys)
+
+        for key in module_keys:
+            # Split the model name and index out of the key
+            _, idx, direction = key.split(":")
+            idx = int(idx)
+
+            # Add the rank
+            ranks[idx] = int(metadata[f"{name}:{idx}:rank"])
+
+            # Insert the weight into the list
+            idx = idx * 2 + (1 if direction == "down" else 0)
+            weights[idx] = nn.parameter.Parameter(safeloras.get_tensor(key))
+
+        loras[name] = (weights, ranks, target)
+
+    return loras
+
+
+def parse_safeloras_embeds(
+    safeloras,
+) -> Dict[str, torch.Tensor]:
+    """
+    Converts a loaded safetensor file that contains Textual Inversion embeds into
+    a dictionary of embed_token: Tensor
+    """
+    embeds = {}
+    metadata = safeloras.metadata()
+
+    for key in safeloras.keys():
+        # Only handle Textual Inversion embeds
+        meta = metadata.get(key)
+        if not meta or meta != EMBED_FLAG:
+            continue
+
+        embeds[key] = safeloras.get_tensor(key)
+
+    return embeds
+
+
+def load_safeloras(path, device="cpu"):
+    safeloras = safe_open(path, framework="pt", device=device)
+    return parse_safeloras(safeloras)
+
+
+def load_safeloras_embeds(path, device="cpu"):
+    safeloras = safe_open(path, framework="pt", device=device)
+    return parse_safeloras_embeds(safeloras)
+
+
+def load_safeloras_both(path, device="cpu"):
+    safeloras = safe_open(path, framework="pt", device=device)
+    return parse_safeloras(safeloras), parse_safeloras_embeds(safeloras)
+
+
+def collapse_lora(model, alpha=1.0):
+
+    for _module, name, _child_module in _find_modules(
+        model,
+        UNET_EXTENDED_TARGET_REPLACE | TEXT_ENCODER_EXTENDED_TARGET_REPLACE,
+        search_class=[LoraInjectedLinear, LoraInjectedConv2d, LoraInjectedConv3d],
+    ):
+
+        if isinstance(_child_module, LoraInjectedLinear):
+            print("Collapsing Lin Lora in", name)
+
+            _child_module.linear.weight = nn.Parameter(
+                _child_module.linear.weight.data
+                + alpha
+                * (
+                    _child_module.lora_up.weight.data
+                    @ _child_module.lora_down.weight.data
+                )
+                .type(_child_module.linear.weight.dtype)
+                .to(_child_module.linear.weight.device)
+            )
+
+        else:
+            print("Collapsing Conv Lora in", name)
+            _child_module.conv.weight = nn.Parameter(
+                _child_module.conv.weight.data
+                + alpha
+                * (
+                    _child_module.lora_up.weight.data.flatten(start_dim=1)
+                    @ _child_module.lora_down.weight.data.flatten(start_dim=1)
+                )
+                .reshape(_child_module.conv.weight.data.shape)
+                .type(_child_module.conv.weight.dtype)
+                .to(_child_module.conv.weight.device)
+            )
+
+
+def monkeypatch_or_replace_lora(
+    model,
+    loras,
+    target_replace_module=DEFAULT_TARGET_REPLACE,
+    r: Union[int, List[int]] = 4,
+):
+    for _module, name, _child_module in _find_modules(
+        model, target_replace_module, search_class=[nn.Linear, LoraInjectedLinear]
+    ):
+        _source = (
+            _child_module.linear
+            if isinstance(_child_module, LoraInjectedLinear)
+            else _child_module
+        )
+
+        weight = _source.weight
+        bias = _source.bias
+        _tmp = LoraInjectedLinear(
+            _source.in_features,
+            _source.out_features,
+            _source.bias is not None,
+            r=r.pop(0) if isinstance(r, list) else r,
+        )
+        _tmp.linear.weight = weight
+
+        if bias is not None:
+            _tmp.linear.bias = bias
+
+        # switch the module
+        _module._modules[name] = _tmp
+
+        up_weight = loras.pop(0)
+        down_weight = loras.pop(0)
+
+        _module._modules[name].lora_up.weight = nn.Parameter(
+            up_weight.type(weight.dtype)
+        )
+        _module._modules[name].lora_down.weight = nn.Parameter(
+            down_weight.type(weight.dtype)
+        )
+
+        _module._modules[name].to(weight.device)
+
+
+def monkeypatch_or_replace_lora_extended(
+    model,
+    loras,
+    target_replace_module=DEFAULT_TARGET_REPLACE,
+    r: Union[int, List[int]] = 4,
+):
+    for _module, name, _child_module in _find_modules(
+        model,
+        target_replace_module,
+        search_class=[
+            nn.Linear, 
+            nn.Conv2d, 
+            nn.Conv3d,
+            LoraInjectedLinear, 
+            LoraInjectedConv2d, 
+            LoraInjectedConv3d,
+        ],
+    ):
+
+        if (_child_module.__class__ == nn.Linear) or (
+            _child_module.__class__ == LoraInjectedLinear
+        ):
+            if len(loras[0].shape) != 2:
+                continue
+
+            _source = (
+                _child_module.linear
+                if isinstance(_child_module, LoraInjectedLinear)
+                else _child_module
+            )
+
+            weight = _source.weight
+            bias = _source.bias
+            _tmp = LoraInjectedLinear(
+                _source.in_features,
+                _source.out_features,
+                _source.bias is not None,
+                r=r.pop(0) if isinstance(r, list) else r,
+            )
+            _tmp.linear.weight = weight
+
+            if bias is not None:
+                _tmp.linear.bias = bias
+
+        elif (_child_module.__class__ == nn.Conv2d) or (
+            _child_module.__class__ == LoraInjectedConv2d
+        ):
+            if len(loras[0].shape) != 4:
+                continue
+            _source = (
+                _child_module.conv
+                if isinstance(_child_module, LoraInjectedConv2d)
+                else _child_module
+            )
+
+            weight = _source.weight
+            bias = _source.bias
+            _tmp = LoraInjectedConv2d(
+                _source.in_channels,
+                _source.out_channels,
+                _source.kernel_size,
+                _source.stride,
+                _source.padding,
+                _source.dilation,
+                _source.groups,
+                _source.bias is not None,
+                r=r.pop(0) if isinstance(r, list) else r,
+            )
+
+            _tmp.conv.weight = weight
+
+            if bias is not None:
+                _tmp.conv.bias = bias
+
+        elif _child_module.__class__ == nn.Conv3d or(
+            _child_module.__class__ == LoraInjectedConv3d
+        ):
+
+            if len(loras[0].shape) != 5:
+                continue
+
+            _source = (
+                _child_module.conv
+                if isinstance(_child_module, LoraInjectedConv3d)
+                else _child_module
+            )
+
+            weight = _source.weight
+            bias = _source.bias
+            _tmp = LoraInjectedConv3d(
+                _source.in_channels,
+                _source.out_channels,
+                bias=_source.bias is not None,
+                kernel_size=_source.kernel_size,
+                padding=_source.padding,
+                r=r.pop(0) if isinstance(r, list) else r,
+            )
+
+            _tmp.conv.weight = weight
+
+            if bias is not None:
+                _tmp.conv.bias = bias
+
+        # switch the module
+        _module._modules[name] = _tmp
+
+        up_weight = loras.pop(0)
+        down_weight = loras.pop(0)
+
+        _module._modules[name].lora_up.weight = nn.Parameter(
+            up_weight.type(weight.dtype)
+        )
+        _module._modules[name].lora_down.weight = nn.Parameter(
+            down_weight.type(weight.dtype)
+        )
+
+        _module._modules[name].to(weight.device)
+
+
+def monkeypatch_or_replace_safeloras(models, safeloras):
+    loras = parse_safeloras(safeloras)
+
+    for name, (lora, ranks, target) in loras.items():
+        model = getattr(models, name, None)
+
+        if not model:
+            print(f"No model provided for {name}, contained in Lora")
+            continue
+
+        monkeypatch_or_replace_lora_extended(model, lora, target, ranks)
+
+
+def monkeypatch_remove_lora(model):
+    for _module, name, _child_module in _find_modules(
+        model, search_class=[LoraInjectedLinear, LoraInjectedConv2d, LoraInjectedConv3d]
+    ):
+        if isinstance(_child_module, LoraInjectedLinear):
+            _source = _child_module.linear
+            weight, bias = _source.weight, _source.bias
+
+            _tmp = nn.Linear(
+                _source.in_features, _source.out_features, bias is not None
+            )
+
+            _tmp.weight = weight
+            if bias is not None:
+                _tmp.bias = bias
+
+        else:
+            _source = _child_module.conv
+            weight, bias = _source.weight, _source.bias
+
+            if isinstance(_source, nn.Conv2d):
+                _tmp = nn.Conv2d(
+                    in_channels=_source.in_channels,
+                    out_channels=_source.out_channels,
+                    kernel_size=_source.kernel_size,
+                    stride=_source.stride,
+                    padding=_source.padding,
+                    dilation=_source.dilation,
+                    groups=_source.groups,
+                    bias=bias is not None,
+                )
+
+                _tmp.weight = weight
+                if bias is not None:
+                    _tmp.bias = bias
+            
+            if isinstance(_source, nn.Conv3d):
+                _tmp = nn.Conv3d(
+                _source.in_channels,
+                _source.out_channels,
+                bias=_source.bias is not None,
+                kernel_size=_source.kernel_size,
+                padding=_source.padding,
+            )
+
+            _tmp.weight = weight
+            if bias is not None:
+                _tmp.bias = bias   
+
+        _module._modules[name] = _tmp
+
+
+def monkeypatch_add_lora(
+    model,
+    loras,
+    target_replace_module=DEFAULT_TARGET_REPLACE,
+    alpha: float = 1.0,
+    beta: float = 1.0,
+):
+    for _module, name, _child_module in _find_modules(
+        model, target_replace_module, search_class=[LoraInjectedLinear]
+    ):
+        weight = _child_module.linear.weight
+
+        up_weight = loras.pop(0)
+        down_weight = loras.pop(0)
+
+        _module._modules[name].lora_up.weight = nn.Parameter(
+            up_weight.type(weight.dtype).to(weight.device) * alpha
+            + _module._modules[name].lora_up.weight.to(weight.device) * beta
+        )
+        _module._modules[name].lora_down.weight = nn.Parameter(
+            down_weight.type(weight.dtype).to(weight.device) * alpha
+            + _module._modules[name].lora_down.weight.to(weight.device) * beta
+        )
+
+        _module._modules[name].to(weight.device)
+
+
+def tune_lora_scale(model, alpha: float = 1.0):
+    for _module in model.modules():
+        if _module.__class__.__name__ in ["LoraInjectedLinear", "LoraInjectedConv2d", "LoraInjectedConv3d"]:
+            _module.scale = alpha
+
+
+def set_lora_diag(model, diag: torch.Tensor):
+    for _module in model.modules():
+        if _module.__class__.__name__ in ["LoraInjectedLinear", "LoraInjectedConv2d", "LoraInjectedConv3d"]:
+            _module.set_selector_from_diag(diag)
+
+
+def _text_lora_path(path: str) -> str:
+    assert path.endswith(".pt"), "Only .pt files are supported"
+    return ".".join(path.split(".")[:-1] + ["text_encoder", "pt"])
+
+
+def _ti_lora_path(path: str) -> str:
+    assert path.endswith(".pt"), "Only .pt files are supported"
+    return ".".join(path.split(".")[:-1] + ["ti", "pt"])
+
+
+def apply_learned_embed_in_clip(
+    learned_embeds,
+    text_encoder,
+    tokenizer,
+    token: Optional[Union[str, List[str]]] = None,
+    idempotent=False,
+):
+    if isinstance(token, str):
+        trained_tokens = [token]
+    elif isinstance(token, list):
+        assert len(learned_embeds.keys()) == len(
+            token
+        ), "The number of tokens and the number of embeds should be the same"
+        trained_tokens = token
+    else:
+        trained_tokens = list(learned_embeds.keys())
+
+    for token in trained_tokens:
+        print(token)
+        embeds = learned_embeds[token]
+
+        # cast to dtype of text_encoder
+        dtype = text_encoder.get_input_embeddings().weight.dtype
+        num_added_tokens = tokenizer.add_tokens(token)
+
+        i = 1
+        if not idempotent:
+            while num_added_tokens == 0:
+                print(f"The tokenizer already contains the token {token}.")
+                token = f"{token[:-1]}-{i}>"
+                print(f"Attempting to add the token {token}.")
+                num_added_tokens = tokenizer.add_tokens(token)
+                i += 1
+        elif num_added_tokens == 0 and idempotent:
+            print(f"The tokenizer already contains the token {token}.")
+            print(f"Replacing {token} embedding.")
+
+        # resize the token embeddings
+        text_encoder.resize_token_embeddings(len(tokenizer))
+
+        # get the id for the token and assign the embeds
+        token_id = tokenizer.convert_tokens_to_ids(token)
+        text_encoder.get_input_embeddings().weight.data[token_id] = embeds
+    return token
+
+
+def load_learned_embed_in_clip(
+    learned_embeds_path,
+    text_encoder,
+    tokenizer,
+    token: Optional[Union[str, List[str]]] = None,
+    idempotent=False,
+):
+    learned_embeds = torch.load(learned_embeds_path)
+    apply_learned_embed_in_clip(
+        learned_embeds, text_encoder, tokenizer, token, idempotent
+    )
+
+
+def patch_pipe(
+    pipe,
+    maybe_unet_path,
+    token: Optional[str] = None,
+    r: int = 4,
+    patch_unet=True,
+    patch_text=True,
+    patch_ti=True,
+    idempotent_token=True,
+    unet_target_replace_module=DEFAULT_TARGET_REPLACE,
+    text_target_replace_module=TEXT_ENCODER_DEFAULT_TARGET_REPLACE,
+):
+    if maybe_unet_path.endswith(".pt"):
+        # torch format
+
+        if maybe_unet_path.endswith(".ti.pt"):
+            unet_path = maybe_unet_path[:-6] + ".pt"
+        elif maybe_unet_path.endswith(".text_encoder.pt"):
+            unet_path = maybe_unet_path[:-16] + ".pt"
+        else:
+            unet_path = maybe_unet_path
+
+        ti_path = _ti_lora_path(unet_path)
+        text_path = _text_lora_path(unet_path)
+
+        if patch_unet:
+            print("LoRA : Patching Unet")
+            monkeypatch_or_replace_lora(
+                pipe.unet,
+                torch.load(unet_path),
+                r=r,
+                target_replace_module=unet_target_replace_module,
+            )
+
+        if patch_text:
+            print("LoRA : Patching text encoder")
+            monkeypatch_or_replace_lora(
+                pipe.text_encoder,
+                torch.load(text_path),
+                target_replace_module=text_target_replace_module,
+                r=r,
+            )
+        if patch_ti:
+            print("LoRA : Patching token input")
+            token = load_learned_embed_in_clip(
+                ti_path,
+                pipe.text_encoder,
+                pipe.tokenizer,
+                token=token,
+                idempotent=idempotent_token,
+            )
+
+    elif maybe_unet_path.endswith(".safetensors"):
+        safeloras = safe_open(maybe_unet_path, framework="pt", device="cpu")
+        monkeypatch_or_replace_safeloras(pipe, safeloras)
+        tok_dict = parse_safeloras_embeds(safeloras)
+        if patch_ti:
+            apply_learned_embed_in_clip(
+                tok_dict,
+                pipe.text_encoder,
+                pipe.tokenizer,
+                token=token,
+                idempotent=idempotent_token,
+            )
+        return tok_dict
+
+
+def train_patch_pipe(pipe, patch_unet, patch_text):
+    if patch_unet:
+        print("LoRA : Patching Unet")
+        collapse_lora(pipe.unet)
+        monkeypatch_remove_lora(pipe.unet)
+
+    if patch_text:
+        print("LoRA : Patching text encoder")
+
+        collapse_lora(pipe.text_encoder)
+        monkeypatch_remove_lora(pipe.text_encoder)
+
+@torch.no_grad()
+def inspect_lora(model):
+    moved = {}
+
+    for name, _module in model.named_modules():
+        if _module.__class__.__name__ in ["LoraInjectedLinear", "LoraInjectedConv2d", "LoraInjectedConv3d"]:
+            ups = _module.lora_up.weight.data.clone()
+            downs = _module.lora_down.weight.data.clone()
+
+            wght: torch.Tensor = ups.flatten(1) @ downs.flatten(1)
+
+            dist = wght.flatten().abs().mean().item()
+            if name in moved:
+                moved[name].append(dist)
+            else:
+                moved[name] = [dist]
+
+    return moved
+
+
+def save_all(
+    unet,
+    text_encoder,
+    save_path,
+    placeholder_token_ids=None,
+    placeholder_tokens=None,
+    save_lora=True,
+    save_ti=True,
+    target_replace_module_text=TEXT_ENCODER_DEFAULT_TARGET_REPLACE,
+    target_replace_module_unet=DEFAULT_TARGET_REPLACE,
+    safe_form=True,
+):
+    if not safe_form:
+        # save ti
+        if save_ti:
+            ti_path = _ti_lora_path(save_path)
+            learned_embeds_dict = {}
+            for tok, tok_id in zip(placeholder_tokens, placeholder_token_ids):
+                learned_embeds = text_encoder.get_input_embeddings().weight[tok_id]
+                print(
+                    f"Current Learned Embeddings for {tok}:, id {tok_id} ",
+                    learned_embeds[:4],
+                )
+                learned_embeds_dict[tok] = learned_embeds.detach().cpu()
+
+            torch.save(learned_embeds_dict, ti_path)
+            print("Ti saved to ", ti_path)
+
+        # save text encoder
+        if save_lora:
+            save_lora_weight(
+                unet, save_path, target_replace_module=target_replace_module_unet
+            )
+            print("Unet saved to ", save_path)
+
+            save_lora_weight(
+                text_encoder,
+                _text_lora_path(save_path),
+                target_replace_module=target_replace_module_text,
+            )
+            print("Text Encoder saved to ", _text_lora_path(save_path))
+
+    else:
+        assert save_path.endswith(
+            ".safetensors"
+        ), f"Save path : {save_path} should end with .safetensors"
+
+        loras = {}
+        embeds = {}
+
+        if save_lora:
+
+            loras["unet"] = (unet, target_replace_module_unet)
+            loras["text_encoder"] = (text_encoder, target_replace_module_text)
+
+        if save_ti:
+            for tok, tok_id in zip(placeholder_tokens, placeholder_token_ids):
+                learned_embeds = text_encoder.get_input_embeddings().weight[tok_id]
+                print(
+                    f"Current Learned Embeddings for {tok}:, id {tok_id} ",
+                    learned_embeds[:4],
+                )
+                embeds[tok] = learned_embeds.detach().cpu()
+
+        save_safeloras_with_embeds(loras, embeds, save_path)

utils/lora_handler.py

@@ -0,0 +1,269 @@
+import os
+from logging import warnings
+import torch
+from typing import Union
+from types import SimpleNamespace
+from models.unet_3d_condition import UNet3DConditionModel
+from transformers import CLIPTextModel
+from utils.convert_diffusers_to_original_ms_text_to_video import convert_unet_state_dict, convert_text_enc_state_dict_v20
+
+from .lora import (
+    extract_lora_ups_down,
+    inject_trainable_lora_extended,
+    save_lora_weight,
+    train_patch_pipe,
+    monkeypatch_or_replace_lora,
+    monkeypatch_or_replace_lora_extended
+)
+
+
+FILE_BASENAMES = ['unet', 'text_encoder']
+LORA_FILE_TYPES = ['.pt', '.safetensors']
+CLONE_OF_SIMO_KEYS = ['model', 'loras', 'target_replace_module', 'r']
+STABLE_LORA_KEYS = ['model', 'target_module', 'search_class', 'r', 'dropout', 'lora_bias']
+
+lora_versions = dict(
+    stable_lora = "stable_lora",
+    cloneofsimo = "cloneofsimo"
+)
+
+lora_func_types = dict(
+    loader = "loader",
+    injector = "injector"
+)
+
+lora_args = dict(
+    model = None,
+    loras = None,
+    target_replace_module = [],
+    target_module = [],
+    r = 4,
+    search_class = [torch.nn.Linear],
+    dropout = 0,
+    lora_bias = 'none'
+)
+
+LoraVersions = SimpleNamespace(**lora_versions)
+LoraFuncTypes = SimpleNamespace(**lora_func_types)
+
+LORA_VERSIONS = [LoraVersions.stable_lora, LoraVersions.cloneofsimo]
+LORA_FUNC_TYPES = [LoraFuncTypes.loader, LoraFuncTypes.injector]
+
+def filter_dict(_dict, keys=[]):
+    if len(keys) == 0:
+        assert "Keys cannot empty for filtering return dict."
+    
+    for k in keys:
+        if k not in lora_args.keys():
+            assert f"{k} does not exist in available LoRA arguments"
+            
+    return {k: v for k, v in _dict.items() if k in keys}
+
+class LoraHandler(object):
+    def __init__(
+        self, 
+        version: LORA_VERSIONS = LoraVersions.cloneofsimo, 
+        use_unet_lora: bool = False,
+        use_text_lora: bool = False,
+        save_for_webui: bool = False,
+        only_for_webui: bool = False,
+        lora_bias: str = 'none',
+        unet_replace_modules: list = None,
+        text_encoder_replace_modules: list = None
+    ):
+        self.version = version
+        self.lora_loader = self.get_lora_func(func_type=LoraFuncTypes.loader)
+        self.lora_injector = self.get_lora_func(func_type=LoraFuncTypes.injector)
+        self.lora_bias = lora_bias
+        self.use_unet_lora = use_unet_lora
+        self.use_text_lora = use_text_lora
+        self.save_for_webui = save_for_webui
+        self.only_for_webui = only_for_webui
+        self.unet_replace_modules = unet_replace_modules
+        self.text_encoder_replace_modules = text_encoder_replace_modules
+        self.use_lora = any([use_text_lora, use_unet_lora])
+
+    def is_cloneofsimo_lora(self):
+        return self.version == LoraVersions.cloneofsimo
+
+
+    def get_lora_func(self, func_type: LORA_FUNC_TYPES = LoraFuncTypes.loader):
+
+        if self.is_cloneofsimo_lora():
+
+            if func_type == LoraFuncTypes.loader:
+                return monkeypatch_or_replace_lora_extended
+
+            if func_type == LoraFuncTypes.injector:
+                return inject_trainable_lora_extended
+                
+        assert "LoRA Version does not exist."
+
+    def check_lora_ext(self, lora_file: str):
+        return lora_file.endswith(tuple(LORA_FILE_TYPES))
+
+    def get_lora_file_path(
+        self, 
+        lora_path: str, 
+        model: Union[UNet3DConditionModel, CLIPTextModel]
+    ):
+        if os.path.exists(lora_path):
+            lora_filenames = [fns for fns in os.listdir(lora_path)]
+            is_lora = self.check_lora_ext(lora_path)
+
+            is_unet = isinstance(model, UNet3DConditionModel)
+            is_text =  isinstance(model, CLIPTextModel)
+            idx = 0 if is_unet else 1
+
+            base_name = FILE_BASENAMES[idx]
+            
+            for lora_filename in lora_filenames:
+                is_lora = self.check_lora_ext(lora_filename)
+                if not is_lora:
+                    continue
+                
+                if base_name in lora_filename:
+                    return os.path.join(lora_path, lora_filename)
+
+        return None
+
+    def handle_lora_load(self, file_name:str, lora_loader_args: dict = None):
+        self.lora_loader(**lora_loader_args)
+        print(f"Successfully loaded LoRA from: {file_name}")
+    
+    def load_lora(self, model, lora_path: str = '', lora_loader_args: dict = None,):
+        try:
+            lora_file = self.get_lora_file_path(lora_path, model)
+
+            if lora_file is not None:
+                lora_loader_args.update({"lora_path": lora_file})
+                self.handle_lora_load(lora_file, lora_loader_args)
+
+            else:
+                print(f"Could not load LoRAs for {model.__class__.__name__}. Injecting new ones instead...")
+
+        except Exception as e:
+            print(f"An error occured while loading a LoRA file: {e}")
+                 
+    def get_lora_func_args(self, lora_path, use_lora, model, replace_modules, r, dropout, lora_bias, scale):
+        return_dict = lora_args.copy()
+    
+        if self.is_cloneofsimo_lora():
+            return_dict = filter_dict(return_dict, keys=CLONE_OF_SIMO_KEYS)
+            return_dict.update({
+                "model": model,
+                "loras": self.get_lora_file_path(lora_path, model),
+                "target_replace_module": replace_modules,
+                "r": r,
+                "scale": scale,
+                "dropout_p": dropout,
+            })
+
+        return return_dict
+
+    def do_lora_injection(
+        self, 
+        model, 
+        replace_modules, 
+        bias='none',
+        dropout=0,
+        r=4,
+        lora_loader_args=None,
+    ):  
+        REPLACE_MODULES = replace_modules
+
+        params = None
+        negation = None
+        is_injection_hybrid = False
+        
+        if self.is_cloneofsimo_lora():
+            is_injection_hybrid = True
+            injector_args = lora_loader_args
+
+            params, negation = self.lora_injector(**injector_args)  # inject_trainable_lora_extended
+            for _up, _down in extract_lora_ups_down(
+                model, 
+                target_replace_module=REPLACE_MODULES):
+
+                if all(x is not None for x in [_up, _down]):
+                    print(f"Lora successfully injected into {model.__class__.__name__}.")
+
+                break
+
+            return params, negation, is_injection_hybrid
+
+        return params, negation, is_injection_hybrid
+
+    def add_lora_to_model(self, use_lora, model, replace_modules, dropout=0.0, lora_path='', r=16, scale=1.0):
+
+        params = None
+        negation = None
+
+        lora_loader_args = self.get_lora_func_args(
+            lora_path,
+            use_lora,
+            model,
+            replace_modules,
+            r,
+            dropout,
+            self.lora_bias,
+            scale
+        )
+
+        if use_lora:
+            params, negation, is_injection_hybrid = self.do_lora_injection(
+                model, 
+                replace_modules, 
+                bias=self.lora_bias,
+                lora_loader_args=lora_loader_args,
+                dropout=dropout,
+                r=r
+            )
+
+            if not is_injection_hybrid:
+                self.load_lora(model, lora_path=lora_path, lora_loader_args=lora_loader_args)
+        
+        params = model if params is None else params
+        return params, negation
+
+    def save_cloneofsimo_lora(self, model, save_path, step, flag):
+        
+        def save_lora(model, name, condition, replace_modules, step, save_path, flag=None):
+            if condition and replace_modules is not None:
+                save_path = f"{save_path}/{step}_{name}.pt"
+                save_lora_weight(model, save_path, replace_modules, flag)
+
+        save_lora(
+            model.unet, 
+            FILE_BASENAMES[0], 
+            self.use_unet_lora, 
+            self.unet_replace_modules, 
+            step,
+            save_path,
+            flag
+        )
+        save_lora(
+            model.text_encoder, 
+            FILE_BASENAMES[1], 
+            self.use_text_lora, 
+            self.text_encoder_replace_modules, 
+            step, 
+            save_path,
+            flag
+        )
+
+        # train_patch_pipe(model, self.use_unet_lora, self.use_text_lora)
+
+    def save_lora_weights(self, model: None, save_path: str ='',step: str = '', flag=None):
+        save_path = f"{save_path}/lora"
+        os.makedirs(save_path, exist_ok=True)
+
+        if self.is_cloneofsimo_lora():
+            if any([self.save_for_webui, self.only_for_webui]):
+                warnings.warn(
+                    """
+                    You have 'save_for_webui' enabled, but are using cloneofsimo's LoRA implemention.
+                    Only 'stable_lora' is supported for saving to a compatible webui file.
+                    """
+                )
+            self.save_cloneofsimo_lora(model, save_path, step, flag)