(
self,
hidden_size: int,
spatial_patch_size: int,
temporal_patch_size: int,
out_channels: int,
use_adaln_lora: bool = False,
adaln_lora_dim: int = 256,
)
| 862 | """ |
| 863 | |
| 864 | def __init__( |
| 865 | self, |
| 866 | hidden_size: int, |
| 867 | spatial_patch_size: int, |
| 868 | temporal_patch_size: int, |
| 869 | out_channels: int, |
| 870 | use_adaln_lora: bool = False, |
| 871 | adaln_lora_dim: int = 256, |
| 872 | ): |
| 873 | super().__init__() |
| 874 | self.layer_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 875 | self.linear = nn.Linear( |
| 876 | hidden_size, spatial_patch_size * spatial_patch_size * temporal_patch_size * out_channels, bias=False |
| 877 | ) |
| 878 | self.hidden_size = hidden_size |
| 879 | self.n_adaln_chunks = 2 |
| 880 | self.use_adaln_lora = use_adaln_lora |
| 881 | self.adaln_lora_dim = adaln_lora_dim |
| 882 | if use_adaln_lora: |
| 883 | self.adaln_modulation = nn.Sequential( |
| 884 | nn.SiLU(), |
| 885 | nn.Linear(hidden_size, adaln_lora_dim, bias=False), |
| 886 | nn.Linear(adaln_lora_dim, self.n_adaln_chunks * hidden_size, bias=False), |
| 887 | ) |
| 888 | else: |
| 889 | self.adaln_modulation = nn.Sequential( |
| 890 | nn.SiLU(), nn.Linear(hidden_size, self.n_adaln_chunks * hidden_size, bias=False) |
| 891 | ) |
| 892 | |
| 893 | self.init_weights() |
| 894 | |
| 895 | def init_weights(self) -> None: |
| 896 | std = 1.0 / math.sqrt(self.hidden_size) |
nothing calls this directly
no test coverage detected