42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
from .convs import InflatedConv3d
|
|
import torch.nn as nn
|
|
|
|
import comfy.ops
|
|
ops = comfy.ops.disable_weight_init
|
|
|
|
|
|
class Downsample3D(nn.Module):
|
|
def __init__(self,
|
|
channels,
|
|
use_conv=False,
|
|
dims=2,
|
|
out_channels=None,
|
|
padding=1,
|
|
dtype=None,
|
|
device=None,
|
|
operations=ops,
|
|
):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.out_channels = out_channels or channels
|
|
self.use_conv = use_conv
|
|
self.padding = padding
|
|
self.dims = dims
|
|
stride = 2 if dims != 3 else (1, 2, 2) # was always 2
|
|
|
|
if use_conv:
|
|
self.op = InflatedConv3d(
|
|
self.channels, self.out_channels, 3, stride=stride, padding=padding).half()
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
def forward(self, hidden_states, **kwargs):
|
|
assert hidden_states.shape[1] == self.channels
|
|
if self.use_conv and self.padding == 0:
|
|
raise NotImplementedError
|
|
|
|
assert hidden_states.shape[1] == self.channels
|
|
hidden_states = self.op(hidden_states)
|
|
|
|
return hidden_states
|