Initial Commit
This commit is contained in:
LucipherDev
@@ -0,0 +1 @@
|
||||
# __init__
|
||||
@@ -0,0 +1,71 @@
|
||||
import torch
|
||||
from typing import Optional
|
||||
from diffusers.models.attention import TemporalBasicTransformerBlock, _chunked_feed_forward
|
||||
from diffusers.utils.torch_utils import maybe_allow_in_graph
|
||||
|
||||
|
||||
@maybe_allow_in_graph
|
||||
class TemporalPoseCondTransformerBlock(TemporalBasicTransformerBlock):
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor, # [bs * num_frame, h * w, c]
|
||||
num_frames: int,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None, # [bs * h * w, 1, c]
|
||||
pose_feature: Optional[torch.FloatTensor] = None, # [bs, c, n_frame, h, w]
|
||||
) -> torch.FloatTensor:
|
||||
# Notice that normalization is always applied before the real computation in the following blocks.
|
||||
# 0. Self-Attention
|
||||
|
||||
batch_frames, seq_length, channels = hidden_states.shape
|
||||
batch_size = batch_frames // num_frames
|
||||
|
||||
hidden_states = hidden_states[None, :].reshape(batch_size, num_frames, seq_length, channels)
|
||||
hidden_states = hidden_states.permute(0, 2, 1, 3)
|
||||
hidden_states = hidden_states.reshape(batch_size * seq_length, num_frames, channels) # [bs * h * w, frame, c]
|
||||
|
||||
residual = hidden_states
|
||||
hidden_states = self.norm_in(hidden_states)
|
||||
|
||||
if self._chunk_size is not None:
|
||||
hidden_states = _chunked_feed_forward(self.ff_in, hidden_states, self._chunk_dim, self._chunk_size)
|
||||
else:
|
||||
hidden_states = self.ff_in(hidden_states)
|
||||
|
||||
if self.is_res:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
norm_hidden_states = self.norm1(hidden_states)
|
||||
if pose_feature is not None:
|
||||
pose_feature = pose_feature.permute(0, 3, 4, 2, 1).reshape(batch_size * seq_length, num_frames, -1)
|
||||
attn_output = self.attn1(norm_hidden_states, encoder_hidden_states=None, pose_feature=pose_feature)
|
||||
else:
|
||||
attn_output = self.attn1(norm_hidden_states, encoder_hidden_states=None)
|
||||
hidden_states = attn_output + hidden_states
|
||||
|
||||
# 3. Cross-Attention
|
||||
if self.attn2 is not None:
|
||||
norm_hidden_states = self.norm2(hidden_states)
|
||||
if pose_feature is not None:
|
||||
attn_output = self.attn2(norm_hidden_states, encoder_hidden_states=encoder_hidden_states, pose_feature=pose_feature)
|
||||
else:
|
||||
attn_output = self.attn2(norm_hidden_states, encoder_hidden_states=encoder_hidden_states)
|
||||
hidden_states = attn_output + hidden_states
|
||||
|
||||
# 4. Feed-forward
|
||||
norm_hidden_states = self.norm3(hidden_states)
|
||||
|
||||
if self._chunk_size is not None:
|
||||
ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)
|
||||
else:
|
||||
ff_output = self.ff(norm_hidden_states)
|
||||
|
||||
if self.is_res:
|
||||
hidden_states = ff_output + hidden_states
|
||||
else:
|
||||
hidden_states = ff_output
|
||||
|
||||
hidden_states = hidden_states[None, :].reshape(batch_size, seq_length, num_frames, channels)
|
||||
hidden_states = hidden_states.permute(0, 2, 1, 3)
|
||||
hidden_states = hidden_states.reshape(batch_size * num_frames, seq_length, channels)
|
||||
|
||||
return hidden_states
|
||||
@@ -0,0 +1,603 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.nn.init as init
|
||||
import logging
|
||||
from diffusers.models.attention import Attention
|
||||
from diffusers.utils import USE_PEFT_BACKEND, is_xformers_available
|
||||
from typing import Optional, Callable
|
||||
|
||||
from einops import rearrange
|
||||
|
||||
if is_xformers_available():
|
||||
import xformers
|
||||
import xformers.ops
|
||||
else:
|
||||
xformers = None
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AttnProcessor:
|
||||
r"""
|
||||
Default processor for performing attention-related computations.
|
||||
"""
|
||||
|
||||
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,
|
||||
scale: float = 1.0,
|
||||
pose_feature=None, # the only difference to the original code
|
||||
) -> torch.Tensor:
|
||||
residual = hidden_states
|
||||
|
||||
args = () if USE_PEFT_BACKEND else (scale,)
|
||||
|
||||
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)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
|
||||
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, *args)
|
||||
|
||||
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, *args)
|
||||
value = attn.to_v(encoder_hidden_states, *args)
|
||||
|
||||
query = attn.head_to_batch_dim(query)
|
||||
key = attn.head_to_batch_dim(key)
|
||||
value = attn.head_to_batch_dim(value)
|
||||
|
||||
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
||||
hidden_states = torch.bmm(attention_probs, value)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states, *args)
|
||||
# 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
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class AttnProcessor2_0:
|
||||
r"""
|
||||
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
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 __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,
|
||||
scale: float = 1.0,
|
||||
pose_feature=None
|
||||
) -> torch.FloatTensor:
|
||||
residual = hidden_states
|
||||
|
||||
args = () if USE_PEFT_BACKEND else (scale,)
|
||||
|
||||
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)
|
||||
|
||||
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, batch_size)
|
||||
# scaled_dot_product_attention expects attention_mask shape to be
|
||||
# (batch, heads, source_length, target_length)
|
||||
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)
|
||||
|
||||
args = () if USE_PEFT_BACKEND else (scale,)
|
||||
query = attn.to_q(hidden_states, *args)
|
||||
|
||||
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, *args)
|
||||
value = attn.to_v(encoder_hidden_states, *args)
|
||||
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // attn.heads
|
||||
|
||||
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)
|
||||
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
# 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 = 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, *args)
|
||||
# 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
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class XFormersAttnProcessor:
|
||||
r"""
|
||||
Processor for implementing memory efficient attention using xFormers.
|
||||
|
||||
Args:
|
||||
attention_op (`Callable`, *optional*, defaults to `None`):
|
||||
The base
|
||||
[operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to
|
||||
use as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best
|
||||
operator.
|
||||
"""
|
||||
|
||||
def __init__(self, attention_op: Optional[Callable] = None):
|
||||
self.attention_op = attention_op
|
||||
|
||||
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,
|
||||
scale: float = 1.0,
|
||||
pose_feature=None, # the only difference to the original code
|
||||
) -> torch.FloatTensor:
|
||||
residual = hidden_states
|
||||
|
||||
args = () if USE_PEFT_BACKEND else (scale,)
|
||||
|
||||
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)
|
||||
|
||||
batch_size, key_tokens, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, key_tokens, batch_size)
|
||||
if attention_mask is not None:
|
||||
# expand our mask's singleton query_tokens dimension:
|
||||
# [batch*heads, 1, key_tokens] ->
|
||||
# [batch*heads, query_tokens, key_tokens]
|
||||
# so that it can be added as a bias onto the attention scores that xformers computes:
|
||||
# [batch*heads, query_tokens, key_tokens]
|
||||
# we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
|
||||
_, query_tokens, _ = hidden_states.shape
|
||||
attention_mask = attention_mask.expand(-1, query_tokens, -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, *args)
|
||||
|
||||
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, *args)
|
||||
value = attn.to_v(encoder_hidden_states, *args)
|
||||
|
||||
query = attn.head_to_batch_dim(query).contiguous()
|
||||
key = attn.head_to_batch_dim(key).contiguous()
|
||||
value = attn.head_to_batch_dim(value).contiguous()
|
||||
|
||||
hidden_states = xformers.ops.memory_efficient_attention(
|
||||
query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale
|
||||
)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states, *args)
|
||||
# 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
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class PoseAdaptorAttnProcessor(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size, # dimension of hidden state
|
||||
pose_feature_dim=None, # dimension of the pose feature
|
||||
cross_attention_dim=None, # dimension of the text embedding
|
||||
query_condition=False,
|
||||
key_value_condition=False,
|
||||
scale=1.0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.pose_feature_dim = pose_feature_dim
|
||||
self.cross_attention_dim = cross_attention_dim
|
||||
self.scale = scale
|
||||
self.query_condition = query_condition
|
||||
self.key_value_condition = key_value_condition
|
||||
assert hidden_size == pose_feature_dim
|
||||
if self.query_condition and self.key_value_condition:
|
||||
self.qkv_merge = nn.Linear(hidden_size, hidden_size)
|
||||
init.zeros_(self.qkv_merge.weight)
|
||||
init.zeros_(self.qkv_merge.bias)
|
||||
elif self.query_condition:
|
||||
self.q_merge = nn.Linear(hidden_size, hidden_size)
|
||||
init.zeros_(self.q_merge.weight)
|
||||
init.zeros_(self.q_merge.bias)
|
||||
else:
|
||||
self.kv_merge = nn.Linear(hidden_size, hidden_size)
|
||||
init.zeros_(self.kv_merge.weight)
|
||||
init.zeros_(self.kv_merge.bias)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
pose_feature,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
scale=None,
|
||||
):
|
||||
assert pose_feature is not None
|
||||
pose_embedding_scale = (scale or self.scale)
|
||||
|
||||
residual = hidden_states
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
assert hidden_states.ndim == 3 and pose_feature.ndim == 3
|
||||
|
||||
if self.query_condition and self.key_value_condition:
|
||||
assert encoder_hidden_states is None
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
|
||||
assert encoder_hidden_states.ndim == 3
|
||||
|
||||
batch_size, ehs_sequence_length, _ = encoder_hidden_states.shape
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, ehs_sequence_length, batch_size)
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
if attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
if self.query_condition and self.key_value_condition: # only self attention
|
||||
query_hidden_state = self.qkv_merge(hidden_states + pose_feature) * pose_embedding_scale + hidden_states
|
||||
key_value_hidden_state = query_hidden_state
|
||||
elif self.query_condition:
|
||||
query_hidden_state = self.q_merge(hidden_states + pose_feature) * pose_embedding_scale + hidden_states
|
||||
key_value_hidden_state = encoder_hidden_states
|
||||
else:
|
||||
key_value_hidden_state = self.kv_merge(encoder_hidden_states + pose_feature) * pose_embedding_scale + encoder_hidden_states
|
||||
query_hidden_state = hidden_states
|
||||
|
||||
# original attention
|
||||
query = attn.to_q(query_hidden_state)
|
||||
key = attn.to_k(key_value_hidden_state)
|
||||
value = attn.to_v(key_value_hidden_state)
|
||||
|
||||
query = attn.head_to_batch_dim(query)
|
||||
key = attn.head_to_batch_dim(key)
|
||||
value = attn.head_to_batch_dim(value)
|
||||
|
||||
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
||||
hidden_states = torch.bmm(attention_probs, value)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class PoseAdaptorAttnProcessor2_0(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size, # dimension of hidden state
|
||||
pose_feature_dim=None, # dimension of the pose feature
|
||||
cross_attention_dim=None, # dimension of the text embedding
|
||||
query_condition=False,
|
||||
key_value_condition=False,
|
||||
scale=1.0,
|
||||
):
|
||||
super().__init__()
|
||||
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.")
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.pose_feature_dim = pose_feature_dim
|
||||
self.cross_attention_dim = cross_attention_dim
|
||||
self.scale = scale
|
||||
self.query_condition = query_condition
|
||||
self.key_value_condition = key_value_condition
|
||||
assert hidden_size == pose_feature_dim
|
||||
if self.query_condition and self.key_value_condition:
|
||||
self.qkv_merge = nn.Linear(hidden_size, hidden_size)
|
||||
init.zeros_(self.qkv_merge.weight)
|
||||
init.zeros_(self.qkv_merge.bias)
|
||||
elif self.query_condition:
|
||||
self.q_merge = nn.Linear(hidden_size, hidden_size)
|
||||
init.zeros_(self.q_merge.weight)
|
||||
init.zeros_(self.q_merge.bias)
|
||||
else:
|
||||
self.kv_merge = nn.Linear(hidden_size, hidden_size)
|
||||
init.zeros_(self.kv_merge.weight)
|
||||
init.zeros_(self.kv_merge.bias)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
pose_feature,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
scale=None,
|
||||
):
|
||||
assert pose_feature is not None
|
||||
pose_embedding_scale = (scale or self.scale)
|
||||
|
||||
residual = hidden_states
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
assert hidden_states.ndim == 3 and pose_feature.ndim == 3
|
||||
|
||||
if self.query_condition and self.key_value_condition:
|
||||
assert encoder_hidden_states is None
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
|
||||
assert encoder_hidden_states.ndim == 3
|
||||
|
||||
batch_size, ehs_sequence_length, _ = encoder_hidden_states.shape
|
||||
if attention_mask is not None:
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, ehs_sequence_length, batch_size)
|
||||
# scaled_dot_product_attention expects attention_mask shape to be
|
||||
# (batch, heads, source_length, target_length)
|
||||
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)
|
||||
|
||||
if attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
if self.query_condition and self.key_value_condition: # only self attention
|
||||
query_hidden_state = self.qkv_merge(hidden_states + pose_feature) * pose_embedding_scale + hidden_states
|
||||
key_value_hidden_state = query_hidden_state
|
||||
elif self.query_condition:
|
||||
query_hidden_state = self.q_merge(hidden_states + pose_feature) * pose_embedding_scale + hidden_states
|
||||
key_value_hidden_state = encoder_hidden_states
|
||||
else:
|
||||
key_value_hidden_state = self.kv_merge(encoder_hidden_states + pose_feature) * pose_embedding_scale + encoder_hidden_states
|
||||
query_hidden_state = hidden_states
|
||||
|
||||
# original attention
|
||||
query = attn.to_q(query_hidden_state)
|
||||
key = attn.to_k(key_value_hidden_state)
|
||||
value = attn.to_v(key_value_hidden_state)
|
||||
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // attn.heads
|
||||
|
||||
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # [bs, seq_len, nhead, head_dim] -> [bs, nhead, seq_len, head_dim]
|
||||
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False) # [bs, nhead, seq_len, head_dim]
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) # [bs, seq_len, 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 attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class PoseAdaptorXFormersAttnProcessor(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size, # dimension of hidden state
|
||||
pose_feature_dim=None, # dimension of the pose feature
|
||||
cross_attention_dim=None, # dimension of the text embedding
|
||||
query_condition=False,
|
||||
key_value_condition=False,
|
||||
scale=1.0,
|
||||
attention_op: Optional[Callable] = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.pose_feature_dim = pose_feature_dim
|
||||
self.cross_attention_dim = cross_attention_dim
|
||||
self.scale = scale
|
||||
self.query_condition = query_condition
|
||||
self.key_value_condition = key_value_condition
|
||||
self.attention_op = attention_op
|
||||
assert hidden_size == pose_feature_dim
|
||||
if self.query_condition and self.key_value_condition:
|
||||
self.qkv_merge = nn.Linear(hidden_size, hidden_size)
|
||||
init.zeros_(self.qkv_merge.weight)
|
||||
init.zeros_(self.qkv_merge.bias)
|
||||
elif self.query_condition:
|
||||
self.q_merge = nn.Linear(hidden_size, hidden_size)
|
||||
init.zeros_(self.q_merge.weight)
|
||||
init.zeros_(self.q_merge.bias)
|
||||
else:
|
||||
self.kv_merge = nn.Linear(hidden_size, hidden_size)
|
||||
init.zeros_(self.kv_merge.weight)
|
||||
init.zeros_(self.kv_merge.bias)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
pose_feature,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
scale=None,
|
||||
):
|
||||
assert pose_feature is not None
|
||||
pose_embedding_scale = (scale or self.scale)
|
||||
|
||||
residual = hidden_states
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
assert hidden_states.ndim == 3 and pose_feature.ndim == 3
|
||||
|
||||
if self.query_condition and self.key_value_condition:
|
||||
assert encoder_hidden_states is None
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
|
||||
assert encoder_hidden_states.ndim == 3
|
||||
|
||||
batch_size, ehs_sequence_length, _ = encoder_hidden_states.shape
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, ehs_sequence_length, batch_size)
|
||||
if attention_mask is not None:
|
||||
# expand our mask's singleton query_tokens dimension:
|
||||
# [batch*heads, 1, key_tokens] ->
|
||||
# [batch*heads, query_tokens, key_tokens]
|
||||
# so that it can be added as a bias onto the attention scores that xformers computes:
|
||||
# [batch*heads, query_tokens, key_tokens]
|
||||
# we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
|
||||
_, query_tokens, _ = hidden_states.shape
|
||||
attention_mask = attention_mask.expand(-1, query_tokens, -1)
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
if attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
if self.query_condition and self.key_value_condition: # only self attention
|
||||
query_hidden_state = self.qkv_merge(hidden_states + pose_feature) * pose_embedding_scale + hidden_states
|
||||
key_value_hidden_state = query_hidden_state
|
||||
elif self.query_condition:
|
||||
query_hidden_state = self.q_merge(hidden_states + pose_feature) * pose_embedding_scale + hidden_states
|
||||
key_value_hidden_state = encoder_hidden_states
|
||||
else:
|
||||
key_value_hidden_state = self.kv_merge(encoder_hidden_states + pose_feature) * pose_embedding_scale + encoder_hidden_states
|
||||
query_hidden_state = hidden_states
|
||||
|
||||
# original attention
|
||||
query = attn.to_q(query_hidden_state)
|
||||
key = attn.to_k(key_value_hidden_state)
|
||||
value = attn.to_v(key_value_hidden_state)
|
||||
|
||||
query = attn.head_to_batch_dim(query).contiguous()
|
||||
key = attn.head_to_batch_dim(key).contiguous()
|
||||
value = attn.head_to_batch_dim(value).contiguous()
|
||||
|
||||
hidden_states = xformers.ops.memory_efficient_attention(
|
||||
query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale
|
||||
)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
@@ -0,0 +1,400 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Optional
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from diffusers.utils import BaseOutput
|
||||
from diffusers.models.attention_processor import Attention
|
||||
from diffusers.models.attention import FeedForward
|
||||
|
||||
from typing import Dict, Any
|
||||
# from cameractrl.models.attention_processor import PoseAdaptorAttnProcessor
|
||||
from .attention_processor import PoseAdaptorAttnProcessor
|
||||
|
||||
from einops import rearrange
|
||||
import math
|
||||
|
||||
|
||||
class InflatedGroupNorm(nn.GroupNorm):
|
||||
def forward(self, x):
|
||||
# return super().forward(x)
|
||||
|
||||
video_length = x.shape[2]
|
||||
|
||||
x = rearrange(x, "b c f h w -> (b f) c h w")
|
||||
x = super().forward(x)
|
||||
x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length)
|
||||
|
||||
return x
|
||||
|
||||
def zero_module(module):
|
||||
# Zero out the parameters of a module and return it.
|
||||
for p in module.parameters():
|
||||
p.detach().zero_()
|
||||
return module
|
||||
|
||||
|
||||
@dataclass
|
||||
class TemporalTransformer3DModelOutput(BaseOutput):
|
||||
sample: torch.FloatTensor
|
||||
|
||||
|
||||
def get_motion_module(
|
||||
in_channels,
|
||||
motion_module_type: str,
|
||||
motion_module_kwargs: dict
|
||||
):
|
||||
if motion_module_type == "Vanilla":
|
||||
return VanillaTemporalModule(in_channels=in_channels, **motion_module_kwargs)
|
||||
else:
|
||||
raise ValueError
|
||||
|
||||
|
||||
class VanillaTemporalModule(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
num_attention_heads=8,
|
||||
num_transformer_block=2,
|
||||
attention_block_types=("Temporal_Self",),
|
||||
temporal_position_encoding=True,
|
||||
temporal_position_encoding_max_len=32,
|
||||
temporal_attention_dim_div=1,
|
||||
cross_attention_dim=320,
|
||||
zero_initialize=True,
|
||||
encoder_hidden_states_query=(False, False),
|
||||
attention_activation_scale=1.0,
|
||||
attention_processor_kwargs: Dict = {},
|
||||
causal_temporal_attention=False,
|
||||
causal_temporal_attention_mask_type="",
|
||||
rescale_output_factor=1.0
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.temporal_transformer = TemporalTransformer3DModel(
|
||||
in_channels=in_channels,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=in_channels // num_attention_heads // temporal_attention_dim_div,
|
||||
num_layers=num_transformer_block,
|
||||
attention_block_types=attention_block_types,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
temporal_position_encoding=temporal_position_encoding,
|
||||
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
encoder_hidden_states_query=encoder_hidden_states_query,
|
||||
attention_activation_scale=attention_activation_scale,
|
||||
attention_processor_kwargs=attention_processor_kwargs,
|
||||
causal_temporal_attention=causal_temporal_attention,
|
||||
causal_temporal_attention_mask_type=causal_temporal_attention_mask_type,
|
||||
rescale_output_factor=rescale_output_factor
|
||||
)
|
||||
|
||||
if zero_initialize:
|
||||
self.temporal_transformer.proj_out = zero_module(self.temporal_transformer.proj_out)
|
||||
|
||||
def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None,
|
||||
cross_attention_kwargs: Dict[str, Any] = {}):
|
||||
hidden_states = self.temporal_transformer(hidden_states, encoder_hidden_states, attention_mask, cross_attention_kwargs=cross_attention_kwargs)
|
||||
|
||||
output = hidden_states
|
||||
return output
|
||||
|
||||
|
||||
class TemporalTransformer3DModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
num_layers,
|
||||
attention_block_types=("Temporal_Self", "Temporal_Self",),
|
||||
dropout=0.0,
|
||||
norm_num_groups=32,
|
||||
cross_attention_dim=320,
|
||||
activation_fn="geglu",
|
||||
attention_bias=False,
|
||||
upcast_attention=False,
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=32,
|
||||
encoder_hidden_states_query=(False, False),
|
||||
attention_activation_scale=1.0,
|
||||
attention_processor_kwargs: Dict = {},
|
||||
|
||||
causal_temporal_attention=None,
|
||||
causal_temporal_attention_mask_type="",
|
||||
rescale_output_factor=1.0
|
||||
):
|
||||
super().__init__()
|
||||
assert causal_temporal_attention is not None
|
||||
self.causal_temporal_attention = causal_temporal_attention
|
||||
|
||||
assert (not causal_temporal_attention) or (causal_temporal_attention_mask_type != "")
|
||||
self.causal_temporal_attention_mask_type = causal_temporal_attention_mask_type
|
||||
self.causal_temporal_attention_mask = None
|
||||
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
|
||||
self.norm = InflatedGroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
self.proj_in = nn.Linear(in_channels, inner_dim)
|
||||
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
TemporalTransformerBlock(
|
||||
dim=inner_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=attention_head_dim,
|
||||
attention_block_types=attention_block_types,
|
||||
dropout=dropout,
|
||||
norm_num_groups=norm_num_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
activation_fn=activation_fn,
|
||||
attention_bias=attention_bias,
|
||||
upcast_attention=upcast_attention,
|
||||
temporal_position_encoding=temporal_position_encoding,
|
||||
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
encoder_hidden_states_query=encoder_hidden_states_query,
|
||||
attention_activation_scale=attention_activation_scale,
|
||||
attention_processor_kwargs=attention_processor_kwargs,
|
||||
rescale_output_factor=rescale_output_factor,
|
||||
)
|
||||
for d in range(num_layers)
|
||||
]
|
||||
)
|
||||
self.proj_out = nn.Linear(inner_dim, in_channels)
|
||||
|
||||
def get_causal_temporal_attention_mask(self, hidden_states):
|
||||
batch_size, sequence_length, dim = hidden_states.shape
|
||||
|
||||
if self.causal_temporal_attention_mask is None or self.causal_temporal_attention_mask.shape != (
|
||||
batch_size, sequence_length, sequence_length):
|
||||
if self.causal_temporal_attention_mask_type == "causal":
|
||||
# 1. vanilla causal mask
|
||||
mask = torch.tril(torch.ones(sequence_length, sequence_length))
|
||||
|
||||
elif self.causal_temporal_attention_mask_type == "2-seq":
|
||||
# 2. 2-seq
|
||||
mask = torch.zeros(sequence_length, sequence_length)
|
||||
mask[:sequence_length // 2, :sequence_length // 2] = 1
|
||||
mask[-sequence_length // 2:, -sequence_length // 2:] = 1
|
||||
|
||||
elif self.causal_temporal_attention_mask_type == "0-prev":
|
||||
# attn to the previous frame
|
||||
indices = torch.arange(sequence_length)
|
||||
indices_prev = indices - 1
|
||||
indices_prev[0] = 0
|
||||
mask = torch.zeros(sequence_length, sequence_length)
|
||||
mask[:, 0] = 1.
|
||||
mask[indices, indices_prev] = 1.
|
||||
|
||||
elif self.causal_temporal_attention_mask_type == "0":
|
||||
# only attn to first frame
|
||||
mask = torch.zeros(sequence_length, sequence_length)
|
||||
mask[:, 0] = 1
|
||||
|
||||
elif self.causal_temporal_attention_mask_type == "wo-self":
|
||||
indices = torch.arange(sequence_length)
|
||||
mask = torch.ones(sequence_length, sequence_length)
|
||||
mask[indices, indices] = 0
|
||||
|
||||
elif self.causal_temporal_attention_mask_type == "circle":
|
||||
indices = torch.arange(sequence_length)
|
||||
indices_prev = indices - 1
|
||||
indices_prev[0] = 0
|
||||
|
||||
mask = torch.eye(sequence_length)
|
||||
mask[indices, indices_prev] = 1
|
||||
mask[0, -1] = 1
|
||||
|
||||
else:
|
||||
raise ValueError
|
||||
|
||||
# generate attention mask fron binary values
|
||||
mask = mask.masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
|
||||
mask = mask.unsqueeze(0)
|
||||
mask = mask.repeat(batch_size, 1, 1)
|
||||
|
||||
self.causal_temporal_attention_mask = mask.to(hidden_states.device)
|
||||
|
||||
return self.causal_temporal_attention_mask
|
||||
|
||||
def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None,
|
||||
cross_attention_kwargs: Dict[str, Any] = {},):
|
||||
residual = hidden_states
|
||||
|
||||
assert hidden_states.dim() == 5, f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}."
|
||||
height, width = hidden_states.shape[-2:]
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
hidden_states = rearrange(hidden_states, "b c f h w -> (b h w) f c")
|
||||
hidden_states = self.proj_in(hidden_states)
|
||||
|
||||
attention_mask = self.get_causal_temporal_attention_mask(
|
||||
hidden_states) if self.causal_temporal_attention else attention_mask
|
||||
|
||||
# Transformer Blocks
|
||||
for block in self.transformer_blocks:
|
||||
hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=attention_mask, cross_attention_kwargs=cross_attention_kwargs)
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
|
||||
hidden_states = rearrange(hidden_states, "(b h w) f c -> b c f h w", h=height, w=width)
|
||||
|
||||
output = hidden_states + residual
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class TemporalTransformerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
attention_block_types=("Temporal_Self", "Temporal_Self",),
|
||||
dropout=0.0,
|
||||
norm_num_groups=32,
|
||||
cross_attention_dim=768,
|
||||
activation_fn="geglu",
|
||||
attention_bias=False,
|
||||
upcast_attention=False,
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=32,
|
||||
encoder_hidden_states_query=(False, False),
|
||||
attention_activation_scale=1.0,
|
||||
attention_processor_kwargs: Dict = {},
|
||||
rescale_output_factor=1.0
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
attention_blocks = []
|
||||
norms = []
|
||||
self.attention_block_types = attention_block_types
|
||||
|
||||
for block_idx, block_name in enumerate(attention_block_types):
|
||||
attention_blocks.append(
|
||||
TemporalSelfAttention(
|
||||
attention_mode=block_name,
|
||||
cross_attention_dim=cross_attention_dim if block_name in ['Temporal_Cross', 'Temporal_Pose_Adaptor'] else None,
|
||||
query_dim=dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
upcast_attention=upcast_attention,
|
||||
temporal_position_encoding=temporal_position_encoding,
|
||||
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
rescale_output_factor=rescale_output_factor,
|
||||
)
|
||||
)
|
||||
norms.append(nn.LayerNorm(dim))
|
||||
|
||||
self.attention_blocks = nn.ModuleList(attention_blocks)
|
||||
self.norms = nn.ModuleList(norms)
|
||||
|
||||
self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
|
||||
self.ff_norm = nn.LayerNorm(dim)
|
||||
|
||||
def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, cross_attention_kwargs: Dict[str, Any] = {}):
|
||||
for attention_block, norm, attention_block_type in zip(self.attention_blocks, self.norms, self.attention_block_types):
|
||||
norm_hidden_states = norm(hidden_states)
|
||||
hidden_states = attention_block(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=norm_hidden_states if attention_block_type == 'Temporal_Self' else encoder_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
**cross_attention_kwargs
|
||||
) + hidden_states
|
||||
|
||||
hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states
|
||||
|
||||
output = hidden_states
|
||||
return output
|
||||
|
||||
|
||||
class PositionalEncoding(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model,
|
||||
dropout=0.,
|
||||
max_len=32,
|
||||
):
|
||||
super().__init__()
|
||||
self.dropout = nn.Dropout(p=dropout)
|
||||
position = torch.arange(max_len).unsqueeze(1)
|
||||
div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
|
||||
pe = torch.zeros(1, max_len, d_model)
|
||||
pe[0, :, 0::2] = torch.sin(position * div_term)
|
||||
pe[0, :, 1::2] = torch.cos(position * div_term)
|
||||
self.register_buffer('pe', pe)
|
||||
|
||||
def forward(self, x):
|
||||
x = x + self.pe[:, :x.size(1)]
|
||||
return self.dropout(x)
|
||||
|
||||
|
||||
class TemporalSelfAttention(Attention):
|
||||
def __init__(
|
||||
self,
|
||||
attention_mode=None,
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=32,
|
||||
rescale_output_factor=1.0,
|
||||
*args, **kwargs
|
||||
):
|
||||
super().__init__(*args, **kwargs)
|
||||
assert attention_mode == "Temporal_Self"
|
||||
|
||||
self.pos_encoder = PositionalEncoding(
|
||||
kwargs["query_dim"],
|
||||
max_len=temporal_position_encoding_max_len
|
||||
) if temporal_position_encoding else None
|
||||
self.rescale_output_factor = rescale_output_factor
|
||||
|
||||
def set_use_memory_efficient_attention_xformers(
|
||||
self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None
|
||||
):
|
||||
# disable motion module efficient xformers to avoid bad results, don't know why
|
||||
# TODO: fix this bug
|
||||
pass
|
||||
|
||||
def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, **cross_attention_kwargs):
|
||||
# The `Attention` class can call different attention processors / attention functions
|
||||
# here we simply pass along all tensors to the selected processor class
|
||||
# For standard processors that are defined here, `**cross_attention_kwargs` is empty
|
||||
|
||||
# add position encoding
|
||||
if self.pos_encoder is not None:
|
||||
hidden_states = self.pos_encoder(hidden_states)
|
||||
if "pose_feature" in cross_attention_kwargs:
|
||||
pose_feature = cross_attention_kwargs["pose_feature"]
|
||||
if pose_feature.ndim == 5:
|
||||
pose_feature = rearrange(pose_feature, "b c f h w -> (b h w) f c")
|
||||
else:
|
||||
assert pose_feature.ndim == 3
|
||||
cross_attention_kwargs["pose_feature"] = pose_feature
|
||||
|
||||
if isinstance(self.processor, PoseAdaptorAttnProcessor):
|
||||
return self.processor(
|
||||
self,
|
||||
hidden_states,
|
||||
cross_attention_kwargs.pop('pose_feature'),
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=attention_mask,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
elif hasattr(self.processor, "__call__"):
|
||||
return self.processor.__call__(
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=attention_mask,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
else:
|
||||
return self.processor(
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=attention_mask,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,244 @@
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from typing import List, Tuple
|
||||
# from cameractrl.models.motion_module import TemporalTransformerBlock
|
||||
from .motion_module import TemporalTransformerBlock
|
||||
|
||||
|
||||
def get_parameter_dtype(parameter: torch.nn.Module):
|
||||
try:
|
||||
params = tuple(parameter.parameters())
|
||||
if len(params) > 0:
|
||||
return params[0].dtype
|
||||
|
||||
buffers = tuple(parameter.buffers())
|
||||
if len(buffers) > 0:
|
||||
return buffers[0].dtype
|
||||
|
||||
except StopIteration:
|
||||
# For torch.nn.DataParallel compatibility in PyTorch 1.5
|
||||
|
||||
def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, torch.Tensor]]:
|
||||
tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
|
||||
return tuples
|
||||
|
||||
gen = parameter._named_members(get_members_fn=find_tensor_attributes)
|
||||
first_tuple = next(gen)
|
||||
return first_tuple[1].dtype
|
||||
|
||||
|
||||
def conv_nd(dims, *args, **kwargs):
|
||||
"""
|
||||
Create a 1D, 2D, or 3D convolution module.
|
||||
"""
|
||||
if dims == 1:
|
||||
return nn.Conv1d(*args, **kwargs)
|
||||
elif dims == 2:
|
||||
return nn.Conv2d(*args, **kwargs)
|
||||
elif dims == 3:
|
||||
return nn.Conv3d(*args, **kwargs)
|
||||
raise ValueError(f"unsupported dimensions: {dims}")
|
||||
|
||||
|
||||
def avg_pool_nd(dims, *args, **kwargs):
|
||||
"""
|
||||
Create a 1D, 2D, or 3D average pooling module.
|
||||
"""
|
||||
if dims == 1:
|
||||
return nn.AvgPool1d(*args, **kwargs)
|
||||
elif dims == 2:
|
||||
return nn.AvgPool2d(*args, **kwargs)
|
||||
elif dims == 3:
|
||||
return nn.AvgPool3d(*args, **kwargs)
|
||||
raise ValueError(f"unsupported dimensions: {dims}")
|
||||
|
||||
|
||||
class PoseAdaptor(nn.Module):
|
||||
def __init__(self, unet, pose_encoder):
|
||||
super().__init__()
|
||||
self.unet = unet
|
||||
self.pose_encoder = pose_encoder
|
||||
|
||||
def forward(self, inp_noisy_latents, timesteps, encoder_hidden_states, added_time_ids, pose_embedding):
|
||||
assert pose_embedding.ndim == 5
|
||||
pose_embedding_features = self.pose_encoder(pose_embedding) # b c f h w
|
||||
noise_pred = self.unet(
|
||||
inp_noisy_latents,
|
||||
timesteps,
|
||||
encoder_hidden_states,
|
||||
added_time_ids=added_time_ids,
|
||||
pose_features=pose_embedding_features,
|
||||
).sample
|
||||
|
||||
return noise_pred
|
||||
|
||||
|
||||
class Downsample(nn.Module):
|
||||
"""
|
||||
A downsampling layer with an optional convolution.
|
||||
:param channels: channels in the inputs and outputs.
|
||||
:param use_conv: a bool determining if a convolution is applied.
|
||||
:param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
|
||||
downsampling occurs in the inner-two dimensions.
|
||||
"""
|
||||
|
||||
def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.out_channels = out_channels or channels
|
||||
self.use_conv = use_conv
|
||||
self.dims = dims
|
||||
stride = 2 if dims != 3 else (1, 2, 2)
|
||||
if use_conv:
|
||||
self.op = conv_nd(dims, self.channels, self.out_channels, 3, stride=stride, padding=padding)
|
||||
else:
|
||||
assert self.channels == self.out_channels
|
||||
self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
|
||||
|
||||
def forward(self, x):
|
||||
assert x.shape[1] == self.channels
|
||||
return self.op(x)
|
||||
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
|
||||
def __init__(self, in_c, out_c, down, ksize=3, sk=False, use_conv=True):
|
||||
super().__init__()
|
||||
ps = ksize // 2
|
||||
if in_c != out_c or sk == False:
|
||||
self.in_conv = nn.Conv2d(in_c, out_c, ksize, 1, ps)
|
||||
else:
|
||||
self.in_conv = None
|
||||
self.block1 = nn.Conv2d(out_c, out_c, 3, 1, 1)
|
||||
self.act = nn.ReLU()
|
||||
self.block2 = nn.Conv2d(out_c, out_c, ksize, 1, ps)
|
||||
if sk == False:
|
||||
self.skep = nn.Conv2d(in_c, out_c, ksize, 1, ps)
|
||||
else:
|
||||
self.skep = None
|
||||
|
||||
self.down = down
|
||||
if self.down == True:
|
||||
self.down_opt = Downsample(in_c, use_conv=use_conv)
|
||||
|
||||
def forward(self, x):
|
||||
if self.down == True:
|
||||
x = self.down_opt(x)
|
||||
if self.in_conv is not None: # edit
|
||||
x = self.in_conv(x)
|
||||
|
||||
h = self.block1(x)
|
||||
h = self.act(h)
|
||||
h = self.block2(h)
|
||||
if self.skep is not None:
|
||||
return h + self.skep(x)
|
||||
else:
|
||||
return h + x
|
||||
|
||||
|
||||
class PositionalEncoding(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model,
|
||||
dropout=0.,
|
||||
max_len=32,
|
||||
):
|
||||
super().__init__()
|
||||
self.dropout = nn.Dropout(p=dropout)
|
||||
position = torch.arange(max_len).unsqueeze(1)
|
||||
div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
|
||||
pe = torch.zeros(1, max_len, d_model)
|
||||
pe[0, :, 0::2, ...] = torch.sin(position * div_term)
|
||||
pe[0, :, 1::2, ...] = torch.cos(position * div_term)
|
||||
pe.unsqueeze_(-1).unsqueeze_(-1)
|
||||
self.register_buffer('pe', pe)
|
||||
|
||||
def forward(self, x):
|
||||
x = x + self.pe[:, :x.size(1), ...]
|
||||
return self.dropout(x)
|
||||
|
||||
|
||||
class CameraPoseEncoder(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
downscale_factor,
|
||||
channels=[320, 640, 1280, 1280],
|
||||
nums_rb=3,
|
||||
cin=64,
|
||||
ksize=3,
|
||||
sk=False,
|
||||
use_conv=True,
|
||||
compression_factor=1,
|
||||
temporal_attention_nhead=8,
|
||||
attention_block_types=("Temporal_Self", ),
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=16,
|
||||
rescale_output_factor=1.0):
|
||||
super(CameraPoseEncoder, self).__init__()
|
||||
self.unshuffle = nn.PixelUnshuffle(downscale_factor)
|
||||
self.channels = channels
|
||||
self.nums_rb = nums_rb
|
||||
self.encoder_down_conv_blocks = nn.ModuleList()
|
||||
self.encoder_down_attention_blocks = nn.ModuleList()
|
||||
for i in range(len(channels)):
|
||||
conv_layers = nn.ModuleList()
|
||||
temporal_attention_layers = nn.ModuleList()
|
||||
for j in range(nums_rb):
|
||||
if j == 0 and i != 0:
|
||||
in_dim = channels[i - 1]
|
||||
out_dim = int(channels[i] / compression_factor)
|
||||
conv_layer = ResnetBlock(in_dim, out_dim, down=True, ksize=ksize, sk=sk, use_conv=use_conv)
|
||||
elif j == 0:
|
||||
in_dim = channels[0]
|
||||
out_dim = int(channels[i] / compression_factor)
|
||||
conv_layer = ResnetBlock(in_dim, out_dim, down=False, ksize=ksize, sk=sk, use_conv=use_conv)
|
||||
elif j == nums_rb - 1:
|
||||
in_dim = channels[i] / compression_factor
|
||||
out_dim = channels[i]
|
||||
conv_layer = ResnetBlock(in_dim, out_dim, down=False, ksize=ksize, sk=sk, use_conv=use_conv)
|
||||
else:
|
||||
in_dim = int(channels[i] / compression_factor)
|
||||
out_dim = int(channels[i] / compression_factor)
|
||||
conv_layer = ResnetBlock(in_dim, out_dim, down=False, ksize=ksize, sk=sk, use_conv=use_conv)
|
||||
temporal_attention_layer = TemporalTransformerBlock(dim=out_dim,
|
||||
num_attention_heads=temporal_attention_nhead,
|
||||
attention_head_dim=int(out_dim / temporal_attention_nhead),
|
||||
attention_block_types=attention_block_types,
|
||||
dropout=0.0,
|
||||
cross_attention_dim=None,
|
||||
temporal_position_encoding=temporal_position_encoding,
|
||||
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
rescale_output_factor=rescale_output_factor)
|
||||
conv_layers.append(conv_layer)
|
||||
temporal_attention_layers.append(temporal_attention_layer)
|
||||
self.encoder_down_conv_blocks.append(conv_layers)
|
||||
self.encoder_down_attention_blocks.append(temporal_attention_layers)
|
||||
|
||||
self.encoder_conv_in = nn.Conv2d(cin, channels[0], 3, 1, 1)
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
"""
|
||||
`torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
|
||||
"""
|
||||
return get_parameter_dtype(self)
|
||||
|
||||
def forward(self, x):
|
||||
# unshuffle
|
||||
bs = x.shape[0]
|
||||
x = rearrange(x, "b f c h w -> (b f) c h w")
|
||||
x = self.unshuffle(x)
|
||||
# extract features
|
||||
features = []
|
||||
x = self.encoder_conv_in(x)
|
||||
for res_block, attention_block in zip(self.encoder_down_conv_blocks, self.encoder_down_attention_blocks):
|
||||
for res_layer, attention_layer in zip(res_block, attention_block):
|
||||
x = res_layer(x)
|
||||
h, w = x.shape[-2:]
|
||||
x = rearrange(x, '(b f) c h w -> (b h w) f c', b=bs)
|
||||
x = attention_layer(x)
|
||||
x = rearrange(x, '(b h w) f c -> (b f) c h w', h=h, w=w)
|
||||
features.append(rearrange(x, '(b f) c h w -> b c f h w', b=bs))
|
||||
return features
|
||||
Reference in New Issue
Block a user