diff --git a/docs/features/custom_pipeline.md b/docs/features/custom_pipeline.md index 887dfa769..463798d6d 100644 --- a/docs/features/custom_pipeline.md +++ b/docs/features/custom_pipeline.md @@ -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 + 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". ) ``` diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index f481154fb..ec895f794 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -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 diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index cfe05255d..7664bb32d 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -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 diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index adcc10669..71689b272 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -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 diff --git a/recipes/README.md b/recipes/README.md index a7c10ca14..c84223510 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -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 | diff --git a/recipes/circlestone-labs/Anima.md b/recipes/circlestone-labs/Anima.md new file mode 100644 index 000000000..bcaf84364 --- /dev/null +++ b/recipes/circlestone-labs/Anima.md @@ -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. diff --git a/tests/diffusion/models/anima/__init__.py b/tests/diffusion/models/anima/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/tests/diffusion/models/anima/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/tests/diffusion/models/anima/test_anima.py b/tests/diffusion/models/anima/test_anima.py new file mode 100644 index 000000000..124171c93 --- /dev/null +++ b/tests/diffusion/models/anima/test_anima.py @@ -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) diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index 709fc7283..fb765e8fd 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -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) diff --git a/tests/e2e/offline_inference/test_anima.py b/tests/e2e/offline_inference/test_anima.py new file mode 100644 index 000000000..877fcfe2d --- /dev/null +++ b/tests/e2e/offline_inference/test_anima.py @@ -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) diff --git a/tests/e2e/online_serving/test_anima.py b/tests/e2e/online_serving/test_anima.py new file mode 100644 index 000000000..d3dc76cce --- /dev/null +++ b/tests/e2e/online_serving/test_anima.py @@ -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) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 7f8b6a4a9..0f51a4320 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -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"): diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 65395265b..338df4833 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -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 diff --git a/tests/helpers/tests/test_check_test_marks.py b/tests/helpers/tests/test_check_test_marks.py new file mode 100644 index 000000000..8b29b6f33 --- /dev/null +++ b/tests/helpers/tests/test_check_test_marks.py @@ -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"]} diff --git a/tests/model_tests/conftest.py b/tests/model_tests/conftest.py index 7760fa25d..28cfa8891 100644 --- a/tests/model_tests/conftest.py +++ b/tests/model_tests/conftest.py @@ -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)) diff --git a/tests/model_tests/diffusion/anima_builder.py b/tests/model_tests/diffusion/anima_builder.py new file mode 100644 index 000000000..ca0f9e29a --- /dev/null +++ b/tests/model_tests/diffusion/anima_builder.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Local random-weight Anima assets and real-checkpoint assets for common tests.""" + +import json +import tempfile +from pathlib import Path + +import torch +from diffusers import AutoencoderKLQwenImage, FlowMatchEulerDiscreteScheduler +from huggingface_hub import hf_hub_download, snapshot_download +from safetensors.torch import save_file +from tokenizers import Tokenizer +from tokenizers.models import Unigram, WordLevel +from tokenizers.pre_tokenizers import Whitespace +from tokenizers.processors import TemplateProcessing +from transformers import PreTrainedTokenizerFast, Qwen3Config, Qwen3Model, T5TokenizerFast + +from vllm_omni.diffusion.models.anima.anima_text_conditioner import ( + ANIMA_TEXT_CONDITIONER_CONFIG, + AnimaTextConditioner, +) +from vllm_omni.diffusion.models.anima.anima_transformer import ANIMA_TRANSFORMER_CONFIG, AnimaTransformer3DModel + +CHECKPOINT_FILENAME = "anima.safetensors" + + +def tiny_anima_builder() -> str: + """Build all components locally, without downloading any model weights or tokenizers.""" + model_dir = Path(tempfile.mkdtemp(prefix="tiny-anima-")) + transformer_config = { + **ANIMA_TRANSFORMER_CONFIG, + "num_attention_heads": 2, + "attention_head_dim": 32, + "num_layers": 1, + "mlp_ratio": 1.0, + "text_embed_dim": 32, + "adaln_lora_dim": 8, + "max_size": (1, 64, 64), + } + conditioner_config = { + **ANIMA_TEXT_CONDITIONER_CONFIG, + "source_dim": 32, + "target_dim": 32, + "model_dim": 32, + "num_layers": 1, + "num_attention_heads": 2, + "mlp_ratio": 1.0, + "target_vocab_size": 8, + "min_sequence_length": 8, + } + transformer = AnimaTransformer3DModel(**transformer_config).to(torch.bfloat16) + conditioner = AnimaTextConditioner(**conditioner_config).to(torch.bfloat16) + save_file( + { + **{f"transformer.{name}": tensor for name, tensor in transformer.state_dict().items()}, + **{f"text_conditioner.{name}": tensor for name, tensor in conditioner.state_dict().items()}, + }, + str(model_dir / CHECKPOINT_FILENAME), + ) + (model_dir / "anima.json").write_text( + json.dumps({"transformer": transformer_config, "text_conditioner": conditioner_config}) + ) + Qwen3Model( + Qwen3Config( + vocab_size=8, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=16, + ) + ).to(torch.bfloat16).save_pretrained(model_dir / "text_encoder") + AutoencoderKLQwenImage( + base_dim=4, + z_dim=16, + dim_mult=[1, 1, 1, 1], + num_res_blocks=1, + latents_mean=[0.0] * 16, + latents_std=[1.0] * 16, + ).to(torch.bfloat16).save_pretrained(model_dir / "vae") + tokenizer = Tokenizer(WordLevel({"": 0, "": 1, "": 2, "Dummy": 3, "prompt": 4}, unk_token="")) + tokenizer.pre_tokenizer = Whitespace() + tokenizer.post_processor = TemplateProcessing(single="$A ", special_tokens=[("", 1)]) + PreTrainedTokenizerFast( + tokenizer_object=tokenizer, pad_token="", eos_token="", unk_token="" + ).save_pretrained(model_dir / "tokenizer") + t5_tokenizer = Tokenizer( + Unigram( + [("", 0.0), ("", 0.0), ("", 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) diff --git a/tests/model_tests/diffusion/config_types.py b/tests/model_tests/diffusion/config_types.py index 634036ed4..8b66af191 100644 --- a/tests/model_tests/diffusion/config_types.py +++ b/tests/model_tests/diffusion/config_types.py @@ -1,10 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Common definitions for controlling what tests run where. """ from collections.abc import Callable from enum import Enum -from typing import NamedTuple, TypeAlias +from typing import Any, NamedTuple, TypeAlias from pytest import MarkDecorator @@ -45,7 +48,9 @@ class DiffusionModelTestOpts(NamedTuple): # HF model name for real-weight tests (advanced_model / full_model level). # For now, whether we use the real weights vs tiny weights in the common tests # depends on the run level. - model: str + # A callable may assemble native assets in a temporary directory; the fixture + # removes that directory after the tests, so it must not return a cache root. + model: str | Callable[[], str] # Creates a tiny model for the given architecture. We should always use tiny # model weights for tests that do not require us to check the model quality. @@ -81,9 +86,12 @@ class DiffusionModelTestOpts(NamedTuple): # set False for pipelines that mandate an image and fail closed without one (SANA-Video I2V). check_i2v_t2v_divergence: bool = True + # Native single-file models: builders/resolvers return the containing directory. + checkpoint_filename: str | None = None + ### Mappings & utils for building offline Omni() instances given a list of enabled accelerations -ACC_OMNI_KWARGS = { +ACC_OMNI_KWARGS: dict[DiffusionAccs, dict[str, Any]] = { DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_use_tiling": True}, DiffusionAccs.CPU_OFFLOAD: {"enable_cpu_offload": True}, DiffusionAccs.LAYERWISE_OFFLOAD: {"enable_layerwise_offload": True}, diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index 083dba0ed..8c5955c8d 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/tests/model_tests/diffusion/model_settings.py @@ -1,4 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from tests.model_tests.diffusion import diff_model_builders +from tests.model_tests.diffusion.anima_builder import CHECKPOINT_FILENAME, real_anima_model, tiny_anima_builder from tests.model_tests.diffusion.config_types import ( DiffusionAccs, DiffusionModelTestOpts, @@ -29,6 +33,12 @@ from tests.model_tests.diffusion.config_types import ( # $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline # ^ Runs all test groups for Flux2KleinPipeline only DIFFUSION_TEST_SETTINGS = { + "AnimaPipeline": DiffusionModelTestOpts( + model=real_anima_model, + builder=tiny_anima_builder, + supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE], + checkpoint_filename=CHECKPOINT_FILENAME, + ), "Flux2KleinPipeline": DiffusionModelTestOpts( model="black-forest-labs/FLUX.2-klein-4B", builder=diff_model_builders.tiny_flux2_klein_builder, diff --git a/tests/model_tests/diffusion/test_common_offline.py b/tests/model_tests/diffusion/test_common_offline.py index 9cecc31fd..b610ed013 100644 --- a/tests/model_tests/diffusion/test_common_offline.py +++ b/tests/model_tests/diffusion/test_common_offline.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from pathlib import Path + import pytest from tests.model_tests.diffusion.case_filtering import get_parametrized_options @@ -47,10 +49,17 @@ def test_pipeline_on_supported_tasks( # since starting the server can take 10+ seconds, even for tiny models. # # NOTE: Be sure to install pytest-subtests if you're running on pytest < 9 + settings = DIFFUSION_TEST_SETTINGS[model_name] + model_path = tiny_model_paths[model_name] + model_kwargs = {} + if settings.checkpoint_filename is not None: + model_path = str(Path(model_path) / settings.checkpoint_filename) + model_kwargs["model_class_name"] = model_name omni = build_omni_from_diff_accelerations( accelerations=accelerations, - model=tiny_model_paths[model_name], + model=model_path, enforce_eager=True, + **model_kwargs, ) try: for task_type in supported_tasks: diff --git a/tests/model_tests/diffusion/test_common_online.py b/tests/model_tests/diffusion/test_common_online.py index ed7f0e2fa..ef4449501 100644 --- a/tests/model_tests/diffusion/test_common_online.py +++ b/tests/model_tests/diffusion/test_common_online.py @@ -7,6 +7,8 @@ online serving stack (CLI arg parsing, subprocess, API routing, response encoding) using tiny models. """ +from pathlib import Path + import pytest from xdist import is_xdist_worker @@ -68,6 +70,10 @@ def test_online_on_supported_tasks( model_path = tiny_model_paths[model_name] server_args = build_server_args_from_diff_accelerations(accelerations) server_args.append("--enforce-eager") + settings = DIFFUSION_TEST_SETTINGS[model_name] + if settings.checkpoint_filename is not None: + model_path = str(Path(model_path) / settings.checkpoint_filename) + server_args.extend(["--model-class-name", model_name]) with OmniServer(model_path, server_args) as server: # TODO: We may want to revisit run_level validation here, diff --git a/tests/model_tests/diffusion/test_native_model_settings.py b/tests/model_tests/diffusion/test_native_model_settings.py new file mode 100644 index 000000000..10db75d05 --- /dev/null +++ b/tests/model_tests/diffusion/test_native_model_settings.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from contextlib import nullcontext +from pathlib import Path +from shutil import rmtree +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from tests.model_tests.diffusion import anima_builder, test_common_offline, test_common_online +from tests.model_tests.diffusion.config_types import DiffusionTasks + +pytestmark = [pytest.mark.cpu, pytest.mark.core_model, pytest.mark.diffusion] + + +@pytest.mark.parametrize("model_name", ["AnimaPipeline", "FluxPipeline"]) +def test_common_runners_resolve_native_checkpoint(model_name, tmp_path, monkeypatch): + """Both entry points receive the file and class; directory models keep their existing arguments.""" + paths = {model_name: str(tmp_path)} + subtests = SimpleNamespace(test=lambda **kwargs: nullcontext()) + build_omni = MagicMock() + monkeypatch.setattr(test_common_offline, "build_omni_from_diff_accelerations", build_omni) + monkeypatch.setattr(test_common_offline, "run_and_validate_text_to_image_request", MagicMock()) + test_common_offline.test_pipeline_on_supported_tasks( + model_name, None, [DiffusionTasks.TEXT_TO_IMAGE], False, False, False, paths, subtests + ) + server = MagicMock() + monkeypatch.setattr(test_common_online, "OmniServer", server) + monkeypatch.setattr(test_common_online, "OnlineOmniClient", MagicMock()) + monkeypatch.setattr(test_common_online, "run_and_validate_online_text_to_image_request", MagicMock()) + test_common_online.test_online_on_supported_tasks( + model_name, None, [DiffusionTasks.TEXT_TO_IMAGE], False, False, paths, "core_model", subtests + ) + if model_name == "AnimaPipeline": + checkpoint = str(tmp_path / anima_builder.CHECKPOINT_FILENAME) + build_omni.assert_called_once_with( + accelerations=None, model=checkpoint, enforce_eager=True, model_class_name=model_name + ) + server.assert_called_once_with(checkpoint, ["--enforce-eager", "--model-class-name", model_name]) + else: + build_omni.assert_called_once_with(accelerations=None, model=str(tmp_path), enforce_eager=True) + server.assert_called_once_with(str(tmp_path), ["--enforce-eager"]) + + +def test_real_anima_assets_leave_cache_intact(tmp_path, monkeypatch): + """The fixture may delete the assembled directory without deleting cached model assets.""" + components = tmp_path / "components" + for name in ("text_encoder", "vae", "tokenizer", "t5_tokenizer"): + (components / name).mkdir(parents=True) + (components / name / "asset").write_text(name) + checkpoint = tmp_path / "real.safetensors" + checkpoint.write_bytes(b"checkpoint") + monkeypatch.setattr(anima_builder, "snapshot_download", lambda *args, **kwargs: str(components)) + monkeypatch.setattr(anima_builder, "hf_hub_download", lambda *args, **kwargs: str(checkpoint)) + directory = Path(anima_builder.real_anima_model()) + try: + assert (directory / anima_builder.CHECKPOINT_FILENAME).read_bytes() == b"checkpoint" + assert (directory / "vae" / "asset").read_text() == "vae" + # Missing scheduler is supported by Anima's default scheduler fallback. + assert not (directory / "scheduler").exists() + finally: + rmtree(directory) + assert checkpoint.read_bytes() == b"checkpoint" + assert (components / "vae" / "asset").read_text() == "vae" diff --git a/tests/model_tests/diffusion/utils.py b/tests/model_tests/diffusion/utils.py index 35577b68f..137eae123 100644 --- a/tests/model_tests/diffusion/utils.py +++ b/tests/model_tests/diffusion/utils.py @@ -1,8 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Utilities for resolving real models to their tiny model equivalents. """ import logging +from pathlib import Path from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS from vllm_omni.diffusion.data import resolve_model_class_name @@ -30,4 +34,7 @@ def resolve_tiny_model_path(model: str) -> str: ) return model - return test_opts.builder() + model_path = test_opts.builder() + if test_opts.checkpoint_filename is not None: + return str(Path(model_path) / test_opts.checkpoint_filename) + return model_path diff --git a/tools/pre_commit/check_test_marks.py b/tools/pre_commit/check_test_marks.py index 6af0ae825..2d84039de 100644 --- a/tools/pre_commit/check_test_marks.py +++ b/tools/pre_commit/check_test_marks.py @@ -33,6 +33,12 @@ HARDWARE_HELPERS = ("hardware_test", "hardware_marks") # The helper implementation is the only file allowed to write pytest.mark.. _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): diff --git a/vllm_omni/config/resolver.py b/vllm_omni/config/resolver.py index 0db436768..76740aa06 100644 --- a/vllm_omni/config/resolver.py +++ b/vllm_omni/config/resolver.py @@ -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) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 5bc5a229b..e19470fa5 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -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, diff --git a/vllm_omni/diffusion/models/anima/__init__.py b/vllm_omni/diffusion/models/anima/__init__.py new file mode 100644 index 000000000..ddbabbcee --- /dev/null +++ b/vllm_omni/diffusion/models/anima/__init__.py @@ -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", +] diff --git a/vllm_omni/diffusion/models/anima/anima_text_conditioner.py b/vllm_omni/diffusion/models/anima/anima_text_conditioner.py new file mode 100644 index 000000000..84c3e4a2c --- /dev/null +++ b/vllm_omni/diffusion/models/anima/anima_text_conditioner.py @@ -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 diff --git a/vllm_omni/diffusion/models/anima/anima_transformer.py b/vllm_omni/diffusion/models/anima/anima_transformer.py new file mode 100644 index 000000000..4c92a4ff3 --- /dev/null +++ b/vllm_omni/diffusion/models/anima/anima_transformer.py @@ -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 diff --git a/vllm_omni/diffusion/models/anima/pipeline_anima.py b/vllm_omni/diffusion/models/anima/pipeline_anima.py new file mode 100644 index 000000000..63fdb7794 --- /dev/null +++ b/vllm_omni/diffusion/models/anima/pipeline_anima.py @@ -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 diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index fadb712df..17d8216db 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -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", diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index b9506f102..f934646dc 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -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 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index d1e53775a..b6aba463b 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -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 diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 14e13cd4a..10ae2166d 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -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",