




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


UpperCamelCase								=					{
    'configuration_vision_encoder_decoder': ['VisionEncoderDecoderConfig', 'VisionEncoderDecoderOnnxConfig']
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase								=					['VisionEncoderDecoderModel']

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase								=					['TFVisionEncoderDecoderModel']

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase								=					['FlaxVisionEncoderDecoderModel']

if TYPE_CHECKING:
     from .configuration_vision_encoder_decoder import VisionEncoderDecoderConfig, VisionEncoderDecoderOnnxConfig

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_vision_encoder_decoder import VisionEncoderDecoderModel

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_vision_encoder_decoder import TFVisionEncoderDecoderModel

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_vision_encoder_decoder import FlaxVisionEncoderDecoderModel

else:
     import sys

     UpperCamelCase								=					_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu





class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					Dict       )				->	Optional[int]:
         super().__init__()
         lowerCAmelCase__          =			nn.Linear(3       ,     4       )
         lowerCAmelCase__          =			nn.BatchNormad(4       )
         lowerCAmelCase__          =			nn.Linear(4       ,     5       )


    def        a				(	self					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					int       )				->	List[Any]:
         return self.lineara(self.batchnorm(self.lineara(SCREAMING_SNAKE_CASE__       )       )       )






class 			__lowerCamelCase	(   UpperCamelCase__   ):




    """simple docstring"""

    def        a				(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					str       ,     *SCREAMING_SNAKE_CASE__					:					Tuple       ,     **SCREAMING_SNAKE_CASE__					:					Any       )				->	Union[str, Any]:
         return (args[0] + 1,) + args[1:], kwargs






class 			__lowerCamelCase	(   UpperCamelCase__   ):




    """simple docstring"""

    def        a				(	self					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					str       )				->	Dict:
         return output + 1






class 			__lowerCamelCase	(   unittest.TestCase   ):




    """simple docstring"""

    def        a				(	self					:					List[str]       )				->	Tuple:
         lowerCAmelCase__          =			ModelForTest()
         lowerCAmelCase__          =			ModelHook()

         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         self.assertEqual(test_model._hf_hook       ,     SCREAMING_SNAKE_CASE__       )
         self.assertTrue(hasattr(SCREAMING_SNAKE_CASE__       ,     "_old_forward"       )       )

         # Check adding the hook did not change the name or the signature
         self.assertEqual(test_model.forward.__name__       ,     "forward"       )
         self.assertListEqual(list(inspect.signature(test_model.forward       ).parameters       )       ,     ["x"]       )

         remove_hook_from_module(SCREAMING_SNAKE_CASE__       )
         self.assertFalse(hasattr(SCREAMING_SNAKE_CASE__       ,     "_hf_hook"       )       )
         self.assertFalse(hasattr(SCREAMING_SNAKE_CASE__       ,     "_old_forward"       )       )


    def        a				(	self					:					Union[str, Any]       )				->	int:
         lowerCAmelCase__          =			ModelForTest()
         lowerCAmelCase__          =			ModelHook()

         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     append=SCREAMING_SNAKE_CASE__       )

         self.assertEqual(isinstance(test_model._hf_hook       ,     SCREAMING_SNAKE_CASE__       )       ,     SCREAMING_SNAKE_CASE__       )
         self.assertEqual(len(test_model._hf_hook.hooks       )       ,     2       )
         self.assertTrue(hasattr(SCREAMING_SNAKE_CASE__       ,     "_old_forward"       )       )

         # Check adding the hook did not change the name or the signature
         self.assertEqual(test_model.forward.__name__       ,     "forward"       )
         self.assertListEqual(list(inspect.signature(test_model.forward       ).parameters       )       ,     ["x"]       )

         remove_hook_from_module(SCREAMING_SNAKE_CASE__       )
         self.assertFalse(hasattr(SCREAMING_SNAKE_CASE__       ,     "_hf_hook"       )       )
         self.assertFalse(hasattr(SCREAMING_SNAKE_CASE__       ,     "_old_forward"       )       )


    def        a				(	self					:					List[str]       )				->	Any:
         lowerCAmelCase__          =			ModelForTest()
         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			test_model(x + 1       )
         lowerCAmelCase__          =			test_model(x + 2       )

         lowerCAmelCase__          =			PreForwardHook()
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     atol=1e-5       )       )

         # Attaching a hook to a model when it already has one replaces, does not chain
         lowerCAmelCase__          =			PreForwardHook()
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     atol=1e-5       )       )

         # You need to use the sequential hook to chain two or more hooks
         lowerCAmelCase__          =			SequentialHook(PreForwardHook()       ,     PreForwardHook()       )
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         assert torch.allclose(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     atol=1e-5       )


    def        a				(	self					:					Any       )				->	Union[str, Any]:
         lowerCAmelCase__          =			ModelForTest()
         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			PostForwardHook()
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE__       ,     output + 1       ,     atol=1e-5       )       )

         # Attaching a hook to a model when it already has one replaces, does not chain
         lowerCAmelCase__          =			PostForwardHook()
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE__       ,     output + 1       ,     atol=1e-5       )       )

         # You need to use the sequential hook to chain two or more hooks
         lowerCAmelCase__          =			SequentialHook(PostForwardHook()       ,     PostForwardHook()       )
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         assert torch.allclose(SCREAMING_SNAKE_CASE__       ,     output + 2       ,     atol=1e-5       )


    def        a				(	self					:					Optional[int]       )				->	int:
         lowerCAmelCase__          =			ModelForTest()
         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			PostForwardHook()
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE__       ,     output + 1       )       )
         self.assertTrue(outputa.requires_grad       )

         lowerCAmelCase__          =			True
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         self.assertFalse(outputa.requires_grad       )


    @require_multi_gpu
    def        a				(	self					:					Optional[Any]       )				->	List[str]:
         lowerCAmelCase__          =			ModelForTest()
         # Everything is on CPU
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # This will move each submodule on different devices
         add_hook_to_module(model.lineara       ,     AlignDevicesHook(execution_device=0       )       )
         add_hook_to_module(model.batchnorm       ,     AlignDevicesHook(execution_device=0       )       )
         add_hook_to_module(model.lineara       ,     AlignDevicesHook(execution_device=1       )       )

         self.assertEqual(model.lineara.weight.device       ,     torch.device(0       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device(0       )       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     torch.device(0       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device(1       )       )

         # We can still make a forward pass. The input does not need to be on any particular device
         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     torch.device(1       )       )

         # We can add a general hook to put back output on same device as input.
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     AlignDevicesHook(io_same_device=SCREAMING_SNAKE_CASE__       )       )
         lowerCAmelCase__          =			torch.randn(2       ,     3       ).to(0       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     torch.device(0       )       )


    def        a				(	self					:					List[str]       )				->	List[str]:
         lowerCAmelCase__          =			ModelForTest()

         # Everything is on CPU
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # This will move each submodule on different devices
         lowerCAmelCase__          =			{"execution_device": 0 if torch.cuda.is_available() else "cpu", "offload": True}

         add_hook_to_module(model.lineara       ,     AlignDevicesHook(**SCREAMING_SNAKE_CASE__       )       )
         add_hook_to_module(model.batchnorm       ,     AlignDevicesHook(**SCREAMING_SNAKE_CASE__       )       )
         add_hook_to_module(model.lineara       ,     AlignDevicesHook(**SCREAMING_SNAKE_CASE__       )       )

         # Parameters have been offloaded, so on the meta device
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         # Buffers are not included in the offload by default, so are on the execution device
         lowerCAmelCase__          =			torch.device(hook_kwargs["execution_device"]       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     SCREAMING_SNAKE_CASE__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_module(model.lineara       )
         remove_hook_from_module(model.batchnorm       )
         remove_hook_from_module(model.lineara       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # Now test with buffers included in the offload
         lowerCAmelCase__          =			{
             "execution_device": 0 if torch.cuda.is_available() else "cpu",
             "offload": True,
             "offload_buffers": True,
         }

         add_hook_to_module(model.lineara       ,     AlignDevicesHook(**SCREAMING_SNAKE_CASE__       )       )
         add_hook_to_module(model.batchnorm       ,     AlignDevicesHook(**SCREAMING_SNAKE_CASE__       )       )
         add_hook_to_module(model.lineara       ,     AlignDevicesHook(**SCREAMING_SNAKE_CASE__       )       )

         # Parameters have been offloaded, so on the meta device, buffers included
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     torch.device("meta"       )       )

         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     SCREAMING_SNAKE_CASE__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_module(model.lineara       )
         remove_hook_from_module(model.batchnorm       )
         remove_hook_from_module(model.lineara       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )


    def        a				(	self					:					Optional[int]       )				->	Union[str, Any]:
         lowerCAmelCase__          =			ModelForTest()

         # Everything is on CPU
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # This will move each submodule on different devices
         lowerCAmelCase__          =			0 if torch.cuda.is_available() else "cpu"
         attach_align_device_hook(SCREAMING_SNAKE_CASE__       ,     execution_device=SCREAMING_SNAKE_CASE__       ,     offload=SCREAMING_SNAKE_CASE__       )

         # Parameters have been offloaded, so on the meta device
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         # Buffers are not included in the offload by default, so are on the execution device
         lowerCAmelCase__          =			torch.device(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     SCREAMING_SNAKE_CASE__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_submodules(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # Now test with buffers included in the offload
         attach_align_device_hook(SCREAMING_SNAKE_CASE__       ,     execution_device=SCREAMING_SNAKE_CASE__       ,     offload=SCREAMING_SNAKE_CASE__       ,     offload_buffers=SCREAMING_SNAKE_CASE__       )

         # Parameters have been offloaded, so on the meta device, buffers included
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     torch.device("meta"       )       )

         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     SCREAMING_SNAKE_CASE__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_submodules(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )


    def        a				(	self					:					Optional[Any]       )				->	str:
         lowerCAmelCase__          =			ModelForTest()

         # Everything is on CPU
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # This will move each submodule on different devices
         lowerCAmelCase__          =			0 if torch.cuda.is_available() else "cpu"
         attach_align_device_hook(
             SCREAMING_SNAKE_CASE__       ,     execution_device=SCREAMING_SNAKE_CASE__       ,     offload=SCREAMING_SNAKE_CASE__       ,     weights_map=model.state_dict()       )

         # Parameters have been offloaded, so on the meta device
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         # Buffers are not included in the offload by default, so are on the execution device
         lowerCAmelCase__          =			torch.device(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     SCREAMING_SNAKE_CASE__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_submodules(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # Now test with buffers included in the offload
         attach_align_device_hook(
             SCREAMING_SNAKE_CASE__       ,     execution_device=SCREAMING_SNAKE_CASE__       ,     offload=SCREAMING_SNAKE_CASE__       ,     weights_map=model.state_dict()       ,     offload_buffers=SCREAMING_SNAKE_CASE__       ,     )

         # Parameters have been offloaded, so on the meta device, buffers included
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     torch.device("meta"       )       )

         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     SCREAMING_SNAKE_CASE__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_submodules(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
