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)