t2v attention_mode

Former-commit-id: 4fd8bd87e0d0a98d5d0a3201d5db9a5ba055b894
2024-03-20 23:04:02 +08:00
parent 4d476bc759
commit dfef484db5
1 changed files with 4 additions and 1 deletions
@@ -649,6 +649,7 @@ class LatteT2V(ModelMixin, ConfigMixin):
            attention_type: str = "default",
            caption_channels: int = None,
            video_length: int = 16,
+            attention_mode: str = 'flash'
    ):
        super().__init__()
        self.use_linear_projection = use_linear_projection
@@ -751,6 +752,7 @@ class LatteT2V(ModelMixin, ConfigMixin):
                    norm_elementwise_affine=norm_elementwise_affine,
                    norm_eps=norm_eps,
                    attention_type=attention_type,
+                    attention_mode=attention_mode
                )
                for d in range(num_layers)
            ]
@@ -775,6 +777,7 @@ class LatteT2V(ModelMixin, ConfigMixin):
                    norm_elementwise_affine=norm_elementwise_affine,
                    norm_eps=norm_eps,
                    attention_type=attention_type,
+                    attention_mode=attention_mode
                )
                for d in range(num_layers)
            ]
@@ -1161,4 +1164,4 @@ if __name__ == '__main__':
    model = LatteT2V.from_config(a)
    ckpt = torch.load(r"E:\下载\t2v.pt", map_location='cpu')['model']
    model.load_state_dict(ckpt)
-    print(model)
+    print(model)