[CI] Fix GLM-Image usage unit tests (#33731)
This commit is contained in:
@@ -47,6 +47,17 @@ class _FakeResponse:
|
|||||||
return data
|
return data
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeBatchResponse:
|
||||||
|
def __init__(self, outputs):
|
||||||
|
self._outputs = outputs
|
||||||
|
|
||||||
|
def raise_for_status(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
def json(self):
|
||||||
|
return self._outputs
|
||||||
|
|
||||||
|
|
||||||
class TestGlmImageARSrtBackend(unittest.TestCase):
|
class TestGlmImageARSrtBackend(unittest.TestCase):
|
||||||
def _server_args(self):
|
def _server_args(self):
|
||||||
return SimpleNamespace(
|
return SimpleNamespace(
|
||||||
@@ -135,16 +146,18 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
|
|||||||
)
|
)
|
||||||
def test_srt_ar_forward_aggregates_usage(self, mock_post, _mock_device):
|
def test_srt_ar_forward_aggregates_usage(self, mock_post, _mock_device):
|
||||||
set_global_server_args(self._server_args())
|
set_global_server_args(self._server_args())
|
||||||
mock_post.side_effect = [
|
mock_post.return_value = _FakeBatchResponse(
|
||||||
_FakeResponse(
|
[
|
||||||
list(range(1025)),
|
{
|
||||||
meta_info={"prompt_tokens": 13, "completion_tokens": 25},
|
"output_ids": list(range(1025)),
|
||||||
),
|
"meta_info": {"prompt_tokens": 13, "completion_tokens": 25},
|
||||||
_FakeResponse(
|
},
|
||||||
list(range(1025)),
|
{
|
||||||
meta_info={"prompt_tokens": 13, "completion_tokens": 25},
|
"output_ids": list(range(1025)),
|
||||||
),
|
"meta_info": {"prompt_tokens": 13, "completion_tokens": 25},
|
||||||
]
|
},
|
||||||
|
]
|
||||||
|
)
|
||||||
stage = GlmImageAR(processor=_FakeProcessor(), vision_language_encoder=None)
|
stage = GlmImageAR(processor=_FakeProcessor(), vision_language_encoder=None)
|
||||||
batch = SimpleNamespace(
|
batch = SimpleNamespace(
|
||||||
prompt="A simple product sketch",
|
prompt="A simple product sketch",
|
||||||
@@ -153,6 +166,7 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
|
|||||||
image_path=None,
|
image_path=None,
|
||||||
num_outputs_per_prompt=2,
|
num_outputs_per_prompt=2,
|
||||||
seed=None,
|
seed=None,
|
||||||
|
extra={},
|
||||||
)
|
)
|
||||||
|
|
||||||
batch = stage.forward(batch, self._server_args())
|
batch = stage.forward(batch, self._server_args())
|
||||||
|
|||||||
@@ -30,7 +30,7 @@ class _RecordingGlmImageAR(GlmImageAR):
|
|||||||
def generate_prior_tokens(self, **kwargs):
|
def generate_prior_tokens(self, **kwargs):
|
||||||
self.initial_seeds.append(torch.initial_seed())
|
self.initial_seeds.append(torch.initial_seed())
|
||||||
output_idx = len(self.initial_seeds)
|
output_idx = len(self.initial_seeds)
|
||||||
return torch.full((1, 4), output_idx, dtype=torch.long), None
|
return torch.full((1, 4), output_idx, dtype=torch.long), None, None
|
||||||
|
|
||||||
|
|
||||||
class _DummySchedulerConfig(dict):
|
class _DummySchedulerConfig(dict):
|
||||||
@@ -105,6 +105,7 @@ def test_ar_stage_generates_one_prior_per_requested_output():
|
|||||||
image_path=None,
|
image_path=None,
|
||||||
num_outputs_per_prompt=2,
|
num_outputs_per_prompt=2,
|
||||||
seed=11,
|
seed=11,
|
||||||
|
extra={},
|
||||||
)
|
)
|
||||||
|
|
||||||
with patch.object(
|
with patch.object(
|
||||||
|
|||||||
Reference in New Issue
Block a user