41 lines
1023 B
Python
41 lines
1023 B
Python
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
|