diff --git a/tests/diffusion/distributed/test_autoencoder_kl_wan.py b/tests/diffusion/distributed/test_autoencoder_kl_wan.py new file mode 100644 index 000000000..2ea1c1214 --- /dev/null +++ b/tests/diffusion/distributed/test_autoencoder_kl_wan.py @@ -0,0 +1,43 @@ +import pytest +import torch + +from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_wan as wan_vae_module +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _DummyOmniAutoencoderKLWan(OmniAutoencoderKLWan): + def __init__(self, *, dtype: torch.dtype): + torch.nn.Module.__init__(self) + self.register_parameter("dummy_weight", torch.nn.Parameter(torch.ones(1, dtype=dtype))) + + +def test_wan_vae_execution_context_handles_fp32(): + model = _DummyOmniAutoencoderKLWan(dtype=torch.float32) + with model._execution_context(): + output = model.dummy_weight + 1 + assert output.dtype == torch.float32 + + +def test_wan_vae_execution_context_handles_bf16(): + model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16) + with model._execution_context(): + output = model.dummy_weight + 1 + assert output.dtype == torch.bfloat16 + + +def test_wan_vae_execution_context_uses_platform_autocast(mocker): + sentinel = object() + platform = mocker.Mock() + platform.create_autocast_context.return_value = sentinel + mocker.patch.object(wan_vae_module, "current_omni_platform", platform) + + model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16) + + assert model._execution_context() is sentinel + platform.create_autocast_context.assert_called_once_with( + device_type=model.dummy_weight.device.type, + dtype=torch.bfloat16, + enabled=True, + ) diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index 027991c3f..35c9434d0 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from contextlib import nullcontext from typing import Any import torch @@ -15,11 +16,38 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor GridSpec, TileTask, ) +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) -class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin): +class OmniAutoencoderKLWan(AutoencoderKLWan): + def _execution_context(self): + try: + first_param = next(self.parameters()) + except StopIteration: + return nullcontext() + + dtype = first_param.dtype + if dtype not in (torch.float16, torch.bfloat16): + return nullcontext() + + return current_omni_platform.create_autocast_context( + device_type=first_param.device.type, + dtype=dtype, + enabled=True, + ) + + def encode(self, x: torch.Tensor, return_dict: bool = True): + with self._execution_context(): + return super().encode(x, return_dict=return_dict) + + def decode(self, z: torch.Tensor, return_dict: bool = True): + with self._execution_context(): + return super().decode(z, return_dict=return_dict) + + +class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): @classmethod def from_pretrained(cls, *args: Any, **kwargs: Any): model = super().from_pretrained(*args, **kwargs) @@ -84,11 +112,12 @@ class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin): """Decode a single latent tile into RGB space.""" self.clear_cache() time = [] - for k in range(len(task.tensor)): - self._conv_idx = [0] - tile = self.post_quant_conv(task.tensor[k]) - decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0)) - time.append(decoded) + with self._execution_context(): + for k in range(len(task.tensor)): + self._conv_idx = [0] + tile = self.post_quant_conv(task.tensor[k]) + decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0)) + time.append(decoded) result = torch.cat(time, dim=2) return result diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index d2d2bb860..a550e576f 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -272,7 +272,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only ).to(self.device) self.vae = DistributedAutoencoderKLWan.from_pretrained( - model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only + model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only ).to(self.device) # Initialize transformers with correct config (weights loaded via load_weights) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 1e8a94eb3..c05ecc9c9 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -217,7 +217,7 @@ class Wan22I2VPipeline( # VAE self.vae = DistributedAutoencoderKLWan.from_pretrained( - model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only + model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only ).to(self.device) # Transformers (weights loaded via load_weights) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py index f116834cf..261f62fb7 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py @@ -24,13 +24,13 @@ from typing import Any, cast import numpy as np import PIL.Image import torch -from diffusers import AutoencoderKLWan from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoTokenizer, UMT5EncoderModel from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan 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 @@ -174,8 +174,8 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress ).to(self.device) # VAE - self.vae = AutoencoderKLWan.from_pretrained( - model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only + self.vae = OmniAutoencoderKLWan.from_pretrained( + model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only ).to(self.device) # Single transformer (TI2V uses dense 5B model, not MoE) diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 4df297fa0..8f1e66747 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -1,12 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from contextlib import nullcontext from enum import Enum from typing import Any import torch +from vllm.logger import init_logger from vllm.platforms import Platform +logger = init_logger(__name__) + class OmniPlatformEnum(Enum): """Enum for supported Omni platforms.""" @@ -113,6 +117,23 @@ class OmniPlatform(Platform): def get_free_memory(cls, device: torch.device | None = None) -> int: raise NotImplementedError + @classmethod + def create_autocast_context( + cls, + *, + device_type: str, + dtype: torch.dtype, + enabled: bool = True, + ): + if not enabled: + return nullcontext() + + try: + return torch.autocast(device_type=device_type, dtype=dtype, enabled=True) + except (RuntimeError, TypeError, ValueError) as exc: + logger.warning("autocast unavailable for device_type=%s dtype=%s: %s", device_type, dtype, exc) + return nullcontext() + @classmethod def supports_cpu_offload(cls) -> bool: return True diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 1d6bea7cb..c40dd6fea 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from contextlib import nullcontext from typing import Any import torch @@ -106,6 +107,24 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): device_props = torch.npu.get_device_properties(device_id) return device_props.total_memory + @classmethod + def create_autocast_context(cls, *, device_type, dtype, enabled=True): + if device_type != "npu": + return super().create_autocast_context( + device_type=device_type, + dtype=dtype, + enabled=enabled, + ) + if not enabled: + return nullcontext() + + # NPU-specific fallback + try: + return torch.npu.amp.autocast(dtype=dtype) + except (RuntimeError, TypeError, ValueError) as exc: + logger.warning("autocast unavailable for device_type=%s dtype=%s: %s", device_type, dtype, exc) + return nullcontext() + @classmethod def get_profiler_cls(cls) -> str: return "vllm_omni.platforms.npu.profiler.NPUTorchProfilerWrapper"