64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Offline e2e smoke test for native Anima single-file checkpoints."""
|
|
|
|
import os
|
|
|
|
import pytest
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.runtime import OmniRunner, OmniRunnerHandler
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
CHECKPOINT_ENV = "ANIMA_CHECKPOINT_PATH"
|
|
COMPONENTS_ENV = "ANIMA_COMPONENTS_PATH"
|
|
CHECKPOINT_REPO = "circlestone-labs/Anima"
|
|
CHECKPOINT_FILENAME = "split_files/diffusion_models/anima-base-v1.0.safetensors"
|
|
COMPONENTS_REPO = "circlestone-labs/Anima-Base-v1.0-Diffusers"
|
|
PROMPT = "A cinematic close-up of a glass teapot on a wooden table."
|
|
|
|
|
|
def _resolve_anima_assets() -> tuple[str, str]:
|
|
checkpoint = os.environ.get(CHECKPOINT_ENV)
|
|
components = os.environ.get(COMPONENTS_ENV)
|
|
if checkpoint and components:
|
|
return checkpoint, components
|
|
|
|
from huggingface_hub import hf_hub_download, snapshot_download
|
|
|
|
if checkpoint is None:
|
|
checkpoint = hf_hub_download(
|
|
repo_id=CHECKPOINT_REPO,
|
|
filename=CHECKPOINT_FILENAME,
|
|
)
|
|
if components is None:
|
|
components = snapshot_download(repo_id=COMPONENTS_REPO)
|
|
return checkpoint, components
|
|
|
|
|
|
@pytest.mark.full_model
|
|
@pytest.mark.diffusion
|
|
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
|
|
@pytest.mark.parametrize("guidance_scale", [1.0, 4.0])
|
|
def test_anima_text_to_image(guidance_scale):
|
|
checkpoint, components = _resolve_anima_assets()
|
|
request_config = {
|
|
"model": checkpoint,
|
|
"prompt": PROMPT,
|
|
"sampling_params": OmniDiffusionSamplingParams(
|
|
height=256,
|
|
width=256,
|
|
num_inference_steps=1,
|
|
guidance_scale=guidance_scale,
|
|
seed=42,
|
|
),
|
|
}
|
|
with OmniRunner(
|
|
checkpoint,
|
|
seed=42,
|
|
model_class_name="AnimaPipeline",
|
|
custom_pipeline_args={"components_path": components},
|
|
) as runner:
|
|
OmniRunnerHandler(runner).send_diffusion_request(request_config)
|