TaffyOfficialandTaffyOfficial
f4d922d8f6
Signed-off-by: TaffyOfficial <2324465096@qq.com> Co-authored-by: TaffyOfficial <2324465096@qq.com>
55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for HunyuanImage3 AR image RGB conversion."""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 import (
|
|
HunyuanImage3Processor,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class _FixedResolutionGroup:
|
|
def get_target_size(self, width: int, height: int):
|
|
return width, height
|
|
|
|
def get_base_size_and_ratio_index(self, width: int, height: int):
|
|
return 1024, 0
|
|
|
|
|
|
def test_ar_processor_uses_pil_rgb_conversion_for_rgba_images():
|
|
captured: dict[str, object] = {}
|
|
|
|
def fake_vit_processor(image: Image.Image):
|
|
captured["mode"] = image.mode
|
|
captured["pixel"] = image.getpixel((0, 0))
|
|
return {
|
|
"pixel_values": torch.zeros((1, 1, 3), dtype=torch.float32),
|
|
"pixel_attention_mask": torch.ones((1, 1), dtype=torch.bool),
|
|
"spatial_shapes": torch.tensor([[1, 1]], dtype=torch.long),
|
|
}
|
|
|
|
def fake_vae_processor(_image: Image.Image):
|
|
return torch.zeros((1, 3, 1, 1), dtype=torch.float32)
|
|
|
|
processor = object.__new__(HunyuanImage3Processor)
|
|
processor.hf_config = SimpleNamespace(
|
|
vit={"num_channels": 3},
|
|
vae_downsample_factor=(1, 1),
|
|
patch_size=1,
|
|
)
|
|
processor.reso_group = _FixedResolutionGroup()
|
|
processor.vision_encoder_processor = fake_vit_processor
|
|
processor.vae_processor = fake_vae_processor
|
|
|
|
processor.process_image(Image.new("RGBA", (1, 1), color=(255, 0, 0, 0)))
|
|
|
|
assert captured["mode"] == "RGB"
|
|
assert captured["pixel"] == (255, 0, 0)
|