diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 61dfc3133fbd..3ebb022df451 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -648,9 +648,6 @@ def set_attention_backend(self, backend: str) -> None: continue processor._attention_backend = backend - # Important to set the active backend so that it propagates gracefully throughout. - _AttentionBackendRegistry.set_active_backend(backend) - def reset_attention_backend(self) -> None: """ Resets the attention backend for the model. Following calls to `forward` will use the environment default, if diff --git a/tests/models/testing_utils/attention.py b/tests/models/testing_utils/attention.py index f31323d1bf52..5d6fc7d573d6 100644 --- a/tests/models/testing_utils/attention.py +++ b/tests/models/testing_utils/attention.py @@ -307,6 +307,23 @@ def teardown_method(self): gc.collect() backend_empty_cache(torch_device) + def test_set_attention_backend_does_not_change_registry(self): + model = self.model_class(**self.get_init_dict()) + initial_backend, _ = _AttentionBackendRegistry.get_active_backend() + model_backend = ( + AttentionBackendName.NATIVE + if initial_backend != AttentionBackendName.NATIVE + else AttentionBackendName._NATIVE_MATH + ) + + try: + model.set_attention_backend(model_backend.value) + active_backend, _ = _AttentionBackendRegistry.get_active_backend() + assert active_backend == initial_backend + finally: + model.reset_attention_backend() + _AttentionBackendRegistry.set_active_backend(initial_backend) + @torch.no_grad() @pytest.mark.parametrize("backend", _ALL_BACKEND_PARAMS) def test_set_attention_backend_matches_context_manager(self, backend):