From 13c0091109506b375a742abccd993be87cf60c90 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=AE=A3=E6=BA=90?= Date: Thu, 12 Sep 2024 15:49:05 +0800 Subject: [PATCH] add swa --- cogvideox/models/transformer3d.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/cogvideox/models/transformer3d.py b/cogvideox/models/transformer3d.py index ef59345..bf81079 100644 --- a/cogvideox/models/transformer3d.py +++ b/cogvideox/models/transformer3d.py @@ -90,6 +90,7 @@ class CogVideoXBlock(nn.Module): ff_bias: bool = True, attention_out_bias: bool = True, swa: bool = False, + window_size: int = 1024, ): super().__init__() @@ -104,7 +105,7 @@ class CogVideoXBlock(nn.Module): eps=1e-6, bias=attention_bias, out_bias=attention_out_bias, - processor=CogVideoXSWAAttnProcessor2_0() if swa else CogVideoXAttnProcessor2_0(), + processor=CogVideoXSWAAttnProcessor2_0(window_size) if swa else CogVideoXAttnProcessor2_0(), ) # 2. Feed Forward @@ -248,6 +249,7 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin): temporal_interpolation_scale: float = 1.0, use_rotary_positional_embeddings: bool = False, swa: bool = False, + window_size: int = 1024, ): super().__init__() inner_dim = num_attention_heads * attention_head_dim @@ -296,6 +298,7 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin): norm_elementwise_affine=norm_elementwise_affine, norm_eps=norm_eps, swa=swa, + window_size=window_size, ) for _ in range(num_layers) ]