Skip to content

Commit f6a3517

Browse files
yiyixuxuclaude
andauthored
Add not_params to ModularPipelineTesterMixin (#14207)
Pipelines sometimes deliberately drop a canonical param — e.g. guidance-distilled pipelines don't accept negative_prompt. Test classes can now declare these in not_params and test_pipeline_call_signature asserts they are absent from the pipeline's inputs, so reintroducing one by accident fails the test. Adopted in the flux2-klein test classes (negative_prompt). Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
1 parent 8eac6ad commit f6a3517

2 files changed

Lines changed: 11 additions & 0 deletions

File tree

tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ class TestFlux2KleinModularPipelineFast(ModularPipelineTesterMixin):
4949

5050
params = frozenset(["prompt", "height", "width"])
5151
batch_params = frozenset(["prompt"])
52+
not_params = frozenset(["negative_prompt"])
5253
expected_workflow_blocks = FLUX2_KLEIN_WORKFLOWS
5354

5455
def get_dummy_inputs(self, seed=0):
@@ -94,6 +95,7 @@ class TestFlux2KleinImageConditionedModularPipelineFast(ModularPipelineTesterMix
9495

9596
params = frozenset(["prompt", "height", "width", "image"])
9697
batch_params = frozenset(["prompt", "image"])
98+
not_params = frozenset(["negative_prompt"])
9799
expected_workflow_blocks = FLUX2_KLEIN_IMAGE_CONDITIONED_WORKFLOWS
98100

99101
def get_dummy_inputs(self, seed=0):

tests/modular_pipelines/test_modular_pipelines_common.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,10 @@ class ModularPipelineTesterMixin:
6969
# of the type of pipeline. They are always optional and have common
7070
# sense default values.
7171
optional_params = frozenset(["num_inference_steps", "num_images_per_prompt", "latents", "output_type"])
72+
# Parameters the pipeline deliberately does NOT accept — e.g. `negative_prompt` on a
73+
# guidance-distilled pipeline. `test_pipeline_call_signature` asserts they are absent,
74+
# so accidentally (re)introducing one fails the test.
75+
not_params = frozenset()
7276
# this is modular specific: generator needs to be a intermediate input because it's mutable
7377
intermediate_params = frozenset(["generator"])
7478
# Output type for the pipeline (e.g., "images" for image pipelines, "videos" for video pipelines)
@@ -176,6 +180,11 @@ def _check_for_parameters(parameters, expected_parameters, param_type):
176180
_check_for_parameters(self.params, input_parameters, "input")
177181
_check_for_parameters(self.optional_params, optional_parameters, "optional")
178182

183+
unsupported_parameters = {param for param in self.not_params if param in input_parameters}
184+
assert len(unsupported_parameters) == 0, (
185+
f"Parameters declared in `not_params` unexpectedly present in the pipeline inputs: {unsupported_parameters}"
186+
)
187+
179188
def test_inference_batch_consistent(self, batch_sizes=[2], batch_generator=True):
180189
pipe = self.get_pipeline().to(torch_device)
181190

0 commit comments

Comments
 (0)