wehub-resource-sync
219 lines
7.4 KiB
Python
219 lines
7.4 KiB
Python
# SPDX-License-Identifier: MIT
|
|
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
|
|
BACKEND_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.insert(0, BACKEND_DIR)
|
|
|
|
# longcat-video is a backend directory, not an importable Python package name.
|
|
from longcat_utils import ( # noqa: E402
|
|
MODEL_KIND_AVATAR,
|
|
MODEL_KIND_BASE,
|
|
attention_overrides,
|
|
avatar_segments_for_duration,
|
|
avatar_segments_for_frames,
|
|
classify_model,
|
|
normalize_model_source,
|
|
normalize_num_frames,
|
|
parse_options,
|
|
validate_dimensions,
|
|
)
|
|
|
|
|
|
SOURCE_DIR = os.path.join(BACKEND_DIR, "sources", "LongCat-Video")
|
|
try:
|
|
import torch
|
|
|
|
sys.path.insert(0, SOURCE_DIR)
|
|
ATTENTION_TESTS_AVAILABLE = (
|
|
os.path.isdir(SOURCE_DIR) and importlib.util.find_spec("triton") is not None
|
|
)
|
|
except ImportError:
|
|
torch = None
|
|
ATTENTION_TESTS_AVAILABLE = False
|
|
|
|
AVATAR_ATTENTION_TESTS_AVAILABLE = ATTENTION_TESTS_AVAILABLE and all(
|
|
importlib.util.find_spec(module) is not None
|
|
for module in ("pyloudnorm", "scipy", "torchvision")
|
|
)
|
|
|
|
|
|
class LongCatUtilsTest(unittest.TestCase):
|
|
def test_parse_options_preserves_colons_and_coerces_scalars(self):
|
|
options = parse_options(
|
|
[
|
|
"use_distill:true",
|
|
"max_segments:4",
|
|
"audio_guidance_scale:3.5",
|
|
"source:https://example.com/model",
|
|
"flag",
|
|
]
|
|
)
|
|
|
|
self.assertEqual(options["use_distill"], True)
|
|
self.assertEqual(options["max_segments"], 4)
|
|
self.assertEqual(options["audio_guidance_scale"], 3.5)
|
|
self.assertEqual(options["source"], "https://example.com/model")
|
|
self.assertEqual(options["flag"], True)
|
|
|
|
def test_classify_model_accepts_only_supported_longcat_models(self):
|
|
cases = {
|
|
"meituan-longcat/LongCat-Video": MODEL_KIND_BASE,
|
|
"https://huggingface.co/meituan-longcat/LongCat-Video": MODEL_KIND_BASE,
|
|
"hf://meituan-longcat/LongCat-Video-Avatar-1.5": MODEL_KIND_AVATAR,
|
|
"other-org/LongCat-Video": None,
|
|
"meituan-longcat/LongCat-Video-Avatar": None,
|
|
"some-org/unrelated-model": None,
|
|
}
|
|
|
|
for model, expected in cases.items():
|
|
with self.subTest(model=model):
|
|
self.assertEqual(classify_model(model), expected)
|
|
|
|
def test_classify_model_reads_local_checkpoint_metadata(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
with open(
|
|
os.path.join(directory, "model_index.json"),
|
|
"w",
|
|
encoding="utf-8",
|
|
) as config_file:
|
|
json.dump({"model_name": "LongCat-Video-Avatar-1.5"}, config_file)
|
|
|
|
self.assertEqual(classify_model(directory), MODEL_KIND_AVATAR)
|
|
|
|
def test_normalize_model_source_handles_huggingface_uri_forms(self):
|
|
self.assertEqual(
|
|
normalize_model_source(
|
|
"https://huggingface.co/meituan-longcat/LongCat-Video/tree/main"
|
|
),
|
|
"meituan-longcat/LongCat-Video",
|
|
)
|
|
self.assertEqual(
|
|
normalize_model_source("huggingface://meituan-longcat/LongCat-Video"),
|
|
"meituan-longcat/LongCat-Video",
|
|
)
|
|
|
|
def test_frame_and_segment_rounding_matches_longcat_temporal_shape(self):
|
|
self.assertEqual(normalize_num_frames(94), 93)
|
|
self.assertEqual(normalize_num_frames(0), 93)
|
|
self.assertEqual(avatar_segments_for_frames(93), 1)
|
|
self.assertEqual(avatar_segments_for_frames(94), 2)
|
|
self.assertEqual(avatar_segments_for_frames(173), 2)
|
|
self.assertEqual(avatar_segments_for_frames(174), 3)
|
|
self.assertEqual(avatar_segments_for_duration(10.0), 3)
|
|
|
|
def test_dimensions_are_bounded_and_aligned(self):
|
|
self.assertEqual(validate_dimensions(0, 0), (832, 480))
|
|
self.assertEqual(validate_dimensions(512, 512), (512, 512))
|
|
with self.assertRaisesRegex(ValueError, "divisible by 16"):
|
|
validate_dimensions(513, 512)
|
|
with self.assertRaisesRegex(ValueError, "must not exceed"):
|
|
validate_dimensions(1920, 1080)
|
|
|
|
def test_attention_backend_validation(self):
|
|
self.assertEqual(
|
|
attention_overrides("sdpa"),
|
|
{
|
|
"enable_flashattn2": False,
|
|
"enable_flashattn3": False,
|
|
"enable_xformers": False,
|
|
},
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "attention_backend"):
|
|
attention_overrides("unknown")
|
|
|
|
|
|
@unittest.skipUnless(
|
|
ATTENTION_TESTS_AVAILABLE,
|
|
"patched LongCat source and torch are required for attention tests",
|
|
)
|
|
class SDPAFallbackTest(unittest.TestCase):
|
|
def test_base_self_attention_matches_reference(self):
|
|
from longcat_video.modules.attention import Attention
|
|
|
|
dim, heads, sequence = 64, 4, 32
|
|
attention = Attention(
|
|
dim,
|
|
heads,
|
|
enable_flashattn2=False,
|
|
enable_flashattn3=False,
|
|
enable_xformers=False,
|
|
enable_bsa=False,
|
|
).float()
|
|
query = torch.randn(2, heads, sequence, dim // heads)
|
|
key = torch.randn_like(query)
|
|
value = torch.randn_like(query)
|
|
|
|
output = attention._process_attn(query, key, value, shape=(1, 1, sequence))
|
|
reference = (
|
|
torch.softmax(
|
|
(query @ key.transpose(-1, -2)) * attention.scale,
|
|
dim=-1,
|
|
)
|
|
@ value
|
|
)
|
|
|
|
self.assertLess((output - reference).abs().max().item(), 1e-4)
|
|
|
|
@unittest.skipUnless(
|
|
AVATAR_ATTENTION_TESTS_AVAILABLE,
|
|
"avatar audio dependencies are required for the avatar attention test",
|
|
)
|
|
def test_avatar_self_attention_matches_reference(self):
|
|
from longcat_video.modules.avatar.attention import Attention
|
|
|
|
dim, heads, sequence = 64, 4, 16
|
|
attention = Attention(
|
|
dim,
|
|
heads,
|
|
enable_flashattn2=False,
|
|
enable_flashattn3=False,
|
|
enable_xformers=False,
|
|
).float()
|
|
query = torch.randn(1, heads, sequence, dim // heads)
|
|
key = torch.randn_like(query)
|
|
value = torch.randn_like(query)
|
|
|
|
output = attention._process_attn(query, key, value, shape=(1, 1, sequence))
|
|
reference = (
|
|
torch.softmax(
|
|
(query @ key.transpose(-1, -2)) * attention.scale,
|
|
dim=-1,
|
|
)
|
|
@ value
|
|
)
|
|
|
|
self.assertLess((output - reference).abs().max().item(), 1e-4)
|
|
|
|
def test_base_cross_attention_remains_block_diagonal(self):
|
|
from longcat_video.modules.attention import MultiHeadCrossAttention
|
|
|
|
dim, heads = 64, 4
|
|
attention = MultiHeadCrossAttention(
|
|
dim,
|
|
heads,
|
|
enable_flashattn2=False,
|
|
enable_flashattn3=False,
|
|
enable_xformers=False,
|
|
).float()
|
|
query = torch.randn(2, 8, dim)
|
|
key_lengths = [5, 7]
|
|
condition = torch.randn(1, sum(key_lengths), dim)
|
|
|
|
first = attention._process_cross_attn(query, condition, key_lengths)
|
|
changed = condition.clone()
|
|
changed[:, key_lengths[0] :] = torch.randn_like(changed[:, key_lengths[0] :])
|
|
second = attention._process_cross_attn(query, changed, key_lengths)
|
|
|
|
self.assertLess((first[0] - second[0]).abs().max().item(), 1e-5)
|
|
self.assertGreater((first[1] - second[1]).abs().max().item(), 1e-3)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|