Nick CaoandClaude
6646a66d41
Signed-off-by: Nick Cao <ncao@redhat.com> Co-authored-by: Claude <noreply@anthropic.com>
96 lines
2.7 KiB
Python
96 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
Expansion tests for LTX-2 two-stage pipelines.
|
|
|
|
Coverage:
|
|
- HSDP with LTX2TwoStagesPipeline (text-to-video)
|
|
- HSDP with LTX2ImageToVideoTwoStagesPipeline (image-to-video)
|
|
"""
|
|
|
|
import os
|
|
|
|
import pytest
|
|
|
|
from tests.helpers.mark import hardware_marks
|
|
from tests.helpers.media import generate_synthetic_image
|
|
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
|
|
|
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
|
|
|
MODEL = os.getenv("VLLM_TEST_LTX2_MODEL", "rootonchair/LTX-2-19b-distilled")
|
|
PROMPT = "A serene lake at sunset with mountains in the background."
|
|
PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
|
|
HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"]
|
|
|
|
|
|
def _cases():
|
|
cases = []
|
|
|
|
# T2V: LTX2TwoStagesPipeline
|
|
cases.append(
|
|
pytest.param(
|
|
OmniServerParams(
|
|
model=MODEL,
|
|
server_args=[
|
|
*HSDP_ARGS,
|
|
"--model-class-name",
|
|
"LTX2TwoStagesPipeline",
|
|
],
|
|
),
|
|
id="t2v_hsdp",
|
|
marks=PARALLEL_MARKS,
|
|
)
|
|
)
|
|
|
|
# I2V: LTX2ImageToVideoTwoStagesPipeline
|
|
cases.append(
|
|
pytest.param(
|
|
OmniServerParams(
|
|
model=MODEL,
|
|
server_args=[
|
|
*HSDP_ARGS,
|
|
"--model-class-name",
|
|
"LTX2ImageToVideoTwoStagesPipeline",
|
|
],
|
|
),
|
|
id="i2v_hsdp",
|
|
marks=PARALLEL_MARKS,
|
|
)
|
|
)
|
|
|
|
return cases
|
|
|
|
|
|
@pytest.mark.parametrize("omni_server", _cases(), indirect=True)
|
|
def test_ltx2_two_stage_hsdp(
|
|
omni_server: OmniServer,
|
|
openai_client: OpenAIClientHandler,
|
|
):
|
|
is_i2v = any("ImageToVideo" in arg for arg in omni_server.serve_args)
|
|
|
|
# The two-stage pipeline generates at the requested resolution then 2x
|
|
# upsamples via LTX2LatentUpsamplerModel, so output dimensions differ
|
|
# from the request. Omit height/width from form_data so the assertion
|
|
# helper skips the dimension check (it only asserts when the key is
|
|
# present). The pipeline falls back to its own defaults.
|
|
form_data = {
|
|
"prompt": PROMPT,
|
|
"model": omni_server.model,
|
|
"num_frames": 9,
|
|
"fps": 8,
|
|
"num_inference_steps": 2,
|
|
"guidance_scale": 1.0,
|
|
"seed": 42,
|
|
}
|
|
|
|
request_config = {
|
|
"model": omni_server.model,
|
|
"form_data": form_data,
|
|
}
|
|
|
|
if is_i2v:
|
|
request_config["image_reference"] = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
|
|
|
|
openai_client.send_video_diffusion_request(request_config)
|