Spaces:
Runtime error
Runtime error
File size: 1,633 Bytes
64671d5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | """
@author: Yanzuo Lu
@author: oliveryanzuolu@gmail.com
"""
import torch.nn as nn
from diffusers.models.resnet import ResnetBlock2D, Downsample2D
class PoseEncoder(nn.Module):
def __init__(self, downscale_factor, pose_channels, in_channels, channels):
super().__init__()
self.unshuffle = nn.PixelUnshuffle(downscale_factor)
self.conv_in = nn.Conv2d(int(pose_channels * (downscale_factor ** 2)), in_channels, kernel_size=1)
resnets = []
downsamplers = []
for i in range(len(channels)):
in_channels = in_channels if i == 0 else channels[i - 1]
out_channels = channels[i]
resnets.append(ResnetBlock2D(
in_channels=in_channels,
out_channels=out_channels,
temb_channels=None, # no time embed
))
downsamplers.append(Downsample2D(
out_channels,
use_conv=False,
out_channels=out_channels,
padding=1,
name="op"
) if i != len(channels) - 1 else nn.Identity())
self.resnets = nn.ModuleList(resnets)
self.downsamplers = nn.ModuleList(downsamplers)
def forward(self, hidden_states):
features = []
hidden_states = self.unshuffle(hidden_states)
hidden_states = self.conv_in(hidden_states)
for resnet, downsampler in zip(self.resnets, self.downsamplers):
hidden_states = resnet(hidden_states, temb=None)
features.append(hidden_states)
hidden_states = downsampler(hidden_states)
return features |