150 lines
4.9 KiB
Python
150 lines
4.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan import (
|
|
DistributedAutoencoderKLHunyuan,
|
|
)
|
|
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan_video_15 import (
|
|
DistributedAutoencoderKLHunyuanVideo15,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class _DummyDistributedAutoencoderKLHunyuan(DistributedAutoencoderKLHunyuan):
|
|
def __init__(self):
|
|
torch.nn.Module.__init__(self)
|
|
self.tile_latent_min_size = 8
|
|
self.tile_sample_min_size = 8
|
|
self.tile_overlap_factor = 0.25
|
|
self.use_spatial_tiling = False
|
|
|
|
@property
|
|
def dtype(self) -> torch.dtype:
|
|
return torch.float32
|
|
|
|
def decoder(self, x: torch.Tensor) -> torch.Tensor:
|
|
return x + 10
|
|
|
|
def encoder(self, x: torch.Tensor) -> torch.Tensor:
|
|
return x + 20
|
|
|
|
def blend_v(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
|
|
return b
|
|
|
|
def blend_h(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
|
|
return b
|
|
|
|
|
|
class _DummyDistributedAutoencoderKLHunyuanVideo15(DistributedAutoencoderKLHunyuanVideo15):
|
|
def __init__(self):
|
|
torch.nn.Module.__init__(self)
|
|
self.tile_sample_min_height = 8
|
|
self.tile_sample_min_width = 12
|
|
self.tile_latent_min_height = 8
|
|
self.tile_latent_min_width = 12
|
|
self.tile_overlap_factor = 0.25
|
|
self.use_tiling = False
|
|
|
|
@property
|
|
def dtype(self) -> torch.dtype:
|
|
return torch.float32
|
|
|
|
def decoder(self, x: torch.Tensor) -> torch.Tensor:
|
|
return x + 10
|
|
|
|
def encoder(self, x: torch.Tensor) -> torch.Tensor:
|
|
return x + 20
|
|
|
|
def blend_v(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
|
|
return b
|
|
|
|
def blend_h(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
|
|
return b
|
|
|
|
|
|
def test_hunyuan_vae_use_tiling_aliases_spatial_tiling():
|
|
# Verify use_tiling property maps to use_spatial_tiling.
|
|
vae = _DummyDistributedAutoencoderKLHunyuan()
|
|
|
|
assert not vae.use_tiling
|
|
|
|
vae.use_tiling = True
|
|
|
|
assert vae.use_spatial_tiling
|
|
|
|
|
|
def test_hunyuan_vae_decode_tiles_round_trip():
|
|
# Validate decode tile split/exec/merge returns expected reconstructed tensor.
|
|
vae = _DummyDistributedAutoencoderKLHunyuan()
|
|
z = torch.arange(144, dtype=torch.float32).reshape(1, 1, 1, 12, 12)
|
|
|
|
tile_tasks, grid_spec = vae.tile_split(z)
|
|
decoded_tiles = {task.grid_coord: vae.tile_exec(task) for task in tile_tasks}
|
|
output = vae.tile_merge(decoded_tiles, grid_spec)
|
|
|
|
assert grid_spec.split_dims == (3, 4)
|
|
assert grid_spec.grid_shape == (2, 2)
|
|
assert grid_spec.tile_spec == {"blend_extent": 2, "row_limit": 6}
|
|
assert len(tile_tasks) == 4
|
|
assert torch.equal(output, z + 10)
|
|
|
|
|
|
def test_hunyuan_vae_encode_tiles_round_trip():
|
|
# Validate encode tile split/exec/merge returns expected latent tensor.
|
|
vae = _DummyDistributedAutoencoderKLHunyuan()
|
|
x = torch.arange(144, dtype=torch.float32).reshape(1, 1, 1, 12, 12)
|
|
|
|
tile_tasks, grid_spec = vae.encode_tile_split(x)
|
|
encoded_tiles = {task.grid_coord: vae.encode_tile_exec(task) for task in tile_tasks}
|
|
output = vae.encode_tile_merge(encoded_tiles, grid_spec)
|
|
|
|
assert grid_spec.split_dims == (3, 4)
|
|
assert grid_spec.grid_shape == (2, 2)
|
|
assert grid_spec.tile_spec == {"blend_extent": 2, "row_limit": 6}
|
|
assert len(tile_tasks) == 4
|
|
assert torch.equal(output, x + 20)
|
|
|
|
|
|
def test_hunyuan_video15_vae_decode_tiles_round_trip():
|
|
vae = _DummyDistributedAutoencoderKLHunyuanVideo15()
|
|
z = torch.arange(192, dtype=torch.float32).reshape(1, 1, 1, 12, 16)
|
|
|
|
tile_tasks, grid_spec = vae.tile_split(z)
|
|
decoded_tiles = {task.grid_coord: vae.tile_exec(task) for task in tile_tasks}
|
|
output = vae.tile_merge(decoded_tiles, grid_spec)
|
|
|
|
assert grid_spec.split_dims == (3, 4)
|
|
assert grid_spec.grid_shape == (2, 2)
|
|
assert grid_spec.tile_spec == {
|
|
"blend_height": 2,
|
|
"blend_width": 3,
|
|
"row_limit_height": 6,
|
|
"row_limit_width": 9,
|
|
}
|
|
assert len(tile_tasks) == 4
|
|
assert torch.equal(output, z + 10)
|
|
|
|
|
|
def test_hunyuan_video15_vae_encode_tiles_round_trip():
|
|
vae = _DummyDistributedAutoencoderKLHunyuanVideo15()
|
|
x = torch.arange(192, dtype=torch.float32).reshape(1, 1, 1, 12, 16)
|
|
|
|
tile_tasks, grid_spec = vae.encode_tile_split(x)
|
|
encoded_tiles = {task.grid_coord: vae.encode_tile_exec(task) for task in tile_tasks}
|
|
output = vae.tile_merge(encoded_tiles, grid_spec)
|
|
|
|
assert grid_spec.split_dims == (3, 4)
|
|
assert grid_spec.grid_shape == (2, 2)
|
|
assert grid_spec.tile_spec == {
|
|
"blend_height": 2,
|
|
"blend_width": 3,
|
|
"row_limit_height": 6,
|
|
"row_limit_width": 9,
|
|
}
|
|
assert len(tile_tasks) == 4
|
|
assert torch.equal(output, x + 20)
|