Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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):
Expand Down
74 changes: 72 additions & 2 deletions tests/models/autoencoders/test_models_autoencoder_mmaudio.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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):
Expand Down Expand Up @@ -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():
Expand All @@ -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):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -23,13 +24,19 @@
BaseModelTesterConfig,
MemoryTesterMixin,
ModelTesterMixin,
TorchCompileTesterMixin,
)


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):
Expand Down Expand Up @@ -79,8 +86,21 @@ class TestMMAudioVocoderModel(MMAudioVocoderTesterConfig, ModelTesterMixin):


class TestMMAudioVocoderMemory(MMAudioVocoderTesterConfig, MemoryTesterMixin):
pass


class TestMMAudioVocoderTorchCompile(MMAudioVocoderTesterConfig, TorchCompileTesterMixin):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

We usually don't do compile tests for autoencoders.

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)
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -23,13 +24,30 @@
BaseModelTesterConfig,
MemoryTesterMixin,
ModelTesterMixin,
TorchCompileTesterMixin,
)


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):
Expand Down Expand Up @@ -87,10 +105,14 @@ 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)

class TestKandinsky6SRLatentUpscalerBankTorchCompile(
Kandinsky6SRLatentUpscalerBankTesterConfig, TorchCompileTesterMixin
):
pass
@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)
Loading