[Model] Add circlestone-labs/Anima (#4083)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

tests/model_tests/diffusion/anima_builder.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Local random-weight Anima assets and real-checkpoint assets for common tests."""
+
+import json
+import tempfile
+from pathlib import Path
+
+import torch
+from diffusers import AutoencoderKLQwenImage, FlowMatchEulerDiscreteScheduler
+from huggingface_hub import hf_hub_download, snapshot_download
+from safetensors.torch import save_file
+from tokenizers import Tokenizer
+from tokenizers.models import Unigram, WordLevel
+from tokenizers.pre_tokenizers import Whitespace
+from tokenizers.processors import TemplateProcessing
+from transformers import PreTrainedTokenizerFast, Qwen3Config, Qwen3Model, T5TokenizerFast
+
+from vllm_omni.diffusion.models.anima.anima_text_conditioner import (
+    ANIMA_TEXT_CONDITIONER_CONFIG,
+    AnimaTextConditioner,
+)
+from vllm_omni.diffusion.models.anima.anima_transformer import ANIMA_TRANSFORMER_CONFIG, AnimaTransformer3DModel
+
+CHECKPOINT_FILENAME = "anima.safetensors"
+
+
+def tiny_anima_builder() -> str:
+    """Build all components locally, without downloading any model weights or tokenizers."""
+    model_dir = Path(tempfile.mkdtemp(prefix="tiny-anima-"))
+    transformer_config = {
+        **ANIMA_TRANSFORMER_CONFIG,
+        "num_attention_heads": 2,
+        "attention_head_dim": 32,
+        "num_layers": 1,
+        "mlp_ratio": 1.0,
+        "text_embed_dim": 32,
+        "adaln_lora_dim": 8,
+        "max_size": (1, 64, 64),
+    }
+    conditioner_config = {
+        **ANIMA_TEXT_CONDITIONER_CONFIG,
+        "source_dim": 32,
+        "target_dim": 32,
+        "model_dim": 32,
+        "num_layers": 1,
+        "num_attention_heads": 2,
+        "mlp_ratio": 1.0,
+        "target_vocab_size": 8,
+        "min_sequence_length": 8,
+    }
+    transformer = AnimaTransformer3DModel(**transformer_config).to(torch.bfloat16)
+    conditioner = AnimaTextConditioner(**conditioner_config).to(torch.bfloat16)
+    save_file(
+        {
+            **{f"transformer.{name}": tensor for name, tensor in transformer.state_dict().items()},
+            **{f"text_conditioner.{name}": tensor for name, tensor in conditioner.state_dict().items()},
+        },
+        str(model_dir / CHECKPOINT_FILENAME),
+    )
+    (model_dir / "anima.json").write_text(
+        json.dumps({"transformer": transformer_config, "text_conditioner": conditioner_config})
+    )
+    Qwen3Model(
+        Qwen3Config(
+            vocab_size=8,
+            hidden_size=32,
+            intermediate_size=64,
+            num_hidden_layers=1,
+            num_attention_heads=2,
+            num_key_value_heads=2,
+            head_dim=16,
+        )
+    ).to(torch.bfloat16).save_pretrained(model_dir / "text_encoder")
+    AutoencoderKLQwenImage(
+        base_dim=4,
+        z_dim=16,
+        dim_mult=[1, 1, 1, 1],
+        num_res_blocks=1,
+        latents_mean=[0.0] * 16,
+        latents_std=[1.0] * 16,
+    ).to(torch.bfloat16).save_pretrained(model_dir / "vae")
+    tokenizer = Tokenizer(WordLevel({"<pad>": 0, "</s>": 1, "<unk>": 2, "Dummy": 3, "prompt": 4}, unk_token="<unk>"))
+    tokenizer.pre_tokenizer = Whitespace()
+    tokenizer.post_processor = TemplateProcessing(single="$A </s>", special_tokens=[("</s>", 1)])
+    PreTrainedTokenizerFast(
+        tokenizer_object=tokenizer, pad_token="<pad>", eos_token="</s>", unk_token="<unk>"
+    ).save_pretrained(model_dir / "tokenizer")
+    t5_tokenizer = Tokenizer(
+        Unigram(
+            [("<pad>", 0.0), ("</s>", 0.0), ("<unk>", 0.0), ("▁", -1.0), ("Dummy", -1.0), ("prompt", -1.0)], unk_id=2
+        )
+    )
+    T5TokenizerFast(tokenizer_object=t5_tokenizer, extra_ids=0).save_pretrained(model_dir / "t5_tokenizer")
+    FlowMatchEulerDiscreteScheduler(shift=3.0).save_pretrained(model_dir / "scheduler")
+    return str(model_dir)
+
+
+def real_anima_model() -> str:
+    """Keep cached real assets intact; return a disposable directory of links."""
+    components = Path(
+        snapshot_download(
+            "circlestone-labs/Anima-Base-v1.0-Diffusers",
+            allow_patterns=["text_encoder/*", "vae/*", "tokenizer/*", "t5_tokenizer/*", "scheduler/*"],
+        )
+    )
+    checkpoint = hf_hub_download(
+        "circlestone-labs/Anima", filename="split_files/diffusion_models/anima-base-v1.0.safetensors"
+    )
+    model_dir = Path(tempfile.mkdtemp(prefix="real-anima-"))
+    (model_dir / CHECKPOINT_FILENAME).symlink_to(checkpoint)
+    for name in ("text_encoder", "vae", "tokenizer", "t5_tokenizer", "scheduler"):
+        if (components / name).is_dir():
+            (model_dir / name).symlink_to(components / name, target_is_directory=True)
+    return str(model_dir)

tests/model_tests/diffusion/config_types.py

@@ -1,10 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """
 Common definitions for controlling what tests run where.
 """
 
 from collections.abc import Callable
 from enum import Enum
-from typing import NamedTuple, TypeAlias
+from typing import Any, NamedTuple, TypeAlias
 
 from pytest import MarkDecorator
 
@@ -45,7 +48,9 @@ class DiffusionModelTestOpts(NamedTuple):
     # HF model name for real-weight tests (advanced_model / full_model level).
     # For now, whether we use the real weights vs tiny weights in the common tests
     # depends on the run level.
-    model: str
+    # A callable may assemble native assets in a temporary directory; the fixture
+    # removes that directory after the tests, so it must not return a cache root.
+    model: str | Callable[[], str]
 
     # Creates a tiny model for the given architecture. We should always use tiny
     # model weights for tests that do not require us to check the model quality.
@@ -81,9 +86,12 @@ class DiffusionModelTestOpts(NamedTuple):
     # set False for pipelines that mandate an image and fail closed without one (SANA-Video I2V).
     check_i2v_t2v_divergence: bool = True
 
+    # Native single-file models: builders/resolvers return the containing directory.
+    checkpoint_filename: str | None = None
+
 
 ### Mappings & utils for building offline Omni() instances given a list of enabled accelerations
-ACC_OMNI_KWARGS = {
+ACC_OMNI_KWARGS: dict[DiffusionAccs, dict[str, Any]] = {
     DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_use_tiling": True},
     DiffusionAccs.CPU_OFFLOAD: {"enable_cpu_offload": True},
     DiffusionAccs.LAYERWISE_OFFLOAD: {"enable_layerwise_offload": True},

tests/model_tests/diffusion/model_settings.py

@@ -1,4 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from tests.model_tests.diffusion import diff_model_builders
+from tests.model_tests.diffusion.anima_builder import CHECKPOINT_FILENAME, real_anima_model, tiny_anima_builder
 from tests.model_tests.diffusion.config_types import (
     DiffusionAccs,
     DiffusionModelTestOpts,
@@ -29,6 +33,12 @@ from tests.model_tests.diffusion.config_types import (
 # $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline
 #   ^ Runs all test groups for Flux2KleinPipeline only
 DIFFUSION_TEST_SETTINGS = {
+    "AnimaPipeline": DiffusionModelTestOpts(
+        model=real_anima_model,
+        builder=tiny_anima_builder,
+        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+        checkpoint_filename=CHECKPOINT_FILENAME,
+    ),
     "Flux2KleinPipeline": DiffusionModelTestOpts(
         model="black-forest-labs/FLUX.2-klein-4B",
         builder=diff_model_builders.tiny_flux2_klein_builder,

tests/model_tests/diffusion/test_common_offline.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+from pathlib import Path
+
 import pytest
 
 from tests.model_tests.diffusion.case_filtering import get_parametrized_options
@@ -47,10 +49,17 @@ def test_pipeline_on_supported_tasks(
     # since starting the server can take 10+ seconds, even for tiny models.
     #
     # NOTE: Be sure to install pytest-subtests if you're running on pytest < 9
+    settings = DIFFUSION_TEST_SETTINGS[model_name]
+    model_path = tiny_model_paths[model_name]
+    model_kwargs = {}
+    if settings.checkpoint_filename is not None:
+        model_path = str(Path(model_path) / settings.checkpoint_filename)
+        model_kwargs["model_class_name"] = model_name
     omni = build_omni_from_diff_accelerations(
         accelerations=accelerations,
-        model=tiny_model_paths[model_name],
+        model=model_path,
         enforce_eager=True,
+        **model_kwargs,
     )
     try:
         for task_type in supported_tasks:

tests/model_tests/diffusion/test_common_online.py

@@ -7,6 +7,8 @@ online serving stack (CLI arg parsing, subprocess, API routing, response
 encoding) using tiny models.
 """
 
+from pathlib import Path
+
 import pytest
 from xdist import is_xdist_worker
 
@@ -68,6 +70,10 @@ def test_online_on_supported_tasks(
     model_path = tiny_model_paths[model_name]
     server_args = build_server_args_from_diff_accelerations(accelerations)
     server_args.append("--enforce-eager")
+    settings = DIFFUSION_TEST_SETTINGS[model_name]
+    if settings.checkpoint_filename is not None:
+        model_path = str(Path(model_path) / settings.checkpoint_filename)
+        server_args.extend(["--model-class-name", model_name])
 
     with OmniServer(model_path, server_args) as server:
         # TODO: We may want to revisit run_level validation here,

tests/model_tests/diffusion/test_native_model_settings.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from contextlib import nullcontext
+from pathlib import Path
+from shutil import rmtree
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from tests.model_tests.diffusion import anima_builder, test_common_offline, test_common_online
+from tests.model_tests.diffusion.config_types import DiffusionTasks
+
+pytestmark = [pytest.mark.cpu, pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize("model_name", ["AnimaPipeline", "FluxPipeline"])
+def test_common_runners_resolve_native_checkpoint(model_name, tmp_path, monkeypatch):
+    """Both entry points receive the file and class; directory models keep their existing arguments."""
+    paths = {model_name: str(tmp_path)}
+    subtests = SimpleNamespace(test=lambda **kwargs: nullcontext())
+    build_omni = MagicMock()
+    monkeypatch.setattr(test_common_offline, "build_omni_from_diff_accelerations", build_omni)
+    monkeypatch.setattr(test_common_offline, "run_and_validate_text_to_image_request", MagicMock())
+    test_common_offline.test_pipeline_on_supported_tasks(
+        model_name, None, [DiffusionTasks.TEXT_TO_IMAGE], False, False, False, paths, subtests
+    )
+    server = MagicMock()
+    monkeypatch.setattr(test_common_online, "OmniServer", server)
+    monkeypatch.setattr(test_common_online, "OnlineOmniClient", MagicMock())
+    monkeypatch.setattr(test_common_online, "run_and_validate_online_text_to_image_request", MagicMock())
+    test_common_online.test_online_on_supported_tasks(
+        model_name, None, [DiffusionTasks.TEXT_TO_IMAGE], False, False, paths, "core_model", subtests
+    )
+    if model_name == "AnimaPipeline":
+        checkpoint = str(tmp_path / anima_builder.CHECKPOINT_FILENAME)
+        build_omni.assert_called_once_with(
+            accelerations=None, model=checkpoint, enforce_eager=True, model_class_name=model_name
+        )
+        server.assert_called_once_with(checkpoint, ["--enforce-eager", "--model-class-name", model_name])
+    else:
+        build_omni.assert_called_once_with(accelerations=None, model=str(tmp_path), enforce_eager=True)
+        server.assert_called_once_with(str(tmp_path), ["--enforce-eager"])
+
+
+def test_real_anima_assets_leave_cache_intact(tmp_path, monkeypatch):
+    """The fixture may delete the assembled directory without deleting cached model assets."""
+    components = tmp_path / "components"
+    for name in ("text_encoder", "vae", "tokenizer", "t5_tokenizer"):
+        (components / name).mkdir(parents=True)
+        (components / name / "asset").write_text(name)
+    checkpoint = tmp_path / "real.safetensors"
+    checkpoint.write_bytes(b"checkpoint")
+    monkeypatch.setattr(anima_builder, "snapshot_download", lambda *args, **kwargs: str(components))
+    monkeypatch.setattr(anima_builder, "hf_hub_download", lambda *args, **kwargs: str(checkpoint))
+    directory = Path(anima_builder.real_anima_model())
+    try:
+        assert (directory / anima_builder.CHECKPOINT_FILENAME).read_bytes() == b"checkpoint"
+        assert (directory / "vae" / "asset").read_text() == "vae"
+        # Missing scheduler is supported by Anima's default scheduler fallback.
+        assert not (directory / "scheduler").exists()
+    finally:
+        rmtree(directory)
+    assert checkpoint.read_bytes() == b"checkpoint"
+    assert (components / "vae" / "asset").read_text() == "vae"

tests/model_tests/diffusion/utils.py

@@ -1,8 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """
 Utilities for resolving real models to their tiny model equivalents.
 """
 
 import logging
+from pathlib import Path
 
 from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
 from vllm_omni.diffusion.data import resolve_model_class_name
@@ -30,4 +34,7 @@ def resolve_tiny_model_path(model: str) -> str:
         )
         return model
 
-    return test_opts.builder()
+    model_path = test_opts.builder()
+    if test_opts.checkpoint_filename is not None:
+        return str(Path(model_path) / test_opts.checkpoint_filename)
+    return model_path