Parallel VAE fails when the degree of parallelism > 1

[gty111]

When generating images with --use_parallel_vae, the program hangs. Because only the last rank of PP or SP execute self.vae.decode , while all the ranks are expected.

xDiT/xfuser/model_executor/pipelines/pipeline_stable_diffusion_3.py

Lines 382 to 392 in e506417

	if is_dp_last_group():
	if output_type == "latent":
	image = latents
	else:
	latents = (
	latents / self.vae.config.scaling_factor
	) + self.vae.config.shift_factor
	image = self.vae.decode(latents, return_dict=False)[0]
	image = self.image_processor.postprocess(image, output_type=output_type)