[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
@@ -1,7 +1,7 @@
 # Custom Pipeline Extension Guide

 Transformer already support Custom Pipeline via
-https://github.com/huggingface/diffusers/blob/main/docs/source/en/using-diffusers/custom_pipeline_overview.md
+<https://github.com/huggingface/diffusers/blob/main/docs/source/en/using-diffusers/custom_pipeline_overview.md>

 This guide demonstrates how to use the newly added features for extending vLLM-Omni's diffusion pipeline with custom functionality.

@@ -20,8 +20,9 @@ Three main features enable custom pipeline extension:
 `WorkerWrapperBase` is a wrapper class that creates `DiffusionWorker` instances with optional extension support. It enables dynamic inheritance, allowing you to add custom methods and functionality to workers without modifying the base worker class.

 **Key capabilities:**
+
 - Dynamic worker class extension via `worker_extension_cls`
- Support for custom pipeline initialization via `custom_pipeline_args`
+- Support for custom pipeline initialization via `custom_pipeline_args["pipeline_class"]`
 - Method delegation to underlying worker
 - Attribute access forwarding

@@ -42,6 +43,7 @@ The `diffusion_load_format` parameter controls the initial diffusion model load.
 `CustomPipelineWorkerExtension` is a mixin class that extends `DiffusionWorker` with the ability to re-initialize the pipeline with a custom implementation.

 **Key method:**
+
 - `re_init_pipeline(custom_pipeline_args)`: Re-initializes the pipeline with custom arguments, properly cleaning up the old pipeline first

 **Location:** `vllm_omni/diffusion/worker/diffusion_worker.py`
@@ -150,6 +152,7 @@ omni = Omni(
    },
    worker_extension_cls=MyCustomExtension,
    # Note: worker_extension_cls is an internal parameter
-    # CustomPipelineWorkerExtension will automatically init pipeline when custom_pipeline_args is provided
+    # CustomPipelineWorkerExtension will automatically init pipeline when
+    # custom_pipeline_args includes "pipeline_class".
 )
 ```
@@ -113,6 +113,7 @@ th {
 | `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — |
 | `HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
 | `HiDreamO1ImagePipeline` | HiDream-O1-Image | `HiDream-ai/HiDream-O1-Image` | ✅︎ | | | | — |
+| `AnimaPipeline` | Anima | `circlestone-labs/Anima` | ✅︎ | ✅︎ | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/circlestone-labs/Anima.md) |

 The `Recipe` column links the corresponding deployment recipe. `Published` is
 preferred when a corresponding `recipes.vllm.ai` page exists; otherwise the
@@ -19,23 +19,23 @@ This folder provides several entrypoints for experimenting with text-to-image di

 ### Supported Models

-| Model | Image Shape  | Peak VRAM (GiB) * | Model Weights (GiB) |
+| Model | Image Shape | Peak VRAM (GiB) * | Model Weights (GiB) |
 | ----- | ----------- | ----------- | ----------------- |
 | `Qwen/Qwen-Image` | 1024 x 1024 | 60.0 | 53.7 |
-| `Qwen/Qwen-Image-2512` |1024 x 1024 | 60.0 | 53.7 |
+| `Qwen/Qwen-Image-2512` | 1024 x 1024 | 60.0 | 53.7 |
 | `Tongyi-MAI/Z-Image-Turbo` | 1024 x 1024 | 24.8 | 19.2 |
 | `stepfun-ai/NextStep-1.1` | 512 x 512 | 71.8 | 28.1 |
 | `meituan-longcat/LongCat-Image` | 1024 x 1024 | 71.2 | 27.3 |
 | `AIDC-AI/Ovis-Image-7B` | 1024 x 1024 | 71.8 | 17.1 |
-| `OmniGen2/OmniGen2` |  1024 x 1024 | 20.1 | 14.7 |
+| `OmniGen2/OmniGen2` | 1024 x 1024 | 20.1 | 14.7 |
 | `stabilityai/stable-diffusion-3.5-medium` | 1024 x 1024 | 20.1 | 15.6 |
 | `black-forest-labs/FLUX.1-dev` | 1024 x 1024 | 33.9 | 31.4 |
 | `black-forest-labs/FLUX.1-schnell` | 1024 x 1024 | 33.9 | 31.4 |
 | `black-forest-labs/FLUX.2-klein-4B` | 1024 x 1024 | 72.7 | 14.9 |
 | `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 |
 | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
-| `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3)  | 160 |
-| `HiDream-I1-Full` | 1024 x 1024 | 63.7  | 57.7 |
+| `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3) | 160 |
+| `HiDream-I1-Full` | 1024 x 1024 | 63.7 | 57.7 |
 | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | 1024 (Raw) / 2048 (Turbo) | — | ~30 |

 !!! info
@@ -112,9 +112,9 @@ python text_to_image.py \
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
 | `--lora-path` | str | — | Path to PEFT LoRA adapter folder or checkpoint file |
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
-| `--lora-backend` | str |`"peft"`| LoRA backend for loading LoRA adapters. Default: peft. Available options: peft, distill |
-| `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0.|
-| `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0.|
+| `--lora-backend` | str | `"peft"` | LoRA backend for loading LoRA adapters. Default: peft. Available options: peft, distill |
+| `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0. |
+| `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0. |
 | `--auxiliary-text-encoder` | str | `None` | Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full). |

 **NextStep-1.1 specific arguments:**
@@ -127,7 +127,7 @@ python text_to_image.py \
 | `--use-norm` | flag | off | Apply layer normalization to sampled tokens |

 > If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
-
+>
 > Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.

 ## More CLI Examples
@@ -216,6 +216,26 @@ python examples/offline_inference/text_to_image/text_to_image.py \
  --output /output.png
 ```

+### Anima Single-File Checkpoints
+
+Anima uses the official model dimensions by default. A checkpoint may provide a JSON file with the same stem (for example, `anima.json` beside `anima.safetensors`) containing `transformer` and `text_conditioner` configuration objects. These override the default dimensions and must match the checkpoint weights. The shared tiny-model tests use this to load small random-weight models. Components in `text_encoder/`, `vae/`, `tokenizer/`, `t5_tokenizer/`, and `scheduler/` beside the checkpoint are discovered automatically.
+
+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`:
+
+```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 \
+  --model-class-name AnimaPipeline \
+  --custom-pipeline-args '{"components_path": "/path/to/models/anima-components"}' \
+  --prompt "A cinematic close-up of a glass teapot on a wooden table." \
+  --seed 42 \
+  --guidance-scale 4.0 \
+  --num-inference-steps 50 \
+  --height 1024 \
+  --width 1024 \
+  --output anima_output.png
+```
+
 ### Krea 2

 Krea 2 is published as diffusers-format checkpoints: `krea/Krea-2-Turbo` (few-step distilled) and `krea/Krea-2-Raw`.
@@ -336,7 +356,7 @@ python text_to_image.py \

 LoRA adapters must be in PEFT format. A typical adapter directory structure:

-```
+```text
 lora_adapter/
 ├── adapter_config.json
 └── adapter_model.safetensors
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

 import argparse
 import functools
@@ -58,6 +58,7 @@ def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]


 parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
+parse_custom_pipeline_args = functools.partial(parse_json_object, flag_name="--custom-pipeline-args")


 def build_text_to_image_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]:
@@ -379,6 +380,18 @@ def parse_args() -> argparse.Namespace:
        default=None,
        help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
    )
+    parser.add_argument(
+        "--model-class-name",
+        type=str,
+        default=None,
+        help="Override the diffusion pipeline class name (e.g. AnimaPipeline).",
+    )
+    parser.add_argument(
+        "--custom-pipeline-args",
+        type=parse_custom_pipeline_args,
+        default=None,
+        help='JSON object passed to native/custom pipelines (e.g. \'{"components_path": "/path"}\').',
+    )
    parser.add_argument(
        "--trust-remote-code",
        action="store_true",
@@ -576,9 +589,13 @@ def main():
        omni_kwargs["trust_remote_code"] = True
    if args.deploy_config:
        omni_kwargs["deploy_config"] = args.deploy_config
-    if use_nextstep:
+    if args.model_class_name:
+        omni_kwargs["model_class_name"] = args.model_class_name
+    elif use_nextstep:
        # NextStep-1.1 requires explicit pipeline class
        omni_kwargs["model_class_name"] = "NextStep11Pipeline"
+    if args.custom_pipeline_args is not None:
+        omni_kwargs["custom_pipeline_args"] = args.custom_pipeline_args
    # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
    # gate is an engine-level config (offline analog of the server's --no-guardrails).
    if args.extra_body and "guardrails" in args.extra_body:
@@ -614,6 +631,10 @@ def main():
        print(f"  LoRA: scale={args.lora_scale}")
    if args.deploy_config:
        print(f"  deploy-config: {args.deploy_config}")
+    if args.model_class_name:
+        print(f"  Model class name: {args.model_class_name}")
+    if args.custom_pipeline_args is not None:
+        print(f"  Custom pipeline args: {args.custom_pipeline_args}")
    print(f"{'=' * 60}\n")

    # Build LoRA request when --lora-path is set
@@ -49,6 +49,7 @@ recipes/
 | [`BosonAI/Higgs-Audio-V3-TTS.md`](./BosonAI/Higgs-Audio-V3-TTS.md) | Online + offline multilingual TTS with voice cloning | 1x H100 80GB |
 | [`BreezeBlue/Breeze-TTS-2.md`](./BreezeBlue/Breeze-TTS-2.md) | Online + offline TTS (plain / voice design / voice clone / voice direction, streaming PCM) | 1x NVIDIA L20 48GB |
 | [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
+| [`circlestone-labs/Anima.md`](./circlestone-labs/Anima.md) | Native single-file text-to-image serving | 1x AMD MI300X |
 | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
 | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
 | [`HiDream-ai/HiDream-O1-Image.md`](./HiDream-ai/HiDream-O1-Image.md) | Text-to-image with shared offline example | 1x or 2x H100 80GB |
@@ -0,0 +1,158 @@
+# Anima
+
+> Native single-file diffusion text-to-image
+
+## Summary
+
+- Vendor: circlestone-labs
+- Model: [`circlestone-labs/Anima`](https://huggingface.co/circlestone-labs/Anima)
+- Task: text2img
+- Mode: Offline inference, Online serving (OpenAI-compatible API)
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to run Anima via vLLM-Omni's native `AnimaPipeline`. Anima is a
+Cosmos-style diffusion transformer text-to-image model distributed as a
+single-file transformer checkpoint, not as a normal Hugging Face model
+directory.
+
+The native path reads the Anima transformer checkpoint directly, converts
+original Cosmos transformer keys when needed, and loads the Cosmos transformer
+and text conditioner into native vLLM-Omni modules. Non-denoiser components
+such as `text_encoder`, `tokenizer`, `t5_tokenizer`, `vae`, and optionally
+`scheduler` must be supplied through a Diffusers-layout components directory.
+
+Native Anima currently supports baseline single-GPU execution. Cache-DiT,
+TeaCache, CPU offload, layer-wise offload, quantization, TP/SP, CFG parallel,
+HSDP, and step execution are not supported by `AnimaPipeline` yet.
+
+## References
+
+- Offline example:
+  [`examples/offline_inference/text_to_image/README.md`](../../examples/offline_inference/text_to_image/README.md)
+- Supported model entry:
+  [`docs/models/supported_models.md`](../../docs/models/supported_models.md)
+- HuggingFace model page:
+  [circlestone-labs/Anima](https://huggingface.co/circlestone-labs/Anima)
+- Diffusers-layout components:
+  [circlestone-labs/Anima-Base-v1.0-Diffusers](https://huggingface.co/circlestone-labs/Anima-Base-v1.0-Diffusers)
+
+## Hardware Support
+
+## ROCm
+
+### 1x AMD MI300X
+
+#### Environment
+
+- OS: Ubuntu 22.04.5 LTS (x86_64)
+- Python: 3.12.13
+- Driver / runtime: ROCm 7.2.53211, HIP runtime 7.2.53211, MIOpen 3.5.1
+- PyTorch: 2.10.0+git8514f05 built with ROCm 7.2.53211
+- GPU: 1x AMD MI300X
+- vLLM version: 0.23.0
+- vLLM-Omni version or commit: 0.1.dev2002+g704724675.rocm
+
+Recommended environment variables:
+
+```bash
+export VLLM_WORKER_MULTIPROC_METHOD=spawn
+export VLLM_ROCM_USE_AITER=0
+export PYTORCH_ROCM_ARCH=gfx942
+export PYTORCH_NVML_BASED_CUDA_CHECK=1
+```
+
+#### Command - prepare assets
+
+Download the official transformer checkpoint and the Diffusers-layout component
+directory:
+
+```bash
+mkdir -p /path/to/models/anima-official
+mkdir -p /path/to/models/anima-components
+
+hf download circlestone-labs/Anima \
+    split_files/diffusion_models/anima-base-v1.0.safetensors \
+    --local-dir /path/to/models/anima-official
+
+hf download circlestone-labs/Anima-Base-v1.0-Diffusers \
+    --local-dir /path/to/models/anima-components
+
+export ANIMA_CHECKPOINT=/path/to/models/anima-official/split_files/diffusion_models/anima-base-v1.0.safetensors
+export ANIMA_COMPONENTS=/path/to/models/anima-components
+```
+
+The `ANIMA_COMPONENTS` directory must be in Diffusers `from_pretrained()`
+layout. Raw auxiliary files such as `qwen_3_06b_base.safetensors` and
+`qwen_image_vae.safetensors` are converter inputs; they are not accepted
+directly as `components_path`.
+
+#### Command - text-to-image
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model "$ANIMA_CHECKPOINT" \
+    --model-class-name AnimaPipeline \
+    --custom-pipeline-args "{\"components_path\":\"$ANIMA_COMPONENTS\"}" \
+    --prompt "A cinematic close-up of a glass teapot on a wooden table." \
+    --seed 42 \
+    --guidance-scale 4.0 \
+    --num-inference-steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output /tmp/anima_output.png
+```
+
+Use 1024x1024, 50 denoising steps, `max_sequence_length=512`, one image per
+prompt, empty negative prompt, and CFG scale 4.0 when matching Anima's default
+Diffusers settings.
+
+#### Verification
+
+Check that `/tmp/anima_output.png` exists and contains a generated image.
+
+#### Notes
+
+- Key flags: `--model-class-name AnimaPipeline` selects the native Anima path;
+  `--custom-pipeline-args` supplies `components_path`.
+- Keep the default diffusion load format. No deploy config is required when a
+  local checkpoint and `--model-class-name AnimaPipeline` are provided.
+- Start with `max-concurrency=1` for correctness and latency validation.
+- Keep requests at the same resolution when comparing runs.
+- Do not enable parallelism, cache acceleration, offload, or quantized
+  checkpoint flags for Anima until support is added to `AnimaPipeline`.
+
+## Online Serving
+
+Anima supports text-to-image generation through the OpenAI-compatible image
+generation API.
+
+### Launch
+
+```bash
+vllm serve "$ANIMA_CHECKPOINT" \
+    --omni \
+    --port 8099 \
+    --model-class-name AnimaPipeline \
+    --custom-pipeline-args "{\"components_path\":\"$ANIMA_COMPONENTS\"}"
+```
+
+### Send requests
+
+```bash
+curl http://localhost:8099/v1/images/generations \
+    -H "Content-Type: application/json" \
+    -d "{
+      \"model\": \"$ANIMA_CHECKPOINT\",
+      \"prompt\": \"A cinematic close-up of a glass teapot on a wooden table.\",
+      \"size\": \"1024x1024\",
+      \"num_inference_steps\": 50,
+      \"guidance_scale\": 4.0,
+      \"seed\": 42
+    }"
+```
+
+The same generation knobs used by other text-to-image recipes apply:
+`num_inference_steps`, `seed`, `height` / `width` through `size`, and optional
+negative prompting.
@@ -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)
@@ -100,6 +100,28 @@ class TestWorkerWrapperBaseInitialization:
            skip_load_model=False,
        )

+    def test_custom_pipeline_args_without_pipeline_class_loads_normally(self, mocker: MockerFixture, mock_od_config):
+        """Test native pipeline args do not trigger custom pipeline re-initialization."""
+        custom_args = {"components_path": "/tmp/anima-components"}
+        mock_worker_init = mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
+
+        wrapper = WorkerWrapperBase(
+            gpu_id=0,
+            od_config=mock_od_config,
+            base_worker_class=DiffusionWorker,
+            custom_pipeline_args=custom_args,
+        )
+
+        assert wrapper.worker_extension_cls is None
+        assert CustomPipelineWorkerExtension not in wrapper.worker.__class__.__bases__
+        assert not hasattr(wrapper.worker, "re_init_pipeline")
+        mock_worker_init.assert_called_once_with(
+            local_rank=0,
+            rank=0,
+            od_config=mock_od_config,
+            skip_load_model=False,
+        )
+

 # -------------------------------------------------------------------------
 # Tests: Worker Extension Functionality
@@ -476,6 +498,12 @@ class TestCustomPipelineWorkerExtension:
            custom_pipeline_args=custom_args,
        )

+        mock_worker_class.assert_called_once_with(
+            local_rank=0,
+            rank=0,
+            od_config=mock_od_config,
+            skip_load_model=True,
+        )
        # Verify re_init_pipeline was called with custom_pipeline_args
        mock_worker_instance.re_init_pipeline.assert_called_once_with(custom_args)

@@ -0,0 +1,63 @@
+# 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)
@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Online e2e smoke test for native Anima single-file checkpoints."""
+
+import json
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer, OpenAIClientHandler, dummy_messages_from_mix_data
+
+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()
+    server_args = [
+        "--model-class-name",
+        "AnimaPipeline",
+        "--custom-pipeline-args",
+        json.dumps({"components_path": components}),
+    ]
+    with OmniServer(checkpoint, server_args) as omni_server:
+        openai_client = OpenAIClientHandler(
+            host=omni_server.host,
+            port=omni_server.port,
+            api_key="EMPTY",
+            run_level="full_model",
+            log_stats=omni_server.log_stats,
+        )
+        _send_text_to_image_request(omni_server, openai_client, guidance_scale)
+
+
+def _send_text_to_image_request(omni_server: OmniServer, openai_client: OpenAIClientHandler, guidance_scale: float):
+    messages = dummy_messages_from_mix_data(content_text=PROMPT)
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 256,
+            "width": 256,
+            "num_inference_steps": 1,
+            "guidance_scale": guidance_scale,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)
@@ -152,6 +152,37 @@ def test_omni_serve_accepts_explicit_model(argv: list[str], mocker: MockerFixtur
    cmd.validate(args)


+@pytest.mark.parametrize(
+    "model_class,exists,native",
+    [
+        ("AnimaPipeline", True, True),
+        ("AnimaModularPipeline", True, True),
+        ("AnimaPipeline", False, False),
+        ("FluxPipeline", True, False),
+        (None, True, False),
+    ],
+)
+def test_serve_native_checkpoint_validation(model_class, exists, native, tmp_path, mocker):
+    checkpoint = tmp_path / "anima.safetensors"
+    if exists:
+        checkpoint.touch()
+    argv = ["serve", str(checkpoint), "--omni"]
+    if model_class is not None:
+        argv.extend(["--model-class-name", model_class])
+    args = _parse_serve_args(argv)
+    detect_model = mocker.patch("vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", return_value=False)
+    validate_llm = mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args")
+
+    OmniServeCommand().validate(args)
+
+    if native:
+        detect_model.assert_not_called()
+        validate_llm.assert_not_called()
+    else:
+        detect_model.assert_called_once_with(str(checkpoint))
+        validate_llm.assert_called_once_with(args)
+
+
 def test_serve_parser_accepts_strategy_config() -> None:
    """``--strategy-config`` must parse onto the ``strategy_config`` dest and be
    forwarded as an explicit kwarg so the engine can overlay the strategy."""
@@ -215,14 +246,23 @@ def test_tracking_namespace_is_picklable_for_spawned_api_workers() -> None:
    assert restored._omni_stage_client_configs == args._omni_stage_client_configs


-def test_serve_validate_rejects_multiple_api_servers_for_diffusion(mocker: MockerFixture) -> None:
+@pytest.mark.parametrize("native_checkpoint", [False, True])
+def test_serve_validate_rejects_multiple_api_servers_for_diffusion(
+    mocker: MockerFixture, tmp_path, native_checkpoint: bool
+) -> None:
    parser = TrackingArgumentParser()
    subparsers = parser.add_subparsers(dest="subcommand")
    cmd = OmniServeCommand()
    cmd.subparser_init(subparsers)
-    args = parser.parse_args(["serve", "fake-diffusion-model", "--omni", "--api-server-count", "2"])
+    argv = ["serve", "fake-diffusion-model", "--omni", "--api-server-count", "2"]
+    if native_checkpoint:
+        checkpoint = tmp_path / "anima.safetensors"
+        checkpoint.touch()
+        argv[1] = str(checkpoint)
+        argv.extend(["--model-class-name", "AnimaPipeline"])
+    args = parser.parse_args(argv)

-    mocker.patch("vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", return_value=True)
+    mocker.patch("vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", return_value=not native_checkpoint)
    validate = mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args")

    with pytest.raises(ValueError, match="not supported for diffusion"):
@@ -303,6 +303,70 @@ class TestFilterDataclassKwargs:


 class TestResolveOmniConfig:
+    @pytest.mark.parametrize("model_class_name", ["AnimaPipeline", "AnimaModularPipeline"])
+    def test_native_anima_checkpoint_uses_default_diffusion_stage_without_model_config(
+        self, tmp_path, mocker: MockerFixture, model_class_name: str
+    ):
+        checkpoint = tmp_path / "anima.safetensors"
+        checkpoint.write_text("dummy")
+        custom_args = {"components_path": "/tmp/anima-components"}
+        mocker.patch(
+            "vllm_omni.config.resolver.StageConfigFactory.create_from_model",
+            side_effect=AssertionError("native Anima checkpoints should not require model config discovery"),
+        )
+
+        resolved = resolve_omni_config(
+            str(checkpoint),
+            trust_remote_code=False,
+            deploy_config_path=None,
+            cli_overrides={"model_class_name": model_class_name, "custom_pipeline_args": custom_args},
+            stage_overrides=None,
+            strategy_config_path=None,
+        )
+
+        assert resolved.config_path is None
+        assert len(resolved.stage_configs) == 1
+        stage_config = resolved.stage_configs[0]
+        assert stage_config.model_stage == "diffusion"
+        assert stage_config.diffusion_config.model_class_name == "AnimaPipeline"
+        assert stage_config.diffusion_config.custom_pipeline_args == custom_args
+
+    @pytest.mark.parametrize(
+        ("model_class_name", "checkpoint_exists", "deploy_config_path"),
+        [
+            ("AnimaPipeline", True, "deploy.yaml"),
+            ("AnimaPipeline", False, None),
+            ("FluxPipeline", True, None),
+        ],
+    )
+    def test_native_single_file_fallback_preserves_config_discovery(
+        self, tmp_path, mocker: MockerFixture, model_class_name, checkpoint_exists, deploy_config_path
+    ):
+        checkpoint = tmp_path / "model.safetensors"
+        if checkpoint_exists:
+            checkpoint.write_text("dummy")
+        create_from_model = mocker.patch(
+            "vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None
+        )
+
+        resolved = resolve_omni_config(
+            str(checkpoint),
+            trust_remote_code=False,
+            deploy_config_path=deploy_config_path,
+            cli_overrides={"model_class_name": model_class_name},
+            stage_overrides=None,
+            strategy_config_path=None,
+        )
+
+        create_from_model.assert_called_once_with(
+            str(checkpoint),
+            trust_remote_code=False,
+            cli_overrides={"model_class_name": model_class_name, "trust_remote_code": False},
+            deploy_config_path=deploy_config_path,
+            strategy_specs=None,
+        )
+        assert resolved.config_path == deploy_config_path
+
    def test_bare_deploy_name_returns_packaged_resolved_path(self, tmp_path, mocker: MockerFixture):
        deploy_name = "qwen3_omni_moe.yaml"
        deploy_path = tmp_path / deploy_name
@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+
+from tools.pre_commit import check_test_marks
+
+pytestmark = [pytest.mark.cpu, pytest.mark.core_model]
+
+
+@pytest.mark.parametrize("kind", ["offline", "online"])
+def test_common_diffusion_marks_are_checked_in_helper(kind, monkeypatch):
+    path = f"tests/model_tests/diffusion/test_common_{kind}.py"
+    assert check_test_marks.get_files_missing_markers([path]) == {}
+
+    read_file = check_test_marks.read_test_file
+    monkeypatch.setattr(
+        check_test_marks,
+        "read_test_file",
+        lambda filename: "" if filename.endswith("case_filtering.py") else read_file(filename),
+    )
+    assert check_test_marks.get_files_missing_markers([path]) == {path: ["Level", "Hardware"]}
+
+
+def test_unrelated_tests_cannot_inherit_diffusion_marks(monkeypatch):
+    path = "tests/test_missing_marks.py"
+    monkeypatch.setattr(check_test_marks, "read_test_file", lambda filename: "get_parametrized_options(settings)")
+    assert check_test_marks.get_files_missing_markers([path]) == {path: ["Level", "Hardware"]}
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import json
 from shutil import rmtree

@@ -12,7 +15,7 @@ def tiny_model_paths(request, run_level, tmp_path_factory):
    """Build or download the tiny models for the selected tests.

    At core_model level, builds tiny models via the builder function.
-    At advanced_model / full_model level, uses the real HF model path.
+    At advanced_model / full_model level, uses the real HF model path or asset resolver.

    NOTE: this is session scoped to avoid churn in tiny model creation,
    but will ensure all the tiny models you need are created for the selected tests
@@ -45,7 +48,11 @@ def tiny_model_paths(request, run_level, tmp_path_factory):
                            built_paths.append(path)
                        else:
                            print(f"Run level is {run_level}; {model_name} will use the full model")
-                            path = settings.model
+                            if callable(settings.model):
+                                path = settings.model()
+                                built_paths.append(path)
+                            else:
+                                path = settings.model
                        model_paths[model_name] = path
                data_file.write_text(json.dumps(model_paths))

@@ -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)
@@ -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},
@@ -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,
@@ -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:
@@ -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,
@@ -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"
@@ -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
@@ -33,6 +33,12 @@ HARDWARE_HELPERS = ("hardware_test", "hardware_marks")
 # The helper implementation is the only file allowed to write pytest.mark.<SKU>.
 _ALLOWED_DIRECT_SKU_FILES = frozenset({"tests/helpers/mark.py"})

+# Common diffusion tests receive per-case marks from this parametrization helper.
+_DELEGATED_MARK_SOURCES = {
+    "tests/model_tests/diffusion/test_common_offline.py": "tests/model_tests/diffusion/case_filtering.py",
+    "tests/model_tests/diffusion/test_common_online.py": "tests/model_tests/diffusion/case_filtering.py",
+}
+
 # Match mark.X since we could also do `from pytest import mark`.
 # \b prevents matching prefixes (e.g., mark.slow vs mark.slow_test).
 HELPER_RE = re.compile(r"(?:" + "|".join(HARDWARE_HELPERS) + r")\s*\(")
@@ -146,6 +152,9 @@ def get_files_missing_markers(
            missing = []
            if has_direct_sku_marker(path, contents):
                missing.append(DIRECT_SKU_MARKER)
+            mark_source = _DELEGATED_MARK_SOURCES.get(_normalize_path(path))
+            if mark_source is not None and "get_parametrized_options(" in contents:
+                contents += "\n" + (read_test_file(str(_repo_root() / mark_source)) or "")
            if not has_level_marker(contents):
                missing.append(MISSING_LEVEL_MARKER)
            if not has_hardware_marker(contents):
@@ -5,6 +5,7 @@ from __future__ import annotations

 from collections.abc import Mapping
 from dataclasses import dataclass, fields, is_dataclass
+from pathlib import Path
 from typing import Any

 from vllm.logger import init_logger
@@ -15,7 +16,7 @@ from vllm_omni.config.endpoint_policy import EndpointRestriction
 from vllm_omni.config.omni_config import VllmOmniConfig
 from vllm_omni.config.stage_config import PipelineConfig
 from vllm_omni.diffusion.data import resolve_model_class_name
-from vllm_omni.diffusion.registry import DiffusionModelRegistry
+from vllm_omni.diffusion.registry import DiffusionModelRegistry, resolve_native_single_file
 from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model

 logger = init_logger(__name__)
@@ -204,17 +205,24 @@ def resolve_omni_config(
    """Resolve registry/deploy inputs through the single public entrypoint."""
    normalized_overrides = _convert_dataclasses_to_dict(dict(cli_overrides or {}))
    normalized_overrides = with_trust_remote_code_override(normalized_overrides, trust_remote_code)
+    native_single_file_model = resolve_native_single_file(normalized_overrides.get("model_class_name"))
+    is_native_single_file = native_single_file_model is not None and Path(model).is_file()
+    if is_native_single_file:
+        normalized_overrides["model_class_name"] = native_single_file_model
    registry_overrides = dict(normalized_overrides)
    _flatten_stage_overrides(registry_overrides, stage_overrides)

    strategy_specs = _load_strategy_specs(strategy_config_path)
-    structured_config = StageConfigFactory.create_from_model(
-        model,
-        trust_remote_code=trust_remote_code,
-        cli_overrides=registry_overrides,
-        deploy_config_path=deploy_config_path,
-        strategy_specs=strategy_specs,
-    )
+    structured_config = None
+    # Native checkpoint files have no HF config or pipeline index to discover.
+    if deploy_config_path is not None or not is_native_single_file:
+        structured_config = StageConfigFactory.create_from_model(
+            model,
+            trust_remote_code=trust_remote_code,
+            cli_overrides=registry_overrides,
+            deploy_config_path=deploy_config_path,
+            strategy_specs=strategy_specs,
+        )
    if structured_config is not None:
        return _build_registered_resolution(structured_config)

@@ -1379,6 +1379,10 @@ class OmniDiffusionConfig:
                    self.model,
                )

+    @property
+    def is_single_file(self) -> bool:
+        return isinstance(self.model, str) and os.path.isfile(self.model)
+
    def _propagate_quantization_from_tf_config(self, tf_config: "TransformerConfig") -> None:
        if tf_config.quant_config is None:
            return
@@ -1483,12 +1487,22 @@ class OmniDiffusionConfig:
        """
        from vllm.transformers_utils.config import get_hf_file_to_dict

+        from vllm_omni.diffusion.registry import resolve_native_single_file
        from vllm_omni.diffusion.utils.hf_utils import (
            get_diffusion_model_index,
            resolve_native_diffusion_model_class,
        )

        assert self.model is not None
+
+        native_single_file_model = resolve_native_single_file(self.model_class_name)
+        if self.is_single_file and native_single_file_model is not None:
+            self.diffusion_load_format = "default"
+            self.model_class_name = native_single_file_model
+            self.diffusers_pipeline_cls = None
+            self.set_tf_model_config(TransformerConfig())
+            return
+
        try:
            config_dict = get_diffusion_model_index(
                self.model,
@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from vllm_omni.diffusion.models.anima.pipeline_anima import AnimaPipeline
+
+__all__ = [
+    "AnimaPipeline",
+]
@@ -0,0 +1,321 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from collections.abc import Iterable
+from types import SimpleNamespace
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention
+
+ANIMA_TEXT_CONDITIONER_CONFIG = {
+    "source_dim": 1024,
+    "target_dim": 1024,
+    "model_dim": 1024,
+    "num_layers": 6,
+    "num_attention_heads": 16,
+    "mlp_ratio": 4.0,
+    "target_vocab_size": 32128,
+    "use_self_attention": True,
+    "use_layer_norm": False,
+    "min_sequence_length": 512,
+}
+
+
+def _rotate_half(hidden_states: torch.Tensor) -> torch.Tensor:
+    hidden_states_1 = hidden_states[..., : hidden_states.shape[-1] // 2]
+    hidden_states_2 = hidden_states[..., hidden_states.shape[-1] // 2 :]
+    return torch.cat((-hidden_states_2, hidden_states_1), dim=-1)
+
+
+def _apply_rotary_pos_emb(
+    hidden_states: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    unsqueeze_dim: int = 1,
+) -> torch.Tensor:
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    return (hidden_states * cos) + (_rotate_half(hidden_states) * sin)
+
+
+class AnimaRotaryEmbedding(nn.Module):
+    def __init__(self, head_dim: int, rope_theta: float = 10000.0) -> None:
+        super().__init__()
+        inv_freq = 1.0 / (
+            rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.int64).to(dtype=torch.float32) / head_dim)
+        )
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        inv_freq_expanded = inv_freq_expanded.to(hidden_states.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos = emb.cos()
+        sin = emb.sin()
+        return cos.to(dtype=hidden_states.dtype), sin.to(dtype=hidden_states.dtype)
+
+
+class AnimaTextConditionerAttention(nn.Module):
+    def __init__(
+        self,
+        query_dim: int,
+        context_dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        inner_dim = num_attention_heads * attention_head_dim
+
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_dim = attention_head_dim
+        self.q_proj = nn.Linear(query_dim, inner_dim, bias=False)
+        self.q_norm = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.k_proj = nn.Linear(context_dim, inner_dim, bias=False)
+        self.k_norm = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.v_proj = nn.Linear(context_dim, inner_dim, bias=False)
+        self.o_proj = nn.Linear(inner_dim, query_dim, bias=False)
+        self.attn = Attention(
+            num_heads=num_attention_heads,
+            head_size=attention_head_dim,
+            softmax_scale=1.0 / (attention_head_dim**0.5),
+            causal=False,
+            num_kv_heads=num_attention_heads,
+            prefix=prefix,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        encoder_position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        encoder_hidden_states = hidden_states if encoder_hidden_states is None else encoder_hidden_states
+        input_shape = hidden_states.shape[:-1]
+        encoder_input_shape = encoder_hidden_states.shape[:-1]
+
+        query = self.q_proj(hidden_states)
+        key = self.k_proj(encoder_hidden_states)
+        value = self.v_proj(encoder_hidden_states)
+
+        query = query.view(*input_shape, self.num_attention_heads, self.attention_head_dim)
+        key = key.view(*encoder_input_shape, self.num_attention_heads, self.attention_head_dim)
+        value = value.view(*encoder_input_shape, self.num_attention_heads, self.attention_head_dim)
+
+        query = self.q_norm(query)
+        key = self.k_norm(key)
+
+        if position_embeddings is not None:
+            if encoder_position_embeddings is None:
+                raise ValueError("encoder_position_embeddings is required with rotary embeddings.")
+            cos, sin = position_embeddings
+            query = _apply_rotary_pos_emb(query, cos, sin, unsqueeze_dim=2)
+            cos, sin = encoder_position_embeddings
+            key = _apply_rotary_pos_emb(key, cos, sin, unsqueeze_dim=2)
+
+        attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
+        hidden_states = self.attn(query, key, value, attn_metadata=attn_metadata)
+        hidden_states = hidden_states.flatten(2, 3).contiguous()
+        return self.o_proj(hidden_states)
+
+
+class AnimaTextConditionerBlock(nn.Module):
+    def __init__(
+        self,
+        source_dim: int,
+        model_dim: int,
+        num_attention_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        use_self_attention: bool = True,
+        use_layer_norm: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.use_self_attention = use_self_attention
+        norm_cls = nn.LayerNorm if use_layer_norm else nn.RMSNorm
+        norm_kwargs = {} if use_layer_norm else {"eps": 1e-6}
+
+        if use_self_attention:
+            self.norm_self_attn = norm_cls(model_dim, **norm_kwargs)
+            self.self_attn = AnimaTextConditionerAttention(
+                query_dim=model_dim,
+                context_dim=model_dim,
+                num_attention_heads=num_attention_heads,
+                attention_head_dim=model_dim // num_attention_heads,
+                prefix=f"{prefix}.self_attn" if prefix else "self_attn",
+            )
+
+        self.norm_cross_attn = norm_cls(model_dim, **norm_kwargs)
+        self.cross_attn = AnimaTextConditionerAttention(
+            query_dim=model_dim,
+            context_dim=source_dim,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=model_dim // num_attention_heads,
+            prefix=f"{prefix}.cross_attn" if prefix else "cross_attn",
+        )
+        self.norm_mlp = norm_cls(model_dim, **norm_kwargs)
+        self.mlp = nn.Sequential(
+            nn.Linear(model_dim, int(model_dim * mlp_ratio)),
+            nn.GELU(),
+            nn.Linear(int(model_dim * mlp_ratio), model_dim),
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        target_attention_mask: torch.Tensor | None = None,
+        source_attention_mask: torch.Tensor | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        source_position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        if self.use_self_attention:
+            norm_hidden_states = self.norm_self_attn(hidden_states)
+            attn_hidden_states = self.self_attn(
+                norm_hidden_states,
+                attention_mask=target_attention_mask,
+                position_embeddings=position_embeddings,
+                encoder_position_embeddings=position_embeddings,
+            )
+            hidden_states = hidden_states + attn_hidden_states
+
+        norm_hidden_states = self.norm_cross_attn(hidden_states)
+        attn_hidden_states = self.cross_attn(
+            norm_hidden_states,
+            attention_mask=source_attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            position_embeddings=position_embeddings,
+            encoder_position_embeddings=source_position_embeddings,
+        )
+        hidden_states = hidden_states + attn_hidden_states
+        hidden_states = hidden_states + self.mlp(self.norm_mlp(hidden_states))
+        return hidden_states
+
+
+def is_conditioner_block_module(name: str, module: nn.Module) -> bool:
+    return "blocks" in name and name.split(".")[-1].isdigit()
+
+
+class AnimaTextConditioner(nn.Module):
+    _repeated_blocks = ["AnimaTextConditionerBlock"]
+    _hsdp_shard_conditions = [is_conditioner_block_module]
+
+    def __init__(
+        self,
+        source_dim: int = 1024,
+        target_dim: int = 1024,
+        model_dim: int = 1024,
+        num_layers: int = 6,
+        num_attention_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        target_vocab_size: int = 32128,
+        use_self_attention: bool = True,
+        use_layer_norm: bool = False,
+        min_sequence_length: int = 512,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        self.config = SimpleNamespace(
+            source_dim=source_dim,
+            target_dim=target_dim,
+            model_dim=model_dim,
+            num_layers=num_layers,
+            num_attention_heads=num_attention_heads,
+            mlp_ratio=mlp_ratio,
+            target_vocab_size=target_vocab_size,
+            use_self_attention=use_self_attention,
+            use_layer_norm=use_layer_norm,
+            min_sequence_length=min_sequence_length,
+            extra_config=kwargs,
+        )
+        self.embed = nn.Embedding(target_vocab_size, target_dim)
+        self.in_proj = nn.Linear(target_dim, model_dim) if model_dim != target_dim else nn.Identity()
+        self.rotary_emb = AnimaRotaryEmbedding(model_dim // num_attention_heads)
+        self.blocks = nn.ModuleList(
+            [
+                AnimaTextConditionerBlock(
+                    source_dim=source_dim,
+                    model_dim=model_dim,
+                    num_attention_heads=num_attention_heads,
+                    mlp_ratio=mlp_ratio,
+                    use_self_attention=use_self_attention,
+                    use_layer_norm=use_layer_norm,
+                    prefix=f"blocks.{i}",
+                )
+                for i in range(num_layers)
+            ]
+        )
+        self.out_proj = nn.Linear(model_dim, target_dim)
+        self.norm = nn.RMSNorm(target_dim, eps=1e-6)
+        self.gradient_checkpointing = False
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    @staticmethod
+    def _prepare_attention_mask(attention_mask: torch.Tensor | None) -> torch.Tensor | None:
+        if attention_mask is None:
+            return None
+        return attention_mask.to(torch.bool)
+
+    def forward(
+        self,
+        source_hidden_states: torch.Tensor,
+        target_input_ids: torch.Tensor,
+        target_attention_mask: torch.Tensor | None = None,
+        source_attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        target_attention_mask = self._prepare_attention_mask(target_attention_mask)
+        source_attention_mask = self._prepare_attention_mask(source_attention_mask)
+
+        hidden_states = self.embed(target_input_ids).to(dtype=source_hidden_states.dtype)
+        hidden_states = self.in_proj(hidden_states)
+
+        position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
+        source_position_ids = torch.arange(source_hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        source_position_embeddings = self.rotary_emb(hidden_states, source_position_ids)
+
+        for block in self.blocks:
+            hidden_states = block(
+                hidden_states,
+                source_hidden_states,
+                target_attention_mask=target_attention_mask,
+                source_attention_mask=source_attention_mask,
+                position_embeddings=position_embeddings,
+                source_position_embeddings=source_position_embeddings,
+            )
+
+        hidden_states = self.norm(self.out_proj(hidden_states))
+
+        if target_attention_mask is not None:
+            hidden_states = hidden_states * target_attention_mask.to(hidden_states).unsqueeze(-1)
+
+        if hidden_states.shape[1] < self.config.min_sequence_length:
+            hidden_states = F.pad(hidden_states, (0, 0, 0, self.config.min_sequence_length - hidden_states.shape[1]))
+
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            if name not in params:
+                continue
+            param = params[name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, tensor)
+            loaded.add(name)
+        return loaded
@@ -0,0 +1,705 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from collections.abc import Iterable
+from types import SimpleNamespace
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.embeddings import Timesteps, apply_rotary_emb
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
+
+ANIMA_TRANSFORMER_CONFIG = {
+    "in_channels": 16,
+    "out_channels": 16,
+    "num_attention_heads": 16,
+    "attention_head_dim": 128,
+    "num_layers": 28,
+    "mlp_ratio": 4.0,
+    "text_embed_dim": 1024,
+    "adaln_lora_dim": 256,
+    "max_size": (128, 240, 240),
+    "patch_size": (1, 2, 2),
+    "rope_scale": (1.0, 4.0, 4.0),
+    "concat_padding_mask": True,
+    "extra_pos_embed_type": None,
+}
+
+
+# NOTE: We import and use diffusers' `apply_rotary_emb` instead of a custom native implementation
+# to prevent numerical drift in bfloat16. Diffusers upcasts queries, keys, and rotary frequency
+# tensors to float32 before computing the rotation, and casts back to bfloat16 at the end.
+# Performing the entire computation in bfloat16 accumulates precision errors across the 28
+# transformer blocks, which is heavily amplified by Classifier-Free Guidance (CFG).
+
+
+class CosmosPatchEmbed(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        patch_size: tuple[int, int, int],
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        self.patch_size = patch_size
+        self.proj = nn.Linear(in_channels * patch_size[0] * patch_size[1] * patch_size[2], out_channels, bias=bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        hidden_states = hidden_states.reshape(
+            batch_size,
+            num_channels,
+            num_frames // p_t,
+            p_t,
+            height // p_h,
+            p_h,
+            width // p_w,
+            p_w,
+        )
+        hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7)
+        return self.proj(hidden_states)
+
+
+class CosmosTimestepEmbedding(nn.Module):
+    def __init__(self, in_features: int, out_features: int) -> None:
+        super().__init__()
+        self.linear_1 = nn.Linear(in_features, out_features, bias=False)
+        self.activation = nn.SiLU()
+        self.linear_2 = nn.Linear(out_features, 3 * out_features, bias=False)
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        emb = self.linear_1(timesteps)
+        emb = self.activation(emb)
+        return self.linear_2(emb)
+
+
+class CosmosEmbedding(nn.Module):
+    def __init__(self, embedding_dim: int, condition_dim: int) -> None:
+        super().__init__()
+        self.time_proj = Timesteps(embedding_dim, flip_sin_to_cos=True, downscale_freq_shift=0.0)
+        self.t_embedder = CosmosTimestepEmbedding(embedding_dim, condition_dim)
+        self.norm = nn.RMSNorm(embedding_dim, eps=1e-6)
+
+    def forward(self, hidden_states: torch.Tensor, timestep: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        timesteps_proj = self.time_proj(timestep).type_as(hidden_states)
+        temb = self.t_embedder(timesteps_proj)
+        embedded_timestep = self.norm(timesteps_proj)
+        return temb, embedded_timestep
+
+
+class CosmosAdaLayerNorm(nn.Module):
+    def __init__(self, in_features: int, hidden_features: int) -> None:
+        super().__init__()
+        self.embedding_dim = in_features
+        self.activation = nn.SiLU()
+        self.norm = nn.LayerNorm(in_features, elementwise_affine=False, eps=1e-6)
+        self.linear_1 = nn.Linear(in_features, hidden_features, bias=False)
+        self.linear_2 = nn.Linear(hidden_features, 2 * in_features, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        embedded_timestep: torch.Tensor,
+        temb: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        embedded_timestep = self.activation(embedded_timestep)
+        embedded_timestep = self.linear_1(embedded_timestep)
+        embedded_timestep = self.linear_2(embedded_timestep)
+
+        if temb is not None:
+            embedded_timestep = embedded_timestep + temb[..., : 2 * self.embedding_dim]
+
+        shift, scale = embedded_timestep.chunk(2, dim=-1)
+        hidden_states = self.norm(hidden_states)
+
+        if embedded_timestep.ndim == 2:
+            shift, scale = (x.unsqueeze(1) for x in (shift, scale))
+
+        return hidden_states * (1 + scale) + shift
+
+
+class CosmosAdaLayerNormZero(nn.Module):
+    def __init__(self, in_features: int, hidden_features: int | None = None) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(in_features, elementwise_affine=False, eps=1e-6)
+        self.activation = nn.SiLU()
+        self.linear_1 = (
+            nn.Identity() if hidden_features is None else nn.Linear(in_features, hidden_features, bias=False)
+        )
+        self.linear_2 = nn.Linear(hidden_features or in_features, 3 * in_features, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        embedded_timestep: torch.Tensor,
+        temb: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        embedded_timestep = self.activation(embedded_timestep)
+        embedded_timestep = self.linear_1(embedded_timestep)
+        embedded_timestep = self.linear_2(embedded_timestep)
+
+        if temb is not None:
+            embedded_timestep = embedded_timestep + temb
+
+        shift, scale, gate = embedded_timestep.chunk(3, dim=-1)
+        hidden_states = self.norm(hidden_states)
+
+        if embedded_timestep.ndim == 2:
+            shift, scale, gate = (x.unsqueeze(1) for x in (shift, scale, gate))
+
+        hidden_states = hidden_states * (1 + scale) + shift
+        return hidden_states, gate
+
+
+class CosmosGELU(nn.Module):
+    def __init__(self, dim: int, inner_dim: int, bias: bool = False) -> None:
+        super().__init__()
+        self.proj = nn.Linear(dim, inner_dim, bias=bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return F.gelu(self.proj(hidden_states))
+
+
+class CosmosFeedForward(nn.Module):
+    def __init__(self, dim: int, mult: float = 4.0, bias: bool = False) -> None:
+        super().__init__()
+        inner_dim = int(dim * mult)
+        self.net = nn.ModuleList(
+            [
+                CosmosGELU(dim, inner_dim, bias=bias),
+                nn.Identity(),
+                nn.Linear(inner_dim, dim, bias=bias),
+            ]
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for module in self.net:
+            hidden_states = module(hidden_states)
+        return hidden_states
+
+
+class CosmosAttention(nn.Module):
+    def __init__(
+        self,
+        query_dim: int,
+        cross_attention_dim: int | None = None,
+        heads: int = 8,
+        dim_head: int = 64,
+        out_bias: bool = False,
+        img_context: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        inner_dim = heads * dim_head
+        cross_attention_dim = query_dim if cross_attention_dim is None else cross_attention_dim
+
+        self.heads = heads
+        self.dim_head = dim_head
+        self.out_dim = inner_dim
+        self.query_dim = query_dim
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=False)
+        self.norm_q = nn.RMSNorm(dim_head, eps=1e-6)
+        self.norm_k = nn.RMSNorm(dim_head, eps=1e-6)
+        self.to_out = nn.ModuleList([nn.Linear(inner_dim, query_dim, bias=out_bias), nn.Identity()])
+        self.attn = Attention(
+            num_heads=heads,
+            head_size=dim_head,
+            softmax_scale=1.0 / (dim_head**0.5),
+            causal=False,
+            num_kv_heads=heads,
+            prefix=prefix,
+        )
+
+        self.img_context = img_context
+        if img_context:
+            self.q_img = nn.Linear(query_dim, inner_dim, bias=False)
+            self.k_img = nn.Linear(query_dim, inner_dim, bias=False)
+            self.v_img = nn.Linear(query_dim, inner_dim, bias=False)
+            self.q_img_norm = nn.RMSNorm(dim_head, eps=1e-6)
+            self.k_img_norm = nn.RMSNorm(dim_head, eps=1e-6)
+            self.img_attn = Attention(
+                num_heads=heads,
+                head_size=dim_head,
+                softmax_scale=1.0 / (dim_head**0.5),
+                causal=False,
+                num_kv_heads=heads,
+                prefix=f"{prefix}.img_attn" if prefix else "img_attn",
+            )
+
+    def _attention(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        encoder_hidden_states = hidden_states if encoder_hidden_states is None else encoder_hidden_states
+        batch_size, seq_len, _ = hidden_states.shape
+        encoder_seq_len = encoder_hidden_states.shape[1]
+
+        query = self.to_q(hidden_states).view(batch_size, seq_len, self.heads, self.dim_head)
+        key = self.to_k(encoder_hidden_states).view(batch_size, encoder_seq_len, self.heads, self.dim_head)
+        value = self.to_v(encoder_hidden_states).view(batch_size, encoder_seq_len, self.heads, self.dim_head)
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        if image_rotary_emb is not None:
+            # We use diffusers' apply_rotary_emb to leverage its internal float32 rotation upcasting
+            # logic, resolving the bfloat16 cumulative precision drift vs. the reference pipeline.
+            query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2)
+            key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2)
+
+        attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
+        hidden_states = self.attn(query, key, value, attn_metadata=attn_metadata)
+        return hidden_states.flatten(2, 3).type_as(query)
+
+    def _image_attention(
+        self,
+        hidden_states: torch.Tensor,
+        img_context: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        img_seq_len = img_context.shape[1]
+
+        query = self.q_img(hidden_states).view(batch_size, seq_len, self.heads, self.dim_head)
+        key = self.k_img(img_context).view(batch_size, img_seq_len, self.heads, self.dim_head)
+        value = self.v_img(img_context).view(batch_size, img_seq_len, self.heads, self.dim_head)
+
+        query = self.q_img_norm(query)
+        key = self.k_img_norm(key)
+
+        attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
+        hidden_states = self.img_attn(query, key, value, attn_metadata=attn_metadata)
+        return hidden_states.flatten(2, 3).type_as(query)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor | None] | None = None,
+        attention_mask: torch.Tensor | tuple[torch.Tensor | None, torch.Tensor | None] | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        if self.img_context:
+            if not isinstance(encoder_hidden_states, tuple):
+                raise ValueError("Expected encoder_hidden_states as (text_context, img_context).")
+            text_context, img_context = encoder_hidden_states
+            text_mask, img_mask = attention_mask if attention_mask is not None else (None, None)
+            query_states = hidden_states
+            hidden_states = self._attention(query_states, text_context, text_mask, image_rotary_emb)
+            if img_context is not None:
+                hidden_states = hidden_states + self._image_attention(query_states, img_context, img_mask)
+        else:
+            hidden_states = self._attention(hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb)
+
+        hidden_states = self.to_out[0](hidden_states)
+        return self.to_out[1](hidden_states)
+
+
+class CosmosTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        cross_attention_dim: int,
+        mlp_ratio: float = 4.0,
+        adaln_lora_dim: int = 256,
+        out_bias: bool = False,
+        img_context: bool = False,
+        before_proj: bool = False,
+        after_proj: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        hidden_size = num_attention_heads * attention_head_dim
+
+        self.norm1 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim)
+        self.img_context = img_context
+        self.attn1 = CosmosAttention(
+            query_dim=hidden_size,
+            cross_attention_dim=None,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            out_bias=out_bias,
+            prefix=f"{prefix}.attn1" if prefix else "attn1",
+        )
+        self.norm2 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim)
+        self.attn2 = CosmosAttention(
+            query_dim=hidden_size,
+            cross_attention_dim=cross_attention_dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            out_bias=out_bias,
+            img_context=img_context,
+            prefix=f"{prefix}.attn2" if prefix else "attn2",
+        )
+        self.norm3 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim)
+        self.ff = CosmosFeedForward(hidden_size, mult=mlp_ratio, bias=out_bias)
+
+        self.before_proj = nn.Linear(hidden_size, hidden_size) if before_proj else None
+        self.after_proj = nn.Linear(hidden_size, hidden_size) if after_proj else None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor | None],
+        embedded_timestep: torch.Tensor,
+        temb: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        extra_pos_emb: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | tuple[torch.Tensor | None, torch.Tensor | None] | None = None,
+        controlnet_residual: torch.Tensor | None = None,
+        latents: torch.Tensor | None = None,
+        block_idx: int | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if self.before_proj is not None:
+            hidden_states = self.before_proj(hidden_states) + latents
+
+        if extra_pos_emb is not None:
+            hidden_states = hidden_states + extra_pos_emb
+
+        norm_hidden_states, gate = self.norm1(hidden_states, embedded_timestep, temb)
+        attn_output = self.attn1(norm_hidden_states, image_rotary_emb=image_rotary_emb)
+        hidden_states = hidden_states + gate * attn_output
+
+        norm_hidden_states, gate = self.norm2(hidden_states, embedded_timestep, temb)
+        attn_output = self.attn2(
+            norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask
+        )
+        hidden_states = hidden_states + gate * attn_output
+
+        norm_hidden_states, gate = self.norm3(hidden_states, embedded_timestep, temb)
+        ff_output = self.ff(norm_hidden_states)
+        hidden_states = hidden_states + gate * ff_output
+
+        if controlnet_residual is not None:
+            if self.after_proj is not None:
+                raise ValueError("controlnet_residual is incompatible with after_proj.")
+            hidden_states = hidden_states + controlnet_residual
+
+        if self.after_proj is not None:
+            return hidden_states, self.after_proj(hidden_states)
+
+        return hidden_states
+
+
+class CosmosRotaryPosEmbed(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        max_size: tuple[int, int, int] = (128, 240, 240),
+        patch_size: tuple[int, int, int] = (1, 2, 2),
+        base_fps: int = 24,
+        rope_scale: tuple[float, float, float] = (2.0, 1.0, 1.0),
+    ) -> None:
+        super().__init__()
+        self.max_size = [size // patch for size, patch in zip(max_size, patch_size)]
+        self.patch_size = patch_size
+        self.base_fps = base_fps
+
+        self.dim_h = hidden_size // 6 * 2
+        self.dim_w = hidden_size // 6 * 2
+        self.dim_t = hidden_size - self.dim_h - self.dim_w
+
+        self.h_ntk_factor = rope_scale[1] ** (self.dim_h / (self.dim_h - 2))
+        self.w_ntk_factor = rope_scale[2] ** (self.dim_w / (self.dim_w - 2))
+        self.t_ntk_factor = rope_scale[0] ** (self.dim_t / (self.dim_t - 2))
+
+    def forward(self, hidden_states: torch.Tensor, fps: float | None = None) -> tuple[torch.Tensor, torch.Tensor]:
+        _, _, num_frames, height, width = hidden_states.shape
+        pe_size = [num_frames // self.patch_size[0], height // self.patch_size[1], width // self.patch_size[2]]
+        device = hidden_states.device
+
+        h_theta = 10000.0 * self.h_ntk_factor
+        w_theta = 10000.0 * self.w_ntk_factor
+        t_theta = 10000.0 * self.t_ntk_factor
+
+        seq = torch.arange(max(self.max_size), device=device, dtype=torch.float32)
+        dim_h_range = torch.arange(0, self.dim_h, 2, device=device, dtype=torch.float32)[: self.dim_h // 2] / self.dim_h
+        dim_w_range = torch.arange(0, self.dim_w, 2, device=device, dtype=torch.float32)[: self.dim_w // 2] / self.dim_w
+        dim_t_range = torch.arange(0, self.dim_t, 2, device=device, dtype=torch.float32)[: self.dim_t // 2] / self.dim_t
+        h_spatial_freqs = 1.0 / (h_theta**dim_h_range)
+        w_spatial_freqs = 1.0 / (w_theta**dim_w_range)
+        temporal_freqs = 1.0 / (t_theta**dim_t_range)
+
+        emb_h = torch.outer(seq[: pe_size[1]], h_spatial_freqs)[None, :, None, :].repeat(pe_size[0], 1, pe_size[2], 1)
+        emb_w = torch.outer(seq[: pe_size[2]], w_spatial_freqs)[None, None, :, :].repeat(pe_size[0], pe_size[1], 1, 1)
+        if fps is None:
+            emb_t = torch.outer(seq[: pe_size[0]], temporal_freqs)
+        else:
+            emb_t = torch.outer(seq[: pe_size[0]] / fps * self.base_fps, temporal_freqs)
+        emb_t = emb_t[:, None, None, :].repeat(1, pe_size[1], pe_size[2], 1)
+
+        freqs = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1).flatten(0, 2).float()
+        return torch.cos(freqs), torch.sin(freqs)
+
+
+class CosmosLearnablePositionalEmbed(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        max_size: tuple[int, int, int],
+        patch_size: tuple[int, int, int],
+        eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        self.max_size = [size // patch for size, patch in zip(max_size, patch_size)]
+        self.patch_size = patch_size
+        self.eps = eps
+        self.pos_emb_t = nn.Parameter(torch.zeros(self.max_size[0], hidden_size))
+        self.pos_emb_h = nn.Parameter(torch.zeros(self.max_size[1], hidden_size))
+        self.pos_emb_w = nn.Parameter(torch.zeros(self.max_size[2], hidden_size))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, _, num_frames, height, width = hidden_states.shape
+        pe_size = [num_frames // self.patch_size[0], height // self.patch_size[1], width // self.patch_size[2]]
+
+        emb_t = self.pos_emb_t[: pe_size[0]][None, :, None, None, :].repeat(batch_size, 1, pe_size[1], pe_size[2], 1)
+        emb_h = self.pos_emb_h[: pe_size[1]][None, None, :, None, :].repeat(batch_size, pe_size[0], 1, pe_size[2], 1)
+        emb_w = self.pos_emb_w[: pe_size[2]][None, None, None, :, :].repeat(batch_size, pe_size[0], pe_size[1], 1, 1)
+        emb = (emb_t + emb_h + emb_w).flatten(1, 3)
+
+        norm = torch.linalg.vector_norm(emb, dim=-1, keepdim=True, dtype=torch.float32)
+        norm = torch.add(self.eps, norm, alpha=np.sqrt(norm.numel() / emb.numel()))
+        return (emb / norm).type_as(hidden_states)
+
+
+class AnimaTransformer3DModel(nn.Module):
+    _repeated_blocks = ["CosmosTransformerBlock"]
+    _hsdp_shard_conditions = [is_transformer_block_module]
+
+    def __init__(
+        self,
+        in_channels: int = 16,
+        out_channels: int = 16,
+        num_attention_heads: int = 32,
+        attention_head_dim: int = 128,
+        num_layers: int = 28,
+        mlp_ratio: float = 4.0,
+        text_embed_dim: int = 1024,
+        adaln_lora_dim: int = 256,
+        max_size: tuple[int, int, int] = (128, 240, 240),
+        patch_size: tuple[int, int, int] = (1, 2, 2),
+        rope_scale: tuple[float, float, float] = (2.0, 1.0, 1.0),
+        concat_padding_mask: bool = True,
+        extra_pos_embed_type: str | None = "learnable",
+        use_crossattn_projection: bool = False,
+        crossattn_proj_in_channels: int = 1024,
+        encoder_hidden_states_channels: int = 1024,
+        controlnet_block_every_n: int | None = None,
+        img_context_dim_in: int | None = None,
+        img_context_num_tokens: int = 256,
+        img_context_dim_out: int = 2048,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        hidden_size = num_attention_heads * attention_head_dim
+        self.config = SimpleNamespace(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            num_layers=num_layers,
+            mlp_ratio=mlp_ratio,
+            text_embed_dim=text_embed_dim,
+            adaln_lora_dim=adaln_lora_dim,
+            max_size=max_size,
+            patch_size=patch_size,
+            rope_scale=rope_scale,
+            concat_padding_mask=concat_padding_mask,
+            extra_pos_embed_type=extra_pos_embed_type,
+            use_crossattn_projection=use_crossattn_projection,
+            crossattn_proj_in_channels=crossattn_proj_in_channels,
+            encoder_hidden_states_channels=encoder_hidden_states_channels,
+            controlnet_block_every_n=controlnet_block_every_n,
+            img_context_dim_in=img_context_dim_in,
+            img_context_num_tokens=img_context_num_tokens,
+            img_context_dim_out=img_context_dim_out,
+            extra_config=kwargs,
+        )
+
+        patch_embed_in_channels = in_channels + 1 if concat_padding_mask else in_channels
+        self.patch_embed = CosmosPatchEmbed(patch_embed_in_channels, hidden_size, patch_size, bias=False)
+        self.rope = CosmosRotaryPosEmbed(
+            hidden_size=attention_head_dim,
+            max_size=max_size,
+            patch_size=patch_size,
+            rope_scale=rope_scale,
+        )
+        self.learnable_pos_embed: CosmosLearnablePositionalEmbed | None = None
+        if extra_pos_embed_type == "learnable":
+            self.learnable_pos_embed = CosmosLearnablePositionalEmbed(
+                hidden_size=hidden_size,
+                max_size=max_size,
+                patch_size=patch_size,
+            )
+        self.time_embed = CosmosEmbedding(hidden_size, hidden_size)
+        self.transformer_blocks = nn.ModuleList(
+            [
+                CosmosTransformerBlock(
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    cross_attention_dim=text_embed_dim,
+                    mlp_ratio=mlp_ratio,
+                    adaln_lora_dim=adaln_lora_dim,
+                    out_bias=False,
+                    img_context=img_context_dim_in is not None and img_context_dim_in > 0,
+                    prefix=f"transformer_blocks.{i}",
+                )
+                for i in range(num_layers)
+            ]
+        )
+        self.norm_out = CosmosAdaLayerNorm(hidden_size, adaln_lora_dim)
+        self.proj_out = nn.Linear(hidden_size, patch_size[0] * patch_size[1] * patch_size[2] * out_channels, bias=False)
+
+        self.crossattn_proj: nn.Module | None = None
+        if use_crossattn_projection:
+            self.crossattn_proj = nn.Sequential(
+                nn.Linear(crossattn_proj_in_channels, encoder_hidden_states_channels, bias=True),
+                nn.GELU(),
+            )
+
+        self.img_context_proj: nn.Module | None = None
+        if img_context_dim_in:
+            self.img_context_proj = nn.Sequential(
+                nn.Linear(img_context_dim_in, img_context_dim_out, bias=True),
+                nn.GELU(),
+            )
+
+        self.gradient_checkpointing = False
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor | None],
+        block_controlnet_hidden_states: list[torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | tuple[torch.Tensor | None, torch.Tensor | None] | None = None,
+        fps: float | None = None,
+        condition_mask: torch.Tensor | None = None,
+        padding_mask: torch.Tensor | None = None,
+        return_dict: bool = True,
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
+        batch_size, _, num_frames, height, width = hidden_states.shape
+
+        if condition_mask is not None:
+            hidden_states = torch.cat([hidden_states, condition_mask], dim=1)
+
+        if self.config.concat_padding_mask:
+            if padding_mask is None:
+                raise ValueError("padding_mask is required when concat_padding_mask=True.")
+            padding_mask_resized = F.interpolate(
+                padding_mask,
+                size=list(hidden_states.shape[-2:]),
+                mode="nearest",
+            )
+            hidden_states = torch.cat(
+                [hidden_states, padding_mask_resized.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)],
+                dim=1,
+            )
+
+        if isinstance(attention_mask, torch.Tensor) and attention_mask.ndim != 2:
+            attention_mask = attention_mask.squeeze(1).squeeze(1)
+
+        image_rotary_emb = self.rope(hidden_states, fps=fps)
+        extra_pos_emb = self.learnable_pos_embed(hidden_states) if self.learnable_pos_embed is not None else None
+
+        p_t, p_h, p_w = self.config.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        hidden_states = self.patch_embed(hidden_states)
+        hidden_states = hidden_states.flatten(1, 3)
+
+        if timestep.ndim == 1:
+            temb, embedded_timestep = self.time_embed(hidden_states, timestep)
+        elif timestep.ndim == 5:
+            if timestep.shape != (batch_size, 1, num_frames, 1, 1):
+                raise ValueError(f"Expected timestep shape [B, 1, T, 1, 1], got {timestep.shape}.")
+            timestep = timestep.flatten()
+            temb, embedded_timestep = self.time_embed(hidden_states, timestep)
+            temb, embedded_timestep = (
+                x.view(batch_size, post_patch_num_frames, 1, 1, -1)
+                .expand(-1, -1, post_patch_height, post_patch_width, -1)
+                .flatten(1, 3)
+                for x in (temb, embedded_timestep)
+            )
+        else:
+            raise ValueError(f"Expected timestep shape [B, 1, T, 1, 1] or [T], got {timestep.shape}.")
+
+        text_context, img_context = (
+            encoder_hidden_states if isinstance(encoder_hidden_states, tuple) else (encoder_hidden_states, None)
+        )
+        if self.config.use_crossattn_projection:
+            if self.crossattn_proj is None:
+                raise ValueError("crossattn_proj is required when use_crossattn_projection=True.")
+            text_context = self.crossattn_proj(text_context)
+        if img_context is not None and self.config.img_context_dim_in:
+            if self.img_context_proj is None:
+                raise ValueError("img_context_proj is required when img_context_dim_in is set.")
+            img_context = self.img_context_proj(img_context)
+        processed_encoder_hidden_states = (
+            (text_context, img_context) if isinstance(encoder_hidden_states, tuple) else text_context
+        )
+
+        controlnet_block_index_map = {}
+        if block_controlnet_hidden_states is not None:
+            n_blocks = len(self.transformer_blocks)
+            controlnet_block_index_map = {
+                block_idx: block_controlnet_hidden_states[idx]
+                for idx, block_idx in list(enumerate(range(0, n_blocks, self.config.controlnet_block_every_n)))
+            }
+
+        for block_idx, block in enumerate(self.transformer_blocks):
+            hidden_states = block(
+                hidden_states,
+                processed_encoder_hidden_states,
+                embedded_timestep,
+                temb,
+                image_rotary_emb,
+                extra_pos_emb,
+                attention_mask,
+                controlnet_block_index_map.get(block_idx),
+            )
+
+        hidden_states = self.norm_out(hidden_states, embedded_timestep, temb)
+        hidden_states = self.proj_out(hidden_states)
+        hidden_states = hidden_states.unflatten(2, (p_h, p_w, p_t, -1))
+        hidden_states = hidden_states.unflatten(1, (post_patch_num_frames, post_patch_height, post_patch_width))
+        hidden_states = hidden_states.permute(0, 7, 1, 6, 2, 4, 3, 5)
+        hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        if not return_dict:
+            return (hidden_states,)
+
+        return Transformer2DModelOutput(sample=hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            if name not in params:
+                continue
+            param = params[name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, tensor)
+            loaded.add(name)
+        return loaded
@@ -0,0 +1,795 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import inspect
+import json
+import logging
+import os
+from collections.abc import Callable, Iterable
+from pathlib import Path
+from typing import Any, cast
+
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils.torch_utils import randn_tensor
+from safetensors.torch import load_file
+from torch import nn
+from transformers import AutoTokenizer, Qwen3Model, T5TokenizerFast
+
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
+from vllm_omni.diffusion.distributed.utils import get_local_device
+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
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload_strategy
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+from vllm_omni.diffusion.utils.size_utils import normalize_min_aligned_size
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+
+logger = logging.getLogger(__name__)
+
+ANIMA_VAE_SCALE_FACTOR = 8
+
+
+def retrieve_timesteps(
+    scheduler: Any,
+    num_inference_steps: int | None = None,
+    device: str | torch.device | None = None,
+    timesteps: list[int] | None = None,
+    sigmas: list[float] | None = None,
+    **kwargs: Any,
+) -> tuple[torch.Tensor, int]:
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        if "timesteps" not in inspect.signature(scheduler.set_timesteps).parameters:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        scheduler_timesteps = scheduler.timesteps
+        num_inference_steps = len(scheduler_timesteps)
+    elif sigmas is not None:
+        if "sigmas" not in inspect.signature(scheduler.set_timesteps).parameters:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        scheduler_timesteps = scheduler.timesteps
+        num_inference_steps = len(scheduler_timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        scheduler_timesteps = scheduler.timesteps
+    if num_inference_steps is None:
+        num_inference_steps = len(scheduler_timesteps)
+    return cast(torch.Tensor, scheduler_timesteps), num_inference_steps
+
+
+_ANIMA_ORIGINAL_TRANSFORMER_RENAMES = {
+    "t_embedder.1": "time_embed.t_embedder",
+    "t_embedding_norm": "time_embed.norm",
+    "blocks": "transformer_blocks",
+    "adaln_modulation_self_attn.1": "norm1.linear_1",
+    "adaln_modulation_self_attn.2": "norm1.linear_2",
+    "adaln_modulation_cross_attn.1": "norm2.linear_1",
+    "adaln_modulation_cross_attn.2": "norm2.linear_2",
+    "adaln_modulation_mlp.1": "norm3.linear_1",
+    "adaln_modulation_mlp.2": "norm3.linear_2",
+    "self_attn": "attn1",
+    "cross_attn": "attn2",
+    "q_proj": "to_q",
+    "k_proj": "to_k",
+    "v_proj": "to_v",
+    "output_proj": "to_out.0",
+    "q_norm": "norm_q",
+    "k_norm": "norm_k",
+    "mlp.layer1": "ff.net.0.proj",
+    "mlp.layer2": "ff.net.2",
+    "x_embedder.proj.1": "patch_embed.proj",
+    "final_layer.adaln_modulation.1": "norm_out.linear_1",
+    "final_layer.adaln_modulation.2": "norm_out.linear_2",
+    "final_layer.linear": "proj_out",
+}
+
+_ANIMA_ORIGINAL_TRANSFORMER_DROP_KEYS = (
+    "accum_video_sample_counter",
+    "accum_image_sample_counter",
+    "accum_iteration",
+    "accum_train_in_hours",
+    "pos_embedder.seq",
+    "pos_embedder.dim_spatial_range",
+    "pos_embedder.dim_temporal_range",
+    "_extra_state",
+)
+
+
+def _anima_vae_scale_factor_from_vae(vae: Any) -> int:
+    vae_config = getattr(vae, "config", None)
+    spatial_compression_ratio = getattr(vae_config, "spatial_compression_ratio", None)
+    if spatial_compression_ratio is not None:
+        return int(spatial_compression_ratio)
+    temporal_downsample = getattr(vae, "temperal_downsample", None)
+    if temporal_downsample is not None:
+        return int(2 ** len(temporal_downsample))
+    return ANIMA_VAE_SCALE_FACTOR
+
+
+def get_anima_post_process_func(od_config: OmniDiffusionConfig) -> Callable[[torch.Tensor], Any]:
+    image_processor = VaeImageProcessor(vae_scale_factor=ANIMA_VAE_SCALE_FACTOR)
+
+    def post_process_func(images: torch.Tensor) -> Any:
+        return image_processor.postprocess(images)
+
+    return post_process_func
+
+
+class AnimaPipeline(nn.Module, DiffusionPipelineProfilerMixin, ProgressBarMixin):
+    """Native Anima model path.
+
+    The Anima transformer checkpoint is distributed as a single safetensors
+    file. This loader builds the native pipeline components and runs the denoise
+    transformer/text-conditioner directly.
+    """
+
+    supports_step_execution = False
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, device: torch.device | str | None = None) -> None:
+        super().__init__()
+        self.od_config = od_config
+        self.parallel_config = od_config.parallel_config
+        self.device = device or get_local_device()
+        self.weights_sources: list[str] = []
+
+        self._raise_unsupported_features()
+
+        self._interrupt = False
+        self._guidance_scale = 1.0
+        self._num_timesteps = 0
+        self._current_timestep = None
+        self._profiler_initialized = False
+
+    def _raise_unsupported_features(self) -> None:
+        pc = self.od_config.parallel_config
+        if pc.cfg_parallel_size > 1:
+            raise NotImplementedError("AnimaPipeline does not yet support CFG parallel. Validate native CFG first.")
+        if pc.sequence_parallel_size is not None and pc.sequence_parallel_size > 1:
+            raise NotImplementedError("AnimaPipeline does not yet support sequence parallelism. Add an _sp_plan first.")
+        if pc.tensor_parallel_size > 1:
+            raise NotImplementedError("AnimaPipeline does not yet support tensor parallelism. Port TP linears first.")
+        if pc.use_hsdp:
+            raise NotImplementedError("AnimaPipeline does not yet support HSDP. Validate sharding after native load.")
+        if self.od_config.cache_backend not in ("none", None):
+            raise NotImplementedError("AnimaPipeline does not yet support TeaCache or Cache-DiT.")
+        if self.od_config.quantization_config is not None:
+            raise NotImplementedError("AnimaPipeline does not yet support quantized checkpoints.")
+        if resolve_offload_strategy(self.od_config) is not OffloadStrategy.NONE:
+            raise NotImplementedError("AnimaPipeline does not yet support CPU offload.")
+
+    @staticmethod
+    def _sub_state_dict(state_dict: dict[str, torch.Tensor], prefix: str) -> dict[str, torch.Tensor]:
+        return {key.removeprefix(prefix): value for key, value in state_dict.items() if key.startswith(prefix)}
+
+    @staticmethod
+    def _convert_original_transformer_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        if "patch_embed.proj.weight" in state_dict:
+            return state_dict
+
+        # The raw Anima checkpoint uses the original training names for a
+        # Cosmos-style denoiser. Convert those keys to the local
+        # AnimaTransformer3DModel module names before strict loading.
+        converted_state_dict = {}
+        for key, value in state_dict.items():
+            new_key = key.removeprefix("net.")
+            if any(drop_key in new_key for drop_key in _ANIMA_ORIGINAL_TRANSFORMER_DROP_KEYS):
+                continue
+            for old_key, new_name in _ANIMA_ORIGINAL_TRANSFORMER_RENAMES.items():
+                new_key = new_key.replace(old_key, new_name)
+            if new_key in converted_state_dict:
+                raise ValueError(f"Duplicate converted Anima transformer key: {new_key}")
+            converted_state_dict[new_key] = value
+        return converted_state_dict
+
+    @staticmethod
+    def _split_single_file_state_dict(
+        state_dict: dict[str, torch.Tensor],
+    ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
+        transformer_state_dict = AnimaPipeline._sub_state_dict(state_dict, "transformer.")
+        text_conditioner_state_dict = AnimaPipeline._sub_state_dict(state_dict, "text_conditioner.")
+        if transformer_state_dict and text_conditioner_state_dict:
+            return transformer_state_dict, text_conditioner_state_dict
+
+        text_conditioner_state_dict = AnimaPipeline._sub_state_dict(state_dict, "net.llm_adapter.")
+        if text_conditioner_state_dict:
+            transformer_state_dict = {
+                key: value for key, value in state_dict.items() if not key.startswith("net.llm_adapter.")
+            }
+            return transformer_state_dict, text_conditioner_state_dict
+
+        raise ValueError(
+            "Anima native loader could not find text-conditioner weights. Expected either "
+            "`text_conditioner.*` Diffusers keys or `net.llm_adapter.*` original Anima keys."
+        )
+
+    def _load_native_denoiser_components(
+        self,
+        state_dict: dict[str, torch.Tensor] | None = None,
+    ) -> tuple[AnimaTransformer3DModel, AnimaTextConditioner]:
+        if state_dict is None:
+            model_path = self.od_config.model
+            if not isinstance(model_path, str) or not os.path.isfile(model_path):
+                raise ValueError(f"AnimaPipeline currently requires a local file path, got: {model_path}")
+
+            state_dict = load_file(model_path, device="cpu")
+        transformer_state_dict, text_conditioner_state_dict = self._split_single_file_state_dict(state_dict)
+        del state_dict
+        transformer_state_dict = self._convert_original_transformer_state_dict(transformer_state_dict)
+
+        config_path = Path(self.od_config.model).with_suffix(".json")
+        config = json.loads(config_path.read_text()) if config_path.is_file() else {}
+        text_conditioner_config = {**ANIMA_TEXT_CONDITIONER_CONFIG, **config.get("text_conditioner", {})}
+        transformer_config = {**ANIMA_TRANSFORMER_CONFIG, **config.get("transformer", {})}
+
+        native_text_conditioner = AnimaTextConditioner(**cast(dict[str, Any], text_conditioner_config))
+        native_text_conditioner.load_state_dict(text_conditioner_state_dict, strict=True)
+        native_text_conditioner.to(device=self.device, dtype=self.od_config.dtype)
+        del text_conditioner_state_dict
+
+        native_transformer = AnimaTransformer3DModel(**cast(dict[str, Any], transformer_config))
+        native_transformer.load_state_dict(transformer_state_dict, strict=True)
+        native_transformer.to(device=self.device, dtype=self.od_config.dtype)
+        del transformer_state_dict
+
+        return native_transformer, native_text_conditioner
+
+    def _component_path(self, name: str, default_model: str) -> str:
+        custom_args = self.od_config.custom_pipeline_args or {}
+        return str(custom_args.get(f"{name}_model", custom_args.get(f"{name}_path", default_model)))
+
+    @staticmethod
+    def _from_pretrained_kwargs(model: str, subfolder: str | None = None) -> dict[str, Any]:
+        kwargs: dict[str, Any] = {"local_files_only": os.path.exists(model)}
+        if subfolder is not None and os.path.isdir(os.path.join(model, subfolder)):
+            kwargs["subfolder"] = subfolder
+        return kwargs
+
+    def _load_outer_components(self) -> tuple[Any, Any, Any, Any, Any]:
+        custom_args = self.od_config.custom_pipeline_args or {}
+        components_model_value = custom_args.get("components_model", custom_args.get("components_path"))
+        if components_model_value is None:
+            if not isinstance(self.od_config.model, str):
+                raise ValueError("AnimaPipeline requires `model` or `components_model` to load outer components.")
+            checkpoint_dir = os.path.dirname(os.path.abspath(self.od_config.model))
+            components_model = checkpoint_dir if os.path.isdir(checkpoint_dir) else self.od_config.model
+        else:
+            components_model = str(components_model_value)
+
+        text_encoder_model = self._component_path("text_encoder", components_model)
+        vae_model = self._component_path("vae", components_model)
+        tokenizer_model = self._component_path("tokenizer", components_model)
+        t5_tokenizer_model = self._component_path("t5_tokenizer", components_model)
+
+        text_encoder = Qwen3Model.from_pretrained(
+            text_encoder_model,
+            torch_dtype=self.od_config.dtype,
+            **self._from_pretrained_kwargs(text_encoder_model, "text_encoder"),
+        ).to(self.device)
+        vae = DistributedAutoencoderKLQwenImage.from_pretrained(
+            vae_model,
+            torch_dtype=self.od_config.dtype,
+            **self._from_pretrained_kwargs(vae_model, "vae"),
+        ).to(self.device)
+        tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_model,
+            **self._from_pretrained_kwargs(tokenizer_model, "tokenizer"),
+        )
+        t5_tokenizer = T5TokenizerFast.from_pretrained(
+            t5_tokenizer_model,
+            **self._from_pretrained_kwargs(t5_tokenizer_model, "t5_tokenizer"),
+        )
+
+        scheduler_model = self._component_path("scheduler", components_model)
+        has_scheduler_config = os.path.isfile(os.path.join(scheduler_model, "scheduler_config.json")) or os.path.isdir(
+            os.path.join(scheduler_model, "scheduler")
+        )
+        if has_scheduler_config:
+            scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+                scheduler_model,
+                **self._from_pretrained_kwargs(scheduler_model, "scheduler"),
+            )
+        else:
+            scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0)
+
+        return text_encoder, tokenizer, t5_tokenizer, vae, scheduler
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]] | None = None) -> set[str]:
+        logger.info("Loading Anima transformer/text-conditioner directly from single-file checkpoint.")
+        state_dict = dict(weights) if weights is not None else None
+        loaded = set(state_dict.keys()) if state_dict is not None else set()
+        if not state_dict:
+            state_dict = None
+        with set_current_diffusion_config(self.od_config):
+            native_transformer, native_text_conditioner = self._load_native_denoiser_components(state_dict)
+        text_encoder, tokenizer, t5_tokenizer, vae, scheduler = self._load_outer_components()
+
+        self.transformer = native_transformer
+        self.text_conditioner = native_text_conditioner
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.t5_tokenizer = t5_tokenizer
+        self.vae = vae
+        self.scheduler = scheduler
+        self.vae_scale_factor = _anima_vae_scale_factor_from_vae(self.vae)
+        self._setup_profiler()
+        return loaded
+
+    def _setup_profiler(self) -> None:
+        if self._profiler_initialized:
+            return
+        self.setup_diffusion_pipeline_profiler(
+            profiler_targets=["vae.decode", "diffuse", "text_encoder.forward"],
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler,
+        )
+        self._profiler_initialized = True
+
+    def _extract_prompts(self, prompts: list[Any]) -> tuple[str | list[str] | None, str | list[str] | None]:
+        """Extract prompt and negative_prompt from OmniPromptType list."""
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts):
+            negative_prompt = None
+        elif prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts]
+        else:
+            negative_prompt = None
+        return prompt, negative_prompt
+
+    def _get_qwen_prompt_embeds(
+        self,
+        prompt: str | list[str],
+        max_sequence_length: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="longest",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids.to(device)
+        prompt_attention_mask = text_inputs.attention_mask.to(device)
+        if text_input_ids.shape[-1] == 0:
+            text_input_ids = text_input_ids.new_zeros((text_input_ids.shape[0], 1))
+            prompt_attention_mask = prompt_attention_mask.new_zeros((prompt_attention_mask.shape[0], 1))
+
+        prompt_embeds = self.text_encoder(
+            input_ids=text_input_ids,
+            attention_mask=prompt_attention_mask,
+            output_hidden_states=False,
+        ).last_hidden_state
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+        prompt_embeds = prompt_embeds * prompt_attention_mask.to(prompt_embeds).unsqueeze(-1)
+
+        return prompt_embeds, prompt_attention_mask
+
+    def _get_t5_prompt_ids(
+        self,
+        prompt: str | list[str],
+        max_sequence_length: int,
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        text_inputs = self.t5_tokenizer(
+            prompt,
+            padding="longest",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        return text_inputs.input_ids.to(device), text_inputs.attention_mask.to(device)
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str],
+        negative_prompt: str | list[str] | None = None,
+        prepare_unconditional_embeds: bool = True,
+        max_sequence_length: int = 512,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> dict[str, torch.Tensor | None]:
+        device = device or self.device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        prompt_embeds, prompt_attention_mask = self._get_qwen_prompt_embeds(
+            prompt=prompt,
+            max_sequence_length=max_sequence_length,
+            device=device,
+            dtype=dtype,
+        )
+        t5_input_ids, t5_attention_mask = self._get_t5_prompt_ids(
+            prompt=prompt,
+            max_sequence_length=max_sequence_length,
+            device=device,
+        )
+
+        negative_prompt_embeds = None
+        negative_prompt_attention_mask = None
+        negative_t5_input_ids = None
+        negative_t5_attention_mask = None
+        if prepare_unconditional_embeds:
+            if negative_prompt is None:
+                negative_prompt = ""
+            if isinstance(negative_prompt, str):
+                negative_prompt = [negative_prompt] * batch_size
+            elif len(negative_prompt) != batch_size:
+                raise ValueError(
+                    f"`negative_prompt` length ({len(negative_prompt)}) must match prompt batch size ({batch_size})."
+                )
+
+            negative_prompt_embeds, negative_prompt_attention_mask = self._get_qwen_prompt_embeds(
+                prompt=negative_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+                dtype=dtype,
+            )
+            negative_t5_input_ids, negative_t5_attention_mask = self._get_t5_prompt_ids(
+                prompt=negative_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+            )
+
+        return {
+            "qwen_prompt_embeds": prompt_embeds,
+            "qwen_attention_mask": prompt_attention_mask,
+            "t5_input_ids": t5_input_ids,
+            "t5_attention_mask": t5_attention_mask,
+            "negative_qwen_prompt_embeds": negative_prompt_embeds,
+            "negative_qwen_attention_mask": negative_prompt_attention_mask,
+            "negative_t5_input_ids": negative_t5_input_ids,
+            "negative_t5_attention_mask": negative_t5_attention_mask,
+        }
+
+    def condition_prompt_embeds(
+        self,
+        qwen_prompt_embeds: torch.Tensor,
+        qwen_attention_mask: torch.Tensor,
+        t5_input_ids: torch.Tensor,
+        t5_attention_mask: torch.Tensor,
+        device: torch.device | None = None,
+        conditioning_dtype: torch.dtype | None = None,
+        output_dtype: torch.dtype | None = None,
+    ) -> torch.Tensor:
+        device = device or self.device
+        conditioning_dtype = conditioning_dtype or self.text_conditioner.dtype
+        output_dtype = output_dtype or self.transformer.dtype
+
+        prompt_embeds = self.text_conditioner(
+            source_hidden_states=qwen_prompt_embeds.to(device=device, dtype=conditioning_dtype),
+            target_input_ids=t5_input_ids.to(device),
+            target_attention_mask=t5_attention_mask.to(device),
+            source_attention_mask=qwen_attention_mask.to(device),
+        )
+        return prompt_embeds.to(dtype=output_dtype, device=device)
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        num_channels_latents: int,
+        height: int,
+        width: int,
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        device = device or self.device
+        dtype = dtype or self.transformer.dtype
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        latent_height = height // self.vae_scale_factor
+        latent_width = width // self.vae_scale_factor
+        shape = (batch_size, num_channels_latents, 1, latent_height, latent_width)
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        # Draw latents in float32 then cast
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        sigma_max = (
+            getattr(self.scheduler.config, "sigma_max", 1.0)
+            if hasattr(self, "scheduler") and hasattr(self.scheduler, "config")
+            else 1.0
+        )
+        return latents.to(dtype=dtype) * sigma_max
+
+    def prepare_timesteps(
+        self,
+        num_inference_steps: int,
+        sigmas: list[float] | None = None,
+        device: torch.device | None = None,
+    ) -> tuple[torch.Tensor, int]:
+        device = device or self.device
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        return retrieve_timesteps(
+            self.scheduler,
+            device=device,
+            sigmas=sigmas,
+        )
+
+    def diffuse(
+        self,
+        prompt_embeds: torch.Tensor,
+        negative_prompt_embeds: torch.Tensor | None,
+        latents: torch.Tensor,
+        padding_mask: torch.Tensor,
+        timesteps: torch.Tensor,
+        do_true_cfg: bool,
+        true_cfg_scale: float,
+    ) -> torch.Tensor:
+        self.scheduler.set_begin_index(0)
+
+        with self.progress_bar(total=len(timesteps)) as pbar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+                self._current_timestep = t
+
+                # Broadcast timestep to match batch size
+                timestep = t.expand(latents.shape[0]).to(dtype=self.transformer.dtype, device=latents.device)
+                timestep = timestep / 1000.0  # normalize to 0..1 range
+
+                # Predict noise
+                latent_model_input = latents.to(self.transformer.dtype)
+
+                # Positive pass
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    timestep=timestep,
+                    encoder_hidden_states=prompt_embeds,
+                    padding_mask=padding_mask,
+                    return_dict=False,
+                )[0]
+
+                # Negative / uncond pass
+                if do_true_cfg and negative_prompt_embeds is not None:
+                    noise_pred_uncond = self.transformer(
+                        hidden_states=latent_model_input,
+                        timestep=timestep,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        padding_mask=padding_mask,
+                        return_dict=False,
+                    )[0]
+                    noise_pred = noise_pred_uncond + true_cfg_scale * (noise_pred - noise_pred_uncond)
+
+                # Scheduler step
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred.to(latents_dtype), t, latents, return_dict=False)[0]
+                if latents.dtype != latents_dtype:
+                    latents = latents.to(latents_dtype)
+
+                pbar.update()
+
+        return latents
+
+    def decode_latents(
+        self,
+        latents: torch.Tensor,
+        output_type: str = "pil",
+    ) -> DiffusionOutput:
+        """Decode final latents."""
+        if output_type == "latent":
+            return DiffusionOutput(
+                output=latents,
+                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+            )
+
+        latents = latents.to(self.vae.dtype)
+        latents_mean = (
+            torch.tensor(self.vae.config.latents_mean)
+            .view(1, self.vae.config.z_dim, 1, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents_std = (
+            torch.tensor(self.vae.config.latents_std)
+            .view(1, self.vae.config.z_dim, 1, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents = latents * latents_std + latents_mean
+        image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
+        return DiffusionOutput(
+            output=image,
+            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+        )
+
+    def forward(
+        self,
+        req: DiffusionRequestBatch,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
+        true_cfg_scale: float = 4.0,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int = 50,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 4.0,
+        num_images_per_prompt: int = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        output_type: str | None = "pil",
+        max_sequence_length: int = 512,
+    ) -> list[DiffusionOutput]:
+        if req.num_reqs != 1:
+            raise ValueError("AnimaPipeline supports one request per forward.")
+        extracted_prompt, negative_prompt = self._extract_prompts(req.prompts)
+        prompt = extracted_prompt or prompt
+        if prompt is None:
+            prompt = ""
+
+        height = req.sampling_params.height or height or 1024
+        width = req.sampling_params.width or width or 1024
+        height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
+        if height is None or width is None:
+            raise ValueError("AnimaPipeline requires resolved height and width.")
+        height = int(height)
+        width = int(width)
+        if req.sampling_params.output_type is not None:
+            output_type = req.sampling_params.output_type
+        output_type = output_type or "pil"
+
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas or sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
+        generator = req.sampling_params.generator or generator
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
+            if req.sampling_params.true_cfg_scale is None:
+                true_cfg_scale = guidance_scale
+        if req.sampling_params.true_cfg_scale is not None:
+            true_cfg_scale = req.sampling_params.true_cfg_scale
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
+
+        self._guidance_scale = guidance_scale
+        self._interrupt = False
+
+        if isinstance(prompt, str):
+            batch_size = 1
+        else:
+            batch_size = len(prompt)
+
+        do_true_cfg = guidance_scale > 1.0
+
+        # Encode prompts
+        enc_out = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            prepare_unconditional_embeds=do_true_cfg,
+            max_sequence_length=max_sequence_length,
+            device=self.device,
+            dtype=self.text_encoder.dtype,
+        )
+        qwen_prompt_embeds = enc_out["qwen_prompt_embeds"]
+        qwen_attention_mask = enc_out["qwen_attention_mask"]
+        t5_input_ids = enc_out["t5_input_ids"]
+        t5_attention_mask = enc_out["t5_attention_mask"]
+        if (
+            qwen_prompt_embeds is None
+            or qwen_attention_mask is None
+            or t5_input_ids is None
+            or t5_attention_mask is None
+        ):
+            raise ValueError("Anima prompt encoding did not produce all required positive embeddings.")
+
+        # Condition prompt embeds
+        prompt_embeds = self.condition_prompt_embeds(
+            qwen_prompt_embeds=qwen_prompt_embeds,
+            qwen_attention_mask=qwen_attention_mask,
+            t5_input_ids=t5_input_ids,
+            t5_attention_mask=t5_attention_mask,
+            device=self.device,
+        )
+
+        negative_prompt_embeds = None
+        negative_qwen_prompt_embeds = enc_out["negative_qwen_prompt_embeds"]
+        if do_true_cfg and negative_qwen_prompt_embeds is not None:
+            negative_qwen_attention_mask = enc_out["negative_qwen_attention_mask"]
+            negative_t5_input_ids = enc_out["negative_t5_input_ids"]
+            negative_t5_attention_mask = enc_out["negative_t5_attention_mask"]
+            if (
+                negative_qwen_attention_mask is None
+                or negative_t5_input_ids is None
+                or negative_t5_attention_mask is None
+            ):
+                raise ValueError("Anima prompt encoding did not produce all required negative embeddings.")
+            negative_prompt_embeds = self.condition_prompt_embeds(
+                qwen_prompt_embeds=negative_qwen_prompt_embeds,
+                qwen_attention_mask=negative_qwen_attention_mask,
+                t5_input_ids=negative_t5_input_ids,
+                t5_attention_mask=negative_t5_attention_mask,
+                device=self.device,
+            )
+
+        # Repeat for num_images_per_prompt
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        if negative_prompt_embeds is not None:
+            _, seq_len, _ = negative_prompt_embeds.shape
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        # Prepare latents
+        latents = self.prepare_latents(
+            batch_size=batch_size * num_images_per_prompt,
+            num_channels_latents=16,  # Cosmos/Anima has 16 channels
+            height=height,
+            width=width,
+            dtype=torch.float32,
+            device=self.device,
+            generator=generator,
+            latents=latents,
+        )
+
+        # Prepare padding mask
+        padding_mask = latents.new_zeros(1, 1, height, width, dtype=self.transformer.dtype)
+
+        # Prepare timesteps
+        timesteps, num_inference_steps = self.prepare_timesteps(
+            num_inference_steps=num_inference_steps,
+            sigmas=sigmas,
+            device=self.device,
+        )
+        self._num_timesteps = len(timesteps)
+
+        # Denoise loop
+        latents = self.diffuse(
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            latents=latents,
+            padding_mask=padding_mask,
+            timesteps=timesteps,
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=true_cfg_scale,
+        )
+
+        self._current_timestep = None
+        return [self.decode_latents(latents, output_type=output_type)]
+
+    @property
+    def guidance_scale(self) -> float:
+        return self._guidance_scale
+
+    @property
+    def num_timesteps(self) -> int:
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self) -> torch.Tensor | None:
+        return self._current_timestep
+
+    @property
+    def interrupt(self) -> bool:
+        return self._interrupt
@@ -19,6 +19,19 @@ from vllm_omni.platforms import current_omni_platform

 logger = init_logger(__name__)

+_NATIVE_SINGLE_FILE_MODELS = {
+    "AnimaPipeline": ("AnimaModularPipeline",),
+}
+
+
+def resolve_native_single_file(model_class_name: str | None) -> str | None:
+    """Return the canonical native pipeline for a single-file model class."""
+    for canonical, aliases in _NATIVE_SINGLE_FILE_MODELS.items():
+        if model_class_name == canonical or model_class_name in aliases:
+            return canonical
+    return None
+
+
 _DIFFUSION_MODELS = {
    # arch:(mod_folder, mod_relname, cls_name)
    "QwenImagePipeline": (
@@ -361,6 +374,11 @@ _DIFFUSION_MODELS = {
        "pipeline_dreamzero",
        "DreamZeroPipeline",
    ),
+    "AnimaPipeline": (
+        "anima",
+        "pipeline_anima",
+        "AnimaPipeline",
+    ),
    "StableDiffusionXLPipeline": (
        "sdxl",
        "pipeline_sdxl",
@@ -387,6 +405,7 @@ DiffusionModelRegistry = _ModelRegistry(
 _NO_CACHE_ACCELERATION = {
    # Pipelines that do not support cache acceleration (cache_dit / tea_cache).
    "NextStep11Pipeline",
+    "AnimaPipeline",
    # π0 is a flow-matching VLA with a self-contained sample_actions loop and no
    # DiT-style ``.transformer`` block list, so cache_dit / tea_cache cannot apply
    # to it; list it here so a stray cache_backend override disables gracefully
@@ -564,6 +583,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
    # `post_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
    # where mod_folder and mod_relname are  defined and mapped using `_DIFFUSION_MODELS` via the `arch` key
    "QwenImagePipeline": "get_qwen_image_post_process_func",
+    "AnimaPipeline": "get_anima_post_process_func",
    "QwenImageEditPipeline": "get_qwen_image_edit_post_process_func",
    "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_post_process_func",
    "GlmImagePipeline": "get_glm_image_post_process_func",
@@ -1568,31 +1568,31 @@ class WorkerWrapperBase:
            gpu_id: GPU device ID
            od_config: OmniDiffusionConfig configuration
            worker_extension_cls: Optional qualified name of worker extension class
-            custom_pipeline_args: Optional arguments for custom pipeline initialization
+            custom_pipeline_args: Optional arguments passed to native pipelines.
+                A ``pipeline_class`` entry triggers custom pipeline initialization.
        """
        self.gpu_id = gpu_id
        self.od_config = od_config
        self.base_worker_class = base_worker_class
        self.worker_extension_cls = worker_extension_cls
        self.custom_pipeline_args = custom_pipeline_args
+        self.uses_custom_pipeline = bool(custom_pipeline_args and "pipeline_class" in custom_pipeline_args)

        # Prepare worker class with extension support
        worker_class = self._prepare_worker_class()

        # Create the actual worker instance
-        # When custom_pipeline_args is provided, skip initial model loading
-        # since re_init_pipeline will handle it. This avoids allocating memory
-        # through CuMemAllocator twice, which causes assertion failures in
-        # sleep mode.
+        # Only dynamic custom pipelines skip initial loading; native pipelines
+        # may also use custom_pipeline_args for model-specific component paths.
        self.worker = worker_class(
            local_rank=gpu_id,
            rank=gpu_id,
            od_config=od_config,
-            skip_load_model=(self.custom_pipeline_args is not None),
+            skip_load_model=self.uses_custom_pipeline,
        )

        # Re-initialize pipeline with custom pipeline if provided
-        if self.custom_pipeline_args is not None:
+        if self.uses_custom_pipeline:
            self.worker.re_init_pipeline(self.custom_pipeline_args)

    def _prepare_worker_class(self) -> type:
@@ -1605,8 +1605,8 @@ class WorkerWrapperBase:
        """
        worker_class = self.base_worker_class

-        # If custom_pipeline_args is provided, use CustomPipelineWorkerExtension
-        if self.custom_pipeline_args is not None:
+        # If custom pipeline initialization is requested, use CustomPipelineWorkerExtension.
+        if self.uses_custom_pipeline:
            # Set worker_extension_cls to CustomPipelineWorkerExtension if not already set
            if self.worker_extension_cls is None:
                self.worker_extension_cls = CustomPipelineWorkerExtension
@@ -181,9 +181,9 @@ class OmniEngineArgs(EngineArgs):
            (e.g. ["text", "audio"]). If None, all modalities supported by
            the model are used.
        log_stats: Whether to log engine statistics. Defaults to False.
-        custom_pipeline_args: Dictionary of arguments for custom pipeline
-            initialization (e.g., ``{"pipeline_class": "my.Module"}``).
-            Passed through to the diffusion stage engine.
+        custom_pipeline_args: Dictionary of arguments passed through to the
+            diffusion pipeline. When it contains ``pipeline_class``, it triggers
+            custom pipeline initialization.
    """

    stage_id: int = 0
@@ -25,6 +25,7 @@ from vllm.entrypoints.launchers.cli_args import make_arg_parser, validate_parsed
 from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm.logger import init_logger

+from vllm_omni.diffusion.registry import resolve_native_single_file
 from vllm_omni.entrypoints.cli.logo import log_logo
 from vllm_omni.entrypoints.openai.api_server import (
    omni_run_server,
@@ -305,7 +306,8 @@ class OmniServeCommand(CLISubcommand):
        from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model

        model = getattr(args, "model_tag", None) or getattr(args, "model", None)
-        if model and is_diffusion_model(model):
+        native_single_file = resolve_native_single_file(getattr(args, "model_class_name", None))
+        if model and ((native_single_file is not None and os.path.isfile(model)) or is_diffusion_model(model)):
            if api_server_count is not None and api_server_count > 1:
                raise ValueError("--api-server-count > 1 is not supported for diffusion models")
            logger.info("Detected diffusion model: %s", model)
@@ -636,6 +638,16 @@ class OmniServeCommand(CLISubcommand):
                '(e.g. \'{"num_inference_steps": 30, "guidance_scale": 7.5}\').'
            ),
        )
+        omni_config_group.add_argument(
+            "--custom-pipeline-args",
+            dest="custom_pipeline_args",
+            type=json.loads,
+            default=None,
+            help=(
+                "JSON object passed to native/custom diffusion pipelines. "
+                'Only args containing "pipeline_class" trigger custom pipeline re-initialization.'
+            ),
+        )
        omni_config_group.add_argument(
            "--usp",
            "--ulysses-degree",