fix(wan): adapt patchify consumers to bare-tensor return (#1063)

WanModel.patchify returns a bare [b,c,f,h,w] tensor, but WanModel.forward and usp_dit_forward unpacked x, (f,h,w) = self.patchify(x). Flatten and derive the grid in the consumers to match. Drops the earlier model_fn_wans2v edit: that path runs WanS2VModel, whose patchify already returns (flattened_x, grid_size), so the original tuple-unpack was correct.

Author: genno-whittlery

diffsynth/models/wan_video_dit.py

@@ -527,8 +527,10 @@ def forward(self,
             clip_embdding = self.img_emb(clip_feature)
             context = torch.cat([clip_embdding, context], dim=1)
 
-        x, (f, h, w) = self.patchify(x)
-        
+        x = self.patchify(x)
+        f, h, w = x.shape[-3], x.shape[-2], x.shape[-1]
+        x = x.flatten(2).transpose(1, 2)
+
         freqs = torch.cat([
             self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
             self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),

diffsynth/utils/xfuser/xdit_context_parallel.py

@@ -81,8 +81,10 @@ def usp_dit_forward(self,
         clip_embdding = self.img_emb(clip_feature)
         context = torch.cat([clip_embdding, context], dim=1)
 
-    x, (f, h, w) = self.patchify(x)
-    
+    x = self.patchify(x)
+    f, h, w = x.shape[-3], x.shape[-2], x.shape[-1]
+    x = x.flatten(2).transpose(1, 2)
+
     freqs = torch.cat([
         self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
         self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),