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_engine.py

@@ -10,6 +10,9 @@ from types import SimpleNamespace
 from typing import Any
 
 import pytest
+import torch
+
+from vllm_omni.diffusion.diffusion_engine import _move_tensor_tree_to_cpu
 
 
 @dataclass
@@ -63,6 +66,80 @@ class MockScheduler:
         return [req.request_id for req in sched_output.scheduled_new_reqs]
 
 
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_move_tensor_tree_keeps_cpu_tensor_identity() -> None:
+    tensor = torch.arange(8, dtype=torch.float32)
+
+    moved = _move_tensor_tree_to_cpu(tensor)
+
+    assert moved is tensor
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_move_tensor_tree_preserves_nested_structure_without_mutating_input() -> None:
+    tensor = torch.arange(4, dtype=torch.float32)
+    nested_tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    sentinel = object()
+    payload = {
+        "tensor": tensor,
+        "list": [nested_tensor, sentinel],
+        "tuple": ({"inner": tensor}, "metadata"),
+        "scalar": 3,
+    }
+
+    moved = _move_tensor_tree_to_cpu(payload)
+
+    assert moved is not payload
+    assert set(moved) == {"tensor", "list", "tuple", "scalar"}
+    assert moved["list"] is not payload["list"]
+    assert moved["tuple"] is not payload["tuple"]
+    assert moved["tuple"][0] is not payload["tuple"][0]
+    assert moved["tensor"] is tensor
+    assert moved["list"][0] is nested_tensor
+    assert moved["list"][1] is sentinel
+    assert moved["tuple"][0]["inner"] is tensor
+    assert moved["tuple"][1] == "metadata"
+    assert moved["scalar"] == 3
+    assert payload["list"][0] is nested_tensor
+    assert payload["list"][1] is sentinel
+    assert payload["tuple"][0]["inner"] is tensor
+    assert payload["tuple"][1] == "metadata"
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_move_tensor_tree_returns_non_tensor_values_unchanged() -> None:
+    value = object()
+
+    moved = _move_tensor_tree_to_cpu(value)
+
+    assert moved is value
+
+
+@pytest.mark.diffusion
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_move_tensor_tree_moves_nested_cuda_tensors_to_cpu() -> None:
+    tensor = torch.arange(8, dtype=torch.float32, device="cuda")
+    other = torch.arange(4, dtype=torch.int64, device="cuda")
+    payload = {"tensor": tensor, "items": [other, ("keep", tensor)]}
+
+    moved = _move_tensor_tree_to_cpu(payload)
+
+    assert moved["tensor"].device.type == "cpu"
+    assert moved["items"][0].device.type == "cpu"
+    assert moved["items"][1][1].device.type == "cpu"
+    torch.testing.assert_close(moved["tensor"], tensor.cpu())
+    torch.testing.assert_close(moved["items"][0], other.cpu())
+    torch.testing.assert_close(moved["items"][1][1], tensor.cpu())
+    assert moved["items"][1][0] == "keep"
+
+
 @pytest.mark.asyncio
 async def test_async_add_req_and_wait_for_response():
     from vllm_omni.diffusion.diffusion_engine import DiffusionEngine