



import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, XLMRobertaTokenizer

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class SCREAMING_SNAKE_CASE__					(		SCREAMING_SNAKE_CASE__			,						SCREAMING_SNAKE_CASE__			,						SCREAMING_SNAKE_CASE__			,						unittest.TestCase   ):




    '''simple docstring'''






    __lowerCamelCase      :			List[Any] =				AltDiffusionPipeline
    __lowerCamelCase      :			int =				TEXT_TO_IMAGE_PARAMS
    __lowerCamelCase      :			str =				TEXT_TO_IMAGE_BATCH_PARAMS
    __lowerCamelCase      :			str =				TEXT_TO_IMAGE_IMAGE_PARAMS
    __lowerCamelCase      :			int =				TEXT_TO_IMAGE_IMAGE_PARAMS






    def       _lowerCAmelCase (	self		):
        torch.manual_seed(0		)
        A			:	Optional[Any]										=       UNetaDConditionModel(
            block_out_channels=(32, 64),		layers_per_block=2,		sample_size=32,		in_channels=4,		out_channels=4,		down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D"""),		up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D"""),		cross_attention_dim=32,		)
        A			:	Dict										=       DDIMScheduler(
            beta_start=0.0_0085,		beta_end=0.012,		beta_schedule="""scaled_linear""",		clip_sample=lowerCamelCase__,		set_alpha_to_one=lowerCamelCase__,		)
        torch.manual_seed(0		)
        A			:	Optional[int]										=       AutoencoderKL(
            block_out_channels=[32, 64],		in_channels=3,		out_channels=3,		down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],		up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],		latent_channels=4,		)

        # TODO: address the non-deterministic text encoder (fails for save-load tests)
        # torch.manual_seed(0)
        # text_encoder_config = RobertaSeriesConfig(
        #     hidden_size=32,
        #     project_dim=32,
        #     intermediate_size=37,
        #     layer_norm_eps=1e-05,
        #     num_attention_heads=4,
        #     num_hidden_layers=5,
        #     vocab_size=5002,
        # )
        # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

        torch.manual_seed(0		)
        A			:	Optional[int]										=       CLIPTextConfig(
            bos_token_id=0,		eos_token_id=2,		hidden_size=32,		projection_dim=32,		intermediate_size=37,		layer_norm_eps=1e-05,		num_attention_heads=4,		num_hidden_layers=5,		pad_token_id=1,		vocab_size=5002,		)
        A			:	List[Any]										=       CLIPTextModel(lowerCamelCase__		)

        A			:	Dict										=       XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""		)
        A			:	Tuple										=       77

        A			:	Any										=       {
            """unet""": unet,
            """scheduler""": scheduler,
            """vae""": vae,
            """text_encoder""": text_encoder,
            """tokenizer""": tokenizer,
            """safety_checker""": None,
            """feature_extractor""": None,
        }
        return components






    def       _lowerCAmelCase (	self,		lowerCamelCase__,		lowerCamelCase__=0		):
        if str(lowerCamelCase__		).startswith("""mps"""		):
            A			:	int										=       torch.manual_seed(lowerCamelCase__		)
        else:
            A			:	str										=       torch.Generator(device=lowerCamelCase__		).manual_seed(lowerCamelCase__		)
        A			:	Dict										=       {
            """prompt""": """A painting of a squirrel eating a burger""",
            """generator""": generator,
            """num_inference_steps""": 2,
            """guidance_scale""": 6.0,
            """output_type""": """numpy""",
        }
        return inputs






    def       _lowerCAmelCase (	self		):
        super().test_attention_slicing_forward_pass(expected_max_diff=3e-3		)






    def       _lowerCAmelCase (	self		):
        super().test_inference_batch_single_identical(expected_max_diff=3e-3		)






    def       _lowerCAmelCase (	self		):
        A			:	List[Any]										=       """cpu"""  # ensure determinism for the device-dependent torch.Generator

        A			:	Tuple										=       self.get_dummy_components()
        torch.manual_seed(0		)
        A			:	Dict										=       RobertaSeriesConfig(
            hidden_size=32,		project_dim=32,		intermediate_size=37,		layer_norm_eps=1e-05,		num_attention_heads=4,		num_hidden_layers=5,		vocab_size=5002,		)
        # TODO: remove after fixing the non-deterministic text encoder
        A			:	List[str]										=       RobertaSeriesModelWithTransformation(lowerCamelCase__		)
        A			:	Union[str, Any]										=       text_encoder

        A			:	Tuple										=       AltDiffusionPipeline(**lowerCamelCase__		)
        A			:	Optional[Any]										=       alt_pipe.to(lowerCamelCase__		)
        alt_pipe.set_progress_bar_config(disable=lowerCamelCase__		)

        A			:	List[Any]										=       self.get_dummy_inputs(lowerCamelCase__		)
        A			:	Dict										=       """A photo of an astronaut"""
        A			:	Tuple										=       alt_pipe(**lowerCamelCase__		)
        A			:	List[str]										=       output.images
        A			:	Any										=       image[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)
        A			:	List[Any]										=       np.array(
            [0.574_8162, 0.6044_7145, 0.4882_1217, 0.5010_0636, 0.543_1185, 0.4576_3683, 0.4965_7696, 0.4813_2733, 0.4757_3093]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2






    def       _lowerCAmelCase (	self		):
        A			:	Optional[int]										=       """cpu"""  # ensure determinism for the device-dependent torch.Generator

        A			:	Dict										=       self.get_dummy_components()
        A			:	Optional[Any]										=       PNDMScheduler(skip_prk_steps=lowerCamelCase__		)
        torch.manual_seed(0		)
        A			:	Dict										=       RobertaSeriesConfig(
            hidden_size=32,		project_dim=32,		intermediate_size=37,		layer_norm_eps=1e-05,		num_attention_heads=4,		num_hidden_layers=5,		vocab_size=5002,		)
        # TODO: remove after fixing the non-deterministic text encoder
        A			:	Optional[Any]										=       RobertaSeriesModelWithTransformation(lowerCamelCase__		)
        A			:	List[Any]										=       text_encoder
        A			:	str										=       AltDiffusionPipeline(**lowerCamelCase__		)
        A			:	List[str]										=       alt_pipe.to(lowerCamelCase__		)
        alt_pipe.set_progress_bar_config(disable=lowerCamelCase__		)

        A			:	Optional[Any]										=       self.get_dummy_inputs(lowerCamelCase__		)
        A			:	Dict										=       alt_pipe(**lowerCamelCase__		)
        A			:	List[str]										=       output.images
        A			:	List[Any]										=       image[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)
        A			:	Optional[Any]										=       np.array(
            [0.5160_5093, 0.570_7241, 0.4736_5507, 0.5057_8886, 0.563_3877, 0.464_2503, 0.518_2081, 0.4876_3484, 0.4908_4237]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2






@slow
@require_torch_gpu
class SCREAMING_SNAKE_CASE__					(		unittest.TestCase   ):




    '''simple docstring'''






    def       _lowerCAmelCase (	self		):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()






    def       _lowerCAmelCase (	self		):
        # make sure here that pndm scheduler skips prk
        A			:	Dict										=       AltDiffusionPipeline.from_pretrained("""BAAI/AltDiffusion""",		safety_checker=lowerCamelCase__		)
        A			:	Optional[Any]										=       alt_pipe.to(lowerCamelCase__		)
        alt_pipe.set_progress_bar_config(disable=lowerCamelCase__		)

        A			:	List[Any]										=       """A painting of a squirrel eating a burger"""
        A			:	Any										=       torch.manual_seed(0		)
        A			:	Union[str, Any]										=       alt_pipe([prompt],		generator=lowerCamelCase__,		guidance_scale=6.0,		num_inference_steps=20,		output_type="""np"""		)

        A			:	Any										=       output.images

        A			:	Optional[int]										=       image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        A			:	Tuple										=       np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2






    def       _lowerCAmelCase (	self		):
        A			:	Optional[Any]										=       DDIMScheduler.from_pretrained("""BAAI/AltDiffusion""",		subfolder="""scheduler"""		)

        A			:	Optional[int]										=       AltDiffusionPipeline.from_pretrained("""BAAI/AltDiffusion""",		scheduler=lowerCamelCase__,		safety_checker=lowerCamelCase__		)
        A			:	Optional[Any]										=       alt_pipe.to(lowerCamelCase__		)
        alt_pipe.set_progress_bar_config(disable=lowerCamelCase__		)

        A			:	List[str]										=       """A painting of a squirrel eating a burger"""
        A			:	int										=       torch.manual_seed(0		)

        A			:	List[Any]										=       alt_pipe([prompt],		generator=lowerCamelCase__,		num_inference_steps=2,		output_type="""numpy"""		)
        A			:	int										=       output.images

        A			:	List[Any]										=       image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        A			:	Tuple										=       np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2







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




    '''simple docstring'''






    def __init__(	self		):
        super().__init__()
        A			:	Union[str, Any]										=       nn.Linear(3,		4		)
        A			:	Union[str, Any]										=       nn.BatchNormad(4		)
        A			:	Optional[Any]										=       nn.Linear(4,		5		)






    def       _lowerCAmelCase (	self,		lowerCamelCase__		):
        return self.lineara(self.batchnorm(self.lineara(lowerCamelCase__		)		)		)






class SCREAMING_SNAKE_CASE__					(		SCREAMING_SNAKE_CASE__   ):




    '''simple docstring'''






    def       _lowerCAmelCase (	self,		lowerCamelCase__,		*lowerCamelCase__,		**lowerCamelCase__		):
        return (args[0] + 1,) + args[1:], kwargs






class SCREAMING_SNAKE_CASE__					(		SCREAMING_SNAKE_CASE__   ):




    '''simple docstring'''






    def       _lowerCAmelCase (	self,		lowerCamelCase__,		lowerCamelCase__		):
        return output + 1






class SCREAMING_SNAKE_CASE__					(		unittest.TestCase   ):




    '''simple docstring'''






    def       _lowerCAmelCase (	self		):
        A			:	Tuple										=       ModelForTest()
        A			:	Any										=       ModelHook()

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






    def       _lowerCAmelCase (	self		):
        A			:	Tuple										=       ModelForTest()
        A			:	Optional[int]										=       ModelHook()

        add_hook_to_module(lowerCamelCase__,		lowerCamelCase__		)
        add_hook_to_module(lowerCamelCase__,		lowerCamelCase__,		append=lowerCamelCase__		)

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






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

        A			:	List[str]										=       PreForwardHook()
        add_hook_to_module(lowerCamelCase__,		lowerCamelCase__		)
        A			:	List[str]										=       test_model(lowerCamelCase__		)
        self.assertTrue(torch.allclose(lowerCamelCase__,		lowerCamelCase__,		atol=1e-5		)		)

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

        # You need to use the sequential hook to chain two or more hooks
        A			:	Any										=       SequentialHook(PreForwardHook(),		PreForwardHook()		)
        add_hook_to_module(lowerCamelCase__,		lowerCamelCase__		)

        A			:	Optional[Any]										=       test_model(lowerCamelCase__		)
        assert torch.allclose(lowerCamelCase__,		lowerCamelCase__,		atol=1e-5		)






    def       _lowerCAmelCase (	self		):
        A			:	Tuple										=       ModelForTest()
        A			:	Any										=       torch.randn(2,		3		)
        A			:	Any										=       test_model(lowerCamelCase__		)

        A			:	List[Any]										=       PostForwardHook()
        add_hook_to_module(lowerCamelCase__,		lowerCamelCase__		)
        A			:	str										=       test_model(lowerCamelCase__		)
        self.assertTrue(torch.allclose(lowerCamelCase__,		output + 1,		atol=1e-5		)		)

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

        # You need to use the sequential hook to chain two or more hooks
        A			:	Dict										=       SequentialHook(PostForwardHook(),		PostForwardHook()		)
        add_hook_to_module(lowerCamelCase__,		lowerCamelCase__		)

        A			:	List[str]										=       test_model(lowerCamelCase__		)
        assert torch.allclose(lowerCamelCase__,		output + 2,		atol=1e-5		)






    def       _lowerCAmelCase (	self		):
        A			:	List[Any]										=       ModelForTest()
        A			:	Tuple										=       torch.randn(2,		3		)
        A			:	Union[str, Any]										=       test_model(lowerCamelCase__		)

        A			:	List[Any]										=       PostForwardHook()
        add_hook_to_module(lowerCamelCase__,		lowerCamelCase__		)
        A			:	List[Any]										=       test_model(lowerCamelCase__		)
        self.assertTrue(torch.allclose(lowerCamelCase__,		output + 1		)		)
        self.assertTrue(outputa.requires_grad		)

        A			:	int										=       True
        A			:	Tuple										=       test_model(lowerCamelCase__		)
        self.assertFalse(outputa.requires_grad		)






    @require_multi_gpu
    def       _lowerCAmelCase (	self		):
        A			:	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
        A			:	str										=       torch.randn(2,		3		)
        A			:	int										=       model(lowerCamelCase__		)
        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(lowerCamelCase__,		AlignDevicesHook(io_same_device=lowerCamelCase__		)		)
        A			:	int										=       torch.randn(2,		3		).to(0		)
        A			:	str										=       model(lowerCamelCase__		)
        self.assertEqual(output.device,		torch.device(0		)		)






    def       _lowerCAmelCase (	self		):
        A			:	List[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
        A			:	int										=       {"""execution_device""": 0 if torch.cuda.is_available() else """cpu""", """offload""": True}

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

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

        A			:	int										=       torch.randn(2,		3		)
        A			:	List[Any]										=       model(lowerCamelCase__		)
        self.assertEqual(output.device,		lowerCamelCase__		)

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

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

        # 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			:	int										=       torch.randn(2,		3		)
        A			:	str										=       model(lowerCamelCase__		)
        self.assertEqual(output.device,		lowerCamelCase__		)

        # 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		):
        A			:	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
        A			:	Tuple										=       0 if torch.cuda.is_available() else """cpu"""
        attach_align_device_hook(lowerCamelCase__,		execution_device=lowerCamelCase__,		offload=lowerCamelCase__		)

        # 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			:	Optional[int]										=       torch.device(lowerCamelCase__		)
        self.assertEqual(model.batchnorm.running_mean.device,		lowerCamelCase__		)

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

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

        # 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			:	List[str]										=       torch.randn(2,		3		)
        A			:	List[str]										=       model(lowerCamelCase__		)
        self.assertEqual(output.device,		lowerCamelCase__		)

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

        A			:	Optional[Any]										=       torch.randn(2,		3		)
        A			:	Tuple										=       model(lowerCamelCase__		)
        self.assertEqual(output.device,		lowerCamelCase__		)

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

        # 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			:	List[Any]										=       torch.randn(2,		3		)
        A			:	Dict										=       model(lowerCamelCase__		)
        self.assertEqual(output.device,		lowerCamelCase__		)

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


