

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







    """simple docstring"""





    def __init__(      self :       List[str]		)->    Union[str, Any]:





        '''simple docstring'''







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


    def 						snake_case__     (      self :       Optional[int],lowercase_ :       List[Any]		)->    Any:





        '''simple docstring'''







        return self.lineara(self.batchnorm(self.lineara(lowercase_		)		)		)




class  A						(							_UpperCAmelCase ):







    """simple docstring"""





    def 						snake_case__     (      self :       List[Any],lowercase_ :       List[str],*lowercase_ :       Optional[int],**lowercase_ :       Tuple		)->    Any:





        '''simple docstring'''







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




class  A						(							_UpperCAmelCase ):







    """simple docstring"""





    def 						snake_case__     (      self :       List[str],lowercase_ :       int,lowercase_ :       Optional[Any]		)->    List[str]:





        '''simple docstring'''







        return output + 1




class  A						(							unittest.TestCase ):







    """simple docstring"""





    def 						snake_case__     (      self :       Union[str, Any]		)->    Tuple:





        '''simple docstring'''







        A__          =   ModelForTest()
        A__          =   ModelHook()

        add_hook_to_module(lowercase_,lowercase_		)
        self.assertEqual(test_model._hf_hook,lowercase_		)
        self.assertTrue(hasattr(lowercase_,'_old_forward'		)		)

        # Check adding the hook did not change the name or the signature
        self.assertEqual(test_model.forward.__name__,'forward'		)
        self.assertListEqual(list(inspect.signature(test_model.forward		).parameters		),['x']		)

        remove_hook_from_module(lowercase_		)
        self.assertFalse(hasattr(lowercase_,'_hf_hook'		)		)
        self.assertFalse(hasattr(lowercase_,'_old_forward'		)		)


    def 						snake_case__     (      self :       Optional[Any]		)->    str:





        '''simple docstring'''







        A__          =   ModelForTest()
        A__          =   ModelHook()

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

        self.assertEqual(isinstance(test_model._hf_hook,lowercase_		),lowercase_		)
        self.assertEqual(len(test_model._hf_hook.hooks		),2		)
        self.assertTrue(hasattr(lowercase_,'_old_forward'		)		)

        # Check adding the hook did not change the name or the signature
        self.assertEqual(test_model.forward.__name__,'forward'		)
        self.assertListEqual(list(inspect.signature(test_model.forward		).parameters		),['x']		)

        remove_hook_from_module(lowercase_		)
        self.assertFalse(hasattr(lowercase_,'_hf_hook'		)		)
        self.assertFalse(hasattr(lowercase_,'_old_forward'		)		)


    def 						snake_case__     (      self :       str		)->    Any:





        '''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(lowercase_,lowercase_		)
        A__          =   test_model(lowercase_		)
        self.assertTrue(torch.allclose(lowercase_,lowercase_,atol=1E-5		)		)

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

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

        A__          =   test_model(lowercase_		)
        assert torch.allclose(lowercase_,lowercase_,atol=1E-5		)


    def 						snake_case__     (      self :       Optional[Any]		)->    Optional[Any]:





        '''simple docstring'''







        A__          =   ModelForTest()
        A__          =   torch.randn(2,3		)
        A__          =   test_model(lowercase_		)

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

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

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

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


    def 						snake_case__     (      self :       Dict		)->    Union[str, Any]:





        '''simple docstring'''







        A__          =   ModelForTest()
        A__          =   torch.randn(2,3		)
        A__          =   test_model(lowercase_		)

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

        A__          =   True
        A__          =   test_model(lowercase_		)
        self.assertFalse(outputa.requires_grad		)


    @require_multi_gpu
    def 						snake_case__     (      self :       Optional[int]		)->    Any:





        '''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(lowercase_		)
        self.assertEqual(output.device,torch.device(1		)		)

        # We can add a general hook to put back output on same device as input.
        add_hook_to_module(lowercase_,AlignDevicesHook(io_same_device=lowercase_		)		)
        A__          =   torch.randn(2,3		).to(0		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,torch.device(0		)		)


    def 						snake_case__     (      self :       Dict		)->    Optional[Any]:





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

        # Parameters have been offloaded, so on the meta device
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        # Buffers are not included in the offload by default, so are on the execution device
        A__          =   torch.device(hook_kwargs['execution_device']		)
        self.assertEqual(model.batchnorm.running_mean.device,lowercase_		)

        A__          =   torch.randn(2,3		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,lowercase_		)

        # Removing hooks loads back the weights in the model.
        remove_hook_from_module(model.lineara		)
        remove_hook_from_module(model.batchnorm		)
        remove_hook_from_module(model.lineara		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)

        # Now test with buffers included in the offload
        A__          =   {
            'execution_device': 0 if torch.cuda.is_available() else 'cpu',
            'offload': True,
            'offload_buffers': True,
        }

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

        # Parameters have been offloaded, so on the meta device, buffers included
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.running_mean.device,torch.device('meta'		)		)

        A__          =   torch.randn(2,3		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,lowercase_		)

        # Removing hooks loads back the weights in the model.
        remove_hook_from_module(model.lineara		)
        remove_hook_from_module(model.batchnorm		)
        remove_hook_from_module(model.lineara		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)


    def 						snake_case__     (      self :       Tuple		)->    Any:





        '''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(lowercase_,execution_device=lowercase_,offload=lowercase_		)

        # Parameters have been offloaded, so on the meta device
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        # Buffers are not included in the offload by default, so are on the execution device
        A__          =   torch.device(lowercase_		)
        self.assertEqual(model.batchnorm.running_mean.device,lowercase_		)

        A__          =   torch.randn(2,3		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,lowercase_		)

        # Removing hooks loads back the weights in the model.
        remove_hook_from_submodules(lowercase_		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)

        # Now test with buffers included in the offload
        attach_align_device_hook(lowercase_,execution_device=lowercase_,offload=lowercase_,offload_buffers=lowercase_		)

        # Parameters have been offloaded, so on the meta device, buffers included
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.running_mean.device,torch.device('meta'		)		)

        A__          =   torch.randn(2,3		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,lowercase_		)

        # Removing hooks loads back the weights in the model.
        remove_hook_from_submodules(lowercase_		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)


    def 						snake_case__     (      self :       List[Any]		)->    Union[str, Any]:





        '''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(
            lowercase_,execution_device=lowercase_,offload=lowercase_,weights_map=model.state_dict()		)

        # Parameters have been offloaded, so on the meta device
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        # Buffers are not included in the offload by default, so are on the execution device
        A__          =   torch.device(lowercase_		)
        self.assertEqual(model.batchnorm.running_mean.device,lowercase_		)

        A__          =   torch.randn(2,3		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,lowercase_		)

        # Removing hooks loads back the weights in the model.
        remove_hook_from_submodules(lowercase_		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)

        # Now test with buffers included in the offload
        attach_align_device_hook(
            lowercase_,execution_device=lowercase_,offload=lowercase_,weights_map=model.state_dict(),offload_buffers=lowercase_,)

        # Parameters have been offloaded, so on the meta device, buffers included
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.running_mean.device,torch.device('meta'		)		)

        A__          =   torch.randn(2,3		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,lowercase_		)

        # Removing hooks loads back the weights in the model.
        remove_hook_from_submodules(lowercase_		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)








from typing import List, Optional, Tuple, Union

import torch

from ...utils import logging, randn_tensor
from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline


lowercase      :       Any									=  logging.get_logger(__name__)  # pylint: disable=invalid-name





class  A__    (							__UpperCAmelCase	):




   """simple docstring"""



   def __init__(					self					,  lowercase					,  lowercase)			->						Optional[Any]:







          '''simple docstring'''

          super().__init__()
          self.register_modules(unet=lowercase					,  scheduler=lowercase)

   @torch.no_grad()
   def __call__(					self					,  lowercase = 1					,  lowercase = 100					,  lowercase = None					,  lowercase = None					,  lowercase = True					,  )			->						Union[AudioPipelineOutput, Tuple]:







          '''simple docstring'''

          if audio_length_in_s is None:
                 a__						:	Optional[int] =				self.unet.config.sample_size / self.unet.config.sample_rate

          a__						:	int =				audio_length_in_s * self.unet.config.sample_rate

          a__						:	Union[str, Any] =				2 ** len(self.unet.up_blocks)
          if sample_size < 3 * down_scale_factor:
                 raise ValueError(
                     F'{audio_length_in_s} is too small. Make sure it\'s bigger or equal to'
                     F' {3 * down_scale_factor / self.unet.config.sample_rate}.')

          a__						:	str =				int(lowercase)
          if sample_size % down_scale_factor != 0:
                 a__						:	List[str] =				(
                     (audio_length_in_s * self.unet.config.sample_rate) // down_scale_factor + 1
                 ) * down_scale_factor
                 logger.info(
                     F'{audio_length_in_s} is increased to {sample_size / self.unet.config.sample_rate} so that it can be handled'
                     F' by the model. It will be cut to {original_sample_size / self.unet.config.sample_rate} after the denoising'
                     ' process.')
          a__						:	List[Any] =				int(lowercase)

          a__						:	int =				next(iter(self.unet.parameters())).dtype
          a__						:	Tuple =				(batch_size, self.unet.config.in_channels, sample_size)
          if isinstance(lowercase					,  lowercase) and len(lowercase) != batch_size:
                 raise ValueError(
                     F'You have passed a list of generators of length {len(lowercase)}, but requested an effective batch'
                     F' size of {batch_size}. Make sure the batch size matches the length of the generators.')

          a__						:	Optional[Any] =				randn_tensor(lowercase					,  generator=lowercase					,  device=self.device					,  dtype=lowercase)

          # set step values
          self.scheduler.set_timesteps(lowercase					,  device=audio.device)
          a__						:	Union[str, Any] =				self.scheduler.timesteps.to(lowercase)

          for t in self.progress_bar(self.scheduler.timesteps):
                 # 1. predict noise model_output
                 a__						:	Dict =				self.unet(lowercase					,  lowercase).sample

                 # 2. compute previous image: x_t -> t_t-1
                 a__						:	Any =				self.scheduler.step(lowercase					,  lowercase					,  lowercase).prev_sample

          a__						:	str =				audio.clamp(-1					,  1).float().cpu().numpy()

          a__						:	List[Any] =				audio[:, :, :original_sample_size]

          if not return_dict:
                 return (audio,)

          return AudioPipelineOutput(audios=lowercase)


