File size: 1,023 Bytes
82ea528 |
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 |
from torch import nn
import comfy.ops
FancyVideoKeys = [
'fps_embedding.linear.bias',
'fps_embedding.linear.weight',
'motion_embedding.linear.bias',
'motion_embedding.linear.weight',
'conv_in.bias',
'conv_in.weight',
]
def initialize_weights_to_zero(m):
if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):
nn.init.constant_(m.weight, 0)
if m.bias is not None:
nn.init.constant_(m.bias, 0)
class FancyVideoCondEmbedding(nn.Module):
def __init__(self, in_channels: int, cond_embed_dim: int, act_fn: str = "silu", ops=comfy.ops.disable_weight_init):
super().__init__()
self.linear = ops.Linear(in_channels, cond_embed_dim)
self.act = None
if act_fn == "silu":
self.act = nn.SiLU()
elif act_fn == "mish":
self.act = nn.Mish()
def forward(self, sample):
sample = self.linear(sample)
if self.act is not None:
sample = self.act(sample)
return sample
|