368 lines
16 KiB
Python
368 lines
16 KiB
Python
from diffusers import UNetSpatioTemporalConditionModel
|
|
from diffusers.models.unets.unet_spatio_temporal_condition import UNetSpatioTemporalConditionOutput
|
|
from diffusers.utils import is_torch_version
|
|
import torch
|
|
from typing import Any, Dict, Optional, Tuple, Union
|
|
|
|
def create_custom_forward(module, return_dict=None):
|
|
def custom_forward(*inputs):
|
|
if return_dict is not None:
|
|
return module(*inputs, return_dict=return_dict)
|
|
else:
|
|
return module(*inputs)
|
|
|
|
return custom_forward
|
|
CKPT_KWARGS = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
|
|
|
|
|
|
class DiffusersUNetSpatioTemporalConditionModelNormalCrafter(UNetSpatioTemporalConditionModel):
|
|
|
|
@staticmethod
|
|
def forward_crossattn_down_block_dino(
|
|
module,
|
|
hidden_states: torch.Tensor,
|
|
temb: Optional[torch.Tensor] = None,
|
|
encoder_hidden_states: Optional[torch.Tensor] = None,
|
|
image_only_indicator: Optional[torch.Tensor] = None,
|
|
dino_down_block_res_samples = None,
|
|
) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
|
|
output_states = ()
|
|
self = module
|
|
blocks = list(zip(self.resnets, self.attentions))
|
|
for resnet, attn in blocks:
|
|
if self.training and self.gradient_checkpointing: # TODO
|
|
hidden_states = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(resnet),
|
|
hidden_states,
|
|
temb,
|
|
image_only_indicator,
|
|
**CKPT_KWARGS,
|
|
)
|
|
|
|
hidden_states = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(attn),
|
|
hidden_states,
|
|
encoder_hidden_states,
|
|
image_only_indicator,
|
|
False,
|
|
**CKPT_KWARGS,
|
|
)[0]
|
|
else:
|
|
hidden_states = resnet(
|
|
hidden_states,
|
|
temb,
|
|
image_only_indicator=image_only_indicator,
|
|
)
|
|
hidden_states = attn(
|
|
hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
image_only_indicator=image_only_indicator,
|
|
return_dict=False,
|
|
)[0]
|
|
|
|
if dino_down_block_res_samples is not None:
|
|
hidden_states += dino_down_block_res_samples.pop(0)
|
|
|
|
output_states = output_states + (hidden_states,)
|
|
|
|
if self.downsamplers is not None:
|
|
for downsampler in self.downsamplers:
|
|
hidden_states = downsampler(hidden_states)
|
|
if dino_down_block_res_samples is not None:
|
|
hidden_states += dino_down_block_res_samples.pop(0)
|
|
|
|
output_states = output_states + (hidden_states,)
|
|
|
|
return hidden_states, output_states
|
|
@staticmethod
|
|
def forward_down_block_dino(
|
|
module,
|
|
hidden_states: torch.Tensor,
|
|
temb: Optional[torch.Tensor] = None,
|
|
image_only_indicator: Optional[torch.Tensor] = None,
|
|
dino_down_block_res_samples = None,
|
|
) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
|
|
self = module
|
|
output_states = ()
|
|
for resnet in self.resnets:
|
|
if self.training and self.gradient_checkpointing:
|
|
if is_torch_version(">=", "1.11.0"):
|
|
hidden_states = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(resnet),
|
|
hidden_states,
|
|
temb,
|
|
image_only_indicator,
|
|
use_reentrant=False,
|
|
)
|
|
else:
|
|
hidden_states = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(resnet),
|
|
hidden_states,
|
|
temb,
|
|
image_only_indicator,
|
|
)
|
|
else:
|
|
hidden_states = resnet(
|
|
hidden_states,
|
|
temb,
|
|
image_only_indicator=image_only_indicator,
|
|
)
|
|
if dino_down_block_res_samples is not None:
|
|
hidden_states += dino_down_block_res_samples.pop(0)
|
|
output_states = output_states + (hidden_states,)
|
|
|
|
if self.downsamplers is not None:
|
|
for downsampler in self.downsamplers:
|
|
hidden_states = downsampler(hidden_states)
|
|
if dino_down_block_res_samples is not None:
|
|
hidden_states += dino_down_block_res_samples.pop(0)
|
|
output_states = output_states + (hidden_states,)
|
|
|
|
return hidden_states, output_states
|
|
|
|
|
|
def forward(
|
|
self,
|
|
sample: torch.FloatTensor,
|
|
timestep: Union[torch.Tensor, float, int],
|
|
encoder_hidden_states: torch.Tensor,
|
|
added_time_ids: torch.Tensor,
|
|
return_dict: bool = True,
|
|
image_controlnet_down_block_res_samples = None,
|
|
image_controlnet_mid_block_res_sample = None,
|
|
dino_down_block_res_samples = None,
|
|
|
|
) -> Union[UNetSpatioTemporalConditionOutput, 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.
|
|
"""
|
|
if not hasattr(self, "custom_gradient_checkpointing"):
|
|
self.custom_gradient_checkpointing = False
|
|
|
|
# 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]
|
|
if len(timesteps.shape) == 1:
|
|
timesteps = timesteps.expand(batch_size)
|
|
else:
|
|
timesteps = timesteps.reshape(batch_size * num_frames)
|
|
t_emb = self.time_proj(timesteps) # (B, C)
|
|
|
|
# `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) # (B, C)
|
|
|
|
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)
|
|
if emb.shape[0] == 1:
|
|
emb = emb + aug_emb
|
|
# Repeat the embeddings num_video_frames times
|
|
# emb: [batch, channels] -> [batch * frames, channels]
|
|
emb = emb.repeat_interleave(num_frames, dim=0)
|
|
else:
|
|
aug_emb = aug_emb.repeat_interleave(num_frames, dim=0)
|
|
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)
|
|
|
|
# encoder_hidden_states: [batch, 1, channels] -> [batch * frames, 1, channels]
|
|
# here, our encoder_hidden_states is [batch * frames, 1, channels]
|
|
|
|
if not sample.shape[0] == encoder_hidden_states.shape[0]:
|
|
encoder_hidden_states = encoder_hidden_states.repeat_interleave(num_frames, dim=0)
|
|
# 2. pre-process
|
|
sample = self.conv_in(sample)
|
|
|
|
image_only_indicator = torch.zeros(batch_size, num_frames, dtype=sample.dtype, device=sample.device)
|
|
|
|
if dino_down_block_res_samples is not None:
|
|
dino_down_block_res_samples = [x for x in dino_down_block_res_samples]
|
|
sample += dino_down_block_res_samples.pop(0)
|
|
|
|
down_block_res_samples = (sample,)
|
|
for downsample_block in self.down_blocks:
|
|
if dino_down_block_res_samples is None:
|
|
if self.custom_gradient_checkpointing:
|
|
if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
|
|
sample, res_samples = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(downsample_block),
|
|
sample,
|
|
emb,
|
|
encoder_hidden_states,
|
|
image_only_indicator,
|
|
**CKPT_KWARGS,
|
|
)
|
|
else:
|
|
sample, res_samples = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(downsample_block),
|
|
sample,
|
|
emb,
|
|
image_only_indicator,
|
|
**CKPT_KWARGS,
|
|
)
|
|
else:
|
|
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,
|
|
)
|
|
else:
|
|
if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
|
|
sample, res_samples = self.forward_crossattn_down_block_dino(
|
|
downsample_block,
|
|
sample,
|
|
emb,
|
|
encoder_hidden_states,
|
|
image_only_indicator,
|
|
dino_down_block_res_samples,
|
|
)
|
|
else:
|
|
sample, res_samples = self.forward_down_block_dino(
|
|
downsample_block,
|
|
sample,
|
|
emb,
|
|
image_only_indicator,
|
|
dino_down_block_res_samples,
|
|
)
|
|
down_block_res_samples += res_samples
|
|
|
|
if image_controlnet_down_block_res_samples is not None:
|
|
new_down_block_res_samples = ()
|
|
|
|
for down_block_res_sample, image_controlnet_down_block_res_sample in zip(
|
|
down_block_res_samples, image_controlnet_down_block_res_samples
|
|
):
|
|
down_block_res_sample = (down_block_res_sample + image_controlnet_down_block_res_sample) / 2
|
|
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.custom_gradient_checkpointing:
|
|
sample = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(self.mid_block),
|
|
sample,
|
|
emb,
|
|
encoder_hidden_states,
|
|
image_only_indicator,
|
|
**CKPT_KWARGS,
|
|
)
|
|
else:
|
|
sample = self.mid_block(
|
|
hidden_states=sample,
|
|
temb=emb,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
image_only_indicator=image_only_indicator,
|
|
)
|
|
|
|
if image_controlnet_mid_block_res_sample is not None:
|
|
sample = (sample + image_controlnet_mid_block_res_sample) / 2
|
|
|
|
# 5. up
|
|
mid_up_block_out_samples = [sample, ]
|
|
down_block_out_sampels = []
|
|
for i, upsample_block in enumerate(self.up_blocks):
|
|
res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
|
|
down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
|
|
down_block_out_sampels.append(res_samples[-1])
|
|
if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
|
|
if self.custom_gradient_checkpointing:
|
|
sample = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(upsample_block),
|
|
sample,
|
|
res_samples,
|
|
emb,
|
|
encoder_hidden_states,
|
|
image_only_indicator,
|
|
**CKPT_KWARGS
|
|
)
|
|
else:
|
|
sample = upsample_block(
|
|
hidden_states=sample,
|
|
temb=emb,
|
|
res_hidden_states_tuple=res_samples,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
image_only_indicator=image_only_indicator,
|
|
)
|
|
else:
|
|
if self.custom_gradient_checkpointing:
|
|
sample = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(upsample_block),
|
|
sample,
|
|
res_samples,
|
|
emb,
|
|
image_only_indicator,
|
|
**CKPT_KWARGS
|
|
)
|
|
else:
|
|
sample = upsample_block(
|
|
hidden_states=sample,
|
|
temb=emb,
|
|
res_hidden_states_tuple=res_samples,
|
|
image_only_indicator=image_only_indicator,
|
|
)
|
|
mid_up_block_out_samples.append(sample)
|
|
# 6. post-process
|
|
sample = self.conv_norm_out(sample)
|
|
sample = self.conv_act(sample)
|
|
if self.custom_gradient_checkpointing:
|
|
sample = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(self.conv_out),
|
|
sample,
|
|
**CKPT_KWARGS
|
|
)
|
|
else:
|
|
sample = self.conv_out(sample)
|
|
|
|
# 7. Reshape back to original shape
|
|
sample = sample.reshape(batch_size, num_frames, *sample.shape[1:])
|
|
|
|
if not return_dict:
|
|
return (sample, down_block_out_sampels[::-1], mid_up_block_out_samples)
|
|
|
|
return UNetSpatioTemporalConditionOutput(sample=sample) |