[Chore]: refactor out unused/redundant params in diffusion pipelines (#1235)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/models/flux/test_flux_pipeline.py

@@ -21,6 +21,8 @@ def _make_flux_sampling(**overrides):
         "true_cfg_scale": 4.0,
         "num_outputs_per_prompt": 0,
         "latents": None,
+        "output_type": None,
+        "max_sequence_length": None,
     }
     values.update(overrides)
     return SimpleNamespace(**values)
@@ -131,7 +133,7 @@ def test_forward_collates_request_prompt_tensors_for_flux(monkeypatch):
                     "negative_prompt_embeds": negative_prompt_embeds_a,
                     "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds_a,
                 },
-                sampling_params=_make_flux_sampling(generator=gen_a, latents=latents_a),
+                sampling_params=_make_flux_sampling(generator=gen_a, latents=latents_a, output_type="latent"),
             ),
             SimpleNamespace(
                 request_id="flux-prompt-b",
@@ -145,12 +147,12 @@ def test_forward_collates_request_prompt_tensors_for_flux(monkeypatch):
                         "negative_pooled_prompt_embeds": [negative_pooled_prompt_embeds_b],
                     },
                 },
-                sampling_params=_make_flux_sampling(generator=gen_b, latents=latents_b),
+                sampling_params=_make_flux_sampling(generator=gen_b, latents=latents_b, output_type="latent"),
             ),
         ]
     )
 
-    outputs = pipeline.forward(batch, output_type="latent")
+    outputs = pipeline.forward(batch)
 
     assert encode_calls[0]["prompt"] is None
     assert encode_calls[0]["prompt_2"] is None

tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py

@@ -1,4 +1,3 @@
-import inspect
 from types import SimpleNamespace
 
 import pytest
@@ -193,6 +192,7 @@ def _make_request_batch_prompt_sampling(**overrides):
         "true_cfg_scale": None,
         "guidance_scale_provided": False,
         "guidance_scale": 1.0,
+        "output_type": None,
     }
     values.update(overrides)
     return SimpleNamespace(**values)
@@ -344,16 +344,3 @@ def test_qwen_edit_validator_excludes_image_placeholders_from_budget(pipeline_cl
 
     with pytest.raises(AssertionError, match="text encoder should not run"):
         pipeline.encode_prompt(prompt="short prompt")
-
-
-@pytest.mark.parametrize(
-    "pipeline_class",
-    [
-        QwenImagePipeline,
-        QwenImageLayeredPipeline,
-        QwenImageEditPipeline,
-        QwenImageEditPlusPipeline,
-    ],
-)
-def test_forward_max_sequence_length_default_is_1024(pipeline_class: type):
-    assert inspect.signature(pipeline_class.forward).parameters["max_sequence_length"].default == 1024

tests/diffusion/models/sd3/test_sd3_pipeline.py

@@ -21,6 +21,7 @@ def _make_sd3_sampling(**overrides):
         "generator": None,
         "latents": None,
         "guidance_scale": 4.0,
+        "output_type": None,
     }
     values.update(overrides)
     return SimpleNamespace(**values)

tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -20,6 +20,12 @@ class _StubTransformer(nn.Module):
         return torch.float32
 
 
+class _StubTextEncoder(nn.Module):
+    @property
+    def dtype(self) -> torch.dtype:
+        return torch.float32
+
+
 class _StubScheduler:
     def __init__(self, timesteps: list[int]) -> None:
         self.timesteps = torch.tensor(timesteps, dtype=torch.int64)
@@ -41,12 +47,30 @@ def _noop_progress_bar(*args, **kwargs):
     yield _Bar()
 
 
+def _stub_encode_prompt(
+    prompt,
+    negative_prompt=None,
+    do_classifier_free_guidance=True,
+    num_videos_per_prompt=1,
+    max_sequence_length=512,
+    device=None,
+    dtype=None,
+):
+    del negative_prompt, do_classifier_free_guidance, device, dtype
+    batch_size = 1 if isinstance(prompt, str) else len(prompt)
+    n = batch_size * num_videos_per_prompt
+    hidden_size = 8
+    prompt_embeds = torch.zeros(n, max_sequence_length, hidden_size)
+    return prompt_embeds, None
+
+
 def _make_pipeline() -> Wan22Pipeline:
     pipeline = object.__new__(Wan22Pipeline)
     nn.Module.__init__(pipeline)
     pipeline.device = torch.device("cpu")
     pipeline.transformer = _StubTransformer()
     pipeline.transformer_2 = None
+    pipeline.text_encoder = _StubTextEncoder()
     pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2), in_channels=4, out_channels=4)
     pipeline.scheduler = _StubScheduler([9, 5])
     pipeline.od_config = SimpleNamespace(flow_shift=5.0)
@@ -61,15 +85,14 @@ def _make_pipeline() -> Wan22Pipeline:
     pipeline._num_timesteps = None
     pipeline._current_timestep = None
     pipeline.check_inputs = lambda **kwargs: None
+    pipeline.encode_prompt = _stub_encode_prompt  # type: ignore[method-assign]
     pipeline.prepare_latents = lambda **kwargs: torch.zeros((1, 4, 1, 8, 8), dtype=torch.float32)
     pipeline.progress_bar = _noop_progress_bar
     return pipeline
 
 
-def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None:
+def test_forward_delegates_denoising_to_diffuse() -> None:
     pipeline = _make_pipeline()
-
-    prompt_embeds = torch.randn(1, 8)
     captured: dict[str, object] = {}
 
     def _fake_diffuse(**kwargs):
@@ -86,8 +109,8 @@ def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None:
             width=None,
             num_frames=1,
             num_inference_steps=2,
-            guidance_scale_provided=False,
-            guidance_scale=None,
+            guidance_scale_provided=True,
+            guidance_scale=1.0,
             guidance_scale_2=None,
             boundary_ratio=None,
             generator=None,
@@ -95,14 +118,16 @@ def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None:
             num_outputs_per_prompt=1,
             max_sequence_length=32,
             latents=None,
+            output_type="latent",
             extra_args={},
         ),
     )
     batch = DiffusionRequestBatch(requests=[mock_req])
 
-    output = pipeline.forward(batch, prompt_embeds=prompt_embeds, output_type="latent", guidance_scale=1.0)
+    output = pipeline.forward(batch)
 
     assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8)))
+    assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8))
     assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps)
     assert captured["guidance_low"] == 1.0
     assert captured["guidance_high"] == 1.0