diff --git a/src/diffusers/models/transformers/prior_transformer.py b/src/diffusers/models/transformers/prior_transformer.py index ace2b529c4f2..f3890446e28e 100644 --- a/src/diffusers/models/transformers/prior_transformer.py +++ b/src/diffusers/models/transformers/prior_transformer.py @@ -318,5 +318,8 @@ def forward( return PriorTransformerOutput(predicted_image_embedding=predicted_image_embedding) def post_process_latents(self, prior_latents): - prior_latents = (prior_latents * self.clip_std) + self.clip_mean + # `clip_std` / `clip_mean` are parameters of this model, not of a submodule, so group offloading onloads + # them only for the duration of `forward`. This runs after the denoising loop, hence the explicit move. + device = prior_latents.device + prior_latents = (prior_latents * self.clip_std.to(device)) + self.clip_mean.to(device) return prior_latents diff --git a/tests/pipelines/kandinsky/test_kandinsky_prior.py b/tests/pipelines/kandinsky/test_kandinsky_prior.py index fc3f4c341545..52d1f2582073 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_prior.py +++ b/tests/pipelines/kandinsky/test_kandinsky_prior.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest import torch from torch import nn from transformers import ( @@ -38,25 +37,6 @@ enable_full_determinism() -# `PriorTransformer` keeps `positional_embedding`, `prd_embedding`, `clip_mean` and `clip_std` as parameters of the -# model itself rather than of a submodule, so group offloading never onloads them: the forward pass then mixes -# onloaded activations with still-offloaded weights. Reproduces at both block and leaf level. -PIPELINE_GROUP_OFFLOAD_XFAIL_REASON = ( - "`PriorTransformer` holds parameters directly on the model (`positional_embedding`, `prd_embedding`, " - "`clip_mean`, `clip_std`), which group offloading never onloads." -) - -# A second, independent gap: the component-scoped test only offloads the denoiser under the names -# `transformer`/`unet`/`controlnet`/`adapter`, and only puts `vae`/`vqvae`/`image_encoder` back on the accelerator. -# A prior pipeline's denoiser is called `prior`, so it matches neither list and is left on CPU while the text -# encoder is onloaded. Fixing this means widening the mixin's component lists, not changing the pipeline. -COMPONENT_GROUP_OFFLOAD_XFAIL_REASON = ( - "The pipeline calls `PriorTransformer.post_process_latents()` after the denoising loop, which reads the " - "`clip_mean` / `clip_std` parameters held directly on the model. Group offloading onloads those only for the " - "duration of `forward`, so by then they are back on the offload device." -) - - class KandinskyPriorPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = KandinskyPriorPipeline required_input_params_in_call_signature = frozenset(["prompt"]) @@ -214,23 +194,3 @@ def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff= class TestKandinskyPriorPipelineMemory(KandinskyPriorPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Kandinsky prior pipeline.""" - - @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True) - @MemoryTesterMixin._USE_STREAM - def test_group_offloading_inference_block_level(self, base_pipe_output, use_stream, expected_max_difference=1e-4): - super().test_group_offloading_inference_block_level( - base_pipe_output, use_stream, expected_max_difference=expected_max_difference - ) - - @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True) - @MemoryTesterMixin._USE_STREAM - def test_group_offloading_inference_leaf_level(self, base_pipe_output, use_stream, expected_max_difference=1e-4): - super().test_group_offloading_inference_leaf_level( - base_pipe_output, use_stream, expected_max_difference=expected_max_difference - ) - - @pytest.mark.xfail(condition=True, reason=PIPELINE_GROUP_OFFLOAD_XFAIL_REASON, strict=True) - def test_pipeline_level_group_offloading_inference(self, base_pipe_output, expected_max_difference=1e-4): - super().test_pipeline_level_group_offloading_inference( - base_pipe_output, expected_max_difference=expected_max_difference - ) diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_prior.py b/tests/pipelines/kandinsky2_2/test_kandinsky_prior.py index f73cceb7169d..5079c278493c 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_prior.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_prior.py @@ -40,25 +40,6 @@ enable_full_determinism() -# `PriorTransformer` keeps `positional_embedding`, `prd_embedding`, `clip_mean` and `clip_std` as parameters of the -# model itself rather than of a submodule, so group offloading never onloads them: the forward pass then mixes -# onloaded activations with still-offloaded weights. Reproduces at both block and leaf level. -PIPELINE_GROUP_OFFLOAD_XFAIL_REASON = ( - "`PriorTransformer` holds parameters directly on the model (`positional_embedding`, `prd_embedding`, " - "`clip_mean`, `clip_std`), which group offloading never onloads." -) - -# A second, independent gap: the component-scoped test only offloads the denoiser under the names -# `transformer`/`unet`/`controlnet`/`adapter`, and only puts `vae`/`vqvae`/`image_encoder` back on the accelerator. -# A prior pipeline's denoiser is called `prior`, so it matches neither list and is left on CPU while the text -# encoder is onloaded. Fixing this means widening the mixin's component lists, not changing the pipeline. -COMPONENT_GROUP_OFFLOAD_XFAIL_REASON = ( - "The pipeline calls `PriorTransformer.post_process_latents()` after the denoising loop, which reads the " - "`clip_mean` / `clip_std` parameters held directly on the model. Group offloading onloads those only for the " - "duration of `forward`, so by then they are back on the offload device." -) - - class KandinskyV22PriorPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = KandinskyV22PriorPipeline required_input_params_in_call_signature = frozenset(["prompt"]) @@ -249,23 +230,3 @@ def callback_inputs_test(pipe, i, t, callback_kwargs): class TestKandinskyV22PriorPipelineMemory(KandinskyV22PriorPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Kandinsky 2.2 prior pipeline.""" - - @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True) - @MemoryTesterMixin._USE_STREAM - def test_group_offloading_inference_block_level(self, base_pipe_output, use_stream, expected_max_difference=1e-4): - super().test_group_offloading_inference_block_level( - base_pipe_output, use_stream, expected_max_difference=expected_max_difference - ) - - @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True) - @MemoryTesterMixin._USE_STREAM - def test_group_offloading_inference_leaf_level(self, base_pipe_output, use_stream, expected_max_difference=1e-4): - super().test_group_offloading_inference_leaf_level( - base_pipe_output, use_stream, expected_max_difference=expected_max_difference - ) - - @pytest.mark.xfail(condition=True, reason=PIPELINE_GROUP_OFFLOAD_XFAIL_REASON, strict=True) - def test_pipeline_level_group_offloading_inference(self, base_pipe_output, expected_max_difference=1e-4): - super().test_pipeline_level_group_offloading_inference( - base_pipe_output, expected_max_difference=expected_max_difference - ) diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py b/tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py index a4d5b8f5c84e..ed80029487f9 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py @@ -16,7 +16,6 @@ import random import numpy as np -import pytest import torch from PIL import Image from torch import nn @@ -46,25 +45,6 @@ enable_full_determinism() -# `PriorTransformer` keeps `positional_embedding`, `prd_embedding`, `clip_mean` and `clip_std` as parameters of the -# model itself rather than of a submodule, so group offloading never onloads them: the forward pass then mixes -# onloaded activations with still-offloaded weights. Reproduces at both block and leaf level. -PIPELINE_GROUP_OFFLOAD_XFAIL_REASON = ( - "`PriorTransformer` holds parameters directly on the model (`positional_embedding`, `prd_embedding`, " - "`clip_mean`, `clip_std`), which group offloading never onloads." -) - -# A second, independent gap: the component-scoped test only offloads the denoiser under the names -# `transformer`/`unet`/`controlnet`/`adapter`, and only puts `vae`/`vqvae`/`image_encoder` back on the accelerator. -# A prior pipeline's denoiser is called `prior`, so it matches neither list and is left on CPU while the text -# encoder is onloaded. Fixing this means widening the mixin's component lists, not changing the pipeline. -COMPONENT_GROUP_OFFLOAD_XFAIL_REASON = ( - "The pipeline calls `PriorTransformer.post_process_latents()` after the denoising loop, which reads the " - "`clip_mean` / `clip_std` parameters held directly on the model. Group offloading onloads those only for the " - "duration of `forward`, so by then they are back on the offload device." -) - - class KandinskyV22PriorEmb2EmbPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = KandinskyV22PriorEmb2EmbPipeline required_input_params_in_call_signature = frozenset(["prompt", "image"]) @@ -233,23 +213,3 @@ def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff= class TestKandinskyV22PriorEmb2EmbPipelineMemory(KandinskyV22PriorEmb2EmbPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Kandinsky 2.2 prior emb2emb pipeline.""" - - @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True) - @MemoryTesterMixin._USE_STREAM - def test_group_offloading_inference_block_level(self, base_pipe_output, use_stream, expected_max_difference=1e-4): - super().test_group_offloading_inference_block_level( - base_pipe_output, use_stream, expected_max_difference=expected_max_difference - ) - - @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True) - @MemoryTesterMixin._USE_STREAM - def test_group_offloading_inference_leaf_level(self, base_pipe_output, use_stream, expected_max_difference=1e-4): - super().test_group_offloading_inference_leaf_level( - base_pipe_output, use_stream, expected_max_difference=expected_max_difference - ) - - @pytest.mark.xfail(condition=True, reason=PIPELINE_GROUP_OFFLOAD_XFAIL_REASON, strict=True) - def test_pipeline_level_group_offloading_inference(self, base_pipe_output, expected_max_difference=1e-4): - super().test_pipeline_level_group_offloading_inference( - base_pipe_output, expected_max_difference=expected_max_difference - )