AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-23 09:09:01 +03:00
Code Activity

[tests] refactor o and p series pipeline tests (#14646)

Browse Source 
* refactor o and p series pipeline tests

* up
This commit is contained in:
Sayak Paul
2026-09-01 08:34:33 +05:30
committed by GitHub
parent a262497cd0
commit 61a11f2191
36 changed files with 1806 additions and 3105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
.ai/references/testing.md
View File

@@ -34,6 +34,10 @@ Follow the style introduced in [#14113](https://github.com/huggingface/diffusers
- `torch.nn.MultiheadAttention` is the common instance: it passes `self.out_proj.weight` straight to `torch.nn.functional.multi_head_attention_forward` instead of calling `self.out_proj`, so the hook on `out_proj` never fires. `SiglipVisionModel`'s attention pooling head wraps one — see `tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py`, whose `image_encoder` is excluded for this reason.
- `HunyuanDiTAttentionPool` (`src/diffusers/models/embeddings.py`) shows the same failure without an MHA module: a plain `nn.Module` that hands its `q_proj` / `k_proj` / `v_proj` / `c_proj` weights to `torch.nn.functional.multi_head_attention_forward`, so all four projections stay offloaded rather than just one. `HunyuanDiT2DModel` opts out of group offloading entirely with `_supports_group_offloading = False`.
- Before adding a skip or an exclusion, confirm the failure still reproduces — several existing skips are stale, having outlived the upstream cause.
- **A migration that surfaces a `src/` gap marks the test `xfail`, it does not patch the pipeline.** Give the marker a module-level name and a `reason` naming the exact gap (`PNDM_*` in `tests/pipelines/pndm/test_pndm.py` is the worked example), and prefer `strict=True` so the marker reports XPASS — and gets deleted — the day the pipeline is fixed. Use `strict=False` only when one mark covers a group whose members do not all fail. Marking a whole test class keeps the mixin's own marks (`@is_memory`, `@require_accelerator`) intact; overriding individual inherited tests drops the decorators they were declared with, so re-declare those too.
- **`from_pipe` tests** (a pipeline that is a variant of an existing one — PAG, AnimateDiff, ...) compose the shared `FromPipeTesterMixin` (`tests/pipelines/testing_utils/from_pipe.py`, exported from `..testing_utils`) in their own test class. It derives the original pipeline from `pipeline_class.__name__`; set `original_pipeline_repo` on the test class to pull it from a repo other than the default for that class. The unittest-era `PipelineFromPipeTesterMixin` in `tests/pipelines/test_pipelines_common.py` is what it replaces.
- **A hardware gap is a conditional skip, not an xfail.** When a test fails only because the runner's cuDNN build has no kernel for an op — `RuntimeError: GET was unable to find an engine to execute this computation`, as Sana's depthwise `Conv2d` hits in bfloat16 — wrap the call in `skip_if_no_cudnn_engine()` (`tests/testing_utils.py`). It skips on that error and re-raises every other `RuntimeError`, so the test still runs wherever the kernel exists.
- **PAG pipelines** also compose `PAGPipelineTesterMixin` (`tests/pipelines/pag/testing_utils.py`) in place of `PipelineTesterMixin`: it adds `test_pag_disable_enable` and `test_pag_inference` on top, driven by `base_pipeline_class` and the `pag_*` knobs on the test class. Keep `test_pag_applied_layers` per pipeline — which layers PAG resolves to is model-specific.
- **`encode_prompt` reading a component that isn't a text encoder or tokenizer?** `test_encode_prompt_works_in_isolation` rebuilds the pipeline with only the components whose names contain `text` or `tokenizer`. When `encode_prompt` also needs another one — a `processor` used for chat templating, say — list it in `text_stack_component_names` on the config class rather than re-implementing the test.
- **IP-Adapter tests** live in their own class decorated with `@is_ip_adapter`, subclassing only the config (not `PipelineTesterMixin`). UNet pipelines that load adapters through the standard `IPAdapterMixin` API compose the shared `IPAdapterTesterMixin` (`tests/pipelines/testing_utils/ip_adapter.py`, exported from `..testing_utils`); pipelines whose IP-Adapter API differs (Flux, for example) keep a bespoke mixin next to their own tests.

4
src/diffusers/pipelines/omnigen/pipeline_omnigen.py
View File

@@ -116,9 +116,7 @@ def retrieve_timesteps(
return timesteps, num_inference_steps
class OmniGenPipeline(
DiffusionPipeline,
):
class OmniGenPipeline(DiffusionPipeline):
r"""
The OmniGen pipeline for multimodal-to-image generation.

39
tests/pipelines/animatediff/test_animatediff.py
View File

@@ -10,7 +10,6 @@ from diffusers import (
AutoencoderKL,
DDIMScheduler,
MotionAdapter,
StableDiffusionPipeline,
UNet2DConditionModel,
)
@@ -22,8 +21,8 @@ from ...testing_utils import (
torch_device,
)
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineFromPipeTesterMixin
from ..testing_utils import (
FromPipeTesterMixin,
IPAdapterTesterMixin,
LoraMemoryTesterMixin,
LoraTesterMixin,
@@ -31,7 +30,6 @@ from ..testing_utils import (
UNetLoraTesterMixin,
)
from .testing_utils import (
FROM_PIPE_SKIP_REASON,
FreeInitTesterMixin,
FreeNoiseSplitInferenceTesterMixin,
MotionPipelineTesterConfig,
@@ -128,32 +126,6 @@ class TestAnimateDiffPipeline(
FreeInitTesterMixin,
FreeNoiseSplitInferenceTesterMixin,
):
def test_from_pipe_consistent_config(self):
original_repo = "hf-internal-testing/tinier-stable-diffusion-pipe"
# create StableDiffusionPipeline
pipe_original = StableDiffusionPipeline.from_pretrained(original_repo, requires_safety_checker=False)
# StableDiffusionPipeline -> AnimateDiffPipeline
pipe_components = self.get_dummy_components()
pipe_additional_components = {
name: component for name, component in pipe_components.items() if name not in pipe_original.components
}
pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components)
# AnimateDiffPipeline -> StableDiffusionPipeline
original_pipe_additional_components = {}
for name, component in pipe_original.components.items():
if name not in pipe.components or not isinstance(component, pipe.components[name].__class__):
original_pipe_additional_components[name] = component
pipe_original_2 = StableDiffusionPipeline.from_pipe(pipe, **original_pipe_additional_components)
# compare the config
original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")}
original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")}
assert original_config_2 == original_config
def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4):
if torch_device == "cpu" and expected_slice is None:
# fmt: off
@@ -244,10 +216,7 @@ class TestAnimateDiffPipelineIntegration:
assert numpy_cosine_similarity_distance(image_slice.flatten(), expected_slice.flatten()) < 1e-3
@pytest.mark.skip(FROM_PIPE_SKIP_REASON)
class TestAnimateDiffPipelineFromPipe(AnimateDiffPipelineTesterConfig, PipelineFromPipeTesterMixin):
"""`from_pipe` forward-pass parity and offload round trip for the AnimateDiff pipeline.
class TestAnimateDiffPipelineFromPipe(AnimateDiffPipelineTesterConfig, FromPipeTesterMixin):
"""`from_pipe` round-trip tests against `StableDiffusionPipeline` for the AnimateDiff pipeline."""
Parked, not deleted: `test_from_pipe_consistent_config` runs for real as a method on the main test class above,
but the forward-pass checks in `PipelineFromPipeTesterMixin` have no pytest-style equivalent yet.
"""
original_pipeline_repo = "hf-internal-testing/tinier-stable-diffusion-pipe"

42
tests/pipelines/animatediff/test_animatediff_controlnet.py
View File

@@ -1,4 +1,3 @@
import pytest
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
@@ -9,14 +8,13 @@ from diffusers import (
ControlNetModel,
DDIMScheduler,
MotionAdapter,
StableDiffusionPipeline,
UNet2DConditionModel,
)
from ...testing_utils import torch_device
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineFromPipeTesterMixin
from ..testing_utils import (
FromPipeTesterMixin,
IPAdapterTesterMixin,
LoraMemoryTesterMixin,
LoraTesterMixin,
@@ -24,7 +22,6 @@ from ..testing_utils import (
UNetLoraTesterMixin,
)
from .testing_utils import (
FROM_PIPE_SKIP_REASON,
FreeInitTesterMixin,
FreeNoiseTesterMixin,
MotionPipelineTesterConfig,
@@ -142,32 +139,6 @@ class TestAnimateDiffControlNetPipeline(
# `get_dummy_inputs` rather than by overriding `num_frames` on the returned dict.
return self.get_dummy_inputs(num_frames=16)
def test_from_pipe_consistent_config(self):
original_repo = "hf-internal-testing/tinier-stable-diffusion-pipe"
# create StableDiffusionPipeline
pipe_original = StableDiffusionPipeline.from_pretrained(original_repo, requires_safety_checker=False)
# StableDiffusionPipeline -> AnimateDiffControlNetPipeline
pipe_components = self.get_dummy_components()
pipe_additional_components = {
name: component for name, component in pipe_components.items() if name not in pipe_original.components
}
pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components)
# AnimateDiffControlNetPipeline -> StableDiffusionPipeline
original_pipe_additional_components = {}
for name, component in pipe_original.components.items():
if name not in pipe.components or not isinstance(component, pipe.components[name].__class__):
original_pipe_additional_components[name] = component
pipe_original_2 = StableDiffusionPipeline.from_pipe(pipe, **original_pipe_additional_components)
# compare the config
original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")}
original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")}
assert original_config_2 == original_config
def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4):
if torch_device == "cpu" and expected_slice is None:
# fmt: off
@@ -198,12 +169,7 @@ class TestAnimateDiffControlNetPipelineLoRAMemory(AnimateDiffControlNetPipelineT
"""LoRA x memory-optimization tests (group offload, CPU offload) for the pipeline."""
@pytest.mark.skip(FROM_PIPE_SKIP_REASON)
class TestAnimateDiffControlNetPipelineFromPipe(
AnimateDiffControlNetPipelineTesterConfig, PipelineFromPipeTesterMixin
):
"""`from_pipe` forward-pass parity and offload round trip for the AnimateDiff ControlNet pipeline.
class TestAnimateDiffControlNetPipelineFromPipe(AnimateDiffControlNetPipelineTesterConfig, FromPipeTesterMixin):
"""`from_pipe` round-trip tests against `StableDiffusionPipeline` for the AnimateDiff ControlNet pipeline."""
Parked, not deleted: `test_from_pipe_consistent_config` runs for real as a method on the main test class above,
but the forward-pass checks in `PipelineFromPipeTesterMixin` have no pytest-style equivalent yet.
"""
original_pipeline_repo = "hf-internal-testing/tinier-stable-diffusion-pipe"

40
tests/pipelines/animatediff/test_animatediff_sparsectrl.py
View File

@@ -1,4 +1,3 @@
import pytest
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
@@ -9,14 +8,13 @@ from diffusers import (
DDIMScheduler,
MotionAdapter,
SparseControlNetModel,
StableDiffusionPipeline,
UNet2DConditionModel,
)
from ...testing_utils import assert_tensors_close, torch_device
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineFromPipeTesterMixin
from ..testing_utils import (
FromPipeTesterMixin,
IPAdapterTesterMixin,
LoraMemoryTesterMixin,
LoraTesterMixin,
@@ -24,7 +22,6 @@ from ..testing_utils import (
UNetLoraTesterMixin,
)
from .testing_utils import (
FROM_PIPE_SKIP_REASON,
FreeInitTesterMixin,
MotionPipelineTesterConfig,
MotionPipelineTesterMixin,
@@ -138,32 +135,6 @@ class TestAnimateDiffSparseControlNetPipeline(
MotionPipelineTesterMixin,
FreeInitTesterMixin,
):
def test_from_pipe_consistent_config(self):
original_repo = "hf-internal-testing/tinier-stable-diffusion-pipe"
# create StableDiffusionPipeline
pipe_original = StableDiffusionPipeline.from_pretrained(original_repo, requires_safety_checker=False)
# StableDiffusionPipeline -> AnimateDiffSparseControlNetPipeline
pipe_components = self.get_dummy_components()
pipe_additional_components = {
name: component for name, component in pipe_components.items() if name not in pipe_original.components
}
pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components)
# AnimateDiffSparseControlNetPipeline -> StableDiffusionPipeline
original_pipe_additional_components = {}
for name, component in pipe_original.components.items():
if name not in pipe.components or not isinstance(component, pipe.components[name].__class__):
original_pipe_additional_components[name] = component
pipe_original_2 = StableDiffusionPipeline.from_pipe(pipe, **original_pipe_additional_components)
# compare the config
original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")}
original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")}
assert original_config_2 == original_config
def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4):
if torch_device == "cpu" and expected_slice is None:
# fmt: off
@@ -229,12 +200,9 @@ class TestAnimateDiffSparseControlNetPipelineLoRAMemory(
"""LoRA x memory-optimization tests (group offload, CPU offload) for the pipeline."""
@pytest.mark.skip(FROM_PIPE_SKIP_REASON)
class TestAnimateDiffSparseControlNetPipelineFromPipe(
AnimateDiffSparseControlNetPipelineTesterConfig, PipelineFromPipeTesterMixin
AnimateDiffSparseControlNetPipelineTesterConfig, FromPipeTesterMixin
):
"""`from_pipe` forward-pass parity and offload round trip for the AnimateDiff SparseControlNet pipeline.
"""`from_pipe` round-trip tests against `StableDiffusionPipeline` for the AnimateDiff SparseControlNet pipeline."""
Parked, not deleted: `test_from_pipe_consistent_config` runs for real as a method on the main test class above,
but the forward-pass checks in `PipelineFromPipeTesterMixin` have no pytest-style equivalent yet.
"""
original_pipeline_repo = "hf-internal-testing/tinier-stable-diffusion-pipe"

42
tests/pipelines/animatediff/test_animatediff_video2video.py
View File

@@ -1,4 +1,3 @@
import pytest
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
@@ -8,14 +7,13 @@ from diffusers import (
AutoencoderKL,
DDIMScheduler,
MotionAdapter,
StableDiffusionPipeline,
UNet2DConditionModel,
)
from ...testing_utils import torch_device
from ..pipeline_params import TEXT_TO_IMAGE_PARAMS, VIDEO_TO_VIDEO_BATCH_PARAMS
from ..test_pipelines_common import PipelineFromPipeTesterMixin
from ..testing_utils import (
FromPipeTesterMixin,
IPAdapterTesterMixin,
LoraMemoryTesterMixin,
LoraTesterMixin,
@@ -23,7 +21,6 @@ from ..testing_utils import (
UNetLoraTesterMixin,
)
from .testing_utils import (
FROM_PIPE_SKIP_REASON,
FreeInitTesterMixin,
FreeNoiseSplitInferenceTesterMixin,
MotionPipelineTesterConfig,
@@ -133,32 +130,6 @@ class TestAnimateDiffVideoToVideoPipeline(
inputs["strength"] = 0.5
return inputs
def test_from_pipe_consistent_config(self):
original_repo = "hf-internal-testing/tinier-stable-diffusion-pipe"
# create StableDiffusionPipeline
pipe_original = StableDiffusionPipeline.from_pretrained(original_repo, requires_safety_checker=False)
# StableDiffusionPipeline -> AnimateDiffVideoToVideoPipeline
pipe_components = self.get_dummy_components()
pipe_additional_components = {
name: component for name, component in pipe_components.items() if name not in pipe_original.components
}
pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components)
# AnimateDiffVideoToVideoPipeline -> StableDiffusionPipeline
original_pipe_additional_components = {}
for name, component in pipe_original.components.items():
if name not in pipe.components or not isinstance(component, pipe.components[name].__class__):
original_pipe_additional_components[name] = component
pipe_original_2 = StableDiffusionPipeline.from_pipe(pipe, **original_pipe_additional_components)
# compare the config
original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")}
original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")}
assert original_config_2 == original_config
def test_latent_inputs(self):
pipe = self.get_pipeline().to(torch_device)
@@ -191,12 +162,7 @@ class TestAnimateDiffVideoToVideoPipelineLoRAMemory(
"""LoRA x memory-optimization tests (group offload, CPU offload) for the pipeline."""
@pytest.mark.skip(FROM_PIPE_SKIP_REASON)
class TestAnimateDiffVideoToVideoPipelineFromPipe(
AnimateDiffVideoToVideoPipelineTesterConfig, PipelineFromPipeTesterMixin
):
"""`from_pipe` forward-pass parity and offload round trip for the AnimateDiff video-to-video pipeline.
class TestAnimateDiffVideoToVideoPipelineFromPipe(AnimateDiffVideoToVideoPipelineTesterConfig, FromPipeTesterMixin):
"""`from_pipe` round-trip tests against `StableDiffusionPipeline` for the AnimateDiff video-to-video pipeline."""
Parked, not deleted: `test_from_pipe_consistent_config` runs for real as a method on the main test class above,
but the forward-pass checks in `PipelineFromPipeTesterMixin` have no pytest-style equivalent yet.
"""
original_pipeline_repo = "hf-internal-testing/tinier-stable-diffusion-pipe"

40
tests/pipelines/animatediff/test_animatediff_video2video_controlnet.py
View File

@@ -1,4 +1,3 @@
import pytest
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
@@ -9,14 +8,13 @@ from diffusers import (
ControlNetModel,
DDIMScheduler,
MotionAdapter,
StableDiffusionPipeline,
UNet2DConditionModel,
)
from ...testing_utils import torch_device
from ..pipeline_params import TEXT_TO_IMAGE_PARAMS, VIDEO_TO_VIDEO_BATCH_PARAMS
from ..test_pipelines_common import PipelineFromPipeTesterMixin
from ..testing_utils import (
FromPipeTesterMixin,
IPAdapterTesterMixin,
LoraMemoryTesterMixin,
LoraTesterMixin,
@@ -24,7 +22,6 @@ from ..testing_utils import (
UNetLoraTesterMixin,
)
from .testing_utils import (
FROM_PIPE_SKIP_REASON,
FreeInitTesterMixin,
FreeNoiseTesterMixin,
MotionPipelineTesterConfig,
@@ -145,32 +142,6 @@ class TestAnimateDiffVideoToVideoControlNetPipeline(
# longer video rather than by passing `num_frames`.
return self.get_dummy_inputs(num_frames=16)
def test_from_pipe_consistent_config(self):
original_repo = "hf-internal-testing/tinier-stable-diffusion-pipe"
# create StableDiffusionPipeline
pipe_original = StableDiffusionPipeline.from_pretrained(original_repo, requires_safety_checker=False)
# StableDiffusionPipeline -> AnimateDiffVideoToVideoControlNetPipeline
pipe_components = self.get_dummy_components()
pipe_additional_components = {
name: component for name, component in pipe_components.items() if name not in pipe_original.components
}
pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components)
# AnimateDiffVideoToVideoControlNetPipeline -> StableDiffusionPipeline
original_pipe_additional_components = {}
for name, component in pipe_original.components.items():
if name not in pipe.components or not isinstance(component, pipe.components[name].__class__):
original_pipe_additional_components[name] = component
pipe_original_2 = StableDiffusionPipeline.from_pipe(pipe, **original_pipe_additional_components)
# compare the config
original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")}
original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")}
assert original_config_2 == original_config
def test_latent_inputs(self):
pipe = self.get_pipeline().to(torch_device)
@@ -213,12 +184,9 @@ class TestAnimateDiffVideoToVideoControlNetPipelineLoRAMemory(
"""LoRA x memory-optimization tests (group offload, CPU offload) for the pipeline."""
@pytest.mark.skip(FROM_PIPE_SKIP_REASON)
class TestAnimateDiffVideoToVideoControlNetPipelineFromPipe(
AnimateDiffVideoToVideoControlNetPipelineTesterConfig, PipelineFromPipeTesterMixin
AnimateDiffVideoToVideoControlNetPipelineTesterConfig, FromPipeTesterMixin
):
"""`from_pipe` forward-pass parity and offload round trip for the AnimateDiff video-to-video ControlNet pipeline.
"""`from_pipe` round-trip tests against `StableDiffusionPipeline` for the AnimateDiff video-to-video ControlNet pipeline."""
Parked, not deleted: `test_from_pipe_consistent_config` runs for real as a method on the main test class above,
but the forward-pass checks in `PipelineFromPipeTesterMixin` have no pytest-style equivalent yet.
"""
original_pipeline_repo = "hf-internal-testing/tinier-stable-diffusion-pipe"

13
tests/pipelines/animatediff/testing_utils.py
View File

@@ -24,19 +24,6 @@ from ..testing_utils import BasePipelineTesterConfig, PipelineTesterMixin
from ..testing_utils.common import BasePipelineOutputMixin
# `PipelineFromPipeTesterMixin` (tests/pipelines/test_pipelines_common.py) still covers `from_pipe` forward-pass
# parity and the model-CPU-offload round trip, but it is unittest-era: its tests call `self.get_dummy_inputs(device,
# seed=0)` and `self.assertLess`, neither of which exists on a `BasePipelineTesterConfig` outside a
# `unittest.TestCase`. Un-skipping the parked classes below without porting the mixin first would error, not fail.
# The mixin is still live for the ten `tests/pipelines/pag/` files and `stable_diffusion_adapter`; rewrite it
# pytest-style when those are migrated, then drop these skips.
FROM_PIPE_SKIP_REASON = (
"`PipelineFromPipeTesterMixin` is still unittest-style and cannot run against `BasePipelineTesterConfig` — "
"these error rather than fail if un-skipped. Port the mixin to pytest (due when `tests/pipelines/pag/` is "
"migrated), then remove this skip."
)
class MotionPipelineTesterConfig(BasePipelineTesterConfig):
"""`BasePipelineTesterConfig` for the AnimateDiff pipelines in this directory."""

53
tests/pipelines/omnigen/test_pipeline_omnigen.py
View File

@@ -1,7 +1,7 @@
import gc
import unittest
import numpy as np
import pytest
import torch
from transformers import AutoTokenizer
@@ -15,15 +15,14 @@ from ...testing_utils import (
slow,
torch_device,
)
from ..test_pipelines_common import PipelineTesterMixin
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
class OmniGenPipelineFastTests(unittest.TestCase, PipelineTesterMixin):
class OmniGenPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = OmniGenPipeline
params = frozenset(["prompt", "guidance_scale"])
batch_params = frozenset(["prompt"])
test_xformers_attention = False
test_layerwise_casting = True
required_input_params_in_call_signature = frozenset(["prompt", "guidance_scale"])
batch_input_params = frozenset(["prompt"])
output_shape = (3, 16, 16)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -54,53 +53,47 @@ class OmniGenPipelineFastTests(unittest.TestCase, PipelineTesterMixin):
scheduler = FlowMatchEulerDiscreteScheduler(invert_sigmas=True, num_train_timesteps=1)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
components = {
return {
"transformer": transformer,
"vae": vae,
"scheduler": scheduler,
"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="cpu").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": 1,
"guidance_scale": 3.0,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
"height": 16,
"width": 16,
}
return inputs
def test_inference(self):
pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
inputs = self.get_dummy_inputs(torch_device)
generated_image = pipe(**inputs).images[0]
class TestOmniGenPipeline(OmniGenPipelineTesterConfig, PipelineTesterMixin):
"""Core pipeline tests for OmniGen. The old `test_inference` only asserted the output shape, which
`PipelineTesterMixin.test_output` now covers against `output_shape`."""
self.assertEqual(generated_image.shape, (16, 16, 3))
class TestOmniGenPipelineMemory(OmniGenPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the OmniGen pipeline."""
@slow
@require_torch_accelerator
class OmniGenPipelineSlowTests(unittest.TestCase):
class TestOmniGenPipelineIntegration:
pipeline_class = OmniGenPipeline
repo_id = "shitao/OmniGen-v1-diffusers"
def setUp(self):
super().setUp()
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
yield
gc.collect()
backend_empty_cache(torch_device)

67
tests/pipelines/ovis_image/test_ovis_image.py
View File

@@ -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 Qwen2Tokenizer, Qwen3Config, Qwen3Model
@@ -27,29 +24,15 @@ from diffusers import (
)
from ...testing_utils import 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
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
class OvisImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class OvisImagePipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = OvisImagePipeline
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",
]
)
test_xformers_attention = False
test_layerwise_casting = True
test_group_offloading = True
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS
output_shape = (3, 16, 16)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -101,47 +84,37 @@ class OvisImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase):
"transformer": transformer,
}
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)
def get_dummy_inputs(self):
return {
"prompt": "a cat",
"negative_prompt": "bad quality",
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 2.0,
"height": 16,
"width": 16,
"max_sequence_length": 32,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
}
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)
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs).images
generated_image = image[0]
self.assertEqual(generated_image.shape, (16, 16, 3))
self.assertTrue(np.isfinite(image).all())
class TestOvisImagePipeline(OvisImagePipelineTesterConfig, PipelineTesterMixin):
def test_inference(self, base_pipe_output):
# `test_output` already pins the shape; this one guards against a NaN/inf output.
assert torch.isfinite(base_pipe_output).all()
def test_guidance_scale_is_set(self):
# The `guidance_scale` property reads `self._guidance_scale`, which `__call__` must initialize.
pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
inputs = self.get_dummy_inputs(torch_device)
pipe = self.get_pipeline().to(torch_device)
inputs = self.get_dummy_inputs()
pipe(**inputs)
assert pipe.guidance_scale == inputs["guidance_scale"]
def test_max_sequence_length_is_used(self):
# `max_sequence_length` should bound the encoded prompt length.
pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
pipe = self.get_pipeline().to(torch_device)
embeds_16 = pipe.encode_prompt(
"a cat", do_classifier_free_guidance=False, device=torch_device, max_sequence_length=16
)[0]
@@ -150,3 +123,7 @@ class OvisImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase):
)[0]
assert embeds_16.shape[1] == 16
assert embeds_32.shape[1] == 32
class TestOvisImagePipelineMemory(OvisImagePipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the OvisImage pipeline."""

465
tests/pipelines/pag/test_pag_animatediff.py
View File

@@ -1,7 +1,4 @@
import inspect
import unittest
import numpy as np
import pytest
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
@@ -10,49 +7,28 @@ from diffusers import (
AnimateDiffPipeline,
AutoencoderKL,
DDIMScheduler,
DPMSolverMultistepScheduler,
LCMScheduler,
MotionAdapter,
StableDiffusionPipeline,
UNet2DConditionModel,
UNetMotionModel,
)
from diffusers.models.attention import FreeNoiseTransformerBlock
from diffusers.utils import is_xformers_available
from ...testing_utils import require_accelerator, torch_device
from ...testing_utils import torch_device
from ..animatediff.testing_utils import (
FreeInitTesterMixin,
FreeNoiseTesterMixin,
MotionPipelineTesterConfig,
MotionPipelineTesterMixin,
)
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
IPAdapterTesterMixin,
PipelineFromPipeTesterMixin,
PipelineTesterMixin,
SDFunctionTesterMixin,
)
from ..testing_utils import FromPipeTesterMixin, IPAdapterTesterMixin, MemoryTesterMixin
from .testing_utils import PAGPipelineTesterMixin
def to_np(tensor):
if isinstance(tensor, torch.Tensor):
tensor = tensor.detach().cpu().numpy()
return tensor
class AnimateDiffPAGPipelineFastTests(
IPAdapterTesterMixin, SDFunctionTesterMixin, PipelineTesterMixin, PipelineFromPipeTesterMixin, unittest.TestCase
):
class AnimateDiffPAGPipelineTesterConfig(MotionPipelineTesterConfig):
pipeline_class = AnimateDiffPAGPipeline
params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_params = TEXT_TO_IMAGE_BATCH_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 = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS
# `num_frames` defaults to 16; height/width default to `unet.sample_size * vae_scale_factor` (8 * 2).
output_shape = (16, 3, 16, 16)
def get_dummy_components(self):
cross_attention_dim = 8
@@ -100,6 +76,7 @@ class AnimateDiffPAGPipelineFastTests(
)
text_encoder = CLIPTextModel(text_encoder_config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
torch.manual_seed(0)
motion_adapter = MotionAdapter(
block_out_channels=block_out_channels,
motion_layers_per_block=2,
@@ -107,7 +84,7 @@ class AnimateDiffPAGPipelineFastTests(
motion_num_attention_heads=4,
)
components = {
return {
"unet": unet,
"scheduler": scheduler,
"vae": vae,
@@ -117,375 +94,51 @@ class AnimateDiffPAGPipelineFastTests(
"feature_extractor": None,
"image_encoder": None,
}
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": 7.5,
"pag_scale": 3.0,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
return inputs
def test_from_pipe_consistent_config(self):
assert self.original_pipeline_class == StableDiffusionPipeline
original_repo = "hf-internal-testing/tinier-stable-diffusion-pipe"
original_kwargs = {"requires_safety_checker": False}
# create original_pipeline_class(sd)
pipe_original = self.original_pipeline_class.from_pretrained(original_repo, **original_kwargs)
class TestAnimateDiffPAGPipeline(
AnimateDiffPAGPipelineTesterConfig,
MotionPipelineTesterMixin,
PAGPipelineTesterMixin,
FreeInitTesterMixin,
FreeNoiseTesterMixin,
):
base_pipeline_class = AnimateDiffPipeline
# AnimateDiff's PAG layers resolve through the motion modules, so the "PAG enabled" leg keeps the pipeline
# default layers and just leaves `pag_scale` at the dummy inputs' 3.0.
pag_enabled_applied_layers = None
# original_pipeline_class(sd) -> pipeline_class
pipe_components = self.get_dummy_components()
pipe_additional_components = {}
for name, component in pipe_components.items():
if name not in pipe_original.components:
pipe_additional_components[name] = component
pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components)
# pipeline_class -> original_pipeline_class(sd)
original_pipe_additional_components = {}
for name, component in pipe_original.components.items():
if name not in pipe.components or not isinstance(component, pipe.components[name].__class__):
original_pipe_additional_components[name] = component
pipe_original_2 = self.original_pipeline_class.from_pipe(pipe, **original_pipe_additional_components)
# compare the config
original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")}
original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")}
assert original_config_2 == original_config
def test_motion_unet_loading(self):
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
assert isinstance(pipe.unet, UNetMotionModel)
@unittest.skip("Attention slicing is not enabled in this pipeline")
def test_attention_slicing_forward_pass(self):
@pytest.mark.skip(
"`AnimateDiffPAGPipeline.check_inputs` rejects a dict `prompt`, so FreeNoise's per-frame prompts cannot be "
"passed to it (the non-PAG `AnimateDiffPipeline` accepts them)."
)
def test_free_noise_multi_prompt(self):
pass
def test_ip_adapter(self):
expected_pipe_slice = None
if torch_device == "cpu":
expected_pipe_slice = np.array(
[
0.5254,
0.5844,
0.4705,
0.4952,
0.4953,
0.5932,
0.5227,
0.4285,
0.5430,
0.4903,
0.4027,
0.4942,
0.3771,
0.3899,
0.6014,
0.4469,
0.4935,
0.5543,
0.5867,
0.5314,
0.3727,
0.5183,
0.6138,
0.5170,
0.5111,
0.4725,
0.5741,
]
)
return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice)
def test_dict_tuple_outputs_equivalent(self):
expected_slice = None
if torch_device == "cpu":
expected_slice = np.array([0.5227, 0.4285, 0.5430, 0.4469, 0.4935, 0.5543, 0.5111, 0.4725, 0.5741])
return super().test_dict_tuple_outputs_equivalent(expected_slice=expected_slice)
@require_accelerator
def test_to_device(self):
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
pipe.to("cpu")
# pipeline creates a new motion UNet under the hood. So we need to check the device from pipe.components
model_devices = [
component.device.type for component in pipe.components.values() if hasattr(component, "device")
]
self.assertTrue(all(device == "cpu" for device in model_devices))
output_cpu = pipe(**self.get_dummy_inputs("cpu"))[0]
self.assertTrue(np.isnan(output_cpu).sum() == 0)
pipe.to(torch_device)
model_devices = [
component.device.type for component in pipe.components.values() if hasattr(component, "device")
]
self.assertTrue(all(device == torch_device for device in model_devices))
output_device = pipe(**self.get_dummy_inputs(torch_device))[0]
self.assertTrue(np.isnan(to_np(output_device)).sum() == 0)
def test_to_dtype(self):
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
# pipeline creates a new motion UNet under the hood. So we need to check the dtype from pipe.components
model_dtypes = [component.dtype for component in pipe.components.values() if hasattr(component, "dtype")]
self.assertTrue(all(dtype == torch.float32 for dtype in model_dtypes))
pipe.to(dtype=torch.float16)
model_dtypes = [component.dtype for component in pipe.components.values() if hasattr(component, "dtype")]
self.assertTrue(all(dtype == torch.float16 for dtype in model_dtypes))
def test_prompt_embeds(self):
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
pipe.to(torch_device)
inputs = self.get_dummy_inputs(torch_device)
inputs.pop("prompt")
inputs["prompt_embeds"] = torch.randn((1, 4, pipe.text_encoder.config.hidden_size), device=torch_device)
pipe(**inputs)
def test_free_init(self):
components = self.get_dummy_components()
pipe: AnimateDiffPAGPipeline = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
pipe.to(torch_device)
inputs_normal = self.get_dummy_inputs(torch_device)
frames_normal = pipe(**inputs_normal).frames[0]
pipe.enable_free_init(
num_iters=2,
use_fast_sampling=True,
method="butterworth",
order=4,
spatial_stop_frequency=0.25,
temporal_stop_frequency=0.25,
def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4):
# Re-recorded: `get_dummy_components` now seeds the `MotionAdapter` like every other component, which it
# did not before the migration, so the adapter's weights (and this slice) changed.
if torch_device == "cpu" and expected_slice is None:
# fmt: off
expected_slice = torch.tensor([0.5132, 0.4380, 0.5327, 0.4619, 0.4955, 0.5457, 0.4980, 0.5015, 0.5652])
# fmt: on
super().test_dict_tuple_outputs_equivalent(
expected_slice=expected_slice, expected_max_difference=expected_max_difference
)
inputs_enable_free_init = self.get_dummy_inputs(torch_device)
frames_enable_free_init = pipe(**inputs_enable_free_init).frames[0]
pipe.disable_free_init()
inputs_disable_free_init = self.get_dummy_inputs(torch_device)
frames_disable_free_init = pipe(**inputs_disable_free_init).frames[0]
sum_enabled = np.abs(to_np(frames_normal) - to_np(frames_enable_free_init)).sum()
max_diff_disabled = np.abs(to_np(frames_normal) - to_np(frames_disable_free_init)).max()
self.assertGreater(
sum_enabled, 1e1, "Enabling of FreeInit should lead to results different from the default pipeline results"
)
self.assertLess(
max_diff_disabled,
1e-3,
"Disabling of FreeInit should lead to results similar to the default pipeline results",
)
def test_free_init_with_schedulers(self):
components = self.get_dummy_components()
pipe: AnimateDiffPAGPipeline = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
pipe.to(torch_device)
inputs_normal = self.get_dummy_inputs(torch_device)
frames_normal = pipe(**inputs_normal).frames[0]
schedulers_to_test = [
DPMSolverMultistepScheduler.from_config(
components["scheduler"].config,
timestep_spacing="linspace",
beta_schedule="linear",
algorithm_type="dpmsolver++",
steps_offset=1,
clip_sample=False,
),
LCMScheduler.from_config(
components["scheduler"].config,
timestep_spacing="linspace",
beta_schedule="linear",
steps_offset=1,
clip_sample=False,
),
]
components.pop("scheduler")
for scheduler in schedulers_to_test:
components["scheduler"] = scheduler
pipe: AnimateDiffPAGPipeline = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
pipe.to(torch_device)
pipe.enable_free_init(num_iters=2, use_fast_sampling=False)
inputs = self.get_dummy_inputs(torch_device)
frames_enable_free_init = pipe(**inputs).frames[0]
sum_enabled = np.abs(to_np(frames_normal) - to_np(frames_enable_free_init)).sum()
self.assertGreater(
sum_enabled,
1e1,
"Enabling of FreeInit should lead to results different from the default pipeline results",
)
def test_free_noise_blocks(self):
components = self.get_dummy_components()
pipe: AnimateDiffPAGPipeline = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
pipe.to(torch_device)
pipe.enable_free_noise()
for block in pipe.unet.down_blocks:
for motion_module in block.motion_modules:
for transformer_block in motion_module.transformer_blocks:
self.assertTrue(
isinstance(transformer_block, FreeNoiseTransformerBlock),
"Motion module transformer blocks must be an instance of `FreeNoiseTransformerBlock` after enabling FreeNoise.",
)
pipe.disable_free_noise()
for block in pipe.unet.down_blocks:
for motion_module in block.motion_modules:
for transformer_block in motion_module.transformer_blocks:
self.assertFalse(
isinstance(transformer_block, FreeNoiseTransformerBlock),
"Motion module transformer blocks must not be an instance of `FreeNoiseTransformerBlock` after disabling FreeNoise.",
)
def test_free_noise(self):
components = self.get_dummy_components()
pipe: AnimateDiffPAGPipeline = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
pipe.to(torch_device)
inputs_normal = self.get_dummy_inputs(torch_device)
frames_normal = pipe(**inputs_normal).frames[0]
for context_length in [8, 9]:
for context_stride in [4, 6]:
pipe.enable_free_noise(context_length, context_stride)
inputs_enable_free_noise = self.get_dummy_inputs(torch_device)
frames_enable_free_noise = pipe(**inputs_enable_free_noise).frames[0]
pipe.disable_free_noise()
inputs_disable_free_noise = self.get_dummy_inputs(torch_device)
frames_disable_free_noise = pipe(**inputs_disable_free_noise).frames[0]
sum_enabled = np.abs(to_np(frames_normal) - to_np(frames_enable_free_noise)).sum()
max_diff_disabled = np.abs(to_np(frames_normal) - to_np(frames_disable_free_noise)).max()
self.assertGreater(
sum_enabled,
1e1,
"Enabling of FreeNoise should lead to results different from the default pipeline results",
)
self.assertLess(
max_diff_disabled,
1e-4,
"Disabling of FreeNoise should lead to results similar to the default pipeline results",
)
@unittest.skipIf(
torch_device != "cuda" or not is_xformers_available(),
reason="XFormers attention is only available with CUDA and `xformers` installed",
)
def test_xformers_attention_forwardGenerator_pass(self):
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)
inputs = self.get_dummy_inputs(torch_device)
output_without_offload = pipe(**inputs).frames[0]
output_without_offload = (
output_without_offload.cpu() if torch.is_tensor(output_without_offload) else output_without_offload
)
pipe.enable_xformers_memory_efficient_attention()
inputs = self.get_dummy_inputs(torch_device)
output_with_offload = pipe(**inputs).frames[0]
output_with_offload = (
output_with_offload.cpu() if torch.is_tensor(output_with_offload) else output_without_offload
)
max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max()
self.assertLess(max_diff, 1e-4, "XFormers attention should not affect the inference results")
def test_vae_slicing(self):
return super().test_vae_slicing(image_count=2)
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline (expect same output when pag is disabled)
components.pop("pag_applied_layers", None)
pipe_sd = AnimateDiffPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
)
out = pipe_sd(**inputs).frames[0, -3:, -3:, -1]
components = self.get_dummy_components()
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).frames[0, -3:, -3:, -1]
# pag enabled
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
out_pag_enabled = pipe_pag(**inputs).frames[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3
def test_pag_applied_layers(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline
components.pop("pag_applied_layers", None)
pipe = self.pipeline_class(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()
# pag_applied_layers = ["mid","up","down"] should apply to all self-attention layers
# Note that for motion modules in AnimateDiff, both attn1 and attn2 are self-attention
@@ -526,7 +179,7 @@ class AnimateDiffPAGPipelineFastTests(
pipe.unet.set_attn_processor(original_attn_procs.copy())
pag_layers = ["mid_block.attentions.1"]
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
# pag_applied_layers = "down" should apply to all self-attention layers in down_blocks
@@ -551,13 +204,19 @@ class AnimateDiffPAGPipelineFastTests(
pipe.unet.set_attn_processor(original_attn_procs.copy())
pag_layers = ["motion_modules.42"]
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
def test_encode_prompt_works_in_isolation(self):
extra_required_param_value_dict = {
"device": torch.device(torch_device).type,
"num_images_per_prompt": 1,
"do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0,
}
return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict)
class TestAnimateDiffPAGPipelineMemory(AnimateDiffPAGPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the AnimateDiff PAG pipeline."""
class TestAnimateDiffPAGPipelineIPAdapter(AnimateDiffPAGPipelineTesterConfig, IPAdapterTesterMixin):
"""IP-Adapter tests for the AnimateDiff PAG pipeline."""
class TestAnimateDiffPAGPipelineFromPipe(AnimateDiffPAGPipelineTesterConfig, FromPipeTesterMixin):
"""`from_pipe` round-trip tests against `StableDiffusionPipeline`."""
original_pipeline_repo = "hf-internal-testing/tinier-stable-diffusion-pipe"

170
tests/pipelines/pag/test_pag_controlnet_sd.py
View File

@@ -13,10 +13,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import unittest
import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
@@ -30,37 +26,30 @@ from diffusers import (
)
from diffusers.utils.torch_utils import randn_tensor
from ...testing_utils import enable_full_determinism, torch_device
from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device
from ..pipeline_params import (
TEXT_TO_IMAGE_BATCH_PARAMS,
TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS,
TEXT_TO_IMAGE_IMAGE_PARAMS,
TEXT_TO_IMAGE_PARAMS,
)
from ..test_pipelines_common import (
from ..testing_utils import (
BasePipelineTesterConfig,
FromPipeTesterMixin,
IPAdapterTesterMixin,
PipelineFromPipeTesterMixin,
PipelineLatentTesterMixin,
PipelineTesterMixin,
MemoryTesterMixin,
)
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class StableDiffusionControlNetPAGPipelineFastTests(
PipelineTesterMixin,
IPAdapterTesterMixin,
PipelineLatentTesterMixin,
PipelineFromPipeTesterMixin,
unittest.TestCase,
):
class StableDiffusionControlNetPAGPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = StableDiffusionControlNetPAGPipeline
params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS
callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union({"add_text_embeds", "add_time_ids"})
output_shape = (3, 64, 64)
def get_dummy_components(self, time_cond_proj_dim=None):
# Copied from tests.pipelines.controlnet.test_controlnet_sdxl.StableDiffusionXLControlNetPipelineFastTests.get_dummy_components
@@ -120,7 +109,7 @@ class StableDiffusionControlNetPAGPipelineFastTests(
text_encoder = CLIPTextModel(text_encoder_config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
components = {
return {
"unet": unet,
"controlnet": controlnet,
"scheduler": scheduler,
@@ -131,123 +120,88 @@ class StableDiffusionControlNetPAGPipelineFastTests(
"feature_extractor": None,
"image_encoder": None,
}
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)
def get_dummy_inputs(self):
generator = self.get_generator(0)
controlnet_embedder_scale_factor = 2
# The conditioning image is drawn from the same generator, which is then handed to the pipeline in the
# state that leaves it — the expected slices below were recorded that way.
image = randn_tensor(
(1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor),
generator=generator,
device=torch.device(device),
device=torch.device("cpu"),
)
inputs = {
return {
"prompt": "A painting of a squirrel eating a burger",
"generator": generator,
"num_inference_steps": 2,
"guidance_scale": 6.0,
"pag_scale": 3.0,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
"image": image,
}
return inputs
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline (expect same output when pag is disabled)
pipe_sd = StableDiffusionControlNetPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
# pag enabled
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3
class TestStableDiffusionControlNetPAGPipeline(
StableDiffusionControlNetPAGPipelineTesterConfig, PAGPipelineTesterMixin
):
base_pipeline_class = StableDiffusionControlNetPipeline
def test_pag_cfg(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# Run on CPU: the expected slice below is CPU-specific.
pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"])
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
64,
64,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array(
[0.4533307, 0.28746086, 0.16543725, 0.786471, 0.5409842, 0.40167668, 0.7050086, 0.69445, 0.42238155]
image = pipe_pag(**self.get_dummy_inputs())[0]
assert image.shape == (1, *self.output_shape), (
f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}"
)
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}"
# fmt: off
expected_slice = torch.tensor([0.4533307, 0.28746086, 0.16543725, 0.786471, 0.5409842, 0.40167668, 0.7050086, 0.69445, 0.42238155])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
def test_pag_uncond(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# Run on CPU: the expected slice below is CPU-specific.
pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"])
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs["guidance_scale"] = 0.0
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
64,
64,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array(
[0.45054412, 0.27958393, 0.15983358, 0.80098593, 0.5432001, 0.4018666, 0.7127423, 0.69932824, 0.42260933]
image = pipe_pag(**inputs)[0]
assert image.shape == (1, *self.output_shape), (
f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}"
)
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}"
# fmt: off
expected_slice = torch.tensor([0.45054412, 0.27958393, 0.15983358, 0.80098593, 0.5432001, 0.4018666, 0.7127423, 0.69932824, 0.42260933])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
def test_encode_prompt_works_in_isolation(self):
extra_required_param_value_dict = {
"device": torch.device(torch_device).type,
"do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0,
"do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0,
}
return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict)
class TestStableDiffusionControlNetPAGPipelineMemory(
StableDiffusionControlNetPAGPipelineTesterConfig, MemoryTesterMixin
):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD ControlNet PAG pipeline."""
class TestStableDiffusionControlNetPAGPipelineIPAdapter(
StableDiffusionControlNetPAGPipelineTesterConfig, IPAdapterTesterMixin
):
"""IP-Adapter tests for the SD ControlNet PAG pipeline."""
class TestStableDiffusionControlNetPAGPipelineFromPipe(
StableDiffusionControlNetPAGPipelineTesterConfig, FromPipeTesterMixin
):
"""`from_pipe` round-trip tests against `StableDiffusionPipeline`."""

149
tests/pipelines/pag/test_pag_controlnet_sd_inpaint.py
View File

@@ -15,9 +15,7 @@
# This model implementation is heavily based on:
import inspect
import random
import unittest
import numpy as np
import torch
@@ -34,26 +32,24 @@ from diffusers import (
)
from diffusers.utils.torch_utils import randn_tensor
from ...testing_utils import enable_full_determinism, floats_tensor, torch_device
from ...testing_utils import assert_tensors_close, enable_full_determinism, floats_tensor, torch_device
from ..pipeline_params import (
TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
TEXT_TO_IMAGE_IMAGE_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class StableDiffusionControlNetPAGInpaintPipelineFastTests(
PipelineLatentTesterMixin, PipelineKarrasSchedulerTesterMixin, PipelineTesterMixin, unittest.TestCase
):
class StableDiffusionControlNetPAGInpaintPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = StableDiffusionControlNetPAGInpaintPipeline
params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
image_params = frozenset({"control_image"}) # skip `image` and `mask` for now, only test for control_image
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
batch_input_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
# The output resolution follows the 64x64 input image.
output_shape = (3, 64, 64)
def get_dummy_components(self):
# Copied from tests.pipelines.controlnet.test_controlnet_inpaint.ControlNetInpaintPipelineFastTests.get_dummy_components
@@ -109,7 +105,7 @@ class StableDiffusionControlNetPAGInpaintPipelineFastTests(
text_encoder = CLIPTextModel(text_encoder_config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
components = {
return {
"unet": unet,
"controlnet": controlnet,
"scheduler": scheduler,
@@ -120,130 +116,83 @@ class StableDiffusionControlNetPAGInpaintPipelineFastTests(
"feature_extractor": None,
"image_encoder": None,
}
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)
def get_dummy_inputs(self):
generator = self.get_generator(0)
controlnet_embedder_scale_factor = 2
# The control image is drawn from the same generator, which is then handed to the pipeline in the state
# that leaves it — the expected slices below were recorded that way.
control_image = randn_tensor(
(1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor),
generator=generator,
device=torch.device(device),
device=torch.device("cpu"),
)
init_image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device)
init_image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device)
init_image = init_image.cpu().permute(0, 2, 3, 1)[0]
image = Image.fromarray(np.uint8(init_image)).convert("RGB").resize((64, 64))
mask_image = Image.fromarray(np.uint8(init_image + 4)).convert("RGB").resize((64, 64))
inputs = {
return {
"prompt": "A painting of a squirrel eating a burger",
"generator": generator,
"num_inference_steps": 2,
"guidance_scale": 6.0,
"pag_scale": 3.0,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
"image": image,
"mask_image": mask_image,
"control_image": control_image,
}
return inputs
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline (expect same output when pag is disabled)
pipe_sd = StableDiffusionControlNetInpaintPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__calss__.__name__}."
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
# pag enabled
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3
class TestStableDiffusionControlNetPAGInpaintPipeline(
StableDiffusionControlNetPAGInpaintPipelineTesterConfig, PAGPipelineTesterMixin
):
base_pipeline_class = StableDiffusionControlNetInpaintPipeline
def test_pag_cfg(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# Run on CPU: the expected slice below is CPU-specific.
pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"])
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
64,
64,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array(
[0.7277897, 0.61666954, 0.54722667, 0.595576, 0.593909, 0.56389576, 0.41761285, 0.50566983, 0.49766505]
image = pipe_pag(**self.get_dummy_inputs())[0]
assert image.shape == (1, *self.output_shape), (
f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}"
)
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}"
# fmt: off
expected_slice = torch.tensor([0.7277897, 0.61666954, 0.54722667, 0.595576, 0.593909, 0.56389576, 0.41761285, 0.50566983, 0.49766505])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
def test_pag_uncond(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# Run on CPU: the expected slice below is CPU-specific.
pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"])
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs["guidance_scale"] = 0.0
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
64,
64,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array(
[0.7349223, 0.60567534, 0.5428778, 0.6091342, 0.60273147, 0.57611704, 0.42401767, 0.5064247, 0.49535546]
image = pipe_pag(**inputs)[0]
assert image.shape == (1, *self.output_shape), (
f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}"
)
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}"
# fmt: off
expected_slice = torch.tensor([0.7349223, 0.60567534, 0.5428778, 0.6091342, 0.60273147, 0.57611704, 0.42401767, 0.5064247, 0.49535546])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
def test_encode_prompt_works_in_isolation(self):
extra_required_param_value_dict = {
"device": torch.device(torch_device).type,
"do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0,
"do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0,
}
return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict)
class TestStableDiffusionControlNetPAGInpaintPipelineMemory(
StableDiffusionControlNetPAGInpaintPipelineTesterConfig, MemoryTesterMixin
):
"""Memory tests (CPU offload, group offload, layerwise casting) for the SD ControlNet PAG inpaint pipeline."""

168
tests/pipelines/pag/test_pag_controlnet_sdxl.py
View File

@@ -13,10 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import unittest
import numpy as np
import pytest
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
@@ -30,37 +27,30 @@ from diffusers import (
)
from diffusers.utils.torch_utils import randn_tensor
from ...testing_utils import enable_full_determinism
from ...testing_utils import assert_tensors_close, enable_full_determinism
from ..pipeline_params import (
TEXT_TO_IMAGE_BATCH_PARAMS,
TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS,
TEXT_TO_IMAGE_IMAGE_PARAMS,
TEXT_TO_IMAGE_PARAMS,
)
from ..test_pipelines_common import (
from ..testing_utils import (
BasePipelineTesterConfig,
FromPipeTesterMixin,
IPAdapterTesterMixin,
PipelineFromPipeTesterMixin,
PipelineLatentTesterMixin,
PipelineTesterMixin,
MemoryTesterMixin,
)
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class StableDiffusionXLControlNetPAGPipelineFastTests(
PipelineTesterMixin,
IPAdapterTesterMixin,
PipelineLatentTesterMixin,
PipelineFromPipeTesterMixin,
unittest.TestCase,
):
class StableDiffusionXLControlNetPAGPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = StableDiffusionXLControlNetPAGPipeline
params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS
callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union({"add_text_embeds", "add_time_ids"})
output_shape = (3, 64, 64)
def get_dummy_components(self, time_cond_proj_dim=None):
# Copied from tests.pipelines.controlnet.test_controlnet_sdxl.StableDiffusionXLControlNetPipelineFastTests.get_dummy_components
@@ -151,114 +141,84 @@ class StableDiffusionXLControlNetPAGPipelineFastTests(
}
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)
def get_dummy_inputs(self):
generator = self.get_generator(0)
controlnet_embedder_scale_factor = 2
# The conditioning image is drawn from the same generator, which is then handed to the pipeline in the
# state that leaves it — the expected slices below were recorded that way.
image = randn_tensor(
(1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor),
generator=generator,
device=torch.device(device),
device=torch.device("cpu"),
)
inputs = {
return {
"prompt": "A painting of a squirrel eating a burger",
"generator": generator,
"num_inference_steps": 2,
"guidance_scale": 6.0,
"pag_scale": 3.0,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
"image": image,
}
return inputs
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline (expect same output when pag is disabled)
pipe_sd = StableDiffusionXLControlNetPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
# pag enabled
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3
class TestStableDiffusionXLControlNetPAGPipeline(
StableDiffusionXLControlNetPAGPipelineTesterConfig, PAGPipelineTesterMixin
):
base_pipeline_class = StableDiffusionXLControlNetPipeline
def test_pag_cfg(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# Run on CPU: the expected slice below is CPU-specific.
pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"])
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
image = pipe_pag(**self.get_dummy_inputs())[0]
assert image.shape == (1, *self.output_shape), (
f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}"
)
inputs = self.get_dummy_inputs(device)
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
64,
64,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array([0.6864, 0.5436, 0.5644, 0.6136, 0.5541, 0.5910, 0.4519, 0.4634, 0.5252])
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}"
# fmt: off
expected_slice = torch.tensor([0.6864, 0.5436, 0.5644, 0.6136, 0.5541, 0.5910, 0.4519, 0.4634, 0.5252])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
def test_pag_uncond(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# Run on CPU: the expected slice below is CPU-specific.
pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"])
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs["guidance_scale"] = 0.0
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
image = pipe_pag(**inputs)[0]
assert image.shape == (1, *self.output_shape), (
f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}"
)
assert image.shape == (
1,
64,
64,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array([0.6843, 0.5381, 0.5675, 0.6109, 0.5493, 0.5988, 0.4477, 0.4679, 0.5242])
# fmt: off
expected_slice = torch.tensor([0.6843, 0.5381, 0.5675, 0.6109, 0.5493, 0.5988, 0.4477, 0.4679, 0.5242])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}"
@unittest.skip("We test this functionality elsewhere already.")
@pytest.mark.skip("We test this functionality elsewhere already.")
def test_save_load_optional_components(self):
pass
class TestStableDiffusionXLControlNetPAGPipelineMemory(
StableDiffusionXLControlNetPAGPipelineTesterConfig, MemoryTesterMixin
):
"""Memory tests (CPU offload, group offload, layerwise casting) for the SDXL ControlNet PAG pipeline."""
class TestStableDiffusionXLControlNetPAGPipelineIPAdapter(
StableDiffusionXLControlNetPAGPipelineTesterConfig, IPAdapterTesterMixin
):
"""IP-Adapter tests for the SDXL ControlNet PAG pipeline."""
class TestStableDiffusionXLControlNetPAGPipelineFromPipe(
StableDiffusionXLControlNetPAGPipelineTesterConfig, FromPipeTesterMixin
):
"""`from_pipe` round-trip tests against `StableDiffusionXLPipeline`."""

171
tests/pipelines/pag/test_pag_controlnet_sdxl_img2img.py
View File

@@ -13,11 +13,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import random
import unittest
import numpy as np
import pytest
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
@@ -30,38 +28,34 @@ from diffusers import (
UNet2DConditionModel,
)
from ...testing_utils import enable_full_determinism, floats_tensor
from ...testing_utils import assert_tensors_close, enable_full_determinism, floats_tensor, torch_device
from ..pipeline_params import (
IMAGE_TO_IMAGE_IMAGE_PARAMS,
TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS,
)
from ..test_pipelines_common import (
from ..testing_utils import (
BasePipelineTesterConfig,
FromPipeTesterMixin,
IPAdapterTesterMixin,
PipelineFromPipeTesterMixin,
PipelineLatentTesterMixin,
PipelineTesterMixin,
MemoryTesterMixin,
)
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class StableDiffusionXLControlNetPAGImg2ImgPipelineFastTests(
IPAdapterTesterMixin,
PipelineLatentTesterMixin,
PipelineTesterMixin,
PipelineFromPipeTesterMixin,
unittest.TestCase,
):
class StableDiffusionXLControlNetPAGImg2ImgPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = StableDiffusionXLControlNetPAGImg2ImgPipeline
params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS
image_latents_params = IMAGE_TO_IMAGE_IMAGE_PARAMS
required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union(
{"pag_scale", "pag_adaptive_scale"}
)
batch_input_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union(
{"add_text_embeds", "add_time_ids", "add_neg_time_ids"}
)
# The output resolution follows the 64x64 conditioning image.
output_shape = (3, 64, 64)
# Copied from tests.pipelines.controlnet.test_controlnet_sdxl_img2img.ControlNetPipelineSDXLImg2ImgFastTests.get_dummy_components
def get_dummy_components(self, skip_first_text_encoder=False):
@@ -153,117 +147,80 @@ class StableDiffusionXLControlNetPAGImg2ImgPipelineFastTests(
# based on tests.pipelines.controlnet.test_controlnet_sdxl_img2img.ControlNetPipelineSDXLImg2ImgFastTests.get_dummy_inputs
# add `pag_scale` to the inputs
def get_dummy_inputs(self, device, seed=0):
def get_dummy_inputs(self):
controlnet_embedder_scale_factor = 2
image = floats_tensor(
(1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor),
rng=random.Random(seed),
).to(device)
rng=random.Random(0),
).to(torch_device)
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device=device).manual_seed(seed)
inputs = {
return {
"prompt": "A painting of a squirrel eating a burger",
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 6.0,
"pag_scale": 3.0,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
"image": image,
"control_image": image,
}
return inputs
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline
pipe_sd = StableDiffusionXLControlNetImg2ImgPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
# pag enable
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3
class TestStableDiffusionXLControlNetPAGImg2ImgPipeline(
StableDiffusionXLControlNetPAGImg2ImgPipelineTesterConfig, PAGPipelineTesterMixin
):
base_pipeline_class = StableDiffusionXLControlNetImg2ImgPipeline
@pytest.mark.skip("We test this functionality elsewhere already.")
def test_save_load_optional_components(self):
pass
def test_pag_cfg(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# Run on CPU: the expected slice below is CPU-specific.
pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"])
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
64,
64,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array(
[0.55155665, 0.4650753, 0.46541628, 0.60965055, 0.55995595, 0.49751496, 0.5937391, 0.5700847, 0.44238678]
image = pipe_pag(**self.get_dummy_inputs())[0]
assert image.shape == (1, *self.output_shape), (
f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}"
)
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}"
# fmt: off
expected_slice = torch.tensor([0.55155665, 0.4650753, 0.46541628, 0.60965055, 0.55995595, 0.49751496, 0.5937391, 0.5700847, 0.44238678])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
def test_pag_uncond(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# Run on CPU: the expected slice below is CPU-specific.
pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"])
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs["guidance_scale"] = 0.0
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
64,
64,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array(
[0.549061, 0.46218234, 0.4675981, 0.6109464, 0.5547849, 0.4960261, 0.60211027, 0.5698843, 0.44092298]
image = pipe_pag(**inputs)[0]
assert image.shape == (1, *self.output_shape), (
f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}"
)
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}"
# fmt: off
expected_slice = torch.tensor([0.549061, 0.46218234, 0.4675981, 0.6109464, 0.5547849, 0.4960261, 0.60211027, 0.5698843, 0.44092298])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
class TestStableDiffusionXLControlNetPAGImg2ImgPipelineMemory(
StableDiffusionXLControlNetPAGImg2ImgPipelineTesterConfig, MemoryTesterMixin
):
"""Memory tests (CPU offload, group offload, layerwise casting) for the SDXL ControlNet PAG img2img pipeline."""
class TestStableDiffusionXLControlNetPAGImg2ImgPipelineIPAdapter(
StableDiffusionXLControlNetPAGImg2ImgPipelineTesterConfig, IPAdapterTesterMixin
):
"""IP-Adapter tests for the SDXL ControlNet PAG img2img pipeline."""
class TestStableDiffusionXLControlNetPAGImg2ImgPipelineFromPipe(
StableDiffusionXLControlNetPAGImg2ImgPipelineTesterConfig, FromPipeTesterMixin
):
"""`from_pipe` round-trip tests against `StableDiffusionXLPipeline`."""

247
tests/pipelines/pag/test_pag_hunyuan_dit.py
View File

@@ -13,11 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import tempfile
import unittest
import numpy as np
import pytest
import torch
from transformers import AutoConfig, AutoTokenizer, BertModel, T5EncoderModel
@@ -29,22 +25,21 @@ from diffusers import (
HunyuanDiTPipeline,
)
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
from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class HunyuanDiTPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class HunyuanDiTPAGPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = HunyuanDiTPAGPipeline
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 = PipelineTesterMixin.required_optional_params
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS
# `transformer.sample_size` (16) * `vae_scale_factor` (8) / 8 -> the dummy transformer generates 16x16 images.
output_shape = (3, 16, 16)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -72,7 +67,7 @@ class HunyuanDiTPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
text_encoder_2 = T5EncoderModel(config)
tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
return {
"transformer": transformer.eval(),
"vae": vae.eval(),
"scheduler": scheduler,
@@ -83,159 +78,88 @@ class HunyuanDiTPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
"safety_checker": None,
"feature_extractor": None,
}
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`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
"use_resolution_binning": False,
"pag_scale": 0.0,
}
return inputs
class TestHunyuanDiTPAGPipeline(HunyuanDiTPAGPipelineTesterConfig, PAGPipelineTesterMixin):
base_pipeline_class = HunyuanDiTPipeline
# HunyuanDiT's denoiser is a transformer: PAG uses the pipeline's default applied layers, turned on by raising
# `pag_scale` (the dummy inputs disable it with `pag_scale=0.0`).
pag_enabled_applied_layers = None
pag_enabled_scale = 3.0
def test_inference(self):
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)
image = pipe(**self.get_dummy_inputs()).images
assert image.shape == (1, *self.output_shape)
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs).images
image_slice = image[0, -3:, -3:, -1]
self.assertEqual(image.shape, (1, 16, 16, 3))
expected_slice = np.array(
[0.56939435, 0.34541583, 0.35915792, 0.46489206, 0.38775963, 0.45004836, 0.5957267, 0.59481275, 0.33287364]
)
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
self.assertLessEqual(max_diff, 1e-3)
@unittest.skip("Not supported.")
def test_sequential_cpu_offload_forward_pass(self):
# TODO(YiYi) need to fix later
pass
@unittest.skip("Not supported.")
def test_sequential_offload_forward_pass_twice(self):
# TODO(YiYi) need to fix later
pass
# fmt: off
expected_slice = torch.tensor([0.56939435, 0.34541583, 0.35915792, 0.46489206, 0.38775963, 0.45004836, 0.5957267, 0.59481275, 0.33287364])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(
expected_max_diff=1e-3,
)
super().test_inference_batch_single_identical(expected_max_diff=1e-3)
def test_feed_forward_chunking(self):
device = "cpu"
# Run on CPU to keep the device-dependent `torch.Generator` deterministic.
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)
image = pipe(**inputs).images
image_slice_no_chunking = image[0, -3:, -3:, -1]
output_no_chunking = self.run_pipe(pipe)
pipe.transformer.enable_forward_chunking(chunk_size=1, dim=0)
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs).images
image_slice_chunking = image[0, -3:, -3:, -1]
output_chunking = self.run_pipe(pipe)
max_diff = np.abs(to_np(image_slice_no_chunking) - to_np(image_slice_chunking)).max()
self.assertLess(max_diff, 1e-4)
assert_tensors_close(
output_chunking, output_no_chunking, atol=1e-4, msg="Forward chunking changed the output."
)
def test_fused_qkv_projections(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
# Run on CPU to keep the device-dependent `torch.Generator` deterministic.
pipe = self.get_pipeline()
inputs = self.get_dummy_inputs(device)
inputs["return_dict"] = False
image = pipe(**inputs)[0]
original_image_slice = image[0, -3:, -3:, -1]
original_output = self.run_pipe(pipe)
pipe.transformer.fuse_qkv_projections()
inputs = self.get_dummy_inputs(device)
inputs["return_dict"] = False
image_fused = pipe(**inputs)[0]
image_slice_fused = image_fused[0, -3:, -3:, -1]
output_fused = self.run_pipe(pipe)
pipe.transformer.unfuse_qkv_projections()
inputs = self.get_dummy_inputs(device)
inputs["return_dict"] = False
image_disabled = pipe(**inputs)[0]
image_slice_disabled = image_disabled[0, -3:, -3:, -1]
output_disabled = self.run_pipe(pipe)
assert np.allclose(original_image_slice, image_slice_fused, atol=1e-2, rtol=1e-2), (
"Fusion of QKV projections shouldn't affect the outputs."
assert_tensors_close(
output_fused, original_output, atol=1e-2, rtol=1e-2, msg="Fusion of QKV projections changed the outputs."
)
assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-2, rtol=1e-2), (
"Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled."
assert_tensors_close(
output_disabled,
output_fused,
atol=1e-2,
rtol=1e-2,
msg="Outputs changed after the fused QKV projections were disabled.",
)
assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), (
"Original outputs should match when fused QKV projections are disabled."
assert_tensors_close(
output_disabled,
original_output,
atol=1e-2,
rtol=1e-2,
msg="Original outputs should match when fused QKV projections are disabled.",
)
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline (expect same output when pag is disabled)
pipe_sd = HunyuanDiTPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
components = self.get_dummy_components()
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
# pag enabled
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 3.0
out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3
def test_pag_applied_layers(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline
pipe = self.pipeline_class(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()
all_self_attn_layers = [k for k in pipe.transformer.attn_processors.keys() if "attn1" in k]
original_attn_procs = pipe.transformer.attn_processors
@@ -265,19 +189,16 @@ class HunyuanDiTPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
assert len(pipe.pag_attn_processors) == 2
@unittest.skip(
@pytest.mark.skip(
"Test not supported as `encode_prompt` is called two times separately which deivates from about 99% of the pipelines we have."
)
def test_encode_prompt_works_in_isolation(self):
pass
def test_save_load_optional_components(self):
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
def test_save_load_optional_components(self, tmp_path, expected_max_difference=1e-4):
pipe = self.get_pipeline().to(torch_device)
inputs = self.get_dummy_inputs(torch_device)
inputs = self.get_dummy_inputs()
prompt = inputs["prompt"]
generator = inputs["generator"]
@@ -325,19 +246,17 @@ class HunyuanDiTPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
output = pipe(**inputs)[0]
with tempfile.TemporaryDirectory() as tmpdir:
pipe.save_pretrained(tmpdir)
pipe_loaded = self.pipeline_class.from_pretrained(tmpdir)
pipe_loaded.to(torch_device)
pipe_loaded.set_progress_bar_config(disable=None)
pipe.save_pretrained(tmp_path)
pipe_loaded = self.pipeline_class.from_pretrained(tmp_path)
pipe_loaded.to(torch_device)
pipe_loaded.set_progress_bar_config(disable=None)
for optional_component in pipe._optional_components:
self.assertTrue(
getattr(pipe_loaded, optional_component) is None,
f"`{optional_component}` did not stay set to None after loading.",
assert getattr(pipe_loaded, optional_component) is None, (
f"`{optional_component}` did not stay set to None after loading."
)
inputs = self.get_dummy_inputs(torch_device)
inputs = self.get_dummy_inputs()
generator = inputs["generator"]
num_inference_steps = inputs["num_inference_steps"]
@@ -361,5 +280,23 @@ class HunyuanDiTPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
output_loaded = pipe_loaded(**inputs)[0]
max_diff = np.abs(to_np(output) - to_np(output_loaded)).max()
self.assertLess(max_diff, 1e-4)
assert_tensors_close(
output_loaded,
output,
atol=expected_max_difference,
msg="Output changed after dropping optional components.",
)
class TestHunyuanDiTPAGPipelineMemory(HunyuanDiTPAGPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the HunyuanDiT PAG pipeline."""
@pytest.mark.skip("Not supported.")
def test_sequential_cpu_offload_forward_pass(self):
# TODO(YiYi) need to fix later
pass
@pytest.mark.skip("Not supported.")
def test_sequential_offload_forward_pass_twice(self):
# TODO(YiYi) need to fix later
pass

126
tests/pipelines/pag/test_pag_kolors.py
View File

@@ -13,10 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import unittest
import numpy as np
import pytest
import torch
from diffusers import (
@@ -32,29 +29,21 @@ from ...testing_utils import enable_full_determinism
from ..pipeline_params import (
TEXT_TO_IMAGE_BATCH_PARAMS,
TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS,
TEXT_TO_IMAGE_IMAGE_PARAMS,
TEXT_TO_IMAGE_PARAMS,
)
from ..test_pipelines_common import (
PipelineFromPipeTesterMixin,
PipelineTesterMixin,
)
from ..testing_utils import BasePipelineTesterConfig, FromPipeTesterMixin, MemoryTesterMixin
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class KolorsPAGPipelineFastTests(
PipelineTesterMixin,
PipelineFromPipeTesterMixin,
unittest.TestCase,
):
class KolorsPAGPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = KolorsPAGPipeline
params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS
callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union({"add_text_embeds", "add_time_ids"})
output_shape = (3, 64, 64)
# Copied from tests.pipelines.kolors.test_kolors.KolorsPipelineTesterConfig.get_dummy_components
def get_dummy_components(self, time_cond_proj_dim=None):
@@ -112,65 +101,27 @@ class KolorsPAGPipelineFastTests(
}
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,
"pag_scale": 0.9,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
}
return inputs
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline (expect same output when pag is disabled)
pipe_sd = KolorsPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
# pag enabled
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3
class TestKolorsPAGPipeline(KolorsPAGPipelineTesterConfig, PAGPipelineTesterMixin):
base_pipeline_class = KolorsPipeline
# fmt: off
expected_pag_slice = torch.tensor([0.26030684, 0.43192005, 0.4042826, 0.4189067, 0.5181305, 0.3832534, 0.472135, 0.4145031, 0.43726248])
# fmt: on
def test_pag_applied_layers(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline
pipe = self.pipeline_class(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()
# pag_applied_layers = ["mid","up","down"] should apply to all self-attention layers
all_self_attn_layers = [k for k in pipe.unet.attn_processors.keys() if "attn1" in k]
@@ -201,7 +152,7 @@ class KolorsPAGPipelineFastTests(
# pag_applied_layers = ["mid.block_0.attentions_1"] does not exist in the model
pipe.unet.set_attn_processor(original_attn_procs.copy())
pag_layers = ["mid_block.attentions.1"]
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
# pag_applied_layers = "down" should apply to all self-attention layers in down_blocks
@@ -212,7 +163,7 @@ class KolorsPAGPipelineFastTests(
pipe.unet.set_attn_processor(original_attn_procs.copy())
pag_layers = ["down_blocks.0"]
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
pipe.unet.set_attn_processor(original_attn_procs.copy())
@@ -225,33 +176,16 @@ class KolorsPAGPipelineFastTests(
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
assert len(pipe.pag_attn_processors) == 2
def test_pag_inference(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
64,
64,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array(
[0.26030684, 0.43192005, 0.4042826, 0.4189067, 0.5181305, 0.3832534, 0.472135, 0.4145031, 0.43726248]
)
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
self.assertLessEqual(max_diff, 1e-3)
def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(expected_max_diff=3e-3)
super().test_inference_batch_single_identical(expected_max_diff=3e-3)
def test_encode_prompt_works_in_isolation(self):
return super().test_encode_prompt_works_in_isolation(atol=1e-3, rtol=1e-3)
class TestKolorsPAGPipelineMemory(KolorsPAGPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the Kolors PAG pipeline."""
class TestKolorsPAGPipelineFromPipe(KolorsPAGPipelineTesterConfig, FromPipeTesterMixin):
"""`from_pipe` round-trip tests against `KolorsPipeline`."""

301
tests/pipelines/pag/test_pag_pixart_sigma.py
View File

@@ -13,15 +13,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import tempfile
import unittest
import numpy as np
import pytest
import torch
from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
import diffusers
from diffusers import (
AutoencoderKL,
DDIMScheduler,
@@ -29,33 +24,30 @@ from diffusers import (
PixArtSigmaPipeline,
PixArtTransformer2DModel,
)
from diffusers.utils import logging
from ...testing_utils import (
CaptureLogger,
assert_tensors_close,
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, assert_mean_pixel_difference, to_np
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class PixArtSigmaPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class PixArtSigmaPAGPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = PixArtSigmaPAGPipeline
params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
params = set(params)
params.remove("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 = PipelineTesterMixin.required_optional_params
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - {
"cross_attention_kwargs"
}
batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS
# `transformer.sample_size` (8) * `vae_scale_factor` (8) / 8 -> the dummy transformer generates 8x8 images.
output_shape = (3, 8, 8)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -85,76 +77,40 @@ class PixArtSigmaPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
return {
"transformer": transformer.eval(),
"vae": vae.eval(),
"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": "A painting of a squirrel eating a burger",
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 1.0,
"pag_scale": 3.0,
"use_resolution_binning": False,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
}
return inputs
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline (expect same output when pag is disabled)
pipe = PixArtSigmaPipeline(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe.__class__.__name__}."
)
out = pipe(**inputs).images[0, -3:, -3:, -1]
# pag disabled with pag_scale=0.0
components["pag_applied_layers"] = ["blocks.1"]
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
# pag enabled
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3
class TestPixArtSigmaPAGPipeline(PixArtSigmaPAGPipelineTesterConfig, PAGPipelineTesterMixin):
base_pipeline_class = PixArtSigmaPipeline
# PixArt's denoiser is a transformer, so PAG resolves per transformer block rather than the UNet's mid/up/down.
pag_enabled_applied_layers = ["blocks.1"]
# `test_pag_inference` builds the pipeline with the class default (`blocks.1`), as it did before the migration.
pag_inference_applied_layers = None
# fmt: off
expected_pag_slice = torch.tensor([0.6499, 0.3250, 0.3572, 0.6780, 0.4453, 0.4582, 0.2770, 0.5168, 0.4594])
# fmt: on
def test_pag_applied_layers(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline
pipe = self.pipeline_class(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()
# "attn1" should apply to all self-attention layers.
all_self_attn_layers = [k for k in pipe.transformer.attn_processors.keys() if "attn1" in k]
@@ -162,189 +118,38 @@ class PixArtSigmaPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
assert set(pipe.pag_attn_processors) == set(all_self_attn_layers)
def test_pag_inference(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
def test_attention_slicing_forward_pass(self, expected_max_diff=1e-3):
# Run on CPU: sliced attention is compared against a full-attention run of the same pipeline.
pipe = self.get_pipeline()
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
8,
8,
3,
), f"the shape of the output image should be (1, 8, 8, 3) but got {image.shape}"
expected_slice = np.array([0.6499, 0.3250, 0.3572, 0.6780, 0.4453, 0.4582, 0.2770, 0.5168, 0.4594])
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
self.assertLessEqual(max_diff, 1e-3)
# Because the PAG PixArt Sigma has `pag_applied_layers`.
# Also, we shouldn't be doing `set_default_attn_processor()` after loading
# the pipeline with `pag_applied_layers`.
def test_save_load_local(self, expected_max_difference=1e-4):
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(torch_device)
output = pipe(**inputs)[0]
logger = logging.get_logger("diffusers.pipelines.pipeline_utils")
logger.setLevel(diffusers.logging.INFO)
with tempfile.TemporaryDirectory() as tmpdir:
pipe.save_pretrained(tmpdir, safe_serialization=False)
with CaptureLogger(logger) as cap_logger:
pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, pag_applied_layers=["blocks.1"])
for name in pipe_loaded.components.keys():
if name not in pipe_loaded._optional_components:
assert name in str(cap_logger)
pipe_loaded.to(torch_device)
pipe_loaded.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(torch_device)
output_loaded = pipe_loaded(**inputs)[0]
max_diff = np.abs(to_np(output) - to_np(output_loaded)).max()
self.assertLess(max_diff, expected_max_difference)
# We shouldn't be setting `set_default_attn_processor` here.
def test_attention_slicing_forward_pass(
self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
):
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
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]
output_without_slicing = self.run_pipe(pipe)
pipe.enable_attention_slicing(slice_size=1)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing1 = pipe(**inputs)[0]
output_with_slicing_1 = self.run_pipe(pipe)
pipe.enable_attention_slicing(slice_size=2)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing2 = pipe(**inputs)[0]
output_with_slicing_2 = self.run_pipe(pipe)
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_tensors_close(
output_with_slicing_1,
output_without_slicing,
atol=expected_max_diff,
msg="Attention slicing (slice_size=1) changed the output.",
)
assert_tensors_close(
output_with_slicing_2,
output_without_slicing,
atol=expected_max_diff,
msg="Attention slicing (slice_size=2) changed the output.",
)
if test_mean_pixel_difference:
assert_mean_pixel_difference(to_np(output_with_slicing1[0]), to_np(output_without_slicing[0]))
assert_mean_pixel_difference(to_np(output_with_slicing2[0]), to_np(output_without_slicing[0]))
def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(batch_size=2)
# Because we have `pag_applied_layers` we cannot directly apply
# `set_default_attn_processor`
def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4):
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
generator_device = "cpu"
if expected_slice is None:
output = pipe(**self.get_dummy_inputs(generator_device))[0]
else:
output = expected_slice
output_tuple = pipe(**self.get_dummy_inputs(generator_device), return_dict=False)[0]
if expected_slice is None:
max_diff = np.abs(to_np(output) - to_np(output_tuple)).max()
else:
if output_tuple.ndim != 5:
max_diff = np.abs(to_np(output) - to_np(output_tuple)[0, -3:, -3:, -1].flatten()).max()
else:
max_diff = np.abs(to_np(output) - to_np(output_tuple)[0, -3:, -3:, -1, -1].flatten()).max()
self.assertLess(max_diff, expected_max_difference)
# Same reason as above
def test_inference_batch_single_identical(
self,
batch_size=2,
expected_max_diff=1e-4,
additional_params_copy_to_batched_inputs=["num_inference_steps"],
):
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(torch_device)
# Reset generator in case it is has been used in self.get_dummy_inputs
inputs["generator"] = self.get_generator(0)
logger = logging.get_logger(pipe.__module__)
logger.setLevel(level=diffusers.logging.FATAL)
# batchify inputs
batched_inputs = {}
batched_inputs.update(inputs)
for name in self.batch_params:
if name not in inputs:
continue
value = inputs[name]
if name == "prompt":
len_prompt = len(value)
batched_inputs[name] = [value[: len_prompt // i] for i in range(1, batch_size + 1)]
batched_inputs[name][-1] = 100 * "very long"
else:
batched_inputs[name] = batch_size * [value]
if "generator" in inputs:
batched_inputs["generator"] = [self.get_generator(i) for i in range(batch_size)]
if "batch_size" in inputs:
batched_inputs["batch_size"] = batch_size
for arg in additional_params_copy_to_batched_inputs:
batched_inputs[arg] = inputs[arg]
output = pipe(**inputs)
output_batch = pipe(**batched_inputs)
assert output_batch[0].shape[0] == batch_size
max_diff = np.abs(to_np(output_batch[0][0]) - to_np(output[0][0])).max()
assert max_diff < expected_max_diff
# Because we're passing `pag_applied_layers` (type of List) in the components as well.
def test_components_function(self):
init_components = self.get_dummy_components()
init_components = {k: v for k, v in init_components.items() if not isinstance(v, (str, int, float, list))}
pipe = self.pipeline_class(**init_components)
self.assertTrue(hasattr(pipe, "components"))
self.assertTrue(set(pipe.components.keys()) == set(init_components.keys()))
@unittest.skip("Test is already covered through encode_prompt isolation.")
@pytest.mark.skip("Test is already covered through encode_prompt isolation.")
def test_save_load_optional_components(self):
pass
class TestPixArtSigmaPAGPipelineMemory(PixArtSigmaPAGPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the PixArt-sigma PAG pipeline."""

238
tests/pipelines/pag/test_pag_sana.py
View File

@@ -12,10 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import unittest
import numpy as np
import pytest
import torch
from transformers import Gemma2Config, Gemma2ForCausalLM, GemmaTokenizer
@@ -27,31 +24,26 @@ from diffusers import (
SanaTransformer2DModel,
)
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
from ...testing_utils import (
assert_tensors_close,
enable_full_determinism,
require_accelerator,
skip_if_no_cudnn_engine,
torch_device,
)
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class SanaPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class SanaPAGPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = SanaPAGPipeline
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",
]
)
test_xformers_attention = False
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS
output_shape = (3, 32, 32)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -115,183 +107,64 @@ class SanaPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
text_encoder = Gemma2ForCausalLM(text_encoder_config)
tokenizer = GemmaTokenizer.from_pretrained("hf-internal-testing/dummy-gemma")
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": "",
"negative_prompt": "",
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 6.0,
"pag_scale": 3.0,
"height": 32,
"width": 32,
"max_sequence_length": 16,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
"complex_human_instruction": None,
}
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 TestSanaPAGPipeline(SanaPAGPipelineTesterConfig, PAGPipelineTesterMixin):
base_pipeline_class = SanaPipeline
# Only the "PAG off reproduces the base pipeline" leg was asserted before the migration.
check_pag_changes_output = False
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs)[0]
generated_image = image[0]
def test_attention_slicing_forward_pass(self, expected_max_diff=1e-3):
# Run on CPU: sliced attention is compared against a full-attention run of the same pipeline.
pipe = self.get_pipeline()
self.assertEqual(generated_image.shape, (3, 32, 32))
expected_image = torch.randn(3, 32, 32)
max_diff = np.abs(generated_image - expected_image).max()
self.assertLessEqual(max_diff, 1e10)
def test_callback_inputs(self):
sig = inspect.signature(self.pipeline_class.__call__)
has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters
has_callback_step_end = "callback_on_step_end" in sig.parameters
if not (has_callback_tensor_inputs and has_callback_step_end):
return
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
self.assertTrue(
hasattr(pipe, "_callback_tensor_inputs"),
f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs",
)
def callback_inputs_subset(pipe, i, t, callback_kwargs):
# iterate over callback args
for tensor_name, tensor_value in callback_kwargs.items():
# check that we're only passing in allowed tensor inputs
assert tensor_name in pipe._callback_tensor_inputs
return callback_kwargs
def callback_inputs_all(pipe, i, t, callback_kwargs):
for tensor_name in pipe._callback_tensor_inputs:
assert tensor_name in callback_kwargs
# iterate over callback args
for tensor_name, tensor_value in callback_kwargs.items():
# check that we're only passing in allowed tensor inputs
assert tensor_name in pipe._callback_tensor_inputs
return callback_kwargs
inputs = self.get_dummy_inputs(torch_device)
# Test passing in a subset
inputs["callback_on_step_end"] = callback_inputs_subset
inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
output = pipe(**inputs)[0]
# Test passing in a everything
inputs["callback_on_step_end"] = callback_inputs_all
inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
output = pipe(**inputs)[0]
def callback_inputs_change_tensor(pipe, i, t, callback_kwargs):
is_last = i == (pipe.num_timesteps - 1)
if is_last:
callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"])
return callback_kwargs
inputs["callback_on_step_end"] = callback_inputs_change_tensor
inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
output = pipe(**inputs)[0]
assert output.abs().sum() < 1e10
def test_attention_slicing_forward_pass(
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]
output_without_slicing = self.run_pipe(pipe)
pipe.enable_attention_slicing(slice_size=1)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing1 = pipe(**inputs)[0]
output_with_slicing_1 = self.run_pipe(pipe)
pipe.enable_attention_slicing(slice_size=2)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing2 = pipe(**inputs)[0]
output_with_slicing_2 = self.run_pipe(pipe)
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",
)
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline (expect same output when pag is disabled)
pipe_sd = SanaPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
assert_tensors_close(
output_with_slicing_1,
output_without_slicing,
atol=expected_max_diff,
msg="Attention slicing (slice_size=1) changed the output.",
)
assert_tensors_close(
output_with_slicing_2,
output_without_slicing,
atol=expected_max_diff,
msg="Attention slicing (slice_size=2) changed the output.",
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
components = self.get_dummy_components()
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
def test_pag_applied_layers(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline
pipe = self.pipeline_class(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()
all_self_attn_layers = [k for k in pipe.transformer.attn_processors.keys() if "attn1" in k]
original_attn_procs = pipe.transformer.attn_processors
@@ -322,18 +195,33 @@ class SanaPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
assert len(pipe.pag_attn_processors) == 2
# TODO(aryan): Create a dummy gemma model with smol vocab size
@unittest.skip(
@pytest.mark.skip(
"A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_consistent(self):
pass
@unittest.skip(
@pytest.mark.skip(
"A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
)
def test_inference_batch_single_identical(self):
pass
def test_float16_inference(self):
# Requires higher tolerance as model seems very sensitive to dtype
super().test_float16_inference(expected_max_diff=0.08)
# Sana's multiscale linear attention runs a depthwise `Conv2d`, which some cuDNN builds have no bfloat16
# engine for. The decorators below repeat the ones the base method is declared with: overriding a test drops
# the marks it inherited.
@pytest.mark.skipif(torch_device not in ["cuda", "xpu"], reason="half-precision inference requires CUDA or XPU")
@require_accelerator
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=str)
def test_half_precision_inference_no_nan(self, dtype):
with skip_if_no_cudnn_engine():
super().test_half_precision_inference_no_nan(dtype)
class TestSanaPAGPipelineMemory(SanaPAGPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the Sana PAG pipeline."""
# Layerwise casting computes in bfloat16, which lands on the same missing depthwise-conv engine as above.
def test_layerwise_casting_inference(self):
with skip_if_no_cudnn_engine():
super().test_layerwise_casting_inference()

146
tests/pipelines/pag/test_pag_sd.py
View File

@@ -14,10 +14,9 @@
# limitations under the License.
import gc
import inspect
import unittest
import numpy as np
import pytest
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
@@ -40,33 +39,26 @@ from ...testing_utils import (
from ..pipeline_params import (
TEXT_TO_IMAGE_BATCH_PARAMS,
TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS,
TEXT_TO_IMAGE_IMAGE_PARAMS,
TEXT_TO_IMAGE_PARAMS,
)
from ..test_pipelines_common import (
from ..testing_utils import (
BasePipelineTesterConfig,
FromPipeTesterMixin,
IPAdapterTesterMixin,
PipelineFromPipeTesterMixin,
PipelineLatentTesterMixin,
PipelineTesterMixin,
MemoryTesterMixin,
)
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class StableDiffusionPAGPipelineFastTests(
PipelineTesterMixin,
IPAdapterTesterMixin,
PipelineLatentTesterMixin,
PipelineFromPipeTesterMixin,
unittest.TestCase,
):
class StableDiffusionPAGPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = StableDiffusionPAGPipeline
params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS
callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union({"add_text_embeds", "add_time_ids"})
output_shape = (3, 64, 64)
def get_dummy_components(self, time_cond_proj_dim=None):
cross_attention_dim = 8
@@ -116,7 +108,7 @@ class StableDiffusionPAGPipelineFastTests(
text_encoder = CLIPTextModel(text_encoder_config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
components = {
return {
"unet": unet,
"scheduler": scheduler,
"vae": vae,
@@ -126,67 +118,28 @@ class StableDiffusionPAGPipelineFastTests(
"feature_extractor": None,
"image_encoder": None,
}
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,
"pag_scale": 0.9,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
}
return inputs
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline (expect same output when pag is disabled)
pipe_sd = StableDiffusionPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
# pag enabled
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3
class TestStableDiffusionPAGPipeline(StableDiffusionPAGPipelineTesterConfig, PAGPipelineTesterMixin):
base_pipeline_class = StableDiffusionPipeline
# fmt: off
expected_pag_slice = torch.tensor([0.23171297, 0.44669262, 0.48407662, 0.29981518, 0.36721927, 0.46788025, 0.46333545, 0.3314417, 0.42078137])
# fmt: on
def test_pag_applied_layers(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline
pipe = self.pipeline_class(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()
# pag_applied_layers = ["mid","up","down"] should apply to all self-attention layers
all_self_attn_layers = [k for k in pipe.unet.attn_processors.keys() if "attn1" in k]
@@ -224,7 +177,7 @@ class StableDiffusionPAGPipelineFastTests(
# pag_applied_layers = ["mid.block_0.attentions_1"] does not exist in the model
pipe.unet.set_attn_processor(original_attn_procs.copy())
pag_layers = ["mid_block.attentions.1"]
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
# pag_applied_layers = "down" should apply to all self-attention layers in down_blocks
@@ -239,7 +192,7 @@ class StableDiffusionPAGPipelineFastTests(
pipe.unet.set_attn_processor(original_attn_procs.copy())
pag_layers = ["down_blocks.0"]
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
pipe.unet.set_attn_processor(original_attn_procs.copy())
@@ -252,52 +205,37 @@ class StableDiffusionPAGPipelineFastTests(
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
assert len(pipe.pag_attn_processors) == 1
def test_pag_inference(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
64,
64,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array(
[0.23171297, 0.44669262, 0.48407662, 0.29981518, 0.36721927, 0.46788025, 0.46333545, 0.3314417, 0.42078137]
)
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
self.assertLessEqual(max_diff, 1e-3)
def test_encode_prompt_works_in_isolation(self):
extra_required_param_value_dict = {
"device": torch.device(torch_device).type,
"do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0,
"do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0,
}
return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict)
class TestStableDiffusionPAGPipelineMemory(StableDiffusionPAGPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD PAG pipeline."""
class TestStableDiffusionPAGPipelineIPAdapter(StableDiffusionPAGPipelineTesterConfig, IPAdapterTesterMixin):
"""IP-Adapter tests for the SD PAG pipeline."""
class TestStableDiffusionPAGPipelineFromPipe(StableDiffusionPAGPipelineTesterConfig, FromPipeTesterMixin):
"""`from_pipe` round-trip tests against `StableDiffusionPipeline`."""
@slow
@require_torch_accelerator
class StableDiffusionPAGPipelineIntegrationTests(unittest.TestCase):
class TestStableDiffusionPAGPipelineIntegration:
pipeline_class = StableDiffusionPAGPipeline
repo_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
def setUp(self):
super().setUp()
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
yield
gc.collect()
backend_empty_cache(torch_device)

157
tests/pipelines/pag/test_pag_sd3.py
View File

@@ -1,7 +1,3 @@
import inspect
import unittest
import numpy as np
import torch
from transformers import (
AutoConfig,
@@ -20,19 +16,19 @@ from diffusers import (
StableDiffusion3Pipeline,
)
from ...testing_utils import (
torch_device,
)
from ..test_pipelines_common import (
PipelineTesterMixin,
from ...testing_utils import assert_tensors_close, torch_device
from ..testing_utils import (
BasePipelineTesterConfig,
MemoryTesterMixin,
check_qkv_fusion_matches_attn_procs_length,
check_qkv_fusion_processors_exist,
)
from .testing_utils import PAGPipelineTesterMixin
class StableDiffusion3PAGPipelineFastTests(unittest.TestCase, PipelineTesterMixin):
class StableDiffusion3PAGPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = StableDiffusion3PAGPipeline
params = frozenset(
required_input_params_in_call_signature = frozenset(
[
"prompt",
"height",
@@ -43,8 +39,8 @@ class StableDiffusion3PAGPipelineFastTests(unittest.TestCase, PipelineTesterMixi
"negative_prompt_embeds",
]
)
batch_params = frozenset(["prompt", "negative_prompt"])
test_xformers_attention = False
batch_input_params = frozenset(["prompt", "negative_prompt"])
output_shape = (3, 32, 32)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -117,64 +113,50 @@ class StableDiffusion3PAGPipelineFastTests(unittest.TestCase, PipelineTesterMixi
"vae": vae,
}
def get_dummy_inputs(self, device, seed=0):
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device="cpu").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`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
"pag_scale": 0.0,
}
return inputs
class TestStableDiffusion3PAGPipeline(StableDiffusion3PAGPipelineTesterConfig, PAGPipelineTesterMixin):
base_pipeline_class = StableDiffusion3Pipeline
# The dummy inputs disable PAG (`pag_scale=0.0`); only the "PAG off reproduces the base pipeline" leg is
# asserted here, as it was before the migration.
check_pag_changes_output = False
def test_stable_diffusion_3_different_prompts(self):
pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
pipe = self.get_pipeline().to(torch_device)
inputs = self.get_dummy_inputs(torch_device)
output_same_prompt = pipe(**inputs).images[0]
inputs = self.get_dummy_inputs(torch_device)
inputs["prompt_2"] = "a different prompt"
inputs["prompt_3"] = "another different prompt"
output_different_prompts = pipe(**inputs).images[0]
max_diff = np.abs(output_same_prompt - output_different_prompts).max()
output_same_prompt = self.run_pipe(pipe)[0]
output_different_prompts = self.run_pipe(
pipe, prompt_2="a different prompt", prompt_3="another different prompt"
)[0]
# Outputs should be different here
assert max_diff > 1e-2
assert (output_same_prompt - output_different_prompts).abs().max() > 1e-2
def test_stable_diffusion_3_different_negative_prompts(self):
pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
pipe = self.get_pipeline().to(torch_device)
inputs = self.get_dummy_inputs(torch_device)
output_same_prompt = pipe(**inputs).images[0]
inputs = self.get_dummy_inputs(torch_device)
inputs["negative_prompt_2"] = "deformed"
inputs["negative_prompt_3"] = "blurry"
output_different_prompts = pipe(**inputs).images[0]
max_diff = np.abs(output_same_prompt - output_different_prompts).max()
output_same_prompt = self.run_pipe(pipe)[0]
output_different_prompts = self.run_pipe(pipe, negative_prompt_2="deformed", negative_prompt_3="blurry")[0]
# Outputs should be different here
assert max_diff > 1e-2
assert (output_same_prompt - output_different_prompts).abs().max() > 1e-2
def test_fused_qkv_projections(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
# Run on CPU to keep the device-dependent `torch.Generator` deterministic.
pipe = self.get_pipeline()
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs).images
original_image_slice = image[0, -3:, -3:, -1]
original_output = self.run_pipe(pipe)
# TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added
# to the pipeline level.
@@ -186,62 +168,31 @@ class StableDiffusion3PAGPipelineFastTests(unittest.TestCase, PipelineTesterMixi
pipe.transformer, pipe.transformer.original_attn_processors
), "Something wrong with the attention processors concerning the fused QKV projections."
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs).images
image_slice_fused = image[0, -3:, -3:, -1]
output_fused = self.run_pipe(pipe)
pipe.transformer.unfuse_qkv_projections()
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs).images
image_slice_disabled = image[0, -3:, -3:, -1]
output_disabled = self.run_pipe(pipe)
assert np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), (
"Fusion of QKV projections shouldn't affect the outputs."
assert_tensors_close(
output_fused, original_output, atol=1e-3, rtol=1e-3, msg="Fusion of QKV projections changed the outputs."
)
assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), (
"Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled."
assert_tensors_close(
output_disabled,
output_fused,
atol=1e-3,
rtol=1e-3,
msg="Outputs changed after the fused QKV projections were disabled.",
)
assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), (
"Original outputs should match when fused QKV projections are disabled."
assert_tensors_close(
output_disabled,
original_output,
atol=1e-2,
rtol=1e-2,
msg="Original outputs should match when fused QKV projections are disabled.",
)
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline (expect same output when pag is disabled)
pipe_sd = StableDiffusion3Pipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
components = self.get_dummy_components()
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
def test_pag_applied_layers(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline
pipe = self.pipeline_class(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()
all_self_attn_layers = [k for k in pipe.transformer.attn_processors.keys() if "attn" in k]
original_attn_procs = pipe.transformer.attn_processors
@@ -270,3 +221,7 @@ class StableDiffusion3PAGPipelineFastTests(unittest.TestCase, PipelineTesterMixi
pag_layers = ["blocks.0", r"blocks\.1"]
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
assert len(pipe.pag_attn_processors) == 2
class TestStableDiffusion3PAGPipelineMemory(StableDiffusion3PAGPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD3 PAG pipeline."""

116
tests/pipelines/pag/test_pag_sd3_img2img.py
View File

@@ -1,9 +1,8 @@
import gc
import inspect
import random
import unittest
import numpy as np
import pytest
import torch
from transformers import (
AutoConfig,
@@ -33,29 +32,30 @@ from ...testing_utils import (
torch_device,
)
from ..pipeline_params import (
IMAGE_TO_IMAGE_IMAGE_PARAMS,
TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS,
)
from ..test_pipelines_common import (
PipelineTesterMixin,
)
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class StableDiffusion3PAGImg2ImgPipelineFastTests(unittest.TestCase, PipelineTesterMixin):
class StableDiffusion3PAGImg2ImgPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = StableDiffusion3PAGImg2ImgPipeline
params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - {"height", "width"}
required_optional_params = PipelineTesterMixin.required_optional_params - {"latents"}
batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS
image_latens_params = IMAGE_TO_IMAGE_IMAGE_PARAMS
required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union(
{"pag_scale", "pag_adaptive_scale"}
) - {"height", "width"}
batch_input_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS
test_xformers_attention = False
# Img2img derives the latents from the input image, so `__call__` takes no `latents`.
optional_input_params = frozenset(
["num_inference_steps", "num_images_per_prompt", "generator", "output_type", "return_dict"]
)
# The output resolution follows the 32x32 input image.
output_shape = (3, 32, 32)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -128,93 +128,49 @@ class StableDiffusion3PAGImg2ImgPipelineFastTests(unittest.TestCase, PipelineTes
"vae": vae,
}
def get_dummy_inputs(self, device, seed=0):
image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device)
def get_dummy_inputs(self):
image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device)
image = image / 2 + 0.5
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device="cpu").manual_seed(seed)
inputs = {
return {
"prompt": "A painting of a squirrel eating a burger",
"image": image,
"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`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
"pag_scale": 0.7,
}
return inputs
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline (expect same output when pag is disabled)
pipe_sd = StableDiffusion3Img2ImgPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
class TestStableDiffusion3PAGImg2ImgPipeline(StableDiffusion3PAGImg2ImgPipelineTesterConfig, PAGPipelineTesterMixin):
base_pipeline_class = StableDiffusion3Img2ImgPipeline
# SD3's denoiser is a transformer, so PAG resolves per transformer block rather than the UNet's mid/up/down.
pag_inference_applied_layers = ["blocks.0"]
# Only the "PAG off reproduces the base pipeline" leg was asserted before the migration.
check_pag_changes_output = False
# fmt: off
expected_pag_slice = torch.tensor([0.741577, 0.5491905, 0.59911674, 0.7702221, 0.65531653, 0.60989463, 0.491042, 0.5380116, 0.5592475])
# fmt: on
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
components = self.get_dummy_components()
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
def test_pag_inference(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["blocks.0"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
32,
32,
3,
), f"the shape of the output image should be (1, 32, 32, 3) but got {image.shape}"
expected_slice = np.array(
[0.741577, 0.5491905, 0.59911674, 0.7702221, 0.65531653, 0.60989463, 0.491042, 0.5380116, 0.5592475]
)
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
self.assertLessEqual(max_diff, 1e-3)
class TestStableDiffusion3PAGImg2ImgPipelineMemory(StableDiffusion3PAGImg2ImgPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD3 PAG img2img pipeline."""
@slow
@require_torch_accelerator
class StableDiffusion3PAGImg2ImgPipelineIntegrationTests(unittest.TestCase):
class TestStableDiffusion3PAGImg2ImgPipelineIntegration:
pipeline_class = StableDiffusion3PAGImg2ImgPipeline
repo_id = "stabilityai/stable-diffusion-3-medium-diffusers"
def setUp(self):
super().setUp()
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
yield
gc.collect()
backend_empty_cache(torch_device)

145
tests/pipelines/pag/test_pag_sd_img2img.py
View File

@@ -14,17 +14,15 @@
# limitations under the License.
import gc
import inspect
import random
import unittest
import numpy as np
import pytest
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
from diffusers import (
AutoencoderKL,
AutoencoderTiny,
AutoPipelineForImage2Image,
EulerDiscreteScheduler,
StableDiffusionImg2ImgPipeline,
@@ -42,36 +40,34 @@ from ...testing_utils import (
torch_device,
)
from ..pipeline_params import (
IMAGE_TO_IMAGE_IMAGE_PARAMS,
TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS,
)
from ..test_pipelines_common import (
from ..testing_utils import (
BasePipelineTesterConfig,
IPAdapterTesterMixin,
PipelineKarrasSchedulerTesterMixin,
PipelineLatentTesterMixin,
PipelineTesterMixin,
MemoryTesterMixin,
)
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class StableDiffusionPAGImg2ImgPipelineFastTests(
IPAdapterTesterMixin,
PipelineLatentTesterMixin,
PipelineKarrasSchedulerTesterMixin,
PipelineTesterMixin,
unittest.TestCase,
):
class StableDiffusionPAGImg2ImgPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = StableDiffusionPAGImg2ImgPipeline
params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - {"height", "width"}
required_optional_params = PipelineTesterMixin.required_optional_params - {"latents"}
batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS
image_latents_params = IMAGE_TO_IMAGE_IMAGE_PARAMS
required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union(
{"pag_scale", "pag_adaptive_scale"}
) - {"height", "width"}
batch_input_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS
# Img2img derives the latents from the input image, so `__call__` takes no `latents`.
optional_input_params = frozenset(
["num_inference_steps", "num_images_per_prompt", "generator", "output_type", "return_dict"]
)
# The output resolution follows the 32x32 input image.
output_shape = (3, 32, 32)
def get_dummy_components(self, time_cond_proj_dim=None):
torch.manual_seed(0)
@@ -117,7 +113,7 @@ class StableDiffusionPAGImg2ImgPipelineFastTests(
text_encoder = CLIPTextModel(text_encoder_config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
components = {
return {
"unet": unet,
"scheduler": scheduler,
"vae": vae,
@@ -127,111 +123,58 @@ class StableDiffusionPAGImg2ImgPipelineFastTests(
"feature_extractor": None,
"image_encoder": None,
}
return components
def get_dummy_tiny_autoencoder(self):
return AutoencoderTiny(in_channels=3, out_channels=3, latent_channels=4)
def get_dummy_inputs(self, device, seed=0):
image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device)
def get_dummy_inputs(self):
image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device)
image = image / 2 + 0.5
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device=device).manual_seed(seed)
inputs = {
return {
"prompt": "A painting of a squirrel eating a burger",
"image": image,
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 6.0,
"pag_scale": 0.9,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
}
return inputs
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline (expect same output when pag is disabled)
pipe_sd = StableDiffusionImg2ImgPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
# pag enabled
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3
def test_pag_inference(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
32,
32,
3,
), f"the shape of the output image should be (1, 32, 32, 3) but got {image.shape}"
expected_slice = np.array(
[0.4508819, 0.49191576, 0.42664427, 0.66448534, 0.5606137, 0.43760118, 0.58251137, 0.5944448, 0.51642907]
)
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
self.assertLessEqual(max_diff, 1e-3)
class TestStableDiffusionPAGImg2ImgPipeline(StableDiffusionPAGImg2ImgPipelineTesterConfig, PAGPipelineTesterMixin):
base_pipeline_class = StableDiffusionImg2ImgPipeline
# fmt: off
expected_pag_slice = torch.tensor([0.4508819, 0.49191576, 0.42664427, 0.66448534, 0.5606137, 0.43760118, 0.58251137, 0.5944448, 0.51642907])
# fmt: on
def test_encode_prompt_works_in_isolation(self):
extra_required_param_value_dict = {
"device": torch.device(torch_device).type,
"do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0,
"do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0,
}
return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict)
class TestStableDiffusionPAGImg2ImgPipelineMemory(StableDiffusionPAGImg2ImgPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD PAG img2img pipeline."""
class TestStableDiffusionPAGImg2ImgPipelineIPAdapter(
StableDiffusionPAGImg2ImgPipelineTesterConfig, IPAdapterTesterMixin
):
"""IP-Adapter tests for the SD PAG img2img pipeline."""
@slow
@require_torch_accelerator
class StableDiffusionPAGImg2ImgPipelineIntegrationTests(unittest.TestCase):
class TestStableDiffusionPAGImg2ImgPipelineIntegration:
pipeline_class = StableDiffusionPAGImg2ImgPipeline
repo_id = "Jiali/stable-diffusion-1.5"
def setUp(self):
super().setUp()
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
yield
gc.collect()
backend_empty_cache(torch_device)

120
tests/pipelines/pag/test_pag_sd_inpaint.py
View File

@@ -15,9 +15,9 @@
import gc
import random
import unittest
import numpy as np
import pytest
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
@@ -26,6 +26,7 @@ from diffusers import (
AutoencoderKL,
AutoPipelineForInpainting,
PNDMScheduler,
StableDiffusionInpaintPipeline,
StableDiffusionPAGInpaintPipeline,
UNet2DConditionModel,
)
@@ -44,32 +45,29 @@ from ..pipeline_params import (
TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS,
)
from ..test_pipelines_common import (
from ..testing_utils import (
BasePipelineTesterConfig,
FromPipeTesterMixin,
IPAdapterTesterMixin,
PipelineFromPipeTesterMixin,
PipelineLatentTesterMixin,
PipelineTesterMixin,
MemoryTesterMixin,
)
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class StableDiffusionPAGInpaintPipelineFastTests(
PipelineTesterMixin,
IPAdapterTesterMixin,
PipelineLatentTesterMixin,
PipelineFromPipeTesterMixin,
unittest.TestCase,
):
class StableDiffusionPAGInpaintPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = StableDiffusionPAGInpaintPipeline
params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
image_params = frozenset([])
image_latents_params = frozenset([])
required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS.union(
{"pag_scale", "pag_adaptive_scale"}
)
batch_input_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union(
{"add_text_embeds", "add_time_ids", "mask", "masked_image_latents"}
)
# The output resolution follows the 64x64 input image.
output_shape = (3, 64, 64)
def get_dummy_components(self, time_cond_proj_dim=None):
torch.manual_seed(0)
@@ -109,7 +107,7 @@ class StableDiffusionPAGInpaintPipelineFastTests(
text_encoder = CLIPTextModel(text_encoder_config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
components = {
return {
"unet": unet,
"scheduler": scheduler,
"vae": vae,
@@ -119,42 +117,39 @@ class StableDiffusionPAGInpaintPipelineFastTests(
"feature_extractor": None,
"image_encoder": None,
}
return components
def get_dummy_inputs(self, device, seed=0):
def get_dummy_inputs(self):
# TODO: use tensor inputs instead of PIL, this is here just to leave the old expected_slices untouched
image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device)
image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device)
image = image.cpu().permute(0, 2, 3, 1)[0]
init_image = Image.fromarray(np.uint8(image)).convert("RGB").resize((64, 64))
# create mask
image[8:, 8:, :] = 255
mask_image = Image.fromarray(np.uint8(image)).convert("L").resize((64, 64))
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device=device).manual_seed(seed)
inputs = {
return {
"prompt": "A painting of a squirrel eating a burger",
"image": init_image,
"mask_image": mask_image,
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 6.0,
"strength": 1.0,
"pag_scale": 0.9,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
}
return inputs
class TestStableDiffusionPAGInpaintPipeline(StableDiffusionPAGInpaintPipelineTesterConfig, PAGPipelineTesterMixin):
base_pipeline_class = StableDiffusionInpaintPipeline
# fmt: off
expected_pag_slice = torch.tensor([0.7173, 0.5821, 0.6031, 0.5765, 0.6412, 0.6558, 0.5803, 0.5675, 0.5246])
# fmt: on
def test_pag_applied_layers(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline
pipe = self.pipeline_class(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()
# pag_applied_layers = ["mid","up","down"] should apply to all self-attention layers
all_self_attn_layers = [k for k in pipe.unet.attn_processors.keys() if "attn1" in k]
@@ -192,7 +187,7 @@ class StableDiffusionPAGInpaintPipelineFastTests(
# pag_applied_layers = ["mid.block_0.attentions_1"] does not exist in the model
pipe.unet.set_attn_processor(original_attn_procs.copy())
pag_layers = ["mid_block.attentions.1"]
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
# pag_applied_layers = "down" should apply to all self-attention layers in down_blocks
@@ -207,7 +202,7 @@ class StableDiffusionPAGInpaintPipelineFastTests(
pipe.unet.set_attn_processor(original_attn_procs.copy())
pag_layers = ["down_blocks.0"]
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
pipe.unet.set_attn_processor(original_attn_procs.copy())
@@ -220,50 +215,41 @@ class StableDiffusionPAGInpaintPipelineFastTests(
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
assert len(pipe.pag_attn_processors) == 1
def test_pag_inference(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
64,
64,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array([0.7173, 0.5821, 0.6031, 0.5765, 0.6412, 0.6558, 0.5803, 0.5675, 0.5246])
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}"
def test_encode_prompt_works_in_isolation(self):
extra_required_param_value_dict = {
"device": torch.device(torch_device).type,
"do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0,
"do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0,
}
return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict, atol=1e-3, rtol=1e-3)
class TestStableDiffusionPAGInpaintPipelineMemory(StableDiffusionPAGInpaintPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD PAG inpaint pipeline."""
class TestStableDiffusionPAGInpaintPipelineIPAdapter(
StableDiffusionPAGInpaintPipelineTesterConfig, IPAdapterTesterMixin
):
"""IP-Adapter tests for the SD PAG inpaint pipeline."""
class TestStableDiffusionPAGInpaintPipelineFromPipe(
StableDiffusionPAGInpaintPipelineTesterConfig, FromPipeTesterMixin
):
"""`from_pipe` round-trip tests against `StableDiffusionPipeline`."""
@slow
@require_torch_accelerator
class StableDiffusionPAGPipelineIntegrationTests(unittest.TestCase):
class TestStableDiffusionPAGInpaintPipelineIntegration:
pipeline_class = StableDiffusionPAGInpaintPipeline
repo_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
def setUp(self):
super().setUp()
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
yield
gc.collect()
backend_empty_cache(torch_device)

144
tests/pipelines/pag/test_pag_sdxl.py
View File

@@ -14,10 +14,9 @@
# limitations under the License.
import gc
import inspect
import unittest
import numpy as np
import pytest
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
@@ -41,33 +40,26 @@ from ...testing_utils import (
from ..pipeline_params import (
TEXT_TO_IMAGE_BATCH_PARAMS,
TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS,
TEXT_TO_IMAGE_IMAGE_PARAMS,
TEXT_TO_IMAGE_PARAMS,
)
from ..test_pipelines_common import (
from ..testing_utils import (
BasePipelineTesterConfig,
FromPipeTesterMixin,
IPAdapterTesterMixin,
PipelineFromPipeTesterMixin,
PipelineLatentTesterMixin,
PipelineTesterMixin,
MemoryTesterMixin,
)
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class StableDiffusionXLPAGPipelineFastTests(
PipelineTesterMixin,
IPAdapterTesterMixin,
PipelineLatentTesterMixin,
PipelineFromPipeTesterMixin,
unittest.TestCase,
):
class StableDiffusionXLPAGPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = StableDiffusionXLPAGPipeline
params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS
callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union({"add_text_embeds", "add_time_ids"})
output_shape = (3, 64, 64)
def get_dummy_components(self, time_cond_proj_dim=None):
# Copied from tests.pipelines.stable_diffusion_xl.test_stable_diffusion_xl.StableDiffusionXLPipelineTesterConfig.get_dummy_components
@@ -129,7 +121,7 @@ class StableDiffusionXLPAGPipelineFastTests(
text_encoder_2 = CLIPTextModelWithProjection(text_encoder_config)
tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
components = {
return {
"unet": unet,
"scheduler": scheduler,
"vae": vae,
@@ -140,67 +132,28 @@ class StableDiffusionXLPAGPipelineFastTests(
"image_encoder": None,
"feature_extractor": None,
}
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,
"pag_scale": 0.9,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
}
return inputs
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline (expect same output when pag is disabled)
pipe_sd = StableDiffusionXLPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
# pag enabled
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3
class TestStableDiffusionXLPAGPipeline(StableDiffusionXLPAGPipelineTesterConfig, PAGPipelineTesterMixin):
base_pipeline_class = StableDiffusionXLPipeline
# fmt: off
expected_pag_slice = torch.tensor([0.5565, 0.5305, 0.4652, 0.4330, 0.4823, 0.4640, 0.5191, 0.4983, 0.4684])
# fmt: on
def test_pag_applied_layers(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
# base pipeline
pipe = self.pipeline_class(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()
# pag_applied_layers = ["mid","up","down"] should apply to all self-attention layers
all_self_attn_layers = [k for k in pipe.unet.attn_processors.keys() if "attn1" in k]
@@ -234,7 +187,7 @@ class StableDiffusionXLPAGPipelineFastTests(
# pag_applied_layers = ["mid.block_0.attentions_1"] does not exist in the model
pipe.unet.set_attn_processor(original_attn_procs.copy())
pag_layers = ["mid_block.attentions.1"]
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
# pag_applied_layers = "down" should apply to all self-attention layers in down_blocks
@@ -249,7 +202,7 @@ class StableDiffusionXLPAGPipelineFastTests(
pipe.unet.set_attn_processor(original_attn_procs.copy())
pag_layers = ["down_blocks.0"]
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
pipe.unet.set_attn_processor(original_attn_procs.copy())
@@ -262,47 +215,34 @@ class StableDiffusionXLPAGPipelineFastTests(
pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False)
assert len(pipe.pag_attn_processors) == 2
def test_pag_inference(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
64,
64,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array([0.5565, 0.5305, 0.4652, 0.4330, 0.4823, 0.4640, 0.5191, 0.4983, 0.4684])
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
self.assertLessEqual(max_diff, 1e-3)
@unittest.skip("We test this functionality elsewhere already.")
@pytest.mark.skip("We test this functionality elsewhere already.")
def test_save_load_optional_components(self):
pass
class TestStableDiffusionXLPAGPipelineMemory(StableDiffusionXLPAGPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the SDXL PAG pipeline."""
class TestStableDiffusionXLPAGPipelineIPAdapter(StableDiffusionXLPAGPipelineTesterConfig, IPAdapterTesterMixin):
"""IP-Adapter tests for the SDXL PAG pipeline."""
class TestStableDiffusionXLPAGPipelineFromPipe(StableDiffusionXLPAGPipelineTesterConfig, FromPipeTesterMixin):
"""`from_pipe` round-trip tests against `StableDiffusionXLPipeline`."""
@slow
@require_torch_accelerator
class StableDiffusionXLPAGPipelineIntegrationTests(unittest.TestCase):
class TestStableDiffusionXLPAGPipelineIntegration:
pipeline_class = StableDiffusionXLPAGPipeline
repo_id = "stabilityai/stable-diffusion-xl-base-1.0"
def setUp(self):
super().setUp()
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
yield
gc.collect()
backend_empty_cache(torch_device)

142
tests/pipelines/pag/test_pag_sdxl_img2img.py
View File

@@ -14,11 +14,10 @@
# limitations under the License.
import gc
import inspect
import random
import unittest
import numpy as np
import pytest
import torch
from transformers import (
CLIPImageProcessor,
@@ -49,37 +48,33 @@ from ...testing_utils import (
torch_device,
)
from ..pipeline_params import (
IMAGE_TO_IMAGE_IMAGE_PARAMS,
TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS,
)
from ..test_pipelines_common import (
from ..testing_utils import (
BasePipelineTesterConfig,
FromPipeTesterMixin,
IPAdapterTesterMixin,
PipelineFromPipeTesterMixin,
PipelineLatentTesterMixin,
PipelineTesterMixin,
MemoryTesterMixin,
)
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class StableDiffusionXLPAGImg2ImgPipelineFastTests(
PipelineTesterMixin,
IPAdapterTesterMixin,
PipelineLatentTesterMixin,
PipelineFromPipeTesterMixin,
unittest.TestCase,
):
class StableDiffusionXLPAGImg2ImgPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = StableDiffusionXLPAGImg2ImgPipeline
params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - {"height", "width"}
batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS
image_latents_params = IMAGE_TO_IMAGE_IMAGE_PARAMS
required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union(
{"pag_scale", "pag_adaptive_scale"}
) - {"height", "width"}
batch_input_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union(
{"add_text_embeds", "add_time_ids", "add_neg_time_ids"}
)
# The output resolution follows the 32x32 input image.
output_shape = (3, 32, 32)
# based on tests.pipelines.stable_diffusion_xl.test_stable_diffusion_xl_img2img_pipeline.get_dummy_components
def get_dummy_components(
@@ -183,101 +178,64 @@ class StableDiffusionXLPAGImg2ImgPipelineFastTests(
# based on tests.pipelines.stable_diffusion_xl.test_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipelineTesterConfig
# add `pag_scale` to the inputs
def get_dummy_inputs(self, device, seed=0):
image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device)
def get_dummy_inputs(self):
image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device)
image = image / 2 + 0.5
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device=device).manual_seed(seed)
inputs = {
return {
"prompt": "A painting of a squirrel eating a burger",
"image": image,
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 5.0,
"pag_scale": 3.0,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
"strength": 0.8,
}
return inputs
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components(requires_aesthetics_score=True)
# base pipeline
pipe_sd = StableDiffusionXLImg2ImgPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
class TestStableDiffusionXLPAGImg2ImgPipeline(StableDiffusionXLPAGImg2ImgPipelineTesterConfig, PAGPipelineTesterMixin):
base_pipeline_class = StableDiffusionXLImg2ImgPipeline
# The expected slice below was recorded against the aesthetics-score configuration.
pag_component_kwargs = {"requires_aesthetics_score": True}
# fmt: off
expected_pag_slice = torch.tensor([0.4566, 0.4907, 0.4374, 0.6633, 0.5626, 0.4494, 0.5771, 0.6011, 0.5245])
# fmt: on
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
# pag enabled
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3
def test_pag_inference(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components(requires_aesthetics_score=True)
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
32,
32,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array([0.4566, 0.4907, 0.4374, 0.6633, 0.5626, 0.4494, 0.5771, 0.6011, 0.5245])
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}"
@unittest.skip("We test this functionality elsewhere already.")
@pytest.mark.skip("We test this functionality elsewhere already.")
def test_save_load_optional_components(self):
pass
class TestStableDiffusionXLPAGImg2ImgPipelineMemory(
StableDiffusionXLPAGImg2ImgPipelineTesterConfig, MemoryTesterMixin
):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the SDXL PAG img2img pipeline."""
class TestStableDiffusionXLPAGImg2ImgPipelineIPAdapter(
StableDiffusionXLPAGImg2ImgPipelineTesterConfig, IPAdapterTesterMixin
):
"""IP-Adapter tests for the SDXL PAG img2img pipeline."""
class TestStableDiffusionXLPAGImg2ImgPipelineFromPipe(
StableDiffusionXLPAGImg2ImgPipelineTesterConfig, FromPipeTesterMixin
):
"""`from_pipe` round-trip tests against `StableDiffusionXLPipeline`."""
@slow
@require_torch_accelerator
class StableDiffusionXLPAGImg2ImgPipelineIntegrationTests(unittest.TestCase):
class TestStableDiffusionXLPAGImg2ImgPipelineIntegration:
repo_id = "stabilityai/stable-diffusion-xl-base-1.0"
def setUp(self):
super().setUp()
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
yield
gc.collect()
backend_empty_cache(torch_device)

141
tests/pipelines/pag/test_pag_sdxl_inpaint.py
View File

@@ -14,11 +14,10 @@
# limitations under the License.
import gc
import inspect
import random
import unittest
import numpy as np
import pytest
import torch
from PIL import Image
from transformers import (
@@ -54,32 +53,29 @@ from ..pipeline_params import (
TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS,
)
from ..test_pipelines_common import (
from ..testing_utils import (
BasePipelineTesterConfig,
FromPipeTesterMixin,
IPAdapterTesterMixin,
PipelineFromPipeTesterMixin,
PipelineLatentTesterMixin,
PipelineTesterMixin,
MemoryTesterMixin,
)
from .testing_utils import PAGPipelineTesterMixin
enable_full_determinism()
class StableDiffusionXLPAGInpaintPipelineFastTests(
PipelineTesterMixin,
IPAdapterTesterMixin,
PipelineLatentTesterMixin,
PipelineFromPipeTesterMixin,
unittest.TestCase,
):
class StableDiffusionXLPAGInpaintPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = StableDiffusionXLPAGInpaintPipeline
params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
image_params = frozenset([])
image_latents_params = frozenset([])
required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS.union(
{"pag_scale", "pag_adaptive_scale"}
)
batch_input_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union(
{"add_text_embeds", "add_time_ids", "mask", "masked_image_latents"}
)
# The output resolution follows the 64x64 input image.
output_shape = (3, 64, 64)
# based on tests.pipelines.stable_diffusion_xl.test_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipelineTesterConfig.get_dummy_components
def get_dummy_components(
@@ -181,108 +177,71 @@ class StableDiffusionXLPAGInpaintPipelineFastTests(
}
return components
def get_dummy_inputs(self, device, seed=0):
def get_dummy_inputs(self):
# TODO: use tensor inputs instead of PIL, this is here just to leave the old expected_slices untouched
image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device)
image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device)
image = image.cpu().permute(0, 2, 3, 1)[0]
init_image = Image.fromarray(np.uint8(image)).convert("RGB").resize((64, 64))
# create mask
image[8:, 8:, :] = 255
mask_image = Image.fromarray(np.uint8(image)).convert("L").resize((64, 64))
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device=device).manual_seed(seed)
inputs = {
return {
"prompt": "A painting of a squirrel eating a burger",
"image": init_image,
"mask_image": mask_image,
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 6.0,
"strength": 1.0,
"pag_scale": 0.9,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
}
return inputs
def test_pag_disable_enable(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components(requires_aesthetics_score=True)
# base pipeline
pipe_sd = StableDiffusionXLInpaintPipeline(**components)
pipe_sd = pipe_sd.to(device)
pipe_sd.set_progress_bar_config(disable=None)
class TestStableDiffusionXLPAGInpaintPipeline(StableDiffusionXLPAGInpaintPipelineTesterConfig, PAGPipelineTesterMixin):
base_pipeline_class = StableDiffusionXLInpaintPipeline
# The expected slice below was recorded against the aesthetics-score configuration.
pag_component_kwargs = {"requires_aesthetics_score": True}
# fmt: off
expected_pag_slice = torch.tensor([0.7893, 0.5446, 0.5826, 0.6441, 0.6660, 0.7566, 0.6605, 0.5838, 0.5160])
# fmt: on
inputs = self.get_dummy_inputs(device)
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}."
)
out = pipe_sd(**inputs).images[0, -3:, -3:, -1]
# pag disabled with pag_scale=0.0
pipe_pag = self.pipeline_class(**components)
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs["pag_scale"] = 0.0
out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
# pag enabled
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1]
assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3
assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3
def test_pag_inference(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components(requires_aesthetics_score=True)
pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"])
pipe_pag = pipe_pag.to(device)
pipe_pag.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
image = pipe_pag(**inputs).images
image_slice = image[0, -3:, -3:, -1]
assert image.shape == (
1,
64,
64,
3,
), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}"
expected_slice = np.array([0.7893, 0.5446, 0.5826, 0.6441, 0.6660, 0.7566, 0.6605, 0.5838, 0.5160])
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}"
@unittest.skip("We test this functionality elsewhere already.")
@pytest.mark.skip("We test this functionality elsewhere already.")
def test_save_load_optional_components(self):
pass
class TestStableDiffusionXLPAGInpaintPipelineMemory(
StableDiffusionXLPAGInpaintPipelineTesterConfig, MemoryTesterMixin
):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the SDXL PAG inpaint pipeline."""
class TestStableDiffusionXLPAGInpaintPipelineIPAdapter(
StableDiffusionXLPAGInpaintPipelineTesterConfig, IPAdapterTesterMixin
):
"""IP-Adapter tests for the SDXL PAG inpaint pipeline."""
class TestStableDiffusionXLPAGInpaintPipelineFromPipe(
StableDiffusionXLPAGInpaintPipelineTesterConfig, FromPipeTesterMixin
):
"""`from_pipe` round-trip tests against `StableDiffusionXLPipeline`."""
@slow
@require_torch_accelerator
class StableDiffusionXLPAGInpaintPipelineIntegrationTests(unittest.TestCase):
class TestStableDiffusionXLPAGInpaintPipelineIntegration:
repo_id = "stabilityai/stable-diffusion-xl-base-1.0"
def setUp(self):
super().setUp()
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
yield
gc.collect()
backend_empty_cache(torch_device)

121
tests/pipelines/pag/testing_utils.py Normal file
View File

@@ -0,0 +1,121 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import pytest
from ...testing_utils import assert_tensors_close, torch_device
from ..testing_utils import PipelineTesterMixin
class PAGPipelineTesterMixin(PipelineTesterMixin):
"""`PipelineTesterMixin` plus the two tests every PAG pipeline shares.
A PAG pipeline is an existing pipeline with perturbed-attention guidance layered on, so the tests all take the
same shape: run the pipeline it derives from, then check that PAG at `pag_scale=0.0` reproduces that output and
that enabling PAG moves it. Subclasses supply the knobs below; anything pipeline-specific (which layers PAG
resolves to, for instance) stays a method on the concrete test class.
"""
# The non-PAG pipeline this one derives from. Required.
base_pipeline_class = None
# `pag_applied_layers` for the "PAG enabled" leg of `test_pag_disable_enable`. `None` keeps the pipeline default.
pag_enabled_applied_layers = ["mid", "up", "down"]
# `pag_scale` for the "PAG enabled" leg. `None` keeps the value from `get_dummy_inputs()`.
pag_enabled_scale = None
# Some PAG pipelines only assert the disabled leg: their dummy denoiser is small enough that PAG's effect is not
# reliably above the tolerance. Set to `False` there.
check_pag_changes_output = True
# `pag_applied_layers` the `test_pag_inference` pipeline is built with. `None` keeps the pipeline default.
pag_inference_applied_layers = ["mid", "up", "down"]
# CPU-specific expected slice for `test_pag_inference`, laid out as the flattened `output[0, -1, -3:, -3:]`
# corner of the `"pt"` output. `None` skips the test.
expected_pag_slice = None
# Extra kwargs for the `get_dummy_components()` call the two tests below build their pipelines from, when the
# PAG comparison needs a configuration other than the default one — the SDXL img2img and inpaint testers pin
# `requires_aesthetics_score=True`, which is what their expected slices were recorded against.
pag_component_kwargs = {}
def get_pag_components(self):
return self.get_dummy_components(**self.pag_component_kwargs)
def get_pag_pipeline(self, components=None, **pag_kwargs):
"""Build the pipeline under test with explicit PAG constructor kwargs (`pag_applied_layers`, ...)."""
components = components if components is not None else self.get_dummy_components()
pipe = self.pipeline_class(**components, **pag_kwargs)
pipe.set_progress_bar_config(disable=None)
return pipe
def test_pag_disable_enable(self):
# Run on CPU to keep the device-dependent `torch.Generator` deterministic.
components = self.get_pag_components()
# base pipeline (expect same output when pag is disabled)
pipe_base = self.base_pipeline_class(**components)
pipe_base.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs()
del inputs["pag_scale"]
assert "pag_scale" not in inspect.signature(pipe_base.__call__).parameters, (
f"`pag_scale` should not be a call parameter of the base pipeline {pipe_base.__class__.__name__}."
)
out = pipe_base(**inputs)[0]
# pag disabled with pag_scale=0.0
pipe_pag = self.get_pipeline(**self.get_pag_components())
out_pag_disabled = self.run_pipe(pipe_pag, pag_scale=0.0)
assert_tensors_close(out_pag_disabled, out, atol=1e-3, msg="PAG at `pag_scale=0.0` changed the output.")
if not self.check_pag_changes_output:
return
# pag enabled
pag_kwargs = {}
if self.pag_enabled_applied_layers is not None:
pag_kwargs["pag_applied_layers"] = self.pag_enabled_applied_layers
pipe_pag = self.get_pag_pipeline(self.get_pag_components(), **pag_kwargs)
extra_inputs = {} if self.pag_enabled_scale is None else {"pag_scale": self.pag_enabled_scale}
out_pag_enabled = self.run_pipe(pipe_pag, **extra_inputs)
assert (out - out_pag_enabled).abs().max() > 1e-3, "Enabling PAG should change the output."
def test_pag_inference(self):
if self.expected_pag_slice is None:
pytest.skip(f"No CPU expected slice pinned for {self.pipeline_class.__name__}.")
if torch_device != "cpu":
pytest.skip("The expected slice is CPU-specific.")
pag_kwargs = (
{}
if self.pag_inference_applied_layers is None
else {"pag_applied_layers": self.pag_inference_applied_layers}
)
pipe_pag = self.get_pag_pipeline(self.get_pag_components(), **pag_kwargs)
image = pipe_pag(**self.get_dummy_inputs())[0]
assert image.shape == (1, *self.output_shape), (
f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}"
)
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), self.expected_pag_slice, atol=1e-3)

234
tests/pipelines/pixart_alpha/test_pixart.py
View File

@@ -14,10 +14,9 @@
# limitations under the License.
import gc
import tempfile
import unittest
import numpy as np
import pytest
import torch
from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
@@ -30,6 +29,7 @@ from diffusers import (
from diffusers.utils.import_utils import is_torch_neuronx_available
from ...testing_utils import (
assert_tensors_close,
backend_empty_cache,
backend_synchronize,
enable_full_determinism,
@@ -39,23 +39,19 @@ from ...testing_utils import (
slow,
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
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
enable_full_determinism()
class PixArtAlphaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class PixArtAlphaPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = PixArtAlphaPipeline
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 = PipelineTesterMixin.required_optional_params
test_layerwise_casting = True
test_group_offloading = True
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS
# `transformer.sample_size` (8) * `vae_scale_factor` (8) / 8 -> the dummy transformer generates 8x8 images.
output_shape = (3, 8, 8)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -87,98 +83,74 @@ class PixArtAlphaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
return {
"transformer": transformer.eval(),
"vae": vae.eval(),
"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": "A painting of a squirrel eating a burger",
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 5.0,
"use_resolution_binning": False,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
}
return inputs
@unittest.skip("Not supported.")
def test_sequential_cpu_offload_forward_pass(self):
# TODO(PVP, Sayak) need to fix later
return
class TestPixArtAlphaPipeline(PixArtAlphaPipelineTesterConfig, PipelineTesterMixin):
def test_inference(self):
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)
image = pipe(**self.get_dummy_inputs()).images
assert image.shape == (1, *self.output_shape)
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs).images
image_slice = image[0, -3:, -3:, -1]
self.assertEqual(image.shape, (1, 8, 8, 3))
expected_slice = np.array([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.483, 0.2583, 0.5331, 0.4852])
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
self.assertLessEqual(max_diff, 1e-3)
# fmt: off
expected_slice = torch.tensor([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.483, 0.2583, 0.5331, 0.4852])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
def test_inference_non_square_images(self):
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)
image = pipe(**self.get_dummy_inputs(), height=32, width=48).images
assert image.shape == (1, 3, 32, 48)
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs, height=32, width=48).images
image_slice = image[0, -3:, -3:, -1]
self.assertEqual(image.shape, (1, 32, 48, 3))
# fmt: off
expected_slice = torch.tensor([0.6493, 0.537, 0.4081, 0.4762, 0.3695, 0.4711, 0.3026, 0.5218, 0.5263])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
expected_slice = np.array([0.6493, 0.537, 0.4081, 0.4762, 0.3695, 0.4711, 0.3026, 0.5218, 0.5263])
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
self.assertLessEqual(max_diff, 1e-3)
@unittest.skip("Test is already covered through encode_prompt isolation.")
@pytest.mark.skip("Test is already covered through encode_prompt isolation.")
def test_save_load_optional_components(self):
pass
def test_inference_with_embeddings_and_multiple_images(self):
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
def test_inference_with_embeddings_and_multiple_images(self, tmp_path):
pipe = self.get_pipeline().to(torch_device)
inputs = self.get_dummy_inputs(torch_device)
prompt = inputs["prompt"]
generator = inputs["generator"]
num_inference_steps = inputs["num_inference_steps"]
output_type = inputs["output_type"]
prompt_embeds, prompt_attn_mask, negative_prompt_embeds, neg_prompt_attn_mask = pipe.encode_prompt(prompt)
inputs = self.get_dummy_inputs()
prompt_embeds, prompt_attn_mask, negative_prompt_embeds, neg_prompt_attn_mask = pipe.encode_prompt(
inputs["prompt"]
)
# inputs with prompt converted to embeddings
inputs = {
embedding_inputs = {
"prompt_embeds": prompt_embeds,
"prompt_attention_mask": prompt_attn_mask,
"negative_prompt": None,
"negative_prompt_embeds": negative_prompt_embeds,
"negative_prompt_attention_mask": neg_prompt_attn_mask,
"generator": generator,
"num_inference_steps": num_inference_steps,
"output_type": output_type,
"generator": inputs["generator"],
"num_inference_steps": inputs["num_inference_steps"],
"output_type": inputs["output_type"],
"num_images_per_prompt": 2,
"use_resolution_binning": False,
}
@@ -187,97 +159,76 @@ class PixArtAlphaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
for optional_component in pipe._optional_components:
setattr(pipe, optional_component, None)
output = pipe(**inputs)[0]
output = pipe(**embedding_inputs)[0]
with tempfile.TemporaryDirectory() as tmpdir:
pipe.save_pretrained(tmpdir)
pipe_loaded = self.pipeline_class.from_pretrained(tmpdir)
pipe_loaded.to(torch_device)
pipe_loaded.set_progress_bar_config(disable=None)
pipe.save_pretrained(tmp_path)
pipe_loaded = self.pipeline_class.from_pretrained(tmp_path)
pipe_loaded.to(torch_device)
pipe_loaded.set_progress_bar_config(disable=None)
for optional_component in pipe._optional_components:
self.assertTrue(
getattr(pipe_loaded, optional_component) is None,
f"`{optional_component}` did not stay set to None after loading.",
assert getattr(pipe_loaded, optional_component) is None, (
f"`{optional_component}` did not stay set to None after loading."
)
inputs = self.get_dummy_inputs(torch_device)
embedding_inputs["generator"] = self.get_generator(0)
output_loaded = pipe_loaded(**embedding_inputs)[0]
generator = inputs["generator"]
num_inference_steps = inputs["num_inference_steps"]
output_type = inputs["output_type"]
# inputs with prompt converted to embeddings
inputs = {
"prompt_embeds": prompt_embeds,
"prompt_attention_mask": prompt_attn_mask,
"negative_prompt": None,
"negative_prompt_embeds": negative_prompt_embeds,
"negative_prompt_attention_mask": neg_prompt_attn_mask,
"generator": generator,
"num_inference_steps": num_inference_steps,
"output_type": output_type,
"num_images_per_prompt": 2,
"use_resolution_binning": False,
}
output_loaded = pipe_loaded(**inputs)[0]
max_diff = np.abs(to_np(output) - to_np(output_loaded)).max()
self.assertLess(max_diff, 1e-4)
assert_tensors_close(
output_loaded, output, atol=1e-4, msg="Output changed after dropping optional components."
)
def test_inference_with_multiple_images_per_prompt(self):
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)
image = pipe(**self.get_dummy_inputs(), num_images_per_prompt=2).images
assert image.shape == (2, *self.output_shape)
inputs = self.get_dummy_inputs(device)
inputs["num_images_per_prompt"] = 2
image = pipe(**inputs).images
image_slice = image[0, -3:, -3:, -1]
self.assertEqual(image.shape, (2, 8, 8, 3))
expected_slice = np.array([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.483, 0.2583, 0.5331, 0.4852])
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
self.assertLessEqual(max_diff, 1e-3)
# fmt: off
expected_slice = torch.tensor([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.483, 0.2583, 0.5331, 0.4852])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
def test_raises_warning_for_mask_feature(self):
device = "cpu"
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()
inputs.update({"mask_feature": True})
with self.assertWarns(FutureWarning) as warning_ctx:
with pytest.warns(FutureWarning, match="mask_feature"):
_ = pipe(**inputs).images
assert "mask_feature" in str(warning_ctx.warning)
def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(expected_max_diff=1e-3)
super().test_inference_batch_single_identical(expected_max_diff=1e-3)
class TestPixArtAlphaPipelineMemory(PixArtAlphaPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the PixArt-alpha pipeline."""
@pytest.mark.skip("Not supported.")
def test_sequential_cpu_offload_forward_pass(self):
# TODO(PVP, Sayak) need to fix later
pass
@pytest.mark.skip("Not supported.")
def test_sequential_offload_forward_pass_twice(self):
# TODO(PVP, Sayak) need to fix later
pass
@slow
@require_torch_accelerator
class PixArtAlphaPipelineIntegrationTests(unittest.TestCase):
class TestPixArtAlphaPipelineIntegration:
ckpt_id_1024 = "PixArt-alpha/PixArt-XL-2-1024-MS"
ckpt_id_512 = "PixArt-alpha/PixArt-XL-2-512x512"
prompt = "A small cactus with a happy face in the Sahara desert."
def setUp(self):
super().setUp()
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
yield
gc.collect()
backend_empty_cache(torch_device)
@@ -296,7 +247,7 @@ class PixArtAlphaPipelineIntegrationTests(unittest.TestCase):
max_diff = numpy_cosine_similarity_distance(image_slice.flatten(), expected_slice)
# Neuron uses bfloat16 internally which has lower precision than float16 on CUDA
atol = 1e-2 if is_torch_neuronx_available() else 1e-4
self.assertLessEqual(max_diff, atol)
assert max_diff <= atol
def test_pixart_512(self):
generator = torch.Generator("cpu").manual_seed(0)
@@ -314,7 +265,7 @@ class PixArtAlphaPipelineIntegrationTests(unittest.TestCase):
max_diff = numpy_cosine_similarity_distance(image_slice.flatten(), expected_slice)
# Neuron uses bfloat16 internally which has lower precision than float16 on CUDA
atol = 1e-2 if is_torch_neuronx_available() else 1e-4
self.assertLessEqual(max_diff, atol)
assert max_diff <= atol
def test_pixart_1024_without_resolution_binning(self):
generator = torch.manual_seed(0)
@@ -419,9 +370,6 @@ class PixArtAlphaPipelineIntegrationTests(unittest.TestCase):
output_type="np",
).images
self.assertEqual(image.shape, (1, 512, 512, 3))
self.assertFalse(np.isnan(image).any(), "Output contains NaN values")
self.assertTrue(
(image >= 0.0).all() and (image <= 1.0).all(),
"Output pixel values outside [0, 1]",
)
assert image.shape == (1, 512, 512, 3)
assert not np.isnan(image).any(), "Output contains NaN values"
assert (image >= 0.0).all() and (image <= 1.0).all(), "Output pixel values outside [0, 1]"

254
tests/pipelines/pixart_sigma/test_pixart.py
View File

@@ -14,10 +14,9 @@
# limitations under the License.
import gc
import tempfile
import unittest
import numpy as np
import pytest
import torch
from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
@@ -30,6 +29,7 @@ from diffusers import (
from ...testing_utils import (
Expectations,
assert_tensors_close,
backend_empty_cache,
enable_full_determinism,
numpy_cosine_similarity_distance,
@@ -37,28 +37,25 @@ from ...testing_utils import (
slow,
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 (
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..testing_utils import (
BasePipelineTesterConfig,
MemoryTesterMixin,
PipelineTesterMixin,
check_qkv_fusion_matches_attn_procs_length,
check_qkv_fusion_processors_exist,
to_np,
)
enable_full_determinism()
class PixArtSigmaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class PixArtSigmaPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = PixArtSigmaPipeline
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 = PipelineTesterMixin.required_optional_params
test_layerwise_casting = True
test_group_offloading = True
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS
# `transformer.sample_size` (8) * `vae_scale_factor` (8) / 8 -> the dummy transformer generates 8x8 images.
output_shape = (3, 8, 8)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -90,94 +87,70 @@ class PixArtSigmaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
return {
"transformer": transformer.eval(),
"vae": vae.eval(),
"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": "A painting of a squirrel eating a burger",
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 5.0,
"use_resolution_binning": False,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
}
return inputs
@unittest.skip("Not supported.")
def test_sequential_cpu_offload_forward_pass(self):
# TODO(PVP, Sayak) need to fix later
return
class TestPixArtSigmaPipeline(PixArtSigmaPipelineTesterConfig, PipelineTesterMixin):
def test_inference(self):
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)
image = pipe(**self.get_dummy_inputs()).images
assert image.shape == (1, *self.output_shape)
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs).images
image_slice = image[0, -3:, -3:, -1]
self.assertEqual(image.shape, (1, 8, 8, 3))
expected_slice = np.array([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.4830, 0.2583, 0.5331, 0.4852])
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
self.assertLessEqual(max_diff, 1e-3)
# fmt: off
expected_slice = torch.tensor([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.4830, 0.2583, 0.5331, 0.4852])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
def test_inference_non_square_images(self):
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)
image = pipe(**self.get_dummy_inputs(), height=32, width=48).images
assert image.shape == (1, 3, 32, 48)
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs, height=32, width=48).images
image_slice = image[0, -3:, -3:, -1]
self.assertEqual(image.shape, (1, 32, 48, 3))
# fmt: off
expected_slice = torch.tensor([0.6493, 0.5370, 0.4081, 0.4762, 0.3695, 0.4711, 0.3026, 0.5218, 0.5263])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
expected_slice = np.array([0.6493, 0.5370, 0.4081, 0.4762, 0.3695, 0.4711, 0.3026, 0.5218, 0.5263])
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
self.assertLessEqual(max_diff, 1e-3)
def test_inference_with_embeddings_and_multiple_images(self, tmp_path):
pipe = self.get_pipeline().to(torch_device)
def test_inference_with_embeddings_and_multiple_images(self):
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(torch_device)
prompt = inputs["prompt"]
generator = inputs["generator"]
num_inference_steps = inputs["num_inference_steps"]
output_type = inputs["output_type"]
prompt_embeds, prompt_attn_mask, negative_prompt_embeds, neg_prompt_attn_mask = pipe.encode_prompt(prompt)
inputs = self.get_dummy_inputs()
prompt_embeds, prompt_attn_mask, negative_prompt_embeds, neg_prompt_attn_mask = pipe.encode_prompt(
inputs["prompt"]
)
# inputs with prompt converted to embeddings
inputs = {
embedding_inputs = {
"prompt_embeds": prompt_embeds,
"prompt_attention_mask": prompt_attn_mask,
"negative_prompt": None,
"negative_prompt_embeds": negative_prompt_embeds,
"negative_prompt_attention_mask": neg_prompt_attn_mask,
"generator": generator,
"num_inference_steps": num_inference_steps,
"output_type": output_type,
"generator": inputs["generator"],
"num_inference_steps": inputs["num_inference_steps"],
"output_type": inputs["output_type"],
"num_images_per_prompt": 2,
"use_resolution_binning": False,
}
@@ -186,80 +159,49 @@ class PixArtSigmaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
for optional_component in pipe._optional_components:
setattr(pipe, optional_component, None)
output = pipe(**inputs)[0]
output = pipe(**embedding_inputs)[0]
with tempfile.TemporaryDirectory() as tmpdir:
pipe.save_pretrained(tmpdir)
pipe_loaded = self.pipeline_class.from_pretrained(tmpdir)
pipe_loaded.to(torch_device)
pipe_loaded.set_progress_bar_config(disable=None)
pipe.save_pretrained(tmp_path)
pipe_loaded = self.pipeline_class.from_pretrained(tmp_path)
pipe_loaded.to(torch_device)
pipe_loaded.set_progress_bar_config(disable=None)
for optional_component in pipe._optional_components:
self.assertTrue(
getattr(pipe_loaded, optional_component) is None,
f"`{optional_component}` did not stay set to None after loading.",
assert getattr(pipe_loaded, optional_component) is None, (
f"`{optional_component}` did not stay set to None after loading."
)
inputs = self.get_dummy_inputs(torch_device)
embedding_inputs["generator"] = self.get_generator(0)
output_loaded = pipe_loaded(**embedding_inputs)[0]
generator = inputs["generator"]
num_inference_steps = inputs["num_inference_steps"]
output_type = inputs["output_type"]
# inputs with prompt converted to embeddings
inputs = {
"prompt_embeds": prompt_embeds,
"prompt_attention_mask": prompt_attn_mask,
"negative_prompt": None,
"negative_prompt_embeds": negative_prompt_embeds,
"negative_prompt_attention_mask": neg_prompt_attn_mask,
"generator": generator,
"num_inference_steps": num_inference_steps,
"output_type": output_type,
"num_images_per_prompt": 2,
"use_resolution_binning": False,
}
output_loaded = pipe_loaded(**inputs)[0]
max_diff = np.abs(to_np(output) - to_np(output_loaded)).max()
self.assertLess(max_diff, 1e-4)
assert_tensors_close(
output_loaded, output, atol=1e-4, msg="Output changed after dropping optional components."
)
def test_inference_with_multiple_images_per_prompt(self):
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)
image = pipe(**self.get_dummy_inputs(), num_images_per_prompt=2).images
assert image.shape == (2, *self.output_shape)
inputs = self.get_dummy_inputs(device)
inputs["num_images_per_prompt"] = 2
image = pipe(**inputs).images
image_slice = image[0, -3:, -3:, -1]
# fmt: off
expected_slice = torch.tensor([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.4830, 0.2583, 0.5331, 0.4852])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
self.assertEqual(image.shape, (2, 8, 8, 3))
expected_slice = np.array([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.4830, 0.2583, 0.5331, 0.4852])
max_diff = np.abs(image_slice.flatten() - expected_slice).max()
self.assertLessEqual(max_diff, 1e-3)
@unittest.skip("Test is already covered through encode_prompt isolation.")
@pytest.mark.skip("Test is already covered through encode_prompt isolation.")
def test_save_load_optional_components(self):
pass
def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(expected_max_diff=1e-3)
super().test_inference_batch_single_identical(expected_max_diff=1e-3)
def test_fused_qkv_projections(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=None)
# Run on CPU to keep the device-dependent `torch.Generator` deterministic.
pipe = self.get_pipeline()
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs).images
original_image_slice = image[0, -3:, -3:, -1]
original_output = self.run_pipe(pipe)
# TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added
# to the pipeline level.
@@ -271,40 +213,56 @@ class PixArtSigmaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
pipe.transformer, pipe.transformer.original_attn_processors
), "Something wrong with the attention processors concerning the fused QKV projections."
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs).images
image_slice_fused = image[0, -3:, -3:, -1]
output_fused = self.run_pipe(pipe)
pipe.transformer.unfuse_qkv_projections()
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs).images
image_slice_disabled = image[0, -3:, -3:, -1]
output_disabled = self.run_pipe(pipe)
assert np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), (
"Fusion of QKV projections shouldn't affect the outputs."
assert_tensors_close(
output_fused, original_output, atol=1e-3, rtol=1e-3, msg="Fusion of QKV projections changed the outputs."
)
assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), (
"Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled."
assert_tensors_close(
output_disabled,
output_fused,
atol=1e-3,
rtol=1e-3,
msg="Outputs changed after the fused QKV projections were disabled.",
)
assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), (
"Original outputs should match when fused QKV projections are disabled."
assert_tensors_close(
output_disabled,
original_output,
atol=1e-2,
rtol=1e-2,
msg="Original outputs should match when fused QKV projections are disabled.",
)
class TestPixArtSigmaPipelineMemory(PixArtSigmaPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the PixArt-sigma pipeline."""
@pytest.mark.skip("Not supported.")
def test_sequential_cpu_offload_forward_pass(self):
# TODO(PVP, Sayak) need to fix later
pass
@pytest.mark.skip("Not supported.")
def test_sequential_offload_forward_pass_twice(self):
# TODO(PVP, Sayak) need to fix later
pass
@slow
@require_torch_accelerator
class PixArtSigmaPipelineIntegrationTests(unittest.TestCase):
class TestPixArtSigmaPipelineIntegration:
ckpt_id_1024 = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS"
ckpt_id_512 = "PixArt-alpha/PixArt-Sigma-XL-2-512-MS"
prompt = "A small cactus with a happy face in the Sahara desert."
def setUp(self):
super().setUp()
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
yield
gc.collect()
backend_empty_cache(torch_device)
@@ -321,7 +279,7 @@ class PixArtSigmaPipelineIntegrationTests(unittest.TestCase):
expected_slice = np.array([0.4517, 0.4446, 0.4375, 0.449, 0.4399, 0.4365, 0.4583, 0.4629, 0.4473])
max_diff = numpy_cosine_similarity_distance(image_slice.flatten(), expected_slice)
self.assertLessEqual(max_diff, 1e-4)
assert max_diff <= 1e-4
def test_pixart_512(self):
generator = torch.Generator("cpu").manual_seed(0)
@@ -349,7 +307,7 @@ class PixArtSigmaPipelineIntegrationTests(unittest.TestCase):
expected_slice = expected_slices.get_expectation()
max_diff = numpy_cosine_similarity_distance(image_slice.flatten(), expected_slice)
self.assertLessEqual(max_diff, 1e-4)
assert max_diff <= 1e-4
def test_pixart_1024_without_resolution_binning(self):
generator = torch.manual_seed(0)

154
tests/pipelines/pndm/test_pndm.py
View File

@@ -13,61 +13,159 @@
# 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 diffusers import PNDMPipeline, PNDMScheduler, UNet2DModel
from ...testing_utils import enable_full_determinism, nightly, require_torch, torch_device
from ...testing_utils import (
enable_full_determinism,
nightly,
require_accelerator,
require_torch,
torch_device,
)
from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
enable_full_determinism()
class PNDMPipelineFastTests(unittest.TestCase):
@property
def dummy_uncond_unet(self):
# `PNDMPipeline.__call__` postprocesses unconditionally — it always runs `image.cpu().permute(0, 2, 3, 1).numpy()`
# and only branches afterwards to decide whether to wrap the result in PIL. There is no `output_type="pt"` path, so
# `get_dummy_inputs` below has to ask for `"np"`, and every shared test that compares outputs with
# `assert_tensors_close` (torch-only) or calls `torch.isnan` on them fails on the numpy array it gets back.
#
# The sibling unconditional pipelines already have a `"pt"` path (`DDIMPipeline`, `DDPMPipeline`); adding one here
# is a `src/` change and out of scope for this test migration, so the affected tests are marked `xfail` rather than
# skipped: whoever adds the `"pt"` branch will see them XPASS and can drop these markers.
NO_PT_OUTPUT = pytest.mark.xfail(
reason="`PNDMPipeline` has no `output_type='pt'` path and always returns a numpy array.",
strict=True,
)
# `PNDMPipeline` samples its initial noise with `randn_tensor(..., device=self.device)`: no `dtype=`, so the noise
# stays float32 and disagrees with a half-precision or layerwise-cast UNet. Same story as above — a `src/` gap, so
# the accelerator-only tests it breaks are marked `xfail`.
UNSUPPORTED_DTYPE = pytest.mark.xfail(
reason="`PNDMPipeline` samples its initial noise without `dtype=self.unet.dtype`, so it stays float32.",
strict=True,
)
# The memory mixin trips over three separate `src/` gaps at once, so its tests carry one marker between them:
# - the numpy output above (`test_group_offloading_inference`, `test_pipeline_level_group_offloading_inference`,
# `test_pipeline_with_accelerator_device_map`),
# - the float32 noise above (`test_layerwise_casting_inference`),
# - `randn_tensor(..., device=self.device)` reading the *pipeline's* device rather than `self._execution_device`,
# which under sequential offload is `meta` (`test_sequential_cpu_offload_forward_pass`,
# `test_sequential_offload_forward_pass_twice`),
# - and no `model_cpu_offload_seq`, which `enable_model_cpu_offload` requires
# (`test_model_cpu_offload_forward_pass`, `test_cpu_offload_forward_pass_twice`).
# `strict=False` because `test_pipeline_level_group_offloading_sanity_checks` never runs the pipeline and so passes
# — it reports XPASS. The class-level marker is what keeps `MemoryTesterMixin`'s own `@is_memory` /
# `@require_accelerator` marks intact; overriding the eight failing tests individually would drop the
# `@require_accelerate_version_greater` gates they are declared with.
UNSUPPORTED_MEMORY_OPTIMIZATIONS = pytest.mark.xfail(
reason=(
"`PNDMPipeline` returns numpy, samples noise without `dtype=`/`_execution_device`, and declares no "
"`model_cpu_offload_seq`."
),
strict=False,
)
class PNDMPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = PNDMPipeline
required_input_params_in_call_signature = UNCONDITIONAL_IMAGE_GENERATION_PARAMS
batch_input_params = UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS
# PNDM is unconditional and samples its own noise: there is no prompt to repeat
# (`num_images_per_prompt`) and no user-suppliable `latents`.
optional_input_params = frozenset(["num_inference_steps", "generator", "output_type", "return_dict"])
# `(height, width, channels)` — the numpy layout, not the `(channels, height, width)` the other configs get
# from `output_type="pt"` (see `NO_PT_OUTPUT` above).
output_shape = (8, 8, 3)
def get_dummy_components(self):
torch.manual_seed(0)
model = UNet2DModel(
block_out_channels=(32, 64),
layers_per_block=2,
sample_size=32,
unet = UNet2DModel(
block_out_channels=(4, 8),
layers_per_block=1,
norm_num_groups=4,
sample_size=8,
in_channels=3,
out_channels=3,
down_block_types=("DownBlock2D", "AttnDownBlock2D"),
up_block_types=("AttnUpBlock2D", "UpBlock2D"),
)
return model
def test_inference(self):
unet = self.dummy_uncond_unet
scheduler = PNDMScheduler()
return {"unet": unet, "scheduler": scheduler}
pndm = PNDMPipeline(unet=unet, scheduler=scheduler)
pndm.to(torch_device)
pndm.set_progress_bar_config(disable=None)
def get_dummy_inputs(self):
return {
"batch_size": 1,
"generator": self.get_generator(0),
# `PNDMScheduler` runs Runge-Kutta warm-up steps, so it needs at least 4 inference steps.
"num_inference_steps": 4,
# `"np"` rather than the usual `"pt"` — see `NO_PT_OUTPUT` above.
"output_type": "np",
}
generator = torch.manual_seed(0)
image = pndm(generator=generator, num_inference_steps=20, output_type="np").images
generator = torch.manual_seed(0)
image_from_tuple = pndm(generator=generator, num_inference_steps=20, output_type="np", return_dict=False)[0]
class TestPNDMPipeline(PNDMPipelineTesterConfig, PipelineTesterMixin):
def test_inference(self):
# Run on CPU: the expected slice below is CPU-specific.
pipe = self.get_pipeline()
image_slice = image[0, -3:, -3:, -1]
image_from_tuple_slice = image_from_tuple[0, -3:, -3:, -1]
image = pipe(**self.get_dummy_inputs()).images
generated_image = image[0]
assert generated_image.shape == self.output_shape
assert image.shape == (1, 32, 32, 3)
expected_slice = np.array([1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0])
expected_slice = np.array([0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0])
assert np.abs(generated_image[-3:, -3:, -1].flatten() - expected_slice).max() < 1e-2
assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2
@NO_PT_OUTPUT
def test_save_load_local(self, tmp_path, base_pipe_output, expected_max_difference=5e-4):
super().test_save_load_local(tmp_path, base_pipe_output, expected_max_difference)
@NO_PT_OUTPUT
def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical()
@NO_PT_OUTPUT
def test_dict_tuple_outputs_equivalent(self):
super().test_dict_tuple_outputs_equivalent()
# The three overrides below re-declare the base methods' skip decorators: overriding a test drops the marks
# the base declared it with, and without them these would run (and xfail for the wrong reason) on CPU.
@NO_PT_OUTPUT
@require_accelerator
def test_to_device(self):
super().test_to_device()
@UNSUPPORTED_DTYPE
@pytest.mark.skipif(torch_device not in ["cuda", "xpu"], reason="half-precision inference requires CUDA or XPU")
@require_accelerator
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=str)
def test_half_precision_inference_no_nan(self, dtype):
super().test_half_precision_inference_no_nan(dtype)
@UNSUPPORTED_DTYPE
@pytest.mark.skipif(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU")
@require_accelerator
def test_save_load_float16(self, tmp_path, expected_max_diff=1e-2):
super().test_save_load_float16(tmp_path, expected_max_diff)
@UNSUPPORTED_MEMORY_OPTIMIZATIONS
class TestPNDMPipelineMemory(PNDMPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the PNDM pipeline."""
@nightly
@require_torch
class PNDMPipelineIntegrationTests(unittest.TestCase):
class TestPNDMPipelineIntegration:
def test_inference_cifar10(self):
model_id = "google/ddpm-cifar10-32"

199
tests/pipelines/prx/test_pipeline_prx.py
View File

@@ -1,6 +1,4 @@
import unittest
import numpy as np
import pytest
import torch
from transformers import AutoTokenizer
from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig, T5GemmaModuleConfig
@@ -11,26 +9,32 @@ from diffusers.models.transformers.transformer_prx import PRXTransformer2DModel
from diffusers.pipelines.prx.pipeline_prx import PRXPipeline
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from ...testing_utils import assert_tensors_close
from ..pipeline_params import TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
class PRXPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
# `T5GemmaEncoder` is instantiated here from a hand-built config rather than loaded from a repo, and transformers v5
# cannot round-trip that through `save_pretrained`/`from_pretrained`, so every test that reloads the pipeline from
# disk is skipped.
T5GEMMA_SERIALIZATION_SKIP_REASON = "Custom T5GemmaEncoder not compatible with transformers v5."
# Both PRX pipelines read `callback_on_step_end`'s inputs out of `locals()` but throw away what the callback
# returns, so a callback that rewrites `latents` (or `prompt_embeds`) has no effect on the denoising loop. Every
# other diffusers pipeline pops those keys back off the returned dict. Fixing it is a `src/` change and out of
# scope for this test migration, so the one shared test that exercises the write-back is marked `xfail`: whoever
# adds the pop-back will see it XPASS and can drop this marker.
CALLBACK_OUTPUTS_IGNORED = pytest.mark.xfail(
reason="`PRX` pipelines discard the dict `callback_on_step_end` returns, so callback edits to `latents` are lost.",
strict=True,
)
class PRXPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = PRXPipeline
params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_params = frozenset(["prompt", "negative_prompt", "num_images_per_prompt"])
test_xformers_attention = False
test_layerwise_casting = True
test_group_offloading = True
@classmethod
def setUpClass(cls):
# Ensure PRXPipeline has an _execution_device property expected by __call__
if not isinstance(getattr(PRXPipeline, "_execution_device", None), property):
try:
setattr(PRXPipeline, "_execution_device", property(lambda self: torch.device("cpu")))
except Exception:
pass
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_input_params = frozenset(["prompt", "negative_prompt", "num_images_per_prompt"])
output_shape = (3, 32, 32)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -102,116 +106,54 @@ class PRXPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
"tokenizer": tokenizer,
}
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)
def get_dummy_inputs(self):
return {
"prompt": "",
"negative_prompt": "",
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 1.0,
"height": 32,
"width": 32,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
"use_resolution_binning": False,
}
def test_inference(self):
device = "cpu"
components = self.get_dummy_components()
pipe = PRXPipeline(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
try:
pipe.register_to_config(_execution_device="cpu")
except Exception:
pass
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs)[0]
generated_image = image[0]
self.assertEqual(generated_image.shape, (3, 32, 32))
expected_image = torch.zeros(3, 32, 32)
max_diff = np.abs(generated_image - expected_image).max()
self.assertLessEqual(max_diff, 1e10)
class TestPRXPipeline(PRXPipelineTesterConfig, PipelineTesterMixin):
@CALLBACK_OUTPUTS_IGNORED
def test_callback_inputs(self):
components = self.get_dummy_components()
pipe = PRXPipeline(**components)
pipe = pipe.to("cpu")
pipe.set_progress_bar_config(disable=None)
try:
pipe.register_to_config(_execution_device="cpu")
except Exception:
pass
self.assertTrue(
hasattr(pipe, "_callback_tensor_inputs"),
f" {PRXPipeline} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs",
)
def callback_inputs_subset(pipe, i, t, callback_kwargs):
for tensor_name in callback_kwargs.keys():
assert tensor_name in pipe._callback_tensor_inputs
return callback_kwargs
def callback_inputs_all(pipe, i, t, callback_kwargs):
for tensor_name in pipe._callback_tensor_inputs:
assert tensor_name in callback_kwargs
for tensor_name in callback_kwargs.keys():
assert tensor_name in pipe._callback_tensor_inputs
return callback_kwargs
inputs = self.get_dummy_inputs("cpu")
inputs["callback_on_step_end"] = callback_inputs_subset
inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
_ = pipe(**inputs)[0]
inputs["callback_on_step_end"] = callback_inputs_all
inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
_ = pipe(**inputs)[0]
super().test_callback_inputs()
def test_attention_slicing_forward_pass(self, expected_max_diff=1e-3):
if not self.test_attention_slicing:
return
# Run on CPU: sliced attention is compared against a full-attention run of the same pipeline.
pipe = self.get_pipeline()
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("cpu")
pipe.set_progress_bar_config(disable=None)
def to_np_local(tensor):
if isinstance(tensor, torch.Tensor):
return tensor.detach().cpu().numpy()
return tensor
generator_device = "cpu"
inputs = self.get_dummy_inputs(generator_device)
output_without_slicing = pipe(**inputs)[0]
output_without_slicing = self.run_pipe(pipe)
pipe.enable_attention_slicing(slice_size=1)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing1 = pipe(**inputs)[0]
output_with_slicing_1 = self.run_pipe(pipe)
pipe.enable_attention_slicing(slice_size=2)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing2 = pipe(**inputs)[0]
output_with_slicing_2 = self.run_pipe(pipe)
max_diff1 = np.abs(to_np_local(output_with_slicing1) - to_np_local(output_without_slicing)).max()
max_diff2 = np.abs(to_np_local(output_with_slicing2) - to_np_local(output_without_slicing)).max()
self.assertLess(max(max_diff1, max_diff2), expected_max_diff)
assert_tensors_close(
output_with_slicing_1,
output_without_slicing,
atol=expected_max_diff,
msg="Attention slicing (slice_size=1) changed the output.",
)
assert_tensors_close(
output_with_slicing_2,
output_without_slicing,
atol=expected_max_diff,
msg="Attention slicing (slice_size=2) changed the output.",
)
def test_inference_with_autoencoder_dc(self):
"""Test PRXPipeline with AutoencoderDC (DCAE) instead of AutoencoderKL."""
device = "cpu"
"""PRXPipeline should also work with an `AutoencoderDC` (DCAE) in place of the `AutoencoderKL`."""
components = self.get_dummy_components()
torch.manual_seed(0)
@@ -240,43 +182,38 @@ class PRXPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
).eval()
components["vae"] = vae_dc
pipe = self.get_pipeline(**components)
pipe = PRXPipeline(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
assert pipe.vae_scale_factor == vae_dc.spatial_compression_ratio
expected_scale_factor = vae_dc.spatial_compression_ratio
self.assertEqual(pipe.vae_scale_factor, expected_scale_factor)
output = self.run_pipe(pipe)
assert output[0].shape == self.output_shape
assert torch.isfinite(output).all()
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs)[0]
generated_image = image[0]
self.assertEqual(generated_image.shape, (3, 32, 32))
expected_image = torch.zeros(3, 32, 32)
max_diff = np.abs(generated_image - expected_image).max()
self.assertLessEqual(max_diff, 1e10)
@unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
@pytest.mark.skip(T5GEMMA_SERIALIZATION_SKIP_REASON)
def test_loading_with_variants(self):
pass
@unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
def test_pipeline_with_accelerator_device_map(self):
pass
@unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
@pytest.mark.skip(T5GEMMA_SERIALIZATION_SKIP_REASON)
def test_save_load_local(self):
pass
@unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
@pytest.mark.skip(T5GEMMA_SERIALIZATION_SKIP_REASON)
def test_save_load_float16(self):
pass
@pytest.mark.skip(T5GEMMA_SERIALIZATION_SKIP_REASON)
def test_save_load_optional_components(self):
pass
@unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
def test_dtype_dict(self):
@pytest.mark.skip(T5GEMMA_SERIALIZATION_SKIP_REASON)
def test_torch_dtype_dict(self):
pass
@unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
def test_dtype_alias(self):
class TestPRXPipelineMemory(PRXPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the PRX pipeline."""
@pytest.mark.skip(T5GEMMA_SERIALIZATION_SKIP_REASON)
def test_pipeline_with_accelerator_device_map(self):
pass

230
tests/pipelines/prx/test_pipeline_prx_pixel.py
View File

@@ -1,6 +1,4 @@
import unittest
import numpy as np
import pytest
import torch
from transformers import Qwen2Tokenizer, Qwen3Config, Qwen3Model
@@ -8,19 +6,29 @@ from diffusers.models.transformers.transformer_prx import PRXTransformer2DModel
from diffusers.pipelines.prx.pipeline_prx_pixel import PRXPixelPipeline
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from ...testing_utils import assert_tensors_close, torch_device
from ..pipeline_params import TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
class PRXPixelPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
# Both PRX pipelines read `callback_on_step_end`'s inputs out of `locals()` but throw away what the callback
# returns, so a callback that rewrites `latents` (or `prompt_embeds`) has no effect on the denoising loop. Every
# other diffusers pipeline pops those keys back off the returned dict. Fixing it is a `src/` change and out of
# scope for this test migration, so the one shared test that exercises the write-back is marked `xfail`: whoever
# adds the pop-back will see it XPASS and can drop this marker.
CALLBACK_OUTPUTS_IGNORED = pytest.mark.xfail(
reason="`PRX` pipelines discard the dict `callback_on_step_end` returns, so callback edits to `latents` are lost.",
strict=True,
)
class PRXPixelPipelineTesterConfig(BasePipelineTesterConfig):
# PRXPixelPipeline is standalone: it inherits from DiffusionPipeline (not PRXPipeline) and always has its own
# image_processor, so it denoises raw RGB in pixel space and supports output_type="pil"/"np" without a VAE.
pipeline_class = PRXPixelPipeline
params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_params = frozenset(["prompt", "negative_prompt", "num_images_per_prompt"])
test_xformers_attention = False
test_layerwise_casting = True
test_group_offloading = True
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_input_params = frozenset(["prompt", "negative_prompt", "num_images_per_prompt"])
output_shape = (3, 32, 32)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -64,108 +72,75 @@ class PRXPixelPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
"prompt_max_tokens": 16,
}
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)
def get_dummy_inputs(self):
return {
"prompt": "",
"negative_prompt": "",
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 1.0,
"height": 32,
"width": 32,
# Pixel-space PRX has no VAE and returns raw (C, H, W) tensors for output_type="pt". The generic
# PipelineTesterMixin tests compare these tensors directly, so default to "pt" here; the PIL/np default
# path is exercised explicitly in test_inference and test_inference_pil_and_np_output.
# path is exercised explicitly in test_inference_pil_and_np_output.
"output_type": "pt",
# 32px is not in the 1024 aspect-ratio bins, so binning must be disabled for these tiny fast tests.
"use_resolution_binning": False,
}
def _build_pipe(self, device="cpu"):
components = self.get_dummy_components()
pipe = PRXPixelPipeline(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
return pipe
def test_inference(self):
device = "cpu"
pipe = self._build_pipe(device)
class TestPRXPixelPipeline(PRXPixelPipelineTesterConfig, PipelineTesterMixin):
@CALLBACK_OUTPUTS_IGNORED
def test_callback_inputs(self):
super().test_callback_inputs()
def test_pixel_space_has_no_vae(self):
# Pixel space: vae_scale_factor is always 1, and the pipeline always carries an image processor
# so postprocessing (and the default output_type="pil") works without any VAE.
self.assertEqual(pipe.vae_scale_factor, 1)
self.assertIsNotNone(pipe.image_processor)
pipe = self.get_pipeline()
# Default output is PIL (no VAE needed: the image processor denormalizes the denoised pixels directly).
inputs = self.get_dummy_inputs(device)
inputs.pop("output_type") # default is "pil"
images = pipe(**inputs).images
self.assertEqual(len(images), 1)
self.assertEqual(images[0].size, (32, 32))
# Raw "pt" output is the denoised RGB tensor at the requested resolution.
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs)[0]
generated_image = image[0]
self.assertEqual(generated_image.shape, (3, 32, 32))
expected_image = torch.zeros(3, 32, 32)
max_diff = np.abs(generated_image.cpu().numpy() - expected_image.numpy()).max()
self.assertLessEqual(max_diff, 1e10)
assert pipe.vae_scale_factor == 1
assert pipe.image_processor is not None
def test_inference_batch(self):
device = "cpu"
pipe = self._build_pipe(device)
pipe = self.get_pipeline()
inputs = self.get_dummy_inputs(device)
inputs["prompt"] = ["", ""]
inputs["negative_prompt"] = ["", ""]
image = pipe(**inputs)[0]
image = self.run_pipe(pipe, prompt=["", ""], negative_prompt=["", ""])
self.assertEqual(image.shape[0], 2)
self.assertEqual(tuple(image.shape[1:]), (3, 32, 32))
assert image.shape[0] == 2
assert tuple(image.shape[1:]) == self.output_shape
def test_inference_with_cfg(self):
device = "cpu"
pipe = self._build_pipe(device)
pipe = self.get_pipeline()
# CFG off.
inputs = self.get_dummy_inputs(device)
inputs["guidance_scale"] = 1.0
out_no_cfg = pipe(**inputs)[0]
self.assertFalse(pipe.do_classifier_free_guidance)
self.assertEqual(out_no_cfg[0].shape, (3, 32, 32))
out_no_cfg = self.run_pipe(pipe, guidance_scale=1.0)
assert not pipe.do_classifier_free_guidance
assert out_no_cfg[0].shape == self.output_shape
# CFG on.
inputs = self.get_dummy_inputs(device)
inputs["guidance_scale"] = 5.0
out_cfg = pipe(**inputs)[0]
self.assertTrue(pipe.do_classifier_free_guidance)
self.assertEqual(out_cfg[0].shape, (3, 32, 32))
out_cfg = self.run_pipe(pipe, guidance_scale=5.0)
assert pipe.do_classifier_free_guidance
assert out_cfg[0].shape == self.output_shape
# Guidance should actually change the output.
max_diff = np.abs(out_no_cfg.cpu().numpy() - out_cfg.cpu().numpy()).max()
self.assertGreater(max_diff, 0.0)
assert not torch.allclose(out_no_cfg, out_cfg)
def test_inference_with_prompt_embeds(self):
device = "cpu"
pipe = self._build_pipe(device)
pipe = self.get_pipeline()
# Precompute embeddings via the public encode_prompt API (CFG on so we get negatives too).
prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = (
pipe.encode_prompt(
prompt="a prompt",
device=device,
device=torch.device("cpu"),
do_classifier_free_guidance=True,
negative_prompt="",
)
)
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs.pop("prompt")
inputs.pop("negative_prompt")
inputs["guidance_scale"] = 5.0
@@ -175,31 +150,29 @@ class PRXPixelPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
inputs["negative_prompt_attention_mask"] = negative_prompt_attention_mask
image = pipe(**inputs)[0]
self.assertEqual(image[0].shape, (3, 32, 32))
assert image[0].shape == self.output_shape
def test_inference_pil_and_np_output(self):
# The default output_type="pil" must work without a VAE: the denoised pixels are denormalized
# directly by the image processor instead of being decoded.
device = "cpu"
pipe = self._build_pipe(device)
pipe = self.get_pipeline()
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs.pop("output_type") # default is "pil"
images = pipe(**inputs).images
self.assertEqual(len(images), 1)
self.assertEqual(images[0].size, (32, 32))
assert len(images) == 1
assert images[0].size == (32, 32)
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs["output_type"] = "np"
images = pipe(**inputs).images
self.assertEqual(images.shape, (1, 32, 32, 3))
self.assertGreaterEqual(images.min(), 0.0)
self.assertLessEqual(images.max(), 1.0)
assert images.shape == (1, 32, 32, 3)
assert images.min() >= 0.0
assert images.max() <= 1.0
def test_non_multiple_size_raises(self):
# height/width must be divisible by vae_scale_factor * transformer patch_size; check_inputs must raise
# a clear ValueError instead of letting the transformer fail on an invalid reshape mid-denoising.
device = "cpu"
components = self.get_dummy_components()
torch.manual_seed(0)
components["transformer"] = PRXTransformer2DModel(
@@ -214,82 +187,45 @@ class PRXPixelPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
bottleneck_size=8,
resolution_embeds=True,
)
pipe = PRXPixelPipeline(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline(**components)
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs["height"] = 31 # vae_scale_factor (1) * patch_size (2) = 2; 31 is not a multiple
with self.assertRaisesRegex(ValueError, "divisible"):
with pytest.raises(ValueError, match="divisible"):
pipe(**inputs)
def test_callback_inputs(self):
device = "cpu"
pipe = self._build_pipe(device)
self.assertTrue(
hasattr(pipe, "_callback_tensor_inputs"),
f" {PRXPixelPipeline} should have `_callback_tensor_inputs` that defines a list of tensor variables its"
" callback function can use as inputs",
)
def callback_inputs_subset(pipe, i, t, callback_kwargs):
for tensor_name in callback_kwargs.keys():
assert tensor_name in pipe._callback_tensor_inputs
return callback_kwargs
def callback_inputs_all(pipe, i, t, callback_kwargs):
for tensor_name in pipe._callback_tensor_inputs:
assert tensor_name in callback_kwargs
for tensor_name in callback_kwargs.keys():
assert tensor_name in pipe._callback_tensor_inputs
return callback_kwargs
inputs = self.get_dummy_inputs(device)
inputs["callback_on_step_end"] = callback_inputs_subset
inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
_ = pipe(**inputs)[0]
inputs = self.get_dummy_inputs(device)
inputs["callback_on_step_end"] = callback_inputs_all
inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
_ = pipe(**inputs)[0]
def test_attention_slicing_forward_pass(self, expected_max_diff=1e-3):
# Overridden: the mixin version calls assert_mean_pixel_difference, which assumes HWC image
# arrays. Pixel-space PRX has no VAE; compare raw (C, H, W) tensors directly ("pt") instead of
# going through PIL.
if not self.test_attention_slicing:
return
# Run on CPU: sliced attention is compared against a full-attention run of the same pipeline.
pipe = self.get_pipeline()
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("cpu")
pipe.set_progress_bar_config(disable=None)
def to_np_local(tensor):
if isinstance(tensor, torch.Tensor):
return tensor.detach().cpu().numpy()
return tensor
generator_device = "cpu"
inputs = self.get_dummy_inputs(generator_device)
output_without_slicing = pipe(**inputs)[0]
output_without_slicing = self.run_pipe(pipe)
pipe.enable_attention_slicing(slice_size=1)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing1 = pipe(**inputs)[0]
output_with_slicing_1 = self.run_pipe(pipe)
pipe.enable_attention_slicing(slice_size=2)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing2 = pipe(**inputs)[0]
output_with_slicing_2 = self.run_pipe(pipe)
max_diff1 = np.abs(to_np_local(output_with_slicing1) - to_np_local(output_without_slicing)).max()
max_diff2 = np.abs(to_np_local(output_with_slicing2) - to_np_local(output_without_slicing)).max()
self.assertLess(max(max_diff1, max_diff2), expected_max_diff)
assert_tensors_close(
output_with_slicing_1,
output_without_slicing,
atol=expected_max_diff,
msg="Attention slicing (slice_size=1) changed the output.",
)
assert_tensors_close(
output_with_slicing_2,
output_without_slicing,
atol=expected_max_diff,
msg="Attention slicing (slice_size=2) changed the output.",
)
@unittest.skip("Slow original-vs-diffusers parity test is optional and intentionally skipped for fast CI.")
def test_prx_pixel_original_parity(self):
pass
def test_encode_prompt_works_in_isolation(self):
extra_required_param_value_dict = {
"device": torch.device(torch_device).type,
"do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0,
}
return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict)
class TestPRXPixelPipelineMemory(PRXPixelPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the PRXPixel pipeline."""

2
tests/pipelines/testing_utils/__init__.py
View File

@@ -7,6 +7,7 @@ from .cache import (
TaylorSeerCacheTesterMixin,
)
from .common import BasePipelineTesterConfig, PipelineTesterMixin
from .from_pipe import FromPipeTesterMixin
from .ip_adapter import IPAdapterTesterMixin
from .lora import LoraMemoryTesterMixin, LoraTesterMixin, UNetLoraTesterMixin
from .memory import (
@@ -26,6 +27,7 @@ from .utils import (
__all__ = [
"BasePipelineTesterConfig",
"PipelineTesterMixin",
"FromPipeTesterMixin",
"IPAdapterTesterMixin",
"LoraTesterMixin",
"LoraMemoryTesterMixin",

208
tests/pipelines/testing_utils/from_pipe.py Normal file
View File

@@ -0,0 +1,208 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
from diffusers import KolorsPipeline, StableDiffusionPipeline, StableDiffusionXLPipeline
from ...testing_utils import (
assert_tensors_close,
require_accelerate_version_greater,
require_accelerator,
torch_device,
)
from .common import BasePipelineOutputMixin
# The repo and constructor kwargs used to build the "original" pipeline a variant is derived from, keyed by that
# original pipeline class.
ORIGINAL_PIPELINE_REPOS = {
StableDiffusionPipeline: ("hf-internal-testing/tiny-stable-diffusion-torch", {"requires_safety_checker": False}),
StableDiffusionXLPipeline: (
"hf-internal-testing/tiny-stable-diffusion-xl-pipe",
{"requires_aesthetics_score": True, "force_zeros_for_empty_prompt": False},
),
KolorsPipeline: ("hf-internal-testing/tiny-kolors-pipe", {"force_zeros_for_empty_prompt": False}),
}
class FromPipeTesterMixin(BasePipelineOutputMixin):
"""`DiffusionPipeline.from_pipe` tests for pipelines that are variants of an existing one.
Composed with `BasePipelineTesterConfig`, which supplies `pipeline_class`, `get_dummy_components()` and
`get_dummy_inputs()`. The pytest-style successor of `PipelineFromPipeTesterMixin` in
`tests/pipelines/test_pipelines_common.py`.
"""
# Set on the test class to pull the original pipeline from a repo other than the default for its class.
original_pipeline_repo = None
@property
def original_pipeline_class(self):
"""The pipeline this one is a variant of — the source `from_pipe` is expected to round-trip through."""
name = self.pipeline_class.__name__.lower()
if "xl" in name:
return StableDiffusionXLPipeline
elif "kolors" in name:
return KolorsPipeline
return StableDiffusionPipeline
def get_dummy_inputs_pipe(self):
inputs = self.get_dummy_inputs()
inputs["return_dict"] = False
return inputs
def get_dummy_inputs_for_pipe_original(self):
"""The dummy inputs, restricted to the parameters the original pipeline's `__call__` accepts."""
original_call_params = set(inspect.signature(self.original_pipeline_class.__call__).parameters.keys())
return {k: v for k, v in self.get_dummy_inputs_pipe().items() if k in original_call_params}
def _split_components(self, components):
"""Split the given components into what the original pipeline expects and what only this one does."""
original_expected_modules, _ = self.original_pipeline_class._get_signature_keys(self.original_pipeline_class)
# components of this pipeline that the original one also expects
original_pipe_components = {}
# components this pipeline doesn't have, but the original one expects
original_pipe_additional_components = {}
# components this pipeline has, but the original one doesn't expect
current_pipe_additional_components = {}
for name, component in components.items():
if name in original_expected_modules:
original_pipe_components[name] = component
else:
current_pipe_additional_components[name] = component
for name in original_expected_modules:
if name not in original_pipe_components:
if name in self.original_pipeline_class._optional_components:
original_pipe_additional_components[name] = None
else:
raise ValueError(f"missing required module for {self.original_pipeline_class.__name__}: {name}")
return (
{**original_pipe_components, **original_pipe_additional_components},
current_pipe_additional_components,
)
def _build_original_pipeline(self, components):
"""Build the original pipeline out of `components`, plus the components only this pipeline has.
Both pipelines are built from the *same* component instances — rebuilding them per pipeline would make the
comparison depend on `get_dummy_components()` being bit-identical across calls.
"""
original_components, current_pipe_additional_components = self._split_components(components)
pipe_original = self.original_pipeline_class(**original_components)
pipe_original.set_progress_bar_config(disable=None)
return pipe_original, current_pipe_additional_components
def test_from_pipe_consistent_config(self):
original_repo, original_kwargs = ORIGINAL_PIPELINE_REPOS[self.original_pipeline_class]
original_repo = self.original_pipeline_repo or original_repo
# create original_pipeline_class(sd/sdxl/kolors)
pipe_original = self.original_pipeline_class.from_pretrained(original_repo, **original_kwargs)
# original_pipeline_class -> pipeline_class
pipe_additional_components = {
name: component
for name, component in self.get_dummy_components().items()
if name not in pipe_original.components
}
pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components)
# pipeline_class -> original_pipeline_class
original_pipe_additional_components = {}
for name, component in pipe_original.components.items():
if name not in pipe.components or not isinstance(component, pipe.components[name].__class__):
original_pipe_additional_components[name] = component
pipe_original_2 = self.original_pipeline_class.from_pipe(pipe, **original_pipe_additional_components)
# compare the config
original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")}
original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")}
assert original_config_2 == original_config
def test_from_pipe_consistent_forward_pass(self, expected_max_diff=1e-3):
components = self.get_dummy_components()
pipe_original, current_pipe_additional_components = self._build_original_pipeline(components)
pipe_original.to(torch_device)
output_original = pipe_original(**self.get_dummy_inputs_for_pipe_original())[0]
# `from_pipe` must not repurpose the original pipeline's attention processors — PAG and friends install
# their own on the derived pipeline only.
original_attn_processor_types = {
name: {k: type(v) for k, v in component.attn_processors.items()}
for name, component in pipe_original.components.items()
if hasattr(component, "attn_processors")
}
pipe = self.get_pipeline(**components).to(torch_device)
output = pipe(**self.get_dummy_inputs_pipe())[0]
pipe_from_original = self.pipeline_class.from_pipe(pipe_original, **current_pipe_additional_components)
pipe_from_original.to(torch_device)
pipe_from_original.set_progress_bar_config(disable=None)
output_from_original = pipe_from_original(**self.get_dummy_inputs_pipe())[0]
assert_tensors_close(
output_from_original,
output,
atol=expected_max_diff,
msg="The outputs of the pipelines created with `from_pipe` and `__init__` are different.",
)
output_original_2 = pipe_original(**self.get_dummy_inputs_for_pipe_original())[0]
assert_tensors_close(
output_original_2,
output_original,
atol=expected_max_diff,
msg="`from_pipe` should not change the output of original pipeline.",
)
for name, expected_types in original_attn_processor_types.items():
component = pipe_original.components[name]
assert {k: type(v) for k, v in component.attn_processors.items()} == expected_types, (
f"`from_pipe` changed the attention processors of `{name}` in the original pipeline."
)
@require_accelerator
@require_accelerate_version_greater("0.14.0")
def test_from_pipe_consistent_forward_pass_cpu_offload(self, expected_max_diff=1e-3):
components = self.get_dummy_components()
# Build the original pipeline before running anything. Both pipelines share one scheduler object, and some
# `__init__`s edit that object in place: `StableDiffusionPipeline`, for one, rewrites a `steps_offset` of 0
# to 1. Building it later would put the two forward passes below on different schedules.
pipe_original, current_pipe_additional_components = self._build_original_pipeline(components)
pipe = self.get_pipeline(**components)
pipe.enable_model_cpu_offload(device=torch_device)
output = pipe(**self.get_dummy_inputs_pipe())[0]
pipe_from_original = self.pipeline_class.from_pipe(pipe_original, **current_pipe_additional_components)
pipe_from_original.set_progress_bar_config(disable=None)
pipe_from_original.enable_model_cpu_offload(device=torch_device)
output_from_original = pipe_from_original(**self.get_dummy_inputs_pipe())[0]
assert_tensors_close(
output_from_original,
output,
atol=expected_max_diff,
msg="The outputs of the pipelines created with `from_pipe` and `__init__` are different.",
)

19
tests/testing_utils.py
View File

@@ -1018,6 +1018,25 @@ def export_to_gif(image: list[PIL.Image.Image], output_gif_path: str = None) ->
return output_gif_path
@contextmanager
def skip_if_no_cudnn_engine():
"""
Skip the enclosing test when cuDNN has no kernel for an op the pipeline runs.
cuDNN does not ship an engine for every (op, dtype, layout) combination, and the set it covers differs between
cuDNN builds and GPU architectures. When none applies it raises `RuntimeError: GET was unable to find an engine
to execute this computation` — for instance for Sana's depthwise `Conv2d` in bfloat16. That is a property of the
runner, not of the code under test, so tests that can hit it are skipped there rather than failed. Any other
`RuntimeError` propagates untouched.
"""
try:
yield
except RuntimeError as e:
if "unable to find an engine" not in str(e):
raise
pytest.skip(f"cuDNN has no engine for this computation on {torch_device}: {e}")
@contextmanager
def buffered_writer(raw_f):
f = io.BufferedWriter(raw_f)