From 630cbe1ee66b2d5d8865bd9ba5413d91b6de0053 Mon Sep 17 00:00:00 2001
From: Jiangyun Zhu <riverclouds.zhu@qq.com>
Date: Thu, 4 Dec 2025 23:23:22 +0800
Subject: [PATCH] [CI] add diffusion ci (#174)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
---
 .buildkite/pipeline.yml                    | 25 +++++++++++++++++++
 .gitignore                                 |  3 ---
 docker/Dockerfile.ci                       | 14 +++++++++++
 tests/single_stage/test_diffusion_model.py | 28 ++++++++++++++++++++++
 4 files changed, 67 insertions(+), 3 deletions(-)
 create mode 100644 docker/Dockerfile.ci
 create mode 100644 tests/single_stage/test_diffusion_model.py

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 7bc41512e..403675563 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -1,6 +1,31 @@
 steps:
+  - label: ":docker: Build image"
+    key: image-build
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+    agents:
+      queue: "cpu_queue_premerge_us_east_1"
+
   - label: "Simple Unit Test"
     commands:
       - ".buildkite/scripts/simple_test.sh"
     agents:
       queue: "cpu_queue_premerge"
+  - label: "Diffusion Model Test"
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/single_stage/test_diffusion_model.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
diff --git a/.gitignore b/.gitignore
index 1c15c4395..d523f9ce2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -156,9 +156,6 @@ cython_debug/
 # cursor
 .cursor/
 
-# docker
-docker/
-
 # scripts
 /scripts/
 
diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci
new file mode 100644
index 000000000..0d513c054
--- /dev/null
+++ b/docker/Dockerfile.ci
@@ -0,0 +1,14 @@
+ARG VLLM_BASE_IMAGE=vllm/vllm-openai
+ARG VLLM_BASE_TAG=v0.11.0
+FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
+ARG APP_DIR=/workspace/vllm-omni
+WORKDIR ${APP_DIR}
+
+COPY . .
+
+# Install vllm-omni into the same uv-managed Python environment used by the base image.
+RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+ENTRYPOINT []
diff --git a/tests/single_stage/test_diffusion_model.py b/tests/single_stage/test_diffusion_model.py
new file mode 100644
index 000000000..dc00a36e5
--- /dev/null
+++ b/tests/single_stage/test_diffusion_model.py
@@ -0,0 +1,28 @@
+import pytest
+import torch
+
+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)
+    # high resolution may cause OOM on L4
+    height = 256
+    width = 256
+    images = m.generate(
+        "a photo of a cat sitting on a laptop keyboard",
+        height=height,
+        width=width,
+        num_inference_steps=2,
+        guidance_scale=0.0,
+        generator=torch.Generator("cuda").manual_seed(42),
+        num_outputs_per_prompt=2,
+    )
+    assert len(images) == 2
+    # check image size
+    assert images[0].width == width
+    assert images[0].height == height
+    images[0].save("z_image_output.png")