fix(diffusion): address Wan auto decode review feedback

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed 2026-08-12 11:49:47 +08:00
1 parent 879d19c4c2
commit c8210b8b5c
7 files changed
+33 -71

No files matched your search

+4 -4
View File
@@ -210,11 +210,11 @@ class DiffusionParallelConfig:
vae_parallel_mode: str = "tile"
"""VAE parallel decode strategy.
- "tile": Patch/tile parallel decode (default). Each rank decodes a subset
of spatial tiles and the results are stitched on rank 0.
- "auto": Select the longer spatial-shard axis for validated Wan tiled
decodes on a full DiT-group topology; otherwise use patch/tile parallel
decode.
- "tile": Patch/tile parallel decode (default). Each rank decodes a subset
of spatial tiles and the results are stitched on rank 0.
- "spatial_shard_height": Spatially-sharded decode that splits decoder
feature maps along height and exchanges halo rows around spatial
convolutions.
@@ -254,9 +254,9 @@ class DiffusionParallelConfig:
assert self.allgather_degree > 0, "AllGather degree must be > 0"
assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0"
assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0"
assert self.vae_parallel_mode in {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}, (
assert self.vae_parallel_mode in {"tile", "auto", "spatial_shard_height", "spatial_shard_width"}, (
"vae_parallel_mode must be one of "
"{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
"{'tile', 'auto', 'spatial_shard_height', 'spatial_shard_width'}, "
f"but got {self.vae_parallel_mode!r}."
)
if self.allgather_degree > 1: