[CI]Add Cosmos3 L2 serving tests (#4535)

Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>

.buildkite/test-ready.yml

@@ -584,6 +584,50 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: "Diffusion · Cosmos3 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_cosmos3.py
+          - vllm_omni/diffusion/models/cosmos3/
+        commands:
+          - |
+            timeout 40m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_cosmos3.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
 
       - label: "Diffusion · Qwen Image Test"
         source_file_dependencies:

tests/e2e/online_serving/test_cosmos3.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Online serving L2 smoke tests for ``nvidia/Cosmos3-Nano``.
+
+These keep Cosmos3 in the regular ``core_model`` diffusion lane while the
+broader image/video similarity checks stay in ``tests/e2e/accuracy``.
+"""
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import generate_synthetic_image, generate_synthetic_video
+from tests.helpers.runtime import (
+    OmniServer,
+    OmniServerParams,
+    OpenAIClientHandler,
+)
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = "nvidia/Cosmos3-Nano"
+PROMPT = "A small warehouse robot moves a blue box across a clean floor."
+NEGATIVE_PROMPT = "blurry, distorted, low quality"
+VIDEO_SHAPE = {
+    "height": 256,
+    "width": 256,
+    "num_frames": 5,
+    "fps": 1,
+    "num_inference_steps": 2,
+    "guidance_scale": 1.0,
+    "flow_shift": 3.0,
+    "seed": 42,
+}
+SERVER_ARGS = [
+    "--model-class-name",
+    "Cosmos3OmniDiffusersPipeline",
+    "--no-guardrails",
+]
+SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _get_diffusion_feature_cases(model: str):
+    """Return a single default Cosmos3 server row for L2 coverage."""
+    return [
+        pytest.param(
+            OmniServerParams(model=model, server_args=SERVER_ARGS),
+            id="default",
+            marks=SINGLE_CARD_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    """Default Cosmos3 T2I smoke through ``/v1/images/generations``."""
+    responses = openai_client.send_images_generations_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "prompt": PROMPT,
+                "negative_prompt": NEGATIVE_PROMPT,
+                "size": "256x256",
+                "n": 1,
+                "response_format": "b64_json",
+                "num_inference_steps": 2,
+                "guidance_scale": 1.0,
+                "flow_shift": 3.0,
+                "seed": 42,
+            }
+        }
+    )
+    response = responses[0]
+    assert response.success, response.error_message
+    payload = response.json_body
+    assert isinstance(payload, dict)
+    assert len(payload["data"]) == 1
+    assert payload["data"][0]["b64_json"]
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_video_generation_modes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    """Default Cosmos3 T2V, I2V, and V2V smoke through the async ``/v1/videos`` endpoint."""
+    image_reference = f"data:image/jpeg;base64,{generate_synthetic_image(256, 256, seed=42)['base64']}"
+    video_reference = f"data:video/mp4;base64,{generate_synthetic_video(256, 256, 5)['base64']}"
+
+    video_cases = [
+        {
+            "form_data": {
+                **VIDEO_SHAPE,
+                "prompt": PROMPT,
+                "negative_prompt": NEGATIVE_PROMPT,
+            },
+        },
+        {
+            "form_data": {
+                **VIDEO_SHAPE,
+                "prompt": "The blue box moves slowly forward from the reference image.",
+                "negative_prompt": NEGATIVE_PROMPT,
+            },
+            "image_reference": image_reference,
+        },
+        {
+            "form_data": {
+                **VIDEO_SHAPE,
+                "prompt": "Continue the same synthetic motion with consistent shapes.",
+                "negative_prompt": NEGATIVE_PROMPT,
+                "extra_params": '{"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}',
+            },
+            "video_reference": video_reference,
+        },
+    ]
+
+    for case in video_cases:
+        request_config = {
+            "model": omni_server.model,
+            **case,
+        }
+        openai_client.send_video_diffusion_request(request_config)

tests/helpers/runtime.py

@@ -2084,6 +2084,9 @@ class OpenAIClientHandler:
         normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None}
         files: dict[str, tuple[str, BytesIO, str]] = {}
         image_reference = request_config.get("image_reference")
+        video_reference = request_config.get("video_reference")
+        if image_reference and video_reference:
+            raise ValueError("Only one of image_reference or video_reference can be provided")
         if image_reference:
             if image_reference.startswith("data:image"):
                 header, encoded = image_reference.split(",", 1)
@@ -2093,6 +2096,15 @@ class OpenAIClientHandler:
                 files["input_reference"] = (f"reference.{extension}", BytesIO(file_data), content_type)
             else:
                 normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference})
+        if video_reference:
+            if video_reference.startswith("data:video"):
+                header, encoded = video_reference.split(",", 1)
+                content_type = header.split(";")[0].removeprefix("data:")
+                extension = content_type.split("/")[-1]
+                file_data = base64.b64decode(encoded)
+                files["input_reference"] = (f"reference.{extension}", BytesIO(file_data), content_type)
+            else:
+                normalized_form_data["video_reference"] = json.dumps({"video_url": video_reference})
 
         result = DiffusionResponse()
         create_url = self._build_url("/v1/videos")