# Copyright 2026 The HuggingFace Team. All rights reserved.

# `TemporalConvLayer` Copyright 2025 Alibaba DAMO-VILAB, The ModelScope Team and 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 functools import partial

import torch

import torch.nn as nn

import torch.nn.functional as F

from ..utils import deprecate

from .activations import get_activation

from .attention_processor import SpatialNorm

from .downsampling import ( # noqa

Downsample1D,

Downsample2D,

FirDownsample2D,

KDownsample2D,

downsample_2d,

)

from .normalization import AdaGroupNorm

from .upsampling import ( # noqa

FirUpsample2D,

KUpsample2D,

Upsample1D,

Upsample2D,

upfirdn2d_native,

upsample_2d,

)

class ResnetBlockCondNorm2D(nn.Module):

r"""

A Resnet block that use normalization layer that incorporate conditioning information.

Parameters:

in_channels (`int`): The number of channels in the input.

out_channels (`int`, *optional*, default to be `None`):

The number of output channels for the first conv2d layer. If None, same as `in_channels`.

dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.

temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding.

groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer.

groups_out (`int`, *optional*, default to None):

The number of groups to use for the second normalization layer. if set to None, same as `groups`.

eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization.

non_linearity (`str`, *optional*, default to `"swish"`): the activation function to use.

time_embedding_norm (`str`, *optional*, default to `"ada_group"` ):

The normalization layer for time embedding `temb`. Currently only support "ada_group" or "spatial".

kernel (`torch.Tensor`, optional, default to None): FIR filter, see

[`~models.resnet.FirUpsample2D`] and [`~models.resnet.FirDownsample2D`].

output_scale_factor (`float`, *optional*, default to be `1.0`): the scale factor to use for the output.

use_in_shortcut (`bool`, *optional*, default to `True`):

If `True`, add a 1x1 nn.conv2d layer for skip-connection.

up (`bool`, *optional*, default to `False`): If `True`, add an upsample layer.

down (`bool`, *optional*, default to `False`): If `True`, add a downsample layer.

conv_shortcut_bias (`bool`, *optional*, default to `True`): If `True`, adds a learnable bias to the

`conv_shortcut` output.

conv_2d_out_channels (`int`, *optional*, default to `None`): the number of channels in the output.

If None, same as `out_channels`.

"""

def __init__(

self,

*,

in_channels: int,

out_channels: int | None = None,

conv_shortcut: bool = False,

dropout: float = 0.0,

temb_channels: int = 512,

groups: int = 32,

groups_out: int | None = None,

eps: float = 1e-6,

non_linearity: str = "swish",

time_embedding_norm: str = "ada_group", # ada_group, spatial

output_scale_factor: float = 1.0,

use_in_shortcut: bool | None = None,

up: bool = False,

down: bool = False,

conv_shortcut_bias: bool = True,

conv_2d_out_channels: int | None = None,

):

super().__init__()

self.in_channels = in_channels

out_channels = in_channels if out_channels is None else out_channels

self.out_channels = out_channels

self.use_conv_shortcut = conv_shortcut

self.up = up

self.down = down

self.output_scale_factor = output_scale_factor

self.time_embedding_norm = time_embedding_norm

if groups_out is None:

groups_out = groups

if self.time_embedding_norm == "ada_group": # ada_group

self.norm1 = AdaGroupNorm(temb_channels, in_channels, groups, eps=eps)

elif self.time_embedding_norm == "spatial":

self.norm1 = SpatialNorm(in_channels, temb_channels)

else:

raise ValueError(f" unsupported time_embedding_norm: {self.time_embedding_norm}")

self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)

if self.time_embedding_norm == "ada_group": # ada_group

self.norm2 = AdaGroupNorm(temb_channels, out_channels, groups_out, eps=eps)

elif self.time_embedding_norm == "spatial": # spatial

self.norm2 = SpatialNorm(out_channels, temb_channels)

else:

raise ValueError(f" unsupported time_embedding_norm: {self.time_embedding_norm}")

self.dropout = torch.nn.Dropout(dropout)

conv_2d_out_channels = conv_2d_out_channels or out_channels

self.conv2 = nn.Conv2d(out_channels, conv_2d_out_channels, kernel_size=3, stride=1, padding=1)

self.nonlinearity = get_activation(non_linearity)

self.upsample = self.downsample = None

if self.up:

self.upsample = Upsample2D(in_channels, use_conv=False)

elif self.down:

self.downsample = Downsample2D(in_channels, use_conv=False, padding=1, name="op")

self.use_in_shortcut = self.in_channels != conv_2d_out_channels if use_in_shortcut is None else use_in_shortcut

self.conv_shortcut = None

if self.use_in_shortcut:

self.conv_shortcut = nn.Conv2d(

in_channels,

conv_2d_out_channels,

kernel_size=1,

stride=1,

padding=0,

bias=conv_shortcut_bias,

)

def forward(self, input_tensor: torch.Tensor, temb: torch.Tensor, *args, **kwargs) -> torch.Tensor:

if len(args) > 0 or kwargs.get("scale", None) is not None:

deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."

deprecate("scale", "1.0.0", deprecation_message)

hidden_states = input_tensor

hidden_states = self.norm1(hidden_states, temb)

hidden_states = self.nonlinearity(hidden_states)

if self.upsample is not None:

# upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984

if hidden_states.shape[0] >= 64:

input_tensor = input_tensor.contiguous()

hidden_states = hidden_states.contiguous()

input_tensor = self.upsample(input_tensor)

hidden_states = self.upsample(hidden_states)

elif self.downsample is not None:

input_tensor = self.downsample(input_tensor)

hidden_states = self.downsample(hidden_states)

hidden_states = self.conv1(hidden_states)

hidden_states = self.norm2(hidden_states, temb)

hidden_states = self.nonlinearity(hidden_states)

hidden_states = self.dropout(hidden_states)

hidden_states = self.conv2(hidden_states)

if self.conv_shortcut is not None:

input_tensor = self.conv_shortcut(input_tensor)

output_tensor = (input_tensor + hidden_states) / self.output_scale_factor

return output_tensor

class ResnetBlock2D(nn.Module):

r"""

A Resnet block.

Parameters:

in_channels (`int`): The number of channels in the input.

out_channels (`int`, *optional*, default to be `None`):

The number of output channels for the first conv2d layer. If None, same as `in_channels`.

dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.

temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding.

groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer.

groups_out (`int`, *optional*, default to None):

The number of groups to use for the second normalization layer. if set to None, same as `groups`.

eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization.

non_linearity (`str`, *optional*, default to `"swish"`): the activation function to use.

time_embedding_norm (`str`, *optional*, default to `"default"` ): Time scale shift config.

By default, apply timestep embedding conditioning with a simple shift mechanism. Choose "scale_shift" for a

stronger conditioning with scale and shift.

kernel (`torch.Tensor`, optional, default to None): FIR filter, see

[`~models.resnet.FirUpsample2D`] and [`~models.resnet.FirDownsample2D`].

output_scale_factor (`float`, *optional*, default to be `1.0`): the scale factor to use for the output.

use_in_shortcut (`bool`, *optional*, default to `True`):

If `True`, add a 1x1 nn.conv2d layer for skip-connection.

up (`bool`, *optional*, default to `False`): If `True`, add an upsample layer.

down (`bool`, *optional*, default to `False`): If `True`, add a downsample layer.

conv_shortcut_bias (`bool`, *optional*, default to `True`): If `True`, adds a learnable bias to the

`conv_shortcut` output.

conv_2d_out_channels (`int`, *optional*, default to `None`): the number of channels in the output.

If None, same as `out_channels`.

"""

def __init__(

self,

*,

in_channels: int,

out_channels: int | None = None,

conv_shortcut: bool = False,

dropout: float = 0.0,

temb_channels: int = 512,

groups: int = 32,

groups_out: int | None = None,

pre_norm: bool = True,

eps: float = 1e-6,

non_linearity: str = "swish",

skip_time_act: bool = False,

time_embedding_norm: str = "default", # default, scale_shift,

kernel: torch.Tensor | None = None,

output_scale_factor: float = 1.0,

use_in_shortcut: bool | None = None,

up: bool = False,

down: bool = False,

conv_shortcut_bias: bool = True,

conv_2d_out_channels: int | None = None,

):

super().__init__()

if time_embedding_norm == "ada_group":

raise ValueError(

"This class cannot be used with `time_embedding_norm==ada_group`, please use `ResnetBlockCondNorm2D` instead",

)

if time_embedding_norm == "spatial":

raise ValueError(

"This class cannot be used with `time_embedding_norm==spatial`, please use `ResnetBlockCondNorm2D` instead",

)

self.pre_norm = True

self.in_channels = in_channels

out_channels = in_channels if out_channels is None else out_channels

self.out_channels = out_channels

self.use_conv_shortcut = conv_shortcut

self.up = up

self.down = down

self.output_scale_factor = output_scale_factor

self.time_embedding_norm = time_embedding_norm

self.skip_time_act = skip_time_act

if groups_out is None:

groups_out = groups

self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)

self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)

if temb_channels is not None:

if self.time_embedding_norm == "default":

self.time_emb_proj = nn.Linear(temb_channels, out_channels)

elif self.time_embedding_norm == "scale_shift":

self.time_emb_proj = nn.Linear(temb_channels, 2 * out_channels)

else:

raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ")

else:

self.time_emb_proj = None

self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)

self.dropout = torch.nn.Dropout(dropout)

conv_2d_out_channels = conv_2d_out_channels or out_channels

self.conv2 = nn.Conv2d(out_channels, conv_2d_out_channels, kernel_size=3, stride=1, padding=1)

self.nonlinearity = get_activation(non_linearity)

self.upsample = self.downsample = None

if self.up:

if kernel == "fir":

fir_kernel = (1, 3, 3, 1)

self.upsample = lambda x: upsample_2d(x, kernel=fir_kernel)

elif kernel == "sde_vp":

self.upsample = partial(F.interpolate, scale_factor=2.0, mode="nearest")

else:

self.upsample = Upsample2D(in_channels, use_conv=False)

elif self.down:

if kernel == "fir":

fir_kernel = (1, 3, 3, 1)

self.downsample = lambda x: downsample_2d(x, kernel=fir_kernel)

elif kernel == "sde_vp":

self.downsample = partial(F.avg_pool2d, kernel_size=2, stride=2)

else:

self.downsample = Downsample2D(in_channels, use_conv=False, padding=1, name="op")

self.use_in_shortcut = self.in_channels != conv_2d_out_channels if use_in_shortcut is None else use_in_shortcut

self.conv_shortcut = None

if self.use_in_shortcut:

self.conv_shortcut = nn.Conv2d(

in_channels,

conv_2d_out_channels,

kernel_size=1,

stride=1,

padding=0,

bias=conv_shortcut_bias,

)

def forward(self, input_tensor: torch.Tensor, temb: torch.Tensor, *args, **kwargs) -> torch.Tensor:

if len(args) > 0 or kwargs.get("scale", None) is not None:

deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."

deprecate("scale", "1.0.0", deprecation_message)

hidden_states = input_tensor

hidden_states = self.norm1(hidden_states)

hidden_states = self.nonlinearity(hidden_states)

if self.upsample is not None:

# upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984

if hidden_states.shape[0] >= 64:

input_tensor = input_tensor.contiguous()

hidden_states = hidden_states.contiguous()

input_tensor = self.upsample(input_tensor)

hidden_states = self.upsample(hidden_states)

elif self.downsample is not None:

input_tensor = self.downsample(input_tensor)

hidden_states = self.downsample(hidden_states)

hidden_states = self.conv1(hidden_states)

if self.time_emb_proj is not None:

if not self.skip_time_act:

temb = self.nonlinearity(temb)

temb = self.time_emb_proj(temb)[:, :, None, None]

if self.time_embedding_norm == "default":

if temb is not None:

hidden_states = hidden_states + temb

hidden_states = self.norm2(hidden_states)

elif self.time_embedding_norm == "scale_shift":

if temb is None:

raise ValueError(

f" `temb` should not be None when `time_embedding_norm` is {self.time_embedding_norm}"

)

time_scale, time_shift = torch.chunk(temb, 2, dim=1)

hidden_states = self.norm2(hidden_states)

hidden_states = hidden_states * (1 + time_scale) + time_shift

else:

hidden_states = self.norm2(hidden_states)

hidden_states = self.nonlinearity(hidden_states)

hidden_states = self.dropout(hidden_states)

hidden_states = self.conv2(hidden_states)

if self.conv_shortcut is not None:

# Only use contiguous() during training to avoid DDP gradient stride mismatch warning.

# In inference mode (eval or no_grad), skip contiguous() for better performance, especially on CPU.

# Issue: https://github.com/huggingface/diffusers/issues/12975

if self.training:

input_tensor = input_tensor.contiguous()

input_tensor = self.conv_shortcut(input_tensor)

output_tensor = (input_tensor + hidden_states) / self.output_scale_factor

return output_tensor

# unet_rl.py

def rearrange_dims(tensor: torch.Tensor) -> torch.Tensor:

if len(tensor.shape) == 2:

return tensor[:, :, None]

if len(tensor.shape) == 3:

return tensor[:, :, None, :]

elif len(tensor.shape) == 4:

return tensor[:, :, 0, :]

else:

raise ValueError(f"`len(tensor)`: {len(tensor)} has to be 2, 3 or 4.")

class Conv1dBlock(nn.Module):

"""

Conv1d --> GroupNorm --> Mish

Parameters:

inp_channels (`int`): Number of input channels.

out_channels (`int`): Number of output channels.

kernel_size (`int` or `tuple`): Size of the convolving kernel.

n_groups (`int`, default `8`): Number of groups to separate the channels into.

activation (`str`, defaults to `mish`): Name of the activation function.

"""

def __init__(

self,

inp_channels: int,

out_channels: int,

kernel_size: int | tuple[int, int],

n_groups: int = 8,

activation: str = "mish",

):

super().__init__()

self.conv1d = nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2)

self.group_norm = nn.GroupNorm(n_groups, out_channels)

self.mish = get_activation(activation)

def forward(self, inputs: torch.Tensor) -> torch.Tensor:

intermediate_repr = self.conv1d(inputs)

intermediate_repr = rearrange_dims(intermediate_repr)

intermediate_repr = self.group_norm(intermediate_repr)

intermediate_repr = rearrange_dims(intermediate_repr)

output = self.mish(intermediate_repr)

return output

# unet_rl.py

class ResidualTemporalBlock1D(nn.Module):

"""

Residual 1D block with temporal convolutions.

Parameters:

inp_channels (`int`): Number of input channels.

out_channels (`int`): Number of output channels.

embed_dim (`int`): Embedding dimension.

kernel_size (`int` or `tuple`): Size of the convolving kernel.

activation (`str`, defaults `mish`): It is possible to choose the right activation function.

"""

def __init__(

self,

inp_channels: int,

out_channels: int,

embed_dim: int,

kernel_size: int | tuple[int, int] = 5,

activation: str = "mish",

):

super().__init__()

self.conv_in = Conv1dBlock(inp_channels, out_channels, kernel_size)

self.conv_out = Conv1dBlock(out_channels, out_channels, kernel_size)

self.time_emb_act = get_activation(activation)

self.time_emb = nn.Linear(embed_dim, out_channels)

self.residual_conv = (

nn.Conv1d(inp_channels, out_channels, 1) if inp_channels != out_channels else nn.Identity()

)

def forward(self, inputs: torch.Tensor, t: torch.Tensor) -> torch.Tensor:

"""

Args:

inputs : [ batch_size x inp_channels x horizon ]

t : [ batch_size x embed_dim ]

returns:

out : [ batch_size x out_channels x horizon ]

"""

t = self.time_emb_act(t)

t = self.time_emb(t)

out = self.conv_in(inputs) + rearrange_dims(t)

out = self.conv_out(out)

return out + self.residual_conv(inputs)

class TemporalConvLayer(nn.Module):

"""

Temporal convolutional layer that can be used for video (sequence of images) input Code mostly copied from:

https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/models/multi_modal/video_synthesis/unet_sd.py#L1016

Parameters:

in_dim (`int`): Number of input channels.

out_dim (`int`): Number of output channels.

dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.

"""

def __init__(

self,

in_dim: int,

out_dim: int | None = None,

dropout: float = 0.0,

norm_num_groups: int = 32,

):

super().__init__()

out_dim = out_dim or in_dim

self.in_dim = in_dim

self.out_dim = out_dim

# conv layers

self.conv1 = nn.Sequential(

nn.GroupNorm(norm_num_groups, in_dim),

nn.SiLU(),

nn.Conv3d(in_dim, out_dim, (3, 1, 1), padding=(1, 0, 0)),

)

self.conv2 = nn.Sequential(

nn.GroupNorm(norm_num_groups, out_dim),

nn.SiLU(),

nn.Dropout(dropout),

nn.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0)),

)

self.conv3 = nn.Sequential(

nn.GroupNorm(norm_num_groups, out_dim),

nn.SiLU(),

nn.Dropout(dropout),

nn.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0)),

)

self.conv4 = nn.Sequential(

nn.GroupNorm(norm_num_groups, out_dim),

nn.SiLU(),

nn.Dropout(dropout),

nn.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0)),

)

# zero out the last layer params,so the conv block is identity

nn.init.zeros_(self.conv4[-1].weight)

nn.init.zeros_(self.conv4[-1].bias)

def forward(self, hidden_states: torch.Tensor, num_frames: int = 1) -> torch.Tensor:

hidden_states = (

hidden_states[None, :].reshape((-1, num_frames) + hidden_states.shape[1:]).permute(0, 2, 1, 3, 4)

)

identity = hidden_states

hidden_states = self.conv1(hidden_states)

hidden_states = self.conv2(hidden_states)

hidden_states = self.conv3(hidden_states)

hidden_states = self.conv4(hidden_states)

hidden_states = identity + hidden_states

hidden_states = hidden_states.permute(0, 2, 1, 3, 4).reshape(

(hidden_states.shape[0] * hidden_states.shape[2], -1) + hidden_states.shape[3:]

)

return hidden_states

class TemporalResnetBlock(nn.Module):

r"""

A Resnet block.

Parameters:

in_channels (`int`): The number of channels in the input.

out_channels (`int`, *optional*, default to be `None`):

The number of output channels for the first conv2d layer. If None, same as `in_channels`.

temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding.

eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization.

"""

def __init__(

self,

in_channels: int,

out_channels: int | None = None,

temb_channels: int = 512,

eps: float = 1e-6,

):

super().__init__()

self.in_channels = in_channels

out_channels = in_channels if out_channels is None else out_channels

self.out_channels = out_channels

kernel_size = (3, 1, 1)

padding = [k // 2 for k in kernel_size]

self.norm1 = torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=eps, affine=True)

self.conv1 = nn.Conv3d(

in_channels,

out_channels,

kernel_size=kernel_size,

stride=1,

padding=padding,

)

if temb_channels is not None:

self.time_emb_proj = nn.Linear(temb_channels, out_channels)

else:

self.time_emb_proj = None

self.norm2 = torch.nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=eps, affine=True)

self.dropout = torch.nn.Dropout(0.0)

self.conv2 = nn.Conv3d(

out_channels,

out_channels,

kernel_size=kernel_size,

stride=1,

padding=padding,

)

self.nonlinearity = get_activation("silu")

self.use_in_shortcut = self.in_channels != out_channels

self.conv_shortcut = None

if self.use_in_shortcut:

self.conv_shortcut = nn.Conv3d(

in_channels,

out_channels,

kernel_size=1,

stride=1,

padding=0,

)

def forward(self, input_tensor: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:

hidden_states = input_tensor

hidden_states = self.norm1(hidden_states)

hidden_states = self.nonlinearity(hidden_states)

hidden_states = self.conv1(hidden_states)

if self.time_emb_proj is not None:

temb = self.nonlinearity(temb)

temb = self.time_emb_proj(temb)[:, :, :, None, None]

temb = temb.permute(0, 2, 1, 3, 4)

hidden_states = hidden_states + temb

hidden_states = self.norm2(hidden_states)

hidden_states = self.nonlinearity(hidden_states)

hidden_states = self.dropout(hidden_states)

hidden_states = self.conv2(hidden_states)

if self.conv_shortcut is not None:

input_tensor = self.conv_shortcut(input_tensor)

output_tensor = input_tensor + hidden_states

return output_tensor

# VideoResBlock

class SpatioTemporalResBlock(nn.Module):

r"""

A SpatioTemporal Resnet block.

Parameters:

in_channels (`int`): The number of channels in the input.

out_channels (`int`, *optional*, default to be `None`):

The number of output channels for the first conv2d layer. If None, same as `in_channels`.

temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding.

eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the spatial resenet.

temporal_eps (`float`, *optional*, defaults to `eps`): The epsilon to use for the temporal resnet.

merge_factor (`float`, *optional*, defaults to `0.5`): The merge factor to use for the temporal mixing.

merge_strategy (`str`, *optional*, defaults to `learned_with_images`):

The merge strategy to use for the temporal mixing.

switch_spatial_to_temporal_mix (`bool`, *optional*, defaults to `False`):

If `True`, switch the spatial and temporal mixing.

"""

def __init__(

self,

in_channels: int,

out_channels: int | None = None,

temb_channels: int = 512,

eps: float = 1e-6,

temporal_eps: float | None = None,

merge_factor: float = 0.5,

merge_strategy="learned_with_images",

switch_spatial_to_temporal_mix: bool = False,

):

super().__init__()

self.spatial_res_block = ResnetBlock2D(

in_channels=in_channels,

out_channels=out_channels,

temb_channels=temb_channels,

eps=eps,

)

self.temporal_res_block = TemporalResnetBlock(

in_channels=out_channels if out_channels is not None else in_channels,

out_channels=out_channels if out_channels is not None else in_channels,

temb_channels=temb_channels,

eps=temporal_eps if temporal_eps is not None else eps,

)

self.time_mixer = AlphaBlender(

alpha=merge_factor,

merge_strategy=merge_strategy,

switch_spatial_to_temporal_mix=switch_spatial_to_temporal_mix,

)

def forward(

self,

hidden_states: torch.Tensor,

temb: torch.Tensor | None = None,

image_only_indicator: torch.Tensor | None = None,

):

num_frames = image_only_indicator.shape[-1]

hidden_states = self.spatial_res_block(hidden_states, temb)

batch_frames, channels, height, width = hidden_states.shape

batch_size = batch_frames // num_frames

hidden_states_mix = (

hidden_states[None, :].reshape(batch_size, num_frames, channels, height, width).permute(0, 2, 1, 3, 4)

)

hidden_states = (

hidden_states[None, :].reshape(batch_size, num_frames, channels, height, width).permute(0, 2, 1, 3, 4)

)

if temb is not None:

temb = temb.reshape(batch_size, num_frames, -1)

hidden_states = self.temporal_res_block(hidden_states, temb)

hidden_states = self.time_mixer(

x_spatial=hidden_states_mix,

x_temporal=hidden_states,

image_only_indicator=image_only_indicator,

)

hidden_states = hidden_states.permute(0, 2, 1, 3, 4).reshape(batch_frames, channels, height, width)

return hidden_states

class AlphaBlender(nn.Module):

r"""

A module to blend spatial and temporal features.

Parameters:

alpha (`float`): The initial value of the blending factor.

merge_strategy (`str`, *optional*, defaults to `learned_with_images`):

The merge strategy to use for the temporal mixing.

switch_spatial_to_temporal_mix (`bool`, *optional*, defaults to `False`):

If `True`, switch the spatial and temporal mixing.

"""

strategies = ["learned", "fixed", "learned_with_images"]

def __init__(

self,

alpha: float,

merge_strategy: str = "learned_with_images",

switch_spatial_to_temporal_mix: bool = False,

):

super().__init__()

self.merge_strategy = merge_strategy

self.switch_spatial_to_temporal_mix = switch_spatial_to_temporal_mix # For TemporalVAE

if merge_strategy not in self.strategies:

raise ValueError(f"merge_strategy needs to be in {self.strategies}")

if self.merge_strategy == "fixed":

self.register_buffer("mix_factor", torch.Tensor([alpha]))

elif self.merge_strategy == "learned" or self.merge_strategy == "learned_with_images":

self.register_parameter("mix_factor", torch.nn.Parameter(torch.Tensor([alpha])))

else:

raise ValueError(f"Unknown merge strategy {self.merge_strategy}")

def get_alpha(self, image_only_indicator: torch.Tensor, ndims: int) -> torch.Tensor:

if self.merge_strategy == "fixed":

alpha = self.mix_factor

elif self.merge_strategy == "learned":

alpha = torch.sigmoid(self.mix_factor)

elif self.merge_strategy == "learned_with_images":

if image_only_indicator is None:

raise ValueError("Please provide image_only_indicator to use learned_with_images merge strategy")

alpha = torch.where(

image_only_indicator.bool(),

torch.ones(1, 1, device=image_only_indicator.device),

torch.sigmoid(self.mix_factor)[..., None],

)

# (batch, channel, frames, height, width)

if ndims == 5:

alpha = alpha[:, None, :, None, None]

# (batch*frames, height*width, channels)

elif ndims == 3:

alpha = alpha.reshape(-1)[:, None, None]

else:

raise ValueError(f"Unexpected ndims {ndims}. Dimensions should be 3 or 5")

else:

raise NotImplementedError

return alpha

def forward(

self,

x_spatial: torch.Tensor,

x_temporal: torch.Tensor,

image_only_indicator: torch.Tensor | None = None,

) -> torch.Tensor:

alpha = self.get_alpha(image_only_indicator, x_spatial.ndim)

alpha = alpha.to(x_spatial.dtype)

if self.switch_spatial_to_temporal_mix:

alpha = 1.0 - alpha

x = alpha * x_spatial + (1.0 - alpha) * x_temporal

return x