diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 5d633ea22..4e62cd665 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -147,7 +147,7 @@ The following tables show which models support each feature: | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | -| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | +| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | **Frame Interpolation Support** diff --git a/tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py b/tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py index 58ce14546..e65313d38 100644 --- a/tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py +++ b/tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py @@ -7,6 +7,9 @@ import torch from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan import ( DistributedAutoencoderKLHunyuan, ) +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan_video_15 import ( + DistributedAutoencoderKLHunyuanVideo15, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -36,6 +39,33 @@ class _DummyDistributedAutoencoderKLHunyuan(DistributedAutoencoderKLHunyuan): return b +class _DummyDistributedAutoencoderKLHunyuanVideo15(DistributedAutoencoderKLHunyuanVideo15): + def __init__(self): + torch.nn.Module.__init__(self) + self.tile_sample_min_height = 8 + self.tile_sample_min_width = 12 + self.tile_latent_min_height = 8 + self.tile_latent_min_width = 12 + self.tile_overlap_factor = 0.25 + self.use_tiling = False + + @property + def dtype(self) -> torch.dtype: + return torch.float32 + + def decoder(self, x: torch.Tensor) -> torch.Tensor: + return x + 10 + + def encoder(self, x: torch.Tensor) -> torch.Tensor: + return x + 20 + + def blend_v(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor: + return b + + def blend_h(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor: + return b + + def test_hunyuan_vae_use_tiling_aliases_spatial_tiling(): # Verify use_tiling property maps to use_spatial_tiling. vae = _DummyDistributedAutoencoderKLHunyuan() @@ -77,3 +107,43 @@ def test_hunyuan_vae_encode_tiles_round_trip(): assert grid_spec.tile_spec == {"blend_extent": 2, "row_limit": 6} assert len(tile_tasks) == 4 assert torch.equal(output, x + 20) + + +def test_hunyuan_video15_vae_decode_tiles_round_trip(): + vae = _DummyDistributedAutoencoderKLHunyuanVideo15() + z = torch.arange(192, dtype=torch.float32).reshape(1, 1, 1, 12, 16) + + tile_tasks, grid_spec = vae.tile_split(z) + decoded_tiles = {task.grid_coord: vae.tile_exec(task) for task in tile_tasks} + output = vae.tile_merge(decoded_tiles, grid_spec) + + assert grid_spec.split_dims == (3, 4) + assert grid_spec.grid_shape == (2, 2) + assert grid_spec.tile_spec == { + "blend_height": 2, + "blend_width": 3, + "row_limit_height": 6, + "row_limit_width": 9, + } + assert len(tile_tasks) == 4 + assert torch.equal(output, z + 10) + + +def test_hunyuan_video15_vae_encode_tiles_round_trip(): + vae = _DummyDistributedAutoencoderKLHunyuanVideo15() + x = torch.arange(192, dtype=torch.float32).reshape(1, 1, 1, 12, 16) + + tile_tasks, grid_spec = vae.encode_tile_split(x) + encoded_tiles = {task.grid_coord: vae.encode_tile_exec(task) for task in tile_tasks} + output = vae.tile_merge(encoded_tiles, grid_spec) + + assert grid_spec.split_dims == (3, 4) + assert grid_spec.grid_shape == (2, 2) + assert grid_spec.tile_spec == { + "blend_height": 2, + "blend_width": 3, + "row_limit_height": 6, + "row_limit_width": 9, + } + assert len(tile_tasks) == 4 + assert torch.equal(output, x + 20) diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_hunyuan_video_15.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_hunyuan_video_15.py new file mode 100644 index 000000000..5011f2a45 --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_hunyuan_video_15.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import torch +from diffusers import AutoencoderKLHunyuanVideo15 +from vllm.logger import init_logger + +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedOperator, + DistributedVaeMixin, + GridSpec, + TileTask, +) + +logger = init_logger(__name__) + + +class DistributedAutoencoderKLHunyuanVideo15(AutoencoderKLHunyuanVideo15, DistributedVaeMixin): + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model + + def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + _, _, _, height, width = z.shape + overlap_height = int(self.tile_latent_min_height * (1 - self.tile_overlap_factor)) + overlap_width = int(self.tile_latent_min_width * (1 - self.tile_overlap_factor)) + blend_height = int(self.tile_sample_min_height * self.tile_overlap_factor) + blend_width = int(self.tile_sample_min_width * self.tile_overlap_factor) + row_limit_height = self.tile_sample_min_height - blend_height + row_limit_width = self.tile_sample_min_width - blend_width + + tiletask_list = [] + for i in range(0, height, overlap_height): + for j in range(0, width, overlap_width): + tile = z[:, :, :, i : i + self.tile_latent_min_height, j : j + self.tile_latent_min_width] + tiletask_list.append( + TileTask( + len(tiletask_list), + (i // overlap_height, j // overlap_width), + tile, + workload=tile.shape[3] * tile.shape[4], + ) + ) + + grid_spec = GridSpec( + split_dims=(3, 4), + grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + tile_spec={ + "blend_height": blend_height, + "blend_width": blend_width, + "row_limit_height": row_limit_height, + "row_limit_width": row_limit_width, + }, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + + def tile_exec(self, task: TileTask) -> torch.Tensor: + return self.decoder(task.tensor) + + def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor: + grid_h, grid_w = grid_spec.grid_shape + result_rows = [] + for i in range(grid_h): + result_row = [] + for j in range(grid_w): + tile = coord_tensor_map[(i, j)] + if i > 0: + tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"]) + if j > 0: + tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"]) + result_row.append( + tile[ + :, + :, + :, + : grid_spec.tile_spec["row_limit_height"], + : grid_spec.tile_spec["row_limit_width"], + ] + ) + result_rows.append(torch.cat(result_row, dim=-1)) + return torch.cat(result_rows, dim=-2) + + def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + _, _, _, height, width = x.shape + overlap_height = int(self.tile_sample_min_height * (1 - self.tile_overlap_factor)) + overlap_width = int(self.tile_sample_min_width * (1 - self.tile_overlap_factor)) + blend_height = int(self.tile_latent_min_height * self.tile_overlap_factor) + blend_width = int(self.tile_latent_min_width * self.tile_overlap_factor) + row_limit_height = self.tile_latent_min_height - blend_height + row_limit_width = self.tile_latent_min_width - blend_width + + tiletask_list = [] + for i in range(0, height, overlap_height): + for j in range(0, width, overlap_width): + tile = x[:, :, :, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width] + tiletask_list.append( + TileTask( + len(tiletask_list), + (i // overlap_height, j // overlap_width), + tile, + workload=tile.shape[3] * tile.shape[4], + ) + ) + + grid_spec = GridSpec( + split_dims=(3, 4), + grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + tile_spec={ + "blend_height": blend_height, + "blend_width": blend_width, + "row_limit_height": row_limit_height, + "row_limit_width": row_limit_width, + }, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + + def encode_tile_exec(self, task: TileTask) -> torch.Tensor: + return self.encoder(task.tensor) + + def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: + if not self.is_distributed_enabled(): + return super().tiled_encode(x) + + logger.debug("Encode running with distributed executor") + return self.distributed_executor.execute( + x, + DistributedOperator( + split=self.encode_tile_split, + exec=self.encode_tile_exec, + merge=self.tile_merge, + ), + broadcast_result=True, + ) + + def tiled_decode(self, z: torch.Tensor) -> torch.Tensor: + if not self.is_distributed_enabled(): + return super().tiled_decode(z) + + logger.debug("Decode running with distributed executor") + return self.distributed_executor.execute( + z, + DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge), + broadcast_result=True, + ) diff --git a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py index 123652edd..1f06799c3 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py @@ -762,15 +762,14 @@ class HunyuanVideo15Transformer3DModel(nn.Module): encoder_hidden_states = torch.stack(new_encoder_hidden_states) encoder_attention_mask = torch.stack(new_encoder_attention_mask) - # Create explicit attn_mask for image tokens when SP auto_pad is active. + max_valid_encoder_tokens = int(encoder_attention_mask.sum(dim=1).max().item()) + if max_valid_encoder_tokens < encoder_attention_mask.shape[1]: + encoder_hidden_states = encoder_hidden_states[:, :max_valid_encoder_tokens] + encoder_attention_mask = encoder_attention_mask[:, :max_valid_encoder_tokens] + if encoder_attention_mask.all(): + encoder_attention_mask = None + ctx = get_forward_context() - if not ctx.sp_active: - max_valid_encoder_tokens = int(encoder_attention_mask.sum(dim=1).max().item()) - if max_valid_encoder_tokens < encoder_attention_mask.shape[1]: - encoder_hidden_states = encoder_hidden_states[:, :max_valid_encoder_tokens] - encoder_attention_mask = encoder_attention_mask[:, :max_valid_encoder_tokens] - if encoder_attention_mask.all(): - encoder_attention_mask = None hidden_states_mask = None if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size diff --git a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py index 6445bfee2..3688aa7c2 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py +++ b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py @@ -11,7 +11,6 @@ from typing import Any import numpy as np import torch -from diffusers import AutoencoderKLHunyuanVideo15 from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler from diffusers.utils.torch_utils import randn_tensor from diffusers.video_processor import VideoProcessor @@ -20,6 +19,9 @@ from transformers import AutoConfig, ByT5Tokenizer, Qwen2_5_VLTextModel, Qwen2To from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan_video_15 import ( + DistributedAutoencoderKLHunyuanVideo15, +) from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -110,7 +112,7 @@ class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, Diff t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only) self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(dtype=dtype, device=self.device) - self.vae = AutoencoderKLHunyuanVideo15.from_pretrained( + self.vae = DistributedAutoencoderKLHunyuanVideo15.from_pretrained( model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only ).to(self.device) diff --git a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py index c1acd1a89..ad6339c8c 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py +++ b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py @@ -11,7 +11,6 @@ from typing import Any, cast import numpy as np import PIL.Image import torch -from diffusers import AutoencoderKLHunyuanVideo15 from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler from diffusers.utils.torch_utils import randn_tensor from diffusers.video_processor import VideoProcessor @@ -27,6 +26,9 @@ from transformers import ( from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan_video_15 import ( + DistributedAutoencoderKLHunyuanVideo15, +) from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -139,7 +141,7 @@ class HunyuanVideo15I2VPipeline( model, subfolder="feature_extractor", local_files_only=local_files_only ) - self.vae = AutoencoderKLHunyuanVideo15.from_pretrained( + self.vae = DistributedAutoencoderKLHunyuanVideo15.from_pretrained( model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only ).to(self.device)