241 lines
11 KiB
Python
241 lines
11 KiB
Python
# coding: utf-8
|
|
|
|
"""
|
|
@File : models6/hm_denoising_motion.py
|
|
@Author : Songkey
|
|
@Email : songkey@pku.edu.cn
|
|
@Date : 1/3/2025
|
|
@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_adapters import InsertReferenceAdapter
|
|
|
|
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
|
|
|
|
|
class HM3DenoisingMotion(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,
|
|
motion_pad_hidden_states: Optional[dict] = None,
|
|
use_motion: bool = False,
|
|
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'down3_{idx}' in control_hidden_states:
|
|
sample += rearrange(control_hidden_states[f'down3_{idx}'], "b c f h w -> (b f) c h w")
|
|
if hasattr(self, 'motion_down') and use_motion:
|
|
sample = self.motion_down[idx](sample, motion_pad_hidden_states[f'down_{idx}'], emb, 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 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 not control_hidden_states is None and f'up3_{i}' in control_hidden_states:
|
|
sample += rearrange(control_hidden_states[f'up3_{i}'], "b c f h w -> (b f) c h w")
|
|
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](sample, reference_hidden_states[f'up_{i}'], num_frames)
|
|
if hasattr(self, 'motion_up') and use_motion:
|
|
sample = self.motion_up[i](sample, motion_pad_hidden_states[f'up_{i}'], emb, num_frames)
|
|
|
|
if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
|
|
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) |