From a8a0981bf1283a96668210c2c9c539654e568bee Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Thu, 13 Aug 2026 20:31:17 +0000 Subject: [PATCH] Compare full pipeline outputs in the modular tester mixin `test_save_from_pretrained` and `test_components_auto_cpu_offload_inference_consistent` compared outputs through `output[0, -3:, -3:, -1]`, which only means "a 3x3 corner of the first image" for a 4-D `(B, H, W, C)` output. Audio `(B, C, samples)` and 5-D video outputs still index without error, so a non-image pipeline gets a silently degenerate comparison or has to override the test. Compare the whole output instead, as `cosmos` already does by hand in its own override and as `test_float16_inference` does. Co-Authored-By: Claude Opus 5 --- .../test_modular_pipelines_common.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index 759da1ac13b0..3292774f6f87 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -366,13 +366,12 @@ def test_components_auto_cpu_offload_inference_consistent(self): cm.enable_auto_cpu_offload(device=torch_device) offload_pipe = self.get_pipeline(components_manager=cm) - image_slices = [] + outputs = [] for pipe in [base_pipe, offload_pipe]: inputs = self.get_dummy_inputs() - image = pipe(**inputs, output=self.output_name) - image_slices.append(image[0, -3:, -3:, -1].flatten()) + outputs.append(pipe(**inputs, output=self.output_name)) - assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3 + assert torch.abs(outputs[0] - outputs[1]).max() < 1e-3 @require_accelerator def test_group_offloading_execution_device(self): @@ -413,13 +412,12 @@ def test_save_from_pretrained(self, tmp_path): pipes.append(pipe) - image_slices = [] + outputs = [] for pipe in pipes: inputs = self.get_dummy_inputs() - image = pipe(**inputs, output=self.output_name) - image_slices.append(image[0, -3:, -3:, -1].flatten()) + outputs.append(pipe(**inputs, output=self.output_name)) - assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3 + assert torch.abs(outputs[0] - outputs[1]).max() < 1e-3 def test_load_expected_components_from_pretrained(self, tmp_path): pipe = self.get_pipeline()