


"""simple docstring"""






UpperCAmelCase    :     Union[str, Any]   						=		'0.18.2'

from .configuration_utils import ConfigMixin
from .utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_inflect_available,
    is_invisible_watermark_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_librosa_available,
    is_note_seq_available,
    is_onnx_available,
    is_scipy_available,
    is_torch_available,
    is_torchsde_available,
    is_transformers_available,
    is_transformers_version,
    is_unidecode_available,
    logging,
)


try:
							if not is_onnx_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from .utils.dummy_onnx_objects import *  # noqa F403
else:
							from .pipelines import OnnxRuntimeModel

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from .utils.dummy_pt_objects import *  # noqa F403
else:
							from .models import (
							    AutoencoderKL,
							    ControlNetModel,
							    ModelMixin,
							    PriorTransformer,
							    TaFilmDecoder,
							    TransformeraDModel,
							    UNetaDModel,
							    UNetaDConditionModel,
							    UNetaDModel,
							    UNetaDConditionModel,
							    VQModel,
							)
							from .optimization import (
							    get_constant_schedule,
							    get_constant_schedule_with_warmup,
							    get_cosine_schedule_with_warmup,
							    get_cosine_with_hard_restarts_schedule_with_warmup,
							    get_linear_schedule_with_warmup,
							    get_polynomial_decay_schedule_with_warmup,
							    get_scheduler,
							)
							from .pipelines import (
							    AudioPipelineOutput,
							    ConsistencyModelPipeline,
							    DanceDiffusionPipeline,
							    DDIMPipeline,
							    DDPMPipeline,
							    DiffusionPipeline,
							    DiTPipeline,
							    ImagePipelineOutput,
							    KarrasVePipeline,
							    LDMPipeline,
							    LDMSuperResolutionPipeline,
							    PNDMPipeline,
							    RePaintPipeline,
							    ScoreSdeVePipeline,
							)
							from .schedulers import (
							    CMStochasticIterativeScheduler,
							    DDIMInverseScheduler,
							    DDIMParallelScheduler,
							    DDIMScheduler,
							    DDPMParallelScheduler,
							    DDPMScheduler,
							    DEISMultistepScheduler,
							    DPMSolverMultistepInverseScheduler,
							    DPMSolverMultistepScheduler,
							    DPMSolverSinglestepScheduler,
							    EulerAncestralDiscreteScheduler,
							    EulerDiscreteScheduler,
							    HeunDiscreteScheduler,
							    IPNDMScheduler,
							    KarrasVeScheduler,
							    KDPMaAncestralDiscreteScheduler,
							    KDPMaDiscreteScheduler,
							    PNDMScheduler,
							    RePaintScheduler,
							    SchedulerMixin,
							    ScoreSdeVeScheduler,
							    UnCLIPScheduler,
							    UniPCMultistepScheduler,
							    VQDiffusionScheduler,
							)
							from .training_utils import EMAModel

try:
							if not (is_torch_available() and is_scipy_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from .utils.dummy_torch_and_scipy_objects import *  # noqa F403
else:
							from .schedulers import LMSDiscreteScheduler

try:
							if not (is_torch_available() and is_torchsde_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from .utils.dummy_torch_and_torchsde_objects import *  # noqa F403
else:
							from .schedulers import DPMSolverSDEScheduler

try:
							if not (is_torch_available() and is_transformers_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from .utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
							from .pipelines import (
							    AltDiffusionImgaImgPipeline,
							    AltDiffusionPipeline,
							    AudioLDMPipeline,
							    CycleDiffusionPipeline,
							    IFImgaImgPipeline,
							    IFImgaImgSuperResolutionPipeline,
							    IFInpaintingPipeline,
							    IFInpaintingSuperResolutionPipeline,
							    IFPipeline,
							    IFSuperResolutionPipeline,
							    ImageTextPipelineOutput,
							    KandinskyImgaImgPipeline,
							    KandinskyInpaintPipeline,
							    KandinskyPipeline,
							    KandinskyPriorPipeline,
							    KandinskyVaaControlnetImgaImgPipeline,
							    KandinskyVaaControlnetPipeline,
							    KandinskyVaaImgaImgPipeline,
							    KandinskyVaaInpaintPipeline,
							    KandinskyVaaPipeline,
							    KandinskyVaaPriorEmbaEmbPipeline,
							    KandinskyVaaPriorPipeline,
							    LDMTextToImagePipeline,
							    PaintByExamplePipeline,
							    SemanticStableDiffusionPipeline,
							    ShapEImgaImgPipeline,
							    ShapEPipeline,
							    StableDiffusionAttendAndExcitePipeline,
							    StableDiffusionControlNetImgaImgPipeline,
							    StableDiffusionControlNetInpaintPipeline,
							    StableDiffusionControlNetPipeline,
							    StableDiffusionDepthaImgPipeline,
							    StableDiffusionDiffEditPipeline,
							    StableDiffusionImageVariationPipeline,
							    StableDiffusionImgaImgPipeline,
							    StableDiffusionInpaintPipeline,
							    StableDiffusionInpaintPipelineLegacy,
							    StableDiffusionInstructPixaPixPipeline,
							    StableDiffusionLatentUpscalePipeline,
							    StableDiffusionLDMaDPipeline,
							    StableDiffusionModelEditingPipeline,
							    StableDiffusionPanoramaPipeline,
							    StableDiffusionParadigmsPipeline,
							    StableDiffusionPipeline,
							    StableDiffusionPipelineSafe,
							    StableDiffusionPixaPixZeroPipeline,
							    StableDiffusionSAGPipeline,
							    StableDiffusionUpscalePipeline,
							    StableUnCLIPImgaImgPipeline,
							    StableUnCLIPPipeline,
							    TextToVideoSDPipeline,
							    TextToVideoZeroPipeline,
							    UnCLIPImageVariationPipeline,
							    UnCLIPPipeline,
							    UniDiffuserModel,
							    UniDiffuserPipeline,
							    UniDiffuserTextDecoder,
							    VersatileDiffusionDualGuidedPipeline,
							    VersatileDiffusionImageVariationPipeline,
							    VersatileDiffusionPipeline,
							    VersatileDiffusionTextToImagePipeline,
							    VideoToVideoSDPipeline,
							    VQDiffusionPipeline,
							)

try:
							if not (is_torch_available() and is_transformers_available() and is_invisible_watermark_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from .utils.dummy_torch_and_transformers_and_invisible_watermark_objects import *  # noqa F403
else:
							from .pipelines import StableDiffusionXLImgaImgPipeline, StableDiffusionXLPipeline

try:
							if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
else:
							from .pipelines import StableDiffusionKDiffusionPipeline

try:
							if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from .utils.dummy_torch_and_transformers_and_onnx_objects import *  # noqa F403
else:
							from .pipelines import (
							    OnnxStableDiffusionImgaImgPipeline,
							    OnnxStableDiffusionInpaintPipeline,
							    OnnxStableDiffusionInpaintPipelineLegacy,
							    OnnxStableDiffusionPipeline,
							    OnnxStableDiffusionUpscalePipeline,
							    StableDiffusionOnnxPipeline,
							)

try:
							if not (is_torch_available() and is_librosa_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from .utils.dummy_torch_and_librosa_objects import *  # noqa F403
else:
							from .pipelines import AudioDiffusionPipeline, Mel

try:
							if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from .utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
else:
							from .pipelines import SpectrogramDiffusionPipeline

try:
							if not is_flax_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from .utils.dummy_flax_objects import *  # noqa F403
else:
							from .models.controlnet_flax import FlaxControlNetModel
							from .models.modeling_flax_utils import FlaxModelMixin
							from .models.unet_ad_condition_flax import FlaxUNetaDConditionModel
							from .models.vae_flax import FlaxAutoencoderKL
							from .pipelines import FlaxDiffusionPipeline
							from .schedulers import (
							    FlaxDDIMScheduler,
							    FlaxDDPMScheduler,
							    FlaxDPMSolverMultistepScheduler,
							    FlaxKarrasVeScheduler,
							    FlaxLMSDiscreteScheduler,
							    FlaxPNDMScheduler,
							    FlaxSchedulerMixin,
							    FlaxScoreSdeVeScheduler,
							)


try:
							if not (is_flax_available() and is_transformers_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
else:
							from .pipelines import (
							    FlaxStableDiffusionControlNetPipeline,
							    FlaxStableDiffusionImgaImgPipeline,
							    FlaxStableDiffusionInpaintPipeline,
							    FlaxStableDiffusionPipeline,
							)

try:
							if not (is_note_seq_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from .utils.dummy_note_seq_objects import *  # noqa F403
else:
							from .pipelines import MidiProcessor




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







					"""simple docstring"""


					def __init__(       self    : List[Any]  ):
												'''simple docstring'''






												super().__init__()
												__UpperCAmelCase   :						Optional[Any]      							=					nn.Linear(3 ,       4  )
												__UpperCAmelCase   :						Optional[int]      							=					nn.BatchNormad(4  )
												__UpperCAmelCase   :						int      							=					nn.Linear(4 ,       5  )






					def lowerCamelCase__							(       self    : List[str] ,       UpperCamelCase    : Dict  ):
												'''simple docstring'''






												return self.lineara(self.batchnorm(self.lineara(UpperCamelCase  )  )  )






class      lowerCamelCase__   (     A				):







					"""simple docstring"""


					def lowerCamelCase__							(       self    : str ,       UpperCamelCase    : Any ,       *UpperCamelCase    : List[Any] ,       **UpperCamelCase    : Dict  ):
												'''simple docstring'''






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






class      lowerCamelCase__   (     A				):







					"""simple docstring"""


					def lowerCamelCase__							(       self    : Any ,       UpperCamelCase    : Dict ,       UpperCamelCase    : List[Any]  ):
												'''simple docstring'''






												return output + 1






class      lowerCamelCase__   (     unittest.TestCase				):







					"""simple docstring"""


					def lowerCamelCase__							(       self    : str  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Dict      							=					ModelForTest()
												__UpperCAmelCase   :						List[str]      							=					ModelHook()

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






					def lowerCamelCase__							(       self    : str  ):
												'''simple docstring'''






												__UpperCAmelCase   :						int      							=					ModelForTest()
												__UpperCAmelCase   :						Optional[Any]      							=					ModelHook()

												add_hook_to_module(UpperCamelCase ,       UpperCamelCase  )
												add_hook_to_module(UpperCamelCase ,       UpperCamelCase ,       append=UpperCamelCase  )

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






					def lowerCamelCase__							(       self    : Any  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Union[str, Any]      							=					ModelForTest()
												__UpperCAmelCase   :						str      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						List[str]      							=					test_model(x + 1  )
												__UpperCAmelCase   :						Optional[int]      							=					test_model(x + 2  )

												__UpperCAmelCase   :						Optional[Any]      							=					PreForwardHook()
												add_hook_to_module(UpperCamelCase ,       UpperCamelCase  )
												__UpperCAmelCase   :						Optional[Any]      							=					test_model(UpperCamelCase  )
												self.assertTrue(torch.allclose(UpperCamelCase ,       UpperCamelCase ,       atol=1e-5  )  )

												# Attaching a hook to a model when it already has one replaces, does not chain
												__UpperCAmelCase   :						Optional[Any]      							=					PreForwardHook()
												add_hook_to_module(UpperCamelCase ,       UpperCamelCase  )
												__UpperCAmelCase   :						Optional[int]      							=					test_model(UpperCamelCase  )
												self.assertTrue(torch.allclose(UpperCamelCase ,       UpperCamelCase ,       atol=1e-5  )  )

												# You need to use the sequential hook to chain two or more hooks
												__UpperCAmelCase   :						Optional[Any]      							=					SequentialHook(PreForwardHook() ,       PreForwardHook()  )
												add_hook_to_module(UpperCamelCase ,       UpperCamelCase  )

												__UpperCAmelCase   :						List[str]      							=					test_model(UpperCamelCase  )
												assert torch.allclose(UpperCamelCase ,       UpperCamelCase ,       atol=1e-5  )






					def lowerCamelCase__							(       self    : Dict  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Optional[int]      							=					ModelForTest()
												__UpperCAmelCase   :						str      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						int      							=					test_model(UpperCamelCase  )

												__UpperCAmelCase   :						Dict      							=					PostForwardHook()
												add_hook_to_module(UpperCamelCase ,       UpperCamelCase  )
												__UpperCAmelCase   :						Any      							=					test_model(UpperCamelCase  )
												self.assertTrue(torch.allclose(UpperCamelCase ,       output + 1 ,       atol=1e-5  )  )

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

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

												__UpperCAmelCase   :						int      							=					test_model(UpperCamelCase  )
												assert torch.allclose(UpperCamelCase ,       output + 2 ,       atol=1e-5  )






					def lowerCamelCase__							(       self    : Optional[int]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						int      							=					ModelForTest()
												__UpperCAmelCase   :						str      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						Optional[Any]      							=					test_model(UpperCamelCase  )

												__UpperCAmelCase   :						Union[str, Any]      							=					PostForwardHook()
												add_hook_to_module(UpperCamelCase ,       UpperCamelCase  )
												__UpperCAmelCase   :						Optional[int]      							=					test_model(UpperCamelCase  )
												self.assertTrue(torch.allclose(UpperCamelCase ,       output + 1  )  )
												self.assertTrue(outputa.requires_grad  )

												__UpperCAmelCase   :						List[str]      							=					True
												__UpperCAmelCase   :						Optional[int]      							=					test_model(UpperCamelCase  )
												self.assertFalse(outputa.requires_grad  )






					@require_multi_gpu
					def lowerCamelCase__							(       self    : Optional[Any]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						str      							=					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
												__UpperCAmelCase   :						List[str]      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						Any      							=					model(UpperCamelCase  )
												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(UpperCamelCase ,       AlignDevicesHook(io_same_device=UpperCamelCase  )  )
												__UpperCAmelCase   :						int      							=					torch.randn(2 ,       3  ).to(0  )
												__UpperCAmelCase   :						Any      							=					model(UpperCamelCase  )
												self.assertEqual(output.device ,       torch.device(0  )  )






					def lowerCamelCase__							(       self    : Optional[int]  ):
												'''simple docstring'''






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

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

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

												__UpperCAmelCase   :						Union[str, Any]      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						Optional[int]      							=					model(UpperCamelCase  )
												self.assertEqual(output.device ,       UpperCamelCase  )

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

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

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

												__UpperCAmelCase   :						str      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						Dict      							=					model(UpperCamelCase  )
												self.assertEqual(output.device ,       UpperCamelCase  )

												# 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 lowerCamelCase__							(       self    : Union[str, Any]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						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
												__UpperCAmelCase   :						Optional[int]      							=					0 if torch.cuda.is_available() else """cpu"""
												attach_align_device_hook(UpperCamelCase ,       execution_device=UpperCamelCase ,       offload=UpperCamelCase  )

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

												__UpperCAmelCase   :						str      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						Any      							=					model(UpperCamelCase  )
												self.assertEqual(output.device ,       UpperCamelCase  )

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

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

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

												# Removing hooks loads back the weights in the model.
												remove_hook_from_submodules(UpperCamelCase  )
												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 lowerCamelCase__							(       self    : Any  ):
												'''simple docstring'''






												__UpperCAmelCase   :						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
												__UpperCAmelCase   :						Any      							=					0 if torch.cuda.is_available() else """cpu"""
												attach_align_device_hook(
												    UpperCamelCase ,       execution_device=UpperCamelCase ,       offload=UpperCamelCase ,       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
												__UpperCAmelCase   :						str      							=					torch.device(UpperCamelCase  )
												self.assertEqual(model.batchnorm.running_mean.device ,       UpperCamelCase  )

												__UpperCAmelCase   :						Dict      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						Optional[Any]      							=					model(UpperCamelCase  )
												self.assertEqual(output.device ,       UpperCamelCase  )

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

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

												__UpperCAmelCase   :						int      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						List[str]      							=					model(UpperCamelCase  )
												self.assertEqual(output.device ,       UpperCamelCase  )

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