




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 			UpperCAmelCase_      (   nn.Module       ):




							'''simple docstring'''



							def __init__(					self       :	Tuple			)			->					int:
										super().__init__()
										SCREAMING_SNAKE_CASE           =					nn.Linear(3       ,     4			)
										SCREAMING_SNAKE_CASE           =					nn.BatchNormad(4			)
										SCREAMING_SNAKE_CASE           =					nn.Linear(4       ,     5			)







							def    _UpperCAmelCase     (					self       :	Any       ,     a       :	Optional[Any]			)			->					Dict:
										return self.lineara(self.batchnorm(self.lineara(a			)			)			)





class 			UpperCAmelCase_      (   A       ):




							'''simple docstring'''



							def    _UpperCAmelCase     (					self       :	str       ,     a       :	Union[str, Any]       ,     *a       :	List[str]       ,     **a       :	Optional[int]			)			->					int:
										return (args[0] + 1,) + args[1:], kwargs





class 			UpperCAmelCase_      (   A       ):




							'''simple docstring'''



							def    _UpperCAmelCase     (					self       :	str       ,     a       :	Optional[Any]       ,     a       :	Optional[int]			)			->					Any:
										return output + 1





class 			UpperCAmelCase_      (   unittest.TestCase       ):




							'''simple docstring'''



							def    _UpperCAmelCase     (					self       :	int			)			->					Optional[int]:
										SCREAMING_SNAKE_CASE           =					ModelForTest()
										SCREAMING_SNAKE_CASE           =					ModelHook()

										add_hook_to_module(a       ,     a			)
										self.assertEqual(test_model._hf_hook       ,     a			)
										self.assertTrue(hasattr(a       ,     """_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(a			)
										self.assertFalse(hasattr(a       ,     """_hf_hook"""			)			)
										self.assertFalse(hasattr(a       ,     """_old_forward"""			)			)





							def    _UpperCAmelCase     (					self       :	Dict			)			->					Any:
										SCREAMING_SNAKE_CASE           =					ModelForTest()
										SCREAMING_SNAKE_CASE           =					ModelHook()

										add_hook_to_module(a       ,     a			)
										add_hook_to_module(a       ,     a       ,     append=a			)

										self.assertEqual(isinstance(test_model._hf_hook       ,     a			)       ,     a			)
										self.assertEqual(len(test_model._hf_hook.hooks			)       ,     2			)
										self.assertTrue(hasattr(a       ,     """_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(a			)
										self.assertFalse(hasattr(a       ,     """_hf_hook"""			)			)
										self.assertFalse(hasattr(a       ,     """_old_forward"""			)			)





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

										SCREAMING_SNAKE_CASE           =					PreForwardHook()
										add_hook_to_module(a       ,     a			)
										SCREAMING_SNAKE_CASE           =					test_model(a			)
										self.assertTrue(torch.allclose(a       ,     a       ,     atol=1E-5			)			)

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

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

										SCREAMING_SNAKE_CASE           =					test_model(a			)
										assert torch.allclose(a       ,     a       ,     atol=1E-5			)





							def    _UpperCAmelCase     (					self       :	List[Any]			)			->					List[Any]:
										SCREAMING_SNAKE_CASE           =					ModelForTest()
										SCREAMING_SNAKE_CASE           =					torch.randn(2       ,     3			)
										SCREAMING_SNAKE_CASE           =					test_model(a			)

										SCREAMING_SNAKE_CASE           =					PostForwardHook()
										add_hook_to_module(a       ,     a			)
										SCREAMING_SNAKE_CASE           =					test_model(a			)
										self.assertTrue(torch.allclose(a       ,     output + 1       ,     atol=1E-5			)			)

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

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

										SCREAMING_SNAKE_CASE           =					test_model(a			)
										assert torch.allclose(a       ,     output + 2       ,     atol=1E-5			)





							def    _UpperCAmelCase     (					self       :	Any			)			->					List[Any]:
										SCREAMING_SNAKE_CASE           =					ModelForTest()
										SCREAMING_SNAKE_CASE           =					torch.randn(2       ,     3			)
										SCREAMING_SNAKE_CASE           =					test_model(a			)

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

										SCREAMING_SNAKE_CASE           =					True
										SCREAMING_SNAKE_CASE           =					test_model(a			)
										self.assertFalse(outputa.requires_grad			)





							@require_multi_gpu
							def    _UpperCAmelCase     (					self       :	Tuple			)			->					Union[str, Any]:
										SCREAMING_SNAKE_CASE           =					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
										SCREAMING_SNAKE_CASE           =					torch.randn(2       ,     3			)
										SCREAMING_SNAKE_CASE           =					model(a			)
										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(a       ,     AlignDevicesHook(io_same_device=a			)			)
										SCREAMING_SNAKE_CASE           =					torch.randn(2       ,     3			).to(0			)
										SCREAMING_SNAKE_CASE           =					model(a			)
										self.assertEqual(output.device       ,     torch.device(0			)			)





							def    _UpperCAmelCase     (					self       :	Union[str, Any]			)			->					Tuple:
										SCREAMING_SNAKE_CASE           =					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
										SCREAMING_SNAKE_CASE           =					{"""execution_device""": 0 if torch.cuda.is_available() else """cpu""", """offload""": True}

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

										# 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
										SCREAMING_SNAKE_CASE           =					torch.device(hook_kwargs["""execution_device"""]			)
										self.assertEqual(model.batchnorm.running_mean.device       ,     a			)

										SCREAMING_SNAKE_CASE           =					torch.randn(2       ,     3			)
										SCREAMING_SNAKE_CASE           =					model(a			)
										self.assertEqual(output.device       ,     a			)

										# 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
										SCREAMING_SNAKE_CASE           =					{
										    """execution_device""": 0 if torch.cuda.is_available() else """cpu""",
										    """offload""": True,
										    """offload_buffers""": True,
										}

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

										# 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"""			)			)

										SCREAMING_SNAKE_CASE           =					torch.randn(2       ,     3			)
										SCREAMING_SNAKE_CASE           =					model(a			)
										self.assertEqual(output.device       ,     a			)

										# 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    _UpperCAmelCase     (					self       :	str			)			->					Union[str, Any]:
										SCREAMING_SNAKE_CASE           =					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
										SCREAMING_SNAKE_CASE           =					0 if torch.cuda.is_available() else """cpu"""
										attach_align_device_hook(a       ,     execution_device=a       ,     offload=a			)

										# 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
										SCREAMING_SNAKE_CASE           =					torch.device(a			)
										self.assertEqual(model.batchnorm.running_mean.device       ,     a			)

										SCREAMING_SNAKE_CASE           =					torch.randn(2       ,     3			)
										SCREAMING_SNAKE_CASE           =					model(a			)
										self.assertEqual(output.device       ,     a			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(a			)
										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(a       ,     execution_device=a       ,     offload=a       ,     offload_buffers=a			)

										# 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"""			)			)

										SCREAMING_SNAKE_CASE           =					torch.randn(2       ,     3			)
										SCREAMING_SNAKE_CASE           =					model(a			)
										self.assertEqual(output.device       ,     a			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(a			)
										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    _UpperCAmelCase     (					self       :	int			)			->					List[Any]:
										SCREAMING_SNAKE_CASE           =					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
										SCREAMING_SNAKE_CASE           =					0 if torch.cuda.is_available() else """cpu"""
										attach_align_device_hook(
										    a       ,     execution_device=a       ,     offload=a       ,     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
										SCREAMING_SNAKE_CASE           =					torch.device(a			)
										self.assertEqual(model.batchnorm.running_mean.device       ,     a			)

										SCREAMING_SNAKE_CASE           =					torch.randn(2       ,     3			)
										SCREAMING_SNAKE_CASE           =					model(a			)
										self.assertEqual(output.device       ,     a			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(a			)
										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(
										    a       ,     execution_device=a       ,     offload=a       ,     weights_map=model.state_dict()       ,     offload_buffers=a       ,     )

										# 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"""			)			)

										SCREAMING_SNAKE_CASE           =					torch.randn(2       ,     3			)
										SCREAMING_SNAKE_CASE           =					model(a			)
										self.assertEqual(output.device       ,     a			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(a			)
										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"""			)			)












import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding




class 			UpperCAmelCase_      (   A       ):




							'''simple docstring'''



							a__						  =     ['''image_processor''', '''tokenizer''']
							a__						  =     '''ViTImageProcessor'''
							a__						  =     ('''CLIPTokenizer''', '''CLIPTokenizerFast''')





							def __init__(					self       :	List[str]       ,     a       :	Dict=None       ,     a       :	List[str]=None       ,     **a       :	List[str]			)			->					Optional[Any]:
										SCREAMING_SNAKE_CASE           =					None
										if "feature_extractor" in kwargs:
													warnings.warn(
													    """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
													    """ instead."""       ,     a       ,     )
													SCREAMING_SNAKE_CASE           =					kwargs.pop("""feature_extractor"""			)

										SCREAMING_SNAKE_CASE           =					image_processor if image_processor is not None else feature_extractor
										if image_processor is None:
													raise ValueError("""You need to specify an `image_processor`."""			)
										if tokenizer is None:
													raise ValueError("""You need to specify a `tokenizer`."""			)

										super().__init__(a       ,     a			)





							def __call__(					self       :	List[Any]       ,     a       :	int=None       ,     a       :	Dict=None       ,     a       :	int=None       ,     a       :	Tuple=None       ,     **a       :	Any			)			->					Any:
										if text is None and visual_prompt is None and images is None:
													raise ValueError("""You have to specify either text, visual prompt or images."""			)

										if text is not None and visual_prompt is not None:
													raise ValueError("""You have to specify exactly one type of prompt. Either text or visual prompt."""			)

										if text is not None:
													SCREAMING_SNAKE_CASE           =					self.tokenizer(a       ,     return_tensors=a       ,     **a			)

										if visual_prompt is not None:
													SCREAMING_SNAKE_CASE           =					self.image_processor(a       ,     return_tensors=a       ,     **a			)

										if images is not None:
													SCREAMING_SNAKE_CASE           =					self.image_processor(a       ,     return_tensors=a       ,     **a			)

										if visual_prompt is not None and images is not None:
													SCREAMING_SNAKE_CASE           =					{
													    """pixel_values""": image_features.pixel_values,
													    """conditional_pixel_values""": prompt_features.pixel_values,
													}
													return encoding
										elif text is not None and images is not None:
													SCREAMING_SNAKE_CASE           =					image_features.pixel_values
													return encoding
										elif text is not None:
													return encoding
										elif visual_prompt is not None:
													SCREAMING_SNAKE_CASE           =					{
													    """conditional_pixel_values""": prompt_features.pixel_values,
													}
													return encoding
										else:
													return BatchEncoding(data=dict(**a			)       ,     tensor_type=a			)





							def    _UpperCAmelCase     (					self       :	Dict       ,     *a       :	Dict       ,     **a       :	Union[str, Any]			)			->					Any:
										return self.tokenizer.batch_decode(*a       ,     **a			)





							def    _UpperCAmelCase     (					self       :	Tuple       ,     *a       :	List[Any]       ,     **a       :	Union[str, Any]			)			->					Dict:
										return self.tokenizer.decode(*a       ,     **a			)





							@property
							def    _UpperCAmelCase     (					self       :	Any			)			->					Optional[Any]:
										warnings.warn(
										    """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""       ,     a       ,     )
										return self.image_processor_class







							@property
							def    _UpperCAmelCase     (					self       :	Optional[int]			)			->					str:
										warnings.warn(
										    """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."""       ,     a       ,     )
										return self.image_processor






