65 lines
1.8 KiB
Python
65 lines
1.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
|
|
|
|
|
def _pipeline_without_init() -> SenseNovaU1Pipeline:
|
|
return object.__new__(SenseNovaU1Pipeline)
|
|
|
|
|
|
@pytest.mark.parametrize("kwargs", [None, {"step_i": 0}])
|
|
def test_combine_cfg_noise_requires_is_it2i(kwargs):
|
|
pipe = _pipeline_without_init()
|
|
out_cond = (torch.ones(1, 2, 3),)
|
|
out_uncond = (torch.zeros(1, 2, 3),)
|
|
|
|
with pytest.raises(ValueError, match="is_it2i"):
|
|
pipe.combine_cfg_noise(
|
|
out_cond,
|
|
out_uncond,
|
|
cfg_scale=4.0,
|
|
cfg_norm="cfg_zero_star",
|
|
kwargs=kwargs,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"kwargs",
|
|
[{"is_it2i": False}, {"is_it2i": False, "step_i": None}],
|
|
)
|
|
def test_cfg_zero_star_requires_step_i(kwargs):
|
|
pipe = _pipeline_without_init()
|
|
out_cond = (torch.ones(1, 2, 3),)
|
|
out_uncond = (torch.zeros(1, 2, 3),)
|
|
with pytest.raises(ValueError, match="step_i"):
|
|
pipe.combine_cfg_noise(
|
|
out_cond,
|
|
out_uncond,
|
|
cfg_scale=4.0,
|
|
cfg_norm="cfg_zero_star",
|
|
kwargs=kwargs,
|
|
)
|
|
|
|
|
|
def test_cfg_zero_star_accepts_step_i():
|
|
pipe = _pipeline_without_init()
|
|
out_cond = (torch.ones(1, 2, 3),)
|
|
out_uncond = (torch.zeros(1, 2, 3),)
|
|
result = pipe.combine_cfg_noise(
|
|
out_cond,
|
|
out_uncond,
|
|
cfg_scale=4.0,
|
|
cfg_norm="cfg_zero_star",
|
|
kwargs={"is_it2i": False, "step_i": 0},
|
|
)
|
|
|
|
assert result.shape == out_cond[0].shape
|
|
assert torch.equal(result, torch.zeros_like(out_cond[0]))
|
|
assert torch.isfinite(result).all()
|