Method __init__

diffsynth/models/svd_unet.py:239–313

(self, add_positional_conv=None)

class SVDUNet(torch.nn.Module):
    def __init__(self, add_positional_conv=None):
        super().__init__()
        self.time_proj = Timesteps(320)
        self.time_embedding = torch.nn.Sequential(
            torch.nn.Linear(320, 1280),
            torch.nn.SiLU(),
            torch.nn.Linear(1280, 1280)
        )
        self.add_time_proj = Timesteps(256)
        self.add_time_embedding = torch.nn.Sequential(
            torch.nn.Linear(768, 1280),
            torch.nn.SiLU(),
            torch.nn.Linear(1280, 1280)
        )
        self.conv_in = torch.nn.Conv2d(8, 320, kernel_size=3, padding=1)

        self.blocks = torch.nn.ModuleList([
            # CrossAttnDownBlockSpatioTemporal
            ResnetBlock(320, 320, 1280, eps=1e-6),                      PushBlock(), TemporalResnetBlock(320, 320, 1280, eps=1e-6),   PopMixBlock(),  PushBlock(),
            AttentionBlock(5, 64, 320, 1, 1024, need_proj_out=False),   PushBlock(), TemporalAttentionBlock(5, 64, 320, 1024, add_positional_conv),   PopMixBlock(320),  PushBlock(),
            ResnetBlock(320, 320, 1280, eps=1e-6),                      PushBlock(), TemporalResnetBlock(320, 320, 1280, eps=1e-6),   PopMixBlock(),  PushBlock(),
            AttentionBlock(5, 64, 320, 1, 1024, need_proj_out=False),   PushBlock(), TemporalAttentionBlock(5, 64, 320, 1024, add_positional_conv),   PopMixBlock(320),  PushBlock(),
            DownSampler(320), PushBlock(),
            # CrossAttnDownBlockSpatioTemporal
            ResnetBlock(320, 640, 1280, eps=1e-6),                      PushBlock(), TemporalResnetBlock(640, 640, 1280, eps=1e-6),   PopMixBlock(),  PushBlock(),
            AttentionBlock(10, 64, 640, 1, 1024, need_proj_out=False),  PushBlock(), TemporalAttentionBlock(10, 64, 640, 1024, add_positional_conv),  PopMixBlock(640),  PushBlock(),
            ResnetBlock(640, 640, 1280, eps=1e-6),                      PushBlock(), TemporalResnetBlock(640, 640, 1280, eps=1e-6),   PopMixBlock(),  PushBlock(),
            AttentionBlock(10, 64, 640, 1, 1024, need_proj_out=False),  PushBlock(), TemporalAttentionBlock(10, 64, 640, 1024, add_positional_conv),  PopMixBlock(640),  PushBlock(),
            DownSampler(640), PushBlock(),
            # CrossAttnDownBlockSpatioTemporal
            ResnetBlock(640, 1280, 1280, eps=1e-6),                     PushBlock(), TemporalResnetBlock(1280, 1280, 1280, eps=1e-6), PopMixBlock(),  PushBlock(),
            AttentionBlock(20, 64, 1280, 1, 1024, need_proj_out=False), PushBlock(), TemporalAttentionBlock(20, 64, 1280, 1024, add_positional_conv), PopMixBlock(1280), PushBlock(),
            ResnetBlock(1280, 1280, 1280, eps=1e-6),                    PushBlock(), TemporalResnetBlock(1280, 1280, 1280, eps=1e-6), PopMixBlock(),  PushBlock(),
            AttentionBlock(20, 64, 1280, 1, 1024, need_proj_out=False), PushBlock(), TemporalAttentionBlock(20, 64, 1280, 1024, add_positional_conv), PopMixBlock(1280), PushBlock(),
            DownSampler(1280), PushBlock(),
            # DownBlockSpatioTemporal
            ResnetBlock(1280, 1280, 1280, eps=1e-5),                    PushBlock(), TemporalResnetBlock(1280, 1280, 1280, eps=1e-5), PopMixBlock(),  PushBlock(),
            ResnetBlock(1280, 1280, 1280, eps=1e-5),                    PushBlock(), TemporalResnetBlock(1280, 1280, 1280, eps=1e-5), PopMixBlock(),  PushBlock(),
            # UNetMidBlockSpatioTemporal
            ResnetBlock(1280, 1280, 1280, eps=1e-5),                    PushBlock(), TemporalResnetBlock(1280, 1280, 1280, eps=1e-5), PopMixBlock(),  PushBlock(),
            AttentionBlock(20, 64, 1280, 1, 1024, need_proj_out=False), PushBlock(), TemporalAttentionBlock(20, 64, 1280, 1024, add_positional_conv), PopMixBlock(1280),
            ResnetBlock(1280, 1280, 1280, eps=1e-5),                    PushBlock(), TemporalResnetBlock(1280, 1280, 1280, eps=1e-5), PopMixBlock(),
            # UpBlockSpatioTemporal
            PopBlock(), ResnetBlock(2560, 1280, 1280, eps=1e-6),        PushBlock(), TemporalResnetBlock(1280, 1280, 1280, eps=1e-5), PopMixBlock(),
            PopBlock(), ResnetBlock(2560, 1280, 1280, eps=1e-6),        PushBlock(), TemporalResnetBlock(1280, 1280, 1280, eps=1e-5), PopMixBlock(),
            PopBlock(), ResnetBlock(2560, 1280, 1280, eps=1e-6),        PushBlock(), TemporalResnetBlock(1280, 1280, 1280, eps=1e-5), PopMixBlock(),
            UpSampler(1280),
            # CrossAttnUpBlockSpatioTemporal
            PopBlock(),        ResnetBlock(2560, 1280, 1280, eps=1e-6), PushBlock(), TemporalResnetBlock(1280, 1280, 1280, eps=1e-6), PopMixBlock(),  PushBlock(),
            AttentionBlock(20, 64, 1280, 1, 1024, need_proj_out=False), PushBlock(), TemporalAttentionBlock(20, 64, 1280, 1024, add_positional_conv), PopMixBlock(1280),
            PopBlock(),        ResnetBlock(2560, 1280, 1280, eps=1e-6), PushBlock(), TemporalResnetBlock(1280, 1280, 1280, eps=1e-6), PopMixBlock(),  PushBlock(),
            AttentionBlock(20, 64, 1280, 1, 1024, need_proj_out=False), PushBlock(), TemporalAttentionBlock(20, 64, 1280, 1024, add_positional_conv), PopMixBlock(1280),
            PopBlock(),        ResnetBlock(1920, 1280, 1280, eps=1e-6), PushBlock(), TemporalResnetBlock(1280, 1280, 1280, eps=1e-6), PopMixBlock(),  PushBlock(),
            AttentionBlock(20, 64, 1280, 1, 1024, need_proj_out=False), PushBlock(), TemporalAttentionBlock(20, 64, 1280, 1024, add_positional_conv), PopMixBlock(1280),
            UpSampler(1280),
            # CrossAttnUpBlockSpatioTemporal
            PopBlock(),        ResnetBlock(1920, 640, 1280, eps=1e-6),  PushBlock(), TemporalResnetBlock(640, 640, 1280, eps=1e-6),   PopMixBlock(),  PushBlock(),
            AttentionBlock(10, 64, 640, 1, 1024, need_proj_out=False),  PushBlock(), TemporalAttentionBlock(10, 64, 640, 1024, add_positional_conv),  PopMixBlock(640),