Skip to content

[tests] refactor h and i series pipelines #14633

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鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sayakpaul merged 3 commits into from
Aug 28, 2026
Merged

[tests] refactor h and i series pipelines #14633

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ff6a1c1
refactor h and i series pipelines
sayakpaul Aug 28, 2026
71deaf9
Merge branch 'main' into h-i-series-pipelines
sayakpaul Aug 28, 2026
9000b76
Merge branch 'main' into h-i-series-pipelines
sayakpaul Aug 28, 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
1 change: 1 addition & 0 deletions src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -951,6 +951,7 @@ def __call__(
encoder_hidden_states_t5=prompt_embeds_t5,
encoder_hidden_states_llama3=prompt_embeds_llama3,
pooled_embeds=pooled_prompt_embeds,
attention_kwargs=self.attention_kwargs,
return_dict=False,
)[0]
noise_pred = -noise_pred
Expand Down
84 changes: 44 additions & 40 deletions tests/pipelines/hidream_image/test_pipeline_hidream.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,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 (
AutoConfig,
Expand All @@ -34,23 +31,26 @@
HiDreamImageTransformer2DModel,
)

from ...testing_utils import enable_full_determinism
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from ...testing_utils import assert_tensors_close, enable_full_determinism
from ..testing_utils import (
BasePipelineTesterConfig,
LoraMemoryTesterMixin,
LoraTesterMixin,
MemoryTesterMixin,
PipelineTesterMixin,
)


enable_full_determinism()


class HiDreamImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class HiDreamImagePipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = HiDreamImagePipeline
params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs", "prompt_embeds", "negative_prompt_embeds"}
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 = PipelineTesterMixin.required_optional_params
test_xformers_attention = False
test_layerwise_casting = True
required_input_params_in_call_signature = frozenset(
["prompt", "height", "width", "guidance_scale", "negative_prompt"]
)
batch_input_params = frozenset(["prompt", "negative_prompt"])
output_shape = (3, 128, 128)

def get_dummy_components(self):
torch.manual_seed(0)
Expand Down Expand Up @@ -107,7 +107,7 @@ def get_dummy_components(self):

scheduler = FlowMatchEulerDiscreteScheduler()

components = {
return {
"scheduler": scheduler,
"vae": vae,
"text_encoder": text_encoder,
Expand All @@ -120,42 +120,46 @@ def get_dummy_components(self):
"tokenizer_4": tokenizer_4,
"transformer": transformer,
}
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": "A painting of a squirrel eating a burger",
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 5.0,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
return inputs

def test_inference(self):
device = "cpu"

components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
class TestHiDreamImagePipeline(HiDreamImagePipelineTesterConfig, PipelineTesterMixin):
def test_inference(self):
# Run on CPU: the expected slice below is CPU-specific.
pipe = self.get_pipeline()

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

# fmt: off
expected_slice = np.array([0.3841, 0.5881, 0.3958, 0.5585, 0.4646, 0.4385, 0.4529, 0.4178, 0.6595, 0.2871, 0.7331, 0.5203, 0.5525, 0.5296, 0.4695, 0.5524])
expected_slice = torch.tensor([0.3841, 0.5585, 0.4529, 0.5123, 0.4314, 0.5122, 0.4885, 0.4554, 0.4653, 0.3369, 0.1799, 0.6126, 0.6202, 0.2871, 0.5525, 0.5524])
# fmt: on

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

def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=3e-4):
super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)


class TestHiDreamImagePipelineMemory(HiDreamImagePipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the HiDream Image pipeline."""


class TestHiDreamImagePipelineLoRA(HiDreamImagePipelineTesterConfig, LoraTesterMixin):
"""LoRA tests for the HiDream Image pipeline."""


def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(expected_max_diff=3e-4)
class TestHiDreamImagePipelineLoRAMemory(HiDreamImagePipelineTesterConfig, LoraMemoryTesterMixin):
"""LoRA x memory-optimization tests (group offload, CPU offload) for the HiDream Image pipeline."""
172 changes: 64 additions & 108 deletions tests/pipelines/hunyuan_image_21/test_hunyuanimage.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest

import numpy as np
import pytest
import torch
from transformers import (
ByT5Tokenizer,
Expand All @@ -33,36 +31,23 @@
HunyuanImageTransformer2DModel,
)

from ...testing_utils import enable_full_determinism
from ..test_pipelines_common import FirstBlockCacheTesterMixin, PipelineTesterMixin, to_np
from ...testing_utils import assert_tensors_close, enable_full_determinism
from ..testing_utils import (
BasePipelineTesterConfig,
FirstBlockCacheTesterMixin,
MemoryTesterMixin,
PipelineTesterMixin,
)


enable_full_determinism()


class HunyuanImagePipelineFastTests(
PipelineTesterMixin,
FirstBlockCacheTesterMixin,
unittest.TestCase,
):
class HunyuanImagePipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = HunyuanImagePipeline
params = frozenset(["prompt", "height", "width"])
batch_params = frozenset(["prompt", "negative_prompt"])
required_optional_params = frozenset(
[
"num_inference_steps",
"generator",
"latents",
"return_dict",
"callback_on_step_end",
"callback_on_step_end_tensor_inputs",
]
)

test_xformers_attention = False
test_layerwise_casting = True
test_group_offloading = True
test_attention_slicing = False
required_input_params_in_call_signature = frozenset(["prompt", "height", "width"])
batch_input_params = frozenset(["prompt", "negative_prompt"])
output_shape = (3, 16, 16)

def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1, guidance_embeds: bool = False):
torch.manual_seed(0)
Expand Down Expand Up @@ -153,7 +138,7 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1,
text_encoder_2 = T5EncoderModel(t5_config)
tokenizer_2 = ByT5Tokenizer()

components = {
return {
"transformer": transformer,
"vae": vae,
"scheduler": scheduler,
Expand All @@ -164,122 +149,93 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1,
"guider": guider,
"ocr_guider": ocr_guider,
}
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": "A painting of a squirrel eating a burger",
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 5,
"height": 16,
"width": 16,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
return inputs

def test_inference(self):
device = "cpu"

components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
class TestHunyuanImagePipeline(HunyuanImagePipelineTesterConfig, PipelineTesterMixin):
def test_inference(self):
# Run on CPU: the expected slice below is CPU-specific.
pipe = self.get_pipeline()

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

expected_slice_np = np.array(
[0.6252659, 0.51482046, 0.60799813, 0.59267783, 0.488082, 0.5857634, 0.523781, 0.58028054, 0.5674121]
)
output_slice = generated_image[0, -3:, -3:].flatten().cpu().numpy()
# fmt: off
expected_slice = torch.tensor([0.6252659, 0.51482046, 0.60799813, 0.59267783, 0.488082, 0.5857634, 0.523781, 0.58028054, 0.5674121])
# fmt: on

self.assertTrue(
np.abs(output_slice - expected_slice_np).max() < 1e-3,
f"output_slice: {output_slice}, expected_slice_np: {expected_slice_np}",
)
generated_slice = generated_image[0, -3:, -3:].flatten()
assert_tensors_close(generated_slice, expected_slice, atol=1e-3)

def test_inference_guider(self):
device = "cpu"

components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
# Run on CPU: the expected slice below is CPU-specific.
pipe = self.get_pipeline()

pipe.guider = pipe.guider.new(guidance_scale=1000)
pipe.ocr_guider = pipe.ocr_guider.new(guidance_scale=1000)

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

expected_slice_np = np.array(
[0.6068114, 0.48716035, 0.5984431, 0.60241306, 0.48849544, 0.5624479, 0.53696984, 0.58964247, 0.54248774]
)
output_slice = generated_image[0, -3:, -3:].flatten().cpu().numpy()
# fmt: off
expected_slice = torch.tensor([0.6068114, 0.48716035, 0.5984431, 0.60241306, 0.48849544, 0.5624479, 0.53696984, 0.58964247, 0.54248774])
# fmt: on

self.assertTrue(
np.abs(output_slice - expected_slice_np).max() < 1e-3,
f"output_slice: {output_slice}, expected_slice_np: {expected_slice_np}",
)
generated_slice = generated_image[0, -3:, -3:].flatten()
assert_tensors_close(generated_slice, expected_slice, atol=1e-3)

def test_inference_with_distilled_guidance(self):
device = "cpu"
# Run on CPU: the expected slice below is CPU-specific.
pipe = self.get_pipeline(**self.get_dummy_components(guidance_embeds=True))

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

inputs = self.get_dummy_inputs(device)
inputs["distilled_guidance_scale"] = 3.5
image = pipe(**inputs).images
image = pipe(**self.get_dummy_inputs(), distilled_guidance_scale=3.5).images
generated_image = image[0]
self.assertEqual(generated_image.shape, (3, 16, 16))
assert generated_image.shape == self.output_shape

expected_slice_np = np.array(
[0.63667065, 0.5187377, 0.66757566, 0.6320319, 0.4913387, 0.54813194, 0.5335031, 0.5736143, 0.5461346]
)
output_slice = generated_image[0, -3:, -3:].flatten().cpu().numpy()
# fmt: off
expected_slice = torch.tensor([0.63667065, 0.5187377, 0.66757566, 0.6320319, 0.4913387, 0.54813194, 0.5335031, 0.5736143, 0.5461346])
# fmt: on

self.assertTrue(
np.abs(output_slice - expected_slice_np).max() < 1e-3,
f"output_slice: {output_slice}, expected_slice_np: {expected_slice_np}",
)
generated_slice = generated_image[0, -3:, -3:].flatten()
assert_tensors_close(generated_slice, expected_slice, atol=1e-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.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()

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

# With tiling
pipe.vae.enable_tiling(tile_sample_min_size=96)
inputs = self.get_dummy_inputs(generator_device)
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",
output_with_tiling = self.run_pipe(pipe, height=128, width=128)

assert_tensors_close(
output_with_tiling,
output_without_tiling,
atol=expected_diff_max,
msg="VAE tiling should not affect the inference results.",
)

@unittest.skip("TODO: Test not supported for now because needs to be adjusted to work with guiders.")
@pytest.mark.skip("TODO: Test not supported for now because needs to be adjusted to work with guiders.")
def test_encode_prompt_works_in_isolation(self):
pass


class TestHunyuanImagePipelineMemory(HunyuanImagePipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the HunyuanImage pipeline."""


class TestHunyuanImagePipelineFirstBlockCache(HunyuanImagePipelineTesterConfig, FirstBlockCacheTesterMixin):
"""First Block Cache tests for the HunyuanImage pipeline."""
Loading
Loading