729 lines
28 KiB
Python
729 lines
28 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
"""Tests for multi-resolution support in HunyuanImage3.
|
||
|
||
Covers three areas:
|
||
|
||
1. ``ResolutionGroup._calc_by_step`` — training-bucket generation,
|
||
verifying that each resolution has correct width×height and that
|
||
indices match expected ratio ordering.
|
||
|
||
2. ``HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS`` — extra resolution matching
|
||
via ``Resolution.append`` and new-entry insertion in
|
||
``ResolutionGroup.__init__``.
|
||
|
||
3. ``get_cached_resolution_group`` and ``ar2diffusion`` — caching
|
||
behaviour and correct ``ResolutionGroup`` construction for the
|
||
AR→DiT stage transition.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
from typing import Any
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
import pytest
|
||
|
||
from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
|
||
HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS,
|
||
Resolution,
|
||
ResolutionGroup,
|
||
get_cached_resolution_group,
|
||
)
|
||
|
||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _make_reso(desc: str | tuple[int, int]) -> Resolution:
|
||
"""Build a ``Resolution`` from a ``"HxW"`` string or ``(h, w)`` tuple."""
|
||
if isinstance(desc, str):
|
||
h, w = desc.split("x")
|
||
return Resolution(int(h), int(w))
|
||
return Resolution(desc[0], desc[1])
|
||
|
||
|
||
def _reso_set(rg: ResolutionGroup) -> set[tuple[int, int]]:
|
||
"""Return all ``(h, w)`` tuples reachable through a ``ResolutionGroup``.
|
||
|
||
For entries that carry ``extra_res`` this includes every alternative
|
||
resolution in that set.
|
||
"""
|
||
out: set[tuple[int, int]] = set()
|
||
for r in rg.data:
|
||
out.add((r.h, r.w))
|
||
out |= r.extra_res
|
||
return out
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 1. _calc_by_step — training-bucket resolution parsing
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_BASE = 1024
|
||
_STEP = _BASE // 16 # 64
|
||
_ALIGN = 16
|
||
_MIN_SIDE = _BASE // 2 # 512
|
||
_MAX_SIDE = _BASE * 2 # 2048
|
||
|
||
|
||
def _build_plain_group(
|
||
base_size: int = _BASE,
|
||
step: int | None = None,
|
||
align: int = _ALIGN,
|
||
) -> ResolutionGroup:
|
||
"""Return a ``ResolutionGroup`` with *no* extra resolutions."""
|
||
return ResolutionGroup(base_size=base_size, step=step, align=align, extra_resolutions=None)
|
||
|
||
|
||
class TestCalcByStep:
|
||
"""Verify ``_calc_by_step`` generates the expected training buckets."""
|
||
|
||
# -- count & range -------------------------------------------------------
|
||
|
||
def test_count_for_base_1024(self):
|
||
"""33 buckets: 1 centre + 2×16 from the stepping loops."""
|
||
rg = _build_plain_group()
|
||
# 1 centre + 16 portrait + 16 landscape
|
||
assert len(rg.data) == 33, f"expected 33 training buckets, got {len(rg.data)}"
|
||
|
||
def test_all_sizes_within_bounds(self):
|
||
rg = _build_plain_group()
|
||
for r in rg.data:
|
||
assert _MIN_SIDE <= r.h <= _MAX_SIDE, f"height {r.h} outside [{_MIN_SIDE}, {_MAX_SIDE}]"
|
||
assert _MIN_SIDE <= r.w <= _MAX_SIDE, f"width {r.w} outside [{_MIN_SIDE}, {_MAX_SIDE}]"
|
||
|
||
def test_all_sizes_aligned(self):
|
||
rg = _build_plain_group()
|
||
for r in rg.data:
|
||
assert r.h % _ALIGN == 0, f"height {r.h} not {_ALIGN}-aligned"
|
||
assert r.w % _ALIGN == 0, f"width {r.w} not {_ALIGN}-aligned"
|
||
|
||
# -- ratio ordering ------------------------------------------------------
|
||
|
||
def test_sorted_by_ratio_ascending(self):
|
||
rg = _build_plain_group()
|
||
ratios = [r.ratio for r in rg.data]
|
||
for i in range(1, len(ratios)):
|
||
assert ratios[i - 1] <= ratios[i], f"ratio[{i - 1}]={ratios[i - 1]:.4f} > ratio[{i}]={ratios[i]:.4f}"
|
||
|
||
def test_ratio_index_matches_numpy_array(self):
|
||
rg = _build_plain_group()
|
||
for i, r in enumerate(rg.data):
|
||
assert math.isclose(rg.ratio[i], r.ratio, rel_tol=1e-9), (
|
||
f"rg.ratio[{i}]={rg.ratio[i]:.6f} != data[{i}].ratio={r.ratio:.6f}"
|
||
)
|
||
|
||
# -- key index positions -------------------------------------------------
|
||
|
||
@pytest.mark.parametrize(
|
||
"idx,expected_hw",
|
||
[
|
||
# First (most portrait) → 1:4 ratio bucket
|
||
(0, (512, 2048)),
|
||
# Centre square bucket
|
||
(16, (1024, 1024)),
|
||
# Most landscape bucket → 4:1 ratio
|
||
(32, (2048, 512)),
|
||
],
|
||
)
|
||
def test_landmark_index(self, idx, expected_hw):
|
||
rg = _build_plain_group()
|
||
r = rg.data[idx]
|
||
assert (r.h, r.w) == expected_hw, f"idx {idx}: expected {expected_hw}, got ({r.h}, {r.w})"
|
||
|
||
def test_center_is_square(self):
|
||
rg = _build_plain_group()
|
||
r = rg.data[16]
|
||
assert r.h == r.w == _BASE
|
||
|
||
def test_first_and_last_are_extremes(self):
|
||
rg = _build_plain_group()
|
||
# first → 512×2048 (ratio 0.25)
|
||
assert rg.data[0].ratio < 0.3
|
||
# last → 2048×512 (ratio 4.0)
|
||
assert rg.data[-1].ratio > 3.9
|
||
|
||
# -- width × height matching ---------------------------------------------
|
||
|
||
def test_every_resolution_has_same_area_product(self):
|
||
"""Symmetry: stepping loop produces the same product set in both
|
||
directions, so (h × w) tuples should appear symmetrically."""
|
||
rg = _build_plain_group()
|
||
products = {(r.h, r.w) for r in rg.data}
|
||
# For every (h, w) there should exist (w, h) or a nearby equivalent
|
||
for h, w in products:
|
||
if h == w:
|
||
continue # square — trivially symmetric
|
||
assert (w, h) in products or any(abs(rr.h - w) <= _STEP and abs(rr.w - h) <= _STEP for rr in rg.data), (
|
||
f"no symmetric counterpart for {h}×{w}"
|
||
)
|
||
|
||
def test_product_of_wh_divisible_by_256(self):
|
||
"""Each bucket's area should be a multiple of 256 (16×16),
|
||
which is the VAE down-sample grid."""
|
||
rg = _build_plain_group()
|
||
for r in rg.data:
|
||
assert (r.h * r.w) % 256 == 0, f"{r.h}×{r.w} = {r.h * r.w} not divisible by 256"
|
||
|
||
# -- custom base_size / step / align -------------------------------------
|
||
|
||
@pytest.mark.parametrize("base_size", [512, 768, 1280])
|
||
def test_count_scales_with_base_size(self, base_size):
|
||
step = base_size // 16
|
||
rg = ResolutionGroup(base_size=base_size, step=step, align=16, extra_resolutions=None)
|
||
expected = 1 + 2 * (base_size // step) # centre + 2×phases
|
||
assert len(rg.data) == expected, f"base_size={base_size}: expected {expected}, got {len(rg.data)}"
|
||
|
||
def test_custom_step(self):
|
||
"""A smaller step produces more buckets."""
|
||
rg_small = ResolutionGroup(base_size=1024, step=32, align=16, extra_resolutions=None)
|
||
rg_default = _build_plain_group()
|
||
assert len(rg_small.data) > len(rg_default.data), "smaller step should produce more buckets"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 2. HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS matching
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_EXTRA_SQUARES = {"512x512", "640x640", "768x768", "896x896"}
|
||
_EXTRA_NON_SQUARES = {"1024x768", "1280x720", "768x1024", "720x1280"}
|
||
|
||
|
||
def _build_extra_group() -> ResolutionGroup:
|
||
return ResolutionGroup(
|
||
base_size=_BASE,
|
||
extra_resolutions=[_make_reso(s) for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS],
|
||
)
|
||
|
||
|
||
class TestExtraResolutionMatching:
|
||
"""Verify extra resolutions are correctly integrated into
|
||
the ``ResolutionGroup``."""
|
||
|
||
# -- square extras (ratio = 1.0) -----------------------------------------
|
||
|
||
def test_square_extras_appended_to_base_square(self):
|
||
"""512², 640², 768², 896² all have ratio 1.0 and should be
|
||
appended to the (1024, 1024) Resolution via ``append()``."""
|
||
rg = _build_extra_group()
|
||
# Find the 1024×1024 entry
|
||
square_reso = next(r for r in rg.data if r.h == 1024 and r.w == 1024)
|
||
expected_extra = {(512, 512), (640, 640), (768, 768), (896, 896)}
|
||
assert square_reso.extra_res == expected_extra, f"square extra_res: {square_reso.extra_res}"
|
||
|
||
def test_base_1024_unchanged(self):
|
||
"""The 1024×1024 entry itself still reports h=1024, w=1024."""
|
||
rg = _build_extra_group()
|
||
r = next(r for r in rg.data if r.h == 1024 and r.w == 1024)
|
||
assert r.h == 1024 and r.w == 1024
|
||
|
||
def test_squares_not_standalone_entries(self):
|
||
"""Square extras must NOT appear as top-level ``Resolution`` entries;
|
||
they live inside ``extra_res`` only."""
|
||
rg = _build_extra_group()
|
||
standalone = {(r.h, r.w) for r in rg.data}
|
||
for square in [(512, 512), (640, 640), (768, 768), (896, 896)]:
|
||
assert square not in standalone, f"{square} should be in extra_res, not a standalone entry"
|
||
|
||
# -- non-square extras (ratio ≠ 1.0) ------------------------------------
|
||
|
||
@pytest.mark.parametrize("reso_str", sorted(_EXTRA_NON_SQUARES))
|
||
def test_non_square_extra_is_standalone_entry(self, reso_str):
|
||
"""1024×768, 1280×720, 768×1024, 720×1280 are appended as
|
||
new ``Resolution`` entries because no training bucket shares
|
||
their exact ratio."""
|
||
r = _make_reso(reso_str)
|
||
rg = _build_extra_group()
|
||
matching = [d for d in rg.data if d.h == r.h and d.w == r.w]
|
||
assert len(matching) == 1, f"{reso_str} should be in rg.data exactly once"
|
||
assert matching[0].extra_res == set(), f"{reso_str} should have empty extra_res, got {matching[0].extra_res}"
|
||
|
||
def test_non_square_extras_empty_extra_res(self):
|
||
"""Non-square extras have no own extra_res."""
|
||
rg = _build_extra_group()
|
||
for r in rg.data:
|
||
if (r.h, r.w) in {(1024, 768), (1280, 720), (768, 1024), (720, 1280)}:
|
||
assert r.extra_res == set()
|
||
|
||
# -- total count ---------------------------------------------------------
|
||
|
||
def test_total_count(self):
|
||
"""33 training + 4 non-square extras = 37."""
|
||
rg = _build_extra_group()
|
||
assert len(rg.data) == 37, f"expected 37, got {len(rg.data)}"
|
||
|
||
def test_extra_count(self):
|
||
"""Exactly 8 extra resolutions (4 square + 4 non-square)."""
|
||
rg = _build_extra_group()
|
||
reso_set = _reso_set(rg)
|
||
extra_in_set = reso_set - {(r.h, r.w) for r in _build_plain_group().data}
|
||
assert len(extra_in_set) == 8, f"expected 8 extra, got {len(extra_in_set)}: {extra_in_set}"
|
||
|
||
# -- index stability -----------------------------------------------------
|
||
|
||
def test_indices_0_32_unchanged(self):
|
||
"""Indices 0–32 remain the training buckets (sorted by ratio)."""
|
||
plain = _build_plain_group()
|
||
rg = _build_extra_group()
|
||
for i in range(33):
|
||
assert (rg.data[i].h, rg.data[i].w) == (plain.data[i].h, plain.data[i].w), (
|
||
f"idx {i}: extra group has ({rg.data[i].h},{rg.data[i].w}), "
|
||
f"expected ({plain.data[i].h},{plain.data[i].w})"
|
||
)
|
||
|
||
@pytest.mark.parametrize(
|
||
"reso_str,expected_idx",
|
||
[
|
||
("1024x768", 33),
|
||
("1280x720", 34),
|
||
("768x1024", 35),
|
||
("720x1280", 36),
|
||
],
|
||
)
|
||
def test_non_square_extra_indices(self, reso_str, expected_idx):
|
||
"""Non-square extras are appended in order at indices 33-36."""
|
||
r = _make_reso(reso_str)
|
||
rg = _build_extra_group()
|
||
idx = next(i for i, d in enumerate(rg.data) if d.h == r.h and d.w == r.w)
|
||
assert idx == expected_idx, f"{reso_str}: expected idx {expected_idx}, got {idx}"
|
||
|
||
# -- ratio lookup via get_base_size_and_ratio_index ----------------------
|
||
|
||
@pytest.mark.parametrize(
|
||
"h,w,expected_idx",
|
||
[
|
||
(1024, 768, 33),
|
||
(1280, 720, 34),
|
||
(768, 1024, 35),
|
||
(720, 1280, 36),
|
||
(1024, 1024, 16),
|
||
(512, 512, 16),
|
||
(640, 640, 16),
|
||
(768, 768, 16),
|
||
(896, 896, 16),
|
||
],
|
||
)
|
||
def test_ratio_index_for_resolutions(self, h, w, expected_idx):
|
||
"""``get_base_size_and_ratio_index`` must return the correct
|
||
ratio index for both training-bucket and extra resolutions."""
|
||
rg = _build_extra_group()
|
||
_, idx = rg.get_base_size_and_ratio_index(w, h)
|
||
assert idx == expected_idx, f"({w},{h}): expected idx {expected_idx}, got {idx}"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Resolution.match — area-based matching within a ratio group
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestResolutionMatch:
|
||
"""Verify ``Resolution.match`` picks the best-fit alternative from
|
||
``extra_res`` based on area proximity."""
|
||
|
||
def test_no_extra_returns_self(self):
|
||
r = _make_reso("1024x768")
|
||
assert r.match(800, 600) == (768, 1024)
|
||
|
||
@pytest.mark.parametrize(
|
||
"target_w,target_h,expected",
|
||
[
|
||
# Exact area match in extra_res
|
||
(512, 512, (512, 512)),
|
||
(640, 640, (640, 640)),
|
||
(768, 768, (768, 768)),
|
||
(896, 896, (896, 896)),
|
||
# Closer to 1024² than to any extra square
|
||
(1024, 1024, (1024, 1024)),
|
||
(1000, 1000, (1024, 1024)),
|
||
# Small square picks 512²
|
||
(256, 256, (512, 512)),
|
||
(400, 400, (512, 512)),
|
||
# 700² → closest to 640² (diff 129600) vs 768² (diff 99856)
|
||
# 768*768 - 700*700 = 589824 - 490000 = 99824
|
||
# 640*640 - 700*700 = 409600 - 490000 = -80400 → abs=80400
|
||
# So 640² is closer in area
|
||
(700, 700, (640, 640)),
|
||
# 800² → closest to 768² (diff 49984) vs 896² (diff 162944)
|
||
(800, 800, (768, 768)),
|
||
],
|
||
)
|
||
def test_square_extra_match(self, target_w, target_h, expected):
|
||
"""Build a Resolution like the 1024² entry with square extras,
|
||
then verify area-based matching."""
|
||
r = Resolution(1024, 1024)
|
||
for s in [(512, 512), (640, 640), (768, 768), (896, 896)]:
|
||
r.extra_res.add(s)
|
||
result = r.match(target_w, target_h)
|
||
assert result == expected, f"match({target_w},{target_h}) = {result}, expected {expected}"
|
||
|
||
def test_match_returns_tuple(self):
|
||
r = Resolution(1024, 1024)
|
||
r.extra_res.add((512, 512))
|
||
result = r.match(512, 512)
|
||
assert isinstance(result, tuple)
|
||
assert len(result) == 2
|
||
assert all(isinstance(v, int) for v in result)
|
||
|
||
def test_match_does_not_mutate_resolution(self):
|
||
r = Resolution(1024, 1024)
|
||
r.extra_res.add((512, 512))
|
||
extra_before = set(r.extra_res)
|
||
r.match(300, 300)
|
||
assert r.extra_res == extra_before
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# ResolutionGroup.get_target_size — end-to-end resolution selection
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestGetTargetSize:
|
||
"""Verify ``get_target_size`` uses ratio-then-area logic correctly."""
|
||
|
||
@pytest.fixture(scope="class")
|
||
def rg(self) -> ResolutionGroup:
|
||
return _build_extra_group()
|
||
|
||
@pytest.mark.parametrize(
|
||
"w,h,expected_min_area,expected_max_area",
|
||
[
|
||
# Square targets should land on square variants
|
||
(1024, 1024, 1024 * 1024, 1024 * 1024),
|
||
(512, 512, 512 * 512, 512 * 512),
|
||
(640, 640, 640 * 640, 640 * 640),
|
||
(768, 768, 768 * 768, 768 * 768),
|
||
(896, 896, 896 * 896, 896 * 896),
|
||
],
|
||
)
|
||
def test_square_target_returns_square(self, rg, w, h, expected_min_area, expected_max_area):
|
||
rw, rh = rg.get_target_size(w, h)
|
||
area = rw * rh
|
||
assert expected_min_area <= area <= expected_max_area, f"({w},{h}) → ({rw},{rh}) area={area}"
|
||
|
||
def test_training_bucket_self_maps(self, rg):
|
||
"""A training bucket should map to itself (or an extra_res variant
|
||
with the same ratio)."""
|
||
for r in rg.data:
|
||
rw, rh = rg.get_target_size(r.w, r.h)
|
||
# Either exact match or the ratio must be preserved
|
||
assert abs(rh / rw - r.ratio) < 0.01, f"({r.w},{r.h}) → ({rw},{rh}) ratio drifted from {r.ratio:.4f}"
|
||
|
||
def test_wild_ratio_still_finds_bucket(self, rg):
|
||
"""Even an unusual aspect ratio should map to some valid bucket."""
|
||
for w, h in [(800, 600), (1200, 900), (1600, 400), (400, 1600)]:
|
||
rw, rh = rg.get_target_size(w, h)
|
||
assert rw > 0 and rh > 0
|
||
assert rw % 16 == 0 and rh % 16 == 0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 3. get_cached_resolution_group & ar2diffusion integration
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestGetCachedResolutionGroup:
|
||
"""Verify ``get_cached_resolution_group`` returns a correctly
|
||
configured ``ResolutionGroup`` and respects the LRU cache."""
|
||
|
||
def test_returns_resolution_group(self):
|
||
rg = get_cached_resolution_group(base_size=1024)
|
||
assert isinstance(rg, ResolutionGroup)
|
||
|
||
def test_base_size_propagated(self):
|
||
rg = get_cached_resolution_group(base_size=1024)
|
||
assert rg.base_size == 1024
|
||
|
||
def test_has_extra_resolutions(self):
|
||
rg = get_cached_resolution_group(base_size=1024)
|
||
reso_set = _reso_set(rg)
|
||
for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS:
|
||
r = _make_reso(s)
|
||
assert (r.h, r.w) in reso_set, f"{s} missing from cached group"
|
||
|
||
def test_cache_same_base_size_same_object(self):
|
||
"""LRU cache: same base_size → same object identity."""
|
||
# Clear any prior cache entries
|
||
get_cached_resolution_group.cache_clear()
|
||
rg1 = get_cached_resolution_group(base_size=1024)
|
||
rg2 = get_cached_resolution_group(base_size=1024)
|
||
assert rg1 is rg2, "same base_size must return cached object"
|
||
|
||
def test_cache_different_base_size_different_object(self):
|
||
get_cached_resolution_group.cache_clear()
|
||
rg1 = get_cached_resolution_group(base_size=1024)
|
||
rg2 = get_cached_resolution_group(base_size=768)
|
||
assert rg1 is not rg2, "different base_size must return different objects"
|
||
|
||
def test_cache_size_limit(self):
|
||
"""Cache is bounded at 4 entries (``maxsize=4``)."""
|
||
get_cached_resolution_group.cache_clear()
|
||
for bs in [512, 768, 1024, 1280, 1536]:
|
||
get_cached_resolution_group(base_size=bs)
|
||
info = get_cached_resolution_group.cache_info()
|
||
assert info.currsize <= 4, f"cache grew to {info.currsize}, max is 4"
|
||
# All 5 unique base_size values caused cache misses
|
||
assert info.misses >= 5, f"expected >=5 misses for 5 unique base_size values, got {info.misses}"
|
||
|
||
def test_resolution_indexing_works(self):
|
||
"""Indices 0..len(rg)-1 must be accessible."""
|
||
rg = get_cached_resolution_group(base_size=1024)
|
||
for i in range(len(rg)):
|
||
r = rg[i]
|
||
assert isinstance(r, Resolution)
|
||
assert r.h > 0 and r.w > 0
|
||
|
||
@pytest.mark.parametrize("base_size", [512, 1024])
|
||
def test_ratio_index_roundtrip_via_group_length(self, base_size):
|
||
"""Every ratio index in [0, len(rg)) must resolve to a valid
|
||
Resolution with height/width within bounds."""
|
||
rg = get_cached_resolution_group(base_size=base_size)
|
||
for i in range(len(rg)):
|
||
r = rg[i]
|
||
assert isinstance(r.height, int) and r.height > 0
|
||
assert isinstance(r.width, int) and r.width > 0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# ar2diffusion — integration tests
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
# Token IDs matching the real prompt_utils values
|
||
_RATIO_0 = 128044
|
||
_RATIO_32 = 128076
|
||
_RATIO_33 = 130103
|
||
_RATIO_36 = 130106
|
||
|
||
|
||
def _fake_ar_output(
|
||
*,
|
||
cumulative_token_ids: list[int] | None = None,
|
||
text: str = "",
|
||
) -> MagicMock:
|
||
"""Minimal stub for a vLLM ``RequestOutput`` with one completion."""
|
||
choice = MagicMock()
|
||
choice.cumulative_token_ids = cumulative_token_ids or []
|
||
choice.text = text
|
||
output = MagicMock()
|
||
output.outputs = [choice]
|
||
return output
|
||
|
||
|
||
class TestAr2DiffusionIntegration:
|
||
"""Verify that ``ar2diffusion`` calls ``get_cached_resolution_group``
|
||
and obtains a correctly configured ``ResolutionGroup`` whose indices map
|
||
to the expected training-bucket or extra-resolution sizes."""
|
||
|
||
def _call_ar2diffusion(
|
||
self,
|
||
token_ids: list[int],
|
||
prompt_kwargs: dict[str, Any] | None = None,
|
||
) -> tuple[Any, Any]:
|
||
"""Invoke ``ar2diffusion`` with a single fake AR output.
|
||
|
||
Returns ``(diffusion_inputs, ar_predicted_sizes)`` so callers can
|
||
assert on the resolved height/width.
|
||
"""
|
||
from vllm_omni.model_executor.stage_input_processors.hunyuan_image3 import (
|
||
ar2diffusion,
|
||
)
|
||
|
||
pk = dict(prompt_kwargs or {})
|
||
pk.setdefault("prompt", "a beautiful sunset")
|
||
pk.setdefault("height", 1024)
|
||
pk.setdefault("width", 1024)
|
||
pk.setdefault("image_base_size", 1024)
|
||
|
||
# Build a prompt-like object that behaves like the real OmniTokensPrompt
|
||
class _FakePrompt:
|
||
def __init__(self, d):
|
||
self._d = d
|
||
|
||
def __getitem__(self, key):
|
||
return self._d[key]
|
||
|
||
def get(self, key, default=None):
|
||
return self._d.get(key, default)
|
||
|
||
def _asdict(self):
|
||
return dict(self._d)
|
||
|
||
prompt = _FakePrompt(pk)
|
||
ar_output = _fake_ar_output(cumulative_token_ids=token_ids)
|
||
result = ar2diffusion(
|
||
source_outputs=[ar_output],
|
||
prompt=[prompt],
|
||
requires_multimodal_data=False,
|
||
)
|
||
return result
|
||
|
||
# -- ratio token ↔ resolution resolution ---------------------------------
|
||
|
||
@pytest.mark.parametrize(
|
||
"ratio_idx,expected_hw",
|
||
[
|
||
# Training buckets
|
||
(0, (512, 2048)),
|
||
(16, (1024, 1024)),
|
||
(32, (2048, 512)),
|
||
# Non-square extras
|
||
(33, (1024, 768)),
|
||
(34, (1280, 720)),
|
||
(35, (768, 1024)),
|
||
(36, (720, 1280)),
|
||
],
|
||
)
|
||
def test_ratio_index_resolves_to_correct_size(self, ratio_idx, expected_hw):
|
||
"""When the AR emits ``<img_ratio_N>``, the DiT input must carry
|
||
the training-bucket width/height for that index."""
|
||
rg = get_cached_resolution_group(base_size=1024)
|
||
reso = rg[ratio_idx]
|
||
assert (reso.height, reso.width) == expected_hw, (
|
||
f"ratio_idx {ratio_idx}: expected {expected_hw}, got ({reso.height}, {reso.width})"
|
||
)
|
||
|
||
@pytest.mark.parametrize("ratio_idx", [0, 16, 32, 33, 34, 35, 36])
|
||
def test_ar2diffusion_extracts_height_width(self, ratio_idx):
|
||
"""End-to-end: simulate AR output containing ``<img_ratio_{idx}>``
|
||
and assert the DiT inputs carry the expected size."""
|
||
rg = get_cached_resolution_group(base_size=1024)
|
||
expected_h = rg[ratio_idx].height
|
||
expected_w = rg[ratio_idx].width
|
||
|
||
token_id = _RATIO_0 + ratio_idx if ratio_idx <= 32 else _RATIO_33 + (ratio_idx - 33)
|
||
result = self._call_ar2diffusion(
|
||
token_ids=[150000, 150001, token_id],
|
||
prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024},
|
||
)
|
||
assert len(result) == 1
|
||
diffusion_input = result[0]
|
||
assert diffusion_input["height"] == expected_h, (
|
||
f"ratio_idx={ratio_idx}: expected height {expected_h}, got {diffusion_input['height']}"
|
||
)
|
||
assert diffusion_input["width"] == expected_w, (
|
||
f"ratio_idx={ratio_idx}: expected width {expected_w}, got {diffusion_input['width']}"
|
||
)
|
||
|
||
def test_no_ratio_token_falls_back_to_prompt_size(self):
|
||
"""If no ratio token is present, height/width from the prompt
|
||
are kept unchanged."""
|
||
result = self._call_ar2diffusion(
|
||
token_ids=[150000, 150001], # no ratio tokens
|
||
prompt_kwargs={"height": 768, "width": 1024, "image_base_size": 1024},
|
||
)
|
||
assert len(result) == 1
|
||
diffusion_input = result[0]
|
||
assert diffusion_input["height"] == 768
|
||
assert diffusion_input["width"] == 1024
|
||
|
||
def test_out_of_range_ratio_index_falls_back(self):
|
||
"""If ratio_idx ≥ len(reso_group), ar2diffusion must keep the
|
||
prompt size (fallback path)."""
|
||
with patch(
|
||
"vllm_omni.model_executor.stage_input_processors.hunyuan_image3._extract_ratio_index",
|
||
return_value=999, # out of range
|
||
):
|
||
result = self._call_ar2diffusion(
|
||
token_ids=[150000],
|
||
prompt_kwargs={"height": 512, "width": 512, "image_base_size": 1024},
|
||
)
|
||
assert len(result) == 1
|
||
assert result[0]["height"] == 512
|
||
assert result[0]["width"] == 512
|
||
|
||
# -- cache behaviour -----------------------------------------------------
|
||
|
||
def test_ar2diffusion_uses_cached_group(self):
|
||
"""Each invocation of ``ar2diffusion`` calls
|
||
``get_cached_resolution_group``, which returns a shared cached object."""
|
||
get_cached_resolution_group.cache_clear()
|
||
rg_before = get_cached_resolution_group(base_size=1024)
|
||
|
||
token_id = _RATIO_0 + 16 # <img_ratio_16> = 1024×1024
|
||
self._call_ar2diffusion(
|
||
token_ids=[token_id],
|
||
prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024},
|
||
)
|
||
|
||
rg_after = get_cached_resolution_group(base_size=1024)
|
||
assert rg_before is rg_after, "get_cached_resolution_group should return the same cached object"
|
||
|
||
def test_correct_reso_group_used_for_base_size(self):
|
||
"""Verify ``ar2diffusion`` constructs or retrieves the
|
||
ResolutionGroup for the correct ``image_base_size``."""
|
||
get_cached_resolution_group.cache_clear()
|
||
|
||
# First request with a non-default base_size
|
||
token_id = _RATIO_0 # <img_ratio_0>
|
||
self._call_ar2diffusion(
|
||
token_ids=[token_id],
|
||
prompt_kwargs={"height": 512, "width": 512, "image_base_size": 512},
|
||
)
|
||
|
||
# The cached group for base_size=512 should exist
|
||
rg512 = get_cached_resolution_group(base_size=512)
|
||
assert rg512.base_size == 512
|
||
|
||
# And base_size=1024 should also be separately accessible
|
||
rg1024 = get_cached_resolution_group(base_size=1024)
|
||
assert rg1024.base_size == 1024
|
||
assert rg512 is not rg1024
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Edge cases
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestEdgeCases:
|
||
def test_resolution_string_parsing(self):
|
||
"""Resolution accepts both "HxW" strings and (h, w) tuples."""
|
||
r1 = Resolution("512x2048")
|
||
assert r1.h == 512 and r1.w == 2048
|
||
|
||
r2 = Resolution(768, 1024)
|
||
assert r2.h == 768 and r2.w == 1024
|
||
|
||
r3 = Resolution(1024) # single int → square
|
||
assert r3.h == 1024 and r3.w == 1024
|
||
|
||
def test_resolution_extra_res_initially_empty(self):
|
||
r = Resolution(1024, 1024)
|
||
assert r.extra_res == set()
|
||
assert r.match(512, 512) == (1024, 1024)
|
||
|
||
def test_resolution_append_deduplicates(self):
|
||
"""Appending the same resolution twice is a no-op on the set."""
|
||
r = Resolution(1024, 1024)
|
||
r.append(Resolution(512, 512))
|
||
r.append(Resolution(512, 512))
|
||
assert len(r.extra_res) == 1
|
||
|
||
def test_all_final_resolutions_are_16_aligned(self):
|
||
"""Every reachable resolution (including extras) must be
|
||
divisible by the VAE down-sample factor of 16."""
|
||
rg = _build_extra_group()
|
||
for h, w in _reso_set(rg):
|
||
assert h % 16 == 0, f"height {h} not 16-aligned in {h}×{w}"
|
||
assert w % 16 == 0, f"width {w} not 16-aligned in {h}×{w}"
|
||
|
||
def test_ratio_array_length_matches_data(self):
|
||
rg = _build_extra_group()
|
||
assert len(rg.ratio) == len(rg.data)
|
||
|
||
def test_resolution_repr(self):
|
||
r = Resolution(1024, 1024)
|
||
rep = repr(r)
|
||
assert "1024x1024" in rep
|
||
|
||
r2 = Resolution(1024, 1024)
|
||
r2.extra_res.add((512, 512))
|
||
rep2 = repr(r2)
|
||
assert "512x512" in rep2
|