diff --git a/setup.py b/setup.py index cf1c7bbe..73392186 100644 --- a/setup.py +++ b/setup.py @@ -31,7 +31,7 @@ def get_cuda_version(): "transformers>=4.39.1", "sentencepiece>=0.1.99", "beautifulsoup4>=4.12.3", - "distvae", + "distvae>=0.1.0", "yunchang>=0.6.0", "einops", "diffusers>=0.33.0", diff --git a/tests/core/test_distvae_integration.py b/tests/core/test_distvae_integration.py new file mode 100644 index 00000000..48d0d397 --- /dev/null +++ b/tests/core/test_distvae_integration.py @@ -0,0 +1,764 @@ +"""xDiT orchestration at the DistVAE boundary. + +DistVAE owns adapter selection, tile planning, and tile distribution algorithms. These tests +cover only the xDiT policy that discovers VAEs, selects its runtime group, forwards CLI settings, +and presents decode failures. +""" + +from types import SimpleNamespace +from unittest import mock + +import diffusers +import pytest +import torch +import torch.nn as nn + +from distvae.utils import ParallelContext +from distvae.vae import VAERowSplitError, parallel +from xfuser import envs +from xfuser.config import FlexibleArgumentParser, xFuserArgs +from xfuser.model_executor.pipelines import base_pipeline +from xfuser.model_executor.models.runner_models import base_model, vae_manager + + +class _TestRunner(base_model.xFuserModel): + def _load_model(self): + raise NotImplementedError + + def _run_pipe(self, input_args): + raise NotImplementedError + + +@pytest.fixture(autouse=True) +def _distributed_log_environment(monkeypatch): + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "1") + + +@pytest.fixture(scope="module", autouse=True) +def _torch_group_norm_for_distvae_tests(): + original = torch.nn.GroupNorm + envs.restore_torch_group_norm_for_distvae() + yield + torch.nn.GroupNorm = original + + +def _runner(**config): + runner = object.__new__(_TestRunner) + defaults = { + "enable_slicing": False, + "enable_tiling": False, + "vae_tile_size_height": None, + "vae_tile_size_width": None, + "vae_tile_overlap_height": None, + "vae_tile_overlap_width": None, + "enable_sequential_cpu_offload": False, + "enable_model_cpu_offload": False, + "use_parallel_vae": False, + } + defaults.update(config) + runner.config = SimpleNamespace(**defaults) + runner.capabilities = base_model.ModelCapabilities( + use_parallel_vae=True, + use_parallel_vae_encoder=True, + enable_tiling=True, + enable_slicing=True, + ) + runner.settings = SimpleNamespace(model_name="test-model", valid_tasks=[]) + runner._vae_manager = vae_manager.VAEManager( + runner.config, runner.capabilities, runner.settings + ) + return runner + + +@pytest.mark.parametrize("add_args", [xFuserArgs.add_runner_args, xFuserArgs.add_cli_args]) +def test_cli_propagates_vae_settings(add_args): + parser = add_args(FlexibleArgumentParser(description="xDiT")) + parsed = parser.parse_args( + [ + "--model", + "test-model", + "--use-parallel-vae", + "--enable_tiling", + "--vae_tile_size_height", + "320", + "--vae_tile_size_width", + "512", + "--vae_tile_overlap_height", + "32", + "--vae_tile_overlap_width", + "64", + ] + ) + + config = xFuserArgs.from_cli_args(parsed) + + assert config.use_parallel_vae is True + assert config.enable_tiling is True + assert config.vae_tile_size_height == 320 + assert config.vae_tile_size_width == 512 + assert config.vae_tile_overlap_height == 32 + assert config.vae_tile_overlap_width == 64 + + +@pytest.mark.parametrize("add_args", [xFuserArgs.add_runner_args, xFuserArgs.add_cli_args]) +def test_vae_tile_settings_document_tiling_requirement(add_args): + parser = add_args(FlexibleArgumentParser(description="xDiT")) + tile_settings = [ + action + for action in parser._actions + if action.dest.startswith("vae_tile_") + ] + + assert len(tile_settings) == 4 + assert all("Requires --enable_tiling." in action.help for action in tile_settings) + + +@pytest.mark.parametrize( + ("config", "message"), + [ + ( + {"vae_tile_size_height": 320}, + "--vae_tile_size_height and --vae_tile_size_width must be provided together", + ), + ( + {"vae_tile_size_width": 512}, + "--vae_tile_size_height and --vae_tile_size_width must be provided together", + ), + ( + {"vae_tile_size_height": 0, "vae_tile_size_width": 512}, + "--vae_tile_size_height must be positive", + ), + ( + {"vae_tile_size_height": 320, "vae_tile_size_width": -1}, + "--vae_tile_size_width must be positive", + ), + ], +) +def test_rectangular_vae_tile_settings_are_validated(config, message): + runner = _runner() + + with pytest.raises(ValueError, match=message): + vae_manager.validate_vae_config( + xFuserArgs(model="test-model", **config), + runner.capabilities, + runner.settings, + ) + + +@pytest.mark.parametrize( + "config", + [ + {"vae_tile_size_height": 320, "vae_tile_size_width": 512}, + {"vae_tile_overlap_height": 32, "vae_tile_overlap_width": 64}, + ], +) +def test_vae_tile_settings_require_tiling(config): + runner = _runner() + + with pytest.raises(ValueError, match="require --enable_tiling"): + vae_manager.validate_vae_config( + xFuserArgs(model="test-model", **config), + runner.capabilities, + runner.settings, + ) + + +def test_rectangular_vae_tile_flag_does_not_enable_tiling(): + runner = _runner(vae_tile_size_height=320, vae_tile_size_width=None) + + assert runner._vae_manager._tiling_flag() is None + + +def test_tile_overlap_flag_does_not_enable_tiling(): + runner = _runner(vae_tile_overlap_height=32, vae_tile_overlap_width=64) + + assert runner._vae_manager._tiling_flag() is None + + +@pytest.mark.parametrize( + ("config", "message"), + [ + ( + {"vae_tile_overlap_height": 32}, + "--vae_tile_overlap_height and --vae_tile_overlap_width must be provided together", + ), + ( + {"vae_tile_overlap_width": 64}, + "--vae_tile_overlap_height and --vae_tile_overlap_width must be provided together", + ), + ( + {"vae_tile_overlap_height": -1, "vae_tile_overlap_width": 64}, + "--vae_tile_overlap_height must be non-negative", + ), + ( + {"vae_tile_overlap_height": 32, "vae_tile_overlap_width": -1}, + "--vae_tile_overlap_width must be non-negative", + ), + ], +) +def test_per_axis_vae_tile_overlap_settings_are_validated(config, message): + runner = _runner() + + with pytest.raises(ValueError, match=message): + vae_manager.validate_vae_config( + xFuserArgs(model="test-model", **config), + runner.capabilities, + runner.settings, + ) + + +def test_zero_overlap_is_valid_for_an_inactive_strip_axis(): + runner = _runner() + config = xFuserArgs( + model="test-model", + enable_tiling=True, + vae_tile_overlap_height=0, + vae_tile_overlap_width=64, + ) + + vae_manager.validate_vae_config( + config, runner.capabilities, runner.settings + ) + + +def test_vae_validator_restores_aiter_groupnorm_for_parallel_decode(): + runner = _runner() + config = xFuserArgs(model="test-model", use_parallel_vae=True) + + with mock.patch.object( + vae_manager, + "restore_torch_group_norm_for_distvae", + return_value=True, + ) as restore: + vae_manager.validate_vae_config( + config, + runner.capabilities, + runner.settings, + ) + + restore.assert_called_once_with() + + +def test_enable_tiling_without_dimensions_keeps_native_window(): + vae = SimpleNamespace(enable_tiling=mock.Mock(), decode=mock.Mock()) + runner = _runner(enable_tiling=True) + runner.pipe = SimpleNamespace(vae=vae) + + with ( + mock.patch.object(vae_manager.vae_tiling, "require_vae_support"), + mock.patch.object( + vae_manager.vae_tiling, "tile_shape", return_value=(512, 512) + ), + mock.patch.object(vae_manager.vae_tiling, "apply_tile_plan") as apply, + mock.patch.object(runner._vae_manager, "_apply_vae_tile_overlap"), + mock.patch.object( + runner._vae_manager, "_check_tiles_against_parallel_vae" + ), + mock.patch.object(runner._vae_manager, "_install_vae_tiled_decode"), + mock.patch.object(runner._vae_manager, "_install_vae_decode_guard"), + ): + runner._vae_manager.enable_options([vae]) + + vae.enable_tiling.assert_called_once_with() + apply.assert_not_called() + + +def test_worker_vae_wrapper_uses_runtime_device_group_with_public_api(): + vae = SimpleNamespace() + device_group = object() + coordinator = SimpleNamespace(device_group=device_group) + + with ( + mock.patch.object( + base_pipeline, "get_vae_parallel_group", return_value=coordinator + ), + mock.patch.object( + base_pipeline, "parallelize_decoder" + ) as parallelize_decoder, + ): + converted = base_pipeline.xFuserVAEWrapper._convert_vae(object(), vae) + + assert converted is vae + parallelize_decoder.assert_called_once_with(vae, device_group) + + +def test_worker_vae_wrapper_accepts_dedicated_raw_process_group(): + vae = SimpleNamespace() + process_group = object() + + with ( + mock.patch.object( + base_pipeline, "get_vae_parallel_group", return_value=process_group + ), + mock.patch.object( + base_pipeline, "parallelize_decoder" + ) as parallelize_decoder, + ): + converted = base_pipeline.xFuserVAEWrapper._convert_vae(object(), vae) + + assert converted is vae + parallelize_decoder.assert_called_once_with(vae, process_group) + + +def test_pipeline_vae_conversion_uses_runtime_vae_group(): + vae = SimpleNamespace() + device_group = object() + coordinator = SimpleNamespace(device_group=device_group) + + with ( + mock.patch.object( + base_pipeline, "get_vae_parallel_group", return_value=coordinator + ), + mock.patch.object( + base_pipeline, "parallelize_decoder" + ) as parallelize_decoder, + ): + converted = base_pipeline.xFuserPipelineBaseWrapper._convert_vae(object(), vae) + + assert converted is vae + parallelize_decoder.assert_called_once_with(vae, device_group) + + +def test_parallel_setup_adapts_runtime_group_for_distvae(): + tiled, sharded = SimpleNamespace(use_tiling=True), SimpleNamespace() + device_group = object() + coordinator = SimpleNamespace( + device_group=device_group, + rank_in_group=1, + world_size=3, + ranks=[4, 7, 9], + ) + runner = _runner(use_parallel_vae=True) + runner.pipe = SimpleNamespace(vae=tiled) + runner.second_pipe = SimpleNamespace(vae=sharded) + + with ( + mock.patch.object( + vae_manager, "get_vae_parallel_group", return_value=coordinator + ), + mock.patch.object(torch.distributed, "get_world_size", return_value=3), + mock.patch.object(torch.distributed, "get_rank", return_value=1), + mock.patch.object( + vae_manager.vae_tiling, + "supports_tile_parallel", + side_effect=[True, False], + ), + mock.patch.object(vae_manager, "mark") as mark, + mock.patch.object( + vae_manager.vae_parallel, + "parallelize_decoder", + return_value="DecoderAdapter", + ) as parallelize_decoder, + mock.patch.object( + vae_manager.vae_parallel, + "parallelize_encoder", + return_value="EncoderAdapter", + ) as parallelize_encoder, + ): + runner._vae_manager.setup_parallel_vae([tiled, sharded]) + + context = mark.call_args.args[1] + assert mark.call_args.args[0] is tiled + assert context == ParallelContext( + group=device_group, + rank=1, + world_size=3, + patch_dim=-2, + global_ranks=(4, 7, 9), + ) + parallelize_decoder.assert_called_once_with(sharded, device_group) + assert parallelize_encoder.call_args_list == [ + mock.call(tiled, device_group), + mock.call(sharded, device_group), + ] + + +def test_tile_overlap_is_applied_through_distvae(): + vae = SimpleNamespace() + runner = _runner( + vae_tile_overlap_height=32, + vae_tile_overlap_width=64, + height=320, + width=1280, + ) + overlap_plan = { + "tile_overlap_factor_height": 0.0625, + "tile_overlap_factor_width": 0.125, + } + + with ( + mock.patch.object( + vae_manager.vae_tiling, + "tile_overlap", + side_effect=[(128, 128), (32, 64)], + ), + mock.patch.object( + vae_manager.vae_tiling, + "tile_overlap_plan", + return_value=overlap_plan, + ) as tile_overlap_plan, + mock.patch.object(vae_manager.vae_tiling, "apply_tile_plan") as apply, + ): + runner._vae_manager._apply_vae_tile_overlap(vae, (320, 1280)) + + tile_overlap_plan.assert_called_once_with( + vae, 32, 64, sample_shape=(320, 1280) + ) + apply.assert_called_once_with(vae, overlap_plan) + + +def test_tile_overlap_tracks_each_run_sample_shape(): + vae = SimpleNamespace() + runner = _runner( + vae_tile_overlap_height=32, + vae_tile_overlap_width=64, + ) + + with mock.patch.object( + runner._vae_manager, "_apply_vae_tile_overlap" + ) as apply: + runner._vae_manager.prepare_run( + [vae], {"height": 320, "width": 1280} + ) + runner._vae_manager.prepare_run( + [vae], {"height": 320, "width": 1280} + ) + runner._vae_manager.prepare_run( + [vae], {"height": 640, "width": 960} + ) + + assert apply.call_args_list == [ + mock.call(vae, (320, 1280)), + mock.call(vae, (640, 960)), + ] + + +def test_tile_overlap_refuses_an_inexact_plan_without_mutating_the_vae(): + vae = SimpleNamespace() + runner = _runner( + vae_tile_overlap_height=33, + vae_tile_overlap_width=65, + height=1024, + width=1024, + ) + + with ( + mock.patch.object( + vae_manager.vae_tiling, "tile_shape", return_value=(512, 768) + ), + mock.patch.object( + vae_manager.vae_tiling, + "tile_overlap_plan", + return_value=None, + ), + mock.patch.object(vae_manager.vae_tiling, "apply_tile_plan") as apply, + ): + with pytest.raises(ValueError) as error: + runner._vae_manager._apply_vae_tile_overlap(vae, (1024, 1024)) + + assert "33x65 pixels" in str(error.value) + assert "512x768" in str(error.value) + apply.assert_not_called() + + +def test_rectangular_tile_shape_is_applied_exactly_through_distvae(): + vae = SimpleNamespace() + runner = _runner(vae_tile_size_height=320, vae_tile_size_width=512) + plan = { + "tile_sample_min_height": 320, + "tile_sample_min_width": 512, + "tile_latent_min_height": 40, + "tile_latent_min_width": 64, + } + + with ( + mock.patch.object( + vae_manager.vae_tiling, "tile_shape_plan", return_value=plan + ) as tile_shape_plan, + mock.patch.object(vae_manager.vae_tiling, "apply_tile_plan") as apply, + ): + assert runner._vae_manager._apply_vae_tile_shape(vae) == (320, 512) + + tile_shape_plan.assert_called_once_with(vae, 320, 512) + apply.assert_called_once_with(vae, plan) + + +def test_rectangular_tile_shape_refuses_a_non_exact_plan_actionably(): + vae = SimpleNamespace() + runner = _runner(vae_tile_size_height=321, vae_tile_size_width=512) + + with mock.patch.object( + vae_manager.vae_tiling, "tile_shape_plan", return_value=None + ): + with pytest.raises(ValueError) as error: + runner._vae_manager._apply_vae_tile_shape(vae) + + message = str(error.value) + assert "--vae_tile_size_height 321" in message + assert "--vae_tile_size_width 512" in message + assert "exactly" in message + + +def test_rectangular_tile_shape_is_applied_to_every_staged_vae(): + first = SimpleNamespace(enable_tiling=mock.Mock(), decode=mock.Mock()) + second = SimpleNamespace(enable_tiling=mock.Mock(), decode=mock.Mock()) + runner = _runner( + enable_tiling=True, + vae_tile_size_height=320, + vae_tile_size_width=512, + ) + runner.pipe = SimpleNamespace(vae=first) + runner.second_pipe = SimpleNamespace(vae=second) + + with ( + mock.patch.object(vae_manager.vae_tiling, "require_vae_support"), + mock.patch.object( + vae_manager.vae_tiling, "tile_shape", return_value=(512, 512) + ), + mock.patch.object( + runner._vae_manager, + "_apply_vae_tile_shape", + return_value=(320, 512), + ) as apply_shape, + mock.patch.object(runner._vae_manager, "_apply_vae_tile_overlap"), + mock.patch.object( + runner._vae_manager, "_check_tiles_against_parallel_vae" + ), + mock.patch.object(runner._vae_manager, "_install_vae_tiled_decode"), + mock.patch.object(runner._vae_manager, "_install_vae_decode_guard"), + ): + runner._vae_manager.enable_options([first, second]) + + assert apply_shape.call_args_list == [mock.call(first), mock.call(second)] + + +def test_single_gpu_rectangular_scalar_vae_installs_undispatched_tiled_decode(): + vae = SimpleNamespace() + runner = _runner( + use_parallel_vae=False, + vae_tile_size_height=320, + vae_tile_size_width=512, + ) + installed = mock.Mock() + + with ( + mock.patch.object( + vae_manager.vae_tile_parallel, "context_of", return_value=None + ), + mock.patch.object( + vae_manager.vae_tiling, "tiled_decode_for", return_value=installed + ) as tiled_decode_for, + ): + runner._vae_manager._install_vae_tiled_decode(vae) + + tiled_decode_for.assert_called_once_with(vae) + assert vae.tiled_decode is installed + + +def test_native_keyed_vae_keeps_its_tiled_decode_when_distvae_returns_none(): + native = mock.Mock() + vae = SimpleNamespace(tiled_decode=native) + runner = _runner(vae_tile_size_height=320, vae_tile_size_width=512) + + with ( + mock.patch.object( + vae_manager.vae_tile_parallel, "context_of", return_value=None + ), + mock.patch.object( + vae_manager.vae_tiling, "tiled_decode_for", return_value=None + ) as tiled_decode_for, + ): + runner._vae_manager._install_vae_tiled_decode(vae) + + tiled_decode_for.assert_called_once_with(vae) + assert vae.tiled_decode is native + + +def test_tiled_decode_install_uses_distvae_context_and_sharing(): + group = object() + context = ParallelContext( + group=group, + rank=0, + world_size=2, + patch_dim=-2, + global_ranks=(0, 1), + ) + vae = SimpleNamespace() + runner = _runner() + installed = mock.Mock() + + with ( + mock.patch.object( + vae_manager.vae_tile_parallel, "context_of", return_value=context + ), + mock.patch.object( + vae_manager.vae_tile_parallel, + "sharing", + return_value=("dispatch", "assemble"), + ) as sharing, + mock.patch.object( + vae_manager.vae_tiling, + "tiled_decode_for", + return_value=installed, + ) as tiled_decode_for, + ): + runner._vae_manager._install_vae_tiled_decode(vae) + + sharing.assert_called_once_with(context) + tiled_decode_for.assert_called_once_with(vae, "dispatch", "assemble") + assert vae.tiled_decode is installed + + +def test_parallel_vae_rank_hint_increases_only_the_sharded_tile_axis(): + vae = SimpleNamespace() + runner = _runner(use_parallel_vae=True) + + def shape_plan(_vae, height, width): + if width != 64 or height % 32: + return None + return {"rows": height // 64} + + def rows(_vae, plan=None): + return 1 if plan is None else plan["rows"] + + with ( + mock.patch.object( + vae_manager.vae_tile_parallel, "context_of", return_value=None + ), + mock.patch.object( + vae_manager, "get_vae_parallel_world_size", return_value=2 + ), + mock.patch.object(vae_manager, "latent_rows", side_effect=rows), + mock.patch.object( + vae_manager.vae_tiling, "tile_shape", return_value=(64, 64) + ), + mock.patch.object( + vae_manager.vae_tiling, "tile_shape_plan", side_effect=shape_plan + ) as tile_shape_plan, + ): + with pytest.raises(ValueError) as error: + runner._vae_manager._check_tiles_against_parallel_vae( + vae, (512, 512) + ) + + assert "--vae_tile_size_height 128 --vae_tile_size_width 64" in str(error.value) + assert tile_shape_plan.call_args_list[0] == mock.call(vae, 64, 64) + + +def test_parallel_vae_accepts_rectangular_tile_with_enough_latent_height(): + vae = SimpleNamespace( + tile_sample_min_height=256, + tile_sample_min_width=64, + tile_latent_min_height=32, + tile_latent_min_width=8, + ) + runner = _runner(use_parallel_vae=True) + + with ( + mock.patch.object( + vae_manager.vae_tile_parallel, "context_of", return_value=None + ), + mock.patch.object( + vae_manager, "get_vae_parallel_world_size", return_value=16 + ), + ): + runner._vae_manager._check_tiles_against_parallel_vae( + vae, native_shape=(512, 512) + ) + + +def test_decode_guard_reports_rectangular_window_and_flags(): + original = mock.Mock(side_effect=RuntimeError("decoder padding failure")) + vae = SimpleNamespace(decode=original) + runner = _runner( + vae_tile_size_height=320, + vae_tile_size_width=512, + ) + + with mock.patch.object( + vae_manager.vae_tiling, "is_tile_padding_error", return_value=True + ) as is_padding: + runner._vae_manager._install_vae_decode_guard( + vae, tile_shape=(320, 512) + ) + with pytest.raises(RuntimeError) as error: + vae.decode(object()) + + message = str(error.value) + assert "320x512" in message + assert "--vae_tile_size_height" in message + assert "--vae_tile_size_width" in message + is_padding.assert_called_once() + + +def test_decode_guard_suggests_complete_tiles_for_incompatible_row_sharding(): + vae = SimpleNamespace( + config=SimpleNamespace(scale_factor_spatial=16), + decode=mock.Mock(side_effect=VAERowSplitError(rows=45, factor=2)), + ) + runner = _runner(use_parallel_vae=True) + + with mock.patch.object( + vae_manager.vae_tiling, "supports_tile_parallel", return_value=True + ): + runner._vae_manager._install_vae_decode_guard(vae) + with pytest.raises(ValueError) as error: + vae.decode(object()) + + assert str(error.value) == ( + "Cannot row-shard 45 latent rows in the VAE decoder: this VAE processes rows " + "in groups of 2, but 45 is not divisible by 2. Use an output height divisible " + "by 32, enable --enable_tiling to distribute complete tiles instead, or disable " + "--use_parallel_vae." + ) + + +def test_decode_oom_hint_reports_rectangular_window(): + vae = SimpleNamespace(use_tiling=True) + runner = _runner() + + with mock.patch.object( + vae_manager.vae_tiling, "tile_shape", return_value=(320, 512) + ): + hint = runner._vae_manager._vae_decode_oom_hint(vae) + + assert "320x512" in hint + assert "--vae_tile_size_height" in hint + assert "--vae_tile_size_width" in hint + assert "no single tile window" not in hint + + +class AiterGroupNorm(nn.Module): + __module__ = "aiter.ops.groupnorm" + + +def _two_d_vae(): + return diffusers.AutoencoderKL( + block_out_channels=[8, 8, 16, 16], + layers_per_block=1, + latent_channels=4, + norm_num_groups=8, + sample_size=32, + down_block_types=["DownEncoderBlock2D"] * 4, + up_block_types=["UpDecoderBlock2D"] * 4, + ) + + +def test_non_aiter_group_norm_restoration_is_a_no_op(): + original = envs._TORCH_GROUPNORM + with mock.patch.object(nn, "GroupNorm", original): + assert envs.restore_torch_group_norm_for_distvae() is False + assert nn.GroupNorm is original + + +def test_aiter_replacement_blocks_recognition_until_xdit_restores_torch_norm(): + original = envs._TORCH_GROUPNORM + vae = _two_d_vae() + assert parallel.decoder_adapter_name(vae) == "DecoderAdapter" + + with mock.patch.object(nn, "GroupNorm", AiterGroupNorm): + assert parallel.decoder_adapter_name(vae) is None + assert envs.restore_torch_group_norm_for_distvae() is True + assert nn.GroupNorm is original + assert parallel.decoder_adapter_name(vae) == "DecoderAdapter" diff --git a/tests/core/test_envs.py b/tests/core/test_envs.py index 97b0d228..b8292235 100644 --- a/tests/core/test_envs.py +++ b/tests/core/test_envs.py @@ -1,21 +1,26 @@ import unittest from unittest.mock import patch -import torch from xfuser import envs +# get_device checks torch.version.cuda and torch.version.hip. Patch those checks directly so each +# test is independent of the PyTorch build used to run it. + + class TestEnvs(unittest.TestCase): - @patch('torch.cuda.is_available', return_value=True) - def test_get_device_cuda(self, mock_is_available): + @patch('xfuser.envs._is_hip', return_value=False) + @patch('xfuser.envs._is_cuda', return_value=True) + def test_get_device_cuda(self, mock_is_cuda, mock_is_hip): device = envs.get_device(0) self.assertEqual(device.type, 'cuda') self.assertEqual(device.index, 0) device_name = envs.get_device_name() self.assertEqual(device_name, 'cuda') - @patch('torch.cuda.is_available', return_value=False) + @patch('xfuser.envs._is_hip', return_value=False) + @patch('xfuser.envs._is_cuda', return_value=False) @patch('xfuser.envs._is_mps', return_value=True) - def test_get_device_mps(self, mock_is_mps, mock_is_available): + def test_get_device_mps(self, mock_is_mps, mock_is_cuda, mock_is_hip): device = envs.get_device(0) self.assertEqual(device.type, 'mps') device_name = envs.get_device_name() @@ -24,10 +29,11 @@ def test_get_device_mps(self, mock_is_mps, mock_is_available): cuda_version = envs.CUDA_VERSION self.assertIsNotNone(cuda_version) - @patch('torch.cuda.is_available', return_value=False) + @patch('xfuser.envs._is_hip', return_value=False) + @patch('xfuser.envs._is_cuda', return_value=False) @patch('xfuser.envs._is_mps', return_value=False) @patch('xfuser.envs._is_musa', return_value=False) - def test_get_device_cpu(self, mock_is_musa, mock_is_mps, mock_is_available): + def test_get_device_cpu(self, mock_is_musa, mock_is_mps, mock_is_cuda, mock_is_hip): device = envs.get_device(0) self.assertEqual(device.type, 'cpu') device_name = envs.get_device_name() diff --git a/tests/core/test_gilbert.py b/tests/core/test_gilbert.py index e712927d..2106d531 100644 --- a/tests/core/test_gilbert.py +++ b/tests/core/test_gilbert.py @@ -8,28 +8,38 @@ ) +def _tridiagonal(n: int) -> torch.Tensor: + band = torch.eye(n, dtype=torch.bool) + band[:-1, 1:] |= torch.eye(n - 1, dtype=torch.bool) + band[1:, :-1] |= torch.eye(n - 1, dtype=torch.bool) + return band + + class TestSlicedGilbertBlockNeighborMapping(unittest.TestCase): """Tests for sliced_gilbert_block_neighbor_mapping with known expected results.""" - def test_identity_mapping_gives_block_diagonal_mask(self): - """Identity linear_to_hilbert and block_m == block_n -> strictly diagonal mask.""" - t, h, w = 2, 2, 2 + # For eight points in a line, each interior point is adjacent only to its immediate predecessor + # and successor. With two points per block, the expected block-neighbor mask is tridiagonal. + # A 2x2x2 grid is unsuitable because every block is adjacent to every other block. + def test_identity_mapping_gives_a_banded_mask(self): + """Identity linear_to_hilbert and block_m == block_n -> diagonal and its two neighbours.""" + t, h, w = 1, 1, 8 block_m = block_n = 2 linear_to_hilbert = torch.arange(8, dtype=torch.int64) mask = _sliced_gilbert_block_neighbor_mapping( t, h, w, block_m, block_n, linear_to_hilbert ) - expected = torch.eye(4, dtype=torch.bool) + expected = _tridiagonal(4) self.assertEqual(tuple(mask.shape), (4, 4)) self.assertEqual(mask.dtype, torch.bool) self.assertTrue( torch.equal(mask, expected), - f"Expected block-diagonal mask, got\n{mask}", + f"Expected a banded mask, got\n{mask}", ) - def test_public_api_identity_mapping_gives_same_block_diagonal_mask(self): - """Public API with identity (lth, htl) yields same block-diagonal mask.""" - t, h, w = 2, 2, 2 + def test_public_api_identity_mapping_returns_tridiagonal_mask(self): + """The public API returns a tridiagonal mask for identity mappings.""" + t, h, w = 1, 1, 8 block_m = block_n = 2 device = torch.device("cpu") linear_to_hilbert = torch.arange(8, dtype=torch.int64) @@ -38,14 +48,20 @@ def test_public_api_identity_mapping_gives_same_block_diagonal_mask(self): t, h, w, block_m, block_n, device, gilbert_mapping=(linear_to_hilbert, hilbert_to_linear), ) - expected = torch.eye(4, dtype=torch.bool) + expected = _tridiagonal(4) self.assertEqual(tuple(mask.shape), (4, 4)) self.assertEqual(mask.dtype, torch.bool) self.assertTrue( torch.equal(mask, expected), - f"Expected block-diagonal mask, got\n{mask}", + f"Expected a banded mask, got\n{mask}", ) + def test_a_volume_small_enough_that_everything_touches_is_all_true(self): + # In a 2x2x2 cube every point neighbors every other point, so the mask is fully connected. + linear_to_hilbert = torch.arange(8, dtype=torch.int64) + mask = _sliced_gilbert_block_neighbor_mapping(2, 2, 2, 2, 2, linear_to_hilbert) + self.assertTrue(bool(mask.all())) + def test_shape_and_dtype_non_square_blocks(self): """Shape and dtype for non-square block grid (qblocks != kblocks).""" t, h, w = 1, 4, 4 diff --git a/tests/core/test_ring_flash_attn.py b/tests/core/test_ring_flash_attn.py index 12bf63df..9e726306 100644 --- a/tests/core/test_ring_flash_attn.py +++ b/tests/core/test_ring_flash_attn.py @@ -1,9 +1,14 @@ import unittest +import pytest import torch import torch.distributed as dist + +# These tests require flash-attn. Skip the module before importing xDiT attention code when the +# package is unavailable. +flash_attn_func = pytest.importorskip("flash_attn").flash_attn_func + from xfuser.core.long_ctx_attention.ring.ring_flash_attn import xdit_ring_flash_attn_func from xfuser.core.long_ctx_attention import xFuserLongContextAttention -from flash_attn import flash_attn_func import os from xfuser.model_executor.layers.attention_processor import ( diff --git a/tests/core/test_sharding.py b/tests/core/test_sharding.py index 61c54015..a5e97084 100644 --- a/tests/core/test_sharding.py +++ b/tests/core/test_sharding.py @@ -6,7 +6,7 @@ Run with: pytest tests/test_sharding.py -v - pytest tests/test_sharding.py::test_shard_transformer_blocks -v # Single test + pytest tests/test_sharding.py::test_shard_component_basic -v # Single test """ import pytest import torch @@ -20,7 +20,7 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from xfuser.core.distributed.sharding import ( - shard_transformer_blocks, + shard_component, shard_dit, shard_t5_encoder, ) @@ -137,17 +137,17 @@ def forward(self, x): # ============================================================================ -# Test shard_transformer_blocks +# Test shard_component # ============================================================================ -def test_shard_transformer_blocks_basic(setup_distributed, simple_transformer_model): +def test_shard_component_basic(setup_distributed, simple_transformer_model): """Test basic FSDP wrapping of transformer blocks.""" model = simple_transformer_model # Shard the model - sharded_model = shard_transformer_blocks( + sharded_model = shard_component( model, - block_attr='blocks', + wrap_attrs=['blocks'], device_id=0, ) @@ -158,14 +158,14 @@ def test_shard_transformer_blocks_basic(setup_distributed, simple_transformer_mo assert hasattr(sharded_model, 'blocks'), "Blocks attribute should exist" -def test_shard_transformer_blocks_with_dtype(setup_distributed, simple_transformer_model): +def test_shard_component_with_dtype(setup_distributed, simple_transformer_model): """Test FSDP wrapping with dtype conversion.""" model = simple_transformer_model # Shard with bfloat16 - sharded_model = shard_transformer_blocks( + sharded_model = shard_component( model, - block_attr='blocks', + wrap_attrs=['blocks'], device_id=0, dtype=torch.bfloat16, ) @@ -178,26 +178,28 @@ def test_shard_transformer_blocks_with_dtype(setup_distributed, simple_transform assert param.dtype == torch.bfloat16, f"Param dtype should be bfloat16, got {param.dtype}" -def test_shard_transformer_blocks_invalid_attr(setup_distributed, simple_transformer_model): +def test_shard_component_invalid_attr(setup_distributed, simple_transformer_model): """Test error handling for invalid block attribute.""" model = simple_transformer_model - - with pytest.raises(ValueError, match="Model does not have attribute"): - shard_transformer_blocks( + + # rgetattr resolves each component of the nested attribute name; the failing getattr raises + # AttributeError when nonexistent_blocks is absent. + with pytest.raises(AttributeError, match="nonexistent_blocks"): + shard_component( model, - block_attr='nonexistent_blocks', + wrap_attrs=['nonexistent_blocks'], device_id=0, ) -def test_shard_transformer_blocks_with_fsdp_kwargs(setup_distributed, simple_transformer_model): +def test_shard_component_with_fsdp_kwargs(setup_distributed, simple_transformer_model): """Test passing additional FSDP kwargs.""" model = simple_transformer_model # Should not raise an error - sharded_model = shard_transformer_blocks( + sharded_model = shard_component( model, - block_attr='blocks', + wrap_attrs=['blocks'], device_id=0, sync_module_states=True, forward_prefetch=True, @@ -292,9 +294,9 @@ def forward(self, x): model = EmptyBlockModel() # Should still work with empty blocks - sharded_model = shard_transformer_blocks( + sharded_model = shard_component( model, - block_attr='blocks', + wrap_attrs=['blocks'], device_id=0, ) @@ -313,9 +315,9 @@ def forward(self, x): model = SingleBlockModel() - sharded_model = shard_transformer_blocks( + sharded_model = shard_component( model, - block_attr='blocks', + wrap_attrs=['blocks'], device_id=0, ) @@ -327,9 +329,9 @@ def test_default_device_id(setup_distributed, simple_transformer_model): model = simple_transformer_model # Should use current device (0 in single-GPU test) - sharded_model = shard_transformer_blocks( + sharded_model = shard_component( model, - block_attr='blocks', + wrap_attrs=['blocks'], device_id=None, # Explicitly pass None ) @@ -341,9 +343,9 @@ def test_no_dtype_conversion(setup_distributed, simple_transformer_model): model = simple_transformer_model original_dtype = next(model.parameters()).dtype - sharded_model = shard_transformer_blocks( + sharded_model = shard_component( model, - block_attr='blocks', + wrap_attrs=['blocks'], device_id=0, dtype=None, # No conversion ) @@ -363,9 +365,9 @@ def test_parameter_count_preserved(setup_distributed, simple_transformer_model): # Count params before sharding original_param_count = sum(p.numel() for p in model.parameters()) - sharded_model = shard_transformer_blocks( + sharded_model = shard_component( model, - block_attr='blocks', + wrap_attrs=['blocks'], device_id=0, ) @@ -384,9 +386,9 @@ def test_end_to_end_forward_pass(setup_distributed, simple_transformer_model): """Test complete forward pass through sharded model.""" model = simple_transformer_model - sharded_model = shard_transformer_blocks( + sharded_model = shard_component( model, - block_attr='blocks', + wrap_attrs=['blocks'], device_id=0, dtype=torch.float32, ) diff --git a/tests/core/test_vae_lifecycle.py b/tests/core/test_vae_lifecycle.py new file mode 100644 index 00000000..190251d1 --- /dev/null +++ b/tests/core/test_vae_lifecycle.py @@ -0,0 +1,111 @@ +from types import SimpleNamespace +from unittest import mock + +import torch +from torch import nn + +from xfuser.model_executor.models.runner_models import base_model +from xfuser.model_executor.models.runner_models.base_model import xFuserModel +from xfuser.model_executor.models.runner_models.vae_manager import VAEManager + + +class TinyVAE(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(4, 4, kernel_size=3, padding=1) + self.decoded = [] + + def decode(self, z): + self.decoded.append(z) + return z + + +class Pipe: + def __init__(self, vae): + self.vae = vae + + +class Staged: + _decoding_vaes = xFuserModel._decoding_vaes + _convert_vae_to_channels_last = xFuserModel._convert_vae_to_channels_last + + def __init__(self, first, second=None): + self.pipe = Pipe(first) + self.second_pipe = Pipe(second) if second is not None else None + self._vae_manager = VAEManager( + config=object(), + capabilities=object(), + settings=SimpleNamespace(model_output_type="image"), + ) + + +def test_decoding_vaes_accepts_any_number_of_pipeline_stages(): + manager = VAEManager( + config=object(), + capabilities=object(), + settings=SimpleNamespace(model_output_type="image"), + ) + first, second, third = TinyVAE(), TinyVAE(), TinyVAE() + + vaes = manager.decoding_vaes( + [Pipe(first), Pipe(second), Pipe(third), Pipe(first)] + ) + + assert vaes == [first, second, third] + + +def test_channels_last_conversion_preserves_decode_output_for_every_stage(): + first, second = TinyVAE(), TinyVAE() + staged = Staged(first, second) + staged._convert_vae_to_channels_last() + sample = torch.randn(1, 4, 8, 8) + + for vae in (first, second): + assert torch.equal(vae.decode(sample), sample) + assert vae.decoded[-1].is_contiguous(memory_format=torch.channels_last) + + +def test_initialize_sets_up_every_parallel_vae_before_enabling_options(monkeypatch): + first, second = object(), object() + events = [] + + class Runner: + initialize = xFuserModel.initialize + _decoding_vaes = xFuserModel._decoding_vaes + + def __init__(self): + self.config = SimpleNamespace( + use_parallel_vae=True, + use_torch_compile=False, + create_config=lambda: (object(), None), + ) + self._vae_manager = mock.Mock() + self._vae_manager.decoding_vaes.side_effect = ( + lambda pipes: [pipe.vae for pipe in pipes] + ) + + def _load_model_checked(self): + return Pipe(first) + + def _get_runtime_state_pipeline(self): + return self.pipe + + def _post_load_and_state_initialization(self, input_args): + events.append("post-load") + self.second_pipe = Pipe(second) + + def _enable_options(self): + events.append("options") + + runner = Runner() + runner._vae_manager.setup_parallel_vae.side_effect = ( + lambda vaes: events.append("parallel") + ) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(base_model, "log", lambda *args, **kwargs: None) + monkeypatch.setattr(base_model, "initialize_runtime_state", lambda *args: None) + + runner.initialize({}) + + assert events == ["post-load", "parallel", "options"] + runner._vae_manager.setup_parallel_vae.assert_called_once_with([first, second]) diff --git a/tests/core/test_xfuser_attn.py b/tests/core/test_xfuser_attn.py index 19dc857c..75b7bf15 100644 --- a/tests/core/test_xfuser_attn.py +++ b/tests/core/test_xfuser_attn.py @@ -1,12 +1,17 @@ import unittest from itertools import product +import pytest import torch import torch.distributed as dist + +# These tests require flash-attn. Skip the module before importing xDiT attention code when the +# package is unavailable. +flash_attn_func = pytest.importorskip("flash_attn").flash_attn_func + from xfuser.core.long_ctx_attention.ring.ring_flash_attn import ( xdit_ring_flash_attn_func, ) from xfuser.core.long_ctx_attention import xFuserLongContextAttention -from flash_attn import flash_attn_func import os from xfuser.model_executor.layers.attention_processor import ( diff --git a/xfuser/compat.py b/xfuser/compat.py index 0e1afc46..7fbaeea1 100644 --- a/xfuser/compat.py +++ b/xfuser/compat.py @@ -19,7 +19,6 @@ logger = init_logger(__name__) - @lru_cache(maxsize=None) def declared_floor(name: str) -> Optional[str]: """The minimum version of ``name`` that setup.py declares xfuser needs. diff --git a/xfuser/config/args.py b/xfuser/config/args.py index 3ad6b610..ee1267ed 100644 --- a/xfuser/config/args.py +++ b/xfuser/config/args.py @@ -123,6 +123,10 @@ class xFuserArgs: enable_sequential_cpu_offload: bool = False enable_tiling: bool = False enable_slicing: bool = False + vae_tile_size_height: Optional[int] = None + vae_tile_size_width: Optional[int] = None + vae_tile_overlap_height: Optional[int] = None + vae_tile_overlap_width: Optional[int] = None # DiTFastAttn arguments use_fast_attn: bool = False n_calib: int = 8 @@ -412,7 +416,38 @@ def add_cli_args(parser: FlexibleArgumentParser): runtime_group.add_argument( "--enable_slicing", action="store_true", - help="Making VAE decode a tile at a time to save GPU memory.", + help="Decode one batch item at a time to reduce GPU memory use. This has no effect " + "when the batch size is 1.", + ) + runtime_group.add_argument( + "--vae_tile_size_height", + type=int, + default=None, + help="Exact output-pixel height of each VAE tile. Requires --enable_tiling. " + "Must be used with --vae_tile_size_width.", + ) + runtime_group.add_argument( + "--vae_tile_size_width", + type=int, + default=None, + help="Exact output-pixel width of each VAE tile. Requires --enable_tiling. " + "Must be used with --vae_tile_size_height.", + ) + runtime_group.add_argument( + "--vae_tile_overlap_height", + type=int, + default=None, + help="Exact VAE tile overlap along the height axis, in output pixels. Requires " + "--enable_tiling. Must be used with --vae_tile_overlap_width. Height and width " + "may differ; use 0 for an inactive strip axis.", + ) + runtime_group.add_argument( + "--vae_tile_overlap_width", + type=int, + default=None, + help="Exact VAE tile overlap along the width axis, in output pixels. Requires " + "--enable_tiling. Must be used with --vae_tile_overlap_height. Height and width " + "may differ; use 0 for an inactive strip axis.", ) runtime_group.add_argument( "--use_fp8_t5_encoder", @@ -634,7 +669,38 @@ def add_runner_args(parser: FlexibleArgumentParser): parser.add_argument( "--enable_slicing", action="store_true", - help="Enable VAE slicing to save GPU memory.", + help="Decode one batch item at a time to reduce GPU memory use. This has no effect " + "when the batch size is 1.", + ) + parser.add_argument( + "--vae_tile_size_height", + type=int, + default=None, + help="Exact output-pixel height of each VAE tile. Requires --enable_tiling. " + "Must be used with --vae_tile_size_width.", + ) + parser.add_argument( + "--vae_tile_size_width", + type=int, + default=None, + help="Exact output-pixel width of each VAE tile. Requires --enable_tiling. " + "Must be used with --vae_tile_size_height.", + ) + parser.add_argument( + "--vae_tile_overlap_height", + type=int, + default=None, + help="Exact VAE tile overlap along the height axis, in output pixels. Requires " + "--enable_tiling. Must be used with --vae_tile_overlap_width. Height and width " + "may differ; use 0 for an inactive strip axis.", + ) + parser.add_argument( + "--vae_tile_overlap_width", + type=int, + default=None, + help="Exact VAE tile overlap along the width axis, in output pixels. Requires " + "--enable_tiling. Must be used with --vae_tile_overlap_height. Height and width " + "may differ; use 0 for an inactive strip axis.", ) parser.add_argument( "--use_int8_gemms", diff --git a/xfuser/envs.py b/xfuser/envs.py index 3652de27..448556df 100644 --- a/xfuser/envs.py +++ b/xfuser/envs.py @@ -213,7 +213,6 @@ def initialize(self): packages_info["has_long_ctx_attn"] = self.check_long_ctx_attn() packages_info["diffusers_version"] = self.check_diffusers_version() packages_info["has_npu_flash_attn"] = self.check_npu_flash_attn() - packages_info["has_distvae"] = self.check_distvae() self.packages_info = packages_info def check_aiter(self): @@ -380,13 +379,6 @@ def check_npu_flash_attn(self): except ImportError: return False - def check_distvae(self): - try: - from distvae.modules.adapters.vae.decoder_adapters import DecoderAdapter - return True - except ImportError: - return False - def get_packages_info(self): return self.packages_info @@ -404,6 +396,19 @@ def _on_rdna4(self): PACKAGES_CHECKER = PackagesEnvChecker() _TORCH_GROUPNORM = torch.nn.GroupNorm + +def restore_torch_group_norm_for_distvae() -> bool: + """Restore torch GroupNorm when xDiT's ROCm setup replaced it with AITER's. + + DistVAE discovers norms to shard by their torch type. This must run before a VAE intended + for sharding is built, while xDiT is still validating its environment. + """ + if torch.nn.GroupNorm.__module__ != "aiter.ops.groupnorm": + return False + torch.nn.GroupNorm = _TORCH_GROUPNORM + return True + + def _setup_rocm_libraries(): if PACKAGES_CHECKER.packages_info.get("has_aiter", False): try: diff --git a/xfuser/model_executor/models/runner_models/base_model.py b/xfuser/model_executor/models/runner_models/base_model.py index 4adadf58..a859ac37 100644 --- a/xfuser/model_executor/models/runner_models/base_model.py +++ b/xfuser/model_executor/models/runner_models/base_model.py @@ -1,9 +1,7 @@ import abc import torch import copy -import argparse import json -import functools from PIL.Image import Image from typing import Callable, List, Optional, Tuple, Generator from dataclasses import dataclass, field @@ -13,11 +11,9 @@ from diffusers.utils import load_image, export_to_video import numpy as np from xfuser.compat import is_diffusers_import_error -from xfuser.config import args, xFuserArgs +from xfuser.config import xFuserArgs from xfuser.envs import ( PACKAGES_CHECKER, - _TORCH_GROUPNORM, - get_platform, _is_hip, _is_cuda, ) @@ -30,10 +26,13 @@ quantize_linear_layers_to_fp8_blockscale, quantize_linear_layers_to_fp4, quantize_linear_layers_to_nvfp4, - convert_model_convs_to_channels_last, _use_aiter_fp8_rdna4, rgetattr, ) +from xfuser.model_executor.models.runner_models.vae_manager import ( + VAEManager, + validate_vae_config, +) from xfuser.core.distributed import ( get_world_group, @@ -50,7 +49,6 @@ from xfuser.core.distributed.attention_backend import AttentionBackendType from xfuser.core.distributed.attention_schedule import AttentionSchedule, create_hybrid_attn_schedule, create_hybrid_gemm_schedule - packages_info = PACKAGES_CHECKER.get_packages_info() MODEL_REGISTRY = {} @@ -250,6 +248,7 @@ class xFuserModel(abc.ABC): def __init__(self, config: xFuserArgs) -> None: self.settings = copy.deepcopy(self.__class__.settings) self._customize_settings(config) + self._vae_manager = VAEManager(config, self.capabilities, self.settings) self._validate_config(config) self._update_model_settings(config) self.config = config @@ -312,6 +311,8 @@ def initialize(self, input_args: dict) -> None: initialize_runtime_state(self._get_runtime_state_pipeline(), self.engine_config) self._post_load_and_state_initialization(input_args) + if self.config.use_parallel_vae: + self._vae_manager.setup_parallel_vae(self._decoding_vaes()) self._enable_options() if self.config.use_torch_compile: @@ -327,13 +328,7 @@ def _enable_options(self) -> None: if getattr(self.config, "use_spargeattn_head_balance", False): log("Enabling Sparge block-sparse head balancing...") - if self.config.enable_slicing: - log("Enabling VAE slicing...") - self.pipe.vae.enable_slicing() - - if self.config.enable_tiling: - log("Enabling VAE tiling...") - self.pipe.vae.enable_tiling() + self._vae_manager.enable_options(self._decoding_vaes()) if self.config.enable_sequential_cpu_offload: log("Enabling sequential CPU offload...") @@ -346,11 +341,17 @@ def _get_runtime_state_pipeline(self): return self.pipe + def _decoding_vaes(self) -> List: + """Forward staged VAE discovery to the VAE manager.""" + return self._vae_manager.decoding_vaes( + [self.pipe, getattr(self, "second_pipe", None)] + ) + def _validate_config(self, config: xFuserArgs) -> None: """ Validate if the model supports requested config """ for key in ModelCapabilities.__annotations__.keys(): config_value = getattr(config, key, None) # Some config options might not be set in the CLI, such as support for specific attention backends. - if isinstance(config_value, int): + if isinstance(config_value, int) and not isinstance(config_value, bool): if not getattr(self.capabilities, key) and config_value > 1: raise ValueError(f"Model {self.settings.model_name} does not support {key}.") else: @@ -413,7 +414,7 @@ def _validate_config(self, config: xFuserArgs) -> None: if not possible_task and self.settings.valid_tasks: raise ValueError(f"Model {self.settings.model_name} requires a task to be specified. Supported tasks: {self.settings.valid_tasks}") if config.dataset_path and not config.batch_size: - raise ValueError(f"Dataset path specified without batch size. Please specify batch size for dataset inference.") + raise ValueError("Dataset path specified without batch size. Please specify batch size for dataset inference.") if self.model_output_type == "video" and not self.fps: raise ValueError(f"Model {self.settings.model_name} produces video output but fps is not set.") @@ -434,18 +435,12 @@ def _validate_config(self, config: xFuserArgs) -> None: f"NVFP4 GEMMs require CUDA capability >= 10.0 (Blackwell). " f"Detected: {torch.cuda.get_device_capability()}" ) - if config.use_parallel_vae: - if not packages_info.get("has_distvae", False): - raise ValueError("DistVAE is not installed. Please install it before using parallel VAE.") - if torch.nn.GroupNorm.__module__ == "aiter.ops.groupnorm": - log("AITER GroupNorm is not supported with parallel VAE. Reverting to torch GroupNorm.") - torch.nn.GroupNorm = _TORCH_GROUPNORM + validate_vae_config(config, self.capabilities, self.settings) if config.distilled_transformer_path or config.distilled_transformer_2_path: if not self.capabilities.supports_distilled_weights: raise ValueError(f"Model {self.settings.model_name} does not support distilled_transformer_path or distilled_transformer_2_path params.") - def _get_compile_mode(self) -> str: # Overrides should return "default" when PACKAGES_CHECKER._on_rdna4(): # CUDA graphs are slow on RDNA4. @@ -584,7 +579,7 @@ def _run_warmup_calls(self, input_args: dict) -> None: for iteration in range(self.config.warmup_calls): log(f"Warmup iteration {iteration + 1}/{self.config.warmup_calls}") self._run_timed_pipe(input_args) - log(f"Warmup complete.") + log("Warmup complete.") def profile(self, input_args: dict) -> Tuple[DiffusionOutput, list, torch.profiler.profiler.profile]: """ Profile the model execution """ @@ -665,7 +660,7 @@ def save_output(self, output: DiffusionOutput) -> None: export_to_video(video, output_path, fps=self.settings.fps) log(f"Output video saved to {output_path}") else: - raise NotImplementedError(f"No output to save.") + raise NotImplementedError("No output to save.") def save_timings(self, timings: list) -> None: timing_file_name = f"{self.config.output_directory}/timings.json" @@ -678,13 +673,14 @@ def save_profile(self, profile: torch.profiler.profiler.profile) -> None: profile.export_chrome_trace(profile_file) log(f"Profile trace saved to {profile_file}", log_from_all_processes=True) - def _prepare_inference_run(self, input_args: dict) -> None: - """Prepare model-specific state before a pipeline invocation.""" + def prepare_run(self, input_args: dict) -> None: + """Prepare model state before a pipeline invocation.""" + self._vae_manager.prepare_run(self._decoding_vaes(), input_args) def _run_timed_pipe(self, input_args: dict) -> Tuple[DiffusionOutput, float]: """ Run a a full pipeline with timing information """ - self._prepare_inference_run(input_args) + self.prepare_run(input_args) start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) torch.cuda.synchronize() @@ -1049,24 +1045,8 @@ def _setup_hybrid_gemm_schedule(self, input_args: dict) -> None: get_runtime_state().set_gemm_schedule(gemm_schedule, total_steps=total_steps) def _convert_vae_to_channels_last(self) -> None: - """ Convert the VAE to channels last """ - convert_model_convs_to_channels_last(self.pipe.vae) - - original_decode = self.pipe.vae.decode - memory_format = torch.channels_last if self.settings.model_output_type == "image" else torch.channels_last_3d - - @functools.wraps(original_decode) - def decode_wrapper(*args, **kwargs): - if args: - args = list(args) - args[0] = args[0].to(memory_format=memory_format) - args = tuple(args) - elif "z" in kwargs: - kwargs["z"] = kwargs["z"].to(memory_format=memory_format) - output = original_decode(*args, **kwargs) - return output - - self.pipe.vae.decode = decode_wrapper + """Forward channels-last conversion for subclass compatibility.""" + self._vae_manager.convert_to_channels_last(self._decoding_vaes()) @abc.abstractmethod def _run_pipe(self, input_args: dict) -> DiffusionOutput: diff --git a/xfuser/model_executor/models/runner_models/causal_wan.py b/xfuser/model_executor/models/runner_models/causal_wan.py index ba0c363a..ae72b470 100644 --- a/xfuser/model_executor/models/runner_models/causal_wan.py +++ b/xfuser/model_executor/models/runner_models/causal_wan.py @@ -1,4 +1,3 @@ -import copy import os from typing import TYPE_CHECKING @@ -32,7 +31,8 @@ class xFuserCausalWanModel(xFuserModel): ring_degree=False, fully_shard_degree=True, use_fp8_gemms=False, - use_parallel_vae=False, + use_parallel_vae=True, + use_parallel_vae_encoder=True, enable_tiling=True, enable_slicing=True, ) diff --git a/xfuser/model_executor/models/runner_models/cosmos3.py b/xfuser/model_executor/models/runner_models/cosmos3.py index 8c79c856..5128ec2e 100644 --- a/xfuser/model_executor/models/runner_models/cosmos3.py +++ b/xfuser/model_executor/models/runner_models/cosmos3.py @@ -3,7 +3,6 @@ from diffusers import UniPCMultistepScheduler from diffusers.pipelines.pipeline_utils import DiffusionPipeline -from xfuser import xFuserArgs from xfuser.model_executor.models.runner_models.base_model import ( DIFFUSERS_FROM_SOURCE, ModelSettings, @@ -14,7 +13,6 @@ DiffusionOutput, _parse_attention_backend, ) -from xfuser.core.distributed.parallel_state import get_vae_parallel_group from xfuser.core.distributed.attention_backend import AttentionBackendType from xfuser.core.utils.runner_utils import log, resize_and_crop_image @@ -41,49 +39,6 @@ } -def _setup_parallel_vae(vae, enable_parallel_encoder=True): - if enable_parallel_encoder: - try: - from distvae.modules.adapters.vae.encoder_adapters import WanEncoderAdapter - # scale_factor_spatial includes patch_size (e.g. 16 = 8x from encoder + 2x - # from patching). The encoder adapter needs just the encoder's own - # downsampling ratio, excluding the patching step. - vae_scale_factor = getattr(vae.config, 'scale_factor_spatial', 8) - patch_size = getattr(vae.config, 'patch_size', None) - if patch_size and patch_size > 1: - vae_scale_factor = vae_scale_factor // patch_size - patched_encoder = WanEncoderAdapter( - vae.encoder, - vae_group=get_vae_parallel_group().device_group, - vae_scale_factor=vae_scale_factor, - ).to(vae.device) - vae.encoder = patched_encoder - log("Parallel VAE encoder enabled.") - except ImportError: - log("DistVAE not available for encoder. Defaulting to single-rank.") - except Exception as e: - raise ValueError(f"Failed to patch VAE encoder: {e}") - try: - from distvae.modules.adapters.vae.decoder_adapters import WanDecoderAdapter - patched_decoder = WanDecoderAdapter( - vae.decoder, vae_group=get_vae_parallel_group().device_group - ).to(vae.device) - # Cosmos3 VAE has patch_size=2 (extra 2x spatial upsampling from - # unpatching). The decoder adapter's scale_factor defaults to 1 in - # its Patchify, which is correct for Wan (patch_size=None). For - # Cosmos3 we need to account for the extra factor so that the - # narrow in _forward crops to the right size. - patch_size = getattr(vae.config, 'patch_size', None) - if patch_size and patch_size > 1: - patched_decoder.patchify.scale_factor = patch_size - vae.decoder = patched_decoder - log("Parallel VAE decoder enabled.") - except ImportError: - log("DistVAE not available for decoder. Defaulting to single-rank.") - except Exception as e: - raise ValueError(f"Failed to patch VAE decoder: {e}") - - @register_model("nvidia/Cosmos3-Super") @register_model("Cosmos3-Super") class xFuserCosmos3SuperModel(xFuserModel): @@ -221,8 +176,6 @@ def _post_load_and_state_initialization(self, input_args: dict) -> None: if self.config.fully_shard_degree > 1: if hasattr(self.pipe.transformer, '_patch_time_embedder_for_fsdp'): self.pipe.transformer._patch_time_embedder_for_fsdp() - if self.config.use_parallel_vae: - _setup_parallel_vae(self.pipe.vae, self.capabilities.use_parallel_vae_encoder) @register_model("nvidia/Cosmos3-Nano") diff --git a/xfuser/model_executor/models/runner_models/flux.py b/xfuser/model_executor/models/runner_models/flux.py index 2e8d32db..ae2d8a6a 100644 --- a/xfuser/model_executor/models/runner_models/flux.py +++ b/xfuser/model_executor/models/runner_models/flux.py @@ -13,30 +13,9 @@ from xfuser.core.utils.runner_utils import ( log, resize_and_crop_image, - quantize_linear_layers_to_fp8, ) from xfuser.core.distributed import get_runtime_state, get_pipeline_parallel_world_size -from xfuser.core.distributed.parallel_state import get_vae_parallel_group -from xfuser import xFuserFluxPipeline, xFuserArgs - - -def _setup_parallel_vae(vae) -> None: - """Parallalizes the VAE decoder using distvae""" - try: - from distvae.modules.adapters.vae.decoder_adapters import DecoderAdapter - - patched_decoder = DecoderAdapter( - vae.decoder, vae_group=get_vae_parallel_group().device_group - ).to(vae.device) - vae.decoder = patched_decoder - log(f"Parallel VAE decoder enabled successfully.") - except ImportError: - raise ValueError( - "DistVAE library is missing or does not support DecoderAdapter. " - "Try installing latest DistVAE from https://github.com/xdit-project/DistVAE." - ) - except Exception as e: - raise ValueError(f"Failed to patch VAE decoder. {e}") +from xfuser import xFuserFluxPipeline @register_model("black-forest-labs/FLUX.1-dev") @@ -80,11 +59,6 @@ class xFuserFluxModel(xFuserModel): }, ) - def _post_load_and_state_initialization(self, input_args: dict) -> None: - super()._post_load_and_state_initialization(input_args) - if self.config.use_parallel_vae: - _setup_parallel_vae(self.pipe.vae) - def _get_compile_mode(self) -> str: if PACKAGES_CHECKER._on_rdna4(): return "default" @@ -149,6 +123,7 @@ class xFuserFluxKontextModel(xFuserModel): enable_tiling=True, enable_slicing=True, use_parallel_vae=True, + use_parallel_vae_encoder=True, fully_shard_degree=True, ) default_input_values = DefaultInputValues( @@ -177,11 +152,6 @@ class xFuserFluxKontextModel(xFuserModel): }, ) - def _post_load_and_state_initialization(self, input_args: dict) -> None: - super()._post_load_and_state_initialization(input_args) - if self.config.use_parallel_vae: - _setup_parallel_vae(self.pipe.vae) - def _load_model(self) -> DiffusionPipeline: from diffusers import FluxKontextPipeline from xfuser.model_executor.models.transformers.transformer_flux import ( @@ -265,6 +235,7 @@ class xFuserFlux2Model(xFuserModel): enable_tiling=True, enable_slicing=True, use_parallel_vae=True, + use_parallel_vae_encoder=True, use_fbcache=True, pipefusion_parallel_degree=True, ) @@ -301,8 +272,6 @@ class xFuserFlux2Model(xFuserModel): def _post_load_and_state_initialization(self, input_args: dict) -> None: super()._post_load_and_state_initialization(input_args) - if self.config.use_parallel_vae: - _setup_parallel_vae(self.pipe.vae) if self.config.use_fbcache: from xfuser.model_executor.cache.diffusers_adapters.flux2 import ( @@ -414,6 +383,7 @@ class xFuserFlux2Klein9BModel(xFuserModel): enable_tiling=True, enable_slicing=True, use_parallel_vae=True, + use_parallel_vae_encoder=True, fully_shard_degree=True, use_fbcache=True, pipefusion_parallel_degree=True, @@ -443,11 +413,6 @@ class xFuserFlux2Klein9BModel(xFuserModel): }, ) - def _post_load_and_state_initialization(self, input_args: dict) -> None: - super()._post_load_and_state_initialization(input_args) - if self.config.use_parallel_vae: - _setup_parallel_vae(self.pipe.vae) - def _get_compile_mode(self) -> str: # CUDA graphs incompatible with FBCache cross-step caching, # and cause pathological re-captures on RDNA4. diff --git a/xfuser/model_executor/models/runner_models/hunyuan.py b/xfuser/model_executor/models/runner_models/hunyuan.py index cc0318d8..f10234b2 100644 --- a/xfuser/model_executor/models/runner_models/hunyuan.py +++ b/xfuser/model_executor/models/runner_models/hunyuan.py @@ -15,13 +15,11 @@ DiffusionOutput, ModelSettings, ) -from xfuser.core.distributed.attention_backend import AttentionBackendType from xfuser.core.distributed.runtime_state import get_runtime_state from xfuser.core.utils.runner_utils import ( resize_and_crop_image, fix_llama_tokenizer_pretokenizer, ) -from xfuser.envs import PACKAGES_CHECKER from xfuser.compile import install_inductor_passes @register_model("tencent/HunyuanVideo") @@ -39,6 +37,7 @@ class xFuserHunyuanvideoModel(xFuserModel): enable_tiling=True, use_hybrid_attn_schedule=True, use_fp8_gemms=True, + use_parallel_vae=True, ) default_input_values = DefaultInputValues( height=720, @@ -121,6 +120,8 @@ class xFuserHunyuanvideo15Model(xFuserModel): enable_slicing=True, enable_tiling=True, use_fp8_gemms=True, + use_parallel_vae=True, + use_parallel_vae_encoder=True, ) default_input_values = DefaultInputValues( height=720, @@ -275,6 +276,8 @@ class xFuserHunyuanvideo15SparseModel(xFuserHunyuanvideo15Model): enable_slicing=True, enable_tiling=True, supports_sparse_attention_backends=True, + use_parallel_vae=True, + use_parallel_vae_encoder=True, ) def _validate_ssta_attention_kwargs(self, attn_param: dict) -> None: diff --git a/xfuser/model_executor/models/runner_models/ideogram4.py b/xfuser/model_executor/models/runner_models/ideogram4.py index b58ea51f..12c66422 100644 --- a/xfuser/model_executor/models/runner_models/ideogram4.py +++ b/xfuser/model_executor/models/runner_models/ideogram4.py @@ -6,7 +6,6 @@ import torch from diffusers.pipelines.pipeline_utils import DiffusionPipeline -from xfuser.core.distributed.parallel_state import get_vae_parallel_group from xfuser.core.utils.runner_utils import log from xfuser.model_executor.models.runner_models.base_model import ( DefaultInputValues, @@ -16,9 +15,6 @@ register_model, xFuserModel, ) -from xfuser.model_executor.models.transformers.transformer_ideogram4 import ( - get_ideogram4_transformer_wrapper_class, -) from xfuser.model_executor.pipelines.pipeline_ideogram4 import ( get_ideogram4_pipeline_class, ) @@ -164,21 +160,6 @@ def _check_load_result( ) -def _setup_parallel_vae(vae) -> None: - try: - from distvae.modules.adapters.vae.decoder_adapters import DecoderAdapter - - vae.decoder = DecoderAdapter( - vae.decoder, - vae_group=get_vae_parallel_group().device_group, - ).to(vae.device) - log("Parallel VAE decoder enabled.") - except ImportError: - log("DistVAE not available for decoder. Defaulting to single-rank.") - except Exception as error: - raise ValueError(f"Failed to patch VAE decoder: {error}") from error - - def _default_guidance_schedule(num_inference_steps: int) -> list[float]: polish_steps = min(3, num_inference_steps) return [7.0] * (num_inference_steps - polish_steps) + [3.0] * polish_steps @@ -251,6 +232,10 @@ def _load_fp8_transformer( model_id: str, subfolder: str, ): + from xfuser.model_executor.models.transformers.transformer_ideogram4 import ( + get_ideogram4_transformer_wrapper_class, + ) + transformer_class = get_ideogram4_transformer_wrapper_class() transformer = transformer_class.from_config( transformer_class.load_config(model_id, subfolder=subfolder) @@ -296,6 +281,10 @@ def _load_fp8_text_encoder(self, model_id: str): return text_encoder def _load_model(self) -> DiffusionPipeline: + from xfuser.model_executor.models.transformers.transformer_ideogram4 import ( + get_ideogram4_transformer_wrapper_class, + ) + model_id = self.config.model transformer_class = get_ideogram4_transformer_wrapper_class() @@ -388,5 +377,3 @@ def _post_load_and_state_initialization(self, input_args: dict) -> None: super()._post_load_and_state_initialization(input_args) self.pipe.transformer._init_sp_state() self.pipe.unconditional_transformer._init_sp_state() - if self.config.use_parallel_vae: - _setup_parallel_vae(self.pipe.vae) diff --git a/xfuser/model_executor/models/runner_models/krea2.py b/xfuser/model_executor/models/runner_models/krea2.py index e66c6454..f3a52ef3 100644 --- a/xfuser/model_executor/models/runner_models/krea2.py +++ b/xfuser/model_executor/models/runner_models/krea2.py @@ -79,7 +79,7 @@ class _Krea2BaseModel(xFuserModel): pipefusion_parallel_degree=False, data_parallel_degree=True, use_cfg_parallel=False, - use_parallel_vae=False, + use_parallel_vae=True, use_fp8_gemms=True, use_fp4_gemms=True, use_hybrid_attn_schedule=True, diff --git a/xfuser/model_executor/models/runner_models/lingbot_video.py b/xfuser/model_executor/models/runner_models/lingbot_video.py index 03d3b5ce..94141a45 100644 --- a/xfuser/model_executor/models/runner_models/lingbot_video.py +++ b/xfuser/model_executor/models/runner_models/lingbot_video.py @@ -6,10 +6,6 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor -from xfuser import xFuserArgs -from xfuser.model_executor.models.transformers.transformer_lingbot_video import ( - xFuserLingBotVideoTransformer3DWrapper, -) from xfuser.model_executor.pipelines.pipeline_lingbot_video import ( xFuserLingBotVideoPipeline, get_lingbot_video_pipeline_class, @@ -22,8 +18,6 @@ DefaultInputValues, DiffusionOutput, ) -from xfuser.core.distributed.runtime_state import get_runtime_state -from xfuser.core.distributed.parallel_state import get_vae_parallel_group from xfuser.core.utils.runner_utils import log @@ -78,32 +72,6 @@ def _load_json_prompt(prompt: str) -> str: return prompt -def _setup_parallel_vae(vae, use_encoder=False): - import torch.distributed as dist - vae_group = get_vae_parallel_group().device_group - log(f"VAE parallel group: world_size={dist.get_world_size(vae_group)}, " - f"rank={dist.get_rank(vae_group)}, vae.device={vae.device}") - try: - from distvae.modules.adapters.vae.decoder_adapters import WanDecoderAdapter - vae.decoder = WanDecoderAdapter(vae.decoder, vae_group=vae_group).to(vae.device) - log("Parallel VAE decoder enabled.") - except ImportError: - log("distvae WanDecoderAdapter not available, skipping parallel VAE decoder.") - return - except Exception as e: - log(f"Failed to patch VAE decoder: {e}") - return - if use_encoder: - try: - from distvae.modules.adapters.vae.encoder_adapters import WanEncoderAdapter - vae.encoder = WanEncoderAdapter(vae.encoder, vae_group=vae_group).to(vae.device) - log("Parallel VAE encoder enabled.") - except ImportError: - log("distvae WanEncoderAdapter not available, skipping parallel VAE encoder.") - except Exception as e: - log(f"Failed to patch VAE encoder: {e}") - - @register_model("robbyant/lingbot-video-moe-30b-a3b") @register_model("LingBot-Video-MoE") class xFuserLingBotVideoMoEModel(xFuserModel): @@ -139,6 +107,7 @@ def _calculate_hybrid_attention_step_multiplier(self, input_args: dict) -> int: use_hybrid_gemm_schedule=True, fully_shard_degree=True, use_parallel_vae=True, + use_parallel_vae_encoder=True, enable_tiling=True, enable_slicing=True, ) @@ -198,9 +167,6 @@ def _post_load_and_state_initialization(self, input_args: dict) -> None: for block in self.pipe.transformer.blocks: if hasattr(block, "_cached_bulk_dtype"): _patch_block_bulk_dtype(block) - if self.config.use_parallel_vae: - _setup_parallel_vae(self.pipe.vae) - # Cache pre-transposed expert weights to eliminate per-call copies self.pipe.transformer.cache_expert_weights() @@ -212,6 +178,9 @@ def _post_load_and_state_initialization(self, input_args: dict) -> None: def _load_refiner(self, input_args): from lingbot_video.scheduling_flow_unipc import FlowUniPCMultistepScheduler + from xfuser.model_executor.models.transformers.transformer_lingbot_video import ( + xFuserLingBotVideoTransformer3DWrapper, + ) log("Loading refiner transformer...") model_name = self.settings.model_name @@ -262,12 +231,6 @@ def _load_refiner(self, input_args): for block in refiner_transformer.blocks: if hasattr(block, "_cached_bulk_dtype"): _patch_block_bulk_dtype(block) - # Enable VAE tiling/slicing for refiner (1080p needs it) - if self.config.enable_tiling: - refiner_pipe.vae.enable_tiling() - if self.config.enable_slicing: - refiner_pipe.vae.enable_slicing() - # FSDP shard the refiner transformer if enabled if self.config.fully_shard_degree > 1: from xfuser.core.distributed.parallel_state import get_fs_group @@ -286,6 +249,9 @@ def _load_refiner(self, input_args): def _build_pipe(self, model_name, transformer_subfolder="transformer", use_i2v=False): from lingbot_video.scheduling_flow_unipc import FlowUniPCMultistepScheduler + from xfuser.model_executor.models.transformers.transformer_lingbot_video import ( + xFuserLingBotVideoTransformer3DWrapper, + ) transformer = xFuserLingBotVideoTransformer3DWrapper.from_pretrained( model_name, torch_dtype=torch.bfloat16, subfolder=transformer_subfolder, @@ -480,15 +446,12 @@ def _compile_model(self, input_args): @register_model("robbyant/lingbot-video-dense-1.3b") @register_model("LingBot-Video-Dense") class xFuserLingBotVideoDenseModel(xFuserLingBotVideoMoEModel): - - def __init__(self, config: xFuserArgs) -> None: - super().__init__(config) - self.settings = ModelSettings( - model_name="robbyant/lingbot-video-dense-1.3b", - output_name="lingbot_video_dense", - model_output_type="video", - fps=24, - fp8_gemm_module_list=["transformer.blocks"], - fp4_gemm_module_list=["transformer.blocks"], - fsdp_strategy=LINGBOT_FSDP_STRATEGY, - ) + settings = ModelSettings( + model_name="robbyant/lingbot-video-dense-1.3b", + output_name="lingbot_video_dense", + model_output_type="video", + fps=24, + fp8_gemm_module_list=["transformer.blocks"], + fp4_gemm_module_list=["transformer.blocks"], + fsdp_strategy=LINGBOT_FSDP_STRATEGY, + ) diff --git a/xfuser/model_executor/models/runner_models/ltx.py b/xfuser/model_executor/models/runner_models/ltx.py index fd337ae7..b5209c57 100644 --- a/xfuser/model_executor/models/runner_models/ltx.py +++ b/xfuser/model_executor/models/runner_models/ltx.py @@ -58,6 +58,7 @@ class xFuserLTX23VideoModel(xFuserModel): ring_degree=True, enable_tiling=True, enable_slicing=True, + use_parallel_vae=True, ) _STG_SCALE = 1.0 @@ -108,6 +109,7 @@ def _load_model(self) -> DiffusionPipeline: ) upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=latent_upsampler) + log("Enabling VAE tiling for LTX-2.3's full-resolution second-stage decode.") second_pipe.vae.enable_tiling() self.second_pipe = second_pipe @@ -115,11 +117,6 @@ def _load_model(self) -> DiffusionPipeline: return pipe - def _enable_options(self) -> None: - super()._enable_options() - if self.config.enable_slicing: - self.second_pipe.vae.enable_slicing() - def _run_pipe(self, input_args: dict) -> DiffusionOutput: from diffusers.pipelines.ltx2.utils import STAGE_2_DISTILLED_SIGMA_VALUES generator = torch.Generator(device="cuda").manual_seed(input_args["seed"]) @@ -221,8 +218,6 @@ def _post_load_and_state_initialization(self, input_args: dict) -> None: super()._post_load_and_state_initialization(input_args) self.upsample_pipe.to(self.pipe.device) self.second_pipe.to(self.pipe.device) - - @register_model("Lightricks/LTX-2") @register_model("LTX-2") class xFuserLTX2VideoModel(xFuserModel): @@ -251,6 +246,7 @@ class xFuserLTX2VideoModel(xFuserModel): enable_tiling=True, enable_slicing=True, use_fp8_gemms=True, + use_parallel_vae=True, ) def _load_model(self) -> DiffusionPipeline: @@ -292,13 +288,6 @@ def _load_model(self) -> DiffusionPipeline: return pipe - def _enable_options(self) -> None: - super()._enable_options() - if self.config.enable_tiling: - self.second_pipe.vae.enable_tiling() - if self.config.enable_slicing: - self.second_pipe.vae.enable_slicing() - def _run_pipe(self, input_args: dict) -> DiffusionOutput: from diffusers.pipelines.ltx2.utils import STAGE_2_DISTILLED_SIGMA_VALUES video_latent, audio_latent = self.pipe( diff --git a/xfuser/model_executor/models/runner_models/qwen.py b/xfuser/model_executor/models/runner_models/qwen.py index ef73c52c..5366ec09 100644 --- a/xfuser/model_executor/models/runner_models/qwen.py +++ b/xfuser/model_executor/models/runner_models/qwen.py @@ -25,6 +25,8 @@ class xFuserQwenImageEditModel(xFuserModel): ring_degree=True, fully_shard_degree=True, use_fp8_gemms=True, + use_parallel_vae=True, + use_parallel_vae_encoder=True, enable_tiling=True, enable_slicing=True, ) @@ -110,6 +112,9 @@ class xFuserQwenImageModel(xFuserModel): ring_degree=True, fully_shard_degree=True, use_fp8_gemms=True, + use_parallel_vae=True, + enable_tiling=True, + enable_slicing=True, ) default_input_values = DefaultInputValues( height=928, diff --git a/xfuser/model_executor/models/runner_models/stable_diffusion.py b/xfuser/model_executor/models/runner_models/stable_diffusion.py index 470283b2..e718d1c2 100644 --- a/xfuser/model_executor/models/runner_models/stable_diffusion.py +++ b/xfuser/model_executor/models/runner_models/stable_diffusion.py @@ -1,6 +1,6 @@ import torch from diffusers.pipelines.pipeline_utils import DiffusionPipeline -from xfuser import xFuserStableDiffusion3Pipeline, xFuserArgs +from xfuser import xFuserStableDiffusion3Pipeline from xfuser.model_executor.models.runner_models.base_model import ( xFuserModel, register_model, @@ -24,6 +24,7 @@ class xFuserStableDiffusionModel(xFuserModel): enable_slicing=True, fully_shard_degree=True, use_fp8_gemms=True, + use_parallel_vae=True, ) default_input_values = DefaultInputValues( height=1024, diff --git a/xfuser/model_executor/models/runner_models/vae_manager.py b/xfuser/model_executor/models/runner_models/vae_manager.py new file mode 100644 index 00000000..4c16fe4e --- /dev/null +++ b/xfuser/model_executor/models/runner_models/vae_manager.py @@ -0,0 +1,439 @@ +"""VAE discovery, configuration, and runtime orchestration for runner models.""" + +import functools +from typing import List, Optional, Tuple + +import torch +from distvae.vae import ParallelContext, VAERowSplitError, latent_rows, mark +from distvae.vae import parallel as vae_parallel +from distvae.vae import tile_parallel as vae_tile_parallel +from distvae.vae import tiling as vae_tiling + +from xfuser.core.distributed.parallel_state import ( + get_vae_parallel_group, + get_vae_parallel_world_size, +) +from xfuser.core.utils.runner_utils import ( + convert_model_convs_to_channels_last, + log, +) +from xfuser.envs import restore_torch_group_norm_for_distvae + + +def _validate_vae_tile_pair( + config, capabilities, settings, kind, minimum, constraint +) -> None: + names = tuple(f"vae_tile_{kind}_{axis}" for axis in ("height", "width")) + flags = tuple(f"--{name}" for name in names) + values = tuple(getattr(config, name, None) for name in names) + asked = tuple(value is not None for value in values) + joined_flags = " and ".join(flags) + + if asked[0] != asked[1]: + raise ValueError(f"{joined_flags} must be provided together.") + for value, flag in zip(values, flags): + if value is not None and value < minimum: + raise ValueError(f"{flag} must be {constraint}, got {value}.") + if not asked[0]: + return + if not config.enable_tiling: + raise ValueError(f"{joined_flags} require --enable_tiling.") + if not capabilities.enable_tiling: + raise ValueError( + f"{joined_flags} configure tiled VAE decoding, " + f"which model {settings.model_name} does not support." + ) + + +def validate_vae_config(config, capabilities, settings) -> None: + """Validate VAE-specific runner configuration without runner state.""" + if config.use_parallel_vae: + if restore_torch_group_norm_for_distvae(): + log( + "AITER GroupNorm cannot be sharded. Restoring torch GroupNorm so DistVAE can " + "identify and shard the GroupNorm layers." + ) + + _validate_vae_tile_pair( + config, capabilities, settings, "size", 1, "positive" + ) + _validate_vae_tile_pair( + config, capabilities, settings, "overlap", 0, "non-negative" + ) + + +class VAEManager: + """Owns all VAE-specific policy used by runner-model lifecycles.""" + + def __init__(self, config, capabilities, settings) -> None: + self.config = config + self.capabilities = capabilities + self.settings = settings + self._overlap_sample_shapes = {} + + def decoding_vaes(self, pipes) -> List: + """Return every unique VAE decoded by the supplied pipeline stages.""" + vaes = [] + for candidate in pipes: + vae = getattr(candidate, "vae", None) + if vae is not None and not any(vae is seen for seen in vaes): + vaes.append(vae) + return vaes + + def enable_options(self, vaes) -> None: + """Apply slicing, tiling, exact plans, and decode wrappers.""" + tiling_flag = self._tiling_flag() + for vae in vaes: + if self.config.enable_slicing: + vae_tiling.require_vae_support( + vae, "slicing", "--enable_slicing" + ) + log(f"Enabling VAE slicing on {type(vae).__name__}...") + vae.enable_slicing() + + applied_shape = None + if self._tiles(vae): + native_shape = vae_tiling.tile_shape(vae) + if tiling_flag is not None: + vae_tiling.require_vae_support(vae, "tiling", tiling_flag) + log(f"Enabling VAE tiling on {type(vae).__name__}...") + vae.enable_tiling() + applied_shape = self._apply_vae_tile_shape(vae) + self._check_tiles_against_parallel_vae(vae, native_shape) + self._install_vae_tiled_decode(vae) + self._install_vae_decode_guard(vae, applied_shape) + + def prepare_run(self, vaes, input_args) -> None: + """Apply shape-dependent VAE options for the current invocation.""" + if self._requested_vae_tile_overlap() is None: + return + sample_shape = (input_args["height"], input_args["width"]) + for vae in vaes: + key = id(vae) + if self._overlap_sample_shapes.get(key) == sample_shape: + continue + self._apply_vae_tile_overlap(vae, sample_shape) + self._overlap_sample_shapes[key] = sample_shape + + def _tiling_flag(self) -> Optional[str]: + """Return the flag asking this run to tile its VAE decode.""" + return "--enable_tiling" if self.config.enable_tiling else None + + def _tiles(self, vae) -> bool: + """Return whether this VAE's decode will be cut into tiles.""" + return self._tiling_flag() is not None or getattr(vae, "use_tiling", False) + + def setup_parallel_vae(self, vaes) -> None: + """Shard VAE decode, and capability-selected encode, across the VAE group.""" + coordinator = get_vae_parallel_group() + vae_group = coordinator.device_group + tile_context = ParallelContext( + group=vae_group, + rank=coordinator.rank_in_group, + world_size=coordinator.world_size, + patch_dim=-2, + global_ranks=tuple(coordinator.ranks), + ) + log( + f"VAE parallel group: world_size={coordinator.world_size}, " + f"rank={coordinator.rank_in_group}", + debug=True, + ) + for vae in vaes: + if self._tiles(vae) and vae_tiling.supports_tile_parallel(vae): + mark(vae, tile_context) + log( + "Parallel VAE will assign complete tiles of " + f"{type(vae).__name__} to ranks instead of sharding rows within each tile." + ) + else: + adapter = vae_parallel.parallelize_decoder(vae, vae_group) + log( + f"Parallel VAE decoder enabled on {type(vae).__name__} via {adapter}." + ) + if self.capabilities.use_parallel_vae_encoder: + adapter = vae_parallel.parallelize_encoder(vae, vae_group) + log( + f"Parallel VAE encoder enabled on {type(vae).__name__} via {adapter}." + ) + + def convert_to_channels_last(self, vaes) -> None: + """Convert every supplied decoding VAE to channels-last exactly once.""" + for vae in vaes: + self._convert_one_vae_to_channels_last(vae) + + def _convert_one_vae_to_channels_last(self, vae) -> None: + if getattr(vae, "_xfuser_decode_channels_last", False): + return + convert_model_convs_to_channels_last(vae) + + original_decode = vae.decode + memory_format = ( + torch.channels_last + if self.settings.model_output_type == "image" + else torch.channels_last_3d + ) + + @functools.wraps(original_decode) + def decode_wrapper(*args, **kwargs): + if args: + args = list(args) + args[0] = args[0].to(memory_format=memory_format) + args = tuple(args) + elif "z" in kwargs: + kwargs["z"] = kwargs["z"].to(memory_format=memory_format) + return original_decode(*args, **kwargs) + + vae.decode = decode_wrapper + vae._xfuser_decode_channels_last = True + + def _requested_vae_tile_shape(self) -> Optional[Tuple[int, int]]: + height = getattr(self.config, "vae_tile_size_height", None) + width = getattr(self.config, "vae_tile_size_width", None) + if height is None or width is None: + return None + return height, width + + def _requested_vae_tile_overlap(self) -> Optional[Tuple[int, int]]: + height = getattr(self.config, "vae_tile_overlap_height", None) + width = getattr(self.config, "vae_tile_overlap_width", None) + if height is None or width is None: + return None + return height, width + + def _apply_vae_tile_shape(self, vae) -> Optional[Tuple[int, int]]: + """Apply and return the exact requested shape, or retain the default tile shape.""" + shape = self._requested_vae_tile_shape() + if shape is None: + return None + height, width = shape + plan = vae_tiling.tile_shape_plan(vae, height, width) + if plan is None: + native = vae_tiling.tile_shape(vae) + native_shown = ( + f" The VAE's default tile shape is {native[0]}x{native[1]}." + if native is not None + else "" + ) + raise ValueError( + f"--vae_tile_size_height {height} with --vae_tile_size_width {width} " + f"is not a shape this VAE ({type(vae).__name__}) can tile exactly." + f"{native_shown} Choose dimensions compatible with the VAE's latent " + "scale and tile stride." + ) + vae_tiling.apply_tile_plan(vae, plan) + log( + f"VAE tile window set to {height}x{width}px " + f"({', '.join(f'{a}={v}' for a, v in sorted(plan.items()))})" + ) + return shape + + def _apply_vae_tile_overlap( + self, vae, sample_shape: Tuple[int, int] + ) -> None: + """Apply exact per-axis output-pixel tile overlap when requested.""" + requested = self._requested_vae_tile_overlap() + if requested is None: + return + overlap_height, overlap_width = requested + plan = vae_tiling.tile_overlap_plan( + vae, + overlap_height, + overlap_width, + sample_shape=sample_shape, + ) + if plan is None: + shape = vae_tiling.tile_shape(vae) + shape_shown = ( + f"{shape[0]}x{shape[1]} pixels" + if shape is not None + else "an unknown pixel shape" + ) + raise ValueError( + f"The requested VAE tile overlap of {overlap_height}x{overlap_width} pixels " + f"is not exact for this VAE ({type(vae).__name__}) at its current tile shape " + f"of {shape_shown}. Height and width are output pixels; use 0 for an inactive " + "strip axis." + ) + vae_tiling.apply_tile_plan(vae, plan) + landed = vae_tiling.tile_overlap(vae) + shown = ( + f"{landed[0]}x{landed[1]}px" + if landed is not None + else f"{overlap_height}x{overlap_width}px" + ) + log( + f"VAE tile overlap set to {shown} " + f"({', '.join(f'{a}={v:g}' for a, v in sorted(plan.items()))})" + ) + + def _check_tiles_against_parallel_vae( + self, vae, native_shape: Optional[Tuple[int, int]] + ) -> None: + """Refuse a tile holding fewer latent rows than row-sharding ranks.""" + if not ( + self.config.use_parallel_vae and self.capabilities.use_parallel_vae + ): + return + if vae_tile_parallel.context_of(vae) is not None: + return + ranks = get_vae_parallel_world_size() + rows = latent_rows(vae) + if ranks < 2 or rows is None or rows >= ranks: + return + shape = vae_tiling.tile_shape(vae) + smallest = self._minimum_vae_tile_shape( + vae, shape, native_shape, ranks + ) + shown = ( + f"{shape[0]}x{shape[1]}px" if shape is not None else "unknown-size" + ) + native_shown = ( + f"{native_shape[0]}x{native_shape[1]}px" + if native_shape is not None + else "unknown" + ) + raise ValueError( + f"A {shown} VAE tile contains {rows} latent rows, but --use_parallel_vae uses " + f"{ranks} ranks and requires at least one latent row per rank" + + ( + f"; the smallest tile shape that satisfies this requirement is " + f"--vae_tile_size_height {smallest[0]} " + f"--vae_tile_size_width {smallest[1]}." + if smallest + else f". No tile shape up to the VAE's default {native_shown} shape has enough " + f"latent rows. Decode without tiling, increase --vae_tile_size_height and " + "--vae_tile_size_width, or use fewer VAE ranks." + ) + ) + + @staticmethod + def _minimum_vae_tile_shape( + vae, + shape: Optional[Tuple[int, int]], + native_shape: Optional[Tuple[int, int]], + min_latent_rows: int, + ) -> Optional[Tuple[int, int]]: + """Find the first exact plan with enough rows, preserving tile width.""" + if shape is None or native_shape is None: + return None + height, width = shape + native_height = native_shape[0] + if height > native_height: + return None + for candidate in range(height, native_height + 1): + plan = vae_tiling.tile_shape_plan(vae, candidate, width) + if plan is None: + continue + rows = latent_rows(vae, plan) + if rows is not None and rows >= min_latent_rows: + return candidate, width + return None + + def _install_vae_tiled_decode(self, vae) -> None: + """Install local or group-dispatched tiled decode when DistVAE provides one.""" + context = vae_tile_parallel.context_of(vae) + if context is None: + if ( + self._requested_vae_tile_shape() is None + and self._requested_vae_tile_overlap() is None + ): + return + installed = vae_tiling.tiled_decode_for(vae) + else: + dispatch, assemble = vae_tile_parallel.sharing(context) + installed = vae_tiling.tiled_decode_for(vae, dispatch, assemble) + if installed is None: + return + vae.tiled_decode = installed + if context is None: + log( + f"VAE tiled decode on {type(vae).__name__}: using DistVAE's local " + "overlap loop." + ) + return + log( + f"VAE tiled decode on {type(vae).__name__}: assigning contiguous groups of " + f"neighboring tiles across {context.world_size} ranks and blending the decoded " + "outputs." + ) + + def _install_vae_decode_guard( + self, vae, tile_shape: Optional[Tuple[int, int]] = None + ) -> None: + """Add actionable OOM and narrow-tile diagnostics to a VAE decode.""" + vae._xfuser_guarded_tile_shape = tile_shape + if getattr(vae, "_xfuser_decode_guarded", False): + return + original_decode = vae.decode + + @functools.wraps(original_decode) + def decode_guard(*args, **kwargs): + try: + return original_decode(*args, **kwargs) + except VAERowSplitError as e: + raise ValueError(self._vae_row_split_hint(vae, e)) from e + except torch.cuda.OutOfMemoryError as e: + raise torch.cuda.OutOfMemoryError( + f"{self._vae_decode_oom_hint(vae)}\n{e}" + ) from e + except RuntimeError as e: + shape = getattr(vae, "_xfuser_guarded_tile_shape", None) + if shape is None or not vae_tiling.is_tile_padding_error(e): + raise + shown = f"{shape[0]}x{shape[1]}" + raise RuntimeError( + f"VAE tiled decode failed with a {shown}px tile shape. At this output size, " + "an edge tile is too small for decoder padding. Try another exact height and " + "width, or remove --vae_tile_size_height and --vae_tile_size_width to use the " + f"VAE's default tile shape.\n{e}" + ) from e + + vae.decode = decode_guard + vae._xfuser_decode_guarded = True + + def _vae_row_split_hint(self, vae, error: VAERowSplitError) -> str: + message = ( + f"Cannot row-shard {error.rows} latent rows in the VAE decoder: this VAE " + f"processes rows in groups of {error.factor}, but {error.rows} is not divisible " + f"by {error.factor}." + ) + scale = getattr(getattr(vae, "config", None), "scale_factor_spatial", None) + actions = [ + ( + f"Use an output height divisible by {scale * error.factor}" + if isinstance(scale, int) and scale > 0 + else f"Use a latent height divisible by {error.factor}" + ) + ] + if self.capabilities.enable_tiling and vae_tiling.supports_tile_parallel(vae): + actions.append( + "enable --enable_tiling to distribute complete tiles instead" + ) + actions.append("disable --use_parallel_vae") + return f"{message} {', '.join(actions[:-1])}, or {actions[-1]}." + + def _vae_decode_oom_hint(self, vae) -> str: + if not getattr(vae, "use_tiling", False): + if self.capabilities.enable_tiling: + return ( + "VAE decode ran out of memory with tiling disabled. Re-run with " + "--enable_tiling to decode in tiles." + ) + return ( + f"VAE decode ran out of memory, and model {self.settings.model_name} does not " + "support VAE tiling." + ) + + shape = vae_tiling.tile_shape(vae) + if shape is None: + return ( + f"VAE tiled decode ran out of memory. The {type(vae).__name__} VAE used by " + f"{self.settings.model_name} does not expose an adjustable pixel-space tile shape." + ) + height, width = shape + return ( + f"VAE tiled decode ran out of memory with a {height}x{width}px tile shape. Choose " + "smaller exact values for --vae_tile_size_height and --vae_tile_size_width, then " + "rerun and compare peak VRAM." + ) diff --git a/xfuser/model_executor/models/runner_models/wan.py b/xfuser/model_executor/models/runner_models/wan.py index 004ab06f..3e1010c8 100644 --- a/xfuser/model_executor/models/runner_models/wan.py +++ b/xfuser/model_executor/models/runner_models/wan.py @@ -1,4 +1,3 @@ -import copy import re import torch from typing import List, Optional @@ -22,15 +21,11 @@ ) from xfuser.core.distributed.runtime_state import get_runtime_state from xfuser.core.distributed.attention_backend import AttentionBackendType -from xfuser.core.distributed.parallel_state import get_vae_parallel_group from xfuser.core.utils.runner_utils import ( log, resize_and_crop_image, resize_image_to_max_area, ) -from xfuser.envs import PACKAGES_CHECKER - - COMMON_FSDP_STRATEGY = { "transformer": { "wrap_attrs": ["blocks"], @@ -77,54 +72,13 @@ def _build_attention_kwargs(config: "xFuserArgs") -> dict: class xFuserWanModel(xFuserModel): """Common lifecycle hooks for Wan runners.""" - def _prepare_inference_run(self, input_args: dict) -> None: + def prepare_run(self, input_args: dict) -> None: + super().prepare_run(input_args) get_runtime_state().reset_vsa_schedule_state( int(input_args["num_inference_steps"]) ) -def _setup_parallel_vae(vae, enable_parallel_encoder: bool = True) -> None: - """ Parallelizes VAE en-/decoder using distvae """ - # Handle encoder - if enable_parallel_encoder: - try: - from distvae.modules.adapters.vae.encoder_adapters import WanEncoderAdapter - vae_scale_factor = getattr(vae.config, 'scaling_factor', 8) - if hasattr(vae.config, 'vae_scale_factor_spatial'): - vae_scale_factor = vae.config.vae_scale_factor_spatial - patched_encoder = WanEncoderAdapter( - vae.encoder, - vae_group=get_vae_parallel_group().device_group, - vae_scale_factor=vae_scale_factor, - ).to(vae.device) - vae.encoder = patched_encoder - log(f"Parallel VAE encoder enabled successfully.") - except ImportError: - log( - "DistVAE library is missing or does not support WanEncoderAdapter. " - "Try installing latest DistVAE from https://github.com/xdit-project/DistVAE. " - "Defaulting to single-rank encoder." - ) - except Exception as e: - raise ValueError(f"Failed to patch VAE encoder. {e}") - # Handle decoder - try: - from distvae.modules.adapters.vae.decoder_adapters import WanDecoderAdapter - patched_decoder = WanDecoderAdapter( - vae.decoder, vae_group=get_vae_parallel_group().device_group - ).to(vae.device) - vae.decoder = patched_decoder - log(f"Parallel VAE decoder enabled successfully.") - except ImportError: - log( - "DistVAE library is missing or does not support WanDecoderAdapter. " - "Try installing latest DistVAE from https://github.com/xdit-project/DistVAE. " - "Defaulting to single-rank decoder." - ) - except Exception as e: - raise ValueError(f"Failed to patch VAE decoder. {e}") - - def _remap_lightx2v_to_diffusers(k: str) -> str: """Remap a LightX2V-format state dict key to the diffusers WanTransformer3DModel naming.""" k = re.sub(r'\.self_attn\.q\.', '.attn1.to_q.', k) @@ -242,8 +196,6 @@ def _calculate_hybrid_attention_step_multiplier(self, input_args: dict) -> int: def _post_load_and_state_initialization(self, input_args: dict) -> None: super()._post_load_and_state_initialization(input_args) - if self.config.use_parallel_vae: - _setup_parallel_vae(self.pipe.vae, self.capabilities.use_parallel_vae_encoder) self.pipe.scheduler.config.flow_shift = input_args["flow_shift"] def _load_model(self) -> DiffusionPipeline: @@ -382,6 +334,8 @@ class xFuserWan22DistilledI2VModel(xFuserWan22I2VModel): use_parallel_vae_encoder=True, cross_attention_backend=True, supports_sparge_attention_backends=True, + enable_tiling=True, + enable_slicing=True, supports_distilled_weights=True, ) default_input_values = DefaultInputValues( @@ -450,7 +404,7 @@ def _validate_args(self, input_args: dict) -> None: ) guidance_scale = input_args.get("guidance_scale") if guidance_scale != 1.0: - log(f"Using guidance_scale=1.0. Other guindance scale values are not supported with this model.") + log("Using guidance_scale=1.0. Other guidance-scale values are unsupported for this model.") def _run_pipe(self, input_args: dict) -> DiffusionOutput: # Guidance is baked into the distilled weights. guidance_scale=1.0 keeps @@ -522,8 +476,6 @@ def _calculate_hybrid_attention_step_multiplier(self, input_args: dict) -> int: def _post_load_and_state_initialization(self, input_args: dict) -> None: super()._post_load_and_state_initialization(input_args) - if self.config.use_parallel_vae: - _setup_parallel_vae(self.pipe.vae, self.capabilities.use_parallel_vae_encoder) self.pipe.scheduler.config.flow_shift = input_args["flow_shift"] def _load_model(self) -> DiffusionPipeline: @@ -624,6 +576,7 @@ class xFuserWan22TI2VModel(xFuserWan21T2VModel): use_hybrid_attn_schedule=True, use_hybrid_gemm_schedule=True, use_parallel_vae=True, + use_parallel_vae_encoder=True, cross_attention_backend=True, supports_sparge_attention_backends=True, enable_tiling=True, @@ -740,6 +693,8 @@ class xFuserWan21VACEModel(xFuserWanModel): enable_tiling=True, enable_slicing=True, fully_shard_degree=True, + use_parallel_vae=True, + use_parallel_vae_encoder=True, ) default_input_values = DefaultInputValues( diff --git a/xfuser/model_executor/models/runner_models/z_image.py b/xfuser/model_executor/models/runner_models/z_image.py index 2b0558de..8ffc1854 100644 --- a/xfuser/model_executor/models/runner_models/z_image.py +++ b/xfuser/model_executor/models/runner_models/z_image.py @@ -59,6 +59,7 @@ class xFuserZImageModel(xFuserModel): fully_shard_degree=True, use_fp8_gemms=True, use_int8_gemms=True, + use_parallel_vae=True, ) settings = ModelSettings( model_name="Tongyi-MAI/Z-Image", @@ -133,9 +134,12 @@ class xFuserZImageTurboModel(xFuserModel): min_diffusers_version = "0.36.0" capabilities = ModelCapabilities( + enable_tiling=True, + enable_slicing=True, use_fp8_gemms=True, use_int8_gemms=True, fully_shard_degree=True, + use_parallel_vae=True, ) default_input_values = DefaultInputValues( height=1024, diff --git a/xfuser/model_executor/pipelines/base_pipeline.py b/xfuser/model_executor/pipelines/base_pipeline.py index 816846e1..f50514c5 100644 --- a/xfuser/model_executor/pipelines/base_pipeline.py +++ b/xfuser/model_executor/pipelines/base_pipeline.py @@ -1,7 +1,7 @@ from abc import ABCMeta, abstractmethod from functools import wraps from xfuser.compat import version_at_least -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional import sys import torch import torch.distributed @@ -9,7 +9,6 @@ from diffusers import DiffusionPipeline from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL -from distvae.modules.adapters.vae.decoder_adapters import DecoderAdapter from xfuser.core.distributed.group_coordinator import GroupCoordinator from xfuser.config.config import ( EngineConfig, @@ -47,6 +46,8 @@ from xfuser.envs import PACKAGES_CHECKER +from distvae.vae import parallelize_decoder + PACKAGES_CHECKER.check_diffusers_version() from xfuser.model_executor.schedulers import * @@ -66,6 +67,12 @@ logger = init_logger(__name__) + +def _vae_process_group(): + group = get_vae_parallel_group() + return getattr(group, "device_group", group) + + class xFuserVAEWrapper: def __init__( self, @@ -99,8 +106,8 @@ def __init__( def _convert_vae(self, vae: AutoencoderKL): """Convert VAE to parallel version""" - logger.info("VAE found, paralleling vae...") - vae.decoder = DecoderAdapter(vae.decoder, vae_group=get_vae_parallel_group()) + logger.info("VAE found; enabling parallel VAE decoding...") + parallelize_decoder(vae, _vae_process_group()) return vae def reset_activation_cache(self): @@ -468,8 +475,8 @@ def _convert_vae( self, vae: AutoencoderKL, ): - logger.info("VAE found, paralleling vae...") - vae.decoder = DecoderAdapter(vae.decoder) + logger.info("VAE found; enabling parallel VAE decoding...") + parallelize_decoder(vae, _vae_process_group()) return vae @abstractmethod