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
3 changes: 3 additions & 0 deletions src/diffusers/utils/accelerate_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
Accelerate utilities: Utilities related to accelerate
"""

import functools

from packaging import version

from .import_utils import is_accelerate_available
Expand All @@ -40,6 +42,7 @@ def apply_forward_hook(method):
if version.parse(accelerate_version) < version.parse("0.17.0"):
return method

@functools.wraps(method)
def wrapper(self, *args, **kwargs):
if hasattr(self, "_hf_hook") and hasattr(self._hf_hook, "pre_forward"):
self._hf_hook.pre_forward(self)
Expand Down
12 changes: 12 additions & 0 deletions tests/others/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -269,6 +269,18 @@ def _capture(target_device):
assert "moved to" in cuda_out, f"Non-MPS target should still emit the CPU-fallback info log, got: {cuda_out}"


class TestApplyForwardHook:
def test_preserves_wrapped_function_metadata(self):
from diffusers.utils.accelerate_utils import apply_forward_hook

@apply_forward_hook
def example(self):
"""Example method docstring."""

assert example.__name__ == "example"
assert example.__doc__ == "Example method docstring."


# Copied from https://github.com/huggingface/transformers/blob/main/tests/utils/test_expectations.py
class TestExpectations:
def test_expectations(self):
Expand Down
Loading