init
This commit is contained in:
songkey-win
@@ -0,0 +1,14 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
@File : __init__.py
|
||||
@Author : Songkey
|
||||
@Email : songkey@pku.edu.cn
|
||||
@Date : 8/14/2024
|
||||
@Desc :
|
||||
"""
|
||||
|
||||
from .hm_denoising_motion import HMDenoisingMotion
|
||||
from .hm_control import HMControlNet
|
||||
from .hm_blocks import HMReferenceAdapter
|
||||
from .hm_denoising_3d import HMDenoising3D
|
||||
@@ -0,0 +1,330 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
@File : models6/hm_blocks.py
|
||||
@Author : Songkey
|
||||
@Email : songkey@pku.edu.cn
|
||||
@Date : 8/14/2024
|
||||
@Desc :
|
||||
"""
|
||||
|
||||
import copy
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.checkpoint
|
||||
import torch.nn.functional as F
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
from einops import rearrange
|
||||
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.attention import Attention, FeedForward
|
||||
from torchvision.models.resnet import BasicBlock
|
||||
from diffusers.models.embeddings import SinusoidalPositionalEmbedding
|
||||
from diffusers.utils import is_transformers_available, logging
|
||||
from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
|
||||
|
||||
from diffusers.configuration_utils import FrozenDict
|
||||
from diffusers.models.attention_processor import (
|
||||
AttnProcessor,
|
||||
AttnProcessor2_0,
|
||||
)
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
class HMReferenceAdapter(ModelMixin, ConfigMixin):
|
||||
@register_to_config
|
||||
def __init__(self,
|
||||
block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
|
||||
down_block_types: Tuple[str] = (
|
||||
"CrossAttnDownBlock2D",
|
||||
"CrossAttnDownBlock2D",
|
||||
"CrossAttnDownBlock2D",
|
||||
"DownBlock2D",
|
||||
),
|
||||
up_block_types: Tuple[str] = (
|
||||
"UpBlock2D",
|
||||
"CrossAttnUpBlock2D",
|
||||
"CrossAttnUpBlock2D",
|
||||
"CrossAttnUpBlock2D"),
|
||||
num_attention_heads: Optional[Union[int, Tuple[int]]] = 8,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.reference_modules_down = nn.ModuleList([])
|
||||
|
||||
if isinstance(num_attention_heads, int):
|
||||
num_attention_heads = (num_attention_heads,) * len(down_block_types)
|
||||
|
||||
for i, down_block_type in enumerate(down_block_types):
|
||||
output_channel = block_out_channels[i]
|
||||
|
||||
self.reference_modules_down.append(
|
||||
SKReferenceAttention(
|
||||
in_channels=output_channel,
|
||||
num_attention_heads=num_attention_heads[i],
|
||||
)
|
||||
)
|
||||
|
||||
self.reference_modules_mid = SKReferenceAttention(
|
||||
in_channels=block_out_channels[-1],
|
||||
num_attention_heads=num_attention_heads[-1],
|
||||
)
|
||||
|
||||
self.reference_modules_up = nn.ModuleList([])
|
||||
|
||||
reversed_block_out_channels = list(reversed(block_out_channels))
|
||||
reversed_num_attention_heads = list(reversed(num_attention_heads))
|
||||
|
||||
output_channel = reversed_block_out_channels[0]
|
||||
for i, up_block_type in enumerate(up_block_types):
|
||||
prev_output_channel = output_channel
|
||||
output_channel = reversed_block_out_channels[i]
|
||||
|
||||
if i > 0:
|
||||
self.reference_modules_up.append(
|
||||
SKReferenceAttention(
|
||||
in_channels=prev_output_channel,
|
||||
num_attention_heads=reversed_num_attention_heads[i],
|
||||
num_positional_embeddings=64 * 2
|
||||
)
|
||||
)
|
||||
|
||||
class SKReferenceAttention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
num_attention_heads: int = 1,
|
||||
norm_elementwise_affine: bool = True,
|
||||
norm_eps: float = 1e-5,
|
||||
num_positional_embeddings: int = 64*2,
|
||||
):
|
||||
super().__init__()
|
||||
self.pos_embed = SinusoidalPositionalEmbedding(in_channels, max_seq_length=num_positional_embeddings)
|
||||
self.attn1 = Attention(
|
||||
query_dim=in_channels,
|
||||
heads=num_attention_heads,
|
||||
dim_head=in_channels // num_attention_heads,
|
||||
dropout=0.0,
|
||||
bias=False,
|
||||
cross_attention_dim=None,
|
||||
upcast_attention=False,
|
||||
out_bias=True,
|
||||
)
|
||||
self.attn2 = Attention(
|
||||
query_dim=in_channels,
|
||||
heads=num_attention_heads,
|
||||
dim_head=in_channels // num_attention_heads,
|
||||
dropout=0.0,
|
||||
bias=False,
|
||||
cross_attention_dim=None,
|
||||
upcast_attention=False,
|
||||
out_bias=True,
|
||||
)
|
||||
# self.ff = FeedForward(
|
||||
# in_channels,
|
||||
# dropout=0.0,
|
||||
# activation_fn="geglu",
|
||||
# final_dropout=False,
|
||||
# inner_dim=in_channels*2,
|
||||
# bias=True,
|
||||
# )
|
||||
self.norm = nn.LayerNorm(in_channels, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
|
||||
self.proj = zero_module(nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1))
|
||||
|
||||
def forward(self, hidden_states, ref_stats, num_frames):
|
||||
h, w = hidden_states.shape[-2:]
|
||||
ref_stats = ref_stats.repeat_interleave(num_frames, dim=0)
|
||||
cat_stats = torch.cat([hidden_states, ref_stats], dim=-1)
|
||||
|
||||
cat_stats = rearrange(cat_stats.contiguous(), "b c h w -> (b h) w c")
|
||||
res1 = self.attn1(self.norm(self.pos_embed(cat_stats)))
|
||||
res1 = rearrange(res1[:, :w, :], "(b h) w c -> b c h w", h=h)
|
||||
|
||||
cat_stats2 = torch.cat([res1, ref_stats], dim=-2)
|
||||
cat_stats2 = rearrange(cat_stats2.contiguous(), "b c h w -> (b w) h c")
|
||||
res2 = self.attn2(self.norm(self.pos_embed(cat_stats2)))
|
||||
|
||||
res2 = rearrange(res2[:, :h, :], "(b w) h c -> b c h w", w=w)
|
||||
# res2 = self.ff(res2)
|
||||
|
||||
return hidden_states + self.proj(res2)
|
||||
|
||||
class SKCrossAttention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channel_in,
|
||||
channel_out,
|
||||
heads: int=8,
|
||||
cross_attention_dim: int = 320,
|
||||
norm_elementwise_affine: bool = True,
|
||||
norm_eps: float = 1e-5,
|
||||
num_positional_embeddings: int = 64,
|
||||
num_positional_embeddings_hidden: int = 64,
|
||||
):
|
||||
super().__init__()
|
||||
self.conv = BasicBlock(
|
||||
inplanes=channel_in,
|
||||
planes=channel_out,
|
||||
stride=2,
|
||||
downsample=nn.Sequential(
|
||||
nn.Conv2d(channel_in, channel_out, kernel_size=1, stride=2, bias=False),
|
||||
nn.InstanceNorm2d(channel_out),
|
||||
nn.SiLU(),
|
||||
),
|
||||
norm_layer=nn.InstanceNorm2d,
|
||||
)
|
||||
|
||||
self.pos_embed = SinusoidalPositionalEmbedding(channel_out, max_seq_length=num_positional_embeddings)
|
||||
self.pos_embed_hidden = SinusoidalPositionalEmbedding(cross_attention_dim, max_seq_length=num_positional_embeddings_hidden)
|
||||
|
||||
self.norm1 = nn.LayerNorm(channel_out, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
|
||||
self.attn1 = Attention(
|
||||
query_dim=channel_out,
|
||||
heads=heads,
|
||||
dim_head=channel_out // heads,
|
||||
dropout=0.0,
|
||||
bias=False,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
upcast_attention=False,
|
||||
out_bias=True,
|
||||
)
|
||||
|
||||
self.norm2 = nn.LayerNorm(channel_out, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
|
||||
self.attn2 = Attention(
|
||||
query_dim=channel_out,
|
||||
heads=heads,
|
||||
dim_head=channel_out // heads,
|
||||
dropout=0.0,
|
||||
bias=False,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
upcast_attention=False,
|
||||
out_bias=True,
|
||||
)
|
||||
|
||||
self.ff = FeedForward(
|
||||
channel_out,
|
||||
dropout=0.0,
|
||||
activation_fn="geglu",
|
||||
final_dropout=False,
|
||||
inner_dim=channel_out*2,
|
||||
bias=True,
|
||||
)
|
||||
|
||||
self.proj = zero_module(nn.Conv2d(channel_out, channel_out, kernel_size=3, padding=1))
|
||||
|
||||
def forward(self, input, hidden_stats):
|
||||
x = self.conv(input)
|
||||
h, w = x.shape[-2:]
|
||||
x = rearrange(x, "b c h w -> (b h) w c")
|
||||
x = self.attn1(self.norm1(self.pos_embed(x)), self.pos_embed_hidden(hidden_stats.repeat_interleave(h, dim=0).contiguous()))
|
||||
x = rearrange(x, "(b h) w c -> (b w) h c", h=h)
|
||||
x = self.ff(self.attn2(self.norm2(self.pos_embed(x)), self.pos_embed_hidden(hidden_stats.repeat_interleave(w, dim=0).contiguous())))
|
||||
x = rearrange(x, "(b w) h c -> b c h w", w=w)
|
||||
x = self.proj(x)
|
||||
return x
|
||||
|
||||
# https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/controlnet.py
|
||||
def zero_module(module):
|
||||
for p in module.parameters():
|
||||
nn.init.zeros_(p)
|
||||
return module
|
||||
|
||||
# https://github.com/huggingface/diffusers/blob/82058a5413ca09561cc5cc236c4abc5eeda7b209/src/diffusers/loaders/ip_adapter.py
|
||||
|
||||
if is_transformers_available():
|
||||
from diffusers.models.attention_processor import (
|
||||
IPAdapterAttnProcessor,
|
||||
IPAdapterAttnProcessor2_0,
|
||||
)
|
||||
|
||||
class CopyWeights(object):
|
||||
@classmethod
|
||||
def from_unet2d(cls, unet: UNet2DConditionModel):
|
||||
# adapted from :https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unets/unet_motion_model.py
|
||||
|
||||
config = dict(unet.config)
|
||||
|
||||
# Need this for backwards compatibility with UNet2DConditionModel checkpoints
|
||||
if not config.get("num_attention_heads"):
|
||||
config["num_attention_heads"] = config["attention_head_dim"]
|
||||
|
||||
config = FrozenDict(config)
|
||||
model = cls.from_config(config)
|
||||
|
||||
model.conv_in.load_state_dict(unet.conv_in.state_dict())
|
||||
|
||||
model.time_proj.load_state_dict(unet.time_proj.state_dict())
|
||||
model.time_embedding.load_state_dict(unet.time_embedding.state_dict())
|
||||
|
||||
if any(
|
||||
isinstance(proc, (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0))
|
||||
for proc in unet.attn_processors.values()
|
||||
):
|
||||
attn_procs = {}
|
||||
for name, processor in unet.attn_processors.items():
|
||||
if name.endswith("attn1.processor"):
|
||||
attn_processor_class = (
|
||||
AttnProcessor2_0 if hasattr(F, "scaled_dot_product_attention") else AttnProcessor
|
||||
)
|
||||
attn_procs[name] = attn_processor_class()
|
||||
else:
|
||||
attn_processor_class = (
|
||||
IPAdapterAttnProcessor2_0
|
||||
if hasattr(F, "scaled_dot_product_attention")
|
||||
else IPAdapterAttnProcessor
|
||||
)
|
||||
attn_procs[name] = attn_processor_class(
|
||||
hidden_size=processor.hidden_size,
|
||||
cross_attention_dim=processor.cross_attention_dim,
|
||||
scale=processor.scale,
|
||||
num_tokens=processor.num_tokens,
|
||||
)
|
||||
for name, processor in model.attn_processors.items():
|
||||
if name not in attn_procs:
|
||||
attn_procs[name] = processor.__class__()
|
||||
model.set_attn_processor(attn_procs)
|
||||
model.config.encoder_hid_dim_type = "ip_image_proj"
|
||||
model.encoder_hid_proj = unet.encoder_hid_proj
|
||||
|
||||
for i, down_block in enumerate(unet.down_blocks):
|
||||
model.down_blocks[i].resnets.load_state_dict(down_block.resnets.state_dict())
|
||||
if hasattr(model.down_blocks[i], "attentions"):
|
||||
model.down_blocks[i].attentions.load_state_dict(down_block.attentions.state_dict())
|
||||
if model.down_blocks[i].downsamplers:
|
||||
model.down_blocks[i].downsamplers.load_state_dict(down_block.downsamplers.state_dict())
|
||||
|
||||
for i, up_block in enumerate(unet.up_blocks):
|
||||
model.up_blocks[i].resnets.load_state_dict(up_block.resnets.state_dict())
|
||||
if hasattr(model.up_blocks[i], "attentions"):
|
||||
model.up_blocks[i].attentions.load_state_dict(up_block.attentions.state_dict())
|
||||
if model.up_blocks[i].upsamplers:
|
||||
model.up_blocks[i].upsamplers.load_state_dict(up_block.upsamplers.state_dict())
|
||||
|
||||
model.mid_block.resnets.load_state_dict(unet.mid_block.resnets.state_dict())
|
||||
model.mid_block.attentions.load_state_dict(unet.mid_block.attentions.state_dict())
|
||||
|
||||
if unet.conv_norm_out is not None:
|
||||
model.conv_norm_out.load_state_dict(unet.conv_norm_out.state_dict())
|
||||
if unet.conv_act is not None:
|
||||
model.conv_act.load_state_dict(unet.conv_act.state_dict())
|
||||
model.conv_out.load_state_dict(unet.conv_out.state_dict())
|
||||
|
||||
# ensure that the Motion UNet is the same dtype as the UNet2DConditionModel
|
||||
model.to(unet.dtype)
|
||||
|
||||
return model
|
||||
|
||||
class InsertReferenceAdapter(object):
|
||||
def __init__(self):
|
||||
self.reference_modules_down = None
|
||||
self.reference_modules_mid = None
|
||||
self.reference_modules_up = None
|
||||
|
||||
def insert_reference_adapter(self, adapter: HMReferenceAdapter):
|
||||
self.reference_modules_down = copy.deepcopy(adapter.reference_modules_down)
|
||||
self.reference_modules_mid = copy.deepcopy(adapter.reference_modules_mid)
|
||||
self.reference_modules_up = copy.deepcopy(adapter.reference_modules_up)
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
@File : models6/hm_control.py
|
||||
@Author : Songkey
|
||||
@Email : songkey@pku.edu.cn
|
||||
@Date : 8/14/2024
|
||||
@Desc :
|
||||
"""
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
|
||||
from einops import rearrange
|
||||
from .hm_blocks import SKCrossAttention
|
||||
|
||||
|
||||
class HMControlNet(ModelMixin, ConfigMixin):
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
embedding_channels: int,
|
||||
input_channels: int = 3,
|
||||
scale_factor: int = 8,
|
||||
cross_attention_dim: int = 320,
|
||||
block_out_channels: Tuple[int] = (128, 320, 640, 1280),
|
||||
):
|
||||
super().__init__()
|
||||
self.scale_factor = scale_factor
|
||||
self.embedding_channels = embedding_channels
|
||||
self.cross_attention_dim = cross_attention_dim
|
||||
self.conv_in = nn.Conv2d(input_channels, block_out_channels[0], kernel_size=3, padding=1, bias=False)
|
||||
|
||||
self.exp_embedding = Timesteps(cross_attention_dim, True, 0)
|
||||
self.exp_proj = TimestepEmbedding(cross_attention_dim, cross_attention_dim)
|
||||
self.face_proj = TimestepEmbedding(1024, cross_attention_dim)
|
||||
|
||||
self.blocks_down = nn.ModuleList([])
|
||||
for i in range(1, len(block_out_channels)):
|
||||
channel_in = block_out_channels[i-1]
|
||||
channel_out = block_out_channels[i]
|
||||
self.blocks_down.append(
|
||||
SKCrossAttention(
|
||||
channel_in=channel_in,
|
||||
channel_out=channel_out,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
num_positional_embeddings=64,
|
||||
num_positional_embeddings_hidden=64
|
||||
)
|
||||
)
|
||||
|
||||
def forward(self, condition, drive_coeff, face_parts):
|
||||
bs, _, video_length, h, w = condition.shape
|
||||
|
||||
condition = rearrange(condition, "b c f h w -> (b f) c h w")
|
||||
conditioning = F.interpolate(condition,
|
||||
size=(h // self.scale_factor, w // self.scale_factor),
|
||||
mode='bilinear',
|
||||
align_corners=False)
|
||||
embedding = self.conv_in(conditioning)
|
||||
|
||||
drive_coeff = rearrange(drive_coeff * 500., "b f c -> (b f c)")
|
||||
drive_embedding = self.exp_embedding(drive_coeff).to(dtype=embedding.dtype)
|
||||
drive_embedding = rearrange(drive_embedding, "(b f c) d -> b f c d",
|
||||
b=bs, f=video_length, d=self.cross_attention_dim)
|
||||
|
||||
face_parts = rearrange(face_parts, "b f c d -> (b f) c d")
|
||||
face_embedding = self.face_proj(face_parts)
|
||||
|
||||
face_embedding = rearrange(face_embedding, "(b f) c d -> b f c d", f=video_length)
|
||||
|
||||
drive_embedding = torch.cat([face_embedding, drive_embedding], dim=2)
|
||||
drive_embedding = rearrange(drive_embedding, "b f c d -> (b f) c d")
|
||||
|
||||
drive_embedding = self.exp_proj(drive_embedding)
|
||||
|
||||
ret_dict = {}
|
||||
for idx, block in enumerate(self.blocks_down):
|
||||
embedding = block(embedding, drive_embedding)
|
||||
ret_dict[f'down_{idx}'] = rearrange(embedding,
|
||||
"(b f) c h w -> b c f h w",
|
||||
f=video_length)
|
||||
return ret_dict
|
||||
@@ -0,0 +1,289 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
@File : models6/hm_denoising_3d.py
|
||||
@Author : Songkey
|
||||
@Email : songkey@pku.edu.cn
|
||||
@Date : 8/14/2024
|
||||
@Desc : 删除实验代码，精简结构
|
||||
adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unets/unet_2d_condition.py
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
from typing import Any, Dict, Optional, Tuple, Union
|
||||
|
||||
from einops import rearrange
|
||||
|
||||
from diffusers.utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers
|
||||
from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel, UNet2DConditionOutput
|
||||
from .hm_blocks import CopyWeights, InsertReferenceAdapter
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
class HMDenoising3D(UNet2DConditionModel, CopyWeights, InsertReferenceAdapter):
|
||||
def forward(
|
||||
self,
|
||||
sample: torch.Tensor,
|
||||
timestep: Union[torch.Tensor, float, int],
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
reference_hidden_states: Optional[dict] = None,
|
||||
control_hidden_states: Optional[torch.Tensor] = None,
|
||||
class_labels: Optional[torch.Tensor] = None,
|
||||
timestep_cond: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
|
||||
down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
|
||||
mid_block_additional_residual: Optional[torch.Tensor] = None,
|
||||
down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
|
||||
encoder_attention_mask: Optional[torch.Tensor] = None,
|
||||
return_dict: bool = True,
|
||||
) -> Union[UNet2DConditionOutput, Tuple]:
|
||||
# By default samples have to be AT least a multiple of the overall upsampling factor.
|
||||
# The overall upsampling factor is equal to 2 ** (# num of upsampling layers).
|
||||
# However, the upsampling interpolation output size can be forced to fit any upsampling size
|
||||
# on the fly if necessary.
|
||||
default_overall_up_factor = 2**self.num_upsamplers
|
||||
|
||||
# upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
|
||||
forward_upsample_size = False
|
||||
upsample_size = None
|
||||
|
||||
for dim in sample.shape[-2:]:
|
||||
if dim % default_overall_up_factor != 0:
|
||||
# Forward upsample size to force interpolation output size.
|
||||
forward_upsample_size = True
|
||||
break
|
||||
|
||||
# ensure attention_mask is a bias, and give it a singleton query_tokens dimension
|
||||
# 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)
|
||||
if attention_mask is not None:
|
||||
# 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)
|
||||
attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
|
||||
# convert encoder_attention_mask to a bias the same way we do for attention_mask
|
||||
if encoder_attention_mask is not None:
|
||||
encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
|
||||
encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
|
||||
|
||||
# 0. center input if necessary
|
||||
if self.config.center_input_sample:
|
||||
sample = 2 * sample - 1.0
|
||||
|
||||
# 1. time
|
||||
t_emb = self.get_time_embed(sample=sample, timestep=timestep)
|
||||
emb = self.time_embedding(t_emb, timestep_cond)
|
||||
aug_emb = None
|
||||
|
||||
class_emb = self.get_class_embed(sample=sample, class_labels=class_labels)
|
||||
if class_emb is not None:
|
||||
if self.config.class_embeddings_concat:
|
||||
emb = torch.cat([emb, class_emb], dim=-1)
|
||||
else:
|
||||
emb = emb + class_emb
|
||||
|
||||
aug_emb = self.get_aug_embed(
|
||||
emb=emb, encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
|
||||
)
|
||||
if self.config.addition_embed_type == "image_hint":
|
||||
aug_emb, hint = aug_emb
|
||||
sample = torch.cat([sample, hint], dim=1)
|
||||
|
||||
emb = emb + aug_emb if aug_emb is not None else emb
|
||||
|
||||
if self.time_embed_act is not None:
|
||||
emb = self.time_embed_act(emb)
|
||||
|
||||
num_frames = sample.shape[2]
|
||||
emb = emb.repeat_interleave(repeats=num_frames, dim=0)
|
||||
|
||||
if len(encoder_hidden_states.shape) == 3:
|
||||
encoder_hidden_states = encoder_hidden_states.repeat_interleave(repeats=num_frames, dim=0)
|
||||
elif len(encoder_hidden_states.shape) == 4:
|
||||
encoder_hidden_states = rearrange(encoder_hidden_states, "b f l d -> (b f) l d")
|
||||
|
||||
if not added_cond_kwargs is None and 'image_embeds' in added_cond_kwargs:
|
||||
if isinstance(added_cond_kwargs['image_embeds'], torch.Tensor):
|
||||
added_cond_kwargs['image_embeds'] = added_cond_kwargs['image_embeds'].repeat_interleave(repeats=num_frames, dim=0)
|
||||
else:
|
||||
added_cond_kwargs['image_embeds'] = [x.repeat_interleave(repeats=num_frames, dim=0) for x in added_cond_kwargs['image_embeds']]
|
||||
|
||||
encoder_hidden_states = self.process_encoder_hidden_states(
|
||||
encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
|
||||
)
|
||||
|
||||
# 2. pre-process
|
||||
sample = rearrange(sample, "b c f h w -> (b f) c h w")
|
||||
sample = self.conv_in(sample)
|
||||
|
||||
# 2.5 GLIGEN position net
|
||||
if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None:
|
||||
cross_attention_kwargs = cross_attention_kwargs.copy()
|
||||
gligen_args = cross_attention_kwargs.pop("gligen")
|
||||
cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)}
|
||||
|
||||
# 3. down
|
||||
# we're popping the `scale` instead of getting it because otherwise `scale` will be propagated
|
||||
# to the internal blocks and will raise deprecation warnings. this will be confusing for our users.
|
||||
if cross_attention_kwargs is not None:
|
||||
cross_attention_kwargs = cross_attention_kwargs.copy()
|
||||
lora_scale = cross_attention_kwargs.pop("scale", 1.0)
|
||||
else:
|
||||
lora_scale = 1.0
|
||||
|
||||
if USE_PEFT_BACKEND:
|
||||
# weight the lora layers by setting `lora_scale` for each PEFT layer
|
||||
scale_lora_layers(self, lora_scale)
|
||||
|
||||
is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
|
||||
# using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets
|
||||
is_adapter = down_intrablock_additional_residuals is not None
|
||||
# maintain backward compatibility for legacy usage, where
|
||||
# T2I-Adapter and ControlNet both use down_block_additional_residuals arg
|
||||
# but can only use one or the other
|
||||
if not is_adapter and mid_block_additional_residual is None and down_block_additional_residuals is not None:
|
||||
deprecate(
|
||||
"T2I should not use down_block_additional_residuals",
|
||||
"1.3.0",
|
||||
"Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \
|
||||
and will be removed in diffusers 1.3.0. `down_block_additional_residuals` should only be used \
|
||||
for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ",
|
||||
standard_warn=False,
|
||||
)
|
||||
down_intrablock_additional_residuals = down_block_additional_residuals
|
||||
is_adapter = True
|
||||
|
||||
res_cache = dict()
|
||||
down_block_res_samples = (sample,)
|
||||
for idx, downsample_block in enumerate(self.down_blocks):
|
||||
if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
|
||||
# For t2i-adapter CrossAttnDownBlock2D
|
||||
additional_residuals = {}
|
||||
if is_adapter and len(down_intrablock_additional_residuals) > 0:
|
||||
additional_residuals["additional_residuals"] = down_intrablock_additional_residuals.pop(0)
|
||||
|
||||
sample, res_samples = downsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
**additional_residuals,
|
||||
)
|
||||
res_cache[f"down_{idx}"] = sample.clone()
|
||||
else:
|
||||
sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
|
||||
if is_adapter and len(down_intrablock_additional_residuals) > 0:
|
||||
sample += down_intrablock_additional_residuals.pop(0)
|
||||
|
||||
if not control_hidden_states is None and f'down_{idx}' in control_hidden_states:
|
||||
sample += rearrange(control_hidden_states[f'down_{idx}'], "b c f h w -> (b f) c h w")
|
||||
if hasattr(self, 'reference_modules_down') and not reference_hidden_states is None and f'down_{idx}' in reference_hidden_states:
|
||||
sample = self.reference_modules_down[idx](sample, reference_hidden_states[f'down_{idx}'], num_frames)
|
||||
|
||||
down_block_res_samples += res_samples
|
||||
|
||||
if is_controlnet:
|
||||
new_down_block_res_samples = ()
|
||||
|
||||
for down_block_res_sample, down_block_additional_residual in zip(
|
||||
down_block_res_samples, down_block_additional_residuals
|
||||
):
|
||||
down_block_res_sample = down_block_res_sample + down_block_additional_residual
|
||||
new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
|
||||
|
||||
down_block_res_samples = new_down_block_res_samples
|
||||
|
||||
# 4. mid
|
||||
if self.mid_block is not None:
|
||||
if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention:
|
||||
sample = self.mid_block(
|
||||
sample,
|
||||
emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
)
|
||||
else:
|
||||
sample = self.mid_block(sample, emb)
|
||||
if hasattr(self, 'reference_modules_mid') and not reference_hidden_states is None and f'mid' in reference_hidden_states:
|
||||
sample = self.reference_modules_mid(sample, reference_hidden_states[f'mid'], num_frames)
|
||||
|
||||
# To support T2I-Adapter-XL
|
||||
if (
|
||||
is_adapter
|
||||
and len(down_intrablock_additional_residuals) > 0
|
||||
and sample.shape == down_intrablock_additional_residuals[0].shape
|
||||
):
|
||||
sample += down_intrablock_additional_residuals.pop(0)
|
||||
res_cache[f"mid"] = sample.clone()
|
||||
|
||||
if is_controlnet:
|
||||
sample = sample + mid_block_additional_residual
|
||||
|
||||
# 5. up
|
||||
for i, upsample_block in enumerate(self.up_blocks):
|
||||
is_final_block = i == len(self.up_blocks) - 1
|
||||
|
||||
res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
|
||||
down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
|
||||
|
||||
# if we have not reached the final block and need to forward the
|
||||
# upsample size, we do it here
|
||||
if not is_final_block and forward_upsample_size:
|
||||
upsample_size = down_block_res_samples[-1].shape[2:]
|
||||
|
||||
if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
|
||||
res_cache[f"up_{i}"] = sample.clone()
|
||||
if hasattr(self, "reference_modules_up") and not reference_hidden_states is None and f'up_{i}' in reference_hidden_states:
|
||||
sample = self.reference_modules_up[i-1](sample, reference_hidden_states[f'up_{i}'], num_frames)
|
||||
|
||||
sample = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
upsample_size=upsample_size,
|
||||
attention_mask=attention_mask,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
)
|
||||
else:
|
||||
sample = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
upsample_size=upsample_size,
|
||||
)
|
||||
|
||||
# 6. post-process
|
||||
if self.conv_norm_out:
|
||||
sample = self.conv_norm_out(sample)
|
||||
sample = self.conv_act(sample)
|
||||
sample = self.conv_out(sample)
|
||||
|
||||
if USE_PEFT_BACKEND:
|
||||
# remove `lora_scale` from each PEFT layer
|
||||
unscale_lora_layers(self, lora_scale)
|
||||
|
||||
# reshape to (batch, channel, framerate, width, height)
|
||||
sample = rearrange(sample, "(b f) c h w -> b c f h w", f=num_frames)
|
||||
|
||||
if not return_dict:
|
||||
return (sample, res_cache)
|
||||
|
||||
return (UNet2DConditionOutput(sample=sample), res_cache)
|
||||
@@ -0,0 +1,241 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
@File : models6/hm_denoising_motion.py
|
||||
@Author : Songkey
|
||||
@Email : songkey@pku.edu.cn
|
||||
@Date : 9/9/2024
|
||||
@Desc :
|
||||
adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unets/unet_motion_model.py
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
from typing import Any, Dict, Optional, Tuple, Union
|
||||
|
||||
from einops import rearrange
|
||||
|
||||
from diffusers.utils import logging
|
||||
from diffusers.models.unets.unet_motion_model import UNetMotionModel, UNetMotionOutput
|
||||
from .hm_blocks import InsertReferenceAdapter
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
class HMDenoisingMotion(UNetMotionModel, InsertReferenceAdapter):
|
||||
def forward(
|
||||
self,
|
||||
sample: torch.Tensor,
|
||||
timestep: Union[torch.Tensor, float, int],
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
reference_hidden_states: Optional[dict] = None,
|
||||
control_hidden_states: Optional[torch.Tensor] = None,
|
||||
timestep_cond: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
|
||||
down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
|
||||
mid_block_additional_residual: Optional[torch.Tensor] = None,
|
||||
return_dict: bool = True,
|
||||
) -> Union[UNetMotionOutput, Tuple[torch.Tensor]]:
|
||||
|
||||
# By default samples have to be AT least a multiple of the overall upsampling factor.
|
||||
# The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
|
||||
# However, the upsampling interpolation output size can be forced to fit any upsampling size
|
||||
# on the fly if necessary.
|
||||
default_overall_up_factor = 2 ** self.num_upsamplers
|
||||
|
||||
# upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
|
||||
forward_upsample_size = False
|
||||
upsample_size = None
|
||||
|
||||
if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
|
||||
logger.info("Forward upsample size to force interpolation output size.")
|
||||
forward_upsample_size = True
|
||||
|
||||
# prepare attention_mask
|
||||
if attention_mask is not None:
|
||||
attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
|
||||
# 1. time
|
||||
timesteps = timestep
|
||||
if not torch.is_tensor(timesteps):
|
||||
# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
|
||||
# This would be a good case for the `match` statement (Python 3.10+)
|
||||
is_mps = sample.device.type == "mps"
|
||||
if isinstance(timestep, float):
|
||||
dtype = torch.float32 if is_mps else torch.float64
|
||||
else:
|
||||
dtype = torch.int32 if is_mps else torch.int64
|
||||
timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
|
||||
elif len(timesteps.shape) == 0:
|
||||
timesteps = timesteps[None].to(sample.device)
|
||||
|
||||
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
||||
num_frames = sample.shape[2]
|
||||
timesteps = timesteps.expand(sample.shape[0])
|
||||
|
||||
t_emb = self.time_proj(timesteps)
|
||||
|
||||
# timesteps does not contain any weights and will always return f32 tensors
|
||||
# but time_embedding might actually be running in fp16. so we need to cast here.
|
||||
# there might be better ways to encapsulate this.
|
||||
t_emb = t_emb.to(dtype=self.dtype)
|
||||
|
||||
emb = self.time_embedding(t_emb, timestep_cond)
|
||||
aug_emb = None
|
||||
|
||||
if self.config.addition_embed_type == "text_time":
|
||||
if "text_embeds" not in added_cond_kwargs:
|
||||
raise ValueError(
|
||||
f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
|
||||
)
|
||||
|
||||
text_embeds = added_cond_kwargs.get("text_embeds")
|
||||
if "time_ids" not in added_cond_kwargs:
|
||||
raise ValueError(
|
||||
f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
|
||||
)
|
||||
time_ids = added_cond_kwargs.get("time_ids")
|
||||
time_embeds = self.add_time_proj(time_ids.flatten())
|
||||
time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
|
||||
|
||||
add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
|
||||
add_embeds = add_embeds.to(emb.dtype)
|
||||
aug_emb = self.add_embedding(add_embeds)
|
||||
|
||||
emb = emb if aug_emb is None else emb + aug_emb
|
||||
emb = emb.repeat_interleave(repeats=num_frames, dim=0)
|
||||
if len(encoder_hidden_states.shape) == 3:
|
||||
encoder_hidden_states = encoder_hidden_states.repeat_interleave(repeats=num_frames, dim=0)
|
||||
elif len(encoder_hidden_states.shape) == 4:
|
||||
encoder_hidden_states = rearrange(encoder_hidden_states, "b f l d -> (b f) l d")
|
||||
|
||||
if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "ip_image_proj":
|
||||
if "image_embeds" not in added_cond_kwargs:
|
||||
raise ValueError(
|
||||
f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'ip_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`"
|
||||
)
|
||||
image_embeds = added_cond_kwargs.get("image_embeds")
|
||||
image_embeds = self.encoder_hid_proj(image_embeds)
|
||||
image_embeds = [image_embed.repeat_interleave(repeats=num_frames, dim=0) for image_embed in image_embeds]
|
||||
encoder_hidden_states = (encoder_hidden_states, image_embeds)
|
||||
|
||||
# 2. pre-process
|
||||
sample = sample.permute(0, 2, 1, 3, 4).reshape((sample.shape[0] * num_frames, -1) + sample.shape[3:])
|
||||
sample = self.conv_in(sample)
|
||||
|
||||
# 3. down
|
||||
down_block_res_samples = (sample,)
|
||||
for idx, downsample_block in enumerate(self.down_blocks):
|
||||
if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
|
||||
sample, res_samples = downsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
num_frames=num_frames,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
)
|
||||
else:
|
||||
sample, res_samples = downsample_block(hidden_states=sample, temb=emb, num_frames=num_frames)
|
||||
|
||||
if not control_hidden_states is None and f'down_{idx}' in control_hidden_states:
|
||||
sample += rearrange(control_hidden_states[f'down_{idx}'], "b c f h w -> (b f) c h w")
|
||||
|
||||
|
||||
if hasattr(self, 'reference_modules_down') and not reference_hidden_states is None and f'down_{idx}' in reference_hidden_states:
|
||||
sample = self.reference_modules_down[idx](sample, reference_hidden_states[f'down_{idx}'], num_frames)
|
||||
|
||||
down_block_res_samples += res_samples
|
||||
|
||||
if down_block_additional_residuals is not None:
|
||||
new_down_block_res_samples = ()
|
||||
|
||||
for down_block_res_sample, down_block_additional_residual in zip(
|
||||
down_block_res_samples, down_block_additional_residuals
|
||||
):
|
||||
down_block_res_sample = down_block_res_sample + down_block_additional_residual
|
||||
new_down_block_res_samples += (down_block_res_sample,)
|
||||
|
||||
down_block_res_samples = new_down_block_res_samples
|
||||
|
||||
# 4. mid
|
||||
if self.mid_block is not None:
|
||||
# To support older versions of motion modules that don't have a mid_block
|
||||
if hasattr(self.mid_block, "motion_modules"):
|
||||
sample = self.mid_block(
|
||||
sample,
|
||||
emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
num_frames=num_frames,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
)
|
||||
else:
|
||||
sample = self.mid_block(
|
||||
sample,
|
||||
emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
)
|
||||
if hasattr(self, 'reference_modules_mid') and not reference_hidden_states is None and f'mid' in reference_hidden_states:
|
||||
sample = self.reference_modules_mid(sample, reference_hidden_states[f'mid'], num_frames)
|
||||
|
||||
if mid_block_additional_residual is not None:
|
||||
sample = sample + mid_block_additional_residual
|
||||
|
||||
# 5. up
|
||||
for i, upsample_block in enumerate(self.up_blocks):
|
||||
is_final_block = i == len(self.up_blocks) - 1
|
||||
|
||||
res_samples = down_block_res_samples[-len(upsample_block.resnets):]
|
||||
down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
|
||||
|
||||
# if we have not reached the final block and need to forward the
|
||||
# upsample size, we do it here
|
||||
if not is_final_block and forward_upsample_size:
|
||||
upsample_size = down_block_res_samples[-1].shape[2:]
|
||||
|
||||
if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
|
||||
if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
|
||||
if hasattr(self,
|
||||
"reference_modules_up") and not reference_hidden_states is None and f'up_{i}' in reference_hidden_states:
|
||||
sample = self.reference_modules_up[i - 1](sample, reference_hidden_states[f'up_{i}'],
|
||||
num_frames)
|
||||
|
||||
sample = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
upsample_size=upsample_size,
|
||||
attention_mask=attention_mask,
|
||||
num_frames=num_frames,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
)
|
||||
else:
|
||||
sample = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
upsample_size=upsample_size,
|
||||
num_frames=num_frames,
|
||||
)
|
||||
|
||||
# 6. post-process
|
||||
if self.conv_norm_out:
|
||||
sample = self.conv_norm_out(sample)
|
||||
sample = self.conv_act(sample)
|
||||
|
||||
sample = self.conv_out(sample)
|
||||
|
||||
# reshape to (batch, channel, framerate, width, height)
|
||||
sample = sample[None, :].reshape((-1, num_frames) + sample.shape[1:]).permute(0, 2, 1, 3, 4)
|
||||
|
||||
if not return_dict:
|
||||
return (sample,)
|
||||
|
||||
return UNetMotionOutput(sample=sample)
|
||||
Reference in New Issue
Block a user