Skip to content

refactor qwenimage pipeline tests to the new mixin structure #14220

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sayakpaul merged 7 commits into from
Jul 18, 2026
Merged

refactor qwenimage pipeline tests to the new mixin structure #14220

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
7fc630b
refactor qwenimage pipeline tests to the new mixin structure
akshan-main Jul 17, 2026
62831a1
parametrize num_layers in qwenimage test dummy components for cache m…
akshan-main Jul 18, 2026
8d37224
drop faster cache test for qwenimage (unsupported by the model)
akshan-main Jul 18, 2026
f7c640e
restore test_inference and drop caching tests for qwenimage
akshan-main Jul 18, 2026
b7914e5
Merge branch 'main' into refactor-qwenimage-pipeline-tests
sayakpaul Jul 18, 2026
6f7e386
drop batch-single override in qwenimage tests, base tolerance covers it
akshan-main Jul 18, 2026
6502fe6
Merge branch 'main' into refactor-qwenimage-pipeline-tests
sayakpaul Jul 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
142 changes: 35 additions & 107 deletions tests/pipelines/qwenimage/test_qwenimage.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest

import numpy as np
import torch
from transformers import Qwen2_5_VLConfig, Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer

Expand All @@ -25,41 +22,28 @@
QwenImageTransformer2DModel,
)

from ...testing_utils import enable_full_determinism, torch_device
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, to_np


enable_full_determinism()
from ...testing_utils import torch_device
from ..testing_utils import (
BasePipelineTesterConfig,
MemoryTesterMixin,
PipelineTesterMixin,
)


class QwenImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class QwenImagePipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = QwenImagePipeline
params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_optional_params = frozenset(
[
"num_inference_steps",
"generator",
"latents",
"return_dict",
"callback_on_step_end",
"callback_on_step_end_tensor_inputs",
]
required_input_params_in_call_signature = frozenset(
["prompt", "negative_prompt", "true_cfg_scale", "height", "width", "guidance_scale", "prompt_embeds"]
)
test_xformers_attention = False
test_layerwise_casting = True
test_group_offloading = True
batch_input_params = frozenset(["prompt"])

def get_dummy_components(self):
def get_dummy_components(self, num_layers: int = 2):
torch.manual_seed(0)
transformer = QwenImageTransformer2DModel(
patch_size=2,
in_channels=16,
out_channels=4,
num_layers=2,
num_layers=num_layers,
attention_head_dim=16,
num_attention_heads=3,
joint_attention_dim=16,
Expand All @@ -75,10 +59,8 @@ def get_dummy_components(self):
dim_mult=[1, 2, 4],
num_res_blocks=1,
temperal_downsample=[False, True],
# fmt: off
latents_mean=[0.0] * 4,
latents_std=[1.0] * 4,
# fmt: on
)

torch.manual_seed(0)
Expand Down Expand Up @@ -115,134 +97,76 @@ def get_dummy_components(self):
text_encoder = Qwen2_5_VLForConditionalGeneration(config).eval()
tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration")

components = {
return {
"transformer": transformer,
"vae": vae,
"scheduler": scheduler,
"text_encoder": text_encoder,
"tokenizer": tokenizer,
}
return components

def get_dummy_inputs(self, device, seed=0):
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device=device).manual_seed(seed)

inputs = {
def get_dummy_inputs(self):
return {
"prompt": "dance monkey",
"negative_prompt": "bad quality",
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 3.0,
"true_cfg_scale": 1.0,
"height": 32,
"width": 32,
"max_sequence_length": 16,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}

return inputs

class TestQwenImagePipeline(QwenImagePipelineTesterConfig, PipelineTesterMixin):
def test_inference(self):

@sayakpaul sayakpaul Jul 18, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's keep this test perhaps in TestQwenImagePipeline?

device = "cpu"
# Run on CPU: the expected slice below is CPU-specific.
pipe = self.get_pipeline()

components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)

inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
image = pipe(**inputs).images
generated_image = image[0]
self.assertEqual(generated_image.shape, (3, 32, 32))
assert generated_image.shape == (3, 32, 32)

# fmt: off
expected_slice = torch.tensor([0.5633, 0.6368, 0.6015, 0.5637, 0.5817, 0.5528, 0.5718, 0.6326, 0.4147, 0.3556, 0.5623, 0.4833, 0.4971, 0.5262, 0.4087, 0.5021])
# fmt: on

generated_slice = generated_image.flatten()
generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]])
self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=5e-3))

def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-1)

def test_attention_slicing_forward_pass(

@sayakpaul sayakpaul Jul 18, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this test already covered?

@akshan-main akshan-main Jul 18, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, it's inherited from the base PipelineTesterMixin. The override only loosens the tolerance to 1e-1 (Qwen's batched vs single output exceeds the 1e-4 default, the same value the old test used).

@akshan-main akshan-main Jul 18, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dropped it. Passes at the base 1e-4 default, so the override wasn't needed.

self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
):
if not self.test_attention_slicing:
return

components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)

generator_device = "cpu"
inputs = self.get_dummy_inputs(generator_device)
output_without_slicing = pipe(**inputs)[0]

pipe.enable_attention_slicing(slice_size=1)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing1 = pipe(**inputs)[0]

pipe.enable_attention_slicing(slice_size=2)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing2 = pipe(**inputs)[0]

if test_max_difference:
max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max()
max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max()
self.assertLess(
max(max_diff1, max_diff2),
expected_max_diff,
"Attention slicing should not affect the inference results",
)
assert torch.allclose(generated_slice, expected_slice, atol=5e-3)

def test_vae_tiling(self, expected_diff_max: float = 0.2):
generator_device = "cpu"
components = self.get_dummy_components()

pipe = self.pipeline_class(**components)
pipe.to("cpu")
pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
pipe.set_progress_bar_config(disable=None)

# Without tiling
inputs = self.get_dummy_inputs(generator_device)
inputs = self.get_dummy_inputs()
inputs["height"] = inputs["width"] = 128
output_without_tiling = pipe(**inputs)[0]

# With tiling
pipe.vae.enable_tiling(
tile_sample_min_height=96,
tile_sample_min_width=96,
tile_sample_stride_height=64,
tile_sample_stride_width=64,
)
inputs = self.get_dummy_inputs(generator_device)
inputs = self.get_dummy_inputs()
inputs["height"] = inputs["width"] = 128
output_with_tiling = pipe(**inputs)[0]

self.assertLess(
(to_np(output_without_tiling) - to_np(output_with_tiling)).max(),
expected_diff_max,
"VAE tiling should not affect the inference results",
assert (output_without_tiling - output_with_tiling).abs().max() < expected_diff_max, (
"VAE tiling should not affect the inference results."
)

def test_true_cfg_without_negative_prompt_embeds_mask(self):
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(torch_device)
pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
pipe.set_progress_bar_config(disable=None)

inputs = self.get_dummy_inputs(torch_device)
inputs = self.get_dummy_inputs()
prompt = inputs.pop("prompt")

prompt_embeds, prompt_embeds_mask = pipe.encode_prompt(
prompt=prompt,
device=torch_device,
Expand All @@ -258,4 +182,8 @@ def test_true_cfg_without_negative_prompt_embeds_mask(self):
inputs["true_cfg_scale"] = 2.0

image = pipe(**inputs).images
self.assertIsNotNone(image)
assert image is not None


class TestQwenImagePipelineMemory(QwenImagePipelineTesterConfig, MemoryTesterMixin):
pass
Loading