
    +h/                     F    S r SSKJr  SSKJr  \" 5       (       a  SSKrS rg)z7
Accelerate utilities: Utilities related to accelerate
    )version   )is_accelerate_availableNc                    ^  [        5       (       d  T $ [        R                  " [        R                  5      R
                  n[        R                  " U5      [        R                  " S5      :  a  T $ U 4S jnU$ )a  
Decorator that applies a registered CpuOffload hook to an arbitrary function rather than `forward`. This is useful
for cases where a PyTorch module provides functions other than `forward` that should trigger a move to the
appropriate acceleration device. This is the case for `encode` and `decode` in [`AutoencoderKL`].

This decorator looks inside the internal `_hf_hook` property to find a registered offload hook.

:param method: The method to decorate. This method should be a method of a PyTorch module.
z0.17.0c                    > [        U S5      (       a6  [        U R                  S5      (       a  U R                  R                  U 5        T" U /UQ70 UD6$ )N_hf_hookpre_forward)hasattrr   r	   )selfargskwargsmethods      Z/home/james-whalen/.local/lib/python3.13/site-packages/diffusers/utils/accelerate_utils.pywrapper#apply_forward_hook.<locals>.wrapper+   sG    4$$)N)NMM%%d+d,T,V,,    )r   r   parse
accelerate__version__base_version)r   accelerate_versionr   s   `  r   apply_forward_hookr      sZ     #$$ z'='=>KK}}'(7==+BB-
 Nr   )__doc__	packagingr   import_utilsr   r   r    r   r   <module>r      s%     1 r   