From f3b4c135ec66cd15541ec4dbb3daf825828acfbb Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 7 Oct 2026 04:04:44 +0000 Subject: [PATCH 1/2] xfail some kandinsky6 tests. --- .../test_models_autoencoder_kandinsky6_sr.py | 31 +++++++- .../test_models_autoencoder_mmaudio.py | 74 ++++++++++++++++++- ...r.py => test_models_kandinsky6_vocoder.py} | 27 ++++++- ...test_models_kandinsky6_latent_upscaler.py} | 31 +++++++- 4 files changed, 158 insertions(+), 5 deletions(-) rename tests/models/autoencoders/{test_models_vocoder.py => test_models_kandinsky6_vocoder.py} (67%) rename tests/models/latent_upscaler/{test_models_latent_upscaler.py => test_models_kandinsky6_latent_upscaler.py} (71%) diff --git a/tests/models/autoencoders/test_models_autoencoder_kandinsky6_sr.py b/tests/models/autoencoders/test_models_autoencoder_kandinsky6_sr.py index f7a27ce31e78..055f91ef0f87 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kandinsky6_sr.py +++ b/tests/models/autoencoders/test_models_autoencoder_kandinsky6_sr.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch from diffusers import Kandinsky6SRVAE @@ -30,6 +31,24 @@ enable_full_determinism() +UNSPLITTABLE_ENCODER_DECODER = pytest.mark.xfail( + reason=( + "`_no_split_modules` keeps the whole encoder and decoder together, " + "preventing the test's required GPU/CPU split." + ), + raises=AssertionError, + strict=True, +) +DISK_OFFLOAD_OUTPUT_DEVICE = pytest.mark.xfail( + reason=( + "`_no_split_modules` keeps each encoder/decoder intact, forcing all-disk dispatch under the test's budgets; " + "the encode/decode hooks leave the output on CPU, but the reference is on GPU." + ), + raises=RuntimeError, + strict=True, +) + + class Kandinsky6SRVAETesterConfig(BaseModelTesterConfig): @property def model_class(self): @@ -98,7 +117,17 @@ def test_segmented_processing_matches_single_pass(self): class TestKandinsky6SRVAEMemory(Kandinsky6SRVAETesterConfig, MemoryTesterMixin): - pass + @UNSPLITTABLE_ENCODER_DECODER + def test_cpu_offload(self, base_model_output, tmp_path): + super().test_cpu_offload(base_model_output, tmp_path) + + @DISK_OFFLOAD_OUTPUT_DEVICE + def test_disk_offload_without_safetensors(self, base_model_output, tmp_path): + super().test_disk_offload_without_safetensors(base_model_output, tmp_path) + + @DISK_OFFLOAD_OUTPUT_DEVICE + def test_disk_offload_with_safetensors(self, base_model_output, tmp_path): + super().test_disk_offload_with_safetensors(base_model_output, tmp_path) class TestKandinsky6SRVAETorchCompile(Kandinsky6SRVAETesterConfig, TorchCompileTesterMixin): diff --git a/tests/models/autoencoders/test_models_autoencoder_mmaudio.py b/tests/models/autoencoders/test_models_autoencoder_mmaudio.py index a12e01ee4049..45f47322a312 100644 --- a/tests/models/autoencoders/test_models_autoencoder_mmaudio.py +++ b/tests/models/autoencoders/test_models_autoencoder_mmaudio.py @@ -13,12 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch from diffusers import MMAudioVAE from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, torch_device +from ...testing_utils import enable_full_determinism, require_accelerator, torch_device from ..testing_utils import ( BaseModelTesterConfig, MemoryTesterMixin, @@ -30,6 +31,28 @@ enable_full_determinism() +MEL_DTYPE = pytest.mark.xfail( + reason="MMAudio's STFT and mel-filter multiplication require float32.", + raises=RuntimeError, + strict=True, +) +LAYERWISE_CASTING_BACKWARD = pytest.mark.xfail( + reason="MMAudio's gain multiplication saves weights that are cast to float8 before backward.", + raises=RuntimeError, + strict=True, +) +NORMALIZATION_BUFFER_OFFLOAD = pytest.mark.xfail( + reason="MMAudio encode/decode bypass the offloading hooks for normalization buffers.", + raises=RuntimeError, + strict=True, +) +INCOMPLETE_DEVICE_MAP = pytest.mark.xfail( + reason="Automatic device maps omit MMAudio's normalization buffers.", + raises=ValueError, + strict=True, +) + + class MMAudioVAETesterConfig(BaseModelTesterConfig): @property def model_class(self): @@ -79,6 +102,16 @@ def output_shape(self) -> tuple[int, ...]: class TestMMAudioVAEModel(MMAudioVAETesterConfig, ModelTesterMixin): + @MEL_DTYPE + @require_accelerator + @pytest.mark.skipif( + torch_device not in ["cuda", "xpu"], + reason="float16 and bfloat16 can only be use for inference with an accelerator", + ) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) + def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): + super().test_from_save_pretrained_dtype_inference(tmp_path, dtype) + def test_latent_shape(self): model = self.model_class(**self.get_init_dict()).to(torch_device).eval() with torch.no_grad(): @@ -89,7 +122,44 @@ def test_latent_shape(self): class TestMMAudioVAEMemory(MMAudioVAETesterConfig, MemoryTesterMixin): - pass + @MEL_DTYPE + def test_layerwise_casting_memory(self): + super().test_layerwise_casting_memory() + + @LAYERWISE_CASTING_BACKWARD + def test_layerwise_casting_training(self): + super().test_layerwise_casting_training() + + @NORMALIZATION_BUFFER_OFFLOAD + @pytest.mark.parametrize("record_stream", [False, True]) + def test_group_offloading(self, base_model_output, record_stream): + super().test_group_offloading(base_model_output, record_stream) + + @pytest.mark.parametrize("record_stream", [False, True]) + @pytest.mark.parametrize( + "offload_type", ["block_level", pytest.param("leaf_level", marks=NORMALIZATION_BUFFER_OFFLOAD)] + ) + def test_group_offloading_with_layerwise_casting(self, record_stream, offload_type): + super().test_group_offloading_with_layerwise_casting(record_stream, offload_type) + + @pytest.mark.parametrize("record_stream", [False, True]) + @pytest.mark.parametrize( + "offload_type", ["block_level", pytest.param("leaf_level", marks=NORMALIZATION_BUFFER_OFFLOAD)] + ) + def test_group_offloading_with_disk(self, tmp_path, record_stream, offload_type): + super().test_group_offloading_with_disk(tmp_path, record_stream, offload_type) + + @INCOMPLETE_DEVICE_MAP + def test_cpu_offload(self, base_model_output, tmp_path): + super().test_cpu_offload(base_model_output, tmp_path) + + @INCOMPLETE_DEVICE_MAP + def test_disk_offload_without_safetensors(self, base_model_output, tmp_path): + super().test_disk_offload_without_safetensors(base_model_output, tmp_path) + + @INCOMPLETE_DEVICE_MAP + def test_disk_offload_with_safetensors(self, base_model_output, tmp_path): + super().test_disk_offload_with_safetensors(base_model_output, tmp_path) class TestMMAudioVAETorchCompile(MMAudioVAETesterConfig, TorchCompileTesterMixin): diff --git a/tests/models/autoencoders/test_models_vocoder.py b/tests/models/autoencoders/test_models_kandinsky6_vocoder.py similarity index 67% rename from tests/models/autoencoders/test_models_vocoder.py rename to tests/models/autoencoders/test_models_kandinsky6_vocoder.py index 225cfdbf4386..cffb7b36b031 100644 --- a/tests/models/autoencoders/test_models_vocoder.py +++ b/tests/models/autoencoders/test_models_kandinsky6_vocoder.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch from diffusers import MMAudioVocoder @@ -30,6 +31,13 @@ enable_full_determinism() +NESTED_UPSAMPLER_OFFLOAD = pytest.mark.xfail( + reason="Block offloading hooks on MMAudio's nested upsampler ModuleLists are never called.", + raises=RuntimeError, + strict=True, +) + + class MMAudioVocoderTesterConfig(BaseModelTesterConfig): @property def model_class(self): @@ -79,7 +87,24 @@ class TestMMAudioVocoderModel(MMAudioVocoderTesterConfig, ModelTesterMixin): class TestMMAudioVocoderMemory(MMAudioVocoderTesterConfig, MemoryTesterMixin): - pass + @NESTED_UPSAMPLER_OFFLOAD + @pytest.mark.parametrize("record_stream", [False, True]) + def test_group_offloading(self, base_model_output, record_stream): + super().test_group_offloading(base_model_output, record_stream) + + @pytest.mark.parametrize("record_stream", [False, True]) + @pytest.mark.parametrize( + "offload_type", [pytest.param("block_level", marks=NESTED_UPSAMPLER_OFFLOAD), "leaf_level"] + ) + def test_group_offloading_with_layerwise_casting(self, record_stream, offload_type): + super().test_group_offloading_with_layerwise_casting(record_stream, offload_type) + + @pytest.mark.parametrize("record_stream", [False, True]) + @pytest.mark.parametrize( + "offload_type", [pytest.param("block_level", marks=NESTED_UPSAMPLER_OFFLOAD), "leaf_level"] + ) + def test_group_offloading_with_disk(self, tmp_path, record_stream, offload_type): + super().test_group_offloading_with_disk(tmp_path, record_stream, offload_type) class TestMMAudioVocoderTorchCompile(MMAudioVocoderTesterConfig, TorchCompileTesterMixin): diff --git a/tests/models/latent_upscaler/test_models_latent_upscaler.py b/tests/models/latent_upscaler/test_models_kandinsky6_latent_upscaler.py similarity index 71% rename from tests/models/latent_upscaler/test_models_latent_upscaler.py rename to tests/models/latent_upscaler/test_models_kandinsky6_latent_upscaler.py index 059041af075d..2865f9d374b7 100644 --- a/tests/models/latent_upscaler/test_models_latent_upscaler.py +++ b/tests/models/latent_upscaler/test_models_kandinsky6_latent_upscaler.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch from diffusers import Kandinsky6SRLatentUpscalerBank @@ -30,6 +31,24 @@ enable_full_determinism() +UNSPLITTABLE_UPSCALER = pytest.mark.xfail( + reason=( + "`_no_split_modules` keeps each Kandinsky6SRLatentUpscaler intact, " + "preventing the test's required GPU/CPU split." + ), + raises=AssertionError, + strict=True, +) +DISK_OFFLOAD_NUMERICS = pytest.mark.xfail( + reason=( + "`_no_split_modules` keeps each upscaler intact, forcing all-disk dispatch and CPU execution under the test's " + "budgets; numerical differences from the GPU reference can exceed the output tolerance." + ), + raises=AssertionError, + strict=False, +) + + class Kandinsky6SRLatentUpscalerBankTesterConfig(BaseModelTesterConfig): @property def model_class(self): @@ -87,7 +106,17 @@ def test_x4_scale(self): class TestKandinsky6SRLatentUpscalerBankMemory(Kandinsky6SRLatentUpscalerBankTesterConfig, MemoryTesterMixin): - pass + @UNSPLITTABLE_UPSCALER + def test_cpu_offload(self, base_model_output, tmp_path): + super().test_cpu_offload(base_model_output, tmp_path) + + @DISK_OFFLOAD_NUMERICS + def test_disk_offload_without_safetensors(self, base_model_output, tmp_path): + super().test_disk_offload_without_safetensors(base_model_output, tmp_path) + + @DISK_OFFLOAD_NUMERICS + def test_disk_offload_with_safetensors(self, base_model_output, tmp_path): + super().test_disk_offload_with_safetensors(base_model_output, tmp_path) class TestKandinsky6SRLatentUpscalerBankTorchCompile( From 076a36499b53ae81f902b810fb0bf2347b98777d Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 7 Oct 2026 05:21:41 +0000 Subject: [PATCH 2/2] remove compile tests --- .../models/autoencoders/test_models_kandinsky6_vocoder.py | 5 ----- .../test_models_kandinsky6_latent_upscaler.py | 7 ------- 2 files changed, 12 deletions(-) diff --git a/tests/models/autoencoders/test_models_kandinsky6_vocoder.py b/tests/models/autoencoders/test_models_kandinsky6_vocoder.py index cffb7b36b031..a29f01b36df1 100644 --- a/tests/models/autoencoders/test_models_kandinsky6_vocoder.py +++ b/tests/models/autoencoders/test_models_kandinsky6_vocoder.py @@ -24,7 +24,6 @@ BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, - TorchCompileTesterMixin, ) @@ -105,7 +104,3 @@ def test_group_offloading_with_layerwise_casting(self, record_stream, offload_ty ) def test_group_offloading_with_disk(self, tmp_path, record_stream, offload_type): super().test_group_offloading_with_disk(tmp_path, record_stream, offload_type) - - -class TestMMAudioVocoderTorchCompile(MMAudioVocoderTesterConfig, TorchCompileTesterMixin): - pass diff --git a/tests/models/latent_upscaler/test_models_kandinsky6_latent_upscaler.py b/tests/models/latent_upscaler/test_models_kandinsky6_latent_upscaler.py index 2865f9d374b7..c96594dbffa9 100644 --- a/tests/models/latent_upscaler/test_models_kandinsky6_latent_upscaler.py +++ b/tests/models/latent_upscaler/test_models_kandinsky6_latent_upscaler.py @@ -24,7 +24,6 @@ BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, - TorchCompileTesterMixin, ) @@ -117,9 +116,3 @@ def test_disk_offload_without_safetensors(self, base_model_output, tmp_path): @DISK_OFFLOAD_NUMERICS def test_disk_offload_with_safetensors(self, base_model_output, tmp_path): super().test_disk_offload_with_safetensors(base_model_output, tmp_path) - - -class TestKandinsky6SRLatentUpscalerBankTorchCompile( - Kandinsky6SRLatentUpscalerBankTesterConfig, TorchCompileTesterMixin -): - pass