|
from dataclasses import dataclass
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
|
from diffusers.models import ModelMixin
|
|
from diffusers.utils import BaseOutput
|
|
from diffusers.utils.import_utils import is_xformers_available
|
|
from einops import rearrange, repeat
|
|
from torch import nn
|
|
|
|
from .attention import TemporalBasicTransformerBlock
|
|
|
|
|
|
@dataclass
|
|
class Transformer3DModelOutput(BaseOutput):
|
|
sample: torch.FloatTensor
|
|
|
|
|
|
if is_xformers_available():
|
|
import xformers
|
|
import xformers.ops
|
|
else:
|
|
xformers = None
|
|
|
|
|
|
class Transformer3DModel(ModelMixin, ConfigMixin):
|
|
_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,
|
|
num_layers: int = 1,
|
|
dropout: float = 0.0,
|
|
norm_num_groups: int = 32,
|
|
cross_attention_dim: Optional[int] = None,
|
|
attention_bias: bool = False,
|
|
activation_fn: str = "geglu",
|
|
num_embeds_ada_norm: Optional[int] = None,
|
|
use_linear_projection: bool = False,
|
|
only_cross_attention: bool = False,
|
|
upcast_attention: bool = False,
|
|
unet_use_cross_frame_attention=None,
|
|
unet_use_temporal_attention=None,
|
|
):
|
|
super().__init__()
|
|
self.use_linear_projection = use_linear_projection
|
|
self.num_attention_heads = num_attention_heads
|
|
self.attention_head_dim = attention_head_dim
|
|
inner_dim = num_attention_heads * attention_head_dim
|
|
|
|
|
|
self.in_channels = in_channels
|
|
|
|
self.norm = torch.nn.GroupNorm(
|
|
num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True
|
|
)
|
|
if use_linear_projection:
|
|
self.proj_in = nn.Linear(in_channels, inner_dim)
|
|
else:
|
|
self.proj_in = nn.Conv2d(
|
|
in_channels, inner_dim, kernel_size=1, stride=1, padding=0
|
|
)
|
|
|
|
|
|
self.transformer_blocks = nn.ModuleList(
|
|
[
|
|
TemporalBasicTransformerBlock(
|
|
inner_dim,
|
|
num_attention_heads,
|
|
attention_head_dim,
|
|
dropout=dropout,
|
|
cross_attention_dim=cross_attention_dim,
|
|
activation_fn=activation_fn,
|
|
num_embeds_ada_norm=num_embeds_ada_norm,
|
|
attention_bias=attention_bias,
|
|
only_cross_attention=only_cross_attention,
|
|
upcast_attention=upcast_attention,
|
|
unet_use_cross_frame_attention=unet_use_cross_frame_attention,
|
|
unet_use_temporal_attention=unet_use_temporal_attention,
|
|
)
|
|
for d in range(num_layers)
|
|
]
|
|
)
|
|
|
|
|
|
if use_linear_projection:
|
|
self.proj_out = nn.Linear(in_channels, inner_dim)
|
|
else:
|
|
self.proj_out = nn.Conv2d(
|
|
inner_dim, in_channels, kernel_size=1, stride=1, padding=0
|
|
)
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def _set_gradient_checkpointing(self, module, value=False):
|
|
if hasattr(module, "gradient_checkpointing"):
|
|
module.gradient_checkpointing = value
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
timestep=None,
|
|
return_dict: bool = True,
|
|
):
|
|
|
|
assert (
|
|
hidden_states.dim() == 5
|
|
), f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}."
|
|
video_length = hidden_states.shape[2]
|
|
hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
|
|
if encoder_hidden_states.shape[0] != hidden_states.shape[0]:
|
|
encoder_hidden_states = repeat(
|
|
encoder_hidden_states, "b n c -> (b f) n c", f=video_length
|
|
)
|
|
|
|
batch, channel, height, weight = hidden_states.shape
|
|
residual = hidden_states
|
|
|
|
hidden_states = self.norm(hidden_states)
|
|
if not self.use_linear_projection:
|
|
hidden_states = self.proj_in(hidden_states)
|
|
inner_dim = hidden_states.shape[1]
|
|
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
|
|
batch, height * weight, inner_dim
|
|
)
|
|
else:
|
|
inner_dim = hidden_states.shape[1]
|
|
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
|
|
batch, height * weight, inner_dim
|
|
)
|
|
hidden_states = self.proj_in(hidden_states)
|
|
|
|
|
|
for i, block in enumerate(self.transformer_blocks):
|
|
hidden_states = block(
|
|
hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
timestep=timestep,
|
|
video_length=video_length,
|
|
)
|
|
|
|
|
|
if not self.use_linear_projection:
|
|
hidden_states = (
|
|
hidden_states.reshape(batch, height, weight, inner_dim)
|
|
.permute(0, 3, 1, 2)
|
|
.contiguous()
|
|
)
|
|
hidden_states = self.proj_out(hidden_states)
|
|
else:
|
|
hidden_states = self.proj_out(hidden_states)
|
|
hidden_states = (
|
|
hidden_states.reshape(batch, height, weight, inner_dim)
|
|
.permute(0, 3, 1, 2)
|
|
.contiguous()
|
|
)
|
|
|
|
output = hidden_states + residual
|
|
|
|
output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length)
|
|
if not return_dict:
|
|
return (output,)
|
|
|
|
return Transformer3DModelOutput(sample=output)
|
|
|
