Initial commit

hyvideo/utils/data_utils.py

+import numpy as np
+import math
+from PIL import Image
+import torch
+import copy
+import string
+import random
+
+
+def align_to(value, alignment):
+    """align hight, width according to alignment
+
+    Args:
+        value (int): height or width
+        alignment (int): target alignment factor
+
+    Returns:
+        int: the aligned value
+    """
+    return int(math.ceil(value / alignment) * alignment)
+
+
+def black_image(width, height):
+    """generate a black image
+
+    Args:
+        width (int): image width
+        height (int): image height
+
+    Returns:
+        _type_: a black image
+    """
+    black_image = Image.new("RGB", (width, height), (0, 0, 0))
+    return black_image
+
+
+def get_closest_ratio(height: float, width: float, ratios: list, buckets: list):
+    """get the closest ratio in the buckets
+
+    Args:
+        height (float): video height
+        width (float): video width
+        ratios (list): video aspect ratio
+        buckets (list): buckets generate by `generate_crop_size_list`
+
+    Returns:
+        the closest ratio in the buckets and the corresponding ratio
+    """
+    aspect_ratio = float(height) / float(width)
+    diff_ratios = ratios - aspect_ratio
+
+    if aspect_ratio >= 1:
+        indices = [(index, x) for index, x in enumerate(diff_ratios) if x <= 0]
+    else:
+        indices = [(index, x) for index, x in enumerate(diff_ratios) if x > 0]
+
+    closest_ratio_id = min(indices, key=lambda pair: abs(pair[1]))[0]
+    closest_size = buckets[closest_ratio_id]
+    closest_ratio = ratios[closest_ratio_id]
+
+    return closest_size, closest_ratio
+
+
+def generate_crop_size_list(base_size=256, patch_size=32, max_ratio=4.0):
+    """generate crop size list
+
+    Args:
+        base_size (int, optional): the base size for generate bucket. Defaults to 256.
+        patch_size (int, optional): the stride to generate bucket. Defaults to 32.
+        max_ratio (float, optional): th max ratio for h or w based on base_size . Defaults to 4.0.
+
+    Returns:
+        list: generate crop size list
+    """
+    num_patches = round((base_size / patch_size) ** 2)
+    assert max_ratio >= 1.0
+    crop_size_list = []
+    wp, hp = num_patches, 1
+    while wp > 0:
+        if max(wp, hp) / min(wp, hp) <= max_ratio:
+            crop_size_list.append((wp * patch_size, hp * patch_size))
+        if (hp + 1) * wp <= num_patches:
+            hp += 1
+        else:
+            wp -= 1
+    return crop_size_list
+
+
+def align_floor_to(value, alignment):
+    """align hight, width according to alignment
+
+    Args:
+        value (int): height or width
+        alignment (int): target alignment factor
+
+    Returns:
+        int: the aligned value
+    """
+    return int(math.floor(value / alignment) * alignment)

hyvideo/utils/file_utils.py

+import logging
+import os
+from pathlib import Path
+import json
+import tarfile
+from collections import defaultdict
+from einops import rearrange
+from typing import List
+import torch
+import torchvision
+import numpy as np
+import imageio
+import PIL.Image
+from PIL import Image
+
+CODE_SUFFIXES = {
+    ".py",  # Python codes
+    ".sh",  # Shell scripts
+    ".yaml",
+    ".yml",  # Configuration files
+}
+
+
+def build_pretraining_data_loader():
+    pass
+
+
+def logger_filter(name):
+    def filter_(record):
+        return record["extra"].get("name") == name
+
+    return filter_
+
+
+def resolve_resume_path(resume, results_dir):
+    # Detect the resume path. Support both the experiment index and the full path.
+    if resume.isnumeric():
+        tmp_dirs = list(Path(results_dir).glob("*"))
+        id2exp_dir = defaultdict(list)
+        for tmp_dir in tmp_dirs:
+            part0 = tmp_dir.name.split("_")[0]
+            if part0.isnumeric():
+                id2exp_dir[int(part0)].append(tmp_dir)
+        resume_id = int(resume)
+        valid_exp_dir = id2exp_dir.get(resume_id)
+        if len(valid_exp_dir) == 0:
+            raise ValueError(
+                f"No valid experiment directories found in {results_dir} with the experiment "
+                f"index {resume}."
+            )
+        elif len(valid_exp_dir) > 1:
+            raise ValueError(
+                f"Multiple valid experiment directories found in {results_dir} with the experiment "
+                f"index {resume}: {valid_exp_dir}."
+            )
+        resume_path = valid_exp_dir[0] / "checkpoints"
+    else:
+        resume_path = Path(resume)
+
+    if not resume_path.exists():
+        raise FileNotFoundError(f"Resume path {resume_path} not found.")
+
+    return resume_path
+
+
+def dump_codes(save_path, root, sub_dirs=None, valid_suffixes=None, save_prefix="./"):
+    """
+    Dump codes to the experiment directory.
+
+    Args:
+        save_path (str): Path to the experiment directory.
+        root (Path): Path to the root directory of the codes.
+        sub_dirs (list): List of subdirectories to be dumped. If None, all files in the root directory will
+            be dumped. (default: None)
+        valid_suffixes (tuple, optional): Valid suffixes of the files to be dumped. If None, CODE_SUFFIXES will be used.
+            (default: None)
+        save_prefix (str, optional): Prefix to be added to the files in the tarball. (default: './')
+    """
+    if valid_suffixes is None:
+        valid_suffixes = CODE_SUFFIXES
+
+    # Force to use tar.gz suffix
+    save_path = safe_file(save_path)
+    assert save_path.name.endswith(
+        ".tar.gz"
+    ), f"save_path should end with .tar.gz, got {save_path.name}."
+    # Make root absolute
+    root = Path(root).absolute()
+    # Make a tarball of the codes
+    with tarfile.open(save_path, "w:gz") as tar:
+        # Recursively add all files in the root directory
+        if sub_dirs is None:
+            sub_dirs = list(root.iterdir())
+        for sub_dir in sub_dirs:
+            for file in Path(sub_dir).rglob("*"):
+                if file.is_file() and file.suffix in valid_suffixes:
+                    # make file absolute
+                    file = file.absolute()
+                    arcname = Path(save_prefix) / file.relative_to(root)
+                    tar.add(file, arcname=arcname)
+    return root
+
+
+def dump_args(args, save_path, extra_args=None):
+    args_dict = vars(args)
+    if extra_args:
+        assert isinstance(
+            extra_args, dict
+        ), f"extra_args should be a dictionary, got {type(extra_args)}."
+        args_dict.update(extra_args)
+    # Save to file
+    with safe_file(save_path).open("w") as f:
+        json.dump(args_dict, f, indent=4, sort_keys=True, ensure_ascii=False)
+
+
+def empty_logger():
+    logger = logging.getLogger("hymm_empty_logger")
+    logger.addHandler(logging.NullHandler())
+    logger.setLevel(logging.CRITICAL)
+    return logger
+
+
+def is_valid_experiment(path):
+    path = Path(path)
+    if path.is_dir() and path.name.split("_")[0].isdigit():
+        return True
+    return False
+
+
+def get_experiment_max_number(experiments):
+    valid_experiment_numbers = []
+    for exp in experiments:
+        if is_valid_experiment(exp):
+            valid_experiment_numbers.append(int(Path(exp).name.split("_")[0]))
+    if valid_experiment_numbers:
+        return max(valid_experiment_numbers)
+    return 0
+
+
+def safe_dir(path):
+    """
+    Create a directory (or the parent directory of a file) if it does not exist.
+
+    Args:
+        path (str or Path): Path to the directory.
+
+    Returns:
+        path (Path): Path object of the directory.
+    """
+    path = Path(path)
+    path.mkdir(exist_ok=True, parents=True)
+    return path
+
+
+def safe_file(path):
+    """
+    Create the parent directory of a file if it does not exist.
+
+    Args:
+        path (str or Path): Path to the file.
+
+    Returns:
+        path (Path): Path object of the file.
+    """
+    path = Path(path)
+    path.parent.mkdir(exist_ok=True, parents=True)
+    return path
+
+
+def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=1, fps=24):
+    """save videos by video tensor
+       copy from https://github.com/guoyww/AnimateDiff/blob/e92bd5671ba62c0d774a32951453e328018b7c5b/animatediff/utils/util.py#L61
+
+    Args:
+        videos (torch.Tensor): video tensor predicted by the model
+        path (str): path to save video
+        rescale (bool, optional): rescale the video tensor from [-1, 1] to  . Defaults to False.
+        n_rows (int, optional): Defaults to 1.
+        fps (int, optional): video save fps. Defaults to 8.
+    """
+    videos = rearrange(videos, "b c t h w -> t b c h w")
+    outputs = []
+    for x in videos:
+        x = torchvision.utils.make_grid(x, nrow=n_rows)
+        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+        if rescale:
+            x = (x + 1.0) / 2.0  # -1,1 -> 0,1
+        x = torch.clamp(x, 0, 1)
+        x = (x * 255).numpy().astype(np.uint8)
+        outputs.append(x)
+
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    imageio.mimsave(path, outputs, fps=fps)

hyvideo/utils/helpers.py

+import collections.abc
+
+from itertools import repeat
+
+import contextlib
+import os
+import random
+
+import numpy as np
+import torch
+import deepspeed
+import torch.distributed as dist
+from torch.utils.tensorboard import SummaryWriter
+
+
+def all_gather_sum(running_value, device):
+    value = torch.tensor(running_value, device=device)
+    dist.all_reduce(value, op=dist.ReduceOp.SUM)
+    return value.item()
+
+
+class EventsMonitor(object):
+    def __init__(self, events_root, rank):
+        self.rank = rank
+        if rank == 0:
+            self.writer = SummaryWriter(log_dir=events_root)
+        else:
+            self.writer = None
+
+    def write_events(self, events):
+        for event in events:
+            name, val, count = event
+            if self.rank == 0:
+                self.writer.add_scalar(name, val, global_step=count)
+
+
+def profiler_context(enable, exp_dir, worker_name):
+    if enable:
+        return torch.profiler.profile(
+            activities=[
+                torch.profiler.ProfilerActivity.CPU,
+                torch.profiler.ProfilerActivity.CUDA,
+            ],
+            schedule=torch.profiler.schedule(
+                skip_first=10,
+                wait=5,
+                warmup=1,
+                active=3,
+                repeat=2,
+            ),
+            profile_memory=True,
+            on_trace_ready=torch.profiler.tensorboard_trace_handler(
+                exp_dir, worker_name=worker_name
+            ),
+        )
+    else:
+        # return empty python context manager
+        return contextlib.nullcontext()
+
+
+def set_reproducibility(enable, global_seed=None):
+    if enable:
+        # Configure the seed for reproducibility
+        set_manual_seed(global_seed)
+    # Set following debug environment variable
+    # See the link for details: https://docs.nvidia.com/cuda/cublas/index.html#results-reproducibility
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+    # Cudnn benchmarking
+    torch.backends.cudnn.benchmark = not enable
+    # Use deterministic algorithms in PyTorch
+    torch.use_deterministic_algorithms(enable)
+
+    # LSTM and RNN networks are not deterministic
+
+
+def set_manual_seed(global_seed):
+    # Seed the RNG for Python
+    random.seed(global_seed)
+    # Seed the RNG for Numpy
+    np.random.seed(global_seed)
+    # Seed the RNG for all devices (both CPU and CUDA)
+    torch.manual_seed(global_seed)
+    # Seed cuda
+    torch.cuda.manual_seed_all(global_seed)
+
+
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+            x = tuple(x)
+            if len(x) == 1:
+                x = tuple(repeat(x[0], n))
+            return x
+        return tuple(repeat(x, n))
+
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+
+
+def as_tuple(x):
+    if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+        return tuple(x)
+    if x is None or isinstance(x, (int, float, str)):
+        return (x,)
+    else:
+        raise ValueError(f"Unknown type {type(x)}")
+
+
+def as_list_of_2tuple(x):
+    x = as_tuple(x)
+    if len(x) == 1:
+        x = (x[0], x[0])
+    assert len(x) % 2 == 0, f"Expect even length, got {len(x)}."
+    lst = []
+    for i in range(0, len(x), 2):
+        lst.append((x[i], x[i + 1]))
+    return lst

hyvideo/utils/lora_utils.py

+import torch
+from safetensors.torch import load_file
+
+
+# load kohya lora for diffusers pipeline
+def load_lora_for_pipeline(
+    pipeline,
+    lora_path,
+    LORA_PREFIX_TRANSFORMER="",
+    LORA_PREFIX_TEXT_ENCODER="",
+    alpha=1.0,
+    device=0,
+    is_parallel=False
+):
+    # load LoRA weight from .safetensors
+    if is_parallel:
+        state_dict = load_file(lora_path, device="cpu")
+        for k, v in state_dict.items():
+            state_dict[k] = v.to(device)
+    else:
+        state_dict = load_file(lora_path, device=device)
+
+    visited = []
+
+    # directly update weight in diffusers model
+    for key in state_dict:
+        # it is suggested to print out the key, it usually will be something like below
+        # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight"
+
+        # as we have set the alpha beforehand, so just skip
+        if ".alpha" in key or key in visited:
+            continue
+
+        if "text" in key:
+            layer_infos = (
+                key.split(".")[0].split(LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_")
+            )
+            curr_layer = pipeline.text_encoder
+        else:
+            layer_infos = (
+                key.split(".")[0].split(LORA_PREFIX_TRANSFORMER + "_")[-1].split("_")
+            )
+            curr_layer = pipeline.transformer
+
+        # find the target layer
+        temp_name = layer_infos.pop(0)
+        while len(layer_infos) > -1:
+            try:
+                curr_layer = curr_layer.__getattr__(temp_name)
+                if len(layer_infos) > 0:
+                    temp_name = layer_infos.pop(0)
+                elif len(layer_infos) == 0:
+                    break
+            except Exception:
+                if len(temp_name) > 0:
+                    temp_name += "_" + layer_infos.pop(0)
+                else:
+                    temp_name = layer_infos.pop(0)
+
+        pair_keys = []
+        if "lora_down" in key:
+            pair_keys.append(key.replace("lora_down", "lora_up"))
+            pair_keys.append(key)
+        else:
+            pair_keys.append(key)
+            pair_keys.append(key.replace("lora_up", "lora_down"))
+
+        # update weight
+        if len(state_dict[pair_keys[0]].shape) == 4:
+            weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32)
+            weight_down = (
+                state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32)
+            )
+            curr_layer.weight.data += alpha * torch.mm(
+                weight_up, weight_down
+            ).unsqueeze(2).unsqueeze(3)
+        else:
+            weight_up = state_dict[pair_keys[0]].to(torch.float32)
+            weight_down = state_dict[pair_keys[1]].to(torch.float32)
+            curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down)
+
+        # update visited list
+        for item in pair_keys:
+            visited.append(item)
+    del state_dict
+
+    return pipeline

hyvideo/utils/preprocess_text_encoder_tokenizer_utils.py

+import argparse
+import torch
+from transformers import (
+    AutoProcessor,
+    LlavaForConditionalGeneration,
+)
+
+
+def preprocess_text_encoder_tokenizer(args):
+
+    processor = AutoProcessor.from_pretrained(args.input_dir)
+    model = LlavaForConditionalGeneration.from_pretrained(
+        args.input_dir,
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True,
+    ).to(0)
+
+    model.language_model.save_pretrained(f"{args.output_dir}")
+    processor.tokenizer.save_pretrained(f"{args.output_dir}")
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input_dir",
+        type=str,
+        required=True,
+        help="The path to the llava-llama-3-8b-v1_1-transformers.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="",
+        help="The output path of the llava-llama-3-8b-text-encoder-tokenizer."
+        "if '', the parent dir of output will be the same as input dir.",
+    )
+    args = parser.parse_args()
+
+    if len(args.output_dir) == 0:
+        args.output_dir = "/".join(args.input_dir.split("/")[:-1])
+
+    preprocess_text_encoder_tokenizer(args)

hyvideo/utils/train_utils.py

+import random
+import torchvision.transforms as transforms
+
+import numpy as np
+import torch
+
+import imageio
+import os
+import PIL.Image
+from typing import Union, Optional, List
+from peft import get_peft_model_state_dict
+
+from hyvideo.modules.posemb_layers import get_nd_rotary_pos_embed
+from hyvideo.vae import AutoencoderKLCausal3D
+
+from pathlib import Path
+from einops import rearrange
+from PIL import Image
+
+from hyvideo.constants import PRECISION_TO_TYPE
+from safetensors.torch import load_file
+
+
+def convert_kohya_to_peft_keys(
+    kohya_dict: dict,
+    kohya_prefix="",
+    peft_prefix: str = "base_model.model",
+    device="cpu",
+) -> dict:
+    peft_dict = {}
+    for k, v in kohya_dict.items():
+        if ".alpha" in k:
+            continue
+        new_key = k.replace(f"{kohya_prefix}_lora_", f"{peft_prefix}.")
+        new_key = new_key.replace("single_blocks_", "single_blocks.")
+        new_key = new_key.replace("double_blocks_", "double_blocks.")
+        new_key = new_key.replace("_img_attn_proj", ".img_attn_proj")
+        new_key = new_key.replace("_img_attn_qkv", ".img_attn_qkv")
+        new_key = new_key.replace("_img_mlp_fc", ".img_mlp.fc")
+        new_key = new_key.replace("_txt_mlp_fc", ".txt_mlp.fc")
+        new_key = new_key.replace("_img_mod", ".img_mod")
+        new_key = new_key.replace("_txt", ".txt")
+        new_key = new_key.replace("_modulation", ".modulation")
+        new_key = new_key.replace("_linear", ".linear")
+        new_key = new_key.replace("lora_down", "lora_A.default")
+        new_key = new_key.replace("lora_up", "lora_B.default")
+        new_key = new_key.replace(
+            "_individual_token_refiner_blocks_", ".individual_token_refiner.blocks."
+        )
+        new_key = new_key.replace("_mlp_fc", ".mlp.fc")
+
+        peft_dict[new_key] = v.to(device)
+    return peft_dict
+
+
+def load_lora(model, lora_path, device):
+    kohya_weights = load_file(lora_path)
+    peft_weights = convert_kohya_to_peft_keys(
+        kohya_weights, kohya_prefix="Hunyuan_video_I2V", device=device
+    )
+    model.load_state_dict(peft_weights, strict=False)
+    return model
+
+
+def black_image(width, height):
+    black_image = Image.new("RGB", (width, height), (0, 0, 0))
+    return black_image
+
+
+def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]:
+    if images.ndim == 3:
+        images = images[None, ...]
+    images = (images * 255).round().astype("uint8")
+    if images.shape[-1] == 1:
+        # special case for grayscale (single channel) images
+        pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+    else:
+        pil_images = [Image.fromarray(image) for image in images]
+
+    return pil_images
+
+
+def get_cond_latents(args, latents, vae):
+    """get conditioned latent by decode and encode the first frame latents"""
+    first_image_latents = latents[:, :, 0, ...] if len(latents.shape) == 5 else latents
+    first_image_latents = 1 / vae.config.scaling_factor * first_image_latents
+    first_images = vae.decode(
+        first_image_latents.unsqueeze(2).to(vae.dtype), return_dict=False
+    )[0]
+    first_images = first_images.squeeze(2)
+    first_images = (first_images / 2 + 0.5).clamp(0, 1)
+    first_images = first_images.cpu().permute(0, 2, 3, 1).float().numpy()
+    first_images = numpy_to_pil(first_images)
+
+    image_transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]
+    )
+    first_images_pixel_values = [image_transform(image) for image in first_images]
+    first_images_pixel_values = (
+        torch.cat(first_images_pixel_values).unsqueeze(0).unsqueeze(2).to(vae.device)
+    )
+
+    vae_dtype = PRECISION_TO_TYPE[args.vae_precision]
+    with torch.autocast(
+        device_type="cuda", dtype=vae_dtype, enabled=vae_dtype != torch.float32
+    ):
+        cond_latents = vae.encode(
+            first_images_pixel_values
+        ).latent_dist.sample()  # B, C, F, H, W
+        cond_latents.mul_(vae.config.scaling_factor)
+
+    return cond_latents
+
+
+def get_cond_images(args, latents, vae, is_uncond=False):
+    """get conditioned images by decode the first frame latents"""
+    sematic_image_latents = (
+        latents[:, :, 0, ...] if len(latents.shape) == 5 else latents
+    )
+    sematic_image_latents = 1 / vae.config.scaling_factor * sematic_image_latents
+    semantic_images = vae.decode(
+        sematic_image_latents.unsqueeze(2).to(vae.dtype), return_dict=False
+    )[0]
+    semantic_images = semantic_images.squeeze(2)
+    semantic_images = (semantic_images / 2 + 0.5).clamp(0, 1)
+    semantic_images = semantic_images.cpu().permute(0, 2, 3, 1).float().numpy()
+    semantic_images = numpy_to_pil(semantic_images)
+    if is_uncond:
+        semantic_images = [
+            black_image(img.size[0], img.size[1]) for img in semantic_images
+        ]
+
+    return semantic_images
+
+
+def load_state_dict(args, model, logger):
+    pretrained_model_path = Path(args.model_base)
+    if not pretrained_model_path.exists():
+        raise ValueError(f"`models_root` not exists: {pretrained_model_path}")
+
+    load_key = args.load_key
+    if args.i2v_mode:
+        dit_weight = Path(args.i2v_dit_weight)
+    else:
+        dit_weight = Path(args.dit_weight)
+
+    if dit_weight is None:
+        model_dir = pretrained_model_path / f"t2v_{args.model_resolution}"
+        files = list(model_dir.glob("*.pt"))
+        if len(files) == 0:
+            raise ValueError(f"No model weights found in {model_dir}")
+        if str(files[0]).startswith("pytorch_model_"):
+            model_path = dit_weight / f"pytorch_model_{load_key}.pt"
+            bare_model = True
+        elif any(str(f).endswith("_model_states.pt") for f in files):
+            files = [f for f in files if str(f).endswith("_model_states.pt")]
+            model_path = files[0]
+            if len(files) > 1:
+                logger.warning(
+                    f"Multiple model weights found in {dit_weight}, using {model_path}"
+                )
+            bare_model = False
+        else:
+            raise ValueError(
+                f"Invalid model path: {dit_weight} with unrecognized weight format: "
+                f"{list(map(str, files))}. When given a directory as --dit-weight, only "
+                f"`pytorch_model_*.pt`(provided by HunyuanVideo official) and "
+                f"`*_model_states.pt`(saved by deepspeed) can be parsed. If you want to load a "
+                f"specific weight file, please provide the full path to the file."
+            )
+    else:
+        if dit_weight.is_dir():
+            files = list(dit_weight.glob("*.pt"))
+            if len(files) == 0:
+                raise ValueError(f"No model weights found in {dit_weight}")
+            if str(files[0]).startswith("pytorch_model_"):
+                model_path = dit_weight / f"pytorch_model_{load_key}.pt"
+                bare_model = True
+            elif any(str(f).endswith("_model_states.pt") for f in files):
+                files = [f for f in files if str(f).endswith("_model_states.pt")]
+                model_path = files[0]
+                if len(files) > 1:
+                    logger.warning(
+                        f"Multiple model weights found in {dit_weight}, using {model_path}"
+                    )
+                bare_model = False
+            else:
+                raise ValueError(
+                    f"Invalid model path: {dit_weight} with unrecognized weight format: "
+                    f"{list(map(str, files))}. When given a directory as --dit-weight, only "
+                    f"`pytorch_model_*.pt`(provided by HunyuanVideo official) and "
+                    f"`*_model_states.pt`(saved by deepspeed) can be parsed. If you want to load a "
+                    f"specific weight file, please provide the full path to the file."
+                )
+        elif dit_weight.is_file():
+            model_path = dit_weight
+            bare_model = "unknown"
+        else:
+            raise ValueError(f"Invalid model path: {dit_weight}")
+
+    if not model_path.exists():
+        raise ValueError(f"model_path not exists: {model_path}")
+    logger.info(f"Loading torch model {model_path}...")
+    state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)
+
+    if bare_model == "unknown" and ("ema" in state_dict or "module" in state_dict):
+        bare_model = False
+    if bare_model is False:
+        if load_key in state_dict:
+            state_dict = state_dict[load_key]
+        else:
+            raise KeyError(
+                f"Missing key: `{load_key}` in the checkpoint: {model_path}. The keys in the checkpoint "
+                f"are: {list(state_dict.keys())}."
+            )
+    model.load_state_dict(state_dict, strict=True)
+    return model
+
+
+class set_worker_seed_builder:
+    def __init__(self, global_rank):
+        self.global_rank = global_rank
+
+    def __call__(self, worker_id):
+        set_manual_seed(torch.initial_seed() % (2 ** 32 - 1))
+
+
+def set_reproducibility(enable, global_seed=None):
+    if enable:
+        # Configure the seed for reproducibility
+        set_manual_seed(global_seed)
+    # Set following debug environment variable
+    # See the link for details: https://docs.nvidia.com/cuda/cublas/index.html#results-reproducibility
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+    # Cudnn benchmarking
+    torch.backends.cudnn.benchmark = not enable
+    # Use deterministic algorithms in PyTorch
+    torch.use_deterministic_algorithms(enable)
+
+    # LSTM and RNN networks are not deterministic
+
+
+def prepare_model_inputs(
+    args,
+    batch: tuple,
+    device: Union[int, str],
+    model,
+    vae,
+    text_encoder,
+    text_encoder_2=None,
+    rope_theta_rescale_factor: Union[float, List[float]] = 1.0,
+    rope_interpolation_factor: Union[float, List[float]] = 1.0,
+):
+    media, latents, *batch_args = batch
+    if len(batch_args) == 3:
+        text_ids, text_mask, kwargs = batch_args
+        text_ids_2, text_mask_2 = None, None
+    elif len(batch_args) == 5:
+        text_ids, text_mask, text_ids_2, text_mask_2, kwargs = batch_args
+    else:
+        raise ValueError(f"Unexpected batch_args.")
+    data_type = kwargs["type"][0]
+
+    # Move batch to device
+    media = media.to(device)
+    latents = latents.to(device)
+    text_ids = text_ids.to(device)
+    text_mask = text_mask.to(device)
+
+    # ======================================== Encode media ======================================
+    # Used for 3D VAE with 2D inputs(image).
+    # Prepare media shape for 2D/3D VAE
+    if len(latents.shape) == 1:
+        if len(media.shape) == 4:
+            # media is a batch of image with shape [b, c, h, w]
+            if isinstance(vae, AutoencoderKLCausal3D):
+                media = media.unsqueeze(2)  # [b, c, 1, h, w]
+        elif len(media.shape) == 5:
+            # media is a batch of video with shape [b, c, f, h, w]
+            if not isinstance(vae, AutoencoderKLCausal3D):
+                media = rearrange(media, "b c f h w -> (b f) c h w")
+        else:
+            raise ValueError(
+                f"Only support media with shape (b, c, h, w) or (b, c, f, h, w), but got {media.shape}."
+            )
+
+        vae_dtype = PRECISION_TO_TYPE[args.vae_precision]
+        with torch.autocast(
+            device_type="cuda", dtype=vae_dtype, enabled=vae_dtype != torch.float32
+        ):
+            latents = vae.encode(media).latent_dist.sample()
+            if hasattr(vae.config, "shift_factor") and vae.config.shift_factor:
+                latents.sub_(vae.config.shift_factor).mul_(vae.config.scaling_factor)
+            else:
+                latents.mul_(vae.config.scaling_factor)
+    elif len(latents.shape) == 5 or len(latents.shape) == 4:  # Using video/image cache
+        latents = (
+            latents * vae.config.scaling_factor
+        )  # vae cache is not multiplied by scaling_factor
+    else:
+        raise ValueError(
+            f"Only support media/latent with shape (b, c, h, w) or (b, c, f, h, w), but got {media.shape} {latents.shape}."
+        )
+
+    cond_latents = get_cond_latents(args, latents, vae)
+    is_uncond = (
+        torch.tensor(1).to(torch.int64)
+        if random.random() < args.sematic_cond_drop_p
+        else torch.tensor(0).to(torch.int64)
+    )
+    semantic_images = get_cond_images(args, latents, vae, is_uncond=is_uncond)
+
+    # ======================================== Encode text ======================================
+    # Autocast is handled by text_encoder itself.
+    # Whether to apply text_mask is determined by args.use_attention_mask.
+    text_outputs = text_encoder.encode(
+        {"input_ids": text_ids, "attention_mask": text_mask},
+        data_type=batch_args[-1]["type"][0],
+        semantic_images=semantic_images,
+    )
+    text_states = text_outputs.hidden_state
+    text_mask = text_outputs.attention_mask
+    text_states_2 = (
+        text_encoder_2.encode(
+            {"input_ids": text_ids_2, "attention_mask": text_mask_2},
+            data_type=data_type,
+        ).hidden_state
+        if text_encoder_2 is not None
+        else None
+    )
+
+    # ======================================== Build RoPE ======================================
+    target_ndim = 3  # n-d RoPE
+    ndim = len(latents.shape) - 2
+    latents_size = list(latents.shape[-ndim:])
+    freqs_cos, freqs_sin = get_rope_freq_from_size(
+        args,
+        model,
+        latents_size,
+        ndim,
+        target_ndim,
+        rope_theta_rescale_factor=rope_theta_rescale_factor,
+        rope_interpolation_factor=rope_interpolation_factor,
+    )
+
+    # ===================================== Pack model kwargs ==================================
+    model_kwargs = dict(
+        text_states=text_states,  # [b, 256, 4096]
+        text_mask=text_mask,  # [b, 256]
+        text_states_2=text_states_2,  # [b, 768]
+        freqs_cos=freqs_cos,  # [seqlen, head_dim]
+        freqs_sin=freqs_sin,  # [seqlen, head_dim]
+        return_dict=True,
+    )
+
+    return latents, model_kwargs, freqs_cos.shape[0], cond_latents
+
+
+def format_params(params):
+    if params < 1e6:
+        return f"{params} (less than 1M)"
+    elif params < 1e9:
+        return f"{params / 1e6:.2f}M"
+    else:
+        return f"{params / 1e9:.2f}B"
+
+
+def set_manual_seed(global_seed):
+    random.seed(global_seed)
+    np.random.seed(global_seed)
+    torch.manual_seed(global_seed)
+
+
+def get_rope_freq_from_size(
+    args,
+    model,
+    latents_size,
+    ndim,
+    target_ndim,
+    rope_theta_rescale_factor=1.0,
+    rope_interpolation_factor=1.0,
+):
+
+    if isinstance(model.patch_size, int):
+        assert all(s % model.patch_size == 0 for s in latents_size), (
+            f"Latent size(last {ndim} dimensions) should be divisible by patch size({model.patch_size}), "
+            f"but got {latents_size}."
+        )
+        rope_sizes = [s // model.patch_size for s in latents_size]
+
+    elif isinstance(model.patch_size, list):
+        assert all(
+            s % model.patch_size[idx] == 0 for idx, s in enumerate(latents_size)
+        ), (
+            f"Latent size(last {ndim} dimensions) should be divisible by patch size({model.patch_size}), "
+            f"but got {latents_size}."
+        )
+        rope_sizes = [s // model.patch_size[idx] for idx, s in enumerate(latents_size)]
+
+    if len(rope_sizes) != target_ndim:
+        rope_sizes = [1] * (target_ndim - len(rope_sizes)) + rope_sizes  # time axis
+    head_dim = model.hidden_size // model.heads_num
+    rope_dim_list = model.rope_dim_list
+
+    if rope_dim_list is None:
+        rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)]
+    assert (
+        sum(rope_dim_list) == head_dim
+    ), "sum(rope_dim_list) should equal to head_dim of attention layer"
+
+    freqs_cos, freqs_sin = get_nd_rotary_pos_embed(
+        rope_dim_list,
+        rope_sizes,
+        theta=args.rope_theta,
+        use_real=True,
+        theta_rescale_factor=rope_theta_rescale_factor,
+        interpolation_factor=rope_interpolation_factor,
+    )
+
+    return freqs_cos, freqs_sin
+
+
+# copy from https://github.com/huggingface/diffusers/blob/ec9bfa9e148b7764137dd92247ce859d915abcb0/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py#L258
+# get kohya lora state dict
+def get_module_kohya_state_dict(module, prefix, dtype, adapter_name="default"):
+    kohya_ss_state_dict = {}
+    for peft_key, weight in get_peft_model_state_dict(
+        module, adapter_name=adapter_name
+    ).items():
+        kohya_key = peft_key.replace("base_model.model", prefix)
+        kohya_key = kohya_key.replace("lora_A", "lora_down")
+        kohya_key = kohya_key.replace("lora_B", "lora_up")
+        kohya_key = kohya_key.replace(".", "_", kohya_key.count(".") - 2)
+        kohya_ss_state_dict[kohya_key] = weight.to(dtype)
+
+        # Set alpha parameter
+        if "lora_down" in kohya_key:
+            alpha_key = f'{kohya_key.split(".")[0]}.alpha'
+            kohya_ss_state_dict[alpha_key] = torch.tensor(
+                module.peft_config[adapter_name].lora_alpha
+            ).to(dtype)
+
+    return kohya_ss_state_dict
+
+
+# get diffusers lora state dict
+def get_module_diffusers_state_dict(module, dtype, adapter_name="default"):
+    diffusers_ss_state_dict = {}
+    for peft_key, weight in get_peft_model_state_dict(
+        module, adapter_name=adapter_name
+    ).items():
+        diffusers_key = peft_key.replace("base_model.model", "diffusion_model")
+        diffusers_ss_state_dict[diffusers_key] = weight.to(dtype)
+
+    return diffusers_ss_state_dict

hyvideo/vae/__init__.py

+from pathlib import Path
+
+import torch
+
+from .autoencoder_kl_causal_3d import AutoencoderKLCausal3D
+from ..constants import VAE_PATH, PRECISION_TO_TYPE
+
+def load_vae(vae_type: str="884-16c-hy",
+             vae_precision: str=None,
+             sample_size: tuple=None,
+             vae_path: str=None,
+             logger=None,
+             device=None
+             ):
+    """the fucntion to load the 3D VAE model
+
+    Args:
+        vae_type (str): the type of the 3D VAE model. Defaults to "884-16c-hy".
+        vae_precision (str, optional): the precision to load vae. Defaults to None.
+        sample_size (tuple, optional): the tiling size. Defaults to None.
+        vae_path (str, optional): the path to vae. Defaults to None.
+        logger (_type_, optional): logger. Defaults to None.
+        device (_type_, optional): device to load vae. Defaults to None.
+    """
+    if vae_path is None:
+        vae_path = VAE_PATH[vae_type]
+    
+    if logger is not None:
+        logger.info(f"Loading 3D VAE model ({vae_type}) from: {vae_path}")
+    config = AutoencoderKLCausal3D.load_config(vae_path)
+    if sample_size:
+        vae = AutoencoderKLCausal3D.from_config(config, sample_size=sample_size)
+    else:
+        vae = AutoencoderKLCausal3D.from_config(config)
+    
+    vae_ckpt = Path(vae_path) / "pytorch_model.pt"
+    assert vae_ckpt.exists(), f"VAE checkpoint not found: {vae_ckpt}"
+    
+    ckpt = torch.load(vae_ckpt, map_location=vae.device)
+    if "state_dict" in ckpt:
+        ckpt = ckpt["state_dict"]
+    if any(k.startswith("vae.") for k in ckpt.keys()):
+        ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items() if k.startswith("vae.")}
+    vae.load_state_dict(ckpt)
+
+    spatial_compression_ratio = vae.config.spatial_compression_ratio
+    time_compression_ratio = vae.config.time_compression_ratio
+    
+    if vae_precision is not None:
+        vae = vae.to(dtype=PRECISION_TO_TYPE[vae_precision])
+
+    vae.requires_grad_(False)
+
+    if logger is not None:
+        logger.info(f"VAE to dtype: {vae.dtype}")
+
+    if device is not None:
+        vae = vae.to(device)
+
+    vae.eval()
+
+    return vae, vae_path, spatial_compression_ratio, time_compression_ratio

hyvideo/vae/autoencoder_kl_causal_3d.py

+# Copyright 2024 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.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+from typing import Dict, Optional, Tuple, Union
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+
+try:
+    # This diffusers is modified and packed in the mirror.
+    from diffusers.loaders import FromOriginalVAEMixin
+except ImportError:
+    # Use this to be compatible with the original diffusers.
+    from diffusers.loaders.single_file_model import FromOriginalModelMixin as FromOriginalVAEMixin
+from diffusers.utils.accelerate_utils import apply_forward_hook
+from diffusers.models.attention_processor import (
+    ADDED_KV_ATTENTION_PROCESSORS,
+    CROSS_ATTENTION_PROCESSORS,
+    Attention,
+    AttentionProcessor,
+    AttnAddedKVProcessor,
+    AttnProcessor,
+)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from diffusers.models.modeling_utils import ModelMixin
+from .vae import DecoderCausal3D, BaseOutput, DecoderOutput, DiagonalGaussianDistribution, EncoderCausal3D
+
+
+@dataclass
+class DecoderOutput2(BaseOutput):
+    sample: torch.FloatTensor
+    posterior: Optional[DiagonalGaussianDistribution] = None
+
+
+class AutoencoderKLCausal3D(ModelMixin, ConfigMixin, FromOriginalVAEMixin):
+    r"""
+    A VAE model with KL loss for encoding images/videos into latents and decoding latent representations into images/videos.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        down_block_types: Tuple[str] = ("DownEncoderBlockCausal3D",),
+        up_block_types: Tuple[str] = ("UpDecoderBlockCausal3D",),
+        block_out_channels: Tuple[int] = (64,),
+        layers_per_block: int = 1,
+        act_fn: str = "silu",
+        latent_channels: int = 4,
+        norm_num_groups: int = 32,
+        sample_size: int = 32,
+        sample_tsize: int = 64,
+        scaling_factor: float = 0.18215,
+        force_upcast: float = True,
+        spatial_compression_ratio: int = 8,
+        time_compression_ratio: int = 4,
+        mid_block_add_attention: bool = True,
+    ):
+        super().__init__()
+
+        self.time_compression_ratio = time_compression_ratio
+
+        self.encoder = EncoderCausal3D(
+            in_channels=in_channels,
+            out_channels=latent_channels,
+            down_block_types=down_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            act_fn=act_fn,
+            norm_num_groups=norm_num_groups,
+            double_z=True,
+            time_compression_ratio=time_compression_ratio,
+            spatial_compression_ratio=spatial_compression_ratio,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        self.decoder = DecoderCausal3D(
+            in_channels=latent_channels,
+            out_channels=out_channels,
+            up_block_types=up_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            norm_num_groups=norm_num_groups,
+            act_fn=act_fn,
+            time_compression_ratio=time_compression_ratio,
+            spatial_compression_ratio=spatial_compression_ratio,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1)
+        self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1)
+
+        self.use_slicing = False
+        self.use_spatial_tiling = False
+        self.use_temporal_tiling = False
+
+        # only relevant if vae tiling is enabled
+        self.tile_sample_min_tsize = sample_tsize
+        self.tile_latent_min_tsize = sample_tsize // time_compression_ratio
+
+        self.tile_sample_min_size = self.config.sample_size
+        sample_size = (
+            self.config.sample_size[0]
+            if isinstance(self.config.sample_size, (list, tuple))
+            else self.config.sample_size
+        )
+        self.tile_latent_min_size = int(sample_size / (2 ** (len(self.config.block_out_channels) - 1)))
+        self.tile_overlap_factor = 0.25
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, (EncoderCausal3D, DecoderCausal3D)):
+            module.gradient_checkpointing = value
+
+    def enable_temporal_tiling(self, use_tiling: bool = True):
+        self.use_temporal_tiling = use_tiling
+
+    def disable_temporal_tiling(self):
+        self.enable_temporal_tiling(False)
+
+    def enable_spatial_tiling(self, use_tiling: bool = True):
+        self.use_spatial_tiling = use_tiling
+
+    def disable_spatial_tiling(self):
+        self.enable_spatial_tiling(False)
+
+    def enable_tiling(self, use_tiling: bool = True):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger videos.
+        """
+        self.enable_spatial_tiling(use_tiling)
+        self.enable_temporal_tiling(use_tiling)
+
+    def disable_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.disable_spatial_tiling()
+        self.disable_temporal_tiling()
+
+    def enable_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.use_slicing = True
+
+    def disable_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.use_slicing = False
+
+    @property
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(
+        self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False
+    ):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor, _remove_lora=_remove_lora)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora)
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+            processor = AttnAddedKVProcessor()
+        elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor, _remove_lora=True)
+
+    @apply_forward_hook
+    def encode(
+        self, x: torch.FloatTensor, return_dict: bool = True
+    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+        """
+        Encode a batch of images/videos into latents.
+
+        Args:
+            x (`torch.FloatTensor`): Input batch of images/videos.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+                The latent representations of the encoded images/videos. If `return_dict` is True, a
+                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+        """
+        assert len(x.shape) == 5, "The input tensor should have 5 dimensions."
+
+        if self.use_temporal_tiling and x.shape[2] > self.tile_sample_min_tsize:
+            return self.temporal_tiled_encode(x, return_dict=return_dict)
+
+        if self.use_spatial_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
+            return self.spatial_tiled_encode(x, return_dict=return_dict)
+
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self.encoder(x)
+
+        moments = self.quant_conv(h)
+        posterior = DiagonalGaussianDistribution(moments)
+
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
+        assert len(z.shape) == 5, "The input tensor should have 5 dimensions."
+
+        if self.use_temporal_tiling and z.shape[2] > self.tile_latent_min_tsize:
+            return self.temporal_tiled_decode(z, return_dict=return_dict)
+
+        if self.use_spatial_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
+            return self.spatial_tiled_decode(z, return_dict=return_dict)
+
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    @apply_forward_hook
+    def decode(
+        self, z: torch.FloatTensor, return_dict: bool = True, generator=None
+    ) -> Union[DecoderOutput, torch.FloatTensor]:
+        """
+        Decode a batch of images/videos.
+
+        Args:
+            z (`torch.FloatTensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+
+        """
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
+        if not return_dict:
+            return (decoded,)
+
+        return DecoderOutput(sample=decoded)
+
+    def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (y / blend_extent)
+        return b
+
+    def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (x / blend_extent)
+        return b
+
+    def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-3], b.shape[-3], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * (x / blend_extent)
+        return b
+
+    def spatial_tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True, return_moments: bool = False) -> AutoencoderKLOutput:
+        r"""Encode a batch of images/videos using a tiled encoder.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+        steps. This is useful to keep memory use constant regardless of image/videos size. The end result of tiled encoding is
+        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+        output, but they should be much less noticeable.
+
+        Args:
+            x (`torch.FloatTensor`): Input batch of images/videos.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`:
+                If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain
+                `tuple` is returned.
+        """
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split video into tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[-2], overlap_size):
+            row = []
+            for j in range(0, x.shape[-1], overlap_size):
+                tile = x[:, :, :, i: i + self.tile_sample_min_size, j: j + self.tile_sample_min_size]
+                tile = self.encoder(tile)
+                tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        moments = torch.cat(result_rows, dim=-2)
+        if return_moments:
+            return moments
+
+        posterior = DiagonalGaussianDistribution(moments)
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def spatial_tiled_decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
+        r"""
+        Decode a batch of images/videos using a tiled decoder.
+
+        Args:
+            z (`torch.FloatTensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+        """
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_sample_min_size - blend_extent
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, z.shape[-2], overlap_size):
+            row = []
+            for j in range(0, z.shape[-1], overlap_size):
+                tile = z[:, :, :, i: i + self.tile_latent_min_size, j: j + self.tile_latent_min_size]
+                tile = self.post_quant_conv(tile)
+                decoded = self.decoder(tile)
+                row.append(decoded)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        dec = torch.cat(result_rows, dim=-2)
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def temporal_tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput:
+
+        B, C, T, H, W = x.shape
+        overlap_size = int(self.tile_sample_min_tsize * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_tsize * self.tile_overlap_factor)
+        t_limit = self.tile_latent_min_tsize - blend_extent
+
+        # Split the video into tiles and encode them separately.
+        row = []
+        for i in range(0, T, overlap_size):
+            tile = x[:, :, i: i + self.tile_sample_min_tsize + 1, :, :]
+            if self.use_spatial_tiling and (tile.shape[-1] > self.tile_sample_min_size or tile.shape[-2] > self.tile_sample_min_size):
+                tile = self.spatial_tiled_encode(tile, return_moments=True)
+            else:
+                tile = self.encoder(tile)
+                tile = self.quant_conv(tile)
+            if i > 0:
+                tile = tile[:, :, 1:, :, :]
+            row.append(tile)
+        result_row = []
+        for i, tile in enumerate(row):
+            if i > 0:
+                tile = self.blend_t(row[i - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :t_limit, :, :])
+            else:
+                result_row.append(tile[:, :, :t_limit + 1, :, :])
+
+        moments = torch.cat(result_row, dim=2)
+        posterior = DiagonalGaussianDistribution(moments)
+
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def temporal_tiled_decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
+        # Split z into overlapping tiles and decode them separately.
+
+        B, C, T, H, W = z.shape
+        overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor)
+        t_limit = self.tile_sample_min_tsize - blend_extent
+
+        row = []
+        for i in range(0, T, overlap_size):
+            tile = z[:, :, i: i + self.tile_latent_min_tsize + 1, :, :]
+            if self.use_spatial_tiling and (tile.shape[-1] > self.tile_latent_min_size or tile.shape[-2] > self.tile_latent_min_size):
+                decoded = self.spatial_tiled_decode(tile, return_dict=True).sample
+            else:
+                tile = self.post_quant_conv(tile)
+                decoded = self.decoder(tile)
+            if i > 0:
+                decoded = decoded[:, :, 1:, :, :]
+            row.append(decoded)
+        result_row = []
+        for i, tile in enumerate(row):
+            if i > 0:
+                tile = self.blend_t(row[i - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :t_limit, :, :])
+            else:
+                result_row.append(tile[:, :, :t_limit + 1, :, :])
+
+        dec = torch.cat(result_row, dim=2)
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        sample_posterior: bool = False,
+        return_dict: bool = True,
+        return_posterior: bool = False,
+        generator: Optional[torch.Generator] = None,
+    ) -> Union[DecoderOutput2, torch.FloatTensor]:
+        r"""
+        Args:
+            sample (`torch.FloatTensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z).sample
+
+        if not return_dict:
+            if return_posterior:
+                return (dec, posterior)
+            else:
+                return (dec,)
+        if return_posterior:
+            return DecoderOutput2(sample=dec, posterior=posterior)
+        else:
+            return DecoderOutput2(sample=dec)
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query,
+        key, value) are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
+    def unfuse_qkv_projections(self):
+        """Disables the fused QKV projection if enabled.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+
+        """
+        if self.original_attn_processors is not None:
+            self.set_attn_processor(self.original_attn_processors)

hyvideo/vae/unet_causal_3d_blocks.py

+# Copyright 2024 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.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from einops import rearrange
+
+from diffusers.utils import logging
+from diffusers.models.activations import get_activation
+from diffusers.models.attention_processor import SpatialNorm
+from diffusers.models.attention_processor import Attention
+from diffusers.models.normalization import AdaGroupNorm
+from diffusers.models.normalization import RMSNorm
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: int = None):
+    seq_len = n_frame * n_hw
+    mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device)
+    for i in range(seq_len):
+        i_frame = i // n_hw
+        mask[i, : (i_frame + 1) * n_hw] = 0
+    if batch_size is not None:
+        mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
+    return mask
+
+
+class CausalConv3d(nn.Module):
+    """
+    Implements a causal 3D convolution layer where each position only depends on previous timesteps and current spatial locations.
+    This maintains temporal causality in video generation tasks.
+    """
+
+    def __init__(
+        self,
+        chan_in,
+        chan_out,
+        kernel_size: Union[int, Tuple[int, int, int]],
+        stride: Union[int, Tuple[int, int, int]] = 1,
+        dilation: Union[int, Tuple[int, int, int]] = 1,
+        pad_mode='replicate',
+        **kwargs
+    ):
+        super().__init__()
+
+        self.pad_mode = pad_mode
+        padding = (kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size - 1, 0)  # W, H, T
+        self.time_causal_padding = padding
+
+        self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs)
+
+    def forward(self, x):
+        x = F.pad(x, self.time_causal_padding, mode=self.pad_mode)
+        return self.conv(x)
+
+
+class UpsampleCausal3D(nn.Module):
+    """
+    A 3D upsampling layer with an optional convolution.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        use_conv: bool = False,
+        use_conv_transpose: bool = False,
+        out_channels: Optional[int] = None,
+        name: str = "conv",
+        kernel_size: Optional[int] = None,
+        padding=1,
+        norm_type=None,
+        eps=None,
+        elementwise_affine=None,
+        bias=True,
+        interpolate=True,
+        upsample_factor=(2, 2, 2),
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_conv_transpose = use_conv_transpose
+        self.name = name
+        self.interpolate = interpolate
+        self.upsample_factor = upsample_factor
+
+        if norm_type == "ln_norm":
+            self.norm = nn.LayerNorm(channels, eps, elementwise_affine)
+        elif norm_type == "rms_norm":
+            self.norm = RMSNorm(channels, eps, elementwise_affine)
+        elif norm_type is None:
+            self.norm = None
+        else:
+            raise ValueError(f"unknown norm_type: {norm_type}")
+
+        conv = None
+        if use_conv_transpose:
+            raise NotImplementedError
+        elif use_conv:
+            if kernel_size is None:
+                kernel_size = 3
+            conv = CausalConv3d(self.channels, self.out_channels, kernel_size=kernel_size, bias=bias)
+
+        if name == "conv":
+            self.conv = conv
+        else:
+            self.Conv2d_0 = conv
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        output_size: Optional[int] = None,
+        scale: float = 1.0,
+    ) -> torch.FloatTensor:
+        assert hidden_states.shape[1] == self.channels
+
+        if self.norm is not None:
+            raise NotImplementedError
+
+        if self.use_conv_transpose:
+            return self.conv(hidden_states)
+
+        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
+        dtype = hidden_states.dtype
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(torch.float32)
+
+        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+        if hidden_states.shape[0] >= 64:
+            hidden_states = hidden_states.contiguous()
+
+        # if `output_size` is passed we force the interpolation output
+        # size and do not make use of `scale_factor=2`
+        if self.interpolate:
+            B, C, T, H, W = hidden_states.shape
+            first_h, other_h = hidden_states.split((1, T - 1), dim=2)
+            if output_size is None:
+                if T > 1:
+                    other_h = F.interpolate(other_h, scale_factor=self.upsample_factor, mode="nearest")
+
+                first_h = first_h.squeeze(2)
+                first_h = F.interpolate(first_h, scale_factor=self.upsample_factor[1:], mode="nearest")
+                first_h = first_h.unsqueeze(2)
+            else:
+                raise NotImplementedError
+
+            if T > 1:
+                hidden_states = torch.cat((first_h, other_h), dim=2)
+            else:
+                hidden_states = first_h
+
+        # If the input is bfloat16, we cast back to bfloat16
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(dtype)
+
+        if self.use_conv:
+            if self.name == "conv":
+                hidden_states = self.conv(hidden_states)
+            else:
+                hidden_states = self.Conv2d_0(hidden_states)
+
+        return hidden_states
+
+
+class DownsampleCausal3D(nn.Module):
+    """
+    A 3D downsampling layer with an optional convolution.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        use_conv: bool = False,
+        out_channels: Optional[int] = None,
+        padding: int = 1,
+        name: str = "conv",
+        kernel_size=3,
+        norm_type=None,
+        eps=None,
+        elementwise_affine=None,
+        bias=True,
+        stride=2,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.padding = padding
+        stride = stride
+        self.name = name
+
+        if norm_type == "ln_norm":
+            self.norm = nn.LayerNorm(channels, eps, elementwise_affine)
+        elif norm_type == "rms_norm":
+            self.norm = RMSNorm(channels, eps, elementwise_affine)
+        elif norm_type is None:
+            self.norm = None
+        else:
+            raise ValueError(f"unknown norm_type: {norm_type}")
+
+        if use_conv:
+            conv = CausalConv3d(
+                self.channels, self.out_channels, kernel_size=kernel_size, stride=stride, bias=bias
+            )
+        else:
+            raise NotImplementedError
+
+        if name == "conv":
+            self.Conv2d_0 = conv
+            self.conv = conv
+        elif name == "Conv2d_0":
+            self.conv = conv
+        else:
+            self.conv = conv
+
+    def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
+        assert hidden_states.shape[1] == self.channels
+
+        if self.norm is not None:
+            hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+        assert hidden_states.shape[1] == self.channels
+
+        hidden_states = self.conv(hidden_states)
+
+        return hidden_states
+
+
+class ResnetBlockCausal3D(nn.Module):
+    r"""
+    A Resnet block.
+    """
+
+    def __init__(
+        self,
+        *,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        conv_shortcut: bool = False,
+        dropout: float = 0.0,
+        temb_channels: int = 512,
+        groups: int = 32,
+        groups_out: Optional[int] = None,
+        pre_norm: bool = True,
+        eps: float = 1e-6,
+        non_linearity: str = "swish",
+        skip_time_act: bool = False,
+        # default, scale_shift, ada_group, spatial
+        time_embedding_norm: str = "default",
+        kernel: Optional[torch.FloatTensor] = None,
+        output_scale_factor: float = 1.0,
+        use_in_shortcut: Optional[bool] = None,
+        up: bool = False,
+        down: bool = False,
+        conv_shortcut_bias: bool = True,
+        conv_3d_out_channels: Optional[int] = None,
+    ):
+        super().__init__()
+        self.pre_norm = pre_norm
+        self.pre_norm = True
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+        self.up = up
+        self.down = down
+        self.output_scale_factor = output_scale_factor
+        self.time_embedding_norm = time_embedding_norm
+        self.skip_time_act = skip_time_act
+
+        linear_cls = nn.Linear
+
+        if groups_out is None:
+            groups_out = groups
+
+        if self.time_embedding_norm == "ada_group":
+            self.norm1 = AdaGroupNorm(temb_channels, in_channels, groups, eps=eps)
+        elif self.time_embedding_norm == "spatial":
+            self.norm1 = SpatialNorm(in_channels, temb_channels)
+        else:
+            self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+
+        self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3, stride=1)
+
+        if temb_channels is not None:
+            if self.time_embedding_norm == "default":
+                self.time_emb_proj = linear_cls(temb_channels, out_channels)
+            elif self.time_embedding_norm == "scale_shift":
+                self.time_emb_proj = linear_cls(temb_channels, 2 * out_channels)
+            elif self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":
+                self.time_emb_proj = None
+            else:
+                raise ValueError(f"Unknown time_embedding_norm : {self.time_embedding_norm} ")
+        else:
+            self.time_emb_proj = None
+
+        if self.time_embedding_norm == "ada_group":
+            self.norm2 = AdaGroupNorm(temb_channels, out_channels, groups_out, eps=eps)
+        elif self.time_embedding_norm == "spatial":
+            self.norm2 = SpatialNorm(out_channels, temb_channels)
+        else:
+            self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
+
+        self.dropout = torch.nn.Dropout(dropout)
+        conv_3d_out_channels = conv_3d_out_channels or out_channels
+        self.conv2 = CausalConv3d(out_channels, conv_3d_out_channels, kernel_size=3, stride=1)
+
+        self.nonlinearity = get_activation(non_linearity)
+
+        self.upsample = self.downsample = None
+        if self.up:
+            self.upsample = UpsampleCausal3D(in_channels, use_conv=False)
+        elif self.down:
+            self.downsample = DownsampleCausal3D(in_channels, use_conv=False, name="op")
+
+        self.use_in_shortcut = self.in_channels != conv_3d_out_channels if use_in_shortcut is None else use_in_shortcut
+
+        self.conv_shortcut = None
+        if self.use_in_shortcut:
+            self.conv_shortcut = CausalConv3d(
+                in_channels,
+                conv_3d_out_channels,
+                kernel_size=1,
+                stride=1,
+                bias=conv_shortcut_bias,
+            )
+
+    def forward(
+        self,
+        input_tensor: torch.FloatTensor,
+        temb: torch.FloatTensor,
+        scale: float = 1.0,
+    ) -> torch.FloatTensor:
+        hidden_states = input_tensor
+
+        if self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":
+            hidden_states = self.norm1(hidden_states, temb)
+        else:
+            hidden_states = self.norm1(hidden_states)
+
+        hidden_states = self.nonlinearity(hidden_states)
+
+        if self.upsample is not None:
+            # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+            if hidden_states.shape[0] >= 64:
+                input_tensor = input_tensor.contiguous()
+                hidden_states = hidden_states.contiguous()
+            input_tensor = (
+                self.upsample(input_tensor, scale=scale)
+            )
+            hidden_states = (
+                self.upsample(hidden_states, scale=scale)
+            )
+        elif self.downsample is not None:
+            input_tensor = (
+                self.downsample(input_tensor, scale=scale)
+            )
+            hidden_states = (
+                self.downsample(hidden_states, scale=scale)
+            )
+
+        hidden_states = self.conv1(hidden_states)
+
+        if self.time_emb_proj is not None:
+            if not self.skip_time_act:
+                temb = self.nonlinearity(temb)
+            temb = (
+                self.time_emb_proj(temb, scale)[:, :, None, None]
+            )
+
+        if temb is not None and self.time_embedding_norm == "default":
+            hidden_states = hidden_states + temb
+
+        if self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":
+            hidden_states = self.norm2(hidden_states, temb)
+        else:
+            hidden_states = self.norm2(hidden_states)
+
+        if temb is not None and self.time_embedding_norm == "scale_shift":
+            scale, shift = torch.chunk(temb, 2, dim=1)
+            hidden_states = hidden_states * (1 + scale) + shift
+
+        hidden_states = self.nonlinearity(hidden_states)
+
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        if self.conv_shortcut is not None:
+            input_tensor = (
+                self.conv_shortcut(input_tensor)
+            )
+
+        output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
+
+        return output_tensor
+
+
+def get_down_block3d(
+    down_block_type: str,
+    num_layers: int,
+    in_channels: int,
+    out_channels: int,
+    temb_channels: int,
+    add_downsample: bool,
+    downsample_stride: int,
+    resnet_eps: float,
+    resnet_act_fn: str,
+    transformer_layers_per_block: int = 1,
+    num_attention_heads: Optional[int] = None,
+    resnet_groups: Optional[int] = None,
+    cross_attention_dim: Optional[int] = None,
+    downsample_padding: Optional[int] = None,
+    dual_cross_attention: bool = False,
+    use_linear_projection: bool = False,
+    only_cross_attention: bool = False,
+    upcast_attention: bool = False,
+    resnet_time_scale_shift: str = "default",
+    attention_type: str = "default",
+    resnet_skip_time_act: bool = False,
+    resnet_out_scale_factor: float = 1.0,
+    cross_attention_norm: Optional[str] = None,
+    attention_head_dim: Optional[int] = None,
+    downsample_type: Optional[str] = None,
+    dropout: float = 0.0,
+):
+    # If attn head dim is not defined, we default it to the number of heads
+    if attention_head_dim is None:
+        logger.warn(
+            f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
+        )
+        attention_head_dim = num_attention_heads
+
+    down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
+    if down_block_type == "DownEncoderBlockCausal3D":
+        return DownEncoderBlockCausal3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            dropout=dropout,
+            add_downsample=add_downsample,
+            downsample_stride=downsample_stride,
+            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,
+        )
+    raise ValueError(f"{down_block_type} does not exist.")
+
+
+def get_up_block3d(
+    up_block_type: str,
+    num_layers: int,
+    in_channels: int,
+    out_channels: int,
+    prev_output_channel: int,
+    temb_channels: int,
+    add_upsample: bool,
+    upsample_scale_factor: Tuple,
+    resnet_eps: float,
+    resnet_act_fn: str,
+    resolution_idx: Optional[int] = None,
+    transformer_layers_per_block: int = 1,
+    num_attention_heads: Optional[int] = None,
+    resnet_groups: Optional[int] = None,
+    cross_attention_dim: Optional[int] = None,
+    dual_cross_attention: bool = False,
+    use_linear_projection: bool = False,
+    only_cross_attention: bool = False,
+    upcast_attention: bool = False,
+    resnet_time_scale_shift: str = "default",
+    attention_type: str = "default",
+    resnet_skip_time_act: bool = False,
+    resnet_out_scale_factor: float = 1.0,
+    cross_attention_norm: Optional[str] = None,
+    attention_head_dim: Optional[int] = None,
+    upsample_type: Optional[str] = None,
+    dropout: float = 0.0,
+) -> nn.Module:
+    # If attn head dim is not defined, we default it to the number of heads
+    if attention_head_dim is None:
+        logger.warn(
+            f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
+        )
+        attention_head_dim = num_attention_heads
+
+    up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
+    if up_block_type == "UpDecoderBlockCausal3D":
+        return UpDecoderBlockCausal3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            resolution_idx=resolution_idx,
+            dropout=dropout,
+            add_upsample=add_upsample,
+            upsample_scale_factor=upsample_scale_factor,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            temb_channels=temb_channels,
+        )
+    raise ValueError(f"{up_block_type} does not exist.")
+
+
+class UNetMidBlockCausal3D(nn.Module):
+    """
+    A 3D UNet mid-block [`UNetMidBlockCausal3D`] with multiple residual blocks and optional attention blocks.
+    """
+
+    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",  # default, spatial
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        attn_groups: Optional[int] = None,
+        resnet_pre_norm: bool = True,
+        add_attention: bool = True,
+        attention_head_dim: int = 1,
+        output_scale_factor: float = 1.0,
+    ):
+        super().__init__()
+        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+        self.add_attention = add_attention
+
+        if attn_groups is None:
+            attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None
+
+        # there is always at least one resnet
+        resnets = [
+            ResnetBlockCausal3D(
+                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,
+            )
+        ]
+        attentions = []
+
+        if attention_head_dim is None:
+            logger.warn(
+                f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}."
+            )
+            attention_head_dim = in_channels
+
+        for _ in range(num_layers):
+            if self.add_attention:
+                attentions.append(
+                    Attention(
+                        in_channels,
+                        heads=in_channels // attention_head_dim,
+                        dim_head=attention_head_dim,
+                        rescale_output_factor=output_scale_factor,
+                        eps=resnet_eps,
+                        norm_num_groups=attn_groups,
+                        spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,
+                        residual_connection=True,
+                        bias=True,
+                        upcast_softmax=True,
+                        _from_deprecated_attn_block=True,
+                    )
+                )
+            else:
+                attentions.append(None)
+
+            resnets.append(
+                ResnetBlockCausal3D(
+                    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,
+                )
+            )
+
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+    def forward(self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None) -> torch.FloatTensor:
+        hidden_states = self.resnets[0](hidden_states, temb)
+        for attn, resnet in zip(self.attentions, self.resnets[1:]):
+            if attn is not None:
+                B, C, T, H, W = hidden_states.shape
+                hidden_states = rearrange(hidden_states, "b c f h w -> b (f h w) c")
+                attention_mask = prepare_causal_attention_mask(
+                    T, H * W, hidden_states.dtype, hidden_states.device, batch_size=B
+                )
+                hidden_states = attn(hidden_states, temb=temb, attention_mask=attention_mask)
+                hidden_states = rearrange(hidden_states, "b (f h w) c -> b c f h w", f=T, h=H, w=W)
+            hidden_states = resnet(hidden_states, temb)
+
+        return hidden_states
+
+
+class DownEncoderBlockCausal3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_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: float = 1.0,
+        add_downsample: bool = True,
+        downsample_stride: int = 2,
+        downsample_padding: int = 1,
+    ):
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlockCausal3D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    temb_channels=None,
+                    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,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    DownsampleCausal3D(
+                        out_channels,
+                        use_conv=True,
+                        out_channels=out_channels,
+                        padding=downsample_padding,
+                        name="op",
+                        stride=downsample_stride,
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+    def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
+        for resnet in self.resnets:
+            hidden_states = resnet(hidden_states, temb=None, scale=scale)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states, scale)
+
+        return hidden_states
+
+
+class UpDecoderBlockCausal3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        resolution_idx: Optional[int] = None,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",  # default, spatial
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        output_scale_factor: float = 1.0,
+        add_upsample: bool = True,
+        upsample_scale_factor=(2, 2, 2),
+        temb_channels: Optional[int] = None,
+    ):
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            input_channels = in_channels if i == 0 else out_channels
+
+            resnets.append(
+                ResnetBlockCausal3D(
+                    in_channels=input_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,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList(
+                [
+                    UpsampleCausal3D(
+                        out_channels,
+                        use_conv=True,
+                        out_channels=out_channels,
+                        upsample_factor=upsample_scale_factor,
+                    )
+                ]
+            )
+        else:
+            self.upsamplers = None
+
+        self.resolution_idx = resolution_idx
+
+    def forward(
+        self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
+    ) -> torch.FloatTensor:
+        for resnet in self.resnets:
+            hidden_states = resnet(hidden_states, temb=temb, scale=scale)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states)
+
+        return hidden_states

hyvideo/vae/vae.py

+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from diffusers.utils import BaseOutput, is_torch_version
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.models.attention_processor import SpatialNorm
+from .unet_causal_3d_blocks import (
+    CausalConv3d,
+    UNetMidBlockCausal3D,
+    get_down_block3d,
+    get_up_block3d,
+)
+
+
+@dataclass
+class DecoderOutput(BaseOutput):
+    r"""
+    Output of decoding method.
+
+    Args:
+        sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            The decoded output sample from the last layer of the model.
+    """
+
+    sample: torch.FloatTensor
+
+
+class EncoderCausal3D(nn.Module):
+    r"""
+    The `EncoderCausal3D` layer of a variational autoencoder that encodes its input into a latent representation.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        down_block_types: Tuple[str, ...] = ("DownEncoderBlockCausal3D",),
+        block_out_channels: Tuple[int, ...] = (64,),
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        double_z: bool = True,
+        mid_block_add_attention=True,
+        time_compression_ratio: int = 4,
+        spatial_compression_ratio: int = 8,
+    ):
+        super().__init__()
+        self.layers_per_block = layers_per_block
+
+        self.conv_in = CausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1)
+        self.mid_block = None
+        self.down_blocks = nn.ModuleList([])
+
+        # 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
+            num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio))
+            num_time_downsample_layers = int(np.log2(time_compression_ratio))
+
+            if time_compression_ratio == 4:
+                add_spatial_downsample = bool(i < num_spatial_downsample_layers)
+                add_time_downsample = bool(
+                    i >= (len(block_out_channels) - 1 - num_time_downsample_layers)
+                    and not is_final_block
+                )
+            else:
+                raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}.")
+
+            downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1)
+            downsample_stride_T = (2,) if add_time_downsample else (1,)
+            downsample_stride = tuple(downsample_stride_T + downsample_stride_HW)
+            down_block = get_down_block3d(
+                down_block_type,
+                num_layers=self.layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                add_downsample=bool(add_spatial_downsample or add_time_downsample),
+                downsample_stride=downsample_stride,
+                resnet_eps=1e-6,
+                downsample_padding=0,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                attention_head_dim=output_channel,
+                temb_channels=None,
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        self.mid_block = UNetMidBlockCausal3D(
+            in_channels=block_out_channels[-1],
+            resnet_eps=1e-6,
+            resnet_act_fn=act_fn,
+            output_scale_factor=1,
+            resnet_time_scale_shift="default",
+            attention_head_dim=block_out_channels[-1],
+            resnet_groups=norm_num_groups,
+            temb_channels=None,
+            add_attention=mid_block_add_attention,
+        )
+
+        # out
+        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
+        self.conv_act = nn.SiLU()
+
+        conv_out_channels = 2 * out_channels if double_z else out_channels
+        self.conv_out = CausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3)
+
+    def forward(self, sample: torch.FloatTensor) -> torch.FloatTensor:
+        r"""The forward method of the `EncoderCausal3D` class."""
+        assert len(sample.shape) == 5, "The input tensor should have 5 dimensions"
+
+        sample = self.conv_in(sample)
+
+        # down
+        for down_block in self.down_blocks:
+            sample = down_block(sample)
+
+        # middle
+        sample = self.mid_block(sample)
+
+        # post-process
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        return sample
+
+
+class DecoderCausal3D(nn.Module):
+    r"""
+    The `DecoderCausal3D` layer of a variational autoencoder that decodes its latent representation into an output sample.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        up_block_types: Tuple[str, ...] = ("UpDecoderBlockCausal3D",),
+        block_out_channels: Tuple[int, ...] = (64,),
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        norm_type: str = "group",  # group, spatial
+        mid_block_add_attention=True,
+        time_compression_ratio: int = 4,
+        spatial_compression_ratio: int = 8,
+    ):
+        super().__init__()
+        self.layers_per_block = layers_per_block
+
+        self.conv_in = CausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1)
+        self.mid_block = None
+        self.up_blocks = nn.ModuleList([])
+
+        temb_channels = in_channels if norm_type == "spatial" else None
+
+        # mid
+        self.mid_block = UNetMidBlockCausal3D(
+            in_channels=block_out_channels[-1],
+            resnet_eps=1e-6,
+            resnet_act_fn=act_fn,
+            output_scale_factor=1,
+            resnet_time_scale_shift="default" if norm_type == "group" else norm_type,
+            attention_head_dim=block_out_channels[-1],
+            resnet_groups=norm_num_groups,
+            temb_channels=temb_channels,
+            add_attention=mid_block_add_attention,
+        )
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio))
+            num_time_upsample_layers = int(np.log2(time_compression_ratio))
+
+            if time_compression_ratio == 4:
+                add_spatial_upsample = bool(i < num_spatial_upsample_layers)
+                add_time_upsample = bool(
+                    i >= len(block_out_channels) - 1 - num_time_upsample_layers
+                    and not is_final_block
+                )
+            else:
+                raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}.")
+
+            upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1)
+            upsample_scale_factor_T = (2,) if add_time_upsample else (1,)
+            upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW)
+            up_block = get_up_block3d(
+                up_block_type,
+                num_layers=self.layers_per_block + 1,
+                in_channels=prev_output_channel,
+                out_channels=output_channel,
+                prev_output_channel=None,
+                add_upsample=bool(add_spatial_upsample or add_time_upsample),
+                upsample_scale_factor=upsample_scale_factor,
+                resnet_eps=1e-6,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                attention_head_dim=output_channel,
+                temb_channels=temb_channels,
+                resnet_time_scale_shift=norm_type,
+            )
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        if norm_type == "spatial":
+            self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels)
+        else:
+            self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6)
+        self.conv_act = nn.SiLU()
+        self.conv_out = CausalConv3d(block_out_channels[0], out_channels, kernel_size=3)
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        latent_embeds: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        r"""The forward method of the `DecoderCausal3D` class."""
+        assert len(sample.shape) == 5, "The input tensor should have 5 dimensions."
+
+        sample = self.conv_in(sample)
+
+        upscale_dtype = next(iter(self.up_blocks.parameters())).dtype
+        if self.training and self.gradient_checkpointing:
+
+            def create_custom_forward(module):
+                def custom_forward(*inputs):
+                    return module(*inputs)
+
+                return custom_forward
+
+            if is_torch_version(">=", "1.11.0"):
+                # middle
+                sample = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(self.mid_block),
+                    sample,
+                    latent_embeds,
+                    use_reentrant=False,
+                )
+                sample = sample.to(upscale_dtype)
+
+                # up
+                for up_block in self.up_blocks:
+                    sample = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(up_block),
+                        sample,
+                        latent_embeds,
+                        use_reentrant=False,
+                    )
+            else:
+                # middle
+                sample = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(self.mid_block), sample, latent_embeds
+                )
+                sample = sample.to(upscale_dtype)
+
+                # up
+                for up_block in self.up_blocks:
+                    sample = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), sample, latent_embeds)
+        else:
+            # middle
+            sample = self.mid_block(sample, latent_embeds)
+            sample = sample.to(upscale_dtype)
+
+            # up
+            for up_block in self.up_blocks:
+                sample = up_block(sample, latent_embeds)
+
+        # post-process
+        if latent_embeds is None:
+            sample = self.conv_norm_out(sample)
+        else:
+            sample = self.conv_norm_out(sample, latent_embeds)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        return sample
+
+
+class DiagonalGaussianDistribution(object):
+    def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
+        if parameters.ndim == 3:
+            dim = 2  # (B, L, C)
+        elif parameters.ndim == 5 or parameters.ndim == 4:
+            dim = 1  # (B, C, T, H ,W) / (B, C, H, W)
+        else:
+            raise NotImplementedError
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=dim)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(
+                self.mean, device=self.parameters.device, dtype=self.parameters.dtype
+            )
+
+    def sample(self, generator: Optional[torch.Generator] = None) -> torch.FloatTensor:
+        # make sure sample is on the same device as the parameters and has same dtype
+        sample = randn_tensor(
+            self.mean.shape,
+            generator=generator,
+            device=self.parameters.device,
+            dtype=self.parameters.dtype,
+        )
+        x = self.mean + self.std * sample
+        return x
+
+    def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        else:
+            reduce_dim = list(range(1, self.mean.ndim))
+            if other is None:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+                    dim=reduce_dim,
+                )
+            else:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean - other.mean, 2) / other.var
+                    + self.var / other.var
+                    - 1.0
+                    - self.logvar
+                    + other.logvar,
+                    dim=reduce_dim,
+                )
+
+    def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar +
+            torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims,
+        )
+
+    def mode(self) -> torch.Tensor:
+        return self.mean

modified/config.py

+import os
+import torch
+import torch.distributed as dist
+from packaging import version
+from dataclasses import dataclass, fields
+
+from torch import distributed as dist
+
+from xfuser.logger import init_logger
+import xfuser.envs as envs
+# from xfuser.envs import CUDA_VERSION, TORCH_VERSION, PACKAGES_CHECKER
+from xfuser.envs import TORCH_VERSION, PACKAGES_CHECKER
+
+logger = init_logger(__name__)
+
+from typing import Union, Optional, List
+
+env_info = PACKAGES_CHECKER.get_packages_info()
+HAS_LONG_CTX_ATTN = env_info["has_long_ctx_attn"]
+HAS_FLASH_ATTN = env_info["has_flash_attn"]
+
+
+def check_packages():
+    import diffusers
+
+    if not version.parse(diffusers.__version__) > version.parse("0.30.2"):
+        raise RuntimeError(
+            "This project requires diffusers version > 0.30.2. Currently, you can not install a correct version of diffusers by pip install."
+            "Please install it from source code!"
+        )
+
+
+def check_env():
+    # https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/cudagraph.html
+    #if CUDA_VERSION < version.parse("11.3"):
+    #    raise RuntimeError("NCCL CUDA Graph support requires CUDA 11.3 or above")
+    if TORCH_VERSION < version.parse("2.2.0"):
+        # https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/
+        raise RuntimeError(
+            "CUDAGraph with NCCL support requires PyTorch 2.2.0 or above. "
+            "If it is not released yet, please install nightly built PyTorch "
+            "with `pip3 install --pre torch torchvision torchaudio --index-url "
+            "https://download.pytorch.org/whl/nightly/cu121`"
+        )
+
+
+@dataclass
+class ModelConfig:
+    model: str
+    download_dir: Optional[str] = None
+    trust_remote_code: bool = False
+
+
+@dataclass
+class RuntimeConfig:
+    warmup_steps: int = 1
+    dtype: torch.dtype = torch.float16
+    use_cuda_graph: bool = False
+    use_parallel_vae: bool = False
+    use_profiler: bool = False
+    use_torch_compile: bool = False
+    use_onediff: bool = False
+    use_fp8_t5_encoder: bool = False
+
+    def __post_init__(self):
+        check_packages()
+        if self.use_cuda_graph:
+            check_env()
+
+
+@dataclass
+class FastAttnConfig:
+    use_fast_attn: bool = False
+    n_step: int = 20
+    n_calib: int = 8
+    threshold: float = 0.5
+    window_size: int = 64
+    coco_path: Optional[str] = None
+    use_cache: bool = False
+
+    def __post_init__(self):
+        assert self.n_calib > 0, "n_calib must be greater than 0"
+        assert self.threshold > 0.0, "threshold must be greater than 0"
+
+
+@dataclass
+class DataParallelConfig:
+    dp_degree: int = 1
+    use_cfg_parallel: bool = False
+    world_size: int = 1
+
+    def __post_init__(self):
+        assert self.dp_degree >= 1, "dp_degree must greater than or equal to 1"
+
+        # set classifier_free_guidance_degree parallel for split batch
+        if self.use_cfg_parallel:
+            self.cfg_degree = 2
+        else:
+            self.cfg_degree = 1
+        assert self.dp_degree * self.cfg_degree <= self.world_size, (
+            "dp_degree * cfg_degree must be less than or equal to "
+            "world_size because of classifier free guidance"
+        )
+        assert (
+            self.world_size % (self.dp_degree * self.cfg_degree) == 0
+        ), "world_size must be divisible by dp_degree * cfg_degree"
+
+
+@dataclass
+class SequenceParallelConfig:
+    ulysses_degree: Optional[int] = None
+    ring_degree: Optional[int] = None
+    world_size: int = 1
+
+    def __post_init__(self):
+        if self.ulysses_degree is None:
+            self.ulysses_degree = 1
+            logger.info(
+                f"Ulysses degree not set, " f"using default value {self.ulysses_degree}"
+            )
+        if self.ring_degree is None:
+            self.ring_degree = 1
+            logger.info(
+                f"Ring degree not set, " f"using default value {self.ring_degree}"
+            )
+        self.sp_degree = self.ulysses_degree * self.ring_degree
+
+        if not HAS_LONG_CTX_ATTN and self.sp_degree > 1:
+            raise ImportError(
+                f"Sequence Parallel kit 'yunchang' not found but "
+                f"sp_degree is {self.sp_degree}, please set it "
+                f"to 1 or install 'yunchang' to use it"
+            )
+
+
+@dataclass
+class TensorParallelConfig:
+    tp_degree: int = 1
+    split_scheme: Optional[str] = "row"
+    world_size: int = 1
+
+    def __post_init__(self):
+        assert self.tp_degree >= 1, "tp_degree must greater than 1"
+        assert (
+            self.tp_degree <= self.world_size
+        ), "tp_degree must be less than or equal to world_size"
+
+
+@dataclass
+class PipeFusionParallelConfig:
+    pp_degree: int = 1
+    num_pipeline_patch: Optional[int] = None
+    attn_layer_num_for_pp: Optional[List[int]] = (None,)
+    world_size: int = 1
+
+    def __post_init__(self):
+        assert (
+            self.pp_degree is not None and self.pp_degree >= 1
+        ), "pipefusion_degree must be set and greater than 1 to use pipefusion"
+        assert (
+            self.pp_degree <= self.world_size
+        ), "pipefusion_degree must be less than or equal to world_size"
+        if self.num_pipeline_patch is None:
+            self.num_pipeline_patch = self.pp_degree
+            logger.info(
+                f"Pipeline patch number not set, "
+                f"using default value {self.pp_degree}"
+            )
+        if self.attn_layer_num_for_pp is not None:
+            logger.info(
+                f"attn_layer_num_for_pp set, splitting attention layers"
+                f"to {self.attn_layer_num_for_pp}"
+            )
+            assert len(self.attn_layer_num_for_pp) == self.pp_degree, (
+                "attn_layer_num_for_pp must have the same "
+                "length as pp_degree if not None"
+            )
+        if self.pp_degree == 1 and self.num_pipeline_patch > 1:
+            logger.warning(
+                f"Pipefusion degree is 1, pipeline will not be used,"
+                f"num_pipeline_patch will be ignored"
+            )
+            self.num_pipeline_patch = 1
+
+
+@dataclass
+class ParallelConfig:
+    dp_config: DataParallelConfig
+    sp_config: SequenceParallelConfig
+    pp_config: PipeFusionParallelConfig
+    tp_config: TensorParallelConfig
+    world_size: int = 1 # FIXME: remove this
+    worker_cls: str = "xfuser.ray.worker.worker.Worker"
+
+    def __post_init__(self):
+        assert self.tp_config is not None, "tp_config must be set"
+        assert self.dp_config is not None, "dp_config must be set"
+        assert self.sp_config is not None, "sp_config must be set"
+        assert self.pp_config is not None, "pp_config must be set"
+        parallel_world_size = (
+            self.dp_config.dp_degree
+            * self.dp_config.cfg_degree
+            * self.sp_config.sp_degree
+            * self.tp_config.tp_degree
+            * self.pp_config.pp_degree
+        )
+        world_size = self.world_size
+        assert parallel_world_size == world_size, (
+            f"parallel_world_size {parallel_world_size} "
+            f"must be equal to world_size {self.world_size}"
+        )
+        assert (
+            world_size % (self.dp_config.dp_degree * self.dp_config.cfg_degree) == 0
+        ), "world_size must be divisible by dp_degree * cfg_degree"
+        assert (
+            world_size % self.pp_config.pp_degree == 0
+        ), "world_size must be divisible by pp_degree"
+        assert (
+            world_size % self.sp_config.sp_degree == 0
+        ), "world_size must be divisible by sp_degree"
+        assert (
+            world_size % self.tp_config.tp_degree == 0
+        ), "world_size must be divisible by tp_degree"
+        self.dp_degree = self.dp_config.dp_degree
+        self.cfg_degree = self.dp_config.cfg_degree
+        self.sp_degree = self.sp_config.sp_degree
+        self.pp_degree = self.pp_config.pp_degree
+        self.tp_degree = self.tp_config.tp_degree
+
+        self.ulysses_degree = self.sp_config.ulysses_degree
+        self.ring_degree = self.sp_config.ring_degree
+
+
+@dataclass(frozen=True)
+class EngineConfig:
+    model_config: ModelConfig
+    runtime_config: RuntimeConfig
+    parallel_config: ParallelConfig
+    fast_attn_config: FastAttnConfig
+
+    def __post_init__(self):
+        world_size = self.parallel_config.world_size
+        if self.fast_attn_config.use_fast_attn:
+            assert self.parallel_config.dp_degree == world_size, f"world_size must be equal to dp_degree when using DiTFastAttn"
+
+    def to_dict(self):
+        """Return the configs as a dictionary, for use in **kwargs."""
+        return dict((field.name, getattr(self, field.name)) for field in fields(self))
+
+
+@dataclass
+class InputConfig:
+    height: int = 1024
+    width: int = 1024
+    num_frames: int = 49
+    use_resolution_binning: bool = (True,)
+    batch_size: Optional[int] = None
+    img_file_path: Optional[str] = None
+    prompt: Union[str, List[str]] = ""
+    negative_prompt: Union[str, List[str]] = ""
+    num_inference_steps: int = 20
+    max_sequence_length: int = 256
+    seed: int = 42
+    output_type: str = "pil"
+
+    def __post_init__(self):
+        if isinstance(self.prompt, list):
+            assert (
+                len(self.prompt) == len(self.negative_prompt)
+                or len(self.negative_prompt) == 0
+            ), "prompts and negative_prompts must have the same quantities"
+            self.batch_size = self.batch_size or len(self.prompt)
+        else:
+            self.batch_size = self.batch_size or 1
+        assert self.output_type in [
+            "pil",
+            "latent",
+            "pt",
+        ], "output_pil must be either 'pil' or 'latent'"
+

modified/envs.py

+import os
+import torch
+import diffusers
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+from packaging import version
+
+from xfuser.logger import init_logger
+
+logger = init_logger(__name__)
+
+if TYPE_CHECKING:
+    MASTER_ADDR: str = ""
+    MASTER_PORT: Optional[int] = None
+    CUDA_HOME: Optional[str] = None
+    LOCAL_RANK: int = 0
+    CUDA_VISIBLE_DEVICES: Optional[str] = None
+    XDIT_LOGGING_LEVEL: str = "INFO"
+    CUDA_VERSION: version.Version
+    TORCH_VERSION: version.Version
+
+
+environment_variables: Dict[str, Callable[[], Any]] = {
+    # ================== Runtime Env Vars ==================
+    # used in distributed environment to determine the master address
+    "MASTER_ADDR": lambda: os.getenv("MASTER_ADDR", ""),
+    # used in distributed environment to manually set the communication port
+    "MASTER_PORT": lambda: (
+        int(os.getenv("MASTER_PORT", "0")) if "MASTER_PORT" in os.environ else None
+    ),
+    # path to cudatoolkit home directory, under which should be bin, include,
+    # and lib directories.
+    "CUDA_HOME": lambda: os.environ.get("CUDA_HOME", None),
+    # local rank of the process in the distributed setting, used to determine
+    # the GPU device id
+    "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")),
+    # used to control the visible devices in the distributed setting
+    "CUDA_VISIBLE_DEVICES": lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None),
+    # this is used for configuring the default logging level
+    "XDIT_LOGGING_LEVEL": lambda: os.getenv("XDIT_LOGGING_LEVEL", "INFO"),
+}
+
+variables: Dict[str, Callable[[], Any]] = {
+    # ================== Other Vars ==================
+    # used in version checking
+    # "CUDA_VERSION": lambda: version.parse(torch.version.cuda),
+    "CUDA_VERSION": "gfx936",
+    "TORCH_VERSION": lambda: version.parse(
+        version.parse(torch.__version__).base_version
+    ),
+}
+
+
+class PackagesEnvChecker:
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(PackagesEnvChecker, cls).__new__(cls)
+            cls._instance.initialize()
+        return cls._instance
+
+    def initialize(self):
+        self.packages_info = {
+            "has_flash_attn": self.check_flash_attn(),
+            "has_long_ctx_attn": self.check_long_ctx_attn(),
+            "diffusers_version": self.check_diffusers_version(),
+        }
+
+    def check_flash_attn(self):
+        try:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            gpu_name = torch.cuda.get_device_name(device)
+            if "Turing" in gpu_name or "Tesla" in gpu_name or "T4" in gpu_name:
+                return False
+            else:
+                from flash_attn import flash_attn_func
+                from flash_attn import __version__
+
+                if __version__ < "2.6.0":
+                    raise ImportError(f"install flash_attn >= 2.6.0")
+                return True
+        except ImportError:
+            logger.warning(
+                f'Flash Attention library "flash_attn" not found, '
+                f"using pytorch attention implementation"
+            )
+            return False
+
+    def check_long_ctx_attn(self):
+        try:
+            from yunchang import (
+                set_seq_parallel_pg,
+                ring_flash_attn_func,
+                UlyssesAttention,
+                LongContextAttention,
+                LongContextAttentionQKVPacked,
+            )
+
+            return True
+        except ImportError:
+            logger.warning(
+                f'Ring Flash Attention library "yunchang" not found, '
+                f"using pytorch attention implementation"
+            )
+            return False
+
+    def check_diffusers_version(self):
+        if version.parse(
+            version.parse(diffusers.__version__).base_version
+        ) < version.parse("0.30.0"):
+            raise RuntimeError(
+                f"Diffusers version: {version.parse(version.parse(diffusers.__version__).base_version)} is not supported,"
+                f"please upgrade to version > 0.30.0"
+            )
+        return version.parse(version.parse(diffusers.__version__).base_version)
+
+    def get_packages_info(self):
+        return self.packages_info
+
+
+PACKAGES_CHECKER = PackagesEnvChecker()
+
+
+def __getattr__(name):
+    # lazy evaluation of environment variables
+    if name in environment_variables:
+        return environment_variables[name]()
+    if name in variables:
+        return variables[name]()
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__():
+    return list(environment_variables.keys())
+

requirements.txt

+opencv-python==4.9.0.80
+diffusers==0.31.0
+accelerate==1.1.1
+pandas==2.0.3
+einops==0.7.0
+tqdm==4.66.2
+loguru==0.7.2
+imageio==2.34.0
+imageio-ffmpeg==0.5.1
+safetensors==0.4.3
+peft==0.13.2
+transformers==4.39.3
+tokenizers==0.15.0
+pyarrow==14.0.1
+tensorboard==2.19.0
\ No newline at end of file

run.sh

+ALLOW_RESIZE_FOR_SP=1 torchrun --nproc_per_node=4 \
+    sample_image2video.py \
+    --model HYVideo-T/2 \
+    --prompt "An Asian man with short hair in black tactical uniform and white clothes waves a firework stick." \
+    --i2v-mode \
+    --i2v-image-path ./assets/demo/i2v/imgs/0.jpg \
+    --i2v-resolution 720p \
+    --i2v-stability \
+    --infer-steps 50 \
+    --video-length 129 \
+    --flow-reverse \
+    --flow-shift 17.0 \
+    --seed 42 \
+    --embedded-cfg-scale 6.0 \
+    --save-path ./results \
+    --ulysses-degree 4 \
+    --ring-degree 1 \
+    --num-videos 1 2>&1 |tee z_logs/bw_image2video_4ka.log

sample_image2video.py

+import os
+import time
+from pathlib import Path
+from loguru import logger
+from datetime import datetime
+
+from hyvideo.utils.file_utils import save_videos_grid
+from hyvideo.config import parse_args
+from hyvideo.inference import HunyuanVideoSampler
+
+
+def main():
+    args = parse_args()
+    print(args)
+    models_root_path = Path(args.model_base)
+    if not models_root_path.exists():
+        raise ValueError(f"`models_root` not exists: {models_root_path}")
+    
+    # Create save folder to save the samples
+    save_path = args.save_path if args.save_path_suffix=="" else f'{args.save_path}_{args.save_path_suffix}'
+    if not os.path.exists(save_path):
+        os.makedirs(save_path, exist_ok=True)
+
+    # Load models
+    hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(models_root_path, args=args)
+    
+    # Get the updated args
+    args = hunyuan_video_sampler.args
+
+    # Start sampling
+    # TODO: batch inference check
+    outputs = hunyuan_video_sampler.predict(
+        prompt=args.prompt, 
+        height=args.video_size[0],
+        width=args.video_size[1],
+        video_length=args.video_length,
+        seed=args.seed,
+        negative_prompt=args.neg_prompt,
+        infer_steps=args.infer_steps,
+        guidance_scale=args.cfg_scale,
+        num_videos_per_prompt=args.num_videos,
+        flow_shift=args.flow_shift,
+        batch_size=args.batch_size,
+        embedded_guidance_scale=args.embedded_cfg_scale,
+        i2v_mode=args.i2v_mode,
+        i2v_resolution=args.i2v_resolution,
+        i2v_image_path=args.i2v_image_path,
+        i2v_condition_type=args.i2v_condition_type,
+        i2v_stability=args.i2v_stability,
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+    )
+    samples = outputs['samples']
+    
+    # Save samples
+    if 'LOCAL_RANK' not in os.environ or int(os.environ['LOCAL_RANK']) == 0:
+        for i, sample in enumerate(samples):
+            sample = samples[i].unsqueeze(0)
+            time_flag = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d-%H:%M:%S")
+            cur_save_path = f"{save_path}/{time_flag}_seed{outputs['seeds'][i]}_{outputs['prompts'][i][:100].replace('/','')}.mp4"
+            save_videos_grid(sample, cur_save_path, fps=24)
+            logger.info(f'Sample save to: {cur_save_path}')
+
+if __name__ == "__main__":
+    main()

scripts/run_sample_image2video_dynamic.sh

+#!/bin/bash
+
+python3 sample_image2video.py \
+    --prompt "An Asian man with short hair in black tactical uniform and white clothes waves a firework stick." \
+    --i2v-image-path ./assets/demo/i2v/imgs/0.jpg \
+    --model HYVideo-T/2 \
+    --i2v-mode \
+    --i2v-resolution 720p \
+    --infer-steps 50 \
+    --video-length 129 \
+    --flow-reverse \
+    --flow-shift 17.0 \
+    --embedded-cfg-scale 6.0 \
+    --seed 0 \
+    --use-cpu-offload \
+    --save-path ./results \
+
+# More example
+#    --prompt "A girl walks on the road, shooting stars pass by." \
+#    --i2v-image-path ./assets/demo/i2v/imgs/1.png \
\ No newline at end of file

scripts/run_sample_image2video_lora.sh

+#!/bin/bash
+
+python3 sample_image2video.py \
+   --prompt "Two people hugged tightly, In the video, two people are standing apart from each other. They then move closer to each other and begin to hug tightly. The hug is very affectionate, with the two people holding each other tightly and looking into each other's eyes. The interaction is very emotional and heartwarming, with the two people expressing their love and affection for each other." \
+   --i2v-image-path ./assets/demo/i2v_lora/imgs/embrace.png \
+   --lora-path ./ckpts/hunyuan-video-i2v-720p/lora/embrace_kohaya_weights.safetensors \
+   --model HYVideo-T/2 \
+   --i2v-mode \
+   --i2v-resolution 720p \
+   --i2v-stability \
+   --infer-steps 50 \
+   --video-length 129 \
+   --flow-reverse \
+   --flow-shift 5.0 \
+   --embedded-cfg-scale 6.0 \
+   --seed 0 \
+   --use-cpu-offload \
+   --save-path ./results \
+   --use-lora \
+   --lora-scale 1.0 \
+
+# More examples
+#    --prompt "rapid_hair_growth, The hair of the characters in the video is growing rapidly. The character's hair undergoes a dramatic transformation, growing rapidly from a short, straight style to a long, wavy one. Initially, the hair is a light blonde color, but as it grows, it becomes darker and more voluminous. The character's facial features remain consistent throughout the transformation, with a slight change in the shape of the jawline as the hair grows. The clothing changes from a simple, casual outfit to a more elaborate, fashionable ensemble that complements the longer hair. The overall appearance shifts from a casual, everyday look to a more stylish, sophisticated one. The character's expression remains calm and composed throughout the transformation, with a slight smile as the hair grows." \
+#    --i2v-image-path ./assets/demo/i2v_lora/imgs/hair_growth.png \
+#    --lora-path ./ckpts/hunyuan-video-i2v-720p/lora/hair_growth_kohaya_weights.safetensors \

scripts/run_sample_image2video_stability.sh

+#!/bin/bash
+
+python3 sample_image2video.py \
+    --prompt "An Asian man with short hair in black tactical uniform and white clothes waves a firework stick." \
+    --i2v-image-path ./assets/demo/i2v/imgs/0.jpg \
+    --model HYVideo-T/2 \
+    --i2v-mode \
+    --i2v-resolution 720p \
+    --i2v-stability \
+    --infer-steps 50 \
+    --video-length 129 \
+    --flow-reverse \
+    --flow-shift 7.0 \
+    --embedded-cfg-scale 6.0 \
+    --seed 0 \
+    --use-cpu-offload \
+    --save-path ./results \
+
+# More example
+#    --prompt "A girl walks on the road, shooting stars pass by." \
+#    --i2v-image-path ./assets/demo/i2v/imgs/1.png \
\ No newline at end of file

scripts/run_train_image2video_lora.sh

+# Root path for saving experimental results
+export SAVE_BASE="."
+echo "SAVE_BASE: ${SAVE_BASE}"
+
+# Path suffix for saving experimental results
+EXP_NAME="i2v_lora"
+
+# Data jsons dir (output_base_dir/json_path in hyvideo/hyvae_extract/README.md) generated by hyvideo/hyvae_extract/start.sh
+DATA_JSONS_DIR="./assets/demo/i2v_lora/train_dataset/processed_data/json_path"
+
+# Master node IP of the machine
+CHIEF_IP="127.0.0.1"
+
+current_datetime=$(date +%Y%m%d_%H%M%S)
+output_dir="${SAVE_BASE}/log_EXP"
+task_flag="${current_datetime}_${EXP_NAME}"
+
+params=" \
+    --lr 1e-4 \
+    --warmup-num-steps 500 \
+    --global-seed 1024 \
+    --tensorboard \
+    --zero-stage 2 \
+    --vae 884-16c-hy \
+    --vae-precision fp16 \
+    --vae-tiling \
+    --denoise-type flow \
+    --flow-reverse \
+    --flow-shift 7.0 \
+    --i2v-mode \
+    --model HYVideo-T/2 \
+    --video-micro-batch-size 1 \
+    --gradient-checkpoint \
+    --ckpt-every 500 \
+    --embedded-cfg-scale 6.0 \
+    "
+
+video_data_params=" \
+    --data-type video \
+    --data-jsons-path ${DATA_JSONS_DIR} \
+    --sample-n-frames 129 \
+    --sample-stride 1 \
+    --num-workers 8 \
+    --uncond-p 0.1 \
+    --sematic-cond-drop-p 0.1 \
+    "
+
+te_params=" \
+    --text-encoder llm-i2v \
+    --text-encoder-precision fp16 \
+    --text-states-dim 4096 \
+    --text-len 256 \
+    --tokenizer llm-i2v \
+    --prompt-template dit-llm-encode-i2v \
+    --prompt-template-video dit-llm-encode-video-i2v \
+    --hidden-state-skip-layer 2 \
+    --text-encoder-2 clipL \
+    --text-encoder-precision-2 fp16 \
+    --text-states-dim-2 768 \
+    --tokenizer-2 clipL \
+    --text-len-2 77 \
+    "
+    
+lora_params=" \
+    --use-lora \
+    --lora-rank 64 \
+    "
+
+export TOKENIZERS_PARALLELISM=false
+
+set -x
+
+#deepspeed --hostfile $hostfile --master_addr "${CHIEF_IP}" \
+# single node, multi gpu
+#deepspeed --include localhost:0,1,2,3,4,5,6,7 --master_addr "${CHIEF_IP}" \
+# single node, single gpu
+deepspeed --include localhost:0 --master_addr "${CHIEF_IP}" \
+train_image2video_lora.py \
+    ${params} \
+    ${val_params} \
+    ${video_data_params} \
+    ${te_params} \
+    ${lora_params} \
+    --task-flag ${task_flag} \
+    --output-dir ${output_dir} \
+    "$@"
\ No newline at end of file

train_image2video_lora.py

+import os
+import sys
+import time
+import warnings
+import json
+from collections import defaultdict
+from dataclasses import dataclass, asdict, field
+from pathlib import Path
+from typing import Optional, Dict, Union
+from functools import partial
+
+warnings.filterwarnings("ignore")
+
+import deepspeed
+import torch
+import torchvision
+import torch.distributed as dist
+from deepspeed.runtime import lr_schedules
+from deepspeed.runtime.engine import DeepSpeedEngine
+from torch.utils.data.distributed import DistributedSampler
+from torch.utils.data import DataLoader
+from einops import rearrange
+
+from hyvideo.config import parse_args
+from hyvideo.constants import C_SCALE, PROMPT_TEMPLATE
+from hyvideo.dataset.video_loader import VideoDataset
+from hyvideo.diffusion import load_denoiser
+from hyvideo.ds_config import get_deepspeed_config
+from hyvideo.utils.train_utils import (
+    prepare_model_inputs,
+    load_state_dict,
+    set_worker_seed_builder,
+    get_module_kohya_state_dict,
+    load_lora,
+)
+from hyvideo.modules import load_model
+from hyvideo.text_encoder import TextEncoder
+from hyvideo.utils.file_utils import (
+    safe_dir,
+    get_experiment_max_number,
+    empty_logger,
+    dump_args,
+    dump_codes,
+    resolve_resume_path,
+    logger_filter,
+)
+from hyvideo.utils.helpers import (
+    as_tuple,
+    set_manual_seed,
+    set_reproducibility,
+    profiler_context,
+    all_gather_sum,
+    EventsMonitor,
+)
+from hyvideo.vae import load_vae
+from hyvideo.constants import PRECISION_TO_TYPE
+
+from peft import LoraConfig, get_peft_model
+from safetensors.torch import save_file
+
+
+def setup_distributed_training(args):
+    deepspeed.init_distributed()
+
+    # Treat micro/global batch size as tuples for compatibility with mix-scale training.
+    world_size = dist.get_world_size()
+    if args.data_type == "video" and args.video_micro_batch_size is None:
+        # When data_type is video and video_micro_batch_size is None, we set the value from micro_batch_size
+        args.video_micro_batch_size = args.micro_batch_size
+
+    micro_batch_size = as_tuple(args.micro_batch_size)
+    video_micro_batch_size = as_tuple(args.video_micro_batch_size)
+    grad_accu_steps = args.gradient_accumulation_steps
+    global_batch_size = as_tuple(args.global_batch_size)
+    if "video" in args.data_type:
+        refer_micro_batch_size = video_micro_batch_size
+    else:
+        refer_micro_batch_size = micro_batch_size
+
+    if global_batch_size[0] is None:
+        # Note: Model/Pipeline parallel is not supported yet. So, data-parallel-size equals to world-size.
+        global_batch_size = tuple(
+            [mbs_i * world_size * grad_accu_steps for mbs_i in refer_micro_batch_size]
+        )
+    else:
+        assert global_batch_size == [
+            mbs_i * world_size * grad_accu_steps for mbs_i in refer_micro_batch_size
+        ], f"Global batch size should be divisible by world size, but got {global_batch_size} and {world_size}."
+
+    rank = dist.get_rank()  # Rank of the current process in the cluster.
+    device = (
+        rank % torch.cuda.device_count()
+    )  # Device of the current process in current node.
+    # Set current device for the current process, otherwise dist.barrier() will occupy more memory in rank 0.
+    torch.cuda.set_device(device)
+
+    # Setup seed for reproducibility or performance.
+    set_manual_seed(args.global_seed)
+    set_reproducibility(args.reproduce, args.global_seed)
+
+    return (
+        rank,
+        device,
+        world_size,
+        micro_batch_size,
+        video_micro_batch_size,
+        grad_accu_steps,
+        global_batch_size,
+    )
+
+
+def setup_experiment_directory(args, rank):
+    output_dir = safe_dir(args.output_dir)
+
+    # Automatically increase the experiment number.
+    existed_experiments = list(output_dir.glob("*"))
+    experiment_index = get_experiment_max_number(existed_experiments) + 1
+    model_name = args.model.replace("/", "").replace(
+        "-", "_"
+    )  # Replace '/' to avoid sub-directory.
+    experiment_dir = (
+        output_dir / f"{experiment_index:04d}_{model_name}_{args.task_flag}"
+    )
+    ckpt_dir = experiment_dir / "checkpoints"
+
+    # Makesure all processes have the same experiment directory.
+    dist.barrier()
+
+    if rank == 0:
+        from loguru import logger
+
+        logger.add(
+            experiment_dir / "train.log",
+            level="DEBUG",
+            colorize=False,
+            backtrace=True,
+            diagnose=True,
+            encoding="utf-8",
+            filter=logger_filter("train"),
+        )
+        logger.add(
+            experiment_dir / "val.log",
+            level="DEBUG",
+            colorize=False,
+            backtrace=True,
+            diagnose=True,
+            encoding="utf-8",
+            filter=logger_filter("val"),
+        )
+        train_logger = logger.bind(name="train")
+        val_logger = logger.bind(name="val")
+
+        ckpt_dir = safe_dir(ckpt_dir)
+    else:
+        val_logger = train_logger = empty_logger()
+
+    train_logger.info(f"Experiment directory created at: {experiment_dir}")
+
+    return experiment_dir, ckpt_dir, train_logger, val_logger
+
+
+def get_trainable_params(model, args):
+    if args.training_parts is None:
+        params = []
+        for param in model.parameters():
+            if param.requires_grad == True:
+                params.append(param)
+    else:
+        raise ValueError(f"Unknown training_parts {args.training_parts}")
+    return params
+
+
+@dataclass
+class ScalarStates:
+    rank: int = 0  # rank id
+    epoch: int = 1  # Accumulated training epochs
+    epoch_train_steps: int = 0  # Accumulated training steps in current epoch
+    epoch_update_steps: int = 0  # Accumulated update steps in current epoch
+    train_steps: int = 0  # Accumulated training steps
+    update_steps: int = 0  # Accumulated update steps
+    current_run_update_steps: int = 0  # Update steps in current run
+    consumed_samples_total: int = 0  # Accumulated consumed samples
+    consumed_video_samples_total: int = 0  # Accumulated consumed video samples
+    consumed_samples_per_dp: int = (
+        0  # Accumulated consumed samples per data-parallel group
+    )
+    consumed_video_samples_per_dp: int = (
+        0  # Accumulated consumed video samples per data-parallel group
+    )
+    consumed_tokens_total: int = 0  # Accumulated consumed tokens
+    consumed_computations_attn: int = (
+        0  # Accumulated consumed computations of attention + mlp
+    )
+    consumed_computations_total: int = 0  # Accumulated consumed computations of total
+
+    def add(self, **kwargs):
+        for k, v in kwargs.items():
+            setattr(self, k, getattr(self, k) + v)
+
+
+@dataclass
+class CycleStates:
+    log_steps: int = 0
+    running_loss: float = 0
+    running_tokens: int = 0
+    running_samples: int = 0
+    running_video_samples: int = 0
+    running_grad_norm: float = 0
+    running_loss_dict: Dict[int, float] = field(
+        default_factory=lambda: defaultdict(float)
+    )
+    log_steps_dict: Dict[int, int] = field(default_factory=lambda: defaultdict(int))
+
+    def add(self, **kwargs):
+        for k, v in kwargs.items():
+            setattr(self, k, getattr(self, k) + v)
+
+    def reset(self):
+        self.log_steps = 0
+        self.running_loss = 0
+        self.running_tokens = 0
+        self.running_samples = 0
+        self.running_video_samples = 0
+        self.running_grad_norm = 0
+        # Must be reset to float to avoid all_reduce type error.
+        self.running_loss_dict = defaultdict(float)
+        self.log_steps_dict = defaultdict(int)
+
+
+def save_checkpoint(
+    args,
+    rank: int,
+    logger,
+    model_engine: DeepSpeedEngine,
+    ema,
+    scalar_state: ScalarStates,
+    ckpt_dir: Path,
+):
+    _ = rank  # Currently not used.
+
+    # gather scalar state
+    scalar_state_dict = dict(**asdict(scalar_state))
+    gather_results_list = [None for _ in range(dist.get_world_size())]
+    torch.distributed.all_gather_object(gather_results_list, scalar_state_dict)
+    gather_scalar_states = {}
+    for results in gather_results_list:
+        gather_scalar_states[results["rank"]] = results
+
+    client_state = {
+        "args": args,
+        "scalar_state": gather_scalar_states,
+    }
+    if ema is not None:
+        client_state["ema"] = ema.state_dict()
+        client_state["ema_config"] = ema.config
+
+    def try_save(_save_name):
+        checkpoint_path = ckpt_dir / _save_name
+        try:
+            model_engine.save_checkpoint(
+                str(ckpt_dir),
+                client_state=client_state,
+                tag=_save_name,
+            )
+            logger.info(f"Saved checkpoint to {checkpoint_path}")
+            return checkpoint_path
+        except Exception as e:
+            logger.error(f"Saved failed to {checkpoint_path}. {type(e)}: {e}")
+            return None
+
+    update_steps = scalar_state.update_steps
+    save_name = f"{update_steps:07d}"
+    save_path = try_save(save_name)
+
+    return [save_path]
+
+
+def main(args):
+    # ============================= Setup ==============================
+    # Setup distributed training environment and reproducibility.
+    (
+        rank,
+        device,
+        world_size,
+        micro_batch_size,
+        video_micro_batch_size,
+        grad_accu_steps,
+        global_batch_size,
+    ) = setup_distributed_training(args)
+    # Setup experiment directory
+    exp_dir, ckpt_dir, logger, val_logger = setup_experiment_directory(args, rank)
+    # Load deepspeed config
+    deepspeed_config = get_deepspeed_config(
+        args,
+        video_micro_batch_size[0],
+        global_batch_size[0],
+        args.output_dir,
+        exp_dir.name,
+    )
+    # Log and dump the arguments and codes.
+    logger.info(sys.argv)
+    logger.info(str(args))
+    if rank == 0:
+        # Dump the arguments to a file.
+        extra_args = {"world_size": world_size, "global_batch_size": global_batch_size}
+        dump_args(args, exp_dir / "args.json", extra_args)
+        # Dump codes to the experiment directory.
+        dump_codes(
+            exp_dir / "codes.tar.gz",
+            root=Path(__file__).parent.parent,
+            sub_dirs=["hymm", "jobs"],
+            save_prefix=args.task_flag,
+        )
+
+    # =========================== Build main model ===========================
+    logger.info("Building model...")
+    factor_kwargs = {"device": device, "dtype": PRECISION_TO_TYPE[args.precision]}
+    if args.i2v_mode and args.i2v_condition_type == "latent_concat":
+        in_channels = args.latent_channels * 2 + 1
+        image_embed_interleave = 2
+    elif args.i2v_mode and args.i2v_condition_type == "token_replace":
+        in_channels = args.latent_channels
+        image_embed_interleave = 4
+    else:
+        in_channels = args.latent_channels
+        image_embed_interleave = 1
+    out_channels = args.latent_channels
+
+    if args.embedded_cfg_scale:
+        factor_kwargs["guidance_embed"] = True
+
+    model = load_model(
+        args,
+        in_channels=in_channels,
+        out_channels=out_channels,
+        factor_kwargs=factor_kwargs,
+    )
+    model = load_state_dict(args, model, logger)
+
+    if args.use_lora:
+        for param in model.parameters():
+            param.requires_grad_(False)
+
+        target_modules = [
+            "linear",
+            "fc1",
+            "fc2",
+            "img_attn_qkv",
+            "img_attn_proj",
+            "txt_attn_qkv",
+            "txt_attn_proj",
+            "linear1",
+            "linear2",
+        ]
+
+        lora_config = LoraConfig(
+            r=args.lora_rank,
+            lora_alpha=args.lora_rank,
+            init_lora_weights="gaussian",
+            target_modules=target_modules,
+        )
+        model = get_peft_model(model, lora_config)
+
+        if args.lora_path != "":
+            model = load_lora(model, args.lora_path, device=device)
+
+    logger.info(model)
+
+    if args.reproduce:
+        model.enable_deterministic()
+
+    # After model initialization, we set different seed for each process.
+    if args.same_data_batch:
+        set_manual_seed(args.global_seed)
+    else:
+        set_manual_seed(args.global_seed + rank)
+
+    ema = None
+    ss = ScalarStates(rank=rank)
+
+    # ========================== Initialize model_engine, optimizer =========================
+    if args.warmup_num_steps > 0:
+        logger.info(
+            f"Building scheduler with warmup_min_lr={args.warmup_min_lr}, warmup_max_lr={args.lr}, "
+            f"warmup_num_steps={args.warmup_num_steps}."
+        )
+        lr_scheduler = partial(
+            lr_schedules.WarmupLR,
+            warmup_min_lr=args.warmup_min_lr,
+            warmup_max_lr=args.lr,
+            warmup_num_steps=args.warmup_num_steps,
+        )
+    else:
+        lr_scheduler = None
+
+    logger.info("Initializing optimizer (using deepspeed)...")
+    model_engine, opt, _, scheduler = deepspeed.initialize(
+        args=args,
+        model=model,
+        model_parameters=get_trainable_params(model, args),
+        config_params=deepspeed_config,
+        lr_scheduler=lr_scheduler,
+    )
+
+    # ====================== Build denoise scheduler ========================
+    logger.info("Building denoise scheduler...")
+    denoiser = load_denoiser(args)
+
+    # ============================= Build extra models =========================
+    # 2d/3d VAE
+    vae, vae_path, s_ratio, t_ratio = load_vae(
+        args.vae, args.vae_precision, logger=logger, device=device
+    )
+
+    # Text encoder
+    text_encoder = TextEncoder(
+        text_encoder_type=args.text_encoder,
+        max_length=args.text_len
+        + (
+            PROMPT_TEMPLATE[args.prompt_template_video].get("crop_start", 0)
+            if args.prompt_template_video is not None
+            else PROMPT_TEMPLATE[args.prompt_template].get("crop_start", 0)
+            if args.prompt_template is not None
+            else 0
+        ),
+        text_encoder_precision=args.text_encoder_precision,
+        tokenizer_type=args.tokenizer,
+        i2v_mode=args.i2v_mode,
+        prompt_template=(
+            PROMPT_TEMPLATE[args.prompt_template]
+            if args.prompt_template is not None
+            else None
+        ),
+        prompt_template_video=(
+            PROMPT_TEMPLATE[args.prompt_template_video]
+            if args.prompt_template_video is not None
+            else None
+        ),
+        hidden_state_skip_layer=args.hidden_state_skip_layer,
+        apply_final_norm=args.apply_final_norm,
+        reproduce=args.reproduce,
+        logger=logger,
+        device=device,
+        image_embed_interleave=image_embed_interleave
+    )
+    if args.text_encoder_2 is not None:
+        text_encoder_2 = TextEncoder(
+            text_encoder_type=args.text_encoder_2,
+            max_length=args.text_len_2,
+            text_encoder_precision=args.text_encoder_precision_2,
+            tokenizer_type=args.tokenizer_2,
+            reproduce=args.reproduce,
+            logger=logger,
+            device=device,
+        )
+    else:
+        text_encoder_2 = None
+
+    # ================== Define dtype and forward autocast ===============
+    target_dtype = None
+    autocast_enabled = False
+    if model_engine.bfloat16_enabled():
+        target_dtype = torch.bfloat16
+        autocast_enabled = True
+    elif model_engine.fp16_enabled():
+        target_dtype = torch.half
+        autocast_enabled = True
+
+    # ============================== Load dataset ==============================
+    if "video" in args.data_type:
+        video_dataset = VideoDataset(
+            data_jsons_path=args.data_jsons_path,
+            sample_n_frames=args.sample_n_frames,
+            sample_stride=args.sample_stride,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            uncond_p=args.uncond_p,
+            args=args,
+            logger=logger,
+        )
+        video_sampler = DistributedSampler(
+            video_dataset,
+            num_replicas=world_size,
+            rank=rank,
+            shuffle=True,
+            seed=args.global_seed,
+            drop_last=False,
+        )
+        video_batch_sampler = None
+        video_loader = DataLoader(
+            video_dataset,
+            batch_size=video_micro_batch_size[0],
+            shuffle=False,
+            sampler=video_sampler,
+            batch_sampler=video_batch_sampler,
+            num_workers=args.num_workers,
+            pin_memory=True,
+            drop_last=False,
+            prefetch_factor=None if args.num_workers == 0 else args.prefetch_factor,
+            worker_init_fn=set_worker_seed_builder(rank),
+            persistent_workers=True,
+        )
+        num_video_samples = len(video_dataset)
+    else:
+        video_dataset = None
+        video_loader = None
+        num_video_samples = 0
+
+    loader = video_loader
+
+    # ============================= Print key info =============================
+    print(f"[{rank}] Worker ready.")
+    dist.barrier()
+    main_loader = video_loader
+
+    try:
+        iters_per_epoch = len(main_loader) // grad_accu_steps
+    except NotImplementedError:
+        iters_per_epoch = 0
+    except TypeError:
+        iters_per_epoch = 0
+
+    params_count = model.params_count()
+    logger.info("****************************** Running training ******************************")
+    logger.info(f"  Number GPUs:               {world_size}")
+    logger.info(f"  Training video samples(total):   {num_video_samples:,}")
+    for k, v in params_count.items():
+        logger.info(f"  Number {k} parameters:   {v:,}")
+    logger.info(f"  Number trainable params:   {sum(p.numel() for p in get_trainable_params(model, args)):,}")
+    logger.info("------------------------------------------------------------------------------")
+    logger.info(f"  Iters per epoch:           {iters_per_epoch:,}")
+    logger.info(f"  Updates per epoch:         {iters_per_epoch // grad_accu_steps:,}")
+
+    logger.info(f"  Batch size per device:     {video_micro_batch_size}")
+    logger.info(f"  Batch size all device:     {global_batch_size:}")
+
+    logger.info(f"  Gradient Accu steps:       {args.gradient_accumulation_steps}")
+    logger.info(f"  Training epochs:           {ss.epoch}/{args.epochs}")
+    logger.info(f"  Training total steps:      {ss.update_steps:,}/{args.max_training_steps:,}")
+    logger.info("------------------------------------------------------------------------------")
+
+    logger.info(f"  Path type:                 {args.flow_path_type}")
+    logger.info(f"  Predict type:              {args.flow_predict_type}")
+    logger.info(f"  Loss weight:               {args.flow_loss_weight}")
+    logger.info(f"  Flow reverse:              {args.flow_reverse}")
+    logger.info(f"  Flow shift:                {args.flow_shift}")
+    logger.info(f"  Train eps:                 {args.flow_train_eps}")
+    logger.info(f"  Sample eps:                {args.flow_sample_eps}")
+    logger.info(f"  Timestep type:             {args.flow_snr_type}")
+
+    logger.info("------------------------------------------------------------------------------")
+    logger.info(f"  Main model precision:      {args.precision}")
+    logger.info("------------------------------------------------------------------------------")
+    logger.info(f"  VAE:                       {args.vae} ({args.vae_precision}) - {vae_path}")
+    logger.info(f"  Text encoder:              {text_encoder}")
+    if text_encoder_2 is not None:
+        logger.info(f"  Text encoder 2:            {text_encoder_2}")
+    logger.info(f"  Experiment directory:      {ckpt_dir}")
+    logger.info("*******************************************************************************")
+
+    # ============================= Start training =============================
+    model_engine.train()
+
+    if args.init_save:
+        save_checkpoint(args, rank, logger, model_engine, ema, ss, ckpt_dir)
+
+    # Training loop
+    start_epoch = ss.epoch
+    finished = False
+
+    ss.current_run_update_steps = 0
+    for epoch in range(start_epoch, args.epochs):
+
+        if video_dataset is not None:
+            logger.info(f"Start video random shuffle(seed={args.global_seed + epoch})")
+            video_sampler.set_epoch(epoch)  # epoch start from 1
+            logger.info(f"End of video random shuffle")
+
+        logger.info(f"Beginning epoch {epoch}...")
+        with profiler_context(
+            args.profile, exp_dir, worker_name=f"Rank_{rank}"
+        ) as prof:
+            # Define cycle states, which accumulate the training information between log_steps.
+            cs = CycleStates()
+            start_time = time.time()
+
+            for batch_idx, batch in enumerate(loader):
+                # broadcast a zero size tensor to indicate starting of step
+                start_flag_tensor = torch.cuda.FloatTensor([])
+                if torch.distributed.is_initialized():
+                    torch.distributed.broadcast(start_flag_tensor, 0, async_op=True)
+
+                # main diff
+                (
+                    latents,
+                    model_kwargs,
+                    n_tokens,
+                    cond_latents,
+                ) = prepare_model_inputs(
+                    args,
+                    batch,
+                    device,
+                    model,
+                    vae,
+                    text_encoder,
+                    text_encoder_2,
+                    rope_theta_rescale_factor=args.rope_theta_rescale_factor,
+                    rope_interpolation_factor=args.rope_interpolation_factor,
+                )
+                cur_batch_size = latents.shape[0]
+
+                cur_anchor_size = max(args.video_size)
+
+                # A forward-backward step
+                with torch.autocast(
+                    device_type="cuda", dtype=target_dtype, enabled=autocast_enabled
+                ):
+                    _, loss_dict = denoiser.training_losses(
+                        model_engine,
+                        latents,
+                        model_kwargs,
+                        n_tokens=n_tokens,
+                        i2v_mode=args.i2v_mode,
+                        cond_latents=cond_latents,
+                        args=args,
+                    )
+                loss = loss_dict["loss"].mean()
+                model_engine.backward(loss)
+
+                # Update model parameters at the step of gradient accumulation.
+                model_engine.step(lr_kwargs={"last_batch_iteration": ss.update_steps})
+
+                # Update accumulated states
+                ss.add(
+                    train_steps=1,
+                    epoch_train_steps=1,
+                    consumed_samples_per_dp=cur_batch_size,
+                )
+                ss.add(consumed_video_samples_per_dp=cur_batch_size)
+
+                # We enable `is_update_step` if the current step is the gradient accumulation boundary.
+                is_update_step = ss.train_steps % grad_accu_steps == 0
+                if is_update_step:
+                    ss.add(
+                        update_steps=1, epoch_update_steps=1, current_run_update_steps=1
+                    )
+
+                if ss.update_steps >= args.max_training_steps:
+                    # Enter stopping routine if max steps reached after this step.
+                    finished = True
+
+                # Log training information:
+                cs.add(
+                    log_steps=1,
+                    running_loss=loss.item(),
+                    running_samples=cur_batch_size,
+                    running_tokens=cur_batch_size * n_tokens,
+                    running_grad_norm=0,
+                )
+                cs.add(running_video_samples=cur_batch_size)
+
+                cs.running_loss_dict[cur_anchor_size] += loss.item()
+                cs.log_steps_dict[cur_anchor_size] += 1
+                if is_update_step and ss.update_steps % args.log_every == 0:
+                    # Reduce loss history over all processes:
+                    avg_loss = (
+                        all_gather_sum(cs.running_loss / cs.log_steps, device)
+                        / world_size
+                    )
+                    avg_grad_norm = (
+                        all_gather_sum(cs.running_grad_norm / cs.log_steps, device)
+                        / world_size
+                    )
+                    cum_samples = all_gather_sum(cs.running_samples, device)
+                    cum_video_samples = all_gather_sum(cs.running_video_samples, device)
+                    cum_tokens = all_gather_sum(cs.running_tokens, device)
+                    # Measure training speed:
+                    torch.cuda.synchronize()
+                    end_time = time.time()
+                    steps_per_sec = (
+                        cs.log_steps / (end_time - start_time) / grad_accu_steps
+                    )
+                    samples_per_sec = cum_samples / (end_time - start_time)
+                    sec_per_step = (end_time - start_time) / cs.log_steps
+                    ss.add(
+                        consumed_samples_total=cum_samples,
+                        consumed_video_samples_total=cum_video_samples,
+                        consumed_tokens_total=cum_tokens,
+                        consumed_computations_attn=6
+                        * params_count["attn+mlp"]
+                        * cum_tokens
+                        / C_SCALE,
+                        consumed_computations_total=6
+                        * params_count["total"]
+                        * cum_tokens
+                        / C_SCALE,
+                    )
+                    log_events = [
+                        f"Train Loss: {avg_loss:.4f}",
+                        f"Grad Norm: {avg_grad_norm:.4f}",
+                        f"Lr: {opt.param_groups[0]['lr']:.6g}",
+                        f"Sec/Step: {sec_per_step:.2f}, "
+                        f"Steps/Sec: {steps_per_sec:.2f}",
+                        f"Samples/Sec: {int(samples_per_sec):d}",
+                        f"Consumed Samples: {ss.consumed_samples_total:,}",
+                        f"Consumed Video Samples: {ss.consumed_video_samples_total:,}",
+                        f"Consumed Tokens: {ss.consumed_tokens_total:,}",
+                    ]
+                    summary_events = [
+                        ("Train/Steps/train_loss", avg_loss, ss.update_steps),
+                        ("Train/Steps/grad_norm", avg_grad_norm, ss.update_steps),
+                        ("Train/Steps/steps_per_sec", steps_per_sec, ss.update_steps),
+                        (
+                            "Train/Steps/samples_per_sec",
+                            int(samples_per_sec),
+                            ss.update_steps,
+                        ),
+                        ("Train/Tokens/train_loss", avg_loss, ss.consumed_tokens_total),
+                        (
+                            "Train/ComputationsAttn/train_loss",
+                            avg_loss,
+                            ss.consumed_computations_attn,
+                        ),
+                        (
+                            "Train/ComputationsTotal/train_loss",
+                            avg_loss,
+                            ss.consumed_computations_total,
+                        ),
+                    ]
+                    # Log the training information to the logger.
+                    logger.info(
+                        f"(step={ss.update_steps:07d}) " + ", ".join(log_events)
+                    )
+                    if model_engine.monitor.enabled and rank == 0:
+                        model_engine.monitor.write_events(summary_events)
+
+                    # Reset monitoring variables:
+                    cs.reset()
+                    start_time = time.time()
+
+                # Save checkpoint:
+                if (is_update_step and ss.update_steps % args.ckpt_every == 0) or (
+                    finished and args.final_save
+                ):
+                    if args.use_lora:
+                        if rank == 0:
+                            output_dir = os.path.join(
+                                ckpt_dir, f"global_step{ss.update_steps}"
+                            )
+                            os.makedirs(output_dir, exist_ok=True)
+
+                            lora_kohya_state_dict = get_module_kohya_state_dict(
+                                model, "Hunyuan_video_I2V_lora", dtype=torch.bfloat16
+                            )
+                            save_file(
+                                lora_kohya_state_dict,
+                                f"{output_dir}/pytorch_lora_kohaya_weights.safetensors",
+                            )
+                    else:
+                        save_checkpoint(
+                            args, rank, logger, model_engine, ema, ss, ckpt_dir
+                        )
+
+                if prof:
+                    prof.step()
+
+                if finished:
+                    logger.info(
+                        f"Finished and breaking loop at step={ss.update_steps}."
+                    )
+                    break
+
+            if finished:
+                logger.info(f"Finished and breaking loop at epoch={epoch}.")
+                break
+
+            # Reset epoch states
+            ss.epoch += 1
+            ss.epoch_train_steps = 0
+            ss.epoch_update_steps = 0
+
+    logger.info("Training Finished!")
+
+
+if __name__ == "__main__":
+    main(parse_args(mode="train"))