





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



       """simple docstring"""





       super().__init__()
       A_          =       nn.Linear(3 ,4 )
       A_          =       nn.BatchNormad(4 )
       A_          =       nn.Linear(4 ,5 )




  def 						__UpperCAmelCase     (							self ,__snake_case ):



       """simple docstring"""





       return self.lineara(self.batchnorm(self.lineara(__snake_case ) ) )





class UpperCAmelCase_				( _SCREAMING_SNAKE_CASE ):


  def 						__UpperCAmelCase     (							self ,__snake_case ,*__snake_case ,**__snake_case ):



       """simple docstring"""





       return (args[0] + 1,) + args[1:], kwargs





class UpperCAmelCase_				( _SCREAMING_SNAKE_CASE ):


  def 						__UpperCAmelCase     (							self ,__snake_case ,__snake_case ):



       """simple docstring"""





       return output + 1





class UpperCAmelCase_				( unittest.TestCase ):


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       ModelForTest()
       A_          =       ModelHook()

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


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       ModelForTest()
       A_          =       ModelHook()

       add_hook_to_module(__snake_case ,__snake_case )
       add_hook_to_module(__snake_case ,__snake_case ,append=__snake_case )

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


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       ModelForTest()
       A_          =       torch.randn(2 ,3 )
       A_          =       test_model(x + 1 )
       A_          =       test_model(x + 2 )

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

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

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

       A_          =       test_model(__snake_case )
       assert torch.allclose(__snake_case ,__snake_case ,atol=1E-5 )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       ModelForTest()
       A_          =       torch.randn(2 ,3 )
       A_          =       test_model(__snake_case )

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

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

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

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


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       ModelForTest()
       A_          =       torch.randn(2 ,3 )
       A_          =       test_model(__snake_case )

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

       A_          =       True
       A_          =       test_model(__snake_case )
       self.assertFalse(outputa.requires_grad )


  @require_multi_gpu
  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       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
       A_          =       torch.randn(2 ,3 )
       A_          =       model(__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(__snake_case ,AlignDevicesHook(io_same_device=__snake_case ) )
       A_          =       torch.randn(2 ,3 ).to(0 )
       A_          =       model(__snake_case )
       self.assertEqual(output.device ,torch.device(0 ) )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





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

       add_hook_to_module(model.lineara ,AlignDevicesHook(**__snake_case ) )
       add_hook_to_module(model.batchnorm ,AlignDevicesHook(**__snake_case ) )
       add_hook_to_module(model.lineara ,AlignDevicesHook(**__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
       A_          =       torch.device(hook_kwargs['''execution_device'''] )
       self.assertEqual(model.batchnorm.running_mean.device ,__snake_case )

       A_          =       torch.randn(2 ,3 )
       A_          =       model(__snake_case )
       self.assertEqual(output.device ,__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
       A_          =       {
           '''execution_device''': 0 if torch.cuda.is_available() else '''cpu''',
           '''offload''': True,
           '''offload_buffers''': True,
       }

       add_hook_to_module(model.lineara ,AlignDevicesHook(**__snake_case ) )
       add_hook_to_module(model.batchnorm ,AlignDevicesHook(**__snake_case ) )
       add_hook_to_module(model.lineara ,AlignDevicesHook(**__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''' ) )

       A_          =       torch.randn(2 ,3 )
       A_          =       model(__snake_case )
       self.assertEqual(output.device ,__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 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       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
       A_          =       0 if torch.cuda.is_available() else '''cpu'''
       attach_align_device_hook(__snake_case ,execution_device=__snake_case ,offload=__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
       A_          =       torch.device(__snake_case )
       self.assertEqual(model.batchnorm.running_mean.device ,__snake_case )

       A_          =       torch.randn(2 ,3 )
       A_          =       model(__snake_case )
       self.assertEqual(output.device ,__snake_case )

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

       A_          =       torch.randn(2 ,3 )
       A_          =       model(__snake_case )
       self.assertEqual(output.device ,__snake_case )

       # Removing hooks loads back the weights in the model.
       remove_hook_from_submodules(__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 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       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
       A_          =       0 if torch.cuda.is_available() else '''cpu'''
       attach_align_device_hook(
           __snake_case ,execution_device=__snake_case ,offload=__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
       A_          =       torch.device(__snake_case )
       self.assertEqual(model.batchnorm.running_mean.device ,__snake_case )

       A_          =       torch.randn(2 ,3 )
       A_          =       model(__snake_case )
       self.assertEqual(output.device ,__snake_case )

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

       A_          =       torch.randn(2 ,3 )
       A_          =       model(__snake_case )
       self.assertEqual(output.device ,__snake_case )

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










import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest





class UpperCAmelCase_				( _SCREAMING_SNAKE_CASE ):
  __lowerCAmelCase   :						int							=   (KDPMaDiscreteScheduler,)
  __lowerCAmelCase   :						Union[str, Any]							=   10


  def 						__UpperCAmelCase     (							self ,**__snake_case ):



       """simple docstring"""





       A_          =       {
           '''num_train_timesteps''': 1_1_0_0,
           '''beta_start''': 0.0001,
           '''beta_end''': 0.02,
           '''beta_schedule''': '''linear''',
       }

       config.update(**__snake_case )
       return config


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       for timesteps in [1_0, 5_0, 1_0_0, 1_0_0_0]:
            self.check_over_configs(num_train_timesteps=__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001] ,[0.0002, 0.002, 0.02] ):
            self.check_over_configs(beta_start=__snake_case ,beta_end=__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       for schedule in ["linear", "scaled_linear"]:
            self.check_over_configs(beta_schedule=__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       for prediction_type in ["epsilon", "v_prediction"]:
            self.check_over_configs(prediction_type=__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.scheduler_classes[0]
       A_          =       self.get_scheduler_config(prediction_type='''v_prediction''' )
       A_          =       scheduler_class(**__snake_case )

       scheduler.set_timesteps(self.num_inference_steps )

       A_          =       self.dummy_model()
       A_          =       self.dummy_sample_deter * scheduler.init_noise_sigma
       A_          =       sample.to(__snake_case )

       for i, t in enumerate(scheduler.timesteps ):
            A_          =       scheduler.scale_model_input(__snake_case ,__snake_case )

            A_          =       model(__snake_case ,__snake_case )

            A_          =       scheduler.step(__snake_case ,__snake_case ,__snake_case )
            A_          =       output.prev_sample

       A_          =       torch.sum(torch.abs(__snake_case ) )
       A_          =       torch.mean(torch.abs(__snake_case ) )

       if torch_device in ["cpu", "mps"]:
            assert abs(result_sum.item() - 4.69_34E-07 ) < 1E-2
            assert abs(result_mean.item() - 6.11_12E-10 ) < 1E-3
       else:
            # CUDA
            assert abs(result_sum.item() - 4.6_93_42_86_50_17_09_72E-07 ) < 1E-2
            assert abs(result_mean.item() - 0.0002 ) < 1E-3


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       if torch_device == "mps":
            return
       A_          =       self.scheduler_classes[0]
       A_          =       self.get_scheduler_config()
       A_          =       scheduler_class(**__snake_case )

       scheduler.set_timesteps(self.num_inference_steps )

       A_          =       self.dummy_model()
       A_          =       self.dummy_sample_deter * scheduler.init_noise_sigma
       A_          =       sample.to(__snake_case )

       for i, t in enumerate(scheduler.timesteps ):
            A_          =       scheduler.scale_model_input(__snake_case ,__snake_case )

            A_          =       model(__snake_case ,__snake_case )

            A_          =       scheduler.step(__snake_case ,__snake_case ,__snake_case )
            A_          =       output.prev_sample

       A_          =       torch.sum(torch.abs(__snake_case ) )
       A_          =       torch.mean(torch.abs(__snake_case ) )

       if torch_device in ["cpu", "mps"]:
            assert abs(result_sum.item() - 20.4125 ) < 1E-2
            assert abs(result_mean.item() - 0.0266 ) < 1E-3
       else:
            # CUDA
            assert abs(result_sum.item() - 20.4125 ) < 1E-2
            assert abs(result_mean.item() - 0.0266 ) < 1E-3




  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       if torch_device == "mps":
            return
       A_          =       self.scheduler_classes[0]
       A_          =       self.get_scheduler_config()
       A_          =       scheduler_class(**__snake_case )

       scheduler.set_timesteps(self.num_inference_steps ,device=__snake_case )

       A_          =       self.dummy_model()
       A_          =       self.dummy_sample_deter.to(__snake_case ) * scheduler.init_noise_sigma

       for t in scheduler.timesteps:
            A_          =       scheduler.scale_model_input(__snake_case ,__snake_case )

            A_          =       model(__snake_case ,__snake_case )

            A_          =       scheduler.step(__snake_case ,__snake_case ,__snake_case )
            A_          =       output.prev_sample

       A_          =       torch.sum(torch.abs(__snake_case ) )
       A_          =       torch.mean(torch.abs(__snake_case ) )

       if str(__snake_case ).startswith('''cpu''' ):
            # The following sum varies between 148 and 156 on mps. Why?
            assert abs(result_sum.item() - 20.4125 ) < 1E-2
            assert abs(result_mean.item() - 0.0266 ) < 1E-3
       else:
            # CUDA
            assert abs(result_sum.item() - 20.4125 ) < 1E-2
            assert abs(result_mean.item() - 0.0266 ) < 1E-3



