Fix Head modulation broadcast for batch size > 1 in WanModel

[HaozheZhang6]

Fixes #1516.

Head.forward only works for batch size 1 when t_mod is 2D (b, dim). self.modulation is (1, 2, dim), so modulation + t_mod broadcasts (1, 2, dim) against (b, dim): the 2 and b dims line up, so it silently mismatches at b=2 and raises at b>=3. The normal forward path hits this — self.head(x, t) passes the per-sample time embedding t of shape (b, dim).

Unsqueeze t_mod to (b, 1, dim) so it broadcasts against (1, 2, dim) -> (b, 2, dim), matching the 3D branch just above. b=1 output is unchanged and the 3D per-token branch is untouched.

Repro on current main (raises at b>=3):

import torch
from diffsynth.models.wan_video_dit import Head
head = Head(dim=16, out_dim=4, patch_size=(1,2,2), eps=1e-6)
head(torch.randn(3, 5, 16), torch.randn(3, 16))
# RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 1

[gemini-code-assist]

Code Review

This pull request modifies the forward method in diffsynth/models/wan_video_dit.py to unsqueeze t_mod along the second dimension (dim=1) before adding it to self.modulation in the else branch, ensuring correct shape alignment for the tensor addition. There are no review comments to evaluate, and I have no additional feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.