(
self,
transformer_args,
num_frames,
time_compressed_rate,
latent_width,
latent_height,
patch_size,
in_channels,
out_channels,
hidden_size,
num_layers,
num_attention_heads,
elementwise_affine,
time_embed_dim=None,
num_classes=None,
modules={},
input_time="adaln",
adm_in_channels=None,
parallel_output=True,
height_interpolation=1.0,
width_interpolation=1.0,
time_interpolation=1.0,
use_SwiGLU=False,
use_RMSNorm=False,
zero_init_y_embed=False,
allow_split_cond=True,
**kwargs,
)
| 826 | |
| 827 | class DiffusionTransformer(BaseModel): |
| 828 | def __init__( |
| 829 | self, |
| 830 | transformer_args, |
| 831 | num_frames, |
| 832 | time_compressed_rate, |
| 833 | latent_width, |
| 834 | latent_height, |
| 835 | patch_size, |
| 836 | in_channels, |
| 837 | out_channels, |
| 838 | hidden_size, |
| 839 | num_layers, |
| 840 | num_attention_heads, |
| 841 | elementwise_affine, |
| 842 | time_embed_dim=None, |
| 843 | num_classes=None, |
| 844 | modules={}, |
| 845 | input_time="adaln", |
| 846 | adm_in_channels=None, |
| 847 | parallel_output=True, |
| 848 | height_interpolation=1.0, |
| 849 | width_interpolation=1.0, |
| 850 | time_interpolation=1.0, |
| 851 | use_SwiGLU=False, |
| 852 | use_RMSNorm=False, |
| 853 | zero_init_y_embed=False, |
| 854 | |
| 855 | allow_split_cond=True, |
| 856 | |
| 857 | **kwargs, |
| 858 | ): |
| 859 | self.latent_width = latent_width |
| 860 | self.latent_height = latent_height |
| 861 | self.patch_size = patch_size |
| 862 | self.num_frames = num_frames |
| 863 | self.time_compressed_rate = time_compressed_rate |
| 864 | self.spatial_length = latent_width * latent_height // patch_size**2 |
| 865 | self.in_channels = in_channels |
| 866 | self.out_channels = out_channels |
| 867 | self.hidden_size = hidden_size |
| 868 | self.model_channels = hidden_size |
| 869 | self.time_embed_dim = time_embed_dim if time_embed_dim is not None else hidden_size |
| 870 | self.num_classes = num_classes |
| 871 | self.adm_in_channels = adm_in_channels |
| 872 | self.input_time = input_time |
| 873 | self.num_layers = num_layers |
| 874 | self.num_attention_heads = num_attention_heads |
| 875 | self.is_decoder = transformer_args.is_decoder |
| 876 | self.elementwise_affine = elementwise_affine |
| 877 | self.height_interpolation = height_interpolation |
| 878 | self.width_interpolation = width_interpolation |
| 879 | self.time_interpolation = time_interpolation |
| 880 | self.inner_hidden_size = hidden_size * 4 |
| 881 | self.zero_init_y_embed = zero_init_y_embed |
| 882 | |
| 883 | # * Custom |
| 884 | self.allow_split_cond = allow_split_cond |
| 885 |
nothing calls this directly
no test coverage detected