From 01bb5672a599bdfceeb200d8acc5ad69a742ce2a Mon Sep 17 00:00:00 2001 From: Hz_Zhang <47402297+HaozheZhang6@users.noreply.github.com> Date: Sun, 5 Jul 2026 10:16:40 -0700 Subject: [PATCH] Fix Head modulation broadcast for batch size > 1 in WanModel --- diffsynth/models/wan_video_dit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diffsynth/models/wan_video_dit.py b/diffsynth/models/wan_video_dit.py index 52f607e10..06e7a18a1 100644 --- a/diffsynth/models/wan_video_dit.py +++ b/diffsynth/models/wan_video_dit.py @@ -279,7 +279,7 @@ def forward(self, x, t_mod): shift, scale = (self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(2)).chunk(2, dim=2) x = (self.head(self.norm(x) * (1 + scale.squeeze(2)) + shift.squeeze(2))) else: - shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(2, dim=1) + shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(1)).chunk(2, dim=1) x = (self.head(self.norm(x) * (1 + scale) + shift)) return x