Add Cosmos3 model support (#3454)

Signed-off-by: Maciej Bala <mbala@nvidia.com>
Signed-off-by: MaciejBalaNV <mbala@nvidia.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>

tests/diffusion/test_diffusion_ipc.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import contextlib
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.ipc import (
+    _SHM_TENSOR_THRESHOLD,
+    _pack_value_if_large,
+    _unpack_if_shm_handle,
+    pack_diffusion_output_shm,
+    unpack_diffusion_output_shm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _large_numel(dtype: torch.dtype) -> int:
+    return (_SHM_TENSOR_THRESHOLD // torch.empty((), dtype=dtype).element_size()) + 1
+
+
+def _cleanup_shm_handle(value: object) -> None:
+    if isinstance(value, dict) and value.get("__tensor_shm__"):
+        with contextlib.suppress(FileNotFoundError):
+            _unpack_if_shm_handle(value)
+
+
+def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None:
+    image = torch.arange(300_000, dtype=torch.float32)
+    video = torch.arange(300_000, dtype=torch.float32) * 2
+    output = DiffusionOutput(output={"image": image, "video": video, "metadata": {"keep": "inline"}})
+
+    pack_diffusion_output_shm(output)
+
+    assert output.output["image"]["__tensor_shm__"] is True
+    assert output.output["video"]["__tensor_shm__"] is True
+    assert output.output["metadata"] == {"keep": "inline"}
+
+    unpack_diffusion_output_shm(output)
+
+    torch.testing.assert_close(output.output["image"], image)
+    torch.testing.assert_close(output.output["video"], video)
+    assert output.output["metadata"] == {"keep": "inline"}
+
+
+def test_diffusion_output_tuple_tensors_round_trip_through_shm() -> None:
+    # LTX2 / DreamID return (video, audio) tuples as DiffusionOutput.output.
+    video = torch.arange(300_000, dtype=torch.float32)
+    audio = torch.arange(300_000, dtype=torch.float32) * 3
+    output = DiffusionOutput(output=(video, audio))
+
+    pack_diffusion_output_shm(output)
+
+    assert isinstance(output.output, tuple)
+    assert output.output[0]["__tensor_shm__"] is True
+    assert output.output[1]["__tensor_shm__"] is True
+
+    unpack_diffusion_output_shm(output)
+
+    assert isinstance(output.output, tuple)
+    torch.testing.assert_close(output.output[0], video)
+    torch.testing.assert_close(output.output[1], audio)
+
+
+def test_diffusion_output_list_tensors_round_trip_through_shm() -> None:
+    frames = [torch.arange(300_000, dtype=torch.float32), torch.arange(300_000, dtype=torch.float32) + 1]
+    output = DiffusionOutput(output=list(frames))
+
+    pack_diffusion_output_shm(output)
+
+    assert isinstance(output.output, list)
+    assert all(isinstance(item, dict) and item["__tensor_shm__"] is True for item in output.output)
+
+    unpack_diffusion_output_shm(output)
+
+    assert isinstance(output.output, list)
+    torch.testing.assert_close(output.output[0], frames[0])
+    torch.testing.assert_close(output.output[1], frames[1])
+
+
+def test_pack_value_keeps_tensor_at_threshold_inline() -> None:
+    tensor = torch.arange(
+        _SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(),
+        dtype=torch.float32,
+    )
+
+    packed = _pack_value_if_large(tensor)
+
+    assert packed is tensor
+
+
+def test_pack_value_packs_large_tensor_and_round_trips() -> None:
+    tensor = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
+    packed = _pack_value_if_large(tensor)
+
+    try:
+        assert isinstance(packed, dict)
+        assert packed["__tensor_shm__"] is True
+        assert packed["shape"] == [tensor.numel()]
+        assert packed["torch_dtype"] == "torch.float32"
+
+        unpacked = _unpack_if_shm_handle(packed)
+        assert isinstance(unpacked, torch.Tensor)
+        torch.testing.assert_close(unpacked, tensor)
+    finally:
+        _cleanup_shm_handle(packed)
+
+
+def test_pack_value_recurses_nested_dicts_and_lists_without_mutating_inline_values() -> None:
+    large = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
+    small = torch.arange(8, dtype=torch.float32)
+    list_tensor = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
+    payload = {
+        "media": {
+            "large": large,
+            "small": small,
+        },
+        "list_value": [list_tensor],
+        "metadata": {"prompt": "keep inline"},
+    }
+
+    packed = _pack_value_if_large(payload)
+
+    try:
+        assert packed is not payload
+        assert packed["media"] is not payload["media"]
+        assert packed["media"]["large"]["__tensor_shm__"] is True
+        assert packed["media"]["small"] is small
+        # Lists are recursed too: the large tensor inside is packed and a new
+        # list is returned, while the input payload is left untouched.
+        assert packed["list_value"] is not payload["list_value"]
+        assert packed["list_value"][0]["__tensor_shm__"] is True
+        assert payload["list_value"][0] is list_tensor
+        assert packed["metadata"] == {"prompt": "keep inline"}
+
+        torch.testing.assert_close(_unpack_if_shm_handle(packed["media"]["large"]), large)
+        torch.testing.assert_close(_unpack_if_shm_handle(packed["list_value"][0]), list_tensor)
+    finally:
+        if isinstance(packed, dict):
+            _cleanup_shm_handle(packed.get("media", {}).get("large"))
+            list_value = packed.get("list_value")
+            if isinstance(list_value, list) and list_value:
+                _cleanup_shm_handle(list_value[0])
+
+
+def test_pack_value_preserves_dtype_shape_and_values_for_bfloat16() -> None:
+    tensor = torch.arange(_large_numel(torch.bfloat16), dtype=torch.float32).to(torch.bfloat16).reshape(1, -1)
+    packed = _pack_value_if_large(tensor)
+
+    try:
+        assert isinstance(packed, dict)
+        assert packed["__tensor_shm__"] is True
+        assert packed["shape"] == list(tensor.shape)
+        assert packed["torch_dtype"] == "torch.bfloat16"
+        assert packed["numpy_dtype"] == "float32"
+
+        unpacked = _unpack_if_shm_handle(packed)
+        assert isinstance(unpacked, torch.Tensor)
+        assert unpacked.dtype == torch.bfloat16
+        torch.testing.assert_close(unpacked, tensor)
+    finally:
+        _cleanup_shm_handle(packed)
+
+
+def test_pack_value_packs_non_contiguous_large_tensor_values() -> None:
+    tensor = torch.arange(_large_numel(torch.float32) * 2, dtype=torch.float32).reshape(-1, 2)[:, 0]
+    assert not tensor.is_contiguous()
+
+    packed = _pack_value_if_large(tensor)
+
+    try:
+        assert isinstance(packed, dict)
+        assert packed["__tensor_shm__"] is True
+        assert packed["shape"] == list(tensor.shape)
+
+        unpacked = _unpack_if_shm_handle(packed)
+        assert isinstance(unpacked, torch.Tensor)
+        torch.testing.assert_close(unpacked, tensor)
+    finally:
+        _cleanup_shm_handle(packed)