188 lines
6.2 KiB
Python
188 lines
6.2 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.models.ernie_image.ernie_image_transformer import (
|
|
ErnieImageEmbedND3,
|
|
_apply_rotary_emb,
|
|
rope,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class TestErnieImageRopePositionEmbedding:
|
|
"""Test ErnieImage RoPE position embedding functionality"""
|
|
|
|
def test_rope_output_shape(self):
|
|
"""Verify rope produces correct output shape"""
|
|
dim = 64
|
|
theta = 256
|
|
pos = torch.tensor([1.0, 2.0, 3.0])
|
|
|
|
out = rope(pos, dim, theta)
|
|
|
|
assert out.shape == (3, dim // 2), f"Expected shape (3, {dim // 2}), got {out.shape}"
|
|
assert out.dtype == torch.float32
|
|
|
|
def test_rope_trigonometric_relationship(self):
|
|
"""Verify rope output satisfies cos^2 + sin^2 = 1"""
|
|
dim = 64
|
|
theta = 256
|
|
pos = torch.tensor([1.0, 2.0, 3.0])
|
|
|
|
out = rope(pos, dim, theta)
|
|
cos_out = out.cos()
|
|
sin_out = out.sin()
|
|
|
|
cos_sq_sin_sq = cos_out**2 + sin_out**2
|
|
assert torch.allclose(cos_sq_sin_sq, torch.ones_like(cos_sq_sin_sq), atol=1e-6)
|
|
|
|
def test_embed_nd3_output_shape(self):
|
|
"""Verify ErnieImageEmbedND3 produces correct output shape"""
|
|
head_dim = 128
|
|
axes_dim = (32, 48, 48)
|
|
theta = 256
|
|
B, S = 2, 100
|
|
|
|
pos_embed = ErnieImageEmbedND3(dim=head_dim, theta=theta, axes_dim=axes_dim)
|
|
ids = torch.randn(B, S, 3)
|
|
|
|
freqs_cos, freqs_sin = pos_embed(ids)
|
|
|
|
expected_dim = sum(axes_dim) // 2
|
|
assert freqs_cos.shape == (B, S, expected_dim), (
|
|
f"Expected shape ({B}, {S}, {expected_dim}), got {freqs_cos.shape}"
|
|
)
|
|
assert freqs_sin.shape == (B, S, expected_dim)
|
|
|
|
def test_embed_nd3_different_batch_elements(self):
|
|
"""Verify different batch elements produce different embeddings"""
|
|
head_dim = 128
|
|
axes_dim = (32, 48, 48)
|
|
theta = 256
|
|
B, S = 2, 100
|
|
|
|
pos_embed = ErnieImageEmbedND3(dim=head_dim, theta=theta, axes_dim=axes_dim)
|
|
|
|
ids = torch.zeros(B, S, 3)
|
|
ids[0, :, 0] = 1.0
|
|
ids[1, :, 0] = 2.0
|
|
|
|
freqs_cos, freqs_sin = pos_embed(ids)
|
|
|
|
assert not torch.allclose(freqs_cos[0], freqs_cos[1]), (
|
|
"Different batch elements should produce different cos embeddings"
|
|
)
|
|
assert not torch.allclose(freqs_sin[0], freqs_sin[1]), (
|
|
"Different batch elements should produce different sin embeddings"
|
|
)
|
|
|
|
def test_apply_rotary_emb_output_shape(self):
|
|
"""Verify _apply_rotary_emb preserves input shape"""
|
|
B, S, H, D = 2, 100, 32, 128
|
|
|
|
x = torch.randn(B, S, H, D)
|
|
freqs_cos = torch.randn(B, S, D // 2)
|
|
freqs_sin = torch.randn(B, S, D // 2)
|
|
|
|
out = _apply_rotary_emb(x, freqs_cos, freqs_sin)
|
|
|
|
assert out.shape == x.shape, f"Expected shape {x.shape}, got {out.shape}"
|
|
|
|
def test_apply_rotary_emb_no_nan_inf(self):
|
|
"""Verify _apply_rotary_emb produces valid values"""
|
|
B, S, H, D = 2, 100, 32, 128
|
|
|
|
x = torch.randn(B, S, H, D)
|
|
freqs_cos = torch.randn(B, S, D // 2)
|
|
freqs_sin = torch.randn(B, S, D // 2)
|
|
|
|
out = _apply_rotary_emb(x, freqs_cos, freqs_sin)
|
|
|
|
assert not torch.isnan(out).any(), "Output should not contain NaN"
|
|
assert not torch.isinf(out).any(), "Output should not contain Inf"
|
|
|
|
def test_apply_rotary_emb_preserves_batch_dimension(self):
|
|
"""Verify _apply_rotary_emb preserves batch dimension for CFG mode"""
|
|
B, S, H, D = 2, 100, 32, 128
|
|
|
|
x = torch.randn(B, S, H, D)
|
|
|
|
freqs_cos = torch.zeros(B, S, D // 2)
|
|
freqs_sin = torch.zeros(B, S, D // 2)
|
|
freqs_cos[0] = 1.0
|
|
freqs_cos[1] = 0.5
|
|
|
|
out = _apply_rotary_emb(x, freqs_cos, freqs_sin)
|
|
|
|
assert not torch.allclose(out[0], out[1]), "Different batch elements should produce different outputs"
|
|
|
|
def test_cfg_mode_text_lens_difference(self):
|
|
"""Verify CFG mode with different text_lens produces different positional embeddings"""
|
|
head_dim = 128
|
|
axes_dim = (32, 48, 48)
|
|
theta = 256
|
|
B = 2
|
|
N_img = 64
|
|
Tmax = 50
|
|
|
|
text_lens = torch.tensor([2, 50])
|
|
|
|
text_ids = torch.cat(
|
|
[
|
|
torch.arange(Tmax, dtype=torch.float32).view(1, Tmax, 1).expand(B, -1, -1),
|
|
torch.zeros((B, Tmax, 2)),
|
|
],
|
|
dim=-1,
|
|
)
|
|
|
|
grid_yx = torch.stack(
|
|
torch.meshgrid(
|
|
torch.arange(8, dtype=torch.float32),
|
|
torch.arange(8, dtype=torch.float32),
|
|
indexing="ij",
|
|
),
|
|
dim=-1,
|
|
).reshape(-1, 2)
|
|
|
|
image_ids = torch.cat(
|
|
[
|
|
text_lens.float().view(B, 1, 1).expand(-1, N_img, -1),
|
|
grid_yx.view(1, N_img, 2).expand(B, -1, -1),
|
|
],
|
|
dim=-1,
|
|
)
|
|
|
|
ids = torch.cat([image_ids, text_ids], dim=1)
|
|
|
|
pos_embed = ErnieImageEmbedND3(dim=head_dim, theta=theta, axes_dim=axes_dim)
|
|
freqs_cos, freqs_sin = pos_embed(ids)
|
|
|
|
assert not torch.allclose(freqs_cos[0], freqs_cos[1]), (
|
|
"CFG mode: uncond and cond should have different positional embeddings"
|
|
)
|
|
assert not torch.allclose(freqs_sin[0], freqs_sin[1]), (
|
|
"CFG mode: uncond and cond should have different positional embeddings"
|
|
)
|
|
|
|
def test_apply_rotary_emb_matches_diffusers_style(self):
|
|
"""Verify _apply_rotary_emb matches diffusers implementation style"""
|
|
B, S, H, D = 2, 100, 32, 128
|
|
|
|
x = torch.randn(B, S, H, D)
|
|
freqs = torch.randn(B, S, D // 2)
|
|
freqs_cos = freqs.cos()
|
|
freqs_sin = freqs.sin()
|
|
|
|
out_vllm = _apply_rotary_emb(x, freqs_cos, freqs_sin)
|
|
|
|
cos_ = freqs_cos.unsqueeze(2)
|
|
sin_ = freqs_sin.unsqueeze(2)
|
|
cos_ = torch.stack([cos_, cos_], dim=-1).reshape(*cos_.shape[:-1], -1)
|
|
sin_ = torch.stack([sin_, sin_], dim=-1).reshape(*sin_.shape[:-1], -1)
|
|
x1, x2 = x.chunk(2, dim=-1)
|
|
rotated = torch.cat((-x2, x1), dim=-1)
|
|
out_diffusers = x * cos_ + rotated * sin_
|
|
|
|
assert torch.allclose(out_vllm, out_diffusers, atol=1e-6), "vllm-omni and diffusers outputs should match"
|