vllm-project--vllm-omni
wehub-resource-sync
145 行
4.3 KiB
Python
145 行
4.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""End-to-end online serving test for Lance (single-stage).
|
|
|
|
Verifies that the Lance pipeline can serve text-to-image and image-edit
|
|
requests via the OpenAI-compatible chat completions API exposed by
|
|
``vllm-omni serve``.
|
|
|
|
Equivalent to running:
|
|
|
|
vllm-omni serve "bytedance-research/Lance" --omni \\
|
|
--pipeline lance --enforce-eager --trust-remote-code --port 8091
|
|
|
|
# text2img
|
|
python3 examples/online_serving/lance/openai_chat_client.py \\
|
|
--prompt "A cute corgi astronaut" --modality text2img
|
|
"""
|
|
|
|
import base64
|
|
import os
|
|
from io import BytesIO
|
|
|
|
import pytest
|
|
from vllm.assets.image import ImageAsset
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.runtime import OmniServerParams
|
|
from vllm_omni.diffusion.models.lance.prompts import (
|
|
VIDEO_PAD,
|
|
VISION_END,
|
|
VISION_START,
|
|
render_lance_prompt,
|
|
)
|
|
|
|
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
|
|
|
MODEL = "bytedance-research/Lance"
|
|
|
|
TEXT2IMG_PROMPT = "A cute corgi astronaut on the moon, cinematic"
|
|
IMG2IMG_PROMPT = "Convert this into a vibrant cartoon-style illustration"
|
|
|
|
# Lance vision block: <|vision_start|><|video_pad|><|vision_end|>
|
|
_VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}"
|
|
|
|
# Lance is single-stage diffusion — no deploy YAML. Pass the engine
|
|
# knobs that used to live in ``vllm_omni/deploy/lance.yaml`` via CLI
|
|
# flags so ``create_default_diffusion`` builds the stage config.
|
|
_LANCE_SERVE_ARGS = [
|
|
"--pipeline",
|
|
"lance",
|
|
"--max-num-batched-tokens",
|
|
"32768",
|
|
"--max-num-seqs",
|
|
"1",
|
|
"--enforce-eager",
|
|
"--trust-remote-code",
|
|
"--no-enable-prefix-caching",
|
|
"--no-async-chunk",
|
|
]
|
|
|
|
test_params = [
|
|
OmniServerParams(
|
|
model=MODEL,
|
|
server_args=_LANCE_SERVE_ARGS,
|
|
stage_init_timeout=300,
|
|
),
|
|
]
|
|
|
|
|
|
def _build_text2img_messages(prompt: str) -> list[dict]:
|
|
"""Build OpenAI-format messages for text2img generation."""
|
|
rendered = render_lance_prompt("t2i", prompt)
|
|
return [
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "text", "text": rendered}],
|
|
}
|
|
]
|
|
|
|
|
|
def _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]:
|
|
"""Build OpenAI-format messages for img2img generation."""
|
|
rendered = render_lance_prompt("image_edit", prompt, vision_token=_VISION_BLOCK)
|
|
return [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": rendered},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:image/jpeg;base64,{image_b64}"},
|
|
},
|
|
],
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.advanced_model
|
|
@pytest.mark.diffusion
|
|
@hardware_test(res={"cuda": "H100"})
|
|
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
|
def test_lance_text2img_online(omni_server, openai_client) -> None:
|
|
"""Lance text2img via the OpenAI-compatible chat completions API."""
|
|
request_config = {
|
|
"model": omni_server.model,
|
|
"messages": _build_text2img_messages(TEXT2IMG_PROMPT),
|
|
"modalities": ["image"],
|
|
"extra_body": {
|
|
"height": 512,
|
|
"width": 512,
|
|
"num_inference_steps": 2,
|
|
"guidance_scale": 0.0,
|
|
"seed": 42,
|
|
},
|
|
}
|
|
|
|
openai_client.send_diffusion_request(request_config)
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.advanced_model
|
|
@pytest.mark.diffusion
|
|
@hardware_test(res={"cuda": "H100"})
|
|
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
|
def test_lance_img2img_online(omni_server, openai_client) -> None:
|
|
"""Lance image_edit via the OpenAI-compatible chat completions API."""
|
|
input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
|
|
buffer = BytesIO()
|
|
input_image.save(buffer, format="JPEG")
|
|
image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
|
|
|
request_config = {
|
|
"model": omni_server.model,
|
|
"messages": _build_img2img_messages(IMG2IMG_PROMPT, image_b64),
|
|
"modalities": ["image"],
|
|
"extra_body": {
|
|
"num_inference_steps": 2,
|
|
"guidance_scale": 0.0,
|
|
"seed": 42,
|
|
},
|
|
}
|
|
|
|
openai_client.send_diffusion_request(request_config)
|