From 0be4fc62c95dc590e199e6a20aca0384f29fec88 Mon Sep 17 00:00:00 2001 From: Yongqi Chen <144848849+BrianChen1129@users.noreply.github.com> Date: Tue, 25 Feb 2025 11:11:17 -0500 Subject: [PATCH] fix train/distill issue (#215) Co-authored-by: Peiyuan Zhang --- fastvideo/models/hunyuan/modules/models.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/fastvideo/models/hunyuan/modules/models.py b/fastvideo/models/hunyuan/modules/models.py index 252ec53e1..936515392 100644 --- a/fastvideo/models/hunyuan/modules/models.py +++ b/fastvideo/models/hunyuan/modules/models.py @@ -523,6 +523,9 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: if guidance is None: guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16) + if mask_strategy is None: + mask_strategy = [[None] * len(self.heads_num) + for _ in range(len(self.double_blocks) + len(self.single_blocks))] img = x = hidden_states text_mask = encoder_attention_mask t = timestep