diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 824dd9a0b..69c81737d 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -145,7 +145,7 @@ The following tables show which models support each feature: | **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | -| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | +| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | **Frame Interpolation Support** diff --git a/docs/user_guide/examples/offline_inference/x_to_video_audio.md b/docs/user_guide/examples/offline_inference/x_to_video_audio.md index cec8d47c5..d5cd0ea46 100644 --- a/docs/user_guide/examples/offline_inference/x_to_video_audio.md +++ b/docs/user_guide/examples/offline_inference/x_to_video_audio.md @@ -47,6 +47,24 @@ python x_to_video_audio.py \ ``` In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled. The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload. +For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP: + +```python +python x_to_video_audio.py \ + --model /path/to/dreamid_omni \ + --prompt "Two people walking together and singing happily" \ + --image-path ./example0.png ./example1.png \ + --audio-path ./example0.wav ./example1.wav \ + --video-negative-prompt "jitter, bad hands, blur, distortion" \ + --audio-negative-prompt "robotic, muffled, echo, distorted" \ + --cfg-parallel-size 2 \ + --num-inference-steps 45 \ + --height 704 \ + --width 1280 \ + --use-hsdp \ + --hsdp-shard-size 2 \ + --output out_dreamid_omni_twoip.mp4 +``` You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni @@ -77,6 +95,9 @@ Key arguments: - `--image-path`: path to the input image list. - `--audio-path`: path to the input audio list, indicate the timbre of the output video. - `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1). +- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks. +- `--hsdp-shard-size`: number of GPUs used for HSDP sharding. +- `--hsdp-replicate-size`: number of HSDP replica groups. - `--num-inference-steps`: number of denoising steps (defaults 45). - `--video-negative-prompt`: negative prompt for video generation. - `--audio-negative-prompt`: negative prompt for audio generation. diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md index 13f2cfe7c..7621a6058 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md +++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md @@ -46,6 +46,24 @@ python x_to_video_audio.py \ ``` In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled. The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload. +For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP: + +```python +python x_to_video_audio.py \ + --model /path/to/dreamid_omni \ + --prompt "Two people walking together and singing happily" \ + --image-path ./example0.png ./example1.png \ + --audio-path ./example0.wav ./example1.wav \ + --video-negative-prompt "jitter, bad hands, blur, distortion" \ + --audio-negative-prompt "robotic, muffled, echo, distorted" \ + --cfg-parallel-size 2 \ + --num-inference-steps 45 \ + --height 704 \ + --width 1280 \ + --use-hsdp \ + --hsdp-shard-size 2 \ + --output out_dreamid_omni_twoip.mp4 +``` You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni @@ -76,6 +94,9 @@ Key arguments: - `--image-path`: path to the input image list. - `--audio-path`: path to the input audio list, indicate the timbre of the output video. - `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1). +- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks. +- `--hsdp-shard-size`: number of GPUs used for HSDP sharding. +- `--hsdp-replicate-size`: number of HSDP replica groups. - `--num-inference-steps`: number of denoising steps (defaults 45). - `--video-negative-prompt`: negative prompt for video generation. - `--audio-negative-prompt`: negative prompt for audio generation. diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py index 497284ceb..544c285ab 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py +++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py @@ -63,6 +63,23 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) + parser.add_argument( + "--use-hsdp", + action="store_true", + help="Enable HSDP for DreamID-Omni fused transformer blocks.", + ) + parser.add_argument( + "--hsdp-shard-size", + type=int, + default=1, + help="Number of GPUs used for HSDP sharding when HSDP is enabled.", + ) + parser.add_argument( + "--hsdp-replicate-size", + type=int, + default=1, + help="Number of HSDP replica groups. Default 1 means pure sharding.", + ) return parser.parse_args() @@ -124,6 +141,9 @@ def main() -> None: parallel_config = DiffusionParallelConfig( cfg_parallel_size=args.cfg_parallel_size, + use_hsdp=args.use_hsdp, + hsdp_shard_size=args.hsdp_shard_size, + hsdp_replicate_size=args.hsdp_replicate_size, ) omni = Omni( diff --git a/tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py b/tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py new file mode 100644 index 000000000..c7e977ecf --- /dev/null +++ b/tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py @@ -0,0 +1,27 @@ +import pytest +import torch.nn as nn + +pytest.importorskip("dreamid_omni.modules.model") + +from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_dreamid_omni_exposes_hsdp_shard_conditions_for_fused_blocks(): + model = object.__new__(FusionModel) + nn.Module.__init__(model) + model.fused_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)]) + model.attn = nn.Identity() + + conditions = getattr(model, "_hsdp_shard_conditions", None) + + assert conditions is not None + assert len(conditions) == 1 + + matched = [] + for name, module in model.named_modules(): + if any(cond(name, module) for cond in conditions): + matched.append(name) + + assert matched == ["fused_blocks.0", "fused_blocks.1"] diff --git a/vllm_omni/diffusion/models/dreamid_omni/__init__.py b/vllm_omni/diffusion/models/dreamid_omni/__init__.py index f14db611e..8c7f05bf8 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/__init__.py +++ b/vllm_omni/diffusion/models/dreamid_omni/__init__.py @@ -1,8 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from .pipeline_dreamid_omni import DreamIDOmniPipeline - __all__ = [ "DreamIDOmniPipeline", ] + + +def __getattr__(name: str): + if name == "DreamIDOmniPipeline": + from .pipeline_dreamid_omni import DreamIDOmniPipeline + + return DreamIDOmniPipeline + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm_omni/diffusion/models/dreamid_omni/fusion.py b/vllm_omni/diffusion/models/dreamid_omni/fusion.py index abca4c947..a7ba745e9 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/fusion.py +++ b/vllm_omni/diffusion/models/dreamid_omni/fusion.py @@ -241,6 +241,12 @@ class FusedBlock(nn.Module): class FusionModel(nn.Module): _layerwise_offload_blocks_attrs = ["fused_blocks"] + @staticmethod + def _is_fused_block(name: str, module) -> bool: + return "fused_blocks" in name and name.split(".")[-1].isdigit() + + _hsdp_shard_conditions = [_is_fused_block] + def __init__(self, video_config=None, audio_config=None): super().__init__() has_video = True