119 lines
4.3 KiB
Python
119 lines
4.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.models.magi2.configuration_magi2 import (
|
|
Magi2MHCConfig,
|
|
Magi2MoEConfig,
|
|
Magi2PreviewConfig,
|
|
)
|
|
from vllm_omni.diffusion.models.magi2.layers import MultiModalityRMSNorm
|
|
from vllm_omni.diffusion.models.magi2.mh_moe import Magi2MultiHeadMoE
|
|
from vllm_omni.diffusion.models.magi2.modeling_magi2 import Magi2PreviewTransformer
|
|
from vllm_omni.diffusion.models.magi2.preview_data_proxy import (
|
|
Magi2DataProxy,
|
|
Magi2PreviewDataProxyConfig,
|
|
)
|
|
from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig, Magi2PreviewSampler
|
|
|
|
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
|
|
|
|
|
def _tiny_config(params_dtype: torch.dtype = torch.float32, *, moe_layers: int = 1) -> Magi2PreviewConfig:
|
|
# The leading layers are multimodal with MoE, the rest single-modality
|
|
# dense, so every layer belongs to one of two kinds.
|
|
moe_indices = tuple(range(moe_layers))
|
|
return Magi2PreviewConfig(
|
|
num_layers=2 * moe_layers,
|
|
hidden_size=16,
|
|
head_dim=8,
|
|
num_query_groups=2,
|
|
video_in_channels=4,
|
|
audio_in_channels=4,
|
|
text_in_channels=4,
|
|
intermediate_factor=2,
|
|
multimodal_layers=moe_indices,
|
|
params_dtype=params_dtype,
|
|
mhc=Magi2MHCConfig(num_streams=2),
|
|
moe=Magi2MoEConfig(
|
|
num_heads=2,
|
|
num_experts=4,
|
|
top_k=2,
|
|
expert_intermediate_size=8,
|
|
shared_expert_intermediate_size=8,
|
|
modality_shared_expert_intermediate_size=8,
|
|
layers=moe_indices,
|
|
),
|
|
)
|
|
|
|
|
|
def _tiny_model(
|
|
seed: int = 11,
|
|
params_dtype: torch.dtype = torch.float32,
|
|
*,
|
|
moe_layers: int = 1,
|
|
) -> Magi2PreviewTransformer:
|
|
model = Magi2PreviewTransformer(_tiny_config(params_dtype, moe_layers=moe_layers))
|
|
generator = torch.Generator(device="cpu").manual_seed(seed)
|
|
with torch.no_grad():
|
|
for parameter in model.parameters():
|
|
parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02)
|
|
for module in model.modules():
|
|
if isinstance(module, MultiModalityRMSNorm):
|
|
module.weight.zero_()
|
|
elif isinstance(module, Magi2MultiHeadMoE):
|
|
module.router.expert_bias.zero_()
|
|
module.router.expert_bias_ema.zero_()
|
|
return model
|
|
|
|
|
|
def _tiny_sampler(model: torch.nn.Module) -> Magi2PreviewSampler:
|
|
return Magi2PreviewSampler(model, Magi2DataProxy(Magi2PreviewDataProxyConfig(time_channel_dim=8)))
|
|
|
|
|
|
def _sampler_tensors(seed: int) -> dict[str, torch.Tensor]:
|
|
generator = torch.Generator(device="cpu").manual_seed(seed)
|
|
return {
|
|
"latent": torch.randn(1, 4, 2, 1, 3, generator=generator),
|
|
"audio_latent": torch.randn(1, 5, 4, generator=generator),
|
|
"txt_feat": torch.randn(1, 3, 4, generator=generator),
|
|
"null_txt_feat": torch.randn(1, 2, 4, generator=generator),
|
|
}
|
|
|
|
|
|
def _longer_text_tensors(seed: int) -> dict[str, torch.Tensor]:
|
|
tensors = _sampler_tensors(seed)
|
|
generator = torch.Generator(device="cpu").manual_seed(seed + 100)
|
|
tensors["txt_feat"] = torch.randn(1, 4, 4, generator=generator)
|
|
return tensors
|
|
|
|
|
|
def test_prepare_model_input_keeps_lengths_on_host() -> None:
|
|
sampler = Magi2PreviewSampler(torch.nn.Identity())
|
|
model_input = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([500.0]), cfg_config=CFGConfig())
|
|
|
|
assert model_input.audio_feat_len == [5, 5]
|
|
assert model_input.txt_feat_len == [3, 2]
|
|
assert model_input.ref_audio_feat_len == [0, 0]
|
|
assert model_input.ref_video_feat_len == [0, 0]
|
|
|
|
positive, negative = Magi2PreviewSampler._split_cfg_model_input(model_input)
|
|
assert positive.txt_feat_len == [3]
|
|
assert negative.txt_feat_len == [2]
|
|
|
|
|
|
def test_pre_adapter_embeds_directly_in_checkpoint_dtype() -> None:
|
|
model = _tiny_model(params_dtype=torch.bfloat16)
|
|
packed = torch.randn(6, 4)
|
|
indices = torch.tensor([0, 1]), torch.tensor([2, 3]), torch.tensor([4, 5])
|
|
|
|
with torch.no_grad():
|
|
hidden = model.pre_adapter(packed, *indices)
|
|
|
|
assert hidden.dtype == torch.bfloat16
|
|
assert hidden.shape == (6, model.config.virtual_width)
|