06e733d3cd
Signed-off-by: weiyan <weiyanlin117@gmail.com> Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
267 lines
8.1 KiB
Python
267 lines
8.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from vllm_omni.diffusion.models.longcat_video.longcat_video_avatar_transformer import (
|
|
LongCatVideoAvatarTransformer3DModel,
|
|
_read_config,
|
|
replace_linear_with_quantized,
|
|
)
|
|
from vllm_omni.diffusion.models.longcat_video.pipeline_longcat_video_avatar import (
|
|
_avatar_model_allow_patterns,
|
|
_build_multi_speaker_ref_target_masks,
|
|
_default_at2v_shape,
|
|
_infer_asset_root_from_path,
|
|
_prepare_multi_speaker_audio_arrays,
|
|
_resolve_num_segments,
|
|
prepare_longcat_video_avatar_model_for_omni,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("resolution", "expected_shape"),
|
|
[
|
|
("480p", (480, 832)),
|
|
("720p", (736, 1248)),
|
|
],
|
|
)
|
|
def test_longcat_video_avatar_at2v_default_shape_uses_resolution_bucket(
|
|
resolution: str,
|
|
expected_shape: tuple[int, int],
|
|
):
|
|
assert _default_at2v_shape(resolution) == expected_shape
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("use_int8", "expected_weight_dir", "unexpected_weight_dir"),
|
|
[
|
|
(True, "base_model_int8/*", "base_model/*"),
|
|
(False, "base_model/*", "base_model_int8/*"),
|
|
],
|
|
)
|
|
def test_longcat_video_avatar_allow_patterns_download_one_weight_set(
|
|
use_int8: bool,
|
|
expected_weight_dir: str,
|
|
unexpected_weight_dir: str,
|
|
):
|
|
allow_patterns = _avatar_model_allow_patterns(use_int8)
|
|
|
|
assert expected_weight_dir in allow_patterns
|
|
assert unexpected_weight_dir not in allow_patterns
|
|
assert "whisper-large-v3/*" not in allow_patterns
|
|
assert "whisper-large-v3/model.safetensors" in allow_patterns
|
|
assert "whisper-large-v3/pytorch_model.bin" not in allow_patterns
|
|
assert "whisper-large-v3/flax_model.msgpack" not in allow_patterns
|
|
assert "vocal_separator/*" in allow_patterns
|
|
|
|
|
|
def test_longcat_video_avatar_prepare_model_adds_omni_metadata(tmp_path):
|
|
model_dir = tmp_path / "LongCat-Video-Avatar-1.5"
|
|
base_model_dir = model_dir / "base_model_int8"
|
|
base_model_dir.mkdir(parents=True)
|
|
(base_model_dir / "config.json").write_text('{"model_type": "longcat_avatar"}', encoding="utf-8")
|
|
(model_dir / "model_index.json").write_text("{}", encoding="utf-8")
|
|
|
|
prepared = prepare_longcat_video_avatar_model_for_omni(str(model_dir), use_int8=True)
|
|
|
|
assert prepared == str(model_dir)
|
|
model_index = json.loads((model_dir / "model_index.json").read_text(encoding="utf-8"))
|
|
assert model_index["_class_name"] == "LongCatVideoAvatarPipeline"
|
|
assert model_index["_diffusers_version"] == "0.38.0"
|
|
transformer_config = json.loads((model_dir / "transformer" / "config.json").read_text(encoding="utf-8"))
|
|
assert transformer_config == {"model_type": "longcat_avatar"}
|
|
|
|
|
|
def test_longcat_video_avatar_infers_repo_root_from_official_asset_path(tmp_path):
|
|
input_json = tmp_path / "assets" / "avatar" / "single_example_1.json"
|
|
input_json.parent.mkdir(parents=True)
|
|
input_json.write_text("{}", encoding="utf-8")
|
|
|
|
assert _infer_asset_root_from_path(input_json) == tmp_path
|
|
|
|
|
|
def test_longcat_video_avatar_falls_back_to_json_parent_without_official_layout(tmp_path):
|
|
input_json = tmp_path / "single_example_1.json"
|
|
input_json.write_text("{}", encoding="utf-8")
|
|
|
|
assert _infer_asset_root_from_path(input_json) == tmp_path
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("value", "audio_duration", "expected"),
|
|
[
|
|
(None, 26.6, 1),
|
|
("1", 26.6, 1),
|
|
(3, 26.6, 3),
|
|
("auto", 3.0, 1),
|
|
("auto", 23.8, 8),
|
|
("auto", 26.6, 9),
|
|
],
|
|
)
|
|
def test_longcat_video_avatar_resolve_num_segments(
|
|
value,
|
|
audio_duration: float,
|
|
expected: int,
|
|
):
|
|
assert (
|
|
_resolve_num_segments(
|
|
value,
|
|
audio_duration=audio_duration,
|
|
num_frames=93,
|
|
num_cond_frames=13,
|
|
save_fps=25,
|
|
)
|
|
== expected
|
|
)
|
|
|
|
|
|
def test_longcat_video_avatar_read_config_filters_non_constructor_metadata(tmp_path):
|
|
config_path = tmp_path / "config.json"
|
|
config_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"_class_name": "LongCatVideoAvatarTransformer3DModel",
|
|
"architectures": ["LongCatVideoAvatarTransformer3DModel"],
|
|
"_diffusers_version": "0.35.1",
|
|
"model_max_length": 512,
|
|
"hidden_size": 4096,
|
|
"depth": 48,
|
|
"num_heads": 32,
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
config = _read_config(config_path)
|
|
|
|
assert config == {
|
|
"hidden_size": 4096,
|
|
"depth": 48,
|
|
"num_heads": 32,
|
|
}
|
|
|
|
|
|
def test_longcat_video_avatar_load_weights_supports_int8_buffers():
|
|
model = LongCatVideoAvatarTransformer3DModel(
|
|
hidden_size=4,
|
|
depth=0,
|
|
num_heads=1,
|
|
caption_channels=4,
|
|
intermediate_dim=4,
|
|
output_dim=4,
|
|
audio_channel=4,
|
|
context_tokens=1,
|
|
)
|
|
replace_linear_with_quantized(model)
|
|
|
|
buffer_name = "t_embedder.mlp.0.weight_int8"
|
|
buffers = dict(model.named_buffers())
|
|
loaded_weight = torch.ones_like(buffers[buffer_name])
|
|
|
|
loaded_params = model.load_weights([(buffer_name, loaded_weight)])
|
|
|
|
assert buffer_name in loaded_params
|
|
assert torch.equal(dict(model.named_buffers())[buffer_name], loaded_weight)
|
|
|
|
|
|
def test_longcat_video_avatar_multi_speaker_para_audio_arrays_are_aligned():
|
|
left = np.ones(3, dtype=np.float32)
|
|
right = np.ones(5, dtype=np.float32) * 2
|
|
|
|
left_out, right_out = _prepare_multi_speaker_audio_arrays(
|
|
[left, right],
|
|
audio_type="para",
|
|
generate_duration=0.5,
|
|
sample_rate=10,
|
|
)
|
|
|
|
assert left_out.tolist() == [1, 1, 1, 0, 0]
|
|
assert right_out.tolist() == [2, 2, 2, 2, 2]
|
|
|
|
|
|
def test_longcat_video_avatar_multi_speaker_add_audio_arrays_are_sequential():
|
|
left = np.ones(3, dtype=np.float32)
|
|
right = np.ones(2, dtype=np.float32) * 2
|
|
|
|
left_out, right_out = _prepare_multi_speaker_audio_arrays(
|
|
[left, right],
|
|
audio_type="add",
|
|
generate_duration=0.1,
|
|
sample_rate=10,
|
|
)
|
|
|
|
assert left_out.tolist() == [1, 1, 1, 0, 0]
|
|
assert right_out.tolist() == [0, 0, 0, 2, 2]
|
|
|
|
|
|
def test_longcat_video_avatar_multi_speaker_masks_use_explicit_bboxes():
|
|
image = Image.new("RGB", (8, 6))
|
|
masks, use_background = _build_multi_speaker_ref_target_masks(
|
|
image,
|
|
{
|
|
"person1": [1, 1, 4, 3],
|
|
"person2": [2, 4, 5, 7],
|
|
},
|
|
)
|
|
|
|
assert not use_background
|
|
assert masks.shape == (3, 6, 8)
|
|
assert masks[0, 1:4, 1:3].sum() == 6
|
|
assert masks[1, 2:5, 4:7].sum() == 9
|
|
assert masks[2, 0, 0] == 1
|
|
assert masks[2, 2, 2] == 0
|
|
|
|
|
|
def test_longcat_video_avatar_multi_speaker_masks_default_to_left_right_split():
|
|
image = Image.new("RGB", (10, 8))
|
|
masks, use_background = _build_multi_speaker_ref_target_masks(image, None)
|
|
|
|
assert not use_background
|
|
assert masks.shape == (3, 8, 10)
|
|
assert masks[0, :, :5].sum() > 0
|
|
assert masks[0, :, 5:].sum() == 0
|
|
assert masks[1, :, :5].sum() == 0
|
|
assert masks[1, :, 5:].sum() > 0
|
|
|
|
|
|
def test_longcat_video_avatar_transformer_accepts_multi_speaker_masks():
|
|
model = LongCatVideoAvatarTransformer3DModel(
|
|
hidden_size=4,
|
|
depth=0,
|
|
num_heads=1,
|
|
caption_channels=4,
|
|
intermediate_dim=4,
|
|
output_dim=4,
|
|
audio_channel=4,
|
|
context_tokens=1,
|
|
)
|
|
hidden_states = torch.randn(1, 16, 5, 4, 4)
|
|
timestep = torch.zeros(1, 5)
|
|
encoder_hidden_states = torch.randn(1, 1, 4, 4)
|
|
audio_embs = torch.randn(2, 5, 5, 12, 4)
|
|
ref_target_masks = torch.zeros(3, 4, 4)
|
|
ref_target_masks[0, :, :2] = 1
|
|
ref_target_masks[1, :, 2:] = 1
|
|
ref_target_masks[2] = torch.where(ref_target_masks[0] + ref_target_masks[1] > 0, 0, 1)
|
|
|
|
output = model(
|
|
hidden_states=hidden_states,
|
|
timestep=timestep,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
num_cond_latents=1,
|
|
audio_embs=audio_embs,
|
|
ref_target_masks=ref_target_masks,
|
|
)
|
|
|
|
assert output.shape == hidden_states.shape
|