Skip to content

Commit 1bcb4fe

Browse files
sayakpaulsayakpaul
committed
xfail some kandinsky6 tests.
1 parent 0a9b872 commit 1bcb4fe

4 files changed

Lines changed: 158 additions & 5 deletions

File tree

‎tests/models/autoencoders/test_models_autoencoder_kandinsky6_sr.py‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16+
import pytest
1617
import torch
1718

1819
from diffusers import Kandinsky6SRVAE
@@ -30,6 +31,24 @@
3031
enable_full_determinism()
3132

3233

34+
UNSPLITTABLE_ENCODER_DECODER = pytest.mark.xfail(
35+
reason=(
36+
"`_no_split_modules` keeps the whole encoder and decoder together, "
37+
"preventing the test's required GPU/CPU split."
38+
),
39+
raises=AssertionError,
40+
strict=True,
41+
)
42+
DISK_OFFLOAD_OUTPUT_DEVICE = pytest.mark.xfail(
43+
reason=(
44+
"`_no_split_modules` keeps each encoder/decoder intact, forcing all-disk dispatch under the test's budgets; "
45+
"the encode/decode hooks leave the output on CPU, but the reference is on GPU."
46+
),
47+
raises=RuntimeError,
48+
strict=True,
49+
)
50+
51+
3352
class Kandinsky6SRVAETesterConfig(BaseModelTesterConfig):
3453
@property
3554
def model_class(self):
@@ -98,7 +117,17 @@ def test_segmented_processing_matches_single_pass(self):
98117

99118

100119
class TestKandinsky6SRVAEMemory(Kandinsky6SRVAETesterConfig, MemoryTesterMixin):
101-
pass
120+
@UNSPLITTABLE_ENCODER_DECODER
121+
def test_cpu_offload(self, base_model_output, tmp_path):
122+
super().test_cpu_offload(base_model_output, tmp_path)
123+
124+
@DISK_OFFLOAD_OUTPUT_DEVICE
125+
def test_disk_offload_without_safetensors(self, base_model_output, tmp_path):
126+
super().test_disk_offload_without_safetensors(base_model_output, tmp_path)
127+
128+
@DISK_OFFLOAD_OUTPUT_DEVICE
129+
def test_disk_offload_with_safetensors(self, base_model_output, tmp_path):
130+
super().test_disk_offload_with_safetensors(base_model_output, tmp_path)
102131

103132

104133
class TestKandinsky6SRVAETorchCompile(Kandinsky6SRVAETesterConfig, TorchCompileTesterMixin):

‎tests/models/autoencoders/test_models_autoencoder_mmaudio.py‎

Lines changed: 72 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,12 +13,13 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16+
import pytest
1617
import torch
1718

1819
from diffusers import MMAudioVAE
1920
from diffusers.utils.torch_utils import randn_tensor
2021

21-
from ...testing_utils import enable_full_determinism, torch_device
22+
from ...testing_utils import enable_full_determinism, require_accelerator, torch_device
2223
from ..testing_utils import (
2324
BaseModelTesterConfig,
2425
MemoryTesterMixin,
@@ -30,6 +31,28 @@
3031
enable_full_determinism()
3132

3233

34+
MEL_DTYPE = pytest.mark.xfail(
35+
reason="MMAudio's STFT and mel-filter multiplication require float32.",
36+
raises=RuntimeError,
37+
strict=True,
38+
)
39+
LAYERWISE_CASTING_BACKWARD = pytest.mark.xfail(
40+
reason="MMAudio's gain multiplication saves weights that are cast to float8 before backward.",
41+
raises=RuntimeError,
42+
strict=True,
43+
)
44+
NORMALIZATION_BUFFER_OFFLOAD = pytest.mark.xfail(
45+
reason="MMAudio encode/decode bypass the offloading hooks for normalization buffers.",
46+
raises=RuntimeError,
47+
strict=True,
48+
)
49+
INCOMPLETE_DEVICE_MAP = pytest.mark.xfail(
50+
reason="Automatic device maps omit MMAudio's normalization buffers.",
51+
raises=ValueError,
52+
strict=True,
53+
)
54+
55+
3356
class MMAudioVAETesterConfig(BaseModelTesterConfig):
3457
@property
3558
def model_class(self):
@@ -79,6 +102,16 @@ def output_shape(self) -> tuple[int, ...]:
79102

80103

81104
class TestMMAudioVAEModel(MMAudioVAETesterConfig, ModelTesterMixin):
105+
@MEL_DTYPE
106+
@require_accelerator
107+
@pytest.mark.skipif(
108+
torch_device not in ["cuda", "xpu"],
109+
reason="float16 and bfloat16 can only be use for inference with an accelerator",
110+
)
111+
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
112+
def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype):
113+
super().test_from_save_pretrained_dtype_inference(tmp_path, dtype)
114+
82115
def test_latent_shape(self):
83116
model = self.model_class(**self.get_init_dict()).to(torch_device).eval()
84117
with torch.no_grad():
@@ -89,7 +122,44 @@ def test_latent_shape(self):
89122

90123

91124
class TestMMAudioVAEMemory(MMAudioVAETesterConfig, MemoryTesterMixin):
92-
pass
125+
@MEL_DTYPE
126+
def test_layerwise_casting_memory(self):
127+
super().test_layerwise_casting_memory()
128+
129+
@LAYERWISE_CASTING_BACKWARD
130+
def test_layerwise_casting_training(self):
131+
super().test_layerwise_casting_training()
132+
133+
@NORMALIZATION_BUFFER_OFFLOAD
134+
@pytest.mark.parametrize("record_stream", [False, True])
135+
def test_group_offloading(self, base_model_output, record_stream):
136+
super().test_group_offloading(base_model_output, record_stream)
137+
138+
@pytest.mark.parametrize("record_stream", [False, True])
139+
@pytest.mark.parametrize(
140+
"offload_type", ["block_level", pytest.param("leaf_level", marks=NORMALIZATION_BUFFER_OFFLOAD)]
141+
)
142+
def test_group_offloading_with_layerwise_casting(self, record_stream, offload_type):
143+
super().test_group_offloading_with_layerwise_casting(record_stream, offload_type)
144+
145+
@pytest.mark.parametrize("record_stream", [False, True])
146+
@pytest.mark.parametrize(
147+
"offload_type", ["block_level", pytest.param("leaf_level", marks=NORMALIZATION_BUFFER_OFFLOAD)]
148+
)
149+
def test_group_offloading_with_disk(self, tmp_path, record_stream, offload_type):
150+
super().test_group_offloading_with_disk(tmp_path, record_stream, offload_type)
151+
152+
@INCOMPLETE_DEVICE_MAP
153+
def test_cpu_offload(self, base_model_output, tmp_path):
154+
super().test_cpu_offload(base_model_output, tmp_path)
155+
156+
@INCOMPLETE_DEVICE_MAP
157+
def test_disk_offload_without_safetensors(self, base_model_output, tmp_path):
158+
super().test_disk_offload_without_safetensors(base_model_output, tmp_path)
159+
160+
@INCOMPLETE_DEVICE_MAP
161+
def test_disk_offload_with_safetensors(self, base_model_output, tmp_path):
162+
super().test_disk_offload_with_safetensors(base_model_output, tmp_path)
93163

94164

95165
class TestMMAudioVAETorchCompile(MMAudioVAETesterConfig, TorchCompileTesterMixin):

tests/models/autoencoders/test_models_vocoder.py renamed to tests/models/autoencoders/test_models_kandinsky6_vocoder.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16+
import pytest
1617
import torch
1718

1819
from diffusers import MMAudioVocoder
@@ -30,6 +31,13 @@
3031
enable_full_determinism()
3132

3233

34+
NESTED_UPSAMPLER_OFFLOAD = pytest.mark.xfail(
35+
reason="Block offloading hooks on MMAudio's nested upsampler ModuleLists are never called.",
36+
raises=RuntimeError,
37+
strict=True,
38+
)
39+
40+
3341
class MMAudioVocoderTesterConfig(BaseModelTesterConfig):
3442
@property
3543
def model_class(self):
@@ -79,7 +87,24 @@ class TestMMAudioVocoderModel(MMAudioVocoderTesterConfig, ModelTesterMixin):
7987

8088

8189
class TestMMAudioVocoderMemory(MMAudioVocoderTesterConfig, MemoryTesterMixin):
82-
pass
90+
@NESTED_UPSAMPLER_OFFLOAD
91+
@pytest.mark.parametrize("record_stream", [False, True])
92+
def test_group_offloading(self, base_model_output, record_stream):
93+
super().test_group_offloading(base_model_output, record_stream)
94+
95+
@pytest.mark.parametrize("record_stream", [False, True])
96+
@pytest.mark.parametrize(
97+
"offload_type", [pytest.param("block_level", marks=NESTED_UPSAMPLER_OFFLOAD), "leaf_level"]
98+
)
99+
def test_group_offloading_with_layerwise_casting(self, record_stream, offload_type):
100+
super().test_group_offloading_with_layerwise_casting(record_stream, offload_type)
101+
102+
@pytest.mark.parametrize("record_stream", [False, True])
103+
@pytest.mark.parametrize(
104+
"offload_type", [pytest.param("block_level", marks=NESTED_UPSAMPLER_OFFLOAD), "leaf_level"]
105+
)
106+
def test_group_offloading_with_disk(self, tmp_path, record_stream, offload_type):
107+
super().test_group_offloading_with_disk(tmp_path, record_stream, offload_type)
83108

84109

85110
class TestMMAudioVocoderTorchCompile(MMAudioVocoderTesterConfig, TorchCompileTesterMixin):

tests/models/latent_upscaler/test_models_latent_upscaler.py renamed to tests/models/latent_upscaler/test_models_kandinsky6_latent_upscaler.py

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16+
import pytest
1617
import torch
1718

1819
from diffusers import Kandinsky6SRLatentUpscalerBank
@@ -30,6 +31,24 @@
3031
enable_full_determinism()
3132

3233

34+
UNSPLITTABLE_UPSCALER = pytest.mark.xfail(
35+
reason=(
36+
"`_no_split_modules` keeps each Kandinsky6SRLatentUpscaler intact, "
37+
"preventing the test's required GPU/CPU split."
38+
),
39+
raises=AssertionError,
40+
strict=True,
41+
)
42+
DISK_OFFLOAD_NUMERICS = pytest.mark.xfail(
43+
reason=(
44+
"`_no_split_modules` keeps each upscaler intact, forcing all-disk dispatch and CPU execution under the test's "
45+
"budgets; numerical differences from the GPU reference can exceed the output tolerance."
46+
),
47+
raises=AssertionError,
48+
strict=False,
49+
)
50+
51+
3352
class Kandinsky6SRLatentUpscalerBankTesterConfig(BaseModelTesterConfig):
3453
@property
3554
def model_class(self):
@@ -87,7 +106,17 @@ def test_x4_scale(self):
87106

88107

89108
class TestKandinsky6SRLatentUpscalerBankMemory(Kandinsky6SRLatentUpscalerBankTesterConfig, MemoryTesterMixin):
90-
pass
109+
@UNSPLITTABLE_UPSCALER
110+
def test_cpu_offload(self, base_model_output, tmp_path):
111+
super().test_cpu_offload(base_model_output, tmp_path)
112+
113+
@DISK_OFFLOAD_NUMERICS
114+
def test_disk_offload_without_safetensors(self, base_model_output, tmp_path):
115+
super().test_disk_offload_without_safetensors(base_model_output, tmp_path)
116+
117+
@DISK_OFFLOAD_NUMERICS
118+
def test_disk_offload_with_safetensors(self, base_model_output, tmp_path):
119+
super().test_disk_offload_with_safetensors(base_model_output, tmp_path)
91120

92121

93122
class TestKandinsky6SRLatentUpscalerBankTorchCompile(

0 commit comments

Comments
 (0)