diff --git a/README.md b/README.md
index ff730c08..89487380 100644
--- a/README.md
+++ b/README.md
@@ -8,12 +8,13 @@ An easy and efficient system for video generation
### Latest News 🔥
+- [2024/09] Support [Vchitect-2.0](https://github.com/Vchitect/Vchitect-2.0) and [Open-Sora-Plan v1.2.0](https://github.com/PKU-YuanGroup/Open-Sora-Plan).
- [2024/08] 🔥 Evole from [OpenDiT](https://github.com/NUS-HPC-AI-Lab/VideoSys/tree/v1.0.0) to VideoSys: An easy and efficient system for video generation.
-- [2024/08] 🔥 Release PAB paper: [Real-Time Video Generation with Pyramid Attention Broadcast](https://arxiv.org/abs/2408.12588).
-- [2024/06] Propose Pyramid Attention Broadcast (PAB)[[paper](https://arxiv.org/abs/2408.12588)][[blog](https://oahzxl.github.io/PAB/)][[doc](./docs/pab.md)], the first approach to achieve real-time DiT-based video generation, delivering negligible quality loss without requiring any training.
+- [2024/08] 🔥 Release PAB paper: [Real-Time Video Generation with Pyramid Attention Broadcast](https://arxiv.org/abs/2408.12588).
+- [2024/06] 🔥 Propose Pyramid Attention Broadcast (PAB)[[paper](https://arxiv.org/abs/2408.12588)][[blog](https://oahzxl.github.io/PAB/)][[doc](./docs/pab.md)], the first approach to achieve real-time DiT-based video generation, delivering negligible quality loss without requiring any training.
- [2024/06] Support [Open-Sora-Plan](https://github.com/PKU-YuanGroup/Open-Sora-Plan) and [Latte](https://github.com/Vchitect/Latte).
-- [2024/03] Propose Dynamic Sequence Parallel (DSP)[[paper](https://arxiv.org/abs/2403.10266)][[doc](./docs/dsp.md)], achieves **3x** speed for training and **2x** speed for inference in Open-Sora compared with sota sequence parallelism.
-- [2024/03] Support [Open-Sora: Democratizing Efficient Video Production for All](https://github.com/hpcaitech/Open-Sora).
+- [2024/03] 🔥 Propose Dynamic Sequence Parallel (DSP)[[paper](https://arxiv.org/abs/2403.10266)][[doc](./docs/dsp.md)], achieves **3x** speed for training and **2x** speed for inference in Open-Sora compared with sota sequence parallelism.
+- [2024/03] Support [Open-Sora](https://github.com/hpcaitech/Open-Sora).
- [2024/02] 🎉 Release [OpenDiT](https://github.com/NUS-HPC-AI-Lab/VideoSys/tree/v1.0.0): An Easy, Fast and Memory-Efficent System for DiT Training and Inference.
# About
diff --git a/examples/open_sora_plan/sample.py b/examples/open_sora_plan/sample.py
index f40431f8..2916c7fc 100644
--- a/examples/open_sora_plan/sample.py
+++ b/examples/open_sora_plan/sample.py
@@ -2,9 +2,10 @@
def run_base():
- # num frames: 65 or 221
+ # open-sora-plan v1.2.0
+ # transformer_type (len, res): 93x480p 93x720p 29x480p 29x720p
# change num_gpus for multi-gpu inference
- config = OpenSoraPlanConfig(num_frames=65, num_gpus=1)
+ config = OpenSoraPlanConfig(version="v120", transformer_type="93x480p", num_gpus=1)
engine = VideoSysEngine(config)
prompt = "Sunset over the sea."
@@ -12,7 +13,7 @@ def run_base():
video = engine.generate(
prompt=prompt,
guidance_scale=7.5,
- num_inference_steps=150,
+ num_inference_steps=100,
seed=-1,
).video[0]
engine.save_video(video, f"./outputs/{prompt}.mp4")
@@ -36,7 +37,26 @@ def run_pab():
engine.save_video(video, f"./outputs/{prompt}.mp4")
+def run_v110():
+ # open-sora-plan v1.1.0
+ # transformer_type: 65x512x512 or 221x512x512
+ # change num_gpus for multi-gpu inference
+ config = OpenSoraPlanConfig(version="v110", transformer_type="65x512x512", num_gpus=1)
+ engine = VideoSysEngine(config)
+
+ prompt = "Sunset over the sea."
+ # seed=-1 means random seed. >0 means fixed seed.
+ video = engine.generate(
+ prompt=prompt,
+ guidance_scale=7.5,
+ num_inference_steps=150,
+ seed=-1,
+ ).video[0]
+ engine.save_video(video, f"./outputs/{prompt}.mp4")
+
+
if __name__ == "__main__":
run_base()
# run_low_mem()
# run_pab()
+ # run_v110()
diff --git a/videosys/models/autoencoders/autoencoder_kl_open_sora_plan.py b/videosys/models/autoencoders/autoencoder_kl_open_sora_plan_v110.py
similarity index 99%
rename from videosys/models/autoencoders/autoencoder_kl_open_sora_plan.py
rename to videosys/models/autoencoders/autoencoder_kl_open_sora_plan_v110.py
index 4368a740..30338fc0 100644
--- a/videosys/models/autoencoders/autoencoder_kl_open_sora_plan.py
+++ b/videosys/models/autoencoders/autoencoder_kl_open_sora_plan_v110.py
@@ -124,7 +124,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
last_ckpt_file = ckpt_files[-1]
config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
model = cls.from_config(config_file)
- print("init from {}".format(last_ckpt_file))
+ # print("init from {}".format(last_ckpt_file))
model.init_from_ckpt(last_ckpt_file)
return model
else:
@@ -778,7 +778,7 @@ def disable_tiling(self):
def init_from_ckpt(self, path, ignore_keys=list(), remove_loss=False):
sd = torch.load(path, map_location="cpu")
- print("init from " + path)
+ # print("init from " + path)
if "state_dict" in sd:
sd = sd["state_dict"]
keys = list(sd.keys())
diff --git a/videosys/models/autoencoders/autoencoder_kl_open_sora_plan_v120.py b/videosys/models/autoencoders/autoencoder_kl_open_sora_plan_v120.py
new file mode 100644
index 00000000..0bb4aa8c
--- /dev/null
+++ b/videosys/models/autoencoders/autoencoder_kl_open_sora_plan_v120.py
@@ -0,0 +1,1139 @@
+# Adapted from Open-Sora-Plan
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# Open-Sora-Plan: https://github.com/PKU-YuanGroup/Open-Sora-Plan
+# --------------------------------------------------------
+
+import glob
+import os
+from copy import deepcopy
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import ConfigMixin, ModelMixin
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+from torchvision.transforms import Lambda
+
+npu_config = None
+
+
+def cast_tuple(t, length=1):
+ return t if isinstance(t, tuple) or isinstance(t, list) else ((t,) * length)
+
+
+class Block(nn.Module):
+ def __init__(self, *args, **kwargs) -> None:
+ super().__init__(*args, **kwargs)
+
+
+class CausalConv3d(nn.Module):
+ def __init__(
+ self, chan_in, chan_out, kernel_size: Union[int, Tuple[int, int, int]], init_method="random", **kwargs
+ ):
+ super().__init__()
+ self.kernel_size = cast_tuple(kernel_size, 3)
+ self.time_kernel_size = self.kernel_size[0]
+ self.chan_in = chan_in
+ self.chan_out = chan_out
+ stride = kwargs.pop("stride", 1)
+ padding = kwargs.pop("padding", 0)
+ padding = list(cast_tuple(padding, 3))
+ padding[0] = 0
+ stride = cast_tuple(stride, 3)
+ self.conv = nn.Conv3d(chan_in, chan_out, self.kernel_size, stride=stride, padding=padding)
+ self.pad = nn.ReplicationPad2d((0, 0, self.time_kernel_size - 1, 0))
+ self._init_weights(init_method)
+
+ def _init_weights(self, init_method):
+ torch.tensor(self.kernel_size)
+ if init_method == "avg":
+ assert self.kernel_size[1] == 1 and self.kernel_size[2] == 1, "only support temporal up/down sample"
+ assert self.chan_in == self.chan_out, "chan_in must be equal to chan_out"
+ weight = torch.zeros((self.chan_out, self.chan_in, *self.kernel_size))
+
+ eyes = torch.concat(
+ [
+ torch.eye(self.chan_in).unsqueeze(-1) * 1 / 3,
+ torch.eye(self.chan_in).unsqueeze(-1) * 1 / 3,
+ torch.eye(self.chan_in).unsqueeze(-1) * 1 / 3,
+ ],
+ dim=-1,
+ )
+ weight[:, :, :, 0, 0] = eyes
+
+ self.conv.weight = nn.Parameter(
+ weight,
+ requires_grad=True,
+ )
+ elif init_method == "zero":
+ self.conv.weight = nn.Parameter(
+ torch.zeros((self.chan_out, self.chan_in, *self.kernel_size)),
+ requires_grad=True,
+ )
+ if self.conv.bias is not None:
+ nn.init.constant_(self.conv.bias, 0)
+
+ def forward(self, x):
+ if npu_config is not None and npu_config.on_npu:
+ x_dtype = x.dtype
+ first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1)) # b c t h w
+ x = torch.concatenate((first_frame_pad, x), dim=2) # 3 + 16
+ return npu_config.run_conv3d(self.conv, x, x_dtype)
+ else:
+ # 1 + 16 16 as video, 1 as image
+ first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1)) # b c t h w
+ x = torch.concatenate((first_frame_pad, x), dim=2) # 3 + 16
+ return self.conv(x)
+
+
+def nonlinearity(x):
+ return x * torch.sigmoid(x)
+
+
+def video_to_image(func):
+ def wrapper(self, x, *args, **kwargs):
+ if x.dim() == 5:
+ t = x.shape[2]
+ x = rearrange(x, "b c t h w -> (b t) c h w")
+ x = func(self, x, *args, **kwargs)
+ x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+ return x
+
+ return wrapper
+
+
+class Conv2d(nn.Conv2d):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: Union[int, Tuple[int]] = 3,
+ stride: Union[int, Tuple[int]] = 1,
+ padding: Union[str, int, Tuple[int]] = 0,
+ dilation: Union[int, Tuple[int]] = 1,
+ groups: int = 1,
+ bias: bool = True,
+ padding_mode: str = "zeros",
+ device=None,
+ dtype=None,
+ ) -> None:
+ super().__init__(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride,
+ padding,
+ dilation,
+ groups,
+ bias,
+ padding_mode,
+ device,
+ dtype,
+ )
+
+ @video_to_image
+ def forward(self, x):
+ return super().forward(x)
+
+
+def Normalize(in_channels, num_groups=32):
+ return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class VideoBaseAE(ModelMixin, ConfigMixin):
+ config_name = "config.json"
+
+ def __init__(self, *args, **kwargs) -> None:
+ super().__init__(*args, **kwargs)
+
+ def encode(self, x: torch.Tensor, *args, **kwargs):
+ pass
+
+ def decode(self, encoding: torch.Tensor, *args, **kwargs):
+ pass
+
+ @property
+ def num_training_steps(self) -> int:
+ """Total training steps inferred from datamodule and devices."""
+ if self.trainer.max_steps:
+ return self.trainer.max_steps
+
+ limit_batches = self.trainer.limit_train_batches
+ batches = len(self.train_dataloader())
+ batches = min(batches, limit_batches) if isinstance(limit_batches, int) else int(limit_batches * batches)
+
+ num_devices = max(1, self.trainer.num_gpus, self.trainer.num_processes)
+ if self.trainer.tpu_cores:
+ num_devices = max(num_devices, self.trainer.tpu_cores)
+
+ effective_accum = self.trainer.accumulate_grad_batches * num_devices
+ return (batches // effective_accum) * self.trainer.max_epochs
+
+ @classmethod
+ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
+ ckpt_files = glob.glob(os.path.join(pretrained_model_name_or_path, "*.ckpt"))
+ if not ckpt_files:
+ ckpt_file = hf_hub_download(pretrained_model_name_or_path, subfolder="vae", filename="checkpoint.ckpt")
+ config_file = hf_hub_download(pretrained_model_name_or_path, subfolder="vae", filename="config.json")
+ else:
+ ckpt_file = ckpt_files[-1]
+ config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
+
+ # Adapt to checkpoint
+ model = cls.from_config(config_file)
+ model.init_from_ckpt(ckpt_file)
+ return model
+
+
+class DiagonalGaussianDistribution(object):
+ def __init__(self, parameters, deterministic=False):
+ self.parameters = parameters
+ self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+ 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).to(device=self.parameters.device)
+
+ def sample(self):
+ x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
+ return x
+
+ def kl(self, other=None):
+ if self.deterministic:
+ return torch.Tensor([0.0])
+ else:
+ if other is None:
+ return 0.5 * torch.sum(torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, dim=[1, 2, 3])
+ 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=[1, 2, 3],
+ )
+
+ def nll(self, sample, dims=[1, 2, 3]):
+ 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):
+ return self.mean
+
+
+class ResnetBlock2D(Block):
+ def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, dropout):
+ super().__init__()
+ self.in_channels = in_channels
+ self.out_channels = in_channels if out_channels is None else out_channels
+ self.use_conv_shortcut = conv_shortcut
+
+ self.norm1 = Normalize(in_channels)
+ self.conv1 = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+ self.norm2 = Normalize(out_channels)
+ self.dropout = torch.nn.Dropout(dropout)
+ self.conv2 = torch.nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ self.conv_shortcut = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+ else:
+ self.nin_shortcut = torch.nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+ @video_to_image
+ def forward(self, x):
+ h = x
+ h = self.norm1(h)
+ h = nonlinearity(h)
+ h = self.conv1(h)
+ h = self.norm2(h)
+ h = nonlinearity(h)
+ h = self.dropout(h)
+ h = self.conv2(h)
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ x = self.conv_shortcut(x)
+ else:
+ x = self.nin_shortcut(x)
+ x = x + h
+ return x
+
+
+class ResnetBlock3D(Block):
+ def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, dropout):
+ super().__init__()
+ self.in_channels = in_channels
+ self.out_channels = in_channels if out_channels is None else out_channels
+ self.use_conv_shortcut = conv_shortcut
+
+ self.norm1 = Normalize(in_channels)
+ self.conv1 = CausalConv3d(in_channels, out_channels, 3, padding=1)
+ self.norm2 = Normalize(out_channels)
+ self.dropout = torch.nn.Dropout(dropout)
+ self.conv2 = CausalConv3d(out_channels, out_channels, 3, padding=1)
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ self.conv_shortcut = CausalConv3d(in_channels, out_channels, 3, padding=1)
+ else:
+ self.nin_shortcut = CausalConv3d(in_channels, out_channels, 1, padding=0)
+
+ def forward(self, x):
+ h = x
+ if npu_config is None:
+ h = self.norm1(h)
+ else:
+ h = npu_config.run_group_norm(self.norm1, h)
+ h = nonlinearity(h)
+ h = self.conv1(h)
+ if npu_config is None:
+ h = self.norm2(h)
+ else:
+ h = npu_config.run_group_norm(self.norm2, h)
+ h = nonlinearity(h)
+ h = self.dropout(h)
+ h = self.conv2(h)
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ x = self.conv_shortcut(x)
+ else:
+ x = self.nin_shortcut(x)
+ return x + h
+
+
+class SpatialUpsample2x(Block):
+ def __init__(
+ self,
+ chan_in,
+ chan_out,
+ kernel_size: Union[int, Tuple[int]] = (3, 3),
+ stride: Union[int, Tuple[int]] = (1, 1),
+ unup=False,
+ ):
+ super().__init__()
+ self.chan_in = chan_in
+ self.chan_out = chan_out
+ self.kernel_size = kernel_size
+ self.unup = unup
+ self.conv = CausalConv3d(self.chan_in, self.chan_out, (1,) + self.kernel_size, stride=(1,) + stride, padding=1)
+
+ def forward(self, x):
+ if not self.unup:
+ t = x.shape[2]
+ x = rearrange(x, "b c t h w -> b (c t) h w")
+ x = F.interpolate(x, scale_factor=(2, 2), mode="nearest")
+ x = rearrange(x, "b (c t) h w -> b c t h w", t=t)
+ x = self.conv(x)
+ return x
+
+
+class Spatial2xTime2x3DUpsample(Block):
+ def __init__(self, in_channels, out_channels):
+ super().__init__()
+ self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=1)
+
+ def forward(self, x):
+ if x.size(2) > 1:
+ x, x_ = x[:, :, :1], x[:, :, 1:]
+ x_ = F.interpolate(x_, scale_factor=(2, 2, 2), mode="trilinear")
+ x = F.interpolate(x, scale_factor=(1, 2, 2), mode="trilinear")
+ x = torch.concat([x, x_], dim=2)
+ else:
+ x = F.interpolate(x, scale_factor=(1, 2, 2), mode="trilinear")
+ return self.conv(x)
+
+
+class AttnBlock3DFix(nn.Module):
+ """
+ Thanks to https://github.com/PKU-YuanGroup/Open-Sora-Plan/pull/172.
+ """
+
+ def __init__(self, in_channels):
+ super().__init__()
+ self.in_channels = in_channels
+
+ self.norm = Normalize(in_channels)
+ self.q = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
+ self.k = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
+ self.v = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
+ self.proj_out = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
+
+ def forward(self, x):
+ h_ = x
+ if npu_config is None:
+ h_ = self.norm(h_)
+ else:
+ h_ = npu_config.run_group_norm(self.norm, h_)
+ q = self.q(h_)
+ k = self.k(h_)
+ v = self.v(h_)
+
+ # compute attention
+ # q: (b c t h w) -> (b t c h w) -> (b*t c h*w) -> (b*t h*w c)
+ b, c, t, h, w = q.shape
+ q = q.permute(0, 2, 1, 3, 4)
+ q = q.reshape(b * t, c, h * w)
+ q = q.permute(0, 2, 1)
+
+ # k: (b c t h w) -> (b t c h w) -> (b*t c h*w)
+ k = k.permute(0, 2, 1, 3, 4)
+ k = k.reshape(b * t, c, h * w)
+
+ # w: (b*t hw hw)
+ w_ = torch.bmm(q, k)
+ w_ = w_ * (int(c) ** (-0.5))
+ w_ = torch.nn.functional.softmax(w_, dim=2)
+
+ # attend to values
+ # v: (b c t h w) -> (b t c h w) -> (bt c hw)
+ # w_: (bt hw hw) -> (bt hw hw)
+ v = v.permute(0, 2, 1, 3, 4)
+ v = v.reshape(b * t, c, h * w)
+ w_ = w_.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q)
+ h_ = torch.bmm(v, w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+
+ # h_: (b*t c hw) -> (b t c h w) -> (b c t h w)
+ h_ = h_.reshape(b, t, c, h, w)
+ h_ = h_.permute(0, 2, 1, 3, 4)
+
+ h_ = self.proj_out(h_)
+
+ return x + h_
+
+
+class Spatial2xTime2x3DDownsample(Block):
+ def __init__(self, in_channels, out_channels):
+ super().__init__()
+ self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=0, stride=2)
+
+ def forward(self, x):
+ pad = (0, 1, 0, 1, 0, 0)
+ x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+ x = self.conv(x)
+ return x
+
+
+class Downsample(Block):
+ def __init__(self, in_channels, out_channels, undown=False):
+ super().__init__()
+ self.with_conv = True
+ self.undown = undown
+ if self.with_conv:
+ # no asymmetric padding in torch conv, must do it ourselves
+ if self.undown:
+ self.conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+ else:
+ self.conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=0)
+
+ @video_to_image
+ def forward(self, x):
+ if self.with_conv:
+ if self.undown:
+ if npu_config is not None and npu_config.on_npu:
+ x_dtype = x.dtype
+ x = x.to(npu_config.replaced_type)
+ x = npu_config.run_conv3d(self.conv, x, x_dtype)
+ else:
+ x = self.conv(x)
+ else:
+ pad = (0, 1, 0, 1)
+ if npu_config is not None and npu_config.on_npu:
+ x_dtype = x.dtype
+ x = x.to(npu_config.replaced_type)
+ x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+ x = npu_config.run_conv3d(self.conv, x, x_dtype)
+ else:
+ x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+ x = self.conv(x)
+ else:
+ x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+ return x
+
+
+class ResnetBlock3D_GC(Block):
+ def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, dropout):
+ super().__init__()
+ self.in_channels = in_channels
+ self.out_channels = in_channels if out_channels is None else out_channels
+ self.use_conv_shortcut = conv_shortcut
+
+ self.norm1 = Normalize(in_channels)
+ self.conv1 = CausalConv3d(in_channels, out_channels, 3, padding=1)
+ self.norm2 = Normalize(out_channels)
+ self.dropout = torch.nn.Dropout(dropout)
+ self.conv2 = CausalConv3d(out_channels, out_channels, 3, padding=1)
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ self.conv_shortcut = CausalConv3d(in_channels, out_channels, 3, padding=1)
+ else:
+ self.nin_shortcut = CausalConv3d(in_channels, out_channels, 1, padding=0)
+
+ def forward(self, x):
+ return checkpoint(self._forward, x, use_reentrant=True)
+
+ def _forward(self, x):
+ h = x
+ h = self.norm1(h)
+ h = nonlinearity(h)
+ h = self.conv1(h)
+ h = self.norm2(h)
+ h = nonlinearity(h)
+ h = self.dropout(h)
+ h = self.conv2(h)
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ x = self.conv_shortcut(x)
+ else:
+ x = self.nin_shortcut(x)
+ return x + h
+
+
+def resolve_str_to_obj(str_val, append=True):
+ return globals()[str_val]
+
+
+class Encoder(nn.Module):
+ def __init__(
+ self,
+ z_channels: int,
+ hidden_size: int,
+ hidden_size_mult: Tuple[int] = (1, 2, 4, 4),
+ attn_resolutions: Tuple[int] = (16,),
+ conv_in="Conv2d",
+ conv_out="CasualConv3d",
+ attention="AttnBlock",
+ resnet_blocks=(
+ "ResnetBlock2D",
+ "ResnetBlock2D",
+ "ResnetBlock2D",
+ "ResnetBlock3D",
+ ),
+ spatial_downsample=(
+ "Downsample",
+ "Downsample",
+ "Downsample",
+ "",
+ ),
+ temporal_downsample=("", "", "TimeDownsampleRes2x", ""),
+ mid_resnet="ResnetBlock3D",
+ dropout: float = 0.0,
+ resolution: int = 256,
+ num_res_blocks: int = 2,
+ double_z: bool = True,
+ ) -> None:
+ super().__init__()
+ assert len(resnet_blocks) == len(hidden_size_mult), print(hidden_size_mult, resnet_blocks)
+ # ---- Config ----
+ self.num_resolutions = len(hidden_size_mult)
+ self.resolution = resolution
+ self.num_res_blocks = num_res_blocks
+
+ # ---- In ----
+ self.conv_in = resolve_str_to_obj(conv_in)(3, hidden_size, kernel_size=3, stride=1, padding=1)
+
+ # ---- Downsample ----
+ curr_res = resolution
+ in_ch_mult = (1,) + tuple(hidden_size_mult)
+ self.in_ch_mult = in_ch_mult
+ self.down = nn.ModuleList()
+ for i_level in range(self.num_resolutions):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_in = hidden_size * in_ch_mult[i_level]
+ block_out = hidden_size * hidden_size_mult[i_level]
+ for i_block in range(self.num_res_blocks):
+ block.append(
+ resolve_str_to_obj(resnet_blocks[i_level])(
+ in_channels=block_in,
+ out_channels=block_out,
+ dropout=dropout,
+ )
+ )
+ block_in = block_out
+ if curr_res in attn_resolutions:
+ attn.append(resolve_str_to_obj(attention)(block_in))
+ down = nn.Module()
+ down.block = block
+ down.attn = attn
+ if spatial_downsample[i_level]:
+ down.downsample = resolve_str_to_obj(spatial_downsample[i_level])(block_in, block_in)
+ curr_res = curr_res // 2
+ if temporal_downsample[i_level]:
+ down.time_downsample = resolve_str_to_obj(temporal_downsample[i_level])(block_in, block_in)
+ self.down.append(down)
+
+ # ---- Mid ----
+ self.mid = nn.Module()
+ self.mid.block_1 = resolve_str_to_obj(mid_resnet)(
+ in_channels=block_in,
+ out_channels=block_in,
+ dropout=dropout,
+ )
+ self.mid.attn_1 = resolve_str_to_obj(attention)(block_in)
+ self.mid.block_2 = resolve_str_to_obj(mid_resnet)(
+ in_channels=block_in,
+ out_channels=block_in,
+ dropout=dropout,
+ )
+ # ---- Out ----
+ self.norm_out = Normalize(block_in)
+ self.conv_out = resolve_str_to_obj(conv_out)(
+ block_in,
+ 2 * z_channels if double_z else z_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ def forward(self, x):
+ hs = [self.conv_in(x)]
+ for i_level in range(self.num_resolutions):
+ for i_block in range(self.num_res_blocks):
+ h = self.down[i_level].block[i_block](hs[-1])
+ if len(self.down[i_level].attn) > 0:
+ h = self.down[i_level].attn[i_block](h)
+ hs.append(h)
+ if hasattr(self.down[i_level], "downsample"):
+ hs.append(self.down[i_level].downsample(hs[-1]))
+ if hasattr(self.down[i_level], "time_downsample"):
+ hs_down = self.down[i_level].time_downsample(hs[-1])
+ hs.append(hs_down)
+
+ h = self.mid.block_1(h)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h)
+
+ if npu_config is None:
+ h = self.norm_out(h)
+ else:
+ h = npu_config.run_group_norm(self.norm_out, h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ return h
+
+
+class Decoder(nn.Module):
+ def __init__(
+ self,
+ z_channels: int,
+ hidden_size: int,
+ hidden_size_mult: Tuple[int] = (1, 2, 4, 4),
+ attn_resolutions: Tuple[int] = (16,),
+ conv_in="Conv2d",
+ conv_out="CasualConv3d",
+ attention="AttnBlock",
+ resnet_blocks=(
+ "ResnetBlock3D",
+ "ResnetBlock3D",
+ "ResnetBlock3D",
+ "ResnetBlock3D",
+ ),
+ spatial_upsample=(
+ "",
+ "SpatialUpsample2x",
+ "SpatialUpsample2x",
+ "SpatialUpsample2x",
+ ),
+ temporal_upsample=("", "", "", "TimeUpsampleRes2x"),
+ mid_resnet="ResnetBlock3D",
+ dropout: float = 0.0,
+ resolution: int = 256,
+ num_res_blocks: int = 2,
+ ):
+ super().__init__()
+ # ---- Config ----
+ self.num_resolutions = len(hidden_size_mult)
+ self.resolution = resolution
+ self.num_res_blocks = num_res_blocks
+
+ # ---- In ----
+ block_in = hidden_size * hidden_size_mult[self.num_resolutions - 1]
+ curr_res = resolution // 2 ** (self.num_resolutions - 1)
+ self.conv_in = resolve_str_to_obj(conv_in)(z_channels, block_in, kernel_size=3, padding=1)
+
+ # ---- Mid ----
+ self.mid = nn.Module()
+ self.mid.block_1 = resolve_str_to_obj(mid_resnet)(
+ in_channels=block_in,
+ out_channels=block_in,
+ dropout=dropout,
+ )
+ self.mid.attn_1 = resolve_str_to_obj(attention)(block_in)
+ self.mid.block_2 = resolve_str_to_obj(mid_resnet)(
+ in_channels=block_in,
+ out_channels=block_in,
+ dropout=dropout,
+ )
+
+ # ---- Upsample ----
+ self.up = nn.ModuleList()
+ for i_level in reversed(range(self.num_resolutions)):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_out = hidden_size * hidden_size_mult[i_level]
+ for i_block in range(self.num_res_blocks + 1):
+ block.append(
+ resolve_str_to_obj(resnet_blocks[i_level])(
+ in_channels=block_in,
+ out_channels=block_out,
+ dropout=dropout,
+ )
+ )
+ block_in = block_out
+ if curr_res in attn_resolutions:
+ attn.append(resolve_str_to_obj(attention)(block_in))
+ up = nn.Module()
+ up.block = block
+ up.attn = attn
+ if spatial_upsample[i_level]:
+ up.upsample = resolve_str_to_obj(spatial_upsample[i_level])(block_in, block_in)
+ curr_res = curr_res * 2
+ if temporal_upsample[i_level]:
+ up.time_upsample = resolve_str_to_obj(temporal_upsample[i_level])(block_in, block_in)
+ self.up.insert(0, up)
+
+ # ---- Out ----
+ self.norm_out = Normalize(block_in)
+ self.conv_out = resolve_str_to_obj(conv_out)(block_in, 3, kernel_size=3, padding=1)
+
+ def forward(self, z):
+ h = self.conv_in(z)
+ h = self.mid.block_1(h)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h)
+
+ for i_level in reversed(range(self.num_resolutions)):
+ for i_block in range(self.num_res_blocks + 1):
+ h = self.up[i_level].block[i_block](h)
+ if len(self.up[i_level].attn) > 0:
+ h = self.up[i_level].attn[i_block](h)
+ if hasattr(self.up[i_level], "upsample"):
+ h = self.up[i_level].upsample(h)
+ if hasattr(self.up[i_level], "time_upsample"):
+ h = self.up[i_level].time_upsample(h)
+ if npu_config is None:
+ h = self.norm_out(h)
+ else:
+ h = npu_config.run_group_norm(self.norm_out, h)
+ h = nonlinearity(h)
+ if npu_config is None:
+ h = self.conv_out(h)
+ else:
+ h_dtype = h.dtype
+ h = npu_config.run_conv3d(self.conv_out, h, h_dtype)
+ return h
+
+
+class CausalVAEModel(VideoBaseAE):
+ @register_to_config
+ def __init__(
+ self,
+ hidden_size: int = 128,
+ z_channels: int = 4,
+ hidden_size_mult: Tuple[int] = (1, 2, 4, 4),
+ attn_resolutions: Tuple[int] = [],
+ dropout: float = 0.0,
+ resolution: int = 256,
+ double_z: bool = True,
+ embed_dim: int = 4,
+ num_res_blocks: int = 2,
+ q_conv: str = "CausalConv3d",
+ encoder_conv_in="CausalConv3d",
+ encoder_conv_out="CausalConv3d",
+ encoder_attention="AttnBlock3D",
+ encoder_resnet_blocks=(
+ "ResnetBlock3D",
+ "ResnetBlock3D",
+ "ResnetBlock3D",
+ "ResnetBlock3D",
+ ),
+ encoder_spatial_downsample=(
+ "SpatialDownsample2x",
+ "SpatialDownsample2x",
+ "SpatialDownsample2x",
+ "",
+ ),
+ encoder_temporal_downsample=(
+ "",
+ "TimeDownsample2x",
+ "TimeDownsample2x",
+ "",
+ ),
+ encoder_mid_resnet="ResnetBlock3D",
+ decoder_conv_in="CausalConv3d",
+ decoder_conv_out="CausalConv3d",
+ decoder_attention="AttnBlock3D",
+ decoder_resnet_blocks=(
+ "ResnetBlock3D",
+ "ResnetBlock3D",
+ "ResnetBlock3D",
+ "ResnetBlock3D",
+ ),
+ decoder_spatial_upsample=(
+ "",
+ "SpatialUpsample2x",
+ "SpatialUpsample2x",
+ "SpatialUpsample2x",
+ ),
+ decoder_temporal_upsample=("", "", "TimeUpsample2x", "TimeUpsample2x"),
+ decoder_mid_resnet="ResnetBlock3D",
+ use_quant_layer: bool = True,
+ ) -> None:
+ super().__init__()
+
+ self.tile_sample_min_size = 256
+ self.tile_sample_min_size_t = 33
+ self.tile_latent_min_size = int(self.tile_sample_min_size / (2 ** (len(hidden_size_mult) - 1)))
+
+ # t_down_ratio = [i for i in encoder_temporal_downsample if len(i) > 0]
+ # self.tile_latent_min_size_t = int((self.tile_sample_min_size_t-1) / (2 ** len(t_down_ratio))) + 1
+ self.tile_latent_min_size_t = 16
+ self.tile_overlap_factor = 0.125
+ self.use_tiling = False
+
+ self.use_quant_layer = use_quant_layer
+
+ self.encoder = Encoder(
+ z_channels=z_channels,
+ hidden_size=hidden_size,
+ hidden_size_mult=hidden_size_mult,
+ attn_resolutions=attn_resolutions,
+ conv_in=encoder_conv_in,
+ conv_out=encoder_conv_out,
+ attention=encoder_attention,
+ resnet_blocks=encoder_resnet_blocks,
+ spatial_downsample=encoder_spatial_downsample,
+ temporal_downsample=encoder_temporal_downsample,
+ mid_resnet=encoder_mid_resnet,
+ dropout=dropout,
+ resolution=resolution,
+ num_res_blocks=num_res_blocks,
+ double_z=double_z,
+ )
+
+ self.decoder = Decoder(
+ z_channels=z_channels,
+ hidden_size=hidden_size,
+ hidden_size_mult=hidden_size_mult,
+ attn_resolutions=attn_resolutions,
+ conv_in=decoder_conv_in,
+ conv_out=decoder_conv_out,
+ attention=decoder_attention,
+ resnet_blocks=decoder_resnet_blocks,
+ spatial_upsample=decoder_spatial_upsample,
+ temporal_upsample=decoder_temporal_upsample,
+ mid_resnet=decoder_mid_resnet,
+ dropout=dropout,
+ resolution=resolution,
+ num_res_blocks=num_res_blocks,
+ )
+ if self.use_quant_layer:
+ quant_conv_cls = resolve_str_to_obj(q_conv)
+ self.quant_conv = quant_conv_cls(2 * z_channels, 2 * embed_dim, 1)
+ self.post_quant_conv = quant_conv_cls(embed_dim, z_channels, 1)
+
+ def get_encoder(self):
+ if self.use_quant_layer:
+ return [self.quant_conv, self.encoder]
+ return [self.encoder]
+
+ def get_decoder(self):
+ if self.use_quant_layer:
+ return [self.post_quant_conv, self.decoder]
+ return [self.decoder]
+
+ def encode(self, x):
+ if self.use_tiling and (
+ x.shape[-1] > self.tile_sample_min_size
+ or x.shape[-2] > self.tile_sample_min_size
+ or x.shape[-3] > self.tile_sample_min_size_t
+ ):
+ return self.tiled_encode(x)
+ h = self.encoder(x)
+ if self.use_quant_layer:
+ h = self.quant_conv(h)
+ posterior = DiagonalGaussianDistribution(h)
+ return posterior
+
+ def decode(self, z):
+ if self.use_tiling and (
+ z.shape[-1] > self.tile_latent_min_size
+ or z.shape[-2] > self.tile_latent_min_size
+ or z.shape[-3] > self.tile_latent_min_size_t
+ ):
+ return self.tiled_decode(z)
+ if self.use_quant_layer:
+ z = self.post_quant_conv(z)
+ dec = self.decoder(z)
+ return dec
+
+ def forward(self, input, sample_posterior=True):
+ posterior = self.encode(input)
+ if sample_posterior:
+ z = posterior.sample()
+ else:
+ z = posterior.mode()
+ dec = self.decode(z)
+ return dec, posterior
+
+ def on_train_start(self):
+ self.ema = deepcopy(self) if self.save_ema == True else None
+
+ def get_last_layer(self):
+ if hasattr(self.decoder.conv_out, "conv"):
+ return self.decoder.conv_out.conv.weight
+ else:
+ return self.decoder.conv_out.weight
+
+ def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+ blend_extent = min(a.shape[3], b.shape[3], 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[4], b.shape[4], 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 tiled_encode(self, x):
+ t = x.shape[2]
+ t_chunk_idx = [i for i in range(0, t, self.tile_sample_min_size_t - 1)]
+ if len(t_chunk_idx) == 1 and t_chunk_idx[0] == 0:
+ t_chunk_start_end = [[0, t]]
+ else:
+ t_chunk_start_end = [[t_chunk_idx[i], t_chunk_idx[i + 1] + 1] for i in range(len(t_chunk_idx) - 1)]
+ if t_chunk_start_end[-1][-1] > t:
+ t_chunk_start_end[-1][-1] = t
+ elif t_chunk_start_end[-1][-1] < t:
+ last_start_end = [t_chunk_idx[-1], t]
+ t_chunk_start_end.append(last_start_end)
+ moments = []
+ for idx, (start, end) in enumerate(t_chunk_start_end):
+ chunk_x = x[:, :, start:end]
+ if idx != 0:
+ moment = self.tiled_encode2d(chunk_x, return_moments=True)[:, :, 1:]
+ else:
+ moment = self.tiled_encode2d(chunk_x, return_moments=True)
+ moments.append(moment)
+ moments = torch.cat(moments, dim=2)
+ posterior = DiagonalGaussianDistribution(moments)
+ return posterior
+
+ def tiled_decode(self, x):
+ t = x.shape[2]
+ t_chunk_idx = [i for i in range(0, t, self.tile_latent_min_size_t - 1)]
+ if len(t_chunk_idx) == 1 and t_chunk_idx[0] == 0:
+ t_chunk_start_end = [[0, t]]
+ else:
+ t_chunk_start_end = [[t_chunk_idx[i], t_chunk_idx[i + 1] + 1] for i in range(len(t_chunk_idx) - 1)]
+ if t_chunk_start_end[-1][-1] > t:
+ t_chunk_start_end[-1][-1] = t
+ elif t_chunk_start_end[-1][-1] < t:
+ last_start_end = [t_chunk_idx[-1], t]
+ t_chunk_start_end.append(last_start_end)
+ dec_ = []
+ for idx, (start, end) in enumerate(t_chunk_start_end):
+ chunk_x = x[:, :, start:end]
+ if idx != 0:
+ dec = self.tiled_decode2d(chunk_x)[:, :, 1:]
+ else:
+ dec = self.tiled_decode2d(chunk_x)
+ dec_.append(dec)
+ dec_ = torch.cat(dec_, dim=2)
+ return dec_
+
+ def tiled_encode2d(self, x, return_moments=False):
+ 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 the image into 512x512 tiles and encode them separately.
+ rows = []
+ for i in range(0, x.shape[3], overlap_size):
+ row = []
+ for j in range(0, x.shape[4], overlap_size):
+ tile = x[
+ :,
+ :,
+ :,
+ i : i + self.tile_sample_min_size,
+ j : j + self.tile_sample_min_size,
+ ]
+ tile = self.encoder(tile)
+ if self.use_quant_layer:
+ 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=4))
+
+ moments = torch.cat(result_rows, dim=3)
+ posterior = DiagonalGaussianDistribution(moments)
+ if return_moments:
+ return moments
+ return posterior
+
+ def tiled_decode2d(self, z):
+ 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 64x64 tiles and decode them separately.
+ # The tiles have an overlap to avoid seams between tiles.
+ rows = []
+ for i in range(0, z.shape[3], overlap_size):
+ row = []
+ for j in range(0, z.shape[4], overlap_size):
+ tile = z[
+ :,
+ :,
+ :,
+ i : i + self.tile_latent_min_size,
+ j : j + self.tile_latent_min_size,
+ ]
+ if self.use_quant_layer:
+ 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=4))
+
+ dec = torch.cat(result_rows, dim=3)
+ return dec
+
+ def enable_tiling(self, use_tiling: bool = True):
+ self.use_tiling = use_tiling
+
+ def disable_tiling(self):
+ self.enable_tiling(False)
+
+ def init_from_ckpt(self, path, ignore_keys=list()):
+ sd = torch.load(path, map_location="cpu")
+ # print("init from " + path)
+
+ if "ema_state_dict" in sd and len(sd["ema_state_dict"]) > 0 and os.environ.get("NOT_USE_EMA_MODEL", 0) == 0:
+ # print("Load from ema model!")
+ sd = sd["ema_state_dict"]
+ sd = {key.replace("module.", ""): value for key, value in sd.items()}
+ elif "state_dict" in sd:
+ # print("Load from normal model!")
+ if "gen_model" in sd["state_dict"]:
+ sd = sd["state_dict"]["gen_model"]
+ else:
+ sd = sd["state_dict"]
+
+ keys = list(sd.keys())
+
+ for k in keys:
+ for ik in ignore_keys:
+ if k.startswith(ik):
+ print("Deleting key {} from state_dict.".format(k))
+ del sd[k]
+
+ miss, unexpected = self.load_state_dict(sd, strict=False)
+ assert len(miss) == 0, f"miss key: {miss}"
+ if len(unexpected) > 0:
+ for i in unexpected:
+ assert "loss" in i, "unexpected key: {i}"
+
+
+ae_stride_config = {
+ "CausalVAEModel_D4_2x8x8": [2, 8, 8],
+ "CausalVAEModel_D8_2x8x8": [2, 8, 8],
+ "CausalVAEModel_D4_4x8x8": [4, 8, 8],
+ "CausalVAEModel_D8_4x8x8": [4, 8, 8],
+}
+
+
+ae_channel_config = {
+ "CausalVAEModel_D4_2x8x8": 4,
+ "CausalVAEModel_D8_2x8x8": 8,
+ "CausalVAEModel_D4_4x8x8": 4,
+ "CausalVAEModel_D8_4x8x8": 8,
+}
+
+
+ae_denorm = {
+ "CausalVAEModel_D4_2x8x8": lambda x: (x + 1.0) / 2.0,
+ "CausalVAEModel_D8_2x8x8": lambda x: (x + 1.0) / 2.0,
+ "CausalVAEModel_D4_4x8x8": lambda x: (x + 1.0) / 2.0,
+ "CausalVAEModel_D8_4x8x8": lambda x: (x + 1.0) / 2.0,
+}
+
+ae_norm = {
+ "CausalVAEModel_D4_2x8x8": Lambda(lambda x: 2.0 * x - 1.0),
+ "CausalVAEModel_D8_2x8x8": Lambda(lambda x: 2.0 * x - 1.0),
+ "CausalVAEModel_D4_4x8x8": Lambda(lambda x: 2.0 * x - 1.0),
+ "CausalVAEModel_D8_4x8x8": Lambda(lambda x: 2.0 * x - 1.0),
+}
+
+
+class CausalVAEModelWrapper(nn.Module):
+ def __init__(self, model_path, subfolder=None, cache_dir=None, use_ema=False, **kwargs):
+ super(CausalVAEModelWrapper, self).__init__()
+ # if os.path.exists(ckpt):
+ # self.vae = CausalVAEModel.load_from_checkpoint(ckpt)
+ # hf_hub_download(model_path, subfolder="vae", filename="checkpoint.ckpt")
+ # cached_download(hf_hub_url(model_path, subfolder="vae", filename="checkpoint.ckpt"))
+ self.vae = CausalVAEModel.from_pretrained(model_path, subfolder=subfolder, cache_dir=cache_dir, **kwargs)
+ if use_ema:
+ self.vae.init_from_ema(model_path)
+ self.vae = self.vae.ema
+
+ def encode(self, x): # b c t h w
+ # x = self.vae.encode(x).sample()
+ x = self.vae.encode(x).sample().mul_(0.18215)
+ return x
+
+ def decode(self, x):
+ # x = self.vae.decode(x)
+ x = self.vae.decode(x / 0.18215)
+ x = rearrange(x, "b c t h w -> b t c h w").contiguous()
+ return x
+
+ def forward(self, x):
+ return self.decode(x)
+
+ def dtype(self):
+ return self.vae.dtype
diff --git a/videosys/models/transformers/open_sora_plan_transformer_3d.py b/videosys/models/transformers/open_sora_plan_v110_transformer_3d.py
similarity index 100%
rename from videosys/models/transformers/open_sora_plan_transformer_3d.py
rename to videosys/models/transformers/open_sora_plan_v110_transformer_3d.py
diff --git a/videosys/models/transformers/open_sora_plan_v120_transformer_3d.py b/videosys/models/transformers/open_sora_plan_v120_transformer_3d.py
new file mode 100644
index 00000000..94eec28a
--- /dev/null
+++ b/videosys/models/transformers/open_sora_plan_v120_transformer_3d.py
@@ -0,0 +1,2313 @@
+# Adapted from Open-Sora-Plan
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# Open-Sora-Plan: https://github.com/PKU-YuanGroup/Open-Sora-Plan
+# --------------------------------------------------------
+
+import collections
+import re
+from typing import Any, Dict, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention import FeedForward, GatedSelfAttentionDense
+from diffusers.models.attention_processor import Attention as Attention_
+from diffusers.models.embeddings import PixArtAlphaTextProjection, SinusoidalPositionalEmbedding
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormSingle, AdaLayerNormZero
+from diffusers.utils import deprecate, is_torch_version
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from einops import rearrange, repeat
+from torch import nn
+from torch.nn import functional as F
+
+from videosys.core.comm import all_to_all_comm
+from videosys.core.parallel_mgr import ParallelManager
+from videosys.core.pipeline import VideoSysPipelineOutput
+
+torch_npu = None
+npu_config = None
+set_run_dtype = None
+
+
+class PositionGetter3D(object):
+ """return positions of patches"""
+
+ def __init__(
+ self,
+ ):
+ self.cache_positions = {}
+
+ def __call__(self, b, t, h, w, device, attn):
+ if not (b, t, h, w) in self.cache_positions:
+ x = torch.arange(w, device=device)
+ y = torch.arange(h, device=device)
+ z = torch.arange(t, device=device)
+ pos = torch.cartesian_prod(z, y, x)
+ if attn.parallel_manager.sp_size > 1:
+ # print('PositionGetter3D', PositionGetter3D)
+ pos = pos.reshape(t * h * w, 3).transpose(0, 1).reshape(3, -1, 1).contiguous().expand(3, -1, b).clone()
+ else:
+ pos = pos.reshape(t * h * w, 3).transpose(0, 1).reshape(3, 1, -1).contiguous().expand(3, b, -1).clone()
+ poses = (pos[0].contiguous(), pos[1].contiguous(), pos[2].contiguous())
+ max_poses = (int(poses[0].max()), int(poses[1].max()), int(poses[2].max()))
+
+ self.cache_positions[b, t, h, w] = (poses, max_poses)
+ pos = self.cache_positions[b, t, h, w]
+
+ return pos
+
+
+class RoPE3D(torch.nn.Module):
+ def __init__(self, freq=10000.0, F0=1.0, interpolation_scale_thw=(1, 1, 1)):
+ super().__init__()
+ self.base = freq
+ self.F0 = F0
+ self.interpolation_scale_t = interpolation_scale_thw[0]
+ self.interpolation_scale_h = interpolation_scale_thw[1]
+ self.interpolation_scale_w = interpolation_scale_thw[2]
+ self.cache = {}
+
+ def get_cos_sin(self, D, seq_len, device, dtype, interpolation_scale=1):
+ if (D, seq_len, device, dtype) not in self.cache:
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, D, 2).float().to(device) / D))
+ t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) / interpolation_scale
+ freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype)
+ freqs = torch.cat((freqs, freqs), dim=-1)
+ cos = freqs.cos() # (Seq, Dim)
+ sin = freqs.sin()
+ self.cache[D, seq_len, device, dtype] = (cos, sin)
+ return self.cache[D, seq_len, device, dtype]
+
+ @staticmethod
+ def rotate_half(x):
+ x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+ def apply_rope1d(self, tokens, pos1d, cos, sin, attn):
+ assert pos1d.ndim == 2
+ if attn.parallel_manager.sp_size == 1:
+ # for (batch_size x nheads x ntokens x dim)
+ cos = torch.nn.functional.embedding(pos1d, cos)[:, None, :, :]
+ sin = torch.nn.functional.embedding(pos1d, sin)[:, None, :, :]
+ else:
+ # for (batch_size x ntokens x nheads x dim)
+ cos = torch.nn.functional.embedding(pos1d, cos)[:, :, None, :]
+ sin = torch.nn.functional.embedding(pos1d, sin)[:, :, None, :]
+
+ return (tokens * cos) + (self.rotate_half(tokens) * sin)
+
+ def forward(self, tokens, positions, attn):
+ """
+ input:
+ * tokens: batch_size x nheads x ntokens x dim
+ * positions: batch_size x ntokens x 3 (t, y and x position of each token)
+ output:
+ * tokens after appplying RoPE3D (batch_size x nheads x ntokens x x dim)
+ """
+ assert tokens.size(3) % 3 == 0, "number of dimensions should be a multiple of three"
+ D = tokens.size(3) // 3
+ poses, max_poses = positions
+ assert len(poses) == 3 and poses[0].ndim == 2 # Batch, Seq, 3
+ cos_t, sin_t = self.get_cos_sin(D, max_poses[0] + 1, tokens.device, tokens.dtype, self.interpolation_scale_t)
+ cos_y, sin_y = self.get_cos_sin(D, max_poses[1] + 1, tokens.device, tokens.dtype, self.interpolation_scale_h)
+ cos_x, sin_x = self.get_cos_sin(D, max_poses[2] + 1, tokens.device, tokens.dtype, self.interpolation_scale_w)
+ # split features into three along the feature dimension, and apply rope1d on each half
+ t, y, x = tokens.chunk(3, dim=-1)
+ t = self.apply_rope1d(t, poses[0], cos_t, sin_t, attn)
+ y = self.apply_rope1d(y, poses[1], cos_y, sin_y, attn)
+ x = self.apply_rope1d(x, poses[2], cos_x, sin_x, attn)
+ tokens = torch.cat((t, y, x), dim=-1)
+ return tokens
+
+
+def get_3d_sincos_pos_embed(
+ embed_dim,
+ grid_size,
+ cls_token=False,
+ extra_tokens=0,
+ interpolation_scale=1.0,
+ base_size=16,
+):
+ """
+ grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
+ [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+ """
+ # if isinstance(grid_size, int):
+ # grid_size = (grid_size, grid_size)
+ grid_t = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size[0]) / interpolation_scale[0]
+ grid_h = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size[1]) / interpolation_scale[1]
+ grid_w = np.arange(grid_size[2], dtype=np.float32) / (grid_size[2] / base_size[2]) / interpolation_scale[2]
+ grid = np.meshgrid(grid_w, grid_h, grid_t) # here w goes first
+ grid = np.stack(grid, axis=0)
+
+ grid = grid.reshape([3, 1, grid_size[2], grid_size[1], grid_size[0]])
+ pos_embed = get_3d_sincos_pos_embed_from_grid(embed_dim, grid)
+ # import ipdb;ipdb.set_trace()
+ if cls_token and extra_tokens > 0:
+ pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+ return pos_embed
+
+
+def get_3d_sincos_pos_embed_from_grid(embed_dim, grid):
+ if embed_dim % 3 != 0:
+ raise ValueError("embed_dim must be divisible by 3")
+
+ # import ipdb;ipdb.set_trace()
+ # use 1/3 of dimensions to encode grid_t/h/w
+ emb_t = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[0]) # (T*H*W, D/3)
+ emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[1]) # (T*H*W, D/3)
+ emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[2]) # (T*H*W, D/3)
+
+ emb = np.concatenate([emb_t, emb_h, emb_w], axis=1) # (T*H*W, D)
+ return emb
+
+
+def get_2d_sincos_pos_embed(
+ embed_dim,
+ grid_size,
+ cls_token=False,
+ extra_tokens=0,
+ interpolation_scale=1.0,
+ base_size=16,
+):
+ """
+ grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
+ [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+ """
+ # if isinstance(grid_size, int):
+ # grid_size = (grid_size, grid_size)
+
+ grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size[0]) / interpolation_scale[0]
+ grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size[1]) / interpolation_scale[1]
+ grid = np.meshgrid(grid_w, grid_h) # here w goes first
+ grid = np.stack(grid, axis=0)
+
+ grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])
+ pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+ if cls_token and extra_tokens > 0:
+ pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+ return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+ if embed_dim % 2 != 0:
+ raise ValueError("embed_dim must be divisible by 2")
+
+ # use 1/3 of dimensions to encode grid_t/h/w
+ emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)
+ emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)
+
+ emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+ return emb
+
+
+def get_1d_sincos_pos_embed(
+ embed_dim,
+ grid_size,
+ cls_token=False,
+ extra_tokens=0,
+ interpolation_scale=1.0,
+ base_size=16,
+):
+ """
+ grid_size: int of the grid return: pos_embed: [grid_size, embed_dim] or
+ [1+grid_size, embed_dim] (w/ or w/o cls_token)
+ """
+ # if isinstance(grid_size, int):
+ # grid_size = (grid_size, grid_size)
+
+ grid = np.arange(grid_size, dtype=np.float32) / (grid_size / base_size) / interpolation_scale
+ pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid) # (H*W, D/2)
+ if cls_token and extra_tokens > 0:
+ pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+ return pos_embed
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+ """
+ embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)
+ """
+ if embed_dim % 2 != 0:
+ raise ValueError("embed_dim must be divisible by 2")
+
+ omega = np.arange(embed_dim // 2, dtype=np.float64)
+ omega /= embed_dim / 2.0
+ omega = 1.0 / 10000**omega # (D/2,)
+
+ pos = pos.reshape(-1) # (M,)
+ out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product
+
+ emb_sin = np.sin(out) # (M, D/2)
+ emb_cos = np.cos(out) # (M, D/2)
+
+ emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)
+ return emb
+
+
+class PatchEmbed2D(nn.Module):
+ """2D Image to Patch Embedding but with 3D position embedding"""
+
+ def __init__(
+ self,
+ num_frames=1,
+ height=224,
+ width=224,
+ patch_size_t=1,
+ patch_size=16,
+ in_channels=3,
+ embed_dim=768,
+ layer_norm=False,
+ flatten=True,
+ bias=True,
+ interpolation_scale=(1, 1),
+ interpolation_scale_t=1,
+ use_abs_pos=True,
+ ):
+ super().__init__()
+ # assert num_frames == 1
+ self.use_abs_pos = use_abs_pos
+ self.flatten = flatten
+ self.layer_norm = layer_norm
+
+ self.proj = nn.Conv2d(
+ in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=(patch_size, patch_size), bias=bias
+ )
+ if layer_norm:
+ self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6)
+ else:
+ self.norm = None
+
+ self.patch_size_t = patch_size_t
+ self.patch_size = patch_size
+ # See:
+ # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161
+
+ self.height, self.width = height // patch_size, width // patch_size
+ self.base_size = (height // patch_size, width // patch_size)
+ self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1])
+ pos_embed = get_2d_sincos_pos_embed(
+ embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale
+ )
+ self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False)
+
+ self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+ self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+ self.interpolation_scale_t = interpolation_scale_t
+ temp_pos_embed = get_1d_sincos_pos_embed(
+ embed_dim, self.num_frames, base_size=self.base_size_t, interpolation_scale=self.interpolation_scale_t
+ )
+ self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False)
+ # self.temp_embed_gate = nn.Parameter(torch.tensor([0.0]))
+
+ # parallel
+ self.parallel_manager: ParallelManager = None
+
+ def forward(self, latent, num_frames):
+ b, _, _, _, _ = latent.shape
+ video_latent, image_latent = None, None
+ # b c 1 h w
+ # assert latent.shape[-3] == 1 and num_frames == 1
+ height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size
+ latent = rearrange(latent, "b c t h w -> (b t) c h w")
+ latent = self.proj(latent)
+
+ if self.flatten:
+ latent = latent.flatten(2).transpose(1, 2) # BT C H W -> BT N C
+ if self.layer_norm:
+ latent = self.norm(latent)
+
+ if self.use_abs_pos:
+ # Interpolate positional embeddings if needed.
+ # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160)
+ if self.height != height or self.width != width:
+ # raise NotImplementedError
+ pos_embed = get_2d_sincos_pos_embed(
+ embed_dim=self.pos_embed.shape[-1],
+ grid_size=(height, width),
+ base_size=self.base_size,
+ interpolation_scale=self.interpolation_scale,
+ )
+ pos_embed = torch.from_numpy(pos_embed)
+ pos_embed = pos_embed.float().unsqueeze(0).to(latent.device)
+ else:
+ pos_embed = self.pos_embed
+
+ if self.num_frames != num_frames:
+ # import ipdb;ipdb.set_trace()
+ # raise NotImplementedError
+ if self.parallel_manager.sp_size > 1:
+ sp_size = self.parallel_manager.sp_size
+ temp_pos_embed = get_1d_sincos_pos_embed(
+ embed_dim=self.temp_pos_embed.shape[-1],
+ grid_size=num_frames * sp_size,
+ base_size=self.base_size_t,
+ interpolation_scale=self.interpolation_scale_t,
+ )
+ rank = self.parallel_manager.sp_size % sp_size
+ st_frame = rank * num_frames
+ ed_frame = st_frame + num_frames
+ temp_pos_embed = temp_pos_embed[st_frame:ed_frame]
+
+ else:
+ temp_pos_embed = get_1d_sincos_pos_embed(
+ embed_dim=self.temp_pos_embed.shape[-1],
+ grid_size=num_frames,
+ base_size=self.base_size_t,
+ interpolation_scale=self.interpolation_scale_t,
+ )
+ temp_pos_embed = torch.from_numpy(temp_pos_embed)
+ temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device)
+ else:
+ temp_pos_embed = self.temp_pos_embed
+
+ latent = (latent + pos_embed).to(latent.dtype)
+
+ latent = rearrange(latent, "(b t) n c -> b t n c", b=b)
+ video_latent, image_latent = latent[:, :num_frames], latent[:, num_frames:]
+
+ if self.use_abs_pos:
+ # temp_pos_embed = temp_pos_embed.unsqueeze(2) * self.temp_embed_gate.tanh()
+ temp_pos_embed = temp_pos_embed.unsqueeze(2)
+ video_latent = (
+ (video_latent + temp_pos_embed).to(video_latent.dtype)
+ if video_latent is not None and video_latent.numel() > 0
+ else None
+ )
+ image_latent = (
+ (image_latent + temp_pos_embed[:, :1]).to(image_latent.dtype)
+ if image_latent is not None and image_latent.numel() > 0
+ else None
+ )
+
+ video_latent = (
+ rearrange(video_latent, "b t n c -> b (t n) c")
+ if video_latent is not None and video_latent.numel() > 0
+ else None
+ )
+ image_latent = (
+ rearrange(image_latent, "b t n c -> (b t) n c")
+ if image_latent is not None and image_latent.numel() > 0
+ else None
+ )
+
+ if num_frames == 1 and image_latent is None and not (self.parallel_manager.sp_size > 1):
+ image_latent = video_latent
+ video_latent = None
+ # print('video_latent is None, image_latent is None', video_latent is None, image_latent is None)
+ return video_latent, image_latent
+
+
+class OverlapPatchEmbed3D(nn.Module):
+ """2D Image to Patch Embedding but with 3D position embedding"""
+
+ def __init__(
+ self,
+ num_frames=1,
+ height=224,
+ width=224,
+ patch_size_t=1,
+ patch_size=16,
+ in_channels=3,
+ embed_dim=768,
+ layer_norm=False,
+ flatten=True,
+ bias=True,
+ interpolation_scale=(1, 1),
+ interpolation_scale_t=1,
+ use_abs_pos=True,
+ ):
+ super().__init__()
+ # assert patch_size_t == 1 and patch_size == 1
+ self.use_abs_pos = use_abs_pos
+ self.flatten = flatten
+ self.layer_norm = layer_norm
+
+ self.proj = nn.Conv3d(
+ in_channels,
+ embed_dim,
+ kernel_size=(patch_size_t, patch_size, patch_size),
+ stride=(patch_size_t, patch_size, patch_size),
+ bias=bias,
+ )
+ if layer_norm:
+ self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6)
+ else:
+ self.norm = None
+
+ self.patch_size_t = patch_size_t
+ self.patch_size = patch_size
+ # See:
+ # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161
+
+ self.height, self.width = height // patch_size, width // patch_size
+ self.base_size = (height // patch_size, width // patch_size)
+ self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1])
+ pos_embed = get_2d_sincos_pos_embed(
+ embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale
+ )
+ self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False)
+
+ self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+ self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+ self.interpolation_scale_t = interpolation_scale_t
+ temp_pos_embed = get_1d_sincos_pos_embed(
+ embed_dim, self.num_frames, base_size=self.base_size_t, interpolation_scale=self.interpolation_scale_t
+ )
+ self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False)
+ # self.temp_embed_gate = nn.Parameter(torch.tensor([0.0]))
+
+ def forward(self, latent, num_frames):
+ b, _, _, _, _ = latent.shape
+ video_latent, image_latent = None, None
+ # b c 1 h w
+ # assert latent.shape[-3] == 1 and num_frames == 1
+ height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size
+ # latent = rearrange(latent, 'b c t h w -> (b t) c h w')
+ latent = self.proj(latent)
+
+ if self.flatten:
+ # latent = latent.flatten(2).transpose(1, 2) # BT C H W -> BT N C
+ latent = rearrange(latent, "b c t h w -> (b t) (h w) c ")
+ if self.layer_norm:
+ latent = self.norm(latent)
+
+ if self.use_abs_pos:
+ # Interpolate positional embeddings if needed.
+ # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160)
+ if self.height != height or self.width != width:
+ # raise NotImplementedError
+ pos_embed = get_2d_sincos_pos_embed(
+ embed_dim=self.pos_embed.shape[-1],
+ grid_size=(height, width),
+ base_size=self.base_size,
+ interpolation_scale=self.interpolation_scale,
+ )
+ pos_embed = torch.from_numpy(pos_embed)
+ pos_embed = pos_embed.float().unsqueeze(0).to(latent.device)
+ else:
+ pos_embed = self.pos_embed
+
+ if self.num_frames != num_frames:
+ # import ipdb;ipdb.set_trace()
+ # raise NotImplementedError
+ temp_pos_embed = get_1d_sincos_pos_embed(
+ embed_dim=self.temp_pos_embed.shape[-1],
+ grid_size=num_frames,
+ base_size=self.base_size_t,
+ interpolation_scale=self.interpolation_scale_t,
+ )
+ temp_pos_embed = torch.from_numpy(temp_pos_embed)
+ temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device)
+ else:
+ temp_pos_embed = self.temp_pos_embed
+
+ latent = (latent + pos_embed).to(latent.dtype)
+
+ latent = rearrange(latent, "(b t) n c -> b t n c", b=b)
+ video_latent, image_latent = latent[:, :num_frames], latent[:, num_frames:]
+
+ if self.use_abs_pos:
+ # temp_pos_embed = temp_pos_embed.unsqueeze(2) * self.temp_embed_gate.tanh()
+ temp_pos_embed = temp_pos_embed.unsqueeze(2)
+ video_latent = (
+ (video_latent + temp_pos_embed).to(video_latent.dtype)
+ if video_latent is not None and video_latent.numel() > 0
+ else None
+ )
+ image_latent = (
+ (image_latent + temp_pos_embed[:, :1]).to(image_latent.dtype)
+ if image_latent is not None and image_latent.numel() > 0
+ else None
+ )
+
+ video_latent = (
+ rearrange(video_latent, "b t n c -> b (t n) c")
+ if video_latent is not None and video_latent.numel() > 0
+ else None
+ )
+ image_latent = (
+ rearrange(image_latent, "b t n c -> (b t) n c")
+ if image_latent is not None and image_latent.numel() > 0
+ else None
+ )
+
+ if num_frames == 1 and image_latent is None:
+ image_latent = video_latent
+ video_latent = None
+ return video_latent, image_latent
+
+
+class OverlapPatchEmbed2D(nn.Module):
+ """2D Image to Patch Embedding but with 3D position embedding"""
+
+ def __init__(
+ self,
+ num_frames=1,
+ height=224,
+ width=224,
+ patch_size_t=1,
+ patch_size=16,
+ in_channels=3,
+ embed_dim=768,
+ layer_norm=False,
+ flatten=True,
+ bias=True,
+ interpolation_scale=(1, 1),
+ interpolation_scale_t=1,
+ use_abs_pos=True,
+ ):
+ super().__init__()
+ assert patch_size_t == 1
+ self.use_abs_pos = use_abs_pos
+ self.flatten = flatten
+ self.layer_norm = layer_norm
+
+ self.proj = nn.Conv2d(
+ in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=(patch_size, patch_size), bias=bias
+ )
+ if layer_norm:
+ self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6)
+ else:
+ self.norm = None
+
+ self.patch_size_t = patch_size_t
+ self.patch_size = patch_size
+ # See:
+ # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161
+
+ self.height, self.width = height // patch_size, width // patch_size
+ self.base_size = (height // patch_size, width // patch_size)
+ self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1])
+ pos_embed = get_2d_sincos_pos_embed(
+ embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale
+ )
+ self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False)
+
+ self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+ self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+ self.interpolation_scale_t = interpolation_scale_t
+ temp_pos_embed = get_1d_sincos_pos_embed(
+ embed_dim, self.num_frames, base_size=self.base_size_t, interpolation_scale=self.interpolation_scale_t
+ )
+ self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False)
+ # self.temp_embed_gate = nn.Parameter(torch.tensor([0.0]))
+
+ def forward(self, latent, num_frames):
+ b, _, _, _, _ = latent.shape
+ video_latent, image_latent = None, None
+ # b c 1 h w
+ # assert latent.shape[-3] == 1 and num_frames == 1
+ height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size
+ latent = rearrange(latent, "b c t h w -> (b t) c h w")
+ latent = self.proj(latent)
+
+ if self.flatten:
+ latent = latent.flatten(2).transpose(1, 2) # BT C H W -> BT N C
+ if self.layer_norm:
+ latent = self.norm(latent)
+
+ if self.use_abs_pos:
+ # Interpolate positional embeddings if needed.
+ # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160)
+ if self.height != height or self.width != width:
+ # raise NotImplementedError
+ pos_embed = get_2d_sincos_pos_embed(
+ embed_dim=self.pos_embed.shape[-1],
+ grid_size=(height, width),
+ base_size=self.base_size,
+ interpolation_scale=self.interpolation_scale,
+ )
+ pos_embed = torch.from_numpy(pos_embed)
+ pos_embed = pos_embed.float().unsqueeze(0).to(latent.device)
+ else:
+ pos_embed = self.pos_embed
+
+ if self.num_frames != num_frames:
+ # import ipdb;ipdb.set_trace()
+ # raise NotImplementedError
+ temp_pos_embed = get_1d_sincos_pos_embed(
+ embed_dim=self.temp_pos_embed.shape[-1],
+ grid_size=num_frames,
+ base_size=self.base_size_t,
+ interpolation_scale=self.interpolation_scale_t,
+ )
+ temp_pos_embed = torch.from_numpy(temp_pos_embed)
+ temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device)
+ else:
+ temp_pos_embed = self.temp_pos_embed
+
+ latent = (latent + pos_embed).to(latent.dtype)
+
+ latent = rearrange(latent, "(b t) n c -> b t n c", b=b)
+ video_latent, image_latent = latent[:, :num_frames], latent[:, num_frames:]
+
+ if self.use_abs_pos:
+ # temp_pos_embed = temp_pos_embed.unsqueeze(2) * self.temp_embed_gate.tanh()
+ temp_pos_embed = temp_pos_embed.unsqueeze(2)
+ video_latent = (
+ (video_latent + temp_pos_embed).to(video_latent.dtype)
+ if video_latent is not None and video_latent.numel() > 0
+ else None
+ )
+ image_latent = (
+ (image_latent + temp_pos_embed[:, :1]).to(image_latent.dtype)
+ if image_latent is not None and image_latent.numel() > 0
+ else None
+ )
+
+ video_latent = (
+ rearrange(video_latent, "b t n c -> b (t n) c")
+ if video_latent is not None and video_latent.numel() > 0
+ else None
+ )
+ image_latent = (
+ rearrange(image_latent, "b t n c -> (b t) n c")
+ if image_latent is not None and image_latent.numel() > 0
+ else None
+ )
+
+ if num_frames == 1 and image_latent is None:
+ image_latent = video_latent
+ video_latent = None
+ return video_latent, image_latent
+
+
+class Attention(Attention_):
+ def __init__(self, downsampler, attention_mode, use_rope, interpolation_scale_thw, **kwags):
+ processor = AttnProcessor2_0(
+ attention_mode=attention_mode, use_rope=use_rope, interpolation_scale_thw=interpolation_scale_thw
+ )
+ super().__init__(processor=processor, **kwags)
+ self.downsampler = None
+ if downsampler: # downsampler k155_s122
+ downsampler_ker_size = list(re.search(r"k(\d{2,3})", downsampler).group(1)) # 122
+ down_factor = list(re.search(r"s(\d{2,3})", downsampler).group(1))
+ downsampler_ker_size = [int(i) for i in downsampler_ker_size]
+ downsampler_padding = [(i - 1) // 2 for i in downsampler_ker_size]
+ down_factor = [int(i) for i in down_factor]
+
+ if len(downsampler_ker_size) == 2:
+ self.downsampler = DownSampler2d(
+ kwags["query_dim"],
+ kwags["query_dim"],
+ kernel_size=downsampler_ker_size,
+ stride=1,
+ padding=downsampler_padding,
+ groups=kwags["query_dim"],
+ down_factor=down_factor,
+ down_shortcut=True,
+ )
+ elif len(downsampler_ker_size) == 3:
+ self.downsampler = DownSampler3d(
+ kwags["query_dim"],
+ kwags["query_dim"],
+ kernel_size=downsampler_ker_size,
+ stride=1,
+ padding=downsampler_padding,
+ groups=kwags["query_dim"],
+ down_factor=down_factor,
+ down_shortcut=True,
+ )
+
+ # parallel
+ self.parallel_manager: ParallelManager = None
+
+ def prepare_attention_mask(
+ self, attention_mask: torch.Tensor, target_length: int, batch_size: int, out_dim: int = 3
+ ) -> torch.Tensor:
+ r"""
+ Prepare the attention mask for the attention computation.
+
+ Args:
+ attention_mask (`torch.Tensor`):
+ The attention mask to prepare.
+ target_length (`int`):
+ The target length of the attention mask. This is the length of the attention mask after padding.
+ batch_size (`int`):
+ The batch size, which is used to repeat the attention mask.
+ out_dim (`int`, *optional*, defaults to `3`):
+ The output dimension of the attention mask. Can be either `3` or `4`.
+
+ Returns:
+ `torch.Tensor`: The prepared attention mask.
+ """
+ head_size = self.heads
+ if self.parallel_manager.sp_size > 1:
+ head_size = head_size // self.parallel_manager.sp_size
+ if attention_mask is None:
+ return attention_mask
+
+ current_length: int = attention_mask.shape[-1]
+ if current_length != target_length:
+ if attention_mask.device.type == "mps":
+ # HACK: MPS: Does not support padding by greater than dimension of input tensor.
+ # Instead, we can manually construct the padding tensor.
+ padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length)
+ padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device)
+ attention_mask = torch.cat([attention_mask, padding], dim=2)
+ else:
+ # TODO: for pipelines such as stable-diffusion, padding cross-attn mask:
+ # we want to instead pad by (0, remaining_length), where remaining_length is:
+ # remaining_length: int = target_length - current_length
+ # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding
+ attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
+
+ if out_dim == 3:
+ if attention_mask.shape[0] < batch_size * head_size:
+ attention_mask = attention_mask.repeat_interleave(head_size, dim=0)
+ elif out_dim == 4:
+ attention_mask = attention_mask.unsqueeze(1)
+ attention_mask = attention_mask.repeat_interleave(head_size, dim=1)
+
+ return attention_mask
+
+
+class DownSampler3d(nn.Module):
+ def __init__(self, *args, **kwargs):
+ """Required kwargs: down_factor, downsampler"""
+ super().__init__()
+ self.down_factor = kwargs.pop("down_factor")
+ self.down_shortcut = kwargs.pop("down_shortcut")
+ self.layer = nn.Conv3d(*args, **kwargs)
+
+ def forward(self, x, attention_mask, t, h, w):
+ x.shape[0]
+ x = rearrange(x, "b (t h w) d -> b d t h w", t=t, h=h, w=w)
+ if npu_config is None:
+ x = self.layer(x) + (x if self.down_shortcut else 0)
+ else:
+ x_dtype = x.dtype
+ x = npu_config.run_conv3d(self.layer, x, x_dtype) + (x if self.down_shortcut else 0)
+
+ self.t = t // self.down_factor[0]
+ self.h = h // self.down_factor[1]
+ self.w = w // self.down_factor[2]
+ x = rearrange(
+ x,
+ "b d (t dt) (h dh) (w dw) -> (b dt dh dw) (t h w) d",
+ t=t // self.down_factor[0],
+ h=h // self.down_factor[1],
+ w=w // self.down_factor[2],
+ dt=self.down_factor[0],
+ dh=self.down_factor[1],
+ dw=self.down_factor[2],
+ )
+
+ attention_mask = rearrange(attention_mask, "b 1 (t h w) -> b 1 t h w", t=t, h=h, w=w)
+ attention_mask = rearrange(
+ attention_mask,
+ "b 1 (t dt) (h dh) (w dw) -> (b dt dh dw) 1 (t h w)",
+ t=t // self.down_factor[0],
+ h=h // self.down_factor[1],
+ w=w // self.down_factor[2],
+ dt=self.down_factor[0],
+ dh=self.down_factor[1],
+ dw=self.down_factor[2],
+ )
+ return x, attention_mask
+
+ def reverse(self, x, t, h, w):
+ x = rearrange(
+ x,
+ "(b dt dh dw) (t h w) d -> b (t dt h dh w dw) d",
+ t=t,
+ h=h,
+ w=w,
+ dt=self.down_factor[0],
+ dh=self.down_factor[1],
+ dw=self.down_factor[2],
+ )
+ return x
+
+
+class DownSampler2d(nn.Module):
+ def __init__(self, *args, **kwargs):
+ """Required kwargs: down_factor, downsampler"""
+ super().__init__()
+ self.down_factor = kwargs.pop("down_factor")
+ self.down_shortcut = kwargs.pop("down_shortcut")
+ self.layer = nn.Conv2d(*args, **kwargs)
+
+ def forward(self, x, attention_mask, t, h, w):
+ x.shape[0]
+ x = rearrange(x, "b (t h w) d -> (b t) d h w", t=t, h=h, w=w)
+ x = self.layer(x) + (x if self.down_shortcut else 0)
+
+ self.t = 1
+ self.h = h // self.down_factor[0]
+ self.w = w // self.down_factor[1]
+
+ x = rearrange(
+ x,
+ "b d (h dh) (w dw) -> (b dh dw) (h w) d",
+ h=h // self.down_factor[0],
+ w=w // self.down_factor[1],
+ dh=self.down_factor[0],
+ dw=self.down_factor[1],
+ )
+
+ attention_mask = rearrange(attention_mask, "b 1 (t h w) -> (b t) 1 h w", h=h, w=w)
+ attention_mask = rearrange(
+ attention_mask,
+ "b 1 (h dh) (w dw) -> (b dh dw) 1 (h w)",
+ h=h // self.down_factor[0],
+ w=w // self.down_factor[1],
+ dh=self.down_factor[0],
+ dw=self.down_factor[1],
+ )
+ return x, attention_mask
+
+ def reverse(self, x, t, h, w):
+ x = rearrange(
+ x, "(b t dh dw) (h w) d -> b (t h dh w dw) d", t=t, h=h, w=w, dh=self.down_factor[0], dw=self.down_factor[1]
+ )
+ return x
+
+
+class AttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+ """
+
+ def __init__(self, attention_mode="xformers", use_rope=False, interpolation_scale_thw=(1, 1, 1)):
+ self.use_rope = use_rope
+ self.interpolation_scale_thw = interpolation_scale_thw
+ if self.use_rope:
+ self._init_rope(interpolation_scale_thw)
+ self.attention_mode = attention_mode
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+ def _init_rope(self, interpolation_scale_thw):
+ self.rope = RoPE3D(interpolation_scale_thw=interpolation_scale_thw)
+ self.position_getter = PositionGetter3D()
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ temb: Optional[torch.FloatTensor] = None,
+ frame: int = 8,
+ height: int = 16,
+ width: int = 16,
+ *args,
+ **kwargs,
+ ) -> torch.FloatTensor:
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+ deprecate("scale", "1.0.0", deprecation_message)
+
+ if attn.downsampler is not None:
+ hidden_states, attention_mask = attn.downsampler(hidden_states, attention_mask, t=frame, h=height, w=width)
+ frame, height, width = attn.downsampler.t, attn.downsampler.h, attn.downsampler.w
+
+ residual = hidden_states
+
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ if attn.parallel_manager.sp_size > 1:
+ if npu_config is not None:
+ sequence_length, batch_size, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+ else:
+ sequence_length, batch_size, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+ else:
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(
+ attention_mask, sequence_length * attn.parallel_manager.sp_size, batch_size
+ )
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ if attn.parallel_manager.sp_size > 1:
+ attention_mask = attention_mask.view(
+ batch_size, attn.heads // attn.parallel_manager.sp_size, -1, attention_mask.shape[-1]
+ )
+ else:
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ if attn.parallel_manager.sp_size > 1:
+ query = query.reshape(-1, attn.heads, head_dim) # [s // sp, b, h * d] -> [s // sp * b, h, d]
+ key = key.reshape(-1, attn.heads, head_dim)
+ value = value.reshape(-1, attn.heads, head_dim)
+ # query = attn.q_norm(query)
+ # key = attn.k_norm(key)
+ h_size = attn.heads * head_dim
+ sp_size = attn.parallel_manager.sp_size
+ h_size_sp = h_size // sp_size
+ # apply all_to_all to gather sequence and split attention heads [s // sp * b, h, d] -> [s * b, h // sp, d]
+ query = all_to_all_comm(query, attn.parallel_manager.sp_group, scatter_dim=1, gather_dim=0).reshape(
+ -1, batch_size, h_size_sp
+ )
+ key = all_to_all_comm(key, attn.parallel_manager.sp_group, scatter_dim=1, gather_dim=0).reshape(
+ -1, batch_size, h_size_sp
+ )
+ value = all_to_all_comm(value, attn.parallel_manager.sp_group, scatter_dim=1, gather_dim=0).reshape(
+ -1, batch_size, h_size_sp
+ )
+ query = query.reshape(-1, batch_size, attn.heads // sp_size, head_dim)
+ key = key.reshape(-1, batch_size, attn.heads // sp_size, head_dim)
+ value = value.reshape(-1, batch_size, attn.heads // sp_size, head_dim)
+ # print('query', query.shape, 'key', key.shape, 'value', value.shape)
+ if self.use_rope:
+ # require the shape of (batch_size x nheads x ntokens x dim)
+ pos_thw = self.position_getter(
+ batch_size, t=frame * sp_size, h=height, w=width, device=query.device, attn=attn
+ )
+ query = self.rope(query, pos_thw, attn)
+ key = self.rope(key, pos_thw, attn)
+
+ # print('after rope query', query.shape, 'key', key.shape, 'value', value.shape)
+ query = rearrange(query, "s b h d -> b h s d")
+ key = rearrange(key, "s b h d -> b h s d")
+ value = rearrange(value, "s b h d -> b h s d")
+ # print('rearrange query', query.shape, 'key', key.shape, 'value', value.shape)
+
+ # 0, -10000 ->(bool) False, True ->(any) True ->(not) False
+ # 0, 0 ->(bool) False, False ->(any) False ->(not) True
+ if attention_mask is None or not torch.any(attention_mask.bool()): # 0 mean visible
+ attention_mask = None
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+ hidden_states = rearrange(hidden_states, "b h s d -> s b h d")
+ hidden_states = hidden_states.reshape(-1, attn.heads // sp_size, head_dim)
+
+ # [s * b, h // sp, d] -> [s // sp * b, h, d] -> [s // sp, b, h * d]
+ hidden_states = all_to_all_comm(
+ hidden_states, attn.parallel_manager.sp_group, scatter_dim=0, gather_dim=1
+ ).reshape(-1, batch_size, h_size)
+ else:
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # qk norm
+ # query = attn.q_norm(query)
+ # key = attn.k_norm(key)
+
+ if self.use_rope:
+ # require the shape of (batch_size x nheads x ntokens x dim)
+ pos_thw = self.position_getter(batch_size, t=frame, h=height, w=width, device=query.device, attn=attn)
+ query = self.rope(query, pos_thw, attn)
+ key = self.rope(key, pos_thw, attn)
+
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # 0, -10000 ->(bool) False, True ->(any) True ->(not) False
+ # 0, 0 ->(bool) False, False ->(any) False ->(not) True
+ if attention_mask is None or not torch.any(attention_mask.bool()): # 0 mean visible
+ attention_mask = None
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ # import ipdb;ipdb.set_trace()
+ # print(attention_mask)
+ if self.attention_mode == "flash":
+ assert attention_mask is None, "flash-attn do not support attention_mask"
+ with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False):
+ hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+ elif self.attention_mode == "xformers":
+ with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True):
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+ elif self.attention_mode == "math":
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+ else:
+ raise NotImplementedError(f"Found attention_mode: {self.attention_mode}")
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ if attn.downsampler is not None:
+ hidden_states = attn.downsampler.reverse(hidden_states, t=frame, h=height, w=width)
+ return hidden_states
+
+
+class FeedForward_Conv3d(nn.Module):
+ def __init__(self, downsampler, dim, hidden_features, bias=True):
+ super(FeedForward_Conv3d, self).__init__()
+
+ self.bias = bias
+
+ self.project_in = nn.Linear(dim, hidden_features, bias=bias)
+
+ self.dwconv = nn.ModuleList(
+ [
+ nn.Conv3d(
+ hidden_features,
+ hidden_features,
+ kernel_size=(5, 5, 5),
+ stride=1,
+ padding=(2, 2, 2),
+ dilation=1,
+ groups=hidden_features,
+ bias=bias,
+ ),
+ nn.Conv3d(
+ hidden_features,
+ hidden_features,
+ kernel_size=(3, 3, 3),
+ stride=1,
+ padding=(1, 1, 1),
+ dilation=1,
+ groups=hidden_features,
+ bias=bias,
+ ),
+ nn.Conv3d(
+ hidden_features,
+ hidden_features,
+ kernel_size=(1, 1, 1),
+ stride=1,
+ padding=(0, 0, 0),
+ dilation=1,
+ groups=hidden_features,
+ bias=bias,
+ ),
+ ]
+ )
+
+ self.project_out = nn.Linear(hidden_features, dim, bias=bias)
+
+ def forward(self, x, t, h, w):
+ # import ipdb;ipdb.set_trace()
+ if npu_config is None:
+ x = self.project_in(x)
+ x = rearrange(x, "b (t h w) d -> b d t h w", t=t, h=h, w=w)
+ x = F.gelu(x)
+ out = x
+ for module in self.dwconv:
+ out = out + module(x)
+ out = rearrange(out, "b d t h w -> b (t h w) d", t=t, h=h, w=w)
+ x = self.project_out(out)
+ else:
+ x_dtype = x.dtype
+ x = npu_config.run_conv3d(self.project_in, x, npu_config.replaced_type)
+ x = rearrange(x, "b (t h w) d -> b d t h w", t=t, h=h, w=w)
+ x = F.gelu(x)
+ out = x
+ for module in self.dwconv:
+ out = out + npu_config.run_conv3d(module, x, npu_config.replaced_type)
+ out = rearrange(out, "b d t h w -> b (t h w) d", t=t, h=h, w=w)
+ x = npu_config.run_conv3d(self.project_out, out, x_dtype)
+ return x
+
+
+class FeedForward_Conv2d(nn.Module):
+ def __init__(self, downsampler, dim, hidden_features, bias=True):
+ super(FeedForward_Conv2d, self).__init__()
+
+ self.bias = bias
+
+ self.project_in = nn.Linear(dim, hidden_features, bias=bias)
+
+ self.dwconv = nn.ModuleList(
+ [
+ nn.Conv2d(
+ hidden_features,
+ hidden_features,
+ kernel_size=(5, 5),
+ stride=1,
+ padding=(2, 2),
+ dilation=1,
+ groups=hidden_features,
+ bias=bias,
+ ),
+ nn.Conv2d(
+ hidden_features,
+ hidden_features,
+ kernel_size=(3, 3),
+ stride=1,
+ padding=(1, 1),
+ dilation=1,
+ groups=hidden_features,
+ bias=bias,
+ ),
+ nn.Conv2d(
+ hidden_features,
+ hidden_features,
+ kernel_size=(1, 1),
+ stride=1,
+ padding=(0, 0),
+ dilation=1,
+ groups=hidden_features,
+ bias=bias,
+ ),
+ ]
+ )
+
+ self.project_out = nn.Linear(hidden_features, dim, bias=bias)
+
+ def forward(self, x, t, h, w):
+ # import ipdb;ipdb.set_trace()
+ x = self.project_in(x)
+ x = rearrange(x, "b (t h w) d -> (b t) d h w", t=t, h=h, w=w)
+ x = F.gelu(x)
+ out = x
+ for module in self.dwconv:
+ out = out + module(x)
+ out = rearrange(out, "(b t) d h w -> b (t h w) d", t=t, h=h, w=w)
+ x = self.project_out(out)
+ return x
+
+
+@maybe_allow_in_graph
+class BasicTransformerBlock(nn.Module):
+ r"""
+ A basic Transformer block.
+
+ Parameters:
+ dim (`int`): The number of channels in the input and output.
+ num_attention_heads (`int`): The number of heads to use for multi-head attention.
+ attention_head_dim (`int`): The number of channels in each head.
+ dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+ cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
+ activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+ num_embeds_ada_norm (:
+ obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
+ attention_bias (:
+ obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
+ only_cross_attention (`bool`, *optional*):
+ Whether to use only cross-attention layers. In this case two cross attention layers are used.
+ double_self_attention (`bool`, *optional*):
+ Whether to use two self-attention layers. In this case no cross attention layers are used.
+ upcast_attention (`bool`, *optional*):
+ Whether to upcast the attention computation to float32. This is useful for mixed precision training.
+ norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
+ Whether to use learnable elementwise affine parameters for normalization.
+ norm_type (`str`, *optional*, defaults to `"layer_norm"`):
+ The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
+ final_dropout (`bool` *optional*, defaults to False):
+ Whether to apply a final dropout after the last feed-forward layer.
+ attention_type (`str`, *optional*, defaults to `"default"`):
+ The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`.
+ positional_embeddings (`str`, *optional*, defaults to `None`):
+ The type of positional embeddings to apply to.
+ num_positional_embeddings (`int`, *optional*, defaults to `None`):
+ The maximum number of positional embeddings to apply.
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ num_attention_heads: int,
+ attention_head_dim: int,
+ dropout=0.0,
+ cross_attention_dim: Optional[int] = None,
+ activation_fn: str = "geglu",
+ num_embeds_ada_norm: Optional[int] = None,
+ attention_bias: bool = False,
+ only_cross_attention: bool = False,
+ double_self_attention: bool = False,
+ upcast_attention: bool = False,
+ norm_elementwise_affine: bool = True,
+ norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen'
+ norm_eps: float = 1e-5,
+ final_dropout: bool = False,
+ attention_type: str = "default",
+ positional_embeddings: Optional[str] = None,
+ num_positional_embeddings: Optional[int] = None,
+ ada_norm_continous_conditioning_embedding_dim: Optional[int] = None,
+ ada_norm_bias: Optional[int] = None,
+ ff_inner_dim: Optional[int] = None,
+ ff_bias: bool = True,
+ attention_out_bias: bool = True,
+ attention_mode: str = "xformers",
+ downsampler: str = None,
+ use_rope: bool = False,
+ interpolation_scale_thw: Tuple[int] = (1, 1, 1),
+ ):
+ super().__init__()
+ self.only_cross_attention = only_cross_attention
+ self.downsampler = downsampler
+
+ # We keep these boolean flags for backward-compatibility.
+ self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
+ self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm"
+ self.use_ada_layer_norm_single = norm_type == "ada_norm_single"
+ self.use_layer_norm = norm_type == "layer_norm"
+ self.use_ada_layer_norm_continuous = norm_type == "ada_norm_continuous"
+
+ if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
+ raise ValueError(
+ f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"
+ f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."
+ )
+
+ self.norm_type = norm_type
+ self.num_embeds_ada_norm = num_embeds_ada_norm
+
+ if positional_embeddings and (num_positional_embeddings is None):
+ raise ValueError(
+ "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined."
+ )
+
+ if positional_embeddings == "sinusoidal":
+ self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings)
+ else:
+ self.pos_embed = None
+
+ # Define 3 blocks. Each block has its own normalization layer.
+ # 1. Self-Attn
+ if norm_type == "ada_norm":
+ self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)
+ elif norm_type == "ada_norm_zero":
+ self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
+ elif norm_type == "ada_norm_continuous":
+ self.norm1 = AdaLayerNormContinuous(
+ dim,
+ ada_norm_continous_conditioning_embedding_dim,
+ norm_elementwise_affine,
+ norm_eps,
+ ada_norm_bias,
+ "rms_norm",
+ )
+ else:
+ self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
+
+ self.attn1 = Attention(
+ query_dim=dim,
+ heads=num_attention_heads,
+ dim_head=attention_head_dim,
+ dropout=dropout,
+ bias=attention_bias,
+ cross_attention_dim=cross_attention_dim if only_cross_attention else None,
+ upcast_attention=upcast_attention,
+ out_bias=attention_out_bias,
+ attention_mode=attention_mode,
+ downsampler=downsampler,
+ use_rope=use_rope,
+ interpolation_scale_thw=interpolation_scale_thw,
+ )
+
+ # 2. Cross-Attn
+ if cross_attention_dim is not None or double_self_attention:
+ # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
+ # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
+ # the second cross attention block.
+ if norm_type == "ada_norm":
+ self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm)
+ elif norm_type == "ada_norm_continuous":
+ self.norm2 = AdaLayerNormContinuous(
+ dim,
+ ada_norm_continous_conditioning_embedding_dim,
+ norm_elementwise_affine,
+ norm_eps,
+ ada_norm_bias,
+ "rms_norm",
+ )
+ else:
+ self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+
+ self.attn2 = Attention(
+ query_dim=dim,
+ cross_attention_dim=cross_attention_dim if not double_self_attention else None,
+ heads=num_attention_heads,
+ dim_head=attention_head_dim,
+ dropout=dropout,
+ bias=attention_bias,
+ upcast_attention=upcast_attention,
+ out_bias=attention_out_bias,
+ attention_mode=attention_mode,
+ downsampler=False,
+ use_rope=False,
+ interpolation_scale_thw=interpolation_scale_thw,
+ ) # is self-attn if encoder_hidden_states is none
+ else:
+ self.norm2 = None
+ self.attn2 = None
+
+ # 3. Feed-forward
+ if norm_type == "ada_norm_continuous":
+ self.norm3 = AdaLayerNormContinuous(
+ dim,
+ ada_norm_continous_conditioning_embedding_dim,
+ norm_elementwise_affine,
+ norm_eps,
+ ada_norm_bias,
+ "layer_norm",
+ )
+
+ elif norm_type in ["ada_norm_zero", "ada_norm", "layer_norm", "ada_norm_continuous"]:
+ self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+ elif norm_type == "layer_norm_i2vgen":
+ self.norm3 = None
+
+ if downsampler:
+ downsampler_ker_size = list(re.search(r"k(\d{2,3})", downsampler).group(1)) # 122
+ # if len(downsampler_ker_size) == 3:
+ # self.ff = FeedForward_Conv3d(
+ # downsampler,
+ # dim,
+ # 2 * dim,
+ # bias=ff_bias,
+ # )
+ # elif len(downsampler_ker_size) == 2:
+ self.ff = FeedForward_Conv2d(
+ downsampler,
+ dim,
+ 2 * dim,
+ bias=ff_bias,
+ )
+ else:
+ self.ff = FeedForward(
+ dim,
+ dropout=dropout,
+ activation_fn=activation_fn,
+ final_dropout=final_dropout,
+ inner_dim=ff_inner_dim,
+ bias=ff_bias,
+ )
+
+ # 4. Fuser
+ if attention_type == "gated" or attention_type == "gated-text-image":
+ self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim)
+
+ # 5. Scale-shift for PixArt-Alpha.
+ if norm_type == "ada_norm_single":
+ self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+
+ # let chunk size default to None
+ self._chunk_size = None
+ self._chunk_dim = 0
+
+ # parallel
+ self.parallel_manager: ParallelManager = None
+
+ def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
+ # Sets chunk feed-forward
+ self._chunk_size = chunk_size
+ self._chunk_dim = dim
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
+ timestep: Optional[torch.LongTensor] = None,
+ cross_attention_kwargs: Dict[str, Any] = None,
+ class_labels: Optional[torch.LongTensor] = None,
+ frame: int = None,
+ height: int = None,
+ width: int = None,
+ added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+ ) -> torch.FloatTensor:
+ # Notice that normalization is always applied before the real computation in the following blocks.
+ # 0. Self-Attention
+ batch_size = hidden_states.shape[0]
+
+ # import ipdb;ipdb.set_trace()
+ if self.norm_type == "ada_norm":
+ norm_hidden_states = self.norm1(hidden_states, timestep)
+ elif self.norm_type == "ada_norm_zero":
+ norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+ hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
+ )
+ elif self.norm_type in ["layer_norm", "layer_norm_i2vgen"]:
+ norm_hidden_states = self.norm1(hidden_states)
+ elif self.norm_type == "ada_norm_continuous":
+ norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"])
+ elif self.norm_type == "ada_norm_single":
+ if self.parallel_manager.sp_size > 1:
+ batch_size = hidden_states.shape[1]
+ # print('hidden_states', hidden_states.shape)
+ # print('timestep', timestep.shape)
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+ self.scale_shift_table[:, None] + timestep.reshape(6, batch_size, -1)
+ ).chunk(6, dim=0)
+ else:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+ self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
+ ).chunk(6, dim=1)
+ norm_hidden_states = self.norm1(hidden_states)
+ norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
+ # norm_hidden_states = norm_hidden_states.squeeze(1)
+ else:
+ raise ValueError("Incorrect norm used")
+
+ if self.pos_embed is not None:
+ norm_hidden_states = self.pos_embed(norm_hidden_states)
+
+ # 1. Prepare GLIGEN inputs
+ cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}
+ gligen_kwargs = cross_attention_kwargs.pop("gligen", None)
+
+ attn_output = self.attn1(
+ norm_hidden_states,
+ encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
+ attention_mask=attention_mask,
+ frame=frame,
+ height=height,
+ width=width,
+ **cross_attention_kwargs,
+ )
+ if self.norm_type == "ada_norm_zero":
+ attn_output = gate_msa.unsqueeze(1) * attn_output
+ elif self.norm_type == "ada_norm_single":
+ attn_output = gate_msa * attn_output
+
+ hidden_states = attn_output + hidden_states
+ if hidden_states.ndim == 4:
+ hidden_states = hidden_states.squeeze(1)
+
+ # 1.2 GLIGEN Control
+ if gligen_kwargs is not None:
+ hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"])
+
+ # 3. Cross-Attention
+ if self.attn2 is not None:
+ if self.norm_type == "ada_norm":
+ norm_hidden_states = self.norm2(hidden_states, timestep)
+ elif self.norm_type in ["ada_norm_zero", "layer_norm", "layer_norm_i2vgen"]:
+ norm_hidden_states = self.norm2(hidden_states)
+ elif self.norm_type == "ada_norm_single":
+ # For PixArt norm2 isn't applied here:
+ # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103
+ norm_hidden_states = hidden_states
+ elif self.norm_type == "ada_norm_continuous":
+ norm_hidden_states = self.norm2(hidden_states, added_cond_kwargs["pooled_text_emb"])
+ else:
+ raise ValueError("Incorrect norm")
+
+ if self.pos_embed is not None and self.norm_type != "ada_norm_single":
+ norm_hidden_states = self.pos_embed(norm_hidden_states)
+
+ attn_output = self.attn2(
+ norm_hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ attention_mask=encoder_attention_mask,
+ **cross_attention_kwargs,
+ )
+ hidden_states = attn_output + hidden_states
+
+ # 4. Feed-forward
+ # i2vgen doesn't have this norm 🤷♂️
+ if self.norm_type == "ada_norm_continuous":
+ norm_hidden_states = self.norm3(hidden_states, added_cond_kwargs["pooled_text_emb"])
+ elif not self.norm_type == "ada_norm_single":
+ norm_hidden_states = self.norm3(hidden_states)
+
+ if self.norm_type == "ada_norm_zero":
+ norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+ if self.norm_type == "ada_norm_single":
+ norm_hidden_states = self.norm2(hidden_states)
+ norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
+
+ # if self._chunk_size is not None:
+ # # "feed_forward_chunk_size" can be used to save memory
+ # ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)
+ # else:
+
+ if self.downsampler:
+ ff_output = self.ff(norm_hidden_states, t=frame, h=height, w=width)
+ else:
+ ff_output = self.ff(norm_hidden_states)
+
+ if self.norm_type == "ada_norm_zero":
+ ff_output = gate_mlp.unsqueeze(1) * ff_output
+ elif self.norm_type == "ada_norm_single":
+ ff_output = gate_mlp * ff_output
+
+ hidden_states = ff_output + hidden_states
+ if hidden_states.ndim == 4:
+ hidden_states = hidden_states.squeeze(1)
+
+ return hidden_states
+
+
+def to_2tuple(x):
+ if isinstance(x, collections.abc.Iterable):
+ return x
+ return (x, x)
+
+
+class OpenSoraT2V(ModelMixin, ConfigMixin):
+ """
+ A 2D Transformer model for image-like data.
+
+ Parameters:
+ num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
+ attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
+ in_channels (`int`, *optional*):
+ The number of channels in the input and output (specify if the input is **continuous**).
+ num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
+ dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+ cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
+ sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**).
+ This is fixed during training since it is used to learn a number of position embeddings.
+ num_vector_embeds (`int`, *optional*):
+ The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**).
+ Includes the class for the masked latent pixel.
+ activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
+ num_embeds_ada_norm ( `int`, *optional*):
+ The number of diffusion steps used during training. Pass if at least one of the norm_layers is
+ `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
+ added to the hidden states.
+
+ During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`.
+ attention_bias (`bool`, *optional*):
+ Configure if the `TransformerBlocks` attention should contain a bias parameter.
+ """
+
+ _supports_gradient_checkpointing = True
+
+ @register_to_config
+ def __init__(
+ self,
+ num_attention_heads: int = 16,
+ attention_head_dim: int = 88,
+ in_channels: Optional[int] = None,
+ out_channels: Optional[int] = None,
+ num_layers: int = 1,
+ dropout: float = 0.0,
+ norm_num_groups: int = 32,
+ cross_attention_dim: Optional[int] = None,
+ attention_bias: bool = False,
+ sample_size: Optional[int] = None,
+ sample_size_t: Optional[int] = None,
+ num_vector_embeds: Optional[int] = None,
+ patch_size: Optional[int] = None,
+ patch_size_t: Optional[int] = None,
+ activation_fn: str = "geglu",
+ num_embeds_ada_norm: Optional[int] = None,
+ use_linear_projection: bool = False,
+ only_cross_attention: bool = False,
+ double_self_attention: bool = False,
+ upcast_attention: bool = False,
+ norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen'
+ norm_elementwise_affine: bool = True,
+ norm_eps: float = 1e-5,
+ attention_type: str = "default",
+ caption_channels: int = None,
+ interpolation_scale_h: float = None,
+ interpolation_scale_w: float = None,
+ interpolation_scale_t: float = None,
+ use_additional_conditions: Optional[bool] = None,
+ attention_mode: str = "xformers",
+ downsampler: str = None,
+ use_rope: bool = False,
+ use_stable_fp32: bool = False,
+ ):
+ super().__init__()
+
+ # Validate inputs.
+ if patch_size is not None:
+ if norm_type not in ["ada_norm", "ada_norm_zero", "ada_norm_single"]:
+ raise NotImplementedError(
+ f"Forward pass is not implemented when `patch_size` is not None and `norm_type` is '{norm_type}'."
+ )
+ elif norm_type in ["ada_norm", "ada_norm_zero"] and num_embeds_ada_norm is None:
+ raise ValueError(
+ f"When using a `patch_size` and this `norm_type` ({norm_type}), `num_embeds_ada_norm` cannot be None."
+ )
+
+ # Set some common variables used across the board.
+ self.use_rope = use_rope
+ self.use_linear_projection = use_linear_projection
+ self.interpolation_scale_t = interpolation_scale_t
+ self.interpolation_scale_h = interpolation_scale_h
+ self.interpolation_scale_w = interpolation_scale_w
+ self.downsampler = downsampler
+ self.caption_channels = caption_channels
+ self.num_attention_heads = num_attention_heads
+ self.attention_head_dim = attention_head_dim
+ self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
+ self.in_channels = in_channels
+ self.out_channels = in_channels if out_channels is None else out_channels
+ self.gradient_checkpointing = False
+ self.config.hidden_size = self.inner_dim
+ use_additional_conditions = False
+ # if use_additional_conditions is None:
+ # if norm_type == "ada_norm_single" and sample_size == 128:
+ # use_additional_conditions = True
+ # else:
+ # use_additional_conditions = False
+ self.use_additional_conditions = use_additional_conditions
+
+ # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
+ # Define whether input is continuous or discrete depending on configuration
+ assert in_channels is not None and patch_size is not None
+
+ if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
+ deprecation_message = (
+ f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"
+ " incorrectly set to `'layer_norm'`. Make sure to set `norm_type` to `'ada_norm'` in the config."
+ " Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect"
+ " results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it"
+ " would be very nice if you could open a Pull request for the `transformer/config.json` file"
+ )
+ deprecate("norm_type!=num_embeds_ada_norm", "1.0.0", deprecation_message, standard_warn=False)
+ norm_type = "ada_norm"
+
+ # 2. Initialize the right blocks.
+ # Initialize the output blocks and other projection blocks when necessary.
+ self._init_patched_inputs(norm_type=norm_type)
+
+ # parallel
+ self.parallel_manager: ParallelManager = None
+
+ def _init_patched_inputs(self, norm_type):
+ assert self.config.sample_size_t is not None, "OpenSoraT2V over patched input must provide sample_size_t"
+ assert self.config.sample_size is not None, "OpenSoraT2V over patched input must provide sample_size"
+ # assert not (self.config.sample_size_t == 1 and self.config.patch_size_t == 2), "Image do not need patchfy in t-dim"
+
+ self.num_frames = self.config.sample_size_t
+ self.config.sample_size = to_2tuple(self.config.sample_size)
+ self.height = self.config.sample_size[0]
+ self.width = self.config.sample_size[1]
+ self.patch_size_t = self.config.patch_size_t
+ self.patch_size = self.config.patch_size
+ interpolation_scale_t = (
+ ((self.config.sample_size_t - 1) // 16 + 1)
+ if self.config.sample_size_t % 2 == 1
+ else self.config.sample_size_t / 16
+ )
+ interpolation_scale_t = (
+ self.config.interpolation_scale_t
+ if self.config.interpolation_scale_t is not None
+ else interpolation_scale_t
+ )
+ interpolation_scale = (
+ self.config.interpolation_scale_h
+ if self.config.interpolation_scale_h is not None
+ else self.config.sample_size[0] / 30,
+ self.config.interpolation_scale_w
+ if self.config.interpolation_scale_w is not None
+ else self.config.sample_size[1] / 40,
+ )
+ # if self.config.sample_size_t > 1:
+ # self.pos_embed = PatchEmbed3D(
+ # num_frames=self.config.sample_size_t,
+ # height=self.config.sample_size[0],
+ # width=self.config.sample_size[1],
+ # patch_size_t=self.config.patch_size_t,
+ # patch_size=self.config.patch_size,
+ # in_channels=self.in_channels,
+ # embed_dim=self.inner_dim,
+ # interpolation_scale=interpolation_scale,
+ # interpolation_scale_t=interpolation_scale_t,
+ # )
+ # else:
+ if self.config.downsampler is not None and len(self.config.downsampler) == 9:
+ self.pos_embed = OverlapPatchEmbed3D(
+ num_frames=self.config.sample_size_t,
+ height=self.config.sample_size[0],
+ width=self.config.sample_size[1],
+ patch_size_t=self.config.patch_size_t,
+ patch_size=self.config.patch_size,
+ in_channels=self.in_channels,
+ embed_dim=self.inner_dim,
+ interpolation_scale=interpolation_scale,
+ interpolation_scale_t=interpolation_scale_t,
+ use_abs_pos=not self.config.use_rope,
+ )
+ elif self.config.downsampler is not None and len(self.config.downsampler) == 7:
+ self.pos_embed = OverlapPatchEmbed2D(
+ num_frames=self.config.sample_size_t,
+ height=self.config.sample_size[0],
+ width=self.config.sample_size[1],
+ patch_size_t=self.config.patch_size_t,
+ patch_size=self.config.patch_size,
+ in_channels=self.in_channels,
+ embed_dim=self.inner_dim,
+ interpolation_scale=interpolation_scale,
+ interpolation_scale_t=interpolation_scale_t,
+ use_abs_pos=not self.config.use_rope,
+ )
+
+ else:
+ self.pos_embed = PatchEmbed2D(
+ num_frames=self.config.sample_size_t,
+ height=self.config.sample_size[0],
+ width=self.config.sample_size[1],
+ patch_size_t=self.config.patch_size_t,
+ patch_size=self.config.patch_size,
+ in_channels=self.in_channels,
+ embed_dim=self.inner_dim,
+ interpolation_scale=interpolation_scale,
+ interpolation_scale_t=interpolation_scale_t,
+ use_abs_pos=not self.config.use_rope,
+ )
+ interpolation_scale_thw = (interpolation_scale_t, *interpolation_scale)
+ self.transformer_blocks = nn.ModuleList(
+ [
+ BasicTransformerBlock(
+ self.inner_dim,
+ self.config.num_attention_heads,
+ self.config.attention_head_dim,
+ dropout=self.config.dropout,
+ cross_attention_dim=self.config.cross_attention_dim,
+ activation_fn=self.config.activation_fn,
+ num_embeds_ada_norm=self.config.num_embeds_ada_norm,
+ attention_bias=self.config.attention_bias,
+ only_cross_attention=self.config.only_cross_attention,
+ double_self_attention=self.config.double_self_attention,
+ upcast_attention=self.config.upcast_attention,
+ norm_type=norm_type,
+ norm_elementwise_affine=self.config.norm_elementwise_affine,
+ norm_eps=self.config.norm_eps,
+ attention_type=self.config.attention_type,
+ attention_mode=self.config.attention_mode,
+ downsampler=self.config.downsampler,
+ use_rope=self.config.use_rope,
+ interpolation_scale_thw=interpolation_scale_thw,
+ )
+ for _ in range(self.config.num_layers)
+ ]
+ )
+
+ if self.config.norm_type != "ada_norm_single":
+ self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
+ self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim)
+ self.proj_out_2 = nn.Linear(
+ self.inner_dim,
+ self.config.patch_size_t * self.config.patch_size * self.config.patch_size * self.out_channels,
+ )
+ elif self.config.norm_type == "ada_norm_single":
+ self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
+ self.scale_shift_table = nn.Parameter(torch.randn(2, self.inner_dim) / self.inner_dim**0.5)
+ self.proj_out = nn.Linear(
+ self.inner_dim,
+ self.config.patch_size_t * self.config.patch_size * self.config.patch_size * self.out_channels,
+ )
+
+ # PixArt-Alpha blocks.
+ self.adaln_single = None
+ if self.config.norm_type == "ada_norm_single":
+ # TODO(Sayak, PVP) clean this, for now we use sample size to determine whether to use
+ # additional conditions until we find better name
+ self.adaln_single = AdaLayerNormSingle(
+ self.inner_dim, use_additional_conditions=self.use_additional_conditions
+ )
+
+ self.caption_projection = None
+ if self.caption_channels is not None:
+ self.caption_projection = PixArtAlphaTextProjection(
+ in_features=self.caption_channels, hidden_size=self.inner_dim
+ )
+
+ def enable_parallel(self, dp_size, sp_size, enable_cp):
+ # update cfg parallel
+ if enable_cp and sp_size % 2 == 0:
+ sp_size = sp_size // 2
+ cp_size = 2
+ else:
+ cp_size = 1
+
+ self.parallel_manager = ParallelManager(dp_size, cp_size, sp_size)
+
+ for _, module in self.named_modules():
+ if hasattr(module, "parallel_manager"):
+ module.parallel_manager = self.parallel_manager
+
+ def _set_gradient_checkpointing(self, module, value=False):
+ if hasattr(module, "gradient_checkpointing"):
+ module.gradient_checkpointing = value
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ timestep: Optional[torch.LongTensor] = None,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ added_cond_kwargs: Dict[str, torch.Tensor] = None,
+ class_labels: Optional[torch.LongTensor] = None,
+ cross_attention_kwargs: Dict[str, Any] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ encoder_attention_mask: Optional[torch.Tensor] = None,
+ use_image_num: Optional[int] = 0,
+ return_dict: bool = True,
+ **kwargs,
+ ):
+ """
+ The [`Transformer2DModel`] forward method.
+
+ Args:
+ hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
+ Input `hidden_states`.
+ encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
+ Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
+ self-attention.
+ timestep ( `torch.LongTensor`, *optional*):
+ Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
+ class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
+ Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
+ `AdaLayerZeroNorm`.
+ cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+ `self.processor` in
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+ attention_mask ( `torch.Tensor`, *optional*):
+ An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
+ is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
+ negative values to the attention scores corresponding to "discard" tokens.
+ encoder_attention_mask ( `torch.Tensor`, *optional*):
+ Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
+
+ * Mask `(batch, sequence_length)` True = keep, False = discard.
+ * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
+
+ If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
+ above. This bias will be added to the cross-attention scores.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
+ tuple.
+
+ Returns:
+ If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+ `tuple` where the first element is the sample tensor.
+ """
+ batch_size, c, frame, h, w = hidden_states.shape
+ # print('hidden_states.shape', hidden_states.shape)
+ frame = frame - use_image_num # 21-4=17
+ if cross_attention_kwargs is not None:
+ if cross_attention_kwargs.get("scale", None) is not None:
+ print.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
+ # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
+ # we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
+ # we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
+ # expects mask of shape:
+ # [batch, key_tokens]
+ # adds singleton query_tokens dimension:
+ # [batch, 1, key_tokens]
+ # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
+ # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)
+ # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
+ attention_mask_vid, attention_mask_img = None, None
+ if attention_mask is not None and attention_mask.ndim == 4:
+ # assume that mask is expressed as:
+ # (1 = keep, 0 = discard)
+ # convert mask into a bias that can be added to attention scores:
+ # (keep = +0, discard = -10000.0)
+ # b, frame+use_image_num, h, w -> a video with images
+ # b, 1, h, w -> only images
+ attention_mask = attention_mask.to(self.dtype)
+ if self.parallel_manager.sp_size > 1:
+ attention_mask_vid = attention_mask[:, : frame * self.parallel_manager.sp_size] # b, frame, h, w
+ attention_mask_img = attention_mask[
+ :, frame * self.parallel_manager.sp_size :
+ ] # b, use_image_num, h, w
+ else:
+ attention_mask_vid = attention_mask[:, :frame] # b, frame, h, w
+ attention_mask_img = attention_mask[:, frame:] # b, use_image_num, h, w
+
+ if attention_mask_vid.numel() > 0:
+ attention_mask_vid_first_frame = attention_mask_vid[:, :1].repeat(1, self.patch_size_t - 1, 1, 1)
+ attention_mask_vid = torch.cat([attention_mask_vid_first_frame, attention_mask_vid], dim=1)
+ attention_mask_vid = attention_mask_vid.unsqueeze(1) # b 1 t h w
+ attention_mask_vid = F.max_pool3d(
+ attention_mask_vid,
+ kernel_size=(self.patch_size_t, self.patch_size, self.patch_size),
+ stride=(self.patch_size_t, self.patch_size, self.patch_size),
+ )
+ attention_mask_vid = rearrange(attention_mask_vid, "b 1 t h w -> (b 1) 1 (t h w)")
+ if attention_mask_img.numel() > 0:
+ attention_mask_img = F.max_pool2d(
+ attention_mask_img,
+ kernel_size=(self.patch_size, self.patch_size),
+ stride=(self.patch_size, self.patch_size),
+ )
+ attention_mask_img = rearrange(attention_mask_img, "b i h w -> (b i) 1 (h w)")
+
+ attention_mask_vid = (
+ (1 - attention_mask_vid.bool().to(self.dtype)) * -10000.0 if attention_mask_vid.numel() > 0 else None
+ )
+ attention_mask_img = (
+ (1 - attention_mask_img.bool().to(self.dtype)) * -10000.0 if attention_mask_img.numel() > 0 else None
+ )
+
+ if frame == 1 and use_image_num == 0 and not (self.parallel_manager.sp_size > 1):
+ attention_mask_img = attention_mask_vid
+ attention_mask_vid = None
+ # convert encoder_attention_mask to a bias the same way we do for attention_mask
+ if encoder_attention_mask is not None and encoder_attention_mask.ndim == 3:
+ # b, 1+use_image_num, l -> a video with images
+ # b, 1, l -> only images
+ encoder_attention_mask = (1 - encoder_attention_mask.to(self.dtype)) * -10000.0
+ in_t = encoder_attention_mask.shape[1]
+ encoder_attention_mask_vid = encoder_attention_mask[:, : in_t - use_image_num] # b, 1, l
+ encoder_attention_mask_vid = (
+ rearrange(encoder_attention_mask_vid, "b 1 l -> (b 1) 1 l")
+ if encoder_attention_mask_vid.numel() > 0
+ else None
+ )
+
+ encoder_attention_mask_img = encoder_attention_mask[:, in_t - use_image_num :] # b, use_image_num, l
+ encoder_attention_mask_img = (
+ rearrange(encoder_attention_mask_img, "b i l -> (b i) 1 l")
+ if encoder_attention_mask_img.numel() > 0
+ else None
+ )
+
+ if frame == 1 and use_image_num == 0 and not (self.parallel_manager.sp_size > 1):
+ encoder_attention_mask_img = encoder_attention_mask_vid
+ encoder_attention_mask_vid = None
+
+ if npu_config is not None and attention_mask_vid is not None:
+ attention_mask_vid = npu_config.get_attention_mask(attention_mask_vid, attention_mask_vid.shape[-1])
+ encoder_attention_mask_vid = npu_config.get_attention_mask(
+ encoder_attention_mask_vid, attention_mask_vid.shape[-2]
+ )
+ if npu_config is not None and attention_mask_img is not None:
+ attention_mask_img = npu_config.get_attention_mask(attention_mask_img, attention_mask_img.shape[-1])
+ encoder_attention_mask_img = npu_config.get_attention_mask(
+ encoder_attention_mask_img, attention_mask_img.shape[-2]
+ )
+
+ # 1. Input
+ frame = ((frame - 1) // self.patch_size_t + 1) if frame % 2 == 1 else frame // self.patch_size_t # patchfy
+ # print('frame', frame)
+ height, width = hidden_states.shape[-2] // self.patch_size, hidden_states.shape[-1] // self.patch_size
+
+ added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
+ (
+ hidden_states_vid,
+ hidden_states_img,
+ encoder_hidden_states_vid,
+ encoder_hidden_states_img,
+ timestep_vid,
+ timestep_img,
+ embedded_timestep_vid,
+ embedded_timestep_img,
+ ) = self._operate_on_patched_inputs(
+ hidden_states, encoder_hidden_states, timestep, added_cond_kwargs, batch_size, frame, use_image_num
+ )
+ # 2. Blocks
+ if self.parallel_manager.sp_size > 1:
+ if hidden_states_vid is not None:
+ hidden_states_vid = rearrange(hidden_states_vid, "b s h -> s b h", b=batch_size).contiguous()
+ encoder_hidden_states_vid = rearrange(
+ encoder_hidden_states_vid, "b s h -> s b h", b=batch_size
+ ).contiguous()
+ timestep_vid = timestep_vid.view(batch_size, 6, -1).transpose(0, 1).contiguous()
+
+ for block in self.transformer_blocks:
+ if self.training and self.gradient_checkpointing:
+
+ def create_custom_forward(module, return_dict=None):
+ def custom_forward(*inputs):
+ if return_dict is not None:
+ return module(*inputs, return_dict=return_dict)
+ else:
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ if hidden_states_vid is not None:
+ hidden_states_vid = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ hidden_states_vid,
+ attention_mask_vid,
+ encoder_hidden_states_vid,
+ encoder_attention_mask_vid,
+ timestep_vid,
+ cross_attention_kwargs,
+ class_labels,
+ frame,
+ height,
+ width,
+ **ckpt_kwargs,
+ )
+ if hidden_states_img is not None:
+ hidden_states_img = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ hidden_states_img,
+ attention_mask_img,
+ encoder_hidden_states_img,
+ encoder_attention_mask_img,
+ timestep_img,
+ cross_attention_kwargs,
+ class_labels,
+ 1,
+ height,
+ width,
+ **ckpt_kwargs,
+ )
+ else:
+ if hidden_states_vid is not None:
+ hidden_states_vid = block(
+ hidden_states_vid,
+ attention_mask=attention_mask_vid,
+ encoder_hidden_states=encoder_hidden_states_vid,
+ encoder_attention_mask=encoder_attention_mask_vid,
+ timestep=timestep_vid,
+ cross_attention_kwargs=cross_attention_kwargs,
+ class_labels=class_labels,
+ frame=frame,
+ height=height,
+ width=width,
+ )
+ if hidden_states_img is not None:
+ hidden_states_img = block(
+ hidden_states_img,
+ attention_mask=attention_mask_img,
+ encoder_hidden_states=encoder_hidden_states_img,
+ encoder_attention_mask=encoder_attention_mask_img,
+ timestep=timestep_img,
+ cross_attention_kwargs=cross_attention_kwargs,
+ class_labels=class_labels,
+ frame=1,
+ height=height,
+ width=width,
+ )
+
+ if self.parallel_manager.sp_size > 1:
+ if hidden_states_vid is not None:
+ hidden_states_vid = rearrange(hidden_states_vid, "s b h -> b s h", b=batch_size).contiguous()
+
+ # 3. Output
+ output_vid, output_img = None, None
+ if hidden_states_vid is not None:
+ output_vid = self._get_output_for_patched_inputs(
+ hidden_states=hidden_states_vid,
+ timestep=timestep_vid,
+ class_labels=class_labels,
+ embedded_timestep=embedded_timestep_vid,
+ num_frames=frame,
+ height=height,
+ width=width,
+ ) # b c t h w
+ if hidden_states_img is not None:
+ output_img = self._get_output_for_patched_inputs(
+ hidden_states=hidden_states_img,
+ timestep=timestep_img,
+ class_labels=class_labels,
+ embedded_timestep=embedded_timestep_img,
+ num_frames=1,
+ height=height,
+ width=width,
+ ) # b c 1 h w
+ if use_image_num != 0:
+ output_img = rearrange(output_img, "(b i) c 1 h w -> b c i h w", i=use_image_num)
+
+ if output_vid is not None and output_img is not None:
+ output = torch.cat([output_vid, output_img], dim=2)
+ elif output_vid is not None:
+ output = output_vid
+ elif output_img is not None:
+ output = output_img
+
+ if not return_dict:
+ return (output,)
+
+ return VideoSysPipelineOutput(video=output)
+
+ def _operate_on_patched_inputs(
+ self, hidden_states, encoder_hidden_states, timestep, added_cond_kwargs, batch_size, frame, use_image_num
+ ):
+ # batch_size = hidden_states.shape[0]
+ hidden_states_vid, hidden_states_img = self.pos_embed(hidden_states.to(self.dtype), frame)
+ timestep_vid, timestep_img = None, None
+ embedded_timestep_vid, embedded_timestep_img = None, None
+ encoder_hidden_states_vid, encoder_hidden_states_img = None, None
+
+ if self.adaln_single is not None:
+ if self.use_additional_conditions and added_cond_kwargs is None:
+ raise ValueError(
+ "`added_cond_kwargs` cannot be None when using additional conditions for `adaln_single`."
+ )
+ timestep, embedded_timestep = self.adaln_single(
+ timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=self.dtype
+ ) # b 6d, b d
+ if hidden_states_vid is None:
+ timestep_img = timestep
+ embedded_timestep_img = embedded_timestep
+ else:
+ timestep_vid = timestep
+ embedded_timestep_vid = embedded_timestep
+ if hidden_states_img is not None:
+ timestep_img = repeat(timestep, "b d -> (b i) d", i=use_image_num).contiguous()
+ embedded_timestep_img = repeat(embedded_timestep, "b d -> (b i) d", i=use_image_num).contiguous()
+
+ if self.caption_projection is not None:
+ encoder_hidden_states = self.caption_projection(
+ encoder_hidden_states
+ ) # b, 1+use_image_num, l, d or b, 1, l, d
+ if hidden_states_vid is None:
+ encoder_hidden_states_img = rearrange(encoder_hidden_states, "b 1 l d -> (b 1) l d")
+ else:
+ encoder_hidden_states_vid = rearrange(encoder_hidden_states[:, :1], "b 1 l d -> (b 1) l d")
+ if hidden_states_img is not None:
+ encoder_hidden_states_img = rearrange(encoder_hidden_states[:, 1:], "b i l d -> (b i) l d")
+
+ return (
+ hidden_states_vid,
+ hidden_states_img,
+ encoder_hidden_states_vid,
+ encoder_hidden_states_img,
+ timestep_vid,
+ timestep_img,
+ embedded_timestep_vid,
+ embedded_timestep_img,
+ )
+
+ def _get_output_for_patched_inputs(
+ self, hidden_states, timestep, class_labels, embedded_timestep, num_frames, height=None, width=None
+ ):
+ # import ipdb;ipdb.set_trace()
+ if self.config.norm_type != "ada_norm_single":
+ conditioning = self.transformer_blocks[0].norm1.emb(timestep, class_labels, hidden_dtype=self.dtype)
+ shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
+ hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+ hidden_states = self.proj_out_2(hidden_states)
+ elif self.config.norm_type == "ada_norm_single":
+ shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1)
+ hidden_states = self.norm_out(hidden_states)
+ # Modulation
+ hidden_states = hidden_states * (1 + scale) + shift
+ hidden_states = self.proj_out(hidden_states)
+ hidden_states = hidden_states.squeeze(1)
+
+ # unpatchify
+ if self.adaln_single is None:
+ height = width = int(hidden_states.shape[1] ** 0.5)
+ hidden_states = hidden_states.reshape(
+ shape=(
+ -1,
+ num_frames,
+ height,
+ width,
+ self.patch_size_t,
+ self.patch_size,
+ self.patch_size,
+ self.out_channels,
+ )
+ )
+ hidden_states = torch.einsum("nthwopqc->nctohpwq", hidden_states)
+ output = hidden_states.reshape(
+ shape=(
+ -1,
+ self.out_channels,
+ num_frames * self.patch_size_t,
+ height * self.patch_size,
+ width * self.patch_size,
+ )
+ )
+ # import ipdb;ipdb.set_trace()
+ # if output.shape[2] % 2 == 0:
+ # output = output[:, :, 1:]
+ return output
+
+
+def OpenSoraT2V_S_122(**kwargs):
+ return OpenSoraT2V(
+ num_layers=28,
+ attention_head_dim=96,
+ num_attention_heads=16,
+ patch_size_t=1,
+ patch_size=2,
+ norm_type="ada_norm_single",
+ caption_channels=4096,
+ cross_attention_dim=1536,
+ **kwargs,
+ )
+
+
+def OpenSoraT2V_B_122(**kwargs):
+ return OpenSoraT2V(
+ num_layers=32,
+ attention_head_dim=96,
+ num_attention_heads=16,
+ patch_size_t=1,
+ patch_size=2,
+ norm_type="ada_norm_single",
+ caption_channels=4096,
+ cross_attention_dim=1920,
+ **kwargs,
+ )
+
+
+def OpenSoraT2V_L_122(**kwargs):
+ return OpenSoraT2V(
+ num_layers=40,
+ attention_head_dim=128,
+ num_attention_heads=16,
+ patch_size_t=1,
+ patch_size=2,
+ norm_type="ada_norm_single",
+ caption_channels=4096,
+ cross_attention_dim=2048,
+ **kwargs,
+ )
+
+
+def OpenSoraT2V_ROPE_L_122(**kwargs):
+ return OpenSoraT2V(
+ num_layers=32,
+ attention_head_dim=96,
+ num_attention_heads=24,
+ patch_size_t=1,
+ patch_size=2,
+ norm_type="ada_norm_single",
+ caption_channels=4096,
+ cross_attention_dim=2304,
+ **kwargs,
+ )
+
+
+OpenSora_models = {
+ "OpenSoraT2V-S/122": OpenSoraT2V_S_122, # 1.1B
+ "OpenSoraT2V-B/122": OpenSoraT2V_B_122,
+ "OpenSoraT2V-L/122": OpenSoraT2V_L_122,
+ "OpenSoraT2V-ROPE-L/122": OpenSoraT2V_ROPE_L_122,
+}
+
+OpenSora_models_class = {
+ "OpenSoraT2V-S/122": OpenSoraT2V,
+ "OpenSoraT2V-B/122": OpenSoraT2V,
+ "OpenSoraT2V-L/122": OpenSoraT2V,
+ "OpenSoraT2V-ROPE-L/122": OpenSoraT2V,
+}
diff --git a/videosys/pipelines/open_sora_plan/pipeline_open_sora_plan.py b/videosys/pipelines/open_sora_plan/pipeline_open_sora_plan.py
index 271345b3..bdef12c4 100644
--- a/videosys/pipelines/open_sora_plan/pipeline_open_sora_plan.py
+++ b/videosys/pipelines/open_sora_plan/pipeline_open_sora_plan.py
@@ -20,18 +20,23 @@
import tqdm
from bs4 import BeautifulSoup
from diffusers.models import AutoencoderKL, Transformer2DModel
-from diffusers.schedulers import PNDMScheduler
+from diffusers.schedulers import EulerAncestralDiscreteScheduler, PNDMScheduler
from diffusers.utils.torch_utils import randn_tensor
-from transformers import T5EncoderModel, T5Tokenizer
+from transformers import AutoTokenizer, MT5EncoderModel, T5EncoderModel, T5Tokenizer
from videosys.core.pab_mgr import PABConfig, set_pab_manager, update_steps
from videosys.core.pipeline import VideoSysPipeline, VideoSysPipelineOutput
+from videosys.models.autoencoders.autoencoder_kl_open_sora_plan_v110 import (
+ CausalVAEModelWrapper as CausalVAEModelWrapperV110,
+)
+from videosys.models.autoencoders.autoencoder_kl_open_sora_plan_v120 import (
+ CausalVAEModelWrapper as CausalVAEModelWrapperV120,
+)
+from videosys.models.transformers.open_sora_plan_v110_transformer_3d import LatteT2V
+from videosys.models.transformers.open_sora_plan_v120_transformer_3d import OpenSoraT2V
from videosys.utils.logging import logger
from videosys.utils.utils import save_video, set_seed
-from ...models.autoencoders.autoencoder_kl_open_sora_plan import ae_stride_config, getae_wrapper
-from ...models.transformers.open_sora_plan_transformer_3d import LatteT2V
-
class OpenSoraPlanPABConfig(PABConfig):
def __init__(
@@ -145,10 +150,10 @@ class OpenSoraPlanConfig:
def __init__(
self,
- transformer: str = "LanguageBind/Open-Sora-Plan-v1.1.0",
- ae: str = "CausalVAEModel_4x8x8",
- text_encoder: str = "DeepFloyd/t5-v1_1-xxl",
- num_frames: int = 65,
+ version: str = "v120",
+ transformer_type: str = "93x480p",
+ transformer: str = None,
+ text_encoder: str = None,
# ======= distributed ========
num_gpus: int = 1,
# ======= memory =======
@@ -160,12 +165,29 @@ def __init__(
pab_config: PABConfig = OpenSoraPlanPABConfig(),
):
self.pipeline_cls = OpenSoraPlanPipeline
- self.ae = ae
- self.text_encoder = text_encoder
- self.transformer = transformer
- assert num_frames in [65, 221], "num_frames must be one of [65, 221]"
- self.num_frames = num_frames
- self.version = f"{num_frames}x512x512"
+
+ # get version
+ assert version in ["v110", "v120"], f"Unknown Open-Sora-Plan version: {version}"
+ self.version = version
+ self.transformer_type = transformer_type
+
+ # check transformer_type
+ if version == "v110":
+ assert transformer_type in ["65x512x512", "221x512x512"]
+ elif version == "v120":
+ assert transformer_type in ["93x480p", "93x720p", "29x480p", "29x720p"]
+ self.num_frames = int(transformer_type.split("x")[0])
+
+ # set default values according to version
+ if version == "v110":
+ transformer_default = "LanguageBind/Open-Sora-Plan-v1.1.0"
+ text_encoder_default = "DeepFloyd/t5-v1_1-xxl"
+ elif version == "v120":
+ transformer_default = "LanguageBind/Open-Sora-Plan-v1.2.0"
+ text_encoder_default = "google/mt5-xxl"
+ self.text_encoder = text_encoder or text_encoder_default
+ self.transformer = transformer or transformer_default
+
# ======= distributed ========
self.num_gpus = num_gpus
# ======= memory ========
@@ -220,23 +242,64 @@ def __init__(
super().__init__()
self._config = config
+ # not implemented
+ if config.version == "v120":
+ if config.num_gpus > 1:
+ raise NotImplementedError("v120 does not support multi-gpu inference")
+ if config.enable_pab:
+ raise NotImplementedError("v120 does not support PAB")
+
# init
if tokenizer is None:
- tokenizer = T5Tokenizer.from_pretrained(config.text_encoder)
+ if config.version == "v110":
+ tokenizer = T5Tokenizer.from_pretrained(config.text_encoder)
+ elif config.version == "v120":
+ tokenizer = AutoTokenizer.from_pretrained(config.text_encoder)
+
if text_encoder is None:
- text_encoder = T5EncoderModel.from_pretrained(config.text_encoder, torch_dtype=torch.float16)
+ if config.version == "v110":
+ text_encoder = T5EncoderModel.from_pretrained(config.text_encoder, torch_dtype=dtype)
+ elif config.version == "v120":
+ text_encoder = MT5EncoderModel.from_pretrained(
+ config.text_encoder, low_cpu_mem_usage=True, torch_dtype=dtype
+ )
+
if vae is None:
- vae = getae_wrapper(config.ae)(config.transformer, subfolder="vae").to(dtype=dtype)
+ if config.version == "v110":
+ vae = CausalVAEModelWrapperV110(config.transformer, subfolder="vae").to(dtype=dtype)
+ elif config.version == "v120":
+ vae = CausalVAEModelWrapperV120(config.transformer, subfolder="vae").to(dtype=dtype)
+
if transformer is None:
- transformer = LatteT2V.from_pretrained(config.transformer, subfolder=config.version, torch_dtype=dtype)
+ if config.version == "v110":
+ transformer = LatteT2V.from_pretrained(
+ config.transformer, subfolder=config.transformer_type, torch_dtype=dtype
+ )
+ elif config.version == "v120":
+ transformer = OpenSoraT2V.from_pretrained(
+ config.transformer, subfolder=config.transformer_type, torch_dtype=dtype
+ )
+
if scheduler is None:
- scheduler = PNDMScheduler()
+ if config.version == "v110":
+ scheduler = PNDMScheduler()
+ elif config.version == "v120":
+ scheduler = EulerAncestralDiscreteScheduler()
# setting
if config.enable_tiling:
vae.vae.enable_tiling()
vae.vae.tile_overlap_factor = config.tile_overlap_factor
- vae.vae_scale_factor = ae_stride_config[config.ae]
+ vae.vae.tile_sample_min_size = 512
+ vae.vae.tile_latent_min_size = 64
+ vae.vae.tile_sample_min_size_t = 29
+ vae.vae.tile_latent_min_size_t = 8
+ # if low_mem:
+ # vae.vae.tile_sample_min_size = 256
+ # vae.vae.tile_latent_min_size = 32
+ # vae.vae.tile_sample_min_size_t = 29
+ # vae.vae.tile_latent_min_size_t = 8
+ vae.vae_scale_factor = [4, 8, 8]
transformer.force_images = False
# pab
@@ -453,6 +516,143 @@ def encode_prompt(
return prompt_embeds, negative_prompt_embeds
+ def encode_prompt_v120(
+ self,
+ prompt: Union[str, List[str]],
+ do_classifier_free_guidance: bool = True,
+ negative_prompt: str = "",
+ num_images_per_prompt: int = 1,
+ device: Optional[torch.device] = None,
+ prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_attention_mask: Optional[torch.FloatTensor] = None,
+ negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
+ clean_caption: bool = False,
+ max_sequence_length: int = 120,
+ **kwargs,
+ ):
+ r"""
+ Encodes the prompt into text encoder hidden states.
+
+ Args:
+ prompt (`str` or `List[str]`, *optional*):
+ prompt to be encoded
+ negative_prompt (`str` or `List[str]`, *optional*):
+ The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
+ instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
+ PixArt-Alpha, this should be "".
+ do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+ whether to use classifier free guidance or not
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ number of images that should be generated per prompt
+ device: (`torch.device`, *optional*):
+ torch device to place the resulting embeddings on
+ prompt_embeds (`torch.FloatTensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+ Pre-generated negative text embeddings. For PixArt-Alpha, it's should be the embeddings of the ""
+ string.
+ clean_caption (`bool`, defaults to `False`):
+ If `True`, the function will preprocess and clean the provided caption before encoding.
+ max_sequence_length (`int`, defaults to 120): Maximum sequence length to use for the prompt.
+ """
+ if device is None:
+ device = getattr(self, "_execution_device", None) or getattr(self, "device", None) or torch.device("cuda")
+
+ if prompt is not None and isinstance(prompt, str):
+ batch_size = 1
+ elif prompt is not None and isinstance(prompt, list):
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ # See Section 3.1. of the paper.
+ max_length = max_sequence_length
+
+ if prompt_embeds is None:
+ prompt = self._text_preprocessing(prompt, clean_caption=clean_caption)
+ text_inputs = self.tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=max_length,
+ truncation=True,
+ add_special_tokens=True,
+ return_tensors="pt",
+ )
+ text_input_ids = text_inputs.input_ids
+ untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+ text_input_ids, untruncated_ids
+ ):
+ removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_length - 1 : -1])
+ logger.warning(
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
+ f" {max_length} tokens: {removed_text}"
+ )
+
+ prompt_attention_mask = text_inputs.attention_mask
+ prompt_attention_mask = prompt_attention_mask.to(device)
+
+ prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)
+ prompt_embeds = prompt_embeds[0]
+
+ if self.text_encoder is not None:
+ dtype = self.text_encoder.dtype
+ elif self.transformer is not None:
+ dtype = self.transformer.dtype
+ else:
+ dtype = None
+
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+ bs_embed, seq_len, _ = prompt_embeds.shape
+ # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+ prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+ prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1)
+ prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
+
+ # get unconditional embeddings for classifier free guidance
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
+ uncond_tokens = [negative_prompt] * batch_size
+ uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption)
+ max_length = prompt_embeds.shape[1]
+ uncond_input = self.tokenizer(
+ uncond_tokens,
+ padding="max_length",
+ max_length=max_length,
+ truncation=True,
+ return_attention_mask=True,
+ add_special_tokens=True,
+ return_tensors="pt",
+ )
+ negative_prompt_attention_mask = uncond_input.attention_mask
+ negative_prompt_attention_mask = negative_prompt_attention_mask.to(device)
+
+ negative_prompt_embeds = self.text_encoder(
+ uncond_input.input_ids.to(device), attention_mask=negative_prompt_attention_mask
+ )
+ negative_prompt_embeds = negative_prompt_embeds[0]
+
+ if do_classifier_free_guidance:
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+ seq_len = negative_prompt_embeds.shape[1]
+
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
+
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+ negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1)
+ negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
+ else:
+ negative_prompt_embeds = None
+ negative_prompt_attention_mask = None
+
+ return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
+
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
def prepare_extra_step_kwargs(self, generator, eta):
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
@@ -480,6 +680,8 @@ def check_inputs(
callback_steps,
prompt_embeds=None,
negative_prompt_embeds=None,
+ prompt_attention_mask=None,
+ negative_prompt_attention_mask=None,
):
if height % 8 != 0 or width % 8 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
@@ -523,6 +725,12 @@ def check_inputs(
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
f" {negative_prompt_embeds.shape}."
)
+ if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
+ raise ValueError(
+ "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
+ f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
+ f" {negative_prompt_attention_mask.shape}."
+ )
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing
def _text_preprocessing(self, text, clean_caption=False):
@@ -685,7 +893,7 @@ def generate(
self,
prompt: Union[str, List[str]] = None,
negative_prompt: str = "",
- num_inference_steps: int = 150,
+ num_inference_steps: int = 100,
guidance_scale: float = 7.5,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
@@ -693,7 +901,9 @@ def generate(
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
+ prompt_attention_mask: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
@@ -702,6 +912,7 @@ def generate(
mask_feature: bool = True,
enable_temporal_attentions: bool = True,
verbose: bool = True,
+ max_sequence_length: int = 512,
) -> Union[VideoSysPipelineOutput, Tuple]:
"""
Function invoked when calling the pipeline for generation.
@@ -773,11 +984,21 @@ def generate(
returned where the first element is a list with the generated images
"""
# 1. Check inputs. Raise error if not correct
- height = 512
- width = 512
+ height = self.transformer.config.sample_size[0] * self.vae.vae_scale_factor[1]
+ width = self.transformer.config.sample_size[1] * self.vae.vae_scale_factor[2]
num_frames = self._config.num_frames
update_steps(num_inference_steps)
- self.check_inputs(prompt, height, width, negative_prompt, callback_steps, prompt_embeds, negative_prompt_embeds)
+ self.check_inputs(
+ prompt,
+ height,
+ width,
+ negative_prompt,
+ callback_steps,
+ prompt_embeds,
+ negative_prompt_embeds,
+ prompt_attention_mask,
+ negative_prompt_attention_mask,
+ )
self._set_seed(seed)
# 2. Default height and width to transformer
@@ -788,7 +1009,7 @@ def generate(
else:
batch_size = prompt_embeds.shape[0]
- device = self._execution_device
+ device = getattr(self, "_execution_device", None) or getattr(self, "device", None) or torch.device("cuda")
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
@@ -796,23 +1017,50 @@ def generate(
do_classifier_free_guidance = guidance_scale > 1.0
# 3. Encode input prompt
- prompt_embeds, negative_prompt_embeds = self.encode_prompt(
- prompt,
- do_classifier_free_guidance,
- negative_prompt=negative_prompt,
- num_images_per_prompt=num_images_per_prompt,
- device=device,
- prompt_embeds=prompt_embeds,
- negative_prompt_embeds=negative_prompt_embeds,
- clean_caption=clean_caption,
- mask_feature=mask_feature,
- )
+ if self._config.version == "v110":
+ prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+ prompt,
+ do_classifier_free_guidance,
+ negative_prompt=negative_prompt,
+ num_images_per_prompt=num_images_per_prompt,
+ device=device,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ clean_caption=clean_caption,
+ mask_feature=mask_feature,
+ )
+ prompt_attention_mask = None
+ negative_prompt_attention_mask = None
+ else:
+ (
+ prompt_embeds,
+ prompt_attention_mask,
+ negative_prompt_embeds,
+ negative_prompt_attention_mask,
+ ) = self.encode_prompt_v120(
+ prompt,
+ do_classifier_free_guidance,
+ negative_prompt=negative_prompt,
+ num_images_per_prompt=num_images_per_prompt,
+ device=device,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ prompt_attention_mask=prompt_attention_mask,
+ negative_prompt_attention_mask=negative_prompt_attention_mask,
+ clean_caption=clean_caption,
+ max_sequence_length=max_sequence_length,
+ )
if do_classifier_free_guidance:
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+ if prompt_attention_mask is not None and negative_prompt_attention_mask is not None:
+ prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
# 4. Prepare timesteps
- self.scheduler.set_timesteps(num_inference_steps, device=device)
- timesteps = self.scheduler.timesteps
+ if self._config.version == "v110":
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
+ timesteps = self.scheduler.timesteps
+ else:
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, None)
# 5. Prepare latents.
latent_channels = self.transformer.config.in_channels
@@ -833,15 +1081,10 @@ def generate(
# 6.1 Prepare micro-conditions.
added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
- # if self.transformer.config.sample_size == 128:
- # resolution = torch.tensor([height, width]).repeat(batch_size * num_images_per_prompt, 1)
- # aspect_ratio = torch.tensor([float(height / width)]).repeat(batch_size * num_images_per_prompt, 1)
- # resolution = resolution.to(dtype=prompt_embeds.dtype, device=device)
- # aspect_ratio = aspect_ratio.to(dtype=prompt_embeds.dtype, device=device)
- # added_cond_kwargs = {"resolution": resolution, "aspect_ratio": aspect_ratio}
# 7. Denoising loop
num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
progress_wrap = tqdm.tqdm if verbose and dist.get_rank() == 0 else (lambda x: x)
for i, t in progress_wrap(list(enumerate(timesteps))):
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
@@ -862,9 +1105,15 @@ def generate(
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
current_timestep = current_timestep.expand(latent_model_input.shape[0])
- if prompt_embeds.ndim == 3:
+ if prompt_embeds is not None and prompt_embeds.ndim == 3:
prompt_embeds = prompt_embeds.unsqueeze(1) # b l d -> b 1 l d
+ if prompt_attention_mask is not None and prompt_attention_mask.ndim == 2:
+ prompt_attention_mask = prompt_attention_mask.unsqueeze(1) # b l -> b 1 l
+ # prepare attention_mask.
+ # b c t h w -> b t h w
+ attention_mask = torch.ones_like(latent_model_input)[:, 0]
+
# predict noise model_output
noise_pred = self.transformer(
latent_model_input,
@@ -873,6 +1122,8 @@ def generate(
timestep=current_timestep,
added_cond_kwargs=added_cond_kwargs,
enable_temporal_attentions=enable_temporal_attentions,
+ attention_mask=attention_mask,
+ encoder_attention_mask=prompt_attention_mask,
return_dict=False,
)[0]
@@ -912,10 +1163,57 @@ def generate(
return VideoSysPipelineOutput(video=video)
def decode_latents(self, latents):
- video = self.vae(latents) # b t c h w
- # b t c h w -> b t h w c
- video = ((video / 2.0 + 0.5).clamp(0, 1) * 255).to(dtype=torch.uint8).cpu().permute(0, 1, 3, 4, 2).contiguous()
+ video = self.vae(latents.to(self.vae.vae.dtype))
+ video = (
+ ((video / 2.0 + 0.5).clamp(0, 1) * 255).to(dtype=torch.uint8).cpu().permute(0, 1, 3, 4, 2).contiguous()
+ ) # b t h w c
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
return video
def save_video(self, video, output_path):
save_video(video, output_path, fps=24)
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: Optional[int] = None,
+ device: Optional[Union[str, torch.device]] = None,
+ timesteps: Optional[List[int]] = None,
+ **kwargs,
+):
+ """
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`List[int]`, *optional*):
+ Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
+ timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
+ must be `None`.
+
+ Returns:
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps