minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] improve: use conv2d width padding in wanvae (#27081)

Browse Source
This commit is contained in:
Mick
2026-06-03 10:14:50 +08:00
committed by GitHub
parent 13852d3f31
commit f0e18be0ba
1 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_dist_utils.py
+5 -4
View File
@@ -202,22 +202,23 @@ class WanDistConv2d(nn.Conv2d):
self.padding: tuple[int, int] self.padding: tuple[int, int]
if self.height_halo_size > 0: if self.height_halo_size > 0:
self._padding = (self.padding[1], self.padding[1], 0, 0) self._padding = (0, 0, 0, 0)
else: else:
self._padding = ( self._padding = (
self.padding[1], 0,
self.padding[1], 0,
self.padding[0], self.padding[0],
self.padding[0], self.padding[0],
) )
self.padding = (0, 0) self.padding = (0, self.padding[1])
self._halo_recv_top_buf: torch.Tensor | None = None self._halo_recv_top_buf: torch.Tensor | None = None
self._halo_recv_bottom_buf: torch.Tensor | None = None self._halo_recv_bottom_buf: torch.Tensor | None = None
self.rank = get_sp_parallel_rank() self.rank = get_sp_parallel_rank()
self.world_size = get_sp_world_size() self.world_size = get_sp_world_size()
def forward(self, x): def forward(self, x):
if any(self._padding):
x = F.pad(x, self._padding) x = F.pad(x, self._padding)
x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange( x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(