Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,945 Bytes
a62eecb |
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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
from typing import Tuple
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as init
from diffusers.models.modeling_utils import ModelMixin
from einops import rearrange
def zero_module(module):
for p in module.parameters():
nn.init.zeros_(p)
return module
class PoseGuider(ModelMixin):
def __init__(
self,
conditioning_embedding_channels: int,
conditioning_channels: int = 3,
block_out_channels: Tuple[int] = (16, 32, 96, 256),
):
super().__init__()
self.conv_in = nn.Conv2d(
conditioning_channels, block_out_channels[0], kernel_size=3, padding=1
)
self.blocks = nn.ModuleList([])
for i in range(len(block_out_channels) - 1):
channel_in = block_out_channels[i]
channel_out = block_out_channels[i + 1]
self.blocks.append(
nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1)
)
self.blocks.append(
nn.Conv2d(
channel_in, channel_out, kernel_size=3, padding=1, stride=2
)
)
self.out = zero_module(
nn.Linear(
block_out_channels[-1]*4,
conditioning_embedding_channels,
)
)
def forward(self, conditioning):
embedding = self.conv_in(conditioning)
embedding = F.silu(embedding)
for block in self.blocks:
embedding = block(embedding)
embedding = F.silu(embedding)
embedding = rearrange(embedding, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
embedding = self.out(embedding)
return embedding
if __name__ == "__main__":
import torch
model = PoseGuider(conditioning_embedding_channels=3072, block_out_channels = (16, 32, 96, 256))
inp = torch.randn((4, 3, 1024, 768))
out = model(inp)
|