



"""simple docstring"""

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):

 def __init__(  self	: str       ):
     super().__init__()
     snake_case_				:  Tuple       =			nn.Linear(3						,  4       )
     snake_case_				:  Optional[int]       =			nn.BatchNormad(4       )
     snake_case_				:  List[Any]       =			nn.Linear(4						,  5       )






 def 		_snake_case      (  self	: Any						,  lowercase_	: str       ):
     return self.lineara(self.batchnorm(self.lineara(lowercase_       )       )       )





class  _UpperCAmelCase (    lowerCAmelCase__):

 def 		_snake_case      (  self	: Tuple						,  lowercase_	: Any						,  *lowercase_	: Union[str, Any]						,  **lowercase_	: Tuple       ):
     return (args[0] + 1,) + args[1:], kwargs





class  _UpperCAmelCase (    lowerCAmelCase__):

 def 		_snake_case      (  self	: List[Any]						,  lowercase_	: Optional[int]						,  lowercase_	: Optional[int]       ):
     return output + 1





class  _UpperCAmelCase (    unittest.TestCase):

 def 		_snake_case      (  self	: int       ):
     snake_case_				:  int       =			ModelForTest()
     snake_case_				:  List[str]       =			ModelHook()

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

 def 		_snake_case      (  self	: str       ):
     snake_case_				:  str       =			ModelForTest()
     snake_case_				:  Union[str, Any]       =			ModelHook()

     add_hook_to_module(lowercase_						,  lowercase_       )
     add_hook_to_module(lowercase_						,  lowercase_						,  append=lowercase_       )

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

 def 		_snake_case      (  self	: Tuple       ):
     snake_case_				:  Dict       =			ModelForTest()
     snake_case_				:  List[Any]       =			torch.randn(2						,  3       )
     snake_case_				:  int       =			test_model(x + 1       )
     snake_case_				:  str       =			test_model(x + 2       )

     snake_case_				:  Any       =			PreForwardHook()
     add_hook_to_module(lowercase_						,  lowercase_       )
     snake_case_				:  Optional[Any]       =			test_model(lowercase_       )
     self.assertTrue(torch.allclose(lowercase_						,  lowercase_						,  atol=1E-5       )       )

     # Attaching a hook to a model when it already has one replaces, does not chain
     snake_case_				:  List[Any]       =			PreForwardHook()
     add_hook_to_module(lowercase_						,  lowercase_       )
     snake_case_				:  Tuple       =			test_model(lowercase_       )
     self.assertTrue(torch.allclose(lowercase_						,  lowercase_						,  atol=1E-5       )       )

     # You need to use the sequential hook to chain two or more hooks
     snake_case_				:  List[str]       =			SequentialHook(PreForwardHook()						,  PreForwardHook()       )
     add_hook_to_module(lowercase_						,  lowercase_       )

     snake_case_				:  List[str]       =			test_model(lowercase_       )
     assert torch.allclose(lowercase_						,  lowercase_						,  atol=1E-5       )

 def 		_snake_case      (  self	: List[Any]       ):
     snake_case_				:  Any       =			ModelForTest()
     snake_case_				:  Dict       =			torch.randn(2						,  3       )
     snake_case_				:  int       =			test_model(lowercase_       )

     snake_case_				:  Dict       =			PostForwardHook()
     add_hook_to_module(lowercase_						,  lowercase_       )
     snake_case_				:  Tuple       =			test_model(lowercase_       )
     self.assertTrue(torch.allclose(lowercase_						,  output + 1						,  atol=1E-5       )       )

     # Attaching a hook to a model when it already has one replaces, does not chain
     snake_case_				:  Optional[int]       =			PostForwardHook()
     add_hook_to_module(lowercase_						,  lowercase_       )
     snake_case_				:  List[Any]       =			test_model(lowercase_       )
     self.assertTrue(torch.allclose(lowercase_						,  output + 1						,  atol=1E-5       )       )

     # You need to use the sequential hook to chain two or more hooks
     snake_case_				:  Optional[int]       =			SequentialHook(PostForwardHook()						,  PostForwardHook()       )
     add_hook_to_module(lowercase_						,  lowercase_       )

     snake_case_				:  Union[str, Any]       =			test_model(lowercase_       )
     assert torch.allclose(lowercase_						,  output + 2						,  atol=1E-5       )

 def 		_snake_case      (  self	: Any       ):
     snake_case_				:  int       =			ModelForTest()
     snake_case_				:  Dict       =			torch.randn(2						,  3       )
     snake_case_				:  int       =			test_model(lowercase_       )

     snake_case_				:  Optional[int]       =			PostForwardHook()
     add_hook_to_module(lowercase_						,  lowercase_       )
     snake_case_				:  Optional[Any]       =			test_model(lowercase_       )
     self.assertTrue(torch.allclose(lowercase_						,  output + 1       )       )
     self.assertTrue(outputa.requires_grad       )

     snake_case_				:  str       =			True
     snake_case_				:  str       =			test_model(lowercase_       )
     self.assertFalse(outputa.requires_grad       )

 @require_multi_gpu
 def 		_snake_case      (  self	: Optional[Any]       ):
     snake_case_				:  int       =			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
     snake_case_				:  Optional[int]       =			torch.randn(2						,  3       )
     snake_case_				:  List[str]       =			model(lowercase_       )
     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(lowercase_						,  AlignDevicesHook(io_same_device=lowercase_       )       )
     snake_case_				:  Any       =			torch.randn(2						,  3       ).to(0       )
     snake_case_				:  Union[str, Any]       =			model(lowercase_       )
     self.assertEqual(output.device						,  torch.device(0       )       )

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

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

     # 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
     snake_case_				:  int       =			torch.device(hook_kwargs['''execution_device''']       )
     self.assertEqual(model.batchnorm.running_mean.device						,  lowercase_       )

     snake_case_				:  Union[str, Any]       =			torch.randn(2						,  3       )
     snake_case_				:  Any       =			model(lowercase_       )
     self.assertEqual(output.device						,  lowercase_       )

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

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

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

     snake_case_				:  List[Any]       =			torch.randn(2						,  3       )
     snake_case_				:  List[Any]       =			model(lowercase_       )
     self.assertEqual(output.device						,  lowercase_       )

     # 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 		_snake_case      (  self	: Union[str, Any]       ):
     snake_case_				:  Any       =			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
     snake_case_				:  Optional[Any]       =			0 if torch.cuda.is_available() else '''cpu'''
     attach_align_device_hook(lowercase_						,  execution_device=lowercase_						,  offload=lowercase_       )

     # 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
     snake_case_				:  Any       =			torch.device(lowercase_       )
     self.assertEqual(model.batchnorm.running_mean.device						,  lowercase_       )

     snake_case_				:  Optional[Any]       =			torch.randn(2						,  3       )
     snake_case_				:  Optional[Any]       =			model(lowercase_       )
     self.assertEqual(output.device						,  lowercase_       )

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

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

     snake_case_				:  Any       =			torch.randn(2						,  3       )
     snake_case_				:  List[str]       =			model(lowercase_       )
     self.assertEqual(output.device						,  lowercase_       )

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

     snake_case_				:  str       =			torch.randn(2						,  3       )
     snake_case_				:  Optional[Any]       =			model(lowercase_       )
     self.assertEqual(output.device						,  lowercase_       )

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

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

     snake_case_				:  Optional[int]       =			torch.randn(2						,  3       )
     snake_case_				:  Any       =			model(lowercase_       )
     self.assertEqual(output.device						,  lowercase_       )

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










"""simple docstring"""

from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import BaseOutput, is_torch_available, is_transformers_available




@dataclass
class  _UpperCAmelCase (    lowerCAmelCase__):
 _lowerCAmelCase     :    Union[List[PIL.Image.Image], np.ndarray]
 _lowerCAmelCase     :    Optional[List[bool]]


if is_transformers_available() and is_torch_available():
     from .pipeline_semantic_stable_diffusion import SemanticStableDiffusionPipeline





