



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







													'''simple docstring'''

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






						def 			_A      (       self ,				_A						):







													'''simple docstring'''

													return self.lineara(self.batchnorm(self.lineara(a_						)						)						)







class   UpperCAmelCase_    (  UpperCamelCase_			):

						'''simple docstring'''




						def 			_A      (       self ,				_A ,				*_A ,				**_A						):







													'''simple docstring'''

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







class   UpperCAmelCase_    (  UpperCamelCase_			):

						'''simple docstring'''




						def 			_A      (       self ,				_A ,				_A						):







													'''simple docstring'''

													return output + 1







class   UpperCAmelCase_    (  unittest.TestCase			):

						'''simple docstring'''




						def 			_A      (       self						):







													'''simple docstring'''

													__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 			_A      (       self						):







													'''simple docstring'''

													__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 			_A      (       self						):







													'''simple docstring'''

													__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 			_A      (       self						):







													'''simple docstring'''

													__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 			_A      (       self						):







													'''simple docstring'''

													__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 			_A      (       self						):







													'''simple docstring'''

													__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 			_A      (       self						):







													'''simple docstring'''

													__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 			_A      (       self						):







													'''simple docstring'''

													__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 			_A      (       self						):







													'''simple docstring'''

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











'''simple docstring'''

import torch

from diffusers import CMStochasticIterativeScheduler

from .test_schedulers import SchedulerCommonTest






class 						a_		(     snake_case    ):
							UpperCAmelCase			:		str        =				(CMStochasticIterativeScheduler,)
							UpperCAmelCase			:		int        =				10


							def    UpperCamelCase     (				self   :						Dict      ,						**a_   :						List[str]			)			->					Any:
										snake_case:    Any      							={
										    'num_train_timesteps': 2_0_1,
										    'sigma_min': 0.0_0_2,
										    'sigma_max': 8_0.0,
										}

										config.update(**a_			)
										return config


							def    UpperCamelCase     (				self   :						List[Any]			)			->					List[Any]:
										snake_case:    Any      							=1_0

										snake_case:    List[str]      							=self.get_scheduler_config()
										snake_case:    List[Any]      							=self.scheduler_classes[0](**a_			)

										scheduler.set_timesteps(a_			)

										snake_case:    Dict      							=scheduler.timesteps[0]
										snake_case:    Union[str, Any]      							=scheduler.timesteps[1]

										snake_case:    List[str]      							=self.dummy_sample
										snake_case:    List[str]      							=0.1 * sample

										snake_case:    int      							=scheduler.step(a_      ,						a_      ,						a_			).prev_sample
										snake_case:    Optional[Any]      							=scheduler.step(a_      ,						a_      ,						a_			).prev_sample

										self.assertEqual(output_a.shape      ,						sample.shape			)
										self.assertEqual(output_a.shape      ,						output_a.shape			)


							def    UpperCamelCase     (				self   :						int			)			->					int:
										for timesteps in [1_0, 5_0, 1_0_0, 1_0_0_0]:
													self.check_over_configs(num_train_timesteps=a_			)


							def    UpperCamelCase     (				self   :						Optional[Any]			)			->					Dict:
										for clip_denoised in [True, False]:
													self.check_over_configs(clip_denoised=a_			)


							def    UpperCamelCase     (				self   :						Tuple			)			->					List[str]:
										snake_case:    List[Any]      							=self.scheduler_classes[0]
										snake_case:    List[Any]      							=self.get_scheduler_config()
										snake_case:    Any      							=scheduler_class(**a_			)

										snake_case:    Dict      							=1
										scheduler.set_timesteps(a_			)
										snake_case:    List[Any]      							=scheduler.timesteps

										snake_case:    Optional[Any]      							=torch.manual_seed(0			)

										snake_case:    Optional[Any]      							=self.dummy_model()
										snake_case:    List[Any]      							=self.dummy_sample_deter * scheduler.init_noise_sigma

										for i, t in enumerate(a_			):
													# 1. scale model input
													snake_case:    Any      							=scheduler.scale_model_input(a_      ,						a_			)

													# 2. predict noise residual
													snake_case:    List[str]      							=model(a_      ,						a_			)

													# 3. predict previous sample x_t-1
													snake_case:    Dict      							=scheduler.step(a_      ,						a_      ,						a_      ,						generator=a_			).prev_sample

													snake_case:    List[Any]      							=pred_prev_sample

										snake_case:    Optional[Any]      							=torch.sum(torch.abs(a_			)			)
										snake_case:    Optional[Any]      							=torch.mean(torch.abs(a_			)			)

										assert abs(result_sum.item() - 1_9_2.7_6_1_4			) < 1E-2
										assert abs(result_mean.item() - 0.2_5_1_0			) < 1E-3


							def    UpperCamelCase     (				self   :						Dict			)			->					Union[str, Any]:
										snake_case:    Dict      							=self.scheduler_classes[0]
										snake_case:    Tuple      							=self.get_scheduler_config()
										snake_case:    str      							=scheduler_class(**a_			)

										snake_case:    List[Any]      							=[1_0_6, 0]
										scheduler.set_timesteps(timesteps=a_			)
										snake_case:    Optional[Any]      							=scheduler.timesteps

										snake_case:    Dict      							=torch.manual_seed(0			)

										snake_case:    Optional[int]      							=self.dummy_model()
										snake_case:    Any      							=self.dummy_sample_deter * scheduler.init_noise_sigma

										for t in timesteps:
													# 1. scale model input
													snake_case:    List[Any]      							=scheduler.scale_model_input(a_      ,						a_			)

													# 2. predict noise residual
													snake_case:    Any      							=model(a_      ,						a_			)

													# 3. predict previous sample x_t-1
													snake_case:    List[str]      							=scheduler.step(a_      ,						a_      ,						a_      ,						generator=a_			).prev_sample

													snake_case:    Optional[Any]      							=pred_prev_sample

										snake_case:    Union[str, Any]      							=torch.sum(torch.abs(a_			)			)
										snake_case:    Tuple      							=torch.mean(torch.abs(a_			)			)

										assert abs(result_sum.item() - 3_4_7.6_3_5_7			) < 1E-2
										assert abs(result_mean.item() - 0.4_5_2_7			) < 1E-3


							def    UpperCamelCase     (				self   :						int			)			->					Tuple:
										snake_case:    List[Any]      							=self.scheduler_classes[0]
										snake_case:    Union[str, Any]      							=self.get_scheduler_config()
										snake_case:    str      							=scheduler_class(**a_			)

										snake_case:    str      							=[3_9, 3_0, 1_2, 1_5, 0]

										with self.assertRaises(a_      ,						msg='`timesteps` must be in descending order.'			):
													scheduler.set_timesteps(timesteps=a_			)


							def    UpperCamelCase     (				self   :						Dict			)			->					Optional[int]:
										snake_case:    Optional[Any]      							=self.scheduler_classes[0]
										snake_case:    Dict      							=self.get_scheduler_config()
										snake_case:    str      							=scheduler_class(**a_			)

										snake_case:    Any      							=[3_9, 3_0, 1_2, 1, 0]
										snake_case:    List[Any]      							=len(a_			)

										with self.assertRaises(a_      ,						msg='Can only pass one of `num_inference_steps` or `timesteps`.'			):
													scheduler.set_timesteps(num_inference_steps=a_      ,						timesteps=a_			)




							def    UpperCamelCase     (				self   :						Optional[Any]			)			->					Tuple:
										snake_case:    Any      							=self.scheduler_classes[0]
										snake_case:    int      							=self.get_scheduler_config()
										snake_case:    Optional[Any]      							=scheduler_class(**a_			)

										snake_case:    List[Any]      							=[scheduler.config.num_train_timesteps]

										with self.assertRaises(
										    a_      ,						msg='`timesteps` must start before `self.config.train_timesteps`: {scheduler.config.num_train_timesteps}}'      ,						):
													scheduler.set_timesteps(timesteps=a_			)




