Skip to content
Merged
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
3 changes: 2 additions & 1 deletion src/diffusers/models/attention_dispatch.py
Original file line number Diff line number Diff line change
Expand Up @@ -418,7 +418,8 @@ def dispatch_attention_fn(
# variable), or we use a custom backend based on whether user is using the `attention_backend` context manager
backend_name, backend_fn = _AttentionBackendRegistry.get_active_backend()
else:
backend_name = AttentionBackendName(backend)
# Skip the `Enum` call for members: older torch versions cannot always trace it in a full graph.
backend_name = backend if isinstance(backend, AttentionBackendName) else AttentionBackendName(backend)
backend_fn = _AttentionBackendRegistry._backends.get(backend_name)

kwargs = {
Expand Down
32 changes: 27 additions & 5 deletions tests/models/autoencoders/test_models_kandinsky6_vocoder.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,11 @@
import torch

from diffusers import MMAudioVocoder
from diffusers.utils import is_torch_version
from diffusers.utils.torch_utils import randn_tensor

from ...testing_utils import enable_full_determinism, torch_device
from ..testing_utils import (
BaseModelTesterConfig,
MemoryTesterMixin,
ModelTesterMixin,
)
from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TorchCompileTesterMixin


enable_full_determinism()
Expand Down Expand Up @@ -104,3 +101,28 @@ 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):
@pytest.mark.xfail(
condition=is_torch_version("<", "2.12"),
reason=(
'With deterministic algorithms on, torch < 2.12 routes `F.pad(mode="replicate")` through `importlib`, '
"which Dynamo cannot trace in a full graph."
),
raises=torch._dynamo.exc.Unsupported,
strict=True,
)
def test_torch_compile_recompilation_and_graph_break(self):
super().test_torch_compile_recompilation_and_graph_break()

@pytest.mark.xfail(
reason=(
"Block-level offloading hooks each `ups[i]` `ModuleList`, but `forward` calls `ups[i][0]` directly, so the "
"hook never onloads the weights."
),
raises=RuntimeError,
strict=True,
)
def test_compile_with_group_offloading(self):
super().test_compile_with_group_offloading()
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,11 @@
import torch

from diffusers import Kandinsky6SRLatentUpscalerBank
from diffusers.utils import is_torch_version
from diffusers.utils.torch_utils import randn_tensor

from ...testing_utils import enable_full_determinism, torch_device
from ..testing_utils import (
BaseModelTesterConfig,
MemoryTesterMixin,
ModelTesterMixin,
)
from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TorchCompileTesterMixin


enable_full_determinism()
Expand Down Expand Up @@ -116,3 +113,19 @@ 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
):
@pytest.mark.xfail(
condition=is_torch_version("<", "2.12"),
reason=(
'With deterministic algorithms on, torch < 2.12 routes `F.pad(mode="replicate")` through `importlib`, '
"which Dynamo cannot trace in a full graph."
),
raises=torch._dynamo.exc.Unsupported,
strict=True,
)
def test_torch_compile_recompilation_and_graph_break(self):
super().test_torch_compile_recompilation_and_graph_break()
6 changes: 6 additions & 0 deletions tests/models/testing_utils/quantization.py
Original file line number Diff line number Diff line change
Expand Up @@ -2074,6 +2074,12 @@ def test_bfl_rejects_unquantized_checkpoint(self):
@is_torch_compile
@require_accelerator
@require_accelerate
# Both schemes store E4M3 tensors (fp8r weights, nvfp4 block scales) and compute in bf16; Triton cannot compile either
# below compute capability 8.9 (e.g. T4).
@pytest.mark.skipif(
torch_device == "cuda" and torch.cuda.get_device_capability() < (8, 9),
reason="torch.compile of BFL models needs FP8 support in Triton (compute capability 8.9+)",
)
class BFLCompileTesterMixin(BFLConfigMixin, QuantizationCompileTesterMixin):
"""
Mixin class for testing torch.compile with BFL (`fp8r` / `nvfp4`) quantized models.
Expand Down
16 changes: 14 additions & 2 deletions tests/models/transformers/test_models_transformer_kandinsky6_sr.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,10 +60,11 @@ def get_init_dict(self) -> dict:
"out_visual_dim": 8,
"time_dim": 16,
"patch_size": (1, 2, 2),
"model_dim": 24,
"model_dim": 32,

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.

To satisfy the requirements of Flex + compile.

"ff_dim": 32,
"num_visual_blocks": 2,
"axes_dims": (4, 4, 4),
# head_dim = sum(axes_dims); Inductor lowers flex attention only for head dims of at least 16
"axes_dims": (4, 6, 6),
}

def _build_dummy_inputs(self, batch_size: int, num_frames: int, height: int, width: int) -> dict:
Expand Down Expand Up @@ -106,6 +107,17 @@ def different_shapes_for_compilation(self):
def get_dummy_inputs(self, height: int = 16, width: int = 16) -> dict:
return self._build_dummy_inputs(batch_size=1, num_frames=2, height=height, width=width)

@pytest.mark.xfail(
reason=(
"The token grid is split into 8x8 blocks and `sliding_tile_mask` sizes its table by `max(...)` of the block "
"counts, so the compiled graph specializes on the input shape and recompiles."
),
raises=torch._dynamo.exc.RecompileError,
strict=True,
)
def test_compile_on_different_shapes(self):

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 can make it work with some modeling-level changes but not a priority for this model at the moment IMO.

super().test_compile_on_different_shapes()


@pytest.mark.skipif(torch_device == "cpu", reason="FlexAttention does not support backward on CPU.")
class TestKandinsky6SRTransformerTraining(Kandinsky6SRTransformerTesterConfig, TrainingTesterMixin):
Expand Down
Loading