455 lines
18 KiB
Python
455 lines
18 KiB
Python
# 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)
|