[diffusion][model] Add LingBot Video dense and MoE support (#5035)
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com> Signed-off-by: wtz2333 <2955110911@qq.com> Co-authored-by: hsliuustc0106 <liuhongsheng4@huawei.com>
This commit is contained in:
wtz2333andhsliuustc0106 authored and GitHub committed
1 parent
2f33df4678
commit
5a6a93dc5f
24 files changed
+3765
No files matched your search
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -0,0 +1,183 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _upstream_sparse_moe_reference(block, hidden_states, padding_mask=None):
|
||||
"""Independent eager reference for the upstream LingBot sparse MoE path."""
|
||||
batch_size, _, hidden_size = hidden_states.shape
|
||||
tokens = hidden_states.reshape(-1, hidden_size)
|
||||
|
||||
logits = F.linear(tokens.float(), block.router.weight.float())
|
||||
if block.router.score_func == "softmax":
|
||||
scores = F.softmax(logits, dim=-1)
|
||||
else:
|
||||
scores = logits.sigmoid()
|
||||
scores_for_choice = scores + block.router.e_score_correction_bias.unsqueeze(0)
|
||||
|
||||
if block.router.n_group is not None and block.router.n_group > 1:
|
||||
experts_per_group = block.router.num_experts // block.router.n_group
|
||||
grouped = scores_for_choice.view(-1, block.router.n_group, experts_per_group)
|
||||
group_scores = grouped.topk(2, dim=-1)[0].sum(dim=-1)
|
||||
group_idx = torch.topk(
|
||||
group_scores,
|
||||
k=block.router.topk_group,
|
||||
dim=-1,
|
||||
sorted=False,
|
||||
)[1]
|
||||
group_mask = torch.zeros_like(group_scores)
|
||||
group_mask.scatter_(1, group_idx, 1)
|
||||
score_mask = group_mask.unsqueeze(-1).expand_as(grouped).reshape_as(scores_for_choice)
|
||||
scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf"))
|
||||
|
||||
top_indices = torch.topk(
|
||||
scores_for_choice,
|
||||
k=block.router.top_k,
|
||||
dim=-1,
|
||||
sorted=False,
|
||||
)[1]
|
||||
top_scores = scores.gather(1, top_indices)
|
||||
if block.router.top_k > 1 and block.router.norm_topk_prob:
|
||||
top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20)
|
||||
top_scores = top_scores.to(tokens.dtype) * block.router.route_scale
|
||||
|
||||
if padding_mask is not None:
|
||||
mask = padding_mask.unsqueeze(-1).to(top_scores.dtype)
|
||||
top_scores = top_scores * mask
|
||||
top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-9)
|
||||
top_scores = top_scores * block.router.route_scale
|
||||
|
||||
routed = torch.zeros(
|
||||
tokens.shape,
|
||||
dtype=torch.float32,
|
||||
device=tokens.device,
|
||||
)
|
||||
for expert_idx in range(block.router.num_experts):
|
||||
token_idx, route_idx = torch.where(top_indices == expert_idx)
|
||||
if token_idx.numel() == 0:
|
||||
continue
|
||||
expert_tokens = tokens[token_idx]
|
||||
h = F.silu(expert_tokens @ block.experts.w1[expert_idx].transpose(-2, -1))
|
||||
h = h * (expert_tokens @ block.experts.w3[expert_idx].transpose(-2, -1))
|
||||
expert_output = h @ block.experts.w2[expert_idx].transpose(-2, -1)
|
||||
weighted = expert_output.float() * top_scores[token_idx, route_idx, None].float()
|
||||
routed.index_add_(0, token_idx, weighted)
|
||||
|
||||
output = routed.to(tokens.dtype).reshape(batch_size, -1, hidden_size)
|
||||
if block.shared_experts is not None:
|
||||
shared = F.linear(hidden_states, block.shared_experts.gate_proj.weight)
|
||||
shared = F.silu(shared) * F.linear(hidden_states, block.shared_experts.up_proj.weight)
|
||||
shared = F.linear(shared, block.shared_experts.down_proj.weight)
|
||||
output = output + shared
|
||||
return output, top_indices, top_scores
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_grouped_mm_matches_per_expert_reference() -> None:
|
||||
from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module
|
||||
|
||||
if not hasattr(torch, "_grouped_mm"):
|
||||
pytest.skip("torch._grouped_mm is unavailable")
|
||||
|
||||
torch.manual_seed(42)
|
||||
block = module.LingBotVideoSparseMoeBlock(
|
||||
hidden_size=16,
|
||||
num_experts=4,
|
||||
top_k=2,
|
||||
moe_intermediate_size=8,
|
||||
score_func="sigmoid",
|
||||
norm_topk_prob=True,
|
||||
n_group=None,
|
||||
topk_group=None,
|
||||
routed_scaling_factor=1.0,
|
||||
n_shared_experts=None,
|
||||
).to(device="cuda", dtype=torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
block.experts.w1.normal_(mean=0.0, std=0.02)
|
||||
block.experts.w2.normal_(mean=0.0, std=0.02)
|
||||
block.experts.w3.normal_(mean=0.0, std=0.02)
|
||||
tokens = torch.randn(23, 16, device="cuda", dtype=torch.bfloat16)
|
||||
counts = torch.tensor([5, 0, 9, 9], device="cuda", dtype=torch.int64)
|
||||
|
||||
expected = block._run_experts_for_loop(tokens, counts)
|
||||
actual = block._run_grouped_experts(tokens, counts)
|
||||
|
||||
torch.testing.assert_close(actual, expected)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_sparse_moe_block_matches_upstream_reference() -> None:
|
||||
from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module
|
||||
|
||||
if not hasattr(torch, "_grouped_mm"):
|
||||
pytest.skip("torch._grouped_mm is unavailable")
|
||||
|
||||
torch.manual_seed(42)
|
||||
block = module.LingBotVideoSparseMoeBlock(
|
||||
hidden_size=16,
|
||||
num_experts=8,
|
||||
top_k=2,
|
||||
moe_intermediate_size=8,
|
||||
score_func="sigmoid",
|
||||
norm_topk_prob=True,
|
||||
n_group=4,
|
||||
topk_group=1,
|
||||
routed_scaling_factor=1.5,
|
||||
n_shared_experts=1,
|
||||
)
|
||||
with torch.no_grad():
|
||||
block.router.weight.zero_()
|
||||
block.router.e_score_correction_bias.copy_(torch.tensor([0.9, 0.8, 1.0, 0.0, 0.7, 0.6, 0.5, 0.4]))
|
||||
for parameter in (
|
||||
block.experts.w1,
|
||||
block.experts.w2,
|
||||
block.experts.w3,
|
||||
block.shared_experts.gate_proj.weight,
|
||||
block.shared_experts.up_proj.weight,
|
||||
block.shared_experts.down_proj.weight,
|
||||
):
|
||||
parameter.normal_(mean=0.0, std=0.02)
|
||||
block = block.to(device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
hidden_states = torch.randn(2, 5, 16, device="cuda", dtype=torch.bfloat16)
|
||||
padding_mask = torch.tensor(
|
||||
[1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
|
||||
device="cuda",
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
expected, expected_indices, expected_scores = _upstream_sparse_moe_reference(
|
||||
block,
|
||||
hidden_states,
|
||||
padding_mask,
|
||||
)
|
||||
actual_indices, actual_scores = block.router(hidden_states.reshape(-1, 16))
|
||||
actual_scores = actual_scores * padding_mask.unsqueeze(-1).to(actual_scores.dtype)
|
||||
actual_scores = actual_scores / (actual_scores.sum(dim=-1, keepdim=True) + 1e-9)
|
||||
actual_scores = actual_scores * block.router.route_scale
|
||||
actual = block(hidden_states, padding_mask=padding_mask)
|
||||
|
||||
corrected_scores = (
|
||||
torch.full(
|
||||
(8,),
|
||||
0.5,
|
||||
device="cuda",
|
||||
)
|
||||
+ block.router.e_score_correction_bias
|
||||
)
|
||||
unrestricted_indices = torch.topk(corrected_scores, k=2, sorted=False)[1]
|
||||
|
||||
assert set(actual_indices[0].tolist()) == {0, 1}
|
||||
assert set(unrestricted_indices.tolist()) == {0, 2}
|
||||
assert torch.equal(actual_indices, expected_indices)
|
||||
torch.testing.assert_close(actual_scores, expected_scores, rtol=0, atol=0)
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
assert torch.isfinite(actual).all()
|
||||
@@ -0,0 +1,266 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _tiny_transformer(**overrides):
|
||||
from vllm_omni.diffusion.models.lingbot_video import LingBotVideoTransformer3DModel
|
||||
|
||||
config = {
|
||||
"patch_size": (1, 1, 1),
|
||||
"in_channels": 2,
|
||||
"out_channels": 2,
|
||||
"hidden_size": 16,
|
||||
"num_attention_heads": 1,
|
||||
"depth": 0,
|
||||
"intermediate_size": 32,
|
||||
"text_dim": 8,
|
||||
"freq_dim": 8,
|
||||
"axes_dims": (4, 4, 8),
|
||||
"axes_lens": (32, 32, 32),
|
||||
}
|
||||
config.update(overrides)
|
||||
return LingBotVideoTransformer3DModel(**config)
|
||||
|
||||
|
||||
def test_joint_position_ids_video_then_text_order():
|
||||
from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import make_joint_position_ids
|
||||
|
||||
positions = make_joint_position_ids(text_len=3, grid_t=1, grid_h=2, grid_w=2, device=torch.device("cpu"))
|
||||
|
||||
assert positions.shape == (7, 3)
|
||||
assert positions[:4, 0].tolist() == [4, 4, 4, 4]
|
||||
assert positions[:4, 1:].tolist() == [[0, 0], [0, 1], [1, 0], [1, 1]]
|
||||
assert positions[4:].tolist() == [[1, 0, 0], [2, 0, 0], [3, 0, 0]]
|
||||
|
||||
|
||||
def test_tiny_transformer_depth_zero_forward_shape():
|
||||
model = _tiny_transformer()
|
||||
hidden_states = torch.randn(1, 2, 1, 2, 2)
|
||||
timestep = torch.tensor([300.0])
|
||||
encoder_hidden_states = torch.randn(1, 3, 8)
|
||||
encoder_attention_mask = torch.ones(1, 3, dtype=torch.long)
|
||||
|
||||
with torch.no_grad():
|
||||
out = model(
|
||||
hidden_states,
|
||||
timestep,
|
||||
encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
assert out.shape == hidden_states.shape
|
||||
assert torch.isfinite(out).all()
|
||||
|
||||
|
||||
def test_packed_attention_uses_sdpa_fallback_without_flash_varlen(monkeypatch):
|
||||
from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module
|
||||
|
||||
monkeypatch.setattr(module, "flash_attn_varlen_func_v3", None)
|
||||
attn = module.LingBotVideoAttention(
|
||||
hidden_size=8,
|
||||
num_heads=2,
|
||||
norm_eps=1e-6,
|
||||
qkv_bias=False,
|
||||
out_bias=False,
|
||||
)
|
||||
captured = {}
|
||||
|
||||
def fake_sdpa_forward(query, key, value, attn_metadata):
|
||||
captured["mask"] = attn_metadata.attn_mask
|
||||
return torch.zeros_like(query)
|
||||
|
||||
monkeypatch.setattr(attn.attn.sdpa_fallback, "forward", fake_sdpa_forward)
|
||||
x = torch.randn(1, 5, 8)
|
||||
rotary = torch.ones(1, 5, 2, dtype=torch.complex64)
|
||||
packed_indices = {
|
||||
"cu_seqlens_kv": torch.tensor([0, 2, 5], dtype=torch.int32),
|
||||
"max_seqlen_in_batch_kv": 3,
|
||||
"attention_mask": module._packed_block_attention_mask([2, 3], x.device),
|
||||
}
|
||||
|
||||
out = attn(x, rotary, packed_indices=packed_indices)
|
||||
|
||||
assert out.shape == x.shape
|
||||
mask = captured["mask"]
|
||||
assert mask.shape == (1, 1, 5, 5)
|
||||
assert mask[0, 0, :2, :2].all()
|
||||
assert mask[0, 0, 2:, 2:].all()
|
||||
assert not mask[0, 0, :2, 2:].any()
|
||||
assert not mask[0, 0, 2:, :2].any()
|
||||
|
||||
|
||||
def test_tiny_transformer_rejects_invalid_rope_dims():
|
||||
from vllm_omni.diffusion.models.lingbot_video import LingBotVideoTransformer3DModel
|
||||
|
||||
with pytest.raises(AssertionError, match="head_dim"):
|
||||
LingBotVideoTransformer3DModel(
|
||||
hidden_size=16,
|
||||
num_attention_heads=1,
|
||||
axes_dims=(4, 4, 4),
|
||||
depth=0,
|
||||
)
|
||||
|
||||
|
||||
def test_transformer_to_keeps_sensitive_modules_in_fp32():
|
||||
model = _tiny_transformer()
|
||||
|
||||
model.to(dtype=torch.bfloat16)
|
||||
|
||||
assert model.patch_embedder.weight.dtype == torch.bfloat16
|
||||
assert model.time_embedder.linear_1.weight.dtype == torch.float32
|
||||
assert model.norm_out_modulation[1].weight.dtype == torch.float32
|
||||
|
||||
|
||||
def test_router_group_limited_topk_uses_bias_corrected_choice():
|
||||
from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module
|
||||
|
||||
router = module.LingBotVideoRouter(
|
||||
hidden_size=2,
|
||||
num_experts=4,
|
||||
top_k=2,
|
||||
score_func="sigmoid",
|
||||
norm_topk_prob=True,
|
||||
n_group=2,
|
||||
topk_group=1,
|
||||
route_scale=2.5,
|
||||
)
|
||||
with torch.no_grad():
|
||||
router.weight.copy_(
|
||||
torch.tensor(
|
||||
[
|
||||
[8.0, 0.0],
|
||||
[7.0, 0.0],
|
||||
[-8.0, 0.0],
|
||||
[-7.0, 0.0],
|
||||
]
|
||||
)
|
||||
)
|
||||
router.e_score_correction_bias.copy_(torch.tensor([0.0, 0.0, 2.0, 2.0]))
|
||||
|
||||
top_indices, top_scores = router(torch.tensor([[1.0, 0.0]]))
|
||||
|
||||
assert set(top_indices[0].tolist()) == {2, 3}
|
||||
assert torch.allclose(top_scores.sum(dim=-1), torch.tensor([2.5]))
|
||||
low_score, high_score = sorted(top_scores[0].tolist())
|
||||
assert low_score < 0.8
|
||||
assert high_score > 1.7
|
||||
|
||||
|
||||
def test_sparse_moe_block_masks_padding_tokens():
|
||||
from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module
|
||||
|
||||
block = module.LingBotVideoSparseMoeBlock(
|
||||
hidden_size=4,
|
||||
num_experts=2,
|
||||
top_k=1,
|
||||
moe_intermediate_size=3,
|
||||
score_func="sigmoid",
|
||||
norm_topk_prob=True,
|
||||
n_group=None,
|
||||
topk_group=None,
|
||||
routed_scaling_factor=1.0,
|
||||
n_shared_experts=None,
|
||||
)
|
||||
with torch.no_grad():
|
||||
block.router.weight.copy_(
|
||||
torch.tensor(
|
||||
[
|
||||
[8.0, 0.0, 0.0, 0.0],
|
||||
[-8.0, 0.0, 0.0, 0.0],
|
||||
]
|
||||
)
|
||||
)
|
||||
block.router.e_score_correction_bias.zero_()
|
||||
block.experts.w1.fill_(0.5)
|
||||
block.experts.w2.fill_(0.5)
|
||||
block.experts.w3.fill_(0.5)
|
||||
|
||||
hidden_states = torch.tensor([[[1.0, 0.0, 0.0, 0.0], [-1.0, 0.0, 0.0, 0.0]]])
|
||||
padding_mask = torch.tensor([1.0, 0.0])
|
||||
|
||||
out = block(hidden_states, padding_mask=padding_mask)
|
||||
|
||||
assert out.shape == hidden_states.shape
|
||||
assert torch.isfinite(out).all()
|
||||
assert not torch.allclose(out[0, 0], torch.zeros_like(out[0, 0]))
|
||||
assert torch.allclose(out[0, 1], torch.zeros_like(out[0, 1]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"counts_list",
|
||||
[
|
||||
[0, 0, 0, 0],
|
||||
[1, 0, 7, 8, 9],
|
||||
[16, 0, 1, 0],
|
||||
[0, 0, 33, 0],
|
||||
],
|
||||
)
|
||||
def test_grouped_padding_matches_reference(counts_list):
|
||||
from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module
|
||||
|
||||
align = 8
|
||||
counts = torch.tensor(counts_list, dtype=torch.int64)
|
||||
num_tokens = int(counts.sum())
|
||||
tokens = torch.arange(num_tokens * 4, dtype=torch.float32).reshape(num_tokens, 4)
|
||||
|
||||
actual = module.LingBotVideoSparseMoeBlock._pad_grouped_tokens(
|
||||
tokens,
|
||||
counts,
|
||||
align,
|
||||
)
|
||||
|
||||
num_experts = len(counts_list)
|
||||
max_len = ((num_tokens + num_experts * align + align - 1) // align) * align
|
||||
aligned_counts = [((max(count, align) + align - 1) // align) * align for count in counts_list]
|
||||
expected_indices = [num_tokens] * max_len
|
||||
source_start = 0
|
||||
write_start = 0
|
||||
for count, aligned_count in zip(counts_list, aligned_counts):
|
||||
expected_indices[write_start : write_start + count] = range(
|
||||
source_start,
|
||||
source_start + count,
|
||||
)
|
||||
source_start += count
|
||||
write_start += aligned_count
|
||||
|
||||
expected_indices_tensor = torch.tensor(expected_indices, dtype=torch.int64)
|
||||
tokens_with_pad = torch.vstack((tokens, tokens.new_zeros((tokens.shape[-1],))))
|
||||
expected_aligned_counts = torch.tensor(aligned_counts, dtype=torch.int32)
|
||||
|
||||
assert actual[0] == tokens_with_pad.shape
|
||||
torch.testing.assert_close(actual[1], tokens_with_pad[expected_indices_tensor])
|
||||
torch.testing.assert_close(actual[2], expected_indices_tensor)
|
||||
torch.testing.assert_close(actual[3], expected_aligned_counts)
|
||||
|
||||
|
||||
def test_tiny_transformer_constructs_moe_and_dense_layers():
|
||||
from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module
|
||||
|
||||
model = _tiny_transformer(
|
||||
depth=2,
|
||||
num_experts=4,
|
||||
num_experts_per_tok=2,
|
||||
moe_intermediate_size=8,
|
||||
decoder_sparse_step=1,
|
||||
mlp_only_layers=(1,),
|
||||
n_shared_experts=1,
|
||||
n_group=2,
|
||||
topk_group=1,
|
||||
routed_scaling_factor=2.5,
|
||||
)
|
||||
|
||||
assert isinstance(model.blocks[0].ffn, module.LingBotVideoSparseMoeBlock)
|
||||
assert isinstance(model.blocks[1].ffn, module.LingBotVideoMLP)
|
||||
assert "blocks.0.ffn.experts.w1" in model.state_dict()
|
||||
assert "blocks.0.ffn.shared_experts.gate_proj.weight" in model.state_dict()
|
||||
|
||||
model.to(dtype=torch.bfloat16)
|
||||
|
||||
assert model.blocks[0].ffn.router.weight.dtype == torch.float32
|
||||
assert model.blocks[0].ffn.experts.w1.dtype == torch.bfloat16
|
||||
@@ -0,0 +1,219 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _make_pipeline():
|
||||
from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import LingBotVideoPipeline
|
||||
|
||||
pipeline = object.__new__(LingBotVideoPipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.default_negative_prompt = "default negative"
|
||||
pipeline.od_config = SimpleNamespace(flow_shift=None)
|
||||
return pipeline
|
||||
|
||||
|
||||
def _make_request_batch(prompt, **sampling_overrides):
|
||||
sampling = OmniDiffusionSamplingParams(**sampling_overrides)
|
||||
request = OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id="lingbot-test")
|
||||
return DiffusionRequestBatch([request])
|
||||
|
||||
|
||||
def test_lingbot_video_pipeline_import_and_registry():
|
||||
from vllm_omni.diffusion.models.lingbot_video import (
|
||||
LingBotVideoPipeline,
|
||||
LingBotVideoTransformer3DModel,
|
||||
get_lingbot_video_post_process_func,
|
||||
)
|
||||
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS
|
||||
|
||||
assert LingBotVideoPipeline is not None
|
||||
assert LingBotVideoTransformer3DModel is not None
|
||||
assert get_lingbot_video_post_process_func is not None
|
||||
assert _DIFFUSION_MODELS["LingBotVideoPipeline"] == (
|
||||
"lingbot_video",
|
||||
"pipeline_lingbot_video",
|
||||
"LingBotVideoPipeline",
|
||||
)
|
||||
assert _DIFFUSION_POST_PROCESS_FUNCS["LingBotVideoPipeline"] == "get_lingbot_video_post_process_func"
|
||||
|
||||
|
||||
def test_component_discovery_declarations():
|
||||
from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline
|
||||
|
||||
assert LingBotVideoPipeline._dit_modules == ["transformer"]
|
||||
assert LingBotVideoPipeline._encoder_modules == ["text_encoder"]
|
||||
assert LingBotVideoPipeline._vae_modules == ["vae"]
|
||||
assert LingBotVideoPipeline.supports_step_execution is False
|
||||
|
||||
|
||||
def test_extra_body_params_include_video_flow_shift_alias():
|
||||
from vllm_omni.model_extras import get_extra_body_params
|
||||
|
||||
params = get_extra_body_params("LingBotVideoPipeline")
|
||||
assert "shift" in params
|
||||
assert "flow_shift" in params
|
||||
assert "batch_cfg" in params
|
||||
|
||||
|
||||
def test_extract_prompt_string_and_mapping():
|
||||
from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import _extract_prompt
|
||||
|
||||
string_req = SimpleNamespace(prompt="a robot waves")
|
||||
assert _extract_prompt(string_req) == ("a robot waves", None)
|
||||
|
||||
mapping_req = SimpleNamespace(prompt={"prompt": "a robot walks", "negative_prompt": "low quality"})
|
||||
assert _extract_prompt(mapping_req) == ("a robot walks", "low quality")
|
||||
|
||||
with pytest.raises(TypeError, match="Unsupported LingBot prompt type"):
|
||||
_extract_prompt(SimpleNamespace(prompt=["not", "supported"]))
|
||||
|
||||
|
||||
def test_check_inputs_accepts_t2v_shapes_and_rejects_invalid_values():
|
||||
from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline
|
||||
|
||||
LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=1)
|
||||
LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=9)
|
||||
|
||||
with pytest.raises(ValueError, match="num_frames"):
|
||||
LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=8)
|
||||
with pytest.raises(ValueError, match="height"):
|
||||
LingBotVideoPipeline.check_inputs(height=190, width=320, num_frames=9)
|
||||
|
||||
|
||||
def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
|
||||
from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
|
||||
|
||||
frames = torch.ones(2, 4, 4, 3)
|
||||
post = get_lingbot_video_post_process_func(SimpleNamespace())
|
||||
|
||||
assert post(frames, SimpleNamespace(output_type="pt")) is frames
|
||||
array = post(frames, SimpleNamespace(output_type="np"))
|
||||
assert isinstance(array, np.ndarray)
|
||||
assert array.shape == (2, 4, 4, 3)
|
||||
|
||||
|
||||
def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
|
||||
pipeline = _make_pipeline()
|
||||
calls = []
|
||||
|
||||
def fake_generate(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return torch.zeros(9, 192, 320, 3)
|
||||
|
||||
pipeline._generate = fake_generate
|
||||
req = _make_request_batch(
|
||||
{"prompt": "a robot arm", "negative_prompt": "bad hands"},
|
||||
height=192,
|
||||
width=320,
|
||||
num_frames=9,
|
||||
num_inference_steps=2,
|
||||
guidance_scale=3.0,
|
||||
seed=42,
|
||||
output_type="pt",
|
||||
extra_args={"flow_shift": 4.0, "batch_cfg": True},
|
||||
)
|
||||
|
||||
out = pipeline.forward(req)
|
||||
|
||||
assert torch.equal(out.output, torch.zeros(9, 192, 320, 3))
|
||||
assert len(calls) == 1
|
||||
call = calls[0]
|
||||
assert call["prompt"] == "a robot arm"
|
||||
assert call["negative_prompt"] == "bad hands"
|
||||
assert call["height"] == 192
|
||||
assert call["width"] == 320
|
||||
assert call["num_frames"] == 9
|
||||
assert call["num_inference_steps"] == 2
|
||||
assert call["guidance_scale"] == 3.0
|
||||
assert call["shift"] == 4.0
|
||||
assert call["batch_cfg"] is True
|
||||
assert call["output_type"] == "pt"
|
||||
|
||||
|
||||
def test_forward_uses_lingbot_guidance_default_when_omitted():
|
||||
pipeline = _make_pipeline()
|
||||
calls = []
|
||||
|
||||
def fake_generate(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return torch.zeros(1)
|
||||
|
||||
pipeline._generate = fake_generate
|
||||
req = _make_request_batch(
|
||||
"a robot arm",
|
||||
height=192,
|
||||
width=320,
|
||||
num_frames=81,
|
||||
num_inference_steps=2,
|
||||
seed=42,
|
||||
)
|
||||
|
||||
pipeline.forward(req)
|
||||
|
||||
assert calls[0]["num_frames"] == 81
|
||||
assert calls[0]["guidance_scale"] == 6.0
|
||||
|
||||
|
||||
def test_forward_prefers_shift_over_flow_shift_and_defaults_negative_prompt():
|
||||
pipeline = _make_pipeline()
|
||||
calls = []
|
||||
|
||||
def fake_generate(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return torch.zeros(5, 192, 192, 3)
|
||||
|
||||
pipeline._generate = fake_generate
|
||||
req = _make_request_batch(
|
||||
"a robot arm",
|
||||
height=192,
|
||||
width=192,
|
||||
num_frames=5,
|
||||
num_inference_steps=2,
|
||||
guidance_scale=1.0,
|
||||
seed=42,
|
||||
extra_args={"shift": 5.0, "flow_shift": 4.0},
|
||||
)
|
||||
|
||||
pipeline.forward(req)
|
||||
|
||||
assert calls[0]["shift"] == 5.0
|
||||
assert calls[0]["negative_prompt"] == "default negative"
|
||||
|
||||
|
||||
def test_forward_rejects_multi_request_batches():
|
||||
pipeline = _make_pipeline()
|
||||
first = OmniDiffusionRequest(
|
||||
prompt="first",
|
||||
sampling_params=OmniDiffusionSamplingParams(seed=1),
|
||||
request_id="lingbot-first",
|
||||
)
|
||||
second = OmniDiffusionRequest(
|
||||
prompt="second",
|
||||
sampling_params=OmniDiffusionSamplingParams(seed=2),
|
||||
request_id="lingbot-second",
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="only supports one request"):
|
||||
pipeline.forward(DiffusionRequestBatch([first, second]))
|
||||
|
||||
|
||||
def test_load_weights_rejects_external_weight_stream():
|
||||
pipeline = _make_pipeline()
|
||||
|
||||
assert pipeline.load_weights([]) == set()
|
||||
with pytest.raises(RuntimeError, match="components are loaded directly"):
|
||||
pipeline.load_weights([("transformer.weight", torch.zeros(1))])
|
||||
Reference in new issue
Block a user