From b4737d30669c94f1a85a8c688e4106fb5a46aa63 Mon Sep 17 00:00:00 2001 From: akshatvishu Date: Wed, 23 Sep 2026 22:34:51 +0530 Subject: [PATCH] [BugFix] Download Anima assets for README CI test (#7928) Signed-off-by: akshatvishu --- .../offline_inference/text_to_image/README.md | 2 ++ tests/diffusion/test_diffusion_engine.py | 11 +++++-- .../offline_inference/test_text_to_image.py | 29 +++++++++++++++++++ vllm_omni/diffusion/diffusion_engine.py | 7 ++--- 4 files changed, 41 insertions(+), 8 deletions(-) diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index 7664bb32d..65bfd47b7 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -222,6 +222,8 @@ Anima uses the official model dimensions by default. A checkpoint may provide a To load Anima, point `--model` to the single-file checkpoint path, pass the native pipeline class name using `--model-class-name`, and supply the converted components directory using `--custom-pipeline-args`: +Download the checkpoint and components using the [Anima recipe](../../../recipes/circlestone-labs/Anima.md#command---prepare-assets), then replace the paths below with their local locations. + ```bash python examples/offline_inference/text_to_image/text_to_image.py \ --model /path/to/models/anima-official/split_files/diffusion_models/anima-base-v1.0.safetensors \ diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index 77a6da5db..b7daa2e95 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -213,8 +213,11 @@ def test_request_started_output_is_emitted_only_for_opted_in_requests() -> None: class TestRequestBatchCapability: pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] - def test_supports_request_batch_uses_registered_model_class(self, monkeypatch: pytest.MonkeyPatch) -> None: - od_config = SimpleNamespace(model_class_name="BatchPipeline", custom_pipeline_args=None) + @pytest.mark.parametrize("custom_pipeline_args", [None, {}, {"components_path": "/tmp/anima-components"}]) + def test_supports_request_batch_uses_registered_model_class( + self, monkeypatch: pytest.MonkeyPatch, custom_pipeline_args: dict[str, Any] | None + ) -> None: + od_config = SimpleNamespace(model_class_name="BatchPipeline", custom_pipeline_args=custom_pipeline_args) monkeypatch.setattr( diffusion_engine_module.DiffusionModelRegistry, @@ -371,14 +374,16 @@ class TestRequestBatchCapability: diffusion_engine_module.supports_request_batch(od_config) registry_load.assert_not_called() + @pytest.mark.parametrize("custom_pipeline_args", [None, {}, {"components_path": "/tmp/anima-components"}]) def test_engine_uses_request_batch_mode_for_single_request_pipeline( self, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, + custom_pipeline_args: dict[str, Any] | None, ) -> None: od_config = SimpleNamespace( model_class_name="SinglePipeline", - custom_pipeline_args=None, + custom_pipeline_args=custom_pipeline_args, streaming_output=False, max_num_seqs=1, ) diff --git a/tests/examples/offline_inference/test_text_to_image.py b/tests/examples/offline_inference/test_text_to_image.py index 6289dcd33..a9783bb1d 100644 --- a/tests/examples/offline_inference/test_text_to_image.py +++ b/tests/examples/offline_inference/test_text_to_image.py @@ -6,9 +6,12 @@ Offline inference tests: text-to-image. See examples/offline_inference/text_to_image/README.md """ +import json +import shlex from pathlib import Path import pytest +from huggingface_hub import hf_hub_download, snapshot_download from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet from tests.helpers.assertions import assert_image_valid @@ -37,12 +40,38 @@ def _skip_readme_snippet(language: str, code: str, h2_title: str) -> tuple[bool, README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH, skipif=_skip_readme_snippet) +def _prepare_anima_snippet(snippet: ReadmeSnippet) -> ReadmeSnippet: + if "--model-class-name AnimaPipeline" not in snippet.code: + return snippet + + assert "/path/to/models/anima-official/split_files/diffusion_models/anima-base-v1.0.safetensors" in snippet.code, ( + "Anima README checkpoint placeholder changed; update _prepare_anima_snippet to match" + ) + + checkpoint = hf_hub_download( + repo_id="circlestone-labs/Anima", + filename="split_files/diffusion_models/anima-base-v1.0.safetensors", + ) + components = snapshot_download( + repo_id="circlestone-labs/Anima-Base-v1.0-Diffusers", + allow_patterns=["text_encoder/*", "vae/*", "tokenizer/*", "t5_tokenizer/*", "scheduler/*"], + ) + argv = shlex.split(snippet.code) + argv[argv.index("--model") + 1] = checkpoint + custom_args_index = argv.index("--custom-pipeline-args") + 1 + custom_args = json.loads(argv[custom_args_index]) + custom_args["components_path"] = components + argv[custom_args_index] = json.dumps(custom_args) + return snippet._replace(code=shlex.join(argv)) + + @pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id) def test_text_to_image(snippet: ReadmeSnippet, example_runner: ExampleRunner): should_skip, reason = snippet.skip if should_skip: pytest.skip(reason) + snippet = _prepare_anima_snippet(snippet) result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER)) for asset in result.assets: assert_image_valid(asset) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index a8385d7c6..babfac099 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -140,13 +140,10 @@ def _func_accepts_parameter(func: object | None, parameter_name: str) -> bool: def _resolve_custom_pipeline_cls(custom_pipeline_args: dict[str, Any] | None) -> type | None: - if custom_pipeline_args is None: + if not custom_pipeline_args or "pipeline_class" not in custom_pipeline_args: return None - try: - pipeline_cls = custom_pipeline_args["pipeline_class"] - except KeyError as exc: - raise ValueError("custom_pipeline_args must include 'pipeline_class'.") from exc + pipeline_cls = custom_pipeline_args["pipeline_class"] if isinstance(pipeline_cls, type): return pipeline_cls