From dfef484db5341b4143bb90a7192f6574c2e81316 Mon Sep 17 00:00:00 2001 From: lb203 <62638829+LinB203@users.noreply.github.com> Date: Wed, 20 Mar 2024 23:04:02 +0800 Subject: [PATCH] t2v attention_mode Former-commit-id: 4fd8bd87e0d0a98d5d0a3201d5db9a5ba055b894 --- opensora/models/diffusion/latte/modeling_latte.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/opensora/models/diffusion/latte/modeling_latte.py b/opensora/models/diffusion/latte/modeling_latte.py index 14d388275..f1b3a4f31 100644 --- a/opensora/models/diffusion/latte/modeling_latte.py +++ b/opensora/models/diffusion/latte/modeling_latte.py @@ -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) \ No newline at end of file + print(model)