Test Cleanup Precision issues (#4812) · lvitol/diffusers@189e9f0 · GitHub
Skip to content

Commit 189e9f0

Browse files
DN6patrickvonplaten
DN6
and
patrickvonplaten
authored
Test Cleanup Precision issues (huggingface#4812)
* proposal for flaky tests * more precision fixes * move more tests to use cosine distance * more test fixes * clean up * use default attn * clean up * update expected value * make style * make style * Apply suggestions from code review * Update src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py * make style * fix failing tests --------- Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
1 parent 38466c3 commit 189e9f0

20 files changed

Lines changed: 238 additions & 43 deletions

‎src/diffusers/models/attention_processor.py‎

Lines changed: 18 additions & 0 deletions

‎src/diffusers/models/autoencoder_kl.py‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,13 @@
2020
from ..configuration_utils import ConfigMixin, register_to_config
2121
from ..loaders import FromOriginalVAEMixin
2222
from ..utils import BaseOutput, apply_forward_hook
23-
from .attention_processor import AttentionProcessor, AttnProcessor
23+
from .attention_processor import (
24+
ADDED_KV_ATTENTION_PROCESSORS,
25+
CROSS_ATTENTION_PROCESSORS,
26+
AttentionProcessor,
27+
AttnAddedKVProcessor,
28+
AttnProcessor,
29+
)
2430
from .modeling_utils import ModelMixin
2531
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder
2632

@@ -228,7 +234,16 @@ def set_default_attn_processor(self):
228234
"""
229235
Disables custom attention processors and sets the default attention implementation.
230236
"""
231-
self.set_attn_processor(AttnProcessor())
237+
if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
238+
processor = AttnAddedKVProcessor()
239+
elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
240+
processor = AttnProcessor()
241+
else:
242+
raise ValueError(
243+
f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
244+
)
245+
246+
self.set_attn_processor(processor)
232247

233248
@apply_forward_hook
234249
def encode(self, x: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput:

‎src/diffusers/models/controlnet.py‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,13 @@
2121
from ..configuration_utils import ConfigMixin, register_to_config
2222
from ..loaders import FromOriginalControlnetMixin
2323
from ..utils import BaseOutput, logging
24-
from .attention_processor import AttentionProcessor, AttnProcessor
24+
from .attention_processor import (
25+
ADDED_KV_ATTENTION_PROCESSORS,
26+
CROSS_ATTENTION_PROCESSORS,
27+
AttentionProcessor,
28+
AttnAddedKVProcessor,
29+
AttnProcessor,
30+
)
2531
from .embeddings import TextImageProjection, TextImageTimeEmbedding, TextTimeEmbedding, TimestepEmbedding, Timesteps
2632
from .modeling_utils import ModelMixin
2733
from .unet_2d_blocks import (
@@ -550,7 +556,16 @@ def set_default_attn_processor(self):
550556
"""
551557
Disables custom attention processors and sets the default attention implementation.
552558
"""
553-
self.set_attn_processor(AttnProcessor())
559+
if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
560+
processor = AttnAddedKVProcessor()
561+
elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
562+
processor = AttnProcessor()
563+
else:
564+
raise ValueError(
565+
f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
566+
)
567+
568+
self.set_attn_processor(processor)
554569

555570
# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attention_slice
556571
def set_attention_slice(self, slice_size):

‎src/diffusers/models/prior_transformer.py‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,13 @@
88
from ..configuration_utils import ConfigMixin, register_to_config
99
from ..utils import BaseOutput
1010
from .attention import BasicTransformerBlock
11-
from .attention_processor import AttentionProcessor, AttnProcessor
11+
from .attention_processor import (
12+
ADDED_KV_ATTENTION_PROCESSORS,
13+
CROSS_ATTENTION_PROCESSORS,
14+
AttentionProcessor,
15+
AttnAddedKVProcessor,
16+
AttnProcessor,
17+
)
1218
from .embeddings import TimestepEmbedding, Timesteps
1319
from .modeling_utils import ModelMixin
1420

@@ -224,7 +230,16 @@ def set_default_attn_processor(self):
224230
"""
225231
Disables custom attention processors and sets the default attention implementation.
226232
"""
227-
self.set_attn_processor(AttnProcessor())
233+
if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
234+
processor = AttnAddedKVProcessor()
235+
elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
236+
processor = AttnProcessor()
237+
else:
238+
raise ValueError(
239+
f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
240+
)
241+
242+
self.set_attn_processor(processor)
228243

229244
def forward(
230245
self,

‎src/diffusers/models/unet_2d_condition.py‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,13 @@
2222
from ..loaders import UNet2DConditionLoadersMixin
2323
from ..utils import BaseOutput, logging
2424
from .activations import get_activation
25-
from .attention_processor import AttentionProcessor, AttnProcessor
25+
from .attention_processor import (
26+
ADDED_KV_ATTENTION_PROCESSORS,
27+
CROSS_ATTENTION_PROCESSORS,
28+
AttentionProcessor,
29+
AttnAddedKVProcessor,
30+
AttnProcessor,
31+
)
2632
from .embeddings import (
2733
GaussianFourierProjection,
2834
ImageHintTimeEmbedding,
@@ -639,7 +645,16 @@ def set_default_attn_processor(self):
639645
"""
640646
Disables custom attention processors and sets the default attention implementation.
641647
"""
642-
self.set_attn_processor(AttnProcessor())
648+
if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
649+
processor = AttnAddedKVProcessor()
650+
elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
651+
processor = AttnProcessor()
652+
else:
653+
raise ValueError(
654+
f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
655+
)
656+
657+
self.set_attn_processor(processor)
643658

644659
def set_attention_slice(self, slice_size):
645660
r"""

‎src/diffusers/models/unet_3d_condition.py‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,13 @@
2222
from ..configuration_utils import ConfigMixin, register_to_config
2323
from ..loaders import UNet2DConditionLoadersMixin
2424
from ..utils import BaseOutput, logging
25-
from .attention_processor import AttentionProcessor, AttnProcessor
25+
from .attention_processor import (
26+
ADDED_KV_ATTENTION_PROCESSORS,
27+
CROSS_ATTENTION_PROCESSORS,
28+
AttentionProcessor,
29+
AttnAddedKVProcessor,
30+
AttnProcessor,
31+
)
2632
from .embeddings import TimestepEmbedding, Timesteps
2733
from .modeling_utils import ModelMixin
2834
from .transformer_temporal import TransformerTemporalModel
@@ -439,7 +445,16 @@ def set_default_attn_processor(self):
439445
"""
440446
Disables custom attention processors and sets the default attention implementation.
441447
"""
442-
self.set_attn_processor(AttnProcessor())
448+
if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
449+
processor = AttnAddedKVProcessor()
450+
elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
451+
processor = AttnProcessor()
452+
else:
453+
raise ValueError(
454+
f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
455+
)
456+
457+
self.set_attn_processor(processor)
443458

444459
def _set_gradient_checkpointing(self, module, value=False):
445460
if isinstance(module, (CrossAttnDownBlock3D, DownBlock3D, CrossAttnUpBlock3D, UpBlock3D)):

‎src/diffusers/pipelines/audioldm2/modeling_audioldm2.py‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,13 @@
2222
from ...configuration_utils import ConfigMixin, register_to_config
2323
from ...loaders import UNet2DConditionLoadersMixin
2424
from ...models.activations import get_activation
25-
from ...models.attention_processor import AttentionProcessor, AttnProcessor
25+
from ...models.attention_processor import (
26+
ADDED_KV_ATTENTION_PROCESSORS,
27+
CROSS_ATTENTION_PROCESSORS,
28+
AttentionProcessor,
29+
AttnAddedKVProcessor,
30+
AttnProcessor,
31+
)
2632
from ...models.embeddings import (
2733
TimestepEmbedding,
2834
Timesteps,
@@ -571,7 +577,16 @@ def set_default_attn_processor(self):
571577
"""
572578
Disables custom attention processors and sets the default attention implementation.
573579
"""
574-
self.set_attn_processor(AttnProcessor())
580+
if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
581+
processor = AttnAddedKVProcessor()
582+
elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
583+
processor = AttnProcessor()
584+
else:
585+
raise ValueError(
586+
f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
587+
)
588+
589+
self.set_attn_processor(processor)
575590

576591
# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attention_slice
577592
def set_attention_slice(self, slice_size):

‎src/diffusers/pipelines/versatile_diffusion/modeling_text_unet.py‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010
from ...models.activations import get_activation
1111
from ...models.attention import Attention
1212
from ...models.attention_processor import (
13+
ADDED_KV_ATTENTION_PROCESSORS,
14+
CROSS_ATTENTION_PROCESSORS,
1315
AttentionProcessor,
1416
AttnAddedKVProcessor,
1517
AttnAddedKVProcessor2_0,
@@ -844,7 +846,17 @@ def set_default_attn_processor(self):
844846
"""
845847
Disables custom attention processors and sets the default attention implementation.
846848
"""
847-
self.set_attn_processor(AttnProcessor())
849+
if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
850+
processor = AttnAddedKVProcessor()
851+
elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
852+
processor = AttnProcessor()
853+
else:
854+
raise ValueError(
855+
"Cannot call `set_default_attn_processor` when attention processors are of type"
856+
f" {next(iter(self.attn_processors.values()))}"
857+
)
858+
859+
self.set_attn_processor(processor)
848860

849861
def set_attention_slice(self, slice_size):
850862
r"""

‎tests/models/test_models_unet_2d.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -263,7 +263,7 @@ def test_output_pretrained_ve_mid(self):
263263

264264
output_slice = output[0, -3:, -3:, -1].flatten().cpu()
265265
# fmt: off
266-
expected_output_slice = torch.tensor([-4842.8691, -6499.6631, -3800.1953, -7978.2686, -10980.7129, -20028.8535, 8148.2822, 2342.2905, 567.7608])
266+
expected_output_slice = torch.tensor([-4836.2178, -6487.1470, -3816.8196, -7964.9302, -10966.3037, -20043.5957, 8137.0513, 2340.3328, 544.6056])
267267
# fmt: on
268268

269269
self.assertTrue(torch_all_close(output_slice, expected_output_slice, rtol=1e-2))

‎tests/models/test_models_unet_2d_condition.py‎

Lines changed: 12 additions & 6 deletions

0 commit comments

Comments
 (0)