887 lines
39 KiB
Python
887 lines
39 KiB
Python
# Copyright 2023 The HuggingFace Team. All rights reserved.
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
from dataclasses import dataclass
|
||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||
|
||
import torch
|
||
from torch import nn
|
||
from torch.nn import functional as F
|
||
|
||
from ..diffusers.configuration_utils import ConfigMixin, register_to_config
|
||
from ..diffusers.loaders import FromOriginalControlnetMixin
|
||
from ..diffusers.utils import BaseOutput, logging
|
||
from ..diffusers.models.attention_processor import (
|
||
ADDED_KV_ATTENTION_PROCESSORS,
|
||
CROSS_ATTENTION_PROCESSORS,
|
||
AttentionProcessor,
|
||
AttnAddedKVProcessor,
|
||
AttnProcessor,
|
||
)
|
||
from ..diffusers.models.embeddings import TextImageProjection, TextImageTimeEmbedding, TextTimeEmbedding, TimestepEmbedding, Timesteps
|
||
from ..diffusers.models.modeling_utils import ModelMixin
|
||
from ..diffusers.models.unet_3d_blocks import (
|
||
get_down_block, get_up_block,UNetMidBlockSpatioTemporal,
|
||
)
|
||
from ..diffusers.models import UNetSpatioTemporalConditionModel
|
||
|
||
# import torch
|
||
# import torch.nn.functional as F
|
||
from torch import nn, einsum
|
||
from inspect import isfunction
|
||
from einops import rearrange, repeat
|
||
|
||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||
|
||
def exists(val):
|
||
return val is not None
|
||
|
||
def default(val, d):
|
||
if exists(val):
|
||
return val
|
||
return d() if isfunction(d) else d
|
||
|
||
class CrossAttention(nn.Module):
|
||
def __init__(self, query_dim=320, context_dim=1024, heads=8, dim_head=64, dropout=0.):
|
||
super().__init__()
|
||
inner_dim = dim_head * heads
|
||
context_dim = default(context_dim, query_dim)
|
||
|
||
self.scale = dim_head ** -0.5
|
||
self.heads = heads
|
||
|
||
self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
|
||
self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
|
||
self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
|
||
|
||
self.to_out = nn.Sequential(
|
||
nn.Linear(inner_dim, query_dim),
|
||
nn.Dropout(dropout)
|
||
)
|
||
|
||
def forward(self, x, context=None, mask=None):
|
||
h = self.heads
|
||
|
||
q = self.to_q(x)
|
||
context = default(context, x)
|
||
k = self.to_k(context)
|
||
v = self.to_v(context)
|
||
|
||
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
|
||
|
||
|
||
sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
|
||
|
||
if exists(mask):
|
||
mask = rearrange(mask, 'b ... -> b (...)')
|
||
max_neg_value = -torch.finfo(sim.dtype).max
|
||
mask = repeat(mask, 'b j -> (b h) () j', h=h)
|
||
sim.masked_fill_(~mask, max_neg_value)
|
||
|
||
# attention, what we cannot get enough of
|
||
attn = sim.softmax(dim=-1)
|
||
|
||
out = einsum('b i j, b j d -> b i d', attn, v)
|
||
out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
|
||
return self.to_out(out)
|
||
|
||
@dataclass
|
||
class ControlNetOutput(BaseOutput):
|
||
"""
|
||
The output of [`ControlNetModel`].
|
||
|
||
Args:
|
||
down_block_res_samples (`tuple[torch.Tensor]`):
|
||
A tuple of downsample activations at different resolutions for each downsampling block. Each tensor should
|
||
be of shape `(batch_size, channel * resolution, height //resolution, width // resolution)`. Output can be
|
||
used to condition the original UNet's downsampling activations.
|
||
mid_down_block_re_sample (`torch.Tensor`):
|
||
The activation of the midde block (the lowest sample resolution). Each tensor should be of shape
|
||
`(batch_size, channel * lowest_resolution, height // lowest_resolution, width // lowest_resolution)`.
|
||
Output can be used to condition the original UNet's middle block activation.
|
||
"""
|
||
|
||
down_block_res_samples: Tuple[torch.Tensor]
|
||
mid_block_res_sample: torch.Tensor
|
||
|
||
|
||
|
||
class ControlNetConditioningEmbeddingSVD(nn.Module):
|
||
"""
|
||
Quoting from https://arxiv.org/abs/2302.05543: "Stable Diffusion uses a pre-processing method similar to VQ-GAN
|
||
[11] to convert the entire dataset of 512 × 512 images into smaller 64 × 64 “latent images” for stabilized
|
||
training. This requires ControlNets to convert image-based conditions to 64 × 64 feature space to match the
|
||
convolution size. We use a tiny network E(·) of four convolution layers with 4 × 4 kernels and 2 × 2 strides
|
||
(activated by ReLU, channels are 16, 32, 64, 128, initialized with Gaussian weights, trained jointly with the full
|
||
model) to encode image-space conditions ... into feature maps ..."
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
conditioning_embedding_channels: int,
|
||
conditioning_channels: int = 3,
|
||
block_out_channels: Tuple[int, ...] = (16, 32, 96, 256),
|
||
feature_channels: int = 1024,
|
||
feature_out_channels: Tuple[int, ...] = (160, 160, 256,256),
|
||
):
|
||
super().__init__()
|
||
|
||
self.conv_in = nn.Conv2d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1)
|
||
|
||
self.blocks = nn.ModuleList([])
|
||
|
||
for i in range(len(block_out_channels) - 1):
|
||
channel_in = block_out_channels[i]
|
||
channel_out = block_out_channels[i + 1]
|
||
self.blocks.append(nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1))
|
||
self.blocks.append(nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2))
|
||
self.conv_out = zero_module(
|
||
nn.Conv2d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1)
|
||
)
|
||
|
||
self.blocks_feature = nn.ModuleList([])
|
||
|
||
for i in range(len(feature_out_channels) - 1):
|
||
channel_in = feature_out_channels[i]
|
||
channel_out = feature_out_channels[i + 1]
|
||
self.blocks_feature.append(nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1))
|
||
self.blocks_feature.append(nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2))
|
||
|
||
|
||
self.conv_out_feature = zero_module(
|
||
nn.Conv2d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1)
|
||
)
|
||
def forward(self, conditioning, feature_embedding):
|
||
#this seeems appropriate? idk if i should be applying a more complex setup to handle the frames
|
||
#combine batch and frames dimensions
|
||
batch_size, frames, channels, height, width = conditioning.size()
|
||
conditioning = conditioning.view(batch_size * frames, channels, height, width)
|
||
|
||
embedding = self.conv_in(conditioning)
|
||
embedding = F.silu(embedding)
|
||
|
||
for block in self.blocks:
|
||
embedding = block(embedding)
|
||
embedding = F.silu(embedding)
|
||
embedding = self.conv_out(embedding)
|
||
|
||
batch_size, frames, channels, height, width = feature_embedding.size()
|
||
feature_embedding = feature_embedding.view(batch_size * frames, channels, height, width)
|
||
for block in self.blocks_feature:
|
||
feature_embedding = block(feature_embedding)
|
||
feature_embedding = F.silu(feature_embedding)
|
||
|
||
feature_embedding = self.conv_out_feature(feature_embedding)
|
||
#split them apart again
|
||
#actually not needed
|
||
#new_channels, new_height, new_width = embedding.shape[1], embedding.shape[2], embedding.shape[3]
|
||
#embedding = embedding.view(batch_size, frames, new_channels, new_height, new_width)
|
||
#
|
||
return embedding+feature_embedding
|
||
|
||
|
||
class DragAnythingSDVModel(ModelMixin, ConfigMixin, FromOriginalControlnetMixin):
|
||
r"""
|
||
A conditional Spatio-Temporal UNet model that takes a noisy video frames, conditional state, and a timestep and returns a sample
|
||
shaped output.
|
||
|
||
This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
|
||
for all models (such as downloading or saving).
|
||
|
||
Parameters:
|
||
sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
|
||
Height and width of input/output sample.
|
||
in_channels (`int`, *optional*, defaults to 8): Number of channels in the input sample.
|
||
out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
|
||
down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "DownBlockSpatioTemporal")`):
|
||
The tuple of downsample blocks to use.
|
||
up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal")`):
|
||
The tuple of upsample blocks to use.
|
||
block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
|
||
The tuple of output channels for each block.
|
||
addition_time_embed_dim: (`int`, defaults to 256):
|
||
Dimension to to encode the additional time ids.
|
||
projection_class_embeddings_input_dim (`int`, defaults to 768):
|
||
The dimension of the projection of encoded `added_time_ids`.
|
||
layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
|
||
cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
|
||
The dimension of the cross attention features.
|
||
transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
|
||
The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
|
||
[`~models.unet_3d_blocks.CrossAttnDownBlockSpatioTemporal`], [`~models.unet_3d_blocks.CrossAttnUpBlockSpatioTemporal`],
|
||
[`~models.unet_3d_blocks.UNetMidBlockSpatioTemporal`].
|
||
num_attention_heads (`int`, `Tuple[int]`, defaults to `(5, 10, 10, 20)`):
|
||
The number of attention heads.
|
||
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
|
||
"""
|
||
|
||
_supports_gradient_checkpointing = True
|
||
|
||
@register_to_config
|
||
def __init__(
|
||
self,
|
||
sample_size: Optional[int] = None,
|
||
in_channels: int = 8,
|
||
out_channels: int = 4,
|
||
down_block_types: Tuple[str] = (
|
||
"CrossAttnDownBlockSpatioTemporal",
|
||
"CrossAttnDownBlockSpatioTemporal",
|
||
"CrossAttnDownBlockSpatioTemporal",
|
||
"DownBlockSpatioTemporal",
|
||
),
|
||
up_block_types: Tuple[str] = (
|
||
"UpBlockSpatioTemporal",
|
||
"CrossAttnUpBlockSpatioTemporal",
|
||
"CrossAttnUpBlockSpatioTemporal",
|
||
"CrossAttnUpBlockSpatioTemporal",
|
||
),
|
||
block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
|
||
addition_time_embed_dim: int = 256,
|
||
projection_class_embeddings_input_dim: int = 768,
|
||
layers_per_block: Union[int, Tuple[int]] = 2,
|
||
cross_attention_dim: Union[int, Tuple[int]] = 1024,
|
||
transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
|
||
num_attention_heads: Union[int, Tuple[int]] = (5, 10, 10, 20),
|
||
num_frames: int = 25,
|
||
conditioning_channels: int = 3,
|
||
conditioning_embedding_out_channels : Optional[Tuple[int, ...]] = (16, 32, 96, 256),
|
||
):
|
||
super().__init__()
|
||
self.sample_size = sample_size
|
||
|
||
print("layers per block is", layers_per_block)
|
||
|
||
# Check inputs
|
||
if len(down_block_types) != len(up_block_types):
|
||
raise ValueError(
|
||
f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
|
||
)
|
||
|
||
if len(block_out_channels) != len(down_block_types):
|
||
raise ValueError(
|
||
f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
|
||
)
|
||
|
||
if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
|
||
raise ValueError(
|
||
f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
|
||
)
|
||
|
||
if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):
|
||
raise ValueError(
|
||
f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
|
||
)
|
||
|
||
if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types):
|
||
raise ValueError(
|
||
f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}."
|
||
)
|
||
|
||
# input
|
||
self.conv_in = nn.Conv2d(
|
||
in_channels,
|
||
block_out_channels[0],
|
||
kernel_size=3,
|
||
padding=1,
|
||
)
|
||
|
||
|
||
|
||
# time
|
||
time_embed_dim = block_out_channels[0] * 4
|
||
|
||
self.time_proj = Timesteps(block_out_channels[0], True, downscale_freq_shift=0)
|
||
timestep_input_dim = block_out_channels[0]
|
||
|
||
self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
|
||
|
||
self.add_time_proj = Timesteps(addition_time_embed_dim, True, downscale_freq_shift=0)
|
||
self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
|
||
|
||
self.down_blocks = nn.ModuleList([])
|
||
self.controlnet_down_blocks = nn.ModuleList([])
|
||
|
||
if isinstance(num_attention_heads, int):
|
||
num_attention_heads = (num_attention_heads,) * len(down_block_types)
|
||
|
||
if isinstance(cross_attention_dim, int):
|
||
cross_attention_dim = (cross_attention_dim,) * len(down_block_types)
|
||
|
||
if isinstance(layers_per_block, int):
|
||
layers_per_block = [layers_per_block] * len(down_block_types)
|
||
|
||
if isinstance(transformer_layers_per_block, int):
|
||
transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
|
||
|
||
blocks_time_embed_dim = time_embed_dim
|
||
self.controlnet_cond_embedding = ControlNetConditioningEmbeddingSVD(
|
||
conditioning_embedding_channels=block_out_channels[0],
|
||
block_out_channels=conditioning_embedding_out_channels,
|
||
conditioning_channels=conditioning_channels,
|
||
)
|
||
|
||
# down
|
||
output_channel = block_out_channels[0]
|
||
controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
|
||
controlnet_block = zero_module(controlnet_block)
|
||
self.controlnet_down_blocks.append(controlnet_block)
|
||
|
||
|
||
|
||
for i, down_block_type in enumerate(down_block_types):
|
||
input_channel = output_channel
|
||
output_channel = block_out_channels[i]
|
||
is_final_block = i == len(block_out_channels) - 1
|
||
|
||
down_block = get_down_block(
|
||
down_block_type,
|
||
num_layers=layers_per_block[i],
|
||
transformer_layers_per_block=transformer_layers_per_block[i],
|
||
in_channels=input_channel,
|
||
out_channels=output_channel,
|
||
temb_channels=blocks_time_embed_dim,
|
||
add_downsample=not is_final_block,
|
||
resnet_eps=1e-5,
|
||
cross_attention_dim=cross_attention_dim[i],
|
||
num_attention_heads=num_attention_heads[i],
|
||
resnet_act_fn="silu",
|
||
)
|
||
self.down_blocks.append(down_block)
|
||
|
||
for _ in range(layers_per_block[i]):
|
||
controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
|
||
controlnet_block = zero_module(controlnet_block)
|
||
self.controlnet_down_blocks.append(controlnet_block)
|
||
|
||
if not is_final_block:
|
||
controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
|
||
controlnet_block = zero_module(controlnet_block)
|
||
self.controlnet_down_blocks.append(controlnet_block)
|
||
|
||
|
||
# mid
|
||
mid_block_channel = block_out_channels[-1]
|
||
controlnet_block = nn.Conv2d(mid_block_channel, mid_block_channel, kernel_size=1)
|
||
controlnet_block = zero_module(controlnet_block)
|
||
self.controlnet_mid_block = controlnet_block
|
||
|
||
|
||
self.mid_block = UNetMidBlockSpatioTemporal(
|
||
block_out_channels[-1],
|
||
temb_channels=blocks_time_embed_dim,
|
||
transformer_layers_per_block=transformer_layers_per_block[-1],
|
||
cross_attention_dim=cross_attention_dim[-1],
|
||
num_attention_heads=num_attention_heads[-1],
|
||
)
|
||
|
||
|
||
|
||
|
||
# out
|
||
#self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=32, eps=1e-5)
|
||
#self.conv_act = nn.SiLU()
|
||
|
||
#self.conv_out = nn.Conv2d(
|
||
# block_out_channels[0],
|
||
# out_channels,
|
||
# kernel_size=3,
|
||
# padding=1,
|
||
#)
|
||
|
||
# ID feature
|
||
# self.CrossAttention = CrossAttention()
|
||
|
||
|
||
@property
|
||
def attn_processors(self) -> Dict[str, AttentionProcessor]:
|
||
r"""
|
||
Returns:
|
||
`dict` of attention processors: A dictionary containing all attention processors used in the model with
|
||
indexed by its weight name.
|
||
"""
|
||
# set recursively
|
||
processors = {}
|
||
|
||
def fn_recursive_add_processors(
|
||
name: str,
|
||
module: torch.nn.Module,
|
||
processors: Dict[str, AttentionProcessor],
|
||
):
|
||
if hasattr(module, "get_processor"):
|
||
processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
|
||
|
||
for sub_name, child in module.named_children():
|
||
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
|
||
|
||
return processors
|
||
|
||
for name, module in self.named_children():
|
||
fn_recursive_add_processors(name, module, processors)
|
||
|
||
return processors
|
||
|
||
def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
|
||
r"""
|
||
Sets the attention processor to use to compute attention.
|
||
|
||
Parameters:
|
||
processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
|
||
The instantiated processor class or a dictionary of processor classes that will be set as the processor
|
||
for **all** `Attention` layers.
|
||
|
||
If `processor` is a dict, the key needs to define the path to the corresponding cross attention
|
||
processor. This is strongly recommended when setting trainable attention processors.
|
||
|
||
"""
|
||
count = len(self.attn_processors.keys())
|
||
|
||
if isinstance(processor, dict) and len(processor) != count:
|
||
raise ValueError(
|
||
f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
|
||
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
|
||
)
|
||
|
||
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
|
||
if hasattr(module, "set_processor"):
|
||
if not isinstance(processor, dict):
|
||
module.set_processor(processor)
|
||
else:
|
||
module.set_processor(processor.pop(f"{name}.processor"))
|
||
|
||
for sub_name, child in module.named_children():
|
||
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
|
||
|
||
for name, module in self.named_children():
|
||
fn_recursive_attn_processor(name, module, processor)
|
||
|
||
def set_default_attn_processor(self):
|
||
"""
|
||
Disables custom attention processors and sets the default attention implementation.
|
||
"""
|
||
if all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
|
||
processor = AttnProcessor()
|
||
else:
|
||
raise ValueError(
|
||
f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
|
||
)
|
||
|
||
self.set_attn_processor(processor)
|
||
|
||
def _set_gradient_checkpointing(self, module, value=False):
|
||
if hasattr(module, "gradient_checkpointing"):
|
||
module.gradient_checkpointing = value
|
||
|
||
# Copied from diffusers.models.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
|
||
def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
|
||
"""
|
||
Sets the attention processor to use [feed forward
|
||
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
|
||
|
||
Parameters:
|
||
chunk_size (`int`, *optional*):
|
||
The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
|
||
over each tensor of dim=`dim`.
|
||
dim (`int`, *optional*, defaults to `0`):
|
||
The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
|
||
or dim=1 (sequence length).
|
||
"""
|
||
if dim not in [0, 1]:
|
||
raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}")
|
||
|
||
# By default chunk size is 1
|
||
chunk_size = chunk_size or 1
|
||
|
||
def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
|
||
if hasattr(module, "set_chunk_feed_forward"):
|
||
module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
|
||
|
||
for child in module.children():
|
||
fn_recursive_feed_forward(child, chunk_size, dim)
|
||
|
||
for module in self.children():
|
||
fn_recursive_feed_forward(module, chunk_size, dim)
|
||
|
||
|
||
|
||
def forward(
|
||
self,
|
||
sample: torch.FloatTensor,
|
||
timestep: Union[torch.Tensor, float, int],
|
||
encoder_hidden_states: torch.Tensor,
|
||
added_time_ids: torch.Tensor,
|
||
controlnet_cond: torch.FloatTensor = None,
|
||
ID_feature: torch.FloatTensor = None,
|
||
image_only_indicator: Optional[torch.Tensor] = None,
|
||
return_dict: bool = True,
|
||
guess_mode: bool = False,
|
||
conditioning_scale: float = 1.0,
|
||
do_classifier_free_guidance = False,
|
||
|
||
|
||
) -> Union[ControlNetOutput, Tuple]:
|
||
r"""
|
||
The [`UNetSpatioTemporalConditionModel`] forward method.
|
||
|
||
Args:
|
||
sample (`torch.FloatTensor`):
|
||
The noisy input tensor with the following shape `(batch, num_frames, channel, height, width)`.
|
||
timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.
|
||
encoder_hidden_states (`torch.FloatTensor`):
|
||
The encoder hidden states with shape `(batch, sequence_length, cross_attention_dim)`.
|
||
added_time_ids: (`torch.FloatTensor`):
|
||
The additional time ids with shape `(batch, num_additional_ids)`. These are encoded with sinusoidal
|
||
embeddings and added to the time embeddings.
|
||
return_dict (`bool`, *optional*, defaults to `True`):
|
||
Whether or not to return a [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] instead of a plain
|
||
tuple.
|
||
Returns:
|
||
[`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] or `tuple`:
|
||
If `return_dict` is True, an [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] is returned, otherwise
|
||
a `tuple` is returned where the first element is the sample tensor.
|
||
"""
|
||
# 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
|
||
batch_size, num_frames = sample.shape[:2]
|
||
timesteps = timesteps.expand(batch_size)
|
||
|
||
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=sample.dtype)
|
||
|
||
emb = self.time_embedding(t_emb)
|
||
|
||
time_embeds = self.add_time_proj(added_time_ids.flatten())
|
||
time_embeds = time_embeds.reshape((batch_size, -1))
|
||
time_embeds = time_embeds.to(emb.dtype)
|
||
aug_emb = self.add_embedding(time_embeds)
|
||
emb = emb + aug_emb
|
||
|
||
# Flatten the batch and frames dimensions
|
||
# sample: [batch, frames, channels, height, width] -> [batch * frames, channels, height, width]
|
||
sample = sample.flatten(0, 1)
|
||
# Repeat the embeddings num_video_frames times
|
||
# emb: [batch, channels] -> [batch * frames, channels]
|
||
emb = emb.repeat_interleave(num_frames, dim=0)
|
||
# encoder_hidden_states: [batch, 1, channels] -> [batch * frames, 1, channels]
|
||
encoder_hidden_states = encoder_hidden_states.repeat_interleave(num_frames, dim=0)
|
||
|
||
# 2. pre-process
|
||
sample = self.conv_in(sample)
|
||
|
||
# tokens = ID_feature["x_norm_patchtokens"]
|
||
# image_features = ID_feature["x_norm_clstoken"]
|
||
# image_features = image_features.unsqueeze(1)
|
||
# ID_feature = torch.cat([image_features,tokens],1) # 1,197,1024
|
||
# ID_feature = torch.cat([ID_feature] * 2) if do_classifier_free_guidance else ID_feature
|
||
|
||
#controlnet cond
|
||
if controlnet_cond != None:
|
||
# print(controlnet_cond.shape)
|
||
# print(ID_feature.shape)
|
||
|
||
# controlnet_cond: torch.Size([1, 14, 3, 256, 256])
|
||
# ID_feature: torch.Size([1, 14, 1024, 256, 256])
|
||
|
||
controlnet_cond = self.controlnet_cond_embedding(controlnet_cond,ID_feature)
|
||
# controlnet_cond: [28, 320, 32, 64]
|
||
# ID feature: [101,1024]
|
||
|
||
#cross attention
|
||
# b, c, h, w = controlnet_cond.shape
|
||
# controlnet_cond = rearrange(controlnet_cond, 'b c h w -> b (h w) c')
|
||
# controlnet_cond = self.CrossAttention(controlnet_cond, context=ID_feature)
|
||
# controlnet_cond = rearrange(controlnet_cond, 'b (h w) c -> b c h w', h=h, w=w)
|
||
|
||
sample = sample + controlnet_cond
|
||
# torch.Size([28, 320, 32, 64])
|
||
# ID feature [101,1024]
|
||
|
||
|
||
image_only_indicator = torch.zeros(batch_size, num_frames, dtype=sample.dtype, device=sample.device)
|
||
|
||
down_block_res_samples = (sample,)
|
||
for downsample_block in 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,
|
||
image_only_indicator=image_only_indicator,
|
||
)
|
||
else:
|
||
sample, res_samples = downsample_block(
|
||
hidden_states=sample,
|
||
temb=emb,
|
||
image_only_indicator=image_only_indicator,
|
||
)
|
||
|
||
down_block_res_samples += res_samples
|
||
|
||
# 4. mid
|
||
sample = self.mid_block(
|
||
hidden_states=sample,
|
||
temb=emb,
|
||
encoder_hidden_states=encoder_hidden_states,
|
||
image_only_indicator=image_only_indicator,
|
||
)
|
||
|
||
controlnet_down_block_res_samples = ()
|
||
|
||
for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks):
|
||
down_block_res_sample = controlnet_block(down_block_res_sample)
|
||
controlnet_down_block_res_samples = controlnet_down_block_res_samples + (down_block_res_sample,)
|
||
|
||
down_block_res_samples = controlnet_down_block_res_samples
|
||
|
||
mid_block_res_sample = self.controlnet_mid_block(sample)
|
||
|
||
# 6. scaling
|
||
down_block_res_samples = [sample * conditioning_scale for sample in down_block_res_samples]
|
||
mid_block_res_sample = mid_block_res_sample * conditioning_scale
|
||
|
||
if not return_dict:
|
||
return (down_block_res_samples, mid_block_res_sample)
|
||
|
||
return ControlNetOutput(
|
||
down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample
|
||
)
|
||
|
||
|
||
@classmethod
|
||
def from_unet(
|
||
cls,
|
||
unet: UNetSpatioTemporalConditionModel,
|
||
controlnet_conditioning_channel_order: str = "rgb",
|
||
conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
|
||
load_weights_from_unet: bool = True,
|
||
conditioning_channels: int = 3,
|
||
):
|
||
r"""
|
||
Instantiate a [`ControlNetModel`] from [`UNet2DConditionModel`].
|
||
|
||
Parameters:
|
||
unet (`UNet2DConditionModel`):
|
||
The UNet model weights to copy to the [`ControlNetModel`]. All configuration options are also copied
|
||
where applicable.
|
||
"""
|
||
|
||
transformer_layers_per_block = (
|
||
unet.config.transformer_layers_per_block if "transformer_layers_per_block" in unet.config else 1
|
||
)
|
||
encoder_hid_dim = unet.config.encoder_hid_dim if "encoder_hid_dim" in unet.config else None
|
||
encoder_hid_dim_type = unet.config.encoder_hid_dim_type if "encoder_hid_dim_type" in unet.config else None
|
||
addition_embed_type = unet.config.addition_embed_type if "addition_embed_type" in unet.config else None
|
||
addition_time_embed_dim = (
|
||
unet.config.addition_time_embed_dim if "addition_time_embed_dim" in unet.config else None
|
||
)
|
||
print(unet.config)
|
||
controlnet = cls(
|
||
in_channels=unet.config.in_channels,
|
||
down_block_types=unet.config.down_block_types,
|
||
block_out_channels=unet.config.block_out_channels,
|
||
addition_time_embed_dim=unet.config.addition_time_embed_dim,
|
||
transformer_layers_per_block=unet.config.transformer_layers_per_block,
|
||
cross_attention_dim=unet.config.cross_attention_dim,
|
||
num_attention_heads=unet.config.num_attention_heads,
|
||
num_frames=unet.config.num_frames,
|
||
sample_size=unet.config.sample_size, # Added based on the dict
|
||
layers_per_block=unet.config.layers_per_block,
|
||
projection_class_embeddings_input_dim=unet.config.projection_class_embeddings_input_dim,
|
||
conditioning_channels = conditioning_channels,
|
||
conditioning_embedding_out_channels = conditioning_embedding_out_channels,
|
||
)
|
||
#controlnet rgb channel order ignored, set to not makea difference by default
|
||
|
||
if load_weights_from_unet:
|
||
controlnet.conv_in.load_state_dict(unet.conv_in.state_dict())
|
||
controlnet.time_proj.load_state_dict(unet.time_proj.state_dict())
|
||
controlnet.time_embedding.load_state_dict(unet.time_embedding.state_dict())
|
||
|
||
# if controlnet.class_embedding:
|
||
# controlnet.class_embedding.load_state_dict(unet.class_embedding.state_dict())
|
||
|
||
controlnet.down_blocks.load_state_dict(unet.down_blocks.state_dict())
|
||
controlnet.mid_block.load_state_dict(unet.mid_block.state_dict())
|
||
|
||
return controlnet
|
||
|
||
@property
|
||
# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
|
||
def attn_processors(self) -> Dict[str, AttentionProcessor]:
|
||
r"""
|
||
Returns:
|
||
`dict` of attention processors: A dictionary containing all attention processors used in the model with
|
||
indexed by its weight name.
|
||
"""
|
||
# set recursively
|
||
processors = {}
|
||
|
||
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
|
||
if hasattr(module, "get_processor"):
|
||
processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
|
||
|
||
for sub_name, child in module.named_children():
|
||
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
|
||
|
||
return processors
|
||
|
||
for name, module in self.named_children():
|
||
fn_recursive_add_processors(name, module, processors)
|
||
|
||
return processors
|
||
|
||
# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor
|
||
def set_attn_processor(
|
||
self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False
|
||
):
|
||
r"""
|
||
Sets the attention processor to use to compute attention.
|
||
|
||
Parameters:
|
||
processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
|
||
The instantiated processor class or a dictionary of processor classes that will be set as the processor
|
||
for **all** `Attention` layers.
|
||
|
||
If `processor` is a dict, the key needs to define the path to the corresponding cross attention
|
||
processor. This is strongly recommended when setting trainable attention processors.
|
||
|
||
"""
|
||
count = len(self.attn_processors.keys())
|
||
|
||
if isinstance(processor, dict) and len(processor) != count:
|
||
raise ValueError(
|
||
f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
|
||
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
|
||
)
|
||
|
||
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
|
||
if hasattr(module, "set_processor"):
|
||
if not isinstance(processor, dict):
|
||
module.set_processor(processor, _remove_lora=_remove_lora)
|
||
else:
|
||
module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora)
|
||
|
||
for sub_name, child in module.named_children():
|
||
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
|
||
|
||
for name, module in self.named_children():
|
||
fn_recursive_attn_processor(name, module, processor)
|
||
|
||
# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
|
||
def set_default_attn_processor(self):
|
||
"""
|
||
Disables custom attention processors and sets the default attention implementation.
|
||
"""
|
||
if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
|
||
processor = AttnAddedKVProcessor()
|
||
elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
|
||
processor = AttnProcessor()
|
||
else:
|
||
raise ValueError(
|
||
f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
|
||
)
|
||
|
||
self.set_attn_processor(processor, _remove_lora=True)
|
||
|
||
# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attention_slice
|
||
def set_attention_slice(self, slice_size: Union[str, int, List[int]]) -> None:
|
||
r"""
|
||
Enable sliced attention computation.
|
||
|
||
When this option is enabled, the attention module splits the input tensor in slices to compute attention in
|
||
several steps. This is useful for saving some memory in exchange for a small decrease in speed.
|
||
|
||
Args:
|
||
slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
|
||
When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If
|
||
`"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is
|
||
provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
|
||
must be a multiple of `slice_size`.
|
||
"""
|
||
sliceable_head_dims = []
|
||
|
||
def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
|
||
if hasattr(module, "set_attention_slice"):
|
||
sliceable_head_dims.append(module.sliceable_head_dim)
|
||
|
||
for child in module.children():
|
||
fn_recursive_retrieve_sliceable_dims(child)
|
||
|
||
# retrieve number of attention layers
|
||
for module in self.children():
|
||
fn_recursive_retrieve_sliceable_dims(module)
|
||
|
||
num_sliceable_layers = len(sliceable_head_dims)
|
||
|
||
if slice_size == "auto":
|
||
# half the attention head size is usually a good trade-off between
|
||
# speed and memory
|
||
slice_size = [dim // 2 for dim in sliceable_head_dims]
|
||
elif slice_size == "max":
|
||
# make smallest slice possible
|
||
slice_size = num_sliceable_layers * [1]
|
||
|
||
slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
|
||
|
||
if len(slice_size) != len(sliceable_head_dims):
|
||
raise ValueError(
|
||
f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
|
||
f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
|
||
)
|
||
|
||
for i in range(len(slice_size)):
|
||
size = slice_size[i]
|
||
dim = sliceable_head_dims[i]
|
||
if size is not None and size > dim:
|
||
raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
|
||
|
||
# Recursively walk through all the children.
|
||
# Any children which exposes the set_attention_slice method
|
||
# gets the message
|
||
def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
|
||
if hasattr(module, "set_attention_slice"):
|
||
module.set_attention_slice(slice_size.pop())
|
||
|
||
for child in module.children():
|
||
fn_recursive_set_attention_slice(child, slice_size)
|
||
|
||
reversed_slice_size = list(reversed(slice_size))
|
||
for module in self.children():
|
||
fn_recursive_set_attention_slice(module, reversed_slice_size)
|
||
|
||
# def _set_gradient_checkpointing(self, module, value: bool = False) -> None:
|
||
# if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D)):
|
||
# module.gradient_checkpointing = value
|
||
|
||
|
||
def zero_module(module):
|
||
for p in module.parameters():
|
||
nn.init.zeros_(p)
|
||
return module
|