"""A 1D upsampling layer with an optional convolution.

def __init__(

self,

channels: int,

use_conv: bool = False,

use_conv_transpose: bool = False,

out_channels: Optional[int] = None,

name: str = "conv",

):

super().__init__()

self.channels = channels

self.out_channels = out_channels or channels

self.use_conv = use_conv

self.use_conv_transpose = use_conv_transpose

self.name = name

self.conv = None

if use_conv_transpose:

self.conv = nn.ConvTranspose1d(channels, self.out_channels, 4, 2, 1)

elif use_conv:

self.conv = nn.Conv1d(self.channels, self.out_channels, 3, padding=1)

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

assert inputs.shape[1] == self.channels

if self.use_conv_transpose:

return self.conv(inputs)

outputs = F.interpolate(inputs, scale_factor=2.0, mode="nearest")

if self.use_conv:

outputs = self.conv(outputs)

"""A 1D downsampling layer with an optional convolution.

def __init__(

self,

channels: int,

use_conv: bool = False,

out_channels: Optional[int] = None,

padding: int = 1,

name: str = "conv",

):

super().__init__()

self.channels = channels

self.out_channels = out_channels or channels

self.use_conv = use_conv

self.padding = padding

stride = 2

self.name = name

if use_conv:

self.conv = nn.Conv1d(self.channels, self.out_channels, 3, stride=stride, padding=padding)

else:

assert self.channels == self.out_channels

self.conv = nn.AvgPool1d(kernel_size=stride, stride=stride)

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

assert inputs.shape[1] == self.channels

"""A 2D upsampling layer with an optional convolution.

def __init__(

self,

channels: int,

use_conv: bool = False,

use_conv_transpose: bool = False,

out_channels: Optional[int] = None,

name: str = "conv",

):

super().__init__()

self.channels = channels

self.out_channels = out_channels or channels

self.use_conv = use_conv

self.use_conv_transpose = use_conv_transpose

self.name = name

conv = None

if use_conv_transpose:

conv = nn.ConvTranspose2d(channels, self.out_channels, 4, 2, 1)

elif use_conv:

conv = LoRACompatibleConv(self.channels, self.out_channels, 3, padding=1)

# TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed

if name == "conv":

self.conv = conv

else:

self.Conv2d_0 = conv

def forward(self, hidden_states: torch.Tensor, output_size: Optional[int] = None, scale: float = 1.0):

assert hidden_states.shape[1] == self.channels

if self.use_conv_transpose:

return self.conv(hidden_states)

# Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16

# TODO(Suraj): Remove this cast once the issue is fixed in PyTorch

# https://github.com/pytorch/pytorch/issues/86679

dtype = hidden_states.dtype

if dtype == torch.bfloat16:

hidden_states = hidden_states.to(torch.float32)

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

if hidden_states.shape[0] >= 64:

hidden_states = hidden_states.contiguous()

# if `output_size` is passed we force the interpolation output

# size and do not make use of `scale_factor=2`

if output_size is None:

hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")

else:

hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")

# If the input is bfloat16, we cast back to bfloat16

if dtype == torch.bfloat16:

hidden_states = hidden_states.to(dtype)

# TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed

if self.use_conv:

if self.name == "conv":

if isinstance(self.conv, LoRACompatibleConv):

hidden_states = self.conv(hidden_states, scale)

else:

hidden_states = self.conv(hidden_states)

else:

if isinstance(self.Conv2d_0, LoRACompatibleConv):

hidden_states = self.Conv2d_0(hidden_states, scale)

else:

hidden_states = self.Conv2d_0(hidden_states)

"""A 2D downsampling layer with an optional convolution.

def __init__(

self,

channels: int,

use_conv: bool = False,

out_channels: Optional[int] = None,

padding: int = 1,

name: str = "conv",

):

super().__init__()

self.channels = channels

self.out_channels = out_channels or channels

self.use_conv = use_conv

self.padding = padding

stride = 2

self.name = name

if use_conv:

conv = LoRACompatibleConv(self.channels, self.out_channels, 3, stride=stride, padding=padding)

else:

assert self.channels == self.out_channels

conv = nn.AvgPool2d(kernel_size=stride, stride=stride)

# TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed

if name == "conv":

self.Conv2d_0 = conv

self.conv = conv

elif name == "Conv2d_0":

self.conv = conv

else:

self.conv = conv

def forward(self, hidden_states, scale: float = 1.0):

assert hidden_states.shape[1] == self.channels

if self.use_conv and self.padding == 0:

pad = (0, 1, 0, 1)

hidden_states = F.pad(hidden_states, pad, mode="constant", value=0)

assert hidden_states.shape[1] == self.channels

if isinstance(self.conv, LoRACompatibleConv):

hidden_states = self.conv(hidden_states, scale)

else:

hidden_states = self.conv(hidden_states)

"""A 2D FIR upsampling layer with an optional convolution.

def __init__(

self,

channels: int = None,

out_channels: Optional[int] = None,

use_conv: bool = False,

fir_kernel: Tuple[int, int, int, int] = (1, 3, 3, 1),

):

super().__init__()

out_channels = out_channels if out_channels else channels

if use_conv:

self.Conv2d_0 = nn.Conv2d(channels, out_channels, kernel_size=3, stride=1, padding=1)

self.use_conv = use_conv

self.fir_kernel = fir_kernel

self.out_channels = out_channels

def _upsample_2d(

self,

hidden_states: torch.Tensor,

weight: Optional[torch.Tensor] = None,

kernel: Optional[torch.FloatTensor] = None,

factor: int = 2,

gain: float = 1,

) -> torch.Tensor:

"""Fused `upsample_2d()` followed by `Conv2d()`.

assert isinstance(factor, int) and factor >= 1

# Setup filter kernel.

if kernel is None:

kernel = [1] * factor

# setup kernel

kernel = torch.tensor(kernel, dtype=torch.float32)

if kernel.ndim == 1:

kernel = torch.outer(kernel, kernel)

kernel /= torch.sum(kernel)

kernel = kernel * (gain * (factor**2))

if self.use_conv:

convH = weight.shape[2]

convW = weight.shape[3]

inC = weight.shape[1]

pad_value = (kernel.shape[0] - factor) - (convW - 1)

stride = (factor, factor)

# Determine data dimensions.

output_shape = (

(hidden_states.shape[2] - 1) * factor + convH,

(hidden_states.shape[3] - 1) * factor + convW,

)

output_padding = (

output_shape[0] - (hidden_states.shape[2] - 1) * stride[0] - convH,

output_shape[1] - (hidden_states.shape[3] - 1) * stride[1] - convW,

)

assert output_padding[0] >= 0 and output_padding[1] >= 0

num_groups = hidden_states.shape[1] // inC

# Transpose weights.

weight = torch.reshape(weight, (num_groups, -1, inC, convH, convW))

weight = torch.flip(weight, dims=[3, 4]).permute(0, 2, 1, 3, 4)

weight = torch.reshape(weight, (num_groups * inC, -1, convH, convW))

inverse_conv = F.conv_transpose2d(

hidden_states, weight, stride=stride, output_padding=output_padding, padding=0

)

output = upfirdn2d_native(

inverse_conv,

torch.tensor(kernel, device=inverse_conv.device),

pad=((pad_value + 1) // 2 + factor - 1, pad_value // 2 + 1),

)

else:

pad_value = kernel.shape[0] - factor

output = upfirdn2d_native(

hidden_states,

torch.tensor(kernel, device=hidden_states.device),

up=factor,

pad=((pad_value + 1) // 2 + factor - 1, pad_value // 2),

)

return output

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

if self.use_conv:

height = self._upsample_2d(hidden_states, self.Conv2d_0.weight, kernel=self.fir_kernel)

height = height + self.Conv2d_0.bias.reshape(1, -1, 1, 1)

else:

height = self._upsample_2d(hidden_states, kernel=self.fir_kernel, factor=2)

"""A 2D FIR downsampling layer with an optional convolution.

def __init__(

self,

channels: int = None,

out_channels: Optional[int] = None,

use_conv: bool = False,

fir_kernel: Tuple[int, int, int, int] = (1, 3, 3, 1),

):

super().__init__()

out_channels = out_channels if out_channels else channels

if use_conv:

self.Conv2d_0 = nn.Conv2d(channels, out_channels, kernel_size=3, stride=1, padding=1)

self.fir_kernel = fir_kernel

self.use_conv = use_conv

self.out_channels = out_channels

def _downsample_2d(

self,

hidden_states: torch.Tensor,

weight: Optional[torch.Tensor] = None,

kernel: Optional[torch.FloatTensor] = None,

factor: int = 2,

gain: float = 1,

) -> torch.Tensor:

"""Fused `Conv2d()` followed by `downsample_2d()`.

assert isinstance(factor, int) and factor >= 1

if kernel is None:

kernel = [1] * factor

# setup kernel

kernel = torch.tensor(kernel, dtype=torch.float32)

if kernel.ndim == 1:

kernel = torch.outer(kernel, kernel)

kernel /= torch.sum(kernel)

kernel = kernel * gain

if self.use_conv:

_, _, convH, convW = weight.shape

pad_value = (kernel.shape[0] - factor) + (convW - 1)

stride_value = [factor, factor]

upfirdn_input = upfirdn2d_native(

hidden_states,

torch.tensor(kernel, device=hidden_states.device),

pad=((pad_value + 1) // 2, pad_value // 2),

)

output = F.conv2d(upfirdn_input, weight, stride=stride_value, padding=0)

else:

pad_value = kernel.shape[0] - factor

output = upfirdn2d_native(

hidden_states,

torch.tensor(kernel, device=hidden_states.device),

down=factor,

pad=((pad_value + 1) // 2, pad_value // 2),

)

return output

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

if self.use_conv:

downsample_input = self._downsample_2d(hidden_states, weight=self.Conv2d_0.weight, kernel=self.fir_kernel)

hidden_states = downsample_input + self.Conv2d_0.bias.reshape(1, -1, 1, 1)

else:

hidden_states = self._downsample_2d(hidden_states, kernel=self.fir_kernel, factor=2)

r"""A 2D K-downsampling layer.

def __init__(self, pad_mode: str = "reflect"):

super().__init__()

self.pad_mode = pad_mode

kernel_1d = torch.tensor([[1 / 8, 3 / 8, 3 / 8, 1 / 8]])

self.pad = kernel_1d.shape[1] // 2 - 1

self.register_buffer("kernel", kernel_1d.T @ kernel_1d, persistent=False)

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

inputs = F.pad(inputs, (self.pad,) * 4, self.pad_mode)

weight = inputs.new_zeros([inputs.shape[1], inputs.shape[1], self.kernel.shape[0], self.kernel.shape[1]])

indices = torch.arange(inputs.shape[1], device=inputs.device)

kernel = self.kernel.to(weight)[None, :].expand(inputs.shape[1], -1, -1)

weight[indices, indices] = kernel

r"""A 2D K-upsampling layer.

def __init__(self, pad_mode: str = "reflect"):

super().__init__()

self.pad_mode = pad_mode

kernel_1d = torch.tensor([[1 / 8, 3 / 8, 3 / 8, 1 / 8]]) * 2

self.pad = kernel_1d.shape[1] // 2 - 1

self.register_buffer("kernel", kernel_1d.T @ kernel_1d, persistent=False)

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

inputs = F.pad(inputs, ((self.pad + 1) // 2,) * 4, self.pad_mode)

weight = inputs.new_zeros([inputs.shape[1], inputs.shape[1], self.kernel.shape[0], self.kernel.shape[1]])

indices = torch.arange(inputs.shape[1], device=inputs.device)

kernel = self.kernel.to(weight)[None, :].expand(inputs.shape[1], -1, -1)

weight[indices, indices] = kernel

r"""

def __init__(

self,

*,

in_channels: int,

out_channels: Optional[int] = None,

conv_shortcut: bool = False,

dropout: float = 0.0,

temb_channels: int = 512,

groups: int = 32,

groups_out: Optional[int] = 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, ada_group, spatial

kernel: Optional[torch.FloatTensor] = None,

output_scale_factor: float = 1.0,

use_in_shortcut: Optional[bool] = None,

up: bool = False,

down: bool = False,

conv_shortcut_bias: bool = True,

conv_2d_out_channels: Optional[int] = None,

):

super().__init__()

self.pre_norm = pre_norm

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

if self.time_embedding_norm == "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:

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

self.conv1 = LoRACompatibleConv(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 = LoRACompatibleLinear(temb_channels, out_channels)

elif self.time_embedding_norm == "scale_shift":

self.time_emb_proj = LoRACompatibleLinear(temb_channels, 2 * out_channels)

elif self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":

self.time_emb_proj = None

else:

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

else:

self.time_emb_proj = None

if self.time_embedding_norm == "ada_group":

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

elif self.time_embedding_norm == "spatial":

self.norm2 = SpatialNorm(out_channels, temb_channels)

else:

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 = LoRACompatibleConv(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 = LoRACompatibleConv(

in_channels, conv_2d_out_channels, kernel_size=1, stride=1, padding=0, bias=conv_shortcut_bias

)

def forward(self, input_tensor, temb, scale: float = 1.0):

hidden_states = input_tensor

if self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":

hidden_states = self.norm1(hidden_states, temb)

else:

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, scale=scale)

if isinstance(self.upsample, Upsample2D)

else self.upsample(input_tensor)

)

hidden_states = (

self.upsample(hidden_states, scale=scale)

if isinstance(self.upsample, Upsample2D)

else self.upsample(hidden_states)

)

elif self.downsample is not None:

input_tensor = (

self.downsample(input_tensor, scale=scale)

if isinstance(self.downsample, Downsample2D)

else self.downsample(input_tensor)

)

hidden_states = (

self.downsample(hidden_states, scale=scale)

if isinstance(self.downsample, Downsample2D)

else self.downsample(hidden_states)

)

hidden_states = self.conv1(hidden_states, scale)

if self.time_emb_proj is not None:

if not self.skip_time_act:

temb = self.nonlinearity(temb)

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

if temb is not None and self.time_embedding_norm == "default":

hidden_states = hidden_states + temb

if self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":

hidden_states = self.norm2(hidden_states, temb)

else:

hidden_states = self.norm2(hidden_states)

if temb is not None and self.time_embedding_norm == "scale_shift":

scale, shift = torch.chunk(temb, 2, dim=1)

hidden_states = hidden_states * (1 + scale) + shift

hidden_states = self.nonlinearity(hidden_states)

hidden_states = self.dropout(hidden_states)

hidden_states = self.conv2(hidden_states, scale)

if self.conv_shortcut is not None:

input_tensor = self.conv_shortcut(input_tensor, scale)

output_tensor = (input_tensor + hidden_states) / self.output_scale_factor

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:

"""

def __init__(

self, inp_channels: int, out_channels: int, kernel_size: Union[int, Tuple[int, int]], n_groups: int = 8

):

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 = nn.Mish()

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)

"""

def __init__(

self, inp_channels: int, out_channels: int, embed_dim: int, kernel_size: Union[int, Tuple[int, int]] = 5

):

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 = nn.Mish()

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:

"""

t = self.time_emb_act(t)

t = self.time_emb(t)

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

out = self.conv_out(out)

hidden_states: torch.Tensor, kernel: Optional[torch.FloatTensor] = None, factor: int = 2, gain: float = 1

) -> torch.Tensor:

r"""Upsample2D a batch of 2D images with the given filter.

assert isinstance(factor, int) and factor >= 1

if kernel is None:

kernel = [1] * factor

kernel = torch.tensor(kernel, dtype=torch.float32)

if kernel.ndim == 1:

kernel = torch.outer(kernel, kernel)

kernel /= torch.sum(kernel)

kernel = kernel * (gain * (factor**2))

pad_value = kernel.shape[0] - factor

output = upfirdn2d_native(

hidden_states,

kernel.to(device=hidden_states.device),

up=factor,

pad=((pad_value + 1) // 2 + factor - 1, pad_value // 2),

)

hidden_states: torch.Tensor, kernel: Optional[torch.FloatTensor] = None, factor: int = 2, gain: float = 1

) -> torch.Tensor:

r"""Downsample2D a batch of 2D images with the given filter.

assert isinstance(factor, int) and factor >= 1

if kernel is None:

kernel = [1] * factor

kernel = torch.tensor(kernel, dtype=torch.float32)

if kernel.ndim == 1:

kernel = torch.outer(kernel, kernel)

kernel /= torch.sum(kernel)

kernel = kernel * gain

pad_value = kernel.shape[0] - factor

output = upfirdn2d_native(

hidden_states, kernel.to(device=hidden_states.device), down=factor, pad=((pad_value + 1) // 2, pad_value // 2)

)

tensor: torch.Tensor, kernel: torch.Tensor, up: int = 1, down: int = 1, pad: Tuple[int, int] = (0, 0)

) -> torch.Tensor:

up_x = up_y = up

down_x = down_y = down

pad_x0 = pad_y0 = pad[0]

pad_x1 = pad_y1 = pad[1]

_, channel, in_h, in_w = tensor.shape

tensor = tensor.reshape(-1, in_h, in_w, 1)

_, in_h, in_w, minor = tensor.shape

kernel_h, kernel_w = kernel.shape

out = tensor.view(-1, in_h, 1, in_w, 1, minor)

out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])

out = out.view(-1, in_h * up_y, in_w * up_x, minor)

out = F.pad(out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])

out = out.to(tensor.device) # Move back to mps if necessary

out = out[

:,

max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),

max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),

:,

]

out = out.permute(0, 3, 1, 2)

out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1])

w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)

out = F.conv2d(out, w)

out = out.reshape(

-1,

minor,

in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,

in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,

)

out = out.permute(0, 2, 3, 1)

out = out[:, ::down_y, ::down_x, :]

out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1

out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1

"""

def __init__(self, in_dim: int, out_dim: Optional[int] = None, dropout: float = 0.0):

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(32, in_dim), nn.SiLU(), nn.Conv3d(in_dim, out_dim, (3, 1, 1), padding=(1, 0, 0))

)

self.conv2 = nn.Sequential(

nn.GroupNorm(32, 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(32, 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(32, 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)