[CI] add diffusion ci

.buildkite/pipeline.yml

@@ -1,6 +1,13 @@
 steps:
   - label: "Simple Unit Test"
     commands:
-      - ".buildkite/scripts/simple_test.sh"
+      - ".buildkite/scripts/setup_env.sh"
+      - source .buildkite/scripts/setup_env.sh && "${VENV_PYTHON}" -m pytest -v -s tests/test_omni_llm.py
     agents:
       queue: "cpu_queue_premerge"
+  - label: "Diffusion Model Test"
+    commands:
+      - ".buildkite/scripts/setup_env.sh"
+      - source .buildkite/scripts/setup_env.sh && "${VENV_PYTHON}" -m pytest -s -v tests/test_diffusion_model.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU

.buildkite/scripts/simple_test.sh → .buildkite/scripts/setup_env.sh

@@ -22,6 +22,7 @@ PATH="${PYTHON_SCRIPTS}:${PATH}"
 if [[ -n "${USER_BASE}" ]]; then
   PATH="${USER_BASE}/bin:${PATH}"
 fi
+export PATH
 
 if ! command -v "${UV_BIN}" >/dev/null 2>&1; then
   "${PYTHON_BOOTSTRAP}" -m pip install --upgrade pip
@@ -49,4 +50,3 @@ VENV_PYTHON="${VENV_DIR}/bin/python"
 
 "${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.11.0
 "${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]"
-"${VENV_PYTHON}" -m pytest -v -s tests/test_omni_llm.py

tests/test_diffusion_model.py

@@ -0,0 +1,17 @@
+import pytest
+
+from vllm_omni import Omni
+
+models = ["Tongyi-MAI/Z-Image-Turbo"]
+
+
+@pytest.mark.parametrize("model_name", models)
+def test_diffusion_model(model_name: str):
+    m = Omni(model=model_name)
+    image = m.generate(
+        "a photo of a cat sitting on a laptop keyboard",
+        height=1024,
+        width=1024,
+        num_inference_steps=9,
+    )
+    image[0].save("z_image_output.png")