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

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
2026-09-19 16:16:38 +00:00
1 parent 698f716012
commit 415005be22
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
@@ -0,0 +1,454 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import json
+from pathlib import Path
+from shutil import rmtree
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from tests.model_tests.diffusion.anima_builder import CHECKPOINT_FILENAME, tiny_anima_builder
+from vllm_omni.diffusion.attention import selector
+from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
+from vllm_omni.diffusion.models.anima import pipeline_anima
+from vllm_omni.diffusion.models.anima.anima_text_conditioner import AnimaTextConditioner
+from vllm_omni.diffusion.models.anima.anima_transformer import AnimaTransformer3DModel
+from vllm_omni.diffusion.models.anima.pipeline_anima import AnimaPipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize(
+    "offload_args",
+    [
+        {"enable_cpu_offload": True},
+        {"enable_layerwise_offload": True},
+        {"enable_distributed_layerwise_offload": True},
+        {"diffusion_offload_config": {"mode": "module", "components": ["dit"]}},
+        {
+            "diffusion_offload_config": {
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"weight_transfer": "rank-local"}},
+            }
+        },
+    ],
+)
+def test_native_anima_rejects_offload(offload_args):
+    config = OmniDiffusionConfig(model="anima.safetensors", **offload_args)
+    with pytest.raises(NotImplementedError, match="offload"):
+        AnimaPipeline(od_config=config, device=torch.device("cpu"))
+
+
+def test_native_anima_initializes_without_offload():
+    config = OmniDiffusionConfig(model="anima.safetensors")
+    AnimaPipeline(od_config=config, device=torch.device("cpu"))
+
+
+def test_anima_registration() -> None:
+    """The native Anima pipeline is registered and importable."""
+    from vllm_omni.diffusion.registry import DiffusionModelRegistry
+
+    assert DiffusionModelRegistry._try_load_model_cls("AnimaPipeline") is not None
+
+
+@pytest.mark.parametrize("model_class_name", ["AnimaPipeline", "AnimaModularPipeline"])
+def test_enrich_config_native_anima_checkpoint(tmp_path: Path, model_class_name: str) -> None:
+    """Native and reference Anima class names resolve to AnimaPipeline."""
+    dummy_checkpoint = tmp_path / "model.safetensors"
+    dummy_checkpoint.write_text("dummy")
+
+    config = OmniDiffusionConfig(
+        model=str(dummy_checkpoint),
+        model_class_name=model_class_name,
+    )
+    config.enrich_config()
+
+    assert config.diffusion_load_format == "default"
+    assert config.model_class_name == "AnimaPipeline"
+    assert config.diffusers_pipeline_cls is None
+
+
+def test_native_anima_component_paths_use_custom_pipeline_args() -> None:
+    """Anima resolves shared and component-specific paths from native args."""
+    pipeline = AnimaPipeline.__new__(AnimaPipeline)
+    pipeline.od_config = SimpleNamespace(
+        custom_pipeline_args={
+            "components_path": "/tmp/anima-components",
+            "vae_path": "/tmp/anima-vae",
+            "text_encoder_model": "/tmp/anima-text-encoder",
+        }
+    )
+
+    assert pipeline._component_path("vae", "/tmp/default") == "/tmp/anima-vae"
+    assert pipeline._component_path("text_encoder", "/tmp/default") == "/tmp/anima-text-encoder"
+    assert pipeline._component_path("tokenizer", "/tmp/default") == "/tmp/default"
+
+
+def test_native_anima_converts_original_cosmos_transformer_keys() -> None:
+    """Original Cosmos checkpoint names map to native module names."""
+    converted = AnimaPipeline._convert_original_transformer_state_dict(
+        {
+            "net.x_embedder.proj.1.weight": "patch",
+            "net.blocks.0.self_attn.q_proj.weight": "q",
+            "net.blocks.0.self_attn.q_norm.weight": "q_norm",
+            "net.blocks.0.mlp.layer1.weight": "mlp",
+            "net.final_layer.linear.weight": "out",
+            "net.accum_iteration": "drop",
+        }
+    )
+
+    assert converted == {
+        "patch_embed.proj.weight": "patch",
+        "transformer_blocks.0.attn1.to_q.weight": "q",
+        "transformer_blocks.0.attn1.norm_q.weight": "q_norm",
+        "transformer_blocks.0.ff.net.0.proj.weight": "mlp",
+        "proj_out.weight": "out",
+    }
+
+
+def test_native_anima_resolves_vae_scale_factor_from_loaded_vae() -> None:
+    """The loaded VAE determines the spatial scale factor."""
+    from vllm_omni.diffusion.models.anima.pipeline_anima import _anima_vae_scale_factor_from_vae
+
+    vae = SimpleNamespace(config=SimpleNamespace(spatial_compression_ratio=16))
+
+    assert _anima_vae_scale_factor_from_vae(vae) == 16
+
+
+@pytest.mark.parametrize("use_config_file", [False, True])
+def test_native_anima_loads_synthetic_checkpoint(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, use_config_file: bool
+) -> None:
+    """Native modules load exact weights from an Anima safetensors checkpoint."""
+    tiny_transformer_config = {
+        "in_channels": 1,
+        "out_channels": 1,
+        "num_attention_heads": 1,
+        "attention_head_dim": 12,
+        "num_layers": 1,
+        "mlp_ratio": 1.0,
+        "text_embed_dim": 4,
+        "adaln_lora_dim": 3,
+        "max_size": (1, 2, 2),
+        "patch_size": (1, 1, 1),
+        "rope_scale": (1.0, 1.0, 1.0),
+        "concat_padding_mask": True,
+        "extra_pos_embed_type": None,
+    }
+    tiny_text_conditioner_config = {
+        "source_dim": 4,
+        "target_dim": 4,
+        "model_dim": 4,
+        "num_layers": 1,
+        "num_attention_heads": 1,
+        "target_vocab_size": 8,
+        "min_sequence_length": 4,
+    }
+    if use_config_file:
+        (tmp_path / "anima.json").write_text(
+            json.dumps({"transformer": tiny_transformer_config, "text_conditioner": tiny_text_conditioner_config})
+        )
+    else:
+        monkeypatch.setattr(pipeline_anima, "ANIMA_TRANSFORMER_CONFIG", tiny_transformer_config)
+        monkeypatch.setattr(pipeline_anima, "ANIMA_TEXT_CONDITIONER_CONFIG", tiny_text_conditioner_config)
+
+    monkeypatch.setattr(selector, "_cached_get_backend_cls", lambda *_args, **_kwargs: SDPABackend)
+    transformer = AnimaTransformer3DModel(**tiny_transformer_config)
+    text_conditioner = AnimaTextConditioner(**tiny_text_conditioner_config)
+    transformer_state = {name: tensor.detach().clone() for name, tensor in transformer.state_dict().items()}
+    text_conditioner_state = {name: tensor.detach().clone() for name, tensor in text_conditioner.state_dict().items()}
+    checkpoint_state = {
+        **{f"transformer.{name}": tensor for name, tensor in transformer_state.items()},
+        **{f"text_conditioner.{name}": tensor for name, tensor in text_conditioner_state.items()},
+    }
+
+    checkpoint_path = tmp_path / "anima.safetensors"
+    save_file(checkpoint_state, str(checkpoint_path))
+
+    pipeline = AnimaPipeline.__new__(AnimaPipeline)
+    pipeline.od_config = SimpleNamespace(model=str(checkpoint_path), dtype=torch.float32)
+    pipeline.device = torch.device("cpu")
+
+    def assert_loaded(loaded_transformer, loaded_text_conditioner):
+        for name, tensor in transformer_state.items():
+            assert torch.equal(loaded_transformer.state_dict()[name], tensor)
+        for name, tensor in text_conditioner_state.items():
+            assert torch.equal(loaded_text_conditioner.state_dict()[name], tensor)
+
+    loaded_transformer, loaded_text_conditioner = pipeline._load_native_denoiser_components(dict(checkpoint_state))
+    assert_loaded(loaded_transformer, loaded_text_conditioner)
+
+    loaded_transformer, loaded_text_conditioner = pipeline._load_native_denoiser_components()
+    assert_loaded(loaded_transformer, loaded_text_conditioner)
+
+    if use_config_file:
+        tiny_transformer_config["num_layers"] = 2
+        (tmp_path / "anima.json").write_text(
+            json.dumps({"transformer": tiny_transformer_config, "text_conditioner": tiny_text_conditioner_config})
+        )
+        with pytest.raises(RuntimeError, match="Missing key"):
+            pipeline._load_native_denoiser_components()
+
+
+def test_tiny_anima_generates_images(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Load every saved component and exercise CFG, decoding, and multiple outputs on CPU."""
+    monkeypatch.setattr(selector, "_cached_get_backend_cls", lambda *_args, **_kwargs: SDPABackend)
+    monkeypatch.setattr(current_omni_platform, "is_available", lambda: False)
+    monkeypatch.setattr(
+        DistributedAutoencoderKLQwenImage,
+        "init_distributed",
+        lambda self: setattr(self, "distributed_executor", SimpleNamespace(parallel_size=1)),
+    )
+    model_dir = tiny_anima_builder()
+    try:
+        config = OmniDiffusionConfig(
+            model=str(Path(model_dir) / CHECKPOINT_FILENAME), model_class_name="AnimaPipeline", dtype=torch.float32
+        )
+        pipeline = AnimaPipeline(od_config=config, device=torch.device("cpu"))
+        pipeline.load_weights()
+        pipeline.eval()
+        outputs = []
+        with torch.inference_mode():
+            for _ in range(2):
+                request = OmniDiffusionRequest(
+                    prompt="Dummy prompt",
+                    request_id="tiny-anima",
+                    sampling_params=OmniDiffusionSamplingParams(
+                        height=32,
+                        width=32,
+                        num_inference_steps=2,
+                        guidance_scale=4.0,
+                        num_outputs_per_prompt=2,
+                        generator=torch.Generator(device="cpu").manual_seed(42),
+                    ),
+                )
+                outputs.append(pipeline(DiffusionRequestBatch([request]))[0].output)
+        assert outputs[0].shape == (2, 3, 32, 32)
+        assert torch.isfinite(outputs[0]).all()
+        assert torch.equal(outputs[0], outputs[1])
+        assert not torch.equal(outputs[0][0], outputs[0][1])
+    finally:
+        rmtree(model_dir)
+
+
+def test_native_anima_profiler_setup_is_deferred_and_idempotent(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Profiler targets are wrapped once, after deferred components are loaded."""
+    pipeline = AnimaPipeline.__new__(AnimaPipeline)
+    pipeline.od_config = SimpleNamespace(enable_diffusion_pipeline_profiler=True)
+    pipeline._profiler_initialized = False
+    pipeline.vae = SimpleNamespace(decode=lambda: "decoded")
+    pipeline.text_encoder = SimpleNamespace(forward=lambda: "encoded")
+
+    monkeypatch.setattr(current_omni_platform, "is_available", lambda: False)
+    pipeline._setup_profiler()
+    wrapped_decode = pipeline.vae.decode
+    wrapped_text_encoder = pipeline.text_encoder.forward
+    pipeline._setup_profiler()
+
+    assert pipeline.vae.decode is wrapped_decode
+    assert pipeline.text_encoder.forward is wrapped_text_encoder
+    assert pipeline.vae.decode() == "decoded"
+    assert pipeline.text_encoder.forward() == "encoded"
+    assert "AnimaPipeline.vae.decode" in pipeline.stage_durations
+    assert "AnimaPipeline.text_encoder.forward" in pipeline.stage_durations
+
+
+def _make_anima_forward_probe():
+    pipeline = AnimaPipeline.__new__(AnimaPipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.vae_scale_factor = 8
+    pipeline.transformer = SimpleNamespace(dtype=torch.float32)
+    pipeline.text_encoder = SimpleNamespace(dtype=torch.float32)
+    pipeline._current_timestep = None
+    pipeline._num_timesteps = 0
+    pipeline._guidance_scale = 0.0
+    captured = {}
+
+    def encode_prompt(**_kwargs):
+        return {
+            "qwen_prompt_embeds": torch.zeros(1, 2, 4),
+            "qwen_attention_mask": torch.ones(1, 2),
+            "t5_input_ids": torch.ones(1, 2, dtype=torch.long),
+            "t5_attention_mask": torch.ones(1, 2),
+            "negative_qwen_prompt_embeds": torch.zeros(1, 2, 4),
+            "negative_qwen_attention_mask": torch.ones(1, 2),
+            "negative_t5_input_ids": torch.ones(1, 2, dtype=torch.long),
+            "negative_t5_attention_mask": torch.ones(1, 2),
+        }
+
+    def condition_prompt_embeds(**_kwargs):
+        return torch.zeros(1, 2, 4)
+
+    def prepare_latents(**kwargs):
+        captured["prepare_latents"] = kwargs
+        return torch.zeros(1, 16, 1, kwargs["height"] // 8, kwargs["width"] // 8)
+
+    def diffuse(**kwargs):
+        captured["diffuse"] = kwargs
+        return kwargs["latents"]
+
+    pipeline.encode_prompt = encode_prompt
+    pipeline.condition_prompt_embeds = condition_prompt_embeds
+    pipeline.prepare_latents = prepare_latents
+    pipeline.prepare_timesteps = lambda **_kwargs: (torch.ones(1), 1)
+    pipeline.diffuse = diffuse
+
+    def decode_latents(latents, output_type="pil"):
+        captured["output_type"] = output_type
+        return DiffusionOutput(output=latents)
+
+    pipeline.decode_latents = decode_latents
+    return pipeline, captured
+
+
+def test_native_anima_forward_uses_official_default_resolution() -> None:
+    """Forward uses Anima's reference resolution when none is requested."""
+    pipeline, captured = _make_anima_forward_probe()
+    req = OmniDiffusionRequest(
+        prompt="a red cube",
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_id="anima-defaults",
+    )
+
+    pipeline.forward(DiffusionRequestBatch([req]))
+
+    assert captured["prepare_latents"]["height"] == 1024
+    assert captured["prepare_latents"]["width"] == 1024
+
+
+def test_native_anima_forward_honors_sampling_output_type() -> None:
+    """Request-level output type controls Anima decoding."""
+    pipeline, captured = _make_anima_forward_probe()
+    req = OmniDiffusionRequest(
+        prompt="a red cube",
+        sampling_params=OmniDiffusionSamplingParams(output_type="latent"),
+        request_id="anima-output-type",
+    )
+
+    outputs = pipeline.forward(DiffusionRequestBatch([req]))
+
+    assert len(outputs) == 1
+    assert isinstance(outputs[0], DiffusionOutput)
+    assert outputs[0].output.shape == (1, 16, 1, 128, 128)
+    assert captured["output_type"] == "latent"
+
+
+@pytest.mark.parametrize("request_count", [0, 2])
+def test_native_anima_forward_rejects_non_single_request_batches(request_count: int) -> None:
+    pipeline, captured = _make_anima_forward_probe()
+    requests = [
+        OmniDiffusionRequest(
+            prompt="a red cube",
+            sampling_params=OmniDiffusionSamplingParams(),
+            request_id=f"anima-{index}",
+        )
+        for index in range(request_count)
+    ]
+
+    with pytest.raises(ValueError, match="supports one request per forward"):
+        pipeline.forward(DiffusionRequestBatch(requests))
+
+    assert captured == {}
+
+
+def test_native_anima_explicit_guidance_scale_drives_cfg_multiplier() -> None:
+    """An explicit guidance scale controls native classifier-free guidance."""
+    pipeline, captured = _make_anima_forward_probe()
+    req = OmniDiffusionRequest(
+        prompt="a red cube",
+        sampling_params=OmniDiffusionSamplingParams(guidance_scale=5.0),
+        request_id="anima-guidance",
+    )
+
+    pipeline.forward(DiffusionRequestBatch([req]))
+
+    assert captured["diffuse"]["do_true_cfg"] is True
+    assert captured["diffuse"]["true_cfg_scale"] == 5.0
+
+
+def test_native_anima_true_cfg_scale_overrides_guidance_multiplier() -> None:
+    """An explicit true-CFG scale overrides the guidance-scale multiplier."""
+    pipeline, captured = _make_anima_forward_probe()
+    req = OmniDiffusionRequest(
+        prompt="a red cube",
+        sampling_params=OmniDiffusionSamplingParams(guidance_scale=5.0, true_cfg_scale=3.0),
+        request_id="anima-true-cfg",
+    )
+
+    pipeline.forward(DiffusionRequestBatch([req]))
+
+    assert captured["diffuse"]["do_true_cfg"] is True
+    assert captured["diffuse"]["true_cfg_scale"] == 3.0
+
+
+def test_native_anima_cfg_equation() -> None:
+    """The denoising loop applies the standard true-CFG equation."""
+    pipeline = AnimaPipeline.__new__(AnimaPipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline._interrupt = False
+
+    class MockProgressBar:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc_val, exc_tb):
+            pass
+
+        def update(self):
+            pass
+
+    pipeline.progress_bar = lambda **_kwargs: MockProgressBar()
+
+    class MockTransformer:
+        dtype = torch.float32
+
+        def __call__(self, hidden_states, timestep, encoder_hidden_states, padding_mask, return_dict=False):
+            # Return double the encoder_hidden_states to simulate transformer output
+            return (encoder_hidden_states * 2.0,)
+
+    pipeline.transformer = MockTransformer()
+
+    class MockScheduler:
+        def set_begin_index(self, index):
+            pass
+
+        def step(self, noise_pred, t, latents, return_dict=False):
+            # Return the noise_pred itself to inspect it
+            return (noise_pred,)
+
+    pipeline.scheduler = MockScheduler()
+
+    prompt_embeds = torch.tensor([[[2.0, 3.0]]])
+    negative_prompt_embeds = torch.tensor([[[1.0, 2.0]]])
+    latents = torch.zeros(1, 16, 1, 16, 16)
+    padding_mask = torch.zeros(1, 1, 16, 16)
+    timesteps = torch.tensor([500.0])
+
+    # Cond = prompt_embeds * 2.0 = [4.0, 6.0]
+    # Uncond = negative_prompt_embeds * 2.0 = [2.0, 4.0]
+    # Expected noise_pred = uncond + true_cfg_scale * (cond - uncond)
+    #                     = [2.0, 4.0] + 4.0 * ([4.0, 6.0] - [2.0, 4.0])
+    #                     = [2.0, 4.0] + 4.0 * [2.0, 2.0]
+    #                     = [2.0, 4.0] + [8.0, 8.0] = [10.0, 12.0]
+    out_latents = pipeline.diffuse(
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_prompt_embeds,
+        latents=latents,
+        padding_mask=padding_mask,
+        timesteps=timesteps,
+        do_true_cfg=True,
+        true_cfg_scale=4.0,
+    )
+
+    expected = torch.tensor([[[10.0, 12.0]]])
+    assert torch.allclose(out_latents, expected)